From 8808f8e23926fb1575694504f03a6e1ac291391a Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Sun, 23 Aug 2026 19:53:19 +0800 Subject: [PATCH 001/105] [Kernel][SM70] Start XQA pipeline follow-up Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> From 680a7dbb4f102fded707ebdaa80d56331da5270f Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 25 Aug 2026 22:11:32 +0000 Subject: [PATCH 002/105] Bump sentence-transformers from 5.2.0 to 6.0.0 Bumps [sentence-transformers](https://github.com/huggingface/sentence-transformers) from 5.2.0 to 6.0.0. - [Release notes](https://github.com/huggingface/sentence-transformers/releases) - [Commits](https://github.com/huggingface/sentence-transformers/compare/v5.2.0...v6.0.0) --- updated-dependencies: - dependency-name: sentence-transformers dependency-version: 6.0.0 dependency-type: direct:production update-type: version-update:semver-major ... Signed-off-by: dependabot[bot] --- requirements/test/cuda.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index 7470b91afd..fa127cd981 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -890,7 +890,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 From 626305e9167e4b437fd582a50440b757280cc001 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 25 Aug 2026 22:11:44 +0000 Subject: [PATCH 003/105] Bump pillow from 10.4.0 to 12.3.0 Bumps [pillow](https://github.com/python-pillow/Pillow) from 10.4.0 to 12.3.0. - [Release notes](https://github.com/python-pillow/Pillow/releases) - [Changelog](https://github.com/python-pillow/Pillow/blob/main/CHANGES.rst) - [Commits](https://github.com/python-pillow/Pillow/compare/10.4.0...12.3.0) --- updated-dependencies: - dependency-name: pillow dependency-version: 12.3.0 dependency-type: direct:production update-type: version-update:semver-major ... Signed-off-by: dependabot[bot] --- requirements/test/cuda.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index 7470b91afd..71ee766bfb 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -616,7 +616,7 @@ 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 From b7c681e6c686d5bb13d053628fa4b0294049db67 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 25 Aug 2026 22:11:49 +0000 Subject: [PATCH 004/105] Bump black from 24.10.0 to 26.5.1 Bumps [black](https://github.com/psf/black) from 24.10.0 to 26.5.1. - [Release notes](https://github.com/psf/black/releases) - [Changelog](https://github.com/psf/black/blob/main/CHANGES.md) - [Commits](https://github.com/psf/black/compare/24.10.0...26.5.1) --- updated-dependencies: - dependency-name: black dependency-version: 26.5.1 dependency-type: direct:production update-type: version-update:semver-major ... Signed-off-by: dependabot[bot] --- requirements/test/cuda.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index 7470b91afd..c4e84d8a8a 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -62,7 +62,7 @@ backoff==2.2.1 # schemathesis bitsandbytes==0.49.2 # 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 From 6cdfc4430aa44bb6e50000ca420173603b147d12 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 26 Aug 2026 10:31:35 +0800 Subject: [PATCH 005/105] [CI/Build] Restore repository pre-commit baseline (#311) Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .pre-commit-config.yaml | 12 +- benchmarks/benchmark_dspark_api.py | 1 + benchmarks/benchmark_dspark_stream.py | 1 + .../benchmark_sm70_flashinfer_fixed_entry.py | 20 +- ...hmark_sm70_flashqla_indexed_state_cache.py | 48 +- .../benchmark_sm70_fp8_fixed_vs_dynamic.py | 17 +- benchmarks/benchmark_sm70_hmma_c2a_mincomm.py | 3 +- .../benchmark_sm70_instruction_latency.py | 3 +- benchmarks/benchmark_sm70_lm_head_top1.py | 18 +- .../benchmark_sm70_long_quality_suffix.py | 6 +- ..._sm70_native_bm32_allp_crossblock_micro.py | 60 +- ...ark_sm70_native_bm32_allp_scratch_micro.py | 73 +- ..._sm70_native_bm32_full_phase_allp_micro.py | 64 +- ...hmark_sm70_native_bm32_full_phase_micro.py | 64 +- ...nchmark_sm70_native_bm32_pv_panel_micro.py | 17 +- benchmarks/benchmark_sm70_raw_hmma_probe.py | 5 +- .../benchmark_sm70_rejection_sampler_micro.py | 34 +- benchmarks/benchmark_sm70_source_inventory.py | 88 +- benchmarks/benchmark_sm70_turbomind_parity.py | 33 +- .../benchmark_sm70_wmma_fragment_probe.py | 38 +- benchmarks/compare_sm70_tensor_dumps.py | 156 +- benchmarks/kernels/benchmark_dspark_markov.py | 1 + .../kernels/benchmark_dspark_metadata.py | 1 + .../benchmark_dspark_sm70_attention.py | 1 + .../benchmark_sm70_fp8_graph_capture.py | 1 + benchmarks/verify_sm70_artifact.py | 168 +- .../marlin_moe_wna16/sm70_marlin_fp8_gemm.cu | 197 +- .../moe/marlin_moe_wna16/sm70_marlin_gemm.cuh | 1517 +++++------ .../sm70_marlin_moe_dispatch.cu | 98 +- .../sm70_marlin_mxfp4_gemm.cu | 128 +- .../sm70_marlin_nvfp4_gemm.cu | 128 +- .../marlin_moe_wna16/sm70_marlin_u4_gemm.cu | 189 +- .../marlin_moe_wna16/sm70_marlin_u4b8_gemm.cu | 167 +- .../marlin_moe_wna16/sm70_marlin_u8_gemm.cu | 211 +- .../sm70_marlin_u8b128_gemm.cu | 191 +- csrc/moe/moe_permute_unpermute_op.cu | 17 +- .../moe_permute_unpermute_kernel.cu | 27 +- .../moe_permute_unpermute_kernel.h | 11 +- .../marlin/sm70_awq_marlin_repack.cu | 32 +- .../marlin/sm70_gptq_marlin_repack.cu | 34 +- .../marlin/sm70_marlin_common.cuh | 600 ++-- .../marlin/sm70_marlin_dispatch.cu | 129 +- .../marlin/sm70_marlin_fp8_gemm.cu | 220 +- csrc/quantization/marlin/sm70_marlin_gemm.cuh | 172 +- .../marlin/sm70_marlin_iterator_utils.cuh | 34 +- csrc/quantization/marlin/sm70_marlin_mma.cuh | 45 +- .../marlin/sm70_marlin_mxfp4_gemm.cu | 212 +- .../marlin/sm70_marlin_nvfp4_gemm.cu | 210 +- .../marlin/sm70_marlin_splitk.cuh | 50 +- .../marlin/sm70_marlin_u4_gemm.cu | 280 +- .../marlin/sm70_marlin_u4b8_gemm.cu | 218 +- .../marlin/sm70_marlin_u8_gemm.cu | 251 +- .../marlin/sm70_marlin_u8b128_gemm.cu | 222 +- csrc/sm70_turbomind/ops/tm_registry_sm70.cu | 75 +- ...6_27b_awq_mtp_target_verifier_fastpaths.md | 36 +- ...0_nvfp4_turbomind_operator_optimization.md | 10 +- docs/design/sm70_tile_runtime_exploration.md | 20 +- flash-attention-v100/include/fused_mma.h | 626 ++--- .../kernel/contiguous_to_paged.cu | 123 +- .../kernel/flash_decode_turboquant.cu | 365 ++- .../kernel/flash_v100_traits.cuh | 227 +- flash-attention-v100/kernel/fp8_kv_utils.cuh | 38 +- .../kernel/fused_mha_backward.cu | 2407 +++++++++-------- .../kernel/fused_mha_forward.cu | 2030 +++++++------- .../kernel/fused_mha_forward_paged_full.cu | 1082 ++++---- .../kernel/paged_kv_utils.cuh | 98 +- .../kernel/paged_to_contiguous.cu | 355 ++- .../kernel/paged_to_contiguous_fixed.cu | 167 +- .../kernel/paged_to_contiguous_old.cu | 140 +- requirements/docs.txt | 1 - .../test_sm70_flash_v100_splitkv3.py | 12 +- .../kernels/moe/test_moe_permute_unpermute.py | 12 +- ...t_deepseek_v4_sm70_qnorm_rope_kv_insert.py | 5 +- .../test_fused_recurrent_packed_decode.py | 12 +- .../test_qwen3coder_tool_parser.py | 2 +- .../spec_decode/test_ddtree_flex_attention.py | 10 +- .../test_ddtree_gpu_runner_positions.py | 12 +- tests/v1/spec_decode/test_ddtree_metadata.py | 3 +- tools/pre_commit/check_forbidden_imports.py | 7 +- tools/pre_commit/check_torch_cuda.py | 10 + .../passes/fusion/matcher_utils.py | 3 +- .../passes/fusion/qk_norm_rope_fusion.py | 3 +- .../passes/fusion/rms_quant_fusion.py | 9 +- .../passes/utility/fix_functionalization.py | 4 +- .../layers/fla/ops/chunk_scaled_dot_kkt.py | 4 +- .../fla/ops/fused_gdn_prefill_post_conv.py | 3 +- .../layers/fla/ops/fused_recurrent.py | 4 +- .../layers/quantization/auto_gptq.py | 12 +- .../layers/quantization/awq_marlin.py | 4 +- .../schemes/compressed_tensors_w4a16_nvfp4.py | 3 +- .../schemes/compressed_tensors_w4a4_mxfp4.py | 8 +- .../layers/quantization/input_quant_fp8.py | 4 +- .../quantization/utils/marlin_utils_fp4.py | 5 +- .../layers/quantization/utils/w8a8_utils.py | 9 +- vllm/model_executor/models/config.py | 4 +- vllm/model_executor/models/qwen2_moe.py | 4 +- vllm/model_executor/models/qwen3_5_mtp.py | 3 +- vllm/tool_parsers/qwen3coder_tool_parser.py | 5 +- .../backends/ddtree_branch_triton.py | 21 +- vllm/v1/attention/backends/fa_utils.py | 1 + vllm/v1/attention/backends/registry.py | 8 +- .../attention/ops/triton_prefill_attention.py | 6 +- .../attention/ops/triton_unified_attention.py | 12 +- vllm/v1/engine/core.py | 16 +- vllm/v1/sample/thinking_budget_state.py | 8 +- vllm/v1/spec_decode/ddtree_metadata.py | 70 +- vllm/v1/worker/block_table.py | 4 +- 107 files changed, 7126 insertions(+), 7519 deletions(-) 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/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_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_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_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_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_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_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_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/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/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/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/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_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_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/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_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_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/requirements/docs.txt b/requirements/docs.txt index a7affa70ea..abf2e75c49 100644 --- a/requirements/docs.txt +++ b/requirements/docs.txt @@ -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/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_moe_permute_unpermute.py b/tests/kernels/moe/test_moe_permute_unpermute.py index 662da0fb62..d68b53fd86 100644 --- a/tests/kernels/moe/test_moe_permute_unpermute.py +++ b/tests/kernels/moe/test_moe_permute_unpermute.py @@ -473,9 +473,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 ) @@ -483,12 +481,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/test_deepseek_v4_sm70_qnorm_rope_kv_insert.py b/tests/kernels/test_deepseek_v4_sm70_qnorm_rope_kv_insert.py index 8dff20c955..66a81a4287 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 @@ -25,10 +25,7 @@ def _make_cos_sin_cache(max_position: int, device: str) -> torch.Tensor: inv_freq = 1.0 / ( 10000.0 - ** ( - torch.arange(0, ROPE_DIM, 2, dtype=torch.float32, device=device) - / ROPE_DIM - ) + ** (torch.arange(0, ROPE_DIM, 2, dtype=torch.float32, device=device) / ROPE_DIM) ) positions = torch.arange(max_position, dtype=torch.float32, device=device) frequencies = torch.outer(positions, inv_freq) 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/tool_parsers/test_qwen3coder_tool_parser.py b/tests/tool_parsers/test_qwen3coder_tool_parser.py index ca2a60d4ce..9167493688 100644 --- a/tests/tool_parsers/test_qwen3coder_tool_parser.py +++ b/tests/tool_parsers/test_qwen3coder_tool_parser.py @@ -1384,7 +1384,7 @@ def test_streaming_multi_param_single_chunk(qwen3_tool_parser, qwen3_tokenizer): @pytest.mark.parametrize( "deltas, expected_content", [ - (["<", "meta charset=\"UTF-8\">"], ''), + (["<", 'meta charset="UTF-8">'], ''), (["macOS"], "macOS"), (["<", "<", "span>text"], "<text"), ], 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/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/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/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..35491a71e0 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") 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/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/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/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/models/config.py b/vllm/model_executor/models/config.py index d48ec981af..de9786c777 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 " diff --git a/vllm/model_executor/models/qwen2_moe.py b/vllm/model_executor/models/qwen2_moe.py index f27c4f3037..ecfbf4f064 100644 --- a/vllm/model_executor/models/qwen2_moe.py +++ b/vllm/model_executor/models/qwen2_moe.py @@ -149,9 +149,7 @@ 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) 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/tool_parsers/qwen3coder_tool_parser.py b/vllm/tool_parsers/qwen3coder_tool_parser.py index 2b21f8f6ae..7a3a63ad53 100644 --- a/vllm/tool_parsers/qwen3coder_tool_parser.py +++ b/vllm/tool_parsers/qwen3coder_tool_parser.py @@ -358,10 +358,7 @@ def extract_tool_calls_streaming( if pos != -1 ] tool_start = min(tool_start_candidates) if tool_start_candidates else -1 - if ( - self.tool_call_start_token_id in delta_token_ids - or tool_start != -1 - ): + if self.tool_call_start_token_id in delta_token_ids or tool_start != -1: self.is_tool_call_started = True # Return any content before the tool call if tool_start > len(previous_text): 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/registry.py b/vllm/v1/attention/backends/registry.py index a96c0dd115..9688a42b64 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" ) 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/engine/core.py b/vllm/v1/engine/core.py index 4e094eb6a6..c1f11a69ba 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -461,9 +461,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 +476,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 +493,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 +605,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/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/worker/block_table.py b/vllm/v1/worker/block_table.py index 616fdd15bd..8329a2c196 100644 --- a/vllm/v1/worker/block_table.py +++ b/vllm/v1/worker/block_table.py @@ -104,9 +104,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, From 1d64f84dce6a42c1496c42e0e4c021ac9f979591 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 26 Aug 2026 10:40:38 +0800 Subject: [PATCH 006/105] [Bugfix][SM70] Bound Gemma4 tower attention memory (#310) Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_v100_migration_control.md | 39 ++++++++++++++ .../multimodal/processing/test_gemma4.py | 46 ++++++++++++++++ vllm/model_executor/models/gemma4_mm.py | 54 +++++++++++++++++-- 3 files changed, 135 insertions(+), 4 deletions(-) diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index 0071129b13..3c43369c88 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -43565,3 +43565,42 @@ 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. 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/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) From f5fc157c49716c3cb9cbe269d607f38d669be479 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 26 Aug 2026 10:56:44 +0800 Subject: [PATCH 007/105] [Model][SM70] Exclude unsafe PP2 TP4 QPN8 gate/up (#313) * [Model][SM70] Prefer accurate PP2 TP4 QPN8 schedules Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> (cherry picked from commit a253e045a60ada82eae27294188d09ddb4dc4cd4) * [Doc] Record DeepSeek V4 QPN8 accuracy audit Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> (cherry picked from commit bdc82e879ff547c16e6acad6321c50e3ccf0c9b4) * [Model][SM70] Preserve bounded QPN8 speed schedules Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --------- Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_v100_migration_control.md | 42 +++++++++++++++++++ .../test_sm70_fp8_qpn8_pp2_tp4.py | 29 ++----------- .../model_executor/layers/quantization/fp8.py | 14 +++---- 3 files changed, 53 insertions(+), 32 deletions(-) diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index 3c43369c88..ccee797c5f 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -43519,6 +43519,48 @@ Interpretation: 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 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. +- 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. That single sample remains visible, but + it is not a greedy-identity gate and does not show that choosing the lowest + operator error improves aggregate task quality. The accepted repair + therefore preserves every numerically bounded speed winner and falls back + only the demonstrated gate/up outlier. 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 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 d65cf59fca..17580656d8 100644 --- a/tests/model_executor/test_sm70_fp8_qpn8_pp2_tp4.py +++ b/tests/model_executor/test_sm70_fp8_qpn8_pp2_tp4.py @@ -55,28 +55,6 @@ def test_pp2_tp4_qpn8_is_default_on_with_explicit_rollback(monkeypatch) -> None: (_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( - "gate_up_proj", - 4, - 4096, - 1024, - output_partition_sizes=[512, 512], - ), - False, - (32, 2, False), - ), - ( - _layer( - "gate_up_proj", - 4, - 4096, - 1024, - output_partition_sizes=[512, 512], - ), - True, - (16, 2, False), - ), (_layer("down_proj", 4, 512, 4096), False, (16, 2, False)), ], ) @@ -97,14 +75,15 @@ def test_pp2_tp4_qpn8_rejects_wrong_tensor_and_concurrency_roles() -> None: fp8._sm70_fp8_qpn8_pp2_tp4_config(concurrent_indexer, gated_silu=False) is None ) - wrong_gate = _layer( + excluded_gate = _layer( "gate_up_proj", 4, 4096, 1024, - output_partition_sizes=[256, 768], + output_partition_sizes=[512, 512], ) - assert fp8._sm70_fp8_qpn8_pp2_tp4_config(wrong_gate, gated_silu=True) is None + assert fp8._sm70_fp8_qpn8_pp2_tp4_config(excluded_gate, gated_silu=False) is None + assert fp8._sm70_fp8_qpn8_pp2_tp4_config(excluded_gate, gated_silu=True) is None wrong_layout = _layer("wo_b", 4, 2048, 4096) wrong_layout.weight_block_size = [64, 128] diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 902501334b..36a7a8097b 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -132,23 +132,23 @@ "qkv_proj": (5120, 3584), } _SM70_FP8_QPN8_PP2_TP4_CONFIGS = { - # Measured B1 winners: (K, N, fused gated-SiLU) maps to - # (split-K, accumulator chains, prefetch codes). + # Real-weight speed winners remain numerically bounded for every admitted + # projection. The shared-expert gate/up role is deliberately absent below + # after its fused activation matched only 53.5% of TurboMind FP16 elements. (4096, 1536, False): (32, 2, False), (1024, 8192, False): (8, 2, False), (2048, 4096, False): (16, 2, False), (4096, 1024, False): (32, 2, False), - (4096, 1024, True): (16, 2, False), (512, 4096, False): (16, 2, False), } _SM70_FP8_QPN8_PP2_TP4_SHAPES = { - # Operator role: accepted (layer TP size, K, N) tuples. The replicated - # indexer wq_b is deliberately 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. 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. "fused_wqa_wkv": {(1, 4096, 1536)}, "wq_b": {(4, 1024, 8192)}, "wo_b": {(4, 2048, 4096)}, - "gate_up_proj": {(4, 4096, 1024)}, "down_proj": {(4, 512, 4096)}, } _SM70_FP8_QPN8_PP2_TP4_WORKSPACE_ELEMENTS = max( From 1b1c9faeb82c30bf3d81158279404cd884bbfb19 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 26 Aug 2026 11:16:04 +0800 Subject: [PATCH 008/105] [Kernel][SM70] Default audited PP2 TP4 decode hot paths (#314) * [Kernel] Tune sparse MLA head grouping for TP4 Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> (cherry picked from commit 4397e7adab5a44f483585ced6c5fce78e9d344c2) * [Kernel] Fuse TP4 Q transform and KV insertion Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> (cherry picked from commit 19e64ab107e247d77bdc92cce39a48f4dd9bf909) * [Kernel] Broadcast MXFP4 W13 input for batch one Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> (cherry picked from commit 089feeae7241daaaf7ccb2a4fb2a821925be5350) * [Kernel] Remove redundant MXFP4 top-6 sorting Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> (cherry picked from commit d1e8c1c40ba1b2d76b2598962fcdcc94f9e390b1) * [Kernel][SM70] Default audited PP2 TP4 hot paths Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> * [Kernel][SM70] Align compact decode defaults Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --------- Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- ...benchmark_sm70_mxfp4_moe_active_experts.py | 99 +++++++++++-- csrc/sm70_turbomind/ops/awq_sm70_gemm.cu | 90 ++++++++---- docs/design/sm70_v100_migration_control.md | 32 +++++ .../kernels/moe/test_moe_permute_unpermute.py | 1 + ...t_deepseek_v4_sm70_qnorm_rope_kv_insert.py | 42 ++++++ tests/models/test_deepseek_v4_sm70_routes.py | 8 ++ tests/quantization/test_sm70_mxfp4_moe.py | 41 +++--- tests/test_envs.py | 22 +++ vllm/envs.py | 24 +++- .../layers/quantization/mxfp4_sm70_moe.py | 48 +++++++ .../sm70/qnorm_rope_kv_fp8_insert.py | 132 ++++++++++++++++++ vllm/models/deepseek_v4/sm70/sparse.py | 6 + 12 files changed, 487 insertions(+), 58 deletions(-) diff --git a/benchmarks/benchmark_sm70_mxfp4_moe_active_experts.py b/benchmarks/benchmark_sm70_mxfp4_moe_active_experts.py index f2ac568233..b79a169e61 100644 --- a/benchmarks/benchmark_sm70_mxfp4_moe_active_experts.py +++ b/benchmarks/benchmark_sm70_mxfp4_moe_active_experts.py @@ -13,6 +13,7 @@ import json import os from dataclasses import dataclass +from pathlib import Path import torch @@ -399,6 +400,7 @@ def make_buffers() -> dict[str, torch.Tensor]: generic = make_buffers() direct = make_buffers() + direct_order = make_buffers() def generic_call() -> None: generic["output"].zero_() @@ -498,11 +500,61 @@ def direct_call() -> None: direct["output"], ) + def direct_order_call() -> None: + route_ids = topk_ids.view(-1) + sm70_ops.mxfp4_moe_dense_stage_sm70_out( + direct_order["gate_up"], + x, + direct_order["compact_offsets"], + route_ids, + w13_ptrs_w, + w13_ptrs_s, + top_k, + 4096, + 512, + 32, + ) + torch.ops._C.silu_and_mul_with_clamp( + direct_order["intermediate"], direct_order["gate_up"], 10.0 + ) + sm70_ops.mxfp4_moe_dense_stage_sm70_out( + direct_order["sorted_output"], + direct_order["intermediate"], + direct_order["compact_offsets"], + route_ids, + w2_ptrs_w, + w2_ptrs_s, + top_k, + 256, + 4096, + 32, + ) + sm70_ops.awq_moe_single_token_weighted_reduce_out( + direct_order["sorted_output"], + topk_weights, + direct_order["token_expert_indices"], + direct_order["output"], + top_k, + 4096, + ) + generic_call() direct_call() + direct_order_call() torch.accelerator.synchronize() initial_equal = torch.equal(generic["output"], direct["output"]) initial_max_abs = float((generic["output"] - direct["output"]).abs().max().item()) + direct_order_initial_equal = torch.equal(direct["output"], direct_order["output"]) + direct_order_initial_max_abs = float( + (direct["output"] - direct_order["output"]).abs().max().item() + ) + initial_sorted_src = torch.argsort(topk_ids.view(-1), stable=True) + direct_order_stage_parity = { + f"{name}_equal": torch.equal( + direct[name], direct_order[name].index_select(0, initial_sorted_src) + ) + for name in ("gate_up", "intermediate", "sorted_output") + } stage_parity = { "gate_up_equal": torch.equal(generic["gate_up"], direct["gate_up"]), "gate_up_max_abs": float( @@ -527,28 +579,52 @@ def direct_call() -> None: generic_graph = _capture(generic_call) direct_graph = _capture(direct_call) + direct_order_graph = _capture(direct_order_call) generic_ms = _time_graph(generic_graph, repeats) direct_ms = _time_graph(direct_graph, repeats) + direct_order_ms = _time_graph(direct_order_graph, repeats) topk_ids.copy_(torch.tensor([route_b], dtype=torch.int32, device=device)) generic_graph.replay() direct_graph.replay() + direct_order_graph.replay() torch.accelerator.synchronize() replay_equal = torch.equal(generic["output"], direct["output"]) replay_max_abs = float((generic["output"] - direct["output"]).abs().max().item()) + direct_order_replay_equal = torch.equal(direct["output"], direct_order["output"]) + direct_order_replay_max_abs = float( + (direct["output"] - direct_order["output"]).abs().max().item() + ) result = { "generic_graph_ms": generic_ms, "direct_graph_ms": direct_ms, + "direct_order_graph_ms": direct_order_ms, "speedup": generic_ms / direct_ms, "projected_savings_ms_per_token": (generic_ms - direct_ms) * 43, + "direct_order_speedup_vs_direct": direct_ms / direct_order_ms, + "direct_order_projected_savings_ms_per_token": (direct_ms - direct_order_ms) + * 43, "initial_bitwise_equal": initial_equal, "initial_max_abs": initial_max_abs, "dynamic_replay_bitwise_equal": replay_equal, "dynamic_replay_max_abs": replay_max_abs, + "direct_order_initial_bitwise_equal": direct_order_initial_equal, + "direct_order_initial_max_abs": direct_order_initial_max_abs, + "direct_order_dynamic_replay_bitwise_equal": direct_order_replay_equal, + "direct_order_dynamic_replay_max_abs": direct_order_replay_max_abs, "initial_stage_parity": stage_parity, + "direct_order_initial_sorted_stage_parity": direct_order_stage_parity, } - if not initial_equal or not replay_equal: + if not all( + ( + initial_equal, + replay_equal, + direct_order_initial_equal, + direct_order_replay_equal, + *direct_order_stage_parity.values(), + ) + ): raise RuntimeError(f"MXFP4 direct top-6 correctness gate failed: {result}") return result @@ -1067,6 +1143,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--seed", type=int, default=17) parser.add_argument("--input-scale", type=float, default=0.01) parser.add_argument("--full-pipeline", action="store_true") + parser.add_argument("--json-out", type=Path) parser.add_argument("--verifier-m8-pipeline", action="store_true") parser.add_argument( "--verifier-tokens", @@ -1141,16 +1218,16 @@ def main() -> int: repeats=args.repeats, seed=args.seed, ) - print( - json.dumps( - { - "benchmark": "sm70_mxfp4_moe_direct_top6_pipeline", - "device": torch.cuda.get_device_name(), - "result": result, - }, - indent=2, - ) - ) + payload = { + "benchmark": "sm70_mxfp4_moe_direct_top6_pipeline", + "device": torch.cuda.get_device_name(), + "result": result, + } + rendered = json.dumps(payload, indent=2) + if args.json_out is not None: + args.json_out.parent.mkdir(parents=True, exist_ok=True) + args.json_out.write_text(rendered + "\n", encoding="utf-8") + print(rendered) return 0 if args.profile_active_once: if args.stage == "both": diff --git a/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu b/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu index 2fedcf4b91..0aa6899e49 100644 --- a/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu +++ b/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu @@ -1074,7 +1074,12 @@ 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() { + const char* raw = std::getenv("VLLM_SM70_MXFP4_MOE_BROADCAST_INPUT_DECODE"); + return raw == nullptr || std::atoi(raw) != 0; } bool mxfp4_moe_grouped_m8_enabled() { @@ -6633,17 +6638,29 @@ void awq_moe_single_token_weighted_reduce_out(torch::Tensor sorted_output, const int hidden_logical_size_i = static_cast(hidden_logical_size); const int sorted_output_row_stride = static_cast(sorted_output.stride(0)); - if (top_k == 8 && (hidden_logical_size_i % 2) == 0 && + if ((top_k == 6 || top_k == 8) && (hidden_logical_size_i % 2) == 0 && (sorted_output_row_stride % 2) == 0) { const int blocks = std::max( 1, ((hidden_logical_size_i >> 1) + kThreads - 1) / kThreads); - awq_moe_single_token_weighted_reduce_half2_kernel<8> - <<>>( - reinterpret_cast(sorted_output.data_ptr()), - topk_weights.data_ptr(), - inv_permuted_idx.data_ptr(), - reinterpret_cast<__half*>(out.data_ptr()), - hidden_logical_size_i, sorted_output_row_stride); + if (top_k == 6) { + awq_moe_single_token_weighted_reduce_half2_kernel<6> + <<>>( + reinterpret_cast( + sorted_output.data_ptr()), + topk_weights.data_ptr(), + inv_permuted_idx.data_ptr(), + reinterpret_cast<__half*>(out.data_ptr()), + hidden_logical_size_i, sorted_output_row_stride); + } else { + awq_moe_single_token_weighted_reduce_half2_kernel<8> + <<>>( + reinterpret_cast( + sorted_output.data_ptr()), + topk_weights.data_ptr(), + inv_permuted_idx.data_ptr(), + reinterpret_cast<__half*>(out.data_ptr()), + hidden_logical_size_i, sorted_output_row_stride); + } } else { const int blocks = std::max(1, (hidden_logical_size_i + kThreads - 1) / kThreads); @@ -8090,7 +8107,8 @@ void mxfp4_moe_gemm_sm70_out_impl( torch::Tensor out, torch::Tensor sorted_input, torch::Tensor expert_offsets, torch::Tensor strided_ptrs_w, torch::Tensor strided_ptrs_s, int64_t num_experts, int64_t k, int64_t n, int64_t group_size, - torch::Tensor b_group_indices, bool compact_grouped_rows = false) { + torch::Tensor b_group_indices, bool compact_grouped_rows = false, + bool broadcast_input_rows = false) { TORCH_CHECK( sorted_input.is_cuda() && sorted_input.scalar_type() == torch::kFloat16, "mxfp4_moe_gemm_sm70: input must be CUDA float16."); @@ -8111,9 +8129,16 @@ void mxfp4_moe_gemm_sm70_out_impl( "mxfp4_moe_gemm_sm70: k must be divisible by group_size."); TORCH_CHECK(sorted_input.dim() == 2 && sorted_input.size(1) == k, "mxfp4_moe_gemm_sm70: input shape mismatch."); - TORCH_CHECK(out.dim() == 2 && out.size(0) == sorted_input.size(0) && + const int64_t total_tokens = + broadcast_input_rows ? out.size(0) : sorted_input.size(0); + TORCH_CHECK(out.dim() == 2 && out.size(0) == total_tokens && out.size(1) == n && out.stride(1) == 1, "mxfp4_moe_gemm_sm70: output must be contiguous [tokens, n]."); + TORCH_CHECK(!broadcast_input_rows || + (sorted_input.size(0) == 1 && total_tokens == num_experts && + compact_grouped_rows), + "mxfp4_moe_gemm_sm70: broadcast input requires one physical " + "row and one logical row per compact group."); TORCH_CHECK(expert_offsets.numel() >= num_experts + 1, "mxfp4_moe_gemm_sm70: expert_offsets too small."); TORCH_CHECK(b_group_indices.is_cuda() && @@ -8123,7 +8148,6 @@ void mxfp4_moe_gemm_sm70_out_impl( "mxfp4_moe_gemm_sm70: B group indices must be contiguous CUDA " "int32."); - const int64_t total_tokens = sorted_input.size(0); if (total_tokens == 0) { return; } @@ -8147,7 +8171,8 @@ void mxfp4_moe_gemm_sm70_out_impl( static_cast(sorted_input.stride(0)), }; desc_A.num = static_cast(num_experts); - desc_A.offsets = expert_offsets.data_ptr(); + desc_A.offsets = + broadcast_input_rows ? nullptr : expert_offsets.data_ptr(); turbomind::gemm::MatrixLayout desc_U{}; const auto order_w = conv_w->order; @@ -8273,8 +8298,12 @@ void mxfp4_moe_dense_stage_sm70_out(torch::Tensor out, torch::Tensor input, "supported."); TORCH_CHECK(input.dim() == 2 && input.size(1) == k, "mxfp4_moe_dense_stage_sm70_out: input shape mismatch."); + 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(0) == input.size(0) && out.size(1) == n, + 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."); @@ -8291,11 +8320,12 @@ void mxfp4_moe_dense_stage_sm70_out(torch::Tensor out, torch::Tensor input, num_experts == 6 && ((k == 4096 && (n == 512 || n == 1024)) || (n == 4096 && (k == 256 || k == 512))); - if (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); + 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); return; } const bool grouped_m8_shape = @@ -8629,6 +8659,9 @@ void mxfp4_moe_single_token_prepare_w13_sm70_out( x.size(0), kTopK); constexpr int kThreads = 256; + const bool broadcast_input = + w13_n == 512 && + vllm::awq_sm70::mxfp4_moe_broadcast_input_decode_enabled(); const int src_row_stride = kPtrRowBytes; const int row_bytes = kPtrRowBytes; if (topk_ids.scalar_type() == torch::kInt32) { @@ -8639,8 +8672,10 @@ void mxfp4_moe_single_token_prepare_w13_sm70_out( w13_ptrs_w.data_ptr(), w13_ptrs_s.data_ptr(), w13_ptrs_w.data_ptr(), w13_ptrs_s.data_ptr(), w13_ptrs_w.data_ptr(), w13_ptrs_s.data_ptr(), - reinterpret_cast<__half*>(compact_input.data_ptr()), nullptr, - expert_offsets.data_ptr(), nullptr, + broadcast_input + ? nullptr + : reinterpret_cast<__half*>(compact_input.data_ptr()), + nullptr, expert_offsets.data_ptr(), nullptr, inv_permuted_idx.data_ptr(), sorted_expert_ids.data_ptr(), kTopK, static_cast(hidden_logical_size), src_row_stride, src_row_stride, @@ -8653,8 +8688,10 @@ void mxfp4_moe_single_token_prepare_w13_sm70_out( w13_ptrs_w.data_ptr(), w13_ptrs_s.data_ptr(), w13_ptrs_w.data_ptr(), w13_ptrs_s.data_ptr(), w13_ptrs_w.data_ptr(), w13_ptrs_s.data_ptr(), - reinterpret_cast<__half*>(compact_input.data_ptr()), nullptr, - expert_offsets.data_ptr(), nullptr, + broadcast_input + ? nullptr + : reinterpret_cast<__half*>(compact_input.data_ptr()), + nullptr, expert_offsets.data_ptr(), nullptr, inv_permuted_idx.data_ptr(), sorted_expert_ids.data_ptr(), kTopK, static_cast(hidden_logical_size), src_row_stride, src_row_stride, @@ -8662,9 +8699,10 @@ void mxfp4_moe_single_token_prepare_w13_sm70_out( } C10_CUDA_KERNEL_LAUNCH_CHECK(); - mxfp4_moe_gemm_sm70_out_impl(gate_up, compact_input, expert_offsets, - w13_ptrs_w, w13_ptrs_s, kTopK, w13_k, w13_n, - group_size, sorted_expert_ids, true); + mxfp4_moe_gemm_sm70_out_impl(gate_up, broadcast_input ? x : compact_input, + expert_offsets, w13_ptrs_w, w13_ptrs_s, kTopK, + w13_k, w13_n, group_size, sorted_expert_ids, + true, broadcast_input); } void fp8_moe_single_token_dense_stage_sm70_out( diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index ccee797c5f..aaf87ea84a 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -43438,6 +43438,38 @@ Interpretation: (`BLOCK_N=2`, `BLOCK_K=1024`) is 0.5--12% slower across the five traced `N=64/256/512/1024/2048` shapes; smaller-K/four-row variants are slower and also change FP32 accumulation. The candidate is not retained. +- The next MXFP4 B1 screen removes redundant route sorting without changing + model arithmetic. The six unique top-k IDs become the compact TurboMind + group IDs directly, fixed offsets remain `[0, 1, ..., 6]`, W13 reads the + single physical input row through the accepted broadcast contract, and W2 + stays in original top-k order. The final weighted reduction therefore keeps + the same route-order FP32 FMA sequence without an inverse permutation. + Initial and changed-route CUDA Graph replays are bitwise at W13, clamped + SwiGLU, W2, and final output. Adding the exact top-k-6 `half2` reducer moves + 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/`. + 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 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 diff --git a/tests/kernels/moe/test_moe_permute_unpermute.py b/tests/kernels/moe/test_moe_permute_unpermute.py index d68b53fd86..dc596f8795 100644 --- a/tests/kernels/moe/test_moe_permute_unpermute.py +++ b/tests/kernels/moe/test_moe_permute_unpermute.py @@ -463,6 +463,7 @@ def test_sm70_single_token_weighted_reduce_matches_moe_unpermute(): ] cases = [ ("topk2_h256_contiguous", 2, 256, 0), + ("topk6_h4096_contiguous", 6, 4096, 0), ("topk8_h2048_contiguous", 8, 2048, 0), ("topk8_h7168_padded_stride", 8, 7168, 64), ] 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 66a81a4287..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 @@ -117,3 +117,45 @@ def test_sm70_kv_rope_insert_is_repeatable_across_physical_blocks(): torch.testing.assert_close(gathered, gathered_outputs[0], rtol=0, atol=0) for q_out in q_outputs[1:]: torch.testing.assert_close(q_out, q_outputs[0], rtol=0, atol=0) + + +def test_sm70_tp4_fused_qnorm_kv_insert_matches_legacy( + monkeypatch: pytest.MonkeyPatch, +): + torch.manual_seed(20260826) + device = "cuda" + num_heads = 16 + block_size = 256 + block_stride = round_up(block_size * HEAD_BYTES, TOKEN_DATA_BYTES) + storage_shape = (3 * block_stride,) + q_input = torch.randn(1, num_heads, HEAD_DIM, dtype=torch.float16, device=device) + kv = torch.randn(1, HEAD_DIM, dtype=torch.float16, device=device) + positions = torch.tensor([1024], dtype=torch.int64, device=device) + slot_mapping = torch.tensor([block_size + 3], dtype=torch.int64, device=device) + cos_sin_cache = _make_cos_sin_cache(2048, device) + + def run(enabled: bool) -> tuple[torch.Tensor, torch.Tensor]: + monkeypatch.setenv("VLLM_SM70_DSV4_QNORM_KV_FUSED_TP4", "1" if enabled else "0") + q = q_input.clone() + storage = torch.full(storage_shape, 0xA5, dtype=torch.uint8, device=device) + cache = storage.as_strided( + (3, block_size, HEAD_BYTES), + (block_stride, HEAD_BYTES, 1), + ) + sm70_qnorm_rope_kv_fp8_insert( + q, + kv, + cache, + slot_mapping, + positions, + cos_sin_cache, + eps=1e-6, + block_size=block_size, + ) + torch.accelerator.synchronize() + return q, storage + + legacy_q, legacy_storage = run(False) + fused_q, fused_storage = run(True) + torch.testing.assert_close(fused_q, legacy_q, rtol=0, atol=0) + torch.testing.assert_close(fused_storage, legacy_storage, rtol=0, atol=0) diff --git a/tests/models/test_deepseek_v4_sm70_routes.py b/tests/models/test_deepseek_v4_sm70_routes.py index 662690f8ab..a983b2b206 100644 --- a/tests/models/test_deepseek_v4_sm70_routes.py +++ b/tests/models/test_deepseek_v4_sm70_routes.py @@ -101,6 +101,14 @@ def test_sm70_sparse_qk_dsplit_uses_graph_workspace(): kwargs = qk_dsplit.call_args.kwargs assert kwargs["partial_qk"].shape == (1, 8, 8, 8, 16) assert kwargs["partial_probs"].shape == (1, 8, 8, 16) + assert kwargs["stage1_block_h"] == 8 + + +def test_sm70_sparse_qk_dsplit_uses_one_tp4_head_group(): + from vllm.models.deepseek_v4.sm70.sparse import _qk_dsplit_block_h + + assert _qk_dsplit_block_h(16) == 16 + assert _qk_dsplit_block_h(8) == 8 def test_sm75_does_not_select_sm70_impl(): diff --git a/tests/quantization/test_sm70_mxfp4_moe.py b/tests/quantization/test_sm70_mxfp4_moe.py index dc48cbf0e9..845e1a3809 100644 --- a/tests/quantization/test_sm70_mxfp4_moe.py +++ b/tests/quantization/test_sm70_mxfp4_moe.py @@ -433,7 +433,7 @@ def test_mxfp4_sm70_active_expert_compaction_replays_dynamic_routes(): ) -@pytest.mark.parametrize(("k", "n"), [(4096, 1024), (512, 4096)]) +@pytest.mark.parametrize(("k", "n"), [(4096, 512), (4096, 1024), (512, 4096)]) @pytest.mark.skipif( not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0), reason="requires NVIDIA V100/SM70", @@ -512,28 +512,35 @@ def test_mxfp4_sm70_tp4_compact_shapes_match_per_expert( torch.accelerator.synchronize() torch.testing.assert_close(actual, reference, rtol=0.0, atol=0.0) - if (k, n) == (4096, 1024): + if k == 4096: direct = torch.empty_like(actual) + legacy_direct = torch.empty_like(actual) direct_input = torch.empty_like(grouped_input) direct_offsets = torch.empty_like(offsets) inverse_indices = torch.empty(num_experts, dtype=torch.int32, device="cuda") direct_expert_ids = torch.empty_like(expert_ids) - torch.ops._C.mxfp4_moe_single_token_prepare_w13_sm70_out( - direct, - direct_input, - x, - expert_ids.view(1, num_experts), - ptrs_w, - ptrs_s, - direct_offsets, - inverse_indices, - direct_expert_ids, - k, - n, - group_size, - k, - ) + for broadcast, output in ((False, legacy_direct), (True, direct)): + monkeypatch.setenv( + "VLLM_SM70_MXFP4_MOE_BROADCAST_INPUT_DECODE", + "1" if broadcast else "0", + ) + torch.ops._C.mxfp4_moe_single_token_prepare_w13_sm70_out( + output, + direct_input, + x, + expert_ids.view(1, num_experts), + ptrs_w, + ptrs_s, + direct_offsets, + inverse_indices, + direct_expert_ids, + k, + n, + group_size, + k, + ) torch.accelerator.synchronize() + torch.testing.assert_close(legacy_direct, reference, rtol=0.0, atol=0.0) torch.testing.assert_close(direct, reference, rtol=0.0, atol=0.0) diff --git a/tests/test_envs.py b/tests/test_envs.py index 78e5e274d6..eb679f1829 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -141,6 +141,28 @@ def test_sm70_concurrency_tuning_envs( monkeypatch.setenv(name, "0") assert environment_variables[name]() is False + name = "VLLM_SM70_DSV4_QNORM_KV_FUSED_TP4" + 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") diff --git a/vllm/envs.py b/vllm/envs.py index e937c52e76..7063f8ae82 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -184,6 +184,7 @@ VLLM_SM70_NVFP4_DENSE_TUNE_MAX_M: int = 16 VLLM_SM70_DSV4_FP16_GEMV: bool = False 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 VLLM_SM70_AWQ_MOE_TUNE_MAX_TOKENS: int = 128 VLLM_SM70_NVFP4_MOE_TUNE_MAX_TOKENS: int = 128 @@ -245,12 +246,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_TOP6_DECODE: bool = True + VLLM_SM70_MXFP4_MOE_DIRECT_ORDER_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 VLLM_SM70_DSV4_SPARSE_MLA_SPLITK_C128: bool = False @@ -1813,6 +1816,9 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_DSV4_MHC_FP32_STAGE": lambda: bool( int(os.getenv("VLLM_SM70_DSV4_MHC_FP32_STAGE", "1")) ), + "VLLM_SM70_DSV4_QNORM_KV_FUSED_TP4": lambda: bool( + int(os.getenv("VLLM_SM70_DSV4_QNORM_KV_FUSED_TP4", "1")) + ), # Skip PP metadata and TP reconstruction only for the exact, replicated # SM70 B1 hidden-state schema validated by the worker on both stages. "VLLM_SM70_PP_STATIC_HIDDEN_TRANSFER": lambda: bool( @@ -2171,7 +2177,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. @@ -2197,7 +2203,17 @@ 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. 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", "1")) + ), + "VLLM_SM70_MXFP4_MOE_BROADCAST_INPUT_DECODE": lambda: bool( + int(os.getenv("VLLM_SM70_MXFP4_MOE_BROADCAST_INPUT_DECODE", "1")) ), # FP8 caller for the generic SM70 TurboMind active-source-group compact # decode path. The backend scheduler keeps source expert group semantics diff --git a/vllm/model_executor/layers/quantization/mxfp4_sm70_moe.py b/vllm/model_executor/layers/quantization/mxfp4_sm70_moe.py index 8cd9473ca1..5408a025ca 100644 --- a/vllm/model_executor/layers/quantization/mxfp4_sm70_moe.py +++ b/vllm/model_executor/layers/quantization/mxfp4_sm70_moe.py @@ -343,6 +343,11 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: ) if envs.VLLM_SM70_MXFP4_MOE_DIRECT_TOP6_DECODE: required_ops += ("mxfp4_moe_single_token_prepare_w13_sm70_out",) + if ( + envs.VLLM_SM70_MXFP4_MOE_DIRECT_TOP6_DECODE + and envs.VLLM_SM70_MXFP4_MOE_DIRECT_ORDER_DECODE + ): + required_ops += ("awq_moe_single_token_weighted_reduce_out",) missing_ops = [name for name in required_ops if not hasattr(torch.ops._C, name)] if missing_ops: raise RuntimeError( @@ -726,6 +731,49 @@ def apply( and layer.expert_map is None and layer.local_num_experts == layer.global_num_experts ) + direct_order = ( + direct_top6 + and envs.VLLM_SM70_MXFP4_MOE_DIRECT_ORDER_DECODE + and topk_ids.dtype == torch.int32 + and topk_ids.is_contiguous() + ) + if direct_order: + route_ids = topk_ids.view(-1) + route_offsets = buffers["compact_expert_offsets"] + sm70_ops.mxfp4_moe_dense_stage_sm70_out( + buffers["gate_up"], + x, + route_offsets, + route_ids, + layer.w13_strided_ptrs_w, + layer.w13_strided_ptrs_s, + _DEEPSEEK_V4_FLASH_TOP_K, + layer.sm70_mxfp4_w13_k_dim, + layer.sm70_mxfp4_w13_n_dim, + layer.sm70_mxfp4_group_size, + ) + self._apply_swiglu(layer, buffers["intermediate"], buffers["gate_up"]) + sm70_ops.mxfp4_moe_dense_stage_sm70_out( + buffers["sorted_output"], + buffers["intermediate"], + route_offsets, + route_ids, + layer.w2_strided_ptrs_w, + layer.w2_strided_ptrs_s, + _DEEPSEEK_V4_FLASH_TOP_K, + layer.sm70_mxfp4_w2_k_dim, + layer.sm70_mxfp4_w2_n_dim, + layer.sm70_mxfp4_group_size, + ) + 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, + ) + return output if direct_top6: sm70_ops.mxfp4_moe_single_token_prepare_w13_sm70_out( buffers["gate_up"], 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 ac88ab5d3d..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 @@ -5,9 +5,13 @@ import torch +import vllm.envs as envs from vllm.models.deepseek_v4.common.ops.cache_utils import ( quantize_and_insert_k_cache, ) +from vllm.models.deepseek_v4.common.ops.fp8_software import ( + fp32_to_fp8_e4m3fn_bits, +) from vllm.triton_utils import tl, triton _HEAD_DIM = 512 @@ -93,6 +97,111 @@ def _sm70_qnorm_rope_kernel( ) +@triton.jit +def _sm70_qnorm_rope_parallel_kv_insert_kernel( + q_ptr, + kv_ptr, + cache_ptr, + slot_mapping_ptr, + positions_ptr, + cos_sin_cache_ptr, + cache_stride0, + eps: tl.constexpr, + num_tokens, + num_heads: tl.constexpr, + cache_block_size: tl.constexpr, + HEAD_DIM: tl.constexpr, + ROPE_DIM: tl.constexpr, + NOPE_DIM: tl.constexpr, + HALF_ROPE: tl.constexpr, + TOKEN_DATA_SIZE: tl.constexpr, +): + """Fuse TP4 Q transform with eight exact packed-KV writer CTAs.""" + token_idx = tl.program_id(0) + worker_idx = tl.program_id(1) + if token_idx >= num_tokens: + return + + if worker_idx < num_heads: + q_position = tl.load(positions_ptr + token_idx).to(tl.int64) + q_pair_idx = tl.arange(0, HALF_ROPE) + q_cos = tl.load(cos_sin_cache_ptr + q_position * ROPE_DIM + q_pair_idx).to( + tl.float32 + ) + q_sin = tl.load( + cos_sin_cache_ptr + q_position * ROPE_DIM + HALF_ROPE + q_pair_idx + ).to(tl.float32) + q_base = q_ptr + token_idx * num_heads * HEAD_DIM + worker_idx * HEAD_DIM + q_offsets = tl.arange(0, HEAD_DIM) + q_values = tl.load(q_base + q_offsets).to(tl.float32) + q_rrms = tl.rsqrt(tl.sum(q_values * q_values, axis=0) / HEAD_DIM + eps) + q_values *= q_rrms + tl.store( + q_base + q_offsets, + q_values.to(q_ptr.type.element_ty), + mask=q_offsets < NOPE_DIM, + ) + + q_even_offsets = NOPE_DIM + q_pair_idx * 2 + q_odd_offsets = q_even_offsets + 1 + q_even = tl.load(q_base + q_even_offsets).to(tl.float32) * q_rrms + q_odd = tl.load(q_base + q_odd_offsets).to(tl.float32) * q_rrms + tl.store( + q_base + q_even_offsets, + (q_even * q_cos - q_odd * q_sin).to(q_ptr.type.element_ty), + ) + tl.store( + q_base + q_odd_offsets, + (q_even * q_sin + q_odd * q_cos).to(q_ptr.type.element_ty), + ) + else: + part_idx = worker_idx - num_heads + slot_idx = tl.load(slot_mapping_ptr + token_idx) + if slot_idx < 0: + return + + block_idx = slot_idx // cache_block_size + pos_in_block = slot_idx % cache_block_size + cache_block = cache_ptr + block_idx.to(tl.int64) * cache_stride0 + token_data = cache_block + pos_in_block * TOKEN_DATA_SIZE + token_scales = ( + cache_block + cache_block_size * TOKEN_DATA_SIZE + pos_in_block * 8 + ) + kv_row = kv_ptr + token_idx * HEAD_DIM + kv_offsets = tl.arange(0, 64) + if part_idx < 7: + input_offsets = part_idx * 64 + kv_offsets + kv_values = tl.load(kv_row + input_offsets) + block_max = tl.maximum(tl.max(tl.abs(kv_values), axis=0), 1.0e-4) + exponent = tl.ceil(tl.log2(block_max / 448.0)) + scale = tl.exp2(exponent) + scaled = tl.clamp(kv_values / scale, -448.0, 448.0) + packed = fp32_to_fp8_e4m3fn_bits(scaled.to(tl.float32)) + tl.store(token_data + input_offsets, packed) + encoded = tl.maximum(tl.minimum(exponent + 127.0, 255.0), 0.0) + tl.store(token_scales + part_idx, encoded.to(tl.uint8)) + else: + tl.store(token_scales + 7, tl.zeros((), dtype=tl.uint8)) + kv_position = tl.load(positions_ptr + token_idx).to(tl.int64) + kv_pair_idx = tl.arange(0, HALF_ROPE) + kv_cos = tl.load( + cos_sin_cache_ptr + kv_position * ROPE_DIM + kv_pair_idx + ).to(tl.float32) + kv_sin = tl.load( + cos_sin_cache_ptr + kv_position * ROPE_DIM + HALF_ROPE + kv_pair_idx + ).to(tl.float32) + kv_even = tl.load(kv_row + NOPE_DIM + kv_pair_idx * 2).to(tl.float32) + kv_odd = tl.load(kv_row + NOPE_DIM + kv_pair_idx * 2 + 1).to(tl.float32) + rotated_even = (kv_even * kv_cos - kv_odd * kv_sin).to(tl.float16) + rotated_odd = (kv_even * kv_sin + kv_odd * kv_cos).to(tl.float16) + bf16_out = (token_data + NOPE_DIM).to(tl.pointer_type(tl.bfloat16)) + tl.store(bf16_out + kv_pair_idx * 2, rotated_even.to(tl.bfloat16)) + tl.store( + bf16_out + kv_pair_idx * 2 + 1, + rotated_odd.to(tl.bfloat16), + ) + + def sm70_qnorm_rope_kv_fp8_insert( q: torch.Tensor, kv: torch.Tensor, @@ -112,6 +221,29 @@ def sm70_qnorm_rope_kv_fp8_insert( assert q.is_contiguous() and kv.is_contiguous() num_tokens, num_heads, _ = q.shape + if envs.VLLM_SM70_DSV4_QNORM_KV_FUSED_TP4 and num_tokens == 1 and num_heads == 16: + cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1) + _sm70_qnorm_rope_parallel_kv_insert_kernel[(num_tokens, num_heads + 8)]( + q, + kv, + cache_2d, + slot_mapping, + positions, + cos_sin_cache, + cache_2d.stride(0), + eps, + num_tokens, + num_heads=num_heads, + cache_block_size=block_size, + HEAD_DIM=_HEAD_DIM, + ROPE_DIM=_ROPE_DIM, + NOPE_DIM=_NOPE_DIM, + HALF_ROPE=_HALF_ROPE, + TOKEN_DATA_SIZE=576, + num_warps=4, + ) + return q + kv_roped = torch.empty_like(kv) _sm70_qnorm_rope_kernel[(num_tokens, num_heads + 1)]( q, diff --git a/vllm/models/deepseek_v4/sm70/sparse.py b/vllm/models/deepseek_v4/sm70/sparse.py index f507372031..1eec99a606 100644 --- a/vllm/models/deepseek_v4/sm70/sparse.py +++ b/vllm/models/deepseek_v4/sm70/sparse.py @@ -40,6 +40,11 @@ logger = init_logger(__name__) +def _qk_dsplit_block_h(num_heads: int) -> int: + """Use one CTA head group for the exact TP4 decode shape.""" + return 16 if num_heads == 16 else 8 + + class DeepseekV4SM70SparseBackend(DeepseekV4FlashMLASparseBackend): supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16] @@ -258,6 +263,7 @@ def _forward_decode( sm70_sparse_attention_paged_fp8_splitk_qk_dsplit( partial_qk=partial_qk, partial_probs=partial_probs, + stage1_block_h=_qk_dsplit_block_h(q.shape[1]), **common_kwargs, ) else: From d9d0da21decfb725f5f9f2c0a7b8175e42d5a0d3 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 26 Aug 2026 11:39:04 +0800 Subject: [PATCH 009/105] [Kernel][SM70] Default safe FP8 prescaled M1 decode (#316) * [Kernel][SM70] Prescale DeepSeek V4 FP8 decode Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> (cherry picked from commit b8d928d803ab841d7e341c6748cb71dc38d3b895) * [Kernel][SM70] Default safe FP8 prescaled M1 decode Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --------- Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- csrc/ops.h | 6 + .../src/turbomind/kernels/gemm/gemm.cu | 44 ++++++- .../kernels/gemm/kernel/sm70_884_8.cu | 24 ++++ csrc/sm70_turbomind/ops/awq_sm70_gemm.cu | 26 +++- csrc/torch_bindings.cpp | 9 ++ docs/design/sm70_v100_migration_control.md | 27 +++++ .../test_sm70_fp8_prefill_exact_dense.py | 97 +++++++++++++++ tests/test_envs.py | 6 + vllm/_sm70_ops.py | 29 +++++ vllm/envs.py | 4 + .../model_executor/layers/quantization/fp8.py | 113 +++++++++++++++++- 11 files changed, 374 insertions(+), 11 deletions(-) diff --git a/csrc/ops.h b/csrc/ops.h index af458e78dc..e2fd08b575 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -220,6 +220,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); 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 0aa6899e49..d555a8be7d 100644 --- a/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu +++ b/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu @@ -3510,10 +3510,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 +5577,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, diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index 3cdfd346fc..38278ee057 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -323,6 +323,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[]"); diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index aaf87ea84a..b7fe95ef9f 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -43678,3 +43678,30 @@ Interpretation: 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. 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/test_envs.py b/tests/test_envs.py index eb679f1829..ffc79b49e4 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -147,6 +147,12 @@ 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 diff --git a/vllm/_sm70_ops.py b/vllm/_sm70_ops.py index a554e2bf4d..5870e325e4 100644 --- a/vllm/_sm70_ops.py +++ b/vllm/_sm70_ops.py @@ -490,6 +490,35 @@ 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, diff --git a/vllm/envs.py b/vllm/envs.py index 7063f8ae82..200f73fa8e 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -159,6 +159,7 @@ 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 @@ -1766,6 +1767,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")) ), diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 36a7a8097b..8f8c975d85 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -168,6 +168,47 @@ _sm70_fp8_qpn8_pp2_tp4_workspaces: dict[tuple[int, torch.dtype], torch.Tensor] = {} +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 @@ -1079,6 +1120,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, @@ -1106,6 +1170,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 @@ -1391,6 +1485,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", @@ -1404,7 +1501,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( From 232e1e4b58ba1e1cf388c969eba5892e15adb9f6 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 26 Aug 2026 12:25:13 +0800 Subject: [PATCH 010/105] [CI/Build] Repair CUDA test dependency updates (#317) Upgrade SentenceTransformers, Pillow, and Black with a regenerated CUDA test lock; add the new pytokens dependency and repair the CrossEncoder model_kwargs API. Keep chardet 5.2 because mbstrdecoder currently requires chardet below 6. Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- requirements/test/cuda.txt | 10 +++++++--- tests/conftest.py | 2 +- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index 7470b91afd..94865c497d 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -62,7 +62,7 @@ backoff==2.2.1 # schemathesis bitsandbytes==0.49.2 # 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 @@ -488,6 +488,7 @@ numpy==2.2.6 # scikit-learn # scipy # segmentation-models-pytorch + # sentence-transformers # soxr # statsmodels # tensorizer @@ -616,7 +617,7 @@ 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 @@ -762,6 +763,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 @@ -890,7 +893,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 @@ -986,6 +989,7 @@ tokenizers==0.22.2 # via # -c requirements/common.txt # -r requirements/test/cuda.in + # sentence-transformers # transformers tomli==2.2.1 # via schemathesis 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: From fe1c13886d7b6d78f94d089a0fb641e96790d6cf Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 26 Aug 2026 12:32:28 +0800 Subject: [PATCH 011/105] [Bench][SM70] Record prefix architecture environment (#318) Track PREFIX_ controls in retained SM70 benchmark environments so memory, scheduling, latency, and numerical configurations remain reproducible. Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- benchmarks/benchmark_sm70_decode.py | 1 + 1 file changed, 1 insertion(+) diff --git a/benchmarks/benchmark_sm70_decode.py b/benchmarks/benchmark_sm70_decode.py index 19f27d6dea..bf6d364c19 100644 --- a/benchmarks/benchmark_sm70_decode.py +++ b/benchmarks/benchmark_sm70_decode.py @@ -169,6 +169,7 @@ def _tracked_env() -> dict[str, str]: "VLLM_FLASH_", "FLASH_ATTN", "FLA_", + "PREFIX_", "CUDA_MODULE_LOADING", "TORCH_CUDA_ARCH_LIST", "TORCHINDUCTOR_", From 5245946aa30634dc05002f4864f0ca2de023422d Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 26 Aug 2026 12:45:09 +0800 Subject: [PATCH 012/105] [Bugfix][SM70] Emit sampler stable ABI module (#320) Make the SM70 sampler target use the stable Python ABI so single-arch Volta source builds package the expected .abi3.so module. Reported-by: Andrew Leech Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- CMakeLists.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/CMakeLists.txt b/CMakeLists.txt index 21dc88d62a..d07300fdff 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -729,6 +729,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 From 2c0aee15dbad37f3f45b8e40dcbcf28e90a206ed Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 26 Aug 2026 12:59:19 +0800 Subject: [PATCH 013/105] [CI/Build] Use latest compatible charset stack (#321) Refresh the CUDA test lock to mbstrdecoder 1.1.5 and the highest chardet release permitted by upstream metadata, 6.0.0.post1. Reported-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- requirements/test/cuda.txt | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index 94865c497d..edcd2df815 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -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 @@ -400,7 +400,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 From 6a424f2aca9d89db1fb132a6870ef6e2f9ce0383 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 26 Aug 2026 13:38:24 +0800 Subject: [PATCH 014/105] [CI/Build] Refresh compatible dependency stack (#323) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Repair Dependabot’s broad update on current main with consistent source inputs and generated locks, preserving the validated Torch/tokenizer/Pillow stack and synchronizing optional dependency consumers. Reported-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .buildkite/test-amd.yaml | 2 +- .buildkite/test_areas/kernels.yaml | 2 +- requirements/common.txt | 26 ++++----- requirements/cpu.txt | 2 +- requirements/cuda.txt | 8 +-- requirements/docs.txt | 6 +-- requirements/rocm.txt | 12 ++--- requirements/test/cuda.in | 22 ++++---- requirements/test/cuda.txt | 75 +++++++++++++------------- requirements/test/rocm.in | 22 ++++---- requirements/test/rocm.txt | 87 +++++++++++++++--------------- requirements/test/xpu.txt | 25 ++++----- requirements/tpu.txt | 2 +- requirements/xpu.txt | 2 +- setup.py | 8 +-- vllm/config/speculative.py | 2 +- 16 files changed, 151 insertions(+), 152 deletions(-) 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/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 abf2e75c49..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 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 edcd2df815..e6a621d712 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,7 +60,7 @@ 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==26.5.1 # via datamodel-code-generator @@ -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 @@ -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 @@ -554,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 @@ -611,7 +609,7 @@ 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 @@ -649,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 @@ -658,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 @@ -695,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 @@ -708,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 @@ -806,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 @@ -847,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 @@ -966,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 @@ -974,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 @@ -1043,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 @@ -1070,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 @@ -1101,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 @@ -1140,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/vllm/config/speculative.py b/vllm/config/speculative.py index b89d8b0210..c74aeb882e 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -999,7 +999,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 From 8700861efac86da8f25d44c9dccbcebd18fdd749 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 26 Aug 2026 13:58:21 +0800 Subject: [PATCH 015/105] [Kernel][SM70] Default guarded packed FP13 GEMV (#325) Repair #322 on current main, retaining only the measured FP13 kernel and adding tensor-value, CUDA-device, shape, dtype, layout, memory, quality, and rollback safeguards. The production path defaults on without model/checkpoint identity admission. Original-commit: 165ec09fa9f9b16f416ee586534d342d6c3f479a Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_v100_migration_control.md | 37 +++ .../test_sm70_deepseek_v4_fp16_gemv.py | 122 +++++++- vllm/envs.py | 4 + .../layers/fused_moe/router/gate_linear.py | 12 +- vllm/model_executor/layers/linear.py | 9 + vllm/models/deepseek_v4/attention.py | 16 + vllm/models/deepseek_v4/compressor.py | 1 + vllm/models/deepseek_v4/sm70/gemv.py | 282 +++++++++++++++++- 8 files changed, 477 insertions(+), 6 deletions(-) diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index b7fe95ef9f..5a80c3c707 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -43705,3 +43705,40 @@ Interpretation: 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. +- Production defaults on through `VLLM_SM70_DSV4_FP13_GEMV=1`; explicit `=0` + is the rollback. Admission checks SM70, CUDA device equality, dtype, shape, + layout, output contract, and the actual weight bits. It never reads a model + name or checkpoint identity, and an incompatible tensor retains FP16. Packing + occurs once after loading; capture and per-token execution perform no packing. + The retained buffers add 549.66 MiB without pipeline parallelism, or at most + 282.34 MiB per rank for the audited PP2 partition. 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/vllm/envs.py b/vllm/envs.py index 200f73fa8e..2d3e04aa45 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -184,6 +184,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 @@ -1817,6 +1818,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")) ), 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..68f99ee3eb 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -369,6 +369,15 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: if not current_platform.is_cuda_alike(): 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 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 From 9bbc554e1d9f63a628d1bc66ecb45ee4632267a2 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 26 Aug 2026 13:59:23 +0800 Subject: [PATCH 016/105] [Doc][SM70] Plan exact-8K FP8 prefill 5500 work (#324) Record the source-matching Qwen3.8 exact-8K baseline, request-wall target, trace budget, bounded shard-consumer design, rejected routes, quality ladder, and rollback requirements so implementation work does not repeat known dead ends. Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_qwen38_fp8_prefill_5500.md | 142 ++++++++++++++++++++ 1 file changed, 142 insertions(+) create mode 100644 docs/design/sm70_qwen38_fp8_prefill_5500.md 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`. From b06f54cee3ab0c2390d9d1c4313221ae429a2790 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 26 Aug 2026 14:07:50 +0800 Subject: [PATCH 017/105] [Bugfix][SM70] Isolate MXFP4 direct-order offsets (#326) Prevent M2-M8 active-expert compaction from poisoning B1 direct-order graph capture by selecting the immutable one-row slot offsets. CUDA kernels, arithmetic, route order, defaults, and rollbacks remain unchanged. Reported-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Original-commit: 052312c66f89bb144cf68f5776d83552c7b9e3d1 Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_v100_migration_control.md | 8 ++++++++ tests/quantization/test_sm70_mxfp4_moe.py | 16 ++++++++++++++++ .../layers/quantization/mxfp4_sm70_moe.py | 14 +++++++++++++- 3 files changed, 37 insertions(+), 1 deletion(-) diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index 5a80c3c707..7bb276c394 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -43456,6 +43456,14 @@ Interpretation: 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 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, diff --git a/tests/quantization/test_sm70_mxfp4_moe.py b/tests/quantization/test_sm70_mxfp4_moe.py index 845e1a3809..9e3e9fa18f 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, + _select_mxfp4_direct_order_offsets, _select_mxfp4_stage_dispatch, validate_mxfp4_sm70_moe_contract, validate_mxfp4_sm70_moe_weight_layout, @@ -210,6 +211,21 @@ def test_mxfp4_sm70_b1_dispatch_selects_six_runtime_experts(monkeypatch): assert count == 6 +def test_mxfp4_sm70_direct_order_ignores_compacted_offsets(): + buffers = { + # 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 + ), + "slot_expert_offsets": torch.arange(7, dtype=torch.int32), + } + + offsets = _select_mxfp4_direct_order_offsets(buffers) + + assert offsets is buffers["slot_expert_offsets"] + torch.testing.assert_close(offsets, torch.arange(7, dtype=torch.int32)) + + def test_mxfp4_sm70_b1_dispatch_rejects_incompatible_permute_fastpath( monkeypatch, ): diff --git a/vllm/model_executor/layers/quantization/mxfp4_sm70_moe.py b/vllm/model_executor/layers/quantization/mxfp4_sm70_moe.py index 5408a025ca..dc39f8426a 100644 --- a/vllm/model_executor/layers/quantization/mxfp4_sm70_moe.py +++ b/vllm/model_executor/layers/quantization/mxfp4_sm70_moe.py @@ -185,6 +185,18 @@ def _select_mxfp4_stage_dispatch( return buffers["expert_offsets"], buffers["dense_expert_ids"], num_experts +def _select_mxfp4_direct_order_offsets( + buffers: dict[str, torch.Tensor], +) -> torch.Tensor: + """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 one row per route, + so it must use the separately maintained slot offsets instead. + """ + return buffers["slot_expert_offsets"] + + def validate_mxfp4_sm70_moe_contract( *, global_num_experts: int, @@ -739,7 +751,7 @@ def apply( ) if direct_order: route_ids = topk_ids.view(-1) - route_offsets = buffers["compact_expert_offsets"] + route_offsets = _select_mxfp4_direct_order_offsets(buffers) sm70_ops.mxfp4_moe_dense_stage_sm70_out( buffers["gate_up"], x, From 89ac768c419494ad172f11099f297e0c38b9593d Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 26 Aug 2026 14:59:16 +0800 Subject: [PATCH 018/105] [Kernel] Skip masked DFlash2 draft history Co-authored-by: OpenAI Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../benchmark_flash_v100_dflash2_swa.py | 149 ++++++++++++++++++ .../kernel/fused_mha_forward_paged.cu | 6 + .../test_sm70_flash_v100_dflash2_swa.py | 78 +++++++++ 3 files changed, 233 insertions(+) create mode 100644 benchmarks/kernels/benchmark_flash_v100_dflash2_swa.py create mode 100644 tests/kernels/attention/test_sm70_flash_v100_dflash2_swa.py 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..e5a9974028 --- /dev/null +++ b/benchmarks/kernels/benchmark_flash_v100_dflash2_swa.py @@ -0,0 +1,149 @@ +# 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() -> 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/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/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..2b8f751dc3 --- /dev/null +++ b/tests/kernels/attention/test_sm70_flash_v100_dflash2_swa.py @@ -0,0 +1,78 @@ +# 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 + ) From f4c3451a2b2205dead8b0a62c968f7303f1cc702 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 26 Aug 2026 15:18:10 +0800 Subject: [PATCH 019/105] [Release] Restore public source snapshot to v1.3.0 (#329) Freeze the public repository at the exact released v1.3.0 source tree before archiving public development. Ongoing work was mirrored to the personal private repository first; history is preserved without force-pushing. Target-tag: v1.3.0 Target-tree: 4373fb80cf67164be8354c8ac098a63ec56200c3 Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .buildkite/test-amd.yaml | 2 +- .buildkite/test_areas/kernels.yaml | 2 +- .pre-commit-config.yaml | 12 +- CMakeLists.txt | 42 - README.md | 65 +- .../analyze_dflash2_selector_alignment.py | 517 -- benchmarks/analyze_dsv4_dspark_confidence.py | 288 - benchmarks/benchmark_dspark_api.py | 1 - benchmarks/benchmark_dspark_stream.py | 1 - .../benchmark_dsv4_dspark_long_context.py | 576 -- benchmarks/benchmark_dsv4_gsm8k_api.py | 231 - benchmarks/benchmark_dsv4_quality_api.py | 635 -- .../benchmark_sm70_cublas_gemm_algorithms.py | 251 - .../benchmark_sm70_custom_all_reduce.py | 70 +- benchmarks/benchmark_sm70_d256_exact_dense.py | 215 - benchmarks/benchmark_sm70_decode.py | 17 +- .../benchmark_sm70_dflash2_gdn_metadata.py | 373 - benchmarks/benchmark_sm70_dflash2_gsm8k.py | 484 - ...benchmark_sm70_dflash2_sparse_rejection.py | 243 - .../benchmark_sm70_flashinfer_fixed_entry.py | 20 +- ...hmark_sm70_flashqla_indexed_state_cache.py | 48 +- .../benchmark_sm70_fp8_fixed_vs_dynamic.py | 17 +- ...hmark_sm70_fp8_prefill_exact_dense_blas.py | 368 - .../benchmark_sm70_fp8_prefill_tactics.py | 486 - benchmarks/benchmark_sm70_fp8_qpn8.py | 729 -- benchmarks/benchmark_sm70_gdn_exactness.py | 396 +- benchmarks/benchmark_sm70_hmma_c2a_mincomm.py | 3 +- .../benchmark_sm70_instruction_latency.py | 3 +- benchmarks/benchmark_sm70_lm_head_top1.py | 18 +- .../benchmark_sm70_long_quality_suffix.py | 6 +- benchmarks/benchmark_sm70_model_tokens.py | 113 +- ...benchmark_sm70_mxfp4_moe_active_experts.py | 232 +- ..._sm70_native_bm32_allp_crossblock_micro.py | 60 +- ...ark_sm70_native_bm32_allp_scratch_micro.py | 73 +- ..._sm70_native_bm32_full_phase_allp_micro.py | 64 +- ...hmark_sm70_native_bm32_full_phase_micro.py | 64 +- ...nchmark_sm70_native_bm32_pv_panel_micro.py | 17 +- .../benchmark_sm70_prefill_gemm_collective.py | 395 - benchmarks/benchmark_sm70_raw_hmma_probe.py | 5 +- .../benchmark_sm70_rejection_sampler_micro.py | 34 +- .../benchmark_sm70_sglang_push_chain.py | 162 - benchmarks/benchmark_sm70_source_inventory.py | 88 +- ...chmark_sm70_tp2_compact_verifier_logits.py | 10 +- ...enchmark_sm70_tp4_fused_norm_collective.py | 305 - benchmarks/benchmark_sm70_turbomind_parity.py | 33 +- .../benchmark_sm70_wmma_fragment_probe.py | 38 +- benchmarks/compare_dsv4_quality_results.py | 425 - benchmarks/compare_sm70_tensor_dumps.py | 156 +- .../kernels/benchmark_device_communicators.py | 74 +- benchmarks/kernels/benchmark_dspark_markov.py | 1 - .../kernels/benchmark_dspark_metadata.py | 1 - .../benchmark_dspark_sm70_attention.py | 1 - .../kernels/benchmark_flash_v100_fp8_kv.py | 1 - .../benchmark_sm70_fp8_graph_capture.py | 1 - .../kernels/benchmark_sm70_nvfp4_qpn2.py | 499 -- .../profile_sm70_fp8_qpn8_gated_pair_ncu.py | 140 - benchmarks/profile_sm70_fp8_qpn8_ncu.py | 194 - benchmarks/verify_sm70_artifact.py | 168 +- cmake/external_projects/vllm_flash_attn.cmake | 6 - .../sm70_flash_attn_d256_gqa_arch.patch | 1673 ---- .../sm70_flash_attn_d256_k_pingpong.patch | 90 - csrc/custom_all_reduce.cu | 125 +- csrc/custom_all_reduce.cuh | 775 +- csrc/libtorch_stable/layernorm_kernels.cu | 27 +- .../marlin_moe_wna16/sm70_marlin_fp8_gemm.cu | 197 +- .../moe/marlin_moe_wna16/sm70_marlin_gemm.cuh | 1519 ++-- .../sm70_marlin_moe_dispatch.cu | 98 +- .../sm70_marlin_mxfp4_gemm.cu | 128 +- .../sm70_marlin_nvfp4_gemm.cu | 128 +- .../marlin_moe_wna16/sm70_marlin_u4_gemm.cu | 189 +- .../marlin_moe_wna16/sm70_marlin_u4b8_gemm.cu | 167 +- .../marlin_moe_wna16/sm70_marlin_u8_gemm.cu | 211 +- .../sm70_marlin_u8b128_gemm.cu | 191 +- csrc/moe/moe_permute_unpermute_op.cu | 17 +- .../moe_permute_unpermute_kernel.cu | 27 +- .../moe_permute_unpermute_kernel.h | 11 +- csrc/ops.h | 675 +- .../marlin/sm70_awq_marlin_repack.cu | 32 +- .../marlin/sm70_gptq_marlin_repack.cu | 34 +- .../marlin/sm70_marlin_common.cuh | 600 +- .../marlin/sm70_marlin_dispatch.cu | 129 +- .../marlin/sm70_marlin_fp8_gemm.cu | 220 +- csrc/quantization/marlin/sm70_marlin_gemm.cuh | 172 +- .../marlin/sm70_marlin_iterator_utils.cuh | 34 +- csrc/quantization/marlin/sm70_marlin_mma.cuh | 45 +- .../marlin/sm70_marlin_mxfp4_gemm.cu | 212 +- .../marlin/sm70_marlin_nvfp4_gemm.cu | 210 +- .../marlin/sm70_marlin_splitk.cuh | 50 +- .../marlin/sm70_marlin_u4_gemm.cu | 280 +- .../marlin/sm70_marlin_u4b8_gemm.cu | 218 +- .../marlin/sm70_marlin_u8_gemm.cu | 251 +- .../marlin/sm70_marlin_u8b128_gemm.cu | 222 +- csrc/sm70_tile_runtime_signal.cuh | 19 +- .../kernels/attention/quantization.h | 1373 +-- .../kernels/gemm/arch/config_sm70_s884.h | 219 +- .../src/turbomind/kernels/gemm/gemm.cu | 1344 ++- .../turbomind/kernels/gemm/iterator_sm70.h | 406 +- .../kernels/gemm/kernel/sm70_884_4.cu | 74 +- .../kernels/gemm/kernel/sm70_884_8.cu | 57 +- .../src/turbomind/kernels/gemm/transform.h | 286 +- .../src/turbomind/kernels/gemm/types.h | 355 +- csrc/sm70_turbomind/ops/LICENSE.v100-skinny | 21 - csrc/sm70_turbomind/ops/awq_sm70_gemm.cu | 1156 +-- csrc/sm70_turbomind/ops/fp8_qpn8_sm70.cu | 1255 --- csrc/sm70_turbomind/ops/nvfp4_qpn2_sm70.cu | 546 -- csrc/sm70_turbomind/ops/nvfp4_qpn4_sm70.cu | 1029 --- .../ops/qwen38_prefill_cutlass.cu | 130 - .../ops/qwen38_prefill_cutlass.cuh | 13 - .../ops/sm70_top20_philox_fragment.cpp | 52 - .../ops/sm70_top20_philox_sampler.cu | 289 - csrc/sm70_turbomind/ops/tm_registry_sm70.cu | 75 +- csrc/torch_bindings.cpp | 198 +- docs/assets/wechat-group-qr-7.png | Bin 529331 -> 0 bytes ...6_27b_awq_mtp_target_verifier_fastpaths.md | 36 +- .../sm70_deepseek_v4_dspark_n7_latest_main.md | 28 - ...ek_v4_dspark_verifier_20ms_long_context.md | 217 - docs/design/sm70_deepseek_v4_flash_bringup.md | 34 - docs/design/sm70_dflash2_acceptance_rd.md | 77 - docs/design/sm70_dflash2_long_verify_decay.md | 465 - docs/design/sm70_dflash2_ngram_hybrid.md | 130 - docs/design/sm70_dflash2_nvfp4_17ms.md | 220 - docs/design/sm70_dflash2_target_graph_20ms.md | 1229 --- ...0_nvfp4_turbomind_operator_optimization.md | 10 +- docs/design/sm70_prefix_anchored_swa.md | 57 - docs/design/sm70_qwen38_130_acceptance.md | 26 +- docs/design/sm70_qwen38_fp8_prefill_5500.md | 142 - docs/design/sm70_qwen38_fp8_prefill_decay.md | 307 - docs/design/sm70_qwen38_nvfp4_decode.md | 711 -- docs/design/sm70_qwen38_qpn8_decode.md | 245 - docs/design/sm70_tile_runtime_exploration.md | 20 +- docs/design/sm70_v100_migration_control.md | 2079 ----- docs/features/speculative_decoding/README.md | 27 - docs/models/supported_models.md | 8 - .../flash_attn_v100/__init__.py | 2 - .../flash_attn_v100/flash_attn_interface.py | 153 +- flash-attention-v100/include/fused_mha.h | 309 +- flash-attention-v100/include/fused_mma.h | 626 +- .../kernel/contiguous_to_paged.cu | 123 +- .../kernel/flash_decode_paged.cu | 1860 +--- .../kernel/flash_decode_turboquant.cu | 365 +- .../kernel/flash_v100_traits.cuh | 227 +- flash-attention-v100/kernel/fp8_kv_utils.cuh | 38 +- flash-attention-v100/kernel/fused_mha_api.cpp | 66 +- .../kernel/fused_mha_backward.cu | 2407 +++-- .../kernel/fused_mha_forward.cu | 2030 ++--- .../kernel/fused_mha_forward_paged.cu | 7866 +++++++++-------- .../kernel/fused_mha_forward_paged_full.cu | 1082 ++- .../kernel/paged_kv_utils.cuh | 98 +- .../kernel/paged_to_contiguous.cu | 355 +- .../kernel/paged_to_contiguous_fixed.cu | 167 +- .../kernel/paged_to_contiguous_old.cu | 140 +- requirements/common.txt | 26 +- requirements/cpu.txt | 2 +- requirements/cuda.txt | 8 +- requirements/docs.txt | 8 +- requirements/rocm.txt | 12 +- requirements/test/cuda.in | 22 +- requirements/test/cuda.txt | 93 +- requirements/test/rocm.in | 22 +- requirements/test/rocm.txt | 87 +- requirements/test/xpu.txt | 25 +- requirements/tpu.txt | 2 +- requirements/xpu.txt | 2 +- setup.py | 14 +- tests/benchmarks/test_dsv4_quality_tools.py | 141 - tests/config/test_prefix_anchored_swa.py | 158 - tests/conftest.py | 2 +- .../test_custom_all_reduce_allocator.py | 224 - .../test_static_tensor_transfer.py | 63 - tests/engine/test_arg_utils.py | 84 +- .../test_sm70_flash_v100_anchored_swa.py | 288 - .../test_sm70_flash_v100_decode_planner.py | 64 - .../test_sm70_flash_v100_grouped_verify.py | 245 - .../test_sm70_flash_v100_smallq_metadata.py | 232 - .../test_sm70_flash_v100_splitkv3.py | 12 +- .../test_sm70_flash_v100_varlen_layout.py | 330 - .../core/test_batched_weight_rms_norm.py | 63 - .../core/test_sm70_dflash2_gemma_rms.py | 121 - tests/kernels/mamba/test_mrv2_mamba_align.py | 196 - .../kernels/moe/test_moe_permute_unpermute.py | 13 +- .../moe/test_sm70_unquantized_moe_config.py | 93 - tests/kernels/test_compressor_kv_cache.py | 225 - .../kernels/test_deepseek_v4_sm70_indexer.py | 323 +- ...t_deepseek_v4_sm70_qnorm_rope_kv_insert.py | 47 +- .../test_fused_recurrent_packed_decode.py | 12 +- .../test_fused_sigmoid_gating_delta_rule.py | 321 - tests/kernels/test_mhc_sm70_fp16.py | 30 +- .../test_sm70_deepseek_v4_fp16_gemv.py | 122 +- tests/kernels/test_sm70_qwen35_gdn_split.py | 127 - .../test_deepseek_v4_swa_rope.py | 137 - .../test_qwen3_5_quantization.py | 136 - .../test_sm70_fp8_qpn8_pp2_tp4.py | 229 - .../model_executor/test_uva_tower_offload.py | 89 - .../multimodal/processing/test_gemma4.py | 46 - tests/models/test_deepseek_v4_pipeline.py | 50 - tests/models/test_deepseek_v4_sm70_routes.py | 54 - tests/quantization/test_nvfp4_mxfp4_gating.py | 91 - .../test_sm70_compressed_tensors_fastpaths.py | 293 - .../test_sm70_fp8_prefill_exact_dense.py | 733 +- .../test_sm70_modelopt_mixed_nvfp4.py | 192 - .../quantization/test_sm70_modelopt_nvfp4.py | 255 - .../test_sm70_modelopt_nvfp4_kernel_oracle.py | 148 - tests/quantization/test_sm70_mxfp4_moe.py | 197 - tests/quantization/test_sm70_nvfp4_qpn2.py | 143 - .../test_sm70_turbomind_adapter.py | 221 - tests/quantization/test_sm70_warmup.py | 356 - tests/test_envs.py | 68 - .../test_qwen3coder_tool_parser.py | 2 +- .../test_deepseek_ocr_resolution_modes.py | 163 - .../v1/attention/test_gdn_metadata_builder.py | 620 +- .../test_indexer_deepseek_v4_slot_mapping.py | 3 - .../test_prefix_anchored_swa_mask.py | 181 - .../attention/test_sm70_flash_v100_policy.py | 567 +- tests/v1/core/test_async_scheduler.py | 3 - tests/v1/core/test_kv_cache_utils.py | 88 - tests/v1/core/test_prefix_caching.py | 98 +- tests/v1/core/test_scheduler.py | 41 - .../core/test_single_type_kv_cache_manager.py | 107 +- .../v1/core/test_swa_inflight_window_free.py | 351 - tests/v1/core/utils.py | 21 +- tests/v1/cudagraph/test_cudagraph_dispatch.py | 461 +- .../test_sm70_mtp_split_cudagraphs.py | 96 - .../offloading_connector/test_scheduler.py | 201 +- tests/v1/kv_connector/unit/test_config.py | 10 - tests/v1/sample/test_sampler.py | 25 - tests/v1/sample/test_topk_topp_sampler.py | 63 - .../spec_decode/test_ddtree_flex_attention.py | 10 +- .../test_ddtree_gpu_runner_positions.py | 12 +- tests/v1/spec_decode/test_ddtree_metadata.py | 3 +- tests/v1/spec_decode/test_dflash2.py | 1310 --- .../spec_decode/test_dflash2_ngram_assist.py | 386 - .../test_dflash2_selector_analysis.py | 102 - .../v1/spec_decode/test_dflash_mrv2_config.py | 199 - tests/v1/spec_decode/test_dspark.py | 167 +- tests/v1/spec_decode/test_eagle.py | 11 +- tests/v1/spec_decode/test_mtp.py | 112 - .../test_mtp_draft_prob_alignment.py | 29 +- .../test_rejection_sampler_utils.py | 125 - tests/v1/worker/test_gpu_model_runner.py | 272 - .../test_gpu_model_runner_v2_kv_zero.py | 87 - tests/v1/worker/test_gpu_warmup_blocks.py | 98 - tests/v1/worker/test_gpu_worker_static_pp.py | 153 - tests/v1/worker/test_kv_block_zeroer.py | 161 - tests/v1/worker/test_mamba_utils.py | 34 +- tools/pre_commit/check_forbidden_imports.py | 7 +- tools/pre_commit/check_torch_cuda.py | 10 - vllm/_custom_ops.py | 34 - vllm/_sm70_ops.py | 972 -- vllm/compilation/cuda_graph.py | 25 +- .../passes/fusion/allreduce_rms_fusion.py | 228 +- .../passes/fusion/matcher_utils.py | 3 +- .../passes/fusion/qk_norm_rope_fusion.py | 3 +- .../passes/fusion/rms_quant_fusion.py | 9 +- .../passes/fusion/sequence_parallelism.py | 75 - .../passes/utility/fix_functionalization.py | 4 +- vllm/config/attention.py | 8 +- vllm/config/offload.py | 3 - vllm/config/speculative.py | 163 +- vllm/config/vllm.py | 282 +- .../device_communicators/cuda_communicator.py | 58 +- .../device_communicators/custom_all_reduce.py | 174 +- .../kv_transfer/kv_connector/v1/__init__.py | 4 - .../kv_transfer/kv_connector/v1/base.py | 15 - .../kv_connector/v1/offloading/scheduler.py | 147 +- .../kv_connector/v1/offloading_connector.py | 11 +- vllm/distributed/parallel_state.py | 119 +- vllm/engine/arg_utils.py | 44 +- vllm/envs.py | 394 +- vllm/forward_context.py | 9 +- vllm/model_executor/kernels/mhc/tilelang.py | 8 +- .../kernels/mhc/tilelang_kernels.py | 4 +- .../layers/attention/attention.py | 77 - .../layers/fla/ops/chunk_scaled_dot_kkt.py | 4 +- .../fla/ops/fused_gdn_prefill_post_conv.py | 3 +- .../layers/fla/ops/fused_recurrent.py | 4 +- .../layers/fla/ops/fused_sigmoid_gating.py | 156 +- .../layers/fused_moe/fused_moe.py | 52 +- .../layers/fused_moe/router/gate_linear.py | 12 +- vllm/model_executor/layers/layernorm.py | 105 - vllm/model_executor/layers/linear.py | 20 +- .../model_executor/layers/logits_processor.py | 100 +- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 701 +- .../layers/quantization/auto_gptq.py | 12 +- .../layers/quantization/awq_marlin.py | 4 +- .../compressed_tensors/compressed_tensors.py | 21 - .../schemes/compressed_tensors_w4a16_nvfp4.py | 3 +- .../schemes/compressed_tensors_w4a4_mxfp4.py | 8 +- .../schemes/compressed_tensors_w4a4_nvfp4.py | 255 +- .../schemes/compressed_tensors_w8a16_fp8.py | 253 +- .../model_executor/layers/quantization/fp8.py | 800 +- .../layers/quantization/input_quant_fp8.py | 4 +- .../layers/quantization/modelopt.py | 159 +- .../layers/quantization/mxfp4_sm70_moe.py | 92 +- .../layers/quantization/nvfp4_sm70_moe.py | 635 -- .../layers/quantization/sm70_turbomind.py | 179 +- .../quantization/utils/marlin_utils_fp4.py | 5 +- .../layers/quantization/utils/w8a8_utils.py | 9 +- .../layers/sparse_attn_indexer.py | 4 +- .../layers/vocab_parallel_embedding.py | 407 +- vllm/model_executor/models/config.py | 4 +- vllm/model_executor/models/deepseek_ocr.py | 60 +- vllm/model_executor/models/dflash_sm70.py | 300 - vllm/model_executor/models/gemma4_mm.py | 54 +- vllm/model_executor/models/interfaces.py | 10 - vllm/model_executor/models/qwen2_moe.py | 4 +- vllm/model_executor/models/qwen3_5.py | 228 +- vllm/model_executor/models/qwen3_5_mtp.py | 3 +- vllm/model_executor/models/qwen3_dflash.py | 801 +- vllm/model_executor/models/qwen3_dflash2.py | 484 - vllm/model_executor/models/qwen3_next.py | 29 +- vllm/model_executor/models/registry.py | 1 - vllm/model_executor/offloader/base.py | 12 - vllm/model_executor/offloader/uva.py | 46 +- vllm/model_executor/warmup/awq_sm70_warmup.py | 254 +- vllm/models/deepseek_v4/attention.py | 19 +- .../models/deepseek_v4/common/ops/__init__.py | 8 +- .../common/ops/fused_compress_quant_cache.py | 373 +- .../common/ops/save_partial_states.py | 207 - vllm/models/deepseek_v4/compressor.py | 190 +- vllm/models/deepseek_v4/nvidia/dspark.py | 83 +- vllm/models/deepseek_v4/nvidia/model.py | 25 - vllm/models/deepseek_v4/sm70/gemv.py | 282 +- vllm/models/deepseek_v4/sm70/indexer.py | 726 +- .../sm70/qnorm_rope_kv_fp8_insert.py | 132 - vllm/models/deepseek_v4/sm70/sparse.py | 6 - vllm/tool_parsers/qwen3coder_tool_parser.py | 5 +- .../processors/deepseek_ocr.py | 157 +- vllm/utils/cache.py | 4 +- vllm/v1/attention/backend.py | 13 - .../backends/ddtree_branch_triton.py | 21 +- vllm/v1/attention/backends/fa_utils.py | 1 - vllm/v1/attention/backends/flash_attn_v100.py | 1468 +-- vllm/v1/attention/backends/flashinfer_sm70.py | 3 - vllm/v1/attention/backends/gdn_attn.py | 697 +- vllm/v1/attention/backends/mla/indexer.py | 8 - vllm/v1/attention/backends/registry.py | 8 +- .../attention/ops/triton_prefill_attention.py | 6 +- .../attention/ops/triton_unified_attention.py | 12 +- vllm/v1/core/kv_cache_coordinator.py | 40 +- vllm/v1/core/kv_cache_manager.py | 76 +- vllm/v1/core/kv_cache_utils.py | 95 +- vllm/v1/core/sched/output.py | 5 - vllm/v1/core/sched/scheduler.py | 57 +- vllm/v1/core/single_type_kv_cache_manager.py | 180 +- vllm/v1/cudagraph_dispatcher.py | 281 +- vllm/v1/engine/core.py | 16 +- vllm/v1/kv_cache_interface.py | 64 +- vllm/v1/request.py | 5 - vllm/v1/sample/ops/topk_topp_triton.py | 39 - vllm/v1/sample/rejection_sampler.py | 23 - vllm/v1/sample/sampler.py | 225 +- vllm/v1/sample/thinking_budget_state.py | 8 +- vllm/v1/simple_kv_offload/manager.py | 4 +- vllm/v1/spec_decode/ddtree_metadata.py | 70 +- vllm/v1/spec_decode/dflash.py | 39 +- vllm/v1/spec_decode/draft_prob_alignment.py | 105 - vllm/v1/spec_decode/dspark.py | 101 - vllm/v1/spec_decode/llm_base_proposer.py | 358 +- vllm/v1/worker/block_table.py | 4 +- vllm/v1/worker/gpu/attn_utils.py | 166 +- vllm/v1/worker/gpu/buffer_utils.py | 6 - vllm/v1/worker/gpu/cp_utils.py | 21 - vllm/v1/worker/gpu/cudagraph_utils.py | 57 +- vllm/v1/worker/gpu/input_batch.py | 8 - vllm/v1/worker/gpu/mamba_align.py | 319 - vllm/v1/worker/gpu/model_runner.py | 182 +- vllm/v1/worker/gpu/model_states/default.py | 1 - vllm/v1/worker/gpu/model_states/interface.py | 11 - .../worker/gpu/model_states/mamba_hybrid.py | 359 +- vllm/v1/worker/gpu/model_states/whisper.py | 1 - vllm/v1/worker/gpu/sample/gumbel.py | 106 +- vllm/v1/worker/gpu/spec_decode/__init__.py | 38 - .../worker/gpu/spec_decode/dflash/__init__.py | 2 - .../gpu/spec_decode/dflash/cudagraph.py | 125 - .../gpu/spec_decode/dflash/speculator.py | 829 -- .../v1/worker/gpu/spec_decode/dflash/utils.py | 101 - .../gpu/spec_decode/dflash2/__init__.py | 2 - .../gpu/spec_decode/dflash2/ngram_assist.py | 162 - .../spec_decode/dflash2/sparse_rejection.py | 261 - .../gpu/spec_decode/dflash2/speculator.py | 621 -- .../gpu/spec_decode/eagle/eagle3_utils.py | 14 +- .../gpu/spec_decode/eagle/speculator.py | 5 - vllm/v1/worker/gpu/spec_decode/eagle/utils.py | 7 - .../spec_decode/rejection_sampler_utils.py | 336 +- vllm/v1/worker/gpu/spec_decode/speculator.py | 371 - vllm/v1/worker/gpu/spec_decode/utils.py | 15 - vllm/v1/worker/gpu/warmup.py | 43 +- vllm/v1/worker/gpu_model_runner.py | 426 +- vllm/v1/worker/gpu_worker.py | 133 +- vllm/v1/worker/mamba_utils.py | 57 - vllm/v1/worker/utils.py | 72 +- 391 files changed, 16841 insertions(+), 68436 deletions(-) delete mode 100644 benchmarks/analyze_dflash2_selector_alignment.py delete mode 100644 benchmarks/analyze_dsv4_dspark_confidence.py delete mode 100644 benchmarks/benchmark_dsv4_dspark_long_context.py delete mode 100644 benchmarks/benchmark_dsv4_gsm8k_api.py delete mode 100644 benchmarks/benchmark_dsv4_quality_api.py delete mode 100644 benchmarks/benchmark_sm70_cublas_gemm_algorithms.py delete mode 100644 benchmarks/benchmark_sm70_d256_exact_dense.py delete mode 100644 benchmarks/benchmark_sm70_dflash2_gdn_metadata.py delete mode 100644 benchmarks/benchmark_sm70_dflash2_gsm8k.py delete mode 100644 benchmarks/benchmark_sm70_dflash2_sparse_rejection.py delete mode 100644 benchmarks/benchmark_sm70_fp8_prefill_exact_dense_blas.py delete mode 100644 benchmarks/benchmark_sm70_fp8_prefill_tactics.py delete mode 100644 benchmarks/benchmark_sm70_fp8_qpn8.py delete mode 100644 benchmarks/benchmark_sm70_prefill_gemm_collective.py delete mode 100644 benchmarks/benchmark_sm70_sglang_push_chain.py delete mode 100644 benchmarks/benchmark_sm70_tp4_fused_norm_collective.py delete mode 100644 benchmarks/compare_dsv4_quality_results.py delete mode 100644 benchmarks/kernels/benchmark_sm70_nvfp4_qpn2.py delete mode 100644 benchmarks/profile_sm70_fp8_qpn8_gated_pair_ncu.py delete mode 100644 benchmarks/profile_sm70_fp8_qpn8_ncu.py delete mode 100644 cmake/patches/sm70_flash_attn_d256_gqa_arch.patch delete mode 100644 cmake/patches/sm70_flash_attn_d256_k_pingpong.patch delete mode 100644 csrc/sm70_turbomind/ops/LICENSE.v100-skinny delete mode 100644 csrc/sm70_turbomind/ops/fp8_qpn8_sm70.cu delete mode 100644 csrc/sm70_turbomind/ops/nvfp4_qpn2_sm70.cu delete mode 100644 csrc/sm70_turbomind/ops/nvfp4_qpn4_sm70.cu delete mode 100644 csrc/sm70_turbomind/ops/qwen38_prefill_cutlass.cu delete mode 100644 csrc/sm70_turbomind/ops/qwen38_prefill_cutlass.cuh delete mode 100644 csrc/sm70_turbomind/ops/sm70_top20_philox_fragment.cpp delete mode 100644 csrc/sm70_turbomind/ops/sm70_top20_philox_sampler.cu delete mode 100644 docs/assets/wechat-group-qr-7.png delete mode 100644 docs/design/sm70_deepseek_v4_dspark_verifier_20ms_long_context.md delete mode 100644 docs/design/sm70_dflash2_acceptance_rd.md delete mode 100644 docs/design/sm70_dflash2_long_verify_decay.md delete mode 100644 docs/design/sm70_dflash2_ngram_hybrid.md delete mode 100644 docs/design/sm70_dflash2_nvfp4_17ms.md delete mode 100644 docs/design/sm70_dflash2_target_graph_20ms.md delete mode 100644 docs/design/sm70_prefix_anchored_swa.md delete mode 100644 docs/design/sm70_qwen38_fp8_prefill_5500.md delete mode 100644 docs/design/sm70_qwen38_nvfp4_decode.md delete mode 100644 docs/design/sm70_qwen38_qpn8_decode.md delete mode 100644 tests/benchmarks/test_dsv4_quality_tools.py delete mode 100644 tests/config/test_prefix_anchored_swa.py delete mode 100644 tests/distributed/test_custom_all_reduce_allocator.py delete mode 100644 tests/distributed/test_static_tensor_transfer.py delete mode 100644 tests/kernels/attention/test_sm70_flash_v100_anchored_swa.py delete mode 100644 tests/kernels/attention/test_sm70_flash_v100_grouped_verify.py delete mode 100644 tests/kernels/attention/test_sm70_flash_v100_smallq_metadata.py delete mode 100644 tests/kernels/core/test_batched_weight_rms_norm.py delete mode 100644 tests/kernels/core/test_sm70_dflash2_gemma_rms.py delete mode 100644 tests/kernels/mamba/test_mrv2_mamba_align.py delete mode 100644 tests/kernels/moe/test_sm70_unquantized_moe_config.py delete mode 100644 tests/kernels/test_sm70_qwen35_gdn_split.py delete mode 100644 tests/model_executor/test_deepseek_v4_swa_rope.py delete mode 100644 tests/model_executor/test_sm70_fp8_qpn8_pp2_tp4.py delete mode 100644 tests/model_executor/test_uva_tower_offload.py delete mode 100644 tests/models/test_deepseek_v4_pipeline.py delete mode 100644 tests/quantization/test_sm70_compressed_tensors_fastpaths.py delete mode 100644 tests/quantization/test_sm70_modelopt_mixed_nvfp4.py delete mode 100644 tests/quantization/test_sm70_modelopt_nvfp4.py delete mode 100644 tests/quantization/test_sm70_modelopt_nvfp4_kernel_oracle.py delete mode 100644 tests/quantization/test_sm70_nvfp4_qpn2.py delete mode 100644 tests/transformers_utils/test_deepseek_ocr_resolution_modes.py delete mode 100644 tests/v1/attention/test_prefix_anchored_swa_mask.py delete mode 100644 tests/v1/core/test_swa_inflight_window_free.py delete mode 100644 tests/v1/cudagraph/test_sm70_mtp_split_cudagraphs.py delete mode 100644 tests/v1/spec_decode/test_dflash2.py delete mode 100644 tests/v1/spec_decode/test_dflash2_ngram_assist.py delete mode 100644 tests/v1/spec_decode/test_dflash2_selector_analysis.py delete mode 100644 tests/v1/spec_decode/test_dflash_mrv2_config.py delete mode 100644 tests/v1/worker/test_gpu_model_runner_v2_kv_zero.py delete mode 100644 tests/v1/worker/test_gpu_warmup_blocks.py delete mode 100644 tests/v1/worker/test_gpu_worker_static_pp.py delete mode 100644 tests/v1/worker/test_kv_block_zeroer.py delete mode 100644 vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py delete mode 100644 vllm/model_executor/models/dflash_sm70.py delete mode 100644 vllm/model_executor/models/qwen3_dflash2.py delete mode 100644 vllm/v1/worker/gpu/mamba_align.py delete mode 100644 vllm/v1/worker/gpu/spec_decode/dflash/__init__.py delete mode 100644 vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py delete mode 100644 vllm/v1/worker/gpu/spec_decode/dflash/speculator.py delete mode 100644 vllm/v1/worker/gpu/spec_decode/dflash/utils.py delete mode 100644 vllm/v1/worker/gpu/spec_decode/dflash2/__init__.py delete mode 100644 vllm/v1/worker/gpu/spec_decode/dflash2/ngram_assist.py delete mode 100644 vllm/v1/worker/gpu/spec_decode/dflash2/sparse_rejection.py delete mode 100644 vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py delete mode 100644 vllm/v1/worker/gpu/spec_decode/speculator.py diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 1bfe200ade..c7ef61b719 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.4.0 + - pip install helion==1.0.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 0259847d99..68e6a5762e 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.4.0 + - pip install helion==1.0.0 - pytest -v -s kernels/helion/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 39de56e352..bd92ec9854 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/.*|csrc/sm70_turbomind/lmdeploy/.*)' +exclude: '(vllm/third_party/.*|flash_qla/.*)' repos: - repo: https://github.com/astral-sh/ruff-pre-commit rev: v0.14.0 @@ -35,11 +35,6 @@ 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: @@ -182,6 +177,11 @@ 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 d07300fdff..7d863a8053 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -396,10 +396,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") "${VLLM_SM70_TURBOMIND_LMDEPLOY_ROOT}/src/turbomind/kernels/gemm/kernel/sm70_884_8.cu" "${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/nvfp4_qpn4_sm70.cu" - "${VLLM_SM70_TURBOMIND_ROOT}/ops/qwen38_prefill_cutlass.cu" - "${VLLM_SM70_TURBOMIND_ROOT}/ops/nvfp4_qpn2_sm70.cu" "${VLLM_SM70_TURBOMIND_ROOT}/ops/awq_sm70_gemm.cu") set_gencode_flags_for_srcs( @@ -713,44 +709,6 @@ if(VLLM_REGISTER_BASIC_ACTIVATION_IN_C) target_compile_definitions(_C PRIVATE VLLM_REGISTER_BASIC_ACTIVATION_IN_C=1) endif() -# Keep the accepted SM70 chunked sampler in a separate wheel module so adding -# it does not perturb the quality- and speed-frozen primary vllm._C module. -if(VLLM_GPU_LANG STREQUAL "CUDA" AND SM70_TURBOMIND_ARCHS) - set(VLLM_SM70_SAMPLER_FRAGMENT_SRCS - "csrc/sm70_turbomind/ops/sm70_top20_philox_fragment.cpp" - "csrc/sm70_turbomind/ops/sm70_top20_philox_sampler.cu") - set_gencode_flags_for_srcs( - SRCS "${VLLM_SM70_SAMPLER_FRAGMENT_SRCS}" - CUDA_ARCHS "${SM70_TURBOMIND_ARCHS}") - define_extension_target( - _sm70_sampler_C - DESTINATION vllm - LANGUAGE CUDA - 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 - CXX_STANDARD_REQUIRED ON - CXX_EXTENSIONS OFF - CUDA_STANDARD 17 - CUDA_STANDARD_REQUIRED ON - CUDA_EXTENSIONS OFF) - target_compile_definitions(_sm70_sampler_C PRIVATE - TORCH_API_INCLUDE_EXTENSION_H - $<$:__CUDA_NO_HALF_OPERATORS__> - $<$:__CUDA_NO_HALF_CONVERSIONS__> - $<$:__CUDA_NO_BFLOAT16_CONVERSIONS__> - $<$:__CUDA_NO_HALF2_OPERATORS__>) - target_compile_options(_sm70_sampler_C PRIVATE - $<$:-UNDEBUG> - $<$:-UNDEBUG;-lineinfo>) - target_link_libraries(_sm70_sampler_C PRIVATE - "${TORCH_INSTALL_PREFIX}/lib/libtorch_python.so") -endif() - if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # # _C_stable_libtorch extension (ops registered via STABLE_TORCH_LIBRARY) diff --git a/README.md b/README.md index f305d6e580..59a93472e6 100644 --- a/README.md +++ b/README.md @@ -1,20 +1,26 @@ # 1Cat-vLLM -> 一猫之下始终相信，V100 不该在今天的大模型浪潮里被轻易宣判“过时”。 +> 一猫之下始终相信，V100 不该在今天的大模型浪潮中被轻易宣判“过时”。 > -> 1Cat-vLLM 是面向 **SM70 / Tesla V100** 的 vLLM 工程分支。项目围绕 -> AWQ、注意力后端、长上下文稳定性、MTP 投机解码、运行时默认值和部署 -> 路径做了成体系的优化，让更多现代模型场景在 V100 上真正变得可用、 -> 好用、能持续部署。 +> 1Cat-vLLM 是一个专注于 **SM70 / Tesla V100** 的 vLLM 工程分支。我们 +> 不追求覆盖所有硬件，而是围绕 Volta 架构补齐并优化现代大模型推理的 +> 关键路径，包括量化推理、注意力后端、长上下文、MTP 投机解码、 +> CUDA Graph、分布式通信、运行时策略和部署工具链，让 Qwen、DeepSeek +> 等新模型在 V100 上从“能跑”走向“可部署、可验证、可持续优化”。 > -> 我们希望把一猫之下在 V100 上的工程经验、优化成果和验证过程贡献给 -> 开源社区，也欢迎继续使用 V100 的个人开发者、工作室和团队一起反馈、 -> 复现和改进。 +> 我们希望把来自真实 V100 环境的工程实现、性能数据和验证过程贡献给 +> 开源社区，也欢迎仍在使用 V100 的个人开发者、工作室和团队参与测试、 +> 复现、反馈和共建，让这批依然有价值的算力继续发挥作用。 -1Cat-vLLM is a **Tesla V100 / SM70** focused vLLM fork for serving modern -Qwen-class AWQ and experimental FP8 models on Volta GPUs. It integrates -TurboMind-derived SM70 kernels, a V100 FlashAttention path, runtime defaults -for long-context serving, and OpenAI-compatible API fixes for common clients. +1Cat-vLLM is a vLLM engineering fork focused on **SM70 / Tesla V100**. +Rather than targeting every hardware platform, it fills and optimizes the +critical inference paths required by modern models on Volta, including +quantized inference, attention backends, long-context serving, MTP speculative +decoding, CUDA Graph execution, distributed communication, runtime policies, +and deployment tooling. + +The goal is to make models such as Qwen and DeepSeek not merely runnable on +V100, but deployable, reproducible, and continuously optimizable. ## Project Focus @@ -25,6 +31,12 @@ for long-context serving, and OpenAI-compatible API fixes for common clients. - **V100 FlashAttention path**: `FLASH_ATTN_V100` decode and prefill backend for Volta GPUs, with SM70 compile-graph, guarded XQA decode, and D=256 paged-prefix low-smem fast paths enabled by default. +- **DeepSeek V4 Flash**: supports running the original, unmodified DeepSeek V4 + Flash weights across eight GPUs. +- **Day-0 Qwen3.8-27B support**: native Day-0 support for Qwen3.8-27B. +- **Quantization support**: FP8, NVFP4, MXFP4, AWQ, and GPTQ inference paths + are included for supported model and hardware combinations. Availability + and production status vary by checkpoint and runtime route. - **Long-context serving**: public profiles default to 256K context where the model and memory budget allow it. - **MTP serving**: Qwen3.6-class MTP speculative decoding remains available as @@ -42,6 +54,7 @@ for long-context serving, and OpenAI-compatible API fixes for common clients. - `tclf90/Qwen3.6-27B-AWQ` - `tclf90/Qwen3.6-35B-A3B-AWQ` - `tclf90/Qwen3.5-122B-A10B-AWQ` for larger 4-GPU setups +- `Qwen/Qwen3.8-27B` for the latest supported Qwen release The launch examples use local paths such as `/path/to/Qwen3.6-27B-AWQ`. Replace them with your local model path or a Hugging Face repository id. @@ -57,6 +70,7 @@ validation. | --- | --- | | 4 x Tesla V100 32 GB | Main public reference target | | 2 x Tesla V100 32 GB | Supported for selected 27B profiles with lower concurrency | +| 8 x Tesla V100 32 GB | DeepSeek V4 Flash original-checkpoint TP8 target | Typical model placement: @@ -65,6 +79,10 @@ Typical model placement: the long-context public default. - `Qwen3.6-35B-A3B-AWQ`: TP4 recommended. - `Qwen3.5-122B-A10B-AWQ`: TP4 supported for larger deployments. +- `DeepSeek-V4-Flash`: the original mixed MXFP4/FP8 checkpoint is supported + on eight 32 GB V100 GPUs with TP8. +- `Qwen3.8-27B`: supported natively from Day 0, with validated FP16 and FP8 + serving paths. Multimodal defaults: @@ -349,17 +367,26 @@ python -m pip install -e . --no-build-isolation 请使用微信扫描下方二维码加入群组： -![1Cat-vLLM 微信交流群二维码](docs/assets/wechat-group-qr-7.png) +![1Cat-vLLM 微信交流群二维码](docs/assets/wechat-group-qr-6.png) + +> 提示：微信群二维码通常 7 天内有效。若扫描失败或提示过期，请联系微信号：YM_isi。 + +## Roadmap + +1Cat-vLLM will continue to prioritize Tesla V100 / SM70, making modern LLM inference more usable, stable, and deployable on existing GPUs. + +- **Near term:** Complete Volta compatibility and validation for DeepSeek V4, GLM5, and other new model families, while improving long-context serving, quantization, multi-GPU parallelism, and runtime stability. +- **Mid term:** Expand the Volta model support matrix with a focus on models up to 300B parameters, backed by reproducible performance, quality, and release validation. +- **Long term:** Build on the Volta foundation and extend hardware support in stages: **Volta → Turing/Ampere → CDNA**, bringing sustainable LLM inference to a broader range of existing accelerators. -> 提示：微信群二维码通常 7 天内有效。若扫描失败或提示过期，请重新打开本页查看最新图片，或关注仓库更新。 ## Repository Notes -- Upstream project: [vLLM](https://github.com/vllm-project/vllm) -- This fork focuses on SM70 AWQ support, V100-oriented attention/runtime - tuning, and experimental FP8/MTP/DFlash validation paths. -- Prebuilt wheels are the public installation path. -- Source builds are for development and kernel work. +- This fork focuses on SM70 quantized inference, V100-oriented attention and + long-context tuning, model-specific runtime and deployment paths, and + continued MTP and DFlash research. +- Production status is route-specific. Use the documented public profiles and + validated release matrices as the source of truth. ## Acknowledgements diff --git a/benchmarks/analyze_dflash2_selector_alignment.py b/benchmarks/analyze_dflash2_selector_alignment.py deleted file mode 100644 index b332106835..0000000000 --- a/benchmarks/analyze_dflash2_selector_alignment.py +++ /dev/null @@ -1,517 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Audit and sweep exact DFlash2 selector proposals from alignment dumps. - -This tool consumes the B1 records emitted by the GPU-runner compact rejection -path when ``VLLM_SPEC_DUMP_ALIGNMENT=1``. It never changes target sampling. -Counterfactual results are one-step overlap proxies on the recorded prefixes; -only an end-to-end run can establish a new acceptance-length result. -""" - -from __future__ import annotations - -import argparse -import glob -import json -import math -from dataclasses import asdict, dataclass -from pathlib import Path -from typing import Any - -import torch - - -@dataclass(frozen=True) -class ProposalConfig: - name: str - proposal_temperature_scale: float = 1.0 - proposal_top_p: float = 1.0 - unary_scale: float = 1.0 - edge_scale: float = 1.0 - future_beta: float = 0.0 - greedy_mix: float = 0.0 - use_cached_logits: bool = False - - -@dataclass -class AlignmentRecord: - path: Path - step: int - temperature: float - top_p: float - target_topk_ids: torch.Tensor - target_topk_logits: torch.Tensor - candidate_ids: torch.Tensor - realized_logits: torch.Tensor - unary_logits: torch.Tensor - lattice_scores: torch.Tensor - draft_sampled: torch.Tensor - num_sampled: int - - -def _compact_probs(logits: torch.Tensor, top_p: float) -> torch.Tensor: - """Apply the compact top-p contract used by the sparse rejection kernel.""" - logits = logits.to(torch.float64) - order = torch.argsort(logits, descending=True, stable=True) - sorted_probs = torch.softmax(logits[order], dim=-1) - cumulative_before = torch.cumsum(sorted_probs, dim=-1) - sorted_probs - keep_sorted = (top_p >= 1.0) | (cumulative_before < top_p) - kept = torch.zeros_like(sorted_probs) - kept[keep_sorted] = sorted_probs[keep_sorted] - kept /= kept.sum() - probs = torch.zeros_like(kept) - probs[order] = kept - return probs - - -def _distribution_overlap( - target_ids: torch.Tensor, - target_probs: torch.Tensor, - proposal_ids: torch.Tensor, - proposal_probs: torch.Tensor, -) -> float: - matches = target_ids[:, None] == proposal_ids[None, :] - proposal_on_target = torch.where( - matches, - proposal_probs[None, :], - torch.zeros((), dtype=proposal_probs.dtype), - ).sum(dim=1) - return float(torch.minimum(target_probs, proposal_on_target).sum()) - - -def _support_mass( - target_ids: torch.Tensor, - target_probs: torch.Tensor, - proposal_ids: torch.Tensor, -) -> float: - supported = (target_ids[:, None] == proposal_ids[None, :]).any(dim=1) - return float(target_probs[supported].sum()) - - -def _transformed_lattice( - record: AlignmentRecord, config: ProposalConfig -) -> torch.Tensor: - unary = record.unary_logits[:, None, :].to(torch.float64) - lattice = record.lattice_scores.to(torch.float64) - edge = lattice - unary - denominator = record.temperature * config.proposal_temperature_scale - return (config.unary_scale * unary + config.edge_scale * edge) / denominator - - -def _backward_messages(lattice: torch.Tensor) -> torch.Tensor: - """Compute normalized log-sum-exp future messages for a KxK chain.""" - num_steps, top_k, _ = lattice.shape - future = torch.zeros((num_steps, top_k), dtype=lattice.dtype) - for step in range(num_steps - 2, -1, -1): - values = lattice[step + 1] + future[step + 1][None, :] - future[step] = torch.logsumexp(values, dim=-1) - future[step] -= future[step].max() - return future - - -def _path_predecessors(record: AlignmentRecord) -> list[int]: - num_steps = record.candidate_ids.shape[0] - if record.draft_sampled.numel() < num_steps + 1: - raise ValueError(f"{record.path}: draft_sampled is missing the anchor row") - predecessors = [0] - for step in range(num_steps - 1): - proposed = int(record.draft_sampled[step + 1]) - matches = torch.nonzero(record.candidate_ids[step] == proposed).flatten() - if matches.numel() != 1: - raise ValueError( - f"{record.path}: proposed token {proposed} is not unique at step {step}" - ) - predecessors.append(int(matches[0])) - return predecessors - - -def _proposal_probs( - record: AlignmentRecord, - config: ProposalConfig, -) -> list[torch.Tensor]: - if config.use_cached_logits: - rows = [ - _compact_probs(row, config.proposal_top_p) for row in record.realized_logits - ] - else: - lattice = _transformed_lattice(record, config) - future = _backward_messages(lattice) - predecessors = _path_predecessors(record) - rows = [ - _compact_probs( - lattice[step, predecessor] + config.future_beta * future[step], - config.proposal_top_p, - ) - for step, predecessor in enumerate(predecessors) - ] - if config.greedy_mix == 0.0: - return rows - if not 0.0 <= config.greedy_mix <= 1.0: - raise ValueError("greedy_mix must be in [0, 1]") - mixed_rows = [] - for probs in rows: - mixed = probs * (1.0 - config.greedy_mix) - mixed[torch.argmax(probs)] += config.greedy_mix - mixed_rows.append(mixed) - return mixed_rows - - -def _target_rows(record: AlignmentRecord) -> list[torch.Tensor]: - rows = [] - for logits in record.target_topk_logits[: record.candidate_ids.shape[0]]: - rows.append(_compact_probs(logits / record.temperature, record.top_p)) - return rows - - -def _load_record(path: Path) -> AlignmentRecord: - payload = torch.load(path, map_location="cpu", weights_only=True) - if payload.get("format") != "dflash2_selector_alignment_v1": - raise ValueError(f"{path}: unsupported alignment format") - candidate_ids = payload["selector_candidate_ids"].to(torch.int64) - cached_ids = payload["draft_candidate_ids"].to(torch.int64) - if not torch.equal(candidate_ids, cached_ids): - raise ValueError(f"{path}: packed selector IDs do not match request-slot cache") - record = AlignmentRecord( - path=path, - step=int(payload["step"]), - temperature=float(payload["temperature"]), - top_p=float(payload["top_p"]), - target_topk_ids=payload["target_topk_ids"].to(torch.int64), - target_topk_logits=payload["target_topk_logits"].to(torch.float64), - candidate_ids=candidate_ids, - realized_logits=payload["draft_realized_logits"].to(torch.float64), - unary_logits=payload["selector_unary_logits"].to(torch.float64), - lattice_scores=payload["selector_lattice_scores"].to(torch.float64), - draft_sampled=payload["draft_sampled"].to(torch.int64).flatten(), - num_sampled=int(payload["num_sampled"].flatten()[0]), - ) - return record - - -def _mean(values: list[float]) -> float: - return float(sum(values) / len(values)) if values else math.nan - - -def _completion_proxy(per_position_overlap: list[float]) -> float: - survival = 1.0 - result = 1.0 - for overlap in per_position_overlap: - survival *= overlap - result += survival - return result - - -def _evaluate(records: list[AlignmentRecord], config: ProposalConfig) -> dict[str, Any]: - num_steps = records[0].candidate_ids.shape[0] - overlaps = [[] for _ in range(num_steps)] - for record in records: - target_rows = _target_rows(record) - proposal_rows = _proposal_probs(record, config) - for step, (target_probs, proposal_probs) in enumerate( - zip(target_rows, proposal_rows) - ): - overlaps[step].append( - _distribution_overlap( - record.target_topk_ids[step], - target_probs, - record.candidate_ids[step], - proposal_probs, - ) - ) - means = [_mean(values) for values in overlaps] - return { - "config": asdict(config), - "num_records": len(records), - "mean_overlap_by_position": means, - "completion_length_proxy": _completion_proxy(means), - } - - -def _sweep_configs() -> list[ProposalConfig]: - configs = [ProposalConfig(name="current", use_cached_logits=True)] - seen: set[tuple[float, float, float, float, float, float]] = set() - - def add( - family: str, - temperature: float = 1.0, - top_p: float = 1.0, - unary: float = 1.0, - edge: float = 1.0, - beta: float = 0.0, - greedy_mix: float = 0.0, - ) -> None: - key = (temperature, top_p, unary, edge, beta, greedy_mix) - if key in seen: - return - seen.add(key) - configs.append( - ProposalConfig( - name=( - f"{family}:t={temperature:g},p={top_p:g},u={unary:g}," - f"e={edge:g},b={beta:g},g={greedy_mix:g}" - ), - proposal_temperature_scale=temperature, - proposal_top_p=top_p, - unary_scale=unary, - edge_scale=edge, - future_beta=beta, - greedy_mix=greedy_mix, - ) - ) - - for temperature in (0.7, 0.8, 0.9, 0.95, 1.0, 1.05, 1.1, 1.2, 1.3): - add("temperature", temperature=temperature) - for top_p in (0.9, 0.95, 0.98, 0.99): - for temperature in (0.8, 0.9, 1.0): - add("nucleus", temperature=temperature, top_p=top_p) - for greedy_mix in (0.1, 0.2, 0.3, 0.4, 0.5): - for temperature in (0.8, 0.9, 1.0, 1.1): - add( - "greedy-mixture", - temperature=temperature, - greedy_mix=greedy_mix, - ) - for unary in (0.75, 1.0, 1.25): - for edge in (0.5, 0.75, 1.0, 1.25, 1.5): - for temperature in (0.85, 1.0, 1.15): - add( - "edge-calibration", - temperature=temperature, - unary=unary, - edge=edge, - ) - for beta in (0.25, 0.5, 0.75, 1.0): - for temperature in (0.8, 0.9, 1.0, 1.1): - for edge in (0.75, 1.0, 1.25): - add( - "future-message", - temperature=temperature, - edge=edge, - beta=beta, - ) - return configs - - -def _baseline_diagnostics(records: list[AlignmentRecord]) -> dict[str, Any]: - num_steps = records[0].candidate_ids.shape[0] - support = [[] for _ in range(num_steps)] - cached_lattice_max_abs: list[float] = [] - for record in records: - target_rows = _target_rows(record) - predecessors = _path_predecessors(record) - for step, target_probs in enumerate(target_rows): - support[step].append( - _support_mass( - record.target_topk_ids[step], - target_probs, - record.candidate_ids[step], - ) - ) - expected = ( - record.lattice_scores[step, predecessors[step]] / record.temperature - ) - cached_lattice_max_abs.append( - float((expected - record.realized_logits[step]).abs().max()) - ) - support_means = [_mean(values) for values in support] - return { - "observed_mean_completion_tokens_per_round": _mean( - [float(record.num_sampled) for record in records] - ), - "candidate_support_mass_by_position": support_means, - "candidate_support_completion_upper_proxy": _completion_proxy(support_means), - "cached_vs_lattice_max_abs": max(cached_lattice_max_abs, default=math.nan), - } - - -def summarize(pattern: str, top_n: int) -> dict[str, Any]: - paths = [Path(path) for path in sorted(glob.glob(pattern))] - if not paths: - raise ValueError(f"no selector alignment dumps matched {pattern!r}") - records = [_load_record(path) for path in paths] - tune = [record for record in records if record.step % 2 == 0] - holdout = [record for record in records if record.step % 2 == 1] - if not tune or not holdout: - tune = records - holdout = records - - configs = _sweep_configs() - tune_results = [_evaluate(tune, config) for config in configs] - holdout_by_name = { - result["config"]["name"]: result - for result in (_evaluate(holdout, config) for config in configs) - } - current_tune = tune_results[0]["completion_length_proxy"] - current_holdout = holdout_by_name["current"]["completion_length_proxy"] - position_policy = [] - position_tune_overlaps = [] - position_holdout_overlaps = [] - for position in range(records[0].candidate_ids.shape[0]): - best = max( - tune_results, - key=lambda result: result["mean_overlap_by_position"][position], - ) - holdout_best = holdout_by_name[best["config"]["name"]] - position_policy.append( - { - "position": position, - "config": best["config"], - "tune_overlap": best["mean_overlap_by_position"][position], - "holdout_overlap": holdout_best["mean_overlap_by_position"][position], - } - ) - position_tune_overlaps.append(best["mean_overlap_by_position"][position]) - position_holdout_overlaps.append( - holdout_best["mean_overlap_by_position"][position] - ) - ranked = sorted( - tune_results[1:], - key=lambda result: result["completion_length_proxy"], - reverse=True, - ) - leaders = [] - for result in ranked[:top_n]: - holdout_result = holdout_by_name[result["config"]["name"]] - leaders.append( - { - "config": result["config"], - "tune_completion_length_proxy": result["completion_length_proxy"], - "tune_delta": result["completion_length_proxy"] - current_tune, - "holdout_completion_length_proxy": holdout_result[ - "completion_length_proxy" - ], - "holdout_delta": ( - holdout_result["completion_length_proxy"] - current_holdout - ), - "tune_mean_overlap_by_position": result["mean_overlap_by_position"], - "holdout_mean_overlap_by_position": holdout_result[ - "mean_overlap_by_position" - ], - } - ) - - return { - "contract": { - "glob": pattern, - "num_records": len(records), - "num_tune_records": len(tune), - "num_holdout_records": len(holdout), - "split": "even/odd sampler step", - "counterfactual_scope": "one-step recorded-prefix overlap proxy", - "end_to_end_claim": False, - }, - "baseline": _baseline_diagnostics(records), - "current": { - "tune": tune_results[0], - "holdout": holdout_by_name["current"], - }, - "position_policy": { - "selection": "best tune overlap independently at each depth", - "rows": position_policy, - "tune_completion_length_proxy": _completion_proxy(position_tune_overlaps), - "tune_delta": (_completion_proxy(position_tune_overlaps) - current_tune), - "holdout_completion_length_proxy": _completion_proxy( - position_holdout_overlaps - ), - "holdout_delta": ( - _completion_proxy(position_holdout_overlaps) - current_holdout - ), - }, - "leaders": leaders, - } - - -def render_markdown(summary: dict[str, Any]) -> str: - contract = summary["contract"] - baseline = summary["baseline"] - current = summary["current"] - position_policy = summary["position_policy"] - lines = [ - "# DFlash2 Selector Alignment Sweep", - "", - ( - f"- Records: `{contract['num_records']}` " - f"(tune `{contract['num_tune_records']}`, " - f"holdout `{contract['num_holdout_records']}`)" - ), - ( - "- Counterfactual scope: one-step overlap on recorded prefixes; " - "this is not an end-to-end acceptance claim." - ), - ( - "- Observed completion tokens/round: " - f"`{baseline['observed_mean_completion_tokens_per_round']:.4f}`" - ), - ( - "- Current overlap proxy (tune/holdout): " - f"`{current['tune']['completion_length_proxy']:.4f}` / " - f"`{current['holdout']['completion_length_proxy']:.4f}`" - ), - ( - "- Fixed-top16 support upper proxy: " - f"`{baseline['candidate_support_completion_upper_proxy']:.4f}`" - ), - ( - "- Cached/lattice max absolute mismatch: " - f"`{baseline['cached_vs_lattice_max_abs']:.3e}`" - ), - ( - "- Position-wise held-out proxy/delta: " - f"`{position_policy['holdout_completion_length_proxy']:.4f}` / " - f"`{position_policy['holdout_delta']:+.4f}`" - ), - "", - "| candidate | tune proxy | tune delta | holdout proxy | holdout delta |", - "| --- | ---: | ---: | ---: | ---: |", - ] - for row in summary["leaders"]: - lines.append( - f"| `{row['config']['name']}` | " - f"{row['tune_completion_length_proxy']:.4f} | " - f"{row['tune_delta']:+.4f} | " - f"{row['holdout_completion_length_proxy']:.4f} | " - f"{row['holdout_delta']:+.4f} |" - ) - lines.extend( - [ - "", - "## Position-wise calibration policy", - "", - "| position | tune overlap | holdout overlap | configuration |", - "| ---: | ---: | ---: | --- |", - ] - ) - for row in position_policy["rows"]: - lines.append( - f"| {row['position']} | {row['tune_overlap']:.4f} | " - f"{row['holdout_overlap']:.4f} | `{row['config']['name']}` |" - ) - lines.append("") - return "\n".join(lines) - - -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument( - "--glob", - default="/tmp/spec_alignment_dflash2_selector_*.pt", - help="Glob for GPU-runner DFlash2 selector alignment records.", - ) - parser.add_argument("--top", type=int, default=20) - parser.add_argument("--out-json", type=Path) - parser.add_argument("--out-md", type=Path) - args = parser.parse_args() - - summary = summarize(args.glob, args.top) - markdown = render_markdown(summary) - if args.out_json is not None: - args.out_json.parent.mkdir(parents=True, exist_ok=True) - args.out_json.write_text(json.dumps(summary, indent=2), encoding="utf-8") - if args.out_md is not None: - args.out_md.parent.mkdir(parents=True, exist_ok=True) - args.out_md.write_text(markdown, encoding="utf-8") - print(markdown) - - -if __name__ == "__main__": - main() diff --git a/benchmarks/analyze_dsv4_dspark_confidence.py b/benchmarks/analyze_dsv4_dspark_confidence.py deleted file mode 100644 index 0445483f4c..0000000000 --- a/benchmarks/analyze_dsv4_dspark_confidence.py +++ /dev/null @@ -1,288 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Calibrate and summarize DeepSeek-V4 DSpark confidence-head dumps.""" - -from __future__ import annotations - -import argparse -import json -import math -from collections import defaultdict -from pathlib import Path -from typing import Any - -import torch - - -def _ece(probs: torch.Tensor, targets: torch.Tensor, bins: int = 15) -> float: - if probs.numel() == 0: - return math.nan - edges = torch.linspace(0.0, 1.0, bins + 1, dtype=torch.float64) - total = probs.numel() - value = 0.0 - for index in range(bins): - if index + 1 == bins: - mask = (probs >= edges[index]) & (probs <= edges[index + 1]) - else: - mask = (probs >= edges[index]) & (probs < edges[index + 1]) - count = int(mask.sum()) - if count: - gap = (probs[mask].mean() - targets[mask].mean()).abs() - value += float(gap) * count / total - return value - - -def _soft_nll(logits: torch.Tensor, targets: torch.Tensor) -> float: - if logits.numel() == 0: - return math.nan - return float(torch.nn.functional.binary_cross_entropy_with_logits(logits, targets)) - - -def _fit_temperature(logits: torch.Tensor, targets: torch.Tensor) -> float: - if logits.numel() < 2: - return 1.0 - # A dense one-dimensional search is deterministic, stable for fractional - # TV labels, and avoids adding scipy/sklearn to the benchmark environment. - temperatures = torch.logspace(-1.0, 1.0, 801, dtype=torch.float64) - scaled = logits.to(torch.float64).unsqueeze(0) / temperatures.unsqueeze(1) - losses = torch.nn.functional.binary_cross_entropy_with_logits( - scaled, - targets.to(torch.float64).expand_as(scaled), - reduction="none", - ).mean(dim=1) - return float(temperatures[int(losses.argmin())]) - - -def _binary_auc(scores: torch.Tensor, labels: torch.Tensor) -> float: - positives = scores[labels == 1] - negatives = scores[labels == 0] - if positives.numel() == 0 or negatives.numel() == 0: - return math.nan - # Proposal blocks are at most five positions in the production contract; - # this exact pairwise definition is small and handles ties correctly. - comparisons = positives[:, None] - negatives[None, :] - return float((comparisons.gt(0).float() + 0.5 * comparisons.eq(0).float()).mean()) - - -def _summary( - logits: torch.Tensor, - targets: torch.Tensor, - temperature: float, - observed_logits: torch.Tensor | None = None, - observed_labels: torch.Tensor | None = None, -) -> dict[str, Any]: - scaled_logits = logits / temperature - probs = scaled_logits.sigmoid() - result: dict[str, Any] = { - "count": int(logits.numel()), - "temperature": temperature, - "predicted_mean": float(probs.mean()) if probs.numel() else math.nan, - "target_mean": float(targets.mean()) if targets.numel() else math.nan, - "ece": _ece(probs, targets), - "brier": ( - float(torch.square(probs - targets).mean()) if probs.numel() else math.nan - ), - "soft_nll": _soft_nll(scaled_logits, targets), - } - if ( - observed_logits is not None - and observed_labels is not None - and observed_labels.numel() - ): - result["observed_count"] = int(observed_labels.numel()) - result["observed_acceptance"] = float(observed_labels.mean()) - result["observed_auc"] = _binary_auc( - torch.sigmoid(observed_logits / temperature), observed_labels - ) - return result - - -def _load_rows(input_dir: Path, rank: int) -> list[dict[str, Any]]: - rows: list[dict[str, Any]] = [] - for path in sorted(input_dir.glob("spec_alignment_*.pt")): - payload = torch.load(path, map_location="cpu", weights_only=True) - if int(payload.get("rank", -1)) != rank: - continue - required = ( - "draft_confidence_logits", - "conditional_acceptance_targets", - "num_draft_tokens", - "output_valid_counts", - ) - if not all(key in payload for key in required): - continue - payload["_path"] = str(path) - rows.append(payload) - return rows - - -def analyze(rows: list[dict[str, Any]], calibration_fraction: float) -> dict[str, Any]: - if not rows: - raise ValueError("no rank-matching confidence alignment dumps were found") - max_position = max(max(row["num_draft_tokens"], default=0) for row in rows) - by_position: dict[int, dict[str, list[torch.Tensor]]] = defaultdict( - lambda: defaultdict(list) - ) - prefix_by_position: dict[int, dict[str, list[torch.Tensor]]] = defaultdict( - lambda: defaultdict(list) - ) - - for row_index, row in enumerate(rows): - logits = row["draft_confidence_logits"].to(torch.float64) - targets = row["conditional_acceptance_targets"].to(torch.float64) - valid_counts = row["output_valid_counts"].to(torch.int64) - offset = 0 - for request_index, count_raw in enumerate(row["num_draft_tokens"]): - count = int(count_raw) - request_logits = logits[offset : offset + count] - request_targets = targets[offset : offset + count] - accepted = max(0, min(count, int(valid_counts[request_index]) - 1)) - split = ( - "calibration" - if row_index < len(rows) * calibration_fraction - else "evaluation" - ) - cumulative_pred_logits = request_logits - cumulative_targets = request_targets.cumprod(dim=0) - for position in range(count): - bucket = by_position[position] - bucket[f"{split}_logits"].append( - request_logits[position : position + 1] - ) - bucket[f"{split}_targets"].append( - request_targets[position : position + 1] - ) - if position <= accepted and position < count: - observed = torch.tensor( - [1.0 if position < accepted else 0.0], dtype=torch.float64 - ) - bucket[f"{split}_observed_logits"].append( - request_logits[position : position + 1] - ) - bucket[f"{split}_observed"].append(observed) - - prefix_bucket = prefix_by_position[position] - prefix_bucket[f"{split}_logits"].append( - cumulative_pred_logits[: position + 1] - ) - prefix_bucket[f"{split}_targets"].append( - cumulative_targets[position : position + 1] - ) - offset += count - - temperatures: list[float] = [] - position_rows: list[dict[str, Any]] = [] - for position in range(max_position): - bucket = by_position[position] - calibration_logits = torch.cat(bucket["calibration_logits"]) - calibration_targets = torch.cat(bucket["calibration_targets"]) - temperature = _fit_temperature(calibration_logits, calibration_targets) - temperatures.append(temperature) - evaluation_logits = torch.cat( - bucket["evaluation_logits"] or bucket["calibration_logits"] - ) - evaluation_targets = torch.cat( - bucket["evaluation_targets"] or bucket["calibration_targets"] - ) - observed_parts = bucket["evaluation_observed"] or bucket["calibration_observed"] - observed_logits_parts = ( - bucket["evaluation_observed_logits"] - or bucket["calibration_observed_logits"] - ) - observed = torch.cat(observed_parts) if observed_parts else None - observed_logits = ( - torch.cat(observed_logits_parts) if observed_logits_parts else None - ) - position_rows.append( - { - "position": position + 1, - "raw": _summary( - evaluation_logits, - evaluation_targets, - 1.0, - observed_logits, - observed, - ), - "calibrated": _summary( - evaluation_logits, - evaluation_targets, - temperature, - observed_logits, - observed, - ), - } - ) - - prefix_rows: list[dict[str, Any]] = [] - for position in range(max_position): - bucket = prefix_by_position[position] - logits_parts = bucket["evaluation_logits"] or bucket["calibration_logits"] - target_parts = bucket["evaluation_targets"] or bucket["calibration_targets"] - prefix_probs = torch.stack( - [ - torch.sigmoid(part / torch.tensor(temperatures[: part.numel()])).prod() - for part in logits_parts - ] - ) - prefix_targets = torch.cat(target_parts) - prefix_rows.append( - { - "position": position + 1, - "count": int(prefix_probs.numel()), - "predicted_mean": float(prefix_probs.mean()), - "target_mean": float(prefix_targets.mean()), - "ece": _ece(prefix_probs, prefix_targets), - "brier": float(torch.square(prefix_probs - prefix_targets).mean()), - } - ) - - return { - "schema_version": 1, - "dump_count": len(rows), - "rank": int(rows[0].get("rank", -1)), - "step_range": [int(rows[0]["step"]), int(rows[-1]["step"])], - "calibration_fraction": calibration_fraction, - "sequential_temperatures": temperatures, - "per_conditional_position": position_rows, - "per_prefix_position": prefix_rows, - "source_files": [row["_path"] for row in rows], - } - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--input-dir", type=Path, required=True) - parser.add_argument("--rank", type=int, default=0) - parser.add_argument("--calibration-fraction", type=float, default=0.7) - parser.add_argument("--output", type=Path, required=True) - args = parser.parse_args() - if not 0.0 < args.calibration_fraction < 1.0: - parser.error("--calibration-fraction must be in (0, 1)") - return args - - -def main() -> None: - args = parse_args() - rows = _load_rows(args.input_dir, args.rank) - result = analyze(rows, args.calibration_fraction) - args.output.parent.mkdir(parents=True, exist_ok=True) - - def json_safe(value: Any) -> Any: - if isinstance(value, float) and not math.isfinite(value): - return None - if isinstance(value, dict): - return {key: json_safe(item) for key, item in value.items()} - if isinstance(value, list): - return [json_safe(item) for item in value] - return value - - args.output.write_text( - json.dumps(json_safe(result), ensure_ascii=False, indent=2) + "\n" - ) - print( - json.dumps({key: result[key] for key in ("dump_count", "rank", "step_range")}) - ) - - -if __name__ == "__main__": - main() diff --git a/benchmarks/benchmark_dspark_api.py b/benchmarks/benchmark_dspark_api.py index a4cfa62140..0279e4e896 100644 --- a/benchmarks/benchmark_dspark_api.py +++ b/benchmarks/benchmark_dspark_api.py @@ -1,5 +1,4 @@ # 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 d40f1f237f..16e7de9ae8 100644 --- a/benchmarks/benchmark_dspark_stream.py +++ b/benchmarks/benchmark_dspark_stream.py @@ -1,5 +1,4 @@ # 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_dsv4_dspark_long_context.py b/benchmarks/benchmark_dsv4_dspark_long_context.py deleted file mode 100644 index d0a1287046..0000000000 --- a/benchmarks/benchmark_dsv4_dspark_long_context.py +++ /dev/null @@ -1,576 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Measure DSpark verification and output health across long contexts. - -The benchmark targets an already-running OpenAI-compatible server. It records -stream timestamps together with vLLM speculative metrics so decode speed, -verification-round cost, and draft acceptance can be compared per context -bucket instead of being conflated with prefill latency. -""" - -from __future__ import annotations - -import argparse -import http.client -import json -import statistics -import time -from pathlib import Path -from typing import Any -from urllib.parse import urlparse - -import regex as re -from transformers import AutoTokenizer - -_PROFILE_INTERVAL_MARKER = "SM70 spec runner profile interval_avg_ms " -_PROFILE_FIELD_PATTERN = re.compile(r"([a-z_]+)=([-+]?\d+(?:\.\d+)?(?:[eE][-+]?\d+)?)") - - -def _parse_lengths(raw: str) -> list[int]: - values = sorted({int(value.strip()) for value in raw.split(",") if value.strip()}) - if not values or values[0] <= 0: - raise ValueError("--lengths must contain positive integers") - return values - - -def _template_token_ids(tokenizer: Any, messages: list[dict[str, str]]) -> list[int]: - try: - encoded = tokenizer.apply_chat_template( - messages, - tokenize=True, - add_generation_prompt=True, - enable_thinking=False, - ) - except (TypeError, ValueError): - # DeepSeek V4's custom tokenizer mode supplies the template in the - # server rather than tokenizer_config.json. The constant allowance is - # only for fitting below the target; server-reported usage remains the - # authoritative prompt length in every result row. - encoded = [0] * 32 - for message in messages: - encoded.extend( - tokenizer.encode(message["content"], add_special_tokens=False) - ) - if isinstance(encoded, dict): - encoded = encoded["input_ids"] - elif hasattr(encoded, "input_ids"): - encoded = encoded.input_ids - if encoded and isinstance(encoded[0], list): - encoded = encoded[0] - return list(encoded) - - -def _filler_unit(index: int) -> str: - return ( - f"Record {index:06d}: unrelated inventory values " - f"{index * 17 % 9973}, {index * 31 % 7919}, {index * 43 % 6151}. " - "This record is only deterministic long-context filler and contains " - "no answer to the final instruction.\n" - ) - - -def _build_messages(unit_count: int, marker: str) -> list[dict[str, str]]: - filler = "".join(_filler_unit(index) for index in range(unit_count)) - instruction = ( - "\nFINAL INSTRUCTION: Ignore the inventory records. Begin the response " - f"with the exact marker {marker} on its own line, then write a concise " - "technical paragraph explaining why bounded CUDA-graph context buckets " - "reduce long-context decode latency. Do not repeat the marker." - ) - return [ - { - "role": "system", - "content": ( - "Follow the final user instruction exactly. Long repeated context " - "must not cause repetition, garbling, or loss of the suffix." - ), - }, - {"role": "user", "content": filler + instruction}, - ] - - -def _fit_messages( - tokenizer: Any, - target_tokens: int, - marker: str, -) -> tuple[list[dict[str, str]], int, int]: - sample_count = 64 - sample_tokens = len( - _template_token_ids(tokenizer, _build_messages(sample_count, marker)) - ) - base_tokens = len(_template_token_ids(tokenizer, _build_messages(0, marker))) - tokens_per_unit = max(1.0, (sample_tokens - base_tokens) / sample_count) - low = 0 - high = max(1, int(target_tokens / tokens_per_unit * 1.2) + 32) - best_messages = _build_messages(0, marker) - best_tokens = len(_template_token_ids(tokenizer, best_messages)) - best_units = 0 - for _ in range(24): - if low > high: - break - middle = (low + high) // 2 - messages = _build_messages(middle, marker) - prompt_tokens = len(_template_token_ids(tokenizer, messages)) - if prompt_tokens <= target_tokens: - best_messages = messages - best_tokens = prompt_tokens - best_units = middle - low = middle + 1 - else: - high = middle - 1 - return best_messages, best_tokens, best_units - - -def _metric_total(text: str, name: str) -> float: - return sum( - float(line.rsplit(" ", 1)[-1]) - for line in text.splitlines() - if line.startswith(name + "{") or line.startswith(name + " ") - ) - - -def _metrics_snapshot(host: str, port: int, timeout: int) -> dict[str, Any]: - connection = http.client.HTTPConnection(host, port, timeout=timeout) - try: - connection.request("GET", "/metrics") - response = connection.getresponse() - raw = response.read() - finally: - connection.close() - if response.status != 200: - raise RuntimeError(f"metrics endpoint returned HTTP {response.status}") - text = raw.decode("utf-8", errors="replace") - positions: dict[int, float] = {} - prefix = "vllm:spec_decode_num_accepted_tokens_per_pos_total{" - for line in text.splitlines(): - if not line.startswith(prefix): - continue - match = re.search(r'position="(\d+)"', line) - if match: - position = int(match.group(1)) - positions[position] = positions.get(position, 0.0) + float( - line.rsplit(" ", 1)[-1] - ) - return { - "rounds": _metric_total(text, "vllm:spec_decode_num_drafts_total"), - "proposed": _metric_total(text, "vllm:spec_decode_num_draft_tokens_total"), - "accepted": _metric_total(text, "vllm:spec_decode_num_accepted_tokens_total"), - "positions": positions, - } - - -def _longest_same_token_run(token_ids: list[int]) -> int: - longest = 0 - current = 0 - previous: int | None = None - for token_id in token_ids: - if token_id == previous: - current += 1 - else: - previous = token_id - current = 1 - longest = max(longest, current) - return longest - - -def _profile_log_offset(profile_log: Path | None) -> int | None: - if profile_log is None: - return None - try: - return profile_log.stat().st_size - except FileNotFoundError: - return 0 - - -def _read_verifier_profile( - profile_log: Path | None, - offset: int | None, - verifier_rows: int, -) -> dict[str, Any]: - if profile_log is None or offset is None: - return {"available": False, "reason": "profile log was not requested"} - try: - with profile_log.open("rb") as profile_file: - profile_file.seek(offset) - text = profile_file.read().decode("utf-8", errors="replace") - except FileNotFoundError: - return {"available": False, "reason": "profile log does not exist"} - - intervals: list[dict[str, Any]] = [] - for line in text.splitlines(): - if _PROFILE_INTERVAL_MARKER not in line: - continue - fields = { - key: float(value) - for key, value in _PROFILE_FIELD_PATTERN.findall( - line.split(_PROFILE_INTERVAL_MARKER, 1)[1] - ) - } - if int(fields.get("num_tokens", -1)) != verifier_rows: - continue - required = ("target_forward", "target_logits", "target_rejection_sample") - if any(name not in fields for name in required): - continue - interval = dict(fields) - for name in ( - "calls", - "interval_calls", - "interval_spec_steps", - "num_tokens", - "num_reqs", - ): - if name in interval: - interval[name] = int(interval[name]) - interval["complete_verifier_ms"] = sum(fields[name] for name in required) - intervals.append(interval) - - # The first/last profiler window can straddle prefill or request teardown. - # Long sweeps run with a short profile interval and retain only complete - # interior windows when enough samples are available. - if len(intervals) >= 4: - steady_intervals = intervals[1:-1] - elif len(intervals) == 3: - steady_intervals = intervals[1:] - else: - steady_intervals = intervals - median_names = ( - "target_forward", - "target_logits", - "target_rejection_sample", - "complete_verifier_ms", - "draft_total", - ) - medians = { - name: statistics.median( - float(interval[name]) for interval in steady_intervals if name in interval - ) - for name in median_names - if any(name in interval for interval in steady_intervals) - } - return { - "available": bool(intervals), - "profile_log": str(profile_log), - "verifier_rows": verifier_rows, - "raw_interval_count": len(intervals), - "steady_interval_count": len(steady_intervals), - "medians_ms": medians, - "intervals": intervals, - } - - -def _stream_request( - *, - host: str, - port: int, - model: str, - messages: list[dict[str, str]], - output_tokens: int, - temperature: float, - top_p: float, - seed: int, - timeout: int, -) -> dict[str, Any]: - payload = { - "model": model, - "messages": messages, - "max_tokens": output_tokens, - "temperature": temperature, - "top_p": top_p, - "seed": seed, - "ignore_eos": True, - "stream": True, - "stream_options": {"include_usage": True}, - "chat_template_kwargs": {"enable_thinking": False}, - } - connection = http.client.HTTPConnection(host, port, timeout=timeout) - encoded = json.dumps(payload, ensure_ascii=False).encode("utf-8") - started_ns = time.perf_counter_ns() - first_token_ns: int | None = None - last_token_ns: int | None = None - content_parts: list[str] = [] - reasoning_parts: list[str] = [] - usage: dict[str, Any] = {} - finish_reason: str | None = None - try: - connection.request( - "POST", - "/v1/chat/completions", - body=encoded, - headers={"Content-Type": "application/json"}, - ) - response = connection.getresponse() - if response.status != 200: - raw = response.read() - raise RuntimeError( - f"request returned HTTP {response.status}: " - f"{raw.decode('utf-8', errors='replace')}" - ) - for raw_line in response: - arrival_ns = time.perf_counter_ns() - line = raw_line.decode("utf-8", errors="replace").strip() - if not line.startswith("data: "): - continue - data = line[6:] - if data == "[DONE]": - break - event = json.loads(data) - if event.get("usage"): - usage = event["usage"] - choices = event.get("choices") or [] - if not choices: - continue - choice = choices[0] - if choice.get("finish_reason") is not None: - finish_reason = choice["finish_reason"] - delta = choice.get("delta") or {} - content = delta.get("content") or "" - reasoning = delta.get("reasoning_content") or delta.get("reasoning") or "" - if content or reasoning: - first_token_ns = first_token_ns or arrival_ns - last_token_ns = arrival_ns - if content: - content_parts.append(content) - if reasoning: - reasoning_parts.append(reasoning) - finally: - finished_ns = time.perf_counter_ns() - connection.close() - - completion_tokens = int(usage.get("completion_tokens") or 0) - decode_seconds = ( - (last_token_ns - first_token_ns) / 1_000_000_000 - if first_token_ns is not None and last_token_ns is not None - else None - ) - return { - "request": payload, - "usage": usage, - "finish_reason": finish_reason, - "wall_seconds": (finished_ns - started_ns) / 1_000_000_000, - "ttft_seconds": ( - (first_token_ns - started_ns) / 1_000_000_000 - if first_token_ns is not None - else None - ), - "pure_decode_seconds": decode_seconds, - "pure_decode_tokens_per_second": ( - (completion_tokens - 1) / decode_seconds - if completion_tokens > 1 and decode_seconds - else None - ), - "content": "".join(content_parts), - "reasoning": "".join(reasoning_parts), - } - - -def _run_case( - *, - tokenizer: Any, - host: str, - port: int, - model: str, - target_tokens: int, - output_tokens: int, - temperature: float, - top_p: float, - seed: int, - timeout: int, - num_speculative_tokens: int, - profile_log: Path | None, -) -> dict[str, Any]: - marker = f"DSV4-CONTEXT-{target_tokens}-OK" - messages, local_prompt_tokens, unit_count = _fit_messages( - tokenizer, target_tokens, marker - ) - profile_offset = _profile_log_offset(profile_log) - before = _metrics_snapshot(host, port, timeout) - result = _stream_request( - host=host, - port=port, - model=model, - messages=messages, - output_tokens=output_tokens, - temperature=temperature, - top_p=top_p, - seed=seed, - timeout=timeout, - ) - after = _metrics_snapshot(host, port, timeout) - rounds = after["rounds"] - before["rounds"] - proposed = after["proposed"] - before["proposed"] - accepted = after["accepted"] - before["accepted"] - accepted_per_position = [ - after["positions"].get(position, 0.0) - before["positions"].get(position, 0.0) - for position in range(num_speculative_tokens) - ] - verifier_profile = ( - _read_verifier_profile( - profile_log, - profile_offset, - num_speculative_tokens + 1, - ) - if num_speculative_tokens - else {"available": False, "reason": "speculative decoding is disabled"} - ) - output_text = f"{result['reasoning']}\n{result['content']}" - output_ids = tokenizer.encode(output_text, add_special_tokens=False) - completion_tokens = int(result["usage"].get("completion_tokens") or 0) - result.update( - { - "target_prompt_tokens": target_tokens, - "local_prompt_tokens": local_prompt_tokens, - "unit_count": unit_count, - "marker": marker, - "output_health": { - "marker_in_content": marker in result["content"], - "replacement_characters": output_text.count("\ufffd"), - "unique_output_tokens": len(set(output_ids)), - "longest_same_token_run": _longest_same_token_run(output_ids), - }, - "speculative_metrics": { - "rounds": rounds, - "proposed_tokens": proposed, - "accepted_tokens": accepted, - "accepted_fraction_of_proposed": ( - accepted / proposed if proposed else None - ), - "mean_emitted_tokens_per_round": ( - 1.0 + accepted / rounds if rounds else None - ), - "round_wall_milliseconds": ( - result["pure_decode_seconds"] * 1000.0 / rounds - if result["pure_decode_seconds"] and rounds - else None - ), - "unconditional_acceptance_per_position": [ - value / rounds if rounds else None - for value in accepted_per_position - ], - }, - "verifier_profile": verifier_profile, - } - ) - result["passed"] = bool( - completion_tokens == output_tokens - and result["output_health"]["marker_in_content"] - and result["output_health"]["replacement_characters"] == 0 - and result["output_health"]["unique_output_tokens"] >= 8 - and result["output_health"]["longest_same_token_run"] < 16 - ) - return result - - -def _summary(records: list[dict[str, Any]]) -> dict[str, Any]: - speeds = [ - float(record["pure_decode_tokens_per_second"]) - for record in records - if record["pure_decode_tokens_per_second"] is not None - ] - round_costs = [ - float(record["speculative_metrics"]["round_wall_milliseconds"]) - for record in records - if record["speculative_metrics"]["round_wall_milliseconds"] is not None - ] - verifier_costs = [ - float(record["verifier_profile"]["medians_ms"]["complete_verifier_ms"]) - for record in records - if record["verifier_profile"].get("medians_ms", {}).get("complete_verifier_ms") - is not None - ] - return { - "cases": len(records), - "passed": all(record["passed"] for record in records), - "pure_decode_tps_min": min(speeds) if speeds else None, - "pure_decode_tps_median": statistics.median(speeds) if speeds else None, - "round_wall_ms_max": max(round_costs) if round_costs else None, - "complete_verifier_ms_median": ( - statistics.median(verifier_costs) if verifier_costs else None - ), - "complete_verifier_ms_max": max(verifier_costs) if verifier_costs else None, - } - - -def main() -> int: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--base-url", default="http://127.0.0.1:8000") - parser.add_argument("--model", required=True) - parser.add_argument("--model-path", type=Path, required=True) - parser.add_argument("--lengths", default="1024,4096,16384,65536,131072,252000") - parser.add_argument("--output-tokens", type=int, default=128) - parser.add_argument("--temperature", type=float, default=1.0) - parser.add_argument("--top-p", type=float, default=1.0) - parser.add_argument("--seed", type=int, default=20260824) - parser.add_argument("--num-speculative-tokens", type=int, default=7) - parser.add_argument("--profile-log", type=Path) - parser.add_argument("--timeout", type=int, default=1800) - parser.add_argument("--out", type=Path, required=True) - args = parser.parse_args() - - if not 0 <= args.num_speculative_tokens <= 7: - raise ValueError("--num-speculative-tokens must be in [0, 7]") - - parsed = urlparse(args.base_url) - if parsed.scheme != "http" or not parsed.hostname: - raise ValueError("--base-url must be an http URL") - host = parsed.hostname - port = parsed.port or 80 - tokenizer = AutoTokenizer.from_pretrained( - str(args.model_path), - local_files_only=True, - trust_remote_code=True, - ) - - args.out.parent.mkdir(parents=True, exist_ok=True) - records: list[dict[str, Any]] = [] - for target_tokens in _parse_lengths(args.lengths): - record = _run_case( - tokenizer=tokenizer, - host=host, - port=port, - model=args.model, - target_tokens=target_tokens, - output_tokens=args.output_tokens, - temperature=args.temperature, - top_p=args.top_p, - seed=args.seed, - timeout=args.timeout, - num_speculative_tokens=args.num_speculative_tokens, - profile_log=args.profile_log, - ) - records.append(record) - contract = { - key: str(value) if isinstance(value, Path) else value - for key, value in vars(args).items() - } - payload = { - "contract": contract, - "records": records, - "summary": _summary(records), - } - args.out.write_text( - json.dumps(payload, ensure_ascii=False, indent=2) + "\n", - encoding="utf-8", - ) - print( - json.dumps( - { - "target_prompt_tokens": target_tokens, - "usage": record["usage"], - "pure_decode_tokens_per_second": record[ - "pure_decode_tokens_per_second" - ], - "speculative_metrics": record["speculative_metrics"], - "verifier_profile": record["verifier_profile"], - "output_health": record["output_health"], - "passed": record["passed"], - }, - ensure_ascii=False, - ), - flush=True, - ) - - return 0 if all(record["passed"] for record in records) else 1 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/benchmarks/benchmark_dsv4_gsm8k_api.py b/benchmarks/benchmark_dsv4_gsm8k_api.py deleted file mode 100644 index 49f19a22b7..0000000000 --- a/benchmarks/benchmark_dsv4_gsm8k_api.py +++ /dev/null @@ -1,231 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Run a deterministic, sequential GSM8K gate against a vLLM API server.""" - -from __future__ import annotations - -import argparse -import hashlib -import http.client -import json -import time -from decimal import Decimal, InvalidOperation -from pathlib import Path -from typing import Any -from urllib.parse import urlparse - -import regex as re - -_NUMBER_RE = re.compile(r"(?"] - - -def _sha256(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as source: - for chunk in iter(lambda: source.read(1 << 20), b""): - digest.update(chunk) - return digest.hexdigest() - - -def _load_jsonl(path: Path) -> list[dict[str, Any]]: - return [ - json.loads(line) - for line in path.read_text(encoding="utf-8").splitlines() - if line and not line.startswith("#") - ] - - -def _last_integer(text: str) -> int | None: - numbers = _NUMBER_RE.findall(text) - if not numbers: - return None - try: - value = Decimal(numbers[-1].replace(",", "")) - except InvalidOperation: - return None - if not value.is_finite() or value != value.to_integral_value(): - return None - return int(value) - - -def _post_completion( - *, - host: str, - port: int, - timeout: int, - payload: dict[str, Any], -) -> tuple[int, dict[str, Any], float]: - connection = http.client.HTTPConnection(host, port, timeout=timeout) - started = time.perf_counter() - try: - connection.request( - "POST", - "/v1/completions", - body=json.dumps(payload).encode("utf-8"), - headers={"Content-Type": "application/json"}, - ) - response = connection.getresponse() - body = response.read() - finally: - connection.close() - elapsed = time.perf_counter() - started - try: - decoded = json.loads(body) - except json.JSONDecodeError as exc: - raise RuntimeError( - f"completion returned non-JSON HTTP {response.status}: " - f"{body.decode('utf-8', errors='replace')}" - ) from exc - return response.status, decoded, elapsed - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--base-url", default="http://127.0.0.1:8000") - parser.add_argument("--model", required=True) - parser.add_argument("--train", type=Path, required=True) - parser.add_argument("--test", type=Path, required=True) - parser.add_argument("--out", type=Path, required=True) - parser.add_argument("--limit", type=int, default=64) - parser.add_argument("--few-shot", type=int, default=5) - parser.add_argument("--max-tokens", type=int, default=256) - parser.add_argument("--seed", type=int, default=42) - parser.add_argument("--timeout", type=int, default=1200) - parser.add_argument("--min-correct", type=int, default=0) - parser.add_argument("--max-invalid", type=int, default=0) - return parser.parse_args() - - -def main() -> int: - args = parse_args() - if args.limit <= 0 or args.few_shot < 0: - raise ValueError("--limit must be positive and --few-shot non-negative") - if args.min_correct < 0 or args.max_invalid < 0: - raise ValueError("quality thresholds must be non-negative") - - parsed = urlparse(args.base_url) - if parsed.scheme != "http" or not parsed.hostname: - raise ValueError("--base-url must be an http URL") - host = parsed.hostname - port = parsed.port or 80 - - train = _load_jsonl(args.train) - test = _load_jsonl(args.test)[: args.limit] - if len(train) < args.few_shot: - raise RuntimeError( - f"training set has {len(train)} rows, fewer than {args.few_shot} shots" - ) - if len(test) < args.limit: - raise RuntimeError(f"test set has only {len(test)} rows, need {args.limit}") - - few_shot = "".join( - f"Question: {row['question']}\nAnswer: {row['answer']}\n\n" - for row in train[: args.few_shot] - ) - rows: list[dict[str, Any]] = [] - started = time.perf_counter() - for index, row in enumerate(test): - prompt = few_shot + f"Question: {row['question']}\nAnswer:" - request = { - "model": args.model, - "prompt": prompt, - "temperature": 0.0, - "top_p": 1.0, - "max_tokens": args.max_tokens, - "stop": _STOP_SEQUENCES, - "seed": args.seed, - } - status, response, elapsed = _post_completion( - host=host, - port=port, - timeout=args.timeout, - payload=request, - ) - choices = response.get("choices") or [] - choice = choices[0] if choices else {} - output = choice.get("text") or "" - expected = _last_integer(row["answer"]) - if expected is None: - raise ValueError(f"GSM8K item {index} has no integral reference answer") - predicted = _last_integer(output) - rows.append( - { - "index": index, - "question": row["question"], - "expected": expected, - "predicted": predicted, - "correct": predicted is not None and predicted == expected, - "invalid": predicted is None, - "status": status, - "elapsed_seconds": elapsed, - "finish_reason": choice.get("finish_reason"), - "usage": response.get("usage"), - "output": output, - } - ) - if status != 200: - raise RuntimeError(f"GSM8K item {index} returned HTTP {status}: {response}") - - wall_seconds = time.perf_counter() - started - correct = sum(bool(row["correct"]) for row in rows) - invalid = sum(bool(row["invalid"]) for row in rows) - completion_tokens = sum( - int((row.get("usage") or {}).get("completion_tokens") or 0) for row in rows - ) - passed = correct >= args.min_correct and invalid <= args.max_invalid - result = { - "contract": { - "base_url": args.base_url, - "model": args.model, - "questions": args.limit, - "few_shot": args.few_shot, - "temperature": 0.0, - "top_p": 1.0, - "seed": args.seed, - "max_tokens": args.max_tokens, - "strictly_sequential": True, - "train_selection": "first_n", - "test_selection": "first_n", - "prompt_format": "gsm8k_question_answer_v1", - "answer_normalization": "last_signed_integral_decimal_v1", - "stop_sequences": _STOP_SEQUENCES, - "min_correct": args.min_correct, - "max_invalid": args.max_invalid, - }, - "input_manifest": { - "train": {"path": str(args.train), "sha256": _sha256(args.train)}, - "test": {"path": str(args.test), "sha256": _sha256(args.test)}, - }, - "correct": correct, - "accuracy": correct / len(rows), - "invalid": invalid, - "wall_seconds": wall_seconds, - "completion_tokens": completion_tokens, - "aggregate_output_tokens_per_second": ( - completion_tokens / wall_seconds if wall_seconds else None - ), - "passed": passed, - "rows": rows, - } - args.out.parent.mkdir(parents=True, exist_ok=True) - args.out.write_text( - json.dumps(result, ensure_ascii=False, indent=2) + "\n", - encoding="utf-8", - ) - print( - json.dumps( - { - "correct": correct, - "samples": len(rows), - "invalid": invalid, - "passed": passed, - } - ), - flush=True, - ) - return 0 if passed else 1 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/benchmarks/benchmark_dsv4_quality_api.py b/benchmarks/benchmark_dsv4_quality_api.py deleted file mode 100644 index 5eba6b30ea..0000000000 --- a/benchmarks/benchmark_dsv4_quality_api.py +++ /dev/null @@ -1,635 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Run paired HumanEval and LongBench-subset gates against a vLLM API.""" - -from __future__ import annotations - -import argparse -import ctypes -import functools -import hashlib -import http.client -import json -import os -import resource -import statistics -import subprocess -import sys -import tempfile -from pathlib import Path -from typing import Any -from urllib.parse import urlparse - -import regex as re -from transformers import AutoTokenizer - -NO_CHAT_DATASETS = {"trec", "triviaqa", "samsum", "lsht", "lcc", "repobench-p"} - - -def _sha256(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as source: - for chunk in iter(lambda: source.read(1 << 20), b""): - digest.update(chunk) - return digest.hexdigest() - - -def _post_json( - *, - host: str, - port: int, - path: str, - payload: dict[str, Any], - timeout: int, -) -> dict[str, Any]: - connection = http.client.HTTPConnection(host, port, timeout=timeout) - try: - connection.request( - "POST", - path, - body=json.dumps(payload, ensure_ascii=False).encode("utf-8"), - headers={"Content-Type": "application/json"}, - ) - response = connection.getresponse() - raw = response.read() - finally: - connection.close() - if response.status != 200: - raise RuntimeError( - f"{path} returned HTTP {response.status}: " - f"{raw.decode('utf-8', errors='replace')}" - ) - return json.loads(raw) - - -def _chat( - *, - host: str, - port: int, - model: str, - content: str, - max_tokens: int, - timeout: int, -) -> tuple[str, dict[str, Any]]: - result = _post_json( - host=host, - port=port, - path="/v1/chat/completions", - timeout=timeout, - payload={ - "model": model, - "messages": [{"role": "user", "content": content}], - "max_tokens": max_tokens, - "temperature": 0.0, - "top_p": 1.0, - "seed": 20260824, - "chat_template_kwargs": {"enable_thinking": False}, - }, - ) - message = result["choices"][0]["message"] - return message.get("content") or "", result - - -def _completion( - *, - host: str, - port: int, - model: str, - prompt: str, - max_tokens: int, - timeout: int, -) -> tuple[str, dict[str, Any]]: - result = _post_json( - host=host, - port=port, - path="/v1/completions", - timeout=timeout, - payload={ - "model": model, - "prompt": prompt, - "max_tokens": max_tokens, - "temperature": 0.0, - "top_p": 1.0, - "seed": 20260824, - }, - ) - return result["choices"][0].get("text") or "", result - - -def _extract_python(text: str, entry_point: str) -> str: - blocks = re.findall(r"```(?:python)?\s*\n(.*?)```", text, flags=re.DOTALL) - for block in blocks: - if re.search(rf"\bdef\s+{re.escape(entry_point)}\s*\(", block): - return block.strip() - if blocks: - return max(blocks, key=len).strip() - return text.strip() - - -def _sandbox_preexec(sandbox_dir: str) -> None: - """Deny filesystem/network access before executing generated code.""" - libc = ctypes.CDLL(None, use_errno=True) - libc.syscall.restype = ctypes.c_long - libc.prctl.argtypes = [ - ctypes.c_int, - ctypes.c_ulong, - ctypes.c_ulong, - ctypes.c_ulong, - ctypes.c_ulong, - ] - libc.prctl.restype = ctypes.c_int - sys_create_ruleset = 444 - sys_add_rule = 445 - sys_restrict_self = 446 - abi = libc.syscall(sys_create_ruleset, 0, 0, 1) - if abi < 1: - raise OSError(ctypes.get_errno(), "Landlock is unavailable") - - class RulesetAttr(ctypes.Structure): - _fields_ = [ - ("handled_access_fs", ctypes.c_uint64), - ("handled_access_net", ctypes.c_uint64), - ] - - class PathBeneathAttr(ctypes.Structure): - _pack_ = 1 - _fields_ = [ - ("allowed_access", ctypes.c_uint64), - ("parent_fd", ctypes.c_int32), - ] - - handled_fs = (1 << 13) - 1 - if abi >= 2: - handled_fs |= 1 << 13 - if abi >= 3: - handled_fs |= 1 << 14 - ruleset_attr = RulesetAttr(handled_fs, 3 if abi >= 4 else 0) - ruleset_fd = libc.syscall( - sys_create_ruleset, - ctypes.byref(ruleset_attr), - ctypes.sizeof(ruleset_attr), - 0, - ) - if ruleset_fd < 0: - raise OSError(ctypes.get_errno(), "landlock_create_ruleset") - - read_only = (1 << 0) | (1 << 2) | (1 << 3) - allowed_paths = ( - ("/usr", read_only), - ("/lib", read_only), - ("/lib64", read_only), - ("/dev/null", (1 << 1) | (1 << 2)), - ("/dev/urandom", 1 << 2), - (sandbox_dir, handled_fs), - ) - for path, allowed_access in allowed_paths: - parent_fd = os.open(path, os.O_PATH | os.O_CLOEXEC) - path_attr = PathBeneathAttr(allowed_access, parent_fd) - if ( - libc.syscall( - sys_add_rule, - ruleset_fd, - 1, - ctypes.byref(path_attr), - 0, - ) - < 0 - ): - raise OSError(ctypes.get_errno(), f"landlock_add_rule({path})") - os.close(parent_fd) - - if libc.prctl(38, 1, 0, 0, 0) < 0: - raise OSError(ctypes.get_errno(), "PR_SET_NO_NEW_PRIVS") - if libc.syscall(sys_restrict_self, ruleset_fd, 0) < 0: - raise OSError(ctypes.get_errno(), "landlock_restrict_self") - os.close(ruleset_fd) - - seccomp = ctypes.CDLL("libseccomp.so.2", use_errno=True) - seccomp.seccomp_init.argtypes = [ctypes.c_uint32] - seccomp.seccomp_init.restype = ctypes.c_void_p - seccomp.seccomp_syscall_resolve_name.argtypes = [ctypes.c_char_p] - seccomp.seccomp_syscall_resolve_name.restype = ctypes.c_int - seccomp.seccomp_rule_add.argtypes = [ - ctypes.c_void_p, - ctypes.c_uint32, - ctypes.c_int, - ctypes.c_uint, - ] - seccomp.seccomp_rule_add.restype = ctypes.c_int - seccomp.seccomp_load.argtypes = [ctypes.c_void_p] - seccomp.seccomp_load.restype = ctypes.c_int - seccomp.seccomp_release.argtypes = [ctypes.c_void_p] - seccomp.seccomp_release.restype = None - context = seccomp.seccomp_init(0x7FFF0000) - if not context: - raise RuntimeError("seccomp_init failed") - deny_with_eperm = 0x00050000 | 1 - for name in ( - b"socket", - b"socketpair", - b"connect", - b"bind", - b"listen", - b"accept", - b"accept4", - b"sendto", - b"sendmsg", - b"recvfrom", - b"recvmsg", - b"kill", - b"tkill", - b"tgkill", - b"pidfd_open", - b"pidfd_getfd", - b"pidfd_send_signal", - b"ptrace", - b"process_vm_readv", - b"process_vm_writev", - b"kcmp", - b"fork", - b"vfork", - b"clone", - b"clone3", - b"unshare", - b"setns", - b"mount", - b"umount2", - b"pivot_root", - b"bpf", - b"perf_event_open", - b"userfaultfd", - b"keyctl", - b"add_key", - b"request_key", - b"open_by_handle_at", - b"name_to_handle_at", - ): - syscall_number = seccomp.seccomp_syscall_resolve_name(name) - if ( - syscall_number >= 0 - and seccomp.seccomp_rule_add(context, deny_with_eperm, syscall_number, 0) - != 0 - ): - raise RuntimeError(f"seccomp_rule_add failed for {name!r}") - if seccomp.seccomp_load(context) != 0: - raise RuntimeError("seccomp_load failed") - seccomp.seccomp_release(context) - - resource.setrlimit(resource.RLIMIT_CPU, (5, 5)) - resource.setrlimit(resource.RLIMIT_AS, (256 << 20, 256 << 20)) - resource.setrlimit(resource.RLIMIT_FSIZE, (1 << 20, 1 << 20)) - resource.setrlimit(resource.RLIMIT_NPROC, (16, 16)) - - -def _execute_humaneval(program: str, timeout: int = 8) -> dict[str, Any]: - with tempfile.TemporaryDirectory(prefix="dsv4-humaneval-") as sandbox_dir: - output_path = Path(sandbox_dir) / "output.log" - try: - with output_path.open("wb") as output_file: - completed = subprocess.run( - ["/usr/bin/python3", "-I", "-S", "-"], - input=program.encode("utf-8"), - stdout=output_file, - stderr=subprocess.STDOUT, - cwd=sandbox_dir, - env={"PATH": "/usr/bin", "HOME": sandbox_dir}, - preexec_fn=functools.partial(_sandbox_preexec, sandbox_dir), - timeout=timeout, - ) - output = output_path.read_bytes()[-2000:].decode("utf-8", errors="replace") - return { - "passed": completed.returncode == 0, - "returncode": completed.returncode, - "output": output, - } - except subprocess.TimeoutExpired as exc: - return { - "passed": False, - "returncode": None, - "output": f"timeout: {exc}", - } - - -def _run_humaneval( - *, - host: str, - port: int, - model: str, - tasks_path: Path, - limit: int, - timeout: int, -) -> dict[str, Any]: - tasks = json.loads(tasks_path.read_text(encoding="utf-8"))[:limit] - records: list[dict[str, Any]] = [] - for task in tasks: - prompt = ( - "Complete the Python function below. Return runnable Python code " - "only, without Markdown or explanation.\n\n" + task["prompt"] - ) - response, raw = _chat( - host=host, - port=port, - model=model, - content=prompt, - max_tokens=768, - timeout=timeout, - ) - code = _extract_python(response, task["entry_point"]) - if re.search(rf"\bdef\s+{re.escape(task['entry_point'])}\s*\(", code): - candidate = code - else: - candidate = task["prompt"] + code - program = candidate + "\n" + task["test"] + f"\ncheck({task['entry_point']})\n" - execution = _execute_humaneval(program) - records.append( - { - "task_id": task["task_id"], - "passed": execution["passed"], - "execution": execution, - "response": response, - "usage": raw.get("usage"), - } - ) - passed = sum(int(record["passed"]) for record in records) - return { - "samples": len(records), - "passed": passed, - "pass_at_1": passed / len(records) if records else None, - "records": records, - } - - -def _truncate_middle(tokenizer: Any, prompt: str, max_input_tokens: int) -> str: - token_ids = tokenizer.encode(prompt, add_special_tokens=False) - if len(token_ids) <= max_input_tokens: - return prompt - half = max_input_tokens // 2 - return tokenizer.decode( - token_ids[:half], skip_special_tokens=True - ) + tokenizer.decode(token_ids[-half:], skip_special_tokens=True) - - -def _select_longbench_rows( - rows: list[dict[str, Any]], limit: int -) -> list[dict[str, Any]]: - rows = [row for row in rows if int(row.get("length", 0)) >= 8000] or rows - if len(rows) <= limit: - return rows - if limit == 1: - return [rows[len(rows) // 2]] - last = len(rows) - 1 - return [rows[round(index * last / (limit - 1))] for index in range(limit)] - - -def _run_longbench( - *, - tokenizer: Any, - host: str, - port: int, - model: str, - data_dir: Path, - longbench_root: Path, - datasets: list[str], - limit: int, - max_input_tokens: int, - timeout: int, -) -> dict[str, Any]: - config_dir = longbench_root / "config" - prompt_formats = json.loads( - (config_dir / "dataset2prompt.json").read_text(encoding="utf-8") - ) - max_output_tokens = json.loads( - (config_dir / "dataset2maxlen.json").read_text(encoding="utf-8") - ) - sys.path.insert(0, str(longbench_root.resolve())) - try: - from eval import dataset2metric # type: ignore - except ModuleNotFoundError as exc: - raise RuntimeError( - "LongBench metric dependencies are missing; install rouge, jieba, " - "and fuzzywuzzy into PYTHONPATH" - ) from exc - - by_dataset: dict[str, Any] = {} - for dataset in datasets: - rows = [] - for source_index, line in enumerate( - (data_dir / f"{dataset}.jsonl").read_text(encoding="utf-8").splitlines() - ): - if not line.strip(): - continue - row = json.loads(line) - row["_source_index"] = source_index - rows.append(row) - selected = _select_longbench_rows(rows, limit) - records: list[dict[str, Any]] = [] - for row in selected: - prompt = prompt_formats[dataset].format(**row) - prompt = _truncate_middle(tokenizer, prompt, max_input_tokens) - if dataset in NO_CHAT_DATASETS: - prediction, raw = _completion( - host=host, - port=port, - model=model, - prompt=prompt, - max_tokens=int(max_output_tokens[dataset]), - timeout=timeout, - ) - else: - prediction, raw = _chat( - host=host, - port=port, - model=model, - content=prompt, - max_tokens=int(max_output_tokens[dataset]), - timeout=timeout, - ) - score = max( - float( - dataset2metric[dataset]( - prediction, - answer, - all_classes=row.get("all_classes", []), - ) - ) - for answer in row["answers"] - ) - records.append( - { - "source_index": row["_source_index"], - "length": row.get("length"), - "score": score, - "prediction": prediction, - "answers": row["answers"], - "usage": raw.get("usage"), - } - ) - scores = [float(record["score"]) for record in records] - by_dataset[dataset] = { - "samples": len(records), - "score": 100.0 * statistics.mean(scores) if scores else None, - "records": records, - } - dataset_scores = [ - float(result["score"]) - for result in by_dataset.values() - if result["score"] is not None - ] - return { - "datasets": by_dataset, - "average_score": statistics.mean(dataset_scores) if dataset_scores else None, - } - - -def main() -> int: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--base-url", default="http://127.0.0.1:8000") - parser.add_argument("--model", required=True) - parser.add_argument("--model-path", type=Path, required=True) - parser.add_argument("--out", type=Path, required=True) - parser.add_argument( - "--humaneval-tasks", - type=Path, - default=Path( - "bench_results/ddtree_humaneval_sweep_20260625/humaneval_tasks.json" - ), - ) - parser.add_argument("--humaneval-limit", type=int, default=32) - parser.add_argument( - "--longbench-data-dir", - type=Path, - default=Path("benchmark-data/longbench/data"), - ) - parser.add_argument( - "--longbench-root", - type=Path, - default=Path("third_party/LongBench/LongBench"), - ) - parser.add_argument( - "--longbench-datasets", - default="hotpotqa,multifieldqa_zh,gov_report,lcc", - ) - parser.add_argument("--longbench-limit", type=int, default=4) - parser.add_argument("--longbench-max-input-tokens", type=int, default=32768) - parser.add_argument("--timeout", type=int, default=1200) - parser.add_argument("--min-humaneval-passes", type=int, default=0) - parser.add_argument("--min-longbench-average", type=float, default=0.0) - args = parser.parse_args() - - parsed = urlparse(args.base_url) - if parsed.scheme != "http" or not parsed.hostname: - raise ValueError("--base-url must be an http URL") - host = parsed.hostname - port = parsed.port or 80 - tokenizer = AutoTokenizer.from_pretrained( - str(args.model_path), - local_files_only=True, - trust_remote_code=True, - ) - - datasets = [ - value.strip() for value in args.longbench_datasets.split(",") if value.strip() - ] - input_manifest = { - "humaneval_tasks": { - "path": str(args.humaneval_tasks), - "sha256": _sha256(args.humaneval_tasks), - }, - "longbench_prompt_config": { - "path": str(args.longbench_root / "config/dataset2prompt.json"), - "sha256": _sha256(args.longbench_root / "config/dataset2prompt.json"), - }, - "longbench_maxlen_config": { - "path": str(args.longbench_root / "config/dataset2maxlen.json"), - "sha256": _sha256(args.longbench_root / "config/dataset2maxlen.json"), - }, - "longbench_datasets": { - dataset: { - "path": str(args.longbench_data_dir / f"{dataset}.jsonl"), - "sha256": _sha256(args.longbench_data_dir / f"{dataset}.jsonl"), - } - for dataset in datasets - }, - } - - human_eval = _run_humaneval( - host=host, - port=port, - model=args.model, - tasks_path=args.humaneval_tasks, - limit=args.humaneval_limit, - timeout=args.timeout, - ) - longbench = _run_longbench( - tokenizer=tokenizer, - host=host, - port=port, - model=args.model, - data_dir=args.longbench_data_dir, - longbench_root=args.longbench_root, - datasets=datasets, - limit=args.longbench_limit, - max_input_tokens=args.longbench_max_input_tokens, - timeout=args.timeout, - ) - passed = bool( - human_eval["passed"] >= args.min_humaneval_passes - and longbench["average_score"] is not None - and longbench["average_score"] >= args.min_longbench_average - ) - result = { - "contract": { - key: str(value) if isinstance(value, Path) else value - for key, value in vars(args).items() - }, - "evaluation_contract": { - "version": "dsv4_api_quality_v1", - "strictly_sequential": True, - "temperature": 0.0, - "top_p": 1.0, - "seed": 20260824, - "human_eval": { - "selection": "first_n", - "limit": args.humaneval_limit, - "endpoint": "chat_completions", - "enable_thinking": False, - "max_tokens": 768, - "execution": "landlock_seccomp_python_isolated_v1", - }, - "longbench": { - "datasets": datasets, - "selection": "length_ge_8000_evenly_spaced_v1", - "limit_per_dataset": args.longbench_limit, - "max_input_tokens": args.longbench_max_input_tokens, - "truncation": "middle_first_half_last_half_v1", - "endpoint_policy": "official_no_chat_dataset_set_v1", - "max_output_tokens": "dataset2maxlen_sha256_manifest", - }, - }, - "input_manifest": input_manifest, - "human_eval": human_eval, - "longbench": longbench, - "passed": passed, - } - args.out.parent.mkdir(parents=True, exist_ok=True) - args.out.write_text( - json.dumps(result, ensure_ascii=False, indent=2) + "\n", - encoding="utf-8", - ) - print( - json.dumps( - { - "human_eval_passes": human_eval["passed"], - "human_eval_samples": human_eval["samples"], - "longbench_average": longbench["average_score"], - "passed": passed, - }, - ensure_ascii=False, - ), - flush=True, - ) - return 0 if passed else 1 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/benchmarks/benchmark_sm70_cublas_gemm_algorithms.py b/benchmarks/benchmark_sm70_cublas_gemm_algorithms.py deleted file mode 100644 index f8c33ed8f1..0000000000 --- a/benchmarks/benchmark_sm70_cublas_gemm_algorithms.py +++ /dev/null @@ -1,251 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Screen legacy cuBLAS Tensor Core algorithms on exact SM70 QKV shapes.""" - -from __future__ import annotations - -import argparse -import ctypes -import json -import statistics -from collections.abc import Callable -from pathlib import Path -from typing import Any - -import torch - -_CUBLAS_STATUS_SUCCESS = 0 -_CUBLAS_OP_N = 0 -_CUDA_R_16F = 2 -_CUBLAS_COMPUTE_32F = 68 -_CUBLAS_TENSOR_OP_MATH = 1 -_CUBLAS_GEMM_DEFAULT_TENSOR_OP = 99 -_CUBLAS_GEMM_ALGO0_TENSOR_OP = 100 - - -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser() - parser.add_argument("--json-out", type=Path, required=True) - parser.add_argument("--device", default="cuda:0") - parser.add_argument("--m", type=int, default=8000) - parser.add_argument("--k", type=int, default=5120) - parser.add_argument("--n", type=int, nargs="+", default=[4096, 3584]) - parser.add_argument("--warmup", type=int, default=3) - parser.add_argument("--iters", type=int, default=10) - parser.add_argument("--trials", type=int, default=5) - parser.add_argument("--seed", type=int, default=20260823) - return parser.parse_args() - - -def _event_trials( - launch: Callable[[], None], warmup: int, iters: int, trials: int -) -> dict[str, Any]: - for _ in range(warmup): - launch() - torch.accelerator.synchronize() - samples: list[float] = [] - for _ in range(trials): - start = torch.Event(enable_timing=True) - end = torch.Event(enable_timing=True) - start.record() - for _ in range(iters): - launch() - end.record() - end.synchronize() - samples.append(float(start.elapsed_time(end) / iters)) - return { - "median_ms": float(statistics.median(samples)), - "mean_ms": float(statistics.mean(samples)), - "min_ms": min(samples), - "max_ms": max(samples), - "samples_ms": samples, - } - - -class _Cublas: - def __init__(self, stream: int) -> None: - self.lib = ctypes.CDLL("libcublas.so.12") - self.handle = ctypes.c_void_p() - self.lib.cublasCreate_v2.argtypes = [ctypes.POINTER(ctypes.c_void_p)] - self.lib.cublasCreate_v2.restype = ctypes.c_int - self.lib.cublasDestroy_v2.argtypes = [ctypes.c_void_p] - self.lib.cublasDestroy_v2.restype = ctypes.c_int - self.lib.cublasSetStream_v2.argtypes = [ctypes.c_void_p, ctypes.c_void_p] - self.lib.cublasSetStream_v2.restype = ctypes.c_int - self.lib.cublasSetMathMode.argtypes = [ctypes.c_void_p, ctypes.c_int] - self.lib.cublasSetMathMode.restype = ctypes.c_int - self.lib.cublasGemmEx.argtypes = [ - ctypes.c_void_p, - ctypes.c_int, - ctypes.c_int, - ctypes.c_int, - ctypes.c_int, - ctypes.c_int, - ctypes.c_void_p, - ctypes.c_void_p, - ctypes.c_int, - ctypes.c_int, - ctypes.c_void_p, - ctypes.c_int, - ctypes.c_int, - ctypes.c_void_p, - ctypes.c_void_p, - ctypes.c_int, - ctypes.c_int, - ctypes.c_int, - ctypes.c_int, - ] - self.lib.cublasGemmEx.restype = ctypes.c_int - self._check(self.lib.cublasCreate_v2(ctypes.byref(self.handle)), "create") - self._check( - self.lib.cublasSetStream_v2(self.handle, ctypes.c_void_p(stream)), - "set stream", - ) - self._check( - self.lib.cublasSetMathMode(self.handle, _CUBLAS_TENSOR_OP_MATH), - "set math mode", - ) - - @staticmethod - def _check(status: int, operation: str) -> None: - if status != _CUBLAS_STATUS_SUCCESS: - raise RuntimeError(f"cuBLAS {operation} failed with status {status}.") - - def gemm( - self, - output: torch.Tensor, - inputs: torch.Tensor, - weight: torch.Tensor, - algo: int, - ) -> int: - m, k = inputs.shape - weight_k, n = weight.shape - if weight_k != k or output.shape != (m, n): - raise ValueError("Incompatible row-major GEMM shapes.") - alpha = ctypes.c_float(1.0) - beta = ctypes.c_float(0.0) - # Row-major C=MxN is column-major C^T=NxM. Compute - # C^T = weight^T(NxK) * inputs^T(KxM) without materializing transposes. - return int( - self.lib.cublasGemmEx( - self.handle, - _CUBLAS_OP_N, - _CUBLAS_OP_N, - n, - m, - k, - ctypes.byref(alpha), - ctypes.c_void_p(weight.data_ptr()), - _CUDA_R_16F, - n, - ctypes.c_void_p(inputs.data_ptr()), - _CUDA_R_16F, - k, - ctypes.byref(beta), - ctypes.c_void_p(output.data_ptr()), - _CUDA_R_16F, - n, - _CUBLAS_COMPUTE_32F, - algo, - ) - ) - - def close(self) -> None: - if self.handle.value is not None: - self._check(self.lib.cublasDestroy_v2(self.handle), "destroy") - self.handle = ctypes.c_void_p() - - -def _run_shape( - args: argparse.Namespace, device: torch.device, cublas: _Cublas, n: int -) -> dict[str, Any]: - generator = torch.Generator(device=device).manual_seed(args.seed + n) - inputs = torch.randn( - (args.m, args.k), generator=generator, device=device, dtype=torch.float16 - ).mul_(0.02) - weight = torch.randn( - (args.k, n), generator=generator, device=device, dtype=torch.float16 - ).mul_(0.02) - reference = torch.empty((args.m, n), device=device, dtype=torch.float16) - torch.mm(inputs, weight, out=reference) - torch.accelerator.synchronize() - - algorithms = [_CUBLAS_GEMM_DEFAULT_TENSOR_OP] + list( - range(_CUBLAS_GEMM_ALGO0_TENSOR_OP, _CUBLAS_GEMM_ALGO0_TENSOR_OP + 16) - ) - results: list[dict[str, Any]] = [] - for algo in algorithms: - output = torch.empty_like(reference) - - def launch(output: torch.Tensor = output, algo: int = algo) -> None: - status = cublas.gemm(output, inputs, weight, algo) - if status != _CUBLAS_STATUS_SUCCESS: - raise RuntimeError(f"cuBLAS algo {algo} returned status {status}.") - - try: - timing = _event_trials(launch, args.warmup, args.iters, args.trials) - launch() - torch.accelerator.synchronize() - difference = (output.float() - reference.float()).abs() - results.append( - { - "algo": algo, - "status": "success", - "timing": timing, - "useful_tflops": float( - 2 * args.m * n * args.k / (timing["median_ms"] * 1e-3) / 1e12 - ), - "max_abs_diff": float(difference.max().item()), - "mean_abs_diff": float(difference.mean().item()), - "equal_elements": int((output == reference).sum().item()), - "elements": output.numel(), - } - ) - except RuntimeError as error: - torch.accelerator.synchronize() - results.append({"algo": algo, "status": "rejected", "error": str(error)}) - return { - "m": args.m, - "n": n, - "k": args.k, - "algorithms": sorted( - results, - key=lambda item: item.get("timing", {}).get("median_ms", float("inf")), - ), - } - - -def main() -> int: - args = _parse_args() - device = torch.device(args.device) - if torch.cuda.get_device_capability(device) != (7, 0): - raise RuntimeError("This benchmark requires SM70/V100.") - torch.accelerator.set_device_index(device) - cublas = _Cublas(torch.cuda.current_stream(device).cuda_stream) - try: - shapes = [_run_shape(args, device, cublas, n) for n in args.n] - finally: - cublas.close() - payload = { - "environment": { - "device": torch.cuda.get_device_name(device), - "torch": torch.__version__, - "torch_cuda": torch.version.cuda, - "m": args.m, - "k": args.k, - "n": args.n, - "warmup": args.warmup, - "iters": args.iters, - "trials": args.trials, - "seed": args.seed, - }, - "shapes": shapes, - } - args.json_out.parent.mkdir(parents=True, exist_ok=True) - args.json_out.write_text(json.dumps(payload, indent=2) + "\n") - print(json.dumps(payload, indent=2)) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/benchmarks/benchmark_sm70_custom_all_reduce.py b/benchmarks/benchmark_sm70_custom_all_reduce.py index 156484f987..88d6398437 100644 --- a/benchmarks/benchmark_sm70_custom_all_reduce.py +++ b/benchmarks/benchmark_sm70_custom_all_reduce.py @@ -51,19 +51,15 @@ def _make_input( elif pattern == "random_small": generator = torch.Generator(device="cuda") generator.manual_seed(seed + rank) - tensor = (torch.rand(size, device="cuda", generator=generator) - 0.5).to(dtype) + tensor = (torch.rand(size, device="cuda", generator=generator) - 0.5).to( + dtype + ) elif pattern == "model_like": generator = torch.Generator(device="cuda") generator.manual_seed(seed + rank) tensor = (torch.randn(size, device="cuda", generator=generator) * 0.03).to( dtype ) - elif pattern == "positive_zero": - tensor = torch.zeros(size, device="cuda", dtype=dtype) - elif pattern == "signed_zero": - tensor = torch.zeros(size, device="cuda", dtype=dtype) - if rank % 2: - tensor = -tensor else: raise ValueError(f"unsupported pattern: {pattern}") return tensor @@ -136,13 +132,7 @@ def _reference_custom_order( def _compare(out: torch.Tensor, ref: torch.Tensor) -> dict[str, Any]: diff = (out.float() - ref.float()).abs() mismatch = out != ref - bit_dtype = torch.int16 if out.element_size() == 2 else torch.int32 - out_bits = out.view(bit_dtype) - ref_bits = ref.view(bit_dtype) - bit_mismatch = out_bits != ref_bits - negative_zero = -32768 if out.element_size() == 2 else -2147483648 nonzero = int(mismatch.sum().item()) - bitwise_nonzero = int(bit_mismatch.sum().item()) first_mismatch: dict[str, Any] | None = None if nonzero: idx = int(torch.nonzero(mismatch, as_tuple=False)[0].item()) @@ -154,10 +144,6 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> dict[str, Any]: } return { "equal": bool(torch.equal(out, ref)), - "bitwise_equal": bitwise_nonzero == 0, - "bitwise_nonzero_count": bitwise_nonzero, - "out_negative_zero_count": int((out_bits == negative_zero).sum().item()), - "ref_negative_zero_count": int((ref_bits == negative_zero).sum().item()), "max_diff": float(diff.max().item()), "mean_diff": float(diff.mean().item()), "nonzero_count": nonzero, @@ -306,15 +292,12 @@ def _make_gemma_rms_norm_inputs( shared_generator = torch.Generator(device="cuda") shared_generator.manual_seed(seed + 1009) - residual = ( - torch.randn( - (tokens, _GEMMA_RMS_NORM_HIDDEN_SIZE), - device="cuda", - dtype=torch.float32, - generator=shared_generator, - ) - * 0.03 - ) + residual = torch.randn( + (tokens, _GEMMA_RMS_NORM_HIDDEN_SIZE), + device="cuda", + dtype=torch.float32, + generator=shared_generator, + ) * 0.03 residual = residual.to(residual_dtype) weight = ( torch.randn( @@ -587,11 +570,6 @@ def _check_graph_multi( for i in range(multi_count) ] refs = [_reference_all_reduce(inp) for inp in inputs] - num_bytes = size * torch.empty((), dtype=dtype).element_size() - algo = _custom_allreduce_algo(ca.world_size, ca.fully_connected, num_bytes) - custom_order_refs = [ - _reference_custom_order(inp, algo, ca.world_size) for inp in inputs - ] graph_timing: dict[str, float] | None = None if timing_iterations: outputs, graph_timing = _time_custom_ar_graph( @@ -600,9 +578,6 @@ def _check_graph_multi( else: outputs = _run_graph(ca, inputs, graph_replays) comparisons = [_compare(out, ref) for out, ref in zip(outputs, refs)] - custom_order_comparisons = [ - _compare(out, ref) for out, ref in zip(outputs, custom_order_refs) - ] max_diff = max(item["max_diff"] for item in comparisons) nonzero_count = sum(item["nonzero_count"] for item in comparisons) first_bad = next( @@ -616,27 +591,14 @@ def _check_graph_multi( return { "mode": "graph_multi", "size": size, - "bytes": num_bytes, + "bytes": size * torch.empty((), dtype=dtype).element_size(), "dtype": str(dtype).replace("torch.", ""), "pattern": pattern, - "custom_allreduce_algo": algo, "multi_count": multi_count, "all_equal": all(item["equal"] for item in comparisons), "max_diff": float(max_diff), "nonzero_count": int(nonzero_count), "first_bad": first_bad, - "custom_order_all_equal": all( - item["equal"] for item in custom_order_comparisons - ), - "custom_order_bitwise_all_equal": all( - item["bitwise_equal"] for item in custom_order_comparisons - ), - "custom_order_bitwise_nonzero_count": sum( - item["bitwise_nonzero_count"] for item in custom_order_comparisons - ), - "custom_order_max_diff": max( - item["max_diff"] for item in custom_order_comparisons - ), "graph_timing": graph_timing, } @@ -873,11 +835,11 @@ def main() -> None: if item["pattern"] not in require_exact_patterns: continue if item["mode"] == "graph_multi": - equal = item["custom_order_bitwise_all_equal"] + equal = item["all_equal"] comparison = item else: - comparison = item["custom_order_comparison"] - equal = comparison["bitwise_equal"] + equal = item["comparison"]["equal"] + comparison = item["comparison"] if not equal: required_failures.append( { @@ -895,11 +857,9 @@ def main() -> None: "custom_allreduce_disabled": False, "fully_connected": ca.fully_connected, "max_size_bytes": ca.max_size, - "sm70_tp4_push_buffer_registered": ( - ca.sm70_tp4_push_buffer_ptrs is not None + "sm70_tp4_m5_ar_threads": os.environ.get( + "VLLM_SM70_TP4_M5_AR_THREADS" ), - "sm70_tp4_m5_ar_threads": os.environ.get("VLLM_SM70_TP4_M5_AR_THREADS"), - "sm70_tp4_push_allreduce": os.environ.get("VLLM_SM70_TP4_PUSH_ALLREDUCE"), "required_exact_patterns": sorted(require_exact_patterns), "required_exact_passed": not required_failures, "required_failures": required_failures[:16], diff --git a/benchmarks/benchmark_sm70_d256_exact_dense.py b/benchmarks/benchmark_sm70_d256_exact_dense.py deleted file mode 100644 index 6b32dbab23..0000000000 --- a/benchmarks/benchmark_sm70_d256_exact_dense.py +++ /dev/null @@ -1,215 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Benchmark one SM70 D256 exact-dense prefill extension. - -Run control and candidate shared libraries in separate processes because both -register the same Torch operator namespace. A saved control tensor provides -the elementwise numerical gate for a candidate process. -""" - -from __future__ import annotations - -import argparse -import hashlib -import json -import statistics -from pathlib import Path -from typing import Any - -import torch - - -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--extension", type=Path, required=True) - parser.add_argument("--query-len", type=int, default=8192) - parser.add_argument("--kv-len", type=int, default=131072) - parser.add_argument("--heads-q", type=int, default=6) - parser.add_argument("--heads-kv", type=int, default=1) - parser.add_argument("--head-dim", type=int, default=256) - parser.add_argument("--warmup", type=int, default=5) - parser.add_argument("--iters", type=int, default=5) - parser.add_argument("--trials", type=int, default=7) - parser.add_argument("--seed", type=int, default=20260824) - parser.add_argument("--reference-out", type=Path) - parser.add_argument("--write-reference", action="store_true") - parser.add_argument("--json-out", type=Path, required=True) - return parser.parse_args() - - -def _digest_bytes(data: bytes) -> str: - return hashlib.sha256(data).hexdigest() - - -def _digest_tensor(tensor: torch.Tensor) -> str: - return _digest_bytes(tensor.contiguous().view(torch.uint8).numpy().tobytes()) - - -def _event_trials( - launch: Any, - *, - warmup: int, - iters: int, - trials: int, -) -> dict[str, Any]: - for _ in range(warmup): - launch() - torch.accelerator.synchronize() - - samples: list[float] = [] - for _ in range(trials): - start = torch.Event(enable_timing=True) - end = torch.Event(enable_timing=True) - start.record() - for _ in range(iters): - launch() - end.record() - end.synchronize() - samples.append(float(start.elapsed_time(end) / iters)) - return { - "median_ms": float(statistics.median(samples)), - "mean_ms": float(statistics.mean(samples)), - "min_ms": min(samples), - "max_ms": max(samples), - "samples_ms": samples, - } - - -def _causal_tflops( - *, - query_len: int, - kv_len: int, - heads_q: int, - head_dim: int, - elapsed_ms: float, -) -> float: - prefix_len = kv_len - query_len - mean_visible_keys = prefix_len + (query_len + 1) / 2.0 - flops = 4.0 * query_len * mean_visible_keys * heads_q * head_dim - return flops / (elapsed_ms * 1e9) - - -@torch.inference_mode() -def main() -> int: - args = _parse_args() - extension = args.extension.resolve() - if not extension.is_file(): - raise FileNotFoundError(extension) - if not torch.cuda.is_available(): - raise RuntimeError("CUDA is required") - if torch.cuda.get_device_capability() != (7, 0): - raise RuntimeError("This benchmark requires SM70/V100") - if args.head_dim != 256: - raise ValueError("The exact operator requires head_dim=256") - if args.query_len <= 0 or args.query_len > args.kv_len: - raise ValueError("Expected 0 < query_len <= kv_len") - if args.query_len % 64 or args.kv_len % 32: - raise ValueError("Query/KV lengths must be divisible by 64/32") - if args.heads_q % args.heads_kv: - raise ValueError("heads_q must be divisible by heads_kv") - if min(args.warmup, args.iters, args.trials) <= 0: - raise ValueError("warmup, iters, and trials must be positive") - if args.write_reference and args.reference_out is None: - raise ValueError("--write-reference requires --reference-out") - - torch.ops.load_library(str(extension)) - op = torch.ops._vllm_fa2_C.sm70_d256_splitd_n32_dense_fwd - - device = torch.device("cuda") - generator = torch.Generator(device=device).manual_seed(args.seed) - query = torch.randn( - (1, args.query_len, args.heads_q, args.head_dim), - dtype=torch.float16, - device=device, - generator=generator, - ) - key = torch.randn( - (1, args.kv_len, args.heads_kv, args.head_dim), - dtype=torch.float16, - device=device, - generator=generator, - ) - value = torch.randn( - (1, args.kv_len, args.heads_kv, args.head_dim), - dtype=torch.float16, - device=device, - generator=generator, - ) - output = torch.empty_like(query) - scale = args.head_dim**-0.5 - - def launch() -> None: - op(query, key, value, output, scale, True) - - timing = _event_trials( - launch, - warmup=args.warmup, - iters=args.iters, - trials=args.trials, - ) - launch() - torch.accelerator.synchronize() - output_cpu = output.cpu() - - comparison = None - if args.reference_out is not None and args.reference_out.exists(): - reference = torch.load( - args.reference_out, - map_location="cpu", - weights_only=True, - ) - if reference.shape != output_cpu.shape or reference.dtype != output_cpu.dtype: - raise RuntimeError("Reference shape or dtype does not match output") - difference = (reference.float() - output_cpu.float()).abs() - comparison = { - "reference_sha256": _digest_tensor(reference), - "max_abs_diff": float(difference.max().item()), - "mean_abs_diff": float(difference.mean().item()), - "equal_elements": int((reference == output_cpu).sum().item()), - "elements": output_cpu.numel(), - } - if args.write_reference: - assert args.reference_out is not None - args.reference_out.parent.mkdir(parents=True, exist_ok=True) - torch.save(output_cpu, args.reference_out) - - median_ms = float(timing["median_ms"]) - payload = { - "environment": { - "device": torch.cuda.get_device_name(), - "torch": torch.__version__, - "torch_cuda": torch.version.cuda, - "extension": str(extension), - "extension_sha256": _digest_bytes(extension.read_bytes()), - "query_len": args.query_len, - "kv_len": args.kv_len, - "heads_q": args.heads_q, - "heads_kv": args.heads_kv, - "head_dim": args.head_dim, - "dtype": "float16", - "causal": True, - "softmax_scale": scale, - "seed": args.seed, - "warmup": args.warmup, - "iters": args.iters, - "trials": args.trials, - }, - "timing": timing, - "causal_tflops": _causal_tflops( - query_len=args.query_len, - kv_len=args.kv_len, - heads_q=args.heads_q, - head_dim=args.head_dim, - elapsed_ms=median_ms, - ), - "output_sha256": _digest_tensor(output_cpu), - "comparison": comparison, - } - args.json_out.parent.mkdir(parents=True, exist_ok=True) - args.json_out.write_text(json.dumps(payload, indent=2) + "\n") - print(json.dumps(payload, indent=2), flush=True) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/benchmarks/benchmark_sm70_decode.py b/benchmarks/benchmark_sm70_decode.py index bf6d364c19..1e50bbfa23 100644 --- a/benchmarks/benchmark_sm70_decode.py +++ b/benchmarks/benchmark_sm70_decode.py @@ -30,13 +30,7 @@ def _module_file(module_name: str) -> str | None: module = sys.modules.get(module_name) if module is not None: return getattr(module, "__file__", None) - try: - spec = importlib.util.find_spec(module_name) - except (AttributeError, ImportError, ValueError): - # Optional extension parents can raise while find_spec imports them. - # Diagnostics must report an unavailable extension, not discard an - # otherwise successful benchmark at result-serialization time. - return None + spec = importlib.util.find_spec(module_name) return spec.origin if spec is not None else None @@ -169,7 +163,6 @@ def _tracked_env() -> dict[str, str]: "VLLM_FLASH_", "FLASH_ATTN", "FLA_", - "PREFIX_", "CUDA_MODULE_LOADING", "TORCH_CUDA_ARCH_LIST", "TORCHINDUCTOR_", @@ -690,10 +683,6 @@ def _sm70_attention_policy(kv_cache_dtype: Any) -> dict[str, Any]: "VLLM_FLASH_V100_XQA_E5M2_PAIR_LOAD", True, ) - e5m2_batch_wide_load = _env_bool( - "VLLM_FLASH_V100_XQA_E5M2_BATCH_WIDE_LOAD", - True, - ) decode_dynamic_partitions = _env_bool( "VLLM_FLASH_V100_DECODE_DYNAMIC_PARTITIONS", True, @@ -809,10 +798,6 @@ def _sm70_attention_policy(kv_cache_dtype: Any) -> dict[str, Any]: "VLLM_FLASH_V100_XQA_E5M2_PAIR_LOAD" ), "e5m2_pair_load_effective": e5m2_pair_load, - "VLLM_FLASH_V100_XQA_E5M2_BATCH_WIDE_LOAD": os.environ.get( - "VLLM_FLASH_V100_XQA_E5M2_BATCH_WIDE_LOAD" - ), - "e5m2_batch_wide_load_effective": e5m2_batch_wide_load, "exact_mtp5_fp8_p1024_dual_cta_policy": ( smallq_decode_use_xqa and mtp5_xqa_dual_cta diff --git a/benchmarks/benchmark_sm70_dflash2_gdn_metadata.py b/benchmarks/benchmark_sm70_dflash2_gdn_metadata.py deleted file mode 100644 index 7c138529f0..0000000000 --- a/benchmarks/benchmark_sm70_dflash2_gdn_metadata.py +++ /dev/null @@ -1,373 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Microbenchmark DFlash2 target GDN metadata construction on SM70.""" - -from __future__ import annotations - -import argparse -import json -import os -import statistics -import time -from dataclasses import dataclass -from types import SimpleNamespace -from typing import Any - -import torch - -from vllm import envs -from vllm.triton_utils import triton -from vllm.v1.attention.backends.gdn_attn import ( - DFlash2GDNGroupDescriptor, - _dflash2_gdn_group_metadata_kernel, - _get_ddtree_gdn_fast_common_buffers, - build_gdn_spec_decode_state_contract, - prepare_dflash2_gdn_group_metadata, -) -from vllm.v1.attention.backends.utils import PAD_SLOT_ID -from vllm.v1.worker.gpu.attn_utils import compute_common_gdn_attn_metadata - - -@dataclass -class _MicroBuilder: - num_spec_state_tokens: int - decode_cudagraph_max_bs: int - spec_state_indices_tensor: torch.Tensor - _ddtree_fast_common_buffers: Any - use_full_cuda_graph: bool = True - vllm_config: Any = None - - def __post_init__(self) -> None: - self.vllm_config = SimpleNamespace( - cache_config=SimpleNamespace(mamba_cache_mode="none") - ) - - -def _percentile(samples: list[float], quantile: float) -> float: - ordered = sorted(samples) - position = (len(ordered) - 1) * quantile - lower = int(position) - upper = min(lower + 1, len(ordered) - 1) - weight = position - lower - return ordered[lower] + (ordered[upper] - ordered[lower]) * weight - - -def _distribution(samples: list[float]) -> dict[str, float]: - return { - "mean": statistics.fmean(samples), - "p50": _percentile(samples, 0.50), - "p90": _percentile(samples, 0.90), - "p99": _percentile(samples, 0.99), - "min": min(samples), - "max": max(samples), - } - - -def _profile_cuda_launches(step: Any) -> int: - from torch.profiler import ProfilerActivity, profile - - with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: - step() - torch.accelerator.synchronize() - return sum( - 1 - for event in prof.events() - if str(getattr(event, "device_type", "")).lower().endswith("cuda") - ) - - -def _measure(step: Any, repeats: int) -> dict[str, dict[str, float]]: - wall_samples: list[float] = [] - gpu_samples: list[float] = [] - start = torch.Event(enable_timing=True) - end = torch.Event(enable_timing=True) - for _ in range(repeats): - torch.accelerator.synchronize() - wall_start = time.perf_counter_ns() - start.record() - step() - end.record() - end.synchronize() - wall_samples.append((time.perf_counter_ns() - wall_start) / 1e6) - gpu_samples.append(start.elapsed_time(end)) - return { - "synchronized_wall_ms": _distribution(wall_samples), - "cuda_event_ms": _distribution(gpu_samples), - } - - -def _run_case( - batch_size: int, - num_groups: int, - width: int, - warmup: int, - repeats: int, -) -> dict[str, Any]: - device = torch.device("cuda") - num_actual_tokens = batch_size * width - common_buffers = _get_ddtree_gdn_fast_common_buffers( - device, - num_actual_tokens, - width, - ) - builders = [ - _MicroBuilder( - num_spec_state_tokens=width - 1, - decode_cudagraph_max_bs=num_actual_tokens, - spec_state_indices_tensor=torch.empty( - (num_actual_tokens, width), dtype=torch.int32, device=device - ), - _ddtree_fast_common_buffers=common_buffers, - ) - for _ in range(num_groups) - ] - block_tables = tuple( - ( - torch.arange( - batch_size * width, - dtype=torch.int32, - device=device, - ).view(batch_size, width) - + group_id * 10_000 - ).contiguous() - for group_id in range(num_groups) - ) - query_start_loc = torch.arange( - 0, - num_actual_tokens + 1, - width, - dtype=torch.int32, - device=device, - ) - query_start_loc_cpu = query_start_loc.cpu() - draft_counts_cpu = torch.full((batch_size,), width - 1, dtype=torch.int32) - common_metadata = compute_common_gdn_attn_metadata( - num_decode_draft_tokens_cpu=draft_counts_cpu, - query_start_loc=query_start_loc, - query_start_loc_cpu=query_start_loc_cpu, - num_spec_state_tokens=width - 1, - legacy_mixed_decode_routing=False, - ) - assert common_metadata is not None - accepted_values = [3, 5, 7, 2] - accepted = torch.tensor( - [accepted_values[index % len(accepted_values)] for index in range(batch_size)], - dtype=torch.int32, - device=device, - ) - seq_lens = torch.full((batch_size,), 4097, dtype=torch.int32, device=device) - spec_mask_cpu = torch.ones(batch_size, dtype=torch.bool) - spec_mask = common_metadata.spec_sequence_masks - - def legacy_step() -> None: - for group_id, builder in enumerate(builders): - contract = build_gdn_spec_decode_state_contract( - block_table_tensor=block_tables[group_id], - seq_lens=seq_lens, - block_size=16, - num_spec=width - 1, - spec_sequence_masks_cpu=spec_mask_cpu, - num_accepted_tokens=accepted, - current_state_block_ids=None, - is_mamba_cache_all=False, - ) - builder.spec_state_indices_tensor[:batch_size].copy_( - contract.spec_state_indices_tensor, - non_blocking=True, - ) - builder.spec_state_indices_tensor[batch_size:].fill_(PAD_SLOT_ID) - - common_buffers.spec_sequence_masks[:batch_size].copy_( - spec_mask, non_blocking=True - ) - common_buffers.spec_sequence_masks[batch_size:].fill_(False) - common_buffers.spec_token_indx[:num_actual_tokens].copy_( - common_metadata.spec_token_indx, - non_blocking=True, - ) - common_buffers.spec_query_start_loc[: batch_size + 1].copy_( - common_metadata.spec_query_start_loc, - non_blocking=True, - ) - spec_num_query_tokens = common_metadata.spec_query_start_loc[-1] - common_buffers.spec_query_start_loc[batch_size + 1 :].fill_( - spec_num_query_tokens - ) - common_buffers.num_accepted_tokens[:batch_size].copy_( - contract.num_accepted_tokens, - non_blocking=True, - ) - common_buffers.num_accepted_tokens[batch_size:].fill_(1) - common_buffers.spec_state_slot_selectors[:batch_size].copy_( - contract.spec_state_slot_selectors, - non_blocking=True, - ) - common_buffers.spec_state_slot_selectors[batch_size:].fill_(1) - - descriptor: DFlash2GDNGroupDescriptor | None = None - - def fused_step() -> None: - nonlocal descriptor - result = prepare_dflash2_gdn_group_metadata( - builders_by_group=[ - (group_id, builder) for group_id, builder in enumerate(builders) - ], - block_tables=block_tables, - common_gdn_metadata=common_metadata, - num_accepted_tokens=accepted, - num_actual_tokens=num_actual_tokens, - descriptor=descriptor, - ) - if result is None: - raise RuntimeError("fused DFlash2 GDN metadata path was not eligible") - _, descriptor = result - - def fused_kernel_only_step() -> None: - if descriptor is None: - raise RuntimeError("fused descriptor has not been initialized") - block = triton.next_power_of_2( - max(num_actual_tokens * width, num_actual_tokens + 1) - ) - _dflash2_gdn_group_metadata_kernel[(num_groups,)]( - descriptor.block_table_ptrs, - descriptor.state_output_ptrs, - descriptor.block_table_strides, - common_metadata.spec_query_start_loc, - accepted, - accepted, - common_buffers.spec_sequence_masks, - common_buffers.spec_query_start_loc, - common_buffers.num_accepted_tokens, - common_buffers.spec_state_slot_selectors, - batch_size, - num_actual_tokens, - WIDTH=width, - PAD_ID=PAD_SLOT_ID, - BLOCK=block, - num_warps=1, - ) - - def fused_with_host_fence_step() -> None: - fused_step() - if ( - common_metadata.spec_query_start_loc[-1].item() - != common_metadata.num_spec_decode_tokens - ): - raise AssertionError("invalid speculative query span") - - legacy_step() - torch.accelerator.synchronize() - expected_states = [ - builder.spec_state_indices_tensor.clone() for builder in builders - ] - expected_common = ( - common_buffers.spec_sequence_masks.clone(), - common_buffers.spec_query_start_loc.clone(), - common_buffers.num_accepted_tokens.clone(), - common_buffers.spec_state_slot_selectors.clone(), - ) - for builder in builders: - builder.spec_state_indices_tensor.fill_(777) - common_buffers.spec_sequence_masks.fill_(True) - common_buffers.spec_query_start_loc.fill_(-1) - common_buffers.num_accepted_tokens.fill_(-1) - common_buffers.spec_state_slot_selectors.fill_(-1) - fused_step() - torch.accelerator.synchronize() - for actual_builder, expected_state in zip(builders, expected_states): - torch.testing.assert_close( - actual_builder.spec_state_indices_tensor, - expected_state, - rtol=0, - atol=0, - ) - for actual, expected in zip( - ( - common_buffers.spec_sequence_masks, - common_buffers.spec_query_start_loc, - common_buffers.num_accepted_tokens, - common_buffers.spec_state_slot_selectors, - ), - expected_common, - ): - torch.testing.assert_close(actual, expected, rtol=0, atol=0) - - for _ in range(warmup): - legacy_step() - fused_step() - torch.accelerator.synchronize() - - legacy_launches = _profile_cuda_launches(legacy_step) - fused_launches = _profile_cuda_launches(fused_step) - fused_with_host_fence_launches = _profile_cuda_launches(fused_with_host_fence_step) - fused_kernel_only_launches = _profile_cuda_launches(fused_kernel_only_step) - legacy = _measure(legacy_step, repeats) - fused = _measure(fused_step, repeats) - fused_with_host_fence = _measure(fused_with_host_fence_step, repeats) - fused_kernel_only = _measure(fused_kernel_only_step, repeats) - return { - "batch_size": batch_size, - "num_actual_tokens": num_actual_tokens, - "num_groups": num_groups, - "state_width": width, - "correctness": "bitwise_equal_after_poison", - "legacy_cuda_launches": legacy_launches, - "fused_cuda_launches": fused_launches, - "fused_with_host_fence_cuda_launches": fused_with_host_fence_launches, - "fused_kernel_only_cuda_launches": fused_kernel_only_launches, - "legacy": legacy, - "fused": fused, - "fused_with_host_fence": fused_with_host_fence, - "fused_kernel_only": fused_kernel_only, - "host_fence_wall_cost_p50_ms": ( - fused_with_host_fence["synchronized_wall_ms"]["p50"] - - fused["synchronized_wall_ms"]["p50"] - ), - "fused_python_overhead_p50_ms": ( - fused["synchronized_wall_ms"]["p50"] - - fused_kernel_only["synchronized_wall_ms"]["p50"] - ), - "wall_speedup": ( - legacy["synchronized_wall_ms"]["p50"] / fused["synchronized_wall_ms"]["p50"] - ), - "gpu_speedup": (legacy["cuda_event_ms"]["p50"] / fused["cuda_event_ms"]["p50"]), - } - - -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument("--batch-sizes", type=int, nargs="+", default=[1, 2, 4]) - parser.add_argument("--num-groups", type=int, default=10) - parser.add_argument("--width", type=int, default=8) - parser.add_argument("--warmup", type=int, default=20) - 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 microbenchmark requires SM70") - os.environ["VLLM_SM70_DFLASH2_FUSED_GDN_METADATA"] = "1" - os.environ["VLLM_SM70_DFLASH2_GDN_METADATA_SHADOW"] = "0" - envs.disable_envs_cache() - - results = [ - _run_case( - batch_size=batch_size, - num_groups=args.num_groups, - width=args.width, - warmup=args.warmup, - repeats=args.repeats, - ) - for batch_size in args.batch_sizes - ] - print( - json.dumps( - {"device": torch.cuda.get_device_name(), "cases": results}, - indent=2, - ) - ) - - -if __name__ == "__main__": - main() diff --git a/benchmarks/benchmark_sm70_dflash2_gsm8k.py b/benchmarks/benchmark_sm70_dflash2_gsm8k.py deleted file mode 100644 index ac92519225..0000000000 --- a/benchmarks/benchmark_sm70_dflash2_gsm8k.py +++ /dev/null @@ -1,484 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Measure target-only or DFlash2 on a fixed local quality subset.""" - -from __future__ import annotations - -import argparse -import ast -import hashlib -import json -import math -import random -import subprocess -import time -from collections import Counter -from pathlib import Path -from typing import Any - -import regex as re -from benchmark_sm70_decode import ( - _diff_spec_metrics, - _hash_ids, - _json_safe, - _module_file, - _module_realpath, - _request_metrics_dict, - _spec_metrics_snapshot, - _tracked_env, -) - -INVALID_ANSWER = -9_999_999 -GSM8K_PROMPT_SUFFIX = ( - "\nPlease reason step by step, and put your final answer within \\boxed{}." -) - - -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser() - parser.add_argument("--model", type=Path, required=True) - parser.add_argument("--draft-model", type=Path) - parser.add_argument("--dataset", type=Path, required=True) - parser.add_argument( - "--dataset-format", - choices=("gsm8k", "turns"), - default="gsm8k", - help=( - "Input schema. 'gsm8k' scores question/answer rows; 'turns' uses " - "the first preformatted turn and leaves task scoring to a separate " - "evaluator." - ), - ) - parser.add_argument("--out", type=Path, required=True) - parser.add_argument("--mode", choices=("target-only", "dflash"), required=True) - parser.add_argument("--num-questions", type=int, default=64) - parser.add_argument("--start-index", type=int, default=0) - parser.add_argument( - "--dataset-order", - choices=("sequential", "zlab-shuffle42"), - default="sequential", - ) - parser.add_argument("--max-tokens", type=int, default=256) - parser.add_argument("--warmup-tokens", type=int, default=32) - parser.add_argument("--temperature", type=float, default=1.0) - parser.add_argument("--tensor-parallel-size", type=int, default=4) - parser.add_argument("--max-model-len", type=int, default=2048) - 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("--gpu-memory-utilization", type=float, default=0.8) - parser.add_argument( - "--target-kv-cache-dtype", - choices=("auto", "fp8_e5m2"), - default="fp8_e5m2", - ) - parser.add_argument("--enforce-eager", action="store_true") - parser.add_argument( - "--draft-sample-method", - choices=("greedy", "probabilistic"), - default="greedy", - ) - parser.add_argument( - "--draft-attention-backend", - choices=("FLASH_ATTN_V100", "TRITON_ATTN"), - default="FLASH_ATTN_V100", - help="Draft-only attention backend; the target remains on FLASH_ATTN_V100.", - ) - parser.add_argument("--seed", type=int, default=0) - parser.add_argument( - "--request-seed", - type=int, - default=0, - help="Sampling seed for every request; use -1 for server-style random seeds.", - ) - parser.add_argument( - "--cuda-profiler-capture", - action="store_true", - help="Wrap the measured generation in cudaProfilerStart/Stop for nsys.", - ) - return parser.parse_args() - - -def _sha256_file(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 _model_weight_files(model: Path) -> dict[str, str]: - index = model / "model.safetensors.index.json" - if not index.is_file(): - weight = model / "model.safetensors" - return {weight.name: str(weight.resolve())} - weight_map = json.loads(index.read_text())["weight_map"] - return { - name: str((model / name).resolve()) for name in sorted(set(weight_map.values())) - } - - -def _answer_value(text: str) -> int: - numbers = re.findall(r"\d+", text.replace(",", "")) - if not numbers: - return INVALID_ANSWER - try: - return int(ast.literal_eval(numbers[-1])) - except (SyntaxError, ValueError): - return INVALID_ANSWER - - -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 _distribution(values: list[float]) -> dict[str, float | int | None]: - return { - "count": len(values), - "mean": sum(values) / len(values) if values else None, - "p50": _percentile(values, 0.50), - "p90": _percentile(values, 0.90), - "p99": _percentile(values, 0.99), - "min": min(values) if values else None, - "max": max(values) if values else None, - } - - -def _load_rows(args: argparse.Namespace) -> list[tuple[int, dict[str, Any]]]: - rows = [ - json.loads(line) - for line in args.dataset.read_text(encoding="utf-8").splitlines() - if line - ] - indices = list(range(len(rows))) - if args.dataset_order == "zlab-shuffle42": - random.Random(42).shuffle(indices) - selected_indices = indices[args.start_index : args.start_index + args.num_questions] - if len(selected_indices) != args.num_questions: - raise ValueError( - f"Requested {args.num_questions} rows at index {args.start_index}, " - f"but the dataset provided {len(selected_indices)}." - ) - return [(index, rows[index]) for index in selected_indices] - - -def _prompt_content(row: dict[str, Any], dataset_format: str) -> str: - if dataset_format == "gsm8k": - return str(row["question"]) + GSM8K_PROMPT_SUFFIX - turns = row.get("turns") - if not isinstance(turns, list) or not turns or not isinstance(turns[0], str): - raise ValueError("turns dataset rows must contain a non-empty string list") - return turns[0] - - -def _summarize_requests(cases: list[dict[str, Any]]) -> dict[str, Any]: - metric_names = ( - "queued_time", - "first_token_latency", - "prefill_time", - "decode_time", - "steady_decode_tps", - "tpot_seconds", - ) - summary = {} - for metric_name in metric_names: - values = [ - float(case["request_metrics"][metric_name]) - for case in cases - if case["request_metrics"] is not None - and case["request_metrics"].get(metric_name) is not None - ] - summary[metric_name] = _distribution(values) - - prefill_tps = [ - case["prompt_tokens"] / case["request_metrics"]["prefill_time"] - for case in cases - if case["request_metrics"] is not None - and case["request_metrics"].get("prefill_time") - ] - summary["prefill_tokens_per_second"] = _distribution(prefill_tps) - summary["prompt_tokens"] = _distribution( - [float(case["prompt_tokens"]) for case in cases] - ) - summary["output_tokens"] = _distribution( - [float(case["output_tokens"]) for case in cases] - ) - return summary - - -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") - - import torch - import vllm._C as vllm_c - import vllm._C_stable_libtorch as vllm_c_stable - from transformers import AutoTokenizer - - import vllm - from vllm import LLM, SamplingParams - - rows = _load_rows(args) - tokenizer = AutoTokenizer.from_pretrained(str(args.model)) - prompt_contents = [_prompt_content(row, args.dataset_format) for _, row in rows] - prompts = [ - tokenizer.apply_chat_template( - [ - { - "role": "user", - "content": prompt_content, - } - ], - tokenize=False, - add_generation_prompt=True, - enable_thinking=True, - reasoning_effort="xhigh", - ) - for prompt_content in prompt_contents - ] - prompt_token_counts = [len(tokenizer.encode(prompt)) for prompt in prompts] - - speculative_config = None - if args.mode == "dflash": - speculative_config = { - "method": "dflash", - "model": str(args.draft_model), - "revision": "dedf8df68adfb1afeaf7b7480c0a0243108177b4", - "num_speculative_tokens": 7, - "kv_cache_dtype": "auto", - "attention_backend": args.draft_attention_backend, - "draft_sample_method": args.draft_sample_method, - "enforce_eager": args.enforce_eager, - } - - engine_kwargs = { - "model": str(args.model), - "tensor_parallel_size": args.tensor_parallel_size, - "dtype": "half", - "kv_cache_dtype": args.target_kv_cache_dtype, - "attention_backend": "FLASH_ATTN_V100", - "max_model_len": args.max_model_len, - "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, - "disable_log_stats": False, - "enforce_eager": args.enforce_eager, - "seed": args.seed, - "speculative_config": speculative_config, - } - if args.cuda_profiler_capture: - engine_kwargs["profiler_config"] = {"profiler": "cuda"} - - load_started = time.perf_counter() - llm = LLM(**engine_kwargs) - load_seconds = time.perf_counter() - load_started - - request_seed = None if args.request_seed == -1 else args.request_seed - warmup_sampling = SamplingParams( - temperature=args.temperature, - top_p=0.95, - top_k=20, - max_tokens=args.warmup_tokens, - seed=request_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) - ) - 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() - spec_after = _spec_metrics_snapshot(llm) - - cases = [] - for ( - dataset_index, - row, - ), prompt_content, prompt_tokens, output, request_spec in zip( - rows, prompt_contents, prompt_token_counts, outputs, request_spec_metrics - ): - result = output.outputs[0] - token_ids = list(result.token_ids) - if args.dataset_format == "gsm8k": - prediction = _answer_value(result.text) - expected = _answer_value(row["answer"]) - correct = prediction == expected - else: - prediction = None - expected = None - correct = None - cases.append( - { - "dataset_index": dataset_index, - "question": row.get("question"), - "prompt_content": prompt_content, - "expected_answer": expected, - "predicted_answer": prediction, - "correct": correct, - "prompt_tokens": prompt_tokens, - "output_tokens": len(token_ids), - "finish_reason": result.finish_reason, - "stop_reason": result.stop_reason, - "text": result.text, - "token_ids": token_ids, - "token_hash": _hash_ids(token_ids), - "spec_decode_metrics": request_spec, - "request_metrics": _request_metrics_dict( - output.metrics, - len(token_ids), - ), - } - ) - - total_output_tokens = sum(case["output_tokens"] for case in cases) - scored_cases = [case for case in cases if case["correct"] is not None] - correct = sum(bool(case["correct"]) for case in scored_cases) - invalid = sum(case["predicted_answer"] == INVALID_ANSWER for case in scored_cases) - per_request_acceptance_lengths = [ - float(case["spec_decode_metrics"]["acceptance_length"]) - for case in cases - if case["spec_decode_metrics"] is not None - and case["spec_decode_metrics"].get("acceptance_length") is not None - ] - per_request_completion_tokens_per_verification_step = [ - case["output_tokens"] / case["spec_decode_metrics"]["num_drafts"] - for case in cases - if case["spec_decode_metrics"] is not None - and case["spec_decode_metrics"].get("num_drafts", 0) > 0 - ] - c_extension = Path(vllm_c.__file__).resolve() - c_stable_extension = Path(vllm_c_stable.__file__).resolve() - payload = { - "contract": { - "source_sha": subprocess.check_output( - ["git", "rev-parse", "HEAD"], - cwd=Path(__file__).resolve().parents[1], - text=True, - ).strip(), - "mode": args.mode, - "dataset": str(args.dataset), - "dataset_sha256": _sha256_file(args.dataset), - "dataset_format": args.dataset_format, - "start_index": args.start_index, - "num_questions": args.num_questions, - "dataset_order": args.dataset_order, - "model": str(args.model), - "model_config_sha256": _sha256_file(args.model / "config.json"), - "model_index_sha256": ( - _sha256_file(args.model / "model.safetensors.index.json") - if (args.model / "model.safetensors.index.json").is_file() - else None - ), - "model_weight_files": _model_weight_files(args.model), - "model_weights_realpath": str((args.model / "model.safetensors").resolve()), - "draft_model": str(args.draft_model) if args.draft_model else None, - "graph": not args.enforce_eager, - "sequential": args.sequential, - "sampling": { - "temperature": args.temperature, - "top_p": 0.95, - "top_k": 20, - "max_tokens": args.max_tokens, - "seed": request_seed, - "ignore_eos": False, - "thinking": True, - "reasoning_effort": "xhigh", - "prompt_suffix": ( - GSM8K_PROMPT_SUFFIX if args.dataset_format == "gsm8k" else None - ), - }, - "engine_kwargs": engine_kwargs, - }, - "runtime": { - "vllm_version": getattr(vllm, "__version__", None), - "vllm_file": getattr(vllm, "__file__", None), - "torch_version": torch.__version__, - "torch_cuda": torch.version.cuda, - "cuda_device_count": torch.accelerator.device_count(), - "device_capabilities": [ - list(torch.cuda.get_device_capability(index)) - for index in range(torch.accelerator.device_count()) - ], - "c_extension": str(c_extension), - "c_extension_sha256": _sha256_file(c_extension), - "c_stable_extension": str(c_stable_extension), - "c_stable_extension_sha256": _sha256_file(c_stable_extension), - "flash_attn_v100_python": _module_file("flash_attn_v100"), - "flash_attn_v100_cuda": _module_realpath("flash_attn_v100_cuda"), - "tracked_env": _tracked_env(), - "load_seconds": load_seconds, - }, - "results": { - "elapsed_seconds": elapsed_seconds, - "total_output_tokens": total_output_tokens, - "aggregate_output_tokens_per_second": ( - total_output_tokens / elapsed_seconds - ), - "questions_per_second": len(cases) / elapsed_seconds, - "accuracy": correct / len(scored_cases) if scored_cases else None, - "invalid_answer_rate": ( - invalid / len(scored_cases) if scored_cases else None - ), - "finish_reasons": dict( - Counter(str(case["finish_reason"]) for case in cases) - ), - "request_metrics": _summarize_requests(cases), - "spec_decode_metrics": _diff_spec_metrics(spec_before, spec_after), - "per_request_acceptance_length": _distribution( - per_request_acceptance_lengths - ), - "per_request_completion_tokens_per_verification_step": _distribution( - per_request_completion_tokens_per_verification_step - ), - }, - "cases": cases, - } - payload = _json_safe(payload) - args.out.parent.mkdir(parents=True, exist_ok=True) - args.out.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") - printable = {"contract": payload["contract"], "results": payload["results"]} - print(json.dumps(printable, indent=2)) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/benchmarks/benchmark_sm70_dflash2_sparse_rejection.py b/benchmarks/benchmark_sm70_dflash2_sparse_rejection.py deleted file mode 100644 index 51cc7ecf35..0000000000 --- a/benchmarks/benchmark_sm70_dflash2_sparse_rejection.py +++ /dev/null @@ -1,243 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Benchmark compact DFlash2 top-k/top-p rejection on one SM70 GPU. - -This isolates target sampling after the TP merge. The separate TP4 compact -logit benchmark measures local top-k and candidate transport. -""" - -from __future__ import annotations - -import argparse -import json -import statistics -from collections.abc import Callable -from pathlib import Path -from typing import Any - -import torch - -from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p -from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import ( - dflash2_sparse_topk_rejection_sample, - rejection_sample, -) - - -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser() - parser.add_argument("--vocab-size", type=int, default=248320) - parser.add_argument("--num-speculative-steps", type=int, default=7) - parser.add_argument("--target-top-k", type=int, default=20) - parser.add_argument("--draft-top-k", type=int, default=16) - parser.add_argument("--top-p", type=float, default=0.95) - parser.add_argument("--warmup", type=int, default=50) - parser.add_argument("--iters", type=int, default=500) - parser.add_argument("--seed", type=int, default=20260823) - parser.add_argument("--device", default="cuda:0") - parser.add_argument("--json-out", type=Path) - return parser.parse_args() - - -def _time_cuda( - operation: Callable[[], Any], - device: torch.device, - warmup: int, - iters: int, -) -> dict[str, float]: - for _ in range(warmup): - operation() - torch.accelerator.synchronize(device) - - events = [ - (torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)) - for _ in range(iters) - ] - for start, end in events: - start.record() - operation() - end.record() - events[-1][1].synchronize() - samples = [float(start.elapsed_time(end)) for start, end in events] - ordered = sorted(samples) - return { - "mean_ms": statistics.fmean(samples), - "p50_ms": statistics.median(samples), - "p90_ms": ordered[int(0.9 * (iters - 1))], - "p99_ms": ordered[int(0.99 * (iters - 1))], - "min_ms": ordered[0], - "max_ms": ordered[-1], - } - - -def main() -> int: - args = _parse_args() - if not 0.0 < args.top_p <= 1.0: - raise ValueError("--top-p must be in (0, 1]") - if not 0 < args.target_top_k <= 64: - raise ValueError("--target-top-k must be in [1, 64]") - if not 0 < args.draft_top_k <= 64: - raise ValueError("--draft-top-k must be in [1, 64]") - - device = torch.device(args.device) - torch.accelerator.set_device_index(device) - capability = torch.cuda.get_device_capability(device) - if capability != (7, 0): - raise RuntimeError(f"Expected SM70, got sm_{capability[0]}{capability[1]}.") - torch.manual_seed(args.seed) - - num_steps = args.num_speculative_steps - num_logits = num_steps + 1 - raw_target = torch.randn( - num_logits, - args.vocab_size, - dtype=torch.float16, - device=device, - ) - target_topk_logits_fp16, target_topk_ids = torch.topk( - raw_target, - k=args.target_top_k, - dim=-1, - ) - target_topk_logits = target_topk_logits_fp16.float() - target_k = torch.full( - (num_logits,), - args.target_top_k, - dtype=torch.int32, - device=device, - ) - target_p_rows = torch.full( - (num_logits,), args.top_p, dtype=torch.float32, device=device - ) - processed_target = apply_top_k_top_p(raw_target.float(), target_k, target_p_rows) - - draft_topk_ids = target_topk_ids[:num_steps, : args.draft_top_k].view( - 1, num_steps, args.draft_top_k - ) - draft_topk_logits = ( - target_topk_logits[:num_steps, : args.draft_top_k] - + torch.randn( - num_steps, - args.draft_top_k, - dtype=torch.float32, - device=device, - ) - * 0.2 - ).view(1, num_steps, args.draft_top_k) - dense_draft = torch.full( - (1, num_steps, args.vocab_size), - -float("inf"), - dtype=torch.float32, - device=device, - ) - dense_draft.scatter_(2, draft_topk_ids, draft_topk_logits) - - draft_sampled = torch.zeros(num_logits, dtype=torch.int64, device=device) - draft_sampled[1:] = draft_topk_ids[0, :, 0] - cu_num_logits = torch.tensor([0, num_logits], dtype=torch.int32, device=device) - pos = torch.arange(num_logits, dtype=torch.int64, device=device) + 32768 - idx_mapping = torch.zeros(1, dtype=torch.int32, device=device) - expanded_idx_mapping = torch.zeros(num_logits, dtype=torch.int32, device=device) - expanded_local_pos = torch.arange(num_logits, dtype=torch.int32, device=device) - temperature = torch.ones(1, dtype=torch.float32, device=device) - top_p_per_req = torch.full((1,), args.top_p, dtype=torch.float32, device=device) - seeds = torch.tensor([args.seed], dtype=torch.int64, device=device) - - def dense_rejection() -> tuple[torch.Tensor, torch.Tensor]: - return rejection_sample( - processed_target, - dense_draft, - draft_sampled, - cu_num_logits, - pos, - idx_mapping, - expanded_idx_mapping, - expanded_local_pos, - temperature, - seeds, - num_steps, - ) - - def sparse_rejection() -> tuple[torch.Tensor, torch.Tensor]: - return dflash2_sparse_topk_rejection_sample( - target_topk_ids, - target_topk_logits, - draft_topk_ids, - draft_topk_logits, - draft_sampled, - cu_num_logits, - pos, - idx_mapping, - temperature, - top_p_per_req, - seeds, - num_steps, - ) - - def dense_topk_topp() -> torch.Tensor: - return apply_top_k_top_p(raw_target.float(), target_k, target_p_rows) - - def dense_finalize() -> tuple[torch.Tensor, torch.Tensor]: - target = dense_topk_topp() - return rejection_sample( - target, - dense_draft, - draft_sampled, - cu_num_logits, - pos, - idx_mapping, - expanded_idx_mapping, - expanded_local_pos, - temperature, - seeds, - num_steps, - ) - - dense_out, dense_count = dense_rejection() - sparse_out, sparse_count = sparse_rejection() - torch.accelerator.synchronize(device) - counts_equal = bool(torch.equal(dense_count, sparse_count)) - valid = torch.arange(num_logits, device=device) < dense_count[0] - tokens_equal = bool(torch.equal(dense_out[0, valid], sparse_out[0, valid])) - - result = { - "device": torch.cuda.get_device_name(device), - "device_capability": list(capability), - "shape": { - "num_logits": num_logits, - "vocab_size": args.vocab_size, - "target_top_k": args.target_top_k, - "draft_top_k": args.draft_top_k, - "top_p": args.top_p, - }, - "correctness": { - "num_sampled_equal": counts_equal, - "valid_tokens_equal": tokens_equal, - "dense_num_sampled": int(dense_count[0].item()), - "sparse_num_sampled": int(sparse_count[0].item()), - }, - "timings": { - "dense_topk_topp_only": _time_cuda( - dense_topk_topp, device, args.warmup, args.iters - ), - "dense_rejection_only": _time_cuda( - dense_rejection, device, args.warmup, args.iters - ), - "sparse_topk_topp_rejection": _time_cuda( - sparse_rejection, device, args.warmup, args.iters - ), - "dense_topk_topp_plus_rejection": _time_cuda( - dense_finalize, device, args.warmup, args.iters - ), - }, - } - text = json.dumps(result, indent=2, sort_keys=True) - print(text) - if args.json_out is not None: - args.json_out.parent.mkdir(parents=True, exist_ok=True) - args.json_out.write_text(text + "\n", encoding="utf-8") - return 0 if counts_equal and tokens_equal else 1 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/benchmarks/benchmark_sm70_flashinfer_fixed_entry.py b/benchmarks/benchmark_sm70_flashinfer_fixed_entry.py index 4eb48de306..2482003308 100644 --- a/benchmarks/benchmark_sm70_flashinfer_fixed_entry.py +++ b/benchmarks/benchmark_sm70_flashinfer_fixed_entry.py @@ -263,7 +263,9 @@ 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) @@ -348,7 +350,9 @@ 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") @@ -613,9 +617,15 @@ 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 e94a7f9ace..f83336bfbc 100644 --- a/benchmarks/benchmark_sm70_flashqla_indexed_state_cache.py +++ b/benchmarks/benchmark_sm70_flashqla_indexed_state_cache.py @@ -72,38 +72,27 @@ 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) @@ -162,9 +151,8 @@ 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 fd9e069e54..7d21e6def3 100644 --- a/benchmarks/benchmark_sm70_fp8_fixed_vs_dynamic.py +++ b/benchmarks/benchmark_sm70_fp8_fixed_vs_dynamic.py @@ -65,9 +65,8 @@ 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: @@ -101,7 +100,8 @@ 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,7 +120,8 @@ 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() @@ -177,7 +178,8 @@ 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: @@ -187,8 +189,7 @@ 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_fp8_prefill_exact_dense_blas.py b/benchmarks/benchmark_sm70_fp8_prefill_exact_dense_blas.py deleted file mode 100644 index 7f515a0438..0000000000 --- a/benchmarks/benchmark_sm70_fp8_prefill_exact_dense_blas.py +++ /dev/null @@ -1,368 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Benchmark the exact-dense SM70 FP8 prefill path with a selected BLAS backend. - -This loads the real TP-local Qwen3.8 gate-up and down-projection weights. Run -each BLAS backend in a fresh process so global ATen and cuBLAS caches cannot -contaminate the comparison. -""" - -from __future__ import annotations - -import argparse -import hashlib -import importlib.util -import json -import math -import os -import statistics -import sys -from collections.abc import Callable -from pathlib import Path -from typing import Any - -import torch -from safetensors import safe_open - -import vllm - - -def _load_vllm_c_extension() -> None: - if "vllm._C" in sys.modules: - vllm._C = sys.modules["vllm._C"] - return - candidate = os.getenv("VLLM_BENCH_C_EXTENSION") - if candidate is None: - importlib.import_module("vllm._C") - return - spec = importlib.util.spec_from_file_location("vllm._C", candidate) - if spec is None or spec.loader is None: - raise RuntimeError(f"Cannot load candidate extension: {candidate}") - module = importlib.util.module_from_spec(spec) - sys.modules["vllm._C"] = module - vllm._C = module - spec.loader.exec_module(module) - - -_load_vllm_c_extension() - -from vllm import _sm70_ops as sm70_ops # noqa: E402 - - -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser() - parser.add_argument( - "--model", - type=Path, - required=True, - ) - parser.add_argument("--json-out", type=Path, required=True) - parser.add_argument("--device", default="cuda:0") - parser.add_argument("--tp-size", type=int, default=4) - parser.add_argument("--tp-rank", type=int, default=0) - parser.add_argument( - "--cases", - nargs="+", - choices=("gate_up", "down", "attention_output"), - default=("gate_up", "down"), - ) - parser.add_argument("--m", type=int, default=8000) - parser.add_argument( - "--backend", - choices=("exact-dense", "resident-turbomind-f16"), - default="exact-dense", - ) - parser.add_argument( - "--gated-silu", - action="store_true", - help="Benchmark gate_up with the fused interleaved SiLU-and-mul output.", - ) - parser.add_argument( - "--blas-library", - choices=("default", "cublas", "cublaslt"), - required=True, - ) - parser.add_argument("--warmup", type=int, default=10) - parser.add_argument("--iters", type=int, default=30) - parser.add_argument("--trials", type=int, default=7) - parser.add_argument("--seed", type=int, default=20260823) - parser.add_argument("--cuda-graph", action="store_true") - return parser.parse_args() - - -def _load_keys(model: Path, keys: list[str]) -> dict[str, torch.Tensor]: - index = json.loads((model / "model.safetensors.index.json").read_text()) - weight_map: dict[str, str] = index["weight_map"] - by_file: dict[str, list[str]] = {} - for key in keys: - by_file.setdefault(weight_map[key], []).append(key) - - result: dict[str, torch.Tensor] = {} - for filename, file_keys in by_file.items(): - with safe_open(model / filename, framework="pt", device="cpu") as handle: - for key in file_keys: - result[key] = handle.get_tensor(key) - return result - - -def _column_shard( - weight: torch.Tensor, - scales: torch.Tensor, - tp_size: int, - tp_rank: int, -) -> tuple[torch.Tensor, torch.Tensor]: - n = int(weight.shape[0]) - if n % tp_size: - raise ValueError(f"N={n} is not divisible by TP={tp_size}.") - begin = tp_rank * (n // tp_size) - end = begin + n // tp_size - return ( - weight.view(torch.uint8)[begin:end].contiguous(), - scales[begin // 128 : math.ceil(end / 128)].contiguous(), - ) - - -def _row_shard( - weight: torch.Tensor, - scales: torch.Tensor, - tp_size: int, - tp_rank: int, -) -> tuple[torch.Tensor, torch.Tensor]: - k = int(weight.shape[1]) - if k % tp_size: - raise ValueError(f"K={k} is not divisible by TP={tp_size}.") - begin = tp_rank * (k // tp_size) - end = begin + k // tp_size - return ( - weight.view(torch.uint8)[:, begin:end].contiguous(), - scales[:, begin // 128 : math.ceil(end / 128)].contiguous(), - ) - - -def _load_case( - model: Path, - case: str, - tp_size: int, - tp_rank: int, - device: torch.device, -) -> tuple[torch.Tensor, torch.Tensor, list[str]]: - if case == "attention_output": - prefixes = ["model.language_model.layers.3.self_attn.o_proj"] - else: - root = "model.language_model.layers.1.mlp" - projections = ("gate_proj", "up_proj") if case == "gate_up" else ("down_proj",) - prefixes = [f"{root}.{projection}" for projection in projections] - keys = [ - key - for prefix in prefixes - for key in (f"{prefix}.weight", f"{prefix}.weight_scale_inv") - ] - loaded = _load_keys(model, keys) - - raw_parts: list[torch.Tensor] = [] - scale_parts: list[torch.Tensor] = [] - fp8_dtype: torch.dtype | None = None - for prefix in prefixes: - weight = loaded[f"{prefix}.weight"] - scales = loaded[f"{prefix}.weight_scale_inv"].float() - fp8_dtype = weight.dtype if fp8_dtype is None else fp8_dtype - if weight.dtype != fp8_dtype: - raise ValueError("Fused gate-up projections must use one FP8 dtype.") - shard = ( - _column_shard(weight, scales, tp_size, tp_rank) - if case == "gate_up" - else _row_shard(weight, scales, tp_size, tp_rank) - ) - raw_parts.append(shard[0]) - scale_parts.append(shard[1]) - - assert fp8_dtype is not None - raw = torch.cat(raw_parts, dim=0) - scales = torch.cat(scale_parts, dim=0) - return ( - raw.to(device).view(fp8_dtype).contiguous(), - scales.to(device).contiguous(), - prefixes, - ) - - -def _event_trials( - launch: Callable[[], None], warmup: int, iters: int, trials: int -) -> dict[str, Any]: - for _ in range(warmup): - launch() - torch.accelerator.synchronize() - samples: list[float] = [] - for _ in range(trials): - start = torch.Event(enable_timing=True) - end = torch.Event(enable_timing=True) - start.record() - for _ in range(iters): - launch() - end.record() - end.synchronize() - samples.append(float(start.elapsed_time(end) / iters)) - return { - "median_ms": float(statistics.median(samples)), - "mean_ms": float(statistics.mean(samples)), - "min_ms": min(samples), - "max_ms": max(samples), - "samples_ms": samples, - } - - -def _digest(tensor: torch.Tensor) -> str: - raw = tensor.detach().contiguous().cpu().view(torch.uint8).numpy().tobytes() - return hashlib.sha256(raw).hexdigest() - - -def _file_digest(path: Path) -> str: - return hashlib.sha256(path.read_bytes()).hexdigest() - - -def _run_case( - args: argparse.Namespace, case: str, device: torch.device -) -> dict[str, Any]: - qweight, scales, prefixes = _load_case( - args.model, case, args.tp_size, args.tp_rank, device - ) - n, k = (int(dim) for dim in qweight.shape) - expected = { - "gate_up": (8704, 5120), - "down": (5120, 4352), - "attention_output": (5120, 1536), - }[case] - if (n, k) != expected: - raise ValueError(f"Unexpected {case} TP-local shape N={n}, K={k}.") - - tm_weight, tm_scales, meta = sm70_ops.fp8_sm70_prepare(qweight, scales, 128) - k_ld, q_ld = (int(value) for value in meta.tolist()) - generator = torch.Generator(device=device).manual_seed(args.seed + n + k) - inputs = torch.randn( - (args.m, k), generator=generator, device=device, dtype=torch.float16 - ).mul_(0.1) - output = torch.empty((args.m, n), device=device, dtype=torch.float16) - if args.gated_silu and case == "gate_up": - output = torch.empty((args.m, n // 2), device=device, dtype=torch.float16) - resident_bytes = 0 - if args.backend == "resident-turbomind-f16": - dense_weight = torch.empty((k, n), device=device, dtype=torch.float16) - sm70_ops.fp8_sm70_dequantize_out(dense_weight, tm_weight, tm_scales, 128) - source_weight = dense_weight.t().contiguous() - f16_tm_weight, f16_meta = sm70_ops.sm70_f16_prepare(source_weight) - f16_k_ld = int(f16_meta[0].item()) - resident_bytes = f16_tm_weight.numel() * f16_tm_weight.element_size() - - def eager_launch() -> None: - sm70_ops.sm70_f16_gemm_out( - output, - inputs, - f16_tm_weight, - f16_k_ld, - args.gated_silu and case == "gate_up", - ) - - else: - workspace = torch.empty((k, n), device=device, dtype=torch.float16) - resident_bytes = workspace.numel() * workspace.element_size() - - def eager_launch() -> None: - sm70_ops.fp8_gemm_sm70_prefill_dispatch_out( - output, - workspace.data_ptr(), - inputs, - tm_weight, - tm_scales, - 128, - k_ld, - q_ld, - args.gated_silu and case == "gate_up", - 3920, - ) - - cutlass_env = "VLLM_SM70_FP8_PREFILL_CUTLASS" - selected_cutlass = os.environ.get(cutlass_env) - os.environ[cutlass_env] = "0" - eager_launch() - torch.accelerator.synchronize(device) - reference = output.clone() - if selected_cutlass is None: - os.environ.pop(cutlass_env) - else: - os.environ[cutlass_env] = selected_cutlass - - launch = eager_launch - graph: torch.cuda.CUDAGraph | None = None - if args.cuda_graph: - for _ in range(3): - eager_launch() - torch.accelerator.synchronize(device) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - eager_launch() - launch = graph.replay - - timing = _event_trials(launch, args.warmup, args.iters, args.trials) - launch() - torch.accelerator.synchronize(device) - output_float = output.float() - reference_float = reference.float() - return { - "case": case, - "prefixes": prefixes, - "m": args.m, - "n": n, - "k": k, - "timing": timing, - "useful_tflops": 2 * args.m * n * k / (timing["median_ms"] * 1e-3) / 1e12, - "output_sha256": _digest(output), - "reference_sha256": _digest(reference), - "output_bitwise_equal": bool(torch.equal(output, reference)), - "output_max_abs_diff": float( - torch.max(torch.abs(output_float - reference_float)).item() - ), - "output_finite": bool(torch.isfinite(output).all().item()), - "resident_bytes": resident_bytes, - } - - -def main() -> int: - args = _parse_args() - device = torch.device(args.device) - if torch.cuda.get_device_capability(device) != (7, 0): - raise RuntimeError("This benchmark requires an SM70/V100 GPU.") - torch.accelerator.set_device_index(device.index or 0) - if args.blas_library != "default": - torch.backends.cuda.preferred_blas_library(args.blas_library) - - results = [_run_case(args, case, device) for case in args.cases] - extension_path = Path(vllm._C.__file__).resolve() - payload = { - "environment": { - "model": str(args.model), - "device": torch.cuda.get_device_name(device), - "capability": list(torch.cuda.get_device_capability(device)), - "torch": torch.__version__, - "torch_cuda": torch.version.cuda, - "blas_library_requested": args.blas_library, - "blas_library_effective": str(torch.backends.cuda.preferred_blas_library()), - "gated_silu": args.gated_silu, - "backend": args.backend, - "cuda_graph": args.cuda_graph, - "cutlass": os.getenv("VLLM_SM70_FP8_PREFILL_CUTLASS", ""), - "vllm_c_extension": str(extension_path), - "vllm_c_extension_sha256": _file_digest(extension_path), - "tp_size": args.tp_size, - "tp_rank": args.tp_rank, - "seed": args.seed, - }, - "results": results, - } - args.json_out.parent.mkdir(parents=True, exist_ok=True) - args.json_out.write_text(json.dumps(payload, indent=2) + "\n") - print(json.dumps(payload, indent=2)) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/benchmarks/benchmark_sm70_fp8_prefill_tactics.py b/benchmarks/benchmark_sm70_fp8_prefill_tactics.py deleted file mode 100644 index 5015fba613..0000000000 --- a/benchmarks/benchmark_sm70_fp8_prefill_tactics.py +++ /dev/null @@ -1,486 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Benchmark TurboMind tactics on Qwen3.8 FP8 prefill projections. - -This diagnostic loads the real TP-local fused QKVZ and QKV weights used by -Qwen3.8-27B-FP8. Run each dispatch policy in a fresh process: TurboMind caches -the first exact GEMM descriptor, so switching from ``default`` to ``measure`` -inside one process would not actually tune the already-cached shape. -""" - -from __future__ import annotations - -import argparse -import hashlib -import importlib.util -import json -import math -import os -import statistics -import sys -from collections.abc import Callable -from pathlib import Path -from typing import Any, Literal - -import torch -from safetensors import safe_open - -import vllm - - -def _load_vllm_c_extension() -> None: - if "vllm._C" in sys.modules: - vllm._C = sys.modules["vllm._C"] - return - candidate = os.getenv("VLLM_BENCH_C_EXTENSION") - if candidate is None: - importlib.import_module("vllm._C") - return - spec = importlib.util.spec_from_file_location("vllm._C", candidate) - if spec is None or spec.loader is None: - raise RuntimeError(f"Cannot load candidate extension: {candidate}") - module = importlib.util.module_from_spec(spec) - sys.modules["vllm._C"] = module - vllm._C = module - spec.loader.exec_module(module) - - -_load_vllm_c_extension() - -from vllm import _sm70_ops as sm70_ops # noqa: E402 - -Case = Literal["gdn_qkvz", "full_qkv"] -Policy = Literal[ - "default", - "measure", - "resident-f16", - "resident-f16-measure", - "exact-dispatch", - "exact-mm", - "resident-mm", -] - - -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser() - parser.add_argument( - "--model", - type=Path, - required=True, - ) - parser.add_argument("--json-out", type=Path, required=True) - parser.add_argument("--device", default="cuda:0") - parser.add_argument("--tp-size", type=int, default=4) - parser.add_argument("--tp-rank", type=int, default=0) - parser.add_argument( - "--cases", - choices=("gdn_qkvz", "full_qkv"), - nargs="+", - default=["gdn_qkvz", "full_qkv"], - ) - parser.add_argument("--m", type=int, nargs="+", default=[8000]) - parser.add_argument( - "--policy", - choices=( - "default", - "measure", - "resident-f16", - "resident-f16-measure", - "exact-dispatch", - "exact-mm", - "resident-mm", - ), - required=True, - ) - parser.add_argument( - "--safe-fast", - choices=("off", "preserve-splits", "preserve-split-count"), - default="off", - help="Optionally constrain measured tactics to the default split contract.", - ) - parser.add_argument("--warmup", type=int, default=10) - parser.add_argument("--iters", type=int, default=50) - parser.add_argument("--trials", type=int, default=7) - parser.add_argument("--seed", type=int, default=20260823) - parser.add_argument("--exact-8k-prescaled-prefill", action="store_true") - parser.add_argument("--cuda-graph", action="store_true") - return parser.parse_args() - - -def _load_keys(model: Path, keys: list[str]) -> dict[str, torch.Tensor]: - index = json.loads((model / "model.safetensors.index.json").read_text()) - weight_map: dict[str, str] = index["weight_map"] - by_file: dict[str, list[str]] = {} - for key in keys: - by_file.setdefault(weight_map[key], []).append(key) - - result: dict[str, torch.Tensor] = {} - for filename, file_keys in by_file.items(): - with safe_open(model / filename, framework="pt", device="cpu") as handle: - for key in file_keys: - result[key] = handle.get_tensor(key) - return result - - -def _column_shard( - weight: torch.Tensor, - scales: torch.Tensor, - tp_size: int, - tp_rank: int, -) -> tuple[torch.Tensor, torch.Tensor]: - n = int(weight.shape[0]) - if n % tp_size: - raise ValueError(f"N={n} is not divisible by TP={tp_size}.") - begin = tp_rank * (n // tp_size) - end = begin + n // tp_size - scale_begin = begin // 128 - scale_end = math.ceil(end / 128) - return ( - weight.view(torch.uint8)[begin:end].contiguous(), - scales[scale_begin:scale_end].contiguous(), - ) - - -def _load_case( - model: Path, - case: Case, - tp_size: int, - tp_rank: int, - device: torch.device, -) -> tuple[torch.Tensor, torch.Tensor, list[str]]: - root = "model.language_model.layers" - if case == "gdn_qkvz": - prefixes = [ - f"{root}.1.linear_attn.in_proj_qkv", - f"{root}.1.linear_attn.in_proj_z", - ] - elif case == "full_qkv": - prefixes = [ - f"{root}.3.self_attn.q_proj", - f"{root}.3.self_attn.k_proj", - f"{root}.3.self_attn.v_proj", - ] - else: - raise ValueError(f"Unknown case: {case}") - - keys = [ - key - for prefix in prefixes - for key in (f"{prefix}.weight", f"{prefix}.weight_scale_inv") - ] - loaded = _load_keys(model, keys) - raw_parts: list[torch.Tensor] = [] - scale_parts: list[torch.Tensor] = [] - fp8_dtype: torch.dtype | None = None - for prefix in prefixes: - weight = loaded[f"{prefix}.weight"] - scales = loaded[f"{prefix}.weight_scale_inv"].float() - if fp8_dtype is None: - fp8_dtype = weight.dtype - elif fp8_dtype != weight.dtype: - raise ValueError("Fused projections must use one FP8 dtype.") - raw, scale = _column_shard(weight, scales, tp_size, tp_rank) - raw_parts.append(raw) - scale_parts.append(scale) - - assert fp8_dtype is not None - raw = torch.cat(raw_parts, dim=0) - scales = torch.cat(scale_parts, dim=0) - qweight = raw.to(device).view(fp8_dtype).contiguous() - return qweight, scales.to(device).contiguous(), prefixes - - -def _configure_policy(policy: Policy, safe_fast: str, max_m: int) -> None: - os.environ["VLLM_SM70_FP8_DENSE_TUNE_MAX_M"] = str(max_m) - os.environ["VLLM_SM70_FP8_TUNE_SMALL_SHAPES"] = "1" if policy == "measure" else "0" - os.environ["VLLM_SM70_F16_DENSE_TUNE_MAX_M"] = str(max_m) - os.environ["VLLM_SM70_AWQ_TUNE_SMALL_SHAPES"] = ( - "1" if policy == "resident-f16-measure" else "0" - ) - os.environ["VLLM_SM70_FP8_SAFE_FAST_SELECTOR"] = "0" if safe_fast == "off" else "1" - os.environ["VLLM_SM70_FP8_PRESERVE_DEFAULT_SPLITS"] = ( - "1" if safe_fast == "preserve-splits" else "0" - ) - os.environ["VLLM_SM70_FP8_PRESERVE_DEFAULT_SPLITS_ONLY"] = ( - "1" if safe_fast == "preserve-split-count" else "0" - ) - - -def _event_trials( - launch: Callable[[], None], warmup: int, iters: int, trials: int -) -> dict[str, float]: - for _ in range(warmup): - launch() - torch.accelerator.synchronize() - - samples_ms: list[float] = [] - for _ in range(trials): - start = torch.Event(enable_timing=True) - end = torch.Event(enable_timing=True) - start.record() - for _ in range(iters): - launch() - end.record() - end.synchronize() - samples_ms.append(float(start.elapsed_time(end) / iters)) - ordered = sorted(samples_ms) - return { - "median_ms": float(statistics.median(samples_ms)), - "min_ms": ordered[0], - "max_ms": ordered[-1], - } - - -def _digest(tensor: torch.Tensor) -> str: - raw = tensor.detach().contiguous().cpu().view(torch.uint8).numpy().tobytes() - return hashlib.sha256(raw).hexdigest() - - -def _file_digest(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as handle: - for chunk in iter(lambda: handle.read(1 << 20), b""): - digest.update(chunk) - return digest.hexdigest() - - -def _run_case( - args: argparse.Namespace, case: Case, m: int, device: torch.device -) -> dict[str, Any]: - qweight, scales, prefixes = _load_case( - args.model, case, args.tp_size, args.tp_rank, device - ) - n, k = (int(dim) for dim in qweight.shape) - if ( - k != 5120 - or (case == "gdn_qkvz" and n != 4096) - or (case == "full_qkv" and n != 3584) - ): - raise ValueError(f"Unexpected {case} TP-local shape N={n}, K={k}.") - - tm_weight, tm_scales, meta = sm70_ops.fp8_sm70_prepare(qweight, scales, 128) - if args.exact_8k_prescaled_prefill: - tm_scales.mul_(256) - k_ld, q_ld = (int(value) for value in meta.tolist()) - generator = torch.Generator(device=device).manual_seed(args.seed + m) - inputs = torch.randn( - (m, k), generator=generator, device=device, dtype=torch.float16 - ).mul_(0.1) - output = torch.empty((m, n), device=device, dtype=torch.float16) - resident_bytes = 0 - dense_workspace_bytes = 0 - reference_output = None - if args.policy in ( - "resident-f16", - "resident-f16-measure", - "exact-dispatch", - "exact-mm", - "resident-mm", - ): - reference_output = torch.empty_like(output) - if args.exact_8k_prescaled_prefill: - sm70_ops.fp8_gemm_sm70_prefill_prescaled_out( - reference_output, - inputs, - tm_weight, - tm_scales, - 128, - k_ld, - q_ld, - ) - else: - sm70_ops.fp8_gemm_sm70_out( - reference_output, - inputs, - tm_weight, - tm_scales, - 128, - k_ld, - q_ld, - False, - ) - dense_weight = torch.empty((k, n), device=device, dtype=torch.float16) - dense_workspace_bytes = int(dense_weight.numel() * dense_weight.element_size()) - if args.policy == "exact-dispatch": - - def launch( - output: torch.Tensor = output, - inputs: torch.Tensor = inputs, - dense_weight: torch.Tensor = dense_weight, - tm_weight: torch.Tensor = tm_weight, - tm_scales: torch.Tensor = tm_scales, - ) -> None: - sm70_ops.fp8_gemm_sm70_prefill_dispatch_out( - output, - dense_weight.data_ptr(), - inputs, - tm_weight, - tm_scales, - 128, - k_ld, - q_ld, - False, - m, - ) - - elif args.policy == "exact-mm": - - def launch( - output: torch.Tensor = output, - inputs: torch.Tensor = inputs, - dense_weight: torch.Tensor = dense_weight, - tm_weight: torch.Tensor = tm_weight, - tm_scales: torch.Tensor = tm_scales, - ) -> None: - sm70_ops.fp8_sm70_dequantize_out( - dense_weight, tm_weight, tm_scales, 128 - ) - torch.mm(inputs, dense_weight, out=output) - - elif args.policy == "resident-mm": - sm70_ops.fp8_sm70_dequantize_out(dense_weight, tm_weight, tm_scales, 128) - resident_bytes = dense_workspace_bytes - - def launch( - output: torch.Tensor = output, - inputs: torch.Tensor = inputs, - dense_weight: torch.Tensor = dense_weight, - ) -> None: - torch.mm(inputs, dense_weight, out=output) - - else: - sm70_ops.fp8_sm70_dequantize_out(dense_weight, tm_weight, tm_scales, 128) - resident_bytes = dense_workspace_bytes - source_f16_weight = dense_weight.t().contiguous() - f16_tm_weight, f16_meta = sm70_ops.sm70_f16_prepare(source_f16_weight) - f16_k_ld = int(f16_meta[0].item()) - - def launch( - output: torch.Tensor = output, - inputs: torch.Tensor = inputs, - f16_tm_weight: torch.Tensor = f16_tm_weight, - ) -> None: - sm70_ops.sm70_f16_gemm_out( - output, inputs, f16_tm_weight, f16_k_ld, False - ) - - else: - if args.exact_8k_prescaled_prefill: - - def launch( - output: torch.Tensor = output, - inputs: torch.Tensor = inputs, - tm_weight: torch.Tensor = tm_weight, - tm_scales: torch.Tensor = tm_scales, - ) -> None: - sm70_ops.fp8_gemm_sm70_prefill_prescaled_out( - output, inputs, tm_weight, tm_scales, 128, k_ld, q_ld - ) - - else: - - def launch( - output: torch.Tensor = output, - inputs: torch.Tensor = inputs, - tm_weight: torch.Tensor = tm_weight, - tm_scales: torch.Tensor = tm_scales, - ) -> None: - sm70_ops.fp8_gemm_sm70_out( - output, inputs, tm_weight, tm_scales, 128, k_ld, q_ld, False - ) - - if args.cuda_graph: - launch() - torch.accelerator.synchronize(device) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - launch() - launch = graph.replay - - timing = _event_trials(launch, args.warmup, args.iters, args.trials) - launch() - torch.accelerator.synchronize(device) - seconds = timing["median_ms"] * 1e-3 - result = { - "case": case, - "prefixes": prefixes, - "m": m, - "n": n, - "k": k, - "tm_meta": {"k_ld": k_ld, "q_ld": q_ld}, - "timing": timing, - "useful_tflops": float(2 * m * n * k / seconds / 1e12), - "dense_workspace_bytes": dense_workspace_bytes, - "resident_f16_bytes": resident_bytes, - "output_sha256": _digest(output), - "output_finite": bool(torch.isfinite(output).all().item()), - } - if reference_output is not None: - difference = (output.float() - reference_output.float()).abs() - result["reference"] = { - "output_sha256": _digest(reference_output), - "equal": bool(torch.equal(output, reference_output)), - "max_abs_diff": float(difference.max().item()), - "mean_abs_diff": float(difference.mean().item()), - "equal_elements": int((output == reference_output).sum().item()), - "elements": output.numel(), - } - del qweight, scales, tm_weight, tm_scales, inputs, output - torch.accelerator.empty_cache() - return result - - -def main() -> int: - args = _parse_args() - if not torch.cuda.is_available(): - raise RuntimeError("CUDA is required.") - device = torch.device(args.device) - if torch.cuda.get_device_capability(device) != (7, 0): - raise RuntimeError("This benchmark requires an SM70/V100 GPU.") - if args.tp_size <= 0 or not 0 <= args.tp_rank < args.tp_size: - raise ValueError("Invalid TP size/rank.") - if any(m < 1 for m in args.m): - raise ValueError("Every M must be positive.") - if args.warmup < 1 or args.iters < 1 or args.trials < 1: - raise ValueError("Warmup, iterations, and trials must be positive.") - - torch.accelerator.set_device_index(device.index or 0) - _configure_policy(args.policy, args.safe_fast, max(args.m)) - results = [_run_case(args, case, m, device) for case in args.cases for m in args.m] - extension_path = Path(vllm._C.__file__) - payload = { - "environment": { - "model": str(args.model), - "vllm_c_extension": str(extension_path), - "vllm_c_extension_sha256": _file_digest(extension_path), - "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), - "device": torch.cuda.get_device_name(device), - "capability": list(torch.cuda.get_device_capability(device)), - "torch": torch.__version__, - "torch_cuda": torch.version.cuda, - "tp_size": args.tp_size, - "tp_rank": args.tp_rank, - "policy": args.policy, - "safe_fast": args.safe_fast, - "warmup": args.warmup, - "iters": args.iters, - "trials": args.trials, - "seed": args.seed, - "exact_8k_prescaled_prefill": args.exact_8k_prescaled_prefill, - "cuda_graph": args.cuda_graph, - "tm_gemm_tune": os.environ.get("TM_GEMM_TUNE"), - "fast_targets": os.environ.get("VLLM_SM70_AWQ_TP2_FAST_TARGETS"), - "fp8_prefill_fast_selector": os.environ.get( - "VLLM_SM70_FP8_PREFILL_FAST_SELECTOR" - ), - }, - "results": results, - } - args.json_out.parent.mkdir(parents=True, exist_ok=True) - args.json_out.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8") - print(json.dumps(payload, indent=2)) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/benchmarks/benchmark_sm70_fp8_qpn8.py b/benchmarks/benchmark_sm70_fp8_qpn8.py deleted file mode 100644 index 29584b4082..0000000000 --- a/benchmarks/benchmark_sm70_fp8_qpn8.py +++ /dev/null @@ -1,729 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Race experimental QPN8 against TurboMind on real TP-local FP8 weights. - -The benchmark is deliberately operator-only: it does not change model -dispatch. It covers the exact Qwen3.8-27B-FP8 TP4 decode shapes at M=1, 2, 4, -and 8, and measures eager launches separately from CUDA Graph replay. - -The production gate/up GEMM has a fused SiLU epilogue. ``gate_up_raw`` is -therefore opt-in and diagnostic only; its raw-GEMM timing is not an end-to-end -replacement claim. -""" - -from __future__ import annotations - -import argparse -import gc -import json -import math -import statistics -from collections.abc import Callable -from pathlib import Path -from typing import Any - -import torch -import vllm._C # noqa: F401 -from safetensors import safe_open - -from vllm import _sm70_ops as sm70_ops - -_KORDER8 = [0, 2, 4, 6, 1, 3, 5, 7, 8, 10, 12, 14, 9, 11, 13, 15] -_DEFAULT_CASES = ("down", "gdn_in", "output", "full_qkv") -_ALL_CASES = (*_DEFAULT_CASES, "gate_up_raw", "gate_up_fused") - - -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser() - parser.add_argument( - "--model", - type=Path, - default=Path("/home/ymzx/models/Qwen3.8-27B-FP8"), - ) - parser.add_argument("--json-out", type=Path, required=True) - parser.add_argument( - "--qpn8-library", - type=Path, - help="Optional standalone library built from fp8_qpn8_sm70.cu.", - ) - parser.add_argument("--device", default="cuda:0") - parser.add_argument("--tp-size", type=int, default=4) - parser.add_argument("--tp-rank", type=int, default=0) - parser.add_argument("--m", type=int, nargs="+", default=[1, 2, 4, 8]) - parser.add_argument( - "--cases", choices=_ALL_CASES, nargs="+", default=list(_DEFAULT_CASES) - ) - parser.add_argument("--split-k", type=int, nargs="+", default=[4, 8, 16, 32]) - parser.add_argument("--nacc", type=int, nargs="+", default=[1, 2]) - parser.add_argument( - "--decoder", choices=("scalar", "fast"), nargs="+", default=["scalar", "fast"] - ) - parser.add_argument( - "--prefetch", choices=("off", "on"), nargs="+", default=["off", "on"] - ) - parser.add_argument("--warmup", type=int, default=10) - parser.add_argument("--iters", type=int, default=100) - parser.add_argument("--trials", type=int, default=5) - parser.add_argument( - "--cache-state", - choices=("warm", "cold"), - nargs="+", - default=["warm", "cold"], - ) - parser.add_argument("--cache-scrub-mib", type=int, default=64) - parser.add_argument("--seed", type=int, default=1234) - parser.add_argument("--relative-l2-limit", type=float, default=5e-3) - parser.add_argument("--cosine-limit", type=float, default=0.999) - return parser.parse_args() - - -def _load_keys(model: Path, keys: list[str]) -> dict[str, torch.Tensor]: - index = json.loads((model / "model.safetensors.index.json").read_text()) - weight_map: dict[str, str] = index["weight_map"] - by_file: dict[str, list[str]] = {} - for key in keys: - by_file.setdefault(weight_map[key], []).append(key) - - result: dict[str, torch.Tensor] = {} - for filename, file_keys in by_file.items(): - with safe_open(model / filename, framework="pt", device="cpu") as handle: - for key in file_keys: - result[key] = handle.get_tensor(key) - return result - - -def _weight_keys(prefix: str) -> tuple[str, str]: - return f"{prefix}.weight", f"{prefix}.weight_scale_inv" - - -def _column_shard_raw( - weight: torch.Tensor, - scales: torch.Tensor, - tp_size: int, - tp_rank: int, -) -> tuple[torch.Tensor, torch.Tensor]: - n = weight.shape[0] - if n % tp_size: - raise ValueError(f"N={n} is not divisible by TP={tp_size}") - begin = tp_rank * (n // tp_size) - end = begin + n // tp_size - scale_begin = begin // 128 - scale_end = math.ceil(end / 128) - raw = weight.view(torch.uint8)[begin:end].contiguous() - return raw, scales[scale_begin:scale_end].contiguous() - - -def _row_shard_raw( - weight: torch.Tensor, - scales: torch.Tensor, - tp_size: int, - tp_rank: int, -) -> tuple[torch.Tensor, torch.Tensor]: - k = weight.shape[1] - if k % tp_size: - raise ValueError(f"K={k} is not divisible by TP={tp_size}") - begin = tp_rank * (k // tp_size) - end = begin + k // tp_size - scale_begin = begin // 128 - scale_end = math.ceil(end / 128) - raw = weight.view(torch.uint8)[:, begin:end].contiguous() - return raw, scales[:, scale_begin:scale_end].contiguous() - - -def _load_case( - model: Path, - case: str, - tp_size: int, - tp_rank: int, - device: torch.device, -) -> tuple[torch.Tensor, torch.Tensor, str]: - root = "model.language_model.layers" - if case == "down": - prefixes = [f"{root}.1.mlp.down_proj"] - shard = "row" - note = "production down_proj TP row shard" - elif case == "gdn_in": - prefixes = [ - f"{root}.1.linear_attn.in_proj_qkv", - f"{root}.1.linear_attn.in_proj_z", - ] - shard = "column" - note = "production GDN qkv+z TP column shards concatenated" - elif case == "output": - prefixes = [f"{root}.1.linear_attn.out_proj"] - shard = "row" - note = "production GDN/full-attention output TP row-shard shape" - elif case == "full_qkv": - prefixes = [ - f"{root}.3.self_attn.q_proj", - f"{root}.3.self_attn.k_proj", - f"{root}.3.self_attn.v_proj", - ] - shard = "column" - note = "production full-attention q+k+v TP column shards concatenated" - elif case in ("gate_up_raw", "gate_up_fused"): - prefixes = [ - f"{root}.1.mlp.gate_proj", - f"{root}.1.mlp.up_proj", - ] - shard = "column" - note = ( - "production gate+up GEMM with fused SiLU-and-multiply" - if case == "gate_up_fused" - else "diagnostic raw gate+up GEMM; production has fused SiLU epilogue" - ) - else: - raise ValueError(f"unknown case: {case}") - - keys = [key for prefix in prefixes for key in _weight_keys(prefix)] - loaded = _load_keys(model, keys) - raw_parts: list[torch.Tensor] = [] - scale_parts: list[torch.Tensor] = [] - fp8_dtype: torch.dtype | None = None - for prefix in prefixes: - weight_key, scale_key = _weight_keys(prefix) - weight = loaded[weight_key] - scales = loaded[scale_key].float() - if fp8_dtype is None: - fp8_dtype = weight.dtype - elif weight.dtype != fp8_dtype: - raise ValueError("concatenated weights have different dtypes") - if shard == "column": - raw, scale = _column_shard_raw(weight, scales, tp_size, tp_rank) - else: - raw, scale = _row_shard_raw(weight, scales, tp_size, tp_rank) - raw_parts.append(raw) - scale_parts.append(scale) - - raw = torch.cat(raw_parts, dim=0) if len(raw_parts) > 1 else raw_parts[0] - scales = torch.cat(scale_parts, dim=0) if len(scale_parts) > 1 else scale_parts[0] - assert fp8_dtype is not None - qweight = raw.to(device, non_blocking=False).view(fp8_dtype).contiguous() - return qweight, scales.to(device).contiguous(), note - - -def _qpn8_prepack(raw_weight: torch.Tensor) -> torch.Tensor: - """[N,K] uint8 -> [tile N/32][group K/16][lane 32][16B].""" - n, k = raw_weight.shape - if n % 32 or k % 16: - raise ValueError(f"QPN8 packing requires N%32=K%16=0, got N={n} K={k}") - device = raw_weight.device - tiles, groups = n // 32, k // 16 - lane = torch.arange(32, device=device) - col = ((lane >> 2) & 3) * 8 + (lane & 3) + ((lane & 16) > 0).long() * 4 - korder = torch.tensor(_KORDER8, device=device) - group = torch.arange(groups, device=device) - kidx = group[:, None] * 16 + korder[None, :] - packed = torch.empty((tiles, groups, 32, 16), dtype=torch.uint8, device=device) - # Bound the temporary expanded index tensors for the 44-MiB projections. - tiles_per_chunk = max(1, 36864 // max(groups, 1)) - for tile_begin in range(0, tiles, tiles_per_chunk): - tile_end = min(tile_begin + tiles_per_chunk, tiles) - tile_count = tile_end - tile_begin - ncol = ( - torch.arange(tile_begin, tile_end, device=device)[:, None] * 32 - + col[None, :] - ) - packed[tile_begin:tile_end] = raw_weight[ - ncol.view(tile_count, 1, 32, 1).expand(tile_count, groups, 32, 16), - kidx.view(1, groups, 1, 16).expand(tile_count, groups, 32, 16), - ] - return packed.view(-1).contiguous() - - -def _qpn8_group_scales(scales: torch.Tensor, n: int, k: int) -> torch.Tensor: - expected = (math.ceil(n / 128), math.ceil(k / 128)) - if tuple(scales.shape) != expected or n % 128 or k % 128: - raise ValueError( - f"QPN8 scales need exact 128x128 blocks, got {tuple(scales.shape)} " - f"for N={n}, K={k}" - ) - return ( - scales.t() - .reshape(k // 128, n // 128) - .repeat_interleave(4, dim=1) - .mul(256.0) - .half() - .contiguous() - ) - - -def _dequantized_weight(qweight: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: - n, k = qweight.shape - expanded = scales.repeat_interleave(128, 0).repeat_interleave(128, 1) - return qweight.float().mul(expanded[:n, :k]) - - -def _error_stats(actual: torch.Tensor, reference: torch.Tensor) -> dict[str, Any]: - actual_f32 = actual.float() - reference_f32 = reference.float() - delta = actual_f32 - reference_f32 - ref_norm = torch.linalg.vector_norm(reference_f32).clamp_min(1e-12) - relative_l2 = torch.linalg.vector_norm(delta) / ref_norm - cosine = torch.nn.functional.cosine_similarity( - actual_f32.flatten(), reference_f32.flatten(), dim=0 - ) - return { - "finite": bool(torch.isfinite(actual).all().item()), - "max_abs": float(delta.abs().max().item()), - "mean_abs": float(delta.abs().mean().item()), - "relative_l2": float(relative_l2.item()), - "cosine": float(cosine.item()), - } - - -def _event_trials( - launch: Callable[[], None], - warmup: int, - iters: int, - trials: int, - cache_scrub: torch.Tensor | None, -) -> dict[str, float]: - for _ in range(warmup): - if cache_scrub is not None: - cache_scrub.add_(1) - launch() - torch.accelerator.synchronize() - samples_us: list[float] = [] - for _ in range(trials): - if cache_scrub is None: - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - start.record() - for _ in range(iters): - launch() - end.record() - end.synchronize() - samples_us.append(float(start.elapsed_time(end) * 1000.0 / iters)) - else: - # Scrub more than V100's L2 immediately before each call, while - # the event pair measures only the tested operator. - starts = [torch.cuda.Event(enable_timing=True) for _ in range(iters)] - ends = [torch.cuda.Event(enable_timing=True) for _ in range(iters)] - for start, end in zip(starts, ends): - cache_scrub.add_(1) - start.record() - launch() - end.record() - ends[-1].synchronize() - samples_us.append( - float( - statistics.mean( - start.elapsed_time(end) * 1000.0 - for start, end in zip(starts, ends) - ) - ) - ) - ordered = sorted(samples_us) - return { - "median_us": float(statistics.median(samples_us)), - "min_us": ordered[0], - "p10_us": ordered[max(0, math.ceil(0.1 * len(ordered)) - 1)], - "p90_us": ordered[min(len(ordered) - 1, math.ceil(0.9 * len(ordered)) - 1)], - "max_us": ordered[-1], - } - - -def _benchmark_eager( - launch: Callable[[], None], - warmup: int, - iters: int, - trials: int, - cache_scrub: torch.Tensor | None, -) -> dict[str, Any]: - timing = _event_trials(launch, warmup, iters, trials, cache_scrub) - launch() - torch.accelerator.synchronize() - return {"timing": timing, "replay_max_abs": None} - - -def _benchmark_graph( - launch: Callable[[], None], - output: torch.Tensor, - warmup: int, - iters: int, - trials: int, - cache_scrub: torch.Tensor | None, -) -> dict[str, Any]: - capture_stream = torch.cuda.Stream(device=output.device) - capture_stream.wait_stream(torch.cuda.current_stream(output.device)) - with torch.cuda.stream(capture_stream): - for _ in range(max(warmup, 3)): - launch() - capture_stream.synchronize() - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph, stream=capture_stream): - launch() - torch.accelerator.synchronize(output.device) - - timing = _event_trials(graph.replay, warmup, iters, trials, cache_scrub) - graph.replay() - torch.accelerator.synchronize(output.device) - first = output.clone() - graph.replay() - torch.accelerator.synchronize(output.device) - replay_max_abs = float((output - first).abs().max().item()) - del first, graph, capture_stream - gc.collect() - return {"timing": timing, "replay_max_abs": replay_max_abs} - - -def _derived_metrics( - timing: dict[str, float], m: int, n: int, k: int, bytes_per_call: int -) -> dict[str, float]: - seconds = timing["median_us"] * 1e-6 - return { - "effective_gbps": float(bytes_per_call / seconds / 1e9), - "useful_tflops": float(2 * m * n * k / seconds / 1e12), - } - - -def _quality_pass( - stats: dict[str, Any], relative_l2_limit: float, cosine_limit: float -) -> bool: - return bool( - stats["finite"] - and stats["relative_l2"] <= relative_l2_limit - and stats["cosine"] >= cosine_limit - ) - - -def _run_case(args: argparse.Namespace, case: str) -> dict[str, Any]: - device = torch.device(args.device) - qweight, scales, note = _load_case( - args.model, case, args.tp_size, args.tp_rank, device - ) - n, k = (int(dim) for dim in qweight.shape) - gated_silu = case == "gate_up_fused" - output_n = n // 2 if gated_silu else n - raw = qweight.view(torch.uint8) - reference_codes = _qpn8_prepack(raw) - reference_group_scales = _qpn8_group_scales(scales, n, k) - if hasattr(torch.ops._C, "fp8_qpn8_prepare_sm70"): - codes, group_scales = sm70_ops.fp8_qpn8_prepare_sm70(qweight, scales) - prepare_matches_reference = bool( - torch.equal(codes.view(-1), reference_codes) - and torch.equal(group_scales, reference_group_scales) - ) - if not prepare_matches_reference: - raise RuntimeError("source QPN8 prepare disagrees with Python reference") - else: - codes = reference_codes - group_scales = reference_group_scales - prepare_matches_reference = None - tm_weight, tm_scales, meta = sm70_ops.fp8_sm70_prepare( - qweight, scales, 128, gated_silu - ) - k_ld, q_ld = (int(value) for value in meta.tolist()) - reference_weight = _dequantized_weight(qweight, scales) - cache_scrub = torch.empty( - args.cache_scrub_mib * 1024 * 1024, - dtype=torch.uint8, - device=device, - ) - torch.accelerator.synchronize(device) - - case_result: dict[str, Any] = { - "case": case, - "note": note, - "shape": {"n": n, "k": k, "output_n": output_n}, - "qweight_dtype": str(qweight.dtype), - "scale_dtype": str(scales.dtype), - "tm_meta": {"k_ld": k_ld, "q_ld": q_ld}, - "packed_bytes": int(codes.numel()), - "group_scale_shape": list(group_scales.shape), - "source_prepare_matches_reference": prepare_matches_reference, - "rows": [], - } - - valid_configs = [ - (split_k, nacc, decoder == "fast", prefetch == "on") - for split_k in args.split_k - for nacc in args.nacc - for decoder in args.decoder - for prefetch in args.prefetch - if (k // 16) % split_k == 0 - and not (prefetch == "on" and decoder != "fast") - and not (gated_silu and split_k == 32) - ] - for m in args.m: - torch.manual_seed(args.seed + m) - inputs = torch.randn((m, k), device=device, dtype=torch.float16).mul_(0.1) - reference_raw = inputs.float().matmul(reference_weight.t()) - if gated_silu: - gate, up = reference_raw.chunk(2, dim=1) - reference = (torch.nn.functional.silu(gate) * up).half() - else: - reference = reference_raw.half() - - tm_out = torch.empty((m, output_n), device=device, dtype=torch.float16) - - def launch_tm( - tm_out: torch.Tensor = tm_out, - inputs: torch.Tensor = inputs, - tm_weight: torch.Tensor = tm_weight, - tm_scales: torch.Tensor = tm_scales, - ) -> None: - sm70_ops.fp8_gemm_sm70_out( - tm_out, - inputs, - tm_weight, - tm_scales, - 128, - k_ld, - q_ld, - gated_silu, - ) - - launch_tm() - torch.accelerator.synchronize(device) - tm_quality = _error_stats(tm_out, reference) - tm_bytes = n * k + 2 * (k // 128) * n + 2 * m * (output_n + k) - for cache_state in args.cache_state: - scrub = cache_scrub if cache_state == "cold" else None - for mode, benchmark in ( - ("eager", _benchmark_eager), - ("graph", _benchmark_graph), - ): - measured = ( - benchmark( - launch_tm, - args.warmup, - args.iters, - args.trials, - scrub, - ) - if mode == "eager" - else benchmark( - launch_tm, - tm_out, - args.warmup, - args.iters, - args.trials, - scrub, - ) - ) - case_result["rows"].append( - { - "backend": "turbomind_current", - "mode": mode, - "cache_state": cache_state, - "m": m, - "config": None, - "quality": tm_quality, - "quality_pass": _quality_pass( - tm_quality, args.relative_l2_limit, args.cosine_limit - ), - **measured, - **_derived_metrics(measured["timing"], m, n, k, tm_bytes), - } - ) - - for split_k, nacc, fast_decoder, prefetch_codes in valid_configs: - qpn_out = torch.empty((m, output_n), device=device, dtype=torch.float16) - - def launch_qpn( - split_k: int = split_k, - nacc: int = nacc, - fast_decoder: bool = fast_decoder, - prefetch_codes: bool = prefetch_codes, - qpn_out: torch.Tensor = qpn_out, - inputs: torch.Tensor = inputs, - codes: torch.Tensor = codes, - group_scales: torch.Tensor = group_scales, - gated_silu: bool = gated_silu, - ) -> None: - if gated_silu: - sm70_ops.fp8_qpn8_gated_pair_sm70_out( - qpn_out, - inputs, - codes, - group_scales, - split_k, - nacc, - fast_decoder, - prefetch_codes, - ) - else: - sm70_ops.fp8_qpn8_gemm_sm70_out( - qpn_out, - inputs, - codes, - group_scales, - split_k, - nacc, - fast_decoder, - prefetch_codes, - ) - - launch_qpn() - torch.accelerator.synchronize(device) - quality = _error_stats(qpn_out, reference) - quality_pass = _quality_pass( - quality, args.relative_l2_limit, args.cosine_limit - ) - config = { - "split_k": split_k, - "nacc": nacc, - "decoder": "fast" if fast_decoder else "scalar", - "prefetch": prefetch_codes, - } - qpn_bytes = n * k + 2 * (k // 128) * (n // 32) + 2 * m * (n + k) - if gated_silu: - qpn_bytes -= 2 * m * (n - output_n) - for cache_state in args.cache_state: - scrub = cache_scrub if cache_state == "cold" else None - for mode, benchmark in ( - ("eager", _benchmark_eager), - ("graph", _benchmark_graph), - ): - measured = ( - benchmark( - launch_qpn, - args.warmup, - args.iters, - args.trials, - scrub, - ) - if mode == "eager" - else benchmark( - launch_qpn, - qpn_out, - args.warmup, - args.iters, - args.trials, - scrub, - ) - ) - case_result["rows"].append( - { - "backend": "qpn8_experimental", - "mode": mode, - "cache_state": cache_state, - "m": m, - "config": config, - "quality": quality, - "quality_pass": quality_pass, - **measured, - **_derived_metrics(measured["timing"], m, n, k, qpn_bytes), - } - ) - - del qpn_out - - del inputs, reference_raw, reference, tm_out - - del ( - reference_weight, - tm_weight, - tm_scales, - codes, - group_scales, - reference_codes, - reference_group_scales, - qweight, - scales, - cache_scrub, - ) - torch.accelerator.empty_cache() - gc.collect() - return case_result - - -def _summarize(cases: list[dict[str, Any]]) -> list[dict[str, Any]]: - summary: list[dict[str, Any]] = [] - for case in cases: - rows: list[dict[str, Any]] = case["rows"] - for m in sorted({row["m"] for row in rows}): - for cache_state in sorted({row["cache_state"] for row in rows}): - for mode in ("eager", "graph"): - current = next( - row - for row in rows - if row["m"] == m - and row["mode"] == mode - and row["cache_state"] == cache_state - and row["backend"] == "turbomind_current" - ) - candidates = [ - row - for row in rows - if row["m"] == m - and row["mode"] == mode - and row["cache_state"] == cache_state - and row["backend"] == "qpn8_experimental" - and row["quality_pass"] - and (row["replay_max_abs"] in (None, 0.0)) - ] - best = min(candidates, key=lambda row: row["timing"]["median_us"]) - current_us = current["timing"]["median_us"] - best_us = best["timing"]["median_us"] - summary.append( - { - "case": case["case"], - "shape": case["shape"], - "m": m, - "mode": mode, - "cache_state": cache_state, - "current_us": current_us, - "best_qpn8_us": best_us, - "speedup": current_us / best_us, - "best_config": best["config"], - "best_effective_gbps": best["effective_gbps"], - "best_useful_tflops": best["useful_tflops"], - "best_quality": best["quality"], - } - ) - return summary - - -def main() -> int: - args = _parse_args() - if args.qpn8_library is not None: - torch.ops.load_library(str(args.qpn8_library.resolve())) - if not torch.cuda.is_available(): - raise RuntimeError("CUDA is required") - device = torch.device(args.device) - if torch.cuda.get_device_capability(device) != (7, 0): - raise RuntimeError("This benchmark requires an SM70/V100 device") - if not hasattr(torch.ops._C, "fp8_qpn8_gemm_sm70_out"): - raise RuntimeError("Missing _C::fp8_qpn8_gemm_sm70_out; build this source tree") - if args.tp_size <= 0 or not 0 <= args.tp_rank < args.tp_size: - raise ValueError("invalid TP size/rank") - if any(m < 1 or m > 8 for m in args.m): - raise ValueError("QPN8 operator supports M=1..8") - if args.warmup < 1 or args.iters < 1 or args.trials < 1 or args.cache_scrub_mib < 1: - raise ValueError("warmup, iters, trials, and cache scrub size must be positive") - - torch.accelerator.set_device_index(device.index or 0) - cases = [_run_case(args, case) for case in args.cases] - summary = _summarize(cases) - payload = { - "environment": { - "model": str(args.model), - "device": torch.cuda.get_device_name(device), - "capability": list(torch.cuda.get_device_capability(device)), - "torch": torch.__version__, - "torch_cuda": torch.version.cuda, - "tp_size": args.tp_size, - "tp_rank": args.tp_rank, - "m": args.m, - "warmup": args.warmup, - "iters": args.iters, - "trials": args.trials, - "cache_state": args.cache_state, - "cache_scrub_mib": args.cache_scrub_mib, - "seed": args.seed, - }, - "summary": summary, - "cases": cases, - } - args.json_out.parent.mkdir(parents=True, exist_ok=True) - args.json_out.write_text(json.dumps(payload, indent=2), encoding="utf-8") - print(json.dumps(summary, indent=2)) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/benchmarks/benchmark_sm70_gdn_exactness.py b/benchmarks/benchmark_sm70_gdn_exactness.py index 165b033254..df3acaf627 100644 --- a/benchmarks/benchmark_sm70_gdn_exactness.py +++ b/benchmarks/benchmark_sm70_gdn_exactness.py @@ -165,15 +165,12 @@ def _make_inputs(args: argparse.Namespace) -> dict[str, torch.Tensor]: ) * args.input_scale ).to(dtype) - g = ( - -torch.rand( - (args.batch, args.seqlen, args.v_heads), - device=device, - dtype=torch.float32, - generator=generator, - ) - * args.gate_scale - ) + g = -torch.rand( + (args.batch, args.seqlen, args.v_heads), + device=device, + dtype=torch.float32, + generator=generator, + ) * args.gate_scale beta = torch.sigmoid( torch.randn( (args.batch, args.seqlen, args.v_heads), @@ -214,7 +211,9 @@ def _make_inputs(args: argparse.Namespace) -> dict[str, torch.Tensor]: "initial_state": initial_state.cpu(), } if args.cu_seqlens: - result["cu_seqlens"] = torch.tensor([0, args.seqlen], dtype=torch.int32) + result["cu_seqlens"] = torch.tensor( + [0, args.seqlen], dtype=torch.int32 + ) return result @@ -566,17 +565,25 @@ def _make_spec_decode_cudagraph_inputs( max_query_len = args.num_spec + 1 graph_rows = args.graph_rows or max_query_len if graph_rows < max_query_len: - raise ValueError(f"--graph-rows must be >= num_spec + 1 ({max_query_len})") + raise ValueError( + f"--graph-rows must be >= num_spec + 1 ({max_query_len})" + ) if args.state_slots < max_query_len: - raise ValueError(f"--state-slots must be >= num_spec + 1 ({max_query_len})") + raise ValueError( + f"--state-slots must be >= num_spec + 1 ({max_query_len})" + ) if args.fixed_query_len is not None and not ( 1 <= args.fixed_query_len <= max_query_len ): - raise ValueError(f"--fixed-query-len must be in [1, {max_query_len}]") + raise ValueError( + f"--fixed-query-len must be in [1, {max_query_len}]" + ) if args.fixed_accepted is not None and not ( 1 <= args.fixed_accepted <= max_query_len ): - raise ValueError(f"--fixed-accepted must be in [1, {max_query_len}]") + raise ValueError( + f"--fixed-accepted must be in [1, {max_query_len}]" + ) q_dim = args.k_heads * args.head_k_dim k_dim = args.k_heads * args.head_k_dim @@ -762,27 +769,7 @@ def _load_or_create_spec_decode_cudagraph_inputs( payload = _torch_load(Path(args.input_file)) if "inputs" not in payload: raise ValueError(f"{args.input_file} does not contain saved inputs") - inputs = payload["inputs"] - available_steps = int(inputs["mixed_qkv_seq"].shape[0]) - if args.steps > available_steps: - raise ValueError( - f"--steps={args.steps} exceeds the {available_steps} saved steps" - ) - if args.steps < available_steps: - inputs = dict(inputs) - for name in ( - "mixed_qkv_seq", - "a_seq", - "b_seq", - "query_start_loc_seq", - "state_indices_seq", - "num_accepted_tokens_seq", - "query_lens", - "z_seq", - ): - if name in inputs: - inputs[name] = inputs[name][: args.steps].contiguous() - return inputs + return payload["inputs"] return _make_spec_decode_cudagraph_inputs(args) @@ -865,7 +852,13 @@ def _add_spec_projection_inputs( def _load_or_create_spec_projection_inputs( args: argparse.Namespace, ) -> dict[str, torch.Tensor]: - inputs = _load_or_create_spec_decode_cudagraph_inputs(args) + if args.input_file: + payload = _torch_load(Path(args.input_file)) + if "inputs" not in payload: + raise ValueError(f"{args.input_file} does not contain saved inputs") + inputs = payload["inputs"] + else: + inputs = _make_spec_decode_cudagraph_inputs(args) return _add_spec_projection_inputs(args, inputs) @@ -913,14 +906,12 @@ def _make_sigmoid_gating_inputs(args: argparse.Namespace) -> dict[str, torch.Ten ) * args.gate_scale ) - inputs.update( - { - "a": a.cpu(), - "b": b.cpu(), - "A_log": a_log.cpu(), - "dt_bias": dt_bias.cpu(), - } - ) + inputs.update({ + "a": a.cpu(), + "b": b.cpu(), + "A_log": a_log.cpu(), + "dt_bias": dt_bias.cpu(), + }) if getattr(args, "decode_segments", False): if args.batch != 1: raise ValueError("decode-segments fixture requires --batch 1") @@ -946,7 +937,9 @@ def _make_sigmoid_gating_inputs(args: argparse.Namespace) -> dict[str, torch.Ten dtype=state_dtype, ) inputs["initial_state"] = initial_state.cpu() - inputs["cu_seqlens"] = torch.arange(0, args.seqlen + 1, dtype=torch.int32) + inputs["cu_seqlens"] = torch.arange( + 0, args.seqlen + 1, dtype=torch.int32 + ) inputs["ssm_state_indices"] = torch.arange( 1, args.seqlen + 1, dtype=torch.int32 ) @@ -1112,7 +1105,9 @@ def run_op() -> tuple[torch.Tensor, torch.Tensor | None]: outputs: dict[str, torch.Tensor | None] = { "out": out.detach().cpu(), - "final_state": final_state.detach().cpu() if final_state is not None else None, + "final_state": final_state.detach().cpu() + if final_state is not None + else None, } if args.save_intermediates: pipeline_outputs = _run_chunk_pipeline(cuda_inputs, cu_seqlens, scale) @@ -1166,17 +1161,12 @@ def run_op() -> tuple[torch.Tensor, torch.Tensor | None]: out_path = Path(args.out) out_path.parent.mkdir(parents=True, exist_ok=True) torch.save(payload, out_path) - print( - json.dumps( - { - "out": str(out_path), - "elapsed_s": elapsed_s, - "timing_summary": payload["timing_summary"], - "output_summaries": payload["output_summaries"], - }, - indent=2, - ) - ) + print(json.dumps({ + "out": str(out_path), + "elapsed_s": elapsed_s, + "timing_summary": payload["timing_summary"], + "output_summaries": payload["output_summaries"], + }, indent=2)) return 0 @@ -1279,17 +1269,12 @@ def run_op() -> tuple[torch.Tensor, torch.Tensor]: out_path = Path(args.out) out_path.parent.mkdir(parents=True, exist_ok=True) torch.save(payload, out_path) - print( - json.dumps( - { - "out": str(out_path), - "elapsed_s": elapsed_s, - "timing_summary": payload["timing_summary"], - "output_summaries": payload["output_summaries"], - }, - indent=2, - ) - ) + print(json.dumps({ + "out": str(out_path), + "elapsed_s": elapsed_s, + "timing_summary": payload["timing_summary"], + "output_summaries": payload["output_summaries"], + }, indent=2)) return 0 @@ -1400,17 +1385,12 @@ def run_op() -> tuple[torch.Tensor, torch.Tensor]: out_path = Path(args.out) out_path.parent.mkdir(parents=True, exist_ok=True) torch.save(payload, out_path) - print( - json.dumps( - { - "out": str(out_path), - "elapsed_s": elapsed_s, - "timing_summary": payload["timing_summary"], - "output_summaries": payload["output_summaries"], - }, - indent=2, - ) - ) + print(json.dumps({ + "out": str(out_path), + "elapsed_s": elapsed_s, + "timing_summary": payload["timing_summary"], + "output_summaries": payload["output_summaries"], + }, indent=2)) return 0 @@ -1532,17 +1512,12 @@ def run_op() -> tuple[torch.Tensor, torch.Tensor]: out_path = Path(args.out) out_path.parent.mkdir(parents=True, exist_ok=True) torch.save(payload, out_path) - print( - json.dumps( - { - "out": str(out_path), - "elapsed_s": elapsed_s, - "timing_summary": payload["timing_summary"], - "output_summaries": payload["output_summaries"], - }, - indent=2, - ) - ) + print(json.dumps({ + "out": str(out_path), + "elapsed_s": elapsed_s, + "timing_summary": payload["timing_summary"], + "output_summaries": payload["output_summaries"], + }, indent=2)) return 0 @@ -1876,17 +1851,12 @@ def run_op() -> tuple[torch.Tensor, torch.Tensor]: out_path = Path(args.out) out_path.parent.mkdir(parents=True, exist_ok=True) torch.save(payload, out_path) - print( - json.dumps( - { - "out": str(out_path), - "elapsed_s": elapsed_s, - "timing_summary": payload["timing_summary"], - "output_summaries": payload["output_summaries"], - }, - indent=2, - ) - ) + print(json.dumps({ + "out": str(out_path), + "elapsed_s": elapsed_s, + "timing_summary": payload["timing_summary"], + "output_summaries": payload["output_summaries"], + }, indent=2)) return 0 @@ -1977,7 +1947,9 @@ def run_op() -> tuple[torch.Tensor, torch.Tensor]: "metadata": _metadata(args), "config": { "op": "flashqla_sm70_gdn_decode_full_route", - "route_components": ("fused_mixed_qkv+gating+l2norm+global_state_update"), + "route_components": ( + "fused_mixed_qkv+gating+l2norm+global_state_update" + ), "batch": args.batch, "k_heads": args.k_heads, "v_heads": args.v_heads, @@ -2012,17 +1984,12 @@ def run_op() -> tuple[torch.Tensor, torch.Tensor]: out_path = Path(args.out) out_path.parent.mkdir(parents=True, exist_ok=True) torch.save(payload, out_path) - print( - json.dumps( - { - "out": str(out_path), - "elapsed_s": elapsed_s, - "timing_summary": payload["timing_summary"], - "output_summaries": payload["output_summaries"], - }, - indent=2, - ) - ) + print(json.dumps({ + "out": str(out_path), + "elapsed_s": elapsed_s, + "timing_summary": payload["timing_summary"], + "output_summaries": payload["output_summaries"], + }, indent=2)) return 0 @@ -2149,10 +2116,13 @@ def launch_one( } comparisons = { "out": _compare_tensor(outputs["eager_out"], outputs["graph_out"]), - "final_state": _compare_tensor(outputs["eager_state"], outputs["graph_state"]), + "final_state": _compare_tensor( + outputs["eager_state"], outputs["graph_state"] + ), } strict_ok = all( - item["torch_equal"] and item["max_diff"] == 0.0 for item in comparisons.values() + item["torch_equal"] and item["max_diff"] == 0.0 + for item in comparisons.values() ) saved_inputs = {name: tensor.cpu() for name, tensor in inputs.items()} payload = { @@ -2190,17 +2160,12 @@ def launch_one( out_path = Path(args.out) out_path.parent.mkdir(parents=True, exist_ok=True) torch.save(payload, out_path) - print( - json.dumps( - { - "out": str(out_path), - "strict_ok": strict_ok, - "elapsed_s": elapsed_s, - "comparisons": comparisons, - }, - indent=2, - ) - ) + print(json.dumps({ + "out": str(out_path), + "strict_ok": strict_ok, + "elapsed_s": elapsed_s, + "comparisons": comparisons, + }, indent=2)) return 0 if strict_ok else 1 @@ -2282,10 +2247,13 @@ def launch_one( } comparisons = { "out": _compare_tensor(outputs["eager_out"], outputs["graph_out"]), - "final_state": _compare_tensor(outputs["eager_state"], outputs["graph_state"]), + "final_state": _compare_tensor( + outputs["eager_state"], outputs["graph_state"] + ), } strict_ok = all( - item["torch_equal"] and item["max_diff"] == 0.0 for item in comparisons.values() + item["torch_equal"] and item["max_diff"] == 0.0 + for item in comparisons.values() ) saved_inputs = {name: tensor.cpu() for name, tensor in inputs.items()} payload = { @@ -2320,17 +2288,12 @@ def launch_one( out_path = Path(args.out) out_path.parent.mkdir(parents=True, exist_ok=True) torch.save(payload, out_path) - print( - json.dumps( - { - "out": str(out_path), - "strict_ok": strict_ok, - "elapsed_s": elapsed_s, - "comparisons": comparisons, - }, - indent=2, - ) - ) + print(json.dumps({ + "out": str(out_path), + "strict_ok": strict_ok, + "elapsed_s": elapsed_s, + "comparisons": comparisons, + }, indent=2)) return 0 if strict_ok else 1 @@ -2401,9 +2364,6 @@ def launch_one( a, b, cuda_inputs["dt_bias"], - beta_dtype=( - torch.float32 if args.beta_output_dtype == "float32" else b.dtype - ), ) core_out, _ = fused_recurrent_gated_delta_rule( q=query, @@ -2511,7 +2471,8 @@ def launch_one( query_lens = inputs["query_lens"].to(torch.int64) live_token_mask = ( - torch.arange(max_query_len, dtype=torch.int64)[None, :] < query_lens[:, None] + torch.arange(max_query_len, dtype=torch.int64)[None, :] + < query_lens[:, None] ) outputs: dict[str, torch.Tensor | None] = { "eager_live_out": eager_out.detach().cpu()[live_token_mask], @@ -2568,7 +2529,6 @@ def launch_one( "conv_width": args.conv_width, "dtype": args.dtype, "state_dtype": args.state_dtype, - "beta_output_dtype": args.beta_output_dtype, "seed": args.seed, "scale": scale, "silu": args.silu, @@ -2597,18 +2557,13 @@ def launch_one( out_path = Path(args.out) out_path.parent.mkdir(parents=True, exist_ok=True) torch.save(payload, out_path) - print( - json.dumps( - { - "out": str(out_path), - "strict_ok": strict_ok, - "elapsed_s": elapsed_s, - "config": payload["config"], - "comparisons": comparisons, - }, - indent=2, - ) - ) + print(json.dumps({ + "out": str(out_path), + "strict_ok": strict_ok, + "elapsed_s": elapsed_s, + "config": payload["config"], + "comparisons": comparisons, + }, indent=2)) return 0 if strict_ok else 1 @@ -2695,9 +2650,6 @@ def launch_split( a, b, cuda_inputs["dt_bias"], - beta_dtype=( - torch.float32 if args.beta_output_dtype == "float32" else b.dtype - ), ) core_out, _ = fused_recurrent_gated_delta_rule( q=query, @@ -2752,9 +2704,7 @@ def launch_mixed_qkv( ) out.copy_(core_out.squeeze(0)) - def run_eager_sequence( - launch_fn: Any, - ) -> tuple[ + def run_eager_sequence(launch_fn: Any) -> tuple[ torch.Tensor, torch.Tensor, torch.Tensor, @@ -2854,7 +2804,8 @@ def time_graph_sequence(launch_fn: Any) -> float: query_lens = inputs["query_lens"].to(torch.int64) live_token_mask = ( - torch.arange(max_query_len, dtype=torch.int64)[None, :] < query_lens[:, None] + torch.arange(max_query_len, dtype=torch.int64)[None, :] + < query_lens[:, None] ) outputs: dict[str, torch.Tensor] = { "split_live_out": split_out.detach().cpu()[live_token_mask], @@ -2920,7 +2871,6 @@ def time_graph_sequence(launch_fn: Any) -> float: "conv_width": args.conv_width, "dtype": args.dtype, "state_dtype": args.state_dtype, - "beta_output_dtype": args.beta_output_dtype, "seed": args.seed, "scale": scale, "silu": args.silu, @@ -2948,18 +2898,13 @@ def time_graph_sequence(launch_fn: Any) -> float: out_path = Path(args.out) out_path.parent.mkdir(parents=True, exist_ok=True) torch.save(payload, out_path) - print( - json.dumps( - { - "out": str(out_path), - "strict_ok": strict_ok, - "timing": timing, - "config": payload["config"], - "comparisons": comparisons, - }, - indent=2, - ) - ) + print(json.dumps({ + "out": str(out_path), + "strict_ok": strict_ok, + "timing": timing, + "config": payload["config"], + "comparisons": comparisons, + }, indent=2)) return 0 if strict_ok else 1 @@ -3224,7 +3169,8 @@ def call_spec_commit( query_lens = inputs["query_lens"].to(torch.int64) live_token_mask = ( - torch.arange(max_query_len, dtype=torch.int64)[None, :] < query_lens[:, None] + torch.arange(max_query_len, dtype=torch.int64)[None, :] + < query_lens[:, None] ) outputs: dict[str, torch.Tensor | None] = { "standard_live_out": standard_out.detach().cpu()[live_token_mask], @@ -3340,18 +3286,13 @@ def call_spec_commit( out_path = Path(args.out) out_path.parent.mkdir(parents=True, exist_ok=True) torch.save(payload, out_path) - print( - json.dumps( - { - "out": str(out_path), - "strict_ok": strict_ok, - "elapsed_s": elapsed_s, - "config": payload["config"], - "comparisons": comparisons, - }, - indent=2, - ) - ) + print(json.dumps({ + "out": str(out_path), + "strict_ok": strict_ok, + "elapsed_s": elapsed_s, + "config": payload["config"], + "comparisons": comparisons, + }, indent=2)) return 0 if strict_ok else 1 @@ -3855,7 +3796,8 @@ def compile_direct_step( query_lens = inputs["query_lens"].to(torch.int64) live_token_mask = ( - torch.arange(max_query_len, dtype=torch.int64)[None, :] < query_lens[:, None] + torch.arange(max_query_len, dtype=torch.int64)[None, :] + < query_lens[:, None] ) outputs: dict[str, torch.Tensor | None] = { "standard_live_output": standard_output.detach().cpu()[live_token_mask], @@ -4044,15 +3986,12 @@ def compile_direct_step( "direct_compile_ssm_state", ] ) - strict_ok = ( - all( - comparisons[name] is not None - and comparisons[name]["torch_equal"] - and comparisons[name]["max_diff"] == 0.0 - for name in strict_keys - ) - and compile_step_error is None - ) + strict_ok = all( + comparisons[name] is not None + and comparisons[name]["torch_equal"] + and comparisons[name]["max_diff"] == 0.0 + for name in strict_keys + ) and compile_step_error is None accepted_hist = torch.bincount( inputs["num_accepted_tokens_seq"][:, 0].to(torch.int64), minlength=max_query_len + 1, @@ -4111,18 +4050,13 @@ def compile_direct_step( out_path = Path(args.out) out_path.parent.mkdir(parents=True, exist_ok=True) torch.save(payload, out_path) - print( - json.dumps( - { - "out": str(out_path), - "strict_ok": strict_ok, - "elapsed_s": elapsed_s, - "config": payload["config"], - "comparisons": comparisons, - }, - indent=2, - ) - ) + print(json.dumps({ + "out": str(out_path), + "strict_ok": strict_ok, + "elapsed_s": elapsed_s, + "config": payload["config"], + "comparisons": comparisons, + }, indent=2)) return 0 if strict_ok else 1 @@ -4297,7 +4231,9 @@ def make_parser() -> argparse.ArgumentParser: sigmoid.add_argument("--v-heads", type=int, default=8) sigmoid.add_argument("--head-k-dim", type=int, default=128) sigmoid.add_argument("--head-v-dim", type=int, default=128) - sigmoid.add_argument("--dtype", choices=("float16", "bfloat16"), default="float16") + sigmoid.add_argument( + "--dtype", choices=("float16", "bfloat16"), default="float16" + ) sigmoid.add_argument( "--state-dtype", choices=("same", "float16", "bfloat16", "float32"), @@ -4358,7 +4294,9 @@ def make_parser() -> argparse.ArgumentParser: model_mixed = subparsers.add_parser("run-model-mixed-qkv-route") model_mixed.add_argument("--out", required=True) model_mixed.add_argument("--input-file") - model_mixed.add_argument("--case-name", default="gdn_model_mixed_qkv_route_smoke") + model_mixed.add_argument( + "--case-name", default="gdn_model_mixed_qkv_route_smoke" + ) model_mixed.add_argument("--batch", type=int, default=1) model_mixed.add_argument("--seqlen", type=int, default=16) model_mixed.add_argument("--k-heads", type=int, default=4) @@ -4476,12 +4414,16 @@ def make_parser() -> argparse.ArgumentParser: ) packed_cg.add_argument("--random-initial-state", action="store_true") packed_cg.add_argument("--capture-null-indices", action="store_true") - packed_cg.set_defaults(func=run_packed_recurrent_cudagraph_case, l2norm_inputs=True) + packed_cg.set_defaults( + func=run_packed_recurrent_cudagraph_case, l2norm_inputs=True + ) conv_cg = subparsers.add_parser("run-conv-update-cudagraph") conv_cg.add_argument("--out", required=True) conv_cg.add_argument("--input-file") - conv_cg.add_argument("--case-name", default="gdn_conv_update_cudagraph_exactness") + conv_cg.add_argument( + "--case-name", default="gdn_conv_update_cudagraph_exactness" + ) conv_cg.add_argument("--batch", type=int, default=2) conv_cg.add_argument("--seqlen", type=int, default=1024) conv_cg.add_argument("--k-heads", type=int, default=4) @@ -4506,7 +4448,9 @@ def make_parser() -> argparse.ArgumentParser: spec_cg = subparsers.add_parser("run-spec-decode-cudagraph") spec_cg.add_argument("--out", required=True) spec_cg.add_argument("--input-file") - spec_cg.add_argument("--case-name", default="gdn_spec_decode_cudagraph_exactness") + spec_cg.add_argument( + "--case-name", default="gdn_spec_decode_cudagraph_exactness" + ) spec_cg.add_argument("--steps", type=int, default=512) spec_cg.add_argument("--num-spec", type=int, default=4) spec_cg.add_argument("--graph-rows", type=int) @@ -4518,12 +4462,6 @@ def make_parser() -> argparse.ArgumentParser: spec_cg.add_argument("--head-v-dim", type=int, default=128) spec_cg.add_argument("--conv-width", type=int, default=4) spec_cg.add_argument("--dtype", choices=("float16", "bfloat16"), default="float16") - spec_cg.add_argument( - "--beta-output-dtype", - choices=("input", "float32"), - default="float32", - help="Materialize split-verifier beta in the input dtype or FP32.", - ) spec_cg.add_argument( "--state-dtype", choices=("same", "float16", "bfloat16", "float32"), @@ -4551,7 +4489,9 @@ def make_parser() -> argparse.ArgumentParser: spec_mixed = subparsers.add_parser("run-spec-mixed-qkv-candidate") spec_mixed.add_argument("--out", required=True) spec_mixed.add_argument("--input-file") - spec_mixed.add_argument("--case-name", default="gdn_spec_mixed_qkv_candidate") + spec_mixed.add_argument( + "--case-name", default="gdn_spec_mixed_qkv_candidate" + ) spec_mixed.add_argument("--steps", type=int, default=512) spec_mixed.add_argument("--num-spec", type=int, default=4) spec_mixed.add_argument("--graph-rows", type=int) @@ -4565,12 +4505,6 @@ def make_parser() -> argparse.ArgumentParser: spec_mixed.add_argument( "--dtype", choices=("float16", "bfloat16"), default="float16" ) - spec_mixed.add_argument( - "--beta-output-dtype", - choices=("input", "float32"), - default="float32", - help="Materialize split-verifier beta in the input dtype or FP32.", - ) spec_mixed.add_argument( "--state-dtype", choices=("same", "float16", "bfloat16", "float32"), @@ -4598,7 +4532,9 @@ def make_parser() -> argparse.ArgumentParser: spec_op = subparsers.add_parser("run-spec-commit-vs-standard") spec_op.add_argument("--out", required=True) spec_op.add_argument("--input-file") - spec_op.add_argument("--case-name", default="gdn_spec_commit_vs_standard_exactness") + spec_op.add_argument( + "--case-name", default="gdn_spec_commit_vs_standard_exactness" + ) spec_op.add_argument("--steps", type=int, default=512) spec_op.add_argument("--num-spec", type=int, default=4) spec_op.add_argument("--graph-rows", type=int) diff --git a/benchmarks/benchmark_sm70_hmma_c2a_mincomm.py b/benchmarks/benchmark_sm70_hmma_c2a_mincomm.py index 1e9baafc70..8eed68503a 100644 --- a/benchmarks/benchmark_sm70_hmma_c2a_mincomm.py +++ b/benchmarks/benchmark_sm70_hmma_c2a_mincomm.py @@ -7,14 +7,13 @@ 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 d75ddb1397..05ec0612f7 100644 --- a/benchmarks/benchmark_sm70_instruction_latency.py +++ b/benchmarks/benchmark_sm70_instruction_latency.py @@ -7,14 +7,13 @@ 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 4e846e1c50..a34e3e0fc1 100644 --- a/benchmarks/benchmark_sm70_lm_head_top1.py +++ b/benchmarks/benchmark_sm70_lm_head_top1.py @@ -111,7 +111,9 @@ 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()) @@ -119,7 +121,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 @@ -187,21 +189,15 @@ 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 f57a802ce3..0a9c06a273 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..ae26060075 100644 --- a/benchmarks/benchmark_sm70_model_tokens.py +++ b/benchmarks/benchmark_sm70_model_tokens.py @@ -20,9 +20,7 @@ from typing import Any SOURCE_ROOT = Path(__file__).resolve().parents[1] -FLASH_V100_ROOT = Path( - os.environ.get("SM70_FLASH_V100_ROOT", SOURCE_ROOT / "flash-attention-v100") -) +FLASH_V100_ROOT = SOURCE_ROOT / "flash-attention-v100" for source_path in (FLASH_V100_ROOT, SOURCE_ROOT): source_path_str = str(source_path) if source_path_str not in sys.path: @@ -502,10 +500,7 @@ def _sm70_tune_policy() -> dict[str, Any]: } -def _sm70_turbomind_policy( - max_num_seqs: int | None = None, - speculative_enabled: bool = False, -) -> dict[str, Any]: +def _sm70_turbomind_policy() -> dict[str, Any]: awq_turbomind = _env_bool("VLLM_SM70_AWQ_TURBOMIND", True) awq_tune_raw = os.environ.get("VLLM_SM70_AWQ_TUNE_SMALL_SHAPES") awq_tune0_pinned = awq_tune_raw == "0" @@ -526,13 +521,7 @@ def _sm70_turbomind_policy( "VLLM_SM70_FP8_PREFILL_EXACT_DENSE", True, ) - fp8_qpn8 = _env_bool("VLLM_SM70_FP8_QPN8", False) - fp8_prefill_visible_dense_mm = _env_bool( - "VLLM_SM70_FP8_PREFILL_VISIBLE_DENSE_MM", - False, - ) nvfp4_turbomind = _env_bool("VLLM_SM70_NVFP4_TURBOMIND", False) - nvfp4_moe_grouped_prefill = _env_bool("VLLM_SM70_NVFP4_MOE_GROUPED_PREFILL", True) mxfp4_turbomind = _env_bool("VLLM_SM70_MXFP4_TURBOMIND", False) fp8_dequant_fallback = _env_bool("VLLM_SM70_FP8_DEQUANT_FALLBACK", True) fp8_moe_dequant_fallback = _env_bool( @@ -543,10 +532,6 @@ def _sm70_turbomind_policy( "VLLM_SM70_UNQUANTIZED_MOE_0DOT3_CONFIG", True, ) - qwen36_mtp_moe_tuned_config = _env_bool( - "VLLM_SM70_QWEN36_MTP_MOE_TUNED_CONFIG", - True, - ) awq_warmup = _env_bool("VLLM_SM70_AWQ_WARMUP", True) awq_moe_disable = _env_bool("VLLM_SM70_AWQ_MOE_DISABLE", False) awq_moe_batched = _env_bool("VLLM_SM70_AWQ_MOE_BATCHED_GEMM", True) @@ -592,25 +577,8 @@ def _sm70_turbomind_policy( "VLLM_SM70_FP8_PREFILL_EXACT_DENSE" ), "fp8_prefill_exact_dense_effective": fp8_prefill_exact_dense, - "VLLM_SM70_FP8_QPN8": os.environ.get("VLLM_SM70_FP8_QPN8"), - "fp8_qpn8_effective": fp8_qpn8, - "fp8_qpn8_native_decode_max_m": 8, - "fp8_qpn8_configured_max_num_seqs": max_num_seqs, - "fp8_qpn8_concurrency_contract": (max_num_seqs is None or max_num_seqs <= 8), - "fp8_qpn8_no_mtp_contract": not speculative_enabled, - "fp8_qpn8_large_m_fallback": "bounded_fp16_weight_workspace", - "fp8_qpn8_large_m_gated_temporary": "M_x_8704_fp16", - "VLLM_SM70_FP8_QPN8_LIBRARY": os.environ.get("VLLM_SM70_FP8_QPN8_LIBRARY"), - "VLLM_SM70_FP8_PREFILL_VISIBLE_DENSE_MM": os.environ.get( - "VLLM_SM70_FP8_PREFILL_VISIBLE_DENSE_MM" - ), - "fp8_prefill_visible_dense_mm_effective": fp8_prefill_visible_dense_mm, "VLLM_SM70_NVFP4_TURBOMIND": os.environ.get("VLLM_SM70_NVFP4_TURBOMIND"), "nvfp4_turbomind_effective": nvfp4_turbomind, - "VLLM_SM70_NVFP4_MOE_GROUPED_PREFILL": os.environ.get( - "VLLM_SM70_NVFP4_MOE_GROUPED_PREFILL" - ), - "nvfp4_moe_grouped_prefill_effective": nvfp4_moe_grouped_prefill, "VLLM_SM70_MXFP4_TURBOMIND": os.environ.get("VLLM_SM70_MXFP4_TURBOMIND"), "mxfp4_turbomind_effective": mxfp4_turbomind, "VLLM_SM70_FP8_DEQUANT_FALLBACK": os.environ.get( @@ -625,10 +593,6 @@ def _sm70_turbomind_policy( "VLLM_SM70_UNQUANTIZED_MOE_0DOT3_CONFIG" ), "unquantized_moe_0dot3_config_effective": (unquantized_moe_0dot3_config), - "VLLM_SM70_QWEN36_MTP_MOE_TUNED_CONFIG": os.environ.get( - "VLLM_SM70_QWEN36_MTP_MOE_TUNED_CONFIG" - ), - "qwen36_mtp_moe_tuned_config_effective": qwen36_mtp_moe_tuned_config, "VLLM_SM70_AWQ_WARMUP": os.environ.get("VLLM_SM70_AWQ_WARMUP"), "awq_warmup_effective": awq_warmup, "VLLM_SM70_AWQ_WARMUP_MAX_M": os.environ.get("VLLM_SM70_AWQ_WARMUP_MAX_M"), @@ -801,15 +765,6 @@ def _sm70_turbomind_policy( and not fp8_moe_batched_w2_dispatch and not f16_dense ), - "accepted_qwen38_fp8_qpn8_default_policy": ( - fp8_turbomind - and fp8_dequant_fallback - and fp8_dense_gated_silu - and fp8_prefill_exact_dense - and fp8_qpn8 - and (max_num_seqs is None or max_num_seqs <= 8) - and not speculative_enabled - ), "accepted_awq_moe_default_policy": ( awq_turbomind and awq_moe_safe_default_selector @@ -855,9 +810,7 @@ def _sm70_turbomind_policy( "route_hit_oracle": ( "Accepted dense route-hit requires logs such as " "`SM70 AWQ TurboMind dense path enabled`, " - "`SM70 FP8 TurboMind W8A16 dense path enabled`, and for the " - "Qwen3.8-27B TP4 QPN8 default " - "`Memory-neutral SM70 FP8 QPN8 path enabled`; when warmup " + "`SM70 FP8 TurboMind W8A16 dense path enabled`, and when warmup " "is relevant `SM70 AWQ warmup finished`. AWQ MoE production " "throughput evidence must keep the default fast route: batched " "GEMM enabled, legacy single-token compact enabled, and no " @@ -932,10 +885,6 @@ def _sm70_attention_policy(kv_cache_dtype: Any) -> dict[str, Any]: "VLLM_FLASH_V100_XQA_E5M2_PAIR_LOAD", True, ) - e5m2_batch_wide_load = _env_bool( - "VLLM_FLASH_V100_XQA_E5M2_BATCH_WIDE_LOAD", - True, - ) decode_fp8_xqa_min_seq_len = _env_int( "VLLM_FLASH_V100_DECODE_FP8_XQA_MIN_SEQ_LEN", 16384, @@ -1047,10 +996,6 @@ def _sm70_attention_policy(kv_cache_dtype: Any) -> dict[str, Any]: "VLLM_FLASH_V100_XQA_E5M2_PAIR_LOAD" ), "e5m2_pair_load_effective": e5m2_pair_load, - "VLLM_FLASH_V100_XQA_E5M2_BATCH_WIDE_LOAD": os.environ.get( - "VLLM_FLASH_V100_XQA_E5M2_BATCH_WIDE_LOAD" - ), - "e5m2_batch_wide_load_effective": e5m2_batch_wide_load, "VLLM_FLASH_V100_DECODE_FP8_XQA_MIN_SEQ_LEN": os.environ.get( "VLLM_FLASH_V100_DECODE_FP8_XQA_MIN_SEQ_LEN" ), @@ -1643,14 +1588,11 @@ def _dump(args: argparse.Namespace) -> int: ) prompts = _load_prompts(args) - if args.repeat_count < 1: - raise ValueError("--repeat-count must be positive") - if args.repeat_count > 1 and len(prompts) != 1: - raise ValueError("--repeat-count greater than one requires one prompt") - if args.repeat_count > 1 and ( - args.cuda_profiler_capture_generate or capture_decode_after_prefix_warmup - ): - raise ValueError("CUDA-profiler capture does not support repeated requests") + if capture_decode_after_prefix_warmup and len(prompts) != 1: + raise ValueError( + "--cuda-profiler-capture-decode-after-prefix-warmup requires " + "exactly one prompt" + ) sampling_params = SamplingParams( max_tokens=args.max_tokens, temperature=args.temperature, @@ -1712,30 +1654,21 @@ def _dump(args: argparse.Namespace) -> int: torch.accelerator.synchronize() torch.cuda.cudart().cudaProfilerStart() - generate_seconds_by_repeat: list[float] = [] - outputs = [] + start = time.perf_counter() try: - for repeat_index in range(args.repeat_count): - if ( - repeat_index - and args.reset_prefix_cache_between_repeats - and not llm.reset_prefix_cache() - ): - raise RuntimeError("Failed to reset the idle prefix cache") - generate_start = time.perf_counter() - if args.sequential_prompts: - for prompt in prompts: - outputs.extend(llm.generate([prompt], sampling_params)) - else: - outputs.extend(llm.generate(prompts, sampling_params)) - generate_seconds_by_repeat.append(time.perf_counter() - generate_start) + if args.sequential_prompts: + outputs = [] + for prompt in prompts: + outputs.extend(llm.generate([prompt], sampling_params)) + else: + outputs = llm.generate(prompts, sampling_params) finally: if args.cuda_profiler_capture_generate or capture_decode_after_prefix_warmup: import torch torch.accelerator.synchronize() torch.cuda.cudart().cudaProfilerStop() - generate_seconds = sum(generate_seconds_by_repeat) + generate_seconds = time.perf_counter() - start metrics_snapshot = _metric_snapshot(llm) import torch @@ -1814,10 +1747,7 @@ def _dump(args: argparse.Namespace) -> int: ], "env": _tracked_env(), "sm70_tune_policy": _sm70_tune_policy(), - "sm70_turbomind_policy": _sm70_turbomind_policy( - int(llm_kwargs["max_num_seqs"]), - speculative_enabled=llm_kwargs.get("speculative_config") is not None, - ), + "sm70_turbomind_policy": _sm70_turbomind_policy(), "sm70_attention_policy": _sm70_attention_policy( llm_kwargs.get("kv_cache_dtype") ), @@ -1848,9 +1778,6 @@ def _dump(args: argparse.Namespace) -> int: "sequential_prompts": args.sequential_prompts, "load_seconds": load_seconds, "generate_seconds": generate_seconds, - "generate_seconds_by_repeat": generate_seconds_by_repeat, - "repeat_count": args.repeat_count, - "reset_prefix_cache_between_repeats": (args.reset_prefix_cache_between_repeats), "total_output_tokens": total_output_tokens, "records": records, } @@ -2600,14 +2527,12 @@ def _parse_args() -> argparse.Namespace: "--cuda-profiler-capture-decode-after-prefix-warmup", action="store_true", help=( - "Prime the prompts with prefix caching while the CUDA profiler is " + "Prime one prompt with prefix caching while the CUDA profiler is " "off, then capture only the second matching llm.generate call. " - "Requires non-sequential prompts and CUDA graphs." + "Requires one non-sequential prompt and CUDA graphs." ), ) parser.add_argument("--max-tokens", type=int, default=16) - parser.add_argument("--repeat-count", type=int, default=1) - parser.add_argument("--reset-prefix-cache-between-repeats", action="store_true") parser.add_argument("--temperature", type=float, default=0.0) parser.add_argument("--top-p", type=float, default=1.0) parser.add_argument("--top-k", type=int, default=-1) diff --git a/benchmarks/benchmark_sm70_mxfp4_moe_active_experts.py b/benchmarks/benchmark_sm70_mxfp4_moe_active_experts.py index b79a169e61..b58ba25399 100644 --- a/benchmarks/benchmark_sm70_mxfp4_moe_active_experts.py +++ b/benchmarks/benchmark_sm70_mxfp4_moe_active_experts.py @@ -13,7 +13,6 @@ import json import os from dataclasses import dataclass -from pathlib import Path import torch @@ -400,7 +399,6 @@ def make_buffers() -> dict[str, torch.Tensor]: generic = make_buffers() direct = make_buffers() - direct_order = make_buffers() def generic_call() -> None: generic["output"].zero_() @@ -500,61 +498,11 @@ def direct_call() -> None: direct["output"], ) - def direct_order_call() -> None: - route_ids = topk_ids.view(-1) - sm70_ops.mxfp4_moe_dense_stage_sm70_out( - direct_order["gate_up"], - x, - direct_order["compact_offsets"], - route_ids, - w13_ptrs_w, - w13_ptrs_s, - top_k, - 4096, - 512, - 32, - ) - torch.ops._C.silu_and_mul_with_clamp( - direct_order["intermediate"], direct_order["gate_up"], 10.0 - ) - sm70_ops.mxfp4_moe_dense_stage_sm70_out( - direct_order["sorted_output"], - direct_order["intermediate"], - direct_order["compact_offsets"], - route_ids, - w2_ptrs_w, - w2_ptrs_s, - top_k, - 256, - 4096, - 32, - ) - sm70_ops.awq_moe_single_token_weighted_reduce_out( - direct_order["sorted_output"], - topk_weights, - direct_order["token_expert_indices"], - direct_order["output"], - top_k, - 4096, - ) - generic_call() direct_call() - direct_order_call() torch.accelerator.synchronize() initial_equal = torch.equal(generic["output"], direct["output"]) initial_max_abs = float((generic["output"] - direct["output"]).abs().max().item()) - direct_order_initial_equal = torch.equal(direct["output"], direct_order["output"]) - direct_order_initial_max_abs = float( - (direct["output"] - direct_order["output"]).abs().max().item() - ) - initial_sorted_src = torch.argsort(topk_ids.view(-1), stable=True) - direct_order_stage_parity = { - f"{name}_equal": torch.equal( - direct[name], direct_order[name].index_select(0, initial_sorted_src) - ) - for name in ("gate_up", "intermediate", "sorted_output") - } stage_parity = { "gate_up_equal": torch.equal(generic["gate_up"], direct["gate_up"]), "gate_up_max_abs": float( @@ -579,52 +527,28 @@ def direct_order_call() -> None: generic_graph = _capture(generic_call) direct_graph = _capture(direct_call) - direct_order_graph = _capture(direct_order_call) generic_ms = _time_graph(generic_graph, repeats) direct_ms = _time_graph(direct_graph, repeats) - direct_order_ms = _time_graph(direct_order_graph, repeats) topk_ids.copy_(torch.tensor([route_b], dtype=torch.int32, device=device)) generic_graph.replay() direct_graph.replay() - direct_order_graph.replay() torch.accelerator.synchronize() replay_equal = torch.equal(generic["output"], direct["output"]) replay_max_abs = float((generic["output"] - direct["output"]).abs().max().item()) - direct_order_replay_equal = torch.equal(direct["output"], direct_order["output"]) - direct_order_replay_max_abs = float( - (direct["output"] - direct_order["output"]).abs().max().item() - ) result = { "generic_graph_ms": generic_ms, "direct_graph_ms": direct_ms, - "direct_order_graph_ms": direct_order_ms, "speedup": generic_ms / direct_ms, "projected_savings_ms_per_token": (generic_ms - direct_ms) * 43, - "direct_order_speedup_vs_direct": direct_ms / direct_order_ms, - "direct_order_projected_savings_ms_per_token": (direct_ms - direct_order_ms) - * 43, "initial_bitwise_equal": initial_equal, "initial_max_abs": initial_max_abs, "dynamic_replay_bitwise_equal": replay_equal, "dynamic_replay_max_abs": replay_max_abs, - "direct_order_initial_bitwise_equal": direct_order_initial_equal, - "direct_order_initial_max_abs": direct_order_initial_max_abs, - "direct_order_dynamic_replay_bitwise_equal": direct_order_replay_equal, - "direct_order_dynamic_replay_max_abs": direct_order_replay_max_abs, "initial_stage_parity": stage_parity, - "direct_order_initial_sorted_stage_parity": direct_order_stage_parity, } - if not all( - ( - initial_equal, - replay_equal, - direct_order_initial_equal, - direct_order_replay_equal, - *direct_order_stage_parity.values(), - ) - ): + if not initial_equal or not replay_equal: raise RuntimeError(f"MXFP4 direct top-6 correctness gate failed: {result}") return result @@ -633,19 +557,16 @@ def benchmark_verifier_m8_pipeline( *, num_experts: int, top_k: int, - verifier_tokens: int, repeats: int, seed: int, route_case: str = "mixed", input_scale: float = 0.01, require_grouped_bitwise: bool = True, ) -> dict[str, object]: - """Compare dense, active-loop, and grouped verifier pipelines.""" - num_tokens = int(verifier_tokens) + """Compare dense-256, active-48 loop, and grouped verifier pipelines.""" + num_tokens = 8 if num_experts != 256 or top_k != 6: raise ValueError("The verifier benchmark requires 256 experts/top-k=6") - if not 2 <= num_tokens <= 8: - raise ValueError("The verifier benchmark requires M in [2, 8]") torch.manual_seed(seed) device = torch.device("cuda") @@ -749,22 +670,11 @@ def make_buffers() -> dict[str, torch.Tensor]: dense = make_buffers() active = make_buffers() grouped = make_buffers() - expert_grouped = make_buffers() def pipeline_call( - buffers: dict[str, torch.Tensor], - *, - active_only: bool, - grouped_verifier: bool, - expert_rows: bool = False, + buffers: dict[str, torch.Tensor], *, active_only: bool, grouped_m8: bool ) -> None: - os.environ["VLLM_SM70_MXFP4_MOE_GROUPED_M8"] = "0" - os.environ["VLLM_SM70_MXFP4_MOE_GROUPED_VERIFIER"] = ( - "1" if grouped_verifier else "0" - ) - os.environ["VLLM_SM70_MXFP4_MOE_GROUPED_M8_EXPERT_ROWS"] = ( - "1" if expert_rows else "0" - ) + os.environ["VLLM_SM70_MXFP4_MOE_GROUPED_M8"] = "1" if grouped_m8 else "0" buffers["output"].zero_() buffers["topk_ids_i32"].copy_(topk_ids) buffers["permuted_idx"].fill_(total_slots) @@ -788,16 +698,8 @@ def pipeline_call( buffers["expert_offsets"].copy_(buffers["expert_offsets64"]) if active_only: stage_offsets = buffers["compact_offsets"] - if grouped_verifier: - if expert_rows: - _compact_mxfp4_active_experts( - buffers["permuted_experts_id"], - buffers["compact_offsets"], - buffers["active_expert_ids"], - ) - stage_expert_ids = buffers["active_expert_ids"] - else: - stage_expert_ids = buffers["permuted_experts_id"] + if grouped_m8: + stage_expert_ids = buffers["permuted_experts_id"] else: _compact_mxfp4_active_experts( buffers["permuted_experts_id"], @@ -847,32 +749,20 @@ def pipeline_call( ) def dense_call() -> None: - pipeline_call(dense, active_only=False, grouped_verifier=False) + pipeline_call(dense, active_only=False, grouped_m8=False) def active_call() -> None: - pipeline_call(active, active_only=True, grouped_verifier=False) + pipeline_call(active, active_only=True, grouped_m8=False) def grouped_call() -> None: - pipeline_call(grouped, active_only=True, grouped_verifier=True) - - def expert_grouped_call() -> None: - pipeline_call( - expert_grouped, - active_only=True, - grouped_verifier=True, - expert_rows=True, - ) + pipeline_call(grouped, active_only=True, grouped_m8=True) dense_call() active_call() grouped_call() - expert_grouped_call() torch.accelerator.synchronize() initial_equal = torch.equal(dense["output"], active["output"]) grouped_initial_equal = torch.equal(dense["output"], grouped["output"]) - expert_grouped_initial_equal = torch.equal( - dense["output"], expert_grouped["output"] - ) stage_parity = { name: torch.equal(dense[name], active[name]) for name in ("gate_up", "intermediate", "sorted_output", "output") @@ -881,26 +771,14 @@ def expert_grouped_call() -> None: name: torch.equal(dense[name], grouped[name]) for name in ("gate_up", "intermediate", "sorted_output", "output") } - expert_grouped_stage_parity = { - name: torch.equal(dense[name], expert_grouped[name]) - for name in ("gate_up", "intermediate", "sorted_output", "output") - } grouped_stage_max_abs = { name: float((dense[name] - grouped[name]).abs().max().item()) for name in ("gate_up", "intermediate", "sorted_output", "output") } - expert_grouped_stage_max_abs = { - name: float((dense[name] - expert_grouped[name]).abs().max().item()) - for name in ("gate_up", "intermediate", "sorted_output", "output") - } grouped_stages_finite = all( bool(torch.isfinite(grouped[name]).all().item()) for name in ("gate_up", "intermediate", "sorted_output", "output") ) - expert_grouped_stages_finite = all( - bool(torch.isfinite(expert_grouped[name]).all().item()) - for name in ("gate_up", "intermediate", "sorted_output", "output") - ) expected_active_ids = sorted({expert for row in route_a for expert in row}) actual_active_ids = active["active_expert_ids"][: len(expected_active_ids)] compact_ids_match = actual_active_ids.cpu().tolist() == expected_active_ids @@ -908,17 +786,14 @@ def expert_grouped_call() -> None: dense_graph = _capture(dense_call) active_graph = _capture(active_call) grouped_graph = _capture(grouped_call) - expert_grouped_graph = _capture(expert_grouped_call) dense_ms = _time_graph(dense_graph, repeats) active_ms = _time_graph(active_graph, repeats) grouped_ms = _time_graph(grouped_graph, repeats) - expert_grouped_ms = _time_graph(expert_grouped_graph, repeats) topk_ids.copy_(torch.tensor(route_b, dtype=torch.int32, device=device)) dense_graph.replay() active_graph.replay() grouped_graph.replay() - expert_grouped_graph.replay() torch.accelerator.synchronize() replay_equal = torch.equal(dense["output"], active["output"]) replay_max_abs = float((dense["output"] - active["output"]).abs().max().item()) @@ -926,21 +801,12 @@ def expert_grouped_call() -> None: grouped_replay_max_abs = float( (dense["output"] - grouped["output"]).abs().max().item() ) - expert_grouped_replay_equal = torch.equal(dense["output"], expert_grouped["output"]) - expert_grouped_replay_max_abs = float( - (dense["output"] - expert_grouped["output"]).abs().max().item() - ) route_b_output = grouped["output"].clone() - expert_grouped_route_b_output = expert_grouped["output"].clone() topk_ids.copy_(torch.tensor(route_a, dtype=torch.int32, device=device)) grouped_graph.replay() - expert_grouped_graph.replay() torch.accelerator.synchronize() route_changes_output = not torch.equal(grouped["output"], route_b_output) - expert_grouped_route_changes_output = not torch.equal( - expert_grouped["output"], expert_grouped_route_b_output - ) result = { "route_case": route_case, @@ -953,37 +819,22 @@ def expert_grouped_call() -> None: "dense_graph_ms": dense_ms, "active_graph_ms": active_ms, "grouped_graph_ms": grouped_ms, - "expert_grouped_graph_ms": expert_grouped_ms, "speedup": dense_ms / active_ms, "grouped_vs_active_speedup": active_ms / grouped_ms, - "expert_grouped_vs_slot_grouped_speedup": grouped_ms / expert_grouped_ms, - "expert_grouped_projected_savings_ms_per_43_layers": ( - grouped_ms - expert_grouped_ms - ) - * 43, "grouped_projected_savings_ms_per_43_layers": (active_ms - grouped_ms) * 43, "projected_savings_ms_per_43_layers": (dense_ms - active_ms) * 43, "initial_bitwise_equal": initial_equal, "grouped_initial_bitwise_equal": grouped_initial_equal, - "expert_grouped_initial_bitwise_equal": expert_grouped_initial_equal, "initial_stage_bitwise_equal": stage_parity, "grouped_initial_stage_bitwise_equal": grouped_stage_parity, - "expert_grouped_initial_stage_bitwise_equal": expert_grouped_stage_parity, "grouped_initial_stage_max_abs": grouped_stage_max_abs, - "expert_grouped_initial_stage_max_abs": expert_grouped_stage_max_abs, "grouped_stages_finite": grouped_stages_finite, - "expert_grouped_stages_finite": expert_grouped_stages_finite, "compact_ids_match": compact_ids_match, "dynamic_replay_bitwise_equal": replay_equal, "dynamic_replay_max_abs": replay_max_abs, "grouped_dynamic_replay_bitwise_equal": grouped_replay_equal, "grouped_dynamic_replay_max_abs": grouped_replay_max_abs, - "expert_grouped_dynamic_replay_bitwise_equal": expert_grouped_replay_equal, - "expert_grouped_dynamic_replay_max_abs": expert_grouped_replay_max_abs, "dynamic_route_changes_output": route_changes_output, - "expert_grouped_dynamic_route_changes_output": ( - expert_grouped_route_changes_output - ), } base_gate_passed = ( initial_equal @@ -992,8 +843,6 @@ def expert_grouped_call() -> None: and replay_equal and route_changes_output and grouped_stages_finite - and expert_grouped_stages_finite - and expert_grouped_route_changes_output ) grouped_bitwise_gate_passed = ( grouped_initial_equal @@ -1005,7 +854,7 @@ def expert_grouped_call() -> None: if not base_gate_passed or ( require_grouped_bitwise and not grouped_bitwise_gate_passed ): - raise RuntimeError(f"MXFP4 verifier correctness gate failed: {result}") + raise RuntimeError(f"MXFP4 verifier M8 correctness gate failed: {result}") return result @@ -1143,15 +992,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--seed", type=int, default=17) parser.add_argument("--input-scale", type=float, default=0.01) parser.add_argument("--full-pipeline", action="store_true") - parser.add_argument("--json-out", type=Path) parser.add_argument("--verifier-m8-pipeline", action="store_true") - parser.add_argument( - "--verifier-tokens", - type=int, - choices=range(2, 9), - default=8, - help="Target verifier row count M for the grouped pipeline gate.", - ) parser.add_argument( "--allow-grouped-numeric-drift", action="store_true", @@ -1161,9 +1002,9 @@ def parse_args() -> argparse.Namespace: ) parser.add_argument( "--route-case", - choices=("mixed", "unique48", "hot6", "all"), + choices=("mixed", "unique48", "hot6"), default="mixed", - help="Verifier expert-overlap distribution.", + help="Verifier M8 expert-overlap distribution.", ) parser.add_argument( "--profile-active-once", @@ -1182,28 +1023,19 @@ def main() -> int: args = parse_args() _require_sm70() if args.verifier_m8_pipeline: - route_cases = ( - ("mixed", "unique48", "hot6") - if args.route_case == "all" - else (args.route_case,) + result = benchmark_verifier_m8_pipeline( + num_experts=args.num_experts, + top_k=args.top_k, + repeats=args.repeats, + seed=args.seed, + route_case=args.route_case, + input_scale=args.input_scale, + require_grouped_bitwise=not args.allow_grouped_numeric_drift, ) - result = { - route_case: benchmark_verifier_m8_pipeline( - num_experts=args.num_experts, - top_k=args.top_k, - verifier_tokens=args.verifier_tokens, - repeats=args.repeats, - seed=args.seed + index, - route_case=route_case, - input_scale=args.input_scale, - require_grouped_bitwise=not args.allow_grouped_numeric_drift, - ) - for index, route_case in enumerate(route_cases) - } print( json.dumps( { - "benchmark": "sm70_mxfp4_moe_verifier_pipeline", + "benchmark": "sm70_mxfp4_moe_verifier_m8_pipeline", "device": torch.cuda.get_device_name(), "result": result, }, @@ -1218,16 +1050,16 @@ def main() -> int: repeats=args.repeats, seed=args.seed, ) - payload = { - "benchmark": "sm70_mxfp4_moe_direct_top6_pipeline", - "device": torch.cuda.get_device_name(), - "result": result, - } - rendered = json.dumps(payload, indent=2) - if args.json_out is not None: - args.json_out.parent.mkdir(parents=True, exist_ok=True) - args.json_out.write_text(rendered + "\n", encoding="utf-8") - print(rendered) + print( + json.dumps( + { + "benchmark": "sm70_mxfp4_moe_direct_top6_pipeline", + "device": torch.cuda.get_device_name(), + "result": result, + }, + indent=2, + ) + ) return 0 if args.profile_active_once: if args.stage == "both": diff --git a/benchmarks/benchmark_sm70_native_bm32_allp_crossblock_micro.py b/benchmarks/benchmark_sm70_native_bm32_allp_crossblock_micro.py index 6069cf2926..10771b3a8b 100644 --- a/benchmarks/benchmark_sm70_native_bm32_allp_crossblock_micro.py +++ b/benchmarks/benchmark_sm70_native_bm32_allp_crossblock_micro.py @@ -103,7 +103,8 @@ 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( - f"formal runs require exactly {K_DEFAULT_LAUNCHES} launches per sample." + "formal runs require exactly " + f"{K_DEFAULT_LAUNCHES} launches per sample." ) @@ -260,8 +261,12 @@ 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") }, @@ -297,11 +302,13 @@ 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,7 +333,8 @@ 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,7 +356,9 @@ 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 " @@ -387,31 +397,40 @@ 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( - f"word_dtype={exactness.get('word_dtype')}, requires uint32 packed fp16" + "word_dtype=" + f"{exactness.get('word_dtype')}, requires uint32 packed fp16" ) if exactness.get("word_count") != expected_words: reasons.append( - f"word_count={exactness.get('word_count')}, requires {expected_words}" + "word_count=" + f"{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(f"mismatch_words={exactness.get('mismatch_words')}, requires 0") + reasons.append( + "mismatch_words=" + f"{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: @@ -422,10 +441,13 @@ def _performance_gate(payload: dict[str, Any], rounds: int) -> tuple[bool, list[ 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 >= {K_REQUIRED_SPEEDUP_PCT}%" + f"median speedup={speedup_pct}, requires " + f">= {K_REQUIRED_SPEEDUP_PCT}%" ) return not reasons, reasons @@ -644,11 +666,14 @@ 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), } @@ -805,7 +830,8 @@ 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 177a128420..4679eff592 100644 --- a/benchmarks/benchmark_sm70_native_bm32_allp_scratch_micro.py +++ b/benchmarks/benchmark_sm70_native_bm32_allp_scratch_micro.py @@ -94,7 +94,10 @@ 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(f"launches-per-sample must be at least {K_DEFAULT_LAUNCHES}.") + raise ValueError( + "launches-per-sample must be at least " + f"{K_DEFAULT_LAUNCHES}." + ) def _graphics_clock_mhz(physical_gpu: int) -> int: @@ -245,7 +248,8 @@ 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", {}) @@ -289,11 +293,13 @@ 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( @@ -320,7 +326,8 @@ 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: @@ -341,7 +348,9 @@ 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 " @@ -371,35 +380,46 @@ 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( - f"word_dtype={exactness.get('word_dtype')}, requires uint32 packed fp16" + "word_dtype=" + f"{exactness.get('word_dtype')}, requires uint32 packed fp16" ) if exactness.get("word_count") != expected_words: reasons.append( - f"word_count={exactness.get('word_count')}, requires {expected_words}" + "word_count=" + f"{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(f"mismatch_words={exactness.get('mismatch_words')}, requires 0") + reasons.append( + "mismatch_words=" + f"{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: @@ -410,10 +430,13 @@ def _performance_gate(payload: dict[str, Any], rounds: int) -> tuple[bool, list[ 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 >= {K_REQUIRED_SPEEDUP_PCT}%" + f"median speedup={speedup_pct}, requires " + f">= {K_REQUIRED_SPEEDUP_PCT}%" ) return not reasons, reasons @@ -479,7 +502,9 @@ 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, @@ -600,11 +625,14 @@ 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), } @@ -631,13 +659,17 @@ 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: @@ -703,7 +735,8 @@ 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 cbbe4b867d..0a222682e1 100644 --- a/benchmarks/benchmark_sm70_native_bm32_full_phase_allp_micro.py +++ b/benchmarks/benchmark_sm70_native_bm32_full_phase_allp_micro.py @@ -87,7 +87,10 @@ 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(f"launches-per-sample must be at least {K_DEFAULT_LAUNCHES}.") + raise ValueError( + "launches-per-sample must be at least " + f"{K_DEFAULT_LAUNCHES}." + ) def _graphics_clock_mhz(physical_gpu: int) -> int: @@ -236,7 +239,8 @@ 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} @@ -262,7 +266,8 @@ def _resource_gate( reasons.append("C++ runtime resource gate did not pass") if baseline.get("threads_per_cta") != 512: reasons.append( - f"baseline threads_per_cta={baseline.get('threads_per_cta')}, requires 512" + "baseline threads_per_cta=" + f"{baseline.get('threads_per_cta')}, requires 512" ) if baseline.get("active_ctas_per_sm", 0) < 2: reasons.append( @@ -276,15 +281,18 @@ def _resource_gate( ) if candidate.get("registers_per_thread", 65) > 64: reasons.append( - f"candidate REG={candidate.get('registers_per_thread')}, requires <=64" + "candidate REG=" + f"{candidate.get('registers_per_thread')}, requires <=64" ) if candidate.get("local_bytes_per_thread") != 0: reasons.append( - f"candidate LOCAL={candidate.get('local_bytes_per_thread')}, requires 0" + "candidate LOCAL=" + f"{candidate.get('local_bytes_per_thread')}, requires 0" ) if candidate.get("static_shared_bytes", 48 * 1024 + 1) > 48 * 1024: reasons.append( - f"candidate shared={candidate.get('static_shared_bytes')}, requires <=49152" + "candidate shared=" + f"{candidate.get('static_shared_bytes')}, requires <=49152" ) if candidate.get("active_ctas_per_sm", 0) < 2: reasons.append( @@ -340,31 +348,40 @@ 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( - f"word_dtype={exactness.get('word_dtype')}, requires uint32 packed fp16" + "word_dtype=" + f"{exactness.get('word_dtype')}, requires uint32 packed fp16" ) if exactness.get("word_count") != expected_words: reasons.append( - f"word_count={exactness.get('word_count')}, requires {expected_words}" + "word_count=" + f"{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(f"mismatch_words={exactness.get('mismatch_words')}, requires 0") + reasons.append( + "mismatch_words=" + f"{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", {}) @@ -373,7 +390,8 @@ def _performance_gate(payload: dict[str, Any], rounds: int) -> tuple[bool, list[ reasons.append("candidate speedup was not reported") elif speedup < K_REQUIRED_SPEEDUP_PCT: reasons.append( - f"candidate speedup {speedup:.3f}% is below {K_REQUIRED_SPEEDUP_PCT:.3f}%" + f"candidate speedup {speedup:.3f}% is below " + f"{K_REQUIRED_SPEEDUP_PCT:.3f}%" ) if pairs.get("count") != rounds: reasons.append(f"pair count={pairs.get('count')}, requires {rounds}") @@ -450,7 +468,9 @@ 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, @@ -564,11 +584,14 @@ 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), } @@ -597,7 +620,9 @@ 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 @@ -606,7 +631,9 @@ 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 @@ -675,7 +702,8 @@ 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 d856e33db1..fa97a71a36 100644 --- a/benchmarks/benchmark_sm70_native_bm32_full_phase_micro.py +++ b/benchmarks/benchmark_sm70_native_bm32_full_phase_micro.py @@ -87,7 +87,10 @@ 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(f"launches-per-sample must be at least {K_DEFAULT_LAUNCHES}.") + raise ValueError( + "launches-per-sample must be at least " + f"{K_DEFAULT_LAUNCHES}." + ) def _graphics_clock_mhz(physical_gpu: int) -> int: @@ -236,7 +239,8 @@ 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} @@ -262,7 +266,8 @@ def _resource_gate( reasons.append("C++ runtime resource gate did not pass") if baseline.get("threads_per_cta") != 512: reasons.append( - f"baseline threads_per_cta={baseline.get('threads_per_cta')}, requires 512" + "baseline threads_per_cta=" + f"{baseline.get('threads_per_cta')}, requires 512" ) if baseline.get("active_ctas_per_sm", 0) < 2: reasons.append( @@ -276,15 +281,18 @@ def _resource_gate( ) if candidate.get("registers_per_thread", 65) > 64: reasons.append( - f"candidate REG={candidate.get('registers_per_thread')}, requires <=64" + "candidate REG=" + f"{candidate.get('registers_per_thread')}, requires <=64" ) if candidate.get("local_bytes_per_thread") != 0: reasons.append( - f"candidate LOCAL={candidate.get('local_bytes_per_thread')}, requires 0" + "candidate LOCAL=" + f"{candidate.get('local_bytes_per_thread')}, requires 0" ) if candidate.get("static_shared_bytes", 48 * 1024 + 1) > 48 * 1024: reasons.append( - f"candidate shared={candidate.get('static_shared_bytes')}, requires <=49152" + "candidate shared=" + f"{candidate.get('static_shared_bytes')}, requires <=49152" ) if candidate.get("active_ctas_per_sm", 0) < 2: reasons.append( @@ -340,31 +348,40 @@ 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( - f"word_dtype={exactness.get('word_dtype')}, requires uint32 packed fp16" + "word_dtype=" + f"{exactness.get('word_dtype')}, requires uint32 packed fp16" ) if exactness.get("word_count") != expected_words: reasons.append( - f"word_count={exactness.get('word_count')}, requires {expected_words}" + "word_count=" + f"{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(f"mismatch_words={exactness.get('mismatch_words')}, requires 0") + reasons.append( + "mismatch_words=" + f"{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", {}) @@ -373,7 +390,8 @@ def _performance_gate(payload: dict[str, Any], rounds: int) -> tuple[bool, list[ reasons.append("candidate speedup was not reported") elif speedup < K_REQUIRED_SPEEDUP_PCT: reasons.append( - f"candidate speedup {speedup:.3f}% is below {K_REQUIRED_SPEEDUP_PCT:.3f}%" + f"candidate speedup {speedup:.3f}% is below " + f"{K_REQUIRED_SPEEDUP_PCT:.3f}%" ) if pairs.get("count") != rounds: reasons.append(f"pair count={pairs.get('count')}, requires {rounds}") @@ -450,7 +468,9 @@ 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, @@ -564,11 +584,14 @@ 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), } @@ -597,7 +620,9 @@ 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 @@ -606,7 +631,9 @@ 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 @@ -675,7 +702,8 @@ 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 0391ec6b96..3394ea1fd6 100644 --- a/benchmarks/benchmark_sm70_native_bm32_pv_panel_micro.py +++ b/benchmarks/benchmark_sm70_native_bm32_pv_panel_micro.py @@ -108,7 +108,9 @@ 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()) @@ -226,7 +228,8 @@ def _resource_gate( ) if baseline.get("threads_per_cta") != 512: reasons.append( - f"baseline threads_per_cta={baseline.get('threads_per_cta')}, requires 512" + "baseline threads_per_cta=" + f"{baseline.get('threads_per_cta')}, requires 512" ) if baseline.get("resident_active_warps") != 32: reasons.append( @@ -255,11 +258,13 @@ def _resource_gate( ) if candidate.get("registers_per_thread", 65) > 64: reasons.append( - f"candidate REG={candidate.get('registers_per_thread')}, requires <=64" + "candidate REG=" + f"{candidate.get('registers_per_thread')}, requires <=64" ) if candidate.get("local_bytes_per_thread") != 0: reasons.append( - f"candidate LOCAL={candidate.get('local_bytes_per_thread')}, requires 0" + "candidate LOCAL=" + f"{candidate.get('local_bytes_per_thread')}, requires 0" ) if candidate_topology.get("ctas_per_group") != 1: reasons.append( @@ -383,7 +388,9 @@ 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_prefill_gemm_collective.py b/benchmarks/benchmark_sm70_prefill_gemm_collective.py deleted file mode 100644 index 807f761ca4..0000000000 --- a/benchmarks/benchmark_sm70_prefill_gemm_collective.py +++ /dev/null @@ -1,395 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Compare GEMM+all-reduce with pipelined GEMM+RS+AG on SM70. - -Launch this benchmark with torchrun. It models the exact Qwen3.8 TP4 -row-parallel down/output projection boundary without loading the model. -""" - -from __future__ import annotations - -import argparse -import hashlib -import json -import os -import statistics -from collections.abc import Callable -from pathlib import Path -from typing import Any - -import torch -import torch.distributed as dist -from torch.distributed._symmetric_memory import enable_symm_mem_for_group - - -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser() - parser.add_argument("--m", type=int, default=8000) - parser.add_argument("--k", type=int, default=4352) - parser.add_argument("--n", type=int, default=5120) - parser.add_argument("--warmup", type=int, default=5) - parser.add_argument("--iters", type=int, default=10) - parser.add_argument("--trials", type=int, default=7) - parser.add_argument("--chunk-counts", type=int, nargs="+", default=[2, 4, 8]) - parser.add_argument("--chunk-axis", choices=("m", "n"), default="m") - parser.add_argument("--skip-symm-mem", action="store_true") - parser.add_argument("--include-gemma-norm", action="store_true") - parser.add_argument("--norm-epsilon", type=float, default=1e-6) - parser.add_argument("--seed", type=int, default=20260823) - parser.add_argument("--json-out", type=Path, required=True) - return parser.parse_args() - - -def _event_trials( - launch: Callable[[], None], warmup: int, iters: int, trials: int -) -> dict[str, Any]: - for _ in range(warmup): - launch() - torch.accelerator.synchronize() - - samples: list[float] = [] - for _ in range(trials): - dist.barrier() - start = torch.Event(enable_timing=True) - end = torch.Event(enable_timing=True) - start.record() - for _ in range(iters): - launch() - end.record() - end.synchronize() - samples.append(float(start.elapsed_time(end) / iters)) - return { - "median_ms": float(statistics.median(samples)), - "mean_ms": float(statistics.mean(samples)), - "min_ms": min(samples), - "max_ms": max(samples), - "samples_ms": samples, - } - - -def _digest(tensor: torch.Tensor) -> str: - raw = tensor.detach().contiguous().cpu().view(torch.uint8).numpy().tobytes() - return hashlib.sha256(raw).hexdigest() - - -def main() -> int: - args = _parse_args() - local_rank = int(os.environ["LOCAL_RANK"]) - torch.accelerator.set_device_index(local_rank) - dist.init_process_group(backend="nccl") - rank = dist.get_rank() - world_size = dist.get_world_size() - if world_size != 4: - raise RuntimeError("This benchmark requires TP4.") - if torch.cuda.get_device_capability(local_rank) != (7, 0): - raise RuntimeError("This benchmark requires SM70/V100 GPUs.") - if args.m % world_size: - raise ValueError("M must be divisible by the world size.") - chunk_extent = args.m if args.chunk_axis == "m" else args.n - if any(chunks <= 0 or chunk_extent % chunks for chunks in args.chunk_counts): - raise ValueError( - f"Every chunk count must be positive and divide {args.chunk_axis.upper()}." - ) - if args.include_gemma_norm: - import vllm.kernels # noqa: F401 - - if args.n != 5120: - raise ValueError("The SM70 Gemma long-prefill norm requires N=5120.") - if not hasattr(torch.ops._C, "sm70_gemma_long_prefill_fused_add_rms_norm"): - raise RuntimeError("The active extension lacks the SM70 Gemma norm op.") - - group_name = dist.group.WORLD.group_name - if not args.skip_symm_mem: - enable_symm_mem_for_group(group_name) - device = torch.device("cuda", local_rank) - generator = torch.Generator(device=device).manual_seed(args.seed + rank) - inputs = torch.randn( - (args.m, args.k), - generator=generator, - device=device, - dtype=torch.float16, - ).mul_(0.1) - weight = torch.randn( - (args.k, args.n), - generator=generator, - device=device, - dtype=torch.float16, - ).mul_(0.01) - baseline_out = torch.empty((args.m, args.n), device=device, dtype=torch.float16) - pipelined_out = torch.empty_like(baseline_out) - plain_sp_mm_out = torch.empty_like(baseline_out) - plain_sp_out = torch.empty_like(baseline_out) - plain_sp_shard = torch.empty( - (args.m // world_size, args.n), device=device, dtype=torch.float16 - ) - chunked_outputs = { - chunks: torch.empty_like(baseline_out) for chunks in args.chunk_counts - } - chunked_n_comm_outputs: dict[int, list[torch.Tensor]] = {} - if args.chunk_axis == "n": - chunked_n_comm_outputs = { - chunks: [ - torch.empty( - (args.m, args.n // chunks), - device=device, - dtype=torch.float16, - ) - for _ in range(chunks) - ] - for chunks in args.chunk_counts - } - residual = None - norm_weight = None - baseline_normalized = None - baseline_residual_out = None - pipelined_normalized_shard = None - pipelined_residual_shard = None - chunked_normalized: dict[int, torch.Tensor] = {} - chunked_residual_out: dict[int, torch.Tensor] = {} - if args.include_gemma_norm: - residual = torch.randn( - (args.m, args.n), - generator=torch.Generator(device=device).manual_seed(args.seed + 10000), - device=device, - dtype=torch.float32, - ).mul_(0.1) - norm_weight = torch.randn( - (args.n,), - generator=torch.Generator(device=device).manual_seed(args.seed + 20000), - device=device, - dtype=torch.bfloat16, - ).mul_(0.01) - baseline_normalized = torch.empty_like(baseline_out) - baseline_residual_out = torch.empty_like(residual) - pipelined_normalized_shard = torch.empty( - (args.m // world_size, args.n), device=device, dtype=torch.float16 - ) - pipelined_residual_shard = torch.empty( - (args.m // world_size, args.n), device=device, dtype=torch.float32 - ) - chunked_normalized = { - chunks: torch.empty_like(baseline_out) for chunks in args.chunk_counts - } - chunked_residual_out = { - chunks: torch.empty_like(residual) for chunks in args.chunk_counts - } - - def norm_out( - output: torch.Tensor, - residual_output: torch.Tensor, - input: torch.Tensor, - input_residual: torch.Tensor, - ) -> None: - assert norm_weight is not None - torch.ops._C.sm70_gemma_long_prefill_fused_add_rms_norm( - output, - residual_output, - input, - input_residual, - norm_weight, - args.norm_epsilon, - ) - - def baseline() -> None: - torch.mm(inputs, weight, out=baseline_out) - dist.all_reduce(baseline_out) - if args.include_gemma_norm: - assert baseline_normalized is not None - assert baseline_residual_out is not None - assert residual is not None - norm_out( - baseline_normalized, - baseline_residual_out, - baseline_out, - residual, - ) - - def pipelined() -> None: - shard = torch.ops.symm_mem.fused_matmul_reduce_scatter( - inputs, - weight, - "sum", - 0, - group_name, - ) - if args.include_gemma_norm: - assert residual is not None - assert pipelined_normalized_shard is not None - assert pipelined_residual_shard is not None - residual_shard = residual.narrow( - 0, rank * (args.m // world_size), args.m // world_size - ) - norm_out( - pipelined_normalized_shard, - pipelined_residual_shard, - shard, - residual_shard, - ) - shard = pipelined_normalized_shard - dist.all_gather_into_tensor(pipelined_out, shard) - - def plain_sp() -> None: - torch.mm(inputs, weight, out=plain_sp_mm_out) - dist.reduce_scatter_tensor(plain_sp_shard, plain_sp_mm_out) - shard = plain_sp_shard - if args.include_gemma_norm: - assert residual is not None - assert pipelined_normalized_shard is not None - assert pipelined_residual_shard is not None - residual_shard = residual.narrow( - 0, rank * (args.m // world_size), args.m // world_size - ) - norm_out( - pipelined_normalized_shard, - pipelined_residual_shard, - shard, - residual_shard, - ) - shard = pipelined_normalized_shard - dist.all_gather_into_tensor(plain_sp_out, shard) - - def make_chunked_all_reduce(chunks: int) -> Callable[[], None]: - if args.chunk_axis == "m": - lhs_chunks = inputs.chunk(chunks, dim=0) - rhs_chunks = (weight,) * chunks - output_chunks = chunked_outputs[chunks].chunk(chunks, dim=0) - else: - lhs_chunks = (inputs,) * chunks - rhs_chunks = weight.chunk(chunks, dim=1) - output_chunks = chunked_n_comm_outputs[chunks] - result_chunks = chunked_outputs[chunks].chunk( - chunks, dim=0 if args.chunk_axis == "m" else 1 - ) - - def launch() -> None: - works = [] - for lhs, rhs, output_chunk in zip(lhs_chunks, rhs_chunks, output_chunks): - torch.mm(lhs, rhs, out=output_chunk) - works.append(dist.all_reduce(output_chunk, async_op=True)) - for work, output_chunk, result_chunk in zip( - works, output_chunks, result_chunks - ): - work.wait() - if args.chunk_axis == "n": - result_chunk.copy_(output_chunk) - if args.include_gemma_norm: - assert residual is not None - norm_out( - chunked_normalized[chunks], - chunked_residual_out[chunks], - chunked_outputs[chunks], - residual, - ) - - return launch - - baseline_timing = _event_trials(baseline, args.warmup, args.iters, args.trials) - plain_sp_timing = _event_trials(plain_sp, args.warmup, args.iters, args.trials) - chunked_launches = { - chunks: make_chunked_all_reduce(chunks) for chunks in args.chunk_counts - } - chunked_timings = { - chunks: _event_trials(launch, args.warmup, args.iters, args.trials) - for chunks, launch in chunked_launches.items() - } - pipelined_timing = None - if not args.skip_symm_mem: - pipelined_timing = _event_trials( - pipelined, args.warmup, args.iters, args.trials - ) - baseline() - plain_sp() - for launch in chunked_launches.values(): - launch() - if not args.skip_symm_mem: - pipelined() - torch.accelerator.synchronize() - - if args.include_gemma_norm: - assert baseline_normalized is not None - baseline_result = baseline_normalized - else: - baseline_result = baseline_out - local_payload = { - "rank": rank, - "baseline": baseline_timing, - "baseline_sha256": _digest(baseline_result), - "plain_sp": None, - "pipelined": None, - "chunked": {}, - } - plain_sp_difference = (baseline_result.float() - plain_sp_out.float()).abs() - local_payload["plain_sp"] = { - "timing": plain_sp_timing, - "speedup": baseline_timing["median_ms"] / plain_sp_timing["median_ms"], - "max_abs_diff": float(plain_sp_difference.max().item()), - "mean_abs_diff": float(plain_sp_difference.mean().item()), - "equal_elements": int((baseline_result == plain_sp_out).sum().item()), - "elements": baseline_result.numel(), - "sha256": _digest(plain_sp_out), - } - if pipelined_timing is not None: - difference = (baseline_result.float() - pipelined_out.float()).abs() - local_payload["pipelined"] = { - "timing": pipelined_timing, - "speedup": baseline_timing["median_ms"] / pipelined_timing["median_ms"], - "max_abs_diff": float(difference.max().item()), - "mean_abs_diff": float(difference.mean().item()), - "equal_elements": int((baseline_result == pipelined_out).sum().item()), - "elements": baseline_result.numel(), - "sha256": _digest(pipelined_out), - } - for chunks, output in chunked_outputs.items(): - chunked_result = ( - chunked_normalized[chunks] if args.include_gemma_norm else output - ) - chunked_difference = (baseline_result.float() - chunked_result.float()).abs() - local_payload["chunked"][str(chunks)] = { - "timing": chunked_timings[chunks], - "speedup": baseline_timing["median_ms"] - / chunked_timings[chunks]["median_ms"], - "max_abs_diff": float(chunked_difference.max().item()), - "mean_abs_diff": float(chunked_difference.mean().item()), - "equal_elements": int((baseline_result == chunked_result).sum().item()), - "elements": baseline_result.numel(), - "sha256": _digest(chunked_result), - } - gathered: list[dict[str, Any] | None] = [None] * world_size - dist.all_gather_object(gathered, local_payload) - if rank == 0: - payload = { - "environment": { - "world_size": world_size, - "device": torch.cuda.get_device_name(local_rank), - "torch": torch.__version__, - "torch_cuda": torch.version.cuda, - "m": args.m, - "k": args.k, - "n": args.n, - "dtype": "float16", - "warmup": args.warmup, - "iters": args.iters, - "trials": args.trials, - "chunk_counts": args.chunk_counts, - "chunk_axis": args.chunk_axis, - "symm_mem": not args.skip_symm_mem, - "include_gemma_norm": args.include_gemma_norm, - "norm_epsilon": args.norm_epsilon, - "seed": args.seed, - "cuda_visible_devices": os.getenv("CUDA_VISIBLE_DEVICES"), - "nccl_algo": os.getenv("NCCL_ALGO"), - "nccl_proto": os.getenv("NCCL_PROTO"), - "nccl_min_nchannels": os.getenv("NCCL_MIN_NCHANNELS"), - "nccl_max_nchannels": os.getenv("NCCL_MAX_NCHANNELS"), - }, - "ranks": gathered, - } - args.json_out.parent.mkdir(parents=True, exist_ok=True) - args.json_out.write_text(json.dumps(payload, indent=2) + "\n") - print(json.dumps(payload, indent=2)) - dist.destroy_process_group() - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/benchmarks/benchmark_sm70_raw_hmma_probe.py b/benchmarks/benchmark_sm70_raw_hmma_probe.py index aec857c44d..a1928d4436 100644 --- a/benchmarks/benchmark_sm70_raw_hmma_probe.py +++ b/benchmarks/benchmark_sm70_raw_hmma_probe.py @@ -32,7 +32,10 @@ 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 ec680a82dc..037b70653a 100644 --- a/benchmarks/benchmark_sm70_rejection_sampler_micro.py +++ b/benchmarks/benchmark_sm70_rejection_sampler_micro.py @@ -174,7 +174,9 @@ 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() @@ -189,7 +191,9 @@ 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) @@ -206,9 +210,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) @@ -243,9 +247,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 @@ -288,7 +292,9 @@ 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 @@ -388,7 +394,9 @@ 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( @@ -487,9 +495,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_sglang_push_chain.py b/benchmarks/benchmark_sm70_sglang_push_chain.py deleted file mode 100644 index 52b38ece3f..0000000000 --- a/benchmarks/benchmark_sm70_sglang_push_chain.py +++ /dev/null @@ -1,162 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Benchmark SGLang's push all-reduce on a verifier-shaped graph chain. - -Run this script from a pinned SGLang checkout and environment. It deliberately -uses the public SGLang JIT communicator so that the comparison includes its -production graph-input registration and two-epoch push protocol. -""" - -from __future__ import annotations - -import argparse -import json -import os -from pathlib import Path - -import torch -import torch.distributed as dist - - -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser() - parser.add_argument("--numel", type=int, default=8 * 5120) - parser.add_argument("--count", type=int, default=128) - parser.add_argument("--warmup", type=int, default=40) - parser.add_argument("--iters", type=int, default=200) - parser.add_argument("--max-push-blocks", type=int, default=80) - parser.add_argument("--seed", type=int, default=0) - parser.add_argument("--json-out", type=Path, required=True) - return parser.parse_args() - - -def _max_rank_float(value: float) -> float: - result = torch.tensor([value], dtype=torch.float64) - dist.all_reduce(result, op=dist.ReduceOp.MAX) - return float(result.item()) - - -def main() -> None: - args = _parse_args() - rank = int(os.environ["RANK"]) - local_rank = int(os.environ["LOCAL_RANK"]) - world_size = int(os.environ["WORLD_SIZE"]) - if world_size != 4: - raise ValueError(f"this verifier microbenchmark requires TP4, got {world_size}") - - torch.accelerator.set_device_index(local_rank) - device = torch.device("cuda", local_rank) - dist.init_process_group(backend="gloo") - - import sglang.srt.distributed.parallel_state as parallel_state - from sglang.jit_kernel.all_reduce import ( - AllReduceAlgo, - _jit_custom_all_reduce_push_module, - ) - from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import ( - CustomAllReduceV2, - ) - - parallel_state._WORLD = parallel_state.init_world_group( - ranks=list(range(world_size)), - local_rank=local_rank, - backend="nccl", - ) - cpu_group = parallel_state._WORLD.cpu_group - nccl_group = parallel_state._WORLD.device_group - if nccl_group is None: - raise RuntimeError("SGLang did not create the NCCL device group") - - bytes_per_input = args.numel * torch.float16.itemsize - communicator = CustomAllReduceV2( - cpu_group, - device, - max_pull_size=bytes_per_input, - max_push_size=bytes_per_input, - max_push_blocks=args.max_push_blocks, - ) - if communicator.disabled: - raise RuntimeError("SGLang push all-reduce is disabled on this topology") - communicator.override_algo = AllReduceAlgo.ONE_SHOT_PUSH - _jit_custom_all_reduce_push_module(torch.float16, world_size) - - generator = torch.Generator(device=device) - generator.manual_seed(args.seed + rank) - source = ( - torch.rand( - (args.count, args.numel), - dtype=torch.float32, - device=device, - generator=generator, - ) - * 0.03 - + 0.01 - ).half() - work = source.clone() - - gathered = [torch.empty_like(source) for _ in range(world_size)] - dist.all_gather(gathered, source, group=nccl_group) - reference = gathered[0].float() - for peer in gathered[1:]: - reference.add_(peer.float()) - reference = reference.half() - - graph = torch.cuda.CUDAGraph() - with communicator.capture(), torch.cuda.graph(graph): - for index in range(args.count): - communicator.custom_all_reduce(work[index]) - torch.accelerator.synchronize() - dist.barrier() - - work.copy_(source) - graph.replay() - torch.accelerator.synchronize() - exact = bool(torch.equal(work, reference)) - difference = (work.float() - reference.float()).abs() - mismatch_count = int((work != reference).sum().item()) - - for _ in range(args.warmup): - graph.replay() - torch.accelerator.synchronize() - dist.barrier() - - start = torch.Event(enable_timing=True) - end = torch.Event(enable_timing=True) - start.record() - for _ in range(args.iters): - graph.replay() - end.record() - end.synchronize() - local_ms = start.elapsed_time(end) / args.iters - max_rank_ms = _max_rank_float(local_ms) - - if rank == 0: - result = { - "implementation": "sglang_v100_one_shot_push", - "world_size": world_size, - "torch_version": torch.__version__, - "cuda_version": torch.version.cuda, - "numel": args.numel, - "bytes_per_input": bytes_per_input, - "allreduces_per_graph_replay": args.count, - "max_push_blocks": args.max_push_blocks, - "warmup_replays": args.warmup, - "timed_replays": args.iters, - "max_rank_ms_per_graph_replay": max_rank_ms, - "max_rank_us_per_allreduce": max_rank_ms * 1000.0 / args.count, - "exact_rank_order_fp32_reference": exact, - "mismatch_count": mismatch_count, - "max_diff": float(difference.max().item()), - "mean_diff": float(difference.mean().item()), - } - args.json_out.parent.mkdir(parents=True, exist_ok=True) - args.json_out.write_text(json.dumps(result, indent=2) + "\n") - print(json.dumps(result, indent=2)) - - communicator.close() - dist.barrier() - dist.destroy_process_group() - - -if __name__ == "__main__": - main() diff --git a/benchmarks/benchmark_sm70_source_inventory.py b/benchmarks/benchmark_sm70_source_inventory.py index 87c7e5f151..f8b13b47fa 100644 --- a/benchmarks/benchmark_sm70_source_inventory.py +++ b/benchmarks/benchmark_sm70_source_inventory.py @@ -502,14 +502,19 @@ 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()) + }, } @@ -592,11 +597,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 @@ -612,9 +617,15 @@ 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 @@ -664,37 +675,32 @@ 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_tp2_compact_verifier_logits.py b/benchmarks/benchmark_sm70_tp2_compact_verifier_logits.py index 16456da153..d7560f8e69 100644 --- a/benchmarks/benchmark_sm70_tp2_compact_verifier_logits.py +++ b/benchmarks/benchmark_sm70_tp2_compact_verifier_logits.py @@ -1,10 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Measure TP full-vocab versus compact top-k verifier-logit transport. +"""Measure TP2 full-vocab versus compact top-k verifier-logit transport. -Run on two or more identical GPUs without loading a model: +Run on two identical GPUs without loading a model: - CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --standalone --nproc_per_node=4 \ + CUDA_VISIBLE_DEVICES=0,1 torchrun --standalone --nproc_per_node=2 \ benchmarks/benchmark_sm70_tp2_compact_verifier_logits.py \ --json-out bench_results/.../tp2_compact_verifier_logits.json @@ -93,8 +93,8 @@ def main() -> None: local_rank = int(os.environ["LOCAL_RANK"]) rank = int(os.environ["RANK"]) world_size = int(os.environ["WORLD_SIZE"]) - if world_size < 2: - raise ValueError(f"This benchmark requires TP>1, got world_size={world_size}.") + if world_size != 2: + raise ValueError(f"This benchmark is TP2-only, got world_size={world_size}.") if args.vocab_size % world_size: raise ValueError("--vocab-size must be divisible by TP size.") if args.top_k <= 0 or args.top_k > args.vocab_size // world_size: diff --git a/benchmarks/benchmark_sm70_tp4_fused_norm_collective.py b/benchmarks/benchmark_sm70_tp4_fused_norm_collective.py deleted file mode 100644 index 38e460417e..0000000000 --- a/benchmarks/benchmark_sm70_tp4_fused_norm_collective.py +++ /dev/null @@ -1,305 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Benchmark SM70 fused reduce-scatter + Gemma RMSNorm + all-gather. - -Launch with torchrun on a fully connected TP4 V100 group. The candidate is a -benchmark-only implementation of the communication topology in NVIDIA's NCCL -Device API fused RMSNorm example, using vLLM's existing IPC buffers and -cross-device signal substrate so it can be measured with NCCL 2.27. -""" - -from __future__ import annotations - -import argparse -import hashlib -import json -import os -import statistics -from collections.abc import Callable -from pathlib import Path -from typing import Any - -import torch -import torch.distributed as dist - -import vllm.kernels # noqa: F401 -from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce -from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator - - -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser() - parser.add_argument("--m", type=int, default=8000) - parser.add_argument("--n", type=int, default=5120) - parser.add_argument("--threads", type=int, nargs="+", default=[256, 512, 1024]) - parser.add_argument("--blocks", type=int, nargs="+", default=[80]) - parser.add_argument("--warmup", type=int, default=3) - parser.add_argument("--iters", type=int, default=5) - parser.add_argument("--trials", type=int, default=5) - parser.add_argument("--epsilon", type=float, default=1e-6) - parser.add_argument("--seed", type=int, default=20260823) - parser.add_argument("--candidate-cuda-graph", action="store_true") - parser.add_argument("--json-out", type=Path, required=True) - return parser.parse_args() - - -def _event_trials( - launch: Callable[[], None], warmup: int, iters: int, trials: int -) -> dict[str, Any]: - for _ in range(warmup): - launch() - torch.accelerator.synchronize() - - samples: list[float] = [] - for _ in range(trials): - dist.barrier() - start = torch.Event(enable_timing=True) - end = torch.Event(enable_timing=True) - start.record() - for _ in range(iters): - launch() - end.record() - end.synchronize() - samples.append(float(start.elapsed_time(end) / iters)) - return { - "median_ms": float(statistics.median(samples)), - "mean_ms": float(statistics.mean(samples)), - "min_ms": min(samples), - "max_ms": max(samples), - "samples_ms": samples, - } - - -def _digest(tensor: torch.Tensor) -> str: - raw = tensor.detach().contiguous().cpu().view(torch.uint8).numpy().tobytes() - return hashlib.sha256(raw).hexdigest() - - -def main() -> int: - args = _parse_args() - local_rank = int(os.environ["LOCAL_RANK"]) - torch.accelerator.set_device_index(local_rank) - dist.init_process_group(backend="nccl") - cpu_group = dist.new_group(backend="gloo") - rank = dist.get_rank() - world_size = dist.get_world_size() - if world_size != 4: - raise RuntimeError("This benchmark requires TP4.") - if torch.cuda.get_device_capability(local_rank) != (7, 0): - raise RuntimeError("This benchmark requires SM70/V100 GPUs.") - if args.n != 5120 or args.m % world_size: - raise ValueError("The candidate requires N=5120 and TP-divisible M.") - if any(threads not in (256, 512, 1024) for threads in args.threads): - raise ValueError("--threads choices must be 256, 512, or 1024.") - if any(blocks < 1 or blocks > 512 for blocks in args.blocks): - raise ValueError("--blocks choices must be in [1, 512].") - variants = [(threads, blocks) for threads in args.threads for blocks in args.blocks] - - device = torch.device("cuda", local_rank) - generator = torch.Generator(device=device).manual_seed(args.seed + rank) - partial = torch.randn( - (args.m, args.n), - generator=generator, - device=device, - dtype=torch.float16, - ).mul_(0.1) - residual = torch.randn( - (args.m, args.n), - generator=torch.Generator(device=device).manual_seed(args.seed + 10000), - device=device, - dtype=torch.float32, - ).mul_(0.1) - weight = torch.randn( - (args.n,), - generator=torch.Generator(device=device).manual_seed(args.seed + 20000), - device=device, - dtype=torch.float16, - ).mul_(0.01) - - baseline_reduced = torch.empty_like(partial) - baseline_normalized = torch.empty_like(partial) - baseline_residual = torch.empty_like(residual) - candidate_normalized = {variant: torch.empty_like(partial) for variant in variants} - candidate_residual = { - variant: torch.empty_like(residual, dtype=torch.float32) for variant in variants - } - - max_size = partial.numel() * partial.element_size() + 1 - pynccl = PyNcclCommunicator(group=cpu_group, device=device) - if pynccl.disabled: - raise RuntimeError("PyNccl is unavailable for the production baseline.") - communicator = CustomAllreduce( - cpu_group, - device, - max_size=max_size, - long_prefill_fusion_enabled=True, - ) - if communicator.disabled: - raise RuntimeError("Custom all-reduce is unavailable for this TP4 group.") - - def norm_out( - output: torch.Tensor, - residual_output: torch.Tensor, - input: torch.Tensor, - input_residual: torch.Tensor, - ) -> None: - torch.ops._C.sm70_gemma_long_prefill_fused_add_rms_norm( - output, - residual_output, - input, - input_residual, - weight, - args.epsilon, - ) - - def baseline() -> None: - pynccl.all_reduce(partial, baseline_reduced) - norm_out( - baseline_normalized, - baseline_residual, - baseline_reduced, - residual, - ) - - def baseline_all_reduce() -> None: - pynccl.all_reduce(partial, baseline_reduced) - - def baseline_norm() -> None: - norm_out( - baseline_normalized, - baseline_residual, - baseline_reduced, - residual, - ) - - def d2d_copy() -> None: - baseline_reduced.copy_(partial) - - def make_candidate(threads: int, blocks: int) -> Callable[[], None]: - def launch() -> None: - os.environ["VLLM_SM70_TP4_LONG_FUSED_NORM_THREADS"] = str(threads) - os.environ["VLLM_SM70_TP4_LONG_FUSED_NORM_BLOCKS"] = str(blocks) - communicator.sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather( - partial, - residual, - weight, - args.epsilon, - normalized_out=candidate_normalized[(threads, blocks)], - residual_out=candidate_residual[(threads, blocks)], - ) - - return launch - - baseline_timing = _event_trials(baseline, args.warmup, args.iters, args.trials) - baseline_ar_timing = _event_trials( - baseline_all_reduce, args.warmup, args.iters, args.trials - ) - baseline_all_reduce() - baseline_norm_timing = _event_trials( - baseline_norm, args.warmup, args.iters, args.trials - ) - copy_timing = _event_trials(d2d_copy, args.warmup, args.iters, args.trials) - candidate_launches = {variant: make_candidate(*variant) for variant in variants} - candidate_graphs: list[torch.cuda.CUDAGraph] = [] - if args.candidate_cuda_graph: - graph_launches: dict[tuple[int, int], Callable[[], None]] = {} - for variant, launch in candidate_launches.items(): - launch() - torch.accelerator.synchronize() - dist.barrier() - graph = torch.cuda.CUDAGraph() - with communicator.capture(), torch.cuda.graph(graph): - launch() - candidate_graphs.append(graph) - graph_launches[variant] = graph.replay - candidate_launches = graph_launches - candidate_timings = { - variant: _event_trials(launch, args.warmup, args.iters, args.trials) - for variant, launch in candidate_launches.items() - } - - baseline() - for launch in candidate_launches.values(): - launch() - torch.accelerator.synchronize() - local_rows = args.m // world_size - residual_reference = baseline_residual.narrow(0, rank * local_rows, local_rows) - candidates: dict[str, Any] = {} - for variant in variants: - threads, blocks = variant - normalized_difference = ( - baseline_normalized.float() - candidate_normalized[variant].float() - ).abs() - candidate_residual_shard = candidate_residual[variant].narrow( - 0, rank * local_rows, local_rows - ) - residual_difference = (residual_reference - candidate_residual_shard).abs() - timing = candidate_timings[variant] - candidates[f"t{threads}_b{blocks}"] = { - "threads": threads, - "blocks": blocks, - "timing": timing, - "speedup": baseline_timing["median_ms"] / timing["median_ms"], - "normalized_max_abs_diff": float(normalized_difference.max().item()), - "normalized_mean_abs_diff": float(normalized_difference.mean().item()), - "normalized_equal_elements": int( - (baseline_normalized == candidate_normalized[variant]).sum().item() - ), - "normalized_elements": baseline_normalized.numel(), - "residual_max_abs_diff": float(residual_difference.max().item()), - "residual_mean_abs_diff": float(residual_difference.mean().item()), - "normalized_sha256": _digest(candidate_normalized[variant]), - "residual_shard_sha256": _digest(candidate_residual_shard), - } - - local_payload = { - "rank": rank, - "baseline": baseline_timing, - "baseline_all_reduce": baseline_ar_timing, - "baseline_norm": baseline_norm_timing, - "d2d_copy": copy_timing, - "baseline_normalized_sha256": _digest(baseline_normalized), - "baseline_residual_shard_sha256": _digest(residual_reference), - "candidates": candidates, - } - gathered: list[dict[str, Any] | None] = [None] * world_size - dist.all_gather_object(gathered, local_payload) - if rank == 0: - payload = { - "environment": { - "world_size": world_size, - "device": torch.cuda.get_device_name(local_rank), - "torch": torch.__version__, - "torch_cuda": torch.version.cuda, - "m": args.m, - "n": args.n, - "dtype": "float16", - "residual_dtype": "float32", - "threads": args.threads, - "blocks": args.blocks, - "candidate_cuda_graph": args.candidate_cuda_graph, - "warmup": args.warmup, - "iters": args.iters, - "trials": args.trials, - "epsilon": args.epsilon, - "seed": args.seed, - "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), - "nccl_algo": os.environ.get("NCCL_ALGO"), - "nccl_proto": os.environ.get("NCCL_PROTO"), - }, - "ranks": gathered, - } - args.json_out.parent.mkdir(parents=True, exist_ok=True) - args.json_out.write_text(json.dumps(payload, indent=2) + "\n") - print(json.dumps(payload, indent=2), flush=True) - - communicator.close() - pynccl.destroy() - dist.barrier() - dist.destroy_process_group(cpu_group) - dist.destroy_process_group() - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/benchmarks/benchmark_sm70_turbomind_parity.py b/benchmarks/benchmark_sm70_turbomind_parity.py index bbc2cd7c75..ed1f27739f 100644 --- a/benchmarks/benchmark_sm70_turbomind_parity.py +++ b/benchmarks/benchmark_sm70_turbomind_parity.py @@ -94,9 +94,8 @@ 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() @@ -109,9 +108,8 @@ 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") @@ -160,7 +158,9 @@ 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,13 +205,8 @@ 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 @@ -275,9 +270,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()) @@ -305,7 +300,9 @@ 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 301d49ef27..26e7c4435d 100644 --- a/benchmarks/benchmark_sm70_wmma_fragment_probe.py +++ b/benchmarks/benchmark_sm70_wmma_fragment_probe.py @@ -15,7 +15,11 @@ 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", @@ -37,16 +41,17 @@ 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( @@ -54,7 +59,9 @@ 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) @@ -62,10 +69,15 @@ 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), @@ -82,7 +94,10 @@ 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, @@ -91,7 +106,8 @@ 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_dsv4_quality_results.py b/benchmarks/compare_dsv4_quality_results.py deleted file mode 100644 index 073d8276de..0000000000 --- a/benchmarks/compare_dsv4_quality_results.py +++ /dev/null @@ -1,425 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Compare matched no-speculation and DSpark quality artifacts.""" - -from __future__ import annotations - -import argparse -import json -from pathlib import Path -from typing import Any - - -def _load_json(path: Path) -> dict[str, Any]: - return json.loads(path.read_text(encoding="utf-8")) - - -def _load_jsonl(path: Path) -> list[dict[str, Any]]: - return [ - json.loads(line) - for line in path.read_text(encoding="utf-8").splitlines() - if line.strip() - ] - - -def _manifest_hashes(value: Any, prefix: str = "") -> dict[str, str]: - if not isinstance(value, dict): - return {} - if isinstance(value.get("sha256"), str): - return {prefix: value["sha256"]} - hashes: dict[str, str] = {} - for key, child in value.items(): - child_prefix = f"{prefix}.{key}" if prefix else key - hashes.update(_manifest_hashes(child, child_prefix)) - return hashes - - -def _compare_humaneval( - reference: dict[str, Any], candidate: dict[str, Any] -) -> dict[str, Any]: - reference_rows = {row["task_id"]: row for row in reference["human_eval"]["records"]} - candidate_rows = {row["task_id"]: row for row in candidate["human_eval"]["records"]} - common = sorted(reference_rows.keys() & candidate_rows.keys()) - regressions = [ - task_id - for task_id in common - if reference_rows[task_id]["passed"] and not candidate_rows[task_id]["passed"] - ] - improvements = [ - task_id - for task_id in common - if not reference_rows[task_id]["passed"] and candidate_rows[task_id]["passed"] - ] - exact_responses = sum( - reference_rows[task_id].get("response") - == candidate_rows[task_id].get("response") - for task_id in common - ) - reference_passed = int(reference["human_eval"]["passed"]) - candidate_passed = int(candidate["human_eval"]["passed"]) - missing = sorted(reference_rows.keys() ^ candidate_rows.keys()) - return { - "reference_passed": reference_passed, - "candidate_passed": candidate_passed, - "samples": len(common), - "exact_response_matches": exact_responses, - "regressions": regressions, - "improvements": improvements, - "missing_or_extra_tasks": missing, - # Individual flips remain visible for diagnosis, but are not token- - # identity gates. Acceptance is based on matched inputs and aggregate - # task quality. - "passed": not missing and candidate_passed >= reference_passed, - } - - -def _compare_longbench( - reference: dict[str, Any], - candidate: dict[str, Any], - tolerance: float, -) -> dict[str, Any]: - reference_sets = reference["longbench"]["datasets"] - candidate_sets = candidate["longbench"]["datasets"] - missing = sorted(reference_sets.keys() ^ candidate_sets.keys()) - results: dict[str, Any] = {} - for dataset in sorted(reference_sets.keys() & candidate_sets.keys()): - reference_rows = { - int(row["source_index"]): row for row in reference_sets[dataset]["records"] - } - candidate_rows = { - int(row["source_index"]): row for row in candidate_sets[dataset]["records"] - } - common = sorted(reference_rows.keys() & candidate_rows.keys()) - row_mismatch = sorted(reference_rows.keys() ^ candidate_rows.keys()) - input_mismatch = [ - index - for index in common - if reference_rows[index].get("answers") - != candidate_rows[index].get("answers") - ] - regressions = [ - { - "source_index": index, - "reference_score": float(reference_rows[index]["score"]), - "candidate_score": float(candidate_rows[index]["score"]), - } - for index in common - if float(candidate_rows[index]["score"]) + tolerance - < float(reference_rows[index]["score"]) - ] - exact_predictions = sum( - reference_rows[index].get("prediction") - == candidate_rows[index].get("prediction") - for index in common - ) - reference_score = float(reference_sets[dataset]["score"]) - candidate_score = float(candidate_sets[dataset]["score"]) - results[dataset] = { - "reference_score": reference_score, - "candidate_score": candidate_score, - "samples": len(common), - "exact_prediction_matches": exact_predictions, - "regressions": regressions, - "missing_or_extra_rows": row_mismatch, - "input_mismatches": input_mismatch, - "passed": not row_mismatch - and not input_mismatch - and candidate_score + tolerance >= reference_score, - } - reference_average = float(reference["longbench"]["average_score"]) - candidate_average = float(candidate["longbench"]["average_score"]) - return { - "reference_average": reference_average, - "candidate_average": candidate_average, - "score_tolerance": tolerance, - "datasets": results, - "missing_or_extra_datasets": missing, - "passed": not missing - and all(result["passed"] for result in results.values()) - and candidate_average + tolerance >= reference_average, - } - - -def _compare_api_quality( - reference: dict[str, Any], - candidate: dict[str, Any], - tolerance: float, -) -> dict[str, Any]: - reference_hashes = _manifest_hashes(reference.get("input_manifest", {})) - candidate_hashes = _manifest_hashes(candidate.get("input_manifest", {})) - manifest_equal = bool(reference_hashes) and reference_hashes == candidate_hashes - reference_contract = reference.get("evaluation_contract") - candidate_contract = candidate.get("evaluation_contract") - contract_equal = ( - bool(reference_contract) and reference_contract == candidate_contract - ) - humaneval = _compare_humaneval(reference, candidate) - longbench = _compare_longbench(reference, candidate, tolerance) - return { - "input_hashes_equal": manifest_equal, - "reference_input_hashes": reference_hashes, - "candidate_input_hashes": candidate_hashes, - "evaluation_contract_equal": contract_equal, - "reference_evaluation_contract": reference_contract, - "candidate_evaluation_contract": candidate_contract, - "humaneval": humaneval, - "longbench": longbench, - "passed": manifest_equal - and contract_equal - and humaneval["passed"] - and longbench["passed"], - } - - -def _compare_gsm8k( - reference: dict[str, Any], candidate: dict[str, Any] -) -> dict[str, Any]: - reference_hashes = _manifest_hashes(reference.get("input_manifest", {})) - candidate_hashes = _manifest_hashes(candidate.get("input_manifest", {})) - manifest_equal = bool(reference_hashes) and reference_hashes == candidate_hashes - contract_keys = ( - "questions", - "few_shot", - "temperature", - "top_p", - "seed", - "max_tokens", - "strictly_sequential", - "train_selection", - "test_selection", - "prompt_format", - "answer_normalization", - "stop_sequences", - ) - reference_contract = { - key: reference.get("contract", {}).get(key) for key in contract_keys - } - candidate_contract = { - key: candidate.get("contract", {}).get(key) for key in contract_keys - } - contract_equal = reference_contract == candidate_contract - reference_rows = {int(row["index"]): row for row in reference["rows"]} - candidate_rows = {int(row["index"]): row for row in candidate["rows"]} - common = sorted(reference_rows.keys() & candidate_rows.keys()) - missing = sorted(reference_rows.keys() ^ candidate_rows.keys()) - expected_samples = reference_contract["questions"] - row_count_matches_contract = ( - isinstance(expected_samples, int) - and len(reference["rows"]) == expected_samples - and len(candidate["rows"]) == expected_samples - and len(reference_rows) == expected_samples - and len(candidate_rows) == expected_samples - ) - - def row_is_consistent(row: dict[str, Any]) -> bool: - expected = row.get("expected") - predicted = row.get("predicted") - return bool( - expected is not None - and bool(row.get("correct")) - == (predicted is not None and predicted == expected) - and bool(row.get("invalid")) == (predicted is None) - ) - - rows_self_consistent = all( - row_is_consistent(row) for row in [*reference["rows"], *candidate["rows"]] - ) - summaries_self_consistent = all( - int(result["correct"]) - == sum(bool(row.get("correct")) for row in result["rows"]) - and int(result["invalid"]) - == sum(bool(row.get("invalid")) for row in result["rows"]) - for result in (reference, candidate) - ) - input_mismatch = [ - index - for index in common - if ( - reference_rows[index].get("question"), - reference_rows[index].get("expected"), - ) - != ( - candidate_rows[index].get("question"), - candidate_rows[index].get("expected"), - ) - ] - regressions = [ - index - for index in common - if reference_rows[index]["correct"] and not candidate_rows[index]["correct"] - ] - improvements = [ - index - for index in common - if not reference_rows[index]["correct"] and candidate_rows[index]["correct"] - ] - exact_predictions = sum( - reference_rows[index].get("predicted") == candidate_rows[index].get("predicted") - for index in common - ) - return { - "input_hashes_equal": manifest_equal, - "reference_input_hashes": reference_hashes, - "candidate_input_hashes": candidate_hashes, - "evaluation_contract_equal": contract_equal, - "reference_evaluation_contract": reference_contract, - "candidate_evaluation_contract": candidate_contract, - "row_count_matches_contract": row_count_matches_contract, - "rows_self_consistent": rows_self_consistent, - "summaries_self_consistent": summaries_self_consistent, - "reference_correct": int(reference["correct"]), - "candidate_correct": int(candidate["correct"]), - "correct_delta": int(candidate["correct"]) - int(reference["correct"]), - "reference_invalid": int(reference["invalid"]), - "candidate_invalid": int(candidate["invalid"]), - "samples": len(common), - "exact_prediction_matches": exact_predictions, - "regressions": regressions, - "improvements": improvements, - "missing_or_extra_rows": missing, - "input_mismatches": input_mismatch, - # Report directional flips without requiring greedy/token identity. - # Matched aggregate correctness and validity are the quality gate. - "passed": manifest_equal - and contract_equal - and row_count_matches_contract - and rows_self_consistent - and summaries_self_consistent - and not missing - and not input_mismatch - and int(candidate["correct"]) >= int(reference["correct"]) - and int(candidate["invalid"]) <= int(reference["invalid"]), - } - - -def _compare_needle( - reference: list[dict[str, Any]], candidate: list[dict[str, Any]] -) -> dict[str, Any]: - reference_rows = {row["sample_id"]: row for row in reference} - candidate_rows = {row["sample_id"]: row for row in candidate} - common = sorted(reference_rows.keys() & candidate_rows.keys()) - missing = sorted(reference_rows.keys() ^ candidate_rows.keys()) - input_mismatch = [ - sample_id - for sample_id in common - if ( - reference_rows[sample_id].get("target_tokens"), - reference_rows[sample_id].get("depth"), - reference_rows[sample_id].get("code"), - ) - != ( - candidate_rows[sample_id].get("target_tokens"), - candidate_rows[sample_id].get("depth"), - candidate_rows[sample_id].get("code"), - ) - ] - final_regressions = [ - sample_id - for sample_id in common - if reference_rows[sample_id]["hit"] and not candidate_rows[sample_id]["hit"] - ] - anywhere_regressions = [ - sample_id - for sample_id in common - if reference_rows[sample_id]["hit_anywhere"] - and not candidate_rows[sample_id]["hit_anywhere"] - ] - reference_hits = sum(bool(row["hit"]) for row in reference_rows.values()) - candidate_hits = sum(bool(row["hit"]) for row in candidate_rows.values()) - reference_anywhere = sum( - bool(row["hit_anywhere"]) for row in reference_rows.values() - ) - candidate_anywhere = sum( - bool(row["hit_anywhere"]) for row in candidate_rows.values() - ) - return { - "reference_final_hits": reference_hits, - "candidate_final_hits": candidate_hits, - "reference_anywhere_hits": reference_anywhere, - "candidate_anywhere_hits": candidate_anywhere, - "samples": len(common), - "final_hit_regressions": final_regressions, - "anywhere_hit_regressions": anywhere_regressions, - "missing_or_extra_rows": missing, - "input_mismatches": input_mismatch, - "passed": not missing - and not input_mismatch - and candidate_hits >= reference_hits - and candidate_anywhere >= reference_anywhere, - } - - -def _paired_paths( - parser: argparse.ArgumentParser, - reference: Path | None, - candidate: Path | None, - label: str, -) -> bool: - if (reference is None) != (candidate is None): - parser.error(f"--reference-{label} and --candidate-{label} must be paired") - return reference is not None - - -def main() -> int: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--reference-api", type=Path) - parser.add_argument("--candidate-api", type=Path) - parser.add_argument("--reference-gsm8k", type=Path) - parser.add_argument("--candidate-gsm8k", type=Path) - parser.add_argument("--reference-needle", type=Path) - parser.add_argument("--candidate-needle", type=Path) - parser.add_argument("--longbench-score-tolerance", type=float, default=0.0) - parser.add_argument("--out", type=Path, required=True) - args = parser.parse_args() - if args.longbench_score_tolerance < 0: - parser.error("--longbench-score-tolerance must be non-negative") - - has_api = _paired_paths(parser, args.reference_api, args.candidate_api, "api") - has_gsm8k = _paired_paths( - parser, args.reference_gsm8k, args.candidate_gsm8k, "gsm8k" - ) - has_needle = _paired_paths( - parser, args.reference_needle, args.candidate_needle, "needle" - ) - if not (has_api or has_gsm8k or has_needle): - parser.error("at least one paired quality artifact must be provided") - - result: dict[str, Any] = { - "contract": { - key: str(value) if isinstance(value, Path) else value - for key, value in vars(args).items() - }, - } - gates: list[bool] = [] - if has_api: - assert args.reference_api is not None and args.candidate_api is not None - result["api_quality"] = _compare_api_quality( - _load_json(args.reference_api), - _load_json(args.candidate_api), - args.longbench_score_tolerance, - ) - gates.append(result["api_quality"]["passed"]) - if has_gsm8k: - assert args.reference_gsm8k is not None and args.candidate_gsm8k is not None - result["gsm8k"] = _compare_gsm8k( - _load_json(args.reference_gsm8k), _load_json(args.candidate_gsm8k) - ) - gates.append(result["gsm8k"]["passed"]) - if has_needle: - assert args.reference_needle is not None and args.candidate_needle is not None - result["needle"] = _compare_needle( - _load_jsonl(args.reference_needle), _load_jsonl(args.candidate_needle) - ) - gates.append(result["needle"]["passed"]) - result["passed"] = all(gates) - - args.out.parent.mkdir(parents=True, exist_ok=True) - args.out.write_text( - json.dumps(result, ensure_ascii=False, indent=2) + "\n", encoding="utf-8" - ) - print(json.dumps(result, ensure_ascii=False), flush=True) - return 0 if result["passed"] else 1 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/benchmarks/compare_sm70_tensor_dumps.py b/benchmarks/compare_sm70_tensor_dumps.py index 649bf9b7c8..9e3de08e8a 100644 --- a/benchmarks/compare_sm70_tensor_dumps.py +++ b/benchmarks/compare_sm70_tensor_dumps.py @@ -79,13 +79,11 @@ 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 @@ -136,37 +134,33 @@ 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, @@ -268,12 +262,10 @@ 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)) @@ -288,25 +280,21 @@ 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( @@ -316,25 +304,21 @@ 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, @@ -344,27 +328,23 @@ 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_device_communicators.py b/benchmarks/kernels/benchmark_device_communicators.py index 4afdf07723..24e22023b9 100644 --- a/benchmarks/kernels/benchmark_device_communicators.py +++ b/benchmarks/kernels/benchmark_device_communicators.py @@ -49,10 +49,9 @@ # Default sequence lengths to benchmark DEFAULT_SEQUENCE_LENGTHS = [16, 64, 128, 512, 1024, 2048, 4096, 8192] -# Defaults preserve the original benchmark contract. CLI overrides make it -# possible to replay an exact model communication shape. -DEFAULT_HIDDEN_SIZE = 8192 -DEFAULT_BENCHMARK_DTYPE = torch.bfloat16 +# Fixed hidden size and dtype for all benchmarks +HIDDEN_SIZE = 8192 +BENCHMARK_DTYPE = torch.bfloat16 # CUDA graph settings CUDA_GRAPH_CAPTURE_CYCLES = 10 @@ -68,20 +67,16 @@ def __init__( device: torch.device, cpu_group: ProcessGroup, sequence_lengths: list[int], - hidden_size: int, - dtype: torch.dtype, ): self.rank = rank self.world_size = world_size self.device = device self.cpu_group = cpu_group - self.hidden_size = hidden_size - self.dtype = dtype # Calculate max_size_override based on largest sequence length max_seq_len = max(sequence_lengths) - max_tensor_elements = max_seq_len * self.hidden_size - self.max_size_override = max_tensor_elements * self.dtype.itemsize + 1 + max_tensor_elements = max_seq_len * HIDDEN_SIZE + self.max_size_override = max_tensor_elements * BENCHMARK_DTYPE.itemsize + 1 # Initialize communicators self.custom_allreduce = None @@ -342,10 +337,7 @@ def benchmark_allreduce_single( try: # Create test tensor (2D: sequence_length x hidden_size) tensor = torch.randn( - sequence_length, - self.hidden_size, - dtype=self.dtype, - device=self.device, + sequence_length, HIDDEN_SIZE, dtype=BENCHMARK_DTYPE, device=self.device ) if not should_use_fn(tensor): return None @@ -413,17 +405,16 @@ def _calculate_speedup_info(comm_results: dict[str, float]) -> str: def print_results( - results: dict[str, dict[str, float]], - sequence_lengths: list[int], - world_size: int, - hidden_size: int, - dtype: torch.dtype, + results: dict[str, dict[str, float]], sequence_lengths: list[int], world_size: int ): """Print benchmark results in a formatted table.""" print(f"\n{'=' * 130}") print("Device Communicator Benchmark Results") - print(f"World Size: {world_size}, Data Type: {dtype}, Hidden Size: {hidden_size}") + print( + f"World Size: {world_size}, Data Type: {BENCHMARK_DTYPE}, " + f"Hidden Size: {HIDDEN_SIZE}" + ) print(f"{'=' * 130}") # Get all communicator names @@ -445,15 +436,15 @@ def print_results( for seq_len in sequence_lengths: if seq_len in results: # Calculate tensor size in elements and bytes - tensor_elements = seq_len * hidden_size - tensor_bytes = tensor_elements * dtype.itemsize + tensor_elements = seq_len * HIDDEN_SIZE + tensor_bytes = tensor_elements * BENCHMARK_DTYPE.itemsize # Format tensor size (MB) tensor_size_mb = tensor_bytes / (1024 * 1024) tensor_size_str = f"{tensor_size_mb:.2f} MB" # Format tensor shape - tensor_shape = f"({seq_len}, {hidden_size})" + tensor_shape = f"({seq_len}, {HIDDEN_SIZE})" row = f"{tensor_shape:<20}{tensor_size_str:<15}" for comm in all_comms: @@ -492,24 +483,9 @@ def main(): "--num-trials", type=int, default=50, help="Number of benchmark trials" ) - parser.add_argument( - "--hidden-size", - type=int, - default=DEFAULT_HIDDEN_SIZE, - help="Hidden dimension of the all-reduce tensor.", - ) - - parser.add_argument( - "--dtype", - choices=("float16", "bfloat16"), - default="bfloat16", - help="All-reduce tensor dtype.", - ) - parser.add_argument("--output-json", type=str, help="Output results to JSON file") args = parser.parse_args() - benchmark_dtype = getattr(torch, args.dtype) # Initialize distributed if not dist.is_initialized(): @@ -530,13 +506,7 @@ def main(): # Initialize benchmark benchmark = CommunicatorBenchmark( - rank, - world_size, - device, - cpu_group, - args.sequence_lengths, - args.hidden_size, - benchmark_dtype, + rank, world_size, device, cpu_group, args.sequence_lengths ) # Run benchmarks @@ -548,7 +518,7 @@ def main(): "Benchmarking sequence length: %s (tensor shape: %s x %s)", seq_len, seq_len, - args.hidden_size, + HIDDEN_SIZE, ) results = benchmark.benchmark_allreduce( @@ -564,13 +534,7 @@ def main(): # Print results (only rank 0) if rank == 0: - print_results( - all_results, - args.sequence_lengths, - world_size, - args.hidden_size, - benchmark_dtype, - ) + print_results(all_results, args.sequence_lengths, world_size) # Save to JSON if requested if args.output_json: @@ -584,8 +548,8 @@ def main(): output_data = { "world_size": world_size, - "dtype": str(benchmark_dtype), - "hidden_size": args.hidden_size, + "dtype": str(BENCHMARK_DTYPE), + "hidden_size": HIDDEN_SIZE, "sequence_lengths": args.sequence_lengths, "num_warmup": args.num_warmup, "num_trials": args.num_trials, diff --git a/benchmarks/kernels/benchmark_dspark_markov.py b/benchmarks/kernels/benchmark_dspark_markov.py index 23e3202761..df2185b35d 100644 --- a/benchmarks/kernels/benchmark_dspark_markov.py +++ b/benchmarks/kernels/benchmark_dspark_markov.py @@ -1,5 +1,4 @@ # 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 40604d9a08..850231e984 100644 --- a/benchmarks/kernels/benchmark_dspark_metadata.py +++ b/benchmarks/kernels/benchmark_dspark_metadata.py @@ -1,5 +1,4 @@ # 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 a9764bddc7..70db0a7892 100644 --- a/benchmarks/kernels/benchmark_dspark_sm70_attention.py +++ b/benchmarks/kernels/benchmark_dspark_sm70_attention.py @@ -1,5 +1,4 @@ # 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_fp8_kv.py b/benchmarks/kernels/benchmark_flash_v100_fp8_kv.py index 34d22ea3b5..8de8b616c3 100644 --- a/benchmarks/kernels/benchmark_flash_v100_fp8_kv.py +++ b/benchmarks/kernels/benchmark_flash_v100_fp8_kv.py @@ -43,7 +43,6 @@ "fp8_b1616_bridge", ) DECODE_VARIANTS = ( - "fp8_b16_scalar", "fp16_b784_xqa", "fp16_b784_scalar", "fp16_b1568_xqa", diff --git a/benchmarks/kernels/benchmark_sm70_fp8_graph_capture.py b/benchmarks/kernels/benchmark_sm70_fp8_graph_capture.py index d31121dd36..ba73f21797 100644 --- a/benchmarks/kernels/benchmark_sm70_fp8_graph_capture.py +++ b/benchmarks/kernels/benchmark_sm70_fp8_graph_capture.py @@ -1,5 +1,4 @@ # 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/kernels/benchmark_sm70_nvfp4_qpn2.py b/benchmarks/kernels/benchmark_sm70_nvfp4_qpn2.py deleted file mode 100644 index 4393434cc6..0000000000 --- a/benchmarks/kernels/benchmark_sm70_nvfp4_qpn2.py +++ /dev/null @@ -1,499 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Race QPN2 against the current SM70 NVFP4 path on real Qwen3.8 weights. - -This is a deliberately narrow verifier microbenchmark. It loads the TP-local -gate/up and down projection shards from one native NVFP4 layer, prepares both -the current TurboMind layout and the QPN2 fragment layout, and measures an M=8 -CUDA-graph replay. Weight loading and preparation are outside the timed -region. - -QPN2 is compiled from an explicitly supplied source file so this benchmark can -evaluate a pinned external implementation before it is admitted to the vLLM -build. The result records the source digest and numerical delta against the -current production operator. -""" - -from __future__ import annotations - -import argparse -import hashlib -import importlib.util -import json -import statistics -import sys -from collections.abc import Callable -from dataclasses import dataclass -from pathlib import Path -from typing import Any - -import torch -from safetensors import safe_open -from torch.utils.cpp_extension import load - - -@dataclass(frozen=True) -class Projection: - name: str - packed: torch.Tensor - scales: torch.Tensor - inverse_global_scale: float - gated_silu: bool - calls_per_round: int - - -QPN2_CONFIGS = { - # (K, N): (split K, independent accumulator chains) - (5120, 8704): (8, 2), - (4352, 5120): (16, 2), -} - - -def _sha256_file(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as source: - for chunk in iter(lambda: source.read(1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() - - -def _load_projection_shards( - model: Path, - layer_index: int, - tp_rank: int, - tp_size: int, -) -> tuple[Projection, Projection]: - path = model / "model.safetensors" - prefix = f"model.language_model.layers.{layer_index}.mlp" - intermediate_size = 17408 - hidden_size = 5120 - if intermediate_size % tp_size or hidden_size % tp_size: - raise ValueError("model dimensions must divide the tensor-parallel size") - intermediate_per_rank = intermediate_size // tp_size - packed_k_per_rank = intermediate_per_rank // 2 - scale_k_per_rank = intermediate_per_rank // 16 - n_start = tp_rank * intermediate_per_rank - n_end = n_start + intermediate_per_rank - packed_k_start = tp_rank * packed_k_per_rank - packed_k_end = packed_k_start + packed_k_per_rank - scale_k_start = tp_rank * scale_k_per_rank - scale_k_end = scale_k_start + scale_k_per_rank - - with safe_open(path, framework="pt", device="cpu") as tensors: - gate_packed = tensors.get_slice(f"{prefix}.gate_proj.weight_packed")[ - n_start:n_end, : - ] - up_packed = tensors.get_slice(f"{prefix}.up_proj.weight_packed")[ - n_start:n_end, : - ] - gate_scales = tensors.get_slice(f"{prefix}.gate_proj.weight_scale")[ - n_start:n_end, : - ] - up_scales = tensors.get_slice(f"{prefix}.up_proj.weight_scale")[ - n_start:n_end, : - ] - gate_global = tensors.get_tensor(f"{prefix}.gate_proj.weight_global_scale") - up_global = tensors.get_tensor(f"{prefix}.up_proj.weight_global_scale") - - down_packed = tensors.get_slice(f"{prefix}.down_proj.weight_packed")[ - :, packed_k_start:packed_k_end - ] - down_scales = tensors.get_slice(f"{prefix}.down_proj.weight_scale")[ - :, scale_k_start:scale_k_end - ] - down_global = tensors.get_tensor(f"{prefix}.down_proj.weight_global_scale") - - gate_up = Projection( - name="gate_up_proj", - packed=torch.cat((gate_packed, up_packed), dim=0).contiguous(), - scales=torch.cat((gate_scales, up_scales), dim=0).contiguous(), - inverse_global_scale=float( - 1.0 / torch.cat((gate_global, up_global)).max().float() - ), - gated_silu=True, - calls_per_round=56, - ) - down = Projection( - name="down_proj", - packed=down_packed.contiguous(), - scales=down_scales.contiguous(), - inverse_global_scale=float(1.0 / down_global.max().float()), - gated_silu=False, - calls_per_round=56, - ) - return gate_up, down - - -def _unpack_nibbles(weight_packed: torch.Tensor) -> torch.Tensor: - """Convert checkpoint [N,K/2] bytes to TurboMind's [K,N] codes.""" - return ( - torch.stack((weight_packed & 15, weight_packed >> 4), dim=-1) - .flatten(start_dim=-2) - .t() - .contiguous() - ) - - -def _qpn2_prepack( - codes: torch.Tensor, scales: torch.Tensor -) -> tuple[torch.Tensor, torch.Tensor]: - """Byte-preserving permutation into v100-skinny's QPN2 layout.""" - n, k2 = codes.shape - k = k2 * 2 - if n % 32 or k % 64: - raise ValueError(f"QPN2 requires N%32=0 and K%64=0, got N={n}, K={k}") - device = codes.device - tiles, groups = n // 32, k // 16 - lane = torch.arange(32, device=device) - column = ((lane >> 2) & 3) * 8 + (lane & 3) + ((lane & 16) > 0).long() * 4 - k_order = torch.tensor( - [0, 2, 4, 6, 1, 3, 5, 7, 8, 10, 12, 14, 9, 11, 13, 15], - device=device, - ) - nibbles = torch.stack((codes & 0xF, codes >> 4), dim=-1).view(n, k) - group = torch.arange(groups, device=device) - k_index = group.view(groups, 1) * 16 + k_order.view(1, 16) - q_codes = torch.empty((tiles, groups, 32, 8), dtype=torch.uint8, device=device) - q_scales = torch.empty((tiles, groups, 32), dtype=torch.uint8, device=device) - # Bound the int64 gather intermediates for the gate/up matrix. - tile_chunk = max(1, 36864 // groups) - for tile_start in range(0, tiles, tile_chunk): - tile_end = min(tile_start + tile_chunk, tiles) - tile_count = tile_end - tile_start - n_column = torch.arange(tile_start, tile_end, device=device).view( - tile_count, 1 - ) * 32 + column.view(1, 32) - nibble_block = nibbles[ - n_column.view(tile_count, 1, 32, 1).expand(tile_count, groups, 32, 16), - k_index.view(1, groups, 1, 16).expand(tile_count, groups, 32, 16), - ] - q_codes[tile_start:tile_end] = nibble_block[..., 0::2] | ( - nibble_block[..., 1::2] << 4 - ) - q_scales[tile_start:tile_end] = scales[ - n_column.view(tile_count, 1, 32).expand(tile_count, groups, 32), - group.view(1, groups, 1).expand(tile_count, groups, 32), - ] - return q_codes.view(-1).contiguous(), q_scales.view(-1).contiguous() - - -def _capture(call: Callable[[], Any]) -> tuple[torch.cuda.CUDAGraph, Any]: - for _ in range(10): - call() - torch.accelerator.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - output = call() - torch.accelerator.synchronize() - return graph, output - - -def _time_graph( - graph: torch.cuda.CUDAGraph, warmup: int, iterations: int, trials: int -) -> dict[str, float | list[float]]: - for _ in range(warmup): - graph.replay() - torch.accelerator.synchronize() - samples = [] - for _ in range(trials): - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - start.record() - for _ in range(iterations): - graph.replay() - end.record() - end.synchronize() - samples.append(float(start.elapsed_time(end)) * 1000.0 / iterations) - return { - "samples_us": samples, - "median_us": statistics.median(samples), - "min_us": min(samples), - "max_us": max(samples), - } - - -def _quality(actual: torch.Tensor, expected: torch.Tensor) -> dict[str, object]: - actual_f = actual.float() - expected_f = expected.float() - delta = actual_f - expected_f - denominator = torch.linalg.vector_norm(expected_f).clamp_min(1e-12) - return { - "finite": bool(torch.isfinite(actual).all()), - "exact": bool(torch.equal(actual, expected)), - "different_fraction": float((actual != expected).float().mean()), - "max_abs": float(delta.abs().max()), - "mean_abs": float(delta.abs().mean()), - "relative_l2": float(torch.linalg.vector_norm(delta) / denominator), - "cosine": float( - torch.nn.functional.cosine_similarity(actual_f, expected_f, dim=1) - .mean() - .item() - ), - } - - -def _fp32_reference( - projection: Projection, x: torch.Tensor, device: torch.device -) -> torch.Tensor: - """Evaluate the checkpoint's native E2M1 codes with FP32 accumulation.""" - packed = projection.packed.to(device) - nibbles = torch.stack((packed & 0xF, packed >> 4), dim=-1).flatten(start_dim=-2) - magnitudes = torch.tensor( - [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], - dtype=torch.float32, - device=device, - ) - weights = magnitudes[(nibbles & 7).long()] - weights = torch.where((nibbles & 8) != 0, -weights, weights) - weights.mul_( - projection.scales.float() - .to(device) - .repeat_interleave(16, dim=1) - .mul_(projection.inverse_global_scale) - ) - raw = x.float().matmul(weights.t()) - if not projection.gated_silu: - return raw - gate, up = raw.chunk(2, dim=1) - return torch.nn.functional.silu(gate).mul_(up) - - -def _run_projection( - projection: Projection, - extension: Any | None, - production: bool, - m: int, - device: torch.device, - warmup: int, - iterations: int, - trials: int, -) -> dict[str, object]: - from vllm import _sm70_ops as sm70_ops - - packed = projection.packed.to(device) - scale_codes = projection.scales.view(torch.uint8).to(device) - qweight = _unpack_nibbles(projection.packed).to(device) - effective_scales = ( - projection.scales.t() - .float() - .mul(projection.inverse_global_scale) - .half() - .contiguous() - .to(device) - ) - k, n = qweight.shape - split_k, accumulator_chains = QPN2_CONFIGS[(k, n)] - x = torch.randn((m, k), dtype=torch.float16, device=device) * 0.1 - - tm_weight, tm_scales, meta = sm70_ops.nvfp4_sm70_prepare( - qweight, effective_scales, 16, False - ) - k_ld, q_ld = (int(value.item()) for value in meta[:2]) - tm_raw = torch.empty((m, n), dtype=torch.float16, device=device) - final_n = n // 2 if projection.gated_silu else n - tm_final = ( - torch.empty((m, final_n), dtype=torch.float16, device=device) - if projection.gated_silu - else tm_raw - ) - - def run_tm() -> torch.Tensor: - sm70_ops.nvfp4_gemm_sm70_out( - tm_raw, x, tm_weight, tm_scales, 16, k_ld, q_ld, False - ) - if projection.gated_silu: - torch.ops._C.silu_and_mul(tm_final, tm_raw) - return tm_final - - qpn_final = torch.empty((m, final_n), dtype=torch.float16, device=device) - if production: - qpn_codes, qpn_scales = sm70_ops.nvfp4_qpn2_prepare_sm70( - packed, projection.scales.to(device) - ) - - def run_qpn2() -> torch.Tensor: - if projection.gated_silu: - sm70_ops.nvfp4_qpn2_gated_sm70_out( - qpn_final, - x, - qpn_codes, - qpn_scales, - projection.inverse_global_scale, - split_k, - accumulator_chains, - ) - else: - sm70_ops.nvfp4_qpn2_gemm_sm70_out( - qpn_final, - x, - qpn_codes, - qpn_scales, - projection.inverse_global_scale, - split_k, - accumulator_chains, - ) - return qpn_final - - else: - if extension is None: - raise AssertionError("external QPN2 extension was not loaded") - qpn_codes, qpn_scales = _qpn2_prepack(packed, scale_codes) - - def run_qpn2() -> torch.Tensor: - qpn_raw = extension.gemm_qpn2( - x, - qpn_codes, - qpn_scales, - projection.inverse_global_scale, - n, - split_k, - accumulator_chains, - ) - if projection.gated_silu: - torch.ops._C.silu_and_mul(qpn_final, qpn_raw) - return qpn_final - return qpn_raw - - tm_graph, tm_output = _capture(run_tm) - qpn_graph, qpn_output = _capture(run_qpn2) - tm_timing = _time_graph(tm_graph, warmup, iterations, trials) - qpn_timing = _time_graph(qpn_graph, warmup, iterations, trials) - tm_graph.replay() - qpn_graph.replay() - torch.accelerator.synchronize() - reference = _fp32_reference(projection, x, device) - saved_us = float(tm_timing["median_us"]) - float(qpn_timing["median_us"]) - return { - "name": projection.name, - "m": m, - "n": n, - "k": k, - "gated_silu": projection.gated_silu, - "calls_per_round": projection.calls_per_round, - "qpn2_config": { - "split_k": split_k, - "accumulator_chains": accumulator_chains, - }, - "turbomind": tm_timing, - "qpn2": qpn_timing, - "speedup": float(tm_timing["median_us"]) / float(qpn_timing["median_us"]), - "saved_ms_per_round": saved_us * projection.calls_per_round / 1000.0, - "quality_vs_turbomind": _quality(qpn_output, tm_output), - "turbomind_quality_vs_fp32": _quality(tm_output, reference), - "qpn2_quality_vs_fp32": _quality(qpn_output, reference), - } - - -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser() - parser.add_argument( - "--model", type=Path, default=Path("/home/ymzx/models/Qwen3.8-27B-NVFP4") - ) - parser.add_argument("--layer", type=int, default=55) - parser.add_argument("--tp-rank", type=int, default=0) - parser.add_argument("--tp-size", type=int, default=4) - parser.add_argument("--m", type=int, default=8) - parser.add_argument("--device", default="cuda:0") - parser.add_argument("--warmup", type=int, default=50) - parser.add_argument("--iterations", type=int, default=1000) - parser.add_argument("--trials", type=int, default=5) - parser.add_argument("--qpn2-source", type=Path) - parser.add_argument( - "--production-library", - type=Path, - help="Load a complete candidate vllm._C containing production QPN2 ops.", - ) - parser.add_argument("--extension-name", default="v100_skinny_qpn2_micro_v1") - parser.add_argument("--json-out", type=Path, required=True) - return parser.parse_args() - - -def _load_production_library(path: Path) -> None: - spec = importlib.util.spec_from_file_location("vllm._C", path) - if spec is None or spec.loader is None: - raise RuntimeError(f"cannot load production library: {path}") - module = importlib.util.module_from_spec(spec) - sys.modules["vllm._C"] = module - spec.loader.exec_module(module) - - -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, - ): - raise RuntimeError("benchmark requires an exact SM70 CUDA device") - if not 1 <= args.m <= 8: - raise ValueError("QPN2 benchmark supports M in [1, 8]") - production = args.production_library is not None - if production: - _load_production_library(args.production_library) - extension = None - else: - if args.qpn2_source is None: - raise ValueError("--qpn2-source is required without --production-library") - extension = load( - name=args.extension_name, - sources=[str(args.qpn2_source)], - extra_cuda_cflags=[ - "-O3", - "--use_fast_math", - "-lineinfo", - "-gencode=arch=compute_70,code=sm_70", - ], - verbose=False, - ) - torch.manual_seed(20260824) - projections = _load_projection_shards( - args.model, args.layer, args.tp_rank, args.tp_size - ) - rows = [ - _run_projection( - projection, - extension, - production, - args.m, - device, - args.warmup, - args.iterations, - args.trials, - ) - for projection in projections - ] - payload = { - "device": torch.cuda.get_device_name(device), - "capability": list(torch.cuda.get_device_capability(device)), - "torch": torch.__version__, - "model": str(args.model), - "model_config_sha256": _sha256_file(args.model / "config.json"), - "layer": args.layer, - "tp_rank": args.tp_rank, - "tp_size": args.tp_size, - "qpn2_source": str(args.qpn2_source) if args.qpn2_source else None, - "qpn2_source_sha256": ( - _sha256_file(args.qpn2_source) if args.qpn2_source else None - ), - "production_library": ( - str(args.production_library) if args.production_library else None - ), - "production_library_sha256": ( - _sha256_file(args.production_library) if args.production_library else None - ), - "warmup": args.warmup, - "iterations": args.iterations, - "trials": args.trials, - "rows": rows, - "projected_total_saved_ms_per_round": sum( - float(row["saved_ms_per_round"]) for row in rows - ), - } - args.json_out.parent.mkdir(parents=True, exist_ok=True) - args.json_out.write_text(json.dumps(payload, indent=2) + "\n") - print(json.dumps(payload, indent=2), flush=True) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/benchmarks/profile_sm70_fp8_qpn8_gated_pair_ncu.py b/benchmarks/profile_sm70_fp8_qpn8_gated_pair_ncu.py deleted file mode 100644 index 3070b16236..0000000000 --- a/benchmarks/profile_sm70_fp8_qpn8_gated_pair_ncu.py +++ /dev/null @@ -1,140 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Paired NCU probe for current and QPN8 fused gate/up kernels.""" - -from __future__ import annotations - -import argparse -import json -from pathlib import Path - -import torch -import vllm._C # noqa: F401 -from benchmark_sm70_fp8_qpn8 import ( - _dequantized_weight, - _error_stats, - _load_case, - _qpn8_group_scales, - _qpn8_prepack, -) - -from vllm import _sm70_ops as sm70_ops - - -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser() - parser.add_argument( - "--model", - type=Path, - default=Path("/home/ymzx/models/Qwen3.8-27B-FP8"), - ) - parser.add_argument("--qpn8-library", type=Path, required=True) - parser.add_argument("--out", type=Path, required=True) - parser.add_argument("--m", type=int, default=1) - parser.add_argument("--warmup", type=int, default=20) - parser.add_argument("--repeats", type=int, default=100) - return parser.parse_args() - - -def _time(launch, warmup: int, repeats: int) -> float: - for _ in range(warmup): - launch() - torch.accelerator.synchronize() - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - start.record() - for _ in range(repeats): - launch() - end.record() - end.synchronize() - return float(start.elapsed_time(end) * 1000.0 / repeats) - - -def main() -> int: - args = _parse_args() - torch.ops.load_library(str(args.qpn8_library.resolve())) - device = torch.device("cuda:0") - torch.accelerator.set_device_index(device.index or 0) - if torch.cuda.get_device_capability(device) != (7, 0): - raise RuntimeError("This profiler requires SM70/V100") - - qweight, scales, _ = _load_case(args.model, "gate_up_fused", 4, 0, device) - n, k = (int(dim) for dim in qweight.shape) - hidden = n // 2 - codes = _qpn8_prepack(qweight.view(torch.uint8)) - group_scales = _qpn8_group_scales(scales, n, k) - tm_weight, tm_scales, meta = sm70_ops.fp8_sm70_prepare(qweight, scales, 128, True) - k_ld, q_ld = (int(value) for value in meta.tolist()) - inputs = torch.randn((args.m, k), device=device, dtype=torch.float16).mul_(0.1) - current_out = torch.empty((args.m, hidden), device=device, dtype=torch.float16) - qpn_out = torch.empty_like(current_out) - - def launch_current() -> None: - sm70_ops.fp8_gemm_sm70_out( - current_out, - inputs, - tm_weight, - tm_scales, - 128, - k_ld, - q_ld, - True, - ) - - def launch_qpn() -> None: - sm70_ops.fp8_qpn8_gated_pair_sm70_out( - qpn_out, - inputs, - codes, - group_scales, - 8, - 2, - True, - True, - ) - - current_us = _time(launch_current, args.warmup, args.repeats) - qpn_us = _time(launch_qpn, args.warmup, args.repeats) - reference_raw = inputs.float().matmul(_dequantized_weight(qweight, scales).t()) - gate, up = reference_raw.chunk(2, dim=1) - reference = (torch.nn.functional.silu(gate) * up).half() - quality = { - "current": _error_stats(current_out, reference), - "qpn8_pair": _error_stats(qpn_out, reference), - } - - torch.cuda.cudart().cudaProfilerStart() - torch.cuda.nvtx.range_push("current_gate_up_fused") - launch_current() - torch.cuda.nvtx.range_pop() - torch.cuda.nvtx.range_push("qpn8_gate_up_fused_pair") - launch_qpn() - torch.cuda.nvtx.range_pop() - torch.accelerator.synchronize(device) - torch.cuda.cudart().cudaProfilerStop() - - result = { - "device": torch.cuda.get_device_name(device), - "m": args.m, - "n": n, - "k": k, - "hidden": hidden, - "current_unprofiled_us": current_us, - "qpn8_pair_unprofiled_us": qpn_us, - "speedup": current_us / qpn_us, - "qpn8_config": { - "split_k": 8, - "nacc": 2, - "decoder": "fast", - "prefetch": True, - }, - "quality": quality, - } - args.out.parent.mkdir(parents=True, exist_ok=True) - args.out.write_text(json.dumps(result, indent=2), encoding="utf-8") - print(json.dumps(result, indent=2)) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/benchmarks/profile_sm70_fp8_qpn8_ncu.py b/benchmarks/profile_sm70_fp8_qpn8_ncu.py deleted file mode 100644 index 7bc5e02c55..0000000000 --- a/benchmarks/profile_sm70_fp8_qpn8_ncu.py +++ /dev/null @@ -1,194 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Profile only the real-shape QPN8 winners with Nsight Compute.""" - -from __future__ import annotations - -import argparse -import json -from pathlib import Path -from typing import Any - -import torch -import vllm._C # noqa: F401 -from benchmark_sm70_fp8_qpn8 import ( - _load_case, - _qpn8_group_scales, - _qpn8_prepack, -) - -from vllm import _sm70_ops as sm70_ops - -_CONFIGS = { - "down": (16, 1, True), - "gdn_in": (32, 1, True), - "output": (16, 1, True), - "full_qkv": (16, 1, True), - "gate_up_raw": (16, 1, True), -} - - -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser() - parser.add_argument( - "--model", - type=Path, - default=Path("/home/ymzx/models/Qwen3.8-27B-FP8"), - ) - parser.add_argument("--qpn8-library", type=Path, required=True) - parser.add_argument("--out", type=Path, required=True) - parser.add_argument("--device", default="cuda:0") - parser.add_argument("--m", type=int, default=1) - parser.add_argument("--warmup", type=int, default=20) - parser.add_argument("--repeats", type=int, default=100) - parser.add_argument("--paired-current", action="store_true") - return parser.parse_args() - - -def _time(launch: Any, warmup: int, repeats: int) -> float: - for _ in range(warmup): - launch() - torch.accelerator.synchronize() - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - start.record() - for _ in range(repeats): - launch() - end.record() - end.synchronize() - return float(start.elapsed_time(end) * 1000.0 / repeats) - - -def main() -> int: - args = _parse_args() - torch.ops.load_library(str(args.qpn8_library.resolve())) - device = torch.device(args.device) - torch.accelerator.set_device_index(device.index or 0) - if torch.cuda.get_device_capability(device) != (7, 0): - raise RuntimeError("This profiler requires SM70/V100") - if args.m < 1 or args.m > 8: - raise ValueError("M must be in [1, 8]") - - prepared: list[dict[str, Any]] = [] - for case, (split_k, nacc, fast_decoder) in _CONFIGS.items(): - qweight, scales, note = _load_case(args.model, case, 4, 0, device) - n, k = (int(dim) for dim in qweight.shape) - codes = _qpn8_prepack(qweight.view(torch.uint8)) - group_scales = _qpn8_group_scales(scales, n, k) - tm_weight, tm_scales, meta = sm70_ops.fp8_sm70_prepare( - qweight, scales, 128, False - ) - k_ld, q_ld = (int(value) for value in meta.tolist()) - inputs = torch.randn((args.m, k), device=device, dtype=torch.float16).mul_(0.1) - qpn_output = torch.empty((args.m, n), device=device, dtype=torch.float16) - current_output = torch.empty_like(qpn_output) - - def launch_qpn( - output: torch.Tensor = qpn_output, - inputs: torch.Tensor = inputs, - codes: torch.Tensor = codes, - group_scales: torch.Tensor = group_scales, - split_k: int = split_k, - nacc: int = nacc, - fast_decoder: bool = fast_decoder, - ) -> None: - sm70_ops.fp8_qpn8_gemm_sm70_out( - output, - inputs, - codes, - group_scales, - split_k, - nacc, - fast_decoder, - False, - ) - - def launch_current( - output: torch.Tensor = current_output, - inputs: torch.Tensor = inputs, - tm_weight: torch.Tensor = tm_weight, - tm_scales: torch.Tensor = tm_scales, - k_ld: int = k_ld, - q_ld: int = q_ld, - ) -> None: - sm70_ops.fp8_gemm_sm70_out( - output, - inputs, - tm_weight, - tm_scales, - 128, - k_ld, - q_ld, - False, - ) - - prepared.append( - { - "case": case, - "note": note, - "m": args.m, - "n": n, - "k": k, - "split_k": split_k, - "nacc": nacc, - "decoder": "fast" if fast_decoder else "scalar", - "launch_qpn": launch_qpn, - "launch_current": launch_current, - "tensors": [ - qweight, - scales, - codes, - group_scales, - inputs, - qpn_output, - current_output, - tm_weight, - tm_scales, - ], - } - ) - - for item in prepared: - item["qpn_unprofiled_us"] = _time(item["launch_qpn"], args.warmup, args.repeats) - if args.paired_current: - item["current_unprofiled_us"] = _time( - item["launch_current"], args.warmup, args.repeats - ) - - torch.cuda.cudart().cudaProfilerStart() - torch.cuda.nvtx.range_push("qpn8_real_shape_winners") - for item in prepared: - if args.paired_current: - torch.cuda.nvtx.range_push(f"current_{item['case']}") - item["launch_current"]() - torch.cuda.nvtx.range_pop() - torch.cuda.nvtx.range_push(f"qpn8_{item['case']}") - item["launch_qpn"]() - torch.cuda.nvtx.range_pop() - torch.accelerator.synchronize(device) - torch.cuda.nvtx.range_pop() - torch.cuda.cudart().cudaProfilerStop() - - public = [ - { - key: value - for key, value in item.items() - if key not in ("launch_qpn", "launch_current", "tensors") - } - for item in prepared - ] - result = { - "device": torch.cuda.get_device_name(device), - "capability": list(torch.cuda.get_device_capability(device)), - "torch": torch.__version__, - "torch_cuda": torch.version.cuda, - "profiles": public, - } - args.out.parent.mkdir(parents=True, exist_ok=True) - args.out.write_text(json.dumps(result, indent=2), encoding="utf-8") - print(json.dumps(result, indent=2)) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/benchmarks/verify_sm70_artifact.py b/benchmarks/verify_sm70_artifact.py index 81dde92497..dfca12798f 100644 --- a/benchmarks/verify_sm70_artifact.py +++ b/benchmarks/verify_sm70_artifact.py @@ -75,14 +75,18 @@ ("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), @@ -101,7 +105,8 @@ ("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), ) @@ -116,7 +121,8 @@ ("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 = ( @@ -132,7 +138,8 @@ ("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 = ( @@ -146,7 +153,9 @@ ("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), @@ -272,7 +281,9 @@ 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}") @@ -328,13 +339,17 @@ 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, @@ -394,12 +409,10 @@ 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", + }], } @@ -516,7 +529,8 @@ def _model_quality_failures(payload: Any) -> list[str]: failures = [] if gate.get("label") != "model-pass": failures.append( - f"model_quality_gate.label is not model-pass ({gate.get('label')!r})" + "model_quality_gate.label is not model-pass " + f"({gate.get('label')!r})" ) if gate.get("default_acceptance") != "model-level gate passed": failures.append( @@ -572,7 +586,8 @@ def _inventory_complete_failures(payload: Any) -> list[str]: failures = [] if gate.get("label") != "inventory-complete": failures.append( - f"inventory_gate.label is not inventory-complete ({gate.get('label')!r})" + "inventory_gate.label is not inventory-complete " + f"({gate.get('label')!r})" ) if gate.get("complete") is not True: failures.append("inventory_gate.complete is not true") @@ -580,7 +595,8 @@ 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( - f"inventory_gate.old_only_env_vars is not empty: {old_only_env_vars!r}" + "inventory_gate.old_only_env_vars is not empty: " + f"{old_only_env_vars!r}" ) if gate.get("old_only_env_vars_count") != 0: failures.append( @@ -631,13 +647,11 @@ 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 = [] @@ -664,12 +678,10 @@ 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 = [] @@ -677,12 +689,10 @@ 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 = [] @@ -690,12 +700,10 @@ 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 = [] @@ -703,12 +711,10 @@ 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 [] @@ -717,7 +723,8 @@ 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 ( @@ -881,7 +888,9 @@ 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", @@ -1008,7 +1017,9 @@ 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) @@ -1023,7 +1034,9 @@ 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: @@ -1044,17 +1057,27 @@ 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") @@ -1062,8 +1085,12 @@ 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 ) @@ -1071,17 +1098,19 @@ 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( @@ -1098,7 +1127,10 @@ 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/external_projects/vllm_flash_attn.cmake b/cmake/external_projects/vllm_flash_attn.cmake index 8918d9e7b0..331556d232 100644 --- a/cmake/external_projects/vllm_flash_attn.cmake +++ b/cmake/external_projects/vllm_flash_attn.cmake @@ -50,12 +50,6 @@ elseif(VLLM_FLASH_ATTN_SM70) COMMAND ${PATCH_EXECUTABLE} --batch --forward -p1 -l -i ${CMAKE_CURRENT_LIST_DIR}/../patches/sm70_flash_attn_d256_splitkv3.patch - COMMAND - ${PATCH_EXECUTABLE} --batch --forward -p1 -l - -i ${CMAKE_CURRENT_LIST_DIR}/../patches/sm70_flash_attn_d256_k_pingpong.patch - COMMAND - ${PATCH_EXECUTABLE} --batch --forward -p1 -l - -i ${CMAKE_CURRENT_LIST_DIR}/../patches/sm70_flash_attn_d256_gqa_arch.patch BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn ) else() diff --git a/cmake/patches/sm70_flash_attn_d256_gqa_arch.patch b/cmake/patches/sm70_flash_attn_d256_gqa_arch.patch deleted file mode 100644 index d6d65439a5..0000000000 --- a/cmake/patches/sm70_flash_attn_d256_gqa_arch.patch +++ /dev/null @@ -1,1673 +0,0 @@ -diff --git a/CMakeLists.txt b/CMakeLists.txt -index acf2219..f5e2732 100644 ---- a/CMakeLists.txt -+++ b/CMakeLists.txt -@@ -84,6 +84,7 @@ string(REPLACE "-O2" "-O3" CMAKE_CUDA_FLAGS_RELWITHDEBINFO "${CMAKE_CUDA_FLAGS_R - - if (FA2_ENABLED) - set(FA2_GEN_SRCS -+ "csrc/flash_attn/src/flash_fwd_d256_gqa_arch_sm70.cu" - "csrc/flash_attn/src/flash_fwd_d256_splitd_sm70.cu" - "csrc/flash_attn/src/flash_fwd_hdim256_causal_prefill_sm70.cu" - "csrc/flash_attn/src/flash_fwd_split_hdim256_causal_prefill_sm70.cu") -@@ -114,7 +115,8 @@ if (FA2_ENABLED) - csrc/flash_attn - csrc/flash_attn/src - csrc/common -- csrc/cutlass/include) -+ csrc/cutlass/include -+ csrc/cutlass/examples/35_gemm_softmax) - - # custom definitions - target_compile_definitions(_vllm_fa2_C PRIVATE -@@ -125,5 +127,16 @@ if (FA2_ENABLED) - # FLASHATTENTION_DISABLE_UNEVEN_K - # FLASHATTENTION_DISABLE_LOCAL - FLASHATTENTION_DISABLE_PYBIND -+ PREFIX_TORCH_EXTENSION -+ PREFIX_QK_FULL_STATS -+ PREFIX_QK_SKIP_APPLY -+ QK_TB_M=128 -+ QK_TB_N=512 -+ QK_WARP_M=64 -+ QK_WARP_N=128 -+ PV_TB_M=64 -+ PV_TB_N=256 -+ PV_WARP_M=32 -+ PV_WARP_N=64 - ) - endif () -diff --git a/csrc/flash_attn/flash_api_torch_lib.cpp b/csrc/flash_attn/flash_api_torch_lib.cpp -index 4b7601a..f5b71a2 100644 ---- a/csrc/flash_attn/flash_api_torch_lib.cpp -+++ b/csrc/flash_attn/flash_api_torch_lib.cpp -@@ -33,6 +33,24 @@ at::Tensor sm70_d256_splitd_dense_splitkv3_fwd( - double softmax_scale, - bool causal); - -+at::Tensor sm70_d256_splitd_dense_state_fwd( -+ const at::Tensor &q, -+ const at::Tensor &k, -+ const at::Tensor &v, -+ at::Tensor &state_max, -+ at::Tensor &state_sum, -+ at::Tensor &out, -+ double softmax_scale, -+ bool causal); -+ -+at::Tensor sm70_d256_gqa_architecture_fwd( -+ const at::Tensor &q, -+ const at::Tensor &k, -+ const at::Tensor &v, -+ at::Tensor &out, -+ double softmax_scale, -+ bool causal); -+ - at::Tensor sm70_d256_splitd_paged_fwd( - const at::Tensor &q, - const at::Tensor &k, -@@ -98,6 +116,17 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { - ops.impl("sm70_d256_splitd_n32_dense_fwd", torch::kCUDA, - make_pytorch_shim(&sm70_d256_splitd_dense_fwd)); - -+ ops.def("sm70_d256_splitd_n32_dense_state_fwd(Tensor q, Tensor k, " -+ "Tensor v, Tensor(a!) state_max, Tensor(b!) state_sum, " -+ "Tensor(c!) out, float softmax_scale, bool causal) -> Tensor(c!)"); -+ ops.impl("sm70_d256_splitd_n32_dense_state_fwd", torch::kCUDA, -+ make_pytorch_shim(&sm70_d256_splitd_dense_state_fwd)); -+ -+ ops.def("sm70_d256_gqa_architecture_fwd(Tensor q, Tensor k, Tensor v, " -+ "Tensor(a!) out, float softmax_scale, bool causal) -> Tensor(a!)"); -+ ops.impl("sm70_d256_gqa_architecture_fwd", torch::kCUDA, -+ make_pytorch_shim(&sm70_d256_gqa_architecture_fwd)); -+ - ops.def("sm70_d256_splitd_n32_dense_splitkv3_fwd(Tensor q, Tensor k, " - "Tensor v, Tensor(a!) partial_out, Tensor(b!) partial_max, " - "Tensor(c!) partial_sum, Tensor(d!) out, float softmax_scale, " -diff --git a/csrc/flash_attn/src/flash_fwd_d256_gqa_arch_sm70.cu b/csrc/flash_attn/src/flash_fwd_d256_gqa_arch_sm70.cu -new file mode 100644 -index 0000000..197453f ---- /dev/null -+++ b/csrc/flash_attn/src/flash_fwd_d256_gqa_arch_sm70.cu -@@ -0,0 +1,1359 @@ -+// SPDX-License-Identifier: BSD-3-Clause -+ -+/*************************************************************************************************** -+ * End-to-end prefix architecture screen for V100/SM70. -+ * -+ * QK GEMM writes FP16 logits and row statistics. PV normalizes logits in its -+ * A-operand transform and writes one reusable FP16 block partial. A float -+ * online accumulator folds that partial into prefix softmax state before the -+ * next block, avoiding resident storage for every block partial. -+ **************************************************************************************************/ -+ -+#include -+#include -+#include -+ -+#include -+#include -+#include -+#include -+#include -+#include -+#include -+#include -+#include -+#include -+ -+#if defined(PREFIX_TORCH_EXTENSION) -+#include -+#include -+#include -+#include -+#include "namespace_config.h" -+#endif -+ -+#include "cutlass/cutlass.h" -+#include "cutlass/device_kernel.h" -+#include "cutlass/epilogue/thread/linear_combination.h" -+#include "cutlass/gemm/kernel/default_gemm.h" -+#include "cutlass/gemm/kernel/gemm.h" -+#include "cutlass/gemm/threadblock/mma_pipelined.h" -+#include "cutlass/layout/matrix.h" -+#include "cutlass/numeric_conversion.h" -+#include "gemm_with_softmax.h" -+ -+namespace { -+ -+using Element = cutlass::half_t; -+ -+#ifndef QK_TB_M -+#define QK_TB_M 128 -+#endif -+#ifndef QK_TB_N -+#define QK_TB_N 128 -+#endif -+#ifndef QK_WARP_M -+#define QK_WARP_M 32 -+#endif -+#ifndef QK_WARP_N -+#define QK_WARP_N 64 -+#endif -+ -+using QKLayoutA = cutlass::layout::RowMajor; -+using QKLayoutB = cutlass::layout::ColumnMajor; -+using QKThreadblockShape = cutlass::gemm::GemmShape; -+using QKWarpShape = cutlass::gemm::GemmShape; -+using QKInstructionShape = cutlass::gemm::GemmShape<8, 8, 4>; -+using QKOutputOp = cutlass::epilogue::thread::LinearCombination< -+ Element, 8, Element, Element>; -+using QKGemm = cutlass::GemmSoftmax< -+ Element, -+ QKLayoutA, -+ Element, -+ QKLayoutB, -+ Element, -+ Element, -+ cutlass::arch::OpClassTensorOp, -+ cutlass::arch::Sm70, -+ QKThreadblockShape, -+ QKWarpShape, -+ QKInstructionShape, -+ QKOutputOp, -+ 2, -+ cutlass::MatrixShape<1, 1024>>; -+ -+#ifndef PV_TB_M -+#define PV_TB_M 64 -+#endif -+#ifndef PV_TB_N -+#define PV_TB_N 128 -+#endif -+#ifndef PV_WARP_M -+#define PV_WARP_M 32 -+#endif -+#ifndef PV_WARP_N -+#define PV_WARP_N 64 -+#endif -+ -+using PVLayout = cutlass::layout::RowMajor; -+using PVAccumulator = Element; -+using PVThreadblockShape = cutlass::gemm::GemmShape; -+using PVWarpShape = cutlass::gemm::GemmShape; -+using PVInstructionShape = cutlass::gemm::GemmShape<8, 8, 4>; -+using PVSwizzle = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>; -+using PVOutputOp = cutlass::epilogue::thread::LinearCombination< -+ Element, 8, PVAccumulator, float>; -+ -+constexpr int kPVAlignment = 8; -+ -+using PVDefaultKernel = typename cutlass::gemm::kernel::DefaultGemm< -+ Element, -+ PVLayout, -+ kPVAlignment, -+ Element, -+ PVLayout, -+ kPVAlignment, -+ Element, -+ PVLayout, -+ PVAccumulator, -+ cutlass::arch::OpClassTensorOp, -+ cutlass::arch::Sm70, -+ PVThreadblockShape, -+ PVWarpShape, -+ PVInstructionShape, -+ PVOutputOp, -+ PVSwizzle, -+ 2, -+ false, -+ cutlass::arch::OpMultiplyAdd>::GemmKernel; -+ -+using PVDefaultMma = typename PVDefaultKernel::Mma; -+using PVIteratorA = typename PVDefaultMma::IteratorA; -+using PVIteratorB = typename PVDefaultMma::IteratorB; -+using PVSmemIteratorA = typename PVDefaultMma::SmemIteratorA; -+using PVSmemIteratorB = typename PVDefaultMma::SmemIteratorB; -+ -+__device__ float const* g_row_max = nullptr; -+__device__ float const* g_row_inv_sum = nullptr; -+__device__ float* g_row_sum_out = nullptr; -+__device__ int g_rows = 0; -+ -+#if defined(PREFIX_QK_FULL_STATS) -+struct ExpRowSumTransformA { -+ using InputFragment = typename PVIteratorA::Fragment; -+ using OutputFragment = cutlass::Array< -+ typename PVSmemIteratorA::Element, InputFragment::kElements>; -+ using ThreadMap = typename PVIteratorA::ThreadMap; -+ static constexpr int kAccessesPerVector = -+ PVIteratorA::UnderlyingIterator::kAccessesPerVector; -+ static constexpr int kContiguousIterations = -+ ThreadMap::Iterations::kContiguous; -+ static constexpr int kStridedIterations = ThreadMap::Iterations::kStrided; -+ -+ float row_max[kStridedIterations]; -+ float row_inv_sum[kStridedIterations]; -+ -+ CUTLASS_DEVICE -+ ExpRowSumTransformA() { -+ auto thread_offset = ThreadMap::initial_offset(threadIdx.x); -+#pragma unroll -+ for (int s = 0; s < kStridedIterations; ++s) { -+ int row = blockIdx.x * PVThreadblockShape::kM -+ + thread_offset.strided() + s * ThreadMap::Delta::kStrided; -+ row_max[s] = row < g_rows ? g_row_max[row] : 0.0f; -+ row_inv_sum[s] = row < g_rows ? g_row_inv_sum[row] : 0.0f; -+ } -+ } -+ -+ CUTLASS_DEVICE -+ OutputFragment operator()(InputFragment const& input) { -+ OutputFragment output; -+ constexpr float kLog2E = 1.4426950408889634f; -+ constexpr int kElementsPerAccess = PVIteratorA::AccessType::kElements; -+ auto const* input_access = -+ reinterpret_cast(&input); -+ auto* output_access = -+ reinterpret_cast(&output); -+#pragma unroll -+ for (int s = 0; s < kStridedIterations; ++s) { -+#pragma unroll -+ for (int c = 0; c < kContiguousIterations; ++c) { -+#pragma unroll -+ for (int v = 0; v < kAccessesPerVector; ++v) { -+ int index = v + kAccessesPerVector * (c + s * kContiguousIterations); -+ typename PVIteratorA::AccessType transformed; -+#pragma unroll -+ for (int e = 0; e < kElementsPerAccess; ++e) { -+ float value = static_cast(input_access[index][e]); -+ transformed[e] = Element( -+ exp2f((value - row_max[s]) * kLog2E) * row_inv_sum[s]); -+ } -+ output_access[index] = transformed; -+ } -+ } -+ } -+ return output; -+ } -+}; -+#else -+struct ExpRowSumTransformA { -+ using InputFragment = typename PVIteratorA::Fragment; -+ using OutputFragment = cutlass::Array< -+ typename PVSmemIteratorA::Element, InputFragment::kElements>; -+ using ThreadMap = typename PVIteratorA::ThreadMap; -+ -+ static constexpr int kAccessesPerVector = -+ PVIteratorA::UnderlyingIterator::kAccessesPerVector; -+ static constexpr int kContiguousIterations = -+ ThreadMap::Iterations::kContiguous; -+ static constexpr int kStridedIterations = ThreadMap::Iterations::kStrided; -+ -+ float row_max[kStridedIterations]; -+ float row_sum[kStridedIterations]; -+ bool valid = true; -+ -+ CUTLASS_DEVICE -+ void set_valid(bool is_valid) { -+ valid = is_valid; -+ } -+ -+ CUTLASS_DEVICE -+ ExpRowSumTransformA() { -+ auto thread_offset = ThreadMap::initial_offset(threadIdx.x); -+#pragma unroll -+ for (int s = 0; s < kStridedIterations; ++s) { -+ int row = blockIdx.x * PVThreadblockShape::kM -+ + thread_offset.strided() + s * ThreadMap::Delta::kStrided; -+ row_max[s] = row < g_rows ? g_row_max[row] : 0.0f; -+ row_sum[s] = 0.0f; -+ } -+ } -+ -+ CUTLASS_DEVICE -+ OutputFragment operator()(InputFragment const& input) { -+ if (!valid) { -+ OutputFragment output; -+ output.clear(); -+ return output; -+ } -+ OutputFragment output; -+ constexpr float kLog2E = 1.4426950408889634f; -+ constexpr int kElementsPerAccess = PVIteratorA::AccessType::kElements; -+ auto const* input_access = -+ reinterpret_cast(&input); -+ auto* output_access = -+ reinterpret_cast(&output); -+#pragma unroll -+ for (int s = 0; s < kStridedIterations; ++s) { -+#pragma unroll -+ for (int c = 0; c < kContiguousIterations; ++c) { -+#pragma unroll -+ for (int v = 0; v < kAccessesPerVector; ++v) { -+ int index = v + kAccessesPerVector * (c + s * kContiguousIterations); -+ typename PVIteratorA::AccessType transformed; -+#pragma unroll -+ for (int e = 0; e < kElementsPerAccess; ++e) { -+ float value = static_cast(input_access[index][e]); -+ float weight = exp2f((value - row_max[s]) * kLog2E); -+ if (blockIdx.y == 0) { -+ row_sum[s] += weight; -+ } -+ transformed[e] = Element(weight); -+ } -+ output_access[index] = transformed; -+ } -+ } -+ } -+ return output; -+ } -+ -+ CUTLASS_DEVICE -+ void finalize() { -+ if (blockIdx.y != 0) { -+ return; -+ } -+ constexpr int kThreadsPerRow = -+ ThreadMap::Detail::WarpThreadArrangement::kContiguous; -+ auto thread_offset = ThreadMap::initial_offset(threadIdx.x); -+ int lane = threadIdx.x & 31; -+#pragma unroll -+ for (int s = 0; s < kStridedIterations; ++s) { -+ float sum = row_sum[s]; -+#pragma unroll -+ for (int offset = 1; offset < kThreadsPerRow; offset <<= 1) { -+ sum += __shfl_xor_sync(0xffffffffu, sum, offset); -+ } -+ int row = blockIdx.x * PVThreadblockShape::kM -+ + thread_offset.strided() + s * ThreadMap::Delta::kStrided; -+ if ((lane % kThreadsPerRow) == 0 && row < g_rows) { -+ g_row_sum_out[row] = sum; -+ } -+ } -+ } -+}; -+#endif -+ -+using PVTransformB = cutlass::NumericArrayConverter< -+ typename PVSmemIteratorB::Element, -+ typename PVIteratorB::Element, -+ PVIteratorB::Fragment::kElements>; -+using PVMma = cutlass::gemm::threadblock::MmaPipelined< -+ typename PVDefaultMma::Shape, -+ PVIteratorA, -+ PVSmemIteratorA, -+ PVIteratorB, -+ PVSmemIteratorB, -+ PVAccumulator, -+ PVLayout, -+ typename PVDefaultMma::Policy, -+ ExpRowSumTransformA, -+ PVTransformB>; -+using PVKernel = cutlass::gemm::kernel::Gemm< -+ PVMma, typename PVDefaultKernel::Epilogue, PVSwizzle, false>; -+ -+void check(cudaError_t result, char const* operation) { -+ if (result != cudaSuccess) { -+ std::cerr << operation << ": " << cudaGetErrorString(result) << "\n"; -+ std::exit(EXIT_FAILURE); -+ } -+} -+ -+void check(cutlass::Status status, char const* operation) { -+ if (status != cutlass::Status::kSuccess) { -+ std::cerr << operation << ": CUTLASS status " << int(status) << "\n"; -+ std::exit(EXIT_FAILURE); -+ } -+} -+ -+struct PVLauncher { -+ typename PVKernel::Params params; -+ dim3 grid; -+ dim3 block; -+ int smem_bytes; -+ -+ PVLauncher(Element* scores, Element* value, Element* output, int rows, int k) { -+ cutlass::gemm::GemmCoord problem(rows, 256, k); -+ PVSwizzle swizzle; -+ auto tiled_shape = swizzle.get_tiled_shape( -+ problem, -+ {PVThreadblockShape::kM, PVThreadblockShape::kN, -+ PVThreadblockShape::kK}, -+ 1); -+ params = typename PVKernel::Params( -+ problem, -+ tiled_shape, -+ {scores, PVLayout(k)}, -+ {value, PVLayout(256)}, -+ {output, PVLayout(256)}, -+ {output, PVLayout(256)}, -+ typename PVOutputOp::Params(1.0f, 0.0f), -+ nullptr); -+ grid = swizzle.get_grid_shape(tiled_shape); -+ block = dim3(PVKernel::kThreadCount, 1, 1); -+ smem_bytes = int(sizeof(typename PVKernel::SharedStorage)); -+ if (smem_bytes >= 48 * 1024) { -+ check(cudaFuncSetAttribute( -+ cutlass::Kernel, -+ cudaFuncAttributeMaxDynamicSharedMemorySize, -+ smem_bytes), -+ "set PV dynamic shared memory"); -+ } -+ } -+ -+ void launch(cudaStream_t stream) const { -+ cutlass::Kernel<<>>(params); -+ } -+}; -+ -+__global__ void prepare_weights( -+ float const* block_max, -+ float const* block_sum, -+ float* weights, -+ float* state_max, -+ float* state_sum, -+ int blocks, -+ int rows) { -+ int row = blockIdx.x * blockDim.x + threadIdx.x; -+ if (row >= rows) { -+ return; -+ } -+ float global_max = -CUDART_INF_F; -+ for (int block = 0; block < blocks; ++block) { -+ global_max = fmaxf(global_max, block_max[block * rows + row]); -+ } -+ float denominator = 0.0f; -+ for (int block = 0; block < blocks; ++block) { -+ float scale = exp2f( -+ (block_max[block * rows + row] - global_max) * -+ 1.4426950408889634f); -+#if defined(PREFIX_QK_FULL_STATS) -+ float scaled_sum = scale / block_sum[block * rows + row]; -+ weights[block * rows + row] = scaled_sum; -+ denominator += scaled_sum; -+#else -+ weights[block * rows + row] = scale; -+ denominator += scale * block_sum[block * rows + row]; -+#endif -+ } -+ float inverse = 1.0f / denominator; -+ for (int block = 0; block < blocks; ++block) { -+ weights[block * rows + row] *= inverse; -+ } -+ state_max[row] = global_max; -+ state_sum[row] = denominator; -+} -+ -+__global__ void merge_partials( -+ __half const* partials, -+ float const* weights, -+ __half* output, -+ int blocks, -+ int rows) { -+ int pair = blockIdx.x * blockDim.x + threadIdx.x; -+ constexpr int kPairsPerRow = 128; -+ int total_pairs = rows * kPairsPerRow; -+ if (pair >= total_pairs) { -+ return; -+ } -+ int row = pair / kPairsPerRow; -+ int column_pair = pair - row * kPairsPerRow; -+ float2 accumulator = make_float2(0.0f, 0.0f); -+ auto const* partials2 = reinterpret_cast<__half2 const*>(partials); -+ auto* output2 = reinterpret_cast<__half2*>(output); -+ for (int block = 0; block < blocks; ++block) { -+ int index = (block * rows + row) * kPairsPerRow + column_pair; -+ float2 value = __half22float2(partials2[index]); -+ float weight = weights[block * rows + row]; -+ accumulator.x = fmaf(weight, value.x, accumulator.x); -+ accumulator.y = fmaf(weight, value.y, accumulator.y); -+ } -+ output2[pair] = __floats2half2_rn(accumulator.x, accumulator.y); -+} -+ -+__global__ void update_prefix_accumulator( -+ __half const* block_output, -+ float const* block_max, -+ float const* block_inv_sum, -+ float* prefix_accumulator, -+ float* prefix_max, -+ float* prefix_sum, -+ int rows, -+ bool initialize) { -+ int row = blockIdx.x; -+ int d = threadIdx.x; -+ if (row >= rows || d >= 256) { -+ return; -+ } -+ __shared__ float scales[3]; -+ if (d == 0) { -+ constexpr float kLog2E = 1.4426950408889634f; -+ float next_max = block_max[row]; -+ float next_mass = 1.0f / block_inv_sum[row]; -+ if (initialize) { -+ scales[0] = 0.0f; -+ scales[1] = next_mass; -+ scales[2] = next_mass; -+ prefix_max[row] = next_max; -+ } else { -+ float old_max = prefix_max[row]; -+ float global_max = fmaxf(old_max, next_max); -+ float old_scale = exp2f((old_max - global_max) * kLog2E); -+ float next_scale = exp2f((next_max - global_max) * kLog2E); -+ scales[0] = old_scale; -+ scales[1] = next_mass * next_scale; -+ scales[2] = prefix_sum[row] * old_scale + scales[1]; -+ prefix_max[row] = global_max; -+ } -+ prefix_sum[row] = scales[2]; -+ } -+ __syncthreads(); -+ int64_t element = int64_t(row) * 256 + d; -+ float block_value = float(block_output[element]); -+ float old_value = initialize ? 0.0f : prefix_accumulator[element]; -+ prefix_accumulator[element] = -+ fmaf(old_value, scales[0], block_value * scales[1]); -+} -+ -+__global__ void prepare_prefix_update( -+ float const* block_max, -+ float const* block_inv_sum, -+ float* prefix_max, -+ float* prefix_sum, -+ float* old_scales, -+ float* block_scales, -+ int rows, -+ bool initialize) { -+ int row = blockIdx.x * blockDim.x + threadIdx.x; -+ if (row >= rows) { -+ return; -+ } -+ constexpr float kLog2E = 1.4426950408889634f; -+ float next_max = block_max[row]; -+ float next_mass = 1.0f / block_inv_sum[row]; -+ if (initialize) { -+ old_scales[row] = 0.0f; -+ block_scales[row] = next_mass; -+ prefix_max[row] = next_max; -+ prefix_sum[row] = next_mass; -+ return; -+ } -+ float old_max = prefix_max[row]; -+ float global_max = fmaxf(old_max, next_max); -+ float old_scale = exp2f((old_max - global_max) * kLog2E); -+ float block_scale = -+ next_mass * exp2f((next_max - global_max) * kLog2E); -+ old_scales[row] = old_scale; -+ block_scales[row] = block_scale; -+ prefix_max[row] = global_max; -+ prefix_sum[row] = prefix_sum[row] * old_scale + block_scale; -+} -+ -+__global__ void apply_prefix_update_half2( -+ __half const* block_output, -+ float* prefix_accumulator, -+ float const* old_scales, -+ float const* block_scales, -+ int rows, -+ bool initialize) { -+ int pair = blockIdx.x * blockDim.x + threadIdx.x; -+ constexpr int kPairsPerRow = 128; -+ if (pair >= rows * kPairsPerRow) { -+ return; -+ } -+ int row = pair / kPairsPerRow; -+ auto const* block2 = reinterpret_cast<__half2 const*>(block_output); -+ auto* accumulator2 = reinterpret_cast(prefix_accumulator); -+ float2 block_value = __half22float2(block2[pair]); -+ float2 old_value = initialize -+ ? make_float2(0.0f, 0.0f) -+ : accumulator2[pair]; -+ float old_scale = old_scales[row]; -+ float block_scale = block_scales[row]; -+ accumulator2[pair] = make_float2( -+ fmaf(old_value.x, old_scale, block_value.x * block_scale), -+ fmaf(old_value.y, old_scale, block_value.y * block_scale)); -+} -+ -+__global__ void finalize_prefix_accumulator( -+ float const* prefix_accumulator, -+ float const* prefix_sum, -+ __half* output, -+ int rows) { -+ int pair = blockIdx.x * blockDim.x + threadIdx.x; -+ constexpr int kPairsPerRow = 128; -+ if (pair >= rows * kPairsPerRow) { -+ return; -+ } -+ int row = pair / kPairsPerRow; -+ int64_t first = int64_t(pair) * 2; -+ float inverse = 1.0f / prefix_sum[row]; -+ output[first] = __float2half_rn(prefix_accumulator[first] * inverse); -+ output[first + 1] = -+ __float2half_rn(prefix_accumulator[first + 1] * inverse); -+} -+ -+__global__ void merge_prefix_accumulator_tail( -+ float const* prefix_accumulator, -+ float const* prefix_max, -+ float const* prefix_sum, -+ __half const* tail_output, -+ float const* tail_max, -+ float const* tail_sum, -+ __half* output, -+ int rows) { -+ int row = blockIdx.x; -+ int d = threadIdx.x; -+ __shared__ float masses[3]; -+ if (d == 0) { -+ constexpr float kLog2E = 1.4426950408889634f; -+ float global_max = fmaxf(prefix_max[row], tail_max[row]); -+ masses[0] = exp2f((prefix_max[row] - global_max) * kLog2E); -+ masses[1] = tail_sum[row] -+ * exp2f((tail_max[row] - global_max) * kLog2E); -+ masses[2] = 1.0f / (prefix_sum[row] * masses[0] + masses[1]); -+ } -+ __syncthreads(); -+ int64_t element = int64_t(row) * 256 + d; -+ float numerator = prefix_accumulator[element] * masses[0] -+ + float(tail_output[element]) * masses[1]; -+ output[element] = __float2half_rn(numerator * masses[2]); -+} -+ -+__global__ void merge_prefix_tail( -+ __half const* prefix_output, -+ float const* prefix_max, -+ float const* prefix_sum, -+ __half const* tail_output, -+ float const* tail_max, -+ float const* tail_sum, -+ __half* output, -+ int rows) { -+ int row = blockIdx.x; -+ int d = threadIdx.x; -+ __shared__ float masses[3]; -+ if (d == 0) { -+ float global_max = fmaxf(prefix_max[row], tail_max[row]); -+ float prefix_scale = exp2f( -+ (prefix_max[row] - global_max) * 1.4426950408889634f); -+ float tail_scale = exp2f( -+ (tail_max[row] - global_max) * 1.4426950408889634f); -+ masses[0] = prefix_sum[row] * prefix_scale; -+ masses[1] = tail_sum[row] * tail_scale; -+ masses[2] = 1.0f / (masses[0] + masses[1]); -+ } -+ __syncthreads(); -+ int64_t element = int64_t(row) * 256 + d; -+ float numerator = float(prefix_output[element]) * masses[0] -+ + float(tail_output[element]) * masses[1]; -+ output[element] = __float2half_rn(numerator * masses[2]); -+} -+ -+using DenseTailRaw = cudaError_t (*)( -+ const void*, -+ const void*, -+ const void*, -+ float*, -+ float*, -+ void*, -+ int, -+ int, -+ int, -+ int, -+ float, -+ cudaStream_t); -+ -+__global__ void fill_pattern(__half* data, size_t elements, uint32_t seed) { -+ size_t index = size_t(blockIdx.x) * blockDim.x + threadIdx.x; -+ if (index >= elements) { -+ return; -+ } -+ uint32_t value = uint32_t(index) ^ seed; -+ value ^= value >> 16; -+ value *= 0x7feb352du; -+ value ^= value >> 15; -+ value *= 0x846ca68bu; -+ value ^= value >> 16; -+ float uniform = float(value & 0xffffu) * (2.0f / 65535.0f) - 1.0f; -+ data[index] = __float2half_rn(uniform); -+} -+ -+struct BlockOperators { -+ int width; -+ QKGemm qk; -+ std::unique_ptr pv; -+}; -+ -+} // namespace -+ -+#if !defined(PREFIX_TORCH_EXTENSION) -+int main(int argc, char** argv) { -+ int block_n = argc > 1 ? std::atoi(argv[1]) : 2048; -+ int iterations = argc > 2 ? std::atoi(argv[2]) : 5; -+#if defined(PREFIX_FULL_ENDPOINT) -+ if (argc <= 3) { -+ std::cerr << "full endpoint requires the exact-tail extension path\n"; -+ return EXIT_FAILURE; -+ } -+#endif -+ constexpr int kRows = 48000; -+ constexpr int kHeadDim = 256; -+ constexpr int kPrefix = 120000; -+ constexpr int kTail = 8000; -+ constexpr int kTotalKV = kPrefix + kTail; -+#if defined(PREFIX_FULL_ENDPOINT) -+ constexpr int kStoredKV = kTotalKV; -+#else -+ constexpr int kStoredKV = kPrefix; -+#endif -+ int blocks = (kPrefix + block_n - 1) / block_n; -+ -+ Element* query = nullptr; -+ Element* key = nullptr; -+ Element* value = nullptr; -+ Element* scores = nullptr; -+ Element* partials = nullptr; -+ Element* output = nullptr; -+ float* qk_norm = nullptr; -+ float* qk_sum = nullptr; -+ float* prefix_accumulator = nullptr; -+ float* prefix_max = nullptr; -+ float* prefix_sum = nullptr; -+ float* old_scales = nullptr; -+ float* block_scales = nullptr; -+#if defined(PREFIX_FULL_ENDPOINT) -+ Element* tail_output = nullptr; -+ Element* final_output = nullptr; -+ Element* exact_output = nullptr; -+ float* tail_max = nullptr; -+ float* tail_sum = nullptr; -+#endif -+ -+ int qk_tiles_n = (block_n + QKThreadblockShape::kN - 1) / -+ QKThreadblockShape::kN; -+ size_t partial_elements = size_t(kRows) * kHeadDim; -+ check(cudaMalloc(&query, size_t(kRows) * kHeadDim * sizeof(Element)), "allocate Q"); -+ check(cudaMalloc(&key, size_t(kStoredKV) * kHeadDim * sizeof(Element)), "allocate K"); -+ check(cudaMalloc(&value, size_t(kStoredKV) * kHeadDim * sizeof(Element)), "allocate V"); -+ check(cudaMalloc(&scores, size_t(kRows) * block_n * sizeof(Element)), "allocate scores"); -+ check(cudaMalloc(&partials, partial_elements * sizeof(Element)), "allocate partials"); -+ check(cudaMalloc(&output, size_t(kRows) * kHeadDim * sizeof(Element)), "allocate output"); -+ check(cudaMalloc(&qk_norm, size_t(qk_tiles_n) * kRows * sizeof(float)), "allocate QK max"); -+ check(cudaMalloc(&qk_sum, size_t(qk_tiles_n) * kRows * sizeof(float)), "allocate QK sum"); -+ check(cudaMalloc( -+ &prefix_accumulator, -+ size_t(kRows) * kHeadDim * sizeof(float)), -+ "allocate prefix accumulator"); -+ check(cudaMalloc(&prefix_max, size_t(kRows) * sizeof(float)), "allocate prefix max"); -+ check(cudaMalloc(&prefix_sum, size_t(kRows) * sizeof(float)), "allocate prefix sum"); -+ check(cudaMalloc(&old_scales, size_t(kRows) * sizeof(float)), -+ "allocate old scales"); -+ check(cudaMalloc(&block_scales, size_t(kRows) * sizeof(float)), -+ "allocate block scales"); -+#if defined(PREFIX_FULL_ENDPOINT) -+ check(cudaMalloc(&tail_output, size_t(kRows) * kHeadDim * sizeof(Element)), -+ "allocate tail output"); -+ check(cudaMalloc(&final_output, size_t(kRows) * kHeadDim * sizeof(Element)), -+ "allocate final output"); -+ check(cudaMalloc(&exact_output, size_t(kRows) * kHeadDim * sizeof(Element)), -+ "allocate exact output"); -+ check(cudaMalloc(&tail_max, size_t(kRows) * sizeof(float)), "allocate tail max"); -+ check(cudaMalloc(&tail_sum, size_t(kRows) * sizeof(float)), "allocate tail sum"); -+#endif -+ size_t query_elements = size_t(kRows) * kHeadDim; -+ size_t kv_elements = size_t(kStoredKV) * kHeadDim; -+ fill_pattern<<<(query_elements + 255) / 256, 256>>>( -+ reinterpret_cast<__half*>(query), query_elements, 0x12345678u); -+ fill_pattern<<<(kv_elements + 255) / 256, 256>>>( -+ reinterpret_cast<__half*>(key), kv_elements, 0x9abcdef0u); -+ fill_pattern<<<(kv_elements + 255) / 256, 256>>>( -+ reinterpret_cast<__half*>(value), kv_elements, 0x31415926u); -+ check(cudaGetLastError(), "initialize deterministic inputs"); -+ -+#if defined(PREFIX_FULL_ENDPOINT) -+ void* tail_handle = dlopen(argv[3], RTLD_NOW | RTLD_LOCAL); -+ if (tail_handle == nullptr) { -+ std::cerr << "load exact-tail extension: " << dlerror() << "\n"; -+ return EXIT_FAILURE; -+ } -+ dlerror(); -+ auto tail_raw = reinterpret_cast( -+ dlsym(tail_handle, "onecat_sm70_d256_dense_state_raw")); -+ char const* symbol_error = dlerror(); -+ if (symbol_error != nullptr || tail_raw == nullptr) { -+ std::cerr << "load exact-tail entry point: " -+ << (symbol_error == nullptr ? "missing symbol" : symbol_error) -+ << "\n"; -+ return EXIT_FAILURE; -+ } -+#endif -+ -+ std::vector operators; -+ operators.reserve(blocks); -+ for (int block = 0; block < blocks; ++block) { -+ int begin = block * block_n; -+ int width = std::min(block_n, kPrefix - begin); -+ BlockOperators operation; -+ operation.width = width; -+ typename QKGemm::Arguments arguments( -+ {kRows, width, kHeadDim}, -+ 1, -+ {query, QKLayoutA(kHeadDim)}, -+ {key + size_t(begin) * kHeadDim, QKLayoutB(kHeadDim)}, -+ {scores, typename QKGemm::LayoutC(width)}, -+ {scores, typename QKGemm::LayoutC(width)}, -+ {Element(0.0625f), Element(0.0f)}, -+ {qk_norm, typename QKGemm::LayoutN(kRows)}, -+ {qk_sum, typename QKGemm::LayoutS(kRows)}, -+ {scores, typename QKGemm::LayoutSoft(width)}); -+ check(operation.qk.initialize(arguments), "initialize QK"); -+ operation.pv = std::make_unique( -+ scores, -+ value + size_t(begin) * kHeadDim, -+ partials, -+ kRows, -+ width); -+ operators.push_back(std::move(operation)); -+ } -+ -+ check(cudaMemcpyToSymbol(g_rows, &kRows, sizeof(kRows)), "set row count"); -+ float const* persistent_max_ptr = qk_norm; -+ check(cudaMemcpyToSymbol( -+ g_row_max, -+ &persistent_max_ptr, -+ sizeof(persistent_max_ptr)), -+ "set persistent PV max pointer"); -+#if defined(PREFIX_QK_FULL_STATS) -+ float const* persistent_inv_sum_ptr = qk_sum; -+ check(cudaMemcpyToSymbol( -+ g_row_inv_sum, -+ &persistent_inv_sum_ptr, -+ sizeof(persistent_inv_sum_ptr)), -+ "set persistent PV inverse-sum pointer"); -+#endif -+ cudaStream_t stream = nullptr; -+ auto launch = [&](cudaEvent_t blocks_done, -+ cudaEvent_t prefix_done, -+ cudaEvent_t tail_done) { -+ for (int block = 0; block < blocks; ++block) { -+ check(operators[block].qk(stream), "launch QK max"); -+#if !defined(PREFIX_QK_FULL_STATS) -+ float* sum_ptr = block_sum + size_t(block) * kRows; -+ check(cudaMemcpyToSymbolAsync( -+ g_row_sum_out, -+ &sum_ptr, -+ sizeof(sum_ptr), -+ 0, -+ cudaMemcpyHostToDevice, -+ stream), -+ "set PV sum pointer"); -+#endif -+ operators[block].pv->launch(stream); -+ prepare_prefix_update<<<(kRows + 255) / 256, 256, 0, stream>>>( -+ qk_norm, -+ qk_sum, -+ prefix_max, -+ prefix_sum, -+ old_scales, -+ block_scales, -+ kRows, -+ block == 0); -+ int pairs = kRows * kHeadDim / 2; -+ apply_prefix_update_half2<<<(pairs + 255) / 256, 256, 0, stream>>>( -+ reinterpret_cast<__half const*>(partials), -+ prefix_accumulator, -+ old_scales, -+ block_scales, -+ kRows, -+ block == 0); -+ } -+ if (blocks_done != nullptr) { -+ check(cudaEventRecord(blocks_done, stream), "record QK/PV blocks done"); -+ } -+#if !defined(PREFIX_FULL_ENDPOINT) -+ int pairs = kRows * kHeadDim / 2; -+ finalize_prefix_accumulator<<<(pairs + 255) / 256, 256, 0, stream>>>( -+ prefix_accumulator, -+ prefix_sum, -+ reinterpret_cast<__half*>(output), -+ kRows); -+#endif -+ if (prefix_done != nullptr) { -+ check(cudaEventRecord(prefix_done, stream), "record prefix merge done"); -+ } -+#if defined(PREFIX_FULL_ENDPOINT) -+ check(tail_raw( -+ query, -+ key + size_t(kPrefix) * kHeadDim, -+ value + size_t(kPrefix) * kHeadDim, -+ tail_max, -+ tail_sum, -+ tail_output, -+ kTail, -+ kTail, -+ 6, -+ 1, -+ 0.0625f, -+ stream), -+ "launch exact causal tail"); -+ if (tail_done != nullptr) { -+ check(cudaEventRecord(tail_done, stream), "record causal tail done"); -+ } -+ merge_prefix_accumulator_tail<<>>( -+ prefix_accumulator, -+ prefix_max, -+ prefix_sum, -+ reinterpret_cast<__half const*>(tail_output), -+ tail_max, -+ tail_sum, -+ reinterpret_cast<__half*>(final_output), -+ kRows); -+#endif -+ }; -+ -+ launch(nullptr, nullptr, nullptr); -+ check(cudaDeviceSynchronize(), "warmup synchronize"); -+ cudaEvent_t start; -+ cudaEvent_t stop; -+ check(cudaEventCreate(&start), "create start"); -+ check(cudaEventCreate(&stop), "create stop"); -+ check(cudaEventRecord(start), "record start"); -+ for (int iteration = 0; iteration < iterations; ++iteration) { -+ launch(nullptr, nullptr, nullptr); -+ } -+ check(cudaEventRecord(stop), "record stop"); -+ check(cudaEventSynchronize(stop), "synchronize stop"); -+ float elapsed_ms = 0.0f; -+ check(cudaEventElapsedTime(&elapsed_ms, start, stop), "elapsed time"); -+ float endpoint_ms = elapsed_ms / float(iterations); -+ constexpr double kFrozenUsefulTflop = 60.0 * 0.101581; -+#if defined(PREFIX_FULL_ENDPOINT) -+ cudaEvent_t phase_start; -+ cudaEvent_t blocks_done; -+ cudaEvent_t prefix_done; -+ cudaEvent_t tail_done; -+ cudaEvent_t phase_stop; -+ check(cudaEventCreate(&phase_start), "create phase start"); -+ check(cudaEventCreate(&blocks_done), "create blocks done"); -+ check(cudaEventCreate(&prefix_done), "create prefix done"); -+ check(cudaEventCreate(&tail_done), "create tail done"); -+ check(cudaEventCreate(&phase_stop), "create phase stop"); -+ check(cudaEventRecord(phase_start, stream), "record phase start"); -+ launch(blocks_done, prefix_done, tail_done); -+ check(cudaEventRecord(phase_stop, stream), "record phase stop"); -+ check(cudaEventSynchronize(phase_stop), "synchronize phase stop"); -+ float blocks_ms = 0.0f; -+ float prefix_merge_ms = 0.0f; -+ float tail_ms = 0.0f; -+ float final_merge_ms = 0.0f; -+ check(cudaEventElapsedTime(&blocks_ms, phase_start, blocks_done), -+ "measure QK/PV blocks"); -+ check(cudaEventElapsedTime(&prefix_merge_ms, blocks_done, prefix_done), -+ "measure prefix merge"); -+ check(cudaEventElapsedTime(&tail_ms, prefix_done, tail_done), -+ "measure causal tail"); -+ check(cudaEventElapsedTime(&final_merge_ms, tail_done, phase_stop), -+ "measure final merge"); -+ check(tail_raw( -+ query, -+ key, -+ value, -+ tail_max, -+ tail_sum, -+ exact_output, -+ kTail, -+ kTotalKV, -+ 6, -+ 1, -+ 0.0625f, -+ stream), -+ "launch monolithic exact reference"); -+ check(cudaDeviceSynchronize(), "synchronize exact reference"); -+ size_t output_elements = size_t(kRows) * kHeadDim; -+ std::vector<__half> candidate_host(output_elements); -+ std::vector<__half> reference_host(output_elements); -+ check(cudaMemcpy( -+ candidate_host.data(), -+ final_output, -+ output_elements * sizeof(__half), -+ cudaMemcpyDeviceToHost), -+ "copy candidate output"); -+ check(cudaMemcpy( -+ reference_host.data(), -+ exact_output, -+ output_elements * sizeof(__half), -+ cudaMemcpyDeviceToHost), -+ "copy exact output"); -+ double squared_error = 0.0; -+ double squared_reference = 0.0; -+ double absolute_error = 0.0; -+ float max_absolute_error = 0.0f; -+ size_t equal_elements = 0; -+ for (size_t index = 0; index < output_elements; ++index) { -+ float candidate_value = __half2float(candidate_host[index]); -+ float reference_value = __half2float(reference_host[index]); -+ float difference = candidate_value - reference_value; -+ float absolute = std::fabs(difference); -+ squared_error += double(difference) * difference; -+ squared_reference += double(reference_value) * reference_value; -+ absolute_error += absolute; -+ max_absolute_error = std::max(max_absolute_error, absolute); -+ equal_elements += candidate_value == reference_value; -+ } -+ double relative_l2 = std::sqrt(squared_error / squared_reference); -+ double mean_absolute_error = absolute_error / double(output_elements); -+ double measured_tflops = kFrozenUsefulTflop / (endpoint_ms * 1.0e-3); -+ std::cout << "{\n" -+ << " \"block_n\": " << block_n << ",\n" -+ << " \"blocks\": " << blocks << ",\n" -+ << " \"full_attention_ms\": " << endpoint_ms << ",\n" -+ << " \"useful_causal_tflops\": " << measured_tflops << ",\n" -+ << " \"prefix_tokens\": " << kPrefix << ",\n" -+ << " \"causal_tail_tokens\": " << kTail << ",\n" -+ << " \"phases_ms\": {\n" -+ << " \"prefix_qk_pv_blocks\": " << blocks_ms << ",\n" -+ << " \"prefix_state_and_partial_merge\": " -+ << prefix_merge_ms << ",\n" -+ << " \"exact_causal_tail\": " << tail_ms << ",\n" -+ << " \"prefix_tail_merge\": " << final_merge_ms << "\n" -+ << " },\n" -+ << " \"numerical\": {\n" -+ << " \"max_abs_diff\": " << max_absolute_error << ",\n" -+ << " \"mean_abs_diff\": " << mean_absolute_error << ",\n" -+ << " \"relative_l2\": " << relative_l2 << ",\n" -+ << " \"equal_elements\": " << equal_elements << ",\n" -+ << " \"elements\": " << output_elements << "\n" -+ << " },\n" -+ << " \"partial_storage_gib\": " -+ << double(partial_elements * sizeof(Element)) / double(1ull << 30) -+ << ",\n" -+ << " \"note\": \"measured serial prefix + exact tail state + final merge endpoint\"\n" -+ << "}\n"; -+ check(cudaEventDestroy(phase_start), "destroy phase start"); -+ check(cudaEventDestroy(blocks_done), "destroy blocks done"); -+ check(cudaEventDestroy(prefix_done), "destroy prefix done"); -+ check(cudaEventDestroy(tail_done), "destroy tail done"); -+ check(cudaEventDestroy(phase_stop), "destroy phase stop"); -+#else -+ double projected_full_ms = endpoint_ms + 4.83; -+ double projected_tflops = kFrozenUsefulTflop / (projected_full_ms * 1.0e-3); -+ std::cout << "{\n" -+ << " \"block_n\": " << block_n << ",\n" -+ << " \"blocks\": " << blocks << ",\n" -+ << " \"prefix_ms_including_merge\": " << endpoint_ms << ",\n" -+ << " \"accepted_tail_ms\": 4.83,\n" -+ << " \"projected_full_ms\": " << projected_full_ms << ",\n" -+ << " \"projected_useful_tflops\": " << projected_tflops << ",\n" -+ << " \"partial_storage_gib\": " -+ << double(partial_elements * sizeof(Element)) / double(1ull << 30) -+ << ",\n" -+ << " \"note\": \"prefix endpoint; prefix/tail state merge excluded\"\n" -+ << "}\n"; -+#endif -+ -+ check(cudaEventDestroy(start), "destroy start"); -+ check(cudaEventDestroy(stop), "destroy stop"); -+ check(cudaFree(query), "free Q"); -+ check(cudaFree(key), "free K"); -+ check(cudaFree(value), "free V"); -+ check(cudaFree(scores), "free scores"); -+ check(cudaFree(partials), "free partials"); -+ check(cudaFree(output), "free output"); -+ check(cudaFree(qk_norm), "free QK max"); -+ check(cudaFree(qk_sum), "free QK sum"); -+ check(cudaFree(prefix_accumulator), "free prefix accumulator"); -+ check(cudaFree(prefix_max), "free prefix max"); -+ check(cudaFree(prefix_sum), "free prefix sum"); -+ check(cudaFree(old_scales), "free old scales"); -+ check(cudaFree(block_scales), "free block scales"); -+#if defined(PREFIX_FULL_ENDPOINT) -+ check(cudaFree(tail_output), "free tail output"); -+ check(cudaFree(final_output), "free final output"); -+ check(cudaFree(exact_output), "free exact output"); -+ check(cudaFree(tail_max), "free tail max"); -+ check(cudaFree(tail_sum), "free tail sum"); -+ dlclose(tail_handle); -+#endif -+ return 0; -+} -+#else -+ -+extern "C" cudaError_t onecat_sm70_d256_dense_state_raw( -+ const void*, -+ const void*, -+ const void*, -+ float*, -+ float*, -+ void*, -+ int, -+ int, -+ int, -+ int, -+ float, -+ cudaStream_t); -+ -+namespace FLASH_NAMESPACE { -+ -+struct Sm70GqaScoreWorkspace { -+ at::Tensor scores; -+ cudaEvent_t completion = nullptr; -+ bool completion_recorded = false; -+ std::mutex launch_mutex; -+ -+ Sm70GqaScoreWorkspace(const at::Tensor& q, int rows, int block_n) -+ : scores(at::empty({rows, block_n}, q.options())) { -+ C10_CUDA_CHECK( -+ cudaEventCreateWithFlags(&completion, cudaEventDisableTiming)); -+ } -+ -+ ~Sm70GqaScoreWorkspace() { -+ if (completion != nullptr) { -+ cudaEventDestroy(completion); -+ } -+ } -+}; -+ -+using Sm70GqaScoreWorkspacePtr = std::shared_ptr; -+ -+std::mutex& sm70_gqa_score_cache_mutex() { -+ static std::mutex cache_mutex; -+ return cache_mutex; -+} -+ -+std::map& sm70_gqa_score_cache() { -+ static std::map cache; -+ return cache; -+} -+ -+Sm70GqaScoreWorkspacePtr get_sm70_gqa_score_workspace( -+ const at::Tensor& q, -+ int rows, -+ int block_n) { -+ std::lock_guard lock(sm70_gqa_score_cache_mutex()); -+ int device = q.get_device(); -+ auto& cache = sm70_gqa_score_cache(); -+ auto& workspace = cache[device]; -+ if (!workspace) { -+ workspace = -+ std::make_shared(q, rows, block_n); -+ } -+ return workspace; -+} -+ -+Sm70GqaScoreWorkspacePtr find_sm70_gqa_score_workspace(int device) { -+ std::lock_guard lock(sm70_gqa_score_cache_mutex()); -+ auto& cache = sm70_gqa_score_cache(); -+ auto found = cache.find(device); -+ return found == cache.end() ? nullptr : found->second; -+} -+ -+at::Tensor sm70_d256_gqa_architecture_fwd( -+ const at::Tensor& q, -+ const at::Tensor& k, -+ const at::Tensor& v, -+ at::Tensor& out, -+ double softmax_scale, -+ bool causal) { -+ constexpr int kQuery = 8000; -+ constexpr int kRows = 48000; -+ constexpr int kHeadDim = 256; -+ constexpr int kHeadsQ = 6; -+ constexpr int kHeadsKV = 1; -+ constexpr int kTail = 8000; -+ constexpr int kMinTotalKV = 40000; -+ constexpr int kMaxTotalKV = 128000; -+ constexpr int kTotalKVStep = 8000; -+ constexpr int kBlockN = 8192; -+ -+ TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda() && out.is_cuda(), -+ "SM70 GQA architecture requires CUDA tensors"); -+ TORCH_CHECK(q.scalar_type() == at::ScalarType::Half -+ && k.scalar_type() == q.scalar_type() -+ && v.scalar_type() == q.scalar_type() -+ && out.scalar_type() == q.scalar_type(), -+ "SM70 GQA architecture requires FP16 q, k, v, and out"); -+ TORCH_CHECK(q.sizes() == at::IntArrayRef({1, kQuery, kHeadsQ, kHeadDim}) -+ && k.dim() == 4 && k.size(0) == 1 -+ && k.size(2) == kHeadsKV && k.size(3) == kHeadDim -+ && v.sizes() == k.sizes() && out.sizes() == q.sizes(), -+ "SM70 GQA architecture only accepts the validated " -+ "Q8000/Hq6/Hkv1/D256 dense shape family"); -+ 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] " -+ "with an 8000-token step, got ", -+ total_kv); -+ const int prefix = total_kv - kTail; -+ const int blocks = (prefix + kBlockN - 1) / kBlockN; -+ TORCH_CHECK(q.is_contiguous() && k.is_contiguous() && v.is_contiguous() -+ && out.is_contiguous(), -+ "SM70 GQA architecture requires contiguous tensors"); -+ TORCH_CHECK(q.get_device() == k.get_device() -+ && q.get_device() == v.get_device() -+ && q.get_device() == out.get_device(), -+ "SM70 GQA architecture tensors must share one device"); -+ TORCH_CHECK(causal, "SM70 GQA architecture requires causal attention"); -+ TORCH_CHECK(std::abs(softmax_scale - 0.0625) < 1.0e-8, -+ "SM70 GQA architecture requires D256 softmax scale 1/16"); -+ -+ const at::cuda::OptionalCUDAGuard device_guard(q.device()); -+ const int qk_tiles_n = -+ (kBlockN + QKThreadblockShape::kN - 1) -+ / QKThreadblockShape::kN; -+ cudaStream_t stream = at::cuda::getCurrentCUDAStream(); -+ -+ constexpr size_t kScratchAlignment = 256; -+ size_t scratch_bytes = 0; -+ auto reserve_scratch = [&](size_t bytes) { -+ scratch_bytes = -+ (scratch_bytes + kScratchAlignment - 1) & ~(kScratchAlignment - 1); -+ size_t offset = scratch_bytes; -+ scratch_bytes += bytes; -+ return offset; -+ }; -+ size_t partial_offset = reserve_scratch( -+ size_t(kRows) * kHeadDim * sizeof(Element)); -+ size_t prefix_accumulator_offset = reserve_scratch( -+ size_t(kRows) * kHeadDim * sizeof(float)); -+ size_t tail_output_offset = reserve_scratch( -+ size_t(kRows) * kHeadDim * sizeof(Element)); -+ size_t qk_norm_offset = reserve_scratch( -+ size_t(qk_tiles_n) * kRows * sizeof(float)); -+ size_t qk_sum_offset = reserve_scratch( -+ size_t(qk_tiles_n) * kRows * sizeof(float)); -+ size_t prefix_max_offset = -+ reserve_scratch(size_t(kRows) * sizeof(float)); -+ size_t prefix_sum_offset = -+ reserve_scratch(size_t(kRows) * sizeof(float)); -+ size_t old_scale_offset = -+ reserve_scratch(size_t(kRows) * sizeof(float)); -+ size_t block_scale_offset = -+ reserve_scratch(size_t(kRows) * sizeof(float)); -+ size_t tail_max_offset = -+ reserve_scratch(size_t(kRows) * sizeof(float)); -+ size_t tail_sum_offset = -+ reserve_scratch(size_t(kRows) * sizeof(float)); -+ int device = q.get_device(); -+ auto score_workspace = find_sm70_gqa_score_workspace(device); -+ if (!score_workspace) { -+ constexpr size_t kRequiredPostWorkspaceHeadroom = 128 * 1024 * 1024; -+ constexpr size_t kScoreBytes = -+ size_t(kRows) * kBlockN * sizeof(Element); -+ size_t free_bytes = 0; -+ size_t total_bytes = 0; -+ C10_CUDA_CHECK(cudaMemGetInfo(&free_bytes, &total_bytes)); -+ size_t required_bytes = -+ kScoreBytes + scratch_bytes + kRequiredPostWorkspaceHeadroom; -+ TORCH_CHECK_WITH( -+ OutOfMemoryError, -+ free_bytes >= required_bytes, -+ "SM70 GQA architecture requires ", -+ required_bytes / (1024 * 1024), -+ " MiB free before its first workspace allocation, including " -+ "128 MiB of downstream headroom, but only ", -+ free_bytes / (1024 * 1024), -+ " MiB remains out of ", -+ total_bytes / (1024 * 1024), -+ " MiB"); -+ score_workspace = -+ get_sm70_gqa_score_workspace(q, kRows, kBlockN); -+ } -+ std::unique_lock launch_lock(score_workspace->launch_mutex); -+ if (score_workspace->completion_recorded) { -+ C10_CUDA_CHECK( -+ cudaStreamWaitEvent(stream, score_workspace->completion, 0)); -+ } -+ at::Tensor scratch = at::empty( -+ {static_cast(scratch_bytes)}, -+ q.options().dtype(at::ScalarType::Byte)); -+ auto* scratch_base = scratch.data_ptr(); -+ -+ auto* query = reinterpret_cast(q.data_ptr()); -+ auto* key = reinterpret_cast(k.data_ptr()); -+ auto* value = reinterpret_cast(v.data_ptr()); -+ auto* score_ptr = -+ reinterpret_cast(score_workspace->scores.data_ptr()); -+ auto* partial_ptr = -+ reinterpret_cast(scratch_base + partial_offset); -+ float* prefix_accumulator_ptr = -+ reinterpret_cast(scratch_base + prefix_accumulator_offset); -+ auto* tail_output_ptr = -+ reinterpret_cast(scratch_base + tail_output_offset); -+ auto* output_ptr = reinterpret_cast(out.data_ptr()); -+ float* qk_norm_ptr = -+ reinterpret_cast(scratch_base + qk_norm_offset); -+ float* qk_sum_ptr = -+ reinterpret_cast(scratch_base + qk_sum_offset); -+ float* prefix_max_ptr = -+ reinterpret_cast(scratch_base + prefix_max_offset); -+ float* prefix_sum_ptr = -+ reinterpret_cast(scratch_base + prefix_sum_offset); -+ float* old_scale_ptr = -+ reinterpret_cast(scratch_base + old_scale_offset); -+ float* block_scale_ptr = -+ reinterpret_cast(scratch_base + block_scale_offset); -+ float* tail_max_ptr = -+ reinterpret_cast(scratch_base + tail_max_offset); -+ float* tail_sum_ptr = -+ reinterpret_cast(scratch_base + tail_sum_offset); -+ -+ C10_CUDA_CHECK(onecat_sm70_d256_dense_state_raw( -+ query, -+ key + size_t(prefix) * kHeadDim, -+ value + size_t(prefix) * kHeadDim, -+ tail_max_ptr, -+ tail_sum_ptr, -+ tail_output_ptr, -+ kTail, -+ kTail, -+ kHeadsQ, -+ kHeadsKV, -+ static_cast(softmax_scale), -+ stream)); -+ C10_CUDA_CHECK(cudaMemcpyToSymbolAsync( -+ g_rows, -+ &kRows, -+ sizeof(kRows), -+ 0, -+ cudaMemcpyHostToDevice, -+ stream)); -+ float const* persistent_max_ptr = qk_norm_ptr; -+ C10_CUDA_CHECK(cudaMemcpyToSymbolAsync( -+ g_row_max, -+ &persistent_max_ptr, -+ sizeof(persistent_max_ptr), -+ 0, -+ cudaMemcpyHostToDevice, -+ stream)); -+ float const* persistent_inv_sum_ptr = qk_sum_ptr; -+ C10_CUDA_CHECK(cudaMemcpyToSymbolAsync( -+ g_row_inv_sum, -+ &persistent_inv_sum_ptr, -+ sizeof(persistent_inv_sum_ptr), -+ 0, -+ cudaMemcpyHostToDevice, -+ stream)); -+ -+ for (int block = 0; block < blocks; ++block) { -+ int begin = block * kBlockN; -+ int width = std::min(kBlockN, prefix - begin); -+ BlockOperators operation; -+ operation.width = width; -+ typename QKGemm::Arguments arguments( -+ {kRows, width, kHeadDim}, -+ 1, -+ {query, QKLayoutA(kHeadDim)}, -+ {key + size_t(begin) * kHeadDim, QKLayoutB(kHeadDim)}, -+ {score_ptr, typename QKGemm::LayoutC(width)}, -+ {score_ptr, typename QKGemm::LayoutC(width)}, -+ {Element(static_cast(softmax_scale)), Element(0.0f)}, -+ {qk_norm_ptr, typename QKGemm::LayoutN(kRows)}, -+ {qk_sum_ptr, typename QKGemm::LayoutS(kRows)}, -+ {score_ptr, typename QKGemm::LayoutSoft(width)}); -+ TORCH_CHECK(operation.qk.initialize(arguments) == cutlass::Status::kSuccess, -+ "initialize SM70 GQA QK block ", block, " failed"); -+ operation.pv = std::make_unique( -+ score_ptr, -+ value + size_t(begin) * kHeadDim, -+ partial_ptr, -+ kRows, -+ width); -+ TORCH_CHECK(operation.qk(stream) == cutlass::Status::kSuccess, -+ "launch SM70 GQA QK block ", block, " failed"); -+ operation.pv->launch(stream); -+ prepare_prefix_update<<<(kRows + 255) / 256, 256, 0, stream>>>( -+ qk_norm_ptr, -+ qk_sum_ptr, -+ prefix_max_ptr, -+ prefix_sum_ptr, -+ old_scale_ptr, -+ block_scale_ptr, -+ kRows, -+ block == 0); -+ constexpr int kPairs = kRows * kHeadDim / 2; -+ apply_prefix_update_half2<<<(kPairs + 255) / 256, 256, 0, stream>>>( -+ reinterpret_cast<__half const*>(partial_ptr), -+ prefix_accumulator_ptr, -+ old_scale_ptr, -+ block_scale_ptr, -+ kRows, -+ block == 0); -+ } -+ merge_prefix_accumulator_tail<<>>( -+ prefix_accumulator_ptr, -+ prefix_max_ptr, -+ prefix_sum_ptr, -+ reinterpret_cast<__half const*>(tail_output_ptr), -+ tail_max_ptr, -+ tail_sum_ptr, -+ reinterpret_cast<__half*>(output_ptr), -+ kRows); -+ C10_CUDA_KERNEL_LAUNCH_CHECK(); -+ C10_CUDA_CHECK(cudaEventRecord(score_workspace->completion, stream)); -+ score_workspace->completion_recorded = true; -+ return out; -+} -+ -+} // namespace FLASH_NAMESPACE -+ -+#endif -diff --git a/csrc/flash_attn/src/flash_fwd_d256_splitd_sm70.cu b/csrc/flash_attn/src/flash_fwd_d256_splitd_sm70.cu -index 7a0f626..1fd4008 100644 ---- a/csrc/flash_attn/src/flash_fwd_d256_splitd_sm70.cu -+++ b/csrc/flash_attn/src/flash_fwd_d256_splitd_sm70.cu -@@ -398,7 +398,11 @@ __device__ __forceinline__ int64_t paged_kv_thread_offset( - + (tid % kThreadsPerRow) * kElemsPerLoad; - } - --template -+template < -+ typename Element, -+ bool PagedKV, -+ bool SplitKV3 = false, -+ bool StoreState = false> - __global__ __launch_bounds__(Sm70D256SplitDTraits::kNThreads, 1) - void sm70_d256_splitd_dense_kernel( - const Element *__restrict__ q, -@@ -859,6 +863,29 @@ void sm70_d256_splitd_dense_kernel( - } - } - } else { -+ if constexpr (StoreState) { -+ if ((lane & 0x0e) == 0) { -+#pragma unroll -+ for (int slot = 0; slot < Traits::kQkRowsPerThread; ++slot) { -+ const int row = FLASH_NAMESPACE::sm70_row_slot< -+ Traits::kQkWarpRows>(slot, lane); -+ const int query_row = query_row_base -+ + group_row_base -+ + n_warp * Traits::kQkWarpRows + row; -+ const int64_t state_row = -+ (static_cast(batch) * query_len + query_row) -+ * heads_q -+ + head_q; -+ // The standalone prefix path stores maxima after applying -+ // softmax_scale. Export the dense-tail state in the same -+ // natural-log coordinate so the two states can be merged -+ // directly with exp(prefix_max - global_max). -+ partial_max[state_row] = row_max[slot] -+ * (softmax_scale_log2 * float(M_LN2)); -+ partial_sum[state_row] = row_sum[slot]; -+ } -+ } -+ } - if ((lane & 0x0e) == 0) { - #pragma unroll - for (int slot = 0; slot < Traits::kQkRowsPerThread; ++slot) { -@@ -1042,6 +1069,154 @@ at::Tensor sm70_d256_splitd_dense_fwd( - return out; - } - -+at::Tensor sm70_d256_splitd_dense_state_fwd( -+ const at::Tensor &q, -+ const at::Tensor &k, -+ const at::Tensor &v, -+ at::Tensor &state_max, -+ at::Tensor &state_sum, -+ at::Tensor &out, -+ double softmax_scale, -+ bool causal) { -+ TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda() -+ && state_max.is_cuda() && state_sum.is_cuda() -+ && out.is_cuda(), -+ "dense-state inputs must be CUDA tensors"); -+ TORCH_CHECK(q.scalar_type() == at::ScalarType::Half -+ && k.scalar_type() == q.scalar_type() -+ && v.scalar_type() == q.scalar_type() -+ && out.scalar_type() == q.scalar_type(), -+ "dense-state requires FP16 q, k, v, and out"); -+ TORCH_CHECK(state_max.scalar_type() == at::ScalarType::Float -+ && state_sum.scalar_type() == at::ScalarType::Float, -+ "dense-state statistics must use FP32"); -+ TORCH_CHECK(q.dim() == 4 && k.dim() == 4 && v.dim() == 4 -+ && q.size(0) == k.size(0) && k.sizes() == v.sizes(), -+ "dense-state q, k, and v shapes are invalid"); -+ TORCH_CHECK(q.size(3) == Sm70D256SplitDTraits::kHeadDim -+ && k.size(3) == Sm70D256SplitDTraits::kHeadDim -+ && q.size(2) % k.size(2) == 0, -+ "dense-state requires D256 and valid GQA heads"); -+ TORCH_CHECK(q.size(1) <= k.size(1) -+ && q.size(1) % Sm70D256SplitDTraits::kBlockM == 0 -+ && k.size(1) % Sm70D256SplitDTraits::kBlockN == 0, -+ "dense-state lengths do not match the exact tiles"); -+ TORCH_CHECK(out.sizes() == q.sizes() && out.is_contiguous() -+ && state_max.is_contiguous() && state_sum.is_contiguous(), -+ "dense-state outputs must be contiguous"); -+ const int64_t rows = q.size(0) * q.size(1) * q.size(2); -+ TORCH_CHECK(state_max.numel() == rows && state_sum.numel() == rows, -+ "dense-state statistic shapes are invalid"); -+ TORCH_CHECK(causal, "dense-state requires causal attention"); -+ -+ const at::cuda::OptionalCUDAGuard device_guard(q.device()); -+ const dim3 block(Sm70D256SplitDTraits::kNThreads); -+ const dim3 grid( -+ q.size(1) / Sm70D256SplitDTraits::kBlockM, -+ q.size(0), -+ q.size(2)); -+ auto stream = at::cuda::getCurrentCUDAStream(); -+ auto kernel = sm70_d256_splitd_dense_kernel< -+ cutlass::half_t, false, false, true>; -+ C10_CUDA_CHECK(cudaFuncSetAttribute( -+ kernel, -+ cudaFuncAttributeMaxDynamicSharedMemorySize, -+ Sm70D256SplitDTraits::kSmemBytes)); -+ kernel<<>>( -+ reinterpret_cast(q.data_ptr()), -+ reinterpret_cast(k.data_ptr()), -+ reinterpret_cast(v.data_ptr()), -+ reinterpret_cast(out.data_ptr()), -+ static_cast(q.stride(0)), -+ static_cast(q.stride(1)), -+ static_cast(q.stride(2)), -+ static_cast(k.stride(0)), -+ static_cast(k.stride(1)), -+ static_cast(k.stride(2)), -+ static_cast(v.stride(0)), -+ static_cast(v.stride(1)), -+ static_cast(v.stride(2)), -+ q.size(1), -+ k.size(1), -+ q.size(2), -+ k.size(2), -+ static_cast(softmax_scale * M_LOG2E), -+ nullptr, -+ 0, -+ 0, -+ nullptr, -+ state_max.data_ptr(), -+ state_sum.data_ptr()); -+ C10_CUDA_KERNEL_LAUNCH_CHECK(); -+ return out; -+} -+ -+// Raw fixed-layout entry point for the task-local end-to-end architecture -+// harness. Keeping this in the same cubin guarantees that the measured tail -+// is the exact production candidate kernel rather than a reimplemented proxy. -+extern "C" cudaError_t onecat_sm70_d256_dense_state_raw( -+ const void *q, -+ const void *k, -+ const void *v, -+ float *state_max, -+ float *state_sum, -+ void *out, -+ int query_len, -+ int kv_len, -+ int heads_q, -+ int heads_kv, -+ float softmax_scale, -+ cudaStream_t stream) { -+ if (q == nullptr || k == nullptr || v == nullptr || state_max == nullptr -+ || state_sum == nullptr || out == nullptr || query_len <= 0 -+ || kv_len < query_len || heads_q <= 0 || heads_kv <= 0 -+ || heads_q % heads_kv != 0 -+ || query_len % Sm70D256SplitDTraits::kBlockM != 0 -+ || kv_len % Sm70D256SplitDTraits::kBlockN != 0) { -+ return cudaErrorInvalidValue; -+ } -+ const dim3 block(Sm70D256SplitDTraits::kNThreads); -+ const dim3 grid( -+ query_len / Sm70D256SplitDTraits::kBlockM, -+ 1, -+ heads_q); -+ auto kernel = sm70_d256_splitd_dense_kernel< -+ cutlass::half_t, false, false, true>; -+ cudaError_t result = cudaFuncSetAttribute( -+ kernel, -+ cudaFuncAttributeMaxDynamicSharedMemorySize, -+ Sm70D256SplitDTraits::kSmemBytes); -+ if (result != cudaSuccess) { -+ return result; -+ } -+ kernel<<>>( -+ static_cast(q), -+ static_cast(k), -+ static_cast(v), -+ static_cast(out), -+ query_len * heads_q * Sm70D256SplitDTraits::kHeadDim, -+ heads_q * Sm70D256SplitDTraits::kHeadDim, -+ Sm70D256SplitDTraits::kHeadDim, -+ kv_len * heads_kv * Sm70D256SplitDTraits::kHeadDim, -+ heads_kv * Sm70D256SplitDTraits::kHeadDim, -+ Sm70D256SplitDTraits::kHeadDim, -+ kv_len * heads_kv * Sm70D256SplitDTraits::kHeadDim, -+ heads_kv * Sm70D256SplitDTraits::kHeadDim, -+ Sm70D256SplitDTraits::kHeadDim, -+ query_len, -+ kv_len, -+ heads_q, -+ heads_kv, -+ softmax_scale * float(M_LOG2E), -+ nullptr, -+ 0, -+ 0, -+ nullptr, -+ state_max, -+ state_sum); -+ return cudaPeekAtLastError(); -+} -+ - at::Tensor sm70_d256_splitd_dense_splitkv3_fwd( - const at::Tensor &q, - const at::Tensor &k, -diff --git a/csrc/cutlass/examples/35_gemm_softmax/gemm_with_softmax.h b/csrc/cutlass/examples/35_gemm_softmax/gemm_with_softmax.h -index 31b2b769..2cb4ba74 100644 ---- a/csrc/cutlass/examples/35_gemm_softmax/gemm_with_softmax.h -+++ b/csrc/cutlass/examples/35_gemm_softmax/gemm_with_softmax.h -@@ -623,6 +623,7 @@ public: - return cutlass::Status::kErrorInternal; - } - -+#if !defined(PREFIX_QK_SKIP_APPLY) - // - // Launch the SoftmaxApplyKernel - // -@@ -646,6 +647,7 @@ public: - if (result != cudaSuccess) { - return cutlass::Status::kErrorInternal; - } -+#endif - - return cutlass::Status::kSuccess; - } diff --git a/cmake/patches/sm70_flash_attn_d256_k_pingpong.patch b/cmake/patches/sm70_flash_attn_d256_k_pingpong.patch deleted file mode 100644 index a215fcd9de..0000000000 --- a/cmake/patches/sm70_flash_attn_d256_k_pingpong.patch +++ /dev/null @@ -1,90 +0,0 @@ ---- a/csrc/flash_attn/src/flash_fwd_d256_splitd_sm70.cu -+++ b/csrc/flash_attn/src/flash_fwd_d256_splitd_sm70.cu -@@ -126,8 +126,16 @@ - Stride<_4, _1>>{})); - - static constexpr int kQElements = size(SmemLayoutQ{}); - static constexpr int kKVElements = size(SmemLayoutKV{}); - static_assert(size(SmemLayoutV{}) == kKVElements); -+ // The second K stage is live only before P is materialized, so it may -+ // alias the beginning of the later P region. Keep the K stages disjoint -+ // and on the same 128-byte bank phase without moving the V/P regions. -+ static constexpr int kKStageElements = 2240; -+ static_assert(kKStageElements >= cosize(SmemLayoutK{})); -+ static_assert( -+ kKStageElements + cosize(SmemLayoutK{}) -+ <= 2 * kKVElements + size(SmemLayoutP{})); - static constexpr int kPElements = size(SmemLayoutP{}); - static constexpr int kExchangeRows = kMmaGroups * kGroupRows; - static constexpr int kTensorSmemBytes = -@@ -521,11 +525,11 @@ - make_gmem_ptr(k + k_batch_offset + head_kv * k_head_stride), - make_shape(kv_len, Int{}), - make_stride(k_row_stride, _1{})); -- auto sK = make_tensor( -+ auto sKFirst = make_tensor( - make_smem_ptr(kv_smem_ptr), typename Traits::SmemLayoutK{}); -- auto tKsKRaw = gmem_k_thread.partition_D(sK); -- auto tKsK = reshape_kv_thread_tensor(tKsKRaw); -- auto tKrKNext = make_fragment_like(tKsK); -+ auto tKsKFirstRaw = gmem_k_thread.partition_D(sKFirst); -+ auto tKsKFirst = reshape_kv_thread_tensor(tKsKFirstRaw); -+ auto tKrKNext = make_fragment_like(tKsKFirst); - auto gKFirst = local_tile( - mK, - Shape, Int>{}, -@@ -543,7 +547,7 @@ - k_row_stride); - tKgKFirst.data() = mK.data() + k_thread_tile_base; - } -- copy_even_tile(gmem_k_copy, tKgKFirst, tKsK); -+ copy_even_tile(gmem_k_copy, tKgKFirst, tKsKFirst); - __syncthreads(); - - for (int n_block = n_block_max; n_block >= n_block_min; --n_block) { -@@ -558,6 +562,11 @@ - - #pragma unroll - for (int d_chunk = 0; d_chunk < Traits::kDChunks; ++d_chunk) { -+ auto sK = make_tensor( -+ make_smem_ptr( -+ kv_smem_ptr -+ + (d_chunk & 1) * Traits::kKStageElements), -+ typename Traits::SmemLayoutK{}); - if (d_chunk + 1 < Traits::kDChunks) { - auto gKNext = local_tile( - mK, -@@ -596,8 +605,15 @@ - acc_s, tSrQ, tSrK, tSsQ, tSsK, qk_tiled_mma, - smem_copy_q, smem_copy_k, smem_thread_q, smem_thread_k); - if (d_chunk + 1 < Traits::kDChunks) { -- __syncthreads(); -- cute::copy(tKrKNext, tKsK); -+ auto sKNext = make_tensor( -+ make_smem_ptr( -+ kv_smem_ptr -+ + ((d_chunk + 1) & 1) * Traits::kKStageElements), -+ typename Traits::SmemLayoutK{}); -+ auto tKsKNextRaw = gmem_k_thread.partition_D(sKNext); -+ auto tKsKNext = -+ reshape_kv_thread_tensor(tKsKNextRaw); -+ cute::copy(tKrKNext, tKsKNext); - __syncthreads(); - } - } -@@ -770,7 +786,14 @@ - } - } - if (n_block > n_block_min) { -- cute::copy(tKrKNext, tKsK); -+ auto sKNextBlock = make_tensor( -+ make_smem_ptr(kv_smem_ptr), -+ typename Traits::SmemLayoutK{}); -+ auto tKsKNextBlockRaw = -+ gmem_k_thread.partition_D(sKNextBlock); -+ auto tKsKNextBlock = -+ reshape_kv_thread_tensor(tKsKNextBlockRaw); -+ cute::copy(tKrKNext, tKsKNextBlock); - __syncthreads(); - } - } diff --git a/csrc/custom_all_reduce.cu b/csrc/custom_all_reduce.cu index e55f515ebb..c294e2acd4 100644 --- a/csrc/custom_all_reduce.cu +++ b/csrc/custom_all_reduce.cu @@ -161,12 +161,12 @@ void sm70_all_reduce_gemma_rms_norm_impl( "SM70 TP4 Gemma RMSNorm prototype residual must be float32."); } else { TORCH_CHECK(residual.scalar_type() == at::ScalarType::Half || - residual.scalar_type() == at::ScalarType::Float, + residual.scalar_type() == at::ScalarType::Float, "SM70 Gemma RMSNorm prototype residual must be float16 or " "float32."); } TORCH_CHECK(weight.scalar_type() == at::ScalarType::Half || - weight.scalar_type() == at::ScalarType::Float, + weight.scalar_type() == at::ScalarType::Float, "SM70 Gemma RMSNorm prototype weight must be float16 or " "float32."); TORCH_CHECK_EQ(inp.dim(), 2); @@ -176,9 +176,8 @@ void sm70_all_reduce_gemma_rms_norm_impl( TORCH_CHECK(normalized_out.sizes() == inp.sizes(), "SM70 Gemma RMSNorm prototype normalized_out shape must match " "inp."); - TORCH_CHECK( - residual_out.sizes() == inp.sizes(), - "SM70 Gemma RMSNorm prototype residual_out shape must match inp."); + TORCH_CHECK(residual_out.sizes() == inp.sizes(), + "SM70 Gemma RMSNorm prototype residual_out shape must match inp."); TORCH_CHECK_EQ(weight.dim(), 1); TORCH_CHECK_EQ(weight.numel(), kHiddenSize); TORCH_CHECK_EQ(residual.get_device(), inp.get_device()); @@ -278,83 +277,6 @@ void sm70_tp4_all_reduce_gemma_rms_norm( reg_buffer_sz_bytes, epsilon); } -void sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather( - fptr_t _fa, torch::Tensor& inp, torch::Tensor& residual, - torch::Tensor& weight, torch::Tensor& normalized_out, - torch::Tensor& residual_out, fptr_t _reg_input_buffer, - fptr_t _reg_output_buffer, int64_t reg_buffer_sz_bytes, double epsilon) { - constexpr int64_t kWorldSize = 4; - constexpr int64_t kHiddenSize = vllm::kSm70GemmaRmsNormHiddenSize; - auto fa = reinterpret_cast(_fa); - const at::cuda::OptionalCUDAGuard device_guard(device_of(inp)); - auto stream = c10::cuda::getCurrentCUDAStream().stream(); - - TORCH_CHECK_EQ(fa->world_size_, kWorldSize); - TORCH_CHECK(fa->fully_connected_); - TORCH_CHECK_EQ(inp.scalar_type(), at::ScalarType::Half); - TORCH_CHECK_EQ(residual.scalar_type(), at::ScalarType::Float); - TORCH_CHECK(weight.scalar_type() == at::ScalarType::Half || - weight.scalar_type() == at::ScalarType::Float); - TORCH_CHECK_EQ(normalized_out.scalar_type(), at::ScalarType::Half); - TORCH_CHECK_EQ(residual_out.scalar_type(), at::ScalarType::Float); - TORCH_CHECK_EQ(inp.dim(), 2); - TORCH_CHECK_EQ(inp.size(1), kHiddenSize); - TORCH_CHECK_EQ(inp.size(0) % kWorldSize, 0); - TORCH_CHECK_LE(inp.size(0) / kWorldSize, - vllm::kSm70LongPrefillMaxTokensPerRank); - TORCH_CHECK(residual.sizes() == inp.sizes()); - TORCH_CHECK(normalized_out.sizes() == inp.sizes()); - TORCH_CHECK(residual_out.sizes() == inp.sizes()); - TORCH_CHECK_EQ(weight.dim(), 1); - TORCH_CHECK_EQ(weight.numel(), kHiddenSize); - TORCH_CHECK_EQ(residual.get_device(), inp.get_device()); - TORCH_CHECK_EQ(weight.get_device(), inp.get_device()); - TORCH_CHECK_EQ(normalized_out.get_device(), inp.get_device()); - TORCH_CHECK_EQ(residual_out.get_device(), inp.get_device()); - TORCH_CHECK(inp.is_contiguous()); - TORCH_CHECK(residual.is_contiguous()); - TORCH_CHECK(weight.is_contiguous()); - TORCH_CHECK(normalized_out.is_contiguous()); - TORCH_CHECK(residual_out.is_contiguous()); - - auto* reg_input_buffer = reinterpret_cast(_reg_input_buffer); - auto* reg_output_buffer = reinterpret_cast(_reg_output_buffer); - const int64_t input_size = inp.numel() * inp.element_size(); - TORCH_CHECK_LE(input_size, reg_buffer_sz_bytes); - if (reg_input_buffer != nullptr) { - AT_CUDA_CHECK(cudaMemcpyAsync(reg_input_buffer, inp.data_ptr(), input_size, - cudaMemcpyDeviceToDevice, stream)); - } else { - reg_input_buffer = inp.data_ptr(); - } - if (reg_output_buffer == nullptr) { - reg_output_buffer = normalized_out.data_ptr(); - } - - auto* input_ptr = reinterpret_cast(reg_input_buffer); - auto* shared_output_ptr = reinterpret_cast(reg_output_buffer); - auto* residual_ptr = reinterpret_cast(residual.data_ptr()); - auto* residual_out_ptr = reinterpret_cast(residual_out.data_ptr()); - const int num_tokens = static_cast(inp.size(0)); - const float epsilon_f = static_cast(epsilon); - if (weight.scalar_type() == at::ScalarType::Float) { - fa->sm70_reduce_scatter_gemma_rms_norm_all_gather<4, float, float>( - stream, input_ptr, shared_output_ptr, residual_ptr, - reinterpret_cast(weight.data_ptr()), residual_out_ptr, - num_tokens, kHiddenSize, epsilon_f); - } else { - fa->sm70_reduce_scatter_gemma_rms_norm_all_gather<4, float, half>( - stream, input_ptr, shared_output_ptr, residual_ptr, - reinterpret_cast(weight.data_ptr()), residual_out_ptr, - num_tokens, kHiddenSize, epsilon_f); - } - if (_reg_output_buffer != 0) { - AT_CUDA_CHECK(cudaMemcpyAsync(normalized_out.data_ptr(), reg_output_buffer, - input_size, cudaMemcpyDeviceToDevice, - stream)); - } -} - void all_reduce_sum2(fptr_t _fa, torch::Tensor& inp_a, torch::Tensor& inp_b, torch::Tensor& out) { auto fa = reinterpret_cast(_fa); @@ -385,17 +307,17 @@ void all_reduce_sum2(fptr_t _fa, torch::Tensor& inp_a, torch::Tensor& inp_b, switch (out.scalar_type()) { case at::ScalarType::Float: { - fa->allreduce_sum2( - stream, reinterpret_cast(inp_a.data_ptr()), - reinterpret_cast(inp_b.data_ptr()), - reinterpret_cast(out.data_ptr()), out.numel()); + fa->allreduce_sum2(stream, reinterpret_cast(inp_a.data_ptr()), + reinterpret_cast(inp_b.data_ptr()), + reinterpret_cast(out.data_ptr()), + out.numel()); break; } case at::ScalarType::Half: { - fa->allreduce_sum2( - stream, reinterpret_cast(inp_a.data_ptr()), - reinterpret_cast(inp_b.data_ptr()), - reinterpret_cast(out.data_ptr()), out.numel()); + fa->allreduce_sum2(stream, reinterpret_cast(inp_a.data_ptr()), + reinterpret_cast(inp_b.data_ptr()), + reinterpret_cast(out.data_ptr()), + out.numel()); break; } #if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) @@ -441,7 +363,8 @@ void top1_argmax(fptr_t _fa, torch::Tensor& input_pair, torch::Tensor& output, } void tile_runtime_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out, - fptr_t _reg_buffer, int64_t reg_buffer_sz_bytes, + fptr_t _reg_buffer, + int64_t reg_buffer_sz_bytes, int64_t tile_numel, int64_t engine_blocks, int64_t compute_iters) { auto fa = reinterpret_cast(_fa); @@ -489,7 +412,8 @@ void tile_runtime_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out, void tile_runtime_all_reduce_engine(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out, fptr_t _reg_buffer, int64_t reg_buffer_sz_bytes, - int64_t tile_numel, int64_t producer_blocks, + int64_t tile_numel, + int64_t producer_blocks, int64_t reducer_blocks, int64_t compute_iters) { auto fa = reinterpret_cast(_fa); @@ -575,10 +499,6 @@ void dispose(fptr_t _fa) { int64_t meta_size() { return sizeof(vllm::Signal); } -int64_t sm70_tp4_push_allreduce_buffer_size() { - return vllm::kSm70Tp4PushAllreduceBufferBytes; -} - void register_buffer(fptr_t _fa, const std::vector& fake_ipc_ptrs) { auto fa = reinterpret_cast(_fa); TORCH_CHECK(fake_ipc_ptrs.size() == fa->world_size_); @@ -589,19 +509,6 @@ void register_buffer(fptr_t _fa, const std::vector& fake_ipc_ptrs) { fa->register_buffer(ipc_ptrs); } -void register_sm70_tp4_push_allreduce_buffer( - fptr_t _fa, const std::vector& fake_ipc_ptrs) { - auto fa = reinterpret_cast(_fa); - TORCH_CHECK_EQ(fake_ipc_ptrs.size(), - static_cast(vllm::kSm70Tp4PushAllreduceWorldSize)); - TORCH_CHECK_EQ(fake_ipc_ptrs.size(), static_cast(fa->world_size_)); - void* ipc_ptrs[vllm::kSm70Tp4PushAllreduceWorldSize]; - for (size_t peer = 0; peer < fake_ipc_ptrs.size(); ++peer) { - ipc_ptrs[peer] = reinterpret_cast(fake_ipc_ptrs[peer]); - } - fa->register_sm70_tp4_push_buffer(ipc_ptrs); -} - // Use vector to represent byte data for python binding compatibility. std::tuple, std::vector> get_graph_buffer_ipc_meta(fptr_t _fa) { diff --git a/csrc/custom_all_reduce.cuh b/csrc/custom_all_reduce.cuh index f98aef77f5..511b7cc668 100644 --- a/csrc/custom_all_reduce.cuh +++ b/csrc/custom_all_reduce.cuh @@ -51,38 +51,9 @@ inline hipPointer_attribute rangeStartAddrAttr = constexpr size_t kSm70Tp2SmallAllreduceBytes = 40 * 1024; constexpr size_t kSm70Tp8HierarchicalAllreduceBytes = 4096 * sizeof(half); -constexpr size_t kSm70Tp4MtpVerifierBytesPerRequest = 5 * 2048 * sizeof(half); constexpr int kSm70Tp8CompletionSignalSlotBase = 2; constexpr int kSm70GemmaRmsNormHiddenSize = 5120; constexpr int kSm70GemmaRmsNormThreads = 1024; -constexpr int kSm70LongPrefillSignalBlocks = - sm70_tile_runtime::kMaxSignalBlocks; -constexpr int kSm70LongPrefillMaxTokensPerRank = 2048; -// DFlash2 verifies eight positions at once. These constants mirror SGLang's -// two-epoch one-shot push collective while keeping its storage separate from -// vLLM's pull-collective metadata and intermediate buffers. -constexpr int kSm70Tp4PushAllreduceWorldSize = 4; -constexpr int kSm70Tp4PushAllreduceBlocks = 80; -constexpr int kSm70Tp4PushAllreduceThreads = 128; -constexpr int kSm70Tp4PushAllreduceEpochs = 2; -constexpr uint16_t kSm70Tp4PushAllreduceSentinel = 0x7f7f; -constexpr int kSm70Tp4PushAllreduceSentinelByte = 0x7f; -constexpr size_t kSm70Tp4PushAllreduceBytes = - 8 * kSm70GemmaRmsNormHiddenSize * sizeof(half); -constexpr size_t kSm70Tp4PushAllreduce8KiBBytes = 4096 * sizeof(half); -constexpr size_t kSm70Tp4PushAllreduceSignalBytes = - ((kSm70Tp4PushAllreduceBlocks * sizeof(uint32_t) + 127) / 128) * 128; -constexpr size_t kSm70Tp4PushAllreduceBufferBytes = - kSm70Tp4PushAllreduceSignalBytes + kSm70Tp4PushAllreduceEpochs * - kSm70Tp4PushAllreduceWorldSize * - kSm70Tp4PushAllreduceBytes; - -inline int sm70_tp4_push_allreduce_blocks(size_t bytes) { - if (bytes == kSm70Tp4PushAllreduceBytes) { - return kSm70Tp4PushAllreduceBlocks; - } - return bytes == kSm70Tp4PushAllreduce8KiBBytes ? 4 : 0; -} inline int sm70_gemma_rms_norm_threads() { const char* raw = std::getenv("VLLM_SM70_TP2_AR_GEMMA_RMS_THREADS"); @@ -93,20 +64,6 @@ inline int sm70_gemma_rms_norm_threads() { : kSm70GemmaRmsNormThreads; } -inline int sm70_tp4_long_fused_norm_threads() { - const char* raw = std::getenv("VLLM_SM70_TP4_LONG_FUSED_NORM_THREADS"); - if (raw == nullptr) return 512; - const int threads = std::atoi(raw); - return threads == 256 || threads == 512 || threads == 1024 ? threads : 512; -} - -inline int sm70_tp4_long_fused_norm_blocks() { - const char* raw = std::getenv("VLLM_SM70_TP4_LONG_FUSED_NORM_BLOCKS"); - if (raw == nullptr) return 80; - const int blocks = std::atoi(raw); - return blocks >= 1 && blocks <= kSm70LongPrefillSignalBlocks ? blocks : 80; -} - inline bool custom_allreduce_current_device_is_sm70() { #ifndef USE_ROCM int device = 0; @@ -129,7 +86,7 @@ inline bool custom_allreduce_current_device_is_sm70() { } inline bool sm70_tp8_hierarchical_custom_ar_enabled(int world_size, - bool fully_connected) { + bool fully_connected) { const char* raw = std::getenv("VLLM_SM70_TP8_HIERARCHICAL_CUSTOM_AR"); return raw != nullptr && std::atoi(raw) != 0 && world_size == 8 && !fully_connected && custom_allreduce_current_device_is_sm70(); @@ -139,30 +96,15 @@ inline bool sm70_tp8_hierarchical_peer(int rank, int peer) { return rank / 4 == peer / 4 || rank + 4 == peer || peer + 4 == rank; } -inline int custom_allreduce_block_limit(int default_limit, int world_size, - bool fully_connected, size_t bytes, - bool tune_sm70_tp4_mtp) { +inline int custom_allreduce_block_limit(int default_limit, + int world_size, + size_t bytes) { const char* raw = std::getenv("VLLM_CUSTOM_ALLREDUCE_BLOCK_LIMIT"); if (raw == nullptr || raw[0] == '\0') { if (world_size == 2 && bytes <= kSm70Tp2SmallAllreduceBytes && custom_allreduce_current_device_is_sm70()) { return 1; } - const char* tune_raw = std::getenv("VLLM_SM70_TP4_MTP_AR_BLOCK_TUNING"); - const bool tuning_enabled = - tune_raw != nullptr && tune_raw[0] != '\0' && std::atoi(tune_raw) != 0; - if (tune_sm70_tp4_mtp && tuning_enabled && - default_limit == defaultBlockLimit && world_size == 4 && - fully_connected && custom_allreduce_current_device_is_sm70()) { - if (bytes == kSm70Tp4MtpVerifierBytesPerRequest) { - return 1; - } - if (bytes == 8 * kSm70Tp4MtpVerifierBytesPerRequest || - bytes == 12 * kSm70Tp4MtpVerifierBytesPerRequest || - bytes == 16 * kSm70Tp4MtpVerifierBytesPerRequest) { - return 8; - } - } return default_limit; } char* end = nullptr; @@ -175,7 +117,8 @@ inline int custom_allreduce_block_limit(int default_limit, int world_size, return static_cast(parsed); } -inline int sm70_tp4_m5_allreduce_threads(int world_size, bool fully_connected, +inline int sm70_tp4_m5_allreduce_threads(int world_size, + bool fully_connected, size_t bytes) { const char* raw = std::getenv("VLLM_SM70_TP4_M5_AR_THREADS"); if (raw == nullptr || raw[0] == '\0') return 512; @@ -196,15 +139,6 @@ inline int sm70_tp4_m5_allreduce_threads(int world_size, bool fully_connected, return static_cast(parsed); } -inline bool sm70_tp4_small_allreduce_pack32(int world_size, - bool fully_connected, - size_t bytes) { - const char* raw = std::getenv("VLLM_SM70_TP4_SMALL_AR_PACK32"); - return raw != nullptr && std::atoi(raw) != 0 && world_size == 4 && - fully_connected && bytes == 5120 * sizeof(half) && - custom_allreduce_current_device_is_sm70(); -} - // Counter may overflow, but unsigned integer overflow is well-defined. using FlagType = sm70_tile_runtime::FlagType; using Signal = sm70_tile_runtime::Signal; @@ -343,7 +277,8 @@ static DINLINE FlagType ld_flag_volatile(FlagType* flag_addr) { } static DINLINE void st_flag_sys_visible(FlagType* flag_addr, FlagType flag) { - asm volatile("membar.sys; st.volatile.global.u32 [%1], %0;" ::"r"(flag), + asm volatile("membar.sys; st.volatile.global.u32 [%1], %0;" + ::"r"(flag), "l"(flag_addr) : "memory"); } @@ -357,47 +292,11 @@ static DINLINE FlagType ld_flag_sys_visible(FlagType* flag_addr) { return flag; } -static DINLINE void membar_sys() { asm volatile("membar.sys;" ::: "memory"); } - -// The TP4 pack32 route waits with volatile peer loads, then executes one -// system fence after the matching flag is visible. This preserves the input -// visibility contract without putting a system fence in every poll iteration. -template -DINLINE void sm70_pack32_barrier_at_start(const RankSignals& sg, - Signal* self_sg, int rank) { - uint32_t flag = self_sg->_flag[blockIdx.x] + 1; - if (threadIdx.x < ngpus) { - const unsigned int peer = threadIdx.x; - auto peer_counter_ptr = &sg.signals[peer]->start[blockIdx.x][rank]; - auto self_counter_ptr = &self_sg->start[blockIdx.x][peer]; - st_flag_sys_visible(peer_counter_ptr, flag); - while (ld_flag_volatile(self_counter_ptr) != flag); - membar_sys(); - } - __syncthreads(); - if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; -} - -template -DINLINE void sm70_pack32_barrier_at_end(const RankSignals& sg, Signal* self_sg, - int rank) { - __syncthreads(); - uint32_t flag = self_sg->_flag[blockIdx.x] + 1; - if (threadIdx.x < ngpus) { - const unsigned int peer = threadIdx.x; - auto peer_counter_ptr = &sg.signals[peer]->end[blockIdx.x][rank]; - auto self_counter_ptr = &self_sg->end[blockIdx.x][peer]; - st_flag_volatile(peer_counter_ptr, flag); - while (ld_flag_volatile(self_counter_ptr) != flag); - } - if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; -} - // This function is meant to be used as the first synchronization in the all -// reduce kernel. Publish the peer flag after a system fence, poll it with a -// volatile load, then execute one system fence after the matching value is -// visible. The post-poll fence preserves input visibility without putting a -// system fence in every spin-loop iteration. +// reduce kernel. The all-reduce input is usually produced by kernels launched +// immediately before this kernel on each rank. Use a system memory fence around +// the peer-visible flag so faster upstream paths do not expose stale input to +// peer ranks. template DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg, int rank) { @@ -408,8 +307,7 @@ DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg, // Write the expected counter value to peer and wait for correct value // from peer. st_flag_sys_visible(peer_counter_ptr, flag); - while (ld_flag_volatile(self_counter_ptr) != flag); - membar_sys(); + while (ld_flag_sys_visible(self_counter_ptr) != flag); } __syncthreads(); // use one thread to update flag @@ -445,38 +343,6 @@ DINLINE void barrier_at_end(const RankSignals& sg, Signal* self_sg, int rank) { #else -template -DINLINE void sm70_pack32_barrier_at_start(const RankSignals& sg, - Signal* self_sg, int rank) { - uint32_t flag = self_sg->_flag[blockIdx.x] + 1; - if (threadIdx.x < ngpus) { - const unsigned int peer = threadIdx.x; - __scoped_atomic_store_n(&sg.signals[peer]->start[blockIdx.x][rank], flag, - __ATOMIC_RELAXED, __MEMORY_SCOPE_SYSTEM); - while (__scoped_atomic_load_n(&self_sg->start[blockIdx.x][peer], - __ATOMIC_RELAXED, - __MEMORY_SCOPE_DEVICE) < flag); - } - __syncthreads(); - if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; -} - -template -DINLINE void sm70_pack32_barrier_at_end(const RankSignals& sg, Signal* self_sg, - int rank) { - __syncthreads(); - uint32_t flag = self_sg->_flag[blockIdx.x] + 1; - if (threadIdx.x < ngpus) { - const unsigned int peer = threadIdx.x; - __scoped_atomic_store_n(&sg.signals[peer]->end[blockIdx.x][rank], flag, - __ATOMIC_RELAXED, __MEMORY_SCOPE_SYSTEM); - while (__scoped_atomic_load_n(&self_sg->end[blockIdx.x][peer], - __ATOMIC_RELAXED, - __MEMORY_SCOPE_DEVICE) < flag); - } - if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; -} - template DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg, int rank) { @@ -530,54 +396,6 @@ DINLINE P packed_reduce(const P* ptrs[], int idx) { return downcast

(tmp); } -// Adapted from SGLang-V100's SM70-capable one-shot push collective at -// haohervchb/sglang-V100@845b9fdf7a7e. SGLang uses positive zero as its empty -// slot sentinel. This variant preinitializes slots with a reserved FP16 NaN -// payload instead, preserving every finite payload bit, including signed zero. -DINLINE void sm70_push_escape_sentinel(half& value) { - auto* bits = reinterpret_cast(&value); - if (*bits == kSm70Tp4PushAllreduceSentinel) *bits = 0x7e00u; -} - -DINLINE bool sm70_push_is_sentinel(const half& value) { - const auto* bits = reinterpret_cast(&value); - return *bits == kSm70Tp4PushAllreduceSentinel; -} - -template -DINLINE void sm70_push_load_volatile_16b(P& value, const void* address, - int offset) { - static_assert(alignof(P) == 16 && sizeof(P) == 16); - const auto* source = reinterpret_cast(address) + offset; - uint4 bits; - asm volatile("ld.volatile.global.v4.b32 {%0, %1, %2, %3}, [%4];" - : "=r"(bits.x), "=r"(bits.y), "=r"(bits.z), "=r"(bits.w) - : "l"(source)); - value = *reinterpret_cast(&bits); -} - -template -DINLINE void sm70_push_store_volatile_16b(const P& value, void* address, - int offset) { - static_assert(alignof(P) == 16 && sizeof(P) == 16); - const uint4 bits = *reinterpret_cast(&value); - auto* destination = reinterpret_cast(address) + offset; - asm volatile("st.volatile.global.v4.b32 [%4], {%0, %1, %2, %3};" - : - : "r"(bits.x), "r"(bits.y), "r"(bits.z), "r"(bits.w), - "l"(destination)); -} - -template -DINLINE P sm70_push_reduce(P (&values)[ngpus]) { - A accumulator = upcast(values[0]); -#pragma unroll - for (int rank = 1; rank < ngpus; ++rank) { - packed_assign_add(accumulator, upcast(values[rank])); - } - return downcast

(accumulator); -} - template DINLINE P packed_reduce_sum2(const P* ptrs_a[], const P* ptrs_b[], int idx) { P local = ptrs_a[0][idx]; @@ -610,113 +428,17 @@ __global__ void __launch_bounds__(512, 1) barrier_at_end(sg, self_sg, rank); } -template -__global__ void __launch_bounds__(512, 1) - sm70_cross_device_reduce_1stage_pack32(RankData* _dp, RankSignals sg, - Signal* self_sg, - half* __restrict__ result, int rank, - int size) { - using P = array_t; - using A = array_t; - auto dp = *_dp; - sm70_pack32_barrier_at_start(sg, self_sg, rank); - for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size; - idx += gridDim.x * blockDim.x) { - reinterpret_cast(result)[idx] = - packed_reduce((const P**)&dp.ptrs[0], idx); - } - sm70_pack32_barrier_at_end(sg, self_sg, rank); -} - -template -__global__ void __launch_bounds__(1024, 1) - sm70_cross_device_reduce_1stage_push(RankData push_buffers, - const half* __restrict__ input, - half* __restrict__ output, int rank, - int packed_size) { - static_assert(ngpus == kSm70Tp4PushAllreduceWorldSize); - using P = typename packed_t::P; - using A = typename packed_t::A; - - auto* local_storage = - const_cast(reinterpret_cast(push_buffers.ptrs[rank])); - auto* local_epochs = reinterpret_cast(local_storage); - const uint32_t epoch = local_epochs[blockIdx.x]; - constexpr int packed_stride = kSm70Tp4PushAllreduceBytes / sizeof(P); - const int epoch_offset = epoch * ngpus * packed_stride; - const int offset = blockIdx.x * blockDim.x + threadIdx.x; - - if (offset < packed_size) { - P value = reinterpret_cast(input)[offset]; -#pragma unroll - for (int element = 0; element < P::size; ++element) { - sm70_push_escape_sentinel(value.data[element]); - } - -#pragma unroll - for (int destination_rank = 0; destination_rank < ngpus; - ++destination_rank) { - auto* destination_base = const_cast( - reinterpret_cast(push_buffers.ptrs[destination_rank])); - void* destination = destination_base + kSm70Tp4PushAllreduceSignalBytes + - (epoch_offset + rank * packed_stride) * sizeof(P); - sm70_push_store_volatile_16b(value, destination, offset); - } - - P peer_values[ngpus]; - while (true) { - bool has_empty_slot = false; -#pragma unroll - for (int source_rank = 0; source_rank < ngpus; ++source_rank) { - const void* source = - local_storage + kSm70Tp4PushAllreduceSignalBytes + - (epoch_offset + source_rank * packed_stride) * sizeof(P); - sm70_push_load_volatile_16b(peer_values[source_rank], source, offset); -#pragma unroll - for (int element = 0; element < P::size; ++element) { - has_empty_slot |= - sm70_push_is_sentinel(peer_values[source_rank].data[element]); - } - } - if (!has_empty_slot) break; - } - - reinterpret_cast(output)[offset] = - sm70_push_reduce(peer_values); - - P empty; -#pragma unroll - for (int element = 0; element < P::size; ++element) { - *reinterpret_cast(&empty.data[element]) = - kSm70Tp4PushAllreduceSentinel; - } -#pragma unroll - for (int source_rank = 0; source_rank < ngpus; ++source_rank) { - void* source = local_storage + kSm70Tp4PushAllreduceSignalBytes + - (epoch_offset + source_rank * packed_stride) * sizeof(P); - sm70_push_store_volatile_16b(empty, source, offset); - } - } - - __syncthreads(); - if (threadIdx.x == 0) { - local_epochs[blockIdx.x] = (epoch + 1) % kSm70Tp4PushAllreduceEpochs; - } -} - // This prototype deliberately stays narrow: it mirrors the FP32 Gemma RMSNorm // reduction order for a [tokens, 5120] FP16 projection. Each CTA handles one // row, retains the normal all-reduce peer order, and applies only its local // residual after the peer reduction. template __global__ void __launch_bounds__(Threads, 1) - sm70_peer_reduce_gemma_rms_norm(RankData* _dp, RankSignals sg, - Signal* self_sg, - const ResidualT* __restrict__ residual, - const WeightT* __restrict__ weight, - half* __restrict__ normalized_out, - float* __restrict__ residual_out, int rank, - float epsilon) { + sm70_peer_reduce_gemma_rms_norm( + RankData* _dp, RankSignals sg, Signal* self_sg, + const ResidualT* __restrict__ residual, + const WeightT* __restrict__ weight, half* __restrict__ normalized_out, + float* __restrict__ residual_out, int rank, float epsilon) { using P = typename packed_t::P; using A = typename packed_t::A; constexpr int kPackedWidth = P::size; @@ -737,14 +459,15 @@ __global__ void __launch_bounds__(Threads, 1) for (int packed_idx = tid; packed_idx < packed_per_row; packed_idx += blockDim.x) { - const P reduced = packed_reduce( - (const P**)&dp.ptrs[0], row * packed_per_row + packed_idx); + const P reduced = + packed_reduce((const P**)&dp.ptrs[0], + row * packed_per_row + packed_idx); const int element_offset = row_offset + packed_idx * kPackedWidth; #pragma unroll for (int i = 0; i < kPackedWidth; ++i) { - const float value = - __half2float(reduced.data[i]) + - sm70_gemma_rms_norm_to_float(residual[element_offset + i]); + const float value = __half2float(reduced.data[i]) + + sm70_gemma_rms_norm_to_float( + residual[element_offset + i]); residual_values[packed_idx * kPackedWidth + i] = value; residual_out[element_offset + i] = value; } @@ -780,140 +503,20 @@ __global__ void __launch_bounds__(Threads, 1) const int column = element_offset + i; const float gemma_weight = sm70_gemma_rms_norm_to_float(weight[column]) + 1.0f; - normalized_out[row_offset + column] = - __float2half_rn(residual_values[column] * inverse_rms * gemma_weight); + normalized_out[row_offset + column] = __float2half_rn( + residual_values[column] * inverse_rms * gemma_weight); } } barrier_at_end(sg, self_sg, rank); } -// Benchmark-only SM70 counterpart of NCCL's fused LSA RMSNorm example: -// reduce-scatter token rows, apply the mixed-dtype Gemma RMSNorm locally, then -// all-gather by writing the owned normalized rows into every peer output. -// Each rank launches one CTA per owned token and therefore moves the same -// asymptotic peer traffic as ring reduce-scatter + all-gather, without -// materializing an intermediate all-reduce tensor. -template -__global__ void __launch_bounds__(Threads, 1) - sm70_peer_reduce_scatter_gemma_rms_norm_all_gather( - RankData* _input_dp, RankData* _output_dp, RankSignals sg, - Signal* self_sg, const ResidualT* __restrict__ residual, - const WeightT* __restrict__ weight, float* __restrict__ residual_out, - int rank, int num_tokens, float epsilon) { - static_assert(ngpus == 4); - using H4 = array_t; - constexpr int kVarianceVectorWidth = 4; - constexpr int kVectorsPerRow = - kSm70GemmaRmsNormHiddenSize / kVarianceVectorWidth; - constexpr int kVectorsPerThread = (kVectorsPerRow + Threads - 1) / Threads; - - __shared__ float inverse_rms; - // Preserve the accepted local GemmaNorm reduction topology. The valid-item - // count below limits participation to the threads in this launch. - using BlockReduce = cub::BlockReduce; - __shared__ typename BlockReduce::TempStorage reduce_store; - - const int tokens_per_rank = num_tokens / ngpus; - const int tid = threadIdx.x; - auto input_dp = *_input_dp; - auto output_dp = *_output_dp; - - barrier_at_start(sg, self_sg, rank); - - for (int local_row = blockIdx.x; local_row < tokens_per_rank; - local_row += gridDim.x) { - const int row = rank * tokens_per_rank + local_row; - const int vector_row_offset = row * kVectorsPerRow; - float4 row_values[kVectorsPerThread]; - float variance = 0.0f; -#pragma unroll - for (int iter = 0; iter < kVectorsPerThread; ++iter) { - const int vector_idx = tid + iter * Threads; - if (vector_idx >= kVectorsPerRow) continue; - const int global_vector_idx = vector_row_offset + vector_idx; - H4 reduced = - reinterpret_cast(input_dp.ptrs[0])[global_vector_idx]; -#pragma unroll - for (int peer = 1; peer < ngpus; ++peer) { - const H4 peer_value = - reinterpret_cast(input_dp.ptrs[peer])[global_vector_idx]; -#pragma unroll - for (int i = 0; i < kVarianceVectorWidth; ++i) { - // NCCL's f16 Sum specialization uses __hadd/__hadd2 at every ring - // step. Preserve that FP16 rounding contract instead of accumulating - // all four inputs in FP32 and rounding only once. - reduced.data[i] = __hadd(reduced.data[i], peer_value.data[i]); - } - } - const float4 residual_value = - reinterpret_cast(residual)[global_vector_idx]; - float4 value; - value.x = __half2float(reduced.data[0]) + residual_value.x; - value.y = __half2float(reduced.data[1]) + residual_value.y; - value.z = __half2float(reduced.data[2]) + residual_value.z; - value.w = __half2float(reduced.data[3]) + residual_value.w; - row_values[iter] = value; - // Keep the public [M, H] residual shape so the compiled model graph does - // not change shape. Only this rank's persistent token shard is valid; - // the next fused boundary reads the same shard and ignores other rows. - reinterpret_cast(residual_out)[global_vector_idx] = value; - // Match the accepted local GemmaNorm exactly: each thread accumulates - // vector indices tid, tid + Threads, ... in x/y/z/w order before CUB. - variance += value.x * value.x; - variance += value.y * value.y; - variance += value.z * value.z; - variance += value.w * value.w; - } - variance = - BlockReduce(reduce_store).Reduce(variance, CubAddOp{}, blockDim.x); - if (tid == 0) { - inverse_rms = rsqrtf(variance / kSm70GemmaRmsNormHiddenSize + epsilon); - } - __syncthreads(); - - const float scale = inverse_rms; -#pragma unroll - for (int iter = 0; iter < kVectorsPerThread; ++iter) { - const int vector_idx = tid + iter * Threads; - if (vector_idx >= kVectorsPerRow) continue; - const int column = vector_idx * kVarianceVectorWidth; - const float4 value = row_values[iter]; - H4 normalized_pack; - normalized_pack.data[0] = __float2half_rn( - value.x * scale * - (sm70_gemma_rms_norm_to_float(weight[column]) + 1.0f)); - normalized_pack.data[1] = __float2half_rn( - value.y * scale * - (sm70_gemma_rms_norm_to_float(weight[column + 1]) + 1.0f)); - normalized_pack.data[2] = __float2half_rn( - value.z * scale * - (sm70_gemma_rms_norm_to_float(weight[column + 2]) + 1.0f)); - normalized_pack.data[3] = __float2half_rn( - value.w * scale * - (sm70_gemma_rms_norm_to_float(weight[column + 3]) + 1.0f)); -#pragma unroll - for (int peer = 0; peer < ngpus; ++peer) { - auto* peer_output = - reinterpret_cast(const_cast(output_dp.ptrs[peer])); - peer_output[vector_row_offset + vector_idx] = normalized_pack; - } - } - // CUB's temporary storage is reused by the next persistent row. - __syncthreads(); - } - - // The normalized rows are written directly into peer IPC buffers. Unlike - // an ordinary final all-reduce barrier, this barrier must publish those - // remote stores before a peer starts its local output copy. - barrier_at_end(sg, self_sg, rank); -} - template __global__ void __launch_bounds__(256, 1) sm70_tile_runtime_reduce_kernel( - RankData* _dp, RankSignals sg, Signal* self_sg, const T* __restrict__ input, - T* __restrict__ staging, T* __restrict__ result, int rank, int packed_size, - int tile_packed_size, int tile_count, int compute_iters) { + RankData* _dp, RankSignals sg, Signal* self_sg, + const T* __restrict__ input, T* __restrict__ staging, + T* __restrict__ result, int rank, int packed_size, int tile_packed_size, + int tile_count, int compute_iters) { using P = typename packed_t::P; using A = typename packed_t::A; @@ -964,10 +567,11 @@ __global__ void __launch_bounds__(256, 1) sm70_tile_runtime_reduce_kernel( template __global__ void __launch_bounds__(256, 1) sm70_tile_runtime_engine_kernel( - RankData* _dp, RankSignals sg, Signal* self_sg, const T* __restrict__ input, - T* __restrict__ staging, T* __restrict__ result, int rank, int packed_size, - int tile_packed_size, int tile_count, int producer_blocks, - int reducer_blocks, int compute_iters) { + RankData* _dp, RankSignals sg, Signal* self_sg, + const T* __restrict__ input, T* __restrict__ staging, + T* __restrict__ result, int rank, int packed_size, int tile_packed_size, + int tile_count, int producer_blocks, int reducer_blocks, + int compute_iters) { using P = typename packed_t::P; using A = typename packed_t::A; @@ -1030,11 +634,10 @@ __global__ void __launch_bounds__(256, 1) sm70_tile_runtime_engine_kernel( template __global__ void __launch_bounds__(256, 1) - sm70_tile_runtime_wait_reduce_kernel(RankData* _dp, RankSignals sg, - Signal* self_sg, - T* __restrict__ result, int rank, - int packed_size, int tile_packed_size, - int tile_count) { + sm70_tile_runtime_wait_reduce_kernel( + RankData* _dp, RankSignals sg, Signal* self_sg, + T* __restrict__ result, int rank, int packed_size, + int tile_packed_size, int tile_count) { using P = typename packed_t::P; using A = typename packed_t::A; @@ -1066,11 +669,9 @@ __global__ void __launch_bounds__(256, 1) } template -__global__ void __launch_bounds__(512, 1) - cross_device_reduce_sum2_1stage(RankData* _dp_a, RankData* _dp_b, - RankSignals sg, Signal* self_sg, - T* __restrict__ result, int rank, - int size) { +__global__ void __launch_bounds__(512, 1) cross_device_reduce_sum2_1stage( + RankData* _dp_a, RankData* _dp_b, RankSignals sg, Signal* self_sg, + T* __restrict__ result, int rank, int size) { using P = typename packed_t::P; using A = typename packed_t::A; auto dp_a = *_dp_a; @@ -1078,8 +679,9 @@ __global__ void __launch_bounds__(512, 1) barrier_at_start(sg, self_sg, rank); for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size; idx += gridDim.x * blockDim.x) { - ((P*)result)[idx] = packed_reduce_sum2( - (const P**)&dp_a.ptrs[0], (const P**)&dp_b.ptrs[0], idx); + ((P*)result)[idx] = + packed_reduce_sum2((const P**)&dp_a.ptrs[0], + (const P**)&dp_b.ptrs[0], idx); } barrier_at_end(sg, self_sg, rank); } @@ -1089,8 +691,9 @@ DINLINE P* get_tmp_buf(Signal* sg) { return (P*)(((Signal*)sg) + 1); } -DINLINE FlagType sm70_tp8_clique_barrier(const RankSignals& sg, Signal* self_sg, - int rank, int signal_slot) { +DINLINE FlagType sm70_tp8_clique_barrier(const RankSignals& sg, + Signal* self_sg, int rank, + int signal_slot) { const int tid = threadIdx.x; const int clique_base = rank < 4 ? 0 : 4; const FlagType flag = self_sg->_flag[0] + 1; @@ -1104,10 +707,9 @@ DINLINE FlagType sm70_tp8_clique_barrier(const RankSignals& sg, Signal* self_sg, return flag; } -static __global__ void __launch_bounds__(512, 1) - sm70_tp8_hierarchical_reduce(RankData* _dp, RankSignals sg, Signal* self_sg, - half* __restrict__ result, int rank, - int packed_size) { +static __global__ void __launch_bounds__(512, 1) sm70_tp8_hierarchical_reduce( + RankData* _dp, RankSignals sg, Signal* self_sg, + half* __restrict__ result, int rank, int packed_size) { using P = typename packed_t::P; using A = typename packed_t::A; @@ -1121,7 +723,8 @@ static __global__ void __launch_bounds__(512, 1) const int partial_slot = (self_sg->_flag[0] >> 1) & 1; const FlagType clique_flag = sm70_tp8_clique_barrier(sg, self_sg, rank, partial_slot); - A* const self_partial = get_tmp_buf(self_sg) + partial_slot * packed_size; + A* const self_partial = + get_tmp_buf(self_sg) + partial_slot * packed_size; const A* const pair_partial = get_tmp_buf(sg.signals[pair_rank]) + partial_slot * packed_size; @@ -1144,9 +747,11 @@ static __global__ void __launch_bounds__(512, 1) const FlagType pair_flag = clique_flag + 1; if (tid < 4) { const int peer = clique_base + tid; - const int completion_slot = kSm70Tp8CompletionSignalSlotBase + partial_slot; + const int completion_slot = + kSm70Tp8CompletionSignalSlotBase + partial_slot; st_flag_volatile(&sg.signals[peer]->end[completion_slot][rank], pair_flag); - while (ld_flag_volatile(&self_sg->end[completion_slot][peer]) != pair_flag); + while (ld_flag_volatile(&self_sg->end[completion_slot][peer]) != + pair_flag); } else if (tid == 4) { st_flag_release(&sg.signals[pair_rank]->end[partial_slot][rank], pair_flag); while (ld_flag_acquire(&self_sg->end[partial_slot][pair_rank]) != @@ -1209,11 +814,9 @@ __global__ void __launch_bounds__(512, 1) } template -__global__ void __launch_bounds__(512, 1) - cross_device_reduce_sum2_2stage(RankData* _dp_a, RankData* _dp_b, - RankSignals sg, Signal* self_sg, - T* __restrict__ result, int rank, - int size) { +__global__ void __launch_bounds__(512, 1) cross_device_reduce_sum2_2stage( + RankData* _dp_a, RankData* _dp_b, RankSignals sg, Signal* self_sg, + T* __restrict__ result, int rank, int size) { int tid = blockIdx.x * blockDim.x + threadIdx.x; int stride = gridDim.x * blockDim.x; using P = typename packed_t::P; @@ -1238,7 +841,8 @@ __global__ void __launch_bounds__(512, 1) // Stage 1 mirrors cross_device_reduce_2stage, but each rank first forms // its local input_a + input_b value before the cross-rank reduction. for (int idx = start + tid; idx < end; idx += stride) { - tmp_out[idx - start] = packed_reduce_sum2(ptrs_a, ptrs_b, idx); + tmp_out[idx - start] = + packed_reduce_sum2(ptrs_a, ptrs_b, idx); } barrier_at_end(sg, self_sg, rank); @@ -1318,8 +922,6 @@ class CustomAllreduce { std::vector graph_unreg_buffers_; // a map from IPC handles to opened IPC pointers std::map ipc_handles_; - RankData sm70_tp4_push_buffers_{}; - bool sm70_tp4_push_buffers_registered_ = false; /** * Signals are an array of ipc-enabled buffers from all ranks. @@ -1400,27 +1002,6 @@ class CustomAllreduce { buffers_[ptrs[rank_]] = d_data; } - void register_sm70_tp4_push_buffer(void** ptrs) { - if (world_size_ != kSm70Tp4PushAllreduceWorldSize || !fully_connected_ || - !custom_allreduce_current_device_is_sm70()) { - throw std::runtime_error( - "SM70 push all-reduce requires fully-connected TP4 on SM70."); - } - for (int peer = 0; peer < world_size_; ++peer) { - if (ptrs[peer] == nullptr) { - throw std::runtime_error( - "SM70 push all-reduce received a null peer buffer."); - } - sm70_tp4_push_buffers_.ptrs[peer] = ptrs[peer]; - } - auto* local_data = - static_cast(ptrs[rank_]) + kSm70Tp4PushAllreduceSignalBytes; - CUDACHECK(cudaMemset( - local_data, kSm70Tp4PushAllreduceSentinelByte, - kSm70Tp4PushAllreduceBufferBytes - kSm70Tp4PushAllreduceSignalBytes)); - sm70_tp4_push_buffers_registered_ = true; - } - RankData* rank_data_for_buffer(cudaStream_t stream, void* buffer, const char* op_name) { RankData* ptrs; @@ -1432,10 +1013,11 @@ class CustomAllreduce { } else { auto it = buffers_.find(buffer); if (it == buffers_.end()) { - throw std::runtime_error( - std::string(op_name) + " buffer address " + - std::to_string(reinterpret_cast(buffer)) + - " is not registered!"); + throw std::runtime_error(std::string(op_name) + + " buffer address " + + std::to_string( + reinterpret_cast(buffer)) + + " is not registered!"); } ptrs = it->second; } @@ -1495,8 +1077,7 @@ class CustomAllreduce { void allreduce(cudaStream_t stream, T* input, T* output, int size, int threads = 512, int block_limit = defaultBlockLimit) { block_limit = custom_allreduce_block_limit( - block_limit, world_size_, fully_connected_, - static_cast(size) * sizeof(T), true); + block_limit, world_size_, static_cast(size) * sizeof(T)); auto d = packed_t::P::size; if (size % d != 0) throw std::runtime_error( @@ -1527,21 +1108,6 @@ class CustomAllreduce { size /= d; auto bytes = size * sizeof(typename packed_t::P); if constexpr (std::is_same_v) { - // The push protocol amortizes peer polling across captured collective - // chains. A lone eager call stays on the ordinary registered-buffer - // pull path. - if (sm70_tp4_push_buffers_registered_ && - status == cudaStreamCaptureStatusActive && - world_size_ == kSm70Tp4PushAllreduceWorldSize && fully_connected_ && - custom_allreduce_current_device_is_sm70()) { - const int push_blocks = sm70_tp4_push_allreduce_blocks(bytes); - if (push_blocks > 0) { - sm70_cross_device_reduce_1stage_push - <<>>( - sm70_tp4_push_buffers_, input, output, rank_, size); - return; - } - } if (sm70_tp8_hierarchical_custom_ar_enabled(world_size_, fully_connected_) && bytes == kSm70Tp8HierarchicalAllreduceBytes) { @@ -1550,20 +1116,10 @@ class CustomAllreduce { return; } } - threads = - sm70_tp4_m5_allreduce_threads(world_size_, fully_connected_, bytes); + threads = sm70_tp4_m5_allreduce_threads(world_size_, fully_connected_, + bytes); int blocks = std::min(block_limit, (size + threads - 1) / threads); - if constexpr (std::is_same_v) { - if (blocks == 1 && threads == 512 && - sm70_tp4_small_allreduce_pack32(world_size_, fully_connected_, - bytes)) { - sm70_cross_device_reduce_1stage_pack32<4><<<1, 512, 0, stream>>>( - ptrs, sg_, self_sg_, output, rank_, bytes / 32); - return; - } - } - // Check environment variable once const char* env_algo = std::getenv("VLLM_CUSTOM_ALLREDUCE_ALGO"); bool force_1stage = false; @@ -1624,11 +1180,11 @@ class CustomAllreduce { template void sm70_allreduce_gemma_rms_norm(cudaStream_t stream, half* input, - const ResidualT* residual, - const WeightT* weight, - half* normalized_out, float* residual_out, - int num_tokens, int hidden_size, - float epsilon) { + const ResidualT* residual, + const WeightT* weight, + half* normalized_out, + float* residual_out, int num_tokens, + int hidden_size, float epsilon) { if (world_size_ != ngpus || !custom_allreduce_current_device_is_sm70()) { throw std::runtime_error("SM70 Gemma RMSNorm prototype requires TP" + std::to_string(ngpus) + " on an SM70 device."); @@ -1642,18 +1198,18 @@ class CustomAllreduce { if (num_tokens <= 0 || num_tokens > kMaxTokens) { throw std::runtime_error( "SM70 Gemma RMSNorm prototype supports tokens in [1, " + - std::to_string(kMaxTokens) + "]. Got " + std::to_string(num_tokens) + - "."); + std::to_string(kMaxTokens) + "]. Got " + + std::to_string(num_tokens) + "."); } - RankData* ptrs = rank_data_for_buffer(stream, input, - "SM70 Gemma RMSNorm prototype input"); + RankData* ptrs = rank_data_for_buffer( + stream, input, "SM70 Gemma RMSNorm prototype input"); if constexpr (ngpus == 4) { // Keep the same 1024-thread CUB reduction topology as vLLM rms_norm. // packed_reduce preserves the cross_device_reduce_1stage // rank order before the FP32 residual add. sm70_peer_reduce_gemma_rms_norm + ResidualT, WeightT> <<>>( ptrs, sg_, self_sg_, residual, weight, normalized_out, residual_out, rank_, epsilon); @@ -1661,11 +1217,11 @@ class CustomAllreduce { } const int threads = sm70_gemma_rms_norm_threads(); -#define VLLM_LAUNCH_SM70_GEMMA_RMS_NORM(THREADS) \ +#define VLLM_LAUNCH_SM70_GEMMA_RMS_NORM(THREADS) \ sm70_peer_reduce_gemma_rms_norm \ - <<>>(ptrs, sg_, self_sg_, residual, \ - weight, normalized_out, \ - residual_out, rank_, epsilon) + <<>>( \ + ptrs, sg_, self_sg_, residual, weight, normalized_out, \ + residual_out, rank_, epsilon) switch (threads) { case 256: VLLM_LAUNCH_SM70_GEMMA_RMS_NORM(256); @@ -1680,63 +1236,12 @@ class CustomAllreduce { #undef VLLM_LAUNCH_SM70_GEMMA_RMS_NORM } - template - void sm70_reduce_scatter_gemma_rms_norm_all_gather( - cudaStream_t stream, half* input, half* shared_output, - const ResidualT* residual, const WeightT* weight, float* residual_out, - int num_tokens, int hidden_size, float epsilon) { - if (world_size_ != ngpus || !fully_connected_ || - !custom_allreduce_current_device_is_sm70()) { - throw std::runtime_error( - "SM70 long-prefill fused norm requires fully connected TP" + - std::to_string(ngpus) + " on SM70."); - } - if (hidden_size != kSm70GemmaRmsNormHiddenSize || num_tokens <= 0 || - num_tokens % ngpus != 0) { - throw std::runtime_error( - "SM70 long-prefill fused norm requires a TP-divisible [M, 5120] " - "input."); - } - const int tokens_per_rank = num_tokens / ngpus; - if (tokens_per_rank > kSm70LongPrefillMaxTokensPerRank) { - throw std::runtime_error( - "SM70 long-prefill fused norm exceeds the token-shard capacity."); - } - - RankData* input_ptrs = - rank_data_for_buffer(stream, input, "long-prefill fused norm input"); - RankData* output_ptrs = rank_data_for_buffer( - stream, shared_output, "long-prefill fused norm output"); - const int threads = sm70_tp4_long_fused_norm_threads(); - const int blocks = - std::min(tokens_per_rank, sm70_tp4_long_fused_norm_blocks()); -#define VLLM_LAUNCH_SM70_TP4_LONG_FUSED_NORM(THREADS) \ - sm70_peer_reduce_scatter_gemma_rms_norm_all_gather \ - <<>>(input_ptrs, output_ptrs, sg_, self_sg_, \ - residual, weight, residual_out, rank_, \ - num_tokens, epsilon) - switch (threads) { - case 256: - VLLM_LAUNCH_SM70_TP4_LONG_FUSED_NORM(256); - break; - case 1024: - VLLM_LAUNCH_SM70_TP4_LONG_FUSED_NORM(1024); - break; - default: - VLLM_LAUNCH_SM70_TP4_LONG_FUSED_NORM(512); - break; - } -#undef VLLM_LAUNCH_SM70_TP4_LONG_FUSED_NORM - } - template void allreduce_sum2(cudaStream_t stream, T* input_a, T* input_b, T* output, int size, int threads = 512, int block_limit = defaultBlockLimit) { block_limit = custom_allreduce_block_limit( - block_limit, world_size_, fully_connected_, - static_cast(size) * sizeof(T), false); + block_limit, world_size_, static_cast(size) * sizeof(T)); auto d = packed_t::P::size; if (size % d != 0) throw std::runtime_error( @@ -1788,28 +1293,29 @@ class CustomAllreduce { } } -#define SUM2_KL(ngpus, name) \ - name<<>>( \ - ptrs_a, ptrs_b, sg_, self_sg_, output, rank_, size); -#define SUM2_CASE(ngpus) \ - case ngpus: { \ - if (force_1stage) { \ - SUM2_KL(ngpus, cross_device_reduce_sum2_1stage); \ - } else if (force_2stage) { \ - SUM2_KL(ngpus, cross_device_reduce_sum2_2stage); \ - } else { \ - if (world_size_ == 2) { \ - SUM2_KL(ngpus, cross_device_reduce_sum2_1stage); \ - } else if (fully_connected_) { \ - if ((world_size_ <= 4 && bytes < 512 * 1024) || \ - (world_size_ <= 8 && bytes < 256 * 1024)) { \ +#define SUM2_KL(ngpus, name) \ + name<<>>(ptrs_a, ptrs_b, sg_, \ + self_sg_, output, rank_, \ + size); +#define SUM2_CASE(ngpus) \ + case ngpus: { \ + if (force_1stage) { \ + SUM2_KL(ngpus, cross_device_reduce_sum2_1stage); \ + } else if (force_2stage) { \ + SUM2_KL(ngpus, cross_device_reduce_sum2_2stage); \ + } else { \ + if (world_size_ == 2) { \ + SUM2_KL(ngpus, cross_device_reduce_sum2_1stage); \ + } else if (fully_connected_) { \ + if ((world_size_ <= 4 && bytes < 512 * 1024) || \ + (world_size_ <= 8 && bytes < 256 * 1024)) { \ SUM2_KL(ngpus, cross_device_reduce_sum2_1stage); \ - } else { \ + } else { \ SUM2_KL(ngpus, cross_device_reduce_sum2_2stage); \ - } \ - } \ - } \ - break; \ + } \ + } \ + } \ + break; \ } switch (world_size_) { @@ -1849,12 +1355,12 @@ class CustomAllreduce { const int packed_size = size / pack; const int tile_packed_size = tile_numel / pack; - const int tile_count = - (packed_size + tile_packed_size - 1) / tile_packed_size; + const int tile_count = (packed_size + tile_packed_size - 1) / tile_packed_size; if (tile_count <= 0 || tile_count > kMaxBlocks) { throw std::runtime_error( "SM70 tile runtime prototype supports tile_count in [1, " + - std::to_string(kMaxBlocks) + "]. Got " + std::to_string(tile_count)); + std::to_string(kMaxBlocks) + "]. Got " + + std::to_string(tile_count)); } auto it = buffers_.find(staging); @@ -1871,12 +1377,14 @@ class CustomAllreduce { blocks = std::max(1, std::min(blocks, tile_count)); compute_iters = std::max(0, compute_iters); -#define TILE_RUNTIME_CASE(ngpus) \ - case ngpus: { \ - sm70_tile_runtime_reduce_kernel<<>>( \ - ptrs, sg_, self_sg_, input, staging, output, rank_, packed_size, \ - tile_packed_size, tile_count, compute_iters); \ - break; \ +#define TILE_RUNTIME_CASE(ngpus) \ + case ngpus: { \ + sm70_tile_runtime_reduce_kernel \ + <<>>(ptrs, sg_, self_sg_, input, staging, \ + output, rank_, packed_size, \ + tile_packed_size, tile_count, \ + compute_iters); \ + break; \ } switch (world_size_) { @@ -1916,7 +1424,8 @@ class CustomAllreduce { if (tile_count <= 0 || tile_count > kMaxBlocks) { throw std::runtime_error( "SM70 tile runtime engine supports tile_count in [1, " + - std::to_string(kMaxBlocks) + "]. Got " + std::to_string(tile_count)); + std::to_string(kMaxBlocks) + "]. Got " + + std::to_string(tile_count)); } auto it = buffers_.find(staging); @@ -1938,13 +1447,15 @@ class CustomAllreduce { const int threads = 256; const int blocks = producer_blocks + reducer_blocks; -#define TILE_RUNTIME_ENGINE_CASE(ngpus) \ - case ngpus: { \ - sm70_tile_runtime_engine_kernel<<>>( \ - ptrs, sg_, self_sg_, input, staging, output, rank_, packed_size, \ - tile_packed_size, tile_count, producer_blocks, reducer_blocks, \ - compute_iters); \ - break; \ +#define TILE_RUNTIME_ENGINE_CASE(ngpus) \ + case ngpus: { \ + sm70_tile_runtime_engine_kernel \ + <<>>(ptrs, sg_, self_sg_, input, staging, \ + output, rank_, packed_size, \ + tile_packed_size, tile_count, \ + producer_blocks, reducer_blocks, \ + compute_iters); \ + break; \ } switch (world_size_) { @@ -1958,7 +1469,8 @@ class CustomAllreduce { template void tile_runtime_wait_reduce(cudaStream_t stream, T* staging, T* output, - int size, int tile_numel, int reducer_blocks) { + int size, int tile_numel, + int reducer_blocks) { if (world_size_ != 2) { throw std::runtime_error( "SM70 tile runtime wait-reduce currently supports only TP2."); @@ -1982,7 +1494,8 @@ class CustomAllreduce { if (tile_count <= 0 || tile_count > kMaxBlocks) { throw std::runtime_error( "SM70 tile runtime wait-reduce supports tile_count in [1, " + - std::to_string(kMaxBlocks) + "]. Got " + std::to_string(tile_count)); + std::to_string(kMaxBlocks) + "]. Got " + + std::to_string(tile_count)); } RankData* ptrs; @@ -2008,13 +1521,13 @@ class CustomAllreduce { constexpr int threads = 256; -#define TILE_RUNTIME_WAIT_REDUCE_CASE(ngpus) \ - case ngpus: { \ - sm70_tile_runtime_wait_reduce_kernel \ - <<>>( \ - ptrs, sg_, self_sg_, output, rank_, packed_size, tile_packed_size, \ - tile_count); \ - break; \ +#define TILE_RUNTIME_WAIT_REDUCE_CASE(ngpus) \ + case ngpus: { \ + sm70_tile_runtime_wait_reduce_kernel \ + <<>>( \ + ptrs, sg_, self_sg_, output, rank_, packed_size, \ + tile_packed_size, tile_count); \ + break; \ } switch (world_size_) { @@ -2043,11 +1556,11 @@ class CustomAllreduce { ptrs = it->second; } -#define TOP1_CASE(ngpus) \ - case ngpus: { \ - cross_device_top1_argmax \ - <<<1, 32, 0, stream>>>(ptrs, sg_, self_sg_, output, rank_); \ - break; \ +#define TOP1_CASE(ngpus) \ + case ngpus: { \ + cross_device_top1_argmax<<<1, 32, 0, stream>>>( \ + ptrs, sg_, self_sg_, output, rank_); \ + break; \ } switch (world_size_) { diff --git a/csrc/libtorch_stable/layernorm_kernels.cu b/csrc/libtorch_stable/layernorm_kernels.cu index ce6a524b36..bc5b367b7d 100644 --- a/csrc/libtorch_stable/layernorm_kernels.cu +++ b/csrc/libtorch_stable/layernorm_kernels.cu @@ -26,25 +26,20 @@ __global__ void rms_norm_kernel( const int64_t input_stride_d4, // input.stride(-4) const int64_t input_shape_d2, // input.size(-2) const int64_t input_shape_d3, // input.size(-3) - const scalar_t* __restrict__ weight, // [hidden_size] or - // [num_groups, hidden_size] - const int64_t weight_stride, // 0 or weight.stride(0) + const scalar_t* __restrict__ weight, // [hidden_size] const float epsilon, const int num_tokens, const int hidden_size) { __shared__ float s_variance; float variance = 0.0f; const scalar_t* input_row; - const scalar_t* weight_row; if constexpr (NUM_DIMS == 2) { // 2D for layernorm normal case [batch_size, hidden] input_row = input + blockIdx.x * input_stride_d2; - weight_row = weight + blockIdx.x * weight_stride; } else if constexpr (NUM_DIMS == 3) { // 3D for q/k norm [batch_size, num_heads, head_size] int batch_idx = blockIdx.x / input_shape_d2; int head_idx = blockIdx.x % input_shape_d2; input_row = input + batch_idx * input_stride_d3 + head_idx * input_stride_d2; - weight_row = weight + batch_idx * weight_stride; } else if constexpr (NUM_DIMS == 4) { // 4D for transformers model_impl qk norm [batch, seq, head, head_dim] int batch_idx = blockIdx.x / (input_shape_d3 * input_shape_d2); @@ -53,7 +48,6 @@ __global__ void rms_norm_kernel( int head_idx = remaining % input_shape_d2; input_row = input + batch_idx * input_stride_d4 + seq_idx * input_stride_d3 + head_idx * input_stride_d2; - weight_row = weight + batch_idx * weight_stride; } auto vec_op = [&variance](const vec_n_t& vec) { @@ -81,7 +75,7 @@ __global__ void rms_norm_kernel( scalar_t* out_row = out + blockIdx.x * hidden_size; auto* v_in = reinterpret_cast*>(input_row); - auto* v_w = reinterpret_cast*>(weight_row); + auto* v_w = reinterpret_cast*>(weight); auto* v_out = reinterpret_cast*>(out_row); for (int i = threadIdx.x; i < hidden_size / VEC_SIZE; i += blockDim.x) { vec_n_t dst; @@ -297,8 +291,7 @@ __global__ void __launch_bounds__(kSm70GemmaLongPrefillThreads, 2) void rms_norm(torch::stable::Tensor& out, // [..., hidden_size] torch::stable::Tensor& input, // [..., hidden_size] - torch::stable::Tensor& weight, // [hidden_size] or - // [num_groups, hidden_size] + torch::stable::Tensor& weight, // [hidden_size] double epsilon) { STD_TORCH_CHECK(out.is_contiguous()); if (input.stride(-1) != 1) { @@ -306,16 +299,6 @@ void rms_norm(torch::stable::Tensor& out, // [..., hidden_size] } STD_TORCH_CHECK(input.stride(-1) == 1); STD_TORCH_CHECK(weight.is_contiguous()); - int64_t weight_stride = 0; - if (weight.dim() == 1) { - STD_TORCH_CHECK(weight.size(0) == input.size(-1)); - } else if (weight.dim() == 2) { - STD_TORCH_CHECK(weight.size(0) == input.size(0)); - STD_TORCH_CHECK(weight.size(-1) == input.size(-1)); - weight_stride = weight.stride(0); - } else { - STD_TORCH_CHECK(false, "rms_norm weight must be 1D or 2D"); - } int hidden_size = input.size(-1); @@ -347,8 +330,8 @@ void rms_norm(torch::stable::Tensor& out, // [..., hidden_size] out.mutable_data_ptr(), input.const_data_ptr(), input_stride_d2, input_stride_d3, input_stride_d4, input_shape_d2, - input_shape_d3, weight.const_data_ptr(), - weight_stride, epsilon, num_tokens, hidden_size); + input_shape_d3, weight.const_data_ptr(), epsilon, + num_tokens, hidden_size); }); }); }); 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 290e531181..1329091679 100644 --- a/csrc/moe/marlin_moe_wna16/sm70_marlin_fp8_gemm.cu +++ b/csrc/moe/marlin_moe_wna16/sm70_marlin_fp8_gemm.cu @@ -35,25 +35,22 @@ class Sm70MoeFp8IteratorB { static int const kPackedMacroN = PackedMacroN_; static constexpr bool kUseMetadataVectorWords = UseMetadataVectorWords_; using Element = cutlass::half_t; - 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."); + 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."); 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 { @@ -86,11 +83,12 @@ 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), @@ -132,7 +130,9 @@ 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,8 +181,9 @@ 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 { @@ -243,9 +244,8 @@ 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,8 +267,9 @@ 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)); } @@ -276,9 +277,10 @@ 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; @@ -292,17 +294,22 @@ 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; @@ -316,17 +323,20 @@ 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>( @@ -356,7 +366,8 @@ 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)); } @@ -375,17 +386,18 @@ struct Sm70MoeFp8GemmSpec { using IteratorA = Sm70MoeGatherIteratorA; template - using IteratorB = Sm70MoeFp8IteratorB; + using IteratorB = + Sm70MoeFp8IteratorB; }; -template using Sm70MoeFp8GemmTraits = - Sm70MarlinMoeGemmTraits, CtaM, - CtaN, CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, - PackedMacroN>; + Sm70MarlinMoeGemmTraits< + Sm70MoeFp8GemmSpec, CtaM, CtaN, CtaK, + Warps, WarpM, WarpN, WarpK, GroupSize, PackedMacroN>; template struct Sm70MoeFp8Launcher { @@ -407,16 +419,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); } }; @@ -434,10 +446,9 @@ 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, "."); @@ -446,45 +457,17 @@ 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 2cd0b71f90..fdc516a3a6 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,14 +68,16 @@ 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}; @@ -150,7 +152,8 @@ 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); } } @@ -159,14 +162,16 @@ 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}; @@ -273,14 +278,16 @@ 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}; @@ -371,16 +378,17 @@ 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}; @@ -392,57 +400,57 @@ sm70_marlin_moe_try_select_cyankiwi_qwen3_5_122b_a10b_awq_8bit_params( // 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); } } } @@ -452,13 +460,13 @@ sm70_marlin_moe_try_select_cyankiwi_qwen3_5_122b_a10b_awq_8bit_params( // 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); } } } @@ -468,13 +476,13 @@ sm70_marlin_moe_try_select_cyankiwi_qwen3_5_122b_a10b_awq_8bit_params( // 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); } } } @@ -484,33 +492,33 @@ sm70_marlin_moe_try_select_cyankiwi_qwen3_5_122b_a10b_awq_8bit_params( // 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); } } } @@ -519,14 +527,16 @@ sm70_marlin_moe_try_select_cyankiwi_qwen3_5_122b_a10b_awq_8bit_params( } 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}; @@ -538,63 +548,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); } } } @@ -604,16 +614,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); } } } @@ -623,13 +633,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); } } } @@ -639,36 +649,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); } } } @@ -676,16 +686,17 @@ 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}; @@ -697,63 +708,63 @@ 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); } } } @@ -763,16 +774,16 @@ 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); } } } @@ -782,13 +793,13 @@ 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); } } } @@ -798,45 +809,45 @@ 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); } } } @@ -844,16 +855,17 @@ 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}; @@ -865,63 +877,63 @@ 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); } } } @@ -931,13 +943,13 @@ 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); } } } @@ -947,23 +959,23 @@ 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); } } } @@ -973,33 +985,33 @@ 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); } } } @@ -1007,16 +1019,17 @@ 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}; @@ -1028,63 +1041,63 @@ 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); } } } @@ -1094,13 +1107,13 @@ 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); } } } @@ -1110,23 +1123,23 @@ 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); } } } @@ -1136,33 +1149,33 @@ 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); } } } @@ -1171,14 +1184,16 @@ 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}; @@ -1189,64 +1204,64 @@ inline bool sm70_marlin_moe_try_select_nvidia_glm_4_7_nvfp4_params( if (ctx.moe_block_size == 8) { // 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_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.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); } } } @@ -1256,19 +1271,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); } } } @@ -1278,13 +1293,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); } } } @@ -1294,42 +1309,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); } } } @@ -1338,14 +1353,16 @@ 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}; @@ -1357,63 +1374,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); } } } @@ -1423,19 +1440,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); } } } @@ -1445,13 +1462,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); } } } @@ -1461,45 +1478,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); } } } @@ -1508,14 +1525,16 @@ 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}; @@ -1527,63 +1546,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); } } } @@ -1593,16 +1612,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); } } } @@ -1612,13 +1631,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); } } } @@ -1628,48 +1647,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); } } } @@ -1677,16 +1696,17 @@ 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}; @@ -1698,69 +1718,69 @@ sm70_marlin_moe_try_select_nvidia_qwen3_next_80b_a3b_thinking_nvfp4_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, 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); } } } @@ -1770,13 +1790,13 @@ sm70_marlin_moe_try_select_nvidia_qwen3_next_80b_a3b_thinking_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=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); } } } @@ -1786,23 +1806,23 @@ sm70_marlin_moe_try_select_nvidia_qwen3_next_80b_a3b_thinking_nvfp4_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, 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); } } } @@ -1812,36 +1832,36 @@ sm70_marlin_moe_try_select_nvidia_qwen3_next_80b_a3b_thinking_nvfp4_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, 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); } } } @@ -1850,14 +1870,16 @@ sm70_marlin_moe_try_select_nvidia_qwen3_next_80b_a3b_thinking_nvfp4_params( } 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}; @@ -1973,14 +1995,16 @@ 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}; @@ -2084,14 +2108,16 @@ 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}; @@ -2180,21 +2206,20 @@ 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; } @@ -2203,8 +2228,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; } @@ -2233,17 +2258,18 @@ 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; } @@ -2257,12 +2283,13 @@ 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; } @@ -2276,44 +2303,46 @@ 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 ", @@ -2328,17 +2357,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 ", @@ -2349,15 +2378,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 " @@ -2419,9 +2448,10 @@ struct Sm70MarlinMoeFp8GroupSizeDispatchLauncher { template torch::Tensor operator()() const { - return dispatch_sm70_marlin_moe_fp8_group_size< - CtaM, CtaN, CtaK, Warps, WarpM, WarpN, WarpK, PackedMacroN>(inner, - group_size); + return dispatch_sm70_marlin_moe_fp8_group_size( + inner, group_size); } }; @@ -2435,59 +2465,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) @@ -2507,8 +2537,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); } @@ -2517,18 +2547,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"); } @@ -2573,8 +2603,8 @@ class Sm70MoeGatherIteratorA { using Element = cutlass::half_t; using Layout = cutlass::layout::RowMajor; using TensorCoord = cutlass::MatrixCoord; - using Fragment = cutlass::Array; + using Fragment = cutlass::Array< + Element, ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; struct Params { int lda; int moe_block_size; @@ -2659,8 +2689,9 @@ 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; @@ -2668,17 +2699,20 @@ 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 @@ -2720,8 +2754,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>; @@ -2729,16 +2763,13 @@ 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; - using IteratorB = - typename Spec::template IteratorB; + 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>; using Mma = Sm70MarlinMmaPipelined< ThreadblockShape, IteratorA, typename MmaCore::SmemIteratorA, IteratorB, typename MmaCore::SmemIteratorB, ElementAccumulator, LayoutC, @@ -2787,11 +2818,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(); @@ -2807,17 +2838,20 @@ 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; @@ -2862,7 +2896,8 @@ 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); } @@ -2871,11 +2906,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 @@ -2921,19 +2956,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; @@ -2967,8 +3002,8 @@ __launch_bounds__(Traits::MmaCore::kThreads, 1) 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; } @@ -2979,8 +3014,8 @@ __launch_bounds__(Traits::MmaCore::kThreads, 1) 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), @@ -3007,8 +3042,8 @@ __launch_bounds__(Traits::MmaCore::kThreads, 1) 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 @@ -3068,9 +3103,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); @@ -3080,11 +3115,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 ec11d6d72e..733d9a4959 100644 --- a/csrc/moe/marlin_moe_wna16/sm70_marlin_moe_dispatch.cu +++ b/csrc/moe/marlin_moe_wna16/sm70_marlin_moe_dispatch.cu @@ -43,13 +43,15 @@ __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]; } @@ -64,9 +66,10 @@ __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; @@ -80,14 +83,15 @@ __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])); } @@ -107,8 +111,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<<>>( + sm70_unpermute_half_metadata_kernel<<< + blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>( reinterpret_cast(b_scales.data_ptr()), reinterpret_cast(out.data_ptr()), total, perm_len); C10_CUDA_KERNEL_LAUNCH_CHECK(); @@ -116,8 +120,9 @@ 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; } @@ -138,13 +143,14 @@ 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<<>>( + sm70_unpack_moe_zp_to_half_kernel<<< + blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>( b_zeros.data_ptr(), reinterpret_cast(b_scales.data_ptr()), reinterpret_cast(out.data_ptr()), total, @@ -223,7 +229,8 @@ 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, @@ -301,7 +308,8 @@ 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."); @@ -356,8 +364,10 @@ 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."); @@ -447,15 +457,17 @@ 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, @@ -498,9 +510,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 { @@ -523,7 +535,8 @@ 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, @@ -546,18 +559,20 @@ 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); } @@ -565,8 +580,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); } @@ -607,9 +622,8 @@ 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 9755b3edbf..76443f12b8 100644 --- a/csrc/moe/marlin_moe_wna16/sm70_marlin_mxfp4_gemm.cu +++ b/csrc/moe/marlin_moe_wna16/sm70_marlin_mxfp4_gemm.cu @@ -51,8 +51,7 @@ 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_; @@ -60,25 +59,22 @@ class Sm70MoeMxfp4IteratorB { static int const kGroupSize = GroupSize_; static int const kPackedMacroN = PackedMacroN_; using Element = cutlass::half_t; - 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."); + 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."); 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; @@ -89,7 +85,9 @@ 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: @@ -110,8 +108,9 @@ 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), @@ -150,7 +149,9 @@ 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 { @@ -183,12 +184,14 @@ 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); } @@ -241,9 +244,8 @@ 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_; @@ -263,13 +265,15 @@ 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; @@ -292,7 +296,8 @@ 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; @@ -315,9 +320,8 @@ 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; @@ -360,15 +364,15 @@ struct Sm70MoeMxfp4GemmSpec { using IteratorA = Sm70MoeGatherIteratorA; template - using IteratorB = - Sm70MoeMxfp4IteratorB; + using IteratorB = Sm70MoeMxfp4IteratorB; }; -template +template using Sm70MoeMxfp4GemmTraits = - Sm70MarlinMoeGemmTraits; + Sm70MarlinMoeGemmTraits; struct Sm70MoeMxfp4Launcher { torch::Tensor& a; @@ -389,15 +393,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); } }; @@ -415,10 +419,9 @@ 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, "."); @@ -426,23 +429,10 @@ 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 b2cc86aa91..a1ea955d34 100644 --- a/csrc/moe/marlin_moe_wna16/sm70_marlin_nvfp4_gemm.cu +++ b/csrc/moe/marlin_moe_wna16/sm70_marlin_nvfp4_gemm.cu @@ -26,8 +26,7 @@ namespace { constexpr int kNvfp4ValuesPerWord = 8; -template +template class Sm70MoeNvfp4IteratorB { public: using Shape = Shape_; @@ -35,25 +34,22 @@ class Sm70MoeNvfp4IteratorB { static int const kGroupSize = GroupSize_; static int const kPackedMacroN = PackedMacroN_; using Element = cutlass::half_t; - 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."); + 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."); 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; @@ -64,7 +60,9 @@ 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: @@ -85,8 +83,9 @@ 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), @@ -125,7 +124,9 @@ 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 { @@ -158,12 +159,14 @@ 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); } @@ -216,9 +219,8 @@ 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_; @@ -238,13 +240,15 @@ 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; @@ -267,7 +271,8 @@ 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; @@ -290,9 +295,8 @@ 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; @@ -335,15 +339,15 @@ struct Sm70MoeNvfp4GemmSpec { using IteratorA = Sm70MoeGatherIteratorA; template - using IteratorB = - Sm70MoeNvfp4IteratorB; + using IteratorB = Sm70MoeNvfp4IteratorB; }; -template +template using Sm70MoeNvfp4GemmTraits = - Sm70MarlinMoeGemmTraits; + Sm70MarlinMoeGemmTraits; struct Sm70MoeNvfp4Launcher { torch::Tensor& a; @@ -364,15 +368,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); } }; @@ -390,32 +394,18 @@ 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 bce2bfb899..cf412f4d0f 100644 --- a/csrc/moe/marlin_moe_wna16/sm70_marlin_u4_gemm.cu +++ b/csrc/moe/marlin_moe_wna16/sm70_marlin_u4_gemm.cu @@ -35,23 +35,21 @@ class Sm70MoeU4ZpIteratorB { static int const kPackedMacroN = PackedMacroN_; static constexpr bool kUseMetadataVectorWords = UseMetadataVectorWords_; using Element = cutlass::half_t; - 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."); + 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."); 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; @@ -87,10 +85,11 @@ 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), @@ -132,14 +131,17 @@ 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; @@ -156,16 +158,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]; @@ -176,8 +178,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]; @@ -185,8 +187,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]; @@ -199,8 +201,9 @@ 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 { @@ -225,8 +228,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( @@ -247,8 +250,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( @@ -257,18 +260,16 @@ 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( @@ -280,8 +281,9 @@ 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)); } @@ -289,7 +291,8 @@ 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; @@ -311,7 +314,8 @@ 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; @@ -333,11 +337,9 @@ 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_; @@ -345,8 +347,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( @@ -364,7 +366,8 @@ 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)); } @@ -383,17 +386,18 @@ struct Sm70MoeU4ZpGemmSpec { using IteratorA = Sm70MoeGatherIteratorA; template - using IteratorB = Sm70MoeU4ZpIteratorB; + using IteratorB = + Sm70MoeU4ZpIteratorB; }; -template using Sm70MoeU4ZpGemmTraits = - Sm70MarlinMoeGemmTraits, CtaM, - CtaN, CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, - PackedMacroN>; + Sm70MarlinMoeGemmTraits< + Sm70MoeU4ZpGemmSpec, CtaM, CtaN, CtaK, + Warps, WarpM, WarpN, WarpK, GroupSize, PackedMacroN>; template struct Sm70MoeU4Launcher { @@ -415,16 +419,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); } }; @@ -442,10 +446,9 @@ 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, "."); @@ -453,45 +456,19 @@ 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 c5e03f0ece..56f4206460 100644 --- a/csrc/moe/marlin_moe_wna16/sm70_marlin_u4b8_gemm.cu +++ b/csrc/moe/marlin_moe_wna16/sm70_marlin_u4b8_gemm.cu @@ -35,25 +35,22 @@ class Sm70MoeU4B8IteratorB { static int const kPackedMacroN = PackedMacroN_; static constexpr bool kUseMetadataVectorWords = UseMetadataVectorWords_; using Element = cutlass::half_t; - 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."); + 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."); 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; @@ -87,8 +84,9 @@ 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), @@ -130,17 +128,19 @@ 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,8 +181,9 @@ 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 { @@ -237,17 +238,16 @@ 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,8 +259,9 @@ 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)); } @@ -268,7 +269,8 @@ 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; @@ -289,7 +291,8 @@ 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; @@ -310,9 +313,8 @@ 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; @@ -340,7 +342,8 @@ 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)); } @@ -359,17 +362,18 @@ struct Sm70MoeU4B8GemmSpec { using IteratorA = Sm70MoeGatherIteratorA; template - using IteratorB = Sm70MoeU4B8IteratorB; + using IteratorB = + Sm70MoeU4B8IteratorB; }; -template using Sm70MoeU4B8GemmTraits = - Sm70MarlinMoeGemmTraits, CtaM, - CtaN, CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, - PackedMacroN>; + Sm70MarlinMoeGemmTraits< + Sm70MoeU4B8GemmSpec, CtaM, CtaN, CtaK, + Warps, WarpM, WarpN, WarpK, GroupSize, PackedMacroN>; template struct Sm70MoeU4B8Launcher { @@ -391,16 +395,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); } }; @@ -418,10 +422,9 @@ 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, "."); @@ -430,43 +433,17 @@ 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 25c86f4f2a..c620f30112 100644 --- a/csrc/moe/marlin_moe_wna16/sm70_marlin_u8_gemm.cu +++ b/csrc/moe/marlin_moe_wna16/sm70_marlin_u8_gemm.cu @@ -35,24 +35,22 @@ class Sm70MoeU8ZpIteratorB { static int const kPackedMacroN = PackedMacroN_; static constexpr bool kUseMetadataVectorWords = UseMetadataVectorWords_; using Element = cutlass::half_t; - 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."); + 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."); 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 { @@ -90,10 +88,11 @@ 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), @@ -135,14 +134,17 @@ 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; @@ -159,16 +161,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]; @@ -179,8 +181,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]; @@ -188,8 +190,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]; @@ -202,8 +204,9 @@ 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 { @@ -266,9 +269,8 @@ 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>( @@ -278,12 +280,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])); } @@ -291,8 +293,9 @@ 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)); } @@ -300,9 +303,10 @@ 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; @@ -317,17 +321,22 @@ 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; @@ -342,17 +351,20 @@ 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>( @@ -383,7 +395,8 @@ 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)); } @@ -402,17 +415,18 @@ struct Sm70MoeU8ZpGemmSpec { using IteratorA = Sm70MoeGatherIteratorA; template - using IteratorB = Sm70MoeU8ZpIteratorB; + using IteratorB = + Sm70MoeU8ZpIteratorB; }; -template using Sm70MoeU8ZpGemmTraits = - Sm70MarlinMoeGemmTraits, CtaM, - CtaN, CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, - PackedMacroN>; + Sm70MarlinMoeGemmTraits< + Sm70MoeU8ZpGemmSpec, CtaM, CtaN, CtaK, + Warps, WarpM, WarpN, WarpK, GroupSize, PackedMacroN>; template struct Sm70MoeU8Launcher { @@ -434,16 +448,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); } }; @@ -461,10 +475,9 @@ 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, "."); @@ -472,45 +485,19 @@ 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 1eabec24dc..511e4d4181 100644 --- a/csrc/moe/marlin_moe_wna16/sm70_marlin_u8b128_gemm.cu +++ b/csrc/moe/marlin_moe_wna16/sm70_marlin_u8b128_gemm.cu @@ -35,25 +35,22 @@ class Sm70MoeU8B128IteratorB { static int const kPackedMacroN = PackedMacroN_; static constexpr bool kUseMetadataVectorWords = UseMetadataVectorWords_; using Element = cutlass::half_t; - 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."); + 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."); 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 { @@ -89,8 +86,9 @@ 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), @@ -132,17 +130,19 @@ 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,8 +183,9 @@ 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 { @@ -245,9 +246,8 @@ 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,8 +269,9 @@ 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)); } @@ -278,9 +279,10 @@ 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; @@ -294,17 +296,22 @@ 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; @@ -318,17 +325,20 @@ 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>( @@ -358,7 +368,8 @@ 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)); } @@ -382,13 +393,13 @@ struct Sm70MoeU8B128GemmSpec { UseMetadataVectorWords>; }; -template using Sm70MoeU8B128GemmTraits = - Sm70MarlinMoeGemmTraits, CtaM, - CtaN, CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, - PackedMacroN>; + Sm70MarlinMoeGemmTraits< + Sm70MoeU8B128GemmSpec, CtaM, CtaN, CtaK, + Warps, WarpM, WarpN, WarpK, GroupSize, PackedMacroN>; template struct Sm70MoeU8B128Launcher { @@ -410,17 +421,16 @@ 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); } }; @@ -438,10 +448,9 @@ 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, "."); @@ -450,43 +459,17 @@ 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 334629cc98..9cd84897ff 100644 --- a/csrc/moe/moe_permute_unpermute_op.cu +++ b/csrc/moe/moe_permute_unpermute_op.cu @@ -28,7 +28,9 @@ 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 @@ -70,15 +72,16 @@ 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) { @@ -87,7 +90,8 @@ 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) { @@ -96,7 +100,8 @@ 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 eaf633af6b..bf86c12abb 100644 --- a/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu +++ b/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu @@ -93,12 +93,10 @@ 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]; @@ -179,8 +177,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; @@ -188,7 +186,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; @@ -224,12 +222,11 @@ 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 fcce90dad4..d87ed201de 100644 --- a/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h +++ b/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h @@ -80,12 +80,11 @@ 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 e2fd08b575..2ee0b700f5 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -116,7 +116,8 @@ std::vector fp8_sm70_prepare(torch::Tensor _kernel, int64_t group_size, bool interleave_gated_silu); -void fp8_sm70_dequantize_out(torch::Tensor out, torch::Tensor _kernel, +void fp8_sm70_dequantize_out(torch::Tensor out, + torch::Tensor _kernel, torch::Tensor _scaling_factors, int64_t group_size); @@ -132,205 +133,131 @@ std::vector nvfp4_sm70_prepare(torch::Tensor _kernel, std::vector sm70_f16_prepare(torch::Tensor _kernel); -torch::Tensor awq_gemm_sm70(torch::Tensor _in_feats, torch::Tensor _kernel, - torch::Tensor _scaling_factors, int64_t group_size, - int64_t k_ld, int64_t q_ld); +torch::Tensor awq_gemm_sm70(torch::Tensor _in_feats, + torch::Tensor _kernel, + torch::Tensor _scaling_factors, + int64_t group_size, + int64_t k_ld, + int64_t q_ld); torch::Tensor sm70_f16_gemm(torch::Tensor _in_feats, torch::Tensor _kernel); -void awq_gemm_sm70_out(torch::Tensor out, torch::Tensor _in_feats, - torch::Tensor _kernel, torch::Tensor _scaling_factors, - int64_t group_size, int64_t k_ld, int64_t q_ld, +void awq_gemm_sm70_out(torch::Tensor out, + torch::Tensor _in_feats, + torch::Tensor _kernel, + torch::Tensor _scaling_factors, + int64_t group_size, + int64_t k_ld, + int64_t q_ld, bool gated_silu); -void awq_gemm_sm70_out_tile_reduce( - torch::Tensor out, torch::Tensor staging, torch::Tensor _in_feats, - torch::Tensor _kernel, torch::Tensor _scaling_factors, int64_t group_size, - int64_t k_ld, int64_t q_ld, int64_t fa_ptr, int64_t tile_numel, - int64_t reducer_blocks, int64_t kernel_reducer_blocks, bool overlap); - -void fp8_gemm_sm70_out(torch::Tensor out, torch::Tensor _in_feats, - torch::Tensor _kernel, torch::Tensor _scaling_factors, - int64_t group_size, int64_t k_ld, int64_t q_ld, +void awq_gemm_sm70_out_tile_reduce(torch::Tensor out, + torch::Tensor staging, + torch::Tensor _in_feats, + torch::Tensor _kernel, + torch::Tensor _scaling_factors, + int64_t group_size, + int64_t k_ld, + int64_t q_ld, + int64_t fa_ptr, + int64_t tile_numel, + int64_t reducer_blocks, + int64_t kernel_reducer_blocks, + bool overlap); + +void fp8_gemm_sm70_out(torch::Tensor out, + torch::Tensor _in_feats, + torch::Tensor _kernel, + torch::Tensor _scaling_factors, + int64_t group_size, + int64_t k_ld, + int64_t q_ld, bool gated_silu); -std::vector fp8_qpn8_prepare_sm70(torch::Tensor qweight, - torch::Tensor scales); - -void fp8_qpn8_dequantize_sm70_out(torch::Tensor out, torch::Tensor codes, - torch::Tensor group_scales); - -void fp8_qpn8_prefill_sm70_out(torch::Tensor out, int64_t dense_weight_ptr, - torch::Tensor input, torch::Tensor codes, - torch::Tensor group_scales, bool gated_silu); - -void fp8_qpn8_dispatch_sm70_out(torch::Tensor out, int64_t dense_weight_ptr, - torch::Tensor input, torch::Tensor codes, - torch::Tensor group_scales, int64_t split_k, - int64_t accumulator_chains, bool prefetch_codes, - bool gated_silu); - -void fp8_qpn8_gemm_sm70_out(torch::Tensor out, torch::Tensor input, - torch::Tensor codes, torch::Tensor group_scales, - int64_t split_k, int64_t accumulator_chains, - bool fast_decoder, bool prefetch_codes); - -void fp8_qpn8_gemm_ba_split_sm70_out(torch::Tensor qkv_out, torch::Tensor z_out, - torch::Tensor b_out, torch::Tensor a_out, - torch::Tensor input, torch::Tensor codes, - torch::Tensor group_scales, - torch::Tensor ba_weight); - -void fp8_qpn8_dispatch_ba_split_sm70_out( - torch::Tensor qkv_out, torch::Tensor z_out, torch::Tensor b_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); - -void fp8_qpn8_gated_pair_sm70_out(torch::Tensor out, torch::Tensor input, - torch::Tensor codes, - torch::Tensor group_scales, int64_t split_k, - int64_t accumulator_chains, bool fast_decoder, - bool prefetch_codes); - -std::vector nvfp4_qpn4_prepare_sm70(torch::Tensor qweight, - torch::Tensor scales); - -std::vector nvfp4_qpn4_prepare_scale_code_sm70( - torch::Tensor qweight, torch::Tensor scale_codes); - -void nvfp4_qpn4_dequantize_sm70_out(torch::Tensor out, torch::Tensor codes, - torch::Tensor scales, double global_scale, - bool use_scale_code); - -void nvfp4_qpn4_prefill_sm70_out(torch::Tensor out, int64_t dense_weight_ptr, - torch::Tensor input, torch::Tensor codes, - torch::Tensor scales, double global_scale, - bool use_scale_code, bool gated_silu); - -void nvfp4_qpn4_dispatch_sm70_out(torch::Tensor out, int64_t dense_weight_ptr, - torch::Tensor input, torch::Tensor codes, - torch::Tensor scales, double global_scale, - bool use_scale_code, bool gated_silu); - -void fp8_gemm_sm70_prefill_prescaled_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); - -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); - -void nvfp4_qpn2_gemm_sm70_out(torch::Tensor out, torch::Tensor input, - torch::Tensor codes, torch::Tensor scales, - double global_scale, int64_t split_k, - int64_t accumulator_chains); - -void nvfp4_qpn2_gated_sm70_out(torch::Tensor out, torch::Tensor input, - torch::Tensor codes, torch::Tensor scales, - double global_scale, int64_t split_k, - int64_t accumulator_chains); - -void nvfp4_qpn2_dispatch_sm70_out(torch::Tensor out, torch::Tensor input, - torch::Tensor codes, torch::Tensor scales, - double global_scale, int64_t split_k, - int64_t accumulator_chains, - torch::Tensor tm_weight, - torch::Tensor tm_scales, - int64_t tm_group_size, int64_t tm_k_ld, - int64_t tm_q_ld, bool gated_silu); - 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, int64_t k_ld, int64_t q_ld, bool gated_silu, int64_t min_prefill_m); -void mxfp4_gemm_sm70_out(torch::Tensor out, torch::Tensor _in_feats, - torch::Tensor _kernel, torch::Tensor _scaling_factors, - int64_t group_size, int64_t k_ld, int64_t q_ld, +void mxfp4_gemm_sm70_out(torch::Tensor out, + torch::Tensor _in_feats, + torch::Tensor _kernel, + torch::Tensor _scaling_factors, + int64_t group_size, + int64_t k_ld, + int64_t q_ld, bool gated_silu); -void nvfp4_gemm_sm70_out(torch::Tensor out, torch::Tensor _in_feats, - torch::Tensor _kernel, torch::Tensor _scaling_factors, - int64_t group_size, int64_t k_ld, int64_t q_ld, +void nvfp4_gemm_sm70_out(torch::Tensor out, + torch::Tensor _in_feats, + torch::Tensor _kernel, + torch::Tensor _scaling_factors, + int64_t group_size, + int64_t k_ld, + int64_t q_ld, bool gated_silu); -void nvfp4_gemv_sm70_raw_out(torch::Tensor out, torch::Tensor _in_feats, +void nvfp4_gemv_sm70_raw_out(torch::Tensor out, + torch::Tensor _in_feats, torch::Tensor _kernel, torch::Tensor _scaling_factors, - torch::Tensor partials, int64_t group_size, + torch::Tensor partials, + int64_t group_size, int64_t split_k); -void nvfp4_gemv_sm70_warp_out(torch::Tensor out, torch::Tensor _in_feats, +void nvfp4_gemv_sm70_warp_out(torch::Tensor out, + torch::Tensor _in_feats, torch::Tensor _kernel, torch::Tensor _scaling_factors, int64_t group_size); -void nvfp4_gemv_sm70_h2_out(torch::Tensor out, torch::Tensor _in_feats, +void nvfp4_gemv_sm70_h2_out(torch::Tensor out, + torch::Tensor _in_feats, torch::Tensor _kernel, torch::Tensor _scaling_factors, - torch::Tensor partials, int64_t group_size, + torch::Tensor partials, + int64_t group_size, int64_t split_k); -void fp8_gemm_sm70_out_auto(torch::Tensor out, torch::Tensor _in_feats, +void fp8_gemm_sm70_out_auto(torch::Tensor out, + torch::Tensor _in_feats, torch::Tensor _kernel, torch::Tensor _scaling_factors); -void fp8_gemm_sm70_out_meta(torch::Tensor out, torch::Tensor _in_feats, +void fp8_gemm_sm70_out_meta(torch::Tensor out, + torch::Tensor _in_feats, torch::Tensor _kernel, - torch::Tensor _scaling_factors, torch::Tensor _meta, + torch::Tensor _scaling_factors, + torch::Tensor _meta, bool gated_silu); -void sm70_f16_gemm_out(torch::Tensor out, torch::Tensor _in_feats, - torch::Tensor _kernel, int64_t k_ld, bool gated_silu); - -void sm70_f16_indexed_rerank_out(torch::Tensor out, torch::Tensor _in_feats, - torch::Tensor _kernel, - torch::Tensor candidate_ids, - torch::Tensor selected_raw, - torch::Tensor selected_packed, - torch::Tensor expanded, torch::Tensor partials, - torch::Tensor barriers, int64_t cta_n, - int64_t split_k); - -void sm70_f16_indexed_rerank_packed_out( - torch::Tensor out, torch::Tensor _in_feats, torch::Tensor _packed_kernel, - torch::Tensor candidate_ids, torch::Tensor selected_packed, - torch::Tensor expanded, torch::Tensor partials, torch::Tensor barriers, - int64_t cta_n, int64_t split_k); - -void sm70_f16_rerank_keys_out(torch::Tensor keys, torch::Tensor logits, - torch::Tensor candidate_ids); - -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_f16_gemm_out(torch::Tensor out, + torch::Tensor _in_feats, + torch::Tensor _kernel, + int64_t k_ld, + bool gated_silu); void sm70_f16_lm_head_top1_out(torch::Tensor values_out, torch::Tensor indices_out, - torch::Tensor _in_feats, torch::Tensor _kernel, - int64_t k_ld, int64_t vocab_start_index, + torch::Tensor _in_feats, + torch::Tensor _kernel, + int64_t k_ld, + int64_t vocab_start_index, int64_t num_vocab_padding); void sm70_f16_lm_head_top1_tc_out(torch::Tensor values_out, torch::Tensor indices_out, torch::Tensor _in_feats, - torch::Tensor _kernel, int64_t k_ld, + torch::Tensor _kernel, + int64_t k_ld, int64_t vocab_start_index, int64_t num_vocab_padding); void sm70_f16_lm_head_top20_tc_out(torch::Tensor values_out, torch::Tensor indices_out, torch::Tensor _in_feats, - torch::Tensor _kernel, int64_t k_ld, + torch::Tensor _kernel, + int64_t k_ld, int64_t vocab_start_index, int64_t num_vocab_padding); @@ -346,20 +273,27 @@ void sm70_sample_packed_top20_out(torch::Tensor sampled_token_out, torch::Tensor sparse_ids_out, torch::Tensor sparse_probs_out, torch::Tensor gathered_pairs, - torch::Tensor exponential, double top_p); + torch::Tensor exponential, + double top_p); void sm70_dynamic_draft_vocab_update_tail_out( - torch::Tensor lru_token_ids, torch::Tensor local_tail_token_ids, - torch::Tensor source_row_indices, torch::Tensor observed_output_ids, - torch::Tensor target_candidate_ids, torch::Tensor base_token_mask, - int64_t full_vocab_size, int64_t local_shard_start, + torch::Tensor lru_token_ids, + torch::Tensor local_tail_token_ids, + torch::Tensor source_row_indices, + torch::Tensor observed_output_ids, + torch::Tensor target_candidate_ids, + torch::Tensor base_token_mask, + int64_t full_vocab_size, + int64_t local_shard_start, int64_t local_shard_end); void sm70_dynamic_draft_vocab_refresh_tail_weight_out( - torch::Tensor local_tail_weight, torch::Tensor source_weight, + torch::Tensor local_tail_weight, + torch::Tensor source_weight, torch::Tensor source_row_indices); -void sm70_f16_gate_mul_out(torch::Tensor out, torch::Tensor _in_feats, +void sm70_f16_gate_mul_out(torch::Tensor out, + torch::Tensor _in_feats, torch::Tensor _gate_weight); int64_t sm70_gemm_import_cache(torch::Tensor device_hint, @@ -368,188 +302,343 @@ int64_t sm70_gemm_import_cache(torch::Tensor device_hint, int64_t sm70_gemm_export_cache(torch::Tensor device_hint, const std::string& path); -std::vector awq_moe_build_strided_ptrs(torch::Tensor tm_weights, - torch::Tensor tm_scales, - int64_t k_ld, - int64_t q_ld, - int64_t num_experts); +std::vector awq_moe_build_strided_ptrs( + torch::Tensor tm_weights, + torch::Tensor tm_scales, + int64_t k_ld, + int64_t q_ld, + int64_t num_experts); -void awq_moe_gemm_sm70_out(torch::Tensor out, torch::Tensor sorted_input, +void awq_moe_gemm_sm70_out(torch::Tensor out, + torch::Tensor sorted_input, torch::Tensor expert_offsets, torch::Tensor strided_ptrs_w, - torch::Tensor strided_ptrs_s, int64_t num_experts, - int64_t k, int64_t n, int64_t group_size, + torch::Tensor strided_ptrs_s, + int64_t num_experts, + int64_t k, + int64_t n, + int64_t group_size, bool gated_silu); -void awq_moe_gemm_sm70_per_expert_dispatch_out( - torch::Tensor out, torch::Tensor sorted_input, torch::Tensor expert_offsets, - torch::Tensor strided_ptrs_w, torch::Tensor strided_ptrs_s, - int64_t num_experts, int64_t k, int64_t n, int64_t group_size, - bool gated_silu); - -void awq_moe_dense_stage_sm70_out(torch::Tensor out, torch::Tensor input, +void awq_moe_gemm_sm70_per_expert_dispatch_out(torch::Tensor out, + torch::Tensor sorted_input, + torch::Tensor expert_offsets, + torch::Tensor strided_ptrs_w, + torch::Tensor strided_ptrs_s, + int64_t num_experts, + int64_t k, + int64_t n, + int64_t group_size, + bool gated_silu); + +void awq_moe_dense_stage_sm70_out(torch::Tensor out, + torch::Tensor input, torch::Tensor expert_offsets, torch::Tensor dense_expert_ids, - torch::Tensor ptrs_w, torch::Tensor ptrs_s, - int64_t num_experts, int64_t k, int64_t n, + torch::Tensor ptrs_w, + torch::Tensor ptrs_s, + int64_t num_experts, + int64_t k, + int64_t n, int64_t group_size); void awq_moe_active_dense_stage_sm70_out( - torch::Tensor out, torch::Tensor input, torch::Tensor permuted_experts_id, - torch::Tensor active_expert_offsets, torch::Tensor active_expert_ids, - torch::Tensor ptrs_w, torch::Tensor ptrs_s, int64_t total_slots, int64_t k, - int64_t n, int64_t group_size); + torch::Tensor out, + torch::Tensor input, + torch::Tensor permuted_experts_id, + torch::Tensor active_expert_offsets, + torch::Tensor active_expert_ids, + torch::Tensor ptrs_w, + torch::Tensor ptrs_s, + int64_t total_slots, + int64_t k, + int64_t n, + int64_t group_size); void awq_moe_single_token_dense_stage_sm70_out( - torch::Tensor out, torch::Tensor input, torch::Tensor expert_offsets, - torch::Tensor sorted_expert_ids, torch::Tensor ptrs_w, torch::Tensor ptrs_s, - int64_t top_k, int64_t k, int64_t n, int64_t group_size); + torch::Tensor out, + torch::Tensor input, + torch::Tensor expert_offsets, + torch::Tensor sorted_expert_ids, + torch::Tensor ptrs_w, + torch::Tensor ptrs_s, + int64_t top_k, + int64_t k, + int64_t n, + int64_t group_size); void awq_moe_single_token_indexed_dense_stage_sm70_out( - torch::Tensor out, torch::Tensor input, torch::Tensor expert_offsets, - torch::Tensor sorted_expert_ids, torch::Tensor ptrs_w, torch::Tensor ptrs_s, - int64_t top_k, int64_t k, int64_t n, int64_t group_size); + torch::Tensor out, + torch::Tensor input, + torch::Tensor expert_offsets, + torch::Tensor sorted_expert_ids, + torch::Tensor ptrs_w, + torch::Tensor ptrs_s, + int64_t top_k, + int64_t k, + int64_t n, + int64_t group_size); void awq_moe_single_token_dense_w13_sm70_out( - torch::Tensor gate_up, torch::Tensor compact_input, torch::Tensor x, - torch::Tensor topk_ids, torch::Tensor w13_ptrs_w, torch::Tensor w13_ptrs_s, - torch::Tensor expert_offsets, torch::Tensor expert_offsets64, - torch::Tensor inv_permuted_idx, torch::Tensor sorted_expert_ids, - int64_t w13_k, int64_t w13_n, int64_t group_size, + torch::Tensor gate_up, + torch::Tensor compact_input, + torch::Tensor x, + torch::Tensor topk_ids, + torch::Tensor w13_ptrs_w, + torch::Tensor w13_ptrs_s, + torch::Tensor expert_offsets, + torch::Tensor expert_offsets64, + torch::Tensor inv_permuted_idx, + torch::Tensor sorted_expert_ids, + int64_t w13_k, + int64_t w13_n, + int64_t group_size, int64_t hidden_logical_size); void awq_moe_single_token_indexed_dense_w13_sm70_out( - torch::Tensor gate_up, torch::Tensor compact_input, torch::Tensor x, - torch::Tensor topk_ids, torch::Tensor w13_ptrs_w, torch::Tensor w13_ptrs_s, - torch::Tensor expert_offsets, torch::Tensor expert_offsets64, - torch::Tensor inv_permuted_idx, torch::Tensor sorted_expert_ids, - int64_t w13_k, int64_t w13_n, int64_t group_size, + torch::Tensor gate_up, + torch::Tensor compact_input, + torch::Tensor x, + torch::Tensor topk_ids, + torch::Tensor w13_ptrs_w, + torch::Tensor w13_ptrs_s, + torch::Tensor expert_offsets, + torch::Tensor expert_offsets64, + torch::Tensor inv_permuted_idx, + torch::Tensor sorted_expert_ids, + int64_t w13_k, + int64_t w13_n, + int64_t group_size, int64_t hidden_logical_size); void awq_moe_single_token_compact_dense_w13_sm70_out( - torch::Tensor gate_up, torch::Tensor compact_input, torch::Tensor x, - torch::Tensor topk_ids, torch::Tensor w13_ptrs_w, torch::Tensor w13_ptrs_s, - torch::Tensor compact_w13_ptrs_w, torch::Tensor compact_w13_ptrs_s, - torch::Tensor expert_offsets, torch::Tensor expert_offsets64, - torch::Tensor inv_permuted_idx, torch::Tensor sorted_expert_ids, - int64_t w13_k, int64_t w13_n, int64_t group_size, + torch::Tensor gate_up, + torch::Tensor compact_input, + torch::Tensor x, + torch::Tensor topk_ids, + torch::Tensor w13_ptrs_w, + torch::Tensor w13_ptrs_s, + torch::Tensor compact_w13_ptrs_w, + torch::Tensor compact_w13_ptrs_s, + torch::Tensor expert_offsets, + torch::Tensor expert_offsets64, + torch::Tensor inv_permuted_idx, + torch::Tensor sorted_expert_ids, + int64_t w13_k, + int64_t w13_n, + int64_t group_size, int64_t hidden_logical_size); -void awq_moe_single_token_exact_layout_prepare( - torch::Tensor topk_ids, torch::Tensor x, torch::Tensor compact_input, - torch::Tensor expert_offsets, torch::Tensor expert_offsets64, - torch::Tensor inv_permuted_idx, int64_t num_experts); +void awq_moe_single_token_exact_layout_prepare(torch::Tensor topk_ids, + torch::Tensor x, + torch::Tensor compact_input, + torch::Tensor expert_offsets, + torch::Tensor expert_offsets64, + torch::Tensor inv_permuted_idx, + int64_t num_experts); void awq_moe_single_token_weighted_reduce_out(torch::Tensor sorted_output, torch::Tensor topk_weights, torch::Tensor inv_permuted_idx, - torch::Tensor out, int64_t top_k, + torch::Tensor out, + int64_t top_k, int64_t hidden_logical_size); void awq_moe_single_token_sm70_out( - torch::Tensor out, torch::Tensor x, torch::Tensor topk_weights, - torch::Tensor topk_ids, torch::Tensor src_w13_ptrs_w_rows, - torch::Tensor src_w13_ptrs_s_rows, torch::Tensor src_w2_ptrs_w_rows, - torch::Tensor src_w2_ptrs_s_rows, torch::Tensor compact_input, - torch::Tensor intermediate, torch::Tensor sorted_output, - torch::Tensor sorted_weights, torch::Tensor dst_w13_ptrs_w_rows, - torch::Tensor dst_w13_ptrs_s_rows, torch::Tensor dst_w2_ptrs_w_rows, - torch::Tensor dst_w2_ptrs_s_rows, torch::Tensor expert_offsets, - torch::Tensor inv_permuted_idx, int64_t w13_k, int64_t w13_n, int64_t w2_k, - int64_t w2_n, int64_t group_size, int64_t hidden_logical_size); - -void fp8_moe_gemm_sm70_out(torch::Tensor out, torch::Tensor sorted_input, + torch::Tensor out, + torch::Tensor x, + torch::Tensor topk_weights, + torch::Tensor topk_ids, + torch::Tensor src_w13_ptrs_w_rows, + torch::Tensor src_w13_ptrs_s_rows, + torch::Tensor src_w2_ptrs_w_rows, + torch::Tensor src_w2_ptrs_s_rows, + torch::Tensor compact_input, + torch::Tensor intermediate, + torch::Tensor sorted_output, + torch::Tensor sorted_weights, + torch::Tensor dst_w13_ptrs_w_rows, + torch::Tensor dst_w13_ptrs_s_rows, + torch::Tensor dst_w2_ptrs_w_rows, + torch::Tensor dst_w2_ptrs_s_rows, + torch::Tensor expert_offsets, + torch::Tensor inv_permuted_idx, + int64_t w13_k, + int64_t w13_n, + int64_t w2_k, + int64_t w2_n, + int64_t group_size, + int64_t hidden_logical_size); + +void fp8_moe_gemm_sm70_out(torch::Tensor out, + torch::Tensor sorted_input, torch::Tensor expert_offsets, torch::Tensor strided_ptrs_w, - torch::Tensor strided_ptrs_s, int64_t num_experts, - int64_t k, int64_t n, int64_t group_size, + torch::Tensor strided_ptrs_s, + int64_t num_experts, + int64_t k, + int64_t n, + int64_t group_size, bool gated_silu); -void fp8_moe_gemm_sm70_per_expert_dispatch_out( - torch::Tensor out, torch::Tensor sorted_input, torch::Tensor expert_offsets, - torch::Tensor strided_ptrs_w, torch::Tensor strided_ptrs_s, - int64_t num_experts, int64_t k, int64_t n, int64_t group_size, - bool gated_silu); - -void fp8_moe_dense_stage_sm70_out(torch::Tensor out, torch::Tensor input, +void fp8_moe_gemm_sm70_per_expert_dispatch_out(torch::Tensor out, + torch::Tensor sorted_input, + torch::Tensor expert_offsets, + torch::Tensor strided_ptrs_w, + torch::Tensor strided_ptrs_s, + int64_t num_experts, + int64_t k, + int64_t n, + int64_t group_size, + bool gated_silu); + +void fp8_moe_dense_stage_sm70_out(torch::Tensor out, + torch::Tensor input, torch::Tensor expert_offsets, torch::Tensor dense_expert_ids, - torch::Tensor ptrs_w, torch::Tensor ptrs_s, - int64_t num_experts, int64_t k, int64_t n, + torch::Tensor ptrs_w, + torch::Tensor ptrs_s, + int64_t num_experts, + int64_t k, + int64_t n, int64_t group_size); -void mxfp4_moe_dense_stage_sm70_out(torch::Tensor out, torch::Tensor input, - torch::Tensor expert_offsets, - torch::Tensor dense_expert_ids, - torch::Tensor ptrs_w, torch::Tensor ptrs_s, - int64_t num_experts, int64_t k, int64_t n, - int64_t group_size); - -void nvfp4_moe_dense_stage_sm70_out(torch::Tensor out, torch::Tensor input, +void mxfp4_moe_dense_stage_sm70_out(torch::Tensor out, + torch::Tensor input, torch::Tensor expert_offsets, torch::Tensor dense_expert_ids, - torch::Tensor ptrs_w, torch::Tensor ptrs_s, - int64_t num_experts, int64_t k, int64_t n, + torch::Tensor ptrs_w, + torch::Tensor ptrs_s, + int64_t num_experts, + int64_t k, + int64_t n, int64_t group_size); void mxfp4_moe_single_token_prepare_w13_sm70_out( - torch::Tensor gate_up, torch::Tensor compact_input, torch::Tensor x, - torch::Tensor topk_ids, torch::Tensor w13_ptrs_w, torch::Tensor w13_ptrs_s, - torch::Tensor expert_offsets, torch::Tensor inv_permuted_idx, - torch::Tensor sorted_expert_ids, int64_t w13_k, int64_t w13_n, - int64_t group_size, int64_t hidden_logical_size); + torch::Tensor gate_up, + torch::Tensor compact_input, + torch::Tensor x, + torch::Tensor topk_ids, + torch::Tensor w13_ptrs_w, + torch::Tensor w13_ptrs_s, + torch::Tensor expert_offsets, + torch::Tensor inv_permuted_idx, + torch::Tensor sorted_expert_ids, + int64_t w13_k, + int64_t w13_n, + int64_t group_size, + int64_t hidden_logical_size); void fp8_moe_single_token_dense_stage_sm70_out( - torch::Tensor out, torch::Tensor input, torch::Tensor expert_offsets, - torch::Tensor sorted_expert_ids, torch::Tensor ptrs_w, torch::Tensor ptrs_s, - int64_t top_k, int64_t k, int64_t n, int64_t group_size); + torch::Tensor out, + torch::Tensor input, + torch::Tensor expert_offsets, + torch::Tensor sorted_expert_ids, + torch::Tensor ptrs_w, + torch::Tensor ptrs_s, + int64_t top_k, + int64_t k, + int64_t n, + int64_t group_size); void fp8_moe_single_token_indexed_dense_stage_sm70_out( - torch::Tensor out, torch::Tensor input, torch::Tensor expert_offsets, - torch::Tensor sorted_expert_ids, torch::Tensor ptrs_w, torch::Tensor ptrs_s, - int64_t top_k, int64_t k, int64_t n, int64_t group_size); + torch::Tensor out, + torch::Tensor input, + torch::Tensor expert_offsets, + torch::Tensor sorted_expert_ids, + torch::Tensor ptrs_w, + torch::Tensor ptrs_s, + int64_t top_k, + int64_t k, + int64_t n, + int64_t group_size); void fp8_moe_single_token_dense_w13_sm70_out( - torch::Tensor gate_up, torch::Tensor compact_input, torch::Tensor x, - torch::Tensor topk_ids, torch::Tensor w13_ptrs_w, torch::Tensor w13_ptrs_s, - torch::Tensor expert_offsets, torch::Tensor expert_offsets64, - torch::Tensor inv_permuted_idx, torch::Tensor sorted_expert_ids, - int64_t w13_k, int64_t w13_n, int64_t group_size, + torch::Tensor gate_up, + torch::Tensor compact_input, + torch::Tensor x, + torch::Tensor topk_ids, + torch::Tensor w13_ptrs_w, + torch::Tensor w13_ptrs_s, + torch::Tensor expert_offsets, + torch::Tensor expert_offsets64, + torch::Tensor inv_permuted_idx, + torch::Tensor sorted_expert_ids, + int64_t w13_k, + int64_t w13_n, + int64_t group_size, int64_t hidden_logical_size); void fp8_moe_single_token_indexed_dense_w13_sm70_out( - torch::Tensor gate_up, torch::Tensor compact_input, torch::Tensor x, - torch::Tensor topk_ids, torch::Tensor w13_ptrs_w, torch::Tensor w13_ptrs_s, - torch::Tensor expert_offsets, torch::Tensor expert_offsets64, - torch::Tensor inv_permuted_idx, torch::Tensor sorted_expert_ids, - int64_t w13_k, int64_t w13_n, int64_t group_size, + torch::Tensor gate_up, + torch::Tensor compact_input, + torch::Tensor x, + torch::Tensor topk_ids, + torch::Tensor w13_ptrs_w, + torch::Tensor w13_ptrs_s, + torch::Tensor expert_offsets, + torch::Tensor expert_offsets64, + torch::Tensor inv_permuted_idx, + torch::Tensor sorted_expert_ids, + int64_t w13_k, + int64_t w13_n, + int64_t group_size, int64_t hidden_logical_size); void fp8_moe_single_token_compact_dense_w13_sm70_out( - torch::Tensor gate_up, torch::Tensor compact_input, torch::Tensor x, - torch::Tensor topk_ids, torch::Tensor w13_ptrs_w, torch::Tensor w13_ptrs_s, - torch::Tensor compact_w13_ptrs_w, torch::Tensor compact_w13_ptrs_s, - torch::Tensor expert_offsets, torch::Tensor expert_offsets64, - torch::Tensor inv_permuted_idx, torch::Tensor sorted_expert_ids, - int64_t w13_k, int64_t w13_n, int64_t group_size, + torch::Tensor gate_up, + torch::Tensor compact_input, + torch::Tensor x, + torch::Tensor topk_ids, + torch::Tensor w13_ptrs_w, + torch::Tensor w13_ptrs_s, + torch::Tensor compact_w13_ptrs_w, + torch::Tensor compact_w13_ptrs_s, + torch::Tensor expert_offsets, + torch::Tensor expert_offsets64, + torch::Tensor inv_permuted_idx, + torch::Tensor sorted_expert_ids, + int64_t w13_k, + int64_t w13_n, + int64_t group_size, int64_t hidden_logical_size); void fp8_moe_single_token_sm70_out( - torch::Tensor out, torch::Tensor x, torch::Tensor topk_weights, - torch::Tensor topk_ids, torch::Tensor src_w13_ptrs_w_rows, - torch::Tensor src_w13_ptrs_s_rows, torch::Tensor src_w2_ptrs_w_rows, - torch::Tensor src_w2_ptrs_s_rows, torch::Tensor compact_input, - torch::Tensor gate_up, torch::Tensor intermediate, - torch::Tensor sorted_output, torch::Tensor sorted_weights, - torch::Tensor dst_w13_ptrs_w_rows, torch::Tensor dst_w13_ptrs_s_rows, - torch::Tensor dst_w2_ptrs_w_rows, torch::Tensor dst_w2_ptrs_s_rows, - torch::Tensor expert_offsets, torch::Tensor inv_permuted_idx, - torch::Tensor sorted_expert_ids, torch::Tensor broadcast_input_indices, - torch::Tensor w2_raw_weight, torch::Tensor w2_raw_scale_inv, int64_t w13_k, - int64_t w13_n, int64_t w2_k, int64_t w2_n, int64_t group_size, - int64_t hidden_logical_size, bool fused_gated_silu, - bool fused_weighted_reduce, bool broadcast_input, bool w2_direct_reduce, - bool indexed_expert_ptrs, bool exact_per_route); + torch::Tensor out, + torch::Tensor x, + torch::Tensor topk_weights, + torch::Tensor topk_ids, + torch::Tensor src_w13_ptrs_w_rows, + torch::Tensor src_w13_ptrs_s_rows, + torch::Tensor src_w2_ptrs_w_rows, + torch::Tensor src_w2_ptrs_s_rows, + torch::Tensor compact_input, + torch::Tensor gate_up, + torch::Tensor intermediate, + torch::Tensor sorted_output, + torch::Tensor sorted_weights, + torch::Tensor dst_w13_ptrs_w_rows, + torch::Tensor dst_w13_ptrs_s_rows, + torch::Tensor dst_w2_ptrs_w_rows, + torch::Tensor dst_w2_ptrs_s_rows, + torch::Tensor expert_offsets, + torch::Tensor inv_permuted_idx, + torch::Tensor sorted_expert_ids, + torch::Tensor broadcast_input_indices, + torch::Tensor w2_raw_weight, + torch::Tensor w2_raw_scale_inv, + int64_t w13_k, + int64_t w13_n, + int64_t w2_k, + int64_t w2_n, + int64_t group_size, + int64_t hidden_logical_size, + bool fused_gated_silu, + bool fused_weighted_reduce, + bool broadcast_input, + bool w2_direct_reduce, + bool indexed_expert_ptrs, + bool exact_per_route); #endif void static_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input, @@ -575,30 +664,27 @@ void all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out, void sm70_tp2_all_reduce_gemma_rms_norm( fptr_t _fa, torch::Tensor& inp, torch::Tensor& residual, torch::Tensor& weight, torch::Tensor& normalized_out, - torch::Tensor& residual_out, fptr_t reg_buffer, int64_t reg_buffer_sz_bytes, - double epsilon); + torch::Tensor& residual_out, fptr_t reg_buffer, + int64_t reg_buffer_sz_bytes, double epsilon); void sm70_tp4_all_reduce_gemma_rms_norm( fptr_t _fa, torch::Tensor& inp, torch::Tensor& residual, torch::Tensor& weight, torch::Tensor& normalized_out, - torch::Tensor& residual_out, fptr_t reg_buffer, int64_t reg_buffer_sz_bytes, - double epsilon); -void sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather( - fptr_t _fa, torch::Tensor& inp, torch::Tensor& residual, - torch::Tensor& weight, torch::Tensor& normalized_out, - torch::Tensor& residual_out, fptr_t reg_input_buffer, - fptr_t reg_output_buffer, int64_t reg_buffer_sz_bytes, double epsilon); + torch::Tensor& residual_out, fptr_t reg_buffer, + int64_t reg_buffer_sz_bytes, double epsilon); void all_reduce_sum2(fptr_t _fa, torch::Tensor& inp_a, torch::Tensor& inp_b, torch::Tensor& out); void top1_argmax(fptr_t _fa, torch::Tensor& input_pair, torch::Tensor& output, fptr_t reg_buffer, int64_t reg_buffer_sz_bytes); void tile_runtime_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out, - fptr_t reg_buffer, int64_t reg_buffer_sz_bytes, + fptr_t reg_buffer, + int64_t reg_buffer_sz_bytes, int64_t tile_numel, int64_t engine_blocks, int64_t compute_iters); void tile_runtime_all_reduce_engine(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out, fptr_t reg_buffer, int64_t reg_buffer_sz_bytes, - int64_t tile_numel, int64_t producer_blocks, + int64_t tile_numel, + int64_t producer_blocks, int64_t reducer_blocks, int64_t compute_iters); void tile_runtime_wait_reduce(fptr_t _fa, torch::Tensor& staging, @@ -606,10 +692,7 @@ void tile_runtime_wait_reduce(fptr_t _fa, torch::Tensor& staging, int64_t reducer_blocks); void dispose(fptr_t _fa); int64_t meta_size(); -int64_t sm70_tp4_push_allreduce_buffer_size(); void register_buffer(fptr_t _fa, const std::vector& fake_ipc_ptrs); -void register_sm70_tp4_push_allreduce_buffer( - fptr_t _fa, const std::vector& fake_ipc_ptrs); std::tuple, std::vector> get_graph_buffer_ipc_meta(fptr_t _fa); void register_graph_buffers(fptr_t _fa, diff --git a/csrc/quantization/marlin/sm70_awq_marlin_repack.cu b/csrc/quantization/marlin/sm70_awq_marlin_repack.cu index b664b2b29f..ff37c46ff5 100644 --- a/csrc/quantization/marlin/sm70_awq_marlin_repack.cu +++ b/csrc/quantization/marlin/sm70_awq_marlin_repack.cu @@ -92,7 +92,8 @@ __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}; @@ -113,8 +114,9 @@ __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; } @@ -153,8 +155,9 @@ __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; } @@ -301,15 +304,16 @@ __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 089ea3220d..db89b2d5af 100644 --- a/csrc/quantization/marlin/sm70_gptq_marlin_repack.cu +++ b/csrc/quantization/marlin/sm70_gptq_marlin_repack.cu @@ -164,8 +164,9 @@ __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; } @@ -214,8 +215,9 @@ __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; } @@ -379,17 +381,19 @@ __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 3ae26c8fcf..0e8ecbf4a4 100644 --- a/csrc/quantization/marlin/sm70_marlin_common.cuh +++ b/csrc/quantization/marlin/sm70_marlin_common.cuh @@ -28,29 +28,31 @@ 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; @@ -63,7 +65,8 @@ 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) { @@ -71,7 +74,8 @@ 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; @@ -86,40 +90,43 @@ 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; @@ -176,8 +183,8 @@ constexpr bool sm70_marlin_volta_epilogue_supported(int cta_m, int cta_n, 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; @@ -185,11 +192,12 @@ 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); @@ -204,7 +212,8 @@ 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, @@ -274,7 +283,8 @@ 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); @@ -290,10 +300,8 @@ 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; } @@ -301,12 +309,14 @@ 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) { @@ -318,8 +328,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}; @@ -327,8 +337,10 @@ 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) { @@ -374,10 +386,9 @@ 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; } @@ -398,8 +409,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; @@ -422,9 +433,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); @@ -436,13 +447,14 @@ 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", @@ -462,7 +474,8 @@ 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, @@ -472,27 +485,31 @@ 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}; @@ -579,14 +596,16 @@ inline bool 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}; @@ -595,144 +614,147 @@ 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}; @@ -741,282 +763,293 @@ inline bool 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; @@ -1024,14 +1057,16 @@ inline bool 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}; @@ -1040,112 +1075,115 @@ 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}; @@ -1208,15 +1246,18 @@ inline bool 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}; @@ -1282,15 +1323,18 @@ inline bool 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}; @@ -1382,16 +1426,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; } @@ -1400,7 +1444,8 @@ 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; } @@ -1419,12 +1464,13 @@ 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; } @@ -1432,25 +1478,27 @@ 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 51aa145884..fe7fa0f408 100644 --- a/csrc/quantization/marlin/sm70_marlin_dispatch.cu +++ b/csrc/quantization/marlin/sm70_marlin_dispatch.cu @@ -43,13 +43,15 @@ __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]; } @@ -64,9 +66,10 @@ __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; @@ -78,18 +81,21 @@ __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."); @@ -97,14 +103,15 @@ 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<<>>( + sm70_unpermute_half_metadata_kernel<<< + blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>( reinterpret_cast(b_scales.data_ptr()), reinterpret_cast(out.data_ptr()), total, perm_len); C10_CUDA_KERNEL_LAUNCH_CHECK(); @@ -136,8 +143,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<<>>( + sm70_unpack_dense_zp_to_half_kernel<<< + blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>( b_zeros.data_ptr(), reinterpret_cast(b_scales.data_ptr()), reinterpret_cast(out.data_ptr()), total, size_n, @@ -195,7 +202,8 @@ 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) { @@ -265,7 +273,8 @@ 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."); @@ -316,8 +325,10 @@ 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."); @@ -410,18 +421,21 @@ 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, @@ -449,7 +463,8 @@ 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); } @@ -462,9 +477,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; } } @@ -478,7 +493,8 @@ 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), @@ -507,26 +523,30 @@ 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); } @@ -534,15 +554,17 @@ 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."); @@ -551,8 +573,9 @@ 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 ", @@ -565,15 +588,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 4dcc8480ad..9520f9546a 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_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_active_split_k; 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,11 +55,12 @@ class Sm70Fp8IteratorB { static int const kPackedMacroN = PackedMacroN_; static constexpr bool kUseMetadataVectorWords = UseMetadataVectorWords_; using Element = cutlass::half_t; - using Fragment = cutlass::Array; + using Fragment = cutlass::Array< + Element, ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; 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, @@ -160,12 +161,13 @@ 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(); } @@ -196,8 +198,9 @@ 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 { @@ -258,9 +261,8 @@ 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>( @@ -282,17 +284,19 @@ 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; @@ -314,9 +318,10 @@ 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; @@ -338,9 +343,8 @@ 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>( @@ -370,7 +374,8 @@ 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)); } @@ -382,29 +387,32 @@ 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; @@ -424,10 +432,10 @@ __global__ __launch_bounds__(Warps * 32, 1) 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), @@ -452,19 +460,18 @@ __global__ __launch_bounds__(Warps * 32, 1) 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; @@ -475,13 +482,14 @@ __launch_bounds__(Warps * 32, 1) 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}; @@ -489,10 +497,10 @@ __launch_bounds__(Warps * 32, 1) 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), @@ -506,8 +514,9 @@ __launch_bounds__(Warps * 32, 1) 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); @@ -516,23 +525,20 @@ __launch_bounds__(Warps * 32, 1) 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::SharedStorage; + using SharedStorage = typename Sm70Fp8GemmTraits< + CtaM, CtaN, CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, + PackedMacroN, UseMetadataVectorWords>::SharedStorage; size_t smem_bytes = configure_sm70_dynamic_smem(kernel); dim3 block(Warps * 32); @@ -552,13 +558,14 @@ torch::Tensor launch_sm70_marlin_fp8_gemm(torch::Tensor& a, torch::Tensor& c, } 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); @@ -566,7 +573,8 @@ torch::Tensor launch_sm70_marlin_fp8_gemm(torch::Tensor& a, torch::Tensor& c, // 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 = @@ -577,9 +585,9 @@ torch::Tensor launch_sm70_marlin_fp8_gemm(torch::Tensor& a, torch::Tensor& c, 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. @@ -587,8 +595,8 @@ torch::Tensor launch_sm70_marlin_fp8_gemm(torch::Tensor& a, torch::Tensor& c, 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; @@ -605,8 +613,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 a968a54db6..426fac1c0f 100644 --- a/csrc/quantization/marlin/sm70_marlin_gemm.cuh +++ b/csrc/quantization/marlin/sm70_marlin_gemm.cuh @@ -30,7 +30,8 @@ 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."); @@ -44,8 +45,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>; @@ -59,10 +60,9 @@ struct Sm70MarlinGemmTraits { cutlass::MatrixShape, ElementA, LayoutA, 1, typename MmaCore::IteratorThreadMapA, 128 / cutlass::sizeof_bits::value>; - using IteratorB = - typename Spec::template IteratorB; + using IteratorB = typename Spec::template IteratorB< + ThreadblockShape, typename MmaCore::IteratorThreadMapB, GroupSize, + PackedMacroN>; 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,16 +152,15 @@ 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(); } @@ -176,57 +175,58 @@ 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 e5ff7e3782..1e134c6c68 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< - typename QwordVector::Type const*>(ptr); + return *reinterpret_cast::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 ded9b11707..7d3d4421bb 100644 --- a/csrc/quantization/marlin/sm70_marlin_mma.cuh +++ b/csrc/quantization/marlin/sm70_marlin_mma.cuh @@ -7,16 +7,17 @@ namespace marlin::sm70 { -template , - typename TransformB_ = cutlass::NumericArrayConverter< - typename SmemIteratorB_::Element, typename IteratorB_::Element, - IteratorB_::Fragment::kElements>, - typename Enable = bool> +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> class Sm70MarlinMmaPipelined : public cutlass::gemm::threadblock::MmaBase { public: @@ -105,10 +106,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), @@ -153,12 +154,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; @@ -221,10 +222,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 8235c2b17c..8a22142725 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_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_active_split_k; 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,8 +72,7 @@ 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_; @@ -81,11 +80,12 @@ class Sm70Mxfp4IteratorB { static int const kGroupSize = GroupSize_; static int const kPackedMacroN = PackedMacroN_; using Element = cutlass::half_t; - using Fragment = cutlass::Array; + using Fragment = cutlass::Array< + Element, ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; 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,8 +94,7 @@ 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; @@ -177,8 +176,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 @@ -197,12 +196,14 @@ 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); } @@ -255,11 +256,9 @@ 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]; @@ -277,12 +276,14 @@ 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; @@ -303,7 +304,8 @@ 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; @@ -324,11 +326,9 @@ 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,25 +361,26 @@ 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; @@ -399,10 +400,10 @@ __global__ __launch_bounds__(Warps * 32, 1) 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), @@ -427,17 +428,16 @@ __global__ __launch_bounds__(Warps * 32, 1) 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,13 +448,14 @@ __launch_bounds__(Warps * 32, 1) 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}; @@ -462,10 +463,10 @@ __launch_bounds__(Warps * 32, 1) 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), @@ -479,8 +480,9 @@ __launch_bounds__(Warps * 32, 1) 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); @@ -489,20 +491,18 @@ __launch_bounds__(Warps * 32, 1) 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::SharedStorage; + sm70_marlin_mxfp4_gemm_kernel; + using SharedStorage = typename Sm70Mxfp4GemmTraits< + CtaM, CtaN, CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, + PackedMacroN>::SharedStorage; size_t smem_bytes = configure_sm70_dynamic_smem(kernel); dim3 block(Warps * 32); @@ -522,19 +522,21 @@ torch::Tensor launch_sm70_marlin_mxfp4_gemm(torch::Tensor& a, torch::Tensor& c, } 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< - CtaM, CtaN, CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, PackedMacroN>; + auto split_kernel = + sm70_marlin_mxfp4_gemm_splitk_kernel; 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 = @@ -545,9 +547,9 @@ torch::Tensor launch_sm70_marlin_mxfp4_gemm(torch::Tensor& a, torch::Tensor& c, 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. @@ -555,8 +557,8 @@ torch::Tensor launch_sm70_marlin_mxfp4_gemm(torch::Tensor& a, torch::Tensor& c, 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; @@ -572,11 +574,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); } }; @@ -588,16 +590,14 @@ 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 debb47d16c..bd33da2e22 100644 --- a/csrc/quantization/marlin/sm70_marlin_nvfp4_gemm.cu +++ b/csrc/quantization/marlin/sm70_marlin_nvfp4_gemm.cu @@ -21,32 +21,31 @@ #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_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_active_split_k; 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_; @@ -54,11 +53,12 @@ class Sm70Nvfp4IteratorB { static int const kGroupSize = GroupSize_; static int const kPackedMacroN = PackedMacroN_; using Element = cutlass::half_t; - using Fragment = cutlass::Array; + using Fragment = cutlass::Array< + Element, ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; 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,8 +67,7 @@ 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; @@ -150,8 +149,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 @@ -171,12 +170,14 @@ 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); } @@ -229,11 +230,9 @@ 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]; @@ -251,12 +250,14 @@ 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; @@ -279,7 +280,8 @@ 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; @@ -302,11 +304,9 @@ 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,26 +340,27 @@ 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; @@ -379,10 +380,10 @@ __global__ __launch_bounds__(Warps * 32, 1) 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), @@ -407,18 +408,17 @@ __global__ __launch_bounds__(Warps * 32, 1) 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,13 +429,14 @@ __launch_bounds__(Warps * 32, 1) 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}; @@ -443,10 +444,10 @@ __launch_bounds__(Warps * 32, 1) 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), @@ -462,8 +463,9 @@ __launch_bounds__(Warps * 32, 1) 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); @@ -472,18 +474,18 @@ __launch_bounds__(Warps * 32, 1) 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::SharedStorage; + sm70_marlin_nvfp4_gemm_kernel; + using SharedStorage = typename Sm70Nvfp4GemmTraits< + CtaM, CtaN, CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, + PackedMacroN>::SharedStorage; size_t smem_bytes = configure_sm70_dynamic_smem(kernel); dim3 block(Warps * 32); @@ -504,19 +506,21 @@ 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< - CtaM, CtaN, CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, PackedMacroN>; + auto split_kernel = + sm70_marlin_nvfp4_gemm_splitk_kernel; 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 = @@ -528,9 +532,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. @@ -538,8 +542,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; @@ -556,11 +560,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); } @@ -574,22 +578,14 @@ 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 463f08b84f..a0fb01910a 100644 --- a/csrc/quantization/marlin/sm70_marlin_splitk.cuh +++ b/csrc/quantization/marlin/sm70_marlin_splitk.cuh @@ -22,7 +22,9 @@ 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() { @@ -74,8 +76,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}; @@ -133,9 +135,10 @@ 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 @@ -165,7 +168,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; @@ -174,7 +177,8 @@ 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); } @@ -216,7 +220,8 @@ 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; } } @@ -258,9 +263,10 @@ 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 @@ -289,7 +295,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; @@ -298,14 +304,15 @@ 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 @@ -340,16 +347,17 @@ 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 e2038fae64..0037ae8187 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_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_active_split_k; 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,11 +54,12 @@ class Sm70U4ZpIteratorB { static int const kPackedMacroN = PackedMacroN_; static constexpr bool kUseMetadataVectorWords = UseMetadataVectorWords_; using Element = cutlass::half_t; - using Fragment = cutlass::Array; + using Fragment = cutlass::Array< + Element, ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; 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, @@ -67,8 +68,7 @@ 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,7 +151,8 @@ 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; @@ -161,8 +162,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 @@ -175,8 +176,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]; @@ -195,8 +196,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]; @@ -209,8 +210,9 @@ 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 { @@ -235,8 +237,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( @@ -257,8 +259,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( @@ -269,15 +271,14 @@ 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( @@ -289,15 +290,17 @@ 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; @@ -311,15 +314,20 @@ 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; @@ -333,19 +341,21 @@ 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_; @@ -353,8 +363,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( @@ -372,7 +382,8 @@ 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)); } @@ -384,30 +395,32 @@ 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; @@ -427,10 +440,10 @@ __global__ __launch_bounds__(Warps * 32, 1) 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), @@ -456,20 +469,19 @@ __global__ __launch_bounds__(Warps * 32, 1) 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; @@ -480,13 +492,14 @@ __launch_bounds__(Warps * 32, 1) 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}; @@ -494,10 +507,10 @@ __launch_bounds__(Warps * 32, 1) 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), @@ -512,8 +525,9 @@ __launch_bounds__(Warps * 32, 1) 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); @@ -522,23 +536,19 @@ __launch_bounds__(Warps * 32, 1) 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::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< + CtaM, CtaN, CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, + PackedMacroN, UseMetadataVectorWords>::SharedStorage; size_t smem_bytes = configure_sm70_dynamic_smem(kernel); dim3 block(Warps * 32); cudaStream_t stream = at::cuda::getCurrentCUDAStream(a.get_device()).stream(); @@ -548,7 +558,8 @@ torch::Tensor launch_sm70_marlin_u4_gemm(torch::Tensor& a, torch::Tensor& c, 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), @@ -558,8 +569,8 @@ torch::Tensor launch_sm70_marlin_u4_gemm(torch::Tensor& a, torch::Tensor& c, } 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 = @@ -572,7 +583,8 @@ torch::Tensor launch_sm70_marlin_u4_gemm(torch::Tensor& a, torch::Tensor& c, // 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 = @@ -584,9 +596,9 @@ torch::Tensor launch_sm70_marlin_u4_gemm(torch::Tensor& a, torch::Tensor& c, 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. @@ -594,8 +606,8 @@ torch::Tensor launch_sm70_marlin_u4_gemm(torch::Tensor& a, torch::Tensor& c, 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,44 +625,38 @@ 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 5de4a3a97e..cf96b3967a 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_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_active_split_k; 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,11 +54,12 @@ class Sm70U4B8IteratorB { static int const kPackedMacroN = PackedMacroN_; static constexpr bool kUseMetadataVectorWords = UseMetadataVectorWords_; using Element = cutlass::half_t; - using Fragment = cutlass::Array; + using Fragment = cutlass::Array< + Element, ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; 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, @@ -67,8 +68,7 @@ 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,7 +148,8 @@ 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; @@ -158,8 +159,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 @@ -189,8 +190,9 @@ 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 { @@ -245,17 +247,15 @@ 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,15 +267,17 @@ 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; @@ -296,7 +298,8 @@ 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; @@ -317,11 +320,9 @@ 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_; @@ -347,7 +348,8 @@ 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)); } @@ -359,29 +361,31 @@ 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; @@ -401,10 +405,10 @@ __global__ __launch_bounds__(Warps * 32, 1) 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), @@ -429,19 +433,18 @@ __global__ __launch_bounds__(Warps * 32, 1) 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; @@ -452,13 +455,14 @@ __launch_bounds__(Warps * 32, 1) 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}; @@ -466,10 +470,10 @@ __launch_bounds__(Warps * 32, 1) 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), @@ -483,8 +487,9 @@ __launch_bounds__(Warps * 32, 1) 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); @@ -493,23 +498,20 @@ __launch_bounds__(Warps * 32, 1) 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::SharedStorage; + using SharedStorage = typename Sm70U4B8GemmTraits< + CtaM, CtaN, CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, + PackedMacroN, UseMetadataVectorWords>::SharedStorage; size_t smem_bytes = configure_sm70_dynamic_smem(kernel); dim3 block(Warps * 32); @@ -529,13 +531,14 @@ torch::Tensor launch_sm70_marlin_u4b8_gemm(torch::Tensor& a, torch::Tensor& c, } 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); @@ -543,7 +546,8 @@ torch::Tensor launch_sm70_marlin_u4b8_gemm(torch::Tensor& a, torch::Tensor& c, // 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 = @@ -554,9 +558,9 @@ torch::Tensor launch_sm70_marlin_u4b8_gemm(torch::Tensor& a, torch::Tensor& c, 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. @@ -564,8 +568,8 @@ torch::Tensor launch_sm70_marlin_u4b8_gemm(torch::Tensor& a, torch::Tensor& c, 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; @@ -582,8 +586,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 f0cc420ed8..543325da82 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_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_active_split_k; 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,11 +55,12 @@ class Sm70U8ZpIteratorB { static int const kPackedMacroN = PackedMacroN_; static constexpr bool kUseMetadataVectorWords = UseMetadataVectorWords_; using Element = cutlass::half_t; - using Fragment = cutlass::Array; + using Fragment = cutlass::Array< + Element, ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; 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, @@ -118,7 +119,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); } } @@ -154,7 +155,8 @@ 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; @@ -164,8 +166,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 @@ -183,8 +185,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]; @@ -203,8 +205,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]; @@ -217,8 +219,9 @@ 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 { @@ -292,12 +295,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])); } @@ -305,17 +308,19 @@ 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; @@ -330,17 +335,22 @@ 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; @@ -355,17 +365,20 @@ 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>( @@ -396,7 +409,8 @@ 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)); } @@ -408,30 +422,32 @@ 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; @@ -451,10 +467,10 @@ __global__ __launch_bounds__(Warps * 32, 1) 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), @@ -480,20 +496,19 @@ __global__ __launch_bounds__(Warps * 32, 1) 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; @@ -504,13 +519,14 @@ __launch_bounds__(Warps * 32, 1) 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}; @@ -518,10 +534,10 @@ __launch_bounds__(Warps * 32, 1) 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), @@ -536,8 +552,9 @@ __launch_bounds__(Warps * 32, 1) 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); @@ -546,23 +563,20 @@ __launch_bounds__(Warps * 32, 1) 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::SharedStorage; + using SharedStorage = typename Sm70U8ZpGemmTraits< + CtaM, CtaN, CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, + PackedMacroN, UseMetadataVectorWords>::SharedStorage; size_t smem_bytes = configure_sm70_dynamic_smem(kernel); dim3 block(Warps * 32); @@ -583,8 +597,8 @@ torch::Tensor launch_sm70_marlin_u8_gemm(torch::Tensor& a, torch::Tensor& c, } 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 = @@ -597,7 +611,8 @@ torch::Tensor launch_sm70_marlin_u8_gemm(torch::Tensor& a, torch::Tensor& c, // 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 = @@ -609,9 +624,9 @@ torch::Tensor launch_sm70_marlin_u8_gemm(torch::Tensor& a, torch::Tensor& c, 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. @@ -619,8 +634,8 @@ torch::Tensor launch_sm70_marlin_u8_gemm(torch::Tensor& a, torch::Tensor& c, 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; @@ -638,14 +653,13 @@ 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); } }; @@ -657,25 +671,22 @@ 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 e4b8316c08..e8c312437d 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_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_active_split_k; 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; - 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< + 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, "SM70 Marlin U8B128 IteratorB expects one contiguous iteration " "per 64-column quant tile."); static_assert(ThreadMap::Delta::kContiguous == kQuantTileN, @@ -152,7 +152,8 @@ 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; @@ -162,12 +163,13 @@ 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(); } @@ -198,8 +200,9 @@ 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 { @@ -260,9 +263,8 @@ 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>( @@ -284,17 +286,19 @@ 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; @@ -316,9 +320,10 @@ 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; @@ -340,9 +345,8 @@ 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>( @@ -372,7 +376,8 @@ 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)); } @@ -384,29 +389,32 @@ 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; @@ -426,10 +434,10 @@ __global__ __launch_bounds__(Warps * 32, 1) 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), @@ -454,19 +462,18 @@ __global__ __launch_bounds__(Warps * 32, 1) 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; @@ -477,13 +484,14 @@ __launch_bounds__(Warps * 32, 1) 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}; @@ -491,10 +499,10 @@ __launch_bounds__(Warps * 32, 1) 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), @@ -508,8 +516,9 @@ __launch_bounds__(Warps * 32, 1) 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); @@ -518,23 +527,20 @@ __launch_bounds__(Warps * 32, 1) 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::SharedStorage; + using SharedStorage = typename Sm70U8B128GemmTraits< + CtaM, CtaN, CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, + PackedMacroN, UseMetadataVectorWords>::SharedStorage; size_t smem_bytes = configure_sm70_dynamic_smem(kernel); dim3 block(Warps * 32); @@ -554,20 +560,23 @@ torch::Tensor launch_sm70_marlin_u8b128_gemm(torch::Tensor& a, torch::Tensor& c, } 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< - CtaM, CtaN, CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, PackedMacroN, - UseMetadataVectorWords>; + auto split_kernel = + sm70_marlin_u8b128_gemm_splitk_kernel; 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 = @@ -578,9 +587,9 @@ torch::Tensor launch_sm70_marlin_u8b128_gemm(torch::Tensor& a, torch::Tensor& c, 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. @@ -588,8 +597,8 @@ torch::Tensor launch_sm70_marlin_u8b128_gemm(torch::Tensor& a, torch::Tensor& c, 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; @@ -606,11 +615,12 @@ 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); } @@ -623,23 +633,21 @@ 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_tile_runtime_signal.cuh b/csrc/sm70_tile_runtime_signal.cuh index 061a8147c2..401b9f99eb 100644 --- a/csrc/sm70_tile_runtime_signal.cuh +++ b/csrc/sm70_tile_runtime_signal.cuh @@ -4,19 +4,17 @@ namespace vllm::sm70_tile_runtime { -// Keep the default TileRT/all-reduce block limit unchanged. The larger signal -// capacity is reserved for long-prefill collective-fusion experiments that use -// one cross-rank barrier per owned token. +// Keep the default all-reduce block limit unchanged, but allow the +// TileRT-style MLP down-proj path to publish one flag per TurboMind N tile. constexpr int kMaxBlocks = 64; -constexpr int kMaxSignalBlocks = 512; constexpr int kMaxRanks = 8; using FlagType = uint32_t; struct Signal { - alignas(128) FlagType start[kMaxSignalBlocks][kMaxRanks]; - alignas(128) FlagType end[kMaxSignalBlocks][kMaxRanks]; - alignas(128) FlagType _flag[kMaxSignalBlocks]; + alignas(128) FlagType start[kMaxBlocks][kMaxRanks]; + alignas(128) FlagType end[kMaxBlocks][kMaxRanks]; + alignas(128) FlagType _flag[kMaxBlocks]; }; struct __align__(16) RankData { @@ -30,13 +28,14 @@ struct __align__(16) RankSignals { #if !defined(USE_ROCM) static __device__ __forceinline__ void store_flag_sys_visible( FlagType* flag_addr, FlagType flag) { - asm volatile("membar.sys; st.volatile.global.u32 [%1], %0;" ::"r"(flag), + asm volatile("membar.sys; st.volatile.global.u32 [%1], %0;" + ::"r"(flag), "l"(flag_addr) : "memory"); } -static __device__ __forceinline__ FlagType -load_flag_sys_visible(FlagType* flag_addr) { +static __device__ __forceinline__ FlagType load_flag_sys_visible( + FlagType* flag_addr) { FlagType flag; asm volatile("ld.volatile.global.u32 %0, [%1]; membar.sys;" : "=r"(flag) diff --git a/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/attention/quantization.h b/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/attention/quantization.h index fbb63b8426..3a3fd5e6ff 100644 --- a/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/attention/quantization.h +++ b/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/attention/quantization.h @@ -14,858 +14,875 @@ namespace turbomind { #define TM_ROUND_USE_CVT_RNI 1 -inline constexpr bool kFuseU4F16Dequant = false; +inline constexpr bool kFuseU4F16Dequant = false; inline constexpr bool kForceIntZeroPoint = false; -template -__device__ T Infinity() { - if constexpr (std::is_same_v) { - return __ushort_as_half((unsigned short)0x7C00U); - } +template +__device__ T Infinity() +{ + if constexpr (std::is_same_v) { + return __ushort_as_half((unsigned short)0x7C00U); + } #if __CUDA_ARCH__ >= 800 - if constexpr (std::is_same_v) { - return __ushort_as_bfloat16((unsigned short)0x7F80U); - } + if constexpr (std::is_same_v) { + return __ushort_as_bfloat16((unsigned short)0x7F80U); + } #endif - if constexpr (std::is_same_v) { - return __int_as_float(0x7f800000U); - } + if constexpr (std::is_same_v) { + return __int_as_float(0x7f800000U); + } - return T{}; + return T{}; } -template -__device__ constexpr T Max(T a, T b) { - if constexpr (std::is_same_v) { - return __hmax(a, b); - } +template +__device__ constexpr T Max(T a, T b) +{ + if constexpr (std::is_same_v) { + return __hmax(a, b); + } #if __CUDA_ARCH__ >= 800 - if constexpr (std::is_same_v) { - return __hmax(a, b); - } + if constexpr (std::is_same_v) { + return __hmax(a, b); + } #endif - if constexpr (std::is_same_v) { - return fmaxf(a, b); - } + if constexpr (std::is_same_v) { + return fmaxf(a, b); + } - if constexpr (std::is_same_v) { - return max(a, b); - } + if constexpr (std::is_same_v) { + return max(a, b); + } - return T{}; + return T{}; } -template -__device__ constexpr T Min(T a, T b) { - if constexpr (std::is_same_v) { - return __hmin(a, b); - } +template +__device__ constexpr T Min(T a, T b) +{ + if constexpr (std::is_same_v) { + return __hmin(a, b); + } #if __CUDA_ARCH__ >= 800 - if constexpr (std::is_same_v) { - return __hmin(a, b); - } + if constexpr (std::is_same_v) { + return __hmin(a, b); + } #endif - if constexpr (std::is_same_v) { - return fminf(a, b); - } + if constexpr (std::is_same_v) { + return fminf(a, b); + } - if constexpr (std::is_same_v) { - return min(a, b); - } + if constexpr (std::is_same_v) { + return min(a, b); + } - return T{}; + return T{}; } -template -inline __device__ Array cvt_f16x4_u8(const Array& src) { - static constexpr uint32_t f16_magic = 0x64000000; - // 01234567 01234567 - // SEEEEEMM MMMMMMMM - // 1MM XXXXXXXX - // (1 + x/2^10) * 2^(e-15) -> e-15=10 -> e=25=16+8+1 -> 01100100b -> 0x64 - Array dst; - dst[0] = __byte_perm((uint32_t&)src, f16_magic, 0x7170); - dst[1] = __byte_perm((uint32_t&)src, f16_magic, 0x7372); - if constexpr (norm) { - for (int i = 0; i < 4; ++i) { - ((Array&)dst)[i] -= __ushort_as_half(0x6400U); +template +inline __device__ Array cvt_f16x4_u8(const Array& src) +{ + static constexpr uint32_t f16_magic = 0x64000000; + // 01234567 01234567 + // SEEEEEMM MMMMMMMM + // 1MM XXXXXXXX + // (1 + x/2^10) * 2^(e-15) -> e-15=10 -> e=25=16+8+1 -> 01100100b -> 0x64 + Array dst; + dst[0] = __byte_perm((uint32_t&)src, f16_magic, 0x7170); + dst[1] = __byte_perm((uint32_t&)src, f16_magic, 0x7372); + if constexpr (norm) { + for (int i = 0; i < 4; ++i) { + ((Array&)dst)[i] -= __ushort_as_half(0x6400U); + } } - } - return (Array&)dst; + return (Array&)dst; } -template -inline __device__ Array cvt_f16x2x2_u8_trans( - const Array& src) { - static constexpr uint32_t f16_magic = 0x64000000; - // 01234567 01234567 - // SEEEEEMM MMMMMMMM - // 1MM XXXXXXXX - // (1 + x/2^10) * 2^(e-15) -> e-15=10 -> e=25=16+8+1 -> 01100100b -> 0x64 - Array dst; - dst[0] = __byte_perm((uint32_t&)src, f16_magic, 0x7270); - dst[1] = __byte_perm((uint32_t&)src, f16_magic, 0x7371); - if constexpr (norm) { - for (int i = 0; i < 4; ++i) { - ((Array&)dst)[i] -= __ushort_as_half(0x6400U); +template +inline __device__ Array cvt_f16x2x2_u8_trans(const Array& src) +{ + static constexpr uint32_t f16_magic = 0x64000000; + // 01234567 01234567 + // SEEEEEMM MMMMMMMM + // 1MM XXXXXXXX + // (1 + x/2^10) * 2^(e-15) -> e-15=10 -> e=25=16+8+1 -> 01100100b -> 0x64 + Array dst; + dst[0] = __byte_perm((uint32_t&)src, f16_magic, 0x7270); + dst[1] = __byte_perm((uint32_t&)src, f16_magic, 0x7371); + if constexpr (norm) { + for (int i = 0; i < 4; ++i) { + ((Array&)dst)[i] -= __ushort_as_half(0x6400U); + } } - } - return (Array&)dst; + return (Array&)dst; } -inline __device__ Array cvt_bf16x4_u8( - const Array& src) { - // 01234567 01234567 01234567 01234567 - // SEEEEEEE EMMMMMMM MMMMMMMM MMMMMMMM - // 1MM... XXXXXXXX - // (1 + x/2^15) * 2^(e-127) -> e-127=15 -> e=142 -> 01000111 -> 0x47 - static constexpr uint32_t f32_magic = 0x47000000; // 32768 - - Array tmp; - tmp[0] = __byte_perm((uint32_t&)src, f32_magic, 0x7604); - tmp[1] = __byte_perm((uint32_t&)src, f32_magic, 0x7614); - tmp[2] = __byte_perm((uint32_t&)src, f32_magic, 0x7624); - tmp[3] = __byte_perm((uint32_t&)src, f32_magic, 0x7634); - - auto& vec = (Array&)tmp; - - Array dst; - PRAGMA_UNROLL - for (int i = 0; i < 4; ++i) { - dst[i] = __float2bfloat16(vec[i] - 32768.f); - } - return dst; +inline __device__ Array cvt_bf16x4_u8(const Array& src) +{ + // 01234567 01234567 01234567 01234567 + // SEEEEEEE EMMMMMMM MMMMMMMM MMMMMMMM + // 1MM... XXXXXXXX + // (1 + x/2^15) * 2^(e-127) -> e-127=15 -> e=142 -> 01000111 -> 0x47 + static constexpr uint32_t f32_magic = 0x47000000; // 32768 + + Array tmp; + tmp[0] = __byte_perm((uint32_t&)src, f32_magic, 0x7604); + tmp[1] = __byte_perm((uint32_t&)src, f32_magic, 0x7614); + tmp[2] = __byte_perm((uint32_t&)src, f32_magic, 0x7624); + tmp[3] = __byte_perm((uint32_t&)src, f32_magic, 0x7634); + + auto& vec = (Array&)tmp; + + Array dst; + PRAGMA_UNROLL + for (int i = 0; i < 4; ++i) { + dst[i] = __float2bfloat16(vec[i] - 32768.f); + } + return dst; } -inline __device__ Array cvt_f32x4_u8(const Array& src) { - // 01234567 01234567 01234567 01234567 - // SEEEEEEE EMMMMMMM MMMMMMMM MMMMMMMM - // 1MM... XXXXXXXX - // (1 + x/2^15) * 2^(e-127) -> e-127=15 -> e=142 -> 01000111 -> 0x47 - static constexpr uint32_t f32_magic = 0x47000000; // 32768 - - Array tmp; - tmp[0] = __byte_perm((uint32_t&)src, f32_magic, 0x7604); - tmp[1] = __byte_perm((uint32_t&)src, f32_magic, 0x7614); - tmp[2] = __byte_perm((uint32_t&)src, f32_magic, 0x7624); - tmp[3] = __byte_perm((uint32_t&)src, f32_magic, 0x7634); - - auto& vec = (Array&)tmp; - PRAGMA_UNROLL - for (int i = 0; i < 4; ++i) { - vec[i] -= 32768.f; - } - return vec; +inline __device__ Array cvt_f32x4_u8(const Array& src) +{ + // 01234567 01234567 01234567 01234567 + // SEEEEEEE EMMMMMMM MMMMMMMM MMMMMMMM + // 1MM... XXXXXXXX + // (1 + x/2^15) * 2^(e-127) -> e-127=15 -> e=142 -> 01000111 -> 0x47 + static constexpr uint32_t f32_magic = 0x47000000; // 32768 + + Array tmp; + tmp[0] = __byte_perm((uint32_t&)src, f32_magic, 0x7604); + tmp[1] = __byte_perm((uint32_t&)src, f32_magic, 0x7614); + tmp[2] = __byte_perm((uint32_t&)src, f32_magic, 0x7624); + tmp[3] = __byte_perm((uint32_t&)src, f32_magic, 0x7634); + + auto& vec = (Array&)tmp; + PRAGMA_UNROLL + for (int i = 0; i < 4; ++i) { + vec[i] -= 32768.f; + } + return vec; } -template -inline __device__ Array cvt_bf16x8_u4( - const Array& src) { +template +inline __device__ Array cvt_bf16x8_u4(const Array& src) +{ #if __CUDA_ARCH__ >= 800 - // 01234567 01234567 - // SEEEEEEE EMMMMMMM - // 1...XXXX - // (1 + x/2^7) * 2^(e-127) -> e-127=7 -> e=134 -> 0100 0011 -> 0x43 - static constexpr uint32_t TEMPLATE = 0x43004300; // nv_bfloat162(128, 128) - static constexpr uint32_t MASK = 0x000f000f; - static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa; - - Array h; - - static_assert(sizeof(Array) == sizeof(Array)); - - uint32_t const& i4s = reinterpret_cast(src); - const uint32_t i4s_4 = i4s >> 4; - const uint32_t i4s_8 = i4s >> 8; - const uint32_t i4s_12 = i4s >> 12; - - asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" - : "=r"(h[0]) - : "r"(i4s), "n"(MASK), "n"(TEMPLATE), "n"(immLut)); - asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" - : "=r"(h[1]) - : "r"(i4s_4), "n"(MASK), "n"(TEMPLATE), "n"(immLut)); - asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" - : "=r"(h[2]) - : "r"(i4s_8), "n"(MASK), "n"(TEMPLATE), "n"(immLut)); - asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" - : "=r"(h[3]) - : "r"(i4s_12), "n"(MASK), "n"(TEMPLATE), "n"(immLut)); - - if constexpr (norm) { - auto result = reinterpret_cast(h.data()); - PRAGMA_UNROLL - for (int i = 0; i < 8; ++i) { - result[i] -= nv_bfloat16(128.f); + // 01234567 01234567 + // SEEEEEEE EMMMMMMM + // 1...XXXX + // (1 + x/2^7) * 2^(e-127) -> e-127=7 -> e=134 -> 0100 0011 -> 0x43 + static constexpr uint32_t TEMPLATE = 0x43004300; // nv_bfloat162(128, 128) + static constexpr uint32_t MASK = 0x000f000f; + static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa; + + Array h; + + static_assert(sizeof(Array) == sizeof(Array)); + + uint32_t const& i4s = reinterpret_cast(src); + const uint32_t i4s_4 = i4s >> 4; + const uint32_t i4s_8 = i4s >> 8; + const uint32_t i4s_12 = i4s >> 12; + + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[0]) : "r"(i4s), "n"(MASK), "n"(TEMPLATE), "n"(immLut)); + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[1]) : "r"(i4s_4), "n"(MASK), "n"(TEMPLATE), "n"(immLut)); + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[2]) : "r"(i4s_8), "n"(MASK), "n"(TEMPLATE), "n"(immLut)); + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[3]) : "r"(i4s_12), "n"(MASK), "n"(TEMPLATE), "n"(immLut)); + + if constexpr (norm) { + auto result = reinterpret_cast(h.data()); + PRAGMA_UNROLL + for (int i = 0; i < 8; ++i) { + result[i] -= nv_bfloat16(128.f); + } } - } - return (Array&)h; + return (Array&)h; #else - return {}; + return {}; #endif } #if TM_ROUND_USE_CVT_RNI -template -inline __device__ T round(float x) { - uint32_t y{}; - if constexpr (std::is_same_v) { - asm("cvt.rni.sat.u8.f32 %0, %1;\n" : "=r"(y) : "f"(x)); - } else if constexpr (std::is_same_v) { - asm("cvt.rni.sat.u16.f32 %0, %1;\n" : "=r"(y) : "f"(x)); - } else if constexpr (std::is_same_v) { - asm("cvt.rni.sat.u32.f32 %0, %1;\n" : "=r"(y) : "f"(x)); - } else if constexpr (std::is_same_v) { - asm("cvt.rni.sat.s32.f32 %0, %1;\n" : "=r"(y) : "f"(x)); - } else { - static_assert(!std::is_same_v, "not implemented"); - } - return y; +template +inline __device__ T round(float x) +{ + uint32_t y{}; + if constexpr (std::is_same_v) { + asm("cvt.rni.sat.u8.f32 %0, %1;\n" : "=r"(y) : "f"(x)); + } + else if constexpr (std::is_same_v) { + asm("cvt.rni.sat.u16.f32 %0, %1;\n" : "=r"(y) : "f"(x)); + } + else if constexpr (std::is_same_v) { + asm("cvt.rni.sat.u32.f32 %0, %1;\n" : "=r"(y) : "f"(x)); + } + else if constexpr (std::is_same_v) { + asm("cvt.rni.sat.s32.f32 %0, %1;\n" : "=r"(y) : "f"(x)); + } + else { + static_assert(!std::is_same_v, "not implemented"); + } + return y; } -template -inline __device__ T round(half x) { - uint32_t y{}; - if constexpr (std::is_same_v) { - asm("cvt.rni.sat.u8.f16 %0, %1;\n" : "=r"(y) : "h"((uint16_t&)x)); - } else if constexpr (std::is_same_v) { - asm("cvt.rni.sat.u16.f16 %0, %1;\n" : "=r"(y) : "h"((uint16_t&)x)); - } else if constexpr (std::is_same_v) { - asm("cvt.rni.sat.u32.f16 %0, %1;\n" : "=r"(y) : "h"((uint16_t&)x)); - } else if constexpr (std::is_same_v) { - asm("cvt.rni.sat.s32.f16 %0, %1;\n" : "=r"(y) : "h"((uint16_t&)x)); - } else { - static_assert(!std::is_same_v, "not implemented"); - } - return y; +template +inline __device__ T round(half x) +{ + uint32_t y{}; + if constexpr (std::is_same_v) { + asm("cvt.rni.sat.u8.f16 %0, %1;\n" : "=r"(y) : "h"((uint16_t&)x)); + } + else if constexpr (std::is_same_v) { + asm("cvt.rni.sat.u16.f16 %0, %1;\n" : "=r"(y) : "h"((uint16_t&)x)); + } + else if constexpr (std::is_same_v) { + asm("cvt.rni.sat.u32.f16 %0, %1;\n" : "=r"(y) : "h"((uint16_t&)x)); + } + else if constexpr (std::is_same_v) { + asm("cvt.rni.sat.s32.f16 %0, %1;\n" : "=r"(y) : "h"((uint16_t&)x)); + } + else { + static_assert(!std::is_same_v, "not implemented"); + } + return y; } #else -template -inline __device__ T round(float x) { - x += .5f; - - uint32_t y{}; - if constexpr (std::is_same_v) { - asm("cvt.rzi.sat.u8.f32 %0, %1;\n" : "=r"(y) : "f"(x)); - } else if constexpr (std::is_same_v) { - asm("cvt.rzi.sat.u16.f32 %0, %1;\n" : "=r"(y) : "f"(x)); - } else if constexpr (std::is_same_v) { - asm("cvt.rzi.sat.u32.f32 %0, %1;\n" : "=r"(y) : "f"(x)); - } else { - static_assert(!std::is_same_v, "not implemented"); - } - return y; -} +template +inline __device__ T round(float x) +{ + x += .5f; -template -inline __device__ T round(half x) { - x += half(.5f); - - uint32_t y{}; - if constexpr (std::is_same_v) { - asm("cvt.rzi.sat.u8.f16 %0, %1;\n" : "=r"(y) : "h"((uint16_t&)x)); - } else if constexpr (std::is_same_v) { - asm("cvt.rzi.sat.u16.f16 %0, %1;\n" : "=r"(y) : "h"((uint16_t&)x)); - } else if constexpr (std::is_same_v) { - asm("cvt.rzi.sat.u32.f16 %0, %1;\n" : "=r"(y) : "h"((uint16_t&)x)); - } else { - static_assert(!std::is_same_v, "not implemented"); - } - return y; + uint32_t y{}; + if constexpr (std::is_same_v) { + asm("cvt.rzi.sat.u8.f32 %0, %1;\n" : "=r"(y) : "f"(x)); + } + else if constexpr (std::is_same_v) { + asm("cvt.rzi.sat.u16.f32 %0, %1;\n" : "=r"(y) : "f"(x)); + } + else if constexpr (std::is_same_v) { + asm("cvt.rzi.sat.u32.f32 %0, %1;\n" : "=r"(y) : "f"(x)); + } + else { + static_assert(!std::is_same_v, "not implemented"); + } + return y; } -#endif +template +inline __device__ T round(half x) +{ + x += half(.5f); -template -inline __device__ To quant(Ti x, B n_bits) { - auto y = round(x); - if constexpr (n_bits < - sizeof(To) * 8) { // saturate operation for sub-byte type - return min(y, To((1 << n_bits) - 1)); - } else { + uint32_t y{}; + if constexpr (std::is_same_v) { + asm("cvt.rzi.sat.u8.f16 %0, %1;\n" : "=r"(y) : "h"((uint16_t&)x)); + } + else if constexpr (std::is_same_v) { + asm("cvt.rzi.sat.u16.f16 %0, %1;\n" : "=r"(y) : "h"((uint16_t&)x)); + } + else if constexpr (std::is_same_v) { + asm("cvt.rzi.sat.u32.f16 %0, %1;\n" : "=r"(y) : "h"((uint16_t&)x)); + } + else { + static_assert(!std::is_same_v, "not implemented"); + } return y; - } } -template -__device__ inline void warp_minmax(Array& stats, const Array& x) { - PRAGMA_UNROLL - for (int i = 0; i < C; ++i) { - stats[0] = Min(stats[0], x[i]); - stats[1] = Max(stats[1], x[i]); - } - if constexpr (sizeof(T) == 2) { - PRAGMA_UNROLL - for (int mask = WarpThreadC / 2; mask > 0; mask /= 2) { - Array tmp; - (uint32_t&)tmp = __shfl_xor_sync(uint32_t(-1), (uint32_t&)stats, mask); - stats[0] = Min(stats[0], tmp[0]); - stats[1] = Max(stats[1], tmp[1]); +#endif + +template +inline __device__ To quant(Ti x, B n_bits) +{ + auto y = round(x); + if constexpr (n_bits < sizeof(To) * 8) { // saturate operation for sub-byte type + return min(y, To((1 << n_bits) - 1)); } - } else { - PRAGMA_UNROLL - for (int mask = WarpThreadC / 2; mask > 0; mask /= 2) { - stats[0] = Min(stats[0], __shfl_xor_sync(uint32_t(-1), stats[0], mask)); - stats[1] = Max(stats[1], __shfl_xor_sync(uint32_t(-1), stats[1], mask)); + else { + return y; } - } } -template -__device__ void warp_stats(Array (¶m)[S], - const Array (&x)[S][C], B n_bits) { - PRAGMA_UNROLL - for (int s = 0; s < S; ++s) { - Array stats{Infinity(), -Infinity()}; +template +__device__ inline void warp_minmax(Array& stats, const Array& x) +{ PRAGMA_UNROLL - for (int c = 0; c < C; ++c) { - warp_minmax(stats, x[s][c]); + for (int i = 0; i < C; ++i) { + stats[0] = Min(stats[0], x[i]); + stats[1] = Max(stats[1], x[i]); + } + if constexpr (sizeof(T) == 2) { + PRAGMA_UNROLL + for (int mask = WarpThreadC / 2; mask > 0; mask /= 2) { + Array tmp; + (uint32_t&)tmp = __shfl_xor_sync(uint32_t(-1), (uint32_t&)stats, mask); + stats[0] = Min(stats[0], tmp[0]); + stats[1] = Max(stats[1], tmp[1]); + } } - const float inv_q_max = fdividef(1.f, float((1 << n_bits) - 1)); - const float scale = ((float)stats[1] - (float)stats[0]) * inv_q_max; - param[s][0] = (P)scale; - param[s][1] = (P)stats[0]; + else { + PRAGMA_UNROLL + for (int mask = WarpThreadC / 2; mask > 0; mask /= 2) { + stats[0] = Min(stats[0], __shfl_xor_sync(uint32_t(-1), stats[0], mask)); + stats[1] = Max(stats[1], __shfl_xor_sync(uint32_t(-1), stats[1], mask)); + } + } +} - if constexpr (kForceIntZeroPoint) { +template +__device__ void warp_stats(Array (¶m)[S], const Array (&x)[S][C], B n_bits) +{ + PRAGMA_UNROLL + for (int s = 0; s < S; ++s) { + Array stats{Infinity(), -Infinity()}; + PRAGMA_UNROLL + for (int c = 0; c < C; ++c) { + warp_minmax(stats, x[s][c]); + } + const float inv_q_max = fdividef(1.f, float((1 << n_bits) - 1)); + const float scale = ((float)stats[1] - (float)stats[0]) * inv_q_max; + param[s][0] = (P)scale; + param[s][1] = (P)stats[0]; + + if constexpr (kForceIntZeroPoint) { #if TM_ROUND_USE_CVT_RNI - // rintf -> cvt.rni.f32.f32 - param[s][1] = (P)(rintf((float)stats[0] / scale) * scale); + // rintf -> cvt.rni.f32.f32 + param[s][1] = (P)(rintf((float)stats[0] / scale) * scale); #else - // roundf -> cvt.rzi.f32.f32(x + 0.5) - param[s][1] = (P)(roundf((float)stats[0] / scale) * scale); + // roundf -> cvt.rzi.f32.f32(x + 0.5) + param[s][1] = (P)(roundf((float)stats[0] / scale) * scale); #endif + } } - } } -template -__device__ void quantize(Array (&dst)[S][C], - const Array (&src)[S][C], - const Array (¶ms)[S], B n_bits) { - PRAGMA_UNROLL - for (int s = 0; s < S; ++s) { - P inv_scale = (P)fdividef(1.f, (float)params[s][0]); - P zero = params[s][1]; +template +__device__ void +quantize(Array (&dst)[S][C], const Array (&src)[S][C], const Array (¶ms)[S], B n_bits) +{ PRAGMA_UNROLL - for (int c = 0; c < C; ++c) { - PRAGMA_UNROLL - for (int i = 0; i < N; ++i) { - const auto v = ((P)src[s][c][i] - zero) * inv_scale; - dst[s][c][i] = quant(v, n_bits); - } - } - } + for (int s = 0; s < S; ++s) { + P inv_scale = (P)fdividef(1.f, (float)params[s][0]); + P zero = params[s][1]; + PRAGMA_UNROLL + for (int c = 0; c < C; ++c) { + PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + const auto v = ((P)src[s][c][i] - zero) * inv_scale; + dst[s][c][i] = quant(v, n_bits); + } + } + } } ////////////////////////////////////////////////////////////////////////////////////////////////// -// generic case for floating point -> floating point / integer -> integer -// conversion -template +// generic case for floating point -> floating point / integer -> integer conversion +template struct ConvertKvCache { - __device__ __host__ ConvertKvCache(float, float) {} - template - __device__ static auto convert(const Array& vi) { - Array vo; - PRAGMA_UNROLL - for (int i = 0; i < N; ++i) { - vo[i] = (To)vi[i]; - } - return vo; - } - template - inline __device__ auto operator()(const Array& vi) const - -> Array { - return convert(vi); - } + __device__ __host__ ConvertKvCache(float, float) {} + template + __device__ static auto convert(const Array& vi) + { + Array vo; + PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + vo[i] = (To)vi[i]; + } + return vo; + } + template + inline __device__ auto operator()(const Array& vi) const -> Array + { + return convert(vi); + } }; // generic case for converting to same type, bypass -template +template struct ConvertKvCache { - __device__ __host__ ConvertKvCache(float, float) {} - template - __device__ static auto convert(const Array& v) { - return v; - } - template - inline __device__ auto operator()(const Array& v) const -> Array { - return convert(v); - } + __device__ __host__ ConvertKvCache(float, float) {} + template + __device__ static auto convert(const Array& v) + { + return v; + } + template + inline __device__ auto operator()(const Array& v) const -> Array + { + return convert(v); + } }; // floating point -> u8 -template +template struct ConvertKvCache { - T inv_scale_; - T zero_; - __device__ ConvertKvCache(T scale, T zero) : zero_{zero} { - // NVCC complains if we put this in the member init list - inv_scale_ = (T)fdividef(1.f, (float)scale); - } - - template - __device__ auto operator()(const Array& vi) const { - Array vo; - PRAGMA_UNROLL - for (int i = 0; i < N; ++i) { - vo[i] = quant((vi[i] - zero_) * inv_scale_, - std::integral_constant{}); + T inv_scale_; + T zero_; + __device__ ConvertKvCache(T scale, T zero): zero_{zero} + { + // NVCC complains if we put this in the member init list + inv_scale_ = (T)fdividef(1.f, (float)scale); + } + + template + __device__ auto operator()(const Array& vi) const + { + Array vo; + PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + vo[i] = quant((vi[i] - zero_) * inv_scale_, std::integral_constant{}); + } + return vo; } - return vo; - } }; -template +template struct ConvertKvCache { - T inv_scale_; - T zero_; - __device__ ConvertKvCache(T scale, T zero) : zero_{zero} { - // NVCC complains if we put this in the member init list - inv_scale_ = (T)fdividef(1.f, (float)scale); - } - - static __device__ Array pack(const Array& vi) { - Array ui = (Array&)vi; - - ui[0] |= (ui[0] >> 12); - ui[1] |= (ui[1] >> 12); - - // 7 6 5 4 3 2 1 0 - // _7_67564_3_23120 - uint32_t uo = __byte_perm(ui[0], ui[1], 0x5140); - - return (Array&)uo; - } - - /// TODO: try cvt.pack.sat.u4 - template - __device__ auto operator()(const Array& vi) const { - static_assert(N % 8 == 0); - Array tmp; - PRAGMA_UNROLL - for (int i = 0; i < N; ++i) { - tmp[i] = quant((vi[i] - zero_) * inv_scale_, - std::integral_constant{}); + T inv_scale_; + T zero_; + __device__ ConvertKvCache(T scale, T zero): zero_{zero} + { + // NVCC complains if we put this in the member init list + inv_scale_ = (T)fdividef(1.f, (float)scale); } - Array vo; - PRAGMA_UNROLL - for (int i = 0; i < N; i += 8) { - (Array&)vo[i] = pack((Array&)tmp[i]); + + static __device__ Array pack(const Array& vi) + { + Array ui = (Array&)vi; + + ui[0] |= (ui[0] >> 12); + ui[1] |= (ui[1] >> 12); + + // 7 6 5 4 3 2 1 0 + // _7_67564_3_23120 + uint32_t uo = __byte_perm(ui[0], ui[1], 0x5140); + + return (Array&)uo; + } + + /// TODO: try cvt.pack.sat.u4 + template + __device__ auto operator()(const Array& vi) const + { + static_assert(N % 8 == 0); + Array tmp; + PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + tmp[i] = quant((vi[i] - zero_) * inv_scale_, std::integral_constant{}); + } + Array vo; + PRAGMA_UNROLL + for (int i = 0; i < N; i += 8) { + (Array&)vo[i] = pack((Array&)tmp[i]); + } + return vo; } - return vo; - } }; -template <> +template<> struct ConvertKvCache { - half scale_; - half zero_; - - __device__ ConvertKvCache(half scale, half zero) { - scale_ = scale; - zero_ = zero; - } - - static __device__ Array cvt_f16x8_u4(const Array& vi) { - Array result; - uint32_t* h = reinterpret_cast(&result); - uint32_t const& i4s = reinterpret_cast(vi); - static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa; - static constexpr uint32_t BOT_MASK = 0x000f000f; - static constexpr uint32_t TOP_MASK = 0x00f000f0; - static constexpr uint32_t MAGIC_NUM_0 = 0x64006400; // `1024` - static constexpr uint32_t MAGIC_NUM_1 = 0x54005400; // `64` - // const uint32_t top_i4s = i4s >> 8; - uint32_t top_i4s = __byte_perm(i4s, 0, 0x4321); - asm("lop3.b32 %0, %1, %2, %3, %4;\n" - : "=r"(h[0]) - : "r"(i4s), "n"(BOT_MASK), "n"(MAGIC_NUM_0), "n"(immLut)); - asm("lop3.b32 %0, %1, %2, %3, %4;\n" - : "=r"(h[1]) - : "r"(i4s), "n"(TOP_MASK), "n"(MAGIC_NUM_1), "n"(immLut)); - asm("lop3.b32 %0, %1, %2, %3, %4;\n" - : "=r"(h[2]) - : "r"(top_i4s), "n"(BOT_MASK), "n"(MAGIC_NUM_0), "n"(immLut)); - asm("lop3.b32 %0, %1, %2, %3, %4;\n" - : "=r"(h[3]) - : "r"(top_i4s), "n"(TOP_MASK), "n"(MAGIC_NUM_1), "n"(immLut)); - asm("sub.f16x2 %0, %1, %2;\n" : "=r"(h[0]) : "r"(h[0]), "r"(MAGIC_NUM_0)); - asm("sub.f16x2 %0, %1, %2;\n" : "=r"(h[1]) : "r"(h[1]), "r"(MAGIC_NUM_1)); - asm("sub.f16x2 %0, %1, %2;\n" : "=r"(h[2]) : "r"(h[2]), "r"(MAGIC_NUM_0)); - asm("sub.f16x2 %0, %1, %2;\n" : "=r"(h[3]) : "r"(h[3]), "r"(MAGIC_NUM_1)); - return result; - } - - static __device__ Array cvt_f16x8_u4_biased( - const Array& vi) { - Array result; - uint32_t* h = reinterpret_cast(&result); - uint32_t const& i4s = reinterpret_cast(vi); - static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa; - static constexpr uint32_t BOT_MASK = 0x000f000f; - static constexpr uint32_t TOP_MASK = 0x00f000f0; - static constexpr uint32_t MAGIC_NUM_1 = 0x54005400; // `64` - static constexpr uint32_t MAGIC_NUM_2 = MAGIC_NUM_1 >> 4; // `64` >> 4 - const uint32_t top_i4s = i4s >> 8; - // uint32_t top_i4s = __byte_perm(i4s, 0, 0x4321); - asm("lop3.b32 %0, %1, %2, %3, %4;\n" - : "=r"(h[0]) - : "r"(i4s), "n"(BOT_MASK), "n"(MAGIC_NUM_2), "n"(immLut)); - asm("lop3.b32 %0, %1, %2, %3, %4;\n" - : "=r"(h[1]) - : "r"(i4s), "n"(TOP_MASK), "n"(MAGIC_NUM_1), "n"(immLut)); - asm("lop3.b32 %0, %1, %2, %3, %4;\n" - : "=r"(h[2]) - : "r"(top_i4s), "n"(BOT_MASK), "n"(MAGIC_NUM_2), "n"(immLut)); - asm("lop3.b32 %0, %1, %2, %3, %4;\n" - : "=r"(h[3]) - : "r"(top_i4s), "n"(TOP_MASK), "n"(MAGIC_NUM_1), "n"(immLut)); - h[0] <<= 4; - h[2] <<= 4; - return result; - } - - template - __device__ static auto convert(const Array& vi) { - Array vo; - PRAGMA_UNROLL - for (int i = 0; i < N; i += 8) { - auto& v = (Array&)vo[i]; - if constexpr (kFuseU4F16Dequant) { - v = cvt_f16x8_u4_biased((Array&)vi[i]); - } else { - v = cvt_f16x8_u4((Array&)vi[i]); - } - } - return vo; - } - - template - __device__ auto operator()(const Array& vi) const { - auto vo = convert(vi); - PRAGMA_UNROLL - for (int i = 0; i < N; ++i) { - vo[i] = vo[i] * scale_ + zero_; + + half scale_; + half zero_; + + __device__ ConvertKvCache(half scale, half zero) + { + scale_ = scale; + zero_ = zero; + } + + static __device__ Array cvt_f16x8_u4(const Array& vi) + { + Array result; + uint32_t* h = reinterpret_cast(&result); + uint32_t const& i4s = reinterpret_cast(vi); + static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa; + static constexpr uint32_t BOT_MASK = 0x000f000f; + static constexpr uint32_t TOP_MASK = 0x00f000f0; + static constexpr uint32_t MAGIC_NUM_0 = 0x64006400; // `1024` + static constexpr uint32_t MAGIC_NUM_1 = 0x54005400; // `64` + // const uint32_t top_i4s = i4s >> 8; + uint32_t top_i4s = __byte_perm(i4s, 0, 0x4321); + asm("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[0]) : "r"(i4s), "n"(BOT_MASK), "n"(MAGIC_NUM_0), "n"(immLut)); + asm("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[1]) : "r"(i4s), "n"(TOP_MASK), "n"(MAGIC_NUM_1), "n"(immLut)); + asm("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[2]) : "r"(top_i4s), "n"(BOT_MASK), "n"(MAGIC_NUM_0), "n"(immLut)); + asm("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[3]) : "r"(top_i4s), "n"(TOP_MASK), "n"(MAGIC_NUM_1), "n"(immLut)); + asm("sub.f16x2 %0, %1, %2;\n" : "=r"(h[0]) : "r"(h[0]), "r"(MAGIC_NUM_0)); + asm("sub.f16x2 %0, %1, %2;\n" : "=r"(h[1]) : "r"(h[1]), "r"(MAGIC_NUM_1)); + asm("sub.f16x2 %0, %1, %2;\n" : "=r"(h[2]) : "r"(h[2]), "r"(MAGIC_NUM_0)); + asm("sub.f16x2 %0, %1, %2;\n" : "=r"(h[3]) : "r"(h[3]), "r"(MAGIC_NUM_1)); + return result; + } + + static __device__ Array cvt_f16x8_u4_biased(const Array& vi) + { + Array result; + uint32_t* h = reinterpret_cast(&result); + uint32_t const& i4s = reinterpret_cast(vi); + static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa; + static constexpr uint32_t BOT_MASK = 0x000f000f; + static constexpr uint32_t TOP_MASK = 0x00f000f0; + static constexpr uint32_t MAGIC_NUM_1 = 0x54005400; // `64` + static constexpr uint32_t MAGIC_NUM_2 = MAGIC_NUM_1 >> 4; // `64` >> 4 + const uint32_t top_i4s = i4s >> 8; + // uint32_t top_i4s = __byte_perm(i4s, 0, 0x4321); + asm("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[0]) : "r"(i4s), "n"(BOT_MASK), "n"(MAGIC_NUM_2), "n"(immLut)); + asm("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[1]) : "r"(i4s), "n"(TOP_MASK), "n"(MAGIC_NUM_1), "n"(immLut)); + asm("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[2]) : "r"(top_i4s), "n"(BOT_MASK), "n"(MAGIC_NUM_2), "n"(immLut)); + asm("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[3]) : "r"(top_i4s), "n"(TOP_MASK), "n"(MAGIC_NUM_1), "n"(immLut)); + h[0] <<= 4; + h[2] <<= 4; + return result; + } + + template + __device__ static auto convert(const Array& vi) + { + Array vo; + PRAGMA_UNROLL + for (int i = 0; i < N; i += 8) { + auto& v = (Array&)vo[i]; + if constexpr (kFuseU4F16Dequant) { + v = cvt_f16x8_u4_biased((Array&)vi[i]); + } + else { + v = cvt_f16x8_u4((Array&)vi[i]); + } + } + return vo; + } + + template + __device__ auto operator()(const Array& vi) const + { + auto vo = convert(vi); + PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + vo[i] = vo[i] * scale_ + zero_; + } + return vo; } - return vo; - } }; -template <> +template<> struct ConvertKvCache { - nv_bfloat16 scale_; - nv_bfloat16 zero_; - __device__ ConvertKvCache(nv_bfloat16 scale, nv_bfloat16 zero) { - scale_ = scale; - zero_ = zero; - } + nv_bfloat16 scale_; + nv_bfloat16 zero_; - template - __device__ static Array convert(const Array& vi) { - Array vo{}; - PRAGMA_UNROLL - for (int i = 0; i < N; i += 8) { - auto& v = (Array&)vo[i]; - auto u = cvt_bf16x8_u4((Array&)vi[i]); - v = (Array&)u; - } - return vo; - } - - template - __device__ Array operator()( - const Array& vi) const { - auto vo = convert(vi); - PRAGMA_UNROLL - for (int i = 0; i < N; ++i) { - vo[i] = vo[i] * scale_ + zero_; + __device__ ConvertKvCache(nv_bfloat16 scale, nv_bfloat16 zero) + { + scale_ = scale; + zero_ = zero; + } + + template + __device__ static Array convert(const Array& vi) + { + Array vo{}; + PRAGMA_UNROLL + for (int i = 0; i < N; i += 8) { + auto& v = (Array&)vo[i]; + auto u = cvt_bf16x8_u4((Array&)vi[i]); + v = (Array&)u; + } + return vo; + } + + template + __device__ Array operator()(const Array& vi) const + { + auto vo = convert(vi); + PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + vo[i] = vo[i] * scale_ + zero_; + } + return (Array&)vo; } - return (Array&)vo; - } }; -template <> +template<> struct ConvertKvCache { + #if 1 - ConvertKvCache impl_; + ConvertKvCache impl_; - __device__ ConvertKvCache(float scale, float zero) : impl_{scale, zero} {} + __device__ ConvertKvCache(float scale, float zero): impl_{scale, zero} {} - template - __device__ auto operator()(const Array& vi) const { - return cast(impl_(vi)); - } + template + __device__ auto operator()(const Array& vi) const + { + return cast(impl_(vi)); + } #else - static __device__ Array cvt_f16x8_u4_biased( - const Array& vi) { - Array result; - uint32_t* h = reinterpret_cast(&result); - uint32_t const& i4s = reinterpret_cast(vi); - static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa; - static constexpr uint32_t BOT_MASK = 0x000f000f; - static constexpr uint32_t TOP_MASK = 0x00f000f0; - static constexpr uint32_t MAGIC_NUM_1 = 0x54005400; // `64` - static constexpr uint32_t MAGIC_NUM_2 = MAGIC_NUM_1 >> 4; // `64` >> 4 - const uint32_t top_i4s = i4s >> 8; - asm("lop3.b32 %0, %1, %2, %3, %4;\n" - : "=r"(h[0]) - : "r"(i4s), "n"(BOT_MASK), "n"(MAGIC_NUM_2), "n"(immLut)); - asm("lop3.b32 %0, %1, %2, %3, %4;\n" - : "=r"(h[1]) - : "r"(i4s), "n"(TOP_MASK), "n"(MAGIC_NUM_1), "n"(immLut)); - asm("lop3.b32 %0, %1, %2, %3, %4;\n" - : "=r"(h[2]) - : "r"(top_i4s), "n"(BOT_MASK), "n"(MAGIC_NUM_2), "n"(immLut)); - asm("lop3.b32 %0, %1, %2, %3, %4;\n" - : "=r"(h[3]) - : "r"(top_i4s), "n"(TOP_MASK), "n"(MAGIC_NUM_1), "n"(immLut)); - h[0] <<= 4; - h[2] <<= 4; - return result; - } - float scale_; - float zero_; - __device__ ConvertKvCache(float scale, float zero) { - scale_ = scale; - zero_ = zero - scale * 64.f; - } - template - __device__ auto operator()(const Array& vi) const { - auto vo = cast(cvt_f16x8_u4_biased(vi)); - using namespace ops; - return vo * scale_ + zero_; - } + static __device__ Array cvt_f16x8_u4_biased(const Array& vi) + { + Array result; + uint32_t* h = reinterpret_cast(&result); + uint32_t const& i4s = reinterpret_cast(vi); + static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa; + static constexpr uint32_t BOT_MASK = 0x000f000f; + static constexpr uint32_t TOP_MASK = 0x00f000f0; + static constexpr uint32_t MAGIC_NUM_1 = 0x54005400; // `64` + static constexpr uint32_t MAGIC_NUM_2 = MAGIC_NUM_1 >> 4; // `64` >> 4 + const uint32_t top_i4s = i4s >> 8; + asm("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[0]) : "r"(i4s), "n"(BOT_MASK), "n"(MAGIC_NUM_2), "n"(immLut)); + asm("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[1]) : "r"(i4s), "n"(TOP_MASK), "n"(MAGIC_NUM_1), "n"(immLut)); + asm("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[2]) : "r"(top_i4s), "n"(BOT_MASK), "n"(MAGIC_NUM_2), "n"(immLut)); + asm("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[3]) : "r"(top_i4s), "n"(TOP_MASK), "n"(MAGIC_NUM_1), "n"(immLut)); + h[0] <<= 4; + h[2] <<= 4; + return result; + } + float scale_; + float zero_; + __device__ ConvertKvCache(float scale, float zero) + { + scale_ = scale; + zero_ = zero - scale * 64.f; + } + template + __device__ auto operator()(const Array& vi) const + { + auto vo = cast(cvt_f16x8_u4_biased(vi)); + using namespace ops; + return vo * scale_ + zero_; + } #endif }; // u8 -> f32/f16/bf16 -template +template struct ConvertKvCache { - T scale_; - T zero_; - __device__ ConvertKvCache(T scale, T zero) : scale_{scale}, zero_{zero} {} - - template - __device__ static auto convert(const Array& vi) { - Array vo; - PRAGMA_UNROLL - for (int n = 0; n < N; n += 4) { - auto& ui = (const Array&)vi[n]; - auto& uo = (Array&)vo[n]; - - if constexpr (std::is_same_v) { - uo = cvt_f16x4_u8(ui); - } else if constexpr (std::is_same_v) { - uo = cvt_f32x4_u8(ui); - } + T scale_; + T zero_; + __device__ ConvertKvCache(T scale, T zero): scale_{scale}, zero_{zero} {} + + template + __device__ static auto convert(const Array& vi) + { + Array vo; + PRAGMA_UNROLL + for (int n = 0; n < N; n += 4) { + auto& ui = (const Array&)vi[n]; + auto& uo = (Array&)vo[n]; + + if constexpr (std::is_same_v) { + uo = cvt_f16x4_u8(ui); + } + else if constexpr (std::is_same_v) { + uo = cvt_f32x4_u8(ui); + } #if __CUDA_ARCH__ >= 800 - else if constexpr (std::is_same_v) { - uo = cvt_bf16x4_u8(ui); - } + else if constexpr (std::is_same_v) { + uo = cvt_bf16x4_u8(ui); + } #endif + } + return vo; } - return vo; - } - template - __device__ auto operator()(const Array& vi) const { - auto vo = convert(vi); - PRAGMA_UNROLL - for (int i = 0; i < N; ++i) { - vo[i] = vo[i] * scale_ + zero_; + template + __device__ auto operator()(const Array& vi) const + { + auto vo = convert(vi); + PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + vo[i] = vo[i] * scale_ + zero_; + } + return vo; } - return vo; - } }; -template +template struct ConvertKvCache { - __device__ static Array cvt_bf16x8_e2m1( - const Array& vi) { - const uint32_t& x = (const uint32_t&)vi; - constexpr uint32_t S = 0x80008000U; - constexpr uint32_t EM = 0x01C001C0U; + __device__ static Array cvt_bf16x8_e2m1(const Array& vi) + { + const uint32_t& x = (const uint32_t&)vi; + + constexpr uint32_t S = 0x80008000U; + constexpr uint32_t EM = 0x01C001C0U; - Array vo; + Array vo; - // clang-format off + // clang-format off vo[0] = (x << 12 & S) | (x << 6 & EM); vo[1] = (x << 8 & S) | (x << 2 & EM); vo[2] = (x << 4 & S) | (x >> 2 & EM); vo[3] = (x << 0 & S) | (x >> 6 & EM); - // clang-format on + // clang-format on - constexpr uint32_t e = (127U - 1U + 127U) << 7U; - constexpr uint32_t ee = e << 16U | e; + constexpr uint32_t e = (127U - 1U + 127U) << 7U; + constexpr uint32_t ee = e << 16U | e; - PRAGMA_UNROLL - for (int i = 0; i < 4; ++i) { + PRAGMA_UNROLL + for (int i = 0; i < 4; ++i) { #if TURBOMIND_ARCH_SM90 - asm("mul.rn.bf16x2 %0, %1, %2;" : "=r"(vo[i]) : "r"(vo[i]), "r"(ee)); + asm("mul.rn.bf16x2 %0, %1, %2;" : "=r"(vo[i]) : "r"(vo[i]), "r"(ee)); #else - asm("fma.rn.bf16x2 %0, %1, %2, %3;" - : "=r"(vo[i]) - : "r"(vo[i]), "r"(ee), "r"(0)); + asm("fma.rn.bf16x2 %0, %1, %2, %3;" : "=r"(vo[i]) : "r"(vo[i]), "r"(ee), "r"(0)); #endif - } + } - return (Array&)vo; - } + return (Array&)vo; + } - __device__ static Array cvt_f16x8_e2m1( - const Array& vi) { - const uint32_t& x = (const uint32_t&)vi; + __device__ static Array cvt_f16x8_e2m1(const Array& vi) + { + const uint32_t& x = (const uint32_t&)vi; - constexpr uint32_t S = 0x80008000U; - constexpr uint32_t EM = 0x0E000E00U; + constexpr uint32_t S = 0x80008000U; + constexpr uint32_t EM = 0x0E000E00U; - Array vo; + Array vo; - // clang-format off + // clang-format off vo[0] = (x << 12 & S) | (x << 9 & EM); vo[1] = (x << 8 & S) | (x << 5 & EM); vo[2] = (x << 4 & S) | (x << 1 & EM); vo[3] = (x << 0 & S) | (x >> 3 & EM); - // clang-format on + // clang-format on - constexpr uint32_t e = (15U - 1U + 15U) << 10U; - constexpr uint32_t ee = e << 16U | e; + constexpr uint32_t e = (15U - 1U + 15U) << 10U; + constexpr uint32_t ee = e << 16U | e; - PRAGMA_UNROLL - for (int i = 0; i < 4; ++i) { - asm volatile("mul.f16x2 %0, %1, %2;" : "=r"(vo[i]) : "r"(vo[i]), "r"(ee)); - } + PRAGMA_UNROLL + for (int i = 0; i < 4; ++i) { + asm volatile("mul.f16x2 %0, %1, %2;" : "=r"(vo[i]) : "r"(vo[i]), "r"(ee)); + } - return (Array&)vo; - } + return (Array&)vo; + } - template - __device__ static auto convert(const Array& vi) { - Array vo; - PRAGMA_UNROLL - for (int i = 0; i < N; i += 8) { - auto& v = (Array&)vo[i]; - if constexpr (std::is_same_v) { - v = cvt_bf16x8_e2m1((Array&)vi[i]); - } else if constexpr (std::is_same_v) { - v = cvt_f16x8_e2m1((Array&)vi[i]); - } else { - static_assert(N != N, "not implemented"); - } - } - return vo; - } + template + __device__ static auto convert(const Array& vi) + { + Array vo; + PRAGMA_UNROLL + for (int i = 0; i < N; i += 8) { + auto& v = (Array&)vo[i]; + if constexpr (std::is_same_v) { + v = cvt_bf16x8_e2m1((Array&)vi[i]); + } + else if constexpr (std::is_same_v) { + v = cvt_f16x8_e2m1((Array&)vi[i]); + } + else { + static_assert(N != N, "not implemented"); + } + } + return vo; + } }; -__device__ inline Array cvt_bf16x4_e4m3( - const Array& vi) { - const uint32_t& x = (const uint32_t&)vi; +__device__ inline Array cvt_bf16x4_e4m3(const Array& vi) +{ + const uint32_t& x = (const uint32_t&)vi; - // 0 7 C 0 - // SEEEEEEEEMMMMMMMSEEEEEEEEMMMMMMM - // SEEEEMMM SEEEEMMM - // SEEEEMMM SEEEEMMM + // 0 7 C 0 + // SEEEEEEEEMMMMMMMSEEEEEEEEMMMMMMM + // SEEEEMMM SEEEEMMM + // SEEEEMMM SEEEEMMM - constexpr uint32_t S = 0x80008000U; - constexpr uint32_t EM = 0x07F007F0U; + constexpr uint32_t S = 0x80008000U; + constexpr uint32_t EM = 0x07F007F0U; - Array vo; + Array vo; - vo[0] = (x << 8 & S) | (x << 4 & EM); - vo[1] = (x << 0 & S) | (x >> 4 & EM); + vo[0] = (x << 8 & S) | (x << 4 & EM); + vo[1] = (x << 0 & S) | (x >> 4 & EM); - constexpr uint32_t e = (127U - 7U + 127U) << 7U; - constexpr uint32_t ee = e << 16U | e; + constexpr uint32_t e = (127U - 7U + 127U) << 7U; + constexpr uint32_t ee = e << 16U | e; - PRAGMA_UNROLL - for (int i = 0; i < 2; ++i) { + PRAGMA_UNROLL + for (int i = 0; i < 2; ++i) { #if TURBOMIND_ARCH_SM90 - asm("mul.rn.bf16x2 %0, %1, %2;" : "=r"(vo[i]) : "r"(vo[i]), "r"(ee)); + asm("mul.rn.bf16x2 %0, %1, %2;" : "=r"(vo[i]) : "r"(vo[i]), "r"(ee)); #else - asm("fma.rn.bf16x2 %0, %1, %2, %3;" - : "=r"(vo[i]) - : "r"(vo[i]), "r"(ee), "r"(0)); + asm("fma.rn.bf16x2 %0, %1, %2, %3;" : "=r"(vo[i]) : "r"(vo[i]), "r"(ee), "r"(0)); #endif - } + } - return (Array&)vo; + return (Array&)vo; } -template -__device__ inline Array cvt_f16x4_e4m3( - const Array& vi) { - const uint32_t& x = (const uint32_t&)vi; +__device__ inline Array cvt_f16x4_e4m3(const Array& vi) +{ + const uint32_t& x = (const uint32_t&)vi; - // 3 F 8 0 - // SEEEEEMMMMMMMMMMSEEEEEMMMMMMMMMM - // SEEEEMMM SEEEEMMM - // SEEEEMMM SEEEEMMM + // 3 F 8 0 + // SEEEEEMMMMMMMMMMSEEEEEMMMMMMMMMM + // SEEEEMMM SEEEEMMM + // SEEEEMMM SEEEEMMM - constexpr uint32_t S = 0x80008000U; - constexpr uint32_t EM = 0x3F803F80U; + constexpr uint32_t S = 0x80008000U; + constexpr uint32_t EM = 0x3F803F80U; - Array vo; + Array vo; - vo[0] = (x << 8 & S) | (x << 7 & EM); - vo[1] = (x << 0 & S) | (x >> 1 & EM); + vo[0] = (x << 8 & S) | (x << 7 & EM); + vo[1] = (x << 0 & S) | (x >> 1 & EM); - constexpr uint32_t e = (15U - 7U + 15U) << 10U; - constexpr uint32_t ee = e << 16U | e; + constexpr uint32_t e = (15U - 7U + 15U) << 10U; + constexpr uint32_t ee = e << 16U | e; - if constexpr (ApplyExponentBias) { PRAGMA_UNROLL for (int i = 0; i < 2; ++i) { - asm("mul.rn.f16x2 %0, %1, %2;" : "=r"(vo[i]) : "r"(vo[i]), "r"(ee)); + asm("mul.rn.f16x2 %0, %1, %2;" : "=r"(vo[i]) : "r"(vo[i]), "r"(ee)); } - } - return (Array&)vo; + return (Array&)vo; } -template +template struct ConvertKvCache { - template - __device__ static auto convert(const Array& vi) { - static_assert(N % 4 == 0); - Array vo; - PRAGMA_UNROLL - for (int i = 0; i < N; i += 4) { - auto& v = (Array&)vo[i]; - if constexpr (std::is_same_v) { - v = cvt_bf16x4_e4m3((Array&)vi[i]); - } else if constexpr (std::is_same_v) { - v = cvt_f16x4_e4m3((Array&)vi[i]); - } else { - static_assert(N != N, "not implemented"); - } - } - return vo; - } + + template + __device__ static auto convert(const Array& vi) + { + static_assert(N % 4 == 0); + Array vo; + PRAGMA_UNROLL + for (int i = 0; i < N; i += 4) { + auto& v = (Array&)vo[i]; + if constexpr (std::is_same_v) { + v = cvt_bf16x4_e4m3((Array&)vi[i]); + } + else if constexpr (std::is_same_v) { + v = cvt_f16x4_e4m3((Array&)vi[i]); + } + else { + static_assert(N != N, "not implemented"); + } + } + return vo; + } }; -template -inline __device__ void StoreQuantParam(T* dst, Array src) { - Store(dst, src); +template +inline __device__ void StoreQuantParam(T* dst, Array src) +{ + Store(dst, src); } -template <> -inline __device__ void StoreQuantParam(half* dst, - Array src) { - if constexpr (kFuseU4F16Dequant) { - src[1] = src[1] - src[0] * __ushort_as_half(0x5400); - } - Store(dst, src); +template<> +inline __device__ void StoreQuantParam(half* dst, Array src) +{ + if constexpr (kFuseU4F16Dequant) { + src[1] = src[1] - src[0] * __ushort_as_half(0x5400); + } + Store(dst, src); } } // namespace turbomind diff --git a/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/arch/config_sm70_s884.h b/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/arch/config_sm70_s884.h index dc8e6c3011..5c7499df34 100644 --- a/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/arch/config_sm70_s884.h +++ b/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/arch/config_sm70_s884.h @@ -19,85 +19,117 @@ namespace turbomind::gemm::sm70_s884 { -template +template struct Sm70_s884 { - static_assert(A::SmemCopyAtom::K == B::SmemCopyAtom::K); - static constexpr int SMEM_M = A::SmemCopyAtom::M / A::SmemCopyAtom::kFragNum; - static constexpr int SMEM_N = B::SmemCopyAtom::M / B::SmemCopyAtom::kFragNum; - static constexpr int SMEM_K = A::SmemCopyAtom::K; - - static constexpr auto MODE_ = - group_axis >= 0 ? Striding::kBlocked : Striding::kFlat; - - static constexpr auto MODE_A = group_axis == 0 ? Striding::kIndexed : MODE_; - static constexpr auto MODE_B = group_axis == 1 ? Striding::kIndexed : MODE_; - static constexpr auto MODE_C = MODE_; - - template - struct Type { - // (TM, TN, TK) = R(MMA_Atom, SmemCopy_Atom) - using MMA_Atom = SM70_MMA_884; - - using Partition = Blocked; - using MMA_Map = - MMA_Map; - - using MMA = Tiled_MMA_v2; - - using IteratorA = - std::conditional_t, - IteratorSm70>; - using IteratorB = - std::conditional_t, - IteratorSm70>; - - using Mainloop = - MainloopSm70; // FusePrefetch_ - - static constexpr int CHUNK_K = - std::lcm(std::lcm(GroupSizeU, GroupSizeV), CTA_K); - - using Scheduler = SchedulerSm70; - - static constexpr int TILE_C_M = TILE_C_M_ == -1 ? CTA_M : TILE_C_M_; - static constexpr int TILE_C_N = TILE_C_N_ == -1 ? CTA_N : TILE_C_N_; - - using Epilogue = gemm::Epilogue_, - Operand_C, MODE_C, SplitK>; - - using Kernel = GemmUniversal; - }; + static_assert(A::SmemCopyAtom::K == B::SmemCopyAtom::K); + + static constexpr int SMEM_M = A::SmemCopyAtom::M / A::SmemCopyAtom::kFragNum; + static constexpr int SMEM_N = B::SmemCopyAtom::M / B::SmemCopyAtom::kFragNum; + static constexpr int SMEM_K = A::SmemCopyAtom::K; + + static constexpr auto MODE_ = group_axis >= 0 ? Striding::kBlocked : Striding::kFlat; + + static constexpr auto MODE_A = group_axis == 0 ? Striding::kIndexed : MODE_; + static constexpr auto MODE_B = group_axis == 1 ? Striding::kIndexed : MODE_; + static constexpr auto MODE_C = MODE_; + + template + struct Type { + + // (TM, TN, TK) = R(MMA_Atom, SmemCopy_Atom) + using MMA_Atom = SM70_MMA_884; + + using Partition = Blocked; + using MMA_Map = MMA_Map; + + using MMA = Tiled_MMA_v2; + + using Mainloop = MainloopSm70, + TransformA, + U, + GroupSizeU, + B, + IteratorSm70, + TransformB, + V, + GroupSizeV, + Stages, + true, + GmemLookahead>; // FusePrefetch_ + + static constexpr int CHUNK_K = std::lcm(std::lcm(GroupSizeU, GroupSizeV), CTA_K); + + using Scheduler = SchedulerSm70; + + static constexpr int TILE_C_M = TILE_C_M_ == -1 ? CTA_M : TILE_C_M_; + static constexpr int TILE_C_N = TILE_C_N_ == -1 ? CTA_N : TILE_C_N_; + + using Epilogue = gemm::Epilogue_, + Operand_C, + MODE_C, + SplitK>; + + using Kernel = GemmUniversal; + }; }; -template -using Config_U4_d = Sm70_s884< - typename GetOperand::Operand, - Transform_Default, VoidOperand, - typename GetOperand::Operand, - Transform_HMMA_SIMT_B, - typename GetOperand::Operand, - kRowMajor, half, raster_order, -1>; - -template -using Config_U4_d_A8x64Swizzle = Sm70_s884< - Operand_A_Swizzle_8x64, Transform_Default, VoidOperand, - typename GetOperand::Operand, - Transform_HMMA_SIMT_B, - typename GetOperand::Operand, - kRowMajor, half, raster_order, -1>; - -template +template +using Config_U4_d = Sm70_s884::Operand, + Transform_Default, + VoidOperand, + typename GetOperand::Operand, + Transform_HMMA_SIMT_B, + typename GetOperand::Operand, + kRowMajor, + half, + raster_order, + -1>; + +template +using Config_U4_d_A8x64Swizzle = Sm70_s884, + Transform_Default, + VoidOperand, + typename GetOperand::Operand, + Transform_HMMA_SIMT_B, + typename GetOperand::Operand, + kRowMajor, + half, + raster_order, + -1>; + +template using Config_U4_g = Sm70_s884, // A Transform_Default, // tarnsform A VoidOperand, // U @@ -106,9 +138,10 @@ using Config_U4_g = Sm70_s884, // A Operand_V_Pack, // V kRowMajor, // order_C half, // Tc - raster_order, 0>; + raster_order, + 0>; -template +template using Config_MXF4 = Sm70_s884, // A Transform_Default, // tarnsform A VoidOperand, // U @@ -117,9 +150,10 @@ using Config_MXF4 = Sm70_s884, // A Operand_V_Pack, // V kRowMajor, // order_C half, // Tc - raster_order, group_axis>; + raster_order, + group_axis>; -template +template using Config_NVF4 = Sm70_s884, // A Transform_Default, // tarnsform A VoidOperand, // U @@ -128,9 +162,10 @@ using Config_NVF4 = Sm70_s884, // A Operand_V_Pack, // V kRowMajor, // order_C half, // Tc - raster_order, group_axis>; + raster_order, + group_axis>; -template +template using Config_E4M3 = Sm70_s884, // A Transform_Default, // tarnsform A VoidOperand, // U @@ -139,21 +174,10 @@ using Config_E4M3 = Sm70_s884, // A Operand_V_Pack, // V kRowMajor, // order_C half, // Tc - raster_order, group_axis>; - -template -using Config_E4M3_Prescaled = - Sm70_s884, // A - Transform_Default, // transform A - VoidOperand, // U - Operand_B_Pack, // B - Transform_HMMA_SIMT_B_PrescaledE4M3, - Operand_V_Pack, // pre-scaled V - kRowMajor, // order_C - half, // Tc - raster_order, group_axis>; - -template + raster_order, + group_axis>; + +template using Config_F16 = Sm70_s884, // A Transform_Default, // tarnsform A VoidOperand, // U @@ -162,6 +186,7 @@ using Config_F16 = Sm70_s884, // A VoidOperand, // V kRowMajor, // order_C half, // Tc - raster_order, group_axis>; + raster_order, + group_axis>; } // namespace turbomind::gemm::sm70_s884 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 4c776b38ab..82f5c511b4 100644 --- a/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/gemm.cu +++ b/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/gemm.cu @@ -25,739 +25,653 @@ namespace turbomind::gemm { -void ExportDispatchCache( - std::ostream& os, - const std::vector>& entries); +void ExportDispatchCache(std::ostream& os, const std::vector>& entries); -void ImportDispatchCache(std::istream& is, +void ImportDispatchCache(std::istream& is, std::vector>& entries, - const std::vector>& kernels); + const std::vector>& kernels); namespace { -template -std::vector ArgSort(size_t size, const Cmp& cmp) { - std::vector idxs(size); - std::iota(idxs.begin(), idxs.end(), 0); - std::stable_sort(idxs.begin(), idxs.end(), cmp); - return idxs; +template +std::vector ArgSort(size_t size, const Cmp& cmp) +{ + std::vector idxs(size); + std::iota(idxs.begin(), idxs.end(), 0); + std::stable_sort(idxs.begin(), idxs.end(), cmp); + return idxs; } -bool GemmTraceEnabled() { - const char* raw = std::getenv("TM_GEMM_TRACE"); - return raw && std::atoi(raw) != 0; +bool GemmTraceEnabled() +{ + const char* raw = std::getenv("TM_GEMM_TRACE"); + return raw && std::atoi(raw) != 0; } -int GemmTraceLimit() { - const char* raw = std::getenv("TM_GEMM_TRACE_LIMIT"); - return raw ? std::max(std::atoi(raw), 0) : 256; +int GemmTraceLimit() +{ + const char* raw = std::getenv("TM_GEMM_TRACE_LIMIT"); + return raw ? std::max(std::atoi(raw), 0) : 256; } -bool GemmTraceFilterAllows(const std::string& desc) { - const char* raw = std::getenv("TM_GEMM_TRACE_FILTER"); - return !raw || !*raw || desc.find(raw) != std::string::npos; +bool GemmTraceFilterAllows(const std::string& desc) +{ + const char* raw = std::getenv("TM_GEMM_TRACE_FILTER"); + return !raw || !*raw || desc.find(raw) != std::string::npos; } -bool Sm70AwqTp2FastSelectorEnabled() { - const char* raw = std::getenv("VLLM_SM70_AWQ_TP2_FAST_SELECTOR"); - return !raw || std::atoi(raw) != 0; +bool Sm70AwqTp2FastSelectorEnabled() +{ + const char* raw = std::getenv("VLLM_SM70_AWQ_TP2_FAST_SELECTOR"); + return !raw || std::atoi(raw) != 0; } // Keep the accepted QKVZ route enabled by default, while allowing an isolated // end-to-end A/B comparison without disabling the other exact small-M routes. -bool Sm70AwqTp4QkvCta64Enabled() { - const char* raw = std::getenv("VLLM_SM70_AWQ_TP4_QKV_CTA64"); - return !raw || std::atoi(raw) != 0; +bool Sm70AwqTp4QkvCta64Enabled() +{ + const char* raw = std::getenv("VLLM_SM70_AWQ_TP4_QKV_CTA64"); + return !raw || std::atoi(raw) != 0; } // Default-on A/B gate for the exact TP4 MTP verifier M=5 routes. -bool Sm70AwqMtpM5FastSelectorEnabled() { - const char* raw = std::getenv("VLLM_SM70_AWQ_MTP_M5_FAST_SELECTOR"); - return !raw || std::atoi(raw) != 0; +bool Sm70AwqMtpM5FastSelectorEnabled() +{ + const char* raw = std::getenv("VLLM_SM70_AWQ_MTP_M5_FAST_SELECTOR"); + return !raw || std::atoi(raw) != 0; } -bool Sm70Mxfp4MoeGroupedM8FastSelectorEnabled() { - const char* grouped = std::getenv("VLLM_SM70_MXFP4_MOE_GROUPED_M8"); - if (!grouped || std::atoi(grouped) == 0) { - return false; - } - const char* raw = std::getenv("VLLM_SM70_MXFP4_MOE_GROUPED_M8_FAST_SELECTOR"); - return !raw || std::atoi(raw) != 0; -} - -bool Sm70Nvfp4Qwen38Tp4M1FastSelectorEnabled() { - const char* raw = std::getenv("VLLM_SM70_NVFP4_QWEN38_TP4_M1_FAST_SELECTOR"); - return !raw || std::atoi(raw) != 0; -} - -// Default-on gate for exact block-FP8 8K prefill GEMM descriptors. -bool Sm70Fp8BlockPrefillFastSelectorEnabled() { - const char* raw = std::getenv("VLLM_SM70_FP8_PREFILL_FAST_SELECTOR"); - return !raw || std::atoi(raw) != 0; -} - -bool Sm70Fp8GroupedBmmDecodeEnabled() { - const char* raw = std::getenv("VLLM_SM70_FP8_GROUPED_BMM_DECODE"); - return !raw || std::atoi(raw) != 0; +bool Sm70Mxfp4MoeGroupedM8FastSelectorEnabled() +{ + const char* grouped = std::getenv("VLLM_SM70_MXFP4_MOE_GROUPED_M8"); + if (!grouped || std::atoi(grouped) == 0) { + return false; + } + const char* raw = std::getenv("VLLM_SM70_MXFP4_MOE_GROUPED_M8_FAST_SELECTOR"); + return !raw || std::atoi(raw) != 0; } struct Sm70AwqTp2FastTarget { - int n; - int k; - int cta_m; - int cta_n; - int cta_k; - int splits; - int swizzle; - bool require_mgroup; - std::string name_contains; + int n; + int k; + int cta_m; + int cta_n; + int cta_k; + int splits; + int swizzle; + bool require_mgroup; + std::string name_contains; }; -std::optional GetSm70AwqTp2EnvFastTarget( - const GemmDesc& desc, const std::string_view desc_str) { - const char* raw = std::getenv("VLLM_SM70_AWQ_TP2_FAST_TARGETS"); - if (!raw || !*raw) { - return std::nullopt; - } - const std::string targets(raw); - size_t begin = 0; - while (begin <= targets.size()) { - const size_t end = targets.find(';', begin); - const auto entry = targets.substr( - begin, end == std::string::npos ? std::string::npos : end - begin); - const size_t sep = entry.find('|'); - if (sep != std::string::npos && entry.substr(0, sep) == desc_str) { - std::string spec = entry.substr(sep + 1); - std::string name_contains; - if (const size_t name_sep = spec.find('@'); - name_sep != std::string::npos) { - name_contains = spec.substr(name_sep + 1); - spec.resize(name_sep); - } - std::replace(spec.begin(), spec.end(), 'x', ','); - std::replace(spec.begin(), spec.end(), ':', ','); - for (char& ch : spec) { - if (ch == ',') { - ch = ' '; - } - } - std::istringstream is(spec); - int cta_m{}; - int cta_n{}; - int cta_k{}; - int splits{}; - int swizzle{}; - int require_mgroup{}; - if (is >> cta_m >> cta_n >> cta_k >> splits >> swizzle >> - require_mgroup) { - if (name_contains.empty()) { - is >> name_contains; +std::optional GetSm70AwqTp2EnvFastTarget(const GemmDesc& desc, + const std::string_view desc_str) +{ + const char* raw = std::getenv("VLLM_SM70_AWQ_TP2_FAST_TARGETS"); + if (!raw || !*raw) { + return std::nullopt; + } + const std::string targets(raw); + size_t begin = 0; + while (begin <= targets.size()) { + const size_t end = targets.find(';', begin); + const auto entry = targets.substr(begin, end == std::string::npos ? std::string::npos : end - begin); + const size_t sep = entry.find('|'); + if (sep != std::string::npos && entry.substr(0, sep) == desc_str) { + std::string spec = entry.substr(sep + 1); + std::string name_contains; + if (const size_t name_sep = spec.find('@'); name_sep != std::string::npos) { + name_contains = spec.substr(name_sep + 1); + spec.resize(name_sep); + } + std::replace(spec.begin(), spec.end(), 'x', ','); + std::replace(spec.begin(), spec.end(), ':', ','); + for (char& ch : spec) { + if (ch == ',') { + ch = ' '; + } + } + std::istringstream is(spec); + int cta_m{}; + int cta_n{}; + int cta_k{}; + int splits{}; + int swizzle{}; + int require_mgroup{}; + if (is >> cta_m >> cta_n >> cta_k >> splits >> swizzle >> require_mgroup) { + if (name_contains.empty()) { + is >> name_contains; + } + return Sm70AwqTp2FastTarget{desc.n, + desc.k, + cta_m, + cta_n, + cta_k, + splits, + swizzle, + require_mgroup != 0, + name_contains}; + } + if (GemmTraceEnabled() && GemmTraceFilterAllows(std::string(desc_str))) { + std::cerr << "[TM_GEMM_FAST_SELECTOR] desc=" << desc_str + << " stage=invalid_env_target entry=" << entry << std::endl; + } + return std::nullopt; } - return Sm70AwqTp2FastTarget{ - desc.n, desc.k, cta_m, cta_n, - cta_k, splits, swizzle, require_mgroup != 0, - name_contains}; - } - if (GemmTraceEnabled() && GemmTraceFilterAllows(std::string(desc_str))) { - std::cerr << "[TM_GEMM_FAST_SELECTOR] desc=" << desc_str - << " stage=invalid_env_target entry=" << entry << std::endl; - } - return std::nullopt; - } - if (end == std::string::npos) { - break; - } - begin = end + 1; - } - return std::nullopt; + if (end == std::string::npos) { + break; + } + begin = end + 1; + } + return std::nullopt; } -std::optional GetSm70AwqTp2FastTarget( - const GemmDesc& desc) { - const std::string desc_str = to_string(desc); - if (Sm70Nvfp4Qwen38Tp4M1FastSelectorEnabled()) { - // Qwen3.8-27B NVFP4 TP4 decode. The exact-MNK registry entries keep - // these small-N tactics out of prefill and unrelated model shapes. - if (desc_str == "sm70_f16_e2m1k16_f16_tnt_fff_1x8704x5120_1" || - desc_str == "sm70_f16_e2m1k16_f16_tnt_fff_1x5120x4352_1") { - return Sm70AwqTp2FastTarget{ - desc.n, desc.k, 8, 32, 64, 3, 4, true, "c8x32_a1x1x64_01"}; - } - } - if (Sm70Fp8BlockPrefillFastSelectorEnabled() && - (desc_str == "sm70_f16_e4m3k128_f16_tnt_fff_8000x4096x5120_1" || - desc_str == "sm70_f16_e4m3k128_f16_tnt_fff_8000x3584x5120_1")) { - return Sm70AwqTp2FastTarget{desc.n, desc.k, 64, 256, 16, 1, 3, true, ""}; - } - if (Sm70Fp8GroupedBmmDecodeEnabled() && - desc_str == "sm70_f16_e4m3k128_f16_tnt_bbb_2x1024x4096_1") { - // Match the accepted dense 1x1024x4096 accumulation tree while launching - // both independent WO-A groups through one blocked descriptor. - return Sm70AwqTp2FastTarget{ - desc.n, desc.k, 8, 128, 64, 7, 2, true, "c8x128_a1x1x64_01"}; - } - const bool awq_fast_selector_enabled = Sm70AwqTp2FastSelectorEnabled(); - if (awq_fast_selector_enabled) { - if (auto target = GetSm70AwqTp2EnvFastTarget(desc, desc_str)) { - return target; - } - } - const char* selector_rerank = std::getenv("VLLM_SM70_DFLASH2_QPN8_RERANK"); - const char* selector_shadow = - std::getenv("VLLM_SM70_DFLASH2_QPN8_RERANK_SHADOW"); - const bool exact_selector_rerank = - (selector_rerank && std::atoi(selector_rerank) != 0) || - (selector_shadow && std::atoi(selector_shadow) != 0); - if (exact_selector_rerank && desc.arch == 700 && desc.type_a == kHalf && - desc.type_b == kHalf && desc.type_c == kHalf && desc.m >= 1 && - desc.m <= 8 && desc.n == 62080 && desc.k == 5120 && desc.num == 1) { - return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 256, 64, - 10, 1, true, "s884_1x4x1"}; - } - if (!awq_fast_selector_enabled) { +std::optional GetSm70AwqTp2FastTarget(const GemmDesc& desc) +{ + const bool awq_fast_selector_enabled = Sm70AwqTp2FastSelectorEnabled(); + const std::string desc_str = to_string(desc); + if (awq_fast_selector_enabled) { + if (auto target = GetSm70AwqTp2EnvFastTarget(desc, desc_str)) { + return target; + } + } + if (!awq_fast_selector_enabled) { + return std::nullopt; + } + if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_5x17408x5120_1") { + return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 64, 64, 1, 0, false, ""}; + } + if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_5x8192x5120_1") { + return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 64, 64, 2, 0, false, ""}; + } + if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_5x5120x3072_1") { + return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 64, 64, 3, 0, false, ""}; + } + if (Sm70AwqMtpM5FastSelectorEnabled()) { + // TP4 MTP verifier, M=5. Both routes are bitwise exact on all TP shards. + if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_5x8704x5120_1") { + return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 64, 64, 2, 4, false, ""}; + } + if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_5x4096x5120_1") { + return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 64, 64, 4, 4, false, ""}; + } + } + if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_1x17408x5120_1") { + return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 256, 64, 3, 3, false, ""}; + } + if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_1x8704x5120_1") { + return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 64, 64, 2, 4, false, ""}; + } + if (Sm70AwqTp4QkvCta64Enabled() + && desc_str == "sm70_f16_u4k128_f16_tnt_fff_1x4096x5120_1") { + return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 64, 64, 4, 4, false, ""}; + } + if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_1x5120x1536_1") { + return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 256, 64, 12, 4, false, "s884_1x4x1"}; + } + if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_1x5120x3072_1") { + return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 256, 64, 7, 0, false, ""}; + } return std::nullopt; - } - if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_5x17408x5120_1") { - return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 64, 64, 1, 0, false, ""}; - } - if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_5x8192x5120_1") { - return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 64, 64, 2, 0, false, ""}; - } - if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_5x5120x3072_1") { - return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 64, 64, 3, 0, false, ""}; - } - if (Sm70AwqMtpM5FastSelectorEnabled()) { - // TP4 MTP verifier, M=5. Both routes are bitwise exact on all TP shards. - if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_5x8704x5120_1") { - return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 64, 64, 2, 4, false, ""}; - } - if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_5x4096x5120_1") { - return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 64, 64, 4, 4, false, ""}; - } - } - if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_1x17408x5120_1") { - return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 256, 64, 3, 3, false, ""}; - } - if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_1x8704x5120_1") { - return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 64, 64, 2, 4, false, ""}; - } - if (Sm70AwqTp4QkvCta64Enabled() && - desc_str == "sm70_f16_u4k128_f16_tnt_fff_1x4096x5120_1") { - return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 64, 64, 4, 4, false, ""}; - } - if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_1x5120x1536_1") { - return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 256, 64, - 12, 4, false, "s884_1x4x1"}; - } - if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_1x5120x3072_1") { - return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 256, 64, 7, 0, false, ""}; - } - return std::nullopt; } -std::optional GetSm70Mxfp4MoeGroupedM8FastTarget( - const GemmDesc& desc) { - if (!Sm70Mxfp4MoeGroupedM8FastSelectorEnabled()) { +std::optional GetSm70Mxfp4MoeGroupedM8FastTarget(const GemmDesc& desc) +{ + if (!Sm70Mxfp4MoeGroupedM8FastSelectorEnabled()) { + return std::nullopt; + } + const std::string desc_str = to_string(desc); + if (desc_str == "sm70_f16_e2m1k32_f16_tnt_bbb_48x512x4096_1") { + return Sm70AwqTp2FastTarget{ + desc.n, desc.k, 8, 128, 64, 5, 0, true, "c8x128_a1x1x64_01"}; + } + if (desc_str == "sm70_f16_e2m1k32_f16_tnt_bbb_48x4096x256_1") { + return Sm70AwqTp2FastTarget{ + desc.n, desc.k, 16, 128, 32, 1, 0, true, "c16x128_a1x1x32_01"}; + } return std::nullopt; - } - const std::string desc_str = to_string(desc); - if (desc_str == "sm70_f16_e2m1k32_f16_tnt_bbb_48x512x4096_1") { - return Sm70AwqTp2FastTarget{ - desc.n, desc.k, 8, 128, 64, 5, 0, true, "c8x128_a1x1x64_01"}; - } - if (desc_str == "sm70_f16_e2m1k32_f16_tnt_bbb_48x4096x256_1") { - return Sm70AwqTp2FastTarget{ - desc.n, desc.k, 16, 128, 32, 1, 0, true, "c16x128_a1x1x32_01"}; - } - return std::nullopt; } -std::optional GetSm70Fp8BlockPrefillPrescaledTarget( - const GemmDesc& desc) { - const std::string desc_str = to_string(desc); - if (desc_str == "sm70_f16_e4m3k128_f16_tnt_fff_8000x4096x5120_1" || - desc_str == "sm70_f16_e4m3k128_f16_tnt_fff_8000x3584x5120_1") { - 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; +bool MatchesSm70AwqTp2FastKernel(const Kernel& kernel, const Sm70AwqTp2FastTarget& target) +{ + const int3 cta = kernel.cta_tile_size(); + if (cta.x != target.cta_m || cta.y != target.cta_n || cta.z != target.cta_k) { + return false; + } + const std::string name = kernel.name(); + const bool is_mgroup = name.find("mgroup") != std::string::npos; + if (is_mgroup != target.require_mgroup) { + return false; + } + return target.name_contains.empty() || name.find(target.name_contains) != std::string::npos; } -bool MatchesSm70AwqTp2FastKernel(const Kernel& kernel, - const Sm70AwqTp2FastTarget& target) { - const int3 cta = kernel.cta_tile_size(); - if (cta.x != target.cta_m || cta.y != target.cta_n || cta.z != target.cta_k) { - return false; - } - const std::string name = kernel.name(); - const bool is_mgroup = name.find("mgroup") != std::string::npos; - if (is_mgroup != target.require_mgroup) { - return false; - } - return target.name_contains.empty() || - name.find(target.name_contains) != std::string::npos; +void MaybeTraceSm70AwqTp2FastSelector(const GemmDesc& desc, const char* stage, const LaunchSpec* spec = nullptr) +{ + if (!GemmTraceEnabled()) { + return; + } + const std::string desc_str = to_string(desc); + if (!GemmTraceFilterAllows(desc_str)) { + return; + } + std::cerr << "[TM_GEMM_FAST_SELECTOR] desc=" << desc_str << " stage=" << stage; + if (spec && spec->kernel) { + std::cerr << " kernel=" << spec->kernel->name() << " splits=" << spec->splits + << " swizzle=" << spec->swizzle; + } + std::cerr << std::endl; } -void MaybeTraceSm70AwqTp2FastSelector(const GemmDesc& desc, const char* stage, - const LaunchSpec* spec = nullptr) { - if (!GemmTraceEnabled()) { - return; - } - const std::string desc_str = to_string(desc); - if (!GemmTraceFilterAllows(desc_str)) { - return; - } - std::cerr << "[TM_GEMM_FAST_SELECTOR] desc=" << desc_str - << " stage=" << stage; - if (spec && spec->kernel) { - std::cerr << " kernel=" << spec->kernel->name() - << " splits=" << spec->splits << " swizzle=" << spec->swizzle; - } - std::cerr << std::endl; +std::optional SelectSm70AwqTp2FastSpec(Context& ctx, + const std::vector& specs, + const Sm70AwqTp2FastTarget& target, + size_t barriers_size, + size_t partials_size) +{ + for (const auto& spec : specs) { + if (!spec.kernel || !MatchesSm70AwqTp2FastKernel(*spec.kernel, target)) { + continue; + } + if (spec.splits != target.splits) { + continue; + } + auto selected = spec; + const auto& actual_desc = ctx.get_desc(*selected.kernel); + const int4 shape{actual_desc.m, actual_desc.n, actual_desc.k, actual_desc.num}; + if (target.swizzle > selected.kernel->GetMaxSwizzle(shape)) { + continue; + } + (void)barriers_size; + (void)partials_size; + selected.splits = target.splits; + selected.swizzle = target.swizzle; + MaybeTraceSm70AwqTp2FastSelector(ctx.desc(), "selected", &selected); + return selected; + } + MaybeTraceSm70AwqTp2FastSelector(ctx.desc(), "no_match"); + return std::nullopt; } -std::optional SelectSm70AwqTp2FastSpec( - Context& ctx, const std::vector& specs, - const Sm70AwqTp2FastTarget& target, size_t barriers_size, - size_t partials_size) { - for (const auto& spec : specs) { - if (!spec.kernel || !MatchesSm70AwqTp2FastKernel(*spec.kernel, target)) { - continue; - } - if (spec.splits != target.splits) { - continue; - } - auto selected = spec; - const auto& actual_desc = ctx.get_desc(*selected.kernel); - const int4 shape{actual_desc.m, actual_desc.n, actual_desc.k, - actual_desc.num}; - if (target.swizzle > selected.kernel->GetMaxSwizzle(shape)) { - continue; - } - (void)barriers_size; - (void)partials_size; - selected.splits = target.splits; - selected.swizzle = target.swizzle; - MaybeTraceSm70AwqTp2FastSelector(ctx.desc(), "selected", &selected); - return selected; - } - MaybeTraceSm70AwqTp2FastSelector(ctx.desc(), "no_match"); - return std::nullopt; +const char* ToString(DispatchPolicy policy) +{ + if ((policy & DispatchPolicy::kPreserveDefaultSplits) + || (policy & DispatchPolicy::kPreserveDefaultSplitCount) + || (policy & DispatchPolicy::kMxfp4MoeGroupedM8Fast)) { + static thread_local std::string text; + auto base = + static_cast((int)policy & ~(int)DispatchPolicy::kPreserveDefaultSplits + & ~(int)DispatchPolicy::kPreserveDefaultSplitCount + & ~(int)DispatchPolicy::kMxfp4MoeGroupedM8Fast); + text = std::string(ToString(base)); + if (policy & DispatchPolicy::kPreserveDefaultSplits) { + text += "|preserve_default_splits"; + } + if (policy & DispatchPolicy::kPreserveDefaultSplitCount) { + text += "|preserve_default_split_count"; + } + if (policy & DispatchPolicy::kMxfp4MoeGroupedM8Fast) { + text += "|mxfp4_moe_grouped_m8_fast"; + } + return text.c_str(); + } + switch (policy) { + case DispatchPolicy::kDefault: + return "default"; + case DispatchPolicy::kMeasure: + return "measure"; + case DispatchPolicy::kReuse: + return "reuse"; + case DispatchPolicy::kAppend: + return "append"; + default: + return "unknown"; + } } -const char* ToString(DispatchPolicy policy) { - if ((policy & DispatchPolicy::kPreserveDefaultSplits) || - (policy & DispatchPolicy::kPreserveDefaultSplitCount) || - (policy & DispatchPolicy::kMxfp4MoeGroupedM8Fast) || - (policy & DispatchPolicy::kSm70Fp8PrefillPrescaled)) { - static thread_local std::string text; - auto base = static_cast( - (int)policy & ~(int)DispatchPolicy::kPreserveDefaultSplits & - ~(int)DispatchPolicy::kPreserveDefaultSplitCount & - ~(int)DispatchPolicy::kMxfp4MoeGroupedM8Fast & - ~(int)DispatchPolicy::kSm70Fp8PrefillPrescaled); - text = std::string(ToString(base)); - if (policy & DispatchPolicy::kPreserveDefaultSplits) { - text += "|preserve_default_splits"; - } - if (policy & DispatchPolicy::kPreserveDefaultSplitCount) { - text += "|preserve_default_split_count"; - } - if (policy & DispatchPolicy::kMxfp4MoeGroupedM8Fast) { - text += "|mxfp4_moe_grouped_m8_fast"; - } - if (policy & DispatchPolicy::kSm70Fp8PrefillPrescaled) { - text += "|sm70_fp8_prefill_prescaled"; - } - return text.c_str(); - } - switch (policy) { - case DispatchPolicy::kDefault: - return "default"; - case DispatchPolicy::kMeasure: - return "measure"; - case DispatchPolicy::kReuse: - return "reuse"; - case DispatchPolicy::kAppend: - return "append"; - default: - return "unknown"; - } -} +void MaybeTraceGemmDispatch(const GemmDesc& desc, DispatchPolicy policy, const LaunchSpec& spec, bool measured) +{ + if (!GemmTraceEnabled()) { + return; + } + const std::string desc_str = to_string(desc); + if (!GemmTraceFilterAllows(desc_str)) { + return; + } + static std::atomic logged{0}; + const int limit = GemmTraceLimit(); + if (limit > 0 && logged.fetch_add(1, std::memory_order_relaxed) >= limit) { + return; + } -void MaybeTraceGemmDispatch(const GemmDesc& desc, DispatchPolicy policy, - const LaunchSpec& spec, bool measured) { - if (!GemmTraceEnabled()) { - return; - } - const std::string desc_str = to_string(desc); - if (!GemmTraceFilterAllows(desc_str)) { - return; - } - static std::atomic logged{0}; - const int limit = GemmTraceLimit(); - if (limit > 0 && logged.fetch_add(1, std::memory_order_relaxed) >= limit) { - return; - } - - std::cerr << "[TM_GEMM_TRACE] desc=" << desc_str - << " policy=" << ToString(policy) - << " measured=" << (measured ? 1 : 0); - if (spec.kernel) { - const int3 cta = spec.kernel->cta_tile_size(); - const int3 mma = spec.kernel->warp_tile_size(); - std::cerr << " kernel=" << spec.kernel->name() << " splits=" << spec.splits - << " swizzle=" << spec.swizzle << " cta=" << cta.x << "x" << cta.y - << "x" << cta.z << " mma=" << mma.x << "x" << mma.y << "x" - << mma.z << " stages=" << spec.kernel->stages() - << " smem=" << spec.kernel->smem_size() - << " regs=" << spec.kernel->info().attr.numRegs - << " max_ctas=" << spec.kernel->info().max_active_ctas; - } else { - std::cerr << " kernel="; - } - std::cerr << std::endl; + std::cerr << "[TM_GEMM_TRACE] desc=" << desc_str << " policy=" << ToString(policy) << " measured=" + << (measured ? 1 : 0); + if (spec.kernel) { + const int3 cta = spec.kernel->cta_tile_size(); + const int3 mma = spec.kernel->warp_tile_size(); + std::cerr << " kernel=" << spec.kernel->name() << " splits=" << spec.splits << " swizzle=" << spec.swizzle + << " cta=" << cta.x << "x" << cta.y << "x" << cta.z << " mma=" << mma.x << "x" << mma.y + << "x" << mma.z << " stages=" << spec.kernel->stages() << " smem=" << spec.kernel->smem_size(); + } + else { + std::cerr << " kernel="; + } + std::cerr << std::endl; } } // namespace struct Gemm::Impl { - Impl() - : props_{GetCudaDeviceProps()}, + + Impl(): + props_{GetCudaDeviceProps()}, arch_{props_->major * 100 + props_->minor * 10}, registry_{props_}, - cache_{registry_.kernels()}, - sm70_fp8_prefill_cache_{registry_.kernels()} { - if (arch_ == 700) { - // V100 decode is dominated by many tiny GEMM/GEMV problems. A - // broader search space consistently finds better launch specs than - // the generic defaults for these SM70 workloads. - tuning_.max_splits = 16; - tuning_.max_waves = 32; - tuning_.swizzle = {0, 1, 2, 3, 4}; - tuning_.top_k = 0; - tuning_.clusters = 0; - tuning_.min_iter = 2; - tuning_.max_iter = 20; - tuning_.max_time = 2.f; - } - if (auto str = std::getenv("TM_GEMM_TUNE")) { - try { - ParseTuningParams(tuning_, str); - } catch (...) { - std::cerr << "[Gemm2] Failed to parse `TM_GEMM_TUNE`, default value " - "will be used.\n"; - tuning_ = {}; - } - } - if (std::getenv("TM_GEMM_WARN_CACHE_MISS")) { - warn_cache_miss_ = true; - } - measurer_.emplace(CreateStoppingCriterion( - tuning_.min_iter, tuning_.max_iter, tuning_.max_time)); - } - - // find launch spec in dispatch cache, dispatch by heuristic on cache miss - 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) { - if (auto spec = sm70_fp8_prefill_cache_.Find(desc); - spec && is_feasible(*spec)) { - return *spec; - } - if (auto fast_target = GetSm70Fp8BlockPrefillPrescaledTarget(desc)) { - 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); - return *fast_spec; - } - } - return {}; - } - if (policy & DispatchPolicy::kMxfp4MoeGroupedM8Fast) { - if (auto fast_target = GetSm70Mxfp4MoeGroupedM8FastTarget(desc)) { - 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; - } - return {}; - } - } - if (policy & DispatchPolicy::kReuse) { - if (auto spec = cache_.LowerBound(desc); spec && is_feasible(*spec)) { - return *spec; - } - if (warn_cache_miss_) { - std::cerr << "Failed to find a feasible kernel in the cache, will " - "dispatch by heuristic: " - << to_string(ctx.desc()) << std::endl; - } - } - - if (auto spec = cache_.Find(desc); spec && is_feasible(*spec)) { - return *spec; - } - - const auto fast_target = GetSm70AwqTp2FastTarget(desc); - auto specs = - Find(ctx, barriers_size, partials_size, fast_target ? 0 : 1, false); - if (!specs.empty()) { - auto selected = specs.front(); - if (fast_target) { - if (auto fast_spec = SelectSm70AwqTp2FastSpec( - ctx, specs, *fast_target, barriers_size, partials_size)) { - selected = *fast_spec; - } - } - cache_.Insert(desc, selected); - return selected; - } - return {}; - } - - std::vector Find(Context& ctx, size_t barrier_size, - 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; + cache_{registry_.kernels()} { - std::vector tmp; - tmp.reserve(feasible.size()); - for (const auto& k : feasible) { - LaunchSpec spec{k}; - tmp.push_back(spec); - } - clusters = Cluster(tmp, ClusteringParam{false, true}); - } - std::vector proxies; - proxies.reserve(clusters.size()); - - for (const auto& c : clusters) { - proxies.push_back(c.front().kernel); - } - - std::vector> specs; - - PopulateParam param{}; - param.max_splits = tuning_.max_splits; - param.max_waves = tuning_.max_waves; - param.swizzle = tuning_.swizzle.at(0); - param.barriers_size = barrier_size; - param.partials_size = partials_size; - - for (int cluster_id = 0; cluster_id < (int)proxies.size(); ++cluster_id) { - auto& kernel = *proxies[cluster_id]; - - auto tmp = ctx.Populate(kernel, param); - for (const auto& s : tmp) { - specs.emplace_back(cluster_id, s); - } - } - - // std::cerr << "#kernel: " << kernels.size() << ", #cluster: " << - // clusters.size() - // << ", #metric: " << metrics.size() << "\n"; - - int64_t mio_max = 0; - int64_t mma_max = 0; - for (const auto& [_, s] : specs) { - auto& [mio, mma] = s.estimated; - mio_max = std::max(mio_max, mio); - mma_max = std::max(mma_max, mma); - } - std::vector mio_ratio; - std::vector mma_ratio; - std::vector avg_ratio; - for (const auto& [_, s] : specs) { - auto& [mio, mma] = s.estimated; - mio_ratio.push_back((float)mio / mio_max); - mma_ratio.push_back((float)mma / mma_max); - avg_ratio.push_back(.5 * (mio_ratio.back() + mma_ratio.back())); + if (arch_ == 700) { + // V100 decode is dominated by many tiny GEMM/GEMV problems. A + // broader search space consistently finds better launch specs than + // the generic defaults for these SM70 workloads. + tuning_.max_splits = 16; + tuning_.max_waves = 32; + tuning_.swizzle = {0, 1, 2, 3, 4}; + tuning_.top_k = 0; + tuning_.clusters = 0; + tuning_.min_iter = 2; + tuning_.max_iter = 20; + tuning_.max_time = 2.f; + } + if (auto str = std::getenv("TM_GEMM_TUNE")) { + try { + ParseTuningParams(tuning_, str); + } + catch (...) { + std::cerr << "[Gemm2] Failed to parse `TM_GEMM_TUNE`, default value will be used.\n"; + tuning_ = {}; + } + } + if (std::getenv("TM_GEMM_WARN_CACHE_MISS")) { + warn_cache_miss_ = true; + } + measurer_.emplace(CreateStoppingCriterion(tuning_.min_iter, tuning_.max_iter, tuning_.max_time)); } - auto idxs = ArgSort(specs.size(), [&](int i, int j) { // - return avg_ratio[i] < avg_ratio[j]; - }); - // for (const auto& i : idxs) { - // auto [cid, s, m] = metrics[i]; - // std::cout << clusters[cid].front().kernel->name() << " s" << s << " " - // << avg_ratio[i] << " " << mio_ratio[i] - // << " " << mma_ratio[i] << " " << m.mio_cost << " " << - // m.mma_cost << "\n"; - // } + // find launch spec in dispatch cache, dispatch by heuristic on cache miss + LaunchSpec Dispatch(Context& ctx, DispatchPolicy policy, size_t barriers_size, size_t partials_size) + { + const auto& desc = ctx.desc(); + const auto is_feasible = [&](const LaunchSpec& spec) { + return spec.kernel && spec.kernel->is_feasible(ctx.get_desc(*spec.kernel)); + }; + if (policy & DispatchPolicy::kMxfp4MoeGroupedM8Fast) { + if (auto fast_target = GetSm70Mxfp4MoeGroupedM8FastTarget(desc)) { + auto specs = Find(ctx, barriers_size, partials_size, 0); + if (auto fast_spec = SelectSm70AwqTp2FastSpec( + ctx, specs, *fast_target, barriers_size, partials_size)) { + return *fast_spec; + } + return {}; + } + } + if (policy & DispatchPolicy::kReuse) { + if (auto spec = cache_.LowerBound(desc); spec && is_feasible(*spec)) { + return *spec; + } + if (warn_cache_miss_) { + std::cerr << "Failed to find a feasible kernel in the cache, will dispatch by heuristic: " + << to_string(ctx.desc()) << std::endl; + } + } + + if (auto spec = cache_.Find(desc); spec && is_feasible(*spec)) { + return *spec; + } - top_k = top_k > 0 ? std::min(idxs.size(), top_k) : (int)idxs.size(); - std::vector ret; - ret.reserve(top_k); - for (int i = 0; i < top_k; ++i) { - const auto& [cluster_id, spec] = specs[idxs[i]]; - // Apply `splits` to all kernels in the cluster - for (const auto& s : clusters[cluster_id]) { - auto tmp = spec; - tmp.kernel = s.kernel; - ret.push_back(tmp); - } + const auto fast_target = GetSm70AwqTp2FastTarget(desc); + auto specs = Find(ctx, barriers_size, partials_size, fast_target ? 0 : 1); + if (!specs.empty()) { + auto selected = specs.front(); + if (fast_target) { + if (auto fast_spec = + SelectSm70AwqTp2FastSpec(ctx, specs, *fast_target, barriers_size, partials_size)) { + selected = *fast_spec; + } + } + cache_.Insert(desc, selected); + return selected; + } + return {}; } - return ret; - } + std::vector Find(Context& ctx, size_t barrier_size, size_t partials_size, int top_k) + { + std::vector feasible = ctx.Filter(registry_.kernels()); + + std::vector> clusters; + { + std::vector tmp; + tmp.reserve(feasible.size()); + for (const auto& k : feasible) { + LaunchSpec spec{k}; + tmp.push_back(spec); + } + clusters = Cluster(tmp, ClusteringParam{false, true}); + } + std::vector proxies; + proxies.reserve(clusters.size()); - template - int Measure(Context& ctx, size_t barriers_size, size_t partials_size, - int top_k, LaunchFunc launch_func, cudaStream_t st) { - // Early exit on exact match - if (cache_.Find(ctx.desc())) { - return 0; - } - // std::cerr << "GEMM: " << desc.m << "x" << desc.n << "x" << desc.k << - // "\n"; + for (const auto& c : clusters) { + proxies.push_back(c.front().kernel); + } + + std::vector> specs; + + PopulateParam param{}; + param.max_splits = tuning_.max_splits; + param.max_waves = tuning_.max_waves; + param.swizzle = tuning_.swizzle.at(0); + param.barriers_size = barrier_size; + param.partials_size = partials_size; + + for (int cluster_id = 0; cluster_id < (int)proxies.size(); ++cluster_id) { + auto& kernel = *proxies[cluster_id]; + + auto tmp = ctx.Populate(kernel, param); + for (const auto& s : tmp) { + specs.emplace_back(cluster_id, s); + } + } - const auto tmp = - Find(ctx, barriers_size, partials_size, tuning_.top_k, false); + // std::cerr << "#kernel: " << kernels.size() << ", #cluster: " << clusters.size() + // << ", #metric: " << metrics.size() << "\n"; - std::vector specs; - for (const auto& spec : tmp) { - // populate swizzle parameters - const auto swis = ctx.Swizzle(spec, tuning_.swizzle); - specs.insert(specs.end(), swis.begin(), swis.end()); + int64_t mio_max = 0; + int64_t mma_max = 0; + for (const auto& [_, s] : specs) { + auto& [mio, mma] = s.estimated; + mio_max = std::max(mio_max, mio); + mma_max = std::max(mma_max, mma); + } + std::vector mio_ratio; + std::vector mma_ratio; + std::vector avg_ratio; + for (const auto& [_, s] : specs) { + auto& [mio, mma] = s.estimated; + mio_ratio.push_back((float)mio / mio_max); + mma_ratio.push_back((float)mma / mma_max); + avg_ratio.push_back(.5 * (mio_ratio.back() + mma_ratio.back())); + } + auto idxs = ArgSort(specs.size(), [&](int i, int j) { // + return avg_ratio[i] < avg_ratio[j]; + }); + + // for (const auto& i : idxs) { + // auto [cid, s, m] = metrics[i]; + // std::cout << clusters[cid].front().kernel->name() << " s" << s << " " << avg_ratio[i] << " " << + // mio_ratio[i] + // << " " << mma_ratio[i] << " " << m.mio_cost << " " << m.mma_cost << "\n"; + // } + + top_k = top_k > 0 ? std::min(idxs.size(), top_k) : (int)idxs.size(); + std::vector ret; + ret.reserve(top_k); + for (int i = 0; i < top_k; ++i) { + const auto& [cluster_id, spec] = specs[idxs[i]]; + // Apply `splits` to all kernels in the cluster + for (const auto& s : clusters[cluster_id]) { + auto tmp = spec; + tmp.kernel = s.kernel; + ret.push_back(tmp); + } + } + + return ret; } - specs = Sampler{*measurer_, tuning_.clusters}.Run(specs, launch_func, st); + template + int Measure( + Context& ctx, size_t barriers_size, size_t partials_size, int top_k, LaunchFunc launch_func, cudaStream_t st) + { + // Early exit on exact match + if (cache_.Find(ctx.desc())) { + return 0; + } + // std::cerr << "GEMM: " << desc.m << "x" << desc.n << "x" << desc.k << "\n"; + + const auto tmp = Find(ctx, barriers_size, partials_size, tuning_.top_k); - // for (const auto& s : specs) { - // std::cout << s.kernel->name() // - // << " swizzle=" << s.swizzle // - // << ", splits=" << s.splits // - // << ", measured=" << s.measured << "ms\n"; - // break; - // } + std::vector specs; + for (const auto& spec : tmp) { + // populate swizzle parameters + const auto swis = ctx.Swizzle(spec, tuning_.swizzle); + specs.insert(specs.end(), swis.begin(), swis.end()); + } + + specs = Sampler{*measurer_, tuning_.clusters}.Run(specs, launch_func, st); - if (!specs.empty()) { - cache_.Insert(ctx.desc(), specs.front()); - } else { - std::cerr << "No valid kernel found for the problem\n"; - return -1; + // for (const auto& s : specs) { + // std::cout << s.kernel->name() // + // << " swizzle=" << s.swizzle // + // << ", splits=" << s.splits // + // << ", measured=" << s.measured << "ms\n"; + // break; + // } + + if (!specs.empty()) { + cache_.Insert(ctx.desc(), specs.front()); + } + else { + std::cerr << "No valid kernel found for the problem\n"; + return -1; + } + + return 0; } - return 0; - } - - /// TODO: move to cuda utils - static std::unique_ptr GetCudaDeviceProps() { - auto props = std::make_unique(); - int device_id = -1; - cudaGetDevice(&device_id); - cudaGetDeviceProperties(props.get(), device_id); - return props; - } + /// TODO: move to cuda utils + static std::unique_ptr GetCudaDeviceProps() + { + auto props = std::make_unique(); + int device_id = -1; + cudaGetDevice(&device_id); + cudaGetDeviceProperties(props.get(), device_id); + return props; + } + + std::shared_ptr props_; + + int arch_; + + Registry registry_; - std::shared_ptr props_; + TuningParams tuning_; - int arch_; - - Registry registry_; - - TuningParams tuning_; - - bool warn_cache_miss_{}; - - std::optional measurer_; + bool warn_cache_miss_{}; - DispatchCache cache_; + std::optional measurer_; - DispatchCache sm70_fp8_prefill_cache_; + DispatchCache cache_; - std::mutex dispatch_mutex_; + std::mutex dispatch_mutex_; }; // implementation of GEMM interfaces -Gemm::Gemm() : impl_{new Impl{}} {} +Gemm::Gemm(): impl_{new Impl{}} {} Gemm::~Gemm() = default; -int Gemm::Run(const Operation& operation, float alpha, const void* A, - const MatrixLayout& Adesc, const void* U, - const MatrixLayout& Udesc, const void* B, - const MatrixLayout& Bdesc, const void* V, - const MatrixLayout& Vdesc, float beta, const void* C, - const MatrixLayout& Cdesc, void* D, const MatrixLayout& Ddesc, - const Workspace& workspace, cudaStream_t stream) { - Context context{*impl_->props_}; - - const auto desc = - context.Init(operation, Adesc, Udesc, Bdesc, Vdesc, Cdesc, Ddesc); - - if (!desc) { - fprintf(stderr, "invalid argument.\n"); - TM_CHECK(0); - return -1; - } - - const auto launch = [=](LaunchSpec spec, cudaStream_t st) { - auto _workspace = workspace; - return spec.kernel->Launch(operation, alpha, A, Adesc, U, Udesc, B, Bdesc, - V, Vdesc, beta, C, Cdesc, D, Ddesc, spec.swizzle, - spec.splits, _workspace, st); - }; - - std::optional dispatch_context_storage; - Context* dispatch_context = &context; - if (operation.dispatch_num_override > 0 && - operation.dispatch_num_override != context.desc().num) { - MatrixLayout dispatch_Adesc = Adesc; - MatrixLayout dispatch_Bdesc = Bdesc; - MatrixLayout dispatch_Ddesc = Ddesc; - dispatch_Adesc.num = operation.dispatch_num_override; - dispatch_Bdesc.num = operation.dispatch_num_override; - dispatch_Ddesc.num = operation.dispatch_num_override; - dispatch_context_storage.emplace(*impl_->props_); - if (dispatch_context_storage->Init(operation, dispatch_Adesc, Udesc, - dispatch_Bdesc, Vdesc, Cdesc, - dispatch_Ddesc)) { - dispatch_context = &*dispatch_context_storage; - } else { - dispatch_context_storage.reset(); - dispatch_context = &context; - } - } +int Gemm::Run(const Operation& operation, + float alpha, + const void* A, + const MatrixLayout& Adesc, + const void* U, + const MatrixLayout& Udesc, + const void* B, + const MatrixLayout& Bdesc, + const void* V, + const MatrixLayout& Vdesc, + float beta, + const void* C, + const MatrixLayout& Cdesc, + void* D, + const MatrixLayout& Ddesc, + const Workspace& workspace, + cudaStream_t stream) +{ + + Context context{*impl_->props_}; + + const auto desc = context.Init(operation, Adesc, Udesc, Bdesc, Vdesc, Cdesc, Ddesc); + + if (!desc) { + fprintf(stderr, "invalid argument.\n"); + TM_CHECK(0); + return -1; + } + + const auto launch = [=](LaunchSpec spec, cudaStream_t st) { + auto _workspace = workspace; + return spec.kernel->Launch(operation, + alpha, + A, + Adesc, + U, + Udesc, + B, + Bdesc, + V, + Vdesc, + beta, + C, + Cdesc, + D, + Ddesc, + spec.swizzle, + spec.splits, + _workspace, + st); + }; + + std::optional dispatch_context_storage; + Context* dispatch_context = &context; + if (operation.dispatch_num_override > 0 && operation.dispatch_num_override != context.desc().num) { + MatrixLayout dispatch_Adesc = Adesc; + MatrixLayout dispatch_Bdesc = Bdesc; + MatrixLayout dispatch_Ddesc = Ddesc; + dispatch_Adesc.num = operation.dispatch_num_override; + dispatch_Bdesc.num = operation.dispatch_num_override; + dispatch_Ddesc.num = operation.dispatch_num_override; + dispatch_context_storage.emplace(*impl_->props_); + if (dispatch_context_storage->Init( + operation, dispatch_Adesc, Udesc, dispatch_Bdesc, Vdesc, Cdesc, dispatch_Ddesc)) { + dispatch_context = &*dispatch_context_storage; + } + else { + dispatch_context_storage.reset(); + dispatch_context = &context; + } + } #if 0 if (operation.reserved) { @@ -773,70 +687,70 @@ int Gemm::Run(const Operation& operation, float alpha, const void* A, } #endif - LaunchSpec spec{}; - - const bool measured = operation.dispatch & DispatchPolicy::kMeasure; - { - std::lock_guard lock(impl_->dispatch_mutex_); - if (measured) { - impl_->Measure(*dispatch_context, workspace.barriers_size, - workspace.partials_size, 1, launch, stream); - } - - spec = impl_->Dispatch(*dispatch_context, operation.dispatch, - workspace.barriers_size, workspace.partials_size); - const bool preserve_default_kernel = - operation.dispatch & DispatchPolicy::kPreserveDefaultSplits; - const bool preserve_default_split_count = - operation.dispatch & DispatchPolicy::kPreserveDefaultSplitCount; - if (spec.kernel && - (preserve_default_kernel || preserve_default_split_count)) { - auto default_specs = - impl_->Find(*dispatch_context, workspace.barriers_size, - workspace.partials_size, 1, false); - if (!default_specs.empty()) { - const auto default_spec = default_specs.front(); - if (preserve_default_kernel) { - spec.kernel = default_spec.kernel; - } - spec.splits = default_spec.splits; - const auto& default_desc = dispatch_context->get_desc(*spec.kernel); - spec.swizzle = std::min(spec.swizzle, spec.kernel->GetMaxSwizzle({ - default_desc.m, - default_desc.n, - default_desc.k, - default_desc.num, - })); - } - } - } - if (spec.kernel && dispatch_context != &context) { - const auto& actual_desc = context.get_desc(*spec.kernel); - spec.swizzle = - std::min(spec.swizzle, - spec.kernel->GetMaxSwizzle({actual_desc.m, actual_desc.n, - actual_desc.k, actual_desc.num})); - } - MaybeTraceGemmDispatch(dispatch_context->desc(), operation.dispatch, spec, - measured); - - if (spec.kernel) { - // std::cout << "[Gemm] dispatch: " << spec.kernel->name() // - // << " split_k=" << spec.splits // - // << " swizzle=" << spec.swizzle << std::endl; - return launch(spec, stream); - } - - TM_CHECK(0) << "No feasible kernel found for the problem: " - << to_string(context.desc()); - - return -1; + LaunchSpec spec{}; + + const bool measured = operation.dispatch & DispatchPolicy::kMeasure; + { + std::lock_guard lock(impl_->dispatch_mutex_); + if (measured) { + impl_->Measure(*dispatch_context, workspace.barriers_size, workspace.partials_size, 1, launch, stream); + } + + spec = impl_->Dispatch(*dispatch_context, operation.dispatch, workspace.barriers_size, workspace.partials_size); + const bool preserve_default_kernel = operation.dispatch & DispatchPolicy::kPreserveDefaultSplits; + const bool preserve_default_split_count = operation.dispatch & DispatchPolicy::kPreserveDefaultSplitCount; + if (spec.kernel && (preserve_default_kernel || preserve_default_split_count)) { + auto default_specs = impl_->Find(*dispatch_context, workspace.barriers_size, workspace.partials_size, 1); + if (!default_specs.empty()) { + const auto default_spec = default_specs.front(); + if (preserve_default_kernel) { + spec.kernel = default_spec.kernel; + } + spec.splits = default_spec.splits; + const auto& default_desc = dispatch_context->get_desc(*spec.kernel); + spec.swizzle = std::min( + spec.swizzle, + spec.kernel->GetMaxSwizzle({ + default_desc.m, + default_desc.n, + default_desc.k, + default_desc.num, + })); + } + } + } + if (spec.kernel && dispatch_context != &context) { + const auto& actual_desc = context.get_desc(*spec.kernel); + spec.swizzle = + std::min(spec.swizzle, spec.kernel->GetMaxSwizzle({actual_desc.m, actual_desc.n, actual_desc.k, actual_desc.num})); + } + MaybeTraceGemmDispatch(dispatch_context->desc(), operation.dispatch, spec, measured); + + if (spec.kernel) { + // std::cout << "[Gemm] dispatch: " << spec.kernel->name() // + // << " split_k=" << spec.splits // + // << " swizzle=" << spec.swizzle << std::endl; + return launch(spec, stream); + } + + TM_CHECK(0) << "No feasible kernel found for the problem: " << to_string(context.desc()); + + return -1; } -int Gemm::Export(std::ostream& os) { return impl_->cache_.Export(os); } +int Gemm::Export(std::ostream& os) +{ + return impl_->cache_.Export(os); +} -int Gemm::Import(std::istream& is) { return impl_->cache_.Import(is); } +int Gemm::Import(std::istream& is) +{ + return impl_->cache_.Import(is); +} -std::vector Gemm::GetTuningSeq() const { return impl_->tuning_.seq; } +std::vector Gemm::GetTuningSeq() const +{ + return impl_->tuning_.seq; +} } // namespace turbomind::gemm diff --git a/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/iterator_sm70.h b/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/iterator_sm70.h index d7f4067595..fc52463377 100644 --- a/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/iterator_sm70.h +++ b/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/iterator_sm70.h @@ -16,256 +16,262 @@ namespace turbomind::gemm { -template -inline __device__ void _Ld(Array& dst, const T* src) { - static_assert(sizeof(Array) <= sizeof(uint4)); - - if constexpr (sizeof(Array) == sizeof(uint4)) { - (uint4&)dst = __ldcs((const uint4*)src); - } else if constexpr (sizeof(Array) == sizeof(uint2)) { - (uint2&)dst = __ldcs((const uint2*)src); - } else if constexpr (sizeof(Array) == sizeof(uint)) { - (uint&)dst = __ldcs((const uint*)src); - } else { - static_assert(!std::is_same_v); - } +template +inline __device__ void _Ld(Array& dst, const T* src) +{ + static_assert(sizeof(Array) <= sizeof(uint4)); + + if constexpr (sizeof(Array) == sizeof(uint4)) { + (uint4&)dst = __ldcs((const uint4*)src); + } + else if constexpr (sizeof(Array) == sizeof(uint2)) { + (uint2&)dst = __ldcs((const uint2*)src); + } + else if constexpr (sizeof(Array) == sizeof(uint)) { + (uint&)dst = __ldcs((const uint*)src); + } + else { + static_assert(!std::is_same_v); + } } -template +template struct GmemIteratorSm70 { - using ThreadMap = Map; - using AccessType = Array; - using Pointer = get_pointer_type; + using ThreadMap = Map; + + using AccessType = Array; + using Pointer = get_pointer_type; + + using Policy = Policy_; - using Policy = Policy_; + static constexpr int ITER_S = Map::kIterS; + static constexpr int ITER_C = Map::kIterC; - static constexpr int ITER_S = Map::kIterS; - static constexpr int ITER_C = Map::kIterC; + static constexpr Striding kMode = mode; + static constexpr bool is_indexed = mode == Striding::kIndexed; - static constexpr Striding kMode = mode; - static constexpr bool is_indexed = mode == Striding::kIndexed; + const char* src_data_; - const char* src_data_; + int src_offset_; + int dst_offset_; - int src_offset_; - int dst_offset_; + int offset_c_; + int offset_s_; - int offset_c_; - int offset_s_; + int src_step_c_; + int src_step_s_; - int src_step_c_; - int src_step_s_; + int src_step_k_; - int src_step_k_; + Predicate pred_; - Predicate - pred_; + bool g_mask{true}; - bool g_mask{true}; + SmemAccessor smem_data_; - SmemAccessor smem_data_; + static constexpr int2 kMK0 = cs2mk(SmemLayout::C0, SmemLayout::S0); + static constexpr int kPeriodC = ceil_div(SmemLayout::C0, Map::kDeltaC); + static constexpr int kPeriodS = ceil_div(SmemLayout::S0, Map::kDeltaS); - static constexpr int2 kMK0 = cs2mk(SmemLayout::C0, SmemLayout::S0); - static constexpr int kPeriodC = ceil_div(SmemLayout::C0, Map::kDeltaC); - static constexpr int kPeriodS = ceil_div(SmemLayout::S0, Map::kDeltaS); + int phases_[kPeriodS][kPeriodC]; - int phases_[kPeriodS][kPeriodC]; + const char* src_data_vec_[ITER_S]; - const char* src_data_vec_[ITER_S]; + using Fragments = AccessType[Map::kIterS][Map::kIterC]; - using Fragments = AccessType[Map::kIterS][Map::kIterC]; + __device__ static constexpr int2 pack(int2 mk) + { + return Packing_v2::apply(mk); + } - __device__ static constexpr int2 pack(int2 mk) { - return Packing_v2::apply(mk); - } + __device__ static constexpr int2 to_cs(int2 mk) + { + return mk2cs(mk.x, mk.y); + } - __device__ static constexpr int2 to_cs(int2 mk) { - return mk2cs(mk.x, mk.y); - } + __device__ GmemIteratorSm70(): smem_data_{Pointer{nullptr}} {}; - __device__ GmemIteratorSm70() : smem_data_{Pointer{nullptr}} {}; + __device__ GmemIteratorSm70(const MatrixData& mat, int2 offset, int2 extent): smem_data_{Pointer{(T*)nullptr}} + { + const int warp_id = threadIdx.x / WARP_SIZE; + const int lane_id = threadIdx.x % WARP_SIZE; - __device__ GmemIteratorSm70(const MatrixData& mat, int2 offset, int2 extent) - : smem_data_{Pointer{(T*)nullptr}} { - const int warp_id = threadIdx.x / WARP_SIZE; - const int lane_id = threadIdx.x % WARP_SIZE; + const Pointer data{(T*)mat.ptr.ptr}; + const int ld = mat.ptr.stride; - const Pointer data{(T*)mat.ptr.ptr}; - const int ld = mat.ptr.stride; + const int2 offsets = Map::get_offset(warp_id, lane_id); - const int2 offsets = Map::get_offset(warp_id, lane_id); + offset_c_ = offsets.x; + offset_s_ = offsets.y; - offset_c_ = offsets.x; - offset_s_ = offsets.y; + // auto src_ptr = reinterpret_cast((T*)data); - // auto src_ptr = reinterpret_cast((T*)data); + if constexpr (pred_.is_active) { + extent = to_cs(pack(extent)); + PRAGMA_UNROLL + for (int s = 0; s < Map::kIterS; ++s) { + PRAGMA_UNROLL + for (int c = 0; c < Map::kIterC; ++c) { + int ss = offset_s_ + s * Map::kDeltaS; + int cc = offset_c_ + c * Map::kDeltaC; + if (ss < extent.y && cc < extent.x) { + pred_.set(s, c); + } + } + } + } - if constexpr (pred_.is_active) { - extent = to_cs(pack(extent)); - PRAGMA_UNROLL - for (int s = 0; s < Map::kIterS; ++s) { PRAGMA_UNROLL - for (int c = 0; c < Map::kIterC; ++c) { - int ss = offset_s_ + s * Map::kDeltaS; - int cc = offset_c_ + c * Map::kDeltaC; - if (ss < extent.y && cc < extent.x) { - pred_.set(s, c); - } + for (int s = 0; s < kPeriodS; ++s) { + PRAGMA_UNROLL + for (int c = 0; c < kPeriodC; ++c) { + phases_[s][c] = SmemLayout::apply(offset_s_ + s * Map::kDeltaS, offset_c_ + c * Map::kDeltaC); + } } - } - } - PRAGMA_UNROLL - for (int s = 0; s < kPeriodS; ++s) { - PRAGMA_UNROLL - for (int c = 0; c < kPeriodC; ++c) { - phases_[s][c] = SmemLayout::apply(offset_s_ + s * Map::kDeltaS, - offset_c_ + c * Map::kDeltaC); - } - } + const int src_offset = is_indexed ? offsets.x : offsets.x + offsets.y * ld; - const int src_offset = is_indexed ? offsets.x : offsets.x + offsets.y * ld; + src_offset_ = src_offset * bitsof / bitsof; - src_offset_ = src_offset * bitsof / bitsof; + src_step_c_ = bitsof * Map::kDeltaC / bitsof; + src_step_s_ = bitsof * Map::kDeltaS * ld / bitsof; - src_step_c_ = bitsof * Map::kDeltaC / bitsof; - src_step_s_ = bitsof * Map::kDeltaS * ld / bitsof; + src_step_k_ = bitsof * cs2mk(Map::kDimC, Map::kDimS * ld).y / bitsof; - src_step_k_ = - bitsof * cs2mk(Map::kDimC, Map::kDimS * ld).y / bitsof; + // initialize for the first tile + if constexpr (is_indexed) { + const int2 cta_cs = to_cs(offset); + for (int s = 0; s < ITER_S; ++s) { + const int ss = cta_cs.y + offset_s_ + s * Map::kDeltaS; + const int idx = (mat.idxs && pred_(s, 0)) ? __ldg(mat.idxs + ss) : ss; + const auto tmp = data + cs2idx({cta_cs.x, idx}, ld); + src_data_vec_[s] = reinterpret_cast((T*)tmp) + src_offset_; + } + } + else { + auto src_data = data + cs2idx(to_cs(pack(offset)), ld); + src_data_ = reinterpret_cast((T*)src_data) + src_offset_; + } + } - // initialize for the first tile - if constexpr (is_indexed) { - const int2 cta_cs = to_cs(offset); - for (int s = 0; s < ITER_S; ++s) { - const int ss = cta_cs.y + offset_s_ + s * Map::kDeltaS; - const int idx = (mat.idxs && pred_(s, 0)) ? __ldg(mat.idxs + ss) : ss; - const auto tmp = data + cs2idx({cta_cs.x, idx}, ld); - src_data_vec_[s] = reinterpret_cast((T*)tmp) + src_offset_; - } - } else { - auto src_data = data + cs2idx(to_cs(pack(offset)), ld); - src_data_ = reinterpret_cast((T*)src_data) + src_offset_; + __device__ constexpr int _src_step_k() const + { + return src_step_k_; } - } - - __device__ constexpr int _src_step_k() const { return src_step_k_; } - - __device__ void ClearSmem(int pipe_iter = 0) { - PRAGMA_UNROLL - for (int s = 0; s < Map::kIterS; ++s) { - PRAGMA_UNROLL - for (int c = 0; c < Map::kIterC; ++c) { - const int pred_s = offset_s_ + s * Map::kDeltaS < Map::kDimS; - const int pred_c = offset_c_ + c * Map::kDeltaC < Map::kDimC; - auto ptr = &smem_data_(offset_s_ + s * Map::kDeltaS, - offset_c_ + c * Map::kDeltaC); - if ((Map::kAlignedC && Map::kAlignedS) || (pred_s && pred_c)) { - turbomind::Store(ptr, Array{}); + + __device__ void ClearSmem(int pipe_iter = 0) + { + PRAGMA_UNROLL + for (int s = 0; s < Map::kIterS; ++s) { + PRAGMA_UNROLL + for (int c = 0; c < Map::kIterC; ++c) { + const int pred_s = offset_s_ + s * Map::kDeltaS < Map::kDimS; + const int pred_c = offset_c_ + c * Map::kDeltaC < Map::kDimC; + auto ptr = &smem_data_(offset_s_ + s * Map::kDeltaS, offset_c_ + c * Map::kDeltaC); + if ((Map::kAlignedC && Map::kAlignedS) || (pred_s && pred_c)) { + turbomind::Store(ptr, Array{}); + } + } } - } } - } - __device__ void Advance() { - if constexpr (!is_indexed) { - if (!g_mask) { - src_data_ -= _src_step_k(); - } + __device__ void Advance() + { + if constexpr (!is_indexed) { + if (!g_mask) { + src_data_ -= _src_step_k(); + } + } } - } - - __device__ void Copy(std::true_type, T* dst, const char* __restrict__ src, - bool mask) { - if (mask) { - AccessType frag; - if constexpr (Policy_::kEvictPolicy != EvictPolicy::kEvictNormal) { - _Ld(frag, (const T*)src); - } else { - Ldg(frag, (const T*)src); - } - turbomind::Store(dst, frag); + + __device__ void Copy(std::true_type, T* dst, const char* __restrict__ src, bool mask) + { + if (mask) { + AccessType frag; + if constexpr (Policy_::kEvictPolicy != EvictPolicy::kEvictNormal) { + _Ld(frag, (const T*)src); + } + else { + Ldg(frag, (const T*)src); + } + turbomind::Store(dst, frag); + } } - } - - __device__ void Fetch(Fragments& frags, bool tile_mask) { - PRAGMA_UNROLL - for (int s = 0; s < Map::kIterS; ++s) { - if constexpr (is_indexed) { - src_data_ = src_data_vec_[s]; - } - - PRAGMA_UNROLL - for (int c = 0; c < Map::kIterC; ++c) { - Copy2(frags[s][c], src_data_ + src_step_c_ * c, - tile_mask && g_mask && pred_(s, c)); - } - - if constexpr (is_indexed) { - src_data_vec_[s] += _src_step_k(); - } else { - src_data_ += src_step_s_; - if (s == Map::kIterS - 1) { - src_data_ -= src_step_s_ * Map::kIterS; - src_data_ += _src_step_k(); + + __device__ void Fetch(Fragments& frags, bool tile_mask) + { + PRAGMA_UNROLL + for (int s = 0; s < Map::kIterS; ++s) { + + if constexpr (is_indexed) { + src_data_ = src_data_vec_[s]; + } + + PRAGMA_UNROLL + for (int c = 0; c < Map::kIterC; ++c) { + Copy2(frags[s][c], src_data_ + src_step_c_ * c, tile_mask && g_mask && pred_(s, c)); + } + + if constexpr (is_indexed) { + src_data_vec_[s] += _src_step_k(); + } + else { + src_data_ += src_step_s_; + if (s == Map::kIterS - 1) { + src_data_ -= src_step_s_ * Map::kIterS; + src_data_ += _src_step_k(); + } + } } - } } - } - - __device__ void Store(Fragments& frags) { - PRAGMA_UNROLL - for (int s = 0; s < Map::kIterS; ++s) { - PRAGMA_UNROLL - for (int c = 0; c < Map::kIterC; ++c) { - // auto dst = &smem_data_(offset_s_ + s * Map::kDeltaS, offset_c_ + c * - // Map::kDeltaC); - - const int i0 = SmemLayout::apply( // - s / kPeriodS * kPeriodS * Map::kDeltaS, - c / kPeriodC * kPeriodC * Map::kDeltaC); - const int i1 = phases_[s % kPeriodS][c % kPeriodC]; - auto dst = &smem_data_.ptr_[i0 + i1]; - - if (pred_(s, c)) { - turbomind::Store(dst, frags[s][c]); + + __device__ void Store(Fragments& frags) + { + PRAGMA_UNROLL + for (int s = 0; s < Map::kIterS; ++s) { + PRAGMA_UNROLL + for (int c = 0; c < Map::kIterC; ++c) { + // auto dst = &smem_data_(offset_s_ + s * Map::kDeltaS, offset_c_ + c * Map::kDeltaC); + + const int i0 = SmemLayout::apply( // + s / kPeriodS * kPeriodS * Map::kDeltaS, + c / kPeriodC * kPeriodC * Map::kDeltaC); + const int i1 = phases_[s % kPeriodS][c % kPeriodC]; + auto dst = &smem_data_.ptr_[i0 + i1]; + + if (pred_(s, c)) { + turbomind::Store(dst, frags[s][c]); + } + } } - } } - } - - __device__ void Copy2(AccessType& frag, const char* __restrict__ src, - bool mask) { - if (mask) { - if constexpr (Policy_::kEvictPolicy != EvictPolicy::kEvictNormal) { - _Ld(frag, (const T*)src); - } else { - Ldg(frag, (const T*)src); - } + + __device__ void Copy2(AccessType& frag, const char* __restrict__ src, bool mask) + { + if (mask) { + if constexpr (Policy_::kEvictPolicy != EvictPolicy::kEvictNormal) { + _Ld(frag, (const T*)src); + } + else { + Ldg(frag, (const T*)src); + } + } } - } }; -template +template struct IteratorSm70 { - template - using Type = GmemIteratorSm70; -}; - -// Exact-shape kernels may promise that every CTA covers a complete M/N/K tile. -// Keep ThreadMap's intrinsic bounds (for example the one-row scale tile), but -// remove runtime edge predicates from otherwise aligned A/B tiles. -template -struct IteratorSm70FullTile { - template - using Type = GmemIteratorSm70; + template + using Type = GmemIteratorSm70; }; } // namespace turbomind::gemm diff --git a/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/kernel/sm70_884_4.cu b/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/kernel/sm70_884_4.cu index ed4dce6012..01d4db3296 100644 --- a/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/kernel/sm70_884_4.cu +++ b/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/kernel/sm70_884_4.cu @@ -13,30 +13,32 @@ using D = cache_policy::Default; namespace { -// Keep shape-specific small-N tactics out of prefill and unrelated CUDA graph -// shapes. -template -class ExactMKernelImpl final : public KernelImpl { - public: - bool is_feasible(const GemmDesc& desc) const noexcept override { - return desc.m == ExactM && KernelImpl::is_feasible(desc); - } +// Keep shape-specific small-N tactics out of prefill and unrelated CUDA graph shapes. +template +class ExactMKernelImpl final: public KernelImpl { +public: + bool is_feasible(const GemmDesc& desc) const noexcept override + { + return desc.m == ExactM && KernelImpl::is_feasible(desc); + } }; -template -class ExactMnkKernelImpl final : public KernelImpl { - public: - bool is_feasible(const GemmDesc& desc) const noexcept override { - return desc.m == ExactM && desc.n == ExactN && desc.k == ExactK && - KernelImpl::is_feasible(desc); - } +template +class ExactMnkKernelImpl final: public KernelImpl { +public: + bool is_feasible(const GemmDesc& desc) const noexcept override + { + return desc.m == ExactM && desc.n == ExactN && desc.k == ExactK + && KernelImpl::is_feasible(desc); + } }; } // namespace -void Registry::sm70_884_4() { - if constexpr (1) { - // clang-format off +void Registry::sm70_884_4() +{ + if constexpr (1) { + // clang-format off using C = Config_U4_d; Add>(); Add>(); @@ -63,11 +65,11 @@ void Registry::sm70_884_4() { Add(std::make_unique>()); Add(std::make_unique>()); - // clang-format on - } + // clang-format on + } - if constexpr (1) { - // clang-format off + if constexpr (1) { + // clang-format off // GroupSizeV=128 using C = Config_U4_g; Add>(); @@ -96,38 +98,30 @@ void Registry::sm70_884_4() { Add>(); Add>(); Add>(); - // clang-format on - } + // clang-format on + } - if constexpr (1) { - // clang-format off + if constexpr (1) { + // clang-format off using C = Config_MXF4; Add>(); Add>(); Add>(); Add>(); Add>(); - // clang-format on - } + // clang-format on + } - if constexpr (1) { - // clang-format off + if constexpr (1) { + // clang-format off using C = Config_NVF4; Add>(); Add>(); Add>(); Add>(); Add>(); - using C32K64L1 = C::Type<8, 32, 64, 1, 1, 1, D, S, 2, true, 1, 16>; - using C32K64L2 = C::Type<8, 32, 64, 1, 1, 1, D, S, 2, true, 1, 16, -1, -1, 2>; - using C32K128L1 = C::Type<8, 32, 128, 1, 1, 1, D, S, 2, true, 1, 16>; - using C32K128L2 = C::Type<8, 32, 128, 1, 1, 1, D, S, 2, true, 1, 16, -1, -1, 2>; - Add(std::make_unique>()); - Add(std::make_unique>()); - Add(std::make_unique>()); - Add(std::make_unique>()); - // clang-format on - } + // clang-format on + } } } // namespace turbomind::gemm 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 1a1a0219c2..7ef3acf80b 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 @@ -11,48 +11,10 @@ using namespace cache_policy; using S = cache_policy::Stream; using D = cache_policy::Default; -namespace { - -template -class Qwen38PrescaledFullTileKernelImpl final : public KernelImpl { - public: - Qwen38PrescaledFullTileKernelImpl() { - this->info_.name += "_sm70_fp8_pscale_full"; - } - - bool is_feasible(const GemmDesc& desc) const noexcept override { - return desc.m == 8000 && (desc.n == 4096 || desc.n == 3584) && - desc.k == 5120 && desc.num == 1 && - KernelImpl::is_feasible(desc); - } -}; - -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() { - if constexpr (1) { - // clang-format off +void Registry::sm70_884_8() +{ + if constexpr (1) { + // clang-format off using C = Config_E4M3; Add>(); Add>(); @@ -67,15 +29,8 @@ void Registry::sm70_884_8() { Add>(); Add>(); Add>(); - - 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 - } + // clang-format on + } } } // namespace turbomind::gemm diff --git a/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/transform.h b/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/transform.h index 48aeaf4be5..1cbfda551a 100644 --- a/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/transform.h +++ b/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/transform.h @@ -13,181 +13,147 @@ namespace turbomind::gemm { struct Transform_Default { - template - __device__ static void apply(Array (&frag)[K][Mf], int k, - Array (&data)[K][Md], S&, int div) { - static_assert(Nf * Mf == Nd * Md); - static_assert(Nd % Nf == 0 && Mf % Md == 0); - static_assert(sizeof(frag) == sizeof(data)); - - // Alignment must be manually enforced for `reinterpret_cast` - auto& frag_k = reinterpret_cast(&)[Md]>(frag[k]); - auto& data_k = data[k]; - - PRAGMA_UNROLL - for (int i = 0; i < std::size(frag_k); ++i) { - frag_k[i] = data_k[i]; + template + __device__ static void apply(Array (&frag)[K][Mf], int k, Array (&data)[K][Md], S&, int div) + { + static_assert(Nf * Mf == Nd * Md); + static_assert(Nd % Nf == 0 && Mf % Md == 0); + static_assert(sizeof(frag) == sizeof(data)); + + // Alignment must be manually enforced for `reinterpret_cast` + auto& frag_k = reinterpret_cast(&)[Md]>(frag[k]); + auto& data_k = data[k]; + + PRAGMA_UNROLL + for (int i = 0; i < std::size(frag_k); ++i) { + frag_k[i] = data_k[i]; + } } - } }; -template +template struct Transform_HMMA_16816 { - template - __device__ static void apply(Array (&frag)[K][Mf], int k, - Array (&data)[K][Md], - Array (&stat)[Ks][Ms], int div) { - static_assert(Nf * Mf == Nd * Md); - static_assert(Nd % Nf == 0 && Mf % Md == 0); - static_assert(Nf * Mf == Ns * Ms * 4); - - auto& frag_k = reinterpret_cast(&)[Md]>(frag[k]); - auto& stat_k = reinterpret_cast(&)[Ns * Ms]>(stat[k / div]); - auto& data_k = data[k]; - - PRAGMA_UNROLL - for (int m = 0; m < Md; ++m) { - auto tmp = ConvertKvCache::convert(data_k[m]); - static_assert(Nd % 8 == 0); - PRAGMA_UNROLL - for (int i = 0; i < Nd; i += 8) { + template + __device__ static void + apply(Array (&frag)[K][Mf], int k, Array (&data)[K][Md], Array (&stat)[Ks][Ms], int div) + { + static_assert(Nf * Mf == Nd * Md); + static_assert(Nd % Nf == 0 && Mf % Md == 0); + static_assert(Nf * Mf == Ns * Ms * 4); + + auto& frag_k = reinterpret_cast(&)[Md]>(frag[k]); + auto& stat_k = reinterpret_cast(&)[Ns * Ms]>(stat[k / div]); + auto& data_k = data[k]; + PRAGMA_UNROLL - for (int s = 0; s < 2; ++s) { - PRAGMA_UNROLL - for (int c = 0; c < 2; ++c) { - const int idx = - (m * Nd + i) / 8 * 2 + s * StatStepS + c * StatStepC; - dequant((Array&)tmp[i + s * 4 + c * 2], stat_k[idx]); - } + for (int m = 0; m < Md; ++m) { + auto tmp = ConvertKvCache::convert(data_k[m]); + static_assert(Nd % 8 == 0); + PRAGMA_UNROLL + for (int i = 0; i < Nd; i += 8) { + PRAGMA_UNROLL + for (int s = 0; s < 2; ++s) { + PRAGMA_UNROLL + for (int c = 0; c < 2; ++c) { + const int idx = (m * Nd + i) / 8 * 2 + s * StatStepS + c * StatStepC; + dequant((Array&)tmp[i + s * 4 + c * 2], stat_k[idx]); + } + } + } + + frag_k[m] = tmp; } - } + } - frag_k[m] = tmp; + template + __device__ static void dequant(Array& x, Array s) + { + Array& _s = (Array&)s; + x[0] = __hfma(x[0], _s[0], _s[1]); + x[1] = __hfma(x[1], _s[0], _s[1]); + } + + __device__ static void dequant(Array& x, Array s) + { + bfloat16_t s1 = __ushort_as_bfloat16((uint16_t)s[0] << 7); + x[0] = __hmul(x[0], s1); + x[1] = __hmul(x[1], s1); + } + + __device__ static void dequant(Array& x, Array s) + { + // half_t s1 = __ushort_as_half(((uint16_t)s[0] + 15 - 127) << 10); + // Adjusted in `AdjustUe8m0ScaleForHalf` + half_t s1 = __ushort_as_half((uint16_t)s[0] << 10); + x[0] = __hmul(x[0], s1); + x[1] = __hmul(x[1], s1); + } + + __device__ static void dequant(Array& x, Array s) + { + auto s1 = __ushort_as_bfloat16(s[0]); + x[0] = __hmul(x[0], s1); + x[1] = __hmul(x[1], s1); + } + + __device__ static void dequant(Array& x, Array s) + { + auto s1 = __ushort_as_half(s[0]); + x[0] = __hmul(x[0], s1); + x[1] = __hmul(x[1], s1); } - } - - template - __device__ static void dequant(Array& x, Array s) { - Array& _s = (Array&)s; - x[0] = __hfma(x[0], _s[0], _s[1]); - x[1] = __hfma(x[1], _s[0], _s[1]); - } - - __device__ static void dequant(Array& x, Array s) { - bfloat16_t s1 = __ushort_as_bfloat16((uint16_t)s[0] << 7); - x[0] = __hmul(x[0], s1); - x[1] = __hmul(x[1], s1); - } - - __device__ static void dequant(Array& x, Array s) { - // half_t s1 = __ushort_as_half(((uint16_t)s[0] + 15 - 127) << 10); - // Adjusted in `AdjustUe8m0ScaleForHalf` - half_t s1 = __ushort_as_half((uint16_t)s[0] << 10); - x[0] = __hmul(x[0], s1); - x[1] = __hmul(x[1], s1); - } - - __device__ static void dequant(Array& x, - Array s) { - auto s1 = __ushort_as_bfloat16(s[0]); - x[0] = __hmul(x[0], s1); - x[1] = __hmul(x[1], s1); - } - - __device__ static void dequant(Array& x, Array s) { - auto s1 = __ushort_as_half(s[0]); - x[0] = __hmul(x[0], s1); - x[1] = __hmul(x[1], s1); - } }; // Used by SM70 MMA struct Transform_HMMA_SIMT_B { - template - __device__ static void apply(Array (&frag)[K][Mf], int k, - Array (&data)[K][Md], - Array (&stat)[Ks][Ms], int div) { - static_assert(Nf * Mf == Nd * Md); - static_assert(Nd % Nf == 0 && Mf % Md == 0); - - auto& frag_k = reinterpret_cast(&)[Md]>(frag[k]); - auto& stat_k = reinterpret_cast(&)[Ns * Ms]>(stat[k / div]); - auto& data_k = data[k]; - - // static_assert(Nf != Nf); - - PRAGMA_UNROLL - for (int m = 0; m < Md; ++m) { - auto tmp = ConvertKvCache::convert(data_k[m]); - PRAGMA_UNROLL - for (int i = 0; i < Nd; i += 2) { - dequant((Array&)tmp[i], stat_k[(m * Nd + i) / Nf]); - } - frag_k[m] = tmp; + template + __device__ static void + apply(Array (&frag)[K][Mf], int k, Array (&data)[K][Md], Array (&stat)[Ks][Ms], int div) + { + static_assert(Nf * Mf == Nd * Md); + static_assert(Nd % Nf == 0 && Mf % Md == 0); + + auto& frag_k = reinterpret_cast(&)[Md]>(frag[k]); + auto& stat_k = reinterpret_cast(&)[Ns * Ms]>(stat[k / div]); + auto& data_k = data[k]; + + // static_assert(Nf != Nf); + + PRAGMA_UNROLL + for (int m = 0; m < Md; ++m) { + auto tmp = ConvertKvCache::convert(data_k[m]); + PRAGMA_UNROLL + for (int i = 0; i < Nd; i += 2) { + dequant((Array&)tmp[i], stat_k[(m * Nd + i) / Nf]); + } + frag_k[m] = tmp; + } + } + + template + __device__ static void dequant(Array& x, Array s) + { + Array& _s = (Array&)s; + + x[0] = __hfma(x[0], _s[0], _s[1]); + x[1] = __hfma(x[1], _s[0], _s[1]); + } + + __device__ static void dequant(Array& x, Array s) + { + // half_t s1 = __ushort_as_half(((uint16_t)s[0] + 15 - 127) << 10); + // Adjusted in `AdjustUe8m0ScaleForHalf` + half_t s1 = __ushort_as_half((uint16_t)s[0] << 10); + x[0] = __hmul(x[0], s1); + x[1] = __hmul(x[1], s1); } - } - - template - __device__ static void dequant(Array& x, Array s) { - Array& _s = (Array&)s; - - x[0] = __hfma(x[0], _s[0], _s[1]); - x[1] = __hfma(x[1], _s[0], _s[1]); - } - - __device__ static void dequant(Array& x, Array s) { - // half_t s1 = __ushort_as_half(((uint16_t)s[0] + 15 - 127) << 10); - // Adjusted in `AdjustUe8m0ScaleForHalf` - half_t s1 = __ushort_as_half((uint16_t)s[0] << 10); - x[0] = __hmul(x[0], s1); - x[1] = __hmul(x[1], s1); - } - - __device__ static void dequant(Array& x, Array s) { - auto s1 = __ushort_as_half(s[0]); - x[0] = __hmul(x[0], s1); - x[1] = __hmul(x[1], s1); - } -}; -// FP8 scales that have absorbed the E4M3 exponent-bias factor (256) during -// one-time weight preparation. -struct Transform_HMMA_SIMT_B_PrescaledE4M3 { - template - __device__ static void apply(Array (&frag)[K][Mf], int k, - Array (&data)[K][Md], - Array (&stat)[Ks][Ms], int div) { - static_assert(std::is_same_v); - static_assert(std::is_same_v); - static_assert(std::is_same_v); - static_assert(Nf * Mf == Nd * Md); - static_assert(Nd % Nf == 0 && Mf % Md == 0); - static_assert(Nd % 4 == 0); - - auto& frag_k = reinterpret_cast(&)[Md]>(frag[k]); - auto& stat_k = reinterpret_cast(&)[Ns * Ms]>(stat[k / div]); - auto& data_k = data[k]; - - PRAGMA_UNROLL - for (int m = 0; m < Md; ++m) { - Array tmp; - PRAGMA_UNROLL - for (int i = 0; i < Nd; i += 4) { - auto& src = reinterpret_cast&>(data_k[m][i]); - reinterpret_cast&>(tmp[i]) = cvt_f16x4_e4m3(src); - } - PRAGMA_UNROLL - for (int i = 0; i < Nd; i += 2) { - auto scale = __ushort_as_half(stat_k[(m * Nd + i) / Nf][0]); - tmp[i] = __hmul(tmp[i], scale); - tmp[i + 1] = __hmul(tmp[i + 1], scale); - } - frag_k[m] = tmp; + __device__ static void dequant(Array& x, Array s) + { + auto s1 = __ushort_as_half(s[0]); + x[0] = __hmul(x[0], s1); + x[1] = __hmul(x[1], s1); } - } }; } // namespace turbomind::gemm diff --git a/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/types.h b/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/types.h index 475f623299..5d084456d5 100644 --- a/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/types.h +++ b/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/types.h @@ -7,244 +7,269 @@ #include #if ENABLE_BF16 - #include +#include #endif namespace turbomind::gemm { -enum class Order : int { - kColMajor = 0, - kRowMajor = 1, +enum class Order : int +{ + kColMajor = 0, + kRowMajor = 1, }; inline constexpr Order kColMajor = Order::kColMajor; inline constexpr Order kRowMajor = Order::kRowMajor; -constexpr Order operator~(Order a) { - return a == kColMajor ? kRowMajor : kColMajor; +constexpr Order operator~(Order a) +{ + return a == kColMajor ? kRowMajor : kColMajor; } -constexpr const char* to_string(Order order) { - switch (order) { - case kColMajor: - return "Col"; - case kRowMajor: - return "Row"; - } - return ""; +constexpr const char* to_string(Order order) +{ + switch (order) { + case kColMajor: + return "Col"; + case kRowMajor: + return "Row"; + } + return ""; } using Pack = uint32_t; -typedef enum MMA_Tag { - HMMA_16816 = 0x100, // sm80+ - HMMA_1688 = 0x200, // sm75 - HMMA_884 = 0x300, // sm70 - HMMA_SIMT = 0x400, // sm75- +typedef enum MMA_Tag +{ + HMMA_16816 = 0x100, // sm80+ + HMMA_1688 = 0x200, // sm75 + HMMA_884 = 0x300, // sm70 + HMMA_SIMT = 0x400, // sm75- } MMA_Tag; -typedef enum Op_Tag { - OPERAND_A = 0x010, - OPERAND_B = 0x020, - OPERAND_U = 0x030, - OPERAND_V = 0x040, - OPERAND_C = 0x050, - OPERAND_D = 0x060, +typedef enum Op_Tag +{ + OPERAND_A = 0x010, + OPERAND_B = 0x020, + OPERAND_U = 0x030, + OPERAND_V = 0x040, + OPERAND_C = 0x050, + OPERAND_D = 0x060, } Op_Tag; -constexpr MMA_Tag get_mma_tag(Pack pack) { - return static_cast(pack & 0xf00); +constexpr MMA_Tag get_mma_tag(Pack pack) +{ + return static_cast(pack & 0xf00); } -constexpr Op_Tag get_operand_tag(Pack pack) { - return static_cast(pack & 0x0f0); +constexpr Op_Tag get_operand_tag(Pack pack) +{ + return static_cast(pack & 0x0f0); } -constexpr int get_pack_num(Pack pack) { return pack & 0x00f; } +constexpr int get_pack_num(Pack pack) +{ + return pack & 0x00f; +} -enum class Striding : int { - kFlat, // [1111,2222,3333] - kRagged, // [11,2222222,333] [0 , 2 , 9 ] - kIndexed, // [xx xxxxxxx xxx], [01, 2345678, 9ab] - kBlocked, // [11][22222][333] +enum class Striding : int +{ + kFlat, // [1111,2222,3333] + kRagged, // [11,2222222,333] [0 , 2 , 9 ] + kIndexed, // [xx xxxxxxx xxx], [01, 2345678, 9ab] + kBlocked, // [11][22222][333] }; -inline const char* to_string(Striding striding) { - switch (striding) { - case Striding::kFlat: - return "f"; - case Striding::kRagged: - return "r"; - case Striding::kIndexed: - return "i"; - case Striding::kBlocked: - return "b"; - default: - return "unknown"; - } +inline const char* to_string(Striding striding) +{ + switch (striding) { + case Striding::kFlat: + return "f"; + case Striding::kRagged: + return "r"; + case Striding::kIndexed: + return "i"; + case Striding::kBlocked: + return "b"; + default: + return "unknown"; + } } -enum class QuantType : int { - kNone = 0, - kK = 1, - kM = 2, - kB = 3, - kDefault = kK, +enum class QuantType : int +{ + kNone = 0, + kK = 1, + kM = 2, + kB = 3, + kDefault = kK, }; -inline const char* to_string(QuantType q) { - switch (q) { - case QuantType::kNone: - return "none"; - case QuantType::kK: - return "k"; - case QuantType::kM: - return "m"; - case QuantType::kB: - return "b"; - default: - return "unknown"; - } +inline const char* to_string(QuantType q) +{ + switch (q) { + case QuantType::kNone: + return "none"; + case QuantType::kK: + return "k"; + case QuantType::kM: + return "m"; + case QuantType::kB: + return "b"; + default: + return "unknown"; + } } -enum class Epilogue : int { - kNone = 0, - kChannelCombination = 0x1, - kGatedSilu = 0x2, - kMoeWeightedReduce = 0x4, - kTileAllReduce = 0x8, +enum class Epilogue : int +{ + kNone = 0, + kChannelCombination = 0x1, + kGatedSilu = 0x2, + kMoeWeightedReduce = 0x4, + kTileAllReduce = 0x8, }; struct QuantDesc { - QuantType type; - int group_size; + QuantType type; + int group_size; - operator bool() const noexcept { return (int)type || group_size; } + operator bool() const noexcept + { + return (int)type || group_size; + } }; -inline std::string to_string(QuantDesc desc) { - if (desc) { - return to_string(desc.type) + std::to_string(desc.group_size); - } else { - return to_string(desc.type); - } +inline std::string to_string(QuantDesc desc) +{ + if (desc) { + return to_string(desc.type) + std::to_string(desc.group_size); + } + else { + return to_string(desc.type); + } } -enum class DispatchPolicy : int { - kDefault = 0, - kMeasure = 1, - kReuse = 2, - kAppend = 3, - kPreserveDefaultSplits = 4, - kPreserveDefaultSplitCount = 8, - kMxfp4MoeGroupedM8Fast = 16, - kSm70Fp8PrefillPrescaled = 32, +enum class DispatchPolicy : int +{ + kDefault = 0, + kMeasure = 1, + kReuse = 2, + kAppend = 3, + kPreserveDefaultSplits = 4, + kPreserveDefaultSplitCount = 8, + kMxfp4MoeGroupedM8Fast = 16, }; -constexpr bool operator&(const DispatchPolicy& a, const DispatchPolicy& b) { - return ((int)a & (int)b); +constexpr bool operator&(const DispatchPolicy& a, const DispatchPolicy& b) +{ + return ((int)a & (int)b); } -constexpr DispatchPolicy operator|(const DispatchPolicy& a, - const DispatchPolicy& b) { - return static_cast((int)a | (int)b); +constexpr DispatchPolicy operator|(const DispatchPolicy& a, const DispatchPolicy& b) +{ + return static_cast((int)a | (int)b); } class Kernel; class Context; struct Tape { - int ctas; - int max_num; - int max_ctas; - char* buffer; - int4* gemm_shapes; - int4* tiled_shapes; - int4* tile_offsets; - int2* iter_k_ranges; - int* tile_ids; + int ctas; + int max_num; + int max_ctas; + char* buffer; + int4* gemm_shapes; + int4* tiled_shapes; + int4* tile_offsets; + int2* iter_k_ranges; + int* tile_ids; }; struct Operation { - DispatchPolicy dispatch; - Epilogue epilogue; - QuantDesc quant_a; - QuantDesc quant_b; - int batch_dim; - int dispatch_num_override; - int active_group_count; - void* reserved; - void* tile_allreduce; + DispatchPolicy dispatch; + Epilogue epilogue; + QuantDesc quant_a; + QuantDesc quant_b; + int batch_dim; + int dispatch_num_override; + int active_group_count; + void* reserved; + void* tile_allreduce; }; struct MoeWeightedReduceParam { - void* out; - const float* sorted_weights; - const int* offsets; + void* out; + const float* sorted_weights; + const int* offsets; }; struct TileAllReduceParam { - void* signals[8]; - void* self_signal; - void* rank_data; - void* output; - int rank; - int world_size; - int tile_numel; - int output_numel; - int reducer_blocks; - int kernel_reducer_blocks; - int producer_grid_x; - int producer_grid_y; - int producer_grid_z; + void* signals[8]; + void* self_signal; + void* rank_data; + void* output; + int rank; + int world_size; + int tile_numel; + int output_numel; + int reducer_blocks; + int kernel_reducer_blocks; + int producer_grid_x; + int producer_grid_y; + int producer_grid_z; }; -inline Operation transpose(Operation o) { - std::swap(o.quant_a, o.quant_b); - o.batch_dim = 1 - o.batch_dim; - return o; +inline Operation transpose(Operation o) +{ + std::swap(o.quant_a, o.quant_b); + o.batch_dim = 1 - o.batch_dim; + return o; } struct MatrixLayout { - DataType type; - Order order; - int rows; - int cols; - int ld; - Pack pack; - int num; - int* offsets; - int* idxs; - int* group_idxs; + DataType type; + Order order; + int rows; + int cols; + int ld; + Pack pack; + int num; + int* offsets; + int* idxs; + int* group_idxs; }; -inline std::ostream& operator<<(std::ostream& os, const MatrixLayout& x) { - os << x.type << " " << to_string(x.order) << " " << x.rows << " " << x.cols - << " " << x.num << " " << x.ld; - return os; +inline std::ostream& operator<<(std::ostream& os, const MatrixLayout& x) +{ + os << x.type << " " << to_string(x.order) << " " << x.rows << " " << x.cols << " " << x.num << " " << x.ld; + return os; } -inline int64_t byte_size(const MatrixLayout& m) { - return byte_size(m.type, (int64_t)m.rows * m.cols); +inline int64_t byte_size(const MatrixLayout& m) +{ + return byte_size(m.type, (int64_t)m.rows * m.cols); } -inline Striding get_mode(const MatrixLayout& m) { - if (m.idxs) { - return Striding::kIndexed; - } else if (m.ld == 0 || m.offsets) { - return Striding::kBlocked; - } - return Striding::kFlat; +inline Striding get_mode(const MatrixLayout& m) +{ + if (m.idxs) { + return Striding::kIndexed; + } + else if (m.ld == 0 || m.offsets) { + return Striding::kBlocked; + } + return Striding::kFlat; } struct Workspace { - void* barriers; - size_t barriers_size; - void* partials; - size_t partials_size; - void* tensormaps; - size_t tensormaps_size; - int* flags; + void* barriers; + size_t barriers_size; + void* partials; + size_t partials_size; + void* tensormaps; + size_t tensormaps_size; + int* flags; }; } // namespace turbomind::gemm diff --git a/csrc/sm70_turbomind/ops/LICENSE.v100-skinny b/csrc/sm70_turbomind/ops/LICENSE.v100-skinny deleted file mode 100644 index 2649c53a63..0000000000 --- a/csrc/sm70_turbomind/ops/LICENSE.v100-skinny +++ /dev/null @@ -1,21 +0,0 @@ -MIT License - -Copyright (c) 2026 v100-skinny contributors - -Permission is hereby granted, free of charge, to any person obtaining a copy -of this software and associated documentation files (the "Software"), to deal -in the Software without restriction, including without limitation the rights -to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -copies of the Software, and to permit persons to whom the Software is -furnished to do so, subject to the following conditions: - -The above copyright notice and this permission notice shall be included in all -copies or substantial portions of the Software. - -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -SOFTWARE. diff --git a/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu b/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu index d555a8be7d..cd6bb22a70 100644 --- a/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu +++ b/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu @@ -37,7 +37,6 @@ #include "src/turbomind/kernels/gemm/types.h" #include "src/turbomind/kernels/gemm/utils.h" #include "custom_all_reduce.cuh" -#include "qwen38_prefill_cutlass.cuh" namespace turbomind { void unpack_awq_gemm(uint4_t* dst, const uint4_t* src, int rows, int cols, @@ -912,7 +911,6 @@ enum class TuneKeyKind : int { kMxfp4Dense = 6, kNvfp4Dense = 7, kMxfp4Moe = 8, - kNvfp4Moe = 9, }; struct DenseTuneKey { @@ -1048,9 +1046,6 @@ turbomind::gemm::DispatchPolicy select_fp8_moe_dispatch_policy( turbomind::gemm::DispatchPolicy select_mxfp4_moe_dispatch_policy( int device, int total_tokens, int n, int k, int num_experts, int group_size, cudaStream_t stream); -turbomind::gemm::DispatchPolicy select_nvfp4_moe_dispatch_policy( - int device, int total_tokens, int n, int k, int num_experts, int group_size, - cudaStream_t stream); bool tune_small_shapes_enabled() { const char* raw = std::getenv("VLLM_SM70_AWQ_TUNE_SMALL_SHAPES"); @@ -1074,12 +1069,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; -} - -bool mxfp4_moe_broadcast_input_decode_enabled() { - const char* raw = std::getenv("VLLM_SM70_MXFP4_MOE_BROADCAST_INPUT_DECODE"); - return raw == nullptr || std::atoi(raw) != 0; + return raw != nullptr && std::atoi(raw) != 0; } bool mxfp4_moe_grouped_m8_enabled() { @@ -1087,16 +1077,6 @@ bool mxfp4_moe_grouped_m8_enabled() { return raw != nullptr && std::atoi(raw) != 0; } -bool mxfp4_moe_grouped_verifier_enabled() { - const char* raw = std::getenv("VLLM_SM70_MXFP4_MOE_GROUPED_VERIFIER"); - return raw != nullptr && std::atoi(raw) != 0; -} - -bool mxfp4_moe_grouped_m8_expert_rows_enabled() { - const char* raw = std::getenv("VLLM_SM70_MXFP4_MOE_GROUPED_M8_EXPERT_ROWS"); - return raw != nullptr && std::atoi(raw) != 0; -} - bool mxfp4_moe_grouped_m8_fast_selector_enabled() { const char* raw = std::getenv("VLLM_SM70_MXFP4_MOE_GROUPED_M8_FAST_SELECTOR"); return raw == nullptr || std::atoi(raw) != 0; @@ -1107,16 +1087,6 @@ bool nvfp4_tune_small_shapes_enabled() { return raw == nullptr || std::atoi(raw) != 0; } -bool nvfp4_moe_grouped_prefill_enabled() { - const char* raw = std::getenv("VLLM_SM70_NVFP4_MOE_GROUPED_PREFILL"); - return raw == nullptr || std::atoi(raw) != 0; -} - -bool nvfp4_qwen38_tp4_m1_fast_selector_enabled() { - const char* raw = std::getenv("VLLM_SM70_NVFP4_QWEN38_TP4_M1_FAST_SELECTOR"); - return raw == nullptr || std::atoi(raw) != 0; -} - bool fp8_moe_single_token_per_expert_dispatch_enabled() { const char* raw = std::getenv("VLLM_SM70_FP8_MOE_SINGLE_TOKEN_PER_EXPERT_DISPATCH"); @@ -1133,11 +1103,6 @@ bool fp8_safe_fast_selector_enabled() { return raw != nullptr && std::atoi(raw) != 0; } -bool fp8_grouped_bmm_decode_enabled() { - const char* raw = std::getenv("VLLM_SM70_FP8_GROUPED_BMM_DECODE"); - return raw == nullptr || std::atoi(raw) != 0; -} - bool awq_reuse_imported_cache_enabled() { const char* raw = std::getenv("VLLM_SM70_AWQ_REUSE_IMPORTED_CACHE"); return raw != nullptr && std::atoi(raw) != 0; @@ -1249,11 +1214,6 @@ int moe_tune_max_tokens() { return raw ? std::max(std::atoi(raw), 0) : 128; } -int nvfp4_moe_tune_max_tokens() { - const char* raw = std::getenv("VLLM_SM70_NVFP4_MOE_TUNE_MAX_TOKENS"); - return raw ? std::max(std::atoi(raw), 0) : 128; -} - int sm70_f16_dense_max_m() { const char* raw = std::getenv("VLLM_SM70_F16_DENSE_MAX_M"); return raw ? std::max(std::atoi(raw), 0) : 64; @@ -1306,19 +1266,6 @@ turbomind::gemm::DispatchPolicy select_dense_dispatch_policy_impl( turbomind::gemm::DispatchPolicy select_dense_dispatch_policy( int device, int m, int n, int k, int group_size, cudaStream_t stream) { - const char* dflash2_rerank = std::getenv("VLLM_SM70_DFLASH2_QPN8_RERANK"); - const char* dflash2_shadow = - std::getenv("VLLM_SM70_DFLASH2_QPN8_RERANK_SHADOW"); - const bool exact_dflash2_rerank = - (dflash2_rerank && std::atoi(dflash2_rerank) != 0) || - (dflash2_shadow && std::atoi(dflash2_shadow) != 0); - if (exact_dflash2_rerank && m >= 1 && m <= 8 && n == 62080 && k == 5120 && - group_size == 0) { - // The sparse reranker reproduces this exact split-K contract. Do not let - // concurrent startup noise choose a numerically different LM-head spec on - // one TP rank. - return turbomind::gemm::DispatchPolicy::kDefault; - } return select_dense_dispatch_policy_impl( device, m, n, k, group_size, stream, TuneKeyKind::kGenericDense, tune_small_shapes_enabled(), false, generic_dense_tune_max_m()); @@ -1357,11 +1304,6 @@ turbomind::gemm::DispatchPolicy select_mxfp4_dense_dispatch_policy( turbomind::gemm::DispatchPolicy select_nvfp4_dense_dispatch_policy( int device, int m, int n, int k, int group_size, cudaStream_t stream) { - if (nvfp4_qwen38_tp4_m1_fast_selector_enabled() && m == 1 && - group_size == 16 && - ((n == 8704 && k == 5120) || (n == 5120 && k == 4352))) { - return turbomind::gemm::DispatchPolicy::kDefault; - } return select_dense_dispatch_policy_impl( device, m, n, k, group_size, stream, TuneKeyKind::kNvfp4Dense, nvfp4_tune_small_shapes_enabled(), false, nvfp4_dense_tune_max_m()); @@ -1369,11 +1311,8 @@ turbomind::gemm::DispatchPolicy select_nvfp4_dense_dispatch_policy( turbomind::gemm::DispatchPolicy select_moe_dispatch_policy_impl( int device, int total_tokens, int n, int k, int num_experts, int group_size, - cudaStream_t stream, TuneKeyKind kind, bool tune_enabled, - int max_tune_tokens = -1) { - const int tune_limit = - max_tune_tokens >= 0 ? max_tune_tokens : moe_tune_max_tokens(); - if (!tune_enabled || total_tokens > tune_limit) { + cudaStream_t stream, TuneKeyKind kind, bool tune_enabled) { + if (!tune_enabled || total_tokens > moe_tune_max_tokens()) { return turbomind::gemm::DispatchPolicy::kDefault; } @@ -1403,12 +1342,6 @@ turbomind::gemm::DispatchPolicy select_moe_dispatch_policy( turbomind::gemm::DispatchPolicy select_fp8_moe_dispatch_policy( int device, int total_tokens, int n, int k, int num_experts, int group_size, cudaStream_t stream) { - if (fp8_grouped_bmm_decode_enabled() && total_tokens == 2 && n == 1024 && - k == 4096 && num_experts == 2 && group_size == 128) { - // The matching dense WO-A projection uses the fixed launch spec selected - // in gemm.cu. Do not let measurement replace its accumulation tree. - return turbomind::gemm::DispatchPolicy::kDefault; - } return select_moe_dispatch_policy_impl( device, total_tokens, n, k, num_experts, group_size, stream, TuneKeyKind::kFp8Moe, fp8_tune_small_shapes_enabled()); @@ -1421,7 +1354,6 @@ turbomind::gemm::DispatchPolicy select_mxfp4_moe_dispatch_policy( total_tokens == 48 && num_experts == 48 && group_size == 32 && ((n == 512 && k == 4096) || (n == 4096 && k == 256)); if (mxfp4_moe_grouped_m8_enabled() && - !mxfp4_moe_grouped_m8_expert_rows_enabled() && mxfp4_moe_grouped_m8_fast_selector_enabled() && exact_grouped_m8) { return turbomind::gemm::DispatchPolicy::kMxfp4MoeGroupedM8Fast; } @@ -1430,15 +1362,6 @@ turbomind::gemm::DispatchPolicy select_mxfp4_moe_dispatch_policy( TuneKeyKind::kMxfp4Moe, mxfp4_tune_small_shapes_enabled()); } -turbomind::gemm::DispatchPolicy select_nvfp4_moe_dispatch_policy( - int device, int total_tokens, int n, int k, int num_experts, int group_size, - cudaStream_t stream) { - return select_moe_dispatch_policy_impl( - device, total_tokens, n, k, num_experts, group_size, stream, - TuneKeyKind::kNvfp4Moe, nvfp4_tune_small_shapes_enabled(), - nvfp4_moe_tune_max_tokens()); -} - WorkspaceHolder& get_workspace(int device, cudaStream_t stream) { thread_local int cached_device = -1; thread_local cudaStream_t cached_stream = nullptr; @@ -2461,7 +2384,8 @@ __global__ void fp8_sm70_dequantize_kernel( } const int words_per_n_tile = k / kFp8QuantValuesPerWord; - const int n_in_tile = static_cast(word_index % kFp8OutputColumnsPerTile); + const int n_in_tile = + static_cast(word_index % kFp8OutputColumnsPerTile); const int64_t tile_word = word_index / kFp8OutputColumnsPerTile; const int k_word = static_cast(tile_word % words_per_n_tile); const int n_tile = static_cast(tile_word / words_per_n_tile); @@ -2470,12 +2394,10 @@ __global__ void fp8_sm70_dequantize_kernel( const int group = k_base / group_size; const uint2 quant = reinterpret_cast(packed_weight)[word_index]; - const auto& quant_lo = - reinterpret_cast&>( - quant.x); - const auto& quant_hi = - reinterpret_cast&>( - quant.y); + const auto& quant_lo = reinterpret_cast< + const turbomind::Array&>(quant.x); + const auto& quant_hi = reinterpret_cast< + const turbomind::Array&>(quant.y); const auto half_lo = turbomind::cvt_f16x4_e4m3(quant_lo); const auto half_hi = turbomind::cvt_f16x4_e4m3(quant_hi); const __half scale = packed_scales[static_cast(group) * n + col]; @@ -2500,8 +2422,10 @@ __global__ void fp8_sm70_dequantize_kernel( __hmul(half_hi[3], scale); } -void fp8_sm70_dequantize_out(torch::Tensor output, torch::Tensor packed_weight, - torch::Tensor packed_scales, int64_t group_size) { +void fp8_sm70_dequantize_out(torch::Tensor output, + torch::Tensor packed_weight, + torch::Tensor packed_scales, + int64_t group_size) { TORCH_CHECK( output.is_cuda() && packed_weight.is_cuda() && packed_scales.is_cuda(), "SM70 FP8 prefill dequant expects CUDA tensors."); @@ -2521,7 +2445,8 @@ void fp8_sm70_dequantize_out(torch::Tensor output, torch::Tensor packed_weight, const at::cuda::OptionalCUDAGuard device_guard(device_of(output)); const int64_t k = output.size(0); const int64_t n = output.size(1); - TORCH_CHECK(k % group_size == 0 && k % kFp8QuantValuesPerWord == 0 && + TORCH_CHECK(k % group_size == 0 && + k % kFp8QuantValuesPerWord == 0 && n % kFp8OutputColumnsPerTile == 0, "SM70 FP8 prefill dequant shape alignment mismatch."); TORCH_CHECK(packed_weight.numel() == k * n, @@ -2534,14 +2459,15 @@ void fp8_sm70_dequantize_out(torch::Tensor output, torch::Tensor packed_weight, "SM70 FP8 prefill tensors must share a device."); const int64_t words = packed_weight.numel() / kFp8QuantValuesPerWord; - const int blocks = static_cast((words + kFp8PrefillDequantThreads - 1) / - kFp8PrefillDequantThreads); + const int blocks = static_cast( + (words + kFp8PrefillDequantThreads - 1) / kFp8PrefillDequantThreads); fp8_sm70_dequantize_kernel<<>>( reinterpret_cast<__half*>(output.data_ptr()), packed_weight.data_ptr(), reinterpret_cast(packed_scales.data_ptr()), - static_cast(k), static_cast(n), static_cast(group_size)); + static_cast(k), static_cast(n), + static_cast(group_size)); C10_CUDA_KERNEL_LAUNCH_CHECK(); } @@ -2876,13 +2802,10 @@ std::vector fp8_sm70_prepare(torch::Tensor qweight, const int64_t k = qweight.size(1); TORCH_CHECK(k % 8 == 0 && n % 8 == 0, "fp8_sm70_prepare: K and N must be multiples of 8."); - const bool channelwise_scales = scales.size(0) == n && scales.size(1) == 1; - const bool blockwise_scales = - scales.size(0) == (n + group_size - 1) / group_size && - scales.size(1) == (k + group_size - 1) / group_size; - TORCH_CHECK(channelwise_scales || blockwise_scales, - "fp8_sm70_prepare: expected [N, 1] channel scales or " - "[ceil(N/128), ceil(K/128)] block scales."); + TORCH_CHECK(scales.size(0) == (n + group_size - 1) / group_size, + "fp8_sm70_prepare: output scale block mismatch."); + TORCH_CHECK(scales.size(1) == (k + group_size - 1) / group_size, + "fp8_sm70_prepare: input scale block mismatch."); const auto converters = turbomind::gemm::GetConverters( turbomind::kHalf, turbomind::kFloat8_e4m3, turbomind::kHalf, true, 70); @@ -2946,21 +2869,12 @@ std::vector fp8_sm70_prepare(torch::Tensor qweight, const bool is_B_s = !is_A_s; const int64_t num_groups = (k + group_size - 1) / group_size; - torch::Tensor group_scales; - if (channelwise_scales) { - group_scales = scales.transpose(0, 1) - .contiguous() - .to(torch::kFloat16) - .repeat({num_groups, 1}) - .contiguous(); - } else { - group_scales = scales.transpose(0, 1) - .contiguous() - .to(torch::kFloat16) - .repeat_interleave(group_size, 1) - .slice(1, 0, n) - .contiguous(); - } + auto group_scales = scales.transpose(0, 1) + .contiguous() + .to(torch::kFloat16) + .repeat_interleave(group_size, 1) + .slice(1, 0, n) + .contiguous(); if (interleave_gated_silu) { group_scales = interleave_gated_silu_cols(group_scales); } @@ -3477,8 +3391,7 @@ void awq_gemm_sm70_out_tile_reduce( void fp8_gemm_sm70_out(torch::Tensor out, torch::Tensor in_feats, torch::Tensor tm_weight, torch::Tensor tm_scales, int64_t group_size, int64_t k_ld, int64_t q_ld, - bool gated_silu, - bool exact_8k_prefill_prescaled = false) { + bool gated_silu) { TORCH_CHECK(in_feats.is_cuda(), "fp8_gemm_sm70: input must be CUDA."); TORCH_CHECK(tm_weight.is_cuda(), "fp8_gemm_sm70: weight must be CUDA."); TORCH_CHECK(tm_scales.is_cuda(), "fp8_gemm_sm70: scales must be CUDA."); @@ -3509,19 +3422,6 @@ void fp8_gemm_sm70_out(torch::Tensor out, torch::Tensor in_feats, "fp8_gemm_sm70: output rows must match input rows."); TORCH_CHECK(out.stride(1) == 1, "fp8_gemm_sm70: output must be row-major contiguous."); - if (exact_8k_prefill_prescaled) { - 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, "fp8_gemm_sm70: gated_silu requires even output dim."); @@ -3602,10 +3502,6 @@ void fp8_gemm_sm70_out(torch::Tensor out, torch::Tensor in_feats, op.dispatch = select_fp8_dense_dispatch_policy( device, static_cast(m), static_cast(n), static_cast(k), static_cast(group_size), stream); - if (exact_8k_prefill_prescaled) { - op.dispatch = - op.dispatch | turbomind::gemm::DispatchPolicy::kSm70Fp8PrefillPrescaled; - } op.epilogue = gated_silu ? turbomind::gemm::Epilogue::kGatedSilu : turbomind::gemm::Epilogue::kNone; op.quant_a = {turbomind::gemm::QuantType::kNone, 0}; @@ -3622,13 +3518,6 @@ void fp8_gemm_sm70_out(torch::Tensor out, torch::Tensor in_feats, TORCH_CHECK(ec == 0, "fp8_gemm_sm70: TurboMind GEMM failed."); } -bool sm70_fp8_prefill_cutlass_gated_silu_enabled( - const torch::Tensor& in_feats, const torch::Tensor& dense_weight) { - const char* raw = std::getenv("VLLM_SM70_FP8_PREFILL_CUTLASS"); - return (raw == nullptr || std::atoi(raw) != 0) && in_feats.size(0) == 8000 && - in_feats.size(1) == 5120 && dense_weight.size(1) == 8704; -} - void fp8_gemm_sm70_prefill_dispatch_out( torch::Tensor out, int64_t dense_weight_ptr, torch::Tensor in_feats, torch::Tensor tm_weight, torch::Tensor tm_scales, int64_t group_size, @@ -3645,19 +3534,6 @@ void fp8_gemm_sm70_prefill_dispatch_out( reinterpret_cast(dense_weight_ptr), {in_feats.size(1), tm_weight.size(1)}, in_feats.options()); fp8_sm70_dequantize_out(dense_weight, tm_weight, tm_scales, group_size); - if (gated_silu && - sm70_fp8_prefill_cutlass_gated_silu_enabled(in_feats, dense_weight)) { - auto gate_up = - at::empty({in_feats.size(0), dense_weight.size(1)}, in_feats.options()); - if (sm70_fp8_prefill_cutlass_out(gate_up, in_feats, dense_weight, false)) { - sm70_silu_and_mul_interleaved_fp16_out(out, gate_up); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - return; - } - } - if (sm70_fp8_prefill_cutlass_out(out, in_feats, dense_weight, gated_silu)) { - return; - } if (gated_silu) { auto gate_up = at::mm(in_feats, dense_weight); sm70_silu_and_mul_interleaved_fp16_out(out, gate_up); @@ -4306,574 +4182,6 @@ void sm70_f16_gemm_out(torch::Tensor out, torch::Tensor in_feats, TORCH_CHECK(ec == 0, "sm70_f16_gemm: TurboMind GEMM failed."); } -template -__global__ void sm70_f16_gather_candidate_rows_kernel( - half* selected, const half* weight, const int64_t* candidate_ids, - int selected_rows, int hidden_size, int64_t weight_stride) { - const int row = blockIdx.x; - if (row >= selected_rows) { - return; - } - const int64_t token = __ldg(candidate_ids + row); - const half* src = weight + token * weight_stride; - half* dst = selected + static_cast(row) * hidden_size; - constexpr int kVectorHalf = sizeof(uint4) / sizeof(half); - const int vectors = hidden_size / kVectorHalf; - for (int idx = threadIdx.x; idx < vectors; idx += Threads) { - reinterpret_cast(dst)[idx] = - __ldg(reinterpret_cast(src) + idx); - } -} - -template -__global__ void sm70_f16_gather_packed_candidate_rows_kernel( - half* selected_packed, const half* packed_weight, - const int64_t* candidate_ids, int selected_rows, int hidden_size) { - // The SM70 HMMA 8x8x4 B layout packs every logical [32, 8] tile as - // eight consecutive uint4 vectors. Copy those vectors directly from the - // already prepared full LM-head layout instead of gathering 5.2 MiB of raw - // FP16 rows and running the generic layout converter every decode step. - constexpr int kRowsPerPack = 32; - constexpr int kColsPerVector = sizeof(uint4) / sizeof(half); - constexpr int kPackedTileElements = kRowsPerPack * kColsPerVector; - const int k_tiles = hidden_size / kColsPerVector; - const int vector_count = selected_rows * k_tiles; - for (int index = blockIdx.x * Threads + threadIdx.x; index < vector_count; - index += gridDim.x * Threads) { - const int selected_row = index / k_tiles; - const int k_tile = index - selected_row * k_tiles; - const int64_t token = __ldg(candidate_ids + selected_row); - const int64_t source = (token / kRowsPerPack) * hidden_size * kRowsPerPack + - static_cast(k_tile) * kPackedTileElements + - ((token % kRowsPerPack) / 4) * 4 * kColsPerVector + - (token % 4) * kColsPerVector; - const int64_t destination = - static_cast(selected_row / kRowsPerPack) * hidden_size * - kRowsPerPack + - static_cast(k_tile) * kPackedTileElements + - ((selected_row % kRowsPerPack) / 4) * 4 * kColsPerVector + - (selected_row % 4) * kColsPerVector; - *reinterpret_cast(selected_packed + destination) = - __ldg(reinterpret_cast(packed_weight + source)); - } -} - -template -__global__ void sm70_f16_extract_candidate_rows_kernel( - half* out, const half* expanded, int rows, int candidates, - int64_t expanded_stride) { - const int idx = blockIdx.x * Threads + threadIdx.x; - if (idx >= rows * candidates) { - return; - } - const int row = idx / candidates; - const int col = idx - row * candidates; - out[idx] = expanded[static_cast(row) * expanded_stride + - row * candidates + col]; -} - -template -__global__ void sm70_f16_rerank_keys_kernel(int64_t* keys, const half* logits, - const int64_t* candidate_ids, - int count) { - for (int index = blockIdx.x * Threads + threadIdx.x; index < count; - index += gridDim.x * Threads) { - const uint16_t raw = - __ldg(reinterpret_cast(logits) + index); - const uint16_t ordered = raw & 0x8000U - ? static_cast(~raw) - : static_cast(raw ^ 0x8000U); - // int64 top-k over this key is exactly lexicographic by FP16 value - // descending, then original vocabulary ID ascending. Candidate IDs are - // unique and below 2^32, so the key is collision-free. - const uint64_t tie = - 0xffffffffULL - static_cast(__ldg(candidate_ids + index)); - keys[index] = - static_cast((static_cast(ordered) << 32) | tie); - } -} - -__global__ void sm70_f16_rerank_topk_kernel(half* values_out, int64_t* ids_out, - const half* logits, - const int64_t* candidate_ids, - int rows, int top_k, - int64_t vocab_start_index) { - constexpr int kCandidates = 64; - const int row = blockIdx.x; - if (row >= rows) { - return; - } - - // One warp owns one candidate row. The composite key is identical to - // sm70_f16_rerank_keys_kernel: FP16 value descending, then original local - // vocabulary ID ascending. Keeping all 64 keys in shared memory lets the - // warp emit the first 16/20 order statistics without launching the generic - // top-k, gather, and ID-offset kernels. - __shared__ uint64_t candidate_keys[kCandidates]; - const int lane = threadIdx.x; - const int base = row * kCandidates; -#pragma unroll - for (int offset = 0; offset < kCandidates; offset += WARP_SIZE) { - const int position = lane + offset; - const uint16_t raw = - __ldg(reinterpret_cast(logits) + base + position); - const uint16_t ordered = raw & 0x8000U - ? static_cast(~raw) - : static_cast(raw ^ 0x8000U); - const uint64_t tie = - 0xffffffffULL - - static_cast(__ldg(candidate_ids + base + position)); - candidate_keys[position] = (static_cast(ordered) << 32) | tie; - } - __syncwarp(); - - for (int rank = 0; rank < top_k; ++rank) { - uint64_t best_key = candidate_keys[lane]; - int best_position = lane; - const uint64_t second_key = candidate_keys[lane + WARP_SIZE]; - if (second_key > best_key) { - best_key = second_key; - best_position = lane + WARP_SIZE; - } -#pragma unroll - for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) { - const uint64_t other_key = - __shfl_down_sync(0xffffffffU, best_key, offset); - const int other_position = - __shfl_down_sync(0xffffffffU, best_position, offset); - if (other_key > best_key) { - best_key = other_key; - best_position = other_position; - } - } - if (lane == 0) { - const uint16_t ordered = static_cast(best_key >> 32); - const uint16_t raw = ordered & 0x8000U - ? static_cast(ordered ^ 0x8000U) - : static_cast(~ordered); - values_out[row * top_k + rank] = __ushort_as_half(raw); - const uint32_t local_id = 0xffffffffU - static_cast(best_key); - ids_out[row * top_k + rank] = - vocab_start_index + static_cast(local_id); - // A valid local vocabulary ID is far below UINT32_MAX, so its key is - // never zero. Zero is therefore an exact in-warp consumed sentinel. - candidate_keys[best_position] = 0; - } - __syncwarp(); - } -} - -template -void launch_sm70_f16_indexed_rerank_gemm( - torch::Tensor out, torch::Tensor in_feats, torch::Tensor selected_packed, - torch::Tensor expanded, torch::Tensor partials, torch::Tensor barriers, - turbomind::gemm::MatrixLayout desc_B, int split_k, cudaStream_t stream) { - constexpr int kCandidates = 64; - constexpr int kThreads = 256; - const int m = static_cast(in_feats.size(0)); - const int k = static_cast(in_feats.size(1)); - const int selected_rows = m * kCandidates; - - turbomind::gemm::MatrixLayout desc_A{ - turbomind::kHalf, - turbomind::gemm::kRowMajor, - m, - k, - static_cast(in_feats.stride(0)), - }; - turbomind::gemm::MatrixLayout desc_U{}; - turbomind::gemm::MatrixLayout desc_V{}; - turbomind::gemm::MatrixLayout desc_D{ - turbomind::kHalf, turbomind::gemm::kRowMajor, m, - selected_rows, static_cast(expanded.stride(0)), - }; - - using namespace turbomind::gemm; - using C = sm70_s884::Config_F16; - using D = turbomind::cache_policy::Default; - using S = turbomind::cache_policy::Stream; - using Gemm = typename C::template Type<8, CTA_N, 64, 1, 4, 1, D, S, 2, true, - 1, 1>::Kernel; - using Sched = typename Gemm::Scheduler; - GemmParam param{ - to_param(in_feats.data_ptr(), desc_A), - to_param(selected_packed.data_ptr(), transpose(desc_B)), - to_param(nullptr, desc_U), - to_param(nullptr, desc_V), - }; - EpilogueParam epi_param{}; - epi_param.c = to_param(expanded.data_ptr(), desc_D); - turbomind::gemm::MatrixLayout desc_P = desc_D; - desc_P.ld = static_cast(partials.stride(0)); - epi_param.partials = to_param(partials.data_ptr(), desc_P); - epi_param.locks = barriers.data_ptr(); - epi_param.combine_mat = - MatrixCombination_v3{to_param(nullptr, MatrixLayout{}), 1.f, 0.f}; - Sched sched{{m, selected_rows, k, 1}, 0, split_k}; - sched.offsets_ = nullptr; - - const auto grid = sched.get_grid_shape(); - const auto block = Gemm::Impl::WARPS * WARP_SIZE; - constexpr int dynamic_smem_size = - static_cast(sizeof(typename Gemm::SharedStorage)); - auto kernel = gemm_kernel; - if constexpr (dynamic_smem_size > (48 << 10)) { - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, - dynamic_smem_size); - } - kernel<<>>(param, epi_param, sched); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - - const int output_count = m * kCandidates; - sm70_f16_extract_candidate_rows_kernel - <<<(output_count + kThreads - 1) / kThreads, kThreads, 0, stream>>>( - reinterpret_cast(out.data_ptr()), - reinterpret_cast(expanded.data_ptr()), m, - kCandidates, expanded.stride(0)); - C10_CUDA_KERNEL_LAUNCH_CHECK(); -} - -template -void launch_sm70_f16_indexed_rerank( - torch::Tensor out, torch::Tensor in_feats, torch::Tensor weight, - torch::Tensor candidate_ids, torch::Tensor selected_raw, - torch::Tensor selected_packed, torch::Tensor expanded, - torch::Tensor partials, torch::Tensor barriers, int split_k, - cudaStream_t stream) { - constexpr int kCandidates = 64; - constexpr int kThreads = 256; - const int m = static_cast(in_feats.size(0)); - const int k = static_cast(in_feats.size(1)); - const int selected_rows = m * kCandidates; - - sm70_f16_gather_candidate_rows_kernel - <<>>( - reinterpret_cast(selected_raw.data_ptr()), - reinterpret_cast(weight.data_ptr()), - candidate_ids.data_ptr(), selected_rows, k, - weight.stride(0)); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - - const auto converters = turbomind::gemm::GetConverters( - turbomind::kHalf, turbomind::kHalf, turbomind::kHalf, true, 70); - const auto* conv_w = converters[0]; - TORCH_CHECK(conv_w, - "sm70_f16_indexed_rerank_out: no compatible TurboMind " - "converter."); - const auto order_w = conv_w->order; - const bool is_A_w = turbomind::gemm::get_operand_tag(conv_w->pack) == - turbomind::gemm::OPERAND_A; - const bool is_B_w = !is_A_w; - turbomind::gemm::MatrixLayout w_desc{ - turbomind::kHalf, - order_w, - selected_rows, - k, - order_w == turbomind::gemm::kRowMajor ? k : selected_rows, - }; - if (is_B_w) { - std::swap(w_desc.rows, w_desc.cols); - w_desc.order = ~w_desc.order; - } - turbomind::gemm::MatrixLayout desc_B = w_desc; - desc_B.pack = conv_w->pack; - if (is_A_w) { - desc_B = turbomind::gemm::transpose(desc_B); - } - TORCH_CHECK(conv_w->Convert(selected_raw.data_ptr(), w_desc, - selected_packed.data_ptr(), desc_B, stream) == 0, - "sm70_f16_indexed_rerank_out: candidate pack failed."); - - launch_sm70_f16_indexed_rerank_gemm(out, in_feats, selected_packed, - expanded, partials, barriers, - desc_B, split_k, stream); -} - -template -void launch_sm70_f16_indexed_rerank_packed( - torch::Tensor out, torch::Tensor in_feats, torch::Tensor packed_weight, - torch::Tensor candidate_ids, torch::Tensor selected_packed, - torch::Tensor expanded, torch::Tensor partials, torch::Tensor barriers, - int split_k, cudaStream_t stream) { - constexpr int kCandidates = 64; - constexpr int kThreads = 256; - const int m = static_cast(in_feats.size(0)); - const int k = static_cast(in_feats.size(1)); - const int selected_rows = m * kCandidates; - const int vector_count = selected_rows * k / 8; - const int blocks = std::min(1024, (vector_count + kThreads - 1) / kThreads); - sm70_f16_gather_packed_candidate_rows_kernel - <<>>( - reinterpret_cast(selected_packed.data_ptr()), - reinterpret_cast(packed_weight.data_ptr()), - candidate_ids.data_ptr(), selected_rows, k); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - - turbomind::gemm::MatrixLayout desc_B{ - turbomind::kHalf, turbomind::gemm::kColMajor, k, selected_rows, k * 32, - }; - desc_B.pack = static_cast( - turbomind::gemm::HMMA_884 | turbomind::gemm::OPERAND_B | 1); - launch_sm70_f16_indexed_rerank_gemm(out, in_feats, selected_packed, - expanded, partials, barriers, - desc_B, split_k, stream); -} - -void sm70_f16_indexed_rerank_out(torch::Tensor out, torch::Tensor in_feats, - torch::Tensor weight, - torch::Tensor candidate_ids, - torch::Tensor selected_raw, - torch::Tensor selected_packed, - torch::Tensor expanded, torch::Tensor partials, - torch::Tensor barriers, int64_t cta_n, - int64_t split_k) { - constexpr int64_t kMaxRows = 8; - constexpr int64_t kCandidates = 64; - TORCH_CHECK(out.is_cuda() && in_feats.is_cuda() && weight.is_cuda() && - candidate_ids.is_cuda() && selected_raw.is_cuda() && - selected_packed.is_cuda() && expanded.is_cuda() && - partials.is_cuda() && barriers.is_cuda(), - "sm70_f16_indexed_rerank_out: tensors must be CUDA."); - TORCH_CHECK(out.scalar_type() == torch::kFloat16 && - in_feats.scalar_type() == torch::kFloat16 && - weight.scalar_type() == torch::kFloat16 && - selected_raw.scalar_type() == torch::kFloat16 && - selected_packed.scalar_type() == torch::kFloat16 && - expanded.scalar_type() == torch::kFloat16 && - partials.scalar_type() == torch::kFloat32, - "sm70_f16_indexed_rerank_out: data tensors must be float16."); - TORCH_CHECK(candidate_ids.scalar_type() == torch::kInt64, - "sm70_f16_indexed_rerank_out: candidate IDs must be int64."); - TORCH_CHECK(barriers.scalar_type() == torch::kInt32, - "sm70_f16_indexed_rerank_out: barriers must be int32."); - TORCH_CHECK(in_feats.dim() == 2 && weight.dim() == 2 && out.dim() == 2 && - candidate_ids.dim() == 2 && selected_raw.dim() == 2 && - selected_packed.dim() == 2 && expanded.dim() == 2 && - partials.dim() == 2, - "sm70_f16_indexed_rerank_out: tensors must be 2D."); - const int64_t m = in_feats.size(0); - const int64_t k = in_feats.size(1); - const int64_t selected_rows = m * kCandidates; - TORCH_CHECK(m >= 1 && m <= kMaxRows, - "sm70_f16_indexed_rerank_out: M must be in [1, 8]."); - TORCH_CHECK(k == weight.size(1) && k % 64 == 0, - "sm70_f16_indexed_rerank_out: invalid K dimension."); - TORCH_CHECK(out.size(0) == m && out.size(1) == kCandidates && - candidate_ids.size(0) >= m && - candidate_ids.size(1) == kCandidates, - "sm70_f16_indexed_rerank_out: expected [M, 64] output and " - "candidate IDs."); - TORCH_CHECK(selected_raw.size(0) >= selected_rows && - selected_raw.size(1) == k && - selected_packed.sizes() == selected_raw.sizes(), - "sm70_f16_indexed_rerank_out: candidate workspaces must be " - "[>=M*64, K]."); - TORCH_CHECK(expanded.size(0) >= m && expanded.size(1) >= selected_rows, - "sm70_f16_indexed_rerank_out: expanded workspace is too " - "small."); - TORCH_CHECK(partials.size(0) >= m && partials.size(1) >= selected_rows && - barriers.numel() >= 64, - "sm70_f16_indexed_rerank_out: split-K workspaces are too " - "small."); - TORCH_CHECK(out.is_contiguous() && in_feats.is_contiguous() && - weight.is_contiguous() && candidate_ids.is_contiguous() && - selected_raw.is_contiguous() && - selected_packed.is_contiguous() && expanded.is_contiguous() && - partials.is_contiguous() && barriers.is_contiguous(), - "sm70_f16_indexed_rerank_out: tensors must be contiguous."); - const int device = in_feats.get_device(); - TORCH_CHECK(out.get_device() == device && weight.get_device() == device && - candidate_ids.get_device() == device && - selected_raw.get_device() == device && - selected_packed.get_device() == device && - expanded.get_device() == device && - partials.get_device() == device && - barriers.get_device() == device, - "sm70_f16_indexed_rerank_out: tensors must share one device."); - TORCH_CHECK(cta_n == 128 || cta_n == 256, - "sm70_f16_indexed_rerank_out: cta_n must be 128 or 256."); - TORCH_CHECK(split_k >= 1 && split_k <= 16, - "sm70_f16_indexed_rerank_out: split_k must be in [1, 16]."); - - const at::cuda::OptionalCUDAGuard device_guard(device_of(in_feats)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - if (cta_n == 128) { - launch_sm70_f16_indexed_rerank<128>( - out, in_feats, weight, candidate_ids, selected_raw, selected_packed, - expanded, partials, barriers, static_cast(split_k), stream); - } else { - launch_sm70_f16_indexed_rerank<256>( - out, in_feats, weight, candidate_ids, selected_raw, selected_packed, - expanded, partials, barriers, static_cast(split_k), stream); - } -} - -void sm70_f16_indexed_rerank_packed_out( - torch::Tensor out, torch::Tensor in_feats, torch::Tensor packed_weight, - torch::Tensor candidate_ids, torch::Tensor selected_packed, - torch::Tensor expanded, torch::Tensor partials, torch::Tensor barriers, - int64_t cta_n, int64_t split_k) { - constexpr int64_t kMaxRows = 8; - constexpr int64_t kCandidates = 64; - TORCH_CHECK(out.is_cuda() && in_feats.is_cuda() && packed_weight.is_cuda() && - candidate_ids.is_cuda() && selected_packed.is_cuda() && - expanded.is_cuda() && partials.is_cuda() && - barriers.is_cuda(), - "sm70_f16_indexed_rerank_packed_out: tensors must be CUDA."); - TORCH_CHECK(out.scalar_type() == torch::kFloat16 && - in_feats.scalar_type() == torch::kFloat16 && - packed_weight.scalar_type() == torch::kFloat16 && - selected_packed.scalar_type() == torch::kFloat16 && - expanded.scalar_type() == torch::kFloat16 && - partials.scalar_type() == torch::kFloat32, - "sm70_f16_indexed_rerank_packed_out: data tensors must have " - "the expected FP16/FP32 dtypes."); - TORCH_CHECK(candidate_ids.scalar_type() == torch::kInt64 && - barriers.scalar_type() == torch::kInt32, - "sm70_f16_indexed_rerank_packed_out: IDs/barriers must be " - "int64/int32."); - TORCH_CHECK(in_feats.dim() == 2 && packed_weight.dim() == 2 && - out.dim() == 2 && candidate_ids.dim() == 2 && - selected_packed.dim() == 2 && expanded.dim() == 2 && - partials.dim() == 2, - "sm70_f16_indexed_rerank_packed_out: tensors must be 2D."); - const int64_t m = in_feats.size(0); - const int64_t k = in_feats.size(1); - const int64_t selected_rows = m * kCandidates; - TORCH_CHECK(m >= 1 && m <= kMaxRows, - "sm70_f16_indexed_rerank_packed_out: M must be in [1, 8]."); - TORCH_CHECK(k == packed_weight.size(1) && k % 64 == 0 && - packed_weight.size(0) % 32 == 0, - "sm70_f16_indexed_rerank_packed_out: invalid packed weight " - "shape."); - TORCH_CHECK(out.size(0) == m && out.size(1) == kCandidates && - candidate_ids.size(0) >= m && - candidate_ids.size(1) == kCandidates, - "sm70_f16_indexed_rerank_packed_out: expected [M, 64] output " - "and candidate IDs."); - TORCH_CHECK( - selected_packed.size(0) >= selected_rows && selected_packed.size(1) == k, - "sm70_f16_indexed_rerank_packed_out: packed workspace must be " - "[>=M*64, K]."); - TORCH_CHECK(expanded.size(0) >= m && expanded.size(1) >= selected_rows && - partials.size(0) >= m && partials.size(1) >= selected_rows && - barriers.numel() >= 64, - "sm70_f16_indexed_rerank_packed_out: split-K workspaces are " - "too small."); - TORCH_CHECK(out.is_contiguous() && in_feats.is_contiguous() && - packed_weight.is_contiguous() && - candidate_ids.is_contiguous() && - selected_packed.is_contiguous() && expanded.is_contiguous() && - partials.is_contiguous() && barriers.is_contiguous(), - "sm70_f16_indexed_rerank_packed_out: tensors must be " - "contiguous."); - const int device = in_feats.get_device(); - TORCH_CHECK( - out.get_device() == device && packed_weight.get_device() == device && - candidate_ids.get_device() == device && - selected_packed.get_device() == device && - expanded.get_device() == device && partials.get_device() == device && - barriers.get_device() == device, - "sm70_f16_indexed_rerank_packed_out: tensors must share one " - "device."); - TORCH_CHECK(cta_n == 128 || cta_n == 256, - "sm70_f16_indexed_rerank_packed_out: cta_n must be 128 or " - "256."); - TORCH_CHECK(split_k >= 1 && split_k <= 16, - "sm70_f16_indexed_rerank_packed_out: split_k must be in " - "[1, 16]."); - - const at::cuda::OptionalCUDAGuard device_guard(device_of(in_feats)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - if (cta_n == 128) { - launch_sm70_f16_indexed_rerank_packed<128>( - out, in_feats, packed_weight, candidate_ids, selected_packed, expanded, - partials, barriers, static_cast(split_k), stream); - } else { - launch_sm70_f16_indexed_rerank_packed<256>( - out, in_feats, packed_weight, candidate_ids, selected_packed, expanded, - partials, barriers, static_cast(split_k), stream); - } -} - -void sm70_f16_rerank_keys_out(torch::Tensor keys, torch::Tensor logits, - torch::Tensor candidate_ids) { - constexpr int64_t kCandidates = 64; - constexpr int kThreads = 128; - TORCH_CHECK(keys.is_cuda() && logits.is_cuda() && candidate_ids.is_cuda(), - "sm70_f16_rerank_keys_out: tensors must be CUDA."); - TORCH_CHECK(keys.scalar_type() == torch::kInt64 && - candidate_ids.scalar_type() == torch::kInt64 && - logits.scalar_type() == torch::kFloat16, - "sm70_f16_rerank_keys_out: expected int64/FP16 tensors."); - TORCH_CHECK(keys.sizes() == logits.sizes() && - keys.sizes() == candidate_ids.sizes() && keys.dim() == 2 && - keys.size(1) == kCandidates, - "sm70_f16_rerank_keys_out: expected matching [M, 64] " - "tensors."); - TORCH_CHECK(keys.is_contiguous() && logits.is_contiguous() && - candidate_ids.is_contiguous(), - "sm70_f16_rerank_keys_out: tensors must be contiguous."); - const int device = logits.get_device(); - TORCH_CHECK( - keys.get_device() == device && candidate_ids.get_device() == device, - "sm70_f16_rerank_keys_out: tensors must share one device."); - const at::cuda::OptionalCUDAGuard device_guard(device_of(logits)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - const int count = static_cast(keys.numel()); - sm70_f16_rerank_keys_kernel - <<<(count + kThreads - 1) / kThreads, kThreads, 0, stream>>>( - keys.data_ptr(), - reinterpret_cast(logits.data_ptr()), - candidate_ids.data_ptr(), count); - C10_CUDA_KERNEL_LAUNCH_CHECK(); -} - -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) { - constexpr int64_t kCandidates = 64; - constexpr int kThreads = 32; - TORCH_CHECK(values_out.is_cuda() && ids_out.is_cuda() && logits.is_cuda() && - candidate_ids.is_cuda(), - "sm70_f16_rerank_topk_out: tensors must be CUDA."); - TORCH_CHECK(values_out.scalar_type() == torch::kFloat16 && - ids_out.scalar_type() == torch::kInt64 && - logits.scalar_type() == torch::kFloat16 && - candidate_ids.scalar_type() == torch::kInt64, - "sm70_f16_rerank_topk_out: expected FP16/int64 tensors."); - TORCH_CHECK(logits.dim() == 2 && logits.size(1) == kCandidates && - candidate_ids.sizes() == logits.sizes(), - "sm70_f16_rerank_topk_out: expected matching [M, 64] " - "candidate tensors."); - TORCH_CHECK(values_out.dim() == 2 && ids_out.sizes() == values_out.sizes() && - values_out.size(0) == logits.size(0) && - (values_out.size(1) == 16 || values_out.size(1) == 20), - "sm70_f16_rerank_topk_out: outputs must be matching [M, 16] " - "or [M, 20] tensors."); - TORCH_CHECK(values_out.is_contiguous() && ids_out.is_contiguous() && - logits.is_contiguous() && candidate_ids.is_contiguous(), - "sm70_f16_rerank_topk_out: tensors must be contiguous."); - const int device = logits.get_device(); - TORCH_CHECK(values_out.get_device() == device && - ids_out.get_device() == device && - candidate_ids.get_device() == device, - "sm70_f16_rerank_topk_out: tensors must share one device."); - const at::cuda::OptionalCUDAGuard device_guard(device_of(logits)); - const int rows = static_cast(logits.size(0)); - if (rows == 0) { - return; - } - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - sm70_f16_rerank_topk_kernel<<>>( - reinterpret_cast(values_out.data_ptr()), - ids_out.data_ptr(), - reinterpret_cast(logits.data_ptr()), - candidate_ids.data_ptr(), rows, - static_cast(values_out.size(1)), vocab_start_index); - C10_CUDA_KERNEL_LAUNCH_CHECK(); -} - void sm70_f16_lm_head_top1_out(torch::Tensor values_out, torch::Tensor indices_out, torch::Tensor in_feats, torch::Tensor weight, @@ -5567,27 +4875,6 @@ void fp8_gemm_sm70_out(torch::Tensor out, torch::Tensor _in_feats, group_size, k_ld, q_ld, gated_silu); } -void fp8_gemm_sm70_prefill_prescaled_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) { - 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_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, @@ -5659,41 +4946,6 @@ 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_f16_indexed_rerank_out(torch::Tensor out, torch::Tensor _in_feats, - torch::Tensor _kernel, - torch::Tensor candidate_ids, - torch::Tensor selected_raw, - torch::Tensor selected_packed, - torch::Tensor expanded, torch::Tensor partials, - torch::Tensor barriers, int64_t cta_n, - int64_t split_k) { - vllm::awq_sm70::sm70_f16_indexed_rerank_out( - out, _in_feats, _kernel, candidate_ids, selected_raw, selected_packed, - expanded, partials, barriers, cta_n, split_k); -} - -void sm70_f16_indexed_rerank_packed_out( - torch::Tensor out, torch::Tensor _in_feats, torch::Tensor _packed_kernel, - torch::Tensor candidate_ids, torch::Tensor selected_packed, - torch::Tensor expanded, torch::Tensor partials, torch::Tensor barriers, - int64_t cta_n, int64_t split_k) { - vllm::awq_sm70::sm70_f16_indexed_rerank_packed_out( - out, _in_feats, _packed_kernel, candidate_ids, selected_packed, expanded, - partials, barriers, cta_n, split_k); -} - -void sm70_f16_rerank_keys_out(torch::Tensor keys, torch::Tensor logits, - torch::Tensor candidate_ids) { - vllm::awq_sm70::sm70_f16_rerank_keys_out(keys, logits, candidate_ids); -} - -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) { - vllm::awq_sm70::sm70_f16_rerank_topk_out(values_out, ids_out, logits, - candidate_ids, vocab_start_index); -} - void sm70_f16_lm_head_top1_out(torch::Tensor values_out, torch::Tensor indices_out, torch::Tensor _in_feats, torch::Tensor _kernel, @@ -6656,29 +5908,17 @@ void awq_moe_single_token_weighted_reduce_out(torch::Tensor sorted_output, const int hidden_logical_size_i = static_cast(hidden_logical_size); const int sorted_output_row_stride = static_cast(sorted_output.stride(0)); - if ((top_k == 6 || top_k == 8) && (hidden_logical_size_i % 2) == 0 && + if (top_k == 8 && (hidden_logical_size_i % 2) == 0 && (sorted_output_row_stride % 2) == 0) { const int blocks = std::max( 1, ((hidden_logical_size_i >> 1) + kThreads - 1) / kThreads); - if (top_k == 6) { - awq_moe_single_token_weighted_reduce_half2_kernel<6> - <<>>( - reinterpret_cast( - sorted_output.data_ptr()), - topk_weights.data_ptr(), - inv_permuted_idx.data_ptr(), - reinterpret_cast<__half*>(out.data_ptr()), - hidden_logical_size_i, sorted_output_row_stride); - } else { - awq_moe_single_token_weighted_reduce_half2_kernel<8> - <<>>( - reinterpret_cast( - sorted_output.data_ptr()), - topk_weights.data_ptr(), - inv_permuted_idx.data_ptr(), - reinterpret_cast<__half*>(out.data_ptr()), - hidden_logical_size_i, sorted_output_row_stride); - } + awq_moe_single_token_weighted_reduce_half2_kernel<8> + <<>>( + reinterpret_cast(sorted_output.data_ptr()), + topk_weights.data_ptr(), + inv_permuted_idx.data_ptr(), + reinterpret_cast<__half*>(out.data_ptr()), + hidden_logical_size_i, sorted_output_row_stride); } else { const int blocks = std::max(1, (hidden_logical_size_i + kThreads - 1) / kThreads); @@ -8125,8 +7365,7 @@ void mxfp4_moe_gemm_sm70_out_impl( torch::Tensor out, torch::Tensor sorted_input, torch::Tensor expert_offsets, torch::Tensor strided_ptrs_w, torch::Tensor strided_ptrs_s, int64_t num_experts, int64_t k, int64_t n, int64_t group_size, - torch::Tensor b_group_indices, bool compact_grouped_rows = false, - bool broadcast_input_rows = false) { + torch::Tensor b_group_indices, bool compact_grouped_rows = false) { TORCH_CHECK( sorted_input.is_cuda() && sorted_input.scalar_type() == torch::kFloat16, "mxfp4_moe_gemm_sm70: input must be CUDA float16."); @@ -8147,16 +7386,9 @@ void mxfp4_moe_gemm_sm70_out_impl( "mxfp4_moe_gemm_sm70: k must be divisible by group_size."); TORCH_CHECK(sorted_input.dim() == 2 && sorted_input.size(1) == k, "mxfp4_moe_gemm_sm70: input shape mismatch."); - const int64_t total_tokens = - broadcast_input_rows ? out.size(0) : sorted_input.size(0); - TORCH_CHECK(out.dim() == 2 && out.size(0) == total_tokens && + TORCH_CHECK(out.dim() == 2 && out.size(0) == sorted_input.size(0) && out.size(1) == n && out.stride(1) == 1, "mxfp4_moe_gemm_sm70: output must be contiguous [tokens, n]."); - TORCH_CHECK(!broadcast_input_rows || - (sorted_input.size(0) == 1 && total_tokens == num_experts && - compact_grouped_rows), - "mxfp4_moe_gemm_sm70: broadcast input requires one physical " - "row and one logical row per compact group."); TORCH_CHECK(expert_offsets.numel() >= num_experts + 1, "mxfp4_moe_gemm_sm70: expert_offsets too small."); TORCH_CHECK(b_group_indices.is_cuda() && @@ -8166,6 +7398,7 @@ void mxfp4_moe_gemm_sm70_out_impl( "mxfp4_moe_gemm_sm70: B group indices must be contiguous CUDA " "int32."); + const int64_t total_tokens = sorted_input.size(0); if (total_tokens == 0) { return; } @@ -8189,8 +7422,7 @@ void mxfp4_moe_gemm_sm70_out_impl( static_cast(sorted_input.stride(0)), }; desc_A.num = static_cast(num_experts); - desc_A.offsets = - broadcast_input_rows ? nullptr : expert_offsets.data_ptr(); + desc_A.offsets = expert_offsets.data_ptr(); turbomind::gemm::MatrixLayout desc_U{}; const auto order_w = conv_w->order; @@ -8262,18 +7494,8 @@ void mxfp4_moe_gemm_sm70_out_impl( op.quant_b = {turbomind::gemm::QuantType::kK, static_cast(group_size)}; op.batch_dim = 0; op.dispatch_num_override = compact_grouped_rows ? 1 : 0; - const bool dynamic_expert_rows = - compact_grouped_rows && num_experts >= 12 && num_experts <= 48 && - num_experts % 6 == 0 && - vllm::awq_sm70::mxfp4_moe_grouped_m8_expert_rows_enabled(); - // The negative active-group contract is valid only when every group owns - // exactly one row (for example compact B1/top-6 and slot-grouped M8). Real - // expert segments can contain multiple rows and have graph-dynamic empty - // tails, so retain the single-group dispatch choice while letting the - // standard offsets scheduler discover their bounds on device. - op.active_group_count = compact_grouped_rows && !dynamic_expert_rows - ? -static_cast(num_experts) - : 0; + op.active_group_count = + compact_grouped_rows ? -static_cast(num_experts) : 0; auto& workspace_holder = vllm::awq_sm70::get_workspace(device, stream); auto& gemm = vllm::awq_sm70::get_gemm(device); @@ -8316,12 +7538,8 @@ void mxfp4_moe_dense_stage_sm70_out(torch::Tensor out, torch::Tensor input, "supported."); TORCH_CHECK(input.dim() == 2 && input.size(1) == k, "mxfp4_moe_dense_stage_sm70_out: input shape mismatch."); - 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), + out.dim() == 2 && out.size(0) == input.size(0) && out.size(1) == n, "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."); @@ -8334,28 +7552,20 @@ void mxfp4_moe_dense_stage_sm70_out(torch::Tensor out, torch::Tensor input, logged_mxfp4_dense_stage, "SM70 MXFP4 MoE CUDA-graph-safe dense-stage path enabled C++ op reached", input, input.size(0), num_experts); - const bool compact_decode_shape = input.size(0) == num_experts && - num_experts == 6 && - ((k == 4096 && (n == 512 || n == 1024)) || - (n == 4096 && (k == 256 || k == 512))); - 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); + const bool compact_decode_shape = + input.size(0) == num_experts && num_experts == 6 && + ((k == 4096 && n == 512) || (k == 256 && n == 4096)); + if (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); return; } const bool grouped_m8_shape = input.size(0) == 48 && num_experts == 48 && ((k == 4096 && n == 512) || (k == 256 && n == 4096)); - const bool grouped_verifier_shape = - input.size(0) == num_experts && num_experts >= 12 && num_experts <= 48 && - num_experts % 6 == 0 && - ((k == 4096 && n == 512) || (k == 256 && n == 4096)); - if ((vllm::awq_sm70::mxfp4_moe_grouped_m8_enabled() && grouped_m8_shape) || - (vllm::awq_sm70::mxfp4_moe_grouped_verifier_enabled() && - grouped_verifier_shape)) { + if (vllm::awq_sm70::mxfp4_moe_grouped_m8_enabled() && grouped_m8_shape) { // Keep every routed slot as an independent one-row group. Repeated experts // reuse the same weight pointer row, while the scheduler retains the exact // arithmetic contract already accepted for compact B1 decode. @@ -8372,237 +7582,6 @@ void mxfp4_moe_dense_stage_sm70_out(torch::Tensor out, torch::Tensor input, } } -void nvfp4_moe_gemm_sm70_out_impl( - torch::Tensor out, torch::Tensor sorted_input, torch::Tensor expert_offsets, - torch::Tensor strided_ptrs_w, torch::Tensor strided_ptrs_s, - int64_t num_experts, int64_t k, int64_t n, int64_t group_size, - torch::Tensor b_group_indices, bool compact_grouped_rows = false) { - TORCH_CHECK( - sorted_input.is_cuda() && sorted_input.scalar_type() == torch::kFloat16, - "nvfp4_moe_gemm_sm70: input must be CUDA float16."); - TORCH_CHECK( - expert_offsets.is_cuda() && - expert_offsets.scalar_type() == torch::kInt32 && - expert_offsets.is_contiguous(), - "nvfp4_moe_gemm_sm70: expert_offsets must be contiguous CUDA int32."); - TORCH_CHECK(strided_ptrs_w.is_cuda() && strided_ptrs_s.is_cuda(), - "nvfp4_moe_gemm_sm70: strided_ptrs must be CUDA."); - TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kFloat16, - "nvfp4_moe_gemm_sm70: output must be CUDA float16."); - TORCH_CHECK(num_experts > 0 && k > 0 && n > 0, - "nvfp4_moe_gemm_sm70: invalid dimensions."); - TORCH_CHECK(group_size == 16, - "nvfp4_moe_gemm_sm70: only group_size=16 is supported."); - TORCH_CHECK(k % group_size == 0, - "nvfp4_moe_gemm_sm70: k must be divisible by group_size."); - TORCH_CHECK(sorted_input.dim() == 2 && sorted_input.size(1) == k, - "nvfp4_moe_gemm_sm70: input shape mismatch."); - TORCH_CHECK(out.dim() == 2 && out.size(0) == sorted_input.size(0) && - out.size(1) == n && out.stride(1) == 1, - "nvfp4_moe_gemm_sm70: output must be contiguous [tokens, n]."); - TORCH_CHECK(expert_offsets.numel() >= num_experts + 1, - "nvfp4_moe_gemm_sm70: expert_offsets too small."); - TORCH_CHECK(b_group_indices.is_cuda() && - b_group_indices.scalar_type() == torch::kInt32 && - b_group_indices.is_contiguous() && - b_group_indices.numel() >= num_experts, - "nvfp4_moe_gemm_sm70: B group indices must be contiguous CUDA " - "int32."); - - const int64_t total_tokens = sorted_input.size(0); - if (total_tokens == 0) { - return; - } - - const at::cuda::OptionalCUDAGuard device_guard(device_of(sorted_input)); - const int device = sorted_input.get_device(); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - - const auto fp4_converters = turbomind::gemm::GetConverters( - turbomind::kHalf, turbomind::kFloat4_e2m1, turbomind::kHalf, true, 70); - const auto fp8_converters = turbomind::gemm::GetConverters( - turbomind::kHalf, turbomind::kFloat8_e4m3, turbomind::kHalf, true, 70); - const auto* conv_w = fp4_converters[0]; - const auto* conv_s = fp8_converters[1]; - TORCH_CHECK(conv_w && conv_s, - "nvfp4_moe_gemm_sm70: no compatible TurboMind converters."); - - turbomind::gemm::MatrixLayout desc_A{ - turbomind::kHalf, - turbomind::gemm::kRowMajor, - static_cast(total_tokens), - static_cast(k), - static_cast(sorted_input.stride(0)), - }; - desc_A.num = static_cast(num_experts); - desc_A.offsets = expert_offsets.data_ptr(); - turbomind::gemm::MatrixLayout desc_U{}; - - const auto order_w = conv_w->order; - const bool is_a_w = turbomind::gemm::get_operand_tag(conv_w->pack) == - turbomind::gemm::OPERAND_A; - const bool is_b_w = !is_a_w; - turbomind::gemm::MatrixLayout weight_desc{ - turbomind::kHalf, - order_w, - static_cast(n), - static_cast(k), - order_w == turbomind::gemm::kRowMajor ? static_cast(k) - : static_cast(n), - }; - if (is_b_w) { - std::swap(weight_desc.rows, weight_desc.cols); - weight_desc.order = ~weight_desc.order; - } - turbomind::gemm::MatrixLayout desc_B = weight_desc; - desc_B.type = turbomind::kFloat4_e2m1; - desc_B.pack = conv_w->pack; - if (is_a_w) { - desc_B = turbomind::gemm::transpose(desc_B); - } - desc_B.ld = 0; - desc_B.num = static_cast(num_experts); - desc_B.group_idxs = b_group_indices.data_ptr(); - - const auto order_s = conv_s->order; - const bool is_a_s = turbomind::gemm::get_operand_tag(conv_s->pack) == - turbomind::gemm::OPERAND_U; - const bool is_b_s = !is_a_s; - const int64_t num_groups = k / group_size; - turbomind::gemm::MatrixLayout scale_desc{ - turbomind::kUint16, order_s, - static_cast(n), static_cast(num_groups), - static_cast(n), - }; - if (is_b_s) { - std::swap(scale_desc.rows, scale_desc.cols); - scale_desc.order = ~scale_desc.order; - } - turbomind::gemm::MatrixLayout desc_V = scale_desc; - desc_V.pack = conv_s->pack; - if (is_a_s) { - desc_V = turbomind::gemm::transpose(desc_V); - } - desc_V.ld = 0; - desc_V.num = static_cast(num_experts); - desc_V.group_idxs = b_group_indices.data_ptr(); - - turbomind::gemm::MatrixLayout desc_D{ - turbomind::kHalf, - turbomind::gemm::kRowMajor, - static_cast(total_tokens), - static_cast(n), - static_cast(out.stride(0)), - }; - desc_D.num = static_cast(num_experts); - desc_D.offsets = expert_offsets.data_ptr(); - - turbomind::gemm::Operation op{}; - op.dispatch = vllm::awq_sm70::select_nvfp4_moe_dispatch_policy( - device, static_cast(total_tokens), static_cast(n), - static_cast(k), static_cast(num_experts), - static_cast(group_size), stream); - op.epilogue = turbomind::gemm::Epilogue::kNone; - op.quant_a = {turbomind::gemm::QuantType::kNone, 0}; - op.quant_b = {turbomind::gemm::QuantType::kK, static_cast(group_size)}; - op.batch_dim = 0; - op.dispatch_num_override = compact_grouped_rows ? 1 : 0; - op.active_group_count = - compact_grouped_rows ? -static_cast(num_experts) : 0; - - auto& workspace_holder = vllm::awq_sm70::get_workspace(device, stream); - auto& gemm = vllm::awq_sm70::get_gemm(device); - const int ec = - gemm.Run(op, 1.f, sorted_input.data_ptr(), desc_A, nullptr, desc_U, - strided_ptrs_w.data_ptr(), desc_B, strided_ptrs_s.data_ptr(), - desc_V, 0.f, out.data_ptr(), desc_D, out.data_ptr(), desc_D, - workspace_holder.workspace, stream); - TORCH_CHECK(ec == 0, - "nvfp4_moe_gemm_sm70: TurboMind batched GEMM failed (ec=", ec, - ")."); -} - -void nvfp4_moe_dense_stage_sm70_out(torch::Tensor out, torch::Tensor input, - torch::Tensor expert_offsets, - torch::Tensor dense_expert_ids, - torch::Tensor ptrs_w, torch::Tensor ptrs_s, - int64_t num_experts, int64_t k, int64_t n, - int64_t group_size) { - TORCH_CHECK(input.is_cuda() && input.scalar_type() == torch::kFloat16, - "nvfp4_moe_dense_stage_sm70_out: input must be CUDA float16."); - TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kFloat16, - "nvfp4_moe_dense_stage_sm70_out: out must be CUDA float16."); - TORCH_CHECK(expert_offsets.is_cuda() && - expert_offsets.scalar_type() == torch::kInt32 && - expert_offsets.is_contiguous(), - "nvfp4_moe_dense_stage_sm70_out: expert_offsets must be " - "contiguous CUDA int32."); - TORCH_CHECK(dense_expert_ids.is_cuda() && - dense_expert_ids.scalar_type() == torch::kInt32 && - dense_expert_ids.is_contiguous(), - "nvfp4_moe_dense_stage_sm70_out: dense_expert_ids must be " - "contiguous CUDA int32."); - TORCH_CHECK(ptrs_w.is_cuda() && ptrs_s.is_cuda(), - "nvfp4_moe_dense_stage_sm70_out: ptr rows must be CUDA."); - TORCH_CHECK(num_experts > 0, - "nvfp4_moe_dense_stage_sm70_out: num_experts must be positive."); - TORCH_CHECK(group_size == 16, - "nvfp4_moe_dense_stage_sm70_out: only group_size=16 is " - "supported."); - TORCH_CHECK(input.dim() == 2 && input.size(1) == k, - "nvfp4_moe_dense_stage_sm70_out: input shape mismatch."); - TORCH_CHECK( - out.dim() == 2 && out.size(0) == input.size(0) && out.size(1) == n, - "nvfp4_moe_dense_stage_sm70_out: out shape mismatch."); - TORCH_CHECK(expert_offsets.numel() >= num_experts + 1, - "nvfp4_moe_dense_stage_sm70_out: expert_offsets too small."); - TORCH_CHECK(dense_expert_ids.numel() >= num_experts, - "nvfp4_moe_dense_stage_sm70_out: dense_expert_ids too small."); - - const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); - static std::atomic logged_nvfp4_dense_stage{0u}; - maybe_log_sm70_moe_route_once( - logged_nvfp4_dense_stage, - "SM70 NVFP4 MoE CUDA-graph-safe TurboMind path enabled C++ op reached", - input, input.size(0), num_experts); - constexpr int kNvfp4LegacyCompactGroups = 8 * 8; - // Cover the validated top-k8 B10 verifier bound. Duplicate expert selections - // remain separate one-row groups, so this is a routed-slot bound rather - // than an active-expert bound. Keep the dense-grouped cutoff independent: - // a disabled compact route above 64 rows must not fall back per expert. - constexpr int kNvfp4MaxCompactGroups = 10 * 8; - const bool compact_decode_shape = - input.size(0) == num_experts && num_experts <= kNvfp4MaxCompactGroups; - if (compact_decode_shape) { - nvfp4_moe_gemm_sm70_out_impl(out, input, expert_offsets, ptrs_w, ptrs_s, - num_experts, k, n, group_size, - dense_expert_ids, true); - return; - } - const bool exact_qwen36_prefill_shape = - input.size(0) > kNvfp4LegacyCompactGroups && num_experts == 256 && - ((k == 2048 && (n == 1024 || n == 512 || n == 256)) || - (n == 2048 && (k == 512 || k == 256 || k == 128))); - if (vllm::awq_sm70::nvfp4_moe_grouped_prefill_enabled() && - exact_qwen36_prefill_shape) { - static std::atomic logged_nvfp4_grouped_prefill{0u}; - maybe_log_sm70_moe_route_once( - logged_nvfp4_grouped_prefill, - "SM70 NVFP4 MoE grouped TurboMind prefill path enabled C++ op reached", - input, input.size(0), num_experts); - nvfp4_moe_gemm_sm70_out_impl(out, input, expert_offsets, ptrs_w, ptrs_s, - num_experts, k, n, group_size, - dense_expert_ids); - return; - } - for (int expert = 0; expert < static_cast(num_experts); ++expert) { - torch::Tensor offsets = expert_offsets.narrow(0, expert, 2); - torch::Tensor expert_idx = dense_expert_ids.narrow(0, expert, 1); - nvfp4_moe_gemm_sm70_out_impl(out, input, offsets, ptrs_w, ptrs_s, 1, k, n, - group_size, expert_idx); - } -} - void mxfp4_moe_single_token_prepare_w13_sm70_out( torch::Tensor gate_up, torch::Tensor compact_input, torch::Tensor x, torch::Tensor topk_ids, torch::Tensor w13_ptrs_w, torch::Tensor w13_ptrs_s, @@ -8627,10 +7606,9 @@ void mxfp4_moe_single_token_prepare_w13_sm70_out( TORCH_CHECK(topk_ids.numel() == kTopK, "mxfp4_moe_single_token_prepare_w13_sm70_out: exact top-k=6 is " "required."); - TORCH_CHECK(group_size == 32 && w13_k == hidden_logical_size && - (w13_n == 512 || w13_n == 1024), + TORCH_CHECK(group_size == 32 && w13_k == hidden_logical_size && w13_n == 512, "mxfp4_moe_single_token_prepare_w13_sm70_out: unsupported " - "MXFP4 W13 shape contract."); + "DeepSeek V4 W13 shape contract."); TORCH_CHECK( compact_input.is_cuda() && compact_input.scalar_type() == torch::kFloat16 && @@ -8677,9 +7655,6 @@ void mxfp4_moe_single_token_prepare_w13_sm70_out( x.size(0), kTopK); constexpr int kThreads = 256; - const bool broadcast_input = - w13_n == 512 && - vllm::awq_sm70::mxfp4_moe_broadcast_input_decode_enabled(); const int src_row_stride = kPtrRowBytes; const int row_bytes = kPtrRowBytes; if (topk_ids.scalar_type() == torch::kInt32) { @@ -8690,10 +7665,8 @@ void mxfp4_moe_single_token_prepare_w13_sm70_out( w13_ptrs_w.data_ptr(), w13_ptrs_s.data_ptr(), w13_ptrs_w.data_ptr(), w13_ptrs_s.data_ptr(), w13_ptrs_w.data_ptr(), w13_ptrs_s.data_ptr(), - broadcast_input - ? nullptr - : reinterpret_cast<__half*>(compact_input.data_ptr()), - nullptr, expert_offsets.data_ptr(), nullptr, + reinterpret_cast<__half*>(compact_input.data_ptr()), nullptr, + expert_offsets.data_ptr(), nullptr, inv_permuted_idx.data_ptr(), sorted_expert_ids.data_ptr(), kTopK, static_cast(hidden_logical_size), src_row_stride, src_row_stride, @@ -8706,10 +7679,8 @@ void mxfp4_moe_single_token_prepare_w13_sm70_out( w13_ptrs_w.data_ptr(), w13_ptrs_s.data_ptr(), w13_ptrs_w.data_ptr(), w13_ptrs_s.data_ptr(), w13_ptrs_w.data_ptr(), w13_ptrs_s.data_ptr(), - broadcast_input - ? nullptr - : reinterpret_cast<__half*>(compact_input.data_ptr()), - nullptr, expert_offsets.data_ptr(), nullptr, + reinterpret_cast<__half*>(compact_input.data_ptr()), nullptr, + expert_offsets.data_ptr(), nullptr, inv_permuted_idx.data_ptr(), sorted_expert_ids.data_ptr(), kTopK, static_cast(hidden_logical_size), src_row_stride, src_row_stride, @@ -8717,10 +7688,9 @@ void mxfp4_moe_single_token_prepare_w13_sm70_out( } C10_CUDA_KERNEL_LAUNCH_CHECK(); - mxfp4_moe_gemm_sm70_out_impl(gate_up, broadcast_input ? x : compact_input, - expert_offsets, w13_ptrs_w, w13_ptrs_s, kTopK, - w13_k, w13_n, group_size, sorted_expert_ids, - true, broadcast_input); + mxfp4_moe_gemm_sm70_out_impl(gate_up, compact_input, expert_offsets, + w13_ptrs_w, w13_ptrs_s, kTopK, w13_k, w13_n, + group_size, sorted_expert_ids, true); } void fp8_moe_single_token_dense_stage_sm70_out( diff --git a/csrc/sm70_turbomind/ops/fp8_qpn8_sm70.cu b/csrc/sm70_turbomind/ops/fp8_qpn8_sm70.cu deleted file mode 100644 index 2dd90101e8..0000000000 --- a/csrc/sm70_turbomind/ops/fp8_qpn8_sm70.cu +++ /dev/null @@ -1,1255 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project -// -// The QPN8 execution layout is derived from dnv2003/v100-skinny (MIT) and -// its block-scale adaptation in haohervchb/sglang-V100. Automatic operator -// dispatch is restricted to accepted SM70 tensor/layout contracts and can be -// disabled with VLLM_SM70_FP8_QPN8=0. -// See LICENSE.v100-skinny in this directory for the retained MIT notice. - -#include -#include - -#include -#include -#include -#include -#include - -#include -#include - -namespace { - -__device__ __forceinline__ void fp8x8_to_half2x4(uint2 q, half2 out[4]) { -#pragma unroll - for (int i = 0; i < 4; ++i) { - const unsigned b0 = (q.x >> (8 * i)) & 0xffu; - const unsigned b1 = (q.y >> (8 * i)) & 0xffu; - const unsigned h0 = ((b0 & 0x80u) << 8) | ((b0 & 0x7fu) << 7); - const unsigned h1 = ((b1 & 0x80u) << 8) | ((b1 & 0x7fu) << 7); - const unsigned packed = h0 | (h1 << 16); - out[i] = *reinterpret_cast(&packed); - } -} - -__device__ __forceinline__ void fp8x8_to_half2x4_fast(uint2 q, half2 out[4]) { - constexpr unsigned kSign = 0x80008000u; - constexpr unsigned kExponentMantissa = 0x3f803f80u; - unsigned permuted[4]; - permuted[0] = __byte_perm(q.x, q.y, 0x0400); - permuted[1] = __byte_perm(q.x, q.y, 0x0501); - permuted[2] = __byte_perm(q.x, q.y, 0x0602); - permuted[3] = __byte_perm(q.x, q.y, 0x0703); -#pragma unroll - for (int i = 0; i < 4; ++i) { - const unsigned value = - ((permuted[i] << 8) & kSign) | ((permuted[i] << 7) & kExponentMantissa); - out[i] = *reinterpret_cast(&value); - } -} - -constexpr int kQpn8PrepareThreads = 256; - -__device__ __forceinline__ int qpn8_col_from_lane(int lane) { - return ((lane >> 2) & 3) * 8 + (lane & 3) + ((lane & 16) ? 4 : 0); -} - -__device__ __forceinline__ int qpn8_lane_from_col(int col) { - return (col & 3) | (((col >> 3) & 3) << 2) | (((col >> 2) & 1) << 4); -} - -__device__ __forceinline__ int qpn8_physical_k(int logical_k) { - const int local = logical_k & 7; - return (logical_k & 8) + (local >> 1) + ((local & 1) << 2); -} - -__global__ void fp8_qpn8_prepack_sm70_kernel( - uint8_t* __restrict__ codes, const uint8_t* __restrict__ qweight, int n, - int k) { - const size_t index = - static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - const size_t numel = static_cast(n) * k; - if (index >= numel) { - return; - } - - const int row = static_cast(index / k); - const int logical_k = static_cast(index % k); - const int tile = row >> 5; - const int lane = qpn8_lane_from_col(row & 31); - const int group = logical_k >> 4; - const int physical_k = qpn8_physical_k(logical_k & 15); - const int groups_k16 = k >> 4; - const size_t packed_index = - (((static_cast(tile) * groups_k16 + group) * 32 + lane) * 16 + - physical_k); - codes[packed_index] = qweight[index]; -} - -__global__ void fp8_qpn8_scale_sm70_kernel(half* __restrict__ group_scales, - const float* __restrict__ scales, - int n_blocks, int k_blocks) { - const int n_tiles = n_blocks * 4; - const int index = blockIdx.x * blockDim.x + threadIdx.x; - const int numel = k_blocks * n_tiles; - if (index >= numel) { - return; - } - - const int k_block = index / n_tiles; - const int n_tile = index - k_block * n_tiles; - group_scales[index] = - __float2half(scales[(n_tile >> 2) * k_blocks + k_block] * 256.0f); -} - -__global__ void fp8_qpn8_channel_scale_sm70_kernel( - half* __restrict__ channel_scales, const float* __restrict__ scales, - int n) { - const int col = blockIdx.x * blockDim.x + threadIdx.x; - if (col < n) { - channel_scales[col] = __float2half(scales[col] * 256.0f); - } -} - -__global__ void fp8_qpn8_dequantize_sm70_kernel( - half* __restrict__ output, const uint8_t* __restrict__ codes, - const half* __restrict__ group_scales, int n, int k, bool channel_scales) { - const size_t word_index = - static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - const size_t word_count = static_cast(n) * k / 16; - if (word_index >= word_count) { - return; - } - - const int groups_k16 = k >> 4; - const int lane = static_cast(word_index & 31); - const size_t tile_word = word_index >> 5; - const int group = static_cast(tile_word % groups_k16); - const int tile = static_cast(tile_word / groups_k16); - const int col = tile * 32 + qpn8_col_from_lane(lane); - const int tiles_n32 = n >> 5; - const half scale = channel_scales - ? group_scales[col] - : group_scales[(group >> 3) * tiles_n32 + tile]; - const half2 scale2 = __halves2half2(scale, scale); - const uint4 packed = reinterpret_cast(codes)[word_index]; - 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 pair = 0; pair < 8; ++pair) { - const half2 value = __hmul2(weights[pair], scale2); - const int k_base = group * 16 + pair * 2; - output[static_cast(k_base) * n + col] = __low2half(value); - output[static_cast(k_base + 1) * n + col] = __high2half(value); - } -} - -__global__ void fp8_qpn8_silu_and_mul_sm70_kernel( - half* __restrict__ output, const half* __restrict__ gate_up, int rows, - int hidden) { - const int row = blockIdx.x; - if (row >= rows) { - return; - } - const half* row_input = gate_up + static_cast(row) * hidden * 2; - half* row_output = output + static_cast(row) * hidden; - for (int col = threadIdx.x; col < hidden; col += blockDim.x) { - const float gate = __half2float(row_input[col]); - const float silu = gate / (1.0f + __expf(-gate)); - row_output[col] = __hmul(__float2half(silu), row_input[hidden + col]); - } -} - -#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)) - -__global__ void fp8_qpn8_ba_split_copy_sm70_kernel( - const half* __restrict__ qkvz, const half* __restrict__ ba, - half* __restrict__ qkv, half* __restrict__ z, half* __restrict__ b, - half* __restrict__ a, int m) { - constexpr int kQkvN = 2560; - constexpr int kZN = 1536; - constexpr int kQkvzN = kQkvN + kZN; - constexpr int kBAN = 24; - constexpr int kOutputN = kQkvzN + kBAN; - const int index = blockIdx.x * blockDim.x + threadIdx.x; - if (index >= m * kOutputN) { - return; - } - const int row = index / kOutputN; - const int col = index - row * kOutputN; - if (col < kQkvN) { - qkv[static_cast(row) * kQkvN + col] = - qkvz[static_cast(row) * kQkvzN + col]; - } else if (col < kQkvzN) { - z[static_cast(row) * kZN + col - kQkvN] = - qkvz[static_cast(row) * kQkvzN + col]; - } else if (col < kQkvzN + kBAN / 2) { - b[static_cast(row) * (kBAN / 2) + col - kQkvzN] = - ba[static_cast(row) * kBAN + col - kQkvzN]; - } else { - a[static_cast(row) * (kBAN / 2) + col - kQkvzN - kBAN / 2] = - ba[static_cast(row) * kBAN + col - kQkvzN]; - } -} - -template -__global__ void fp8_qpn8_sm70_kernel( - const uint8_t* __restrict__ codes, const half* __restrict__ group_scales, - const half* __restrict__ input, half* __restrict__ output, - half* __restrict__ z_output, const half* __restrict__ ba_weight, - half* __restrict__ ba_output, half* __restrict__ b_output, - half* __restrict__ a_output, int ba_n, int qkv_n, int n, int k, int m, - bool channel_scales) { - __shared__ float partials[SplitK][M1Only ? 32 : 256]; - - const int lane = threadIdx.x & 31; - const int warp = threadIdx.x >> 5; - const int tile = blockIdx.x; - if constexpr (FusedBA) { - const int qpn_tiles = n >> 5; - if (tile >= qpn_tiles) { - constexpr int kBARowsPerBlock = 2; - constexpr int kBAThreadsPerRow = 256; - const int ba_block = tile - qpn_tiles; - const int ba_group = threadIdx.x / kBAThreadsPerRow; - const int ba_thread = threadIdx.x % kBAThreadsPerRow; - const int ba_warp = ba_thread >> 5; - const int ba_row = ba_block * kBARowsPerBlock + ba_group; - float value = 0.0f; - const half2* input2 = reinterpret_cast(input); - const half2* weight2 = reinterpret_cast( - ba_weight + static_cast(ba_row) * k); - for (int pair = ba_thread; pair < k / 2; pair += kBAThreadsPerRow) { - const float2 x = __half22float2(__ldg(input2 + pair)); - const float2 weight = __half22float2(__ldg(weight2 + pair)); - value = fmaf(x.x, weight.x, value); - value = fmaf(x.y, weight.y, value); - } -#pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) { - value += __shfl_down_sync(0xffffffffU, value, offset); - } - if (lane == 0) { - partials[ba_group][ba_warp] = value; - } - __syncthreads(); - if (ba_warp == 0) { - value = lane < 8 ? partials[ba_group][lane] : 0.0f; -#pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) { - value += __shfl_down_sync(0xffffffffU, value, offset); - } - if (lane == 0 && ba_row < ba_n) { - if constexpr (SplitOutputs) { - if (ba_row < ba_n / 2) { - b_output[ba_row] = __float2half(value); - } else { - a_output[ba_row - ba_n / 2] = __float2half(value); - } - } else { - ba_output[ba_row] = __float2half(value); - } - } - } - return; - } - } - 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_warp = groups_k16 / SplitK; - const int group_begin = warp * groups_per_warp; - const int tiles_n32 = n >> 5; - const uint4* code_ptr = reinterpret_cast(codes) + - static_cast(tile) * groups_k16 * 32 + lane; - const half* scale_ptr = group_scales + tile; - - float accum[NAcc][8]; -#pragma unroll - for (int chain = 0; chain < NAcc; ++chain) { -#pragma unroll - for (int i = 0; i < 8; ++i) { - accum[chain][i] = 0.0f; - } - } - int loaded_scale_group = -1; - half loaded_scale = - channel_scales - ? __ldg(group_scales + tile * 32 + qpn8_col_from_lane(lane)) - : __float2half(0.0f); - uint4 prefetched = make_uint4(0, 0, 0, 0); - if constexpr (PrefetchCodes) { - prefetched = __ldcs(code_ptr + static_cast(group_begin) * 32); - } - -#pragma unroll 4 - for (int group = group_begin; group < group_begin + groups_per_warp; - ++group) { - const int scale_group = group >> 3; - if (!channel_scales && scale_group != loaded_scale_group) { - loaded_scale = - __ldg(scale_ptr + static_cast(scale_group) * tiles_n32); - loaded_scale_group = scale_group; - } - - const uint4 packed = - PrefetchCodes ? prefetched - : __ldcs(code_ptr + static_cast(group) * 32); - uint4 next = make_uint4(0, 0, 0, 0); - if constexpr (PrefetchCodes) { - if (group + 1 < group_begin + groups_per_warp) { - next = __ldcs(code_ptr + static_cast(group + 1) * 32); - } - } - half2 weights[8]; - if constexpr (FastDecoder) { - fp8x8_to_half2x4_fast(make_uint2(packed.x, packed.y), weights); - fp8x8_to_half2x4_fast(make_uint2(packed.z, packed.w), weights + 4); - } else { - fp8x8_to_half2x4(make_uint2(packed.x, packed.y), weights); - fp8x8_to_half2x4(make_uint2(packed.z, packed.w), weights + 4); - } - - const half2 scale2 = __halves2half2(loaded_scale, loaded_scale); -#pragma unroll - for (int i = 0; i < 8; ++i) { - weights[i] = __hmul2(weights[i], scale2); - } - - uint4 input01 = make_uint4(0, 0, 0, 0); - uint4 input23 = make_uint4(0, 0, 0, 0); - if (row < m) { - const half* input_row = input + static_cast(row) * k; - 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); - 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]); - if constexpr (PrefetchCodes) { - prefetched = next; - } - } - -#pragma unroll - for (int chain = 1; chain < NAcc; ++chain) { -#pragma unroll - for (int i = 0; i < 8; ++i) { - accum[0][i] += accum[chain][i]; - } - } - - if constexpr (M1Only) { - if ((lane & 17) == 0) { -#pragma unroll - for (int pair = 0; pair < 2; ++pair) { -#pragma unroll - for (int offset = 0; offset < 2; ++offset) { - const int i = pair * 4 + offset; - const int output_col = - offset | (((lane >> 1) & 1) << 1) | (pair << 2); - partials[warp][quadpair * 8 + output_col] = accum[0][i]; - } - } - } - } else { -#pragma unroll - for (int i = 0; i < 8; ++i) { - const int output_row = (i & 2) | ((lane & 16) ? 4 : 0) | (lane & 1); - const int output_col = - (i & 1) | (((lane >> 1) & 1) << 1) | ((i >> 2) << 2); - partials[warp][output_row * 32 + quadpair * 8 + output_col] = accum[0][i]; - } - } - __syncthreads(); - - constexpr int kOutputElements = M1Only ? 32 : 256; - for (int element = threadIdx.x; element < kOutputElements; - element += blockDim.x) { - float value = 0.0f; -#pragma unroll - for (int k_warp = 0; k_warp < SplitK; ++k_warp) { - value += partials[k_warp][element]; - } - if constexpr (M1Only) { - const int output_col = tile * 32 + element; - if constexpr (SplitOutputs) { - if (output_col < qkv_n) { - output[output_col] = __float2half(value); - } else { - z_output[output_col - qkv_n] = __float2half(value); - } - } else { - output[output_col] = __float2half(value); - } - } else { - const int output_row = element >> 5; - const int output_col = element & 31; - if (output_row < m) { - output[static_cast(output_row) * n + tile * 32 + output_col] = - __float2half(value); - } - } - } -} - -template -void launch_fp8_qpn8_sm70(const uint8_t* codes, const half* group_scales, - const half* input, half* output, int n, int k, int m, - bool channel_scales, cudaStream_t stream) { - fp8_qpn8_sm70_kernel - <<<(n / 32), (32 * SplitK), 0, stream>>>( - codes, group_scales, input, output, nullptr, nullptr, nullptr, - nullptr, nullptr, 0, n, n, k, m, channel_scales); -} - -void launch_fp8_qpn8_ba_split_sm70(const uint8_t* codes, - const half* group_scales, const half* input, - half* qkv_output, half* z_output, - const half* ba_weight, half* b_output, - half* a_output, int ba_n, int qkv_n, int n, - int k, bool channel_scales, - cudaStream_t stream) { - fp8_qpn8_sm70_kernel<16, 2, true, false, true, true, true> - <<<(n / 32 + (ba_n + 1) / 2), 512, 0, stream>>>( - codes, group_scales, input, qkv_output, z_output, ba_weight, nullptr, - b_output, a_output, ba_n, qkv_n, n, k, 1, channel_scales); -} - -template -__global__ void fp8_qpn8_gated_pair_sm70_kernel( - const uint8_t* __restrict__ codes, const half* __restrict__ group_scales, - const half* __restrict__ input, half* __restrict__ output, int hidden, - int k, int m, bool channel_scales) { - __shared__ float partials[2][SplitK][M1Only ? 32 : 256]; - - const int lane = threadIdx.x & 31; - const int warp_in_block = threadIdx.x >> 5; - const int projection = warp_in_block / SplitK; - const int warp = warp_in_block - projection * SplitK; - const int hidden_tiles = hidden >> 5; - const int tiles_n32 = hidden_tiles * 2; - const int tile = blockIdx.x + projection * hidden_tiles; - 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_warp = groups_k16 / SplitK; - const int group_begin = warp * groups_per_warp; - const uint4* code_ptr = reinterpret_cast(codes) + - static_cast(tile) * groups_k16 * 32 + lane; - const half* scale_ptr = group_scales + tile; - - float accum[NAcc][8]; -#pragma unroll - for (int chain = 0; chain < NAcc; ++chain) { -#pragma unroll - for (int i = 0; i < 8; ++i) { - accum[chain][i] = 0.0f; - } - } - int loaded_scale_group = -1; - half loaded_scale = - channel_scales - ? __ldg(group_scales + tile * 32 + qpn8_col_from_lane(lane)) - : __float2half(0.0f); - uint4 prefetched = make_uint4(0, 0, 0, 0); - if constexpr (PrefetchCodes) { - prefetched = __ldcs(code_ptr + static_cast(group_begin) * 32); - } - -#pragma unroll 4 - for (int group = group_begin; group < group_begin + groups_per_warp; - ++group) { - const int scale_group = group >> 3; - if (!channel_scales && scale_group != loaded_scale_group) { - loaded_scale = - __ldg(scale_ptr + static_cast(scale_group) * tiles_n32); - loaded_scale_group = scale_group; - } - - const uint4 packed = - PrefetchCodes ? prefetched - : __ldcs(code_ptr + static_cast(group) * 32); - uint4 next = make_uint4(0, 0, 0, 0); - if constexpr (PrefetchCodes) { - if (group + 1 < group_begin + groups_per_warp) { - next = __ldcs(code_ptr + static_cast(group + 1) * 32); - } - } - half2 weights[8]; - if constexpr (FastDecoder) { - fp8x8_to_half2x4_fast(make_uint2(packed.x, packed.y), weights); - fp8x8_to_half2x4_fast(make_uint2(packed.z, packed.w), weights + 4); - } else { - fp8x8_to_half2x4(make_uint2(packed.x, packed.y), weights); - fp8x8_to_half2x4(make_uint2(packed.z, packed.w), weights + 4); - } - const half2 scale2 = __halves2half2(loaded_scale, loaded_scale); -#pragma unroll - for (int i = 0; i < 8; ++i) { - weights[i] = __hmul2(weights[i], scale2); - } - - uint4 input01 = make_uint4(0, 0, 0, 0); - uint4 input23 = make_uint4(0, 0, 0, 0); - if (row < m) { - const half* input_row = input + static_cast(row) * k; - 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); - 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]); - if constexpr (PrefetchCodes) { - prefetched = next; - } - } - -#pragma unroll - for (int chain = 1; chain < NAcc; ++chain) { -#pragma unroll - for (int i = 0; i < 8; ++i) { - accum[0][i] += accum[chain][i]; - } - } - if constexpr (M1Only) { - if ((lane & 17) == 0) { -#pragma unroll - for (int pair = 0; pair < 2; ++pair) { -#pragma unroll - for (int offset = 0; offset < 2; ++offset) { - const int i = pair * 4 + offset; - const int output_col = - offset | (((lane >> 1) & 1) << 1) | (pair << 2); - partials[projection][warp][quadpair * 8 + output_col] = accum[0][i]; - } - } - } - } else { -#pragma unroll - for (int i = 0; i < 8; ++i) { - const int output_row = (i & 2) | ((lane & 16) ? 4 : 0) | (lane & 1); - const int output_col = - (i & 1) | (((lane >> 1) & 1) << 1) | ((i >> 2) << 2); - partials[projection][warp][output_row * 32 + quadpair * 8 + output_col] = - accum[0][i]; - } - } - __syncthreads(); - - constexpr int kOutputElements = M1Only ? 32 : 256; - for (int element = threadIdx.x; element < kOutputElements; - element += blockDim.x) { - float gate = 0.0f; - float up = 0.0f; -#pragma unroll - for (int k_warp = 0; k_warp < SplitK; ++k_warp) { - gate += partials[0][k_warp][element]; - up += partials[1][k_warp][element]; - } - if constexpr (M1Only) { - const float silu = gate / (1.0f + __expf(-gate)); - output[blockIdx.x * 32 + element] = __float2half(silu * up); - } else { - const int output_row = element >> 5; - const int output_col = element & 31; - if (output_row < m) { - const float silu = gate / (1.0f + __expf(-gate)); - output[static_cast(output_row) * hidden + blockIdx.x * 32 + - output_col] = __float2half(silu * up); - } - } - } -} - -template -void launch_fp8_qpn8_gated_pair_sm70(const uint8_t* codes, - const half* group_scales, - const half* input, half* output, - int hidden, int k, int m, - bool channel_scales, cudaStream_t stream) { - fp8_qpn8_gated_pair_sm70_kernel - <<<(hidden / 32), (64 * SplitK), 0, stream>>>( - codes, group_scales, input, output, hidden, k, m, channel_scales); -} - -} // namespace - -std::vector fp8_qpn8_prepare_sm70(torch::Tensor qweight, - torch::Tensor scales) { - TORCH_CHECK(qweight.is_cuda() && scales.is_cuda(), - "fp8_qpn8_prepare_sm70: tensors must be CUDA tensors"); - TORCH_CHECK(qweight.scalar_type() == at::ScalarType::Float8_e4m3fn, - "fp8_qpn8_prepare_sm70: weight must be float8_e4m3fn"); - TORCH_CHECK(scales.scalar_type() == torch::kFloat32, - "fp8_qpn8_prepare_sm70: scales must be float32"); - TORCH_CHECK(qweight.dim() == 2 && scales.dim() == 2, - "fp8_qpn8_prepare_sm70: tensors must be 2D"); - TORCH_CHECK(qweight.is_contiguous() && scales.is_contiguous(), - "fp8_qpn8_prepare_sm70: tensors must be contiguous"); - TORCH_CHECK(qweight.get_device() == scales.get_device(), - "fp8_qpn8_prepare_sm70: tensors must share one device"); - - 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 || block_scales, - "fp8_qpn8_prepare_sm70: expected channel scales [N, 1] or " - "block scales [N/128, K/128]"); - - const at::cuda::OptionalCUDAGuard device_guard(device_of(qweight)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - auto codes = torch::empty( - {k, n}, - torch::TensorOptions().device(qweight.device()).dtype(torch::kUInt8)); - auto group_scales = torch::empty( - {channel_scales ? 1 : k / 128, channel_scales ? n : n / 32}, - torch::TensorOptions().device(qweight.device()).dtype(torch::kFloat16)); - - const int64_t weight_numel = n * k; - const int weight_blocks = static_cast( - (weight_numel + kQpn8PrepareThreads - 1) / kQpn8PrepareThreads); - fp8_qpn8_prepack_sm70_kernel<<>>( - codes.data_ptr(), - reinterpret_cast(qweight.data_ptr()), static_cast(n), - static_cast(k)); - - const int64_t scale_numel = group_scales.numel(); - const int scale_blocks = static_cast( - (scale_numel + kQpn8PrepareThreads - 1) / kQpn8PrepareThreads); - if (channel_scales) { - fp8_qpn8_channel_scale_sm70_kernel<<>>( - reinterpret_cast(group_scales.data_ptr()), - scales.data_ptr(), static_cast(n)); - } else { - fp8_qpn8_scale_sm70_kernel<<>>( - reinterpret_cast(group_scales.data_ptr()), - scales.data_ptr(), static_cast(n / 128), - static_cast(k / 128)); - } - C10_CUDA_KERNEL_LAUNCH_CHECK(); - return {codes, group_scales}; -} - -void fp8_qpn8_dequantize_sm70_out(torch::Tensor out, torch::Tensor codes, - torch::Tensor group_scales) { - TORCH_CHECK(out.is_cuda() && codes.is_cuda() && group_scales.is_cuda(), - "fp8_qpn8_dequantize_sm70_out: tensors must be CUDA tensors"); - TORCH_CHECK(out.scalar_type() == torch::kFloat16 && - codes.scalar_type() == torch::kUInt8 && - group_scales.scalar_type() == torch::kFloat16, - "fp8_qpn8_dequantize_sm70_out: dtype mismatch"); - TORCH_CHECK(out.dim() == 2 && codes.dim() == 2 && group_scales.dim() == 2, - "fp8_qpn8_dequantize_sm70_out: tensors must be 2D"); - TORCH_CHECK(out.is_contiguous() && codes.is_contiguous() && - group_scales.is_contiguous(), - "fp8_qpn8_dequantize_sm70_out: tensors must be contiguous"); - TORCH_CHECK(out.get_device() == codes.get_device() && - out.get_device() == group_scales.get_device(), - "fp8_qpn8_dequantize_sm70_out: tensors must share one device"); - - 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 || block_scales, - "fp8_qpn8_dequantize_sm70_out: scale shape mismatch"); - - const at::cuda::OptionalCUDAGuard device_guard(device_of(out)); - const int64_t word_count = k * n / 16; - const int blocks = static_cast((word_count + kQpn8PrepareThreads - 1) / - kQpn8PrepareThreads); - fp8_qpn8_dequantize_sm70_kernel<<>>( - reinterpret_cast(out.data_ptr()), - codes.data_ptr(), - reinterpret_cast(group_scales.data_ptr()), - static_cast(n), static_cast(k), channel_scales); - C10_CUDA_KERNEL_LAUNCH_CHECK(); -} - -void fp8_qpn8_prefill_sm70_out(torch::Tensor out, int64_t dense_weight_ptr, - torch::Tensor input, torch::Tensor codes, - torch::Tensor group_scales, bool gated_silu) { - TORCH_CHECK(input.is_cuda() && out.is_cuda(), - "fp8_qpn8_prefill_sm70_out: input and output must be CUDA"); - TORCH_CHECK(input.scalar_type() == torch::kFloat16 && - out.scalar_type() == torch::kFloat16, - "fp8_qpn8_prefill_sm70_out: input and output must be float16"); - TORCH_CHECK(input.dim() == 2 && out.dim() == 2 && dense_weight_ptr != 0, - "fp8_qpn8_prefill_sm70_out: invalid input or workspace"); - TORCH_CHECK(input.is_contiguous() && out.is_contiguous(), - "fp8_qpn8_prefill_sm70_out: tensors must be contiguous"); - - const int64_t m = input.size(0); - const int64_t k = input.size(1); - const int64_t n = codes.size(1); - TORCH_CHECK(codes.dim() == 2 && codes.size(0) == k, - "fp8_qpn8_prefill_sm70_out: code shape mismatch"); - TORCH_CHECK(out.size(0) == m && out.size(1) == (gated_silu ? n / 2 : n), - "fp8_qpn8_prefill_sm70_out: output shape mismatch"); - - auto dense_weight = torch::from_blob( - reinterpret_cast(dense_weight_ptr), {k, n}, input.options()); - fp8_qpn8_dequantize_sm70_out(dense_weight, codes, group_scales); - if (!gated_silu) { - at::mm_out(out, input, dense_weight); - return; - } - - auto gate_up = at::mm(input, dense_weight); - constexpr int kThreads = 256; - fp8_qpn8_silu_and_mul_sm70_kernel<<(m), kThreads, 0, - at::cuda::getCurrentCUDAStream()>>>( - reinterpret_cast(out.data_ptr()), - reinterpret_cast(gate_up.data_ptr()), - static_cast(m), static_cast(n / 2)); - C10_CUDA_KERNEL_LAUNCH_CHECK(); -} - -void fp8_qpn8_gemm_sm70_out(torch::Tensor out, torch::Tensor input, - torch::Tensor codes, torch::Tensor group_scales, - int64_t split_k, int64_t accumulator_chains, - bool fast_decoder, bool prefetch_codes) { - TORCH_CHECK(input.is_cuda() && out.is_cuda() && codes.is_cuda() && - group_scales.is_cuda(), - "fp8_qpn8_gemm_sm70_out: tensors must be CUDA tensors"); - TORCH_CHECK(input.scalar_type() == torch::kFloat16 && - out.scalar_type() == torch::kFloat16, - "fp8_qpn8_gemm_sm70_out: input and output must be float16"); - TORCH_CHECK(codes.scalar_type() == torch::kUInt8, - "fp8_qpn8_gemm_sm70_out: codes must be uint8"); - TORCH_CHECK(group_scales.scalar_type() == torch::kFloat16, - "fp8_qpn8_gemm_sm70_out: group scales must be float16"); - TORCH_CHECK(input.dim() == 2 && out.dim() == 2 && group_scales.dim() == 2, - "fp8_qpn8_gemm_sm70_out: expected 2D tensors"); - TORCH_CHECK(input.is_contiguous() && out.is_contiguous() && - codes.is_contiguous() && group_scales.is_contiguous(), - "fp8_qpn8_gemm_sm70_out: tensors must be contiguous"); - TORCH_CHECK(input.get_device() == out.get_device() && - input.get_device() == codes.get_device() && - input.get_device() == group_scales.get_device(), - "fp8_qpn8_gemm_sm70_out: tensors must share one device"); - - const int64_t m = input.size(0); - const int64_t k = input.size(1); - const int64_t n = out.size(1); - TORCH_CHECK(m >= 1 && m <= 8, "fp8_qpn8_gemm_sm70_out: M must be in [1, 8]"); - 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 || block_scales, - "fp8_qpn8_gemm_sm70_out: scale shape mismatch"); - TORCH_CHECK(split_k == 4 || split_k == 8 || split_k == 12 || split_k == 16 || - split_k == 32, - "fp8_qpn8_gemm_sm70_out: unsupported split_k"); - TORCH_CHECK((k / 16) % split_k == 0, - "fp8_qpn8_gemm_sm70_out: K/16 must be divisible by split_k"); - TORCH_CHECK(accumulator_chains == 1 || accumulator_chains == 2, - "fp8_qpn8_gemm_sm70_out: accumulator_chains must be 1 or 2"); - TORCH_CHECK(!prefetch_codes || fast_decoder, - "fp8_qpn8_gemm_sm70_out: prefetch experiment requires the " - "fast decoder"); - const bool exact_shape_split = split_k == 12; - TORCH_CHECK(!exact_shape_split || - (accumulator_chains == 2 && fast_decoder && !prefetch_codes), - "fp8_qpn8_gemm_sm70_out: split_k 12 requires the " - "fast decoder, two accumulator chains, and no prefetch"); - - const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - const auto* code_ptr = codes.data_ptr(); - const auto* scale_ptr = - reinterpret_cast(group_scales.data_ptr()); - const auto* input_ptr = - reinterpret_cast(input.data_ptr()); - auto* output_ptr = reinterpret_cast(out.data_ptr()); - - // Keep the M=1 reduction order restricted to the two tuned output/down - // projections. Extending it to the other split variants changed the frozen - // random-sampling token stream without an end-to-end decode win. - if (m == 1 && accumulator_chains == 2 && fast_decoder && !prefetch_codes && - (split_k == 12 || split_k == 16)) { - if (split_k == 12) { - launch_fp8_qpn8_sm70<12, 2, true, false, true>( - code_ptr, scale_ptr, input_ptr, output_ptr, static_cast(n), - static_cast(k), 1, channel_scales, stream); - } else { - launch_fp8_qpn8_sm70<16, 2, true, false, true>( - code_ptr, scale_ptr, input_ptr, output_ptr, static_cast(n), - static_cast(k), 1, channel_scales, stream); - } - C10_CUDA_KERNEL_LAUNCH_CHECK(); - return; - } - -#define VLLM_LAUNCH_QPN8(SPLIT, NACC, FAST, PREFETCH) \ - launch_fp8_qpn8_sm70( \ - code_ptr, scale_ptr, input_ptr, output_ptr, static_cast(n), \ - static_cast(k), static_cast(m), channel_scales, stream) - - if (prefetch_codes) { - if (split_k == 4 && accumulator_chains == 1) { - VLLM_LAUNCH_QPN8(4, 1, true, true); - } else if (split_k == 4 && accumulator_chains == 2) { - VLLM_LAUNCH_QPN8(4, 2, true, true); - } else if (split_k == 8 && accumulator_chains == 1) { - VLLM_LAUNCH_QPN8(8, 1, true, true); - } else if (split_k == 8 && accumulator_chains == 2) { - VLLM_LAUNCH_QPN8(8, 2, true, true); - } else if (split_k == 16 && accumulator_chains == 1) { - VLLM_LAUNCH_QPN8(16, 1, true, true); - } else if (split_k == 16 && accumulator_chains == 2) { - VLLM_LAUNCH_QPN8(16, 2, true, true); - } else if (split_k == 32 && accumulator_chains == 1) { - VLLM_LAUNCH_QPN8(32, 1, true, true); - } else { - VLLM_LAUNCH_QPN8(32, 2, true, true); - } - } else if (fast_decoder) { - if (split_k == 4 && accumulator_chains == 1) { - VLLM_LAUNCH_QPN8(4, 1, true, false); - } else if (split_k == 4 && accumulator_chains == 2) { - VLLM_LAUNCH_QPN8(4, 2, true, false); - } else if (split_k == 8 && accumulator_chains == 1) { - VLLM_LAUNCH_QPN8(8, 1, true, false); - } else if (split_k == 8 && accumulator_chains == 2) { - VLLM_LAUNCH_QPN8(8, 2, true, false); - } else if (split_k == 12) { - VLLM_LAUNCH_QPN8(12, 2, true, false); - } else if (split_k == 16 && accumulator_chains == 1) { - VLLM_LAUNCH_QPN8(16, 1, true, false); - } else if (split_k == 16 && accumulator_chains == 2) { - VLLM_LAUNCH_QPN8(16, 2, true, false); - } else if (split_k == 32 && accumulator_chains == 1) { - VLLM_LAUNCH_QPN8(32, 1, true, false); - } else { - VLLM_LAUNCH_QPN8(32, 2, true, false); - } - } else if (split_k == 4 && accumulator_chains == 1) { - VLLM_LAUNCH_QPN8(4, 1, false, false); - } else if (split_k == 4 && accumulator_chains == 2) { - VLLM_LAUNCH_QPN8(4, 2, false, false); - } else if (split_k == 8 && accumulator_chains == 1) { - VLLM_LAUNCH_QPN8(8, 1, false, false); - } else if (split_k == 8 && accumulator_chains == 2) { - VLLM_LAUNCH_QPN8(8, 2, false, false); - } else if (split_k == 16 && accumulator_chains == 1) { - VLLM_LAUNCH_QPN8(16, 1, false, false); - } else if (split_k == 16 && accumulator_chains == 2) { - VLLM_LAUNCH_QPN8(16, 2, false, false); - } else if (split_k == 32 && accumulator_chains == 1) { - VLLM_LAUNCH_QPN8(32, 1, false, false); - } else { - VLLM_LAUNCH_QPN8(32, 2, false, false); - } -#undef VLLM_LAUNCH_QPN8 - - C10_CUDA_KERNEL_LAUNCH_CHECK(); -} - -void fp8_qpn8_gemm_ba_split_sm70_out(torch::Tensor qkv_out, torch::Tensor z_out, - torch::Tensor b_out, torch::Tensor a_out, - torch::Tensor input, torch::Tensor codes, - torch::Tensor group_scales, - torch::Tensor ba_weight) { - TORCH_CHECK(qkv_out.is_cuda() && z_out.is_cuda() && b_out.is_cuda() && - a_out.is_cuda() && input.is_cuda() && codes.is_cuda() && - group_scales.is_cuda() && ba_weight.is_cuda(), - "fp8_qpn8_gemm_ba_split_sm70_out: tensors must be CUDA tensors"); - TORCH_CHECK(qkv_out.scalar_type() == torch::kFloat16 && - z_out.scalar_type() == torch::kFloat16 && - b_out.scalar_type() == torch::kFloat16 && - a_out.scalar_type() == torch::kFloat16 && - input.scalar_type() == torch::kFloat16 && - group_scales.scalar_type() == torch::kFloat16 && - ba_weight.scalar_type() == torch::kFloat16 && - codes.scalar_type() == torch::kUInt8, - "fp8_qpn8_gemm_ba_split_sm70_out: dtype mismatch"); - TORCH_CHECK(qkv_out.is_contiguous() && z_out.is_contiguous() && - b_out.is_contiguous() && a_out.is_contiguous() && - input.is_contiguous() && codes.is_contiguous() && - group_scales.is_contiguous() && ba_weight.is_contiguous(), - "fp8_qpn8_gemm_ba_split_sm70_out: tensors must be contiguous"); - TORCH_CHECK(qkv_out.get_device() == z_out.get_device() && - qkv_out.get_device() == b_out.get_device() && - qkv_out.get_device() == a_out.get_device() && - qkv_out.get_device() == input.get_device() && - qkv_out.get_device() == codes.get_device() && - qkv_out.get_device() == group_scales.get_device() && - 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; - 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(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]"); - 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"); - TORCH_CHECK(ba_weight.dim() == 2 && ba_weight.size(0) == ba_n && - ba_weight.size(1) == k, - "fp8_qpn8_gemm_ba_split_sm70_out: b/a weight shape mismatch"); - - const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - launch_fp8_qpn8_ba_split_sm70( - codes.data_ptr(), - reinterpret_cast(group_scales.data_ptr()), - reinterpret_cast(input.data_ptr()), - reinterpret_cast(qkv_out.data_ptr()), - reinterpret_cast(z_out.data_ptr()), - reinterpret_cast(ba_weight.data_ptr()), - reinterpret_cast(b_out.data_ptr()), - reinterpret_cast(a_out.data_ptr()), - static_cast(ba_n), static_cast(qkv_n), static_cast(n), - static_cast(k), true, stream); - C10_CUDA_KERNEL_LAUNCH_CHECK(); -} - -void fp8_qpn8_dispatch_sm70_out(torch::Tensor out, int64_t dense_weight_ptr, - torch::Tensor input, torch::Tensor codes, - torch::Tensor group_scales, int64_t split_k, - int64_t accumulator_chains, bool prefetch_codes, - bool gated_silu); - -void fp8_qpn8_dispatch_ba_split_sm70_out( - torch::Tensor qkv_out, torch::Tensor z_out, torch::Tensor b_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; - 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, - "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() && - ba_staging.is_cuda() && input.is_cuda() && codes.is_cuda() && - group_scales.is_cuda() && ba_weight.is_cuda(), - "fp8_qpn8_dispatch_ba_split_sm70_out: tensors must be CUDA"); - TORCH_CHECK(qkv_out.scalar_type() == torch::kFloat16 && - z_out.scalar_type() == torch::kFloat16 && - b_out.scalar_type() == torch::kFloat16 && - a_out.scalar_type() == torch::kFloat16 && - qkvz_staging.scalar_type() == torch::kFloat16 && - ba_staging.scalar_type() == torch::kFloat16 && - input.scalar_type() == torch::kFloat16 && - group_scales.scalar_type() == torch::kFloat16 && - ba_weight.scalar_type() == torch::kFloat16 && - codes.scalar_type() == torch::kUInt8, - "fp8_qpn8_dispatch_ba_split_sm70_out: dtype mismatch"); - TORCH_CHECK( - qkv_out.is_contiguous() && z_out.is_contiguous() && - b_out.is_contiguous() && a_out.is_contiguous() && - qkvz_staging.is_contiguous() && ba_staging.is_contiguous() && - input.is_contiguous() && codes.is_contiguous() && - group_scales.is_contiguous() && ba_weight.is_contiguous(), - "fp8_qpn8_dispatch_ba_split_sm70_out: tensors must be contiguous"); - TORCH_CHECK( - qkv_out.numel() == m * qkv_n && z_out.numel() == m * z_n && - b_out.numel() == m * ba_n / 2 && a_out.numel() == m * ba_n / 2 && - qkvz_staging.numel() == m * n && ba_staging.numel() == m * ba_n, - "fp8_qpn8_dispatch_ba_split_sm70_out: output shape mismatch"); - TORCH_CHECK( - qkv_out.get_device() == z_out.get_device() && - qkv_out.get_device() == b_out.get_device() && - qkv_out.get_device() == a_out.get_device() && - qkv_out.get_device() == qkvz_staging.get_device() && - qkv_out.get_device() == ba_staging.get_device() && - qkv_out.get_device() == input.get_device() && - qkv_out.get_device() == codes.get_device() && - qkv_out.get_device() == group_scales.get_device() && - qkv_out.get_device() == ba_weight.get_device(), - "fp8_qpn8_dispatch_ba_split_sm70_out: tensors must share one device"); - TORCH_CHECK(codes.dim() == 2 && codes.size(0) == k && codes.size(1) == n, - "fp8_qpn8_dispatch_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_dispatch_ba_split_sm70_out: scale shape mismatch"); - TORCH_CHECK(ba_weight.dim() == 2 && ba_weight.size(0) == ba_n && - ba_weight.size(1) == k, - "fp8_qpn8_dispatch_ba_split_sm70_out: b/a weight mismatch"); - TORCH_CHECK(m <= 8 || dense_weight_ptr != 0, - "fp8_qpn8_dispatch_ba_split_sm70_out: large-M requires the " - "dense QPN8 workspace"); - - if (m == 1) { - 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 " - "enabled.\n"); - std::fflush(stderr); - }); - fp8_qpn8_gemm_ba_split_sm70_out(qkv_out, z_out, b_out, a_out, input, codes, - group_scales, ba_weight); - return; - } - - fp8_qpn8_dispatch_sm70_out(qkvz_staging, dense_weight_ptr, input, codes, - group_scales, 16, 2, false, false); - at::mm_out(ba_staging, input, ba_weight.t()); - const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); - constexpr int threads = 256; - const int64_t elements = m * (n + ba_n); - const int blocks = static_cast((elements + threads - 1) / threads); - fp8_qpn8_ba_split_copy_sm70_kernel<<>>( - reinterpret_cast(qkvz_staging.data_ptr()), - reinterpret_cast(ba_staging.data_ptr()), - reinterpret_cast(qkv_out.data_ptr()), - reinterpret_cast(z_out.data_ptr()), - reinterpret_cast(b_out.data_ptr()), - reinterpret_cast(a_out.data_ptr()), static_cast(m)); - C10_CUDA_KERNEL_LAUNCH_CHECK(); -} - -void fp8_qpn8_gated_pair_sm70_out(torch::Tensor out, torch::Tensor input, - torch::Tensor codes, - torch::Tensor group_scales, int64_t split_k, - int64_t accumulator_chains, bool fast_decoder, - bool prefetch_codes) { - TORCH_CHECK(out.is_cuda() && input.is_cuda() && codes.is_cuda() && - group_scales.is_cuda(), - "fp8_qpn8_gated_pair_sm70_out: tensors must be CUDA tensors"); - TORCH_CHECK(out.scalar_type() == torch::kFloat16 && - input.scalar_type() == torch::kFloat16 && - group_scales.scalar_type() == torch::kFloat16 && - codes.scalar_type() == torch::kUInt8, - "fp8_qpn8_gated_pair_sm70_out: dtype mismatch"); - TORCH_CHECK(out.is_contiguous() && input.is_contiguous() && - codes.is_contiguous() && group_scales.is_contiguous(), - "fp8_qpn8_gated_pair_sm70_out: tensors must be contiguous"); - TORCH_CHECK(out.dim() == 2 && input.dim() == 2 && group_scales.dim() == 2, - "fp8_qpn8_gated_pair_sm70_out: expected 2D tensors"); - TORCH_CHECK(out.get_device() == input.get_device() && - out.get_device() == codes.get_device() && - out.get_device() == group_scales.get_device(), - "fp8_qpn8_gated_pair_sm70_out: tensors must share one device"); - - const int64_t m = input.size(0); - const int64_t k = input.size(1); - const int64_t hidden = out.size(1); - 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(channel_scales || block_scales, - "fp8_qpn8_gated_pair_sm70_out: scale shape mismatch"); - TORCH_CHECK(split_k == 4 || split_k == 8 || split_k == 16, - "fp8_qpn8_gated_pair_sm70_out: split_k must be 4, 8, or 16"); - TORCH_CHECK((k / 16) % split_k == 0, - "fp8_qpn8_gated_pair_sm70_out: invalid split_k for K"); - TORCH_CHECK( - accumulator_chains == 1 || accumulator_chains == 2, - "fp8_qpn8_gated_pair_sm70_out: accumulator_chains must be 1 or 2"); - TORCH_CHECK(!prefetch_codes || fast_decoder, - "fp8_qpn8_gated_pair_sm70_out: prefetch requires fast decoder"); - - const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - const auto* code_ptr = codes.data_ptr(); - const auto* scale_ptr = - reinterpret_cast(group_scales.data_ptr()); - const auto* input_ptr = - reinterpret_cast(input.data_ptr()); - auto* output_ptr = reinterpret_cast(out.data_ptr()); - - // The frozen gate/up route is split-8. Preserve the generic reduction order - // for split-4/16 rather than instantiating unaccepted M=1 variants. - if (m == 1 && split_k == 8 && accumulator_chains == 2 && fast_decoder && - !prefetch_codes) { - launch_fp8_qpn8_gated_pair_sm70<8, 2, true, false, true>( - code_ptr, scale_ptr, input_ptr, output_ptr, static_cast(hidden), - static_cast(k), 1, channel_scales, stream); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - return; - } - -#define VLLM_LAUNCH_QPN8_GATED_PAIR(SPLIT, NACC, FAST, PREFETCH) \ - launch_fp8_qpn8_gated_pair_sm70( \ - code_ptr, scale_ptr, input_ptr, output_ptr, static_cast(hidden), \ - static_cast(k), static_cast(m), channel_scales, stream) - - if (prefetch_codes) { - if (split_k == 4 && accumulator_chains == 1) { - VLLM_LAUNCH_QPN8_GATED_PAIR(4, 1, true, true); - } else if (split_k == 4 && accumulator_chains == 2) { - VLLM_LAUNCH_QPN8_GATED_PAIR(4, 2, true, true); - } else if (split_k == 8 && accumulator_chains == 1) { - VLLM_LAUNCH_QPN8_GATED_PAIR(8, 1, true, true); - } else if (split_k == 8 && accumulator_chains == 2) { - VLLM_LAUNCH_QPN8_GATED_PAIR(8, 2, true, true); - } else if (split_k == 16 && accumulator_chains == 1) { - VLLM_LAUNCH_QPN8_GATED_PAIR(16, 1, true, true); - } else { - VLLM_LAUNCH_QPN8_GATED_PAIR(16, 2, true, true); - } - } else if (fast_decoder) { - if (split_k == 4 && accumulator_chains == 1) { - VLLM_LAUNCH_QPN8_GATED_PAIR(4, 1, true, false); - } else if (split_k == 4 && accumulator_chains == 2) { - VLLM_LAUNCH_QPN8_GATED_PAIR(4, 2, true, false); - } else if (split_k == 8 && accumulator_chains == 1) { - VLLM_LAUNCH_QPN8_GATED_PAIR(8, 1, true, false); - } else if (split_k == 8 && accumulator_chains == 2) { - VLLM_LAUNCH_QPN8_GATED_PAIR(8, 2, true, false); - } else if (split_k == 16 && accumulator_chains == 1) { - VLLM_LAUNCH_QPN8_GATED_PAIR(16, 1, true, false); - } else { - VLLM_LAUNCH_QPN8_GATED_PAIR(16, 2, true, false); - } - } else if (split_k == 4 && accumulator_chains == 1) { - VLLM_LAUNCH_QPN8_GATED_PAIR(4, 1, false, false); - } else if (split_k == 4 && accumulator_chains == 2) { - VLLM_LAUNCH_QPN8_GATED_PAIR(4, 2, false, false); - } else if (split_k == 8 && accumulator_chains == 1) { - VLLM_LAUNCH_QPN8_GATED_PAIR(8, 1, false, false); - } else if (split_k == 8 && accumulator_chains == 2) { - VLLM_LAUNCH_QPN8_GATED_PAIR(8, 2, false, false); - } else if (split_k == 16 && accumulator_chains == 1) { - VLLM_LAUNCH_QPN8_GATED_PAIR(16, 1, false, false); - } else { - VLLM_LAUNCH_QPN8_GATED_PAIR(16, 2, false, false); - } -#undef VLLM_LAUNCH_QPN8_GATED_PAIR - C10_CUDA_KERNEL_LAUNCH_CHECK(); -} - -void fp8_qpn8_dispatch_sm70_out(torch::Tensor out, int64_t dense_weight_ptr, - torch::Tensor input, torch::Tensor codes, - torch::Tensor group_scales, int64_t split_k, - int64_t accumulator_chains, bool prefetch_codes, - bool gated_silu) { - // Keep the M decision inside the opaque operator. AOTInductor compiles one - // dynamic M=1..8192 range for this model, so a Python shape branch traced at - // M=1/2 would otherwise be incorrectly reused by large-M prefill. - if (input.size(0) <= 8) { - if (gated_silu) { - fp8_qpn8_gated_pair_sm70_out(out, input, codes, group_scales, split_k, - accumulator_chains, true, prefetch_codes); - } else { - fp8_qpn8_gemm_sm70_out(out, input, codes, group_scales, split_k, - accumulator_chains, true, prefetch_codes); - } - return; - } - fp8_qpn8_prefill_sm70_out(out, dense_weight_ptr, input, codes, group_scales, - gated_silu); -} - -#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. -TORCH_LIBRARY_FRAGMENT(_C, ops) { - 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( - "fp8_qpn8_dequantize_sm70_out(Tensor(a!) out, Tensor codes, " - "Tensor group_scales) -> ()"); - ops.impl("fp8_qpn8_dequantize_sm70_out", torch::kCUDA, - &fp8_qpn8_dequantize_sm70_out); - ops.def( - "fp8_qpn8_prefill_sm70_out(Tensor(a!) out, int dense_weight_ptr, " - "Tensor input, Tensor codes, Tensor group_scales, bool gated_silu) -> " - "()"); - ops.impl("fp8_qpn8_prefill_sm70_out", torch::kCUDA, - &fp8_qpn8_prefill_sm70_out); - ops.def( - "fp8_qpn8_dispatch_sm70_out(Tensor(a!) out, int dense_weight_ptr, " - "Tensor input, Tensor codes, Tensor group_scales, int split_k, " - "int accumulator_chains, bool prefetch_codes, bool gated_silu) -> ()"); - ops.impl("fp8_qpn8_dispatch_sm70_out", torch::kCUDA, - &fp8_qpn8_dispatch_sm70_out); - ops.def( - "fp8_qpn8_gemm_sm70_out(Tensor(a!) out, Tensor input, Tensor codes, " - "Tensor group_scales, int split_k, int accumulator_chains, " - "bool fast_decoder, bool prefetch_codes) -> ()"); - ops.impl("fp8_qpn8_gemm_sm70_out", torch::kCUDA, &fp8_qpn8_gemm_sm70_out); - ops.def( - "fp8_qpn8_gemm_ba_split_sm70_out(Tensor(a!) qkv_out, Tensor(b!) " - "z_out, Tensor(c!) b_out, Tensor(d!) a_out, Tensor input, Tensor codes, " - "Tensor group_scales, Tensor ba_weight) -> ()"); - ops.impl("fp8_qpn8_gemm_ba_split_sm70_out", torch::kCUDA, - &fp8_qpn8_gemm_ba_split_sm70_out); - ops.def( - "fp8_qpn8_dispatch_ba_split_sm70_out(Tensor(a!) qkv_out, Tensor(b!) " - "z_out, Tensor(c!) b_out, Tensor(d!) a_out, Tensor(e!) qkvz_staging, " - "Tensor(f!) ba_staging, int dense_weight_ptr, Tensor input, Tensor " - "codes, Tensor group_scales, Tensor ba_weight) -> ()"); - ops.impl("fp8_qpn8_dispatch_ba_split_sm70_out", torch::kCUDA, - &fp8_qpn8_dispatch_ba_split_sm70_out); - ops.def( - "fp8_qpn8_gated_pair_sm70_out(Tensor(a!) out, Tensor input, " - "Tensor codes, Tensor group_scales, int split_k, " - "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); -} -#endif diff --git a/csrc/sm70_turbomind/ops/nvfp4_qpn2_sm70.cu b/csrc/sm70_turbomind/ops/nvfp4_qpn2_sm70.cu deleted file mode 100644 index 641949e4b4..0000000000 --- a/csrc/sm70_turbomind/ops/nvfp4_qpn2_sm70.cu +++ /dev/null @@ -1,546 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project -// -// The QPN2 execution layout is derived from dnv2003/v100-skinny (MIT). -// See LICENSE.v100-skinny in this directory for the retained MIT notice. - -#include -#include - -#include -#include -#include -#include -#include - -#ifndef VLLM_NVFP4_QPN2_STANDALONE -void silu_and_mul(torch::Tensor& out, torch::Tensor& input); - -void nvfp4_gemm_sm70_out(torch::Tensor out, torch::Tensor input, - torch::Tensor weight, torch::Tensor scales, - int64_t group_size, int64_t k_ld, int64_t q_ld, - bool gated_silu); -#endif - -namespace { - -constexpr int kPrepareThreads = 256; - -__device__ __forceinline__ int qpn2_col_from_lane(int lane) { - return ((lane >> 2) & 3) * 8 + (lane & 3) + ((lane & 16) ? 4 : 0); -} - -__device__ __forceinline__ int qpn2_logical_k(int physical_k) { - const int local = physical_k & 7; - return (physical_k & 8) + ((local & 3) << 1) + (local >> 2); -} - -__global__ void nvfp4_qpn2_prepack_codes_kernel( - uint8_t* __restrict__ output, const uint8_t* __restrict__ weight, int n, - int k) { - const size_t index = - static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - const size_t numel = static_cast(n) * k / 2; - if (index >= numel) { - return; - } - - const int physical_byte = static_cast(index & 7); - size_t outer = index >> 3; - const int lane = static_cast(outer & 31); - outer >>= 5; - const int groups_k16 = k >> 4; - const int group = static_cast(outer % groups_k16); - const int tile = static_cast(outer / groups_k16); - const int column = tile * 32 + qpn2_col_from_lane(lane); - const int logical_k0 = group * 16 + qpn2_logical_k(physical_byte * 2); - const int logical_k1 = group * 16 + qpn2_logical_k(physical_byte * 2 + 1); - const int k_bytes = k >> 1; - const uint8_t packed0 = - weight[static_cast(column) * k_bytes + (logical_k0 >> 1)]; - const uint8_t packed1 = - weight[static_cast(column) * k_bytes + (logical_k1 >> 1)]; - const uint8_t code0 = - static_cast((packed0 >> ((logical_k0 & 1) * 4)) & 0x0f); - const uint8_t code1 = - static_cast((packed1 >> ((logical_k1 & 1) * 4)) & 0x0f); - output[index] = static_cast(code0 | (code1 << 4)); -} - -__global__ void nvfp4_qpn2_prepack_scales_kernel( - uint8_t* __restrict__ output, const uint8_t* __restrict__ scales, int n, - int k) { - const size_t index = - static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - const size_t numel = static_cast(n) * k / 16; - if (index >= numel) { - return; - } - - const int lane = static_cast(index & 31); - size_t outer = index >> 5; - const int groups_k16 = k >> 4; - const int group = static_cast(outer % groups_k16); - const int tile = static_cast(outer / groups_k16); - const int column = tile * 32 + qpn2_col_from_lane(lane); - output[index] = scales[static_cast(column) * groups_k16 + group]; -} - -__device__ __forceinline__ half2 fp8e4m3_to_half2(uint8_t value) { - const unsigned short bits = - ((static_cast(value) & 0x80u) << 8) | - ((static_cast(value) & 0x7fu) << 7); - const half converted = - __hmul(__ushort_as_half(bits), __ushort_as_half(0x5c00)); - return __halves2half2(converted, converted); -} - -__device__ __forceinline__ void dequant_e2m1x8(unsigned packed, half2 scale, - half2 output[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) { - output[index] = __hmul2(*reinterpret_cast(&values[index]), scale); - } -} - -#define VLLM_SM70_QPN2_MMA(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 nvfp4_qpn2_sm70_kernel(const uint8_t* __restrict__ codes, - const uint8_t* __restrict__ group_scales, - const half* __restrict__ input, - half* __restrict__ output, int n, int k, - int m, float global_scale) { - __shared__ float partials[SplitK][256]; - - const int lane = threadIdx.x & 31; - const int warp = threadIdx.x >> 5; - const int tile = blockIdx.x; - 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_warp = groups_k16 / SplitK; - const int group_begin = warp * groups_per_warp; - const uint2* code_ptr = reinterpret_cast(codes) + - static_cast(tile) * groups_k16 * 32 + lane; - const uint8_t* scale_ptr = - group_scales + static_cast(tile) * groups_k16 * 32 + lane; - const half2 global_scale2 = __float2half2_rn(global_scale * 16384.0f); - - 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_warp; - ++group) { - const uint2 packed = __ldcs(code_ptr + static_cast(group) * 32); - const half2 scale = __hmul2( - fp8e4m3_to_half2(__ldg(scale_ptr + static_cast(group) * 32)), - global_scale2); - half2 weights[8]; - dequant_e2m1x8(packed.x, scale, weights); - dequant_e2m1x8(packed.y, scale, weights + 4); - - uint4 input01 = make_uint4(0, 0, 0, 0); - uint4 input23 = make_uint4(0, 0, 0, 0); - if (row < m) { - const half* input_row = input + static_cast(row) * k; - 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); - const unsigned* b = reinterpret_cast(weights); - VLLM_SM70_QPN2_MMA(accum[0], a0[0], a0[1], b[0], b[1]); - VLLM_SM70_QPN2_MMA(accum[1 % NAcc], a0[2], a0[3], b[2], b[3]); - VLLM_SM70_QPN2_MMA(accum[2 % NAcc], a1[0], a1[1], b[4], b[5]); - VLLM_SM70_QPN2_MMA(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]; - } - } - -#pragma unroll - for (int index = 0; index < 8; ++index) { - const int output_row = (index & 2) | ((lane & 16) ? 4 : 0) | (lane & 1); - const int output_col = - (index & 1) | (((lane >> 1) & 1) << 1) | ((index >> 2) << 2); - partials[warp][output_row * 32 + quadpair * 8 + output_col] = - accum[0][index]; - } - __syncthreads(); - - for (int element = threadIdx.x; element < 256; element += blockDim.x) { - float value = 0.0f; -#pragma unroll - for (int k_warp = 0; k_warp < SplitK; ++k_warp) { - value += partials[k_warp][element]; - } - const int output_row = element >> 5; - const int output_col = element & 31; - if (output_row < m) { - output[static_cast(output_row) * n + tile * 32 + output_col] = - __float2half(value); - } - } -} - -template -__global__ void nvfp4_qpn2_gated_sm70_kernel( - const uint8_t* __restrict__ codes, const uint8_t* __restrict__ group_scales, - const half* __restrict__ input, half* __restrict__ output, int hidden, - int k, int m, float global_scale) { - __shared__ float partials[2][SplitK][256]; - - const int lane = threadIdx.x & 31; - const int warp_in_block = threadIdx.x >> 5; - const int projection = warp_in_block / SplitK; - const int warp = warp_in_block - projection * SplitK; - const int hidden_tiles = hidden >> 5; - const int tile = blockIdx.x + projection * hidden_tiles; - 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_warp = groups_k16 / SplitK; - const int group_begin = warp * groups_per_warp; - const uint2* code_ptr = reinterpret_cast(codes) + - static_cast(tile) * groups_k16 * 32 + lane; - const uint8_t* scale_ptr = - group_scales + static_cast(tile) * groups_k16 * 32 + lane; - const half2 global_scale2 = __float2half2_rn(global_scale * 16384.0f); - - 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_warp; - ++group) { - const uint2 packed = __ldcs(code_ptr + static_cast(group) * 32); - const half2 scale = __hmul2( - fp8e4m3_to_half2(__ldg(scale_ptr + static_cast(group) * 32)), - global_scale2); - half2 weights[8]; - dequant_e2m1x8(packed.x, scale, weights); - dequant_e2m1x8(packed.y, scale, weights + 4); - - uint4 input01 = make_uint4(0, 0, 0, 0); - uint4 input23 = make_uint4(0, 0, 0, 0); - if (row < m) { - const half* input_row = input + static_cast(row) * k; - 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); - const unsigned* b = reinterpret_cast(weights); - VLLM_SM70_QPN2_MMA(accum[0], a0[0], a0[1], b[0], b[1]); - VLLM_SM70_QPN2_MMA(accum[1 % NAcc], a0[2], a0[3], b[2], b[3]); - VLLM_SM70_QPN2_MMA(accum[2 % NAcc], a1[0], a1[1], b[4], b[5]); - VLLM_SM70_QPN2_MMA(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]; - } - } - -#pragma unroll - for (int index = 0; index < 8; ++index) { - const int output_row = (index & 2) | ((lane & 16) ? 4 : 0) | (lane & 1); - const int output_col = - (index & 1) | (((lane >> 1) & 1) << 1) | ((index >> 2) << 2); - partials[projection][warp][output_row * 32 + quadpair * 8 + output_col] = - accum[0][index]; - } - __syncthreads(); - - for (int element = threadIdx.x; element < 256; element += blockDim.x) { - float gate = 0.0f; - float up = 0.0f; -#pragma unroll - for (int k_warp = 0; k_warp < SplitK; ++k_warp) { - gate += partials[0][k_warp][element]; - up += partials[1][k_warp][element]; - } - const int output_row = element >> 5; - const int output_col = element & 31; - if (output_row < m) { - // Match the existing SM70 silu_and_mul contract: round both GEMM - // outputs to FP16 before the activation, round SiLU to FP16, then use - // FP16 multiplication. This makes fusion numerically equivalent to - // QPN2 GEMM followed by the current activation kernel. - const half gate_half = __float2half(gate); - const half up_half = __float2half(up); - const float gate_float = __half2float(gate_half); - const half silu = __float2half(gate_float / (1.0f + expf(-gate_float))); - output[static_cast(output_row) * hidden + blockIdx.x * 32 + - output_col] = __hmul(silu, up_half); - } - } -} - -template -void launch_qpn2(const uint8_t* codes, const uint8_t* scales, const half* input, - half* output, int n, int k, int m, float global_scale, - cudaStream_t stream) { - nvfp4_qpn2_sm70_kernel<<<(n / 32), (32 * SplitK), 0, stream>>>( - codes, scales, input, output, n, k, m, global_scale); -} - -template -void launch_qpn2_gated(const uint8_t* codes, const uint8_t* scales, - const half* input, half* output, int hidden, int k, - int m, float global_scale, cudaStream_t stream) { - nvfp4_qpn2_gated_sm70_kernel - <<<(hidden / 32), (64 * SplitK), 0, stream>>>( - codes, scales, input, output, hidden, k, m, global_scale); -} - -void check_qpn2_tensors(const torch::Tensor& out, const torch::Tensor& input, - const torch::Tensor& codes, const torch::Tensor& scales, - bool gated_silu) { - TORCH_CHECK( - out.is_cuda() && input.is_cuda() && codes.is_cuda() && scales.is_cuda(), - "NVFP4 QPN2 tensors must be CUDA tensors"); - TORCH_CHECK(out.scalar_type() == torch::kFloat16 && - input.scalar_type() == torch::kFloat16 && - codes.scalar_type() == torch::kUInt8 && - scales.scalar_type() == torch::kUInt8, - "NVFP4 QPN2 dtype mismatch"); - TORCH_CHECK(out.is_contiguous() && input.is_contiguous() && - codes.is_contiguous() && scales.is_contiguous(), - "NVFP4 QPN2 tensors must be contiguous"); - TORCH_CHECK(out.dim() == 2 && input.dim() == 2, - "NVFP4 QPN2 input and output must be matrices"); - TORCH_CHECK(out.get_device() == input.get_device() && - out.get_device() == codes.get_device() && - out.get_device() == scales.get_device(), - "NVFP4 QPN2 tensors must share one device"); - const int64_t m = input.size(0); - const int64_t k = input.size(1); - const int64_t n = gated_silu ? out.size(1) * 2 : out.size(1); - TORCH_CHECK(m >= 1 && m <= 8 && out.size(0) == m, - "NVFP4 QPN2 requires M in [1, 8]"); - TORCH_CHECK(k > 0 && k % 64 == 0 && n > 0 && n % 32 == 0, - "NVFP4 QPN2 shape alignment mismatch"); - TORCH_CHECK(codes.numel() == n * k / 2 && scales.numel() == n * k / 16, - "NVFP4 QPN2 packed tensor size mismatch"); -} - -} // namespace - -std::vector nvfp4_qpn2_prepare_sm70(torch::Tensor weight_packed, - torch::Tensor weight_scale) { - TORCH_CHECK(weight_packed.is_cuda() && weight_scale.is_cuda(), - "nvfp4_qpn2_prepare_sm70 expects CUDA tensors"); - TORCH_CHECK(weight_packed.scalar_type() == torch::kUInt8 && - weight_scale.scalar_type() == at::ScalarType::Float8_e4m3fn, - "nvfp4_qpn2_prepare_sm70 expects uint8 weights and E4M3 scales"); - TORCH_CHECK(weight_packed.dim() == 2 && weight_scale.dim() == 2 && - weight_packed.is_contiguous() && weight_scale.is_contiguous(), - "nvfp4_qpn2_prepare_sm70 expects contiguous matrices"); - TORCH_CHECK(weight_packed.get_device() == weight_scale.get_device(), - "nvfp4_qpn2_prepare_sm70 tensors must share one device"); - - const int64_t n = weight_packed.size(0); - const int64_t k = weight_packed.size(1) * 2; - TORCH_CHECK(n > 0 && n % 32 == 0 && k > 0 && k % 64 == 0, - "nvfp4_qpn2_prepare_sm70 shape alignment mismatch"); - TORCH_CHECK(weight_scale.size(0) == n && weight_scale.size(1) == k / 16, - "nvfp4_qpn2_prepare_sm70 scale shape mismatch"); - - const at::cuda::OptionalCUDAGuard device_guard(device_of(weight_packed)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - auto codes = torch::empty_like(weight_packed); - auto scales = torch::empty({n, k / 16}, torch::TensorOptions() - .device(weight_packed.device()) - .dtype(torch::kUInt8)); - - const int code_blocks = - static_cast((codes.numel() + kPrepareThreads - 1) / kPrepareThreads); - nvfp4_qpn2_prepack_codes_kernel<<>>( - codes.data_ptr(), weight_packed.data_ptr(), - static_cast(n), static_cast(k)); - const int scale_blocks = static_cast( - (scales.numel() + kPrepareThreads - 1) / kPrepareThreads); - nvfp4_qpn2_prepack_scales_kernel<<>>( - scales.data_ptr(), - reinterpret_cast(weight_scale.data_ptr()), - static_cast(n), static_cast(k)); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - return {codes, scales}; -} - -void nvfp4_qpn2_gemm_sm70_out(torch::Tensor out, torch::Tensor input, - torch::Tensor codes, torch::Tensor scales, - double global_scale, int64_t split_k, - int64_t accumulator_chains) { - check_qpn2_tensors(out, input, codes, scales, false); - TORCH_CHECK(split_k == 8 || split_k == 16 || split_k == 32, - "NVFP4 QPN2 split_k must be 8, 16, or 32"); - TORCH_CHECK((input.size(1) / 16) % split_k == 0, - "NVFP4 QPN2 K/16 must be divisible by split_k"); - TORCH_CHECK(accumulator_chains == 1 || accumulator_chains == 2, - "NVFP4 QPN2 accumulator_chains must be 1 or 2"); - - const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - const auto* code_ptr = codes.data_ptr(); - const auto* scale_ptr = scales.data_ptr(); - const auto* input_ptr = - reinterpret_cast(input.data_ptr()); - auto* output_ptr = reinterpret_cast(out.data_ptr()); - const int n = static_cast(out.size(1)); - const int k = static_cast(input.size(1)); - const int m = static_cast(input.size(0)); - -#define VLLM_LAUNCH_QPN2(SPLIT, NACC) \ - launch_qpn2(code_ptr, scale_ptr, input_ptr, output_ptr, n, k, \ - m, static_cast(global_scale), stream) - if (split_k == 8 && accumulator_chains == 1) { - VLLM_LAUNCH_QPN2(8, 1); - } else if (split_k == 8) { - VLLM_LAUNCH_QPN2(8, 2); - } else if (split_k == 16 && accumulator_chains == 1) { - VLLM_LAUNCH_QPN2(16, 1); - } else if (split_k == 16) { - VLLM_LAUNCH_QPN2(16, 2); - } else if (accumulator_chains == 1) { - VLLM_LAUNCH_QPN2(32, 1); - } else { - VLLM_LAUNCH_QPN2(32, 2); - } -#undef VLLM_LAUNCH_QPN2 - C10_CUDA_KERNEL_LAUNCH_CHECK(); -} - -void nvfp4_qpn2_gated_sm70_out(torch::Tensor out, torch::Tensor input, - torch::Tensor codes, torch::Tensor scales, - double global_scale, int64_t split_k, - int64_t accumulator_chains) { - check_qpn2_tensors(out, input, codes, scales, true); - TORCH_CHECK(split_k == 8 || split_k == 16, - "NVFP4 QPN2 gated split_k must be 8 or 16"); - TORCH_CHECK((input.size(1) / 16) % split_k == 0, - "NVFP4 QPN2 gated K/16 must be divisible by split_k"); - TORCH_CHECK(accumulator_chains == 1 || accumulator_chains == 2, - "NVFP4 QPN2 gated accumulator_chains must be 1 or 2"); - - const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - const auto* code_ptr = codes.data_ptr(); - const auto* scale_ptr = scales.data_ptr(); - const auto* input_ptr = - reinterpret_cast(input.data_ptr()); - auto* output_ptr = reinterpret_cast(out.data_ptr()); - const int hidden = static_cast(out.size(1)); - const int k = static_cast(input.size(1)); - const int m = static_cast(input.size(0)); - -#define VLLM_LAUNCH_QPN2_GATED(SPLIT, NACC) \ - launch_qpn2_gated(code_ptr, scale_ptr, input_ptr, output_ptr, \ - hidden, k, m, \ - static_cast(global_scale), stream) - if (split_k == 8 && accumulator_chains == 1) { - VLLM_LAUNCH_QPN2_GATED(8, 1); - } else if (split_k == 8) { - VLLM_LAUNCH_QPN2_GATED(8, 2); - } else if (accumulator_chains == 1) { - VLLM_LAUNCH_QPN2_GATED(16, 1); - } else { - VLLM_LAUNCH_QPN2_GATED(16, 2); - } -#undef VLLM_LAUNCH_QPN2_GATED - C10_CUDA_KERNEL_LAUNCH_CHECK(); -} - -#ifndef VLLM_NVFP4_QPN2_STANDALONE -void nvfp4_qpn2_dispatch_sm70_out(torch::Tensor out, torch::Tensor input, - torch::Tensor codes, torch::Tensor scales, - double global_scale, int64_t split_k, - int64_t accumulator_chains, - torch::Tensor tm_weight, - torch::Tensor tm_scales, - int64_t tm_group_size, int64_t tm_k_ld, - int64_t tm_q_ld, bool gated_silu) { - if (input.size(0) <= 8) { - if (gated_silu) { - nvfp4_qpn2_gated_sm70_out(out, input, codes, scales, global_scale, - split_k, accumulator_chains); - } else { - nvfp4_qpn2_gemm_sm70_out(out, input, codes, scales, global_scale, split_k, - accumulator_chains); - } - return; - } - - if (!gated_silu) { - nvfp4_gemm_sm70_out(out, input, tm_weight, tm_scales, tm_group_size, - tm_k_ld, tm_q_ld, false); - return; - } - auto gate_up = - torch::empty({input.size(0), out.size(1) * 2}, input.options()); - nvfp4_gemm_sm70_out(gate_up, input, tm_weight, tm_scales, tm_group_size, - tm_k_ld, tm_q_ld, false); - silu_and_mul(out, gate_up); -} -#endif - -#ifdef VLLM_NVFP4_QPN2_STANDALONE -// Compile the exact production kernels as a task-local operator-race library -// before paying for a complete vLLM rebuild. Production registers these -// operators centrally in torch_bindings.cpp. -TORCH_LIBRARY_FRAGMENT(_C, ops) { - ops.def( - "nvfp4_qpn2_prepare_sm70(Tensor weight_packed, Tensor weight_scale) -> " - "Tensor[]"); - ops.impl("nvfp4_qpn2_prepare_sm70", torch::kCUDA, &nvfp4_qpn2_prepare_sm70); - ops.def( - "nvfp4_qpn2_gemm_sm70_out(Tensor(a!) out, Tensor input, Tensor codes, " - "Tensor scales, float global_scale, int split_k, " - "int accumulator_chains) -> ()"); - ops.impl("nvfp4_qpn2_gemm_sm70_out", torch::kCUDA, &nvfp4_qpn2_gemm_sm70_out); - ops.def( - "nvfp4_qpn2_gated_sm70_out(Tensor(a!) out, Tensor input, Tensor codes, " - "Tensor scales, float global_scale, int split_k, " - "int accumulator_chains) -> ()"); - ops.impl("nvfp4_qpn2_gated_sm70_out", torch::kCUDA, - &nvfp4_qpn2_gated_sm70_out); -} -#endif diff --git a/csrc/sm70_turbomind/ops/nvfp4_qpn4_sm70.cu b/csrc/sm70_turbomind/ops/nvfp4_qpn4_sm70.cu deleted file mode 100644 index 0395289403..0000000000 --- a/csrc/sm70_turbomind/ops/nvfp4_qpn4_sm70.cu +++ /dev/null @@ -1,1029 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// Qwen3.8 NVFP4 M=1 decode kernels for NVIDIA Volta SM70. - -#include -#include - -#include -#include -#include -#include - -namespace { - -constexpr int kPrepareThreads = 256; -constexpr float kFp4Bias = 16384.0f; - -struct SplitHalfScale { - half hi; - half lo; -}; - -SplitHalfScale split_half_scale(float value) { - const half hi = __float2half_rn(value); - const half lo = __float2half_rn(value - __half2float(hi)); - return {hi, lo}; -} - -__device__ __forceinline__ int qpn_col_from_lane(int lane) { - return ((lane >> 2) & 3) * 8 + (lane & 3) + ((lane & 16) ? 4 : 0); -} - -__device__ __forceinline__ int logical_k_from_physical(int physical) { - const int local = physical & 7; - return (physical & 8) + ((local & 3) << 1) + (local >> 2); -} - -__global__ void nvfp4_qpn4_prepack_codes_kernel( - uint8_t* __restrict__ codes, const uint8_t* __restrict__ qweight, int k, - int n) { - const size_t index = - static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - const size_t packed_numel = static_cast(k) * n / 2; - if (index >= packed_numel) { - return; - } - - const int physical_byte = static_cast(index & 7); - size_t outer = index >> 3; - const int lane = static_cast(outer & 31); - outer >>= 5; - const int groups_k16 = k >> 4; - const int group = static_cast(outer % groups_k16); - const int tile = static_cast(outer / groups_k16); - const int col = tile * 32 + qpn_col_from_lane(lane); - const int physical_k0 = physical_byte * 2; - const int logical_k0 = logical_k_from_physical(physical_k0); - const int logical_k1 = logical_k_from_physical(physical_k0 + 1); - const uint8_t lo = - qweight[static_cast(group * 16 + logical_k0) * n + col] & 0x0fU; - const uint8_t hi = - qweight[static_cast(group * 16 + logical_k1) * n + col] & 0x0fU; - codes[index] = lo | (hi << 4); -} - -__global__ void nvfp4_qpn4_prepack_scales_kernel( - half* __restrict__ packed_scales, const half* __restrict__ scales, int k, - int n) { - const size_t index = - static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - const size_t numel = static_cast(k / 16) * n; - if (index >= numel) { - return; - } - const int lane = static_cast(index & 31); - size_t outer = index >> 5; - const int groups_k16 = k >> 4; - const int group = static_cast(outer % groups_k16); - const int tile = static_cast(outer / groups_k16); - const int col = tile * 32 + qpn_col_from_lane(lane); - packed_scales[index] = __hmul(scales[static_cast(group) * n + col], - __float2half_rn(kFp4Bias)); -} - -__global__ void nvfp4_qpn4_prepack_scale_codes_kernel( - uint8_t* __restrict__ packed_scale_codes, - const uint8_t* __restrict__ scale_codes, int k, int n) { - const size_t index = - static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - const size_t numel = static_cast(k / 16) * n; - if (index >= numel) { - return; - } - const int lane = static_cast(index & 31); - size_t outer = index >> 5; - const int groups_k16 = k >> 4; - const int group = static_cast(outer % groups_k16); - const int tile = static_cast(outer / groups_k16); - const int col = tile * 32 + qpn_col_from_lane(lane); - packed_scale_codes[index] = scale_codes[static_cast(group) * n + col]; -} - -__device__ __forceinline__ void fp4x8_to_half2x4(unsigned x, half2 out[4]) { - constexpr unsigned kSign = 0x80008000U; - constexpr unsigned kExponentMantissa = 0x0E000E00U; - unsigned values[4]; - values[0] = ((x << 12) & kSign) | ((x << 9) & kExponentMantissa); - values[1] = ((x << 8) & kSign) | ((x << 5) & kExponentMantissa); - values[2] = ((x << 4) & kSign) | ((x << 1) & kExponentMantissa); - values[3] = ((x << 0) & kSign) | ((x >> 3) & kExponentMantissa); -#pragma unroll - for (int i = 0; i < 4; ++i) { - out[i] = *reinterpret_cast(&values[i]); - } -} - -__device__ __forceinline__ void fp4x16_to_half2x8(uint2 packed, half2 out[8]) { - fp4x8_to_half2x4(packed.x, out); - fp4x8_to_half2x4(packed.y, out + 4); -} - -__device__ __forceinline__ half nvfp4_scale_code_to_half(uint8_t scale_code, - half scale_hi, - half scale_lo) { - // NVFP4 group scales are positive, normal E4M3FN values. Adding the bias - // delta while shifting the E4M3 payload produces the exact FP16 value. - const unsigned half_bits = - (static_cast(scale_code) << 7U) + 0x2000U; - const half raw_scale = *reinterpret_cast(&half_bits); - const half correction = __hmul(raw_scale, scale_lo); - return __hfma(raw_scale, scale_hi, correction); -} - -template -__global__ void nvfp4_qpn4_dequantize_sm70_kernel( - half* __restrict__ output, const uint8_t* __restrict__ codes, - const void* __restrict__ packed_scales, half scale_hi, half scale_lo, int n, - int k) { - const size_t word_index = - static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - const size_t word_count = static_cast(k) * n / 16; - if (word_index >= word_count) { - return; - } - - const int lane = static_cast(word_index & 31); - size_t outer = word_index >> 5; - const int groups_k16 = k >> 4; - const int group = static_cast(outer % groups_k16); - const int tile = static_cast(outer / groups_k16); - const int col = tile * 32 + qpn_col_from_lane(lane); - half scale; - if constexpr (UseScaleCode) { - scale = nvfp4_scale_code_to_half( - reinterpret_cast(packed_scales)[word_index], scale_hi, - scale_lo); - } else { - scale = reinterpret_cast(packed_scales)[word_index]; - } - half2 weights[8]; - fp4x16_to_half2x8(reinterpret_cast(codes)[word_index], weights); - const half2 scale2 = __halves2half2(scale, scale); -#pragma unroll - for (int pair = 0; pair < 8; ++pair) { - const half2 value = __hmul2(weights[pair], scale2); - // fp4x8_to_half2x4 pairs nibbles (0,4), (1,5), (2,6), (3,7) - // within each 8-value word rather than adjacent physical K values. - const int physical_k0 = (pair & 3) + ((pair & 4) ? 8 : 0); - const int logical_k0 = logical_k_from_physical(physical_k0); - const int logical_k1 = logical_k_from_physical(physical_k0 + 4); - output[static_cast(group * 16 + logical_k0) * n + col] = - __low2half(value); - output[static_cast(group * 16 + logical_k1) * n + col] = - __high2half(value); - } -} - -__global__ void nvfp4_qpn4_silu_and_mul_sm70_kernel( - half* __restrict__ output, const half* __restrict__ gate_up, int rows, - int hidden) { - const int row = blockIdx.x; - if (row >= rows) { - return; - } - const half* row_input = gate_up + static_cast(row) * hidden * 2; - half* row_output = output + static_cast(row) * hidden; - for (int col = threadIdx.x; col < hidden; col += blockDim.x) { - const float gate = __half2float(row_input[col]); - const float silu = gate / (1.0f + __expf(-gate)); - row_output[col] = __hmul(__float2half(silu), row_input[hidden + col]); - } -} - -#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 nvfp4_qpn4_sm70_kernel(const uint8_t* __restrict__ codes, - const void* __restrict__ packed_scales, - half scale_hi, half scale_lo, - const half* __restrict__ input, - half* __restrict__ output, int n, int k, - int m) { - __shared__ float partials[SplitK][32]; - - const int lane = threadIdx.x & 31; - const int warp = threadIdx.x >> 5; - const int tile = blockIdx.x; - 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_warp = groups_k16 / SplitK; - const int group_begin = warp * groups_per_warp; - const uint2* code_ptr = reinterpret_cast(codes) + - static_cast(tile) * groups_k16 * 32 + lane; - const size_t scale_offset = - static_cast(tile) * groups_k16 * 32 + lane; - const half* scale_ptr = - reinterpret_cast(packed_scales) + scale_offset; - const uint8_t* scale_code_ptr = - reinterpret_cast(packed_scales) + scale_offset; - - float accum[NAcc][8]; -#pragma unroll - for (int chain = 0; chain < NAcc; ++chain) { -#pragma unroll - for (int i = 0; i < 8; ++i) { - accum[chain][i] = 0.0f; - } - } - - uint2 prefetched = make_uint2(0, 0); - if constexpr (PrefetchCodes) { - prefetched = __ldcs(code_ptr + static_cast(group_begin) * 32); - } - -#pragma unroll 4 - for (int group = group_begin; group < group_begin + groups_per_warp; - ++group) { - const uint2 packed = - PrefetchCodes ? prefetched - : __ldcs(code_ptr + static_cast(group) * 32); - uint2 next = make_uint2(0, 0); - if constexpr (PrefetchCodes) { - if (group + 1 < group_begin + groups_per_warp) { - next = __ldcs(code_ptr + static_cast(group + 1) * 32); - } - } - half2 weights[8]; - fp4x16_to_half2x8(packed, weights); - half scale; - if constexpr (UseScaleCode) { - const uint8_t scale_code = - __ldg(scale_code_ptr + static_cast(group) * 32); - scale = nvfp4_scale_code_to_half(scale_code, scale_hi, scale_lo); - } else { - scale = __ldg(scale_ptr + static_cast(group) * 32); - } - const half2 scale2 = __halves2half2(scale, scale); -#pragma unroll - for (int i = 0; i < 8; ++i) { - weights[i] = __hmul2(weights[i], scale2); - } - - uint4 input01 = make_uint4(0, 0, 0, 0); - uint4 input23 = make_uint4(0, 0, 0, 0); - if (row < m) { - const half* input_row = input + static_cast(row) * k; - 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); - 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]); - if constexpr (PrefetchCodes) { - prefetched = next; - } - } - -#pragma unroll - for (int chain = 1; chain < NAcc; ++chain) { -#pragma unroll - for (int i = 0; i < 8; ++i) { - accum[0][i] += accum[chain][i]; - } - } - if ((lane & 17) == 0) { -#pragma unroll - for (int pair = 0; pair < 2; ++pair) { -#pragma unroll - for (int offset = 0; offset < 2; ++offset) { - const int i = pair * 4 + offset; - const int output_col = offset | (((lane >> 1) & 1) << 1) | (pair << 2); - partials[warp][quadpair * 8 + output_col] = accum[0][i]; - } - } - } - __syncthreads(); - - for (int element = threadIdx.x; element < 32; element += blockDim.x) { - float value = 0.0f; -#pragma unroll - for (int k_warp = 0; k_warp < SplitK; ++k_warp) { - value += partials[k_warp][element]; - } - output[tile * 32 + element] = __float2half(value); - } -} - -template -__global__ void nvfp4_qpn4_gated_sm70_kernel( - const uint8_t* __restrict__ codes, const void* __restrict__ packed_scales, - half scale_hi, half scale_lo, const half* __restrict__ input, - half* __restrict__ output, int hidden, int k, int m) { - __shared__ float partials[2][SplitK][32]; - - const int lane = threadIdx.x & 31; - const int warp_in_block = threadIdx.x >> 5; - const int projection = warp_in_block / SplitK; - const int warp = warp_in_block - projection * SplitK; - const int hidden_tiles = hidden >> 5; - const int tile = blockIdx.x + projection * hidden_tiles; - 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_warp = groups_k16 / SplitK; - const int group_begin = warp * groups_per_warp; - const uint2* code_ptr = reinterpret_cast(codes) + - static_cast(tile) * groups_k16 * 32 + lane; - const size_t scale_offset = - static_cast(tile) * groups_k16 * 32 + lane; - const half* scale_ptr = - reinterpret_cast(packed_scales) + scale_offset; - const uint8_t* scale_code_ptr = - reinterpret_cast(packed_scales) + scale_offset; - - float accum[NAcc][8]; -#pragma unroll - for (int chain = 0; chain < NAcc; ++chain) { -#pragma unroll - for (int i = 0; i < 8; ++i) { - accum[chain][i] = 0.0f; - } - } - uint2 prefetched = make_uint2(0, 0); - if constexpr (PrefetchCodes) { - prefetched = __ldcs(code_ptr + static_cast(group_begin) * 32); - } - -#pragma unroll 4 - for (int group = group_begin; group < group_begin + groups_per_warp; - ++group) { - const uint2 packed = - PrefetchCodes ? prefetched - : __ldcs(code_ptr + static_cast(group) * 32); - uint2 next = make_uint2(0, 0); - if constexpr (PrefetchCodes) { - if (group + 1 < group_begin + groups_per_warp) { - next = __ldcs(code_ptr + static_cast(group + 1) * 32); - } - } - half2 weights[8]; - fp4x16_to_half2x8(packed, weights); - half scale; - if constexpr (UseScaleCode) { - const uint8_t scale_code = - __ldg(scale_code_ptr + static_cast(group) * 32); - scale = nvfp4_scale_code_to_half(scale_code, scale_hi, scale_lo); - } else { - scale = __ldg(scale_ptr + static_cast(group) * 32); - } - const half2 scale2 = __halves2half2(scale, scale); -#pragma unroll - for (int i = 0; i < 8; ++i) { - weights[i] = __hmul2(weights[i], scale2); - } - - uint4 input01 = make_uint4(0, 0, 0, 0); - uint4 input23 = make_uint4(0, 0, 0, 0); - if (row < m) { - const half* input_row = input + static_cast(row) * k; - 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); - 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]); - if constexpr (PrefetchCodes) { - prefetched = next; - } - } - -#pragma unroll - for (int chain = 1; chain < NAcc; ++chain) { -#pragma unroll - for (int i = 0; i < 8; ++i) { - accum[0][i] += accum[chain][i]; - } - } - if ((lane & 17) == 0) { -#pragma unroll - for (int pair = 0; pair < 2; ++pair) { -#pragma unroll - for (int offset = 0; offset < 2; ++offset) { - const int i = pair * 4 + offset; - const int output_col = offset | (((lane >> 1) & 1) << 1) | (pair << 2); - partials[projection][warp][quadpair * 8 + output_col] = accum[0][i]; - } - } - } - __syncthreads(); - - for (int element = threadIdx.x; element < 32; element += blockDim.x) { - float gate = 0.0f; - float up = 0.0f; -#pragma unroll - for (int k_warp = 0; k_warp < SplitK; ++k_warp) { - gate += partials[0][k_warp][element]; - up += partials[1][k_warp][element]; - } - const float silu = gate / (1.0f + __expf(-gate)); - output[blockIdx.x * 32 + element] = __float2half(silu * up); - } -} - -template -void launch_qpn4(const uint8_t* codes, const void* scales, half scale_hi, - half scale_lo, const half* input, half* output, int n, int k, - int m, cudaStream_t stream) { - nvfp4_qpn4_sm70_kernel - <<<(n / 32), (32 * SplitK), 0, stream>>>( - codes, scales, scale_hi, scale_lo, input, output, n, k, m); -} - -template -void launch_qpn4_gated(const uint8_t* codes, const void* scales, half scale_hi, - half scale_lo, const half* input, half* output, - int hidden, int k, int m, cudaStream_t stream) { - nvfp4_qpn4_gated_sm70_kernel - <<<(hidden / 32), (64 * SplitK), 0, stream>>>( - codes, scales, scale_hi, scale_lo, input, output, hidden, k, m); -} - -} // namespace - -std::vector nvfp4_qpn4_prepare_sm70(torch::Tensor qweight, - torch::Tensor scales) { - TORCH_CHECK(qweight.is_cuda() && scales.is_cuda(), - "nvfp4_qpn4_prepare_sm70: tensors must be CUDA"); - TORCH_CHECK(qweight.scalar_type() == torch::kUInt8 && - scales.scalar_type() == torch::kFloat16, - "nvfp4_qpn4_prepare_sm70: expected uint8 codes and FP16 scales"); - TORCH_CHECK(qweight.dim() == 2 && scales.dim() == 2 && - qweight.is_contiguous() && scales.is_contiguous(), - "nvfp4_qpn4_prepare_sm70: tensors must be contiguous matrices"); - TORCH_CHECK(qweight.get_device() == scales.get_device(), - "nvfp4_qpn4_prepare_sm70: tensors must share one device"); - const int64_t k = qweight.size(0); - const int64_t n = qweight.size(1); - TORCH_CHECK(k > 0 && k % 128 == 0 && n > 0 && n % 32 == 0, - "nvfp4_qpn4_prepare_sm70: shape alignment mismatch"); - TORCH_CHECK(scales.size(0) == k / 16 && scales.size(1) == n, - "nvfp4_qpn4_prepare_sm70: scale shape mismatch"); - - const at::cuda::OptionalCUDAGuard device_guard(device_of(qweight)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - auto packed_codes = torch::empty( - {k, n / 2}, - torch::TensorOptions().device(qweight.device()).dtype(torch::kUInt8)); - auto packed_scales = torch::empty_like(scales); - const int64_t code_numel = packed_codes.numel(); - const int code_blocks = - static_cast((code_numel + kPrepareThreads - 1) / kPrepareThreads); - nvfp4_qpn4_prepack_codes_kernel<<>>( - packed_codes.data_ptr(), qweight.data_ptr(), - static_cast(k), static_cast(n)); - const int64_t scale_numel = packed_scales.numel(); - const int scale_blocks = - static_cast((scale_numel + kPrepareThreads - 1) / kPrepareThreads); - nvfp4_qpn4_prepack_scales_kernel<<>>( - reinterpret_cast(packed_scales.data_ptr()), - reinterpret_cast(scales.data_ptr()), - static_cast(k), static_cast(n)); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - return {packed_codes, packed_scales}; -} - -std::vector nvfp4_qpn4_prepare_scale_code_sm70( - torch::Tensor qweight, torch::Tensor scale_codes) { - TORCH_CHECK(qweight.is_cuda() && scale_codes.is_cuda(), - "nvfp4_qpn4_prepare_scale_code_sm70: tensors must be CUDA"); - TORCH_CHECK(qweight.scalar_type() == torch::kUInt8 && - scale_codes.scalar_type() == at::ScalarType::Float8_e4m3fn, - "nvfp4_qpn4_prepare_scale_code_sm70: expected uint8 weights and " - "float8_e4m3fn scale codes"); - TORCH_CHECK(qweight.dim() == 2 && scale_codes.dim() == 2 && - qweight.is_contiguous() && scale_codes.is_contiguous(), - "nvfp4_qpn4_prepare_scale_code_sm70: tensors must be contiguous " - "matrices"); - TORCH_CHECK(qweight.get_device() == scale_codes.get_device(), - "nvfp4_qpn4_prepare_scale_code_sm70: tensors must share one " - "device"); - const int64_t k = qweight.size(0); - const int64_t n = qweight.size(1); - TORCH_CHECK(k > 0 && k % 128 == 0 && n > 0 && n % 32 == 0, - "nvfp4_qpn4_prepare_scale_code_sm70: shape alignment mismatch"); - TORCH_CHECK(scale_codes.size(0) == k / 16 && scale_codes.size(1) == n, - "nvfp4_qpn4_prepare_scale_code_sm70: scale shape mismatch"); - - const at::cuda::OptionalCUDAGuard device_guard(device_of(qweight)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - auto packed_codes = torch::empty( - {k, n / 2}, - torch::TensorOptions().device(qweight.device()).dtype(torch::kUInt8)); - auto packed_scale_codes = torch::empty( - {k / 16, n}, - torch::TensorOptions().device(qweight.device()).dtype(torch::kUInt8)); - const int64_t code_numel = packed_codes.numel(); - const int code_blocks = - static_cast((code_numel + kPrepareThreads - 1) / kPrepareThreads); - nvfp4_qpn4_prepack_codes_kernel<<>>( - packed_codes.data_ptr(), qweight.data_ptr(), - static_cast(k), static_cast(n)); - const int64_t scale_numel = packed_scale_codes.numel(); - const int scale_blocks = - static_cast((scale_numel + kPrepareThreads - 1) / kPrepareThreads); - nvfp4_qpn4_prepack_scale_codes_kernel<<>>( - packed_scale_codes.data_ptr(), - reinterpret_cast(scale_codes.data_ptr()), - static_cast(k), static_cast(n)); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - return {packed_codes, packed_scale_codes}; -} - -void nvfp4_qpn4_dequantize_sm70_out(torch::Tensor out, torch::Tensor codes, - torch::Tensor scales, double global_scale, - bool use_scale_code) { - TORCH_CHECK(out.is_cuda() && codes.is_cuda() && scales.is_cuda(), - "nvfp4_qpn4_dequantize_sm70_out: tensors must be CUDA"); - TORCH_CHECK(out.scalar_type() == torch::kFloat16 && - codes.scalar_type() == torch::kUInt8, - "nvfp4_qpn4_dequantize_sm70_out: output must be FP16 and codes " - "must be uint8"); - TORCH_CHECK((use_scale_code && scales.scalar_type() == torch::kUInt8) || - (!use_scale_code && scales.scalar_type() == torch::kFloat16), - "nvfp4_qpn4_dequantize_sm70_out: scale dtype mismatch"); - TORCH_CHECK(out.dim() == 2 && codes.dim() == 2 && scales.dim() == 2 && - out.is_contiguous() && codes.is_contiguous() && - scales.is_contiguous(), - "nvfp4_qpn4_dequantize_sm70_out: tensors must be contiguous " - "matrices"); - TORCH_CHECK(out.get_device() == codes.get_device() && - out.get_device() == scales.get_device(), - "nvfp4_qpn4_dequantize_sm70_out: tensors must share one device"); - const int64_t k = out.size(0); - const int64_t n = out.size(1); - TORCH_CHECK(k > 0 && k % 128 == 0 && n > 0 && n % 32 == 0, - "nvfp4_qpn4_dequantize_sm70_out: shape alignment mismatch"); - TORCH_CHECK(codes.numel() == k * n / 2 && scales.numel() == k * n / 16, - "nvfp4_qpn4_dequantize_sm70_out: packed tensor size mismatch"); - - const at::cuda::OptionalCUDAGuard device_guard(device_of(out)); - const int64_t word_count = k * n / 16; - const int blocks = - static_cast((word_count + kPrepareThreads - 1) / kPrepareThreads); - const SplitHalfScale split_scale = - split_half_scale(static_cast(global_scale) * kFp4Bias); - const half zero_scale = __float2half_rn(0.0f); - if (use_scale_code) { - nvfp4_qpn4_dequantize_sm70_kernel - <<>>( - reinterpret_cast(out.data_ptr()), - codes.data_ptr(), scales.data_ptr(), - split_scale.hi, split_scale.lo, static_cast(n), - static_cast(k)); - } else { - nvfp4_qpn4_dequantize_sm70_kernel - <<>>( - reinterpret_cast(out.data_ptr()), - codes.data_ptr(), scales.data_ptr(), zero_scale, - zero_scale, static_cast(n), static_cast(k)); - } - C10_CUDA_KERNEL_LAUNCH_CHECK(); -} - -void nvfp4_qpn4_prefill_sm70_out(torch::Tensor out, int64_t dense_weight_ptr, - torch::Tensor input, torch::Tensor codes, - torch::Tensor scales, double global_scale, - bool use_scale_code, bool gated_silu) { - TORCH_CHECK(input.is_cuda() && out.is_cuda(), - "nvfp4_qpn4_prefill_sm70_out: input and output must be CUDA"); - TORCH_CHECK(input.scalar_type() == torch::kFloat16 && - out.scalar_type() == torch::kFloat16, - "nvfp4_qpn4_prefill_sm70_out: input and output must be FP16"); - TORCH_CHECK(input.dim() == 2 && out.dim() == 2 && dense_weight_ptr != 0 && - input.is_contiguous() && out.is_contiguous(), - "nvfp4_qpn4_prefill_sm70_out: invalid input or workspace"); - const int64_t m = input.size(0); - const int64_t k = input.size(1); - TORCH_CHECK(codes.dim() == 2 && codes.size(0) == k, - "nvfp4_qpn4_prefill_sm70_out: code shape mismatch"); - const int64_t n = codes.size(1) * 2; - TORCH_CHECK(out.size(0) == m && out.size(1) == (gated_silu ? n / 2 : n), - "nvfp4_qpn4_prefill_sm70_out: output shape mismatch"); - - auto dense_weight = torch::from_blob( - reinterpret_cast(dense_weight_ptr), {k, n}, input.options()); - nvfp4_qpn4_dequantize_sm70_out(dense_weight, codes, scales, global_scale, - use_scale_code); - if (!gated_silu) { - at::mm_out(out, input, dense_weight); - return; - } - - auto gate_up = at::mm(input, dense_weight); - constexpr int kThreads = 256; - nvfp4_qpn4_silu_and_mul_sm70_kernel<<(m), kThreads, 0, - at::cuda::getCurrentCUDAStream()>>>( - reinterpret_cast(out.data_ptr()), - reinterpret_cast(gate_up.data_ptr()), - static_cast(m), static_cast(n / 2)); - C10_CUDA_KERNEL_LAUNCH_CHECK(); -} - -void nvfp4_qpn4_gemm_sm70_out(torch::Tensor out, torch::Tensor input, - torch::Tensor codes, torch::Tensor scales, - int64_t split_k, int64_t accumulator_chains, - bool prefetch_codes) { - TORCH_CHECK( - out.is_cuda() && input.is_cuda() && codes.is_cuda() && scales.is_cuda(), - "nvfp4_qpn4_gemm_sm70_out: tensors must be CUDA"); - TORCH_CHECK(out.scalar_type() == torch::kFloat16 && - input.scalar_type() == torch::kFloat16 && - codes.scalar_type() == torch::kUInt8 && - scales.scalar_type() == torch::kFloat16, - "nvfp4_qpn4_gemm_sm70_out: dtype mismatch"); - TORCH_CHECK(out.is_contiguous() && input.is_contiguous() && - codes.is_contiguous() && scales.is_contiguous(), - "nvfp4_qpn4_gemm_sm70_out: tensors must be contiguous"); - const int64_t m = input.size(0); - const int64_t k = input.size(1); - const int64_t n = out.size(1); - TORCH_CHECK(m == 1 && out.size(0) == 1, - "nvfp4_qpn4_gemm_sm70_out: M must be 1"); - TORCH_CHECK(k > 0 && k % 128 == 0 && n > 0 && n % 32 == 0, - "nvfp4_qpn4_gemm_sm70_out: shape alignment mismatch"); - TORCH_CHECK(codes.numel() == k * n / 2 && scales.numel() == k * n / 16, - "nvfp4_qpn4_gemm_sm70_out: packed tensor size mismatch"); - TORCH_CHECK(split_k == 4 || split_k == 8 || split_k == 10 || split_k == 16 || - split_k == 17, - "nvfp4_qpn4_gemm_sm70_out: unsupported split_k"); - TORCH_CHECK((k / 16) % split_k == 0, - "nvfp4_qpn4_gemm_sm70_out: invalid split_k for K"); - TORCH_CHECK(accumulator_chains == 1 || accumulator_chains == 2, - "nvfp4_qpn4_gemm_sm70_out: accumulator chains must be 1 or 2"); - - const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - const auto* code_ptr = codes.data_ptr(); - const auto* scale_ptr = - reinterpret_cast(scales.data_ptr()); - const auto* input_ptr = - reinterpret_cast(input.data_ptr()); - auto* output_ptr = reinterpret_cast(out.data_ptr()); - const half zero_scale = __float2half_rn(0.0f); - -#define LAUNCH_QPN4(SPLIT, NACC, PREFETCH) \ - launch_qpn4( \ - code_ptr, scale_ptr, zero_scale, zero_scale, input_ptr, output_ptr, \ - static_cast(n), static_cast(k), static_cast(m), stream) -#define DISPATCH_NACC_PREFETCH(SPLIT) \ - do { \ - if (accumulator_chains == 1 && !prefetch_codes) { \ - LAUNCH_QPN4(SPLIT, 1, false); \ - } else if (accumulator_chains == 2 && !prefetch_codes) { \ - LAUNCH_QPN4(SPLIT, 2, false); \ - } else if (accumulator_chains == 1) { \ - LAUNCH_QPN4(SPLIT, 1, true); \ - } else { \ - LAUNCH_QPN4(SPLIT, 2, true); \ - } \ - } while (0) - if (split_k == 4) { - DISPATCH_NACC_PREFETCH(4); - } else if (split_k == 8) { - DISPATCH_NACC_PREFETCH(8); - } else if (split_k == 10) { - DISPATCH_NACC_PREFETCH(10); - } else if (split_k == 16) { - DISPATCH_NACC_PREFETCH(16); - } else { - DISPATCH_NACC_PREFETCH(17); - } -#undef DISPATCH_NACC_PREFETCH -#undef LAUNCH_QPN4 - C10_CUDA_KERNEL_LAUNCH_CHECK(); -} - -void nvfp4_qpn4_gated_sm70_out(torch::Tensor out, torch::Tensor input, - torch::Tensor codes, torch::Tensor scales, - int64_t split_k, int64_t accumulator_chains, - bool prefetch_codes) { - TORCH_CHECK( - out.is_cuda() && input.is_cuda() && codes.is_cuda() && scales.is_cuda(), - "nvfp4_qpn4_gated_sm70_out: tensors must be CUDA"); - TORCH_CHECK(out.scalar_type() == torch::kFloat16 && - input.scalar_type() == torch::kFloat16 && - codes.scalar_type() == torch::kUInt8 && - scales.scalar_type() == torch::kFloat16, - "nvfp4_qpn4_gated_sm70_out: dtype mismatch"); - TORCH_CHECK(out.is_contiguous() && input.is_contiguous() && - codes.is_contiguous() && scales.is_contiguous(), - "nvfp4_qpn4_gated_sm70_out: tensors must be contiguous"); - const int64_t m = input.size(0); - const int64_t k = input.size(1); - const int64_t hidden = out.size(1); - const int64_t n = hidden * 2; - TORCH_CHECK(m == 1 && out.size(0) == 1, - "nvfp4_qpn4_gated_sm70_out: M must be 1"); - TORCH_CHECK(k > 0 && k % 128 == 0 && hidden > 0 && hidden % 32 == 0, - "nvfp4_qpn4_gated_sm70_out: shape alignment mismatch"); - TORCH_CHECK(codes.numel() == k * n / 2 && scales.numel() == k * n / 16, - "nvfp4_qpn4_gated_sm70_out: packed tensor size mismatch"); - TORCH_CHECK(split_k == 4 || split_k == 8 || split_k == 10 || split_k == 16, - "nvfp4_qpn4_gated_sm70_out: unsupported split_k"); - TORCH_CHECK((k / 16) % split_k == 0, - "nvfp4_qpn4_gated_sm70_out: invalid split_k for K"); - TORCH_CHECK(accumulator_chains == 1 || accumulator_chains == 2, - "nvfp4_qpn4_gated_sm70_out: accumulator chains must be 1 or 2"); - - const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - const auto* code_ptr = codes.data_ptr(); - const auto* scale_ptr = - reinterpret_cast(scales.data_ptr()); - const auto* input_ptr = - reinterpret_cast(input.data_ptr()); - auto* output_ptr = reinterpret_cast(out.data_ptr()); - const half zero_scale = __float2half_rn(0.0f); - -#define LAUNCH_QPN4_GATED(SPLIT, NACC, PREFETCH) \ - launch_qpn4_gated( \ - code_ptr, scale_ptr, zero_scale, zero_scale, input_ptr, output_ptr, \ - static_cast(hidden), static_cast(k), static_cast(m), \ - stream) -#define DISPATCH_GATED_NACC_PREFETCH(SPLIT) \ - do { \ - if (accumulator_chains == 1 && !prefetch_codes) { \ - LAUNCH_QPN4_GATED(SPLIT, 1, false); \ - } else if (accumulator_chains == 2 && !prefetch_codes) { \ - LAUNCH_QPN4_GATED(SPLIT, 2, false); \ - } else if (accumulator_chains == 1) { \ - LAUNCH_QPN4_GATED(SPLIT, 1, true); \ - } else { \ - LAUNCH_QPN4_GATED(SPLIT, 2, true); \ - } \ - } while (0) - if (split_k == 4) { - DISPATCH_GATED_NACC_PREFETCH(4); - } else if (split_k == 8) { - DISPATCH_GATED_NACC_PREFETCH(8); - } else if (split_k == 10) { - DISPATCH_GATED_NACC_PREFETCH(10); - } else { - DISPATCH_GATED_NACC_PREFETCH(16); - } -#undef DISPATCH_GATED_NACC_PREFETCH -#undef LAUNCH_QPN4_GATED - C10_CUDA_KERNEL_LAUNCH_CHECK(); -} - -void nvfp4_qpn4_gemm_scale_code_sm70_out(torch::Tensor out, torch::Tensor input, - torch::Tensor codes, - torch::Tensor scale_codes, - double global_scale, int64_t split_k, - int64_t accumulator_chains, - bool prefetch_codes) { - TORCH_CHECK(out.is_cuda() && input.is_cuda() && codes.is_cuda() && - scale_codes.is_cuda(), - "nvfp4_qpn4_gemm_scale_code_sm70_out: tensors must be CUDA"); - TORCH_CHECK(out.scalar_type() == torch::kFloat16 && - input.scalar_type() == torch::kFloat16 && - codes.scalar_type() == torch::kUInt8 && - scale_codes.scalar_type() == torch::kUInt8, - "nvfp4_qpn4_gemm_scale_code_sm70_out: dtype mismatch"); - TORCH_CHECK(out.is_contiguous() && input.is_contiguous() && - codes.is_contiguous() && scale_codes.is_contiguous(), - "nvfp4_qpn4_gemm_scale_code_sm70_out: tensors must be " - "contiguous"); - TORCH_CHECK(out.get_device() == input.get_device() && - out.get_device() == codes.get_device() && - out.get_device() == scale_codes.get_device(), - "nvfp4_qpn4_gemm_scale_code_sm70_out: tensors must share one " - "device"); - const int64_t m = input.size(0); - const int64_t k = input.size(1); - const int64_t n = out.size(1); - TORCH_CHECK(m == 1 && out.size(0) == 1, - "nvfp4_qpn4_gemm_scale_code_sm70_out: M must be 1"); - TORCH_CHECK(k > 0 && k % 128 == 0 && n > 0 && n % 32 == 0, - "nvfp4_qpn4_gemm_scale_code_sm70_out: shape alignment " - "mismatch"); - TORCH_CHECK(codes.numel() == k * n / 2 && scale_codes.numel() == k * n / 16, - "nvfp4_qpn4_gemm_scale_code_sm70_out: packed tensor size " - "mismatch"); - TORCH_CHECK(split_k == 4 || split_k == 8 || split_k == 10 || split_k == 16 || - split_k == 17, - "nvfp4_qpn4_gemm_scale_code_sm70_out: unsupported split_k"); - TORCH_CHECK((k / 16) % split_k == 0, - "nvfp4_qpn4_gemm_scale_code_sm70_out: invalid split_k for K"); - TORCH_CHECK(accumulator_chains == 1 || accumulator_chains == 2, - "nvfp4_qpn4_gemm_scale_code_sm70_out: accumulator chains must " - "be 1 or 2"); - - const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - const auto* code_ptr = codes.data_ptr(); - const auto* scale_code_ptr = scale_codes.data_ptr(); - const SplitHalfScale split_scale = - split_half_scale(static_cast(global_scale) * kFp4Bias); - const auto* input_ptr = - reinterpret_cast(input.data_ptr()); - auto* output_ptr = reinterpret_cast(out.data_ptr()); - -#define LAUNCH_QPN4_LUT(SPLIT, NACC, PREFETCH) \ - launch_qpn4( \ - code_ptr, scale_code_ptr, split_scale.hi, split_scale.lo, input_ptr, \ - output_ptr, static_cast(n), static_cast(k), \ - static_cast(m), stream) -#define DISPATCH_QPN4_LUT(SPLIT) \ - do { \ - if (accumulator_chains == 1 && !prefetch_codes) { \ - LAUNCH_QPN4_LUT(SPLIT, 1, false); \ - } else if (accumulator_chains == 2 && !prefetch_codes) { \ - LAUNCH_QPN4_LUT(SPLIT, 2, false); \ - } else if (accumulator_chains == 1) { \ - LAUNCH_QPN4_LUT(SPLIT, 1, true); \ - } else { \ - LAUNCH_QPN4_LUT(SPLIT, 2, true); \ - } \ - } while (0) - if (split_k == 4) { - DISPATCH_QPN4_LUT(4); - } else if (split_k == 8) { - DISPATCH_QPN4_LUT(8); - } else if (split_k == 10) { - DISPATCH_QPN4_LUT(10); - } else if (split_k == 16) { - DISPATCH_QPN4_LUT(16); - } else { - DISPATCH_QPN4_LUT(17); - } -#undef DISPATCH_QPN4_LUT -#undef LAUNCH_QPN4_LUT - C10_CUDA_KERNEL_LAUNCH_CHECK(); -} - -void nvfp4_qpn4_gated_scale_code_sm70_out( - torch::Tensor out, torch::Tensor input, torch::Tensor codes, - torch::Tensor scale_codes, double global_scale, int64_t split_k, - int64_t accumulator_chains, bool prefetch_codes) { - TORCH_CHECK(out.is_cuda() && input.is_cuda() && codes.is_cuda() && - scale_codes.is_cuda(), - "nvfp4_qpn4_gated_scale_code_sm70_out: tensors must be CUDA"); - TORCH_CHECK(out.scalar_type() == torch::kFloat16 && - input.scalar_type() == torch::kFloat16 && - codes.scalar_type() == torch::kUInt8 && - scale_codes.scalar_type() == torch::kUInt8, - "nvfp4_qpn4_gated_scale_code_sm70_out: dtype mismatch"); - TORCH_CHECK(out.is_contiguous() && input.is_contiguous() && - codes.is_contiguous() && scale_codes.is_contiguous(), - "nvfp4_qpn4_gated_scale_code_sm70_out: tensors must be " - "contiguous"); - TORCH_CHECK(out.get_device() == input.get_device() && - out.get_device() == codes.get_device() && - out.get_device() == scale_codes.get_device(), - "nvfp4_qpn4_gated_scale_code_sm70_out: tensors must share one " - "device"); - const int64_t m = input.size(0); - const int64_t k = input.size(1); - const int64_t hidden = out.size(1); - const int64_t n = hidden * 2; - TORCH_CHECK(m == 1 && out.size(0) == 1, - "nvfp4_qpn4_gated_scale_code_sm70_out: M must be 1"); - TORCH_CHECK(k > 0 && k % 128 == 0 && hidden > 0 && hidden % 32 == 0, - "nvfp4_qpn4_gated_scale_code_sm70_out: shape alignment " - "mismatch"); - TORCH_CHECK(codes.numel() == k * n / 2 && scale_codes.numel() == k * n / 16, - "nvfp4_qpn4_gated_scale_code_sm70_out: packed tensor size " - "mismatch"); - TORCH_CHECK(split_k == 4 || split_k == 8 || split_k == 10 || split_k == 16, - "nvfp4_qpn4_gated_scale_code_sm70_out: unsupported split_k"); - TORCH_CHECK((k / 16) % split_k == 0, - "nvfp4_qpn4_gated_scale_code_sm70_out: invalid split_k for K"); - TORCH_CHECK(accumulator_chains == 1 || accumulator_chains == 2, - "nvfp4_qpn4_gated_scale_code_sm70_out: accumulator chains " - "must be 1 or 2"); - - const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - const auto* code_ptr = codes.data_ptr(); - const auto* scale_code_ptr = scale_codes.data_ptr(); - const SplitHalfScale split_scale = - split_half_scale(static_cast(global_scale) * kFp4Bias); - const auto* input_ptr = - reinterpret_cast(input.data_ptr()); - auto* output_ptr = reinterpret_cast(out.data_ptr()); - -#define LAUNCH_QPN4_GATED_LUT(SPLIT, NACC, PREFETCH) \ - launch_qpn4_gated( \ - code_ptr, scale_code_ptr, split_scale.hi, split_scale.lo, input_ptr, \ - output_ptr, static_cast(hidden), static_cast(k), \ - static_cast(m), stream) -#define DISPATCH_QPN4_GATED_LUT(SPLIT) \ - do { \ - if (accumulator_chains == 1 && !prefetch_codes) { \ - LAUNCH_QPN4_GATED_LUT(SPLIT, 1, false); \ - } else if (accumulator_chains == 2 && !prefetch_codes) { \ - LAUNCH_QPN4_GATED_LUT(SPLIT, 2, false); \ - } else if (accumulator_chains == 1) { \ - LAUNCH_QPN4_GATED_LUT(SPLIT, 1, true); \ - } else { \ - LAUNCH_QPN4_GATED_LUT(SPLIT, 2, true); \ - } \ - } while (0) - if (split_k == 4) { - DISPATCH_QPN4_GATED_LUT(4); - } else if (split_k == 8) { - DISPATCH_QPN4_GATED_LUT(8); - } else if (split_k == 10) { - DISPATCH_QPN4_GATED_LUT(10); - } else { - DISPATCH_QPN4_GATED_LUT(16); - } -#undef DISPATCH_QPN4_GATED_LUT -#undef LAUNCH_QPN4_GATED_LUT - C10_CUDA_KERNEL_LAUNCH_CHECK(); -} - -void nvfp4_qpn4_dispatch_sm70_out(torch::Tensor out, int64_t dense_weight_ptr, - torch::Tensor input, torch::Tensor codes, - torch::Tensor scales, double global_scale, - bool use_scale_code, bool gated_silu) { - // Keep the dynamic-M decision inside the opaque operator so one compiled - // graph can safely serve both M=1 decode and large-M prefill. - if (input.size(0) == 1) { - if (gated_silu) { - TORCH_CHECK(use_scale_code, - "QPN4 gated decode requires E4M3 scale codes"); - nvfp4_qpn4_gated_scale_code_sm70_out(out, input, codes, scales, - global_scale, 8, 2, false); - } else if (use_scale_code) { - nvfp4_qpn4_gemm_scale_code_sm70_out(out, input, codes, scales, - global_scale, 17, 1, false); - } else { - nvfp4_qpn4_gemm_sm70_out(out, input, codes, scales, 17, 1, false); - } - return; - } - nvfp4_qpn4_prefill_sm70_out(out, dense_weight_ptr, input, codes, scales, - global_scale, use_scale_code, gated_silu); -} - -#ifdef VLLM_QPN4_STANDALONE -TORCH_LIBRARY_FRAGMENT(_C, ops) { - ops.def("nvfp4_qpn4_prepare_sm70(Tensor qweight, Tensor scales) -> Tensor[]"); - ops.impl("nvfp4_qpn4_prepare_sm70", torch::kCUDA, &nvfp4_qpn4_prepare_sm70); - ops.def( - "nvfp4_qpn4_prepare_scale_code_sm70(Tensor qweight, Tensor scale_codes) " - "-> Tensor[]"); - ops.impl("nvfp4_qpn4_prepare_scale_code_sm70", torch::kCUDA, - &nvfp4_qpn4_prepare_scale_code_sm70); - ops.def( - "nvfp4_qpn4_gemm_sm70_out(Tensor(a!) out, Tensor input, Tensor codes, " - "Tensor scales, int split_k, int accumulator_chains, " - "bool prefetch_codes) -> ()"); - ops.impl("nvfp4_qpn4_gemm_sm70_out", torch::kCUDA, &nvfp4_qpn4_gemm_sm70_out); - ops.def( - "nvfp4_qpn4_gated_sm70_out(Tensor(a!) out, Tensor input, Tensor codes, " - "Tensor scales, int split_k, int accumulator_chains, " - "bool prefetch_codes) -> ()"); - ops.impl("nvfp4_qpn4_gated_sm70_out", torch::kCUDA, - &nvfp4_qpn4_gated_sm70_out); - ops.def( - "nvfp4_qpn4_gemm_scale_code_sm70_out(Tensor(a!) out, Tensor input, " - "Tensor codes, Tensor scale_codes, float global_scale, int split_k, " - "int accumulator_chains, bool prefetch_codes) -> ()"); - ops.impl("nvfp4_qpn4_gemm_scale_code_sm70_out", torch::kCUDA, - &nvfp4_qpn4_gemm_scale_code_sm70_out); - ops.def( - "nvfp4_qpn4_gated_scale_code_sm70_out(Tensor(a!) out, Tensor input, " - "Tensor codes, Tensor scale_codes, float global_scale, int split_k, " - "int accumulator_chains, bool prefetch_codes) -> ()"); - ops.impl("nvfp4_qpn4_gated_scale_code_sm70_out", torch::kCUDA, - &nvfp4_qpn4_gated_scale_code_sm70_out); - ops.def( - "nvfp4_qpn4_dequantize_sm70_out(Tensor(a!) out, Tensor codes, Tensor " - "scales, float global_scale, bool use_scale_code) -> ()"); - ops.impl("nvfp4_qpn4_dequantize_sm70_out", torch::kCUDA, - &nvfp4_qpn4_dequantize_sm70_out); - ops.def( - "nvfp4_qpn4_prefill_sm70_out(Tensor(a!) out, int dense_weight_ptr, " - "Tensor input, Tensor codes, Tensor scales, float global_scale, bool " - "use_scale_code, bool gated_silu) -> ()"); - ops.impl("nvfp4_qpn4_prefill_sm70_out", torch::kCUDA, - &nvfp4_qpn4_prefill_sm70_out); - ops.def( - "nvfp4_qpn4_dispatch_sm70_out(Tensor(a!) out, int dense_weight_ptr, " - "Tensor input, Tensor codes, Tensor scales, float global_scale, bool " - "use_scale_code, bool gated_silu) -> ()"); - ops.impl("nvfp4_qpn4_dispatch_sm70_out", torch::kCUDA, - &nvfp4_qpn4_dispatch_sm70_out); -} -#endif diff --git a/csrc/sm70_turbomind/ops/qwen38_prefill_cutlass.cu b/csrc/sm70_turbomind/ops/qwen38_prefill_cutlass.cu deleted file mode 100644 index bba83cf9b4..0000000000 --- a/csrc/sm70_turbomind/ops/qwen38_prefill_cutlass.cu +++ /dev/null @@ -1,130 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -#include "qwen38_prefill_cutlass.cuh" - -#include -#include - -#include -#include - -#include -#include -#include -#include - -namespace vllm::awq_sm70 { -namespace { - -// Matches CUTLASS's generated -// cutlass_tensorop_f16_s884gemm_f16_128x256_32x2_nn_align8 operation. The -// profiler reports its logical instruction tile as 16x16x4, while the SM70 -// DefaultGemmUniversal specialization is built from native m8n8k4 MMA ops. -using Sm70Fp8PrefillCutlassKernel = - typename cutlass::gemm::kernel::DefaultGemmUniversal< - cutlass::half_t, cutlass::layout::RowMajor, - cutlass::ComplexTransform::kNone, 8, cutlass::half_t, - cutlass::layout::RowMajor, cutlass::ComplexTransform::kNone, 8, - cutlass::half_t, cutlass::layout::RowMajor, float, - cutlass::arch::OpClassTensorOp, cutlass::arch::Sm70, - cutlass::gemm::GemmShape<128, 256, 32>, - cutlass::gemm::GemmShape<64, 64, 32>, cutlass::gemm::GemmShape<8, 8, 4>, - cutlass::epilogue::thread::LinearCombination, - cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<8>, 2, - cutlass::arch::OpMultiplyAdd>::GemmKernel; -using Sm70Fp8PrefillCutlassGemm = - cutlass::gemm::device::GemmUniversalAdapter; - -void maybe_log_sm70_fp8_prefill_cutlass_route(int64_t m, int64_t n, int64_t k) { - const char* raw = std::getenv("VLLM_SM70_PROFILE_TRACE"); - if (raw == nullptr || std::strcmp(raw, "1") != 0) { - return; - } - const unsigned bit = n == 8704 ? 1u - : k == 4352 ? 2u - : k == 1536 ? 4u - : n == 4096 ? 8u - : 16u; - static std::atomic logged_shapes{0}; - const unsigned previous = - logged_shapes.fetch_or(bit, std::memory_order_relaxed); - if ((previous & bit) == 0u) { - std::cerr << "SM70 block-FP8 exact-8K CUTLASS projection route M=" << m - << " N=" << n << " K=" << k << std::endl; - } -} - -} // namespace - -bool sm70_fp8_prefill_cutlass_out(torch::Tensor out, torch::Tensor in_feats, - torch::Tensor dense_weight, bool gated_silu) { - const char* raw = std::getenv("VLLM_SM70_FP8_PREFILL_CUTLASS"); - if ((raw != nullptr && std::atoi(raw) == 0) || gated_silu || - in_feats.dim() != 2 || dense_weight.dim() != 2) { - return false; - } - const int64_t k = in_feats.size(1); - const int64_t n = dense_weight.size(1); - const bool exact_qkv = k == 5120 && (n == 4096 || n == 3584); - const bool exact_gate_up = k == 5120 && n == 8704; - const bool exact_down = k == 4352 && n == 5120; - const bool exact_attention_output = k == 1536 && n == 5120; - if (in_feats.size(0) != 8000 || (!exact_qkv && !exact_gate_up && - !exact_down && !exact_attention_output)) { - return false; - } - - TORCH_CHECK(out.is_cuda() && in_feats.is_cuda() && dense_weight.is_cuda(), - "SM70 block-FP8 prefill CUTLASS GEMM expects CUDA tensors."); - TORCH_CHECK(out.scalar_type() == at::ScalarType::Half && - in_feats.scalar_type() == at::ScalarType::Half && - dense_weight.scalar_type() == at::ScalarType::Half, - "SM70 block-FP8 prefill CUTLASS GEMM expects float16 tensors."); - TORCH_CHECK( - out.is_contiguous() && in_feats.is_contiguous() && - dense_weight.is_contiguous(), - "SM70 block-FP8 prefill CUTLASS GEMM expects contiguous tensors."); - TORCH_CHECK(out.dim() == 2 && dense_weight.size(0) == k && - out.size(0) == in_feats.size(0) && out.size(1) == n, - "SM70 block-FP8 prefill CUTLASS GEMM shape mismatch."); - - const at::cuda::OptionalCUDAGuard device_guard(device_of(in_feats)); - const int m = static_cast(in_feats.size(0)); - const int problem_n = static_cast(n); - const int problem_k = static_cast(k); - auto* dense_ptr = - reinterpret_cast(dense_weight.data_ptr()); - auto* input_ptr = - reinterpret_cast(in_feats.data_ptr()); - auto* output_ptr = - reinterpret_cast(out.data_ptr()); - maybe_log_sm70_fp8_prefill_cutlass_route(m, problem_n, problem_k); - typename Sm70Fp8PrefillCutlassGemm::Arguments arguments{ - cutlass::gemm::GemmUniversalMode::kGemm, - {problem_n, m, problem_k}, - 1, - {1.0f, 0.0f}, - dense_ptr, - input_ptr, - output_ptr, - output_ptr, - 0, - 0, - 0, - 0, - problem_n, - problem_k, - problem_n, - problem_n}; - Sm70Fp8PrefillCutlassGemm gemm; - const cutlass::Status status = - gemm(arguments, nullptr, at::cuda::getCurrentCUDAStream()); - TORCH_CHECK(status == cutlass::Status::kSuccess, - "SM70 block-FP8 prefill CUTLASS GEMM failed: ", - cutlassGetStatusString(status), "."); - return true; -} - -} // namespace vllm::awq_sm70 diff --git a/csrc/sm70_turbomind/ops/qwen38_prefill_cutlass.cuh b/csrc/sm70_turbomind/ops/qwen38_prefill_cutlass.cuh deleted file mode 100644 index 23ce6a1f7f..0000000000 --- a/csrc/sm70_turbomind/ops/qwen38_prefill_cutlass.cuh +++ /dev/null @@ -1,13 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -#pragma once - -#include - -namespace vllm::awq_sm70 { - -bool sm70_fp8_prefill_cutlass_out(torch::Tensor out, torch::Tensor in_feats, - torch::Tensor dense_weight, bool gated_silu); - -} // namespace vllm::awq_sm70 diff --git a/csrc/sm70_turbomind/ops/sm70_top20_philox_fragment.cpp b/csrc/sm70_turbomind/ops/sm70_top20_philox_fragment.cpp deleted file mode 100644 index bc8534fcdf..0000000000 --- a/csrc/sm70_turbomind/ops/sm70_top20_philox_fragment.cpp +++ /dev/null @@ -1,52 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -#include -#include - -#include - -void sm70_sample_sorted_top20_philox_out( - torch::Tensor sampled_token_out, torch::Tensor sparse_ids_out, - torch::Tensor sparse_probs_out, torch::Tensor top_values, - torch::Tensor top_indices, const std::optional& generator, - int64_t vocab_size, double top_p); - -void sm70_sample_sorted_top20_philox_token_out( - torch::Tensor sampled_token_out, torch::Tensor top_values, - torch::Tensor top_indices, const std::optional& generator, - int64_t vocab_size, double top_p); - -void sm70_sample_chunked_top20_philox_token_out( - torch::Tensor sampled_token_out, torch::Tensor global_values, - torch::Tensor local_indices, torch::Tensor global_positions, - const std::optional& generator, int64_t vocab_size, - double top_p, int64_t chunk_size); - -// Keep these registrations in a wheel-safe sidecar. Adding the sampler must -// not relink or perturb the speed- and quality-frozen primary vllm._C module; -// Python loads this fragment only for the explicitly admitted SM70 route. -TORCH_LIBRARY_FRAGMENT(_C, ops) { - ops.def( - "sm70_sample_sorted_top20_philox_out(Tensor(a!) sampled_token_out, " - "Tensor(b!) sparse_ids_out, Tensor(c!) sparse_probs_out, " - "Tensor top_values, Tensor top_indices, Generator? generator, " - "int vocab_size, float top_p) -> ()"); - ops.impl("sm70_sample_sorted_top20_philox_out", torch::kCUDA, - &sm70_sample_sorted_top20_philox_out); - - ops.def( - "sm70_sample_sorted_top20_philox_token_out(Tensor(a!) " - "sampled_token_out, Tensor top_values, Tensor top_indices, " - "Generator? generator, int vocab_size, float top_p) -> ()"); - ops.impl("sm70_sample_sorted_top20_philox_token_out", torch::kCUDA, - &sm70_sample_sorted_top20_philox_token_out); - - ops.def( - "sm70_sample_chunked_top20_philox_token_out(Tensor(a!) " - "sampled_token_out, Tensor global_values, Tensor local_indices, " - "Tensor global_positions, Generator? generator, int vocab_size, " - "float top_p, int chunk_size) -> ()"); - ops.impl("sm70_sample_chunked_top20_philox_token_out", torch::kCUDA, - &sm70_sample_chunked_top20_philox_token_out); -} diff --git a/csrc/sm70_turbomind/ops/sm70_top20_philox_sampler.cu b/csrc/sm70_turbomind/ops/sm70_top20_philox_sampler.cu deleted file mode 100644 index a09c094b98..0000000000 --- a/csrc/sm70_turbomind/ops/sm70_top20_philox_sampler.cu +++ /dev/null @@ -1,289 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// Exact sparse Philox sampling for the SM70 Qwen3.8 batch-one top-20 route. - -#include - -#include -#include -#include -#include -#include -#include -#include -#include - -#include -#include -#include -#include -#include - -namespace { - -constexpr int kTopK = 20; -constexpr int kWarpSize = 32; -constexpr int kNativeBlockSize = 256; -constexpr int kNativeUnroll = 4; -constexpr int kMaxGeneratorOffsetsPerCurandCall = 4; - -__device__ __forceinline__ bool top1_better(float value, int token, - float best_value, int best_token) { - return value > best_value || (value == best_value && token < best_token); -} - -__device__ __forceinline__ float native_exponential_for_index( - int64_t index, int64_t total_threads, at::PhiloxCudaState philox_args) { - auto [seed, offset] = at::cuda::philox::unpack(philox_args); - const int64_t values_per_round = total_threads * kNativeUnroll; - const int64_t round = index / values_per_round; - const int64_t within_round = index - round * values_per_round; - const uint64_t subsequence = - static_cast(within_round % total_threads); - const int component = static_cast(within_round / total_threads); - - curandStatePhilox4_32_10_t state; - curand_init(seed, subsequence, offset, &state); - float4 random = make_float4(0.f, 0.f, 0.f, 0.f); - for (int64_t iteration = 0; iteration <= round; ++iteration) { - random = curand_uniform4(&state); - } - const float uniform = reinterpret_cast(&random)[component]; - return at::transformation::exponential(uniform, 1.f); -} - -template -__global__ void sm70_sample_sorted_top20_philox_kernel( - int64_t* sampled_token_out, int64_t* sparse_ids_out, - float* sparse_probs_out, const float* top_values, - const int64_t* top_indices, const int64_t* local_indices, - const int64_t* global_positions, int64_t local_candidate_count, - int chunk_size, int64_t vocab_size, int64_t total_threads, float top_p, - at::PhiloxCudaState philox_args) { - const int lane = threadIdx.x; - __shared__ float selected_values[kTopK]; - __shared__ int64_t selected_ids[kTopK]; - __shared__ float selected_random[kTopK]; - - int64_t lane_token64 = -1; - if (lane < kTopK) { - if (ChunkedIndices) { - const int64_t position = global_positions[lane]; - if (position >= 0 && position < local_candidate_count) { - const int64_t local_index = local_indices[position]; - if (local_index >= 0 && local_index < chunk_size) { - lane_token64 = (position / kTopK) * chunk_size + local_index; - } - } - } else { - lane_token64 = top_indices[lane]; - } - } - const bool lane_valid = - lane < kTopK && lane_token64 >= 0 && lane_token64 < vocab_size; - const int lane_token = lane_valid ? static_cast(lane_token64) : -1; - const float lane_raw_value = lane < kTopK ? top_values[lane] : -FLT_MAX; - const float lane_value = isfinite(lane_raw_value) ? lane_raw_value : -FLT_MAX; - const int valid_count = __popc(__ballot_sync(0xffffffff, lane_valid)); - int lane_rank = 0; -#pragma unroll - for (int source = 0; source < kTopK; ++source) { - const int other_token = __shfl_sync(0xffffffff, lane_token, source); - const float other_value = __shfl_sync(0xffffffff, lane_value, source); - if (lane_valid && other_token >= 0 && other_token < vocab_size && - top1_better(other_value, other_token, lane_value, lane_token)) { - ++lane_rank; - } - } - if (lane_valid) { - selected_values[lane_rank] = lane_value; - selected_ids[lane_rank] = static_cast(lane_token); - } else if (lane >= valid_count && lane < kTopK) { - // Preserve the prior fallback for malformed candidate ids. - selected_values[lane] = -FLT_MAX; - selected_ids[lane] = lane; - } - __syncwarp(); - - if (lane < kTopK) { - selected_random[lane] = native_exponential_for_index( - selected_ids[lane], total_threads, philox_args); - } - __syncwarp(); - - if (lane == 0) { - const float max_value = selected_values[0]; - float probability_sum = 0.f; - for (int index = 0; index < kTopK; ++index) { - const float probability = __expf(selected_values[index] - max_value); - selected_values[index] = probability; - probability_sum += probability; - } - - float cumulative_probability = 0.f; - float filtered_probability_sum = 0.f; - for (int index = 0; index < kTopK; ++index) { - const float probability = selected_values[index] / probability_sum; - const bool keep = index == 0 || cumulative_probability < top_p; - cumulative_probability += probability; - selected_values[index] = keep ? probability : 0.f; - filtered_probability_sum += selected_values[index]; - } - - float best_score = -FLT_MAX; - int64_t sampled_token = selected_ids[0]; - for (int index = 0; index < kTopK; ++index) { - const float probability = - selected_values[index] / filtered_probability_sum; - if (EmitMetadata) { - sparse_probs_out[index] = probability; - sparse_ids_out[index] = selected_ids[index]; - } - const float score = probability / selected_random[index]; - if (score > best_score) { - best_score = score; - sampled_token = selected_ids[index]; - } - } - sampled_token_out[0] = sampled_token; - } -} - -template -void launch_sm70_sample_sorted_top20_philox( - torch::Tensor sampled_token_out, int64_t* sparse_ids_out, - float* sparse_probs_out, torch::Tensor top_values, - const int64_t* top_indices, const int64_t* local_indices, - const int64_t* global_positions, int64_t local_candidate_count, - int chunk_size, const std::optional& generator, - int64_t vocab_size, double top_p) { - const auto* properties = at::cuda::getCurrentDeviceProperties(); - const int64_t blocks_per_sm = - properties->maxThreadsPerMultiProcessor / kNativeBlockSize; - const int64_t required_blocks = - (vocab_size + kNativeBlockSize - 1) / kNativeBlockSize; - const int64_t grid_blocks = std::min( - required_blocks, properties->multiProcessorCount * blocks_per_sm); - const int64_t total_threads = grid_blocks * kNativeBlockSize; - const uint64_t counter_offset = - ((vocab_size - 1) / (total_threads * kNativeUnroll) + 1) * - kMaxGeneratorOffsetsPerCurandCall; - - auto* cuda_generator = at::get_generator_or_default( - generator, at::cuda::detail::getDefaultCUDAGenerator()); - at::PhiloxCudaState philox_args; - { - std::lock_guard lock(cuda_generator->mutex_); - philox_args = cuda_generator->philox_cuda_state(counter_offset); - } - - sm70_sample_sorted_top20_philox_kernel - <<<1, kWarpSize, 0, at::cuda::getCurrentCUDAStream()>>>( - sampled_token_out.data_ptr(), sparse_ids_out, - sparse_probs_out, top_values.data_ptr(), top_indices, - local_indices, global_positions, local_candidate_count, chunk_size, - vocab_size, total_threads, static_cast(top_p), philox_args); - C10_CUDA_KERNEL_LAUNCH_CHECK(); -} - -} // namespace - -void sm70_sample_sorted_top20_philox_out( - torch::Tensor sampled_token_out, torch::Tensor sparse_ids_out, - torch::Tensor sparse_probs_out, torch::Tensor top_values, - torch::Tensor top_indices, const std::optional& generator, - int64_t vocab_size, double top_p) { - TORCH_CHECK(sampled_token_out.is_cuda() && sparse_ids_out.is_cuda() && - sparse_probs_out.is_cuda() && top_values.is_cuda() && - top_indices.is_cuda(), - "sm70_sample_sorted_top20_philox_out: tensors must be CUDA"); - TORCH_CHECK(sampled_token_out.scalar_type() == torch::kInt64 && - sparse_ids_out.scalar_type() == torch::kInt64 && - sparse_probs_out.scalar_type() == torch::kFloat32 && - top_values.scalar_type() == torch::kFloat32 && - top_indices.scalar_type() == torch::kInt64, - "sm70_sample_sorted_top20_philox_out: dtype mismatch"); - TORCH_CHECK(sampled_token_out.numel() == 1 && - sparse_ids_out.numel() == kTopK && - sparse_probs_out.numel() == kTopK && - top_values.numel() == kTopK && top_indices.numel() == kTopK, - "sm70_sample_sorted_top20_philox_out: shape mismatch"); - TORCH_CHECK( - sampled_token_out.is_contiguous() && sparse_ids_out.is_contiguous() && - sparse_probs_out.is_contiguous() && top_values.is_contiguous() && - top_indices.is_contiguous(), - "sm70_sample_sorted_top20_philox_out: tensors must be contiguous"); - TORCH_CHECK(vocab_size >= kTopK && top_p > 0.0 && top_p <= 1.0, - "sm70_sample_sorted_top20_philox_out: invalid parameters"); - - const at::cuda::OptionalCUDAGuard device_guard(device_of(top_values)); - launch_sm70_sample_sorted_top20_philox( - sampled_token_out, sparse_ids_out.data_ptr(), - sparse_probs_out.data_ptr(), top_values, - top_indices.data_ptr(), nullptr, nullptr, 0, 0, generator, - vocab_size, top_p); -} - -void sm70_sample_sorted_top20_philox_token_out( - torch::Tensor sampled_token_out, torch::Tensor top_values, - torch::Tensor top_indices, const std::optional& generator, - int64_t vocab_size, double top_p) { - TORCH_CHECK( - sampled_token_out.is_cuda() && top_values.is_cuda() && - top_indices.is_cuda(), - "sm70_sample_sorted_top20_philox_token_out: tensors must be CUDA"); - TORCH_CHECK(sampled_token_out.scalar_type() == torch::kInt64 && - top_values.scalar_type() == torch::kFloat32 && - top_indices.scalar_type() == torch::kInt64, - "sm70_sample_sorted_top20_philox_token_out: dtype mismatch"); - TORCH_CHECK(sampled_token_out.numel() == 1 && top_values.numel() == kTopK && - top_indices.numel() == kTopK, - "sm70_sample_sorted_top20_philox_token_out: shape mismatch"); - TORCH_CHECK( - sampled_token_out.is_contiguous() && top_values.is_contiguous() && - top_indices.is_contiguous(), - "sm70_sample_sorted_top20_philox_token_out: tensors must be contiguous"); - TORCH_CHECK(vocab_size >= kTopK && top_p > 0.0 && top_p <= 1.0, - "sm70_sample_sorted_top20_philox_token_out: invalid parameters"); - - const at::cuda::OptionalCUDAGuard device_guard(device_of(top_values)); - launch_sm70_sample_sorted_top20_philox( - sampled_token_out, nullptr, nullptr, top_values, - top_indices.data_ptr(), nullptr, nullptr, 0, 0, generator, - vocab_size, top_p); -} - -void sm70_sample_chunked_top20_philox_token_out( - torch::Tensor sampled_token_out, torch::Tensor global_values, - torch::Tensor local_indices, torch::Tensor global_positions, - const std::optional& generator, int64_t vocab_size, - double top_p, int64_t chunk_size) { - TORCH_CHECK( - sampled_token_out.is_cuda() && global_values.is_cuda() && - local_indices.is_cuda() && global_positions.is_cuda(), - "sm70_sample_chunked_top20_philox_token_out: tensors must be CUDA"); - TORCH_CHECK(sampled_token_out.scalar_type() == torch::kInt64 && - global_values.scalar_type() == torch::kFloat32 && - local_indices.scalar_type() == torch::kInt64 && - global_positions.scalar_type() == torch::kInt64, - "sm70_sample_chunked_top20_philox_token_out: dtype mismatch"); - TORCH_CHECK(sampled_token_out.numel() == 1 && - global_values.numel() == kTopK && local_indices.dim() == 2 && - local_indices.size(1) == kTopK && - global_positions.numel() == kTopK, - "sm70_sample_chunked_top20_philox_token_out: shape mismatch"); - TORCH_CHECK( - sampled_token_out.is_contiguous() && global_values.is_contiguous() && - local_indices.is_contiguous() && global_positions.is_contiguous(), - "sm70_sample_chunked_top20_philox_token_out: tensors must be contiguous"); - TORCH_CHECK(chunk_size > 0 && - local_indices.size(0) * chunk_size == vocab_size && - vocab_size >= kTopK && top_p > 0.0 && top_p <= 1.0, - "sm70_sample_chunked_top20_philox_token_out: invalid parameters"); - - const at::cuda::OptionalCUDAGuard device_guard(device_of(global_values)); - launch_sm70_sample_sorted_top20_philox( - sampled_token_out, nullptr, nullptr, global_values, nullptr, - local_indices.data_ptr(), global_positions.data_ptr(), - local_indices.numel(), static_cast(chunk_size), generator, - vocab_size, top_p); -} diff --git a/csrc/sm70_turbomind/ops/tm_registry_sm70.cu b/csrc/sm70_turbomind/ops/tm_registry_sm70.cu index 7d692f848f..f4102ef0e3 100644 --- a/csrc/sm70_turbomind/ops/tm_registry_sm70.cu +++ b/csrc/sm70_turbomind/ops/tm_registry_sm70.cu @@ -1,38 +1,39 @@ -// 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 38278ee057..bd5b2c45e2 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -165,13 +165,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "Tensor? qzeros_or_none, bool inplace) -> Tensor"); // conditionally compiled so impl registrations are in source file - #ifdef ENABLE_SM70_TURBOMIND +#ifdef ENABLE_SM70_TURBOMIND ops.def("silu_and_mul_interleaved(Tensor! result, Tensor input) -> ()"); - ops.impl("silu_and_mul_interleaved", torch::kCUDA, &silu_and_mul_interleaved); + ops.impl("silu_and_mul_interleaved", torch::kCUDA, + &silu_and_mul_interleaved); ops.def( - "awq_sm70_prepare(Tensor _kernel, Tensor _scaling_factors, Tensor " - "_zeros, " + "awq_sm70_prepare(Tensor _kernel, Tensor _scaling_factors, Tensor _zeros, " "int group_size, bool interleave_gated_silu) -> Tensor[]"); ops.impl("awq_sm70_prepare", torch::kCUDA, &awq_sm70_prepare); @@ -193,7 +193,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.def( "fp8_sm70_dequantize_out(Tensor(a!) out, Tensor _kernel, " "Tensor _scaling_factors, int group_size) -> ()"); - ops.impl("fp8_sm70_dequantize_out", torch::kCUDA, &fp8_sm70_dequantize_out); + ops.impl("fp8_sm70_dequantize_out", torch::kCUDA, + &fp8_sm70_dequantize_out); ops.def( "mxfp4_sm70_prepare(Tensor _kernel, Tensor _scaling_factors, " @@ -236,128 +237,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "bool gated_silu) -> ()"); ops.impl("fp8_gemm_sm70_out", torch::kCUDA, &fp8_gemm_sm70_out); - 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( - "fp8_qpn8_dequantize_sm70_out(Tensor(a!) out, Tensor codes, " - "Tensor group_scales) -> ()"); - ops.impl("fp8_qpn8_dequantize_sm70_out", torch::kCUDA, - &fp8_qpn8_dequantize_sm70_out); - - ops.def( - "fp8_qpn8_prefill_sm70_out(Tensor(a!) out, int dense_weight_ptr, " - "Tensor input, Tensor codes, Tensor group_scales, bool gated_silu) -> " - "()"); - ops.impl("fp8_qpn8_prefill_sm70_out", torch::kCUDA, - &fp8_qpn8_prefill_sm70_out); - - ops.def( - "fp8_qpn8_dispatch_sm70_out(Tensor(a!) out, int dense_weight_ptr, " - "Tensor input, Tensor codes, Tensor group_scales, int split_k, " - "int accumulator_chains, bool prefetch_codes, bool gated_silu) -> ()"); - ops.impl("fp8_qpn8_dispatch_sm70_out", torch::kCUDA, - &fp8_qpn8_dispatch_sm70_out); - - ops.def( - "fp8_qpn8_gemm_sm70_out(Tensor(a!) out, Tensor input, Tensor codes, " - "Tensor group_scales, int split_k, int accumulator_chains, " - "bool fast_decoder, bool prefetch_codes) -> ()"); - ops.impl("fp8_qpn8_gemm_sm70_out", torch::kCUDA, &fp8_qpn8_gemm_sm70_out); - - ops.def( - "fp8_qpn8_gemm_ba_split_sm70_out(Tensor(a!) qkv_out, Tensor(b!) " - "z_out, Tensor(c!) b_out, Tensor(d!) a_out, Tensor input, Tensor codes, " - "Tensor group_scales, Tensor ba_weight) -> ()"); - ops.impl("fp8_qpn8_gemm_ba_split_sm70_out", torch::kCUDA, - &fp8_qpn8_gemm_ba_split_sm70_out); - - ops.def( - "fp8_qpn8_dispatch_ba_split_sm70_out(Tensor(a!) qkv_out, Tensor(b!) " - "z_out, Tensor(c!) b_out, Tensor(d!) a_out, Tensor(e!) qkvz_staging, " - "Tensor(f!) ba_staging, int dense_weight_ptr, Tensor input, Tensor " - "codes, Tensor group_scales, Tensor ba_weight) -> ()"); - ops.impl("fp8_qpn8_dispatch_ba_split_sm70_out", torch::kCUDA, - &fp8_qpn8_dispatch_ba_split_sm70_out); - - ops.def( - "fp8_qpn8_gated_pair_sm70_out(Tensor(a!) out, Tensor input, " - "Tensor codes, Tensor group_scales, int split_k, " - "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("nvfp4_qpn4_prepare_sm70(Tensor qweight, Tensor scales) -> Tensor[]"); - ops.impl("nvfp4_qpn4_prepare_sm70", torch::kCUDA, &nvfp4_qpn4_prepare_sm70); - - ops.def( - "nvfp4_qpn4_prepare_scale_code_sm70(Tensor qweight, Tensor " - "scale_codes) -> Tensor[]"); - ops.impl("nvfp4_qpn4_prepare_scale_code_sm70", torch::kCUDA, - &nvfp4_qpn4_prepare_scale_code_sm70); - - ops.def( - "nvfp4_qpn4_dequantize_sm70_out(Tensor(a!) out, Tensor codes, Tensor " - "scales, float global_scale, bool use_scale_code) -> ()"); - ops.impl("nvfp4_qpn4_dequantize_sm70_out", torch::kCUDA, - &nvfp4_qpn4_dequantize_sm70_out); - - ops.def( - "nvfp4_qpn4_prefill_sm70_out(Tensor(a!) out, int dense_weight_ptr, " - "Tensor input, Tensor codes, Tensor scales, float global_scale, bool " - "use_scale_code, bool gated_silu) -> ()"); - ops.impl("nvfp4_qpn4_prefill_sm70_out", torch::kCUDA, - &nvfp4_qpn4_prefill_sm70_out); - - ops.def( - "nvfp4_qpn4_dispatch_sm70_out(Tensor(a!) out, int dense_weight_ptr, " - "Tensor input, Tensor codes, Tensor scales, float global_scale, bool " - "use_scale_code, bool gated_silu) -> ()"); - ops.impl("nvfp4_qpn4_dispatch_sm70_out", torch::kCUDA, - &nvfp4_qpn4_dispatch_sm70_out); - - ops.def( - "fp8_gemm_sm70_prefill_prescaled_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_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[]"); - ops.impl("nvfp4_qpn2_prepare_sm70", torch::kCUDA, &nvfp4_qpn2_prepare_sm70); - - ops.def( - "nvfp4_qpn2_gemm_sm70_out(Tensor(a!) out, Tensor input, Tensor codes, " - "Tensor scales, float global_scale, int split_k, " - "int accumulator_chains) -> ()"); - ops.impl("nvfp4_qpn2_gemm_sm70_out", torch::kCUDA, &nvfp4_qpn2_gemm_sm70_out); - - ops.def( - "nvfp4_qpn2_gated_sm70_out(Tensor(a!) out, Tensor input, Tensor codes, " - "Tensor scales, float global_scale, int split_k, " - "int accumulator_chains) -> ()"); - ops.impl("nvfp4_qpn2_gated_sm70_out", torch::kCUDA, - &nvfp4_qpn2_gated_sm70_out); - - ops.def( - "nvfp4_qpn2_dispatch_sm70_out(Tensor(a!) out, Tensor input, " - "Tensor codes, Tensor scales, float global_scale, int split_k, " - "int accumulator_chains, Tensor tm_weight, Tensor tm_scales, " - "int tm_group_size, int tm_k_ld, int tm_q_ld, bool gated_silu) -> ()"); - ops.impl("nvfp4_qpn2_dispatch_sm70_out", torch::kCUDA, - &nvfp4_qpn2_dispatch_sm70_out); - ops.def( "fp8_gemm_sm70_prefill_dispatch_out(Tensor(a!) out, " "int dense_weight_ptr, Tensor _in_feats, Tensor _kernel, " @@ -382,63 +261,40 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "nvfp4_gemv_sm70_raw_out(Tensor(a!) out, Tensor _in_feats, " "Tensor _kernel, Tensor _scaling_factors, Tensor(b!) partials, " "int group_size, int split_k) -> ()"); - ops.impl("nvfp4_gemv_sm70_raw_out", torch::kCUDA, &nvfp4_gemv_sm70_raw_out); + ops.impl("nvfp4_gemv_sm70_raw_out", torch::kCUDA, + &nvfp4_gemv_sm70_raw_out); ops.def( "nvfp4_gemv_sm70_warp_out(Tensor(a!) out, Tensor _in_feats, " "Tensor _kernel, Tensor _scaling_factors, int group_size) -> ()"); - ops.impl("nvfp4_gemv_sm70_warp_out", torch::kCUDA, &nvfp4_gemv_sm70_warp_out); + ops.impl("nvfp4_gemv_sm70_warp_out", torch::kCUDA, + &nvfp4_gemv_sm70_warp_out); ops.def( "nvfp4_gemv_sm70_h2_out(Tensor(a!) out, Tensor _in_feats, " "Tensor _kernel, Tensor _scaling_factors, Tensor(b!) partials, " "int group_size, int split_k) -> ()"); - ops.impl("nvfp4_gemv_sm70_h2_out", torch::kCUDA, &nvfp4_gemv_sm70_h2_out); + ops.impl("nvfp4_gemv_sm70_h2_out", torch::kCUDA, + &nvfp4_gemv_sm70_h2_out); ops.def( "fp8_gemm_sm70_out_auto(Tensor(a!) out, Tensor _in_feats, " "Tensor _kernel, Tensor _scaling_factors) -> ()"); - ops.impl("fp8_gemm_sm70_out_auto", torch::kCUDA, &fp8_gemm_sm70_out_auto); + ops.impl("fp8_gemm_sm70_out_auto", torch::kCUDA, + &fp8_gemm_sm70_out_auto); ops.def( "fp8_gemm_sm70_out_meta(Tensor(a!) out, Tensor _in_feats, " "Tensor _kernel, Tensor _scaling_factors, Tensor _meta, " "bool gated_silu) -> ()"); - ops.impl("fp8_gemm_sm70_out_meta", torch::kCUDA, &fp8_gemm_sm70_out_meta); + ops.impl("fp8_gemm_sm70_out_meta", torch::kCUDA, + &fp8_gemm_sm70_out_meta); ops.def( "sm70_f16_gemm_out(Tensor(a!) out, Tensor _in_feats, Tensor _kernel, " "int k_ld, bool gated_silu) -> ()"); ops.impl("sm70_f16_gemm_out", torch::kCUDA, &sm70_f16_gemm_out); - ops.def( - "sm70_f16_indexed_rerank_out(Tensor(a!) out, Tensor _in_feats, " - "Tensor _kernel, Tensor candidate_ids, Tensor(b!) selected_raw, " - "Tensor(c!) selected_packed, Tensor(d!) expanded, Tensor(e!) partials, " - "Tensor(f!) barriers, int cta_n, int split_k) -> ()"); - ops.impl("sm70_f16_indexed_rerank_out", torch::kCUDA, - &sm70_f16_indexed_rerank_out); - - ops.def( - "sm70_f16_indexed_rerank_packed_out(Tensor(a!) out, " - "Tensor _in_feats, Tensor _packed_kernel, Tensor candidate_ids, " - "Tensor(b!) selected_packed, Tensor(c!) expanded, " - "Tensor(d!) partials, Tensor(e!) barriers, int cta_n, int split_k) " - "-> ()"); - ops.impl("sm70_f16_indexed_rerank_packed_out", torch::kCUDA, - &sm70_f16_indexed_rerank_packed_out); - - ops.def( - "sm70_f16_rerank_keys_out(Tensor(a!) keys, Tensor logits, " - "Tensor candidate_ids) -> ()"); - ops.impl("sm70_f16_rerank_keys_out", torch::kCUDA, &sm70_f16_rerank_keys_out); - - ops.def( - "sm70_f16_rerank_topk_out(Tensor(a!) values_out, " - "Tensor(b!) ids_out, Tensor logits, Tensor candidate_ids, " - "int vocab_start_index) -> ()"); - ops.impl("sm70_f16_rerank_topk_out", torch::kCUDA, &sm70_f16_rerank_topk_out); - ops.def( "sm70_f16_lm_head_top1_out(Tensor(a!) values_out, " "Tensor(b!) indices_out, Tensor _in_feats, Tensor _kernel, int k_ld, " @@ -644,14 +500,6 @@ 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( - "nvfp4_moe_dense_stage_sm70_out(" - "Tensor(a!) out, Tensor input, Tensor expert_offsets, " - "Tensor dense_expert_ids, Tensor ptrs_w, Tensor ptrs_s, " - "int num_experts, int k, int n, int group_size) -> ()"); - ops.impl("nvfp4_moe_dense_stage_sm70_out", torch::kCUDA, - &nvfp4_moe_dense_stage_sm70_out); - ops.def( "mxfp4_moe_single_token_prepare_w13_sm70_out(" "Tensor(a!) gate_up, Tensor(b!) compact_input, Tensor x, " @@ -728,7 +576,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "bool exact_per_route) -> ()"); ops.impl("fp8_moe_single_token_sm70_out", torch::kCUDA, &fp8_moe_single_token_sm70_out); - #endif +#endif #endif @@ -814,14 +662,6 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) { "int reg_buffer, int reg_buffer_sz_bytes, float epsilon) -> ()"); custom_ar.impl("sm70_tp4_all_reduce_gemma_rms_norm", torch::kCUDA, &sm70_tp4_all_reduce_gemma_rms_norm); - custom_ar.def( - "sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather(int fa, Tensor inp, " - "Tensor residual, Tensor weight, Tensor! normalized_out, Tensor! " - "residual_out, int reg_input_buffer, int reg_output_buffer, int " - "reg_buffer_sz_bytes, float epsilon) -> ()"); - custom_ar.impl("sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather", - torch::kCUDA, - &sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather); custom_ar.def( "all_reduce_sum2(int fa, Tensor inp_a, Tensor inp_b, Tensor! out) -> ()"); custom_ar.impl("all_reduce_sum2", torch::kCUDA, &all_reduce_sum2); @@ -849,12 +689,8 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) { custom_ar.def("dispose", &dispose); custom_ar.def("meta_size", &meta_size); - custom_ar.def("sm70_tp4_push_allreduce_buffer_size", - &sm70_tp4_push_allreduce_buffer_size); custom_ar.def("register_buffer", ®ister_buffer); - custom_ar.def("register_sm70_tp4_push_allreduce_buffer", - ®ister_sm70_tp4_push_allreduce_buffer); custom_ar.def("get_graph_buffer_ipc_meta", &get_graph_buffer_ipc_meta); custom_ar.def("register_graph_buffers", ®ister_graph_buffers); diff --git a/docs/assets/wechat-group-qr-7.png b/docs/assets/wechat-group-qr-7.png deleted file mode 100644 index a1aaf540dc2978225cfd4efd9c42673c9e53a053..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 529331 zcmeFZg;N~O7xxPU4-yFO?(R;2V8Md>;_mK{00F|{1a}C!0fOtY1PvZ!ad&suOXPRo zs`vf}w;rmtcBW@wrn>u_{+`b{&FA;ZvgoKps4y@v=<;%J)nH)YpI~4R?~&o2Th^}% zfX`pBZfdfUFcrfjJI^O?Ep+59m6TwZo{y1XV8d--;D6up{2_Y&z`(r7g@t+X{0;lN zEf?;8SK*&>U;OVF@&5OXFYVf&fAbBD{96ePZ`ggXZUc#|Z??YfS)OmI+55Wb;Ss5*&HSnA zemMnT8JTAc6`96;fTy8=MUi_0hyCwxpXcMQ^2l(X{QqhCcQyFK?$pu$r`0p(!>-rW z)HLC{|NGMJl+*v+1pnXN|MkiLJt_X1C;!>b|M5Zpd5Hf!#D5;*e=G8T*6^P-{AUgS zS;Jom{-09%PbvMUl>Qg4{$~yUzp;ilhVXy+!$EMIN1&n# zyQk)RTj+s?lX_5=+d^|Y`d7yHw`xd`3+cizhD*p1&{tGT{m&bJFM!v82seLMR(K{T z?GWCNAU;4(H7+YM0U7(oFbML~hA-nuKiog-wz*>bO0#B~V`*l>bJeL8%=Q#*`s7GH zF4QL4y0__J_z3low@}Np?{NACmj##_o}=cJ{Yzzz5n;W=Ah2Dsf7ZckAl4gDu0tGW z7h_za2GwVWT=B7^xjhf2?dJQdx(BZ1E21%CAqV}dX3q*{Ncr`Vh|5BPVeeD(>HRKj zq%U|F&d@?qSUV&CmRRs<_(nkF?_c$jdvkGj>6H^g$cJ~ER@o+!_QK1&^;CcwmiKRy|=5a0Uq?3e;&|7P+5M;kSlq&rx}P#1?lg!tW=F+rlv%`lN3s;0Fp`b zY+2ZG=RmG$bC36Jsu=R}G%xy#wU&k8jvk92n#VT(yukhqP0CdO0c2{Wb6xvq*;yR}54z+x5;?(12sxHXz z(F^TjX&vfu)*SQX5u&*F(+Nq#{E$aDXGGm3IqLpA(x2=}<)hnz`B2JlmiIA8|M|H= z`j3}v&+oDPL;@#2nCwN2n^X&I!~y5YU8@^vSNFL`cY7szrPRDD5VUaAwUExU)y?J^ z*;ATd)3e`=NK-15I#9w}%mxxoE@{>q#E+NJk&oJq3Gy-!2p544sV{31Ui$KD2~AxK5ivMt3tD3OYk(l&66Hzo%acU?MKXWh=Q6##jTl&ryEkRp- z{<)#4XREH%O-=aIrAJgByyZL42;CTm*nsS38@`$=xkui$XtYg^G@m={G&o)AqS(iE z@5d99gggxyd7berUMX1*Xw{H9X;*or%OGp8Ch3>Vn;OQ+Fq$xAvux5ZU6f=bMe<JB71g8iJE7M_maC_-=;ukh^F=ach(;cDQ6VoVTo38>e!Ro7-8^Rwu7_7(>Y? z8SIts8qXbtzn-38b3)a{(&OS=i4^bI+KnS#NY1rz-itW07Pj^x#f>6Nn6t7-y zB%_V6V@G*^`{4C;aY{>N>c}1);~Q^Cu|(L=aa`LM^VhHno%(dj5u;PKtpOQx!Wq2k z|9lT0Ec9qjaWMW*Pk6nWHZRm*Dqyngf5j~?o5u`Xap%ooEz0mM)YIn`f+~Ax*`~f5 z%P6dt)EHx9qFePq_hTaxI6t9tAMIYnVKyqBsGz(q^I-9p%i5sZA2ffpTiDn1iDjtl z0)WYwIVH(8RK~h$QHq#SD9Sx!W|y&zb;fa0LLN<%;ij65(!!O{f(xFgCQ+S>e3&(1 zkhWl*6x@V`0&Yw)ij6^RT(1jYBwl%JG!jX?onK)cmPA%jB6Mks5qcWZm!z2 ztMFDS+I&N7F$%qnGEx%XEx)%d*lR zvN{JF?!yaeCwqR3T=@En$utl4tJjSXUu<>Oh-WE}GFgEGaAqTw`cl;t^l8I1cC9vc zhGejk$2@=5d`pw|FKB#2!3$Jsm;-zKB0opmQyN|wDIUqr+j#F)QXl~e(bsCOsD~Z>3u6M9A95~-@fKlt-P1Emf05*K^kSzE?zqC`Uu=EyKhHv@K@OyiE z$uDLY4dy`Kh%aP-iIBaEH7k{;6Y{5%zC`-k_pF6fbGd{v#aQ+8=Ey2HjG{KJhx`{g zi*!Gd`wcczf*ks%6e;yuu-DE$_c|Ddeq|iDuN|#GOORbH$DvWeb!VrQzrUS=w7IMy z`exTL>dQ1+ts5B_9CyNn_|Qr&RX+f-H`tJ_MDvLt^sUSvt_){RqGiZmq}ms{y097Cg-8l7^!8N^1D9gKrmeOt%=+{8O7WqCIpC|NP{!kL)QV_IYgE8_@Cg^CgT)+1r zPHL_1cw->NDDcce4Z@I9@VMC5C-3I`0i(-31MzL&iudL5N@tLuH2)A_*vQYUxh&Ai zUY8BZUiYi2cCvEgJa9TQ{!C&kbkVL*I(=nyGw1QV_~}OcX=~|WVQJK_T7%z~Kp7}w z&(tPz?XUR|v*(nJr0g_K0g-x;4148}%&8Xl7Td3jQ9`*2b&dWt{r=HW(D4tXtHXu4 zxj7&s#Qz4WZNzCADRFdpFgIuiZay#TPK|FnPE)olS#@RS0fS2e<1cMWmbe$cEiUem zYUW)H<5dh;9V>>jhhu&=;!8EW&A*JfoOv02OB>$&4bDP`waI&NxS>a&3^-mL(KNa;2MmTIcu^wk<CVhzin0#zIbwJz?R7vUFJ zVUlyRD=y^u5zF{g|xVLA?w4r(=tVQ)b>aKk&xY zn=@U}vi7(i(eGwX;ygoH!f2~g@Lel)t*y3R=I&p|1nfS`7pb&$dMSA$JI&v!@F)IV{+$0Mp& zWBgr4JJBdUx@1DqadB@VzxuDgvklG!IXHYmO_Oinsj<{C+H;w|YSg^F5cP{dKi(Y5 znC2$lqpcT19jD;JbhEB`E$aqaN#7kKJ}z32TFX+gB;Yg3#I=7gtelJJ)Ggf@JCH~r zs;hOw6D$r8^fibPKeR>-_?!0+u2Vb<2UjpxpcIgIadC=%Zi393^lrGb^=8AMn&JjJzC- z@{%^?_7o$aM#te83hY|Fx2iJuMr4warAa&><2NSwX3$;hP7bFFKQKWs5@>3p)tRDc z$5*s|Lg$1x(+lo-c<5AdBF+DJmExCYY$V2(yoT5EwRdf8W9Q#4iMF&d~NzFP_vH4stAW{pj5Tx%0OAkMW0H&5gmt-HvhP$ zMLM;A3c=SJPc55TyQ)0`ur!@j)X9s5GU(iejLi7KFe<&C_X^Vyx@j+~seKiT3Oj!X znXou9mGGHTM1xFmG^|^-L3NT>2svB(34WeNRJ5)A66-E)pF%l7Bv!malzeK#g)hUI z?<_wcId9B&XmsN`+>eX&VIb%6!@A%7^+}D@$h3nPNzD6Rcbo#f^#jiat>}Q+#_}F> z!%8yJ&VX z{_-u93pXB>-1m=}jGTUqd_;-XMtS)Bi2%woRJ^@tIxflB4Svhe%4MvFZle(s_LguJ z7MIwil?v7E*&N(mN0MDahqDQw0UDP0PV(~YGlj%{si#hDOj z$SX-VX;aoEt~M+csf|9|g}SzK?jKO756ty~qLaE;wLcupfp&O!10~*EnO=0G5E>1# zIA)?YDad{DF<%vNB+h#YFg*BK`%jArBbNk1Uc^-X)fr5%wd-_Q3z_B2>ibP5x-7Q3 zO`LVXjIW99>VDb|*oNQfYZsE%{pfD2^92!Mr0Ibu#|m2_hQ|DocCv=Ea>yr8! zUMjC_T8-=`p56K@%KHPV&Kl<&_Z#?<{3S1J!_-TITY|PS`2a)s1JdHmyPz9b{Y>Bu zAe;AeMxy+aTsFQ3ODpK~GSfnIzujlYD;+%G?Yp_pIcw96@Fv4vckz46tn;g;uy&@W z&|mqlEOtGV@vlbBrVG>PsC#Fl$AtDh-fs@fpt1T-5Pz4rnmj@8=TYb#z%MzTu{Y=v zl^^iUR3sbTvUc|mHEGM9eTdD5eeY z5aXiCUuQ zf2G72@EXDFZ-~739CIxZCn5Zq;1uVG+EbwkQURYr-_C{3`pUj33H4`Wt^6t|?>8ad zAt3mZ{pUk2_;I@v^Jr;STwvM>$^q5-UYBn-{S`Q|LQbzo_e%j>YgHu6bjoMPC+}`& zIA);H1qx}Ez3Y6JbGi@SWKEAt4fGj^Qx}raMVCm@qo$)ST*`i9P#)@5uB|%Z&rK=*7p(s84%{XZ6%8IlHIY#wiTfzD$(v_MRJDyT2JAs(X(}3 zjEjx4gu#h~O`aZWCt&chzf@}YqQ#G`tF*hzJ76Ld86S%)o%!K>dqn5aWtxjuRP;nr?#0e zT$5XWwJBA>s7BL4lz>FmZLd{lgIIA0`VOpWHv!@Y4pjqI0+)S?7i}!f&IxO=-lJwA*$E(!_=Y`eBz`@0HdYXb9!?TNEYyfCFB zrCXZ%hGo~#ebF@uWEKU6f%XMUT6O1+-Ge@_8&{?xMS5h&KXXv(R<>-_XFx5at%Sr426?PfJUUPZGR=1wc9vkJN^u8SC9vEo$bBR~1T!?crS z@`$^H-CSyG^~nHbA?(*fIck?1QuwG$0OA_Hk^hgW)Ahc@r~9YJSYfZjuj;ETrnb>< zJ5C8`PL%Gao($CaPq!%EltqfW`4Uf&1li}1S`nd zbudR_v`1UB61!plQftk(s+3LZrn@AdIKbJ}bsj@>yw9p+K9Jgd^!tLI=_$zJ7=5AZ zvG<_uA+UHUXV}VxCi!Xmm1abCm&d8)c=pgVLd$xI%XLaN?%9b1z^jXX5c~x!56op*GS1zpfQpb-e2<8gM)@3n&4w?DU(T7dL#OT}dym!rr(q$(dbYJo zpkO|P6OJ1#`#UrtBq^D|LNDD;p}Uc>7jM(fk}z2%lxhv_I?vc=ov~Il_B=*wogV?- zd4{OCa?RQWRX(`QH($~x|AvU#Y36#b`pK55$1ui3pR4XotC9h099QM$SDuirz~S&N z6-4`wn`md+MgbS68w14TK|h_y?eTJ9+!WpNu!_)Lchnob{zu7<JWx=xW8Hf(R(ha9zAR-Y?UyoVty4S%_e#}IcqG{>J5xw9Flp&=oJu}dTSXD52xTC z2(A`GW^0-2S|@duO`4R&K{H$i%w~2n_@v~AYavWIN~6i!IsSEnS}S(Vsy3CNAk9v-_ca` z@U6Ac)&<)M%w!*cK6M^wh^XN{tWg?Lz6#@%ebrF{oR(@CGrU*)6`8ft+UY1VAhqF| zhs{q*xy7iJ#x7`X2;!S&8kU^sjL6#kM0>Ms{HD9lyO|LWv{5_H(Yl(-C%x;DOgF+| zY4&|f!;=JH+1D?eW4)8q$|9g{c-T)zkD~fWr&JTgt-NJg%I0xhXj|SxMQ=lNBO6CHbJfL#gs477QHC~3VVU;YJIIZ}ee^H;g<+BuKC-^!9gOGucrWN@ zT6H@{S1ETjDrVuUZ^3S-+LqRi~XPo$^B-fXvR|%k?TrsY_DG7sC(d8ePk}dFJyIqb zUOqbMs65_v$INr~=B8rwMQEJ#Nj8-Awke?Y=>){0EC%)TZ1okMh<6vK9DHe|KnGcE z@`g0fY5&@Q+i>5!VR|W0vt;k|)yc*#cZyEsDVqh63MD#MUGkXrnNiZI(?1^Qh4qZ8 z@JW~clp(h#!+5@Vt0tt8?${Nr;EF<$y0J+5?Ry5(q|BqSt+|&|L4sTJ<@ahq!RaRp zP~_JeOz{5t`nqwAts~u@eQl}oSwD*bpgj3(#@lJ(8>N!AiBv_AdUcZ4N4FAxaQYzL z*^G7D+t^WUupVd`mZ-(mK1gOLKZHLu%@6kwVe>0ze#7X9B&;dnx*na7@wVKQ^L5k) z-{pRfh}Lw4U3^XoS-crvp>o#)cd(17Q9aisbNHA{jp>ePW7Wv@l)#5!%F0{pEw0L`Dqzua;5w1W3P9U0 z&^CAAsXDu^?#hG8`Yz)0xNqomU^bT#%D{4LRZ8H9@`(I6az1$6v6m+S@3%M$EM7BvRZTv2pBkbnKG1FFzX8lc-qCrL?t- zt!>C$Fbp4IjjCgj)z`GPUN>eM@+{bU@#)r$?wm5CaAm{67L#z`$|D6gybuN4psXES z{lx42QBJM5Dv32 zlPMn!pO~+qhqwMGF&_@K2SYnG$`|giH#}SRJ-l*zY#9lHn?(HZ-G%}aS(SU+U|rfW zjdy-r8KL8I)?e|!>p%gccez7QD!s!m#fmmLHz_X29jWrA1q|)p`7MrD&K&J8)CsCE zG5COZ4l(@;O)AM%cSN*uNj`PbE7CSJhLdX)j^%T-i^4)hY|+L@*(zHMLPa4mS!*>u zRWJzUs`0vsxw4B?_diWAory3xE2o8(*m0)#!ys@*b9?Ji49)j5?2NyU z-HG3|c!`xLkETy#;w@72nZL}Ww8*6_*N>Xbp&KXnA^)5?S#OSSz>G2KMp*LHAZv_A zqLJluN~KyC-NgDhnaki;y>8bLjt_5+(4zb!vD7@UIjuwfg{%+1Wlv5T?{B4lFNZ?2 z=cFh$6S&v6SnoK&AH6|0fiL?m-o=MByt~^sk)L>I(zMC@F~50!zwTBxB$L^P`RqLe zbU}g@VpFGD>7KDszjflzol~i;FbLigvRsY}j-Z;*Wst8JF5k+mG5m&iW8YIW=PVeF zw9RYv0lzzQt>WcyY;~GXrTR#dIEr*LW4gCX#WyPw6`2RuEOtH{crTq|;#B_>gkrn1 z*5*|ete|w498N2jfhHbiX!on)QK*X8&H6Kvtvz0!S(Al`4{39|Uz>9QT=B=!L%@O! z&<=B6bc6xi)(Dk)t*J9#-kf~85&ZG#&xWyvGoWs4X|FRHN0h+pb*x)HbMvO|^k$f^ zir8S0G0W;8Dg?$v<;Ad1+)&#?>_X_ZFy>st^caQCgy~A1MoSbWe@6eOipbwQtD+aV zy3&8v1;A^$TN;#9=YE{a&iC%X88~R?Fuj+ zGD0=jL*36$T-9`Yz05&ap)R|w?m}E_sUVP&)~fVTdqb~Vohb)9-YE90Gn8l9+pqM% ztYXk{4!;r^M?qaV7AwV}Es%%yFOzHxduEbE9F~8p2^vto>+&jkS(Nk?Lh|&aa=+=% zM55gvYBkoB03&w?IjPu~hg5M}tPAH3NZj?nz72-<={{i!tEOTbr%Q3*nvJ4nd&@ix zI~^xg2#K>rCkIms*?01AFUxt{txg~}JekAzlPwun<@Ut%ViZwF(5Yoi?3NGj7pH<~ zWdJ^F%)RN|2C|cYXbk!o)VF5nWX^W>X7FZ}x7e51pr$rk%MPShzy6IADapo1J+$nt zH}70|NDrmyeymYo%x%_Sv4SJ!a7Q&b@~j70uvTYUlMn^b{t%_7;k96d+3b%E9wr=; z{XL5_UTjF}aQFc9dPiZmKISsjW6V{e3a9-GJ%=7$@VWF&!P!mp2iaPWsIrNME(c=4=a=Ucogr;gr2EJ9954N$&(^)mf7ywProyi+s7vc^>Q)!xj2 zSlJ-OR}5975mE4yLmJGT?TDLsGJ2k0Uo7ephUzJ__vV0K?B6CLz;kw_&C=xEU!lbD zkw#$7*qaYd;M-`kh45dSHU6t|Y;h6YHE_Ew1T38YK~(q4F-()t(2U@_+%KqDWeI zl8BD0luvT9Vx($TReYy|B9E{Ivq$-)-Aa5}G``&!kEEiJ`_Oy S)DV-pB>9rAiH zcS9+G%eyZ%K%N&A(YK#4H*ENDHd;FB*|7^z!BJ^?L@*`fqW7CX_O|D280XH0h&!)1 z2}f9Gr_`nWRHrk*kbf;mvI`Su!+eKFtq?#9xi4Q=tL`ytD$Rh02)ugzQ@*%MHzQTr z=L$RRBxdV$>w|Wm&IV4tIU&H}y5Zj%A^MCpmtG`G|E6d6sVEc@^-Si5yP{~VCGpDT zpM`84c1zNBO#^?~`uooZX(_Uo&A6ti2zUu^8}R5^sbphrskj7b60^rT&I|k9ZghM5 zk8M{!H2#3r-E5B5Y+0AD5z}vXkI&F#lGW*V*%NIt&GOMrcoJ*(0XL!w3Tk$|0Vpq> zi$6IPp*D1|Z2|qk1cnCdxV**4l(KX1Yr1WPdTscI zTR0|%t!h$=M`B|zl^sdBM2h+d-d}wuilwjR&Kh~YKmuYbPCs`*3swHi{>hF+I_4FJ zzz5R1Sp|K^G{R2ipwVU8&gF&|bg5iBK=_Ycd|9pudTok#DgVaV_c@x}o4J`w27iPL94__<;k8hqA4arAr6CW5R3w+)mT^uWYn=@tUg<>^7DXK9)mYmtV=Ot9d2|g+J1K(KXSR zc{q)5$-+y;-12+czzi}+31c6ze#%{qLAqV9Kq!bI_wbsJbmHofAO z- zCN+P`g@l?vu9=Zx6CbCOlc?3v!16Yk<`Lkf1T z&qN$=o1moo7ngin_^BrtWnkZ$MFKO;EEVRJpSD^ownKV_ngfF0=DwM^Y;Y|Sl4TFP zx7SML>1FSEPMr!a7Qgbl62E`-)YYvTlYe-GE|cgYrmLhq0Gw%5-*Luk5T=004iMgz zHo*(%03a}=;VFYSB=#=rVHX)3o?}*vH~F(0mH3d!ZUS$B)>8AF zrh`34u}rsI8+blcP%t8Pw;wXs@TzqUi@OX%8~FTFQ}fa_)qdL5u5hu3X4ZO;s5bQF z-gL;2VfVaCOKHQnY^bNS=onl^+vha4vIO`nDhnTYEE}8bEYTCu8u>&E32gfi3Vi&mh{}kSxpZNu>^HF6-bs- z9wkO-;B^`wAi0fRZqm78I?Y-qFE6F8f8w>&4_^zCBE<>*e>+~|qb=2Y2f zJ5w=r@^lW&xnI|Y+(N$@f_UJ;B{%$xC}(H~+vM29U}%z(VApwq;h`%QjGsh6$0 zJ0JKBZX$WDzC;pJm%HrU%1pZG;{EZP*p#;VclwTu>QgJzK;kmC0IgBs?h-#~Y zQxpyP&kpj~N60)$@aMvdY5!As&!GH5{ckXNBuzm=OwF;{ZfT*p;2gtSS?y6jRvz}v zW)Xy4myG+Bcmh$DFz;%fY|q zYSTdq|#u`wn>dtg|0EvRy}{5P)s=E z5>q87rveXjro>;n5j@nDp=zONq|L@58EeO`#rt*Ep7e^Q9FLl1pd9Y=baR#CRmQ^I zP=D6&vZL83fJMLii$Y*JbUG3|oIY5q<$S>Yn5Av-b*`ee6Ul?|EGFhk{QCG2d`G@d zPKJho+3fei6G7bX@|(6z<=udXg)V4#${aYPld%IxbvTKlKrdjAX2t!q14jYOBKh3R+qC)Z@z=0E7^Guo(zjR z1lRN9_3P~(l&!)Wth|qydjhn?mp+4Fg4#}EUgl4{sUIX#mvCgxMh9xTyLt4#J*|xd z_6amC5r~g?0LT~mk7CH@TY~N;qxJb!b{Bq~mFvcOkZ*B6@}!U=rYS$2+QE6G+01BP ze`H%tH&MP!9^<_Q9lz+g{AlVgA6ngRCu%?xRp31FP;ZmcTcQCc|Cu4Bqd}pq|Kfkx zC*v8SIbomZ{xwA48%7%{$wcGUZYJ+NhI?h+%^CapAb`2G9*(tj*+nSalZby5;@$VO zyp&CQe*jCiSU{v9mKi%r{b@d8l_%B9&T-9RH_uD#Bsr6HDqNpWK5o8 z?oB1{igg)k-+NndzTISq^|PBX8=Ut3ocG)0(R+ zVXsuIk0$G9oEADU`VCG{*{$=_`xPlIbSJ|JA)&}K=zg$e{|%R;#-<`WCs_3*lz3<2 z4)>;?S(1nbVG((b@xA5U?!oDf!#c#Xzn;}m^3O2P9J4-a4yTo4i9e^$CDfJ62h8(w z-8L$h4W`;GSJiy)>F$y@swLTtzh1YF)k~@&60PS%!p|tM zJy0%o1%-c~ z=OhROmg9s^bHcKKH}F=ii`5QQvOm4}KI1%?ejBN*kjS1GEra$CT8M`{FM%T;tTg}T z;_K<;iR^{YlU%38+Jj!h(IEh6U(*&z)*gxnoyBG5EI~r@uu(!|Y`r?^KI3$?N6@{! zh?hyKBPx6b7GBK@D&h|~n2Nv|9Mb7b%Ld1VbjMkFoI=3+cR%G8F^B~|5phPnk9`n3 z0sDn-kM*$(4iPY+@GS1+x=2Zd!T8lx;k~BA#cEzxPv!Y!!+kh!SSfzJ{350!@P3)P zzs4aW&(rVV?&=6~ze#!4uU!>H3}QLOT~TPfI_W(Rtc|>j-U0<9HgVy13y}7@Si@0U z^=HyP?x+Ms(zDhB&2%k}e=hSdF^v%k)sT4nbdlpCBI@O6G3bi11i)Sb2w|kD%TFyA z;H`)gbDdebOgrUt;F!21HiMNiBHiC&tlN;h6k@t^_m#P!-x^^Q2lcSnF3_Rg`xEdv zSaXME@H?q(+6(0SCX|%Q>d}35?;zG|M6;nmL-cbo|C(QkdIq2DaLFeTh!bddugUwA z{JRzt7R4X)ca{jYkLu5Od`OEA_^}el=>E>3sZ_JxBf}`}6o35bl%KtAx!M zFB}}^%>S^_a!U^aUpJ>z@~r_M^fuOF#oG)O0g3=Imh8;R;j>xiUt%Rk;>O}UNapy_ z(wv{8`vtO!dM)@mTNMp>l7SxA)vIbI?z#@@tXgO6kn8(#6*B*uUU9BHqx(~j5>S{d@yJ ziJFH|B=U#-nw)5txTMN0pOd?#r?%*%n`7gJi1JatBs32pqMW>ysf_N?+_RPNhCW%( z>>M|m#-StRB-&ediFrGXCLe;?Xv-A1heNWx6q7_N6m5cLKd3$5YRU&-&8)4UwdH~5 zg}HOPAs93i90^2AJihUE>G`<}?FS=y(LLR-o(0`(X@N8NiNeEspO%~+_m=KYyu2y= zu7~hPasu413NVSwH-i(w^WB50Yh1r@J!)*j=pD@Y0N~Y^j7OcT?AIJ3aNXYX%iDLF z3@SR|vNrW8q1d@!umir5Q15WR8eEO1g-CwijRmAGA++MHakm5*2$XE*wCY6&ZQNp2|Z4$Be zAJl!JdA*=>*oh9LimufgrM@BHBJ_=Vf=;C3mltN58%X2S=n51?;V@%Xn%2Wr~pw?h#hZ)1&AX^ldONZdQNm4XkOvy_Q2 zEi8Rvd~P5Q4NqYf!baicsV8~)N0&$(KDYanm($&@ip4zL)7BqTi@xJ}&6?B|NysQ* z zgMe`Dv$`ez^o$om>Tc`>x)mR`R|*~}D2k<3l(5uk+ml7)6Q#(W;`=D7UKI1Vh_pqr z2bd%c5YhI>Y0>C${|g}k8J_n^VA5KN{`Q~Zrf~A1p%g4qjFytnFRq(vik3bFZKYnI zi@<=z?y8CECIfyf!7~qlgXUR*$zr+)I5mfjsm}8P;Pe#osOscwRulQq?(lZOUQ53s zu1WM;ux@-boj`06jcr96lb&wC;*kl0h1(fC$k+$>eZ%;+M+QDQcbH;?BxI@|q?-0ZEU!wIbi2{P^S(eH2i~>jHGe+g02klK z^_lLot+Vl(UH>0gW*AgvbL-3Q}dO6GgStxAUa&Pg1^vwAkc9!oY?8DkX-NM-)xB85hXMMI@kklCErcb-hK~DuLM*+7p#aCjj zvebTlA%T#q4g4kNwyS4{b$6YJ6i&X>;^BSX%^&Lx%)9A%TpX1+*sFwy=|`>lwXXxW z4O6DxtieT7xuALzCmgeq7c5(RFhgxMww32_avRfbaOepA2Vw$#CuLGnH~s)OZF5i(mHJoq*AdDEh4}Gy81u zhw~S<Q-gsSGyJkvwzzy@D%RNZKPI4;cL0$-}AJ6uQ(3XJ%aQ#*j553jCQ#)bj!f z#fMzqT%|y65%6X+_&Yn@bW{aJTZTIU3)b($NQa~1_QD9oi)Kf}pYCVQQmmEBIA)cd zK5Teaee2R%)Qi>&c@UjoPm^o26&-9MP#tX`uusE?UPEm$!fb%LxIdQ;ZtTn-ry2SV zA;LL+m!GnG)^V`C937*$b8&0uEOm`EH#-Ge){h~UKxHS7Yu#LZ1kQ%nZl5(;vJR!v z*;WJBI`!`|D+7=F!pQ^9Rh`&nlJ3Qqo|erj40_ag53B|<2{|o&KU-L+_}9f)2w=(W zruZe7=#NDY!QxQ=3Vipqk)h)@fY9`33TkR(>4pXiucUunh0VWIvfpzmxecG~h0Y_0c^0^42Cw$9W)5PxOFJ@`!GpaFMe|~5_G#~6m)r(lOJ9@ z@ZS1>&$l`oU*QnjR-;-dY&a9r$?f5m2KK(LzBBc2nH3Zguq7qW7pzgv)Hx)Xe>*7| z^#)&Fu)vV>lWXTO_77iQAs-*T#iCkoo-Eu%O0n{3%^^=44Zh2X{0Pzs4T(=wHL8%V zsBpy5d^eVXpK9~Wy*bXe2D+?M12^tspBq5(EuOz-2TwE?zOmRzuW|TWG-&H?-xmn? z?)!*|k-6*r9MUX^Is2}YxoCWJdCzz2+oYErcy@6;Gk#%?nkXZ-q!wXP z@+XY4tO2;r9j6zp0vl0m?!fl1PdZ^cS)`Bn=;-^rvEfnaFIvizPXFIRXzAy48=#Se zZ!^_9ww$2O6jcHDrJ`>SzkU9bj zod6M~Lh|P?@m^wkzn5G%$Rf#o!4Kr*Q>g&co`pO}ofCBD`yqa~+wep$)CS3B@We>h za*~glRpSkxVM(G>({Q4PMK{cC;?j-BwQiJ4ySS->G?#K!;w_H=>S5tO5 zE>X=sT3)9q@nVc)#fD zyZ|w?d~_g@9|C2e7A&NJgttv&ouJ-w~|0SKEZn zboeE|(y^sG$Fl|(f6I~2KE#o)x$w>+dljz)atnf_2YSC-nR{Tfn*hL9-Ma^~KCORu z+v{1Lbh%r@jBVKv7_sb3h1m_{oI9Vrka5tpcQxRze0iw3DMmLDD3-Py2AYc8iogLf59of)=TlKi1hL>WkSv8A% zGy-pMo(>oxsdec5i~e6cXO@CsctF3E&NU~t-0oEt5dn#{9w1LmjRNVRCvEI(<%?)v zr>oIP5$-I3W{CaU6*amE61-(!B-pyCk~2K{uxcWArfIpo2B&cN-LUS)adHY_l+0)0 zw=)iTn1ehNZz*W30#^t;g!HK5^m`i)>5jT0`pkD6>O1H@w-=fgKzoS>G0YHUkcmC; ztlS6eTyBKnLg+0C;LAk3k6{Mr8jc*9?B=C$GX$;%-j^y^YUr$G+J)JoTil&8x_Af_ zuoYNLG4H|kxT_6Ih4Xf#$t4(fE<=D#Pmt;Zd`NL#UPrsS4xHGuB26uONcmK(=pS|Cc@4v1&>a*#)t9sR1tEwZ|F$-A&&h1%SsKTg0 ziMO0~OLS~Xh1}vxt2TuCiLL9cxR2A+ap#YZ1`qo0TLn(QapO5kG#mnNrN+>*$&??^uFfsc3>iye6vUQ zs6j-C*?D$3Um+41 zzaNj&ls`Q>jzlZ~@OZ*%U<#cM)1^w`XV(Va0_2mqsTv@tMkK+!H@mgP)^(_ zSs#{*_;}JG8=6NB#11S3%G}Kwk1S6aiwv;?lOTDwR@=bTzHrBy(KVYh?GiTin-=WI zkSg~Yuv4YHg2y3X}(gzZ`+jHtj#a#x)+S#J|i6 zkA^Hzs>!5orj6J@ibIa7w(Ga3dn1l6X$D5N_5yf#@+K!>zMr4-PgImU$DBj0Z0Tx2 zKBi9e+eNUz%k@?Ic6CSp@y|jB$7g3~i<8ox&D=%s^MQ`%?GD`ryvS}=_h)~D8Jskq z#T3u&HGEbUcf`R2Mbs_&lAqZ4zXGvqX3up$^|z0#@=_kt7+^(V+{j^X)RHl*D_8(STHvC=Yt2Z(M7y~S)5Y( zTrG5t9~(WJ*z_`m4o~DUvOEyxOyR5+IN~Y_SaED;Pw?=MOS2g8%x{zWGQb8@<(nLW z=c1;;B)bB8!=ShdEh*VlZq5+MDpHmY2HPwkIiI|M#NfA1P*>hZJE}k%ZwuIPVaLN4 zrtKXow(!`>=+q#U{1e&SzQYA}pI38*;?bFt`iQUp54;)EN$y*YAee*1xEfmnkG5oC zRQD+}r&^5%mktLLvu4dQBq3)6X!B0#k&!kq8qHm&oomlKiZOg^*q)q~A@I0&c@BK; zz+hJ(CY`ABWUF>IUQ#LxA9744mFlmTq%=rq?FRlcQYoB6KK zX$2^cebt7t@_5=>ba$FCpKt2_`k8#nfqke-w7=bqZgatN|60GE=)KXca;(aObefZ7 zxUT(6zv5Hh^lV#hkA}cwp4a(W4tkhfWMmW0s$g+uU)SoTVQ-&=>tlCnGsgiF37Koq z{jo9prkDM7a@;J&Sk;C7YWMw>C4u{|$f*!h%MLRZ=bO5S3q*ezzpiSjtQw`y@`TTHSvXst>OP*hf$j zf{oJrL)?82YgN&Zk0oFJZoXOhMqJ&QP}kF-V71tLmyEbg#;6ae=fG*|@D2Tyoxo;+ zE*$ylqxO7wJ@+Z1UuTXb zYnC9Ym5Rl_Khp0P*<%WFC%2|lXWMqA_XwxP2-2+I8&lD%5&b4t_!!{F0%*xEet!^n zdmbS49!YLO!9)(ejK}Y+3oxBszqM^}G4W`Eg2q%IGP@r45j6E7D)JBv7iUan2oPTGX#AqT>A;370dm?~!Q9?LxWD_ITk3EvKxBL|HdJAqgZ?El>~FVLQDG!o zl@{i1En>=>MB12(mb-NzlhmlW=V}K@l!7MiT+7CeAN<9NE0p&ceG%sfX1J2ss1>Qh!_=oNatwK5aTRIg^NFr=^KmOP#;n0<|ys{1_SOM zDbI3m%Hb02?Q)>iQ^nMMKk52x^yT|~dd)w?%}d}1oP8G}vij?w>5TUIHbt=f08eet zmc`0kAmZm!d=U-L4Ye5@p$1A+rPdkQzB$orlT0}zBz<`!e2c1SJYb-ycEMVvK!LT` zSxw^*!u(~Tw)j5GwQx4+)5r|?9-$-P86tn!wTVn=fEd*sJm<=v>t&f7LVkn8t{q0a zz$i4eXNuDNiWtoRd|qBB810Fj0=8hX@9gm()&@x$F~Px=LktBIC&dr$HmMOvj?e2| z_x@s|!3_PlKjXPbyTZb`Sff}4TrXCw?fav0PqXD~JJ~?_?9`zjB|}_MoC`O=HFz6{%{w^)A4%F709Z zn(>_-%A)FFBd;ymBM*}rC z=JtH-L~W6dw$(F{W zL@TPTQPHg$oMW-KfX5Lob_?f0@&2%Lyld(D&INmQ!$H57l}cNxdE8_ z&%ye1a}Ne#k1y{hZHw{vgHC(c9=%!)pHO&8zC%b#8kUOVnHdw@ir-K3Mx=~9iSn_? z4vHzp(o{=elzvx5X>ABlC^uDRh z&FWg|F4Bk19M{MVp7-`8mx4^$`};bcuYHfh`=@uC;)vca-Y+MG0d~;Z zmL$ATNGpKe(s~WrF$|IT2knlBJ_eB)1G^tvT>ctp_Y=3X=AWiLp03=u$A&8)i%9L0;2MD&ZK<-V#d<&Z4;amc`~c0b#2eeF<|==!SP+B%Jc!9g*%ZLA=WTH8Fbi$Cr5JpUD^073Wm}>2r|cM*SyrL!o)P^ zbUfksL(G_?(H&0PBtj5nkDQF<#vEBl-jHz|j(Q`ucCINlN#CG7c%rKuQ$$OSYF{*q z>efFsUeu+ao zug|Lzp6HT^scB<$N%`2n=_GS&HGxky?~u2zvmjvLPd*&XO63KHTXzRvucZsgUPC$u z)jc^)ISRmgH7f~VBn=w|6p*&M$93NKVr16kyCd zkiglQAp^z0W^nrUzEQ0lommJsj7ylH>!-O%q~P%^mt*Th^5elLrpwa1fm{ZTv`IJs zN89UV`7)w)`2f78g+|4uG@|vY+19}^D6>1TSxVB&7JoDz&kBv$H6(O0q1G^2wp(1_ zjaV}Zc9xN8MnXbSLCmzIqun%%U&E_DG;|sbs;1{PxO6rSdcbTVK9~KEI_D+HeEF(1 zzEzL&s}JWr_-PV9J0rD+{TU>uIoR(o*6DcccsuqEe79lOuTK}u%OD#3ejs$Z*uEw4g?KSx zOs93m>Y|PMPZO)O*Em{+1#9g5jhJLmN%L}|NLb=N!w4!$3S9Ay^WD;1!#-57cPMXZ zmhCr1%pS)uYb0U5MH)V}I=745jf(5tg5cMoS2rvGnmG5__QzkJ&?c_Sb<=_C86h{M zdDivlrxKE9p&rQ80LOhWKvodA^+$g>ct1`BqOldZD_E1^I^W z8Z}B<(r~^q+ryYK0$HWm^z?rK{QycBErr}Djr_ql`BN4M`IApb^TG_mfx z!sWn3mCHXT0t+03T*N59ekc(VPvW*+$1@*j4|A@?u~lkCV#I7KL^~7cR}lKe z74cU6(41p}^Tpu#k-xn_hu=$xE`Rb3g5vSuav&sFkIcEjfEPf^u!87SZJl6@vE6$6 zoK1f|C(+)O2U~Fcu;_rbH_@!wE(_zuU9qcy3lYXULaJFySro5;mLkzf{`BF?PnUBp z4Q%8*ZA>AXC~1~paILfSI8iKxHy{cadAf}DeyGWMJ+XP)v$^muhrArg`>uIV{ZpId zZbUQd{civL1H<1;%TG#Fot#t7%Z9Q6tZOFptMtzDk~)|ee>S?GxJV1ICgW7gp3a`M zPiGcS#_#}GYXFWo0y{r9zB>XigJphHQ5i1JG8~v=JsDwhaah^TA?Z9H4*-Ul_|Qbd z*Sd9@~mw_NvoE9UFP(Emv z!Xa}XLDxE~C+P}H7Txiu@3H@?q{Z)j_2SaOKjPr2Oc`9m`?_*fF@)CcLy9sgblXrK5T73l7y# zemeq$e>GM9U;xv{oWOnbyH&7Y_<;pq4sgjyuc$rb|uX&V_R}#xk@2EL<2vyPiQDW_RJMQ3$HcIDhE}FBRO#CCq`e z@l7vA;wd)M6*_1$ToQHhdfQEs&*Cr7&^bH3Oql3$Ay^51rY}MFyJ!}}Jqt;Kn8YC` zW~PHsb^Jr$pwVabWf4_76@xU;7 ze1lb~)2~oV(gdtuR>eyR(fqXs>1$08J#eAayY`p0{u?8nA9J@ZE&HbyAvkuZ@f7EGNg7#lsM5ze4D@Fd3wCt8F)=ca<{LjBzEGVxE=|qyyGdsc za6Ih%2EkRtl!OXRtJMwtCrfoAxr#$Xytb3tvXA~<@kK=VBm~KW&`oc$;;bjbOX--M z*`2l{4g*iZjE0=~Jp|RdIqc4n9*=^wiXo~F;P{h(3;FutxTa8EMduv=!VtfHFmAvjK7igvJPj6k=Gz z`Bhi1jkv>M^Grb#8=Oi7aw1>e++56`QF0n z{Ez`XHFrfJPV37jow%^fG=Dv3Mt@|Ad$=7Jggp=9a!M^XZ+Jls30Wzxq9cAnu)y#+ z8@k<{b-VB6sc=zI1bU1loVAi3zo=5&afg&K{|D^3O8@;_*+iC5| zh=?eQoD|P?Ici0?@g?KWFe#&EYT3;9!)%pAB59+%A~nJ=!BC?hD{8QCDd zw%t62???2KW6`kDCx}BW6$|Yc{+Zmx*F9lLW_AjrLPIWSUOiZ2E1^@TGK zOq&VCE%7tErs-Mpx);L!>&vvxi(j(#Nd~>>=VJe*S^Y@1dco(12QK8ME!L}OI(Rzh z8M6ZEHAaViL$hwj13nwIKm+!=7hCl8u&Ho|t3)-+g;S9~3I~M~+XoaKDPm8gLeyFs z9i-+b_-Mziu|0Vp;BM4rG;ENH z>%;auwNQ5C@Bh)9&6;UgIE)G>$dE15YZas{$Sry*RQy{BHJzc}Bpi0le%e(I;)qzt zmHeY~)JX|1hMJS%?pzLG|DMp`@oU%(V=L8XeL?TKgn7W(z;eYRE^17ZuNK!JiZulS zLl_&}_dzCRy35j4t*zQWvDu=e#LiK>VrriElRZ!Z!1UAck^rDlZ8?Wu^=6PRGXd|@AIM%Ag^oYC$pgBhIwhYHRC zufz}*4|a*qN3=$5Js#?leMrmH#anz1NI+DC-5sTIe90x+Hucua7GTF3h4$F`d-SsG zbt_r|XQRzF>gYp;+kh|fsMbfbkZxzbNlI_~+ffr3lGZ9A(J1Jgk_?BTP4>-wc8(IV zkxu~jGhWi1O;6!x`LtK-$gi;KwEqGtn~;wSp95kj2>V~~Ed;hEWQSKN-}vU)r@eyo zPUA#FC`3&|Xc@_kxceC{UgS0qZ_1&vA=nVcEI?84eL_0hyRuPV5$Qce#WY(LwPRES zoJj=!xsS+yaL&WtdFZC81$RUuLg4xJJT_rm1$Rj};nJBkWz4u{QI>Fju3Ph`8nV0nZs>lsKykVYiAb+Aff~Mo%BXe!#^VWFL-z4SK4pkT zd!D2XxSYHE0HW z4U+$EGa(H4vhNJH9sSj$-7yw2hB3sX8MmE~BKN@4M z>7|0vE~k}aM*VSonIy$A|LbbFLh`Bi(*xAu+9dbbMY7l3*4E(lbP^^>yi-8i;HhKu zx1S8bS06%VhlewyMfr5GkYS%7pQA{+kVMdrEoPY*ir6O3Ks^)#H}*_2-fPNnuG|$n z;!Qw=a(aGz8y6p0AD>!!{LuZwAnJ=p7VtJ;c-a+J80TwY-=>XRKX(i19skqK!$l>v zh4T7xnq5~)^yi#|M%v?o+IH$n;colYCWoX$-nflHFCtTDmtAaYHT!)ok_)zH5pVJ1 z{3$mr89FRxEAPd0?~JIx66frKm>s?}rv#^kRe-VTd9BW*v^c-l>(OP}`^Y1;eU^^O zJ6^?_+uRYIOX~!_Jql8xeJ%pc*CrbGE^lfuT52!xOs)CSFb3pwzixux3)TuzGzLmR zUY8CX>)pFp&&&V4))$C%Pv+mK|67hCgKrV#i12J>Q|_}58B)Qp@s>LfWJC`8)mY~* zlIUE)pum!svSUuv)S3YRfEbg*bh-OWFC&sN*R1;I>c zbewU;>d1MDnqGs9ZjcT^lZ_=_!UkHG9ap=k%y(lCN{H;H-WGGS8V$asg({Krk@paf z=6F5M>5nDDIR75@piT{v3%+-Q_e2cK#*c#jV8To zeaqG`VNqZ|4h_8l)_9f_EIUIcUvNf2UD4H=VXHtIW+bc zo4U2K&87zgm04$Q@~(4&Q13;GZQ{=0V2*1`w7|j09@0jp{-*2Jg+!PWnUEc-p`T3V ziNcSflTI39hJon&6Y96QTHy%8;+QXw`HqaS>*snMC*lSx4UkK4w{_2wE?MlFNzK%6 zn89yP0~U41m%^_*7U!yOOVg!$_xD%YKZ= z9RqdudDW^BRCN-lxd2Mi1?@2>cki0hF-E9;47byrb{aazqzW-MFEsQPjm8GSGQrG7W?-i!~&m3V!ZQ!hqJ)>_huyU4M^ui zA|A98lsYi6Zm-XgAu{L>FG}mU6|t>n)j!Ynz&ytZcF+3B|@bA_GN66 zG1KOX`o(z`rKEA;ZnZ%Ap+}SMMz8gBS#)+gE=@Tg?C9db%Ab*R!y!)5j5}g)GLoqM zQUCUK5Hoy~n+MlL>d|3^l2Y(aQ_c6u3bZ0Yem11XDvQqohtfM=Pntimb8rBC)#*}00a{^r_^jkFxp^qj1Lwp}<+rp)5k z!(dJiA8mY-&+QWl>dQE?VXo>i;>y5M;Z(P}zhaKip zPDUwIgXRsouJw*CvD#b3VuR2OM5GT|J_EYU^5{lDFg9Gtxv{`x^M#Tu(O21^!!>Sx zI%F40qHCv~*OxJxj>5OMewrh@izeiCPig)2;We8#uclL*JNLjO{I)v40o6I`5(iZMY!+1lJ$KM7{fxW#^&PJn_8?7kBekyg{QZ5D*E zv5}}uS&AhaaguQqrYzqz8?OX1`-PV_^shB0M-)75bbVf6m_9V2Z?K4rO+U?HBJG*_ zZPKlr21DPI|b05r;)}6bY+Niw!y1!XVX|63uLC}lv@>p{E^01^)dwm|E z7~Jr(sv;8JzBX7XP(6E0T5k0}c#iPFd4q)epWy$)fw%j9g9#G&u~d)0CAvBuf|IeM z#7~iY$ve_a#lU0cWh-Bfzi-AQ%5u_EMyzCsHnX}hwNf8(RFe}Or^}4nKRw>=*?-}9 zj>h~_-8-3tH9e)M0%*4jPe>ILtrKms@-S)_{odw+c=Zg1fER!p)gAO|W)|NR%r`NC zp=nEP4*88?k!uE@zeBBpsEXqb5lh)ddS+)!*YCwk+iHAFuJqKIJBt{M!>dQ~Pk2N% zU7BP}^p*dzIMQJ>fanm`lqTG{R-Q#`}f;oBQAH)BOxS%sq{@?$W1VX8x@ zx{mgBri}4I6Q_<)_E#8qZ{y_?Wr%quZEt{14V3(u+oA0po+<;w^+7w%nTf}}HhAD+ zU@7TM;k8Y5qf^|t5 z32sDrXIELgD7?jdAXn%h*#YbPiGJO_2t(ztzHP|0_vpp^#PA-~eMVb{Wvyf7^)$1n zazmz#q;}n+hu-lKad!!t-ShNNX}ag?=#i&9qwD-eBltg%!iV^x9(sW*O-RZ<)pkiEZeqlWk(~?xS z7G7Vx=KU5WO1@pMZzc7dmUmYF5v3m&?#qftn)|UEC0|tQeq~wt9!l9OK@I0FJ6Pfo9t~<+8Am&6s0C zN!8bHw~3-x`7@=hCBdrC_vH39AyT8cI?&^n!$~dcr527*Z}zw#v8ln181WB3o4=Lq z&sYDd|I=$e{XN5AWVrvXY`BJy@?XT-RO)bpU_cvfNGLE#;Rolzi3HH7z0oao7tIA_ zrl<>Fr@AAR#fCB(&aXA4iWw@Y>9>a}2GpkCO|?)qxRWnolmooTw3nEe>f;Q{U4N`- zVURYD0mC^>(C@_YA?*u9@r@7{e^U(+$JRIQdq%z%3_{?n*IGQqtGZXIJ;Osf?g+zLVllQud>Rrl=f99K!0> zE;q@DU(hW=H;~a}FuWjtTbytV{%lByH-z%e5C0+z&?tRJ_q~zl;U5Ic^g7k)%{HmQ zv3-rZ+IS7d?+AQlUT_#SPHqgf{>7o&;Z=JJ%LSXdy9hde0X&m)F@t#ZGYxd=BVXM| z%Gm#ZSlKO|#XXYtB}T zLmAtp+XmXVx0i{r*-ZkYyT({FQUz25m{BcVO^`hYy#u@*S2AS*qT-V2;iw3Dqt-bP zPS{C1bW&5<=$}OMX9pF|tT2f`754!}%{Cv?$CAjD?{Y;x)i9u2HDQFq9pf?TA$%%F zqer|nCy2C9G_s&4vI}1ceAWl_E}TwOX*Mj6%|8yDUT~%if`Mwh!w7D zJ59+)SP5{FC12MywLNoh=saHXoG%Nmf8k}A9k*X}h9F!!=c;DeS=y>byKi}j{wyaq z7xbEDQD-nSpIQDd_5C3oV*X|LFtJ8M{>c%90C0Y5)cs%_H)jzII6CDs68VAOsg*pO z#C%bEp-+=O;6EKyk}k>0WuN+wP)Jj+w*1Bn$Nv7uX*grJU18WUo8Oy-lHE z{3@)P#=`ghXoXowUSNx3BgIQ0H;O-Lv22zk>HAeCKvOZM30Nn{pF0gc-P9j>xV0A@ zpExzm43fT~KuL6tsy%ksRHen#gY8O@!{W5LROpXl$rpD4tRo4uDkf`j5aJN+@r)g& z?9&nf_RoqyS+=uouRm}d7)&34&eA=I?NNvwWw;whtFW(i!VgR5e|A(zwTAx>8hOKf zgyDN(r*-&0fc1h9>XUptFV9t-K-w9qaLn(Zu#Vn36MJx4*I(OKhENmwwHYr$o9eZP!37lR)MvQf!D_?Og6(wEm+4)$DpEeW zNBYTZ>s^cGrn`v@O3aJ-S|n6n%8gdm!X{NQ zG}GE1rYHW?L;%I~e8m^-IP&Xs&V;Y&dfanN4Uf6#<#o4+u9iRsS6~~Q4^hR;Dbl;v zH0U}|Gm86B1={U)-qWBC(oKp#b?tYDzZ~|Otkv32Vv8ZFH2TV`EqHj+x6KDv>CGZ> zx3m(Nm$D5Gf!;~nmEKb{?IuZ5(hBAOmh&GP7t{y(IfoLj`6pj>LWOtIhV4;JqqjJa zP+?|$yhpSn(xy3ML;$KpNkP*>v%9uTRGEjmfWFMVQ*(Y18J~lu675MEV`5#uqxQx_ zyzX&fGxM!HaZ<6P8-QezTepVr)>_7$a}DP_9yQNV;NfCevYbA!!lE{^!wQ9p(kdZN z{B5gjT$P+Ptl=W_E`>_1X?AzqjExg*^n~6xOx@GEI~koF^ke$Gsi>X$po>XKd|j_2 z;*tT8YbbT8lwN~h|qiTc{<{&k4+icM(*Ngj$Th&`K;|z}w;8!qk%Wf;h%w2cA z7uc680VOO>W{72tH^*e%A~~!vIGFcM`J1VD@nLrsM_A+Fh?}5qy`QFCku~L)Ynd?J zuv#M~!$y6P1%4$xJbVKJyQ^Avj0jAjUwOVN-Dk?gJAH} znlbPlm|}vzonT_BYpP0)eipRP0ZHJKW+-x%HH76jYHjc0^O3z|-=U4w=oCEkRSlxJ z9}mX2jHGii-)h}S|Dl(4oQ(T=xpVKcNOfIuC8Bg~e!4iYg^Ob8par-q(o&vzr@{xU9l6O&458 zz&fNm@-Ca^bx~;j5_wF!jI!MPFv%cs$O@ zXZHty+zxa7>@deIO%Yuaf;cSzN{UX5;iY^>icI)jz``!yAyfWu}DPySxu`J`su z`%z@433G6tfm7f<2iZ3^Xpq=*Ti2#d+r(@nIYL&30(~;$c&P=wR1oV>!Y!Kl{@3qf zb2;@}@?-u+7ZY!9LP|rrabwguKAx7331N$d{3b7hi+$_wuj}jY56E&YvV{xt?a&x1 zrkd#5p8}sJc_rtPpiw^1<}~JbaVSktrF?~6tjmr?yIP~FjuX#Cu1PMF$A!xha-rZ) z0TPaGud%{EqOiy4WFt^exF(czNl17+aKlTl0KFmC*RcUBk2MJF6cIQo{Fo??i?wBJ~>4~X3r_U9}MEZ)FCpgbcR0Vo4 zuukA*6EO?;WVWhuaq5A1)?L-{uobTJc9mqpkTgQCa=dHLh5yBpVjSw>T|D}ID>s^d zDC|(+ZCl_?K;YeE5pE^48u>X`&Cf@L_kBw|TCcq0^=1ABRU@2!L!|69|0^F{F2TOR zoeu`h${em7&MvVl@a5v|V{8hAXctG|azpfi+tUqKi(eXZ>0!`^F{cA~ae!J@^|I^h zJcH?f^N-u8IyFG!1?f+zRSz_b$Rh0HEMIm zeOEURFAuGn7^BXSjcN6Q$aQM-F0##@Z!O`q*yHy-M`UMYrwB1`L|BpWr}cu9k@gFTdYkaR)4gfzW>ayxn(kJmgM>v5M>zY|}HbTq|?Ap~svlh!|x4GSIz4WAM zQx=@W+}~*;hn#fcD>fBLb~RFz3av2>^j!^G`U#w$Qr<&@L!N{koLvj{(@Us{`glOa zS6DaH5NxJSyDcB@rPWYKw8jr=PxYK$%cRn~EOPxwM|*Mp1(YcssP?&JFYy09@0LnS zQP*1BGlHz==f%Q7xKtvlsq#{jEuWCj#c=VR0-c$?2?a1572`pHhxPmg`M(O=Ir`tL z0bP#8_-~2-K?eI2wtpwWLij=;jjJBmel&>ASL9igAxT^vUdCO-V7559a2WnX*C>r-)ec+@AztN^Bh;sOWa6deB# z&+XeLqQ;$DRmo)VGGOBm+;)gs-0wvauiRjV;^eL$fHAIK6H@0ySX820Yl2O({3oEqP!2S<-uk#u%r)vSYPyv1~`rEo9K zS|WU(Y2Om*9KP`z23Scx?}(_h-QO$71yFn>6YWjC4e{(d0T)8o4kJ+_VZyuTyfti> z;p<)SG^Jd$0V9w%qNpMx80=DZmQ%*7wpvxt|9F7R52p4x(}$Ji9|6L%2>M$a5hEj$ zIaQ8pSvaQBFQ@ew7fOQNX;x$g^TEkuJmo&)Z6AUl7UbApPdhs&VVUH4r|1mmZsE<) zcj(ku_S17(+{0Incwb=k`7eVtMk|5?4g=88-s1uFsEyNAs8r4tHtB5V~g^KQaTLwhi%WPq!R{u2!8;r;x_? zy-;K+wxbCN>|Ai#Uv~x}d^Um&|CE>b?|X1@3`Wbx=Y*I@Cdbp6**Q0JY7lnpM5NX0 zWM+y~HL{*U%eu14$4?dnC|#{(Npgl5gzbm7G>daT0ye%))sXyv(h8;E4?6+`6;{O6 zFQsb+?2=jCKVZkvL1k5(me$iVzn*P%H*e7CN6IsFn59SHUp?J1LHz8nyEGorcT%QL z3NL>POz@f;-9J(L)bo|PK-P7vY?KI3LF507wmg4v08gIr#DCluu%9MXu)-Zq0>0*I zf2iWl^}Am(Ds>tGWWoIt-%`~>@FugxP**eQhvDC9`8G?4^5p_v-g`l}nD3N036mV+L~d3d`g7;!(C`=n$h!(6OV z5`bjzNiZ2WQjJQH`?(>u{fiHZOmf>LJGM|^l$awxZk+Mpi(y38_b8pMN6lY(_|{j8 zYr{!}gAePfM!RI6Z#Ob^r)PA|0woSom)*O?OOx9-6W!JG#rymySKzhK=BCr=9YVBgLP zigq~dd|zGlybEjdF!}YraM1o64%)h`r2h$rM(A%g8xrY+6njcK50%{%Y<4^Wqlqsq zT6}sks*dHGvl(edpHIlraIqAbP%q_GcZIe;B#n;)WCFhQ0rhtOGYZGIVB zK#e`6nE@mrPIQ^CDpOhz6#EJ|Zl@|DHkfQn*BuCKm{zD$kUhBe)tG|A2fl7`eIdQV zy)<)pVFvn6gqJ4LQ#xUaKN@q%teI(kMB|+kl*me!9YLP&?m;z(yUpxLHV8->Z$Ic} z_PiG4V!3Enki=8EdhYM7iP**Y8!Le3PmTFjv(}TNZ3xFK25zQ)(xnp-uajzsfjkH3TJHS;IUp1bdt0CaqH4o;UT!w3M#C}JiiVD)gxbcg_ zT(4LUCu;c?0taH2S9k)rkOd63s|g=>zpu_V&Xjr;&i+q$e*PaPr(rhA;V{>RpK5>Ul>(+^JSk{)vwA9JX7`vk?UvLqPt*@mv%#1wGH`DpU|vC0T(r*eS&$>oF-(#p6=Rv` z%53kzZkDN7#)0)1OyL`S;-=^6GHFz|=24}iWuygQEg=GwFIhF=ONol@5P`NHI(VSI zCD)G;zGJGL35r;2VrE_smvj^w{uD&PDaKty_8?u?iPk|KvVw9TyJ#NJpII-miAX znRQyeCHKCILbyQXS69x{S9h_?taKxxD$KPMvfW+qV*7?L&t~YEkUQgGf-f9=n`(95 z>Lbe0!UX;tlX&W0x6$&Ym9RsbA2BBG>zy2)*{5%E`4t)rm@gBWm5@p>yu2_9H>dUL zIea$r3=*oEc}^-OR^(z@fU(UNVa|P)-l7FW$yL@=@2icES)Y)D=_D_|f5f>sE)P&n zBRrB%uB--w*Ovqy^ovarVnIDzXc7Ms$~+XC8?t_mPEF^r!n;Kg#Ze@@UZwz6~eUUaIn zRu)2Ku!Nx5FFM#0U<(VFQ{Hz zn{fQuCUQ}uHSU?g$@8=x_APA?R~sW?6nFEcU7$j!UcrjSky>&!Gf&r_HHcaiSds#Z zripJUjXv-1(>~DSnmA2y=tV$!7=W1&+cAV!3*L3oMv0BO!o0%i8aOemTVASjR$dqH zA(PH(oK0g_{+}=P9~uavn4jsm1Q$BLaU%6}4{79sbxixO; z9#ysZy2Ofe+tYfE0Wmx@sD|gf>+Ht4y3cJJ8txE%qm|T1Eq-b41+3yMJtDEm*Kpom z@q+EtYB~!lND{vK&aSa@h1AD=Bp#n`DBC+TBt`}=wg6DbM}R~!?!l@#^o$mWZxg<& zENSq;2>3$%D=u9gj#%f*Lyv@5hv*>~Mjz*RTs9drrU*awVM|{#FVNY(+0_>~ENcF_Bxj5?bOe@!GZ_|Y zY|(*9vMart!z3`YsTY2^U%VD&+!A*+7XuwPvi}RV!>)8#e~274JzDC6)Q1~m%bl%s z;&zW=<;ep%)#P~7o`*+zf*cSu;r5~*6KPQ@5^s&u zvCOfa7s>h`8H5qy4eev+-lz{2KP>J4V8nS6pNmq)v4^-jBt>X-IHujI3T`rGSI`1} zR6Blb1k{es0!}v-S=djGEf4;8*$J006LjH-?Djb1CYxne=+G-xv7-QU1XtLeRcQPC zmbzxe6U9AfLfx=!M~MHAtZxd>Y+Is@I~}uQ+jcr;$F|Kswr%5&txh_&ZKq?~wmLU^ zpMCbZ-@UK%ajlwjj8QeJN}4>Rz83D*EUb=R8v0Os_wj_)5?R+h?xfX0h}^&y`^z#b z89utAIrop-_hYXlS=pSN%nJG&QnKII4mFbzvS%2U2ow~ta+?ec--HU~MkfyL_D7zN zmpM^1g2{XFxfFGcfwZp=7k!@mZf{SMfv=gw6;hKN%Q<3POnni3tK!$w7`wQkCVz~L zj<>IzEPfy56S*gd@MM6fxJRaA7cFmZ%giiFGu1N=l)aVdz215qit-w-onK_MuCQhm zp5~^c0=hJ8WMkxXdAUqK+C!UNK3#%c&x90pfbacW5~$Z{Zp%!X>`bkEulIm3gOq3r zV?e9x$M&ar30r13OeDehYP&pO^$lXFCoTt2a_2GQnz}#yk%Q{}TK&J;-ro*)amItp}dDd231l~}EB0cZi?wX+mj!daSh%OR2v~;biy2PX! z+^6#C)SkBjEpzP_12AXp+S@yxlQzy}r_m0Cf+bu7joLGMwhPuOITR<{k+4XmI`uY{ zSp_`Z6rxD-n9?T@@M*{RPz^;|qRNaxU9l1reCdsLWR3bypJ_-^F9h1(?)woPM23 z4yk9_!Nj(t8Jp&{9peE_4c@_Rj;}YX92Z%>&wD$@&aKUUHxomxgn z4hljZO+lR@cH9WzGW0AY(Q`__ns&MmBPzYD3UkM~nNir8-O?X}#^}c;8he?s!jVl+ zE6b23GX}k3kW>M>tsWT^yAvlnZ+oH&bZw0>=_pXE<6``8GZd~DX9NTU0Qwa5>^(^R z9XFX{WUo+vQAs-$o2IeI-H8{kn$qk0i*%HeHceexN$GjE5#U;v43j-m>)O0+yw@M( zGkWvaF0b=^?ylzro5FwP(Ptj!)te>L`oAHqMGVrEClIY)hjxuyvO6o>0y84Q9}HT! zF4_^K%``S-UX0T)?YT`{*~|6SqbOcD7?P5nzudZ!#*&h`Z_5plfUiCJ(H7~= zm*{i=TlMR&^3dL^$v_U_8L5uMmB*xLWq;+r6eM1fjUp-FQ?g@2F#^ZAOH;@WzuNEk!pBkT{NET zP<-CQabMe+SKD-1`F%*1fVca=WiJ}Sn6{7p!v3iM8eXX6HQ9swSny`A0G3s2o3 zsUfjg>TT@-J;;0N+2g16sqY`l5S)LL_|~tF#q0CbYOH^P{KFopIbWBNpcEED65^(I zdPL#)4p(|@@S4?h2H%5N%jAlEPm)yH5JL(|> z(I&INOY8=6p~3XGhOR*dI+Op+aj&sp4QAF2f|_!sy?vM{#>9?xp+_8hB$SLsqlh+f zW_K34@=|Hamdj-)$1x;&&*;9=PkwtppepaihtxlEfJCn` z>yjTzspwVsQF_XH+7}=_EK%}{J?2ES8_39xhsr&ucZpl{GLO#poVeJ|1#4TnzPo{+^i13;e%Q8b#vvGN{kZevC*&@ zri1%lmGM~wa~ncivhCd7c!fWIv@HKa_&e!zsHXAX&{zE{^l6%{tN(e;Gl)U35;K>= zyu)`}p$D_QCDVe2bY>YTDM@ZK3;7Il8HJ(b^nR7y6;fJM=!VQkph z@5~sc;+_f&)$;aSoMZZNNw_z-qOZ4HyHwb*Wq7aWL{&~ckIE9wXpPY!Txo-D<%SXX zl=PYxpW8wQ$~li`gfjBBs!yv!dknFxCRD}6M85M_)pm?$(QM6!z|NWAG*CpfeC~rY zYFvfiu}Os^og>3lSI}PT@We6OlD-(UuUm*V=Om$lYI*o-RvfB(K-}KTnA6UisJ|+q zG+`=!7DwlwEwy^LZq(5@r1g-An~e3kCijKo{hZi2P<{GX&B(Jtr4r-JSO+G*LIT;X zPde~+1dYl~YXsGen_r*hEkDylX*4yXXUPEl3<)CoW(9s3_TLn|H_2xT-kZeo;U9sI z3x=dS#5&+fm3s?B$I^itH2S8Mz4C$Qhds1T zSI-XwXVoY^kGtmZY`{X7j`nY?fJY>!hDy&lbE|KcO;+pbDL3EypC}3N3wMu$AS6WN zX?~_eLZ+aNlmfT;b!-%Z6nYbjaa1#Ia0|3~CnIe%x;VK)O>+QWDU+vSLu^zW(6o!% zG?xoE+JBGTx^x_!(mD1?SbD6_>r9i0Hf~s}mcuv8Tx*l70IKQ$>|{T@nN-Ptdgaq9 zsT8z}i>#;1itmPx<^1M2-T%xotR4lxDH~jz-ZN*?Zcn|12!oJyUgTKgcLDxHE=&AF z4V%P1h5rBLnc08<6*t-DX}yJ`3b=ZP8i(@;`^C&K}AHL!kmnYXghJ%drTKIQvbURP(M0*Hy7x}oRfRVLPZ zOME03e|>#5U^iMANK9=7nKV`npRsb(!MgP5<=K6Ylz_DoeJ%2j8-vookcKubT2gbS z5yg1M#%iY)j8fPBjpn3grw_3HI@T?8T{7K+^_+sxkeQ^Rj`TzDBwGF!Ucjm)-vNBFRX5{*%Mu~o zADDm9?bx1W)xUGJTYbWWdx}hw`Uc7U+G6{G1&!}WG7)9XxZfEG)@>K=O<`i?847jO zE63m}-rD*{e&P)&Rcb8(px$hnFFY*LogiRxv#MtNSF?R zq#e>bDx0VYKm*C8&uyX_l(vUpHt=+|0|{m|b3xwG@-k0UwNEKDZktY-#D`%|krzGlZEN2_ZFO`19A>PH-ew{mWjz${7J4;K7Lqry!mV@tOHjTs z6_s9f`)iY5hP=U~iN_fIoM0I-!k5CY#UEEJv(XbZ4l;| za%jGt8tbs%a{?#C&KbqWE&>gde%$N)I*|gbb(jrys z-m5`7H2OjVv}oO1Hsi1lVa;2rfu<-#OX!V4*UF;_?6C?!sVtXu8o&kFCVEc>b+h3V zd;Z<$4ai2Gh-S7$!1n^((WIGFEw3G%HtB{Al@inMlhd#*%IA>1tV z6;nf13pX5DVfj~T&-#ImpH-=2VOtMyAUCG}IJ3`jz&|5Imi*y)$#u3b#6nMWzLR&l z>y{&j2(zd08aHOt@k|pVj84O!#IUpB;z@iAoFegt|P%2U(auytZEHJIV zVWGUtPT3<-osf7O9auj+Soak<+m3Y*HKpfEs3Ut~(RNRdF7HXCyVuRzo1GJ921dsD zz(heu&fA^%YTiNRrCTNJTrB^eqtEFc$4FOn`VG%a3bhiz*Q)|823fgmFBOaLH470c z@qhxPo#qf4v1iJ$>zd$0)H!TO6~SIVu+U|+8fR+dQ%3Rc#hYCTcsJiWWU%Cci|dwo zuk$@F(OE9(a0c;vpTnv+xsosM;J$2l7(Fh6L&Z|XN?__^6un>_Gue6X!x*=%7UuSMf}S$@R2{I4t@`- z(m%0Xg|x*^j~mV^{z(nfkEqYc;gI={8BP=Mf#{IV$Q!Ulh5oZr{9c$~YxAu++NsH; zqB;Q6R$TiUC=MJHs!)064rwK^!8IiW;UM{FwAf($A7Bgjt65ucH7H$*6ryV7`#|c- zqX#J1V!>Y~EO8fP(9BeD;XS)-~`=<}B2zbIpBl5qKV1(MsD4cWaQZCKIWs zDQ$M;1zF>^+5Su*1+_BcYrD7ut0c{S<>L} z+@pwbalma$L|8)P6otWDryoDI$Dh+=ymuFeXN{yJ6kM@RN!4KJvU-YxKmLCSkU$n2=;xoJ`k@bmZz8d7hxacNI^!Zbbs16XEZuw)ev0hrZamt^g1$VM?XE=GkS30 zpGU#C4v6XIW72n84qA?HDeygsS)L3&Ka;A1L-ODb--7o;o+bq~)ttX=hp)sWpeZm% zC-C~G^>vp^PTn6}({jO$RUo9hw=M=xG0@8b)?1g2-^Vce{`^ruSkw|JxaQ+>ofI?#JdLQ`ni(u<6iu2DhfKU83Giqqx5YKe^&znA9 zLSeggswf|P;K0ZLijU``IRhS9zcl$e00@`p&i%QBb-i6#hy*GVxG zu^Fr|wjOA4@t{tN8{pNUJ;LKjSvzYl4D@PW1PxG1u94rfV86?#i>$@?Lhg^eoT?E< zK-sOtW3yYpHPxr#DM)y>OV{t4B`MpN$=9-UfC85}W~kyJpph9{l$BL1hR?dxlM z4E)H}HX&Iu=<|WGD)$|8X_2`B)aHG9$8r&oyS@>PfZb}!Audi%&Q@m}kINz*dc9j} z;t2+>h^*!bMt;aZ=|d6D%~|u+^a(;i7CuuB3(e;48DZ`LWp+Lh_(8t8F3)`ajli;k zf&yFn+kwc#qwX`WOv*F%v-EIWm^#X_l7~x>Y7FezFb7N-w@u9{>SE@Gh7n4@Pe#(LDJLbVv?W(~2Z}|A*&9--zKC81{iM!uDxFOB+pUgMc=z4xcG*(-$Ed14pbF~kvYc|cVt}& zS*UBDt!I<1S{NVQY!4}K{lu;+enyPScCo3AOWutCa+swR7akIizBZP9W8O+7VaK09 zqjONXYx}=kjQ?+m-k;`VxW+<>kRJ^bPdQ2`Lw1U2h9%FpJwK7%IJVKZBCD&EV}~lI z!fTEV&1xyHJ9Oqp)-#@sVIjpxXCk(e)62CrNOt)Im!2N7c9Ry66tuuXFPSF82c10} z`N93xG|p)!Mk@pBVRU-zKC+wV0HQ_@jJ z*OpgKGg2Xr^LvLLp(taVt%9ase!ZLt3{~n)mlR&G!;WP`(-P~HvMk4yr|9Q0T?b$e zif2sO=LUy3PB$lRRRB9_*?ttCRzB_^dLJOFPtPh*_)B?EUu)#W9|Dm<@>)1H^&z+) z4_uV&Y4p4-6~u?X8K2}Zog9{Et@OK^xk>aE@n))UW?4IYl>S?c`fR1l`Ab&kn5WbI zEl&6ZAaXg-5PU>Aw-w~^`UOjWDqi*KTd)MVOn>Hz<0KPMXTTXyRCWR?#!|s|^QS}5 zX2Lb31@paLBNkqu@+#l*ebxnGDq*tF$Q7X&c7=FH(9dd#pP5QyJWdU1`*!NtdX05% zMDd?VJg_20j>%n8kSnZG)&th#FC6UMYL=|W*d-75nD}(=^?11g@A8-VJRP9a za3?~x@|Sa><-?JUX>PV3G;m^MG0aX&rsPuKrBV2GhZ*bBNCFYtgckRQZGL&l@(PvA z)RjQQlM`LR)u<`do@#!so4I^L@nXE&W5y6N`881)s|NY3dxF%V8tB#p4SA;9x}Vi7UzZ6>z&}o zw~NrOod8SCFc@0bu{t*rO^b)hpt;=~?>jKRj(BLq*A`|1sB^`!`cI$KXPE{0FFOe| zIgasvPOi8icyro4#z->ZXQAD1ZFJ9}e-;$S5qn^6;piR&|C)Fhtfrk|SJNMk6DBA2 ze}Sq3O2`;|)m5mBn$zdfdv(SVzE&zlJ+7j5>aEaN=K^DKQG3(UGfD&X(Q%_{DlJQ$ zM}djJYIHKIpqw)^Gdm+V2kv!z3Y44Aek0v66UKgg>~V$JVMxCS^cr`2B7klbBXT_E zeu~=?TpymMIp+gU+?TM~9emnFjW8ZC=J@V|U{clRZPR7J&Y|M)FVPJOr3}yoL?%mT zSL{aJq1ifwWv3-C1vV12-}wGK8L{qd0&;en;68?YFjdS16q$u5j*J-VzEUH12 zUB5KP6-OEHQWygMl+1AT=P#a>;A4|+h9VN=zubgf6Ji^*eCF9djuy^}{v~QYf90fo z>{$N;QC`*10?w6_7pug8D`F=-alRtlgP>QHyor0wsTlo0dBU0@&Ot{u$z12R?+0De zLLp3;V;YTvpBBOFctw`%Y-%r2Kc->L{*D`XQ>O8F=5@=5HsYCzO^P&cJSGZLCr@X7 zOe{rT%(F}%WF`&K0S*7&4mbPu?-bU~BAFJL2ixagWOy<&N7T~TI-hPH889*u2O@BK zZGhx+Sc3WNHsTJtYXM)?f<54vbSwKq^MXu}Y*jN_Q(v;s<~k^iKGw_*fTsf}!ATiY zJADegITMJFh!2m5fJtv@S`$(*VnVV|#=5Il=j3*;#2I|uX4a!+dvD(Zj?w~N3gla3 zz{j=Tn~5S`-N*h|mD6SGMr-Msw(2G;7e6Xk^^^PJ!JEMFva?G3qQ>_udL2an3*N;n zVjUZ+^B{?}Cxy+=@nX9Ke+J&geMe&3EQTv9fGe`s_g#^H(9a9?sK);%Qyydg!m-#g zjlb(&Ti_zyEYb+FeHhn7hc2)^fmd$5P-X?);2jXu0T>XR0r}i72HK~=TC;sa-Jlq+bTZ5`QC1J!%ZNhh9##FaxYa-cCNAt(r5&$;yy7qX-ia_ zP+1@k(BcW79-x<8+#lC@RI~KEe4_rJ`%GC#`)MqEK5tKna(v5O&SNz@aXt?tT|rWke`Gi;cK;{y{>uFD2%Dk) z%?zHZk^HVWZ<>uVvax1&%Lt{RiVugTx2PFvsVkyDtI5NV`W_x0avPBFpHP-&LVouI zdvQ!JYRjD4tk{(K^6{nr5 zndMZ-odYX)W!}+5bj66*h>1vQdeQGG#JountLuXLPaN=So{43|2i%fV($P%6yo&Et z|H`106ra+|OekqS-kqBe<(qs}mThv-LReYu;mwGeoSfQ!%s{Ph2jWl4bv_YeQP{oG z3t7pqqtnmGAeji%KfJC3x?1V_~*(5iN9f zo&o>4y!fPGy4CdQyagL248}~JiHmuQ?nXTrN{OYmhvaBKq9Ae6=LT5{-}{`87^3p% z3k5k-&$k$e!JWZ+?O>J4c{Ti{i(r+1xNJ!G|8KoXORoBUE<&5En?Bt1&9a{c8lAc9 zu)aOug8AlPB7+Y=_uWn>0H!YvE0$ei0t7~6lEV#>KZsq}@9U*Vv_?qyJa^2Om#D(I z#dZhnna!{i3TiqIQQAm~538~sU+tDZuR2(d71?1Lv@#S%`WcMergq$vxNZCB1D7dG z(%TLVo9qn)Q)mxUP@w*@F}uH>S9}SF9k%VN`2)`Zi&BHiq+l-L0ne1x`~;p4WBAQH zGHoOpNOw$gW&!lG&i0kS&vw_z`-k8OJ6u+?wS=P4R-EIzYBWrvsUgTdP7ajW^wGCvi${j^aL8`f|#vWp_%2^Z3Z zj1kxjk^#}VjGM-h3mn1k3K2L0kOt`OE0wJ^!Ol4@b%oyxz-umN0OSC$22+?k64^4BLZ_?I!P>NuHL_F!F9_64K!@b`& zJH@UzY{lnDdV#_T?Wzv}#O>Qpzu9v=ynb!;Y`p$AB3B7NzqZYEZOwn>vmgkzSdQ`c z)B@-M#NNPG9()1;kv!wyx~Go65vWNAVT*#6nIY%VgkV8k1cOEMj{BkenwiF;&8q~s zM=2VR0x?6O)M`*jnBu8?z*+URJI`_I-?gw)ZKAJ(@v+x3Tp09JoNl(%jGb5J#dt{R zF{J7oeTg6ZAmA^5(U{P=g!>bPMnFTt@1a>Fgzvy>LQJP=LttJ(X3B$16i#4{zT=n^29-sPfOC(Rq<&K$NjIB zU@H0(jQBZvO*j8GTzqUn3mEa{t zSKx1jFr1Iznx1p>c52?6&|k9yXfDmmae5ewCRyfy%Wix9YnNWWTQle##vhV)QOrLO z4AxK1l~YTr4RXk+ENmqK?Hd(@fgLxnC9Ckyb|aGgK3>}N;S+8kQgXV00M0VYH&ayG z*jtM{B0pyOCK94<6FVei$32nowTwi=Udh)8!f~BwT>78bH3e#i?sckb>}M5J9P_54 zB(gi{ernF>r>B0CAsZE&!fW&3(#)hH_O3JOo-=odVVi<2+bn{;^6UX5>_;}`4mFxF zFW3WFs8Tdyd}7tEYPk&WKAYC5?Ef0Do>2c{Ruq0_O8>~$%qBXj8?m|p z*m?KXAK*ccJNsiKeOsmShR`TB>h#jH1(>E^V~@HyvvxdSnn<0#^=);a0V2rsJuv&B zD}dM}(xK5aSD8_oq!Z{Lql-J8b{t(xCg)%D_|E@4xVUwb4=C9W<7M=Kv4KxJ#7m0$ zU5`132YP?RU^~n&AGbzHlM}x!Q8n#xIaILi95g0;^~7Zi02^<)!Yj7@2_FdR55{tr zYY_c@wcyLA4-#ARK#WNhXYr0cYVYA1)$|q_j)q8O`hAjnoou@7aQ`8x55jUG!J!=#Ag%zn;?}s!xJ3xmRL zz>LyPS%~lm2!zw()XNO~)T(laDufK^v$9!l!q`F-828JnTpRY|)$Byoap6BC=)_K8 zR2hY~&sDC$+wsI>yFace+}&tBw%s*s=smB_`dug-F(&M%`gmilY@HgW2Qq%f(hF zOITc&wPMlOr}SMiN?I4A1uJH9$pz4ua>nK%ZJtQ8(4ZKzWrWUa8tFKGA8lH5j-fwe zS<(tTRe`HQIV-%JM+6FTRts0W>ImMr@2NIgx8O=;`kS`Fsx<^Ao%=JtXVLSxM#K~$ zf-8d&NOK|#JX%#3T9++*kxmo-FYa;(KI<#KK2bK1e}nNKVG_UE!zwaii(E-Y?ieZC zjrV6-DfAl{e8|8XKxiH>e+Sa3q z9TZpe3R*~W{c-*E>(_8ZKPF70ZuY(o%&w*HuF+}9w^)yEIePsga~i@;H07Tzf>yO= zTy!26I2xalq6@RCF!X2rr3Ib0iVpKIE5{9_&F4Ees7pxnrqcs1*w}PirJSAl3cD_- z9%E|ugJjbjwKMRXsS-l`Y$K#v>I5G%%sN4?Nzbh1my^+#ldC`LkJoeL_HNCt9}SDT zuWM#{Ek5ra{IHwehrgDa_kJ2&tnW+EJ@PW)ZGs{}AUfib$>7Dx9==!i&4yP@Ebr$3 ztjT(sh3CNJ>1xFYWz=pj!iEL)g@ABcXAp{KX%@ORFbmfEWqreav~c`l?1H^VvrEEZ7A(rA2x5#q*gwwY{L*V=su|f3VP)u*^ZX|kTe&Bq3BRBMV8`Njt zb!#;(otTipcPFcL?$m~f>NfuqpJJVI5TB_=hRL1PJ9jtBZX(laKBDbJ^jLSB$`Ovt zIF5LPGd&n@RJt@rpUHM#M69<>JO6^2$qX5i&j2xuF1e?iKy0WgxMYhhosYk&uNhKA zg6clSC6%GwCN_L3OaVDDQ@5yR;Or$yXy()9_2B` zm}957{_L9@!L3~?3s~K&{Tm7}U4JU+9ozW>+2)N$yu!zlFUY*;!qh0lM2Em&+#$jQ z#}a=TrNhevQmGq_5mFjdX)6t#@x|NpxT4_i4_EAm27*jPuMTZV{KebV)ip2Afu4d6 zev1;HDo)rI37M_>Tad6Ud{WqAQx|)#G+DNc1K~9ApPVb)mIkZ(+K^a7%cAIr_{gm z6GDgo z&H_Nzj(GUZU@#7j$rbKF4#sv2#sbT>ZmMmc<0%D>V&a#op)a`dgx$)B0Vpja8Q4@* zuShu`{o#}+Ev3rJ#}#to-v5)b)c?)%@xs9Odk!ChYC#}`lz4@qeidBR zQVnnRSlfoy4BlAO6pc4tvxK4aOG4OU2Zog|TOi|aW+tFotxFS(IKIRBBFs4k1((R- z8B1U_z60;p-kWjM@sNgQy_a%Vok$whox_Of4wsphf+FsmxPv5esry?WtD%$PlHKJBW9v>@}x zciI${&+wc{W2R>FV;$HO$di(g+T6e!SkfX#GRVT*csZ>KH6_w^m@&`T3>J(4 zgvXJ7>Z!iUYjZm@i?2U6ySp0J{&VTO(VDd(qkI_DO^GTPUN_V9 zFfK*u|ENa2&yEn3M;z=%TFtbOga_>3I`KH$Tkil!XZ)}?BQv3X3*}f|xKdMc={U$_@Fvm#o!;pjv&w}fvo(>YkOaayR zHa$#6<2)=Ot4C+UMIfiROJ1p-!0BGRq)IF;!|^`WmySUt$YcpoJIwwDCry_4;<318 z_>e;|J^JreM*ZuR%%@eg_ym{F!>f82Xn$=tQs9>$@NNOY>ge5S{4_VIn<_uzyez8` zB3*g-No2kjH-)|5(@RbM!4giOgM^k2*q*UW1uAa3S$%0gyc*YO55b6UXl2d8duz-v z^XDeF!%m&{8YV{ zfu~1!^zCQZUmTY z__nZ6jDR%uxk)t^;ju-TDhdp4H^BhOZcnp#8K&b(Uug;u2u3SDc-SiV6UJ#FwIA3% zuoe1rL*+BxvNk8!#1Ymy;8^8V*;Ruojo<$9=QZ)YHoSnn_g-h$#H^6A(*C^|!u{Y@ za5rn*j=}z?I((%^`!z(oi{gHH*5{;(HvGf#Sm0LrsMzQ5-&*m1aR^oM^V*}LEQGNA z4NG(RAkYv_HJZ5?gy5!I_~*SUq@=;^H{y0>5VI1U-zUrwJ3Z3TVIj3m%i%lf@Awp3 z=axvAhEG~1NhqIL@#ez8q>sc+1y&sU(L8}k5b--yCVh;34l+ox5)WsJBc=7EZs^`M zX@bUneT*mGhAt=X5tOgu`?f|d!|*@SxlLNj>CbawfaM5sOqkGQU((+r%+m%}aCI?m zhz`r%R(I!=z6FE!yL7hJMcb_5h>SCarMJZ($HuSzY&p7$Cx{?71c5_xk}MasW^Z4YwJn!S#M8kT+xxlwZIZJtP)ma2dvhD6@`ldWV3#Q9 z84_jXQ}qTRf36lc*i@xHPRre1+d6eV(kDT_6|@y|*4^#cl#z|}|4<&@axk&3sh^qm zEPs8yZZta>5Gf?sa0<|DsVlPu^fg5}@TV^TM$)f!@!dEMK%~rFx~oAW4szf~Vu=LD7{yebIe!cDGSK)5SVrn*xvQ~Q zaqMAN2I1`Oz4)6zuvT=Lbu73VcImcE+2VorbM4o0vy9q+?lgo!8=)+SgG?5FB`H5k zB()pIVfE|!ZJh{}xqP`qr=@#d?e&$GXG34_68b*&`9e$26kt}_aZ|-z$n#5`>O<2% zjETz50tWM4H$90x?ms6Wq~HcZKS;g0h{;?nG28+j1;|f`ILwC@ISm!Z%^2NO8Ei+3 z684b8Qb4kY)Wie~beW*hR#aDlrN~Bv8Mw~(MV$FoO)g;F3t0x+$@tTxnLuuSbH?-x zTZoZ4J|6=gKm9|O?Y9V!TPZP#$J$935BhW8q=50*C@I=)F44UUub1|A95&Rr+lj?~ z68nZJCJ|r04LXE+1Y#vJ-N^nN-;0Z2j zxS)QfYTGM@v7Sqrf%p^RzkJpdxAJ4*$^hG7*LH?_Td~_5BgrJIH0n zYSSe%3IZU?43nM6kpt!Nz1C|Qqk69#6f!Cgcsnn#YRAd zwQi)RuuI16NxRCQ!&|mL935db3k&}RdADz$AfF@3c53twYBhr~MwRI0BA8pZ;jOq0 zU%NrM2)PncCfhMYQ}qEzHD?DCaTuY^dqmrdvYU7l^o45nuq)it)aG$W-F6Ts91J0=uIc4w{q&P3J1X z(>t4ah!pDZ6s%w-nD~!J=MU``2S|py<3My(?hJEXY(vB@Le$`2m7KCdSRs9>F_Pjn z6->!=`?wZ^C>8k!;ij~w-u{_EpnG!GfaHIne@NjoR(DIyy#HaS{Y_giU+E~pMZ0k_ zmlswL;mAS7v+3NWK#lX^D@3rwh)duIi||C9AK=6?Obm@z5|xDgnclT~h!xS3XaR+7 z+OFGJrbr8QI40)Rle0>SzUuP+K&!8lK+ZIMRz4y-3+*?&ejgEL;cBs`B@JJO0YcI2z4@0&% zD=CZe>=u(=^zP#_;|0I_7=!t!n0T(8oRAI!!xYj3ehG zs6LEW?I-B?)V{17{T-=qYcwE(-MnM{B;DY#kkO1lW*E1hclY=|wiVsNA+)a-4(lPh z-R`QqaG>&Fza|-~e?%N>x>5q!JJSdvJoqsnEkW4D9{P$!J_oEJux_lc&e`fRKHZ_2 z9XL*Xr?@gAYSu8)F*p@WWiQzM0(ZaAb}29xIwW2fur+bT7;`9~&~z0rXfWB(t>D() z#?zz2SD#HFhIz2xiwg)vz#E#~ z$*iHSWepQfwi2<6J%Wa;b0unVl~tWxc4{;uiZ|i&_9%?h0{}Ver9{q-p-Tjbsx`Uo zC)Or6hs1U?{TC~(4%>$X1+W}OZ&t1DFuw{ysU&Fx=;G)^i;H)zHH+> zKBV6Ll;(KAs47u7wmvk!t&hme=@P8Fq@-W*fUcuDZTWZ~mF+3wH~>_e$LVz8dLdYM`wt5N}P3(V~9Cvd4%E19B7QQ3g^a|#w=8T2@| zJq}E@;V>&Aq(7i!B9E>il>@IR+K?L(0@3k%JaQ6|SvjvuQjOjxq`mi=&h`EbQnPRas)RO4 zPaC z$y&@`yz#P^7k?IndP+f*3s%WRTsY6#A>@sd{|sDjVErN>zfL+M|J?~Kp?-|hhjRaI ziU?lGhjd9p}E#b62}+Q8K51`e^W@ zB<_NIht;mwZi1WWJPDa}=-mCo4jD;3<95za0dHy{OH5LG;X24@#g5DN=)z8uwO}_M z*eI*H)Hrtur?y%##dKtAW|ge>GCC2B!vMd)zew)g>tiI4K8*L>(kyDMkd)8)DIhm> z3fiAGu}r&BdeU*#K_8e?xozg$mGe`=1IwWsv+p_Yd1hbQwIE9`oM6r2-{?n%O!>I* zF5JNXHz8A^2;zm_6hw3s*@iepQIkbg%?A7dvsLK|wNt9xA*GM$3(k1ONFq?E;c73( z@Hu5;q8h}trYR8nBSL$99bbpqz0F$MUv2E+5cZ0zO2WxZnzIoLfpN>p39iNn^ByRO zF{QCZ5H!MPcmh~!zZvU#8#6%Nee*g!U*&)5LEPx*5KlU?)gujmub*M-vS?jgHM&Y} zkO|qLCpke5n(M+GthP2L?^hW&k+nw2W^gJvn+VbN_Se;@$l8|pu@nDdpm+KcjW7zl z)pCSp{GbjDh5|3E`q$eqf(J&fyDV-MlxgXfL%kO7Y}CYBA!MqLyx5n*T^g_mCJZ;N z$uRTpCIk=BI2=pc(j~idX?)zxmxK%*cPM+Yg+Ot2#xn>F(Vx91qA$6HFuw*UK33)! z9=RB7w4Ca?{#IgtrB`B0KGV@db5kA}eD_@7 zv|StHwS@`lj?9EUyUrTEui&`sgr_Dz!*#B-i8+hI^OZv7x8Sxeg{dA}s2x|dcGHuq zBk$G^140g}_ucCL`A?l#!V+>ymTg{kVgCDEn$O7eqsepe*shn;yzXLU@L>X@DpOj# z?%2)saWt~ZXM0S~=TC4d^yB?xo9okdV-=rxb$JFuR|J>A2Gz3N#kXN9(-WjU4Ns5o?u^NI-&dGEbmG?&cAtSG_$lS{En+O7%9r#Dc9{h*wD4+h zD+juIcN(J)Gwla?QNo)O^QXzW+4He-+{#JrGPi5+#Qf{{E*SbLRKpH_6CJ{ap2vCQ znVrw$MAOxsgFxy`t4%bbpN=-VI(btv@ zQDl7QheX6sZ&Nu>Zng>7;)hOG<@`*M2;xDsq~o17D|rsIiTKb4CagIx0G=BKTRQhCP382 z1W0At}kS?PT9mN4;H~ zVXz|6&k1n*tKyd5?rV-7qaJZ(wqN%20%ZEKH&`7#f7xe`uF{t0A3sfI@yv}Nb7F9f zPSnzWt`FlU`gYgxWoR0pF{R&|h11O2WZN$c{OIY>aS=!K1V6Ckp}NDlYlU`ryaH9$ z9(U5S=Rj|ZXxA~M*Xss@v9%~I;r=yEHfr=h-Ud|k`XRxyf`-@1u=LwdIFqUv>QRO8gHO=?*Y;7)z1=l-agvZWTlU(f2kTHFXImaI>#p9j2{PGzzVRPCJDhv{ z@|n6dBz=HR#O2k8`yB$ zwhgtWg_lmWpBEY`GzvQY{y%)ZRahL_(glcw;O|+PJ$k z?(XjH)44b2o`2?Hp8D&4g8KHZUA5M#RgZD-fj{@th836R0=2sXs;MiCwNfN0OcErK zt`2R=w%hry;QFL$-8W_DZ$iIWkj*BlijR~aI`eW(1rM3`lbC!pm!pz8IKaNk-wVAd zO$hbw$PgQ!W(w#ml6SU1my>%br7VNrUEJz{u9Y;Dh+Lity@h6HyZI zqjMn0j?VAP$b&f|-G} zT2$7UwMi_Ay2z2MT2JPOUCR>9G&P4v67c7^>Lv%TXjz+9A|~6{_*rYQHUM-#vodCa^OQb+13qx37+$t@ zJ|DG3yV0QTFhP9_XWccVd)X^vDeW(nQ&6-d-Qa^YgBu*JVu$c^xRW6`G32c&N=Q%1 zY)8^*s!3|!pY8^fEjmxvD4Gcof9g)7I3ot*I(@A8KNtQ_e`cmU_Vl;)pvMGr-@sCW z?hkUJj@j{YW=ROZMdPTuHyIZlIcv{Ildl-5vGZ$>WMb>yj&J{n@7;B;E6BL~?BUZE z?L(pGre~m6t?0$WZ6o!dhdIh)=?vtA4RwrgijM`^8^frA(ZkzWGNP3%&4SD7$T1 zlRq#24{o9T1GmUfTK&Df7pQQ;BAu`>pYh=~P+lOB9tR|u_j>w^3H(I41sR^79Ei@L zUM0n_i^SX7@9K!8UDO}>> zMpr05D?(jm!E-#cWt;V*dkF9PKR?zyi0xRN;Swi#TlHb(2}bO=`V z&h)-(o{*~(H~feNnD6}7r{^(!#>3spu|J&TDIMDV>{>GjtUaA>@4v74my0R_Vb|EI zqf7s(jzBptuM`0ux70pq~f6In*u?}keW`6jm(@7 z-VAK59ED%1?ppvPKC?PLd#>*CkBA)X4$QP0wYD=la(KnY`zkMLBA9E}8Wt-H)KuSOeV^?-bx(j!=FR3}&{$+` z()8$C#}ku&)KuPe5<9wOwJ56XEuncb6G~jr2`YiQ#efaSziQI+ENs$o3{(D6*mr~d z?DJ}#@b{lef*{TYrE3%XIJ@v!J zc|t9Q%b>u7CcH07-8OpDS9ciQaZI$!XcwN&4X2$wd zmoDz=c0LuJ=g&4t%6divb+*LyOinJX`a zs-e=GaDPx(vCICiGBf$ZkK@uV;r!B7Q7LfcQC zIY(>12lP?fAAiXxc05SA*m6!NY?wd}ld!lw&fF&LelOqRl#D{&?p#IKiQ&Egd_et0 z3j-%Pel|>wgbKge`m8O7u-Y`cIhucxV@VTH!#~QXg_AHidyEDh(WLmZIC*|OedM%P zVL-Ynm3aW5Sh<;2=Y&bdkOuiblucxn%N%5)Jlrb#+9KHsYyHGOZ{(QWUYRO$kQA0Ulr%I=H zNas@usk(zKiU&%|!LVku<(6=%p4&0n&O-gy3mWf+DXN6&hTRhVU4jpLWwo-*IR&;K z;zFLE&)e?T-qs-FcCPfma3G~rl`c&JN`0DCxDNMbKK!~yK;^}opA{}(IVm@WnDg_v~z z4gN^Net^xX-AI|Kh4U~xIRO!t@R~Jf^61?q`AJa!deQ{@z;bej{}l_I>*fHbTpxmg zKPVs{biB5w#rkeacEHc}1ho1MkHFLgH5kDXs5mDHI|Nra5TxWW9+y3XJtN5^Y)IyU zKHj&lADYXhF;_tKFUqDdA3f>`L-Y-7vv3UlOqyVNt+EtMSCP(9FT{c@Y{FS)8UkFV zJg2?Rv`*zY?`wZ|RCnbcMrUi5^D0U#z03808U!s@;&{Y;ZN_^Kde9$nxwtuQtMitf zU~e3j<~2D{f7>BeR%sl1Q3rFp&*RldPIG9L{y!;2AM`LsS7t~4(L5D&7HWZyQ~ATu zKIj~fH8)q$BMlK9lLHnrk`LaF23d5%jJ&$Dtf?D{UO7pbtavuzGj!s*aRPXDoC=hdruo|02S)n4|{tBmB~?fN*?@t$rx0slgQzQ%s3w{@V*JC zZB~4kV_iYjUIp++*H15OGXXxacx|_nxvy73>fPCopqU{t>OoJPNf?3Q%JP>co~elJ z>dpDO0k)4$($e-jDC?Sl;?I2f<`pJ<9^saI(n&f$Ci}A)p{y@%=qjHzmqMNNZ&r1q z`-BYIt^d}6{xtT=f8-7&N%dU+JxA_R0TgNOYL{m%AzC^+T>(?wh;-IZFb{2mRdoXA zAS~%Fd$M|ce1A`MZ=O_Z`yCiE#La%)q<_v@*tSG(uSljM_GV~dXBXYh9otrW|M2i| z^YEd4R~c~YXknia(4}6X1Cq}1l8yIYZfK=K(WMma8t)MYnCgZdF3fbp!JGccv~4pLUNS zucr`b5^hSfw^op)&c^fvwYvUSe{lY(KjgWNP5yft)WH>*GAx~+X1ZX#MG<$AVUBVq z&S8%NOxouHa+yTU6wN*E&WyOWP^un7jVUy>qEm8ccab8US&(ZTSI=5A*p(MIGuCRP z_9+zR3b^I_yyNdUsd#K0UO(EPC-UyED8ipQx(aSm+vdVtsj!#jIaDX2d-(G(AWSQi z(i}X``Mv2A97E}^`(Dd)vYnru0f&vu%=Y_d<-s22(wE(T^!$`{++S9oTks)to2T+B zj6HmDW-s2mH>GN$TRy%tU>0l#)Yr!&xz*8Wxn0b;%yf0H^m*H^Wb9*guowCf=E^|t zA7EY@5m#qGjd!sVUBee>`Ngu<_DHfgQ{TU^I=W-~_2||q zp=Wix)r70jt8#IF1+;QRPRrN4*?|5eI=!0gqUiFn+S#E7{(WVQ=8LOaW#i8ki{i`B zYzk3ePR(ZjZht!4-arXi>zrIHz0NS}%J07zG0g?}pWx?(&HmLOdLLnenHzGtZoyHl zdVI|&z6n*Q{;)yQ)N$*b@Um=z4;StXlpOwHR)jY3x_00G#}xmE$Ln=wi|6f8V#kZn zr=FYjx7|P6BL3`Rz&zT1xZK;LL@c@oZZ=aNjHF8|-PDh_N<|zvmf-f^TAP?q93F*j zZ0Sf@FH&2znQ5qH%sxEzxq+MQK-;vAr7;u;lp1fY2$Z$QvuB5We96R9rlF-^VIH4} zGv@TT1_$$Ik`b)k7WVxCD%pg@ zNW5(R+gM#3eiddZ;!ZglFQ}QuA-0%7UJk=k%YZ64j<{Su#HkDSaP$ zzj$h19I77h^4E=EshHEXXJRJcEPlYrED_ZyHVKEs*??B7$9w5~sISLQ)|jp7W+_(* z8Yx6G9he-^k(3-jVe*EbmS7-q`k7EzpnxYK-jXsg^L=QmClEaM2^^RL_R&dOBKDEvWWSt9H}n48cTd%iJ2=oEDBQGv zVCPda zMq1Y~E9(I&V`(q#0`c@ACUkDtnhtHh?S^vCa|7|Qt?IQGZA=qv&RmUz0);-7e-vS+ z!-P-4G-M!@=Ks9_cSgLqKe(Ru89MTis_x&?o&8T<1qo&c|JE)0(%^TP*Hn8N7J~(= z2|{qj4Ci#gN!MMC0b~FMjp!{y5yBbjwLbU~zNk76b&cZMx=#u~v}lIO?852A9q0Uc z{Edy(g#)*i`}6e>FFp%iK_1`d^2MOGPm+2~OZUP1nZ`o)Dwy}@YZ;@a&&X%dQFn)W z*v4bQ940&h0v_8J8&u~4{xJjc4Q}^9VB|PhGLATWStq$GNl)?ybJMVr*Ad?A- zKR$hT=_6rb^9Wa<3_4u%GZ$>r%}@U!C_BwIJ2~(f@_FDINs*RL>TD}Lbx>s<@%w%> z&n`TRSbWbWN+nI9x1xdqX9K6c=OG30{^K|2=7t6dN$;kJ{?!f;c@mcT1weobE6YZC zSgs9#2^X)7Pao&MUvYZ7IMw(B_!jY|rLhb5c6kzy=N?Wxj7^vvM=s}5U5OK}JGMQI z>Y5SNUuMM58*Dny_2cu}k!dkjvf~7xn|O@hfQx5M_ct+gc9CYYUpm8$4&TQ}jPDYhr=_Yw7M(UR}UH-Ax9@VCNDEbxMxy&+D&_Ds~fS1xwlg0Z)nL^K z>f^zpxzcpX3tFL|y7)Wi4Ks4}yG4GXojx$Y+qn0eXemen*8DnXWr43XnI4TbU*ycr z+i&ll&pSG5+BtC}Kgs>X?)0^{yEnIjR{-)_aHJLxw4ItMr?Mzy&_>;?U78!)T-Tpb zzIfdBjoaFm0id*V3N6-l_d6jLbEjJg_WQoMGJ6rty-GaP4v7IFAbT8Vj-;5F(w4F3 z9T@iQbTw)*u`!pTmgOA9 z0wnvl(TocRZEeczv zh|?4bW4;ASg4}fCVQ}p@=w&c>*-;pC=$CcsXh-AfZR?VXX2LYbe%M^ZMFP7T zQ_$nWWBcjL(R-oq*Ztl*eA$=QGYQ9|>|b&ml_N?#s_Lbk4_5iy_&nf{$D86BOMcGv zop8BU6fkw@thB9Gg}3Bhx2Dr?8?aqA-zi2Z+#&frO=q4FMcwci*jxu=>ovZu$TG7U zG2-HES0{iMccIH0p=8D=0IeSZ|4u~+iGkXM%ug|Gf0Y*?5wHOccL>WSsbV1sB0pJ# z$J~y;xya{?IR4-#!_6WyoP{QKWJO~BEHAZT)sa5IB(2){WSB3&fa2%q#!}9>fq1%* zR$SWW^l@8}D!NmPABL@tBjoU1{6v?Rv>DgVGPSmqAtV^wALSr{a|#9MaAmy8uB~9l`lguZ?w`~V1pE*3F&9VZMSpUIIi1!k=EP;x!OD%X8K&}?1nh)$GKvCyD9RYJGjS&6crTD z*V>+AGhaT>bC;JJBEvoxUjNSay@>^Yb~?SmKlSZ-y9C_@%`&7I;>8}lQB!D&GQFDo z@1VM?{VC{Fd5+EgmN;^7odHUS2QA-mH%v862ceD|yNZS0A9`UE5xv2^uY?Eb3B-aI z`K_-Xde4Z}c8$W83U`&PCHT!qg$8)Pgcr|rr8oKt z?1rs0=FOJ}hTDZgs2~AJ4P(B2Y~?q0G{0u4hj1`bIxu$UGz}R z!tfU9v_6X{4X3kd#1zJ&wA<@5k})sU-)!n)0GZ4DZR*G>rC)_6u)06L3d|R;{qPFW zV7*KmT8-?B`>NuYqS?}@`<2#Bq0x~;6=!1us#)Y^@*97$yaUvK)@|Y0RsncmTrwWN zm86D%=6~9Z*l8(a@w74&Mb@*DE)8>uphND5iG81$nx;Y-R@|!|l@ID${61G5bJ?A* zdEU2qe3q&_D~nLLi#9_@0D7heuV)Yghx&QrzP0@OOeTX+4<7cZ(}U@v(A7`FQS9N6}3Jh<@|xM`XQApD2*rvy)B(UJTvxgA6a zA}u3~N|@d<&NC?nH!cfIBys`y?h^%K& z>Vgo_h-rc_6Qv1GJG4|L|A6nEJNmj6Xu!tIEF@UT$=!Wa=%aHs3vX471&L5o!S$!y z)qOCtLe$D<2l~H~-Us(Tf*=0ZvcJ01PAE(;>ot%4^mD)WXW>bv*C{5SyV^{Ol31xB z2eKd!YqEQi9b)3l@1AHj-dE8GK=ZUblgv}UmjS=l`!%nQ=avq5D-oX+6hWmH@3R=`H#Nu*3)!&bslOl{ZI{cWR3 zW>qH}!k40|SY)*ysS9IG+TtUYF;GTA{JhN_E%1wwA@qnH2L>j$OSi|V$5uI-eTKGt z$p>7cEf#ZhR@v6FF(w3~DLm9kEHHE{Q-<7#sUDg{O*WJHu>pHHFuGURm!-ETLtZOC zFfWrfztK-3D)>`hAF5G_99E^4>z&R)4@k!52sA^VMeP$9SA;Gjyk-#an>WF@J9nGw zd>JKf9S>4)webl~AKe|wqgAkL^Qv<@UoEFwX0~6`Z}N1B+{L^Gk{3}<$6CxV;nE|; zRz!4xc0-qxj6zh0WzGKkTAQhWUh75YTD`y6<#-VoNf`1)h|tR?GZCKKFyxqUOGA9O z7*#omWbzBg3f)X02iEo9X#n<~5gCXW&{9I;*FOE1;10q2QvFw2#G(%B*C)fFOMA)> zp^gas^e*)aLvoS#2V03BrWUA9EWaTbR!TAFhHwp&u=!L+vO1RltQ7@ie`Y0wTABx8 z*HOzV46)u$PNIS#?22~abgUzayYDPNLYmk z35P?KMPVJU&HXF(6F>d>%1W;+fBN`CP`JB|v2j(+)ZtB1AXdb7!JtboS1*Fj#EI@+ zpdRB)ykFdFZ-n;15~>e#>N#QN3Kgppid@?H%ycEP!!DUnqD>}X<)fS>cNT(B&taJY z-}JGFT zV~=qr|MB|ffjt+FrcYD`1Lr$>A=pLc+hHa|r_kGF^d%|8*BLg{tOM;o=@v#u`eww| zNWC}uxJsvXo|cXrlW{I#wF`Jz%G0~N6d;Gs;I=uR4{p+kYY6Q-N*}}&_@tU=vpK|-^N(1)|l#64jLlpcn4s2Sld@eUG3_E zd3QPq1v#g-1H0cqh9wHanATZaLmK^rUskt%xJ|s*3ED%XX!a0OS-G2|PsKqN!`zKc z{}&x(hRTrt!+1`QdH#2Pv{(h~4op1ZgD7}lP#H?#)#Hxv31o%WXBy8*94dxqU7lRo znRTJ}YazeqQ^5N}^!o!yNCWrMlVpHQATHY(eWrxP)AL19!HviR{>S5_Jtq+X!ApI1a*GTiL2YuFO;jKXQN9gU?a+WJsBUO zGN9!}^GdlRzI{QB;}FS)xHQCa_7|~hAWjg$J5hb1sUgW8Pfkt5U6!~+fi*!E!>flseq~wx-fTJfh+rBYGmBZgX#RHCn}0*Ai0ILr08egic@MnZkTFq zGM=4F%;6WE9Hjw^uUoaJ{jt!8@2q)iCHZo!7wtC&xrz5SYWaywAI9xX50YhDD=Wt&P^6SmpdIgU8xBZ`E zi@hsAdlgB;6qF3QsQFKuUG^>QdLc{cx&t22-=jFWTwNPd6nKlnRYWxghE`uM<5N9C zxp>FOTbV=^Go zEEpB(5P5*LkGFxSu$b+rOx)i4ZO0q-WsLd8t1k+TiPjYi%T#!p44Q7QsWikB+6LUs z>zP}}h^GD#uG+yaTF)0sf$I|QbrWC_eZXc>wm69#wGs~h;!@1(%X%dl-!iV!vj=2}-Sn!;bVOb&rgT`@(X`U+rdBzDDy z!j|U3UO4zRZnD8Mv`>&nJ#udFR2B-xcF56|C*8u=H2#j^WH?g_&SK1b?G{x_zl2=p z&c(XhB&Y(aWAqD=;7koQmQZTKIVYGn?H4phe%`?5CLVa2#ToQhz<9X#qQQFnA65Z`VXUoqg;vNE+cM1;W|A^>lES`n`k-j8jVb1j&(u{ucljPK}Tbe#!Q?6ibxfqp@`H zx5vLq1E&%E)2R&W=Li!#IN?&qY}YKhbQ`3w5buzy{vR;t%@2A2 zg))5S{>lY50^z_cNDwi~v&>?*#q8LzXvHRSu}<8H@;M8NAN*gX?VrZjt}_K1hpowy2$aZ7B!FI?2l6j8T0 z7TUI0>w*~5(dGJ=Ayh*=$Tn-X@$~ciuTiQPAooXnmPUvXbdtrGMQ^_%PnikM1TV;`}fjf7#ZGxhEZZCiE}!x~b2ZI~};4cca@ z9^{Xl|I%2AIUr^{f6zSWufBx89AY9$Lq637N<4CbP767hCanTvn8ysL`oZ(fC5-GZW1WaH2FeTBF~-R& zWP)8y*O`ijM|)j<96IN4l4{iQFEEzWdUweUq&K!ki0VS0yqWv%x7_U&g|1|;RFtNP zjcY*d81hHg1t|CVnhZ(oUVanJMp~jDS94t8*e_#Jhc-hpo0db;J6kkyY5dujs?jH5 zF4^&+^vkGe|BR{x4yK!D+wARa;Y3_8aks$9{WkI0`&j8ohgmj5M72^I>6XQeVg8CTpXLII?A6sYSvyx{a_yt*Z zqdsjoGgkvmb1uPl;Js{ue$EOhSvqMvu8C}u*s=8gIw9nPVrxQUvHeeM;~^aR9ZRNl zk07r$18!D}c2K>WU5@A0U&hw`y9_?%edabtu49jlt+L_iemZM8_dWMo?>T%0w5Ip_!r4@h *sqpr1pR?l%rQ;?LZRYJ@fqlI zUWh`Hoo(-`TovkMJc62;jcp1%wy_Ou#uo9&%dbFH{>_LiX;`H6V`oPAZ~IPY>~&qM za@Z>^zAq-Sz`dyq87_|G-+tkrd%lh=09_6eJ`KXcM;g|Nc-y@g1p(Oj>~BUCyS06> zlV03;^ZjunH02bkXhV1NOhlaBmJ>#b60UiuUz32-_{UQAO78P*PD(9}i{9)pwzd_0 zCu;@`R++(`a~rnFII_v)wf|xz<9VRCu??A~|HbDrRKYJoD@Y>UhvG%Ziu(BiXvEwe zToxe^!Biej*WXSP^?e`vBk^wnLyGq^$Xna}?%n+E`|<5INI!PLbngT)i1az%pY$e; zZXKTCM)U=R#m>cAqM=-KD~K846dGl7NJH!K~e!qDMC z-IVK*WkK74yGee0%v>VGAb>X9p0VRhc_FZ!j!)eZsGhktD;ZD}7Q|js=qNmHhHBPi zC0^Pc7biCGFwPTsV(ZSE5uVoLPj~{xBgFrXW3P?=PKbue*}@pG9nthCMsJ>B5>Q>? zi|~&S1cV)fuD`Se1S%iN>dq5B@nsGS0a3B#`nIc>V&c ze1EFkhFD!)-TlhUQqg`}ZufdqC}fgWs$PkP{0MR{*-pE7dlM9?dhsVp!-?6k$|@;j zv0LnSBggKuY^gw+`HgBTGkS>hX&R@LJ5H8;k8~zo%$I>4I>Jo6ZhX4$&U{Whzu)0( zt^ji*8w`HNcfEuZ1~ZyI>=IXmw~klr81tHL>ldI`YWZoNali771i$}2%Qj74| zgtRf`SFyu%b2}2yEWTZkIpu@nujJ$u8lnRI`K+WOX`a1Pek4pEPn*boHh@FA2#!z~ zd$5q74zR9}@F`zFs0^XP`L6Mlex(7SpntUDb<I_8+kv3nU+v0J(r{3{2Tj}?DY#n3D#BR6lbuJLCI1mKm=T()YWeBbi!AENh>W`qTTzz8lfwPx1~H!6ZF|8%EpC4h(9Yb zST0Hx4ZM;|()-bSgWwn^?fTSzyW{KwH)pg$0q8m>|6=?ynl0cwmMdoZF20>Z%&rmg z7r_cptcSpUQ3QRcHMsv}dK8tw?FX-!TcMX8lNVk_DI_K+GxY9wRQg{>Yz^dGvh zecwUJbl&%LegUJczFu74J?V5h(In@-k(Cl*4v{wfoB}0on}xP16Pqks3FhL!op@XF zLR;C|&d&Hfs}MA#EGp1$H14>O7M;DTK42pgA%TyRv^xC|6QraMejpln)e0WbPbZO! z9e>^Y_49$~a|}-|9gcG?_dO44^-l3!k7!Nn?ynUpys4y`@#H@24QaT>)FDX8?h|jd z$zNd3N4Eb`5Z|EKs)f>2Zg!6-w`8s&fWxqh_} z%*aX5Yn=&M7_}-v&`EeiO%p#CfbNo^Yd85!KCp`Cc5i!jUp7n)IEf%l)1W_w8h-7X zU7TBtuKhrQ=2jYAIE#$2h%?Wy&E?R~(KJnh^&Cl+bGZoP$mie-tn_{yS+dIMw8@HDGBP$(C#Ab@RgPRtTwBD zXeiYw6_LX(8+Qzf1vZ*gO11EFh^N>+wkds;9WM}YE5qSNf`WL5b!^L-pn=XkR~__A z-0|A=aRPIuI{e?))6JN;TpSz*KJn^K6=B0)7{4Nsxwg({`Scxb)Ha%6cwY$}FQU&~ zm>52OgZdxeSo{Y_rA;pT^*69ostULv%M{n&`C!+Nw_`^O6zT2iS5wHcgf0p8H?M5Vf)f3BtJu5H$EA|m8|ON zAFv!d)9?*kkV>!}mt6)?7XFaD;+%-uP2r1r_+%fvEgt3VUV-cjqwt>J$O)jJN9P>J)S+WvuE?NmCAG22&!o^pjCA`NsQ_?%YMC)^v@ zu++8$$rT{bY2ZFm0@=W%z3G>J;sk-AY2%$NhDwh6jyx68j}AW`?%eEL26fB!O|LNU#ev5BrD?fvqQWIy4 z6%4!_McXsjxm+8B{CupOd^A*h+rpTMKAt%@KUnYE4E;D^$vOjmGNvYSih9|SDWG(y zMd{)VZ6cGo!AD0%z)`rO{>p8Ypc2O+Sv=)>ugansnGrfCXOa1-yh3L!aIv#TKKkc zfZ}?Jtft{8rhWG2;*gj8bF!7Sx7cXtWT2oW>z+QaZZKu%XA1Tk9heP23!@Vp`C4=7 zkBfdd8e?5P40|3lo-@;3+kG}@UytLc*N@kZKpmV=)p45}J}rTA_zu7N!W^dX#j3CC1)A|D0UX2-UyeU+!taAObWV&;-OFzJ6Epe7&82 z-%)(~&d}Rg&OJ;RzKg&-|AGdSvK+m!w*FMU4v~}fxf2#EA92Q+DDATxwj$U@E4@T4 zDvpl2)9Civ5@iG@08p{o*zU5;BWFT)ea+O~%ON5skUaL0j*eg$D#Fmf9{B9F?s+{n zWX(Iq2bzV>V)A$My2twYw&6}}Whe*68nOFXS9A5-r*WS>sK%%Mv3(QQw26LWztn;{ z!%{ACX&)3CB-W7~Uc-88GMt3Y5n-}#BPae;44P2kZl3<_jlPiG46>wv(nUM}JC@9F zTans=gmbRSn@s4H7Fj{QRkPT~*;_`TfqgFC(jfZO42%IN40$igt<$3?yeV~xtW@%1 z_!)v7oZkm{xy_q$tn2X3^-QUSF8kBQHrf(6sJiQ;=b##M;Vv!{oS8*bHXHm9!uS;&(ex9rH~ z#sp^z(Z^F40IE&nXPo^y8M@p0|NT&{wPZ}~+SUq20ry|WvKX?7zCz~B(?*@*7!sqD?_0QjIz^`g=r+3QV0ZKmUsrb;nF>RWfltX9NSfoiJave0V- z6bdAMwW<;Rr$GbtXO*vP|Db{EC176hhU6Q=+i;%|F@Ck8V!Y5LA)3+D{cyU)Uw`;` zrW9y^dVT~g^Z;HqLdPkEt^@bc_YW?9RQVRxHX%&nnTKP4UBHNC`puLkC+t@&I~8caeG74oPat0c9(*#(`hZ1R*v%hH*50Tm zpMcyHaL0GvDpngG5B_n!JcWYL^tk2}KZU}uLs1XWQYj+-ciEcUDnJtoT6UVEKu93Gj!U+;lzcd1+ z`9({N_-$`o;-q*KL7aCN3RNi@Xc5YHmHMDD#NgX4)YZngL~$0>wj}Ijm}NPu3Pn;g z7zHXW)b9je8FN$#j~YHfXwdwKf9q^PqQ`coF>7288O_%HGCaA;JbW?PV51iC;KKR5 z$y`3*b#E_1-0|uMpU1@)6m^!eMMctxtvzwstTHZB4L)ir@#9K})V1a~d7vaABze)J z-@AO3kgp$KP9IFy&PV3}>S;vC=WRR$;<7t;9PHnSM*lGm%blrq`^(KHR)Bqg83`yh zB672z;j=IlErRyWaXf6!`}or(4`W0BwQc=bK&OiU76hesL_#O&HrC~l2@h70Pz=JM z2l~g*>dC;2#o|J4JZ=so@NGdee*(5uL=U56*3d@AM5Mk>}*P zEF;BdCvrm@J3AX3+p!!w61@P;X2d8gm&gSyj)-D^f-0ttGaX$<&KnI~erS#9_RztaIM})lxXmW55EdY8GNGM`Vmyy5~vNzxgYDpTJ;m8d5g*N2Wiq75Bu^ z11d`p1kaw*h9{YX?E)+!VoPl2n>EvNmL_Tksa?}p66)UBSLXl!;~C|=qxHR>gbrt6 zGk*eG!UUTo`YS8kf<%rmwq+KjPT7aXoUT7hv&-##jwOYjP!$YV@?<+{BP65zaHKvy zQae4OFv>h2t{2Bo&4Y$48MCR38x}2rk-W-g}Z7I+hl*)2{^9}cG{N<&}#KK|slf{>mrfQ9)@q*j+lvuB z#QO?l5=X+lESMj+qjokwHFhS#M6sxIe5$A;PR$;y4)?b45oc?D7Gx4$BAmP{e=!jl)==gCg5K>dHm5^KM;xt#{3?f zun8;Jzz1uK%VvF-iVDB|!z=n4pkM=z+&Z3hIoGEWvkhbo5TyDsLy6ItD5Qo21RMhd zRk}GR@Ennbe>?T9w3ZT%6gm}*C>G#w>S*>D!1ZR)hHLic3!dTB+5>rnETTIwn4EWj z_o9OrJP6U3X-R#iwo*APBcDIWzZMClANGb!a!aToZH;y^F($9owRAwq?@(h$6iQB`vlG8x>FUR zEN1F{ZHco1-s5pU**HTi{!evpfB1*0YR==#`dgGq(ttsUbbR>rw1wRPvk_3%iS`1i z*Z!~-j)G;9H8lB)pkssB1deESaMO$kHW-SoTju@F?!BK|7(4z3Du)^}xzVrVWtXXE zh^XfyWZy5}rw{g1IM>w9DKxEI=?Y1ViqFhhOz3E^lLsmI2;cn3?->n%n%bhLZL zRxv2B>BmMI6n@52plaQV`^R=(4bWC{lj_W;FVaBh4W*W9zcN|pzYwIZI=XO#hiI$+ zHg)yHQJY?L{xP3;-G9wvRt~=WAc!Wi~QA7bzCdA9!> z7p5C6?jn!S*%A?XEPo?|GNzf8^<)kb9?BDcT_(#Y}6I^5ZPx{RTGrxH{u6*Al ze%k~=4e~g!I3Gj6$g5EHslJ7~{bJ?tM_~uKS~PVTXv9gP9l*T#ydLZO?Z&HLyiTo6 z#?gOBGlylL8ug+qj(fig3~j<0p0-1GxQs!g2k_u1aT7&@Oj$V3)`O6<^$`X=b@9YH zssz9b9ZrdeiK+=!3JKyE3fS%pY6;e!&})~M`7U@0y)*^+lB#30quE7KgqJoAq>SBT z(hiG4nCd0Aq^$<3<_5V_xPD#pkU!Ohb2G54l2`zuU&%k-fQ&Y1hzi#=^Zm!2=M}3K zyY`N?Y5MnXKM}G(jvX(;EcXxaI;@@}R0nu|*3(|@^ex*}Pn9!$bNZ;2-;84=(lMP} zW}lG;vmqnCZB>%8r>y@U;-FLU2f{5aj5PhLG-SiL@b2qtQnZG~ z6ZoAlWz>XzpmExqoQvE$7w2Wky0W)M?_DUT&t{ zo6af^&+6_ZA1$JGfO7(JsXCX|B#iTu&n4PQrAVIgZ=bD@m+9J@L42!bca@B76jL7| zzi}eCFUYlC^1Y3(EH5wLi{*>G^i7lF&XXl6^W@z%2IEWhve0vV#snTvL%g{d@q~SE z3Kldp_i6xISW&T=`geLNt6E^G=r^jg_YRpdc6<^zzol^s$IuWgMZc>R4^Q~oyFNktKU15KKZuzWlGlGsutA1Jd6*~_{4@;5iAl}$ z69iIx>xVC;*S8~Ak019(Mm`GFf*o{MOb}W({65fObl&5-_e|?CNXv}q*?xV#H>-se zRJ3da-s}5b^l+`Y2jTzV_IhcW7aD=74b2g&HcUcXvyLw?bzz=4l|0*9OSe)tG+02* z=@&;}Z;(T_kZitb{Dxc=zKkH&osh>pe|AJV_Iv0;|e~NHPkt# z;OprA^RtUeu(JO%+o!+;KAYvi-Eb9-Z!gxJ66ShA>GF6IF{VUxz`4lsvBb=lCS`MI zZgox;V;F+z!67-eVFwP<4lFz24)tdDoalXMFJ#;Gk6ujzk9oPVZOkQu4U-uc&4~7r z$k23B0X=Sq=}Z`i7HMk^P+UtBz=Wb1!s1bK;L7w%RK?DWuU z=5A)qn}pgZUtY6usbfw@+&$D-yeMhCKwc@O3;d;HlOUq$Kvh5^rP@FI&3hmOu&>ZV z0RAyeQuqMXihug)<(K-%8HQk4*S4-;*+(KsLYA^=i|m@q#-RPX-qp+xV}WH47Mr~G zU!`=0QQ-)_oJVz65b`gCeD4xHYTA{e`p!p)w{d@?XGDaXKw#8sLW_sU9;e@;Qy(RJ z=^L0xMOF9^a{04+XT_j5$s$dev9P%i5bC)DD^fHS7dd&;haj;{^a+lHcRLe{mroai ztQ>$US96#g@jCwXdoYB&{_>)!J{Nwng`g0BFP}j4{zBUQm$BU_JxNBRx%l^x|A)P| z3X7{-qD2FN0Kp0F?i$?P-Q5%1-6gm~aBJKlcw@mKxI2xzdvM5Ye%X7U``y?3c3#)& zZ_U-UYL2QJHRqV^Pw!^v%Z>)tExoR6)+oYHNNFh%jG+8}_To_m;#l3j%^z8058<37T<9#;1 z_Cwe8DA<0kN_Agj^=a+59U-EJHVz`|3n>eQvFpW#@qXn}q!Wb+68nJ-hHe{I#>J+o zrt=1`BeBED@{~;;Z$>5tv;>405Wao40xzD1nog9!+8ZAo*}PwQ7_n+;@;yRDkZEx zPiDWzo{(W{2p~85+B-}rH&M8G$$h%4Es}u5d~8UTRWVtMKS=A48GaIt1wLR>8M@v1 zYDo=VGO!5d@v$C1P0SP+`q_!Qg6%2vmsv}qg~U5c5Y$mI|Gu)40o%Z3!KEX{R_JUO z<(J8=Qn8@rymW2PLUM-N(}a2aaf@DmZyK^vNq+VY*wk)2vF01|$K^?U1yceEzLOd{ zxZ~jjCQ!`lCY9;4O~N&ho25pS&vWPji$;gbZFYw4(-mmVyPL}D{(LsK*yYP*y zVz=#~18Y}8!#^ebxTUP+6(wl^S-GQsp_xKr`lJRp(tDINX@jC0am@mq&6zJM^Rv-T`& zH>DRgx2oxntiF!7fM7mwJGIiyy=EOmte@{=2vRaP32*cBsC_72!aIDJ+V?L!LVKTj zds|n6oVK9+=7Mh~%i@56l+UM zk7ofhr1uPr)%04 zL7mh4M33>)r%~{9Fm_o2ZePxCZ8>d0d-Xa|ex!?0g99l|=PPcn(EIv6Llkr~tmiGS zzUSV5e!KC4h*ujLvmd(xO(n6U&Hdo5>f>Y19lAco6y`h;Fk5y`Bvl)DR$xWbZuEX3 zx!Jtcap~1#xC9<((|C<^O%1BF0Lgoy-#fc8Fu!I zkIZ$j^NJghPWm7J(5FPo-ko8%o~f)Zl~=2)!tw2fUj|3bMul2&$1(uzWwezro*XB} zZEIa&doKAn*g*ntuF9r$rl}5IW_E7l)M{$HYxYy$3fj%`VFPhtf`Mst%!b#;!FYGFE>+k5=Yxdx&$ioa&L-0Z@l?J0 zh_>@K50|0k^RtI+Rml|0lw!t3_rUPW7tS@CrW?TZMCZm9JVWtWn!Rr9zHgYP+3rQ& zdlPP^79J*~Za_X*dQ-Glvn+JdMJE|Jd17w6Z|1ANyznEq(Pgd8dEg}W0W;!k^*GE0 zI+Cv5FWl0~>|DUQK$ew(aE*Im?KRr~uqx2`lltnb+^0sD&DvhglZ;M}7GFaz}HiL}_c+DETQz}L{s$0nEJ2i9BuQ+9Y;OZQIdNhg&_ zBPs~Si6|E!z&7V*O{AA4=+bv~eA?!I&h78@daxOCh|y{O)BjHo;ndlbZG9Lf)*kjx zQ@?%S!>y^*2+-$=5V%LsIzMD)ra|xzGKX||K+!%D#L4dd1I?9^LX}(b>pkT_=f)z6 ztWS!2S+F?r`^;$1`6o*iXo>4bd^p^38pPH&btP20x)Qc@fZk0-RljEZIXz=eh`@e> z?ke@M-Jmd=OMNb&(W>kDNhP+2*#J`a(esMidI!(XKg1@RSqr`(iGyt33En>{$6akN zQa^||Gpt9dSySTrOM5kQmx=PWYD($+pi&()pCjH`Ai+J{H$l&t4)pcz7&r}4XvAIT zRRX1kvsy6U>&g`0-pGNVt5|frVEJn)7V-k=8-Q_|$l+tiHM@-~w=>eJs?Gz;vAf6M zS=3mz1c+p?&g5ZcHh zvv#7|8cVFHn&8$3rlaR&lGr<*YI!Z+K7VR$>csR&E#60>$6xi|i0ATfG;<64Iuyky zFv0X-;B+|Kop~ApMDsM-bnw>wz1VX`Hh{!LpeNAkvDtBb?sWK-3w9njW?R1o>f&~r zO`ox?Jy?k)xUN5M@9cIQ( zYpQn8YOQ(GWE0H^gqTJ^2?(&6obD-SRt(zuSqEJ8uq=XpKIM!nZkn(K)z>qCWI*jz zKtpaBHJWx?FDuaOl{{!KHJ0$2PM%G$roflN+OL+z?LsK3%eE{gm+E!ZRd zj1QLsjr_H`-h+0Lai3AHB5s!VW3ezR1?#eh)?^yl$_I;JmEJ9Bcunw5lrsUG=#9=m z?f_y9@sWI?L^z754l4pnCS;=D-BN5 zxzsfW`GdA%Y98Jh?S%T=DDHaR+cBGU-g?vZo2jxQGI)V*pWe;PRABp`N7Jy0KV45u zZ3C-S13L7pj<0LJSo#+ooo$Y8R(|`>gq_80KiS$FxEkwD5~pO>_AzET4wmmqeTG_m zQJ2r}BJa#J_D-nR+RSsLzUy&^*#X}_)?6rMAo04N@%-t5XLvO5x-8m1<+nFBvosry z`+%r?*1$hfKjPQ25H;-Atw{}n9J7pPVtiP0maXagqn-WRuhMpX-h-yDqA}C;>6WZ- z)}-0hPSBFib(2E!49ML1nPE8Fz+(MwnH$8GrLT<+BI5DcSe(9E`Pr1PAIr!<$mip7 zC}pi}Jvg4?of$#*-tM|r#|86Cm=qU7T{HD9=4sRF(o8*ovXXubc=1|HMufJ`Xb)Ov zTLP*=3Gl9edxpPT+rH=)ox!RMKZ7&E$h15wSnwO8;ft`7R4t7CE4_J(*D0u_2X`qf zw2^;bfq)I?L~{yvCKv03Ocv|xmVBjFA{pm?5GAIU!!6=TS2!FxPAmgEy2EesAt@`h zEh*sO@eW3YvaLRk*G-hd;d3lOif)MiR#z{4*h}lglfVT4%ulaM?x&o&R(qfF_^)GV z>eX65A#tr-1R|G5jT%$c&rN0XUS~?cuNCt1J-%O>vEGML>w-c3y2(N49hY@a?UobhvN60 zkSz=D^*Ry)rJNBQDNN{_5j^#t`*EhSN-pO&O|T?hbsw*v`pmonr&hf?sga*v3QPxX zneeSJbk&_eM4dl(qU}fOq_xN2d9|TY&A7kAX`vx%o z=)&uBZ=#~2>6#>B7|v}N0AJ(v!Bi3+n5@bED6{JKAV<5_>*wpZ+K z13G*fj0#JAPGJ^zJ6WW=w;`SYYI+p8=lH*_bik(o&kfSj4FMUMOD*(>n$&D>aTCcPq_G+=pz4U$|>v4^(t&wD8%t`}Qo5_)+Of_7{NOL*JA6%aU^}_qqO) z+rsU$(#8Fcx*wA268}m3KY|0VX16ldKS@thcmS+qK(O58=T;OdJLMgE6p25>hDP*l} z?bz-Ptg}lm^4JJ!Zv!z2N{7Qiryn#u1d`b=j!%F#>1XJ=!;tXR8oc-2Dw^pAmKCp- z2ckf)ct{27jpm%VOm{r-jnnD31(NgL1lPQv<+>2Vf!d{8(e7Twok)D$T!EhhD_gf8 z4kt;`82pQ%+Xrkq2|1i(X9jp~Ph2{9eMgJp_`PjvJsMt!RY)~DOUVfhY^v#)#_5~* zf2o2%QDdpfAkNa%<)en|S^1&O#92DXht2&E@gB)AWI^PdY6DC>YsgDxFEso?9R+VG zr!8wO2I*^Ha2wZ2e$LXAz_6)yzrYJM-Qxup;^~dv<$n5xCdXhoQxp@MULRU|`uX2%rp%>?c zeesR;3&vputYbwR<^6%Zzf4`f4m*aAF1nuVWfqw)$L@*3VA_eaoKM4x#i@dH!`rMS zqZCSF8!c`ZMV^A^PJlqB1*R8f7)95(>PjR*xLE;cN)CR)u%`88*Ido1DMDHVY@p!U zF@M5&x;|x=S|6^(5ij_*J#5L-3)IMuOZMq?tpXihhfYx;|*pLvPI>Q zXwd8L+jH2T$@G8Q6gy6+>Lqrd$myS?MFAB8J}ttbu5SVIM-c_;@+H$dWIZUsY>-8( zxs2Q}Q4x1?Lzg9D(+>Hb91ql1uY#OcEaNvgIS{YS#Txs}R#$%+*|?;+&Z0X)8<3!> zzchunDrx+5b*h=d53hp>Bea8=g5=f>ldqgz?KgU^Zy>ki?BH?}^IVbZ>X*>mIfB@- zL#6D}rTi_DLV(>=r9=IY6;YecQ0z_mAtpmq0f^@aWz5dLEg~ecYTn|ejX5|j!p}VP z7ks-E-b>f~exrg!Cp_eQ@mBePW+?lzL{2qZqQjP!-fxHPFE_{T>G4K;KW7zGBm6Xx z2;b4pWxLziA74ZSRge)i;k(LIwk%6c97dBp){OO{FuhMh+M(~PXKiY(t=Hp$ z6R3l3ZfCsl;B4`#CR4Mat>JS!8&w?$=&VRbkRtV|4T$CP5?XbaUu*jm5@-Lx-f5Yx z9ab%Nr5grrj=EdI8O%2(?Qv-u}_)Mt7;$6w|nrTAkGy9Z% zYNo52vV)6VQm>3nl(3?dn9p|stoL@`W6WT0!zheNtfN3Img2`Xi1JT$vjJ9#?DWX+ zbi1(9*-tFY_RpsF@TGXi25WB2IX|UbGAw+;{b2Uwh90jYCOb8;NkMwrQ^{Q8+|Oax zHTVD&4=c~`hSZnX@72bhFYfn)^=lL?Wi4%`GXLQlo5I1IMH8$t>OWU(f&=a)@6Gen ze7=)KbzMnTRJGtU$UI`g&8y=Oy&!#I_>68P{RGNyaU=Y;K%!xzd~})S3_BExv9W_N zWfVo-?>t(|fJK+-4>~UH1?(4Xwr|ku4omEj8OmGe4iI2cGK3)EcuqA|jR0Vks-`(; z-ZMv0T`7VU+Fl>j)Cw$V*&;t()`dGPD)Whjs)9p+&TkYtG}+$3ub|Xu&}H-TIXW^2 z$19i|FuAWXIwu{mpDH@nM=skg@gwm^UruOTpDN{2`Cwf+z>&5nF`HFwfP5646~^V` zJu_uJGLBiT>(wGJ!pCmGGgCvte)34@7B}%@RO|A$+E-X#@!cR>9V_O>-fz9-E^R*w z{!UVIQNR!*-ns2qw(?J|s46M6@_nualM$)2pZ1Cym_%J69*4Bz1n9#|lomI==9Ko! z8ps{F{1m~RBz@y7rJA&Us;y&_uyZwA)7sYlk}>tYXb3pS_tULyaXLFFSjd4Hs2#&q zspHipjIhL(CLivHmJ;j=!MxHJ{5~hpFZ`9+l6#X?EE?`LG}=s@S1gq(e2Q0REy5Yp z&F@qw_3Qa$=q~<~VzGtg?!7%^8@_GcA_mZw=+2>8j}I8>51SQ-Z)-1K9qXryJe0vF zix=9<^iGz`Z7j;iefe6>YLuH%fVoUt@;_OCG??jFB{vc`_`4vl!3`w*B)wReu~C)K zkWO zEbY07>M@CY>Y6CbHVz6H{z~Uz_dy$?bgPl}rG-_Cz3?=WzP7I3DmY3 z(IZ*xq*wf@%)|mA<*`z0Bq+Vr`Lrq44D>Yb_Nn3A(%jp-Ifh;YS?}n-`~=z?KLMwe z@E1R!1Y0R#1X~dcV^2!t7)8OG{?M1$N&3rAV@|ib0tdfSgq6czWsltJfoAx@1DgX4 zXq=Wm-eM?obA#gb@Z~DEQ_~Do`z7;ibo@bUwj+q$XJa|6WOkWF(5Ihv6xyx#V@>^f zkL@AJ=h1Te{Pq5A)H4(YcXfrzkUh{8*;0iIQd81&Pm!Wb;^)r_?X1DX?j-wsPtri>F;kGvYmk5rLjXVnH}uib!RV=hKA|?JYQQFT)o#rnzH;w z?B=3ysX|eTTPNOsOr*q6_xpF&ymiHn-`<3uGFkO+OE_Sn99UO*#7LB9#yC-xeU4E4 z5|ct%a?C0dU^AV;`Y4@0uot(?m0~2nR~l*^X*V|_?FC@)rMx_s5Ld$DjP##Ks)26R zEED^83w|pE!D8d`kX}{UO_*_I3oIC2FXAi7PaIfi`ZMNpF2ZgDdM0>Qv$N@0gR| z|I{R^9P!WRu8;1@=yDp{#Wl{1;dmLJz(sm-Y9}^U?=WGGUK)>FLtF0=C3^xqPB3-V8vWdD5QRbmES}upkCN{2zOR+d&vm zR;T|j13hwUu0p z|7PtU1D&zJGk#(fNH_K$+y3kRZ|xBU4DfA5jI#WH8YayJBsFI3?qj zl${~eQH5p`($-xoL7P_+O$=^1vepKkWr*ONb`mkc$OWBl6=4C$H|&z<*cvIC`j z{p*e}@9MeQ*OG+t=qjX4U^S&tIM1zdMs5T^eFBLBY+|IdgOZ@$k? z=&JrN3jF5|_~$hkZ+Gc{R!jc@(f{K%?pxitt*rU~`D9&G@ZJAA@PF<4e@*>gTl`<= z`VUzAGg1GV$p3db{SS8i9{@56VqE)I3*di#^AGa;|2nl=sw~w07XX`(A;BBMnxZxi z4-Y}&c5wf+S-8SbN*mXmOKWB3{x`_{6Y{}m|JF^If%@@3apFIp9vlKi-_Rb3U8Bsa z92TG7%f-w=A-CZ-R%jkZuorFmpDflYDE*xn z{311K9-NXr73$EBr8Yt+jH^eb-~OU$<4Xz^E=luLL;q(oPjsAi?!xCV?#-7*G-1*%d9o#R}u#p zu>g0LV<$j%_UX|j_4W1jixXSkpM6a1qWxe!CruF07HOAG8L}_P9THk);6ILFWGvu(C6~E-R!yRI*vjOBl%E;!uf*pIVKwh|>|^+4m_@oD zGt{~~xK+p&9MqQ-m_HM3}lJo*o z4-rY7Ar}q(+(DH%bjcY0D66=-d_nEnZ_9WV0y@_6t-Pth&Dc%`S@6JdPa~c~=>uGi z%+?bp?2?jhvZXAv5ZMA;8Ou=168r^1C|KM)VP_IB~eoxwEN};q9I>X0}_)!H+->mBm!+;K}R1&Xzgk3&JWb&3nF^%Ci2;l^xF}~(l4}> zW2PLcdIrNKmnf9N*%7y>gnBzlGZ}H0CU-vEc^VYe?2L8NPQ_M`XXRAupjtxEhP=Jk zOx){1T*u3)J*Vk~%dM)wuByQKwLn~3Br#I7peX(2`JdNw@VD?Ek`}y9pZK{IN!M{B z+j&!7@k57+jc_0MkqK#C^2l z%dBD}PNvsV$tq4PWl}3b)6}!fU)%vRTp)ss3?lI_iil`&97W2x*56lc0;nuJ%nE2$ za7n;WdDJK}$7(3AhiuUg24p+r0_GGs=mngIaJ@Wf`Z20WJ!O`$&7o)IuWu(J z8X>;!{PWa_j!6U-fu3{hHg8vkYFBGCtFR~U0+8nijRs@0<4#vRIrnbn1&xj3Zh__b zw);;}U_l-ivo8f;836y+eSb`m&P_!7`_YN#Q7VDU2+#2j!mQ!9l9^TV4M_%SxK;vD z!b{xujPP5s&>yEWt0PUfBerE8SReRs1pCZQotoNKPgl?V?%LLVGo~&W`oFj?=%;$1 zQPip370#;~0M5Wm-iVNGty2#B3kZLgN>5%?b5+O&4(?8|?&9i(R9ciqik%#J-+tB% zzhn1A%Qy5sZLIeDUBRI3+}h*HC)#@$yGxWqFg{+mle?q!IQ&-y)kO^v9F6Tu%e0Tg zZug}DAEbq|6Z7krT=$%6&ZC;48|+u)uISNFLwHA=(t#`Pw)9JS{nLP-+Nh2&Ruq$0 zfbqWs3&H;y0lA;aqv~eC|7C&bX>q_w7wY0p~zi>f=&t?CH6D>F<*&<<&{HNYtTHN z0e47OODvcL*!O}@d9483{|{IvhzT%VA&if_FwX7BTydbF=O73;chWp);U8D*sjFJB ze>@2Uy_Q!$(vEccpdsJA$z!lrdye%IyP5<)`rJa@ z1!pz-Uyoya9_jz)Tqg(3C3{`dSlG+1DO1ZAsWHO#?K^^66l8?{Tv6u6%8YM%<*x^K zzkT-ppj+Fimnw65oiT68r#nF7nmfu-rxFh!Du3{2#n#zn3 zTsDB+X(=!Tzvj&OwsGt|M|5A`1SSW-I=<&%7rDFCa7EY_dj(6Q?#|B(euR&wHd9m7 z+Wi>_XS(?rpqjjg{fqa)4!9~K`m3B$a>MfsEITQ|w(H z70Y)DD?;;46=){Xiw*AH6#DWjF0?#M))mp6BhlTPzNOR7AMi9!?4J(fKQV50&8V&9 zNG{8J+d{>muyPG)x4J@Xe5xfNH_2mBTKMDPSBdX(Nmk7f zo~j3Z$-wVnH-gIZb3H}ebgdV7$=#5<$<><^QnHwa43npBqcaGC+qJ(b7z7l3Vy|&a z7Au^8_)vWvWf(O*Df8?31h|mSxsmpvo%-_tmgyVT?6?A}d%1HfL&-@fLj6WrYPJNFhRAi zT&YhU+FsPUDWu|;u1pigHjfy)iW~aifDyB~q;&(96kW=B9?BUNy@bZc!973pc`WgN zG)uP^TJNUfIg93fttEP5xy<`&ZokERTG9Df(Fm{hJ4L#9_mG*Qg*?r{oKSkPZwIyh z-6BQ7ecQyJ=LvIY4f6^^&*w|9$X2!<>AlP*b<+sz5xHsAJZNfiGOts4uOF=6Axe0K zayyNGa^y6xQNEiNCKNDEbcsromlS>h^$kZn*sFoT=++A3o6FByC%SrYSRO|Sl#eSn zr>-u?f^4FN1T0+P@TEv2oA*LE_pU?-(S09Jf<*LBF+#&xKa7^3q$0;mR_M1F!O48& zdl*to+1MOCZI(X2xw3^-`UvD_-P?rfYJOm3%URI>?1;n~BnP`=Dw4uAf30u8dL1x1 z6YsL7h5^TiS~|iuT~?w|9`(!c&*MU-!83TTT3)tyl5xSJE{3aRQNHAqUfHC+MylpRj&xQ$l;-<6(^D{Vj+D}NP9jbT_okJy} z&qxQ@{|X{;Aiq)*7q^U&N=fa#_G=#clt*-HOtkmO0w8&So$B#l4J;WR5Aev&YHszv4OQ@Gp<^RC=j?J5lY!6XzO>r z>*m3nKaV-iG>#IA8gvN`$8mJYbI}W^O(7<`w$KKz(j!V_eW*|96XDRKiXwN6wp4Vz zTZ-R*16ll$_7>;Cjfja`jI~{cz9b`hkV;d^gp6H(YG41fuWPxIHt{%3eS(jQR=7h? zS>vto`(e(xZUB`2ynjyO+5Wfv8j9nm6y6huk~HLeftKk>VRey|Y*z3c4X zz2ILVq%!u|H$gm}cGA0|aGWM_82>8aGv|pruHF&!?xx(?lIXWG(QYCJ5Yv?OXRYGl zFkLk|sPybYKK4y|IQ+fnc@Sw<`;i&;W$t3@D7ROIxR?w#_aXIRG;&ig@}byF*zqnd zznO-t>!%RI$0EZk;;cl8bAHB~L3EP%UPWCT4t)o$tHvSujv=U=7;c~iK3`Mh#&?`5 z;9|OddTb5+?0D|A7sidy5gK~D|F|>2bQ~gHde2-S3@!SQHa##8E7<7%O~UjoQL?lD zviyxc%xo^48d`)WPSn|Fi^?y5>4QMzG*Jz~DfS^8S3v@xa9Vq4A-y*&_XZ;4rsWEp zZ?rKR%hG?Y^695D*!=9f)-ZidP_L*nG`qDdC2Ol8%^7Pa6dZ$Rz{Z^CZr+l-Kv1X?@BC`k|=<3uH&YHXGvEw3ODY3WBdY?#OWc4 zx#Ea)9$>7L8m;_~9~CBQ6v`#8p>d?47J9weY`Cd-Y_!-eonUigp}XCM0c<2?yr` zh(bSmD}>u$((nBQ2NrYAc;y62K>0lvDbxHwpzaO|WVsBpdD2>uSy2gP0mtf@Q)-*m zaA@H^JR@94XsY81`>lkw8L|w?d>cZt+%Xx)-a7<=E5s?kCgiX+ikhv|pxaVPet9D5 zuL8uXJ=$g-ZgqfMzL|U2<$Q<){6@7Q#P9X8P}Hl*N5z{;GBc_ocg{5x zaAvTK{I<24$8QO7ZGX&@$J#Z(Tq3XSW7R1_6UVr31z~7{xSr+}8(PCO&kY;4mt0&X zGj=o-KUhnAKQ3G}sNR?%sf47c9365W1#-K6kNb;InGS3;X&2sQELcplWe7#l)Mvlz zUJXf0K+`H||D-*MZYrD8+=$jpR{fdLi*GEDiwY^|$#F>E7hJZqRP{4rqzf(PeOc~n zd6X)U{V{_74P7`E2S_VM6(5dw?DhRvDYCOAoYt)%J)F28V!@?Mn9){!KgOn&HC^r1)r;lp8~GYQ-9?d6zqDPLXWnjj8h_Kf>mDNPu3mZyKx5 zVCBW+vbe*I_kkaH?H8|nenx^wWGbxbin}>-|5LWFnP2VA#&_HX&DNqa! z-Gc7+N2oktzDujsd{=Y+3{&;0T=NTUOX6)3p@7)}!(lMNjE zF#y=OgonR)_EZr*rHTdUL)JLq19m|=6dmUvtwlXhLGvD-98Lqha(`iLDe4G$jtQhDRZg()H~H~d7>NEeciqDv?^Nj z^8$g^epw5BfbnkBf|B-8S>?`Z>ye>4T(_L3Dz!p2KyDfmgk_#IDvCjY5h}F*X1 zfh!g6w=FdLVywyEOCkesLN~5;J9c%skOiV8H)V>z`h5PA3Yym+)|RazIK5-8t}b__ z5FSZCRZ}^#W+B4R2zG=`F#nG5XdI2dw@g$}SYDCp7>Y?wuk3d6IN+b-;M`KUmxXdn zM03=+{4#i8wkM;jvqWbRh8^PRF1XVrovG3X1#iLV{kg=_%+SoJWd~&e&I4d?7;Cp z0J$L@r&qX8ZALCbeQ+X43ZtshlK$sgX#lJl_z6>f$el(l9&McQu1K_aJl8qWj{#kE z8_WQeVs(v*jd?%XeE1~kr3|&Yj0l0r-#1`1)>WFPsRYt^sEl1(p0?+IB~Yo6wMIf@ z(3xc0`sDO=@590uM&k~euQhW-vodt1nI(jFR?n+m`VfaRX{yta<7(e?x%beJe5H}Q zRTKJj*B>?Fx+)b_3XYi!dD!@!eIbP{ctqnaqD{qW6evI4z`zOB&gO- zQUU;Np21B9`S>>u*A(PIB_Ig%XGAY!LSGdYOLG|Eow@`Q9;=3L+(OZ2Xfo0zE%0Rs*)W^C6vD)%|(zU zLE=UsMU3s>*$S0!PsyMh*cxa^@~bP!>^mG1qf=&y_@HQi7A1~}Yex_nxGAOnB&u*( z>AXk6v^tp2x4#vFb#{Wz9gMJn#)VDyCQ9Ug-#nVGk;Xfn@QyUUZ7;*x2TOEzK%aCb zhdbNwx!(|7ciC$!>y5&^<6&yN&3%{)Gl@nFZ3?U{zDK}7Yar_v+O5Jm_3(C^OF~8$?s)et<<)M zOp?~it|sse(&4eeRiud7+F>*wz4+xk#*vAPtcP%_DS$OK_BES@B^Vw`{FySQOQ{Kk%!&v{9I6FgViqwNQTH z-~jNedcBAx8A-wtIm(7_>}JVVl0M3KO7g@X$hEc@5e|@QXl9uFEFh&F68`1NqH1P_ z%^|dGBTa_{Fx}HTE$q?22)%YQLpHJD3XkXh2_`4H`x3@4-j3>Q!Ym&ywbP|UB<1Ll zKTcbDx-Fh#7<^}Vd`yTC_T44J^!w(0z$y+AoHKzT=hSd{iE;pfA`IjTIXbuLkHg)A zFf+^aN39Pm@U??`WTlv&T*#-%bX;K&sVtq`uFI-S({^>g8DUL5Ci{8+knFs0v96K& zWDT8rc#$!&N-`#ADL&l_Q#T`|E`AUa1U^`%p$?dmj!7=XBf?2R%2V+TPeaXh+jRV6 zDw;BAkW8&!r;%nayr2e3KzT0Qn9)w0(H#nUF!EKIxbdJsR+jFg-Jo8+JDJqM=`RU= zDX*72p3fDR8ujdfvIQwIv|v`VT{1$_9;RGH2nhwDUht`2kZ!L8C>zEbPQRhlOHrDA z7#cu4g#;uRRYv8H@e~?=N>3rhp$0gTHM2kqeiN-}c2@qV=0G;}Lm#z+O*@!}Zf}h^k&b|+~c7tW) z4qaJ*O*vifkCnM@z6v35xQmd8NCF!-+&S0WPR!8^S4*+Zv+v<|W2h9pmuDL4!WRtS zZg5c@t`&l3Aq?!h`ta^_my_iO8?&f>|& zwX{I6_n;_2O^G4V&7vJ0*A}>f<4*YMWg$1}0gYMwxjW2`r`v?D%ex$+PtVQ+3;ON^ zuSY$J^b^GN6IV)3E0WBcFZ`g$Q&bXu=?CSi`NF_1p9y)3BK^`5XK(>~W1K1)ON$$V z4yQ8jP&y)F6!o^c71Mf>w~ZVB{T=HO;XMXm;^2G(l(h{Rl%%{iDTogANk5<%b!sM! z9g7o+4D_UARhSz?ZGo{HA_XvGn=3kVGaG0KOpWiEKT{BJ@x?BY4~_H0T~BQgo#%g4 z{M53?f=QhOCgCV~c1~0TG|-8&O|bJ*Y2?)(o#P2}-p&#&=Q`;xdIu>E_@ZljFehDD z6Lt(#ydJ)Es5GK~R9iPO2do*(TrMwk%x?wB8^+&=3B6~*AFI@%pLmU_U8+ zRG9827L^J39;C>ajmX;pD8ih)n3VwNe7KEw&DQT>kV8q_MLP+9?Bi*BeMG1W_|UL`;DT zb5TdPDDuwSZ8xu~ZHT1gl7Mi4xMrUA6-wj=x$rPcza7w}LG~Te7+6tcVQLB|5{)n$ zi#xzp%j7vEOHE1DNKbqm1}MBBhEx{6gf9#bagB4e zL50D}PF7iTwhiy!9*MI%h*+D=%OAMOg-73DFEl!{Y|bi$Nl7Di#3kh@QbYywJG@{* zZ|zK8>0{PcsIA)QZM`H`tRC*mMzcurp`pqUs%_I9YS;1SWU7q{B`4(1`XBA;XS;IB zN@al4o*&`_arf!Z5{^;kv|@zQwV_M6j&P?EM=&lq!^uQ2(#yHX_HOa{;1ky=~jUbU48dnt{{S)IRUREa5v&HK8UWx>frellGE{M5{hE`0>!_#Lv z!Ibs&oVBxEKw9Vrhm_AN5Ipg%P)5bp$HG!fPPa%vbwd!Ku^JHHSjRX1r-$iTq%(*@ zRdy+@1a8`j&?@b&l=^l)yKlnmdH$Ilzu%(t;U}C;WRO2D#m7R_1v+O;=f#SsBgJmUgY~9ENe7bbTIwV(8D_mJ|ua>y@ErZR5hES+$?-cIGAC$vHU+w7aR% z;hKzpT+us{6^k!LcShXzvk^NZ*0&t^8+ckn+|A%#K*J`jad5?y|WQs$>g~dHJDpFa9F&a+q zVR{N~h02-x9WfxSK-ebgYEIKoo(>|Kok#Rw!FP25M%!zaXoaOYbg%%bAj5Z!kYI1}bok-DNl`a_knyg5~~ z4`_jLS_$9pl?-UG!7{-zd_@Hu_KU)JbHca~l#rVGgGcRwntLB!RB6VIBJ{^+iz8t5 zC?~&9T-;#N)}o4X-{ge`zo*cJ$~xMLmy1y#vW$HPiFO1TII@90`@2~zOc|>P2Obqb z+Ak4lyc=z+lq~RDG!B{!P#A#jFWg9B!PAFQxrp|K^k6m}3*2aDBn2aL%c51$wHQ=0 zwP_Oz0kq`Tq&P!1(+y}0ZeDbpW7S{@zpN~u^mIeg4NDYf!Pq zRG`WZZmyBvOR1v_JPpU$ZK4pEkR?Wlw){qVYGq*eAnZQcsGMs+Vr$6IMorQaK{FxF z4u0pOwEV#`wb|%BMW`ph*lgr5QM%{f~;BYkw!mxc(!R^K7<&{DB+Fxsw1xi#NB4#$tbtjj7^nUeZwq2L zT4SB7iWIUOQn+Uav}_OfNJA=~5Zeu9QxM^@52t7;9_t_!C$}1HlkFr=YO@-VauUqu zI~yZl9qt?<2q1}%1%GPdjlsS!83G5}K!Xm6usJPYK{|^Pl3G8ZmM2m7*dWi7Lg!&- z$57R&K=O9Vz>@7ZTeW=Y^?Fegmo5~#YHbNGKo)L}a_gmf-%>i#;@CKH)_hLFbLi}-#C8o>Xh4Yd8APbYN*n{sj$_`iZ6dt zz8Z3(9JKDjL3{a<5Ek%3v9~}O{i@4F=pBxU>u=k^wyNAfyh@wGI(36TB#<~+jqt zaZFo`%t4LhY)yX0LuJz|j%eYR35v#Ob1GsK2$s9=WY2#TWY}lj&dz<(jdQB$kAa=g zH3r~O1@sAi*pF#Zw9KUrfE1d#t0eyY%JAh#^NEj4aa*_&|5b>gBap5~KI2ZM?_6g~l&JFukMQ`xoH%8U zAX$%15apF)e1x=*+<2D7#LzIw=bLdK5&=HUs0)_}Edjul{#Eq~6@dv&sMe5sdM^c2 z1tXH2(g;P>masBu7m6-h+4LtWkLW*my{Ov~Ki=zeT~BiDWVuf0{TMrtFz%llioBvfvqA(_=lBr@1JAkDkZ`>t9jCz6_Xg= zxvAT8MoURyqt-+K6Gagh>76o{I6yCaNJ+1D87Y`ZYd(VuUH}Qv&V!WX|mK+>M{)FNT>DP(dsf#lZ7Wrhf5|Kz7)CHX(&@j zW(P8r#D_ACRA`uqAzS1VlpLxOO+=>)gr|h9%U4}>m7Ow_ z^coH_{Q)b?!xoJy21fndi@|h|PYT;bzit8c%27+?ECF6m6+MFV)~-vcvMz{8&<#Up z(v`9|iRdPuMaT^Ac0fqGv!boImvwGBx*_)P;z>_>QUh^#%t2fDIBf0F`kEDu&v)5- zrIxw~2x22`h3FDs$DBoRAJR}t&k5TonK6VW$hZmf4SX;Oq-ZuiVoorWmhdjIhX~j* zpH?LLS{PK*jyQ+EQfbG={HwGO*a8St^t%Nbu}nCLfT-1WHKA!nb7wBoC~w@4KjTY; zL*_IbCttT~>ZH~{HazVekKTVD3*Pk*}E4-T9h|HLs!a9Z4nZv}BJWScD7 z8Yi&fbcu@#iBP4c4kl6(@kSD~avpndv~-{?Dd;6QJc+?2ojj?Vi?K{@n@m#{g6PJZ zN-)a7c1JADnK`$CKADKXB$Ej`yzqrDv{~`(Z+|;auf0s^oFDX{2ieD{9laXeb(TZWD?|tKFyJeF~qD32+>fFMSvbd2|ZR$ie&1v zDPkjTqqcwtEPO@<6yhD6ncb^yf6*=?easjfdbJI;0V3|4H>J=}V1 zk$_c<`jbET69$UJ#3vF4|Ll(ij)YA2J~YHOUXnDPq!te0;|>z%*bf3aU^cLZNH4t_ zn5xTH-Ef4b8&eVKDvE6cRMLq#wOu-SM65YuOL>M+37;HT$;2O91gYYZcjLC~u=R$m z-eSy&4pAn=0MVDYC$&7}75#$ZQIC3*H=kAK&PMeKdembPp{hVY!es^=IoQgiA&Zj0 zfu7w{&;h*zG21-{qz?pW&M28(q8vWjr}%uIA5fHJrq4Bd)aZdDp7aS5RjjdlUDhZH zVOyqskc?ewJBy83`Ki`3$sJzhsFhf&NY z9a-#BOTPDOs3FtJ>hqufe0>IJ^2tpPA@t{e{^xxtOLbB_w7UB0tJCEatvPgzgucnA zhD3B*Lzzob9O*@ACTVeRwmoBI-0+0ma~2!Ofv*-a7D<32{P~OYim@Upvk=%C2vqdD zHQKSPIEjGK-v$52CmF{%9>cmbH^4t}SG+I(q10IKQl=sy6fKcu6LzfG^Gw^72@F!e|X5!!PqE~ImB3d z5mV{G4}P$pQSkw4%(b|_XU*YNuX>evs#zc%9BeVc%De92#KB_ac6R7JDxP@w*&rBY?N(GODo8i7^Ugy zuA77^o#!a4Eszcxeks%|eCFBA2Z;E=4*z`LE{1nBHZZw2r7%MvEn3h|O#Q9<`GFtz zPGk{*+Xj5@P({&stV)2&SSrCx_XT3O{m>8nkjVe>kAM6(e&aVXRvhBi15Us7aD#yH z&8@0#M7YgCEiKXGH`UQNj5e6KRLh-}yVCuYcAL`i88*i%pG@K$*mfE^ zglDwvDj7yRsx!j)x{}9phF>1anYJiHtYSt3npiuvf-Q6u{#|v|RnK_FGiZxUQHDzD zIU+X-jH=P|Rbf6#2IAl$qq9L$D2@+qGz>%4BOO*y1(cj1fCfXKH}g|NVp?R{%aiD@ z|N5`@hA^5iQH5uD;~U@j%x6B+&R&6{PV2rfGUth8M;JY09Yg?7H=vUK_L#>! zMj{xlIYUMc=6KRAr9H8RIF~_BT3rXZZ+g?4_!nxk2pSvnD}33)gnjSx+!06K?scO&Fr= zi8Zxu`%5MSLjaP_Mu14bgue%=QkbOZ)((A1(F)GA3ff>7US`Q%kuLFsCp^KgYdreV zkEWqRgQ85^sh9)2+Nhty(;21B@Sd8h#X2+&^DwizvAr(S3M&CZ2hA%?gFrn5^>tB8XfYB-4N^wPMcd!vkE`ZL?j*wk`#;i9SLCFo} zL5U6g+R*?d-moMkh}F<8M|;HvRzplFVDm9Hm$>OIUU}`$`!epO4{X5doG4-%J#`hm z-TWRgSO6OG3iKCgf?Vj}Rwkfj-INMz z%a%J20-g|yNq8m>()i99@G-VJd;BTKSZ|c~)eRFz?!hI$!snPG*!t!9Xnc~VwpLfT zD~Y((I1;6ZfypDv!H^%t#3eos;*{o$vggZ;EN5wy&SQ@#9Pi7Cnmkd1A@TRlIn?cm zv3zdr%-9Ho=*A)#+WseQx9ngO<`q|50r`0AfOe~@LRK>F!3R4DKpXbqOechCx z^iqOM-2w<@)J=1k581FyP$urAI26gm&?Z1;y43y9XxuDHuQBJeFDcBZOwnLSkftG; z4r6%+Vag{#rC5QOwko>R^vq-|w!L&Lw{tg>`3@-Ab#JUQYOxo9x{y7t-sNK&aJ1US zHl_<coJS8^n(i;o9gOui8pp4tfAz^!}FGQI7Gtj*wZt;a) zyEz1k_JzU$$ogbQ5@NZJax3+Y3OD!|sFmXUz2E!2+C_^9ZZiN9?FhXyJ{m~>cTagF z2YErbOf+5A2B7sDIW!omB~zcL2_59ysB+CU*I3r(2Zh$omUb=gd+CF!y=9{Y!v&dK zBr;3=^KiOkS2metQ^u|m`fa}B%;I678?VUIGhNpp!PJ0 zvS^5g&RAuv1~Q=V0UG9{`*k`ov-{eBMk{-CpZw$}bL#{Jr6h8X+YbF*-}PP3e)hBV ze|)gjj_T6EA~rpliMf1fx~ZE>2|CM64wbw!oF`w|4v9>8vw#<8`>7?gO|MPoO=A zHr(@Po_cbnT&7E2B;J&-@@SmbG!fx|t*iLnNt}TgfIk?;F-}Nb{#C_zP@YCYQJ0%^ z+@5qe;Hk$msEF!-D z`@f&gfo=_kfrX`@BQ7+I^=JdD+*7JN@{x~x%2S>K&(NQSVt4Q=EQNq}%qoqKnFA0S zMm)iWj}++Mx&bQ^OH48+Fv3$%WD_Rn}V~e7^+1X#EL@PX-*}91Km_a93~<$-*mLK^AaM4rEOBhi>1-*W(LLTZ3O(<6 z&$A3};gp`-ja@U?iYh(L^lQ?bX3R-leQb0bP$hsk3$aKsDvAd4rF}5!qQ)HGAoLQi zCzAGH!-skXWY|JLMVHav4xQ&b=Q&0mM!AdK%f6}MH5Q`P%6*;Q8%S9pD8*c-H86_& zX)a#%nFdN)&yT8bP&ZXnCs7XMWjbRn!C!c$nVhr*uknHxyugrXs4+Zxd|>iG5CaX) z6jql+0~ZGo+jS~;5J8HfCmZostK&u`r3P>u9Ew;GY}cIGGMYm-mqB|pe|;KHnLL+A z2{a^;8uZbW{7#xI>d^`%LE1fP@O0}LR!$+X#Sp0IcZ)S>nREP&U*k9jW=wbRtZ{fU zxdAgpzK08OxC9x;z_apRA_TP%a}RM#N+m{NRVxI}7;@fxVcf|mS==Ybk6u1CQmR%w z@hEXmt0jOJ1)-P5)?;j2xvV6haoH)QM=yK+ZIgglJaNPc;e@`nUtV#=6(%|^A>+F# zBDxp?@LN@}K#{gHtI#Oq0zaXzgxM_xGXGGNMj0S%r(8xu4V1P^L(BBBmIDTTUn z*kvs=^ej~rCg~W`DH?l5tV0GplQWAmp4a%&j6}MnQWXi(8!?y0*4ZR1-&%9lhxMZL z5UlV-$vRUsYsh}&$x4V-2ouNq-uJ%9(064_Pj!2=Hq4 zmd@u)5`MMCE{uMQ>Et#Jn5Bt+37Q08^84-MQA;MB)Rw7<^Bh6zGx2$+h%4*zs!x~UtUWA#@t`2&xW>DOd*Sw+j90$7?hB6&n**b$5R5LyqO$x=7^1e~rpnpmVG zd8Xj&$8UY>TPbE#(U(NmmewEfh(~xzUX$p=G)z%Xtf9=MWFugpSVYame3@78F?vf?-bM)UZF z?jg8*ZXkY+H{yE?^v3rL_wso%F3KUzE0`E53Omlp!CRomFm5$WoaLyZh$9LAUQww7 z&$zlMgOy81_>65cAGEZ^O8}f{0o}~+ef+~*1U>#adU`CA+aQzUB!!y~I9#mR2IYlb zrZr{^zB4J)8(DGreFktV=@v-g4CKV6+6De;xZs}_p@4an_Ht;8_4*(NyZmrKW=p&r z6B>O(oMvN~$LC5Fk=xm-l3BGjTFv@e#I+Z zA*}Ul88&Ll$A__suyRhvfV`O^Rq3hQc=m7l=p(4XA?>9uweeSl2-CGlai?e|IXau} zhbJ18P-0^Qbt}#=ftwu`R>@#1EpQe3TCfswjOENyn$R=GI%`Ce={&K5oKnBiO z4qGErTj(Tg(ISjB>9;!PQriK?>yDo@{A!LC%%hVaw%+#i2N>G$lE&2yOzs6Ox}Q>b z_PW4d&*9w4$(^1)EUllXM;RhKcw(uJWaSV6XCYV;4;#w##s@y|frtSj0?c?K&GrsX z(h1M-%*`Sbo*5+ZOn$1kRD>sqhDvNfjc0CWoqS_pV@Rh-MVWeqw|idv;um|v)+G^w zSeZsDyP$=>LP|NIX9x~tA#+Gm6^#{O>&Aj-(JzbW1`3OZ5-Nhiu)vCbC+ayFcA}3< zq9ekX=FJW~{irVvG$9lLko8!Dm{*IfFhc`9=^i&2uf(e!;<$xn3Vp4kTD2B=1|n*p zH;AD}pwGJmqi#$Ax(VPtL^D(pXB@Ll#xk;WAthrYVp8x= zLfQ?YZZ6SarJ{Uk%22?8!0;QpTy{*-aA{xa$#xq*? zA5xI{Xo$|h5dv5zj7* zdM=&W7<*kNVU|OGRU91wX#@c+eIS^aa+#3fBOxDuOC|Q&by*>H=}2En`WDZ4Iw;dC zj#k$P!U2ZTRz>`hBrwgG%~YtNu1kGV=&0ktKMf%tflh8Gn|-Y;GLbcNHOtXncxo(y z#7->IhY{=HGVw&9huakD5+BQe`k^Y>2^oEy@eEv}IXyj^7$pm~DV1+#EO_&9$1Dwo zZ1IYkR}f7xsY-z$gI+_V;i$yY0B*Wzokv02)i6Ch|QCR>`bRBoSz zcfSLmA4io2Y&s9hAkf(PF5OfhnOLcp{3K;?_jb_SG)Xx$Q7$$+( zBLfkat-pTgj@A0S~*t)!af3Mvym<0YEbxkg2w^IA3~!*8p0)W*B%KRO#0K;fC1H?i@t-C8~+GsXt$;2AWcHj{xs1Y(pPg{Wd^ih*pQ02-Xtl8F!8 z9I9(UHxXqNX%uOjQ^F^3sfY|Q<&; z3UdjUcuF|H1v;EbRM%m6)E{`RUQ2=$vZ#tJtysjHbcSJ7=R#mhB2dxqmTcBLuTchc zj!?)dVG!q3q)o4l>@FQ5;WxhVjXVuM!P6jf@C~Bbcrs}Yj0;4-p?X!xoRVAOS$RPW zfvbqqP49oW0OV87g%c#_f+9n5DjhXD#c6!bgS6q?Vt=72> zOe4$tvG%Z~qnb79oTfxJzIwsT)E4G6p{FUAt&gT$7DY`EvwGu`C#IFibVXKBSpU>t z&Ovnm06+jqL_t&`efuOcEeuWB{Yo~gbzlTaANtUT#$-a}5MlDv!A$F?qsnx7DmZjt z!^bpB@Md+w49vt4H}{U0SD*H@rx8j?WSL*)UeF3rIMrYO`q!JAO3QS^GS&`H_aG^1 zor*r;b8ejzBe5}1Xj5UE@r}Kv`dm4l7H|<`|kr zrb`AQritnXZKE-Ia&up(MKNPCdO8r~frah&h!Lc(`VA!4GDg=d`$fjc1r)%)*+jlPvl=6)=Oh!*%T9>`cP zOzmD-smILTM2e5$O}R{tnIAH(;idRaRd@wyr6yZEp9iQMhRGsN6Tg@<50t4D4~;-D z>=KPbQCuVbDk@b)er5^rXJeZI0qF?hqA0_kxg1M;J7%(lbc#We&Ytf3VPx0_PYZEE zkt`T80qimSc2sy%q8SBthgQ=8oYQ0)#lf!38<@5!wx>Z1Iz7`;j2Ssp&5CGZRm6rs zOHyDr!wXyaz&FLQm(@>Jc%M^Uvq>%G@r^|rdk{j&E}MjHuvZri11$i0B*ZkAO3j#U zmc|gh>T-H>C#UC!q4Ojt@Wdyj(xG&TCk*LN6E}_sL@At2k$t*JBGc385GkrigrS6m zV)q>alnzf@`YW&aM^-9sr8CnQ8yuXib(_;$@0D-J(o8-+chyx_8Txb;f#O)(fd+hy zo^9twErXSpU_JFncftn&z_FFy0L-)$0@G`h=rf!~$@S6=suV+yQe9G{^*UX5J+T}F zO3~2S*}Wh^5dMmcDspsksw$PEibN`gp(55Lz5a|_FJ(SfLPxI{c!Y|Ny6gfbwT%gg zbKs-YAtA>(Qk)aKQ4Cvk4gN|g?7Zf1-+4ZWfbH6j^61yl^h;8DY0ONK==4W^PyoD>e!vrjXQn#muvsYYb ztzBrs{lAc~-g*xU(Tp(eItD#kaV< zhJJ39*U3ZcctP&aM7g=F!D*Q=@h(PxhaxVP*HHwTIGc_U2ZIAY$2sv;Tn_&6WSj~w zBGXyEDsgQ1xD0dtFGWhb3ru>{Z*_514}SU}J}-3!tWdB{V|&tezl^?cRw*E;Gl zJO=9C$Pi?bgho*e30czpoYD=Th+>_d0uKD;2;+quKG2#=caSGj9ya;eCpc3KB{HQB zin*KQm}0tT)r!LS3tB|%xPb^g&Cz#>TPeXbR%`Ke)c>Eod;8HXtI7r*5K*HS{t3H- zdiz~)Lt;qUt@3$6;wc(W!EQiQ6g&XZU_4=i*y1558iI0&t%-(2O-%Fxx=An^?mnZJ z+H`^;-l{R;k?&Xa(8I-g54);%)!w^SJ@aDKTyu>%4$l~C&N=2<>n)ygquF)hZg5G` z4CiGbxyf* zFFjT+UYc9Z3cQ6_jWT&1k8g(Hf01$ZA>gVo!RqHe_jQEMEsK>y) z%r(X$Kc6nYB}k@GywuwG&+){ z1vggkHMMEtlJ*Q@6gIHO1V;p}qIX1~lKM2Nq!9O9^>YyzV!hbwxw-*AQwFwFf^+_t z9yN5-xa2UbS}6iI2Z3h4H^=1X&-dS+jXJTrmi2nehX5SX`2Jg#anA(^)n z^MNjzflL@?JH#4od!sbFph2L~saeGAm|4UnJVRh+MWE=2USIWLtCIFZ&d?o=LJb{p zpoTenX4yj>?^ioLX$UT>a0_T4K*p)15|vzXWb*w|@4Ry9uM4-PBNtefB@*~T;d+sy zT_sDI^NKhxz2&~>62MkAMBBfY-^y~shiME`(S@12Hr|y!D2>=~64xikHzQ^-82*_b> zmFZ`4yn>q@AH+E5?v9wgPA%``ds=dvK*K|CQ5RzETP;G4$Kxr+G82FvA$t0x5onGc zUL`YeHI%8KkAEHIoZbk??djALcwM9rq35Z1EwVHo!E-PdVEn?HNA7nTVlIr68{n*$ z-)Fp)Z+p!_%;19t=fnV30vnAY=7q=$K13(qJZ0sfBf z_zn!E@t$y|BOE4e2pN3frjAl^6NzCrTHGNlS&G1=(1$ojEv#_Kuas$FpksUy;v+}U zxhnX%&ogRDsx#cC7@n@W@jc~}!ueM{j((lq*Y)`xp2SgLCBh{*;Q9XdzaMDmJ)Cq; z`Hx~UO#Hc(jYbJanU(a_k6L*{1Rqy}h6?_3s;-Ps?IwsfzxmB_UH^`V!moVgEA5a{ z+jkZuQh?WB2nwT?sMZ%aTmwvkSci$o0dIyf$4Eh@Bk!YhLu+O{M zrYNGl(q`V(@lcoh(oP>nv6uPL&BDd*YOPX$l6xW;shYY^f)7!$0oGsWp%U@}Y+w5;<-JuJLxU zRDmyqH2SEmKK5c24uTwW1X^50qeMj7SFSvcz={}8Of=%0OMVsHQzjqOxs}{I1bysd zA6q8si656h!4Mo251^b9O=qmW;w!!aCQg+QMB(inq2aa)bPN}kwwb&@?&odo@fnxo z174)`wyL}$P#_BqD)33&XWqd)34`qt7;>~O2v7nxbw@%EA&HDW*y?KlrqdSht<`-#(&oxazbkHkTPYF*BIHXP`X^+nhqQ8vFL5-XkJQ(X zU70Cx_S!S=$gcf!_=E)iLc=9-V&GCiKq8bN4ArQO70y9Ja3Wipg9GrHOpqJF-ir5< z{A8DWJ7fsbciAg-!NXWbJgFpBRoDfM?5FF<3YepZfcc%yT#3MBO$0QQaATJp1js~T z9VzDc!4H1W4Tv*7ZaupItiUIffdaOMeYUSuoj6fReM!JHuLBqjV1*%8;I)b4Ee41U zR`HAOO9$eHj;1S91&GWC<`3KdDxvYp5N;YV)M)C$7Ns}D?A+_CptR=j;hf3FB{Ujw zMsN1g;?z+I!SUJGvc}77=6>{+Un*gB%HTmcWYz=)i)Qw9I%{8!fJ02ac2jID;x#5% zMzfJCSrNx28DL-{;YwX1emo+u7L?)GVX+QVMDhqtWKuQ37FFXE3m{6;-wKkd0tt-jK zBqb-ZUsx#`;j9D`E_0K5NE;VJRh;;ESVDr~_{de*>aIomApl*CU+|}udUP8;tB%qM z+~fd!ghnrcr*g<&D8rDFVk5-#*Ck7Nv7NrTTQxeRwhHFRw6Anno>C_K(W6Y>La2@a z0fM`RTy~{VUu+$hz-r0O<>+->CHE92OIgdhyO`$8(U-#b>!KdEOUh+ksxJ@p%W!_f zNq4h@n}z-a$+9GPzHX3-nZLB>j_^&m*l>g7=dYIBmUW!HknFaSWCFoEBc?O&n4rO5 z8s z(32}IG(I3R!y=F|ddZ*ng*}kH;0-$MBzn?o()5uC9dwstd)+;cmyMr3ftJ$}q@X_n zvh4#8&-sFS9-?{@ji_i0w5HArhElWYgDrBaT+58~VTb zo4|Fc+dpuOts>W@^dYtR>uZGz00F)4gsmOUc-+GY})NqLsHZ$Bes| z{ysJ?tuSFTR)MW{Q@D9oqf;hljeZ|hJ(p;jk7=CFUshBs(wOB-mNGH(T%0dJOp3@9 z6$(7Fa1q;pr4oi(RII3mxn(h8D`QUu6Gbf2h1w|j$SLBuOkC-!9i?DqATn3Sk5NFE zNJ%dQ7`n3rYAzug1(p?8adtH-IjXxOqOdi$*swJBWXZyjEr=CcOIF@aAs@=fb%`3b zha9y{k^5Bl)LS&1XYv&X?iAGVJCWo#raa(9n7%gPW4JnA{6j}coq z?GkH=v1m~0l3veBe5|5X-Yid^~WL z`pUZii>Hr%y+HBNkA4&lkfnv{o-xunx((Z@NRn+Ti6sEifh~f{vsrQ}rbwP8lRTAw zK97}r+k;oadw3!#_3^rhar#!Wk7WL^UI#xG}?=)dc`zRN4?bn^-f%}GIqEz?k{A1dx6=_nfJ z_&UJ8xgm}e_Xv<30pcn?gpPHw#i<9*X#_iQ0xzfN$9Brtp@1{PH9CqJJkcvyfdc+I zav!Htm5Y_=Jo=|Hrwqs0V}IR=!JUG|08Xn5lI=hhZ!1XL>8!JCx?L$s@V z;t_j~`09Z zd22zsd1htr2?1GQq8)*}%R1Zo+q#xq1b+C+d%2!os%t@`Slh80AFqAlsfF!0cy8%? zLUp!`o>bD;F(q}!>t3M2j{*Sfic^V!qvgIB@Wq6sKZnW$JOQ>Z5bvW`$%ushi6@@m zFBpOgq;^R>C<8AVaDfPCM{dd*9seV^afGYDSG;wpIMwk`7a4_*Ke-rNsfD*g{mQpT zd57rPOGyDEAekoWq;vG7q_NSZ6Cc`9tk2x9>x^s0Z;X@PM6}J5D@BXfwLBb#r zQcy-uaFIwEG=Nf13_}tVL%rM3HYJOOc`l*H89hGg?lR_}5wQ@qxk;R>Dv93dzN{e7 z5j|PbFn9O6-~Dd%@-eQr(}79oiBMks>Q_TXC8FIjZWCUHAvI{OZ}iv8h~ZNL*{{(f zP8SmX@t8~GvLcGvqI&IXUn>C?KVFIQpMmWrk2#?sajEvysX??;mE5*4%;=<}#!!+) zD?D8y$dQZ1l|}{lS_CR$0TzksgktX1blq9U5k?a9u!=r7ARGF~+IwU#M+3pk?u8V5 zL>JdcYmzoDH(ZiOIogSqU-4}s~)R!bV?=O*Ii~#Pwh?<7^jR+(4#|`#N9y{ zLyZbDaD^B*Eg2ZGvRNSf;2;=aB_=&p_Octpbw>~;Q9F$|gh{|5k#sFl33JxzIK(F! zv=o6&Akgf06K<>#&mmycRyDpG!5vNNW*Wt?ZZ^%vyi;pvvyMqI6Nwh|rWsSRX~Qv7 z;LJlaLu8U5f*_69W=7E^a-UP=6JbVKUCk+#sI88(Vve`IH?5LI5TO*n zhr;TaIDI4;E4=me&Ue1k%VQ|@)g`Plb!m(hd*!w<2n40z#zzqIyv5N=s5K zRZ+&LX9SSr^i6kCu$qS_VU>ko>WG!oKsb35PetV9hpcSt$Z&l4_slcTSP@d2GKi@` zRi{KaO_*5{36(Qyc$_S3@u#6>7OgCgD3GP$ZeR3tMobxM1VfDpMaDA9jp_M}|21Nj z8P06+^0U_$L8NFzpSq`-_#n{2DC<1WpaK!;C|yzs|6P5y+VVAD^EDp(7(JjZVX#5t zbH_Kl;SB&`a>7YqSSTA(JUmARB1M895;+0_m>8?$6qImSGTg1lOyb?xpk1A(CHSL}QrZQx4Q(Tgl2GM( z;NXuK8JYy=^u5Ownz&6q{pnBp_?c69ir4b_E)Tlul;M;F_<)I$AzEzE*nwivO6?#J z7~@Jujrf3fMZ7%aai(mgR zIAapK;%>Upc9t=&>Wg@nLtBc#bwpr&i(W??s&f$-1Pt7!&kWq2g#6i`{h8t0bfv{| zFuyKq{xHfT%hJcBmBmUXdo!Pn|98CO9bTHs%#j(32r~r!N@ff6*lNekf*Z>oOVo1NKS3Z`;N#j<@&x$ zR-dI+Y-Bh}y`tw4xx1R&%$9LeDvn|RZZH=ZGGQ@eZJ}Kwx1JXzU9i$!(Kl=8@WD`{ zdc39{xml;GTb}wT&M*3Utx;5q5Cs|ftcOy^R#do@s9vs?&a3TS`h;Jk#2m_yxR;np z^ruXMnKDZZx56}%^g1dwStEaJcS3iVC#}R=WqY^Z{ed3D5HmWMB|7{1T8Mxg-l$8NC50e{c#kiJ71x7j# zEZ!U@cT=G%&)Qh@1`WkzypFGW)vLU)Faf}8$%JVY4{|ZhVPbc{;$a0&ly`(-puj2~ zH<}&Mb|j)fkAnh}jCkhyxQ42nXk%0i*wG-xwDki&@B@AeP)BZE4ZI@}e_eP}C?cbq zzD}`KVNvi$v-Ht)Z>0!aUj&-{USHo;^Zf`Iud~{=+ifN=IT_PcGIJQ&%>+jC*~eAS ztZy#C(I{`?M`LB+)a0bQiANo!DM3XGJ{q$SC($MCW_LwJK||uQMc+(j;TE#pBhi_1 zvhuLhh^o}n(X@vr&LkqBqz4G^eeZjjEX^uc&|Q7lDCM2> zlbIvdS9DQW7+{XGXc0ET05K5>fHCvJ6$}%-RFyb63M!%#|GR9V=hmZ=B|A;1p`N9i zdU{b$mg`^Cdf37iYkEt3PvN{)Z6S_7vY;dbv7je~4m8m&^>u?Jf6y24t`0Y$1sb07 z*xiGPK}pQW;@~O!1r9Ip6y$}UU+#q-@YMq#Kk6c8;@L}qE92x|R#IOh6;CA?sG*|l zC8l!%oBq^u6g-ZIwWshqvZeQ(uDI>Cga9ibh)kz|^1i(z^hhy^%4Ty9QqMg|RRw_w zA_`-xU;<%XY&|f|jj4`|jc!9%Rj5Z>z0_UB7BN8ro9|rG7&TU!=adv2)ZIrIg5sfv z9&*$xoYm1FAObfu5@SCXZ)y+;)MloagS@Nk#c96Qg3?BwA4&S?M?b2s{KsoVl3OpL zGXYx!dQlkLm>DS^3u__5F;|xr^}R)mvD#bftX6>!J+Tc3ved&rbxgDjt7Ux}FTE0w zg{WJJ8sG&KOMQfF98-skaTOVT)Qwc}W&%J*7-|$ey2k{RN>pW_3RL->;fnfP-8Cm6 zk+8E}z8?XDwb|W+GTE<{aVF`k|BUPgddm)Dx?+e= zhMZRCA9%ymTOh~Dguo-*ndOmsk? z9GUI#G{@kn7ol`e*O5fr^kssIj5li$W>g@-DxD^Lk&WY$lVX+jci>x z2!@ss))p3Rae)skJ#3f#UJM>gL>fQxniYlxs1yY&-r>9ORCh^TDoBBMz3W|8W0p1E zXA^h{LVAiTtSBZ19Mhfao#JolL+ymqWvyD03{U{WH@)diRH8-#elHpVIAKW*D@b<= z#ptOQCX^T60MXaGyMEHf&H~$AVwy+Xe$L7r5!lgq-`l=}3IdJH^kJLRST`#?m!|IO zO(^{0c_4)0Y)msOc<^hdz%RZgkbnSf+Xr497b?iOz@YTCU;DKP0@o#u4~Ot4*AsqW zKo-|i+jG}^L62LS^Fm1m!QHe)G0UCIbAyy5An z#gW=>Sh7Sio-si}sJmIfj1$0}QcOqBcy(m5b}5A0*{R}Qb_|(b0v@red-c?aGMDh7 zFnV{vj2P1^wnPVahup?%L52~TyrIT>-4ZiS&mH7}p?|QauR910HFOa(>ZZ97%2@T4 zn^>o4$0`*zu4r7H0!)$V?G>8K63Oyu13&^}n2bZG*Rm2dmMxcF=U+VCIEEUS@Cdav zGL7FUKzXsD5=tnsMHg1U9I-)3M>#M@Fv_Jpi`&GK0hta~3rz45KAjO@ZXDKYx9Ei-dWHxL4&wUvyyrbW?RF_v3)Dv+ zebkfWgv2hBwiJQuh(NR7>*%%Wd=P<5qAvf%U;Kp>(Ufo8*JvnLF=u$e(!7;P$P8qz zGC5@y@fnpClqLhKaF3o8Q=P@4StL`St7af|i%V05x)u|LFRU!%5SwXS&5TYB9T6y& z6uuooPgbPtwZT-Hd6drjn#uGf$PrtOu2PS6-iS6iSuL5p%(X1?p_)9+{O4-awz`cJ zLXkP{723u@07_uzrmW|(xb(A`{(mOeUE(`Fl303HnWSs_p0UC;Spkbi-|qS!o3`pX z&O*>5P>v=-6vg9W&+^pg>s(AXTG7EW#q&C`{h}AW$p2;0-SUX9^|DeyBS-}QumAe5 z3t&zcMMCt9>K!K&gfJ0!Z1=OP7G!})``(VE)5;|=OpT#o%ha=y^HfUyW*LWGFPuOA z@sEG^cYn7p44JG^7U5r>KuScfb-~qg%}PFL}R1DEgo`_dXgoX11%6jo5Mcrf>QtU$NRt)slOS zX_z>AjONLzc0umdbf2^&g%HM|7ShOIA{F%c?s4$R#V5+6bUKE>Bv!Z!Fig8Fi=Ofp zVPGRkJh}^c9M9&fwsmpKaxbsksHj*EBn@V%Vk1vSg*miJ$5Rte@>MS?yQI-EKXV zD}V&9#EdObQDlYVQVMvX&5ycJbr~&7DY&te2xMfrWDc^41AmQ}tFTx135jL*){DUP zMWEU5_4Qpf--CcD!GLbjZ)!IIm|YCtS=(7*+`fI=`okn)zDKFunm{UIe779=gQW2K2)!goS_su>5(+V%CX!pXZpnUW z91ZuVe%fUAjSe(4_h1q!nQN(`P=FvUC{10=^pcy058wRF-z*=HhkH(_FB2>-=@6k< z>3B$P+XsQ51TIi#(O_xIeHOC76dsUQ_d%WQhUsm|!cabrC0LSMN9_wvb7_ zia*XoLn%TJ^+mQMu*{-%u*EIC@DafF0^E`T$0HHD46@&t<~g>NB6Cpc;+~RqsAy4i zN!?XS<}{pd!V#qK^x)32+b*62LnPP%u}9|HzX0gevK~01Ld+NsMUn9;E`V5oy;8Wu zj2)lxj+ht2cY88=QCa;t1d3;)#YtH1ZPzi&GNa z5r>o{ckb)x5WBEd(J_tW;*d+2t81~TWDCg*Ke7%LNY$-p8FrU1>eujeZRKOgOF z3$mh~c7cRA#s;Nu?7d@jB+=XU8yg+ln3x?*Y}+;`m~i5m*mfqiZQFJ-$;7se)AL{F z#d)9i^ZE3w)nB@+YS*rNS5WjLovEGt7x+8g8s(*5gqnS!OF z7eil3VAc)$xpCoZM%CVf3t+2mq1PS{+@r@6u-EJl6%UP!BlPj_oykFsn$hJ}VMKgs zY7b7QWLItU3gRG!l01)uD|MjdqEXRq)HM}dw>9Q-;p~lmNF(it=`>9C!71uDI5bGr znDqt@r{(NNw`B(T2HWG2D#dfN7A?oAqv3n1FcV>{HT+g8UBZ+HMI>U+$8d*}&QQuV>p3BW< z2pNesGGJ}BlX>+)r1TMoEmBE{sk7WHkf$wwRAr%-jTDJBo9VBl%LU^3gIJ&+-c0)R z(J^&9`)}!lRExSEC4o!H1KaXi9DD|Gi@Z^vaA(_PeibJqZ5OR2$#V0cn+KCCsJh6h z_w=Zk89)f6k5$*QjTb@68*`yO@S$Qu0o>l!4GKkQfgktK(DU2EPD$3aJj<((Jj}Od zHKWjEu)LrH1+t<%Gc6R7?OtI*GD?S90In9xDV7<60d^0|gPB=c>g1ax5_~Kq`Dq%gc`#~={Rv~5(LoYF ziNDU@vV^RSEmwuBqE~+Ztivg1z6Y3I#3T5yc^T&gi5&xoW4pp(vv&MQT*-9Gwz$L@;|5iH zqbe5i;W^xU@yFx2=&U8PwfZm||LnHV5k6Pt93%7TPW`QB0`P)1v*`{9mh+{=>j4ztd&<~N$f-eNNr)|P(1(Pr`4a0 zk0{_27Ygwt$c);CpHp@LPrWSP1tb@(2-F|PI@2YT=2a<6qs6L5Nb;lAOi_8+t!l;Yn2Tvm_(E;6>>&9!|U_s1R z38YJVTx`Ay=TRb9gW&ccR>za3czEZ)+5Tb77a3@uTq7WM+W3Oojx$|i8n6BGTq|_! zu?PkNMx{u&qR=Az+^l+A80%z`W7_r5Q2 z9tex0pFbW}@GuYXft-dN)Ba$(KI}&xt&mDC^HnA=s0lSxn!M&}m?FgkG(TMyS$m7J zqK)$qiMB4)Br6(nW;GF$+z&x$_L0l?jmzVUGSyyRp&TCQh=&nZ{Mpi2TqzC+Jql-~ zI$3@dk!&?i&v0a?dUo`tTy0`$*c1wxbqVK-B7+>5lxEZ|={05B(LE_5O9IxPsCv-c zwgdB*dD_ZHLAf-2t6ty*slm}TCxzq=Ad8z@h)2&=RlDyf5JhRYHXTz?rYhGt{c8o$ zob^$E|3v*=y%oS=TUyP^OFf@Ju%ML;%1K+~L#_$PqBY#GF-;}CrJ?nd$C(df(qRDV zk5Pssm5$35@!u)j4@y$B17=p)X3dVZt1{MA!2Ish3|U#p-Z$sT4R5}KASnYbQkEVJ zY;zoK-@a0byDgIY;ulFkbKyDtM3rVnB)z5o4eHdCnqH(fVpitmmY4fSW{nu31f~Vc z$Lf^G10f#f4gm>R)3DDA1_6>Z9)*Mq*a&AgCpo_&tw&Rp7Ix-xdI>k6xWD%ahdDEH zxzn;!qNXgmoQ}(ISzC1$2cY{KjAwXnvD{LO-Wtwu*VYok7*ja^gVw zBhIZ4*s6$z&s9j+%qvEcuMv+eSuuACV-(=^1}z*Q8=%qKRs{3EJV=do9uze}j9R?m^TR3tNFRa1OY>6?a`qdUeYT)j2ZnI zic#Xk2cW!`3^{dAS-C7u9`_oE{YMV8NXQx^$WCJN+jgs z>k1}&v+U7OMMfuhCYF0jg>}_4*D)uP5EpCHhDoZcHz5^jsw zWy@%y9NaOY-5{w@?HyZ(yXD?#*Tw0&z(der$B!z@#IJo*BOuYnIS*@^7?Eo+C!0cs zqtNb=(9(^|A9a@r>z%Lns_9wZVZp5&8(wi_cEGhesxw{aa%cT})MZX0sBK}>3I}%Y zk4p1L7l0v*oeAT^R|s5aw#QW_Sr z=c;??@g?0DF1cYNP4CjTemQIS`+vqD0vqE4rlb4lCMqTjo+%kRGf$V8sahJr=@(tv zd5pOQ3EV2N$Kv2Xxde8vk62`gHa|Gs>DbPlbfzcCTbfocC=QK+E$LgSiH#a;r1_vJ z&^&XpNqa?=ctkeP3$!A2x(c1}j2MiK9GfDdh`ap0eXe1CC9v`z5P? zC7?%=t`9o-&kpj(yzIaUQ#@pjstO@lHEWMQ&(U@TdJKTpKha7YN3-9%`E|d^J)5M!Ud#4aj;j=+Yq|oOwD*sE$BS)EOoyuijG3%F|KI|J9u3dJZf23=-7?D`s zX`lZ^o#~sb=trY`4g|uC2Mw~cYn$}i+ll&x=MGX`1od*bC`B2H+P-=!~Ti z(BFZhO&|hv%CXLJ5`o7@BdJLZurs>!v$UOmV`B`p4pLEGuG+D3Ci7A(f^03EG@t8V zRd@Y;u>2s86TDyniIMm* zyg|jLMw{b@N)$e}zl3~k0>OP6+#7qJ5@>#}spgRLUbeV^26_=aK9Zo0y^otPRtp@2 zwtgCpJi*W@gUMcJ)3!Y09C|_q_x8?-Kacnw6RP{F_mZS@$M_AQ$HLA~{Sg2XeJx6ia=Jd$6iaN41jB+}oZRpTs?GE@ z^b{lF-y1CZJU)T7aG(}pF_W(+pQ24-kt|})cOe^z>i1E=J4!5+Q{7qb^kyoLgEJa~ ztvFnVHge;2PtruOHJ<=NnRz74$4wn3A*y5725t8OOz8fg*5*OJY*jNA9J(|kq8^~+ zS<+d}DEc%bPN3~Z28)kDU*M3_4B4Eu`MfqwkL-7D(;9ile=~~gG>bce?JuFrIMa$N z7av9(0ZsTaRwy}AGD*U9f4*NnD*eWcO1u{P69PFx^DA_k-ztOv3M%!H0ScGKXa}1@ zN56Lbc+|Lqg4&L$g^K+IV#YEL&vkpW>R-RYXH?bzJ%%FsIkXY?5X05!XZ9=p?cV}> z5VVM7IG&Rr5PoqGP49Z#(ZNVolzO1mleP_`AU`j&$)MrEq-g;~@wf`7a+2eR=+5td zbZ#8V&8Mi%qa@8oB2U7jjIUHkjWjDK{?ZC=F(T0dR-r$V7{78rd5p;`rVyUL?Elp< zYco3i@UNHG!qOLd+`;MQexp#>t!m+IQu6)a##K22v(p`tuw~Eba*rn_!81FvB0w%V z90{Wq>o$i@bKi@uONN5_I>D=go-s+$)RyJA~L%%4gN*WQLSKdIp!84j&Tn zlw6#m4oP+XJsQiq2Le_!WNVm=`%4Clvr}Wibu^!>5>yp51gvDOO|DN;M=c!g3-Tm2 zYW@!)7RQ&TF`a>lZ<=ezr-q-N)sU-B3Lg$k6*ZGu-`mDSB;8!uk#m{M&GVK2?nMdI zD;G7Voy^S7v)Y8V(p9rur^Z*7vxt4)sZFkXQW{1s#C|Ra$EN>+Rk31{W`3BZF-w`D zA20RFh`~-@Sy5O~%K)6^F7&;iJPCpfg<~7`@;#^xD9p--LW*hygvjN&2Y*E=_b00) zG;J?0kZp@!@K9ewb$o7B-ZB!g@6lszl_=KMVi{BAi1FMD$eiM$p zaopB5evl=FhUt}GICQzVf|=ni<6yV5HLqC2X-OH|o4pAR@JY8jp|vI$bc z70U0XB(J#0j#7^4#XBd?c3<*$T4iF%`kcS})~FF3q^9&0Z*1;tes&Hb;4P9+%MufGVMr%6Di_!Yc>$r$W<)Z|X_oP;x0fS9AcN)E+&FU83 znqP8YNXm_MSYSnoe~2q08dXKtF84HrVMbb4w?G89hQW!4pnZlRa$aPto;u=5|fgMpMw=pttN?f_}xB0=O+Ok{Q z!0Z@|oR76xtbZwiXv=x+K z8b}k(zBO5@y&gOU^DTM)W>|>qr=E} z4?Ua58>J6xA}bbg!}1~`C8R654j86{l6aDdS#wG%lErFqRf3ip9! zCO`x`IY!|jvTjcy2A>KgDlOP)6rJ9GKqs+AViSojo4BqA%Q81sXy_Q9!lXRzH%ojT znWq|XR>am>E1Dr z12NdmL8h`FBQU=XSjzFTIkoq#zNi_N@zAM>=jVjE9`FIFKT>^t1&)#mpBJ>SQ$+d+ zBSdT@Pt)bnbQoOtLLRdPZ{UJL8Q78ut2BvRjJ7Fk>astMvx_?15E`U>%?qxeLhnyO z*gpwUhw#Kjd+1LYy~-?MD6@0!o?#ITQ81lj9mP+Or=;4|b~X3H;+ARd!J*bU!xcAL zn^)o?zjRlV=fi!D1Uln@Ikz4uL31iL(lVaiJD@9M7_Ae!rq1bKwXiH0GU%K`IZBUBSKWbV$!g;69Hwq4-cY zOj086;e%}HUHXgQz#O&$y=(tRWoXu6&uVwV)Dkk0_XUKAZO;|8#RU&!vWp_N2YmW1 z@gEUkWT8JqhI6S0BGie#M0$o<7WWZ|DGI*3=%kX&rncY<1WVoqL2hLIIo!be2! zKc$pF!-#<3o<$~eVxUd=v;|;<#+G|(V2G){D{RcxPH@A`zDNrrH7EsMvCh%z$5)6@ zv`x93*g`XdSu5j{1#eMc(}6-%5i(UhW>(Q#1+GH-&y8^rcGuY@7bQ zoW^vwv}D@Je%SR139W5Hv^e7ZRV(;?Z9FIj6T;4iGe*5c4Xh=mk-NdzV3#JCIT0p# zLOiKoQLMvMiHTp-B4i3GqWuPydEMc~Z9x&loeOu_0`%soKiaqBRfZwSv!2MtI2q{6T3HMl`mA zpKG9(NfwH7`jnRRL48W?$e-X%ALr^wmXHwk zSC$su2;ldfcLgnDXj{_FC7IGx!j>O*;->j4sa zZzxn`WFIF2yiBFUMIkvOrARn}ZymI;m?z3AZX!to+cNzTnK0md9+UNwybuG{m~RWW z!O_)n>Uxv7POP*kbF%ik&=y-gGP${woPYX;3Hma2fL!?l+-`*V-Zx;n(s!*|^)^iCC3`4-LYG30x(O zfce4c5OB9v{Evvaob@y?tgkpfrqA(OB}*jj{-|6+Jet}UN*P+?W<}$jA_<2|Ysc+l z!p40yCN*l9=VYYonZ=2WOWifl2*I}|r>@OIl^e)S?&S*)Iwzw-b+q?FbB{$^>2C0e z<*2e3Hv+J5nF-9qn`ii~K`h(?8HZeK^pYYkLo1yG-nLQmflH@W!Oqg{s`K5Bn`#AB z`P&AL-v+qm%|`R0xt)cUCPU@{5%rZfCVG17+soUoRRr(lT@bf;LZEfe4^AdZ|SK>mgBbcngC?`J|*| zOSGhTGXDIZJf|yS9*sXGj@*|uhTO#Y<_=E z{$BfMT(Kk6(k$#-n{u$yki>vG$*~+ch8e_(OA||@!(%x0FOTkiAG)l=q0Mqsl<^W% zYBo49NobM-`ZJTRWU#IjEC@476^a3}dH$<5)xbaq;GecKVrW8kF$P|yLQ0)83~m4U za5z4W6DgWzC&J>>D%J}0FADi!VVDkOysbvFfvU2v%|v2jeO|f z;bb$i5%V->D&b7{`uGBWQoJ>aox=nq6tc$<=_WGgv~R+M{t{Wb6mrV3JkI&at6zE6 zzyH%iW{iQ+m#hnHcyR7Hk#M-U3?M(XM{1rP3T^rT{{RzbTaQM7c&o z2=R@4-@MIwOeO3r%)@~qS9tR1hvvQzI4-dcYT6Rnl0rb9tp`e!nr7v=st0+#T_ha9 znN#>6ao`!xjW1O;?@;yZ`03F-n}zkYp5XF1f>8tE7NuwoTFaFAhi5|RgCG>JxSheL zob-#IZNa(Q4*5}{T5wdEar9%>0bE?`l zEXa;h)dv|+!vm>V&e>7T(O8NP>t>M5z5&UOgSrvymlWylC_{sp#6f9h2k$q`pHA`E zDHBH<9uHz@7(fIlE0a|ySgdc;-ob(zI?1=EtBuGc!3q0e85L_hAn);H2(3*l*((No zmJZy{dK;&}IUOh&q+%7myr4d8q4mjE(!8uMiyE>469&OG$golnljIxZ_Ydc5#1lGh zp2X$6w0Ywuh>97?(R!Le;nCPR$?;h}+Gft+zbpD55C#L4F5!gc)*^p)jmT)~GY88B z0z$rGbRZ&oeEv?wyxN%QOBJ0D>VHCvP4}QPrqmC(g?Q2&lIdy`{_*QM5&g#EcW%X^W*nVVW#Dy;c4ikt z!kF(&!YYV2GP6592r|z=+cyvh!M{hI!Yd8kC5?L8DgxU~%|SQR^~fJUmiNwmWm!&@ zrMrP;_I(458osg-igBxfAD!O0WEiw#kvsUNBbZ=`zc~cT9F3~88O%l<&|uKLwC%N) z=$-|TyyeV}))*y(N?P5E3f_rI1e;O!f71+UILf=ipGpMpLO7r#S6W6V%3FEb2lT7+ zs{d6YH+1vCY8${D zRsJR&yN0_qJn4qNF~t2&4(=cV&uP_8e8^tB=)wH@0KN~@SF}5*{A|m5#V-kW!4MPI zZ&{bm#5uYDRveRu-Ej0QIr7I6 zhLb{bJvv8x;t%>aA?NGOA{S&2wr7>dEx9O~F;8qH-7Cp|?Vwkfe`EVdezis6&?uV*yZQ00g%xT-It`NudA|sIzCOQ{C*OOCq=K^^e}7VIq+(*v zw#dvmzA^eCZHhe~RYrlGnS6@j$m$PnD@|kSMBb1KpJ@1!?~d`%-XJKoVd5S5A$Z)braYhvt9!j3Un2`PrHikvs!q}?p7#<6F(?B$<;+Z;zj z;ZG+hzeYU*^}fRBNBd@avOY5$7&E%ut`J`M%gU$n`gEq^*Fg-K7~;J4zA1< zSFMTjVze4diB9flR8ytM4hZg2N$t#IlAoeZ{RkuxuDEojANDv%ZUl|7yJ7v_!4XwF-=`S}-wk}`9ZGpapQf`GwY9k89Ix+NSwJvqE#ER?TU zd;WgTX!TP7aDSm56m&0RUCyI&FWr(idyta<9@+4uZJCcZiHBS%b*Px-b3{_W zA4a#yCX$VXl1r&x%23f9nY8V;u3@QKZrt|o6pbq?)qP8U7Q~cUZk#&4gVFeSP7Z35 z)!d+#u<>*`@|0c-*74V(zOb*Z0rW7ecW?W!Z#1Tq#m=@+5HiGOzs>F;NuGScd#TcD zCNFie+O3ucNeezS*3`s+^oiVx$1=^dDMWEPn_3Tg} zT^&ZJ!UOX$9PU(>BPd+Lx*u*7Dg0+*^F&N$c<1)Rnc$T(#6w9nb0dYq1FEDzs^jotm^JDQP>4#a9CBnG^lY_B;8-Cs!UXBJAw#G*H zZOA5;yRW4GN+GwMTh4pi4k0;kl;iu|gpa-Q>7}DBm-2_G3X+@frEct1za8fB6XC_MYmu6+0x$x`thB-QZ8Pj%|JL(ebQ;3V=#^z`i!fJzaaTIN z$d4QjU}G}exG5FGfeq5dE6J;F7IznB));0=YSuB($?am+xFX0nUqC8o_M>v2`D^z>SQiF>T1zmRi3{_m%yXb`;VP+L<%EtzfI&bjhDFgq?k~R_Q zgkL?cFTCenXlo_PQ}C{Ig504QE;WsfqbyZ(VxL+K!`XtUKG7sxBtO0oo^3u}d-crR zB7R(2jqS^^c<#iHlP}Jf$9F~T;j4Xe*~wSSRYyB&(!=%XV&DlTVwR=CqGRRrv7|Dd zDHA_ii991@=NjEfjfQR%89EeYsTlLwomIF>C0UUgjgC|-Uh2!63nE?BCQSwrB%(2C zv{=sv_FYyViI;~FXIJq=~iC+07X zEu;}5^F;asUK?0!N7w3!CHV%);3J0O3QY3v&=AORAU@F%nRvK!xe*>PyN6XLX9%L^fGJ&df3Qj6?FL4af|U zM}W1%&I&y4^O`B#$2)FsB|bwQuw)~IDQiJ%=aJhsQbZ_6CmNP<5k{+7oz83Y7Ud%i zTliJ79<1mlV$Eker+r7iPKKRz8qa$I5z$ zSj*l%L1c@iAmH+E-KQHY^lQgo(z3E8dj>Lj3?*1in4cLuJg5u)o%2bI6#G`6vI##I zFNaHie&+~;tkbC@nn0sqcP$h0WqTW}dgI!qAQQG~z?W^+v@{6Q$a6wz}Bq!P33(gTPXTy1)Ug!Cd0pBzaN6 znVn@F4-g(Fd||f2+h+$)Z-e7JEg-Eaa{#V`a@7)#8Rgh;B$n^ghbM1gk?Ko9t6!Sz zh*NiI28TQ+5vb;&IleoDbk{^CGs-{z z2xD%Z--KBo@01`)p;>Z(=rayY`Mr!zq9n5{iZ}){`gMT^p?W|be|!ucZt+O?W0?g9 z=`a-IAcj13nqW34x&lcI#;@wa?AGa?1KRyTzjn-?a=P0@WOAYfCX`Jg)Y)1aOcKL? zywHavI>JYl=fa`ZWKNR`6p2Eda3-?F`O7m4di2vrcM`N`a6{zj*8?Q%4FT=Euwma5NT1 z6p%N+y21|D50FXmOBM@hnA9N%Aeh=U{TU(EyzW^-5=lsnOkMOy;`a+zg!zR9^zivZ zRaPeWuwpu=OrRwft9EEes4Sr|0qhRmFNAr+KvkA2I3{RPj$v%ozCzP!!`O?LoY`QA zt-HzOhl6BblMD_RZwew3RV@fvYas&hB^ZO-9BK_`a$q8kd{MR6@nswHm8Fz-WZKAl zzoK}FZ}kY7VJ_k5V*Mq0oHKMaU7)5y;8hq`k6`*Fcld?YBG|nC-q}3@6CQC^i0J*nb1$6c!Uc%#pd+-|spNlZ+q&X#fBf zx6xABp4D<OWENipe1q69Xag~$FN@+Zqk$x4i zifP!_V(zGOI`oDmX#F@gGeZ>(cesQYizv-M@7~QM zQrXKlXp+yV6<&3Xv&P!Or3HgC#Q+2-&N0>kX9bYUGMo-*Li&++$$Gz0@sFqJDi`9j zCAe+klLyS{wpmyA{a0TC&XJi-^)x!@^^d(F-eB`!**%&_bo+%{-vv`=W>u*|ky74w zT8TB%oT@(HrHhy`aBTg6CwQZVIqv7naWQMk0WEQ%=Uv%*ZU`|sJv-4mZIJ@N;coC- zi#Q=Zxxy^JA6L;=Dp~k4xCLt}vPaCjG|Q~(p++C%+i1C!^Q=id;(S%^}r} zX9FUot$X6e7!z-;Tq4DG7|?3c2q=TydLBGAUalHuj~|*)Y8bGu))$3i@YuGTDbZ`W zfWbMdAp>vsw*M)>7iaC2H6h^zB?3)#(9r{^u*DRL1hHM7om?3wkOW#!(1{2K zIq@|YONlZUM_t!-uh67=_L$Au+}%jsP7Pkl?F+y-tT@ID8wYNfAyDBUzLbB$iY1jP zLoq>v*Js9Kb&QUVMB5!GLw1g`snl7>5T(ruCe}aNYOph!YY63hrxu)O^g1Cpcsi2G zSd=LGTvr{LnSOU2TBTBr8(lPkUZXkg25y^H=Fg&jZZX{7D-ykfg}EgcbO~IIs|ZQL zSoB4^q;Dou{1jUtu+m}BX5hWS^v6=ZjBI1JSFyMB~Js!0CDkruD zFh-8~hZ*10fc~ME2Y0bEEko*9FbcL-@YCkQekPZTg>1%+XYFN(bZ;2bO5)!O_qpRb z@<^5|tdd`b4OdReeFndgXFKbQ^i;wRCRNA^BF8I7f|i)J)pFTEDV67`&Qh?6OoM~D zVftRjDQ4|Zar-4#uK*#fVet9mo0C|g>LaZAzVKIET(7@V*Iqr#sXvO_Vhg|Hf3*^^ z_iPPM`Z2s7qmyV(-k{m+5~dFqnxp$PeN{rFNQKX%&-NJL{V^RiT%6U!|8kV?)WrDx{6QQMG;0bs$#&?L;K zD$}Xf1%(|@({X;YOIw1a%qWpNh$(y1SG|;YIcVvm0C>}3cfaCbMY^_j9^wpDpW(IHPu*M>fpNEw~dy2IK}Sdd|3&(2vwxtZqew? zVm$W4$*eGe-mr3aaOkn2`gt|}?DxTkRp;~mxaEYb1HH#!yk@15ud&v~?WS+#R3M_; z=vN=1z^^^IE^rU(kMul2ND0SVQc#r*V%YRmQdXeR=D>7;a&Jx>{*=%;W5sQPHiKvG<%WAdpg z-3RsAUpEQaiD~0ios=0dCDT)T=(O@wKVaoVjp9Ze!4E>}1}L2Dm}>ecmME_z`?lZ` z4Co-lbKo*?Kx61b@rpa~EdkP4s6|H=y$IBC((ewdh|N#=^OUu8rOY@OI(Onlj;)cr zh3}hgR~G?xSlUqGJqLhTeA+A((CH zrfs=S3VZ|KTaWgY_Q=uIfkgT`HGb(`k5+vIelrW3t-h_$SXMvXYdlHA@OX0-sBj^) zm4C79tZAr)tzA08aEK4>%$ zT;P76?}5kETDk-gJ+6{HGq;e`>3v6VT)gI{K&6-Nn_O} zTlTgM;tleb1|i*ibiaA;uV?imE(U4Bk3tAp~9RuUl z5h|{>9(iz& z1CLwmAIs7kb4f*OiIM=?r%Jp>L6^|d^{?!%uJQveIKTbxrqQd~suu80R`Y+u&5jsV0lakcD>DBA zyG2a8^}I1CpstVXi1%2fwvCiN?&#FKFN3AUYA8}{-yjIws^&46r3EFVPB6<>Ofk6} zj~N$UcWB?OQYF{@doIC14+wGqxzfM;5>#T&zu*4zWkmquOLxuke{R_S&j!#JqcuS|55Az9B4}dx~R2BM@EMK{eu4AGA+8J{x2n=fi8FQXY=K*j@EyP z1Py2l0p$Nt-G6ba@d0r->^FxSYyM}{{#V6+zjMO;pD_h(HG2IQx~>=-t`6aC*$(0{+;{_{5e@BbiBXXd~EgFqek|L!FG|Jq5wvE3ZdaZdm| zommQw>ka=?h&y_JfQ;{<4X>cBwR^8XaG&psbMv?4SO3-wncaI4KVG{(3H}&91sNXQ z34Y!RcJIk{L1q5}+2P7k0T1^(z&%q#YXk8Bfq70NOasmwM9L@7i=Fy21wNis2N>ra zLQKnn!2xPc(&ooSfc5vdJ1u+;mWfgspj`V)L%SG8u#&9{Z^7sls$cZedhN1-JtZAb zIT=dK2oacp#w}j6z&UW-e%8_)jfLU}n(GL43T!(8zAg517F%N0YF;3JM;ZYqQ|f;Z z{MZ#Vc34a}m~yq-wzv-s7$H!(A*N4zu6?~qPBL-)B6sJq_$5|+&h(;W1EAoo-6+2)|c+I9mQ$grw1bJKVArr+w(4u9ifEmND;E= z*4=iRrhon*pFuwQb^&*84~x_qoRZQ!sv-jcW(h6-R&vQA+;Mh_F5~Z_w_zha(w{c( z3UBEt!4IYUvq8ViG3K;_G>$Bq6$BFJFoBK$%BErAjO!*Z5hNh?Z-@g%JYC8Qf^*u4kc8v zK`7iT;{;^FiC_^!oxoZZ0#`{%CP$ycn8pUWXgLY_R*Hl$7>BGCs^vLoQN`7)FQ*1! z`oJMGv~SIXgG@V6N&`{zsUJ|bolAIUnT`eIL|WD0m1q--Rm>RonHc#grFf~&RC=L4 zS5bdo;LmcaP?yH3YF790O5Ux{Q9YtLi-~t|KK^-f3#t*?Cz@L9nox)_?UQ_x6r*v? z4)?68f+PRqh}S@-u(?h#&xWPkaOR}fJl@H}6B!(sF!QEv01o4Lo&%Sy$@6u_jsI@yJf+u6k41 zh*@NrYbS@wYw3I|>0yxtq&%1;l`{otQ61~eG0G)7r67f@b$^_8|EWx!Z?Bk;cMDbg zAtaK#kg0*7ZY^*?SckJeme+2W0R9#FR$H!|ooph1V~jj%tjtkC-^9O z+-ieQ51$I-g75Q!nd=J0ByKZ|yBlQoc7lOmN9cj%^3Nc(;LWGAPO_=y%MG<#j^IGq zx%ieR8Azw-`CR^I1JBQ6?ltD~Dds<+3v1z2c=u#`~N{fq$PRwGs zRND8(ry3T|>j=ZQo0RT1ac0k1S%aJS>bJ4#1^aUdElv5fF%PZ}uHagudhUArnu1`t z871&c^NR|UxZW{q4P3wvPb;h{-bAH74{ioW*+5JpWkyM1EOp6uw~-DZ9BW?)CoFCaA(7uFB(Fva$IV7AqZCFek69q|T?I<) zHz{eSn^IY7T%@<>ZSRUb9L)un)9X4RAa`5I}SX* z4}F=qb=u@Z7*xOjZ+sQT!_AZ*QxGAY9081Qy+E=+Ytov z2)I~w(DAS(%lShb6HC_}Nce8N`RnuDZPVxRpW+ro-$Ip7yA+dEozf%dpA0?`o_tj6 zKCP&q!q&B|y&fU}*n-TZDMaB~LSNTiO`nR=?f)70 zf<9HCgp8d@u1(vYJZ*`ZH5+6Q&kKoZRO5Gx6@|$<-_q@xJT@J4t#X6Hy|Xi&uT)jZ zzM$;ey$HpP$QkyGZvFZ={3Vnq*J)WsuaSP8LB-8+6E=1sw4pZI9i?i6)daq15raoT z?37zvl@kQsK!9n3m*WQ=msydtNsx$1QmC$JBp2nk=~QX7)>aT96=Y$TOi-U&j9;7V z?GqbWwe9--{X*8_RsD+THMDo-M(W3Jt*IMgad$v%p*ZQmlF^Z-CM`c)NtX=6Dn!Wx z%N`(In=}HvJ-bqP)PPb<7iUL$rGU~ZORw^+^+mhTIA`3cd$mPc3JupWTizjRA`6FbtC#10VQLm%M|#aee}eni8=U?yzzgJmprEJt z8fo+ZLG76*MttAnXjY4kbm!k-l8`aeccM;+fUq;RmTNe$R-Q1PiOQ*t8kG-YgC`yj z(G-grY{M4~lkYXn)CT2MA$)kFAjljE#<(DnLL@)z+vJhpE-~v1?ChDU3OndcFH{|w z4!doh*A$;jbZ@xkX~anI>buo47=xg$P<6QE^kr!kQ&xeOjqMVjZ`VD5mEq^q%<@(# z-JFaaBtQv%ph;z(Ip=NJYUa~`=4~=M58whyX(~NrQhblA| z4-M%b0tJ2p+A5*e$I}j;h91ypTK1O%naXA-d{FAo7@DxVhhO!7Zo!Q5f@6n0OtR|s z%)d9i$GJZL2z!)RPdVNmU@(6VY>t6J;ABJXfiDXK)dvp}Cvv@O#?#eq z?mEhpnc%5c<3N!vu^4n*Suw-qZty!Sj)Efy-NuF<-LQ570YZz;oFsX>@=sxmLKB+Y z>n(o7w^GMMHPCk6hr!c4!QXdhaTaWZDd1*J@cgl7l%-&!DoUx)*fy5Qf6b-hU*=_L z(H0}k3-RSO7=;?b-D&c5AOe2Wece7yqixNVgJOZb8y|&X%eUjtX2>8uV_MT;t~U(U zlCycd0lfY;`8oWt^elt-hXo#P?ucI}T@8FUyoOC+mFIVkXJ_yKJ_<}r=+TKo;X`kAeb@2t`Oe!ljvuTp zLN2y00D@SsKlq7vtYQ-z*1c*zsz0XdW)LYS4#1vBPsw-{gKmvr| z$dKP0perhnZsbZ$+(=Kw7MW8nUGoK>$_f_{ByEF&09KEe)N z7`148cTpsk$!%578%W7wN0R+~;YkBVs@sR2rQd*!wo+G>S51L&Mt}TiO6(VBeM-pP zF4`L-1{1z}U`7|j;$!2L!s#c7kT1V_hZ1^!TpSKOQs{BsLmoXH15TJZ92CWiZ;nW zTNUrmsu-fT<~+nr@)nme^t`vq>vdD&rl5GArh+3+`__kX?v4{c*H|zf|G?Py-TX?J zHouNo>pSc`9s$~*ClCm*=-n5I9yMQ&tB7}sQi#`(qBZbvAa(krK(6yAAoF4apRS06 zs89ab#oSD`P(f~r+_!PhVsgV0&yOCxs`$H%HWTw=L6x>ynn7v^hJlGRFwsbO#TsGO z?FV0cD&WnC)<91kaJkmn0tF2n@I1PihMC^5EWW>28^G%I&KtJ^8m5hK zR%R?zi%e|y)v>0cs2|U$0a#nSDl`TL?j#`LS9-{&-GRB*IgxL7dcUOcPT8y;iG4zy zP(pngOi`j<^#kb#D_S!%(0ch^ERK8RdUFr|8}sPCxT{o9q8k?qvE_x0ew^T6*T<;u ze)+1eBFtlu8eV+ll$8_l9o)k=OzCsKpBZw(_Jml zYcEwT@`3?S5)4bf=;(hI&OPQUtKWC4`%)fL94YV<@BnmeeZU|zm1kN$pI#qJ&0`+T zs7+lJ==pf3A;SEoZPhV?>Suw2F#Y%qViir;(=N^U37bLbAZGJKImFdnQ!Dp0ns@xT?DD(ogDX<$U-upE_T{0RPTQs4+6plmBs&B4wwtJqH8G0d7Sffpq1VWZp?5-e>Z*HX@<5@6%1ny#ua5jz5=OZ*<4U$-W12&s! z(L+iD_;FD+498$P$AERB9vj}s zY1xWqqOLo8yoCqpKFTRz$6=9GD%v9R-odr&>g|h=vDaWASwf7_rd;UdC<_&oV)Uxs z%C40Xbj-oSJ%))N!-%@3;8;3(jCpOzdv#mF7DGI?d)B^-T~;4uq2Ul^qROQtI7FOf zDGjY86&IIK9h0HiUKv?d<929VHuupmOziqBF`-$_C{w!Xygo^4m#Fl>Iw-9E_`V9; z&ZSWtbpJwqB{93xF+tt@4c*W}gIE*8K5p(Jql@sq)z}CPN=g!7YtK&Xo5j+E)ih#Srwh*j7eDwNJclLYpBm6~ERS zxZbaIQ7%z~R`C#6$vWm5Zfw=C<4@Y{$6y(d$b z%8C6-hHqz;o6MB&pKbW@xaf)mg%zaYI&UAi4XcxGG8sl^p1y(ir2n*qkZN4-8-c~99w0HA;v$348=-_bd;ZnDDP>m zDB7v!h((KO*gKl6y;YWWn;f0y26HBxh0S7Wl9ylyHzI0G4udW{b>Q=&RX+cs96468 zC6fO==Vu%4x@N!Jbkci;%Ydb{AQcblc9r|Y@S9-a&!5$$*M1_ncCZ7yq}dk{ZTs=)RG%B3WcNK!-Fxo^#a{-<6!EzaVT7m zb8VmH1$(znRRmMd^oELhK63jW&n3*O8lv%>70hG5Nx1<#7F+P;(9wTu*rNpKz_ zA!fV9C!^7v<7wl%4ZR_eY7t1XsYx8qC4VxZY#ek?P5eG`y9#gSg!?DL*vAaxJ(W5# z$e)fS{(ECmasbOt(>QmDn#_Sd`M26M8-olJJeUmPi|>A(Op$@UqIDCVGW!<8YLqc_ z3t{fCEK|{Dh!*1@E#<+d*V(Nk7CuI-0BkjQ%ycq7`Eak)B8WKfJi87(EfWv$%D(DE~6H|N$ zH6R+SvX07^g-%>_#pWkAFu*nH(Sp)@G~H#U_4^j_o8{MepK>qbSqt;44x zIZjrtGwj}jeq#InoAd1i!PA5fbbG*q>+I{q_*%I*vF@?Xx5EiXw4!Wd2@5&)cYSyb zUe?oJti>Zc&td}`M-J97uA3Nco%>piz(im+8KBkFgY!1vV?$N84b^?nUI;Ci2{Dm! z*r(TlV)$}Tk}?l|U~jCWRG;7g_smm$0qaws2R37Tjv&Q|78k^e}Bpru&dCp==-R}vmUnAj(0#|l% zllCEGB{>eF5h_b@DSM=u90>X4#h>!er7lJ>4MmZ>KUI8f9bC1wM;7V8+{N2=9395= zw6l739tFW=NJh+EJ>}NPkOKCbV|R^BOS;u1iRbHi>Jg_JE#62Ukb#j75xDYSr>_^` zv&i3z+78%6dt!=t2sXnW*VcKC{wE9Iw~CXG=<9`Iq0FCXZG~}e@-o?}7Uz{NPwPKk zZlfvQDr?RUNx8X<_1TBn@{y~+2j2FX#*A_fe?_r3)x=(mG(Meytzxgbt0v_$ZuVqBP~=7U8W?Z7nPic3 zM!Z(mY6cyv?CroKReVCwkNe3 z4uKX|&IL+{tL-T^VMv@dM*Gwn0|ys@;8}?NTMn(0#U%*iOt~AWzZ`oOs{1H0)<*)v~0Ct1mWH@v>mYvqp@%HBRO+8@lug z;I^`B^sw|@4RH#0weZIn-;S(t225WMQFtVL)ei4@?h}7Lp@=q)rehkVk64ZmHhtD)S~FOVq}z&^ajGI5h+63kmCoHQ&ea-DRd44Wcl@Ld z6|nhQUHBF;^^c6@q;&*i=w?7(gbBuvvV=JkJm=&A5^R+`0x0%i6(f&}@8-K!?KB_z zQBS#*IR_NPHk!T2AlxC6QPwtEQFK*drc@vQqeJO#=8eG5JvzNpoTe0|r%Fn}{+fi!W!+D|a~|wsKZC*$YCgmMvL56Tv&s@k5NMV+$>7Y(L`hzh4w=ZG72qAOCh3Q z5s@tjKO9I(^haTtXf&Z%xxq$j`8L4YwglzK`7@^?(8ky)>sAsah73zYKgM^#9M@Oa z_A2tkgY@1tt0v)h67v&DuaDLA?g>?y^s#GwF{y5h#5$_gsPSIa4~yfm+$$xuwg*Xp z8iR(M8N90w>|=)lC1o`8vSkoY0!vp(5Vjs$<)k|ut-LpdoLICyhz#&xMQU52sw5%0d%6SYnvLkhv@B@I@$##5?q$v7>7~KQu`+7N_E~}HOTpFN z#JGws^kjTzjPld4p%aLFaHS_qV!f(mc9(M0*JYz9Z}pm)iC*+pDO0MF3prSRkEA#X z%v6p7d#!cF(K?7RXR*O+fY!ec7eN>~a?X{});~G) zNrj+xU!+A|fhD)@L~aKj2c^|!Uw-Z9oDxstDbvr|rfhm_VxYe|;m>G(Js>s_z4olWN}`?F4{U3e zcrY;Pg4R4)JKJI@n_>@L$Nm?_^lBQ@ef{0WI}kBevKQi!^6bHBE1l z9A)<=N?*H6Azw9QmUIiV6!Rr-Oh%D_JH?WL5khn=H3!XBNbz!d`{fW@v(WY3+b=8e zh!`=NgKf-d<$Nk!yQ8HLC^tjL5ttC<+n6Q!kf6b*YllPchwNVT(NdxuqWEfPV~uw} znOh)f{Qi%opf|W>K^mF01o;ol6$4n&S08&+Yf zPuUU}3|X(?ll5ASji)m~;-T)5)PqY9PtPL4=zqPh5F+z)mEba5?aP!E>fN`An6iOl z<5K~t0)i<1Unt(3u^*pakJUX~eleCN|1nyw#*<#JYSp;6Z~tk`%ix_Su?uo!&KxG< z_M93Im)-z_!+_ptM@-QVd>(P-lG$g%3U}jZ^3l<%S*)5eHAW@nRCm3*va=($@y95f zLIu+~FLuwdV#>=AIwf9hQ{N0ks+`Iv`y-$bbSr|h6JDR;dX>yHyby-qsk(+=)-zzj z_(uSzkWrBtB}oK_rOe&2dx#{OZ&ZJylsr#NHKVe;Z>bCk0bIv^n@3NZX4kFTk~|dM zj)tYVy`VkvnHXUdTcb#@^_o(+;$Ursm@u;lrCMJ~v{_v=l5sC_a?J{1`PD?*95E+= zm8>GI*DeSWw~!ww*roiSOb&ONRV2akN>L(ij5)w!Ey^;i}N=fZC zuO5PN&rIu5jF9Zd3<+t1`@&tdNhJVifECk`B41fRSjPPgZxPcM(>TN@vTF*_ zaJl-J#tCcD_8km2F22{Vw(!ssF@*#7GmMU!OgNpwH~YAW%4y*&-O$v*Si~2&yIiGA z>xG6nARkN$n+MEy7;oRLWl$29S4Qk~++-V>Xz^ z*P3R?G6_uX)N46yORf^O)#oNMgcfqx3b!CB33S`$p4MYIMe9)1r`6j1i`Y-q?4TNCMGk?y#(1x@uhehi1IWh-FDZ z(q#FgzRkQMYJ8HhAKFq!iIunZSmPu+1?vo=YD|l&+cV6A)7jtFE`~&_SHjaRRe3l_ zFqY0}1P&upNP293Gu%dij7IY@B2`57TE2YSj;2WM&>NG*JPJJ0wHDFW4$^|iiAC@3>MG!)deyb5@D6;m!_EWpK#J>Xz@%eeqw}p9`z2u$lz?k!6uG{>v;UD z6B-DIsDG}6L*`G)HzkDaYnkChoT0jo>--1}M=t6D;~1&4Y>I;A0ozd|Q>0b1gKeNN zz1>`@ewNJBlaHI4XknfXW^X>#B4(2l`Gk8rySal6cUT+*BHyH4Ef72XdTeYJhomGE zKymqNv6?mPGxGhRdkM{PSb}SYETFGUee(6pUbUXwyN4g-0m%fIX(*J>-%_cb393XT zuoI%M^}66w$%0A0UY+b*$=R}HNUZ@CO`aiUM7U4y3jVlR7E3az3~j$eqJ3~PQmv_a z>s;1%)q)?fZ9PQgKNTBXb&uQWc`I=5;sI#RTT|3LFHZ>d4o)q#s2ji?pOYfJ+8t{b z6s~M=T_W#Fc~&Px?58T!XyH(P>hr_4!ugp4#y%Q5)cgS8bN>|^lrJN|boH|d-Nc(E zv$JnkQ)=wYxqBZWXk5<9o;R1KtG^#N|8&`X6de*^GGRb=;nNHXrQU0kOVM|XWrB*F zH1sbS+hMWYXyg7p`};tCj({Vvl_^@;oV4Oz{dJr$QGsk`B%kZ^#wCCjGSgi)o78@N zh}D!FI!sS7VpDaw#%>EG$+N@+#TGMmG0`|o12BDeA3BG&Q&40e_rv8l4+?#5`2)ql z^t)3Z0AgdY7}|??OIa!XrwHCY>Fgiswt$mRnx~$xI~komK)L$v&vzV&a~TgFGdT-( zKuPnkJQw`*oQOz7_-e5eNWVR*>o`>v5ypr$_$z(0pjU~7>Vz3qg@u~rI{x{)|N6Vz zRi!6dZ{a%v{3~yf;SdNKHH3)Ln(3e*a0gA1Arhe;oY#l&w6Wm%DRBKC^SS=!RS07o zTnE4heR+2U94lFq@5)~Cu%dMTCJmFR#?!Bgs-*~gq-TzbscYIGHFIyMD0tdq`TpPq zp*ws8CMyA zY>x3*^uBX_=wW8R{xU!8_Ho?%gZ(C4q>9@wISSiy&RGhIr(Nk6yQa=BxM>kYNlhbt(6nv3hjS%)YZ1Z&872B_^fztcb>rRT25kI%$wJ({vJ)g%&>sOc@ZD*8< zjsw^11av=sPw{w@?x-XYJofUoW%e}yxpzk~3Avj}Stm%lIN&6uZpqU;Xt#AmYUeJN z@XB`p4+3+2Q%Nl0Dh}-LM88*b;oFAR52gPG!HkV*RcBZ1h9EJ9QNxL7td-pOjoz}a z9Vwe7->f_4ygBXsWBF9|u+y4p^p8Mjd8;a$(^CKb`PUc0%D@slUscX1!@b=i&G98T zmTRuIm48rAd``M4kti}l@T&mQ53e7Zgif79feiWa&tpZ|B7$&DwT;VkUlUWEYR9JB z>o)td2X}Q`2WRQUQ|oT|h)erIe@%k0l`65f>KcRR8*dJN zsI5dfIBR1aLj{umo+b9x&lPs_L0#;A6-&3>_8wx z>zB`h!@JPAgL#9ejLuPS2|5*?(RG0KrvYdi!=$fu6svfv$jEk`=2zlB?ltdc-ed?r zLMWFReW){l^MYKLFG|mHP+jgP8;vnygB-^bE>_W1S4rediTv=Go(EKavFVoC>GjV~A+aCfYlh19 z4w*t5mh3)L3NL~cp$RPGW>QTRc@?xQZZ^#uF6l+1xqgO>!u!+d$yUl~{XIe?+kTx?6u>VmSY59`SJ0X`6;rvU5 zfUCNuDURs~JKi#80Gh*GyVT`V+>n6Yw!eUTe_i|;J?SDpkpj^_Y@5gUroIoYaVBD- z`)Ry8kOY~E@7>?3Fx6NT1oBrKp8fqe(xH<=%O{7FGGv`!;pH5+(x)%}QJz84Qy(Fc zu%v#I#WSt2O9G7M4O%=ie^2BGON}bj7J6-I4Ktg{0nx~CE%ViArSy>Bt(kkT+x9+4 z#XOyP`eivlF9K)m>U^d|ho=cg!KoLe%BRMH*U?3lm$S~vUL=3r%UJ$oI0dpYULbwg z6PG4(vnICKDjLV2QeC#kzqFX()v*eSM7pX&Wcrk?(MKhGteiGHE+xd)u5U@REzGSV z+PJE)&$_7KvWn_G*6!z)nxAc3)rwq+iFL%T#&0miz}97A;(~8$$}5r}R)!jqN*~}b zdTmN!v0WN3ehe=56+}=AOiVkL@FZU3gtveY@`*zQ=EC2=CY!i?Ga>r22!g^k>d)b; z;mw$=gbL(0?lUOK+9!acQRHn;pKjb1<%Y3TOJ)DMYmEu@b}7@4GHMikguWrAFR5US z>&{FWUe3NF-#pB~^J6kKMT273CaIRrGE;X3b^Dy1GHJ47a|@>nvVAvh84o?HZG$-W zH{*9Vi#m}UJ&&1@eb%I})%*u6Qn?SPgL!R)=_-ei41CdI2Ylr6)T7Bduvw^^WDWNz z;oZHumBtaNp)y>2yW@)`DbwFR*et6^l8crcf>Sjbg*2xlu*F75h(r77Zv|d6us`{_PRsC3*P0=Auqa7mpML)cQb$;_NqiA}a%;JN&0ZeQZp6=RQbe zq>nsJ{e!tLEiYix#Bh7Dt0LxokIIJ}@<4j~X-BdhbnU{lu(C*8Q&&?qcNO&bzblNH zsU8}p2}+>M;$jX9f)+f>%a=YjBz5vPaaj0NbHyAJEwoj=6VbYyWL~(y#8sVC5APAq zLArj@{;({uu$O}tU45Anli#=>68V-~2gKZ4!Q8fN^fNFqcRL{|HB=h}G^kstB5h9V zUX-sQUPyuugiO9VhymZ=#H0?#=N1yWrUO$b8RAulXqa_IE0$$bG+m6UsEMmATVf)3 z5iV}WRoqT8n5978rL<*Ld$`%*GpiZIt-)=jBvx%pvh)@C1uiTYf;!v7lxliZ5|`gc z6JvCm=j8}}v;IakyWYzTiVErpe>@-kX#bwCgpqv=Eq+8e$2ksN*)h(aixhi?KcL(_ zt2&Y5eVh-?e$|gfLo=?~43S)M?UrD6g=G^bs(?^y@4XZyu{VkxGx$2SPyhDV8`Op4 zfF%zvKTkx&jk7;)Vp0^LjMf=A_n$t|WXVqQu#d9m2OVSqEawsn6GSiwIGK$PZ(BC> z9XUIqaw7YUb=*+uc^e`_ zkAkRKtA`4U6oB;XTiB}|x>2alMjQShMw~mDi@&x<_YPI(JsZCo{ZifFFL+aBft-AL zt_cYtT;yE(NaBwE@BzBfp&EW`Oed`3SH;NA6ND~M(%p*mI0d9)NrS{5^eD0cISQ>= zh?3YT_XK1!(Rrwy48M3{-d$sgM5OTZ*fk?`-0_|Wf={t;?rgEB#n}70M4T*7h}Q;A z=WQ~Wl|?9YfZZ?X2}4FaL$7>mgD2>%t~vW|dvU-pRe-tQq5&Nyx0wkX8w(JVlLv+) zusA(?r0Od8`*awv;T6-4ujqpUK6X(oL;skQ8&kvtz$u5>GN{eE{ro_`SpJ3)J=$T0#AZ4i%SE+U9EK6@e2>zfU_F?k# zq-4!GLiaiQ_Mu9L)yFTuA9r+F#jSQshB=Sbv+hS7_tf?P#&%1?_jlVdzb#A`Y^)S5mch)6%KPM>z1b>%AP42UCO38;m7Eq`fydHB^_ zRC?C*x4Q1@F+hx-x7tXYcpS6u0+O8m=?I`O_YTC3nOfXdvJvHUXTWaa@!tc)NMDYM z#GoC+s z%2QI4ClUBdI3S%n?Bh0~5yv|Z+WH%Rw}lxqG0}tKWYo6I<_y*B85TfC zY8WW~_QaW&S9YoiGo$@_ImV$Sx}WY=-)X4L)S8hd)Xc}nJa=iQLoG0-w&%%t&gIsv zT%*p9TK0qUksfTwG4_pVcu19S^$M7k6bz_&%mcprRgE>V2YW$to>wmCRqdV)NcnS^ zQ$`sw7J&et-)atFP?Hv3`5=Pd`n$6W6J&vGvq`_ib~&$qnV*YZO$fZdskQ~G3G=8#jzArnhB!YV)7x8G+8Vh49Zb#I7_ofBl-5$e=RZQRzP{{kW| zLgkn7bQ(%w48E7t%3}fvY#-_IzbEifER7WI4YK)-N(?SY$=bmWtcnSnlvRx1(YI@C z@fCN=o5Ll^E7JY%eB=|y^Oe2`%rOMHV(-$5BQ(iQU^S|IbYw6+XG~1;z*&n z>=X;r(&0NNw+M35&p);DgoFnzlU?3}(_oQi{xN)jEUJJbjx>uG?jSsnv@rNq1C5u5 zyl2>8p>j_jNpk~n_v9)dlzN|j_x<-+4u~iLf%R^hp^*XNf!jZP-pa#DC=7(_S#Oce zk(5HLhR#H@JMPHblTOA4;yjsCo5FPE=gYQ89S4-EvA`=z;fR;Pq*AFjr#z*5-(EqB z0R0{Yi0{&GcCL+eWnt{f?6bLjUx*g?H4h1 zKZ!sO8>xE-g^#vzP&Hn9?Qqb(Q?)-`j%Rc+hC`n9BcjC5!5Ob1?~H1Ku`H?U2TiK_ zSIeS4sXNqDy^dl}7xd^en{pg|n_owgIXXNl5uhHv%WM?p5u1YOxXb_6AQ7G_qSaE|G#h!8S3e|9YRZ(=Q#o1Y#}Y~OYX^UJd5U>c z0c#e)3DbbLfm3jJGra|5_0H*+rBV0mh#oVZxrFG}OMOH~8qnyMv~XsyK%I%= z>QN!x=j3dsnm=f!q{oFs@MVQ_M!Go_Vu9AM9mA;;%o|j5v{# z%m~y|#=#wo+`{{xyPvr3>O4GXHub_xSA%V2vRj`n z_4W39#JjZPK4n_l6jDXNBM9Cq?@NxfQs$`2!%w(jVzwdvLz76Nxj2!vi;9veXF-gepOwXjT3ZRmjKZpSQsmWzxTBxh?&nKPIQ*&LU zgdD2zG19aB5*^oj%f7-H@teXi=d~J&(yd@FgzR8NJYGPc&a6jE#l~}Myb!aA`-*b^ zNSQAo=xXbX8p}08)KIbO9>Ym_X1~i2NxBlR(our6TxR6-v60DCEj*Xp$s&^%-R&oI zmI`CJdZvdF{XawAF?V|GCD*U&Y{FISceK=BPaTD@3Xym-vV_Sv+K?l7&MERU=n){R z6FMeX+KGH|G-yF!Cn1wWN5*X&qqB3CUBJ~5#5|FW>4602NI(Y`K24c|-l(d&VM16h z1lOG?f0xrI-g0+3#Ic5nE)-VDregHhVcsjVk=_hdb~ZoWA-QHgjg~iObr-Ag@Na9` z@wu<$NdNlgxy>ikjTssMCy?}E6Q+JMu+F|#JI0;tf+@J3q8;Fzkk3XJS;xnk`>^%_ zPMtQb10SuzJdi_57EVNN+_*is9#KS;fJ2H1EE?Z4Gpq|p!bCrY<3w+uX8_ieQU(=~ z0};-R%=BOTP56cdnWKnTqNcrM;o)~Fo{R_)mVKkNYKSl>pRnCcELfMWdLc&F6LZw} zk;q1EpAE@K-&&2C>*}gKiL%*5&1E*@1h=AwJjI|*_lNoql~Y}EfpfPZ!^)lu@dY>H zV>`6y0hLl6rye#ip=NW9L#%8dzV-7|ae7nad2W@V(Bxa7lxvB+g-f056N%9{V4gLu zEe_yrn5R9Bk%qon_D@8TA22N*{B^Bvi&@+{1ljNcRgm~V2YZAC(OFncAUepWmTSDo1c6mFs1b=T-8?Ht?ckD> z6Exk*lqv3+Nczy$t-ccnNCWY>`3Ez(pE=$qRDDEt5D_>| zMzN4y8xO{W`pM5JmSr(9t_P7$yZQ$js41@UKd%9pFA2os{>`k|L9{HWY z|Nkg*#D$(#%zW}s)zGYje=7gWQ~mY-bFcrt^{zr;*S&=1)^XEn{J)FNHVvo$-$hu6 zu96#YCD@1bK;M`b-(mT#?mxd1Ob<1EIo}OlqT;qwY&`<_pb)sqnE!pO|9DDX47^W^ zXPg@g$8#`TPZYsP_K3O9=xl#8X5W)wM*LMUeOq39_dep~%}>D>lyU0j=MLq+V*JF= zzjUkp0$VfEnu6<3nJ}<(Pik!T2TyDJ(EIwL`?E}aTH#^RF&oF5qCIID;ac_yhZ%cp z!=KGnCw_V(kX@j6GX9kMUJF^<)Q;=^@9GVPYBs9mQ;|qWv+sTa68oOD@)Xqi(>>vh$hUlF;hw zo!QTq`I`5eYWmhH)4UUij&_E@_SCy8d0=d%7G4sgxa3-5 zuDi+(YQ_O~i(HN{I29NzmV5~hcl-P+@FB;e>mHE6P!=o$`2MGKNiwa{l=&N%I4)oV zQwzl=vxai-&FkYwq;k&0u!4`CYvVVkMPYp~Lu6lJ-f%vqdVW$Gg@4YZe{Kvw_I{~v z_hc=)O?L1so#Wm#3Ya*g4HwT_7=wIDFO(y?``xwoG8^GpSE9M}R`uQ?+dFDQtvS*$ zofASOyyk5tWXRs>^^CkVi2gTct20trulS>t@6UY${TjfnP;futbo)KP%7o`@8(^QR zHnRj^Arh<5l-Cl^lip<0k3ast)W7@fhT9gGvLd|ZTq6B)k)$Pc&$!qqtQXwSgEDJ+ zGXyvbNBgnG)tQK^v+IfkadLK>kOCD*DsshB16wE6csj2EPY6TyT=6*VDlbCmS$f$0 zxRoJXb1$J5Q6y>p`_U9-!&^H3aAk=R5* z1B2t@>7N(bfk<_zF_-Uns6KR`gXU=qWz_{;UXXmolK2+0ub<67yNQQ&2VBwe&kD9m zZk7Nh-UegQHZTQkehUJuN8bV41}c*qo`#L}qR5Wt&7zUh@jM`-27WVKTlJ5xYviL9 zJJ!)VfJFk2LP<+LK8TCJQ~bSe$q_m!{GE2`(hKmJ5I%i??<*E`v@1z>wN1og^>)C0 zdI)!Pz2+)5c8J|iN$x@|8X?t!32;9hx!c_X#L-y-P?NK4hF)pJoVfcDZbZ%{`MLFAyMlnOR#x^%(B4LvsS}liqS6)kjjHXEB$u%8^dNj@G~)5ik_ea= z%i&_$h?FsFxR(7DNkF9Lt_?zQV;IcslR$i33a)KZaejZFV=SB6>j4&1-r*6zSuu(d zoV3y#ld|>A|Fou|4_$F9++a{s;*QGoeqiFQSj#LN&Fl-j>E$tnw$U!+UA0H-I>yGteSeMoZ}g&UPhEOo%wj_nwe<=#)z3S-!%Y;ZkLXL`29 z!AzW|WEMJU6p%0RVAz%n#GBI*Bl zT_IiD69NDjU!r5saVqVP{~6y^FerHK^p??(kB9ciA*d~J_VCOpP3DzXfYbAXUibZ^ z$KHiLW9trW7|w9GbbtKV!a*k>=zploJ^42UylwsMI8UG;PJ}HxiXT9XmG1JF5F&+7Tx9v2|e8f9))$H9)P$q z|2~k@K2DVm4tA&IS|Tq^nb}2EgB8gxvi%euetA5bUidefQ$oa`{xPDbu0`Sc69NX zz)=v)C0=J_M-A`?vPioTxID2KWJUc26PTk-lhJpQU!(?Q?%~=dBzqO6K|3|>K z`QDSu+abJ-qqMnyl;47T&y=101Y6Bqu9Kq?z=A;9O)1OpAjjp86NsOxv53HAkbi~4 zR=IwfBR443qmttBa`or=tP+i1$pfGTdU`&4&r^BB_cLN_n20)7u_`a@J@avM6$=|% z9E`3$^;h5z-Tj71;q|@5`99u$lu<>-58)u^FR278}O_>#E8NOL-3|B`JpN> zeyG)JBVql5D9EtVthXGJn%rBlQpyYyX7|@?qx8_< zPQdNAQ0BQc6i3bNO}a-TASddPhecmtRG1(KW-L@!t>T{eX(I!YRd~-K-0bN6 zM<{rg`G3ACd8sjs_6~Frj0;b4{6DheBoy0c=|i$w~rO@?l+RG;1;w`TnPDk!rc@t`jYr2 zMck_xz)|F)fT~02U{?w_r~of!vI>J3wH`~qK!~{``dq!NX)ksEB?(*XNv%#!@HMM3~G+Z zo7IE8Rvj{LJVSofe)gXh%shqT#^32{AvMV`1m+VWzFh0@Nt;@|sY6xg14%^|^A(6~ znmzv1dLImMS`vRYe1hWn3v^YAJehlsUwIcxRmVw6vcp(1Q-j$@B$2Q8hAE<{X(v_` zN{7Q@Xl>uX?O`k1g*Jgf?#Y~C+k#OQsF?DX5j{*rvVu#y10FI`aZg%W%Fq}V{`j6* z%mv||3gyH&X+I_z-|WVNCD?*945)iu zkjL0Hmh!7wv(N)Z&x#`AG1yj)c90N8wrMLsB#&z3Sc?gIK08LoTk49DGyp%OX_W)X$PJVqAwE8e!#bY?>6ut;oH{O{T6hEjsvE|L-4Aa8Hba^9fX7O zwoV8GLJC3CZ1Jg*5iOVvgwp-#Lz_Hv28&Vc-K*B?q*Zysud)NFv9|8ZsMGp^?RGu! z^f4;26Xphl#pTA;84#|CNpMsURz6f;;etCZJ2zsuH|AoG_1o`u0eTNI;}L4oCWWXp ziq~xnnpsGGGi~N>28dWMu6u6CHVi1_tfTV4aMq@P!rY7n!d9KpQ)jd{z_5 zolq)izE%`1~Y8c|kD7(Au54x|fSF(d4htW+;@Vhx(H< z{^a~pBX6gN3^k?}kB8c#p_zYdc8tYxVz(KJ04erIfLYfTO-1NN*wzd%o6ufqBP9nK2wCEIw1sx+~O!nOlvyfffZ|mc{uLEE%#Q17W)8RM#CIv3%3hX zc{%$?f8=JQDy=OwH6}7VlrqLV9zNkv@%EEdlm!qyOKQT{OsvDfBGp-hYfH?MIN&P+k8y*(arSeX!&4Sho&-uef1qm# zJkFV>?!>7IeIjAl!*!6zh+uqmQj~79d2|&X9M9Dj_~UW&Rfgj?0t>!;#AN2}5I_*f z-FdBB$HXScB-@K@e-A}Ove0Z;biY3 zOpY#kBpE~=n*EvKz`phQsEl?Omu2;|L$W#FxXXw9vX3b1^?typsk2Fzgke{t`zbfy z=UdQjYb73wm4y0j$cZIFBm42);17Tjd+d5V^^r2w^=CXaQ?AsT_;dA6Qi`8T>-H8*@lhBa3%jglK|Fr9K%m!dw1bmmz#d;#BKoGlbbbeUDQwL;g(n zQ8^}bF-;*%W76o$@4zYo1`5gOCsigHAg&>KrxQS%BPaC-*rr^MkX$2HB9wl*)kY~2 zYF@V(+HH7`A$%o7_T{6kGQ019)chA)?-W?a_x*v!wv)!TbK*2M8e1o3V>h;K+qP}n zjomhCo2E%~C*R-wKi&Iy9&$2g&+I*GeYC)cEyxm~7?Hz_(zc;H2Va}eKDf~^+egX-K3-Q$ZJ zA(aBeCMVp9q;>TNXGMEZ@n9I1Z#c=9ntEwPiCN|v$yN6a!`}^W&kch=g14hew0@5q z_s2{MK_r`rk|6pm1fovOPaVo?9Ha1^DiSW`k-`BFWFWI@Mm}I{{Vv-*b5J^oL(r5J zR(CLTLs#w)*APU`sY!+B>eqn`C>xhF(F{eIMY)^WuEl{JXo_n_OO$-_L)VhsWvJ_id- zab~lU2Hm~^^rpHZzlK&L#sRSW$ zQeJ?mMQv5*aQ|8RT}dNFK+n zg^Cv;OtznDK$^)vlu*7a=4T!45=zvC!xp0_eHox^3!NoRu0y?NL=`6e$nM8Y??h(C2rl*u&Sr&c0Myi%Sg1iNl&}8gV#q$N{ zR$koIT(i&C+mJ?|N%#ehD=y`vOB9@mm8`X~_N_?apUaSa;r>f#EONQ;VSVwr=C9az zAB87xCqLTp>x@7$l`&b6;gz04MC0YupZj`Dz!ux7XDQ;@F$`^KD> zdP>ie&@buk(z;%Q>gS&bGYtW{Io(pvn~ac4BdBcHM~Xiq!9sxp3iEDDY+WZ<39!}~ zSV5FgNgB!ccVnBPLlDJ)Uevz(CfE3p@x^3PkH3$4!L9%hFE(_y&iIRE>h|eCH8Zcc zinCSuSl1FLPgf>{+DACUwa@lsKy8MQK<{WIuX0|f39GuXGHK%oo#%PxN5{xPQ&*I_ zfqfh;5tZwef~so4=E{n2>PXxT7$wyAs#p(WUT(P{*!Uy(c$siOyL_f5ADSWe#Bp$(e)~qITb1Oqxb5 z)QPjz(3vbg_k07P*X-d1nP%96Ve32{F5D?xFx&kCjSXEG`v5Gy5{?mwrwt%fO#4&Y8~UBdjAp zutuS5mcIM>&?I_R=NH4Cy%zb4QW}_LK+$H=NyW6WRR9@s2#%W|n$<#Lqo;fbsRGUc z;V)D>pJF2>Woz$kY<`BN+=@K48Ix2n{ZCj0l^qz;eA#%yls>B1gbIxcEi0$!S*^7$ zo&%hbSvH}fR!$XiOlI%F4_Wq>=JN({*07uGx50MKl0QgnSoFDqwBn3#{*EOdEn>BE zY!-TdF$-vkf=4kJujQXGHB!Fv?^~yJ*$snK&@zP%4uiAAY{77)HzgOEY+NL$kbqtv zpNI6p9?1+7uzT?GCUTAsX>0h81`q}cs}A@?AD#C#`v}$CT8@cjCzt7lB;2K!&Zif%;@E>F{QtJ=JQ{FUrq29};b3xGi)LEk`Q{ z3!bQl2teVClXWcRDr`SK`_s`_yN`)$?BOz(WoAPyBk-!Q=SR@VS)BpnlI z&O~SnN5jB;ol?)ADjy!XLQbe}G?p4ywx1{lcFmL#m6CNm7?NG)nU%bZX{FtXC866( z&4=rF{B)!-6;7BCpIjlzwVsh>q9gYjhT%(&I3BA%r!Y7O`l)lI6%>@*M?a%5ipO7m zyC@qY@`7k6M4e&Hra)_k@y)2v1oSN=q?FGZSJ~+IUmil~;S?SDym|gpRgG{8yf|16 zFECTb;3i=u@;MH=P7b|%U=x>-TNVY#D?0Af>FhF)oLL-8Y8{RhA6{Jyk;C3ASdDW1 z*VgaA5U?o}gd3j%Hy-XQv6GuL*K$D}IqG zC4Jb&l+R4MW*L!MjilzDcuj|IEKrMOY=VGzo*|McnjAX?1;?I#D-UdittNR6HPXxo zYo7EeU2!Y0EzsS!o`6Bw+H-RoOb{acX|5`cv(N)=VF9M1Xeb&;2oV=GhD)oFvM{~g167#A;TVKzY?c(oH@~iEv)0&^fni;UCKJ$SZZGxGc(kTXnnbq#>oX2y%3i-`JQ=QeuB5y z5=Db>caj5q4Mjk>cj++KRt6zBXA%<~!z{swnO~8wP0w21y=ak0{LeQlA4+v5RV%7{ zB1S|+!(1gBYU+baq`iC6eA+k3IBW_H-$Q=5#N3v0O{wyuW^6j4u`@s~DOIXr#;zmT z69h-|MLkRxr{4r*htz^CnL{L)GtNB`#5+#M4!1cH=JJooP7-x9qxV;E#Ei3nO3V+4 z97beUno@kCrlHhvP>)R%gIfeEt*J?4Q2(WZD=|o+k~OZFR)s|SKf?1;3b2Od8NJl} zdZ#{NA-6JjnXspRCDUHDHIwTWN{pJP+}EWxBsZRzPRB;7LDCK|C&%P+CKBm|DSgo5 zNfpT0%KB*HdcV=zfPgJK9BS+n!Dv)pHP;!%`5>I_fvip&l_h%RG}xl_bt;rm*_7(` zaj_jVL};whPPkcH#OM2Wpbp%k#kFXkxNl0yKARE;_hXe`d{X7_Rn_s*fjT*)7W<*# zwE4~yb)vvZP*pRMb)=rE;do&Pbq!dDa-fUKVk#9iwK>)lI&llTkE3A?F1L5_)(Hea zHfde5$k4>1lEJdB&lxi&u60TiD6^Xotc;IM*c@Yx?^V&SSO!L{BmC&~5iRupofnUk z>=)vozI4rR@25v6rnRMiDBmiamBR5lj~udvs$5kTOA@Ct&0tdTQYs#NnVlN9Fx(%l zm$aC3Pno8vljM6a6C$Bcl{5y&r%>yzXw`rCf2594$k;RSi9>PKO%b|_pzQdpc+vpI zsE?`Tu3l6ifM1ZuH)#-?0C6%DW>hbSaro`8k~0{enFCmM#1d8y60kcQhuMb0V6haG-MwQ+Jdu!<3LeDHAS^^2`s7tac-Af=C5@uOIm+4V0fGQZVI zAd{4fx`#;ExO?FIUlu?u*r|;>=f@kzPL%P_&bMOBM8_;;51Da*{^^$6 z!2puh`NX(N%lA(K;iC5y=Amj5a$OoIGMrr>%}DB5Kb<&vD^jE&=a7zREQ_JI z44Ip$b(ao|9xb?hbXQ;lA(tpm^O0G6cE+>5IhSHTp_R{_+fecjP2M7OQT)dyIyR2(=ayfz9k~^VdAnKLYiDBbx%-Itk(5Z4TRmm( z!lfcAJ(r~G;}L^JO-p3ChFvb04xc$UyD*qkDisdHEa?YU(vEp#G;1n5UlQ_UIdS#5 z18P2V!E-)pJ3-(W2!z1pi48{42&pg`35m>_%a-tu!h13NOX<}ZMX245(~8DX0)q`* z{Am)N)(<5UXOZ0sG;UHJU+Kd0eL3c<(aH4jcSQbB8Cf88?R|63ikUEqe6ot~SWA9? zvH6RUhZ!W(5$+ttu(4LjUyZ}M@pptO2A{GsbVgv^DXUWWa)EN7BuE<|+T~p=)1*@Y zloylem|}Fhc#=S3bpA!e3AVelJcCBzk)LTI%=G>ea8K>6RAl{Fdeoy=rhZ0*OClG- z1E};N5V05wiCqwi(`kg}dNG1{$Zmj~E|^Wtr;3+-LI#<~Y_=UP!LsgFAu+;1?XQLb z>5G8fE@HIsI7>ZRic7O7oC=CU5O9zwTy%xW**a9RzAR=stXO~@a>UYi z)P#V~F`>m45$CA6UkHtpVt&BFI?}Qm-Yj;1vVxUuSr_M}dTg~zGWh#`-czzmU}S!7 z?j{V@<|%pCy%52icwjNu6r|b#JaNJEYK4}Ow?$PBOgtnfqk1;X3`McMViXh0dp6}8 zKu{4>K#go`Y8$%0z5dR)n%YmswsPR)x3dwWFa?~5Pn#1W?NhWlvfK(|AO5Hx-wSTI z0Kq!-fX2iAjD+op>GWqQQV6i;kGQBB{XEAZQBnhuLdBz5KwQp&bi_40gg%DuYRt#aN}BiRVM(ux*`m4D}YF* zD0tR|CGPeAE%;aBMZYjuYl?L{Xc2T$U*2$&KT zl)X?8HWt_%3cmb0JO%{BFQeEhgD_y`jSA|N95^+464zM{M{Jn!KVv;MP3m3i)#ox0Y-_Tc~+3y-N~ z0?Y~E0c;28vr09DA|lv4=@)g3XgBYdN|fVQ0`#cYw z{{(}srw8ko#jwMNIMHy2-n;R75NXR-w-u!*NHaSw&;d?JA8aF#)Jftzib=2R@RKnI zxA^(=C1d>9&gQkIZw_H)=*8H$(+nQDg7VgLArG)7(L7}=!c11|Lb3`4> z>4qL#Lh*QUfYyQpTIXMl@vFQ}`<0YuJpGMfy%c7s>SRTC+5p!74zgBqphvL@>ie%Q zl@=|e?o{K|ZHvI-?4rpqP@`l;Oxd>eGS{w#vPq9s3a^MvgD){hM6y&-=zlq7^K}ZM zk5l(8qzhV^IiFOe_8sm8)WX}$;y_c$GHIxiDJ zX)hyW0;O<6i1eJi%q}+sEeGtAgQ zh#!-`#S^qFu2)7l#7+1Vr~0)>UXHmk{L=fWjFY4pjhD25Y&8mKOD>uypmjFOz+I>5 z`T&wPJFX7xo;qC@6Ci$}5Gaq?(G2>V60!DLeS@P;~sW_%j`la$~GRG=64iUZBz+68Sil)c0*NcW{Cd(R|Fv}WUP`xct8{~b+KICG zKrHE}?Fr@81H1f@b)7_%r6l3`l4LSh^kmXoQHBwxcH{-ufE@2>9*|W^G4p`OWW^uJ zstPB58yXu(8<4TfsGPK@K9Y%-@<`%P;!)u zRCKlr8Ev#=G$pp#lwI$*$`dV)R`^K+BB*jYs`_-9ynt-q6m85rDzsI`X4Z^3Z;U5y zgYbA*PQFSGHz?WfoTyTF%}&D+u3>Fhuf#~=7}&hw*kMC#nzoW+d;`) znxN~#JjB)!Ll3sFr@tU7){N967+;c6aB+(UyYxcuxApS z52-G%Kpk}U#n_FOx6(%_0!(0aZ2+3Da zA>Wkc24Q_di*+fTFGsJS}0Ty3GbIdU-aUgyC$dwz4kWL*4)&)vB}yhA37RY)5m zZf8?bX3|I~Bxwy)&lgi@cRfkWVhduyn6|kUJc+2j&`(QgZDr{CNxBzs5F5jZ{gNeH zrV{(jlBTgxaeUxe8NK*p)vgcl1qRt6T%B%2j9Px0WGjLptZDk3%%lthZ-t+CZ8t z1BF`?o?8@Zc{7*w?*^!bxA9W7V7k;yNRA4{o7s(6xfq@P!jy2-%t-FxAf5P?$d#hK zqG!O~PU4UCoL1ZjPY0&}WfuAD7sC@lZhY9mJhiK%Ja6;KBjc zT7O^ln#oLqlR2S3d^6YEpB4(VMhw-pup@W@d?n+(qajN_brJ|n3R4r&!XgYLfq z;rBI)V5eyPugmYRH8QUQ$l_b*#SJkIbg9i-F4Nil(|vFZOJMF~M+M(l9jfqQC^kX) z2~&4Abhlr^Y6A35wtbUkZ#a&Mx$f(_FBWva4_9mlO(RIo7)_%?Hf!DI0AKF|1$TsJ&;a3u6obkiWSRj;w6a{H zx?wAFj=RCe(TcCjuGP=QE*svqF^0Zj5gZ92vA+w(3G@BE%|6$Nh3*np!STBo&88^o$BLqsA zZ~5C*-LIA^LZ}t1-o)}1CUP8D%!D5SvF>@U>pT!bZ9N!hk3=ShANrGl4lopAPsgeY zZ6f@22i$%JKuGkjZ5HS-=Ne?-q2PsUt8_u;Jl?@f))hx9o5npnnes#VE~CHChLURg zb_-1Zt*8bzMuO<)&cpB&Wr3(rEkQE^OP-3W`G9@Oh-E0f1xB|P3eW^S8iftN;;ME< zo(EB=Z+eD=55pHA^+ogw1R0`~G%J?>4;a(o4X6`U5l(EVtuQ8vo&kb&=|6q~K*hEo>GU`B*UG8HN7$M+(g zoE*DTia#cXT;}OizqE-Nzr14CBC~c57S4dI`doJzEL~rJ;VOPAw{z`@t$qBjn*yph zBIx(dG*&|Neze>?QH_2HjfxAM0B=JKo#7Vd#~0Wg1=ItJ)Z6D5&;=?_{~k*pRFkq7#+~ zWNPmfRnm@tP?(mv13+2%=Y97>hHd>>{&|B>CzO}zNctz}LRP&7+20Bq7ty5}d#A7* zR)jgI@h#6x{`KY!9aY*U*%G^W7spWEIKtjcZSYc?aej~J>3L%y|ART$C@BDtwdtkN ztlMtU;;`JAU`E_3Fw%r{;pWYj|f*V@PbIQ=2L3PvIL(zX} z!XIGD`K@0A`hTtru|nZB`W5&0z6*$OB!x`FD)Y7`fBQ9DVKs2is9Wx4TY2l#6P9c) zV3P%<)19CU#hPIB&UKB|N$%gq|QqHF3_yR2|B7M@1GSm1yu<*w-#sG8sXv9A5*b^z=4SCQyxoqD>~Q z%SAN3N|!Lt&q~yaxzo;wZjf~)nBd<(y7qNTEl5RVRSmoJW8#`)kF5uRg-Xr#bbQZA zOoAvvv#ZBnhBkERy)3y5aE$?z>+yIL_UbAwN>zjEy;R#wnScR;@Vg>hHL zX$m)DTwLKp&Hl9Ue|Lr!H|Y7sQ|+*B_`TD*BPX|ORcx=z@SPf@o|Y@cKuxy3;t{%^ zG0(LP=Vbb!*q{H5Es+mCMNFX41;Wm+@Zzp&F22yo<+L#hf0;r!4dX8zqY`VE%QNe} z;m0jxNYw$qD5Es`aq1;f@Vf_#wbRC%BV$n1dJ?>-GxBMKILkvO)(YqcV(WWF`;A(T z{6Gu)23?|u6KGk|It4n)B$gkM#;;V>$?1z6)Cd`~v_!7HQSb{Pl{aD@VfN{x9HQVvI>IMP^n9V_+9pB6*e|i;M9T|r7sS1Mv-+W2 z;!sunaCrXlo`s4?3KlX9c{!w*oIV;Um^`wQ9*L1LZ!6f#A3t^{1X+I7_brJ0&2L%Q`d$CCXui73@^7uCrmNOk{fK5y z?S)hA2C+`)*&ScusIt~%DNH7=n%WlXd}J|RRTn{+Ibwa4GBI2vo)c$U$t`AH2o9Ii zC~~OcV>VJX{|bgR0yoUF${lJ(xTr)SMPoF=BoU1oC*7YEMTK+*8JQC-aI=-fbxZR{ z;c_l>Z4jG+LNP7$N^c7tdJH*S9(CbkX^If8A((EDP9bWqSRL#srzA4xa8^3=Sh=2z zJP!Hu12*s?7CYn?Nc@DRo_;bo%`sA!3Z1xfE2hXy$SIa>B2^T{!IO|J_TYxerA%+i zg@j~(>^$GQlCq2+2pyTGZ6=^Ao@KlemtrXLX$+nPN4fZ#j@#6Q@9Y!jsvQl5lcre= zg@d0#pU+WL8N$fw2l{8BONa2(jBvgw%Ino&UKv!~YECh)*lk=`27AdP3RHVSsf7nc zoZLwIzBs;G6x9!r2%HsCb!cAR!7EJ|=Rc!h2+O&m`TnN!#~B9Qu0zG7^y!y1u(w5) zVue#LhE}P^@sai@5|j}`BVi}yWz@kAt_l5En^IjtMk_XHxe_fU!#JZ^KC^mX&p+gg z`LB*5u6vXY2(yGxYV^9C6*SYfh+V0>9I29OTbeHNU9=#5hPEQ7V(KE9*_D>kA{za| zQv1a7G6HhepqTlitVW4MKI|!yv>uqS5p0ruS{`#%8S$yg_@bL~E)3?QYF=PT2+alP=mT8OhlCBT;c6pJ1G`1bfjUZ zX>1hM7>-qrz&ZA@;e@4N=zWisp|bGC@~1J)e;TpH3V}SU-y}=DN|$KMW*UF})l$xG zmRTtXT4sDr-Mzt_mLI|g>dj>i9Zy8rs z>CeN(#$VGJ_n+SXPx#=rq>1##Y8#>7L-qTQZb8t9qxf-D>L47Oy=%or?{mn`Z-=ve z7BzPw4+`g+*Jl1uVjDL}OF55|FL0R=vGolLpKl#M&eGhn&EOB0J6+9u`oc@c+C9?s zssJ4hvqX9y-Zl2L^4z$!d04IHq&m~?C`?Ecwrt?C7>8hvPA2prGjrHaxg_W0;;d-WPfe4sFML*d;%e-1!DMMT zTkSJJ{B${c@3?&*&ea976P?sW)WuDi0nPvKhOA>)ul>-=Njg&jw%)B3OtOt!?-l0v zf(hCSImIpNFfptN9rj@+;ltx;LiMlfpd7VWLy=O|&A?OycG?Nrp&+i3xlh2q z1nveujNZ>4mrc;z`<+TfaKrvieUz;1P<54@|BS{$bG`Bu77gWG<&!?$;Ias3l=Ka* zO(Tr!EJE$`F`N$sN zQjTO2=ZSjXKQUFjTa=r6H-nkVk2S&*oPI8jy3((9@|xxQ&JLys?Y^-M3?Fm2fQu28 zFp+z|M{-U)XYU{I5M)`I?vj101T>V|71~$~)%l+mIDSTkZ!<2ZS39`)XB~+)^BL2? zeD>AH56Y%`;6yBi;6Sd1fAO#B@LLQarQT4Kt_J~m##v|r%A@*wS6XzrPQx)qkHz}? z_vs$50oi@O_y4*#`0jPPUsG4PJPVD+Fk$HuE%W+J2=3nzE_jXGROCtHib*JzTt`^l zv+(W%d5@${gOJ>}EWt<>)rWu5eJp#vwF?4$i}8JO#tTp9r;wyg&MJDk4{aPvQaCbQO4#X_LjiyS7C z=QtUd!IJBJ;PX&+$?l|_a=wcRKBLKN3ng%ySP?K3! z`CY(-t_{C0mkP$f+^W1_^r3xeXr_15p==1hzchenWT=nfF$f8?>uDoOpe4SKHsq+g z`1gSwHS+ln*_h#3&_4(*4iZ*&4GY}f{DmvtCS1(JWN3Z!r$@-P59(z)=ZoXYe^Y1Z~CK?9bz3Hl)0 zO(<19uCZrW4HFqet#uuNvaC0iFIVbhzdI-RA>hxNyq3GhSR9>!*;uc@5Wqe&JUZc4 z)c-rg8oLpg|L9I%b??-^KbtwwP;!}!IjubayCVF%YXOKF3t&=Rq1KchDQ7cGKk=4t z2jYz0BZ>9_>s?`Ju)4tC1Ne_m=mAe7?_?Pb2B}*3s?-u2>}nVAw*vrvAopWc$ULP5 zJhICYM+;N17lJO5+4UYA3fM%=HHFSN$1Sr) zDi?J661x~gnpJ*}gkiV=$*0HWG~p;?)*cFS9oVO})~R-Z+W^Lr|1?TB?E_>@{+Mo$ zTKP;LUI_1ZC;o)n;5BCW-31@%j|46~FguFg<0J5+_5d&Ql1=LeD9Y)g2-0VOVDtb$ zomeg;^rlK)ouwgT5g9}zfv?E08aH}V5`zBcS$UHP*OA@1X9!S|R6UZFN^v-UQTxX)wxzA(<=CXZazW*KlO zxls%4S2gO&{8ACJnegr52&68P6zj=HyY!njohy0Y-GIq$`o8Z8r&IUTTCvT;S+jg5 z{I%^1)q_*|GP~EtZ$6&!@8qdEj|oRlEYt&lm08Lbkf*|2gvWp5KhHf5q?ySzED_*Q zw{>^_0ZjM9iRO`@``t3g|ZdRomW}BHEe=Ekf_wu ztuX;*iC^Ks2j7mzt|M-)yaP1x1;1J=fW7^rI0Zd7fg(oI4F8N#`aB&-k9{ZH#E5ZM z1MBt;wX89yn8Hrj=~UR+aWps*v2eQCgHj+kBoUDo4x`WrdkpDNDHwmc(p|hfEL1pb zTGtLv=ye0BsB>;gpV(ZfK>r<%YvP5P12tiOEjvJ+9xF-SrQOb{ahYKte@IP|;Wfn| zN2FBDV5V+1xk)`2e$t#Q)FXGi%8$fVHo5oRgBK;T*_(kuIf@(=XL{K@sG2XA8@tHB zHb&Nam*CRwD?|i;Ap}76*8TZmiWxnf=q@cc>KYDs^kW#W{tL zw+%zjvRqZI+Y@C-O~1me*heBMZZXNqqmTYU7-IOr4jcOIt)dRH1%(c!x7S>Cmagy< zc~C}TQU0?dbK+7bbgqvVf@Ym|5ayI_#gVDXVm6N@u2|%bvxHKkg)OCcbPljsJ)>3C z{{q_!gJfZdm}#$$&eeA5tpYI~<#89dDs&^BmFq}v$O2yOpmS-4SuJIhlsR?LEowIP zDuCj{8ljdIs_xm)@FnNxdqD&bnFU}FSy|>d5jgUovoJNc*VySzluLy!`Isf@08V?0 zicoZ^321JiDIZOo)Rg2Tg{9jFW!#lC54vv>u|=Fs)n<8puJDHq(S*+k)_C#-5 z*vN}y9nvkjt_=ZS)`=NlzjCn=#jih>Ot1!_N1LP}jSJmc7D=NfwuC_w2|i~8Df<{y zjmbMHd`f);rtYA9Yzq($2KgtEq6 zmRPuQM5gKPMp8&^uc%|A=*kB#Zcw)S=~;joi;l+p0v4&{g6};$dGIPV%8|}hY25`p ziD8-FVpp!QO24|C?$q7vb(73vUPG48O82{)a?Og@A1aRYx$WsQQge(h0=L~`CO+l! zBYcqc$mfkpUw9%g_j}`maf(Ka#Pg+Dt3RA2^_DBuBHd7;tYSuM^p0d3f@!msm54iN z^A{4W6Fnw3kFGekdAzKYiZrIGm%>2gq@{*6;j);O>;YafXTC|pULzq?s&uaMQQvIb z=6?;k7IKi*YdzlmoFhK)<(mpNlJ^`N<-G7 z5hYFtCreQWa9Q&|rq778@%o z{rkLmfEPpXNYlSXavGwlmPQ7z)|_o%5G#F*tDJYyF%}g}ky+6`iUA@)(bM!q9vR2|h>ecdX=+v)K&BBE`Ebf)QzVhM>valFm(C zFWWiIbO6sL8dfg`6{_F?OL{n}<%NZUiK9zUIxH%Yd#P6|0>t(qVHaUfp~lZb3ABY< zTjr^@yMZe;D8dJNoArN(M6;00)>2n?Mu01QTTDdPg&Cxkxbu-1n}y_H6(mZX&~$1I zzuRmU%aF~tgJa!ug;d-eU4~=|#;g(KpGFcZG2Ndqz+2U50hE)b;kKrWA;mO)FVCr7 zHTyRaR@|5DKX2#3QqBn+AGJxStk1hviF zQL8{csjghw>KsP+EwvdJ;{w%1L}aodVWy$P)Wgfa#f0Zd`Dn$WUF8BehU^I$nAiE0 z550uVIWTA=d@?ezOL8HWxn>%dlgMSI97JtVu&zREn6r^lcA;VZ&ACt71J7$f93a(4 zkR|056TEduU_l+HGr)NLHGBk8C`3Q9uv3Ihdz6URUg zEiz@~^F@hc<1mzrzB*|?^c-EF2RaY7$WR>LVM~5Kze+q+F)smqsn(G!lKj+^LqX|* zYB6*>Z*c!O{j`sDQ)JYEmn{<#Wj_IPg7!CrxX(D*p6%?`wl1)5O%tRFAY>}L zsL!gw5iDP{pwLif?_fC8TOhuMS4}0dRN0zNvMS0Vq9)*p{VRMq;~n4!WH5d{bA|-d zS=dljICKKBF0+DbBFn!fWki<2?Ku%0@_8M-rz|3EYd@)Jo5!k-;VbD1{F)0I$6ZH9 z?dV-8jEKZntvDRbA{P(~3ekk#B}>mfErJYHN8p+s8Vm6+?0-mqcft!`3S6RR(mx`n z4_{1tk}x@><3)O~76IgEU!TT#yOGU!C5k*d{97D^$p8n}D$<+_)`mrf;jdpE1BR#X z$<{FR?be^pbStA)nWp=>tb|t+D*udIH& z^KU%X;(tgWN(@;K3C@1JPfFmM^@VLgtE2fE{cwFf>$U5VDX1}U8^*>*=pwQDa#(>) z1xM2UXnIZ8@9il}%e?U(&jTpJeEojfroM4^!ah{W2uJf(DsjB>9QW0O{&%ot3Io+f zp77b+`^QEED6D4clEV}Yxvu=5wN~#r?x|4nOX4b5k|LyeX@TFHlzYf%!8PTdXqTh`f*#iW;6e{^j+K=axd;Ec?F`%6O zB6?tJLfFm}=fAVFAO5c+;BUv39LdjLKsxBZ_t>BQ9p-bb2ZJ7pXq_hB6_00PoOgM}+e}{aro1v%xgO<>8cU%?}71`gK9T^;&x^s=rDz zUuj5IsP<7|B2GNq>k87o)uQ~-+isOLd^TxmHl8IMMoNx9ih`xOQ*}KmFqi_#T@m#{ zg9IvwT{9ArE<0_AcCxoJ|@>}9)c@=B@R&l&1PP=5a`g!M@jp{9v`QNbNb2D2W+P3kCBe zP`T><;j@JC>y`G}4!)5Wg>LJ8bzQ5fxkXiG7j#C3sn`mM zTtt+EHiCN)B-|Vp%~*knQ8=8y5_ZQ=t^u!4H?{&t;!9RF>4ZqwVIzUtQl5X{`m8-h z`8Cq|0gwX(ZeYH?&_Y8YITL{+T`^j8k{S84_(_W51$JYDmzH?=J3b#s%oH*cUvj>u zd}* z=zciUJXyQikcCJ|JNp1Bq3f$2R#fyYQ51ht0*p9m>p3Lu^WXbxQhrq77$4Gnc7+(> zT*4B+)lX-XB`!Q~H`+wJ)7<)M+(YvB+{UMARLPpqK=hL_MdnR`*B@!>x|)>|arfrj zX073;G4|;{YsJ2_>4Xj5_=aIsDYr!wY#I5}>eyBH%S^RlD)=cu{@j%XMbwJ=)(WmY z7SPxBW8C^{6UP)q8FV2D;)Z`n#Dpvt+qWAH)sLgdbVYjPOpD>Yx_cg-aM54)AAiX< z_7qzRiB+Q|xOyiO99s$9JfrYBnLa#i#O2ot!&1zD5A~3$X4A1Zf7jXl9ACW#tyDkV z)zy+KI_MH!m`|ErVXFvnJTz1HWUkl>+@FD%bS2PwO0APMh0taX@apL z?Qy-nZ?GGBzzckMA$k!UY-j8t-4}4dAA=9M#BO^_J3Yl0MU>vVF4z2R9yyoIoQgC^ z%vbL$76qqHJU{HwseYuqkzwE>7vFg|VfY9?rC1h_I#)wa=48L>m?`Y>>ucsjP7Om% z6{=Ke3@Nyzb|L?J`Y4d@b0mtt$TTyajg!f&-AIa?4$5BPL=MWMki;&cv1-;|!I?NZbm_uzsC8A0Svpng*AU)42-LEc zlz>dzi)Cd)Q}3LvA$E46>aYAr(4CKMT)i|taD$MU8Z;hi!&3nr!`_p{L%C^>$+A(^ z7oA@I&px`(lX-yCXdRFa=q2c zV2IMw$gHqB=gnmz{ajyLzxWKq$Rz_wL;hn^C(B%G-4 ze_bjj4Bi{!S0d-ut|g$ovCl`SRhjyinjR-!=oL+o8dmK6`W7f9LpXmE6R`4K!y|`X z3{y0n=}M|e^K1-W-9)=1j!FQ%6KyF(y1EObAJ1%KH8K<7uweiToz=9ro z!wyqr6H8pXI%5FSDqQr8=4^dod6<@NglKx!l9$Wc&eOEX18)#VL}mvkSgG{CE1QJ~ zvQXx;JaFdv@4~bYWJ_0-uZ;8T>~|hwI!45dzN`{eE@exGm6*Yr(-cM zlx60${VvdoHc@Hp+ZEi~v`aHOHRj?5{<{)pDk$C=g>===>TCaXfa#b)r4}VP^(s4f zAWrwX%pM-syI8g6t|jEMICMXVos%Bg;Ir=%rK=+%Ch0|xRdJpr+@&Y=X@iMm&CU8K zT;Fd;)nY#yZf@wA4t8{wrtlM?~t zJ1mR;zo3E%CX5DW_OOj#)#6-N%Ac(>Y&N-n(VdZfb0&PBG?aVqFg+{Er|Gn;+doA2 zdrNABS<*Xd1RU?tqcTggM90fM4Ap@%o$1Oxc&=v{x|GjCo#86ZX zwoF#dHjcpTJg6zeQslS418@3U5|r}trut>u&oE~0AcOlkI-8CF$EqJdZV2o-toIvb zCl=>rcEBgZUxg%I*a*2XEDN}Ik=^BiwP$ZhM+u#1ngq!OrphywZ!S}<#F!zJs2nT8 zY@N=Vu^|df1$I2qrw#wE8iE7gf^w6YA)>bfAZXE8^O8xv2O{SrOsFY@IEbLd5}ZNQ zEWpoRTOGWu1hMD?4OxSubBsPY*kn>Hl+v;1i@M&i*~1UD{HB_UkXMDnmmgYQ>cg+7cy$ni*Lr*ijcA-Ji@r?;80Msw;eb zenR=3cT}_Fous2Te?JQnL^+|&JbRBxfb4he49cnsMezwK4Ly@8!mdxqO2cJ^-T1-KzCmr~dQ;|su6Q0$T3jVM7=5GE?j`eR2hAk@fH zAn)oFBqA1IO5qsIO&aMq+1-tMTuLIbzkAVojY|1DZ?G7o|4S?+d_k186m#dx%B-kgaFXN3YB+d0Ty&^eEO3BPKzwyUGKA4`Q*Y)e!t>lzdXO*JW4 z)~w=in1OvvIr!~gi%CYTZ+o&xrVR7|&cXu;EX69A5`4~T_*c---wl0+j1Bgnm=rcY z=tz5SRm1L%+Cjb?@=c3XgUd}p1OyVfSVxI)LSs-9i&BN3P+2mpok7W zVP_V`h3mhuOxhU1AO$-Ax&+L%*8sWfO)yjtmMR&Sd4ItLr^yuAS0}73z?8W>J8m++ zx@ZBsR~WhQUL8Rec^nB78<B`ae%h7}TlUm8|UQ20mcR^`+Q z;e(aHvdc*g;FL`&>0pUUtLU%jzb8wN8DCo0HJd?pCLA3=H>8B8Vbo7VPl;C`rRV0? z(S%CpTM$K@<8Z`1-R1pCThYzT92CK-9E%q#8uHQ&*nrdCF_sgca_oa*p~0Y`?yJVS zz9xqvaGk^pei+NhI3Xc&85>I*=&jKg5YqZ)9vxU<}vYNv&-->lkcM$eW&W2u9ghsz8(_|gdJ ziwupU_AW|Kwb0((t(>O;pe%ogT@A_ zmJ%$%e;myoh#G$Jd)k#Ga0<35+9b6}pF1JucEXkdbc;WqSc-N>z}hyDb-0&2budwEdiOPtrFNdoEq z&lKgrTs0ZgjwBq`h`&II3iwMdEaikT@y59}%g+gu9K<mNN*6?5fBj)vGN6pNo3u zj^1B@H@9_{^EX&Xh-`iX10lZYs4?Z*e@M9wXg>KXe<`|X7>bFbNd%%&niE&#{m=AhQS87-@6)Hz- zAg}YF`7HquW_Z|1qr>@fbz*h<+o0mR2fy^9?=m-K)0W5Yx3a?6K%>zSbPheY{-;HYVFuOWLz6s+aOkqTk;zv3FF>lW;1nnUDl}_K%vh_OgOv;WRXu4 zGQH+G9_lai9ydL1@vnKTGz5*_*|uMvdZBr%K~s*yRS8ZpP0oQAPDP0K+ex;WO+GRQ z<(aQw3U3N8YdL`7{pk>wMx8{nO~;EK)1ukkN4IB@D?7%F#HHuG_Ck);;!3N7<;N`) z!y`|tTaEMiwW|oYe^=#mv{KXC&;R#4^W!}PsJ>;`)R5G`TJ!zzT)ul3WU?d#Rvydc zhGpF(2Xas2U4&G254=Bbl?ye4vh>;1vk*4j&%TENv8h@^Y>HYrQAwE8`pM? zrdGZ`{pjfV_tk?`jSvTS9JxMC|gb}>j;H7yj;})aG=6I3N*2v*l|d0-CS_$t>7hkD3gppiZcwk zhz)g_0QQWsn8mqfEjf!uxvFDbkJ5%!k>s{_!G%iz09Hq)hQy`0HB|Jv**zdrjR0=` zQ5!C*f&@X=${scE(PFy{{TqN;Ud+hP1(!qS3TqSpL=D4jilEg-D*U_#x(oBXlXhS# z?AMX9tJJ21bF_PPL4^EKMJr1j;uwWOPY!+5Vj(kKx$pO-_FlwDTsPK8lK}8{H_aq= z43YANm#!i>o0Jit@hE8p+=JjE=^VwvHe9Lj-=)vEQi%&L{Q>R)#~EO|pYAVrSWH10 z7?ri;=I7&lbBLMaWWSF-RBo;9IOjZ2z`)dy9S@)y)7f`0DZxRnws+XBHH+Aqte}y? zBT^~GjWj+o3kh}Vucbdsg?%v z2Da;H$}hL=(lVyI7rW{!itv5&haP*n*zSy$!D##P$GI{9d#8m}q%Q->mX_i#SrOAb z_taOU|6>6_zT&9$%y><^@G<{nkqRtNo)Oo`r|Th+g~&Y4jYJg1Yq33sVaF^<5JDhH zt0AL^nuXKT3O~0IK|qpgSPQ|Yi{*}!sGj_!e&1NEHfcNhXTlKaK;cl&@P08j(2p$c z7T64l!R*HsZp+XW$*F3uJ_@|j@lO-d7}7;ILd^WSeU>Q`_qz+fANLr}LM?l8uHx59 zG@w{v2Mm4LpLDRDx|7ET)uc}iw*9X>50__aWEDc8z#joKo z!Hy^Wf>8-RQO5N3vI{$*%Z5MgDz}l~yh$A@x7OAcjGfmailhAc&{33=H%Jx+{TFXz znq6n=tfx1m1BHE9**vK~o(Y=4>c!b1g>}>L-sXQ{V79!#ZYsf%5F&Gq-A0-(vGO89 zpvr%#f+;af@fwhc(r+jV<<|?spy~cf2sTNu7AaM(pebUeZx_9&QA0&nfQkIqxwxvX3jdEE~LV0inDTE;0yU@6+c>s?pGu(02mxV#D zhKm#*BC~jzKm~9+&nM|%w~h!pfuzOgul)y{g#ld0C}x=jKW9*aBB;4(*t1|kjU(XK zwd)HvIWy&8($TUBF4$GS$`r{S3%N&#PBKD=XEJqkiu}3N{6>tsM_ay_Rb0JV!6t*O z)o_qRK)WG_RiGu)vH^ZR^An(7B3Xdx|xfi+Pv&U8Bt7H>Ppypm0H? zsD-}Z39Do%xF-&;>HnL$!&Q(`f6g^ka3|$mgFOjG=d0FIru(iegt~c;ZiVILP={J7 zUkqTGs_F@lUZ^Fsi(;$<4A~$n3d5PPNqJMpo9Gj_luN8+Z{%cLF^19c&F0Yv z3&k34wZ%yKTjl7EaJh_}$GrD*s!?3Qq(qJwpz0l?2iDf*>r!f@z};$(O@oT&&hDLK zsfG!N1gIlb+|Ebp4YXq;%5dmOl>R!$l7Nx4By}S#sco6tsky;1%AZbfB{KSEk-+83 zIu#Bot5Dtb#fS76of}2V+E`0jil9*D4Nh?Fg!4yPtg4TYAq-e$u8mD$C}%Q}V)T!r z2G|FrduJ^kp-5O@Q2ipPJ6dFowqWRRF{4EMC2I9w5*?|img6?e>e)vURaBNx z>BOVi+*!g%)+6(IjdqgD7gcESV_GE;t&$xjsd#;_;Z^0;6;jr<-7A!K)v<%0{IzD? zQc;cvl(ZgR^u{|1i|b@BE*pt>f657)FS8hX_&y5mtU6kDVTDD&F%MRmju1ARPu(NB z6kMY#nlIbaf`&ySA^I{W6{-+(3tJDJsoEs1GqV|uI!8irYl*Ki$2COdV>KRz@oX5p zvkaz9_f3O#+y{Tz@v35g6U6x2U8&qXG}fTi$Z1RQ$2Hhsg2sNO0wg4q>V7*Js!Xn| zYygC-G+X&e$+FL}3dvMuVV8|ykyQmQQAoz6PeqhPvyQ&-VN5;!+~6qVg*oZ&|A#FQ z$Yy_7BHnSxd?qafW$KtNzn;=Y)Mos!b8hs%Up*ybWemPUd zah_BhdWBl5y|`BTD&C(2@w3V{RC%Tk=M4aXaF0{}eeIr(<9p*nj`PJJ-c;cIDr22S zIY@)5#yXA>r?_auu-$Z%G&S#npf`6-Pojk zlU788Uxx`#?Hzb9l`Q#9BT5V>)_N`4I|(MR&`*{lnvDWJLvXi*G%B&ecD=^#xo;O1 zU-?uU3I@Aqu++0D=nZ>PTK%&{E%)Ngv|!sQwqbA_gf zD$@8pu37vudnm&vjb8)ei?2}ufZ4?YZG_GXHyLwNI`&@_E=D--1EYrpFVptZY&?UR zq^Hx#mr7ifrc>1(1N&rsd=f;MeAXro23gg#4dw9rquJ66NI7%y)d?zQQEq@U?C;0q ziMa0t%F-qImdP)*VX)j&$qwrE2n9~E?eiN$O7r|Up!U2&kP$`KS#Yg0Xgv< z=;B?$jI-3>G@>ytp7s|!XdBWE5@zucH;dGjw*nsL#{+O&Qg)v|HL|VMXMFC>R7}O4 z0$a%UEti$8lW}c+>3Q4eaY}=V<+=Jb)pR&zX8{hsK}f(=^qH}JZpVR}#z=hWB4T9v znOi6C+Fkbm;OW;ddfToEn4wjE6FLq)3#1@pl$sOmSatko+bp318efH?u*+vA9!HE= z+5pYiB(fLTcL)6AIrrmKIHrk-fpuPtLAyb+15ZJ^=w23--k}2+k^J<;fdR`eQ$fNi zXnJ_0Ii;fPqiM-MKE2nu?{})(lpSxtQS>&}qf_)Tg5(u^Bf1)Q!d_ybs_HUtNVFi)GqKHtnwvGFJA3 zQR*I69?E+mS=a&lKRaD6ubvf^5@fCOj%hH}F}>G!#j`=V{aySUC+ zuDl20A=0rE2>(kH4?-laU3ER;IRU|Qe!T&%T1VOt)p zOi@nq2&P#glXu-^c~Fe*p+?}g$}H5=(rdoa6EWQw!w(rNrH6+j+m==g0SQt9_cLXJV1tX=c3S9!Gy(MFUM%!45Yg@;$ zkMe4HavA^_O!&-@UfE$MghE7Usm*7c3m}L1Z zi1;FHQJiT-ji&%>yf|*DUcF6QC59y$$&0X6%6rlz%~C#2hP_s*|IFTgul6ifJslW> zI1>q-^@lc^IJC$gJPZ9Qa%kUXGAlx}E6~~O=eqYn2peA(hU}-}WQh2)0NV1U36}QW zs61(nS=zuRv%w|JWsQci*k8TLRKaT+nj%;kOFp+`KvdUO8<`aBLZwgp5hbHv8O9Vd z{aY9h(cf?M?=0uPU)VEcSZxG60tw)?(WuRTd*>|}k4DDOlJdkm>0GAB+v|&?G$KuD zi*(iSR{nxcMTb<-Xfl|-5yL|Rb7CB#vma@1t}uzJ0mphjpN+6ohvy00ucn`{yx?u~ zDw8EaaQnp~(?*8m&qLyZ_U;ylM=^i$-@ohdXKiKJpUdUbXhlQ0R->XZOrSkgJ1n6w zPvAbY;Ce{&r3d9MUGVuJ`&SjI7(z z;3%dMY0<9$2eNFnZ`4mutvK?5Ts?dF>NEulKrdt;sK1jP`k zQJ>c?$|51wB*imU@11+1UTfPNcL7(c;mIf=Pv(Q#G13V4>tRA~n?5?bVPKd9Do^em z$@?ghW6OUvETFzw<+Ff6k|1U9fG3Xeli~062ap+DG-8~lA7eZse=(MucSy}qMieA>THo6x(M+{=g7a%J z=Ka|)0s3|!0uGK~RA2(nU(Dp4Ls7w=5T$6;9x(|5%4Qabqw=N`fYe2RRMODOGt81{ zOhHIj~L-G*m9#>=-RAKm7CJjlZPyw{KUjTk1cXC zh9eiAYiVU+`Z`jt|L_C{=vODXJE(VK1W)l96fo*nA!8lkcA1+lnZdrVGR=naq^25M zv#JVU+u~+!R$VK6_<@3Cj?VW{K)8?P;kj|!h$W1USH3JOrxnM1r-VA6}!uZy>k*~+9{uHLiaa7N`7lS}T=AFDXcIJOV4!``*GX5^Fs=Irp` zth}>FqT145ka$5?|F%tGinPWK9d{Z=q?u|m3Ui6%?jT=Ed)EH$Y$%gOu>8BPkA z$q2(CjhWq(M+sv!xIdk5c<)#RHWJ;89IaNui>@d}y%KHIU|OU{hlBmoL&rCac{>WP z$&6v*g0+?^1WT+5i-U7tkG1Hb9N{EyMkrJ|n-MoW!@boZjmDIm3$Mif-kPA)W(2@! zM!<6l{JmwuW{K^Sth4@OXlA3ZU`(B3xF7<|_Fy5}wdvV;mk9OxE-9x}%=725csmzb z03^WC8+%lp_<{32on~@CWovblE6&9<)}-Pf@&>L>55b#*yaiZ%^mQCfNyPF?LRYsZ z#{nbE+B5(4x_0cGu0~_7V+fVy#!KqM3GtZ#N-aWaQ0*so$%<0 zbmzsi89Cj4Q&h392<6X7TXS`>nkl)`XZ6dpC(q;qdWDQW^6)Vm*t5|MTm1)$>_Rn) zNmHn_L_vzMbIkH2LaJtMLGcSrTyeQh?O-{Rs0V|0gf*Qm%F!Su7`dNjXr;Nz0uwRR z!S2i#%~7Lpx=(bfRIHI+Z6~&nx&4D|yh~sUBFq^=X#=69CEQZ$=8vfpxg{1}sG~G1 z0V%5W#V1@{CUOV#rUiFp%dAQaD<-;8B1WiJ>es(6u{9w;~OMN6+`@-9HVuW+!p;Nbux7 zQ?e^0WbYNROFWn5Pl6Sihl#uM@N9W`CPlQB=-z4bE%zvbWE9>ca0eqQ9# zO4fXt^k9tfeoAMI-J2q;KYczfE41&jU1yn;DGI$Vb}D$>`ZcYwFUi06QHCTecy5j) za?jG~Mi}1qcZ#O1k;`WC$a8XbH<`#P5Omz07NtS>#E{C+6wIdh;3u9HJC}&jzcML$ z%Tsymro}S9g|pcIWsM_|KrXyysm>U*MCbuwgYVN&5jW*4GHH?*7_6KW2}xhe7lt{) z4#Em?cWbBw>O{mDn|k|MEMUu>>CJ9J+OH|;R@Fh?tGEbfSbUcgRQ8p{1&vR*%?XXp)-i}1LnaI?*-Msh`{~ASE(x^m=qA4S%R$re)w9fDG2&Fv zx-x*0>JCS@6Ov38i#(u!q!nD9$DaL!n*_C*;=!eA&iVb2kxE62p*Esz*?h0Tm_h#| zW{qYJ;&jGMY3dmWl6`)ZF~YTD)BIqJlR$Ho4N&fLf1y$P>ys$yT^tM4nZXIEz+^E2 zlj}f88o7@&Ki?BdJJnRyeN=;9#@>|eLCAso*8#8CQ_0%YACTg2ZD$sv8W?kYRS-jE z7?HmY{E5Doj9jZS5}Jk(S0&(`{iXe)NKOQ`C-YsGI2!HeR1sW=bWWsVLK>0{yuEE( z2PjJis4s9kWjut%@3_X@-=#)enByt=DZgZA&^r`hxxj*&70^>+GVU532-z5vF)%h% z>XIj2%?=jJjc!4x6LUm2;gh7ymczp_5v!!z)4nnm*o0G+QWZRJxKfd~9KS!|UAK6t zf+?ebVcm&_d!^-ugdW9w3{eTMsBy%6hXbd{gOW?L%gMFZ$)L$$E!$WtXHw_?On?0= zgViBA*$kJJDew@skR`j!C;5&!EHe=;e*eeb+2J=vV}wuV;&XTez&?WU+q)RH-OqI7 z5a21b?}p;%05hfSlKu&7&z)h*uMRMBVM3LDr32N{cBC-B;(q{9P%j|6@<*ml;85%7 zS(`Oep1eN>ejli6z5+GLN5IFGDWQ0qFZwpt#i?*ZIzMJkuFE?wkGo_KP-!GWO(dp> zfAd`cKvcWKS;lL;7vP6fIhGuJyy9_Pqn|@#jKb@~gA=6an}v>mu`2f?@&VltQWls& zvOE0fKCvl$HI;g#AJ_oUrRZdcr_pGq#S20*`>%voj?hE??`gJ6)cjvBn^bcei*pEq zENbUy!u9rOg$a**uSPxMlvI%Opy>*&JR8VepE?1TJOE>!|Ys8Zqc2{D^CSgm=(w=eVUQ3|o$NnLu5r_qi-}IX@!=?B1r9&|{cy0Qpe_Z% z+cci2JeTYIw1bes2@O+`LUSk``EsG4#d(S zE6a*EFb`$;;J1gK`*dL;2v0W>b*Trvl;r?pivacmI3HWsNBaUNswyn7TX1Qi7GQK^ z8|bhc((pNvq>2J&j-X>b5H9+2LRyi?=L2V#nG^mG?2pgyf1d(_bX#vh%B*xG8|7*S zcJ|xf-yd}YNV>E|a=$ob?aY{yq}260`G33vhV!-!wyrbl`R+EL3jxg5PGD=voEA1? z@3ZDw9$t1hS%wlT_zq-@_E2>lx;P4RuBCSP6b+MXYqwAiO8_d~)@NPzH{iLdb}J~q zv6j-z9i|4VZ|S>CahT-jlLvpc*XsdYN?DR{sGt|sg5 zN8!NtFTf$t8wQtq&P-jnbVkF>0hnz|T|Q<9VhP zNjkr&1q!t;n4ss#>EC5*4*&}6axi|W*v5BXjd-yKQV?{2X90NU0T|PpKJCB;Wahi{=&q{JU)QYOZhYLNZG9v?2F-s ztp~}4+H?Ah?-5|{yG&o#R389m^7$SXd0x3VgY#v(3O7yxL(-0;JV#Tp>v;bLY;US zXMc-?t*6{DOd!@WvXJx0?qKNmy8q?aw3geio^KacQ7L5ar!gFaMGFNQgmr6cq9#9o zg!i;dP9ee*AcuhwP`jChl%xK&eH|r><2y`Jbyd0!-mdgQBh{sh-Pl((4V^U9i_}NT zbsEN%(*tTq6)>~f#qOAxv6g(-%0v@GpYmj9N@1+sN`dmo7KR6IT2u}iGR&b9&M#EkIb{(7$2>78 zUW@{e7FlSWc>1Y8nQtSr{?1r^1)mw5yfYxck?(+L7ZM=YVy8jGk+Q85o^d2EsWF^w|W7FglkwC?>{+_$_Jsn8+-HOQ6#?9d4u7;{h|$5?d&f^~-eU?;)9SLX?5&Qd7!I zAI_-4eJFM{+MlKW&+kVUdBe6*zVIG6zXAWI$&Err4Eqm;077Bu#&CU%5c>A0cz1yC_-ubC$cQ~f5L z;ggL($=cbm{ZNy1$M(MZcpFCowaV-`Z=xkoc`gwlv=3bBsWhr9yQx>2V=OX7snd+D z@?5$6hbimZ0Oh&79aI)V)UJ4I`}ICWOj8|8E|pa}S2;r%+nlX<>8cLuwWRRt*TGL7 z;0!Ecevk+ZugV=g2w@rl0N@}(NJ&-$EB-=NzK9l{j2FqDMUBIrMUNDonsbxQ(kgyJ zp9KkK)M6faFa!EO^UokfWi z>T1l$xMrr-S)^5upG8=~I~en*K~{G_B41Qadz$ew`;r5YN>qgeIP{hPZ-kEGf35}5+6 z3{o1**K7C`doi-b>04!Z%=#Vrq4qsEl&6kDuIcozNGwb{+)Ob))ocFogeH>Yz@&KA zs;#s-w*S3JNbQ6m{L#`C2rgzBjB8)A!7u-pWq%mt*!W*VS zELA5^m?I}@a&sn*{uZfoT$pX}=eYa}$DMkf z9+eIe<(*s^73bq}@ z=2$dCFlE zWzV)&wom2RYXL_X;wedN#bFi$hXy0$QldVGuIcq!760*iv7f{Af&AY|cC zdz4)Xt#L)D(nvbs#Q!LVt#LNoH=+@}A5%evGAEh1!u7Z#x@FF3MO0K(W>hgAKlu@q zkQ<52oHHNJ*V*RvW<4L~`0e1C8%V%a1rwA#Bk)HZ>;8VhGOkVbU;El5jmC!Bk^wR` zO^FNe>`0h-KI54LmBVeYq&sJ6l$@y$sfyVJwt=8=Fw&em?J%HL5Lj#A)Ta7y`9;;? zjh&_qW?1dhGM%&m%hE3ii*Qt}s)kZ56uQ#u2`KpDNgouDayd!l{_kZXRqUt~d=c1G z)CG0tiZ9X^7)x@7QLuP2A9YepA>4(y6$-372Jd^IQOEPvf9SUij#vh3BAj&<1nI#f zETK)`nu1}=%F*>2G8_vc*Nqmaec`+O_jn4mp>XGC7(+s7|8dFiz#XO&zh&eliDF6x z)NR*7^6T8AuP}w8q&5{Nl6d39Y@x*+yPt9;ODH>9i#32vBqr!Ifd}Ih1oVd(E{BPh zVmXa))5RODSB$UfPzobQ}A%QA<7K!edM$Qa9hlcXD= zmLBbtSCUI2rbt61Nm$Zm>SBhZW(nX~D%ldj2S?YZ8`x%*B=J5am=UZ3Ek@!ZWE9W; z2lLYP8Zj5|cpV%A-ZeQPbzo3cm)7b-?eW&|o@z>g#3XIXnARxMmN6g(t+!k)&Vl{T zMTHMG5eu9=QVHiv9CJ_{5gm6-1ZE^I)bzMCy=u~qfaOyP@YV^oxb zw(eoOzoK5)#=-EbH*2?pJlU=Ly@}V5X?5nDGWvvsn(t%~RPE5Tex~`zM|>r^*?(Z& z$uZs}@MYYannFtV#g|j`G&y$t5&-Yh@%CulJ-HeXj<&3Ji03>r66elP<`^<&&+8&~ikIh~X(LP!QLnMnVQjTg*G(&k+ z(wJ#b-8Bv|nI2D9YS4WE*c;BcYF|X}T^pwJQ%6=H7>nPRSj)vaRY&10*vK%mF@Ehp zu)GR=0nWhbP(UOUO>H~5cgDC|$KmlAr31n3!G#D6;#nH@P@*kW8le&Al>O}BVx?OB z)a~H}f9Mn&4oC`iUdMP9@JAhs`f43R0_%P<)#gdNb%+fHP;mh32n z-ENzy_yvqH5>k8@9JFq=fSf|dN{&2mtexC$>;o{}UVR^AA`6GOG@WuBBWAcgR*j&) ze7)(jc&-teEL0$%ED*H=o1>~+<0ukUtiVuFR>;pqebHO~w7|DL53l^_&Lw*UVD~q4 z)U3L)_9L=JQ2ze99AJ2#NcycGH*Q~?hQJT>o*$Pg7m@4GPCe-c3Rz7gAmOok5dX(Q z0$%k3l>H5nC~@{O85FP-`5)92GmVnhZ6o8sJIdbb)_x>DmNO(LIM_dr-P9#g3wWM? z<@Fma&E82T(8A~?5^&(?O6ojW?(0^O zSj8B^g=UNyRKIn4qhHL*rQU!FjC_Q#ECc^htTSJCg*Si_`U22}>cr#b78%meM-SX1 zd_OY*Ywo%(b|pB?n8evhHU%zU_XCi9^BmrV$;RQko|uP8Px`CKf(a_F@}Iuv<^CZJ z)F1XP z&}0``p}`d@03eTbm#8C7R-Vkz4-z1=oO?{U{WA&&cU9Rl5&6vJ5^~frX!CN|GlSYO zAqd_?DBhR<4p{h;^MvsO3F-4i@tTbuGJ#`64Qe>Z!)JUCGnhAs+tR5zWAzhRVd7IRv zlbWmtAyPq0GRq0MqFYQ;N4#p4!yMkP6ni5pw3;bj$0>zCS*^W?P_9nlIW%OuX2hDe zQyoU=cEHJgV%q2cf}mMAx&$PtET4cR4y`lBxqDoQ$60^6S~_4o;VK|nh^l} z42hOdy8H_@eJ)iT&qN#)Y0l#FI$L3P_&}Heg)i1m`fvW@%7NWK&+Cnt?mP>cv@j}^ zKHzSFBt|7wfy$$4>rM>)-=;IRqV>NhIE#1|!N4#)Aj1cm4E_wG(0(cZw>8`9pqtT#aBob`aP0lyP{V(onn&Yzh2w) zXs+HR;n<@Bt0lx34(d;~H!;fZE;%|Z8Q*Igbi=}kyvw|dkaYPV+>9r$$aN39nhz8r zspiQ?&bdQoAhzFT*mu_ng-tX1YSt10Nh zpE9R?aQyb(fia3KZ~+qqvZ#O$cb7>W+5X@;^O)kKQwnYBj;#A3(Ih%^j_3!{w$8hw_i9AJ1cqgYQ1D90QaP z{6@n6p%B(SdCdw#_hrqb&$5u>3$k}AG+awel{k9G#P0nu%ONw?FehNFg7^v-#sgOk zpd!=IKr4HiLwos1cgUwNZQuSz(IESiu<%G8ei04If2p)PT$7c`o;%6Y5HLBFp~^#T72ktS09`62mXBCC2DB zgEd=Ou8fd2Q16`ie&s?#kju?FQlqDbS>c*bA|VT~v;+pr#*aB_z2x6ji#L^B3-LxU zm~-S_$lc`R0$ifky^CCqzy402UPR+v@0a`STca)ZxBu)DQh~!TLj~LDsl^VN!}bpY z7u|SLb>sochclWM-{7$!)$2R~4Q8TFR`*%RF12QsAxDlM%JG?M41V`*=o>Sh$%{50 zZ%KoOK%%ZqWNak?Bgp#m9&~|Y`!1& z<1(?;mylS&Y0BZCaiF+# zwEMn4lk>lC^!lRdG$Xbkactll1aoeA!Vze6iTgYhENa>G zRD;`Cnhb1&A0%UJF2ay1?$rhNsNmcS5g3%+et$Ss`*hz62Q@;eS7=Q|U$kL)Y3<;( zxDjHktEg0TZF&Wf|C}*RCAzX!8%%i~vE`P>m!w-6J+@3X#CQ4cLt=qS_(tc8VTXe3 z{pK=-5+B(-h*^AcaVy6Q!AHzAWz!aEX~sA|(6E%-zL~OMOQva@s-Mm7V?o~v#RIc} z3wXTnxCQTVsv&_O5U)`oNa&*$%$m(!%Z}PFK<9`0&i5joOAb=%Y=<^2A18S`hiGlc zB?{Uy!8x1yBDGi~Tr%7#)#FMjC(}O8<|Iup_x7}$Sg>0@X^P9bf}4y>$4P3b?pVwE zhmpk|`D-Z8IW$`#yVQs>RR1L7NZNvhPGkfwg@n&DbU3C1R->}x^IA2ub`tgW@aoVc z%gnt@%0RJW(vClP^SQb3`9UZv1g#M~~vncKsIN;j52#+ONzlJti z$7yKBhn_4^{vNe--M+>Y;9n*;H;dlG;soylY;j z?Uaejg38-=moU!LdTpvL-NamkwvC`{O)|src_KSM-Kfg;xUI%d_d7WK*UF=M>}t#I zF`J#C%xP0(j^iuyTaSMC&tLQHQakMl1dF_M>;GkeT!Mrueplq_ap$!D=qhd z+{dyWuWp02%^{ou2S?tRilstCW7r_{*qmE=+Kd|&E8^pHx?KCir}%Jm?o;W<^}8+i zeAML#Q-oG`=?`#=%N0v>$`|dO!}=Is}r#w)FIO8AT?L& z=NZ^jOPkPVivuV5i-R%q3qVkZBMLaetXwRzj|`K<;}6#)@o;oVyhy@-{nN?To;lK% zbVOL;q_e*2Otw=|ez+a{fqVW%;bmN5J6XqGRsP%mkP4PClPq8l++0g)`rN)La$AQZ zXI|;l4#-R2mw1po74&8Xg1F(`;{x?Z)O(&q42U3CI@}(pjTHiI>^UK@fDDI&F0@Ou zKkbV8Lfh)vTIS5ZKRrEuHAysaLnR^#EPNwQ)5+%`$VY&vJ?Jz-NYyiD=3EnKIsIS2 z6!`W@<0qPQCOC9dEoqAN2H!I4_u~Y&YiuOLYiO=2YpS^9|3N|MHOq&23CA0cm3PM4 zOKe69RUOcp)s?Ri{|`|D1u&VWr_?wkJ7BoUELLiRvPvHQZh3tB@HGY^s?OG&Iu090 zyN;SI51JUDr%^{6Lmq7t`me`CZbL<5S$+5`OM9(F6mkeLg6WX-kflY)6>HeAsFaM$G*2%@kbOHTG|t#N7_XEUGE{hXsdBXe-A=XXK3IGE)e;d$^MQ1nEcPSBd4n z4$(TYSMFm$52eg9GDUuJCyV(e1?ZmH(n_9~T&prq{gXci#lx68RMVzy`zv-tjQA^9 zD{9H+30U2q6A#5AA6T@_H*1RV%@Zluy(H+6sm}*iYj(c2~(|S95@QO;{y~r$BDm?dz{)59p5#i)_5IZ@G9n5LPNv2 z<(N{bDuXab{fIY4T~8 z-I@sK0p>k^XXNQ*lpXKkMp>-cx0f5w?I)l%|5}u&;LH&hYDpCCn$#kZivsEy*p+s9 zWwl~X85A~8PO+U-z+)Wv!ODGIUEaW%BBKgQyNRBOfXzc*NOezZ zkAzsSF$x|{?YoDSA8#wtJG3^XRJd9g4S{{7!j~DY(!N&_a>g|iqO*saXfHqya1^V# zO0!dc`uMS9x8qkW75hN5EpheV+u(#Z$~Aw$mo98LRDv6-c3rq{*%REb0c zLr1~#g}M0ys(1zLI0o|pifT^&Y@%J7)rs}hVSDLX|)LOgjDS^RFGVB@+a=l zcGho>dMqD8r+UCG4*s}`9<;DtgQM3+U!^z1m=+r=Xo14tvqINn-} zRqHU&W#0ix`GNWz5=*k^P>f5x(pYtXs|Nis14NK#A%l%5Sm3aF+HoV=3*lQmMx zw#^U|1X{Li@QM4;jnQY!yT=)UUBDzNe{>AdARGj;H}}?P-5GsQ05#=i3>qbSj_P6A zCrh_#EJJ=5&>>d|GGp<+M67;3@tWq%3fsW)?Eqf5^ZaMrUi<)s@ka%?!JhVIcTg|& z2#CN?z6|n0{|&?Uir#+!vM;!$#%D>>;dt`mgP6NX`o?`xx~eGj>npPVcKMpmpDuvh zv0_u(g%uiMlO~8^VOR$Ot;a3N!y`k$4|1lohg}*w_<82S)^qFwww993vK)ycTsqlI zs+#hN4LNjl{HMx6oB(dJQ?%#vW+|NunNRzjujnt5I&;*g3APzNm;XoBUkAnYyicGo z5Fog_y9HPz5S-v{!QCB#y9al74<0nQI{`v)cXxM(d-n5vfA9U{)=*Hj3tNZXIWy06 z_tQPfGf!j<$A}F%FFabL<-1|UG<{1ZY;*(I7xLgk1T$yN>ZECBLP4rabG1<{d9t4Z z1gLX5m;yZm9N}9Gqv-CXYlAVbWU?ig(enj~HEXkGIXMhzCKbP&5AS*amhrxYJNxTW zWpX!I#ZnI1@?`LNcUL1IbJA;(O*tQnzmUu8)6hUk<8`}ErS2G|WGuFe(A7ZDWk$@Chl zD!GUG0Iq(LWr+c|e~vZo3SciC+2{=aw5VjIn@v6jv>aEp#T?T8ZqTU<`av~1B>>WE z?W#F#T-HK05wmheD4Fv)c%>x6{?Azb5YEtgK}$Hcr90RhX=Yc|1$y^Z ziGXnfuShYHq@dm~q!1$5l8M`UrI{Q<%&lnv$2O$ViYVG>5{H;9|4Z-F1hj> z-r6%05F+$c)>@0$SAfn&xNgpLFW+p2k(W5z_vtW{)VFzezh77?)oId~BDyJx2OnYa zl+dkU=#QdXri&>PH~k}e1OdkKw`h#v_U>tiL8ol6Pn{V~HhW)mE`GaI9M!}NSGGb&k+D&CONKoHB=VLwOHGrFYwNh{r80;|m4N>|dJFJKHHp5($kFfY~4^V&+_ z@)_zIS}bo~pM z<}5A}%BgZ!+t>^Ye|ukS8#iCB{q^>)Dwyjh7pB?9`U#e08DT07a}R=nZcmYhB}qD6 z4E=*-IQ7pWBVj5jNCZhNY?i#Rhy^HwQoLU`s1F@y9@BlZUez7vEvd_^`#EiUfgXbd zPW|IV)}!3_OK!RN1eCbcc=QI$8g=1K@HO)?6#QtPp5y>c4Jvg zl@R0^*K(IXpA%`hI(Nw5)JCpzrp|!utc4M0YHQ{L(r{zZ_Iw6)QyV3~z8I2@AH9td zI@+IylqTqB#M9A1E~meqQPazPZT(J5e9{IE;dDaQUEQt1(Mzg@k5YXv3y_)C)T?)-44-%^OOn;`;HbGDyM-DsZtMNw=>q<>BI7Gf4NNq81f2PE}ZNP%lbH$sN#VN^Y zdUHc%y|myH-^eSndAu$O%nEP_Dv4w619Y@UBrj8-*KgrmN8^4t7S5jcVb8mOH0&n& z1H%JfM$%a}Ixn)23;_RMIWjmYH@kTUY*B+e!oR#22O&Z^wV;^}mgD`1h@4JpTa9VPXom8m#d>@{ zDdz4~Y(26+MO>%+*58Q$YP%u#{VXuF?gX`{beF7a?pVNC;0@nE&^EK0!yaxiy(e7c+j5i&YNHqy4wse}u~B_POO%L#1|bHS zjBLfVb7hxB$t!Ver)Icd_Lrq~pd`~HVob8&yB#+75<^<53D?VtrjaqbIMUL$3NUZN z^#s?(dv(Z|dAQ|mB}r6vn-H+N-odj8{SJZ|jVDkmGBr8-k*!P=@6PF17&6Izn7EDh zOtSs?HdsqA-d-h;;aQPm1S4Ov_!s7%U7w^HOob}rI}T(sk3k9h`JE$&RzrADgDVnA zR~Av7fqbT-MH&op7kRys69bZjKDjF2!=P?O41+q$Ua&+u(2T>TLS7w`X-$F(U#%5J zzabkP_2aU}iX^m!s#)2LE1>(ly8#QBQEHIXnmzj}-v~NoIBdeo`kAeqFc*X%{EMaJ zTm>`e>m9QtdcJig=s9wFifkN-8SuX9bvQMNb;kIAeBx`X%(ECN9Q2H2&Bl?J&2qsc zyf=bd2xy3=w0LK&mr|D=7mN}Q9* z>LYTGYehCtueLq+l%5XcLrUCIP1rFKh+1~68$wdAEEsyU-BaW3U=ed9qJ`$R&Too9 za#yL*_ME$LBBQ3`p}WQ({XSk~OE!*7Q`$pa6AD_(j3BcGTQO`_J3d|P-k5zEMh>Wp z>UM?T8mX%ra2rht4|t3>E4}4Jvq!!ui=Hfz0z3Tvw=qn3g}uf0ZGJ8 zYiggl+qmJ`$gHjn!HyttR#u?ba2G6uRO9{{>^O;Gx;7hSzueZcLhdT<7U3;%@&NcE)OfU8AP7O?JN>0}>6fP(*whbR#;Ar`~k03xAH`F?w5>}Q7N+(UTmwB7Yd+3277eU~5O6UtO zqw$pnMVj%75i~{3xP#H`m4jKGnjnX4k^3Ed4}=h3P@mtggXV%)rS96AtUQNS5Z^Sb z$jg*1bBKqMktT`76?>-fuO}=d_fQX|k<4xk$OYHhfvgBQh+kpftBU@L+red_LciH< zy9-<`)-IqepmCAHX+)jgd15dAS@PJRSMd028CO#aR^>K9S>Xg&Vr*`RCW2)E8= zz~8=d%f*n zT1>1bzs8gRtEvaO36#^QS};T4JXm)SJD~OklpsrVX;`5~i1#~>%d;u(r*1OG$Gz6B zC%9_$Mo0X_u6bSUDUHoiKl-aa*8RE3HJA-2#ReLY^@DuKlN;Z>C9MHh$*Qp#5@Me{uqVO8`yW3vs@2%!OZSxL3uzSzU>y|;p# z$vRmM9R3!o zGvBVa^?dQyifAMc`-1R-ibyx!5N`mkMj^HPe9~EbYOxpmzAm6O zcZdd`YaC`W`&0qoCX`rtD{bCpX>6-~s8IO=NTb?~VZ}n1IrDt${@Ud(+42U|zu#4m z{zUD1%`4T<^8qq@6edw-FBcrJPVc8)YDV)L0;*Ok^}m2&0ZS)^c!^Ad2s`&aLf;(; z-+seC>R(^gO^l+Pk>dQ_+C7N`XRovN`9jrfSmmI4nxxh5S4zQ0m$Jf#d>gLVS1@HT~_=B8lS9+tg5!tGLjI zQKnej($>+UGASp=!E?^@;f)sOKhrC{%m2OAT2`6uEEhwb(kF7C%5ojUTz-;GeCXhtgd^1C zXY#`tiFIFn9CIr+-5$GT&wcy{Sn>CO+l1+9Bqb#=07Qy zXr_BcY6rhO#J)d_zpjjr=Grb-zYE|hz95Xp20UxaWwRu|j=!CZgDq5`K}@nFn-#mi zDQ+e)y1qjWZtRvDv(+PTJ!Y?2h(#LIIo>uUWfy_fwjMn)Ym4J(l;V@bWS7B5sigzk zYPvazZaP_IGQ%m9-N{9$4kSQd{reT@y><&p11e#5M15zrC8@*fW3@n)1e29b0N6!HA2eW3Y3$ z?H61ALr-K=ucMB+tPG54YuBIh1h(IZA+m6?+%FedcOH&pTJG9h>&r*u+`c!yy^SFl zmcWsmMxCC3&GESBmgX=GkSD9=a~~aq|B%yDC4kSy*$~9o@)WHbS}{*fEB=>t;YB3wsF`iLd8t!hBR_q8Dq}kpoG3wH?^(1wR7jWEC|=?Y(L)+0mAUYA@fJ z2hMRy?-ov8XhO)p(#>OiZck8sX$yGm5SK5gyc><>D}8z|%x_Wja)rhkOH^I*eSQx9XMaWiD~a z=b6V%14Ce`xkhkTD_lM=Dd4Zm-X3Ap*q_z6gFHC7eLm^oOaa>{{c#z8YPJ$DM}D+X zLMoTAhaz8)0^YArT2!BWZwr*l1cWMC2MjC<%C&?3R&Y59L9p@%h}g=u+zzo&EM9HF z5i1)AKy6gr!Q-1#Y`!;>RpLEk)`I?*yySP?9XxoJ&Lvs*#kFi=Nc z0}OjP@|tYl>axSNgjKX=E3}$$Tx+jU`4HPdbxaEpV5Q9Y>ZQCdn~H+5Hf6x`Z7G(3 z0I|*?R%BHjGkns~;DVIHOJp|Pus#L4OT4A~A2CxXPxa~&0(tEtTI@iL21_iAdJa{sVJg0&#IKDplb95?wMcS5rEio~hc zSCDI<;MfP?I(0Bne$-#Qpm3vOz@UCu3?n0ZrZ#<_IVQy@x9u z@lO#PC7o+7jo%+H&7pZ?Ay7kYAO?iO*NzK6AIc*+A$;n*PzAK?OWCFiPH6BLV}-&M zSYQ4BZ`gU+zljU8f9f{iin?YE#_|LeeO<%%Z>^(_Nen!7(H@7J!>Eg=wh60A^jG?b zdNtGG*-%X#UNCBCKuS^YC{{x$pL%-IQQKEYfg#H2K}D9zDmGHYV$UP#?6{G15thwP zvEq;9@xX;Od_VLO5>giMp%S7!;)3*$vg zld1GaWE47+Tg^K$#NuQj%F(Lt*j*&K5PXL*g=ES-rH-hT&t5GVup695h0+E2{M$*V6yv_S$}Pci11o_*k0PefrpMVoK4BP zwN-1A$rtD<6lgO)&F03DF$m6sy3Bq%(nYnYvs~#+OI}|>pAr}mEM%pb=>{pe9F|^R ztddjuK3I+^R-XX9M-G&Ic6YbzB?t;}$)TbUwNaWs#v+m< znzrCbEBxW0%8{Lqkl9W!_rmNy*KrY@cB^*aS7HsXU9Ux@ z$+;H(p|MtFL&TTj!fhuD3c~lZ$ua#YD&44;VH%J&6p-hVj$oqm(`R~9GCMZ-zf$7A zI_?&2q4UxE*sx7=4P4gALVJ$v6Io+gZZC*q@@u=gE;sDwA)6A)5k;V&-kovIu`b8g zP2fW~-K`i8#$-_@aMq?Vhr2*5aC0-(F;FEK4IzsCDP#Fzm%5Gy@Z}AQpD-zv^}p5G z>3MRSJ7W@HnIfZM`WzUbIh(Z8aZ-|5Syka^^j9;8Y8%R6N+2K)Ww0eEML1JYkZ$V}WMTcrKAh?-ngk(1E-R0A zs>ND{ad>Bm5IP%ib`mdNEwr5VPX!x)&}D;(qW7gG<;3r=P@*kUz#Dbh%pW&Jv^^a8 z{7>c@Z7e*#5U+P{$?vpNt2^E@pY!?P7fL>r2~!XB1&1s@idax@x26>}pHB(s0c}j3 zf8+>ZTB%Mc24ZOAAYr(3naOPwsx(mu>t87&e8I57Eyh_48K5)B40mGK3zN-t##)=& zSS%5d)uS1>rn)dK#P*(O38Jkpsu-^CM>iK7Wa|Cjg6RaDW1sj36WfX!$?Y2`4C*Y* z%aXrlx-83i-MgkdgyrjglfnyFk6D48iJ7KWj-p!MOD^N4t+#bar}}|}O~RgOa(D0b z?5@Yt&1H78D$j|2oy(_-aS-IsOj&$2077Kmc4F}Vx#Z@C!i9p=0h!n=+Yl2hH{tsu z&%#xb(^Jw>|DMraw)An9haO(aQjQ`mQ9hZoNWi=ItA^g${_MDfYH0M2k;^FkR7Q2B z#}9+e6RLn+GV|&gxb%R=_i!0j;J>=Zx5vD4WWb$dTg_yJ z3vhNBs;c+lGkhNNsGD)sIp>?8aT!zcw#2_AtW?p+lS%V=@kP8hvuk=t9dw50HFGv7 z-hv=ATlRc@*1N@Hpkzq5S3&#G>VJzBjWJ!2D3R5rmVptZZBB0S6~nIq#xXJBMfFCc zr&xi*Dj~U8eQxSdapoWABt!djL4u-8#K(MLI*bCDyuXX*D7qgV=$suoqU^8%;JYs79=_ahArv)L;RpqCfoQDYk#NTJnr~I+lCMdJoZTf0?U_L#=T|i!3AwCdQGUdZK|!k ze!P!thi>s&uGb;k+k{hk+G}ZklY3-S&F#rh&^d;($o%xtgQ3?impnF;#3j=3{4?kT z<#*-y7c1i$x3nF2-XM#Uh4ap7-VV$k=*UV2aH9fg9KMtjZq0qfOPEc@XaByEJuY|Z z%i&y0i%|%*1dh$HEQi8)x>s>x?qr_6CZyw)rj>od8&Cdw97miUq5l^{$(^9p*qmY1 z%SBQHeJiS=$*ZkNeL zSL-q)`;Svwqiyu9^v`)TOk(OI@s2dIO9!Qvn{@ zFL@)1z{=S*S6pzYa|9Fq^AhSyH|k5p_)`Vy`vt1R@*!}$-E{78R6ka~KY_GN zvOd{(FP07LEXlOQK7_?KgX4=zcvs63rSwkHWUOWnPY05Qu$^rkUv=q}w}Fta-7p7Q zo(tcj5??DiYWSaH<+0hqU1 zT?Rwn7CRa0tv>ON`TUBxekT7zU9O*Uw-2qU^=sKfS8bigZN7g?y#E5ZqfJdd%T2_}=ILnd zqGXhk9VV~u*Yh1`E&^nkmdTwn7-TCmUsW@5AJcO!t)Bc(7)x6I?}qL~UwB(Tf6Q5D z0(`{EF#=rEgTDuU*M5i5zkq2>7O1StwE3_~k|gUT$LpTZIJtg-yl#+c>1BI;%6PeU z{<^mQ5WQ|OXvce|c*J6Jzlf4dEN+2+yRW)_&A*;__R+ihV~r)~*nSIbYhk0>^X*#j z{hFKz|1v|}UH^4f|CD#CVi6P5EFvEKI(cK&2F+(OGb%xY4RxauJ~I+N1M1!bfY&W5 z+XKv;qASW{$FE04M;7~}U=}SWwv*|A*>-`a0>LNi&OK}I=C6_dk5FS*eLqwy>-tne zbMHNJX-W1Mz-@bGr1^a9-%qn^@{8Yq6yjGO=^`mZjX}w+gCtw+eGBQ-FF9_jCk+3yJTDC54IOuQ=IkU#9a9a zWW9qIUL3UMp=t{qL^K=j9AoIE;Dwl^g3%X+)0MNVM4h=`DRViU7-Q1uki zl7ufNVj*&Rno!?|slhLwP9e>d{-wijrhufu3YwrXhpmCsWo33>iw_Ow6x`D+_&P1v zAQJ4ckKFs9$N>6iQ*D|>D!Mk39zt~s*l)MVkCpqr1E2pi{lMY<>iX-d+u9LV71#dD z>H1q`qLU$)S(@iruN+nJKllq^jqu6t$W*kCAD+MTluY=C^#Nt%mQh95bicN@)UE7c zAdoC?99Q?oQA_u%o)ymiYrssruf*O>YK$!qqy28;>M4a%s;YYpr3adbl1IQC{`FVW zxL2(hyuKU~LKfpWuL>2TJ}`e+Pt5`X4_kq$a&L)4v#8#_g$+g|vyIiSjYJ{4mfJcq z6L+*bceQEy$Z(HdfaZ$L^_S80xXD&!cN6`-NnQh97pLfAVcqH$(Q&zsBLen18Q}$x z=VT;S7e&u_Qbn>3DeUGaLef`zIfJPaspKA38H29Jzy5a6r2$b2{O2+xn^F2~MVjk! zzpEE)6F|Q3nOxO#;2wr^#O*j*N^7U1)7EM!V@Cvpg$y!i{NmP3|LIi~yhWX2@dO01 zBbjv7^K#kN+GEQPhI1-I0yFsMIG8M5{b*RmRRn`K3*?D%LVv%nKaj6d=MGL`Z6YvB9pbSE04SnhH;}Q zq24)2$KQ85MAegK$YqjWjd1IFHnAH^f0HqA(8=&@D4`zb{ecTx;%o29od20YMuOK_ z0!7xN=NIxy%s+ZiB&VSmOrpqef@aR|K1?Ue%_Mm{(C+~Ci@4N7*%gxjixFO-E$2_^ z0ufH12?EHYFldSme{#efA*a)OQ}bMfeTVeEOp~kAmpu681Bb)brlLR9F3a;Ue5|SH z$SVR}OTz&*0FU<3>sD4T<1FtdLqbutMTE)~@8)UJP3IgriuL!>M>KPioY0-iRMZG~ z1O2r)8hLuR`+RDFqO&PX<2m{|1HHR_3f_kxuM^Puu1{Ae7a5dg>0g3;bIhxde@Qex zvdVFI%E}!FlUgoHR#AAmPg>Xly({TZ#ofIp`M19nn2OAlQWHK|f9Ca&DRgmq?*# zz-b{gZt+3=6I9CA$k1B@jGDccdNNJ!pU8^}+8gFf_=eWut{z^In)I4AD~*T$u+G(Z zt?MMiIIIAtp*auw%5alVI^V0gUT&9#W+XX1SIb?BhRBoO*9tJ?arbJkp!L* zp2gQm($Yi9Qj+YtTN*jz(l{S2O4OyP^B`bIFJAKgbbV~H|f=rCMmh5kld)5 z=y?v4K9DlkL%QnB82S}oR8t-8{?--_Jx*pKJy=u%DIpiyjrM&?$eoe!Cr4@tm1q0S zLjGE`tom7dw|(5bXfJvu^g!_!hA6eT?K2uJXj~^zzABCal@$W!K9h$~IHO-;rS_Oq zS&svk?lUserk7s^ge0Q)I+G?x#WvWX=!oJAbF6#Ap2E_nb;fJc9N1(Jqpg1AFPnQ9 z+2a>d&`x8pOjM4YH&b@m82b&c06J#m01+^%ATnh00u&&c+(cRxtjv|nbA;A`FjL?A zwq~}X-K+0CK~h>M4iycR%~VV*mEt9#dBlO|gSFx#t-D{BV}y_haMzgXmo()Xh_^XK zTve#t9(0ShMxvSb`8CQ5sFnM0Ke&XO=upA=?iyW_yLmKVjOztZ5_+7~ne<86>z-RLVe*(HU@&QZoJa(2Y^0`$O8$!zPfc6A z^b6Va$5eYDxREU{)65`V!cwd^G*UdXFGk6A^};tdq-8jwk*OwO6T=Wn5bD@_@5!*> zqOF<6WZ#KQqAKTKd$e_2BX-~+>Gm=4knOtHLDuqLd zwYOSP-Q_hkqGbeb0(%p^>%@dKhbnbCm6@`Wi^=pD1_k;?;QO)LMq&82pxVJ`#z?sI zzk~GS&D;=0BsOfMAG~HTUniq>;G)d>b((`h^3CuKtwbqyRVC!oeZxf9n}2ooVI`As zSkNUr4Z0h1AKDSF*t2$26VRTqm@zN z&I(WNJ30>{JBQ9yrm>SD!WQ)f4y6=4|FGxeCwy#TCFvQ8{<8M4tN*rJpVo_&4>jjX z1@h#^VPm!O~dDZLZZ1WM_Ge_J_%k0P$ zakYlK^gXdos0+I1mfUA<`RC6yEBCv82cf=0e#0?>u?^qWaP@Y-rtxX;Y`r~Izw;y& zl7;*bY8#m_iW5kB(MehVJ)WTiN!UtNf+=NlaK1j?ALNJd(Lsb1Cpxd(181XnBe35y z?r@-^H~=5B_B-Ga-465(#$Cg4uH4)Pxl2M2-Z1G<2e zpcI7>5k$P7duuiI3+}I%01LOXre%suICcmR=mG`qj0NX@4p_Fac5(^xDvmxm11*xP z>C+j1-L2b|TL-3MpW<4$yI1`de<9SG)ZBXAS~4 zB7!RJ0&zJ}&@YLCkBOYf@F!1gk4vSb91ZY8x@g=8OP*RG(GiQiI<85zYId0bIr$Li zYvY>Lc^H(0+Mv4K$ktc~&UkqxPYY*iZiUwRnog4oTJTnMc=@yrcDDAO2qZCi`u@lG zex2nDhi}tk?At-?(jb6zLQl_`Q{>7RSx}q9t#Yc^f=S!v55tbIv=4A4$~DXR0=H!f zl}4Ev;7|Plm>^Go?NojZn|ds8`wY5yhw2i8gP~q<;$wu%F zIdeximzj+ITM`!p{G6=a1Tx)kD~wX>qP8VQvcpyw+wIk<3pD3Q?I^R-mfZU(NhTl7 zBL=g#y=VO>6yts=y*rTJJDfES%kRW=#`O2H^Rzw(ZM^n-j^kK**17TZZWnxi1fYWR zTMs%3iSIlG;JNy4?ssQi;ecQw+HR1f{59$fIIXVx0Ba<#eNc)Ct9a%##64>dAW87| zi1?Oz1%{lKh2xagVEF0+4HD5s1k~Od^q)Bt&K)9Hl$9-I-T-xt4O|7q#h))2<;H|g zA{$ns8xnpu5^isvg)qjhzH_m=!UsOtw+*{3*!Nc{ar%$R`j_fF+M72hw6^CVeTUvs z=Ya9oGcZ#z7nXV!nz)}lW4|Pni;m>$R`BPha&wmbba8e|S@V#aUs{x=w|>@wW$N}t zn)8CyN@zIn%=G|3t>L($j@&8T+ephV+x+?Z_LX?;#f|*Lbyrhbhfb?KjMf6nu7B19 zk@YTH*{tPq*VYDhPnLiW8|YJ3!mWuy6jD^Sxl1}>hHX9=;P zEV$OgkO>(|0`-+``jvqBPwKtul7G+Cf7G@0XEDcVy?bft473>SPjy8_Ml>znY2}00*+ z*yDw3VKxw_0ZKJR&y5eHbUyi4Z2Q-<&oTaXXI!UI>t(Ws=U-UHPRQUC48Xt#GBxaK z4uRVql;XAIVF;50DfL6cMK<){)j~VnW*uy<-tXaN%?pU(H!c`qR5b6{QFDRobOCl% znK3QFOH9GZLAvMIp3TpbIZOh`&JWN-Vj?^g8wF_Vn^)f($?~y&uU}hZJ3-v_z2V=C z(n#o+!Wt(=GAjqCHlU$h+EDIYTJAjVQIhY-|JvpbN}ttXeIN3fiv;g-$-6GR`61}fIDu#B&dc}huD{_H^nB1A7o_S6en!c4Rj`Pa6Mqs8aETZAq{%;e^y ztwCCv8Kc>#uB(G?Mi`Uw?#kC|eNH{o^;F)))*o+q^_2>Ow1*PNGqAy$RvUjCj$*3K{v+@#DAaa$jwN;8=?S^=zY9c8^zNS*tEN#=Ey!K@w$N zM)%$jVdVG0dKcw#EHUD7&}AF>v}WNIu;KK-XvBj4Va%m=o@;rPCODLlYaBCAAci2I z@g0@50|{CaB2xMW=ua|M->>KBHPTaiaj;orMDgAI;j}+Q!ub%9;nld)J-%ZlaPU&- z*vbW!l%i^|4A$W7B`!K{!KFx;_-{En$|lm?JpG}=|LmA-dN-^+DBWc zZ#bN?UoyjMcxo-e$W~~Fn7TrXe4Q1-9{RakR$UpkVyK@k6VJ?!$A&+jf*ECchWA#mv2D`Zd zzOgc1SXtzCXa>&;OQ3P!wajqt2Spv@Lcb{Jd9hmEB?ON=YMsWCzP)gmZbK2K8Umpf z+TmX{gs{NC)j*QKBSZk6Rz@23FCHD9rF=^VX$|gz%DU4sFU9tYQ5m`7{~lF4DwKe!C4h5g@*qt2NeKAYHgQzw=dcP)EO z^mv)hsTP40*8Ip8+{x74=TTLjNkBfCkbC2WBXq7Rr6s(oWpLC7h9o?lT{Ud+1$lpx1Xa)BRRjJ(bAyC zevn7gvr_)WVqHVmu@LK})UUOd;(BZ+LU7pz@7at1-znn3UGpD~1uk0uzhfUZ;B@u8 zCH*f`{Qs!o|M&YcCSJ@XGQrvTAD8-{CW(;3zD(a}#dW!9b0>>j zee}6VZ3JK10MQiS5&<^x?yJG`-t}p~)1K3R4}c&1?+H2aZv1BBi%c6!cJS+f)lFnz zbyK+IS^}7v{lCFg4QKOl>VTMvBY>OfmV=D0f1)ab_lG-Chvy)yDE?b_pB-AKNj z?qKF^0h~>pJywpbJ#YyU23`PBJ>G-U^MAkmzn`XLHQvWk+77b z`ui_VQZ4f{2b5>ZNDTTB2a#<8!4HI1rH&CUnSeA2&t92F>g>7U2PTud4l1YqkxZ_B zUrQ+H%3SYt(GiImxsUqn9D6GKr9m{mhkwRYR|_NM2hSJu*=$+AhR567DSmy zSK}qcW%kDX*+>|o*H7cu?v+5M-D5^8^sj3pvybHBWDY&+*6ZF}Hnv@b_<`0?)SHo2VM3b(%HQ@M{gNoL+wf`U#}8s z{byH(|0r)fB$)Q;UJ2J?5d|d>r4i0n(8q5=1Ou9kj^g6y7h3Bc0POJU0@G@DVRjq|?7|`%!&;?`?sJ4$fIv3Zj>+b- z&oW&|=i9SC>VrR+{>eA`ZN~xT2KCEvV?MM5!xHPicb-3Y-f0tA&b5O>s{GoTym%8R zJhp$<-^X&q68E0p>YbaqQD{-v-KyCQl*tjUP6=7i%IsvKA8eeegj|COQoH9yyX6pn zq}cO^m~%KZNcNFitGZ4ezCaa_Rs_ISzjlG^x5D7FjgQ(Egu6u@P+b;K=Xc-F?>0%_ zHl6aIgvDtDQ_G{jN7j?*QhKm^Ecq3|TCo-?;@8mg)w3+rFQai1yaO$;ezEuNsJ&Ra zz8fEX^R#%jqN4g3boYnBv^EQ+WVVy)V8#_II6m(uU3Lq z2b6+Wl;`e)ylg}0^USA5Kjw81$?D+PqO2L8yLT(spDMvje2B>l&?zZ@W^s6c%5i0+ z_B-n;43=^`jG&mNhD|t%-nrl@jcbWeh}%0g#;kMNzKz{cbg*IlgKYc0wEWs`w)Eh_OT@dKK(U9}xIUYmTHUD!!<~P1#6b zrMeU+i|&b<0u}_vYT_LTl#Hk+1fW>GpMk%Mp{+qSRcNtnKCo zws;E!TYxi1_`SI&P?wKg=$AD^C^0Mto0>`bR^ga|NZ72a;}`t}cC3FU!aml7Q1diU zc6VSVb>3`Gdu|^CGS;{J2ghFS+hK0j6C?ao|II|x+wIL@Pt(}qtKUeIcjW8{!E}aD zR8+&}x1QhPOL!eu$<_6e%H}~ZB$ZZhW({hcs84yQPqm=uTEHUW`y7irJXg}U$LFqh zFDRJybKUx@C$Q9_`t7LNRgvBRhjAG=8Nd5Z%YD&Y(;(DeK^`q0xRU9*ks-ddm1eAR zg(XP-(!Pld=&auzFjjlcwgrH%dndtp9wrIS%Ed-eFibu-|4Uu&oH0>jfDP1GH_TQyEDnlCm}OX_*2%3 zoaocVMf#La`oy7s&Ed(byQ%aMYkLtCkLB=?963iPpmqLD-?fYpTUq~(GWRc1v}P8j zAa0a(e~Jk6@_TgDB($hxXQ{l3-Q;O==I`H$UsF9f9bD5M9K>JeC|$iDx&9bm0FKPI zp`jEz#H_4&*e`>=Cxc+8bQf6r3yZm#BPj+u?_U6dQJOP5E?aM;s4Fc=-m;w}(U~Uf z{3{=ZEY^hl}f$z4BFeg09gD}Co87cVV@l_z2!|7!rX|# zf)TLxA`2m~kWew1=Mk?1t!KCiwtxHUtv_IGQsvfBPM(kz5%luVvAF#HYJEVphDbbY zNitCUr|$yjBqvTc3zf})W9FL}^4jnNq#(iAA7>uEyN>hZcE25Wy_IV*w7;%lX^_c| z&ikIuKVgkyvF{pUuo`BCpk*3pV3t^LGmWY0*4eo{Z-L&78v1d={4KTUJGZc_C1D!- z@rGlPwz~ikM_;Pu zwE$Bp!1gtyyCDkb7N4F26vdcOjWlXPG6T6U!?`yt)wh43<`~w*M`I=sBS23-LA9Cm z_bG8|xSzyROG0nX7tkB8~eC}q8pPc zpAYhT9D!*$nn4*p+rJh#ziy8{F%h2Ma;x@Qm0?zsOP9t>w0?&D{6qQi=oIPTE}=#; z>#_143lK-WPo#*Imiu3Bqv&W;F^&mgq)CSJ=h{Hc6^Q z8-{K+<%O73`JE_8`|v;o$lD6 zuG)UvP11^|8dt$5uP;J({>S(k3+_1~I_9;AxR+OwS;9?4k_25~$Ia4S>KqS53?=i% zk{ZzjN&Jff-M{HSc+##GSi|0=nKo5m5jwT=$mEW}7?wL8L%t$Gg$ChX{suhS_HW)* z;+>fw*SRb25p-Rm=(Hd8^KOHy+a&k)+wKemE?P98jq2yt>g^8HdC|_J8vwy`f;gDP z57-{4vm!6XC_2WJB3b47?v#44htAdigYFUBboJfo64Re5nL^f=h)LDx%*p%08OSgF z_znYd4_MJOjxe=1go^lC!o$MfD!9!qU>rq=LH>`Mn)^kS!UO<2X0OAP;d4kxX7Dd( z*a}Ctg{cI_>IAmj(Ujb5uHBtU?mWu(F$PrK7_|`u_nia=XW(j+gU^0lh^1rP6dEBj zg4<*A^KH}`xL*4c&EWTvqCy*06(|HuFFX*6rx{%ihKFirqr1iSy@NwR`dp>5pHgAY z?tRbiO{f-w?sdH)p2!Gk_#Y2>P&0){1}b{GmJEWd@3YsYXHTXZip{4>6-R5O9N7QP zL{GmUggZb=A}gP_fVqftlB<>9Y6E7p`p*t$KyNV4T`1=maL@R62ams01@ko#Rii^Q zIx<0|@PK+M=+q+BI6A0BDZl>vY4iyq#Z?yvkH|WlwQf@sRIh3qvNE{IJpYty+MD}F@~D|oE^dM8x-z&iuq`7okHfwVU-4Ij!F ziD0P0;%eubE-Y5@M7_=OCvWm(h++ z2AiY%<=xTV&g_8slk$&;aliB7qc@-mN};1za6`6a>iRpv*;?(026fzR2mT-G&a$b^ zuwA>jOR(Y+9E!WUOK^9W;tmB`+}+(Bio08JcPQ=@C|aDp;o0+id-o66laHBUCPQ-F z*D1$ZZRQ)g3CWp{^0QR?B|*#9h($6D5o21NhQ^s_Q}1Kf?w0l={H!xE@ zqfkdWg%Yh>uaec2TEuUmeqIwMcrHJ3a$dSM0UtYIh;dGh6Ay-^F=I{C(CG*oKoN~Q zjqT^KMvNHou=ijAp&=4!!D;DM_6-6pm*Q+?>W<`Z6@@Banw=h4^83`a|FQt!6zUO! zL~W3wDN%IPo)S?#h9@OwI2wJy_m~83&al@??T*s8oI2A~+{odX>QD+2Mrso^8x;?o zer_Y*c8gz?lxrN!b~jsH5TAhjpV^T2IaU||pL!?WeyY@+ST%h={^+s(po3{y_`}J$ zxgybEk6e5fv0eBDztP3*G$O7pmb|*+@e`N4O4Mw*KAh^y7sM<^vkBwbnp&(fmXkfb zWAeY<`%zQ#;dSv*c^2J_2_@^#9kmSyO>xCwo2!C~V{cTD;rO3f%w50s206%ubtO#u z(FS@-g=?o{Rdt?Ec2}x2YfxQgzW1u-6Yf45M|TQ(g84CuqHIhSDu~VYFuF}R3fA%5 zy2~dTo(j!5^I5Ykij-F0x+*y*CsR)sq06dcCNaeNBx&ugm4xZfTV2IC?n^@Hz6^^G zpU2ZEHR_{ecc=vyEjcjyg;NRXN{Bb}WKB4s=%h<}57S;IYmK08 zRzqUUEzzGfau|)qsrZa4^ek+-8Iv7ES@QJ@3llQ4D7F=U05#N5gG#>(xzE{Z6KL>3bDArjjBe`$_l`bg`Y0#a zfgTL7%;Np*moxq97hD;b3g?INBTO__0R9yk?)`d3doaOT9s|e$)fnm0g0Z@`F_v}a zRJjMByu)W-jn-+K)FY}ansftgK?AOri4IIzR`ug^1`0aA$?(HjU}l{M zZLPR2;WzP}Uu59UyNd(cFU)qZk(f&x!lzX4Y_=xkwt8gIj$+jMkrA(k$tTU335UCb zQERegdWtKMPCS*Pzo|>-tRtJVO$JZrz7r+oyGD}pfJNBgPTObXzl=?pbFvFkWfZ^X zk-eehr)MW!Fvx~WXE5-r!qb~cD|uaHhX-arRqnSgkZVa7Y4R}@vfl)x@W9>r_9q2H zFE`I2CZ93bll3%&bODjjeayt@3EF(8N$g&8wjNh}*EM#gx{~MFArNM_Q@k;th8Y=P z17qBFm%@});)_74XOf5rPxrKw=)$TbsL)gZ2|wUAMiE&COhj{PELwHNiYrUr6B63FddmzEk#Jn zE|{1&Jec5Ddt{ZqX`zuwRJ)cyXB$XL-cn}AYzY7mGf3G~lpGFS93(o!Mc`+SsF|< zy421Ga{NlY2jo$b2eZq}QX{wx7`5rHk{Pj1PdjBv& zk8z^gsVIHgJkVybs9Hr3ncz#+B}>dvf^*Xq+y*huG8=8;VKi3JKsfll0kWNaU2)c7 z{C?Qd3E-?QC%BKI->9S_{{`DqY?BF`(SrRrWtv$gR+M8eaPIt_h6`11flV|*>+_~L zb#U4{_wQA4TjSh%s8K_QjEI`mLT5}6;nFOlZLrPv6_TUnUU#2(iCFRO?{$s5BYTBS zK5;oeH{@ZJ&NZ2v09TxXa~#BZ)YFG?Me3(CX8=7((cFv(;63g8Zk~;N!)!Q%`iQQW z(Y7e|G=yvq9EwR?ZR6oud+Cunl*_w!Qrs6+=vZLnv$*Kx2@^9kLVRH_76a(LarA}*HI03ZkB_`>od>|= z*Wg?~XF)UkbTl%_;##c58-E<0`fn9ZFDYwH;gqJ!&vS5hV38$Ie3&Y0K)Jgo2LTn4 z%|YhV(CVVR91d?i?eI7zJHAT5=rNlcVq6iG^V|JOf0~o(nK>n$1#dEI;o!=G;}Em za03G_1Tupobyx~r))*jafV$4%K!r$N&h{N$RJ^XU@ozOWc1Ub{sm?l=MHYO~ks}OR z!R$l){%M6`SfyJzoTmIkc*O&p-b9UE6drT(x={H{o=vN<<-s%4!8h;YT0NL*MS7VC z1AB;iYCE8?6cwUGB9*&;(zAf&(Y~qQKK)(<%6}YgaCCGRZ>Jl0Z0_xJVz{I?*@$cy zSR6>a=U#;BSWYslLv`r>-Bk zp8Y%!lI5qs7;zTBSyG~sOvWsq5witoYXLYN+u45{m&S|Cc zRr|n=o!$f9wH#QcbZa%j%^UyraboGuUZ4GY33z(Hmw%VljKML1Te@X_-(!9ziea)# zN#a}l2VUC9+}$z0pkJXZHnFqkqe*I&O>*0>8;Fn6?#xYxB$W`3yXq?xLy72w2 z`S1AidV8kSO=7OJyks-905YaR5cK2~<{W$~jOm8JI!I)od#i8zo+|XhGp$ zT5WEnXzx? zpIel%J^dp4;*Jk*+&OslO|1o6%N!29#ACFLS5=qG@V4E*7aLaTdn;m=Vv1{#Jr%A_ zl=X=B=(|wfrxdUXz}Y{aA6H-NwS9tfGRy5-+)&)qdi@;1%A}>V(T1}GUYURBVj{cm z!nl8Ixv`#Pd`ie*Ik&oH680--dwv+^x6uKNYlEi4hgkH*-&6@ruj>x<0;QmLv!?y~?# zM+)8p3oFGX!LjU+8R^R^m@Jd)OSkTv_WO{^z03hVx3J_l(_@R{3zw_hDs9|==8O*` za>(um_d8KO_rJ>pp$FeHwLIG#a-uG6>0C80Ululm|$ z4Am=nuZSJ(&GL-ukV`Qr7ortrjncurtJretiFxL9phl1%_@pWu8F=jODnXLV#SsSY zr7h;6apB&AwNETZqGuBfsVa_tcu}{?gx^kc`})%6C_fwPib*O+K0U8GzW#LlMWDux zU^f)nrxEZm{-7_lLB451(8eZx?1P^UoU=VX-45K8+nHBcI18-PeY~NpS(|NT$zI+* z9KT-eu=r=L$Be0{-dKGJuvZ*ijeS|BYSr|VVj2{d?A!oPAV{YzqYh&`+fhWMVfA)> zn^TUT6LmY{d=Dttt@+yGr+L%_Wype~863Rl2_ z8?oox=$hX%&0qkL$*}hb+uVLc1RdxDmEnRr~_TeJ0*s&*x}a$p>5>tgp~j$ zE=6s(dpfC_pS?fy_Wnbf>-yltoIq2+m5m}U5PuN|TwUNq-rN%#d#p}MlAzxYbt&I( zl6R!t0Vuw8DH>b=nskXRb=-QnIpcXSn#l!oiQl!q0RskG5VG z)1&WaaeT%azP;2{P4>(n?xTLUv9s&FBZ<8F(;PM$f0)KQNOqxChC1Z-da;zvq@ zvD@oLYBX0>@M_on<--$R7BWuLBNXST1xyzwsc~wN`y^JY=}W?_-k-JzOcD3ZXgI1; z^XUxW{ZT!RjACT0dl`)SQ)8Dxp`+Y&llMpcJI(5ea3P;oDDi?eZHH(Y=}Y z=7b$;!8Vqx8XE`imt{|;PsJ`rT+i~vqI$l$0ky~IF#UFaVk)mifO8W>rZo$_J5m2= zh|#y4FH55~>5|&7hWc5VSEti1V?v-Ue2LX6Fyn`*{slw20j>TlvhFk#uJj6xv*=7T z2>~ZkJw4N04Sleh^|RPN)|nNkO>0v!@y&CxNo0Nmwr4t)v^nx?rcq*P4C#IKPpTBX z>*WWTNCt>v{4Z5ERxvuOnAq*`74it4&?HFB@Zr^h`63eM^y3f8s$q{uE6Nu8El>#^ zmAR+^pCQ>INeO&71jgDpyVTWj#xWOKZ6FR~^9!vTbH-#qczl~qK_g=)2akb`#RRnIyo#p)QLv@ zgjwp5KX`S$>qb1hKSMb%vRy`Djg2xLZwUZy3jf9v9nDNXdQy&ONacp!uE?08&3l$l9M_e=sQSr?zYpWxL} z0N&({=}Ylt$#E=hIsbtKBd}egu^gSs%BqJmOLA9Z!bSQC&zHzo0aZHAM%Man8MJJ> z_;+)_QV@%BWw}K>ri(Z(WKDBxe1(1^(6gc-a5U3kS;V-7aix+N;H)@o4N}tw$CI#m zjbH3G*#0EyT*MkK%UyLqZTt;wNnmXGgJE43)~{@hQ(twpAJ* z&-ZCtTzt)Ay)-W2%G6Yg-U+1R=sn-Sin4$5cutO`t~`pxh5L-o+E|+wQrflO1THX; zH{CLeSRf z?EH%AuK$~!n6pUxXORrz#h+C+DigJ-4~weEDI4G zkpV@rNg%{QSqsjwBk|y;wNq?&2lW-7w6IAeY(woQlQRnfWXLo5JG?o@ky%szL|jmt zB2A%eDYTF zc=hMz1TL^&4czH3CtxZwu!LN7iE#uaLY0L}lYlhxIIf^&L>J zCc}9J;QNgLc37TX-9pn}4 z=h>z2H9*7lkQvy)emA{UJ-LHRGWT%!L=mEZ4#dF~^2g3plQ*aiJ5w0XsE=6NJ3RnoQ&^yM`le$se zN#ip@Jb!@ggKF#pDp?-+8NFA7S>v#TU$$e^>FH^_(1#I`8X<_5K!(d&Dh*+YeADnM zDdH%o*w3foxje{7vj@ zN=Gts889C=)EY8`*$DMG80@Vdf{&6^!)c?;HkQ;8=?Yf|O4Z;mTTaPS>U9fg4^wXMIe#fK7i_G%9spRrFp55k~;ZG9!Afx!|L6?%s^($M~P9LoIgP0zl z_~X?G_G-zx(YPy+&V+rmVyXi%h11{Ex4pL`xd_6^PS(EE>?}ld^GOL%8mUFCEFyyq%`!%J0u%5TlL7L~S5wvIY@ zl$?pPG73NymF~^;O|$?h!2x9@h`$=K5&1vP6hw9P+jM@lyLu)~f^#?D)!g#1=Js-! zeHBp_LkiqCObes(p1t}#$7BLdGjPyNUd0ughD0Rn#`F8oD)i}kMh9W^1)=Y6Fhh?* z=%pKI{nr%7A*n`l$A+AWzZ=iG1^3xz`_>O0DTHP0IcPlcBoX(tElcbr(y4G#*ttOJ zk=#>b7)|5=@#$(xkUI!KrgnvbT#AK`@%Um4humP)YNpr+YE^B)?>(v6S_i%QLGXcD z;9aSu{gc?3o=P=fnLcgEDO**y>W+0xf9uR$c1KqkcDc|(T9w5>_vc^2{pGGbN7mHv zkgLRO%+$>0+1-B34WWFC_jqCw=IVMHLS#r>sUcvBMvryKRBVL_J3nk4Ju~RSoW^O_ zUiO&$tZKk2p3bH3xo^1oRI&W-a6_o=J1S72a4l|2RzmRHiOFZJN%J6^-bRq$WI7E< zB(qn|*Igc^P@gYL&)lL}DjO)}%to&BfXj*#E8bS9A*|SK&5FAAcdnsip_YqH-R`+y z+%E0hDNI;IRv9b8%RKBTp^G^fIq1_q!%2nu&qYE$mF_EV{LKLhso^PA8JJXG@s}8W z*qB;?!@hMBJFd=N)1xAZD47Ia1P~c^b8cTO`6Vn|%$eGdNoo{*s*e@k&#B}&EMdpH6rD)*x@05K zjv9T<$#8~rHp9^&QUj*hAw1E?>PS4Jfx!Y!Pf>b)LHvQH^!u1!1GE*yW1qj$slgo@ zzWNX)AW*P;Z5drr{*I?OjNKPZ+4Fdxc)!C&Z$tzYR(Kfl1D_IKF>xIDB<`znU*e`D zmXCqORr(!6kdl>#8QIj3hm#^QGPNwmf318}zRCf%0^yC*vNpDRDNj`ZdkMUL%QJ$Y z&)GTVDgdu%tMv);v&2MAnE!-SRk$Qne;Vg}AseAiRXFUVwTU)%SWiR1g*H|$OLg;e zPQlW|UvZzP?41%1a}C4h?->|8+B<7Dyx3zdWBT*zJV>g(+A5QyK-?(&P#l0Tf%H&epk*WEVFqPcnt(qYjwTij4WYE9sdHlcEfHa!4|z%ikNg zcB&HuhivRdxsQ?9o)5yC;{1}i(Yx0_!&3uAy#NeWcPZF(5ph1IW-!d0-_HH>ye z0pPU>eqCYu#7ynJy?JX=*&Qoe7yvX6xp`mIQ6U6%bjh-_E=wjIwUoiZfK<4IhV|7! zpu*25F`8=T(ZIZ`D>c^f{0WlWH``Klm6Ak%!U&e$#9 zdg??cZmp${eqM+qz?%?yah^&h3+!5!he)NZui+{hc3i`TGc5z4^2qpYGtLb zplIRPO0!r>i9#g4vzY7e6-jE8FUCm>1lDb?V`oq)a6ggL;rScbQpQwfN7Kl{0l%1V zsAtPkvzw6FT3982vBb9vOw7x5o}y(`2wpt8<1w|$rj-_bZl_>&nZgcNDD-c>e;(}= z(vih`h|8@)S9@6vQ5{E|3VmF3>ZNE2Tj@2YV7_T*dbNrpwzrJAnddFAioP5e6VD z?bg^MWD~+LKr>=$QhX7w>N&TBq&99J5T!JhwdPe=buk@lZ^ng-(g7lx1!fO_b>|DF zWLb+yO@&2aT>q&5{Jqnf;ix`oOI251QPC%UG60P^bZ~VjX6?Gvn9Ao6ZO{|1y*_KAz*-+Ah@CXKOkNUQ%`srj`CKl48nS9O_ zR!Hv*8>e1#%y5cWppXEa9P`vV`N{UX7=D{bX|JUqHh{+y+`7L^ z{~?N!7E1O2yJXrJP0&%SQIQ5`&A@1_iUhq|3ipb0F^f3HKCQQKch*~px%jpWQ%-KQ!U$`+;c>djCKJ7z!1{(9(t zFChd#W-AA>4BHkFHtt!;nFUXxm2BGaXwbHf~irFs13Cb*PJ0|YoLZJ2mc=?-Aulaca*^x9h$54 znNn~i$hSPr-<~wgtZR1HDYW97h+adOXtcM}W@0J^M+f28d zjRU`VfJGUmFnB=b+6?5Ufv_KRBx4ISqv}@!n)0ms);)aJV3YS>k+RlobIA!d}V9DlqeNS>(Gw!RRy19Y)P24~Z1$SH;t#Dt_ntV_k&MaY7HjNrk_< zSLN~eg(v^#9IZ=0QJw@X+9vS`E+mIqxP!U=QFJlC9 ze@=0&l@qTU;q2alWKUO+%w*1WW^#WSJ_;XHQ!yOCjddR{2Rl(eCQngq5W^ix9WU&* z@(RQ9`$i-grk?*$Up2LWqQ{;MlT4uo!fu4`(!^qcYs0oLli1{hHk9Ut<0QEKgBGx{ z$KGCrmdvgcyxlOS;Ii)`?F{f9I2gd-2U# zTY`)C*j2xgpS1Op1vyHaj4(=XgKw!?EKbXVqD+bN#tz4nI#7Ntak2_Z^Cu;xJBO5v z!{(5|4u@7nhyoPx9w4cCKnq z-e2WvWKo&9K_xE=!qlZWU;aU2ToQ=@__?yb-zHwgu@l!JC%(_zKcv#ylrDzdSMJ2k z0ci@u!xqF;hA=!>Dg}{*?9b^IP=EpBG^)@Oov+9pv~L}D72<6YZ&J)35~(~gS|c6I z^*lzn-BIB0aw%OrS)+(b}Dr;4j*N@cs#Z^(3&<)7V6Q zU7DzRpwh6h7g5$d#W>|pK&C|qhYFOTKtQ0SUzJp60eYkZCc-f|QfSMKp;ZgXpvYB} zMI>xsivNhlP*)Tu&iwKGvmM>qpym2^O`qu@yYnc_{fNM|=^yWMrrD3?nJqtZKR1#F zNt1WhV4697KGCUG7}HUKGigFPMRMn&#@#Y_b}E=#^(x+WKf{?aoY_7d9PYcNF15Xm32AxFR8dKan|N}F8KjEK)tQ_2x%EK59?wCD=``;arX(zbJs7{I4U2{Epw>H7oG6dggxkV4)QXIlo(HaB517fH@x9evMOH?BijNEjOeg zY8sCbG`Y0qWsoDvopd9x&4;OanFI(YQGTH>gX3`u0?Byb2KD5mLtQ_dgZ1SuqK+n| zD?Jf*HcNfIv-pPZ?bIgz07We&OX#n$&r0;ov~x;K#9$5Bk1r7i>R~E~zYd~@3wIX1 znwK>7?7<)PNyr3S>+J&(vBDr|ldbc!^R6>Wv1ubg$Kz?o`)R_6NrW4attl(jx@Jiz zz}A(6Sm0GVij{^}q*Inn5cyT+mxmz`$gi{5|l9}ME18 zbZnoBX2c)#9>}p#XvzJ+tE+aU&(*W|ppijz;a{7GR(;@t&ER>(X!D>+6pyk8 zrze*{2e2|sgZhZw26RaBFQmvLpdhvnV~N!x4xZGPg^#pvFC8V@#h(>R1<~dMU1rTv zHjk*1sAlSt0U8LNYVx^s1uMC(fqV5b=7knb;YVV^Si4Ulu`P2_gc-{1#VYA|RN72@ z&d%a$BJz~zu*rXtl~a84kLyQBio){DpVk~a&>X*RY<&2ccOmHn=z#v)eh<7j$x)d= zH!UP8koOOC6BAo{7iuSZTYI?VAh4R5S?&BLgPr3x5|g+n_D5TilF-_81_B7X?P(gX z7kBUP08)rP)!s?6B#>h7_zdjQ>$8`T!ayEO=}KCKHGO!0+>hIKoe@;N_+qIuoP zD^g>NOU$8pgJ_ti-$>^oQ*s;SNX4tP zI_xo#K9D!8Pkx_~kH2NGE#5f~NaLTUh!ZfCYSVq-E251sFoGiWyPL};gyMRDOl#K* zcBK7YXM8^wu;2hs>;}*wR4FY9XfNf$o1Soeo|@DBQe!d3!dW^~8`6nMw^$bWFN3+l zwCHo1er|Z^==zV}@?wORCM|@5H=-3q{$ATJTQuJgT)aXQr@ChGGNfTNA)x;K!~cuK z!j*}K;njh0#WcWXjQaSuwfAiy?|l*6+iNn`1~Qn1s^i3-ob7Y>K26C%GFhOKeNmzL zV&mKKcAykpy6k&5C*6Br2hekrOQXXqkxy&6sk(7$y?=oD>Y$Q5%+wI9V67NDyM6 zYPy_$h{WT*y_$(f&Yj~7nk7q4P5;`OEOw<6k3AhSwRZ{+gHrX`+DiwtA~l1V8@Y8){6bLSf082?)~MD(xMdS8xe)| z%W-?$m#Mcf-gaV3DdzTrM6Ur8RAIPQGGcC2epY@yrLZS)39|NrH}L+o<)TIqVi}7 zo2fweSE#_kbG?@_Ot9{%D+8r8p7YCXLxJK9^JW~X;U!4yohL|865{6lh6LkdHhf13 zBngo+hyB*$Q9^4Ci5W3r0A*baIvrSAq~_3*#prR!mdt)B7^?w=X#0gRrod2juNGTA z4lPl&)fNjHTHpM#{DuBmE8G?iasCi$?yIOd88HDD%4i2%URi8j-B#y79v(vfAett`EKk#mf zFi7H!!0z@gN#1dyIihz>EWG=Y6PO}~H@atqTc6~eo<8)Mq@Hpra{k;rd5)5@2*6&Z zm^-|PTMLmi7De}iH5DU3GVbGsc&8=+*A@t<+tNd34332UT!@q}SE4Q!3Q!BoFQ)BJ z5(7rj`YdeDE9V@YAf33Trn?VIUrbncHG>~>y?1s~FvuB;;5VQqksDO^}=-#gau$@ON z{fGov7Z+s(=rBy_H>t%NO}pr8B_*GAzdk7?E)-6*bvc=L34uNT3i0FU@n)ZiV^8Pyaxxr(C!g zi3GDmY14y>JV3+mD%xx$If^P178oMNKA*Rgnas??z^V(6+pWUww0iq~YO2OQRO^u{ z6OBhKnOF=OF-K`cFeybr+w$#5vt z0xvtRvhFbCLz9DW7+d%={&RR0Eq24O)eBP{3;^2>* zbW2r8i78haH*;rw4e9r+>5W6gHgO5pWr~nM)549y`fO(qpF^Y)z00yim!fM%~W zDJeA^mEpygZNrYPoESsRp}enPdvD;WDmIYCBQ(KXhL|M$x?;_8kgg`pw0ltOxp;(2 z^@!rNL^3{^cxamlgC73;qv(MCol+Ry@`X%f7z&^QZ)n@xl0-Y;hC)y`EG%A^lc@zV zy-M&`d;rkUPR-Q|le3W5vHj<$4b{Q((?&s{mHW~2Z2NDnQl|N=3J1HHt|CQq9M(PY zThaZDdd~|bWLfE)7ROGLxh1Eh4O8CKJCamo<&QSX$mEFL$6@Bz-$nl-#^HiCG7D(} zjDNO(ZM)C#EdMx4JhZ-=h|I8JXy^l}o9+r!2;uIMi?`mMX_#0_Gbz!mq{+vXd!8Uz zT%(_#D;T_BpHs<#JhkHfEVYNq{VQJG_Kjs5T*}}Y`Ld4C zifYxW$xs&57yfy10kf#>p{+%5cMj8xF1J0iMmdP%?K~2D;>UVXp(P^5^g@XSHFHYU zt@7v%Gnh#XUBSj#qi=$a&EVtW9=2ENI?HeGXft(Z0GauYwCfH*PJT(gv5eC25LUrx zdG8xW40Ti{=&vLJceHKsP&m1HxacX^jC@qnx@AkagN@;ny_;A_Ix6Aj3#?>bp$gfy z!G1PU{KtMiD!%c^;>zIPKwhL@^A4O%K#)6{K9?9ok=|}R{Kx1wm*BX7f&5U&h<2MD z_PeGv*T%pq$NdC1w&2{%Vh6YPQ8G=n%)mk3e;f07;_Hz0v-JdZZ_+n!jF3JK;Ubok zl0WZPMWk;si951^8DwQ#eFMQX##2$>vJPx#Wi$c8`_b&LQFQ$Zo*R0PL}0Oa4}3Ui z9Lwh8+oopyoyB+tT$aO=c=gfwGr7xkgL9`fwQSQG3|6n5od?JCL>vyqL4_uDpB@c8 zw;z7a+JWXNCW-b}gV2RewmtncB+Bx9li=P5)ozYSSK2wmg*X`5%xyaNB0s-we>Ea; zNmMby+0C=Ljv|))?dEcgq396k0%Yhp`57(+^9cuvX{9~}AsRuz9cHn)^H0xLO&>}7 zv$D$L%Iq6Kd@e)}WCnUPQ!SP9Y3pSLP;=SklG-8wJ!Pi^ez-Mv_2hE8{{;;go$2Bn0sk=iUUtoLmwf=cq9T7U(V73W2 zuE=aCa;o;eHW>O?;PsPUlt_0cQ1Wu0m!(;g8^EY9Q~Eg0OIog?!#Wr>$}AA_&VVdsMq=y{{HV z?L$ad^JK&h{EEVNV0S!t3E0bQc9iKMRnOr%WyQlt5q)B264!C+1??4DrEgEc@EVNr&B(zD8Mm`#h%eUBdqPjZyj}%G&#LPOk8wv7 z8}YS7&ikG8DRvQ9vSy|E#TS}h(hSoeMxA@GB z++Y_U&i+N~w*R_K+2O=L{%gcJwRSMIa`2s#Pt$UAD=-ZXcOb(mH)Em;i%7y_GeD4{ z9OlSXZLDk0Kznkkb9vy#_s&DY7)QeRv?~37zBrOFT|qa0nVVy8engTiiClcWMH{&k zGbp!g(59DG1yiyYjCi$XE%7hxgYp8#$_ypb1QX-c#V=LxE}6)X1NZwDgb2*qmiR}8;i*txVB&=oMl6P@m}pLEf*?*t-aR$R>pKIm_I z#sCr3oH_aNZ#rhan(fr37k0&Gvu&B)*24;PgfgYQTt&H^6`>3J+C5$4MG+L2)H4AxNO07Xx!pq8G+dD^nnnK2l&I)_ELTp5K8Z=@+!ruMtX8Ti1*cKax?Vh}V>c$68m8p&S`xcDQz68S?Cixbl!`@?hwlqFmyy7!`6GDcJo6w9K7C^0knG6laPW9s=^sF_4HB;4r$P2=*>ol)4w!qM2q325Mi48DQ z2h{*6BB8bVWhelq&VJg>^LatKxU1N_VX_teNFa|tD}R=??FUs|xJ%UKf-sGN{*k)w zyczuqL6LfU+YUaS^K^wTD=VedCp8H(1`4Ps%m9twKbvUFoKS2nL_`FOy_yx}5(v6><4aqGuk)S3k3)LAFk5hOmNHT}4T58>}#OU#GQv3q7n6f2~(-Ucn_zEZ$EQ~;s zl@@AzwIWbx%>@0=BXjDjtOiJ|t`skCQ05p*Yv^xDWS=v$s3mA#yUIM$&VHhE*Y5`CntSQ#D`6d*!|lCY)~M#Y${Fay0* zTP3^fI3?zq>ynjZ5E2xvq22@|_-}n0Fz2K$L>0cPj04deJ^oQ2&FH`3-SY!^%qvpm zgc3as>NBX~9J^}=YYfF$ip?1aT@oHqAud-4Z<&^6BaSK)V&%L}Bxu*wAvkHTozWzD z4+9E;*ovgEJGA7O6g2D>=~D!7Wu;_Vrg?e!btOa0ssdx#w12Qpi(R1wH(Wd*5d{4` z*eIyAf0(-Bq~kt_D}6>Fm=7)rS)hA3KzDS%D(aK|y|67EhoEhJWy|MS2LE~x=fv^D z<$G(F%Q^TBt2E5@lR0O(I1x-UJtV?hyk;oV9ro+T?IR;1?4$-BDEuNk3~hHn`opHv zQh;T&n}o3;Lb>E3KE|91sW!F-N`ZPwVET=&sB%^5fWf!WGkZgl>Aq^?$Rs3!5_;6B zd73%8wxTsUE^=ErRVZvaqJd8$m39LF@aQD_9TJ@;tr0bcj%_f>+CT>S)XcA+ zgS>HfpF0Oros{>y^BA}Uq%>~3rYO(i*^+fxhK-CPxJCh@D`Nc$nnvvOw8{_ol4Yn$ zmEY8gNY@|}{wfz11#_W_|BbPqR^H3%K^12>Zm?gXU#M-xAQ*wv3 z4LK}eHCXmm1W0DD!Vvc7kH=8!HO{cAwkH`+q*PJR){M3LzA#9caA`rwhTZDBV9VgO zC0t?Ic?9)b=|^9ysqW64z>r}HsU1dyD*eMWAExUknph5&IkwTC+9#>JLfPQVd{7{%Xf$DglA*-WpZ5RMs#CyxaM=oj6d)!|*AxDf%A z&J_kn-T-CqlR$Jl-HOnmHq|F~K6Y@u!a@T@Of{O#Mll%P_rJXcKFV39>L-G3Vhp_h z?<*ZP!1;k6*|uo%|IXjYfev`Kmlc-R?hQ$w{_rnJdOhMIAG6G|zkVdJzSMSeiTfMP z)&4tuJQVMBP0!<-WIilBmd1TkIQ;P22oHwOyR}I_+|b}cWmo&>WxURoi`;Ca_+Q@e z&Ul0TY1oCC84{i~ZZN8_S6^im;S(rcS_-wcejpDA*sfB|h4r zAX|=N(p9rXXYkzO-ZdX1ZzbX8xnFP8u5CBeKZv8>ds6vYd`W?I1pK@4aq1hhk`S-9&No|86z|&I zHjQGmhR9!-mZD$R&#J0vAvj!gfhPWzmuJrhNW{aB{5Q93DODlEewIv__~bqFl|;ig z)?GR`-*{mKL*B`0118G5zuF!=a>GMWw9c(_2&7$enPB|=g9NU@2J02NjKdDhJ^nR_ z)n;qBt{UwC9tR%gPP>9@4@+wc0~)t=ug^^%=|v9%bXv!(FRIq^??DdW!gey{zSCOo z8-l4*vOKt^Bq`hPT?bx>6Q+lHmPLmKJsUb=hfMx?t-K)O4Z z?p#Xg5=4+iy1N9HF29r#g1(2}%=y00f7zFK9%~iO7Za3XJ`ep+hO8eaq@)4{^Er(rJi`Q{q##D8)hJ-R;FV$$X-d4 z{Jjc=BqlJ9m|+15k-s1IRk}{(3=mu#)U)7zeZKH`pe@}{mOpFy zDiv_IYJ*oYQd@BUa^e312CfG{ino2jIBYIKob=uNv<-;Jsem9#Ta9@~;fS{c8C+!9 zi;(~E5CL2CR1S%=P&_GeYQsLie@&AG+`9}JH0gkYyHB8lz8)G@-4L;Ye>#%yyXj)xK#RL(JAD%fmzKq zuve=E*Gk8;FcWy+#QO%`P|{f?%;or5Tba84;lD2odBMl~|1r#Kk%YZrP)ad{NQKJB zpt3;)0nO4=LZDaOc-$m<#q52Q>;lwq^d)O2P^M}DK!aD&T=Ed@fV-{tXWgd}%T5LN zE*3Whz_f8tzQy_&2}-t39rUWC;krDIhMyq-u2e{F2mEc=i1?3|hr_}%0f#7u4<}&O zTopQRGvAZrqzQt&5;f1&=S$#HO2zWp6t*%oArh#_ju_zCo$72w({$PKdVBnPzLP$0 z9+C(gl#vko&P4rpA&MsEkxbB+YvJ)ll5Vb8HJ(BF^NC1{s-VWHb8H(3k5CFZQbpPH zJXi~SeA$shu&G#?j+vPtquv^5ERvw%wew3ZC_y=bUF?%Xv1Oz7x7x9J#M}k``lfsn zdL>mcrcAwO3^1|taOSI%Wb`0J&c9Be!MHaXmq+9C38B=6f`;>6wyDH;3CEr_C;Ng* z$|V4l5iN5^*q=vFWn)r}KN_?%T>>&7*zqqF8lc`z{8XfyQRXs1K+c2l0yVH>^#i{q z@$ZJI>^c@#vP|V^7nehLMjl+i0jpvB)RE5B5Y;4BJX;1kk5~U$A^J%CqCHD8GoB{& z-m0S#6B`ee7zGFQDewK3pHcM%aLJW@YGmS&@;dT-)}3;Nea{59zxCOU5yl+8#P%=i zyds((P(ghEopaK*8OQMfp+|TKa)X>ZB^GFQkHf%xT-5VJZ#?q3S%4ICZtO zaHZ}owRANi^t2g|h`a}O`@YjcQ|*cH>~V(YrrB{nm~?Hp_dWyWXocuz8;8E}<1fG0 z&j5|t+Y(3gxbBNCD$C=SNlIg38QGF-L@cbiSo~-A-o5Lnll`nt5*Hn8qx0k*j~4qe z+`vjcdJ z!^3TEJ;(gv7an2ae}OgPL?Qlr%|}w5AVGYE{Bk?Y=&_r4AI^>cq}=kak_m}B5`ez? zw&v++LlY?*&a43AP0A@HVSKlQiRbS9EP|KJuGmgEl1792_e1hDG`C{^i^dQ97qe2+TGCAB9j+Q|Q(pV5W+lQ%WD%GK3CWy26 z_P6$;F-tuhR(4U1rBH};+0{u6p`bRCAwq- z|6d3($hT%&*c1XGNUbZH{eqz9qKnP}za?8rW2rH=a??Y&$wTj~_I$+s=al5_bJ}4B z<)kUkHy3(Wv|DmP)N4z}$rMZ_Gg3t<2jijYNK&SojoMeDbW;$cS+AF9t^062#$xcf z*uU)`TGL~ppEishXH#WS2x^T$AY93r^G)q5;g#^1xWiLk(#N%kr=YWs4)z3MMQbzr zLB!Ut$BE>zTUd>E_t`7cO%J-UtrX8!Me+aMlV;TPoa)t<%G55En~W^r<4|DG#t;^P z6KwJJ#v=I%v=Ti!yWK|3IQ!=qSj9^rnU6|WJGmq7hS3+j@tCl+dX(=oHLf7Oc$H}fS;iHmj3RQZ{)?YT$X$0lX;@&zB zZ3`MbTb3*&CcU|W=lAQe4(EG2$wVEiR2Q*cfkWwxVg87Nu0%_ zlan45en65+mS;f3dJVkD^m-n--}WIzn3}dy^!zdt8iFFvb^SI&CB|W&!+n6tv>vv& zJJn%wPS>Wa;t-&i^(@}z7%|-w_vOhGbDFj>JAQ(=SqWh_>GEDt$-mN{C6?4>vto-2 z{Gc;3<2Xk?RlnIumPojU98(Ie1%bF^XXo&l$XnZ4xU*a!0ZM-~eCHcH=d1*?6vMQg zlVH)uc{k&WsHoP?dReK6M4rQ6h&f199t=B7*?4HRV=t@+G^9BMm80}JCsPd#)x?bW zp`sGARPb?ZTM zLl5VHky`N;_C_2rAdKkl<6|_VTN|GO8Q*pdGLSXgWep%8uMvIThz zk&4c-{7|w2Ux-EoDOwey55UGi6~x+#i%3Ub7GSqbQAi*6#yCQ1sE_7hCYgI&C!71jXlN=nN=~iiOweXDT?hoGf8Vq%HY5N@v6ZV2N*o zjcMdjp%j%$F`TsCKHO$gxRNz^KA$3CGzP2qCl2j_vgHfs#?#S6K13YuWx4Fn_BJc{ z+pH1E<@**{4h8+hZjVYNlDfI3j|<1U*rn808(%6Jtd|K|rH-f>Vz}hJ5`>B6ovczJ ze<3NE=3DNs1A{O1D@x{%M}_H|^YF>ue@$`{9x=|LPw%MDA>z;^g>qizFk#nG7Wn

F3ZR;8@jVg6`P!el0R8#BIo4FDAxb*bACnoANxhA^HZrYj zX+=2gUtqq8;{vmWZ!zah!Xjs0LF1-uU9S5$Sv4A+8ihw*Ov)`+dyuzLQL%|>o zCO5eq88Hb?i%+qdOgJfh6LMqcH#mH|EIBgKqt?zmL6=yza#jnV_A7N8Ouu*sh%j{6 z$sd~uU*~v7t3yEy{N%+*tm2M72@Xs3mb8?$qi1C9_+cD_WNX$zdaH^Ah>`7SL}vW-4pzQoJ1s&-V-J zP9!|Ei=vLZ-x-ZV&*|z$A)BhXV&MALYTp4ay{zT=cJV!8-nmwO&qR0bY7N7US+LJFgH)LbB5uX_Ozj)!lT##5jC-p!tw_!k_tK~ zX;t${87EIJ=$Gw5RBTL^Kq~2U{jUxB2W_g)ip4YcG2gXXL=y!YJjjX+(K~pR9Z(AO z3_4WSs)Zp54A10t;649lOex5oR;l8%H$xDvwOlaFhwhS#)2Vj!}@e~wRdrGbvCU`yk8sha1D=F#$Yj1Oati` zY7(4UzRd|tH2&-ChsnC!nCm6 zM2RrfkwACedrKI;C{B8bL;6muR5({vW7;9U5~;iqX0CpxoEiP8_?t)`&&u+}K>_U-4UF*Y7r} zvd#Ey(pI(W?+V~y6mb44d^68ng=iHIj&}k8SpA3YUOdk919|3_>)Edn`=!j<M-)0z~$ zcf@&v2m0G{b!PZp99S_u`Qa`D4gKYxg-t|4>^^60v!40M*J>Qw(!3nxXc1P1R)Ihn z2l14wt$Pu0B^ zahP@D8r9TC2+j}J-={pb)^@*MyB9vt;UWb8P}&T;ec-3dT3y@E&j0m=-&>X1#IQ`^ z2{6asDJo-1{TN51s8V>9i_LnhOdHZpuhcvkX9-e8RRdAd)VaqG@ArY3zZq#gEWZ=O zOEs6!Z$^b?%lISWGsu<;3G|p zmcW#iFV=GZTAM>(o5H{06QwHVrx)BTP9n0q-kN(2le`GH7_}Vs>KlLKBt5+I(q*P^ z1py&r&TWpsDy-0QOOYw>Q0;|%mA=WJ@f5-pw^srdD{^n^#zF-!2(~;){q_qJdaGA{H^CQ zM3E)T_sD{yv1iOtWG655(T|FJz}&jEa&_mXP+NlNm_rIWupuqe@xzbSMm-^@rpVzr z+A1usK}op4ef6+p#_VwZmR#tgdJd=8cG9(}&|$mK%RN{PIY(-oVk3UQT5f>*7Dl8s zNs+Ex|MTy!=4ESp@=Cf4^**%;^ZCud;0qQXD{9UfxD}E~fbJyVCW>PPsgmJaDGB$Wo@nSIi88G{}dc2A& zYay=r)Cr7mSvg&^ou>P0Yb}1&8|5!g+DB_t*USVelf+g&PFnbOE=gtsEM3bnN13aU zzBJ#W6u*Nv2`Du9c$}kY-_4+FIe$#w{!qh~5et3VqwN}YyB7R| z@2{BL4nM81J^5NFRjiHgnm(za-VdEWgDx&89-R zF-P!%caQ?Rq>)4ic_Sg{o7DyJ1)T-YwiJ-QD*>L2j~E5sx+Sum`%x^io&D(^7MzC_ zTjn2W;=8{yu+DGt8ET*>qwcD6_rQK}WH@s(n+v(H3VZk~0EQ{8o?QKQ z`A6dMp{j}xD?o{&u-C*kWT}nTKI`**<9Vm}@2-B*zW||(TG=rGD5Fa!W)%1oqCfO7 zAZ$9B_-A7jIGz3dh?2N0wY*wX53VcEYcrvp?wt8$57X+fkK)`u5BMBC8DQhd zI@--h=pS8VgyCI@2^WSa!ss!~%{0WtwuFaA^3bvL-mVcX<^nLjQ}Yt0@!_VW4ySED z$4_7ok$UfnKXrJ@Dn#GhO7}wcc9N4n2dbePhf1tYyL^K1n^ZHCB0-nc8wwoSOB_aQ zNi)~fx?W6pH<}+?szohKWuHo`-~BDO^wZ}ivj)E#F|uASsb;qST>l+P{NVd70o?Ka z@dSF4C14OgRJGM=*|46yw-z(6P4eAWIBevg{4YSZ8-1O@pk3Y|mFGK!W^#W%u0a3r z-_eLpIK~)y=8~>e)qi<(lxZSgPBeu--FP?Uxqw_>J#$(w8c$5g6* zU}Sl*9IfS`nG)JjgID6hI%3eSk(Z0b-97LZ%^?fACc>BwqmSU}yZqKi2f>3Dd8_?Y zCr`<&LQo5=@W9()DQI2FB>_N`k4+TcrJK)q)sRejY*)56+cofTrzOH#o_9Jg)Xb?n zC)(rKEX0SF7daa<3x+m-XW1}=d@_>swL_-<4i-n0Q3iNUYx-G;&g7EUj|KWaYviH@ ze=5MhD-1LiQ4Qq2+Qe>C&!r5t=8C*;$+yNa1LoH_QE4ozFnc+>%8_0}ZuzQ$PUc6l zO}8gj3h6dmJ9^R-)PDf~ewQpT9p9DgElQii-K63ttaR~Nl5xtn z+q2cF^!4ITc76Z~nVH!V2m2RbWo(3&P4BrHzOu8DJj}DfXW^||P+=lf{%VxFja!tP z)nafK^W_mu`hM_pmo}EF+qVa@gTl96cq3E6Sj9~Z4FBQ>#uEnqV{eX__&`NLU@U{q zx(E#UaFetR9ZPz*%)H))PO&cH|NNW7u~FZbIH{K4U_IrihI4B&nFcrZtF16^gVg7> z4ZfKu;wOkQg?*mZb|&tpvmOw$wC=%1LBlHl47*l^nhv)7_ZwNueH)m0=d2QQUn&$0 zl5e=O>L?m+6KgL)(L{x@ov(O#LjWY`RaRHHn>s?eFzBI1_{)F&fAm}0c2&b_qS#Lc zvmye#>>*CHV*Uzr6~YEjY@-3$38aW40W_#Z$INmH!baG3&fXf~j{zVv$_Q=7uc=2; zx*YY0xkbruD`Q3NP1i0K7My&=2-gmxF#NbgndcWBLb&twG)2ASMMH{^ZplL62`mFN zIPC}g=a6ePA|HO*uj8v`kKk$itY{rpo2dYQw)(&(-2-W+{9NBE1X*QRG>_B8<(9_k z_@>omG`K)#-l0HPP2nh{-zla3aPOT@BQysU{f`NMJvzTyAbG*@j;&y`t**BHXjn;}`NR3o zDP$%-Nvls_1hP1t=$&gXI>vdplYhmmPI%d+G;Sd3)_#5sUdpBtug~gPh=3xO^|6?+ z0YL=z=Qx7Y$4EE+P{@422FJ~50+R>qKKN8MlUQ}Mm&b;EB4c!k#({wla#{PsnB}J^UxF8z0MneI<*6P=WweC)o!5t~M4-mPxtCAAo-O!FGj4xh`WIwoDt zsDiTS>}=j_vUfnFLxG;Vo?M}uv?LXoN+{+7r!l@0I7W%_3T)$Uw?p9ZS5aU(;6KTo z?2rt+-ChO~&K?M=Le2`lbe32qof&iPWVCl(rada3%gH{Coi|fCe&z>a!VwuO)_WaZ zB!n%liL-Zx^zP`;y@98BnU8}0)%!}(rRusX8Poiy`q$cdxF8N>N#_jp9{K=947QR= zhWX{SiX72WHs5trsGeu!v409&YHu{Cej8CJGa}%X`d2f*~lfq|>bL>}^k6((?|nMH*#T>#IFa z%|fZ>Z0MylY<|{V8Y3>zSAd)K4ZIc%C%%N&*$8%^4+|9R-2?%e zOalYK+|vQdS{9hzEs;^bh)5BUYY0G;`+T)hM^|Djmc2MqXr zRLwTpw3^BcGDd{Q;V9pVCA1;ZuTE|wcC85Ba^tl=)<>1#ied(uQil$k>+!K4z=q9y zBE-7P<>MbbdqcTFx@4EiKGw=~$2oQoLisg=0_zY+KWVW9;{>a4kT0u7Dcw{|vRpQu zm)Mfl`enLk1sHiwfLOTbB9@46^wNa5ad!5Qs;2wt-(qf+<1(um>R*=Sil|PN> z3XT#mR0Mih&-%!=G7Byay!}K^MIdg%*RKL03O9JxCn_;k>iMOhI>!Kq$`6%)9HJSC&zw?$>m)YO?#i;su9;<^*c}pV}Q1VHJ2`DoRf(9 zAktecs)iJUFZMF2Yv%Y>F>$5XJWq9m(pLsFg%Qb6XUL?z^}yq;ylk=nh4cdI^>{j~ zE49E}{u|u|Gi(-u1fy6-EMlaffcO-Z$rBb;km4nyU?p=!6|2tqVTpFL4Jk1|cY9GB zKmxuRnejsfu{TmMM$A##y_UQzQdpXqR50Rxu|X6=)OZ=GN=CR@j-uPOHWHyez1Ik% z-151XuiPVJVjBt@Fo@KIfVs~pBABfzG5!!@|ASCNN$eWU zM|5t1N0NxTHKE~HU{YOBbbeXzEog)K;m!Be>##X6(T73-;s-3Zmu8RaY&9VuZE8?~BkL8?J;{?xX?cT$!Kt9G3R9qSNy8q>IS0(le|B2LZ9 zR^w%our#xvySXKL=-5& zU>VE7m;|>}*M8H`e&Cn4BH%~`kTgHKDZ7TPU8^FDLWW-dK4isHw^s&5y9PC)lY`6~xJj^~LqSEW4frc2cb|m%V3A!YA7Y1!mV6Vq@Vp-*&yOejoGA&m3Z zX3-^Z>8mCd4(u$&tJH}C(aAo)0q$&iPUN~!Ig;ERf4sXi;0KG=!L`9M^!dErd5^8j z|EB`q^_v2D{7MId##s7{W_XBbSdNY3;kT;G6dhGLf$`R0lOR@k>p6|MN@RgTDU>WT zgbLh$du`>OM`8VmK$$K|MVZ6w0!(K9(YJ=SWHViC;3!hWUmo_Ic5H~L z23QX6AjgpHxlWoCNjGX)B8KtseW~pA6n5LpGbJgkd1-5M7?B1rD$ z(#bM{(BSda>tD@pD<`eKlb6Ll#8HzeD6KgV=)*~7^|K}DKFMtym*u2flcI{7d^GNxGG`q-I%-4TU%N&5_as`Gvrc{9iKTQn zeX_BQmehJBiN0#3Nhu95A!j>%U?E8w-4D=pc|1oqQkMd?@jV9Vh?I!AEL+(yKN4rP zaE}6|&zhM%D6xD?=PrsvRb3m&qmT=-A7%2!8a!4ZKu31~IQ`h^EP3_m6C*q zQ(E5?Zlkl+`I1 zYz{Ph+UjDVq0V?J}@npP4^BHS~;JaCrF@#V$%Qx})Wge5(axeTU zo};^350_`SX~f+XB|n=|i>$G4DpKh0lmGEis1ZLJg60%NQCsWj(wXQ`Qf%BGbI@>z zO+NUx>KmuEfz2>lznfhga)z&`&&&I8DV|s^=-}I4V9b59!O$pG`FKnfALn*3kwFhy zN3Olz5g3V$Lct~Yam`gN;T3t~mmn%RS5rFtVL*^(o7oyES((Abbir!2F6!qLD~O9}r(SWeiQJUloG!jqKCJ?FpbxdQm#` zq!14Np40Ye2j3*dZZvA+wg-&jB#dxWH`6f(=%#l3YI>!{ji5+&^eajO(TO$!+X|Os zP%%0`ZjP77NN*7lL6kTXM-fd_M^Vg2d%n5F5xs3;)o=YbFL8eYxIgw(~s&nmUt_7+7x=w_rQ?J>OSC1P0*dj<(cib4mZ$X8Az~! z4{P(HE=fIe&`IXH9qNlpOd|Vz#a8@Yhz#9wYEf{*qg6?_G3Y>uqM*r%w@(V#mlYb|NZpH{#GPpM$6LK*l`n?*;o=j7{j%Go`KB5ZkLFqTCaWKcny`Yxv4fkn_s~d5lOlU1g){(Ob>d=dBu;<;IQ|IjB!; z$Gt1bJhMIS`Bdo2TfV81^mt_86zeER!BPsW$A#o62&^9S^_pU9O#yc7dPntZ6G3^p z4vlw7rn&hIB3j%2nBBA)bgguD*&+Ae05Z67Z*fH8UGEHW7&GKo;@tb+h}!1GIQ8_u zOswfwZ)G7aPbat$J!ID^9ChF1-QNqRk{^5w6oh_L7>luq>~_|lo=u4Bv`u389IuT% z$LZlGc!#;5Mu{3U>Mx9~BsA7kjmIse4UTrAF7S(%H!B)@Rx4>sO&JQf6$)`SQT}g2 zUUnDC%bXd}s3+BQ*sqY?+k%z;#)um2mdyI>YJ_%$Ac?nU&@3|b#F7^|3k_X0 zDB0!T^qw6Af#4b{Pf&i2!Q9FW*TX^#XjE0AK8i{BrJ*AS1q+(JW$jutvAHkvP>HjL z6tVMbH>QFej1i9xtnzlt^g2af;p@D%lVd5>d1L=^SV_W5vw12|4T&bAt(L)g9koK( zR=3q*0Apid!b~2v*PI_&5*tbDu;Z7Qa>FMA{h)~o*8P^AD{;(D$D~tAOMQ$ zr{52%ltd7fx6TV9SRi-h<4E4;4vMoN>Z4jCvF`YTr6(5xxGm*Ezvbi23IBp>NsL zWem+VQ}!D^f+_9Ib;Zqy!Ys3@Ip?fB=;K24qPKvm1%@UnsaY#{lvTGNv&zuD`PRj% zM%EXsm^iY3_&wfuof%Zds|GFgWUyNthDRVCl-Pe4tT4VDrJ{kqrhY!Y6)5OBd(;>K zf>^Xqz1@tupp)`tuUOL3N@UR;PO)9PgHf-*}L87cgg~=uo3%hHvCHm z)8gqagHvvS>)sinVFfb-Phs&lF}768PClu?{lx(J1u99@xMe4{cfakJwdXS8Lq=?H zA~l8P*eYs#^vq$MUf;sTqBc)sa2lSmO&Pw zu{{;2sp%1VG@Tt+vZN;q+I?w3|Jtis=(==a|d5MX>Yh`hCh}!(t6V-(;61&uh;x2@< zrZYn6YO>^xm6AR#__AaOyDQVldncO~sd_g| zb}B8f@d)Ooy8Ia@w8#6+3ma$mjPL)Yh$tl_u0W*ZNj@2#q&*LaeZ;nIJwrI zE?Z}7;zFHgs$#WQAz;#gYzLl>3BAn!d%?=b;6lKCup_XKd$T?QEC&}ABtpoS!aL2LVrQFbO7 z^@7Ub9pL^nd7Vb-yoDm1UqO?}@{R7lnqwg`rHSKzU6wI|{n3M_$H!L7~&4X5~b1zlg zy?st8!YG*)!QVnke#4KwR1;qo+Bl>P;LdwWhl??NK)!gbjTU1>^~1iF zo=kA%b(|AwuW8l}mc=C^c9n)?kggmH9Xc9|tTylI=ZRpx(^NWI(OCf2QfHWy+Y>0h zB}+U!NvrbyN`LmFw%|x?ir8N`qFcA*SEB=$Nh_nF{y#V{JrJ0{(8WCm#65EOjU5A= z`^z|&{%tm!VgH_e-g!F>MT^ds{A6Wh;plM-8n@5MqAAXbWVOk{(*;X05-}#2&rj;m zaC)1uOk=(4!(cHJOIGdG#LQDwzny97C?k-95&_Nc0y_^c@vkFmL37O9r zSoTQT=YH*~Qgyy?e@8xPfYi#OeT=y&B#ZlOxU#=M|zIxEK9@ zemWF+q5kEJ798OyexN8`XaJ)kn2M8NQvEsL2|Gc`={wz5XqWqXDAW|VN4Z@~HIq0o zcCP!V06K9Ud`JELC-B3&S+z;Cs)S~d6R5o`o_blMbe&T6T>gmB48%p0iKJc>w?rAZ zcm}S(Ctn4OS~3c6NlO1Twx85sWI$BCY#}8tyT>;>@^xUtpU34jg8>e7f(iNtp8TlA z|6&>f3;8}VOET}5ci-c!5@|_X0eh;sYTFFH4}7*{`%e%gqraZ`2p z-}en2(q(sXtpUD4@1x@PQjf=#n`iCnq;L72 zJG;JUvU>|*Y>r9DM8A>^JU$tiu=vfN@8Migj+kK?V~I9=L{bVyBG44rr|?zW z1lUk}F}=M^0&wpIpm)kO_fAL}%l99ZdgwSBhATHszz%}IQeBP_ zS=A=xkIPuVF!$1fuxt8I@Q;N>Jz*pAY}DKgdb!6J(!=uw*js%+@%2evoT*n$TuU3@ z*H`|gpaqZCUs5@JV1)X8k#by5K~?*6`DXf%4mW9RGH71atYQ_&Cjsc;j>dBUA^yqB zH*DZ$KTGQ~+}7TqIbQH(kyp?{;g5ZCq65S^?~q>!zVjx?t;+fCRY1X^)UW=(X z%yE)=KI2yIp0CW8*>KCs^lFJ%;n^}8aI@`4L!*X`jJXF3evB01$byT{$o&A9Jjn6n zZ$nL4`Lbbzv$3Q^l<5b?Hx()ZV-kF0-Wf-um@zV&;8S#C5>A%igyWGfy%S2XOCRF2 z!v^Qg&->U$Ez(yTB+6Jwy@qe(JFsVV5?0({OU0)Fzwj%a&l+n4+rE8@+yhYlF>4cM zxJ3)b-_fV=M}Ok_z6m&cP_#4-zU#UyArbfsZ$YwQPDixq&B|3pS8 z1Zi9@UMF7ovFiZx7Sw|l`xj_ZOUm)SqhPWxLXFTy@nL!$x1fDZZ6+V#qMNf^|Cb7Kgx2_2(9^OL}Z>G1J^DvWJUkpvC)LC=2M_RcqOvbPT*@9ZW z+U+s*F2y74&xYowkR(<|q}K-n5yDR-8IXvWPn!T-r z1M7K=fC}3?f9qp}da|1H^h^@}g!}AYn47j4fs-{9)QVBr?+x7vW(5ktI1!8L(=Y3o z!v5WbZ98QUvZ-ljJN*lf8%oxN&&^~g)2I#sMIV8z%*AXBR-r9$*Sd2<6fQ_bDI+FU zN9V&cVA2XX1@?GZZli{m?ceclFTa-lce3L6^VTPv)wrH_)@RKwm)Z+P@U!+;tf!o| zBcQD;ME3uJuBK#FRz6U)^KK2iTel`g<+rF)qV@z3P=rE9ZcCuIU+&>abHK2V-WFY* zWoO2-;bd`JIhn2v!|XjO1}EKH9>Bdg^GBhosJ$i2_XuX=Z|?ro4#v>V`y?J{JsP?M zA4NtbVn5SvEElb$qE-`bp^4J{@UnyVq8%#{cn+56h*VS|l)gojdMz(7W-v|UaSdn-OXA)f~>N>xv9FRq`opV5ticiq%?-dtEHaj~qwF%2u!%L%m&u?+k z;!H}odzQN6#4l_=fa2$s~M;%-J)M0jM%oS0_s#IC8f0jF{C|JG5jup z6-}Qd2a(Cfv!P~E&Jst_a;}_sbo{U$G>yub;JV_PA(2)p=AQ|L%5+uTb8JbeV)#oq z%~?K-spyUL=*hj?2DTMO<-8xcP2$)O6AJs8SnQ&Q>qtu)eH7x&Sk(M8{0P0lFg}(sB7Van&;3y&q$w^wP|$#l!e%>(O%v~0+Sd-W%-&<88#VHC#z9T1S#?tW)`K2t`z%_20pZ+H=wFbkIPT_PIY%n z?J&T?toy~wtPgRlW#RLQe6F|Wo&&c9G1RNX^=VflLRES&oV-*#I}L_`Y>BdKPNB-1 z4!+zMTp^fDICdzAz3J4OafRIjT4Kfuz*-U@nEDg&KyWYmX^|l-I5v66G4{yYqyPx&gO&2#v?aBGIHRTB(gOcU3 zavYPh<8Y|5u)TH&;F!AI5Yd(A#-uxbHIiXz$w(L~nn&C;6-7exq;z%msuXo}_Pe_Lton-y$c%aa`GEFma(T0ZAh+K zjQ=&9e0cvRRs}OHXV|{w8>@^un+)=?c)okjLC%Xjir(4l0(>m2=UA>h97SkDr z@TBcCLk+mlW8We-zqAi{QujN6vs9Bp)Rl+{tqPMh`BX-x>aXVX`<&=JYsIe%U&IplPAQ|F zj;dZeX5vVBM#-F#l`smvOFh_L;%7+UYc(?WLMi=&EjyEZ4xG;)&~)f#IsN|J|2e+p zrA`V6eOxGggi5W$&CD>jj56GS4-z|y*v8|t-l~a7WI8TI&WaZ@U@Vmv&usQfWXo62 z|E8iTFyp0R7XoO6my%*-*cT966QW}OxnEjY=JtdI?0uyB=$l$Bgu-->Wh$XkP{950XV>@nz0h)% zPFwwrhJ2o<&id?h#Oj7mpE1pOeTiLv(uBjgc(?G)y&moYr3!OWI0p!hUFz z{v+`2FhWznmd}15QZOzbMI0NsfnVeN-`9W7|Go1ZMCm_9u?|`K*c5VI&Cri9{2yq5 zRE_uct?bKN5CH7aQ)BKPIx31FlSz=ek|b*uHyK`suwWLGT@R!5{J#spApCKhcALyL z;z*|7t(r*v-TFDDcJQ7fd#v#omEScUzN+NytUR!SG0gAbAeSSF{1?A)KK@Mhy8@id zFOPN5688B%Q_b`?h2_5<%H)4>bQNKbtV{{Mu2p&TqWam%f8zoAb zX_YO&LOh|gxaeC@wh%?RU4Htf!Y#ZlU+cI3^<iCxgHTpz~S${zu>^YK#0SFoD3C zOFG#vl6M_G=y}CYAqxINmHp?ZFK^92f7GwgI2`M%9_;Tp6KCalx>(V?33onEBr`D( zLfXS3kVNVg%Fj{%;~Sycwm`@OrEzV1c zM1ZFcz>N*YQoI-5qBoNQTeHE6o|$?|PVVCK^6sNPMzlsK$3NiR^k3fo+?^?#m3Rcf zUn=LXEduDgj4*8_N|4^LdUMRm-ZA-<-K3cKJ900*$}UIfwU5wu_e$;gZU+%>RJq+L4=m)#Fj zvd-7XUH7;LuLu2c&c1NvZC|ce0lLQT6jDU|fW@*SPu6;*TZjHADAJO$KO!&8W3zc^A7|qv#4W)cMqUP}r?zRk8)|*Xm!f5m4_tbMyM*WG0Zs0I zpg-svV4S7uOM=vqtHH|cO>EOE;(qJw70&m~L%3}>$=)byf`-fi=edGvKY(UEO0@>E zF{n^|sf+wWhQ&fT?Dl!r8Ftuy%JS@8m9^22zmyiNf@JgA})_4^?z6K$qtd~EPlf7++jmf6`g z_8i6K9)LrAex*AwCPWB1MF-Y#L4R_&qq}3}PjP0X?u@q(bsb)Ob}!_Xe5de#WW8ln z8(jcyjT5xE7I$}-;_gmx*WzBF5Zv9}-GW0)ad(PK@#4i_ptQZim36L->5C4Pq8yftHEKO4KucYkhaYNnd zF&tRGBQ!fSjeced*H#m%Ac!+IMu@S8c4mT$0w2kq^iBP3aET4@cU523C;*JhiA=Fo zlm_A18Zowk8))34shPK~(ynC%oh9mM#gUc(O^c0L@N4v)#0CqdGA2KICUxC7BrYHz zifxg)-S!`_kn=q7F4(-?2lc(CFMsrN#3IGNr!n#2Db#tdJ^*>hWPjh`GxU^RU_uw$zTRsI<)XvJlMTpoKi16b;Jh7*AcHH<~#)AUt@k~=y)%+p`Y zyvsR(R@e}SMX@_T-&p&p9H#O9dF4C5joBAqm5{&wNj>kYfJC(iGpTp>0elawPIT2S zJoX>))P9T`+p!%a;1_yHjVJO0bVfsL@RP07)iY+VMzr;T>@Z_#4i5xd|Al_+> zww%l__-j6Cw}}pYS=aPIA&VNj|D2x4p*Txi(v{1F3p^~LX8%NxVmV`?EwM%NBxLN*s7TReE?<2HW2w7w6bWs60#19Xi@|+AIj0d6S z!IS)V;wUr!g@&ziGsx{tY3&JcAX}hLHoOz_<)o4i+f2{v7B6`0Br-OB!RBJRw@%!; z-+I|XFWb`mqRbK2$ncj1J$7L}y8r)ze+uAXS;DvL>4tMH5GQj>#HW<teBuCvvZ{#u?Ai7AweqWR-tFs8P*h4a{dQ>pHc$nC4VJ2OK!W={~ z2l*()lJ}B$<=3-(xo&k47)(%TPjcjj$~emEGmLEdlQ|s6l!ixru}SxXoBJYSL-(A# zoua9g-*jTj+0fT=UKv`O4Z3ASr0;#vbfJh&(0q+0m45HzwXtrGHdwFV3XQ%GWH-)Mm;{JRqGMx1D%zgJaiaDh#mSHkrn7egAG{#b;Vg zM0L?xc_SQocuts~f9b+PR^d3I0%cq3=aZ%&B~Vf_7EFuxVA5_EEeU|Lvc&@^Kx^m| zIpxXX?Z?!~iXm-FMx8|@I##?^B7~tD5{HCtUF6RD|C<*Wkh0cgA}d3d z2LpCtnHB4lGcjb!m{N6=Ul?<3kUTFxS8{j27bRBg zY_wVC2_nA4s#Q?^&`Kzw)hf7-vYRaX^H}uvw8+wSuCXo#f9je)VM=>Cz~UIyC{adpvN6g+iYUnzF$Ec zkD|Z*i(Wqe7Potxs9JXbRUV@^4KFfE!Rh3pu#q<+qsbDxAA0hrdBR{e+jx>g8ovzB zNqy<1ui^6kqy!gIF=%BB&wP_k(3%ZVx0`B3=Z0F943lt0!<2aAmI21cPqi9dnPC6GGdQ-vc{TksSmc+zDGp z=`*}U8*n*PMo9cEnXH0cC4(s^X)CoSbf_1xo2$*~wDsTX4~Ej=Hg!-UQK%N{^)*g> z7e;2=c6QKNOF4^)hUGjI&H=zX~@v}$VIC~fm_Ms{|l|!vgarKV{~pbc40n0o zE3{D~yAC|M9ws#3yMRZ(xL2H62)-$)>A2$m^Q@8m{D~4Ge_Rg8iYtC#t?6N{f$Vm# z80NgMwn3pESqh4e-k#%sok3)ON5b2&o0mVJt{~@6Ag`q2PQOQQq`nav2iX68%>O!# zwKHtsY#aJwj35}H9|=QuBz+G8C&4LrEtmXSQ(yj{69}N1SNnDz;dQt4KmSuCmOb(= zl@3YM|8vk8AgcNNyYqAJi(cMHWItPY$lQy|n~lYLS@&%p9M-M2-vduTa5~4gc;CFu ztEDga9HogHi z)lL5@R!o$D>vkCvJ;dK;Yyk@Xv6iw ztxCbK$+16V@8MGkjol=JZv`qB0nm?=ugdKknC<#$?8DzT3kx7uVyyI;pwPrK?Y3eA zY1d8(h3u++Ql>Eyhc8`#Gd<>rgw!-g}~Vxox=sysTh>RQ6UyPyR8tE59)j zabS7+^Y=1xI(hQ2&$R*L!qLV}kB^V2!J4A!q>26s8);p<@WVhqg5EIr9+8 z`yFTjFf-*JRWsqU>N0X*7a-{`V|xuas1HU#a_LAk^|1m}`{uyuyn1-yb(QFl+EgiJ z#;pr&!5b(|s8@LXq_RvPdzP}&lhy6kJJ*9eKdio~MSlGI`=7UL7+t>M!|NH;;%}_} zmwk5M?l=t=uq_P-5`y?ajjHa+r0>cT0t44MQhd3+F-mUMeeBDGWst=6DOEkOWywQD zmDr_&9OExyKH1@oa3@dw;CG@0q1JMeHJtb~5S9E|sjf!WAGaeqO%s!@StpZYy>->a zv_I}lv+}dhPYg&=er2SD8RL~*46k$$ zOn~9Y&m`-#9)7&9`cUvgWcyUjrH()~P89auyk>{4Jvv9X4KJ9pbC?uS5T8nei+^9vTQQ*j zve|%}_=(+z-|G^srZ~-?H-Pxg=3^)_&7g;iroIP~Cu@tE^k1Ot5V#H`kR`u+v6)Fs9E&C?ztoV9t3T=`+jJPA+CP3>JtlV*cCts3*s?J6Vg&w55wRn$d-EnC@P6-=l&-|)Gbe5!Omlgc-`D6=+fYz+B zmG0`0nmKhpKK3310i0zTP_wMTE?03~UOx%4jLk$jBcAgcHT!#HwgWH(=ry?NF6Foh z-1Or`fqRSQB@oUMth(?KBtdkhLkZw;vuXjV{u9s;m?`{pMa(gQ)`6;M2;!UG(WP~x z$MtoCN`_d^S-IfaIgdtYS_Qp!e*P_x#{bRG zX4m)`LbwF>1lB932=S=x#{P!0Cl#L9Sn!GM!^@YH#-g>}1GMr7z!^nW=!|Nh*p@1b zKB_3EZjk=x>){}Y1kN|76`CAommCip)xC7pY(~~%7BiV7%qkO7I8usGQgNyEr_c@% zDS*{(Zm#_E>tDd%0AT)WO+M$oYV2y;nY>1Bm{Tw8Gn%=RyU!-IlI_63f8PNT%aOEI zw}sys99SVx;*awY3d5F~P@EnGrOvlPz^}#DKkosdR1*$zm#mlqPHXzUV4*+y0IRtP z+9t4C2|i!52X8k(!-9VSAj~!CYF(VXhY<}Xm_&3l0;NlZrX*xz6C98nsi1`0tjPWGP%`6d7eUNG zV{fV0w=Uq?WeC2aB_-S^ymk$~BX!#wAIw0P(#NJvy6pj$CxU)&{wOS^?4H-`rso$) z?zYNDMT?CmQbGTY-CZEtyPr?_q>E6Tu8I`LVcV!1qwQCsjH5XaE;Pq?N&VD6C9AKT zc%v1*Q5_C%!I1erttVC%Z6XJ#Pkmt!(bQRP&omN(<569jVe#?e*Y*NLktITl)EDOp zu`u^;Di|}lr*UXY(O$d79}cq@WTkagHE>}iv?k+;Xok^qjO$bxyzVqWco}SLKirhvwDsQ#a*wHH(Aa&-9z1ec$$hKUwi{SlROU zZUg-X-0lpGq7z_FBw&5EBSn^qM{h%0`|#M{z6K&Tm;aY)ky$_@zarRqjG*s@ndyF5 zD;ZHf*FTJ3b+KEA(R$d6pcp}8*@Go@D3=guBab6b4JAl*p+Q{9H=$!ee^k_pL6bM* z;l$*!9l`Y`7~dsj!tA&G8w;D}eue-d{HE;IzsZ5tPL6hQPg5xy(H78)6Hk z>x~wVvG6fNq`rX_Od~4T7|5A*@IN#85{RZ$M#vt<<%xKLNG*kmAM<XI~08M4nv6{9^t-bG%NOGc>5ecT_A*SUml;Zx#weD{l@;+U>vTBnVn z$#z>T*s_eZ9&}*s)w1YdiI?Vj@~pKMWI!_X{h5G5*s7~SH;c+KpCU6J2lIRXXI(OS zd<2RGHclP+QHYXU|6-1J>=5Pe{85nu!4y_`1PBV|6y{ENnxJC4ux@6OBw?5e625M! zEUi4RrG_biQZ0495f#ZE@3zi{EI~u*r+?9ZJKD%a`HBh8=6os-}B-c8eck zRFMY)|3v;|QvkYQ>ZvU?=-w?C3b;AwT*O+#1GO#vs%fJ6 z`3SX##H%eu3ZBLqtg9=s-W_-y_Y*DjR{K*cq}c3+YlN_&2=HEp<>SEqL~@ol6{T&b z7Uh;|a+WD4{xSmWb2c;az_3nqv8DxWux@NNek_ZpYD;ccp2Cx48Mg<^Y@?_ICt7Cg zhj3=iP%gNW>nZd7Uye3-QDy&W9k#G=zBAG2Z!q@E6;c(IP1ae%a}>YKDuCeXV7p9k z;=fKEJJUi|)BDP6Cu3ONRaqv;!^R)MMVT9^*(J;kVKex?_x|)$T)GfmM4pHh;^ePC z6%CoA@4HUzJJKRh3~!<%^ZxqZ;dv(7lNSXnf?KevnWaXX)r8%S!WtY6y-s_hiLV%V ziL-RqGb<;DhN--2k8_YR2Lfp(P=^Um^IL`1HoA2(S3aF38nahpol6PWh8oZr#eU`W zr=o8jHv9;V#j*SLv<;G0{)V~eq0**F_(iacw7#cANv46@30J<~N;m8_VI6Bt@a^L? z>JUdT!AxoX(jjGVIDKKy6Wqbw4VG0|5f^AhUnMa=8m~1zw_(Fyf>mXWU5f4W`{ECZdvJp#w(d1-oa>N9Oo3<1;W%c0OSt7P(%Sl( zcKq?+Xi3w3_LyD!Aj>{|JYL9#4=cvqpV@svr)oGl^mQ&K}al2%nL^X z7V{B#WJo4n83@w!>mz6;o2qKw8*ovo+^n=;pc)9UyH5Zg=ee}hnqnv}HG%TC$}*ev zA{IC~$SRzVE^P2yz7RexyO{}GNK7od?3v)Y4Qg2LMawlyRw;^o3Um|>Op8!%_9|Iw zsm4Z&J)u_JqLFx5sBDc#az3SP>uNw8OQ9f8t zG?MSzFw7p^Ew!Cn z`UT~4<5dx8Qguhrs3&^#<|dYs9SFnPQIZTR{nUkF!pVJ&S!sO^(o_ke2R{6OHZBy{ zFQW>o8V?sQ#nZ>&2?~;=6I8v(<z$i=C$uqkGR12Dy!q>_=>U_^!$%O2ODAVk6}E0;1GlyIR0&Q z{QKdh2lhXPRf6D^FX5Dwz6czw#W@&C%QJrU?!ytA=mK&$ipB~>*cND%hT%WH()(F@ z5uZltTQJ~sDw#r}6~YJ6E;G?m(WK(q;j|a)v@5(VizH~%YWUlE%ZMq?-(5&+wSoA1 zLBs7%^@zN1UqCQvUzl|MOQsB@mNP}!Fj0HD+`w9U`k$Xk*D``VZmjXqk*zmR2Abqx*uhSii|J)CY>{pcb zx!XF2>7S@5;}@cbxg6upmG%|Z827T2%BC&gjctu1lAZ+Ia* zZ}Xo9ql$y_SC4XkA5!E8)%^x>P5<6z{eI~r(pi4MEqV%{ZHkI+J#V!CH1O76e?tbZ zTQd{GhE9s?+$1&y9oK7_GzFWtsR))UDc2)NQ0UFD6*4y#V>z zgdMPA)D(2-B%HxX^xXlLTJ15EDH6b;{{#5aZ(wXms(2;j6x-Y-{TSlox9IgPRsHAw z(UaAqAl9jwXFS5fDf<&Lw1wsmXAEpIIvh_*`V93w9lva%bczz}-$viB=fHUXeHBR} z+Ov}0$8f_(*9y6rw8HyxEv8maoMlgH*Ch5Oj-XTJziFZgBU#pnlxTafhOGf*g$}ga zA(r43bAO<{@h2rIaQJ5JgEdr?Xh@*%2IJZGucDc(ncrT-+jtcNmi=&$GRb&ore!!~ zZa}3{F28ylWIO+azyl_B9}c$*0ji=u#BYXik=rFccAYAlntB8ET>)kiO}}_Xdv~OM zb6p=w`X_Yc)F%fhPXC2_tz#x5T4@AM*BWwP{H+0Vom`=k=(*wCqOjn_5a-8oK7dnD z7*Ok)csMaqPbhRO0LSP!c`slo{O@h!1I6nuvf6ewm{v?rw^IJ!XQv%;VC~Kas1+RA zINM78`(7pf!13>v!xViz$ad8;==EmoEu>uSf-xGzUK`mMp%@7gqq@EP2pyORDr1?jPFkksP8T`zGpmN+ z>C^4Yc3%clz`Sdy#`*#4S{rKzZad>NVdnzHnbj01cD!I0gxCuGU_DXfFWDgQ`_@9i zmty|7-k-y^A~dx_Y?qJCc5VsUti&3?AS2qN5SP;3iom2`nBGz|w#sRgMr#3HAHt~j z9txR1P0OsGJ2?+FUH*=ky+&Z*W+3pKt6k%F1cmIcWGLO%g4Mo;@lv!r<{$H`nRw>` zv#dPLpZZzN-U>c&(KlPQ4argx?V_nIAMTXq-WCJUNaF@dl@B)kgxp4*w@iYckJPG8 zFB?xAb2Dn3=@xdSD1gT8`Y?m-8K&)yjBvJX*6F1-f0;W8n!Wn_TMzC{fI976(spm(63hS0%%b&^CVa?y3jkZecn1EGw!2s(aM*$fNw%5&TrMkyLjYDVc` zy0>IwW@X&&I3MMF_h*RW?|gYm0^cjo#5KMz64-l}@au`JseG2p9CVmv%>TW&%@>v5 z0LZ-HV=1mLifFeY-uQVEe##`U#eqv zl`MUtFmw&@7jPpZ(M75J65y5=HohrEQ*0Iv>=9eTj{gQY8tz{KT5}o$&ilGl!7f%Y zawYmDzf%(y!e1XdMn}MIg|Z?y!yc=@>a9C{Aht?zUWlSM{{8Uel;gQr#Ry#a#y9Qh zpW_7g!f|nukZcC1X0&72j5?B4w7#l%)&+?zcuF<48^HTBr$mJrmm7YHc2_c42JSIq z*uJOw0V8jF72bXOe%9ZJha8u72T0`odyahh|J#L09_;Q9brY47o3ufeZ%AeC5t{&& zuLvJa46c>I-moaH;IhxbmEwoJQ2?uQ9C_H1c^PylkW9?~TN_s?9fGggr%>j9+-#?6 zPQf;K>XBu^2yH-tWy$&3m2=LM!7=?ufXy)AQYW8_)nt91G7yIQCvrDtKY5w0Slx*E z!H|8o*6^uNOG9LZSYA@zpl6y$9{e|m(=Qjg#KyXG?G?tO3#F+p*;PBZj%*T_=C$12 zIHO*XIls&w+j8jzLNVp^svnN_heGp-UEn5FT`#^0wzpt05UzL5oYoEBS6q0%tuwux zaELtQy>w3$qzgTiL?y`(m9RU1%RJ@T%ja5zhDO*%_c#VcwQQIKJq=Mj_Bs77r28G# zukb?4P^H26>=i+M7@T<^k3VQBvNu z;bxRI{ot8H`H(pIX|X>~d(`TI@1xnvsN|PqfOx$4fW>sQQJAy}vp@La z1T+p6SPzQ0QvPP3)mW zIM||*Qf~qUgd65;SWo|ECfAxBUL*VIrTWXA=j?nfwba&R@1f&vxA5>Q?>_mqzXeBh zX#u_%Cl1REo|qhZ*vcjiSh~FyC^fLdp6)+>;e4CR8ojQ;qg9x;z+uRNG&-(EH!HMG zr(9n!ho5P)u5$UV`&838%dWD|+EP=b6&>GAQKE_C?eliOt3|+)6ka%zcM>wWg%Pq> zS{ly9tWoKV3yJ{=dDkzdD;CAh1xKU|*BMc#-=7WBOOo!#vG;(?I9)<3*W5`Y!}cN+ zO`I`pzYmpIikLqA=QXj)uDt|W`%FP(FV^DSi*M(%OsVuMchoTVC8Rkfh480n+K{vG z=~pZyDbdkc%c(857y?)_*OkFZ!JbAHn}SH63|Sl$Y;j`jT=aunQQu`M8Xh7EGIWEu zo9TY;Wp!KKwRpo@OP)B$s&AMbU6l~N7AP5d@SD{NOvk~emOvt_u27kS&JgLJ!{-?g zK+P9jLDS7AmMM{hMHF2=-Bs8m(xOrP5MJdM z?bajXa++zKWX@3Od|Qt^8SbX;8@3huVZdlR;*ehgBeH5Tz#*aYpTT%wjXEh(LVXCa zeC{ao`DCv>Hg2WkT**kQJJW2qUK8^s+)mjPhnjljCuP=>Go%jRb(=0ILV)cA1mbQuitAEA9~`MIUcb47;D4jf-M6~vcoDf zN7C!@CYd-`$E4^g92SMmC;AOD(FQGjY{LJj5WZMkaBTEhLgZ5KC>6 zxeqjjY-W=UqNF-;5pnPHVwIlob5+~Scu!;c$HgUaVMxrW1f88{PSyt=jmJram|f5E zOxiZ=At=UxG79=d7-I~X6nNtB)TsC?+s|IZVf|SY;*n!zw3^2e3z8@*Y)z^2BRRh* zhLWA?Yqh#d>F$m5je3$!xa`|U)F|A>9Yf6I?K|ToRH!5hei-_W#IRfN*O3-0n_otU zBlDh|hFfMYWm&LG<@6|Rz3Rat?``!&MjP^4Brd-9d~vh{23m3b-}`r(B|G~`;H*H| zORXxlO%l%TnBOLXS$du_F@6&<8!FtYehEAiv@O!RFgjyown`<CLg7TCuq?FpE*M@#$ypNW++bnF}DNKyXDNf{Qk!06D?qNXRb zk+Nx4n)Fat?oMT*0iAR!3F$iahT><30pG$cqAc2pIFShMi?CQqmc`<>RQ$MAPKb(~ z%2Pcq%GvB(YS@j|mLLDJ^_HM#uk=d~P>K)r~PLBd5Y~(@lc#!!&-SA>&Hgl*RtFiK87k5zmvvwf8m-^=*>%H9t{5-@ozG-8g2Oj+2 zt%IZqX$6ui4a4+*2Cexw>OVuL&azc(#7elGtP%5FaIY(v?$QI`>}qs%!knH(Er$E* zm!P2FrIoNV(NQK@=GPwlj~}s-+DEgwMiRNV*#nqp@|j!c6@on*l66pGm@J*;C~q}s z)y)Sc6MVEpeCS5~`q28#+ior3G7wn=HfoHqO46{CqPc!M&8YkB8^(LV_6+kM>sL!C zqiAG2WKU8S5y~w(Di#%VH9AS5Qt-S&S9;n0@?z&z#Pv^Pl%2#fQ0<*}A`JT0sr*z9 z3Ct-ySs@UT@rb!`v&>&9or`dRAuxe!+0v>sC_V=#+Nw^lz?raFxRI$fx^USo$64_= z_gS-BYJ~N8=QD;HmO{hiU`mLhZ6n46OG{-<7<8~w3ozRqu)Eywyy&%tWe!l0=&b|));P@HI7v1M8nOr@5UZ$RJ7lH0E7NyE zScB&vMyhjCD3qxS^@p{d5Ojk@uCA+NkrVc2nL-$o;%+wQG!Trtl)jwB*@tK%0AYd@ z#AhpQ9!BFf^wK#{=Df#;a4gfmpEWXl-!OmHXzkJplisd!Gjg$=X@E^jjz}w&-tC2QqFzumZQ-Tq1W$83oglx?n%W2FqhVX4`Ed+Pup5p9otwfewiRJ1Y zD@dGx$dLPDGRv&zkWV|BH`B-mhbU(}%O;R4gtiatwHUvax`QH<8o`~I1LL8pAQn~| zknbPZP>PNKN#c}JA!VpknPN4Za+p(at2Q@)DT7ePFrrUAu~ya6M?WRFGzxAA_tkzF zmI`Ny4<|e>j??Leswhh$EW6#asPtH;6DbmP<0QTo|KYO5P~^v@f>4? zLzz&@OwsRSvjREJT6b}7_Lf1vbu6JmuA;>X-$#I$8h>6*<>nIdv=W1zr+gND4Sfv%??F4I1QLN^3k5^9WG!5phf=> z)H#5-d|yh$#3;4Fr_>E5o@eXj$8ug#6L0K*r|C2>$h2En z2SMq7aN<(-{{vZ)K+xDqCvBsd!lDi725lxAQV5(7d$l8x z$x>DljJTGjvMRZAZIfz^9Q$)%%gPhZpZmsyw5?ZUaouX{t5mN!2oVpxu^AB8)&MQ9?zV z+MXX%yD0Ol(>^t6RdZ+7Zah$~zDZOt7v|w=Z)=dqY@$kwGaK5DYgFkg(I@PmC`bI_X{?nPE(bdP!fAbL;L^zTCi$>Q=jDoS7Z4MzL?xgt>G{%x**eIdAjdHZ= zZx65Sw#q5RVS_#N(P$Yvl+r=Xpj9Zwxj<>!|G00*87`jLgdg9HW%Ocki%+EpBIVwT z!u9444Z)qIr4-R{}$VhF(7WFtj>ziY(roO+nT^ldwVtgbS3uY``h;wVFP#^GblxMVaN1T?$Z zBRz-iiZ2k$Do_jaZIR(k!|Mhqt8j1bO7+_0w69Cz%o9{IGi2zRvnS1In`|7_pQ2=cUpjKbvCieQv})wDACV)Q z@vx!geBaG5=*3xihBhQ=KK+F7!IVgUp4-NFPa!3)l=(d+ZAX`9h7O(407wMR5g84? zk06NM%bO8f85%7+hn97MTIjnflNusf>qFht!fg7aJ5{q_M~P;uLi7QRg6)d2Qj+0N zzamz|2gfne|8{v(mUIo56liZn9%RyXwdueNm*aTfAh*GYd}_eV3_Lx1M>#1WhtR{z zZrwGWA-d}FyqLx$mxcEFW$v9k8hyyf3n}|vGm`>bmqW_D$pl}8s=ei|un^|UpS~b7j0gZY^NJu0dM4Y*kVltdbgqOulX_Ic~6~!)j!|tz>aQuz<31yWr2a zY4UU1an91rwqDaK4$KCgYwoN{!RZdI%S#_nXnap^;`$ZJXgg|_erdLLa~E^bns0_Q zOK;V<%O8^5+&WDIH2@DmW8cX!m|fIbjUJH=0&zEv7kRd#`njhGow8L1HHmzp4X7Mm`^C7a9 zsmh4vxAbgIrY-x$l@%K?gJkHl;{Z7_NUeB`i6}u(wyW_Z^O*`O|kq z^_Ff`X?efyg`HD1-+d6A+mtaLT402^G5!IF6Ve8ulAzPI3SO>cRxDYUbRrX@YKPmE z!Vi9R^a8IAnYHV2$K%NswY%b-Iue#o6E&v&8k=DzGQk zI2(tv%6s?-WM*2MT7wJGq!jI(*NDJUwB`ShmAg>YyWc7o)I6V&PR+~{`-5vw_vci>O!?_#d$aGF1+o~{ zV6j}Z8MWgQMH>iT;kYv;yy22DH~eNn(2I)GU_XTxKNU3k8DNX-A6Ul@spuwKF#M6H z9mouVE&k;eB~tDR64|TZGt}Vq0uc;1C9C(JPC(=Ob4~s4Sti||Ej+D^X? z$KzQhbDe$!NGvC!cNw5uB5^ows5AnLo8&UScHD3-Gh72N0zcVXN#r!Q9#J$R#Cvwy zWr16AvzuBW??`opK4_{Cd`{fLsyX}`(o4KmIp$&6AZsu*eGwK>Q18aNzLG?& zCfmVvHjIv(6N#1XjhVn>Q2#komy~dRomc`VfD#<#{2qT_7RC~n|1bi)poeZ)vQ1q- z6YcEesH~HEhjk(HDH$(B?G`9mt~A=Kq;Q&XwqdUWM#@m!;?cvgr}z_BEc$*j@HHIK zMR@k`ThI75V)iPEJ+UstBBCZ|ae_=>ftJacTPBY;e?}%f&a|<|G~a?#15tIQQ}Pr)Vl|_FP@MXjmia zN&C_`8j7?i>;gJWcZiN+Iec<{@z@Ict!65r{(n==#B^me2rFaWrV}$`5O$ZzX4D4VwJLf|*!%Ba?JZ?k z^TIIwK?|{LKeJaf6KJ4%aQY&47p!}j<(*)26SqFy4^pu5d1U6RA{lpMkQxjqXwEuv zM;xRs$w2Tt(m!nWAXrv;$+q;MwB@vA^&(5@J|O!QO#>!FnZ4l5L#g^fF0g}aH_jLU zXF`tQ`KmK8Q%{bt2rrr5v2m~;nRDCv3ovpPBBJSh!%Y`v z?Ah#zxSa;08c@vaMzMb;_}%|QMZru?Cj>o47 zjx$eyE`p7Cu=00{SIR(EMB)s6p9UWt z*CwBCGc_(z&CltgpgD7Bb{kR2ET!4P2yz#Tlj4A47-73hXxZTsie*0iAQz$2By#=J zNKKfcWOxsO%5cs0fiOlb10$`hlUlXPQ9A@gW-mF$#OP-Yi`jzW7Wn?@Cydr9aK5vS#D-1RFtjp0-D%=3H(c7$OY z40rOvsdF@$4gGlJ>W27YSTYy7<5~@95(U zx2N4PXI&;3e6u|-o9!U@M0w@|7HPD7dzm%wWL{!eWG#NWk#Bn;@tWb8ETC!wLR-({ zGi$QHV6!w=PK2iG3u^??$fQ1I+OIm*Js99P`{Mkf3gh$o_-M(A3~MAzd8iNHFvQN- zs?l!6DZ#Pz+ea!+elW>5s%Bzr&MCCtK$A7WP~52@VnD|QIvZHvPl5dAa(;dN`C%?M zh3-=R!^^_k7ArLsSiOB`n!?{(KLfI7vZ_it=<^U+9KaWB=HCD_wHgou0g;?4bfr5c z{T&dA1MD|6#bnD64^=~e*{EnGN-_Jk#q&*WTmKTXqWc>h(8}RH!|jZF4cLBP0UHRc z3Cn+|gUX-Bm(~80!<9s2$eCbF@jJ4JyfWT9EP%rwHEZ> z_Gdr?Mm+`wKGMIl?U69`Nqo6IDauxaAEt=qN?U9zv&Z20^6tfor*fOrlh zQwbNkd>De|sP$;^*9_QdQ3NzmV-S?&dZ?Nq{pJrq1b8hp=dVUj$5#4N0_nRwEZ+F8 z7~mRxQ@~6!yRrY8@}3#Q*Egm@l#gCh_dQBnx{dDtybn|Y01Md&eOF#Vx@@CYj##dW z;1-iTL$dFe1G4|V{1meiN^2O7NVz{C`TlYSK&blgNH7L(hNhkm+7(aocL$bKpb|v1k zl;;db=xX0%`!R0MY;YqB%Z8MIAkfDJd}=2iymCaRZ*QaOL!eEtLC(^1do7O!)5ewu zpz<;LSXbnfwL_bBSB3p{%UVEeDr^!&`Sv{Hs%EAnPon(Tk`z?$-oS+0@Yj4|JiL;@ z8f49FT_S?pmL2hbSO8+E#}ynwHTLbwJw0_3^SF$ur$V_x_+|99 zxIu-xwOZrBBIPa5g_6R{S5JWZ`9>Q^uvS7PtN5*!_+(TZDtR#uG6w+&R$plaX-{;m zql~oq0ZsEYenAWD(i29gYHUcxUp80&sWI_KuI4x(K@-mj}Pe1@q3I(tMPMp4HE2PMSfM$3ZkZj|Q>@bU0w{ztH<%;!BK%;;I7`7NT zAcx)4wdsF7(={YxYvIS3mp*I}&e$n5pC3of=I07&7xnEteIuCN!T)+%544Vy-HZgc zZj0XnC5I6f(-Od7u)joFx!&&uvuPtzCG5Y!BKSd&Xd3n41l@S+(#7`^fQ+#Jq~DB= zs?JH5@y8JRP1_Azmr02;bGHvfe}RwB5pebrYOM!JLCEHFe7r2;4}trza*ClkVWxpM zBkw+~4dt8jP%|^3FQxu50LYrJr|rafh61~TSvWHn*6qrE+*b;m_6wtM zXdOgrBSwZUlyw5h@w2-Bq_*A`T`Hxsp}F>9rHkEe>zw#ze&x5lv^_9BSXz0dTIhD1 zCasbgHtHXQ@XIb>Tti4C@J*2TdHZTS=PG-9IWcRELb`Grwuw1V~M%gfJK3Bky zx_{#)`9f|#DWw%2^lIH&*s zYlo}(^@!bF07CVR0V#weGQ58hdzx&{dj+xz@c0j!^>sSDs4)+T`Wkpwe!ulRCJuTV z;Edi^8nY$BjcT%F8RrxP;c*i^#^#iAHw|J~INBiJs#U$2;6|NMpWb#ZX#yUyPY_TC zE3J9ZY17}crpCS(#~8m{*WKMZ#6;8b|0Z=db2xiPd8+p6rTa|2CAfg9!UbZMn>WgB?civYdQ%rUT81{2kPG>kl8Tune6P~b&r z@8y7vh`=%qBf%Q#glc(-PO)KPQJdvUKIBc*Ygfz9KB_p-tJ2$^w?yJ*D9rNL)>yj3MN;7G>G9oX0bRmH&(tHe~vU8aAlReIQ^tTMCo4P)uM3r#;n6 z-iTO|Vw6XPpiPbaL7aE%SoGHt!{!XMhC#uw5+{X*X6h*pL~ z`2?dq*OS|tJ2i#iS{UkRU#pX6!KqhpPOMBaHwGLfB0ZyZiGotzzCu(i3n#9l?-MJF z-!={T3}^l^-f$UCSAX<;v43BhVr2PX3K9L5Y$YU~cI^Q3isfz`m6_cIYPYgxHIGOZI#BKG<5y}UBT zl$`c~5JfIATJDKegB4w*b{RS;*zyn2s%0SY&M@Y>k;7<0eXc0bEhTpXJb|Dve5I6x zG^N)E1FCotgm6!0wUzok>uStVrXRcU=Havu0^fiv`QY~sa4T-cOYcMJ$aC;50@7i| zHTLkxhzzx(TbjaiuyYf~95!q?2X*Zr`ia6kQsxlCTJ96jG1zQHBg4DLT~4XlQbumr zi`3FSuoEkphLQ5Xu&T~_>A40WRx;_W=DPbf2u5-(CfI>ubK2z9YcsB5qI^v07YH)i zoc?84XWZ~Vs5;A_IKr)4<3WQH+}+(FxVyW%yF>8cE`z%dE(sD`0|a+>cMXB$YtF4( zbx!r)se-28X|}BOtPd*~x`JSG-F-MBb)J)vsDluqwem)$(ZUG*+$>Ke#3b}M#-@J+ z$jG4JYI=I!Tc)Vxs#Q%&`cgzv$D{1$+RpU_N4pESmZ&%kSv*}Q6R@hR9d<7Worv^L zSPE7V>+t@E-{jVmF@6-`dqT+*)s?N@jn3*r^=l9qW2bkb!4LunbP3lFca<8;7-UxF z5JB_hI&GrX4+SbzAir}#@9&32hFRDb zlcEVywXQ!Nr1_~>;z`K06?y8TsN$-L?u{5)(M3?sm3T*a`V$l@(jbS8N%r54V`Ejg!O_W8i zM@3QwhbCpdEPD8yS;)42InIUAWv4Bh%g}Gy)*P4)v%kM+lJzTw+KA!j1)cPG()E$S z-{%raQYIP3^k~Xkxy_kGz+kOTu&>%E2CZc$%(vK_n21c(K1zKY4^V57X?TQ#S04eJ zKWy%5@_p}^n~|cpy+Tt3uRnuH$4abrs$^VX2%VES%`N!&1a-EHZ_|1YZ|}Moe%MJF z71_@WH;v}7Dj*O78?WEX&2&)B7`-U&h&2 z*ug7H&!5PEH>33me82h?lw}I6S>L1~P6{xYx1D+L8r2GtoHbEG5MYo|)Kfne41m@Y zY^>r$YbP@7&^s8!oDB+sK&mI%?VMp{8`;+ngUsl)Ss{I*mAvgON617p} z$yOu0#S$#fan-|@3cx@)#rO&V&J#tA2dV~Pqmyu3X!IVT5I*PTf28LHxosV<5`H3X-V-?S?H z%4czW1&{tY&*_NBT(hf1jb+FElfLeGTj5igFbKTcv&H$3iymEPIj#*biddmHkm~xm z`!vY;_6=ut-&0));{kQng~7o9zt`fepX!7Fmg_@6ZQKO(FD6>(&qK`}Drx_)Cf_2gbYq6f((`~x?Mly*GdgX~tQAqvYoL4Qra_C~bB+Xmpvvr} z9TYI0_}wBAR-eVu%V+q5343k1Yivfd^Hv{nD;2{>kLS4b#x3odE_ zfb{ZfHTMSgc4`m+KG42-K2iRnHwO_d)D*UUMc-uy?u(zO51)qfr!wrf-wf@ z&SkA1M5WE-$@Tg~ra6Noc%(B$e-AsrPeca?m3h;_g-W+%R67IEp^n8uKe%pdYpI%@ z6Ajx2RJl6lIozt7{wDq8{bT!5e>X~8g|fkE?65DDW2N7wg#Tcer|&Ad%JFnMk-HS`>ITz*%HOaRiVW?6GL{o2BewS%=S#a*c7V7^DtXevI6YO6 z^Xn5{ns{E>Xtt0JVG6c&`@8J8+vo5s8%x${_C6E2^N|+k`MjlY|NCwxx7*;*MAjrb z$-0Zfj80o(qx>BnEYye5cl7?}n`q_bRNd!3kPzCZ<{{Upx;eMFW+^L6-amgEnp+aT z;Gq{)U^i-IH2+2;mhExa$6SJuLT+TQ@W|AJf#cEE3lgC!1gxUexn0&n#N$VwyB8GX zD@8J@gc=n~JtPtaqQpp}MK)RYEsrf#kD>`UjWcw&aH~ zILTuKfY^y~%@t!~S>%0!D*TE*VZt_TQf#{$i%dS?r5lvVFrlh}W^$h@sYH~}Z5y$| zz!n|6vDEJAxs@!%vSgd;41bTt%6b3j;WL+Zl$h_5#fE&^XO_V+5wC*Ty7k5F=dAYF z(?s116N>}l7*`AKYgpV@Q7dGHFj{jtrqzIy73!^k8{UZ)=lM!S*>DXtx6YY5YE4B7 zG1Nh!PmlKoi@a*AGg**LgWuE3(>6j!1I`9qJhDAYs1|l@;T<;MTOxeVoVGGHX5r@_ z=d%_ozSLz@vbDaH*O`Y|Uxzap@)^ceIoNf>>c8?_3gIY0@p?VNcLuZ{O~8qN0KuE< z&0PG*)d{#QENngzJ*vz|<)EN6v_?|r_7yJE7%KcB)+XJYGDjHn+wG?r8k1f`Ox!O2 za1Cr=l|D%?r823KFH4B6H%|rkaidAI%_qy$eKY~9X)i7834|j-7MYLi9jQ(x0D7ME51BX8ShV`};Ze_*^ z4H~Te@~8n;+}W40+I>hY_gJc6IDO?xf+f{M72(*y6@bEwW?(z3DxIdL(hOJZ!Q+bl z5@rmWt8kFAcmu(V6*0c)6<9g02>DFgJDVcyl8SEJYT;NN@Cb}FF^OtNTPOdr_Ru}iq~rTWdJYjK z%Cd&vPBFimWIf2X{OOvo4g`hIU-4ZxZ8rirobIrz5i0k0@vopq9fTAv$jTUFs9lL> zmY@KM7V_x}${p|vSMVUCjzrE8BNXv{L_NvjPHe5??ATxS+miwIbfsrEh28M6<6bO7 z1m4dzD8DQ~7v)-#9)fZ#g(XDl9Yk{$Ci3U-?qz$u%RvPOEJfZ-N(Z0gR-^5BW*yd6)(|9}U&>Ob z*THH{_3fh)&Mo6qZ*pgR8W%zBXyf0!iQyq8S`j&6dT!I5wcVr&5LH1bQSb7Mb}b5; zA-mvW5Y(6U2xBR%I&2suGjon|iE-7eR!g!-SrG>75v_izw$pH4zJmmHYgKI*?6&4h zc#%kjk9G4xoymqBb*q&P(WH?iuc=7BupD|8!KB{oh@rbCMRlSpi9~EsfS;0JA4dA~ zj$Y_G`krbQt~7o`7wOpG!ytNK{}p02M!GOXn}1~#{Z2X#`41=9x*R$2LbU0P z>EYO!iX6*0)Q@TkKN3oTD&Xmk6mF*tT#Igs)tsv$e&3aZ{uiI!iE zM9tDUhQ$jLI%M+KpVLmcbxvFr!raxgALpgzbaMsI{K#3fciKvsa z&zA3Of0IPz5T^l-O%d=~>rW$FgcdEOl7Oh-At#?%LFNJvSJx zL<=a@x|=mclDx=0ZKgzC2AfZHwyKFl<;&>365XL+m7yrwoWhu&ubD^flB0QR1XRhW zt@IW+Ls|e_!?+~rfV%hX>x@YjQH4+#96TQMXG3jA#`{67r?!ABQ@gf$yvX`+i|3;p z^)aX|+_CeoC5X}qKbEcdghC#A((7-S^Cl{tIOaY|rO1Y{Vy9e_bvfNFo5;L3*;eQ_ z1#eExidz^xkRj)kUrN2wim4O)wf9LP;p}&haSJB6Ty>D-L84Bs&!dKD?yIlp_Pp$Jn?krM z9#T3E+akme>=h*tUN`dGrE-*AmliM3w;eaOS3$6ktNX)MEU2}*LlqC+?3Yu}gmCU$ zAG|dQY8~4SAV--93!X`aG}}H_%AaOhsbWuh&=jW*1V#7;Y!GJgvcp=!NA;<2xk<+2 zjq~Dtf?_CNzDC%7&KQ}DHWl(5Vnpg2*B)D6&SSF*8*?HXO{xI$Lf8qUfVXg)>pQQ8 z;0j{R?*)gxQ*($}@XX^Q^xSue;o`caiAGaB%hW|LZw{3dNtJAw5w|*S|9UhnZ&<(! z(CR5*A8$-h4kE`HTmOZq>>|N{v5CII(xJr3)*}Zy<>p)PUSdlOu?UDQR86e{0N!nQ$TzC*kclJqZr>l9<(i8GAeZZuL0rgXI@9?vSOVUQlc*SRC+LCwJX8L&N zU61RTpz~$>krzBR(YJ0?!hjX-Nr>kX$izNbYwp@gwTZ&SeGD|3A(-(iJcIc-=%GQB zSgI>Vwtg{&*-@iI1-BT^6dZ>?u;^$2&!P*@2wR||wV;V?exV>c9f##Jqb8qiPK=N= zbtl&@-dmi(=crJ2LjZ*m(kcRUo=Z>*R_=48CR{%enRdZH=j4-H>Bk8gr{_A86rl%Y z3i+0on_zIp$jb{RkKzzDhkPeWAf>SW5VSBL@u3bL79O3iIfOwuKg!PFr;(W3fnP4~ z1xW@^pksPrjJauc8MAVfN3rBy5h1pNvm0a?H(2W3VAGWGZRuK|KJ`1UL?iCOOV%!m zqN&JHL1ET~N$F*=j_Ks6WYxyRBv6Xg4JD>dd0Iyd>}0C<%XQTcFtEN4~P-zT$3A9 zZGU(&_zCHEZX>8)A_ZLZ7o{)i_hvQt&*OWQ?Ms$wX)I z({fFy1<4TfaZgYe+^w(~SW{ZQQ>UJRuk~Dl2N*4UP|d=OQljV^>15qupONvMCmiu0 z!+1tYvk7XNC;N-{PRb&yi~p#S4fKQe^<3}u>;+cc}K z@D5B4VqB2pp@EK^Y@1{vk(8ks_Xh-Mi{x!q?-V**ao{kGB-y??LP4G-9nf_8qME%) z)N=9qFE+l`>bFE4Y%wxj&)u4Paif1;qt1$1@ull|;K%*!a_*p^yBbzHXVnIab%AVh zu=p#kHF^USrt%B7ak$Vbo7WQ=F09J(EtoRYn|RI()9W5GW##J(IUkwep$BY0!~Ac# zOS9nFL1q;VQ7Qf4{ zWX3x+-1r)J7IunxA#}ve6(d;<;tI8kvhMJ`7cjzMO4_~QsK!_k723~Kv(ZZOameYs z~HOpj<2Omdzf`Z!|=hDKr>#PBEvZekm$bJ$O<|_ zwwwpoNW#O&j`dRtsM;_zt~g%m)GGy(4*ck*c=3OqWz;H(ljrv05v`mNf^t&#kEw+|PLnYwdU zSn*mj_=j4MgsoZM#X2!qjaZ{Cx6!&U9GwXs6kWM67fZw6m&i2zk-@V0B1)rH&Qq{< zfw?LUz{D=H_mmvy=WAlEo7=|FMw$CI`w?4!c(6+#A`77NWsUZc=fVB~06j-PE3N+OXsV+E{Us)+IZ^NWmBi{G#f~Ts9VryQs;CO_+_XYSCfCk%heM#alwSd)J z`FsOU_;w2HfCMIJIx;R1Q4R{jK0=|p1e`aHPqllD)@E*Bdwj5snB(XQQaO4vN79^Vj>!` zy62;+m0G#zi}#9ASe9%;bj{^mT0orYVPa}THkg`? zf+T73wd+e?9C=M@v|!bbs#_GwxaMuq?sf ze5ZE}KYl?m#Bp2up_OY$o!v?q%A9j!lrllDsUHwtnZ{fp&7lc(05A8RyL1~a>>oJJ zH|RFGvsg6TI@>XEB7l;1?WdvdqCIL=XN8F$T>Qw$*-&**HXT(n0aw?d< zHjEkj?P!8yY8}8*m!m3u`~j#TVwv0bF|4Wv@e|+Y6GBA|CzNQpbL4z0<tIt{W0Wf|ZL0AwZqH7WZL%rlUE!=2WWyESkJslZjs>shy~&724IXT@5M$XqVr7HLNiR7*AqgHi4nx%MvO**P2U= z5_0VD-!Fl2xGXSODcJ5j8top=M@7JoE+7uPpUDbUQ2cRcB&6hq&k4=39ENVqn&NA~ zq7BD>0~>z~SJr30Xo*=ETykK)_U~T#pN4yVWD+<_uqIAFHH>R?Aog7)k;ONH^K>B9 z4fXqUVqrH<2BwEyE#l*6W_TTelqbr#DLcnijxx^upNin9TP1O}pN*WPb3S1t;d;scN zz(b93*(laD;}KhE9p)17inaQa?p&jl`dNqdz9-=dQ}r_hQ_TPU6Aqibq-k22lmo#Qxbd!6$|57 z^JPrREb(h7HKnkYc(?-vC6S*iTEswrav)$b*oG`pU-ZiH8lg1 z0+kF3X?Ak`I6Mcci4ArtRODh48}kT<=jdkl8vU+#c7|%fUSpAOpsT^i5TMGV(E~u% z4$9z~7boswDom_q02*cpgd|z|(`kkk>YhwEoVnSh9!17H{I>VpdCa8vyF7>3Cfgr& zINRsh6I;zna5J#vC+n4PAcsE2Xk+{0B!V|K4Ib!@i%d{Qb2KU*R+VaA)KsBW4nM*> z9QyiwZfVNEYi1Ho*J8_onRgVJ@N(O9Z<2s#Y~8&c;*hbY-6D6m+)o*+?Fe^W)&`Zd z#8-bGHZ{?#Xu zSxGApJ>!76&m?qen3m0&rC%&A36&-|6;&1GB6#v+R!gQLf@Id}m+vuXyu5Z*q#d1g zlxp$I@s7kDB35{Rv0b#aO?uIrQ$P$v&nl)CD#_?{QIntWDTYJWI7JPXEq|gHAmPf< zaF;PEZp)MKKD&sgf|nt<7D)#2D#CS@?pzH!?l2=1@Uzi>dkq|C@5;A?esbZCU0644 zY%Q#zYsazBgZUfQJy_pP65{!vh6r;1 za2H3UH#fg(-J-Uzc>*+|xSKUnD~^+&{nOq%XTh8GxUxn}2m2>1#fD*1HUDRe15(ns zX~Grq&RU|`!g}UZ(c3r+o)W|lhDsC@O+z|rUDCp@#Rw`IAhr2J{KiHGUyub$e<;|O zj1fr_0bzuR`lB~ngbLEdsKuFhoVrRSQ<`7hfh4K)*7xm9WcBv=zFVVH-sAa5-56vo zO2}|la_(m1Qn86vN^8Tx^AzojF|4kBm6WvW`SQD@lLsEn!QrHnmsVGgoH-Nj*kVmy zz_DeE#r`|7^U^6r#Fr%u4MLzB4iR2c6Gq2O3(??uNHV+1;*fGTmMExbki~u@@@#d7 z>Jo9vXzB{srJ*3I4qb_1ZazcPm>a>z;;698k&x<|CfaiSwGtKiFr96yVyG?7VBUvG zor-QpdC1E@Gc|I49NxliW#1Kt;Qf1JlORYa8~t772&c>aS!PhD?pkb~$+IP8z0)Y8 zbg#I=#$$VaA?g?6Y_pU|x{6|S-H+7{D!7XYf1B6*&J=@`DCsOKtFy_euv3EyPYC7} zy>@cOa|xw5SJ9M^FZq$y>%P)LOb@os$n6 zR?mn_zLVU?$bpUED}_OG^?U~oZ2^^G8-8nQP??b&!B>wSbTr+a)B6V|_|*C#`!8G* zlWHJpOfji_V*e+WIyMlK7$rU zuYKy8nRA*;0Wa;nIx@;Kh;Du1aA7Oi67t9^0;k9)xV<+uk|HvHGF)RYf2)lhaV&I} zAB;OBqP3}D2h7~!R?NMTNm*1EmTWe6;e7MsN-|CFw#HH(p`&Qj+)8zru}3#*ZqJAJ zAmA}ZVR@d`(f|Q)&H~|F0qZ!t_)mxb#H6jYv-dla`*n?S zj!zymH+g&LF}|=|MR6A-Il!JXBzRJq+6_oyXNP`2)PErg6^VhJiar1;Zy$igvDzB^ z%P!BdsTwYtu_igO_fz*Z>^};6%ZXA_`?>?OA~E-uqGD3HxG|B}8FSl*HKHZi1fWR+E&=Kiz~Q>YVKlG# zTky70@Ox>??vrv!9%UfK+xNYM5}Wz~$yYgb*`FXI_BQ|yy&)jUt7v2!5v~zJHt&@c z#yK8J2+p>$a^EA70P2v$H6EEcQ=Q)3{Y*v0X58n0etjoJUy12uAopnKudoFU($txM z+^f-{@)ut|zYNj{&wS48*A-u`JBW__S1<4|ul~b|THzFGn)8JCyyTn1%5JMdQ1{om zelBX09w4tN;KJV6T-qTl=toi79MH0SY4ExLky{v4+9o$THJ?Iie3&k@&GYS0j7WAO zeZF7VUXB6t#0gbs%cg44yI z=qqbX7Cq$S%wQU_%dx56DonsIo^@rsLKW7|k42pCfDVD+>y4o0uVaQROLUKSM%h`( z60MHk`p4=uF;MT`FHU(l(BHFt;P%SqDVwHIar{F^v4uo;V=lgW*6`XgN30woc|xf zX71^|$6%QnwP{SqE@v`H_>k2(@D1G+@!?jXUB3XVa~EuatpT%CpWfavQ_{boyF9G*Zp&B3mR!(s=nLf zMZ`BizRO5FG!Y5`NHtM^AH=~Y8V&`_t4jNhpc-ck8zEzBMqq`OJktkJ!a9?oMxd^%hoF=v_PnWZFgoa8fabo@2vlPvu?jydmlxHz67LMBQQZVmX1~%$j0V zgAblZrlH!%Q(ykxVOn#==*O37Xt&h(0Nm?zVF*(t&mn+59+uD=9*u&d`Lf)c|G>!d z{d9P5`5gOa5M~4tJEFz?z@;=Euh+LnYwm838fm^K;Y-)V*!@Cv1-SR->+JMV%EP~Q zJ>G&i+b;x+5X$$<0k0(^fla4XZvu)2j=^})q_Gn_04g;>k#dOh(k^2Dd_5H2X5D9r zZ902x{NF?pF{Nc7n?Tk&yPBI=FOPbbZC(<}1-5r2GUBOKG;SXrVA3%c7yo*{KO^eu z>hyi5Q-)g8UCb&jJA2gUhwtQX!yNPBinWLcjN@6*tgIAOFzW!sinw@W)yl`Bf+5lyev2e6nNqTe0MjnCfz{%1){HOw@&ah#x5yOZJhbvV~On6unG-ewOZbck702B z?1bn4@S4l2CEEAV*|tGkJGJc)4V8vhYk?m%M||G63i8hHu^}yp_cw1a!f#||G}QCH zWr{oku=3i5p6R#_(#PHVR4ii*MpfPs;zt%aH5ZvAPu1H!fV@`_cnK(`>0rff4IQ8h@xc*`yMT(#+AcknBO=Hf0jcf>dM zq3fJ~M0yk5^GooZU>m{3dXV|GtjfO5$`Qp5eoc9IR;b1>;IijtXg$1A{UI}!((?S= zQc}-Cf}os~5s4PvwpcZRYFBIHtC9oi^m5Tsfys>^?bulEPu7D;0k6Ghg9FmxdiZ~d zKmH`N)mC4A_$&Q)e6&HeYz(>-F{)BVRWw^Jl{{3Z1a6)(!GaYd#b%JZm=2C2I!OYy zN(%-XJ&9Z`jhs{l-IrEI73{iSzeUnB)oBb`$#0(WJ=*j>>$HDsyjpjk>R!xtul@0~ zA78!6bQV`Fd>zM|Eq59UH=^%dl3AyE;gvI(LSjHy#F6t0f7(^O&s^yOIot@I4IaT^_) z1|<-;GyX6~x+KfxX)y91I>C%A5IdSQ5RC1o)ipXgJv^!S97Zgv1_KEooM3`S@2$XZ z-#Yi?AC;JsOUU5yg8wYPP1ltntDC2jw6*{H`^T`g96{*&w-u>fU_>EpIFz$pc`0eY zt}fcFP*tP)f$p5ML+A>F1vJ?}P5{w$W8X8`x2IEA5Pz9Bs{MKlD8Nc?1eB&bAR;IVt8oe3D{)Vt}HG(VXS-JH3{R z4%#(xqAZPs+yx(TpQBAR3MO?**@#nRGZ19@dwfQM7kpoOg*e3VRn?07*e zsi{G_1^`YEe0`WjR;anZ=YR1}?a!H@-mh^!s8$R)H1I;?C(Us_;mIeQkpI5gV@1~9 zei=IC?Bj&cTxPCax9GIMBTJZ=5D2<(RXCcj6J$g>pxgY{1=59Cz4W$nwKrgbnu`h+ zQdomfrZBQ+vAwXe7sP6K6y1t;FH zcLaXKtqX~mi-|6aERZjg9WNnsexSv-tZrFOocMS}gfVAIwMpc7yzCnM3%!d_2 zk``>`FkT_Ah{hWGJkT?#T1p4sNGNbN=|IJ=*{-&zAvLVb5jEi+TyP&DcXMAJ^!&b# zQTL&dw;`kDzAA0|FdGf3caNNZUq6Vc*d-**Bqx#2+OW(bS$3vGWZ_xKU^X%Bih@*v zKBfqk&n7Vp9~pA(1r2k>LG6M}q03{Fj%%HMY8KPuj{*qM5alGYl{kiHI`@2DWC_Ne z)*R!kVW(sMkJAcIE%T-ptn%bVNKZc*hv67yU_20iywpf(LoOJzft4~jO)CGO>oG$p zkbiEJ`4QRi>cf41S$PhJFx4q`54CJZLXq$5=a zNPaL$8V!SG%e@dVY-*0NR6nXI7GEfW4tOKGwc#l*CSnmmc@>3wxll@D7{MshM=!<%1~ZbgZ!8!~+2TaqZ99*AG!E^@juaIBWRgim zG|uBn488O-Hp?hXvw@4$*cFPB_%D5!dnCh`akkW%LMU9agcGjJ+*4glUyaYqWC2cf zH>0CGtz2`IkK*N2P$u1-;TXMvA!_-d256ZDw`moxfQQinBn0}MpK>$`>pnTSq*%wk zJkQ-%=4XM-jY4D?BjD~RD4E?T(fM7K$wTQYhZ^fd5YQL?gnjr3nDHMyF4F8MFG}Uu z=b8sos!(L+#6KSI*{a|Ow5@XRG>E@ldO5&8}Y;+ZFBrGNio)3c-WP*RYxd1;w5_A4&+YOY-ff1PJ%jq)sXy@8XHmK;ybOJJ*LM_F7RvK^gAo6 zvJ9j@a|tA34NHz7b9{SPst*(+)aacrWt3`e3X!hrhFVTi>D38G9};`UQ0wm$X)BE3 zr$lOwJi}0cs^~@fe#-)tpx6GHF{BDCH<#1%cHumf>cxoP+*#DCR9?ayJYs*OO!3on z>@KWm_2)1gh#yj6%57v;_k*rDYoavZvK;h-DP%E-V7W?q;AVW`*?R>rqTRh}A$j{w z9}r-@xC;{}bartGmZ7EKmMu~hdue5lH4$xJNe`2W>7W_YtQo2a$zmY23{TXRY@4Zi z#GoMJOcWPzy9(PVvjrKHIN z`Y53j;E`>mh5|IwQ}11Z#c@SIsCqv`PNyXESRFl~HB>MIMFuDD?4Hn7F3qx10V?fOh$boF1^v{y0|aZwRE-SE>vxlV1xEs=YNmRJgYVoj8-G5ki@S!LuIO8 zhd?d#P+;F2{MDL|uB9+31v_%u_OYx`3+m~Z3LPP;ednZ;2gSx6B4x(RIc@K$DF;-x z)^=?Vgek9v&#^tvu?p_W|8RZ=a{rs$dR6?XQbs+Tv(zNV?jE4RD-7s^1z(O!6u3V4 zqV(3COt-fg(~r}-*(;lF`0T}-FPI4X-L5YHlc5Midk6VE`mhy2rZ8I|p=PV#dH=FE z)i3{!nTm~nS;+`-U<3iztu5%6l8Zu-P{|wgx&-ehZ zW1`#Hb6>{N`V06g3-r!@n~ZeK_s?+SkGnO0YkXML{{oDg{Kw>t$cNTymdQZza>;s{ zp7yskv?l(S?KDilQ78}?k)xJehHP0Y7X~oe#St-pA@h%By`Ls`S+X)wK=mH&!l@!X zEl%rVavrzF`-1aXvI;Sci1Z3RQNZSVeYLI+FvWVqW2)w&d<2%?OX^6Ln+ZtuqxWrq zD&cSazn6OBn;7kF%ixTmEzp8dF23-pjViqyC2cYJuuQG?)%!LL`Xg^J#ksVJX(N!5 ztj}x!SYZve&98@Q#O}miJ%yB?md#OGJSGJn`?#J z(A;FWYnZ>lcD-pxC~)oq2<97;G-SlZj)4W-C(wJkaYh%+-|l{zUn+ms zdfwIq0Q&6+v45Nn9CXKcITvb4a>hdXWl5;7#`O64B91V_< zMg>DF7YQ*Eb5SmP#YDYh+lrJP@Wj`8Wu>Vja_h4=szo;AT)vKPQ1R3X7Ax55ph;V7 z01rP2@13)20XJtb%vw3YkN=fRP>xlInb{=ynX-vKB5>nh+@0sea|iAgxJ1V0M?Cy7%u6_?pPByW+*0F05RRL zzl<0U8H^!5x7Ztr6~O#%U;*?adQOQ}bdum2+SVKT)didoN!*aFKWO9RE%jNzKe$@> z(rr4E@AG8xz6%Tsl3B^qnsi8kxMv*S1e*~~7TqwgI3(@${1%W~6(ibaTFprNN+7vy*& zz9VkRiF2aY6TRbZO!hy)s<(x`bw<)phy@=guLfh1HzVedQS}E0w`Rmvs3=0H$QxFy za>(j7CB{}hkOZL3x?HqsQ8MWwH$2-Q<65P~;BEq_fv?5Vg&{>-rV?-{9LW9*xaazg zt=L^s!sbv&C2(P2Z5FwK<)l%WQ6(0|(eYI0R(pGqT|uPJ#+Z1`fqyZ>Ee>Qjs~F_?TnyTlC6j zHI(!YJV)2pFi@H|U9e#c%8v{{Jkbh7WIR8S-6C=e_dt&PHJ_Xi?<9w#a&9qpE#Q0t zM||#6ZL+ZOXG%0!b2ls2EEp&a5yPbL&)YMvxu>q&qb)IMszqPY zv1_P_tI$u8t~^8u+3Hn?MWM?La!Ay3GH^mx;@?|5C&8SSh(wW;rqvSsvGdt)x8__d z#?WfNh^RakWXd$w_Py_Mnc%GsMKX|lYAIz(XmRy&`xyvSzXs#Lk&fci!j*(rT+l*` zEm9@W-rLdd=o+RIrpKw#bEfpc+x-sByqX`{R+Yod?Kmy3x6Atg6_E?$g4(|dYx!G- z1?2eSe3!eSM=Vi|RT3hJ*TQs%bH+=;N)tV+jG^Mr7wlw~=xUQJ(RrNi*b#Qawduwq zbjHnyfb>UtB{xBSD%6vKIX3;?`N3`36HgU~Q;CCg9OF4oQw*65hOypqy}?U6SrQ&Z z&Plz5FZ`E^z5~}B-D^&fC-Cje_{ry|y792iS2mwzP^9iPyhY4O_o*oK*agkq{2FuT zWk=1woa=hLwn4npI_B?__GZEC$z5r5NkmF;UqV&zs#-vxF78%E%wWD)3|;gI4z%%- z7B$7!7}Il!J*VuN6D?+%KAbm$H3L_rj~)t66DJK`A+;C-KFp&n8{Ht*Z7NE7wN}e^ zx%mZes)pUHJPS1j0avmNrWLKhYb6E{JlxOq27C8-*D$UFI<`bd)t?c$TDG*(Kj|;S znn&p@p!HyIzPvZwH&3d~F8Js&bSbT=^{(Yc^nCQ>rJRPp`O_T=i@lrwabo@ZyIEed zv?}y1TwX6ufh>H2&fMhZIr+z4dS)9NHTAH>3d?3FE6*~Cbal_@KXf~gtQhIV*wApC z8AjtD0fQLib+kUq8?J>zY~q5blD{t6+W1PFZaA5hX&i}4ykXcrJkBq~RtmrK1GoAZ zUy;@Jg*MijVmHB6JeJZ0=@WQqi&r9-HhUKx@<~U3)D?rtI&Ah072_mZFi8?Oevm>Kx&W$F9}VW6ppcrW&U`F;GF3ryo1Q@ z4}=zW@Jf4$HBGZ7h1-V**RY836>#$yGC?(CgtYz1OK0UKljnj;39?I(sK;0)@JTo& ziIo0Lan-4Px?s75BljRyC1~Ts?XL>BF0Ula5kQVoc44705SyD5{+33}{SlEcn%o^Y zj@YsXR>XhQ8DZn4z35m?BvpTS(+Oy%;Xl23XIUxss`oX1Ib`|yKK=jA{OuphB=pXX z4O0EmfJZ#NeC%nC(+dDmkxoA3d>-VE&}CiOhGmMeJId;~9PauOQT9vTBSZRo)3*CO zodYyjr?FqW8aGV7BtjvsWK%>c%y2E>Zu0l&jfV}G=+Qobbt1~bPe|b509#oxLN7|9 z{6$hHTnUZQ@omIs*uuv>ZspEO&^(S=JmQA>7t^=(UE^_N)Zu~SWX7IDS{udlT&Z3Y z4kE(0vvO}6AXswyT^EKS=7n`)@D|&h-6>d+%rYkQ98Du5Ht} z7#5_0@|OQN=hSn77Fxd==7=EqRvd>7-GWR>l()x>lC#r*FVfiO(>tQ3%P2+p2QXT5 zlI$k=-!bK&{!vO7xDg6&3YTzna&KmO2~#XP5<4|=X}l8Y=UCN*^v4sqDKl@`#|&F4 zFIV|O+F3<*h_MbS94;}r6I!N{>l0(`D)-X7Yd}9nM2cef*Peoc1{xKmR6FbSia{VY zFzr<`B9TI2H+DH@6OaXQR4ZB1p5EknZ|ZXXR3|FxnZW-SYk?aZsa^sF$E@{$=u0zl*)~H8AAR60S317<1@C`8xAIxp+dysDY-2X|d z&R}J{czE=$aL@hfQupFtGGF7?X<%POV$$_`JC2PEW%yGrG#9OsmYcYHVh4$KlS|F8 zSSY5$rT!FFv?*&7oC@LKbuY6xk{i%6@VIbGHJKs<(atE5pyeY>=(e9M?RPZ;rK+nb zaiiz+Tkcd-+mMPHv7!luZ=EqMFlmAq7cT^!_CX9>e;@6L!*h*HA^O`|N26xCAHWR# zVGmZLE^5rV#Y=e?RGGp5fFi+x;g5Y=Z;^Vs&?;cP8TwTx^lluS;y+^~*k*1raRco* zMGd_HdtW!P)7emnuDflGE8O_T$rmZKZC`p&f>r@eKtEMKk{M;$#54<@l9r*7+h5?$ zy?|DH>&w5lOZCBO59aYtKS#^GCn*X`Hcp$`GDyz91I!f_ThS@26%VV~%f=Th zC_x_apJ39E3%1U%h56Wb7OnsrDqYwe5!$(dd|t8M)il!H+v__`lpkKOJtm?IOK2bm zg3X!`ofGJx&a@czB1xyA)5E@?;cgm2Y0v}G6@v-pY^e1$ zdg@r>JKT(yXCqy3M~JB75k_f8rkM8X8E=@&8UObQl1n5g^QO!hd-ZB$`7{6^@j(0| zQT-rRh5I}ljaVco2(;PW8hV0`%>)`YLcZ|r9mTVqrEa~X15hpnnm>F2h3}572yGY* zUq$RGNKeO$eE|hZ9&zdz%x>OZgCIZ3c-2(fx_oc~L@ZzY$%rYd=^pj~7){@$x0xRx zw0+KK|9XjD@*Fj%_P^@~bA1o9f}}?IqdBW1I8id^zyh@w29j6JV_I&EGf4l(+Ia`V z8Mp6#clF-e>K4&E(Id)Y3DHY*tBW9dN%XpERwwEby+#lcEqd<~5uyaqqb4Hf@t!#| zzccfm|IX}xhFR9`^L+37x~|XFv29TJ8%@FvotGWh@vjIHN7P37a16Zb1nk`dsq|eq zHuW!haOmbTYXpdu;@DNwf;rreZMCId6=w8+PUQ1{`cRvtY++uxaimI1~O4d=+#%% zS3EPH#7$+LjyJQJJbt>p{%of=H@<2p5}6X$)6yLNH1m`NGCAh28L92a*yXr#XSKV( z-7#FoPs`BdQ2_}db_YiH>ik)*Jl3U9@RP@n_;k?e9L~DRJHL4n6Ln;9@Pd znO#rOa0MF_(To6Gd6V}3XY_-=$v@`|bkA+q_mZQyGzs$gb}3}7mR;6pH#vJWQ;M*s zO3rq2XLTQZEil%gujBf|5DXGr&g(H<1<*q6dOv44J81J`PGqAQ5Gkc~GJvO%#<-Zp zp7A9Ge!VpzIk~O>US7mr=~*HLch`|C=phx;Bd;>*a1xOb3lgFmddz`-VmC*U65LO& z6<=TypQnTA-3#xtmEmVJn-rL`Z(j^?#*wZ3!GJgUr?n8({a{YH^*J3$+@!Co#26V+E0yZ^;9Aegn=Uok zXrmP`WalhX0XKf~s^U(9L3YQ}SaJ2G{L-_$%>+lgC?$65 z#{-H?Uer|bs02jwq#f%i&6XF&uHL5`l@Mpgw)^Gsw@69?qnQ6l0be1XTWex_GpVuQ z-0PCEPm}BcKWjXmz2~z|jf6KcVEn-+WA>a_OxIif&dbBXt7fjjNmRsdn zECn4=16VVJxz)$^xjbgx7?iT8}XiC~kzLn3@N0mng8>#k#w zgQTojOp#>am*01i+-L*o6{>$|t?a+eC7y*xQfT^ar``!jT^y}lboo2!Vz^(UK3kZ328L%2uRMA)k8)yr=E z@sSMvaWMnMq%qWdZqRUSoP+qah$M%j(!_Inqt`3m7S%kWSQgEwz%eO>44D^Od`k*M zB26Vs4G$&|4^KPS+Nz!(%&PLDZ9wx3G*SLf*=S2ms&V`4E)192;qFmXKP8U9Cph_# z;tj|k?*Pct>rs{U8y&M+|32)oR?_vvI0CjmuP)yAFs&ixRh~gQ%CO_&|9&j2lJeKD zkEURN{*7ac7M|^&9-`1kz+9DH$UUVnHNR8$F!R>qpUnVs1qWd~O1QRylb`)GxSm^L zB&;e-R86@+sm4ggz)?j$pPt?n2O!4^z}-1m;KU1$|{~MQ=t| zTMU)pF|b;0aa&8QdlPVp-G1*)7GUQa?jkLrMwSd2mSIVk=Q;PNsRh8MjQAohk2E*G zJbOKMS?EAk`OgeI1FNt zV)fJ;z_9OzX1c%1P8&Mr2;bA`md4r87vuZAsVl0jfACgtHBBJRjx6b^r-y~3&^>OU z;XS0?H2Dv{fu&jgm(=HFRFt-}Sp{j0ZQMSsLWbxkf@a45rQov4@%1&ma_i2?@fZ z&47izX>YXAbWVcq{$%SE#+$GPdz(-C=e@@7(2$X*30(KvA1DrSs={R2B%VlG9J|Ex z<${ou_!AM_lRP$O$&%#qS|a8C5IL^YsO(&pA0{A(0-m3wg)M6DP_$o!>0PhAlWB*< zF9Pu={BB^g`FJJYM>X@jzdaP6{@kN{@~?XPBluTPS(}3fh*IBnj<3^l6<7EcABnxE zSdh&?Xd{IZy71Ip3?7_U2d}tPt;-62B+gEEC(PUCWZb0Afe3*S-jA+rm@E&rUlABJ z7kYM5b4gAozN1!qZ||@G>_{#VOQR~tXxKpWnk81^q9ZT8^F8&7_78Ern@7p#ho z>**Y|-eop#_@mhT2nmvVC^GgSrN?cA_$)<+6*v5xItE&x#99ZfVZQT>G~VD|9Q%iT z!YnQ)$p-lO3F)3-v^J%@pXa8qwmjEUbOn9`k!9;Rb~#_4jc{FR@qU#bqA3F`RCN%Q z6UCL0r^&^srZ6+pd@C$^I^DLnOd7JJwF?Rh6FLLz88f)f+Xse9Ek5KJka{cEoL5xj zu?R%+Squr!>5I~TxMwp3Mkd5k@t|#;ffV~_FYp1A$om2v`fzhDlW)5goeaS?tkWlY z9IQ+cOb){b0ryVoC!4&j$g=a(c#t^De> z({VY2pH~TpoHvW6;SRV}r_jm6KUrAYSrKMds+%xSrF@U@Y>xtpA~x60cZ!!0atAF#n)fbFW-O>l0IvwjEPH_nsi2KurVf5+#HTa3J4AMn>GUTZozqPACF&kjt3f}f%HF3H=l7d8Su8PgYVH>gI7imHDz zPwb{l0RD@xtkmi8IbcJV5EgOo&lfB#bX3r*r|28%TWYdp)|`@bSu}!`-fCD^(uf^5 z0lPCo4+po=n=)4~h>no8+Hk?NTB+*o3LReFs%Zwkj}}{yVnT}JD<*{6eT~jE#~x9O zipXLO`QdoNlZPn&G$)lx?xf7Q4FLz|>Eb1=>&b-r$xQo3fq zF5h?CZf}zJlaJafCU7r^XUBlP)XjSh0>csIK>vG43dwFITQokUoZVA40XL85zZ3I8 z1Y+~IqZ2Q0wz;$&2VXcw#%4^sl#hMtob95PzD)+m(*|@7@@FB{7VsJ(*smpQ~&&oX-BU#;?}N@sQmH0vXxE4ap_eI z$r5qpljzT^P$~U|F=+UfFN%;FWuNM2ooZtuIm2|qF-3@_!R1$JveCU4C;6jPShxjnAfQ~RyuXHrJal81+Y1fsTr(Br}S1Uf8xuV^yru?*bs!4&}Abb zf?gas`j0fs-V}2t$X-gWP4C!cI+4sdcx5n+ij_2|Lgvj_RZWg4vbHHls_ojF zZ(V;6SwAY=X524#%brf;!um4CC72#Lh7TAw9E*C(Vg}RuxAsuy+^;l9Rj=1&)n{f# zh%>w2K$7trqWEQ=cHB zRT-0T29+^l6WgM&_d&icQz*YBbtC2$5c6cxCM{zwqM;)`5FqF>2x=^{eVWx|N_&s3 z_csUbFAe&@{(8NIl4+a})m^C;11~s=i1KkGq6@~5+D?^2P_q}$B3E|HK|+lG=x#(Q z#&6MkA0=CzBYTe`^@Bm=mVY9RtN=dTrE0GJg3W+M-wVu3pwq0%#{P-B#E=_1x0=0N zhhGg5r{&wB5mpzijZoV(0dY}1+~Zly=S~+@q-WoPLIp&GxSZFaTigh2h;r7HO!*Af zOm7^B>-uf_^O2aEx+!vHX2zxqX3u*i_&vP=ohLli_eXKE@fjTMeVNJQQ5eVnn{712 zBhwuHt~FU0+uIn~g73sgob^0GKToi$I>R6+XZ`tmAgO2hK_Sm$d}z|^Z3(nzCP~S+ zvDoBKHOV0BDQRfAoHh^ZzeW4QEsP@A$91ct|7SE7sbvGE$C$`{8iynQlgDRg4)srHvy z%LGbq??MJuapmkhHS@^h^2x=+$QSf0!kH}*Ngy_x*H%+IcI*+%xMHLOZlq0k z#8egDn~rbk^Sk%#Lf0vq2z4UJ!;fwI9p0!UkU{o^Z1LV>laj9=EO9gdcpSfvTGbzO z-!`*`BX2hS+pUt&ztt;;zq1)}2FG_t@!(e0Bb0}RP;blyDV|VP&8RNYV@gpu_vKle z&ol*Vz7IMEc*RBD%;!6Pg|X$%miUi2(nD&cK&2K+@oc^YwCHe^kBCDwPZoa4jfL`V zdZ`u^`Px&AI6u2Nr)P+~ffHHRi_WY`v1ga=jZN7Ny>ibkfKQ8T-Pp`L6qQ@`Jtw3_ zMx@$b?pbs*NO6j5dd}>>`1@NcRgaJ?oAp&mBQqTzBxspULIfmb@&em=>;~#k2DWcq z2m_DG1QumsCr)dR&A~IhvR9Jn9_cUeT-zU~#k2k6V1E=Y-jJB6u=Y5UL1!FNCH^@v zKoN`^2R{|>qr?{CbdzyFj;R=|M~9Q0xNgLR(L!**Yt+FK#@cH=GkWzcdp#~{~OY*D5eR;CO~f$BENX{s#>1_WQ%ZPw^keh&C8d-P5DIukn{M z44KB8J#a>sIk&Pz9)Q~Ot5mb0fEOLSptG%zc(pAD@$Sf^NvAY5dbyO$`5R3rP?bDV zOy6zVS4qYDCT#8xHF3q;n_kR?VR$hlm^x`%&VBgC#?)$idPq4*xHGXL)q@kZVTj*G z>-4GL6^ntG-7K&uoSiex=kpg;nup)Dy>~hA6!%I*_h~(4Xsr5ua7qn&t1C3@=_4A@ z*2e}OwLwQh@PQXFaN~lJgpHuR1q})JC-bunOU6`@_Z2If&jf5K*HP=t)rd3$k-Td4 zy_p&;noSSc_TJ-_qS8Vp2UG6MYKg{me5Lmpy-?>B)V$G|*2V{ghgT^k@e$%;4 zTYr3rU2tPCZiLIkRsuVR0;KV|S=Y~{BWOBLnv*=kOjo0elhUQdpQ{sev4b-CihIhW>-!}<7<{< z=&XR@EPkGzrY~#6J!?o)k2{mcOh~a~p3*m>$abC1*9Pg%Li^rX)FQD*!0CZpmuBPV zwq1M)o}?vmBJ@eo7E9v`&6*+3Z4P1!pP3{ikN6X@HZP-t zF>SKOMr()R7TOKdVHYDgyD7rQ?=2ZwfYbvtP`;mkqIsrE%!%0SC~w~B>v0=bPHj31 zeWFC2Mf}<&?VM|O<)sy(ufT?zXOikjcKC3K?VpHVomf&K6s$Tp0v{&tFRF(5*$YS= zbEG-3zuZU;L-LAO_zlD)$RE5uzka6s?-%dqhyf|wGCcdu{Liyd!U5Q4a*Sm}Uk{p{ z?ym-Hl%^=kidUofy(RWzLR+8oLk)Ab2aLjo>f=#t{j>9lY1-Z#BOs}=STzTEDxLR^ ziUU7U*T|~fszXc)SKpwOokmYQ$#8JIQQUI=aCU%qlBveP7o}4OE@l=X&}B}p(cWaY zW1kmnrBj*gi8@Xupr&Vm1@9mR4rvq#;|rqe^eJ)%uo(085A<&yU9$6wccb-__$lF) zrtVkRnd7LGr#4p3edRz|4}=5%HwW0uRf3Xhn_kPYcLxn@hBeF@HH{^YXVf(CJ9?Om zzlsP#62$?V+I#6g7&A#-3XQI;e(;5faPab|)QV@{$+ATw`tKdO{IF?Ms8Rp#65uuf zHeJtjSE=JYij$B?L}JAHT=HQz7Fiii7E}c)A3>cS=pAe)>8=$As@eLA%jkYQll8%; zG)ao1k4tFqwk-@dbt~--=x~+)Lmn@@-1N?=@)@v&Z%Drk{^_Hw@(YNSr_i3Tzbg59 z_tm8Pl%*Gbbuk|MB+spZ(^PXh;)rnBLj*gde-ydG3@nU$f(+>eVvxL<159TLS=h ziG0NX`dC%5BGZiJ(wDBS_pLPE%jIujW+_UPkVf3dSUvGn;lere{q9W+z+<%klt4q_ z@CM7ZDQ|Rm4>wQCa>S)_LSw?PmnZ26SiZDkbH5l5iZn}n8fN+(5KTTjo*&w6K~Hcr zn*B@9>Qohz0NZ*rUJ};bZ89a;dg*rRJRXXvfdGMD5K` z;s<)(P(jNu_zd zYo93q4y%J2jmN%HUwsU|m=`vN9{1p7;GNHlMyq~(q&YFYaPs%)TyhIvMgW9Q#F;0x12#h z+DZQ4A)D$aF`GRw9@8=4>swA<97#t)Yuze;MVks3-7)=P$pl3#hF?!0YYfkLaWAOk ztQwFHR_DAsy+(M$Y-LY?bjtntj)7A9jh?EwsJmDD>O_ZlKtmYM8ZhB70e_cA1hKT( zWQFn4i>dDbhhh35abiU$ZDBT{O?q}G;ELU@$sp;JQupI-jWS5{`tJ1&K)e4OEKFiD z5&vVyl1neIJ|NMU$PUp=VTg5nM{$*4(rK6Xv^A&hY0Nuz-MgYRuj3p?4WE*{zts~l!1(QIz;<-dUNyU*n4Q*U^r^;Ge|I0* z1A*W{-YY0sM&OE|0Q98l_*#nkUUK&;&@CUh;j%<@s8Z#8(Q)Ty z60osr9>IL9@uOZK)QcO>nnQG9Z3rDdSv{pS>I2a}?SdCM>TsnKQAU-wsM!oqI9Dez z#MX({kx*!qq};Q4b`E4;6yB|PC$g3>ZM?SIU?``_`(A*94J>c1`~k0$^~F>ISqCaL>(&pdc1}s!AKk7iO<>Yj$ayPczoWJ%?3LD*iT6RLtaMP97-$p)!;`_bb)WYfoOb)gq?Cd zm!aS8j-K*1wYzw{wkUkTE$s2;CmA*7q$2jxAxl@(N)$Gm7S=D5Z)R^W$}rqPf9NOm zcQM>E8dKzJG}<6}7w3G41*B) zVb+mB@9ojr+lWo@^M%ysJ!^}Pa~E}>6i+h_9f(gxC!p z@6pf;l&`j&4Jdefdhwl(|JL&^QFx`W8;3nqMx?yIw7wC|q^09sR8;IU@_EEw$Ov*~ zrRf}WK<5<50)9#Le1@L59B{u9!undS(iLjyWj|ru+J+6Wnfl*k{82sMOY(1{eMg zV-@?tu4-nWfnzH@XVezVOta<>H-foQmr;4CC@(Aav)QFL_Gb?Nj z{I6B*-RsuK#shNZU2-`D!zYA$sn&(-jbuirhqgCdjh~*##SY&C};=UL! z2-=Tx*;GTz@NL;ea7Ig8@k@U@lRrY(O$HOq?+A+_)3$*Nz-w6yYKnWh;F%;86UtRz z4*j)jGW#%`v11g2lHdLknh&(vLut(YuNZ>8BYFypvv+NBsh6YDCh6e=rC+AU~6Y z>$!Yt9t4rsq8kwy$Ma?KK>p`oeEl4@rV+_0Uk6 zOY>ZR!5W-JWtwIei58hiLa~#h(0_5xno!E$?V1$?l}w`&H{z(GzL)4|`73#IcGWz& z$29qgQmt)?EV-kbGm#!~1>JN|0sL$W5%_1qwF{8JTp~dT6Jt@zo1#YJ8<4h&=x_}Y-+d!tLE>~t5=wD7&g|= zs0GK3mQl)CvVC4IoX29L_Xe4bfLb)lMSJ24O7QUad6O=af=Te@Sox*RTA~9GH%DJL zbpoc{^ymm6CCeKw%PS>O?05`z6;Qi`Q)egxEQh4;J^hLsCnp@n?#6wlM=@V&I!4V)B2oQ3EJ#l?^%Me-Q$h^bI< zkNDUb8JNg&i;wWYsLBfBC{7;ID2kI1lSt_pmWIINaWDoX zn9^#HI`)ZdgBFjL3Eg96WB|Po=Uf`D3Pg89Ap97=MU}1G*pre2xwB=o&mN~~nS=x5 zk&6_VR}&}bo-k(&YkUkXTFOz)VnNEYb#3omWvDE5olJ-@tK(m?6R@+JSSHAVQY^ex zV@$^e+gLCB+$l`G@5Dc@MThYEX)yLt)C3pLI@VEh9%k{g!p#ib#tnZUq< zvc~=yTxW8JWcizpx{3g+4c*Ra#+*aCha9p5A`httdP^sb9?Qf(oA2cyq=wYzEIx zJ_x&hW(s!{rD$AT_(j(Dq1<4v0|ec-TECE_5RQ>KuqvP(({UW6L|FsrlrGkZbSrl1 z!0@L>U6&I-B96JQyq4!xLBjM+1M`NPudh<%MC$(cg(@3nW9ToglmqPq3+u5BeOE*$ z_b~o)jP3v4V*HQGR{rrXdiieAh&hqdbc$`8qJuTX0aKP`zvd( zL@cC?8VqvV2b|0<(L51l_pdhrvtZ}x2!DG$M_ZKXt3Vs1IIsDu7hHV-{Y1F-m-v?n z0QL-+74M|_T2R(%;gJI8VmcywgiHNh`qUIbAu{sB zjly(-;;|IckK(Z$Rgk#w+Z-6e?EN{oOkaHug%chJsL}yD9u{^7E;A4oA*B(#->G+ zf8mpVjG=L0ws*?gipESOUVNw;-TS!Jau!KYeH*! z*6`H?Gvy$~UGy~CBuh%`1JeR8mYx2<4GWX1Jt4`9U-Isbp0zD@*>z+h*IwP*00{4M zifQWP+$y|&`V{E;_<8&1C?!)vqbTOwacZG~7ZQ%(AFILIRCnx{yxbJHzx?g12_4Q>pAY2AMHv)h2w^HHQ3m?8_^nTL z*88o({`eIg?-k!djb2o$B+u^(s_ ze@Ne1Bupj1fmQ{7Vf}Z)`i0%V9JR*K!eewAWof5-*_aA8Bimsq&?fEmx!s7I_U25Hwb)u7lEkJ+bzw^^k;U9@>Y<=JNJe!C8hb!_Jv;#PJ zQJYyybtHOwd^r(HHas}jAO0Ltts3Ym(7fMfmj5h>Z^sAM5SY7meW#XJL8~Y`Vjkc> z8#;czo%f8TRfK*<#N0Xv-(K*$3qiPx!w*g+j%|&&c@~TvPO|zj*0|l!&`N> z(I!pIuq8jcB`p^5Dt;_67&W^lF1a?_ftrnmrRLj(={Mr3opzbTSf+PIDuk5<9%^oPB($)bmmV~F8$G+~=~>gJrbDalXY zSr`R2TAI% z`0e3hWM`^Pa4%@!5@;YtM{gbMm(opTSrNl<-MYN=il8!#0?%L2L62d6!$VFb>^N7(6#D{cx%G%mfGE%Gm@pOp84kDna17kSYekC^4)*J|3 zg z#@jfd`L8;r43U%G;M6Yze*R@;xOn!FJ|nypQC_F^3E%01unCd)`u;>h^|sv$3=5*b zS=T;&ukO4}Hu`(qmuuQ5NMnE0eeUV<>Y<1erVI(eDlh;#W>g&()jOsBdslC2C-U(o zK{$6s4)09r!p?u#fFkclNAqjxv9+sWve3A0AGi7ep_!@)ST2g|_SCY?bd^N0bKz~7 z&UA%u{e+PA3+#cii~V9oW(EttX;dI%MMxVm_^csXTt2QCmTBcb4ISw!b`u7>Dj`A`Hfv)pDmhp44D(h(!g)(P&kik|HQMNMdN`M=P~9@it6b zeYsL%-{W5DbG$h>xX*1#51SWgDQkc{Cv^M=dSvoqSIqz!y{}IzuZ{caEj7GYqlS&D zTBPa_gp_fxJv@U^si43JEn%fNot`+Z2Xk{;QY^#0($3d7Hu^g6Euu9me#Fg~Qcj*7 z9q()@o8Y&JaCwFm6|~}ySz?QM3{M_c&q_|4!rBLMWdk4Dkq|{p`kG(Uz@9zwRdadeTlr zG*RtnAO|^ImrgmK+vRrO%ohf4l1t~hm*tG^X>i8B${yLG5JH$FejpmkSz^$9wkaZ( z$yxX+=Kj8OFOivuf#rB*Z$g1zHqst#>1#=E)|(hD$4Tt8CKM{Ht%wOqfFFJGeE13P zIlI-`jKQveh}*93u(4Df&N6>zDz+Vx{d%3Ftd4f7fZNi#f$kA7jbsL4^xeJyl+ z`G0XJ{wI~9+w*P>AT4h}D7bh_#;qw*LS9;r8V2 zNuK_k-q!tp{38FKe+}{YFxIKro9x&2!Y?oRmV}Ske?5=8UjKF|%u3q)$MVAsHG9y5 zo*%K`Us|Cra>-H)W&5MKKi2P5T%Z07x%PU1yjJM59Tx8+{zHGnk)&Vi40_R=W}8|w zawJ+(^UYVa=uuAR9vw~(@x`zB;{NI$`;?pE6wa$p8thet5h(Hpb@c3Czlp16e$&&g zS%OXMjIP;snVJ&)i+y<&yK}uGjU}h~=9@Q)VL*k`Kcm-hgQBne1=;Dp`G6hD`vK+@ z*>b!<(ICa9p<$$so_skqXzmHr_%4i-} zU}fj<77z&y_mGR_kNG#at0Xo`AWWEclQCsha zv)@Y%ikGecSyjOOqtPv(u|EJHuOPVynx2*%d7%)p2Z<%S+Fj;e3`0fZvL>fV5`a_<)|H_ycG(|91qCsg>sVt`Rb2ZSsm%K)Kl z6cyOy4??!E4OMsaftv8XWfY|`d3$K-X`0xnhLKHV*-<1vDqa$Qa3;jcc%t!imOO~ z91xh(_Ibtqj#~1D-*r;y+NkRd)J z>8|_gH7}ZftOg%Bv>MdzfT2SBaZmSd0s8bAfOI5hdYw52;g-nms8Bqwr2ZQCqX6B6 zS8AZ1UI6C}%T(QEJ_?A{`i*mXH*NuYD2`r{-fI4+5tBa;A2a>GNE0L23`x&) z|4rR^HD&H!K)3rAEJdCpE#0DV58@XobkI1+@(jq)Sf8e6rNUo`ec>?(KH0{X7dPdP z61gSWa;m^x0||U z1|}}yU;ybLc1dNRf@k<+ZYLD*{{#ZM9}$J_aFmmLFsSd5XN#&pfAY!b*Hz$1{bR4X zz8qZ%ZB}8G(M~!cjB@C913Hpd*akiT191>Gpc42GR<6RVUOWZH9996k=?8S2Q%iSt z_1~>rFVt-UUPx&JH>U0}j-qYerOyEu=mFqdxpfUpVk>}7QbCtqF8+f%EkB@o`^Vnp$A6)2)ZCRH5*qg=bQ)j= zG9HXJnkT@FnU(3U-E>O+rk`(+VQCjPoJa2-%8W5y7l$t_BNp@eZ9ytL`WVuGpax(R z-XOxVPUHYgDFE!C{3Sr*1i%lF#}>e+0FT)x^)Ycul+98zu%`gXaYsqA7LBc&0xU&A zf{_cL23l5{$>?6jSgfjZclMr~M_3XuRM5dSFoPBrQm1;cBMb3@y|J`r7MT6(qqW3Q zcfIy9%yn=Zl3rW2_Rn0VPl1G+*kZQWq`||E{lvSF|308d+Mh-)b?63tqA@ATQsWnH zCkwba8T|6lp-j^{CTAtPC-cNz_{Vl>t@||>H)ZJ`V?ZM(Ltwpq*M0`b??-2|x+WD% zV-xPyS`(W}hkuJd0+)+BZvE17HMRH8+IP1=kHF>xH0U)q!krcaupz8?zFQ_&k+E%o z*bpXGfm{QgUj{3iWGL6Xbyc&dxAZkAZ4 z@;c~LvyoA=+3adrzn=Zi(Pjho_!DtX35Hh>T$SA(MQqhLhE-^xKPi96$i>zbP@@Fr{}Y3~*Z0 z88o@OQXx1oUh{PU!0_NlRUniVHNn-uMd7Bd;22NWF#t`Jiq-Th zE685lY_RrS5>k*LfI1Ml`-<2LyFQW^2?a!)bi@SdIO^8l0MRuTAhBl`E7D2T7?%(I zL?9WG+;vR28|pfc`89y&LH08v`2BuU;zK(r_91$#wfRt>k@6BAllH6YC^V({u&NCa z%?-V2eJK5ZSOB(~S>YtvrW*KGDjsifg)Je91K&^ykEE72fnXO9qoz`QZg}NeeE*;Id7tCy00*srLX0 zq?v#-)8II6nm+0F6J1hll!rDa9)}3Gf!wJOah+9HfP&R3z@hlsBM_l|BU`{^BZX

re;%}CYsVSxNG1&Qf-7y2u?>!e|h+hbpIVY{_&FKK;C1vp-zcjx2Uqwm{mA*xE# zMmA_F67mdQU(0h#lwG+qM)?VUTN$sObf3oKFtA9O znL>jgQ5FQTp9?ZXS0wyZV)m8n1?M(MZ+T}j+gy-I=h3K`=+I)vX%_Jh#uXgX_cdJ+ zxqr)JTu&pXK~Z$ov>>#y^rM0$dw<0tNk1$5v-V!jq&zJ5#~;g0(cFD&MJs{c3q6Y$ zY^v+&vLK462^g6B)c%Kt!B*nJv>Nl|z?X!K>CW^~lxf8A z(>!WSG!!^G%TUa_)!jYG=UXDxZCAxbzF zWn;Aw-J4RLVihcUdvQ#8JE0D<43^qcVtXOmMUesbJ-(y{dFFfB{R71P5ADEC8lLy~ z!P!(BLV1u(xuOiN&Fj-r3z^z; zRiR3&Yfp*@k1fVy{AYacdQsX1$xIdrC5{DY7`P~(uikoYv8S6!@8nN5G ztX+-SWpsQ?Au>?2?;90sVg+v4Jx7!!)g3S_iNehOCA6%| z)5nf|F0(t-4m0+2^u~md6kLr^Mw=z_qO2FxVg!YiX|x9d43ZD&SyL#_s2NhdDMk$}k?FWpAHrwds;=T(5zo^_nw)%HZ~geupd*55g}SX1Pr@#*-IOfG;80{o=lKM2$mo>b{z45gujm=`CG2H>Eq$r0rh?N*Pa|Iu zC0*rd7}gvTZsvsqE{YE8n|oMF2%DZhF}_nd~LBJ&T( zA@477vXz93NdEat3unrb8ZE=0L32{=mEWmkN@DphhisQ%Vs9B^W8(q9b<;UccM*YE z{fsoDOWgD^hm!DL9!>)2pkM`xq~L0&G-Yo~@9g_Cdv3*>rUe(v zRJXuvbX{;1qP{00a2}iqtgmXHngqkvUK&JC_fu$EzB-(u2m<4>U6rFc%M3JzD8jVp zcIXP?ybN#NnHIe?ODSOR))WpwHhE4ksi4V~3b;(_l1y6c_x2$E!UTl#WY(4(2QQo?BKhr?9nE zg;o~klzI2<#rv@Y#uDDWVi8RxSTs^*4c}ti)OLz<^8fqxL(ooKJX56&XTpYmn-7 zJ*oBcdc|N}Nr}q0YBj%cV|Bk{NZ*-{Xxfm(`d+id}@>55)11^_vlZ7Flt}EI6?myGhjnm zsr8*KhW)Wpld%n)Hc%-3Nwknobe0@f8IhskXaoHUCWlE%!!G{i-X8$PB|KV`{(wzY zL(m~)e7;#Vl7tScA?pO|Z1k7mXDW zQUNu~HNbQC?q7aQBY3K|H10jB8KO?7E{#Fnfx%{6>XCKBeO9K&2M(+%QbbYw1cESe zse5j|&rM=@)RwfolYR+3vt-5v-13LntlW%|J#}<%u-ZzC8?2s>OUh~`{M-_A@~!@c zZzBb@FpOmm9Xo_r_EjmZ%@oIniI#Hwfo6Uc9!Tg8ySQU~d^p`=^}@-g^%^*|E)B&_ z`y0fk5-khr13}j*lOJA1&~++(6k&^|xrr7IdH4?6VRHe(Cx$UPpOj})?Bq@Q{n^XX z>JT<}u1_%b)5p`HgEiFqNl-GaKYts?IK|Nr@}J91MSWhyb9(UQ+8xd9nzLLjp6^X( zr4gS&rJ&MRsH?*vO>s!It$h~T3g3QXJk((_$KXOf(0tVF+$vA^?6! zoU2MddWS{e2Ma*OO)(6Ha5gdeu1W|Llj;%P!(X%g-XTozU8v%p=QLc6Cs5&TP{TlK zIkmcQl`7k{ySPXdwhdEjtF1Y>Cpj~qH7Hm%g}4tHz)QRoCq^|5(6P!ZQOg`W~UisVp64@9NhN3|cxh zQ3jT;gruxD=Wo&l{mK(Th<~3vZ2h6Sk>!3M zE5ecNH`MNd>lD-YFTO$DJ912AS`@*B{Y3qGAqn*C^o#H^%0Wve#x2dldp5Sx<8>~n z-^bO28Lyn|F4~~|e@V|mB2&2Ee$mX??m}0q=^tS7#3!WtqzM`gMG~~?{ufi<7#?RA zwL7ud*j8iPwms8e;xuZk#z|w_Y;4<3V_S`lwn>{adB5+R>-?KvGtc$xJ$tXUZn+mt z>rY@vaIAI%e>j#{NU@rKJAy*R1&Sy4-`@;kP%rSf_rA`zUwLULx@|wO;)(CB2CRx? zzBiXRd60n+81opLa%*liv=dwQ_xugS$@<;XwX>?DQ$R=$OFiD@+F*Q~RRe2!dL4(d zS-l*lGD@-3BP=6GmZS4$u8%!qb?}n?QR9kqv~8AE85yb}n$=uYvgVwU*ttAB819sC zQiX5p0W$C9W=eL4Tb;8JhO1M=ybu@Gu8(ciX^j0@)m-kg%~Y>$QZGtuK@Tx2w*OU> z9Lb@Sry*&H^{HmMaH{6}*v~4C-SF5uNAqRApk>;RtAJV9Wf&>Hz227wPSzj`^oTOk z)bxl7Ynp+P6fbc9S^YP^+UJq(^xdXE_|Cgi`(B7e-^7djJ@h8UU|!-xeM~R{{{{!V z$k}tnHF%P8ksbHX43|^=r4_Xd+-WI)&4?bd6gP=mu{TFWx_$5^5D^C@s@C@Wu{HYS zL;oNvUys_UJO|OTn{J$rkZ44ZHk(eY4byN*(_QjfKKZl@%vX-7?=iCyn-t0Gdm^kU zbs;8j=scPCA`eKWyQg71ufy-3B& z7_MS7>PjZd1!@%5iwW5QDLC~G(9|ZqJ`;p}{R~C%vedG$L~bWBMi$x`v%S(5$ntjP z#&iYI5N8PSr~t$%c8R9dW(Z8;N}d12(4nZYOwyi=mK9mz&|e8Y%fN=KEiLR%u|SUg z69Opx*kXdqnl@H0GbLQrLD09DPHNtk#B@H!RKuaHln^FiX*8)SzMPXvXM#)Wi0DaV zQ&i!CEJjHa?s6%LjH9+tAX$v$xUA3}8^xX)O$$3rFr-9I(XD43D)f{o(t*eX*Uw z#;)l#tG=2{#}@=Xv!f0fJM{WFII#OrMvipw#f4_D{hkv%r!}2&( zqd1h$oDTj$Q7NM+t$98w&QCojfNSZdF9kHokD7f~=7CCMyf`dGN`%Li6JL?IFU|}S z$A$xy)hA6Fh=sJS%>3uyP7nbB571$Y-@!c^Xxo%h=}-d$%~#DDZMwQbxUtfeCFbVe@vSifafp@M{9}DD9XhC6Pu=KI?=})|2 zNx+ga?tecQw3h7|KkQA(71aQSSf|~z$E_-8Oq4zVlR4cW5)V;>LmpK&P4A_F(HOH; z&Wj098p0nLhrk4kZXTw=lToMsvacmX0Rklj$vPw&$o>Rah(fHAu5xW*I_C5XW6rKA zB$8p{mZTq8G;G0^?9rTw?pKIGVOSipsYRcdisXBX%Q|B3QCyVmfiksPl0{O(rl@~! zXuTCM`O=02#A9Miz0*<4Ru$Z2Jo^?ONI6dheIC;*w2TWG9+Cyo=GB{^;I*)b*jd>j z6+n72Yb;u~EJU9VxpTzd2Ne_)mk@$|7*^d_OPW@A^7kej^Je7ypbFFSQq}%Nes8ZO zn+iQopQV-0CWg$()YWM&%L!*38Mnh4M>E)z&8gFXU(8s{{6V!v>Jb37IoQ}x)zFXN zLEfS%L5j5OOnTQq}kJGnn(E(>urN;#qJSd`R7hX-w=$gO1Ppwc0d*zj;Ep01mwcrSHgKtYBS%vJj@=c1g z7aK+8W7u{YaR@9bbm(xm&Q$H2;>-*>JTmRhSDN-l2TJSs&h8S)2uyV< zvlGUTL}NOyA=Ib}-wp+%3o?8`kZpovxmSH1t{Y)bHJ&ccI*+KQRG^5o{=~CUmfH1C zY%OYbds!5crRDo#Rl{PyboUYWgp;Ka9x@wsgwLCC!%piT#!prf z@lY`%tKnDvh6O0ze3*}uwA_zk--cBopj?<7G-l<%=vIPCM_yAI^W+i+z?dD~xK=8W zNRND)Pb_M~i0k86D?j8PR$>I6?D2rv1q4a|U#<05>9S)p;60H+fhzy-JIGcu0=GDs zc34?S0^(Bn0eWm#=ZaY|d4wW7PLVxr29h)b&M{i71!lfFoG(ei4lJp2VwaRCRo?DVT(o1}M6J63&2IL}M$?0b>uQ+=F8_hhJ!zv!vWUwmr?Bhi5Gx+b4-v1KI|m^LE9{6kK| zh8c?q&fu`FIH@YJ<}Dak2#Gf#+yizJh^-`PtioUuqOuTlD*gm0@`L5l zOuZthc)O<7-!#k5I9LqY@dk^6ZjHak_j!Kt4b=tU7>cu`^?z9`Ot`^uTLOL|@--=A za@I8m$U6rd;#)zrJTt!O)KO~v+w&oI3!jo_crf{_JyZEZyJ z^&`=m1uzsCzYK6fprV!U-`B#2o3ItqqVpg=k9UyVmPm3xLTV;82#=9KhSTg3x~(!a zex6@;x39!i<%ktph;Z{ku-4D zhl(LTfFd7BC|-RMe(hITT(KKblIfTnlTd-FCa!jM=u} zPMIHMe3GMzO)g7Ay{2nXqLDBYRuH74V+K%^0PGUkDW*?!fq;fR{9C(hoE_YeiAY?@ zWL5!CxJIf1$g#7|2B<6TYggzcJ3+`A3H@$3ccwWCKCo;!hLuutH%Es*MW5UMl^3VHX3%~bJSll^k;YOQd@ne%*hTYZ!o_8 z8f0Zv2+9H9(I1hFjva@|q;XtaK zQ!4G^OiEHz{K}5kPx;P}Tt5{iLbcg(bPB=1Sunr8hp~O_%14P-tRa)JK zC@@mXW5=Q%yGS^ZU}(%F_zVvk8WR607O5@NDXetrI#y_*T+(Hv%-b)ps)vY8i{I%2 zb!yM*)5UdLDvcDAE8utT?7AY$yKExp2l>y)GnrU3I=^2f4e(KO`OMB{*NN2nAFQ9q*Qk5_LmMNoc8Y4%~v5oQoYX3P6oL_OHnhuXrC+g~X=a7gD}{eX$^ zInfh?3Qk~BpzUbrZOb6H3D1v1qLxRzWR&KwAdSH4IZIe9s=nMO?KwaTgbPZ+`t=Nr z9J{A|j8RyYD1n0~b+LqFLK6o~5lw=rqkr!%&MzmF%Rjb{Byb#2nXQ>#!RnACPyqUz zYPLe&K(q37L5#;H8I;oY7_LI$o~&zw2sm~sq1sErgrz&};V zF+#LbP{NK=jw*6EiINe1?YT^2Bv20Z2bo;tJaNfS6Mq=rIG@b}2+OJDN)V^gK7<$4KP1ILvT;14Vj3eh8PG6o(p1s}Ntg#j#{yBlOaN*128r}#3JuVxQNSXN;PQS;38J060h3F2ZOyDTrRRL_Fj zZ3bkQA0sl$m~8)Lg#(M`P&-jTGfjYv@JV||@7;8}b+!NfEJj{0RHP&k;#0zeMXvbT zW`>KRJ01^wtkBzPb^J68;>of9mjeB#()fsJrEP>KS%e)|u3#Mr8b}1JU<)|frtZG5 zU^9x@ITLm|NtQE_P+vEFy59K;;b~-GUc?G_&$!hD_x&S>3QU@jq{36vxCAx~k*R-S zM z(3c-R+xL)kdDrJ5D(sl;)Z8dAjJr)(~V^Uq@s zvNkguYj~`4_mteuA4~`L^29nV7zZx%vO25{Rb z`JP#Tlgg?OhJ2GIA`R{;++`SFS3&wDUeb+xfex#4sjIJ+_@Av^c)lCAxJc1IC7UNq zys?}0D)9;h2TrePc<76{wWgh5*lrx4<84GRYY3ym6i1AI6pA|=3ma%ptG0@DDMT7Anc>FFWGSN%cOeJwj z7#oT|0SV$5Yay(oY1exE;d1HV(V4~5U7dhSyj`SzDXU^Z;5KDaefF=2{O?$28Xm%g zFI=udgxr!>2WEzy5vwDJGX6**an1wpKdRYxy{1y8$uW z$gkD}b)*kkXgag1g@mq%{XJEyc32rMo9l&$;{Vo#i(GW9bYE#qhHuGqGANLtD`vf$ z=e{RDH414c86&uuRmQ+k;U=nIiM%`7jz5HgjF(4`<`YAPwImZ|5xcZbkr&zk6c$3o zzt}j*aOeucR`gqwcYKjrG3BIcy`+kU%SYI-|Gi6+36L)4g}$Ttgaino>E3c@Y(l~& zS_an=ykZtjC?{MM7vb=YZ{EE|fjl-?tWtEpRJ2?MA@?GNS48`;E3Des5!+WUQyev+ z!?v(P+P>iDGLcK0E54XcrlI*777em*T8xxvQyAGt2ch^NYk1X{d5sFLDdTET-CB+` ztfX`CrLgf`22%k1WN!RQA#hwDUniGHy%*&sZhBKsQa6md=e$k#wvB-nt>{NqM#x+Q zFR|`7QwRV1EG_Rpr)8b+aK@{&1rom@F5y2p*^VNJ6o%Q*MriGgXCN|~kHAyKjCB}O zLTjYJ{KPU8k#tshXwv7)jpqnC;X>Fjw&Tw9WIB<|1jyYOk2ux)A|e6r?S4Nd+xKT& za4^OQg*ru(Q8TG&0vA<(OOx8KQ>Pmo!wR;o-k3seLALr`;-pFx*++H<={##kaS++P zmskAu75)r{YI}dH5QbrBe(G4NcCf&*#4yLWaNwgkz`c%MNwF)-^{%+=c-hCZ-IP8q zABK1wd{t8weGl@5k%}BIF30ffgSqk?6>^DJhSs1eqV=IOh_#_vX=7A_9rF8&03mze z`9h@ke8tnnpT>$qOtw!|J#5EWIU9oF75*IHt~Jbah_{6d>pU%%i7rL%G^891>|5(} z2z-@iSU-8+Mp_O=PQjN6!s;%8tp5+;%gFkc=O+P<&|S#ukXr2mD{%tXLV-XIBd?M zv#HtTOgNz5o>hmb+{1n%k`!U7V`vjW3I+q$f(o%zMqQR_n)re^eh}~s(=)-ZkES(E zDu4UCS_fh+`$tts;wzhQ1GRCr?awd(3T_3CKI3CPf%;mdt)Ztu2%}qkhi+EqoGW8U zYq+rckWx4$OOvU@%SUJ0n$RcP$P;mJcNl&z>?(#famaAWo$3Yw!1Q zBZB_X8%(e6nSex!yZS|##hFyS{ZV?3F$wX9VFL5|do_nfwe6Z*2+1*zg`XEdZM7@o)2L~u&g`B-{<@3szS5W4D1nMx4T0m4Piy>*~bhZiIYB!BS`9PW}Gjs;hz(IgYk!9}&DXDlp z>IM(#ICE{|gPnEQqK8w>VGuC8!{Cb+_B@@##!=Q8;TFaga;^f0M+Xv`T<8R(@Lx5h z#G}&d8vF2|hU@}(1 zDi`|%KP<9uTSxgQe}9U-d71UzgtB?^tn~GJmVzsvAQ9ffysI!ZsPU*LqnFV#_5EXo z&K+Oh{IjgWykr*2Sy}ot;N6i$oim^1@pHD2#{<)okp@_qqtfojFxU1`_O{Zb8yUb@ zJ+4w6)?N#thX&k^eRSlYX!uVIVZFgFTof5R-roxA7toj=I>3+{loZdDT%j}2qu$O_ z5hqX;4HIAa8>BERrD|YDTVOYjC&f{BUF3--0+wQ#rM0EHPz+89r&R1;tPd@L(HSl} z_%%^YCC^A_A`>@y&fti;Zxo~ka|!UbLh?4=dxlO3V?MN7h>>fBBv1q!YH8pLZl%6V zA?HY)hIb`{4#nhV{f50xh>X1rvs1OdLkLrd59JMZKzCw`s>lWMllV_R`DMqRfDrK&6#gKs#`rZM z5UM__*+IVyUnL6ZYzR4{PKr~!2>mX=pcgf0xY8#Uq1k(8h@ejYa=k@mxp9F_t{5FqB zV*H}-28pKf(LUF)-GE`)k7Datt(@Ghc`XwALv6&?zM|7Y>3BrYz#8KB^7PxS?>%YT zy>~DGYD_%&-~kpERiY^V8h+-4TRe8!Feg=>3a=;1gU4(?V@jsJ5*$uCSV0@Zyjc9V zxQ2r<26W*i$i4xaiA`ILiI(c-6{@c-XJQbcG><+)3N>gl;@ES-K$K|n32o(M3su|I zIBYRNAsOLf_vV%rUAof(9XhiH3j?|b!}A1fToUb0ysqVn&9$O6U?t^K#~4y9*%~JQ zfAf+HD(L$w4}6I%gC|$m^)`R2N=9?rFzTo!dNvM0oxBTHZ5*+F9hr5epNX87C3csc zSH#2>J)zB=13kbCxZ*ZMd!CzVNB;%8q(!XCl?uHj*Jx$Z;qxy=%bk%1C~(mcq-#>D zwCsO|52(m$TUI;$JcV@&qZ)Iiz2Gezbt80-*ylk5Xs6N4-UNi|9I+Ow<~(5yr5_s~ zTlQP&VBU!xEhipH1F{48RZO9$jq~nZEN~$;&7tIvrlu_8P5icC7lio)<=*qguLP3` z3crctsikR}JRT=1@Ym@L?j$FEPig)zY^VwMI0aq4JfVRu#I2mkEx6Z zj2Yh~8g-4E~j$P-8e!BlaQbuz8d z&Zx+h^K#W$Jeww`SA_oSNMINw6$Cv*@YVA zR87MZgxO)BXnFLT`I2icW<&8%)4221_E97LmZ%-^{uv^{82ym=0?B${x8xFYG{s;q8wc^E;7>}7S7NUNpUC&6Xl zMGvCgP7F-&FQDglXFR0XZ7T?GpZ@GnwBkVC(*(6@kptm+yekBXXNy6Rn?YDqNXmx9 z6%VTOqJRm1w-wK#sa z$xyiwt_z9za?E6hSE)L)ivrDOU+mDvusn79rgA8Y10maZ&yb{C#vKoRA4*+hS7smW zCjAwt&gOTj^e(FEbXil=Hbx>b>uxSsp|46NgAKK-kFcL@UgsG)TUQ)}`9frZ+dCm- z{$E|90qfr4x}9*6;ThXE7Z)HBSN2)7Xdi84In7?Z{j(nwUfM)Dlf!6Xg$Trc?@ov# zPfa?%%21`^yw>m1?$Ufk$Kc%D1u|;=QSP3575gR_aD-dy8Jg4BT6(`OasU{0x6`Fr zZhrsAAw3gI=J{xB6l!cgG^h|8M~Oda!wKyJ<+{e&kVpTOK%QU4(K|^+H11*n7My}T z^7&uy+eX)<8qWnt~x z^sSeQx175;i`gYb@a`Y~HwY!Q)8HebX=3;A`Z|G{NNJ)`L z%B^&tlNevr2mvA%7haLwE=n3}dOSn3TiW*VoLC{mjHl&NTPgNBD1l1xsH-S|A-0t9 z8A`>QI$dcbpBpBsBH8^%6S)QRoc+(_Yz5q#q>HHSm^taLh1-Gs^@sXOy`H-iy0@!z zw$!^Oj13&QI<<$aZnsc!1ze5{!|vcY#)d$I^Z6EQj9nn% z9@h`yg~~)v7Nbx9?$UhDu~9Ot@$v>@LaQtsILBM1t>G@@g9eWE1*XzVifXFwJOVCe ziG$x`oKX=Bj8uRkkrnhoWdo}O1)@D#nEZ^P!3_>LQxZBAXsl$W4Aa?51z}jSdu3j7MfNxb zM~Qq%6P~xR>CT;`88X2J2{CK!@WPYPQxHtcb;RcP>X|_yw@A&xSKTAXk?#7f-V{XC z*pawbP1r%Oi@jPt?@Wr8<{`*^g0xl+dU$%ub&M1EI%t+}-FEOa(9 z5ygDy+ClR*!4lH+Lb?@8Sa)CY+lb|1gI^?pQqcjSwXN0K%ytqM%Kay3`orVUdTVv@ zOK^3F3nL_^+L*2l3$xP-5W|P z&ZZ%Y>3fD2YRTs{M#e0*Skk9$h!}?wP}_5M62Y4Jt6d#N1eQ(3$JF^2AC>TQibzKA zm$E0;qa~7H6F3-TFf7%nK(5s&JBe<>`;vc|-Bv#4pVsV(yg8O1) z3j=c*pn2GpPBQHae=@S7RHw$TO*e$(L?{3zVcE@T$oP|XlYR)@QHp@ageK(_$EB5b z1t)2VDZNFE@Mwo%S+~C!jC#KA$;2Jx`WrdD>^8D@Tos4V(ZvQ~SpQmXcXW+!-!uDm zE#iJl@;kV7@80JQ5@3|$o9^a<8YznpT|r9Sv%XYA$K}=^4{uVOF_dktHT^H5lZkH( z0#)AwJhwhN2S&$?3T~*1`vz}2apNh~yBU7upL@gRyY=7)v|qe7mp7e&B^ zB%~R+N;___$Mfux8p7n3TBXRf*E!8Z(H7R+rq3A>ejAk@2+N&GC}s`lz&k=E#!HBL z;1ZTId7q`$9juwY6k1IPK+B0iS5v1XDYC)X(~&Dn=(6#Ac7YuQ#n^!p&sl4V4?3MEa**iY zzXv4g4ki+RfX@&*{Ldie6^W8Lf{cP&-)*)gKgPZo-VB=>_T`XCZZUN=QtxrwEQNNfPTnSHzs5P zh7p|y0A+PIO6T_mu22VAe{Ou35Z6ofW-1C>)e-kMlUShM=cvA150Jq<56Lg6hX={I zt9PM9LcDh`2ce5p?%Z`sMHJc1Qc?d@5tRYB4GF}H1Qc2rtmIc`Ub!NyAWQj6I0kQQOA{=y9S+HygE z+v5y(Fka!YF7o&&K%Qm#+lI^Jt(_Xz{a6$Zymm&8^BbevA>!AYf`iw4Uijm$Rv%i` z@us_D8CL$5G{7Q?IAHCPEE#O;fLp(2&fu~6fWF~cV9%2`9i7KkkCjK5{yuGcf5Z>$ zklyGSmgHgWNC4W07Fr4Awq&8(?8ts}J2n8JBDeO>h<(<`mdIIjTG3MpNL%zEn(H;Y zA2c_Tc)0z&wZ3U@^K>~)S#8QbAn&1={1t2wPo<~Tj=pXD6Wnhs>%W~KjtHG%ayx-u zg`$5r@_1$-Fsl;*HT$(chB>>H#(y=sDD-`#?Q^}5hQK6Y=O-T}Vw0EgEt!FyIZ%B! zY&uDtZ}aZfxbb1b-bJ$VwZ?|8_vP9$#Jx{HE?H{`-sn@OarZXZ5@VJ6?}gN z{EE{0E_Qaj{#-z4*;#2S9Tt@nwENcDZ=FG^jO7hRTqQ{psXI1s<-T`*U)TTHaoN^b z=^nlG>k&Wx&0wucVOiG-p&AO3{(`K39VN7Tv?y(fjCteAhzrc5ZQABS1X0&#zoMp? zrog5RdykF@fAwgV4m?7X4=3f#^c~`1u1E|5%tT-ePn;&T?ZmI2?*mz*)v1@x@!LxklGl@Woa78y-Xc$h_hbBx0WY| zB9%qJX6cvlmaKjb5LtOtybl+9HO*W9v-wkWvPk<06i?jUqz|H^ee1T1q-)ETxPVubx9q z-(6v*RMBd~qPgXu!x_5Si|})Ao9aGPO$u&2n?n&3tt5WKCuCu1>zC5-9NjFyUvrW( zfrG_Ys?+~Uy?R4{T-)3^`v{lcg6zeot&T${6HA0bJdHM>pufUkMr=pC4=Vi(rZ?~m z1f^ZVnk~;qwFV5f+9U*agk9g{;OQhJaCx8Nl9cs6+4Wj>?z<{5BAAk3@ijaN_-QaI z%K-r_?hj}a0K8#I+XzzgXwL-8KI^dWZLTcGefh$I0k9vO;N@|Nyi8$bvOW{m^{3II zU<{q7VHsT#rC*{jN3nDM(8#T;sN7}*Mn99sy*m-FK6`7UF#hab@?GiNUP^V_>G7k_ zaZTKeN(5o+HlKf{f|lo%q=_q{qjMP-n*V2>l>UqI)}!Y>=Uqa=_VxFqMH;H~f=)&# zO`UqM)dtgDz^uqlL$A2Asj`HN5k5J-d7Hd_mumqzOsmSLET?5G`cx1+SXpmP@DG}2 zw3jqd1DZS4ffAS))}c%(S01|%oCc{vj`{S~XlgwtxYv(B>u_a+E{XpbC|W-tCx8_M z*c|q!P(*o?=uPJrIV+TqpaeqV)j=u^C!yp50i08SvryD)SJ!%J?LjLplm$4#I*IZc z2Ol0P`v3zK?AUwBg{ZLAU4?Njm)JfIqW?vDqccUbMC43(pi5nZrh1~+pA-mr5BP$! z9;PpFEd!9B%f9|8^XNJt9oGL>!=@95J{D@TCQflH`2ZLIoI~UO4pKg4SeFJ)mAU*UVU_(A@&q*PA&jD zk=-1>n_XWo*^$e6iW?D2KcVlr(ayRU-L!Dmri2>l?|=}k2?Oo~xRW5|a>NZAEn8~} zI2E>GtC7h zZ#c%t>k620SNVsxt4abkvnCA&1 zAaz5_>?gqNJeUpAT)lVr=9!J2fqI{`_>M+{rb)ZQq3c3!!~W^EcBCK_4^~ieEgd{C z0F(=5^ucgnxsYb{i5a_(zq9%7q5KDiSLk(5fRAo#3AQ2V{;y!MU%~Gsl}Gy|(Ac%d zjODjiT0HYkxzGuZ0aKcM=UqzeMK9jVSJfIU`HIUCg;F?!ey~6cJJhc5M1Yi-rZlGj zKyD#_{N_xr_fLTmC^`NnTai?*=|*EfXek;;ADw!3$ysAPz#J>yEJ!a+A1z2{H(>eG z*)Y$p&IB3(ZM4^7Ov1#mTZLY8U029vzC}Z2Iq1pO((zekT|xyrjkovjMz7yH`vtGH zEtD~K@;o{_*~v5EnD2>Cd6V9lB-3AONSFW*k7Klyldoa=wR9ieJ4fe=n^Du!pWKDK zE5at*{^&>iNFo+_z`|4nzVkXIf)chezWws9DYy9jP4RaBOeaYQRy+;hI~|P-43e55 z-KkrL^?1AhWGoyzj5yV~WA40YnZ+m5x2~$93DE$~8a$OFyOUl9s_ZXM9DAqJ3-2Mt zc}ja^N1a-p%XUn@atpqT8>M zqFv#eK#bQxvDd-BFV7#WS2DSove6S0g^>TR7eLJT5Ak0e5kbz2=ZFeLIzhl^8GPW7 zq|%29`Ya)ocT{`RxX!)*%Gxjb88VpUhlvi0os<)je@uF*H(6d?q6d zOK6TgN{qu8M<}aQv0eLtczi(i%)jA=-msiG{>5Yi)yK_ij2C58HUw>)PqRf_bxZbl zRPV2-jjb#9UE_Eq|5Y0pX2mhxlmX2B=RC+;@H>?;eihiLh+tp`sibLnj*gK$+OkW_ zi06--Q+?yeEHTHXI3S7y5F2((L=JYolq`zdwL*r6o8C@}_%b#ZuEC~qg^=_)C&~%z zdW&M};_fSfpV{}VUte1D{w$W+(}Sb1d{+m_j?u>;1h-?!B$+!oPAib{u1s%>jn@~f zr1XL{+sx-ptM={38mp*%Qk#|hYjXk4xX&kA(rNHvLp>^P8?%L%&C{3#zD4BP73RuesuES;wnkg24vDeW1_y0SU3TF=mJ85GP!uIif0v zIh`^2K5%w>%goyZuO6PB_%Sqo?9bJbt1uK9KQXt(km)(>Y}_}{OO3&EauP@ z+sui%zG$(Ku+VilsM~ptcC_y%wh3dOwoe~3{fJP28(Oo;VujX&7Xe$zY>s(6TMdGl zcVh9g*G_+q?Q{MBC4{VobL|u^q(9~L-Hx|4mfJ7u*-av|pY$&|7WgbxaP(ONUfaRP z@Qr^l(utE!u?Se%arGOu9uGm)MngqS;0Wi%%7@Ho2qZy+7w86~0RWTreotO$pK)Un zE7FC67cTc8O++^H3g4OtRyaK~r2wbnee@+SL)^^5M~?lTSq-n-yt=ZMixkBgIT9-p zlOpBWO(PYXhf$$zA_bC7D)KW7c!5>tz!mV|#)S^Yf-r@+Nh9@*w-9$AQC+OUXxIx%wdv<@=xt zGEP7L6VE2VE8sv|7k>onBJHcbUX$nl=QQ3rB+b41$%X1S1k#jYRXGbyE_w1XRQP%~WL^T9zv3MSKo6~OLtEmlHV^LBQEkZqc^Mo{o8f1x`GG7H=C zRxz0cvs;c6pv299{+{W*ZSDNc`B`((1);B$Pd!_!rGS-`HkwVJn_rUh9uy(Uu(mzPDzkSxO_+~ zi!G!apdM>$`YgB#z->}Utx*Ird{8P{C~!!$!zH|K+q5Qu*867qD+@U}7%&;oiUjUG zqp*%#C$!2oB}-*0S*XBFcG$`GJ}XZ&Xu^snv`qQMU^3Lt&=^99_F1TZ z@n?b^&u>exu(WX$!3HTIdvtp+x(7~kZNnnukcO8YtF3s9lzw~4A^QeP{ipI#ln$7a z#`5cQ^e0LrcWKn}dHeM9NA&2Lsn&BTUdcBM@n_RRTq3zq)_bM~leGm=>Gza4Z4u+T@(h}-P9nj58f}S|9 zW3RMlBu>L8NeWzn+T-I|PwA!|XDZ%FcjL;;W_^frG#ua(34liA6W%i|M#f5g)nHP7 z6g8F;Lg|H}MM3?mL{T8Z0n@Dn$Dn`mLoAM=1&IYTT6aoxP@SUh1BVVKOn*x$=Wj;_ z0op`r>A_tnd)^do*kC^xj*^ie!3^~hYgZE@HqSZt^(E6tifZMkKAWMjk*X*_uE)4B zPag&eo}FvyJ`pXhr8T8HqMFMHyPJ?`jcpzWq=`0xJ+|ugaZ^Oopv)C{s>&_@TsG24gnL ze++>w7yVTNsI_{aXJRN{O7PV_3+Egr*4e3QuIP!>MHoh_W*O>Ds(kh~F>I#7%Rb1Fi<( z;jsf_g9$+9*otyaL>T(*&F#P?g0OzTe22UBP#Y`%c?gMY7ot23yl3fID0=N3`X}_F z)XOMa=~;Y*Jvgo(S%{w?r_VcIZQZKFNmcNah7HBj>Cg%e^*D$<^CN(uekdIT3rK8c zuN;L?{j|@iHaRj5BAqhD9oL`dd%m{2ZhQTFd`lfkFZI*8+a9Xp{n}dSrqI^*>EBNw z=rkl~{DhAKRd8BnQQwv$lm+FXM-U^1$q12>9*X-9A^}BxrFIZx$l6U9$z|kuGX&m=F4dc4+&xp)18tMV+V;2>u)sk@CLzI__W_W0n48XS`EpbK2(Wa>g>~6|Oz(7ssMiXZYvZe}{3qnBg zK-@#zi%9Bny|8Ek$FdrSpbEfXSdYomv8Q_+`?3kLBu?Lu1e)_f3xSBpI%{T9^ANGv z4L-sg+XcX@V{`X>^iTN=Tk0m)JguR3F9d*3+n)FV9)fo+K}0hPEf|F}XDMk4kr)Yx z+{BcRsz{&;fA>`2gf$b`NW?K`0*Aabf3e07hwG44{z=HifApX{ylKEu;Os12J?5GeZFz-6 zW!*r`nFNDH9-BA1-u_-+Z}|r^*q|YJCsP?r;BJ)6kten~I`PJHd7ip3A)FsI5pC&{GttHAof7&Wh9mf}|bZ@^^{W9UtYG+ZUwcrv;vEsSje0ufGFNzf= z(0?9(+_5VfYkNKynVhqaUjgHk=?e$;iS4kN>IpPN2l0(OjR3^{DSfL+OrbUnUzjI9 z!lrDc{vs2cAsyQ^CuE7?J=)wmBop=PJ5ebjw`zN@R*4cP#?fN&Sx>rpouV#3&vGe0 zU?rNAIsznBph5~P3V5a61p0DfIZr)7R1K3HY56sL|!!=%~E>e(_dbEIwN+p7+825*9cf!}vtytWcdQF9lAj#68@O2}zjGQYvsA{e~8&5O`!Dh|h_E`I=U zS7{db8H_rHeU?YGile6V{%g1JZmh`rW7tZw_L|`3a~CO31qX8yEejXS#qS_ed+705YhJ?1pKV30s)WlQc4B9M*YE zb7tm?I251qh`!Ny2ps&pHs}0RvldZ*^p|ex*BF>%9!^I30T9bSNY2ro63}cuAAtU= zv3V4}sC|T|X3r|GF!?6B$um*u-^VHpEsL89TgsBTx)#%2Mxv*Cvy@-oO$nTChw)zo z3i}JWD`PXCWMr6~GYR+n*rww~KK-81^M|+(@@3}!5q;gmcus6Nl%uF1hX^zY8&ATS{TL>*(7~kk1u0BsZhhzpxL@!8vP@Y#8sBBulk?`NRT!) z>)&Z+!k15Q9>ZYOLlx zpHwxobb#ZWTBco=qYMbnyN$NeA2vHQzxY2qon=5((blzT>5}da>F(xm=R7>ynul*>T%~~*$NiA`ZMEKHdGd8b&6bo7WjkeGw{q`KU zpm;eej?oraj1iZqj}HpvATD%dlzHd0FDLM4`Tw5sFFP$Hty0biFV-5#AoxB2bOl0~ z_zCId(B!ogZJYu@k@sKx#B52pcghCZTfu3Zlb)4=;QKT_{nnd)7$5P`6+IxqvA(;p zjo4~*`?~>O+o_DwSU&w)_^nfaoDB`V$}W)tZ)5#p3wc4YL5VuK%6t05KkyuH49Ut* z$_GXqe)%k# zBM5rc3{=J{1GB|`UkKg#{o-TGQ7FzoYo)Y~ZKm%GpJJcQzHa`lk|!(K=c$nr^1^@W zBIqN`qjs*8hP=ju&DFDj;fOCEj6C!HQehOKc=vC?Lj^!=@|Ilrq3+BYA+p8=MLW%9 zT*u=g+TOnAT9dzC;OTE+{gv($Ii~|uM zsCf1B9h6FyRf>#iKW(4EZ22z-^fuD#pQ3G^AbR4#&o^nQ-#`5Ua9?`kt?AmHO|CfZ z<&ZlldBU=U=d1cOnDzI6jYVDLBe-RHfRF5NJ;?8Uv>U-^&>HWosJw;n{_Q7za*Ab`_WR_2v7$ZAk4!M4I}^uKt=oOb;>1Cd@2ekd=woZ< z5yNxVcvJ+~L*xqtghWbSMncv*MkvX~1Uzu&IQojH=GZ?g66svSLKPMmu*IqE5?KgU z)i;)Nr$B$Mqcbl1>?5!O^E!-PdQG}=4i&Vag%R%q{{p|fBu1(yv*rrKXoUUYtv6Rt ztk`=ZBG_G*X8`m4E}2q*>Z)xf@#^!x=Q@4^9Bx$h`VD@609e%g7qg=xIh~cCsWVm#UK$)iRElbk-NG`uK^CWP#vFA2tPdO!|un zVT`VyVd40CkB18;pCFXykMG+_;fGQ#cQU*Tpyk)DTnjFCh3{Gdz+rD|Iwe1N1RDW% zV#KtCd(iH|s=;)43@bSX{N0h6s*c5W3MyqvgQzQZ@1KUOq}T9P-pVCC%&D#VguM&4 znb0(ZZ}_?+-|Wgkor;CiQ2!c`{y3w8ATt~0)(ut~Dkq2)S})|OGSf|%Hx_d4I;A_Q z)i5PR)ed-V@?}Tu@AqzJgjDWr&nw*44R_u3L#h258N;=&dDpbnkXqZmF>CW^NfQ{- z&4h^to0-aPvCM(?f+na5Tf^eUw)p*Nrc z-}59aoI(w26RWsu8_5p$_&;ZyA>~yTEBLCxx5emNicJu-e3c&Jb`2Hcx1!%Q8!=O+ z+UFBCGb-cM$k(H?yXp9GP-dKY<@N>NgC_!ufcirP7vIK68w0j-;x-nyY;F6@mBs-* zOa|Hqy3Ug%F8je51dJS`hU{As4z7!cuJ%RC9+d7R0hShv125&aldY;sZd=}DI*uj{ zB2#U~csMJ?444^+YjKoR;^MCF={!iEmh&eN#apo6<50p=U=Tq5sHOIN`C+Gu#9~nP z_6z>~B;S*HnlpINJqkq+NmvnC!-G&M+k&l|eqz2p0Tm6*t&CQO+J_RcYTG}}U#=IS zE2;pIj8bHd277f|EL1RH><-yR8kOYl0+Uwhw@E7phhkGAy$@UpuzGNCva3>}>MZnY zKjeR~x;}LT7%0V z3iFG| zt!6geBF93zWXR@N;*fEjp!sbltlYYL^BN59c|~eZmqe?Z#FLQvet%dcAx02O-J3v& zc*XXuXn0}5>qH+eNduY3FN6XU*)VCsSl`L1-52iLSSMl=bT`TM(u_gRD<3WabY_Rl zQD2D3fk+m+NricQnlauHXG{$O%IS&z#F-%tU^pEtZMF-T~v6Azc zQF42O5zHlfHKlQAM|`ZE!4N}L!r9yrTJQ%Z`!LV%m2wtlyUe(a73N~Df#karB_Uk6 z{>im=Bn3K5ANMYCgeY*^FrC5X0>bYIQO6IlIWrYl3zjG{SKFYHR7tYp+#;d;DwDvW z47-R7jdjkKkR<|MdjoaBd|>$h{5Hk*{TyS30`!R^0{1its}J};1bj&D&;`}IX=S#^ zZ_`@V!_>P`_L-3uU}ca52K$KjYf_WrOGG4Ui5i4zXv)Q@bk|s@&*N!c74`d=Lqp~_ z2#gZq8_1c4%x%kZFtv=A(IShBY9Rd?(8aq5{^k&WfuDY1A93Xf=(Tx7Gt; zC3m7l@_dL+;ir0gz)%aT^Hb2TYUqEnE#*DR2H>b(zZFOuBvzTOZEc_}A_m2R?Hc7* zSp5(oPfD9U^n0O^3V${=o5EG4WcP-Slt$rFLq+uw!IsD*NhC(jj z9&%fIHu~1@Xn2Y7Y!D0?s2~V+TSg^ru2-GHyjMgGUsU)Eee0P{>r=TQLu4usQIiZPSVP8BbOG!Va zi>28zQMRet(%7wi`UKa_%3jCO!)a?1ac?Wd7zDY1(vMN@bi@rIT-$`zkaiiJVwHyc zsyt80p(IT3j?Mr|poZUXBAVZ?1k-pkSZy>gxFu2VQx@6JayXH+^X~*}vx>fYs9m~g zflYH{gsb+$;D12F8+1jH`z%~*yPZ#P%J1UN)JUQ_S5a0AE}V12u-mZ6v)Az&)?#sC zOGv6`trZfX`GpcDIpf&|TE|72=UI98T>Io6;&CwC>ksL+Nh2}-f^T|^Z@SzaLjekT>?gGP%W5F6>c0r%qccDG`(V*%kjT4yOjQ_;& z+kq{*Gw8iH5N*2a?0xR^pOOh7&Vr8cAP72)hU`%d39bos7*oOXH>)~=zn0W3%5Y^o zhZK*(J6ZfKnX9MK-%)Qhr5?gw^{RxwUR27zzTtu-McTH}qw)Wd9pMgAoi-`2_MRk_ zHIhVT;J}2HHt@TS3PTC^+-dY_Vhi36le`m+?9H%G*d;dHs8rUcEd8doR%6zfZo!&? z3tb{$^^fM$*_q@xwnLH&7d9KrLPeU|%Z6e-UwS2Rk5l-0ff?U>)%40SRI1oe*{~RB zNp^cFgdaBPNE1eD0ytP}&TRpaN;=|c#Y@3LXO%gu`JigX zl)m9#eQ~tK0^&NhA0!vdJP{$cvK*Xb(arPd(b6le_)^=mN1Mrg$1b&$AfvgB1x<)X zLvTZpt=d4Z#-gOl8k3F2ZIOd(oG9U>I*0Ll7tJxT(oQa`?gSix*k9U@dVsCN92ZM5 zo*k*Pjv*9AEveSSv7P-CkP74aWN6?MayNoBx8;uz zOMbKK`dBl3lFKt=^|59L`ptogh64*-G)sAT60BGsGt_vqI%Lb>p?c-JhB|+`Wq@m{ ziS;@gGqE zyFh$9Wj2~PImI_2PmI3cHDvjt3HHo3-y$2$b;{-7A>7>aVUfM26V_5 z9w?(Um>s$J2a2Tr8Z+0z^CS5B)lB^tw0sONWOL|UHi1ymImU@$ezxJEIvLGqf{8Cj zupZo+1QIa_vamHxw+uUk;?egc$$re!|59(13W^In)j`pae=C!#VZuVV#=fh`Ew6kT zZIKNe%S3FZG$mN`mM6K!Zu6j?n&&q`#OmZr(z@rqANuIe)z7r2ds#~;>ub>ZaQnZ! ziDN{M_fbAduWLRN)jBOWdWA8+m$g2-CMth6bYu|TfyQXa+?^3z29vXO53ABt;)yFk zH7{)2=n;|W;isTL$Ffi8_JeCAc-@yk5JH|)oUbFAP#{=BH4nBeqQ0K7g&mL6pwTr` z66$8zUP8K{-pT~?*x`isiENbbB??R|_{Tw7X>zE>Ug+{rtM~#_n9FizvV}8xlt$mY zN&oH+G$or7m0C1nslT*Dntv(&%@Z}7Ch*VG(@M6C z!7kAhOsCyAW_cbpi8_XWlBXz(uphNLuiJN$v?iB~onT=rg%1^rK#(Mmt^T0_RXE5C zYw!S{Gv*+ErA6Rrm_3;NO+3hf^Cl!iMO(2yizE(0&T!FkIPc_vl=7hlL*vtl1v`jr zMd4+-lory2hv2DPX9@DHr&y%M!=X@^L9Z`KVr+XPnZ|o#7%!Ut)(`3T#j;B2pIal zL|ero*PlR8{4N;^kFc+=)v+%9h5Ia&L;Y(6F{BYI!x(HDFA|$Rk(fwQqF>bzj_+aj zHVhN#PT_)Yh&}|aO`Yiy713GE@6TfXPm}lp?;&ck-F7=HVyVE4^|gGRwNc8$##!=EPb%qj(>re*mfe;xrb)`ttBA5K?~OUhhi^2!e42{cfsS zy!5aBoYu2JkgTkvMwBlE1;9l4w!3EPv6E0fMvtM~LZrk#%rZlj@7x%kBGf~mm7(X; zm&Xl7vJf2945gRmL;Dg#3PVXYa?fB==n)Oe8P<1fl@b>(`HDJ(p9TU>w-gb@)?G5u zxQY0+=XsJmo_$ahbTov>B_j?8EEGaDamW+D_Y_H$Oi9RtsMXkwUa5Yt_cyg~V-c3i zG&wIf7ZUrVDTW6+x zFKTGeNZwO-^v);J@?-I!=NBe9M5HhLzH>1U$uy@Ub>osO@uv4#!M!D(0Qse1^gqSK zBh>T}u%}=KA+JS%?vWPMQPtsLbiz`LkW;ajvC(9t@W|8$)lV-gmXuo5juFd)jd|)X zLo5wBmuIHz)DY>od^qm~-KQR7>=TDjht7%QVXD`?korSI<%R@py-1UhRAbc6K&N-s zti8V+SHZll4a+jl#VxG-2r&Mp(5`qS3X*qwu`+&C;)@mi)GaPB5xB7~u&9CbnC$Bb z{pFJxym7h^*3^-V#;gawyNls}oxf&tGIVsRJ!FoaO}}!bgW9ZNk)s12TWMcCQ#89K zwCeUM>;EOJGbo!yu-G2=G8?uoN)l*kSef+22p+~-q0KDtDHH$} zc`$1-ZQZ!WUFeMWxaML9L$1I(>YbjQ zzoocmgeMlMW{c(QSX2k1Rj5VVWOp7h8DPvQD2UnM4(!k{ z@sylmogd}8zY7Y6z4PH{`wJ$$>V^X(Yd~8lRfzfq1AXobNFx| zPHGNZ87YNZOC<}ripc#E@%DAtUVVCKvX48KmW}|1xj3=OD%~V4FX+LZ_aRx5BAgaW z_wH~m{E%4?iR@s`@;rl@qJnBy@Zg6P*<#tVZBq=Q%IG#)*yS4u75u(BAZ7N6a0 z5)voGGDat^#65QsY%dutq9*l9)n`{uz2PO-cit93&dmBXM#9EIEiTe?&i?&adfZat z)a!*_Y5{(!F`D@Ehv_(W?j(&n-tQuX>AGj!RT|LSiFHPp(JXaVEFMG>qcC~RGGau7 zyJrk;I}6#KkP%mB@P?}c)6_f178qG;A4PO7sF2kYcvWycE4VyXjbLc~2xsc(B>5;J zT78V&?<+zeVIp}-HPj|&0Sz4AEVtNPhI+XiH6<)mN5u-rJ#XM3U{uTf4G&}-*Ehp@ z-+eJ|54i33MPVjXELD^izBg-unIP^GEk_l3NMjW#sUmjPy@yg^6ccBED?Ai9?A2sZ z1mp^oZ|vlti&0Qek>JqXFM>QZsSPkslu#fzSG4Z?I4F%VsAVY65#rc_MN_DN?SjaR zS!##uo$id445tAqv8Wp9ep`P4GXoaI2nu^otvS@zFMO)U5q(|3(0M!JKvdLm)EGE) zn8#t6fT49g9eMomIJ&7N#Uf1-8Ys3Nphw;F=8Bxq$QuxPlqB+NvoGcztz3+%&A-_A zD`=T|CH%a@<;#0q_zC>u zlHRZk7PpfzI!1Z;q=noZIM#tiEa}ow2E~mn%K0%jN>12wZ6u%@7*W0FmqTUT^ks-` z`#$!={Vw@2DX9M&N+quxZvsvs4pqy1wZQa|gewh{k9T@DpTE>%{)DbU{3yCh@NmzKE6>FBr~r$OpK{ zSfT%c^14Gq^N{Q#8j(T9p3}lqAS?V-#Rugt&lO3+;+w`jQw}IMyx`C|(=R`rf4n_H z1WX93HaKLo;7@;m!H!Zdr>#TbuYS#n;ah`H|3DfR*#~Xf@$iB?99!owg4Jft`BfmAgod&qJsvTcgtRssLR3zjWj%L zIi1T=ry@Tpl&0Ql!fxnaN0cAuxW;`J(qyZ*+}hjtIj2-gJF0LG_cqS(=!JMY-kzRW z>#kNB08r}<(>%Xl5+k%|1yt0~6k-$WG&K_wg>F{*%YZ*0GLdBu5E)+5K$j5d~PtQQOSU8>7 z2x!%b_fU1?WeQNEnqSVdR=dg!qS*KKskyd8&TUsXe{|IK??!F}>{CC*>E}&gvcW47 zA<5aWW|KTzZh{#{U>f~P89oUR`O9|?eg=RrD~vYMO1wshCDoM%&R!iTALZ6S<`}Hh z))AR3Fx#()NM!C&Yzboo+9`d8oq$QW14Tnahrv2sn9c%fT!xY@Y%KUphXk7L#D}nbb^~$E3QMv$0ap5k{!+2U~Y#H2$Y;^IYG%iIBs|(^s-lx#& zXpj?UDr#@wTyzcV(BNxV#+wB;74c=_&1!Exys|zc07LHrr%CbCQ_~$PRo7(lI;z;c zcLTvBRHF!UEyo&)_6Z2kTYW>K2vrl*TDkpw92I*oB4Pei8;Cv0nmSnU?GZumOvlNMLjT0`Kwz zgAvBQwWmjbGevGbRHI}*(xj9gOUOrfp@w_=VYxcsk|!qUJz*Ss8v|0kcx{sggIi?M z)y?)Ux@Vtp$k*qMb(PP5G_n^;`p;RGbgi+tsjaQafG}!<-$Ee9883g_Q?b{5v4)z( z=)ZtW^i_)NSHhyQP;9PQ+~g-x|3^e?Q=}&+|8oQX9B;14I3KPrFa5)eDiGGLYe@yg z$g1^tD+jp=*}@4dExwl>2~ zVs(u+fQ1QGXP}A1e8SL*ljgwq8>CwI&0zai+8c7>5?a>k%ik|X;XibuT?2OKa7{+2 zi&s4k-(ve<$sfgAK@2=|_l~zG9YP6s2E@YhgyWk32EB*-6fje4Nm9b^HlWNS{j-Uv z2Zz1tD)Hg6JOE{Oe#Et4cLZ8R5~b1KaF#DrNg)@qr#~Yyhm_CE_zK70El+PTSaFdV z*UQ4FA{T2YDQS|>7Bci-$qM--p)iko^s)y#1C|_qP8+o>{u7=*x+UM=APdjHV6NpP z^VkE`=Y6rsMbigQM2QQWghNg+vRujL_Cre)Y0$U*RE(rvLRgo|MDaMn&5OUspI?t1 z%;p*3hO!)%h^+bu8kxcQH(oq^F)1<60O^U7Fkq_mo-CKr81)Es7K?uf@xmTiRN9j3 z`M)@uz{ZlPiU@g9eoEWCX^bu)v&FMHsk%+FNUaakpiTFxD*d7Fx{9=fwqPN8#LTEU zI$)WXXO>efw=CKMG^79h0|~@?H$38EzJ`W5GG#$pt#DPusv7y|<&+i%*`gHmHV0A~`OK!|c=l90o^Cmv7Q?;4^#KW*oRPMy#nEQH!~3YNj9f>c(dOhd_MR6#4FZ@wX=kdSE^7@3q$#n>PfIQ z%Cvwv!*1>vO@ho(=3$9o7n7`LL+@O>CoHU!TT$=0lB1dOdE5xABGvkw zMwWPb2GcN(Em>fBwDDO69B*7prHn~O)fh5UELs{X!Gfa3RzH?_*Y%cDV>2CiBC_E* z-jVgZ6eP;f_(qISd3|Hl$>9^RaU~`7OP1TaT12+PM;kyLvFgCFxv41;Wq6clF;vQ1 zK@Ck23{D_8sGkC~w}y0T=W(O8Zu`BrGEo4A3PM{*RGG3fp>BF3F7s70e`yT4H>$Lu z(pY(yr3Lsd9n&(ouWp$b$y>2KK8Yp>P1BfXqBimjz5ac6@IN0$QN6q{`uk1YsabDY zn_tLbu#Cp{gIQ!{^)i`oGIu;Pq$ZhJW))Z27m9aLl8z_^N?{X#8)cAmF-P)2Ex7#1 zlES?bz#d(1P<}0GN~cQti${WL%|k7kXRu*aLyUipaVKekA)F+kTM|wYkp$+}=9S5` z_%uh0V@CVhI1tgv2Q7sq5Qjs^I4gGzGnNw*rq<&ByrJ zQ+DJdJJ+u`6k4dupLCCiY2joGUSD#Bv{=G|Bbj=o*Lo(CwGw~h+yDxqCLxV);1)eL z&RR8gTR#q4Rz(5*5sy_UUjfgNBCY9(&GF|8rlbWj7zCRzRsBiu~_bi6v#!5>>3(y`dV zI$qabQpv^Ni-@nYTk0fh#tp-0U;VVYRok!BW|7qv#vyoI60jXC!Na@PhMa+%yINot z#1DGrGu=FxG~Yc~U9K==;n?h(UCh;^CrHr9f4krC{br;xT=L>R#u4!ag3HwbWr+_8 zG+kz2AR72ULKx_PDC^2h34(>t29ZvV1zC*XJ{OJb6zUYXB!{5N5{(n>fhPKy499A^ zwQ0lk(xc9=&ky=Ol2M?<5B>C1?+lD?V0Fy%*s#e<^I9-t)a!G#V024%p-bctEtg!EoIINkoD7tTr~NG^&)vOxdyQTa_JWzQxsR==gk$|5(oiAgf?)6)}S z%1|gOP)?V{j*$yCTBgUU{2NKWe+4g2aF7r zM+&dMaPWP~nOCdw!c=PTs>$&BX?m@fmP5p?5J*^itptEtPap1xFF#(t;v)#m2(w&8 z=c&wmJ$rIuXqBDh3gYM2*-=`^BzK1b@k=n$1Iq)w~KGl%R5oK5)zDuaK(-8 zeAsX@#}c?11oo9&aM)x!_puVzk8KcUsnA_=pZ~dxB%eG1#mv+6>h9CEYbP3+ z)M@KfL{C z3e@Rih#OHAHjEuKHb#%Y1N%?BDc?QY3DFaSzKz`?e>1>Ed#^&=oZIoelqJv2T>zy_!{*dY%_>DG zRcvt^Ywm|TRq|%MXrJe@$+jzGWB0!7T_?+2w@LDXT>weD>zt(b+?Bp{ zjL84Yh2yNa-=-2r^jujYI6im~^5(wv72XuQUDdM)>6Ec_a6jZ#{l`|nr@0F$j$bR) z|D;ius6QPT?e0EPVN^RmOZms})7&bg?O=SiYuN7dm-z67Vd38fz7LxtuLCG-ciAtq zfaxyTEb}0s;!44FLg}-bI(fTEieJ+ua8zvOxbMMc8>)JIcbqObte1UkCd-)FQ8jQr zV!4n(C;*12NcVAbv;smS4Y&(Qe_t61PBtdQ}3ioJ==Hm9}Z+Gg3lK4 zE)=Z*c|%WK2es-mX2<+w#bQo`lQw^&QIgVS1agP77v4tj50*&H z>G31cP~7-ze94jzx2s8$Fs1G?l(=VE*1v1=7|a|l{Z5A{s-`>9Q$pd5eP#a%M+$zR zGZ>Vw?C$u-Y8RTQUM=pl`R9an`{!KiHfP_79^XqhY2jGi1M4Elx3@cUU$XuU-{?0tBl^)nPBRy5L1 z_K%k`BhNIFd1~_UbjraS_B&n3>@g ze;Uez=CcZx0`@}}aEiSjCRJle>LeDW#-^ssI=&VyYc9{pjU~h)Xm=pfW}4QtI>Ij4 zH4w~dR2+U~3>Pe*HBAaiypmEc*Ow@t|Y%N&$b%Z{l)=K&-p&6S9ED>0Xt) zX082YYgVji`WRc02Z21;7|XdhLpwH20Y33pB+Sxr6Lmr5aOzjIab1Kb_(sdWYr%0t z<_Cdt)@YzmTrzjX(v7B1ao_An_0yn|CYkh(No`t@^yj_<6hKg9Yw0GU?1lNep48NT zKPt=fY~i`q>3Dk#DoEkE_7psScY1L|`N{2_hWoznytYg%y$RkeAP{t)EyBY8_jU;S z0!ch42o)In1$WUq0jvY;ZT~8mAwX($UKP10Msaf^IfVLYMOG0+;Qx{*E6E z+l#AM8QnLekK^(_w{6vJqRjR0^3v8j7S52-AZ`L?I%ry^n zx}7afZ!cB`H?6;Vd%pgRu&Y#lidTAKkU!VAy4}bEd&}pVVBd>)cMs2D#qOKsV!xYi z>l_vvakDQ3q{DF|j&$49So~w%MWyF{|DEMQCU_YQ;yV&6Q5*0sbhWhcp{+Y!QZRjK zI+Hb7IICMgntE5i@!PU==3{u+1n%d_@V$g?uY_yvA_(S0_cJ91|U*A^hN zH4V{`7pq~a1PWo*wnU=ED{fX^X9CG0WRZcDm2k#6Mp|pjne0>lX18+@3^cSnVe934 zMkyoJRK3WoKp^SN9D$SO9Ssw#t!_OEj4_2~hQI zV1q7Nh}HE%JtG!)eu9P>nMm`n#l7CfI8%{IQRPLt#vRMv$}23mC2udLj~ z7&AmJiENP~E`OVU0Rb*Ciw1W$!EH&kHo%j(Lz0?M5H0#gL|9*%c)?{rt27@EeK|;c zo0D-;M*AC`5kiX4Y9#X!ITyKWeqqHt1}uU46!#~(|5 zICt@lTcO(ewPIlQPoGT|^S0Wbka%A>=Y~`HRk(+3h@36emf#1UMJH`Jbr>db#t0lm ztv?)9cQqlomKycHpy{;2f_^UA{nWlryI&CAX5Y#+HFwG^r^>*iGOpOqgPbm{X6O|_ z^wrdcrAYkfI8nV=S^dzLR0;W_iqo*Ggl8=IzQpogDeN=V7+k9j3lgB%w!rBg`J}QT zG<6_o4}`0q%5cQSg}lXkfW8dJ47zH`eZ|KD~#YCY7Rt@@RbN$7)O{Mwh)svioK3+;O*$vu9A$1LM4k&~dGHd5+N_>tnqooQ$hZjY6lXK3lu> zfKQ!X?IU^U0oDyYqYx207gss29AR!H^{UIN_PoBtNQpeXPVr|_Ry5XOCS@u&G;;Bw z*ln0`!+{&s2 zSPl5iSItWIPNj`X&8f}h&6Sf<8^WL)1JQp;8&>*hSi zo^4@*+3;)`9PbAzpYf`S>!r=5{%r!@yUY2y_;CN5Xp^o#XVq+- z7o7RL)AeP4!3pQ}y`tpO_j6veBEP@Ae~@0i+WQqP?cnb)B7@os{J%V(+k*t?GOZu$aAM5gCq!_ITlfv@whHm)(B z!P>$D1{_lbS9z#rEYs&8}PZI!!_QPlS z2@qKXe{3gv`Xf;NF8jgVtVyqVy6N=N@zUP<+;w!QvG?ABKtg?fu99v$;(=uC6pimY zWYOz!H~blAMf>yZy21A!uMe%EnIfNkFTGJrV5jf@yWAJ4=U(^cmBrxj&;#F?+ZeuU zgEgum9XD+NGu&e2LfpE0wTNn%DgUJso<~dm7qi)q@a)iLj)9xeQP8=8gjkBW9n7nM3oTum+JLD+^d@uzYiZhjFcZHLo!Vb*J z!8*0$mjD*<0+#F=s@|z>&+8Y^>W;5IHOu1qXoY@25MZez5dS5uWgih@zPVNTDjghoGoD3$SYZf(=RjsW4E$U~83VsPb z)UROtuD)ZW!0~U~H+SV?q%*XEa$L=FOo)&@%)m@#Pxknk3@O zS^IHbXXtl_I$k)Bh`{Oi_3wQ8U8B}f=`M}+?6s3pyhfS7r?-y&z8Nl!zm#a;#$}3h zHcnj#eba&oP_%YG0ZqIa2B$bEO*LDnWlYnwe~psA^po3fwkQ#?W@hXchLyS!Nn%@q zTxmCQsEWXm!6}eTNACZuR)8UV-Iw0`a)@UT0G4n1uu*V8K1HC=EPF--vQw^wetDZo{J{(_ z(GIZuFtzE}d08&@m+>%BI>ZCjG0hX8lqH>7%y)F+qb-vxl;+zeQHE$du?)S*})Ku%yQBD$_EUt)5_NgP>P1Fk?7sf z?8$5*lw{;-P7*3t*>GtAC8pfK7WUH~f2=V26N6rtc>_;thT2x&ecnI7-Y#N_hi){$ ze9(DPX5zW>(#pOSBqfEQIEUU6_lT@zi_2|A0Xo-f`rYENd7%BzcM5-%K)rLxs%EW_ z!&Z`uQ05l!L#r?j*7#1cw4SKKPD-3IZNaWDN&2iT{3Hz1AH}l(^A~ zP|$oqoIMM#jgYoxvmm73Z`|u)X8I(aCyPM4gH&j+G|v{AV8KxiQgkxfGY!uBR*GhX zgxh3oX4}ru?UFSA-JLnd;`f1=86^Y0Z$b?H?qp*k{;L|86tw+0@dgLoTHU2yER7li z_^xTqGd?4042S~$>~BdWrjw<0_ccx=yPtv2-C-PpD?cX1xC3pA_&lhLgd)&HoYLM% zS?sTy#FpEl6734@iKlJu@04<3T6B~gkk5tpe+cyF%PkTa-&}U|o`HK~d;y9Cke^$+@TA>)|cdVNzNoEC^ zzr;May%*g@O_g0~V2m~@vXxjmbPCXS+upfnf8TFlRKRxUtNh6}VG2L`C;<*7)JOHV zKcHIXb=#8qS`s&GJ;R?lJF(I1Qmc1A!uO47$pJrwd?KIKP3Myd;XWQsIBMwu=$hDCqNa09Y^;>o5nY}~ zvSH3ijpcWH%k4kyAitu*bK(`T-!nvX0~tkbQT4n!&>e_wOhBy>(r45v*_tpqy&TrZ z4rAqw!q3^5Yf+rBIGWwIeo9G9sHH~$8kveCMBA^W`*guU6X1a59iiixqR$w@)vTOk zez0Wu?zsJV^FGALrz2+#S%U>GWunt){Pa{lYRO7Ts_95MQwkAH4LZ&`OaMzJv zbQd&4t6y$5)_a@OJmUqGUV*&!TWGv_SLx8eWNi+Xt_zWPBOb@W$#DZ+QOjuDe29uG zi)tfe597VxEx}4-<)%qUGwp1wac4$NV1-*81^adt6*i}VM;@8m506cU{iO3gEdFCs zCh^x|j%bac@2}Fmueh8A-zS}mr=sgVfN>`ThozH~YL3N) zdNk-O$zQ^5TUTgrZhRW%&V8H`NYz8)R|-p$wWx8{HIa`I+Q5tx5*Kq~lw0^fYdmQxE$L94mE<))M)Lz#OAhC;`2 z8RW<-p&lEG(*6{7YEQKIT`iS5g*$!YNluI&3wipWIC)LBP)Vwmf$|~+J5)EW#x9P| z>hb3P(R5Z}aW!qT#@z|-?gV%D!EJD7Ah=5)K!D)R4DRl(0fM^)x8UvseSrkR9=`wD zXB^J8deN);6aV>oHlKzq@Z`7aoghad7-`wUn4p{_^Exrkhz5tTxX`+4wv;!(nV~(3T0tz$+XKGZKl{Qi#9MoBsoaIL4CP1g5gV%<>AlP|L2I z#+4Ey?vS2;+5k#}diIf!w$ui*r~J^TEXTfCUXT53Hy8|RZpG%%M3SJ>oQN@)qY|-H z=7G#bJr3iX14_^S%!b;JmV$cXcvIT`*FhDXN)wb0srdJqUMh|2JE)Y~i@^poj{7vNl(6^67H_b*qAp;} zSJ*68UWed(M%}-b#&0zgPw|$(DWLCp$V#i)|6w!`oh)D@Cs7c>$A$skm`x+@8eg`@ z!U5LkfNdgiqE3CmLX=@K%&$s}{(A@CH$A@9X0RnB8O9+g^VE$UBr6SZk;pVVCfN<( zB`eia>Ex`exC}kKwilrh|QP$it1Ye@%C8h=~3GAx?2)jK}Ap zj`66?L3$H1(Tj=tDYEGndc09rI3SLAbQ+u5rAWpgJ%}rm5G4Ne@3EK4{Hx!w7>|@$ zM@Q>ylD^a{*WcX9a8xYcZp8HY!BvMY*j- z^~2OP>@x9H;aYZQlQJ=LWoQN(>(IfQlgdjnX;T=dj!7~zF9jd+=(NZ#JE|{07sT&~fr8tA-lKwbP5H6;g5?>tBvVAXxR!=K1S(+H|;h0pZtQPaLu&T zgZ8wAf1Kwp+qZjCoZf!FU5HFwfRl1Sl0Q-~-nQTF5TmRUgh!D2SjR7{w z$W#APVL;(6Mn6ead_UpxC^pGE3HuC#l2JUZm1A+~zP>)J{QXKVQ>^(O#fToe1%mPb0nb(690{>b6qTX~$z60^v6Ze@%}v7gYB<;#a<&l00p-5|AwRNt-% z+1T&_q2CoW>b|EC7j2+bV_rl?)S28-9du?B@FT`ry+f|gBNEA>HW3p;6u);s4aHx_ zgx)%3z@5pR&#N51-@mBXus=kSR-(ctNGLs|Dms|If(PLdfbC_{9@PsZ3O)tmf==WKz6q>+x$a{?hsY>=;|;w*J- zN>wQCI}@!`jo z9D{39-0y&n>a-de;U4Q`UKA>((Q)0w~)wyx*rLUhe%_3=z`+ccb>5- z@UsI?EH^$I;NEay7R#-B6sM#d;Uw6*6e{EiIMVqN22jLTTjc02{^?8oUGG47D6gsQ zGl@T;6NR`$K1sFx@Nn4Xv|>4Ly!fY~>MzPE>CHSq!<>rqX7MO$7QuPw3OXe52G&Uy zo`>uCbfI@y)rD0om0%Sytys8-IdAv8z8mklpn{9nsIBOqzV)uH2hK0V@b#Sd_d<-` zmx=nAZps9Q(12oTx|s+Y3cnpM7XM@GI+^Q_IsK@YlJEy7k4~2oQMZ?im$T~nWqM_g zrKj)d9F0a-@feA2_rPC%)s$!6oF3$ZZyFVBf%Ri26#Zi4*@V=G z*aCGZZGrZ)Hm`HEH>w`m(g3epbYnn;7|^*Is4wf#g{jg59c6Qb^v&z!>XE|5hV=~- znO_u7c;v~e^29solPFQjRMS$^_3RPU<8-aJjFTTd-mHCv1nU5zBnw~ju7~h7rrJRe zbI3~9=s^&h2Nm+a`Scz>NA~dq`{4?mSjo3WmYFGQ<|M*_48x68XCuv^9eApYr4!Z( zWAZ20x}HcwSBgaTA83NrwC4nG39D~*9LP=YO2T7$c zTi^uAm!oHsiQq@jNH5k;wphM%ISQjdqvab@(H(+;ijJm8#TAk=3I^K!s#t)^m9_v(?EmMWH`*-7EmX3eNtJo-hXqhD9>}h9nJ1Y z>j_{Y+dKM3Bm0#LV%160MptsWwg*PT@i$(uKdpock0P2eb_OjIPL6CvFp=$S&_it; z5(m+oyBpLh%O`Cdzdl(07yHn;cc#98_63A5Fw+;@;ISJ{bd+CBK-wK)e$^B!35Pms$JA(u)z(?PSh zICusv3;7D3gO&KScqC$kbo%!4g^Q}-N*)P<&5w20-M;ng*|+Bg{3zt^@!hi88R{oT zx!;)0_1Y8&lC`*_qkc0)V!;E;osS()2;TH^dKgOo{JV?&mDy4y7i5m_xM00eK zH!v-@eL*$#9a)mL&80j|&+PDlGuS&L&wIjoxgQU-hw2nE;!HPcH0>eqGKLVw{~#O@ zXP^l`eb+9=Yku2Qk0M8O;vf6D-to4Vl7H%;Jff(0V!?K%dqJne!T;fV3d~rA%aWe& z9sOXjeh)W#2!@zz{GXJ?FI#t%uDOdp4sPko9ytoS1?&q*V`8`~?{a-UR%?I93+46l z2#%gj@PFfQ^@~jxcvkQuH;7Jcbwno|)Zq4SF7(j9f{w4@@)5zPNe`V~GRO+Q`?K(; zB1dZDIj`%#@1ZS_lA!T7(qtZVWbQ)0Tooy}v$O&Y#(5>5s;|RI^wIgsLmCCcc2Lba zMljO!DZS#!T5T=V`m!TN^5Cwg1nhv?5$kU0AHQ1v(0`-1txI7pkBR(Kao?bgDa$k& zixo0x?1Pzymkl1X;6>p*2%D-v_0#sWqTVZH6nifQInTVK-yLy<3J}X!A^&x4w+^=1 zZ^0lO5WHD8<2iGDQ!q)s27j#Ohw!+USgQDe^(3?1q#^vsk__s@^V1GXk;!gPPN+iQ zANXHt1?Z2iBrnNfK~H>)#EA6|`4=vkl(y|O#>yNSv&*V&Cx&@G3@Z=;;un*6d8|9N z#yxlTH`$$}#w%o9%xtisU}VTqCcvlgr6%wx3VvxN2nsE6TfLS7!6c$!=kSzkpmKp$ zMQu`U0U^jFfoCC}I19^t#9VV#)G!7~Wh(`b(3HS>)uYYFFPU%MmHQikp3?oxEzqh| zZt_~RU5ku7BU3Sw(&ae{%#XSLI)m4-sGC_=W~Q;#rI;nxVU(fj@T;kgSReCmbW;0k zl;zNLr^CYLNWmSKZRL2L!BQBilRcb6H5H$eH0x7mm>4fy%Ea7L&;ZLVRI?bZmeh#y*II9T+LYD$C$2K1v%Vd*$s;+J;I(3=;Q_wbWfgSfTg5Tj^B z)-NG5N(oNjW)E5>cgZV@$e%=mbb0H!A!2ET7D<7Mj>*j_%rF{Cskxj>7@DQ3JCqW^YqY@U5IPg7UkG^cC9kQois+*ZDt!hWx5CP2e@R^11)G+wfu%lFbtA z6iR#X-GZX=wv(6#DuxrB{$~91huZ~B-MFtO&y6q^T~)fm-eL4%o|(jb>hWfxIKWPL zXMC$B+Jc>RQSXNjxj99`uvWAC-?rS@Gv#S{#K_~t4LB0}OXCkqh~W{rtqs^L2GBHj zT5C%~g4<59!dPEJqv&Nkl6!Ew8UDrT&&7>%zQ%Aamp|+2p(-+#gHU)KI%Q8JZq-#j z`^&^||7XH(1q#KZW8E*&2UH=z(eC0$z)1Orf3i^@_@pPU)uBFT-lFrgu~miD`sQ<#k!^KP+bd7nFD@|L5e> zSiuBwpL=Hvp}su9qmzCGs}#o04&f_g1F5UZge8kQlrpUXir^_oxaLIFPfie{d7Bgx)n6)Bx}@V53{v^DTXAoN={^KQr*m`)7- zL3?}{cm&!gnGfOHw{2dOpQ0_BzsB->5s}sUcg=NMx_}%+b0;sKy@bb?B=! zS_edTYigK97|^XmGZE7&BJZ~{!j-kl03KmE*M>`_jqxG3u=`FpzkJ_hcuK6Kk3n7}C2){`mu7bfTOeQ)5m)JhM2dYUCc$73Rs~)Dw5m&?a4Fo!6RY$-5U7bzYcWD$ zPpm8zMJ)8Dy8q#N%I4Su_*cRtehTeO{tp)P;!gtR%Df2KjWgq@z)ZcT_;jMU&QV~v z8yu?VL9l3_k-eLtMJgqfnQWe&>7EiMr2tz=tEwer%S9LWl{_81!Hq)=cw*djB24?WE{0bWjb>Kg;O<~v{ZXbK6(u`MAvl`*&7Ia z_qs!YAeDbm_!RA}(G}q&m%A)V#2|^8TxfTo`LhoBXS7QJhwB}J5X;hOuv@jm`yK@M zz6YV}X8(+rlTt_y1DV(Z=2K%>86G<#;$oV6uB8-Mv>wi|DSMpMS7vp056XeK1% z_5pF=S5eLsPiRSz0Nb~!+=zmt3~Sv&-0{cVYVD#Jg5_KV;@cm;48;84Xcfb)(}Z#` z^lmzGn1YLAq#`sCqH!4}o|<8kHOImX)jaecW@gkf5?^}UL6mx9?nMD#c8U{3&Q6LGw|W9aEv@_) z64bfz_;MCmOxYF|M6y!9TPl`+YPI1GoXD9wsZg!=+@yy`#P-XcW||?u8al=}diu7{ z=U!?F3NuOkdu=~yi}8!=cYJtU`e1y(iM_gpT15A)aju~(R&a=5*!36L$qbsx!VU)J zQ|fYM3&j)kX;iY1p=%W~v0d}>1^g7v3dehNe9rvjO($O^&Fq`LC80qVxSTCo0IUP2 zl%j%g!(bQFMunlctS@7%+V}T&pWxlT0JLyY7C?ssl+I_^l*$QR=?TtNrN+>6+&;%A zku&nHsmBR!xM6KK`B&iJyo@Hg7IKHvG4JkIfWf&D8FXl9%>$FaLXFD#MC0KbL-llR zUPBay1V`;M2LGUCUt6x#GkmIO;%eHOA4YK*^;G03L@Z?ovLmAafYPqvc^HDn*sx`5 zp{Y82l7_#X$HaLEHOvkOv!G3P>t?&KjC@wJIk7D7ogn$v*7(xvYT^m_M zhHl&Q*oKECN0=~roB(vAn_e?9w4gCrN)M>ia@qB}=U-`C#G3qxiO{E<8ZVjiZZ@2C1%!9jBTV(9QqxC!jPf<|V=)pAy;7X4B5*xDO(eU;%`VVAU5#|8cTP3m!8&7S4)!7_F{I`X&hlXKXqh)uJ{VP z1YDyU#x|oERc6os_mM^k$<$YIJW$kGe_yEs$qsr)p-TOjQUP?KXa1HUI<2M6KtJ@_ z)Ho%NgNU1|&#kP=V4KRlhLp?~jx?P4DwuKl$74fs?%b_|32|fKX%Y^iNl4-^8Ie=G2lmqLl6nUc0)If!wK zeOQXflRtQL4yMxJz6KM_uvin@RwtvyWQI0a{zCwdE(DpD2wlI`F$Iyj{|~s8)PD~t;@7FrsmI&j?XN4Qz7O@2uf8@Aa6a;1(xjkqJagIC((*XB zu%0pI+JCf_-;Z+xtN9ke7xxeL+HyxI^EJ}LMH^=j99>tE_Zo#h7V zE=aUwNFi(k<3A3@1nXqzWckW*WziKE6>5wIwBP_lMxnH*^U3OfK@>Z5Zu5xkz{gE@ zC2fOgx@HSSnnC*6{!bQ;>3!qWaot|XBGVMWIi$w11Cn}aB_W4}j|jHINw)?EObjku zyR(q~U;s@q5>l~)b!cSfkZv-2WO;iX!C~kGkcH4S$@^nbMGSmcE_hsp%Fa6aG;Z)T zUSJ&R5V1IXpX^R>JCO?D!llvOX3|zFMT)5?d)Dai#IRNq`l&$Ik+OVQ*5jBl#=$R> zARSzSpTbo-Z^Tw%g!NiS)tpRZ>QiM$GlQ#Z!~nGk8rXD)mW9hjL~EB$sZ-ql#{4&Ahx~ZXu%~7-iAxBb?P&I;vI}G1=bNy^fUllZ zZfTz0{s27rG%qXPqHhfDFlTsotw9JAHie~L|Mku6ef z?%fAqmHXJhn(}^uL1?DEyQabPTi;4x+U=iMm|qrU(!Ec|d_7j)>(1?kuiqHT?Y^Bg zO>#n8wkVGe&)~7el90%qaog5}x}ay_F-iIJKYGvwlqu;@aOjimD;wxVn>>KTW?`h5 zm{<+b)D*&M;Kk#gIAF+1!1Y~3tPM-w{e@amI-}#T&SE&|F=gF_ZHtJdZRtUfa#VyP zMg}zu^&|%rhi*+Syjh$@v=CLsU7;iK8ra*n;t@NJc8g;>`Y<4*Y}5$F_P4hVjJi8W`78-^zHge2@0(^VuFIbl z0@!DTBZyWLpL@#I=ivQ}65@m~CV5MLl>N7aG1)NG9eP1J0Sy4SOanXE9T#H|K};maw>!YR2sgY!aS6!r&ync?_Vw2Xul{Dyl=5ir8%sA7;x$G^`5 zio1D;oPt70{K#!$DsOVkYZ9MeUVqa>dLU4S0q8 z`3M{dXUIK*s@1g|GMjP2i*1m0^+e5moMH*lJNMeMiV40MnK4P}#x z`!7Ity=Z}1NxpGR!nq#?BChdmH{cU!Meh7eYNN3vuMT;9m4h?Tr-{qx@9=EP6a$*& zu4}$`?lz-nhe!YVgBQiL1iC@FW>k=4tllDSZbxt`K>FyW%C1oV_F3*5+*f736~qL@ zkF5uRAtrZbhw!S(^QU~rSss@|nHTb{-n%&eV^-ls8MeWKF!{w28G6Ei| zB_!VN(LWJXFx061`S|9oTATyFG&(~E=2`D%HAgg8ot`a}G3HxMy)SE*L-44?G_PR* zYCk)uDOwpT&Am2hmhj#awB4G(yN;p9+I+k006aQs#>>a04?q3DVE|{99(}|1p#GbK z3w3gPE-9iXc@P1shy5qM8H$vhf!sO%L)7FV&andJ(M*B7lsY=(*3yAb-&-GZSs+!!6HlvlMDup zTscPVq67fA2;sJ#7Y{9frFsFGHI(~#ouTM>-KMN1dVt~oF!g7zB|^|Lg&HO~yEcOU zUUnp4!2c^D>=)4>hur%yS8|I3-smN|*7`{;)@EJS+b;i{*C)OxQ8(71ehRCJe$%|{ z6!>}z0!RPflmRZpTpT(rTaj^HSi@PxK>&R)RFi2<(D*NP^+3jNOa!o;9Fik~_XRr4 z`wXQ>koz4Z8v#W&TKvZLsL^pLSqGA9oeyfZE>hg!gh=Mbhhzv{J!-dc{^S=CcIA1@ z{Q3xsX<|Qy)1{ZGs`7%-&s(j4z!WX)U?UD|tdaR2)aC_Hn7jnQu^q^c!^PrBE#3W4{0zGKZ#OTv{O>3wO1o3 z1GgpXR~Pp}2@#~$F)FK1THc%9E2dvATTQh4%8?(ztJtn}b#p^0<;qT3*3Gfa11&XN8X>ir?DQhV~#BDd3(e{R|{7ZHO3N&a+i z{}tCK|BU(1#x9*aW9hk;HlKUMm^+I~ECWc*?grdPWTdBh9#hOp|6z@-%8L|sqt<)X zQPLa7-E5MjbCps;!z%u6tq|J{!}pr5!8zqItyvJ3;(h%y$I};99An!0M6xxt_Gl}n zg!4YK*ATlI5=NqN=Htyek2QKkcI`M!yJ+E!L7^v1aA(1C)2PF(F(==dA$5Yh&!+UE zq`v&uB?Du8rxw{~8A-v|2xynMr(GXc?I>nP-Bk&3Y291R7KfHvi$5Pgav8nHRmWuk zkZcYvGk1Vo!p*E_ua3H~{OwU&F5cVSnEHiwjcLs1TCoz)WGl1=KCtJ4Gkjr+IVHHm zd0l{w%#k|FwddJZv9D?m9snmM$NWAU=U_k&WiI_%y27Q~|5NBWakme{)cHeNVRDs6xP za-&x+Ph2rl>tfVdm!<+Fu&iL))cH2S;;u+k24&Rr5;NG3aP0Zt9;2wBcvz$6NWI$Z@ksQ3@@3ZI~mZ&g~eE)qh$5#EXVIC^NKt~rs3DG^1?vM zsRi$HBk8VnBKh9;hB0hh6!HQ$NI_$n3k(ja3r@z&?v8tD7kx*wn9B-+5ke-b}8jQcli{(|^A+xLMe zXrz!XHjRAa%YV)J&dhi?5*j1?g6i)Mb{~`%+L_wi__7X~1{ahuHvuG%DQFQXu)4FP z6O>*9$3WJhjP>7C|B*E19vSS5k$P=*J1>c{3t{;E+Vp~i{XTl53{LX(p%J2GF>L_j z(D|#vW5y>Q{Q;ZU=o$;%HEKWA12c))rdx|)hY>9f>{=Lhs~@J();HXzrOFNIg{{={ z4}4@8`sRVH1~t5LmZnxYZg_G+D32H23mM@=jeQxb_&Y}51&biOh6T-k524)2vTT>8 z6Y}J7mPqNdT%g2$N%1i}ME;e*Z1WK+>7Z-2;Xs07j{I-v^R@Jo7)K21zuYJiq|dsu z%dXTLKuJcQ{C_U5H$X?$CXNa+jCI%X!14@+w2(6%zC{(hg#rd^sAoz+sJlQ^yLw?< zsFnM=n{Covu-RlKL;2L)yC@)5(o97d>a`a`&K}pqWUa-pz(OW@)Tyxfe6c50)@eKnPMiq(&+=qfhP zs+h(o3-kpPY?85Cw$rHoMd))R<(^iU=Rox|Dj4ifFR?%YsFHieiX94=)E}K@@>BA0 zt3A(?y<(ic5+APL8W>C<`fsDOJTS0pbfA?9egI6Ia);$8vys^6L(h^K0!i0U%QEhH zjIFv1pEKHjP?_O-N-a=fr^Y8c;QdX&dq49roV2!Oe+a3GC%0kjP;3JHN&d_w9B=%24JG{TsuLsViz0)VluO$U#TuocbA zaAP);e1`3$G4B`i%ruIZ1yy?m-AHEjHW_v2iJ#2dZ^Zza7Xn^vr7A886H2k1oWHOC}Cq)td1&yk2x{9ecL2VnQpG2TCFn zBGAvtaebVW^>?V+?a-0;difB$%B`l{C?l}~by2BP>uWwCxGmbIQ9~kxqEe{)n8jF_ z@+*aY)(}c$>u5|LkVsnne!$m+;Y*}(u1S#AdO_FMH?GWQO{j`$GZBlWOG7=aL00#8P?J%ML{N%f=*W{5k zmVD!|*7Xlw|JfGPl0QvPD2jej=(p~^8_v9QZc}(7q>Y!NdEi_CHdiN6Z7)t<78D9{ zpv(7`vRfvgi6ZL%)E#`qkp=3;gNw&e*H!xpO3?2)`L@G^lhM3+2*dEAfnwKMWkpL| zR$C=}q5j(vaEnjv&Tjo})69(v)$^l^C|G0g`>O&)Tm zJXGX-19B?LmewW@gn8+?Z_6J$4pi<+@@So?SxWn3OxSKSnLSSB6GL!Yy=(Mr?C|kA zoE)p%g~P1*SQW~2T1kptd>u{{Moxfy{X{(ouq42YBa&$#6xpgj`{YomjZ!g>qf8%S z3y3iy)%w6tK-iFsagh)s=Fq7W`%ZM~?HDR0w1hFRcJa#0?eyO>{~FWe3DiQ_WP+=@nA zvfm@Up&_!bJ5R3KTF7v#5tL@7yPpuWm9^N#&3IPCLQhmecT=IGrc$y_AY_=2t2wn1 zkV_L%KV;1`xJ}A<_2G7;rrB4@@hA)nOY2-FjRe^!tXXLp$*^_%G+XX>q_DWUgO3N2 zJ1Dx7$_NtfZCiJ>KBW<*`Ue(7!EpEfI$P>po?^!e8~1EUT=+cglppwCB}i;9m}oGp zp!wd^3fE|WB$-^ohVyX+w|zTEwDrKff=As;@cKDP{@V3r#I<8RW&aMj?vq2~2VH+-DdWB@7jCf@ES|aJnf%b=Ku*%pK3jJ!T z=r$i8LK`z&o%>3%;}FBXFFB>$U&;&vkIO{lLjt^&gz9;M^5_W8M)$_>-3$#p3xUQm z7SFL;_N6x;7(dS%c6cD>o$bdV#@A45Jk4O>Fu6G689o1m=f^KUQk{B^^S(4GNk~9U zA}d98m}O0PSlI<$s1YHO5y%=mA+m8WL3o&w-N;wRvri=Q`?#2%MFK4I+$buQ1HYY> zBz*elWSNM^bkKCdO5!@?b-Fj?>Qg`N-J%g9R^VuK`Yc{kVp zj`!nLrO;_SO&_dp|MTgEh1eZN@TJ(;Ona{qT!%2>X{RB+cRy}H4}|X}XEW-qVyXl0 z0T?qV7|YE7NzH6j?mg1+UJ=ewTG;H4mw=hw!isaMDc4hUnmywgpyas$=@1qfu2TAF z5^yt|ZsHN`I_2i7{i8nMknL4uk&&ytJsh6fWuPQ+FqtS z9AyJ%z1{WEu}{DI40zDx$WJ{3a*?pKG)g#;!->vPf#S`+UPVSlxYWz!o^SO-PY*}u zJO~$GXaMj963}Biv@A6o+Fi&rWJl-pc^!Uz3vvU{Hm~DxGB^dkKc-v+*cBz)pgS8VDLh=AQrSy21(I!J{V-+sPGDuCeP%WTk)@ve;OaMNEOD1P zb6@1}y5AN)JV>_vq%c~db;{Vzlt(^oRM*x+ENgI>d)T})T5%`UIz z=2V)dApmJx|5(ZIXl6^xb=)-?-4`CixpS%2gB++af6!Q`Qe8;G!GAwq&?CCt{Ruz+ zX$V)5yii#}g|yA8qLf%$mc=Lv0lzHkl5p1}X$WU!FFmeQ-@ z0X8?va=N?K{(O*Wz0KJgj+BHr1QA(}ge>Rgl8S0)Jc&8%xOx42Dn)gD(Q`)J_p{iL zXTDPR%k3mDUQOi5Ed|PVqo|5g%^&i0j8n;!-(P>46U?MX=k(LO$he%g z&Cp+gWq2D|#*}-*bosoIWf6Gyod54M|GVzWzm2O*jO~{r!Q=IyL7!J_fj8ypK&iHS zCEsJ=_6K{?Uloc|9NuBwouXi=3U(70GIhIvVc<^ojG1EAh+Mtu1!*S@h6Q@|4KA5_ zz%_l}@=83V_4l$@v%~l`9D~mYXha3O?(1)b0>*>5OU&)pUmrUx zL0#bL!ERyw**@>%qhumA%%>@(yJ;1=mB}*k*2oq^59kfxRK?uk*J8Z4@#Pig>+rU! z$?Tm^wBNWGOun!{5xBT=oJ{0*z_IB0^TF8Ef9TbdDJ*kHPf(r&E!Xtc`NGGJ;}aTI z{k4aHhxZ&@ND2+#UZKt)J8_XI3KFHo@l~KtD}V4NHHtE1gySoi=fFges#(jND(Wbq zj^7ETh6RqdH<>3$9C{D1xTD^Q6l5_0V}mI#Xien*%fzetTbwaIVf%HB?JchS?>hYe zv}q5C0g(~C=`fc;Q3Z{qH!Xs0xzuxQl4pQYTS$yZCMU&l4uKEXf>fJ%7k44d)|jqk zQMAZsQ+;;|lQTNGeeD_B06F2S*&ZxUq|P3Yi7l_|U|iXt?>DTg}ur z>TpQpB`I9mL2z}oHUWO;I0eqWbp9S_X|>uMsK>j7IClNB@^PGBS56sfDU3Iw+o725i#IL!>%xQ;3@zELg>u}-N9<89=-hp7=Dj9xfj%n>p>8J;k##uFg zrKyQJA@S7IqRQ}isanIBGHja#msYt?pAYVW)hP1&^MIk(ev{Q^$oDCTHZH0L_a(l3L|XMP zFbhgiNP~S=!}bBes8lD5X-p`?xT?Ju(XViJ+K zV)QkSnAmgn)vR#Qi&_^5*E(ds^m_IEtf8SiJOmG#C&f`)GBSAl@V-*`6|IkOoDBL& zv1~-PW%cr#e`Ki1)EDTaw0RmcvixsCKG_lr^YNAgSg0Q-G@$jqzDm>MhH*`KYc=hF zIOTI&SGYg&_@(-vy)#>lc4W09=!UEiH8)U9o&eKF&anUUo6z42Unw^u0{w$?xKGo0 z{+ox$SYv_FTw7%tBoujaOQ0-cIAg4k(nP_Qc^m^`1fJ&X4SD+Wj6G+JSwLUrl?|NL z8ae=}>HwQ(99hV|rz`_>&UQB*>mgt~zO6#iRewG0AjP(pdv{RFniDg??LL*OJ1X7~Dbw`iD?f{9~{-xx=6pZtj2=frXrI;fa?iBaC~G?KB!x z!ospy)8RYex##R^gi%;(Hwr!A@#w1nqA!zzt9H>qU}|MTsmIW7sF~kCorr{k_%3EO z@$FF9S>c=?yDG%HAR{)m=}sWTJJF)Hfbj`HDNo20EOog@VnC7!gSR~c!F1zf7Qq$0 zm|$SDoYMRraeiuDo7lI@EI)}RKhstG+lZul;`Bs88-FNt|Z%pr-t`IOwA zhC76lC1eydz3b*S_SfPAlJU-{)d#|9@y{|JRS<3)DH%Ks@++gU7=Bj>#-jz6W7zGK zB+r@(ho}tRdy|sV4A3x!DG`>;ee6Rshs-9x`xpadEi7@|yqVj&I!0nokPL_r2Y|9> zq15@4^w?xw(6D+o)rWo&MxVJ$x1ccuN6w{~0N;cxyd1B{h`6jCRmaSs&cL}+$dPI^ z=^6F7%ph~-{9b@nW3@9Eu1A^qvWg`UXO9gFUfB)fWD3}U2NM@lMgWfDG@9&bUmmDO z8YLBG5b2G@R?CMV!1VqAh5wfN@iW1Bl%&~yRm@bV?C8;_3*iGemqO1|K+uRRnxTo0 z=P6f#MQ%B|lwIBdT`KH_pmL{t4j5K4~c0WsxcGoCQcC(j+x-|ghu^Uf(Nqgv^ z1;oBR_R##lDFTmRC!q;N$H93hDRxYTO63s)TwVS9xjtA1Hnu{IE)-J^0>|zrH;$pv z!62YM#3xa#%-HD8o#WK-kboDW`mKAQdU6V;;&W-OliZhDoFkg;C|s11Zm=kbasDLv zR|RRGu)?^eicsyuF~u(jXtsV>c87>}ac_8;_Lnejk_D?8ZOTEp86Ryy)s=?+$YrWI zhiz&1XJZ%Dc?~Io-S6o~ga{D4jH-&Nh$#AA0ONAX%pLA=W~^(s59H@UOyB&k#@@sMOUZR!k-qbMK&)KMl8NF~lNRq8>ulba{ywi32$7Ln)SQQ148` zD#`)gn0KRB>$WMQ{3YJqkDtjfhNH8P1ES!MQt}P`tmTkBS_P@vOysIUCRFTe6@W+> ztBq7o)r1u;u(eE%(Q-y#t(JrsGKDxkrb?n)Uxmgs)Swlp-_iGN?;=MEB0!36?>GymD+H|*D_7nS+Phh=X0QvEGU&)H6hWGkD zS2R`u#1F6aJdt=8P$86XQYcuHF<%?^DP|=hcTmqngFLQV_W|n|0TbHFUn%xFMt!}y z1tqK*Citp0IRb+)R9bwcX^ughN}z*TsK%yYo&<;#1^Ln@P&OB@_4LEhyw8N*LKh#P z$wXe&^BX@r5+nJ4Da2KGm^8H+o+V6@tCiehSFGEiiF6{S!qk#=h3=$q=0gt@$Mu4DBWL_ii&0$wZ0)TW0-d~#qNjjITs2jDe#3%{LvZ3HHMi^Vqc<-XcQQ$ za%jVr1C3q5sr-kd3_ZgWeU3)%28l^5epAAK!XsMXgn-h;eN>6*ZqkaBF8!V+9^S$> z9}0~fJw+|Tzo7Yr8-%A><`C(oOxF<}0g6oI>T42N>etXuiWcBFB&6tqLx!CJ>UG*n8>SKp>ZFe5u! z&F5ayPdCdYctCu9c+EAYQpic?@RU63ME8x=h8?M$T#??>KYgjo9QCT`7hA8< z6tRbHSf)4XC&p7EbU=tu8nM%Gwho*!qC7hv$qPn#JOwsH#_{ES7`21u(fwV-BjiBH zVv{NM)`qUuMEQUy2%To+AmJw3hP)zcOHHnfND+m|((RV4EX3<%o7Ik%#EVxCJdq_sSWnK8ciUKGKLz;wue~G^d)((A zLDJEUfFfFI$?BDnG{<-+w`6?nQtG%>^6`VY50XgusSvwaOVpmKRn9m!3)PQzfNy53 zyt2&w8MjBlhpvDGknrp}2$&%fOF)?UAA?(W9Pu?|?z0U>@COcU(+j@WyjFEzjFkLv z*?+0JH)L}-t7|x8oCCo8j|9S(jG<=9vKhkBwUq#>w_p+@?ZV3950n#9A0r zDjkJLBs${HeEFhotMkP}C)Y2(`KiFCij-wj&i-W)1N5#3*M)c+^qAyynFbhQNmp)T z@zrMnA>d4t9d~OJVkT>T;6HI~q|4xddhq+cf%m>|D0Qz{8Tu!$p2>|r#;g6wJYkY< zh7w8sExW$z-STA;*q#IQMxNc{T=CjA6X>XrOw_c~*?&VLCQrSw%d}zd1B*>9me^I^@gM508 z46RsVxFyi)5}@kn=$yT4dl%d56N$sriH`WBY%HtPB~*beh_$ya5c`dV(VB(74y}yc zP;ZOC3_+CYBcMmS1IKc%d@<$eo=e__%5=P@_+vIpn=$uozdBW$(WIs#ZRwc-z)U|z zc{$2BdbN)UjNIjrFyyCTpsF2TM&t92msDObC5ZbGb0Ioa2#K_=fz|BL&dIIstK*EM z!qrmFR8fp8loZg-z&M8)xm?X;$c{Q!wG;H%mEW?&Mq=x)W_}{gLT7xBNZAgr*-{lD zh4Zuoy_vUL&fyr)BVY;bz(ij|1=_cnw-=4*e!2X&hEHedU#6aTLP9@7-5$Pk@I`^rO}5%cC1EKoXIQ@APQ3%vNJ;THLts+*C+Hv`r{no3d;d6!wH!_Mb~&0U?+IGI z9NBfz6dRG@k;hS_*S_I_9^T}7o4H6FnlbymeZ}lcH#F$+8LHP{!h(ee7{zruO(oDJ zNe!YYNsH11fbJisS|)fykfHTy0X+}P#$Taw5|ZaF%l(EL=W8Iee(Of@>%HQ;3^%FK z#>YaX5Rky+U`<+RH%Y=42i24BI?wASC3`W)hNN)yoBv5kUWZ@@K!ZwO22hz0TApOJ z@pEuvtAi$X^41jxxZgVzrKvpv+cE>iXSA&z9US;4rs7x$~thc#P4U{MRzFK zJ#i?{rAd2~DG>(>NFA~XL|T6*_7SG5x2iSJ&vD04Lh+W53o%@|bWd{G#YUqis=hB! zb@#{pyvp0__M14hVgXFzsgH>TLLBlw>o@5P+Xg0yo9MsLq+boqmp=|BHk5U3{2xr4 zJF;|1Uf69_MfkBn_(j4DLzGap74He!`F*O-RhR6KAa@fbW`pN}GAfOsY+W7ZLmE!jx%SjTbd`k+hm-B>@Bs)5e zI)5Ji7J_V^tPtIXV_rgk33UGVW)k|%;E(%((k+XI--PgSPB9w^L)qdgN*wvDISIjV zzuYK!r_^!|s(2{QQSjXmEes)L7Ni$ithniSoaV%{9MR8Pii}nO#6t}!U z$}uk@N06Fj#$Hqap$!erdF+2AEcNhU)Xj;1_8MQZ-c)X&#+N_uY$B)N=|#qqw|~Fi z^lUm3|3Q`fw#Th6P)tg_Z^;Tq9s;2VYST&{hP3Q|nL{6}JPvA_d*zA>izBf8XTPZn zAAS)2_sg;VHUI57->hq+S*qn*Yw~_?3)2khW%$Pc8Y+a(qZGad1qC;#t7q3b|KYWY zAK#2pxw`iaMcp1WFq2nKRL-8z@9=e`%?xO8Y5>7SzZ_(0Sbl(d72sfPh_TW--MYH&RN`*JL~t<-bTG>J6w%?83ymMcsC zO)|3GdEMu%vK;Nd2r?Hgq`rsuU^ds{gKh^EXxAm@JI+th-(sy&IIe|gI5yPU!Lx)> zO17^ho=p5SQ*4N130c>?f#s-Ju%$wfd%l06GzI7QT8~9}5(@ghYC=mB%vFUMj%l}6 zo&TAdaH1`g9aZu1-J-5qtqJXZJ>t!$vCf(prB?RDg?f#?s8Xw2ixVYO7hI$raP5lB z3n&XEosZc}aw|G_gk&;D?RKg1g%~4p7MVm3w&@v=)n}vyzhgCEg04LC*{Bp!Vkb9` zA?hOIoVuywkuTRNJVi9+10Q&b|FfTgrgHk;eH+Uo>qARyzjKLvIh-#)SAU{ty;Qck z72X^I^zAW=hzG;t4_qw2nB&uFgcYO>kn{ReEQC=`RPy9b(7oevg{%<=k|M-t$_h%Y z_Z5#jQ>m#8gN8Ees&(d3ldxTUc^tS0XP^|EyY6ZHV)uWy>Af4dHazno?k*Q3XS(hZ z-{SKf7Hkj^VI!=xAY(hse2RYco6zr=lM%v|zh<-jt78H`!jVEBxwFsmMnLWwIR{lS zq=NhqCOc%&hlo4t{#Sg7YBbj=lkB6i`Z&l=IGX*Lg+*PN{JkV^Z20(jiO-e(d4E?78ByWH2b@$dhwOG zNhDxer`|j5@&xuNn4VUmI%OFB=$tRjU z_&eza^J0~M8=J69dIs_b7j?gkK|)63WDrpZ?|efpCo!Ce(Z{Z`INHJSzxhPanLYO? zIj>FzKLtss`-Z}cGk>frL@(kqa_hU+l^*t-bRp1rnaA2?Lgd}jzApo z0*-*SBO6E=;e54v$O{?jctarnUg+M!8g>fk_4O1d<{lLvDwn$hwU~8T5-O`t%4K_ajW^m-1)mU$&siG48*xmvfAEiaNakkW5oBSGIN5D+7Ev$PIk# zVeAEVu2o`mRw5>CGB{#vB|TiTwZ3bk5OncJJ4ZZQHgxNn_i_#75)Bwrw=F8#iri+h&tI zjnM}0o$qhG|7X_9Ju_#{g}pz!Eq#4cOmrp~>(Mr4OFB?756DGQLLw7jFy@!)JWfs)ypDxrfssxN-{f}%UuQ0a3X ziepv&0ahh=fp3K_+^CRorZsWKS_@%C_L`G<-u35Vy#VHUS;oDj(AaljT9V3|C1+vQ zCbqfO37c}3QJB1Xi_=M0)#_{2D9mqDj6%GcA^gHt0wKEMB}E(5)NMR{XUf%<4897j zq7=hRzkf1NM)%^Qsigdz8>AabkLVeB(zA!iU0Sb7nD|PnO$`6U_>eXIjY?anhVQkq znu2I0jyq2GOX|mo20>j;ouK{+IqtQiGWn|!3gmw>kwEp+)Mpv~1q%mQ7~*TMQhTD{kHFaynxU$Cx}fjd%b1p@$6HlK+~83s{Pk z_NYjklTyzv-Fh@Yo%bxhSu|}jCwEGJE@;4!1csw0{5;KeS$N! zT!Ajsg*2&t9d9tqMP<^W>f)&`Wd>I3g@Be6SL3b)jaP?cZZThfQvognd@+VtdBR?1 zr1HvlXW`|cEQYMm0x#do{@FIMIHO<`1!4N$-4btUXPQ5!do2c7X&AdAj~K~>E3=^3$e^Ffsy;i5rlQkCC4Can z$DQKwY$O$Qeo{>PWSs2C6BeZCbP?(pRgqq@K0_Ou%j&vZ-2WET2`wW9w06JaLuzQ+eoiJzf&t1 zSGJpAg)HY{rcn<_ecFVTuh!J%@#PIzOZS^1D8!&d}JA;;9+eVcVQDTq)UVLRf= zXYw0K^BQfRtj~(T(08NY@M+&wXYT*b=#y>Ei9?WDD(~eJ!q=9znT+$Z?M`LBmln^P z>TvW5CX&pl;(%v{j7-}_os|eJ2YZ45Be~yQG|Qq17{ulh0=}57$({wjTzJUa=Jeo< z?AiRB4yAo7H}`oJ^@OiU&WFc`MWwWbL|td)ximk(8_g7uBI_TTc!s8^cEqAmF8MDt zV5_f1y?N}o&ya~Q3(S+=<+mPqdmQw`hZQ8PSo&EYxl4pHoMtoAOj6sFqYqk`NIB7} z<@L$c06IMjtSf=!8jt$FGIsxHdeGhH@tg#XJ-Bd+t)Kqge`Kqc!UO~IuB9+MIm=XD z)&KL*sG4zq9>R~+)i7a1$8BosS(tCkfL5>;}C!G%@nIP)>`xY&iNDg{_npCMI^mIp}QhNh9c> zBL0f-2gsQuFQf}aqI|LPgrKzBNND$Cf(8pqZ=>Ix-E`$ymeOm+;6sR$E+yq2mzdQNQ6Tke=K871GwDNcdV#WB# z=h7*CVh2s?MJ^f#0D7H@)-69mK*kr>86ZF`zyPWxVspWul#uYRIr%>VQ}2W!s^L0h zjF!8Jv^~jr+o%clUz;>!}nmd_<$qe47G8;7M;uCpX5G--j7uoiQPhX#;OD)WqvmP? zw&&B&ZQD%Zz-j{^n{z)>=)&|6Il}}Ol>+MLWldr+F=c&77e;;Oan8m(v;;CpL*F*iEt?3!Sj94#s^S7o^aW| zpj>e~p|akp5B>zixdkr6gd3Ufi8?6p12$0n07~uxo5#f?S4&z^)h0^dAtM>`m*4#h zRbpHv5h#A9yu;75pfEd}i&J13=7`(l@g*f0BpaVpBl%%W80kge4+=yD{*xU zleKfK4G9%K@|BU`h^(^LsH8(fN;+O6ZK3Y2IgX#I6rqQr zP)MHQ2%RUlpXYTQ?#aaVbl?bmU*g2qKA4A*xvG9Oa-iKPo)t#NWkF-09!9osgv6aW zxecVr^IXT!wT8)XvTOzRQ8X~6*3p!hth5SYX$|6Y3p zf~r#}0V`hRhsYzL;No2VL@~^9hF;`ZZr!H9-MqW{>7hOFUwhl>*sP22acS>*wboPh z>`{RQ;2@Z%=dlk=r&%FxHrXO-hqU6+TevyT|KnGoqH=4;pUW=vPpLU=F&4Yff2Yz3 z`xO7zv*7I{w-%#!4yMmbmWr8vhz+d@_bUK{F?)Y~Y-c!R^=_-h>>QyAxz(%M6#_S+>7cg{; zJklowzrco;v3cf`iEv6FjJKD!KqlD!eHM(bS@7bb;C8jmWlMv6u4bjGN8>rG!_t>; z$E)xa5ZC!SXpK0bJcBHGHA$@<<}Hkk=FrVO0}xUD$7nMpi;mTU{snbDXI(GkQgQYO z-+oM0w$kMIFHdU6|eg|=EuE)ikyk@uHIEdxNR^0-}in|E^Pzv7OBy#^PzYFB|A5mcih zNoamd{-@v&QTNX-q1AE}()D#>pkSdWxoDgbT8F5i@gG14uFVyW_K#2Lvc~w;ZX-SJ zom6?K&C!Weow%@~^RE2b^xDEsGiYLKo(M(ZR=)yAB#dZH&~aZ@I-4vZ_PN$Nh0@9F(3{tO4cE5?^P=@@c0(_ull_1ofi zRgfx}qPURv{Pvgj22h0`SsQ_*CV%?@n*NW9+W%Egn&K zPO{R;6!b@5$$?CLzl(O0%1aK*EPAG z1mzJ(3kjcb*n+XzW* z$K}&9@XPDLn6ECO&>lu{2yuBKqh;}gg%u9G#ijX9AgD9YNOJRX0_TY2ij8t+eQ-Ct zX?=2+`to{N`L7SQKcI6M-ruSH1hY&G$=K)$PT40R(;UBOE33;DOSQK*!0kRX=NWOj zO`o09&hAzosf%;#@{F|aP9auy)2gaJOuS-EuCcFX$0(BVd*Bf!@o6R}z&sVYBr>`& z=@N35TYvDnc>pNFTjq8K#=?uSo@9)V`VN!fV!gXKu_e~3{{=k_jR5k-X5qOd+N2`g z%(1_^VfkJ4BkY$hO~Y%42%2B1yDaOrteELAqytdj*hquRG@D9R z0+-h}KAr~c<7otfHWh&;mJ~{^XETR=N$YrdQ{0o$G@)#AU-S26v1+T)@uvHVWH_GV zR<&GS^XK`tr`+ryi>{zlZSEsG%1tQ>_v)L;Is*g8$mnl}x^#47nXV~CoxiXD2>-h{ z@bLMwYx}>`{Ja+OuA?%sR69Zr$4u^QM>XE;Q#ENnI$Irat;Pf5|L$2L++@ zp!#uTs;qfpe;(&Y9!dA)jYXtH6V1o!Wgg-C`(t-mcLVVCVy>G0nW^(GHU0X) zRw5OL(7q7SYesz)`|RFy;4sGvVEx5J+ccUH#cD}1Ed3RH>{2x%bP3(Yi^FQr|_Pku^>_~^aUi}TTx*&`s>`6E` zcTY)2c{0T3S56sB+(^(rau@B;a&(2cGsiu^=vhV^kxU?bccEDBM(dDXUeEqG280+N zyiW;s9&#e6^Dxm*CqovWZk7H>x>|fM)3ThRDqeG)`puWJx7b8e(f&EoVp3_X(Jxm$ zpkxpTUb7-}{A4g$uahOe{5MJ*I-d@NL@Q6Uo+i^e^#18P)b>Cjc;2hb5vu z(wfbl7>AuR^;*wG2O&r5S*M+=CKz)LZ(zIJonbL2TuNTpM~FDFq0V*$tT3iaO=1U5d5`xi?SK+;8|$H*;Zr}O;{wJHr2$X_A`oZliOkFh z^JH!SMaV-JGJ~i#K(K^2^2wKV`UQJ}gS5_`_mSMmMl`i=@cM|jI!n8_ItNG1x5y7q z7rWnXY^0q>lKNCried*<8QYIU_ZmrLD)4e8fK~n}ZZ_}OdGuWM9XZ)z(T-u2ho;3d zywmTN@vSkIfP)?k=Eu4%bUi4N9l()Y|NYIGY4{G!pljd?8j5xotJNt+zpnC}xtLoniEMx%%hL6k%)8AN; zJhL2JYn_FiL+LX{c*RE;;=#2XEea+&o7HBpvTomP70+MUBGR3cwJL@g#MA zGXFqx-_*VFTC-kIR#NSUBq9$BN*t}iFx(ZKDGh_@7wG6}Rn%OD(1(ND6z>%r7{<)6 zf-B<7+s#hAf?Tt0MuZu?%Lp=mTjV`etVUQ~7{aVHzxeGA^Gf8Jb&QpjQpwR8pF$+k zL&Ozk7UeOAC=v_#i5#0QGp;?3?I)wAEtUY27L|Ly&63I}UDJt&2c#z5TLH1DWduJTOt zNs}ac+yUz+7Wipimm$Nf5ve56blG{45{QL$SSc?ZkVx`SXvM|7^ofB$AhLuT>e^Bg zM6XnGSiNC4;Bn)~3v}r%e6!ztZ6ajHGf=!wUlsvN*p;aG^*xC_bXwbL|-+&a? zu0F6!0^)Gsw02_x0ss>*w+9KjclUXHw-1>YQ2Y4)>3Mux9TCpJB^cp3p+9;F9&2pV z?xsYrP(2%Lt~Z!lGL9(Tv(E6nFem?57CsGq5NnNF;0fzow=l~qcoZd)ba;5eMqB#p zOS;2`t*=ql%-F6rzi}Gk-h3G8YAdPv-)L-Z18Ox{I&zv5XMXl;yjz0{KZ zmlmA^4CABKLIuwIe}C-O_7G_B)3O#x#8S@XV`{*tK_V_={Z1zvC?go~F5;$*jfvBt zwZ+hLcakcOZm0;NQlMdRC7^u-xKaXdT?BbF-^Gh{h!umlvwt{|#ivZV%Rg(cEnAOL>#vq9*>Kt5HE}f6Zdot?eMACp1o#0<(lv&(Y0B!n zKDE2K6kjep!bys27#c^p#YM=72{(v#LaaA`g81^xX~WWMh$EL>fxnC1a;~U(NTM-A zs68;zDDYU+;lxWwF6PFjFD|k!K@m9M>e@I1I5SwjJRgf>vgmQiK^p{E5DGAX_~n`V zx0jd1HJ0C?5evd5?bwASIypp+sk9dnm>LjJ4-| zZH7YZ!y$ZQJPYyCiqn9M!71i*t}lt6RkBYRzBQDp0Nu%7wU%!Ti9kOt1{0>kZP4qGAg;5qTyu9up^PSfN#dv@5;NKZiGpM+356 zxzRP6P0<7v45K7q;OLCKqI*c=pt=lsTh4xKrxnZH#H*^9a?OQK(o+m<;-Z$&fb1+Z zNJ*r4&M_JV-n}79q@Immf9puAVMXle5iKTk!%zl2Vt`K0mC7l$oS`#aNKuJ4lBS^D zu6y4z%Of|uP*Q|k4b#wW0D5~*B&58&uilQn|1)0BPaYVfy|k&X)_p! z3J`mnClO$LZFK~f!>(~fCwK$(O#ESt-xX%1h0bInzu22NN(qjj;TtymcZ1e zMUk47qb9R)2z3T0;EWGy$rSs3(#^xu z(~1Um;a|&lIglKAU`)~t2#``wBXzS-8pI;Vu&?fjNU8cIKi9^<$FCnHVFsP^?Bs5J z-SN*S<_^WNqL#7sDRJ~u2Zu;X$H;w5hF?zW!OE8M&$IojmwWVk3bU7~vji zE1>cipA7a4&#PjS5j%cQo`Lg+s&y#a`x0-cl;h? z4&Yn-Pd@K-zQ6qoo$9b+;ti_$n*U`rj4V1V@M_oT-)9CsXNGa5q9Np!SBhP~cd9se>s9OT6 zcloj4@yE!2-#O-PwLD0A?vl&C0-%Izl)?@wD}_a?lBchRO_g<^nTT4$9){#51WZfG!*-h>GLO6?NBV8|)}%$b@C#Mz2ef!b@`XF%KTY!8i76 z${&G?b)79N5FFPC_%WrKgQ-F7mfi6oJL4S;v$9^ul@0?F>C+wG1wo^In*<$)VudQ^ z2Z=bXZ~OjH;xPb{&g%QrXi}*5V#>buoEsrh!p-YZVdz&B1_-bd#&XNo-gT8Z*IrQ= z&{Eh?nj=r4CZdIk6E}OC(E6HK4b@w`<~U73Yut-6(Np0_SJMW;w*$^e82`Bk|7eQ^ zrB5C|Z{9T#F5c!L{0a%INQ4Jz^E+zlyMF)I{YM$;PmB8GSfB8`rXC<<;$vn&T;<~JbWP03IT_lG0{>7=Ij z6QDf%1q{k?NBEJb$`Oem80Fm(DbvoB3m9dnl8tP7L*?0kI2pvhhQLNDA&`YAt$zu} zbNGsaPPo*0i!lB8tv}Zf1`L&&d1P(YH0>W)Ft{n#Fp5$jj`xuq10But`^I zu9->c<7q0>=*4D3uYMVTeI-oMkYUx$EwwDrv2ZhyR!iw|%2eEF6;~ZHeMmaJz*b|F z5Omflo;%B|2``SoVwrm~hEVBwU5KYHB-u-Hcmnv)cJRY|1?27_=mDhuLn+`sD>&GK zM0Fr`&anMB7J+1R+%KI7?L;0S#7-IhSw%D@+-<+m_6b0?x zl&`9L_2LzNr}36WR7NQETsVosvN;T;crD3SM@Pw4;W-gj}4T&I^0S%o|2+a^m1@>^Q)D$5`|94_6i}ndoMMbl6;gS2B z&}{~dEGc-VLkOO@kx^27th5G7hYB)z2<&;>$XN}1`IS}TK&;?xH<(1kSqchfH^?AD zdY!t|6WCFXF@26ivFR;VHHtgf1G$;5rCe|Fipk9(-f>3nUJJ{;uwSBFdO24L4hRJ8 zCaLy(Lvll=;+*Zss_ezEaZS;pFRUb9f~~UOshSyXE*~$fA`UBCAUb0|({Pf1O<#!G zm7{B9lkRlQG-z{2H8!+d20QstO=77bM&uX=)jK#Ltd2I-|G3Oy2pLk+jv*;s`zLW81`iqV6g8hF-B``AB?+-`MYTLD zn2^f$7NsQ!uWHq~Gbd|#W~g#9N1O}5%#48YPkOY3Z2KDF-{Sk6|AYft2}~0!_T0ZgoUMVe_&SGR#6pk^iRS5j>b1wy5SpXR ze08#YD)s-iuuT-@0jfK;0Y=8ow#Im9{+Ryb^~tUq^-ct_hiplN9XcXz8cIzVr^~$i z8Lp?^IJa=HB9^O7b6-jUWjW(i9UK}hR5nARO?Qik@;HLQSUSOUG!?Y-p4Ooy!ZeFDr)X3Fjungd06lJ3U6MEEC05NGU|_ z5k-zD7M~SaSbYch@X~|D%yiobRvjsw!w-{6-w&k-TBSHl{to(bu*~BMUvYps;NWaH z=Cp~H#W#OZ4aBSy?rXF5i$XXZb^S&yzj{X5fsvbuc+bYN`Z<1V2q7FtwT7Os*A-8h4(l)rsPh*)`_#JWN)U^ zA~sP$Bk*Mn?bo#la@1LEJ+~Xey)6Twm@Kjv-q#CT7EO zen!C^o}MppZaSJWW*s3FY{Lg5nuXM-Tfuu%~#Dsw@Gs6Z~)Y_^RH{ z+Sl2r^n*cv6Z}0nVD#?0zGq={1P;**CZ7_$5tdEV^-L_Sa#b(ZsQFK)a%O8JaPgY> z_13pN{s=-+&*Y%MAya*z3)=JMWh<+|3F>w1fP^!k3*om^P-ikTg0Y1_MLkI!s!2ld zJG95M;a&UQ2OijKnre>e5E~ai#G!P+D}(l6_;O$gL}Spo zN2!GIT^j=KYtZ=*lk(cRzrs;H99RDVq~{E=*7&MbPe>Vm{`*U)^;940)N`-Rj(dBv zv6$n--MKGDlT_o~UDZDN7fRO$b9Q!VW$6d;CQF#ffg^f=9eY%jFUcY-vf$jF^jUaz5OR znGcSDO`B&C_76d*y60`a>4cxE+vM1mT}|D?nswn1QOHlDBpD9BDhhn;@Yn@uTV&MG zQRBYc5mbirm+0o(d$x%c8SOV-raZsZ39&|SV?#y&TQx3ywICvJ!zpjT;wf4lb>0Az z8C_O6$94RUO2ADj&|(t*q7_HDaYC@j8-w2OLw^6!$t1!CQP^}RcwXD`@&eMbm0c3! zon(BJx9~x|ezpT!?M=m^9M zPHz(_-(D`8I&TxBo1m}CK(1)vb-6LJO7j@69y!YBfJb8t5MK+da2C3<(%}>`xLm%c z*yn3~03RA@mGqUEl*+#*j5m!VtIY`BDysC2WOJuxSFBNG9#88$uc;8p_tb|@88oOe zz1e+}H_N+W#!xY+nEhd-(Q1dMDO=6ZD95R^JDH}V=$qV16^+~Z*LfG{_}oN!VWvzJ zOnjFrGBc_u;*~b5OL4*A&opEm_}Mw?ljcX=`=z&Se;t;pwLYZMSm@O}1DtBV0VYHP z)%2=Z(HYj44=caFbw|xDQeH^;yud~9I`kR?nxhfaysfIUzTO%NcD=Tk<6Gax1E~X6 z&#gzY7T#R_T0YRba{a+;>tLy=bw|#VwBP()5;}CsL=L1w8s8XJ(WpNGoCX)qX7LvN zWG$x8hexk-00##smjbv_m#l7yityMYXz&8b6!dFJ zB^vPY1ZX}}VU6@$c%Z>pcb_tGLyd@RU9Lx1**R#hx+1?r@Q7z4rZRu^&N#uH+G?Vn zq~Db8!B8dG-7Coy9O@wvcQPsNw9yXFj-}rPvW_c$sLv1c6e65>sPM0}0!0K0)B zG->q#p`Rc*4$6-l2#R%QXYr6Fo@C66Uhv!`|C~2t!=$2te36p&ZEws?@8oUuGB8!4 zMfU@A)Eid0j)$zBf{5V1bEPD_NTOsZ!Q#7_3Lo|g`Bg?VSEV1ofy{v+A_aXU?Hy_mY@tRILA?{{vVP2#U4M%D`rM;!i|6-J&n z8mIc%`iINk2A&_)oa9MdG;V#0ja73HPD)DUam>ns*OF{z;p2Y*>h-jY5GCv&k?(YK^U&Cu}3yemK${sSjJ8Pu*UwyoOpX}6T zW|ytnt9_ddVBB3OeEW`TUDo~gZmFO)(=0^$a+l+wit_zvI~2{GX-%>K0 z%k%I353}^ZzI^yy-8(;E_}FdhO_?B0hWO#xZJM7RsfPANGld*)Tvz(K1+S>y#OS&c zef6?;B4y=?$?<1^)7=#Gc^teIQuF_P;9iOa#yzRNi%6VG^{(}^o9mCEX{&5MuEcr9 zEVOw=b&!&0Ym=DMEAyFY`&?fn&`GJ2TK+!^K&Vt8xiNLx+z7`i)=oO}km<1o(dpOg zLx-gvi+bGr;V|)+J3U(1%FMt?Z5q_yB?8w6sL#{HP=716uGc;`ox06 z7y`GLYWcA#^lcCbJAv%ipeu@GD7B7Z&)~TGBBXpw>~S{c%4uAqeYj8Pieu1g#iTl} z-!UQ|kbY$rOeM`aMEH0h>R%`74qDyf^rD6N;D0Mn)77KPNO|X`TJ&j}Z-MC6nao#X zwVtn&F;;dFC>GVuLwyEXVo!U7h`F%poSq&hNSJ5aE1Vg|tqA7;C5}v3xWYLIGrH4K!KW$QMZrMB+_>4i z_t2{fU(b`KfKnhpLbTW|po4Wc#!)w$h#WM6_%xz+%^eL{N|O_X^p6!a+p1iZYxB!Z za{Z@mQSQpPz(UT7Cj!fmsM}Bj$w<(+`^bI+vkLdt~by;?Jc(I^`t zI*Gp;2|~fh&B83aSQl^bOhsyM)IG zaYK>=+SurQ&g?#oO=i#%5Q?@-_(c1wA>0KH-~E;|3Eijbgd=hM!vecef6vi=S$SIL zZqcNf<9f2dgn5+A%%QeX6VYR+h=(va7rY7KhtDoGRe)30=xiF$} z83nze?g-8ba?pM>*J&ght8BUOi+C3FmyjuY5ZsJkA2gE=GZz^i8H6>jkYIbpRDvuQ z0*Xy+wYsEI77OC0!ThRzx?>h^CFgpqwqrPR6CiMoe2|C*XJwSvboAGXj<_ zc*&olxFSM-h}LjB3#BPNCmxYwv1Z+D#OO zbIO=X^Em_}>uPZI@ANvOAk}c)>=!VgqYGNUc)d$z2t}$QDsKRTB-F{a)Sb@|2P8*s zM?gQ`zL_z#*7{pEujkm$jIM`&9J1kYqURDJX)w~M zB1XuEPHi?p1q&cT@}?nIN#QA!RoxWvzh&bzo{f%)v<>7Y#E91K7N(V&Smn%;_^YAu znK#K_;|L?_4|Gl>IiZZ4RDI4Li{`^fK~O{~Vgk-o7nW4$8zr)vQAcMEwT=W_^v94; zo=kRjT?NLXzrl=*B8Zogb#%+By%e#uCY@oiRr2Ns&tRi+oE|dAlR2(a*t`rr*A%tL zM4=%u5fuCnl9%553&wxMbdcgIW1B>=iQp9$I~|T!+XSn|Aq%4?n`Opm&L6jN05XJ?=f_JvR}rSZ{jzxrH;7uWE5h= zRXdS_E$~9pEtf9bz6!KZMRUt(0!osL0x&#%Pzu#UGg{IN8Lol%DC*jmg|^HBo}%gu z*T@aMxg{7eUN;%q_tA@OxEkCe6GAwU)i z%v+rdp3+7sT(Uz5GQ^@hs$%)%$Ax#oU=9NHEP4{J>U$!4ZjPg6-*5j?M`TVL<$xbN zWYNmv`q%*$gS&6qeuPD!oMm> zn5YjEY7P7gqI4+Y7IvthoKlF_+v8U6CxqPZK&a_};&fC~@5O}|g*oFI@-OPk;@!E2 zjr8t_r!Eg=`Gzk+aeG`AVnP=EEYSl9P7bd(q3tA zsQ{0ba(sy~SkTe-RNsz6sm-dveL|g3Un(8kF2;o+O{HR#o=6otpbo`T#PO;j-En_C z1OuxHjbOMIKB=3boHd15n^C4vRh=EB8B+_lP+HBPW~WY(F8sSOXD4l_OVzW&cQB?A z7h!Dsn`mK6XED0 zJiJQ>D$*r&38h{e$ed7@WEAQ|UWH9OBhb9ax^T$?#g|h?Ldb6uD_>g5Sv6t~!SW%XVO z;DiS)1u+>5{xb$kFsRT9Xp|JxqF1MJo0h{to@enj^xG$0rcm~jr%w+?GuY)qvA$ox zXEVlr0u8b%Bsb|dxXPavqemhN5XU06)sH$?g643)?$t8uK*&$>lZz#co5)kmGBoNJ zjej<8>FD>q(JZzFW(zJt6mj7`RGadcNyVSKjgvb2?Zvyq-K|Kc-&&Y`jF6X+0{;>4 zUR80^h*6^k+A~coq%_So|LY^Jl!7FIDq*QpjUZ4Nm4z|S`oq3xM*ba&`J{%U={D7g z1%m6ZyUfvKL=X{-iy8?+D2gJaY`f+q+Qq;ek`QCZo+A7xK~(exM<1%S5yX>&Ji&b9 zjAttWoLLBpMDaBVzgEAEi)|hWez@DP6)4|LLCCS$Mo>t5Iyn2-iGDVfKtUq)50%ol zFGe4Cy~Yi2iztsZ#b4fdiT0;T(h|rZcq=f;gA@5$IxnIv$`z0#$wl?WQ#FP7gIQxu z)l&NPZ<_3@V>FfT9xFfdfA;bj#xRn5u0qR{3QXOf^HOfiQ*6q+qhm;QE6J%Ur*`8p z9UVV2tIGBPdTnLezs#ePIt9WU#+li!yzC6Hrn>~AYNzypsx$mR(1mJw#|V{Bf!=8aXnJyGM0{@bS-gYQ3=|86Bfsr>Sm5DD>*At%kX6|{`uxQAy?iv z76HB#f>ecZgrrxRfT)O5)Wb3dcD;i(OlZt2MX%vtwb0DF{z!D1@V>0XsPml&3BlA` zKE{X=8Y+@I@nV*UtdhMA;ejQy3xy6!>ICkl;$O@bS6%YjDACxH2ug7w?ome!A2Pz@v1Dl9gsVH_LP zCX`qqanj@@>(=c-#vRGE{K@fAgE+~Z&(fsHT!27`o1Z=S!on3l=*JM&Qjh4W752P2 z=DdJ^42nBlcFKj%=!NcA;V_V~u*9TwZ~4g-H%H-&+r&RdiFbvZuyNHsUP~d+k-(iX z$kNVWUS-A15rg)#>5H5+?bGRjT@*8;o^e)c`X_=Bya^1en*gW8A;|2hQo}ww1MDIY z{yeLE9wp?Ji7XLTz7hy6F1riruAvrzi!R7K{53w)dAfkS`>F^`(ht6sHS?SnJ{cpB zkKitNa%R-IB{<_Cdei4*yobP;HI7x|J`Ea~Gzt_n&kT8=V%=ibQR=Wkt z>~m~Z$;#vh?RDGdye@lE0m=(NnntMME^Xdla{XhR29!7O&f=u-*R^O0SIB>`@IiK} z6yYn?S~VuUx{i%~aEH|tTgpi_D@h`62O@mwPpV`t>Cs*K|96nQMsWA&pIB)jzce%w zA}J6c(SSe(u7uk9u3p+F`Yn+t{laSZb@fw?sKa zNUy;(w@_NC5i#a~k5}7SmsVjzd#6DC2K!W$QTU2bxY^MCdp*oi$27&IsT*nyPt|(@SepEfQ`ZcDiPtQAGIH^c_=<7J6 zfCCdiSr9~DqKk(v@IxqUb5Ptm&Udp z#w}%GLHL(PT<-nN+XPJZ6bMMGE`GrA>;p=e3&=Hva7n2PBx@J-N_5 z!3=0*0bAU?K3q40##gALpHM`XCg4C^=7CU8Jw! zFacj!;6K2>gM}(@o^SH!*ZUhhx!)}7uQEmWu*5YLua2Jp0Yv0_5WDa?&RR7t6%kho zPkEan-*O&5ehI+s`v4+3;Q#rLmv;hJ)Bm36m~kKSToyb2!ovr1##?TCNypd`)As5N z)&f`8m>*-^VEgZKJDK~e7@vyb@j0hNE!t#wPf5=Geu;RE2MKx2{9LBb2hpV{lI(Y> zBzf$LuFp0A{}JX#fy2T`w}X^d=Su}X7DqrFwj7>LIxsg)6bOjShA+|G2gZMI1NyOt zoxpd;wzY`<9O%7+L=mUhZ=Qo<#qMPyMafPo&f9MvMx5+hEJnQp9BF7Y%P1L?PulDe z(k>l!GZ(AOz#L0B*VG49%*3m-qykzdo0?5vKZ-g_u;25>@lIxd!QluWolX-<#1jzx-327c$O^h*{ScHo$JZl7tS?1xCu1;s6t%)|d>nP>> zoyvO}PVU77QjWRQe^E_NHH@U+*jm6QnkWv*3h!(HlsXxY6HLjXk0 z6wB;jli^TmIZ~D9BGCb1Z<)Z7}%SG+Ub zY>FO_*mR9z&)ca_$g&HjIkQ2Kg2O!8-9OKizI^Dhyd#KgnMBAb5R4XL%KheHW$#SG zH7W(R8H$bgJ%NSIt!<*~X5g!sYlB)|BKjPJ4!{EX*whKdeRH--ig?!Du2`P@^qJuw zYWpLQ?Ap zmJWU$5a8t^Y%~wq$0y}3IdW1_p9A^F6TnG6+npMaex5-aX#q4|DPPEaFF3o3c$OLD z*H`P5T>CulQ&ohYTO0gYNs|a;i;kW@uud$k^QdR-9GigIa_6?-@nx*=C0Z?RpF#8i z56L_$5P@scwOnxOOjN?4&-ne`_*Om4DqDjWNi!+U0O?yU*n-1?vky10FzI$aG2pJ` z!WYF|UXofa8}uKS%)EP;dgy!8|6}f~y4w2UtsUIm-Q7L7YjCGnaHqHxcQ5Y6-K99R zxKrFI#jVhyEl^6&=6}X}cRs*7a+QnhjL6t)t>2vUnY>~xF;z-G)K!$xQ@}!zx%Aot z*q*cNY)yT*R30P=Il1NO6P>QJn?NSI189hxG9k8Pr$?x^Q^>#;bAC!8S$zA-i*ep% zKFeyxdHX7Yask?scC~bYRwBeRx(}gd`k6%H>T!umqGZUVPYv zoAiW(;x9966@I88t>CPF+g!#^?1~lA3ps4Xdb$C8CAq(6SQ7Wj?a8apyhqK|#sVGP z-ZJ9g;>py&bB|mxZxbcuXyzSI3AJ`=2r{6b)t{;)c4T1jGTv3R=^noSe!0XY61|CH zRd8^nn)J9od7;!pZA!x?$LAy<>Qqha$|xi)kFl>(r*u;KDO>Ybs#EpkKBk@03{I}Z z3cT^pRcZ}h!U0wFW5?9Ftb!)FDDM%bH@~v3q!Wi@qX&2Azed$-hIiq6>e#OAO4)-$ zX$Bkh)oeG?IuVDrrFv-Ec2cz>UZS@Wh9cQD%0>iPa=HXY1hmQX+u#zb&$HO4=5rqn z2Z{GrD+hg|oz(NWeKwmBbm+_$(L#l=3s*pn^ytvBQCWN|WN9x<5?2&qDt*x(+U)8e zZX`QnS_TIh2H9A;g#I4zhycfzz?&bxUq~B;)FDU|S3+XWUvF4vg{eYVX^#7uL(#{i za=?-uOMJ4;PvZU>J5aqg6rx+W-FgH($xs*>b!^LN!whKSkGx0vpT9a^*ZFu#`>Jm& zBAqQ}gjh<#%BHEDZ>8s$ykY*g+JGO}&zFphX zgI(Qve3`ZU<#CAV)>KcQ(uFZvzlT}E7E7ri70S8&iA^(oa(9~*fPFg16+zdC&ppzv zXO56AAU@TvJ1CFr{;sqR#kEG|Dox;6>z=hfvw=R&rf5f}fBi)!hD4ec8fP-Yg7w8k zVMW&iS~*w!usWl_sUp&~A00i(y7p;)D1+b*y?}%(+yR77*Qg4|*HBN_Ldf`Sj01$2 zcbcL{uOZtnus2CKY(=fniepll%g#7JXB=8|31XQajV4&O8+wbgX>Q4AVBpZoD9HK4 z@MW+W-!tEK%%2OZP*v!&%@Q#WhWps6<%xC1X`Pei-rfFKk_;RFMphX}bruIb2u#J3 zV>#kp6V`~Y#0MXSC{>A;l4-AJ`#C&p9ERj&6sf;4m~Krf0|!5ow*#W;S$~e41m-8H zHXy=z172QdHtIy!CF~>~>hZvrH7pmR6eQOygPoc3x6C+&XAL%Pa1W8;l5!LKe{NlVEjHcY}-xMufc(A zQ*XS+9e0YISRF>AS`f$dp$12V8-zrwpsseEsRh!OF`^W-*3_4yQ&8q?CFX5}C(-gg zB%~r1IUvOB1JHG4K8dx8#AgEV(i@`c{2gcKn`!=oCQbfvGz+LkhdqwEtXOWTJK%9C zRrH!oq2cduJkXXXD36bY`g>55VUZ`IMvyT0AW7)b`USmGlz^lb$cxao)?<)j#2!*nM(bMIRwlTdPJ1e=?232#b#+3`eCA zaf1d{Cs;OK>#dyxbxg8k#A(JGBALsMBw-q!u^5e%}io9LT|{L5tz8%H%~*wfkq5p!_Il z6BD9#$K@7Xoo<42M^?g54{+7?3N|hPM?LoVFnTx8p^}qz8%>i@ za!bp3iO&q&{3EXzcjV$bK)OErPDvllhsmAml-#a$rCr$JCLp~DOaDcyMS#AH7?YoT z4=>8JT@G!W)=C2M=PueV1oYY#M?NZ9$(SK;9?JEj6#p&xXj1OT5Aw51mASMS{RM%0I-ST`|PEall3UrvHA|ZwbhMYh@BaU&~vS!r{1nX!`%|g7GlpQ#k z3;XdLm&o}bjl7#prkfL_Z)V0|G)|UM!ZS~=4X2^Z`uMSH!EubJ4G&&9=A0B}s)g=W zR7SrOUpBWpHUnkGN%wzmMqgjg3a6Yd=-Eqf9<}8RE8c9lN|`FS!bd%i3W#YzI1>rc zN}a%AX$SCO1py|0jJz;mm7h4*W>c?CG4?|iDSn(vr8Xn4WiH0Cu3dYGdGE=NrnO#~ zo&PT8UfAJla4CZy_FmC_UWnDWK;yr;F2K!iL@^vskB@n+) zgK$a{s)E~Eysp~FcN`w0h5Ew0lnRdCVLorJeeyZZwF@YT4Fek2yl|=#6)FQafYaS8 zv*K1Sdk^Tj~p!;si*%WNdiqTDY1n^ zJdwTY;liXp<*rd%UQf7vk&L%PraN2D0F7k7wyu#BUyi}=^NITeO-$c)(%VCf`iFF6 zQ&1aq^|pu1b13z5oKotEAfk`_NG^+6@54Y&qHfYDi2&}e*xrR7t5O(?stbZTk}|R_ zII1jCb64zVqs{h}sPk$%5{n~nbZq;p#%~Fp#ZX9?sCTb)J2BID1P@(AbXFYOXB;3! z0MF+chD;w|0Ik4XnqIKR=x>x z36pHr%*3`NTVp1zYahQdJAp6$?H|W_Hw?D23n*QwX9CoE3hKrf&Rn{8iGH~d|2&;Q@Rp0jIhVT@{NmZ}GpynKj4_}oepI<-l7*;a^cClM= z^GoB-cK&;0|i#|zuLUs0@Yzyc2+&gmyDVcep_T$~d zn4;`)Z}h~I>-T};+ItMI3$J@z zD=yDpG7E_nhv^;){{oJVm6<-^m-OK;Fmev`HRcq!4WamtGGh>=mh%~vaSUpWHE3V~ zzZKZ7_DZOs?+~MfhPZFlt^@N?{hr|MSYDZ`P8Y)tZth?XT0=Ne%MPk8|JOZ|#eUKR zyKIU2-QGa#>NJXvy9ZKh)dt(h`_>H^JdVvP*dWV#02Wt;q$pI2cy%#`$Gr50j~O$n z$A88HR9<;kM-R}T;zK+Q4WHTL!JlH$u*oMlAP_mLbs zp+5XV{CVz~vaCinXpe_7Nq5vWah|3QU4XxXn`3xJRl&?(eUltNh8!J{@_W4(__bTu z0LH({va7L3R6z*#TYALOpS?T1`V%jg;AMGcFK!jo2EqH$Z-H@cPAyX512eRim-O~H zjNw)aHNgZW)7d0Fr9jnV%Yd;)3FEN3Y$TP)4)P%kD2PbgSY$C+drWg2*5X5q^^wN-oK5SXi+XNroQ`(Os@?p{*fC_F z3S>nBuvE73pb@g(9D}PV9h>!Nhk`du|3j}&Ouh26Z>~mY+;`T>;@`@_r4#r$WdSJ` zwrbUU<&O)QpkI*_lKQGwGw=%ui#D@Q&%z}#ovsvA^m=$%Ao$>mO&__BR-WiTMFM}Y z8Xc!*wvP!v*da}?xNM)U)@<=1{EJPHG++7XBQEK?FfFrT3MM$Ywpa_3S5Q_|k?ft4 zgk@}u&GuvDaFP^7=&#EhM&`L`ByM|5P69CGz*~ z*UU9k3c2v-+@{k%P{%QAWa$SY%$Xr=iV=NDxc-A1q1=;N^-QEEMwj-t+`r(=ale(O z7-@NFuO9RE<>mKhNzPKe^6Vy+M+r^-Hq1hQCH%9$48v{>upntyh<%a&Fh3Sxp_zmG z{QK|6O%_1tY0e6u&XilhQ}$C%%Dn2SG!s*fjiC+!Cu^I^1 zpHFaZSc5%K7``5ymrK>S4>4!3(s-OXbVcuc_D~*BkY>y$L?X^z=TPBh0g;` z=vjhzVJzsmfygK1YZAyzzmggk5D0O{mLY;wA5t`_x@f+54ZH*G<$K58ucjcMn+Mt8 zcv=lNAAi&khOuJT3vwq)N}Fz1Nf~Ig4WkaNT*FHOFOBm8OH7lQe!1@zDN!!kKg&D~ zkLB2Fl&1m_ke@BMI3#Q#&Tv12GVVbcJfQj}0ch@jf$9f?6^o(0<&KtP7RG3Aj zJJUpgS$kynbTb?j@Yk@Y=nvU%w8)5IG!mTPtofg)men{I!i}kQ3AH0IH%>jlhu9LI zNsvi0pVytsp*d7EBVwP;POIn^+Lm1&K#czmsXG?1i%FtnNw+K}Xe&aWH!Ut?>W_Z2 zrj`g{S5U2^FKS3Dth{7uO0J)dlerg%1T(j5bi6@VyE!OeMwxvpu~ZUV*aFtXO8b4Y z^;!46F0(-AbBj?1om2~LLTz{Q{KqZ%YtZlg^sC#tZbmedf?4PwXheoN)uG~fkm^~w z(>KmSnl`OWHLfzSL4rHVQ$AFpPL@BM&uxe`*3Lv(x+zX}=233FcfWYA1`WIW>NGBa zOzV7#1H8Q*Zn`!(W@dIN0$-oRr$2G6u$@1Ws=UY6AnM3Vg{cTd{%4oGy?Mur^MFs@ z(b%yXi9I#589UN=ClKkbjwf#zb~fzev0RzX76t@FERbsqnnr+o@ z2gt=!J~9zW0F(b3fe)!^jDOZLD8SQ|i;k>S2nI1Z#1PaDx!JIrL;(i9oXIkYx|@pRhy)XP7HmGM&>y-?Rt|SXlTE4;!%J1C zOPQLfB`jXLCSTA324~rG+w3AsSqwYGQewa3ftjo0`L*`>#hVddNRe@PR-k=2#$?9S zK|X|MHwXVrAs3(N3aC?fE7#vMI&yUi;N&InjNN#OvTKw4!jV*iXgL-z39atk!wz?eV<2F)6R&05%1jyy&cHl;!zLB%URI^c-Z; zZ5X>VaTJf;vdwll%NTYZ!|1ty_=`8mUqP;SL{c;Ig3Qsk=67QPs`qu zmZmaeww=uyh;+y4?eh=2hz?N)I%h(c!W5cNp1%G2YcBOM{&Z$D@aRjzr|@(M6q;+< z7U>hI^HFmz_XIZIYdfEx{*y2DQ6?d-7)VPda7(QD zVj8#>v1iD(W&whL;&~j_z6>|5k5OOo;^^4)jxq&VfaMC{D~(M9Buh1+ue$rHc^Bpt z%1Y|%&TX?D+a3@-@5PhAuW24JCIdcOF29F}K4BxF_~~bdL8f^5^PO8!^lZEk20aUJ z$x%hIbd>LbaZtTBiXB2^9vPz4hyt059C?(5*;P2vt5iOOTJatH7sqoNQ}hB*rOAqU zwn^T;(jx#p8$yE8@3dRmOG*K{Bt96=x4)!}tvZ5qA|m3NWuqfb==Nc>iSDc20iFjcw7-%!pT=+Klen)MqfVN7 z0caQROm+@f!YIc#VL*MPXAGsai9u{ zY*Dahs0kp^V@W<=(Y+VbmIpYqIJyhR2UEsx`!g}~@@nL$Lur%*ngKENwGN{;44*bq zb=60s5T?=P0fDRq^8Co9;k4~ z+M<4JGMjXFcua^4>P~=>O0@fRXtX72lW*SvMIzwRC?Zo`@A$qS zK+OYnoXG}t(Ej-J$90Zz#>&xVe=$Kym2qqbD3g69Wo-j26(WONfA?Frk;%au{1`&^ zP>v0F=C}QthQmuTZPp_p>ZAci8Q!Mfox)GwFj|~yC6Bt+ya2**E8@N7C`Th0j(;0` z+W_z@zCF)=69dd2$CD|8icj$FPVchSLfL$_-4(}Q@?Co-ExMkEFaO$Z5kYM_s=cPDq~ zrrx4gr@w&?Kt&`_-yc{NWqjEIFaw+YADF(DW(bTeXbkWbdC92F4r58OA6%gi=YD%@ z|B4R-4%Rc2U*y-{xL4lWN&eXY6hOL9``Tq*4&W2&4+ycU9%8$paO+^{qSY0^>D0{W z52eJ)Tw`d<=a_}N4(+S%tESaR2I5tsp8>h|mAw+dDf4q!R=|>S3{u2>4$1XS35L}s zWlZ+sKaYEb;Uq`J*ACLChe~8^zy17I7jD5{@cmr*<7=yGws~-OT920|?dhn|bxBsm zPGj$nb1T(x(}&tz{E^oH-3D6jKb*^iU(YW|Zv)Fn-x&DTt6aafsNBr?)l(ywRFEMc z-#`4jZwl#8rzu8|hNVVhXxcnHPoFR<;V@7{K2Wh?@tzlsGCZ z+&Y>6*_;*nrkd*>_>2?=`ejH6qtRM{yVOM!;n$*w`WEkXY5SJdT|<&hC7zO59{4Lk zkDQ=TfBQu8pFiWg^p&yux(ADd5%N{x;tyb|$os)%2hmGw4Ra7%@a}uW@G4Xx zvD%_i*|yjUh(%F4!%iEk*o|%@rSFHOV)^Fv4$Q1@m~Nbwdr+_7nB3$^a!B^#hrD1$ z1ZmaSyKOJYB$7XgcA(`SIUnb^fd#{-mq+(8=U;6Qf`9T&GE_l)SbK* zTe_HDJGNK@=CGDoIahm_C2l}ZI#g;_PKAIrB|zZL#yg2)VBllMsLHU2q@rC0@|AlM zyI-apc&&y`Re?2;w{L_1>W5kkov<-gW^e6=dzK|X()mlK#6UZQry6POhk$4OTY1o$ zWh^@MaV(ta!IQUcRdfT60Jl#pE>y!zedH_dEYWtD_*gePF4mtb{7kLyPY>pUs{(F|hsXuR7c)9krQu+$& zWnGJ)U9wq})iJsj6bp?1uJCH7*fW2uxQcCOHlJ95ZHEvs2s~I1Ef*gWS_fu!srf~C z8uuVbZF|@61wkMI|7iuQJg?unh+xTV!J|L*<6Mrlod2833;~;YX$9o)>@S}1-n{^c z=5P59Z8By~3Yg}P{;|3mxwy{q{;obGD$JNQ)Lu{E4mlKkE@KT$%xllZh*%pf8-c#- z8T!=shA2bnZp z{&KrzsDlE{v<0-FxK-Ty3fvlPj)#>4{2GG(qDwH1#VT(fkiOgZ6**>_t^ zT?V^3{wBCwb>E{iApJpYo_r#6g8F&Df+-_-<_Q)FQxVn);MEPb36Mz52CTamd^q~Y zu#S?tHrst7aC)C~2)`w#H2lhvgKM+$6m_TNTinCYD7HoQ5TK7g-x|q(4e2x zjqNo&pgr}N&KqHMKot+OKYViu7)Y$0<)mPeSp^VdMjQ^0Sv16Cy~Y@id&yCM)UHJ# zRs3(5P~_lM9=5A&C27URD&OVPHp_l4Kn&8EAd>6Pa!C2mmi{mQUIMGB;`*$PzIL)J z)@Hk|>C>iA2GrU~K&7?mXVlauN28gbgJq$v^FQ7SH2Mr;uY9)x&tzXs4AiLdDqoXF z#ImM^o)0Ni-ZZBM4hLXK`71N@V>GBrr>>H%++I#k`o9k+K|j~>!^{DyVmARgdF;T& zm0kRD`=}P?o?U@^G*=_jRn9Q{!bItM@_<3cXj-VQ!1>f;N=vqhM;w>Zt%6L?TV6(h z^j6iGYB9<|l00+ydN682N&w@!&pY+U`hWXy?Em~;_?Y-DQ5_O?&oWm*#D#kiiI{n&qY2l6 zFuX4sVxefr!9xPqnwzC}HdYB=WR~o4rIE$Jw*jA>qu*V$#B>_lUBDmzL|=`WyD)G0YF7SFQy89*aa@+&yU_F0z1^xT=`i-v!^kkPc#G28MH#^*^Jz& zfIwc;!X~wo?Vi9Z_z}RR!P)o?7mMC@C1X>f?$u%0 z00y>+rY3tTJqW}_epHkY1ce^Nzrrw+rQq@MWg%V(z&LnNNit#3qg0jYFUQ{Ihg0UO zx?Wl{c#HSY+4y=abMV{sHTTU${5{s@4>Xtp5%qLkPz^eeKK9An2UK6t#@ZC{(HgL{ znvmYAFg1y)h{Sl=N&vV7lOth{2Wz?JO`^p?+0y&Zo?xZUf{20MSgyXSAuiw`qt1Td z2wTAsU=t}+?3a(vyPY`dznwzxUbl%MIW@8>_JGAp#0RwhW?ag{p-c|>}7 zZB&LE+21au3KqE=%}ArxG~-2n-YseTyReh|Zc`q)AoJ_;ZOA)z z`5N{9{dl%dS5aKgK(*rg(%AuNwUPX6L6#?R7Ab1#OxMD%;G7^p$@&ENo2jpTkCaD* zSqC&j*;AXye_K{bAG2ZS)(|1aVUxy}zPh5b%>r1ye*~`zbTZ@V0%J(7I`MN{13@Gv zY_10HPm8J;9Rc1I632mXUwcGaiunaPV&}2XIoyGFX(|8mLp0t?RZRbO(F5r%UVK}( zzVTM7y67c;mAO*q{^0V#uAm(;97Ws*VxE@H956qQG?moM1)3U>@f8Xiqm&4opZo44 zQ_)3+4ZFZRGJz>LeOkVp`L9bG=Hbqy&_p0Cu10@fm?VQ$)7L+ARFAQ1bVu;h?}Qz& ztNc;%&|w|@1XIn7>+J9-C2Uu6*k^i!8B6&Wj!-aXYc5fX%wRP6mq5g^4Z=r4AUx&^ zm%zd*^my4hL#|@pZpzB~H|?lK+uqlsL?QQe=H4F3Nq0!i5Ubu?4@kV0N0uEP;?Qjl zGLL@YzTFJc(I!bBHrbK!&JAHx9z;xv&|P2UD<`3pfBm{;X-%H}N1eYwNmafc-Gy)@ z!4-6sXG!m{>z-M2BtXZYkbMqPaYg?@R<6r1%>-xz%)nCzOB30+EkUrG98F+Wk4AGV z$81ldJa5z1|Jwo(j>G**0HXGVQy8wo>LBft5XO<6SGL0tRMe?=nGMZM$ zb=P$8H0vG(h-?&9O8vBP^gDaB`$S6px{8Mq%4XC-@-v(JB1OUvx?xG@>WK0Eylt*Q z=(OvR1-NP(RjyVwvTr34VnY~B9#ty*0L?s8;|1uSG|em+4qXH~SF)3@+|OTL<(aoM zjV)~C&3W}p+KZCDV1nirCn8AhcUj}||HlF-$Qo^2(<+h-N^z}!Xz6@Fuaz*FbqgFR z%#Cm2UA_5Trf77GnC9A~1R{&_XijY;0zfWX9RGqX)q}3?Gn)c!uAFP`G@P-CYyx4A zO|tpiHu96SBJnG+z-iv8(aoED5*CC}D2oHuS`zyn_b7Y7Q7zzG&~IxA zS7SU91=EeO(d-L_YY#Z}0KP817|5}>idCJm3SxH1KLZ~v=+_SEPEtxY^Ji|-_-IUJ18?Jvc8c1R;lh4Crn;#qUUUlsa^qP!~Ia|Y2Y(<`-HnH~~2?e?g2ps)%af&g}5olM6tk-3IAl68q{5jfo(Rfv{FW$@CUz*lYSNE?fD1- z66+F;nuX6(l{pWfah6endasTP7`h~U7x+irPxl%Z5=?vAKa$7@e8{!U_8(_bm%`tN z%tlpfP7RE|mG8d!dioDH3Qt_VGQ~=%-U`+*V#Ylx_eDrgiV<0y?h1aK=?$jU^2;Oyj44GVXgv=N#u){j*E~!qSm%=>PvM>9oWvgr zBVSkViW1+C&`*lG*`_`l0GHIIzT$jI-s%ReZ`-`xRr8}8L zZrNjx(JAHtRE~ehyGtIoU*MH6kl<0!*FX}8zaw1m3(@{@VvOdHtLC6_$ozfaqd%te z%qZ*PbTS`#65AFs+18$G2m{(@Ilyh$F_DiGD};MSSWCM`q-+ZImG09Dj7!*_VbaD% z)N}cWqb2nY_a>p66LPUe9dB|~J%>3$ILPRTNMtl+M(uK5=q8utI7_|(t*VpWY{6Z= z&8Z0kgojkhGsyZmb|UjwwLQ#Aj!uRv+eQMvI-R2lumsFEvhAk{!>s?Led6c`7>asX z=nF6RV+Le$0*N048G9A_;|(PgopX)dDcP>pVGlcKUZZJF`(eL|=A8EX1$`QROSKk1 z)59sROor%H+|SC^oQ4T$vg*--Ihg6_J9crFd8n1O477Q1T8ryiwCU(A>ye`t2rz?6 zxD87br7RUvDP4(&x-c=FCuTD0mPp@%9#Ve$B|;wO;YZ)qsPj@V>##PMeyZST*DF9Z zVFbs);L^B`*5akE>zqZW3GE6@kl#MYfrupfy_2xC_H(TYIlot&)*TNcC2_*5qgBXe zsq&aSN?|5#SL^F{pknb1i{#x}v=a{na&T{9z`)JKG87;bItkN{7)?LGL7OoKCZS>r zcC50J9`eJJBx&3V#jWG>)+lkRa*?*UgUeRQEyx2yaOvp24Z`xmw2klTh?+(m1EN4*ieEy&%@ddVFc}499^1`&>QoKvhlEUz<1n6C% zX7K^YJ^RhG9IhyEvZ@&Lj7QZ!Qbwa7@b{<6#vUv>4><1V%4@3NFNj#MjAVs2U!+1C zNr^*x2&(za^3~WzV~`S~M)aQ(AF6V!*}vX=yFc}|f$o-?n1t$=%Mr@5=2*;dEgYk% zdBt+{7*@fcFW?Hgp6t}yl&F)RJtQ2?Oj9E*#wir9(f6Q2N+5UO5X>8giT@((zh5Kl z-B@HX-h*(w*nDYdHm&X(N`(-fRMoRLxm@IF5qs(9R~c9Q{f*c7q?25a8Vq}$H2oe` zPK%!3N_wNkIMC~V z!zwcc<~Fn}AwT@Qn8H6NQqh+AC5w1HDYx;9J#j3$xXO=xX3&-?4If=JluV`mN$eVh z=e@rfbT!4n)X0>$(`{@CWrSePBy|A^eIQHaK~9#Inr}@QZLs*vV{ANy(-yvCP2Q_3 zN)mnTsDKKGJ0sF7p-Q0<5$~(QdgCubB;2@!A;T4HM1_hWkP;WE)q!d$#vIe|bsSsk z%T@E@QOKr0SkYXy`BEiU(1ek@fb?MGY~`*sBkVY~A~8JC;h}kX$hVHhQ$Cj<504jC zQ8E3r_GozT0pl9HaCMB(UyNW^7nfbS7#ZfYD8GenM=Gn72u2-uXoN}OG1(nO-Q+Fe z1NbA0FjQi6c}pN(WPO~=_%=kP13Ww}B^5`p0=K;EQ-$rSd3l)jOGJhuUo)=bd0*;YtOfFvuP;U_W`!e$c-O;w^b9%& zX7{Nk&`2m?=ECmWYDAjRmAKgFTsk*bzRv^m9ZkKbcm=W_?8FCMk&57d<6CaGd` zFB6J2PHBcAL3trG6uukC$0V{~6$ofJv5hEbXlGn>uf9WpJS6IR6~vCRhK9WGA6#*h7!E@?^CGTe;pl%F#^(@}v}lH>5`|=mT)%c=#-l!y6}EA#BaZpN%&r z`27OU{f$!9;GkR!s>20^!uAVQsIAg;q5TNLGh zTeMqdGZ}*gl%96%(tc6#s%tO%IVAQ?LEts&AN=#%yVXc7vMsS@!2s&qwBxjXvCCRz zr2}VWj1q%|vhs^)Gz$9Nm2TF?n)kD~Y-)dfgq4 zMpnjTcc9OY_0hL_&d}h5-!d~5dd2X^i_3P5B^rFpmhv=sJ*rTy(RDTz0v}SctXs~9$lUkQiQdz1qy3h&5j~vge>Gimcaqsa zV2CZ>Pnjq_4%{(KTX$!96?{D(B0=KdBbM$NYfp(Q`Z#X|>8Uf`q98rr15p#%$VD>c zRST()P*YN2(j$ayUqvp^B1mvpm67GJi*WOO`{jkEECRdjWf6rE&O{=1U1IOLP&J($ zV-b67Av?s>TIp3mAf2N#Aa=tMB@a0hhArl&cU`jyc0h}5+86ot=7pI@6%^?*R=SXY zax#z-C1HX1eh!M}el3K}I8$}bfIF!|m&%@J_uiv!lmNMU77MvRc;owODgh7EfaIiA zV@QKQ0frEvP`DeXQ4nV(mq^zp`Me$KQwGM;Ysgp+$h1!9)aH!!`9p02q|VYuYB;l# zv^7*U^;%V%${8k={ahI&-Ol@-@jjR^1KdvWFU&T*=se`AUam_CQa{L@&uC+Z_Tsud z;M4ix8}7xSX8W9&8mT8MT(vFuY&AAcgy@Pze${aYo|Hd4&2ot=wG5Ffs%v4hwNSd( zPxE%Asot56O4j+?hWCTV-P~h$+(r~C&W4bZ>Y_6hs0>VHzMg{&KE|x5d}>VyV7Eb_ zoh5id0Cg~k^f47!#Maoo;|cZloH7*7u)3?6Zw1b{0S#I#2nsmgVZ#|`J2+i1~g5L-3+-D81Q3Txz5LMA+?>6=p znpPnPU9%K(k>mvvrIaP!XgUqkDTUjZvUczvZQ9Sc$;URhq@zHcq{MU9j)^i3)Ne&k z?522Y>3@Ya$KMW;%RhSf1$;%+@gPDdW-`(pDuA7w01-DTp<`#{WL(Rl51aMJh50+s zLb-V&<%R{!rbNzf7p9Baboq!jZ#cSBg7%{w{@2+Ccsy!D5?)9dIb zlLG;$R=A`` z`i_w(8dp*=0XJ%0yrTB_REL06={z+9?!@+im;ekCrVqt4N{kB}zX@;M%T~#$?DA&q z$)akev>P6V@e_%S)Vud-G9Ei%@n&FQ@9eqfm}iVK#%M@}-u1aa$9zUJ>tVYvWy6x} zz+L;A@H()(2p91jS`phl3_^KQyxSanrnz*8xa#R^Yv9=fPlIRod2X-#RSG?Kz4Kgz zyX7Bg0$r~MB@{7kmKKh%2gXu==!SSPZ%Q!z#bLRMRz$qA!znP^X^*;dh@E0lDr-hH zx?v-hKF8vpTvF~_Y__h~^X6s7$Hg-UjD;jKQFOZBVVG0y@>A;HIk}gGmpOu6<#PRb z1W<_((swio)7ufS-^Ix}fl8w4KNgz2aMq+)zA<0&iiHRyXM8j4)=jPcc02cM--^JJ8pKTrV zPX5{zKEjcYYfbv3^C>B$t>eluutI5B5Z$;#lf&pYV@1H>$e34-#0WewFbol};Yx@E z%3rPJDC~~!^NNbrZX2Hl@-8XVKZuyF1l7JXm%MoM&Xg^1t!MSq(5Ih0Sk7;bSlzG~ z2|D4K;~U4x#X5T?x*TB|L0LAtFusuDnUuHJn?z7`(~_vjb?8==$JwK$MY&cM$wBG8 z7I!d-Y;2GW9G5>!DQqo|IBZ&CHFMVA7{)(>&sGTaP@Ia@04ru3G$mdE_ z31jQRO@=zUvoZeI=C}C_nE$ddXT*v-qh()LajNMmpPW|J!1#1=A40ugB*ks0uWYGc zmE~!?7+bMn1%bormNFULce=Hd1aY+7Xee-G;;}y9igvFstXKy`LQ;);R%%)A%L4xP zHS;xV@mZIB!Uzpgg+=s@g8^+|YMaI(uut|GaAR=h;(KOXpI);4iJ$_-Po9*v>=Z|R zLkXjOuUH+UJJw{Z?vFG202@sg%HS_U9Xja{mJwo5Ink8|ZEL1*s6>zFp*?n7B%Byn zukldGsX+5f{bg&tb-_DUBJT)`uO%JVZ@kNN2|le&C>={RW~ja5@1(tdQeCFU@`G2A zC%lQJRNTs$!$*=y?cqg^7#}xiKY?z*8S)G)l_^qEfbLSIZ3|{Ci3Et*H4bQ@$VLJg zUm2C_XwfGFO!=B^cUvb10-pMvZ~6j+D@k%zdM@ReMX4`kRYyr}`|lQv?6lH?k(b-? z>W}7cR&hnA3(dA-jc|gHqlc5lIYmcg!n@nFoUe}!oY01L)%P;HXF^Aw9){-ET_#E* z*VrBvYTegt`b=JJz(%!X$s1V84~x6*l`htLQHJr;lm-tujr$%1yw4J2IQ2!Jvr>8{ ziUoS6OTQxjvcCT?5lxXlhHRnic~qBoLd%s?5H0s|x!L5?UKezeUX<~O(Ki05j}MxD zzaDiU(tGZ*4$sfBLNC|ie-;XrZd5Qkw;dX^YqKOH3CUpNx~UAH2>SEy+{z*X8U`Y; zCqoEbI3`i@Y8htXfZB#`V>WYRW1`4VFI+fx8Sk_<<3q$`{8&a}a5h@fnnzsSwZ26!?p*197`FjuNlz zYKdaw!%4ze++0jz*+>NkE2kyl%iwn%b@0ja-Lbn7^0BB zCzDf{;>(P`3Jr9vHqYBaG?inJ&SI~h&-Z*jl!uO7byX1yBHg3|73wouA&$8qd$MUf4F!w_) z+wj^1wA<{BB!Xg=zPmxP5oS7e0&Kh6@80}FmcuT@kXE%7rv0gJb@Ivs$lM}2_X!}c zOiE~9TBE*nnCOhx_^EO+6PhYe87#2&#V?KVB$A9VayGYdLz7g|D&7iS0M^uZQ*vu2 z9oxs`*JJI<}6~g|=O`Y1{#ON${__&rfxB!~Jv9*L%~r>sMZgDYRD`Z;f{&vp!tO!1HI z7kqP`P^Kcs;gh);QgDQx*MWQ@X)sQW0tAtRJnYr?km(W+izBKmxI`gJZ){R7lV$Ji zR4#^~JGLg6Di9we^QGxIF~Q}LZvrKR0F8S{31zAI3w}mo7-WV6X2=Z7^ni2?yr~&5 zFv^B4NU#UQ_g%I?!UglziF6otxZfQFq}IB>c$4PgnEiK?=FCTb-l_uX#IuvCR;Kg9 zta4!3(&g)H<`!MuTXmiRjj9EU)j#R*FMu(R5j&(2>OWI!aAaQxZx)()Ofsy#fz$}J z-)B<5>TZv9|L|`5_r5 z!j<$(&{sD9?7?;xf1vBg!M=P*z1pP?UXT_%2FY&N^9?@n-2o5-UhH1R0d!ftBn8v` zB)1=kFZBOweSTz}`5**e-JC$gL3Vz*+WG08>t&y7b72b0>5J{J5g zj?AbBD2n?&U}Cfh*7j;$nSNu`kP!bS^Qw?ErUi|wZ@v&KF^Dsnz})uYtIZgq0W1Z* z-0x)HzOoTlSr^7!cKW4E*S>m?7zty zsIxLLHtt^5bxwjZxl=OO^zP0yLN*D_WwN@bL>Wl})Axoce>L}RA92Ex>%`FWj+)v> zg>u8?4|r#~ReCIU;n&j!E|*byH@*8hZI25L^VOII>d#0BS=WNDNi$xH6NP#=W?%NG zUP?OA9gTNmMgkp1LJi2|tu+H#zw;kB=jnRP*3eCa`JfbwBd+%{U}^wd z!u6-N>8TculEWDL`|7K-^yYBOxx|Gf_zEdxK73Pqt#l&U+qd2*4Jf(}*{*C}5m{?& z3i zCG6X+_%|!rT%0#Pf(vPVOE)$H{9U43FpXS2r=_nbip+vH93t`riZnE?{?o>rr8PSj z$b=TZ|K*H}9Rt2Gy5=@y9-}`cR8KVH!Z+vhD-m>(KDocU#>JV{+fS|)yP~AH99!iNU#4WGT^>In+W%G^U6+Oryp~u^>ps6=kVs6Zn5}mPXfh)c&{Koos zg@uuH!#0N+x|H2_$iQ^(W4@bu{`2L>&tK-;53IxgL)Tdb#lZwj8h3Yhcb8zn-QC@S zySoGpu)qSrAy^3RzPP))2M-nq&Mj}x1*g(*?B{%% zNQj+j-+CD04<6#q!yH)o)yI$yD}C`u;r_!*xHDvxPoJ0PW5#` zHL+HTb&;5Bd06zmu3@Y;S#!Fk1&+xnD<5@j;|5%q-pFW;HTax>?~rU4!@a^a5>F8D zmqZ-gj{Ck6gRE0DCLe}h`c>kAt{S{V^n@SSB}eMuaLa6c6Ew@9{d4fO7_ea(;|Gp7 zy3?-l#~t*r;}VXDX!>rz?b>F$gS_4LMMkg20j}DS@EFnYmCbv+`ECLpUr5uO-I}zi zmeOdpAy40jwdLdG?52N?KaI6+lIHB6PB)D`RI|$#M-eG7yHH9_2cfqu!q|6Q_Hj0WoxBfd~LiDAm zGCpWXlTaECnY^}+quPrwt-39GQgrcHKr8tum&4d%9PX;A+KW}!*dX%EdImz1m_#{Y z%u-?<(QL>Wie5%>i09J6C>@oP_G?W$44VntIYN-KmU}D1oKfm^U9HF&2ftAz9(gWf z-VD%0_g;i6+3p-$k^st*o(EPHRCP)9iSTp}zwE~K1#}Q&dP7Wh1azf5nqoyJro;f* z2O@{ujN1Ph+A&8kFhEx<-SF2Evbq_RYm@LfV;cSD{a{kD67#PU+4KijpoOoBC@)%3 zo>beMbi8n1^6VbwqzX0^VI#Hy`V^mg-$LR^pqF2sTc{xHN+-=UfMWWFeG#FQX3{P% zcD852isf@{(kJ zhZ#fZ9`I)nu7b?N6s8e2SxG~^&*%*Xe|VM#iA<}HLr2?WR#=!6Gs!R@fbEo*;tSQw$k$+q^8Kl0_P7azp(1Lm%PiifgV3TAuTv4}W!`MFVonp= zU{op4B-%9L_}ql&u1;Kqio9;HPGXn?d~E-0;eog2s{bJ|hEBnn|s8BYa+!c4zKy(2{DLw0>ylB3uI?Zls4_>2(u8Rrz9u$L?wrKGQCI zhR*;*7t7K?_Dc>s@iEfwVPy%cB%dL~S`u-r`BW}__iTSNUz(n27dbu{o8vYaRn|hn z8ey8Rlis_6ePJ_c79$JtyNE zV^mkcQMB%%0NOaZ2^^opHXa?ICW@6mjZnbB$VqPI1jitS7E2D#`};1wBOo(Tt7W{> zeNX*TBeaQwTcfK&(uD6^4yq2yya58ULNhXC$nylJYWV>ec*6p3DS7U=q2fuHzfQqw z$x;$}g5@aPiX&j`8%QV{=8(z;ye>(hjAFyX~%gI2kh0WP`gDz@dzn z2>a4=GlNV+_Hj_N&OM^0{jO2N_F421#mK`?YdR6d=OjD) zyL?Dou?33pjHjVdsIrzh^HtdK6YFjsyn|hgCP`xC+U9fS{eS zL32wPoV1M|bLDE7k)lKWY5agn*?sRbgvVsC&~HK_n^~k#@!FHfFyp z0C(2tRCx}7;F0huhn$7 zoW$lybO<4m!rSc0SXi{Uhw0&#FegkVU>)Jr$^_4IlLIOXTg2hrY*HV0~^jgk6!7;ipvyRE& zE)=$1RMz+JVRUe(zxvL7U6mGk3se#&o$oXVcR7xTL?nf4hc^{|5emE!>Z$e&ha~7k z+KNC|5whk8au;6y{lwiZZNWG#(SpEb7n$Bq^jk-$?tw+??~u4%pgz#G$ftd10*7j| z`Z?0)MzT-Qjy3ZBn8r}3Mf;$)dEH>e+5TQ0sAP)Mm@?>=ck&GMR#YJ}I&P9N;C352$O!%azrayIHl`i`MOnqq56s&)n zK?kX*gerWf&tjTnlyvI@T)$ z3V!(E;IwASHSNLSsNnOVYgRQXSR#Pmy@CW<*}S%xD;XEfoJ09VVkn2Q>xvOcMF*`Q z&9g>~LE9+xS^CNEJbPR8?UNb=2c<4)Kg>x_xWzC6;xG)lUu4yEA#i9$OI9`K63PwX z^s86OIZEm(5~G|)Q*qGBENXnZKZmInm{F!AOE}v_7LRiu8m3%t`eRtNqL!1B)T@#v zDTd>*YtW!KW=cri7BM{>s{bUU*#RmC^FoGyKh*dH+>CR8B7Xx=P*{0mS?0(>1xf4Z zqGFQ~tSlVj*|ZGB!=4IINkdZyN}uT-D{(2BO5PUvIHM*nt>bNvABHzg)U$jsme9E; zW8OwzR`g+N;~2G+5VC1O`vYmQaQoXfG*J?(5iWMM8xK0K~ea>EGMUcBUqF#)^06Fq{n0^&*z1jQkNqmsc+}=myvg69<<8%{ zx$~`CC(G!_kC+C}T544ttSiAiM;`@9xE*25$P-u$^woHd0eizDf81Ky7puI#i_qTI zlg#IuIa)iW<$9Qa>(_doH!tsRg@pz??F8OmZeLX@mh7+;|A8$%%$rGp2h{a6B=4w6 z9E3sy_KnjGiTAc&UnH`(KXG~AlwSO;PtD4>_k(NC%}-ng2b4NCS*^>0RK}MkAq5~Y zS5+!X$X{p4U%%Hpeuu*$3%mgb2!bIl-EIJ^+0T5Zp5U}a_7b7GGmBK3O+~i01F(M8 zy`6cTf3XQCU$~avrfp{N@9oeWJs>k*JdgCW`5Iba{m#rzt7dN*Hk$da@Q@It`^DrB zHQ6-+$B3D?=k@-l=%IJS$XY18cyUS@nXQ?RGkVW1ad+<}F6vMB2QXo{&K|gnJc<$` zzoYoTA2iP|^3aPe3O1~GHpxqOf@Ol@XK+C3GRBA?XLaEovX*6JG1iOlT}$v;%Qv0T zAEGe+ww6fbZ~m!1-mkgakGZ-%Pm5guivIe~F_Y2j#itoz=FL_A)KWR^M}E9K=+nKs z#m$YLr(Xu6}J$lUD8x;Y3E(NP4H(1IMqaa&P7UO14@jo z`=*1md>nBv=BoTh|m|1 z-shswN^Pa63-;p&$e+k;tq}Ru3?mp*qL+%kk{JhWq~Sm=HFU5IG2|4IveXaGT1Jg% z-|>?v_>Kx}Z{06X9YvRjYyUP$%c+iz4IgB{J#kEym+6RL+eovD^OSQWJnWK33^iS# z;R9r5*r&z2jToIFvJEBLP#2?7A4170+?D=;q~~@MuK7?M*eer?e+wYbOg0JWHa{6m zc)iK!*+|+V^l2Opb9;0#(hFU!LcbPzv)u^tmjc$(o`o-% zO?~^? zRJ{_f*NI*GMPqd}Cxxy3n98O1L?aDbjQ>kgUj`{kI!mM=(Y!Ona1vJ1JV4AIA4=14 zZ@DSuqQJyt*;`Rm>UA~$RUh2MJ$gFz8$^M=rSPDxDM^jA8+s9Vn3TriBTYN(H3Vzk z@C>e?_YdCrZjmk==F<^GYi1a2ih8sl|D2K(_=A;{y^XJ4(yTc|!)E*KV*8ITl~Sm_ zt5F?he!e(eTLS@y^*&fidfI;7+~y7ZLh$M#_RoXAa~Jin?3nIz%lQ2rIHEJy1ozyN zen=9F)IL&aD{S|pm%3{PGoGLg?P|^|Kzx`RmIZ`!L7gjh;bTqKY-J}nS6eIP5OPv} zDW%1Eez z*j)~h z*ZiEMhXHI8)@=Q!Y7~{RsAteN9jMNDDR??*P>n)CNU$k#Ju>&)Fq`xrE42QN zHX1m@MprnmOxx^o9YNH>7A>&IwHZI_L9evv{M%^gn|I#dzQ54q51^FYbP{%Uy(=yi zYyVQ31|{cK4~J77T(s12D4mzm3UD^_y*5;%qV$9S(|{s5vXp9Ka>*L@(b3Z+C~^P3 zi3+DJZruyGM+oj<53RfohAZMo^e0?#ot7wX>5B=t zwMQ+v(e5ri976SouM_qhvbtcFD(6|&c0px(@*5bu80;8tKFw7?*LKNe1FUWI47bS0 z4J-;#^KL^~nl$lmx^;HmKGms?!bli%YdBq^gxL(M8vnG0OHoGr;gEf{IbikqqPMm| zN7g2rV)`UC?=MYufXPDt%_dA)BGxDLx0Q5B3h7W%>n_Jr zXpcS!k1#d3!$AA6`OBy3OLzBb1ar2R2>KHLFh@0wYx8dORN_hD-y+M35LTVfR~GS6 zN`B(Waa^JTCgh~gk*v$52p;j8&G^0{mK*8SdwhM!#4f>GV0H zVaa&$J$=Xx{jJ}=RlgMEP@@fE^&-GWyHFcFCN>DsQa*Lo7X{a1HK6%q<4yzg>tRGz zw4Pq&2lg)xVoXRTz?TU?9XM|6>o8*pHH#p`g9>UxToq#IDBUg5#xLXq3?uBSIQ5r zl5y(G?%Dk)guxdX7Vgl7W?twc1-Q4xvl2fo9g9g&It}!=c44wxb5?4ie+Cqp7cxW` z!g2|zP%IGRC6zb=`^XTAMNgdvMn$P*c6}MLmZv+ijzmst-Z%P4gngzd?nXjU0LM=jIMgbY{RZg&NDAths&FY8?ysIq7!M8esW zJI_D2Agc@SbR-hG9UVhAiuA9JL^ghYGUO~)lU^k|QKsYUAUU~#EM^P*CfB#~ef~mt zm@*&Y8G(!tW@V&^6P2e3wUO;+!o8O?Jg&GVX>-l>nPEb=oo5%F93WC@o0u8Zj3RGQhl@Gttyc!|6N+XH?I-3-{;6UM~t#Wcxu6(+{WZtrq z!cTHQ*>eC;{49Z8%4Vg`a4Fxh#RRN{O4gue_(X3zEy1VSzUCG-Fi-jQ(to*w21(o` zqL+|>Ab1y@qX!1UwQm6yLS6uO+A1?C61^KriDK8zWMwxaBf%CD0;ZJN2eM#q-X5Ap zsn40zg7#}@&3+ui<^?k(cV?Yk+TdQrzh0cn1wV6ariNsN{dhOP!ztu;rDyj8k121wZ_yyJT*tL?W2XJNH z9ZnH!M)T{#q_CJKTQig0GFgci(g__rTIb?Nn`9?Y&K0q7P|2z3H}YXH>VjyCt-~;3 zvLx2=L?cKL(|{;9@oLjsBD{Ll-wSJ)wGOj>tTliTffFYZ0K#BowVk1udAO0g15og; z?*vz!E0Jh=1%z-FXr$+FpQVKJ=0IZ~SG8~W4}Z6fQ;Pvs+08v)dGj@UjZH$R<@vLB z=H5{RP^m@BIYAH!0g=R2Glfy=c7wiT{_#{F=I?d`$%&K&4$?=vc89k)F+h>sASQrI;F%-Lo1+GgNK<;R zuo=5fwC>R54x`eLg-nd>WZ5>{sSt05s330dA7%!^N2 zBKU1^sK~6FO|fcCAvz=DsRYaBS(9}2k1Cq)@#7_AE){qw0+cw)bx}|?RhZ3g^rv_u zVmPx8%&S50T==<;jNvKNQc0{5KrOQ)PJ@lR1xeyirzD}us0yc1?(toVF&Dt-AXh3Y zYX1)^1}7_~{dw3S7w9mKKwORn>7cp(Sc8kDbl0YfpUcMqAk0V91j?BX?$LC}IzL14 zWd++>Z#uDt7b4TF!YP4rzeqgUnj@+g!CZ@V5ChmUYW*7c9>(kf_xqjH z>4#a+a8F9|o%jF!Rs%$!eoDL+h{NiRY^N#I5bzP@fk~{+p`0}Zs}b4c|Z#``(U$vS{+Z) z{ec-7XEeUCok-|>`7hnR&bc#YwH>;madc_n*Hq;$T6oAp8Kp`3oCJuwy>b{SK1?A) zRSVxDzk_h=bdcPmJpc0C9$qGbD#qyh75Q8L``h+=03ZvblEldzT|Fr6Q-Jr)P|4&W zajIJ%Rx5T7hZ7i(D@6HKw(Vh8t5aS!6&kzZxj~iExNh?$HK6Mu^IFyE z`dn^>e8Rxm*xp_cl%7skry~leGJ!8AnW4qi?O0pvei`o}fsAgtfG~d+RxG6^lz{3v zAMZYA?%rb_>0pH!FSFXy>#+q#CYKGyH|f-fyWI$al@5Zsu_AitePgvj*cDV~SLpUp zgm*grHpF!KG%sDn6eOR{P{`1A2XNQ;(nGTl3yuvlKOsr`z}V>z9;m5Kgk&mwY7T*n z$Sv_f0o;iAnO2NnhBwZ}>l>-FW;2zPLr~^I=ojEEYsBLSz)cS9pkSUml+S#thyVbi_q6u+VAVGZu0?8Xfgr*~Ndp(&&)sJe59F z5@Hg0*$Ad4aSe<(81KtPLsYT9&as!seNIe~J$Ep@(>A4z+oSSZtTfSPjMhgWQI~3M z-lmP+k#qVipudK8YBu%$DE2Pf>pCDn4Mq>Tk0UE4HDk-K7t8BTXc$2W)Xu*B#5%XI z2n6!|yE+01+-C|~ZGVWCTV+Du4^7{qXkOYQ^Z(qZvpGD>S?FQ*{WI$W+1or)^#@Ov zw$60o5Siae3S?;^a2&SR&C1c5xiA{k+wpzB4Wa#pGIlX^=~uIcCg5fY1j(%w7qDr>;t&X=reh(!&(UVC}7Ch+z|=!_!iBKkL!lq>GB6= zl8Jm zTP;|k;-JAMZVtqZ z(V=iWNDdalJRM^Yvtuo^U+8wOna=Nzg%1+96TFGi%|yv8Zm|R}f8XbuPIwE4$)~K9 zAk4fV^&(*Sm@%fNWaxqc&PmZr8k_BRV5#)RHss!xlo35InkDN4T>79Repe@Sebl;) zNRnBENVc50C+0WJc)#0&#fqi9~((JFl)9=C_ddYq7WC54Wn8O>2otGcsh@^o{Ek{n7ZqZ~m&F zWhHfR1NtA8$E#Ku#}OX3Oo`xLWD~P9?<<*k;MAY(8JrkpL=-UV@lNp~KkL4|1Z+22 z4Sp?q|C`UiSCE55X;uxj+|;0~Le(Dp8+_n;yL=Ds;2%mXt+WZj8t=%hkv1-%lO@M{r;LG5D9fPzo zWJybO>K`_JC&17=LB^d^)@Ku-_=uaqP~6zK}-hbH+leVughic zZ|5*K{bDcu?_fj!FQ|F0ez}K_*k0Q~t1)>VO^bz%Px9gP(h*r_)_xOuNg#mUz5CP0 zad`WwYI{xKw_HisGRD#5KYGTXx4eXlc@>5L&=|6yT=;$lk#>6letAUzz+xuky9v+6p zRHBCSaz94FQ?tu+v!~!5n#w~Jrr%LYyRM?D|6Holoo7a3%Hf3q7$4Q;s(tI(E{(eL zhaC-tpZ;|41#WpwCUYi=KR(&d6yN4$oYo!TQ6%wN{p2uo*n5#_rHLij4?2f$T+NimS4SORC<1I zm04W#x>nPKe#MB_sG~{wTyVLI9@5IYo;4YM(P@SLDKx{+cVnN-=c~|$_yV3I@Z=M> z*e;?fh0E|N!s#Hp7ZXH zQ8VeiO0=p3Ay=+r^=$2#SyX1r&&u0>mAx%ZNCcotI^5Dd$cxsPMzMNW8oRhn&;xg% zAH9|=wh3^FyqGTD+!A)sUqI+HeX5~x>jK2xz?RfhBsiO-Y*m%u8~TvBMF&jpgp^h} zD3@s~c=?rbK?vcSl%&5ZQVHH#XrN{zpRO}wr!(AGYHZnY61G>-V{!C$@GRsc8O^Ax zfN;^3aGS(W>iLL@g?#bTTrH03=_H>%us!d8vH;%f;mK4+S)9nCZQ@|b2yL@qRZ6bG zpj07>e-(N4;<31+GJZ0kvq5}6rVA1d0ry-XY znbT|gSnIN4C%U`YpUbO@`UV0n;Dt3(A8l9scI((zYw;&e&@F!mH>(oaYD(6J!23`~ z{>g|a8dlV zl%qc=NLPXdWDkp54h!c-MDd+*Hl|Wp-%iFW2qHAKuC%l-f9h4qRQ4sXPzkceVl86( z^uLq(|3v$Gbv1k53N;;m<|%!=C(>}-N#WY)^oZ&ehGILnr~VXW`T$#*a7%_9y&g%y z&Fg`{4%IWr>b(1K6ydDD8ROg+W`52dYtmL)P83z}Gq(g&1{hzji@>4vZuXhB>tjdK z$UTBdsF`Cse}0_=#>+ej%S_t;GQqXrIXi45y^JJba9|GFPU;h()QkeJqIDu!0wyLw z-EN^njOtAD9g3QDi#f;>GCz;f;(w|GfAtKk@UR&~K!JQej2K61a4LbAha}-0U~$zo z6Q%r=&N|X-hkA%nm_%P`t>V%&G2x3XY(x6{bHFr zF`-ebfaqGJHSDya+&GwDGwm#MONs*zUd@=WW_X)7h{D1;hXQymEeE{u40Pc%nndHy z#HEKwM4Od%*X`CzGLw9;7s4P$g8JQN*O{zC&hJ!aowiOA`<$OPp|+Fli#I+OVD8(U zqfla)0QZX=;C!G+vST6{fvBOJ6$oL}`WDHUP+W|0*G)8X3M1!_?ourYDUHR}*rxqj z?|j;GcC8or2r=>7F{2RXf^Nu*bDJ~%kmD|OtpMA}uVe|!(n?KsOoa%qO`hslhG7pFSuYqxI7_<1Fh{WCYya%yeR zi&NT}!g)#;_})iZdnFdGi#vrXN)L&!LIR<@W1x=3l9VHa-8CSQF{V5IQ8D@JaUf3% zWS%m}Dq6GAXtrdesqg@H=c<<7@&+23u7onoGaCRdK|;Zxa-N@$P5B}s!#2Sdle*YYw=kH4I*XEIfm<* zJ#9feGQ}=s0lp(oc3_*vstVe-Ex3#X9V3h{gJ_{LO48cI68p$<->x3N?D)b^P{Rh# z9sD`MaqmW8f&aGJCXH<-9C1vk%~Xw5$%k;~RuafX zr9l`$*~wW)Tf-j9^p_nk#XKScy&L(CouOFg9uhoz4!AHT|YCnXzoMS18$lm-)c#;j>{nMO#LHDINL#%qkHc#36+KV?N= zTdRDj=kV@^_4XUi&I+S=P-&y5g!XTkgg$eh|n&EfpUeY*oU zuED67l`>qi@dUlqpRu~<~pB9lG@?YNP*Owyu>4kvalo6%xU z%*AHIMjJ_;!*yfz#`G*@Vh)(02TH{3s;$CgbdUz-gM#&_6u-=VyGh>z-1l76N9XGX zzPMF>y$7M*2;1}Xb;li;YQ{D1WFcb`(19ZH=Fr&zSjOFjEikl<3&rEG)Vfwe;NGJR zWSV5nDuSq|akVAEg%dJ)5};L32dt~n_SO<<8W*(FPLSe5W$-MmnAb+|a2i)*xDWr5 znPsgLK!7?9yYMbPM8+03L4f_Ub9&qV0OByYq14aQ^&!$Ne`w?xh#q#9wzWC--Pb!) zq)Ln(fS7|^gdp=?UbBcE1@CxrEt^e^>B~7fQR;M_I_u?`UCPqB zip3_*#ImV)5WR_={Yf8By>rt|YSVSSa>D0G=Eee*-VR>9D_;zZPOa0dy%ct1%daTD z&vQZ$Hm5c-q4yB=&75BjhPl4G0xR?7(UJ#=&EcRCrx59lPitxAb=5X{L#Q@Ql^!S( zW)xvj56=Di1Sx3{vBF7#sX4AcyC79r8l+R6z?Lf4YAZ95*F0wGrafUTI2)jBW?B_e z6?~OQ<`S-w?&O?1L$R$m0$F8DC9k6nOY%Ky%OF;=C#y}Bt*qoQf{$naPWwEXqmv#y zP2JM?lCgt_@dY>CyT5#MLtS9_Sgu(U3GKy99Xz}7F>H|K3KWqWu_`*0 zRvwQ7DOOS=(wxAS+@Y}46f|KL=04Fna}wWgHzsJt^!F8sn(ZszlLF&Pu@W_^3y))$ z%^YZthZGu-3Vaw3o||!N%~*kG8DpDX>g0HytXrpfQ&~m67RU2d-uL)4<&Zl#tNbOlLwP9Y&wPlk_hpC@?gcjVSQcsMt{45!cM2h>-;2WI!W! zbyShqK~77L!UKkO>r}=QgDFRDo?1*c`hiHqpUhQOGf{%wAGc-K(;xiNn`I-Fa3YEx zC5srHPmjlbNJ!R(+nR|U@x(q5KDFD$;Hh=ztozT|3wh>zZLqf#tK&zv=_|czkJKX2 zdOeAuI_SdTnO0lOQIj1_qISIx-ZhdP@uD%Xz!RMI^^8lknt2y2x$J%%V)G=XzM(JL zv}B1)rs+gv!z=@rmolMs#yMCA&av%TZJ8CY{=wkM3cK-?KlD~(!_LA4^cEwe zPvtV`H3>zuUG?%X9cj_X*AdlR@&tA|M4Dm&%vXoP6qo}p_yUd2#Y#V{GE_WH6KMYhip@YGC;mQw%U~fCX3J!8JkY~AOA$<->iY`6xpeResqY3F;ESR~wpzSi@OsQmoHKtUoh6R`& zOdz65M!_?#O?CWHPa(+eKO6sZ6gAmd@J{WyX6G#+550ikNSzjiiEAtZIWZoA#KIL1 zD(mIj62>}NMoy8(A%7W)&IL)Le()}R#QbOEMC5%yRGjJM<@?*^vuRAJ9iiSQqBZz& z^x&(Ati2e0-C%vorF49l>;vdzi^l}#&%*_K|K96fFYTHRHl^CGAlVGt=H zImT8ts9Y6%Esyv(b@8^qaY~gTw8}D?LqXb*7nxcQ!P?{W?B8eO(8MLl^9kSMtiO{N}U}?_8m%-omA)|5#ql3NaW3?g6X)#yl1sY-F>8g zPes#=0#)EMe8epu60&}lq;NU!SDoO2jeV~;^ZDURPltP$Rg1ftWN{VzS>!>rAy0r; zr%1sa&-Yg>kBT3NCZU=PsClfRE`LRAgV)I9XZ6r$F3U0({?@hqEaLElzxh@fQsF_b zr<&*xL@0|s8-wv1HXo_aDP7a^AqE95_#w*}>=&7eMLB;{r9OiBi+k`r>;k1xq7y8Q z+!n_KL-zeN3u?0D5MRe>w#j~F1w)K@H}v3CvZ{XrHl$n)(-KV=dnoP(&Q%^CnaeMr zQu3ov^Rl3`WY{fLnMaWo{X*C#%?F6m)^hCH(dY4=m+=ofm;0cv-?~mC_wT=$NxqTW zJWOJdh9|ND5zgH8=0BATbkwXwLprIZ% z1l)1av*_B!U zg{{qv+}zLiqmZ|w4<_w*j}clkZEyT;L(kt-z7J_@QBQ4jB}T3$|EjKK=4uym3m8Al zG|%vmZqReRZzl_<$vX$0SWKTgk3^=Y1dF{5j#6Zub z{9U8jPyof<66G4R1KDcdi5SBkB$Z{(Q8&^ZCTu1`;1`mWgZb{&OXz?t`ni(Bh?AS0 zc9b}31FyuV8$487<3tHJR%<%kOm2+~20O@<<#Rv7!*FK4di`)Q9_~c^&X2`0XfU z>!sG%z`pGhp+D(WxvYpnu)rksc#JiaJQe|K6lcr8%#9z`-Fqce$e*KN70mA6xX_SWBJ_S8Kr9fPppG?-|@> zM|H`6?H8+jWYSac+cryzhD4D6n~8c>*dTMGjhOwl@C}0!1OpPV zNPmI-@dOnP@@&E7?YGMqmOK-2iuZU$zR$g;$6IicT*OfiNRM5>>k6(^eH=k3I*G=Y zF3=&zl)^zl3qs>U7y31>ep}L;IS_s}H@GX4<5mb-cQo~&q*jmb0g=8TP>E>c!7ryG zI^|(hjxk#k6WR*i%>=ks6M2J__4Oe0$syoFULqEzwCow`x#Q`j9cgOcV8 zdL+#Tn*?CJo_~(JAEWHR9P4>mMS4ScMx8Sp(`!7JrsC*ddGn2Dn@)iPOObOY`W?4j zVJ;u<#E(%X-7i~3bn|?+{KTG$#ZIB)aKylBUKcoD^lx-<8QuZc?=iH#GQ2m7WJ!Ag zQlnS0-(J-oxt%d~Z3gI7OgUHNk0Y|(*YP2eBe8u)aD3YB&9XFT*)K$-AwsA4A=mg9 zEeRD;tBh|8s;@X@cW^9g!OS$6e5Ln+qUh+XUPr1u>*v_91NX)i&!0daY!rD}s6y@^ z%;;#Kjq14v>$}gna${k@E{LJug5r7f#w_N2ye^_d6`rbU> z*WIX;Y0lpck?ukP0z5Cqm}sYBx7=d;+I=a7WR#T|CepatFMN9{mmi|V#5pm3#rz>_ z0*fH47P>LJ2w)r}p@y1F+6k@H&#yY>yblAWh zk!Lxo_FUi$GYIb!PG3Em3yxZ1@7r6Co*%UCHU5sZY#wvLC`&jRI2=qI-e{x67lpd} zG7+mN!f=yYrQycg0(LsG7gF`FeUyv`h*=cvdkSK6)vFs@y+}QNu#oNwvy?NR4y%g9 z!(7WWl{CMSJQ3V4hrBK;uUQYFoDW!Y9Tc@m|A^8b0eL#iHGKA+Kt<}QenUtmXGOn) zyfcHkJ{?XwgEq9&hP|BjoAOuyEGE7#wQ7j`!60fdA$SG1l!$0_ z(Zw_ffF-6gg)SCg-p0guf?C(ie_Hs~?fC6qr=#hmdq#)v-PT;N{f*`Og6Yk{UTDcD zi0+0X7yD*vH;7vJK6Wu;1kQb_6}iXyO4R&`6QZS4sLFH8Svcsf_n%yoWN{jv^j^{L1-=yeK$$o_hQaw?S5@Yci9n@U)Yu@8rH4*j^i*Xp9=o&A#)AcSZ)7jQxjvE7ID?u=5x_p2dQs-Y#Ejd zUq=6(M*Zewq*@M2%=AMvFfixS9a*vGjU)4@K^e0}9+sh81<$ z`Ye!PJ;<&US-Em+T|__=x2?IQ@%IeR zL`|1u(n>pfv{bQ&2r+@j!qx^@lMFGK!ys(w{@vN^9v!FZ{UVNR6zKi$)&KNWs;ZKH zyKX~eE^gqovu`DAoT&KRX#Dbb4dex?Y+eAI>pw<-xudh-G@v%5FmgTV9h6Y=Kq zj(K+B$S|55l)GF2JU8prYTesvKW4Lt1ot|5P)j1ImuA9iUvl~3ofbuYXSLW@Cu z%L?6^Ez2hAWLl6agBnnNLmoeKw?^)$v}uRX>X ziF17dc7hrUM*f^QlTbwKM(M!)R(Eyn8>g5(oCRkh5~(Jj0F~Ilj#NxR4>uU1NyG92 zaLutavWU>JE8^3D;w-hGRC8jXEw22O8f!S6GO=Lt&`RYq zgc96Bu0@%Z#(|dgCu>IIuy`J-HC#TWs~J*L;-IW+7T4XAbzSsHiO=6Bf4L3t6zYp+ z7L9YMbi0dWMzQDDV)8`|eML)Ciu;{lDsBljVc| zwUr`Nokh~+)PH|%@PU;#qMe~eFLPVv_d1KhsHG}l=IQ-ni@y8C??2QGd^DUYZmJm zG`P|wYhnREpL%8i%3sKSNhja~9*OqMsuE2M!kyI2lFgw|_8jHB*G&=$1|cbQ5B^ou zx=7RaLuHK>6iV%%NjIK0h;7kIxy<(3^vEeGLdZ6qd#veQF_ct9=3w&Cu2z{e>ewNa zt3is3TrD3p*f$>CnCfk4m*4i4Rlk>i11e#dyOsZ_mj1wupPCA5rE!WLg!% ze&P-@YQB+(QWc$G+6x!D4Z}eQ57#5zE8a4}bFqJWTu2ZhgxI%G20VawY{#@HjOwpv z7H)UVgr%t3o6}a=0PwV3jHAor#bBDzu-1(<7`6{N6#RBpN;hnQB=&7A_Y`U;mI^*` zPxQtuc3|xh9J7ZE*Db-Mo)4Ys?JmE~TZ-0f>BLYhU3bcA=nc!!R|laG_vxYEI3;0V zL8l9IYVU?@K=iP-xq}xKd9GTW_jId*8W!{Kdy|)V+{~Vw+5cE^Gl8mIHi5dx<_?7> z%;NE;ioB2xqt>Z4*SG{Kl&dF1J%cM=Ffi^gi2<`_JMVAx!!>eUxwBuKK<%U#nD8=mtAWDncoR)FJBmoYI)RPx0=3^$P zwpwdwHi7#J1YZO9sbty)-Nr0(ihrFPMiFS8qN0wvGpH84?6pNDRH2l`)H(@xAv>L( zCs3KIv_zr*J_v8xn>-g4ifpq?4kei1qz?==l`PI+Zy!g!7ERAlr~%Cu$+SXIDB7X+ z8hB0zz=BE*?V5t_lZboLbcA8`fcGWsZkR%ly<{z}DzRn$b`LX9m(WAdd+8YI6s}3> zl?EEwU2nAmZb) zQGBF_{;?&ok9J`{_*HqT;wqdl4>id-p66MdHp`&}WUTA&I*vcnB!&s%S^voQuz*NOLF?!&!E zXM!&E#r9{0T>DNFs2Glm6SiFxpQ;4Kih{EV=yYsOVjXHQYNBfjc0t##)VlgLiJ0)G z`9&1%c}a|OD!?K;X>_O=uSe^F)fnagDX6{L&%VvL-Au>O9mvBDSCqKWx~^k?&fvcd z$tpjK$vxwFz52h%I_KcX!miuLo|qHc#>BR5+Y{Ti?POxxp4fIWv7O0GJon_S_ul%x zx__m*tNWa;lhaQ>d#|;BYd&IKes4^FGi!h1zMn|l%W=-@jHU6ruFJaEoupt>deIoa z9*6@JNYG!$kgz{w&1mDcwBq8<7dAbe;+!A%+UVU5dUkZ)3eocT1S?dX(%xO;c03`J zR7rwq3wGndTwG?AWd%9WKo~icoFL%8s&D|J~*~jD> zrMExer)$`XP6njN5ttpRoaC@2-E^Ds6OUYv$=5XH8_T{IG;-)WDWwZp=>BvuGXmDY z!i={C6E+EL7wwL6$)a!K>R}y{kO93lq_t2fC}6c`ad%%fTqX77m$lmYF4ZZ@>)XhF#IkypJvW}LrOd|Xp zjNa<1+oElLa(9vl*0W7l>!)EWdSK^lo&Nn-oPfRX2bjs4WqEB~b~iQxPYY{{gBf`l za6wWL$6uJHgdRQyD#bbhbh|L1N{X8rM4P#{C2eLrA-MW$16I}ZA=_YESfnVO*z@Ge z+VK9^5b~TDA@@Z6@T);3*p!)e3Dk|cy8P=Sg)@+omLZ!9N0+fu8FF9xbVlGY2fC_+ z3c>u1$=Y;|Xt=wK8U%S7njExis5E!uE>K&pnByeJlM!fEj~<$NQ5hc0e$w*@kC^&f z93ToSouFVFf1xG)I%AkHYri+dt8`@GbHvPY0zu5c8A?qOl0fcnL8>J^S*$w_GVV+w zOEeH>K#xNW@h0Bwh(0zWUe)}&&kgrtF*e-{7{qM9$V{>mccZ{TlOHA1HZ5?C4*28p zp4C(4r6uxUkNl-UQ1h7fwXqOOfFegZWJL643V}lMk|+xV#}U3JY!2JetW}UEwRV^V zPv4>41Z^;DC$|Rgeo`sQ1$lUh^v#HBt_kji~(-JVfm7v9PO193Z zZ>}Kf=hRV1>;1l?3FEBI*$LzEtci6fZk_P`wS=6Qw{s-%T)K+r*^vgRg*Jq&;twENzUN<7q^b1okS62{w5x zZwecSI)n}M{g#xs8YnXRVsmEK!)v(QoS;*yeX>ER)fpG2;$-N2)Tsy7N9oj1MZS(` zm^RCl&(dE4A|X5$bYY0RG-Iw9BN+jcECp#{t^U2ke2$rY(l7#SJJr{Ziw0A+_Pvbm zUTWXA4MkkhA~u&FtMcMjW5g&-SIkkaVPONHX~-@Wp(!!6Uyp_(-XxlGG0NEcq7gS0 zt*wv!hB>vv3oHgpLckKXqTE&2@L+;0L9lH+$U)O!9;9xo9jsuP-``k|k)f{nx~4$OCf%r*XY-!zso~=VrWkAPN8DKBtJlPRk+cv&Et@T8C2OR+(_(s zTT?*v7a0S;ET+H;hqGbVGBxn4o9H+EGxmd);;nBE*B7D5}g4 z9XOnM$Jnz8KLxKf2gD4#vEK-tFZdKucxF*5Y3xvXIP(yLSQ6+1mbDZhxQMn889}Ty z$xvK9X;bE=8mICgR7iO9zB8@7>WsWkrJ z&1XWW(HUZG^HUa*s1LsO2=4wJuX_VrZ%@Fum)#Zy{|tLStrV2(Y1Z=CyggLQ>QvYy(kG6d( zI$INxB7(=9CdF_#N}p;?#Ba`ZQ;Av$YSj1>i@B@B?txYM zYI8jhAy-m;qoD)&cm3+?6_AO*h}!C~T&pPpPIzyJaxDHZK4mIB~1!oy^`&{r|x z%84Ar*hzHF4ma#ij;G`p!Y{i1YG-?45fMj|!6ZA;9Bx5Em=XfU!*lMu-2M#HDihn_ zsN#;e=idcC%Ra;F=SeINOClZfVV`wm~$R5;d=&|WLESdN|`ns_m^>-yfm|!Y06wk4=diAjV2qK&)oQ;RA-&`!+ zJQVU!0v}L$n<-3l%<2#CfchIgkNEZq-2C9SUFAN8mHt~0q--mOo#>B_D|iESqhxY_ zm-=t`!JayPN7Xp(gBE~5slONgG=Wi|HXURNnr*h`Qm6H&+Njh!Bt{czt5ch`ngkOG zo}hCFGqowJ9n&?AZt6|j#*Rh1#{I?4LeMGA&iFKU5Sj^cEZPIKAkM9Yr$KJLf?Vba zQ>JzkD%7X!7F4!^Va;)fS_S%&9W?uBT|cT`$E~^KN34z>S?G8!6bI&L(qx zm;r7uI%t!UB!8;4m1&x2*EomHgsjWe+sg$}f*~epQ-Vu*np#`TW(u;Q0VtDB zz_YWSG48>6i zyoSW%{!-)?%ZL)Vi&W@e(GsJ3-A{!xhH9fM*^GK<7FN-a)zwdz%XO=C zX8D@>W^P*c#$THC+tS^Q9Rt|}q6xgPe^^cJr#h-EmYQqrxXtTj*+RFdEdTqsB?m@W zal_)g24}ZO5;a>KgE=f(G8`Tl-^~+p6!+xjPHchxeG<4s6CTfHmmL3Os!c+#taYC-cG<`zdJ(2?jySGdh^Z6y5i(G-lbgp$ijQqa@qW7y)f@ zdOh$$EMW#+nf5Blwpc<1d7%M1%~T?|5^@<}U~Ab0R(Cmb-7N}Vr?j0HQusFLCl;aT zc_fo&>??;YQp(*`g%XC#c#i#Dd=zW~!)l&;$d8*|?tv3FiF!|YoXq>pk&bv$<&2g% zH95qPXi3Hcg778ZHZUe)T64hrIyZLdJT2^u2uf~b3`-N7GocxIG$_g@#*VaPipoT@ zB-O|TNU>?iElc;NAd+fLr^e+hdCM&{aSw!Cd!;ih)C;M3^gU}XT(~u!p3%V zYaEf8oLR|>0_qIR88@yR7UK$jvoNWNFd=H$*Nkq3?ztq#*kJROfS_=WUeid6)TG1XM#$izGsZ3_(Z7nMm6y3y>n78h+!MhT;i)Gt=dL(>`)<_3Pn ziDb=1Y0CuPaC-6X^3!a*-G}KQrt!>{5Mk6&D+BX_&egR>cXc6G6$ld+pRk!$%ohQC!?;jHBGw7s4SVcF0%kv3c& z)CV|J!FF7t2*MlX88?xaN(w1!#LPkVHy}e_NeykSpB}>x*xp+r5tZQ>5v$V02qrqg zMHX;RqKHh9iC|k;EDq<6hNLrJ(Dts6bMhFIjq{V$Cj8!=GhV>}p+SX&;>j>Aav!Nm zH_)y|S!9X!SI-_eDo2Qx39yfizLJP6 zQyX7kjMM6-o&08=v}}2T8g0cRa}UHQUwtA;b%6XKdOwMNYBwV1^Tj0qA~Gb+XzjDo zn*F~1DY8HdJPo*aZEj!sNrIWo^KDCd78~Q5Ui|sg`RB&C+%wGl4$jw}BP4kD8QeRf zYx)=lIT#2(Zo{t>x4(+6&V?1*eeoc(@||q+nY6Qa&R+5030CkuSILPIW5nsaI)@;^ z%i+-EukQiirVqE_0$aUIQUSm@k(5uOS7yXl9ds|WSG_*H$- zS&g4HQk+xBPL6-O{nBK0U3J)!e{LC zO%5Ss<+4%ctUo7Am>Q{JQ`xX@+aHJcfFIj#z-k=ZAE@pmO9wqAbNd%+$n>R9W?7<(UUHR`F9Q4V(e(@d55E;Qge(Dyit`jyre7F96qDAB|KZ7R44Iog@y3;z}qI-Y0}_0gjh%3}3NFid>~zjcPkv z>-@I%K#@EG!sU&a#J**dx$pBf$@W&_XtA(0hBP%rYM^-}*f_>Dw?U!kl-vqt;3 z0OIBSR6vIzFdtTHh~u%jxno{8=WJ#qsEMkpBX?i`Fq=C70n(1Nx=Vdj{pfz~-e29> zn9xvYA{2ZN-rz~tA|v-gUAC&?VCysPelwz(&U@o5=+LEEABa!hxzPWcln<^hls#_& zC~t^j9E2+}*SgTpdp!zXpUKmk3~S13RBXsRb@854CZGO|xvza;X8?DZqx{RMJiYA2 z3ozT|gIij!B(b)!q15tIpGJOn!DjRz=69`ok4=y=6v?OafIfS>WF@K9nMH6%nioJy zIEgY>I~*@DEaolOj_H7ht|}GdU$LBwY`ghKriFe|bUr65yv^!e+KTHOV{Lc=|9*X$ zt*)&F^s^Pp1f!?LV_bT2T1l6s&fX^-v{g%$OrltV3aS&|-Rud-b$tEwNjI;jrd_?8 z;Ys{j?fuijm1+?}_daa!7E0$CZkAkABnh^j%)cD!v9XwU93WRBN*w2(O+B_q^K(Zz zaxk1@lN7b=+{wuUxPet&!pviIUp$P~)){(^8s_RdeSvWthab4x<+B%mhA#^4e|>EF zXIIb-a#>X4Cn<_a^#X!q2k_lKRoXsb{`f>5ls2M6_Y-f&YDRs!k$Ng`+o`oRAh$>f zb0toe^PeK~Z@cGGi{2`s(c(I4myKSt@WXe~0jRA{g-^WaL(=fwYwyh~FH3X8EXyk@Ez*BVn%$--o7;dguHs~)79nNUr${^=Cgl$)vGw{ z(K)dj9TRw86gZO4U@go<__S?sga}^Sp6}ho`*%N_n`S@)wooiTqZ}`BSzv}FI&Xa$ z!PcI+Iw%)*x?w$9>=uzBlJKRF&HSHeAF64a@ZlY~#8W3UAp`hR$o zB)9s;u|2*N{2U3bpY90^4s4MU_N>omO_@>p)_WHVYZ5id5tnxNI1txUN{^0EyJY8r zC!JX-xwY}Ja~t-JEIxlaeah`9H* z5jR#0+G2yE+QnQ7F&MN8(b_%a_A(uPRaYp)DkiJ=I^tM{BKs2*u!aUbf)0MtkCEgz zhjO*~os*4Vtz6(t%MJ&obd#0-+yKkWp#&7dv^;%2g&x3TKr;thU?|xSUE#;b2vbI{ zOfeJJ{Sr$#X~OUT*N*>B{&h4#D<&Y-NXu*5MxS&lR}ojpfbUkQt$8d7E~jXt!byAT z+)gJn-R>BLlh2aGOi`%!?TGO&_SgL|r_F+mhvT^Jtiq(1)-@we3X&!~sZrr0)vyx# zvRW+DN-B3b#-X?bi)V<<=Hua<@R@-mIaU#qaK!~0t$?TA1Y`S5i65g|hqu@6ecr$7 zoSA1rJH(_s(o}HD5oZCUkpDRRL5479kOQv)ItGEezMw}R#Xn~JwCe1TPtzQy8i^Z+ zc=ydGQ`+K^>g&W1h+_XvohN~JB*OzC zY)4>zFW#8*w(B(8q0x@w2#6+SCY{`lMb~OO|Cu9CJsON-4z66+Iw{(4 zW!7oi5p9k?1mx4STJ8w$1ogIoEYC;ex*HAE)cjPM6pjMmK(OPbwwkBJONPj59}?8!MWShK9yiJF9p7DxjnA1bMPWyrx7O^=l1(o%6Wo#S7#$?qtb|WRjExo) zrG~PT)U1_MYd9#MqfdNjvRi^@vuZjP%g+$8SWAZ~82C}IwG)E06osR07Q9If1-YR9 zjG1eiVJ)S(G7Z!ZB0fDTOwDNE#D)ytV*-!rJHG(Q1*fkOlo5^P$+ohwWGY94Jjp^l zc*9LJPO&<~T0RB~RvL59W%b`KXzOjaJHeu=r}y%a(NLRc5lr6<9npb3OO5M4juI4j zx34Ab3Vk6@mlVsHB_O9FTdkZ#cB-j1^iSSOKrexfCkb!iBHgw)oT#&9nF6N{1sz$RVel{n~nAdxtwU} zN5Nxrhf{U(!e3%7_j>zxl_Km##);xW7G1<}ds3z7DATg9WQ~Y}Q!R%|%W-w(&gmca ztFzENXV-CI>|x91)$ql{h;zg-ny5`cjpiV)&MILCY@Nj7UPur@f~r7zbyh{?J9_iSWRyoj)MjZ>H8@SQ4c;siXb6Ycddn zZIaey4cez4I4+m9eX+J1R_`IHr>c-Gim#`Cr5`n^9GFf4#j3kt7D{Nigo$@hPqDr~ zMB(QyO{b0DigQ5h#L0n;+}ScGLAd{pPC{RbYjIXWd!DoG8Ll~EQ-dOlC6_b&Yh-~$ zmPD(rbL6ks*qI(Tn9f(Q_upM5u}|jL-%D_w)wEL)lCT>K6O-qZZ24_Bd>~|^k^)UB zL?3OvlDsV7sG(SuBh^UT)-|Ml(w#&!XT4*t+D>|8)i`?@NfOidCP%WIG7_+?Qs=!| zWxAPlTY9!iD@(#kd~PAJ}QDs5vG- z@}ziO+$52JZdHw;Cu4p>wgoUiV!uitB(t2c8NWP1@Y8c~loyzSpA;WrD#C(~tw9?9 zk>2XUr@^c8p+TKxDH3hDnfZo()EnE<`6&Mk_u(^Xxp*pQFiV4gqn09zwsTy}+~>Aux&- zgdI|R7KLuc!(+p^+*Of(M(72O98u-kO6Iv-fe|I)l z58LS`t+mtX{oP+7i-i2I$Rc8u$I6jm82S*8EI52J)5KlSQ7!3~(VvSbN_t7NgoxqF zUP7_-PrdW}9J7pYq8*q0Sbq0|BQ=XzUgek;UnAgRv4rxE#YpPql+bcMd=78k$F@l6 zVV_O5LE||DFvj0S#5-h|&XN&NN}g2g9Y9VB06yL);ex_U)A;WTFUJ`My$ylz@o`Yw zF+86zPY?tM7}4>w$%4&{WAD^nS&z$x(HDN=JgNI{?73wm$)3GZmMc$@kzWlwsHXaJp()eZeduq?z;q<@^o9MqQ4~L3JBgKw zo>1YT%Fz2{4716NWKBFb_wxaX}HPD}&U} zZ1sE60GyYP&>nUCTigOg3h*@ZL==Hb=6z<)pA+~dS2P68B#o3xfm9z#J~Hb_%rEnPZuadJkc}Pc zPw31X=sg$W^+@X3jRWRtgT+9Vah|8E!gw-|VQ8ePMrt6GVjL zyJ?185Ww31V7BsBbnTyqVD9NZT~)tC2aKU@o*>Xpt-&6)>WxZa#RfFS0j4*rDwYjO z{r=oK+PhJKc$>uWPB z_6LznhK4(mj6(eVXD)8aa76!8RzG?n(;}~85OM~+OVM>fxC-Q@8Tzu?nMiG63^^RW?E1}{&E$2i2B_dQ+AlJe$Q%Ofk`Qlj$x zAF=fN5+$+xnfGOb#MELIU&7|y;N_2a{*9SNwH|#Qi8t-78im_)7>BxGJg!SEG@1kP zqT6qY<|^KUu#{_6^Z00~GsSTPaX z1+{X)iVvd&r;baZuz&Ygdb%VlXH`>Z$#*2^`kdw7Z+j2p`-SGeNa)(w|K8NK{_=a3 zERDD&O3u7kAzOwMWE}Y9c6iB%C-OC8ExC;^b91tV(IvejJW251qgbqpVP`V%s`29R z?njWWorzJwt3kmx(eub(0;o~%;P|RGAO63Eky5v^qYw(B!>SrcqHpVD zx9b-9^aL5&RYwI)5jouVW29fCk^c-J6B6Cu(%PpPEb~z^>%((sjl(KRz{8gWoBiC* zgoGV0E~3JS_x2$pUi>mhqQ!?suLRmw`Gf_|kB80(h$Wri^dM;T!2|id{L;G>x*ulv z%c!GavTM&jHMKQuOjoP5OaO>fRQNzt0?$?R;&ENMvl|3A7%ocS9%CEA%ra0dRKq~{ zculWt;V|+fk>z7=BmZ9(z&4rW^Y~$6(iIPeH~LMr&)Anln{iFfqb99~l+bm0<`|n= z)NbFmjTW1S6^$y1%Z41%by}0jHs%hWWi^dGO3D(~gvq~268})6n^bN7B?*tKP@q@9MWaw~=IIhHIrO`TpJUmboM^Zc`Ycq9OUW^Jl)V zh6h(u0b4MMaYw0%n3#=fg@6{>Gi*QAgBR62c#2ihlu02r^dK{F;@0zyKkBCEujlz* zA2U~yqpmyXu~`0(tg})>GBPxJk?Ptc5%d&Wg>)3?R!`wQ{EB~%!yai7hmGNP#9_t` zbzRro%u^DM<#k2WtLN&aNLq%oNjtcNQ!J_su;(Ei5# zJT(x(WX&Oa5Q>kP@v*h>wEgAm#Hv#h8)B-aghMwq0z&e@9gPJ2h1p^|i?d!@0y8=0 z)J?BoYR~ZMQF{}2cO{IxHh5dc6^b>Lv6^(DYIQ%;_=#41Y&!WL-yc=rdOq={q=;+^%Dk+hAVog~`j?3*HbR9sX|rJn473)q zklWsN9DaTogi=>oAy>e$Ow8boCsK4@MJ2$<-v3HUn#se@vl}g0dy1nAgMGKOe!vk{ z0ncS%QQu01JJ~v9S8Z*#sd?X(6W4DZ{8@u=h^*slg$t_tSE$ld;!7gH3eH;fDxLhY%-c}$}%MlU{@M@7p@^oHW;a=!(Vlp8u;q;C6OlQ zjn=Kt53*X}iu4z*iqGk1mUC1jhZSkG9tfI8>hGFG>SN=DAS*@K(OX?~lLjgent&5y zkI*BOTy+sYHb&{Yd7X{g4pDpI;N0?IAq(hc4?Jatz)qq)2m;Cs1>CJ3G`t|ufgEWa znzsHgu2fj4(DDMKlw~QVFBNXIs{Q443ZU)sjTS14Tw{ZxRy&gotl}dxY{Tv;PEF@D z(e;F!L4;sS6;&XaRE0(I)VQZCaMHq5Y3)TBXV@FxznyMkLS4*s#Rw0zB0N+_d!94# z%+F&-ls?Ir%pD+_yPd}oGNVc4PDwhLQSMF2l$Sf#>+8-gj>S%aqU?R(Hln`AzukfhFicG)Oad7S zW^lSc6mzek$^s#WKXl+E*E%igvJv|MBBM!Szc9E+WlVDwBsA;=6CNPwP0_AHP#E3j zK_9iiXArVO<<=-fOq05_t<0v$5w*ZjOj?hMqS>l9CJj-w$@>apnr}2;haV&R^O6ia zwP+M;9cK7L7BQ`aOi!xFX+-vUvXV@@f=|3tE)?D|2Y%hZev7}|XZ#ohJ!IDw-}d6* z!#rM|!pE`@rB9ly>^t%b79aYms!OxYcnw2wVlpf1eu`d!T$0edUY!U2=2j`{^8wU% zB66eMuMg6$h1?{BI*MXQ>K$?Q)AuzF#sukF*Hhp|W=VaV!x{9{qiXxMJaKlTSzZwP z8j3@~ag≥v5@tQ3!4A_zh)rrk;cPVoV_f?&6)YeUK^F+{VEoLzSy3I0I33na5sc z0Y%ydJtUN0MWm`K^qR75&iAx#S|F0rGFEe?hODheT|`y!8f>}u44VqE5F%h<7bNFW z^eNb*l8&f&m&}B)6T3~@szd8Xxe*S8qaWJ8U=tlwrc0{Nqh$9E$8ZW0qlwht!z4|5 zD5jjZc&Y~cl$*c)8vjI<5JHlUfgxwtjJ3NgCWfL)M)vfG%Br%vPz>L|++WjpMpQyO zXcJ{7YZ7pTB7#sZjYE3bk;X1zgQJUptDJfGNcIwTdhW{uM*B7J*?bh^%Jf?==^V-_ zj0>=e;|InGN)-I}VXFhg%b7boQoyy}|9kF#ykw08`Lyg8y>UiY^8b3~j$@0^ixb`u{&ZaCEZIx>*eI_4uOMpZXaH8@omN*hOkT%~v?~hk9`W{FOcW zA42*5%cuURMgM?E0F>jm%>^ImW%hr84uwQ0wTFFQOY8n!o^ZHgv>yG#D^d*Pk~>R% zDYki65k+`z%M`V1SyqoRD0}?(-Ul$M__OMtp(8&+3%HA_bY?1f0mbc4{5Q|CcKV-7 zsa!bV>CNZDRv#!-7yzHjKZ5~XOg;;lG=~Ud6`FS!Gy|s~4rc~oq`E+MKD(FOy!+lC zy@G;+JgnyL*9UEj;Nm!XW+xWzO9=FyuM$vlP5Frl`J_R@mio|w}CNFi$bBFAKGlPz%7RB(J z+|MehVGG()v1QfweLU7thwcE{r#}>vc-4Y6W0Wh>EUkNyAs^97SE-+TSw|Ns*G#?RdTJ`cyre z!qs54^+%2eF5-BuEk#2>w+ugkq>9(0=wu4ZsIa()M{39+D@B@ z(*rRySmQ{{np)MNrA}03`>$PYkSq8C#>|Z6!Ym~?i#2UK+I-A{w$ca`tY;YvdMRpE z>X|IZ2R$);tSx1m{eI-UqKQnQ0{0Cq{84+t6{m^8tP69r78nocVm%e4^~%UkJ?u6s z)#y51x%h#PBt#dI0vsu;dhXkwmtivEJSk*ZSMTeLo$@vi9}&Si%liNbp4Hm|4G`Rqz$#_QcM+H!xY@qO`5;oHSN!6nBYnc2MS&q{LAZc3q zUYU~+iD^4`s>Mb*FxEsPz6f)uw%qW0y#fZu|M>wsv61iJ z;d+4hLgsJVE{njGJE{O>n=nW+iU=A#lZ?3Ty6OF7|5N&HD6K6aw%*6q5vhJW2>g2h zCY&P&*p;UaERVsUaW1U`cxRA9pb0z*LW{tVPDo^j#F)?kx9q zN969!V$7yI(!6;{x9@SZ8;j(10x{)hg9lN>UIQ5w(Bv0a4bEbDEVMODPS8;6x|(sX z3(9#9IVasfxTQCc(Fp3`!f~78_J}RZsAT^`FZYXwQ&nv<*og||Vv_KWsK;qMESW)l zI% zotpC~I2)1T61KRpn>PGFjBvq!(V__Qd7u1l8cG}~t>4z;$d z&1J0tJE239juSv!FMTU(Npr=`VycsvzTvQtK_PAo)wyB`_d)A1vIxKn{^F}=pScM~ z#!5^>)kpi3$0KKbR%QHy=wPCNGS8!k^%m%bdAwe=rGE*z9SqAR_xbC#9{`T%0i?c; zUYHt0vPwISj7fR5l)@{So|6v<0MpP7VoGHx4m#o#iB^3@Yurs}|MI-Atw^u&sxXQ> zH`$BUhIni`07l*m<122>!LWJOw99slUk8+UM8y_W!!Ul+< z8=%&YL(_d7`PQrSyLFJv*KR|eb1!7uZ8LT?EKGc7I`*U!9%FtO5}~#;LCcAy-L)$P z!;J-YO!u~Wz)qi)8xJvSue=E;N4Wr-9wqek#cO>y-{mL2l#>50lD4)Tte!h4clA6H zPJ_qr2uRhr00{v2pYmCTTlOGA-(bZIQu)M-U_W*sz++5@BTjUsevm;cdI8XV?Es?fSzaq9FwAx8?s7L6=}*)>(CL{s zLoSl#jyH_sI?m{}n2gKmKV7J5$#L7u4HBy2Cizx_T04^dg+Q$v2x?#+<&(l&ie5kM z1Hw3UTC7v|Fa_|~B0n=FqGLbrrHI3m2FayFk%G?Nyf}QfHKQ0p^j?Ueq2Lr^0*`k< zdFMe}(L}QL|I~Je?9uBTr!!ZonU9uH)Hrc-xFmz-SPJLzZ$l_KU%1Z;pl^kqWdw`$ zR87UAn3=J2vCINWZwpM3O*o9_qdkjJe?6iifo)pf6k9|I;nvSIoZ9;X?8gY@DGyzx zYI;Sd{bdcKHNdx_xW1+P1cWTdc`D<<2{<1DF(;jK3yGmxE5_Jx2PUmbGi$>aD^BW4 z^UB96SSK;F3xNhsx=rS>OJTxT%Mk+{VI5hz1X+@y_cEpbyom7|$Y4m73mFYqxJhlV zY8X^UZCdAXZ%sT;J&diTs^FA}>>&SmNf?Te*4xAP7U7?cxI<_|*a{^*^TxLsq|8>| zFI3u+4QpRFwd>p+8E*1aTmXy&f72h;uXbUU!!u3t8b_)$;!a6hKTHkO9^}bg&D%1~ z6@OeyB}-*CCfbopq#QE2;CDnewHkM>S#q~2;y*~EaBN(-=~k3l$L88ow+uY;5_H=- zK$6lo+mt|t^094<7BC?VQXS$#9vE}>=gG~sE2C!fn{{c8IRCTmfgbF0(3G>rS;;`qA-MXC<3kQTc#D7JZFEsL4K(~2sRsu1kU zOcwi2M5VLo90-HQVii{1^RO?eO$NddlK+zkX;l#^$d$`DF%G4vtq$M!dmM;-6~#gu z#Sc_uNL13$Dx*g~mdBV-nAtaDCT3%@^?qEknFVxEUHBck+pPII>~-y%$WRiG}q0IOZ7J;6P?8`hzM| z9%r6$^T4IfV@%4ZD2&}vfsck3rSdt;Eio`FK{W3mp}YO1CR`YJ6FIs_h>Ijwkx?l} z$)HFP3RmIu5fq*uVzlX2Wwf9`RF<*NCYv%N$ovl~w(ga?z9<5$-D2!B;0pNGO*m$yPCkNwWp?SNUFW;>QT(6O%P#Q!Ui6&ZO^EQTfO;N&MRz z=aBkcQEMWh(jJm3_0((FY10y3ViB>zp!(;PEc_6nXPk)Uu9>QZ+Q}%C$;@zER5Ye6 zoE-w7tgC*rC%=T}t$%S?tqxD0#=f4LLZNBYAc#kHHg}d{u?8je3W^(y436E0;uzwv zl@9lHf~e>61j#Y}kWJ~_oj>ePqk@dsow{~Vd_qM=3%mDj=_q(jPWS+ocNr~&97^KmpAd>el|RaZ*;0n`;!X!*&rvy*iG}KrBbDc6urDOBGHW4>qu0*k zl5oQNA{J@QpfX$TDL{7navfkV!hP1g~@Dp4RcjzC#0wY7uCvA9g3nd4~cNZ z4p8@r{Ej=IsCr0*#3HAN6tow-b{A40jXZz2Z`H}4!O(pT^n1}`bLlyVAu-%b8=>LT zT%o_Sr@xRIZ&Yop*hRgcqc}D*GMY1JQ*7)fyUAvI8&B z{V&ejrQ`J8Fjj5+Qg0|MV&4;n#J#5&q+ezpv_(vQXHMiHL~#cqB+NgTEcdpm*5|Cl z(Tm}&jxpvS*Pc!Yo(>*+tilO7PJ$G%D4$ovz#D+Uc{f^_q&n_J--+ckM;j;g=*ia$ zP9nmhc;Qv7qsIs7a7m*01*TZFL+r#>!F-&WbxY^Sh2Zkv6~rYU5sM=_px2rSj)p1} ziq=&v7|ZZvQiYm^D8@?8^9JYRm4fM}Un){7Gk#-*tL>j~OLLTCtgZB^i@+?5rYHhZbj zKtW%-)o_u9L~7X6s$8en^~XZsHFL@Ua7r%x&Q$WUfN>Fx_vAID*t%0!YQA2VlurJTF*5e;GE zf<}=OZML~i@U^?F>Je@J?v2Wjla_#ko6#Yul?l71` zp5aTdu&M|h%J_vD6ix4nz*5nBpN!vFQnCu>p70H@lIBK$7E?tcmpg5lpcqjgzUL+P zIsh3@->pfYD-7M%6&WO4qs6IkCVua4@9)xE zE|S6;sQwAoH%1lc(?#1{HSq6wl8~4}4HrO2Q7*ed9~(ovq|`5c}-|0Dl{YC5&ko4>E}xBWG$M_`VO>8A~PHD&-}qy*vCY z<{T19scU-Kkhvi0R7hR4ks1Qw0~?NnE0H$D3XBRx854rl(+heK@)}s5u)XjjiEH71 zo6`P+R+Xyk+XooPlH&(eDu1Bs;j>Ekyzv|0SVGFzU1%cELE!J9g2Hc45MdekFteiD zQspd*H9);^gFsJnO7B%mh3rQ5SEjs4VeByH*{Ci`rZZB5cqR~1k|s|-bb%Nx=7P#5 zm1Uf)q>NVA54|laK zrbxq~0=Y}pVDuc9v5m94H07kJf8M|T^rwDWK0So-G~|K^>*w{XSIMOsXx2?n<=JO8 z9>Eg&ev4P@lR2xS4%7Yp2_f<)^gvokI)!^iDDN5VYVE_dR`N>zzCOxr1bL9@+;%lb z_}#uK|2WF1d+t!t{pGCem(Rm=@HyuozV8*xC$fD%F!-kFZ0zqRq1__ix78WmmGj|n zxR=taP@;}LH2+`+(7B;M?^@DK>F7cu1@VceFwWYZku07}-3nc=vl72YX*zPUC$S*v zQWB2)VWm3||31*(J#cROHC}2*rd81jZsfI6O}_nhIOKg9t!QG`H?f?uortI+m&oQz zPLuc*V^Ts(VfZG#a%eTc{A%$Di^DAyihW|g>ihTM=Q}NovT9sR82te>2Mz6*^`%E> zPn0>@OIFbs1X_;Vk8>c5la$cf2(WvA4%h6Cmm(v_`CAsihc*&S7%*@gLc*N@2qE%EdQ+rxfp%LXk$c!I$Lm>97X~ z?lTiy)bfXmB1n8oj0qJ9$wtF(pB%o=FfGx>4D<&oQwmT=Ovr~%V$Y8_?uq?;N84?! z?|e}~xaRl7x1T+bPo5T=nK+1|N-0JeuaJZEU%|O}djgj^ArP&)G$E)t$0$J!1Yr#j zg_pk&Tfes%;@v-^2$ar|tES>^#e-F>2d% zIR68HAWO-xfXT+{2uu~nwm*p}@HG|jG`?YWb7Pdx4H$3np6@u8zB`o8@6B>!ED zGE%E7kc*N-I;O$fkh==#O{$m-NQ^;PxTY?l@L}C<>@_dSKk(;B2oIk6043*LJX>QA zHcs4DX;*`@58p*~lgkoqK>OpwcKck_uDjgRSt_&e@0sbk)USv=mneERF_`$V%!X)s zCi+`3q4Mf}zb9+IA*R_la@-OAeJ{8Jx@{VW_9XD%kQ#))#d+A-WH#E99YWO=aH#qA z{!B*y?&a!I-SX=(WYT=qq)ycEaoV&_=LBwa^%JkCKh5=adR=Z(oFg*G+Rkapf#Xio z?LJW-SB>&FNJhSyuLJrJHi%RMYKiO zXzBQ(CBi3%&q%sr?1YcT7OlwVh|dJVv4+ZRQsg-pw!+sy1B@)&F%jo>#c*@ z`olMEoS?06X7T&Rv*L>Ae12KV#h{Cht# zw9oK(hw_txo02^-P<dJ)V3>_-6CLf2be1GE% ziDVCPb7Rv{luaTPn{m=9%tih@EA&o5JS_3QUt-lfxT(7a;ypN+RB?uwq3A<8x_2x% zc0UqPf1*J&HJ^Y!<~HDM6f6Scbfqh0jq}}we%ilD@gJH&KgJiL>JRFVlxo#gPmpP3 zxN2W~i`^=c6>q!Wu4r9!InWf$8eO4nz-0UW{u1fa*L1QsPQ!1QQU2RkAIm3>lEbGs z8l;RQ-n%#K%`&=Y$=1#1o`2;UkQe%?qL~?ck9dg8W0+EdmuTG9m!cc#B8oIQc3mx> z|Fw%#+K_xOlc-R%OQ;3r`P@+;*K-B*`D}Ahg!ZWLEMpX}lB{>m>-Oeg*NhhEN8jCY zyFfM(6P>M-Mo5rsoKWj6=iv(8w6+tTsx4j$rv^H%GUs}m-tQ> z#}JF$BQ$mKPP=Y8My3op0rf#MP!gSUM`1MQKtWf#{dm+hA98G-K`$Zp{ugHeV{eN0 z0z2|ZiU|$+!_|r6P||!YIn7Xz4V<=*`-mDl=?u$F2^U+u=}EH+J=DY(XaC9DuBNr zdq}MYr1%YuaF&3(`hal=&$xTfuK+yaEQwIo?bnB2)~racsml1<6f19?9#pcKvvUuo zrAAb*k056UDiFt7?7_5rX8JShtMHfP`~g^%7nihY1E+o)OJadficibB|P(%7F*&MCcG@G=M|;%jzMKFk%* zF=A&dFtFStEvz+^!GK*3M?Zd;NIy@m_6@bY@r+bf90I~7H$%q8tr+JAJcwbuIUC&i zJXy-8(+A7tvHm$S%Em&`Z+_4OrMrWw_3OM)d54*`Q3{PiSDwQj^SFtp|Mve4mPYsy z&=P54jxMJ2lZK{ce=Wlon4VJq37TU^kZ^YJbL_K9lp2y|Z#A0HtB|>TP^|0nbJ$Du z>78^)TlZqEouSBJo^9(*nNm19oD=}DUwyOjl|TjQvWi`lV6h;cv9>@v=$ z^wVpZlndr~#-*e8FmB|-*%HF}6!ss6(j`T)H>tPvqMKUs>8*9;`Fs+7MpfxG+#bO! zvRiik9j|J_`ApAw?v1uoyyX3q>QrqxhB&eu?M&8*QV@RAH^uDjhIbjc4rx@I4+WXD zLHLy!s;5IdD~<`mTj#!zHrM>y!4-eTYrHYC@!65U=1gstN)mz@~Jh}XDF5>hH@ZtcVDsyTOK zZZc5CgSi(%hFgVe_~IV$<1YJQZKrk4=shdj7Y#y$$Zc?uM)A1><1ym+!;J2P!Lg$1Q-nkHL z1mVJO9iYHZbRQtC^OAeRx{zsBJZG{om?*}`BZM)Oe0#Rp2yR)sC?9HVpo|07N&zHQ zS)64xDOl}jUe@ucbBO6CDIN-y;p;Q=9`8vP4Ei**G3)1=3Jla|@$5a5StV61Tb zq!vDQ17Y9*MGw`PY0c)b+$vo~8#JnJm@+Ff&IC2OwRM*EZ&fX|e0|EqkMuBxvCccS zp;&1-G}%d2IfS0`%f|C8{Bn`VyuZv@vJ;db9U1V{ML`&R1@R4ISOO@!^>=6WJMecC6bE^+@l zf}_Z){G__ru(=|$;X(*j+wcMZ7p@$v|4WYjwC^lHrbt{$uJZX@IjCx(apQ#0SBd#& z#JZ8N!}1WsB3(v{YLHPR9se_TO35BUsX8F%s} z`eF~NocDFAFaa^21zvd%dzLBX+^3TCrE6&NCTk>h0OQY0Pf~)me%XyS@wc~hV3*8@ zdieR0&1XA42Qykhs#^puN?Dq_|cRQSyjZACw_+K(4Bk~(JH&@MJ|IXb?!4d~P z@e>0BZLMu=Y-9G%E)}9120XSNPUfdB@m~!(nJJyYG(_MF23DvIv$FMsQc(_{k|lxO zzGvkL{SiY3Qg<$8%XWifnKan-LS2~3XY8c9gO{b-N~Jj)FQ^I6D*qgupxS-YVOMb4 zcTBq5GySlY=Xr>2fg@1sqK(Vfqy%f?D2{q7or^PWh*QQTct2FBZh7ZCVt-DL&*MWN z&QldWZgej-Fso=wZc9Kkq4=AgcRGvyxlVDxB~3CximQ?EdiH$$sy@c848t^~Gf;$Y z`f5zuK4=1A{Bc`Pi7SidHwKw~Vma61*n>N);E<|(UghfXf0=2>VVj8)HB$|PvI>W@P$ zYX%|^+0sQSIG>k3OmN4QG~h1J*lf5oW^TaBC|t$pGg^6 zy$T>8i9*&3G0s5IO6&y9BVj)0K^hR%VF4@ld(2xV z(_gvTli}etHoAOOI5g4_K8**U{_b+^Mmyp9pexqY(~##c*0r3{>{rbZOCKvO7xZxWX`*jvBC! z%~gwcxz4HU{m3z5>SfVQiXhEx)%875)D^L<52`j=q61Iy)_B<37Oesq986|lW$818 z##0!uV&LLBOQ?n*6sgGHJwqf{kbJw`u_QnVll`d&zL^E=}^~ zx_OPvmI%BaNjIPo{x~y$n`i{BfD=ZFfEV2-MazJz7%5$%HESwwz@7Fp-=-^q4_TAF zI)#owgH>K9WB7hmmX#DFSi$?Bf%|R#UpOlc;-d$GS@?76Ur0DVE{m7*R_y%pk)jR- z_439Z`WU^W<2!}b{36D5Zqm4Ns+^u4jy8#l;ql+1Nrx_%8;`pQ>%`j{D2KkuP9GMW zwc$VF>Pk;wT&9W`C(^jk{)X_c*kTyxB$t#ZjIoyd54$@(bTfEXEWyX2!o>-Avc+px z6IP#(>^G)?u*I0FM>e@|Q{eNMLUx*n;nsxRvy>h&o+j^DM~@mN98<1Cg^#LkS}N_T zY!ODjwa62LnApJUXDtvpoWoWVa^(SS)f**dOjOXhbU@6#=w_|IZ{)jwUsy>!Cz9wL zbKCHIv?=YYZ7Il9uRxcEiqL)3zGd;fteYa;rYPa#-3G_>VrTOzS#VKb%Xzxg~b;qf14dgdZi8?jvVd2M4%GIqG*GNMU0S1RbT1J0Z|y9*oV(t~IqdL#1n zqSLleNsLw_dF6EdS&%yBnv`&>ZXu>3TV4jbeO);-y`3X>_qWm=+LMbngncAFn3YX- z#TrE+(g_*x5lzf)U^v)_l#lw2Xj7-_uUZ}S>7mMna=f>3Z@$GrG}QbE!mlL^G#%#! zvrJICLA zoYtG&d``qlf%@k@BhQ~PmM=?7ii5V2Ag6@6bPBG{9*$OMMp5*Jt?AtfRMZfg6`Z&2 z9l0st0nWfF+^*O9o^&?j1y-8nK!m!4w$5|dgM-`VdT?$7=rhxDb!W{)Dj3um>*oyncMGPgYo6QLgM9WIZaTAsh-}BpN z5EJq{(`|G^c1R?3I~FYCt!|`?Zjvy&OeNusjhBb5xggs3@Ts^`AEh?!XBdW1lSq6^|`w+8M zG&sDfYYk|-sgLGRg_)LBu1AqAUc>WyJr86PJYE*Yn?;3~%Fju$oy-4I<7HI^@ z=2n~bUdhgs^NSxFPjiBnX_wi`$Y(P5Wh+j*P)-r-&G}C8zPHYx|HX>)% zc+Xn?Ffxb>*z-V!-hVW1WdHQ3lQKdI35p`mn*bin0tS z{#>_vONci5E+X?0D)NPMZ5(b>nF^Y=ehOy(O#0^a}9#i8pX)T}F z;UpCbBC2E3nWDJYY?uLNTob(h_y+5;So>EVFL~w2;fg3ixF|Qd<&yM(A@Ua*b`anG z9vOxg-3*+ZbPQUmetMoJ@sKg3>|dx`1DiE!mUT%^rX!~^c+yb#MW z+d1#&CEoibLyUiU_b`(%YhX(RiH5Y<59MlaUd!uW)m!py~YY+U1^Y2*O^cext! zlB*y-`g6Pl6x@H7Jn=rX43n(SqWaM#ivc_r0cI$_)=a7H_+@$e=LL6cFTRoO_cl09 z?6LofjzF8b)RoR5>czoPH32b{UaOyY^ltGEd(%-SjpqMEpl`vUV~}p9nd4P-y-^Ue zjA|V8IrLYNjyB9g&7vcdT^W(oz3LFA_oB^CBl0zB60lOV5D$_AJw2<$uk5LMUWy$q zD>~ivO}JV|{6wc9vX(*h4RdjOQTBwXm;Vsg_5^z;{Jrxcs=hryeMiEiXyn2on2$Nt z;DC4N$uLn)Z!PbiwTAZwb4g+V&|V*lc`H~F{U7diE|PLDCS__yH_rk-V3l< zW1G33Jx$QzHe}lZ~ns5<+)u3j62;48VdfR4kmM>T|wFs(u*7 z5_<3%swjlgfJgoMJSk2!=QaYLn3lB5~08aPLV2iL{O;2JBll0<)WK?t~qtpgq?Zw}n}=R(n@ z&o4pn6kWqoN~^$DgkKr>lMKXN<&#dM%w#Io(~jKO0_neA zSxu3}z?AW&BO<}jtk*zfCC1g!5(k_w8hM#HkcaHAC`XIBZcl(sxsGU|B>uqm{qN`J ze}JV%wfS^Y?s50sCv`adS!Ao)F}3cJIA#Ln)VELPlM+8Xor(a1xZlcxcx;F@X~?|5 zW*nCEl6SH{%(6lk0@e%yltav|usB{rgau1{e#cGOp7X4_Fd$=%=|!haw*Ou9l|g*X z-m0#rI|qjYY9)m}2V$gdGI~l<94?9$pftq+%RtFHrApQoKz&XX^uy09>N_c zvf$33%T6~dxw>ckqJX{v*Pj71V4F?_M0=pHlY2)`F^`#@g3BxG`-PSQr7yq-D_hQu z($Ck=nm=hLNE>?*=1`_)9Pqqi+VAeybUxl@R;L|rY+seW31vofO(c{J+U%9ke*hm+ z0-;E4foXJ(9-@RrubZg7$uHgeVZr|)(pm{yfA}t4g2dKZI{o7epg}Vbxk;?7gNI#c z#Nc+qE`$wz_|g4b{SppV(mmw%`0hWx+*NOPlyal+)#&YxL=OR8nB|!Jz{TKVCNk!H z1g?w@PGWbRyEaxLbd;=s85XHr>K8dB5~ilOgm!yVO=UO@t>pq59LxxIghr>XYHSXB z{<+D-mgVuP-(X-pZ65p`Ja?VwwhHw2-6ecLz>=-HX!g?9IT{&H}Mq>62CSP;oyD}PYCyJ#8aCx@~`1ub5m!ha-OL(&8 zRm@Bxy8QdM04{@D<5}X1A{#K^(Flxjzl{Vj(a$3oupY$V>-({9WCNStO#c{^baMlg z168}F1t63;$B6Jpd^1n&xDWXJh0s~jc9p1PzxL}5+ySss$VR9G+Kv<JWZs}PKJawa%Qq>cSd)PwFQ|V6~~Pw zg!eljA0l3mqCTwsG#~AcE$N4$#0(A_d}V+CURo?plk$9`5SLy4LDUZ4sgsm`PT(}c zE?a$-fRb#4U~L56+A<<5y?b2z@b!^4`%?nNS5K*PMp(E}@DXf_EebS58k3jH8UWb8 zewqdxLI|1Y2M%0?iyy|A^1sDPu!YkO`2p|8x!>MBAHKoLJB_r(Y|7`oT&Hvur}6f> z{x(B9Mlf}~t@`0%lv*(UORKw{OWU6Roo~NX>Pq^S>VOWb}Q8Jgcn#Q6dM2po*GTSZsIgma|UxF zqon4bG~ax&e4eay99g!JVogq}R562opuE|rynJVF8<)c!W8j8yr2|$B%51!n5O<|K zlZ|(N(<=`Dr)S1(LLvyO#_Iltd8N$_^jQgB2m%QT^q|(=TM!9deELbWSn~DtMamAG z3dP!GR7j|x;44wZfZbixXR9ERYfDt87@Kc6B|gAtCk6eEct_Xx+vH}A)P&2+l}4E4zgFB_imLL) z7*RSj62sY5V78B11Z)`%U4E5fhzM8Po=S<#DxR>Ph>RdFh}H)AM_A`31qd3Z=~9H- zFj-$ub7DcCu>0UpAJSrm^F&%~L3eaxSwx{c0S8AjczuRh%L+Fjlqgzt9k|JinUb%k zbY4tWn0e?nT9clhBuhOpkyr}>c)rYiL;Xh{y{3k^L3~dw2n{8jwmdtGUuCH}s0%vf zkfgzN~efgECxsPO9rC}XlDu9tE%l=K#yJ&ui zQiuu+;>-POz3$EP{G z+yfI3=rGc8@L2Ob7#-BS(vVs!%BPeT%HnWJw3Hb|B(Y;EScuGdjmqf7FX&_;INbS}RXOZ|zUxdKqkyw_nY;0%Ck`8FL<5E3Y%cGv z=3rUH!Fo<+)iN!Mxi}9G%mpNhyToUyCs6AV&Ir9)FF01|aVc4O0s=k-J=5B2>TVr7 zH_8u3fh4elLgLGimDT9~RL1pOnvla7Dsn`3 z9!R2R`$c3>ICGtR`bCo}@=fjBW$t z@PsWa2h-8I6DIvZV?fmk!mLB*2GO9yffFV5b!Ox@t(k( z>+o*DHV(Qlw*p=Hn=U=_{%es5S`V4yphDn5 zYC_TuBqk)ZjPm0YOaCv;3dKS10q@+neKZZ1J!({c;b;=Y;f}px+$akU?_63Io@{-b zZ!sZNQ5elKh4h-KD{FUVnQUj~3-0l%8FAJ~gj*yQlk!n)swYQECi;}_Kx0enM0DX` z3)aZ^J9iHy&Sx8E*VTc(8XU*osT=m?L=HlO3GY}=6mFZ9lFZ>;G@p8mLpvp&x#~X_ z+yVdB1)#k_w3aSd%T<|QqjE^oLHZt9lV^|j?{ej5JgS|QZr7MxPI#lrTvWZyQqD%H zPui+_xADc1WwzPyYdoWF(v<>Unq1at%Zb_pZOyr3;)=2BG@*DrG6#UKib>If9Bwig zjbw2lWUM?NJx$lpy=cxGhyFCX1BvA4oiY{e{!r~`tbCTCwskO|tQtWwu_}HQaLvxmIE!l*|WXKMds(qu58hX)WUJ1DJtf?GgGas zCUhuFjCus1@bQd*Ay!WYBBEjpUf~-}GrVj|+WydX+@~X5+bL@f;yHR7H};29!B#CO zk2h;nF`Q+=AnEWJH#&xzq^z1B4JertuPxk9GwKi;Ob7MN3j@oTytpk}!uzjYAU{;nQfFwx%>;DGQxY8*49o2&eCH z{o84Xv9_dgil!?vG`yzbNEsVZ(-dDsRGQ^jTd+b!*v5pzE+6>Hdf!P;8A2s@PGOy_ zud%P~1j7Q1*RYgWA~cWS&+!P&=c(ZgZe~HsZA-zdO#vbxxz)))-a{5DQ(L@h7T4-v zDsUbdII_c0ImeicW6(`PVOP(kY+0Ib74x=rOb}d;RotMGaTMW!+2kXKy$V>?L^Kei3bt2m z>sl2<;EMw&=7ob2pA4c@$(5}8J%gl>^+xTgapR!&1{Ww(OXhGC0_wnwO zGrv;#DfApuulV;uax{d2u7wJF>{_LJ%`OkhdQfzASmrs@vhHn@8L?cICzeUm;N7KI zV^|Bn67e#&Bs62u(Wm~ZbJKi~Pu1S;6VxJ87e*A7;!~QIJ;;1<{-2!1%WY1m4#?e$ zo9e~-n6cP)yU6d$bhp*OuQdWKzdXSPny1u%4{yHNndoG95Pt4G3N~Ff8MUGh`@QY( z3me&cg>OB&UWRT;vP)!zT`2Nc&Qc6cGZ^1qZJ1& zu7*`vWQv+Z@FnDnDoa3Ui-p1z4v1&v08wZpIfX0y(T?n}RrBz{*B=~v4Rc?#op^j9 zp+oU3u0ju%8mY|z=v`eGUlYuIfYog7IYPH_m})C$ZPg20$tXUFg^(8Rari50{?QZq z%vasKG=gao<5Wic#)jr+j`8bh-V&OcJCHvd67aIPH`YllaNKhTzfx7?hTIDX_f&x(>?*= z_rCH>9*^siG)as1U7Vh*S`7AMF}R=4%M#uDZ{!8f%YOD~XJ%-;%}X0W9giYu#^>re zK)8lsz$@!WmxA`J&mGbv&Q``^Oe-tiha8jJC!axcJ9!arFPdyz-##2fVX5`uRFd$n z?xJJXONsRkiCGNp0g&w%E`3a14JKZ8$0e8TKWvWH@jOujH0?q6&K4ws$3A1(M&-@a z$iQvvTdRMM|R}#>?4L;YaZ+M#(XL^gENXz)BRJLX8Smp6ORU9_ z4OCvf%oA^-cIK*eNjwv(4AxLnO_h6Wgu@C*6|mzJ8TIRFbPaf1khp!H7$J7<35McP z$^c2ty$S+d0#H=}7J|OVDS)MBTKo{WIay|(hki7{2sP7t{@s>oIyl>7bdy1t@b+T@ z#p`i#In`k)7xR!+LXkn27^3}I#VOI=M92C&rlU`maH6nbyw#N|o5o*5h=p}hs*{)@DE|UC5`RMKQ?L$7Zp%CZ1^D z1cJI$7eS=GPa5KWt(Yh(nWDZAc4DIvFK&`Ac*q_GxpOj&u3n z&VvCU9pj)&#UxdgiHL7O$eu-jUu$5AZQ$b4=nSwDQ=JX4REq^sN^wY8mNH+%enkJA zCdP2NXd30vK=WCLNo7q!-~|_qO#o~@^)V;-V&`dVqkG%FU=xrX6xX+9n}^QpXu$o= zZRZ)EZmIET%;F|wKFw+MS=REek7mLMo8>4D`nn8rzMQ1+PQq`qj59ZoCQ7(2d)&FO z4s3QMLh;O%b2{;$3+tGxOV-+eiK`Ew?UYUT)QK7SR;y=Y{b`|2iTe-8fB(Z`ian{m`UM?& zP-%uTT0GncW~HZ_uIMH}(I1s@O%|>qu0~&A%`Y=Q_@EftO?nb)%*BcSE>^j6fP?Fcm|f|Cp1*9j0h^p zDaz!@#si|>X*Q*121nc*cx{^Pl&~JI=Niq(Y0v>}(AibzSLxFnrTDEv_yvhSrorVs zO1>-tk^;=E`QkkAD!;3AgU>u>SgJerm z`$+x2I8_Id;5z(8PvJO_IeR%a@T`+l9B(r&Vc+ zn8oSC7wq|fqG5eRg$**|#DafoMTXYSy{7FnH%ggjWW^|=EEffxYrlB^RadV|zgR|W zz4ip3!OdC3!t<$zcOG<5sxW3o*V8LEepNq^7R78xE;=bmOLK9r#9Lf;|NPg_QX6J= z{#qpb%|f)h?@Q1a^n(3q!GCa-l=r4(>5w_h;r`WXd=psM)&fM|&hPz}+tKio+arh1 z*u@WLSO}s5$pkn3^yV5p?w1LJw&!%1hZ_+($9o$uUF84!2YG=+j7OL&SO}gRD63>5 z$eJ<8zUTJ8Bt`mp#2&>3MDA0y#+0_Vffse*?cKHw@qn|f3olU+`2dt9yAi`T=}!Qp zQvm0fuu<|lcc?vH9Ad?thNboi;gY1e_ah5Fx`qVfb$?ZemCPKoa zah>L-XYJ-%WQT%CTKQ^Bl@SJgzH#$k6F-}OB_eo)kQYWlI*R4KD3hjSbJeDSb=!HPy-4 z8pXNMkH1Tztu;+aGxYy`bt=6bUF=6!WQrJl^RAF6sJ8wlnbcIKqBXxi)&~bH<@mVU z_SZ)=GEz^Kl05*_S~MTCxK5v2;{h~ohVsGN$Opn2<5KP8P-*_#lzNMuKe%;{<*h2` z;Czj?2sGD!p%6a`##%ip*^wP!UE{E;{nV9cMTezXJJx9`LJlV()cVt*Q~?t~Si$0X z1cfoGl88U`@U~3j53l{pg1s1L&26jdZxZav&i2OYX+LpT&F>{X6?8$Xp)c14l8g_l z>9bVM3^BBf;iH9;rv!*77=qMoe}PyC(dHm9?GAGLv5GoQA>!M~+{$pqoEdd*QSMp{8y2DL11|}~3%uQ1D#F|d^>$Hn@1jpe} zTw6)kL^WjUqx@c|^iG3@1@vY$*JLvxp^hT(Bpn6v8822mYiNT<@KVe$uI(fNOqNx% zAdneo?n}y^(Py)fx!g)Z6?(VKuH>4N>IS|ULLLKz0>exKw@E|QmRqScEF`x=tvU6K zchKE$wOg9@s=S@RU4DD~kk<=qQsBCZg5@!tWu?Ar6m%NY-LSK-SlCuOAFVIOA@)9R zG@ldIaE#;5Xn$T?4CT2AgI5)AJ&_^-R6T{gl11cI6UVP6jLWlXh!$gxY0NgCm-Nin{< zI|h4Maf&nEO7ThzYcvvdHSNYv$1Y*B;pSeZPHRGBW;=teIEbPQgk(1TDI+B9+QZZ? z0&0_+lbUw3r&A_a6%%%8NfD!cH`IVFq#MTN5nCl9 z`8Gx*N@ii|jG)w!M8p5gMd!)(7J)l5V2L>XLa&V3o-t3Ns92o9Pmo}Upa-936rM51 zyn;?mAPC$A%#nWbOfEeqtjTL^#tR~UlkwyZl1Lbm){;lHLx0Nym0$4UXgz`H`+GV% zM#k&&@}b@wN4PPJGQH!%n>}uOs`3+rJ$2Hzc;Y^M4|6%5`Mx5GFbfGk{j|{SIyJ~D z!mdQ1wbE9if@WNarG(w8|2Sjy%U=qCW?sOkFeJ3Lolug=V*J=vM1C+m%G23b&0I+^ zzl;;}B2Epw4G8{D@GD3B&0F4eJ5_cJ6qVWavjj!G!6Y2XklHaV2yQy`|9 z^fIAC$neZ?X~iPSzC_?>>=l)VBAO&zPcG>W5i#`f>G#Cs?j%h$@(=P=RYWomMKf56f>VF@ zCsO2_l_gxeH_}SoQNA?j*p~r6)!bOmkzt8LP&|5f$L8D1M5sMbbz7SkD)OJ*K@SdQ zl#*MBCmDjf^^%!On(W7;>SSK^;c$_lYHkhZ1hC{n)U4yjv7`k7dvYVQJ6{)1Zo{T{ z=ah06JxoXW?cZi>I(!WswyId3m^2~^vU{kBKYV&eWFO+eg0eP7Dh=A*tY~tqH9EU1 zQ(@&pB!qo5)GiU!D3nnlDT95?J{NWrHf(Xmrj!u{FHk(G9j$c>9TEA@UFq=Sd|ti4L;d}d2)BMm!r1f2!gO@Zfy!!w zy%i}%5Bh$J(#XJwxF$ow=;b_TS)Y$(#C(V7hps~nvg)vcA7otT)>PGa?=sf1Fbc=Z z-%s%9Ff!wop}piH+MqxzSkdDG5;$5)9Il3hmEKqTBiR#0Kn4!-jfDEFC?kL^KKm~v z6=pAU{ol4#?0?Ul^wK*TzD*Br{AuB^7RY%66}Ohpq~m4m)!bW9l`zUie|Ca8mbYf% zh|nASsK5v#{K4=hiwH9-<4%u@u$G`&-&amfhr&RstfNsfhE6HHR@ezGbw}44MHeyP zsV<|LSQL%deHTfcWVgyDy2HRH3X~>7z8Ol;9x_~&5|KvW>lz2%=SXfMF>Y1SKNACD z<1w;kE!G;$t)}?Yd1RxOl>)w@$6@Yg7oY(t0p;=(AV1wDPg!uRdC|D>sC$P+?qX8! z%xvPz6Qy*2q{uDFc;K-#1$|n z@Pxu_9i{z1&w-Lwl~vn!osTra$J}^qW&YL0)WgngXpXy-pzfSU49D}`S%jp)ZwE)b z*k0@Ujhh7IWTX}x#kP8LE&c`ZqA~+Y>6tG|W{y0zYvW-CoC?j1VBwM$Wv$xq2&FRS z1yRKF;z1r7UR1qHlIZ}7YgJ~!=#^VOn(i@ZZ3=aDnfc*~zD2j{NG-9BiaKqo0^2tq zA2#0BtZbNe)|~(+5-g4nn`ktI^@5Nig87hN?tNbXe;OLH6#)v;ie&l^5?ut!)`-Y- zAT>MzZbae4QEM}*j~HrB1L~iAdK;BhrAviUGKbh>)_8w7tYvK%{ zdI?smz!fGAGd=#o;Cm`4LIc|!h2R2S3kn3AD5meRQyLjb+NCN zm6Q>r(lpfD`Lt1eEdNhSxYdtyzL+o6d~@}+!Sa>$Cln9ZyGt3O!*R6u2ME+baKSnM z<1qxe&Yrr6ef@|yp5{_6vahUi5F;8uQm`=?h!%)C;|u7bei4sV>e2N%>$ z{=WZ)`11oYmHB@L5dkC3c&>Y2{|x?-I7yR95rx%I`tL4(#<4vV+;JU+);ACKz3jl^ zR{9`kwbjKz=cTc*Q)*UeJoi_(-1qK%bxOfovtL|ve~+agkm1LYDD0=i&ku>uiJOL< zOKgOavmXRhEicX%aqAD(s33&4|JbM9a9zQ8pVSd0QTvxjN1OFp&Nov3t{et+fIWxa)u%Nfm z^RCaT8bYpEP-)IAtGiB|#A&7u_S*+b6nv19)cC5{&wjpt4+JLcy75*42;aJP-2Z#6 ziUX`-agsWcy}$-MI)n3}wX&Iws!!T#XPW7w>xMY6SbpPZD<5#P_!$8naLbj-^&k-; ze&3n^ciTvOH8w+#qu@b>4+d}}9P^tYDP^bS%=?HAi9fr)$qcF4#Tllb93ho1|hjbfAk zMf~ya{LncOhs1M&CB-0(60b-HjiZeu;Vn^|)9I|!YxF`k&@<8HA|M`zI4fac$vL95 z)5P;A+7jULbTz&F`>*iVY1h@jhG_85KZU=Fy}zD4+7tmpwXh$D|LraY{cva;_#<*o z{c=dH>Mh?1U`It~Ky6<4?Z--FHc90vdfrXwGk4&Ecd_@164Q48F%a!zb;eOb>#jBV zdYk=1C^(2ywF3|&_fAvQek!aEuw1>K8zA0Ft)S zh;iYoI`HU{Zc85W0}>qG$t(MmPUX9%WuF%RUU2V-sQUrn0m6U@u6i(ux~M1>pzQlG z1W*a@aZ0)_PZId|dt{T0L7xaM{I>0$sbM5rsp7YaF1rz^R8#QGdH3OccPlDt-n8dU zHwxwA-n7vT!7yNUS-~WsG^J`u#%}$T|4nnHRmj$`w_f~lTI202&E-xRZzS8)`;DxZ z;2(xYP8kL2a5uh&xR!+sB6a9q&dJI+vBexSDoJvT(<-p)|HD^&lXRT=&YUwny;WNN zTzb7DMD0*kP5W5hbmk&SK}Lob(OJ%qIy-hio`qPqOc=Nj|M%EtO;Gz)ivfm4Yur|I z@r*cfJe>T%3hA~8WR|2K>A+zfA`Nx7#)?RelHB?0UfAkb+<<_fPb`;9E=nh+QJ-Vd zKR&It3%;zkA8vOyx~2k}3P(O~_!G zq6H!bTnaBI8ZSt0gjP-Dzv7lnVVoJl8`X%|!b4#uT4i-D_H3HbqZD7>N>#wi9#(ef zK{8{QDOwFv^x&gM=}z-u(PkhFZnhe3y)vP`B3~jArIlP~V4H%LD8HJ{8co0Y5AC4j|!8h z?y5qB;qpV8-&pZM18XQxW_NxO|*lLZHy^N)O?q3AoGY^_(Jd!t2Y7XMk%?R1s}d zu*zyS!ZYsZa4bDz4!Hx<$j^v>^_w=*9TahWbA2WnIjN5}JW{J0^b(~QJ`EoyMXnH3 z6aNap1-se`&$-Z#!2c&I;|DZrwAB;qRu}yER!ccs8cS7&Ars9<&l*h{FMbj!c)fnA zJ5Iq=U4ZriOu$vi5R)tHuf)H#pn{662`o``HK&tLV1PE0@o~?@@&^$-(eZ2zC3{U^ z;=*@3(XSt#CqLY6#Zg()sd4hyCW32cG$?uYju=YIwXUxrX8uf+m1*K9SztsCEV!mb zo%~_!?{J*xf#)&bT{Kp@B-?V|xgU0bNL|AvlLo(;G>Y3g zbqS(t&1I3C%Jf=ufH2S3oDgEuOA>J&#A ze9H0WVtLcllHW!vY6mR-QBRJf%}e;gKqqT~fsp#(ppL#uhc^)m9qU%VZV661EJ2QD z1#!ykPhW>592>RkJ^Vrk!{M8=h{9OYv7TroWB8!5LSP_J{n$kvFae;p>SvvK&vl7O4-j2;U)Ou8R)#3hH;P7|rjDw6- zo-ZPFl{GrSH(N=}8c=ZO6E5)76UeNIvNC`AwB!V)FKj64Y?EJcG30_@dO+gW;;T^i zBFI0T(YNSy`0Fh&V7emjx(+6zk?A#b1X)x*RC0w*8>{pN$Ej@C@FtMV*4dF9uu|Ql z23GS68~YU9W-rFTxn(7f#Fys~Za=N4(z+q}^d3ZFP>_y{;g(yn(8e^UoFz#>RY)(< zI@yXh{(2x#=m-60PkxMN*^5L?!7t&+#l+QjPA8HVGPNX`)QJ!h#cZWs|orP zJ#o>gk{g?O;@sU_dQ&&n<>3^ba~WJ0AUAQY=}1Gs)8%L9GBP%$6XjFND!R&3xn(eT zOZZ;LMprn0j#l2#8(n7HdfQ8&{KNgaLGa$0?*0g=8Pt&+@8j|iEe#Q^@XCeR$IXQ* zL1vH0Tr|b@!kZ?v)nANcM-?j5;GZX?120dP@Yb@|)7RN(b;3!t%pJ2{Gt8~KN)Pep zeFkAEI9^fRNW3Ki9_&OF;uYDW4H}qOJUMlnaite3Ptq*SAJ*IWCz;tmDOxTH*b-Pv0Bp`tYXvTJTO4w z_2-*j3>WW6I%Z6kfh9u}z1}XmD;~iOj^L7cNbSQa_#$&z1Sr@UcR$)v#N4*aG`=N^ zBn;YJx(XAK`Mi23cf7QS>KXCJ-T#vbDVSb+!G!Cfkb;jcH|>P))YhG0~8Iwa&Tajh)0?Qz+gxlLBLgy29|6 zR&Vo2uxdnNX0ZJJAck@bO5GKVe6{I`JMqU;LgoMB>Z}5yYWuDa!vI4`gET{TBPre8 z-QC?K-QC^YNT+~wNT+m33kVW&>$`alzJu=&j+kNZ>-xu9zqQu5vXyqsL=doK^Uf83 z6XVJ`gp9g?UAPpO616OhZqsaoX``6c2$ku^FWc?8=bEK(+3)VNTmp&>7_G9%f2v{C zOp{eTAoE_HGinweZ8u$BUC|lIT5>u_X5w0uR~Yg6S^R|4Jz-76q^iHow|aqx=Sj-! zn_)d3$u!qPjveg#n^h>q1@^)LIX?E5B|CeqA;a2dM+^l942!IOc)VvKVmPUV-=1PO z6jSidpI`E~!yQU0gyxora#5gNq3dVnH6gIpV>B-Gm zlKg7TLV22w>ZlRnWtW~qBZB&AX8{K%Ymw2WHTwM`rAUQZBi=jF8TNdMTz*RSJKd~% zXAO?OSHw{VET}QvOU-aRqDg)KU#wOZLCo<18iN=c`__VdVz67TB9T zhFDf0X=~zXgGQ@U^EUY~cJ=oqC=Ou(qbm=srv?oLxI&L6EIMf0+zVHuV3&iuGWp9& zXN0zJYpVAR$);i%zL@`UDfllQpH`3lAefa@sI-ih_G3@S)4?(%{<_9azgY+CH?|#o z*8r|`Z(>Qno*CUoMfQ^^lw_XchlXH$ z6CEFsQQ02AZY%a+?W{T8pMQ}|6~QneDJTwQo4sY2K@wNmSY)JAt-w!IS55zRX2csl zN$!EqN`b*1E(1zOUr!+iS8EkM2o z_g#eH9|Gw$&0HbuIjsATRiMV&@NeRmA27mxC; zUKA5*&7TH-n*_Z!Euj(3tXNK^Wi>+fZSHLnX1&sRG2VkTX&W?Kz(_i#`8vgrC; z)l?A^B2;z)iXut&nPo7Gt!8Mltnr;87scwwF;n8yt#1&DwjPZtNNOShBv;%6w^0UP zj-9Y*g+P;7!G=!cCZuS_l#odY^R@oO?y0``_5rzX2oI6WsuwdAL1 zj9pj>ubEO9wxFwOxG+@KF3UBb3#nZ%kqv-Tg$sD4QV`O#1qsg7~@d zsLh$tEjw6wWN;XdKR4l>@F84h1bQ39oZnMZe_^Bz=D%Sa83HrLuxMHLdqfu+vt<)_ zqvYYVNEXnNRu3^@Pg<$JZTM=OC(DXN8Wr!MS1aPL%=JG~Q_lqT^Zb=f=B4rj-6MG7 zE!pZ>Ab@I%D#wTaSIIbgM53Zi26PN`vkqu*#D{Z+mVTipa1aYyYC$b@_ zo{5@Fu4}Y*hf;m2F4d(pqp-N(rXB!MY8~JfP|5*+zii=;bT8uZK^Cxt6X=PsfV~8T z!ih3Ghvgpp9IeahFrzRke)fkEAqUW%T| zu6rvObKHTSUgYI*jFG_I;WkyIL{tDTf4kf|b4)q*mWjiCFrdVKH-kU6p2+ zr|b>ybQe72a8wlKm|8O9bfc4Pp1BRNs_G9k%vhx0 zI0(n=qyCBKJ&LPM^&O6>4f82@^zlr0cjCqquK&J5*lAoYy*TRfB$b7oP^LR!Pc8ja z9-SYW-T^tCW(g@La1KP4Rm$zow&2>)Bkxkw>G(5Z8)3^UvCqZRuahr?MMtB}y*F<# zLl8qQ%PY$FAhp?v5P4%d$~-+@_&kSL^@3co(<5;8^5F;}Dl6rQIf^~%s4NpI;gI7> zE;rGK25Z_jRLbD=UO>sqmZ$FWCE)PCyJ=S1%MCd$W&O^Br${y-cTryGg$%kC7cFru zuLv%gMKSePO=9gd)TI|`UR;e)EXICA8W*~pb9o~EZ$IqLvcd(?TaPYq*gvXuZzclm zWc*t@d9+|VFq_>kNjn17uJl_;WBb9y!{~!d~<@4rO7AL_TB&o2WFpE{#gh z<%t^tGQ&T~5bq>I>is&ZfYxPJgdd$oI{0xWnD;k|m^FB}f5?Ku#@8YZAACDX3=P&+ z!G>}D(endc{($m?{UbC}J%}O00rDSSqA~W+k+O@!=e>v?l|0oJlCm4uo7UL*x(VLtqvA9sX-kwUABVZK)#p?Kd&?U^8{#B z-XMHygQcd!)#{QBZ2t|SmOXq_wu8m}XD;cjK#UE|OD-=OlvPXZNuF~JV*}@F{Z`6j z5xz)ei5#%!&jPV(s!kINDT@`?OzD8Rs#F+6O2OEoPR4(gKsf%Lw>hUoGh1wic+>#JXP4&vp%^C`!o0hHxX|k z7{e&6^^-y+A%?qikCg5~QoWT#IU*#ce(^Y;IJ$`!A{atb;{`gASScUtqPnSQnNHez z1QLSmxNC{cr>U#jTnhuZq)ng2X)d~q!9GAHxjP_joF*U;!ue@V_D4q;F$2v-ko|K| z^bY#+MYgFW1(FDoW;X8t!%u09cfzkO3PzS#l3(M;LviG|MD)uK&aYU?j+9eNz-)?W zx1nZ!rTd80_hw4cG_?X^(k(WIR%Jw>JW$AP0K{Fr>wCTfB%X3l$&s2eEn=7Xd&7pJ zpuFVuwib%^t$S{sKQ1!%=x$eck6cVHO>Mo1?vrn5<%X&`FcV~ACj8RxkesuZ{?|6h z^>b`%qvU`cNys6FOkcDbFXa$*sslWu1Rue%*$nk7s3X0VXj9XS&{qLXtD^QOZi4!q zz39r{Uk^>6`mb_oDV+9x4Gk;nAJKQl){*0qd0~oE`|F-jn|-4m4H^GO_B52szvbiO$`PiyAt;N5Wrhn%=+IH#g z?f%P;v5bu-{{h(y-=sy-j%s#ffc28^;{hErCBes|A)+Wp%Q?Ou`JDjVzhmp{d%I>E zaqry}BlIv!=~tt`uqp*c)%4fidoTD|3Sn1}*vQSVM5Jduu&mxFM-3U%sJMn zx-~t>9y$3kU}@izT6oGG$RHIAlTSv14Hhi3e?;zY#hA+sl+`Ri_Z&pZJgYa2U3O*TCygxv8#y2Ic^(Hz@X!#Shkh zp@g3@DgJpXhVCH|rTWMf7-%zrsYa1LLlS7Q6tL7c6J_zM04?8umtEEOXwn| ztRo#Ubj{yOo61t0nl7)j8O0@RV62})m=yn8B?O^`L=Gm33?6KLe==iEk}i~9M;E+y zP5SlWat-DBXjenE>zwFc;Tuh}HFHVOS%F=z^3B6H^XG_Rq4MWRY_NtBLkwji?{G)* z=rZE3+!L3cr=r;J9ggZ$ig|f^Z0ClLk6E8WGZuDy;C?1^QYf-UbQy<`#|Um=HoBC`y>m=zO68`q}-X(91RcTtJ)Xj`J&5AaUyWr@z07=CH3vTc$$Xy5xO& zPm7|)O~=c@bJF)stJGRs)B&3t!hN!!q%CT!i2lq^RhRoDWDf7v#^*F+vg$f2$({xg z?}{=x_}hgq6J>0s6BPfx`LR|XxLJgB*+P?!6r?$V#&fAk$PFN4hq)%_dEuNcBBF@%Y5T-DEg zvZM>*3T2A35=QQNp1+ebjDqeVLjthS_&%5FYm-|FRr8Nto=tVF7%@Cj8;a$cr4M!( zls>+v=Y1!BxYwu(Wo4fGV)mwT*+-~}(B!XZ-*rEOY;tNVzsx23kFYJrW#c9HmvJ_6 zV_VAv8U0;+dT953K%M-xLF}E8lA6LS*+hcLIQ{%5Z}xXGA9-&&@wZv3YSc+*uAoYb zSA34gC+Ar*<&9p2sTE@a<*qB3%NPuLk~b~ith zp+`BOnVKBa+@}aiy0^=EducSKi?4#TNfs-cq4OBoxFNL&ela$a))M2~wR&nb7;y~! zc-0K|niu5=a@$gqPugVDu||iUi#`zcEOTPI*(DNlnRo^tvKe=Uaofpa_F42Q)D}to zY~;^SUk;p}+{5IBl4}8GVOiWZxNR<+(Ig!toG-iy$v*HXT^3D0t*rnpkVv|@Y5R=P zOQiJ+g-1;tHu$*TQOsdm>1u@Im_%_dW`1=!VQsCazxb*aP4tP<_dwuziWOA{CVFRfb0@bv+X|U5tEPtZ3nvjTB zm_n(dyc5q_?6l81^P}hAjtj_+b|8=AWELFWcLZH#8u&ho?0d+2ye`MmknxXWx_ic0 z(4#M%$vk#d4O2=9;XSJ$6D`ES$b{L=(_q6;_6f#+A(RfdkrFcW!~I}|0i!yf7Mt3N z!AN1YVIuKf8vibj(V6osT*cH&-R9~JN-*uT*;T?)?&=(vPyA^2yG|sNL2nE1E;i7< zBXwDR_tuEKHm`4h@!a`;ro!N z>>gpXD6|wA-*B0t|BiA784Z)Chk9E=TFM<))Ch%^gf?7Ssp=sf(u*VR_v~ZZQOtyg z&A!v$s1+;}nWZ=|u&FY^G{&DpA(r(`N`Vq)F0|1f_wQ48IQnF~xi*aloG%Ewcu1m- zuP}xU`oYzXoV;GM2`MDOm$2wLQv!>dA5vfQ6`sTFv~b>&{15-S8}F;9%jb*1-%G>G zp=^$?WNSGx$E0TTBm#^Ju_Q?221DFC6^3tbmljIM?<3pxWCy?hJ&aN&-I+SIPTXzi zF=zxb$H2Tj*_)qActrccW+i2`4eA?k*WI(_dB#Q|=n?aeeUmw?=Uzg37>Vyy2HBVW z!mk1n_rZ^5X60<80#pkoT!;%cgbmM1JzP6em6$}-4@ez#kYEK3wfaxuOfa2t#C>aI zs=~HX=$(rj79+xhwz7Wt&4x+Cuc~>F&&oKxaITJHe!?F{Mj$ZpG`GcF8YW zs6&-AZ^zYZmf!?@gjK?jjz5O>;dy`G&JHJV-{NxBR6Py>rDnwWhtcq|@te&pnk<3& zuX3o8=-=Wuvrn1_G_hun#pB2Rg_9{YVX?sHH+B>JcgbM!jXtV(Mc=z_9#^W$$(+CA zKs*sU^{OonCVCNubf#B2m3tN2%#b>P1$RwIEK$2KsBs_*XFI8+AwX7jvHvf4@)6SoB4A}a_A^T;>k0Jh}NgvUV!_1YvbwjO+Wvd*;on-NwO;!ML9Vow1!vl}(Y5noZv=4*i<^ z3Mu->6&(GYfBq|L+}QQH+T`~5rD3NKRLhUH82-d7p#^_?Rco)TI- z$_=|NTN7GdSuBMz1XEwEkilylVJ-?vo9I?qB7#IDMu9n7H(V7$LCdlj^HmX6^1Ud! z`)JQ5ge#Kh?4gy{uR4V{9j`~JI>Yo)n1A_P#6)_}t>_%FA`S`Rhw;QO5sb7r{)?KS zZrM)32>MS)p6 zqTc-9=l|cOvCR&#C}|(G!7E9J_arQe{bFvgg{Adh9-=GVw%YC>U$x%EQabzwIi|cU zm-E(wcmu;wvs&r5!F;+0wU2ZH6LiZHb?aux^6FlTE}sZ;QYpL+BmAL0}-O>5gce%EcBY;w={{Gpzt7873)fo-)FemVS(i zhi4DYnt3Uk^5B40qBRm@jJ_Esi0MPhkkbC~Lf(x#{?kQ6{89(?TRnJh&5MC*9yozP z3&*?buy$C|6?w!%C>-qm-9?6~P@6Y>`%Jhv)3d2d$2@*UV4UkSGiGgeurh9N zrOmtsOg77qztQoL?fe2}L97I`r4DC+YIYCw#ZZpba6$msM@r@VQ%zJUP+=5#FwS0* z=Ex4mjS!=mylz`(B981TD8@-SDE&~iG+u#h-_5i-s(+pSfRephK>Cug9aBKG**5AD zAFA^?<GGujOFq5H5oQgDPpNgnp>_Wd z6o>{;)7z-AC%qirjv?5$V+b)hFXlmfu+)sXrMjBAo)OO^KV7KZuc3%Xe>1fQtz654 zoYgU#-_qv5OBPP31wOf0YbomlaPF@@c8!1RGpKMW)q>E%Bl<>OyF>|G0+ENmZEIvs zxS3gm)Q`V9-(Ld$EM9TBDKgJMzN0FW=ie7<4Pd(=Bezm^M(b}K-~j@l{%wk4(*s^Fs3)OpxWHFfMB+V` zkev|6qQKu9pZ@*lYnjeDe>!rkq2fe;Hx$s84`h_0W0Avp$85G$t>s20hg%F6?6Y`x z%U{bs?f?nTS0r9{^23OXiGpBHkFQRPH9rI-5)a zC^~vCKgSV$PV_s@B;~h_u?3@DG&p(InJ0b3GeD`w7yru@1YFW=l%OyEfK6$JZgYeX zL_`LQ5`Rn;RNAiB{~6RB4k^g6Jw$AW{Rl45%45ma|Kr+~yU1LIQ62W%Om)77LW70w zOLi~d=d-)SU|~-;zh|Temc~Z_^d_Z7X6oxk!!&^;s&8P1sKO^eQ#I?o&X7-$P1D5o zhWhz^L@T^esIu!(v~5OpW+A!zedl*a=*!>7FGb#A@*>LIhJ6;i*IkE&0C%Idh)jM!;OsZB`(9w?KpE>3gYt(3QU=M@O&yb;Io zQhZwdP$BNXtc}^=VodC_F7|!u3jPd;I90pww&T!l*RYM-U^BY6#<1=Bv;DKtOl9-H zK2ljTB^JETRRnFb<~=JEx}QBR+E(Ghd;G+_N`WPY;zPaOqA2i z*q!Kgk}L^qy2`OjQqAQQvDXMl@)_u6k<~z_cYp@iEjQUY=V-Y0pJT!W@ zf4>z&D(7%$8peQMZ|*iwc<}A@r1nisz)Jx6y#GHKO0K!Av5yddpsB5kTjRn-w5qi#gk_T#X8nB-?MaH|+8 zSBJ;$#C#rcJ~$`M;FuESMq9c~7C!1;z|~Kwj6yKmKmbh8K%82w*X31S=*$IW$)P;`xR&-VY2#&V#a6I1e3M|7^ifI9KYvp+p zM;s|J##MDn%iEJEQbw++|E!FiJg%v_M`ek2UcktR9ZEgaNHuAz*ffJgA%(>cxIGz}%r1 z(G4RHLH)SW$1po4>soWb_uu`}M~=rkw&PxT@Owmo!5?>}Jvb%)zkPUf)rHDS&gKSE z7El{0!q&~a)Z5F0)8tEw9Vh;2$j_6Ljh}E6&fb^fSr;Ibpy^)T1tKqenT1rr{6vli z%=x#uXj<>n7*6ia-AhEy;(iG;yD?8e74TEvG`LZda&oK>XH+ z83N*kT~eL4H9?ibUK;u)F_t?K9d5BKg>*z1}h{W<#2SuN2zr|lS!+1E110M*u91*PHu5XU4~^v zYP(`0iA#)1z%PDQ5=qj;GI9h%y&VPLA4%8>wCOsA;LU4q&s&GhRrHWr(typ`sFC+t z#(E*kXm8l1xAf%o|MU#&`E_H~=2{EZQ_;zU{yon<4v8`HudDAQ;A0RWH)t?3yoXOX zr+uc3h^+m{ai_9MZyYeTWCM!t6!JIE{3XybrVLg7|t|E(0vTcR@Is(N+&#a$~LsRW=jS$ zv5Br_fjK!O2D9d-mtBl)77Mp%%Xd6xPIKPjFN2C*_ES-<@Gu}N` zP%JfK)KAjoZkr1WxX|C_7!hfno9PGrC3Gk`3dt>Me1Lc3wGW{PuNyf5(LS(?1CeDd z3O6Pc0^n7&wEP%DklJYTXD1sA%odKe0~CtklsQ%D+z7m~(~;3wqOv-q@Zf-7xDvwZ z;W_hcY2Pbqx7Ovt>22-)+=E3}`fkL&xKX)GOs6^RN&aya*q0~bq-*urQGob8l*g-n z;iUbjV1b;_CF)8mIw<^n=D!!jyj4>(+Zm7YF_reP%ZoA<(V9*6dT44bdWNL%o-iVB zSMH`_w|yWZ@~xPU+Ye}_j?&ozj0C|i$M$Z$(cQQe-wBZ z7OBnD=1|Y4>zF(g_*1mMBLO=E3Inyuda9+$2W>BUg$!1`t&YHlrg#n2e#QAIE&uYh zz7y`y70Gdo4}5RIXs`Eei6b?%FMjZJCczE_@4zj%To8OlsN zw%ECz+vCHJ5k(4(E3x)HB(Qlw_}*dD{Kwnv2TFl411__J$67v(_beru5gZ z1!sm667ICCa>m+fQ=whFHEh& zuH1^JGrPRMNFa6K`o~Y;a{tBYFaAiax{fvgWG%Ofua5RGIooEh-7gy2^1*rj1B!{$ zlCYAJfQiPY@2>#R!7=_bJaHK|+OJ%ix$52LN2JQ0xOW?#h#kzV*!(n%u}jf!BNlUt z`M-w?BYI0pJZK=RWw*&5jhT`aG$U6!oVR0+@lb8N8rRKpNuqs{M3ee`9Kg(*nQN?- zdHsV351Tmb+A1-&qEY&Jn)ENAJFFp7QeV3703T;etY{lIl#x|xu~9Sc-eq?70r5m* z_(Cmxd7d$D#2ExM1x7r9SN6ug4oi)UMU)BsNh}!efV$ead&3?l4MCd~$QCc8bw#d^&Xim+X2*h+R&&G1B3!fXk z`!9%gpAgc+5y*IAi;-K{_HB1BFeCBhNKHl1PZE#Z^HH$Fj$OkXuKdjJ@MjxKI<_r+zt$sKD{W(jQw~P5Go|`Zdd$x)8#ntxTPku9t zc&nRyo+IF<9%s1)7Mt11(Ic59uASwx81p}!OqLBYI5|uKrTlcB@>f#}!zo8Qqh9yP z)?u`oVKU{=t@(kcQC6$Zd-osa4$HB%w~dHHx`T^M5aAvE3w;EdPc1rM9}BgoAn7L6 z0LFK|jEOwYJt%44-Qnf4ad&z+qnE>H=?7aI3O)o*QfDPwFVFFg|9}TV#wTfZ-v+!O z5H`=F*Cxg~6GSJh+fFq0Yi;M3*Xb`a_`!nYm|g*wCcCTB5si z6ZIop=R@mrrGM*Mr$*nZ_xQr@yMfqeYtYxVTM`lyV{TEim}iEdxBGYyA> z8k0%_h)ulW2o&#L@4GkNmeuS(Y)!VYdcC^SHJ}}HLVtXMv;WvJ<#ry6tRL=tBZCxZ1udYa4(3l$HsJWZ^b}&Dqmn@ZP}iiBk4`Mx&E!;tlL( zG8PgOvSQOKUjM-4O$keXKzFUV0`{I1%+RIdb$!6WBXj!Bmo>78vHbcmK zo0Um31i4$!XUhvIMB*<*Ux1QvogBi{mFbM}Gml7s6QpaNskQm<@!#_8K+zd?nAWa_ zr|wv&vVOuUvvNAUT=*r~t3@*7{2)Rpe*np}o4pHhC^W z)q|h2q<(No=J_c-2EY)BMiy1L&_YE@@aDc@+bF@mQ;_yK!Er2~kCgI6Q3gfQV8Awv z4CfqAXhqJ@0iyhqp;)4=PsuRmD4+^2d92DjA5P1`1Fr5{MsU_Yn37yw@v-QqYDbvKNNXdPEwp&} zr8WXcLe0!-b!|%~D(?|rrmE6#@{g51uOAs%!i2qo_YzWzZmHkJ{p^&#=^I)Ij6{Y@ zKXJGtr61F--G%`8VlFx#cmf`AJ3gcx%>Tr*iWPC&90L;^V5$ffzC+^D1tY*7-Ud;w z1b1I<-&#@vOr2l7KjMS9S|SGxAaw7T{+*fCw6guliR~L{zSQl6=!Dm2jnt-QPfK7M z!N6w38r9UUunc6Ma~9>*5$-iG$dT8?M}$$v`4P#URl&Hmwan<78Cn#hMJ8?>k*mSi zCaB#}sI64UP5Wb+y}b0)T-SWf|NUXTvJBmH^M0{rVN&+a29(arubD5&NvnBfvFv5g zFss}Y!#o*Y3q)ohiUflhx$51tCx1u62%?p;1cP$x2!jUmQOL|yZKl&~Omg8AdKFv5 z#^^7&IB88e0+ep8+v>{jOu=aP?nJY(%;D?iDGI6TzfVT>ENCpI8l7VMI*IkFPAL z8UJ*G%0oRi+|o&aU>PWWySz__U$S_u4sh+L-A!*|FU> zoS5_M`&=A7Q7*)vySSgEz!P>_`-4!JQg*nMvIQiS4U?xqS_f!mi3BSQyh#qEZ!!~S zwkE@UQnRXl$r*v~SRtN2(?5FNkO%UMKv>RTu_Aqm)pW|eYC)N|0|tbAw-CP?o*y5p zv$mC4QbI*7fn1~zHB(;#lR#p@af~T$#(0MoQc!ZjGDMNzzGv{&BA$bSSiGyV**-WG zHIY$oBPhf~3_WTzQAP{Z_Qc2pN|Rr%4)IWmdhvShiC#09P!r2TC~cepE54@*i9mNa zLzfNJSMg!OHI_bDLXVC;xDyPeS+H{9o|~N6}2mGmrM(%iXpW2;#;JrIyn* zNoi%XM``lPBS}a&{8DS#4btstp_2}1UMghE1AcYXMXTwiNlr1YCLmN)ibuX6W97x7rOnB*!BE2RqE!x3{z7kQ$4#6e z;rc%nF=@>X1DC0suyjrhME*f4tgW7cNtOsK(}PSVolpGhKlW{P=Y;hN z^|BH6VaOft!dOmveB&KdiX=u8mnfvCEe*p9)7Mb_7U7k@fE;JYFl0~&$~j$wxo{*h zCVx_HF}}P#xX+pX(tA~s$=Hw+*z=| z$4R${@Nw5m>)|otXm*?QSQ5LJOXQ?y$C$S&IO5%Gl;B!qxu;I%1{ZmH9Q|6E9 z?L3_(x8cyxs?)a7Rm&kcC6)0u9M4>dEADfKGExP$aq-zR(9?Sr5=Gy?lsZ+l{{Qlu zoyZL=3oOP`f)laLuG-!%NE=Ao_D?CIl#@APh{Z~anPmZyF3As-O^HsT5FCYb1UR5< z>g>=z=BHP=A4$!XVjlhE*!L3C$3HiIRNG`!`#pW3mZCi+5)PYVs;o9yY~9DBm8XeB z>hfHYXcniE_y^1^{kVB3Iv;iX9?XD|*)k@la|63~Azr#jgC)~QA@md;HZnCQICG%i zjFDq#ARqKalY$8T#r8V$j9ueUl9NDsd$P%fu8f$yZQ!iXd<`RKW{FA`+|!)y!qQ;l zmxdh8G8j5Nt*IKn=OA~ouNt&a;8JL$^pEuQ^i^^aH~d=70Ug@uJ3|H5XNaq7{o^>!7D8x875 z%`S;Q{zRbtl{nS;^{pEVBL}n=L>A<^M8BZPmbPQmc0wzPegYUWq8%FEOxg?xGW(xD zeviDmu}o>)sAk+w=Rl?}+Qn841@zGzKjZFlDOjCH>&gZ-&2cKt-jVaERXDh5%Fp5i zH7`a;{`B=@jwFzYIctpc(rB@~Af1b8k!N~Y2)AWHl+gb2PteRJ&JateD}V;uAS3g% zgl&?D{{4xL$-7DP+mZ(75apR@30{Sl?i5(+9O7slDkKd{VaJNPZk$Zgz}8PP(z@d+ zx=jYDUmZL(QwW;=R(;$sAZ~Wi5jtwOGixfzNtC>M(lo_JOsSKNv7K~HG(w<`EHx#3 zZbY~~#?Ka%+AqiS;Qr=4x*B>?MV!p0GPUhMA>Je9E}KQwDsMnC)cUd**hN=SICO$F*ysL6rO){H3-s%+Wanm}Ffe zu|`U`UIw*~+Wj;gRO-tj%^{Qi&Sfwd0UDtAP^_fB87CG6R{Gg$ml27{8 zB0c+*t6SQKT{s7W#2r)>{^0O8Dm0Wd&S~x1_SR)<68f>*^AG+2I29eT8g(?uHo~P% z+1mbN8slBUF2=gY+iP7(^!ell3&g<5%$RS<;UCj+b9r;55^)2y)E%QHiYK{VkAJs5 z!eEIN!N5@A61xmVdwAY|TMwO6Steg60_NyWmYRQVOlZln1PPQ0{c?A6A?A|CKXqQ= zsH!NR6_|o{0DrmEkMs_}B)^p6X{Vm0F2*6ik~N=?NfV@b^CTdPJy0M#SW%Je4mU)M zylZ8?X6UV`Xutv6p=1p$R#cAH;-AFJQE7PG{10Zidp z@f>>18<*_q;YZ2_i=&izx@;T;WrcGDKWCt;g|C%iGFD7agl{J(!c5qS5(Y6E}e zgKhcY2sPiHm+9R>USBu=$IumdPL=lkGmH(j{STpA`>(auX8AHhY%OdSL0YXHbtk>-e%&+}D%P>sE3mwbu4)D5cVG zJ%|6v=AzzC0s-)ED-VCn39~fwRigV5QgeUDw!d){$)!UsHg?hQwO=Kfot=gCDdCAb zsWWmfo{yKN=)l%lU;7i+2p8FB{9!?0 z=UfN`W1yey)&{`20t@Fez~&8f(U))u>B~|sbUOg&;SpF0$0^9XQ-Bj#b`AYmGm;`r z>bn<$NVpVFB9=>l&+qy)7=AJo<);VQ+JEO=T<#V*6aKA49oao5=8e7fkMw{^arD^R zX1kn&=c+8b_!XGQahm4Tv!Xp6Ef^4Tv~fTE)&`cLU(xxFFecu^t0o=6?N@!gvK5`w zpgxmo8C=33+17c>VBqA@w=qe)T>X2m{O^8V9R1kzuEgqA4_&nQCU~w+`;kDgxgDvf zqfp?B1HK_v(hP`bSmeBjfec`cT_?1T_>T+hu>T^IBQA9P$&U&0$r*TyHhH}r@3o+q zuT+d)IRip1PL$qKu#c}3oVn~lnL3nj0<46{!JmZy{ex(>>O@)GL=wy-qnlG-f6n>u z`myiVH}$Wxy*JClR)LD!=5XRweyXmV+e~i#ijPKrtxPu`knk>z@P=SPaxP80BFGfl z>ZM44XpS3OkQ15Uo#lKO+3VgZeVaD1F%89*yv1o%MzzmG>q>ulaSHUQkwQY+J^fH3 zTMV|RmXe>FSe8gNFFB5-XJX~dVo;StY(~}B4xpEFh2LT@~{ST-0#Cm;iR@T|PIBiPSmb8sls zG2X+^AlNTNl}5nmO2*Lts=X~Gel#9svslGmXEqD)aNH{UPN~oP*|v#(Z$a_V{xq{N zi4oKTL$(_yH>VYl`s?MQ$y$sBQyd~Z6pOzf15>8#x0oh45(AYpa~e*r+r}GTPhcy} zgQ4R5vq*h?9Q_sB(clwcLUFuwS?sT)@~(XflRBz>Oomihay`q&l-$lhBmi^=9~Ww@ zEvOVG0&;cF`uBj<=)M#`hW}|fX~mNy`=<3Vo#>>C_!l%lQBK8)>D07+t(bfCt z@R7sj+qsUpYjWi@I?F;;$y1rP>j0geLjX9l@@4Gtv+p`awiG!^OkUHUySH(kK19N4 zYH3)U*pNvZ&&U4z-#6dy-fkT~bJ+WKa82wSFhWI^{_!=fzhiSpvw%4Sa=|E=Vk zN4>GhL}wZJ$XBGPKH^xrSO ze%aa)nV8ZXS9& zgl1gFqo(rscW0sh%7C2=07L+9ap8HggVRSU$SKLTqA$hRk!-fNTnPU&T-xs%sjD6~ zEf9M)xnM-1Ze{^@SzJ(9cqq8fA!59r{RZ7U;pZSRIa{W)fDb;l2+lBkN(^s53YPEr z<<1i4*Rm*?XLZ6RnE*rB?iws_C<__^VZxf5;si|me*HD~Wj~V0=eV~FQ|6Hu@8wOs z0;GJIyO_$M5(s5|r_Pt9nteO=0kn>#!Etl3jx=T02z#Nw-?2EgFN0~QIBQEZC^DzH zq9fX=&r}VU5jSWE>s^^yycjVd7q}%p{TjHnw*hJjD)m4{1B>Y5Ar%kxoO*opiIHf- z0Op0*<`u6f3iBh{or~3d-tH?+4zyRrjk6D^B*uOqX}TA@veAiS*!n_gm$Z`HI}LG&vAUkH+3%ojM=Dq!N<`Vzb|@?jRmCS^|5X% z*?!X;2+AmVwMs}&x4@SI?3kLT4!`gI{lJJP27-XZz1-=gr>R;ZyuZBIiCapSYEfCT zdu{e2SbKWQ+L{%+*Or&@d6I^9_D~WVA{U+LR2erxc_!+R9BAc~w`3lqpn}|l#3i+< zaEBbuwrRo9=2S~AQ+TuDE6nV8y$_k!mN*46lhw@?JfO7U;K4Mvu2`RK(INg8q}kiq zg_Qt;CWYk+3Y%3dNdq@I>l2N&$@9;J;ZQhGoPpS{DRxQH&^k6mR$;h-JDy%+PrJgB z^hn#zFw>{OWikAT2%{%ds849L_7bJhNwlO?)*|u8MEPP?T6dOP2~OlQ3T8|sU7>RP zvb*qrVeed3#sWsH78wO&KQ0gUs@@-|l)DG((KjBC2RJkp^VF-o&~&Sgn~FBdmf_|$ z+!VSk)ZXW!VhXtG87BV8qjhe*+OZfQI#kVRxs`8{1~y(yis}VGrk&7yZL50~U7q+& z7@Nj|3z3|LExag97s=wRsHKqeHtexQDst=U3709vH+x-EVE3a?$)e#`${&L*i!Zt@ z)!iNBfL^u{(Zp@<=iwC=d_`;BnR>5WDR9SJBa>=K#|+mA)^Zf2JP>PQ>KFd=!exgq-*Z*hD8=_M1$nS^;N6g0a0U(r#4`zx zM!2q7WI5VJD3Mu{Y?_EqQ#FV>~si0)T@fHIr@w55)Y0&fq32j}-R09pa&A8nXFLu2ad@zDuOa&?KwMZj zeHGPTz=v-jlJ|iyG3n>nSDRv$kmyyCf}z-d1Q9x{iFVG_9;BbUdIF4)iqMA@%Jat+ zctvHIEVwE;xyi9R)f*7#ZZnthMSzCJW}~|7AT^!dslLa zOGkYFN7h*e#TBhx8X9-k;I6@40yOUKu8m8u5FkKscXyZI5&}VjySoJsZh;_yFsJXG znVPTWCshrnyPC7l-tSuLSp_RRI-{Kt`$9?{qo>2xXh4Oj2Xp?xD;+*rxpP{nyuDIB z7@R4L7;!^_Lv}<>2@B7ZENa7Q|M~c&MI(0t2NSQ5zxwF7+`x4h#g;zGLh841DXE_n zcLmMZOp0}jo-H8*S*1tpE1T^TtzReJxcPC}0saGW98czE{H=D|k$kyqHEUab%Pg0} zv$D&8=fUTq{fe4+c;C~^LSq{AbL(Zlrwu1Q_%_Pmf6QGkemUXcy|$9Q<;!~E&)?(f ziN7`5&5{~QHZQ0|@YB6x|HdT*(zA=)v~+cDZZN1wRQ$dA-QP0hxAhWrsLQ4~ zy=qi%ix2H}N($$6@)-o>zN;NSmBKOyKMEKI7%>e@I=poA2kI=dm zrb-RS;fXa!?9FVC%G?Ti$R;@>o%T=DcxCtDYayEhN!m}5qh(vZCfq8X!6?oWXdv{* z&{~+VAw0{vn|F$wv)06mVVk|{j+V7qxhRUcyl>;Bw%>F0D!}Bt6MFh3o_nH?;?!k;Y>HgO0eLFj3w3*TQmLS_&>Yn15 zS=`XYT5j0+SvqgFC;}#oUQG}*m=#t?{ivMTpNbZ0EB^N+^jLzuyJi$AV}!92{&a`* zu@A@zB~=xT_96`R9;Cc_H%|x=l)KZVAStt1l9FQ5t&r_uxyHpOC2tLkk%$wWM`TBl zh)gJFWikbmpyXen3Z`<7DX56I%Uz4!V>gWOnl;@q@5bt-o+=Ta&!5qKWWV{TA35$+ zrtX7HE$I6B+~p6jB@e>IoRN(UsvNA#DuTp$jBZRySYEJ8b}qklrn;*!PwgkGwTbmR z(6cV*sHYj-*?Ma{XG~xPPsK^MQA4e~g=%o7S_JQ}NM+!__!KhI0jE<4tN37mO!rP~ zV`$+PM6Lrb)h^^4n_P6vuoC*ORXojH^j8Ik%~f2)1{^s&uHaAe5av43dvsD!HC$P7 zuDaH*DxbjS1USHCUI)IpYRwj|6Hi{y+9OmrHeS|{p?ZQyP#G@M^Pg|OP{+j z*7|3l(If%do6)5@*PGbQnCsTpbs6BA)vO)H-t%r&%w(k)DNkZXXtB{VoZ%b(v`6cC z=31q7hcqH`MJfy(s9WIyD07haI!&eCiIKr!ZrqEyf_51EwIMa*QX%pm#Ax@Q*|fsE z!^kqu7F-j0aQSW~kIURzr&*JqPI39^^>2W@IA7O=GC1fELUov759LWQ(f6WEp1v10 z8Sw}&NOQWHxVo-r;B$lDCsz1R@85I%>k|iPQzE2CbIw-w&UNg;*30xDX(7q*Cskr- zH%2_INqy%0D1c;~vlC3jO+H}=TsqWjN z(Kip>$9)dkte>p5mv4sQGWazOuna7!fnjz0x9TN=g4_VaML7<`3)f>QODnve#qOz)vw77*@<_Ls0Y{!STp~r+%(|chN z!U~*t#!cG$v$ut%y>T)!k2-oLcdoC0Tu&GYge*3rqB7ua@hyIRfebZ%x;QabGhE=x zc+pQ(`|d9bNuw}EP|DATQQzxf9|o*0hzZ2GnrT+s9_ZB8mj`s!uX1H8e5mI4Jhmst zP|UzqT{D81mc6jIrJk`nLN^9hZpqIoZRTGU9~+I5uc<5e;5!QOsIt-P$%F`jXu_I- zAbOcAj`v+Y=E@7ys$2GBF_W-WYdaFb;{H@2IQZ(yoe*q7qL|B|IYvzz_oDg(*zwz9 zS$-dh=7G0tS-*!Ci)y{j$x6|nhpOb&UP40_PC3kHV5TZFTRw1pgz?$Qq0R0+M8sp&qeg7P1GP3J6qx8{`K8WQo;fr2IhNa zmyGGAg_p4gKM$4+!Ggxh1(S*3r5fsne(b~pV0X$-3o=M%=ghq~LsF)`KSf|j)*P+P z#jO_oB2!5m4r9^s6=SrXuVA`n?tw(CDk?7oB-a2rh=0{ZT|a}DRZRS{D@2Mp&#A2_nbq*c6dR%8@^lAw2OVu4?Z)IHHRemW3A8q(EarvP>`}2NM3x<|i(Os);CK{@^phQTKB^jFHP!Gu7RKlcJ!R=zL{YtW zI6BeS&HPi~fqxJ+s1Y!Ij^BL*AD>Y2GE4cQ@vSKN&rfa?#^Nn=t~-Q`=t%|xGY&4n zJ%L|9KE&u7e1xc|RK}akFwui8q2(E6q}HX}!V}`#%$jQ!rW4H|dUfO)oUR4MLwK5~ ziZH@fUF1yeKUOXF{QaL^mOq`J3&4(~ETiG*WG2q0KfO*}`exn$GtlW_jCP`V|CBh@ zl1X<^`a;*@wnKu>Ib4|(+83#eF`%1QkP~Eo(N6^8rAh1ERA$N!0-oB_IA>l$ zdnA==NI@-{$PI;gfQ&4N%+_Qyq6`@h?_H9Rz$ww?WMtP)7yf$3LS7C>mfaDZ3}mIn z$bKy}iAhGXk|}vv!pjKahcJzQ`vp?h%obc)bdPy&!VXkp6Jn!Mxw9>tX*ud$ym-@UA%3_~85=#OTT z8ff%iRNgP2vrvs_W8inF?!Bwum58rXb?eXr>kT*cwfH3O#(Buq`$Y!JWcU-XQwdtG znH?n8n36h5;kC{)$$5U3**iTIlW?HA6JThuxWFZ&D$FQfz|;_agsTJQBlD)U`*Tl} zpR279^!Es6+O;!4Y~nN|yRFmyQ4xgQLVLQBYE5u@lnDCTU&XmS5A=1`X-V6bk1 zT$Nc~GA@?D@50fyGDrg*z#a~7ej(MO!ui< zSsO)kvur&eWGQBrlitqc98)giZSXlAl`X!n%P2Ctv+q{xVwnZg;(=y)g_Ss4UeUht zPV5l++F?lif?s>DqD)8L6vvE-=Hm)aJ-YBGNabxzji8s3k69Dd8XQ{m`dexT*Y81h zOJ&pza5xW>&NfKGv&sVS{AeHpYW|q7lhc5jS_ml75FpG4oy9|cg!)rLig7T4({NVK zI?bkQ^)Nl)l$^X@1vpzKTmQ!|7*PWGzs)DB4H_q8d{9ony=gFS@&EpF-YoEIPpw)Q zK4@*!4c^GlOBs_(Q_Axj|KnQ!&!x#Dp`(x%6AkG9q$rTb`}|BuKv-fT z&RXdD4)%-^kseyV&2{_%UGj7s0dODcU% z9O^|Uk!DrJLgVlQjce^@nxE&(V%PF*yyXqc#X-;UjSzK=UQo7#6rJomHV$S1J>27m zJgj`YMB(#t-9wQ{W7L9K-$A>t2(N3rSLq+?k_C9d{}R6svb*fa?)Mbos_JrV_jHr& z##Xr{3*nW2Q_-YCw|7s#@(ZK5;aM#McyF@>^6t5Z)k-FjxgBehV1^sB?DuMW%5f%V zxRr#M7z6DU0L}2PTb;`@tsVyaug*GKNJnmQoXkb2+l;^~EfM<(*`16W1!L$=(0dQ! zM!0J({SCMMOnu*+v|?P^43@2Dfo*A)Bu{r{ydG=PS*X>sMNQl_+QLEA1s}diYT%8M z3E{^IU67|?o$dAbO~b#fW|yd@hO*pU2OkpCPoEB3EB`fINMV)!c)tJidfz;zn##bt z7h07(BEh6~&Awwe)-4D_xW=2zDcJy`+0+W5@+=zQ0zp9Fk82>^;hfDgTZyG4mf;Uz zQ-TlfDA92Khu;WW6?vr#VL)H=)?g~VG=Czin?c62f>uwl%$eQsnWZw2=d&)9^?a71 z5cJ_mbX`1Ba;W^$T$LFLg0o;{J|p|#QTK{Yp=EKxawdm1^u8Zu1}Q%Em1-Y2$PVaK zz)0h}2$GaQBA1!RBu> z90;wk_Ygmo7AtjF$D#)x_tHBfn}Xi`1#sez<7O8Cf_$Gxd6g60i)8ooO@Vjr@F4 zw4jSoKe!0ZB6Tc=)BmU~0J?B==KZjZ(&5QkLmc|odcfIXW4JbAF?S1;#)c-mOF;h~ zl3dKi*2hcdGz7t1KKKc~P~te5Hfu!4*G+HQ+@eO`y)%G^Z5^7xFzPJqjt$?^;rSOj zqyG9$Ww%&;7`lV58#=m80#HhpGa3L*a)~sH`USwK*N5n2dT?wdD~~37rlI%`-J+gi z`m}@9J*u;R0?}Mhyz{eX>5jEDfTvwgVpJwL@cF(+4wd_&yKu$*9X{+bVVNM%p89+n zH`#^xR}P^{rZy(%T<8Eq<|snFuU>Gl+UgyPwJtXIj&Jr{U}^xSyopV0Ir}#vuFM6d z_UfvNB6wbhQ(WP$jy7*W_Rt%}Sk9d}07dv`N$lB43kq#oKWNC2 z1L32`vrLQVn{V+K8a^HKd4h<=54rw@boY4O?!EMwlt4|XM?yp71YcoR(A(p=_0Ymm zdJ+FEj~o_t;tT+cPUSb3^evM!=`+jeKFTFf?TUB=+;yK0I6HAOvD_2Fw-~%o+7T22 zNskoRGINS0VIa9avF&;14%3ehw;WqKm~@ccyoO+`DsP00C*5iHgFvc|r0 z4UH;%_?4G2-yuwU4c*cTOAGqEZp?$#9DBbSdtLs8j8iC_77&C4(ZdcWS6SfdzY@Ma zTGj4ys=1L>&>AL{?0EsY48UD8dq2t7jGPj6M|SB||55BJ7z;jf-c46qifY_pOshh1 zj-L0^<$d>coGT1ZM?)212kE*7iA{ChWxLmm>35`>g}Ka^hI^>bw4|eHV%L zKncQ*r(qaOpL-k~S1B4Y5?>chtsU;Y?xO7&6G5q zWw(9BhdL@q8c^X-rKKj8WKctAb?gi8z>oinj$5ik52yfto9I3 z^#n8^Od0))36oKd<9|9xR*InTJ3=$P0?C0$orFepc&4}}qDruNJKN>NPBgIFz$xZI z#VIx6vJ|-HgI%n$JqoB7 zIJGENgV*(n)?ZoYwEPaFo-05GzV=JKij))W12JZe!qeP%Fkhh(zN3ynruwpPdn7r8 z2|uiT@Qv>R5KRvohDaVKELtRT9YpV$Li;c5+n#XV>sSivPB~pAIbQ&S0%SH#44=OO z9n=lY)QM7zZ%}+bIes?WpHt${Sum6-OhTnqk_ug>q~VLJjg%Rm8W0h~Eacm&>*0wp zX!+uy|I-}WC(3eoCXBDIh^KdrhiSOFW8Wb*1Poe?7PoNZ$cn5`4U~;q0sWOI{euX& z)hb*et-K`6SV67`qsU|x;NPbQ6e@-L0tcT)qxow9Vt5u864CMT^zVqAxF(KTI7+~8 z-mb6X+1y>3yD+?(lLLeUO`XU^rGISquwJjQn$ZY4vG=2P5d{l?!d^)4*`1d0pU03i z^IM68%&$ll1%bjP&~oRpfU$D@31)>9DHZ$W9d^QCay|~f&!iS(u=#dxMb`f0dd?}( zK#z%cih<=HA>0z-b70D+86?qdSkq*1G9t4Dj|5`C^&rtH$YHfcvY!I+CljX~L{;>9 z01ituQ9}8M0Lc;pCLiZ|KlzWP``eBL3TCpRv@9jNgx-$%5g&C9AH(6S9A%!#5y)K< z5psNVxSbT7lYg>-xI)Z@YUKTRID#iM3hRUGEw;kVd*w4aBk^X3zmDibU)R5tr?s$- z#qm!XbZ^N&hRJ1~W%K3k?8p$Az|-({F5IA2XKRH6OSG+}I3VROg(;JD1DUncKL}gU zr)Zb1eDnxj@38107ZSNAHp!ZKnizV7SSr`#qgpHhiQ~gIGgDvu(7|j%Anzy7h3!!b z({I+w!Ys|JE33BRCuo;EtS{ir_Hd<576(hXP=%!QZR41 z?%pJN0rI|O%AoeWU4d>|%>!Yr_D9rms8;~U_1t*RYCWf_j3Fe?3e{SZ3kK%(<+B$g zHB+$W4-kJjY|6)8u-i)fbjVXMOJpYvw=p+3>n13EYB22wFevf}JdqbQ^~V$tC1SWZ z$w#R~6CyaT6i6AKkCellCv8XOZ#DHpoAk|)v1d#;0n0rBFUn70n10)v3e$6X>;j;Q^Pf-WyF%DW7pEXga}$^0}gMt0gy>1xchkb zF=PQ0|CB@Amt@P4+l8+vC;OwuHv}=VQpYrhysaCaaHPK}?YG?jVF6T{ywd8$)f0Xt z58JyVN@daGic$TpdY&R_l3`6MaV3^^MHP-avWGo`pk5_HqWK;yXdo?49%Zk(r=TjZ zkW}D*A#9JvLV3k^Ol`^~%s`-^h%2HK)#yeA!1pIxX05W7cu<@eq|$_*SNKvlUfFHf$Sd?~qqKa~yh(vZfEYc~ zy4M_zfjAfi5xvixEkzoi1m*LO{h|F_m8n(9^fJ~L%6|p#yE$F&3|p60KhE0~Y;N*V z9#ZffM)nBRI80QxE_*B?%l786JrIW5(ArkS4(mucD4<(zyuD{%8=MX4Jr3!;xw!WX zcIWcYl`Is=>eM?(njYx>(^?;}d1`8BcMxUICoO;#S;@He=IovDh5jy18j%WV-#T$U zh+oI+FXfJgSjV1j-M7sUuC*NO?M^kV%ewED5uEW)UJG_EJ9N&D=SwTR!WTcm8Izf| z*t2c$x{r7D#$b;ys$Sazw2E?5CjxBNU32&>z+hZDOKTmbQI`4i% zDwUGqs)kzLw|=aQQ2|~#Isd88Ev{>L1CG@qV@*Mb=^2-qmVnzh!p_HSleIiVNJ-D5 zK>MOnVb9VNbrFkaDiS@Q8xm_zg1;rd#0@#Z&1fHB`1_#cxaS$a`p;}4k?8wudZmVQ zZbM#=@^z+GAI0xGi+={ohXS<4H(#D%t=BR@WDE?;Z`NqoUhW_{_?-MjOW{0* zyIYO@!TzkAUVSF6*X*TJFOAa6ahGs%qU93^8$E?+?Sf(Ouq3r+8Ks<{3uDaS0Aw%s zMk==_nufDLH8W8nM>`LA3!l^}9y(1J`6v zK=XQIPS)7RR7u)xIMsIg`T=9j`38%Nw-?| z>f%^MHf(ZtkI9OqhE>f&8ZnkZC8|xtOJVoHTD!Xm!~qB`9p2XLKeh}b9R|p#tEddI z^yB#(&$QQ6$!8ZU6QicPna!)?3dmK}8QWZ*VLU?DoDiw3t?8!RLx_F{IbKOhXS$p) z6n~?_BJ0ANPef1%6Q05B_pz=KZKEV{9ZQaSU-*Rcz0z9y3YI2sPJkC8$ZLp&5=qwb#BnJa@yWbGKOTHx|>J%^QlG>uT~w(**f|4LE>)5zIS{j zWcFYjRt|TrC zn-&L%{XS1zwX9pr3@Jp>@X?=SZp6457yhdfkr5v-N3}}B0tp6w%%dhr!b*MMN0Zs zcL)F*@P9RI_V32etRFOvdIa!T6XQ$ zeXq&>aT_dF5Pds$2u$mv^-2sl6*UMPkuAk*PwJ*4YvXI{=_MyDUNHo7(8eG74pZqw zOvT(y8U<@Sl(S%u;64)4wGU`nn0V1}l6=_Q|J-TnsNETZH04Hh`8!cO<4tw5yHw{r zH8RVWH0Er72$e@fq-%K-VNd@3^HT<_d-5$xn7B&bRm-JJ@g`a#8ppU+7Wv${m81o| z#ON6s7<@`%x60&Lc&byY&W=v@7+%f6+BYzZ)kkdt^ECX9BSO+zc-tx~sj9FAGXX*t zXbRrqlOCJEWsoox4WPqyR@qg`QlZqnpZ*1>&@n&6p}=zrSJQ-Ds$ok%cm|mKTBqU( zMMvRBwv8GWy1u$T844Mr1(ot({#aPpEVTKhQaEJZ>p^YK!Y$(%LIlStmSB}$I#(d6 z{xONkd?I~-^P@J+BSCFC&GOtRNZ2;nQom7@0%q`1{h8Hr9-PUGrrFPU*1SaMmVykX zP~qa`O;O4W+IlxIrcKLvMcG}Mh;y~QZ`1DLEjt!BM+(vc6_ATPYKIS>#SZjWH#7&~ z%@;xLELIeD_EZUpvCA(IRpt_jqBdq1V_+i!ET8~!FI1yZT@xmaP$TB?k*Gaf6`7Jg z5c>%QZ>OR@i-q_GMo8?ZgN9^_sj7!83>vdkdGrQCZz4Oe-~f&PBR?;^?T4!zl~4`J z@CqjPmn%N{C;dc06i$e|DH%#&^&CF24Eh>1m$QT0>oM1|)ega&ejqEXTL0VOs@{LY zJ-)`+E3}{p4jv^VY(Y4hZ3%EQ=8;X z`5%WhYuD%68-+agUWb8Nsk~DDG}uV6&DTo<3+Z(VFVJ0s%tN3#eT@Cs!oq{MpK_=s z^C%R9HD>b~v@age9w|7N?YWAg{9#{2^(TN-VzwzCLb%HPB^nU~8oVgR-`K1;Uh#O) z#iH;06kSrI$pQ1)eblMoo8_#`?BoCY7w~^RCOYS!q6)X7_5>wuyE37Wp_wo6N+Fog zXLIcS>DegeX706jc!qn13aWoTHji4aExa6<{i9bt`X_8J#a!VQIK-2`^Em`{M#Gg) za{eD6K-XV|Yn_~IiWRaL80x0JhgBr(*M~dAf%I2dt_kq}HS@U{Gpi}~KmSO19VrJx zn7pPRdIl+Sp8gGR#j*%ZKE4y{0sYZPIwgN};5F5s5uDtk@;7K9DP6oNSudX*7TA8* zXT+%gIr$fn_8F|<&=GKey#rVZHdV)E6RHu~yS~NOQ1(Sgzdhq7aNM0Gb2h}14G74t zU?;QZ5B6J@{WTkK@t!ir^ldq7Lo=2OzRFXK=`Bs4oc1*FrvC3e2gF|)-}mYNytX+f zV;F@Jj&P^7x@z2DJJTryv;iqN+-BMG3G%hox3&n`3?F^<%zyeWC$B7s)IPJSS1`)F zijB9U-2fv8&jNNz$LY~zgzei2v>l^zz=WMZ6pI8FE$+DrTh+x00Q?>*bo&HtJx=+w z^#I_qRf_VqKu;aZ*uPi^taCuc0C3-S#n1 zG;r}C3L>_X97+wjH2FK1(XH(~w&2Uq>u-#& z<*j%>I}4$tvhD@Dfp3`F5F;nGk8ATzafp3w-Q90H4dnAK^gnuhm3BXa|MD9u)Z5H1 zQ+4eU)N}4T@I2RH3zebKyE@Jjv@!o#=Q#T1u+dXMvT?i1`KSoUZ3*Ytr) zw2mKw>g;j`IrtO*6o@V0HkPLHXR0%uJ5@d{jii1v`h0q0#rT+i|9#jhl|9CS4PDL6 z6gW_Qy|kHm_Ul_@B>jR!$R^{6%di)M_j6AOOQpQ0a^DKAo2mQ!y5XHtW*AUQ@W?+= zc3Ed}dTPUM!HD)Z^4TM9Z)U-B_dyfgPVt+TXWoZ}5{ku)*BJF&){1q(Tiz?^fCuZa z(ujCXI2_GR(p+*INGhKNVy(rpU7z8@3J4Yd1+2dA+He$)osu0H2{$f(>io&FHun~@ zvo%eN+O*t0(G!gf+uXB;ey8={aCw@7J>=G$(GKgS9JaUv{u>w z4qgA1=YbuqTeL8wWI~G3R`-b@p@qFkwY8KltNpqHQyT)6sftqxJ{E2R$BU$2upV_qEHgQuPy$ z;QCP1URZU3)JV3R#~#92yJ1~#*yYBL7SS&Sghj(l;@}@R-F&w6)53#i`*$Q{lq1Cl z6z5t?CPqhV2XzlN&TJ#;UxQ3iJ4$_5cxAP6^Dnis-Ctd{g&y#~YEUa@p>;u*k3#H~ zKME|2f943`3E}yvx64nh`K#Af=-z*S(1E*iImt#n^hIb!>XED_)Ckm&dl)#bq>T&z zP8|vc;ifW>kgX%?pk;Rid@)osuvL8l%SbjCjTQfq7AF!2|EySjVwLjrfP@FG1E8L$6ch+(;Eq!ANzu+k; z%wUr;#@16o@PUzvbDET9<}|L?D~8r)R5vV z$ET|ZU8+30-9w0BhUR1y2~(K3XdqHn1Ix}WiMqE*Bb2RNHk6w4_acLFQ#u~I<*up z>}+MeO5=e>%hiIBRTH66uE8EsA~O-}*TcJHFB4LC(zN>0u@2?CLj-a4rax?ZguOGu zBcUTksv*2)H&Xh!5}D3TaA;g5^ExVFG}ct%M;&X76<@j%2jm>D6kzhhdBN$4iz27^ z&AlVds~^}sAa+PU7%+oXBkJfTP@9go^~z12CQ=VhYB|b-NErFoD>(<9a8CFpBFoIM z^|1vTkjc8!*`0yA+x!=DJ4q)_$IiFo4ms~%>))RK$ww#B7NK${R3kZm`d(mas?DOT zAgNMct}zB<%vZ&1fKReWlB`pRt=mdmTl-WQbAZd^U0_v{+Ovno2edCMGg35#0xZ5s;~Eof&Vt-| zZ*$Tf44NyN=*Qy;2&P!Z2}aK&HmXRgkOKSKl+T@(y++r?bdm6A&Fg$lQ`425PnvqY z>_SV~7{N<$sSkbB#$_(iU4O$EFRyk|Pq}(n(LTRb02#Ohb6w0J7Or)0ZV!=)xb<3) zMVM)1;ZyciD9P}tNzif*Mvt`+q-{6PhoXj`=I&;4dqXfn8o-t`vwoFHtAY?h>10dr ze7;Mak^w5+Bk6A@T6TiK9V#0{WgAV$7LH+^_Jc$t}Upf>dpt3(j?0#g#XHBFW9{( zVsjghNSW?GdI_?A=OhAdR!mIL_{IJ9#31e{2@LF4fqQKq^?-YMfn%3PBz`{}-7u$) zFv{MMFjR-Gz8Ya!uj10T8K880STdE#+MOC1&PCq+bFu|Z0R zE)Aw^TU)73E3qp#knF<*gA4VoE_PiLcn!3xAOZAW>7^BqoB$ud9(zT~^P5(Sm*{O=5Q zUkWB{W3Ft)pFva}$C) zz2CTzzq{$iS8}FucyPgc$B}bKrr}q(HF}ttV2aV z>pU;$*oxxf!BsR*LKdk_R2Ll-TY3xDehQE(>XrmiGi+Z4>x6z8v_*zx5{eJeoGr&$ zr*KeTKlg2#c^mUgvGf{=L_K9Yk1m|+%uu%4Lc9;@hVS8yTBmhwkpWSfR-RtNhCl-G zXbWMH-k%+mj#h~~rEom-k6-J(X2kSILfC~MXqnL(h#$zpxV!szsi`Enndh3W&XT@Z zbNOPSJJem<{rG{8Ytjry0xL+jQC*(%ay`^(E?f)Iup_g8jTAC$(ZntSgJkXA&_3+q zzw=tRwIlgEw}8Dvi`Iw|<@OqSsci9+sd1{%;4iQUC*6RW@L}lDlzF?!?n)H!*wJc6 z*lzVX@`f|Ove0sPw_%l;$iOl4dlrPz?! zkjbp!`L`A^J!X`g&_HM)I~RD_vNGN(U(FDv4811Z_-)5|#WI#qLW{D5ttkT)FnXiW-jW&IV_GDiEmuq8o*x&WU` zdL@E5cW;~LF}1cqvyHi zRq|uSWIF7#NOgFY}r89nR#FpxPD38?d zuTV~BS`1UNt5V0yGgT^!+>@x!u%+=W!ruB=#$qc4!i10n36j*vC8 zhNQIgPAqtOxA&7-ig{D_Cz&=24jMn_4h>toX*I8ruub`^h!H-{aX`NoUikFQOUaxh z`h0|7gkUqmq(zA|zApMVGv=6vCLy|=+D^>kse%MXZ5WWn9MtmU7|!P78?ZfdGS5A^eR-rTNHVTSc6Kq_2o8O{l?B%L-A(MO7sAuRa{Bbd-K&daxrJ zd$aQC$j##c9_6v48$_(czLmI=3K>wShStm*h;>E1&8z0D%_cJ8> z;E7#7?72l7pbL3LzFQ|Po~zy#GkXALU%<9=1evY#=Vwt7=zU^m56=>N>kW!dU?x0? z9ZSkQI!VKfV=295p*SfC?W0bl6Q}Xh*Bwrr)#f?L1QDO~FDL3D+hYGn<%HCR3WweY zV}rStGxDB3Y8MmONh0N5e~`Maai{MV^F{pFnQTMS_&sD=b*{AR=uGFtSDz{$WOTXe z#1e50Lm45$$dtZ!#d^Ac#D%Fah6&k0T+|2ycmw%_A6$+1KKh(+u65vks^dAB!FF_t zANbA6xqYh|<(8}(d~##0`{ZF!%@amw!4vi_4o@tW_w(HVFTYGBk^aT46GQWEWB(Tg zEj-5ZER}p35!~5W(hxe*;DTcMVONm|bc!p}eK?HAF?VZg`JADj79blH(>Ge21mc-3 z5VVQD%j_(4fXetVf>P zQaN*ja*#X5&l`#*!QCV@L)q(?)#CG9+zCIcP>yS&H#8OfXa8NJl%Y4nkK7FUYJ}is zUeF2pqaMpDb){134dxyeq+BLJ;|<~+?;~oVw@n%?OJE;!Bc=8iq9g~ zDaCu0GWiEhx}p!%6P-mkTmSHAh{DaKi>X-U#U#cIeIEv6=Y9GkQxp-5S&V4vnHIna zBJG2##ZBS^O4s{!LY9*|PIx)1`;C{%$x|JZd2T^_V0gbi$sc+B9>e%ivC$QLy%`BE zxQ8YQ1{f;DaaBcVZ;BcV2VpVMFuzIck7L(jZ!7JxdvIe`VW+s-ugTDc&l`WaMcT{a5p+nrjf@0N18?$R+w&P6b~1MzmjaiF!b0cL)4A4W5xMZwBlIr+RP)uLYxpZgY{tr79b=0U&$)? zrs%ue#`N59P*K- zl};0l2eV6H!!X-|A~A#aP;F~(!qmpoD0yTowVWz=ca|(FP-wNe2SXm%dALF4VMo}r zi#cv+t@4@N*4ISg{Zb~EVZwZ>N_+H{$NxD=v4J}4C@`HEl+6G0a3_UdASnq~&!{Y` zgW`Ari%fkug}p9JmBvV=)B=s-A0^{mT!Ho<7;cdy!qXq7-q#pyj#=DcHy=fFp0qr0 zmr~r+e$V(eVdFd6-4rzs5G9&ol^t(2;OJ$E+=J1Kw`~PF4OCHhTV7wBUEb%!n4DQY z6Y4DH^lPeP5+?V@81K_SXQQw&wA*Mv@@=_esr7kGwGZAlN8{jMoQr44`@B=*ma(sS zI_W+u^L#iFnB*OYeD%N^$xqRRRbyRacJ8kG^m1x;qM0$4(cdWQedHeb@((c40V;vm z(%oIJx2L-g=$`|T?z|6M$)cBV$D5afVk~r@mQgC98&L~uZKl=Cz znwZdg`#b%?pr#sA$hXxBji|u8e#*r)kR=W1c=p-W&#IN@(W6fG( zTQ}c-YjG1fJAn{7Df0&I1>hZh^`5%g@mrIHvG59GHQp&rQKi97#tfsz8RELTK|P7x zlXmmzf3wwqz)2%CB0a9QL>2Z6iaKugkAI&Fcf+0RRC!Y2megYYC70FoIvOtcbM9fJ z%Dx!O&E1|}=i1`iPmkl{;}h`+4N|Fh2pqnW6~y{d{c7IdECnw!5%Od}Vz+`}A_F@y zG#p@<0AOHE?iZ@3o<|eg7F6#Z=a>s1*Ijn51ILZF=4_%37%y+zD5g>#Q`5g9glj2sWvsDmGCjC2AJK?o|kP`N<;44|J= zJ8gjO4e)rdj0_??&yP5Jo)+uJKyoi9gQQX=n)Df&PsdOiiwj^+71oAh9_+OV>;aN9 z{Np9YdmRdT@0hU(yrGAYOL&PC-KcLTI zAfwr8?cO}4yn9S}_wf0hx(%ut4G3@F<)z8>k4lTc3bean(nEH?k7``e0=~Je+x9mA zDpG?jyCurMGcTS70&_<4K{UE8J3&K98Q`cDjDM>IFC?ie0t4B3xQ68~!f(y0&=VL! z>RSCGL^ghSFKfLad#v7bh7o@P994kncWTd1EeUPfh2pjl$>0^U?C^`mFxv6Yf9F31 z**W;cn%R2QIne#Pq2X8L0=GUCSP(t)$b!M;o18Dd`)hv5r)<`b^i%S7&4j(Z1vujm zsamSAq$aRt0OH20Eqck2v=j4(I({bb%(R;i{RiZ?lFkABU4dL;N?x)Fo1d@kqFZqD z03nuEq(xS?i5r1i8A=uM6SN^PP#P9C+fpOCE$vNpEN5fG<*%iW8i%xz+=Py_Yyo{N zv}zjV>K(Ub94`QzcY6RS#p(rbkEnpAul5g7&jZ5c zl~4PY@%gzP5H5aQ`HfIwet#9oSF=px8TMm827#X>gQ#bKSFN4mDhO%gSk(uYfGcTb za;z^`cyf}tf;8>U8wdwJKmk@f^du*mwvV2$pSYT!ql+_aL8}dqc$0rI2MEDD7$BD} zTB4TqHn4z`fG<3zA2OMLvF8cF+mmBxTMX)?;)Dhyy?kvgPbHsm~^4c|D&e_fCVpjWcxIVsjDWNxCY%aRP7{kEHj;C5t+jl&S=0VQ^}N5!aF({!lFj<~yZBcqQ1`V7`O!V0X5tul%K)3} z4#(A1kM$?iG`R;}zPaIR#W)&~0kSdG%S5_9He*0lXuV6@>o~}Bv;JYC)L}9i>0kN2 z)?G|Br3GGs=#%_Ks5z>13-*D{@7GX9(IS~9mteGCR08F6@B@06As0Sw%B9@ zELtCk-pNFOZ8z(!R|VX)Tjr23EoH`|J+@TXXXs@G%STVNTanMV@`B7>pOyt`!C8jS z-aT#I_&Eki630L#Q?V9IiUGCznv6cpZYZur{jfzP6>iXK2_G_#gN!I1edd zH&@I4QWhIipoPVde_P70N>xOAK&QOH5l?@ND4vn;wLErX=2afT?_b1#zkwiM$+ zpXoiT7SShujMh5h>G6&$Z%iUiv#j_`u!gT7ecvW3DJzdi3!x=+C{7-1mOgR{2k2kh zL9{bs|BI@(42q-c!bQrlD-D_%dH%WZ9Ct zE1E+9mWqPEt4uuO{NEQ;4^Dw>VZ+%=)&-&dl91SpquIzjB0(*<2BX zd&UkAYXqCG^AWjyWJ2!USo%|$pC5|;r&gdpRxa9R0!GgbNhFHlt9|wOWbFtzLva|n zE>s8Tb<-(oJ6>|~>!$vz=wYB+{_NnVq&_Q;;C#_wpj>F~WWdRe`~a_5pc#a!iB?Iz zqQBA?%`ky(7^#VGLG+1f&x=QoymYS=ZE-oTYN$EC{&>o0@k61JC{@F8=**nh{BCB-t+V{Kq=u!{;J46bF(v?S6aAgEJB=najmb-l&j{%V^R z!A)##jtsYD$K)&*Xj$NDPk^jWwY?UORD4nF-Vm?ueWQu2@Xd#(C|I}{86yVa{WgYG zNQWV`X)zJ*$edwgZIxi#a!ZjYE2xyMWCr1v2FML#NQKcf>CGdq4;1W$`_pcRXTYMC z2}=J{`Zo_5Z)fx_EA!dPZx#zmNn-_-cjz<=^xPUSh{8*uh~hVdW0N+obxr1Q?(VCu z;;swex^8-?;fSOjHhQQ-I?51nDQ}A@N~*yudiDj_s+CH0g@ixDpg7;{Ycsx5JNP89 z=dgLCIwJ-{I>He#4Qq_`g2yA=r3&q_M_q+M!{ze=12bi5A}}yB+un{?!*q@Ku~#{i zF7ywxNAepFQy{b6sW>^TsZt3DiAi9fK$S__47KmPRJ6KXp9k`=$1c2dVgGyTsL&60 za&#$4TT^mUA;iI*e1XVqnd09}9*)>LFo(X_GEh2Wszhn%J?+u<#-jz7lQ6jJEHIEI z6WidfP~UrsL8XxSSG9cN@y9)?2RD!Qdu*R)Ba)pE96N+zI%vti$~BHM?Gxf~#(Iwq z9N@-`pvJ~Wmlp=XivM${psOSr2Z7DiX~Pr)OlR6-A!mV*nvL2!-v$>)vsC*$n!aP) zY{eo@S+xYF;8&Vm{l`d+6J4@FrsT$(2^tYj@emHKJ(otg-@0IQEoD3lMsBv$p@R4L zgZAyG@uF5EoF#Lj(N|p8Ps4CsXKgd-Qx3L5QYPmAfU+;qB@(ozwgYzo0s@ z=HNrg5R`4K=0R7sBCjKN41f+c24SWCu2=7(>IBes4E&<#`tHw$+ef0EUp6I`^?6IR z=hXg^C#8cG>)0wHIQ6*h0aKwzpsuT+2;(A}7Y{94fmc5+n7q2|`OBe>AeNRRO6MXK zS_hM%$AB6{M6N;0XeKR)Em9PB^%HkP_-$4gU^o51+Nfj!Z@@i!as4wWH~WfnscF)w zH~6vpm~l6nT?VQLNe|_O-_ZYEfW4Owd9KCnyc}Si`F?wRuF|;sxy1P2_m4oL9KpKo zHnRdrk;t?LW)umxHI5=%U08X_!&X|TsRVLL2NT*Nj{#>&%6E0Jxm9>VipmReuqHWJ zDgt5NK2uAX2q1%9{sNp=o&l*&FH%XSsymrRr+(jJ*WbH0Wu`H};CjOjf$}Q=f?^9( z_5@GuAO)KJ+LB(P+HLr@z(;u-jF&!d=Hm}EP+x#sI-OcczhgAs%sq(OfkUWv9@D6B zA(Q%rBb3WBJhR!jVU?qlVdho}ZUVmb^Iye2>Her5^p^#J7f)=9{;lLzK+_t;+teW5 zmcBe)+Koa>^--5~d4@ek9Fw?Oz2Uy4BFToI$^T=}ZjHR0q7hIWjem#YF#2Q0oWNu> z7|16b1X?6OS3wkQA8BmjpU13i%2RlxU!%l#Or^~o_l|FoaHkrF-=k5ffu;k4V^MnN zZ9nCcVE`7b@xL4JKQrid6_LI~x~&D|G1Rw*CMSp5zy@=WsLt<8nq8zXkm?6M=n2zW z(uUUnILCC*mz3ew4saMR9vDi>S>Sshk}2c90?wXAbV{mFG5EI6IhiwQQ%V*3##@uUp4v)y6|v(|i65vESP>~7$NwR<6uD(nbQy8?yG*eB~iA8L11ES$xa>flz@6#s!o z=SCCS@5xkk3u0aIhen@FWzqP&m=HBsosF9HG>`NZN#fqR>RJgfCLy!wcebAv+^#42 z3w#UCkWJDUa0KK^*ik@~3WB=?zpomRIM|G(Uo#gw2Jhk%bW7MiFD6IjIdNRTwwaks zEPkxL2ih7py85yh9Q@7Sc?f-J!po6JBcJld#BN8i(bOpl2MA4hDTd8gLZ4?u+l~1Z zKqjswB{B*>3Q0$IdRscVLX=xne_3&w6)79V_pHJbIeuST*V-gP-V@${2ac3dE43KYKd z&olmAqUUG9S-p^s+T6{rr^He}BbY(i%jnLNOb()3(GlIZQu{AdIT?9eo)kyE{~sc`Ti&S&|cQTlP(^TrA9}<@lz0= z9quHm#|vpS5&p6=O1GHrov|`ro|W3-icRZs++1-&>AU3Susy${sk{<_eR@QYdYSu` zL(8sgt8WY5_cNP$PRX%MIT_jQy9Ar_k5VRz&pg?f{pCu<o{EhtGWAUv58$XY<_lQod;e3oXHcHA}*$E0&rf zI{sDGKx2w(%3)W;MQq-JU^n0;e+`58v)lM>-dL4DH2Aqx_pr2ELdJ4Dirhi)p!b8n z9`^?C&4|xt&W0~O+E>2p`}5%BUB*S+YKL#qPzF~xfvx%(!5Docpj2DkZXhGf@bgU5 z6;lAH#(u^(Iq(>M7hqPhm8_DiiL%G>FM4Du@@1ZTk`=9{drEh|02KPMQ8 z@$ckps_ux+Pqt@2LPXi4AG3~_7mE}F@cssO>*59(UgGk-r3OHDM$xS$@?1U_)rATT z?`t=n?-72#c1eUsG5iZgVM2pHbR9l~fzShhGfDVFSUjLIz4qK?KK(fnS2hMd7F035 z^kA3j*@fs8&Oi?1v+N$OBCn~GpwUAYIDDD6FS9<98B;WbQd!q90+yz@jXgbg`qExl z9#aZXf1TptB~>g?s(fi`Qhs(oF{=M);~Z@|xS}L|iV(k}kUeRLDFjV~e#|hj4oAFv zHsl(#wZ*hh|2?;aSe->-Tg<7{1Q%>$=#3O!5Y5Wd|4v%t>o<}ub3|oI0zo8HtEix8 z&5{JLG?B7tU^G+>?+nM0N)m?gFSDY2)cfGcjc69iuyktJ(!E7Lc2SLSCfHfTosr#BJ#TWn;=)6mT4?d_%e%U?MK? zcN_0QGw^7HH!B9sz$@ao`%;;&O}du?-BWh3xadt0Ee(U!bVCf6zttVgZfLDE!-kW( zJs)Ay%+F!$l3+GF)Hmb6&z&Zm7dK*-Qw;$Gq^~y4t%+_E_HAlgl1_XQ6x>=|_$UV3qK7@JK`|bI436}d;2g$bGI}Co`Fo1oZtz1Tyejv;+3TB6@s1f9d4e177c4mb^798Yq5zY1Mst@#PS09khPBCGSGi!Tegl9-e1_@Dx z)eNd#q@j z&w4CZ+rXeC4$ri>tE9iXWupB?Ll8Fa35R?=P|wb0EVpL?=q4OW*SXa9miO$?H1VqJ z*?9&WqE0dT$P}h>y=3g`x-dr~oCjQSxnoeb7aYPv6i4x0L0>#8HBq{bMY=iI;@`eWYP#@Ds{)XY>z;_3yHqE4=W28l9@$pK|C;0#xgf&G>*YJ% zZ}n1KKtRj%E)7o#(~}NfJ}TYeyU-<@A_}xA@>BMVp|03zbw;nkcP`O<%h67&MjvI2 zO&fjbY!;h-3puv<>yivB;X7O5F;Y`y>Nq)$6te))Z542>JDM?5z#Zoh7MA=k`G<`) zVy^F!z&v*;FbxJZ{z>E&4V+mU6LAMXZG(3RYwzWJ8?5Pyx?d_>U;Xq^UA+(*T}oiPic zzj~Ub4mnX3Up!E~=*kNh`T)|OUqAYZ*ooC>1v00}o!h#{xOVsFf7>HA_1$&CAw;ZQ zYQ`2mn<{udPs^KuPT*LH^Cef&2&2Bd9u;lMorkkPH(fZ6=1t05VYCVvqVr&VfOb-e z{5 z6pliq4p*-f@q#<(8?jO)I!00k1`aecD&3*W;6>HPF#9ac&h z=RA=#F=*IAgi)lGV2pE!0b*LN0tyDnLm5@tf5+Z`Lk%RJu_?$JxE3|hV`wd7)B(FD zM9~>EI^IV@DLSWs(&b2hRgq=50l~y4uy=6J_SXEXKpDLZG5FAaXR8rGUN4>#{i$GP z+w04NSA}qdCi1n~H&2BibSX_#F95}Q3b=d-P!6gd<;~tRz4m^X5oI5K~%4wPwo&=>x5MP1qbaw-D$Kxg!k9!ZNZ6cn8_}S5L$D$b<%@HkRP>= zw62jERH_d7tbbkshh9t!O`XS3KXL3FZIAnAC>tIM`$3F^l~$05#M=lo`sAW0J@9#u${4-vi?ehG1`0c(*7cG$rV zpsgwiHz43ObY88(*naDhCX}oTS|{v|MlTm16X}-L>i>7xkqTc-&JcV_k27hHgnP)9 zHEE|Zbwd=NFBx&gnvk#hT8#h-yrK)6Gt|jG_Tem-)bAs%eo!}swzI#X3kw^?;K-*w zX8@=^L~Sto@~5eBqX_TK>JQT114BV9qJW` z%OD_r+f?gBA6^mpUN|aSGv4+W5S~4>Z?69pz3Bg0{8Fa;S&mpaA(DZ7*QOwva=5h3 z5dfHt5BU~V5vX;k+7bCv*6$5~N~TgZ9)is?Zt1GCk~?u|<&wSgSF> zKVhu8^$zRP#Zz&t>a;Hnuz0qwu1{t&#Fnvw z`j*M?4n(@^8+jBb0L6x&HgayzB;D1vzqeNG8Qs(&&k(Pp^VbjQIr2s~iJu;{;Dz!7 zmIQD^k^tA(pAG7-G(l|t*9*XP=fC@^C$T3lbOzE=uB*L%RdB>!1f*_3Q$VSJbh06W zurD`BKX^h)#8DKHAN#ZKP$9J}B7X1Izz{D>74>%hNi3H<&?efbH^+;&385NZE(T=$y z77qZ?ooLd?&c{pym67p{*Yw#!5o4kk8Wjw%{{;tJM?oe)XXZN+=wnWPs+|BaVHFZW zQ+r=*6(`p)A{l#IJg;f^4^X7dF;lV<(Bq5q81NS!WggdGnjIa>@{-K2qwC0ix$=^a z$guk@Za``fGCf%^L*OI1H0VbKO_BE5nn-#eK6L%ra1Re|vK~a#BL7bp_nP_|LXat- z$YuHvDwA@m=;(?CII%f+hB44`LJ$1l3xo1eBsO>Q1m*{`8_`HaDM9BUk#Zh)kS#{; zJ%de5?TmCm>@oa+C5u~zx)N_u6*vI8_BU}x!JvN|=*sUKz}5gll)FA103yax3+oAn zt+RP2mni);AW+ifH~@urpN@aHY^_wZJ-& zfTu;m0s-IdoJ|KNg_y+bXi7lxV^TTlFYilM9LODUBa#>D{1O_$r_c{Sr85hT9=?`)G(_Pj-HE0XU1c1X?Gpyf0olv*pEH0xX^! z;X(x0-5m9(f*V=kDj(S5fuIwT$)Cj{MwpPCX{CXbb0AiE$-WUgr(9^w1+p3IEw*XR zmY2X#GeB4zhSbgm&TS!{0}UW%)!)58YhRkCDaPqb5y<$mJ^0cMd5P7gS3}(777dmf zIV8JH5%Lw|O#sTtpQY_~nD#I6fk{Sz}4l9ZtF;zY-qYXVy#_&N@JbnZGMEKvOIgfY@7^_8&2x z6`kj`(VmE@Ke(F!6jc3Uc>M#w0s(GA-08d83hYvU0Lk~aj%Vad<`w4iixVtiMw zy)0JRNHmFf$;-k{wi6gZh)scib ztUt$4r1)5{{AV1c%7QPElA(IM_2c%8!Oj?7XZ-@EH`O`$?ajGSIryon*@}_ z>4Z;r;wxX_$@tfO^J#M?MYT&<6j3FQVD?obmW_3htw!f~NUh&S*PABi`#9K#3+2-l zV{`W~F>C^Z(O&%X44IMoRu_$!>*LVc;A9Z#sZdbN1d+|E^20ik9`GeS&!#X_WQLt- zF`@kucTUc^xCs64$r&Nk!XIr*76bj!hvW|HG&ClBcikmA@x^uDoH5|Hly#+aOARK zQ2sCrrUMiu>|Fi9C+g;=>H|_?{5N{|RueJHfwUQ(;0EE8!Xq! z`0|V9A46|oX=98dwqZ}<-`H73bq{gN_W{%eW@5*=D&5*rDwRZep*34>Jfo=N=6aYZ zy+@3JMCI^z=jI>hIfKFK7^zB;1uBEIte*-P)=#JsQ4!{biur1w9uZx8a* zM_TX=Civb)l-0BNMjr#=+1$qn-v^Bp43go*KHG!3Rz#55sBBRh)Z-Z4L1R%Jd6@kc zwBy%qpGuKxI(d)+2DYkETZfUVTnsH!-zeyb*)NcucaqzsOWR%DHB9NzGFS=1)y42p zqjc?e>UPiLaVxWJ*5@L&=Q2Tp&_sWDnp*#Zqyav*XP17rZWyG=j75Vf;M+FREWA{z zF)62U?xOu#M!xnkQVL(OQ9giBvF+ltg=s&Jhd6t(h^W zoObX_VnA`aKI61J-viGQ1rCORJ5THFNK*gFCfO`ew8si2rP-AV-93aJbE+TRZ5>?N zN@*S;hueR3idT8sIjNAn;Ji=&QzV=l=>UG^+}fVM2jp|r%ub$rLjBG=)L5LS!PiW~X*n#$<2QC0KDd2N%JA(K1bf&I$!S^-ou`7_`RAkl3S>LVKn_artNO4T)!$y4=Rc| zlgk*$^MV=Cor$E(F_Cm4HC6v|td$ghzMZ6fXPxyU(HCg0bKVfO+WO32iSNs=V7#ZJ zjOHGe5s@bw&Vkys<6i`&GGo^f;w^ETynQy_rSD?9f9pqdsNRAN*OS^v<7Y^B@!1%4 zAXToPBwNiQV;N!y-5FRFm-Z6K66^EOB-P#0=6u_!G2Z2W2;Xd0;7A{^M@hOI9QRlz zyp5n->pRjC2{U@U)%C`^U({sLSqb;gqyg$Q7Hs@*i2`bU*MFta8 zQ~LpgQxy|z)5$G+4lq>%({;pD?1m4RYBoopB={TdhsA014=^-hS6i4+e*u`V6G#?l zqxsw{O0v~SuylrnOU}%~<2v9x2t>cQWmL0PBF>r;=c5tx^+z#w1DrXjUJ{y2-73p1 zV>MEMTvrQ4(UXK%6_y+NJMA*_8DK62>h8JgCjqW@;lmL@l(gEtFd{mUOwP{$*;4PK z>ZuX`T2M32yPqA5&x~N!eQ{Io# zLL!SGl8R0dValkZf+2NlWQy+lkLj`$7Kr6x3k9yYne-X49gp6|*m6DjCnIs`KCFR9On>NHdhaLCP$#%(g$53;JijQHw4N8vb2h zMlHX;!Cro8@I(DbKLi=^O^DzHF;v|)tdB(%eL+< zojPw~6!T(>Y$99ybV+djfGG?5c1vs7reWeQM`(YXJ*?WuHei+bOD~+t4nWiUOxTi0 zLqRe=9iW!d(u|S}d6l#mA-HJQzn89(_K`=SE1w4w%U38h>xDDfa9jD^Yg`l9R#L{5 zaI%SB9S_g1R-=@y%!g6GaK{xZoLcVZuO*k@1nj;sVbw?rXb8r_;-I!d4r?KnGegx2q_`G(M#faDh2qMJ#>K&NnqLk%$)R8 zJ9kA{&7DGEj~oq}pMcy!@1N9pn~!a4^Xk7dYvENqyAKnpzi@OqpOvzKhHZfBNBhO! zPMk9UOlZ5@%r9-CLhnF&0bM%3|AMEGHld_x%sJc(VZ>J$o_}gSl2wCNVsJr55btCW zWe=T6*FE`9dbC6ehbmmG2W+nd7(o{V>*;ChH7}?3mQxO$tM>#p537x_md6i%XJQrP zqZ)U)VfanMrL56An?Rqh2}RuGOnU!b^{xQ3+@|4aJ0eX^kTn*+SfVIv)8eM6SmC@i zQ&n>M&UN1J!~Kw_8(gB`xpzYMArZnP%MON3#w@?{)L+)TBikC-1%w!Cq%R z@51=WgNMVWkd~N~KT0-iXbyruAIldl6P=@2pKMe<%jOxRFHaKw(YK6_TfpqrSYKg9 z^0wfES&}UtDJ>ad;T(HCuz9SH9EAC^Fv(I;u7`1%hgQXQtgCZeG2QbBY_4(Y87pD} z03vL|e9azBnJJCS6C>pfYu_MRk)E}gL;FGK8F0a}H1F$nB8%fh`r59valK@7X!5-? zQOWai&|x(M4&5a-r52FF>lkjq;ME@Eh^K*gE}#52eJFMXB-1IotUF0o5lgT}9cRrc zYY4Km(6UzMbmg~&4?>?uTa|ABhv6tr>Z6`nbCN&!HObl?l|)t(?=DFmeui5YQy*#8 z8XHpZPF@?-9KCCweFYf#(2~%8R1dF*!lc%cu5Yzm_N&J}ML2d!ka0+R^OhY=nlwrd z%Zd-%BOh=O|{L|}5ReWKwGO0x~yA@H+_MO#1^lHM#75oY`1uL~}v zmJaY6 z>(_7c#R*zuC1qC_<-n1@HYpFn>SVoU%4Fqqfw#L87BQP<{c3d}d5H*-I z%hnd|VaaG$4EI~c4=SL6g9bHGKo8dcx-_x*UG$3gl4IGHBo{w@x-M-b3f#0?4m3nu5 z6Mg1}MiKbK_oS>!BoG7@X0|$KfaNPLR2;Wghq}t-j1flDroJ%*%?&U6^Q}rHXR&_< zr{mACjUAM!;V~Dxr&-QlRDiQ^Rk_OQ!h939mYM94O(|DF3_bv;_Z7s^C{aC^m=BX! zL#@c?c?kbfLjw6H94=#A3tOYxN%*i!K-;XVI_Nl;05^|&M1GAR1g#F&hm3Pk%{<;x z`IEE3uNXZv_Idr(Gjx32c(_a>DEj;$dN7=i|24!ilz@SU=@h}Hv8~XH8Y40Ea4q}0 zge=Q{>!D5xfE!5cLhxSx10|$#^L5u%TDR~}Q!vUAalR75=h3SB?+lI4m}7eW6I2Y=a& z+CYqPUy<7P+0@kS3<0ZRo(dV2LH0%D4E-hHk7^{}YFUqRumm#cZ>)61o#)uC(5qrp zK}#VTSV?tg*EHpb0k}j8!|8;1D;T=M@HA0E?4y+4<_Ah(dKk%2tNtuOT!~(6(}KIO z8Nmp)t4og1_8cPA!$dRGmK#TJ~cGTD*) zjcDz7oeGLkq`hWXh3VJUSyxO>Drxexgh*O`C1`vh{84PnBUb6&RHhEMd9b}oe^MiL zDS6F?WsD*jN^tY4$o!F~JtDv+>{^t6hv9u&JqeCa9GbA=VX{XKTNt~YTf`J<{<$w> zHx$8#U?yV=Nh@ZXaHdkvgwD;kOWKL$X-Fl`!%)RF$_bm|_A?c1;BZKa)CWc5N-e8e zGf-Mu90BB$?qsXi|IAL7K0=g_oua>Z^Z>AHi5`aG-9Ef6e27tNp!yy9dv()7-e}{V zAgNN}IX-?5Pd(m!Rm~`h?+KA5vyp&&XuLG_Ro>{Ei1v(*-50%^FmV zNbxi6j$f~qA+xGk>==eU0XU+cCFM4h98IF0kgeU)Gw#!;wPAm&BK zD!#`hAKWIG$5Dzz?=kk3s;gfNKc|otSyxgl|H$>^S)>Rn9zSS#u&e`w^u1bq!>8XO z93IYi`YhMMt4eUen~EMk+&xC`7gZB0{Z@Tnp*}LfIQx3P0C=k{FJzighgV;BurQza zzOfN92!4|U-qs87MpO&rH)66}cA0Y5ej|CGym!5e$#S1if5}Nw%2tl*YX;F&Vy*=J z&`c=DO~r_&Qka$tpc{e1(j5fCJjwET?)rA~yV+ zrYUqp=>rEyd!_t%M0^I_I3frkL@`JqEAusR8Q9$8{pm*&@T%q|@pBNRZDQdDFZ0C1 zGtCOBEEEz#27R3-kYz0eGiw8nZ!BBB?y7j!OA(>&X#1gwXMg(D>WUBZmF(MhhmMF@ z^sIbLDLADk_OZ+k&fbM+ayu@ODNHAV+=s4=&q2CYgO?29nWwB3?vwcIeGoX?hWMDV;RT8B-W#Ks( zBd75P9ij#mkdHG3BbP|IC_Jb8!cBtWFxY)buHm<)jB?Tz7%0UY_~jVAP6dPNc&$rX zDUIVRJ$e%CMHB0qxEabU$odMd6xfO%d(31^s!WxTT$iNdkMsNMIidoSK*k?L&5A3A zytc*mB#Uvo{zeSWFaz-R87a(%%TF(Oy*sx5p~zc}AUX-wa(QeDIXIken9@zAGV1+D zSSmiqVo&UHMYEY}(+O=7b&|=7_q_9SnwqRBV0H0G21zto1A*Sm2v;EPtBHiCWd(PHI=|gTXp=78;v<`VVbs z?DIks`PDO;&bKJiKpw^Y%Gcuodt8yxveW3NbG@EB`_s05eo2ht*r}EGk!y662RZPq zUiSkZ;o{*Q^BoZ|E;P9EcnE`1#*|K~bl$a6?h>ZZf5xp$wJ4<(SGDVTgzKE|1wW-m ziZ|k0B@J+vdu-#-Hqg*TTun!t%i+#?a5bc(PM^J4d(Hdq6oEn9hKP{>af%=OVBE z)LPe?;qh_ai_*Qo<8bvbWZgg-3qcQRz8J=nzs6|!u=M@_)i;usRjmfD?i`~_qR^_* zfGV^)kkVEIZ~~*#^~8Z{P$cqzdx^~A2{t7abcp63hcCt9JjsjSOfl}g$*%}VR2s=2 z3aLg?hVz^C2(Xkw7Q;85}Jqg zjwD0l)0TK*jI*u$qc*We-1}I;Rs?JRn<;B>$?#yqLy3nT%Yk+7+m0{ z5Z%)$<@4`ETJ}Y0O!Ydv<#t>@H4bbd5->81*lwbmdS*%#D15^uClRzO8t0NF=ZZ{~ zG`4mahQ$EwssxEMgob~M;Es%s?L-23HFt)NvU|;Og^8J^c0bxV6wPt+v8OoIsfz~d3W7{f#QP+EsXI(8v2P($b>A*Te zL%866C2y#|A570yF>-q@J~gj)S}Y|fEfSJqW7BrpyUKer|Cq>xbhMQ-KUzbcg{u^5 zDEN`FBXs!;-6oI1dz0dKRM|%`E^JA!iRfxRJHLe|{&h$0u%#Za623GmAx*{{biM~Z z)6uUd%UYuF>(tXg_bXY%zPPd_*mMIJ8tOVl8>H{4XFpBuLEQgV8Ra~$Vh;G?SS_a7 zh^wO}bdvn$if;cgm9f>iqC-A|g-Ur>HwHgEsV)kqgatjs<3J4N(SH9zT*x(%*wD?oj;9a>G;X@KUmbgtK*jbfNH=V~lbA zt-!{p)~mIZ?)L~5k%*b9TCvjwn4>iJ^bA=H9Yk`*$PE+45nyo|Jttxc8{8*w$S8F;0CfNj+(Z>c=-ZzuF4>g3C;QiS{wy_2GJJ1u+) z9h1FCbgdXnwzB{1brCG$kOXaM%M$Uw0<@+jCW0#{j^*C!10UV zyh5!(mF_7a{`nsQG|%PA>rWY&)+H-S^0j|&G_aqRgY2{(Z!NjS}oCbc1_<2NNwe|50CR0^ z+hS*7Kh_&Sd}m(>B+D7{#5z=A*6-Hm%kl4peY}j7e7xLnq7S^(HJD#WN4^V2>UOKC z#6oqjgL(;zwPI0Qp*;Q>(rgN7<#Vvwqg;LS};V>h`Ol5X#21*?5?mt_BgWzi>P-VGJ+;~ z>$}_Qp`iRxoeVQfC93(D`>yIAlgR|$x)%fX{ z7&ru;d-D|5i0YT|tW_73JgBk?o@U4f`} zX}k{XQL+7-CW>CY;PX*qRSoMs z6c^6EpGX~jfYn_vJwqpYzn+FKr&V|uQ?1~h9kXQyJeMG)=Di`cYs>?FIXB;`++M_nsSmvuVX9>DE$VJ^l+jLpsKTGb(Er;u->a3L>pMpk01h=_f&AGY zN)C*?AyNv!P zS%}TboMOdTgj6TJlrtU{Oo*%=CO1`B$fs++!RvnwUYn*#)Ru4mz68wIWJJTrFIlE- zmFGr~AqAU2A;2UiQ6aztQ(z_|r6I9Mh`_Tz!%NhYkp)X&V8oNriR#SsvM{yF?e?o*yhJNa(oz;b)F>5Ef4UDqa$_13=7#m9{w&q&L(lnq!TU^PjXf59OG1PdNLv39L@y}Z3e1Tm=nRm6Hj_;Yzcb|ND>_l;_ zKk@1}l{uPC2~tAixXQObKW%yg#EJis%b6Wrq6+!jo8-IxsX}68&Iy1x=2WfFOQBz*D)ee(FHCW-3{$Tajq{rRA(4s$r53N%rqJ?xz}$D!aFnn!leqrMqn~ z_lYwAtW~ZelDeB7>ENN|yNlxoo(7><_g=SsPVCpo) zWvSKT5zB{#plOf=x;GY4Ht#73of%*sA68Ps>Q0xUJD2yqA7!hn=mG!$OSD?4f1>C0 zBtO$Wrn9(3@%rH-1%{991A9vkkUeJsaMrs0kusNjC9)97H&zx~N%inMO|yu1D+6Xf zQ2+qI%NmeB-cs5>VC|*4^hT7qDle6R6f`AkgfyW<;40pb8UVz-PLwkBkgakU4?yo&G_*|khPRvn`kI%@CH01F>#^#1u@zh{7I{DIU;db{uF+|i4_hMUTnhm3=vYNq@}m>99z#f?uSE??trJ1(dU!OmHzunH$VKHh8Y^1It#?&U4g?LBC)dlc+aZ(* z`&EYjhp4j*iX+_Cb&v#?;0^-}GPt`0x8UyX!68_1cXxLS?(Xgmfe%Gi z^f}0!8cv~N4iYh7EF1&DpyU|LUqTD%OkHi%c(OeVe^;ybAfrRBy|b3P5SDFjR+L|W zf72lGgpGv{>dZD=zUBeUWR0C`uOsD#>iO03#Bcgk{T*-~vaLHcx_g6=Q_Ma=nJ9I3TM;qR=hLoVuF?D9upX9UgMNnQXKg_&r%xL`#iIjp8T8Uo+ zk$ESYrCVu$O5(MJy}yr9agI0X@aw-X7(CD1dN#kRJoVf*9$OpHo>XPA4jYN%ZDy%3 zmOX_mNX5EZj+Ji!>8+M8%|Z9M0?i#8n}2-1g1kCDn=ic|^}tQ-^6Vs^vG}o0EpQ}F zG8cO6pc*%Bf#7y$Md6z>eMnyvXGvfQ7S~ z+Z7Q-^;f+-H|a>@ecJjDWxHz&w3%`(i*-$CaIG`6!L73O?*WJwbZC~aQwhH-2|x_t zKQ%&=wQHdpLB7P=&*BuQAEic9vLaGW zdC60PiT(FYJc&QX*yDH5d-USWT?)ytehn#`M8(S9n41f_TA-2^%0@g%62gN!J?bvt)5^cey1h=CV^FY_czwQ44`V^Hq114V zwjz&fBViNd3(^SLJ5ZQETD16QiNA2BnpSNSeGb#s2v@1~Qb{oKC8+lzeO>FJmndP6 zho>`xczzgzyB=Y(XO_3##LfkH3xkB92$7>OW065mq1|8m5wuh`1L}LyAX^mm5&*wzYQnSen%v(KA~vrq>>gU zkB#Fos=%#T zvI!)ywf{1bAK3W9W+!Kom27WZxlPLS7MNU+@&rYn9{g< z!07K3%Qk|7RMVfERR}`Hu&QZnJ;w{-nrv?C$>D{R(U6qHzha#k5e$nSG|JJA6#Z39 z+`?^m4Hy_zAY&mz7o>x%gXVO^Lg8TQ$ccV!Xj5Gzt)R*v3y#J**cfK$6UHM7`o$0x zXyDd|qa>q_lS99T^Y_UV+b?GNRz5>w0Vm^WilNNw`?8NF)}Rq@I%-8p^Pj4N9j1RG zdV)#F(%C*o#F&GWj6T6jTj$Ev@;Y*u7j%-{iMg?1sh=^;7nx2VNfw5~j7T+B+eQ3Q zvXG?vMAe?SMk+eviLIK0ze*2eWdeBnXwEW3Ppcz>2O3|VzY%kf38oU`ln5QO4T{0@ zo4|Nzxqn8g9piWj$Pl_>|3r_dJ4#`}k4Ow5f^G9ZjP;SwBIlV)D}at93ejSKRD z;A?O`oZ_g2)m})C?agb;KH5stVZ@x=B|%PT1s#&0Di06Ud$7>No&w_>%%@JvJsrFN z*P*_h)Z8~)*9Cvg2RA}!9d^EX+W=&tmz#xGfb}9lff=(fG2*#fp>^>)@84bWg@bF0 z8ydyKr2MZv8tboYJvoN@hHx^nRH%%pndOe*>A1d?r-)0%d-G2_xIKp@rrHQmQ`gWa zE0IE5u1p-14Wh!}j_&u{9Hp087_9$F)*1S>uc*hp(};(RCc=~DDqEwlr%7KgNeM0Q z7?fb$jawrbKZf_mtk2v&IYOD3xTwEo4A*(8J$s<))_856sJ}3ixNI0s+v7aic;k%y z!I>toSXYA%CH6n}LF=Hc9C*7o;{D{rW5_e>ZkrU`fx^z$H_M;8fQ>SN(eGc9nbZ8& zazfmDrER9z!!j`s%`$|2wqWE%vUMM$#$}w)@ z17zX0>{;S!?_lu&@L9A$*y$3cLaOO~=tb6}ERiHC5~O^WXG` z`I*`NDC|!>^Kbj4P@e!3x-OTLtz zcleho>IFzFIv;LunRmICw_86W=5cYek3cyK&~=tLR=L@i`&@JRE?`^KyLVWxoaOwn z*nXkp(-|b|&k9V%Qq4|z z_DdC>dv8^37ZC_|SS81p<{v=prG+iY7~F^mH@^J5h}18m2y4#ZSRJNhvooGeHTRIy zV_*LDS#Y=qpyLk-79+WMBu4x6j|QO&Ur>+dFVN2?&5oPx4NaZ1v6QgmN~ zV_sQ~2k1IJHB1_PWRT3jA?|xSM z#G>Q!4bUvA#IDGDw`yX-etEuOOu5!Hm6B6Wr_$IQtbkppIM^7r_$6JgplUN@_SIN~ zo{dl-Fg|<)V1i6^m=t;KNs>uGMC}*v^G1opuEj}NJ35<*x+N?XJc}}Vt z=+nr|z9t#2oD)c{wjg-JPeyOw*%7ZgJyc=Js`{lQN=+qI{;=V&H>&sJfEEdP7j>Kh zHZB?VW!*d2CHU(|hcc|)KD_>}S$Eq*eC_REs?U*P)=O}xgWiUR00G-)wnWO90j_{Er14KDrGDEc@aP~BTs=idNo?q|Hq}UZ z0Fa%NZmQ}d=-aApc1fnnT%^lN;6T{WYrkuYX1tvNQXNTJ3TJ*$CtOETcOqyKe@)6?tn8r5MIGW94Sj5BIvM3HLKWJp0M-FM!uYT5_X8 z?4`qcJ@~c?fBItf@9!=HhonQBx}x2p*-0vbppUbM6_n{U0g~pwB^Q#y*6BcRg2W zE3N)XUt|{(hLmh}_lM_|UNR-=Rc8_Xnt}LtiKA?yq4M?@=ImpDiDWM8X4md*O?@Qo zdA1NDR(EZ(kwLBC>ytoUZ8)*hHT@g>vSi zyg7T3&0rOy1i}Y?*Q7#%T22-S5tYZR`8kP;_SM)1A`0W|s44Jp?l3@0=Z?tW1=FGE zJARmitWf{LxuTMdl2Uf6Q4pR%FS_%yDC>T@&#fh;-pAnRAM#IH-j zQz$q~BNr4k3B?%jRO0tBIGXAligqsYQzF{vWF`_4;4Vf1*|}q|v@p}uBJVT~E^FuT zK?IeMx!AYgL#B=i2yB^8^H8|qXpL+^Fn4#}cV9j-Ku_*OxiUb{l?6!S80mAev^b-M z7D@HhO=9)cgP*my=Mxw=~&xv(GQe{%J!aL0piNnB@97?G%69tWd zDOmi6@Myv8wU%Kxl9hAuB4dB%F$3DN8Ds}8 z`1`Fg4;gMMq6)B<=rQ4_&IiN2M(aEA_)pjTdX(Z{79Zn`D6?9L9{4f-BCqgKr9zfX zcr%DG%OiT@L5X}IxR6jCO-y!)5z25+FJ3QU34) zDeEX09^&aVM?D#RRfb9dR&BrR;j-$XlDkM8AQrmCYDC0#QTb(jBmY#+E6cGnV&qy+Ktf^%({ zx0%qi7W(O0;Nz*}swEmyI0{|xEuj)yQ2&NlcKy3c>D?AF+~3@DK*;k+>E!|a0#Q=j z?BxxzO;YqN*kp8PVeU)zYT_dhkfft~)Q8A5111I+lziB?|K9(*ewC5WZvIi?*hF!Q z?{%rkY-ykb-l30&d?O>>D>~Ma;CHC^0S||}i(#_=_nkagR2b;Z7m^Q&s<>z*P=q`N zZsTrT2C&M1xB^OgzkJJnT+|r)JG*#2)mScf@-~_O-&Okp433-scq&4HV`+4jP#}JX zfx#u<~hF9rLY zwvY()48O@voV~|)y^_CuekT* z{PX8rUK>ZAM<0_bXqPEr8fhE}0!1C@oFjN>YLzjM)FJ5|iiAY<;D!yoEk913l4o#% zk^>r5Lq}Z>MDkaK8hh(7*TGUfGWv%Px-QUgy3(6R>&Ry_k2@^dh^=CV6aoN6&VZwqmR z!HEvScMS*eF9RHM1HF3Vk=HSPMZAAY0bW9c3xf~k+OEMvhOz#jPC75GyO1)|7)GGT z#QUkmm+z&TiubcWFO=&EN-i=2i!3qC>yzPfPTNFdyZ7jyYp3ejl81iW?pm4Fk0qQ^ zuEjfPKevliqJvOWcQ9)1tES6_S$cG3)|fi}2rF-!VrA%R5T%3i?6zsGiWY9q)AC?# zPE~!7bq!d=6vKxul~dcXJXOCFh_AkYuH-VIUpX#GPOIm)?l8Dd;l0m*c#rogk6x%O z8SE9ZgZDz_Pt(Md1--t@-iyJpp0pT6D7Hb0oDmKN*+Jo&J5k> z0J3+b<0qK<8dmK>3x-;6VHfSW7wxTYAo_p<=BuEM=vWzL3!Vh{MVn^tEN(5|M>?u5 z{}URjex&!IDCs)sMf5zN&Q@&iAKbAUF^40mE1?Hg|($oAVo5O@HD zaera#;E<6zb_t2X2$Q`p`ESj9-v6$1ORd$mYj4hRT3X*1VhkeVPD36z&Z5Rcx{M8P z{`@=tW9W#&Yi+=Yyb08ZwC@9xDo~`Zt=BI@wJ`XqIakKLA6p0=uo{>l*U5ofaKZVT zE%yN#UuJZW+{0M@NQs4G>#1vN;_$Ec%->o3B1i}t=F^$?ey^|k;UxVC%T)v&JmAvI znR3uSiVlU~_g9ar)p#zwDJDQe$bQ0<&EEW50KKZf;lXmiu<4fvV1Kxv_SvD15x83w zEBg2YU@asoJ=FDUtXnEdcj}tmJORlod+%LBgA?tQ@eUHy1z{@Oe--ZU2}Q98L_lg1 z3|_6>L+EzgSJ5AS@(x2qPl)TjyNZB`Koq)UDSg9>v_VM)6vGNLm?z4o20nH2bOD78 zB%LS3I_LYPvf;8(tnu$B0niAB2exrJ+c>ZNTO-6XV7UlyD%y0AG=_W-!3`#G&Nin< zyg(zrC@wMmybBM6UQ3mmF~g7F}sYt#y|D`Mn5>;G8KtQTGq$fJ@dUOjaD zGM4)+o$8U5-)C|N^qjR?u&)Om2bc_YU-R``zlmqk)Y+!8R`g9gtVT(7A?OREI46D6 z_+T1%*q9|VXZO1`3umvsIFB-6*QDgxwpZ(7$jnL3oJ6FlV;o1)sQTjNvAgfyDcUR=ibqTh@z zbg2O%tdwjzjLWf|cM%$$J8v};d%|)PQl^B>nCkI$u$C&8DGqG;VA@pMG0rPM-AzTQ z)knZD93gU!irxpLU!zY6Ni7Vxm*Xay&uzI+Jed*_O;V!;V!_E5`F-AC#@sT>sd@2` z?+$F^)Jr%$V%L*wGn5VE)szZ$k8gtB*30JB7jAHfLu`m1HV-*;GtIoD;Og;&omHMg zvlQkZudSzXEQ*1~?PKYG#_Pgo2+{H&I^^)7U?M^J?RKCi36%rLV$4<672t%aT*-g# z`}%F=s{oM0yOa>6m}*2&ucR=tC`U+ZqA!|3=_>T_Ir4&Y35+&6v&mIVVy4C*u$H`?Dq*^uS$ zKxn1Pn(vDZf~ima1SUaH71u3&?Y{eWhv~#c@yvp0J0fQ9&>^u|w!^ty)$VGuwK5!7 zJmYK=ep7K>>Kylim*jwC75CeNM8?bEpub}Jg{?dACuU}$Boa|aY0lP+zpm;CE#A-_ z^TRrokyZ>b@QI+AvhUL{Bp2^+`C3LG1<^1vJFbwu!x0D4sTVU1l4@M@Y*|G6If}_8 zIh>PqgrXwHuI9C$u*G!ihaZDlq!>frWcOq3QFxVF^H7*$MuD#|f9!m=aecOrHF|id zl+Iai+NWUFcVbtsrW7MJLUAsL|F}z5Zmw4Kx*KB^?x48M6$`9p52Hh!>F&fn4LFVv z|K=3{){`Dm!yE+y;+uO1i0uQ!=W)1dA(0xEKkZumWu?L#C(%cc9VrK@Bajvd6$Wmv z98&fkC2&N?=Exg|`Ago|@E2d&G|TpC9D?cIKpgGBifNp#dN}=7eO(7y65`L)mCdU> zV52L^3^fCZE#`;$uPk5oO#d&%ISHs5Hg{pr<&r$U`%b@M{hCfrn2%`+ear6;$=&Jz0=_? zrP>h+j(A?We5<+b4s1WA=2k}W$c#?Uj#5k7hfTs7sR@>j^mmQOaV>_X#gD@mdBR9- z9ZUu`2cru!x(WGjL2yon%H{lGyK&o3U~7qeNM4YLpcu@{dNH@s^yA_yvUKD(=wOM5 zIK_}-c7l?ta?C`7akFfU32Uv{PJZiVGj79hFecorX7uOP!K#>vr=XssECHF&wjAuV>`>w0rn!yBn z$Oz_)G|G&TRJj~I674yOm(3Fh#H*Un6G$QW(%bV(!n3qEum|ly4dF;kU}a0n-O0rdy> zX1*p;V{Vd{G_ED`nncm8&$3l~b-j6V5k&(}#>hG>TxN zdSbkAA^LLW-|0+R;ke_yXrNN~;+&`7ijL6j?N@A=RHXg%ZkN7$z_l=h@;HMK4~)&hXsR>3=ArAN^%Nr3m6?4| z_W#&4h>EG;>CGMkOYR_I1s^N3=R7#O8rbp_J!fSXKM1Ccu@$b!N8WN%j>n;F1Ialj z@l_k3Nen4+293$qvxbT!)8zXehl3LEx}?qIK0Q;jeOCxAQ<;vir6O#&x9XnEXm5;7 zu(D$hIVnhKiD}lP2wrDkv_qn#Rjoq9?5Kb@MJ{z%YB!dTf&b;;u4-(z{A!6)UpKf{ z;$)XB7met7Y!#J*?8r`mSKncea`u`P!}fb$Xpcf}KfF}rk2Na6Md-Q&RE#u@a&=561`8sfDLG10mBNG+ z^cZo1QA?0tX=?z>h#XCNYtZV&NJbI5vgeI?VApW2o@?C%IM{w6J)%;TVdC`?xw~MX!t1`~W?32^)Qtw^sJe2jy z5I7NL)-!hhngn|<1fo9J0|QFZJ}9n6j5Lf3WJP`wnA|X7)+wdHxzK@fSR`KVV~8(rFPT?~f}> z@_J^;dbm|%C{Uu7qZkEj@y9aa z_RZZ{()n!vOnqtEC(U=jHC6SQO=1erlj~vmji(o*IE9T$!D#m7^{AzJB^9j^I#=$6 zJ3rx5Lh>*pB}SsaIH4Y+vzEEmf`-{n;6!ndnwD3gnSZi=ur>}1o63>uw&(`BJbPjk zzY%(tJeMt;i+z_cd!o@qywaZDd!FEIS(b|P7#7ZmlO=htw$IYq%t~;SZrSZ0;As3C zFsS0G(_Jf-1QObn^Mjrjm&;CM;CWaf?2jYLlg7M~pnT-54RBlal zL6Z1-uw@aP(cu|h&tq|o$T^fv=j&$BA4U3##ffFSj6uxg+H_I&E;5dVwL}XkfBPC7>2sgtxiiY z#T&4rqUHJAxm4_^mc#y`xDJ^7h*13h7RCWyvkg!BbgH6cG5s{TYEf8;U0G>gT3Rj){bkd9h>)N)Ls2e=-H~ZO63UgP_??+ua5AW zV7(tU}+t;k*Gm?;q^Q%oC^Xs|xOOFI~g=w6~&GkE_38GYJ6FOAO+vL^%c zv%q@XY8B4aq(Ttk8u8}4ulwoxKdd{J@6nX0xiJ_{k0@O26#EZ^m8Z?#VEoF|e}~CR zZ4y>HZigg;MIg?W-MC*Y+G=Km4)1K7Qc$RJFp5_I&7B(NIB&Ha=$Z8~AmZI_ex!?S zI=O><%w=eQ zmNh$NOF2CZbdUzNAxb`*BSN1xE1wCJBm(|dn2V#2R#m2(6$8*=D$Q9Ch`a|z$Kuu0 z9If=Y^OzXn3A=@@X(Khx2=X>HejyEYqO?NV6UzCGFyS0Mo9AK^tf0=sP=Saa9(6D# z^z&AE*LZU8N6BBZTpJXO)gM%$`N)6qY_Z;kpqg??2XY4|)*|eI?vctJd zkKq))?&c_L!~#wo3-Wwt0Dx)(m_<1S@#H<^_>!eZ225fFQha>rEv9rHI>)mjlB^h9 z7zvj|gojJYYSW^|+CQd6pJfpP%m8% z$Ntx$H!R)3%-#+*3iKNW>(-{{&ZfcVO40m}{ntN$ZBnyRi!+90Trte7N(Hgc#X18k zfnXqoP1~tRedrON@8SK4XZq@hi%{<>hxe=Qnb(#Pi$xWVgs4-`efOE}u(HUR3un9w z$3x=oQR+ZxNgGh*xp`FT`&jzDSnT5$C5~iWgBE&B`VrMTa)(%L!Gm8hjR#k=yJsc9O6 z4@y9FQSN!HHrnBmJgk$(8{O@=|MGX0f8y=fS(QKwldt(7xb)|oXL%QSvv)I#3nlXT zKj{%xGMSJCW9Z0FW(4!I@WtKR!#wlwu zEBG1mm@$=h}$2DvFQi7Km!%F&A^0hgoO*!TnND%zw>#? z;3e5VWYk#gC*WLs@-SU%RwGC^i?s3p_FdLCU&1&}%T}xJzvQT05si-DbjWayrO(oC z)2ywV+z$$Lpd!&zzURCH&3E}r?tk^M=MvF_y;>hDYi0!6FyhK{8W6D z=SIKlxK`c5Vb|GaP4l*m`UMx7ZaAS-8$)Xq&i0eb!x?vZ<_{}r3T>-+6T)H9GMyu!oR_>t@z1=NZ3}A;e3_OLk5&i>~9O(s%Dp3oag4v_9lqw&l^1%SR${7NT>^wOqGBE0`={x@OX?wYhykJAy$0qaTBG&IlIz&-1Q zF0GQI4XFO#E#0@@|G_O1``m1YmVnO%!Mizif+;-z?svbrb0~%}4mMIAqdVNN*y@MD zE6RHV?i4msb5w###UJZyCT!Mjd#!Krk74`yY68~vmziRp$K;k>PSHi*=jRaGb57o`g({0qaqX$CScB^@2bU82F>B#ZN@#n{~f-dvE*d z&F6%X^SD`3Af72c&!WA=Z3)CArl=6XpAmbKu;?_AKKGGEF-YR+(W?j&T!s_UJRpWY zIoJ)gMK3fA+ue-~fLl1Ky8%AYb&uRo*-6ZtH*`9la}-((qqC_dR>{dU_)jGNz;>;+ z;Cy%jF8NRAH7Dm=fKlqm=3V!Rs=sQH#R7%tfIHxa^e;O5U*!H!923kR}c)zB@Cs zv&0&*1yyAu)8qKotECse)=1q^eN5VhO4ba3m%!#RBj~&FmpcyKqfur9ODBJ?kPX9} zgIN3oimQM8)IostUYafcN{+D{lzzyXLJ~h89<_A%7Ve<+b(WGFx~erGM=ngMGfw?H zA{rToJo;~ojI^04^0UtHLEm&AE|cr`IU(X!uq*!%Qk4S7uJ`gxm6ujyZRg9GeaphT z#MOJZ{nI$L=#sI)2Lqj7%8px_;B+@aB|gpmjq3-Y&~GivsQfd*w9e3LvV>9%hq}i9 z62Eiz{2UjtXKZ|D)F04@I>OOJcFJzY91-=37p^ps5E>&m-INa=c%a!UlEJG${u^}g z1JU*)Cm-y5@M`_gEknu+vrGat)^NLuKF>i)WdKGlPwp9c2hb$r2$=wy%`1l1IVlHdgJ^yU} z;xymSM93^Uu`&$T7fhI)w@>la$o3t^pn^Y<#f+zk>wKiXlDP!&@`FFp|(rs@U1N6|!{qjO`J zFy^zWE5dH2QEg>8`0CMW1vbeYP$1#m|6ZqsoszO0;jeX9&S%y?C+KXi<8pwEYmUXM zHoUC8BC?%^!>T2frBtV|TQR&*-dR0P!#jrbNt!-u@*L@>CWyk}3S!i4==efAk!wOQPJASgrOKCN2^)yzB6*p<{8|h+k{umjgWx<|gex_`vZyexI zMS3wbRaL^?TR-Ma%m!hD=k-nn27*g^b zPtG!Fd1RFl1UpU5c0QqwXKLL< zN=c#UxC%3q!86=Z6CNV*abXn2&3XKl#z>%Ir%2z8!YWdpiV5CnKOyR(O3x%Xakk7X zU^Sx!)FAVgXmB{kJUsuqVWA&MNcJm!*Kt38e{3(_Byu0k_3CuP@2+B%&VLVcO#HBQ z2fOAlc7!s!V=2|Dctm(DD$BeJ!zv5>iTgqF*wMVD1TEZKd5^nBuzuiM`ikzD^51Pa z40S>31~l-Ic@nzp0bH0TBAqIM@hJ!)Wln4yW$iAlLrmZV<2AA2F>I!hD)KWK0wiu! zjQ+6!nI+pC$K%TD{Epar`qHfQ1P?^rMGsA<(N8uUN-u?);Y~f@ zCBZ8JjC2Oy(Ef%ZkBQVnH2Ht2Q>RkaM5#sX_#ti;z_r5`#IGjhw5Nm}aWI48NNxDy z&o2&Rwk30*BSQMDCyB3Q*0n6O_3%@}f8I7n5u76|mqWYD4W5a8;%VB6a8fVeiLu&) zsL>}N1gVvT(|#z4O74ZScJ-GV=w7|+!6N)SrkevBHWX9{jiqy1QdM7>EutP)UKz#F zY%bLS_$^R*-fX!KSd0E$n0N~GLLE*`m3mWutk3l{B<%wHuq3HICGz*i)?_5TaTvT5 ziNGI}Y4XRcl+U!jNLx@aeT;wETL{tG{{E$Ls+U3hLrMMkQoLY7d@n*u(FEu$MzuJG zgzQ~NzMN`%iY^1s4g?dfNuZkdN(LS>2^Pv}iu^D$B*K~muY-oq9m4%tI#LY^l_mH? zv@lC)PLR&GspS}cF)i&p9+UwZ16NaW_e+h%l}ydz-*JtFBaDma2smA*nCd)|dm2rJrmeBgFPr8ePtOc^7NPHKKCE#BJ{)pjj(-^v76(^mhtS{3?SzH(R$XyK zW7SKc!4ZL+_J%V+`W;0pn)2E$c|7ONEcYXe5(zI=vnC=e}wTkxRo0Z;pX#$Ipn)&VVpHmN2FJtm z3UabHgQTv;<&tCl@qRzk$c~blw|QKizw#21VbTQ(tIui5u8%pZgr=L@|1#hwub&xKBvM)y^6tP}kYbf2FY6CYyu&bbXF zXSR0FmRYQVKx4p$baI&KL`p1#B)ZUYBRA+&7IO(sQ@C{xooV^jlfWaXY?@05 zNH>U+%awjOR#vToijFaEjKvV-2FZ~zut=(5oNIMV;wPhJJ>b*XJSXihdZ2PtiBIhEr#IY@V`lSAuSd@eNaJ8%#jq=aCAy*_G=kYC z>o9z|-<2~)UIdL$-+$~l+2E~@+&%lw~1HR&qhr2^+Un8|Cvi8~AKuTxQ*HL8Zb zkgMVW$QPtDX)-xStof|muTJeu8jn|hFXg%lls?F5q7KqAgiP47dVjiJ&M`N@{cGD| zx@4n^@)=s0J$H}tcrTJ5uwN~%QDKOvJTrsEM$=5Tc6hw+e8+DL14*=8E_J+-ksK74 zIy6kzmEzyt#>p9}2=1#0rf62NWglqYDPg32nQT&<4(#7+;=nsPFbf%vb zn%GoWd-rilNv0mNi4k{lx$nmQ^hsA(=JCXvGt|juO6U!UAzP`=b59m0fb&x!b*23Z z$%ibE)scXkEfovbHdrZ)(K7CafkGv9DkJNuj0Ih!W0zsf+h5=T9abJ|q6x1qo}9>l zPJc%y)vS`7F{dY|3dO@wDASE&IJM=Y?S?08r~wOE&I^7t0UxTMNe*zIl%dO9`5;s> z5bqE0Z2RQb8PW9>g094;1CJ;Et_vP7O$qWAs@9POwE6GZ8NAIg!a1oP{T@UZw~qd= z9)AEdOdaXhDJkq&IkeJ9-97T%G4`bSUp!p$I&0WaXit<5^`XdOny5xIwD{ux%mK3i zik$<@Yb^Pg+-&uYtol`G#ubP_u0S^kK}If%HR?F|3#DGO$zF0&wjZe7&Fp3tQIxtj zsJYi|jm3C7>bw6xE-YDMy1FMOg{9)aq*02`ai%V&{(Jk@VIOM;w3fZ?np=IQln_U_ zANCuDeB(c2;m?_kg+s%-*1;=oxA97Q!t;p>i+`-T%{z2$Od+3)0!0+tj0I!UY3B_j zdOJBL921$RUI5m>;z&;S2Gk=haL<>s{y}OiV&sHydgC$^5L5VvI*`_FNAs+v$jY@=KaC8?F`M;bFlHIC+L8tNJPc{GkSvZs?r z1^h4`1sZx_QL(|6`$-ntY^TZSX<+9|P+f$Pt5OzNJSGtiSs_6igIjV?E`B~|KWMmm)xUrmam;}?>sd_wSftKK=#p^NCRI2MU;QT zzwi42KLi;3F}PJ>QA@Z87P0PrG%nv@q1k*wF9E*>ujjt&D&z?bE=!dV-dj{ ziTbYq-d{8p88flQp0-du%s4s2v?vtc+db2ScWb><08ICd7_@aQu$5?Qj_4gHWC%tC zW8OSr8n2~|0}moF1}@Vdlj$N%Ht)|VZo((OyZ!j(YL}_~w0_iAt6>fZTzq$6c~p_k zMf#W`!FNk8Ebzy#C@q~>D?1BK6x#DTJ?E#h(4V(g^sP7QOfQ&Et!^A0V;mw1dinr?_v(p)RIM*D|}M@k3%x(wHhKL zdFuPdxz6$Ga54mrx-{oA^n)HZHp?p7Y2=DoCL*fF>(3pCb(@dUFkkS^Qp(NdXe|-p z=Edc_+FMgqsuLWIKF~{J6R;~6TY?avx9&b|y!w~wKYtex7kIRy;+ZUVq?R14opK#9 z=ARo4;$2t_*zE^NfJljgCsP=Qj^t0TQquevD=L{Y^orxYV|wb^oV!AHrZQXQ$WQ7j zHHyk3zr}Hx=3n6JL)$r0raLeoMTnPTB~T9(NJ$rdAxxhAVLXFGNSXCYJYWcpy@611&~({r#d}#wmI`y3p-RKEov^RoU_Bm{ zbaW9jA2or)U7v~kd>nk#c>Qy8)4P5h*g?BsN`1T`4jtyn;6KNK;Dr>W4#eoWjo#0& zqJD%VQw^SzHni0zH2S@?6LG7{}SL z$Y->X4N7V90d~yw^l5sB1Re&WNW3w*>hV_16qNmclZZxL&senCjVwvuO7vT}OZz9k zx;1um5l$4i)kN2#=*zgk<{&TZ{jfC4BJd5KL_!WK)d(F zQ)SO%}a3oOlfbf`m6PTmG+_48Vr;c()}1l5v|9atn?obU~9Ka^0)T3IMI(D z*6o(|IeRhjNA#d<$)T@&!l0a*8HN_iM0~2psLx@~q#5HbcgUZ?o&4l&YmtVK7wWOB zcvWdwj%gg>qNsLx{!1}h_XPgJ=M@9>j=I{Yu*N-+xPU|=d<$fhI_HM|V$%q9R1F4D zIvw1iMWo?NLC40TaFxgyGO+OHGN3A_x_*jVB!A<+90;AF3csibNCo^XE~cm~FY_%8-Bb?x=e( z%6Ql2{xcAoc>KrG+qqWV?H_qYVNKfw6BUM=YV?|=ZOX7|n%stNV&66#N6Z9TRUV{L z2(#vbx@vRSLVR28vpB@+%D1 zcJ>^T)mM(R2}Qt;7R`IBr(gj8qFPc0(veVTY4@j&`|*Y6Mx5HpLV!W;8xk*%>j^k%T4`;XG;;%)dxzC3Ua;HxoEo^z_5yZJsGeFO17O>-Wr$_`w@h9xVuIA zR7aW!0>e!3%kwx=l)n9~6@UP#z$)*aC}r#A_QzQFAXh%}jF5EC`aRNne^PfK>zIbt z#dZ2*6E?PCXs(fk`qUCN}#sU%i4_39l#BH(T(cV zPF|NW%Do7p)`WQu4`tm2Or+b*!dgLYtP6CWrq?)8&$mFwg^olUn`Zi(QfL3|L?<+S z+YWUUHVIL!JkYzPW0N6q$ApEJJJjH=J~f>H+A(>q{D0VctDrdIHp~}iaCe8nb#S-A zA-EF=HZVwVhXi+bcXxLQ?(PtRL(l{X36K!ZWKZpGeY@xO+jmYNX^;k4NB<_+E2`6Eg8p}GWs0V2< zJv!F&zmBj8ju!i%+%K5Cby(sc|9e0s@eB$J#-=fE9=^5XJFS+XpSL8j$|2F&N#N~!RyPW&D@32JXmFhgfB(1I=koIQ^Th2 zT5A(1G0WC9qcfWfqR!Y%IE)N+zN(|qK54M^V08A+^~W`)L7+MKP^!0K>Jd>B^!x%v8 z|8nb0j*t1o6)&e5|8z}A%(FOrMu`*8%p+$+f~8U>O8tdpSTxftkVL&ND=qIr9gb!Z z?|1XN38lH3ACxO8B9`zs6{3mV0X`$g!@ADuWMg?|o(af4+qQ`jwlBw|!{o4*WasZU zD%f2`k}wNnj!X>;|B+xP(&K)OD}?Tx|J0+qTjHOt9Y#b8LW#7Ee|C$^-Gyn^Rr>{J z&3u+Bz6qTD15N!RhDCHK_qyngMlsX$9k^wie7xP^z!gk4tGHnnN_}i~glStp{(y}T zpLZyidpqUZ>$?B#v{2Jwaye8Y=X4r51O~=*kWbX<(wG^wR6>CXIfU&fxIU#!&OqGV z57$)OO%JgEiHj`GTDB(%6V10WxRKY!8i%Kj0LUQ`nHZ3G9h0F-6dp*jU>`ziak*!E&33j6fv(mPklHvNnQ`})Q| zDU8fphc79rhxgp(X*70XpiP9NKCXY9G+4r>%FjYwOWY$#Ck zgwJ6@8Xmc!_98wyVf%%2QqZ%2%^UlJ04I+l^poQ9s@QuBQLz76wx{O@D=r#dU9ASp zx9jzu^UXu{q$+Vky(&_jZ|K{HWTUm(PQ3DC=c_l{TUFr{3}O^Avywc~wpe?E+!e-{ zwtn~ODPS$h7xP{%7gl$D)=>PG#2Ex2!Uh}H`Z2#vG{7oRN)+wai+Vg*(up<@rxP2DkeyI7MCba*39jX89 zVZgbI|Bn-@!z^b`U;9>T|Bs7+eSL#}_RTw$zqa=OKH-0V*nixb|K73xctrmDDE`OG zwE5q6?0*eKus{AgD*oq)|KGvmzk|nrC+Po}Pyf5}_>WiOziY++oKOG%S_D&yn7w>@ zfiU0xfr+jzaew~`d=)dnHvgY20Ply_m%UdiP7^ac2qm>-rF`jhAPm=)xPnvykt=if zQG!T~JG;Y(H8*(ZFX1IK&wf47P$nkAW_7HPZT z8Lrb`g?WvC>PHuVypV*9iP_%j$vS$7Q0Wv2spv)w%7#mi5^v|LhsJC+wcD(=6g+U& z;Y2E{UtXeqEAvKb;LUr-=VEmFFJp6K(pithaW44d51k(wIJ=9cU>Q!T69q>SLsZGI8q$S?Pu71yq zs_!zg(O0sKG(~75tLLfKO9KHoD|X(LrYqC9Ix?Qy?KSG@yeRWA=U?=)^;!m8d8`4O zbmCVD#OZq21k{#klZX?%q}TOng9?CvIgj@Da#rbSa0u&b8v%10S6Jb|)7QUyuf3hk zk@ji5Pp?l>?|d_X0l&UA{_bf+q|4ka0;y->l{0q!ZhYBpe8a&{Ar?VX1D8nMYh^ZD zWF-C3!c8dFGAmcl0TfrQ0m;MmC11__969JXR)v{&UH)Fi{H${x{vd1Gu)I}eoxg!s zqK+euQR;Jh%?FD|;(`dc6w(@4GuOOQw;e;g`D+HRWpso`{CN1ZE|WebR(6gdX-nz~ z%U1O9JS0M$?yFf3?)Q)yI-68l`oy?m3z<6}(Yd*I!mKj|v08Y&DQ3Zy)XI)y zAgjd;fy$BJe5)TX8~M0#40X=KofIO8x`7?V^)Vyd-VN3sPv3G0Am;wI@Xo>mQ{0@< zFG{Q3B^|A3DKq;hi} zfDNCiRC9bok9a<(0N#Mido)&Ag-VE#TKLTeHogmoo`_lXtkxD7k){~JqQ=-igZO|l zNh-#j4ntp?8r3#bNmI+PBL20IoiZ!OzQi=|Gd)x{{B8_D`8q3My zI8&gEaU?Q=Ujb%tJY7vTh@PpUq1{`VO|M7C)%iaXcZ6mfT4#MWi?W(S2~%x zLo+uZkf)Wc2zDcnth85@KoAjltNdlD3ba((41>#kMH+C6E^s$p(?q~Pa{DwR;VN!L zY&~x0RCFW>U9G~&uCE_8nFwc2T5Zs0BBB=Qiw73pu#M7zZwk<%JB}gst zz%iUe(yp`8X0O^R#F643K7Y>lk-3Jqc*Ty&;XG($qOGRP$&szq%Ja?V4GMC0U(sXp z)nDCzArptoE*#fKKgxL)Cf;ze^`L@)oFp6Hu54Jx#kk zG2hQ%U(r!+pEWxfZH?RrfFH%br6RoB)Stu7!;T(Qu*_@sY@+a;mpgSzum7O)8**i- z0b-hlEnLjmHvBF;efGMVqHV0iD}7Z!z41kWIGdDQ*R4n4 z<5$yelqWtWkaO z_okXh`hT-wU4X*Zr@~Lyp9ijC;In-n%D-RsUh;!UdKSam2r42`Mk*6C!hH;2smNSh z<JExUM(2)wW zVW?BHq&V8CjaO@ZF{CD^R-Ue`E?YU@ zL^iP1AmAqW7vx&ymM99-9tgqU;xN&(XuWa)#W#2L(*WD$^DYan+BFAI?;Km$7ryqo z&7CQio9B~+mw611i6LKJxKeLeBhI8$w3g>60_%n!k*qWW{hXdhp`-wS6`|iy# zuAlXL&-EB77?roS+mk9g)FwJL{OB<*7p^1c9{X(tzCvEs zoUsrVs!eu%Ifsm3$;Rr@K+`f}wyRw(SFa^vDiF}KyR-g}9!maCU&U7*o3a(WUAhiJ zP(YhlB^6j-yV`_F`9=Pf2t|-K57^PV&Sx@q=G-S1@}1T`XgVp*#(xLXJSX*j?F^Xd z%=0;n-nD0E$8vOtwB3Xma9mA{G*R;_;730L0#cO$->Ft$5AQ(yn{`dqjF z==pwB0lF(lTdADk|E4ECWewy5WT$&t7?p{8V4v*;?8BVHU^S@}G}|sPm_)eVVAG$1 ze3@OEFUHe{te7td!ODZofkU~he8qG*F&Y!)x!?5p5IJp=X(JtWad-`8 zu!^1*6=m1x*|BG9S?_XrT~1#-Cwzaz77FG-L--x=+#;pR_}Z#IFMNdAspCK2_a?XJ z{VE`mqQ}QFD$FEb{KR9KSb4a6gpy~xI;arnAym39G5zuc<3axG6YrEqcIq4mJm}r! z9T%|A6Wzl6cPWt#!o4keHKS1UsG)z@6n{OxKerRT$qgxc9bSCp0FnEW5*N3+!ZrXO z(X~vb)^Ek{6yF~5dfF`-xC0I_av45c)ua`A@3JVDLoEFJu5c-khu%ORPFQGRfy^Wh? zmZ)1)EcOUm5ePJj_rV&0spRzZFKQNm8Ibe%{ zRq4_@%^!;;F(PPT8G3iV$37l-NfRt9s}>fkfFFB^0 zz89Ep)aHG!s3rq$(CFS8LLxsCT15DHo$?QYTRKnovW-iM%f(*m#ExpZX53Z}s(Y5UhN$><6=Bbaovu18zY@CU%Y5>Dw!(U5 zbW-0F^5EM|t3=P>E0w9bQ#*$DJ2NM7aKl$@V4E(iQr5wny4HEO58dO@c01`(`#c9v z3A;S5>TKeQ5jgr$BHC(sa_Yb9jIl2dyXn@`{d1<=h;r-j>`-h3hm^;bSStC2c}E)* z@O*`}htD8f?%77AFblc+KbVlaY@^i%-VIdQ>Fc+QINx~w>>X5zm_GRQMd59ts|zw% zeh5qq`f1oCQE!0hT*u2H()LSl6BB(dq1TKEvPTRrPNZ3 zwB^Ghvs;Us)8EDj`j#EaSUOxMOR+CdqEZ?GIbsh3cx#M&hvqvltUL1gTCepd9e+Df zY@OiiW6Jb6fB(z5m-6zbL3fE@8;(37`q0%Wp@N{e;Y!6l?4UfOw&Lj3_PdkcwP;L7 z;wAMCO$grCO#}oWI?`wwJpN5`au$6s5|ExQOm4{LXF!aJ-$(z~YZd?NO!uot;fy}N z#nqaIo{RF{TL1bFl(X%>V;24xI6c{fH5FJHBPCR`?)pD#YT~|4>GSN(1l&yb7+XxB z65DmN8dC>k5GU3iv%Y`I=XL#N=T|xd|Aw#v49>a?=xQ4j1W;eUnacF-#SHNc3ldHr zwfM$Sm_es=j*&@GFs7hZ(g2iWt1?svn&{#~RiDSd&RMj}&&F@DG;-hzeS`s%23}vj zW4IHW4LoF>x747rDEvC-{`4n>*nN75WkB{$q{xG12GgjSjYtjei`2F4im+O$wt5|V zruQFS`~1y&8*AYu#F>X}WK&t>-;#Qb?tbY}rQ76CwYa3%M1vrnlFZRA*GD z-i|{H?cpEudPATX6r*==AB^ZQ`#$P6GSUNL68hcc>nA(eODq|_)c`zdc{Pi_FzEs) z*Sf2ohJESSwa&L!JCsR*^1#sQ?E{M{^8BlQJ{@IZW(XFYhP1q<%%7V`sUMAKe}E2& zlXOC5AruM|ozCDvJY+DSFos{hZD@2uT&3vzo|))a45JV3UXEaUvJijHQ${&o_8?cR zO-DIzbROH(QiVJ6sYrde6SSmPqmzCU8`%PKa2P%Gj=9ZE`;E>+!nXkFZ`~HB>3xB%9YI;PQ4uRz^`WWLrE95z&Mq{PBRu;1-n#jRtxs;(^_B!_2G1yVzVMzw~YY z02;I{W5;wl9pHzIs6_C5@u2e52GYkY#$-}ZuBUpZS)!vKpv%Yc9d8Xw9Unf zgcIGG`Fh$y^Z^0Tt&=tLVOFidDeE@#Dd2)AWi9;y&sYeU{owgl=Wn&xzSQq4zPtWq zbNo~+B%0C5@~kk3Nrp9b%lfg3I=5+qN>rURG&;fVB1f^wuPN>>L&}TTcxLT6EKw!W zmtRHxQJA_Lt59lD96@H4AS-FgWJHh~+nCn+QK754O|u&-G6+bL?EGi_b)OcBH|ZW0 z&Wz^4+M@wJ%{A(Pj2>1-?}oA3c%(GZh0%^U^}n^072IFkjxr6j-Y{J;P5So7gWJ?s ztOqlC-ddo$gvGwW?Q`2Uc}NOzqQ}$J%)mY6gyZ3*I^FqECXT^U$I@X@*CD*xf5mt=VaE8sWpPbPu`zn~&DEQ@ zSA;zQ40Zad!Z$Z)_iz5bN6jTkh(eDPH|0-wk1#FzBVhtQ(Kq%Kt#onmD#rn;mWHK9 zSBeHV+n@9h0RYI!Ip)8ojFRHj1;*3uA~cT`NPa-yN^X`hC$Ys9EX`FMUcm%qMxJf0 zR`1`-5D>aJ*TY07Qkcac$W)5~$cN>C=RU79$CJx{+EjBT6)%k}EGf&3yp~j3mz)nv zPOsRKxTgej9{0GK1i@CxxWVslT2hF>hxLeh^f%aK(-!moqu#TuiS&>*JI4F7B?b)b z{doYAK|Q;~A6GbYz5K}Ogf#)MShMYZyY;^1nh)oh*ko@Wyu8kV&hwgZr$wbXH}E2- zq`HMYnF`KN&%||f$-%zdzDhy0$gc+da%q%ppr%!k$xqH;3sINLKTweToyKW7Inf0N zZ#|cov5&wk79m;6AwI`zzUbi?-6^S3mZ)s;uTY#mPfoqx_}sMth^Hbo?`qNdP7AL>m8&bU>SGqBNz_=x730Kd#{hrpR)c z&{&oOP6ki&j(P3X&_QTYJT?=1R-eEIkf1^dZq64wcG${5$m2T&y0G2E>diiublMA0neLv^wJGSnU>>+{9s`j-SS0RZT(eJ=& z*>d$Wp)|rnb)AH=N$!Xl_82tEQj$~_n$_^|hlGh{|3327GJ8(80=O8$U}ZmfJjWuB zpl@{=z)O&?tIB&bIz#%4sh?c0dV4=orT*q9^^Xo zd$d1Pr2=V19{ui}(BZZ{{!*gYzF{D*J55wK&~R~te~?*rR8Xxtfwz4~GBlnBNOglg zrih=OZ`p*wDf$*S1l+Dq+kI{D`d;>8Fpz?$J9@mxVdXXjg3m^aEsK@ZI3A7Dk-48M z-MK5<$&mvu;R1lwi3*Nq!Z=?vJn-_WUtpTPllpZ`%R4#fax81-V@}~WD(k3!no zg@f9HNi7Bzs&R9iK~}^L(td0>g>H`iDSBnL{E8)XhlNFr3l=Kk=hRd}mamh558vkg*#R8_DaTf+M)VUznqjNK7b(qC5xBTI$`IeHemEu`+! zZfLD{*O~`>CKvvl_RR6mT+(XqK*7QXNLf}5xiJHXo7Hfm&`NB6R<}KQOaaPd1FfRCE{HDkqd-B#tuE6p>_uljEYLra>3igER4Fu@y0I*^(!Dp ziqp?m&^QmJ9ETbu!50+4V-q6Wzx>)`wqucm4}z+1daP(Zl99VjJESf^~#<+0Z{7L1lJ}SxwX&a5{DkuDB5Z2JvZXq9?h}?<7%~EWFOFkYM-kBd96RnH}6uVz+eG*KVa?2Q+5;c)FCnw@|=cTpRl%Ovv zmPyB@qvaolK5qF{n{fwc{F|JQf6+W_gCy9sz z6)lz@PNgz0ADRT^Y6Bhw#mJ&V;xYp7a>S+CJ-g+C$_PP~#kbLRDTs4P9inl{v8s=1 zb>K{>W)h-C1kq{P3?$U5?~s{kFDe-{okZ9Fl#YJBW?~v}1lQxgRL9-}hmoWof$l=h z?b~(#3$xD=#q8UepV1G~zxi|0$OZzsTWi9MyYgQz8$bQs55-v>sOdiTU}MR=V6Z`6 z4VKGqcvh;-i7mype^e=U%j>zzcM%ZE2|83&&@fAUZE=^;Its0zABii)LI3y{j)d@Q z@42W>dxCZIZ@k&-_42oIvhH?Ai8}C{Ftj`_3!skYE2#@Jk1ku-Y1mE0>XtMUDzv`R zAi;G!eIu6Zm(nuwj0D?pzpkoXn<5*eq5%CeE{4QbJBH=w?yHD`5RIUO=aL12Yi3H4 z_VLf7Z^(kPGJ?}nJ0ogY^^|}e=zNk#h~Fw^@`jt^(RhBM8QRXsF%6-M7Z_|3*>bUz zWXGF$^^B+R_9WHi(RvbtYzWjwI`>5`In_`7qV=%iB>zoNMK!;QnCbq|1t*2&M`GG3 zteL6!Lgs)Qme;&h|5)OvHHZaV=5gzmVfLu-TR#luy>Wh7VF~)U3x_{%$MnaF@meN6 zJBieauNWsAQ$+No-^yreFVRHsglS+V*bo^t3}5-N;WR}54ACWJWWqU34$7ua6%nv~ zb9#Z8W0{&@L%6N5FPZpeCAG5ZD(d_T^S6U@)P#W7-vdR-hh<|`Pt_*yo8q_ z6~fWZy7*jc!3=+%67I18qdP{&H^H5*9O~2=@b1s&m4(LYNc_ZDProAVSxd%JDFza% z9qDf3MJHAzER#+qHlzazY;)k`ulr(gMoto=kN^XucmdLv-M}?J)$QZckdYAag%^gk z(VMAND4Gav8*#l*F)y6rV;E#GQe=)-6ND-9W{{kdnG#HK=i}~fWtyG8v#8VQMlIoI z{j(dsu-0nHic&4>kz+h=Ft|uGxIJmMO9T?$l+Mau>ai3e&vWfS;)g|$np#w6D+!Hv z+f2cACjP#4ALekm9kP^>SzIEeektu!%&eA`f%)&08 z?PJMu8)oOV#LMyVmBr{{A2oRQ0ivYkIf_I8Dy}w%&pQZrW_TQYO@~muGZb*iYgZ9YuFcl^O+gzNSotatmZPKKSc1TTM zcFtilTwd_Y^{;fDPWIhoc%tku`pY8Ers{{9EXB%%%WBbvS=L*shFvgaOpMiosU9P7 z2(LB9Sz+ftqfX08R~Z7y$qDf1MWvB`C7}6u9fw|MSvKzRh0VGT=Y5SsLq09uuLfeMVf>`w*vgV^sLDUm1(*S{tkMWYQ7LZA3#>QgrC z?p5JrcS{VC%)a9TIe6^{sdbvge71tXJ^Cd##g9_qef-|(yZk1-O94b0t zMdgZKy3*4$=LoehqxCbJ85X7-=Oeg~!rP^lC|BCXMe84Pjfkrb1=UkY22B3%KtD*Z zj>DBfI8@deg_4ZQLfEpY1nfi9ohe3T$bksTrx1UNFy`GUe+Ey!V8@|4Ej8iH!D^*& z9+_%c1SDwX{^I*8ps+{JMO4WBPb$*`r36QSKD8mthi0;-P)rB#=_KTa&$6~JuW~#C z3nEEPTp?G_z#|RBgM;e12ObX{HF>5|+caO+BrG(62c;P3imB}L7}I2nZA3gpNFMVmTFS`?!HP*41<-e!km4G$iq@0eSh}fWGNb}uPji?zWw1$ z4TuI0A6Rn}zt{Q~^MVht%#CFu?5P!X}Q|I;2K#y;^!OrjW`#5f}8=svQ{=;u#)$xJ5PCg_J@ zASaw1==*voyG<*O;H=hYCke}R0lU>NU0!msyFaUlban^X-e6+1zjmv_5swPU5~J1D z4Hph*X98aBSg+mt2Yak-$wxe(VFpAbPA=}Y2{FL>a!beoA7++{L{OHzG!vu?81j>Z*%-64 z1}4s2EZ0@>f_ZJe6~QMfPoS&0Ur;IU!Wsp;&FVf%$tu}$6aw5oX_5HVpj+}l+~FTt zwW*K76@#{#pQ3n&+Rko7O1p?k6_Mnbt~@Y8aYn`6q;33YD!#>)0tg83Rt-PlXn82d z%rNk?t9&@E+%l5EtWRW@h;8+G;^615<*l2J4sm*H!|g}EFmlzyjXfH?wE(2)wrQd< zm_LuMOuW(j%h~*r!uN)SGZFY8K=()lUojEF>k3FY* zIiMVIUX~zdT)?$S9qXy8*CR%aLCnaoqn8@PRbpz96%I))RVvblj_m#0j4>ka)h%rx zbcSA$u^>)o&vj$+jGw4=Qz^l6(Blzj>1mf<#ZlM!=|9h=~RX{1VDUdz_R~oqB<&p**bQ;roM0 zsKc(5H%?Uv?xpiV-QTRC9J~e|uBUA-1>Iln-kYIx)#K$3%7B7L8i1kp&_Rk59)gf% z5;w`Xn@OBI+%|yNhnMplyD{%0P`$Ven4*f3;`P=x>Wcz~zB=yRkyDw_D0sLW?+UaC z7IMN|brB3>AghC%i$WoS#&69hOkSg6<<*Tdj_fU~UE-R2kkz$D7W(7m!C{5dgg+I-43v@O0^mO`=*aOU@EQhpz%ua+)RCa7m|9nyZ7 z<%zZCpKrBf2@VsINZ0#7#p$nKO@q&gsbh!iC*oR2V9=z^$ydrqshSgG{yxB$rDuJ! z%R(2W4n(nt!9&EfLx^h22q4Lf+I)!O!3*zD*;U3#o8?bG!Z&Rq<7M+|lY|Gi4&Dgd zMI45%iH|SD%Z#THcg~(}=W=f1yRBPj>PJ!56n;=-84+o!Y5Iouu|VesOF2`lh_PVo zT9k5-6mHH6(Xo0D=}A~wS+=ly4yRUCQic$di5jkz@j}D8dYRJr1_46fGpwa)Ne3@; zM9|T;YW9Q|i`%0KevVMzWf(49 zFA2d_3h1xOP67!-Q?7aM1oMq8Y&bD1^(gX)2;wt9vFHfWe9P2B8`IYPT9lRoiM}5R z1N1LI%(Gz<&k>exLphoVlsEF@9Vc0O7nMpW_ILi@`y**_I&=;Dv2}V)#=83Z+m+&S zTvs?v7pQ!@U(<wv8skM}kmzx;ASkAjk2a%`%*ePQNhk5CZPUJ&!I6)3U-4;XxdVumx3M2rAF`x* zsvURlnbj<{=TI%>6%8m%9=9j&7yAHnY=?*|d34T-RxQvHli-##Ckb#pOIU_gRB1a{ zh{?-xVeoQp@aOY`-;#Hmw;&=_&#+<@B&OZoP%uI>nMf{?tpX>*nsQcM5fw!_oLhEC zPdZ&IQGFU5FHS9jezN-J&cXDXmM^lba4+|oLGQ#Gn9S^>G?cLDERSoRSuxy&@PlF| zEbtLyqIsIVGK3x3BGU~#?13wq7xTq}wgDR-0G(UB5OjBO$(muDhok(38?cXN`ZasE zlz|y%Mw^WnEe0-D4OLbGRX74H39klbWr}X_7604tYQ)M zldI4i4^8#r?u zSHc#0(z=Np9>I+UD`{?y`cM^GZ6Ocq4czV-HKMzkX_hy*y(fStNsxC^)@h^3Pka~n z(ysA`xQ`tv!Sr;bw%>fj&$+L^-!m51Ej4QYli7?%_LRnB7JgMv45PBsz7Dh1`%Rf$ zM~T1B8=tT(R#Fe(h-8|1t0zkod(N6!Tsp2)3hukDc{y=X%-v4H)bv@%5wKAaSpMud zoiJxCaX9UKrul^#@Eh~W`O}_c@;Rm^Ea++P zwI`#x<730^#2{ZzkD0P3+u&{Y>OnBTnzLebooy=nmKhb$rz=*A>>e{Pf>{KXe{6u7 zQoUt`!G%BJmhNsz@A@ie6MpjoOoSfATWUtl!+D1UR8%DcX_rA4$v}`qsp%UYI-e~s z{Ux_qpGF>)e8Xs2hOSZDJ2K~C98cGbWxAC8TibzZ2qXRwv;Ezef z5~s%f&RA3-#ML?uh|%Sb^n+B zG=}$s9qoahY?9k~aLL?5&E#HgCusc3i~ZT+bo`!Z?Tsu#(fLYc zn;|Uv=;H{bsf4dx^9*4IKIz_K^}RSK6tj6UCaotTBP_lN2$haQIn&_!kY-KwPURK$ z+*DL01~t`8R!$&9pVmYRet>slkWU1wYQ%i0$=b&hc`J|Qj7thf) zQc*^a7~#j}hc72)OhO_eqJJIh)dAhNE)q23@!3YFkF2$GMIKY>x|hh`q>l>}+oHm3 z4y~*EjF9${dCIWHw7Zq}9=W%QfQA#b}5j7}1gVTL9X<%o04cct7 zuUi+Ev4>7%s<|4Syy1XVUhpb@pgnJ>`&YsEr(Z2e9pWlSsunztYfuv@0yaI;q)VB4 z;&>dGdh()-1-n7qn$Av+5>-})fAL!4$(*^Lzsvzx2egUmpy#YN50S9$q1g_Es z?))4*q43%FYa22oIf{wYB7QO!qF&sChpb zzZ9Zj`tBnVAb5lT!Gd+|2&&L6BbhYRCq-AgDo=Oa0BEUMkn?hcrLLi6mH3+!bw6go&P zxjH{A@l;k08weA|%QnTnq%CQ%b1A>1fnv6%tGaGjWAgpo7@v@Y_ZFda5uvnvts{qQ zg8YV-uqx%v68pT;u5ys2Jh;i9YLs%Tn4`n*jm(p?g0Zu8U0adhy1Uy*-&Y2gCO?I1 zT5N+x!Heu9_mqg7Mr+|MQM%R+GDcdoDJ5dh|s|n4@l=D z!BI!~p)Cl!-Z2Z^MGuijgil|P0hv0*8w&>CcXwm|e1G6+4g#~1B_sb%d#T7HT{o6+90Gn5o`bw>Kkpo zG=thPF}oPqj=ZxR}Ue{pp;pDB9I@u zS0^3!YezvMRlnMEwBbpm+2jM`eVr=ZX9OnlIerjR+tOjQV(DUvR5*QJAzG#CzSU&o< zd-+W2cC|}tIL&SkeCd%}a^)+hNNU3p^88di>*DU|+EWB?6yXsHT z4x?$#k+q>iB9SO1X^3*YICO7S0wi(8eAQ^(kVl2gDetmIdGLg1waOS(Q;0d3$oljA z6AwEk=J729P*gj0_q z;)X-i`r98l^#HnW%Qo zyPO8+qo4(H(7G?D3m9q}=&eN-;N4Cy-W zcuNG9vOCQD`Ia09Jby^UZX3DT?!(qyA+55|kUGd{bg@0gl)xwa_x*da6^A3DVJIgP zs!v0;^M_t|)xo$pXG+>28Unfic`)AxDx~7uEoaMPPU<~nF}biJNK}^m@RSktVvlbq z{Q@o-Rv@Jh1pu};XI<-i?fFLumR(oCG5n+Pi7pu2W?`%@(|x8hjN# zN*z{l% zX`zzIU+~E&WZ!FJ%?*v}UC7JHBbI;F3qn;)q?|>-e7lF+Jp2Q|>=#9?;SaA(SWQlm z@JdP2vw&37fd00rX#5-Nd;~Zv!wjnD9DUf0ndLxy5^P5|BBvu{|GfDg13>vMInCEI z1y25=A{1RdkNs~FKf2BloUPw6elNKsnOVA*dmCD@BCkGE4t=^Z?_jd?=bm#iMuzMJ zy59}g(Qw1Q=R`&6K^3wJZeimfgp%=2Qgp00w9c#{5Un8?DK0R2ar{~I%g7F_knn3^ zP$>5NuW0XaLGV-%IbnT9wb2=u$wiJ>hJ9_-97XJ*UOZn7xgVXbXtd{YRYy#|tevh7 zxc?*6y#LYHSx&Q1R$Q{r!g9JHDRzUsVG$|Unah5Ro=eIobtoAraIEZ%{YTFI1577r zi;5ymsVHlUNNJlwV$2M0sxN1Xk{281DvF{atSp4de|Z&ho#dcESzN;v9zSHU<ph%22-QmhPsADAlEQjLY-0@;VFDMncDG3`o@ zK7GP?;4-FV5JYl*<4sVR_&w%l=ckwV*5ANN$?v}*7j*vl^#&;?pr_s4gaS8U0Aw75 z4;gKBbc@5n6WTX(ahH))q!lDP{Nj&i9U6qRqISXDbaJ4;h;wh3q> z2ngz*c5NNl{TR2Bx@_Fg;6~FV0*HST>4vw>U}pca9{tuH_W5HuFi{YQq-v0#tY<>Y z#DdUvyme(_w@$$Vka+`eMk}t!c2)oQV@YwCtnrZ>7X$A0>1&}X_A4xXaL?eEO3uK^ zd)4i?rw;l0cO?asIRHur9K~XDCCTCgm!L9aWjIBQiPCht0T1p%vvp!LT~1eo_`083GLUQzcLlTrNL(P~`h&x% zIh<)t??ggBeBABZTvZ2V=FJu>7CgPfrM0#|SGgkgCEh}aI_;y=fnk6|uo%&GEPk-_ zU*gNp&gNJKkie<>SzYoAufM7FwRnLs7DE@%g`^?dj?;biYO{pDBaJI*h3jSz>7};H zqfP^mr*Yv=t?M|NW#Q_#9!okldFr`lf)M{yVpQtD>6BV3g~pvyL<-Py_L-wK!Q4`4 za=)TlJ7E>Mq<_sLqLhx>fK+&S6_9oWE*qh&UVnG@?Q*$ebr4(ck&#F#SejMtH}=c% z6afhbB`>xQw`^pNci&c+lM+}AEf1BgT`Ljx4~~%N{wK4k`ACGXLA>>Ku_V`TZ&ky+ z{mb`y8|=BzPs{+{yokI3>g$qEswD7${+DMnjJLdEf{O`A@j(MU&K|t`dF}zCj+5T+ zT=(&QZnlGQ3j3}^8g2vZK^zRB7LpKGG&9`KikAGVfM+NO`iO9bU&rZA_OH zoR5-8z<1+9rfMg=fuO zq30uUQ3Y*4^C*?tDpN2=2xqL^m-}Y5A?v6ZIso87@qz;%0Z2u*UjHf{<}C4V1GFmo z9Ysh;h)5#qt{RQCP*TzEzUzkuo$ z2R;aXEd)?V#OEMYLZvrviE#vzzlxaPj}Zifp`bJQ_P2?xyxhu0Tye)}$4gC662xhz z*;|sA=-CbgJ?bhd$jjdsYu#~Fgz&@^K!VIn{=y9b^CGj_lQCVWHPFmlcO&H;K36SP zxoDU4MTu$%6o&`a3>++;3{GE-(0dZBIyZ5^!lI2Iodc`s@P|=Fgud@On|h|nvAH26 zWg$0y^h@)lfT#5YXCgSJG3#0{Pt*NB0L4H$zdUswjqlM=KF=aqoxeiU8WkU8z2vi~ znN4mBgbC{Z{`bFyCL()NXhy9f3rt_I_Oj?W* ziu3!>o`E$37jLDj&`$3exRe1kNbRc{c2<@3aX;5^d^muaUO51nMHzu3t;1B^O~Iv$KV zHvBAUhzc2iXm=6ma>!A13TDUCqwWaFsphgw%P*);pgVa*(*omXM$* zdIaU&W2T0%#5$E(fkuY^Msv7#Ft#8BH9b%GrZr^BC~7=$hj4W;1nnqtxZE6le->#s zTHrz|}~G(nn6gGV}3XRu0JrX` z-lF1U3z$DRk}NGg5G3F>EER5wMr-g9XJPNdq+SABuAqi9KI8C2(j>(DBuO)pOLVDR z3N7T^7#Z@WHSn1F2+UC-W^LHouK0GDWvpfrz;EFJQ8OinKVE?C0<}(Z9zA*_fo{pv zEwHTDTU40jRT#TWT2FX^l}gN(EDku3Vs@X$_Hq;1MiB^3h~K@yCMqwqu&`HCC)-*;d1yZk8*K#;fL6 zZ_8^CG?T&;zckX#)n*nrK7aI1ABKPs#@D#P8S3~(MM_<|m=4EKQ73k=biYOmpd!y< zE2@)79-PZoV zu;YBU890>H7w^{mQ5=;vr4X%nA(S>8$C%>ij8X5<6!ELk)q1C(`)I!xDS0kEgdH9V zqO;a<1;od~Ilte0>~jznQ4XNR;$U+{6T6d-+x#zHSFyxIM_nUaWncZ5CFi$jT25b?q%N7{ol+)+Q4VNsU2XtfO7NLB zlsbg7jTUn*vTQkbT_*A${wShH7pyR^y#0k zaDh1Z#Uh_BO zK@Rc&06+jqL_t(7Iky~}#X7WdT*bW7xc_k}G(t{g(PRoY1x$LRbaZH>lVWJUj!>J7UnU%AccQOlB z0-JYh2}kpo&mqWz+SU6EKb0?k81Hk>)2B~4(=9g1jWaVfg9e<83C+rAm6;JkNV^3f zA(n?XC|7^kARngukddn=Po5xbd4+|%O<}c5ls79r1F>7ZMrO#USY(oLm2fO3@b>g3 zUGtLQM{G`t+U#gv#^g#4O@tm_u-+Tnus8_0l^;WjfD~bRj8nX`o(UV(1fCXf1aP*f zctjg5?dSxOiEqdm?;xY<^fd2OOwiubaPXh=?>fre8no8kYW~vaEoH*^%0=W~D z*vOz;$v8kXp(S9c3&e`WaYL}tEPI)00T(GPvDAu=lNvsOIAd!=5TnQzd{ZVxI5LTY zZ(qnigl14brDlVJuReDO0-4d=@EDUnHOQFx+{C%_%(Rxa0h_4!=_w-`%B>r_xa* z08yoLZF2ZIJX`|EctLVD7^3bJ@fS68hxvjNZ{&fW8Xs0`6=I${FgWRi^V!nbhA&?7 zgwb<}B&oQIsK66|Nm!3LXN)0=I1nUCTGrzVp!=WR*-1uvOQ{=`q`{}6&sDi4;NoZ$ zW{n|6oyR00yWbGQbcw{#s0_jy@`fe@#En5*4zSJD6K3YpB3uqPY20A4E}YyhbsXzf zO;RdywT$UK;zNed4IIuRHn4X-3dE@|16H7)n)`(BWxy@mAN%=#x8j)j<;*1f6jQt_ zlULm1*t5JpA5E-TPCoxRZx=Af6IyJ^stGd7hGQSdbR!R06PkVM>S>lOvlVvauMqM-T69YkqnGk2_I8t%Ox{#_jOK>Loc|Ee zWA`+K8M-GCf&}JRv9m(GAgX2*>TxXDahM2~OfTN*vo+7PwI!QwwG11;CVZz7c ziSt73f{&RwlP(Wv#TE^OFrp}U<&j;DjD`@_Jj3#et6l8WiO0~tq3c3Z% zH&!ko3uGM~I>LH3gQisooCJ<}df%T?sF^kaSsMAK|!}IWy9I^Zis8Erdb9+=g8r9>h;p6yP4}`Eool7%{iJnJ_UJ*el#x=y%c z;lYHU8O`m5)u*hvRD=-xfNU7~G!W&!AlB}Oh ziozw>BYrZ2fCEHN#=qcIEqQ3Nd?763`h)!MTXycipWOhqZPfLr{& zQgORNVMzJ4a}grS+{W#<;Zc$+?{Y3Kero2TN$0%X@%O*~-M$uuzhad}t<;e$)$*h8 z>Mf_3vbu@R%Cy-{A~D%B$}}bwnMBBy4<9yQ<&HUS81>)!o^{KgmQ6kwMYK2b1f5;# z+W26!M=E66l{kvIz3&k+&lS2#B1LZaE=HJ2LCFa}{Vgsnxl)-qV@S);SRcD9k>q3# zlm!pXx>8+1m2(7{o0}VJ-f8ZUA?8vXz%g3j`M-ut5}=$b&SuBqWh%wOuH|YLD8jT{ z2|z0+2q2mvlErTJL76812`hhgj>(*g%E|F(B_K;9c~k7ulEu*a!9*}5D5=&nL0HXM z)ran=l^>}f)}m5}5j^B7CgCn@IJ~m7U?S)W4Z)%t5l@0;-3FJ6_uhLChZwYiw&~pF!{b85Z>kudsL0 zGR9oFr5*yNJEBVi!llQR+6|jnPCB-&NsrwWndk=-^Fl!QKmdlB18YR%(xmQmG>&+M zX*pR)UI;CC*PX5%1Ir5E`q02z%S^7MXl_NUj70?Ba5xEEZ#l^!E7r^^Hgrx(??lG| zb1?!Ng6HKhzDI3nXaSB^qwZw#Y$fxbW7e>gCL=%#iI(U%DZ_)7t|6SUQvUV$Jn3FF zDiPmdFC1Q3#^N<=LgS3}EHn;_&i4yRh<=TJmL-5Cs%2nWt3ny#h8jXr^EVEhY%0s4 zQH!wmqEpAjm7>1U3|05co%I7vm>NQBHz62e!bV10kXE}j?{{V^`SIh&__BnfM{{6I zST^i3=c`cXg-Xaz!YqQdRXxr+nMHEy|98ciiXlH5pE-%i^IVcH@gDMZ{2P`8&_*W@ zRPg0zW?_vK-J|K2Qt7qjn zaoZgE_JjjaQ&$_sRNTxm&mkM4A_R?}?f|@Un1;sN0MGg$x&(kEAaHe0sqrY=aAq$g z!z+cb$BMG5BYmK4B;zZ*>0j`)s@5G*ght9zM_WRHo-tk|^KxGt55mZ>! zO-*T5slEET#5i?aeS8k@x?|Kq_4Mge=a5^Yw_=AM#TP)|f(J*Mny%uDg{vw%nqR8F zQvq=d^zgccFzU<@tRxu1EE|{zZ+C`VR)~2t>W}~;a6KlDgU8{dcbRq-V3A5^Pqy?V zDVF4p8M#FmI&IuGN`-t;=((E1BOi{(6^E0>LcLo{G0rrN*5F(xM&98qG>NqsL`utX z>q_Ac;Uw-#!Dt!2j=32vEiv=f!R!r0AP#b8vm01H;FIAGdHLzj_06fj#w5%E-!*u` zfuf(~D9w+cHS>>64X)>g}M&G0eGECr` zwGkiF+pk%Wb1my!I{5KF|KrW`$3On@C4VOWppv12u(~W@(?I~IKI^F@JQ-l!hY!or zXGELJUD=zhL&4fsQn)vUC}slTEyem#4ZJu>tRid;5au)FaVcAGRWBU%#46>aZPi3n za!r$UV%ZDRrbZhTS~MX5B&w}Rwfe85On5An?;z&C~rno99G*KM3+5i5p z4Ze8z=~2Q*cP03vQ6Whl z4j)Pme2p*x8_hl?!}BCzYSyH^On}2c$Q$)&MF|8`x}X8!MrZ^mu0HJao`DC=z&EAm zLCx&odj`rtC9?Wk)l;LLF>XB040j-Z`Q~567*Yz=T-DhDU`WyQrs4Gi!dsY3FRL9h zJntY>==;e6r*}ebcZlk}cnHJ^f$0E{5Ql^?oo%9a?sjx_sE~KY05OWDNZ6xS9VQ&r zgr%#@L~y3l#gPNQj16&`K*%wv*K+KvN--wSA2&|-36Bm)r@NU#<0o*vKb=jAT0Y3C zcU^%mq0!A2=RI!h(!~sufKl`@)@Q#Mpi{BtU1FH99YpOavqG(61c1m10wygP<6+>d z7COazqs3doh>ci}#H<`L0xZMakohL}tV$+Wnb?7px`{N!j0Vn1kF@C4h6+q=%6It4e@=3#3nfoPnm-7mdIH+;UyHZUN-PuRQD!{K6+`0# z+T9gB z=l1s20igtWcj|RQWoXMV0a@Dw^)@_OsHTI5GW_)E)BG$neEt)Vb4kem{>xa=VlRX; zAZin5D83%76it|DIh2-?ro4Flz%;Y|F_BHJf_V>A^ej%9q)3u8#Lzs^Gr?S#8H#?{ zKA+QSz42DVaT>B^uEi23F)A9m&=3Vn7^g^2mu|XaqST2_sduJGg0krZ5pu{(fGSu% zIB@wRx6Wc;7!Bpa2#1f!hjBwPog+>|xrOY*G#pIxNS-q|zB0fUpbtHI;HG=x1qHJJ zO#v0RBOKqpL0J-wX8p8>{J{qwP-Ilu!i4&DAfEpsc{J?Zk3A^$o{6ckOOW#(_5C7ggf$UuN3aq!JGrhn1o zn(hX}J(ns&8hK8R%6Jl!t7d_in}H$B5m+DX3(croLeWuSooP%hEmoM*5SWm`N)CTW zdZa>9->8{#=MW~2wbx&9&&_7_1kEDTz~OL)3M=NqAEs_3=aFz?!#T}E&Wr#8r|1|e zz?dHSX1Q9rg5ra?R>kN6B=3SH1Nl%APUXKVj#TiFZ=4inA&YKSPBU(Zhh{3`b$Kf- z@b%Kd=sGX0ABJXKWt5CkU=85&_MZL-F!g40v>@PT7hsJ{!g?IN8Na=~<*Jt67SCGjVDW*s_N|p4>B^ z2gS5^JRx(bx50$R5OvmbAR-RO?)1dkOBm6g2?)QBry}JNE32Vs3rX2gFr~}q z70YZE)h3E&4b!W#>xw{R%<)7;jWsp4yfh-nv;oEIZ<@e%;5yxaWAc-e94J!A7PGxT z-PV8!A%Tq#^GMebOrF~mMZqFn4O$yeo_iN{%$fJmJxee*eC}~_9t*ssBBd4&51lNo z0sI%5o?V1SYxQ4>Gy>FT201_AJF#crp)>H!QhsPlJN}*l8Bk1>LM2b{{lQUDc~ywL zbtu-ALS@>)<9JZPl}HDPQI!;e%6^|S?+s_$H`GqeDdpn;|22;lcLE6E_$PLlhJ@oJ zaF+D*qr=0g=~yC(sh1*0i~9$8!ZJY1m8Fw|3<1X{COW%4q7C^L+qDMMyF>z;-Vj61 zHjZ_CQpCBGYPO2kQQBX6)AdTvs|~V8IH%0vz$wgCOhUS;z~_={A`b=bHsEWBdx=qZ z7gwdT76<$ik@I|J$RhOp1r95|zc|0~xuh`)K`P_{bI=^*)~aOEG9(*0pB6g&dRhcP znXjhL;wPC3nP7VqqahqUIs`*}z+6?B_6S*uX0k?o#f;Xw^rRwEYpgJ5bR}4a4;FrI z7-eXcIi(3>@;O5;UOC37Hj$=SK6BoX7KbDOefKu^DtT57x;TcD%*Fw?Pl=UbJIH21 zY<(2RVnOQC7KRM$EN;>J>jRJxut-sG~ED zW{6^@Py2%_$0iCXMM8k)P}7r#z^Ekjrt0w&WH@w7xSo3m_9=^K z@4b!+r;&=O{Ar4;>zaz`oK|z4BC4v3frt@I~D z0MViAlPxVLM3M=h9jIbAgr?pRQcn?6q7Y;ddYoLM;&63}JM*A1igVhyP2IdW@*Pwp z;W(>uptRG<2Lyc9nQ)A`(-TI!AK;&YU+08LJu$F6@u!bYQ%wH!>MQD8@+czbx4RnM zek-_c;&c?UMxDbhg1$Al_CfZYg%$FKBBDK#k48A)sYnjO`s5EBjX-IMPD_v9od93P z%C@;UIPpAiGFmE`2J5AWVrZ?aXb9~W>SQP)@~0_FIDQTg0LZ9ESD&mWd@%fC%MGKn zk`z~_Ule@diLhsj%^!3KL5eI1AmOZ{&@fqXy%vi{^4*PC$ziL}uJ4V}&C;lq*(d=G zoki-h>bB>t+jUBGlGYHh4Y9iM#%Mn9M5NYjX*3rc&QCsf9pq1=ugfuZ?IuZvZytN$ zkuJ_ODTr*9rCIaV4{y6GfByOBm%N%OD8SJ4OsMUU5gOR=^$}z=>PBa+W<^f;#_e{B zIpsDj^0`VLp!M&OjGiK9=Ja$61bCR3Z&fQ6PEN?}GO_0;!jOd_E<=KqfrKXyhg3eV z9wuuIu#~O>IF_8EYEEq>I|~|Esg|WaS4pPr=JSD_+kuNu-cHou1EB?Wd{~X%NzMs4EW1QN1f8_Xi@tVU+=wMfYnXa zQxQ}&RZ@S6s%kqn)M&3aa`*=niv@`Iff9LQ7XXEB@L_-E))0^HB2gU%QuY)%(J=j3g>>1ZvY+A?mZ@m zm9z4gZ0P}-P;Pi+D<*!5^LPmapV1j5&$MJ(0Pbq85$+$pTNA){MUV}PfkmgOdjk)x z5TRZ1bHG=9+8qSmfpl(~vw{&wjvc(i47N9zs((^d?<$z)JGEHZnB%;f(dB zSz6@zDY51G)Fu3e^oRL4He%Koix3b1#RbsdARP4$-O<#cptT*yyCriNU%GrPMVL^H zdTS{ty+APC<{7fE1*I}r2_6e@wK%5V1|sJz05QssT;U)%kMJlo<;bJB#O_h93V(va z!52=;1{B(;98UNiiQ$uMkEo)`ooq8F&jbkDXqF`b|NT5)r9qW&9ugxUr9dF-0gewI z@7btm6GdL>EikiOU9vlPf++)x>d_h#Ol$lhU%KMoX9Amh`uD!~Js;C8?12pd zxJc3RZ^)zkdz)2xe2BLuIo1p0SS%cm5uEQc4MpmSLDPDa^X5nhn9f`>D!k@vFY87% z4j(f0{GWty25tHXL{ld=ZfppWCX469T%yb-4*&i1s_Y<`b7*$tLEfnM$PjZpiQp|l z6gXNkV9k+U!ME=eYD-9r@GYd0Nz3ids1v~y^moxz%wT8`LO$!p&fGKba2eR&ejhHk z?=-T~reaEgI;g@b#|rV2Cr{MDPG2uJno8=FRRKv>?)@2RVii`$m0wE~Yv)R*NJUd^ z?Xg-+R}Ce6=~TE>d81ApC$!il+1Y~=RQYG!VgIXN{fb>G8ZjUoAUswO>Bgi5mLW#P zi`eW;04Sf~1zNp8WbK1KfbA#pk5i_h(dwTqYcx4pgIbfn#bM z^Ph+1Au+!(8g*&M?jCJ(*Gd@I!h(#*$CXFsbeA~!)3^P7?XyJ34JjL90b(2+d ztxu(%JTY|gj~_qwWk-zeVL}L&KuO43m~KVB&`YUH+b5rV;=f`-7&3v(%{s@C1ZZG| zU_;5$cpI$86v?=M$es@hZgQDR__S};@t?(<7loC8D?ZG@tOP{;_Q+%Q3Ijdr(d zmL8u)!4^W95u;Z=*|TlMXtB{HLzn~M;n2o__oWLE%&&R=^D)t!sC^c5+ft+CRaklJKIY5M!De<2_eApZ=x+! z$h;-cyc?p+C7j$V1iA?6#VpAbr6 zhWLq-UCS~LF)JW>6%JUVW1S&uhNSK9r!X>dWj68M*FE}oB2EDW@KwvLx~~S-7)&`7 z-ObGnUFLFNwVU;|FI~Y}kW^%875FyvW-eqYv_Jge4`*Qsm?=JJicvFV*_%)*4n%LEIgeJxIi@`J_DH}K^PB@@nr?(D4FoYI>cv*+cq9GQgQBMW}PtdZBMo7XO z0m76h!eUoNeb;~~8+^FEy_J6|fdp%`M|_;po=oNl?KIQp{~el>h9nsd zyR8AiT5=rl+dBcrnTp3H0_)d=EQSwutOMNJo`HwRz_VKX5YRj3yPN^#uM@)YqX?>f z3aAq4(OZqmt3;~8YTT4Lw{!Ox$AAK@Vk-JT6iS?Qo!QjUQ<_ytNrpqx5GXg`vuW>m zI6L%G%;KWL)*OMDI35MaVaJ%}#BzFdbwfLc42imfM9+DOg(D*8mYWCIm;ebRlZW^o z$!q8y8p0CG+H6vVSKB$4Dyo-c`b+C5*VFp4FoqQH(&x|ccq6Jku#7rV$|bo z#O5)hoz0lEW(XmKBhojsvyMh*LO7tT2gkp!A%^fn%MBclVzC$`$D{Rc$jUKsTJk)x zUkpi?H7^xz67Et)bJC>KE!${D{kmuc;s7^vB{@dhFR24exG+KGSfhkBb|cs_ah!Lf zwd?hadf(%J>d-WICQTBAAYLkVx8wi;JH5E90T+70wh~3wg_(PLe0hN51(*+BkjQrl z%b!g@iujgX4J<<&$sTj&>`=}?Kr|_5#bTB*R+>l5$ykN142w=Pq^5CIsT2`9_sl*; zd`XtH@ogqzFvk-^Qe=e~k`@Tp=RL;=Q`FeCT5fJ``hGF6LqL-O@|qa+eA9c?vEGRB zHjA5Cq=@ZXUh*T08>3y}h+6KK}S)Wx*mL;N46Qq!gKP+u!;JWl6a7<9l;c zL)dcR#vF62tqo0BGC8jS)C-yQjo7O_0?&s-f}ZaYezw$4F-0dvDuR+&6Q@>qEdL>p zp$J*eb4)!-ZYBOs^IUK$3ft6=#Ta2xCTSKExo)%k~gA<(((lR?D0X-&_?D8#n*{}!& zqNzf*l9jBO5H3+0=RG!G)VbHFdSc;u_2s8(Cdt;7-YCq7!Qm=8uZFo2*(Fwttf<=r zOo3CcYXc?(!J0&3+L~>;T7FU#vd2syN|8reK+Ls`U>|+?<(C!%CWp-?Adp2QX$m`|GH=5w3lbV?&Z!c4)pe`@1*IpQ2L#4 zj$UyUGCl!f9TlC364Du9$l1_AZU`(V4Kb}D(n2I-v}{X}U)w|)j$?_gAq}T_hxzyaqOT~HhtD@lUNAHAqG@(xt((v>cJibreDEQVG@Cfd}ydb;KkmMaUu z-AJl4M>xW0UP&q4R)|UGtx2)EEE_TcAz#b-rYDdC2XCGq>f|?K^C)T2L-a^DqG<|h zv^E5ML0LKGSESvrXZAxDorHT`lvT~`)2eodbAzF45{YRRkIO_0kGxSCvq6~1hhl+% zV1hi0dW#TvqZ{`&k8m{1grk|pK}%qCS}UNbkqymdk>^6-YA>b4!f_F|aV$}f?xXIP zE=em?Vy;I*W5~}|K-2*`iAXh!k#3SlFdbdyLdTS7)_wDe>&1eKWroM=5Q$uR{s z3$(0nd^>&Q4@`EsWWiXmRgY6ioWW#Hql;O2q+VjVgc)q^CB{lxZ1Ge1GMg;=*~)e% z^WJ;!nR;)wAp6&{TCRq;atrpCg5u@!C+r-~}02kL( z3F%@*-pI@}*+l-1#zw=DbDLCEuh?ctF{!hJT>P8}Mr|}( zEXS@}yyVHw-!t&-GqAt?e)}Bm2w61RNoc)6g@h$ z9I*;I4l5iJtRvGQ%_~LU>4L^|#*iUE%khs8K8_O(hd`%-bh&&8oFNEKjtEea<0NyU zFq9|y&T_rv_(V+}vxFeDTm_PcdWr2(uMsvKB+HUKYxo`wrC#4*77c_&-g+#?b#&Dl zn!qZHo`m30m~-urB0rzX8(qh+AugF3%S^(hckG@Lpa>Li zDbg46sW48SByZLj#q>zUqa-jbSBafPh-vY=>bbaE^Rj17@qvsYH{kG121Mj;;R5PT zZ(7R_0i#qH)mv#mU<#a9R6q?HD2|b|gXEc_okax*Nsm`tEg}MmSs`e>{1@q{LklPC z(qgltXG!zi1n~;fD1^`~8@rm^!gy2c5f_f$+`X z(1jp&tkpp&bO|%xEG=jW4@iphcuJ6`POJ@u&RcJWR=ZVEXpL4b2ymDpfAq=M z!(SaB)|%mPrcJ%F57{Q7E{7JFBwOchAU2&zVN_P}jiR-zkFB6iA5EX*NqQ+-R~gMQ zZfG+v5Jc+kCuwl}SM>Cjw0SQCk6MQIe$T+0$-oO&+c#r!xAHDzpt7#~c_|gITZ;>g zbF0%GUkUb#<-Nr{MlDt=jXDe6W)PiEdN*3{dzwn^yd!yLdjTJm$r5Q~I~rQ-2`bv#p9l>8&c;hvgruA!!ep(|G| zisS*hT(YZYx9GC6=@Vvj_>2OP1T$+H2$QSx4l_wx710`VK*%Gy2~l*tlQGvdi^MX3 z(>w5(+YrYD96BLuMuXjm&10civV@w<#4FR0^AzQZGlr-e;*3SB$y=UG-a^PFLpGFa zQ^8s=py05}F;j08h^+qq_U8m0Q?fNBOXMfu&OT@IpABs$sP;RWf?fd|<_zj4FkdV7Fuz`!j>d8>gsa2*Dh1r= zTC|<2R!4~fDE<;P(y1~FyJ8idev7trXFFB8E{e8y*{<12p7+!l>{O{t#{`g1S5~tY zQCCFflmFs#>Ym=~w8lFo;P(~Bvk23eTT`*sZ&$=zp=Y%;kZU`$sLH8gi=_3LMy<8x z+kH*5gx|8Uy#~_bhbUiJXQ*9BCl|n#ot!B8_TSy*2eqDC#PB+m9J#B zbyu@3q-W(nH@^m}a`b}gD(G`I81Gw~r8w)PsE%Qm@U8I7QJ;0u5axzi%jGl?eU4#g zgKf?9wM()a&P>)RUxCIx_n$Hgpf%Ys_tNsy=xe=KY+@L(TRi7$XKQnY+{`u$-2u8^ zHLYzyTZFaQ*)~q)cMt4T?U!HJr8^3vXdA+moJpK!%MY+M&?q{}tlzpT$ewqjQBKRR zdu;cL3U}YA>+XBa>{OY-o`A23T9`J=Y;o3gebs5A_`fUj3!l`=4X|R&F?rNvAwb^CZdhRL+q1kUC zb(@{viRM&i=a`3AW6Y-`|GH~N>ZsOA*Gos~qAI_VooW$qZq9Yl(brDpTo_d; zyz16$IkjB%fBoxUji{W)d%YuNdmGCcrUZfSAz=6WJyf;%;eD=~UAMxXJ>498N1wOh zXu4^3!|gq$hhH15?v*`~yA@Wf&-1`5zjQam&UUZtM%k`o{<_=ar+21#l-6Kn<`LCa zxu?-X{dp=pJoQ!p~l1y0FH)OCph>$9<(HPAay=^|zJ}IF3#`}o0xt4d^A&5P9TQ1`)p$n@)^(k> z<{Q_9+S!VAs!kht#in}sr@KNLx2}|4q3W);pMDo;%36mpPDiKwEfHX^DdCjg7;CD} zN!0JmLeUmaopyG*tIcvIzjVvB!u8maZU|MMaaFziCc5CK{8TA_5@nwJ0w`U{TFa_6 z*@&j2I&I<2+N;4CZ-pnk+lbEBYE`}ET0D^!TgTKrWlmyj&)1V^9;;H<1zdx5`mBYi zGG$87xcVAJ3#saL(TvPx*eLp1W@qQJ?NqPYBVVl;o0JwqXU{Mj&IxR~=(A$I-FAJI z=vwZmO1<9-Un`dG?2M>HF?39WW!_58*Fx%43#`SQKu33>a)RxhyQz<8dYpmRU}dJG z2Gt{=x*gq!x+vPB%r0e`jiU0clRy3GPkr{tNU16<&elVhTZLN^wRXyv?q*nzb=TOt zA=FDZwElKC0pC6=)=}mJQ=*Q#Qf|e*`0F%2)stxx+m+Zr#^imom2X-bPN#xu1{3IO ze|0S9%)*DD-@=|Nv;q~bVxy?#;{WSk|2mCVtSRWp*aGN@(;7&xo$ZuwYg(IScC@rw zekE$H2)q32x*41Y+_W}x*Jr0D=k{4^^*F^^0R7c>euHSHilJGneDU<`?m}+To1m_~ z=CpLvJxh8*>a(w9HpT|kF%QsMsP{5EDl@14_P4*4=u}B)e|UFz+hNY5Xr3SQ z%qagnv&%mfi?jIiNv1e^HJQhA1$rz^Yx6{WF^UPGPTLstEUCw;PYhG5oTAOGuXE~J zV4iNz&w3>sx>3ES()3!J*-v0<=0WO-&w{m*v4A13Vo{)Rb`7hA_w?_sC)RerrU zijHQmboi~$Vi@W&YwCUVa%RT+^Hi&%*M)x2Ddw|i-}gSl40)DE)hkouvo_wBuRy(2 zU=l^xi%3&kU|oj21v&lNw^0e0HiKN6eJJ__%H~n_1Fs9bHO}fwP*9vFVvaEckPUI(1;f|?0Ex-S4$PK6XyY2MeHgy|palS8H zPK6)lSFuO}Nv6-5ecg32H}+~(tnT`(N(-RbX;8%+bxQ5r`mAJUProhZc`L1>DmHP& z-^J3Zn|6+B?Q~2pRcYS43W{(F&)J%-R_W9JEpdt!bD?$&_^sG{t*LrXsnt@kx-K1N zjTGJtXG(UeMEQrJmK)lPt_XWI@4A?YsF4}k)T)c_NOjR>X*Owp{_~&b1^PraCHq?8 zUdpRpy2j_{uxGp9Q$7F0s+)h2pZu>(T%DWK#y&^gJ#4u$#n7n|b2QymX){yi45~!U z&UpK*e3ISkDp|9gDslQP(OtS97He$^G|YZ(NO--5#kfI@>^IIDM^TM_=n}uCglDNDgz9 zUrtAheB!ComQX`0_>L~pZp)pme2tf1v0A9gl)rse>Zq@IZcU<`>(}n%9k*hQrK4hnT4Hv3Rj*jjhDn^y61Ss~c2vEVucykKhkE(bay}I( zvtFjPQkGz}I*=`)X=x*fGM!PkC| z-ukWg+Uz~G7nGib^>wgX26`P!qm zlC7ds=GM_EIf;p<0oKv{f*x8`$|;P3?_ob&BX;7ywsX4ZrJU-vQ+PV6{VCkn zn(fpr!ueX=7G4ewq@2Etq3&iZC*WJf-VUe0L_5=07*%WqmRW;c zbX^E@s=!j}=t?Qn7ItM?t|82#?S;7V^-`H}I6W3atrb{*QmdWHx7ur`%;~)h49G#XJth)-;t4q|fUVRA!D&wSM|cyF@LQA5n$Nsa74$yp&%g zHQ3pZP3GscsLrex#UtPE$Oy zP<{oXu913d<+Qwe1k`nzRVo%}eFqNd9Z<9$mkq^%sTStg2O+1cDvgYxM%C!Hx;6rW#1y z&CRgb0_Cs0;l_9H{ZGM&x5)9 zau0@X!QDWo`V{zRQ!je?F{g@ueqQZOXr4XiQB(luargs*8P($HNUt!`>nN)8`EA3f zmpYyH>vaCjRqsaQ&|~;?_mvK{+yblFUK-|seOmtMi3a}0{Qr7i`Nu#0(cf5lm+C)( z*xUTKifv|g)>YlL_3gDI-j>|a_~`Z0Dt#DFuQd>k?!Ylm~4xC&JL_0)8u?q#~`sEWO+ z&inW4DiiA5Pg-mb3m4?%mk1<*K098l{~1?}_u%L#@^1cZ;U2bsGHWyQR(>nFMBA`) z>Y5me?igK+U6Mil|F{3V(?@--jr=z0s5_ea=Dmk(IDdNbUw*1gT-mlTx`(%_-y?j= zdU#)Ja#}^tPCd-OK1LrF_*t2)!?})T*EjpjhruMD*jzrRz}Y&UoRg^hq^r_=^>b=J z^mJWRrAg^?Oy+iRYW?^`;Y}p1mhMP_bp3XrJ&a`TlT)Q3%w>35cvwXyfiDxV`~5OE zwZ2@STYUGpZpyQF?Izwk(EJeiu+e!cbRV4eyzZWr?`tpg=b2ZVot**?8~KTu#nupt zt(Z#`e@DSj&fMP%uTvf8>vsycE^7G-_15iw{bGGR`whd$q1Mg-o5YsEtbyk7+G!zO zzkW&#HQsUBspUGl?x;VdksnZ|p-r6i)%!~S60EQN2K4Oh{0l;w3B5Yc_0k*aTrU;q zeInt?JaewV!>&)Wx3kaMxfNUc^~36XJDP%-g($iY3`lx&^tD8v&C#4XiTQd8oUfCZ zT9dBX-h=xWzwJ-@t6Ce&`TN_cQXtRf@w#|gHf?gcLOV{~TcYyUu4}oRJ|Av}Zgu-w z;Z|CoFUnt|EuOl25l@9H&@p{YFnvC|70iou`;))_{qMb<|K%@#>0e*JkMp!xa*gH9b)$j@{9LK<>>0SioQcqgehs zLG~;pzadN=fdhzfV<#fPfQ(nVk^VyKz`*}D-2hsTzgdb{? zSE@S!V17E8CvrRpUNO(2=%$-Im-?OYUmrDQ9zzn{7smnuOd@@w55FjCC0g5{pXFN2 zXK?^8zWlj3${18c354)h+&cA#f1&etg%3*pje^aXzqF^tMzgOil+_V|o5IL^u-;N4 z8d%isbc1GnuB@e!m99~PJjRckKgD_GGi(ab1=#wyUba@%{SUU1%;kD8DAjSa<3Zmv zxl0cEW@ySCk!Hj=nA9Sbn&e`a(X;_J$NZ?fEA%1dJ36zlD6fmHv@%n)QDs{Q5X1GO z!BMLLS6H;3HCbB;28h5PEu$Fjm<)AGMM(4k*Xt$p2ud0LGc0#&s?aECvt>ZFdISbJ zpHL32vkLII>QGQCuv_6?7G2AVpWc?Cd(vRz)Rt}zQ70IX=pd!46j1L}32_OY_}a3S zPH9qwpqr?T@OM}w1`Nd9oxuvHm7c-w#;lMy7tSJV81=n8x)e4w*%(G@AuxenPdG$+ zJU3z^x=3URHA$Nyl)J!h%ww4s0`w@$u&J4rYL1lmr0%dJwbv`xz1IUaDyp`HLmfHz zYHaN7_;MdSUdj%al{$^ZyzT4nKD;BA1bIPKc_N&>nA=fGszpF!fj3Fegw#ou!ArrGGd@A#yQqd!R8ZxeSV z(@5V9CwO?BX0g!q)XGX+vP&&7?h>Ib{J_>sCrV^hwZn!Bq?$W^g{b2d_*A%V{IiMT zj`sAV2U|E|m2)-|j5``TUG7C*Fmpdt7S64d9JmOKCz>YyP3Fxe=*E^@MOg~LU$m5u zPOcsa%6yH(A++#i1ST$0-2arBKabc(SDS zRa4MnXkSJCi_PfW?2`B5Hz*f6+_hdPkQ_c6v(R(GTpE6nd}g|^Tt-eBte;Cy_ml>W zz;^`98MBH{I;7ArT!ZD><~GL4tM$^J=1u-PWxG;g{dc~WPhjnP>6ON($DJ&PWpaff zdit3wTY-gv|5T@l>F_X(PS_;mnV_5rPuzcGJe>?4RSY zRcfEIW0%TwNWo3Y?umnb4GAcIsk|&Ir4%wx{Xl+A4RV*qcx`1cT4sj^MoSMr;rg|uQAB7Q}Ab%Gx=7 zM$Vop$8qQP$ena6YZ463=IDAY9~|Zno~Ik>&U#lsbl<`UWr1C#9d7AX8O>q7Y7;(y zOm65qiVxNl77NTq3IhR-$R5-lsZRNWs#7mFRJ&<2M|+Xd!*D)?ANjS@9dl|HAF{wy z9PjUL0NZm0K9hu7kX4&yLU_xy(d5S&rY4uLF>-nH%wa`Q5mNL}T*#=StH9TT*3mH7 zn=yqCMuY=zbSu^QH1(JKj;&Uh)hl_o^_wBXkNVx~Uni5#M6OA*Z%QgLX4EStS$)(1iW?yo?Wf*jpkl&2?k#Qk;I{Ygs>WW4%IP75d)7$BWmr&2}Y` z@lJaE2~-`Lmd>B9rNg32bT4l9qw;z(@3s<2;-#!2#DO?%XB|EUHgiO3xpm=O*eCj` zE%cW6ofb69*7Tu)?i^X)hp0C3ZU$*PD}iUz?;Oy z;%`!v_TsHkN^Xczwgu3q>c@UAuv;kP$%=D>BOIfK3B~hvP!!HBc`#lucWe)Bi9l+&l!#$TbGiuKODo}uw2lwf z{uLv(9iwA@IEVRP>@r`i`#_Nxn<^_fwjWRRZ|B&5zqu!LfBV9zJh5LuLvn~vYgYNKxao&oAdM^q&@*rmov3-ykNXnlcsazd^ZcDu9aBgb612)W-Tcs(!eF5M3yWo$Od)Smu^8TmxQL0?i z&g|<0Of2^OxA#Z!28-zW=ZL?_^mLShto|=7_WxCHTyFv(z_hjgI)Cd|PlX!A-_58o z&qp*F?9+hak0YW1g$5Lw`%!a0YDiPV3V*t18dlJ-f`%0|te|0q|C<$*&S~l|O5aeU R -- -- -- -- +- 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: -- 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: +- 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 diff --git a/docs/design/sm70_deepseek_v4_dspark_n7_latest_main.md b/docs/design/sm70_deepseek_v4_dspark_n7_latest_main.md index 304038abdc..70563273f3 100644 --- a/docs/design/sm70_deepseek_v4_dspark_n7_latest_main.md +++ b/docs/design/sm70_deepseek_v4_dspark_n7_latest_main.md @@ -307,33 +307,6 @@ scoped greedy output remains byte-identical to control, and the quality prompt stops naturally after 76 coherent tokens. All task-owned services were stopped after validation. -## M=8 Real-Expert Grouping Candidate - -The fixed 48 one-row grouped verifier still reloads the same expert weights -when adjacent verifier rows select an overlapping expert. A default-off -candidate compacts those rows into the real expert segments already produced -on device, while retaining a fixed 48-entry graph contract. Empty tail groups -remain device-visible zero-row entries, so the route adds no host readback. -TurboMind keeps the single-group dispatch tactic but uses its ordinary -device-offset scheduler for the variable row bounds. - -An exclusively owned V100 CUDA Graph gate covered mixed overlap, 48 distinct -experts, and six hot experts. W13, activation, sorted W2 output, final output, -and changed-route replay were bitwise identical to the admitted slot-grouped -route in every case: - -| M=8 route distribution | Slot-grouped | Real-expert grouped | 43-layer projection | -|---|---:|---:|---:| -| Mixed, 34 unique experts | 0.42345 ms | 0.30437 ms | -5.120 ms | -| 48 unique experts | 0.42407 ms | 0.43113 ms | +0.304 ms | -| Six experts with eight rows each | 0.42183 ms | 0.15401 ms | -11.516 ms | - -The source gate released its GPU after completion. The route remains behind -`VLLM_SM70_MXFP4_MOE_GROUPED_M8_EXPERT_ROWS=1` until a synchronized TP8 -full-model profile proves that production routing overlap transfers the mixed -microbenchmark saving and the paired acceptance/quality gates pass. Evidence -is retained under `/data/models/dsv4-verifier-20ms-expert-grouped-m8-r2`. - ## Experiment Log | Date | Source | Test | Result | Decision | @@ -367,7 +340,6 @@ is retained under `/data/models/dsv4-verifier-20ms-expert-grouped-m8-r2`. | 2026-08-03 | candidate | Exact-selector seeds 4210-4219 | TPOT 31.282 -> 25.533 ms; emitted/chunk +4.0% | Accept speed and acceptance gates | | 2026-08-03 | candidate | Greedy equality and natural stop | 256-token SHA256 equal; coherent 71-token stop | Accept quality gate and enable by default | | 2026-08-03 | candidate | Scoped-policy full-model smoke | Target 32.426 ms; matched emitted/chunk -0.23%; greedy equal; natural stop | Accept final scoped selector | -| 2026-08-24 | candidate | Real-expert-grouped M=8 CUDA Graph gate | Three route distributions bitwise; mixed projects -5.120 ms/43 layers; distinct-48 projects +0.304 ms | Admit only to matched TP8 profile and quality gate | ## Artifacts diff --git a/docs/design/sm70_deepseek_v4_dspark_verifier_20ms_long_context.md b/docs/design/sm70_deepseek_v4_dspark_verifier_20ms_long_context.md deleted file mode 100644 index 66ebe2a582..0000000000 --- a/docs/design/sm70_deepseek_v4_dspark_verifier_20ms_long_context.md +++ /dev/null @@ -1,217 +0,0 @@ -# SM70 DeepSeek V4 DSpark Verifier and Long-Context Ledger - -This ledger tracks the 2026-08-24 TP8/V100 campaign for -`/data/models/DeepSeek-V4-Flash-0731`. It complements the historical N7 ledger -and prevents short route-hit smokes from being reused as verifier, long-context, -or quality evidence. - -## Acceptance Contract - -- Eight V100 GPUs, TP8, one request, DeepSeek V4 Flash DSpark with seven draft - tokens, temperature 1.0, top-p 1.0, and the production strict rejection - sampler. -- Primary speed gate: at least 100 pure decode token/s on the pinned 120-token - Chinese HTML prompt with 3,500 generated tokens. TTFT and prefill are reported - separately. -- Primary latency gate: complete verifier at most 20 ms per speculative round. - Complete verifier means target forward, target logits, and strict - rejection/sampling. Draft work, accepted-state commit, CPU scheduling, and - total round wall are reported separately rather than hidden in that number. -- Long-context gate: report raw contexts 1K, 4K, 16K, 64K, 128K, and 252K with - identical output length and sampling. Report verifier/round cost, acceptance - by position, emitted tokens/round, and pure decode separately. -- Quality gate: compare no-speculation and DSpark on deterministic GSM8K, - HumanEval, a multilingual/multitask LongBench subset, and 128K/252K needle - retrieval. A speed result cannot promote a route that loses the paired quality - or output-health gates. -- Shared-machine rule: launch only after an event-driven free-GPU gate, never - stop unrelated owners, avoid frequent polling, and terminate every task-owned - worker after the artifact is complete. - -## Reproduced Baselines - -Current latest-main source plus the admitted SM70 route repairs reproduces the -historical endpoint target. Three 3,500-token runs measured 118.481, 115.975, -and 116.288 token/s; median throughput is 116.288 token/s. The third run -accepted 2,816 of 4,781 proposed draft tokens (0.589), emitted 5.123 tokens per -round, and had per-position unconditional acceptance -`0.8873/0.7818/0.6911/0.6047/0.5066/0.3880/0.2635`. - -The matched no-speculation route measured 66.014 token/s and 15.1483 ms/token -at 1,024 input plus 256 output tokens. A synchronized DSpark profile measured: - -| Component | Steady median per round | -|---|---:| -| Target forward, M=8 | 34.295 ms | -| Target logits | 0.326 ms | -| Strict rejection/sample | 0.499 ms | -| Complete verifier | 35.120 ms | -| Draft GPU work | 4.419 ms | -| Total round wall | 43.505 ms | - -The 20 ms goal therefore requires target-forward work, not sampler-only tuning. - -## Admitted Source Gates - -The following changes have passed focused current-source gates but still require -their stated end-to-end promotion gate: - -1. FP32 mHC staging now covers verifier rows M=1..8. M=8 falls from 0.05383 to - 0.01257 ms per call and projects 3.507 ms saved across 85 calls. Four V100 - numerical cases were bitwise exact. -2. The M=8 graph can use a bounded 2,048-token short-context bucket instead of - reserving the whole maximum context. Speculative buckets require an explicit - `VLLM_SM70_DSV4_DECODE_CONTEXT_BUCKETS`; CPU dispatcher tests pass, while the - synchronized TP8 transfer remains required. -3. A default-off `VLLM_SM70_DSV4_PRIVATE_COMPRESSOR_STATE=1` route stores - single-request SM70 compressor C4/C128 intermediates in bounded private - device rings when prefix caching, pipeline parallelism, KV transfer, and - parallel drafting are absent. Seven V100 ring/workspace tests and eleven CPU - route tests pass. It cannot become the default before the long-context - quality and capacity gates finish. -4. A default-off `VLLM_SM70_INDEXER_DECODE_CUBLAS=1` route gathers long-context - paged FP8 keys once and uses FP32 cuBLAS scores for bounded row shapes. It - retains per-head ReLU and FP8 scale semantics. The captured-graph test, - shorter dynamic replay, graph-tail mask, and workspace-width gate pass on - V100; the incomplete long-context quality gate keeps it opt-in. - -5. A default-off `VLLM_SM70_MXFP4_MOE_GROUPED_VERIFIER=1` route extends the - single-launch, one-row-per-slot MXFP4 contract from M=8 to M=2..M=8. This is - required before a shorter or confidence-scheduled DSpark block can reduce - verifier latency; the previous M=2..M=7 path launched one operation per - compact expert and erased the benefit of verifying fewer rows. - -The focused M=5 and M=6 CUDA Graph gates covered mixed overlap, all-distinct -slots, and six hot experts. Slot grouping and real-expert grouping were bitwise -at every pipeline stage and after a changed-route replay in all six cases. The -mixed-route results are: - -| Verifier rows | Active-expert loop | Slot grouped | Real-expert grouped | -|---:|---:|---:|---:| -| M=5, 22 unique of 30 slots | 1.98837 ms | 0.18552 ms | 0.18000 ms | -| M=6, 26 unique of 36 slots | 1.32993 ms | 0.34895 ms | 0.27260 ms | - -All-distinct expert grouping is slightly slower than slot grouping -(`+0.121 ms` projected over 43 layers at M=5 and `+0.210 ms` at M=6), so the -full-model route must report the observed production overlap. Evidence is -retained under `/data/models/dsv4-verifier-20ms-variable-grouped-r1`; the -single V100 was released after the gate. - -The 0731 checkpoint declares an official DSpark block size of five and stores -`mtp.2.confidence_head.proj.weight`. The runtime instantiates and loads that -optional head only for explicit confidence-threshold scheduling or alignment -diagnostics and fails cleanly if a requested projection is absent. Ordinary -serving performs no confidence projection. DeepSpec computes conditional step -probabilities with a sigmoid, -uses their cumulative product for prefix reliability reporting, and truncates -at the first conditional probability below a threshold. Before production -confidence scheduling, this runtime must first collect calibration data on the -pinned datasets and profile M=2..M=8 service; a raw uncalibrated threshold is -not an acceptance or latency proof. Static N=4/N=5 trials use the same strict -rejection sampler and therefore preserve the target distribution while the -hardware width crossover is measured. - -The indexer crossover gate uses FP32 scores, signed head weights, top-k 512, -M=8, H=64, D=128, and CUDA Graph replay: - -| Compressed keys | Fused paged Triton | Gather + FP32 cuBLAS | 21-layer projection | -|---:|---:|---:|---:| -| 1,024 | 0.06854 ms | 0.02096 ms | -0.999 ms | -| 2,048 | 0.13455 ms | 0.02519 ms | -2.297 ms | -| 4,096 | 0.23433 ms | 0.04738 ms | -3.926 ms | -| 8,192 | 0.45660 ms | 0.07680 ms | -7.976 ms | -| 16,384 | 0.85016 ms | 0.14493 ms | -14.810 ms | - -All 40 rows retained the exact top-k set; maximum logit absolute difference was -`1.221e-4`. The opt-in crossover is 1,024 compressed keys. Generic batches, -noncontiguous inputs, more than eight rows, and non-ReLU semantics retain the -existing fused paged kernel. - -The production block-size candidate verifies M=5 rows. Its matched crossover -gate retained the exact top-k set in all 30 rows and measured: - -| Compressed keys | Fused paged Triton | Gather + FP32 cuBLAS | 21-layer projection | -|---:|---:|---:|---:| -| 1,024 | 0.03548 ms | 0.01905 ms | -0.345 ms | -| 2,048 | 0.06881 ms | 0.02196 ms | -0.984 ms | -| 4,096 | 0.13532 ms | 0.03410 ms | -2.126 ms | -| 8,192 | 0.26614 ms | 0.05806 ms | -4.370 ms | -| 16,384 | 0.58173 ms | 0.10066 ms | -10.103 ms | -| 65,536 | 1.92835 ms | 0.32707 ms | -33.627 ms | - -The maximum logit absolute difference was `1.221e-4`. The microbenchmark owner -released its V100 immediately after writing the artifact. - -The speculative SM70 graph context candidate uses explicit raw widths 2K, 4K, -16K, 64K, and 128K when the maximum is 256K. Longer requests retain the generic -full-context graph. `VLLM_SM70_DSV4_DECODE_CONTEXT_BUCKETS` is required to -enable the candidate; an explicit MTP bucket override, including an empty -value, remains authoritative. Automatic non-speculative selection uses only -the compressed-index operator configuration and never architecture or -checkpoint identity. - -## Long-Context Diagnosis - -DeepSeek V4 does not use the repository's classic Flash-V100 dense-attention -decode route. Its production log identifies packed-FP8 SM70 sparse MLA with C4, -C128, SWA-128, and top-k 512. Sparse attention after selection is nearly fixed -work; the C4 indexer scan grows with compressed history and was the expected -linear decay source. The opt-in private compressor rings address capacity and -memory growth, while opt-in gather-once FP32 cuBLAS addresses the repeated -index scan. - -`benchmarks/benchmark_dsv4_dspark_long_context.py` records each completed -context incrementally. It uses server-reported prompt usage, separates TTFT -from pure decode, snapshots speculative counters around each request, and gates -suffix-marker retrieval, replacement characters, token diversity, and repeated -token runs. When given the TP0 server log it also records interior synchronized -profiler windows for target forward, target logits, strict rejection/sampling, -and their summed complete-verifier cost. The verifier row count and acceptance -positions follow the configured speculative width instead of assuming N=7. A -zero speculative width records the matched no-speculation long-context speed -and output-health curve without fabricating verifier counters. - -## Quality Gate - -`benchmarks/benchmark_dsv4_quality_api.py` provides deterministic HumanEval and -LongBench-subset API gates. HumanEval code executes with Landlock filesystem -isolation, seccomp network/process restrictions, a private temporary directory, -and CPU, address-space, output-file, and process limits. The sandbox has proven -that normal Python works while `/home`, `/data`, global `/tmp`, networking, -forking, and signalling other processes are denied. - -The first paired LongBench subset is `hotpotqa`, `multifieldqa_zh`, -`gov_report`, and `lcc`. Official LongBench metrics are loaded from the pinned -local checkout; `jieba==0.42.1`, `rouge==1.0.1`, and `fuzzywuzzy==0.18.0` are -isolated under `/data/models/dsv4-quality-deps` rather than installed into the -shared runtime environment. Every quality artifact includes SHA-256 hashes for -the HumanEval tasks, LongBench prompt/max-length configs, and each selected -dataset. LongBench records retain their source-row indices so the matched -no-speculation and DSpark scores cannot silently use different samples. -`benchmarks/compare_dsv4_quality_results.py` then rejects any reference-pass to -candidate-fail transition in HumanEval, GSM8K, an individual LongBench row, or -either final/anywhere needle retrieval. Aggregate scores alone cannot hide a -per-sample regression. - -## Rejected or Deferred Paths - -- A TP8 M=8 64 KiB hierarchical all-reduce was slower than NCCL in isolation - and saved only 0.028 ms across 87 joined calls, within noise. It was fully - reverted and must not be repeated without a different communication design. -- The expert-row grouped MXFP4 candidate is tracked separately. Its bitwise - microbenchmarks justify a full TP8 trial but not promotion by themselves. -- Disabling per-head ReLU, changing index-cache layout, reducing index top-k, - or weakening strict rejection are outside the quality contract. - -## Artifact Roots - -- 100+ endpoint: `/data/models/v100-dsv4-0731-tp8-dspark-prob-n7-latest-c4-20260824-r6` -- no-speculation: `/data/models/v100-dsv4-0731-tp8-nospec-latest-c4-20260824-r1` -- synchronized breakdown: `/data/models/v100-dsv4-0731-tp8-dspark-prob-n7-profile-latest-c4-20260824-r1` -- mHC source gate: `/data/models/dsv4-verifier-20ms-mhc-source-gate-r2-cuda128` -- private rings: `/data/models/dsv4-verifier-20ms-private-ring-source-gate-r1` -- indexer crossover: `/data/models/dsv4-verifier-20ms-indexer-crossover-r3` -- indexer source gate: `/data/models/dsv4-verifier-20ms-indexer-source-gate-r1` -- M=5 indexer crossover: `/data/models/dsv4-verifier-20ms-indexer-m5-r1` -- M=5/M=6 grouped verifier: `/data/models/dsv4-verifier-20ms-variable-grouped-r1` -- rejected TP8 all-reduce: `/data/models/dsv4-verifier-20ms-tp8-ar-screen-r1` diff --git a/docs/design/sm70_deepseek_v4_flash_bringup.md b/docs/design/sm70_deepseek_v4_flash_bringup.md index 60be98350c..750f8a896b 100644 --- a/docs/design/sm70_deepseek_v4_flash_bringup.md +++ b/docs/design/sm70_deepseek_v4_flash_bringup.md @@ -55,40 +55,6 @@ kernel module and userspace NVML versions differ. Local validation is limited to static and CPU tests until reboot; GPU correctness and performance run on the remote V100 host after the PR is merged and freshly cloned. -## PP Topology Follow-up (2026-08-22) - -The original bring-up named PP2 x TP4 as a required deployment candidate, but -the retained full-model evidence covers TP8 only. This follow-up removes the -configuration blocker on current `onecat/main` at -`7aede2cf010d92815c9d7bff25867b4fa009b6cb`; the full runtime evidence gap -remains explicit below. - -- The roughly 160 GB checkpoint cannot fit pure TP4 or pure PP4 on four - V100-32GB GPUs. The memory-feasible eight-GPU candidates are PP2 x TP4 and, - if four pipeline stages are required, PP4 x TP2. -- The default 43-layer partitions are `22+21` for PP2 and `11+11+11+10` for - PP4. DSpark target layers 40-42 remain entirely on the final pipeline stage - in both layouts, matching the runner's final-stage-only drafter ownership. -- PP intermediate tensors must retain the four mHC streams with shape - `[tokens, 4, 4096]`; collapsing them to a two-dimensional hidden state at a - stage boundary would change model semantics. -- The first runtime gate is PP2 x TP4 with the same model, FP8 MLA KV, DSpark - N7, 1024/256 request, sampling, graph mode, and fast-stack flags as the - accepted TP8 reproduction. The TP8-only hierarchical all-reduce is disabled - for each TP4 group. - -The first current-main PP2 x TP4 attempt failed before worker launch. DSpark's -draft config inherited target `pipeline_parallel_size=2`, so generic model -validation required the final-stage-local `DSparkDraftModel` to implement -`SupportsPP`. The runner intentionally constructs the complete drafter only on -the final target PP rank; partitioning that three-layer draft again would not -match its ownership model. The candidate therefore gives DSpark a draft-local -PP size of one while retaining the target TP/PP topology and draft TP size. It -also rejects a custom layer partition if any configured target layer falls -outside the final target pipeline stage. Exact configuration regressions pass -for PP2 x TP4 and PP4 x TP2. Post-fix runtime route, quality, speed, and cleanup -evidence remain pending. - ## Upstream Reuse Reuse these official implementations and semantics: diff --git a/docs/design/sm70_dflash2_acceptance_rd.md b/docs/design/sm70_dflash2_acceptance_rd.md deleted file mode 100644 index c190844ec0..0000000000 --- a/docs/design/sm70_dflash2_acceptance_rd.md +++ /dev/null @@ -1,77 +0,0 @@ -# SM70 DFlash2 acceptance research without retraining - -> **Status:** the default-off alignment capture and offline analyzer are -> development tooling. They do not change the production proposal policy and -> do not claim an acceptance or performance improvement by themselves. - -## Scope and frozen contracts - -- Date: 2026-08-25. -- Integration base: `onecat/main@34403018d917054dd7765d5e820ad29c8d342348`. -- Branch: `codex/v100-v100-dflash2-acceptance-rd-20260825-113155`. -- No draft or target retraining and no checkpoint edits. -- Target validation remains the complete MRV2 block-eight workload: seven draft - tokens plus one target bonus row. DDTree index reuse and wider verifier trees - are outside this campaign. -- Target sampling remains exact standard rejection at temperature 1.0, top-p - 0.95, and top-k 20. A candidate proposal distribution is admissible only if - the selector draw and the rejection sampler consume the same exact `q`. - -The score gate remains the retained 128-row mixed32 dataset with SHA256 -`85ce10d84735ec981c001663d5748b1939a81892ef3ccd4549aec66177795607`. -Its private filesystem location is intentionally not committed. -It uses sequential B1, a fixed seed, xhigh reasoning, at most 2,048 generated -tokens, and natural EOS. Corrected MBPP-32 and WikiText-2 PPL retain their -existing independent gates. The historical native-grouped result is 83/96 on -GSM8K/MATH-500/HumanEval, 25/32 on corrected MBPP, request-mean completion -length 4.238430, and pooled acceptance length 3.749104. - -The production performance gate remains the PR #288 practical contract: TP4 -V100, 256K maximum context, `max_num_batched_tokens=4096`, prefix caching, -Mamba align mode, FP8 E5M2 target KV, FP16 draft KV, Flash-V100, probabilistic -DFlash2, tool/reasoning parsers, and FULL CUDA Graph. Its short-context complete -round is 18.465--18.603 ms. Acceptance work may not increase verifier rows or -regress that round by more than measurement noise. - -## First target and promotion gate - -The first research target is at least +0.5 request-mean completion tokens per -verification round on the unchanged mixed32 contract. Promotion additionally -requires: - -1. no aggregate task-score loss against the paired target-only and current - DFlash2 controls; -2. no corrected-MBPP or WikiText PPL regression; -3. exact rejection-distribution tests and normalized proposal rows; -4. unchanged eight-row target validation and at most 0.05 ms incremental - selector latency on SM70; -5. no material acceptance or latency regression at 1K, 32K, 128K, or 256K. - -## Selector headroom audit - -The first implementation is diagnostic-only. With both compact sparse -rejection and `VLLM_SPEC_DUMP_ALIGNMENT=1`, the speculator keeps the checkpoint -top16 IDs, unary logits, full `7 x 16 x 16` selector lattice, and exact realized -proposal rows. The target sampler records its top20 rows and the actual strict -rejection outcome. Disabled diagnostics allocate no shadow lattice and add no -copy or sampler operation. - -`benchmarks/analyze_dflash2_selector_alignment.py` reports, by draft depth: - -- current overlap `sum(min(p, q))`; -- target probability mass covered by the fixed top16 candidate support; -- the gap between support mass and current overlap; -- one-step counterfactual sweeps over proposal temperature, proposal nucleus, - unary/edge calibration, and backward log-sum-exp future messages. - -Even sampler steps tune the candidates and odd steps are a held-out report. -Counterfactual values after a changed path are explicitly only recorded-prefix -overlap proxies, not end-to-end acceptance claims. Survivors must be rerun on -the frozen dataset. - -## Initial verification - -- Focused Ruff lint and format: passed. -- Diagnostic analyzer synthetic checks: passed. -- CPU DFlash2 suite with CUDA hidden: 77 passed, 12 expected CUDA skips. -- No GPU performance or end-to-end acceptance result has been claimed yet. diff --git a/docs/design/sm70_dflash2_long_verify_decay.md b/docs/design/sm70_dflash2_long_verify_decay.md deleted file mode 100644 index 7a8f161214..0000000000 --- a/docs/design/sm70_dflash2_long_verify_decay.md +++ /dev/null @@ -1,465 +0,0 @@ -# SM70 DFlash2 long-context verification decay - -## Scope - -This task measures, attributes, and reduces the context-length growth of one -MRV2 DFlash2 verification round on V100. The historical study started on -private DFlash2 verifier head `e8daa778f55cf16df311f8ce8dd032e2281f1680`; -the public change set contains only the five retained verifier commits rebased -onto current `main`. It does not replace or retune DFlash1, `dflash_ddtree`, -Eagle, or MTP. - -DFlash2 names the measured workload, not a runtime activation identity. The -kernel entry admits only the explicit SM70, E5M2 paged-KV, Q1-Q8, Hq6/Hkv1, -D256, page-layout, causal-mask, and workspace contracts. It does not inspect a -model name, checkpoint, `model_type`, or architecture identity. - -Quality promotion remains score based. Greedy token identity is diagnostic; -the production gates are probabilistic task scores, PPL, acceptance length, -and output health against the same target-only checkpoint. - -## Frozen workload - -- Target: local Qwen3.8-27B mixed NVFP4/FP8 checkpoint, TP4. -- Draft: `incoai/Qwen3.8-27B-DFlash2`, FP16 draft KV. -- Target KV: `fp8_e5m2`. -- Attention: `FLASH_ATTN_V100` target and draft routes. -- Execution: FULL CUDA Graph, sequential batch one. -- Sampling: temperature 1.0, top-p 0.95, target top-k 20, probabilistic - proposals, fixed seed. -- Selector top-K: checkpoint contract 16; speculative tokens: 7. -- Context sweep: 1K, 32K, 128K, and 256K boundary. - -All comparisons must keep model, TP/GPU set, graph mode, KV dtypes, sampling, -and route gates identical. Prefill/TTFT are reported separately from steady -decode and verification-round wall time. - -## Existing evidence to reproduce - -The parent task recorded the following unprofiled paired measurements: - -| Context | Prior complete round | Grouped complete round | Prior target phase | Grouped target phase | -| --- | ---: | ---: | ---: | ---: | -| 32K | 36.311 ms | 26.361 ms | 30.248 ms | 20.452 ms | -| 128K | 69.855 ms | 40.778 ms | 63.708 ms | 34.647 ms | - -From 32K to 128K, the grouped complete round grows by 14.417 ms and its target -phase grows by 14.195 ms. This strongly localizes the observed growth to target -verification, but it does not yet distinguish paged attention service, -dependency gaps, launch geometry, communication placement, or trace overhead. - -The retained grouped-attention microbenchmark reported 0.0686/0.3564/1.2145/ -2.3962 ms at 1K/32K/128K/256K. These values must be re-run and reconciled with -the full-model critical-path wall before using them as a closed decomposition. - -## Closed trace attribution - -The retained 32K and 128K Nsight Systems traces use the same grouped one-pass -kernel and target CUDA Graph route. Summing the partial and combine kernels, -then dividing by the number of target graph replays across all four TP ranks, -gives the following service time per target verification graph: - -| Context | Partial average | Partial calls | Combine average | Attention per target graph | -| --- | ---: | ---: | ---: | ---: | -| 32K | 0.313606 ms | 1,344 | 0.008282 ms | 5.150214 ms | -| 128K | 1.198762 ms | 1,088 | 0.008295 ms | 19.312912 ms | - -There are sixteen full-attention layers per target graph. The attention service -therefore grows by 14.162698 ms between 32K and 128K, while the measured target -phase grows by 14.270586 ms. Grouped attention explains 99.24% of target-phase -growth and 95.83% of complete-round growth in this pair. Draft, selector, -sampling, TP exchange, and host residuals are not material causes of this -context-length slope. - -Both traces launch the partial kernel with grid `(2, 40, 1)`, block size 256, -128 registers per thread, and 39,424 bytes of static shared memory. The fixed -80-CTA grid maps one CTA to each V100 SM even as the KV length quadruples. The -kernel's `__launch_bounds__(256, 2)` and resource arithmetic permit two resident -CTAs per SM (65,536 registers and 78,848 shared-memory bytes), so the first -bounded experiment is 40 versus 80 context splits. This is a hypothesis about -exposed parallelism, not yet a performance claim: an exact-shape paired -microbenchmark must distinguish occupancy benefit from memory-bandwidth -saturation and extra split-reduction cost. - -The experiment follows the same split-KV principle used by the official -[FlashInfer scheduler](https://github.com/flashinfer-ai/flashinfer/blob/main/include/flashinfer/attention/scheduler.cuh), -which derives its maximum grid from CUDA occupancy, and by -[Flash-Decoding](https://princeton-nlp.github.io/flash-decoding/). Recent -[PersistentKV](https://arxiv.org/abs/2606.26666) results likewise motivate -adaptive sequence splitting and page-aware scheduling for long-context decode. -These sources guide the search space; only same-machine measurements can -promote a V100-specific default. - -## Independent upstream audit - -`sglang-V100` is derived from this project's V100 work. It remains useful for -route and result cross-checks, but it is not independent evidence for an -optimization decision. The bounded search therefore uses primary upstream -implementations and papers, pinned at the audited revisions: - -- Official FlashAttention revision - [`0251105a`](https://github.com/Dao-AILab/flash-attention/tree/0251105a2fb19d2957484b7f023cd8c115286ced) - packs `seqlen_q * qheads_per_kv_head` into the M dimension in - [`pack_gqa.h`](https://github.com/Dao-AILab/flash-attention/blob/0251105a2fb19d2957484b7f023cd8c115286ced/hopper/pack_gqa.h), - then computes dynamic Split-KV occupancy from the packed M work in - [`flash_prepare_scheduler.cu`](https://github.com/Dao-AILab/flash-attention/blob/0251105a2fb19d2957484b7f023cd8c115286ced/hopper/flash_prepare_scheduler.cu). - This is the direct algorithmic source for combining the six Q heads that - share the Qwen3.8 verifier KV head. Its SM80/SM90 implementation is not - copied; only the exact PackGQA scheduling principle is ported to Volta WMMA. -- FlashInfer revision - [`c18a974a`](https://github.com/flashinfer-ai/flashinfer/commit/c18a974a88b8c7da88be0338a7478636f1dffa17) - explicitly keeps head-dimension-256 prefill within the 64-KiB SM70/SM75 - shared-memory limit. Its release JIT successfully compiled on this V100, - but the grouped Q8/H6/D256 paged launch still ended in an unspecified launch - failure during CUDA Graph setup. The direct backend route is therefore - rejected; its packing and scheduler layout remain useful source material. -- TensorRT-LLM revision - [`958d651b`](https://github.com/NVIDIA/TensorRT-LLM/tree/958d651b18800da1da0677160e0277f701c0ff84) - supports multi-block and multi-query-token XQA, but - `supportConfigHMMA` rejects `SM < 80`. ONNX Runtime XQA likewise explicitly - requires compute capability 8.0 or newer. Neither is a drop-in V100 kernel; - both independently confirm head-group packing plus context splitting as the - intended speculative-verification topology. TensorRT-LLM's speculative - multi-block scheduler further caps the packed grid at one SM wave and avoids - scaling short histories, independently matching the bounded 80-CTA policy - tested below. -- [PersistentKV](https://arxiv.org/abs/2606.26666) maps native paged GQA work - by `(request, KV head, sequence split)`, explicitly reusing each K/V tile - across the grouped query heads. Its compact workqueue is valuable for ragged - concurrent batches; the frozen single-request Q8 verifier has no empty work - items to compact. The head-group mapping is directly applicable, while the - workqueue is retained as a later concurrency experiment rather than being - conflated with the B1 result. -- xFormers Split-K combines split states with LSE and supports small causal - query counts plus experimental paged attention. Its current Triton route - requires SM80, so it is algorithmic cross-check material, not a new V100 - dependency. -- [FlashAttention-3](https://papers.neurips.cc/paper_files/paper/2024/file/7ede97c3e082c6df10a8d6103a2eebd2-Paper-Conference.pdf), - [LeanAttention](https://arxiv.org/abs/2405.10480), and - [FlashDecoding++](https://arxiv.org/abs/2311.01282) provide the exact - split-softmax reduction, Stream-K work distribution, and GEMM-oriented - decode alternatives. They motivate bounded follow-ups, not unmeasured - performance claims on Volta. - -TinyFA revision `bc755512` was also inspected and excluded: its fast path -requires SM75-era asynchronous-copy machinery, lacks E5M2 paged KV support, -and its paged decode mapping does not pack all six query heads into one KV -load. This audit leaves native SM70 PackGQA as the smallest dependency-closed -candidate. - -## Split-count microbenchmark - -The first candidate doubles the maximum context splits from 40 to 80. With -two GQA head groups this changes the fixed partial launch from 80 to 160 CTAs, -matching the kernel's two-block-per-SM launch bound. No attention arithmetic, -mask, FP8 decode, or MRV2 scheduling code changes. - -An A/B/A CUDA-Graph run on physical GPU 1 used the production Qwen3.8 target -shape (Q8/H6/KV1/D256, E5M2 KV, page size 3,296), 40 warmups and 200 measured -rounds per point: - -| Context | Split 40 A | Split 80 | Split 40 A2 | Split 80 versus A/A2 | -| --- | ---: | ---: | ---: | ---: | -| 1K | 0.058368 ms | 0.067584 ms | 0.058368 ms | +15.79% / +15.79% | -| 32K | 0.319488 ms | 0.261120 ms | 0.322560 ms | -18.27% / -19.05% | -| 128K | 1.204736 ms | 0.955392 ms | 1.206272 ms | -20.70% / -20.80% | -| 256K | 2.377728 ms | 1.864704 ms | 2.378752 ms | -21.58% / -21.61% | - -At sixteen full-attention layers, the operator delta projects to -0.934 ms at -32K, -3.989 ms at 128K, and -8.208 ms at 256K per verification graph. The 1K -operator regression projects to +0.147 ms, or about +0.74% of the retained -approximately 20 ms complete short-context round. This remains a projection -until the TP4 full-model A/B is complete. - -The split-80 extension passes all six targeted grouped-verifier tests: five -random physical-page cases against the FP32 reference, and CUDA Graph replay -while mutating runtime sequence length. These are operator correctness gates; -acceptance length, PPL, and task scores remain required before promotion. - -## Full-model 32K and 128K confirmation - -A TP4 FULL-Graph Nsight run used the same 128K prompt, 128 generated tokens, -probabilistic sampling seed, model/KV dtypes, and route gates as retained -split-40 run `v80`. The only implementation change was the split-80 Flash-V100 -extension: - -| Phase | Split 40 | Split 80 | Delta | -| --- | ---: | ---: | ---: | -| Complete verification round | 40.778072 ms | 36.683712 ms | -4.094360 ms (-10.04%) | -| Target | 34.646866 ms | 30.557332 ms | -4.089534 ms (-11.80%) | -| Draft | 4.604373 ms | 4.627743 ms | +0.023370 ms | -| Draft to target | 0.230538 ms | 0.225575 ms | -0.004963 ms | -| Target to draft | 1.296296 ms | 1.273062 ms | -0.023233 ms | - -The partial kernel changed from grid `(2, 40, 1)` to `(2, 80, 1)`. Its mean -service time fell from 1.198762 ms to 0.944041 ms per full-attention layer -(-21.25%); combine rose from 0.008295 ms to 0.010430 ms. Across sixteen layers, -the net attention prediction is -4.041 ms, explaining 98.8% of the measured -target-phase improvement. The earlier operator projection was -3.989 ms, only -about 0.10 ms from the full target result. - -The fixed diagnostic request retained identical accepted-token counts -`[17, 17, 17, 16, 16, 15, 15]`, mean acceptance length 7.647059, and all 128 -output token IDs. This is trajectory evidence, not a substitute for the -required score/PPL gates. - -The matching 32K TP4 FULL-Graph trace independently confirms the shorter -context point: - -| Phase | Split 40 | Split 80 | Delta | -| --- | ---: | ---: | ---: | -| Complete verification round | 26.361447 ms | 25.238442 ms | -1.123005 ms (-4.26%) | -| Target | 20.451644 ms | 19.468326 ms | -0.983318 ms (-4.81%) | -| Draft | 4.365014 ms | 4.257568 ms | -0.107446 ms | -| Draft to target | 0.260419 ms | 0.235765 ms | -0.024654 ms | -| Target to draft | 1.284369 ms | 1.276783 ms | -0.007586 ms | - -The split-80 partial kernel averages 0.248999 ms at 32K, versus 0.313606 ms -for split 40. The measured target reduction is within 0.05 ms of the operator -projection. Across the paired full-model endpoints, the 32K-to-128K complete -round growth falls from 14.416625 ms to 11.445270 ms (-20.61%), while target -growth falls from 14.195222 ms to 11.089005 ms (-21.88%). The remaining slope -is still serious: target accounts for 96.89% of the optimized complete-round -growth, so subsequent work remains scoped to grouped paged attention. - -## Physical-page and tile-balanced follow-up - -Two address-equivalent scheduling changes pass the same six operator tests: - -- Specialize the actual FP8 kernel-page sizes (1,648 and 3,296) at compile - time, replacing runtime integer division with constant arithmetic. The - Qwen3.8 hybrid cache reports a 1,648-token logical attention block, while the - grouped verifier receives the 3,296-element physical FP8 view; the final - trace proves the `` specialization is selected. -- Divide the sequence by complete 32-token WMMA tiles, distributing the - remainder across splits. This removes per-split padded tail tiles without - changing visibility, loaded KV bytes, softmax, or WMMA arithmetic. - -Repeated CUDA-Graph A/B/A runs show no 32K regression and a further -approximately 1.7% operator reduction at 128K versus plain split 80. The final -TP4 FULL-Graph endpoints are: - -| Phase | Split 80 at 32K | Final at 32K | Split 80 at 128K | Final at 128K | -| --- | ---: | ---: | ---: | ---: | -| Complete verification round | 25.238442 ms | 25.237987 ms | 36.683712 ms | 36.482220 ms | -| Target | 19.468326 ms | 19.455442 ms | 30.557332 ms | 30.354906 ms | -| Draft | 4.257568 ms | 4.269400 ms | 4.627743 ms | 4.606652 ms | - -At 128K, the partial kernel falls from 0.944041 ms to 0.926568 ms per -full-attention layer; sixteen layers predict -0.280 ms and the measured target -delta is -0.202 ms. At 32K, the partial kernel is 0.245912 ms and complete wall -is unchanged within 0.001 ms. The final 32K-to-128K complete-round slope is -11.244233 ms, 22.01% below split 40 and 1.76% below plain split 80. The final -target slope is 10.899464 ms, 23.22% below split 40. Target attention still -accounts for 96.93% of the residual complete-round slope. - -A logical-work roofline proxy narrows the remaining operator bottleneck. The -two padded 32-row GQA groups perform about 2.15/8.59 GFLOP per layer at -32K/128K; the measured partial times correspond to only 8.73/9.27 logical -TFLOP/s. Their requested E5M2 K/V payload corresponds to about 136/145 GB/s -before accounting for duplicate-group L2 reuse. Both rates stay nearly flat as -context grows and are far below V100 peak figures. This is not an NCU hardware -counter claim, but it rules out treating the residual as unavoidable saturated -HBM traffic and prioritizes flat-QK/PV utilization, softmax synchronization, -and GQA packing over further page-address arithmetic. - -The final 128K diagnostic request preserves the split-80 accepted-token counts -and 7.647059 mean acceptance length. At 32K, the single-request diagnostic -changes from 6.142857 to 6.095238 (-0.047619), inside the 0.05 diagnostic gate; -dataset-level acceptance and model-quality gates remain authoritative. - -An attempted compact mapping for the two-head GQA tail passed the six operator -correctness tests but was 6.3-7.0% slower than split 80 across the micro sweep. -The run also observed a late external-worker ownership race, so its absolute -timings are not promotion evidence; the uniformly negative direction was -sufficient to revert the experiment completely. Nsight Compute hardware -counters are unavailable to the current user (`ERR_NVGPUCTRPERM`), so no NCU -counter claim is made. - -A second exact-address experiment replaced sixteen redundant page divisions, -modulos, and block-table loads per D=256 row with one half-warp calculation and -broadcast. It preserved register count and passed all six operator correctness -tests, but regressed the one-pass kernel by 3.51%/2.36%/2.47%/2.63% at -1K/32K/128K/256K in a production-page-size A/B/A run. On this Volta path, the -independent address arithmetic is hidden better than the added shuffle -dependency. The candidate was reverted completely. - -## Short-context floor and rejected split-K QK - -The missing 1K TP4 FULL-Graph endpoint now uses the same final binary and -sampling/configuration contract as the 32K and 128K traces: - -| Phase | 1K mean per verification round | -| --- | ---: | -| Complete | 21.634854 ms | -| Target | 16.348011 ms | -| Draft | 3.760873 ms | -| Draft to target | 0.242986 ms | -| Target to draft | 1.282984 ms | - -The grouped partial and combine kernels average 0.057043 and 0.005948 ms per -full-attention layer. Across sixteen layers they contribute about 1.007854 ms -to the target graph. Subtracting this measured service leaves approximately -15.340 ms of context-independent target work and 20.627 ms of complete-round -work. Thus long-attention work explains the decay, but eliminating it entirely -would not by itself satisfy a sub-20-ms short-context goal; the fixed target, -draft, and return path remain a separate optimization problem. - -An exact split-K experiment divided the D=256 QK accumulation between two -warps and joined the FP32 partials before the unchanged softmax. Limiting -unrolling removed all local stack spills (128/126 registers for the one-/two- -pass kernels), and all six grouped-verifier correctness tests passed. A/B/A -medians nevertheless changed by -0.40% at 32K, +0.11% at 128K, and +0.28% at -256K versus the second baseline. There is no stable long-context gain, so the -experiment was reverted and is not part of the performance result. - -## Pack-GQA48 follow-up - -The accepted grouped baseline launches two 256-thread CTAs per context split: -one four-head group and one two-head tail. Both CTAs load and decode the same -single-KV-head page. Following the official FlashAttention PackGQA mapping, -the follow-up packs all six heads and eight verifier tokens into one 48-row, -512-thread CTA. QK and PV remain Volta WMMA, online softmax remains FP32, and -the split outputs retain the same FP16-partial/FP32-LSE reduction contract. -The production partial kernel uses 128 registers, no local stack, and 50,944 -bytes of opt-in dynamic shared memory. - -The first packed prototype retained 128 tokens per split. It improved the -32K/128K/256K operator by approximately 20-21%, but regressed 1K by 19-21% -because only eight large CTAs were launched. Reducing the packed split quantum -to 64 tokens recovered the 1K parallelism, but the single-query utility shape -then crossed its strict FP32-reference mean-error gate. The retained dispatch -therefore keeps 128-token splits for Q1, forces one split for Q2-Q8 at total KV -length at most 128, and uses 64-token splits otherwise. This is a numerical -dispatch boundary, not a sampling or acceptance heuristic. - -The final production-page A/B/A CUDA-Graph medians are: - -| Context | Accepted grouped A/A2 | Pack-GQA48 | Delta versus A/A2 | -| --- | ---: | ---: | ---: | -| 1K | 0.057344 / 0.057344 ms | 0.046080 ms | -19.64% / -19.64% | -| 32K | 0.257024 / 0.257024 ms | 0.205824 ms | -19.92% / -19.92% | -| 128K | 0.939008 / 0.937984 ms | 0.745472 ms | -20.61% / -20.52% | -| 256K | 1.838080 / 1.838080 ms | 1.448960 ms | -21.17% / -21.17% | - -At the ultra-short 128/256-token fixed shapes the packed route costs about -0.002-0.003 ms more than the accepted grouped kernel; at 512 tokens it is -already 10.9-12.5% faster. These sub-millisecond boundary points do not affect -the frozen 1K-or-longer promotion gate, but remain covered by correctness -tests. - -The expanded operator suite covers Q1/Q3/Q4/Q5/Q6/Q8, 128/256/512/1K and -longer boundaries, random physical pages, one-/two-pass equivalence, and CUDA -Graph replay. All thirteen collected cases pass. The maximum absolute error is -unchanged from the accepted grouped baseline at every audited shape; the largest -positive mean-error delta is `1.60e-7`, below the per-shape `5e-7` guard. - -The TP4 FULL-Graph endpoints confirm that the operator gain lands on the full -target critical path without a short-context regression: - -| Context/phase | Accepted grouped | Pack-GQA48 | Delta | -| --- | ---: | ---: | ---: | -| 1K complete | 21.634854 ms | 21.303621 ms | -0.331233 ms (-1.53%) | -| 1K target | 16.348011 ms | 16.043125 ms | -0.304886 ms (-1.86%) | -| 32K complete | 25.237987 ms | 24.390177 ms | -0.847810 ms (-3.36%) | -| 32K target | 19.455442 ms | 18.633571 ms | -0.821871 ms (-4.22%) | -| 128K complete | 36.482220 ms | 33.334558 ms | -3.147662 ms (-8.63%) | -| 128K target | 30.354906 ms | 27.208276 ms | -3.146630 ms (-10.37%) | - -The traced packed partial averages 0.733521 ms per full-attention layer and its -combine averages 0.010906 ms. The diagnostic request keeps the exact 128 output -token IDs, accepted counts `[17, 17, 17, 16, 16, 15, 15]`, and mean acceptance -length 7.647059. Seventeen verification steps produce 128 output tokens, so -the conservative steady-round projection is about 225.9 token/s at 128K, -versus 206.4 token/s for the accepted grouped endpoint. - -The 32K-to-128K complete-round slope falls from 11.244233 to 8.944381 ms -(-20.45%), and the target slope falls from 10.899464 to 8.574705 ms (-21.33%). -The 32K/128K geometric-mean complete round falls from 30.343662 to 28.513782 -ms (-6.03%). Diagnostic mean acceptance changes by +0.030303 at 1K, +0.047619 -at 32K, and exactly zero at 128K, all inside the frozen 0.05 boundary. - -WikiText PPL is exactly equal to the accepted grouped candidate across all -16,376 prompt-logprob positions. Against target only, the maximum PPL delta is -0.003952 and the mean absolute prompt-logprob delta is 0.007519, exactly the -same quality envelope recorded before Pack-GQA. Mixed task-score gates remain -required before promotion. - -## Score and PPL promotion evidence - -The final split-80/page-specialized/tile-balanced binary was evaluated with -probabilistic proposals, target sampling at temperature 1.0/top-p 0.95/top-k -20, fixed seed, a 2,048-token cap, and natural EOS. The mixed dataset supplies -32 GSM8K, 32 HumanEval, and 32 MATH500 cases. Its embedded MBPP prompts have an -incompatible function-name contract and are excluded; the independent -`mbpp-tests-32` dataset supplies the valid MBPP result. -The tested extension SHA256 is -`4a925265cfed934ebd7772be2a1eccaa64163d5324f87d31ba8ee99d10c3a044`. - -| Suite | Target only | Prior grouped | Final candidate | -| --- | ---: | ---: | ---: | -| GSM8K | 29/32 | 30/32 | 30/32 | -| HumanEval | 28/32 | 28/32 | 27/32 | -| MATH500 | 26/32 | 25/32 | 26/32 | -| MBPP, independent tests | 25/32 | 25/32 | 26/32 | -| Valid total | 108/128 | 108/128 | 109/128 | - -The suite-level movement is one HumanEval case down and a net two MATH/MBPP -cases up; the valid aggregate is therefore one case above both controls, not a -quality-regression claim. WikiText prompt-logprob/PPL output is exactly equal -to the prior grouped candidate across all eight samples (zero maximum and mean -difference). Against target only, the maximum PPL difference is 0.003952 and -the mean absolute prompt-logprob difference is 0.007519. - -Acceptance also remains inside the frozen 0.05 gate. On the 128-request mixed -run, request-mean completion tokens per verification step change from 4.238430 -to 4.218840 (-0.019589), and pooled acceptance length changes from 3.749104 to -3.714972 (-0.034132). On independent MBPP they change from 3.527945 to -3.509266 (-0.018679), while pooled acceptance changes from 3.421634 to -3.431016 (+0.009382). The final runs produce 101,312 and 28,070 output tokens, -respectively; 103/128 and 24/32 requests terminate naturally before the cap. - -## Measurement order - -1. Reproduce low-overhead full-model B1 baselines at 1K/32K/128K/256K. -2. Record synchronized target, draft, selector/sampling, communication, and - host residual timing from the same requests. -3. Capture short Nsight Systems graph-node traces at 32K and 128K, aligned by - replay ordinal across all four TP workers. -4. Classify target graph nodes into attention, quantized GEMM, GDN/norm, - LM-head/selector, collectives, and dependency gaps. -5. Build the smallest exact-shape microbenchmark for the confirmed growing - kernel family. Run Nsight Compute only after the runtime route is proved. - -## Promotion gates - -- Reduce geometric-mean 32K/128K complete verification-round wall without a - regression greater than 1% at either length. -- Do not increase 1K complete-round wall by more than 1%. -- Preserve request-mean acceptance length within 0.05 and pooled acceptance - within sampling noise. -- Preserve the existing task-score and WikiText PPL gates before enabling a - numerical or scheduling change by default. -- Validate the 256K boundary for capacity, output health, and latency after a - candidate passes at 128K. - -## Current status - -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. - -The final source audit rebased all five patches without semantic change onto -`main@acc0f6fb92`, after the generic prefix-anchored and E4M3 batch-XQA merges. -The complete changed-file static suite passes and the combined SM70 extension -builds and links. Removing a redundant translation-unit-wide dynamic-shared -alignment attribute preserves the packed verifier layout while restoring all -121 pre-existing XQA kernel SASS instruction streams exactly; none are changed -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. diff --git a/docs/design/sm70_dflash2_ngram_hybrid.md b/docs/design/sm70_dflash2_ngram_hybrid.md deleted file mode 100644 index bb45e80d93..0000000000 --- a/docs/design/sm70_dflash2_ngram_hybrid.md +++ /dev/null @@ -1,130 +0,0 @@ -# SM70 DFlash2 + ngram Hybrid - -## Purpose - -Add an optional prompt-ngram assistant to the MRV2 DFlash2 path. An ngram hit -must skip the DFlash2 query/selector while preserving DFlash2 context-KV state, -and a miss must fall back to the existing DFlash2 implementation unchanged. -The target verifier and the probabilistic rejection-sampling contract remain -authoritative, so the optimization cannot change the target distribution. - -This work is based on `onecat/main` at -`d62ef5cb20b48de93a91562e777ac48985f44b76` and is isolated on -`agent/v100-dflash2-ngram-hybrid-20260825-050018`. - -## Upstream audit - -- vLLM prompt lookup: PRs #12193, #22437, #24986, and #29184 provide the CPU - KMP and GPU-vectorized implementations. MRV2 still has no combined ngram + - model-drafter route. -- SGLang: PRs #17260, #21243, and #22737 show that overlap scheduling requires - complete request-token state and explicit accepted-token indexing. A stale - host output list is not a valid lookup source. -- llama.cpp: its comma-separated speculative configuration gives draftless - ngram proposers priority, falls back to DFlash, and still calls `process()` on - every proposer so model-drafter state remains synchronized. This is the - closest reference architecture for the first implementation here. -- Arctic Suffix Decoding and SAM-Decoding are useful follow-ups, but their - confidence policy, external dependency, and longer trees are intentionally - outside this first block-8 implementation. - -## Initial contract - -- The feature is opt-in and only valid for `method=dflash` with a DFlash2 - checkpoint. Standalone `ngram`, Eagle, MTP, DFlash1, and `dflash_ddtree` - routing is unchanged. -- Lookup reads the authoritative MRV2 request-token state and supports normal - synchronous and overlap scheduling without rebuilding history in Python. -- Structured-output requests bypass the ngram assistant until grammar-aware - proposal masking is proved correct. Tool/reasoning parsers without a grammar - are unaffected. -- A hit returns at most the configured DFlash2 draft width (seven tokens in the - official block-8 setup). The normal DFlash2 proposer handles misses. -- Context K/V materialization always runs. Only the DFlash2 query, candidate - projection, selector, and selector walk may be skipped. -- Greedy and probabilistic modes are supported. In probabilistic mode the - ngram proposal is represented as a one-hot draft distribution inside the - existing sparse target-rejection interface, preserving target sampling. -- Prefix-cache hits may rebuild missing DFlash draft K/V as before; this change - must not corrupt either target or draft cache state. - -## Test plan - -1. Unit tests: configuration/routing isolation; KMP hit, miss, truncation, and - overlap cases; per-request mixed-source state; one-hot sparse draft support; - grammar bypass; prefix/state transitions. -2. CPU microbenchmarks: lookup at 1K, 32K, 128K, and 256K contexts. -3. V100 correctness: eager before CUDA graph; batch 1/2/4; hit/miss/mixed; - compare proposed tokens, accepted trajectory, target output, and draft K/V. -4. V100 performance: report lookup, context-KV, query/selector, target verify, - full round, emitted tokens per round, and pure-decode tokens/s for baseline - DFlash2 and DFlash2+ngram at short and long contexts. -5. Quality: compare target-only, DFlash2, and hybrid on coding/general/tool - datasets under the same sampling configuration. Report scores, completion - counts, wall time, hit rate, conditional acceptance length, and failures. - -## Promotion gates - -- No statistically meaningful quality regression against target-only or the - existing DFlash2 route. -- DFlash2-miss acceptance trajectory remains unchanged in deterministic tests. -- Ngram hits actually skip query/selector work in a trace. -- Hybrid pure-decode throughput is non-regressing at every measured context; - otherwise the feature remains opt-in while the losing shape is investigated. - -## Status - -- 2026-08-25: Draft PR #287 implements opt-in MRV2 host lookup over the - authoritative UVA request-token state, full-hit query/selector bypass, - all-hit batch application, and one-hot dense/sparse rejection caches. - Structured-output batches bypass the assistant, and intermediate chunked - prefill materializes draft context K/V before returning without lookup. -- CPU split-history KMP microbenchmark (median/P95): 1K `2.66/2.74 us`, 32K - `50.05/53.22 us`, 128K `190.72/201.03 us`, and 256K `369.57/393.17 us`. - Focused validation reports `93 passed, 9 skipped` for the CPU DFlash2/ngram - set, `10 passed` for MRV2 routing, and `24 passed` for the V100 ngram/AOT - fullgraph set. -- PR #288 first restored the missing production optimization closure to main. - On that restored source, no-assist practical coding measured `18.660 ms` per - complete round, `135.70 tok/s`, and acceptance length `2.532`. Repeated runs - of the historical MBPP item 28 stabilized at `18.484--18.616 ms`, - `224.05--225.66 tok/s`, and acceptance length `4.184`. Older 27--31 ms - branch measurements predate this closure and are invalid baselines. -- With ngram `[5,5]` enabled under the identical TP4 practical contract, the - coding request measured `18.980 ms`, `131.14 tok/s`, and acceptance length - `2.494`; only about 2.2% of eligible rounds were full hits. The matched MBPP - item measured `18.700 ms`, `217.18 tok/s`, and acceptance length `4.071`. - Thus low-hit short requests currently pay roughly `0.1--0.3 ms` per round - and do not pass the default-enable performance gate. -- The 32-case MBPP natural-stop run (16K output cap, fixed per-item seeds) gave - the pure-DFlash control `65,934` output tokens in `389.542 s`, `19,767` - verification rounds, `19.707 ms` wall time per round, and acceptance length - `3.336`. Hybrid produced `53,710` tokens in `327.613 s`, `16,700` rounds, - `19.618 ms` per round, and acceptance length `3.216`. Cumulative full-hit - rate reached about 10.1% and lookup averaged `0.016--0.018 ms`; skipped - queries slightly reduced round cost, but the lower sampled acceptance path - left raw aggregate output throughput at `163.94` versus `169.26 tok/s`. -- EvalPlus on the mapped 31 MBPP cases reports pure DFlash Base `30/31` and - Plus `28/31`; hybrid reports Base `31/31` and Plus `28/31`. This is no score - regression, but it does not override the failed throughput gate. The - probabilistic one-hot proposal preserves the target distribution; it is not - expected to reproduce the same sample-by-sample random trajectory. -- A fresh hybrid 32K chunked-prefill request measured `10.688 s` cold and - `1.416 s` on an identical-prefix hit, versus the restored no-assist evidence - of `10.58--10.65 s` and `1.405 s`. The context-only skip log was present and - no cache corruption occurred. -- Promotion decision: keep `ngram_assist` opt-in. The merge audit removes the - mixed-hit hazard: an n-gram draft is now applied only when every active - request has a full-width hit and the complete DFlash2 query/selector can be - skipped. Mixed-hit batches retain the unchanged DFlash2 proposals because - overriding a row after paying the full query cost cannot improve latency and - can alter sampled acceptance. A future default-on policy still needs matched - evidence that full-query skips repay lookup/probe overhead; blindly reducing - the ngram length is not justified by the current data. -- After the paired run was stopped, fresh graph captures on GPUs 4--7 began - failing inside the draft paged-attention capture with - `cudaErrorStreamCaptureInvalidated`. A clean pre-ngram main worktree failed - at the same point, so this is not an ngram source regression. The NVLink - topology requires resetting all eight GPUs, which was deliberately not done - while the user-facing API remained live on GPUs 0--3. Do not count failed - startups as performance samples. diff --git a/docs/design/sm70_dflash2_nvfp4_17ms.md b/docs/design/sm70_dflash2_nvfp4_17ms.md deleted file mode 100644 index f29e7abc9d..0000000000 --- a/docs/design/sm70_dflash2_nvfp4_17ms.md +++ /dev/null @@ -1,220 +0,0 @@ -# SM70 Qwen3.8 NVFP4 DFlash2 17 ms Campaign - -> **Status:** planning record only. No source optimization or 17 ms result is -> claimed by this document. The historical baseline below is retained to make -> the next implementation reproducible. - -## Scope and frozen baseline - -This campaign targets the complete batch-one DFlash2 speculative round measured -with a Qwen3.8-27B NVFP4 workload on four V100 GPUs. Its historical baseline is -`onecat/main@34403018d917054dd7765d5e820ad29c8d342348`; the planning record was -integrated on `main@cadcf1d899b6d7511f815e7ee939b1e4676aff19`. It does not -credit route-hit logs as performance evidence. - -The retained benchmark artifacts identify a mixed NVFP4/channel-FP8 target -with a BF16 LM head and a DFlash2 draft at revision -`dedf8df68adfb1afeaf7b7480c0a0243108177b4`. Private filesystem locations are -intentionally not recorded. The frozen workload uses TP4 on an owned four-V100 -group, probabilistic DFlash2 with seven draft tokens, selector K=16, target -top-k=20, FP8 E5M2 target KV, FP16 draft KV, Flash-V100 for target and draft -attention, and FULL target and draft CUDA Graphs. Model/checkpoint labels are -benchmark evidence only and must not activate an optimized runtime route. - -Two runtime contracts are deliberately kept separate: - -1. The localization contract uses a 32K maximum model length, 512 maximum - batched tokens, and disables prefix caching, Mamba alignment, and parsers. - It is used only for short unprofiled repeats, graph-node traces, and - exact-shape microbenchmarks. -2. The promotion contract uses a 256K maximum model length, 4096 maximum - batched tokens, prefix caching, Mamba alignment, and the Qwen tool and - reasoning parsers. A speedup is not accepted until it reproduces under this - practical configuration. - -The current accepted practical baseline on PR #288 is `18.465--18.537 ms` per -complete round for 512-token coding requests and `18.587--18.603 ms` for -1,024-token coding requests. The high-acceptance MBPP item 28 baseline is -`18.567 ms`, acceptance length `4.686`, and `251.60 token/s`, with a natural -EOS and EvalPlus base/plus `1/1`. A fresh current-source 16-prompt GSM8K run -measures request-mean acceptance `4.45732`, pooled acceptance `4.07740`, and -`19.3688 ms` per round with diagnostic counters enabled. Diagnostic timing is -not a production baseline. - -## Acceptance and runtime-binary audit - -The apparent NVFP4 acceptance regression was reproduced on the frozen 64-row -GSM8K order, but it is not caused by NVFP4. With the same current source and -the same official stochastic contract, FP8 measured request-mean/pooled -acceptance `4.49887/3.89772` and 60/64 accuracy; NVFP4 measured -`4.53642/3.79374` and 59/64. The request mean is slightly higher for NVFP4. -Its lower pooled value comes from a longer generated trajectory and is not a -valid quantization comparison by itself. The corresponding request-mean -completion tokens per verification round are `4.49761` and `4.53214`; this -closely related metric was previously mislabeled as request-mean acceptance. -All future records report both fields explicitly. - -Both current paths are nevertheless below the repaired August 22 FP8 baseline -of `5.38880/4.56444` (`5.38884` request-mean completion tokens per round). -This is a real regression rather than only output-length mix: current FP8 is -lower on 61 of the 64 matched prompts, with paired -request-mean delta `-0.89123` and bootstrap 95% interval -`[-1.02152,-0.75597]`. Restricting each run to outputs of at most 512 tokens -still moves the mean from `5.60920` to `4.63777`. - -Single-variable runs exclude the target quantized operators and compact -sampler as the main cause. Disabling NVFP4 QPN2 on the fixed 16 rows measured -`4.37420` request mean instead of recovering the current `4.57520` control. -Disabling FP8 target QPN8 on all 64 rows measured `4.49095/3.80785`, and -disabling sparse target rejection measured `4.51820/4.02970`; none approaches -the repaired baseline. All completed 64-row variants retained 60/64 FP8 -accuracy. - -The causal fault is binary provenance. The low-acceptance runs loaded -`_C_stable_libtorch` SHA256 `fe986a...9ba8c`, even though the checked-out source -contains the batched-weight RMSNorm repair. That binary predates the per-layer -weight implementation: its focused GPU suite fails 10 of 13 cases, silently -uses weight row zero for the other draft layers, and mismatches 81--96% of -elements on multi-layer fixtures. The repaired binary SHA256 -`7689e5...d449` passes all 13 cases bitwise. The low result also numerically -matches the historical pre-repair `4.5395/3.9725` baseline. - -The source now verifies the loaded stable extension once with deterministic -nonzero per-layer inputs. A conforming binary continues to use the single -grouped kernel after warmup. A stale binary emits an explicit warning and uses -the bitwise per-layer fallback, so a wheel/source or worktree/symlink mismatch -cannot silently lower acceptance again. The new fallback tests pass for both -runtime behaviors; the full CPU DFlash2 suite reports 70 passed and 12 skipped. - -The repaired-binary FP8 closure uses the same frozen 64-row GSM8K contract and -the dense selector baseline. It records request-mean acceptance `5.37699`, -pooled acceptance `4.62217`, and request-mean completion tokens per round -`5.37814`, with 61/64 correct and 62/64 natural stops. Against the repaired -August 22 baseline, paired request-mean delta is `-0.01180` with bootstrap 95% -interval `[-0.08421,0.06282]`; every position-wise pooled acceptance value is -higher. The artifact records stable-extension SHA256 `7689e5...d449`. This -closes the large acceptance regression: the stale-binary run was lower by -`0.87812` request-mean tokens on the same rows, with bootstrap 95% interval -`[0.75099,1.00344]` after repair. Its concurrently collected throughput is not -a standalone performance baseline because NVFP4 compilation occupied the -other TP4 group. - -The matching repaired-binary NVFP4 run records request-mean/pooled acceptance -`5.39851/4.57051`, request-mean completion tokens per round `5.40178`, and -59/64 correct with 62/64 natural stops. Its request mean is `+0.02152` versus -the repaired current-source FP8 run; the paired 95% interval -`[-0.10021,0.14795]` does not distinguish the two quantizations. The lower -NVFP4 pooled value is explained by its longer output trajectory (35,366 versus -32,658 tokens), not lower per-request acceptance. Relative to the stale NVFP4 -run, the repair adds `0.86209` request-mean tokens with 95% interval -`[0.74547,0.98007]` and improves 62 of 64 matched rows. - -The experimental draft-MLP QPN8 route is rejected independently. On the same -fixed 16 requests it changed request mean `4.57520 -> 4.51123` and pooled -acceptance `4.39123 -> 4.09891`; the request delta `-0.06397` exceeds the -`-0.05` gate. Its code and benchmark hook were removed rather than retained as -a nominal 17 ms result. - -The corrected 16K coding-quality pair uses the same NVFP4 target, MBPP-32 -prompts, request seed, graph mode, E5M2 KV, and natural-EOS sampling on both -sides. The dataset tests score no-DFlash `31/32` and DFlash2 `32/32`. -EvalPlus maps 31 of those rows: base is `30/31` versus `31/31`, and plus is -`27/31` versus `27/31`. At the plus level, each route alone passes one -different row; exact paired McNemar `p=1.0`. Both routes naturally stop on 31 -rows and reach the 16K cap on one row. This provides no evidence of a DFlash2 -quality regression under the repaired runtime. - -The first score incorrectly reported one DFlash result as invalid syntax even -though the complete executable answer followed ``. The extractor had -searched illustrative reasoning-era code fences before checking an unfenced -final answer. The scorer now gives the final answer precedence, its regression -fixture parses the affected `comb_sort` result, and both routes were rescored -with the same corrected extractor. The corrected DFlash results are dataset -tests `32/32`, EvalPlus base `31/31`, and EvalPlus plus `27/31`. - -The opt-in selector-QPN8 candidate also passes the fixed 64-row acceptance -gate. Relative to the repaired dense selector, request-mean acceptance changes -`5.39851 -> 5.43951`, pooled acceptance `4.57051 -> 4.58679`, and GSM8K -`59/64 -> 60/64`. The paired request delta is `+0.04101` with bootstrap 95% -interval `[-0.02309,0.10669]`; 23 rows improve, 15 fall, and 26 tie. It -therefore has no acceptance-loss signal. - -Its first clean directional speed signal is positive but not promotable: the -dense run on physical GPUs 4--7 measures request-mean/median complete round -`19.083/19.046 ms`, while QPN8 on GPUs 0--3 measures `18.077/18.029 ms`. -Aggregate output is `218.58 -> 229.41 token/s` and mean steady decode is -`282.68 -> 300.12 token/s`. Because these are different physical groups, they -are not a valid paired speed result. QPN8 also lowers available KV capacity -from 529,616 to 459,817 tokens through its prepared rerank layout. An unrelated -eight-GPU job occupied both groups before the same-GPU repeat; the candidate -remains explicit opt-in and is not a default or a credited 17 ms result. - -## Acceptance gates - -A short-context candidate is accepted only when all of the following hold: - -- the no-diagnostic promotion contract measures a mean and median complete - speculative round at or below `17.0 ms` across at least three independent - steady-state requests; the p90 must be at or below `17.5 ms`; -- the same fixed prompt, seed, sampling parameters, output cap, CUDA Graph - shapes, and GPU set are used for the baseline/candidate pair; -- paired request-mean acceptance does not fall by more than `0.05`, the - per-position counters remain healthy, and no stale-buffer, graph-replay, or - draft-KV mismatch appears; -- the existing GSM8K, MATH-500, HumanEval, corrected MBPP, and WikiText PPL - gates do not regress versus target-only and the accepted DFlash2 baseline; -- every retained performance optimization has an exact engine-contract - admission predicate, a rollback switch, focused numerical coverage, and an - unprofiled endpoint win. Once these gates pass, the optimized path is the - default; rollback remains for diagnosis rather than permanent opt-in. - -Long-context work is evaluated at 32K, 128K, and 256K. The candidate must not -make complete-round latency or pure-decode throughput worse by more than 1% -at any length. Numerical probes require finite outputs and a dtype-appropriate -absolute/relative error envelope; official-sampling quality and acceptance -must remain non-regressing. Bitwise output or greedy-token identity is not a -promotion requirement when a changed but valid reduction order explains the -difference. The 17 ms target applies to the short round; long-context results -are reported as an explicit decay curve because the verifier attention cost -necessarily grows with resident context. - -## Trace-first implementation sequence - -1. Reproduce the current no-diagnostic NVFP4 baseline in this worktree and - record source SHA, extension hashes, route logs, GPU clocks, and per-request - round statistics. -2. Capture the smallest generation-only Nsight Systems trace that contains - steady CUDA Graph replays. Split every round into draft, draft-to-target, - target verification, and target-to-draft phases, then expand the target - phase into exact kernel categories and launch counts. -3. Microbenchmark only the largest current-source bucket at its production - M/N/K, TP, dtype, and graph shapes. A microbenchmark result is directional - evidence, never an endpoint claim. -4. Implement the smallest exact change that removes the measured bottleneck. - Rejected DDTree index reuse, skipped verification, generic auxiliary-stream - overlap, variable K, and approximate reranking are not retried: they either - changed output/acceptance or increased the complete round. -5. Promote a default only after focused CPU/GPU numerical tests, graph replay, - three clean short repeats, the practical 256K configuration, the quality - suite, and the 32K/128K/256K decay sweep pass. - -The first trace should determine whether the remaining roughly `1.6 ms` is in -NVFP4 QPN2/QPN8 projections, draft non-causal KV work, target verification, -or scheduler/metadata boundaries. No saving is assigned to a hypothesis -before that trace and its exact-shape microbenchmark exist. - -## Draft PR record - -Purpose: reduce the complete Qwen3.8-27B NVFP4 DFlash2 round to 17 ms on SM70 -without changing the sampling contract, acceptance, or task quality. - -Test plan: unprofiled paired endpoint repeats; Nsight graph-node phase -breakdown; exact-shape microbenchmarks; focused numerical and CUDA Graph replay -tests; GSM8K/MATH-500/HumanEval/MBPP/PPL quality gates; 32K/128K/256K decay -sweep. - -Test result: the runtime-provenance acceptance repair and paired MBPP quality -gate pass. Selector-QPN8 passes the acceptance gate and has a cross-group -directional speed signal, but remains opt-in pending a same-GPU endpoint pair. -The 17 ms endpoint and long-context performance gates remain pending; no -rejected or unpaired speed candidate is credited. diff --git a/docs/design/sm70_dflash2_target_graph_20ms.md b/docs/design/sm70_dflash2_target_graph_20ms.md deleted file mode 100644 index 3e87b8aac3..0000000000 --- a/docs/design/sm70_dflash2_target_graph_20ms.md +++ /dev/null @@ -1,1229 +0,0 @@ -# SM70 DFlash2 target-verifier graph: 20 ms control log - -## Scope and frozen baseline - -- Date: 2026-08-23. -- Integration base: `c62c1dcd833458054daa12e134695aa19a4ac609` - (`codex/v100-dflash2-gdn-metadata-20260822-122723`, Draft PR #257). -- Worktree: - `/home/ymzx/桌面/1cat-vllm/worktrees/v100-dflash2-target-graph-20ms-20260823-084427`. -- Branch: `codex/v100-dflash2-target-graph-20ms-20260823-084427`. -- Task cache: `/data/minimax-h3/task-cache/v100-dflash2-target-graph-20ms`. -- Frozen workload: Qwen3.8-27B-FP8 target, official DFlash2 draft, TP4, - batch one, block eight (seven draft tokens), Flash-V100, target E5M2 KV, - draft FP16/auto KV, and CUDA Graph. -- Baseline node trace: - `/data/minimax-h3/task-cache/v100-dflash2-pr257-gdn-metadata/profiles/latest-e2e/dflash2-fused-viewcache-b1-nodes-o128-v5.sqlite`. - -The 30-round, four-rank steady-state baseline (one edge round removed from -each side) is: - -| Phase | Critical-path wall | Nodes or launches per rank | -|---|---:|---:| -| Draft graph | 4.046 ms | - | -| Draft to target | 4.934 ms | 153 kernels | -| Target graph | **24.740 ms** | **2612 nodes** | -| Target to draft | 2.838 ms | - | -| Complete round | 36.300 ms | - | - -The target-graph objective is at most 20.000 ms under this same diagnostic. -Performance acceptance requires an improvement in the complete round. The -final quality contract uses probabilistic task scores, paired acceptance, and -PPL against target-only; greedy token identity is diagnostic evidence rather -than a production gate. The superseding measurements are recorded at the end -of this document. - -## Resource-use decomposition - -The target graph contains 23.365 ms of rank-average GPU kernel service inside -a 24.740 ms critical span, or 94.7% activity-envelope coverage. This rules out -a large idle bubble as the primary cause, but it does not imply efficient SM -or memory-pipeline utilization within each small kernel. Nsight Compute metrics -must remain separate from graph-span and Nsight Systems service time. - -| Kernel category | Rank-average service | Launches per rank | -|---|---:|---:| -| TurboMind FP8 dense GEMM | 10.831 ms | 256 | -| Copy/cast elementwise | 3.387 ms | 941 | -| Other kernels | 2.869 ms | 240 | -| TP all-reduce/communication | 2.316 ms | 128 | -| Other Torch/Triton elementwise | 2.494 ms | 739 | -| RMSNorm/residual | 0.616 ms | 128 | -| LM-head/sample/TP gather | 0.416 ms | 48 | -| Dense GEMV/GEMM/compressor | 0.257 ms | 65 | -| Fill/mask | 0.179 ms | 67 | - -The graph is therefore busy but fragmented: 1680 copy/cast/elementwise nodes -consume 5.881 ms, while the average node is only about 8.9 microseconds. The -20 ms target requires recovering at least 4.740 ms (19.2%); launch reduction -and work fusion are first-order requirements, not optional cleanup. - -## Ordered experiments and stop gates - -1. A/B the repaired, default-off DFlash2 packed GDN verifier with all other - flags frozen. It preserves the current gating values, recurrent arithmetic - order, and FP32-state contract while removing packed-QKV rearrangement and - the final output copy. Record route-hit logs, graph wall, graph nodes, - complete-round wall, exact output tokens, and acceptance trajectory. -2. If the packed route wins but remains above 20 ms, use its new node trace to - isolate residual GDN state/copy traffic. Compare with SGLang-V100's fused - recurrent verifier and all-layer state commit without importing unrelated - scheduler or Eagle/MTP changes. -3. Profile the exact TP4 all-reduce shape. The current 2.316 ms/rank service is - a second independent target; any replacement must preserve collective - ordering and exact target hidden states. -4. Only after the verifier contract is stable, assess a DFlash2-safe variant of - v100-skinny's small-M QPN8 projection work. It is not accepted on the basis - of its non-speculative result. - -Every failed experiment must be reverted or kept behind a default-off gate. -Do not report profiler-instrumented latency as unprofiled throughput. Do not -use target-only runs, and do not occupy a partial TP4 group or terminate an -unrelated process. - -## External references - -- SGLang-V100 fixed source and trace are compared locally at - `/data/minimax-h3/task-cache/v100-dflash2-pr257-gdn-metadata/sources/sglang-v100` - and - `/data/models/v100-dflash2-20260820/sglang-audit/perf-rootcause/sglang-dflash2-single1-step20-v2.nsys-rep`. -- v100-skinny is pinned at `5b589c0dc81223e0ba65bcb3e755874723f8b515`; - its 219.1 tok/s result is a Qwen3.8 mixed-NVFP4/FP8 MTP result, not a DFlash2 - target-verifier baseline. - -## Results - -### Accepted target-graph reduction - -The matched TP4 node trace at source `500130882ba68fb7bac3a0b9e4eb5872647fb045` -uses the same model, draft, request, sampling, KV dtypes, and graph policy as the -frozen baseline: - -- SQLite: - `/data/minimax-h3/task-cache/v100-dflash2-target-graph-20ms/profiles/dflash2-fused-gdn-norm-split-gemma-dynamic-b1-nodes-o128-gpu0123-v3.sqlite`. -- Result JSON: - `/data/minimax-h3/task-cache/v100-dflash2-target-graph-20ms/results/dflash2-fused-gdn-norm-split-gemma-dynamic-b1-nodes-o128-gpu0123-v3.json`. -- Four-rank synchronized target graph, with the first and last rounds removed: - **19.317 ms mean, 19.322 ms p50, 19.444 ms p90, 19.477 ms p99**, and - **1,257 nodes/rank**. -- Rank-average kernel service is 18.552 ms, or 96.04% of the critical span. -- Relative to the frozen 24.740 ms / 2,612-node baseline this removes - **5.423 ms (21.92%)** and **1,355 nodes (51.88%)**. The p99 is also below - the 20 ms objective. -- The benchmark endpoint inside this profiled run reports 125.060 steady - decode token/s; this is trajectory evidence, not an unprofiled throughput - claim. All 128 output - token IDs and hash `fe0300...` match the baseline exactly. Probabilistic - acceptance changed slightly from 4.129 (31 draft rounds) to 4.000 (32 draft - rounds), so the Gemma suffix fusion remains a Type-B candidate pending a - distribution-quality gate rather than being accepted from one trajectory. - -The reduction is cumulative: the one-pass GDN output norm first reached -22.646 ms / 2,084 nodes; fused GDN split materialization reached 21.625 ms / -1,892 nodes; the dynamic-shape-correct Gemma residual RMS fusion reached the -19.317 ms / 1,257-node result above. A prior `v2` launch omitted the installed -`flash_attn_v100_cuda` extension directory from `PYTHONPATH` and failed during -graph capture; it is not a performance result. - -### Remaining graph-boundary cost - -Across 31 steady boundaries per rank, the draft-graph end to target-graph -start interval is 5.776 ms mean (5.582 ms p50). Only 1.304 ms is covered by -non-graph GPU kernels, leaving about 4.47 ms of launch, synchronization, and -CPU queueing bubbles. Every boundary submits 153 kernels: - -- one TP `cross_device_reduce_1stage`: 0.785 ms; -- 20 `DeviceScanKernel`, 20 `DeviceScanInitKernel`, 20 `compute_cuda_kernel`, - and 21 `indexSelectSmallIndex` calls: 0.276 ms service but much larger - serialized submission cost; -- remaining input, slot, block-table, GDN metadata, and elementwise kernels: - about 0.243 ms service. - -The scan/select work arrives in five serial groups with scalar D2H ordering -points. Therefore the next boundary objective is not a faster individual -3-microsecond kernel. It is to construct persistent target metadata once and -capture or batch-submit the five repeated groups while preserving every -buffer/update dependency. - -The same parser applied to the retained SGLang-V100 trace gives a useful -matched structural target. Its draft-to-target interval is 3.398 ms mean with -30 non-graph kernels and 0.115 ms GPU service, versus this branch's 5.776 ms, -153 kernels, and 1.304 ms service. SGLang still pays a 2.715 ms host -`cudaStreamSynchronize`, so it is not a zero-overhead endpoint. The immediate -vLLM gap is nevertheless concrete: remove 123 repeated launches and move the -roughly 0.785 ms TP reduction adjacent to target input/embedding preparation -into the target replay dependency chain. This comparison comes from -`/data/models/v100-dflash2-20260820/sglang-audit/perf-rootcause/sglang-dflash2-single1-step20-v2.sqlite`. - -### Preliminary probabilistic quality pair - -A fixed 16-question sequential GSM8K pair used graph TP4, target E5M2 KV, -draft FP16 KV, block eight, and official `temperature=1.0/top_p=.95/top_k=20` -sampling. The control artifact is -`quality/gsm8k-16-control-gemma-off.json`; the candidate artifact is -`quality/gsm8k-16-candidate-gemma-on.json` under the task cache. - -- Control: 68.75% accuracy, aggregate acceptance length 4.470. -- Candidate: 75.00% accuracy, aggregate acceptance length 4.693. -- Eight of 16 full output-token trajectories match. Eight diverge under - probabilistic sampling, as expected from the accepted one-FP16-ULP numeric - bound; this small sample cannot establish a quality improvement. - -The Gemma fusion stays default-off until paired prompt perplexity and broader -dataset gates show no regression. The GDN norm and split stages retain their -Type-A classification and can be defaulted independently of that decision. - -The first deterministic distribution probe scores 1,850 prompt tokens from 32 -fixed GSM8K questions in eager target prefill (Graph is irrelevant to prompt -logprobs). Weighted perplexity is 4.20119 for the decomposed control and -4.20172 for the Gemma fusion, a +0.00053 / +0.013% change. All 32 next-token -argmaxes match. Mean absolute prompt-logprob difference is 0.00173 and the -worst token is 0.02524. This is small enough to continue dataset testing but -confirms that the kernel is Type B, so it remains default-off. - -### Fused Flash-V100 small-query metadata candidate - -The five repeated scan groups are generated by -`FlashAttnV100MetadataBuilder._update_smallq_decode_metadata`: each of five -target KV groups performs four `repeat_interleave` scans, then materializes and -copies block-table and sequence-length temporaries. The -`VLLM_SM70_DFLASH2_FUSED_SMALLQ_METADATA` candidate writes the persistent -block-table, sequence-length, and query-boundary buffers directly with one -Triton launch per group. It is limited to a DFlash2 target on SM70; draft, -DDTree, Eagle, MTP, and CPU metadata retain the existing path. - -- Exact V100 tests pass for B1 and mixed B3/B4 layouts, negative block IDs, - zero-length padded requests, and graph-token padding. -- The unchanged CPU persistent-buffer, padding, and overflow tests pass. -- A realistic `q=8`, three block columns, five-group microbenchmark reports - 1.278 ms legacy versus 0.114 ms fused wall time and 1.235 versus 0.100 ms GPU - service: **1.164 ms saved / 11.2x** for this isolated constructor. - -The matched TP4 end-to-end trace accepts the candidate as a Type-A default-on -optimization (explicit environment value `0` remains an opt-out): - -- Synchronized draft-to-target falls from 5.720 ms to 1.911 ms (-66.6%). -- Per-rank non-graph work falls from 153 nodes / 1.304 ms GPU service to 23 - nodes / 0.195 ms. All 20 `DeviceScan`, 20 `DeviceScanInit`, and 20 generic - scan-compute launches disappear; `indexSelectSmallIndex` falls from 21 to 1. -- The five replacement launches total 0.016 ms of GPU service. -- The full speculative round falls from 31.791 ms to 27.951 ms (-12.1%), while - the target graph remains effectively unchanged at 19.261 ms / 1,257 nodes. -- Profiled single-request steady decode rises from 125.06 to 142.37 token/s - (+13.8%). The 128-token output hash, aggregate acceptance length 4.0, and - per-position accepted counts `[28, 23, 18, 14, 8, 3, 2]` match exactly. - -The accepted trace is -`profiles/dflash2-fused-gdn-norm-split-gemma-smallqmeta-b1-nodes-o128-gpu0123-v1.sqlite` -under the task cache. Its 1.911 ms synchronized boundary is also below the -3.398 ms boundary in the retained SGLang-V100 audit trace, though that external -trace is a structural reference rather than a fully matched throughput run. - -### Next boundary: target-to-draft - -The accepted trace leaves a stable 2.738 ms target-to-next-draft interval on -rank 0 (33 steady intervals). A representative `M=8` interval and the aggregate -kernel service separate it as follows: - -| Stage | Representative wall / service | -| --- | ---: | -| Dense target LM head, local FP16 shard | 0.947 ms | -| TP full-vocabulary all-gather | 0.139 ms | -| Full-vocabulary top-k/top-p | 0.606 ms | -| Rejection statistics, rejection, and resample | 0.225 ms | -| Target-hidden-to-draft KV precompute | about 0.60 ms | -| Post-update and next-draft metadata | about 0.22 ms | - -The next candidate must therefore attack the first four rows as one semantic -unit. For the fixed no-penalty `top_k=20, top_p=0.95` contract, each TP rank can -compute its exact local top 20, exchange only 20 `(score, token-id)` pairs per -row, merge the global top 20, and perform rejection/residual sampling on that -sparse support. The repository already contains the SM70 TurboMind FP16 -LM-head top-20 epilogue, but the `_C` binary linked into this worktree predates -that op. The first TP4 microbenchmark was therefore rejected at route checking -without recording a timing; an incremental SM70 build is required before this -candidate can pass the microbenchmark gate. - -### Compact target sampling microbenchmarks and gated integration - -The production-shaped TP4 communication probe uses eight verifier rows, -vocabulary 248,320, target top-K 20, and four V100s. It compares the existing -full-vocabulary all-gather plus global top-K with local top-K followed by two -compact all-gathers and a global merge: - -- Exact top-K token IDs and values match. -- The full path moves 993,280 bytes per rank; the two-gather compact path moves - 1,920 bytes, a 517x reduction. -- Despite that byte reduction, p50 is 0.3318 ms for full gather plus top-K and - 0.3292 ms for the compact path. Local top-K (0.1208 ms), two collectives - (0.0911 ms), and compact merge (0.0896 ms) consume the saved transport time. -- Packing values and IDs into one collective is slower at 0.3538 ms. - -Compact communication alone is therefore rejected as a performance feature. -The retained artifact is -`results/tp4-compact-verifier-logits-m8-v248320-k20.json`. - -The second probe fuses compact target top-p, DFlash2 acceptance, log-domain -`relu(p-q)` recovery, and token-keyed Gumbel resampling into one Triton program -per request. At the real B1/block-eight shape it is token/count exact against -the dense rejection path for block 4/8 and `top_p` 1.0/0.95. Its p50 is: - -| Isolated stage | p50 | -| --- | ---: | -| Dense top-K/top-p | 0.8366 ms | -| Dense rejection only | 0.2202 ms | -| Dense top-K/top-p plus rejection | 0.8663 ms | -| Compact top-p plus sparse rejection | **0.0901 ms** | - -Combining this result with the measured TP4 compact-candidate transport gives -an expected real saving of roughly 0.5-0.6 ms per complete round. The artifact -is `results/dflash2-sparse-rejection-b8-v248320-k20-q16-v1.json`. - -The candidate is integrated behind default-off -`VLLM_SM70_DFLASH2_SPARSE_TARGET_REJECTION=1`. Its route is limited to SM70, -MRV2 DFlash2, one active decode request, probabilistic proposals, target -`top_k=20`, positive temperature, `0 < top_p <= 1`, and no grammar, min-p, -penalty, logit bias, bad words, NaN reporting, or logprob request. Every other -configuration computes the unchanged dense logits and uses the shared -rejection sampler. DFlash2 now keeps the 16 proposal IDs and FP32 realized -scores in persistent request-slot order in addition to the dense fallback -cache; a V100 poisoned/reordered-slot test verifies the mapping. - -Validation completed before the full-model run: - -- DFlash2 CPU suite: 46 passed, 11 CUDA-only skipped. -- Strict route and fallback tests: 13 passed. -- Sparse dense-equivalence plus selector/cache CUDA tests: 6 passed. -- Ruff, formatting, compileall, and `git diff --check`: passed. - -The matched single-request Graph/Nsight run remains pending because both TP4 -GPU groups became occupied by unrelated work after these tests. No external -process was interrupted. - -### Residual target-graph service after the 27.951 ms round - -The accepted 1,257-node target graph has 18.357 ms of rank-0 GPU service inside -its 19.261 ms critical span. Its largest exact buckets per replay are: - -| Target graph bucket | Nodes | GPU service | -| --- | ---: | ---: | -| TurboMind FP8 GEMM | 256 | 10.730 ms | -| vLLM one-stage TP4 all-reduce | 128 | 2.313 ms | -| recurrent GDN verifier | 48 | 1.213 ms | -| Flash-V100 partition/reduce | 32 | 0.802 ms | -| fused Gemma residual/RMS suffix | 127 | 0.643 ms | -| generic elementwise | 144 | 0.541 ms | - -The FP8 GEMMs split into three stable launch shapes: 64 calls at 4.394 ms, -128 calls at 3.924 ms, and 64 calls at 2.411 ms. They already match the retained -SGLang target GEMM total closely, so replacing them is a separate QPN8 quality -and throughput project rather than the next low-risk leaf. - -The TP4 collective remains the clearest next `>1 ms` gap. vLLM launches ten -512-thread CTAs and averages 18.07 microseconds per call. The matched -SGLang-V100 trace uses its JIT one-shot push kernel with eighty 128-thread CTAs -and averages 12.36 microseconds. Across 128 calls this structural difference is -about 0.73 ms of service; the broader matched audit found a 0.7-1.2 ms -per-graph TP reduction gap depending on critical-rank accounting. The next -microbenchmark should compare the exact `[8,5120]` FP16 target shape under one -captured 128-call chain, first sweeping the existing vLLM CTA/thread policy, -then porting SGLang's two-epoch push buffer only if the policy sweep cannot -close the gap. Bitwise FP32 accumulation order, CUDA-Graph registration, and -all 128 collective dependencies are hard gates. - -### Sparse rejection end-to-end acceptance - -The matched Graph/Nsight run completed after the GPU group became available: - -- Result: - `results/dflash2-sparse-target-rejection-b1-nodes-o128-gpu0123-v1.json`. -- SQLite: - `profiles/dflash2-sparse-target-rejection-b1-nodes-o128-gpu0123-v1.sqlite`. -- Full synchronized round: **27.166 ms**, down from 27.951 ms by 0.785 ms. -- Target-to-draft: **1.994 ms**, down from 2.738 ms by 0.744 ms. -- Draft, draft-to-target, and target graph remain 4.004, 1.977, and - 19.191 ms respectively. -- All 128 output IDs and hash `fe0300...` match the accepted baseline. The - run records acceptance length 4.129 and per-position counts - `[27, 22, 18, 14, 9, 4, 3]`. - -This is a Type-A result for the narrow sparse route contract. The environment -gate remains explicit until the default policy and paired dataset result are -landed together. - -### TP4 push all-reduce audit - -The existing vLLM CTA/thread-policy sweep did not improve the 128-collective -chain. A pinned SGLang-V100 checkout at `845b9fdf7a7e` measured its two-epoch -one-shot push collective at 0.846-0.886 ms for the same 128-call -`[8,5120]` FP16 chain, versus 1.854 ms for vLLM's pull path. A first vLLM port -measured 0.873 ms and matched the FP32 rank-order reference numerically across -eager, graph, and 128-call graph replay. - -Two issues prevented accepting that first port: - -1. Its temporary `_C` reused objects from the independent QPN8 worktree. Both - push-on and push-off runs first diverged from the accepted output at token - 24, proving that those runs were not a clean source comparison. -2. SGLang's positive-zero empty-slot sentinel converts an all-positive-zero - reduction to negative zero. A dedicated signed-zero probe found 40,960 - bit mismatches despite zero numerical error. This can perturb later finite - arithmetic and is below the required bitwise gate. - -A release/acquire CTA-ready protocol fixed every payload bit for exact, -random, model-like, positive-zero, and signed-zero inputs over a 128-call graph -chain, but regressed the chain to 1.884-1.901 ms. It is rejected. The accepted -microbenchmark candidate instead preinitializes the two data epochs with a -reserved FP16 NaN payload and retains SGLang's fine-grained polling. Finite -inputs, including both signed zeros, are never rewritten. - -The NaN-sentinel candidate passed the strict TP4 microbenchmark gate: - -- Artifact: - `results/vllm-push-ar-nan-sentinel-correctness-timing-m8-h5120-v1.json`. -- Exact-integer, rank-marker, random-small, model-like, positive-zero, and - signed-zero inputs are bitwise equal to the fixed rank-order FP32 reference - in eager, one-call Graph, and 128-call Graph-chain modes on every rank. -- The 128-call chain is **0.850-0.856 ms**, or about 1.00 ms faster than the - retained 1.854 ms vLLM pull chain. One call averages 6.64-6.68 microseconds - inside the chain. -- The task-owned `_C` contains no QPN8 operator and has SHA256 - `275594c0b38a358c1683efa1d6351e1c572d6da992789d3b5210a74ecbfce1e8`. - -The first unprofiled full-model run with that clean extension completed at -176.56 steady decode token/s. Its 128 token IDs, hash `0cc12...`, acceptance -length 4.467, and per-position accepted counts -`[28, 23, 14, 13, 11, 8, 7]` exactly match the retained push-off control built -with the current CUDA toolchain. They differ from the older `fe030...` release -extension trajectory at token 24; both push-on and push-off current-toolchain -builds make that same transition, so it is not attributed to the collective. -The result is -`results/dflash2-sparse-push-ar-nan-pure-b1-o128-gpu0123-v1.json`. - -A same-binary push-off rerun was invalidated when another task claimed the TP4 -group during model initialization, before KV-cache allocation; it produced no -generation or performance result. The push route therefore remains -default-off until that paired rerun and a clean Nsight complete-round trace are -recorded. The microbenchmark establishes an expected roughly 1 ms graph saving, -but the 176.56 token/s endpoint is not yet used as the paired throughput claim. - -### DFlash2 target QPN8 backport - -The accepted FP8 QPN8 operator and production route were cherry-picked from -`db978179de` and `e0547f1b4a`. The audited integration admits only explicit -QPN8 opt-in, SM70/TP4, block-FP8 layout, exact projection shapes, and bounded -concurrency. It does not inspect a model, architecture, or checkpoint identity; -unsupported tensor or runtime contracts retain TurboMind. - -The target LM head is not a QPN8 shape and remains on the dense FP16 path. The -only admitted target projections are fused gate/up, down, and GDN/full-attention -output. GDN input and full-attention QKV also remain TurboMind. This preserves -the DFlash2 selector and target-sampling contracts while attacking the three -measured FP8 verifier buckets. - -- The imported target-only evidence remains documented in - `sm70_qwen38_qpn8_decode.md`: WikiText perplexity and all fixed GSM8K, MMLU, - and C-Eval metrics pass, and the operator race passed 960 rows. -- The DFlash2-specific runtime-isolation tests pass for target-only, DFlash2, - DFlash1, and a generic non-DFlash speculative configuration. Together with - environment and warmup coverage, 70 focused CPU tests pass. -- A task-owned combined extension contains both the bitwise-gated push - collective and all six QPN8 operators. Its SHA256 is - `512b44cd49e99f26e0f091c68c11bf8c17e5575ffe2c42ba9af1de69938a29e9`. - -No DFlash2 throughput claim is made from the target-only result. The required -next gate is a production-shape M=8 operator smoke followed by paired -single-request DFlash2 generation and a clean Nsight phase trace. The working -hypothesis is that QPN8 plus push can remove enough of the 19.2 ms target graph -to approach a 23 ms complete round, but only the new trace may accept or reject -that hypothesis. - -### Quality-safe 23.678 ms checkpoint - -The next accepted stack keeps the sharded target-hidden projection disabled -and collapses the five full-attention metadata refreshes into one pointer-table -kernel. The paired single-request node traces measure 24.424136 ms for the -control and **23.678170 ms** for the candidate. Draft-to-target falls from -1.857847 to 1.213916 ms. All token IDs, token hash, acceptance length, and -per-position acceptance counts are exact. This grouped metadata route proved -the measured workload, but remains default-off in the audited source so the -kernel contract can be rolled out independently of a model identity. - -The retained target graph has 1,254 nodes and about 16.055 ms of rank-critical -GPU service. Its largest buckets are 4.235 ms gated-pair QPN8, 3.186 ms QPN8 -`(16,1)`, 2.059 ms QPN8 `(16,2)`, 1.277 ms TP4 push all-reduce, 1.205 ms -recurrent GDN, and 0.734 ms Flash-V100 partition attention. This decomposition -is the active short-context optimization budget; the 20 ms complete-round goal -requires at least another 3.679 ms from the measured checkpoint, not merely a -lower isolated kernel number. - -Caching the prepared GDN metadata tensor tuple removes repeated empty-tensor -allocation and tuple assembly. A ten-group/four-layer microbenchmark falls -from 50.519 to 2.663 microseconds and the targeted metadata test passes. The -post-convolution Q/K/V packer is bitwise exact in seven V100 tests and improves -the production-shape graph from 7.475 to 4.116 microseconds, projecting about -0.161 ms over 48 GDN layers. It stays default-off because the packed-GDN route -already subsumes the same rearrangement. - -The graph-internal text-embedding experiment is rejected and fully removed. -Full target graph capture failed in the AOT Qwen signature with -`AttributeError: 'NoneType' object has no attribute 'size'` after omitting the -previously compiled `inputs_embeds` tensor. It generated no tokens and has no -performance result; the failed run must never be counted toward the 20 ms -goal. - -### QPN8 plus TP4 push fusion rejection - -An exact fused epilogue was implemented for the 128 target out/down projection -calls. It rounded the local FP32 QPN8 accumulator to FP16, pushed all four -ranks' packets, accumulated ranks 0 through 3 in FP32, and rounded the result -to FP16. A TP4 128-call CUDA-Graph chain passed every per-rank bitwise check -against both the separate QPN8-plus-push path and the fixed rank-order -reference. - -The performance gate failed: the separate chain was 2.696 ms median, while the -fused chain was 3.249 ms, or 21.06 versus 25.38 microseconds per call. Keeping -the 512-thread QPN CTA resident while one warp polls peer packets increases SM -and NVLink contention instead of hiding communication. The entire production -route, environment switch, bindings, and tests were removed. Only the -standalone benchmark artifact is retained at -`results/qpn8-fused-push-tp4-m8-n5120-v1.json`; this negative 0.552 ms must not -appear in an endpoint speed claim. - -### 2026-08-24 external implementation audit - -The one permitted upstream-head refresh fetched SGLang-V100 at -`0083b9fd83a601b1fcd9a691f7240be4e6be111e`. Its new native long-context work -adds grouped q=1 exact split-KV decode, grouped multi-token verification, and -exact E5M2 byte-to-FP16 expansion. The implementation assigns a Qwen3.8 TP4 -verification layer to two GQA subgroups and at most 40 context splits, instead -of scanning the same cache independently for all eight query rows. - -The same redundant-read diagnosis was independently merged in -[vLLM-Metal PR #534](https://github.com/vllm-project/vllm-metal/pull/534). -That implementation shares each KV block across two verification-window rows -while retaining each row's original accumulation order and reports bitwise -parity. This is stronger evidence than copying a benchmark number: it -identifies redundant KV traffic as a kernel-shape problem specific to -multi-token verification and supplies a quality-preserving scheduling pattern. - -`dnv2003/v100-skinny` remains a useful QPN8 geometry inventory. Its graph -sweep reports `(8,2)` around 12.27 microseconds for M=8/K=1536/N=5120, while -the active target trace uses `(16,1)` around 24.9 microseconds. The kernels and -decoder contracts are not assumed identical. A task-owned sweep of the exact -current operator is therefore required before changing the table; this is now -the highest-value short-context microbenchmark after rejecting QPN/AR fusion. - -### Long-context verification microbenchmark - -The pinned SGLang TileLang 0.1.8 compiler required its matching -`apache-tvm-ffi==0.1.8.post2` and the CUDA 12.8 toolkit still mounted at -`/mnt/nvme3n1p2/usr/local/cuda-12.8`. These are off-tree measurement tools, -not new runtime dependencies. The production implementation will be native -Flash-V100 code. - -On the same V100 and shape `B=1, Q=8, Hq=6, Hkv=1, D=256, page=16, E5M2`, the -measured per-layer medians are: - -| Prefix | Flash-V100 scalar rows | Flash-V100 XQA rows | SGLang grouped verify | -| ---: | ---: | ---: | ---: | -| 1K | 0.238 ms | 0.315 ms | **0.122 ms** | -| 32K | 1.390 ms | 1.151 ms | **0.527 ms** | -| 128K | 5.119 ms | 3.561 ms | **1.957 ms** | -| 256K | 9.064 ms | 6.481 ms | **3.803 ms** | - -The actual current policy selects XQA for E5M2 at and above 16K. Relative to -that production long-context family, grouped verification is 2.18x, 1.82x, -and 1.70x faster at 32K, 128K, and 256K. Across the target's 16 full-attention -layers, the isolated projection is roughly 9.99, 25.67, and 42.84 ms removed -per complete verification round at those lengths. Short-context dispatch must -still use the existing scalar path when the captured workspace is small: the -accepted 8K-capacity trace measures only about 45.5 microseconds per layer, -below the grouped kernel's fixed cost. - -At 1K, both grouped and XQA outputs are within 1.526e-5 maximum and about -1.6e-6 mean absolute error of the FP32 reference. This is currently -`B-pending`, not a default-enablement result: the port must align block tables, -causal boundaries, E5M2 scale semantics, partition order, and shuffled physical -pages; then pass non-growing numeric bounds, greedy trajectory, acceptance, -PPL, and the fixed GSM8K/MATH-500/HumanEval/MBPP/MT-Bench quality suites. - -The first packed-GDN node trace on GPUs 4-7 produced no result. A separate -task claimed the group during initialization, the reservation guard reported -`TP4_HANDOFF_COLLISION_ABORT`, and the run exited before model generation. -No external process was stopped and no timing from that attempt is retained. - -### 2026-08-24 verifier follow-up: rejected seams and two bounded candidates - -The complete-round acceptance gate remains **strictly below 20 ms**. The -latest retained quality-safe result is still 23.678170 ms, so the unresolved -measured gap is at least 3.679 ms. Isolated microbenchmarks below are not -endpoint evidence and must not be presented as completion of that gate. - -Two short-path seam experiments were rejected before integration: - -- The existing fused pull-all-reduce plus Gemma RMS operator measured - 3.005-3.013 ms for a TP4 128-pair graph chain, versus 1.463-1.469 ms for the - current push-all-reduce plus RMS sequence. It was also not bitwise equal on - two ranks. No source route was enabled. -- Reusing the existing paged-prefill operator for the Q=8 verifier measured - 0.436/14.872/57.935/115.838 ms per layer at 1K/32K/128K/256K. Six head CTAs - still serially scan the whole prefix, so merely sharing the eight query rows - does not solve context parallelism. This route is rejected. - -The Flash-V100 split-KV framework was then extended narrowly to E5M2. The -implementation keeps the existing FP16 routes unchanged, dispatches the -generic paged loader with `KV_DTYPE=fp8_e5m2`, preserves separate K/V scales, -and uses the existing split merge order. The source compiles successfully; -the task build has SHA256 `e457f06f...9b0f23`. The production-shape graph -microbenchmark, non-unit scales, shuffled physical pages, and endpoint route -remain pending because all eight GPUs are occupied by external TP4 jobs. - -For the short-context target graph, the local -`Qwen3.8-27B-NVFP4` checkpoint was audited rather than assuming a foreign -checkpoint contract. Its config hash is `1b3c7186...3a5c`; layers 0-55 use -native NVFP4 only for the MLP trunk, while attention, GDN, layers 56-63, and -the full-vocabulary LM head remain FP8. This makes it a useful quality-gated -candidate: the DFlash2-sensitive head is not reduced to four bits, while most -of the verifier's MLP weight traffic is halved. The already accepted no-MTP -route for this exact checkpoint measures 14.017 ms/token, but that number is -not a DFlash2 verification-round result. - -The existing current-tree TurboMind NVFP4 implementation is the control. A -prior real-weight M=1 QPN4 experiment saved a projected 1.090 ms across the 56 -gate/up and 56 down calls, but it is M=1-only and therefore does not establish -an M=8 verifier gain. A new bounded benchmark, -`benchmarks/kernels/benchmark_sm70_nvfp4_qpn2.py`, now loads the real TP-local -layer-55 shards and compares M=8 CUDA-graph replay against the production -TurboMind route. The pinned v100-skinny QPN2 source compiled into a task-only -extension with SHA256 `09d96053...633a4`; no production dispatch has been -changed. The next free-GPU action is this M=8 race, followed by a full DFlash2 -boot only if the measured aggregate saving and numerical gate are positive. - -The mixed checkpoint may become a default performance profile only after -paired FP8-target comparisons pass WikiText PPL, fixed GSM8K, MATH-500, -HumanEval, MBPP, and MT-Bench gates, plus acceptance-length and per-position -acceptance checks. A low-bit microbenchmark alone cannot authorize that -change. - -### 2026-08-24 QPN2 production admission and split-KV rejection - -The generic Flash-V100 E5M2 split-KV follow-up is now rejected rather than -pending. With production shape `B=1,Q=8,Hq=6,Hkv=1,D=256`, non-unit scales, -and shuffled physical pages, its per-layer medians were 0.124/1.936/7.791/ -15.769 ms at 1K/32K/128K/256K. The 1K result is numerically healthy (maximum -absolute error `1.526e-5` versus FP32), but the long-context results are worse -than the retained XQA path. The generic split kernel still rereads one KV -stream per query head; it does not implement SGLang's grouped-GQA sharing and -therefore cannot be enabled. The raw artifact is -`results/flash-v100-splitkv-verify-e5m2-b1-q8-g6-d256-v1.json`; the task-only -binary SHA256 is `e457f06f...9b0f23`. - -The real-weight QPN2 operator race is positive at exact verifier width M=8, -but its control is the mixed NVFP4 checkpoint's existing TurboMind route, not -the FP8+QPN8 target used by the accepted 23.678170 ms complete-round baseline. -Across layer-0/rank-3 gate/up and down projections, replacing TurboMind with -QPN2 projects **3.814523 ms removed per complete mixed-target round**: -gate/up falls from 96.020 to 40.681 microseconds and down from 31.687 to -18.909 microseconds. Relative L2 error versus the FP32 checkpoint oracle is -`0.000569` for QPN2 gate/up versus `0.000598` for TurboMind, and `0.000300` -for QPN2 down versus `0.000291` for TurboMind. The retained artifacts are -`results/nvfp4-qpn2-vs-turbomind-real-layer0-rank3-m8-oracle-v1.json` and -`results/nvfp4-qpn2-vs-turbomind-real-layer55-rank0-m8-oracle-v1.json`. - -That 3.814523 ms value must not be subtracted directly from 23.678170 ms. -Against the actual FP8+QPN8 baseline buckets, the current evidence projects -only about 1.8 ms, with a plausible 1.8-2.5 ms range pending the production -fused-gate measurement. QPN2 alone therefore does not prove the sub-20 ms -gate; a new complete single-request DFlash2 trace must measure the movement, -and any remaining 1-2 ms gap must be removed from that trace's actual -critical-rank buckets. - -A production QPN2 implementation is linked behind -`VLLM_SM70_NVFP4_QPN2=1`. Admission uses TP4 and the checkpoint-native gate/up -or down tensor layout; it does not inspect model or speculator identity. M<=8 -uses QPN2, while larger dynamic M dispatches inside the opaque C++ operator to -the existing TurboMind layout. The switch remains default-off. - -The exact production source compiles for SM70, its four routing tests pass, -and a full `_C` candidate contains QPN8, graph-only push, and all four QPN2 -operators. Candidate path: -`build/push-qpn8-qpn2-graphonly-v3/_C.abi3.so`, SHA256 -`c133500e...3f068b5`. The earlier task-only standalone hash is obsolete after -the final production-source edits and is intentionally not used as endpoint -evidence. - -The production candidate passed the real-weight M=8 CUDA-Graph oracle on a -V100. Gate/up is 96.759 versus 36.801 microseconds and down is 31.830 versus -19.794 microseconds against TurboMind, for a 4.031684 ms mixed-target -projection. QPN2 relative L2 versus the FP32 checkpoint oracle is `0.000569` -for the fused gate/up output and `0.000300` for down. The result is -`results/nvfp4-qpn2-production-layer0-rank3-m8-v1.json`, SHA256 -`cfc0634d...1a5a6c`. This passes the operator gate but does not alter the -correct-baseline warning above: only a complete DFlash2 trace can prove the -net movement from 23.678170 ms. - -The first full-service launch exposed a checkpoint contract mismatch rather -than a QPN2 failure. The mixed checkpoint declares calibrated E4M3 KV scales, -which upstream vLLM correctly refuses to reinterpret as E5M2 scales. This is -the behavior retained by upstream PR #45040. A task-cache-only model overlay -therefore keeps every original weight/tokenizer file by symlink but removes -the incompatible `kv_cache_scheme` from its copied config, making the explicit -E5M2 run use the same unit-scale storage contract as the accepted FP8 target. -The original checkpoint is unchanged. Overlay config SHA256 is -`5b72ecd7...40f6e52`; the benchmark now records both that hash and the resolved -weight path. - -The overlay launch passed the KV contract, but an unrelated TP4 job acquired -GPUs 4-7 between the availability check and worker initialization. Free -memory fell to 24.41 GiB, below the matched 0.8 utilization contract, so vLLM -refused startup before model loading. No endpoint or timing result is claimed -from either failed launch. The next valid action remains the matched complete -single-request trace on a genuinely free TP4 group; the hard acceptance gate -is strictly below 20.000 ms. - -The next owned launch reached DFlash2 warmup and exposed a second checkpoint -contract rather than a timing result: the mixed checkpoint quantizes -`lm_head`, while DFlash2 candidate selection intentionally requires an -unquantized target LM head. The final task-cache overlay therefore maps only -`lm_head.weight` to the accepted FP8 target's BF16 `[248320,5120]` tensor and -removes `lm_head` from both the mixed quantization target list and its FP8 scale -entry. All other target weights still resolve to the original mixed -checkpoint. The combined overlay config and index SHA256 values are -`1b418636...01e8a` and `dacc00bf...836e`; the BF16 sidecar SHA256 is -`ddff1d66...b698ff`. The original FP8 and mixed model directories remain -unchanged. - -The launch audit also found that the production checkpoint selects -`CompressedTensorsW4A4Fp4` because its config contains dynamic NVFP4 input -activation metadata. The first QPN2 integration had been attached to the -W4A16 scheme, so its real-weight microbenchmark remained valid but the full -model did not actually dispatch QPN2. The exact admission and fused -gate/up/down dispatch now live in the W4A4 scheme, and the W4A16 file is back to -its original content. The routing, FP8-prefill, and DFlash2 targeted suites -pass **78 tests**. No endpoint speed claim is made until the newly queued -single-request node trace contains the QPN2 kernels and measures a complete -round strictly below the 20.000 ms gate. - -Because the default safetensors iterator visits every tensor in each referenced -shard, deleting only the index entry was insufficient: the original monolithic -body still exposed `lm_head.weight_scale` during loading. A streaming -task-cache rewrite now removes exactly `lm_head.weight` and -`lm_head.weight_scale` from the body without decoding or re-encoding any other -tensor. The slim body is 21,296,297,056 bytes with SHA256 -`9f49b437...389267`; all 1,951 retained tensor dtype/shape records match and -the first, middle, and last bytes of every retained tensor match the source. -The final index SHA256 is `94f500bf...c7ef2c`, and the benchmark records that -index plus every actually referenced weight realpath. This replaces neither -the original checkpoint nor the DFlash2 unquantized-LM-head guard. - -The first valid mixed-target node trace is a rejection, not a speed result to -retain. Its complete round is **40.775579 ms**: draft 3.941496 ms, -draft-to-target 1.228315 ms, target 33.567780 ms, and target-to-draft -2.037988 ms. The target graph has 1,462 nodes. QPN2 itself matches its -microbenchmark, consuming 2.521294 ms for 56 gate/up calls and 1.518048 ms for -56 down calls. The regression is instead 144 residual channel-FP8 projections -falling through to SM70 Marlin for **21.110017 ms**. The single-request output -has acceptance length 3.823529; neither its 91.451 tok/s steady decode nor that -one-prompt acceptance value is a quality or performance acceptance claim. - -The missing dependency was already implemented and DCO-signed in local commit -`5123a6a7b2` (`[Kernel] Recover Qwen3.8 NVFP4 SM70 decode`): its narrow -compressed-tensors W8A16 channel-FP8 adapter reuses the accepted TurboMind and -QPN8 operators. Only that adapter, its fused-SiLU delegation seam, and focused -tests were ported here; unrelated attention and tuning changes from the commit -were not duplicated. This also fixes the QPN2 production seam: without -`CompressedTensorsLinearMethod.apply_fused_silu_and_mul`, the full model ran a -separate Triton activation even though the scheme exposed the fused operator. -The combined compressed-FP8, QPN2, FP8-runtime, and DFlash2 suites pass **83 -tests**. A new complete trace must prove removal of the 21.110 ms Marlin bucket -and still meet the strict sub-20.000 ms gate. - -### 2026-08-24 mixed-target recovery and remaining 1.807 ms - -The channel-FP8 adapter and QPN2 route are now present in a clean combined -extension. The retained `v8` complete single-request trace uses the mixed -NVFP4 target, QPN2 MLP, QPN8 residual projections, graph-only TP4 push, -sharded target-hidden projection, E5M2 target KV, FP16 draft KV, and an -unquantized target LM head. Across 140 steady rank-cycles it measures: - -| Phase | Mean wall | -| --- | ---: | -| Draft graph | 3.914734 ms | -| Draft to target | 1.097309 ms | -| Target graph | 15.367610 ms | -| Target to draft | 1.729242 ms | -| Complete round | **22.108896 ms** | - -The target graph has 1,254 nodes. Its largest service buckets are QPN8 -projections (3.327 ms), QPN2 gate/up (2.511 ms), QPN2 down (1.568 ms), push -all-reduce (1.294 ms), recurrent GDN (1.200 ms), Flash-V100 attention -(0.738 ms), and the 127 Gemma suffix kernels (0.676 ms). Both the draft -selector and target rejection still execute an approximately 0.98 ms dense -FP16 LM head. - -The draft-to-target interval contained six four-byte device-to-host copies and -stream synchronizations. Five were accidental: the five Flash-V100 groups -recomputed `seq_lens_cpu` even though `InputBatch.seq_lens_cpu_upper_bound` -already held the required contract. Passing that tensor through -`MambaHybridModelState.prepare_attn()` removes exactly those five syncs. The -new focused unit test passes. The corrected `v10` trace, with the sharded -target-hidden projection explicitly restored after reboot, measures: - -| Phase | Mean wall | -| --- | ---: | -| Draft graph | 3.914056 ms | -| Draft to target | **0.790517 ms** | -| Target graph | 15.357529 ms | -| Target to draft | 1.745331 ms | -| Complete round | **21.807433 ms** | - -The complete median/min/max are 21.823373/21.557190/22.653044 ms. This is the -current accepted local best, but it **does not pass** the strict sub-20.000 ms -gate: the measured deficit remains 1.807433 ms. The one remaining D2H fence -is the batch-level GDN metadata ordering assertion. A retained historical -GPU-assert experiment measured 24.093540 ms versus 23.679850 ms for its -same-generation synchronized control, so it is rejected; the final ordering -edge will not be removed on the basis of a theoretical sync saving. - -At the exact production LM-head shape `M=8,K=5120,N_local=62080`, a new V100 -probe measures `torch.mm` at 1.018880 ms p50 and the existing SM70 dense -full-logits operator at 0.812032 ms p50. The 0.206848 ms per-call gain projects -to about 0.414 ms for the two heads in one complete round. The probe's top-16 -support and values are exact, but full-logit non-bitwise behavior still makes -this a quality-gated candidate rather than accepted endpoint evidence. Raw -artifact: `results/lmhead-dense-sm70-vs-torch-m8-n62080-k5120-v11.json`. - -The next bounded candidate preserves the accepted push transport and fuses -its following FP32 residual/Gemma RMSNorm suffix. It targets the 128 push -collectives plus 127 RMS kernels as a 128-node chain. It must first beat the -current production-shaped TP4 chain microbenchmark, keep the reduced residual -bitwise exact, and satisfy the FP16 normalized-output bound. Only then may it -be combined with the LM-head candidate in a full Graph/Nsight run; neither -microbenchmark may be subtracted from 21.807433 ms to claim completion. - -That push-suffix candidate is now rejected and fully removed from the -production source. The first centralized-row implementation measured -2.544-2.550 ms for the 128-pair TP4 graph chain versus 1.475-1.480 ms for the -current push-plus-Triton sequence. A second implementation fused each CTA's -partial square sum and distributed normalization across all five row CTAs; it -reduced the candidate to 1.638 ms but the matched control remained -1.470-1.472 ms. The best fused form is therefore still about 11.3% slower. -Both versions kept residual output bitwise exact; the second version was -normalized-output bitwise exact on three ranks and differed in two FP16 values -on the fourth, with maximum absolute error 0.000244. Artifacts are -`results/tp4-push-vs-push-fused-gemma-rms-m8-h5120-v2.json` and `v3.json`. -No endpoint run or speed claim is permitted from this failed branch. - -### 2026-08-24 synchronized-fence removal and 20.746 ms endpoint - -The final batch-level GDN metadata assertion was audited against both the -metadata builders and `ModelCudaGraphManager.run_fullgraph()`. All metadata -kernels, copies, and graph replay are enqueued on the same current CUDA stream; -the device `.item()` checked a debug invariant but supplied no ordering edge. -The assertion is separately gated by `VLLM_SM70_DFLASH2_GDN_SYNC_ASSERT`. -After the exact token/acceptance trajectory passed, the production default was -changed to off; it remains an explicit opt-in metadata audit. A structural CPU -invariant remains unconditional. - -The `v17` single-request Graph/Nsight trace disables only that debug fence and -enables the already quality-gated dense FP16 LM-head route. Across 140 steady -rank-cycles it measures: - -| Phase | Mean wall | -| --- | ---: | -| Draft graph | 3.756561 ms | -| Draft to target | 0.156985 ms | -| Target graph | 15.288068 ms | -| Target to draft | 1.543990 ms | -| Complete round | **20.745603 ms** | - -The complete median/min/max are 20.747498/20.523095/20.939874 ms. This is a -1.061830 ms improvement over `v10`, but it still misses the strict sub-20 ms -gate by 0.745603 ms and is therefore not an accepted endpoint. The output -token SHA256 remains -`22db7b2e81df17d3afdc6982fc1425769335c6cc7931da10e64bd6bac3459f28`; -the answer is correct and the complete acceptance contract is unchanged at -3.685714 with per-position counts `[28,22,17,12,8,5,2]`. Steady decode is -174.016 tokens/s. Result/report/sqlite SHA256 values are respectively -`1ee462b4...78a2e`, `7f686d51...f3ba`, and `1f82d151...3302`. - -The remaining two local LM-head passes are now the bounded next target. With -the real rank-0 FP16 shard `[62080,5120]`, channel-QPN8 local top-64 contains -every dense top-16 token in 512/512 random rows. Re-evaluating those 64 rows -with the original FP16 weight via `index_select` plus batched dot products -reduces the complete local selector from 0.904192 to 0.518144 ms p50. The -reranked top-1 matches 511/512 rows; maximum/mean absolute logit differences -are 0.00390625 and 6.68e-6. Artifact SHA256 is -`39fde704...718a7`. These random-hidden results authorize only a production -candidate, not a quality or endpoint claim. - -The implementation is isolated behind -`VLLM_SM70_DFLASH2_QPN8_RERANK=0`. It admits only SM70, TP4, the exact local -LM-head shape, FP16 weights, zero vocabulary padding, top-k 16/20, and at most -eight rows. QPN8 supplies only a coarse top-64 support; the returned logits -are always recomputed from the original FP16 rows. The same shared target LM -head serves both the draft selector and sparse target rejection. An eager-only -shadow gate executes the candidate, logs real-hidden support/set/top-1 -agreement, and returns the dense result so the audit trajectory cannot change. -The path remains default-off until four-rank real-weight coverage, real-hidden -shadow coverage, dataset quality/acceptance, and a new complete Graph trace all -pass. Eighteen focused CPU routing tests and static checks currently pass. - -The stricter four-shard `v18` probe uses local top-20, candidate width 64, -1,024 random rows per TP rank, and the same real checkpoint weights. All -4,096 dense top-20 sets are fully contained in the QPN8 support (minimum -recall 1.0 on every rank). Fixed split-8/non-fast QPN8 is 0.425-0.427 ms; -complete QPN8 plus original-FP16 rerank is 0.499-0.531 ms versus -0.893-0.923 ms for the dense selector. The reranked top-1 matches 4,081/4,096 -rows, while exact local top-20 set agreement is 3,971/4,096. The maximum -returned-logit difference is one FP16 step (0.00390625). Thus the coarse -support gate passes, but reduction-order differences still require real-hidden -and model-quality gates; `v18` alone does not authorize endpoint use. The four -result SHA256 prefixes are `b07428ec`, `85e94de8`, `e56ed9b0`, and `fb4c1016`. -The complete CPU DFlash2 targeted file now passes 52 tests with 12 GPU tests -skipped under an explicitly hidden CUDA environment. - -The `v19` eager shadow then exercised the same candidate on real hidden states -while always returning the dense local result. Across four TP ranks it -observed 8,564 local rows (draft top-16 and target top-20), with **zero** dense -tokens missing from the QPN8 top-64 support. Original-FP16 rerank set agreement -is 8,406/8,564 (98.1551%) and local top-1 agreement is 8,547/8,564 (99.8015%); -the largest ordered logit difference is 0.015625. This passes the coarse -support gate and quantifies the remaining reduction/tie-order effect. Because -shadow is eager-only and deliberately returns dense outputs, neither its -throughput nor its four truncated GSM8K answers are candidate quality evidence. -Log/result SHA256 values are `6dad2b23...482e` and `29c8ff88...753f`. The next -required artifact is the active single-request Graph/Nsight trace; the hard -gate remains a measured complete mean strictly below 20.000 ms. - -### 2026-08-24 exact-rerank endpoint and sub-20 ms replication - -The first active Graph launch, `v20`, is rejected before timing. Startup, -compilation, and QPN8 preparation succeeded, but the draft selector's TP -all-gather rejected a non-contiguous top-16 tensor. The implementation had -sliced the first 16 columns from one persistent `[8,20]` output, retaining a -row stride of 20. It now owns separate graph-stable `[8,16]` and `[8,20]` -value, position, and ID buffers. Row-only views preserve contiguity without a -runtime copy. Six layout cases (top-16/top-20 by 1/7/8 rows), the target -candidate route, Ruff, and `git diff --check` pass. `v20` produced no timing -or quality result. - -The corrected `v21` single-request probabilistic CUDA Graph/Nsight trace -crosses the strict complete-mean gate for the first time. Across 140 steady -rank-cycles on physical GPUs 4--7 it measures: - -| Phase | Mean wall | -| --- | ---: | -| Draft graph | 3.332724 ms | -| Draft to target | 0.158264 ms | -| Target graph | 15.281917 ms | -| Target to draft | 1.163312 ms | -| Complete round | **19.936216 ms** | - -The complete median/min/max are 19.956518/19.751641/20.129946 ms. Steady -single-request decode is 180.686 tokens/s. The 128-token SHA256 is exactly -`22db7b2e81df17d3afdc6982fc1425769335c6cc7931da10e64bd6bac3459f28`, -the answer is correct, and acceptance is exactly unchanged at 3.685714 with -per-position counts `[28,22,17,12,8,5,2]`. Result/report/sqlite SHA256 values -are `3d22d740...e1248`, `27175b2f...09c7`, and `5fe2773c...60e0`. - -An unchanged `v22` replication then ran on physical GPUs 0--3. Its 140 -steady rank-cycles measure 3.339601 ms draft, 0.155429 ms draft-to-target, -15.330372 ms target, 1.165597 ms target-to-draft, and **19.990999 ms** complete -mean. The median/min/max are 20.011181/19.785692/20.212695 ms, and steady -decode is 181.177 tokens/s. Token SHA256, correct answer, acceptance length, -and every per-position acceptance count again match the dense control. The -result/report/sqlite SHA256 values are `06fa010e...19b`, `af23b773...031`, and -`c34083f8...f62e`. - -Thus the complete mean is below 20.000 ms in two independent runs and on both -physical TP4 groups; this is direct Graph/Nsight evidence rather than a -microbenchmark projection. The margin is still only 0.009--0.064 ms and the -upper tail crosses 20 ms, so the candidate remains default-off pending a -larger stability margin and the multi-dataset/PPL quality gate. The next -bounded microbenchmark checks whether the QPN8 support top-64 can skip ordering -before exact FP16 reranking; candidate order is not consumed by the rerank. - -The `v23` sorted/unsorted support experiment rejects that optimization. On the -real rank-0 shard, removing QPN8 top-64 ordering changes complete rerank p50 -from 0.513024 to 0.501760 ms. In the production fixed split-8 core the saving -is only about 4.1 microseconds per LM-head call, or about 8 microseconds per -complete DFlash2 round. This is below run-to-run noise and changes tie order, -so production retains sorted support. The sorted/unsorted artifact SHA256 -values are `5eb1ad6f...31de` and `da3499cc...fb6`. - -The accepted next micro-optimization removes the 5.2 MiB FP16 -`index_select` temporary and its batched matrix multiplication. A fixed-grid -Triton kernel directly computes each selected original-FP16 vocabulary-row dot -product. For the production `[8,62080,5120]`, top-64 to top-20 shape, the -complete selector p50 falls from 0.517120 to 0.464896 ms. All 512/512 top-20 -sets and top-1 IDs match the gathered-bmm reference; maximum ordered-logit -difference is 0.001953125. The selected configuration is `BLOCK_K=1024`, four -warps, with a graph-stable 512-program grid masking inactive rows. The -artifact SHA256 is `b898c410...0631`. Dedicated V100 tests for one, seven, and -eight rows pass, as do the focused CPU routing/layout tests and static checks. - -Two unchanged full Graph/Nsight runs then validate the direct indexed-dot -endpoint on both TP4 groups: - -| Run / physical GPUs | Draft | D-to-T | Target | T-to-D | Complete mean | Median | Max | Steady decode | -| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | -| `v25` / 4--7 | 3.324374 | 0.157883 | 15.272063 | 1.139965 | **19.894284 ms** | 19.909407 | 20.087468 | 182.516 tok/s | -| `v26` / 0--3 | 3.332363 | 0.157819 | 15.313486 | 1.145328 | **19.948997 ms** | 19.960269 | 20.148380 | 181.216 tok/s | - -Each result contains 140 steady rank-cycles. Both target and draft are FULL -CUDA Graphs and the indexed-dot path is captured without a graph break. Both -runs reproduce the baseline 128-token SHA256 -`22db7b2e81df17d3afdc6982fc1425769335c6cc7931da10e64bd6bac3459f28`, -the correct answer, acceptance length 3.685714, and per-position counts -`[28,22,17,12,8,5,2]` exactly. The `v25` result/report/sqlite SHA256 values are -`c82ef281...0952`, `f9c26647...6c81`, and `7c7645ef...006a`; the corresponding -`v26` values are `04498a79...0cbc`, `0ebd0b86...4244`, and -`3de27587...9d1`. Parsed round artifacts have SHA256 -`a49fe2cd...be3b` and `bb9d09dc...5e2d`. - -The strict complete-mean and median target is therefore independently -replicated below 20 ms after the final kernel change. It is not yet an -all-sample tail claim: the observed maxima remain 20.087--20.148 ms. The -rerank and GDN debug-fence defaults remain conservative until paired -multi-dataset generation and prompt-logprob/PPL gates pass; a further -0.2--0.4 ms safety-margin reduction remains the immediate performance target. - -### 2026-08-24 paired quality rejection and dense-order correction - -The first paired 16-question GSM8K quality run (`v27`) rejects the active -QPN8 reranker despite its exact single-prompt trace. Both variants are -sequential single-request FULL Graph runs with probabilistic draft sampling, -temperature 1.0, top-p 0.95, top-k 20, identical shuffled rows and request -seed. Dense control scores 11/16 with aggregate acceptance 4.490716; the -candidate scores 9/16 with acceptance 4.412224. The acceptance delta is --0.078493, below the fixed -0.05 gate, and only 5/16 token streams match. -Artifacts `gsm8k16-qpn8-indexeddot-v27-{control,candidate}.json` have SHA256 -`ced400d1...aee` and `e4c89772...4eeb`; the failed comparison has SHA256 -`9650f531...045c`. Consequently the candidate remains default-off and its -sub-20 traces are not accepted as a default quality-safe endpoint. - -The follow-up `v28` microbenchmark corrects the reference used during rerank -development. It compares against the production TurboMind -`sm70_f16_gemm_out`, not `index_select+bmm`. Across 1,024 random rows on the -four real TP shards, QPN top-64 contains every production dense top-20 token. -Depending on reduction configuration, 99.969--99.994% of candidate logits are -bitwise equal, yet candidate-order top-20 set equality is only 250--253/256 -rows per shard. This isolates a second source of drift: `torch.topk` tie -selection over a shuffled 64-element support does not reproduce its selection -over the original 62,080-element vocabulary order. - -`v29` restores that dense order by scattering the 64 re-evaluated FP16 logits -into a persistent full-shard buffer initialized to negative infinity before -local top-k. All four shards then match production dense output exactly for -all 1,024 rows: ordered token IDs, top-20 sets, top-1 IDs, and returned FP16 -values. Complete local selector p50 is 0.5007--0.5038 ms, versus -0.4608--0.4628 ms without dense-order restoration and 0.894--0.930 ms for the -production dense selector. The four result SHA256 values are -`93ebf045...15ec`, `6ec9f526...c064`, `9e4e9868...2a`, and -`904790b6...bb3`. Focused CPU tests (12) and V100 tests (4) pass. A new -same-source paired GSM8K run is required before this correction can re-enter a -complete Graph/Nsight trace. - -The same-source `v30` paired run rejects the corrected active path at the -acceptance gate. Both variants are sequential single-request FULL Graph runs -over the same 16 shuffled GSM8K questions, probabilistic draft sampling, -temperature 1.0, top-p 0.95, top-k 20, and identical request seeds; the only -environment difference is `VLLM_SM70_DFLASH2_QPN8_RERANK=0/1`. Control and -candidate both score 11/16 with zero invalid answers, but pooled acceptance -changes from 4.581867 to 4.498670 (delta -0.083197), below the fixed -0.05 -gate. Seven of 16 token streams are identical. Candidate mean steady decode -is 249.052 token/s versus 244.506 token/s for control, but that speed is not -accepted because the paired acceptance gate fails. Control, candidate, and -comparison SHA256 values are respectively `a3289717...db1`, -`d446ec34...9776`, and `6b32738b...71f`. The reranker remains default-off and -must not be included in a quality-safe sub-20 ms claim. - -This result also narrows the remaining numerical mechanism. The earlier -real-hidden shadow observed zero dense top-k tokens missing from QPN8 top-64, -so the failure is not coarse-support recall. The candidate recomputes only -the selected rows with a different reduction schedule from the production -TurboMind dense GEMM; real-hidden ordered logits can differ by up to 0.015625. -Restoring dense vocabulary tie order fixes exact random-row tests but cannot -restore those production reduction values. The next selector candidate must -either execute the selected rows with the same Tensor-Core reduction contract -as the dense kernel or pass a larger paired distribution/acceptance gate; it -cannot be promoted from random-hidden equality alone. - -### 2026-08-24 exact context-projection overlap rejection - -The default-off `v31` experiment moves the position-independent DFlash2 -target-hidden combine, TP4 sharded context projection, and fused five-layer K/V -projection onto an auxiliary CUDA stream before target sampling. The main -stream waits on one event only when RoPE and cache slots are available. It is -numerically exact in the full single-request probabilistic run: the 128-token -SHA256 remains `22db7b2e...9f28`, the answer remains correct, and acceptance is -exactly 3.685714 with per-position counts `[28,22,17,12,8,5,2]`. Both target -and draft remain FULL CUDA Graphs and there is no deadlock or graph break. - -Performance rejects the overlap. Across 140 steady rank-cycles, the trace -measures 3.802137 ms draft, 0.254715 ms draft-to-target, 16.041994 ms target, -1.748872 ms target-to-draft, and **21.847718 ms complete**. Complete -median/min/max are 21.858843/21.612554/22.046909 ms, and observed steady decode -is 165.686 token/s. This is about 1.10 ms slower than the 20.745603 ms exact -`v17` baseline, so the gate remains default-off. - -The dual-stream timeline explains why subtracting the isolated projection cost -was invalid. The 0.824 ms dense LM head occupies the SMs first; the 0.118 ms -context projection cannot start until its tail. It then overlaps the many -small target top-k/rejection kernels and stretches the main sampling tail by -about 0.25 ms. The auxiliary all-gather and K/V projection finish before the -event wait, but the contention and cross-stream launch schedule cost more than -the hidden work. Result, qdstrm, report, sqlite, and parsed-round SHA256 values -are `7a75a159...181`, `0cfa07d8...17d`, `83a05726...99e`, -`1891c169...6c1`, and `6b87003b...869`. No speed from this experiment is -accepted; a future overlap attempt needs a fused single-stream schedule or an -explicitly prioritized kernel, not another generic auxiliary stream. - -### 2026-08-24 native grouped verifier and long-context result - -The retained verifier optimization does not reuse DDTree indices or skip target -attention. It implements the MRV2 block-eight semantics directly in -Flash-V100: eight query rows share each E5M2 K/V page, accumulate all six GQA -heads, and write the eight verifier outputs in one pass. Admission requires -SM70, FP16 Q, E5M2 KV, B1, Q=8, Hq=6, Hkv=1, D=256, causal decode, and the -original single-request metadata. All other shapes keep the previous path. -The dedicated CUDA-Graph test poisons and replays the sequence length and cache -pages; six V100 cases pass against the FP32 oracle. - -The paired same-process CUDA-Graph microbenchmark at page size 3,296 measures: - -| Context | Native one-pass | Native two-pass | Previous XQA | -| ---: | ---: | ---: | ---: | -| 1K | 0.068608 ms | 0.097280 ms | 0.230400 ms | -| 32K | 0.356352 ms | 0.544768 ms | 1.024000 ms | -| 128K | 1.214464 ms | 1.864704 ms | 3.021312 ms | -| 256K | 2.396160 ms | 3.690496 ms | 6.011392 ms | - -The microbenchmark is -`results/flash-v100-grouped-paired-page3296-v72.json`, SHA256 -`aeccb306...657d9`. It establishes the isolated attention saving; endpoint -claims use the full-model traces below. - -Matched 32K and warmed 128K single-request FULL-Graph runs preserve the output -token stream and complete acceptance counters within each on/off pair: - -| Input | Grouped | Complete round | Target phase | Steady decode | -| ---: | :---: | ---: | ---: | ---: | -| 32K | off | 36.310865 ms | 30.248113 ms | 77.537 tok/s | -| 32K | on | **26.361447 ms** | **20.451644 ms** | **88.392 tok/s** | -| 128K | off | 69.854780 ms | 63.708214 ms | 106.428 tok/s | -| 128K | on | **40.778072 ms** | **34.646866 ms** | **182.294 tok/s** | - -At 32K the complete round falls 27.40% and the target phase 32.39%; at 128K -they fall 41.62% and 45.62%. The 128K pure-decode result rises 71.28%. The -repetitive fixed long-context prompt has acceptance 6.095 at 32K and 7.647 at -128K in both variants; these values are paired route checks, not corpus -acceptance estimates. The four phase artifacts and SHA256 values are -`v77`/`2d705aa9...0c0e`, `v76`/`de6d922d...4235`, -`v79`/`e61e858f...cc1c`, and `v80`/`76164c8c...8deb`. - -The short fixed 128-token trace is neutral rather than a claimed win: grouped -is 21.117097 ms versus 21.060181 ms off. The long-context admission threshold -therefore remains `max_model_len >= 32768`; the feature can be disabled with -`VLLM_FLASH_V100_DFLASH2_GROUPED_VERIFY=0`. Dataset generation below is the -deciding endpoint gate, where grouped improves aggregate throughput despite -mixed prompt and output lengths. - -### 2026-08-24 replicated sub-20 ms short-round bundle - -Before the grouped-attention addition, the complete candidate bundle was -replicated on both physical TP4 groups under FULL target and draft CUDA Graphs: - -| Run | Draft | D-to-T | Target | T-to-D | Complete mean / median / max | -| --- | ---: | ---: | ---: | ---: | ---: | -| `v52-r2` | 3.374848 | 0.156565 | 15.292271 | 1.073660 | **19.897343 / 19.895756 / 20.098654 ms** | -| `v53` | 3.385334 | 0.157777 | 15.320752 | 1.079131 | **19.942993 / 19.948144 / 20.145704 ms** | - -Thus mean and median are below 20 ms twice; this is not an all-sample-tail -claim because both maxima remain slightly above 20 ms. Phase artifacts have -SHA256 `eb8feded...5482` and `412ad1a0...7599`. These traces authorize the -performance candidate, while the larger score/PPL run below authorizes its -sampling and model-quality behavior. - -### 2026-08-24 probabilistic 2K score and PPL gate - -Per the final acceptance decision, greedy token-for-token equivalence is no -longer a production gate. Quality is evaluated with probabilistic DFlash2, -temperature 1.0, top-p 0.95, target top-k 20, reasoning effort `xhigh`, fixed -seed, `max_tokens=2048`, and natural EOS. Every run is sequential B1 and uses -the same target checkpoint, TP4, E5M2 target KV, FP16 draft KV, Flash-V100, and -FULL Graph. The target is the local mixed NVFP4/FP8 checkpoint with a BF16 -LM-head; it must not be described as the official unquantized checkpoint. - -The first 128-row dataset contains 32 GSM8K, 32 MATH-500, 32 HumanEval, and 32 -raw-MBPP prompts. The raw MBPP prompts copied the speed benchmark and omitted -the function name and tests, so their 0--2/32 results are invalid as task -scores and excluded. The remaining 96 scored cases are: - -| Variant | GSM8K | MATH-500 | HumanEval | Included total | -| --- | ---: | ---: | ---: | ---: | -| Target-only | 29/32 | 26/32 | 28/32 | **83/96 (86.458%)** | -| Dense verifier DFlash2 | 30/32 | 24/32 | 28/32 | 82/96 (85.417%) | -| Native grouped DFlash2 | 30/32 | 25/32 | 28/32 | **83/96 (86.458%)** | - -Target-only versus grouped has three wins on each side and McNemar exact -two-sided p=1.0. Dense versus grouped has two dense-only and three grouped-only -wins. Grouped therefore matches target-only aggregate score and improves one -case over the dense-verifier DFlash2 sample; no task-score regression is -observed. - -The corrected MBPP prompt includes the requested function name and all three -held-out tests, following the Google MBPP evaluation contract. Target-only and -grouped both score **25/32 (78.125%)**. They have two exclusive wins each, 23 -shared passes, five shared failures, and McNemar exact p=1.0. Target has eight -length finishes and grouped six. The grouped MBPP run overlapped an unrelated -process that appeared after reservation, so its throughput is discarded; -pass/fail scoring completed without OOM or engine error. - -Across the valid 96-case mixed run, performance is: - -| Variant | Aggregate output | Mean request steady decode | Mean completion tokens / verify step | -| --- | ---: | ---: | ---: | -| Target-only | 67.128 tok/s | 68.008 tok/s | n/a | -| Dense verifier DFlash2 | 176.341 tok/s | 218.436 tok/s | 4.2243 | -| Native grouped DFlash2 | **193.484 tok/s** | **231.126 tok/s** | **4.2384** | - -Grouped is +9.72% in aggregate output and +5.81% in mean steady decode versus -the dense verifier. Its pooled acceptance is 3.749104 versus 3.749153 dense; -the difference is negligible, while the request-mean completion length is -+0.0141. The summary artifact is -`quality/dflash2-mixed32-quality-2k-v83-summary.json`, SHA256 -`9ca37f6a...dab7`. Corrected MBPP score artifacts have SHA256 -`92908b64...08f3` and `326e559b...339e`. - -The supplementary WikiText-2 gate scores 16,376 prompt tokens in eight equal -2,048-token windows. Token-weighted/geometric PPL is 5.497872 target-only and -5.497376 grouped, a relative **-0.0090%** change. Mean absolute prompt-logprob -difference is 0.007519 and the maximum is 1.101103. The generic comparison tool -labels this as failure only because its historical bounds require every -logprob and every per-window PPL to agree within `1e-6`; that exactness contract -is intentionally superseded here by the user-approved score/PPL contract. -Aggregate PPL does not regress. Artifact SHA256 is `d2c54784...e0c7`. - -The official DFlash2 model card publishes lossless-distribution semantics plus -acceptance/throughput, but no GSM8K, MATH-500, HumanEval, or MBPP task accuracy. -Its published acceptance lengths are 5.46/5.28/4.39/4.79 respectively under -H200, SGLang, FlashAttention-3, block eight, maximum 4,096 tokens, and the full -official datasets. They are useful reference points, not a direct comparison -to this 32-row-per-suite V100/mixed-checkpoint/max-2K gate. See -. - -### Audit disposition and isolation - -The score, PPL, paired acceptance, microbenchmark, short-trace, and long-context -results remain evidence for the measured workload. They do not install a model -profile. Every new leaf is default-off and independently selected by explicit -environment opt-in plus hardware, algorithm configuration, dtype, tensor -layout, cache/graph mode, sampling, and bounded-concurrency contracts. - -The earlier clean 8-token TP4 auto-profile smoke remains historical route -evidence only; automatic profile injection was removed during merge audit. -Unsupported requests execute the previous dense/general path rather than -changing routing semantics. diff --git a/docs/design/sm70_nvfp4_turbomind_operator_optimization.md b/docs/design/sm70_nvfp4_turbomind_operator_optimization.md index 77c5d96789..5626b03d2e 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_prefix_anchored_swa.md b/docs/design/sm70_prefix_anchored_swa.md deleted file mode 100644 index 3c0207f63b..0000000000 --- a/docs/design/sm70_prefix_anchored_swa.md +++ /dev/null @@ -1,57 +0,0 @@ -# SM70 Prefix-Anchored Sliding-Window Attention - -Prefix-anchored sliding-window attention is an explicit, model-agnostic -inference-engine option for bounded decode KV usage on V100/SM70. The request -prefix remains globally visible, while generated tokens attend only to the -most recent configured window: - -```text -visible(q, kv) = kv <= q and (kv < prefix_len or q - kv < window) -``` - -Enable it through `AttentionConfig`; checkpoint metadata and model architecture -names do not activate or gate the route: - -```bash -vllm serve \ - --attention-config '{"prefix_anchored_decode_window": 4096}' -``` - -The option is off by default. It changes attention semantics after generation -exceeds the configured window, so it is not an exact-output optimization for a -full-attention workload. - -## Current Engine Contract - -The engine admits the route only when all of these operator/runtime contracts -hold: - -- NVIDIA SM70 and the `FLASH_ATTN_V100` backend; -- causal, standard full decoder attention with fp16 KV; -- equal key and value head sizes, with no attention sinks or per-layer window; -- decode and prefill context-parallel sizes both equal to one; -- no speculative decoding, KV connector, or KV-cache offloading. - -Prefix caching is disabled, and full CUDA graphs are reduced to piecewise -graphs. Unsupported combinations fail during configuration or layer -construction. Once enabled, missing or mismatched per-request prefix metadata -is a hard runtime error; the engine never evicts gap blocks and then silently -runs an unmasked kernel. - -## Validation Status - -Validation on the current `main` source includes: - -- 32 CPU tests covering the mask, generic engine admission, fail-closed - metadata, cache-spec registration, and middle-gap block reclamation; -- a clean SM70 source build of the Flash-V100 extension; -- 19 V100 operator tests for paged decode and paged prefill, including GQA, - multiple head sizes, multi-partition decode, full-prompt and continuation - shapes, a negative control, and the default-off path. -- same-toolchain comparison against the current `origin/main` build: all 55 - default-off scalar decode kernel instances and all 64 default-off paged - prefill kernel instances are SASS-identical. - -No model end-to-end or throughput claim is attached to this route yet. Treat -performance and workload-level quality as separate acceptance gates before -enabling it in a production configuration. diff --git a/docs/design/sm70_qwen38_130_acceptance.md b/docs/design/sm70_qwen38_130_acceptance.md index 7da048b4c2..30fe226511 100644 --- a/docs/design/sm70_qwen38_130_acceptance.md +++ b/docs/design/sm70_qwen38_130_acceptance.md @@ -131,28 +131,22 @@ issue, not an idle-GPU or missing graph-shape issue. ## Wheel And Compatibility - Wheel: `1cat_vllm-1.3.0-cp312-cp312-linux_x86_64.whl`. -- Final release source: `6ada86ed64af6d1a7b3cb0f34df237fd86f06d48`. -- SHA256: `2bdb14a9c44f83ee6a766d88ed0d85b11390d6f5d65747e8dbe80a8e2d5d63e0` - (68,764,308 bytes). +- SHA256: `e4ffba3bf61e16dcba0eeeafac4036b7cead5abe047e4474ecacb7e5abb35da3` + (75,765,498 bytes). - A cloned environment imports vLLM 1.3.0 from site-packages with Torch 2.10.0+cu128. `_moe_C.moe_permute_sort_workspace_size` is present. - The wheel contains `_C`, `_C_stable_libtorch`, `_moe_C`, Flash-V100, vLLM FA2, and `flash_qla_sm70_gdn_strided.so`. - A clean-cache MTP4 plus FP8-KV boundary run does not create `TORCH_EXTENSIONS_DIR`, proving FlashQLA does not invoke NVCC at runtime. -- The final 2026-08-17 wheel uses source-controlled Python from the release - source above and the accepted SM70 native payload at source `383bbdd166`. - Every `.so` member and `vllm-rs` is byte-identical to that accepted payload; - the precompiled packaging path preserves both Flash-V100 and the SM70 - FlashQLA extension. In an isolated Python 3.12 install, all native modules - and FlashQLA load successfully; the MTP mixed-batch guard and Flash-V100 - policy suites pass 87 tests. -- The accepted runtime evidence remains Qwen3.8-27B-FP8 TP4 with the `fp8` - shorthand resolving to E5M2. In both no-MTP and MTP4 lanes, repeated - fixed-seed greedy responses were byte-identical and official sampling - stopped naturally without replacement characters. MTP's first request may - still JIT Triton helper kernels; this is a cold-latency concern, not a - numerical-output difference. +- The 2026-08-17 release candidate was packaged from the accepted SM70 build + at source `383bbdd166`, installed into a fresh Python 3.12 environment, and + served Qwen3.8-27B-FP8 TP4 with the `fp8` shorthand resolving to E5M2. In + both no-MTP and MTP4 + lanes, repeated fixed-seed greedy responses were byte-identical and official + sampling stopped naturally without replacement characters. MTP's first + request may still JIT Triton helper kernels; this is a cold-latency concern, + not a numerical-output difference. - Qwen3.6-35B-A3B-AWQ TP2 final-wheel checks pass. Official 1K/256 and 4K/1024 pure decode are 95.63 and 95.29 tok/s (10.457 and 10.494 ms TPOT), with complete, non-corrupted outputs and the TurboMind AWQ MoE route. diff --git a/docs/design/sm70_qwen38_fp8_prefill_5500.md b/docs/design/sm70_qwen38_fp8_prefill_5500.md deleted file mode 100644 index 8cf39d62ad..0000000000 --- a/docs/design/sm70_qwen38_fp8_prefill_5500.md +++ /dev/null @@ -1,142 +0,0 @@ -# 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..e36b263378 100644 --- a/docs/design/sm70_qwen38_fp8_prefill_decay.md +++ b/docs/design/sm70_qwen38_fp8_prefill_decay.md @@ -123,315 +123,8 @@ The extra address-generation/scalar-publication issue cost exceeds the saved P-load replay. Do not retry this layout spelling or the earlier shuffle/repack form. -## 2026-08-24 Exact-D256 K-Stage Ping-Pong - -The accepted exact-N32 arithmetic order is unchanged. The candidate alternates -K D64 panels between two non-overlapping shared-memory stages. The second K -stage borrows bytes from the later V/P allocation: K is dead before V is -loaded and P is materialized, so those lifetimes do not overlap. A next-K -fragment is published to the stage that the current QK phase is not reading, -so the pre-overwrite full-CTA barrier is no longer required. The visibility -barrier before the next QK phase remains. - -The K layout's logical size is 2048 half elements, but its padded physical -`cosize` is 2188. The final stage stride is therefore 2240 half elements: it -is larger than the physical span, 16-byte aligned for `STS.128`, and keeps -both stages on the same 128-byte shared-bank phase. The final lifetime-aliased -layout leaves dynamic shared memory at 45,568 bytes. Dense/split-KV3 kernels -remain at 254/253 registers per thread with zero stack or spill. - -SASS for the exact dense nonpaged specialization changes as follows: - -- static `BAR.SYNC.DEFER_BLOCKING` count: 15 to 12; -- HMMA step counts: 128 each, unchanged; -- `LDG`, `LDS`, `STS`, `SHFL`, and `MUFU` counts: unchanged; -- `FADD` and `IADD3` static counts: 15 to 12 and 71 to 68; -- full SASS from the clean repository-patch build is byte-identical to the - measured lifetime-alias artifact. - -The initial fully-disjoint v5 stage layout was screened across the length -curve. Same-build, separate-process A/B/A CUDA-event measurements use FP16, -`Hq=6`, `Hkv=1`, D256, causal attention, five warmups, five queued calls per -sample, and seven samples. The control column is the mean of the two bracketing -control medians except for the first 8K chunk, where the first process was a -visible clock-ramp outlier and the hot second control is reported. - -| Q | KV | Control | Ping-pong | Latency change | Exact | -|---:|---:|---:|---:|---:|---:| -| 8192 | 8192 | 4.8259 ms | 4.8302 ms | +0.09% | bitwise | -| 8192 | 32768 | 32.5750 ms | 32.2132 ms | -1.11% | bitwise | -| 8192 | 65536 | 68.2283 ms | 67.7216 ms | -0.74% | bitwise | -| 8192 | 131072 | 142.3325 ms | 140.9176 ms | -0.99% | bitwise | -| 15680 | 125440 | 251.6402 ms | 250.2695 ms | -0.54% | bitwise | - -A second 128K run from the clean v5 build measured -`142.1020 -> 141.5512 ms` (`-0.39%`) and was again bitwise exact. The final -lifetime-alias layout was then compared with v5 in both execution orders, -each bracketed by fresh controls. Across the two orders, pooled control, v5, -and lifetime-alias medians are `142.4068`, `141.8076`, and `141.7669 ms`: -`-0.42%` for v5 and `-0.45%` for the final layout. The `0.03%` difference -between candidates is noise-sized; the final layout is selected because it -does not enlarge shared memory or alter the original V/P addresses. - -Direct paged and split-KV3 gates at `Q4096/KV32768` are bitwise equal over -6,291,456 elements each; the dense gate covers 12,582,912 elements. The gain -is small but length-directed: the hot first chunk is neutral while long-prefix -calls consistently improve. - -Closed intermediate variants must not be repeated: - -- 2048-half stage spacing overlaps the 2188-half K physical span and changes - almost every output element; -- 2188-half spacing removes overlap but misaligns the second stage for - `STS.128` and raises a CUDA misaligned-address fault; -- 2192-half spacing is exact but loses address/bank-phase efficiency; -- 2304-half spacing is exact but is slower than the minimal same-phase 2240 - spacing; -- allocating the two 2240-half stages as fully disjoint storage is exact and - performance-equivalent, but unnecessarily grows shared memory to 46,336 - bytes; the lifetime-alias form retains the original 45,568-byte envelope. - -Nsight Compute recognized the exact mangled kernel but the host driver denied -performance-counter access with `ERR_NVGPUCTRPERM`. No privilege or machine -security setting was changed. CUDA-event timing and static SASS are therefore -the timing and structural authorities for this experiment. - -The matched end-to-end TP4 128K A/B/A gate also passes. It fixes Qwen3.8-27B -FP8 weights, 131,072 input tokens, 256 sampled output tokens, E5M2 KV, chunk -8192, no MTP, prefix caching, Mamba align, Flash-V100, CUDA graphs, and -`temperature=1.0/top_p=0.95/top_k=20/seed=20260824`. The control is the mean -of the two bracketing baseline runs. - -| Metric | Bracketed control | Lifetime alias | Change | -|---|---:|---:|---:| -| prefill | 47.9233 s | 47.6687 s | -0.531% | -| prefill throughput | 2735.0 tok/s | 2749.6 tok/s | +0.534% | -| TTFT | 47.9498 s | 47.6937 s | -0.534% | -| pure decode | 5.1557 s | 5.1547 s | neutral (+0.020% speed) | - -All three runs produce the same 256-token output hash -`2db7ef09...503325a`. Every rank reports the same route counts, including 32 -exact dense D256 calls, 544 prefix-paged calls, and 544 E5M2 bridge calls. -This promotes the small dependency-chain win, but it also establishes that -barrier removal alone is not a double-digit long-context solution. - -## 2026-08-24 Q8000 Split-KV3 Tail-Wave Experiment - -The matched runtime does not present `Q=8192` to long prefix-attention calls. -Although `max_num_batched_tokens` is 8192, Mamba/attention page alignment makes -the repeated full chunks `Q=8000`; a 40K diagnostic observed -`Q=(1,8000,6,256)` and `KV=(1,40000,1,256)`. The first Q8192 experiment was -therefore a useful operator screen but had zero end-to-end route hits. Its -first control measurement was also invalidated by a second TP4 job starting on -GPUs 4-7 between warmup and measurement (`47.90 -> 64.45 s`). The uncontended -candidate/control-B pair used the same unsplit route and differed by only -0.15%, confirming that run did not measure split-KV3. - -The corrected candidate admits exactly `Q=(1,8000,6,256)` behind the -default-off -`VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3_Q8000_EXPERIMENTAL` gate. It leaves -the existing production Q4096 policy unchanged and does not admit Q8192 or -other nearby shapes. Three independent KV partitions turn 750 Q tiles into -2250 CTAs at Q8000, avoiding the short last wave of the 72-SM V100. Bracketed -CUDA-event operator results from the final K-stage binary are: - -| Q | KV | Unsplit | Split-KV3 | Throughput change | Causal TFLOP/s | -|---:|---:|---:|---:|---:|---:| -| 8000 | 40000 | 40.1736 ms | 38.6519 ms | +3.94% | 44.05 -> 45.78 | -| 8000 | 64000 | 66.5214 ms | 64.6753 ms | +2.85% | 44.33 -> 45.60 | -| 8000 | 96000 | 102.8820 ms | 99.4545 ms | +3.45% | 43.95 -> 45.47 | -| 8000 | 128000 | 138.9150 ms | 134.3961 ms | +3.36% | 43.87 -> 45.35 | - -The three FP32 partial-output buffers plus max/sum state require 141.72 MiB -per rank. Outputs are repeat-stable but not bitwise equal to unsplit because -the partition merge changes reduction order. Across the four shapes, maximum -absolute difference is `1.526e-5` to `3.052e-5`; mean absolute difference is -`7.51e-7` to `1.38e-6`. - -The corrected full-model A/B/A uses the same TP4 128K official-sampling -contract as the accepted K-stage endpoint. Control A/candidate/control B -prefill times are `47.5940/47.1074/47.6612 s`. The bracketed control is -`47.6276 s` or `2752.0 tok/s`; split-KV3 is `47.1074 s` or `2782.4 tok/s`, -giving `-1.092%` latency and `+1.104%` throughput. TTFT improves 1.078%. -Decode remains neutral: `5.1533 -> 5.1617 s` and -`49.483 -> 49.402 tok/s`. Every rank reports 384 split-KV3 kernel hits, which -matches 12 eligible long chunks times 16 full-attention layers times two -requests. KV capacity remains 2,424,439 tokens per rank and no OOM fallback -occurs. - -Quality does not yet justify a default-on promotion. Every official-sampling -lane is internally stable between warmup and measurement, and the candidate -is token-for-token identical to control B for both 256-token outputs. However, -identical unsplit control A selects a different stable sampled stream despite -the same official seed. The candidate therefore shows no unique output -divergence, but this cross-process sampling instability cannot prove -equivalence for a non-bitwise route. - -A follow-up deterministic greedy gate used three exact-128K natural-text -requests: two identical Chinese prompts and one Python task. Both control and -candidate reproduce the duplicate prompt exactly within their own process, -and all outputs are coherent; both Python outputs stop normally and return the -same correct `add(a, b)` implementation. The two routes are not token-for-token -equal, though. The Chinese outputs first differ at token 7 and align for -126/128 tokens; the Python outputs first differ at token 20, with all 67 -control tokens aligned inside the 70-token candidate stream. Candidate prefill -times are `47.1360/47.2292/47.3380 s` versus control -`47.6704/47.8330/47.9274 s`. This unbracketed quality run is directionally -consistent with the formal A/B/A but is not a replacement performance claim. -All four candidate ranks report exactly 576 split-KV3 hits, matching three -requests times 12 eligible chunks times 16 full-attention layers; controls -report zero. - -The semantic smoke passes stability and basic task correctness, but strict -greedy identity fails. Keep Q8000 split-KV3 explicit and default-off until a -fixed-text logprob/perplexity or dataset-level quality gate establishes that -the classified Type-B reduction-order drift does not reduce model quality. - -## 2026-08-25 Q8000/KV40K..128K GQA-packed Attention Architecture - -This experiment freezes the final long-prefill call at causal FP16 -`Q=8000`, `KV=128000`, `Hq=6`, `Hkv=1`, and `D=256` on one -V100-SXM2-32GB. Its metric is 6.094872576 useful causal TFLOPs divided by the -complete Attention elapsed time. It is neither whole-model TOPS nor prompt -tokens/s; the acceptance gate is at most 101.581 ms, or at least 60.0 useful -causal TFLOP/s. - -Tile and barrier tuning of the fused BM64/BN32 kernel was already exhausted: -the long-shape NCU record showed one CTA/SM, 12.5% occupancy, 62.14% of cycles -without an eligible warp, and only 9.06% DRAM throughput. The replacement is a -different scheduling and dataflow architecture: - -- split a fully visible 32K..120K prefix from the exact causal 8K tail; -- pack all six GQA query heads into GEMM M=48,000; -- run wide SM70 Tensor-Core QK GEMMs that emit FP16 scores plus row max/sum; -- transform scores into normalized probabilities during the PV - global-to-shared load, avoiding a materialized probability matrix; -- fold one reusable FP16 PV partial into a FP32 online prefix - numerator/max/sum state; -- export max/sum from the accepted exact 8K tail and merge the two online - softmax states. - -The retained score block is BN8192. Final frozen-extension Torch A/B/A for the -score-cache layout measures `140.00213 -> 100.76979 ms`, or -`43.53414 -> 60.48313` useful causal TFLOP/s: 28.0227% lower latency and a -1.38933x speedup. All 12,288,000 outputs are finite; max/mean absolute error -is `2.2888e-4 / 2.5430e-5` and relative L2 is `6.8629e-3` versus the exact -FP32-accumulator route. BN7168 reaches only 58.42057 TFLOP/s and BN8000 only -59.92825 TFLOP/s, so neither is promoted or rounded into a pass. - -The same bounded workspace was then generalized across the twelve observed -8K-chunk shapes. Each row is a strict A/B bracket against the active split-KV3 -control: - -| KV tokens | split-KV3 ms | architecture ms | latency change | useful TFLOP/s | -|---:|---:|---:|---:|---:| -| 40,000 | 38.41638 | 29.81530 | -22.3891% | 59.34862 | -| 48,000 | 47.25811 | 36.38477 | -23.0084% | 59.44005 | -| 56,000 | 56.53999 | 42.68715 | -24.5010% | 59.87584 | -| 64,000 | 65.17897 | 49.14552 | -24.5991% | 60.00842 | -| 72,000 | 74.00670 | 55.44960 | -25.0749% | 60.27745 | -| 80,000 | 82.74432 | 61.77792 | -25.3388% | 60.46783 | -| 88,000 | 91.98507 | 68.37111 | -25.6715% | 60.38797 | -| 96,000 | 100.88397 | 74.72862 | -25.9262% | 60.51241 | -| 104,000 | 109.80642 | 81.09448 | -26.1478% | 60.61108 | -| 112,000 | 118.42901 | 87.71959 | -25.9307% | 60.51602 | -| 120,000 | 127.35232 | 93.99757 | -26.1909% | 60.65749 | -| 128,000 | 135.93617 | 100.34193 | -26.1845% | 60.74103 | - -All 147,456,000 candidate output elements are finite. Across the family, the -worst max absolute difference is `5.0354e-4` and worst relative L2 is -`6.8983e-3` versus the exact FP32-accumulator route. The three shortest points -remain slightly below 60 because the 43-TFLOP/s exact causal tail is a larger -fixed fraction, but they are still 1.288x..1.325x faster than split-KV3. - -The architecture trace is compute-dense rather than launch-starved: 49 timed -kernels occupy 95.630493 ms of a 95.809850-ms GPU span, or 99.81%. Prefix QK -accounts for 54.22% and prefix PV 39.74%; the exact tail is 4.89%. The next -compute ceiling is therefore QK/PV instruction efficiency, not additional -host launch fusion. - -Real-model integration exposed a separate memory contract. Retaining all -partials or the entire workspace passes the operator gate but leaves too -little memory for the following 80-MiB TP all-reduce. The final layout retains -only the 750-MiB FP16 score matrix per device and packs about 100 MiB of other -state into one transient byte slab. Before first use it requires enough -driver-visible free memory for both allocations plus 128 MiB of downstream -headroom; otherwise it raises a typed OOM before touching the caching allocator -and falls back to the exact route. - -The route is default-on for its exact engine contract. The legacy -`VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL=0` setting is the -rollback, and the route rejects CUDA graph capture. It admits only `Q=8000`, -`KV=40000..128000` in 8000-token steps, `Hq=6/Hkv=1/D=256`, FP16, causal -attention, and scale 1/16. The earlier -endpoint-only `.875` TP4 -Qwen3.8-27B-FP8 control/candidate/control prefill times are -`46.45658 / 45.47797 / 46.11195 s`. Relative to the `46.28426-s` bracketed -control, the candidate lowers latency by 1.7421% and raises prompt throughput -from 2831.89 to 2882.10 token/s (1.7729%). Every candidate rank logs 16 -architecture hits, with no architecture OOM/fallback, and all three runs emit -the same 32 output token IDs and SHA256 -`df4fee7f5f0126fe6b391fe77b4fc19667831de5ef55fd69c28c2f52a3d7086e`. - -The endpoint-only `.88` run also succeeds at 45.53427-s prefill with the same -hash and 16 hits/rank. Its memory profiler left enough headroom, so the route -passed rather than exercising the fallback branch. - -The widened TP4 Qwen3.8-27B-FP8 gate brackets the candidate with `46.11195-s` -and `46.09222-s` controls. Relative to their `46.10208-s` mean, the -`41.51191-s` candidate lowers prefill latency by 9.9565% and raises prompt -throughput from 2843.08 to 3157.46 token/s (11.0575%). Every rank logs exactly -192 family-route hits with no architecture OOM/fallback. Both controls and the -candidate emit the same 32 token IDs and SHA256 -`df4fee7f5f0126fe6b391fe77b4fc19667831de5ef55fd69c28c2f52a3d7086e`. - -The 2026-08-26 merge audit rebuilt the four-patch vendored FA2 stack from its -locked clean baseline with CUDA 12.8 and `sm_70`. The production QK/PV kernels -again compiled at 254/119 registers with zero spill. On the rebuilt extension, -the 40K endpoint measured `39.29395 -> 30.27579 ms` (`1.29787x`) with all -outputs finite, max/mean absolute difference `5.4932e-4 / 4.4392e-5`, and -relative L2 `6.6098e-3`. The 128K endpoint measured -`136.07628 -> 100.20284 ms` (`1.35801x`) with all outputs finite, -max/mean absolute difference `2.2888e-4 / 2.5484e-5`, and relative L2 -`6.9166e-3`. Both points pass the FP16 merge envelope of max absolute error -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. - ## Artifacts -- D256 GQA architecture artifacts are retained outside Git; their private - filesystem location is intentionally not committed. -- Final operator A/B/A: - `results/torch-architecture-scorecache-final-v1-aba.json` under that task - root. -- Shape-family operator A/B: - `results/torch-architecture-shapefamily-v1-aba.json` under that task root. -- Shape-family TP4 model gate: - `results/tp4-128k-architecture-shapefamily-{candidate,control-after}-0875.json`, - bracketed with the endpoint gate's - `results/tp4-128k-architecture-scorecache-final-control-b-0875.json`. -- Final TP4 `.875` model gate: - `results/tp4-128k-architecture-scorecache-final-{control-a,candidate,control-b}-0875.json`. -- Final high-memory `.88` route-success run: - `results/tp4-128k-architecture-scorecache-final-preflight-fallback-088.json`; - despite the retained filename, the preflight passed and no fallback ran. -- K-stage task root: - `/data/minimax-h3/task-cache/qwen38-fp8-128k-flashattention-20260824`. -- Final clean FA2 binary: - `build/formal-v6-final-cmake/_vllm_fa2_C.abi3.so` under that task root. -- Endpoint A/B/A JSON: - `results/tp4-128k-{baseline-a-v6,candidate-v6,baseline-b-v6}.json`. -- Q8000 operator sweep: - `results/splitkv3-q8000-kv{40000,64000,96000,128000}-v6.json`. -- Q8000 endpoint A/B/A: - `results/tp4-128k-splitkv3-q8000-{control-a,candidate,control-b}.json`. -- Q8000 deterministic quality gate: - `results/tp4-128k-splitkv3-q8000-quality-{control,candidate}.json` and the - matching files under `logs/`. -- Shape diagnostic and the zero-hit Q8192 negative run: - `logs/tp4-splitkv3-shape-diag.log` and - `results/tp4-128k-splitkv3-{control-a,candidate,control-b}.json`. - Root: task-local `qwen38-fp8-prefill-decay-20260815` artifact directory. - Nsight Systems: `profiles/qwen38-fp8-tp4-i128k-chunk15680-r2.nsys-rep`. diff --git a/docs/design/sm70_qwen38_nvfp4_decode.md b/docs/design/sm70_qwen38_nvfp4_decode.md deleted file mode 100644 index e937564413..0000000000 --- a/docs/design/sm70_qwen38_nvfp4_decode.md +++ /dev/null @@ -1,711 +0,0 @@ -# SM70 Qwen3.8-27B-NVFP4 Decode Recovery - -Date: 2026-08-23; acceptance updated 2026-08-24 - -## Decision - -The accepted Qwen3.8-27B-NVFP4 TP4, no-MTP decode route keeps every eligible -FP8 and NVFP4 projection on an SM70 native path. Channel-FP8 projections use -the memory-neutral QPN8 decode kernels for the accepted gate/up, down, and -output shapes. Remaining channel-FP8 projections, including the LM head, use -TurboMind W8A16. NVFP4 projections use TurboMind W4A16. No accepted result in -this document credits a Marlin fallback. - -The recovery route also adds E4M3 KV support to Flash-V100 XQA and selects a -p64 partition for the exact batch-one, G6/D256 layout. This is the change that -makes the frozen native-sampler result exceed 70 tok/s. The p64 route is -restricted to E4M3, batch one, G6/D256; the existing E5M2 and FP16 policies -are unchanged. The later no-MTP 80 tok/s acceptance adds exact QPN4, TP4 -all-reduce, QPN8 specialization, and an exact-Philox chunked sampler sidecar; -that acceptance is recorded at the end of this document. - -## Frozen Contract - -- Model: Qwen3.8-27B-NVFP4, config SHA256 - `1b3c71868d1299e52df6fc907deb202d5132b1ef0f72aae0ef6d15185dd53a5c`. -- Hardware/runtime: four V100-SXM2-32GB GPUs, TP4, Python 3.12, Torch - 2.10.0+cu128, CUDA 12.8. -- Engine: compressed-tensors, FP16 activation/output, E4M3 KV, - `FLASH_ATTN_V100`, `max_model_len=262144`, `max_num_batched_tokens=8192`, - `max_num_seqs=1`, prefix caching, aligned Mamba cache, chunked prefill, and - `FULL_AND_PIECEWISE` CUDA graphs. -- Request: input 1024, natural output cap 256, no `ignore_eos`, no MTP. -- Sampling: temperature 1.0, top-p 0.95, top-k 20, request seed 20260815, - engine seed 0. The 2026-08-23 recovery baseline uses the unmodified native - sampler. The 2026-08-24 acceptance uses the separately registered, - exact-Philox chunked sampler described below. -- Pure decode uses the 255 steady token intervals and excludes TTFT/prefill. - -## Route Map - -| Layer family | TP-local shape | Accepted decode route | -|---|---|---| -| FP8 gate/up | K5120 x N8704 | fused QPN8 gate/SiLU/up | -| FP8 down | K4352 x N5120 | QPN8 split-16 | -| FP8 output | K1536 x N5120 | QPN8 split-12 | -| FP8 GDN input/full-attention QKV and LM head | model shapes | TurboMind W8A16 | -| NVFP4 gate/up | K5120 x N8704 | TurboMind N32, lookahead 1 | -| NVFP4 down | K4352 x N5120 | TurboMind N32, lookahead 2 | -| E4M3 G6/D256 attention | B1, page 1568 | Flash-V100 XQA p64 | - -Channelwise `[N,1]` FP8 scales are admitted and packed without retaining a -second permanent weight copy. The existing QPN8 model/shape/concurrency gates -remain in force. The NVFP4 selector is exact-shape gated by -`VLLM_SM70_NVFP4_QWEN38_TP4_M1_FAST_SELECTOR`; setting it to zero restores -dynamic tuning. Setting `VLLM_FLASH_V100_DECODE_PARTITION_SIZE=128` restores -the previous E4M3 attention partition. - -## Pure Decode Result - -| Milestone | Steady decode | TPOT | -|---|---:|---:| -| Original mixed checkpoint route | 29.35 tok/s | about 34.07 ms | -| Channel-FP8 TurboMind | 57.285 tok/s | 17.457 ms | -| E4M3 XQA p256 | 62.615 tok/s | 15.971 ms | -| Exact channel-FP8 QPN8 | 65.522 tok/s | 15.262 ms | -| E4M3 XQA p128 | 67.880 tok/s | 14.732 ms | -| NVFP4 N32 selectors | 69.991 tok/s | 14.288 ms | -| Clean final binaries, native sampler, E4M3 p128 | 69.904 tok/s | 14.305 ms | -| Clean pre-merge binaries, native sampler, E4M3 p64 | 71.732 tok/s | 13.941 ms | -| Merged-source confirmation, native sampler, E4M3 p64 | **71.342 tok/s** | **14.017 ms** | - -Both p64 requests generated all 256 tokens, contained no EOS, and finished by -length. Relative to the immediately preceding clean-p128 control, the -pre-merge p64 run saves 0.365 ms/token and improves pure decode by 2.61%. -After merging `onecat/main` at `675a12dedc`, rebuilding Flash-V100, and -repeating the frozen request under an exclusive four-GPU reservation, p64 -measures 71.342 tok/s at 14.017 ms/token. This is the conservative 2026-08-23 -baseline and is superseded by the 2026-08-24 acceptance below. - -Sampled token identity is not an attention quality criterion: one-output-ULP -changes can flip a low-margin random sample. The final p64 stream is coherent -and, independently, reproduces an earlier accepted full-length stream. The -operator gate below is the numerical criterion. - -## Numerical and Operator Gates - -The E4M3 XQA p64/p128 race uses page size 1568, G6/D256, checkpoint-style -E4M3 bytes, K/V scales 0.04/0.25, and the scalar paged decoder as reference. - -| Sequence length | p64 | p128 | p64 gain | Maximum error | -|---:|---:|---:|---:|---:| -| 1025 | 46.633 us | 63.614 us | 16.981 us | 7.63e-6 | -| 1152 | 45.896 us | 59.346 us | 13.450 us | 7.63e-6 | -| 1280 | 46.072 us | 55.625 us | 9.553 us | 7.63e-6 | -| 2049 | 45.807 us | 56.751 us | 10.945 us | 7.63e-6 | - -Every p64 and p128 output is within one representable FP16 output ULP of the -scalar reference. The focused GPU regression covers p64/p128/p256 with unit -and non-unit K/V scales and passes 6/6. - -The retained NVFP4 selectors also passed independent FP32-oracle checks: -relative L2 is `2.701e-4` for gate/up and `2.924e-4` for down, with cosine -approximately one. The QPN8 source already passed its model-quality and -operator gates recorded in `sm70_qwen38_qpn8_decode.md`. - -## Per-Token Profile - -The latest short Nsight Systems node trace uses the merged p64 route. It is -composition evidence; its 14.635 ms traced TPOT is not used as the absolute -speed result because the unprofiled frozen request measures 14.017 ms. The -trace captures 63 decode replays on each TP rank and uses the 61 middle steps -for steady statistics. Graph-node kernel coverage is 93.48%. - -| Component | Time/token | -|---|---:| -| TurboMind NVFP4 gate/up | 2.762 ms | -| QPN8 split-16 projections | 2.234 ms | -| TurboMind NVFP4 down | 1.706 ms | -| TP all-reduce | 1.697 ms | -| QPN8 split-12 projections | 0.895 ms | -| E4M3 XQA p64 | 0.587 ms | -| QPN8 fused gate/up | 0.508 ms | -| TurboMind FP8 dense/LM head | 0.407 ms | - -Across rank-token samples, the replay interval is 14.621 ms and GPU activity -union is 14.055 ms, or 96.136% of the interval. Idle gaps total 0.565 ms/token. -The stream still launches 1140.8 kernels per rank per token; half are shorter -than 5 us. The grid-limited static occupancy ceiling places 69.40% of service -below 25% occupancy and 28.50% at 25-50%. Nearly continuous GPU work therefore -does not imply high useful compute utilization: batch-one launch geometry and -the serial projection/communication chain remain the main limit. - -Fifty-millisecond NVML samples report 99.47% mean GPU busy-window duty and -48.26% memory-active-window duty. Per-GPU power averages 175.9-180.2 W and -peaks at 221.4 W; runtime allocation peaks at 29.06-29.36 GiB/GPU. Model loading -accounts for 5.77 GiB/GPU, while the configured cache budget reports 21.29 -GiB/GPU; the remaining allocation includes CUDA graphs, workspaces, state, and -runtime context. - -NVML memory duty is not achieved HBM bandwidth. Current trace durations imply -effective minimum packed-weight rates of 451.2 GB/s for NVFP4 gate/up, 365.1 -GB/s for NVFP4 down, and 717.4 GB/s for QPN8 down, with useful arithmetic rates -of 1.805, 1.461, and 1.435 TFLOP/s/GPU respectively. These omit scales, -activation/output traffic, caches, and implementation-internal work. A current -NCU capture was attempted, but the host rejected non-root performance-counter -access with `ERR_NVGPUCTRPERM`; exact current SM, Tensor Core, occupancy, and -DRAM counters therefore require administrative profiling permission. The -previously accepted QPN8 counter evidence remains in -`sm70_qwen38_qpn8_decode.md`, but it is not credited as current NVFP4 or p64 -XQA counter evidence. - -## Rejected Experiments - -- NVFP4 N64 and K32 candidates lost to N32 and were removed. The retained - exact shapes use split 3/swizzle 4 with shape-specific lookahead. -- QPN8 split 17/20/24 lost. Split 12 won only for K1536 x N5120, saving about - 12.7 us/token across the 64 output projections. -- The first compact top-k20 Python sampler and first fused CUDA top20 - candidate did not preserve native Philox sampling and were removed. Both - 2026-08-23 p64 results use the native sampler. They are distinct from the - later canonicalized, generator-state-preserving implementation accepted on - 2026-08-24. -- Replacing only the full-vocabulary sort while retaining both native - full-vocabulary softmax operations saved 53.9 us in isolation, insufficient - for a stable acceptance margin. A one-full-softmax hybrid stayed - experimental because the p64 attention route solved the target without - changing sampling math. - -## Build, Tests, and Evidence - -- Final `_C` SHA256: - `e0ea14d0e40330b08a9951e67634e50b722540d5d1bbb48500f690323ab07624`. -- Final p64 Flash-V100 SHA256: - `b418fed86b9c1ab9297c8795c24732818239b9a3aaca5ec9efb60933853d8ce7`. -- Focused CPU policy/dispatch tests: 12/12 passed (nine FP8/QPN8 and three - E4M3 p64 policy cases). -- Focused E4M3 XQA GPU numerical tests: 6/6 passed. -- Ruff lint/format, Python byte compilation, shell syntax, and - `git diff --check` pass for the changed files. -- Retained task evidence is under - `.artifacts/qwen38_nvfp4_speed_20260823/`, notably the final p128/p64 JSON - results, the merged confirmation - `final_merged_qwen38_nvfp4_tm_e4m3_xqa_p64_native_sampler_full_graph_i1k_o256.json`, - `e4m3_xqa_p64_vs_p128_clean.json`, and the parsed per-token Nsight tables - under `profiles/`. - -## 2026-08-24 No-MTP 80 tok/s Acceptance - -### Accepted deployment and route - -The accepted deployment is deliberately a three-binary composition. It uses -the compatible primary `_C` module with SHA256 -`a0a0cd9ddeccc73fa3d920c7a869450c4b33d001f97637c35b75b966d89ad36d`, -the production CMake C++17 sampler sidecar with SHA256 -`cdbfdd87dfa9119e52acc88d5787063202561a879a63334145a5616344a549ae`, -and Flash-V100 p64 with SHA256 -`b418fed86b9c1ab9297c8795c24732818239b9a3aaca5ec9efb60933853d8ce7`. -MTP remains disabled. - -The exact FP8 gate/up, down, and output shapes stay on the proven-faster QPN8 -route. GDN input, full-attention QKV, and the LM head stay on TurboMind W8A16. -The accepted batch-one NVFP4 shapes use the native QPN4 decode route; unsupported -or non-admitted cases retain the TurboMind fallback. Thus every eligible FP8 -sublayer still uses TurboMind unless the frozen shape has a measured faster -QPN8 specialization. - -The sampler sidecar replaces a full-vocabulary sampling fragment with an -80-chunk top-20 reduction while preserving canonical sort order, top-p math, -Philox draws, and generator state. It is built separately so sampler operator -registration does not relink the quality- and speed-frozen primary `_C`. -`setup.py` extracts either CPython-SOABI or abi3 sidecar names, and -`vllm/_sm70_ops.py` loads the bundled module or the explicitly configured -`VLLM_SM70_SAMPLER_LIBRARY`. - -### Speed result - -All absolute results use the frozen TP4, input-1024/output-256 contract and -255 steady decode intervals. TTFT and prefill are excluded. - -| Binary/measurement | Steady decode | TPOT | Disposition | -|---|---:|---:|---| -| Merged native-sampler baseline | 71.342 tok/s | 14.017 ms | baseline | -| Hand sidecar repeat A2 | 80.177 tok/s | 12.472 ms | accepted | -| Hand sidecar repeat A3 | 80.164 tok/s | 12.474 ms | accepted | -| Hand sidecar repeat A4, physical GPUs 4-7 | 80.026 tok/s | 12.496 ms | accepted | -| Production CMake C++17 sidecar A2 | **80.624 tok/s** | **12.403 ms** | accepted | - -The hand-sidecar accepted repeats preserve the frozen 256-token SHA256 -`0b2d335ddce9b282e45eea1b6c86525bc61eeb5ba1655e8228e6ef3bd1ce823b`. -The production C++17 sidecar changes a low-margin full-model random trajectory, -so random token identity is not treated as the sole quality gate. A direct -cross-build sampler diagnostic covers 100 independent seeds and 256 sequential -draws: hand and production sidecars produce identical token selections and -identical final generator-state SHA256 -`4257d205138503840fea92fe6b15ddfa276df82c315513da4bbd785393c98c96`. - -### Output-quality gate - -The quality contract is the frozen first 250 GSM8K test questions, five-shot -prompting, greedy generation, maximum 256 output tokens, and per-question -record retention. It is intentionally independent of the random performance -prompt. - -| Route | Correct | Accuracy | Invalid | -|---|---:|---:|---:| -| Frozen pre-optimization baseline | 226/250 | 90.4% | 0 | -| Hand-sidecar candidate | 227/250 | 90.8% | 0 | -| Production CMake C++17 sidecar | **226/250** | **90.4%** | **0** | - -The production result therefore equals the frozen accuracy baseline with no -invalid outputs. Its JSON contains all questions, generated texts, extracted -answers, labels, correctness flags, token IDs, and per-item output hashes. The -result file SHA256 is -`143b2b345b830a858bff2568f9cf51148c8185acb42bb129e2dfc2d421a196d2`. - -### Accepted-route trace and resource use - -The latest Nsight Systems trace uses the accepted compatible main library and -the behavior-equivalent hand sidecar. It captures 63 graph replays on each of -four ranks and reports the middle 61 steps. Its 13.465 ms node-traced TPOT is -composition evidence; the unprofiled 12.474 ms result remains the absolute -speed evidence. Graph-node kernel coverage is 94.88%. - -| Critical component | GPU service/token | -|---|---:| -| QPN8 split-16 projections | 2.213 ms | -| QPN4 fused gate/up | 2.165 ms | -| TP4 pack32 all-reduce | 1.587 ms | -| QPN4 down | 1.441 ms | -| QPN8 output projection | 0.915 ms | -| Flash-V100 E4M3 XQA p64 | 0.618 ms | -| QPN8 fused gate/up | 0.496 ms | -| TurboMind FP8 dense/LM head | 0.410 ms | - -The mean replay interval is 13.430 ms and GPU activity union is 12.887 ms, -or 95.958%; idle gaps total 0.543 ms/token and the largest mean gap is only -9.538 us. The route still launches 1061.9 kernels per rank per token. The -grid-limited occupancy ceiling assigns 38.46% of service below 25% occupancy, -47.86% at 25-50%, and 13.63% at 50-75%, showing that batch-one kernel geometry -and the serial projection/communication chain remain the main headroom. - -NVML's coarse windows report 100% GPU-busy duty, 53.25% memory-active duty, -and 56.36% of the 300 W power limit. Runtime memory is about 28.97 GiB/GPU; -SM and memory clocks are 1530 and 877 MHz. These are duty indicators, not -achieved SM, Tensor Core, or HBM throughput. Nsight Compute counters remain -unavailable because the driver returns `ERR_NVGPUCTRPERM`. A payload-only -model estimate is about 491 GB/s/GPU and useful arithmetic is about 4.33 -TFLOP/s across TP4, but neither value is a hardware-counter measurement. - -### Reproducibility boundary and evidence - -The production sidecar is reproducible through the CMake target -`_sm70_sampler_C` and its focused microbenchmark selects exactly the reference -tokens across five distributions, 100 explicit-seed trials, and 100 default -generator trials. It reduces the measured fragment from 102.427 to 62.972 us. -The final Python regression run passes all 94 tests in the sampler, NVFP4 -admission, and SM70 TurboMind adapter files. Ruff lint/format, Python byte -compilation, the sidecar wheel-name gate, and `git diff --check` also pass. - -Freshly relinking the primary `_C` was tested separately. The first cold run -measured 77.314 tok/s, while the matching second run measured 80.181 tok/s at -12.472 ms/token. Both used a current-source C++17 sampler sidecar and produced -the exact same 256 token IDs as the accepted production C++17 run. The earlier -roughly 77-80 tok/s spread was therefore startup-state variation, not a decode -route or output regression. - -The clean source snapshot is `7bd1d783e4861af5f5396a721156c44f623f88b0`. -Its primary `_C` SHA256 is -`c5d7c6a9a0fa714c1e97427c3d16369404ae5af6ee80f48b165595081cc43e56`, -and its C++17 sampler SHA256 is -`3058051f59949d88f5982d0c5bd80121f3c75a0f5307ca901fa6aebdec4cbc74`. -The primary and sampler `.nv_fatbin` SHA256 values are respectively -`dbba740219f00db9a528310fa84c0defb97c03f5cf99b519af5975609eb926e6` -and -`9a89e8cc22d1f8e36378005c08e2b32b315d79728989d6b4b2fe44fedfa5cb2c`; -both are byte-identical to the accepted binaries. The full primary ELF differs -because merged main adds two unrelated ModelOpt NVFP4 MoE host entry points, -and full sidecar hashes include host build-ID differences. Those differences -do not enter this dense batch-one graph. - -The clean second-run result SHA256 is -`4e6ed3306772efca1a9c84b3681d50a268c644b8d6327df0fab2341581236bb1`. -Its serialized token-array SHA256 is -`f50e0ebab44ae350c7921bf91ba709fa896dcfb253c06842323b80de5b8bde32`, -exactly matching the accepted C++17 result. It records TP4, physical GPUs 0-3, -custom all-reduce enabled, both FP8 and NVFP4 TurboMind enabled, -`speculative_config=None`, and `spec_decoding_metrics=null`. Because both -device binaries and the full sampled stream are identical, the existing -production C++17 GSM8K result, 226/250 with zero invalid outputs, remains the -quality gate; repeating the same 250-question run would add no new path -coverage. Keeping the sampler in its own sidecar remains the packaging -boundary, but a current-source primary rebuild is now accepted. - -CMake component installation was also checked without rebuilding unrelated -wheel targets. It installs `vllm/_C.abi3.so` and -`vllm/_sm70_sampler_C.cpython-312-x86_64-linux-gnu.so`, removes their build-tree -RPATHs, and preserves both accepted fatbin hashes. The installed-file SHA256 -values are `d7ed490ef41e55f3858a45b3a92d62b2d527343fa14b04053cbe65ba16646749` -and `8ee8466359dc4432234dc3daa92918b2226879caad563d23a85b980c2a7c755e`. -A loader probe against the staged files registered the QPN4, QPN8, and -`_C::sm70_sample_packed_top20_out` schemas successfully. A broad wheel rebuild -of unrelated extensions is not required to establish this route's packaging -boundary. - -Primary evidence paths are: - -- `.artifacts/qwen38_nvfp4_speed_20260823/results/candidate_qpn4_fold_qpn8_m1_arpack32_official_cxx17_sidecar_chunk80_i1k_o256_a2.json` -- `.artifacts/qwen38_nvfp4_speed_20260823/results/candidate_qwen38_nvfp4_clean_head_cxx17_chunk80_i1k_o256_a2.json` -- `.artifacts/qwen38_nvfp4_speed_20260823/results/qwen38_nvfp4_gsm8k_250_official_cxx17_sidecar.json` -- `.artifacts/qwen38_nvfp4_speed_20260823/results/sampler_stream_hand.json` -- `.artifacts/qwen38_nvfp4_speed_20260823/results/sampler_stream_cxx17.json` -- `.artifacts/qwen38_nvfp4_speed_20260823/wheel_stage_clean_head_7bd1d783/` -- `.artifacts/qwen38_nvfp4_speed_20260823/profiles/candidate_qwen38_nvfp4_sidecar_chunk80_nsys_nvml_i1k_o64_per_token.md` -- `.artifacts/qwen38_nvfp4_speed_20260823/profiles/candidate_qwen38_nvfp4_sidecar_chunk80_resource_summary.md` - -## Quality-safe GDN decode fusion - -The next no-MTP batch-one optimization keeps the mixed checkpoint route map -unchanged and removes work at the GDN boundary. The exact Qwen3.8 TP4 decode -shape can issue its channel-FP8 QKV/Z projection and FP16 b/a projection from -one QPN8 launch, writing the four destination tensors directly. A separate -one-pass Triton kernel performs the 12-by-128 gated RMSNorm without changing -the accepted arithmetic. The accepted pair is opt-in through both -`VLLM_SM70_GDN_QPN8_BA_SPLIT=1` and -`VLLM_SM70_GDN_RMSNORM_ONEPASS=1`; the split projection cannot activate -without its paired RMSNorm gate. Other operator shapes and prefill retain the -existing operators. These are hardware/layout/operator gates; the measured -checkpoint identifies the evidence workload and does not select the route. -The archived JSON predates this audit naming cleanup and records the former -Qwen-prefixed flag names; kernel arithmetic is unchanged. The final source -also rechecks the loaded QPN8 code/scales, workspace pointer, bias, dtype, -contiguity, and same-device contracts before dispatch. - -The matched endpoint used the same four V100-SXM2-32GB GPUs, TP4, input 1024, -output cap 256, official random sampling, E4M3 KV, Flash-V100, full CUDA graph, -and no MTP. Each result is the stable median of three sequential requests and -pure decode excludes TTFT/prefill. - -| Route | Pure TPOT | Steady decode | Output gate | -|---|---:|---:|---| -| Current control | 12.249652 ms | 81.634970 tok/s | frozen 256-token stream | -| QPN8 QKV/Z plus b/a only | 11.941616 ms | 83.740763 tok/s | reject: first divergence at token 202 | -| QPN8 QKV/Z plus b/a and one-pass RMSNorm | **11.921648 ms** | **83.881019 tok/s** | pass: 3/3 exact | - -The accepted pair improves steady decode by 2.751% and TPOT by 2.678% versus -the matched control. All three accepted requests reproduce the control's full -256-token SHA256 -`8b37337f4c393711cb8550db6bae909b1e85de8df1cf5ba8c60d8c000749c0a2`. -The one-pass RMSNorm 48-layer microbenchmark is bitwise exact and saves -55.946 us/token. The QKV projection is exact on all layers; b/a differs only -at normal FP32-reduction roundoff (`6.368e-8` relative L2, maximum absolute -error `2.384e-7`) and the paired endpoint restores exact sampled output. - -The post-rebase admission recheck uses the final source-built extension and -the default Qwen3.5 model wrapper rather than a diagnostic GDN boundary. The -C++ split-route oracle fired on all four TP ranks. Its matching latest-main -control/candidate stable medians are `12.261727/11.956949 ms` and -`81.554583/83.633383 tok/s`, a 2.549% decode-throughput improvement. All -three candidate requests exactly match all three latest-main control -256-token streams with SHA256 -`ca77db3b032a1600a8567adea706108c1bd8c5472b3ace2318c41ab17c66c1f9`. - -A fused Q/K RMSNorm, MRoPE, and KV-write experiment reached 85.451 and -85.015 tok/s on its two stable repeats but is rejected: it diverged from the -control at token 1 and stopped after 208 output tokens. That source and its -route flag are intentionally absent from the submitted change. - -Primary local evidence paths are: - -- `.artifacts/qwen38_nvfp4_speed_20260823/results/gdn_endpoint_control_r3_i1k_o256.json` -- `.artifacts/qwen38_nvfp4_speed_20260823/results/gdn_endpoint_ba_split_rms_r3_i1k_o256.json` -- `.artifacts/qwen38_nvfp4_speed_20260823/results/current_main_gdn_control_r3_i1k_o256.json` -- `.artifacts/qwen38_nvfp4_speed_20260823/results/current_main_gdn_route_recheck_r3_i1k_o256.json` -- `.artifacts/qwen38_nvfp4_speed_20260823/results/gdn_rmsnorm_production_op_bound_48layers.json` -- `.artifacts/qwen38_nvfp4_speed_20260823/results/qpn8_fused_ba_split_bound_48layers.json` -- `.artifacts/qwen38_nvfp4_speed_20260823/results/qknorm_mrope_cache_gdn_ba_rms_r3_i1k_o256.json` - -## 2026-08-25 Long-context E4M3 page800 decode - -### Scope and route - -This optimization targets the exact SM70, E4M3, grouped-query decode route -used when aligned Mamba state makes the attention page 800 tokens. The -production KV allocation is interleaved as `[blocks, 2, 800, 1, 256]`; its K -and V views each have strides `[409600, 256, 256, 1]`. Admission additionally -requires G6/D256 and the existing dual-CTA batch route. Other page sizes, -layouts, head dimensions, KV dtypes, batch routes, and prefill keep the prior -implementation. - -The specialized kernel uses the compile-time 800-token page and physical -block stride, replaces repeated general page division with the bounded -one/two-page comparison for a 256-token partition, and loads adjacent shared -V values as aligned `half2`. Each output dimension retains the same token and -FMA order as the fallback. The strict layout predicate is checked before -dispatch; setting `VLLM_FLASH_V100_E4M3_PAGE800_FASTPATH=0` restores the -general-stride kernel. The fast path is enabled by default on an admitted -route. `VLLM_FLASH_V100_E4M3_PAGE800_FASTPATH_TRACE=1` emits a once-per-process -route diagnostic. - -The measured contract uses Qwen3.8-27B-NVFP4, four V100-SXM2-32GB GPUs with -TP4, compressed-tensors NVFP4 weights, FP16 activations, E4M3 KV, Flash-V100, -FlashQLA, prefix caching, aligned Mamba cache, FULL decode graphs, and no MTP. -The workload is NVIDIA SPEED-Bench revision -`487aa718444e816458d1a0a52bfce7a454285cf4`: the official 32K `low_entropy` -JSONL has SHA256 -`c3ee8a3f63b6cce18d063a2ff9992b6b96cb72c2b6408970fa78d2335b542f8e`. -The explicitly derived 64K workload pairs two official rows and has SHA256 -`0dcc877ecc7ccf66450c1f168ff7493dcf50f52da89d321dba394144a70e9800`. -Performance sampling is temperature 1.0, top-k 20, top-p 0.95, request seed -20260822, natural EOS, and a 512-token cap. Pure decode excludes queueing, -TTFT, and prefill; each cell has one warmup and three measured repeats. - -### Long-context result - -The matched rollback/candidate A/B was measured on source `05d5aa4e57`. B1 -references are 52.052 token/s at 32K and 40.152 token/s at 64K. The accepted -16K scaling references are 70.736% at B8 and 52.281% at B16. - -| Context | Batch | Rollback token/s | Candidate token/s | Candidate TPOT | Gain | Efficiency vs B1 | -|---:|---:|---:|---:|---:|---:|---:| -| 32K | 8 | 292.03 | **326.47** | 24.505 ms | **11.79%** | 78.40% | -| 32K | 16 | 392.84 | **451.92** | 35.405 ms | **15.04%** | 54.26% | -| 64K | 8 | 209.94 | **244.51** | 32.719 ms | **16.47%** | 76.12% | -| 64K | 16 | 260.36 | **311.74** | 51.325 ms | **19.74%** | 48.52% | - -B8's two-context geometric-mean gain is 14.11%, B16's is 17.36%, and the -four-cell geometric mean is 15.72%; no cell regresses. Relative to the 16K -references, 64K loses no efficiency at B8 and 3.76 percentage points at B16, -within the ten-point budget. Every one of the 144 candidate performance -requests completed with its full 512 tokens and without empty or replacement- -character output. Route logs show FULL graphs and the expected long-context -`dual_cta_split` B8 and `dual_cta` B16 dispatches on all four ranks. - -The endpoint reports that the exact D256 prefill extension is unavailable and -uses its established fallback. That warning is common to rollback and -candidate and is excluded from the pure-decode metric. This result must not be -cited as a prefill improvement. - -### Counter and correctness evidence - -An Nsight Compute rollback/candidate comparison on B16, 64K, p256 attributes -the endpoint gain to less address/control and shared-memory pressure rather -than added occupancy. Kernel duration falls 22.47%, executed instructions -11.92%, shared-bank conflicts 48.16%, MIO-throttle stalls 95.92%, long- -scoreboard stalls 8.61%, and wait stalls 11.06%. Eligible warps per active -cycle rise 0.565 to 0.699. Peak-relative DRAM and SM utilization rise from -24.49% to 28.11% and from 39.11% to 44.79%, respectively; registers rise only -from 166 to 168 and dynamic shared memory is unchanged. Thus the retained -change follows the measured copy/addressing bottleneck and does not depend on -an occupancy increase. - -The production-layout operator A/B is bitwise equal to the rollback path for -B8/B16 and p64/p128/p256. At p256, its direct kernel gains are 22.10% at -32K/B8, 21.37% at 32K/B16, 22.59% at 64K/B8, and 22.84% at 64K/B16. The two -focused GPU suites pass 15/15 and 14/14; the production interleaved cases use -`torch.equal`, not a tolerance. All 55 environment tests, Ruff lint/format, -the SM70 extension build, and `git diff --check` also pass. - -Because the production path uses official stochastic sampling, independent -server processes are not expected to reproduce every generated text. The -quality gate therefore pairs rollback and candidate over the same 32 GSM8K -main/test prompts at B8 and B16, temperature 1.0/top-k 20/top-p 0.95, natural -EOS, and sampling seeds 20260822, 20260823, and 20260824. This protocol was -fixed before the two additional seeds ran; both sides were rerun rather than -repeating only the candidate. - -| Batch | Rollback numeric EM | Candidate numeric EM | Delta | -|---:|---:|---:|---:| -| 8 | 71/96 | **74/96** | +3 | -| 16 | 74/96 | **74/96** | 0 | - -All 384 quality requests succeed, and the candidate has no empty or Unicode -replacement output. The first seed alone was one item lower in each batch; -the predeclared three-seed aggregate is recorded to make that sampling -variance visible rather than discarding it. The final source patch rebases -without conflict onto `onecat/main@d62ef5cb20`. Its rebuilt Flash-V100 ELF has -SHA256 `70816f9655e36fcd04b07d75be1b3abafabc81a30c261d4fd99b06389cbf82a2`. -The full SM70 SASS dump is byte-identical to the measured binary, with SHA256 -`b30031fa4f29181b8e1c757bc408c608bae97df380893e613b96a2cf7d3e9a51`; -the differing full-ELF hash is build metadata, not device code. - -## 2026-08-25 No-MTP Long-Context Decode Acceptance - -### Contract and route - -The long-context gate keeps the accepted Qwen3.8-27B-NVFP4 computation and -sampling paths unchanged. It uses four V100-SXM2-32GB GPUs with TP4, -compressed-tensors weights, FP16 activations/output, E4M3 KV, TurboMind for -the admitted FP8/NVFP4 projections, Flash-V100 attention, full CUDA Graph, -prefix caching, aligned Mamba state, and no MTP. Each request ends at an exact -128K or 256K context after 64 generated tokens. Pure decode is the 63 steady -token intervals and excludes prefill and TTFT. The acceptance thresholds are -60 tok/s at 128K and 45 tok/s at 256K. - -The new Flash-V100 route is restricted to batch one, G6/D256, 1,568-token -interleaved paged KV, and E4M3. It selects p64 below 12,288 tokens, p256 below -49,152, p512 below 98,304, p896 below 196,608, and p1664 thereafter. A single -long-wave partition kernel and matching reducer select p512/p896/p1664 from -device-resident sequence lengths, so CUDA Graph replay does not need a host -readback or three long inactive launch pairs. The existing short p64/p256 -launch remains separate. Other batch sizes, KV formats, head layouts, page -sizes, explicit partition overrides, and attention routes are unchanged. - -Full CUDA Graph capture has two semantic batch-one variants at the p512 -boundary. Contexts below 49,152 tokens replay the default graph, which contains -only the p64/p256 launch and reducer pair. Contexts at or above 49,152 tokens -replay the long-context graph, where `batch_context_routing` admits the wave -launch and reducer pair. The dispatcher uses its existing host scheduler -context length to select an already-captured graph; the attention kernel still -reads the live sequence length from device memory. This keeps the 1K decode -graph free of inactive long-wave launches without adding a device-to-host -sequence-length readback. - -E4M3 conversion uses a 256-entry exact half-bit lookup table initialized in -CTA shared memory. It adds 512 bytes of dynamic shared memory and avoids the -long integer conversion chain on every KV element. The final SM70 cubin uses -168 registers/thread and 32 bytes of stack for the 192-thread long-wave -partition kernel; total dynamic shared memory is about 44.8 KiB, retaining two -CTAs/SM. Its reducer uses 32 registers/thread and 80 bytes of static shared -memory. The route is enabled with -`VLLM_FLASH_V100_XQA_E4M3_G6_P64_P256_AUTO=1`, -`VLLM_FLASH_V100_XQA_E4M3_G6_WAVE_PARTITIONS=1`, and -`VLLM_FLASH_V100_XQA_E4M3_G6_MERGED_WAVE_LAUNCH=1`. Setting the wave flag to -zero restores the prior p64/p256 policy. Setting the merged-launch flag to zero -retains the wave thresholds but restores the separate long-wave launches. -Setting the auto flag to zero restores the original explicit-partition route. - -### Exact endpoint result - -The control and candidate use the same model, projections, sampler, TP/GPU -set, max length, input/output lengths, attention backend, KV type, graph mode, -and no-MTP state. Only the Flash-V100 E4M3 attention partition policy changes. - -| Final context | Route | Prefill | TTFT | Pure TPOT | Steady decode | Target | -|---:|---|---:|---:|---:|---:|---:| -| 128K | p64/p256 control | 283.296 s | 283.327 s | 24.654 ms | 40.561 tok/s | 60 tok/s | -| 128K | LUT merged-wave | 283.127 s | 283.154 s | **16.172 ms** | **61.834 tok/s** | pass, +3.06% | -| 256K | p64/p256 control | 1085.628 s | 1085.673 s | 36.422 ms | 27.456 tok/s | 45 tok/s | -| 256K | LUT merged-wave | 1085.519 s | 1085.568 s | **19.851 ms** | **50.376 tok/s** | pass, +11.95% | - -At 128K, throughput improves 52.45% and TPOT falls 34.40%. At 256K, -throughput improves 83.48% and TPOT falls 45.50%. Both requests emitted all -64 tokens, finished by length, and reported no corruption. The complete token -streams are exactly unchanged from control: 128K SHA256 is -`66a130ea9d68faaf40f7c7cfb942c579f759b9535394cce958699bd620df2d6f`; -256K SHA256 is -`e1bc9ab8a979e43cda22348c0ff0056b5a7c3212ba57e830b731f3cf51906524`. -This is an attention-only optimization and does not credit a sampler change. - -The post-integration 1K regression guard keeps the accepted TurboMind -projection routes, FlashQLA GDN BA-split route, official sampler, disabled -compile cache, and 256 output tokens. After excluding the first cold request, -the accepted GDN BA-split control is 83.761/83.721 tok/s and the -graph-specialized candidate is 83.184/83.187 tok/s, a 0.66% median delta. All -three candidate repetitions are token-for-token identical to one another and -to that accepted GDN control; their compact token-list SHA256 is -`c243818da5cced8b8653df4c441f713e978f8dfe7d8343d1b78c0eb3aeffe8ef`. -The short default graph therefore removes the inactive long-wave pair without -changing the admitted output stream or causing a material short-context -regression. A cache-reload diagnostic is not used as quality evidence. - -### Operator and resource evidence - -A same-GPU operator A/B uses the production interleaved K/V layout. At 128K, -the merged route takes 370.087 us versus 377.429 us for separate long-wave -dispatch; at 256K it takes 599.839 us versus 611.925 us. The selected explicit -p896/p1664 controls take 362.175/593.954 us respectively. Merged-route eager -and CUDA Graph outputs are bit-exact to those explicit controls. The small -remaining dispatch delta is preferable to three long kernel/reducer pairs in -the captured full-model graph. - -The final 256K Nsight Systems CUDA-profiler-range trace uses the same -production interleaved layout and final extension. The active merged p1664 -partition kernel takes 0.576603 ms and its reducer takes 0.033280 ms. The -device-routed short guard kernel/reducer take 0.004544/0.003552 ms, for -0.617979 ms across Flash attention kernels. The old p256 trace takes -1.391781 ms in the partition kernel and 0.316769 ms in the reducer, or -1.708550 ms total. The retained route therefore cuts traced Flash attention -time by 63.83% (2.76x), while remaining bit-exact to the explicit p1664 -control. The trace is an isolated per-rank attention operator capture, not a -full-model TPOT decomposition. - -During the exact decode windows NVML reports 100% GPU-busy duty on every rank. -At 128K, memory-active duty averages about 48.2%, allocation is -29.45-29.75 GiB/GPU, and mean power is 200.7-215.8 W. At 256K, memory-active -duty averages about 45.5%, allocation is unchanged, and mean power is -211.6-226.8 W. SM and memory clocks remain 1530/877 MHz. NVML memory-active -duty is not achieved HBM bandwidth; exact SM, Tensor Core, occupancy, and DRAM -counters remain unavailable because Nsight Compute returns -`ERR_NVGPUCTRPERM` on this host. - -The production-layout correction is material: treating the unbound K/V views -as contiguous produced full-model output corruption and is rejected. A QK -software pipeline slowed the p896/p1664 kernels by roughly 30-38%; a split -reducer did not beat the retained route; paired E4M3 conversion measured -409.958/683.428 us at 128K/256K and lost to the shared LUT. All three -experiments are absent from the admitted path. - -### Post-acceptance kernel headroom - -Three additional changes retain the same admitted route and the real unbound -interleaved K/V views. First, the long-wave CTA stages its partition page IDs -once instead of reloading them for every tile. Second, a strict host-side -stride gate specializes the production `(blocks, 2, 1568, 1, 256).unbind(1)` -layout: the physical block stride is `2 * 1568 * 256`, while token and head -strides remain 256. Any other view uses the generic stride path. Third, the -fixed-layout E4M3 PV loop reads adjacent shared-memory values as `half2` while -retaining one FP32 FMA per dimension and the original token accumulation -order. It only changes lane ownership and restores dimensions explicitly at -the output store. - -Each stage was measured as a four-arm control/candidate/candidate/control A/B -on one otherwise idle V100. Every arm contains five repetitions for each of -three page-table layouts; the table reports the median of the ten samples per -layout followed by the median across layouts. The exact partition-boundary -companions at 131,071 and 262,080 tokens pass as well. - -| Exact context | Stage | Control | Candidate | Kernel delta | -|---:|---|---:|---:|---:| -| 128K | partition page IDs | 373.076 us | 368.604 us | -1.199% | -| 256K | partition page IDs | 602.017 us | 594.632 us | -1.227% | -| 128K | fixed interleaved stride | 370.102 us | 364.585 us | -1.491% | -| 256K | fixed interleaved stride | 598.257 us | 587.909 us | -1.730% | -| 128K | shared-V `half2` PV | 362.739 us | 348.265 us | -3.990% | -| 256K | shared-V `half2` PV | 584.745 us | 558.515 us | -4.486% | - -All 12 sequence/layout cases have one output SHA256 across all four arms and -are bit-exact to their explicit p896/p1664 controls (`max_abs=0`). The final -partition kernel remains at 168 registers/thread and two CTAs/SM; the PV -change removes the prior 32-byte stack frame. Relative to the fixed-stride -binary, its static SASS instruction count falls from 2,424 to 2,248 and shared -loads fall from 118 to 90, with the same 57 FFMA instructions. These are -static cubin counts, not NCU throughput counters. The timed binary SHA256 is -`43f742cc228100dd0a58aed37bf1cb6007447129bcf3f1d28148bee71b6ff2a5`. -The clang-formatted final build SHA256 is -`5e66a5331aaf83e5e15080c5122c78047b94d3cb47a721572613bbbba9172d47`; -its target-kernel SASS is byte-identical to the timed build (normalized SASS -SHA256 `0684d453072ed630aa68477b618ab4c090eb18c87422a6de90af3fd43740a1ee`) -and its resource record is unchanged. The final build is bit-exact in eager -and CUDA Graph replay at 128K/p896 and 256K/p1664 (`max_abs=0`); all eight -parameterized page-1568 wave-boundary regressions pass. - -Applying the three independently matched operator ratios to the measured -5.921/9.597 ms attention slices, while leaving the measured 10.251/10.253 ms -non-attention residuals unchanged, projects 15.784 ms (63.354 tok/s) at 128K -and 19.151 ms (52.216 tok/s) at 256K. These are decomposition-based -projections, not replacement end-to-end measurements; the accepted measured -endpoint remains 61.834/50.376 tok/s. - -A packed pair-LUT initialization experiment is rejected because it is about -1.0-1.2% slower at 128K and 0.78% slower at 256K. Replacing constant division -with compare/subtract address normalization passed 2,801,664 exhaustive host -address checks and all GPU exactness cases but was statistically neutral. A -scalar PV-accumulator rewrite removed the stack frame but was neutral or -slower at 256K. None of these rejected variants is present in the admitted -path. - -Primary retained evidence is: - -- `.artifacts/qwen38_nvfp4_speed_20260823/results/e4m3_xqa_auto_single_p12288_long_i16k_64k_128k_256k_o64.json` -- `.artifacts/qwen38_nvfp4_speed_20260823/results/e4m3_xqa_lut_merged_wave_p12288_long_i128k_256k_o64.json` -- `.artifacts/qwen38_nvfp4_speed_20260823/telemetry/e4m3_xqa_lut_merged_wave_p12288_long_i128k_256k_o64.nvml.csv` -- `.artifacts/qwen38_e4m3_xqa_long_route/results/e4m3_xqa_lut_separate_same_gpu4.json` -- `.artifacts/qwen38_e4m3_xqa_long_route/results/e4m3_xqa_lut_merged_same_gpu4.json` -- `.artifacts/qwen38_e4m3_xqa_long_route/results/e4m3_xqa_wave_boundaries_graph_workspace.json` -- `.artifacts/qwen38_e4m3_xqa_long_route/results/page_ids/` -- `.artifacts/qwen38_e4m3_xqa_long_route/results/page_ids_fixed_stride/` -- `.artifacts/qwen38_e4m3_xqa_long_route/results/page_ids_fixed_stride_pv_half2/` -- `.artifacts/qwen38_e4m3_xqa_long_route/results/page_ids_packed_lut/` -- `.artifacts/qwen38_e4m3_xqa_long_route/results/page_ids_fast_address/` -- `.artifacts/qwen38_e4m3_xqa_long_route/profiles/e4m3_xqa_final_merged_wave_p1664_l262144_nsys.nsys-rep` diff --git a/docs/design/sm70_qwen38_qpn8_decode.md b/docs/design/sm70_qwen38_qpn8_decode.md deleted file mode 100644 index 21ae39aafb..0000000000 --- a/docs/design/sm70_qwen38_qpn8_decode.md +++ /dev/null @@ -1,245 +0,0 @@ -# SM70 Qwen3.8-27B-FP8 QPN8 Decode Acceptance - -Date: 2026-08-23 - -## Decision - -The memory-neutral QPN8 dense route is available as an opt-in for the -Qwen3.8-27B-FP8 TP4 no-MTP model and projection shapes when configured -`max_num_seqs` is at most eight. Set `VLLM_SM70_FP8_QPN8=1` to request it. -The model-quality gate passes, but pure-FP8 default promotion still requires a -same-source speed rerun with the same number of steady decode intervals. An -explicit request with missing QPN8 operators fails closed. - -The gate checks the Qwen3.8 text architecture, 5120 hidden size, 17408 -intermediate size, 64 layers, full-attention interval 4, head dimension 256, -and TP4. Qwen3.6 and both 35B acceptance routes remain unchanged. -Configurations admitting more than eight live sequences retain TurboMind at -load time until a measured M>8 QPN8 route exists. MTP configurations also -retain TurboMind until their verification widths pass speed and quality gates. -These guards prevent the functional large-M dequantization fallback from -becoming a concurrency or speculative-decode regression. - -This acceptance is a no-MTP, M=1 decode speed result with numerical tests up -to M=8 and model-quality batches up to eight. M>8 remains functional through -one bounded 85 MiB FP16 weight workspace. Large-M fused gate/up also has a -transient `M x 8704` FP16 GEMM output, so that fallback is not an accepted -high-concurrency memory or speed result. - -## Contract - -- Model: `Qwen3.8-27B-FP8`, TP4 on physical V100-SXM2-32GB GPUs 4-7. -- Runtime: Python 3.12, Torch 2.10.0+cu128, CUDA 12.8, source checkout and - source-built SM70 extension; no wheel changes are used. -- Production speed case: no MTP, E5M2 KV, `FLASH_ATTN_V100`, CUDA graphs, - `max_model_len=262144`, input 1024, output 256, and one live request. -- Official sampling: `temperature=1.0`, `top_p=0.95`, `top_k=20`, request seed - 20260815. Pure decode excludes TTFT/prefill and uses the 255 steady token - intervals. -- Quality case: identical baseline/candidate datasets and seeds, 4096-token - rolling-PPL windows, and generation/log-likelihood batches 1, 2, and 8. - -## Source Route - -Checkpoint-native block-FP8 weights are repacked once at model load into one -QPN8 `[K, N]` code tensor plus grouped FP16 scales. The old TurboMind tensor -is replaced rather than retained, avoiding the approximately 6 GiB/rank -duplicate-weight experiment. The matched candidate reported 19.89 GiB -available KV memory versus 19.80 GiB for the control. - -The opt-in model and shape gate admits only: - -| Projection | TP-local K | TP-local N | Decode schedule | -|---|---:|---:|---| -| MLP gate/up with fused SiLU×up | 5120 | 8704 | split-K 8, two accumulator chains, prefetch | -| MLP down | 4352 | 5120 | split-K 16, one accumulator chain | -| GDN/full-attention output | 1536 | 5120 | split-K 16, one accumulator chain | - -GDN input and full-attention QKV remain on TurboMind. The M decision is inside -the opaque C++ operator so AOTInductor cannot incorrectly specialize a Python -M=1 branch for a later prefill. M=1-8 calls QPN8 directly; larger M -materializes one layer at a time in the shared weight workspace and calls FP16 -GEMM. The large-M gate/up temporary described above is bounded by the engine's -`max_num_batched_tokens`, not by the 85 MiB weight workspace. - -## Pure Decode Result - -| Route | Pure TPOT | Steady decode | Change | -|---|---:|---:|---:| -| Matched TurboMind control | 16.985327 ms | 58.874 tok/s | baseline | -| QPN8 accepted subset | 15.808000 ms | 63.259 tok/s | +7.448%, -1.177327 ms/token | -| Complete-source TurboMind control | 16.961295 ms | 58.958 tok/s | source A/B baseline | -| Complete-source QPN8 candidate | 16.347973 ms | 61.170 tok/s | +3.752%, -0.613322 ms/token | - -The matched operator-integration pair generated all 256 tokens and finished by -length. Their sampled token streams first differ at output token 91; sampled -token identity is not the quality gate because both trajectories are coherent -and the acceptance criterion is PPL plus fixed reasoning, knowledge, and -Chinese datasets. - -The complete-source control generated 256 tokens and finished by length. The -complete-source QPN8 request used the same prompt, sampling contract, -and source extension, but encountered the model's normal stop token 248044 at -token 233. Its TPOT therefore covers 232 steady intervals rather than the -control's 255. The earlier 63.259 tok/s matched operator-integration result is -retained as evidence of headroom. Because 61.170 tok/s uses 232 intervals while -the control uses 255, another same-source equal-interval run is required before -pure-FP8 default promotion. - -## Nsight Systems Per-Token Decomposition - -Node tracing adds profiler overhead, so these rows explain composition rather -than replace the unprofiled TPOT above. - -| Steady replay metric | Control | QPN8 | Delta | -|---|---:|---:|---:| -| Replay interval | 17.418 ms | 16.437 ms | -0.981 ms | -| GPU activity union | 16.870 ms | 15.894 ms | -0.976 ms | -| Total idle gaps | 0.548 ms | 0.543 ms | -0.005 ms | -| Kernel launches/rank/token | 1085.5 | 1085.9 | unchanged | - -The trace gain is therefore GPU kernel service, not a hidden CPU wait or -launch-gap reduction. Dense service changes from 10.335 ms/rank/token to: - -| Candidate dense component | Time | Launches/rank/token | -|---|---:|---:| -| Fused QPN8 gate/up | 4.018 ms | 64 | -| QPN8 down and output projections | 3.040 ms | 128 | -| Remaining TurboMind GDN/full-QKV input | 2.332 ms | 64 | -| Total dense | 9.391 ms | 256 | - -This accounts for about 0.944 ms/token of the trace improvement. TP -communication is 1.680 to 1.637 ms and LM-head/sample/gather remains about -1.33 ms, so neither explains the main win. - -## Resource Release and Peak-Band Residency - -Whole-request NVML averages are not used as achieved FLOPS or bandwidth. -Nsight Systems establishes how long grids can occupy the machine, and Nsight -Compute supplies achieved per-kernel counters. - -The grid-limited occupancy ceiling moves from 93.56% of service below 25% -occupancy in the control to 48.56% below 25%, 32.11% at 25-50%, and 19.26% at -50-75% with QPN8. This is a duration distribution, not an average-utilization -claim. - -For the accepted QPN8 kernels, NCU reports: - -| Shape | Registers/thread | Achieved occupancy | SM throughput | DRAM throughput | Tensor-pipe elapsed | -|---|---:|---:|---:|---:|---:| -| Down | 50 | 48.90% | 27.20% | 82.74% | 9.50% | -| Output | 50 | 48.17% | 27.69% | 71.32% | 8.23% | -| Fused gate/up | 64 | 42.33% | 29.98% | 87.74% | 10.08% | - -Counting the useful M=1 dense multiply-add work, throughput is 1.177 to -1.295 TFLOP/s per GPU, or 4.71 to 5.18 TFLOP/s across TP4. -Dense-service-weighted Tensor Core pipe activity is 7.97% to 9.15%, -equivalent to about 9.99 to 11.47 TFLOP/s per GPU, or 39.96 to 45.88 TFLOP/s -across TP4, relative to the V100's nominal FP16 Tensor Core peak. These answer -different questions: the first is model-useful arithmetic, while the second -includes tensor instructions spent inside the split-K implementation. Neither -is a whole-request average. - -Across the 9.391 ms candidate dense service, 7.059 ms resides in the 25-50% -SM-throughput band and 2.332 ms remains below 25%; no dense kernel reaches -50% SM throughput. For DRAM, 6.172 ms resides at 75-90% and 3.219 ms at -50-75%; no dense kernel reaches 90%. Dense-service-weighted achieved -occupancy rises from 9.96% to 36.25%, but the duration bands above are the -primary evidence. - -The result is not “the V100 cannot reach peak compute.” QPN8 makes a much -larger fraction of dense time sustain high memory duty and occupancy, while -the unchanged approximately 0.54 ms idle budget shows that host launch work -is not the present first-order limit. - -## SGLang-V100 Comparison Boundary - -The current [sglang-V100 README](https://github.com/haohervchb/sglang-V100/blob/main/README.md) -now reports a Qwen3.8-27B-FP8 target-only, E5M2 KV, TP4 row at 58.2 tok/s for -1K decode and 50.8 tok/s for 25K decode. Its table uses one cold request and -256 greedy output tokens. The linked -[target-only audit](https://github.com/haohervchb/sglang-V100/blob/main/benchmark/qwen38_27b_fp8_target_e5m2_v100_20260822/README.md) -records 17.180 ms TPOT, or 58.21 tok/s, at 1,024 input tokens. The 63.259 -tok/s QPN8 result here is 8.67% above that published number, so the public -evidence points to this route being faster. It is not a strict head-to-head -because this acceptance uses fixed random sampling and a different -prompt/client harness. A formal claim still requires running SGLang's -target-only command with the same prompt, sampling, graph, and pure-decode -accounting used here. - -The published 136.6 tok/s DFlash2-8 row is speculative and is not a no-MTP -comparison. - -## Quality Gate - -The gate compares the prior TurboMind route and QPN8 with identical prompts, -few-shot examples, seeds, and batching. Lower is better for the three -WikiText metrics; higher is better for the remaining accuracy metrics. - -| Dataset and metric | Samples | Control | QPN8 | Result | -|---|---:|---:|---:|---| -| WikiText word perplexity | 62 documents | 8.790013 | 8.789904 | pass, lower | -| WikiText byte perplexity | 62 documents | 1.501519 | 1.501515 | pass, lower | -| WikiText bits/byte | 62 documents | 0.586422 | 0.586419 | pass, lower | -| GSM8K 5-shot strict exact match | 128 | 69.531% | 70.312% | pass, +1 item | -| GSM8K 5-shot flexible extraction | 128 | 73.438% | 73.438% | pass, equal | -| MMLU abstract algebra 5-shot | 100 | 67.0% | 67.0% | pass, equal | -| MMLU high-school computer science 5-shot | 100 | 89.0% | 89.0% | pass, equal | -| MMLU professional law 5-shot | 100 | 63.0% | 66.0% | pass, +3 items | -| MMLU three-subject macro | 300 | 73.0% | 74.0% | pass | -| C-Eval computer network 5-shot | 19 | 73.684% | 73.684% | pass, equal | -| C-Eval high-school Chinese 5-shot | 19 | 73.684% | 73.684% | pass, equal | -| C-Eval advanced mathematics 5-shot | 19 | 57.895% | 57.895% | pass, equal | - -No evaluated metric regresses. This closes the quality side of the opt-in -route; it does not replace the missing equal-interval performance rerun. - -## Build and Tests - -- The complete source `_C` target built 47/47 objects for CUDA 12.8 and SM70; - the resulting 45 MiB extension registers all six prepare, dequantize, - prefill, dispatch, GEMM, and fused-gate QPN8 operators. -- The operator race passed 960 rows. Maximum relative L2 was `2.851e-4`, - minimum cosine was `0.99999988`, and maximum absolute error was `9.77e-4`; - CUDA Graph replay deltas were zero. -- Seventeen focused Python tests pass, covering opt-in/explicit-off policy, - model and shape gates, workspace reuse, M=1 and M>8 opaque dispatch, fused - gate/up dispatch, and warmup registration. -- Ruff lint/format, changed-line clang-format, Python byte compilation, shell - syntax checks, and `git diff --check` pass. - -## Next Optimization Order - -1. Extend QPN8 with a single-launch two-dimensional row-tile grid and race - M=9, 11, and 16 against TurboMind. Do not compose an 8-row launch with a - separate tail launch, and do not assume the batch-one winner stays best. -2. Race M=17, 32, and 64 and keep a measured route map. The external - [v100-skinny QPN study](https://github.com/dnv2003/v100-skinny/blob/main/docs/qpn_race_notes.md) - also finds a crossover to WMMA/TurboMind above the small-batch band. -3. Before high-concurrency acceptance, replace or chunk the large-M fused - gate/up temporary and record peak memory under the supported batched-token - limit. -4. Add a bounded large-M route for GDN input/full-QKV, then quality-gate it. - Those 64 remaining TurboMind launches consume 2.332 ms/token. -5. Continue shape-specific work on fused gate/up (4.018 ms) and down/output - (3.040 ms), measuring time above 90% DRAM and 50% SM rather than optimizing - a whole-run average. -6. Revisit the 1.637 ms communication and 1.33 ms LM-head only after the dense - residue falls. The current trace does not justify prioritizing Python launch - overhead. -7. Repeat the decomposition at 25K and 70K. SGLang's target-only 70K trace - reports 10.60 ms/token in FP8 projections and 10.35 ms/token in attention, - so the 1K optimization order must not be hard-coded for long context. - -The ordering is also consistent with external evidence, without treating a -newer-architecture implementation as a V100 drop-in. The -[v100-skinny M sweep](https://github.com/dnv2003/v100-skinny/blob/main/docs/REPRODUCE.md) -uses a measured SIMT/QPN/WMMA dispatch map rather than one kernel for every M. -The [CUTLASS grouped scheduler](https://github.com/NVIDIA/cutlass/blob/main/media/docs/cpp/grouped_scheduler.md) -shows how a bounded persistent CTA work queue can retain residency across many -independent small tiles; it becomes relevant only if independent problems can -be grouped without crossing transformer dependencies. The -[FlashInfer paper](https://arxiv.org/abs/2501.01005) similarly makes dynamic, -load-balanced scheduling compatible with a static CUDA Graph. For this route, -those are design constraints to test under concurrency, not reasons to replace -the measured SM70 kernels speculatively. diff --git a/docs/design/sm70_tile_runtime_exploration.md b/docs/design/sm70_tile_runtime_exploration.md index 1a594ba074..0fe732e7a1 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: -- TileRT execution-model blog: -- TileRT 1000 TPS blog: -- DistFuse tile-wise GEMM/all-reduce paper: -- TokenWeave paper: -- CUTLASS-native Distributed GEMM write-up: +- 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 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 7bb276c394..74b8425639 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -24,576 +24,6 @@ Goal: sampling, and normal MoE architecture wherever possible. - Do not carry unsafe quality-regressing routes into the main path. -## Active MRV2 DFlash2 campaign, 2026-08-20 - -- Integration base: `onecat/main@7aede2cf010d92815c9d7bff25867b4fa009b6cb`. -- Review chain starts at Draft PR #252 on - `codex/v100-dflash-mrv2-base-20260820-041115`. -- Frozen sources: vLLM MRV2 DFlash base `9842d7014502`, vLLM DFlash2 PR - `66e5414c6d75`, SGLang `c14312a66420`, and sglang-V100 `5526ef1c6a82`. -- Route contract: `method=dflash` is MRV2-only; the experimental V1 flat path - is removed. `method=dflash_ddtree` retains the existing V1 DDTree scheduler, - payload, verifier, and diagnostics through a narrow model compatibility - adapter. Existing Eagle and MTP implementations are not replaced. -- Selector top-K is checkpoint-owned and fixed at 16 for Qwen3.8 DFlash2; - official bring-up uses seven draft tokens (block size eight). Target sampling - top-k 20 is a separate sampling parameter. -- PR1 evidence so far: focused routing/config/DDTree tests pass, and the V1 - DFlash input-preparation test passes on V100 after supplying the KV-group - initialization contract used in production. Two unrelated DDTree metadata - assertions fail identically on the frozen base and are recorded as pre-existing. -- Draft PR #253 is the PR2 review scope on - `codex/v100-dflash2-sm70-20260820-044011`, based only on PR1. Its worktree is - `/home/ymzx/桌面/1cat-vllm/worktrees/v100-dflash2-sm70-20260820-044011`. -- PR3 evaluated a gated SM70 dense-FP16 selector fusion on top of PR2. Its - required M=7 gate was negative, so the final PR #254 merge scope retains only - the correctness/performance evidence below; no selector implementation, - runtime flag, duplicate weight layout, or benchmark harness enters `main`. - -### PR2 DFlash2 correctness evidence, 2026-08-20 - -- Checkpoint contract is frozen at - `incoai/Qwen3.8-27B-DFlash2@dedf8df68adfb1afeaf7b7480c0a0243108177b4`. - The retained local `model.safetensors` SHA256 is - `67fc76d68dc5a9415511a4f394ef744d67510cd20e93b37cc2cc7d28e4bab65c`. - Its architecture is `DFlash2DraftModel`, selector K is 16, block size is 8, - attention is non-causal with sliding window 2048, and checkpoint layer IDs - `[5,19,33,47,61]` map to target hidden boundaries `[6,20,34,48,62]`. -- The dependency closure now contains the official grouped convolution, - selector codebooks, TP-aware candidate top-K, DFlash2 MRV2 walk/cache, FP32 - proposal scores, rejection `log1p`, and a separate selector compile-cache - tag. FlashInfer top-K is capability-gated and is never called on SM70. -- The SM70-only numeric adaptation keeps FP32 residual state, explicit BF16 - round-to-nearest-even points, scaled FP16 attention and gate-up projections, - and per-row power-of-two SwiGLU transport. The seventh selector slot uses a - separate Volta tail kernel; other architectures retain the upstream walk. -- Flash-V100 graph capture now binds DFlash's K+1 non-causal query directly to - persistent paged-prefix metadata. The causal small-query decode branch is - unavailable to a non-causal draft. Target E5M2 and draft FP16 KV pages are - unified without truncating hybrid Mamba physical padding. -- Single-request eager and CUDA Graph runs produced the exact same 16 token IDs, - hash `6de4c3f326c76ecc1b31b8026501ddf22f73a34b53b2fc46c622dee6b4a2d995`, - candidate lattice, acceptance trace, and draft KV. The graph artifact is - `/data/minimax-h3/task-cache/v100-dflash2-20260820/pr2/dflash2-graph-greedy16-fixed.json`; - the matched eager artifact is `dflash2-eager-noncausal-fixed16.json`. -- Matched batch-2 eager/graph runs generated 96 and 61 tokens with identical - hashes and all acceptance counters. CUDA Graph batch-4 generated mixed - lengths 64/61/64/64 without OOB, stale-buffer, NaN, or replay corruption; - each result matched the corresponding eager prefix. Retained artifacts are - `dflash2-{eager,graph}-gsm8k-mixed-b2-2x96-greedy.json` and - `dflash2-graph-gsm8k-mixed-b4-4x64-greedy.json` under the PR2 cache directory. -- On four fixed GSM8K prompts with the Qwen3.8 xhigh chat template and official - target sampling, greedy drafting emitted 421 tokens with mean accepted length - 4.1068 and per-position acceptance - `[.8544,.6505,.4854,.4175,.3301,.2233,.1456]`. Probabilistic drafting emitted - 452 tokens with mean 4.2407. These short V100/FP8 results are correctness and - sampling diagnostics, not a like-for-like replacement for the published - H200/BF16 acceptance figures. -- The fixed 32-to-16 graph diagnostic reported 80.11 steady decode tok/s versus - 18.72 eager with exact token/acceptance equality. The matched batch-2 run was - 116.67 tok/s graph versus 43.53 eager. These are short diagnostic shapes and - are not a target-only or dense-selector paired performance gate. -- DFlash1 remains on the official MRV2 speculator for the local 9B, 27B, and - 35B-A3B checkpoints. A real 9B-AWQ/FP16-draft eager smoke produced eight - tokens, hash `2a205303e4c2b54b62d7562643abcaa425952d7dbcf3f3c1a4b6a4e3e3634be9`, - accepted length 3.5, and confirmed Flash-V100 routing. The artifact is - `/data/minimax-h3/task-cache/v100-dflash2-20260820/pr2/dflash1-9b-awq-mrv2-eager-route.json`. -- Existing MTP was not replaced. A real Qwen3.8 TP4 MTP4 probabilistic smoke - confirmed native MTP, local argmax, Flash-V100, and E5M2 routes; it emitted 32 - tokens with accepted length 4.25 and per-position acceptance - `[1.0,1.0,.75,.5]`. The artifact is - `/data/minimax-h3/task-cache/v100-dflash2-20260820/pr2/dflash2-mtp4-qwen38-regression-eager.json`. -- Final changed-path gate: SM70 arguments plus KV layout, DFlash2/MRV2 routing, - and warmup reservation total `113 passed, 1 skipped` on CPU; the dedicated - V100 Mamba align kernels add `5 passed`. The new warmup matrix covers MRV2 - DFlash K+1, DDTree tree-budget lookahead, unchanged Eagle/MTP/DSpark values, - and align/non-align Mamba speculative tails. Earlier V100 rejection sampling - remains `14 passed` and probabilistic selector score-cache `1 passed`; Ruff - passes every changed Python file. The DDTree targeted matrix remains - `71 passed, 1 failed`, with - the one shared-storage assertion reproduced identically on PR1. Three stale - MTP step-index helper failures also reproduce on PR1; their 19 unaffected - cases pass. Four Eagle tests that require the gated Llama-3.1 config stop at - Hugging Face 401 on both PR1 and PR2; the local DFlash/Eagle compatibility - test itself passes. -- MRV2 now implements the upstream hybrid Mamba `align` state migration on the - GPU without changing the V1 DDTree state-selection path. The new-request - state column is seeded from the resolved `MambaSpec.block_size`, not the - smaller global cache block size; this is required when E5M2 target pages and - FP16 DFlash pages share a hybrid pool. Upstream hash-layout fix #51180 is also - included so aligned Mamba groups retain the finest divisible hash unit. -- Prefix correctness is accepted for the real Qwen3.8-27B-FP8 TP4 route. With - the resolved target Full/Mamba block size 3296 and draft SWA block size 1648, - a 3505-token prompt produced a 3296-token hit on the second request. Eager - kept the exact 16-token hash `a2c583fc17c7...`, accepted length 6.6667, and - per-position counts `[3,3,3,2,2,2,2]`; TTFT changed from 1.117s to 0.303s. - CUDA Graph kept the same token and acceptance trace with TTFT - 1.119s to 0.174s. Artifacts are `dflash2-prefix-tp4-{eager,graph}-4k.json` - under the PR2 cache directory. -- A paired target-only TP4 probe hit 3136 tokens with exact output and TTFT - 1.010s to 0.167s (`qwen38-prefix-tp4-target-only-eager-4k.json`). This - separates target Mamba-state restoration from DFlash proposal correctness. - No separate cross-request draft-prefix index or persistence layer was added; - the accepted same-process path uses the existing unified hybrid block cache, - while an ordinary miss rebuilds draft context KV during DFlash proposal. -- Numeric classification for the DFlash non-causal paged attention route is - `B-accept`, with its layout component separately passing Type A. Aligned - DFlash shapes at sequence lengths 8/69/512/2048 show semantic-boundary - paged-vs-dense maxima - `0/2.44e-4/1.22e-4/6.10e-5`, all below the FP16 `2e-3` bound and not growing - with length. Scaling V by 256 reproduces raw differences up to 0.0625, which - become at most 2.44e-4 after the model's required `/256` range restoration. - Reversing every physical KV page and applying the matching block-table - permutation is bitwise exact for every shape, proving page addressing, - layout, mask, and KV-source alignment. The bounded cross-kernel residual is - mechanically explained by the source arithmetic: for D != 256, dense - Flash-V100 scales QK by `log2(e)` and evaluates `exp2f`, whereas paged - Flash-V100 retains natural-log scores and evaluates `expf`. These are - mathematically equivalent reductions with different FP16 rounding. The - release extension was built from paged-kernel source SHA256 - `9be6b237fcd1c89653c26460054c1c25edd6e7d8f9685f7b9f4b83a2f1ba5e58`, - identical to the integration source; an older independently built extension - produced the same case-by-case results. - The retained evidence is - `/data/minimax-h3/task-cache/v100-dflash2-20260820/pr2/dflash2-noncausal-paged-attention-exactness.json`. - Together with exact eager/graph model tokens, candidates, acceptance trace, - and draft KV, the attention numeric gate no longer blocks PR3. This does not - relax the separate fused-selector performance and exact-token gates. - -### PR3 SM70 selector-fusion disposition, 2026-08-21 - -- Draft PR #254 evaluated reuse of the TurboMind dense-FP16 LM-head kernel - behind `VLLM_SM70_DFLASH2_FUSED_SELECTOR`. The experiment preserved the - checkpoint contract—local top-20 candidates reduced to the checkpoint-owned - global top-16, with the official seven-token draft block—but the production - implementation and flag were removed after the negative gate. -- The original top-20 epilogue assigned one row to each of four warps although - its CTA has `M=8`. The fixed epilogue assigns rows round-robin, covering the - seventh DFlash2 row and M=8 tail. A graph-safe operator variant accepts - persistent FP32 value and INT64 token partial workspaces; the allocation-owning - legacy entry remains available for non-graph callers. -- Equal FP16 logits exposed a second exactness issue: the fused kernel's - deterministic token-ID tie break did not reproduce `torch.topk`'s dense-shard - ordering. The accepted eager experiment scatters the fused top-20 into a - persistent shard-shaped FP16 buffer and runs local dense-order top-16 before - exchanging only 16 values and token IDs per TP rank. This restores exact - ordered IDs and values, not merely the same candidate set. -- The measured experimental SM70 extension SHA256 was - `077ca806e8fad6b452dbe59cd0da336ff88c9341386c8480bbbe009b8029c7b2`. - Persistent-workspace tests cover M=1/4/5/7/8, TP4 vocab offsets and padding, - equal-score ordering, and CUDA Graph replay. The M-shape set passed three - consecutive runs after the test stopped exercising the allocation-owning - legacy entry. -- The real Qwen3.8 TP4 selector microbenchmark used M=7, local N=62080, K=5120, - FP16 hidden/weight, local top-16 and TP global top-16. Ordered IDs and FP16 - values were exact with maximum difference zero. The dense - `torch.mm + local16 + TP global16` mean was 1.1585 ms; the production-shaped - fused path was 1.2254 ms, a 5.78% regression. Evidence is - `/data/minimax-h3/task-cache/v100-dflash2-20260820/pr3/sm70-top20-workspace-dense-order-m7-tp4.json`. -- The dedicated experimental-branch DFlash2 suite passed `30 passed, 11 skipped` without a - visible GPU and `41 passed` on V100. Ruff passes all eight changed Python - files, the rebuilt SM70 extension links successfully, and `git diff --check` - is clean. -- Earlier direct-top20 graph results that appeared much faster are rejected. - The dense 256-token artifact had hash `6477ffd9...` and accepted length - 2.2768, while the direct fused artifact had hash `1f7f1cdf...` and accepted - length 3.3462. The speed change therefore measured a different sampling and - acceptance trajectory, not an optimization of the same computation. -- Preparing the duplicate LM-head layout also raised model memory by about - 1.19 GiB per rank (10.23 versus 9.04 GiB in the final compiled-graph runs), - directly reducing long-context KV capacity. The fused selector is therefore - eager-only. Under the accepted Flash-V100 compiled CUDA Graph policy, an - explicit request is normalized to zero before workers fork; no duplicate - layout is prepared and startup prints the fallback decision. -- A paired post-rebuild TP4 graph check compared selector-off with a requested - selector that was normalized before worker startup. Both emitted the exact - same 32 token IDs, hash `471733c0...`, accepted-token count 13, mean accepted - length 1.6842, and per-position acceptance vector. Artifacts are - `dflash2-pr3-dense-rebuilt-graph-1k-32-r1.json` and - `dflash2-pr3-selector-normalized-graph-1k-32-r1.json` in the PR3 cache. -- Decision: do not merge the selector implementation; keep the official - draft-token default at seven. Because the required M=7 micro gate is already negative, - the 32K/128K confidence-interval sweep and 3/5/7 default scan are intentionally - not run. A future graph-capable fusion must first beat dense M=7 without - duplicate-weight capacity loss and then re-enter the full long-context gate. - -### DFlash acceptance root-cause repair, 2026-08-22 - -- A same-machine 64-prompt GSM8K audit held the FP8 target, E5M2 target KV, - FP16/auto draft KV, TP4, official probabilistic sampling, seven draft tokens, - Flash-V100, CUDA Graph, prompt order, and single-request sequencing fixed. - Before the repair, vLLM averaged 4.5395 completion tokens per verification - step versus 5.2489 for `haohervchb/sglang-V100`; the paired prompt delta was - 0.7094 with a 95% interval of [0.5882, 0.8306]. This excludes target FP8 as - the explanation for the acceptance gap. -- The MRV2 backport copied the grouped DFlash K-normalization call from upstream: - input `[L, num_ctx, num_kv_heads, head_dim]` with weights `[L, head_dim]`. - It omitted upstream dependency commit `02a1f23711`, whose stable RMSNorm - kernel uses the outer layer index to select each weight row. The local kernel - still treated every weight as one-dimensional and silently applied row zero - to all five draft layers. Checkpoint rows are not interchangeable: layers - 1-4 differ from layer zero by 0.468-0.621 mean absolute weight and - 0.383-0.443 relative L2. -- SGLang-V100 is unaffected because both paths preserve layer ownership: its - sequential fallback calls each layer's `k_norm`, and its fused Triton - materializer indexes `layer_id * k_norm_weight_stride_layer`. This is the - implementation difference that explains why SGLang acceptance was normal. -- Commit `c53b7974fb` restores the exact missing dependency semantics while - retaining the local SM70 RMSNorm additions. One-dimensional weights keep - stride zero; two-dimensional weights require `[input.size(0), hidden]` and - select the row from the outermost input dimension. Invalid rank and shape are - rejected instead of silently producing the wrong draft KV. -- The V100 kernel regression passes 13/13 cases for FP16, BF16, and FP32, - 3D/4D inputs, the real DFlash `[L, context, kv_heads, 128]` layout, a - non-power-of-two hidden size, and shape validation. Batched output is bitwise - identical to looping the original one-row RMSNorm. The repaired extension is - `7689e5e298be64a83c9f9743e08a073e857a5a4c89a79837396d05b37303d449`. -- The final 64-prompt repaired run averages 5.3888 per request and 4.5644 pooled, - versus 4.5395/3.9725 before repair and 5.2489/4.3571 for SGLang-V100. The - repaired-minus-control paired delta is 0.8493 with a 95% interval of - [0.7226, 0.9760]; it wins 60 of 64 matched prompts. Repaired per-position - acceptance is `[85.833,70.569,57.920,46.887,38.032,31.342,25.861]%`, above - SGLang-V100 at every position and within 0.0712 of the public official 5.46 - GSM8K mean. -- Acceptance recovery raises same-route aggregate output throughput from - 97.080 to 108.826 tok/s (+12.10%) and mean steady decode from 115.394 to - 134.940 tok/s (+16.94%). The repaired vLLM path remains slower than the - separately measured SGLang-V100 aggregate 156.063 tok/s, so acceptance is - closed but the remaining execution-speed gap is a distinct optimization item. -- Evidence is retained in - `/data/models/v100-dflash2-20260820/pr3/acceptance-rootcause/` as - `dflash2-probabilistic-graph-{zlabshuffle42,fixed-batchedrmsnorm-zlabshuffle42}-randomseed-first64-officialprompt-e5m2-o4096.json`; - the SGLang comparison is - `/data/models/v100-dflash2-20260820/sglang-audit/sglang-dflash2-graph-zlabshuffle42-randomseed-first64-officialprompt-e5m2-o4096.json`. - Earlier four-prompt and selector-pair acceptance values predate this repair - and remain route/equality diagnostics only; do not cite them as DFlash2 - acceptance baselines. A BF16 target rerun is intentionally unnecessary - because the causal A/B kept the same FP8 target and changed only K-norm row - selection. - -### DFlash2 single-request execution-cost root cause, 2026-08-22 - -- The repaired 64-prompt single-request runs close acceptance but not execution - cost. For requests with at least 512 completion tokens, the round-weighted - vLLM decode cost is 40.7148 ms/verification versus 26.8166 ms for the same - SGLang-V100 target, draft, TP4, E5M2-KV, block-8, sampling, prompt order, and - CUDA-Graph route. The 13.8982 ms gap is not caused by acceptance: vLLM's - per-request mean is 5.3888 versus SGLang's 5.2489. -- Matched CUDA-Graph node traces show that the one-pass DFlash2 proposal is not - the bottleneck. The vLLM and SGLang draft-model-plus-selector graphs take - 3.976 and 3.915 ms respectively under tracing. Their dense draft LM head, - five-layer draft forward, non-causal attention, selector top-K, and selector - walk therefore have effectively the same cost on this route. -- The raw traced round decomposition is: - - vLLM: draft graph 3.976 ms; draft-to-target transition 18.148 ms with 431 - eager kernels; target graph 24.845 ms; target-to-next-draft transition - 2.786 ms with 63 eager kernels. - - SGLang: draft graph 3.915 ms; draft-to-target transition 3.386 ms with 30 - eager kernels; target graph 19.622 ms; target-to-next-draft transition - 4.164 ms with 125 eager kernels. - - These traced wall times include launch-tracing overhead and must not replace - the unprofiled 40.7148/26.8166 ms round results. They are valid for - composition. SGLang captures its target LM head and TP all-gather inside the - target graph (about 1.11 ms combined), while vLLM executes those after its - target graph, so raw graph boundaries are not a perfectly semantic split. -- The dominant pre-verification gap is a hybrid-KV metadata amplification. The - target has 48 GDN and 16 full-attention layers, while the draft has five - layers. vLLM's unified grouping therefore creates ten padded GDN groups and - four padded full-attention groups. GDN metadata does not support cached - block-table updates, so `build_gdn_spec_decode_state_contract` repeats five - boolean-index selections for every GDN group. The trace contains exactly 50 - each of CUB reduce, compact-init, and select kernels per round, plus 281 other - metadata/input kernels: 431 launches and 18.148 ms of traced wall time for - only 2.452 ms of GPU service. SGLang keeps one persistent linear-attention - metadata object shared by all GDN layers; its corresponding transition has - 30 kernels and 0.117 ms of GPU service. -- SGLang also has a genuinely cheaper verifier graph: 1,130 nodes and 19.095 ms - GPU service versus vLLM's 2,612 nodes and 23.547 ms. Dense target FP8 GEMM is - already equal (10.798 versus 10.877 ms), and full-attention service is close, - so neither TurboMind GEMM nor Flash-V100 attention is the first DFlash2 - target. The main graph differences are: - - vLLM executes 1,422 generic vectorized/unrolled/elementwise nodes totaling - about 4.625 ms; SGLang executes 211 such nodes totaling about 0.677 ms. - - vLLM custom all-reduce costs about 2.350 ms per target graph versus 1.113 ms - for SGLang's one-stage push kernel. - - SGLang fuses sigmoid-gating computation, recurrent verification, and - intermediate-state writes. After rejection sampling, two all-layer fused - gather/scatter kernels commit the accepted SSM and convolution state. Its - target-to-draft interval also uses fused target-hidden-to-draft KV - materialization. -- Three narrow attempts to remove the repeated vLLM work were rejected and - fully reverted: reusing the DDTree one-kernel updater, reusing only its - pure-spec state path, and bypassing only the state-contract boolean - selections. All visible metadata regressions passed (34 focused tests), but - every shortcut changed the fixed-seed token trace at token 12, from the - 235-token/40-round control to a 309-token/57-round trajectory. This proves - that deleting the allocations/indexing before establishing explicit graph - buffer ownership is unsafe. The current evidence is consistent with a hidden - padding, alias, scratch-lifetime, or stream-order dependency; that precise - mechanism remains an inference rather than a proven field mismatch. -- Implementation decision and gates: - 1. Add a DFlash-only persistent GDN verify metadata/state object. Every live - and padded field is overwritten each epoch, and a debug shadow path compares - it with the existing generic builder before graph replay. - 2. Replace the ten repeated state contracts with one DFlash-specific fused - metadata kernel that writes shared structural metadata plus group-specific - state block IDs. Do not route Eagle, MTP, or DDTree through it. - 3. Port SGLang's fused target-verify GDN and two all-layer state-commit - kernels, including the per-step FP16 store/reload boundary needed for - recurrent-state numerical equivalence. - 4. Then port/capture fused target-hidden-to-draft KV materialization and tune - the one-stage TP4 all-reduce. Sparse rejection-sampling work is secondary. - 5. Do not reintroduce selector fusion until a new design clears the M=7 gate: - the measured proposal graphs are already equal, and the duplicate-layout - fused selector is slower and larger. - The first performance gate is an unprofiled round cost at or below 32 ms with - the fixed-seed token hash unchanged; the follow-up target is 28-29 ms before - claiming parity with SGLang's 26.8166 ms reference. Greedy tokens and accept - trajectories must be exact; probabilistic acceptance may differ by at most - 0.05 on the fixed corpus, with batch 1/2/4 mixed-length graph replay and - poisoned padding/state buffers included. -- Evidence is retained in - `/data/models/v100-dflash2-20260820/pr3/perf-rootcause/` as the vLLM baseline, - graph-node report, SQLite export, and rejected A/B JSON files. The matching - SGLang report is - `/data/models/v100-dflash2-20260820/sglang-audit/perf-rootcause/sglang-dflash2-single1-step20-v2.{qdstrm,nsys-rep,sqlite}`. - -### PR257 shared-metadata and packed-GDN checkpoint, 2026-08-22 - -- Draft PR #257 is isolated on - `codex/v100-dflash2-gdn-metadata-20260822-122723`; it is stacked on PR #254 - and changes only the MRV2 DFlash2 target-verification path. Eagle, MTP, - DFlash1, and `dflash_ddtree` retain their existing dispatch. Both candidate - leaves remain default-off while the TP4 profile and statistical-quality - gates are pending. -- The first leaf adapts vLLM PR #52297 at fixed head - `5f4c9678d1bb791ffb3ad23a9ae4129db9eb299d`: request classification, token - indices, and query offsets are computed once per batch, while cache-group - state block IDs remain group-owned. A same-hash/AAL 64-token Graph B1 probe - improved steady decode from 84.79 to 91.10 tok/s and decode time from 0.7430 - to 0.6915 s. The eager candidate regressed, so this evidence supports only - the production Graph route. An eight-request probabilistic probe changed the - aggregate acceptance trajectory and is not accepted as a compute-only speed - comparison. -- This port does not yet remove the dominant state-contract fan-out. Each of - ten padded GDN cache groups still performs group-specific block-table work - plus repeated boolean compaction for accepted-state metadata. The next - metadata step must be judged by the measured draft-to-target interval and - kernel count, not by the upstream H200 claim alone. Upstream reports its - builder moving from about 900 to 300 microseconds and B1 throughput +61%; the - V100 result remains an independent gate. -- A second DFlash2-only leaf removes packed-QKV rearrangement and the final GDN - output copy while preserving the existing split gating materialization. The - first implementation used `q * rsqrt(sum(q^2)+eps)` where the accepted - recurrent verifier uses `q / sqrt(sum(q^2)+eps)`. Although mathematically - equivalent, that changed at most one FP16 ULP in the operator test and moved - a fixed-seed probabilistic decision at output token 111. This explains why - the earlier fully fused and packed candidates followed the same alternate - trajectory. -- The repaired packed kernel now explicitly selects the recurrent verifier's - launch geometry and arithmetic order, including `/sqrt` normalization and - the same decay exponential helper. Gating remains precomputed in its original - dtype. On a physical V100, all eight production-shape parity cases - (`Hq=4`, `Hv=12`, `K=V=128` per TP4 rank; B1/B2; draft 3/7; FP16/FP32 - recurrent state) are bitwise equal for both output and persistent state. The - per-step low-precision state reload case also passes, and the complete shared - fused-sigmoid/recurrent file reports 27/27 passing tests. -- Before the arithmetic repair, the fully fused greedy one-request diagnostic - kept the exact token hash and acceptance length while improving steady decode - 99.15 to 103.45 tok/s (+4.34%) and decode time 1.2809 to 1.2277 s (-4.16%). - That number is retained only as evidence that this kernel family can reduce - work; it is not the accepted result for the repaired packed leaf. The repaired - leaf still needs a matched unprofiled B1 run and a control/candidate Nsight - pair before default enablement. -- The statistical quality gate is broader than greedy identity. Greedy hashes - remain a sensitive diagnostic, while default enablement also requires bounded - fixed-text logprob/perplexity change and no regression on fixed subsets of - GSM8K, MATH-500, HumanEval, and MBPP. The retained datasets and SHA256 values - are `GSM8K 3730d312...`, `MATH-500 35dc4108...`, - `HumanEval 2f2871a1...`, `MBPP e9e9efa2...`, and WikiText-2 validation - `204929b7...`; all live under - `/data/models/v100-dflash2-20260820/datasets/` and are not repository assets. -- The next Nsight pair must preserve TP4, Flash-V100, target E5M2 KV, draft - FP16 KV, block eight, single-request Graph execution, and the fixed request. - It will report draft Graph, draft-to-target metadata/input preparation, - target verifier Graph, rejection/target-to-draft, full round time, Graph node - count, rank-critical GPU service, and TP GPU-sum service. The existing control - trace lacks supported replay NVTX ranges and is therefore not fed through the - generic per-token parser; control and candidate will be recaptured with the - same replay markers rather than manufacturing a partial comparison. -- A stricter re-audit of the retained endpoint JSONs intersects the long-output - requests by `dataset_index` and requires both implementations to emit at - least 512 tokens. Across those 11 matched requests, vLLM uses 40.7194 ms per - verification round (`decode_time / num_drafts`) and SGLang uses at most - 26.7101 ms (`request_elapsed_seconds / spec_verify_ct`), leaving a gap of at - least 14.0093 ms. The SGLang field includes prefill and is therefore an upper - bound on its decode-round cost. The earlier 40.7148/26.8166 ms values use the - independently selected 13/14-request long-output subsets and remain useful - aggregate references, but they are not a strictly paired endpoint statistic. -- Re-parsing the existing node traces by complete round and taking the slowest - TP rank for every interval gives the following diagnostic wall table. These - values include graph-node tracing overhead and must not replace the endpoint - numbers above. - - | traced phase | vLLM mean / p50 / p90 / p99 (ms) | SGLang mean / p50 / p90 / p99 (ms) | vLLM / SGLang launches per rank | - | --- | --- | --- | ---: | - | draft Graph | 4.065 / 4.028 / 4.073 / 4.088 | 3.961 / 3.928 / 4.123 / 4.193 | 185 / 158 | - | draft to target | 18.382 / 17.660 / 21.258 / 22.690 | 3.517 / 3.423 / 3.660 / 4.066 | 431 / 30 | - | target Graph | 24.835 / 24.880 / 24.943 / 24.995 | 19.723 / 19.600 / 19.939 / 20.526 | 2612 / 1130 | - | target to next draft | 2.791 / 2.793 / 2.806 / 2.809 | 4.216 / 4.188 / 4.332 / 4.402 | 63 / 125 | - | complete traced round | 49.860 / 49.145 / 52.410 / 54.266 | 31.212 / 31.038 / 31.905 / 31.906 | n/a | - - Per-phase critical-rank means are diagnostic and are not strictly additive; - the complete-round row is measured directly from draft-Graph start to the - next draft-Graph start. -- The 431-kernel vLLM draft-to-target interval now has an exact launch - accounting. Repeated GDN state-contract compaction and its generic - indexing/copy/scan helpers consume 423 launches and about 1.276 ms of - rank-average GPU service; block-table, slot, RoPE, and input mapping consume - seven launches and about 0.044 ms; the final TP4 reduce/synchronization is one - launch and averages about 0.598 ms. That reduce is primarily a rank-skew wait - at this boundary (0.128 ms p50, 1.339 ms p90, 6.241 ms p99), not a stable - 0.598 ms compute kernel. SGLang reaches the target with 30 launches and about - 0.115 ms GPU service. This makes removal of the ten repeated group contracts - the first optimization, before tuning the boundary collective itself. -- Inside the target Graph, both implementations spend essentially the same - time in 256 target FP8 GEMMs (10.850 versus 10.570 ms of rank-average GPU - service). vLLM's excess is concentrated in 1,422 generic vectorized, - unrolled, and elementwise copy nodes (4.621 ms) versus SGLang's 211 nodes - (0.667 ms), plus the TP reduction path (2.350 ms on vLLM's slowest rank versus - 1.603 ms on SGLang's slowest rank). Attention and KV-cache service are close. - SGLang includes roughly 0.979 ms of target LM-head GEMM and 0.129 ms of TP - all-gather inside its target Graph, whereas vLLM launches the corresponding - work after the graph; semantic comparisons must retain this boundary - difference. -- With the repaired corpus-wide pooled acceptance length of 4.5644, the - current 40.7148 ms round cost corresponds to about 112 emitted tokens/s - before request-level residuals. A 32 ms round is only an intermediate gate - (about 143 tokens/s), and SGLang parity near 26.8 ms is about 170 tokens/s at - the same pooled acceptance. Reaching the requested 200 tokens/s without - relying on a higher acceptance length requires approximately 22.82 ms per - round. The optimization sequence is therefore persistent/fused GDN metadata, - removal of target-Graph copy/state fan-out, then the TP4 reduction path; - proposal/selector work is not on the critical gap. -- SGLang PR #26520's final-state recomputation path is not selected for this - B1 lane. Its own acceptance-aligned result is 2.5% lower throughput than the - cached-state control. vLLM already writes speculative intermediate states to - state-pool slots, so the nearer-term path is shared/persistent metadata and a - numerically equivalent verifier kernel. ReplaySSM remains a later high-batch - research lane, not a substitute for the current single-request verification - cost gate. -- The first persistent pointer-table implementation has now passed an isolated - SM70 microbenchmark at the production block-eight width and ten target GDN - cache groups. It overwrites group-specific state IDs plus every shared live - and padded field in one Triton launch and remains bitwise equal to the legacy - contract after poisoning all persistent outputs. On an otherwise-idle V100, - the B1 legacy fan-out uses 400 CUDA launches and 5.360 ms synchronized-wall - p50 versus one launch and 0.354 ms for the fused path, a 15.13x speedup and - 5.006 ms local saving. B2 is 5.360 versus 0.358 ms and B4 is 5.346 versus - 0.354 ms, both about 15x. A separate lower-bound probe measures the raw fused - kernel at 0.056 ms wall / 0.040 ms GPU p50; roughly 0.305 ms remains in Python - eligibility, pointer validation, and metadata-view construction. These are - microbenchmark results, not yet an endpoint round-cost claim. The benchmark - is `benchmarks/benchmark_sm70_dflash2_gdn_metadata.py`; an unprofiled TP4 - fixed-trajectory run remains required before default enablement. -- The persistent descriptor now also caches only the shape-specific - `GDNAttentionMetadata` views. Their backing state IDs, accepted counts, - selectors, masks, and query offsets are still overwritten by the fused - kernel on every replay. A poison-and-replay test proves that a same-shape - call reuses the view objects while observing new tensor contents, and a B2 - to B1 transition proves that a shape change rebuilds the views. On the same - B1/ten-group/block-eight microbenchmark, the cached fused path is 0.131 ms - synchronized-wall p50 versus 0.383 ms before view caching; the remaining - Python delta over the 0.053 ms raw-kernel wall is about 0.078 ms. Reinstating - the production scalar host fence costs only 0.028 ms p50, so it is not the - source of the former multi-millisecond Draft-to-Target gap and remains in - place pending the end-to-end trajectory gate. -- The first matched TP4 endpoint trace now validates that the microbenchmark - saving survives the real MRV2 DFlash2 path. Both runs use source - `11a190397abd7209ed1d3aeeeb8ed691e06087f7`, the same fixed shuffled GSM8K - request, Graph target and draft, block eight, probabilistic sampling, target - E5M2 KV, and draft FP16 KV. The control omitted the parent fast-path gate and - therefore exercised the generic metadata builder; the candidate enabled both - `VLLM_SM70_DFLASH2_VERIFY_FASTPATH=1` and - `VLLM_SM70_DFLASH2_FUSED_GDN_METADATA=1`. - - Control: 4.820868 s decode / 103 verification rounds = 46.8045 ms per - round, 88.99 steady-decode tok/s, and 86.11 aggregate tok/s. - - Fused candidate: 3.573145 s / 102 rounds = 35.0308 ms per round, 120.06 - steady-decode tok/s, and 114.20 aggregate tok/s. - - This is an 11.7737 ms (`25.15%`) reduction in traced endpoint round cost, - `+34.92%` steady-decode throughput, and `+32.61%` aggregate throughput. - The 430 output token IDs and SHA256 - `c3aa4bdf534d580213871cb17d31ee997db9cf2ae19ae4766646fe92d463a225` - are exactly equal and the answer remains correct. Acceptance length does - not regress: it moves from 4.1748 to 4.2157 (`+0.0409`). - - The candidate trace contains 412 fused metadata launches across TP4, - corresponding to 103 captured steps per rank (the request reports 102 - draft/verification rounds). Generic CUB `DeviceReduce` and `DeviceSelect` - disappear (`20,800 -> 0` each); `DeviceScan` falls from 20,008 to 19,816, - leaving scans owned by other paths. CUDA runtime fan-out falls from - 27,060 to 2,492 stream synchronizations, 90,524 to 16,328 async copies, - and 205,752 to 88,928 ordinary kernel launches across the four workers. - This proves that the wall saving is primarily removal of the repeated - host-synchronized metadata fan-out rather than the fused kernel's own - 1.528 ms aggregate GPU service. - - These endpoint values were collected under low-overhead Graph-level - Nsight tracing and are the accepted matched A/B for this leaf. The short - node-level trace below updates the semantic Draft-to-Target and - target-verifier phase table; it is used for composition only because node - tracing inflates absolute latency. - Artifacts are - `/data/minimax-h3/task-cache/v100-dflash2-pr257-gdn-metadata/results/dflash2-fused-viewcache-b1-graph-o512-{v2,v4}.json` and - `/data/minimax-h3/task-cache/v100-dflash2-pr257-gdn-metadata/profiles/latest-e2e/dflash2-fused-viewcache-b1-graph-v4.{qdstrm,nsys-rep,sqlite}`. -- A follow-up 128-token node trace at source - `9b8af0dbbe1d48214a555a96bbe762d26524e87c` supplies the missing semantic - phase split. The trace contains 32 complete draft-to-draft intervals; after - dropping the first and last edge intervals, 30 steady rounds remain. For - every phase and round the table takes the slowest of the four TP ranks, which - matches the earlier vLLM/SGLang audit methodology. - - | traced phase | old vLLM mean (ms) | fused vLLM mean / p50 / p90 / p99 (ms) | SGLang mean (ms) | fused launches per rank | - | --- | ---: | ---: | ---: | ---: | - | draft Graph | 4.065 | 4.046 / 4.048 / 4.078 / 4.090 | 3.961 | 185 | - | draft to target | 18.382 | 4.934 / 4.855 / 5.049 / 6.426 | 3.517 | 153 | - | target Graph | 24.835 | 24.740 / 24.747 / 24.892 / 24.907 | 19.723 | 2,612 | - | target to next draft | 2.791 | 2.838 / 2.840 / 2.849 / 2.934 | 4.216 | 63 | - | complete traced round | 49.860 | 36.300 / 36.257 / 36.424 / 37.827 | 31.212 | n/a | - - The persistent/fused metadata path therefore removes 13.448 ms (`73.16%`) - and 278 launches per rank (`64.50%`) from Draft-to-Target. The remaining gap - to SGLang in that phase is 1.417 ms and 123 launches, while the complete - traced-round gap is 5.088 ms. The dominant residual is now the target Graph: - 24.740 versus 19.723 ms (`+5.017 ms`), not the 4.046 ms draft Graph. -- The remaining 153 Draft-to-Target eager kernels per rank are stable across all - 30 steady rounds. They include one fused GDN metadata launch (0.0037 ms GPU - service), one TP4 one-stage reduce (0.2357 ms), 40 CUB scan/init launches - (0.1161 ms), 22 direct-copy launches (0.0686 ms), 20 index-select launches - (about 0.0926 ms), and 20 `compute_cuda_kernel` launches (0.0614 ms), - plus small slot/block/RoPE/input-mapping kernels. Critical-rank GPU service is - only 0.989 ms on average inside the 4.934 ms wall interval, so the next - Draft-to-Target leaf should share/capture the residual common-attention and - input-mapping fan-out rather than tune the 3.7-us fused kernel. In parallel, - the larger end-to-end opportunity is the verifier Graph's generic - elementwise/copy/state fan-out and TP4 reduction path. - Artifacts are - `/data/minimax-h3/task-cache/v100-dflash2-pr257-gdn-metadata/results/dflash2-fused-viewcache-b1-nodes-o128-v5.json` and - `/data/minimax-h3/task-cache/v100-dflash2-pr257-gdn-metadata/profiles/latest-e2e/dflash2-fused-viewcache-b1-nodes-o128-v5.{qdstrm,nsys-rep,sqlite}`. -- `dnv2003/v100-skinny` is useful as an optimization inventory, but its published - 219.1 tok/s result is a Qwen3.8 mixed-NVFP4/FP8 MTP-k7 route on its 1Cat-vLLM - 1.2.2 fork, not a DFlash2 result and therefore not a direct baseline. Its - chain-MTP GDN builder removes 21 device synchronizations and about 70 copies - for a reported 1.4 ms/step saving; PR257's persistent pointer table is the - stronger version of the same principle and now has direct DFlash2 evidence. - The fork's lagged CUDA-event/NVTX phase profiler is worth adapting to MRV2 so - verify, rejection, and proposal can be measured without adding a per-round - synchronization. Its partition pin is not active leverage at this run's - `max_model_len=8192`, and its FP16-KV win is specific to a route where FP8 KV - fell back to a scalar path; Flash-V100 reports the optimized E5M2 cache path - here, so any KV change still requires a matched A/B. The QPN8 verifier idea is - relevant, but the independently adapted block-FP8 implementation in Draft PR - #258 currently rejects speculative configurations. It must pass a dedicated - DFlash2 M=8 contract and quality gate before stacking onto this PR. - Main implementation priority: 1. Add a decoupled SM70 TurboMind backend for AWQ and FP8 base GEMM. @@ -1896,100 +1326,6 @@ Source behavior: - Decode path reads vLLM paged KV cache directly. - Prefill has direct and fallback paths, including paged KV gather fallback. -D256 8K-by-40K..128K GQA architecture checkpoint, 2026-08-25: - -- Frozen single-rank shape is causal FP16 Q/K/V, `Q=8000`, `KV=128000`, - `Hq=6`, `Hkv=1`, `D=256` on one V100-SXM2-32GB. The metric is useful - causal Attention FLOPs divided by complete Attention elapsed time; it is not - model TOPS or prompt throughput. -- The architecture separates a fully visible 32K..120K prefix from the causal - 8K tail. Six GQA heads are packed into GEMM M=48,000. QK writes FP16 logits and - row statistics; PV normalizes logits during its global-to-shared operand - transform. One reusable FP16 block partial is folded into a FP32 online - prefix numerator/state before the next block. The exact causal-tail kernel - exports max/sum state and a final kernel merges prefix and tail. -- The final BN8192 score-cache layout retains only the 750-MiB FP16 score - matrix per device. All other approximately 100 MiB of state is packed into - one transient aligned byte slab. A per-device event and mutex serialize - score reuse across streams. Before its first allocation, the route requires - driver-visible free memory for both workspaces plus 128 MiB of downstream - headroom; failure raises a typed OOM without touching the caching allocator - or launching a kernel, so the Python path can fall back safely. -- The final frozen-extension Torch A/B/A measured the existing exact operator - at `140.00213 ms / 43.53414 TFLOP/s` and the complete architecture at - `100.76979 ms / 60.48313 TFLOP/s`: `28.0227%` lower latency and `1.38933x` - speedup. All 12,288,000 outputs are finite; versus the exact operator, - max absolute difference is `2.2888e-4`, mean absolute difference - `2.5430e-5`, and relative L2 `6.8629e-3`. Peak allocated memory is - `1.11960 GiB`. -- The generalized operator admits all twelve observed `Q=8000` chunk shapes, - `KV=40000..128000` in 8000-token steps. Strict per-length A/B measured - `59.34862..60.74103` useful causal TFLOP/s. Against the active split-KV3 - control, every point is faster: latency is `22.3891%..26.1909%` lower and - speedup is `1.28848x..1.35485x`. All 147,456,000 candidate elements are - finite; the worst max absolute difference is `5.0354e-4` and worst relative - L2 is `6.8983e-3` versus the exact FP32-accumulator route. -- Smaller score blocks do not pass the target: BN7168 is - `104.32751 ms / 58.42057 TFLOP/s`; BN8000 is - `101.70283 ms / 59.92825 TFLOP/s` and must not be rounded to 60. BN8192 is - therefore retained while model KV reservation supplies downstream headroom. -- QK and PV compile at 254 and 119 registers/thread with no spill, stack, or - local storage. The complete operator is loadable through `torch.ops`; the - initial namespace registration mismatch was caught by a static load test and - fixed before timing. Synchronous device-symbol updates were changed to - current-stream asynchronous copies so the integrated measurement includes no - artificial whole-device fence. -- Integration is shape-family bounded and default-on. The legacy - `VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL=0` setting is the - rollback. It rejects CUDA-graph capture and falls back to the exact dense - kernel on workspace OOM. - The final external FA2 patch applies cleanly after the existing D256 - pipeline, split-KV3, and K-ping-pong patches, and a fresh CUDA 12.8 / Torch - cu128 SM70 build passes its operator-schema load test. The full Flash-V100 - policy file passes 112/112 tests on a real V100, including direct typed - architecture-OOM fallback to the exact dense route; Ruff and - `git diff --check` pass. -- Early real-model memory tests at `gpu_memory_utilization=0.88` rejected the - fully persistent and pre-preflight score-cache layouts after their first - successful route call: the following 80-MiB PYNCCL all-reduce had only - 62-74 MiB free. This is a downstream-headroom failure, not an Attention - compute failure. -- The final matched `.875` TP4 Qwen3.8-27B-FP8 - control/candidate/control prefill times are `46.45658 / 45.47797 / - 46.11195 s`. Against the `46.28426-s` bracketed control, the candidate is - `1.7421%` lower latency and `1.01773x` faster (`2831.89 -> 2882.10` - prompt token/s). Every candidate rank logs 16 architecture hits and no - architecture OOM/fallback. All three runs produce the same 32 output token - IDs and SHA256 - `df4fee7f5f0126fe6b391fe77b4fc19667831de5ef55fd69c28c2f52a3d7086e`. -- A final `.88` run also completes in `45.53427 s`, with the same token hash - and 16 route hits/rank. Its profiler left enough headroom for the route, so - this validates high-memory route success rather than directly exercising - the preflight fallback branch. -- The widened TP4 Qwen3.8-27B-FP8 gate uses a `46.11195-s` control before the - candidate and a `46.09222-s` control after it. Against their `46.10208-s` - mean, the `41.51191-s` candidate lowers prefill latency by `9.9565%` and - raises prompt throughput from `2843.08` to `3157.46` token/s (`11.0575%`). - Every rank logs exactly 192 family-route hits with no architecture OOM or - fallback. Both controls and the candidate emit the same 32 token IDs and - SHA256 - `df4fee7f5f0126fe6b391fe77b4fc19667831de5ef55fd69c28c2f52a3d7086e`. - Token identity is route-health evidence rather than a promotion - requirement. The stronger numerical gate covers 147,456,000 finite output - elements with worst max absolute error `5.0354e-4` and relative L2 - `6.8983e-3`, inside the accepted FP16 envelope of `1e-3` and `1e-2`. - Artifacts and the experiment ledger are retained outside Git; their private - filesystem location is intentionally not committed. -- The 2026-08-26 merge audit replayed all four patches from the locked clean - vendored baseline and rebuilt `_vllm_fa2_C` with CUDA 12.8 for `sm_70`. - QK/PV again compiled at 254/119 registers with zero spill. Current-source - direct A/B measured `39.29395 -> 30.27579 ms` (`1.29787x`) at KV40K and - `136.07628 -> 100.20284 ms` (`1.35801x`) at KV128K. Both outputs were fully - finite. KV40K had max absolute error `5.4932e-4` and relative L2 - `6.6098e-3`; KV128K had max absolute error `2.2888e-4` and relative L2 - `6.9166e-3`. These pass the FP16 merge envelope without imposing greedy or - bitwise identity and justify keeping the validated shape route default-on. - Latest target state: - No `FLASH_ATTN_V100` enum. @@ -42335,1418 +41671,3 @@ Interpretation: formatter changes from current main and keep large C/CUDA formatter work separate from accepted kernel repairs. Small policy substitutions still need their own current-source static/argument-equivalence checks before merge. - -## 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 - `MIXED_PRECISION`: dense projections use FP8, routed/shared experts use - `W4A16_NVFP4`, and 19 `mtp.*` tensors are present and unquantized. Exact SM70 - admission now requires both FP8 and NVFP4 TurboMind gates; unsupported model - shapes, TP8, EP, and all-to-all fail closed rather than taking a generic - fallback. -- The native Qwen3.6-35B-A3B MoE contract is H2048, expert I512, 256 experts, - top-k 8, replicated experts, and TP1/2/4. B1-B10 decode uses compact routed - slot groups. B11-B18 CUDA Graph/MTP decode and larger prefill use persistent - buffers with the full 256-expert grouped TurboMind operation. The grouped - prefill policy is default-on and may be disabled only for diagnosis with - `VLLM_SM70_NVFP4_MOE_GROUPED_PREFILL=0`. -- A real layer-0 TP4 checkpoint oracle matched the previous per-expert route: - M4096 W13/W2 were bitwise identical and improved 20.39x/9.48x. Full-grouped - B9 and B18 had cosine 1.0, maximum absolute error at most 1.221e-4, and - 21.59x-35.21x isolated-stage speedups. A rejected 144-group experiment - coalesced repeated expert IDs even though the forced-one-row scheduler - requires routed slots to remain independent; its roughly 0.495 W13 - relative-L2 error and catastrophic W2 output do not establish a general - 64-group correctness limit. The corrected slot-independent B9/B10 evidence - is recorded in the concurrency section below. -- Matching TP4, FP8-E5M2 KV-cache, Flash-V100/FlashQLA, full CUDA Graph, - input-4096/output-1024, max-length-8192 runs measured AWQ at 0.3813 s prefill - and 113.71 token/s steady decode. Default-grouped NVFP4 measured 0.4216 s and - 116.99 token/s: prefill is 10.56% slower than AWQ while decode is 2.89% - faster, placing both in the same performance class. Grouping reduced NVFP4 - prefill from 1.5479 s to 0.4216 s (3.67x). -- MTP4 resolved the bundled draft as `Qwen3_5MoeMTP`, captured FULL CUDA Graphs - through B18, and measured 0.8379 s prefill plus 174.76 token/s steady decode - under the same model contract. It accepted 792/924 draft tokens, with mean - acceptance length 4.43, and improved no-MTP decode by 1.49x. MTP prefill is - slower because the drafter participates. That original run still exposed - first-request sampler/Mamba JIT, which the later cold-start work below moves - ahead of request processing. -- The accepted cold-start patch warms the exact Mamba speculative postprocess, - all-greedy rejection sampler, and unquantized MTP-drafter MoE signatures. For - the checkpoint's 256-expert/top-8 draft, token counts 9, 33, and 257 cover the - sorted-assignment, unaligned-buffer, and SM70 small-to-large tile boundaries. - A fresh-cache ShareGPT16 run then produced no inference-time Triton JIT - warning: 97.69 end-to-end output token/s, 242.72 serial-prefill token/s, - 120.09 pure-decode token/s, and 8.827 ms mean TPOT. Moving these compilations - into startup adds about 4.33 s to model load versus the preceding - Mamba/sampler-only candidate. -- MTP quality is gated by dataset accuracy, validity, and repetition health, - not token equality with no-MTP. On pinned GSM8K test records 0-127 with - five-shot chat prompting, thinking disabled, greedy sampling, and a 512-token - cap, no-MTP scored 122/128 and three full-forward MTP4 runs scored 121, 122, - and 121. The post-rebase final at `d3b1ea3678` scored 121/128 with zero invalid - or repetitive records and the same wrong set as the preceding full MTP run: - `[2, 12, 37, 87, 89, 102, 119]`. Record 37 is a cap-induced truncation and is - correct at 1024 output tokens, so effective MTP4 accuracy is 122/128. A - standard-GDN diagnostic scored only 108/128 and is rejected. -- The bounded B10 source at `5422173e99` passes the same pinned GSM8K-128 - MTP4 gate at 122/128 (95.3125%), with zero invalid answers, zero repetitive - records, and wrong indices `[2, 12, 87, 89, 102, 119]`. It generated 19,434 - output tokens in 135.967 s. This is dataset-level evidence that the B10 - extension does not regress the accepted MTP4 quality band; no-MTP token - equality is not used as its oracle. -- A final default-grouped, no-MTP natural-text gate used two identical Chinese - prompts and one Python prompt. All three requests reached EOS with `stop`; - the duplicate outputs were token-for-token identical, and the code response - produced a correct `add(a, b)` implementation. This supersedes the earlier - 128-token truncated quality smoke for route acceptance. -- Reproducible JSON, logs, numeric oracles, and the rejected compact experiment - are under - `bench_results/modelopt_mixed_nvfp4_moe_20260822/{results,logs,telemetry}`. - The final dataset artifact is - `results/tp4_mtp4_fp8kv_gsm8k_chat_nothink_5shot_128_coldstart_rebased_final_graph.json`; - the all-JIT-warm ShareGPT artifact is - `results/tp4_gpu0123_mtp4_fp8kv_sharegpt16_seed20260822_all_jit_warmup_reserved_graph.json`. - The bounded-B10 dataset artifact is - `bench_results/qwen36_nvfp4_concurrency_scaling_20260824/results/candidate_final_mtp4_gsm8k_chat_nothink_5shot_128.json`. - These results are the accepted 35B evidence; earlier short route-hit and - quality smokes must not be substituted for the matched speed contract. - -## 2026-08-24 Qwen3.8 exact-D256 long-prefill K ping-pong - -- The Qwen3.8-27B-FP8 TP4 128K trace remains attention-led: exact-dense and - direct-paged D256 attention account for 35.55% and 6.22% of profiled wall. - The exact long shape sustains only 38.67% tensor-pipe active with 62.14% of - cycles lacking an eligible warp; DRAM is only 9.06%. This is an SM70 - shared-memory/dependency problem, not HBM saturation. -- Added a third vendored-FA2 patch that alternates the four D64 K panels across - two non-overlapping shared stages. The second stage borrows the later V/P - allocation because K and V/P lifetimes do not overlap. Its 2240-half stride - is above the padded 2188-half K `cosize`, aligned for 128-bit stores, and on - a common 128-byte bank phase. This removes three static full-CTA barriers - while preserving every HMMA/load/store count and the exact N32 arithmetic - order. -- Shared memory remains 45,568 bytes. Dense and split-KV3 remain at 254/253 - registers per thread with zero stack or spill. A clean CUDA 12.8 SM70 build - has byte-identical SASS to the measured candidate. The first clean-clone - build failed only because the CUTLASS submodule was absent (`cute/tensor.hpp` - not found); initializing pinned CUTLASS `62750a2b...` repaired the build. -- Exact dense A/B/A improves 0.74%-1.11% at Q8192/KV32K-128K and 0.54% at the - historical Q15680/KV125440 shape; a separate clean-build 128K run improves - 0.39%. The first Q8192/KV8192 chunk is neutral. Dense, randomized-paged, and - split-KV3 output gates are all bitwise equal to their matching controls. - Two order-reversed 128K comparisons pool to `142.4068 ms` control, - `141.8076 ms` fully-disjoint v5, and `141.7669 ms` lifetime alias. The final - form is `-0.45%` versus control and preserves the smaller shared envelope. -- Rejected 2048-half spacing because K stages overlap, 2188 because `STS.128` - becomes misaligned, and 2192/2304 because they are slower than the final - same-bank-phase layout. The exact fully-disjoint 2240 form is superseded - because it grows shared memory without a measurable speed advantage. Nsight - Compute counter collection is unavailable to this user - (`ERR_NVGPUCTRPERM`); no driver security setting was changed. -- The first full-model attempt failed before any request during FULL-graph XQA - metadata warmup. The source passed the new 19th - `batch_context_max_seq_len` argument, while the stale system Flash-V100 - binary exposed the old 18-argument ABI. Rebuilding the unchanged current - Flash-V100 source produced SHA `ac07dc9d...5580a8`; its 19-argument docstring - and a CUDA-graph E5M2 XQA scalar/padded bitwise gate pass. This harness repair - was used identically for both control arms and the candidate. -- Matched TP4 Qwen3.8 128K E5M2-KV, chunk-8192, official-sampling A/B/A passes. - Baseline A/candidate/baseline B prefill times are - `47.8992/47.6687/47.9475 s`. The bracketed control is `47.9233 s` or - `2735.0 tok/s`; the candidate is `47.6687 s` or `2749.6 tok/s`, giving - `-0.531%` latency and `+0.534%` throughput. TTFT changes - `47.9498 -> 47.6937 s` (`-0.534%`). Pure decode is neutral at - `5.1557 -> 5.1547 s` and `49.460 -> 49.470 tok/s`. -- All arms return 256 tokens with identical hash `2db7ef09...503325a`. - All four ranks report the same route counts: exact-dense D256 32, - prefix-paged 544, E5M2 bridge 544, scalar decode 176, and XQA decode 256. - The endpoint promotion gate is satisfied; the result is deliberately cited - as a small dependency-chain gain, not as completion of the broader 128K - utilization target. - -## 2026-08-24 Qwen3.8 Q8000 split-KV3 long-prefill experiment - -- A runtime-shape diagnostic closed the initial Q8192 assumption. With - chunk-8192 plus Mamba align, long prefix calls are actually Q8000; the first - Q8192-gated endpoint run had zero split-KV3 hits and is not performance - evidence. One control was additionally invalidated when another TP4 job - started between its 47.90-second warmup and 64.45-second measurement. -- Added an exact-shape, default-off - `VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3_Q8000_EXPERIMENTAL` policy. Existing - production Q4096 admission is unchanged; Q8192 and nearby shapes remain - rejected. The policy suite passes 89/89 and Ruff lint/format checks pass. -- Bracketed Q8000 operator sweeps at KV 40K/64K/96K/128K improve throughput by - 3.94%/2.85%/3.45%/3.36%. At KV128K the kernel moves from 138.9150 to - 134.3961 ms and 43.87 to 45.35 causal TFLOP/s. Workspace is 141.72 MiB per - rank. The FP16 outputs are repeat-stable but non-bitwise versus unsplit; - maximum absolute error is at most `3.052e-5` and mean error at most - `1.38e-6`. -- Matched TP4 128K A/B/A measures control A/candidate/control B at - `47.5940/47.1074/47.6612 s`. The bracketed control is `47.6276 s` or - `2752.0 tok/s`; candidate is `47.1074 s` or `2782.4 tok/s`, for `-1.092%` - latency and `+1.104%` throughput. TTFT improves 1.078%; decode is neutral. - All ranks report exactly 384 Q8000 split-KV3 hits and unchanged surrounding - prefill/decode route counts. KV capacity remains 2,424,439 tokens per rank. -- Each lane's warmup and measured 256-token stream is internally exact. The - candidate and control B are also token-for-token identical, but identical - unsplit control A selects a different stable sampled stream under the same - official seed. This cross-process sampling instability is not equivalence - evidence for a non-bitwise route. -- A deterministic greedy follow-up uses two identical Chinese exact-128K - prompts and one exact-128K Python task. Both routes are internally stable and - coherent, and both code responses stop normally with the same correct - `add(a, b)` implementation. Cross-route strict identity fails: Chinese first - differs at token 7 with 126/128 tokens aligned; Python first differs at token - 20 with all 67 control tokens aligned in the 70-token candidate response. - Candidate ranks each report the expected 576 split-KV3 hits and controls - report zero. Keep the Q8000 gate default-off until fixed-text - logprob/perplexity or dataset-level evidence proves no quality loss from the - classified Type-B reduction-order drift. Evidence is under - `/data/minimax-h3/task-cache/qwen38-fp8-128k-flashattention-20260824/`. - -## 2026-08-23 Qwen3.8-27B-FP8 TP4 no-MTP QPN8 decode acceptance - -- Integration base is `7aede2cf01`; the owned source worktree and Draft PR are - `codex/v100-qwen38-fp8-qpn8-operator-race-20260822-114943` and PR 258. -- Added a memory-neutral SM70 QPN8 weight layout and direct decode operators for - the accepted TP-local gate/up `(K,N)=(5120,8704)`, down `(4352,5120)`, and - output `(1536,5120)` projections. GDN input and full-attention QKV remain on - TurboMind. The old layout is replaced at load time rather than retained. -- Opt-in admission uses the engine contract only: SM70 block-FP8 layout, TP4, - exact projection shapes, no MTP, and configured `max_num_seqs <= 8`. - Native QPN8 handles M=1-8. - Wider-concurrency and MTP configurations retain TurboMind until their own - performance and quality gates pass; this is not a batch-one-only dispatch. -- The opaque C++ M dispatch prevents an AOTInductor M=1 trace from being reused - for prefill. M>8 remains functionally correct through one bounded 85 MiB FP16 - weight workspace, but fused gate/up additionally creates an `M x 8704` FP16 - temporary and is not accepted as a high-concurrency performance route. -- Complete-source no-MTP E5M2 TP4 A/B at input 1024 and output cap 256 measured - TurboMind at `16.961295 ms/token`, `58.958 tok/s`, and QPN8 at - `16.347973 ms/token`, `61.170 tok/s` (`+3.752%`). The candidate stopped - normally at token 233; the control finished all 256 tokens. An earlier - matched operator-integration run measured `15.808000 ms/token`, - `63.259 tok/s` (`+7.448%`) and remains headroom evidence, not the conservative - complete-source observation. Because these arms contain 255 versus 232 - steady intervals, respectively, this is not a pure-FP8 default-promotion - result; an equal-interval same-source rerun remains required. -- Nsight Systems attributes about `0.944 ms/token` of the trace improvement to - dense service (`10.335 -> 9.391 ms`), while launch count and GPU idle time are - unchanged. The QPN8 dense service spends 6.172 ms in the 75-90% DRAM band and - 7.059 ms in the 25-50% SM band; no dense kernel sustains at least 90% DRAM or - 50% SM. Useful dense arithmetic rises from 1.177 to 1.295 TFLOP/s per GPU. -- The quality gate passed without a measured regression: WikiText PPL/byte PPL - and bits/byte improved slightly, GSM8K flexible accuracy was equal and strict - accuracy improved by one item, all measured MMLU subjects were equal or - better, and all measured C-Eval subjects were equal. Sampled token identity - is therefore not used as the sole acceptance criterion. -- Full evidence and artifact paths are recorded in - `docs/design/sm70_qwen38_qpn8_decode.md`. Next work is a measured single-launch - M=9/11/16 row-tile route, followed by larger concurrency crossovers; do not - compose an eight-row launch with a separate tail or promote the current - large-M temporary as an accepted throughput implementation. - -## 2026-08-23 Qwen3.8-27B-NVFP4 TP4 decode recovery - -- Recovered the mixed NVFP4/channel-FP8 checkpoint without crediting Marlin: - accepted FP8 projection shapes use QPN8, remaining FP8 projections and the - LM head use TurboMind W8A16, and NVFP4 uses TurboMind W4A16. -- Added channelwise FP8 TurboMind packing, exact NVFP4 N32 selectors, E4M3 - Flash-V100 XQA, QPN8 split-12 for the K1536 x N5120 output projection, and - an exact B1/G6/D256 E4M3 p64 attention partition. -- Frozen TP4 I1024/O256, E4M3 KV, official random sampling, no-MTP, and full - CUDA graph decode first measured 71.732 tok/s at 13.941 ms/token with the - native sampler. After merging current `onecat/main` at `675a12dedc` and - rebuilding Flash-V100, the frozen request reconfirmed **71.342 tok/s** at - **14.017 ms/token**. The clean p128 control is 69.904 tok/s at - 14.305 ms/token. -- The p64 operator sweep saves 9.55-16.98 us per attention launch from sequence - lengths 1025-2049. All p64/p128 cases stay within one FP16 output ULP of the - scalar reference; the focused GPU suite passes 6/6. -- Compact and fused top-k20 sampler experiments were removed because their - sampled trajectories were not deterministic against the native route. They - are not part of either accepted p64 result. -- Post-merge checks pass: 12/12 focused CPU policy/dispatch cases, 6/6 E4M3 - XQA GPU numerical cases, Ruff lint/format, and `git diff --check`. The final - `_C` and Flash-V100 SHA256 values are `e0ea14d0...7624` and - `b418fed8...dce7`. -- The latest merged p64 Nsight Systems trace captures 63 decode replays on all - four ranks. Steady replay/GPU-union time is 14.621/14.055 ms, GPU activity - union is 96.136%, and the route launches 1140.8 kernels/rank/token. NVML - samples show 99.47% GPU busy-window duty but only 48.26% memory-active duty; - 69.40% of service has a grid-limited occupancy ceiling below 25%. Current NCU - counters remain unavailable because the driver rejects non-root counter - access with `ERR_NVGPUCTRPERM`; do not present NVML duty or effective payload - rates as achieved SM or HBM throughput. -- Full contract, route map, staged results, profile table, numerical evidence, - rejected paths, rollback controls, and artifacts are in - `docs/design/sm70_qwen38_nvfp4_decode.md`. - -## 2026-08-24 Qwen3.8-27B-NVFP4 TP4 no-MTP 80 tok/s acceptance - -- The frozen input-1024/output-256 TP4 contract now passes the requested - no-MTP speed gate. The production CMake C++17 sampler-sidecar run measures - **80.624 tok/s** at **12.403 ms/token**. Three behavior-equivalent hand - sidecar repeats measure 80.177, 80.164, and 80.026 tok/s, including a repeat - on physical GPUs 4-7. Pure decode excludes TTFT and prefill. -- The deployment retains the proven-faster exact-shape QPN8 projections. GDN - input, full-attention QKV, and the LM head continue through TurboMind W8A16; - non-admitted NVFP4 cases retain TurboMind fallback. The accepted batch-one - NVFP4 shapes use native QPN4, E4M3 attention uses Flash-V100 XQA p64, and TP4 - communication uses the pack32 one-stage all-reduce. MTP is disabled. -- The quality gate passes at the frozen baseline: the production sidecar scores - **226/250 (90.4%)** on the first 250 GSM8K questions with five-shot greedy - prompting and has zero invalid outputs. The baseline is also 226/250; the - hand-sidecar result is 227/250. Per-question prompts, texts, token IDs, - extracted answers, labels, correctness, and hashes are retained in the JSON. -- The production sidecar and hand-built validation library select identical - tokens for 100 independent seeds and 256 sequential Philox draws and finish - with the same generator state. This exact implementation supersedes the - earlier rejected compact/fused sampler experiments, which did not preserve - native Philox behavior. -- The accepted compatible main `_C`, production sampler sidecar, and - Flash-V100 SHA256 values are `a0a0cd9...d36d`, `cdbfdd87...9ae`, and - `b418fed8...dce7`. The sampler is intentionally a separate CMake/wheel module - so registering it does not relink the frozen main library. -- The latest 63-replay/rank Nsight Systems trace shows a 13.430 ms mean replay - interval, 12.887 ms GPU activity union (95.958%), 0.543 ms idle, and 1061.9 - launches/rank/token. QPN8 split-16, QPN4 gate/up, TP all-reduce, QPN4 down, - QPN8 output, XQA, QPN8 gate/up, and TurboMind FP8 service contribute 2.213, - 2.165, 1.587, 1.441, 0.915, 0.618, 0.496, and 0.410 ms/token respectively. -- NVML reports 100% busy-window duty, 53.25% memory-active duty, about 28.97 - GiB allocated per GPU, and 56.36% of the 300 W limit. Grid-limited occupancy - ceilings put 86.32% of service below 50% occupancy. These are not achieved - SM/Tensor/HBM counters; Nsight Compute remains blocked by - `ERR_NVGPUCTRPERM`. -- Current-source rebuild closure is complete at source snapshot `7bd1d783e4`. - A clean primary `_C` plus clean C++17 sampler measured 77.314 tok/s in the - cold A1 observation and **80.181 tok/s** at **12.472 ms/token** in A2. A2 - exactly matches the accepted C++17 256-token stream and records - `speculative_config=None` plus `spec_decoding_metrics=null`. -- The clean primary and sampler SHA256 values are `c5d7c6a9...3e56` and - `3058051f...bc74`. Their `.nv_fatbin` SHA256 values, `dbba7402...26e6` and - `9a89e8cc...b2c`, are byte-identical to the accepted binaries. Full ELF - hashes differ only across host/link content: merged main adds two unrelated - ModelOpt NVFP4 MoE host entry points, while the sampler includes host - build-ID differences. These do not enter the accepted dense batch-one graph. -- The clean A2 result SHA256 is `4e6ed330...bb1`; its serialized token-array - SHA256 is `f50e0eba...e32`, exactly matching the accepted production C++17 - run. Device-binary and sampled-stream identity carry forward the already - passed **226/250 (90.4%)**, zero-invalid GSM8K gate. The separate sampler - module remains the packaging boundary, but current-source primary rebuilds - are no longer rejected merely because the full ELF hash changes. -- Focused CMake installation puts the modules at `vllm/_C.abi3.so` and the - CPython-3.12 SOABI sampler filename, strips build-tree RPATHs, and preserves - both accepted fatbins. The installed SHA256 values are `d7ed490e...6749` and - `8ee84663...55e`. Loading both staged files registers QPN4, QPN8, and the - packed top-20 sampler schemas. This closes the relevant install boundary - without rebuilding 50 unrelated wheel-target objects. -- Final focused regressions pass 94/94 target sampler, NVFP4-admission, and - SM70 TurboMind-adapter tests; the post-merge four-file run including the - adjacent ModelOpt route passes 106/106. Ruff lint/format, Python byte - compilation, the sidecar wheel filename gate, and `git diff --check` also - pass. -- Full results, route details, quality artifacts, trace/resource tables, and - exact evidence paths are in `docs/design/sm70_qwen38_nvfp4_decode.md`. - -## 2026-08-24 Qwen3.8-27B-FP8 exact-8K CUTLASS projections - -- The accepted starting point is TP4 V100, input 8000, output 1, FP8-E5M2 KV, - Flash-V100/FlashQLA, FULL_AND_PIECEWISE CUDA Graph, bounded exact-dense FP8 - projections, and the fused TP4 reduce-scatter/Gemma norm/all-gather boundary. - Its request-wall median was `1.616075 s`, or `4950.27 prompt tok/s`. -- The fused graph trace put `506.081/254.947/93.510 ms` into 64 gate/up, down, - and attention-output FP16 GEMMs. Their current useful rates were only about - `90.2/89.5/86.1 TFLOP/s`, despite the V100 SXM2's 125-TFLOP/s Tensor Core - peak. QKVZ plus QKV retained another `241.625 ms` at exact 8K. -- NVIDIA CUTLASS profiler selected its generated SM70 - `f16_s884gemm_f16_128x256_32x2_nn_align8` operation for all five exact TP4 - shapes. It uses CTA `128x256x32`, warp `64x64x32`, native `m8n8k4`, two - stages, and split-K one. The compiled kernel uses 220 registers/thread, - 256 threads, and no stack or local spill. -- Real-checkpoint CUDA Graph operator tests passed bitwise equality. The key - complete-boundary medians improved gate/up `8.67876 -> 8.26086 ms`, down - `4.26622 -> 4.02442 ms`, and attention output `1.54365 -> 1.48992 ms`. - Gate/up reuses the same graph-captured full projection temporary and existing - interleaved SiLU kernel; it adds no resident weight or second workspace. -- The final source-matching accepted extension SHA256 is - `6a7e8493ca04bd7738c9b016a6a0b4914b2c582598c920f99bb67626ae3e11c1`. - It is default-on only at `M=8000` for TP4 shapes - `(K,N)={5120x4096,5120x3584,5120x8704,4352x5120,1536x5120}`. Decode, - tails, other M values, and other shapes retain existing dispatch. Rollback is - `VLLM_SM70_FP8_PREFILL_CUTLASS=0`. -- The source-matching final full-model run produced six steady request-wall - samples `1.559085/1.561460/1.560481/1.562664/1.564697/1.564643 s`. - Their `1.562062 s` median is `5121.44 prompt tok/s` with `0.132%` CV. The - pure prefill median is `1.547100 s`, or `5170.96 tok/s`. This saves - `54.013 ms`, raises wall throughput `3.46%`, and clears the 5000-tok/s - requirement by `121.44 tok/s`. -- Route proof contains exactly 20 one-time messages, five admitted shapes on - each of four TP ranks. The accepted 766 graph-address registration count and - fused collective backend remain intact. Seven of seven official-sampling - outputs remain token `[1061]`, text `" This"`. -- The side-stream two-workspace FP8 dequant prototype is rejected: placing its - two dequantizations across the RMSNorm window changed - `7.98956 -> 8.00819 ms` (`0.99767x`) despite bitwise equality. The - cuBLAS-algorithm-106 candidate - was also rejected. Both implementations were removed; do not repeat them as - substitutes for the accepted CUTLASS shape route. -- Final evidence is under task-cache - `qwen38-fp8-prefill-utilization-20260823/results/`, especially - `full_model_cutlass_all_projections_final_source_tp4_fused_graph_i8000.raw.json`. - -## 2026-08-24 Qwen3.6-35B NVFP4 MTP4 concurrency scaling - -- Audit disposition: the B10 compact-routing extension has its own matched - speed and GSM8K evidence and remains enabled. The broader high-width tune - limit, exact-shape drafter tile, all-reduce block policy, and eight-warp - sampler remain opt-in because their matched formal curve and combined - dataset-quality gate are incomplete. - -- The formal workload uses ModelOpt mixed FP8/NVFP4, TP4 on V100 GPU0-3, - Flash-V100/FlashQLA, FP8-E5M2 KV cache, MTP4 probabilistic sampling with - local argmax, max length/batched tokens 8192, and 48 fixed ShareGPT requests - totaling 9,115 prompt and 10,623 output tokens. Forty-eight requests give - C16 three complete waves and remove the tail-underfill artifact in the - earlier 16-request localization sweep. -- Current-main output throughput at C1/C2/C4/C8/C12/C16 is respectively - 98.265/146.124/276.600/414.315/491.270/542.982 tok/s. Scaling efficiency - relative to C1 is 100.00%/74.35%/70.37%/52.70%/41.66%/34.54%. Mean MTP - acceptance length remains 2.887-2.955 and draft acceptance remains - 47.18%-48.87%, so the scaling loss is not an acceptance collapse. -- The old AWQ scheduling pair is rejected for this route. Shared-expert stream - overlap plus `all_reduce_sum2` regressed corrected C16 output 4.10%; isolated - overlap and sum2 regressed 4.66% and 4.06%. The pair also regressed C1 from - 98.265 to 94.880 tok/s (3.45%). Leave both switches default-off for NVFP4. -- MTP4 verifier width is five times request concurrency. Before this work, - production warmup stopped at NVFP4 MoE M15 while captured C4/C8/C12/C16 - shapes require M20/M40/M60/M80, or 160/320/480/640 routed rows. -- Clean independent-process graph microbenchmarks on one exclusively claimed - V100 show measured routed-MoE W13/W2 speedups at M20/M40/M60/M80 of - 1.107x/1.311x, 1.200x/1.179x, 1.195x/1.181x, and 1.177x/1.117x. Shared-dense - W13/W2 speedups are 1.742x/1.261x, 2.749x/1.181x, 2.004x/1.179x, and - 1.519x/1.348x. All outputs are finite; maximum accumulation-order difference - is 6.104e-5, which still requires the dataset-level quality gate. -- The high-width candidate raises only the NVFP4 MoE tune limit to 640 routed - rows, derives warmup widths from the already-captured CUDA Graph sizes, and - builds balanced 256-expert offsets above 32 verifier tokens. The dense-NVFP4 - M80 extension was isolated and withdrawn; its default remains M16. In a - 48-request C16 candidate/control/candidate sandwich, routed-MoE tuning - measured `535.635/524.376/539.113 tok/s`; the candidate endpoint mean is - `+2.48%`, summed pure decode is `+2.85%`, and P50 TPOT improves about `4.0%` - with stable MTP acceptance. - Production keeps the prior 128-row default; reproducing this candidate - requires `VLLM_SM70_NVFP4_MOE_TUNE_MAX_TOKENS=640`. -- Two earlier low-width target-MoE candidates are rejected. Directly reusing - sorted expert IDs instead of the compact metadata copy regressed the C1 - sandwich; splitting B9/B10 into 64-slot compact chunks was numerically - healthy but `41.5%/10.3%` slower. The old conclusion that every direct - compact execution above 64 slots was invalid was too broad: that experiment - coalesced repeated expert IDs even though TurboMind's forced-one-row - scheduler requires each routed slot to remain an independent group. -- A C1 Nsight graph-node trace identified the bundled unquantized MTP drafter - MoE as a separate `2.825 ms` per-cycle hotspot. An initial oracle incorrectly - used the checkpoint-global I512 shape. The first candidate startup log proved - that TP4 shards it to local I128 (`E=256,N=128` config lookup); the candidate - therefore did not route-hit and its full-model run was stopped before - measurement. The global-I512 speedups are not TP4 selection evidence. The - corrected local E256/H2048/I128/top-k8 M1-M16 graph oracle is bitwise equal - at every width. The isolated M1 operation improved `1.290x` by retaining - N32/K64 and reducing to two warps; M2-M16 select M8/N128/K32/W4/S3 and - improve `1.543x/2.656x/2.879x/2.697x/2.625x` at the formal - C2/C4/C8/C12/C16 widths. -- The route-hit candidate initially exposed request-time old-tile MoE - signatures. The exact E256/H2048/I128/top-k8 TP4 contract can warm tuned - M1-M16 plus M33/M257 - and explicitly covers legacy M1/M2/M9/M10 naive/aligned signatures; a fresh - 48-request run then produced no inference-time `fused_moe_kernel` JIT. - Unmatched shapes retain the old `(9,33,257)` warmup set, and the tuned route - requires `VLLM_SM70_MTP_MOE_TUNED_CONFIG=1`. -- Full-model evidence rejects the isolated M1 result. A physical-GPU4-7 C1 - candidate/control/candidate sandwich measured endpoint-mean summed pure - decode `112.233` versus control `115.451 tok/s` (`-2.79%`) and endpoint-mean - P50 TPOT `8.874` versus `8.758 ms` (`+1.33%`). M1 therefore retains the exact - 0.0.3 tile; balanced single-GPU kernel means are not a substitute for TP4 - rank-imbalance and synchronization behavior. -- The narrowed M2-M16 candidate passes C2. Candidate/control/candidate output - was `154.865/147.180/163.480 tok/s`; endpoint mean is `+8.15%`. Summed pure - decode improves `+4.47%`, P50 TPOT is effectively flat (`+0.14%`), and mean - acceptance is lower than control rather than artificially favorable. M17+ - retains the old prefill tile. -- Every C1 graph all-reduce in the Nsight trace is the exact 10,240-element - FP16 verifier payload. A TP4 operator sweep found that one CTA is faster for - this 20-KiB payload, while eight CTAs are faster only for the exact - C8/C12/C16 verifier payloads; C2/C4 retain the default. A physical-GPU4-7 C1 - candidate/control/candidate sandwich with the one-CTA launch measured - endpoint-mean pure decode `116.405` versus `112.842 tok/s` (`+3.16%`), output - `103.070` versus `100.112 tok/s` (`+2.96%`), and P50 TPOT `8.678` versus - `8.750 ms` (`-0.82%`). The source heuristic is limited to fully connected - SM70 TP4, the exact MTP4 payloads, and ordinary all-reduce; the global block - override still takes precedence and - the policy is enabled only by - `VLLM_SM70_TP4_MTP_AR_BLOCK_TUNING=1`; unset/zero retains the legacy policy. -- The C1 graph also spends about `0.509 ms` per cycle in the combined - top-k/top-p kernel over five 248,320-vocabulary verifier rows. Keeping the - existing 8192/4096 tiles and masking algorithm but launching eight rather - than four warps improves isolated B5/B10/B20/B40/B60/B80 means by - `1.737x/1.943x/1.881x/1.732x/1.689x/1.646x`. Every candidate had zero finite - mask or retained-value mismatches against the four-warp result. The opt-in - source change is restricted to combined top-k plus top-p on SM70, vocabulary - 248,320, and measured row counts 5/10/20/40/60/80. Enable it with - `VLLM_SM70_TOPK_TOPP_8_WARPS=1`; unset/zero keeps Triton's default launch. -- Corrected slot-independent compact metadata is finite and route-correct on a - real TP4 layer through B20. B9 is bitwise equal to the dense route. B10 has - cosine at least `0.99999988`, relative L2 at most `1.0115e-4`, and maximum - absolute error `2.38e-7`; its isolated graph stage is `1.091x` faster. The - production candidate stops at B10/80 routed slots and prewarms the exact - B1-B10 metadata signatures before CUDA Graph capture. -- A physical-GPU4-7 C2 candidate/control/candidate sandwich accepts B10. - Output was `154.608/147.135/160.890 tok/s`; candidate-mean output improves - `7.21%`. Pure decode was `91.962/89.762/94.566 tok/s`, a `3.90%` candidate - mean gain. Candidate mean acceptance length was lower than control - (`2.909` versus `2.963`), so acceptance does not explain the win. Extending - compact routing through B20 is rejected by a matching C4 sandwich: candidate - mean output/pure decode regressed `1.90%/2.48%`, with P50 TPOT about `1.18%` - worse. B11+ therefore keeps the 256-expert dense-grouped route. -- The final B10 48-request curve on one exclusive physical-GPU4-7 claim is: - - | concurrency | output tok/s | efficiency vs C1 | pure decode tok/s | serial prefill tok/s | P50 TPOT ms | mean acceptance | - | ---: | ---: | ---: | ---: | ---: | ---: | ---: | - | 1 | 97.693 | 100.00% | 114.247 | 564.795 | 8.522 | 2.890 | - | 2 | 150.736 | 77.15% | 90.851 | 423.341 | 10.207 | 2.946 | - | 4 | 246.399 | 63.05% | 71.322 | 596.885 | 13.559 | 2.884 | - | 8 | 350.834 | 44.89% | 52.398 | 629.601 | 20.006 | 2.969 | - | 12 | 444.117 | 37.88% | 46.313 | 631.019 | 22.274 | 2.909 | - | 16 | 530.810 | 33.96% | 44.734 | 720.400 | 23.478 | 2.900 | - -- Relative to the preceding final-source single-run curve, output/pure-decode - movement was C1 `-2.73/-3.08%`, C2 `+3.13/+2.41%`, C4 - `-0.73/-0.18%`, C8 `-2.93/-2.99%`, C12 `+4.21/+3.93%`, and C16 - `-4.47/-6.38%`. This alternating high-width movement is not attributable to - the bounded B10 route; B11+ still uses the same dense-grouped main shape and - only drain-tail widths can reach B9/B10. Treat B10 as a low-concurrency win, - not as high-concurrency recovery. -- The next target-side candidate must compact GPU-resident active expert IDs - for B20/B40/B60/B80 and use a device-count guard with an exact full-group - fallback. It may reduce padded grouped scheduling without a host readback or - a probabilistic correctness assumption. A persistent tile scheduler remains - the larger follow-up if guarded active-group bucketing cannot beat its extra - CUDA Graph launch. B10 has passed the pinned GSM8K-128 MTP4 gate at 122/128, - with zero invalid or repetitive records; the next quality gate belongs to - any high-concurrency scheduler candidate that survives full-model A/B/A. - -### Matched no-MTP scaling control - -- A matched performance-only control disabled MTP while preserving the B10 - source, model, TP4 physical GPU4-7 claim, 48 fixed ShareGPT requests, - sampling, FP8-E5M2 KV cache, TurboMind ModelOpt mixed FP8/NVFP4 route, - Flash-V100/FlashQLA, max length/batched tokens 8192, and - `FULL_AND_PIECEWISE` CUDA Graph policy. Every point reported - `speculative_config=None` and null speculative metrics, generated the exact - expected 9,115 prompt and 10,623 output tokens, and had no timed-request - Triton JIT warning. This control is not an MTP quality oracle; dataset-level - MTP quality remains gated independently. -- The no-MTP curve is: - - | concurrency | output tok/s | efficiency vs C1 | pure decode tok/s | serial prefill tok/s | P50 TPOT ms | - | ---: | ---: | ---: | ---: | ---: | ---: | - | 1 | 92.260 | 100.00% | 111.256 | 454.959 | 9.007 | - | 2 | 141.946 | 76.93% | 84.615 | 432.510 | 10.197 | - | 4 | 210.221 | 56.96% | 61.679 | 385.377 | 12.333 | - | 8 | 305.539 | 41.40% | 46.288 | 330.254 | 20.327 | - | 12 | 388.944 | 35.13% | 41.609 | 288.067 | 25.511 | - | 16 | 448.918 | 30.41% | 35.490 | 688.532 | 29.957 | - -- Relative to this no-MTP control, MTP4 improves output throughput by - `5.89/6.19/17.21/14.82/14.19/18.24%` at C1/C2/C4/C8/C12/C16 and improves - pure-decode throughput by `2.69/7.37/15.63/13.20/11.30/26.05%`. - MTP4 scaling efficiency is `0.00/0.22/6.09/3.49/2.75/3.55` percentage - points higher. Disabling MTP therefore neither recovers the absolute - high-concurrency throughput nor improves scaling; verifier/drafter work is - not the primary scaling bottleneck. Continue on the shared target decode - path, especially high-width routed-MoE scheduling and TP4 synchronization, - while retaining the independent MTP dataset-quality gate. -- The first completely cold C1 process also exposed a startup-only anomaly: - FULL decode graph capture took 230.21 s, total graph capture took 294 s, and - recorded load time was 380.46 s. Subsequent independent processes captured - FULL decode graphs in less than two seconds and total graphs in 35-68 s. - This one-time cold-cache cost was outside every timed throughput interval; - investigate it separately rather than attributing it to steady-state decode. -- Reproducible JSON and route logs are under - `bench_results/qwen36_nvfp4_no_mtp_scaling_20260824/{results,logs,telemetry}`. - -## 2026-08-24 Qwen3.8-27B-NVFP4 GDN decode fusion - -- Frozen contract: TP4 on four V100-SXM2-32GB GPUs, input 1024, output cap - 256, official random sampling, E4M3 KV, Flash-V100, full CUDA graph, and no - MTP. Pure decode excludes TTFT/prefill and each endpoint is the stable median - of three sequential requests. -- The quality-safe QPN8 QKV/Z plus FP16 b/a split projection and one-pass GDN - RMSNorm improve the matched control from `12.249652 ms / 81.634970 tok/s` - to `11.921648 ms / 83.881019 tok/s`. All three requests exactly reproduce - the control's 256 token IDs and SHA256 - `8b37337f4c393711cb8550db6bae909b1e85de8df1cf5ba8c60d8c000749c0a2`. -- A final post-rebase default-wrapper admission recheck reports - `12.261727 ms / 81.554583 tok/s` for the matching control and - `11.956949 ms / 83.633383 tok/s` for the candidate. The C++ route oracle - fires on all four ranks and all three 256-token candidate streams exactly - match the current-main control hash - `ca77db3b032a1600a8567adea706108c1bd8c5472b3ace2318c41ab17c66c1f9`. -- The GDN RMSNorm 48-layer operator gate is bitwise exact and saves - `55.946 us/token`. QKV is exact on all 48 layers; b/a has relative L2 - `6.368e-8` and maximum absolute error `2.384e-7`. The split projection is - admitted only when its one-pass RMSNorm pair is also enabled. Runtime - selection uses SM70, exact tensor/layout, QPN8 workspace, bias, TP, and - no-speculation contracts; the measured model/checkpoint is evidence only. - Large-M and unsupported shapes retain the original paths. Archived evidence - uses the pre-audit Qwen-prefixed flag names; final source uses generic GDN - flags and leaves the measured kernels unchanged. -- QPN8 QKV/Z plus b/a alone reaches `83.740763 tok/s` but diverges at token - 202, so it is not an independent accepted configuration. It is accepted only - together with the exact one-pass RMSNorm, whose paired endpoint restores the - complete frozen output. -- The Q/K RMSNorm plus MRoPE plus KV-write candidate reaches stable repeats of - `85.451/85.015 tok/s` but diverges at token 1 and stops at 208 tokens. Its - source, schema, flag, and rejected QPN4 communication experiments were - removed before submission. -- Next target: quantify no-MTP pure-decode decay at exact final context lengths - 16K, 64K, 128K, and 256K, reporting TTFT/prefill separately and preserving - the same model, quantization, TP, attention, KV, sampling, and graph routes. - -## 2026-08-24 MRV2 DFlash2 score-gated verifier promotion - -- The native Flash-V100 grouped verifier keeps all eight MRV2 target rows and - does not reuse DDTree indices. Its one-pass paged-GQA kernel is 0.356/1.214/ - 2.396 ms at 32K/128K/256K versus 1.024/3.021/6.011 ms for the previous XQA - verifier. Matched FULL-Graph endpoints reduce complete rounds from 36.311 to - 26.361 ms at 32K and 69.855 to 40.778 ms at 128K; paired output and - acceptance counters match. -- Two independent short-context traces of the complete pre-grouped bundle - measure 19.897/19.943 ms mean and 19.896/19.948 ms median. Their maxima are - 20.099/20.146 ms, so this is a replicated mean/median sub-20 result, not an - all-tail claim. -- The final quality contract is probabilistic sampling at T=1/top-p=.95/top-k - 20/xhigh, maximum 2,048 tokens, and natural EOS. Greedy token identity is not - a hard gate. Grouped matches target-only at 83/96 across GSM8K, MATH-500, and - HumanEval, and at 25/32 on corrected MBPP. Paired McNemar p is 1.0 for both - aggregates. WikiText PPL changes 5.497872 -> 5.497376 (-0.0090%), so no - measured score or PPL regression remains. -- The invalid raw-MBPP speed prompts omitted function names/tests and are - permanently excluded from intelligence claims. Official DFlash2 publishes - acceptance and throughput but not task pass@1; local target-only is the - quality baseline, while official acceptance is reference-only. -- Audit disposition: the acceleration leaves remain explicit opt-ins and use - hardware, tensor layout, algorithm configuration, cache/graph, and bounded - concurrency contracts rather than model or checkpoint identity. The - recorded profile remains evidence for the measured workload, not a runtime - model whitelist. Full evidence is in - `docs/design/sm70_dflash2_target_graph_20ms.md`. - -## 2026-08-24 DeepSeek V4 DSpark verifier and long-context campaign - -- The active TP8/V100 campaign, exact baseline contracts, admitted focused - gates, rejected paths, long-context sweep, and paired quality requirements - are recorded in - `docs/design/sm70_deepseek_v4_dspark_verifier_20ms_long_context.md`. -- Latest-main DSpark has reproduced a 116.288 token/s median on the historical - 3,500-token endpoint contract. The current complete verifier is 35.120 ms per - round; the active target is at most 20 ms without acceptance or dataset loss. -- Do not cite indexer, mHC, compressor-ring, or MoE microbenchmarks as the final - verifier result. Promotion requires a synchronized TP8 full-model profile, - long-context 1K-252K evidence, paired no-speculation quality, and task-owned - GPU cleanup. -- Source-audit disposition: incomplete-risk leaves are retained only as - explicit opt-ins. `VLLM_SM70_INDEXER_DECODE_CUBLAS=1` enables gather-plus-FP32 - cuBLAS, `VLLM_SM70_DSV4_PRIVATE_COMPRESSOR_STATE=1` enables private state - rings, and speculative context buckets require an explicit list. Default - serving does not compute or retain confidence logits unless a threshold or - alignment diagnostic consumes them, and a requested missing confidence - projection fails before use. -- Runtime gates express hardware, operator configuration, tensor/layout, - graph, cache, and concurrency contracts. The architecture-name check was - removed; measured checkpoint details remain evidence only and never select - an optimization route. - -## 2026-08-24 PR #276 cache-block reuse audit - -- The reported corruption root cause is valid: a reused full-attention cache - block can retain non-finite values in slots that its new owner has not yet - written, and arithmetic masking does not make `0 * NaN` finite. The MRV2 - runner had omitted the cache-zero lifecycle hook already used by the legacy - runner. -- MRV2 now initializes the shared `KVBlockZeroer` metadata and consumes the - scheduler's `new_block_ids_to_zero` before either attention or a zero-token - early return. Work is proportional to newly allocated blocks, uses the - cache-group/operator layout metadata, and introduces no per-step sequence - length readback or request-by-layer launch loop. -- The PR's per-step tail walker and sampled-token vocabulary clamp are - rejected. The former performs device-to-host synchronization plus many - small zero launches on every step; the latter turns corrupted output into - an arbitrary boundary token and hides the producer. The proposed GDN PAD - replication and last-column gather clamp are also rejected: the align-mode - allocator already guarantees a live current block plus its speculative - scratch window, so those fallbacks would conceal a broken block-table - contract and repeat an unrelated state block. -- Route selection and correctness depend only on allocation lifecycle, - cache-group layout, block IDs, tensor shape/dtype, and operator contracts. - No model, checkpoint, or architecture identity participates in the fix. -- Focused validation covers MRV2 zeroer initialization and zero-before-return - ordering, dense and multi-pool interleaved zeroing, Mamba align allocation, - and the existing GDN state-window contracts. No full-model end-to-end run - was required for this source-level repair. - -## 2026-08-24 Qwen3.8-27B NVFP4 no-MTP concurrency scaling - -- The heading names the matched performance workload, not an activation - identity. Runtime admission does not inspect a model name, checkpoint, or - architecture: it uses only SM70 capability, FP8 format, batch and tensor - shapes, paged-KV layout, graph context, and sampler operator settings. -- Contract: four V100s with TP4, compressed-tensors NVFP4 weights, FP16 - activations, E4M3 FP8 KV cache, Flash-V100 attention, prefix caching, Mamba - align mode, no speculation, and CUDA Graph decode sizes 1/2/4/8/16. The - throughput workload is official SPEED-Bench `low_entropy` at 1K and 16K, - 512 output tokens, one warmup plus three measured repeats, temperature 1.0, - top-k 20, top-p 0.95, fixed seed 20260822, and natural EOS. Reported decode - throughput uses only the interval in which every request is actively - decoding, so TTFT and prefill are excluded. -- The admitted E4M3 G6/D256 XQA route covers exact B2-B16 only. It uses paired - 128-bit KV loads, packed exact E4M3-to-half2 conversion, the existing - partition workspace, and the established long-context graph variants. At - 16K, B8 captures `dual_cta_split` and B16 captures `dual_cta`; the baseline - CUDA Graph remains selected below 16K. `VLLM_FLASH_V100_E4M3_BATCH_XQA=0` - restores scalar paged attention. The separate - `VLLM_FLASH_V100_E4M3_BATCH_XQA_OPTIMIZED=0` switch retains XQA while - restoring scalar KV loads, the original E4M3 conversion, and baseline CTA - routing. B1 keeps the original converter regardless of this batch-only - optimization. -- The exact 248,320-column, no-spec B8/B16 sampler contract now selects eight - Triton warps for combined top-k/top-p masking. The operation and reduction - math are unchanged; `VLLM_SM70_TOPK_TOPP_B8_B16_8_WARPS=0` restores Triton's - launch heuristic. Existing MTP verifier-row tuning remains independently - opt-in. -- Development A/B on base `a8ce98ba34` measured the following median pure - decode rates. The candidate includes batched XQA, eight-warp sampling, and - the long-context graph variants. - - | Context | Batch | Control tok/s | Candidate tok/s | Speedup | Candidate scaling efficiency | FP8 reference efficiency | - |---:|---:|---:|---:|---:|---:|---:| - | 1K | 1 | 72.82 | 72.94 | 1.002x | 100.0% | 100.0% | - | 1K | 8 | 402.51 | 465.92 | 1.158x | 79.8% | 87.4% | - | 1K | 16 | 642.37 | 774.18 | 1.205x | 66.3% | 72.3% | - | 16K | 1 | 63.03 | 63.12 | 1.001x | 100.0% | 100.0% | - | 16K | 8 | 186.43 | 356.50 | 1.912x | 70.6% | 67.7% | - | 16K | 16 | 222.03 | 526.59 | 2.372x | 52.1% | 49.1% | - -- Thus the 16K NVFP4 scaling efficiency exceeds the same-criterion FP8 - reference by 4.2% at B8 and 6.2% at B16. The remaining short-context gap is - 8.6% at B8 and 8.2% at B16; do not repeat the attention experiment to close - it. At 1K the next target is non-attention concurrent work, primarily GDN - state traffic and dense projection scheduling. -- A clean production revalidation on base `febf1fc962` reproduced the result - with one warmup and three measured repeats per cell. Only the steady window - in which every request was decoding is reported; endpoint throughput, TTFT, - and prefill are excluded from the table. - - | Context | Batch | Control decode tok/s | Candidate decode tok/s | Speedup | Candidate scaling efficiency | FP8 reference efficiency | - |---:|---:|---:|---:|---:|---:|---:| - | 1K | 1 | 72.65 | 72.62 | 1.000x | 100.0% | 100.0% | - | 1K | 8 | 402.98 | 465.24 | 1.154x | 80.1% | 87.4% | - | 1K | 16 | 641.95 | 772.44 | 1.203x | 66.5% | 72.3% | - | 16K | 1 | 63.10 | 63.00 | 0.998x | 100.0% | 100.0% | - | 16K | 8 | 186.97 | 356.52 | 1.907x | 70.7% | 67.7% | - | 16K | 16 | 222.61 | 527.00 | 2.367x | 52.3% | 49.1% | - - The B8 geometric-mean efficiency gap to the FP8 reference is 2.2%, and the - B16 gap is 1.0%. Both therefore reach the revised near-FP8 scaling target; - at 16K they exceed FP8 by 4.4% and 6.5%, respectively. Every timed request - succeeded and produced its full 512-token cap without empty or replacement- - character output. The original implementation was rebased without conflict - onto `onecat/main@6fc9fe9492`; the Flash-V100 and sampler source blobs were - byte-identical to the measured tree. The final audit rebased the same two - commits onto `onecat/main@f6a5b57b64` after generic prefix-anchored SWA was - merged. `git range-diff` reports both patches as exactly equivalent. The - combined Flash-V100 extension builds for SM70, 24 focused E4M3/graph/sampler - policy tests pass, and all 32 prefix-anchored SWA CPU regressions remain - green. The audit also prevents E4M3-only Flash-V100 graph variants from - being captured for an explicitly selected FlashInfer backend. A - same-toolchain `cuobjdump` comparison against the prefix-anchored main build - finds all 115 common XQA kernel instances SASS-identical, with zero missing - or changed instances; the candidate adds only six batch-optimized instances. -- GPU operator validation covers B2/B16 at partition sizes 64/128/256 plus - B8/B16 long-context routing. All eight comparisons stay within one FP16 ULP - of scalar paged attention with mean absolute error at most `1e-5`. Separate - B8/B16 full-vocabulary tests show bitwise-identical top-k/top-p masks between - the eight-warp schedule and its rollback path. -- Matched GSM8K main/test quality uses 32 prompts per B1/B8/B16 with the same - official stochastic sampling and 512-token natural-EOS cap. Changed batches - move numeric exact match from 22/32 to 23/32 at B8 and remain 23/32 at B16; - all requests complete with no empty or replacement-character output. The - unchanged B1 route moves 25/32 to 24/32 while only 21/32 texts reproduce - across independent processes, so that one-sample movement is recorded as - stochastic run variance rather than attributed to the candidate. -- The first metadata-only long-context probe was invalid: the CUDA Graph - dispatcher still admitted specialized variants for E5M2 only. It produced - no E4M3 route trace and must not be cited. The accepted implementation gates - both metadata and dispatcher creation; capture logs prove eight FULL graphs - and explicit B8/B16 long-context route hits. - -## 2026-08-25 DFlash2 practical-performance closure - -- PRs #252, #253, #257, #266, and #284 are already merged. The missing - performance closure was twelve post-#284 files that remained only in the - measured production worktree, not an open historical PR. -- A matched single-request coding probe with FP8 E5M2 KV, 256K maximum - context, 4096 maximum batched tokens, prefix caching, Mamba align, CUDA - Graph, and tool/reasoning parsers measured `18.66--18.83 ms` per complete - round (`131.2--139.9 tok/s`) on that production tree. The clean - `onecat/main@05d5aa4e57` route hit the same visible QPN8, sharded projection, - grouped verifier, and CUDA Graph logs but measured `27.84--27.96 ms` - (`90.5--92.6 tok/s`). Route-hit logs alone are therefore not an acceptable - performance proof. -- Draft PR #288 restores the exact missing source closure on latest main. The - decode hot path enables fused GDN group metadata for `mamba_cache_mode=align` - and makes disabled GDN dump hooks return before CUDA runtime queries. The - prefill half skips unused DFlash queries on intermediate 4096-token chunks - and keeps the small-M TurboMind projection out of large-M prefill shapes. -- Focused pre-commit passes. On the latest-main worktree, DFlash2/GDN CPU - suites report `110 passed, 32 skipped`; the V100 align-metadata replay and - AOT fullgraph predicate tests both pass. The first end-to-end launch produced - no timing result because an external TP4 job acquired GPUs 4--7 between API - startup and worker memory admission. Wait for the cards and rerun the exact - practical contract before merging; do not lower memory utilization or cite - that failed launch as implementation evidence. -- The queued clean rerun on PR #288 used the exact practical contract: TP4 on - V100, FP8 E5M2 target KV, FP16 draft KV, 256K maximum context, 4096 maximum - batched tokens, prefix caching, Mamba align, probabilistic DFlash2, CUDA - Graph, and the Qwen tool/reasoning parsers. Startup logs prove QPN8 exact - rerank, the TP4 sharded context projection, the one-pass grouped verifier, - target and draft graph capture, and, critically, fused GDN metadata for all - 10 cache groups. -- After one explicitly excluded first-request JIT warmup, four 512-token - single-request runs of `给我写个自我介绍网页吧` measured `18.465--18.537 ms` - per complete round and `133.05--146.63 tok/s`; acceptance length was - `2.457--2.718`. Three 1,024-token practical coding runs measured - `18.587--18.603 ms`, `132.62--139.93 tok/s`, and acceptance length - `2.470--2.606`, reproducing the measured production worktree. -- The historical high-acceptance MBPP item 28 reproduced `251.60 tok/s`, - `18.567 ms` per complete round, acceptance length `4.686`, and a natural-EOS - 328-token completion. Its complete generated text has the same SHA256 as the - previous production result, and EvalPlus reports base `1/1` and plus `1/1`. - This closes both the roughly 18 ms and the 220+ token/s restoration gates. -- A fresh 32K chunked-prefill request logged the context-only DFlash path and - measured `10.58--10.65 s` cold; the identical prefix hit measured `1.405 s`. - The older public server's nominal cold result may have retained an identical - prefix, so only the fresh-server numbers are treated as strict cold evidence. - -## 2026-08-25 Qwen3.8-27B NVFP4 DFlash2 17 ms campaign - -- This section is an active planning record, not evidence that a 17 ms result - or source optimization already exists. The document itself is safe to merge - while implementation and measurements remain pending. -- The next isolated campaign starts from - `onecat/main@34403018d917054dd7765d5e820ad29c8d342348` and targets a stable - batch-one complete speculative round of at most `17.0 ms`. The frozen target - is the local mixed NVFP4/channel-FP8 Qwen3.8-27B checkpoint with BF16 LM - head, TP4 on V100 GPUs 0--3, E5M2 target KV, FP16 draft KV, seven draft - tokens, probabilistic sampling, and FULL target/draft CUDA Graphs. -- The accepted PR #288 practical baseline is `18.465--18.603 ms`, so the - remaining measured gap is approximately `1.5--1.6 ms`. A current-source - no-diagnostic graph-node trace is required before choosing the next source - change. Historical DDTree index reuse, skipped verification, generic - auxiliary-stream overlap, variable K, and approximate reranking remain - rejected. -- Development may use a 32K/512 minimal runtime only to localize kernels. The - promotion gate is the 256K/4096 configuration with prefix caching, Mamba - alignment, and tool/reasoning parsers enabled. Acceptance may not fall by - more than `0.05`; the existing scored datasets and PPL gates must not - regress. A separate 32K/128K/256K sweep protects long-context decay. -- A retained speedup becomes default-on after the engine-contract admission, - rollback, numerical, official-sampling quality, and matched endpoint gates - pass. A valid reduction-order change need not be bitwise or greedy-token - identical, but it must remain finite, within a dtype-appropriate error - envelope, and non-regressing on acceptance and scored quality. -- The full frozen contract, sequence, and pending Draft PR test record are in - `docs/design/sm70_dflash2_nvfp4_17ms.md`. -- The acceptance audit found a runtime/source mismatch rather than an NVFP4 - regression. Current FP8 and NVFP4 request-mean acceptance values are - `4.49887` and `4.53642` (request-mean completion tokens per round are - `4.49761` and `4.53214`), but both runs loaded stale stable-ABI RMSNorm SHA - `fe986a...9ba8c`. Its - batched-weight suite fails 10/13 cases and reproduces the old row-zero - DFlash K-norm defect; repaired SHA `7689e5...d449` passes 13/13. Target - QPN2/QPN8 and sparse-rejection ablations do not recover acceptance. -- DFlash now performs a one-time nonzero bitwise capability check before using - grouped per-layer K normalization. A stale binary is made correctness-safe - through an explicit per-layer fallback and warning. The two new behavior - tests and the full CPU DFlash2 suite pass (`70 passed, 12 skipped`). The - draft-MLP QPN8 experiment is removed after exceeding the `-0.05` paired - acceptance gate. -- The repaired-binary FP8 formal rerun closes at request/pooled acceptance - `5.37699/4.62217`, 61/64 GSM8K, and stable SHA `7689e5...d449`. Relative to - the repaired August 22 formal baseline, paired request delta is `-0.01180` - with bootstrap 95% interval `[-0.08421,0.06282]`; every pooled position is - higher. Relative to the stale current-source run, repair adds `0.87812` - request-mean tokens with 95% interval `[0.75099,1.00344]`. The concurrent - throughput is excluded from standalone speed evidence. -- The matching repaired NVFP4 formal run closes at request/pooled acceptance - `5.39851/4.57051` and 59/64 GSM8K. Its request mean is `+0.02152` versus - current repaired FP8 (paired 95% interval `[-0.10021,0.14795]`), while its - longer 35,366-token trajectory explains the lower pooled value. Repair adds - `0.86209` over stale NVFP4 with interval `[0.74547,0.98007]` and improves 62 - of 64 rows. Its GSM8K pass/fail set is exactly the same as the repaired - August 22 FP8 reference (59 pass, 5 fail). -- The repaired 16K executable MBPP-32 pair passes the quality gate. Dataset - tests are no-DFlash `31/32` versus DFlash2 `32/32`; EvalPlus base is - `30/31` versus `31/31`, and plus is `27/31` versus `27/31`. The plus - discordance is one win per route with exact McNemar `p=1.0`; both runs stop - naturally on 31 rows and hit the 16K limit on one. A scorer extraction bug - that preferred reasoning-era example fences over valid unfenced final code - caused the initial DFlash syntax false negative. Final-answer-first - extraction is regression-checked, and both sides use the corrected scorer. -- The selector-QPN8 candidate passes acceptance but not yet performance - promotion. Dense-to-QPN8 request/pooled acceptance is - `5.39851/4.57051 -> 5.43951/4.58679`; the paired request delta is - `+0.04101` with 95% interval `[-0.02309,0.10669]`, and GSM8K is - `59/64 -> 60/64`. A cross-physical-group run shows directional complete - round `19.083 -> 18.077 ms`, aggregate `218.58 -> 229.41 token/s`, and - steady decode `282.68 -> 300.12 token/s`, but it is not accepted as a speed - pair. The prepared rerank layout reduces available KV capacity from 529,616 - to 459,817 tokens. An unrelated eight-GPU job blocked the same-group repeat, - so the route remains opt-in and default-off. - -## 2026-08-25 Qwen3.8-27B NVFP4 long-context decode - -- The accepted candidate specializes only the SM70 E4M3 G6/D256 dual-CTA - decode route whose production K/V views come from an interleaved - `[blocks, 2, 800, 1, 256]` allocation. It replaces general page arithmetic - with a compile-time page800/physical-stride path and uses aligned shared-V - `half2` loads without changing per-output FMA order. Other layouts, dtypes, - dimensions, routes, and prefill retain the fallback. -- `VLLM_FLASH_V100_E4M3_PAGE800_FASTPATH=0` is the rollback. The default-on - path also has a once-per-process trace flag and is protected by exact shape - and stride checks. Production-layout B8/B16 p64/p128/p256 tests are bitwise - equal to rollback; the two focused GPU suites pass 15/15 and 14/14. -- The matched TP4, no-MTP, E4M3-KV, FULL-graph SPEED-Bench `low_entropy` - pure-decode A/B improves 32K B8/B16 by 11.79%/15.04% and 64K B8/B16 by - 16.47%/19.74%. The four-cell geometric mean is 15.72%, no cell regresses, - and 64K B16 scaling efficiency is 48.52%, only 3.76 percentage points below - the accepted 16K reference. All 144 candidate requests finish with healthy - 512-token output. -- NCU on 64K/B16 p256 attributes the 22.47% kernel reduction to 11.92% fewer - instructions, 48.16% fewer shared-bank conflicts, and 95.92% less MIO- - throttle stall. Peak-relative DRAM/SM utilization rises 24.49% to 28.11% - and 39.11% to 44.79%; occupancy is unchanged. This is copy/addressing and - shared-memory relief, not an occupancy-only claim. -- The official-sampling quality gate uses the same 32 GSM8K main/test prompts - at B8/B16 with temperature 1.0, top-k 20, top-p 0.95, natural EOS, and three - predeclared sampling seeds. Candidate/control aggregate numeric EM is 74/96 - versus 71/96 at B8 and 74/96 versus 74/96 at B16; all 384 requests succeed - with no candidate empty or replacement-character output. The initial single - seed's one-item deficits remain recorded rather than being silently rerun. -- The measured patch rebases without conflict onto `onecat/main@d62ef5cb20`. - The final-source rebuild and measured binary have byte-identical complete - SM70 SASS dumps. - The server's common D256 prefill fallback warning is outside the pure-decode - metric, so this change closes long decode only and must not be cited as a - prefill improvement. - -## 2026-08-25 generic PP profiling and compact MXFP4 audit - -- PR #283 mixed accepted engine work with rejected workload experiments. Its - QPN8 candidates had repeated baseline-correct-to-candidate-wrong quality - regressions, while the metadata-free PP transfer recovered only about - `0.034 ms/token` in the full endpoint. Neither path is promoted to main. -- The retained PP repair is model-independent. A non-last pipeline stage can - return `IntermediateTensors` during a dummy profile/capture; the runner now - selects the semantic `hidden_states` entry instead of treating the wrapper - as a tensor. A CPU unit test also proves that dictionary insertion order - cannot select a residual tensor accidentally. -- The retained MXFP4 route is admitted only by the existing SM70 operator, - group size 32, six one-row groups, and exact `(K,N)` tensor shapes. Runtime - code does not inspect checkpoint paths, model names, `model_type`, or - architecture identity, and the compact grouped route remains default-off. -- On a V100, both new shapes are bitwise equal to the per-expert fallback and - pass CUDA Graph replay. Median operator time changes from `0.139827` to - `0.037478 ms` for `(4096,1024)` and from `0.064696` to `0.021371 ms` for - `(512,4096)`, or `3.731x` and `3.027x` respectively. Final-source reruns on - latest main measured `0.151521` to `0.040428 ms` (`3.748x`) and `0.065004` - to `0.022538 ms` (`2.884x`). The full CUDA 12.8 - extension builds with target architecture 7.0 only; 28 CPU checks and the - changed-file static gate pass. - -## 2026-08-25 Qwen3.8-27B-NVFP4 no-MTP long-context decode - -- Frozen acceptance is TP4 on four V100-SXM2-32GB GPUs, E4M3 KV, - Flash-V100, TurboMind FP8/NVFP4 projections, full CUDA Graph, official - temperature-1/top-p-.95/top-k-20 sampling, no MTP, and exact final contexts. - Pure decode excludes prefill/TTFT and uses 63 intervals from 64 output - tokens. Required rates are 60 tok/s at 128K and 45 tok/s at 256K. -- The admitted batch-one G6/D256, page-1568 E4M3 route uses an exact shared - conversion LUT and device-selected p64/p256/p512/p896/p1664 waves with - thresholds 12,288/49,152/98,304/196,608. A merged long partition/reducer pair - replaces three inactive long launch pairs during graph replay. Admission is - based only on hardware, dtype, batch, head, page-layout, graph, and explicit - partition contracts; it does not inspect a model or checkpoint identity. -- Batch-one full CUDA Graph replay is specialized at 49,152 tokens. The - default graph contains only the p64/p256 launch pair; the semantic - long-context graph admits the merged wave pair. Selection uses scheduler - context metadata and does not synchronize a device sequence length to the - host. Eager execution continues to route from the live context. -- The matched control is 40.561 tok/s (24.654 ms) at 128K and 27.456 tok/s - (36.422 ms) at 256K. The candidate reaches **61.834 tok/s (16.172 ms)** and - **50.376 tok/s (19.851 ms)** respectively, exceeding both targets. Prefill - remains separately reported at 283.127/1085.519 s and TTFT at - 283.154/1085.568 s. -- Both candidate requests return all 64 tokens without corruption and exactly - preserve their control streams: hashes are `66a130ea9d68...` at 128K and - `e1bc9ab8a979...` at 256K. No sampling path or MTP state changes. -- The disabled-compile-cache post-integration 1K/256-output regression guard - reaches 83.184/83.187 tok/s on its two hot repetitions versus - 83.761/83.721 tok/s for the accepted GDN BA-split control, a 0.66% median - delta. All candidate repetitions preserve that control's exact token stream - (`c243818da5cc...`). The first cold request remains separately reported, and - cache-reload diagnostics are excluded from quality evidence. -- A final isolated 256K Nsight Systems trace measures 0.576603 ms for the - merged long partition kernel and 0.033280 ms for its reducer. Including the - inactive device-routed short guard pair gives 0.617979 ms of Flash kernels, - down 63.83% from the old p256 trace's 1.708550 ms. This is operator evidence, - not a full-model TPOT decomposition. -- Exact decode windows show 100% NVML GPU-busy duty, about 48.2%/45.5% - memory-active duty at 128K/256K, 29.45-29.75 GiB allocation per rank, and - 200.7-226.8 W mean per-rank power at 1530/877 MHz SM/memory clocks. Memory - duty is not HBM throughput. NCU counters are blocked by - `ERR_NVGPUCTRPERM` and are not inferred. -- A post-acceptance operator pass retains the same route and arithmetic. It - stages partition page IDs once, specializes only the actual page-1568 - interleaved unbound-view strides, and consumes shared V in adjacent `half2` - pairs while preserving per-dimension FP32 FMA order. Four-arm same-GPU A/B - gives staged exact-context kernel reductions of 1.20%, 1.49%, and 3.99% at - 128K, and 1.23%, 1.73%, and 4.49% at 256K. All 12 sequence/layout cases are - bit-exact with one hash per case across all arms. The final kernel retains - 168 registers and two CTAs/SM, removes its 32-byte stack frame, and reduces - static SASS from 2,424 to 2,248 instructions versus the fixed-stride stage. - The formatted final build is eager- and CUDA-Graph-exact at 128K/256K, and - all eight page-1568 wave-boundary GPU regressions pass. -- Scaling the measured attention slice by those independently matched ratios, - while freezing the measured non-attention residual, projects 63.354 tok/s - at 128K and 52.216 tok/s at 256K. This is an operator-decomposition - projection, not a new endpoint measurement; 61.834/50.376 tok/s remain the - accepted measured rates. -- Reject the contiguous-KV assumption, QK software pipeline, split reducer, - and paired E4M3 converter. The real unbound K/V views are interleaved; - violating that layout corrupts model output. The QK pipeline slows the long - kernels by about 30-38%, and the paired converter loses to the exact shared - LUT. Detailed evidence and rollback flags are in - `docs/design/sm70_qwen38_nvfp4_decode.md`. -- The merge audit rebases this route onto `main@1ce39810b` and retains the - existing page-800 E4M3 specialization beside page 1568. The combined CUDA - 12.8 build contains only SM70 cubins; three host CUDA Graph routing checks, - six page-800 V100 fallback comparisons, and eight page-1568 wave-boundary - comparisons pass. Both GPU suites require bitwise equality. Runtime source - admission remains limited to device capability, dtype, batch/head geometry, - cache strides, page size, graph semantics, and partition thresholds; no - model, checkpoint, `model_type`, or architecture identity is consulted. - -## 2026-08-26 generic in-flight KV eviction audit - -- PR #293 mixed a valid engine-level cache lifetime repair with - architecture-keyed serving recipes, model-specific generation defaults, and - OCR prompt/resolution behavior. The architecture/model identity defaults are - not admitted to main; they are unrelated to inference-engine route safety. -- The retained repair tracks tokens scheduled in batches whose outputs have - not settled yet. Sliding-window, chunked-local, prefix-anchored SWA, and - Mamba eviction now derive their safe freeing boundary from processed tokens, - rather than the scheduler's optimistic computed-token boundary. This avoids - recycling a KV block while an asynchronous or pipeline-parallel batch can - still read it, including when speculative tokens later roll back. -- Recycling-aware admission reserves for the same executor concurrency bound: - pipeline depth for PP, asynchronous queue depth otherwise. The contract uses - scheduler concurrency, attention/cache geometry, token state, and block - lifetime only; it does not inspect a model, checkpoint, `model_type`, or - architecture identity. -- Final-source validation passes 17 focused cache/accounting cases, all 11 - async-scheduler unit cases, and a simple CPU-offload store/load round trip. - All changed-file pre-commit hooks pass. This is a cache lifetime correctness - repair; no throughput claim is attached. - -## 2026-08-26 DeepSeek V4 PP2 TP4 no-DSpark decode recovery - -- The frozen endpoint contract is DeepSeek V4 Flash on eight V100-SXM2-32GB - GPUs, PP2 with TP4 ranks confined to GPUs 0--3 and 4--7, input 1024, output - cap 256, model-default sampling, no speculation/DSpark, E5M2 KV, sparse MLA, mHC, - TurboMind FP8/MXFP4 projections, and full CUDA Graph replay. Pure decode is - reported separately from prefill and TTFT; the target remains 100 tok/s. -- The latest uninstrumented control reaches 59.160 tok/s, or 16.903 ms/token. - A 26-step Nsight Systems trace is perturbed to 18.939 ms/token but localizes - the per-stage service leaders to 6.434 ms of FP8 dense work, 2.647 ms of - MXFP4 MoE, 2.582 ms of FP16 GEMV, 1.714 ms of mHC, 1.695 ms of sparse MLA, - 1.694 ms of routing, 1.688 ms of Q/KV work, and 1.476 ms of TP all-reduce. - Roughly 9 ms pipeline rows are dependency waits and are not added as PP - transfer cost. -- PR #299 extends the fully connected SM70 TP4 - push all-reduce to the exact 8-KiB FP16 CUDA Graph payload. Same-binary TP4 - A/B/B/A operator timing moves 9.286--9.425 us to 3.041--3.174 us. Both paths - are bitwise equal to the explicit rank-0-through-rank-3 FP32 accumulation - oracle. A stronger changing-input graph test covers 43 collective nodes, - eight input patterns, 64 replays, and all four ranks with zero mismatches; - every GPU claim is released after the bounded test. The route is default-on - only for SM70, fully connected TP4, active CUDA Graph capture, FP16, and the - exact 8-KiB or 80-KiB payload; `VLLM_SM70_TP4_PUSH_ALLREDUCE=0` rolls back. -- Matched full-model control and candidate medians are 59.160 and 61.272 - tok/s, or 16.903 and 16.321 ms/token. The candidate therefore recovers - 3.57% and 0.583 ms/token. This is accepted performance and operator evidence, - while two independent unchanged controls select different stable token - streams because of pre-existing cross-process TurboMind autotuning. Greedy - identity is diagnostic rather than an acceptance gate; the exact collective - oracle isolates the changed operator. -- The next arithmetic target is the FP8 dense/WO-A launch family, which is the - largest trace category. Each of the 43 layers currently dispatches the two - one-row `(K=4096,N=1024)` groups separately. The replacement prepares - two persistent TurboMind pointer rows and dispatches both groups through one - blocked GEMM only for batch-one decode; all other shapes and token counts - retain the existing loop. -- The accepted arithmetic contract fixes the grouped descriptor to the dense - `8x128x64`, split-K 7 schedule. An environment-injected selector screen and - a source-built rerun both pass eager and changing-input CUDA Graph checks for - exact-small, random-small, and model-like inputs with zero mismatches. The - measured source-built extension SHA256 is - `133724eb67e70e3b4ff098c5f32c00e3155cd7c5e8f7aeb6dff84934ac294e9e`. - The merge audit changes the absent-environment default to the measured - explicit `=1` behavior; it does not change the selected tactic or kernel. - Median isolated graph time changes from `74.943` to `22.450 us` per layer, - projecting `2.257 ms/token` of service reduction before model overlap. - Independently tuned flat/grouped schedules differ by up to one FP16 ULP and - remain rejected. -- The matched full-model pair keeps the source, extension, push all-reduce, - topology, lengths, and cache contract fixed. The control median is - `57.436 tok/s` (`17.411 ms/token`) and the grouped candidate reaches - `59.209 tok/s` (`16.889 ms/token`), recovering `0.521 ms/token` or 3.09%. - This is the measured endpoint reduction; the larger isolated service - projection is not used as a TPOT claim. Both sides are internally stable - across three requests. Their different hashes reflect unrelated - cross-process FP8 tactic selection and are not treated as a numerical gate. -- The grouped route is therefore default-on only for the exact engine - contract already proven above: SM70 block-FP8 `[128,128]`, two groups, - `K=4096`, `N=1024`, batch-one decode, the source-built grouped and pointer - operators, and the fixed `8x128x64` split-K-7 tactic. Other shapes, multi-row - calls, missing operators, and non-SM70 devices retain the two-launch path. - `VLLM_SM70_FP8_GROUPED_BMM_DECODE=0` is the rollback. Admission never reads - model name, checkpoint, `model_type`, or architecture identity. -- A multi-row FP16 GEMV screen is rejected. The only bitwise-exact candidate - (`BLOCK_N=2`, `BLOCK_K=1024`) is 0.5--12% slower across the five traced - `N=64/256/512/1024/2048` shapes; smaller-K/four-row variants are slower and - also change FP32 accumulation. The candidate is not retained. -- The next MXFP4 B1 screen removes redundant route sorting without changing - model arithmetic. The six unique top-k IDs become the compact TurboMind - group IDs directly, fixed offsets remain `[0, 1, ..., 6]`, W13 reads the - single physical input row through the accepted broadcast contract, and W2 - stays in original top-k order. The final weighted reduction therefore keeps - the same route-order FP32 FMA sequence without an inverse permutation. - Initial and changed-route CUDA Graph replays are bitwise at W13, clamped - SwiGLU, W2, and final output. Adding the exact top-k-6 `half2` reducer moves - 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/`. - 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 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 - -- Development is isolated on - `agent/v100-dflash2-ngram-hybrid-20260825-050018` and rebased onto the - restored `onecat/main@d62ef5cb20`; Draft PR #287 contains only the MRV2 - DFlash2 dependency closure. -- The opt-in assistant reads the authoritative UVA token state, uses the same - reverse-KMP policy as standalone vLLM n-gram lookup, skips DFlash2 query and - selector only on a full seven-token hit, and preserves context-KV - materialization. Misses remain on the unchanged DFlash2 route. -- V100 unit coverage is 24/24, including probabilistic one-hot rejection, - mixed rows, overlap state, and the CUDA override kernel. Practical TP4 graph - startup with prefix cache, FP8 E5M2 KV, 256K max context, 4096 batched-token - limit, and tool/reasoning parsers is proven. The earlier 27--31 ms clean-main - measurements predate the merged production closure and are invalid as a - speed baseline; collect `ngram_assist=false/true` numbers from the same - restored build. -- The restored paired result is now complete. No-assist practical coding is - `18.660 ms` per round, `135.70 tok/s`, acceptance `2.532`; ngram `[5,5]` is - `18.980 ms`, `131.14 tok/s`, acceptance `2.494`, with only about 2.2% full - hits. The feature therefore remains opt-in for short/low-hit traffic. -- On the fixed-seed, 16K-cap, natural-stop MBPP32 workload, pure DFlash records - `389.542 s`, `19,767` rounds, `19.707 ms/round`, acceptance `3.336`, and - EvalPlus Base/Plus `30/31` and `28/31`. Hybrid records `327.613 s`, `16,700` - rounds, `19.618 ms/round`, acceptance `3.216`, and Base/Plus `31/31` and - `28/31`. Its cumulative full-hit rate is about 10.1%, but lower sampled - acceptance leaves aggregate output throughput `163.94` versus - `169.26 tok/s`; the quality gate passes and the throughput gate does not. -- Hybrid 32K prefill is `10.688 s` cold and `1.416 s` on an identical-prefix - hit, matching restored no-assist evidence (`10.58--10.65 s`, `1.405 s`) - closely enough to rule out a material prefill regression in this probe. -- Post-run capture failures on GPUs 4--7 reproduce unchanged on the clean - pre-ngram main worktree at the same draft paged-attention capture. They are - runtime/GPU-state evidence, not an ngram regression or speed sample. A full - NVLink-group reset was intentionally deferred because GPUs 0--3 host the - live user API. -- The merge audit removes the mixed-hit quality/performance hazard. N-gram - drafts are now applied only when every active request has a full-width hit - and the complete DFlash2 query/selector is skipped. If any request misses, - the batch keeps the unchanged DFlash2 proposals because overriding a row - after paying the full query cost cannot improve latency and can alter - 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 - -- 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. -- 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. -- 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. -- 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. - Three matched endpoint runs improve QPN8 by about `0.22%` (`0.034 ms/token`) - without changing its output behavior. Non-admitted configurations use the - 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 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. -- 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. That single sample remains visible, but - it is not a greedy-identity gate and does not show that choosing the lowest - operator error improves aggregate task quality. The accepted repair - therefore preserves every numerically bounded speed winner and falls back - only the demonstrated gate/up outlier. 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 - -- PR #297 originally wrapped one named vision module by invoking the layer - offloader a second time. That form is not admitted: it is model-specific, - loses the tower prefix needed by `--cpu-offload-params visual`, and violates - the prefetch backend's single-call contract. -- The accepted path is an engine-level component hook on every module recorded - by the existing multimodal tower lifecycle. UVA qualifies relative parameter - names with the recorded module path before applying exact segment matching; - layer-prefetch and no-op backends inherit a single-component no-op. Parameters - already visited through `make_layers()` are not copied or counted twice. -- Admission uses only the configured offload backend and budget, the explicit - parameter-path selector, and the module's generic tower role. It never reads - a model name, checkpoint path, `model_type`, or architecture identity. With no - CPU-offload budget the global offloader remains a no-op. -- The submitted V100 evidence reports about 0.78 GiB additional device memory - per rank, a 152,056-token KV pool, and successful 151,296-token context on - two 16-GiB V100 ranks. Text-only decode does not execute the tower; image - prefill accepts the expected PCIe/UVA cost. This is capacity evidence, not a - general throughput claim, so the route remains explicitly configured rather - than becoming an unconditional default. - -## 2026-08-26 PR #299 8-KiB TP4 push all-reduce acceptance - -- The existing two-epoch SM70 TP4 push collective now also admits the exact - 8-KiB FP16 decode payload. Four CTAs cover its 512 packed 16-byte elements; - the existing 80-KiB verifier payload retains its 80-CTA launch. Buffer - sizing already uses the larger payload, so this adds no allocation growth. -- Same-binary TP4 A/B/B/A timing improves from 9.286--9.425 us to - 3.041--3.174 us. A dynamic CUDA Graph gate covers 43 collective nodes, - eight changing input patterns, 64 replays, and all four ranks with zero - element mismatches. Both paths are bitwise equal to an explicit fixed-order - rank-0-through-rank-3 FP32 accumulation oracle. -- The matched PP2 x TP4 endpoint improves from 59.160 to 61.272 token/s - (`+3.57%`) and from 16.903 to 16.321 ms/token (`-0.583 ms/token`). Independent - unchanged controls themselves select different stable greedy streams due to - pre-existing cross-process autotuning, so token hashes are diagnostic only; - the exact operator gate isolates this collective without imposing greedy - identity. -- `VLLM_SM70_TP4_PUSH_ALLREDUCE` is therefore default-on. Runtime admission is - limited to SM70, fully connected TP4, active CUDA Graph capture, FP16, and - the exact 8-KiB or 80-KiB payload. Every other device, topology, dtype, size, - 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. -- Production defaults on through `VLLM_SM70_DSV4_FP13_GEMV=1`; explicit `=0` - is the rollback. Admission checks SM70, CUDA device equality, dtype, shape, - layout, output contract, and the actual weight bits. It never reads a model - name or checkpoint identity, and an incompatible tensor retains FP16. Packing - occurs once after loading; capture and per-token execution perform no packing. - The retained buffers add 549.66 MiB without pipeline parallelism, or at most - 282.34 MiB per rank for the audited PP2 partition. diff --git a/docs/features/speculative_decoding/README.md b/docs/features/speculative_decoding/README.md index bda540b88c..768e9f78d4 100644 --- a/docs/features/speculative_decoding/README.md +++ b/docs/features/speculative_decoding/README.md @@ -117,33 +117,6 @@ 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. - -```bash -vllm serve \ - --speculative-config '{ - "method": "dflash", - "model": "your-org/your-dflash2-selector-model", - "num_speculative_tokens": 7, - "draft_sample_method": "probabilistic", - "ngram_assist": true, - "prompt_lookup_min": 5, - "prompt_lookup_max": 5 - }' -``` - -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. - #### Suffix decoding | Key | Type | Default | Meaning | diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 9ffaa6954b..0654a59caf 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -539,14 +539,6 @@ See [this page](generative_models.md) for more information on how to use generat These models primarily accept the [`LLM.generate`](./generative_models.md#llmgenerate) API. Chat/Instruct models additionally support the [`LLM.chat`](./generative_models.md#llmchat) API. -!!! note "DeepSeek-OCR image resolution" - `DeepseekOCRForCausalLM` accepts explicit `mm_processor_kwargs` with - `image_mode` set to `tiny`, `small`, `base`, `large`, or `gundam`. - `min_crops` and `max_crops` may also be set with - `1 <= min_crops <= max_crops <= 9`. The default remains the existing - `gundam` processing contract; no mode is selected from a model or - checkpoint identity at runtime. - | Architecture | Models | Inputs | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) | | ------------ | ------ | ------ | ----------------- | -------------------- | ------------------------- | | `AriaForConditionalGeneration` | Aria | T + I+ | `rhymes-ai/Aria` | | | diff --git a/flash-attention-v100/flash_attn_v100/__init__.py b/flash-attention-v100/flash_attn_v100/__init__.py index b3560a3ca8..5d38754d54 100644 --- a/flash-attention-v100/flash_attn_v100/__init__.py +++ b/flash-attention-v100/flash_attn_v100/__init__.py @@ -10,7 +10,6 @@ flash_attn_decode_paged_xqa_available, flash_attn_decode_qk_scores, flash_attn_func, - flash_attn_grouped_verify_paged, flash_attn_lse, flash_attn_prefill_paged, flash_attn_prefill_paged_bfla, @@ -31,7 +30,6 @@ "flash_attn_decode_paged_xqa", "flash_attn_decode_paged_xqa_available", "flash_attn_decode_paged_wmma", - "flash_attn_grouped_verify_paged", "flash_attn_turboquant_decode_paged", "flash_attn_turboquant_decode_paged_available", "flash_attn_bhmd_func", 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..a28f7571e4 100644 --- a/flash-attention-v100/flash_attn_v100/flash_attn_interface.py +++ b/flash-attention-v100/flash_attn_v100/flash_attn_interface.py @@ -18,13 +18,11 @@ DEFAULT_DECODE_PARTITION_SIZE = 256 VALID_DECODE_PARTITION_SIZES = (256, 512, 1024) -E4M3_XQA_VALID_DECODE_PARTITION_SIZES = (64, 128, 256, 512, 896, 1024, 1664) _decode_plan_cache = {} _decode_workspace_cache = {} _xqa_staged_rescale_workspace_cache = {} _turboquant_decode_workspace_cache = {} _prefill_splitkv3_workspace_cache = {} -_grouped_verify_workspace_cache = {} logger = logging.getLogger(__name__) @@ -54,12 +52,6 @@ class _PrefillSplitkv3Workspace: softmax_lse: torch.Tensor -@dataclass -class _GroupedVerifyWorkspace: - partial_out: torch.Tensor - partial_lse: torch.Tensor - - def maybe_contiguous(x): return x.contiguous() if x is not None and not x.is_contiguous() else x @@ -160,7 +152,6 @@ def _get_decode_plan( workspace_seq_capacity_hint: int | None = None, active_num_partitions: torch.Tensor | None = None, partition_size_hint: int | None = None, - valid_partition_sizes: tuple[int, ...] = VALID_DECODE_PARTITION_SIZES, ) -> _DecodePlan: batch_capacity = batch_size_hint or block_table.shape[0] num_heads = q.shape[1] @@ -193,7 +184,6 @@ def _get_decode_plan( _validate_decode_partition_size( int(partition_size_hint), "partition_size_hint", - valid_partition_sizes, ) if partition_size_hint is not None else _get_decode_partition_size( @@ -503,35 +493,6 @@ def _get_prefill_splitkv3_workspace( return workspace -def _get_grouped_verify_workspace(q: torch.Tensor) -> _GroupedVerifyWorkspace: - device_index = q.device.index if q.device.index is not None else -1 - key = ( - q.device.type, - device_index, - _workspace_stream_id(q.device), - q.dtype, - ) - workspace = ( - _grouped_verify_workspace_cache.get(key) if _can_cache_workspace(q) else None - ) - if workspace is None: - workspace = _GroupedVerifyWorkspace( - partial_out=torch.empty( - (80, 8, 6, 256), - dtype=torch.float16, - device=q.device, - ), - partial_lse=torch.empty( - (80, 8, 6), - dtype=torch.float32, - device=q.device, - ), - ) - if _can_cache_workspace(q): - _grouped_verify_workspace_cache[key] = workspace - return workspace - - def _get_decode_partition_size( max_seq_capacity: int, head_dim: int, @@ -568,13 +529,11 @@ def _select_default_decode_partition_size( return DEFAULT_DECODE_PARTITION_SIZE -def _validate_decode_partition_size( - value: int, - name: str, - valid_partition_sizes: tuple[int, ...] = VALID_DECODE_PARTITION_SIZES, -) -> int: - if value not in valid_partition_sizes: - raise ValueError(f"{name} must be one of {valid_partition_sizes}, got {value}") +def _validate_decode_partition_size(value: int, name: str) -> int: + if value not in VALID_DECODE_PARTITION_SIZES: + raise ValueError( + f"{name} must be one of {VALID_DECODE_PARTITION_SIZES}, got {value}" + ) return value @@ -938,8 +897,6 @@ def flash_attn_decode_paged( workspace_seq_capacity_hint: int | None = None, active_num_partitions: torch.Tensor | None = None, partition_size_hint: int | None = None, - anchor_lens: torch.Tensor | None = None, - anchored_window: int = 0, ): if softmax_scale is None: softmax_scale = q.shape[-1] ** -0.5 @@ -948,16 +905,9 @@ def flash_attn_decode_paged( block_table = maybe_contiguous(block_table) seq_lens = maybe_contiguous(seq_lens) out = maybe_contiguous(out) - anchor_lens = maybe_contiguous(anchor_lens) window_size_left, window_size_right = window_size if window_size_left < -1 or window_size_right < -1: raise ValueError(f"Invalid window_size={window_size}; values must be >= -1") - if anchored_window < 0: - raise ValueError("anchored_window must be non-negative") - if (anchored_window > 0) != (anchor_lens is not None): - raise ValueError( - "anchor_lens and a positive anchored_window must be provided together" - ) batch_capacity = q.shape[0] num_heads = q.shape[1] head_dim = q.shape[2] @@ -1010,8 +960,6 @@ def flash_attn_decode_paged( float(v_scale), int(window_size_left), int(window_size_right), - anchor_lens, - int(anchored_window), ) @@ -1023,49 +971,6 @@ def flash_attn_decode_paged_xqa_staged_available() -> bool: return hasattr(flash_attn_v100_cuda, "decode_paged_xqa_staged_fwd") -def flash_attn_grouped_verify_paged( - q: torch.Tensor, - k_cache: torch.Tensor, - v_cache: torch.Tensor, - block_table: torch.Tensor, - seq_lens: torch.Tensor, - softmax_scale: float | None = None, - out: torch.Tensor | None = None, - kv_cache_dtype: str = "fp8_e5m2", - k_scale: float = 1.0, - v_scale: float = 1.0, - one_pass: bool = False, -) -> torch.Tensor: - """Exact grouped q8/H6/D256 DFlash2 verifier prototype 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. - """ - if softmax_scale is None: - softmax_scale = q.shape[-1] ** -0.5 - q = maybe_contiguous(q) - block_table = maybe_contiguous(block_table) - seq_lens = maybe_contiguous(seq_lens) - out = maybe_contiguous(out) - workspace = _get_grouped_verify_workspace(q) - return flash_attn_v100_cuda.grouped_verify_paged_fwd( - q, - k_cache, - v_cache, - out, - block_table, - seq_lens, - workspace.partial_out, - workspace.partial_lse, - float(softmax_scale), - kv_cache_dtype, - float(k_scale), - float(v_scale), - bool(one_pass), - ) - - def flash_attn_decode_paged_xqa( q: torch.Tensor, k_cache: torch.Tensor, @@ -1082,7 +987,6 @@ def flash_attn_decode_paged_xqa( workspace_seq_capacity_hint: int | None = None, active_num_partitions: torch.Tensor | None = None, partition_size_hint: int | None = None, - batch_context_routing: bool = False, ): if not flash_attn_decode_paged_xqa_available(): raise RuntimeError("flash_attn_v100 CUDA extension lacks XQA decode") @@ -1112,22 +1016,6 @@ def flash_attn_decode_paged_xqa( workspace_seq_capacity_hint=workspace_seq_capacity_hint, active_num_partitions=active_num_partitions, partition_size_hint=partition_size_hint, - valid_partition_sizes=( - E4M3_XQA_VALID_DECODE_PARTITION_SIZES - if kv_cache_dtype in ("fp8", "fp8_e4m3") - and q.ndim == 3 - and q.shape[1:] == (6, 256) - and ( - q.shape[0] == 1 - or ( - os.getenv("VLLM_FLASH_V100_E4M3_BATCH_XQA", "1") == "1" - and 1 < q.shape[0] <= 16 - ) - ) - and k_cache.dtype == torch.uint8 - and v_cache.dtype == torch.uint8 - else VALID_DECODE_PARTITION_SIZES - ), ) _assert_decode_launch_covers_seq_lens( plan, @@ -1186,22 +1074,6 @@ def flash_attn_decode_paged_xqa( int(window_size_right), ) - batch_context_max_seq_len = 0 - if batch_context_routing: - live_max_seq_len = int(max_seq_len_hint or 0) - workspace_max_seq_len = int(workspace_seq_capacity_hint or 0) - if _cuda_graph_capture_active(): - batch_context_max_seq_len = max( - live_max_seq_len, - workspace_max_seq_len, - ) - else: - batch_context_max_seq_len = live_max_seq_len - if batch_context_max_seq_len <= 0: - batch_context_max_seq_len = workspace_max_seq_len - if batch_context_max_seq_len <= 0: - batch_context_max_seq_len = plan.launch_num_partitions * plan.partition_size - return flash_attn_v100_cuda.decode_paged_xqa_fwd( q, k_cache, @@ -1221,7 +1093,6 @@ def flash_attn_decode_paged_xqa( float(v_scale), int(window_size_left), int(window_size_right), - int(batch_context_max_seq_len), ) @@ -1361,8 +1232,6 @@ def flash_attn_prefill_paged( v_scale: float = 1.0, causal: bool = True, window_size: tuple = (-1, -1), - anchor_lens: torch.Tensor | None = None, - anchored_window: int = 0, ): if softmax_scale is None: softmax_scale = q.shape[-1] ** -0.5 @@ -1372,16 +1241,9 @@ def flash_attn_prefill_paged( block_table = maybe_contiguous(block_table) seq_lens = maybe_contiguous(seq_lens) out = maybe_contiguous(out) - anchor_lens = maybe_contiguous(anchor_lens) window_size_left, window_size_right = window_size if window_size_left < -1 or window_size_right < -1: raise ValueError(f"Invalid window_size={window_size}; values must be >= -1") - if anchored_window < 0: - raise ValueError("anchored_window must be non-negative") - if (anchored_window > 0) != (anchor_lens is not None): - raise ValueError( - "anchor_lens and a positive anchored_window must be provided together" - ) q_ = q.permute(0, 2, 1, 3).contiguous() out_ = out.permute(0, 2, 1, 3).contiguous() if out is not None else None @@ -1400,8 +1262,6 @@ def flash_attn_prefill_paged( causal, int(window_size_left), int(window_size_right), - anchor_lens, - int(anchored_window), ) return _copy_bhmd_to_bmhd_out(out_, out_original) @@ -1573,8 +1433,6 @@ def flash_attn_prefill_paged_bhmd( causal, int(window_size_left), int(window_size_right), - None, - 0, ) @@ -1663,7 +1521,6 @@ def flash_attn_prefill_paged_d256_bm32_allp_pair_scratch_splitkv3( "flash_attn_decode_paged", "flash_attn_decode_paged_xqa", "flash_attn_decode_paged_xqa_available", - "flash_attn_grouped_verify_paged", "flash_attn_decode_paged_wmma", "flash_attn_decode_qk_scores", "flash_attn_turboquant_decode_paged", diff --git a/flash-attention-v100/include/fused_mha.h b/flash-attention-v100/include/fused_mha.h index 0987368cba..e778ed0d8f 100644 --- a/flash-attention-v100/include/fused_mha.h +++ b/flash-attention-v100/include/fused_mha.h @@ -8,132 +8,243 @@ #include 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_); - -at::Tensor flash_attention_qk_scores(const at::Tensor& q, const at::Tensor& k, - const float softmax_scale, - const bool is_causal); + 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_ +); + +at::Tensor flash_attention_qk_scores( + const at::Tensor& q, + const at::Tensor& k, + const float softmax_scale, + const bool is_causal +); at::Tensor flash_attention_decode_paged( - const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_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& active_num_partitions, - const float softmax_scale, const int partition_size, - const int launch_num_partitions, const std::string& kv_cache_dtype, - const float k_scale, const float v_scale, const int window_size_left, - const int window_size_right, const std::optional& anchor_lens, - const int64_t anchored_window); + const at::Tensor& q, + const at::Tensor& k_cache, + const at::Tensor& v_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& active_num_partitions, + const float softmax_scale, + const int partition_size, + const int launch_num_partitions, + const std::string& kv_cache_dtype, + const float k_scale, + const float v_scale, + const int window_size_left, + const int window_size_right +); at::Tensor flash_attention_decode_paged_xqa( - const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_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& active_num_partitions, - const float softmax_scale, const int partition_size, - const int launch_num_partitions, const std::string& kv_cache_dtype, - const float k_scale, const float v_scale, const int window_size_left, - const int window_size_right, const int batch_context_max_seq_len); + const at::Tensor& q, + const at::Tensor& k_cache, + const at::Tensor& v_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& active_num_partitions, + const float softmax_scale, + const int partition_size, + const int launch_num_partitions, + const std::string& kv_cache_dtype, + const float k_scale, + const float v_scale, + const int window_size_left, + const int window_size_right +); at::Tensor flash_attention_decode_paged_xqa_staged( - const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_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, at::Tensor& online_rescales, - const at::Tensor& active_num_partitions, const float softmax_scale, - const int partition_size, const int launch_num_partitions, - const std::string& kv_cache_dtype, const float k_scale, const float v_scale, - const int window_size_left, const int window_size_right); - -at::Tensor flash_attention_grouped_verify_paged( - const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, - std::optional& out_, const at::Tensor& block_table, - const at::Tensor& seq_lens, at::Tensor& partial_out, - at::Tensor& partial_lse, const float softmax_scale, - const std::string& kv_cache_dtype, const float k_scale, const float v_scale, - const bool one_pass); + const at::Tensor& q, + const at::Tensor& k_cache, + const at::Tensor& v_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, + at::Tensor& online_rescales, + const at::Tensor& active_num_partitions, + const float softmax_scale, + const int partition_size, + const int launch_num_partitions, + const std::string& kv_cache_dtype, + const float k_scale, + const float v_scale, + const int window_size_left, + const int window_size_right +); at::Tensor flash_attention_decode_paged_wmma( - const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, - std::optional& out_, const at::Tensor& block_table, - const at::Tensor& seq_lens, const float softmax_scale, - const std::string& kv_cache_dtype, const float k_scale, - const float v_scale); + const at::Tensor& q, + const at::Tensor& k_cache, + const at::Tensor& v_cache, + std::optional& out_, + const at::Tensor& block_table, + const at::Tensor& seq_lens, + const float softmax_scale, + const std::string& kv_cache_dtype, + const float k_scale, + const float v_scale +); at::Tensor flash_attention_decode_qk_scores( - const at::Tensor& q, const at::Tensor& k_cache, - const at::Tensor& block_table, const at::Tensor& seq_lens, - const float softmax_scale, const int partition_size, - const std::string& kv_cache_dtype, const float k_scale); + const at::Tensor& q, + const at::Tensor& k_cache, + const at::Tensor& block_table, + const at::Tensor& seq_lens, + const float softmax_scale, + const int partition_size, + const std::string& kv_cache_dtype, + const float k_scale +); 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 +); at::Tensor flash_attention_prefill_paged( - const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, - std::optional& out_, const at::Tensor& block_table, - const at::Tensor& seq_lens, const float softmax_scale, - const std::string& kv_cache_dtype, const float k_scale, const float v_scale, - const bool is_causal, const int window_size_left, - const int window_size_right, const std::optional& anchor_lens, - const int64_t anchored_window); + const at::Tensor& q, + const at::Tensor& k_cache, + const at::Tensor& v_cache, + std::optional& out_, + const at::Tensor& block_table, + const at::Tensor& seq_lens, + const float softmax_scale, + const std::string& kv_cache_dtype, + const float k_scale, + const float v_scale, + const bool is_causal, + const int window_size_left, + const int window_size_right +); std::vector flash_attention_prefill_paged_d256_bm32_allp_pair_scratch( - const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, - std::optional& out_, std::optional& softmax_lse_, - const at::Tensor& block_table, const at::Tensor& seq_lens, - const float softmax_scale); + const at::Tensor& q, + const at::Tensor& k_cache, + const at::Tensor& v_cache, + std::optional& out_, + std::optional& softmax_lse_, + const at::Tensor& block_table, + const at::Tensor& seq_lens, + const float softmax_scale +); std::vector flash_attention_prefill_paged_d256_bm32_allp_pair_scratch_splitkv3( - const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, - std::optional& out_, std::optional& softmax_lse_, - at::Tensor& split_tmp_out, at::Tensor& split_tmp_row_max, - at::Tensor& split_tmp_row_sum, const at::Tensor& block_table, - const int64_t actual_n, const float softmax_scale); + const at::Tensor& q, + const at::Tensor& k_cache, + const at::Tensor& v_cache, + std::optional& out_, + std::optional& softmax_lse_, + at::Tensor& split_tmp_out, + at::Tensor& split_tmp_row_max, + at::Tensor& split_tmp_row_sum, + const at::Tensor& block_table, + const int64_t actual_n, + const float softmax_scale +); at::Tensor flash_attention_prefill_paged_bfla( - const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, - std::optional& out_, const at::Tensor& block_table, - const at::Tensor& seq_lens, const at::Tensor& bfla_block_mask, - const int bfla_mask_block_n, const float softmax_scale, - const std::string& kv_cache_dtype, const float k_scale, const float v_scale, - const bool is_causal, const int window_size_left, - const int window_size_right); + const at::Tensor& q, + const at::Tensor& k_cache, + const at::Tensor& v_cache, + std::optional& out_, + const at::Tensor& block_table, + const at::Tensor& seq_lens, + const at::Tensor& bfla_block_mask, + const int bfla_mask_block_n, + const float softmax_scale, + const std::string& kv_cache_dtype, + const float k_scale, + const float v_scale, + const bool is_causal, + const int window_size_left, + const int window_size_right +); at::Tensor flash_attention_prefill_paged_splitkv( - const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, - std::optional& out_, const at::Tensor& block_table, - const at::Tensor& seq_lens, const float softmax_scale, - const std::string& kv_cache_dtype, const float k_scale, const float v_scale, - const bool is_causal, const int window_size_left, - const int window_size_right, const int split_kv_tokens, - const int max_seq_len_hint); + const at::Tensor& q, + const at::Tensor& k_cache, + const at::Tensor& v_cache, + std::optional& out_, + const at::Tensor& block_table, + const at::Tensor& seq_lens, + const float softmax_scale, + const std::string& kv_cache_dtype, + const float k_scale, + const float v_scale, + const bool is_causal, + const int window_size_left, + const int window_size_right, + const int split_kv_tokens, + const int max_seq_len_hint +); void flash_attention_fp8_e5m2_paged_kv_to_fp16( - const at::Tensor& key_cache, const at::Tensor& value_cache, - const at::Tensor& block_table, const at::Tensor& seq_lens, - at::Tensor& key_out, at::Tensor& value_out, const float key_scale, - const float value_scale); + const at::Tensor& key_cache, + const at::Tensor& value_cache, + const at::Tensor& block_table, + const at::Tensor& seq_lens, + at::Tensor& key_out, + at::Tensor& value_out, + const float key_scale, + const float value_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); + 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 +); #endif diff --git a/flash-attention-v100/include/fused_mma.h b/flash-attention-v100/include/fused_mma.h index 407289c504..2d7c89f883 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,409 +15,385 @@ 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"); -} - -__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"); -} - -__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"); -} - -__device__ __forceinline__ void load_matrix_sync( - 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) { + fragment& frag, + const half* smem_ptr, unsigned ldm) { asm volatile( - "wmma.load.c.sync.aligned.row.m16n16k16.f32 " + "wmma.load.a.sync.aligned.row.m16n16k16.f16 " "{%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]) + : "=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"); - } else { + : "memory" + ); +} + +__device__ __forceinline__ void load_matrix_sync( + fragment& frag, + const half* smem_ptr, unsigned ldm) { asm volatile( - "wmma.load.c.sync.aligned.col.m16n16k16.f32 " + "wmma.load.a.sync.aligned.col.m16n16k16.f16 " "{%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]) + : "=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"); - } + : "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.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) { - 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.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.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 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" + ); + } } __device__ __forceinline__ void load_matrix_sync( - 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"); + 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 float* smem_ptr, - unsigned ldm, layout_t layout) { - if (layout == mem_row_major) { + fragment& frag, + const half* smem_ptr, unsigned ldm) { asm volatile( - "wmma.load.c.sync.aligned.row.m32n8k16.f32 " + "wmma.load.a.sync.aligned.col.m32n8k16.f16 " "{%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]) + : "=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"); - } else { + : "memory" + ); +} + +__device__ __forceinline__ void load_matrix_sync( + fragment& frag, + const half* smem_ptr, unsigned ldm) { asm volatile( - "wmma.load.c.sync.aligned.col.m32n8k16.f32 " + "wmma.load.b.sync.aligned.row.m32n8k16.f16 " "{%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]) + : "=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"); - } + : "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.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) { - 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 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" + ); + } } __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.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.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"); + 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 float* smem_ptr, - unsigned ldm, layout_t layout) { - if (layout == mem_row_major) { + fragment& frag, + const half* smem_ptr, unsigned ldm) { asm volatile( - "wmma.load.c.sync.aligned.row.m8n32k16.f32 " + "wmma.load.b.sync.aligned.row.m8n32k16.f16 " "{%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]) + : "=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"); - } else { + : "memory" + ); +} + +__device__ __forceinline__ void load_matrix_sync( + fragment& frag, + const half* smem_ptr, unsigned ldm) { asm volatile( - "wmma.load.c.sync.aligned.col.m8n32k16.f32 " + "wmma.load.b.sync.aligned.col.m8n32k16.f16 " "{%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]) + : "=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"); - } + : "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" + ); + } } __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) @@ -436,6 +412,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 a6561485e7..1c29b0678d 100644 --- a/flash-attention-v100/kernel/contiguous_to_paged.cu +++ b/flash-attention-v100/kernel/contiguous_to_paged.cu @@ -4,72 +4,89 @@ #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 int token_idx = blockIdx.x; - const int64_t slot_idx = slot_mapping[token_idx]; + 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 +) { - if (slot_idx < 0) { - return; - } + const int token_idx = blockIdx.x; + const int64_t slot_idx = slot_mapping[token_idx]; - const int64_t block_idx = slot_idx / block_size; - const int64_t block_offset = slot_idx % block_size; + if (slot_idx < 0) { + return; + } + + 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..f336475a35 100644 --- a/flash-attention-v100/kernel/flash_decode_paged.cu +++ b/flash-attention-v100/kernel/flash_decode_paged.cu @@ -4,8 +4,6 @@ #include #include #include -#include -#include #include #include #include @@ -20,8 +18,7 @@ namespace { int kv_cache_dtype_code_from_string(const std::string& kv_cache_dtype) { - if (kv_cache_dtype == "auto" || kv_cache_dtype == "float16" || - kv_cache_dtype == "bfloat16") { + if (kv_cache_dtype == "auto" || kv_cache_dtype == "bfloat16") { return flash_v100::KV_CACHE_DTYPE_FP16; } if (kv_cache_dtype == "fp8" || kv_cache_dtype == "fp8_e4m3") { @@ -52,8 +49,6 @@ constexpr int kXQARouteP1024Sawtooth = 4; constexpr int kXQARouteP256Sawtooth = 5; constexpr int kXQARouteP1024SawtoothMid = 6; constexpr int kXQARouteP1024SawtoothFinal = 7; -constexpr int kXQARouteRangeSeqLens = 8; -constexpr int kXQARouteWaveLongSeqLens = 9; // The dense 256/128-half strides alias Volta shared-memory banks in the // WMMA A/B fragment loads. Both padded strides remain 16-byte aligned. constexpr int kXQATC256WidePaddedQStride = 264; @@ -150,41 +145,6 @@ bool xqa_g6_dual_cta_enabled() { return value != nullptr && value[0] == '1'; } -bool xqa_e4m3_batch_enabled() { - const char* value = std::getenv("VLLM_FLASH_V100_E4M3_BATCH_XQA"); - return value == nullptr || value[0] != '0'; -} - -bool xqa_e4m3_batch_optimized_enabled() { - const char* value = std::getenv("VLLM_FLASH_V100_E4M3_BATCH_XQA_OPTIMIZED"); - return value == nullptr || value[0] != '0'; -} - -bool xqa_e4m3_page800_fastpath_enabled() { - const char* value = std::getenv("VLLM_FLASH_V100_E4M3_PAGE800_FASTPATH"); - return value == nullptr || value[0] != '0'; -} - -bool xqa_e4m3_page800_fastpath_trace_enabled() { - const char* value = - std::getenv("VLLM_FLASH_V100_E4M3_PAGE800_FASTPATH_TRACE"); - return value != nullptr && value[0] == '1'; -} - -void trace_xqa_e4m3_page800_fastpath(const int batch_size, - const int partition_size, - const bool active) { - if (!active || !xqa_e4m3_page800_fastpath_trace_enabled()) { - return; - } - static std::atomic traced{false}; - if (!traced.exchange(true, std::memory_order_relaxed)) { - TORCH_WARN("Flash-V100 E4M3 page800 fast path active: batch=", batch_size, - ", partition_size=", partition_size, - ", standard interleaved Hkv=1, PV=half2"); - } -} - bool xqa_e5m2_g6_dual_cta_enabled() { const char* value = std::getenv("VLLM_FLASH_V100_XQA_E5M2_G6_DUAL_CTA"); return value == nullptr || value[0] != '0'; @@ -206,11 +166,6 @@ bool xqa_e5m2_pair_load_enabled() { return value == nullptr || value[0] != '0'; } -bool xqa_e5m2_batch_wide_load_enabled() { - const char* value = std::getenv("VLLM_FLASH_V100_XQA_E5M2_BATCH_WIDE_LOAD"); - 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)); @@ -258,57 +213,6 @@ bool xqa_g6_p1024_sawtooth_enabled() { return value == nullptr || value[0] != '0'; } -bool xqa_e4m3_g6_p64_p256_auto_enabled() { - const char* value = std::getenv("VLLM_FLASH_V100_XQA_E4M3_G6_P64_P256_AUTO"); - return value != nullptr && value[0] == '1'; -} - -bool xqa_e4m3_g6_p64_p256_auto_trace_enabled() { - static const bool enabled = [] { - const char* value = - std::getenv("VLLM_FLASH_V100_XQA_E4M3_G6_P64_P256_AUTO_TRACE"); - return value != nullptr && value[0] == '1'; - }(); - return enabled; -} - -int xqa_e4m3_g6_p256_begin() { - const char* value = std::getenv("VLLM_FLASH_V100_XQA_E4M3_G6_P256_BEGIN"); - return value == nullptr ? 12288 : std::max(1, std::atoi(value)); -} - -int xqa_e4m3_g6_dual_cta_begin() { - const char* value = std::getenv("VLLM_FLASH_V100_XQA_E4M3_G6_DUAL_CTA_BEGIN"); - return value == nullptr ? 32768 : std::max(1, std::atoi(value)); -} - -bool xqa_e4m3_g6_wave_partitions_enabled() { - const char* value = - std::getenv("VLLM_FLASH_V100_XQA_E4M3_G6_WAVE_PARTITIONS"); - return value != nullptr && value[0] == '1'; -} - -bool xqa_e4m3_g6_merged_wave_launch_enabled() { - const char* value = - std::getenv("VLLM_FLASH_V100_XQA_E4M3_G6_MERGED_WAVE_LAUNCH"); - return value != nullptr && value[0] == '1'; -} - -int xqa_e4m3_g6_p512_begin() { - const char* value = std::getenv("VLLM_FLASH_V100_XQA_E4M3_G6_P512_BEGIN"); - return value == nullptr ? 49152 : std::max(1, std::atoi(value)); -} - -int xqa_e4m3_g6_p896_begin() { - const char* value = std::getenv("VLLM_FLASH_V100_XQA_E4M3_G6_P896_BEGIN"); - return value == nullptr ? 98304 : std::max(1, std::atoi(value)); -} - -int xqa_e4m3_g6_p1664_begin() { - const char* value = std::getenv("VLLM_FLASH_V100_XQA_E4M3_G6_P1664_BEGIN"); - return value == nullptr ? 196608 : std::max(1, std::atoi(value)); -} - bool decode_partition_size_overridden() { return std::getenv("VLLM_FLASH_V100_DECODE_PARTITION_SIZE") != nullptr; } @@ -363,102 +267,6 @@ bool xqa_split_reduce_enabled() { return value != nullptr && value[0] == '1'; } -enum class XQABatchContextRoute : int { - kDisabled = -1, - kBaseline = 0, - kDualCta = 1, - kDualCtaSplit = 2, -}; - -bool xqa_batch_context_routing_enabled() { - const char* value = std::getenv("VLLM_FLASH_V100_XQA_BATCH_CONTEXT_ROUTING"); - return value == nullptr || value[0] != '0'; -} - -bool xqa_batch_context_routing_trace_enabled() { - const char* value = - std::getenv("VLLM_FLASH_V100_XQA_BATCH_CONTEXT_ROUTING_TRACE"); - return value != nullptr && value[0] == '1'; -} - -XQABatchContextRoute select_xqa_batch_context_route(const int batch_size, - const int max_seq_len, - const int partition_size) { - if (!xqa_batch_context_routing_enabled() || batch_size < 4 || - max_seq_len <= 0) { - return XQABatchContextRoute::kDisabled; - } - if (batch_size <= 4) { - if (max_seq_len <= 8191) { - return XQABatchContextRoute::kBaseline; - } - if (max_seq_len <= 12287) { - return XQABatchContextRoute::kDualCta; - } - return partition_size == 1024 ? XQABatchContextRoute::kBaseline - : XQABatchContextRoute::kDualCtaSplit; - } - if (batch_size <= 8) { - if (max_seq_len <= 4095) { - return XQABatchContextRoute::kBaseline; - } - return max_seq_len <= 12287 ? XQABatchContextRoute::kDualCta - : XQABatchContextRoute::kDualCtaSplit; - } - if (batch_size <= 12) { - if (max_seq_len <= 2047) { - return XQABatchContextRoute::kBaseline; - } - return max_seq_len <= 16000 ? XQABatchContextRoute::kDualCta - : XQABatchContextRoute::kDualCtaSplit; - } - return max_seq_len <= 12287 ? XQABatchContextRoute::kBaseline - : XQABatchContextRoute::kDualCta; -} - -const char* xqa_batch_context_route_name(const XQABatchContextRoute route) { - switch (route) { - case XQABatchContextRoute::kBaseline: - return "baseline"; - case XQABatchContextRoute::kDualCta: - return "dual_cta"; - case XQABatchContextRoute::kDualCtaSplit: - return "dual_cta_split"; - default: - return "disabled"; - } -} - -void trace_xqa_batch_context_route(const int batch_size, const int max_seq_len, - const int partition_size, - const int block_size, - const XQABatchContextRoute route) { - if (!xqa_batch_context_routing_trace_enabled() || - route == XQABatchContextRoute::kDisabled) { - return; - } - const int batch_class = batch_size <= 4 ? 0 - : batch_size <= 8 ? 1 - : batch_size <= 12 ? 2 - : 3; - const int context_class = max_seq_len <= 4095 ? 0 - : max_seq_len <= 8191 ? 1 - : max_seq_len <= 12287 ? 2 - : max_seq_len <= 16000 ? 3 - : 4; - const unsigned long long bit = - 1ULL << (batch_class * 15 + context_class * 3 + static_cast(route)); - static std::atomic traced_routes{0}; - const unsigned long long previous = - traced_routes.fetch_or(bit, std::memory_order_relaxed); - if ((previous & bit) == 0) { - TORCH_WARN("Flash-V100 XQA batch/context route active: batch=", batch_size, - ", max_seq_len=", max_seq_len, - ", partition_size=", partition_size, ", block_size=", block_size, - ", route=", xqa_batch_context_route_name(route)); - } -} - int xqa_block16_layout_mode() { const char* value = std::getenv("VLLM_FLASH_V100_XQA_BLOCK16_LAYOUT"); if (value == nullptr) { @@ -535,10 +343,6 @@ __device__ __forceinline__ bool xqa_seq_len_route_active( return seq_len >= route_seq_len_begin && seq_len < route_seq_len_end; } else if constexpr (SEQ_LEN_ROUTE == kXQARouteP1024SawtoothFinal) { return seq_len >= route_seq_len_final; - } else if constexpr (SEQ_LEN_ROUTE == kXQARouteRangeSeqLens) { - return seq_len >= route_seq_len_begin && seq_len < route_seq_len_end; - } else if constexpr (SEQ_LEN_ROUTE == kXQARouteWaveLongSeqLens) { - return seq_len >= route_seq_len_begin; } return true; } @@ -630,115 +434,22 @@ __device__ __forceinline__ uint4 fp8_e5m2_vector_to_half8(const uint64_t raw) { fp8_e5m2_pair_to_half2_bits(static_cast(raw >> 48))); } -__device__ __forceinline__ uint16_t fp8_e4m3fn_to_half_bits(const uint8_t raw) { - const uint16_t sign = static_cast(raw & 0x80u) << 8; - const uint8_t magnitude = raw & 0x7fu; - const uint8_t exponent = magnitude >> 3; - const uint8_t mantissa = magnitude & 0x07u; - if (magnitude == 0) { - return sign; - } - if (exponent == 0) { - // E4M3 subnormals are exact fp16 normals: mantissa * 2^-9. - const uint16_t magnitude_bits = - mantissa < 2 - ? 0x1800u - : (mantissa < 4 - ? static_cast(0x1c00u | ((mantissa - 2) << 9)) - : static_cast(0x2000u | ((mantissa - 4) << 8))); - return sign | magnitude_bits; - } - if (magnitude == 0x7fu) { - return sign | 0x7e00u; - } - return sign | static_cast((exponent + 8) << 10) | - static_cast(mantissa << 7); -} - -__device__ __forceinline__ uint32_t -fp8_e4m3fn_pair_to_half2_bits(const uint16_t raw_pair) { - return static_cast( - fp8_e4m3fn_to_half_bits(static_cast(raw_pair))) | - (static_cast( - fp8_e4m3fn_to_half_bits(static_cast(raw_pair >> 8))) - << 16); -} - -__device__ __forceinline__ uint32_t -fp8_e4m3fn_pair_to_half2_bits_fast(const uint16_t raw_pair) { - const uint8_t raw0 = static_cast(raw_pair); - const uint8_t raw1 = static_cast(raw_pair >> 8); - if ((raw0 & 0x7fu) == 0x7fu || (raw1 & 0x7fu) == 0x7fu) { - return fp8_e4m3fn_pair_to_half2_bits(raw_pair); - } - - // Moving a finite E4M3 encoding into the corresponding fp16 sign, - // exponent, and mantissa fields represents exactly value / 256. A packed - // half2 multiply restores both values without per-byte exponent branches. - const uint32_t expanded = (static_cast(raw_pair & 0x0080u) << 8) | - (static_cast(raw_pair & 0x007fu) << 7) | - (static_cast(raw_pair & 0x8000u) << 16) | - (static_cast(raw_pair & 0x7f00u) << 15); - union { - uint32_t u; - __half2 h2; - } converter; - converter.u = expanded; - converter.h2 = __hmul2(converter.h2, __float2half2_rn(256.0f)); - return converter.u; -} - -__device__ __forceinline__ uint4 -fp8_e4m3fn_vector_to_half8(const uint64_t raw) { - return make_uint4( - fp8_e4m3fn_pair_to_half2_bits(static_cast(raw)), - fp8_e4m3fn_pair_to_half2_bits(static_cast(raw >> 16)), - fp8_e4m3fn_pair_to_half2_bits(static_cast(raw >> 32)), - fp8_e4m3fn_pair_to_half2_bits(static_cast(raw >> 48))); -} - -__device__ __forceinline__ uint4 -fp8_e4m3fn_vector_to_half8_fast(const uint64_t raw) { - return make_uint4( - fp8_e4m3fn_pair_to_half2_bits_fast(static_cast(raw)), - fp8_e4m3fn_pair_to_half2_bits_fast(static_cast(raw >> 16)), - fp8_e4m3fn_pair_to_half2_bits_fast(static_cast(raw >> 32)), - fp8_e4m3fn_pair_to_half2_bits_fast(static_cast(raw >> 48))); -} - -__device__ __forceinline__ uint4 fp8_e4m3fn_vector_to_half8_lut( - const uint64_t raw, const uint16_t* __restrict__ lut) { - return make_uint4( - static_cast(lut[static_cast(raw)]) | - (static_cast(lut[static_cast(raw >> 8)]) << 16), - static_cast(lut[static_cast(raw >> 16)]) | - (static_cast(lut[static_cast(raw >> 24)]) << 16), - static_cast(lut[static_cast(raw >> 32)]) | - (static_cast(lut[static_cast(raw >> 40)]) << 16), - static_cast(lut[static_cast(raw >> 48)]) | - (static_cast(lut[static_cast(raw >> 56)]) << 16)); -} - template + int KV_DTYPE = flash_v100::KV_CACHE_DTYPE_FP16> __device__ __forceinline__ uint4 load_xqa_tc_kv_vector( const void* __restrict__ kv_cache, const int* __restrict__ page_ids, const int copy_idx, const int panel_d_stride_uint4, const int tile_page_offset, const int kv_tile_start, const int block_size, const int kv_head_idx, const int64_t block_stride, const int64_t token_stride, const int64_t head_stride, - const int panel_offset, const uint16_t* __restrict__ e4m3_lut = nullptr) { + const int panel_offset) { const int row = copy_idx / panel_d_stride_uint4; const int vec_col = copy_idx % panel_d_stride_uint4; const int token_offset = tile_page_offset + kv_tile_start + row; - static_assert(BLOCK_SIZE == 0 || BLOCK_SIZE == 16 || BLOCK_SIZE == 784 || - BLOCK_SIZE == 800 || BLOCK_SIZE == 1568 || - BLOCK_SIZE == 1648 || BLOCK_SIZE == 3296, + static_assert(BLOCK_SIZE == 0 || BLOCK_SIZE == 16 || BLOCK_SIZE == 784, "Unsupported paged-KV block-size specialization"); - static_assert(!CONTIGUOUS_HKV1_LAYOUT || BLOCK_SIZE == 16 || - BLOCK_SIZE == 800 || BLOCK_SIZE == 1568, - "The fixed-stride Hkv=1 layout requires a specialized page"); + static_assert(!CONTIGUOUS_HKV1_LAYOUT || BLOCK_SIZE == 16, + "The contiguous Hkv=1 layout requires 16-token pages"); int logical_block; int block_offset; if constexpr (BLOCK_SIZE == 16) { @@ -747,18 +458,6 @@ __device__ __forceinline__ uint4 load_xqa_tc_kv_vector( } else if constexpr (BLOCK_SIZE == 784) { logical_block = token_offset / 784; block_offset = token_offset - logical_block * 784; - } else if constexpr (BLOCK_SIZE == 800) { - logical_block = token_offset >= 800; - block_offset = token_offset - logical_block * 800; - } else if constexpr (BLOCK_SIZE == 1568) { - logical_block = token_offset / 1568; - block_offset = token_offset - logical_block * 1568; - } else if constexpr (BLOCK_SIZE == 1648) { - logical_block = token_offset / 1648; - block_offset = token_offset - logical_block * 1648; - } else if constexpr (BLOCK_SIZE == 3296) { - logical_block = token_offset / 3296; - block_offset = token_offset - logical_block * 3296; } else { logical_block = token_offset / block_size; block_offset = token_offset % block_size; @@ -766,9 +465,8 @@ __device__ __forceinline__ uint4 load_xqa_tc_kv_vector( const int physical_block = page_ids[logical_block]; int64_t physical_offset; if constexpr (CONTIGUOUS_HKV1_LAYOUT) { - constexpr int64_t kHeadDim = 256; - constexpr int64_t kBlockStride = - BLOCK_SIZE == 16 ? 16 * kHeadDim : 2 * BLOCK_SIZE * kHeadDim; + constexpr int kHeadDim = 256; + constexpr int kBlockStride = 16 * kHeadDim; physical_offset = static_cast(physical_block) * kBlockStride + static_cast(block_offset) * kHeadDim + panel_offset; @@ -782,28 +480,17 @@ __device__ __forceinline__ uint4 load_xqa_tc_kv_vector( const uint4* cache_vec = reinterpret_cast(kv_cache); return __ldg(&cache_vec[physical_offset / 8 + vec_col]); } else { - static_assert(KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP8_E4M3 || - KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP8_E5M2, - "XQA only supports fp16, FP8 E4M3, and FP8 E5M2 KV"); + static_assert(KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP8_E5M2, + "XQA only supports fp16 and FP8 E5M2 KV"); const uint64_t* cache_vec = reinterpret_cast(kv_cache); - const uint64_t raw = __ldg(&cache_vec[physical_offset / 8 + vec_col]); - if constexpr (KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP8_E4M3) { - if constexpr (E4M3_SHARED_LUT) { - return fp8_e4m3fn_vector_to_half8_lut(raw, e4m3_lut); - } else { - return fp8_e4m3fn_vector_to_half8(raw); - } - } else { - static_assert(!E4M3_SHARED_LUT, - "The E4M3 conversion LUT requires E4M3 KV"); - return fp8_e5m2_vector_to_half8(raw); - } + return fp8_e5m2_vector_to_half8( + __ldg(&cache_vec[physical_offset / 8 + vec_col])); } } template + bool E5M2_PAIR_LOAD = false> __device__ __forceinline__ void load_xqa_tc_kv_panel( __half* __restrict__ shared_kv, const void* __restrict__ kv_cache, const int* __restrict__ page_ids, const int valid_kv_tile_rows, @@ -811,15 +498,11 @@ __device__ __forceinline__ void load_xqa_tc_kv_panel( const int tile_page_offset, const int kv_tile_start, const int block_size, const int kv_head_idx, const int64_t block_stride, const int64_t token_stride, const int64_t head_stride, - const int panel_offset, const int copy_thread_idx = threadIdx.x, - const uint16_t* __restrict__ e4m3_lut = nullptr) { + const int panel_offset, const int copy_thread_idx = threadIdx.x) { uint4* shared_vec = reinterpret_cast(shared_kv); - if constexpr (FP8_PAIR_LOAD) { - static_assert(KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP8_E4M3 || - KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP8_E5M2, - "Paired XQA loads require FP8 KV"); - static_assert(!E4M3_SHARED_LUT, - "Paired E4M3 conversion does not use the shared LUT"); + if constexpr (E5M2_PAIR_LOAD) { + static_assert(KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP8_E5M2, + "Paired XQA loads are only valid for FP8 E5M2 KV"); const int pair_stride = panel_d_stride_uint4 / 2; const int pair_count = valid_kv_tile_rows * pair_stride; for (int pair_idx = copy_thread_idx; pair_idx < pair_count; @@ -835,18 +518,6 @@ __device__ __forceinline__ void load_xqa_tc_kv_panel( } else if constexpr (BLOCK_SIZE == 784) { logical_block = token_offset / 784; block_offset = token_offset - logical_block * 784; - } else if constexpr (BLOCK_SIZE == 800) { - logical_block = token_offset >= 800; - block_offset = token_offset - logical_block * 800; - } else if constexpr (BLOCK_SIZE == 1568) { - logical_block = token_offset / 1568; - block_offset = token_offset - logical_block * 1568; - } else if constexpr (BLOCK_SIZE == 1648) { - logical_block = token_offset / 1648; - block_offset = token_offset - logical_block * 1648; - } else if constexpr (BLOCK_SIZE == 3296) { - logical_block = token_offset / 3296; - block_offset = token_offset - logical_block * 3296; } else { logical_block = token_offset / block_size; block_offset = token_offset % block_size; @@ -854,12 +525,11 @@ __device__ __forceinline__ void load_xqa_tc_kv_panel( const int physical_block = page_ids[logical_block]; int64_t physical_offset; if constexpr (CONTIGUOUS_HKV1_LAYOUT) { - constexpr int64_t kHeadDim = 256; - constexpr int64_t kPhysicalBlockStride = - BLOCK_SIZE == 16 ? 16 * kHeadDim : 2 * 800 * kHeadDim; - physical_offset = - static_cast(physical_block) * kPhysicalBlockStride + - static_cast(block_offset) * kHeadDim + panel_offset; + constexpr int kHeadDim = 256; + constexpr int kBlockStride = 16 * kHeadDim; + physical_offset = static_cast(physical_block) * kBlockStride + + static_cast(block_offset) * kHeadDim + + panel_offset; } else { physical_offset = static_cast(physical_block) * block_stride + static_cast(block_offset) * token_stride + @@ -871,16 +541,10 @@ __device__ __forceinline__ void load_xqa_tc_kv_panel( const int shared_offset = row * kv_smem_stride_uint4 + vec_pair * 2; const uint64_t raw_lo = static_cast(raw.x) | (static_cast(raw.y) << 32); - shared_vec[shared_offset] = - KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP8_E4M3 - ? fp8_e4m3fn_vector_to_half8_fast(raw_lo) - : fp8_e5m2_vector_to_half8(raw_lo); + shared_vec[shared_offset] = fp8_e5m2_vector_to_half8(raw_lo); const uint64_t raw_hi = static_cast(raw.z) | (static_cast(raw.w) << 32); - shared_vec[shared_offset + 1] = - KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP8_E4M3 - ? fp8_e4m3fn_vector_to_half8_fast(raw_hi) - : fp8_e5m2_vector_to_half8(raw_hi); + shared_vec[shared_offset + 1] = fp8_e5m2_vector_to_half8(raw_hi); } } else { const int copy_count = valid_kv_tile_rows * panel_d_stride_uint4; @@ -889,18 +553,16 @@ __device__ __forceinline__ void load_xqa_tc_kv_panel( const int row = copy_idx / panel_d_stride_uint4; const int vec_col = copy_idx % panel_d_stride_uint4; shared_vec[row * kv_smem_stride_uint4 + vec_col] = - load_xqa_tc_kv_vector( + load_xqa_tc_kv_vector( kv_cache, page_ids, copy_idx, panel_d_stride_uint4, tile_page_offset, kv_tile_start, block_size, kv_head_idx, - block_stride, token_stride, head_stride, panel_offset, e4m3_lut); + block_stride, token_stride, head_stride, panel_offset); } } } template + int KV_DTYPE = flash_v100::KV_CACHE_DTYPE_FP16> __device__ __forceinline__ void load_xqa_tc_kv_panel_and_zero( __half* __restrict__ shared_kv, const void* __restrict__ kv_cache, const int* __restrict__ page_ids, const int valid_kv_tile_rows, @@ -908,8 +570,7 @@ __device__ __forceinline__ void load_xqa_tc_kv_panel_and_zero( const int tile_page_offset, const int kv_tile_start, const int block_size, const int kv_head_idx, const int64_t block_stride, const int64_t token_stride, const int64_t head_stride, - const int panel_offset, const int copy_thread_idx, - const uint16_t* __restrict__ e4m3_lut = nullptr) { + const int panel_offset, const int copy_thread_idx) { const int copy_count = valid_kv_tile_rows * panel_d_stride_uint4; uint4* shared_vec = reinterpret_cast(shared_kv); constexpr int kLoadStages = 4; @@ -921,12 +582,10 @@ __device__ __forceinline__ void load_xqa_tc_kv_panel_and_zero( const int copy_idx = copy_base + stage * NUM_THREADS; if (copy_idx < copy_count) { staged[stage] = - load_xqa_tc_kv_vector( + load_xqa_tc_kv_vector( kv_cache, page_ids, copy_idx, panel_d_stride_uint4, tile_page_offset, kv_tile_start, block_size, kv_head_idx, - block_stride, token_stride, head_stride, panel_offset, - e4m3_lut); + block_stride, token_stride, head_stride, panel_offset); } } #pragma unroll @@ -1004,7 +663,7 @@ __device__ __forceinline__ float dot_qk_cache(const __half* __restrict__ q_ptr, } template + int SEQ_LEN_ROUTE = kXQARouteAllSeqLens> __global__ void flash_attention_decode_partition_kernel( const __half* __restrict__ q, const void* __restrict__ k_cache, const void* __restrict__ v_cache, __half* __restrict__ tmp_out, @@ -1022,12 +681,7 @@ __global__ void flash_attention_decode_partition_kernel( const int64_t v_head_stride, const float softmax_scale, const float k_scale, const float v_scale, const int window_size_left, const int window_size_right, const int route_seq_len_begin, - const int route_seq_len_end, const int route_seq_len_final, - const int* __restrict__ anchor_lens, const int anchored_window) { - // The anchored decode-window mask variant is only generated for the fp16-KV - // configuration; every other instantiation keeps its original code path. - static_assert(!ANCHORED_SWA || KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP16, - "ANCHORED_SWA requires an fp16 KV cache"); + const int route_seq_len_end, const int route_seq_len_final) { const int batch_idx = blockIdx.x; const int head_idx = blockIdx.y; const int partition_idx = blockIdx.z; @@ -1064,17 +718,6 @@ __global__ void flash_attention_decode_partition_kernel( : seq_len - 1; const int part_start = max(start_token_idx, min_token_idx); const int part_end = min(start_token_idx + PARTITION_SIZE, max_token_idx + 1); - - // Anchored decode window: keys below the per-request prompt length stay - // globally visible; later (generated) keys must fall inside the sliding - // window over generated positions. Keys inside [gap_lo, gap_hi) are masked. - int gap_lo = seq_len; - int gap_hi = seq_len; - if constexpr (ANCHORED_SWA) { - const int anchor_len = anchor_lens[batch_idx]; - gap_lo = anchor_len; - gap_hi = max(anchor_len, query_pos - anchored_window + 1); - } const int q_per_kv = num_heads_q / num_heads_kv; const int kv_head_idx = head_idx / q_per_kv; const int lane = threadIdx.x % kWarpSize; @@ -1134,21 +777,9 @@ __global__ void flash_attention_decode_partition_kernel( float score = dot_qk_cache(q_shared, k_cache, k_index, lane); if (lane == 0) { - if constexpr (ANCHORED_SWA) { - const int token_idx = part_start + token_local; - if (token_idx >= gap_lo && token_idx < gap_hi) { - // Masked gap key: never contributes (nulled gap blocks included). - scores_shared[token_local] = -1.0e30f; - } else { - score *= score_scale; - scores_shared[token_local] = score; - local_max = fmaxf(local_max, score); - } - } else { - score *= score_scale; - scores_shared[token_local] = score; - local_max = fmaxf(local_max, score); - } + score *= score_scale; + scores_shared[token_local] = score; + local_max = fmaxf(local_max, score); } } @@ -1156,14 +787,6 @@ __global__ void flash_attention_decode_partition_kernel( float local_sum = 0.f; for (int i = threadIdx.x; i < part_tokens; i += blockDim.x) { - if constexpr (ANCHORED_SWA) { - const int token_idx = part_start + i; - if (token_idx >= gap_lo && token_idx < gap_hi) { - // Exact zero weight for masked keys (no reliance on expf underflow). - scores_shared[i] = 0.f; - continue; - } - } const float p = __expf(scores_shared[i] - part_max); scores_shared[i] = p; local_sum += p; @@ -1204,8 +827,7 @@ template + bool PARTITION_PAGE_IDS = false, bool E5M2_PAIR_LOAD = false> __global__ void __launch_bounds__(NUM_THREADS, MIN_BLOCKS_PER_SM) flash_attention_decode_xqa_tc_partition_kernel_256_wide( const __half* __restrict__ q, const void* __restrict__ k_cache, @@ -1249,9 +871,6 @@ __global__ void __launch_bounds__(NUM_THREADS, MIN_BLOCKS_PER_SM) QK_SW_PIPELINE ? SmemLayout::kQKStride / 8 : kv_smem_stride_uint4; constexpr int qk_panel_d_stride_uint4 = kQKPanelDim / 8; constexpr int kAccumsPerThread = D / kWarpSize; - constexpr bool kPVHalf2 = CONTIGUOUS_HKV1_LAYOUT && - (BLOCK_SIZE == 800 || BLOCK_SIZE == 1568) && - KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP8_E4M3; static_assert(GROUP_SIZE == 4 || GROUP_SIZE == 6 || GROUP_SIZE == 8, "Wide D=256 TC XQA kernel supports q_per_kv in {4, 6, 8}"); static_assert(NUM_THREADS >= GROUP_SIZE * kWarpSize, @@ -1261,9 +880,6 @@ __global__ void __launch_bounds__(NUM_THREADS, MIN_BLOCKS_PER_SM) NUM_THREADS == 8 * kWarpSize)), "The QK software pipeline requires a six- or eight-warp G6 " "kernel"); - static_assert( - !E4M3_SHARED_LUT || KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP8_E4M3, - "The shared conversion LUT is only valid for E4M3 KV"); const int batch_idx = blockIdx.x; const int kv_head_idx = blockIdx.y; const int partition_idx = blockIdx.z; @@ -1279,27 +895,13 @@ __global__ void __launch_bounds__(NUM_THREADS, MIN_BLOCKS_PER_SM) route_seq_len_final)) { return; } - int partition_size; - if constexpr (PARTITION_SIZE == -1) { - static_assert(SEQ_LEN_ROUTE == kXQARouteWaveLongSeqLens, - "Runtime wave partitions require the wave-long route"); - partition_size = seq_len < route_seq_len_end - ? 512 - : (seq_len < route_seq_len_final ? 896 : 1664); - } else if constexpr (PARTITION_SIZE == 0) { - const int partition_size_begin = - route_seq_len_end > 0 ? route_seq_len_end : route_seq_len_begin; - partition_size = seq_len < partition_size_begin ? 64 : 256; - } else { - partition_size = PARTITION_SIZE; - } - const int start_token_idx = partition_idx * partition_size; + const int start_token_idx = partition_idx * PARTITION_SIZE; if (seq_len <= 0 || start_token_idx >= seq_len) { return; } const int runtime_num_partitions = active_num_partitions[0]; const int seq_num_partitions = - (seq_len + partition_size - 1) / partition_size; + (seq_len + PARTITION_SIZE - 1) / PARTITION_SIZE; const int effective_num_partitions = min(max_num_partitions, max(runtime_num_partitions, seq_num_partitions)); if (partition_idx >= effective_num_partitions) { @@ -1314,19 +916,12 @@ __global__ void __launch_bounds__(NUM_THREADS, MIN_BLOCKS_PER_SM) const int tid = threadIdx.x; const int warp_id = tid / kWarpSize; const int lane_id = tid % kWarpSize; - const int part_tokens = min(partition_size, seq_len - start_token_idx); + const int part_tokens = min(PARTITION_SIZE, seq_len - start_token_idx); const int num_k_tiles = (part_tokens + kXQATCBlockN - 1) / kXQATCBlockN; const int* block_table_seq = block_table + batch_idx * max_num_blocks; extern __shared__ char smem_raw[]; auto& smem = *reinterpret_cast(smem_raw); - uint16_t* e4m3_lut = nullptr; - if constexpr (E4M3_SHARED_LUT) { - e4m3_lut = reinterpret_cast(smem_raw + sizeof(SmemLayout)); - for (int raw = tid; raw < 256; raw += NUM_THREADS) { - e4m3_lut[raw] = fp8_e4m3fn_to_half_bits(static_cast(raw)); - } - } __half* sQ = smem.q; __half* sK = smem.k_buffer(0); __half* sV = smem.v_buffer(); @@ -1364,18 +959,10 @@ __global__ void __launch_bounds__(NUM_THREADS, MIN_BLOCKS_PER_SM) int partition_start_page = 0; int partition_page_offset = 0; if constexpr (PARTITION_PAGE_IDS) { - int partition_page_count; - if constexpr (BLOCK_SIZE == 800) { - partition_start_page = start_token_idx / 800; - partition_page_offset = start_token_idx - partition_start_page * 800; - partition_page_count = 1 + (partition_page_offset + part_tokens > 800); - } else { - partition_start_page = start_token_idx / block_size; - partition_page_offset = - start_token_idx - partition_start_page * block_size; - partition_page_count = - (partition_page_offset + part_tokens + block_size - 1) / block_size; - } + partition_start_page = start_token_idx / block_size; + partition_page_offset = start_token_idx - partition_start_page * block_size; + const int partition_page_count = + (partition_page_offset + part_tokens + block_size - 1) / block_size; for (int idx = tid; idx < partition_page_count; idx += NUM_THREADS) { smem.page_ids[idx] = __ldg(&block_table_seq[partition_start_page + idx]); } @@ -1399,14 +986,6 @@ __global__ void __launch_bounds__(NUM_THREADS, MIN_BLOCKS_PER_SM) start_page = tile_token_start / 784; tile_page_offset = tile_token_start - start_page * 784; page_count = (tile_page_offset + valid_k_rows + 783) / 784; - } else if constexpr (BLOCK_SIZE == 800) { - start_page = tile_token_start / 800; - tile_page_offset = tile_token_start - start_page * 800; - page_count = 1 + (tile_page_offset + valid_k_rows > 800); - } else if constexpr (BLOCK_SIZE == 1568) { - start_page = tile_token_start / 1568; - tile_page_offset = tile_token_start - start_page * 1568; - page_count = (tile_page_offset + valid_k_rows + 1567) / 1568; } else { start_page = tile_token_start / block_size; tile_page_offset = tile_token_start - start_page * block_size; @@ -1446,12 +1025,11 @@ __global__ void __launch_bounds__(NUM_THREADS, MIN_BLOCKS_PER_SM) const int producer_tid = tid - kConsumerWarps * kWarpSize; if (producer_tid >= 0) { load_xqa_tc_kv_panel_and_zero( + kProducerThreads, KV_DTYPE>( smem.k_buffer(0), k_cache, smem.page_ids, valid_kv_tile_rows, qk_panel_d_stride_uint4, qk_smem_stride_uint4, tile_page_offset, kv_tile_start, block_size, kv_head_idx, k_block_stride, - k_token_stride, k_head_stride, 0, producer_tid, e4m3_lut); + k_token_stride, k_head_stride, 0, producer_tid); } __syncthreads(); @@ -1460,14 +1038,12 @@ __global__ void __launch_bounds__(NUM_THREADS, MIN_BLOCKS_PER_SM) const int panel_offset = panel_idx * kQKPanelDim; if (producer_tid >= 0 && panel_idx + 1 < kNumQKPanels) { load_xqa_tc_kv_panel_and_zero( + kProducerThreads, KV_DTYPE>( smem.k_buffer((panel_idx + 1) & 1), k_cache, smem.page_ids, valid_kv_tile_rows, qk_panel_d_stride_uint4, qk_smem_stride_uint4, tile_page_offset, kv_tile_start, block_size, kv_head_idx, k_block_stride, k_token_stride, - k_head_stride, panel_offset + kQKPanelDim, producer_tid, - e4m3_lut); + k_head_stride, panel_offset + kQKPanelDim, producer_tid); } if (warp_id < kConsumerWarps) { const int tile_n = warp_id * WMMA_N; @@ -1491,12 +1067,11 @@ __global__ void __launch_bounds__(NUM_THREADS, MIN_BLOCKS_PER_SM) for (int panel_idx = 0; panel_idx < kNumQKPanels; ++panel_idx) { const int panel_offset = panel_idx * kQKPanelDim; load_xqa_tc_kv_panel( + KV_DTYPE, E5M2_PAIR_LOAD>( sK, k_cache, smem.page_ids, valid_kv_tile_rows, qk_panel_d_stride_uint4, qk_smem_stride_uint4, tile_page_offset, kv_tile_start, block_size, kv_head_idx, k_block_stride, - k_token_stride, k_head_stride, panel_offset, threadIdx.x, - e4m3_lut); + k_token_stride, k_head_stride, panel_offset); for (int idx = tid + valid_kv_tile_rows * qk_panel_d_stride_uint4; idx < kXQATCBlockN * qk_panel_d_stride_uint4; idx += NUM_THREADS) { @@ -1641,11 +1216,11 @@ __global__ void __launch_bounds__(NUM_THREADS, MIN_BLOCKS_PER_SM) const int valid_kv_tile_rows = min(kXQATCBlockN, valid_k_rows - kv_tile_start); load_xqa_tc_kv_panel( + KV_DTYPE, E5M2_PAIR_LOAD>( sV, v_cache, smem.page_ids, valid_kv_tile_rows, pv_panel_d_stride_uint4, kv_smem_stride_uint4, tile_page_offset, kv_tile_start, block_size, kv_head_idx, v_block_stride, - v_token_stride, v_head_stride, panel_offset, threadIdx.x, e4m3_lut); + v_token_stride, v_head_stride, panel_offset); for (int idx = tid + valid_kv_tile_rows * pv_panel_d_stride_uint4; idx < kXQATCBlockN * pv_panel_d_stride_uint4; idx += NUM_THREADS) { const int row = idx / pv_panel_d_stride_uint4; @@ -1671,29 +1246,13 @@ __global__ void __launch_bounds__(NUM_THREADS, MIN_BLOCKS_PER_SM) prob = __half2float(sP_row[token]); } const __half* sV_row = sV + token * SmemLayout::kKVStride; - if constexpr (kPVHalf2) { #pragma unroll - for (int d_iter = 0; d_iter < (kPVPanelDim / (2 * kWarpSize)); - ++d_iter) { - const int local_d2 = lane_id + d_iter * kWarpSize; - const float2 value = __half22float2( - reinterpret_cast(sV_row)[local_d2]); - const int acc_idx = - panel_idx * (kPVPanelDim / kWarpSize) + d_iter * 2; - out_acc[acc_idx] = fmaf(prob, value.x, out_acc[acc_idx]); - out_acc[acc_idx + 1] = - fmaf(prob, value.y, out_acc[acc_idx + 1]); - } - } else { -#pragma unroll - for (int d_iter = 0; d_iter < (kPVPanelDim / kWarpSize); - ++d_iter) { - const int local_d = lane_id + d_iter * kWarpSize; - const int acc_idx = - panel_idx * (kPVPanelDim / kWarpSize) + d_iter; - out_acc[acc_idx] = - fmaf(prob, __half2float(sV_row[local_d]), out_acc[acc_idx]); - } + for (int d_iter = 0; d_iter < (kPVPanelDim / kWarpSize); ++d_iter) { + const int local_d = lane_id + d_iter * kWarpSize; + const int acc_idx = + panel_idx * (kPVPanelDim / kWarpSize) + d_iter; + out_acc[acc_idx] = + fmaf(prob, __half2float(sV_row[local_d]), out_acc[acc_idx]); } } } @@ -1722,25 +1281,11 @@ __global__ void __launch_bounds__(NUM_THREADS, MIN_BLOCKS_PER_SM) static_cast(batch_idx) * tmp_out_stride0 + static_cast(head_idx) * tmp_out_stride1 + static_cast(partition_idx) * tmp_out_stride2; - const float output_scale = KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP16 - ? inv_row_sum - : inv_row_sum * v_scale; - if constexpr (kPVHalf2) { -#pragma unroll - for (int acc_idx = 0; acc_idx < kAccumsPerThread; ++acc_idx) { - constexpr int kAccumsPerPanel = kPVPanelDim / kWarpSize; - const int panel_idx = acc_idx / kAccumsPerPanel; - const int panel_acc_idx = acc_idx % kAccumsPerPanel; - const int d_iter = panel_acc_idx / 2; - const int component = panel_acc_idx & 1; - const int d = panel_idx * kPVPanelDim + d_iter * (2 * kWarpSize) + - thread_in_row * 2 + component; - tmp_out_ptr[d] = __float2half(out_acc[acc_idx] * output_scale); - } - } else { - for (int d = thread_in_row; d < D; d += kXQATC256WideThreadsPerRow) { - tmp_out_ptr[d] = __float2half(out_acc[d / kWarpSize] * output_scale); - } + for (int d = thread_in_row; d < D; d += kXQATC256WideThreadsPerRow) { + const float output_scale = KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP16 + ? inv_row_sum + : inv_row_sum * v_scale; + tmp_out_ptr[d] = __float2half(out_acc[d / kWarpSize] * output_scale); } if (thread_in_row == 0) { const int64_t stats_index = @@ -1752,529 +1297,7 @@ __global__ void __launch_bounds__(NUM_THREADS, MIN_BLOCKS_PER_SM) } } -// Exact grouped selector-verifier kernel for the admitted TP4 tensor contract. -// -// Independent-row decode scans the same paged KV prefix once for every -// verifier token. This kernel instead assigns one packed GQA group to each of -// eighty context splits. Every CTA handles all six query heads for all eight -// verifier tokens, so K/V is shared across the 48 rows before the exact -// split-softmax states are merged. -// -// The diagnostic two-pass variant intentionally uses: -// 1. QK computes the exact FP32 max/sum for each query row. -// 2. QK is recomputed, then half P x half V is accumulated with Volta WMMA. -// 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; -constexpr int kGroupedVerifyHeads = 6; -constexpr int kGroupedVerifyHeadDim = 256; -constexpr int kGroupedVerifyHeadsPerCta = 6; -constexpr int kGroupedVerifyRows = - kGroupedVerifyMaxQ * kGroupedVerifyHeadsPerCta; -constexpr int kGroupedVerifyBlockN = 32; -// 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; -// 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 -// unchanged eighty splits. The single-query utility variant keeps the old -// 128-token split contract for its stricter FP32-reference error envelope. -constexpr int kGroupedVerifyMinTokensPerSplit = 64; -constexpr int kGroupedVerifySingleQueryMinTokensPerSplit = 128; -constexpr int kGroupedVerifyShortContextMaxTokens = 128; -constexpr int kGroupedVerifyThreads = 512; -constexpr int kGroupedVerifyWarps = kGroupedVerifyThreads / kWarpSize; -constexpr int kGroupedVerifyQKWarps = - (kGroupedVerifyRows / 16) * (kGroupedVerifyBlockN / 16); -constexpr int kGroupedVerifyOutputTiles = - (kGroupedVerifyRows / 16) * (kGroupedVerifyHeadDim / 16); -constexpr int kGroupedVerifyOutputTilesPerWarp = - kGroupedVerifyOutputTiles / kGroupedVerifyWarps; - -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]; - } 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]; -}; - -static_assert(sizeof(GroupedVerifySmem) <= 64 * 1024, - "packed grouped verifier must fit Volta's 64 KiB opt-in budget"); -static_assert(kGroupedVerifyOutputTiles % kGroupedVerifyWarps == 0, - "output tiles must divide evenly across warps"); - -template -__device__ __forceinline__ int grouped_verify_active_splits( - const int total_kv) { - constexpr int kMinTokensPerSplit = - SINGLE_QUERY ? kGroupedVerifySingleQueryMinTokensPerSplit - : kGroupedVerifyMinTokensPerSplit; - const int active_splits = - min(kGroupedVerifySplits, - max(1, (total_kv + kMinTokensPerSplit - 1) / kMinTokensPerSplit)); - if constexpr (SINGLE_QUERY) { - return active_splits; - } - return total_kv <= kGroupedVerifyShortContextMaxTokens ? 1 : active_splits; -} - -__device__ __forceinline__ void grouped_verify_qk( - const __half* __restrict__ shared_q, const __half* __restrict__ shared_k, - float* __restrict__ shared_scores, const float qk_scale) { - const int warp_id = threadIdx.x / kWarpSize; - if (warp_id >= kGroupedVerifyQKWarps) { - return; - } - - const int m_tile = warp_id / (kGroupedVerifyBlockN / 16); - const int n_tile = warp_id % (kGroupedVerifyBlockN / 16); - volta::fragment - q_fragment; - volta::fragment - k_fragment; - volta::fragment score_fragment; - volta::fill_fragment(score_fragment, 0.0f); - -#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); - // 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); - volta::mma_sync(score_fragment, q_fragment, k_fragment, score_fragment); - } -#pragma unroll - for (int i = 0; i < score_fragment.num_elements; ++i) { - 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); -} - -__device__ __forceinline__ void grouped_verify_scale_output_fragment( - volta::fragment& fragment, - const float* __restrict__ row_scale, const int tile_row_start) { - const int lane = threadIdx.x & 31; - const int row = (lane & 1) + ((lane >> 2) & 1) * 8 + ((lane >> 4) & 1) * 4; - const float first_scale = row_scale[tile_row_start + row]; - const float second_scale = row_scale[tile_row_start + row + 2]; - fragment.x[0] *= first_scale; - fragment.x[1] *= first_scale; - fragment.x[2] *= second_scale; - fragment.x[3] *= second_scale; - fragment.x[4] *= first_scale; - fragment.x[5] *= first_scale; - fragment.x[6] *= second_scale; - fragment.x[7] *= second_scale; -} - -template -__global__ -__launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_e5m2_partial_kernel( - const __half* __restrict__ q, const void* __restrict__ k_cache, - const void* __restrict__ v_cache, const int* __restrict__ block_table, - const int* __restrict__ seq_lens, __half* __restrict__ partial_out, - float* __restrict__ partial_lse, const int query_len, - const int max_num_blocks, const int page_block_size, - const int64_t k_block_stride, const int64_t k_token_stride, - 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) { - 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) { - return; - } - - const int total_kv = seq_lens[0]; - if (total_kv <= 0) { - return; - } - const int active_splits = - grouped_verify_active_splits(total_kv); - if (split_id >= active_splits) { - return; - } - const int total_tiles = - (total_kv + kGroupedVerifyBlockN - 1) / kGroupedVerifyBlockN; - const int base_tiles = total_tiles / active_splits; - const int extra_tiles = total_tiles - base_tiles * active_splits; - const int split_tile_start = - split_id * base_tiles + min(split_id, extra_tiles); - const int split_tiles = base_tiles + (split_id < extra_tiles ? 1 : 0); - const int split_start = split_tile_start * kGroupedVerifyBlockN; - 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 tid = threadIdx.x; - const int warp_id = tid / kWarpSize; - const int lane_id = tid % kWarpSize; - - extern __shared__ char grouped_verify_smem_raw[]; - GroupedVerifySmem& smem = - *reinterpret_cast(grouped_verify_smem_raw); - __half* shared_q = smem.storage.compute.q; - __half* shared_kv = smem.storage.compute.kv; - float* shared_scores = smem.storage.compute.scores; - __half* shared_probs = smem.storage.compute.probs; - - constexpr int kVecsPerRow = kGroupedVerifyHeadDim / 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 head_idx = head_start + local_head; - if (token_idx < query_len && head_idx < kGroupedVerifyHeads) { - shared_q_vec[idx] = __ldg( - q_vec + (token_idx * kGroupedVerifyHeads + head_idx) * kVecsPerRow + - vec_col); - } else { - shared_q_vec[idx] = make_uint4(0, 0, 0, 0); - } - } - if (tid < kGroupedVerifyRows) { - smem.row_max[tid] = kXQANegInf; - smem.row_sum[tid] = 0.0f; - smem.row_scale[tid] = 1.0f; - } - __syncthreads(); - - constexpr int kPanelStrideVec = kGroupedVerifyHeadDim / 8; - constexpr int kSharedStrideVec = kGroupedVerifyHeadDim / 8; - - volta::fragment - output_fragments[kGroupedVerifyOutputTilesPerWarp]; -#pragma unroll - for (int fragment_idx = 0; fragment_idx < kGroupedVerifyOutputTilesPerWarp; - ++fragment_idx) { - volta::fill_fragment(output_fragments[fragment_idx], 0.0f); - } - - // The conservative baseline computes exact per-split FP32 max/sum in a - // separate pass. The candidate below fuses this state update with P x V. - if constexpr (TWO_PASS) { - 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); - for (int idx = tid + valid_k_rows * kSharedStrideVec; - idx < kGroupedVerifyBlockN * kSharedStrideVec; - idx += kGroupedVerifyThreads) { - reinterpret_cast(shared_kv)[idx] = make_uint4(0, 0, 0, 0); - } - __syncthreads(); - - grouped_verify_qk(shared_q, shared_kv, shared_scores, qk_scale); - __syncthreads(); - -#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 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] - : kXQANegInf; - const float tile_max_lane = warp_reduce_max(score); - const float tile_max = __shfl_sync(0xffffffffu, tile_max_lane, 0); - const float probability = - visible ? __expf(fmaxf(score - tile_max, -80.0f)) : 0.0f; - const float tile_sum_lane = warp_reduce_sum(probability); - const float tile_sum = __shfl_sync(0xffffffffu, tile_sum_lane, 0); - if (lane_id == 0 && tile_sum > 0.0f) { - const float old_max = smem.row_max[row]; - const float old_sum = smem.row_sum[row]; - const float new_max = fmaxf(old_max, tile_max); - smem.row_sum[row] = - old_sum * __expf(fmaxf(old_max - new_max, -80.0f)) + - tile_sum * __expf(fmaxf(tile_max - new_max, -80.0f)); - smem.row_max[row] = new_max; - } - } - __syncthreads(); - } - } - - // Recompute QK for the conservative path, or consume it once while updating - // the online state for the fused path. Both form half P and use identical - // Volta WMMA P @ V arithmetic. - 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); - for (int idx = tid + valid_k_rows * kSharedStrideVec; - idx < kGroupedVerifyBlockN * kSharedStrideVec; - idx += kGroupedVerifyThreads) { - reinterpret_cast(shared_kv)[idx] = make_uint4(0, 0, 0, 0); - } - __syncthreads(); - - grouped_verify_qk(shared_q, shared_kv, shared_scores, qk_scale); - __syncthreads(); - - if constexpr (TWO_PASS) { - for (int idx = tid; idx < kGroupedVerifyRows * kGroupedVerifyBlockN; - 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 head_idx = head_start + local_head; - const int kv_idx = tile_start + col; - const bool visible = - token_idx < query_len && head_idx < kGroupedVerifyHeads && - 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); - } - __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 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] - : kXQANegInf; - const float tile_max_lane = warp_reduce_max(score); - const float tile_max = __shfl_sync(0xffffffffu, tile_max_lane, 0); - const float old_max = smem.row_max[row]; - const float new_max = fmaxf(old_max, tile_max); - const float probability = - visible ? __expf(fmaxf(score - new_max, -80.0f)) : 0.0f; - const float tile_sum_lane = warp_reduce_sum(probability); - 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] = - __float2half_rn(probability); - if (lane_id == 0) { - if (tile_sum > 0.0f) { - smem.row_sum[row] = smem.row_sum[row] * exp_diff + tile_sum; - smem.row_max[row] = new_max; - } - smem.row_scale[row] = exp_diff; - } - } - __syncthreads(); -#pragma unroll - for (int fragment_idx = 0; - fragment_idx < kGroupedVerifyOutputTilesPerWarp; ++fragment_idx) { - const int output_tile = warp_id + fragment_idx * kGroupedVerifyWarps; - const int m_tile = output_tile / (kGroupedVerifyHeadDim / 16); - grouped_verify_scale_output_fragment(output_fragments[fragment_idx], - smem.row_scale, m_tile * 16); - } - } - - 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); - for (int idx = tid + valid_k_rows * kSharedStrideVec; - idx < kGroupedVerifyBlockN * kSharedStrideVec; - idx += kGroupedVerifyThreads) { - reinterpret_cast(shared_kv)[idx] = make_uint4(0, 0, 0, 0); - } - __syncthreads(); - -#pragma unroll - for (int fragment_idx = 0; fragment_idx < kGroupedVerifyOutputTilesPerWarp; - ++fragment_idx) { - const int output_tile = warp_id + fragment_idx * kGroupedVerifyWarps; - const int m_tile = output_tile / (kGroupedVerifyHeadDim / 16); - const int d_tile = output_tile % (kGroupedVerifyHeadDim / 16); - volta::fragment - probability_fragment; - volta::fragment - value_fragment; -#pragma unroll - 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); - volta::load_matrix_sync( - value_fragment, - shared_kv + k_offset * kGroupedVerifyHeadDim + d_tile * 16, - kGroupedVerifyHeadDim); - volta::mma_sync(output_fragments[fragment_idx], probability_fragment, - value_fragment, output_fragments[fragment_idx]); - } - } - __syncthreads(); - } - - // The compute buffers are dead. Reuse their storage for the dense FP32 - // partial output, then normalize and write only real query/head rows. - __syncthreads(); - float* shared_output = smem.storage.output; -#pragma unroll - for (int fragment_idx = 0; fragment_idx < kGroupedVerifyOutputTilesPerWarp; - ++fragment_idx) { - const int output_tile = warp_id + fragment_idx * kGroupedVerifyWarps; - const int m_tile = output_tile / (kGroupedVerifyHeadDim / 16); - const int d_tile = output_tile % (kGroupedVerifyHeadDim / 16); - volta::store_matrix_sync( - shared_output + m_tile * 16 * kGroupedVerifyHeadDim + d_tile * 16, - output_fragments[fragment_idx], kGroupedVerifyHeadDim, - volta::mem_row_major); - } - __syncthreads(); - - for (int idx = tid; idx < kGroupedVerifyRows * kGroupedVerifyHeadDim; - 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 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) * - kGroupedVerifyHeads + - head_idx) * - kGroupedVerifyHeadDim + - d); - partial_out[output_idx] = __float2half_rn(shared_output[idx] * scale); - } - } - if (tid < kGroupedVerifyRows) { - const int token_idx = tid / kGroupedVerifyHeadsPerCta; - const int local_head = tid % kGroupedVerifyHeadsPerCta; - 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) * - kGroupedVerifyHeads + - head_idx; - partial_lse[lse_idx] = - sum > 0.0f ? smem.row_max[tid] + logf(sum) : kXQANegInf; - } - } -} - -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) { - 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]; - __shared__ float final_max; - __shared__ float final_inv_sum; - - if (threadIdx.x < kGroupedVerifySplits) { - const int64_t lse_idx = - (static_cast(threadIdx.x) * kGroupedVerifyMaxQ + token_idx) * - kGroupedVerifyHeads + - head_idx; - split_lse[threadIdx.x] = - threadIdx.x < active_splits ? partial_lse[lse_idx] : kXQANegInf; - } - __syncthreads(); - - if (threadIdx.x == 0) { - float max_lse = kXQANegInf; - for (int split = 0; split < active_splits; ++split) { - max_lse = fmaxf(max_lse, split_lse[split]); - } - float sum = 0.0f; - for (int split = 0; split < active_splits; ++split) { - if (split_lse[split] > -1.0e20f) { - sum += __expf(fmaxf(split_lse[split] - max_lse, -80.0f)); - } - } - final_max = max_lse; - final_inv_sum = sum > 0.0f ? 1.0f / sum : 0.0f; - } - __syncthreads(); - - for (int d = threadIdx.x; d < kGroupedVerifyHeadDim; - d += kGroupedVerifyThreads) { - float accumulator = 0.0f; - for (int split = 0; split < active_splits; ++split) { - if (split_lse[split] > -1.0e20f) { - 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) * - kGroupedVerifyHeads + - head_idx) * - kGroupedVerifyHeadDim + - d); - accumulator = - fmaf(weight, __half2float(partial_out[partial_idx]), accumulator); - } - } - out[(token_idx * kGroupedVerifyHeads + head_idx) * kGroupedVerifyHeadDim + - d] = __float2half_rn(accumulator); - } -} - -template +template __global__ void flash_attention_decode_reduce_kernel( const __half* __restrict__ tmp_out, const float* __restrict__ max_logits, const float* __restrict__ exp_sums, const int* __restrict__ seq_lens, @@ -2283,9 +1306,7 @@ __global__ void flash_attention_decode_reduce_kernel( 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 route_partition_size_begin, - const int route_seq_len_begin, const int route_seq_len_end, - const int route_seq_len_final) { + const int64_t out_stride1) { const int batch_idx = blockIdx.x; const int head_idx = blockIdx.y; @@ -2294,25 +1315,8 @@ __global__ void flash_attention_decode_reduce_kernel( } const int seq_len = seq_lens[batch_idx]; - if (!xqa_seq_len_route_active(seq_len, route_seq_len_begin, - route_seq_len_end, - route_seq_len_final)) { - return; - } - int partition_size; - if constexpr (PARTITION_SIZE == -1) { - static_assert(SEQ_LEN_ROUTE == kXQARouteWaveLongSeqLens, - "Runtime wave partitions require the wave-long route"); - partition_size = seq_len < route_seq_len_end - ? 512 - : (seq_len < route_seq_len_final ? 896 : 1664); - } else if constexpr (PARTITION_SIZE == 0) { - partition_size = seq_len < route_partition_size_begin ? 64 : 256; - } else { - partition_size = PARTITION_SIZE; - } const int num_partitions = - min(max_num_partitions, (seq_len + partition_size - 1) / partition_size); + min(max_num_partitions, (seq_len + PARTITION_SIZE - 1) / PARTITION_SIZE); (void)active_num_partitions; if (seq_len <= 0 || num_partitions <= 0) { @@ -3041,8 +2045,7 @@ void launch_flash_attention_decode_paged( const int window_size_left, const int window_size_right, cudaStream_t stream, const int route_seq_len_begin = 0, const int route_seq_len_end = 0, const int route_seq_len_final = 0, - const bool launch_reduce = true, const int* anchor_lens = nullptr, - const int anchored_window = 0) { + const bool launch_reduce = true) { const int batch_size = q.size(0); const int num_heads_q = q.size(1); const int num_heads_kv = k_cache.size(2); @@ -3056,41 +2059,22 @@ void launch_flash_attention_decode_paged( const size_t reduce_shared_mem = static_cast(2 * max_num_partitions) * sizeof(float); - const bool use_anchored = anchor_lens != nullptr && anchored_window > 0; - // Second kernel version: the anchored decode-window mask is a separate - // template instantiation, generated only for the fp16-KV configuration; - // the non-anchored instantiations stay untouched. - const auto launch_partition = [&](auto anchored_tag) { - constexpr bool kAnchored = decltype(anchored_tag)::value; - flash_attention_decode_partition_kernel - <<>>( - reinterpret_cast(q.data_ptr()), - k_cache.data_ptr(), v_cache.data_ptr(), - reinterpret_cast<__half*>(tmp_out.data_ptr()), - max_logits.data_ptr(), exp_sums.data_ptr(), - block_table.data_ptr(), seq_lens.data_ptr(), - active_num_partitions.data_ptr(), batch_size, max_num_blocks, - max_num_partitions, num_heads_q, num_heads_kv, block_size, - q.stride(0), q.stride(1), tmp_out.stride(0), tmp_out.stride(1), - tmp_out.stride(2), max_logits.stride(0), max_logits.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), - softmax_scale, k_scale, v_scale, window_size_left, - window_size_right, route_seq_len_begin, route_seq_len_end, - route_seq_len_final, anchor_lens, anchored_window); - }; - if constexpr (KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP16) { - if (use_anchored) { - launch_partition(std::true_type{}); - } else { - launch_partition(std::false_type{}); - } - } else { - TORCH_CHECK(!use_anchored, - "anchored decode window requires an fp16 KV cache"); - launch_partition(std::false_type{}); - } + flash_attention_decode_partition_kernel + <<>>( + reinterpret_cast(q.data_ptr()), + k_cache.data_ptr(), v_cache.data_ptr(), + reinterpret_cast<__half*>(tmp_out.data_ptr()), + max_logits.data_ptr(), exp_sums.data_ptr(), + block_table.data_ptr(), seq_lens.data_ptr(), + active_num_partitions.data_ptr(), batch_size, max_num_blocks, + max_num_partitions, num_heads_q, num_heads_kv, block_size, + q.stride(0), q.stride(1), tmp_out.stride(0), tmp_out.stride(1), + tmp_out.stride(2), max_logits.stride(0), max_logits.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), + softmax_scale, k_scale, v_scale, window_size_left, window_size_right, + route_seq_len_begin, route_seq_len_end, route_seq_len_final); if (!launch_reduce) { return; @@ -3104,7 +2088,7 @@ void launch_flash_attention_decode_paged( 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), 0, 0, 0, 0); + out.stride(0), out.stride(1)); } template @@ -3166,8 +2150,7 @@ template + bool PARTITION_PAGE_IDS = false, bool E5M2_PAIR_LOAD = false> void launch_flash_attention_decode_paged_xqa_tc_256_wide( const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, at::Tensor& out, const at::Tensor& block_table, const at::Tensor& seq_lens, @@ -3178,9 +2161,6 @@ void launch_flash_attention_decode_paged_xqa_tc_256_wide( cudaStream_t stream, const int route_seq_len_begin = 0, const int route_seq_len_end = 0, const int route_seq_len_final = 0, const bool launch_reduce = true) { - static_assert(!E4M3_SHARED_LUT || - KV_DTYPE_OVERRIDE == flash_v100::KV_CACHE_DTYPE_FP8_E4M3, - "The shared conversion LUT requires an E4M3 specialization"); const int batch_size = q.size(0); const int num_heads_q = q.size(1); const int num_heads_kv = k_cache.size(2); @@ -3190,8 +2170,7 @@ void launch_flash_attention_decode_paged_xqa_tc_256_wide( QK_SW_PIPELINE, XQATCQKPipelineSmem256WideLayout, XQATCSmem256WideLayout>; - const size_t shared_mem = - sizeof(SmemLayout) + (E4M3_SHARED_LUT ? 256 * sizeof(uint16_t) : 0); + const size_t shared_mem = sizeof(SmemLayout); #define LAUNCH_XQA_PARTITION(KV_DTYPE) \ do { \ auto partition_kernel = \ @@ -3199,7 +2178,7 @@ void launch_flash_attention_decode_paged_xqa_tc_256_wide( PARTITION_SIZE, GROUP_SIZE, PADDED_SMEM, NUM_THREADS, \ MIN_BLOCKS_PER_SM, BLOCK_SIZE, CONTIGUOUS_HKV1_LAYOUT, \ ALIGNED_PADDED_SMEM, KV_DTYPE, SEQ_LEN_ROUTE, QK_SW_PIPELINE, \ - PARTITION_PAGE_IDS, FP8_PAIR_LOAD, E4M3_SHARED_LUT>; \ + PARTITION_PAGE_IDS, E5M2_PAIR_LOAD>; \ cudaFuncSetAttribute(partition_kernel, \ cudaFuncAttributeMaxDynamicSharedMemorySize, \ shared_mem); \ @@ -3212,7 +2191,7 @@ void launch_flash_attention_decode_paged_xqa_tc_256_wide( PARTITION_SIZE, GROUP_SIZE, PADDED_SMEM, NUM_THREADS, \ MIN_BLOCKS_PER_SM, BLOCK_SIZE, CONTIGUOUS_HKV1_LAYOUT, \ ALIGNED_PADDED_SMEM, KV_DTYPE, SEQ_LEN_ROUTE, QK_SW_PIPELINE, \ - PARTITION_PAGE_IDS, FP8_PAIR_LOAD, E4M3_SHARED_LUT> \ + PARTITION_PAGE_IDS, E5M2_PAIR_LOAD> \ <<>>( \ reinterpret_cast(q.data_ptr()), \ k_cache.data_ptr(), v_cache.data_ptr(), \ @@ -3229,16 +2208,9 @@ void launch_flash_attention_decode_paged_xqa_tc_256_wide( route_seq_len_end, route_seq_len_final); \ } while (0) - if constexpr (KV_DTYPE_OVERRIDE == flash_v100::KV_CACHE_DTYPE_FP8_E4M3) { - static_assert(!(FP8_PAIR_LOAD && E4M3_SHARED_LUT), - "Paired E4M3 conversion and the shared LUT are exclusive"); - TORCH_CHECK(k_cache.scalar_type() == at::kByte, - "E4M3 XQA requires uint8 KV cache"); - LAUNCH_XQA_PARTITION(flash_v100::KV_CACHE_DTYPE_FP8_E4M3); - } else if constexpr (FP8_PAIR_LOAD) { - static_assert(!E4M3_SHARED_LUT, "The E4M3 conversion LUT requires E4M3 KV"); + if constexpr (E5M2_PAIR_LOAD) { TORCH_CHECK(k_cache.scalar_type() == at::kByte, - "Paired XQA loads require uint8 FP8 KV cache"); + "Paired XQA loads require FP8 E5M2 KV cache"); LAUNCH_XQA_PARTITION(flash_v100::KV_CACHE_DTYPE_FP8_E5M2); } else if (k_cache.scalar_type() == at::kByte) { LAUNCH_XQA_PARTITION(flash_v100::KV_CACHE_DTYPE_FP8_E5M2); @@ -3252,22 +2224,16 @@ void launch_flash_attention_decode_paged_xqa_tc_256_wide( return; } - if constexpr (PARTITION_SIZE == -1) { - TORCH_CHECK(!use_split_reduce, - "Runtime wave partitions do not support split reduction"); - } if (use_split_reduce) { - if constexpr (PARTITION_SIZE != -1) { - launch_flash_attention_decode_xqa_split_reduce( - out, seq_lens, tmp_out, max_logits, exp_sums, launch_num_partitions, - split_reduce_dim_tile, stream); - } + launch_flash_attention_decode_xqa_split_reduce( + out, seq_lens, tmp_out, max_logits, exp_sums, launch_num_partitions, + split_reduce_dim_tile, stream); } else { const dim3 reduce_grid(batch_size, num_heads_q, 1); const dim3 block(kThreadsPerBlock); const size_t reduce_shared_mem = static_cast(2 * launch_num_partitions) * sizeof(float); - flash_attention_decode_reduce_kernel<256, PARTITION_SIZE, SEQ_LEN_ROUTE> + flash_attention_decode_reduce_kernel<256, PARTITION_SIZE> <<>>( reinterpret_cast(tmp_out.data_ptr()), max_logits.data_ptr(), exp_sums.data_ptr(), @@ -3275,34 +2241,11 @@ void launch_flash_attention_decode_paged_xqa_tc_256_wide( reinterpret_cast<__half*>(out.data_ptr()), batch_size, launch_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), - route_seq_len_end > 0 ? route_seq_len_end : route_seq_len_begin, - route_seq_len_begin, route_seq_len_end, route_seq_len_final); + max_logits.stride(1), out.stride(0), out.stride(1)); } C10_CUDA_KERNEL_LAUNCH_CHECK(); } -template -void launch_flash_attention_decode_paged_xqa_e4m3_g6_page1568( - const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_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& active_num_partitions, const float softmax_scale, - const float k_scale, const float v_scale, const int launch_num_partitions, - cudaStream_t stream, const int route_seq_len_begin = 0, - const int route_seq_len_end = 0, const int route_seq_len_final = 0, - const bool launch_reduce = true) { - launch_flash_attention_decode_paged_xqa_tc_256_wide< - PARTITION_SIZE, 6, true, kXQATCG6DualCtaThreads, 2, 1568, - FIXED_INTERLEAVED_HKV1_LAYOUT, false, SEQ_LEN_ROUTE, false, true, false, - flash_v100::KV_CACHE_DTYPE_FP8_E4M3, true>( - q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, max_logits, - exp_sums, active_num_partitions, softmax_scale, k_scale, v_scale, - launch_num_partitions, false, 8, stream, route_seq_len_begin, - route_seq_len_end, route_seq_len_final, launch_reduce); -} - void launch_flash_attention_decode_paged_xqa_tc_256_staged( const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, at::Tensor& out, const at::Tensor& block_table, const at::Tensor& seq_lens, @@ -3403,7 +2346,7 @@ void launch_flash_attention_decode_paged_xqa_tc_256_staged( reinterpret_cast<__half*>(out.data_ptr()), batch_size, launch_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), 0, 0, 0, 0); + max_logits.stride(1), out.stride(0), out.stride(1)); } C10_CUDA_KERNEL_LAUNCH_CHECK(); } @@ -3437,158 +2380,6 @@ void launch_flash_attention_decode_qk_scores( } // namespace -at::Tensor flash_attention_grouped_verify_paged( - const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, - std::optional& out_, const at::Tensor& block_table, - const at::Tensor& seq_lens, at::Tensor& partial_out, - at::Tensor& partial_lse, const float softmax_scale, - const std::string& kv_cache_dtype, const float k_scale, const float v_scale, - const bool one_pass) { - TORCH_CHECK(q.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda(), - "grouped verify q/K/V must be CUDA tensors"); - TORCH_CHECK(block_table.is_cuda() && seq_lens.is_cuda(), - "grouped verify metadata must be CUDA tensors"); - TORCH_CHECK(partial_out.is_cuda() && partial_lse.is_cuda(), - "grouped verify workspaces must be CUDA tensors"); - TORCH_CHECK(q.dtype() == torch::kFloat16, "grouped verify q must be fp16"); - TORCH_CHECK(kv_cache_dtype == "fp8_e5m2", - "grouped verify prototype supports fp8_e5m2 KV only"); - TORCH_CHECK( - k_cache.dtype() == torch::kUInt8 && v_cache.dtype() == torch::kUInt8, - "grouped verify E5M2 cache must use uint8 storage"); - TORCH_CHECK( - block_table.dtype() == torch::kInt32 && seq_lens.dtype() == torch::kInt32, - "grouped verify block_table/seq_lens must be int32"); - TORCH_CHECK(partial_out.dtype() == torch::kFloat16 && - 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(1) == kGroupedVerifyHeads && - q.size(2) == kGroupedVerifyHeadDim, - "grouped verify q must have shape [1..8, 6, 256]"); - 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 && - k_cache.size(3) == kGroupedVerifyHeadDim, - "grouped verify KV must have shape [blocks, page, 1, 256]"); - TORCH_CHECK(block_table.dim() == 2 && block_table.size(0) == 1, - "grouped verify block_table must have shape [1, blocks]"); - TORCH_CHECK(seq_lens.dim() == 1 && seq_lens.size(0) >= 1, - "grouped verify seq_lens must cover one sequence"); - TORCH_CHECK(q.is_contiguous(), - "grouped verify q must be contiguous [M, H, D]"); - TORCH_CHECK(block_table.is_contiguous() && seq_lens.is_contiguous(), - "grouped verify metadata must be contiguous"); - TORCH_CHECK(k_cache.stride(-1) == 1 && v_cache.stride(-1) == 1, - "grouped verify KV head dimension must be contiguous"); - 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, - 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]"); - TORCH_CHECK(k_scale > 0.0f && v_scale > 0.0f, - "grouped verify E5M2 K/V scales must be positive"); - - at::Tensor out = out_.has_value() ? out_.value() : torch::empty_like(q); - TORCH_CHECK(out.is_cuda() && out.device() == q.device() && - out.dtype() == torch::kFloat16 && out.sizes() == q.sizes() && - out.is_contiguous(), - "grouped verify out must be contiguous fp16 and q-shaped"); - TORCH_CHECK( - q.device() == k_cache.device() && q.device() == v_cache.device() && - q.device() == block_table.device() && - q.device() == seq_lens.device() && - q.device() == partial_out.device() && - q.device() == partial_lse.device() && q.device() == out.device(), - "all grouped verify tensors must be on the same device"); - - c10::cuda::CUDAGuard device_guard(q.device()); - const auto* properties = at::cuda::getCurrentDeviceProperties(); - TORCH_CHECK(properties->major == 7 && properties->minor == 0, - "grouped verify prototype supports SM70 only"); - cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); - - const dim3 partial_grid(1, kGroupedVerifySplits, 1); - const size_t partial_shared_mem = sizeof(GroupedVerifySmem); -#define LAUNCH_GROUPED_VERIFY_PARTIAL(TWO_PASS, PAGE_SIZE, SINGLE_QUERY) \ - do { \ - auto partial_kernel = \ - (void*)flash_attention_grouped_verify_e5m2_partial_kernel< \ - TWO_PASS, PAGE_SIZE, SINGLE_QUERY>; \ - const cudaError_t smem_status = cudaFuncSetAttribute( \ - partial_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, \ - partial_shared_mem); \ - TORCH_CHECK(smem_status == cudaSuccess, \ - "Failed to set packed grouped-verifier shared memory: ", \ - cudaGetErrorString(smem_status)); \ - const cudaError_t carveout_status = cudaFuncSetAttribute( \ - partial_kernel, cudaFuncAttributePreferredSharedMemoryCarveout, 100); \ - TORCH_CHECK(carveout_status == cudaSuccess, \ - "Failed to set packed grouped-verifier shared carveout: ", \ - cudaGetErrorString(carveout_status)); \ - flash_attention_grouped_verify_e5m2_partial_kernel \ - <<>>( \ - reinterpret_cast(q.data_ptr()), k_cache.data_ptr(), \ - v_cache.data_ptr(), block_table.data_ptr(), \ - seq_lens.data_ptr(), \ - reinterpret_cast<__half*>(partial_out.data_ptr()), \ - partial_lse.data_ptr(), static_cast(q.size(0)), \ - static_cast(block_table.size(1)), \ - static_cast(k_cache.size(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), softmax_scale * k_scale, \ - 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); \ - } \ - } while (0) - - const bool single_query = q.size(0) == 1; - if (one_pass && single_query) { - DISPATCH_GROUPED_VERIFY_PARTIAL(false, true); - } else if (one_pass) { - DISPATCH_GROUPED_VERIFY_PARTIAL(false, false); - } else if (single_query) { - DISPATCH_GROUPED_VERIFY_PARTIAL(true, true); - } else { - DISPATCH_GROUPED_VERIFY_PARTIAL(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()), \ - static_cast(q.size(0))) - if (single_query) { - LAUNCH_GROUPED_VERIFY_COMBINE(true); - } else { - LAUNCH_GROUPED_VERIFY_COMBINE(false); - } -#undef LAUNCH_GROUPED_VERIFY_COMBINE - C10_CUDA_KERNEL_LAUNCH_CHECK(); - return out; -} - at::Tensor flash_attention_decode_paged( const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, std::optional& out_, const at::Tensor& block_table, @@ -3597,8 +2388,7 @@ at::Tensor flash_attention_decode_paged( const float softmax_scale, const int partition_size, const int launch_num_partitions, const std::string& kv_cache_dtype, const float k_scale, const float v_scale, const int window_size_left, - const int window_size_right, const std::optional& anchor_lens, - const int64_t anchored_window) { + const int window_size_right) { TORCH_CHECK(q.is_cuda(), "q must be on CUDA"); TORCH_CHECK(k_cache.is_cuda() && v_cache.is_cuda(), "k/v cache must be on CUDA"); @@ -3684,30 +2474,6 @@ at::Tensor flash_attention_decode_paged( TORCH_CHECK(window_size_left >= -1 && window_size_right >= -1, "window sizes must be >= -1"); - TORCH_CHECK(anchored_window >= 0 && anchored_window <= INT_MAX, - "anchored_window must fit in a non-negative int32"); - TORCH_CHECK((anchored_window > 0) == anchor_lens.has_value(), - "anchor_lens and a positive anchored_window must be provided " - "together"); - const bool use_anchored = anchor_lens.has_value(); - const int* anchor_lens_ptr = nullptr; - if (use_anchored) { - const at::Tensor& anchors = anchor_lens.value(); - TORCH_CHECK(anchors.is_cuda(), "anchor_lens must be on CUDA"); - TORCH_CHECK(anchors.device() == q.device(), - "anchor_lens must be on the q device"); - TORCH_CHECK(anchors.dtype() == torch::kInt32, "anchor_lens must be int32"); - TORCH_CHECK(anchors.dim() == 1 && anchors.size(0) == batch_size, - "anchor_lens must have shape [B]"); - TORCH_CHECK(anchors.is_contiguous(), "anchor_lens must be contiguous"); - TORCH_CHECK(kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP16, - "anchored decode window requires an fp16 KV cache"); - TORCH_CHECK(window_size_left == -1 && window_size_right == -1, - "anchored decode window cannot be combined with " - "sliding-window attention"); - anchor_lens_ptr = anchors.data_ptr(); - } - at::Tensor out = out_.has_value() ? out_.value() : torch::empty_like(q); TORCH_CHECK(out.is_cuda(), "out must be on CUDA"); TORCH_CHECK(out.dtype() == torch::kFloat16, "out must be fp16"); @@ -3717,12 +2483,11 @@ at::Tensor flash_attention_decode_paged( auto stream = at::cuda::getCurrentCUDAStream().stream(); c10::cuda::CUDAGuard device_guard(q.device()); -#define LAUNCH_TYPED(HDIM, PARTITION, KV_DTYPE_CODE) \ - launch_flash_attention_decode_paged( \ - q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, max_logits, \ - exp_sums, active_num_partitions, softmax_scale, launch_num_partitions, \ - k_scale, v_scale, window_size_left, window_size_right, stream, 0, 0, 0, \ - true, anchor_lens_ptr, static_cast(anchored_window)) +#define LAUNCH_TYPED(HDIM, PARTITION, KV_DTYPE_CODE) \ + launch_flash_attention_decode_paged( \ + q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, max_logits, \ + exp_sums, active_num_partitions, softmax_scale, launch_num_partitions, \ + k_scale, v_scale, window_size_left, window_size_right, stream) #define LAUNCH_BY_KV_DTYPE(HDIM, PARTITION) \ do { \ @@ -3798,7 +2563,7 @@ at::Tensor flash_attention_decode_paged_xqa( const float softmax_scale, const int partition_size, const int launch_num_partitions, const std::string& kv_cache_dtype, const float k_scale, const float v_scale, const int window_size_left, - const int window_size_right, const int batch_context_max_seq_len) { + const int window_size_right) { TORCH_CHECK(q.is_cuda(), "q must be on CUDA"); TORCH_CHECK(k_cache.is_cuda() && v_cache.is_cuda(), "k_cache and v_cache must be on CUDA"); @@ -3811,10 +2576,8 @@ at::Tensor flash_attention_decode_paged_xqa( TORCH_CHECK(q.dtype() == torch::kFloat16, "q must be fp16"); const int kv_dtype_code = kv_cache_dtype_code_from_string(kv_cache_dtype); TORCH_CHECK(kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP16 || - kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP8_E4M3 || kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP8_E5M2, - "XQA decode supports fp16, fp8_e4m3, and fp8_e5m2 KV cache " - "only"); + "XQA decode supports fp16 and fp8_e5m2 KV cache only"); if (kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP16) { TORCH_CHECK(k_cache.dtype() == torch::kFloat16 && v_cache.dtype() == torch::kFloat16, @@ -3822,7 +2585,7 @@ at::Tensor flash_attention_decode_paged_xqa( } else { TORCH_CHECK( k_cache.dtype() == torch::kUInt8 && v_cache.dtype() == torch::kUInt8, - "FP8 XQA requires uint8 K/V tensors"); + "fp8_e5m2 XQA requires uint8 K/V tensors"); } TORCH_CHECK(k_scale > 0.f && v_scale > 0.f, "XQA K/V scales must be positive"); @@ -3859,12 +2622,9 @@ at::Tensor flash_attention_decode_paged_xqa( const int q_per_kv = num_heads_q / num_heads_kv; TORCH_CHECK(q_per_kv == 4 || q_per_kv == 6 || q_per_kv == 8, "XQA decode supports q_per_kv in {4, 6, 8}, got ", q_per_kv); - TORCH_CHECK(partition_size == 64 || partition_size == 128 || - partition_size == 256 || partition_size == 512 || - partition_size == 1024 || - ((partition_size == 896 || partition_size == 1664) && - kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP8_E4M3), - "Unsupported XQA decode partition_size: ", partition_size); + TORCH_CHECK( + partition_size == 256 || partition_size == 512 || partition_size == 1024, + "Unsupported XQA decode partition_size: ", partition_size); TORCH_CHECK(launch_num_partitions > 0, "launch_num_partitions must be positive"); TORCH_CHECK(tmp_out.dtype() == torch::kFloat16, @@ -3888,345 +2648,6 @@ at::Tensor flash_attention_decode_paged_xqa( TORCH_CHECK(out.sizes() == q.sizes(), "out must have same shape as q"); TORCH_CHECK(out.stride(-1) == 1, "out last dim must be contiguous"); auto stream = at::cuda::getCurrentCUDAStream().stream(); - if (kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP8_E4M3) { - const bool e4m3_batch_allowed = - q.size(0) > 1 && q.size(0) <= 16 && xqa_e4m3_batch_enabled(); - const bool use_large_partition = - q.size(0) == 1 && - (partition_size == 512 || partition_size == 896 || - partition_size == 1024 || partition_size == 1664) && - k_cache.size(1) == 1568 && k_cache.size(2) == 1; - TORCH_CHECK((q.size(0) == 1 || e4m3_batch_allowed) && q_per_kv == 6 && - (partition_size == 64 || partition_size == 128 || - partition_size == 256 || use_large_partition), - "E4M3 XQA supports B=1, or B=2..16 when " - "VLLM_FLASH_V100_E4M3_BATCH_XQA=1; q_per_kv=6 and D=256 are " - "required. Page-1568/Hkv=1 B1 additionally supports partition " - "sizes 512, 896, 1024, and 1664"); - const bool e4m3_batch_optimized = - q.size(0) > 1 && xqa_e4m3_batch_optimized_enabled() && - k_cache.size(1) >= 256 && k_cache.size(1) % 16 == 0 && - k_cache.size(2) == 1 && k_cache.scalar_type() == at::kByte; - const XQABatchContextRoute e4m3_batch_route = - e4m3_batch_optimized && batch_context_max_seq_len > 0 && q.size(0) >= 4 - ? select_xqa_batch_context_route( - q.size(0), batch_context_max_seq_len, partition_size) - : XQABatchContextRoute::kDisabled; - const bool e4m3_dual_cta = - e4m3_batch_route == XQABatchContextRoute::kDualCta || - e4m3_batch_route == XQABatchContextRoute::kDualCtaSplit; - const bool e4m3_split_reduce = - e4m3_batch_route == XQABatchContextRoute::kDualCtaSplit; - constexpr int64_t kPage800HeadDim256Elements = 800 * 256; - const bool e4m3_page800_contiguous = - xqa_e4m3_page800_fastpath_enabled() && k_cache.size(1) == 800 && - k_cache.size(2) == 1 && k_cache.size(3) == 256 && - k_cache.stride(0) == 2 * kPage800HeadDim256Elements && - k_cache.stride(1) == 256 && k_cache.stride(2) == 256 && - k_cache.stride(3) == 1 && - v_cache.stride(0) == 2 * kPage800HeadDim256Elements && - v_cache.stride(1) == 256 && v_cache.stride(2) == 256 && - v_cache.stride(3) == 1; - trace_xqa_batch_context_route(q.size(0), batch_context_max_seq_len, - partition_size, k_cache.size(1), - e4m3_batch_route); - trace_xqa_e4m3_page800_fastpath(q.size(0), partition_size, - e4m3_dual_cta && e4m3_page800_contiguous); - const int split_reduce_dim_tile = xqa_split_reduce_dim_tile(); - -#define LAUNCH_E4M3_DUAL_CTA(PARTITION, PAGE_SIZE, CONTIGUOUS) \ - launch_flash_attention_decode_paged_xqa_tc_256_wide< \ - PARTITION, 6, true, kXQATCG6DualCtaThreads, 2, PAGE_SIZE, CONTIGUOUS, \ - false, kXQARouteAllSeqLens, false, true, true, \ - flash_v100::KV_CACHE_DTYPE_FP8_E4M3>( \ - q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, max_logits, \ - exp_sums, active_num_partitions, softmax_scale, k_scale, v_scale, \ - launch_num_partitions, e4m3_split_reduce, split_reduce_dim_tile, stream) - -#define LAUNCH_E4M3_BATCH_XQA(PARTITION) \ - do { \ - if (e4m3_dual_cta) { \ - if (e4m3_page800_contiguous) { \ - LAUNCH_E4M3_DUAL_CTA(PARTITION, 800, true); \ - } else { \ - LAUNCH_E4M3_DUAL_CTA(PARTITION, 0, false); \ - } \ - } else if (e4m3_batch_optimized) { \ - launch_flash_attention_decode_paged_xqa_tc_256_wide< \ - PARTITION, 6, true, kXQATC256WideThreads, 1, 0, false, false, \ - kXQARouteAllSeqLens, false, true, true, \ - flash_v100::KV_CACHE_DTYPE_FP8_E4M3>( \ - q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, \ - max_logits, exp_sums, active_num_partitions, softmax_scale, k_scale, \ - v_scale, launch_num_partitions, false, split_reduce_dim_tile, \ - stream); \ - } else { \ - launch_flash_attention_decode_paged_xqa_tc_256_wide< \ - PARTITION, 6, true, kXQATC256WideThreads, 1, 0, false, false, \ - kXQARouteAllSeqLens, false, false, false, \ - flash_v100::KV_CACHE_DTYPE_FP8_E4M3>( \ - q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, \ - max_logits, exp_sums, active_num_partitions, softmax_scale, k_scale, \ - v_scale, launch_num_partitions, false, 8, stream); \ - } \ - } while (0) - - const bool use_p64_p256_auto = - q.size(0) == 1 && partition_size == 64 && k_cache.size(1) == 1568 && - k_cache.size(2) == 1 && !decode_partition_size_overridden() && - xqa_e4m3_g6_p64_p256_auto_enabled(); - constexpr int64_t kPage1568HeadDim256Elements = 1568 * 256; - const bool fixed_interleaved_hkv1_layout = - k_cache.size(3) == 256 && - k_cache.stride(0) == 2 * kPage1568HeadDim256Elements && - k_cache.stride(1) == 256 && k_cache.stride(2) == 256 && - k_cache.stride(3) == 1 && - v_cache.stride(0) == 2 * kPage1568HeadDim256Elements && - v_cache.stride(1) == 256 && v_cache.stride(2) == 256 && - v_cache.stride(3) == 1; - if (use_p64_p256_auto) { - const int p256_begin = xqa_e4m3_g6_p256_begin(); - const int p512_begin = xqa_e4m3_g6_p512_begin(); - const bool use_wave_partitions = xqa_e4m3_g6_wave_partitions_enabled() && - batch_context_max_seq_len >= p512_begin; - if (use_wave_partitions) { - const int p896_begin = xqa_e4m3_g6_p896_begin(); - const int p1664_begin = xqa_e4m3_g6_p1664_begin(); - TORCH_CHECK(p256_begin < p512_begin && p512_begin < p896_begin && - p896_begin < p1664_begin, - "Invalid E4M3 wave-partition thresholds: ", p256_begin, - ", ", p512_begin, ", ", p896_begin, ", ", p1664_begin); - - const int p64_launch_num_partitions = std::max( - 1, std::min(launch_num_partitions, (p256_begin + 62) / 64)); - const int p256_launch_num_partitions = (launch_num_partitions + 3) / 4; - const int p512_launch_num_partitions = (launch_num_partitions + 7) / 8; - const int p896_launch_num_partitions = - (launch_num_partitions + 13) / 14; - const int p1664_launch_num_partitions = - (launch_num_partitions + 25) / 26; - const int p256_mid_launch_num_partitions = std::max( - 1, std::min(p256_launch_num_partitions, (p512_begin + 254) / 256)); - const int short_launch_num_partitions = - std::max(p64_launch_num_partitions, p256_mid_launch_num_partitions); - const int p512_route_launch_num_partitions = std::max( - 1, std::min(p512_launch_num_partitions, (p896_begin + 510) / 512)); - const int p896_route_launch_num_partitions = std::max( - 1, std::min(p896_launch_num_partitions, (p1664_begin + 894) / 896)); - const int wave_long_launch_num_partitions = std::max( - {p512_route_launch_num_partitions, p896_route_launch_num_partitions, - p1664_launch_num_partitions}); - const bool use_merged_wave_launch = - xqa_e4m3_g6_merged_wave_launch_enabled(); - - launch_flash_attention_decode_paged_xqa_tc_256_wide< - 0, 6, true, kXQATC256WideThreads, 1, 0, false, false, - kXQARouteShortSeqLens, false, false, false, - flash_v100::KV_CACHE_DTYPE_FP8_E4M3>( - q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, - max_logits, exp_sums, active_num_partitions, softmax_scale, k_scale, - v_scale, short_launch_num_partitions, false, 8, stream, p512_begin, - p256_begin, 0, false); - if (use_merged_wave_launch) { - if (fixed_interleaved_hkv1_layout) { - launch_flash_attention_decode_paged_xqa_e4m3_g6_page1568< - -1, kXQARouteWaveLongSeqLens, true>( - q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, - max_logits, exp_sums, active_num_partitions, softmax_scale, - k_scale, v_scale, wave_long_launch_num_partitions, stream, - p512_begin, p896_begin, p1664_begin, false); - } else { - launch_flash_attention_decode_paged_xqa_e4m3_g6_page1568< - -1, kXQARouteWaveLongSeqLens>( - q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, - max_logits, exp_sums, active_num_partitions, softmax_scale, - k_scale, v_scale, wave_long_launch_num_partitions, stream, - p512_begin, p896_begin, p1664_begin, false); - } - } else { - launch_flash_attention_decode_paged_xqa_e4m3_g6_page1568< - 512, kXQARouteRangeSeqLens>( - q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, - max_logits, exp_sums, active_num_partitions, softmax_scale, - k_scale, v_scale, p512_route_launch_num_partitions, stream, - p512_begin, p896_begin, 0, false); - launch_flash_attention_decode_paged_xqa_e4m3_g6_page1568< - 896, kXQARouteRangeSeqLens>( - q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, - max_logits, exp_sums, active_num_partitions, softmax_scale, - k_scale, v_scale, p896_route_launch_num_partitions, stream, - p896_begin, p1664_begin, 0, false); - launch_flash_attention_decode_paged_xqa_e4m3_g6_page1568< - 1664, kXQARouteLongSeqLens>( - q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, - max_logits, exp_sums, active_num_partitions, softmax_scale, - k_scale, v_scale, p1664_launch_num_partitions, stream, - p1664_begin, 0, 0, false); - } - - const dim3 reduce_grid(q.size(0), q.size(1), 1); - const dim3 reduce_block(kThreadsPerBlock); -#define LAUNCH_E4M3_WAVE_REDUCER(PARTITION_SIZE, SEQ_ROUTE, MAX_PARTITIONS, \ - PARTITION_BEGIN, ROUTE_BEGIN, ROUTE_END, \ - ROUTE_FINAL) \ - do { \ - const size_t reduce_shared_mem = \ - static_cast(2 * (MAX_PARTITIONS)) * sizeof(float); \ - flash_attention_decode_reduce_kernel<256, PARTITION_SIZE, SEQ_ROUTE> \ - <<>>( \ - reinterpret_cast(tmp_out.data_ptr()), \ - max_logits.data_ptr(), exp_sums.data_ptr(), \ - seq_lens.data_ptr(), active_num_partitions.data_ptr(), \ - reinterpret_cast<__half*>(out.data_ptr()), q.size(0), \ - MAX_PARTITIONS, q.size(1), 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), PARTITION_BEGIN, ROUTE_BEGIN, \ - ROUTE_END, ROUTE_FINAL); \ - } while (0) - LAUNCH_E4M3_WAVE_REDUCER(0, kXQARouteShortSeqLens, - short_launch_num_partitions, p256_begin, - p512_begin, 0, 0); - if (use_merged_wave_launch) { - LAUNCH_E4M3_WAVE_REDUCER(-1, kXQARouteWaveLongSeqLens, - wave_long_launch_num_partitions, 0, - p512_begin, p896_begin, p1664_begin); - } else { - LAUNCH_E4M3_WAVE_REDUCER(512, kXQARouteRangeSeqLens, - p512_route_launch_num_partitions, 0, - p512_begin, p896_begin, 0); - LAUNCH_E4M3_WAVE_REDUCER(896, kXQARouteRangeSeqLens, - p896_route_launch_num_partitions, 0, - p896_begin, p1664_begin, 0); - LAUNCH_E4M3_WAVE_REDUCER(1664, kXQARouteLongSeqLens, - p1664_launch_num_partitions, 0, p1664_begin, - 0, 0); - } -#undef LAUNCH_E4M3_WAVE_REDUCER - C10_CUDA_KERNEL_LAUNCH_CHECK(); - - static bool traced_wave_partitions = false; - if (xqa_e4m3_g6_p64_p256_auto_trace_enabled() && - !traced_wave_partitions) { - TORCH_WARN( - "Flash-V100 XQA E4M3 G6 wave-partition route active; " - "thresholds p256/p512/p896/p1664=", - p256_begin, "/", p512_begin, "/", p896_begin, "/", p1664_begin, - ", launch partitions short/p512/p896/p1664=", - short_launch_num_partitions, "/", - p512_route_launch_num_partitions, "/", - p896_route_launch_num_partitions, "/", - p1664_launch_num_partitions, - ", merged_long=", use_merged_wave_launch, - ", converter=shared-lut"); - traced_wave_partitions = true; - } - return out; - } - const int dual_cta_begin = - std::max(p256_begin, xqa_e4m3_g6_dual_cta_begin()); - const int p64_launch_num_partitions = - std::max(1, std::min(launch_num_partitions, (p256_begin + 62) / 64)); - const int p256_launch_num_partitions = (launch_num_partitions + 3) / 4; - const int p256_mid_launch_num_partitions = std::max( - 1, - std::min(p256_launch_num_partitions, (dual_cta_begin + 254) / 256)); - const int short_launch_num_partitions = - std::max(p64_launch_num_partitions, p256_mid_launch_num_partitions); - const int reduce_num_partitions = - std::max(short_launch_num_partitions, p256_launch_num_partitions); - - launch_flash_attention_decode_paged_xqa_tc_256_wide< - 0, 6, true, kXQATC256WideThreads, 1, 0, false, false, - kXQARouteShortSeqLens, false, false, false, - flash_v100::KV_CACHE_DTYPE_FP8_E4M3>( - q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, max_logits, - exp_sums, active_num_partitions, softmax_scale, k_scale, v_scale, - short_launch_num_partitions, false, 8, stream, dual_cta_begin, - p256_begin, 0, false); - launch_flash_attention_decode_paged_xqa_tc_256_wide< - 256, 6, true, kXQATCG6DualCtaThreads, 2, 1568, false, false, - kXQARouteLongSeqLens, false, false, false, - flash_v100::KV_CACHE_DTYPE_FP8_E4M3, true>( - q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, max_logits, - exp_sums, active_num_partitions, softmax_scale, k_scale, v_scale, - p256_launch_num_partitions, false, 8, stream, dual_cta_begin, 0, 0, - false); - - const dim3 reduce_grid(q.size(0), q.size(1), 1); - const dim3 reduce_block(kThreadsPerBlock); - const size_t reduce_shared_mem = - static_cast(2 * reduce_num_partitions) * sizeof(float); - flash_attention_decode_reduce_kernel<256, 0> - <<>>( - reinterpret_cast(tmp_out.data_ptr()), - max_logits.data_ptr(), exp_sums.data_ptr(), - seq_lens.data_ptr(), active_num_partitions.data_ptr(), - reinterpret_cast<__half*>(out.data_ptr()), q.size(0), - reduce_num_partitions, q.size(1), 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), p256_begin, 0, - 0, 0); - - static bool traced_p64_p256_auto = false; - if (xqa_e4m3_g6_p64_p256_auto_trace_enabled() && !traced_p64_p256_auto) { - TORCH_WARN( - "Flash-V100 XQA E4M3 G6 device-side p64/p256 route active; " - "p256 threshold=", - p256_begin, ", p64 launch partitions=", p64_launch_num_partitions, - ", dual-CTA threshold=", dual_cta_begin, - ", short/mid launch partitions=", short_launch_num_partitions, - ", long p256 launch partitions=", p256_launch_num_partitions, - ", routed launch partitions=", reduce_num_partitions); - traced_p64_p256_auto = true; - } - } else if (q.size(0) == 1 && partition_size == 64) { - launch_flash_attention_decode_paged_xqa_tc_256_wide< - 64, 6, true, kXQATC256WideThreads, 1, 0, false, false, - kXQARouteAllSeqLens, false, false, false, - flash_v100::KV_CACHE_DTYPE_FP8_E4M3>( - q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, max_logits, - exp_sums, active_num_partitions, softmax_scale, k_scale, v_scale, - launch_num_partitions, false, 8, stream); - } else if (q.size(0) == 1 && partition_size == 128) { - launch_flash_attention_decode_paged_xqa_tc_256_wide< - 128, 6, true, kXQATC256WideThreads, 1, 0, false, false, - kXQARouteAllSeqLens, false, false, false, - flash_v100::KV_CACHE_DTYPE_FP8_E4M3>( - q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, max_logits, - exp_sums, active_num_partitions, softmax_scale, k_scale, v_scale, - launch_num_partitions, false, 8, stream); - } else if (q.size(0) == 1 && partition_size == 512) { - launch_flash_attention_decode_paged_xqa_e4m3_g6_page1568<512>( - q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, max_logits, - exp_sums, active_num_partitions, softmax_scale, k_scale, v_scale, - launch_num_partitions, stream); - } else if (q.size(0) == 1 && partition_size == 896) { - launch_flash_attention_decode_paged_xqa_e4m3_g6_page1568<896>( - q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, max_logits, - exp_sums, active_num_partitions, softmax_scale, k_scale, v_scale, - launch_num_partitions, stream); - } else if (q.size(0) == 1 && partition_size == 1024) { - launch_flash_attention_decode_paged_xqa_e4m3_g6_page1568<1024>( - q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, max_logits, - exp_sums, active_num_partitions, softmax_scale, k_scale, v_scale, - launch_num_partitions, stream); - } else if (q.size(0) == 1 && partition_size == 1664) { - launch_flash_attention_decode_paged_xqa_e4m3_g6_page1568<1664>( - q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, max_logits, - exp_sums, active_num_partitions, softmax_scale, k_scale, v_scale, - launch_num_partitions, stream); - } else if (partition_size == 64) { - LAUNCH_E4M3_BATCH_XQA(64); - } else if (partition_size == 128) { - LAUNCH_E4M3_BATCH_XQA(128); - } else { - LAUNCH_E4M3_BATCH_XQA(256); - } -#undef LAUNCH_E4M3_BATCH_XQA -#undef LAUNCH_E4M3_DUAL_CTA - C10_CUDA_KERNEL_LAUNCH_CHECK(); - return out; - } // Default only the shape with exact direct and model-level evidence. const bool padded_smem_enabled = xqa_padded_smem_enabled(); const bool use_padded_smem = @@ -4281,20 +2702,6 @@ at::Tensor flash_attention_decode_paged_xqa( kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP8_E5M2 && !decode_partition_size_overridden() && xqa_e5m2_g6_dual_cta_enabled() && xqa_e5m2_g6_split_reduce_enabled(); - const bool batch_context_page_supported = - k_cache.size(1) == 16 || - (k_cache.size(1) >= 256 && k_cache.size(1) % 16 == 0); - const bool use_e5m2_g6_batch_context_route = - batch_context_max_seq_len > 0 && q.size(0) >= 4 && q_per_kv == 6 && - batch_context_page_supported && k_cache.size(2) == 1 && - k_cache.scalar_type() == at::kByte && - kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP8_E5M2 && - !decode_partition_size_overridden(); - const XQABatchContextRoute batch_context_route = - use_e5m2_g6_batch_context_route - ? select_xqa_batch_context_route(q.size(0), batch_context_max_seq_len, - partition_size) - : XQABatchContextRoute::kDisabled; const int e5m2_g6_dual_cta_seq_len = use_e5m2_g6_dual_cta ? at::cuda::getDeviceProperties(q.get_device())->multiProcessorCount * @@ -4322,27 +2729,13 @@ at::Tensor flash_attention_decode_paged_xqa( use_e5m2_g6_dual_cta && xqa_e5m2_partition_page_ids_enabled(); const bool use_e5m2_pair_load = use_e5m2_partition_page_ids && xqa_e5m2_pair_load_enabled(); - // The B1 E5M2 route already amortizes page-table lookups across a p256 - // partition and converts two half8 vectors from one aligned 128-bit load. - // Reuse that exact load path for the batch/context route without changing - // its partition, softmax, PV, or reduction order. Restrict the page size so - // one partition always fits the existing shared page-id capacity. - const bool use_e5m2_batch_wide_load = - use_e5m2_g6_batch_context_route && partition_size == 256 && - k_cache.size(1) >= 256 && k_cache.size(1) % 16 == 0 && - (batch_context_route == XQABatchContextRoute::kDualCta || - batch_context_route == XQABatchContextRoute::kDualCtaSplit) && - xqa_e5m2_batch_wide_load_enabled(); const bool use_g6_dual_cta = use_g6_p1024_auto || use_g6_p1024_sawtooth || use_mtp5_dual_cta || use_e5m2_g6_dual_cta || - batch_context_route == XQABatchContextRoute::kDualCta || - batch_context_route == XQABatchContextRoute::kDualCtaSplit || (xqa_g6_dual_cta_enabled() && (use_padded_smem || use_g6_dual_cta_dense)); - const bool use_split_reduce = - use_g6_p1024_auto || use_g6_p1024_sawtooth || use_e5m2_g6_dual_cta || - batch_context_route == XQABatchContextRoute::kDualCtaSplit || - (use_g6_dual_cta && xqa_split_reduce_enabled()); + const bool use_split_reduce = use_g6_p1024_auto || use_g6_p1024_sawtooth || + use_e5m2_g6_dual_cta || + (use_g6_dual_cta && xqa_split_reduce_enabled()); const bool supports_block16_index = use_g6_dual_cta && k_cache.size(1) == 16; const bool supports_block16_contiguous_layout = supports_block16_index && k_cache.size(2) == 1 && @@ -4369,17 +2762,6 @@ at::Tensor flash_attention_decode_paged_xqa( k_cache.stride(0), ",", k_cache.stride(1), ",", k_cache.stride(2), ",", k_cache.stride(3), "]"); } - trace_xqa_batch_context_route(q.size(0), batch_context_max_seq_len, - partition_size, k_cache.size(1), - batch_context_route); - static bool traced_e5m2_batch_wide_load = false; - if (xqa_batch_context_routing_trace_enabled() && use_e5m2_batch_wide_load && - !traced_e5m2_batch_wide_load) { - TORCH_WARN( - "Flash-V100 XQA batched E5M2 partition page reuse and paired " - "128-bit loads active"); - traced_e5m2_batch_wide_load = true; - } static bool traced_block16_layout = false; if (xqa_block16_layout_trace_enabled() && block16_layout_mode != 0 && !traced_block16_layout) { @@ -4700,22 +3082,12 @@ at::Tensor flash_attention_decode_paged_xqa( launch_num_partitions, use_split_reduce, split_reduce_dim_tile, stream); } else if (partition_size == 256) { - if (use_e5m2_batch_wide_load) { - launch_flash_attention_decode_paged_xqa_tc_256_wide< - 256, 6, true, kXQATCG6DualCtaThreads, 2, 0, false, false, - kXQARouteAllSeqLens, false, true, true>( - q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, - max_logits, exp_sums, active_num_partitions, softmax_scale, k_scale, - v_scale, launch_num_partitions, use_split_reduce, - split_reduce_dim_tile, stream); - } else { - launch_flash_attention_decode_paged_xqa_tc_256_wide< - 256, 6, true, kXQATCG6DualCtaThreads, 2, 0, false>( - q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, - max_logits, exp_sums, active_num_partitions, softmax_scale, k_scale, - v_scale, launch_num_partitions, use_split_reduce, - split_reduce_dim_tile, stream); - } + launch_flash_attention_decode_paged_xqa_tc_256_wide< + 256, 6, true, kXQATCG6DualCtaThreads, 2, 0, false>( + q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, max_logits, + exp_sums, active_num_partitions, softmax_scale, k_scale, v_scale, + launch_num_partitions, use_split_reduce, split_reduce_dim_tile, + stream); } else if (partition_size == 512) { launch_flash_attention_decode_paged_xqa_tc_256_wide< 512, 6, false, kXQATCG6DualCtaThreads, 2, 0, false>( diff --git a/flash-attention-v100/kernel/flash_decode_turboquant.cu b/flash-attention-v100/kernel/flash_decode_turboquant.cu index 41627bd78d..79c7630612 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,28 +77,32 @@ __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; @@ -122,9 +126,10 @@ __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; @@ -149,20 +154,32 @@ __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; @@ -174,7 +191,8 @@ __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; @@ -250,8 +268,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; @@ -266,15 +284,22 @@ __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; @@ -288,11 +313,13 @@ __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[]; @@ -315,8 +342,7 @@ __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; } @@ -324,8 +350,9 @@ __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; @@ -342,13 +369,18 @@ __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); @@ -366,35 +398,68 @@ 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(), @@ -409,8 +474,7 @@ 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, @@ -419,8 +483,7 @@ 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"); @@ -437,26 +500,23 @@ 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( - 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"); + 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), @@ -477,67 +537,65 @@ 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: @@ -559,15 +617,14 @@ 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 eaa191e2c3..909ea97db4 100644 --- a/flash-attention-v100/kernel/flash_v100_traits.cuh +++ b/flash-attention-v100/kernel/flash_v100_traits.cuh @@ -3,116 +3,137 @@ #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; -}; - -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; -} - -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; + 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; +}; -#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; +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); + } - uint4 kv_val = make_uint4(0, 0, 0, 0); + const int64_t global_row_offset = block_row_offset + n_block * kBlockN; - if (token_offset < num_tokens && vec_col < d_stride_uint4) { - const int global_token_idx = start_token_idx + token_offset; + 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_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 = + ((int64_t)block_table[virtual_page_idx]) * ((int64_t)page_stride) + + page_offset * ((int64_t)row_stride) + + col_offset; - const int64_t physical_offset_uint4 = physical_offset_halves / PER_UINT4; + return physical_offset; +} - kv_val = __ldg(&kv_vec[physical_offset_uint4 + vec_col]); +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; } - - 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 dfa54a3eae..693b46a93f 100644 --- a/flash-attention-v100/kernel/fp8_kv_utils.cuh +++ b/flash-attention-v100/kernel/fp8_kv_utils.cuh @@ -14,7 +14,9 @@ __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); @@ -36,7 +38,8 @@ __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; } @@ -49,8 +52,9 @@ __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; @@ -60,7 +64,8 @@ __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]); } @@ -69,31 +74,36 @@ __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_api.cpp b/flash-attention-v100/kernel/fused_mha_api.cpp index ec1f942353..bbd31094df 100644 --- a/flash-attention-v100/kernel/fused_mha_api.cpp +++ b/flash-attention-v100/kernel/fused_mha_api.cpp @@ -9,41 +9,59 @@ namespace py = pybind11; PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.doc() = "FlashAttention-2 implementation optimized for Volta"; - m.def("fwd", &flash_attention_forward, - "FlashAttention-2 Forward Pass (Volta)"); - m.def("qk_scores_fwd", &flash_attention_qk_scores, + m.doc() = "FlashAttention-2 implementation optimized for Volta"; + m.def("fwd", &flash_attention_forward, "FlashAttention-2 Forward Pass (Volta)"); + m.def( + "qk_scores_fwd", + &flash_attention_qk_scores, "Debug FlashAttention QK score dump before softmax (Volta)"); - m.def("bwd", &flash_attention_backward, - "FlashAttention-2 Backward Pass (Volta)"); - m.def("decode_paged_fwd", &flash_attention_decode_paged, + m.def("bwd", &flash_attention_backward, "FlashAttention-2 Backward Pass (Volta)"); + m.def( + "decode_paged_fwd", + &flash_attention_decode_paged, "FlashAttention decode over paged KV cache (Volta)"); - m.def("decode_paged_xqa_fwd", &flash_attention_decode_paged_xqa, + m.def( + "decode_paged_xqa_fwd", + &flash_attention_decode_paged_xqa, "FlashAttention XQA decode over paged KV cache (Volta)"); - m.def("decode_paged_xqa_staged_fwd", &flash_attention_decode_paged_xqa_staged, + m.def( + "decode_paged_xqa_staged_fwd", + &flash_attention_decode_paged_xqa_staged, "Staged FlashAttention XQA decode over paged KV cache (Volta)"); - m.def("grouped_verify_paged_fwd", &flash_attention_grouped_verify_paged, - "Exact grouped DFlash2 verification over paged KV cache (Volta)"); - m.def("decode_paged_wmma_fwd", &flash_attention_decode_paged_wmma, - "FlashAttention single-query decode through paged-prefill WMMA order " - "(Volta)"); - m.def("decode_qk_scores_fwd", &flash_attention_decode_qk_scores, + m.def( + "decode_paged_wmma_fwd", + &flash_attention_decode_paged_wmma, + "FlashAttention single-query decode through paged-prefill WMMA order (Volta)"); + m.def( + "decode_qk_scores_fwd", + &flash_attention_decode_qk_scores, "Debug scalar paged decode QK score dump before softmax (Volta)"); - m.def("decode_turboquant_paged_fwd", &flash_attention_turboquant_decode_paged, + m.def( + "decode_turboquant_paged_fwd", + &flash_attention_turboquant_decode_paged, "FlashAttention decode over TurboQuant paged KV cache (Volta)"); - m.def("prefill_paged_fwd", &flash_attention_prefill_paged, + m.def( + "prefill_paged_fwd", + &flash_attention_prefill_paged, "FlashAttention prefill over paged KV cache (Volta)"); - m.def("prefill_paged_d256_bm32_allp_pair_scratch_fwd", + m.def( + "prefill_paged_d256_bm32_allp_pair_scratch_fwd", &flash_attention_prefill_paged_d256_bm32_allp_pair_scratch, "Fixed causal D256 BM32 ALL_P pair-scratch paged prefill (SM70)"); - m.def("prefill_paged_d256_bm32_allp_pair_scratch_splitkv3_fwd", + m.def( + "prefill_paged_d256_bm32_allp_pair_scratch_splitkv3_fwd", &flash_attention_prefill_paged_d256_bm32_allp_pair_scratch_splitkv3, - "Fixed causal D256 BM32 ALL_P pair-scratch three-way split-KV paged " - "prefill (SM70)"); - m.def("prefill_paged_bfla_fwd", &flash_attention_prefill_paged_bfla, + "Fixed causal D256 BM32 ALL_P pair-scratch three-way split-KV paged prefill (SM70)"); + m.def( + "prefill_paged_bfla_fwd", + &flash_attention_prefill_paged_bfla, "BFLA sparse FlashAttention prefill over paged KV cache (Volta)"); - m.def("prefill_paged_splitkv_fwd", &flash_attention_prefill_paged_splitkv, + m.def( + "prefill_paged_splitkv_fwd", + &flash_attention_prefill_paged_splitkv, "FlashAttention split-KV prefill over paged KV cache (Volta)"); - m.def("fp8_e5m2_paged_kv_to_fp16", &flash_attention_fp8_e5m2_paged_kv_to_fp16, + m.def( + "fp8_e5m2_paged_kv_to_fp16", + &flash_attention_fp8_e5m2_paged_kv_to_fp16, "Expand paged FP8 E5M2 K/V into a preallocated FP16 paged workspace"); } diff --git a/flash-attention-v100/kernel/fused_mha_backward.cu b/flash-attention-v100/kernel/fused_mha_backward.cu index 1f0ed13072..f3e98fed57 100644 --- a/flash-attention-v100/kernel/fused_mha_backward.cu +++ b/flash-attention-v100/kernel/fused_mha_backward.cu @@ -14,1235 +14,1176 @@ 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); - } - } - - 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); - } - -#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); +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 (thread_in_row == 0) { - sRowDot[row] = thread_dot; + 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); + } } - } - if (tid < valid_q_rows) { - sLse[tid] = lse_ptr[tid]; - } - __syncthreads(); + 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; - 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); + uint4 q_val = make_uint4(0, 0, 0, 0); + uint4 do_val = make_uint4(0, 0, 0, 0); - if constexpr (IS_CAUSAL) { - if (start_col >= start_row + valid_q_rows) { - continue; - } + 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(); - const uint4* v_vec = reinterpret_cast(v_ptr + start_col * D); - uint4* sV_vec = reinterpret_cast(sV); + 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); + } -#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; + #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; + } } + + if (tid < valid_q_rows) { sLse[tid] = lse_ptr[tid]; } __syncthreads(); - 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; + 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); - 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; + if constexpr (IS_CAUSAL) { + if (start_col >= start_row + valid_q_rows) { continue; } + } - if (global_tile_idx >= total_tiles_dov) break; + 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(); - const int tile_m_idx = global_tile_idx / num_tiles_n_dov; - const int tile_n_idx = global_tile_idx % num_tiles_n_dov; + 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 = tile_m_idx * WMMA_M; - const int tile_n = tile_n_idx * WMMA_N; + 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; - if (tile_m >= valid_q_rows || tile_n >= valid_k_rows) continue; + if (global_tile_idx >= total_tiles_dov) break; - fragment a_frag; - fragment b_frag; - fragment acc_frag; + const int tile_m_idx = global_tile_idx / num_tiles_n_dov; + const int tile_n_idx = global_tile_idx % num_tiles_n_dov; - fill_fragment(acc_frag, 0.0f); + const int tile_m = tile_m_idx * WMMA_M; + const int tile_n = tile_n_idx * WMMA_N; -#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); - } + if (tile_m >= valid_q_rows || tile_n >= valid_k_rows) continue; - store_matrix_sync(sdOV + tile_m * S_STRIDE + tile_n, acc_frag, S_STRIDE, - mem_row_major); - } - __syncthreads(); + fragment a_frag; + fragment b_frag; + fragment acc_frag; - const uint4* k_vec = reinterpret_cast(k_ptr + start_col * D); - uint4* sK_vec = reinterpret_cast(sK); + fill_fragment(acc_frag, 0.0f); -#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(); + #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); + } - 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; + store_matrix_sync(sdOV + 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(); + + 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; } - } - store_matrix_sync(sS + tile_m * S_STRIDE + tile_n, acc_frag, S_STRIDE, - mem_row_major); - } - __syncthreads(); + __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); + } + __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(); - 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]; + 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(); } - __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); - -#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); + 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" + ); } - __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; - - 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(); - - 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); +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); - if constexpr (IS_CAUSAL) { - if (start_kv >= start_col + valid_q_rows) { - continue; - } - } + #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; - const uint4* q_vec = reinterpret_cast(q_ptr + start_col * D); - uint4* sQ_vec = reinterpret_cast(sdO); + uint4 k_val = make_uint4(0, 0, 0, 0); + uint4 v_val = make_uint4(0, 0, 0, 0); -#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 (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(); - 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; + 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; } } - } else { -#pragma unroll - for (int i = 0; i < acc_frag.num_elements; ++i) { - acc_frag.x[i] *= softmax_scale; + + 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; } } - } - 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); + 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); + } + __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; } + } + } + __syncthreads(); -#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 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; - const __half* current_o_ptr = o_ptr + start_col * D; + 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; - 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; - } - } + const int tile_m_idx = tile_idx / num_tiles_n_dv; + const int tile_n_idx = tile_idx % num_tiles_n_dv; - if (tid < valid_q_rows) { - sLse[tid] = lse_ptr[start_col + tid]; - } - __syncthreads(); + const int tile_m = tile_m_idx * WMMA_M; + const int tile_n = tile_n_idx * WMMA_N; - 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(); + 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); - 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; + #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); } - if (valid1) { - sP[row1 * BLOCK_M + col1] = p_h2.y; - sdS[row1 * BLOCK_M + col1] = ds_h2.y; + __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; } - } - } - __syncthreads(); + __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; - - 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(); + 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; - __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; + 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 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); + 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" + ); } - __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<<>>( - 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 - <<>>( + 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()), @@ -1250,104 +1191,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); - } + 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 + ); + } } 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 fc9cfb4c7a..1bc0c75448 100644 --- a/flash-attention-v100/kernel/fused_mha_forward.cu +++ b/flash-attention-v100/kernel/fused_mha_forward.cu @@ -16,1109 +16,1133 @@ 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; - } + 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); - } + 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; } - 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]); +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); } - 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 (start_col + valid_k_rows <= min_key_pos) { - continue; + if (window_size_right >= 0) { + max_key_pos = min(max_key_pos, + causal_q_offset + start_row + valid_q_rows - 1 + + window_size_right); } - - 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 (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; } - __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); - } + const uint4* q_vec = reinterpret_cast(q_ptr); + uint4* sQ_vec = reinterpret_cast(sQ); -#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; + #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]); } - store_matrix_sync(sS + tile_m * S_STRIDE + tile_n, acc_frag, S_STRIDE, - mem_row_major); - } + sQ_vec[row * q_stride_uint4 + vec_col] = q_val; } __syncthreads(); - 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; - - 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_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))); - } + if (start_col + valid_k_rows <= min_key_pos) { + continue; } -#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]); - } + const uint4* k_vec = reinterpret_cast(k_ptr + start_col * D); + uint4* sK_vec = reinterpret_cast(sK); -#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 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 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); + 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); + } } + __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 uint4* v_vec = reinterpret_cast(v_ptr + start_col * D); + uint4* sV_vec = reinterpret_cast(sV); - float row_sum = - __shfl_sync(mask, thread_sum, row_leader, 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; - if (thread_in_row == 0) { - sRowSum[row] = exp_diff * sRowSum[row] + row_sum; - sRowMax[row] = new_max; + 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(); - 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 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(); } + } - if (tail_col >= 0) { - sP_row_h[tail_col] = tail_value; - } + 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" + ); + } -#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 (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 (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; - } +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 num_tiles_k_pv = (p_tile_capacity + WMMA_K - 1) / WMMA_K; + 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); - 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; + if constexpr (IS_CAUSAL) { + if (start_col >= causal_q_offset + start_row + valid_q_rows) { + return; + } + } - 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 tid = threadIdx.x; + const int warp_id = tid / MAX_THREADS_PER_WARP; + const int lane_id = tid % MAX_THREADS_PER_WARP; - const int tile_m = tile_m_idx * WMMA_M; - const int tile_d = tile_d_idx * WMMA_N; + 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; - if (tile_m >= valid_q_rows) continue; + extern __shared__ char smem_raw[]; + init_smem(smem_raw); + auto& smem = *reinterpret_cast(smem_raw); - fragment a_frag; - fragment b_frag; - fragment acc_frag; + __half* sQ = smem.q; + __half* sK = smem.reuse_kv.k; + float* sS = smem.reuse_sp.s; - load_matrix_sync(acc_frag, sO + tile_m * O_STRIDE + tile_d, O_STRIDE, - mem_row_major); + 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; -#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; + const uint4* q_vec = reinterpret_cast(q_ptr); + uint4* sQ_vec = reinterpret_cast(sQ); - 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); + #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]); } - store_matrix_sync(sO + tile_m * O_STRIDE + tile_d, acc_frag, O_STRIDE, - mem_row_major); - } - __syncthreads(); + sQ_vec[row * q_stride_uint4 + vec_col] = q_val; } - } - - 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; - } -} -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; + 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]); } - } else { -#pragma unroll - for (int i = 0; i < acc_frag.num_elements; ++i) { - acc_frag.x[i] *= softmax_scale; + 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); } - } - 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]; - } + __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; + } + 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); + 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); - } +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 + ); + } } -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; +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..7a4d1056aa 100644 --- a/flash-attention-v100/kernel/fused_mha_forward_paged.cu +++ b/flash-attention-v100/kernel/fused_mha_forward_paged.cu @@ -3,7 +3,6 @@ #include #include #include -#include #include #include #include @@ -24,2064 +23,2296 @@ 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 namespace { int kv_cache_dtype_code_from_string(const std::string& kv_cache_dtype) { - if (kv_cache_dtype == "auto" || kv_cache_dtype == "bfloat16") { - return flash_v100::KV_CACHE_DTYPE_FP16; - } - if (kv_cache_dtype == "fp8" || kv_cache_dtype == "fp8_e4m3") { - return flash_v100::KV_CACHE_DTYPE_FP8_E4M3; - } - if (kv_cache_dtype == "fp8_e5m2") { - return flash_v100::KV_CACHE_DTYPE_FP8_E5M2; - } - return -1; + if (kv_cache_dtype == "auto" || kv_cache_dtype == "bfloat16") { + return flash_v100::KV_CACHE_DTYPE_FP16; + } + if (kv_cache_dtype == "fp8" || kv_cache_dtype == "fp8_e4m3") { + return flash_v100::KV_CACHE_DTYPE_FP8_E4M3; + } + if (kv_cache_dtype == "fp8_e5m2") { + return flash_v100::KV_CACHE_DTYPE_FP8_E5M2; + } + return -1; } } // namespace -#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 128 #define KV_STAGE_N_256_LOW_SMEM 1 -#define WARPS_256_LOW_SMEM 16 +#define WARPS_256_LOW_SMEM 16 #define BLOCK_N_256_LOW_SMEM_SCALAR_QK 32 #define KV_STAGE_N_256_LOW_SMEM_SCALAR_QK BLOCK_N_256_LOW_SMEM_SCALAR_QK -#define LOW_SMEM_PAGE_SIZE 16 +#define LOW_SMEM_PAGE_SIZE 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 - : (LOW_SMEM ? (LOW_SMEM_BM32 ? BLOCK_M_256 : 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 ? (LOW_SMEM_SCALAR_QK ? BLOCK_N_256_LOW_SMEM_SCALAR_QK - : 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 KV_STAGE_N = - (D == 256 && LOW_SMEM) - ? (LOW_SMEM_SCALAR_QK ? KV_STAGE_N_256_LOW_SMEM_SCALAR_QK - : KV_STAGE_N_256_LOW_SMEM) - : BLOCK_N; - 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 PIPELINE_S_ELEMENTS = - (D == 256 && LOW_SMEM && D256_SW_PIPELINE_QK && D256_SW_PIPELINE_PV) - ? BLOCK_M * S_STRIDE - : 1; - static constexpr int PIPELINE_PAGE_ELEMENTS = - (D == 256 && LOW_SMEM && D256_SW_PIPELINE_QK && D256_SW_PIPELINE_PV) - ? BLOCK_N / LOW_SMEM_PAGE_SIZE - : 1; - static constexpr bool PIPELINE_SWIZZLED_Q = - D == 256 && LOW_SMEM && D256_SW_PIPELINE_QK && D256_SW_PIPELINE_PV; - static constexpr int Q_ELEMENTS = - PIPELINE_SWIZZLED_Q ? BLOCK_M * D : BLOCK_M * Q_STRIDE; - // The D256 low-SMEM candidate changes only the output accumulator pitch - // from 264 to 268 floats. This preserves all WMMA/softmax arithmetic and - // targets the largest 8-way shared-memory load/store groups. - static constexpr int OUTPUT_EXTRA_PAD = - (D == 256 && LOW_SMEM && D256_OUTPUT_STRIDE_268) ? 4 : 0; - static constexpr int O_STRIDE = D + PAD + OUTPUT_EXTRA_PAD; - static constexpr int PER_UINT4 = 8; - static constexpr int LOW_SMEM_PAGE_COUNT = - (D == 256 && LOW_SMEM) ? (BLOCK_N / LOW_SMEM_PAGE_SIZE) : 1; - struct alignas(128) SmemLayout { - alignas(16) __half q[Q_ELEMENTS]; + 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 ? (LOW_SMEM_BM32 ? BLOCK_M_256 : 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 ? (LOW_SMEM_SCALAR_QK ? BLOCK_N_256_LOW_SMEM_SCALAR_QK : 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 KV_STAGE_N = (D == 256 && LOW_SMEM) ? (LOW_SMEM_SCALAR_QK ? KV_STAGE_N_256_LOW_SMEM_SCALAR_QK : KV_STAGE_N_256_LOW_SMEM) : BLOCK_N; + 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 PIPELINE_S_ELEMENTS = + (D == 256 && LOW_SMEM && D256_SW_PIPELINE_QK + && D256_SW_PIPELINE_PV) + ? BLOCK_M * S_STRIDE + : 1; + static constexpr int PIPELINE_PAGE_ELEMENTS = + (D == 256 && LOW_SMEM && D256_SW_PIPELINE_QK + && D256_SW_PIPELINE_PV) + ? BLOCK_N / LOW_SMEM_PAGE_SIZE + : 1; + static constexpr bool PIPELINE_SWIZZLED_Q = + D == 256 && LOW_SMEM && D256_SW_PIPELINE_QK + && D256_SW_PIPELINE_PV; + static constexpr int Q_ELEMENTS = + PIPELINE_SWIZZLED_Q ? BLOCK_M * D : BLOCK_M * Q_STRIDE; + // The D256 low-SMEM candidate changes only the output accumulator pitch + // from 264 to 268 floats. This preserves all WMMA/softmax arithmetic and + // targets the largest 8-way shared-memory load/store groups. + static constexpr int OUTPUT_EXTRA_PAD = + (D == 256 && LOW_SMEM && D256_OUTPUT_STRIDE_268) ? 4 : 0; + static constexpr int O_STRIDE = D + PAD + OUTPUT_EXTRA_PAD; + static constexpr int PER_UINT4 = 8; + static constexpr int LOW_SMEM_PAGE_COUNT = + (D == 256 && LOW_SMEM) ? (BLOCK_N / LOW_SMEM_PAGE_SIZE) : 1; + struct alignas(128) SmemLayout { + alignas(16) __half q [Q_ELEMENTS]; union { - alignas(16) __half k[KV_STAGE_N * KV_STRIDE]; - alignas(16) __half v[KV_STAGE_N * KV_STRIDE]; + alignas(16) __half k [KV_STAGE_N * KV_STRIDE]; + alignas(16) __half v [KV_STAGE_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 pipeline_s[PIPELINE_S_ELEMENTS]; - alignas(16) int pipeline_page_idx[PIPELINE_PAGE_ELEMENTS]; - alignas(16) int pipeline_page_offset[PIPELINE_PAGE_ELEMENTS]; - // WMMA accumulator loads and stores require 256-bit alignment. - alignas(32) float o[BLOCK_M * O_STRIDE]; - alignas(16) float row_max[BLOCK_M]; - alignas(16) float row_sum[BLOCK_M]; - alignas(16) int page_idx[LOW_SMEM_PAGE_COUNT]; - alignas(16) int page_offset[LOW_SMEM_PAGE_COUNT]; - }; - - static constexpr size_t TOTAL_SMEM = - ((sizeof(SmemLayout) + 127) & ~size_t(127)); - static_assert(offsetof(SmemLayout, o) % 32 == 0, - "WMMA accumulator storage must be 256-bit aligned"); + alignas(16) __half p_strict[P_STRICT_ELEMENTS]; + alignas(16) float pipeline_s[PIPELINE_S_ELEMENTS]; + alignas(16) int pipeline_page_idx[PIPELINE_PAGE_ELEMENTS]; + alignas(16) int pipeline_page_offset[PIPELINE_PAGE_ELEMENTS]; + // WMMA accumulator loads and stores require 256-bit alignment. + alignas(32) float o [BLOCK_M * O_STRIDE]; + alignas(16) float row_max[BLOCK_M]; + alignas(16) float row_sum[BLOCK_M]; + alignas(16) int page_idx[LOW_SMEM_PAGE_COUNT]; + alignas(16) int page_offset[LOW_SMEM_PAGE_COUNT]; + }; + + static constexpr size_t TOTAL_SMEM = ((sizeof(SmemLayout) + 127) & ~size_t(127)); + static_assert(offsetof(SmemLayout, o) % 32 == 0, + "WMMA accumulator storage must be 256-bit aligned"); }; -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(); } __device__ __forceinline__ int pipeline_swizzled_q_row_slot(int row) { - return (row & 3) | ((row & 8) >> 1) | ((row & 4) << 1); + return (row & 3) | ((row & 8) >> 1) | ((row & 4) << 1); } -__device__ __forceinline__ int pipeline_swizzled_matrix_a_offset(int row, - int col) { - const int tile = col / WMMA_K; - const int within_tile = col - tile * WMMA_K; - const int plane = within_tile >> 3; - const int inner = within_tile & 7; - return tile * WMMA_M * WMMA_K + plane * WMMA_M * 8 + - pipeline_swizzled_q_row_slot(row) * 8 + inner; +__device__ __forceinline__ int pipeline_swizzled_matrix_a_offset( + int row, + int col) { + const int tile = col / WMMA_K; + const int within_tile = col - tile * WMMA_K; + const int plane = within_tile >> 3; + const int inner = within_tile & 7; + return tile * WMMA_M * WMMA_K + plane * WMMA_M * 8 + + pipeline_swizzled_q_row_slot(row) * 8 + inner; } __device__ __forceinline__ void load_pipeline_swizzled_matrix_a_fragment( fragment& frag, - const __half* __restrict__ matrix, int k_offset) { - const int lane = threadIdx.x & 31; - const int row = (lane & 3) + ((lane >> 4) & 1) * 4 + ((lane >> 2) & 1) * 8; - const int slot = pipeline_swizzled_q_row_slot(row); - const int tile_offset = (k_offset / WMMA_K) * WMMA_M * WMMA_K; - uint32_t address = static_cast( - __cvta_generic_to_shared(matrix + tile_offset + slot * 8)); - uint32_t* words = reinterpret_cast(&frag); - asm volatile("ld.shared.v4.u32 {%0, %1, %2, %3}, [%4];" - : "=r"(words[0]), "=r"(words[1]), "=r"(words[2]), "=r"(words[3]) - : "r"(address) - : "memory"); - address += WMMA_M * 8 * sizeof(__half); - asm volatile("ld.shared.v4.u32 {%0, %1, %2, %3}, [%4];" - : "=r"(words[4]), "=r"(words[5]), "=r"(words[6]), "=r"(words[7]) - : "r"(address) - : "memory"); + const __half* __restrict__ matrix, + int k_offset) { + const int lane = threadIdx.x & 31; + const int row = + (lane & 3) + ((lane >> 4) & 1) * 4 + ((lane >> 2) & 1) * 8; + const int slot = pipeline_swizzled_q_row_slot(row); + const int tile_offset = (k_offset / WMMA_K) * WMMA_M * WMMA_K; + uint32_t address = static_cast( + __cvta_generic_to_shared(matrix + tile_offset + slot * 8)); + uint32_t* words = reinterpret_cast(&frag); + asm volatile( + "ld.shared.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(words[0]), "=r"(words[1]), "=r"(words[2]), "=r"(words[3]) + : "r"(address) + : "memory"); + address += WMMA_M * 8 * sizeof(__half); + asm volatile( + "ld.shared.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(words[4]), "=r"(words[5]), "=r"(words[6]), "=r"(words[7]) + : "r"(address) + : "memory"); } -template +template __device__ __forceinline__ void pipeline_qk_slice( - const __half* __restrict__ sQ, float* __restrict__ score, - const __half* __restrict__ k_cache, const int* __restrict__ page_idx, - const int* __restrict__ page_offset, int64_t k_block_stride, - int64_t k_token_stride, int64_t k_head_stride, int kv_head_id, - int start_col, int valid_k_rows, int start_row, int valid_q_rows, - int causal_q_offset, int tile_n, int k_begin, bool load_partial, - bool finalize, float softmax_scale, int window_size_left, - int window_size_right, float neg_inf, int anchor_len = 0, - int anchored_window = 0) { - if (tile_n >= valid_k_rows) { - return; - } - - const int page_slot = tile_n >> 4; - const int block_offset = page_offset[page_slot]; - const int physical_block_idx = page_idx[page_slot]; - const __half* k_tile_base = k_cache + - (int64_t)physical_block_idx * k_block_stride + - (int64_t)block_offset * k_token_stride + - (int64_t)kv_head_id * k_head_stride; - - fragment a_frag; - fragment b_frag; - fragment acc_frag; - if (load_partial) { - load_matrix_sync(acc_frag, score + tile_n, S_STRIDE, mem_row_major); - } else { - fill_fragment(acc_frag, 0.0f); - } - -#pragma unroll - for (int k_tile = 0; k_tile < K_TILES; ++k_tile) { - const int k_offset = k_begin + k_tile * WMMA_K; - if constexpr (SWIZZLED_Q) { - load_pipeline_swizzled_matrix_a_fragment(a_frag, sQ, k_offset); - } else { - load_matrix_sync(a_frag, sQ + k_offset, Q_STRIDE); - } - load_matrix_sync(b_frag, k_tile_base + k_offset, k_token_stride); - mma_sync(acc_frag, a_frag, b_frag, acc_frag); - } - - if (finalize) { - const int lane_id = threadIdx.x & 31; - 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; - -#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 + 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 = 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; - } - bool is_anchor_valid = true; - if constexpr (ANCHORED_SWA) { - is_anchor_valid = (global_n < anchor_len) || - (global_q_pos - global_n < anchored_window); - } - acc_frag.x[i] = - (is_valid && is_causal_valid && is_window_valid && is_anchor_valid) - ? acc_frag.x[i] * softmax_scale - : neg_inf; - } - } - - store_matrix_sync(score + tile_n, acc_frag, S_STRIDE, mem_row_major); -} + const __half* __restrict__ sQ, + float* __restrict__ score, + const __half* __restrict__ k_cache, + const int* __restrict__ page_idx, + const int* __restrict__ page_offset, + int64_t k_block_stride, + int64_t k_token_stride, + int64_t k_head_stride, + int kv_head_id, + int start_col, + int valid_k_rows, + int start_row, + int valid_q_rows, + int causal_q_offset, + int tile_n, + int k_begin, + bool load_partial, + bool finalize, + float softmax_scale, + int window_size_left, + int window_size_right, + float neg_inf) { + if (tile_n >= valid_k_rows) { + return; + } -template -__device__ __forceinline__ void pipeline_pv_tile( - const __half* __restrict__ sP, float* __restrict__ sO, - const __half* __restrict__ v_cache, const int* __restrict__ page_idx, - const int* __restrict__ page_offset, int64_t v_block_stride, - int64_t v_token_stride, int64_t v_head_stride, int kv_head_id, - int sub_start, int sub_valid_k_rows, int tile_d) { - fragment a_frag; - fragment b_frag; - fragment acc_frag; - load_matrix_sync(acc_frag, sO + tile_d, O_STRIDE, mem_row_major); - -#pragma unroll - for (int tile_k = 0; tile_k < 2; ++tile_k) { - const int k_offset = tile_k * WMMA_K; - if (k_offset >= sub_valid_k_rows) { - break; - } - const int token_offset = sub_start + k_offset; - const int page_slot = token_offset >> 4; - const int block_offset = page_offset[page_slot] + (token_offset & 15); + const int page_slot = tile_n >> 4; + const int block_offset = page_offset[page_slot]; const int physical_block_idx = page_idx[page_slot]; - const __half* v_tile_ptr = v_cache + - (int64_t)physical_block_idx * v_block_stride + - (int64_t)block_offset * v_token_stride + - (int64_t)kv_head_id * v_head_stride + tile_d; - if constexpr (SWIZZLED_P) { - load_pipeline_swizzled_matrix_a_fragment(a_frag, sP, k_offset); + const __half* k_tile_base = + k_cache + (int64_t)physical_block_idx * k_block_stride + + (int64_t)block_offset * k_token_stride + + (int64_t)kv_head_id * k_head_stride; + + fragment a_frag; + fragment b_frag; + fragment acc_frag; + if (load_partial) { + load_matrix_sync( + acc_frag, score + tile_n, S_STRIDE, mem_row_major); } else { - load_matrix_sync(a_frag, sP + k_offset, P_STRIDE); + fill_fragment(acc_frag, 0.0f); } - load_matrix_sync(b_frag, v_tile_ptr, v_token_stride); - mma_sync(acc_frag, a_frag, b_frag, acc_frag); - } - store_matrix_sync(sO + tile_d, acc_frag, O_STRIDE, mem_row_major); -} -template -__global__ void __launch_bounds__( - KernelConfig::THREADS_PER_BLOCK, - 2) - flash_attention_forward_kernel_paged( - const __half* __restrict__ Q, const void* __restrict__ K_cache, - const void* __restrict__ V_cache, __half* __restrict__ Out, - float* __restrict__ softmax_lse, const int* __restrict__ block_table, - const int* __restrict__ seqused_k, const int B, const int H, - const int M, const int N, const int* __restrict__ bfla_block_mask, - const int bfla_mask_block_n, const int64_t bfla_mask_stride_b, - const int64_t bfla_mask_stride_h, const int64_t bfla_mask_stride_q, - const int64_t bfla_mask_stride_k, const int page_block_size, - const int num_kv_heads, const int64_t k_block_stride, - const int64_t k_token_stride, 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 softmax_scale, - const float k_scale, const float v_scale, const int window_size_left, - const int window_size_right, float* __restrict__ split_tmp_out, - float* __restrict__ split_tmp_row_max, - float* __restrict__ split_tmp_row_sum, const int split_kv_tiles, - const int* __restrict__ anchor_lens, const int anchored_window) { - static_assert(!ANCHORED_SWA || - (IS_CAUSAL && KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP16), - "ANCHORED_SWA requires causal attention and an fp16 KV cache"); - using Config = KernelConfig; - using Traits = FlashV100Traits; - - 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; - constexpr bool USE_SW_PIPELINE = D256_SW_PIPELINE_QK || D256_SW_PIPELINE_PV; - if constexpr (USE_SW_PIPELINE) { - static_assert(D == 256 && LOW_SMEM && !LOW_SMEM_CONTIG_FAST && - !LOW_SMEM_SCALAR_QK && !LOW_SMEM_BM32 && !SPLIT_KV && - BLOCK_M == 16 && BLOCK_N == 128 && WARPS_PER_BLOCK == 16, - "software pipeline is specialized for D256 BM16 BN128"); - } - 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 / num_kv_heads; - 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; - - const int actual_N = seqused_k[batch_id]; - int num_n_tiles = (actual_N + BLOCK_N - 1) / BLOCK_N; - const int valid_q_rows = min(BLOCK_M, M - start_row); - const int causal_q_offset = max(actual_N - M, 0); - - int anchor_len = 0; - if constexpr (ANCHORED_SWA) { - anchor_len = anchor_lens[batch_id]; - } - - int max_key_pos = actual_N - 1; - if constexpr (IS_CAUSAL) { - max_key_pos = - min(max_key_pos, start_row + valid_q_rows - 1 + causal_q_offset); - } - if (window_size_right >= 0) { - max_key_pos = min(max_key_pos, start_row + valid_q_rows - 1 + - causal_q_offset + 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, start_row + causal_q_offset - 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 __half* q_ptr = Q + q_head_linear * M * D + start_row * 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; - - const int max_num_blocks_per_seq = - (N + page_block_size - 1) / page_block_size; - const int* block_table_seq = block_table + batch_id * max_num_blocks_per_seq; - - const int64_t k_row_stride = k_token_stride; - const int64_t v_row_stride = v_token_stride; - - extern __shared__ __align__(128) 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; - int* sPageIdx = smem.page_idx; - int* sPageOffset = smem.page_offset; - - 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; - sRowSum[tid] = 0.0f; - } - constexpr int O_UINT4_PER_ROW = D / 4; - constexpr int O_UINT4_STRIDE = O_STRIDE / 4; - for (int idx = tid; idx < BLOCK_M * O_UINT4_PER_ROW; - idx += THREADS_PER_BLOCK) { - const int row = idx / O_UINT4_PER_ROW; - const int col = idx % O_UINT4_PER_ROW; - reinterpret_cast(sO)[row * O_UINT4_STRIDE + col] = - make_uint4(0, 0, 0, 0); - } - __syncthreads(); - - 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]); - } - if constexpr (Config::PIPELINE_SWIZZLED_Q) { - const int k_tile = vec_col >> 1; - const int plane = vec_col & 1; - const int slot = pipeline_swizzled_q_row_slot(row); - sQ_vec[k_tile * (2 * BLOCK_M) + plane * BLOCK_M + slot] = q_val; - } else { - sQ_vec[row * q_stride_uint4 + vec_col] = q_val; - } - } - __syncthreads(); - - int cross_block_first_n_tile = 0; - if constexpr (D256_SW_PIPELINE_QK && D256_SW_PIPELINE_PV) { - const bool use_cross_block_pipeline = - valid_q_rows == BLOCK_M && actual_N > 0 && page_block_size == 784 && - bfla_block_mask == nullptr && window_size_left < 0 && - window_size_right < 0; - if (use_cross_block_pipeline) { - const __half* k_cache_h = reinterpret_cast(K_cache); - const __half* v_cache_h = reinterpret_cast(V_cache); - - if (tid < Config::LOW_SMEM_PAGE_COUNT) { - const int global_token_idx = tid * LOW_SMEM_PAGE_SIZE; - const int virtual_block_idx = global_token_idx / page_block_size; - sPageIdx[tid] = __ldg(&block_table_seq[virtual_block_idx]); - sPageOffset[tid] = - global_token_idx - virtual_block_idx * page_block_size; - } - __syncthreads(); - - if (warp_id < BLOCK_N / WMMA_N) { - const int valid_k_rows = min(BLOCK_N, actual_N); - pipeline_qk_slice( - sQ, sS, k_cache_h, sPageIdx, sPageOffset, k_block_stride, - k_token_stride, k_head_stride, kv_head_id, 0, valid_k_rows, - start_row, valid_q_rows, causal_q_offset, warp_id * WMMA_N, 0, - false, true, softmax_scale, window_size_left, window_size_right, - NEG_INF, anchor_len, anchored_window); - } - __syncthreads(); - - // Keep the hot loop on a compile-time steady-state path. The - // existing exact path below drains the final block, avoiding a - // loop-carried has_next predicate in every softmax/PV panel. - const int steady_n_tiles = num_n_tiles - 1; - for (int block_n = 0; block_n < steady_n_tiles; ++block_n) { - const int start_col = block_n * BLOCK_N; - const int valid_k_rows = BLOCK_N; - const int next_start_col = start_col + BLOCK_N; - const int next_valid_k_rows = min(BLOCK_N, actual_N - next_start_col); - const bool odd_block = (block_n & 1) != 0; - float* const score_current = odd_block ? smem.pipeline_s : sS; - float* const score_next = odd_block ? sS : smem.pipeline_s; - int* const page_idx_current = - odd_block ? smem.pipeline_page_idx : sPageIdx; - int* const page_offset_current = - odd_block ? smem.pipeline_page_offset : sPageOffset; - int* const page_idx_next = - odd_block ? sPageIdx : smem.pipeline_page_idx; - int* const page_offset_next = - odd_block ? sPageOffset : smem.pipeline_page_offset; - - if (tid < Config::LOW_SMEM_PAGE_COUNT) { - const int page_token_offset = tid * LOW_SMEM_PAGE_SIZE; - const int global_token_idx = next_start_col + page_token_offset; - const int virtual_block_idx = global_token_idx / page_block_size; - page_idx_next[tid] = __ldg(&block_table_seq[virtual_block_idx]); - page_offset_next[tid] = - global_token_idx - virtual_block_idx * page_block_size; + #pragma unroll + for (int k_tile = 0; k_tile < K_TILES; ++k_tile) { + const int k_offset = k_begin + k_tile * WMMA_K; + if constexpr (SWIZZLED_Q) { + load_pipeline_swizzled_matrix_a_fragment( + a_frag, sQ, k_offset); + } else { + load_matrix_sync(a_frag, sQ + k_offset, Q_STRIDE); } - __syncthreads(); + load_matrix_sync( + b_frag, k_tile_base + k_offset, k_token_stride); + mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } - for (int sub_start = 0; sub_start < valid_k_rows; - sub_start += P_SUB_TILE) { - const int sub_valid_k_rows = - min(P_SUB_TILE, valid_k_rows - sub_start); - const int row = warp_id; - const int thread_in_row = lane_id; - const unsigned mask = 0xFFFFFFFFU; - float* sS_row_f = score_current + row * S_STRIDE + sub_start; - - 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) { - const int vc = thread_in_row + j * THREADS_PER_ROW; - if (vc < vec_cols) { - const float4 v4 = sS_vec4[vc]; - thread_max = fmaxf(thread_max, - fmaxf(fmaxf(v4.x, v4.y), fmaxf(v4.z, v4.w))); - } - } -#pragma unroll - 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)); - } - - const float row_max = __shfl_sync(mask, thread_max, 0); - 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; - int vc_base = thread_in_row; - -#pragma unroll 4 - for (int j = 0; j < vecs_per_thread; - ++j, vc_base += THREADS_PER_ROW) { - if (vc_base < vec_cols) { - const float4 v4 = sS_vec4[vc_base]; - const float e0 = __expf(fmaxf(v4.x - new_max, -80.0f)); - const float e1 = __expf(fmaxf(v4.y - new_max, -80.0f)); - const float e2 = __expf(fmaxf(v4.z - new_max, -80.0f)); - const float e3 = __expf(fmaxf(v4.w - new_max, -80.0f)); - thread_sum += (e0 + e1) + (e2 + e3); - const int p_col = vc_base * 4; - __half2* p_half2 = reinterpret_cast<__half2*>( - sP + pipeline_swizzled_matrix_a_offset(row, p_col)); - p_half2[0] = __float22half2_rn(make_float2(e0, e1)); - p_half2[1] = __float22half2_rn(make_float2(e2, e3)); + if (finalize) { + const int lane_id = threadIdx.x & 31; + 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; + + #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 + 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; } - } -#pragma unroll 4 - for (int c = tail_start + thread_in_row; c < sub_valid_k_rows; - c += THREADS_PER_ROW) { - const float e = __expf(fmaxf(sS_row_f[c] - new_max, -80.0f)); - thread_sum += e; - sP[pipeline_swizzled_matrix_a_offset(row, c)] = __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); - } - const float row_sum = __shfl_sync(mask, thread_sum, 0); - if (thread_in_row == 0) { - sRowSum[row] = exp_diff * sRowSum[row] + row_sum; - sRowMax[row] = new_max; - } - -#pragma unroll 4 - for (int c = tail_start + thread_in_row; c < P_SUB_TILE; - c += THREADS_PER_ROW) { - if (c >= sub_valid_k_rows) { - sP[pipeline_swizzled_matrix_a_offset(row, c)] = - __float2half(0.0f); + bool is_window_valid = true; + if (window_size_left >= 0) { + is_window_valid = + global_n >= global_q_pos - window_size_left; } - } - - if (block_n > 0 || sub_start > 0) { - float4* sO_vec = reinterpret_cast(sO + row * O_STRIDE); -#pragma unroll 4 - for (int ov = thread_in_row; ov < D / 4; ov += THREADS_PER_ROW) { - float4 v = sO_vec[ov]; - v.x *= exp_diff; - v.y *= exp_diff; - v.z *= exp_diff; - v.w *= exp_diff; - sO_vec[ov] = v; + if (window_size_right >= 0) { + is_window_valid = + is_window_valid + && global_n <= global_q_pos + window_size_right; } - } - __syncthreads(); + acc_frag.x[i] = + (is_valid && is_causal_valid && is_window_valid) + ? acc_frag.x[i] * softmax_scale + : neg_inf; + } + } - if (sub_start == 0) { - if (warp_id < BLOCK_N / WMMA_N) { - pipeline_qk_slice( - sQ, score_next, k_cache_h, page_idx_next, page_offset_next, - k_block_stride, k_token_stride, k_head_stride, kv_head_id, - next_start_col, next_valid_k_rows, start_row, valid_q_rows, - causal_q_offset, warp_id * WMMA_N, 0, false, true, - softmax_scale, window_size_left, window_size_right, NEG_INF, - anchor_len, anchored_window); - } else { - const int pv_warp = warp_id - BLOCK_N / WMMA_N; - pipeline_pv_tile( - sP, sO, v_cache_h, page_idx_current, page_offset_current, - v_block_stride, v_token_stride, v_head_stride, kv_head_id, - sub_start, sub_valid_k_rows, pv_warp * WMMA_N); - pipeline_pv_tile( - sP, sO, v_cache_h, page_idx_current, page_offset_current, - v_block_stride, v_token_stride, v_head_stride, kv_head_id, - sub_start, sub_valid_k_rows, - (pv_warp + BLOCK_N / WMMA_N) * WMMA_N); - } - } else { - pipeline_pv_tile( - sP, sO, v_cache_h, page_idx_current, page_offset_current, - v_block_stride, v_token_stride, v_head_stride, kv_head_id, - sub_start, sub_valid_k_rows, warp_id * WMMA_N); - } - __syncthreads(); + store_matrix_sync(score + tile_n, acc_frag, S_STRIDE, mem_row_major); +} + +template +__device__ __forceinline__ void pipeline_pv_tile( + const __half* __restrict__ sP, + float* __restrict__ sO, + const __half* __restrict__ v_cache, + const int* __restrict__ page_idx, + const int* __restrict__ page_offset, + int64_t v_block_stride, + int64_t v_token_stride, + int64_t v_head_stride, + int kv_head_id, + int sub_start, + int sub_valid_k_rows, + int tile_d) { + fragment a_frag; + fragment b_frag; + fragment acc_frag; + load_matrix_sync( + acc_frag, sO + tile_d, O_STRIDE, mem_row_major); + + #pragma unroll + for (int tile_k = 0; tile_k < 2; ++tile_k) { + const int k_offset = tile_k * WMMA_K; + if (k_offset >= sub_valid_k_rows) { + break; } - } - cross_block_first_n_tile = steady_n_tiles; - } - } - - int first_n_tile = cross_block_first_n_tile; - int last_n_tile = num_n_tiles; - if constexpr (SPLIT_KV) { - const int partition_id = blockIdx.y; - const int tiles_per_partition = split_kv_tiles > 1 ? split_kv_tiles : 1; - first_n_tile = partition_id * tiles_per_partition; - last_n_tile = min(num_n_tiles, first_n_tile + tiles_per_partition); - } - - for (int block_n = first_n_tile; block_n < last_n_tile; ++block_n) { - const int start_col = block_n * BLOCK_N; - if (start_col >= actual_N) break; - const int valid_k_rows = min(BLOCK_N, actual_N - start_col); - - if (start_col + valid_k_rows <= min_key_pos) { - continue; - } - - if (bfla_block_mask != nullptr && bfla_mask_block_n > 0) { - const int mask_q_idx = start_row / bfla_mask_block_n; - const int mask_k_idx = start_col / bfla_mask_block_n; - const int keep_tile = - __ldg(bfla_block_mask + (int64_t)batch_id * bfla_mask_stride_b + - (int64_t)kv_head_id * bfla_mask_stride_h + - (int64_t)mask_q_idx * bfla_mask_stride_q + - (int64_t)mask_k_idx * bfla_mask_stride_k); - if (keep_tile == 0) { - continue; - } - } - - const int partial_block_size = - (block_n == num_n_tiles - 1 && actual_N % BLOCK_N != 0) - ? (actual_N % BLOCK_N) - : -1; - - uint4* sK_vec = reinterpret_cast(sK); - const int64_t row_stride_uint4 = k_row_stride / PER_UINT4; - const int start_page = start_col / page_block_size; - const int page_offset = start_col % page_block_size; - const bool single_page_tile = - (page_offset + valid_k_rows) <= page_block_size; - const bool two_page_tile = - !single_page_tile && - (page_offset + valid_k_rows) <= (page_block_size * 2); - const int first_page_rows = - single_page_tile ? valid_k_rows - : min(valid_k_rows, page_block_size - page_offset); - const int second_page_rows = valid_k_rows - first_page_rows; - const int physical_block_idx0 = __ldg(&block_table_seq[start_page]); - const int physical_block_idx1 = - second_page_rows > 0 ? __ldg(&block_table_seq[start_page + 1]) : -1; - const bool four_page_aligned_tile = - D == 256 && page_block_size == 16 && page_offset == 0 && - BLOCK_N == page_block_size * 4 && valid_k_rows == BLOCK_N && - k_block_stride == (int64_t)page_block_size * k_token_stride && - v_block_stride == (int64_t)page_block_size * v_token_stride; - const int physical_block_idx2 = - four_page_aligned_tile ? __ldg(&block_table_seq[start_page + 2]) : -1; - const int physical_block_idx3 = - four_page_aligned_tile ? __ldg(&block_table_seq[start_page + 3]) : -1; - const bool four_page_contiguous_tile = - four_page_aligned_tile && - physical_block_idx1 == physical_block_idx0 + 1 && - physical_block_idx2 == physical_block_idx0 + 2 && - physical_block_idx3 == physical_block_idx0 + 3; - bool low_smem_contiguous_page_tile = false; - if constexpr (LOW_SMEM) { - static_assert(D == 256, "low-smem paged prefill is D=256 only"); - static_assert(KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP16, - "low-smem paged prefill is fp16-KV only"); - static_assert(BLOCK_N % LOW_SMEM_PAGE_SIZE == 0, - "low-smem BLOCK_N must be page aligned"); - - constexpr int LOW_SMEM_PAGE_COUNT = Config::LOW_SMEM_PAGE_COUNT; - if (tid < LOW_SMEM_PAGE_COUNT) { - const int page_token_offset = tid * LOW_SMEM_PAGE_SIZE; - if (page_token_offset < valid_k_rows) { - const int global_token_idx = start_col + page_token_offset; - const int virtual_block_idx = global_token_idx / page_block_size; - sPageIdx[tid] = __ldg(&block_table_seq[virtual_block_idx]); - sPageOffset[tid] = - global_token_idx - virtual_block_idx * page_block_size; + const int token_offset = sub_start + k_offset; + const int page_slot = token_offset >> 4; + const int block_offset = + page_offset[page_slot] + (token_offset & 15); + const int physical_block_idx = page_idx[page_slot]; + const __half* v_tile_ptr = + v_cache + (int64_t)physical_block_idx * v_block_stride + + (int64_t)block_offset * v_token_stride + + (int64_t)kv_head_id * v_head_stride + tile_d; + if constexpr (SWIZZLED_P) { + load_pipeline_swizzled_matrix_a_fragment( + a_frag, sP, k_offset); } else { - sPageIdx[tid] = -1; - sPageOffset[tid] = 0; + load_matrix_sync(a_frag, sP + k_offset, P_STRIDE); + } + load_matrix_sync(b_frag, v_tile_ptr, v_token_stride); + mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + store_matrix_sync(sO + tile_d, acc_frag, O_STRIDE, mem_row_major); +} + +template +__global__ void __launch_bounds__( + KernelConfig::THREADS_PER_BLOCK, 2) +flash_attention_forward_kernel_paged( + const __half* __restrict__ Q, + const void* __restrict__ K_cache, + const void* __restrict__ V_cache, + __half* __restrict__ Out, + float* __restrict__ softmax_lse, + const int* __restrict__ block_table, + const int* __restrict__ seqused_k, + const int B, + const int H, + const int M, + const int N, + const int* __restrict__ bfla_block_mask, + const int bfla_mask_block_n, + const int64_t bfla_mask_stride_b, + const int64_t bfla_mask_stride_h, + const int64_t bfla_mask_stride_q, + const int64_t bfla_mask_stride_k, + const int page_block_size, + const int num_kv_heads, + const int64_t k_block_stride, + const int64_t k_token_stride, + 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 softmax_scale, + const float k_scale, + const float v_scale, + const int window_size_left, + const int window_size_right, + float* __restrict__ split_tmp_out, + float* __restrict__ split_tmp_row_max, + float* __restrict__ split_tmp_row_sum, + const int split_kv_tiles +) { + using Config = KernelConfig; + using Traits = FlashV100Traits; + + 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; + constexpr bool USE_SW_PIPELINE = + D256_SW_PIPELINE_QK || D256_SW_PIPELINE_PV; + if constexpr (USE_SW_PIPELINE) { + static_assert(D == 256 && LOW_SMEM && !LOW_SMEM_CONTIG_FAST + && !LOW_SMEM_SCALAR_QK && !LOW_SMEM_BM32 + && !SPLIT_KV && BLOCK_M == 16 && BLOCK_N == 128 + && WARPS_PER_BLOCK == 16, + "software pipeline is specialized for D256 BM16 BN128"); + } + 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 / num_kv_heads; + 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; + + const int actual_N = seqused_k[batch_id]; + int num_n_tiles = (actual_N + BLOCK_N - 1) / BLOCK_N; + const int valid_q_rows = min(BLOCK_M, M - start_row); + const int causal_q_offset = max(actual_N - M, 0); + + int max_key_pos = actual_N - 1; + if constexpr (IS_CAUSAL) { + max_key_pos = min(max_key_pos, + start_row + valid_q_rows - 1 + causal_q_offset); + } + if (window_size_right >= 0) { + max_key_pos = min(max_key_pos, + start_row + valid_q_rows - 1 + causal_q_offset + + 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, start_row + causal_q_offset + - 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 __half* q_ptr = Q + q_head_linear * M * D + start_row * 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; + + const int max_num_blocks_per_seq = (N + page_block_size - 1) / page_block_size; + const int* block_table_seq = block_table + batch_id * max_num_blocks_per_seq; + + const int64_t k_row_stride = k_token_stride; + const int64_t v_row_stride = v_token_stride; + + extern __shared__ __align__(128) 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; + int* sPageIdx = smem.page_idx; + int* sPageOffset = smem.page_offset; + + 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; + sRowSum[tid] = 0.0f; + } + constexpr int O_UINT4_PER_ROW = D / 4; + constexpr int O_UINT4_STRIDE = O_STRIDE / 4; + for (int idx = tid; idx < BLOCK_M * O_UINT4_PER_ROW; + idx += THREADS_PER_BLOCK) { + const int row = idx / O_UINT4_PER_ROW; + const int col = idx % O_UINT4_PER_ROW; + reinterpret_cast(sO)[row * O_UINT4_STRIDE + col] = + make_uint4(0, 0, 0, 0); + } + __syncthreads(); + + 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]); } - } - __syncthreads(); - - if constexpr (LOW_SMEM_CONTIG_FAST) { - low_smem_contiguous_page_tile = - valid_k_rows == BLOCK_N && - k_block_stride == (int64_t)page_block_size * k_token_stride && - v_block_stride == (int64_t)page_block_size * v_token_stride && - sPageIdx[0] >= 0; -#pragma unroll - for (int i = 1; i < LOW_SMEM_PAGE_COUNT; ++i) { - const int linear_offset = sPageOffset[0] + i * LOW_SMEM_PAGE_SIZE; - const int expected_page_delta = linear_offset / page_block_size; - const int expected_page_offset = - linear_offset - expected_page_delta * page_block_size; - low_smem_contiguous_page_tile = - low_smem_contiguous_page_tile && - sPageIdx[i] == sPageIdx[0] + expected_page_delta && - sPageOffset[i] == expected_page_offset; + if constexpr (Config::PIPELINE_SWIZZLED_Q) { + const int k_tile = vec_col >> 1; + const int plane = vec_col & 1; + const int slot = pipeline_swizzled_q_row_slot(row); + sQ_vec[k_tile * (2 * BLOCK_M) + plane * BLOCK_M + slot] = + q_val; + } else { + sQ_vec[row * q_stride_uint4 + vec_col] = q_val; } - } else { - low_smem_contiguous_page_tile = false; - } + } + __syncthreads(); - const __half* K_cache_h = reinterpret_cast(K_cache); + int cross_block_first_n_tile = 0; + if constexpr (D256_SW_PIPELINE_QK && D256_SW_PIPELINE_PV) { + const bool use_cross_block_pipeline = + valid_q_rows == BLOCK_M && actual_N > 0 + && page_block_size == 784 && bfla_block_mask == nullptr + && window_size_left < 0 && window_size_right < 0; + if (use_cross_block_pipeline) { + const __half* k_cache_h = + reinterpret_cast(K_cache); + const __half* v_cache_h = + reinterpret_cast(V_cache); + + if (tid < Config::LOW_SMEM_PAGE_COUNT) { + const int global_token_idx = tid * LOW_SMEM_PAGE_SIZE; + const int virtual_block_idx = + global_token_idx / page_block_size; + sPageIdx[tid] = + __ldg(&block_table_seq[virtual_block_idx]); + sPageOffset[tid] = + global_token_idx + - virtual_block_idx * page_block_size; + } + __syncthreads(); - if constexpr (LOW_SMEM_SCALAR_QK) { - uint4* sK_vec = reinterpret_cast(sK); - 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) { - const int page_slot = row >> 4; - const int block_offset = - sPageOffset[page_slot] + (row - page_slot * LOW_SMEM_PAGE_SIZE); - const int physical_block_idx_direct = sPageIdx[page_slot]; - const __half* k_row_ptr; - if constexpr (LOW_SMEM_CONTIG_FAST) { - k_row_ptr = - low_smem_contiguous_page_tile - ? K_cache_h + (int64_t)sPageIdx[0] * k_block_stride + - (int64_t)(sPageOffset[0] + row) * k_token_stride + - (int64_t)kv_head_id * k_head_stride - : K_cache_h + - (int64_t)physical_block_idx_direct * - k_block_stride + - (int64_t)block_offset * k_token_stride + - (int64_t)kv_head_id * k_head_stride; - } else { - k_row_ptr = K_cache_h + - (int64_t)physical_block_idx_direct * k_block_stride + - (int64_t)block_offset * k_token_stride + - (int64_t)kv_head_id * k_head_stride; + if (warp_id < BLOCK_N / WMMA_N) { + const int valid_k_rows = min(BLOCK_N, actual_N); + pipeline_qk_slice( + sQ, sS, k_cache_h, sPageIdx, sPageOffset, + k_block_stride, k_token_stride, + k_head_stride, kv_head_id, 0, valid_k_rows, start_row, + valid_q_rows, causal_q_offset, warp_id * WMMA_N, 0, false, + true, softmax_scale, window_size_left, window_size_right, + NEG_INF); } - const uint4* k_row_vec = reinterpret_cast(k_row_ptr); - k_val = __ldg(&k_row_vec[vec_col]); - } - sK_vec[row * kv_stride_uint4 + vec_col] = k_val; - } - __syncthreads(); + __syncthreads(); + + // Keep the hot loop on a compile-time steady-state path. The + // existing exact path below drains the final block, avoiding a + // loop-carried has_next predicate in every softmax/PV panel. + const int steady_n_tiles = num_n_tiles - 1; + for (int block_n = 0; block_n < steady_n_tiles; ++block_n) { + const int start_col = block_n * BLOCK_N; + const int valid_k_rows = BLOCK_N; + const int next_start_col = start_col + BLOCK_N; + const int next_valid_k_rows = + min(BLOCK_N, actual_N - next_start_col); + const bool odd_block = (block_n & 1) != 0; + float* const score_current = + odd_block ? smem.pipeline_s : sS; + float* const score_next = + odd_block ? sS : smem.pipeline_s; + int* const page_idx_current = + odd_block ? smem.pipeline_page_idx : sPageIdx; + int* const page_offset_current = + odd_block ? smem.pipeline_page_offset : sPageOffset; + int* const page_idx_next = + odd_block ? sPageIdx : smem.pipeline_page_idx; + int* const page_offset_next = + odd_block ? sPageOffset : smem.pipeline_page_offset; + + if (tid < Config::LOW_SMEM_PAGE_COUNT) { + const int page_token_offset = + tid * LOW_SMEM_PAGE_SIZE; + const int global_token_idx = + next_start_col + page_token_offset; + const int virtual_block_idx = + global_token_idx / page_block_size; + page_idx_next[tid] = + __ldg(&block_table_seq[virtual_block_idx]); + page_offset_next[tid] = + global_token_idx + - virtual_block_idx * page_block_size; + } + __syncthreads(); + + for (int sub_start = 0; sub_start < valid_k_rows; + sub_start += P_SUB_TILE) { + const int sub_valid_k_rows = + min(P_SUB_TILE, valid_k_rows - sub_start); + const int row = warp_id; + const int thread_in_row = lane_id; + const unsigned mask = 0xFFFFFFFFU; + float* sS_row_f = + score_current + row * S_STRIDE + sub_start; + + 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) { + const int vc = + thread_in_row + j * THREADS_PER_ROW; + if (vc < vec_cols) { + const float4 v4 = sS_vec4[vc]; + thread_max = fmaxf( + thread_max, + fmaxf( + fmaxf(v4.x, v4.y), + fmaxf(v4.z, v4.w))); + } + } + #pragma unroll + 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)); + } + + const float row_max = + __shfl_sync(mask, thread_max, 0); + 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; + int vc_base = thread_in_row; + + #pragma unroll 4 + for (int j = 0; j < vecs_per_thread; + ++j, vc_base += THREADS_PER_ROW) { + if (vc_base < vec_cols) { + const float4 v4 = sS_vec4[vc_base]; + const float e0 = + __expf(fmaxf(v4.x - new_max, -80.0f)); + const float e1 = + __expf(fmaxf(v4.y - new_max, -80.0f)); + const float e2 = + __expf(fmaxf(v4.z - new_max, -80.0f)); + const float e3 = + __expf(fmaxf(v4.w - new_max, -80.0f)); + thread_sum += (e0 + e1) + (e2 + e3); + const int p_col = vc_base * 4; + __half2* p_half2 = reinterpret_cast<__half2*>( + sP + pipeline_swizzled_matrix_a_offset( + row, p_col)); + p_half2[0] = + __float22half2_rn(make_float2(e0, e1)); + p_half2[1] = + __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) { + const float e = __expf( + fmaxf(sS_row_f[c] - new_max, -80.0f)); + thread_sum += e; + sP[pipeline_swizzled_matrix_a_offset(row, c)] = + __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); + } + const float row_sum = + __shfl_sync(mask, thread_sum, 0); + if (thread_in_row == 0) { + sRowSum[row] = exp_diff * sRowSum[row] + row_sum; + sRowMax[row] = new_max; + } + + #pragma unroll 4 + for (int c = tail_start + thread_in_row; + c < P_SUB_TILE; c += THREADS_PER_ROW) { + if (c >= sub_valid_k_rows) { + sP[pipeline_swizzled_matrix_a_offset(row, c)] = + __float2half(0.0f); + } + } + + if (block_n > 0 || sub_start > 0) { + float4* sO_vec = reinterpret_cast( + sO + row * O_STRIDE); + #pragma unroll 4 + for (int ov = thread_in_row; ov < D / 4; + ov += THREADS_PER_ROW) { + float4 v = sO_vec[ov]; + v.x *= exp_diff; + v.y *= exp_diff; + v.z *= exp_diff; + v.w *= exp_diff; + sO_vec[ov] = v; + } + } + __syncthreads(); + + if (sub_start == 0) { + if (warp_id < BLOCK_N / WMMA_N) { + pipeline_qk_slice< + Q_STRIDE, S_STRIDE, D / WMMA_K, IS_CAUSAL, + true>( + sQ, score_next, k_cache_h, page_idx_next, + page_offset_next, k_block_stride, + k_token_stride, k_head_stride, kv_head_id, + next_start_col, next_valid_k_rows, start_row, + valid_q_rows, causal_q_offset, + warp_id * WMMA_N, 0, false, true, + softmax_scale, window_size_left, + window_size_right, NEG_INF); + } else { + const int pv_warp = + warp_id - BLOCK_N / WMMA_N; + pipeline_pv_tile( + sP, sO, v_cache_h, page_idx_current, + page_offset_current, v_block_stride, + v_token_stride, v_head_stride, kv_head_id, + sub_start, sub_valid_k_rows, + pv_warp * WMMA_N); + pipeline_pv_tile( + sP, sO, v_cache_h, page_idx_current, + page_offset_current, v_block_stride, + v_token_stride, v_head_stride, kv_head_id, + sub_start, sub_valid_k_rows, + (pv_warp + BLOCK_N / WMMA_N) * WMMA_N); + } + } else { + pipeline_pv_tile( + sP, sO, v_cache_h, page_idx_current, + page_offset_current, v_block_stride, + v_token_stride, v_head_stride, kv_head_id, + sub_start, sub_valid_k_rows, warp_id * WMMA_N); + } + __syncthreads(); + } - 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; - } - if constexpr (ANCHORED_SWA) { - is_valid = - is_valid && ((global_n < anchor_len) || - (global_q_pos - global_n < anchored_window)); - } - - 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 * softmax_scale : NEG_INF; - } - } else if constexpr (D256_SW_PIPELINE_QK) { - constexpr int QK_TILES = BLOCK_N / WMMA_N; - if (warp_id < QK_TILES) { - pipeline_qk_slice( - sQ, sS, K_cache_h, sPageIdx, sPageOffset, k_block_stride, - k_token_stride, k_head_stride, kv_head_id, start_col, - valid_k_rows, start_row, valid_q_rows, causal_q_offset, - warp_id * WMMA_N, 0, false, true, softmax_scale, window_size_left, - window_size_right, NEG_INF, anchor_len, anchored_window); + cross_block_first_n_tile = steady_n_tiles; } - } 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; + int first_n_tile = cross_block_first_n_tile; + int last_n_tile = num_n_tiles; + if constexpr (SPLIT_KV) { + const int partition_id = blockIdx.y; + const int tiles_per_partition = + split_kv_tiles > 1 ? split_kv_tiles : 1; + first_n_tile = partition_id * tiles_per_partition; + last_n_tile = min(num_n_tiles, first_n_tile + tiles_per_partition); + } - const int tile_m = tile_m_idx * WMMA_M; - const int tile_n = tile_n_idx * WMMA_N; + for (int block_n = first_n_tile; block_n < last_n_tile; ++block_n) { + const int start_col = block_n * BLOCK_N; + if (start_col >= actual_N) break; + const int valid_k_rows = min(BLOCK_N, actual_N - start_col); - if (tile_m >= valid_q_rows || tile_n >= valid_k_rows) { + if (start_col + valid_k_rows <= min_key_pos) { continue; - } + } - const int page_slot = tile_n >> 4; - const int block_offset = sPageOffset[page_slot]; - const int physical_block_idx_direct = sPageIdx[page_slot]; + if (bfla_block_mask != nullptr && bfla_mask_block_n > 0) { + const int mask_q_idx = start_row / bfla_mask_block_n; + const int mask_k_idx = start_col / bfla_mask_block_n; + const int keep_tile = __ldg( + bfla_block_mask + + (int64_t)batch_id * bfla_mask_stride_b + + (int64_t)kv_head_id * bfla_mask_stride_h + + (int64_t)mask_q_idx * bfla_mask_stride_q + + (int64_t)mask_k_idx * bfla_mask_stride_k); + if (keep_tile == 0) { + continue; + } + } - fragment a_frag; - fragment b_frag; - fragment acc_frag; - fill_fragment(acc_frag, 0.0f); + const int partial_block_size = (block_n == num_n_tiles - 1 && actual_N % BLOCK_N != 0) + ? (actual_N % BLOCK_N) : -1; -#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; + uint4* sK_vec = reinterpret_cast(sK); + const int64_t row_stride_uint4 = k_row_stride / PER_UINT4; + const int start_page = start_col / page_block_size; + const int page_offset = start_col % page_block_size; + const bool single_page_tile = + (page_offset + valid_k_rows) <= page_block_size; + const bool two_page_tile = + !single_page_tile && + (page_offset + valid_k_rows) <= (page_block_size * 2); + const int first_page_rows = + single_page_tile ? valid_k_rows + : min(valid_k_rows, page_block_size - page_offset); + const int second_page_rows = valid_k_rows - first_page_rows; + const int physical_block_idx0 = __ldg(&block_table_seq[start_page]); + const int physical_block_idx1 = + second_page_rows > 0 ? __ldg(&block_table_seq[start_page + 1]) : -1; + const bool four_page_aligned_tile = + D == 256 && page_block_size == 16 && page_offset == 0 + && BLOCK_N == page_block_size * 4 + && valid_k_rows == BLOCK_N + && k_block_stride == (int64_t)page_block_size * k_token_stride + && v_block_stride == (int64_t)page_block_size * v_token_stride; + const int physical_block_idx2 = + four_page_aligned_tile ? __ldg(&block_table_seq[start_page + 2]) : -1; + const int physical_block_idx3 = + four_page_aligned_tile ? __ldg(&block_table_seq[start_page + 3]) : -1; + const bool four_page_contiguous_tile = + four_page_aligned_tile + && physical_block_idx1 == physical_block_idx0 + 1 + && physical_block_idx2 == physical_block_idx0 + 2 + && physical_block_idx3 == physical_block_idx0 + 3; + bool low_smem_contiguous_page_tile = false; + if constexpr (LOW_SMEM) { + static_assert(D == 256, "low-smem paged prefill is D=256 only"); + static_assert(KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP16, + "low-smem paged prefill is fp16-KV only"); + static_assert(BLOCK_N % LOW_SMEM_PAGE_SIZE == 0, + "low-smem BLOCK_N must be page aligned"); + + constexpr int LOW_SMEM_PAGE_COUNT = + Config::LOW_SMEM_PAGE_COUNT; + if (tid < LOW_SMEM_PAGE_COUNT) { + const int page_token_offset = tid * LOW_SMEM_PAGE_SIZE; + if (page_token_offset < valid_k_rows) { + const int global_token_idx = + start_col + page_token_offset; + const int virtual_block_idx = + global_token_idx / page_block_size; + sPageIdx[tid] = __ldg( + &block_table_seq[virtual_block_idx]); + sPageOffset[tid] = + global_token_idx + - virtual_block_idx * page_block_size; + } else { + sPageIdx[tid] = -1; + sPageOffset[tid] = 0; + } + } + __syncthreads(); - const __half* k_tile_ptr; if constexpr (LOW_SMEM_CONTIG_FAST) { - k_tile_ptr = - low_smem_contiguous_page_tile - ? K_cache_h + (int64_t)sPageIdx[0] * k_block_stride + - (int64_t)(sPageOffset[0] + tile_n) * - k_token_stride + - (int64_t)kv_head_id * k_head_stride + k_offset - : K_cache_h + - (int64_t)physical_block_idx_direct * - k_block_stride + - (int64_t)block_offset * k_token_stride + - (int64_t)kv_head_id * k_head_stride + k_offset; + low_smem_contiguous_page_tile = + valid_k_rows == BLOCK_N + && k_block_stride + == (int64_t)page_block_size * k_token_stride + && v_block_stride + == (int64_t)page_block_size * v_token_stride + && sPageIdx[0] >= 0; + #pragma unroll + for (int i = 1; i < LOW_SMEM_PAGE_COUNT; ++i) { + const int linear_offset = + sPageOffset[0] + i * LOW_SMEM_PAGE_SIZE; + const int expected_page_delta = + linear_offset / page_block_size; + const int expected_page_offset = + linear_offset + - expected_page_delta * page_block_size; + low_smem_contiguous_page_tile = + low_smem_contiguous_page_tile + && sPageIdx[i] == sPageIdx[0] + expected_page_delta + && sPageOffset[i] == expected_page_offset; + } } else { - k_tile_ptr = K_cache_h + - (int64_t)physical_block_idx_direct * k_block_stride + - (int64_t)block_offset * k_token_stride + - (int64_t)kv_head_id * k_head_stride + k_offset; + low_smem_contiguous_page_tile = false; } - load_matrix_sync(a_frag, sQ + tile_m * Q_STRIDE + k_offset, - Q_STRIDE); - load_matrix_sync(b_frag, k_tile_ptr, k_token_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 __half* K_cache_h = reinterpret_cast(K_cache); + + if constexpr (LOW_SMEM_SCALAR_QK) { + uint4* sK_vec = reinterpret_cast(sK); + 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) { + const int page_slot = row >> 4; + const int block_offset = + sPageOffset[page_slot] + + (row - page_slot * LOW_SMEM_PAGE_SIZE); + const int physical_block_idx_direct = + sPageIdx[page_slot]; + const __half* k_row_ptr; + if constexpr (LOW_SMEM_CONTIG_FAST) { + k_row_ptr = + low_smem_contiguous_page_tile + ? K_cache_h + + (int64_t)sPageIdx[0] * k_block_stride + + (int64_t)(sPageOffset[0] + row) + * k_token_stride + + (int64_t)kv_head_id * k_head_stride + : K_cache_h + + (int64_t)physical_block_idx_direct + * k_block_stride + + (int64_t)block_offset + * k_token_stride + + (int64_t)kv_head_id * k_head_stride; + } else { + k_row_ptr = + K_cache_h + + (int64_t)physical_block_idx_direct + * k_block_stride + + (int64_t)block_offset * k_token_stride + + (int64_t)kv_head_id * k_head_stride; + } + const uint4* k_row_vec = + reinterpret_cast(k_row_ptr); + k_val = __ldg(&k_row_vec[vec_col]); + } + sK_vec[row * kv_stride_uint4 + vec_col] = k_val; + } + __syncthreads(); + + 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 * softmax_scale : NEG_INF; + } + } else if constexpr (D256_SW_PIPELINE_QK) { + constexpr int QK_TILES = BLOCK_N / WMMA_N; + if (warp_id < QK_TILES) { + pipeline_qk_slice( + sQ, sS, K_cache_h, sPageIdx, sPageOffset, + k_block_stride, k_token_stride, k_head_stride, + kv_head_id, start_col, valid_k_rows, start_row, + valid_q_rows, causal_q_offset, warp_id * WMMA_N, 0, + false, true, softmax_scale, window_size_left, + window_size_right, NEG_INF); + } + } 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; + } + + const int page_slot = tile_n >> 4; + const int block_offset = sPageOffset[page_slot]; + const int physical_block_idx_direct = sPageIdx[page_slot]; + + 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; + + const __half* k_tile_ptr; + if constexpr (LOW_SMEM_CONTIG_FAST) { + k_tile_ptr = + low_smem_contiguous_page_tile + ? K_cache_h + + (int64_t)sPageIdx[0] * k_block_stride + + (int64_t)(sPageOffset[0] + tile_n) + * k_token_stride + + (int64_t)kv_head_id * k_head_stride + + k_offset + : K_cache_h + + (int64_t)physical_block_idx_direct + * k_block_stride + + (int64_t)block_offset + * k_token_stride + + (int64_t)kv_head_id * k_head_stride + + k_offset; + } else { + k_tile_ptr = + K_cache_h + + (int64_t)physical_block_idx_direct + * k_block_stride + + (int64_t)block_offset * k_token_stride + + (int64_t)kv_head_id * k_head_stride + + k_offset; + } + + load_matrix_sync( + a_frag, + sQ + tile_m * Q_STRIDE + k_offset, + Q_STRIDE); + load_matrix_sync(b_frag, k_tile_ptr, k_token_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; + } + + acc_frag.x[i] = + (is_valid && is_causal_valid && is_window_valid) + ? acc_frag.x[i] * softmax_scale + : NEG_INF; + } - 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; + store_matrix_sync( + sS + tile_m * S_STRIDE + tile_n, acc_frag, S_STRIDE, + mem_row_major); + } } - bool is_window_valid = true; - if (window_size_left >= 0) { - is_window_valid = is_window_valid && - global_n >= global_q_pos - window_size_left; + __syncthreads(); + } else { + if constexpr (KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP16) { + const __half* K_cache_h = reinterpret_cast(K_cache); + const uint4* k_page0_vec = reinterpret_cast( + K_cache_h + (int64_t)physical_block_idx0 * k_block_stride + + (int64_t)page_offset * k_token_stride + + (int64_t)kv_head_id * k_head_stride); + const uint4* k_page1_vec = + two_page_tile && second_page_rows > 0 + ? reinterpret_cast( + K_cache_h + (int64_t)physical_block_idx1 * k_block_stride + + (int64_t)kv_head_id * k_head_stride) + : nullptr; + + #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) { + if (four_page_contiguous_tile) { + k_val = __ldg( + &k_page0_vec[row * row_stride_uint4 + vec_col]); + } else if (single_page_tile) { + k_val = __ldg(&k_page0_vec[row * row_stride_uint4 + vec_col]); + } else if (two_page_tile) { + if (row < first_page_rows) { + k_val = __ldg( + &k_page0_vec[row * row_stride_uint4 + vec_col]); + } else { + const int row_page1 = row - first_page_rows; + k_val = __ldg( + &k_page1_vec[row_page1 * row_stride_uint4 + vec_col]); + } + } else { + const uint4* k_vec = reinterpret_cast(K_cache_h); + const int global_token_idx = start_col + row; + const int virtual_block_idx = + global_token_idx / page_block_size; + const int block_offset = global_token_idx % page_block_size; + const int physical_block_idx_slow = + __ldg(&block_table_seq[virtual_block_idx]); + const int64_t physical_offset_halfs = + (int64_t)physical_block_idx_slow * k_block_stride + + (int64_t)block_offset * k_token_stride + + (int64_t)kv_head_id * k_head_stride; + const int64_t physical_offset_uint4 = + (physical_offset_halfs / PER_UINT4) + vec_col; + k_val = __ldg(&k_vec[physical_offset_uint4]); + } + } + sK_vec[row * kv_stride_uint4 + vec_col] = k_val; } - if (window_size_right >= 0) { - is_window_valid = is_window_valid && - global_n <= global_q_pos + window_size_right; + } else { + for (int idx = tid; idx < valid_k_rows * D; idx += THREADS_PER_BLOCK) { + const int row = idx / D; + const int col = idx % D; + const int global_token_idx = start_col + row; + const int virtual_block_idx = global_token_idx / page_block_size; + const int block_offset = global_token_idx % page_block_size; + const int physical_block_idx_slow = + __ldg(&block_table_seq[virtual_block_idx]); + const int64_t physical_offset = + (int64_t)physical_block_idx_slow * k_block_stride + + (int64_t)block_offset * k_token_stride + + (int64_t)kv_head_id * k_head_stride + col; + sK[row * KV_STRIDE + col] = + flash_v100::load_kv_cache_half( + K_cache, physical_offset, k_scale); } + } + __syncthreads(); - bool is_anchor_valid = true; - if constexpr (ANCHORED_SWA) { - is_anchor_valid = (global_n < anchor_len) || - (global_q_pos - global_n < anchored_window); - } + 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; - acc_frag.x[i] = (is_valid && is_causal_valid && is_window_valid && - is_anchor_valid) - ? acc_frag.x[i] * softmax_scale - : NEG_INF; - } + 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; - store_matrix_sync(sS + tile_m * S_STRIDE + tile_n, acc_frag, S_STRIDE, - mem_row_major); - } - } - __syncthreads(); - } else { - if constexpr (KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP16) { - const __half* K_cache_h = reinterpret_cast(K_cache); - const uint4* k_page0_vec = reinterpret_cast( - K_cache_h + (int64_t)physical_block_idx0 * k_block_stride + - (int64_t)page_offset * k_token_stride + - (int64_t)kv_head_id * k_head_stride); - const uint4* k_page1_vec = - two_page_tile && second_page_rows > 0 - ? reinterpret_cast( - K_cache_h + - (int64_t)physical_block_idx1 * k_block_stride + - (int64_t)kv_head_id * k_head_stride) - : nullptr; - -#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) { - if (four_page_contiguous_tile) { - k_val = __ldg(&k_page0_vec[row * row_stride_uint4 + vec_col]); - } else if (single_page_tile) { - k_val = __ldg(&k_page0_vec[row * row_stride_uint4 + vec_col]); - } else if (two_page_tile) { - if (row < first_page_rows) { - k_val = __ldg(&k_page0_vec[row * row_stride_uint4 + vec_col]); - } else { - const int row_page1 = row - first_page_rows; - k_val = - __ldg(&k_page1_vec[row_page1 * row_stride_uint4 + vec_col]); - } - } else { - const uint4* k_vec = reinterpret_cast(K_cache_h); - const int global_token_idx = start_col + row; - const int virtual_block_idx = global_token_idx / page_block_size; - const int block_offset = global_token_idx % page_block_size; - const int physical_block_idx_slow = - __ldg(&block_table_seq[virtual_block_idx]); - const int64_t physical_offset_halves = - (int64_t)physical_block_idx_slow * k_block_stride + - (int64_t)block_offset * k_token_stride + - (int64_t)kv_head_id * k_head_stride; - const int64_t physical_offset_uint4 = - (physical_offset_halves / PER_UINT4) + vec_col; - k_val = __ldg(&k_vec[physical_offset_uint4]); + 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); } - } - sK_vec[row * kv_stride_uint4 + vec_col] = k_val; - } - } else { - for (int idx = tid; idx < valid_k_rows * D; idx += THREADS_PER_BLOCK) { - const int row = idx / D; - const int col = idx % D; - const int global_token_idx = start_col + row; - const int virtual_block_idx = global_token_idx / page_block_size; - const int block_offset = global_token_idx % page_block_size; - const int physical_block_idx_slow = - __ldg(&block_table_seq[virtual_block_idx]); - const int64_t physical_offset = - (int64_t)physical_block_idx_slow * k_block_stride + - (int64_t)block_offset * k_token_stride + - (int64_t)kv_head_id * k_head_stride + col; - sK[row * KV_STRIDE + col] = flash_v100::load_kv_cache_half( - K_cache, physical_offset, k_scale); - } - } - __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; + #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; - 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); - } + 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; -#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; - } - - bool is_anchor_valid = true; - if constexpr (ANCHORED_SWA) { - is_anchor_valid = (global_n < anchor_len) || - (global_q_pos - global_n < anchored_window); - } - - acc_frag.x[i] = (is_valid && is_causal_valid && is_window_valid && - is_anchor_valid) - ? acc_frag.x[i] * softmax_scale - : NEG_INF; - } + 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; + } - store_matrix_sync(sS + tile_m * S_STRIDE + tile_n, acc_frag, S_STRIDE, - mem_row_major); - } - __syncthreads(); - } - - uint4* sV_vec = reinterpret_cast(sV); - const int64_t v_row_stride_uint4 = v_row_stride / PER_UINT4; - if constexpr (!LOW_SMEM) { - if constexpr (KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP16) { - const __half* V_cache_h = reinterpret_cast(V_cache); - const uint4* v_page0_vec = reinterpret_cast( - V_cache_h + (int64_t)physical_block_idx0 * v_block_stride + - (int64_t)page_offset * v_token_stride + - (int64_t)kv_head_id * v_head_stride); - const uint4* v_page1_vec = - two_page_tile && second_page_rows > 0 - ? reinterpret_cast( - V_cache_h + - (int64_t)physical_block_idx1 * v_block_stride + - (int64_t)kv_head_id * v_head_stride) - : nullptr; - -#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) { - if (four_page_contiguous_tile) { - v_val = __ldg(&v_page0_vec[row * v_row_stride_uint4 + vec_col]); - } else if (single_page_tile) { - v_val = __ldg(&v_page0_vec[row * v_row_stride_uint4 + vec_col]); - } else if (two_page_tile) { - if (row < first_page_rows) { - v_val = __ldg(&v_page0_vec[row * v_row_stride_uint4 + vec_col]); - } else { - const int row_page1 = row - first_page_rows; - v_val = __ldg( - &v_page1_vec[row_page1 * v_row_stride_uint4 + vec_col]); - } - } else { - const uint4* v_vec = reinterpret_cast(V_cache_h); - const int global_token_idx = start_col + row; - const int virtual_block_idx = global_token_idx / page_block_size; - const int block_offset = global_token_idx % page_block_size; - const int physical_block_idx_slow = - __ldg(&block_table_seq[virtual_block_idx]); - const int64_t physical_offset_halves = - (int64_t)physical_block_idx_slow * v_block_stride + - (int64_t)block_offset * v_token_stride + - (int64_t)kv_head_id * v_head_stride; - const int64_t physical_offset_uint4 = - (physical_offset_halves / PER_UINT4) + vec_col; - v_val = __ldg(&v_vec[physical_offset_uint4]); + acc_frag.x[i] = (is_valid && is_causal_valid && is_window_valid) + ? acc_frag.x[i] * softmax_scale + : NEG_INF; } - } - sV_vec[row * kv_stride_uint4 + vec_col] = v_val; + + store_matrix_sync(sS + tile_m * S_STRIDE + tile_n, acc_frag, S_STRIDE, mem_row_major); } - } else { - for (int idx = tid; idx < valid_k_rows * D; idx += THREADS_PER_BLOCK) { - const int row = idx / D; - const int col = idx % D; - const int global_token_idx = start_col + row; - const int virtual_block_idx = global_token_idx / page_block_size; - const int block_offset = global_token_idx % page_block_size; - const int physical_block_idx_slow = - __ldg(&block_table_seq[virtual_block_idx]); - const int64_t physical_offset = - (int64_t)physical_block_idx_slow * v_block_stride + - (int64_t)block_offset * v_token_stride + - (int64_t)kv_head_id * v_head_stride + col; - sV[row * KV_STRIDE + col] = flash_v100::load_kv_cache_half( - V_cache, physical_offset, v_scale); + __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 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); - - // ================================================================ - // Patch 0108 (surgical) — sS/sP union race fix, D-conditional. - // ISSUE-0007 / issues/0090. The race (W(__half2) P-commit racing - // R(float4) sS reads with no cross-warp barrier) exists ONLY for - // D != 256, where sP = reuse_sp.p aliases sS byte-for-byte. - // - // D == 256 is race-IMMUNE: sP = p_strict, a SEPARATE buffer - // (SmemLayout::p_strict, line ~134; selected at line ~491). So the - // D==256 arm below is the ORIGINAL monolithic block VERBATIM — no - // barrier, no variable hoisting. This is the ONLY difference from - // archived/patches/0011: 0011 split the guard and hoisted vars for - // EVERY D, which perturbed the D=256 register/spill allocation and - // cost the immune path +0.47..+2.79% for zero benefit. Guarding the - // whole transform under `if constexpr (D != 256)` makes the D=256 - // codegen byte-identical to baseline (proved via SASS diff) => the - // immune path pays EXACTLY zero. The D!=256 arm is 0011's - // two-phase-commit fix, verbatim (barrier ALONE, no riders). - // ================================================================ - if constexpr (D == 256) { - 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 - 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 = __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); - __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) { - 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); - - const __half2 p01 = __float22half2_rn(make_float2(e0, e1)); - const __half2 p23 = __float22half2_rn(make_float2(e2, e3)); - if constexpr (D256_SW_PIPELINE_PV) { - // D256 owns a separate probability buffer, so the - // values can leave registers before row reduction. - const int base_offset = vc_base * 2; - sP_half2[base_offset] = p01; - sP_half2[base_offset + 1] = p23; - } else { - half_buffer[h2_idx++] = p01; - half_buffer[h2_idx++] = p23; - } - } - } - -#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 = __expf(fmaxf(v - new_max, -80.0f)); - thread_sum += e; - if constexpr (D256_SW_PIPELINE_PV) { - sP_row_h[c] = __float2half_rn(e); - } else { - 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; - } - - if constexpr (!D256_SW_PIPELINE_PV) { - h2_idx = 0; - vc_base = thread_in_row; -#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); + uint4* sV_vec = reinterpret_cast(sV); + const int64_t v_row_stride_uint4 = v_row_stride / PER_UINT4; + if constexpr (!LOW_SMEM) { + if constexpr (KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP16) { + const __half* V_cache_h = reinterpret_cast(V_cache); + const uint4* v_page0_vec = reinterpret_cast( + V_cache_h + (int64_t)physical_block_idx0 * v_block_stride + + (int64_t)page_offset * v_token_stride + + (int64_t)kv_head_id * v_head_stride); + const uint4* v_page1_vec = + two_page_tile && second_page_rows > 0 + ? reinterpret_cast( + V_cache_h + (int64_t)physical_block_idx1 * v_block_stride + + (int64_t)kv_head_id * v_head_stride) + : nullptr; + + #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) { + if (four_page_contiguous_tile) { + v_val = __ldg( + &v_page0_vec[row * v_row_stride_uint4 + vec_col]); + } else if (single_page_tile) { + v_val = __ldg(&v_page0_vec[row * v_row_stride_uint4 + vec_col]); + } else if (two_page_tile) { + if (row < first_page_rows) { + v_val = __ldg( + &v_page0_vec[row * v_row_stride_uint4 + vec_col]); + } else { + const int row_page1 = row - first_page_rows; + v_val = __ldg( + &v_page1_vec[row_page1 * v_row_stride_uint4 + vec_col]); + } + } else { + const uint4* v_vec = reinterpret_cast(V_cache_h); + const int global_token_idx = start_col + row; + const int virtual_block_idx = + global_token_idx / page_block_size; + const int block_offset = global_token_idx % page_block_size; + const int physical_block_idx_slow = + __ldg(&block_table_seq[virtual_block_idx]); + const int64_t physical_offset_halfs = + (int64_t)physical_block_idx_slow * v_block_stride + + (int64_t)block_offset * v_token_stride + + (int64_t)kv_head_id * v_head_stride; + const int64_t physical_offset_uint4 = + (physical_offset_halfs / PER_UINT4) + vec_col; + v_val = __ldg(&v_vec[physical_offset_uint4]); + } + } + sV_vec[row * kv_stride_uint4 + vec_col] = v_val; } - } - - 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 = D / 4; - 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; + } else { + for (int idx = tid; idx < valid_k_rows * D; idx += THREADS_PER_BLOCK) { + const int row = idx / D; + const int col = idx % D; + const int global_token_idx = start_col + row; + const int virtual_block_idx = global_token_idx / page_block_size; + const int block_offset = global_token_idx % page_block_size; + const int physical_block_idx_slow = + __ldg(&block_table_seq[virtual_block_idx]); + const int64_t physical_offset = + (int64_t)physical_block_idx_slow * v_block_stride + + (int64_t)block_offset * v_token_stride + + (int64_t)kv_head_id * v_head_stride + col; + sV[row * KV_STRIDE + col] = + flash_v100::load_kv_cache_half( + V_cache, physical_offset, v_scale); } - } } - } else { - // ---------------------------------------------------------------- - // ISSUE-0007 race-class fix (patch 0011, - // patches/0011-paged-race-fix-issue0007/ — PoC: poc/race_poc.cu). - // - // For D != 256, sP aliases sS byte-for-byte (union reuse_sp; - // sP = reuse_sp.p selected above with p_stride == S_STRIDE). The - // softmax phase reads sS as float4 (max pass, exp pass, scalar - // tail) and then commits P into the SAME union bytes as - // __half2/__half. Row r's half P row occupies union bytes - // [2*S_STRIDE*r, +2*S_STRIDE) — inside the FLOAT score row of row - // r/2 — so the P commit of warp w lands exactly on the score row - // still being read as float4 by warp w/2. The register staging - // (half_buffer) orders read-then-write only WITHIN one thread; - // cross-warp there was NO ordering until the __syncthreads() at - // the end of this phase, i.e. AFTER both: a W(__half2) x R(float4) - // data race. Evidence: compute-sanitizer racecheck 3 errors / - // 87,712 hazard instances on the PoC extraction of this exact - // access pattern, and the REAL kernel 10/10-runs bitwise - // non-deterministic at D=128 (q=512/kv=4096 and q=1024/kv=8192). - // - // Fix, proved racecheck-clean AND bit-equal in the PoC: two-phase - // P commit. Phase A = all sS reads + reductions with P staged in - // registers exactly as before; ONE uniform block-scope - // __syncthreads(); phase B = all sP stores. The former single - // guarded region is split in two so the barrier sits in UNIFORM - // control flow, and the per-thread state below is hoisted to - // survive across the barrier. No arithmetic is changed or - // reordered. D == 256 keeps its original barrier-free structure: - // sP = p_strict there, a separate buffer that never aliases sS - // (constexpr guard on the barrier below). - // ---------------------------------------------------------------- - const int row = tid / THREADS_PER_ROW; - const int thread_in_row = tid % THREADS_PER_ROW; - __half* sP_row_h = sP + row * p_stride; - __half2* sP_half2 = reinterpret_cast<__half2*>(sP_row_h); - 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 exp_diff = 1.0f; - __half2 half_buffer[20]; - int tail_col = -1; - __half tail_value = __float2half(0.f); - - if (tid < valid_q_rows * 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; - - 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 - 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); - // 0011: exp_diff hoisted above the phase split (used by the - // sO rescale in phase B). Same expression, same position in - // the arithmetic order. - exp_diff = __expf(old_max - new_max); - - float thread_sum = 0.0f; - int vc_base = thread_in_row; - int h2_idx = 0; - -#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); - - const __half2 p01 = __float22half2_rn(make_float2(e0, e1)); - const __half2 p23 = __float22half2_rn(make_float2(e2, e3)); - if constexpr (D256_SW_PIPELINE_PV) { - // D256 owns a separate probability buffer, so the - // values can leave registers before row reduction. - const int base_offset = vc_base * 2; - sP_half2[base_offset] = p01; - sP_half2[base_offset + 1] = p23; - } else { - half_buffer[h2_idx++] = p01; - half_buffer[h2_idx++] = p23; - } - } - } - -#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 = __expf(fmaxf(v - new_max, -80.0f)); - thread_sum += e; - if constexpr (D256_SW_PIPELINE_PV) { - sP_row_h[c] = __float2half_rn(e); - } else { - 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; - } - } // end phase A: every sS read of this sub-tile is done; P is - // staged in registers (half_buffer / tail_col / tail_value). - - // ISSUE-0007 fix (patch 0011): THE one uniform __syncthreads() - // between the sS read phase and the sP commit phase — the only - // behavioral change of this patch. UNIFORMITY: this point is at - // the body scope of the sub_start loop; its trip count - // (valid_k_rows, softmax_sub_tile) and every enclosing control - // transfer (kernel early returns on blockIdx.z/M; block_n loop - // break/continue on start_col/min_key_pos/bfla block mask) are - // block-uniform — the same uniformity the pre-existing - // __syncthreads() at the end of this loop body already requires. - // All threads of the block reach this barrier exactly once per - // sub_start iteration. D == 256 is immune (sP = p_strict, - // separate buffer) and keeps its barrier-free codegen. - if constexpr (D != 256) { - __syncthreads(); + __syncthreads(); } - if (tid < valid_q_rows * THREADS_PER_ROW) { - // phase B: commit P to the union bytes. Bit-identical values - // (straight from the phase-A registers), same store order, - // same addresses as the pre-fix code. - if constexpr (!D256_SW_PIPELINE_PV) { - int h2_idx = 0; - int vc_base = thread_in_row; -#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 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; - if (tail_col >= 0) { - sP_row_h[tail_col] = tail_value; - } - } + 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); + + // ================================================================ + // Patch 0108 (surgical) — sS/sP union race fix, D-conditional. + // ISSUE-0007 / issues/0090. The race (W(__half2) P-commit racing + // R(float4) sS reads with no cross-warp barrier) exists ONLY for + // D != 256, where sP = reuse_sp.p aliases sS byte-for-byte. + // + // D == 256 is race-IMMUNE: sP = p_strict, a SEPARATE buffer + // (SmemLayout::p_strict, line ~134; selected at line ~491). So the + // D==256 arm below is the ORIGINAL monolithic block VERBATIM — no + // barrier, no variable hoisting. This is the ONLY difference from + // archived/patches/0011: 0011 split the guard and hoisted vars for + // EVERY D, which perturbed the D=256 register/spill allocation and + // cost the immune path +0.47..+2.79% for zero benefit. Guarding the + // whole transform under `if constexpr (D != 256)` makes the D=256 + // codegen byte-identical to baseline (proved via SASS diff) => the + // immune path pays EXACTLY zero. The D!=256 arm is 0011's + // two-phase-commit fix, verbatim (barrier ALONE, no riders). + // ================================================================ + if constexpr (D == 256) { + 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 < 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 = D / 4; - 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(); + #pragma unroll + 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]); + } - const int num_tiles_k_pv = (p_tile_capacity + WMMA_K - 1) / WMMA_K; + #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)); + } - 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 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); + __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) { + 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); + + const __half2 p01 = + __float22half2_rn(make_float2(e0, e1)); + const __half2 p23 = + __float22half2_rn(make_float2(e2, e3)); + if constexpr (D256_SW_PIPELINE_PV) { + // D256 owns a separate probability buffer, so the + // values can leave registers before row reduction. + const int base_offset = vc_base * 2; + sP_half2[base_offset] = p01; + sP_half2[base_offset + 1] = p23; + } else { + half_buffer[h2_idx++] = p01; + half_buffer[h2_idx++] = p23; + } + } + } - const int tile_m_idx = global_tile_idx / num_tiles_n_pv; - const int tile_d_idx = global_tile_idx % num_tiles_n_pv; + #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 = __expf(fmaxf(v - new_max, -80.0f)); + thread_sum += e; + if constexpr (D256_SW_PIPELINE_PV) { + sP_row_h[c] = __float2half_rn(e); + } else { + tail_col = c; + tail_value = __float2half_rn(e); + } + } - const int tile_m = tile_m_idx * WMMA_M; - const int tile_d = tile_d_idx * WMMA_N; + #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); + } - if (tile_m >= valid_q_rows) continue; + float row_sum = + __shfl_sync(mask, thread_sum, row_leader, THREADS_PER_ROW); - fragment a_frag; - fragment b_frag; - fragment acc_frag; + if (thread_in_row == 0) { + sRowSum[row] = exp_diff * sRowSum[row] + row_sum; + sRowMax[row] = new_max; + } - load_matrix_sync(acc_frag, sO + tile_m * O_STRIDE + tile_d, O_STRIDE, - mem_row_major); + if constexpr (!D256_SW_PIPELINE_PV) { + h2_idx = 0; + vc_base = thread_in_row; + #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 - 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; + #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); + } + } - load_matrix_sync(a_frag, sP + tile_m * p_stride + k_offset, p_stride); - if constexpr (LOW_SMEM) { - const __half* V_cache_h = reinterpret_cast(V_cache); - const int token_offset = sub_start + k_offset; - const int page_slot = token_offset >> 4; - const int block_offset = sPageOffset[page_slot]; - const int physical_block_idx_direct = sPageIdx[page_slot]; - const __half* v_tile_ptr; - if constexpr (LOW_SMEM_CONTIG_FAST) { - v_tile_ptr = - low_smem_contiguous_page_tile - ? V_cache_h + (int64_t)sPageIdx[0] * v_block_stride + - (int64_t)(sPageOffset[0] + token_offset) * - v_token_stride + - (int64_t)kv_head_id * v_head_stride + tile_d - : V_cache_h + - (int64_t)physical_block_idx_direct * - v_block_stride + - (int64_t)block_offset * v_token_stride + - (int64_t)kv_head_id * v_head_stride + tile_d; + 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 = D / 4; + 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; + } + } + } } else { - v_tile_ptr = V_cache_h + - (int64_t)physical_block_idx_direct * v_block_stride + - (int64_t)block_offset * v_token_stride + - (int64_t)kv_head_id * v_head_stride + tile_d; + // ---------------------------------------------------------------- + // ISSUE-0007 race-class fix (patch 0011, + // patches/0011-paged-race-fix-issue0007/ — PoC: poc/race_poc.cu). + // + // For D != 256, sP aliases sS byte-for-byte (union reuse_sp; + // sP = reuse_sp.p selected above with p_stride == S_STRIDE). The + // softmax phase reads sS as float4 (max pass, exp pass, scalar + // tail) and then commits P into the SAME union bytes as + // __half2/__half. Row r's half P row occupies union bytes + // [2*S_STRIDE*r, +2*S_STRIDE) — inside the FLOAT score row of row + // r/2 — so the P commit of warp w lands exactly on the score row + // still being read as float4 by warp w/2. The register staging + // (half_buffer) orders read-then-write only WITHIN one thread; + // cross-warp there was NO ordering until the __syncthreads() at + // the end of this phase, i.e. AFTER both: a W(__half2) x R(float4) + // data race. Evidence: compute-sanitizer racecheck 3 errors / + // 87,712 hazard instances on the PoC extraction of this exact + // access pattern, and the REAL kernel 10/10-runs bitwise + // non-deterministic at D=128 (q=512/kv=4096 and q=1024/kv=8192). + // + // Fix, proved racecheck-clean AND bit-equal in the PoC: two-phase + // P commit. Phase A = all sS reads + reductions with P staged in + // registers exactly as before; ONE uniform block-scope + // __syncthreads(); phase B = all sP stores. The former single + // guarded region is split in two so the barrier sits in UNIFORM + // control flow, and the per-thread state below is hoisted to + // survive across the barrier. No arithmetic is changed or + // reordered. D == 256 keeps its original barrier-free structure: + // sP = p_strict there, a separate buffer that never aliases sS + // (constexpr guard on the barrier below). + // ---------------------------------------------------------------- + const int row = tid / THREADS_PER_ROW; + const int thread_in_row = tid % THREADS_PER_ROW; + __half* sP_row_h = sP + row * p_stride; + __half2* sP_half2 = reinterpret_cast<__half2*>(sP_row_h); + 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 exp_diff = 1.0f; + __half2 half_buffer[20]; + int tail_col = -1; + __half tail_value = __float2half(0.f); + + if (tid < valid_q_rows * 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; + + 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 + 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); + // 0011: exp_diff hoisted above the phase split (used by the + // sO rescale in phase B). Same expression, same position in + // the arithmetic order. + exp_diff = __expf(old_max - new_max); + + float thread_sum = 0.0f; + int vc_base = thread_in_row; + int h2_idx = 0; + + #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); + + const __half2 p01 = + __float22half2_rn(make_float2(e0, e1)); + const __half2 p23 = + __float22half2_rn(make_float2(e2, e3)); + if constexpr (D256_SW_PIPELINE_PV) { + // D256 owns a separate probability buffer, so the + // values can leave registers before row reduction. + const int base_offset = vc_base * 2; + sP_half2[base_offset] = p01; + sP_half2[base_offset + 1] = p23; + } else { + half_buffer[h2_idx++] = p01; + half_buffer[h2_idx++] = p23; + } + } + } + + #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 = __expf(fmaxf(v - new_max, -80.0f)); + thread_sum += e; + if constexpr (D256_SW_PIPELINE_PV) { + sP_row_h[c] = __float2half_rn(e); + } else { + 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; + } + } // end phase A: every sS read of this sub-tile is done; P is + // staged in registers (half_buffer / tail_col / tail_value). + + // ISSUE-0007 fix (patch 0011): THE one uniform __syncthreads() + // between the sS read phase and the sP commit phase — the only + // behavioral change of this patch. UNIFORMITY: this point is at + // the body scope of the sub_start loop; its trip count + // (valid_k_rows, softmax_sub_tile) and every enclosing control + // transfer (kernel early returns on blockIdx.z/M; block_n loop + // break/continue on start_col/min_key_pos/bfla block mask) are + // block-uniform — the same uniformity the pre-existing + // __syncthreads() at the end of this loop body already requires. + // All threads of the block reach this barrier exactly once per + // sub_start iteration. D == 256 is immune (sP = p_strict, + // separate buffer) and keeps its barrier-free codegen. + if constexpr (D != 256) { + __syncthreads(); } - load_matrix_sync(b_frag, v_tile_ptr, v_token_stride); - } else { - 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); + + if (tid < valid_q_rows * THREADS_PER_ROW) { + // phase B: commit P to the union bytes. Bit-identical values + // (straight from the phase-A registers), same store order, + // same addresses as the pre-fix code. + if constexpr (!D256_SW_PIPELINE_PV) { + int h2_idx = 0; + int vc_base = thread_in_row; + #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 = D / 4; + 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(); + + 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); + if constexpr (LOW_SMEM) { + const __half* V_cache_h = + reinterpret_cast(V_cache); + const int token_offset = sub_start + k_offset; + const int page_slot = token_offset >> 4; + const int block_offset = sPageOffset[page_slot]; + const int physical_block_idx_direct = + sPageIdx[page_slot]; + const __half* v_tile_ptr; + if constexpr (LOW_SMEM_CONTIG_FAST) { + v_tile_ptr = + low_smem_contiguous_page_tile + ? V_cache_h + + (int64_t)sPageIdx[0] * v_block_stride + + (int64_t)(sPageOffset[0] + token_offset) + * v_token_stride + + (int64_t)kv_head_id * v_head_stride + + tile_d + : V_cache_h + + (int64_t)physical_block_idx_direct + * v_block_stride + + (int64_t)block_offset * v_token_stride + + (int64_t)kv_head_id * v_head_stride + + tile_d; + } else { + v_tile_ptr = + V_cache_h + + (int64_t)physical_block_idx_direct + * v_block_stride + + (int64_t)block_offset * v_token_stride + + (int64_t)kv_head_id * v_head_stride + tile_d; + } + load_matrix_sync(b_frag, v_tile_ptr, v_token_stride); + } else { + 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(); } - store_matrix_sync(sO + tile_m * O_STRIDE + tile_d, acc_frag, O_STRIDE, - mem_row_major); - } - __syncthreads(); + } - } - if constexpr (SPLIT_KV) { - const int partition_id = blockIdx.y; - const int num_partitions = gridDim.y; - const int64_t row_base = - ((int64_t)batch_head_id * num_partitions + partition_id) * M + - start_row; + if constexpr (SPLIT_KV) { + const int partition_id = blockIdx.y; + const int num_partitions = gridDim.y; + const int64_t row_base = + ((int64_t)batch_head_id * num_partitions + partition_id) * M + + start_row; + + for (int i = tid; i < valid_q_rows * D; i += THREADS_PER_BLOCK) { + const int row = i / D; + const int col = i - row * D; + split_tmp_out[(row_base + row) * D + col] = + sO[row * O_STRIDE + col]; + } + + if (tid < valid_q_rows) { + split_tmp_row_max[row_base + tid] = sRowMax[tid]; + split_tmp_row_sum[row_base + tid] = sRowSum[tid]; + } + return; + } - for (int i = tid; i < valid_q_rows * D; i += THREADS_PER_BLOCK) { - const int row = i / D; - const int col = i - row * D; - split_tmp_out[(row_base + row) * D + col] = sO[row * O_STRIDE + col]; + 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) { - split_tmp_row_max[row_base + tid] = sRowMax[tid]; - split_tmp_row_sum[row_base + tid] = sRowSum[tid]; - } - return; - } - - 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); - } + const float sum = fmaxf(sRowSum[tid], 1e-24f); + softmax_lse_ptr[tid] = sRowMax[tid] + logf(sum); + } } 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_default_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; } -template +template void launcher_flash_attention_forward_paged_impl( - const torch::Tensor& Q, const torch::Tensor& K_cache, - const torch::Tensor& V_cache, torch::Tensor& Out, - torch::Tensor& softmax_lse, const torch::Tensor& block_table, - const torch::Tensor& seq_lens, const int* bfla_mask_ptr, - int bfla_mask_block_n, int64_t bfla_mask_stride_b, - int64_t bfla_mask_stride_h, int64_t bfla_mask_stride_q, - int64_t bfla_mask_stride_k, float softmax_scale, bool is_causal, - float k_scale, float v_scale, int window_size_left, int window_size_right, - cudaStream_t stream, const int* anchor_lens = nullptr, - int anchored_window = 0) { - using Config = KernelConfig; - - const int B = Q.size(0); - const int H = Q.size(1); - const int M = Q.size(2); - const int page_block_size = K_cache.size(1); - const int num_kv_heads = K_cache.size(2); - const int max_num_blocks = block_table.size(1); - const int N = max_num_blocks * page_block_size; - const int64_t k_block_stride = K_cache.stride(0); - const int64_t k_token_stride = K_cache.stride(1); - const int64_t k_head_stride = K_cache.stride(2); - const int64_t v_block_stride = V_cache.stride(0); - const int64_t v_token_stride = V_cache.stride(1); - const int64_t v_head_stride = V_cache.stride(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"); - - const bool use_anchored = anchor_lens != nullptr && anchored_window > 0; - TORCH_CHECK(!use_anchored || is_causal, - "anchored decode window requires causal attention"); - if constexpr (KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP16) { - if (use_anchored) { - auto anchored_kernel = flash_attention_forward_kernel_paged< - D, LOW_SMEM, LOW_SMEM_CONTIG_FAST, LOW_SMEM_SCALAR_QK, LOW_SMEM_BM32, - false, true, KV_DTYPE, D256_OUTPUT_STRIDE_268, D256_SW_PIPELINE_QK, - D256_SW_PIPELINE_PV, true>; - cudaFuncSetAttribute((void*)anchored_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, smem); - anchored_kernel<<>>( - reinterpret_cast(Q.data_ptr()), K_cache.data_ptr(), - V_cache.data_ptr(), reinterpret_cast<__half*>(Out.data_ptr()), - softmax_lse.data_ptr(), block_table.data_ptr(), - seq_lens.data_ptr(), B, H, M, N, bfla_mask_ptr, - bfla_mask_block_n, bfla_mask_stride_b, bfla_mask_stride_h, - bfla_mask_stride_q, bfla_mask_stride_k, page_block_size, num_kv_heads, - k_block_stride, k_token_stride, k_head_stride, v_block_stride, - v_token_stride, v_head_stride, softmax_scale, k_scale, v_scale, - window_size_left, window_size_right, nullptr, nullptr, nullptr, 0, - anchor_lens, anchored_window); - return; - } - } else { - TORCH_CHECK(!use_anchored, - "anchored decode window requires an fp16 KV cache"); - } - - auto kernel = - is_causal - ? (void*)flash_attention_forward_kernel_paged< - D, LOW_SMEM, LOW_SMEM_CONTIG_FAST, LOW_SMEM_SCALAR_QK, - LOW_SMEM_BM32, false, true, KV_DTYPE, D256_OUTPUT_STRIDE_268, - D256_SW_PIPELINE_QK, D256_SW_PIPELINE_PV> - : (void*)flash_attention_forward_kernel_paged< - D, LOW_SMEM, LOW_SMEM_CONTIG_FAST, LOW_SMEM_SCALAR_QK, - LOW_SMEM_BM32, false, false, KV_DTYPE, D256_OUTPUT_STRIDE_268, - D256_SW_PIPELINE_QK, D256_SW_PIPELINE_PV>; - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, - smem); - - if (is_causal) { - flash_attention_forward_kernel_paged< - D, LOW_SMEM, LOW_SMEM_CONTIG_FAST, LOW_SMEM_SCALAR_QK, LOW_SMEM_BM32, - false, true, KV_DTYPE, D256_OUTPUT_STRIDE_268, D256_SW_PIPELINE_QK, - D256_SW_PIPELINE_PV><<>>( - reinterpret_cast(Q.data_ptr()), K_cache.data_ptr(), - V_cache.data_ptr(), reinterpret_cast<__half*>(Out.data_ptr()), - softmax_lse.data_ptr(), block_table.data_ptr(), - seq_lens.data_ptr(), B, H, M, N, bfla_mask_ptr, bfla_mask_block_n, - bfla_mask_stride_b, bfla_mask_stride_h, bfla_mask_stride_q, - bfla_mask_stride_k, page_block_size, num_kv_heads, k_block_stride, - k_token_stride, k_head_stride, v_block_stride, v_token_stride, - v_head_stride, softmax_scale, k_scale, v_scale, window_size_left, - window_size_right, nullptr, nullptr, nullptr, 0, nullptr, 0); - } else { - flash_attention_forward_kernel_paged< - D, LOW_SMEM, LOW_SMEM_CONTIG_FAST, LOW_SMEM_SCALAR_QK, LOW_SMEM_BM32, - false, false, KV_DTYPE, D256_OUTPUT_STRIDE_268, D256_SW_PIPELINE_QK, - D256_SW_PIPELINE_PV><<>>( - reinterpret_cast(Q.data_ptr()), K_cache.data_ptr(), - V_cache.data_ptr(), reinterpret_cast<__half*>(Out.data_ptr()), - softmax_lse.data_ptr(), block_table.data_ptr(), - seq_lens.data_ptr(), B, H, M, N, bfla_mask_ptr, bfla_mask_block_n, - bfla_mask_stride_b, bfla_mask_stride_h, bfla_mask_stride_q, - bfla_mask_stride_k, page_block_size, num_kv_heads, k_block_stride, - k_token_stride, k_head_stride, v_block_stride, v_token_stride, - v_head_stride, softmax_scale, k_scale, v_scale, window_size_left, - window_size_right, nullptr, nullptr, nullptr, 0, nullptr, 0); - } + const torch::Tensor& Q, + const torch::Tensor& K_cache, + const torch::Tensor& V_cache, + torch::Tensor& Out, + torch::Tensor& softmax_lse, + const torch::Tensor& block_table, + const torch::Tensor& seq_lens, + const int* bfla_mask_ptr, + int bfla_mask_block_n, + int64_t bfla_mask_stride_b, + int64_t bfla_mask_stride_h, + int64_t bfla_mask_stride_q, + int64_t bfla_mask_stride_k, + float softmax_scale, + bool is_causal, + float k_scale, + float v_scale, + 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 M = Q.size(2); + const int page_block_size = K_cache.size(1); + const int num_kv_heads = K_cache.size(2); + const int max_num_blocks = block_table.size(1); + const int N = max_num_blocks * page_block_size; + const int64_t k_block_stride = K_cache.stride(0); + const int64_t k_token_stride = K_cache.stride(1); + const int64_t k_head_stride = K_cache.stride(2); + const int64_t v_block_stride = V_cache.stride(0); + const int64_t v_token_stride = V_cache.stride(1); + const int64_t v_head_stride = V_cache.stride(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_paged< + D, LOW_SMEM, LOW_SMEM_CONTIG_FAST, + LOW_SMEM_SCALAR_QK, LOW_SMEM_BM32, false, true, + KV_DTYPE, D256_OUTPUT_STRIDE_268, + D256_SW_PIPELINE_QK, D256_SW_PIPELINE_PV> + : (void*)flash_attention_forward_kernel_paged< + D, LOW_SMEM, LOW_SMEM_CONTIG_FAST, + LOW_SMEM_SCALAR_QK, LOW_SMEM_BM32, false, false, + KV_DTYPE, D256_OUTPUT_STRIDE_268, + D256_SW_PIPELINE_QK, D256_SW_PIPELINE_PV>; + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + smem); + + if (is_causal) { + flash_attention_forward_kernel_paged< + D, LOW_SMEM, LOW_SMEM_CONTIG_FAST, LOW_SMEM_SCALAR_QK, + LOW_SMEM_BM32, false, true, KV_DTYPE, D256_OUTPUT_STRIDE_268, + D256_SW_PIPELINE_QK, D256_SW_PIPELINE_PV> + <<>>( + reinterpret_cast(Q.data_ptr()), + K_cache.data_ptr(), + V_cache.data_ptr(), + reinterpret_cast<__half*>(Out.data_ptr()), + softmax_lse.data_ptr(), + block_table.data_ptr(), + seq_lens.data_ptr(), + B, + H, + M, + N, + bfla_mask_ptr, + bfla_mask_block_n, + bfla_mask_stride_b, + bfla_mask_stride_h, + bfla_mask_stride_q, + bfla_mask_stride_k, + page_block_size, + num_kv_heads, + k_block_stride, + k_token_stride, + k_head_stride, + v_block_stride, + v_token_stride, + v_head_stride, + softmax_scale, + k_scale, + v_scale, + window_size_left, + window_size_right, + nullptr, + nullptr, + nullptr, + 0 + ); + } else { + flash_attention_forward_kernel_paged< + D, LOW_SMEM, LOW_SMEM_CONTIG_FAST, LOW_SMEM_SCALAR_QK, + LOW_SMEM_BM32, false, false, KV_DTYPE, D256_OUTPUT_STRIDE_268, + D256_SW_PIPELINE_QK, D256_SW_PIPELINE_PV> + <<>>( + reinterpret_cast(Q.data_ptr()), + K_cache.data_ptr(), + V_cache.data_ptr(), + reinterpret_cast<__half*>(Out.data_ptr()), + softmax_lse.data_ptr(), + block_table.data_ptr(), + seq_lens.data_ptr(), + B, + H, + M, + N, + bfla_mask_ptr, + bfla_mask_block_n, + bfla_mask_stride_b, + bfla_mask_stride_h, + bfla_mask_stride_q, + bfla_mask_stride_k, + page_block_size, + num_kv_heads, + k_block_stride, + k_token_stride, + k_head_stride, + v_block_stride, + v_token_stride, + v_head_stride, + softmax_scale, + k_scale, + v_scale, + window_size_left, + window_size_right, + nullptr, + nullptr, + nullptr, + 0 + ); + } } -template +template __global__ void flash_attention_forward_paged_splitkv_merge_kernel( const float* __restrict__ split_tmp_out, const float* __restrict__ split_tmp_row_max, - const float* __restrict__ split_tmp_row_sum, __half* __restrict__ Out, - float* __restrict__ softmax_lse, const int B, const int H, const int M, - const int num_partitions) { - static_assert(D == 256, "split-KV paged prefill merge is D=256 only"); - constexpr int BLOCK_M = BLOCK_M_256_LOW_SMEM; - constexpr int THREADS = 512; - 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; - const int valid_q_rows = min(BLOCK_M, M - start_row); - const int tid = threadIdx.x; - - __shared__ float sFinalMax[BLOCK_M]; - __shared__ float sFinalSum[BLOCK_M]; - - if (tid < valid_q_rows) { - const int row = start_row + tid; - float final_max = NEG_INF; -#pragma unroll 1 - for (int p = 0; p < num_partitions; ++p) { - const int64_t state_idx = - ((int64_t)batch_head_id * num_partitions + p) * M + row; - final_max = fmaxf(final_max, split_tmp_row_max[state_idx]); - } - - float final_sum = 0.0f; -#pragma unroll 1 - for (int p = 0; p < num_partitions; ++p) { - const int64_t state_idx = - ((int64_t)batch_head_id * num_partitions + p) * M + row; - const float part_sum = split_tmp_row_sum[state_idx]; - if (part_sum > 0.0f) { - final_sum += - __expf(fmaxf(split_tmp_row_max[state_idx] - final_max, -80.0f)) * - part_sum; - } - } - sFinalMax[tid] = final_max; - sFinalSum[tid] = final_sum; - softmax_lse[(int64_t)batch_head_id * M + row] = - final_max + logf(fmaxf(final_sum, 1e-24f)); - } - __syncthreads(); - - for (int idx = tid; idx < valid_q_rows * D; idx += THREADS) { - const int row_local = idx / D; - const int col = idx - row_local * D; - const int row = start_row + row_local; - const float final_max = sFinalMax[row_local]; - const float inv_sum = 1.0f / fmaxf(sFinalSum[row_local], 1e-24f); - - float acc = 0.0f; -#pragma unroll 1 - for (int p = 0; p < num_partitions; ++p) { - const int64_t state_idx = - ((int64_t)batch_head_id * num_partitions + p) * M + row; - const float part_sum = split_tmp_row_sum[state_idx]; - if (part_sum > 0.0f) { - const float scale = - __expf(fmaxf(split_tmp_row_max[state_idx] - final_max, -80.0f)); - const int64_t out_idx = state_idx * D + col; - acc = fmaf(scale, split_tmp_out[out_idx], acc); - } - } - Out[((int64_t)batch_head_id * M + row) * D + col] = - __float2half_rn(acc * inv_sum); - } + const float* __restrict__ split_tmp_row_sum, + __half* __restrict__ Out, + float* __restrict__ softmax_lse, + const int B, + const int H, + const int M, + const int num_partitions +) { + static_assert(D == 256, "split-KV paged prefill merge is D=256 only"); + constexpr int BLOCK_M = BLOCK_M_256_LOW_SMEM; + constexpr int THREADS = 512; + 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; + const int valid_q_rows = min(BLOCK_M, M - start_row); + const int tid = threadIdx.x; + + __shared__ float sFinalMax[BLOCK_M]; + __shared__ float sFinalSum[BLOCK_M]; + + if (tid < valid_q_rows) { + const int row = start_row + tid; + float final_max = NEG_INF; + #pragma unroll 1 + for (int p = 0; p < num_partitions; ++p) { + const int64_t state_idx = + ((int64_t)batch_head_id * num_partitions + p) * M + row; + final_max = fmaxf(final_max, split_tmp_row_max[state_idx]); + } + + float final_sum = 0.0f; + #pragma unroll 1 + for (int p = 0; p < num_partitions; ++p) { + const int64_t state_idx = + ((int64_t)batch_head_id * num_partitions + p) * M + row; + const float part_sum = split_tmp_row_sum[state_idx]; + if (part_sum > 0.0f) { + final_sum += + __expf(fmaxf(split_tmp_row_max[state_idx] - final_max, + -80.0f)) * part_sum; + } + } + sFinalMax[tid] = final_max; + sFinalSum[tid] = final_sum; + softmax_lse[(int64_t)batch_head_id * M + row] = + final_max + logf(fmaxf(final_sum, 1e-24f)); + } + __syncthreads(); + + for (int idx = tid; idx < valid_q_rows * D; idx += THREADS) { + const int row_local = idx / D; + const int col = idx - row_local * D; + const int row = start_row + row_local; + const float final_max = sFinalMax[row_local]; + const float inv_sum = 1.0f / fmaxf(sFinalSum[row_local], 1e-24f); + + float acc = 0.0f; + #pragma unroll 1 + for (int p = 0; p < num_partitions; ++p) { + const int64_t state_idx = + ((int64_t)batch_head_id * num_partitions + p) * M + row; + const float part_sum = split_tmp_row_sum[state_idx]; + if (part_sum > 0.0f) { + const float scale = + __expf(fmaxf(split_tmp_row_max[state_idx] - final_max, + -80.0f)); + const int64_t out_idx = state_idx * D + col; + acc = fmaf(scale, split_tmp_out[out_idx], acc); + } + } + Out[((int64_t)batch_head_id * M + row) * D + col] = + __float2half_rn(acc * inv_sum); + } } -template +template void launcher_flash_attention_forward_paged_splitkv_impl( - const torch::Tensor& Q, const torch::Tensor& K_cache, - const torch::Tensor& V_cache, torch::Tensor& Out, - torch::Tensor& softmax_lse, torch::Tensor& split_tmp_out, - torch::Tensor& split_tmp_row_max, torch::Tensor& split_tmp_row_sum, - const torch::Tensor& block_table, const torch::Tensor& seq_lens, - float softmax_scale, bool is_causal, float k_scale, float v_scale, - int window_size_left, int window_size_right, int split_kv_tokens, - int max_seq_len_hint, cudaStream_t stream) { - static_assert(D == 256, "split-KV paged prefill is D=256 only"); - static_assert(LOW_SMEM || !LOW_SMEM_CONTIG_FAST, - "contiguous low-smem loads require LOW_SMEM"); - static_assert(LOW_SMEM || !LOW_SMEM_SCALAR_QK, - "scalar low-smem QK requires LOW_SMEM"); - static_assert(!LOW_SMEM || KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP16, - "low-smem split-KV supports FP16 KV only"); - using Config = KernelConfig; - - const int B = Q.size(0); - const int H = Q.size(1); - const int M = Q.size(2); - const int page_block_size = K_cache.size(1); - const int num_kv_heads = K_cache.size(2); - const int max_num_blocks = block_table.size(1); - const int N = max_num_blocks * page_block_size; - const int64_t k_block_stride = K_cache.stride(0); - const int64_t k_token_stride = K_cache.stride(1); - const int64_t k_head_stride = K_cache.stride(2); - const int64_t v_block_stride = V_cache.stride(0); - const int64_t v_token_stride = V_cache.stride(1); - const int64_t v_head_stride = V_cache.stride(2); - - split_kv_tokens = std::max(split_kv_tokens, Config::BLOCK_N); - max_seq_len_hint = std::max(max_seq_len_hint, 1); - const int split_kv_tiles = - std::max(1, (split_kv_tokens + Config::BLOCK_N - 1) / Config::BLOCK_N); - const int max_kv_tiles = - std::max(1, (max_seq_len_hint + Config::BLOCK_N - 1) / Config::BLOCK_N); - const int num_partitions = - std::max(1, (max_kv_tiles + split_kv_tiles - 1) / split_kv_tiles); - - const int grid_x = (M + Config::BLOCK_M - 1) / Config::BLOCK_M; - const dim3 grid(grid_x, num_partitions, 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"); - TORCH_CHECK(split_tmp_out.size(2) == num_partitions, - "split_tmp_out partition mismatch"); - TORCH_CHECK(split_tmp_row_max.size(2) == num_partitions, - "split_tmp_row_max partition mismatch"); - TORCH_CHECK(split_tmp_row_sum.size(2) == num_partitions, - "split_tmp_row_sum partition mismatch"); - - auto kernel = is_causal - ? (void*)flash_attention_forward_kernel_paged< - D, LOW_SMEM, LOW_SMEM_CONTIG_FAST, LOW_SMEM_SCALAR_QK, - false, true, true, KV_DTYPE> - : (void*)flash_attention_forward_kernel_paged< - D, LOW_SMEM, LOW_SMEM_CONTIG_FAST, LOW_SMEM_SCALAR_QK, - false, true, false, KV_DTYPE>; - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, - smem); - - if (is_causal) { - flash_attention_forward_kernel_paged - <<>>( - reinterpret_cast(Q.data_ptr()), K_cache.data_ptr(), - V_cache.data_ptr(), reinterpret_cast<__half*>(Out.data_ptr()), - softmax_lse.data_ptr(), block_table.data_ptr(), - seq_lens.data_ptr(), B, H, M, N, nullptr, 0, 0, 0, 0, 0, - page_block_size, num_kv_heads, k_block_stride, k_token_stride, - k_head_stride, v_block_stride, v_token_stride, v_head_stride, - softmax_scale, k_scale, v_scale, window_size_left, - window_size_right, split_tmp_out.data_ptr(), + const torch::Tensor& Q, + const torch::Tensor& K_cache, + const torch::Tensor& V_cache, + torch::Tensor& Out, + torch::Tensor& softmax_lse, + torch::Tensor& split_tmp_out, + torch::Tensor& split_tmp_row_max, + torch::Tensor& split_tmp_row_sum, + const torch::Tensor& block_table, + const torch::Tensor& seq_lens, + float softmax_scale, + bool is_causal, + float k_scale, + float v_scale, + int window_size_left, + int window_size_right, + int split_kv_tokens, + int max_seq_len_hint, + cudaStream_t stream +) { + static_assert(D == 256, "split-KV paged prefill is D=256 only"); + using Config = KernelConfig; + constexpr int KV_DTYPE = flash_v100::KV_CACHE_DTYPE_FP16; + + const int B = Q.size(0); + const int H = Q.size(1); + const int M = Q.size(2); + const int page_block_size = K_cache.size(1); + const int num_kv_heads = K_cache.size(2); + const int max_num_blocks = block_table.size(1); + const int N = max_num_blocks * page_block_size; + const int64_t k_block_stride = K_cache.stride(0); + const int64_t k_token_stride = K_cache.stride(1); + const int64_t k_head_stride = K_cache.stride(2); + const int64_t v_block_stride = V_cache.stride(0); + const int64_t v_token_stride = V_cache.stride(1); + const int64_t v_head_stride = V_cache.stride(2); + + split_kv_tokens = std::max(split_kv_tokens, Config::BLOCK_N); + max_seq_len_hint = std::max(max_seq_len_hint, 1); + const int split_kv_tiles = + std::max(1, (split_kv_tokens + Config::BLOCK_N - 1) / Config::BLOCK_N); + const int max_kv_tiles = + std::max(1, (max_seq_len_hint + Config::BLOCK_N - 1) / Config::BLOCK_N); + const int num_partitions = + std::max(1, (max_kv_tiles + split_kv_tiles - 1) / split_kv_tiles); + + const int grid_x = (M + Config::BLOCK_M - 1) / Config::BLOCK_M; + const dim3 grid(grid_x, num_partitions, 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"); + TORCH_CHECK(split_tmp_out.size(2) == num_partitions, + "split_tmp_out partition mismatch"); + TORCH_CHECK(split_tmp_row_max.size(2) == num_partitions, + "split_tmp_row_max partition mismatch"); + TORCH_CHECK(split_tmp_row_sum.size(2) == num_partitions, + "split_tmp_row_sum partition mismatch"); + + auto kernel = is_causal + ? (void*)flash_attention_forward_kernel_paged< + D, true, LOW_SMEM_CONTIG_FAST, + LOW_SMEM_SCALAR_QK, false, true, true, KV_DTYPE> + : (void*)flash_attention_forward_kernel_paged< + D, true, LOW_SMEM_CONTIG_FAST, + LOW_SMEM_SCALAR_QK, false, true, false, KV_DTYPE>; + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + smem); + + if (is_causal) { + flash_attention_forward_kernel_paged< + D, true, LOW_SMEM_CONTIG_FAST, LOW_SMEM_SCALAR_QK, false, true, true, + KV_DTYPE> + <<>>( + reinterpret_cast(Q.data_ptr()), + K_cache.data_ptr(), + V_cache.data_ptr(), + reinterpret_cast<__half*>(Out.data_ptr()), + softmax_lse.data_ptr(), + block_table.data_ptr(), + seq_lens.data_ptr(), + B, + H, + M, + N, + nullptr, + 0, + 0, + 0, + 0, + 0, + page_block_size, + num_kv_heads, + k_block_stride, + k_token_stride, + k_head_stride, + v_block_stride, + v_token_stride, + v_head_stride, + softmax_scale, + k_scale, + v_scale, + window_size_left, + window_size_right, + split_tmp_out.data_ptr(), split_tmp_row_max.data_ptr(), - split_tmp_row_sum.data_ptr(), split_kv_tiles, nullptr, 0); - } else { - flash_attention_forward_kernel_paged - <<>>( - reinterpret_cast(Q.data_ptr()), K_cache.data_ptr(), - V_cache.data_ptr(), reinterpret_cast<__half*>(Out.data_ptr()), - softmax_lse.data_ptr(), block_table.data_ptr(), - seq_lens.data_ptr(), B, H, M, N, nullptr, 0, 0, 0, 0, 0, - page_block_size, num_kv_heads, k_block_stride, k_token_stride, - k_head_stride, v_block_stride, v_token_stride, v_head_stride, - softmax_scale, k_scale, v_scale, window_size_left, - window_size_right, split_tmp_out.data_ptr(), + split_tmp_row_sum.data_ptr(), + split_kv_tiles + ); + } else { + flash_attention_forward_kernel_paged< + D, true, LOW_SMEM_CONTIG_FAST, LOW_SMEM_SCALAR_QK, false, true, false, + KV_DTYPE> + <<>>( + reinterpret_cast(Q.data_ptr()), + K_cache.data_ptr(), + V_cache.data_ptr(), + reinterpret_cast<__half*>(Out.data_ptr()), + softmax_lse.data_ptr(), + block_table.data_ptr(), + seq_lens.data_ptr(), + B, + H, + M, + N, + nullptr, + 0, + 0, + 0, + 0, + 0, + page_block_size, + num_kv_heads, + k_block_stride, + k_token_stride, + k_head_stride, + v_block_stride, + v_token_stride, + v_head_stride, + softmax_scale, + k_scale, + v_scale, + window_size_left, + window_size_right, + split_tmp_out.data_ptr(), split_tmp_row_max.data_ptr(), - split_tmp_row_sum.data_ptr(), split_kv_tiles, nullptr, 0); - } - - const dim3 merge_grid(grid_x, 1, B * H); - const dim3 merge_block(512); - flash_attention_forward_paged_splitkv_merge_kernel - <<>>( - split_tmp_out.data_ptr(), split_tmp_row_max.data_ptr(), - split_tmp_row_sum.data_ptr(), - reinterpret_cast<__half*>(Out.data_ptr()), - softmax_lse.data_ptr(), B, H, M, num_partitions); + split_tmp_row_sum.data_ptr(), + split_kv_tiles + ); + } + + const dim3 merge_grid(grid_x, 1, B * H); + const dim3 merge_block(512); + flash_attention_forward_paged_splitkv_merge_kernel + <<>>( + split_tmp_out.data_ptr(), + split_tmp_row_max.data_ptr(), + split_tmp_row_sum.data_ptr(), + reinterpret_cast<__half*>(Out.data_ptr()), + softmax_lse.data_ptr(), + B, + H, + M, + num_partitions); } constexpr int D256_BM32_PHASE_BLOCK_M = 32; @@ -2100,41 +2331,41 @@ constexpr int D256_BM32_PHASE_PANELS = constexpr int D256_BM32_PHASE_PROBABILITY_ELEMENTS = D256_BM32_PHASE_PANEL_M * D256_BM32_PHASE_SOFTMAX_N; -template +template struct alignas(16) D256BM32PhaseSharedStorage { - __half query[D256_BM32_PHASE_BLOCK_M * D256_BM32_PHASE_D]; - float score[D256_BM32_PHASE_BLOCK_M * D256_BM32_PHASE_BLOCK_N]; - __half probability_top[PROBABILITY_PANELS * - D256_BM32_PHASE_PROBABILITY_ELEMENTS]; - __half probability_bottom[PROBABILITY_PANELS * - D256_BM32_PHASE_PROBABILITY_ELEMENTS]; - float row_max[D256_BM32_PHASE_BLOCK_M]; - float row_sum[D256_BM32_PHASE_BLOCK_M]; - float row_exp_diff[PROBABILITY_PANELS * D256_BM32_PHASE_BLOCK_M]; - int page_idx[D256_BM32_PHASE_PAGE_SLOTS]; - int page_offset[D256_BM32_PHASE_PAGE_SLOTS]; - uint64_t k_tile_ptr[D256_BM32_PHASE_PAGE_SLOTS]; - uint64_t v_tile_ptr[D256_BM32_PHASE_PAGE_SLOTS]; - int block_index; - int batch_id; - int kv_head_id; - int actual_n; + __half query[D256_BM32_PHASE_BLOCK_M * D256_BM32_PHASE_D]; + float score[D256_BM32_PHASE_BLOCK_M * D256_BM32_PHASE_BLOCK_N]; + __half probability_top[PROBABILITY_PANELS + * D256_BM32_PHASE_PROBABILITY_ELEMENTS]; + __half probability_bottom[PROBABILITY_PANELS + * D256_BM32_PHASE_PROBABILITY_ELEMENTS]; + float row_max[D256_BM32_PHASE_BLOCK_M]; + float row_sum[D256_BM32_PHASE_BLOCK_M]; + float row_exp_diff[PROBABILITY_PANELS * D256_BM32_PHASE_BLOCK_M]; + int page_idx[D256_BM32_PHASE_PAGE_SLOTS]; + int page_offset[D256_BM32_PHASE_PAGE_SLOTS]; + uint64_t k_tile_ptr[D256_BM32_PHASE_PAGE_SLOTS]; + uint64_t v_tile_ptr[D256_BM32_PHASE_PAGE_SLOTS]; + int block_index; + int batch_id; + int kv_head_id; + int actual_n; }; -template +template struct alignas(16) D256BM32PhaseSplitSharedStorage : D256BM32PhaseSharedStorage { - int split_end_block; + int split_end_block; }; -template +template struct D256BM32PhaseSharedStorageSelector { - using Type = D256BM32PhaseSharedStorage; + using Type = D256BM32PhaseSharedStorage; }; -template +template struct D256BM32PhaseSharedStorageSelector { - using Type = D256BM32PhaseSplitSharedStorage; + using Type = D256BM32PhaseSplitSharedStorage; }; using D256BM32PhaseSharedStorageSingle = D256BM32PhaseSharedStorage<1>; @@ -2147,8 +2378,8 @@ static_assert(sizeof(D256BM32PhaseSharedStorageAllP) == 41936, "D256 BM32 all-P shared layout changed unexpectedly"); static_assert(sizeof(D256BM32PhaseSharedStorageAllP) <= 48 * 1024, "D256 BM32 phase shared storage exceeds 48 KiB"); -static_assert(sizeof(D256BM32PhaseSplitSharedStorage) == - 41952, +static_assert(sizeof(D256BM32PhaseSplitSharedStorage< + D256_BM32_PHASE_PANELS>) == 41952, "D256 BM32 split shared layout changed unexpectedly"); using D256BM32PhaseMatrixAFragment = @@ -2161,1928 +2392,2143 @@ using D256BM32PhaseAccumulatorFragment = fragment; __device__ __forceinline__ int d256_bm32_phase_swizzled_row_slot(int row) { - return (row & 3) | ((row & 8) >> 1) | ((row & 4) << 1); + return (row & 3) | ((row & 8) >> 1) | ((row & 4) << 1); } -__device__ __forceinline__ int d256_bm32_phase_matrix_a_offset(int row, - int column) { - const int tile = column / WMMA_K; - const int within_tile = column - tile * WMMA_K; - const int plane = within_tile >> 3; - const int inner = within_tile & 7; - return tile * WMMA_M * WMMA_K + plane * WMMA_M * 8 + - d256_bm32_phase_swizzled_row_slot(row) * 8 + inner; +__device__ __forceinline__ int d256_bm32_phase_matrix_a_offset( + int row, + int column) { + const int tile = column / WMMA_K; + const int within_tile = column - tile * WMMA_K; + const int plane = within_tile >> 3; + const int inner = within_tile & 7; + return tile * WMMA_M * WMMA_K + plane * WMMA_M * 8 + + d256_bm32_phase_swizzled_row_slot(row) * 8 + inner; } __device__ __forceinline__ void d256_bm32_phase_stage_query( - const __half* __restrict__ source, __half* __restrict__ destination) { - constexpr int HALF_PER_UINT4 = sizeof(uint4) / sizeof(__half); - constexpr int VECTORS_PER_ROW = D256_BM32_PHASE_D / HALF_PER_UINT4; - constexpr int VECTORS_PER_PANEL = D256_BM32_PHASE_PANEL_M * VECTORS_PER_ROW; - constexpr int QUERY_VECTORS = D256_BM32_PHASE_BLOCK_M * VECTORS_PER_ROW; - - const uint4* source_vectors = reinterpret_cast(source); - uint4* destination_vectors = reinterpret_cast(destination); - -#pragma unroll - for (int index = threadIdx.x; index < QUERY_VECTORS; - index += D256_BM32_PHASE_THREADS) { - const int row = index / VECTORS_PER_ROW; - const int vector_column = index % VECTORS_PER_ROW; - const int panel = row / D256_BM32_PHASE_PANEL_M; - const int panel_row = row - panel * D256_BM32_PHASE_PANEL_M; - const int k_tile = vector_column >> 1; - const int plane = vector_column & 1; - const int slot = d256_bm32_phase_swizzled_row_slot(panel_row); - destination_vectors[panel * VECTORS_PER_PANEL + - k_tile * (2 * D256_BM32_PHASE_PANEL_M) + - plane * D256_BM32_PHASE_PANEL_M + slot] = - __ldg(source_vectors + index); - } + const __half* __restrict__ source, + __half* __restrict__ destination) { + constexpr int HALF_PER_UINT4 = sizeof(uint4) / sizeof(__half); + constexpr int VECTORS_PER_ROW = D256_BM32_PHASE_D / HALF_PER_UINT4; + constexpr int VECTORS_PER_PANEL = + D256_BM32_PHASE_PANEL_M * VECTORS_PER_ROW; + constexpr int QUERY_VECTORS = + D256_BM32_PHASE_BLOCK_M * VECTORS_PER_ROW; + + const uint4* source_vectors = reinterpret_cast(source); + uint4* destination_vectors = reinterpret_cast(destination); + + #pragma unroll + for (int index = threadIdx.x; index < QUERY_VECTORS; + index += D256_BM32_PHASE_THREADS) { + const int row = index / VECTORS_PER_ROW; + const int vector_column = index % VECTORS_PER_ROW; + const int panel = row / D256_BM32_PHASE_PANEL_M; + const int panel_row = row - panel * D256_BM32_PHASE_PANEL_M; + const int k_tile = vector_column >> 1; + const int plane = vector_column & 1; + const int slot = d256_bm32_phase_swizzled_row_slot(panel_row); + destination_vectors[panel * VECTORS_PER_PANEL + + k_tile * (2 * D256_BM32_PHASE_PANEL_M) + + plane * D256_BM32_PHASE_PANEL_M + slot] = + __ldg(source_vectors + index); + } } __device__ __forceinline__ void d256_bm32_phase_stage_query_partial( - const __half* __restrict__ source, __half* __restrict__ destination, + const __half* __restrict__ source, + __half* __restrict__ destination, int valid_rows) { - constexpr int HALF_PER_UINT4 = sizeof(uint4) / sizeof(__half); - constexpr int VECTORS_PER_ROW = D256_BM32_PHASE_D / HALF_PER_UINT4; - constexpr int VECTORS_PER_PANEL = D256_BM32_PHASE_PANEL_M * VECTORS_PER_ROW; - constexpr int QUERY_VECTORS = D256_BM32_PHASE_BLOCK_M * VECTORS_PER_ROW; - - const uint4* source_vectors = reinterpret_cast(source); - uint4* destination_vectors = reinterpret_cast(destination); - -#pragma unroll - for (int index = threadIdx.x; index < QUERY_VECTORS; - index += D256_BM32_PHASE_THREADS) { - const int row = index / VECTORS_PER_ROW; - const int vector_column = index % VECTORS_PER_ROW; - const int panel = row / D256_BM32_PHASE_PANEL_M; - const int panel_row = row - panel * D256_BM32_PHASE_PANEL_M; - const int k_tile = vector_column >> 1; - const int plane = vector_column & 1; - const int slot = d256_bm32_phase_swizzled_row_slot(panel_row); - uint4 value = {0, 0, 0, 0}; - if (row < valid_rows) { - value = __ldg(source_vectors + index); - } - destination_vectors[panel * VECTORS_PER_PANEL + - k_tile * (2 * D256_BM32_PHASE_PANEL_M) + - plane * D256_BM32_PHASE_PANEL_M + slot] = value; - } + constexpr int HALF_PER_UINT4 = sizeof(uint4) / sizeof(__half); + constexpr int VECTORS_PER_ROW = D256_BM32_PHASE_D / HALF_PER_UINT4; + constexpr int VECTORS_PER_PANEL = + D256_BM32_PHASE_PANEL_M * VECTORS_PER_ROW; + constexpr int QUERY_VECTORS = + D256_BM32_PHASE_BLOCK_M * VECTORS_PER_ROW; + + const uint4* source_vectors = reinterpret_cast(source); + uint4* destination_vectors = reinterpret_cast(destination); + + #pragma unroll + for (int index = threadIdx.x; index < QUERY_VECTORS; + index += D256_BM32_PHASE_THREADS) { + const int row = index / VECTORS_PER_ROW; + const int vector_column = index % VECTORS_PER_ROW; + const int panel = row / D256_BM32_PHASE_PANEL_M; + const int panel_row = row - panel * D256_BM32_PHASE_PANEL_M; + const int k_tile = vector_column >> 1; + const int plane = vector_column & 1; + const int slot = d256_bm32_phase_swizzled_row_slot(panel_row); + uint4 value = {0, 0, 0, 0}; + if (row < valid_rows) { + value = __ldg(source_vectors + index); + } + destination_vectors[panel * VECTORS_PER_PANEL + + k_tile * (2 * D256_BM32_PHASE_PANEL_M) + + plane * D256_BM32_PHASE_PANEL_M + slot] = value; + } } __device__ __forceinline__ void d256_bm32_phase_load_matrix_a( - D256BM32PhaseMatrixAFragment& fragment, const __half* __restrict__ matrix, + D256BM32PhaseMatrixAFragment& fragment, + const __half* __restrict__ matrix, int k_offset) { - const int lane = threadIdx.x & 31; - const int row = (lane & 3) + ((lane >> 4) & 1) * 4 + ((lane >> 2) & 1) * 8; - const int slot = d256_bm32_phase_swizzled_row_slot(row); - const int tile_offset = (k_offset / WMMA_K) * WMMA_M * WMMA_K; - uint32_t address = static_cast( - __cvta_generic_to_shared(matrix + tile_offset + slot * 8)); - uint32_t* words = reinterpret_cast(&fragment); - asm volatile("ld.shared.v4.u32 {%0, %1, %2, %3}, [%4];" - : "=r"(words[0]), "=r"(words[1]), "=r"(words[2]), "=r"(words[3]) - : "r"(address) - : "memory"); - address += WMMA_M * 8 * sizeof(__half); - asm volatile("ld.shared.v4.u32 {%0, %1, %2, %3}, [%4];" - : "=r"(words[4]), "=r"(words[5]), "=r"(words[6]), "=r"(words[7]) - : "r"(address) - : "memory"); + const int lane = threadIdx.x & 31; + const int row = + (lane & 3) + ((lane >> 4) & 1) * 4 + ((lane >> 2) & 1) * 8; + const int slot = d256_bm32_phase_swizzled_row_slot(row); + const int tile_offset = (k_offset / WMMA_K) * WMMA_M * WMMA_K; + uint32_t address = static_cast( + __cvta_generic_to_shared(matrix + tile_offset + slot * 8)); + uint32_t* words = reinterpret_cast(&fragment); + asm volatile( + "ld.shared.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(words[0]), "=r"(words[1]), "=r"(words[2]), "=r"(words[3]) + : "r"(address) + : "memory"); + address += WMMA_M * 8 * sizeof(__half); + asm volatile( + "ld.shared.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(words[4]), "=r"(words[5]), "=r"(words[6]), "=r"(words[7]) + : "r"(address) + : "memory"); } -__device__ __forceinline__ int d256_bm32_phase_accumulator_row(int lane, - int element) { - const int row_base = - (lane & 1) + ((lane >> 2) & 1) * 8 + ((lane >> 4) & 1) * 4; - return row_base + ((element >> 1) & 1) * 2; +__device__ __forceinline__ int d256_bm32_phase_accumulator_row( + int lane, + int element) { + const int row_base = + (lane & 1) + ((lane >> 2) & 1) * 8 + ((lane >> 4) & 1) * 4; + return row_base + ((element >> 1) & 1) * 2; } -__device__ __forceinline__ int d256_bm32_phase_accumulator_column(int lane, - int element) { - const int column_base = ((lane >> 1) & 1) * 2 + ((lane >> 3) & 1) * 8; - return column_base + (element & 1) + ((element >> 2) & 1) * 4; +__device__ __forceinline__ int d256_bm32_phase_accumulator_column( + int lane, + int element) { + const int column_base = ((lane >> 1) & 1) * 2 + ((lane >> 3) & 1) * 8; + return column_base + (element & 1) + ((element >> 2) & 1) * 4; } __device__ __forceinline__ void d256_bm32_phase_spill_pair_scratch( float* __restrict__ score, D256BM32PhaseAccumulatorFragment& accumulator_top, D256BM32PhaseAccumulatorFragment& accumulator_bottom) { - const int warp = threadIdx.x >> 5; - const int lane = threadIdx.x & 31; - const int warp_pair = warp >> 1; - const int warp_in_pair = warp & 1; - const int pair_column = warp_pair * 32 + (lane & 15) * 2; - const int lane_row = lane >> 4; - -#pragma unroll - for (int element_pair = 0; element_pair < 4; ++element_pair) { - const int element = element_pair * 2; - const int offset = (warp_in_pair * 16 + element_pair * 2 + lane_row) * - D256_BM32_PHASE_BLOCK_N + - pair_column; - const uint32_t address = - static_cast(__cvta_generic_to_shared(score + offset)); - asm volatile("st.shared.v2.u32 [%0], {%1, %2};" - : - : "r"(address), - "r"(__float_as_uint(accumulator_top.x[element])), - "r"(__float_as_uint(accumulator_top.x[element + 1])) - : "memory"); - asm volatile("st.shared.v2.u32 [%0+4096], {%1, %2};" - : - : "r"(address), - "r"(__float_as_uint(accumulator_bottom.x[element])), - "r"(__float_as_uint(accumulator_bottom.x[element + 1])) - : "memory"); - } + const int warp = threadIdx.x >> 5; + const int lane = threadIdx.x & 31; + const int warp_pair = warp >> 1; + const int warp_in_pair = warp & 1; + const int pair_column = warp_pair * 32 + (lane & 15) * 2; + const int lane_row = lane >> 4; + + #pragma unroll + for (int element_pair = 0; element_pair < 4; ++element_pair) { + const int element = element_pair * 2; + const int offset = + (warp_in_pair * 16 + element_pair * 2 + lane_row) + * D256_BM32_PHASE_BLOCK_N + + pair_column; + const uint32_t address = static_cast( + __cvta_generic_to_shared(score + offset)); + asm volatile( + "st.shared.v2.u32 [%0], {%1, %2};" + : + : "r"(address), + "r"(__float_as_uint(accumulator_top.x[element])), + "r"(__float_as_uint(accumulator_top.x[element + 1])) + : "memory"); + asm volatile( + "st.shared.v2.u32 [%0+4096], {%1, %2};" + : + : "r"(address), + "r"(__float_as_uint(accumulator_bottom.x[element])), + "r"(__float_as_uint(accumulator_bottom.x[element + 1])) + : "memory"); + } } __device__ __forceinline__ void d256_bm32_phase_reload_pair_scratch( const float* __restrict__ score, D256BM32PhaseAccumulatorFragment& accumulator_top, D256BM32PhaseAccumulatorFragment& accumulator_bottom) { - const int warp = threadIdx.x >> 5; - const int lane = threadIdx.x & 31; - const int warp_pair = warp >> 1; - const int warp_in_pair = warp & 1; - const int pair_column = warp_pair * 32 + (lane & 15) * 2; - const int lane_row = lane >> 4; - -#pragma unroll - for (int element_pair = 0; element_pair < 4; ++element_pair) { - const int element = element_pair * 2; - const int offset = (warp_in_pair * 16 + element_pair * 2 + lane_row) * - D256_BM32_PHASE_BLOCK_N + - pair_column; - const uint32_t address = - static_cast(__cvta_generic_to_shared(score + offset)); - uint32_t first_word; - uint32_t second_word; - asm volatile("ld.shared.v2.u32 {%0, %1}, [%2];" - : "=r"(first_word), "=r"(second_word) - : "r"(address) - : "memory"); - accumulator_top.x[element] = __uint_as_float(first_word); - accumulator_top.x[element + 1] = __uint_as_float(second_word); - asm volatile("ld.shared.v2.u32 {%0, %1}, [%2+4096];" - : "=r"(first_word), "=r"(second_word) - : "r"(address) - : "memory"); - accumulator_bottom.x[element] = __uint_as_float(first_word); - accumulator_bottom.x[element + 1] = __uint_as_float(second_word); - } + const int warp = threadIdx.x >> 5; + const int lane = threadIdx.x & 31; + const int warp_pair = warp >> 1; + const int warp_in_pair = warp & 1; + const int pair_column = warp_pair * 32 + (lane & 15) * 2; + const int lane_row = lane >> 4; + + #pragma unroll + for (int element_pair = 0; element_pair < 4; ++element_pair) { + const int element = element_pair * 2; + const int offset = + (warp_in_pair * 16 + element_pair * 2 + lane_row) + * D256_BM32_PHASE_BLOCK_N + + pair_column; + const uint32_t address = static_cast( + __cvta_generic_to_shared(score + offset)); + uint32_t first_word; + uint32_t second_word; + asm volatile( + "ld.shared.v2.u32 {%0, %1}, [%2];" + : "=r"(first_word), "=r"(second_word) + : "r"(address) + : "memory"); + accumulator_top.x[element] = __uint_as_float(first_word); + accumulator_top.x[element + 1] = __uint_as_float(second_word); + asm volatile( + "ld.shared.v2.u32 {%0, %1}, [%2+4096];" + : "=r"(first_word), "=r"(second_word) + : "r"(address) + : "memory"); + accumulator_bottom.x[element] = __uint_as_float(first_word); + accumulator_bottom.x[element + 1] = __uint_as_float(second_word); + } } __device__ __forceinline__ void d256_bm32_phase_sync_warp_pair(int warp_pair) { - const int barrier_id = warp_pair + 1; - asm volatile("bar.sync %0, 64;" : : "r"(barrier_id) : "memory"); + const int barrier_id = warp_pair + 1; + asm volatile("bar.sync %0, 64;" : : "r"(barrier_id) : "memory"); } __device__ __forceinline__ float d256_bm32_phase_make_probability_row( - const float* __restrict__ score_row, __half* __restrict__ probability, - int probability_row, float* __restrict__ row_max, - float* __restrict__ row_sum, int state_row, int panel, float neg_inf) { - const int lane = threadIdx.x & 31; - const float* panel_score = score_row + panel * D256_BM32_PHASE_SOFTMAX_N; - float thread_max = neg_inf; - if (lane < D256_BM32_PHASE_SOFTMAX_N / 4) { - const float4 values = reinterpret_cast(panel_score)[lane]; - thread_max = fmaxf(fmaxf(values.x, values.y), fmaxf(values.z, values.w)); - } -#pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) { - thread_max = - fmaxf(thread_max, __shfl_down_sync(0xffffffffU, thread_max, offset)); - } - const float panel_max = __shfl_sync(0xffffffffU, thread_max, 0); - const float old_max = row_max[state_row]; - const float new_max = fmaxf(old_max, panel_max); - const float exp_diff = __expf(old_max - new_max); - - float thread_sum = 0.0f; - if (lane < D256_BM32_PHASE_SOFTMAX_N / 4) { - float4 values = reinterpret_cast(panel_score)[lane]; - values.x = __expf(fmaxf(values.x - new_max, -80.0f)); - values.y = __expf(fmaxf(values.y - new_max, -80.0f)); - values.z = __expf(fmaxf(values.z - new_max, -80.0f)); - values.w = __expf(fmaxf(values.w - new_max, -80.0f)); - thread_sum = (values.x + values.y) + (values.z + values.w); - const int column = lane * 4; - __half2* first_pair = reinterpret_cast<__half2*>( - probability + d256_bm32_phase_matrix_a_offset(probability_row, column)); - __half2* second_pair = reinterpret_cast<__half2*>( - probability + - d256_bm32_phase_matrix_a_offset(probability_row, column + 2)); - *first_pair = __float22half2_rn(make_float2(values.x, values.y)); - *second_pair = __float22half2_rn(make_float2(values.z, values.w)); - } -#pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) { - thread_sum += __shfl_down_sync(0xffffffffU, thread_sum, offset); - } - const float panel_sum = __shfl_sync(0xffffffffU, thread_sum, 0); - if (lane == 0) { - row_sum[state_row] = exp_diff * row_sum[state_row] + panel_sum; - row_max[state_row] = new_max; - } - return exp_diff; + const float* __restrict__ score_row, + __half* __restrict__ probability, + int probability_row, + float* __restrict__ row_max, + float* __restrict__ row_sum, + int state_row, + int panel, + float neg_inf) { + const int lane = threadIdx.x & 31; + const float* panel_score = + score_row + panel * D256_BM32_PHASE_SOFTMAX_N; + float thread_max = neg_inf; + if (lane < D256_BM32_PHASE_SOFTMAX_N / 4) { + const float4 values = + reinterpret_cast(panel_score)[lane]; + thread_max = fmaxf(fmaxf(values.x, values.y), + fmaxf(values.z, values.w)); + } + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + thread_max = fmaxf( + thread_max, __shfl_down_sync(0xffffffffU, thread_max, offset)); + } + const float panel_max = __shfl_sync(0xffffffffU, thread_max, 0); + const float old_max = row_max[state_row]; + const float new_max = fmaxf(old_max, panel_max); + const float exp_diff = __expf(old_max - new_max); + + float thread_sum = 0.0f; + if (lane < D256_BM32_PHASE_SOFTMAX_N / 4) { + float4 values = reinterpret_cast(panel_score)[lane]; + values.x = __expf(fmaxf(values.x - new_max, -80.0f)); + values.y = __expf(fmaxf(values.y - new_max, -80.0f)); + values.z = __expf(fmaxf(values.z - new_max, -80.0f)); + values.w = __expf(fmaxf(values.w - new_max, -80.0f)); + thread_sum = (values.x + values.y) + (values.z + values.w); + const int column = lane * 4; + __half2* first_pair = reinterpret_cast<__half2*>( + probability + d256_bm32_phase_matrix_a_offset( + probability_row, column)); + __half2* second_pair = reinterpret_cast<__half2*>( + probability + d256_bm32_phase_matrix_a_offset( + probability_row, column + 2)); + *first_pair = __float22half2_rn(make_float2(values.x, values.y)); + *second_pair = __float22half2_rn(make_float2(values.z, values.w)); + } + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + thread_sum += + __shfl_down_sync(0xffffffffU, thread_sum, offset); + } + const float panel_sum = __shfl_sync(0xffffffffU, thread_sum, 0); + if (lane == 0) { + row_sum[state_row] = exp_diff * row_sum[state_row] + panel_sum; + row_max[state_row] = new_max; + } + return exp_diff; } -__device__ __forceinline__ float d256_bm32_phase_make_split_probability_row( - const float* __restrict__ score_row, __half* __restrict__ probability, - int probability_row, float* __restrict__ row_max, - float* __restrict__ row_sum, int state_row, int panel, float neg_inf, +__device__ __forceinline__ float +d256_bm32_phase_make_split_probability_row( + const float* __restrict__ score_row, + __half* __restrict__ probability, + int probability_row, + float* __restrict__ row_max, + float* __restrict__ row_sum, + int state_row, + int panel, + float neg_inf, bool has_visible_value) { - if (!has_visible_value) { - const int lane = threadIdx.x & 31; - if (lane < D256_BM32_PHASE_SOFTMAX_N / 4) { - const int column = lane * 4; - __half2* first_pair = reinterpret_cast<__half2*>( - probability + - d256_bm32_phase_matrix_a_offset(probability_row, column)); - __half2* second_pair = reinterpret_cast<__half2*>( - probability + - d256_bm32_phase_matrix_a_offset(probability_row, column + 2)); - const __half2 zero = __float2half2_rn(0.0f); - *first_pair = zero; - *second_pair = zero; - } - return 1.0f; - } - return d256_bm32_phase_make_probability_row(score_row, probability, - probability_row, row_max, row_sum, - state_row, panel, neg_inf); + if (!has_visible_value) { + const int lane = threadIdx.x & 31; + if (lane < D256_BM32_PHASE_SOFTMAX_N / 4) { + const int column = lane * 4; + __half2* first_pair = reinterpret_cast<__half2*>( + probability + d256_bm32_phase_matrix_a_offset( + probability_row, column)); + __half2* second_pair = reinterpret_cast<__half2*>( + probability + d256_bm32_phase_matrix_a_offset( + probability_row, column + 2)); + const __half2 zero = __float2half2_rn(0.0f); + *first_pair = zero; + *second_pair = zero; + } + return 1.0f; + } + return d256_bm32_phase_make_probability_row( + score_row, probability, probability_row, row_max, row_sum, + state_row, panel, neg_inf); } __device__ __forceinline__ void d256_bm32_phase_scale_accumulator_rows( - D256BM32PhaseAccumulatorFragment& accumulator, float first_row_scale, + D256BM32PhaseAccumulatorFragment& accumulator, + float first_row_scale, float second_row_scale) { - accumulator.x[0] *= first_row_scale; - accumulator.x[1] *= first_row_scale; - accumulator.x[2] *= second_row_scale; - accumulator.x[3] *= second_row_scale; - accumulator.x[4] *= first_row_scale; - accumulator.x[5] *= first_row_scale; - accumulator.x[6] *= second_row_scale; - accumulator.x[7] *= second_row_scale; + accumulator.x[0] *= first_row_scale; + accumulator.x[1] *= first_row_scale; + accumulator.x[2] *= second_row_scale; + accumulator.x[3] *= second_row_scale; + accumulator.x[4] *= first_row_scale; + accumulator.x[5] *= first_row_scale; + accumulator.x[6] *= second_row_scale; + accumulator.x[7] *= second_row_scale; } __device__ __forceinline__ void d256_bm32_phase_scale_accumulators( D256BM32PhaseAccumulatorFragment& accumulator_top, D256BM32PhaseAccumulatorFragment& accumulator_bottom, const float* __restrict__ row_exp_diff) { - const int row = d256_bm32_phase_accumulator_row(threadIdx.x & 31, 0); - d256_bm32_phase_scale_accumulator_rows(accumulator_top, row_exp_diff[row], - row_exp_diff[row + 2]); - d256_bm32_phase_scale_accumulator_rows( - accumulator_bottom, row_exp_diff[D256_BM32_PHASE_PANEL_M + row], - row_exp_diff[D256_BM32_PHASE_PANEL_M + row + 2]); + const int row = d256_bm32_phase_accumulator_row(threadIdx.x & 31, 0); + d256_bm32_phase_scale_accumulator_rows( + accumulator_top, row_exp_diff[row], row_exp_diff[row + 2]); + d256_bm32_phase_scale_accumulator_rows( + accumulator_bottom, + row_exp_diff[D256_BM32_PHASE_PANEL_M + row], + row_exp_diff[D256_BM32_PHASE_PANEL_M + row + 2]); } __device__ __forceinline__ void d256_bm32_phase_update_pv_panel( const __half* __restrict__ probability_top, const __half* __restrict__ probability_bottom, - const __half* __restrict__ value_k0, const __half* __restrict__ value_k16, - int64_t value_token_stride, int d_offset, int valid_columns, + const __half* __restrict__ value_k0, + const __half* __restrict__ value_k16, + int64_t value_token_stride, + int d_offset, + int valid_columns, D256BM32PhaseAccumulatorFragment& accumulator_top, D256BM32PhaseAccumulatorFragment& accumulator_bottom) { - D256BM32PhaseMatrixAFragment a_fragment; - D256BM32PhasePVMatrixBFragment b_fragment; - - load_matrix_sync(b_fragment, value_k0 + d_offset, value_token_stride); - d256_bm32_phase_load_matrix_a(a_fragment, probability_top, 0); - mma_sync(accumulator_top, a_fragment, b_fragment, accumulator_top); - d256_bm32_phase_load_matrix_a(a_fragment, probability_bottom, 0); - mma_sync(accumulator_bottom, a_fragment, b_fragment, accumulator_bottom); - if (valid_columns > WMMA_K) { - load_matrix_sync(b_fragment, value_k16 + d_offset, value_token_stride); - d256_bm32_phase_load_matrix_a(a_fragment, probability_top, WMMA_K); + D256BM32PhaseMatrixAFragment a_fragment; + D256BM32PhasePVMatrixBFragment b_fragment; + + load_matrix_sync(b_fragment, value_k0 + d_offset, value_token_stride); + d256_bm32_phase_load_matrix_a(a_fragment, probability_top, 0); mma_sync(accumulator_top, a_fragment, b_fragment, accumulator_top); - d256_bm32_phase_load_matrix_a(a_fragment, probability_bottom, WMMA_K); + d256_bm32_phase_load_matrix_a(a_fragment, probability_bottom, 0); mma_sync(accumulator_bottom, a_fragment, b_fragment, accumulator_bottom); - } + if (valid_columns > WMMA_K) { + load_matrix_sync( + b_fragment, value_k16 + d_offset, value_token_stride); + d256_bm32_phase_load_matrix_a(a_fragment, probability_top, WMMA_K); + mma_sync(accumulator_top, a_fragment, b_fragment, accumulator_top); + d256_bm32_phase_load_matrix_a(a_fragment, probability_bottom, WMMA_K); + mma_sync(accumulator_bottom, a_fragment, b_fragment, + accumulator_bottom); + } } __device__ __forceinline__ void d256_bm32_phase_store_output( const D256BM32PhaseAccumulatorFragment& accumulator, - __half* __restrict__ output, const float* __restrict__ row_sum, - int row_offset, int d_offset) { - const int lane = threadIdx.x & 31; -#pragma unroll - for (int element = 0; element < accumulator.num_elements; ++element) { - const int row = row_offset + d256_bm32_phase_accumulator_row(lane, element); - const int column = - d_offset + d256_bm32_phase_accumulator_column(lane, element); - const float inverse_sum = 1.0f / fmaxf(row_sum[row], 1e-24f); - output[row * D256_BM32_PHASE_D + column] = - __float2half_rn(accumulator.x[element] * inverse_sum); - } + __half* __restrict__ output, + const float* __restrict__ row_sum, + int row_offset, + int d_offset) { + const int lane = threadIdx.x & 31; + #pragma unroll + for (int element = 0; element < accumulator.num_elements; ++element) { + const int row = row_offset + + d256_bm32_phase_accumulator_row(lane, element); + const int column = d_offset + + d256_bm32_phase_accumulator_column(lane, element); + const float inverse_sum = 1.0f / fmaxf(row_sum[row], 1e-24f); + output[row * D256_BM32_PHASE_D + column] = + __float2half_rn(accumulator.x[element] * inverse_sum); + } } __device__ __forceinline__ void d256_bm32_phase_store_output_partial( const D256BM32PhaseAccumulatorFragment& accumulator, - __half* __restrict__ output, const float* __restrict__ row_sum, - int row_offset, int d_offset, int valid_rows) { - const int lane = threadIdx.x & 31; -#pragma unroll - for (int element = 0; element < accumulator.num_elements; ++element) { - const int row = row_offset + d256_bm32_phase_accumulator_row(lane, element); - if (row < valid_rows) { - const int column = - d_offset + d256_bm32_phase_accumulator_column(lane, element); - const float inverse_sum = 1.0f / fmaxf(row_sum[row], 1e-24f); - output[row * D256_BM32_PHASE_D + column] = - __float2half_rn(accumulator.x[element] * inverse_sum); - } - } + __half* __restrict__ output, + const float* __restrict__ row_sum, + int row_offset, + int d_offset, + int valid_rows) { + const int lane = threadIdx.x & 31; + #pragma unroll + for (int element = 0; element < accumulator.num_elements; ++element) { + const int row = row_offset + + d256_bm32_phase_accumulator_row(lane, element); + if (row < valid_rows) { + const int column = + d_offset + + d256_bm32_phase_accumulator_column(lane, element); + const float inverse_sum = 1.0f / fmaxf(row_sum[row], 1e-24f); + output[row * D256_BM32_PHASE_D + column] = + __float2half_rn(accumulator.x[element] * inverse_sum); + } + } } __device__ __forceinline__ void d256_bm32_phase_store_unnormalized_partial( const D256BM32PhaseAccumulatorFragment& accumulator, - float* __restrict__ output, int row_offset, int d_offset, int valid_rows) { - const int lane = threadIdx.x & 31; -#pragma unroll - for (int element = 0; element < accumulator.num_elements; ++element) { - const int row = row_offset + d256_bm32_phase_accumulator_row(lane, element); - if (row < valid_rows) { - const int column = - d_offset + d256_bm32_phase_accumulator_column(lane, element); - output[row * D256_BM32_PHASE_D + column] = accumulator.x[element]; - } - } + float* __restrict__ output, + int row_offset, + int d_offset, + int valid_rows) { + const int lane = threadIdx.x & 31; + #pragma unroll + for (int element = 0; element < accumulator.num_elements; ++element) { + const int row = row_offset + + d256_bm32_phase_accumulator_row(lane, element); + if (row < valid_rows) { + const int column = + d_offset + + d256_bm32_phase_accumulator_column(lane, element); + output[row * D256_BM32_PHASE_D + column] = accumulator.x[element]; + } + } } -template -__device__ __forceinline__ void -flash_attention_forward_paged_d256_bm32_phase_body( - const __half* __restrict__ Q, const __half* __restrict__ K_cache, - const __half* __restrict__ V_cache, __half* __restrict__ Out, - float* __restrict__ softmax_lse, const int* __restrict__ block_table, - const int* __restrict__ seqused_k, int H, int M, int max_num_blocks_per_seq, - int num_kv_heads, int64_t k_block_stride, int64_t k_token_stride, - int64_t k_head_stride, int64_t v_block_stride, int64_t v_token_stride, - int64_t v_head_stride, float softmax_scale, - float* __restrict__ split_tmp_out, float* __restrict__ split_tmp_row_max, - float* __restrict__ split_tmp_row_sum, int split_actual_n) { - constexpr float NEG_INF = -1e30f; - static_assert(!SPLIT_KV || (ALL_P && PAIR_SCRATCH), - "BM32 split-KV requires the accepted all-P pair-scratch body"); - __shared__ __align__(16) typename D256BM32PhaseSharedStorageSelector< - SPLIT_KV, ALL_P ? D256_BM32_PHASE_PANELS : 1>::Type shared; - - const int start_row = blockIdx.x * D256_BM32_PHASE_BLOCK_M; - const int tid = threadIdx.x; - const int warp_id = tid >> 5; - const int lane_id = tid & 31; - const int valid_q_rows = min(D256_BM32_PHASE_BLOCK_M, M - start_row); - - { - const __half* q_ptr = - Q + static_cast(blockIdx.z) * M * D256_BM32_PHASE_D + - start_row * D256_BM32_PHASE_D; - if (valid_q_rows == D256_BM32_PHASE_BLOCK_M) { - d256_bm32_phase_stage_query(q_ptr, shared.query); - } else { - d256_bm32_phase_stage_query_partial(q_ptr, shared.query, valid_q_rows); - } - } - if (tid < D256_BM32_PHASE_BLOCK_M) { - shared.row_max[tid] = NEG_INF; - shared.row_sum[tid] = 0.0f; - } - if (tid == 0) { - shared.block_index = 0; - shared.batch_id = blockIdx.z / H; - shared.kv_head_id = (blockIdx.z % H) / (H / num_kv_heads); - if constexpr (SPLIT_KV) { - shared.actual_n = split_actual_n; - } else { - shared.actual_n = seqused_k[shared.batch_id]; +template +__device__ __forceinline__ +void flash_attention_forward_paged_d256_bm32_phase_body( + const __half* __restrict__ Q, + const __half* __restrict__ K_cache, + const __half* __restrict__ V_cache, + __half* __restrict__ Out, + float* __restrict__ softmax_lse, + const int* __restrict__ block_table, + const int* __restrict__ seqused_k, + int H, + int M, + int max_num_blocks_per_seq, + int num_kv_heads, + int64_t k_block_stride, + int64_t k_token_stride, + int64_t k_head_stride, + int64_t v_block_stride, + int64_t v_token_stride, + int64_t v_head_stride, + float softmax_scale, + float* __restrict__ split_tmp_out, + float* __restrict__ split_tmp_row_max, + float* __restrict__ split_tmp_row_sum, + int split_actual_n) { + constexpr float NEG_INF = -1e30f; + static_assert(!SPLIT_KV || (ALL_P && PAIR_SCRATCH), + "BM32 split-KV requires the accepted all-P pair-scratch body"); + __shared__ __align__(16) + typename D256BM32PhaseSharedStorageSelector< + SPLIT_KV, + ALL_P ? D256_BM32_PHASE_PANELS : 1>::Type shared; + + const int start_row = blockIdx.x * D256_BM32_PHASE_BLOCK_M; + const int tid = threadIdx.x; + const int warp_id = tid >> 5; + const int lane_id = tid & 31; + const int valid_q_rows = + min(D256_BM32_PHASE_BLOCK_M, M - start_row); + + { + const __half* q_ptr = + Q + static_cast(blockIdx.z) * M * D256_BM32_PHASE_D + + start_row * D256_BM32_PHASE_D; + if (valid_q_rows == D256_BM32_PHASE_BLOCK_M) { + d256_bm32_phase_stage_query(q_ptr, shared.query); + } else { + d256_bm32_phase_stage_query_partial( + q_ptr, shared.query, valid_q_rows); + } + } + if (tid < D256_BM32_PHASE_BLOCK_M) { + shared.row_max[tid] = NEG_INF; + shared.row_sum[tid] = 0.0f; } - } - __syncthreads(); - - const int total_n_blocks = - (shared.actual_n + D256_BM32_PHASE_BLOCK_N - 1) / D256_BM32_PHASE_BLOCK_N; - if constexpr (SPLIT_KV) { if (tid == 0) { - shared.block_index = static_cast(blockIdx.y) * total_n_blocks / - D256_BM32_PHASE_SPLIT_PARTS; - shared.split_end_block = (static_cast(blockIdx.y) + 1) * - total_n_blocks / D256_BM32_PHASE_SPLIT_PARTS; + shared.block_index = 0; + shared.batch_id = blockIdx.z / H; + shared.kv_head_id = + (blockIdx.z % H) / (H / num_kv_heads); + if constexpr (SPLIT_KV) { + shared.actual_n = split_actual_n; + } else { + shared.actual_n = seqused_k[shared.batch_id]; + } } __syncthreads(); - } - D256BM32PhaseAccumulatorFragment accumulator_top; - D256BM32PhaseAccumulatorFragment accumulator_bottom; - fill_fragment(accumulator_top, 0.0f); - fill_fragment(accumulator_bottom, 0.0f); - - for (;;) { + const int total_n_blocks = + (shared.actual_n + D256_BM32_PHASE_BLOCK_N - 1) + / D256_BM32_PHASE_BLOCK_N; if constexpr (SPLIT_KV) { - if (shared.block_index >= shared.split_end_block) { - break; - } - } else { - if (shared.block_index >= total_n_blocks) { - break; - } - } - const int start_col = shared.block_index * D256_BM32_PHASE_BLOCK_N; - const int valid_k_rows = - min(D256_BM32_PHASE_BLOCK_N, shared.actual_n - start_col); - - if (tid < D256_BM32_PHASE_PAGE_SLOTS) { - const int token_offset = tid * D256_BM32_PHASE_PAGE_SIZE; - if (token_offset < valid_k_rows) { - const int global_token_idx = start_col + token_offset; - const int virtual_block_idx = - global_token_idx / D256_BM32_PHASE_PAGE_BLOCK_SIZE; - shared.page_idx[tid] = - __ldg(&block_table[shared.batch_id * max_num_blocks_per_seq + - virtual_block_idx]); - shared.page_offset[tid] = - global_token_idx - - virtual_block_idx * D256_BM32_PHASE_PAGE_BLOCK_SIZE; - shared.k_tile_ptr[tid] = reinterpret_cast( - K_cache + (int64_t)shared.page_idx[tid] * k_block_stride + - (int64_t)shared.page_offset[tid] * k_token_stride + - (int64_t)shared.kv_head_id * k_head_stride); - shared.v_tile_ptr[tid] = reinterpret_cast( - V_cache + (int64_t)shared.page_idx[tid] * v_block_stride + - (int64_t)shared.page_offset[tid] * v_token_stride + - (int64_t)shared.kv_head_id * v_head_stride); - } else { - shared.page_idx[tid] = -1; - shared.page_offset[tid] = 0; - shared.k_tile_ptr[tid] = 0; - shared.v_tile_ptr[tid] = 0; - } - } - __syncthreads(); - - if (warp_id < D256_BM32_PHASE_BLOCK_N / WMMA_N && - warp_id * WMMA_N < valid_k_rows) { - const int tile_n = warp_id * WMMA_N; - const __half* k_tile = - reinterpret_cast(shared.k_tile_ptr[warp_id]); - if constexpr (PAIR_SCRATCH) { - d256_bm32_phase_spill_pair_scratch(shared.score, accumulator_top, - accumulator_bottom); - } else { - store_matrix_sync(shared.score + tile_n, accumulator_top, - D256_BM32_PHASE_BLOCK_N, mem_row_major); - store_matrix_sync( - shared.score + D256_BM32_PHASE_PANEL_M * D256_BM32_PHASE_BLOCK_N + - tile_n, - accumulator_bottom, D256_BM32_PHASE_BLOCK_N, mem_row_major); - } - asm volatile("" ::: "memory"); - - D256BM32PhaseAccumulatorFragment qk_top; - D256BM32PhaseAccumulatorFragment qk_bottom; - { - D256BM32PhaseMatrixAFragment a_fragment; - D256BM32PhaseQKMatrixBFragment b_fragment; - fill_fragment(qk_top, 0.0f); - fill_fragment(qk_bottom, 0.0f); - -#pragma unroll - for (int k_offset = 0; k_offset < D256_BM32_PHASE_D; - k_offset += WMMA_K) { - load_matrix_sync(b_fragment, k_tile + k_offset, k_token_stride); - d256_bm32_phase_load_matrix_a(a_fragment, shared.query, k_offset); - mma_sync(qk_top, a_fragment, b_fragment, qk_top); - d256_bm32_phase_load_matrix_a( - a_fragment, - shared.query + D256_BM32_PHASE_PANEL_M * D256_BM32_PHASE_D, - k_offset); - mma_sync(qk_bottom, a_fragment, b_fragment, qk_bottom); - } - } - asm volatile("" ::: "memory"); - - if constexpr (PAIR_SCRATCH) { - d256_bm32_phase_reload_pair_scratch(shared.score, accumulator_top, - accumulator_bottom); - const int partner_warp = warp_id ^ 1; - if (partner_warp * WMMA_N < valid_k_rows) { - d256_bm32_phase_sync_warp_pair(warp_id >> 1); + if (tid == 0) { + shared.block_index = + static_cast(blockIdx.y) * total_n_blocks + / D256_BM32_PHASE_SPLIT_PARTS; + shared.split_end_block = + (static_cast(blockIdx.y) + 1) * total_n_blocks + / D256_BM32_PHASE_SPLIT_PARTS; } - } else { - load_matrix_sync(accumulator_top, shared.score + tile_n, - D256_BM32_PHASE_BLOCK_N, mem_row_major); - load_matrix_sync(accumulator_bottom, - shared.score + - D256_BM32_PHASE_PANEL_M * D256_BM32_PHASE_BLOCK_N + - tile_n, - D256_BM32_PHASE_BLOCK_N, mem_row_major); - } - - store_matrix_sync(shared.score + tile_n, qk_top, D256_BM32_PHASE_BLOCK_N, - mem_row_major); - store_matrix_sync(shared.score + - D256_BM32_PHASE_PANEL_M * D256_BM32_PHASE_BLOCK_N + - tile_n, - qk_bottom, D256_BM32_PHASE_BLOCK_N, mem_row_major); + __syncthreads(); } - __syncthreads(); - const int causal_q_offset = max(shared.actual_n - M, 0); -#pragma unroll 1 - for (int index = tid; - index < D256_BM32_PHASE_BLOCK_M * D256_BM32_PHASE_BLOCK_N; - index += D256_BM32_PHASE_THREADS) { - const int row = index / D256_BM32_PHASE_BLOCK_N; - const int column = index - row * D256_BM32_PHASE_BLOCK_N; - const int global_m = start_row + row; - const int global_n = start_col + column; - const bool is_valid = global_m < M && global_n < start_col + valid_k_rows; - if constexpr (IS_CAUSAL) { - if (is_valid && global_n <= global_m + causal_q_offset) { - shared.score[index] *= softmax_scale; + D256BM32PhaseAccumulatorFragment accumulator_top; + D256BM32PhaseAccumulatorFragment accumulator_bottom; + fill_fragment(accumulator_top, 0.0f); + fill_fragment(accumulator_bottom, 0.0f); + + for (;;) { + if constexpr (SPLIT_KV) { + if (shared.block_index >= shared.split_end_block) { + break; + } } else { - shared.score[index] = NEG_INF; + if (shared.block_index >= total_n_blocks) { + break; + } } - } else { - if (is_valid) { - shared.score[index] *= softmax_scale; - } else { - shared.score[index] = NEG_INF; + const int start_col = + shared.block_index * D256_BM32_PHASE_BLOCK_N; + const int valid_k_rows = min(D256_BM32_PHASE_BLOCK_N, + shared.actual_n - start_col); + + if (tid < D256_BM32_PHASE_PAGE_SLOTS) { + const int token_offset = tid * D256_BM32_PHASE_PAGE_SIZE; + if (token_offset < valid_k_rows) { + const int global_token_idx = start_col + token_offset; + const int virtual_block_idx = + global_token_idx / D256_BM32_PHASE_PAGE_BLOCK_SIZE; + shared.page_idx[tid] = + __ldg(&block_table[shared.batch_id + * max_num_blocks_per_seq + + virtual_block_idx]); + shared.page_offset[tid] = + global_token_idx + - virtual_block_idx * D256_BM32_PHASE_PAGE_BLOCK_SIZE; + shared.k_tile_ptr[tid] = reinterpret_cast( + K_cache + + (int64_t)shared.page_idx[tid] * k_block_stride + + (int64_t)shared.page_offset[tid] * k_token_stride + + (int64_t)shared.kv_head_id * k_head_stride); + shared.v_tile_ptr[tid] = reinterpret_cast( + V_cache + + (int64_t)shared.page_idx[tid] * v_block_stride + + (int64_t)shared.page_offset[tid] * v_token_stride + + (int64_t)shared.kv_head_id * v_head_stride); + } else { + shared.page_idx[tid] = -1; + shared.page_offset[tid] = 0; + shared.k_tile_ptr[tid] = 0; + shared.v_tile_ptr[tid] = 0; + } } - } - } - __syncthreads(); + __syncthreads(); - if constexpr (ALL_P) { -#pragma unroll - for (int panel = 0; panel < D256_BM32_PHASE_PANELS; ++panel) { - const int panel_start = panel * D256_BM32_PHASE_SOFTMAX_N; - const int valid_panel_columns = - min(D256_BM32_PHASE_SOFTMAX_N, valid_k_rows - panel_start); - if (valid_panel_columns > 0) { - __half* probability_top = - shared.probability_top + - panel * D256_BM32_PHASE_PROBABILITY_ELEMENTS; - __half* probability_bottom = - shared.probability_bottom + - panel * D256_BM32_PHASE_PROBABILITY_ELEMENTS; - float* row_exp_diff = - shared.row_exp_diff + panel * D256_BM32_PHASE_BLOCK_M; - float top_exp_diff; - float bottom_exp_diff; - if constexpr (SPLIT_KV && CHECK_SPLIT_EMPTY) { - const int first_global_n = start_col + panel_start; - const bool top_has_visible_value = - warp_id < valid_q_rows && - (!IS_CAUSAL || - first_global_n <= start_row + warp_id + causal_q_offset); - const bool bottom_has_visible_value = - D256_BM32_PHASE_PANEL_M + warp_id < valid_q_rows && - (!IS_CAUSAL || first_global_n <= start_row + - D256_BM32_PHASE_PANEL_M + - warp_id + causal_q_offset); - top_exp_diff = d256_bm32_phase_make_split_probability_row( - shared.score + warp_id * D256_BM32_PHASE_BLOCK_N, - probability_top, warp_id, shared.row_max, shared.row_sum, - warp_id, panel, NEG_INF, top_has_visible_value); - bottom_exp_diff = d256_bm32_phase_make_split_probability_row( - shared.score + (D256_BM32_PHASE_PANEL_M + warp_id) * - D256_BM32_PHASE_BLOCK_N, - probability_bottom, warp_id, shared.row_max, shared.row_sum, - D256_BM32_PHASE_PANEL_M + warp_id, panel, NEG_INF, - bottom_has_visible_value); - } else { - top_exp_diff = d256_bm32_phase_make_probability_row( - shared.score + warp_id * D256_BM32_PHASE_BLOCK_N, - probability_top, warp_id, shared.row_max, shared.row_sum, - warp_id, panel, NEG_INF); - bottom_exp_diff = d256_bm32_phase_make_probability_row( - shared.score + (D256_BM32_PHASE_PANEL_M + warp_id) * - D256_BM32_PHASE_BLOCK_N, - probability_bottom, warp_id, shared.row_max, shared.row_sum, - D256_BM32_PHASE_PANEL_M + warp_id, panel, NEG_INF); - } - if (lane_id == 0) { - row_exp_diff[warp_id] = top_exp_diff; - row_exp_diff[D256_BM32_PHASE_PANEL_M + warp_id] = bottom_exp_diff; - } - } - // Publish lane 0's online-softmax state to this warp before - // it advances to the next panel. - __syncwarp(); - } - __syncthreads(); - -#pragma unroll - for (int panel = 0; panel < D256_BM32_PHASE_PANELS; ++panel) { - const int panel_start = panel * D256_BM32_PHASE_SOFTMAX_N; - const int valid_panel_columns = - min(D256_BM32_PHASE_SOFTMAX_N, valid_k_rows - panel_start); - if (valid_panel_columns > 0) { - const __half* probability_top = - shared.probability_top + - panel * D256_BM32_PHASE_PROBABILITY_ELEMENTS; - const __half* probability_bottom = - shared.probability_bottom + - panel * D256_BM32_PHASE_PROBABILITY_ELEMENTS; - const float* row_exp_diff = - shared.row_exp_diff + panel * D256_BM32_PHASE_BLOCK_M; - if constexpr (SPLIT_KV) { - d256_bm32_phase_scale_accumulators( - accumulator_top, accumulator_bottom, row_exp_diff); - } else if (shared.block_index != 0 || panel != 0) { - d256_bm32_phase_scale_accumulators( - accumulator_top, accumulator_bottom, row_exp_diff); - } - const int page_slot_k0 = panel_start >> 4; - const __half* value_k0 = - reinterpret_cast(shared.v_tile_ptr[page_slot_k0]); - const __half* value_k16 = nullptr; - if (valid_panel_columns > WMMA_K) { - const int page_slot_k16 = (panel_start + WMMA_K) >> 4; - value_k16 = reinterpret_cast( - shared.v_tile_ptr[page_slot_k16]); - } - d256_bm32_phase_update_pv_panel(probability_top, probability_bottom, - value_k0, value_k16, v_token_stride, - warp_id * WMMA_N, valid_panel_columns, - accumulator_top, accumulator_bottom); + if (warp_id < D256_BM32_PHASE_BLOCK_N / WMMA_N + && warp_id * WMMA_N < valid_k_rows) { + const int tile_n = warp_id * WMMA_N; + const __half* k_tile = + reinterpret_cast(shared.k_tile_ptr[warp_id]); + if constexpr (PAIR_SCRATCH) { + d256_bm32_phase_spill_pair_scratch( + shared.score, accumulator_top, accumulator_bottom); + } else { + store_matrix_sync(shared.score + tile_n, accumulator_top, + D256_BM32_PHASE_BLOCK_N, mem_row_major); + store_matrix_sync( + shared.score + D256_BM32_PHASE_PANEL_M + * D256_BM32_PHASE_BLOCK_N + tile_n, + accumulator_bottom, D256_BM32_PHASE_BLOCK_N, + mem_row_major); + } + asm volatile("" ::: "memory"); + + D256BM32PhaseAccumulatorFragment qk_top; + D256BM32PhaseAccumulatorFragment qk_bottom; + { + D256BM32PhaseMatrixAFragment a_fragment; + D256BM32PhaseQKMatrixBFragment b_fragment; + fill_fragment(qk_top, 0.0f); + fill_fragment(qk_bottom, 0.0f); + + #pragma unroll + for (int k_offset = 0; k_offset < D256_BM32_PHASE_D; + k_offset += WMMA_K) { + load_matrix_sync(b_fragment, k_tile + k_offset, + k_token_stride); + d256_bm32_phase_load_matrix_a( + a_fragment, shared.query, k_offset); + mma_sync(qk_top, a_fragment, b_fragment, qk_top); + d256_bm32_phase_load_matrix_a( + a_fragment, + shared.query + + D256_BM32_PHASE_PANEL_M + * D256_BM32_PHASE_D, + k_offset); + mma_sync(qk_bottom, a_fragment, b_fragment, qk_bottom); + } + } + asm volatile("" ::: "memory"); + + if constexpr (PAIR_SCRATCH) { + d256_bm32_phase_reload_pair_scratch( + shared.score, accumulator_top, accumulator_bottom); + const int partner_warp = warp_id ^ 1; + if (partner_warp * WMMA_N < valid_k_rows) { + d256_bm32_phase_sync_warp_pair(warp_id >> 1); + } + } else { + load_matrix_sync(accumulator_top, shared.score + tile_n, + D256_BM32_PHASE_BLOCK_N, mem_row_major); + load_matrix_sync( + accumulator_bottom, + shared.score + D256_BM32_PHASE_PANEL_M + * D256_BM32_PHASE_BLOCK_N + tile_n, + D256_BM32_PHASE_BLOCK_N, mem_row_major); + } + + store_matrix_sync(shared.score + tile_n, qk_top, + D256_BM32_PHASE_BLOCK_N, mem_row_major); + store_matrix_sync( + shared.score + D256_BM32_PHASE_PANEL_M + * D256_BM32_PHASE_BLOCK_N + tile_n, + qk_bottom, D256_BM32_PHASE_BLOCK_N, mem_row_major); } - } - // Protect the current V pointers and block index from the next - // iteration while slower warps finish the last PV panel. - __syncthreads(); - } else { -#pragma unroll - for (int panel = 0; panel < D256_BM32_PHASE_PANELS; ++panel) { - const int panel_start = panel * D256_BM32_PHASE_SOFTMAX_N; - const int valid_panel_columns = - min(D256_BM32_PHASE_SOFTMAX_N, valid_k_rows - panel_start); - if (valid_panel_columns > 0) { - const float top_exp_diff = d256_bm32_phase_make_probability_row( - shared.score + warp_id * D256_BM32_PHASE_BLOCK_N, - shared.probability_top, warp_id, shared.row_max, shared.row_sum, - warp_id, panel, NEG_INF); - const float bottom_exp_diff = d256_bm32_phase_make_probability_row( - shared.score + - (D256_BM32_PHASE_PANEL_M + warp_id) * D256_BM32_PHASE_BLOCK_N, - shared.probability_bottom, warp_id, shared.row_max, - shared.row_sum, D256_BM32_PHASE_PANEL_M + warp_id, panel, - NEG_INF); - if (lane_id == 0) { - shared.row_exp_diff[warp_id] = top_exp_diff; - shared.row_exp_diff[D256_BM32_PHASE_PANEL_M + warp_id] = - bottom_exp_diff; - } + __syncthreads(); + + const int causal_q_offset = max(shared.actual_n - M, 0); + #pragma unroll 1 + for (int index = tid; + index < D256_BM32_PHASE_BLOCK_M * D256_BM32_PHASE_BLOCK_N; + index += D256_BM32_PHASE_THREADS) { + const int row = index / D256_BM32_PHASE_BLOCK_N; + const int column = index - row * D256_BM32_PHASE_BLOCK_N; + const int global_m = start_row + row; + const int global_n = start_col + column; + const bool is_valid = global_m < M + && global_n < start_col + valid_k_rows; + if constexpr (IS_CAUSAL) { + if (is_valid && global_n <= global_m + causal_q_offset) { + shared.score[index] *= softmax_scale; + } else { + shared.score[index] = NEG_INF; + } + } else { + if (is_valid) { + shared.score[index] *= softmax_scale; + } else { + shared.score[index] = NEG_INF; + } + } } __syncthreads(); - if (valid_panel_columns > 0) { - if constexpr (SPLIT_KV) { - d256_bm32_phase_scale_accumulators( - accumulator_top, accumulator_bottom, shared.row_exp_diff); - } else if (shared.block_index != 0 || panel != 0) { - d256_bm32_phase_scale_accumulators( - accumulator_top, accumulator_bottom, shared.row_exp_diff); - } - const int page_slot_k0 = panel_start >> 4; - const __half* value_k0 = - reinterpret_cast(shared.v_tile_ptr[page_slot_k0]); - const __half* value_k16 = nullptr; - if (valid_panel_columns > WMMA_K) { - const int page_slot_k16 = (panel_start + WMMA_K) >> 4; - value_k16 = reinterpret_cast( - shared.v_tile_ptr[page_slot_k16]); - } - d256_bm32_phase_update_pv_panel( - shared.probability_top, shared.probability_bottom, value_k0, - value_k16, v_token_stride, warp_id * WMMA_N, valid_panel_columns, - accumulator_top, accumulator_bottom); + if constexpr (ALL_P) { + #pragma unroll + for (int panel = 0; panel < D256_BM32_PHASE_PANELS; + ++panel) { + const int panel_start = + panel * D256_BM32_PHASE_SOFTMAX_N; + const int valid_panel_columns = + min(D256_BM32_PHASE_SOFTMAX_N, + valid_k_rows - panel_start); + if (valid_panel_columns > 0) { + __half* probability_top = + shared.probability_top + + panel * D256_BM32_PHASE_PROBABILITY_ELEMENTS; + __half* probability_bottom = + shared.probability_bottom + + panel * D256_BM32_PHASE_PROBABILITY_ELEMENTS; + float* row_exp_diff = + shared.row_exp_diff + + panel * D256_BM32_PHASE_BLOCK_M; + float top_exp_diff; + float bottom_exp_diff; + if constexpr (SPLIT_KV && CHECK_SPLIT_EMPTY) { + const int first_global_n = start_col + panel_start; + const bool top_has_visible_value = + warp_id < valid_q_rows + && (!IS_CAUSAL + || first_global_n + <= start_row + warp_id + + causal_q_offset); + const bool bottom_has_visible_value = + D256_BM32_PHASE_PANEL_M + warp_id < valid_q_rows + && (!IS_CAUSAL + || first_global_n + <= start_row + + D256_BM32_PHASE_PANEL_M + + warp_id + causal_q_offset); + top_exp_diff = + d256_bm32_phase_make_split_probability_row( + shared.score + + warp_id * D256_BM32_PHASE_BLOCK_N, + probability_top, warp_id, shared.row_max, + shared.row_sum, warp_id, panel, NEG_INF, + top_has_visible_value); + bottom_exp_diff = + d256_bm32_phase_make_split_probability_row( + shared.score + + (D256_BM32_PHASE_PANEL_M + warp_id) + * D256_BM32_PHASE_BLOCK_N, + probability_bottom, warp_id, shared.row_max, + shared.row_sum, + D256_BM32_PHASE_PANEL_M + warp_id, panel, + NEG_INF, bottom_has_visible_value); + } else { + top_exp_diff = d256_bm32_phase_make_probability_row( + shared.score + + warp_id * D256_BM32_PHASE_BLOCK_N, + probability_top, warp_id, shared.row_max, + shared.row_sum, warp_id, panel, NEG_INF); + bottom_exp_diff = + d256_bm32_phase_make_probability_row( + shared.score + + (D256_BM32_PHASE_PANEL_M + warp_id) + * D256_BM32_PHASE_BLOCK_N, + probability_bottom, warp_id, shared.row_max, + shared.row_sum, + D256_BM32_PHASE_PANEL_M + warp_id, panel, + NEG_INF); + } + if (lane_id == 0) { + row_exp_diff[warp_id] = top_exp_diff; + row_exp_diff[D256_BM32_PHASE_PANEL_M + warp_id] = + bottom_exp_diff; + } + } + // Publish lane 0's online-softmax state to this warp before + // it advances to the next panel. + __syncwarp(); + } + __syncthreads(); + + #pragma unroll + for (int panel = 0; panel < D256_BM32_PHASE_PANELS; + ++panel) { + const int panel_start = + panel * D256_BM32_PHASE_SOFTMAX_N; + const int valid_panel_columns = + min(D256_BM32_PHASE_SOFTMAX_N, + valid_k_rows - panel_start); + if (valid_panel_columns > 0) { + const __half* probability_top = + shared.probability_top + + panel * D256_BM32_PHASE_PROBABILITY_ELEMENTS; + const __half* probability_bottom = + shared.probability_bottom + + panel * D256_BM32_PHASE_PROBABILITY_ELEMENTS; + const float* row_exp_diff = + shared.row_exp_diff + + panel * D256_BM32_PHASE_BLOCK_M; + if constexpr (SPLIT_KV) { + d256_bm32_phase_scale_accumulators( + accumulator_top, accumulator_bottom, + row_exp_diff); + } else if (shared.block_index != 0 || panel != 0) { + d256_bm32_phase_scale_accumulators( + accumulator_top, accumulator_bottom, + row_exp_diff); + } + const int page_slot_k0 = panel_start >> 4; + const __half* value_k0 = + reinterpret_cast( + shared.v_tile_ptr[page_slot_k0]); + const __half* value_k16 = nullptr; + if (valid_panel_columns > WMMA_K) { + const int page_slot_k16 = + (panel_start + WMMA_K) >> 4; + value_k16 = + reinterpret_cast( + shared.v_tile_ptr[page_slot_k16]); + } + d256_bm32_phase_update_pv_panel( + probability_top, probability_bottom, + value_k0, value_k16, v_token_stride, + warp_id * WMMA_N, + valid_panel_columns, + accumulator_top, accumulator_bottom); + } + } + // Protect the current V pointers and block index from the next + // iteration while slower warps finish the last PV panel. + __syncthreads(); + } else { + #pragma unroll + for (int panel = 0; panel < D256_BM32_PHASE_PANELS; + ++panel) { + const int panel_start = + panel * D256_BM32_PHASE_SOFTMAX_N; + const int valid_panel_columns = + min(D256_BM32_PHASE_SOFTMAX_N, + valid_k_rows - panel_start); + if (valid_panel_columns > 0) { + const float top_exp_diff = + d256_bm32_phase_make_probability_row( + shared.score + + warp_id * D256_BM32_PHASE_BLOCK_N, + shared.probability_top, warp_id, shared.row_max, + shared.row_sum, warp_id, panel, NEG_INF); + const float bottom_exp_diff = + d256_bm32_phase_make_probability_row( + shared.score + + (D256_BM32_PHASE_PANEL_M + warp_id) + * D256_BM32_PHASE_BLOCK_N, + shared.probability_bottom, warp_id, + shared.row_max, shared.row_sum, + D256_BM32_PHASE_PANEL_M + warp_id, panel, + NEG_INF); + if (lane_id == 0) { + shared.row_exp_diff[warp_id] = top_exp_diff; + shared.row_exp_diff[ + D256_BM32_PHASE_PANEL_M + warp_id] = + bottom_exp_diff; + } + } + __syncthreads(); + + if (valid_panel_columns > 0) { + if constexpr (SPLIT_KV) { + d256_bm32_phase_scale_accumulators( + accumulator_top, accumulator_bottom, + shared.row_exp_diff); + } else if (shared.block_index != 0 || panel != 0) { + d256_bm32_phase_scale_accumulators( + accumulator_top, accumulator_bottom, + shared.row_exp_diff); + } + const int page_slot_k0 = panel_start >> 4; + const __half* value_k0 = + reinterpret_cast( + shared.v_tile_ptr[page_slot_k0]); + const __half* value_k16 = nullptr; + if (valid_panel_columns > WMMA_K) { + const int page_slot_k16 = + (panel_start + WMMA_K) >> 4; + value_k16 = + reinterpret_cast( + shared.v_tile_ptr[page_slot_k16]); + } + d256_bm32_phase_update_pv_panel( + shared.probability_top, + shared.probability_bottom, + value_k0, value_k16, v_token_stride, + warp_id * WMMA_N, + valid_panel_columns, + accumulator_top, accumulator_bottom); + } + __syncthreads(); + } + } + + if (tid == 0) { + ++shared.block_index; } __syncthreads(); - } } - if (tid == 0) { - ++shared.block_index; - } - __syncthreads(); - } + { + const size_t q_head_linear = static_cast(blockIdx.z); + if constexpr (SPLIT_KV) { + const size_t partial_row_base = + (q_head_linear * D256_BM32_PHASE_SPLIT_PARTS + + static_cast(blockIdx.y)) + * M + + start_row; + float* partial_out = + split_tmp_out + + partial_row_base * D256_BM32_PHASE_D; + d256_bm32_phase_store_unnormalized_partial( + accumulator_top, partial_out, 0, + warp_id * WMMA_N, valid_q_rows); + d256_bm32_phase_store_unnormalized_partial( + accumulator_bottom, partial_out, + D256_BM32_PHASE_PANEL_M, warp_id * WMMA_N, valid_q_rows); + if (tid < valid_q_rows) { + split_tmp_row_max[partial_row_base + tid] = + shared.row_max[tid]; + split_tmp_row_sum[partial_row_base + tid] = + shared.row_sum[tid]; + } + } else { + __half* out_ptr = + Out + q_head_linear * M * D256_BM32_PHASE_D + + start_row * D256_BM32_PHASE_D; + float* softmax_lse_ptr = + softmax_lse + q_head_linear * M + start_row; + if (valid_q_rows == D256_BM32_PHASE_BLOCK_M) { + d256_bm32_phase_store_output( + accumulator_top, out_ptr, shared.row_sum, 0, + warp_id * WMMA_N); + d256_bm32_phase_store_output( + accumulator_bottom, out_ptr, shared.row_sum, + D256_BM32_PHASE_PANEL_M, warp_id * WMMA_N); + } else { + d256_bm32_phase_store_output_partial( + accumulator_top, out_ptr, shared.row_sum, 0, + warp_id * WMMA_N, valid_q_rows); + d256_bm32_phase_store_output_partial( + accumulator_bottom, out_ptr, shared.row_sum, + D256_BM32_PHASE_PANEL_M, warp_id * WMMA_N, valid_q_rows); + } - { - const size_t q_head_linear = static_cast(blockIdx.z); - if constexpr (SPLIT_KV) { - const size_t partial_row_base = - (q_head_linear * D256_BM32_PHASE_SPLIT_PARTS + - static_cast(blockIdx.y)) * - M + - start_row; - float* partial_out = split_tmp_out + partial_row_base * D256_BM32_PHASE_D; - d256_bm32_phase_store_unnormalized_partial( - accumulator_top, partial_out, 0, warp_id * WMMA_N, valid_q_rows); - d256_bm32_phase_store_unnormalized_partial( - accumulator_bottom, partial_out, D256_BM32_PHASE_PANEL_M, - warp_id * WMMA_N, valid_q_rows); - if (tid < valid_q_rows) { - split_tmp_row_max[partial_row_base + tid] = shared.row_max[tid]; - split_tmp_row_sum[partial_row_base + tid] = shared.row_sum[tid]; - } - } else { - __half* out_ptr = Out + q_head_linear * M * D256_BM32_PHASE_D + - start_row * D256_BM32_PHASE_D; - float* softmax_lse_ptr = softmax_lse + q_head_linear * M + start_row; - if (valid_q_rows == D256_BM32_PHASE_BLOCK_M) { - d256_bm32_phase_store_output(accumulator_top, out_ptr, shared.row_sum, - 0, warp_id * WMMA_N); - d256_bm32_phase_store_output(accumulator_bottom, out_ptr, - shared.row_sum, D256_BM32_PHASE_PANEL_M, - warp_id * WMMA_N); - } else { - d256_bm32_phase_store_output_partial(accumulator_top, out_ptr, - shared.row_sum, 0, - warp_id * WMMA_N, valid_q_rows); - d256_bm32_phase_store_output_partial( - accumulator_bottom, out_ptr, shared.row_sum, - D256_BM32_PHASE_PANEL_M, warp_id * WMMA_N, valid_q_rows); - } - - if (tid < valid_q_rows) { - const float sum = fmaxf(shared.row_sum[tid], 1e-24f); - softmax_lse_ptr[tid] = shared.row_max[tid] + logf(sum); - } - } - } + if (tid < valid_q_rows) { + const float sum = fmaxf(shared.row_sum[tid], 1e-24f); + softmax_lse_ptr[tid] = shared.row_max[tid] + logf(sum); + } + } + } } -template -__global__ -__launch_bounds__(D256_BM32_PHASE_THREADS, 2) void flash_attention_forward_paged_d256_bm32_phase_kernel( - const __half* __restrict__ Q, const __half* __restrict__ K_cache, - const __half* __restrict__ V_cache, __half* __restrict__ Out, - float* __restrict__ softmax_lse, const int* __restrict__ block_table, - const int* __restrict__ seqused_k, int H, int M, int max_num_blocks_per_seq, - int num_kv_heads, int64_t k_block_stride, int64_t k_token_stride, - int64_t k_head_stride, int64_t v_block_stride, int64_t v_token_stride, - int64_t v_head_stride, float softmax_scale) { - flash_attention_forward_paged_d256_bm32_phase_body( - Q, K_cache, V_cache, Out, softmax_lse, block_table, seqused_k, H, M, - max_num_blocks_per_seq, num_kv_heads, k_block_stride, k_token_stride, - k_head_stride, v_block_stride, v_token_stride, v_head_stride, - softmax_scale, nullptr, nullptr, nullptr, 0); +template +__global__ __launch_bounds__(D256_BM32_PHASE_THREADS, 2) +void flash_attention_forward_paged_d256_bm32_phase_kernel( + const __half* __restrict__ Q, + const __half* __restrict__ K_cache, + const __half* __restrict__ V_cache, + __half* __restrict__ Out, + float* __restrict__ softmax_lse, + const int* __restrict__ block_table, + const int* __restrict__ seqused_k, + int H, + int M, + int max_num_blocks_per_seq, + int num_kv_heads, + int64_t k_block_stride, + int64_t k_token_stride, + int64_t k_head_stride, + int64_t v_block_stride, + int64_t v_token_stride, + int64_t v_head_stride, + float softmax_scale) { + flash_attention_forward_paged_d256_bm32_phase_body< + IS_CAUSAL, ALL_P, PAIR_SCRATCH, false>( + Q, K_cache, V_cache, Out, softmax_lse, block_table, seqused_k, + H, M, max_num_blocks_per_seq, num_kv_heads, k_block_stride, + k_token_stride, k_head_stride, v_block_stride, v_token_stride, + v_head_stride, softmax_scale, nullptr, nullptr, nullptr, 0); } -template -__global__ -__launch_bounds__(D256_BM32_PHASE_THREADS, 2) void flash_attention_forward_paged_d256_bm32_splitkv3_partial_kernel( - const __half* __restrict__ Q, const __half* __restrict__ K_cache, - const __half* __restrict__ V_cache, const int* __restrict__ block_table, - int H, int M, int actual_n, int max_num_blocks_per_seq, int num_kv_heads, - int64_t k_block_stride, int64_t k_token_stride, int64_t k_head_stride, - int64_t v_block_stride, int64_t v_token_stride, int64_t v_head_stride, - float softmax_scale, float* __restrict__ split_tmp_out, +template +__global__ __launch_bounds__(D256_BM32_PHASE_THREADS, 2) +void flash_attention_forward_paged_d256_bm32_splitkv3_partial_kernel( + const __half* __restrict__ Q, + const __half* __restrict__ K_cache, + const __half* __restrict__ V_cache, + const int* __restrict__ block_table, + int H, + int M, + int actual_n, + int max_num_blocks_per_seq, + int num_kv_heads, + int64_t k_block_stride, + int64_t k_token_stride, + int64_t k_head_stride, + int64_t v_block_stride, + int64_t v_token_stride, + int64_t v_head_stride, + float softmax_scale, + float* __restrict__ split_tmp_out, float* __restrict__ split_tmp_row_max, float* __restrict__ split_tmp_row_sum) { - flash_attention_forward_paged_d256_bm32_phase_body( - Q, K_cache, V_cache, nullptr, nullptr, block_table, nullptr, H, M, - max_num_blocks_per_seq, num_kv_heads, k_block_stride, k_token_stride, - k_head_stride, v_block_stride, v_token_stride, v_head_stride, - softmax_scale, split_tmp_out, split_tmp_row_max, split_tmp_row_sum, - actual_n); + flash_attention_forward_paged_d256_bm32_phase_body< + true, true, true, true, CHECK_SPLIT_EMPTY>( + Q, K_cache, V_cache, nullptr, nullptr, block_table, nullptr, + H, M, max_num_blocks_per_seq, num_kv_heads, k_block_stride, + k_token_stride, k_head_stride, v_block_stride, v_token_stride, + v_head_stride, softmax_scale, split_tmp_out, + split_tmp_row_max, split_tmp_row_sum, actual_n); } -template +template void launch_flash_attention_forward_paged_d256_bm32_phase_kernel( - const torch::Tensor& Q, const torch::Tensor& K_cache, - const torch::Tensor& V_cache, torch::Tensor& Out, - torch::Tensor& softmax_lse, const torch::Tensor& block_table, - const torch::Tensor& seq_lens, float softmax_scale, cudaStream_t stream) { - const int H = Q.size(1); - const int M = Q.size(2); - const int num_kv_heads = K_cache.size(2); - const int max_num_blocks_per_seq = block_table.size(1); - const int64_t k_block_stride = K_cache.stride(0); - const int64_t k_token_stride = K_cache.stride(1); - const int64_t k_head_stride = K_cache.stride(2); - const int64_t v_block_stride = V_cache.stride(0); - const int64_t v_token_stride = V_cache.stride(1); - const int64_t v_head_stride = V_cache.stride(2); - const dim3 grid((M + D256_BM32_PHASE_BLOCK_M - 1) / D256_BM32_PHASE_BLOCK_M, - 1, Q.size(0) * H); - const dim3 block(D256_BM32_PHASE_THREADS); - flash_attention_forward_paged_d256_bm32_phase_kernel - <<>>( - reinterpret_cast(Q.data_ptr()), - reinterpret_cast(K_cache.data_ptr()), - reinterpret_cast(V_cache.data_ptr()), - reinterpret_cast<__half*>(Out.data_ptr()), - softmax_lse.data_ptr(), block_table.data_ptr(), - seq_lens.data_ptr(), H, M, max_num_blocks_per_seq, num_kv_heads, - k_block_stride, k_token_stride, k_head_stride, v_block_stride, - v_token_stride, v_head_stride, softmax_scale); + const torch::Tensor& Q, + const torch::Tensor& K_cache, + const torch::Tensor& V_cache, + torch::Tensor& Out, + torch::Tensor& softmax_lse, + const torch::Tensor& block_table, + const torch::Tensor& seq_lens, + float softmax_scale, + cudaStream_t stream) { + const int H = Q.size(1); + const int M = Q.size(2); + const int num_kv_heads = K_cache.size(2); + const int max_num_blocks_per_seq = block_table.size(1); + const int64_t k_block_stride = K_cache.stride(0); + const int64_t k_token_stride = K_cache.stride(1); + const int64_t k_head_stride = K_cache.stride(2); + const int64_t v_block_stride = V_cache.stride(0); + const int64_t v_token_stride = V_cache.stride(1); + const int64_t v_head_stride = V_cache.stride(2); + const dim3 grid((M + D256_BM32_PHASE_BLOCK_M - 1) + / D256_BM32_PHASE_BLOCK_M, + 1, Q.size(0) * H); + const dim3 block(D256_BM32_PHASE_THREADS); + flash_attention_forward_paged_d256_bm32_phase_kernel< + IS_CAUSAL, ALL_P, PAIR_SCRATCH><<>>( + reinterpret_cast(Q.data_ptr()), + reinterpret_cast(K_cache.data_ptr()), + reinterpret_cast(V_cache.data_ptr()), + reinterpret_cast<__half*>(Out.data_ptr()), + softmax_lse.data_ptr(), block_table.data_ptr(), + seq_lens.data_ptr(), H, M, max_num_blocks_per_seq, + num_kv_heads, k_block_stride, k_token_stride, k_head_stride, + v_block_stride, v_token_stride, v_head_stride, softmax_scale); } -__global__ -__launch_bounds__(256) void flash_attention_forward_paged_d256_bm32_splitkv3_merge_kernel( +__global__ __launch_bounds__(256) +void flash_attention_forward_paged_d256_bm32_splitkv3_merge_kernel( const float* __restrict__ split_tmp_out, const float* __restrict__ split_tmp_row_max, - const float* __restrict__ split_tmp_row_sum, __half* __restrict__ Out, - float* __restrict__ softmax_lse, int M) { - const int start_row = blockIdx.x * D256_BM32_PHASE_BLOCK_M; - const int row_local = threadIdx.x >> 3; - const int lane_in_row = threadIdx.x & 7; - const int row = start_row + row_local; - if (row >= M) { - return; - } - - const size_t q_head_linear = static_cast(blockIdx.z); - const size_t first_state = - (q_head_linear * D256_BM32_PHASE_SPLIT_PARTS) * M + row; - const float max_0 = split_tmp_row_max[first_state]; - const float max_1 = split_tmp_row_max[first_state + M]; - const float max_2 = split_tmp_row_max[first_state + 2 * M]; - const float sum_0 = split_tmp_row_sum[first_state]; - const float sum_1 = split_tmp_row_sum[first_state + M]; - const float sum_2 = split_tmp_row_sum[first_state + 2 * M]; - const float merged_max = fmaxf(max_0, fmaxf(max_1, max_2)); - const float weight_0 = - sum_0 > 0.0f ? __expf(fmaxf(max_0 - merged_max, -80.0f)) : 0.0f; - const float weight_1 = - sum_1 > 0.0f ? __expf(fmaxf(max_1 - merged_max, -80.0f)) : 0.0f; - const float weight_2 = - sum_2 > 0.0f ? __expf(fmaxf(max_2 - merged_max, -80.0f)) : 0.0f; - const float denominator = - weight_0 * sum_0 + weight_1 * sum_1 + weight_2 * sum_2; - const float inverse_denominator = 1.0f / fmaxf(denominator, 1e-24f); - - if (lane_in_row == 0) { - softmax_lse[q_head_linear * M + row] = - merged_max + logf(fmaxf(denominator, 1e-24f)); - } - - const size_t output_base = (q_head_linear * M + row) * D256_BM32_PHASE_D; - const size_t partial_base_0 = first_state * D256_BM32_PHASE_D; - const size_t partial_base_1 = (first_state + M) * D256_BM32_PHASE_D; - const size_t partial_base_2 = (first_state + 2 * M) * D256_BM32_PHASE_D; -#pragma unroll - for (int column = lane_in_row; column < D256_BM32_PHASE_D; column += 8) { - const float numerator = weight_0 * split_tmp_out[partial_base_0 + column] + - weight_1 * split_tmp_out[partial_base_1 + column] + - weight_2 * split_tmp_out[partial_base_2 + column]; - Out[output_base + column] = - __float2half_rn(numerator * inverse_denominator); - } + const float* __restrict__ split_tmp_row_sum, + __half* __restrict__ Out, + float* __restrict__ softmax_lse, + int M) { + const int start_row = blockIdx.x * D256_BM32_PHASE_BLOCK_M; + const int row_local = threadIdx.x >> 3; + const int lane_in_row = threadIdx.x & 7; + const int row = start_row + row_local; + if (row >= M) { + return; + } + + const size_t q_head_linear = static_cast(blockIdx.z); + const size_t first_state = + (q_head_linear * D256_BM32_PHASE_SPLIT_PARTS) * M + row; + const float max_0 = split_tmp_row_max[first_state]; + const float max_1 = split_tmp_row_max[first_state + M]; + const float max_2 = split_tmp_row_max[first_state + 2 * M]; + const float sum_0 = split_tmp_row_sum[first_state]; + const float sum_1 = split_tmp_row_sum[first_state + M]; + const float sum_2 = split_tmp_row_sum[first_state + 2 * M]; + const float merged_max = fmaxf(max_0, fmaxf(max_1, max_2)); + const float weight_0 = + sum_0 > 0.0f ? __expf(fmaxf(max_0 - merged_max, -80.0f)) : 0.0f; + const float weight_1 = + sum_1 > 0.0f ? __expf(fmaxf(max_1 - merged_max, -80.0f)) : 0.0f; + const float weight_2 = + sum_2 > 0.0f ? __expf(fmaxf(max_2 - merged_max, -80.0f)) : 0.0f; + const float denominator = + weight_0 * sum_0 + weight_1 * sum_1 + weight_2 * sum_2; + const float inverse_denominator = + 1.0f / fmaxf(denominator, 1e-24f); + + if (lane_in_row == 0) { + softmax_lse[q_head_linear * M + row] = + merged_max + logf(fmaxf(denominator, 1e-24f)); + } + + const size_t output_base = + (q_head_linear * M + row) * D256_BM32_PHASE_D; + const size_t partial_base_0 = first_state * D256_BM32_PHASE_D; + const size_t partial_base_1 = + (first_state + M) * D256_BM32_PHASE_D; + const size_t partial_base_2 = + (first_state + 2 * M) * D256_BM32_PHASE_D; + #pragma unroll + for (int column = lane_in_row; column < D256_BM32_PHASE_D; column += 8) { + const float numerator = + weight_0 * split_tmp_out[partial_base_0 + column] + + weight_1 * split_tmp_out[partial_base_1 + column] + + weight_2 * split_tmp_out[partial_base_2 + column]; + Out[output_base + column] = + __float2half_rn(numerator * inverse_denominator); + } } void launch_flash_attention_forward_paged_d256_bm32_splitkv3_kernel( - const torch::Tensor& Q, const torch::Tensor& K_cache, - const torch::Tensor& V_cache, torch::Tensor& Out, - torch::Tensor& softmax_lse, torch::Tensor& split_tmp_out, - torch::Tensor& split_tmp_row_max, torch::Tensor& split_tmp_row_sum, - const torch::Tensor& block_table, int actual_n, float softmax_scale, + const torch::Tensor& Q, + const torch::Tensor& K_cache, + const torch::Tensor& V_cache, + torch::Tensor& Out, + torch::Tensor& softmax_lse, + torch::Tensor& split_tmp_out, + torch::Tensor& split_tmp_row_max, + torch::Tensor& split_tmp_row_sum, + const torch::Tensor& block_table, + int actual_n, + float softmax_scale, cudaStream_t stream) { - const int H = Q.size(1); - const int M = Q.size(2); - const int num_kv_heads = K_cache.size(2); - const int max_num_blocks_per_seq = block_table.size(1); - const int64_t k_block_stride = K_cache.stride(0); - const int64_t k_token_stride = K_cache.stride(1); - const int64_t k_head_stride = K_cache.stride(2); - const int64_t v_block_stride = V_cache.stride(0); - const int64_t v_token_stride = V_cache.stride(1); - const int64_t v_head_stride = V_cache.stride(2); - const dim3 partial_grid( - (M + D256_BM32_PHASE_BLOCK_M - 1) / D256_BM32_PHASE_BLOCK_M, - D256_BM32_PHASE_SPLIT_PARTS, Q.size(0) * H); - const dim3 partial_block(D256_BM32_PHASE_THREADS); - const int64_t total_n_blocks = - (static_cast(actual_n) + D256_BM32_PHASE_BLOCK_N - 1) / - D256_BM32_PHASE_BLOCK_N; - const int64_t last_split_begin = - (D256_BM32_PHASE_SPLIT_PARTS - 1) * total_n_blocks / - D256_BM32_PHASE_SPLIT_PARTS * D256_BM32_PHASE_BLOCK_N; - const bool check_split_empty = - total_n_blocks < D256_BM32_PHASE_SPLIT_PARTS || - last_split_begin > static_cast(actual_n) - M; - if (check_split_empty) { - flash_attention_forward_paged_d256_bm32_splitkv3_partial_kernel - <<>>( - reinterpret_cast(Q.data_ptr()), - reinterpret_cast(K_cache.data_ptr()), - reinterpret_cast(V_cache.data_ptr()), - block_table.data_ptr(), H, M, actual_n, max_num_blocks_per_seq, - num_kv_heads, k_block_stride, k_token_stride, k_head_stride, - v_block_stride, v_token_stride, v_head_stride, softmax_scale, - split_tmp_out.data_ptr(), - split_tmp_row_max.data_ptr(), - split_tmp_row_sum.data_ptr()); - } else { - flash_attention_forward_paged_d256_bm32_splitkv3_partial_kernel - <<>>( + const int H = Q.size(1); + const int M = Q.size(2); + const int num_kv_heads = K_cache.size(2); + const int max_num_blocks_per_seq = block_table.size(1); + const int64_t k_block_stride = K_cache.stride(0); + const int64_t k_token_stride = K_cache.stride(1); + const int64_t k_head_stride = K_cache.stride(2); + const int64_t v_block_stride = V_cache.stride(0); + const int64_t v_token_stride = V_cache.stride(1); + const int64_t v_head_stride = V_cache.stride(2); + const dim3 partial_grid( + (M + D256_BM32_PHASE_BLOCK_M - 1) / D256_BM32_PHASE_BLOCK_M, + D256_BM32_PHASE_SPLIT_PARTS, + Q.size(0) * H); + const dim3 partial_block(D256_BM32_PHASE_THREADS); + const int64_t total_n_blocks = + (static_cast(actual_n) + D256_BM32_PHASE_BLOCK_N - 1) + / D256_BM32_PHASE_BLOCK_N; + const int64_t last_split_begin = + (D256_BM32_PHASE_SPLIT_PARTS - 1) * total_n_blocks + / D256_BM32_PHASE_SPLIT_PARTS * D256_BM32_PHASE_BLOCK_N; + const bool check_split_empty = + total_n_blocks < D256_BM32_PHASE_SPLIT_PARTS + || last_split_begin > static_cast(actual_n) - M; + if (check_split_empty) { + flash_attention_forward_paged_d256_bm32_splitkv3_partial_kernel + <<>>( + reinterpret_cast(Q.data_ptr()), + reinterpret_cast(K_cache.data_ptr()), + reinterpret_cast(V_cache.data_ptr()), + block_table.data_ptr(), H, M, actual_n, + max_num_blocks_per_seq, num_kv_heads, k_block_stride, + k_token_stride, k_head_stride, v_block_stride, + v_token_stride, v_head_stride, softmax_scale, + split_tmp_out.data_ptr(), + split_tmp_row_max.data_ptr(), + split_tmp_row_sum.data_ptr()); + } else { + flash_attention_forward_paged_d256_bm32_splitkv3_partial_kernel + <<>>( reinterpret_cast(Q.data_ptr()), reinterpret_cast(K_cache.data_ptr()), reinterpret_cast(V_cache.data_ptr()), - block_table.data_ptr(), H, M, actual_n, max_num_blocks_per_seq, - num_kv_heads, k_block_stride, k_token_stride, k_head_stride, - v_block_stride, v_token_stride, v_head_stride, softmax_scale, + block_table.data_ptr(), H, M, actual_n, + max_num_blocks_per_seq, num_kv_heads, k_block_stride, + k_token_stride, k_head_stride, v_block_stride, + v_token_stride, v_head_stride, softmax_scale, + split_tmp_out.data_ptr(), + split_tmp_row_max.data_ptr(), + split_tmp_row_sum.data_ptr()); + } + + const dim3 merge_grid( + (M + D256_BM32_PHASE_BLOCK_M - 1) / D256_BM32_PHASE_BLOCK_M, + 1, + Q.size(0) * H); + flash_attention_forward_paged_d256_bm32_splitkv3_merge_kernel + <<>>( split_tmp_out.data_ptr(), split_tmp_row_max.data_ptr(), - split_tmp_row_sum.data_ptr()); - } - - const dim3 merge_grid( - (M + D256_BM32_PHASE_BLOCK_M - 1) / D256_BM32_PHASE_BLOCK_M, 1, - Q.size(0) * H); - flash_attention_forward_paged_d256_bm32_splitkv3_merge_kernel<<< - merge_grid, 256, 0, stream>>>(split_tmp_out.data_ptr(), - split_tmp_row_max.data_ptr(), - split_tmp_row_sum.data_ptr(), - reinterpret_cast<__half*>(Out.data_ptr()), - softmax_lse.data_ptr(), M); + split_tmp_row_sum.data_ptr(), + reinterpret_cast<__half*>(Out.data_ptr()), + softmax_lse.data_ptr(), M); } void launcher_flash_attention_forward_paged_d256_bm32_phase( - const torch::Tensor& Q, const torch::Tensor& K_cache, - const torch::Tensor& V_cache, torch::Tensor& Out, - torch::Tensor& softmax_lse, const torch::Tensor& block_table, - const torch::Tensor& seq_lens, float softmax_scale, bool is_causal, + const torch::Tensor& Q, + const torch::Tensor& K_cache, + const torch::Tensor& V_cache, + torch::Tensor& Out, + torch::Tensor& softmax_lse, + const torch::Tensor& block_table, + const torch::Tensor& seq_lens, + float softmax_scale, + bool is_causal, cudaStream_t stream) { - const bool use_all_p = - env_flag_default_enabled("VLLM_FLASH_V100_PREFILL_D256_BM32_ALL_P"); - const bool use_pair_scratch = - use_all_p && env_flag_default_enabled( - "VLLM_FLASH_V100_PREFILL_D256_BM32_PAIR_SCRATCH"); + const bool use_all_p = + env_flag_default_enabled("VLLM_FLASH_V100_PREFILL_D256_BM32_ALL_P"); + const bool use_pair_scratch = + use_all_p + && env_flag_default_enabled( + "VLLM_FLASH_V100_PREFILL_D256_BM32_PAIR_SCRATCH"); + + if (is_causal) { + if (use_all_p) { + if (use_pair_scratch) { + launch_flash_attention_forward_paged_d256_bm32_phase_kernel< + true, true, true>( + Q, K_cache, V_cache, Out, softmax_lse, block_table, + seq_lens, softmax_scale, stream); + } else { + launch_flash_attention_forward_paged_d256_bm32_phase_kernel< + true, true, false>( + Q, K_cache, V_cache, Out, softmax_lse, block_table, + seq_lens, softmax_scale, stream); + } + } else { + launch_flash_attention_forward_paged_d256_bm32_phase_kernel< + true, false, false>( + Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, + softmax_scale, stream); + } + return; + } - if (is_causal) { if (use_all_p) { - if (use_pair_scratch) { - launch_flash_attention_forward_paged_d256_bm32_phase_kernel( - Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, - softmax_scale, stream); - } else { - launch_flash_attention_forward_paged_d256_bm32_phase_kernel( + if (use_pair_scratch) { + launch_flash_attention_forward_paged_d256_bm32_phase_kernel< + false, true, true>( + Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, + softmax_scale, stream); + } else { + launch_flash_attention_forward_paged_d256_bm32_phase_kernel< + false, true, false>( + Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, + softmax_scale, stream); + } + } else { + launch_flash_attention_forward_paged_d256_bm32_phase_kernel< + false, false, false>( Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, softmax_scale, stream); - } - } else { - launch_flash_attention_forward_paged_d256_bm32_phase_kernel( - Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, - softmax_scale, stream); - } - return; - } - - if (use_all_p) { - if (use_pair_scratch) { - launch_flash_attention_forward_paged_d256_bm32_phase_kernel( - Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, - softmax_scale, stream); - } else { - launch_flash_attention_forward_paged_d256_bm32_phase_kernel( - Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, - softmax_scale, stream); - } - } else { - launch_flash_attention_forward_paged_d256_bm32_phase_kernel( - Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, - softmax_scale, stream); - } + } } -template +template void launcher_flash_attention_forward_paged( - const torch::Tensor& Q, const torch::Tensor& K_cache, - const torch::Tensor& V_cache, torch::Tensor& Out, - torch::Tensor& softmax_lse, const torch::Tensor& block_table, - const torch::Tensor& seq_lens, float softmax_scale, bool is_causal, - float k_scale, float v_scale, int window_size_left, int window_size_right, - cudaStream_t stream, const int* bfla_mask_ptr = nullptr, - int bfla_mask_block_n = 0, int64_t bfla_mask_stride_b = 0, - int64_t bfla_mask_stride_h = 0, int64_t bfla_mask_stride_q = 0, - int64_t bfla_mask_stride_k = 0, const int* anchor_lens = nullptr, - int anchored_window = 0) { - if (anchor_lens != nullptr && anchored_window > 0) { - TORCH_CHECK(KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP16, - "anchored decode window requires an fp16 KV cache"); - TORCH_CHECK(is_causal, "anchored decode window requires causal attention"); - TORCH_CHECK(window_size_left < 0 && window_size_right < 0, - "anchored decode window cannot be combined with " - "sliding-window attention"); - TORCH_CHECK(bfla_mask_ptr == nullptr, - "anchored decode window cannot be combined with a " - "block sparsity mask"); - launcher_flash_attention_forward_paged_impl( - Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, - bfla_mask_ptr, bfla_mask_block_n, bfla_mask_stride_b, - bfla_mask_stride_h, bfla_mask_stride_q, bfla_mask_stride_k, - softmax_scale, is_causal, k_scale, v_scale, window_size_left, - window_size_right, stream, anchor_lens, anchored_window); - return; - } - if constexpr (D == 256 && KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP16) { - const int M = Q.size(2); - const int page_block_size = K_cache.size(1); - const bool use_low_smem = - M > 1 && page_block_size >= 16 && (page_block_size % 16) == 0 && - env_flag_default_enabled("VLLM_FLASH_V100_PREFILL_D256_LOW_SMEM"); - if (use_low_smem) { - const bool use_d256_bm32_phase = - env_flag_default_enabled("VLLM_FLASH_V100_PREFILL_D256_BM32_PHASE") && - page_block_size == D256_BM32_PHASE_PAGE_BLOCK_SIZE && - M >= D256_BM32_PHASE_BLOCK_M && bfla_mask_ptr == nullptr && - window_size_left < 0 && window_size_right < 0; - if (use_d256_bm32_phase) { - launcher_flash_attention_forward_paged_d256_bm32_phase( - Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, - softmax_scale, is_causal, stream); - return; - } - const bool use_low_smem_contig_fast = - page_block_size == 16 || - env_flag_enabled("VLLM_FLASH_V100_PREFILL_CONTIG_FAST"); - const bool use_low_smem_scalar_qk = - env_flag_enabled("VLLM_FLASH_V100_PREFILL_D256_SCALAR_QK"); - const bool use_low_smem_bm32 = - env_flag_enabled("VLLM_FLASH_V100_PREFILL_D256_BM32"); - // Default only for the measured hybrid-cache page size. Other - // page sizes retain the prior layout unless explicitly enabled. - const bool use_low_smem_output_stride_268 = - page_block_size == 784 - ? env_flag_default_enabled( - "VLLM_FLASH_V100_PREFILL_D256_OUTPUT_STRIDE_268") - : env_flag_enabled( - "VLLM_FLASH_V100_PREFILL_D256_OUTPUT_STRIDE_268"); - const bool use_sw_pipeline = - page_block_size == 784 && - env_flag_enabled("VLLM_FLASH_V100_PREFILL_D256_SOFTWARE_PIPELINE"); - const bool use_sw_pipeline_qk = - page_block_size == 784 && - (use_sw_pipeline || - env_flag_default_enabled( - "VLLM_FLASH_V100_PREFILL_D256_SW_PIPELINE_QK")); - const bool use_sw_pipeline_pv = - page_block_size == 784 && - (use_sw_pipeline || - env_flag_default_enabled( - "VLLM_FLASH_V100_PREFILL_D256_SW_PIPELINE_PV")); - if (use_low_smem_contig_fast) { - if (use_low_smem_scalar_qk) { - if (use_low_smem_bm32) { - launcher_flash_attention_forward_paged_impl( + const torch::Tensor& Q, + const torch::Tensor& K_cache, + const torch::Tensor& V_cache, + torch::Tensor& Out, + torch::Tensor& softmax_lse, + const torch::Tensor& block_table, + const torch::Tensor& seq_lens, + float softmax_scale, + bool is_causal, + float k_scale, + float v_scale, + int window_size_left, + int window_size_right, + cudaStream_t stream, + const int* bfla_mask_ptr = nullptr, + int bfla_mask_block_n = 0, + int64_t bfla_mask_stride_b = 0, + int64_t bfla_mask_stride_h = 0, + int64_t bfla_mask_stride_q = 0, + int64_t bfla_mask_stride_k = 0 +) { + if constexpr (D == 256 && KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP16) { + const int M = Q.size(2); + const int page_block_size = K_cache.size(1); + const bool use_low_smem = + M > 1 && page_block_size >= 16 && (page_block_size % 16) == 0 + && env_flag_default_enabled( + "VLLM_FLASH_V100_PREFILL_D256_LOW_SMEM"); + if (use_low_smem) { + const bool use_d256_bm32_phase = + env_flag_default_enabled( + "VLLM_FLASH_V100_PREFILL_D256_BM32_PHASE") + && page_block_size == D256_BM32_PHASE_PAGE_BLOCK_SIZE + && M >= D256_BM32_PHASE_BLOCK_M + && bfla_mask_ptr == nullptr && window_size_left < 0 + && window_size_right < 0; + if (use_d256_bm32_phase) { + launcher_flash_attention_forward_paged_d256_bm32_phase( + Q, K_cache, V_cache, Out, softmax_lse, block_table, + seq_lens, softmax_scale, is_causal, stream); + return; + } + const bool use_low_smem_contig_fast = + page_block_size == 16 || + env_flag_enabled("VLLM_FLASH_V100_PREFILL_CONTIG_FAST"); + const bool use_low_smem_scalar_qk = + env_flag_enabled("VLLM_FLASH_V100_PREFILL_D256_SCALAR_QK"); + const bool use_low_smem_bm32 = + env_flag_enabled("VLLM_FLASH_V100_PREFILL_D256_BM32"); + // Default only for the measured hybrid-cache page size. Other + // page sizes retain the prior layout unless explicitly enabled. + const bool use_low_smem_output_stride_268 = + page_block_size == 784 + ? env_flag_default_enabled( + "VLLM_FLASH_V100_PREFILL_D256_OUTPUT_STRIDE_268") + : env_flag_enabled( + "VLLM_FLASH_V100_PREFILL_D256_OUTPUT_STRIDE_268"); + const bool use_sw_pipeline = + page_block_size == 784 + && env_flag_enabled( + "VLLM_FLASH_V100_PREFILL_D256_SOFTWARE_PIPELINE"); + const bool use_sw_pipeline_qk = + page_block_size == 784 + && (use_sw_pipeline + || env_flag_default_enabled( + "VLLM_FLASH_V100_PREFILL_D256_SW_PIPELINE_QK")); + const bool use_sw_pipeline_pv = + page_block_size == 784 + && (use_sw_pipeline + || env_flag_default_enabled( + "VLLM_FLASH_V100_PREFILL_D256_SW_PIPELINE_PV")); + if (use_low_smem_contig_fast) { + if (use_low_smem_scalar_qk) { + if (use_low_smem_bm32) { + launcher_flash_attention_forward_paged_impl< + D, KV_DTYPE, true, true, true, true>( + Q, K_cache, V_cache, Out, softmax_lse, block_table, + seq_lens, bfla_mask_ptr, bfla_mask_block_n, + bfla_mask_stride_b, bfla_mask_stride_h, + bfla_mask_stride_q, bfla_mask_stride_k, + softmax_scale, is_causal, k_scale, v_scale, + window_size_left, window_size_right, stream); + } else { + launcher_flash_attention_forward_paged_impl< + D, KV_DTYPE, true, true, true, false>( + Q, K_cache, V_cache, Out, softmax_lse, block_table, + seq_lens, bfla_mask_ptr, bfla_mask_block_n, + bfla_mask_stride_b, bfla_mask_stride_h, + bfla_mask_stride_q, bfla_mask_stride_k, + softmax_scale, is_causal, k_scale, v_scale, + window_size_left, window_size_right, stream); + } + } else { + if (use_low_smem_bm32) { + launcher_flash_attention_forward_paged_impl< + D, KV_DTYPE, true, true, false, true>( + Q, K_cache, V_cache, Out, softmax_lse, block_table, + seq_lens, bfla_mask_ptr, bfla_mask_block_n, + bfla_mask_stride_b, bfla_mask_stride_h, + bfla_mask_stride_q, bfla_mask_stride_k, + softmax_scale, is_causal, k_scale, v_scale, + window_size_left, window_size_right, stream); + } else { + launcher_flash_attention_forward_paged_impl< + D, KV_DTYPE, true, true, false, false>( + Q, K_cache, V_cache, Out, softmax_lse, block_table, + seq_lens, bfla_mask_ptr, bfla_mask_block_n, + bfla_mask_stride_b, bfla_mask_stride_h, + bfla_mask_stride_q, bfla_mask_stride_k, + softmax_scale, is_causal, k_scale, v_scale, + window_size_left, window_size_right, stream); + } + } + } else { + if (use_low_smem_scalar_qk) { + if (use_low_smem_bm32) { + launcher_flash_attention_forward_paged_impl< + D, KV_DTYPE, true, false, true, true>( + Q, K_cache, V_cache, Out, softmax_lse, block_table, + seq_lens, bfla_mask_ptr, bfla_mask_block_n, + bfla_mask_stride_b, bfla_mask_stride_h, + bfla_mask_stride_q, bfla_mask_stride_k, + softmax_scale, is_causal, k_scale, v_scale, + window_size_left, window_size_right, stream); + } else { + launcher_flash_attention_forward_paged_impl< + D, KV_DTYPE, true, false, true, false>( + Q, K_cache, V_cache, Out, softmax_lse, block_table, + seq_lens, bfla_mask_ptr, bfla_mask_block_n, + bfla_mask_stride_b, bfla_mask_stride_h, + bfla_mask_stride_q, bfla_mask_stride_k, + softmax_scale, is_causal, k_scale, v_scale, + window_size_left, window_size_right, stream); + } + } else { + if (use_low_smem_bm32) { + launcher_flash_attention_forward_paged_impl< + D, KV_DTYPE, true, false, false, true>( + Q, K_cache, V_cache, Out, softmax_lse, block_table, + seq_lens, bfla_mask_ptr, bfla_mask_block_n, + bfla_mask_stride_b, bfla_mask_stride_h, + bfla_mask_stride_q, bfla_mask_stride_k, + softmax_scale, is_causal, k_scale, v_scale, + window_size_left, window_size_right, stream); + } else { + if (use_low_smem_output_stride_268) { + if (use_sw_pipeline_qk) { + if (use_sw_pipeline_pv) { + launcher_flash_attention_forward_paged_impl< + D, KV_DTYPE, true, false, false, + false, true, true, true>( + Q, K_cache, V_cache, Out, + softmax_lse, block_table, seq_lens, + bfla_mask_ptr, bfla_mask_block_n, + bfla_mask_stride_b, + bfla_mask_stride_h, + bfla_mask_stride_q, + bfla_mask_stride_k, softmax_scale, + is_causal, k_scale, v_scale, + window_size_left, + window_size_right, stream); + } else { + launcher_flash_attention_forward_paged_impl< + D, KV_DTYPE, true, false, false, false, + true, true, false>( + Q, K_cache, V_cache, Out, softmax_lse, + block_table, seq_lens, bfla_mask_ptr, + bfla_mask_block_n, bfla_mask_stride_b, + bfla_mask_stride_h, + bfla_mask_stride_q, + bfla_mask_stride_k, softmax_scale, + is_causal, k_scale, v_scale, + window_size_left, window_size_right, + stream); + } + } else if (use_sw_pipeline_pv) { + launcher_flash_attention_forward_paged_impl< + D, KV_DTYPE, true, false, false, false, + true, false, true>( + Q, K_cache, V_cache, Out, softmax_lse, + block_table, seq_lens, bfla_mask_ptr, + bfla_mask_block_n, bfla_mask_stride_b, + bfla_mask_stride_h, bfla_mask_stride_q, + bfla_mask_stride_k, softmax_scale, + is_causal, k_scale, v_scale, + window_size_left, window_size_right, + stream); + } else { + launcher_flash_attention_forward_paged_impl< + D, KV_DTYPE, true, false, false, false, + true>( + Q, K_cache, V_cache, Out, softmax_lse, + block_table, seq_lens, bfla_mask_ptr, + bfla_mask_block_n, bfla_mask_stride_b, + bfla_mask_stride_h, bfla_mask_stride_q, + bfla_mask_stride_k, softmax_scale, + is_causal, k_scale, v_scale, + window_size_left, window_size_right, + stream); + } + } else { + launcher_flash_attention_forward_paged_impl< + D, KV_DTYPE, true, false, false, false>( + Q, K_cache, V_cache, Out, softmax_lse, + block_table, seq_lens, bfla_mask_ptr, + bfla_mask_block_n, bfla_mask_stride_b, + bfla_mask_stride_h, bfla_mask_stride_q, + bfla_mask_stride_k, softmax_scale, is_causal, + k_scale, v_scale, window_size_left, + window_size_right, stream); + } + } + } + } + } else { + launcher_flash_attention_forward_paged_impl< + D, KV_DTYPE, false, false, false, false>( Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, bfla_mask_ptr, bfla_mask_block_n, bfla_mask_stride_b, bfla_mask_stride_h, bfla_mask_stride_q, bfla_mask_stride_k, softmax_scale, is_causal, k_scale, v_scale, window_size_left, window_size_right, stream); - } else { - launcher_flash_attention_forward_paged_impl( - Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, - bfla_mask_ptr, bfla_mask_block_n, bfla_mask_stride_b, - bfla_mask_stride_h, bfla_mask_stride_q, bfla_mask_stride_k, + } + } else { + launcher_flash_attention_forward_paged_impl< + D, KV_DTYPE, false, false, false, false>( + Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, + bfla_mask_ptr, bfla_mask_block_n, bfla_mask_stride_b, + bfla_mask_stride_h, bfla_mask_stride_q, bfla_mask_stride_k, + softmax_scale, is_causal, k_scale, v_scale, window_size_left, + window_size_right, stream); + } +} + +at::Tensor flash_attention_prefill_paged_splitkv( + const at::Tensor& q, + const at::Tensor& k_cache, + const at::Tensor& v_cache, + std::optional& out_, + const at::Tensor& block_table, + const at::Tensor& seq_lens, + const float softmax_scale, + const std::string& kv_cache_dtype, + const float k_scale, + const float v_scale, + const bool is_causal, + const int window_size_left, + const int window_size_right, + const int split_kv_tokens, + const int max_seq_len_hint +) { + TORCH_CHECK(q.dtype() == torch::kFloat16, "q must be fp16"); + const int kv_dtype_code = kv_cache_dtype_code_from_string(kv_cache_dtype); + TORCH_CHECK(kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP16, + "split-KV paged prefill supports fp16 KV cache only"); + TORCH_CHECK(k_cache.dtype() == torch::kFloat16, "k_cache must be fp16"); + TORCH_CHECK(v_cache.dtype() == torch::kFloat16, "v_cache must be fp16"); + TORCH_CHECK(q.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda(), + "Tensors must be on CUDA"); + TORCH_CHECK(block_table.is_cuda() && seq_lens.is_cuda(), + "block_table and seq_lens must be CUDA tensors"); + TORCH_CHECK(q.stride(-1) == 1 && k_cache.stride(-1) == 1 && + v_cache.stride(-1) == 1, + "Last dim must be contiguous"); + TORCH_CHECK(k_cache.stride(1) % 8 == 0 && k_cache.stride(2) % 8 == 0 && + v_cache.stride(1) % 8 == 0 && v_cache.stride(2) % 8 == 0, + "Paged KV strides must be divisible by 8 half elements"); + + 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 num_kv_heads = k_cache.size(2); + const int page_block_size = k_cache.size(1); + + TORCH_CHECK(D == 256, "split-KV paged prefill supports D=256 only"); + TORCH_CHECK(H % num_kv_heads == 0, + "num_heads must be divisible by num_kv_heads"); + TORCH_CHECK(M > 1, "split-KV paged prefill requires prefill M > 1"); + TORCH_CHECK(page_block_size >= 16 && (page_block_size % 16) == 0, + "split-KV paged prefill requires page block size multiple of 16"); + TORCH_CHECK(window_size_left >= -1 && window_size_right >= -1, + "window sizes must be >= -1"); + + at::Tensor out_fp16 = out_.has_value() ? out_.value() : torch::zeros_like(q); + auto softmax_lse = 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."); + + const bool use_low_smem_contig_fast = + page_block_size == 16 || + env_flag_enabled("VLLM_FLASH_V100_PREFILL_CONTIG_FAST"); + const bool use_low_smem_scalar_qk = + env_flag_enabled("VLLM_FLASH_V100_PREFILL_D256_SCALAR_QK"); + const int block_n = + use_low_smem_scalar_qk + ? BLOCK_N_256_LOW_SMEM_SCALAR_QK + : BLOCK_N_256_LOW_SMEM; + const int split_tokens_rounded = std::max(split_kv_tokens, block_n); + const int split_kv_tiles = + std::max(1, (split_tokens_rounded + block_n - 1) / block_n); + const int max_hint = std::max(max_seq_len_hint, 1); + const int max_kv_tiles = std::max(1, (max_hint + block_n - 1) / block_n); + const int num_partitions = + std::max(1, (max_kv_tiles + split_kv_tiles - 1) / split_kv_tiles); + + if (num_partitions <= 1) { + launcher_flash_attention_forward_paged<256, flash_v100::KV_CACHE_DTYPE_FP16>( + q, k_cache, v_cache, out_fp16, softmax_lse, block_table, seq_lens, + softmax_scale, is_causal, k_scale, v_scale, window_size_left, + window_size_right, stream); + return out_fp16; + } + + auto tmp_options = torch::dtype(torch::kFloat32).device(q.device()); + auto split_tmp_out = + torch::empty({B, H, num_partitions, M, D}, tmp_options); + auto split_tmp_row_max = + torch::empty({B, H, num_partitions, M}, tmp_options); + auto split_tmp_row_sum = + torch::empty({B, H, num_partitions, M}, tmp_options); + + if (use_low_smem_contig_fast) { + if (use_low_smem_scalar_qk) { + launcher_flash_attention_forward_paged_splitkv_impl< + 256, true, true>( + q, k_cache, v_cache, out_fp16, softmax_lse, split_tmp_out, + split_tmp_row_max, split_tmp_row_sum, block_table, seq_lens, softmax_scale, is_causal, k_scale, v_scale, window_size_left, - window_size_right, stream); - } + window_size_right, split_kv_tokens, max_seq_len_hint, stream); } else { - if (use_low_smem_bm32) { - launcher_flash_attention_forward_paged_impl( - Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, - bfla_mask_ptr, bfla_mask_block_n, bfla_mask_stride_b, - bfla_mask_stride_h, bfla_mask_stride_q, bfla_mask_stride_k, - softmax_scale, is_causal, k_scale, v_scale, window_size_left, - window_size_right, stream); - } else { - launcher_flash_attention_forward_paged_impl( - Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, - bfla_mask_ptr, bfla_mask_block_n, bfla_mask_stride_b, - bfla_mask_stride_h, bfla_mask_stride_q, bfla_mask_stride_k, + launcher_flash_attention_forward_paged_splitkv_impl< + 256, true, false>( + q, k_cache, v_cache, out_fp16, softmax_lse, split_tmp_out, + split_tmp_row_max, split_tmp_row_sum, block_table, seq_lens, softmax_scale, is_causal, k_scale, v_scale, window_size_left, - window_size_right, stream); - } + window_size_right, split_kv_tokens, max_seq_len_hint, stream); } - } else { + } else { if (use_low_smem_scalar_qk) { - if (use_low_smem_bm32) { - launcher_flash_attention_forward_paged_impl( - Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, - bfla_mask_ptr, bfla_mask_block_n, bfla_mask_stride_b, - bfla_mask_stride_h, bfla_mask_stride_q, bfla_mask_stride_k, + launcher_flash_attention_forward_paged_splitkv_impl< + 256, false, true>( + q, k_cache, v_cache, out_fp16, softmax_lse, split_tmp_out, + split_tmp_row_max, split_tmp_row_sum, block_table, seq_lens, softmax_scale, is_causal, k_scale, v_scale, window_size_left, - window_size_right, stream); - } else { - launcher_flash_attention_forward_paged_impl( - Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, - bfla_mask_ptr, bfla_mask_block_n, bfla_mask_stride_b, - bfla_mask_stride_h, bfla_mask_stride_q, bfla_mask_stride_k, - softmax_scale, is_causal, k_scale, v_scale, window_size_left, - window_size_right, stream); - } + window_size_right, split_kv_tokens, max_seq_len_hint, stream); } else { - if (use_low_smem_bm32) { - launcher_flash_attention_forward_paged_impl( - Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, - bfla_mask_ptr, bfla_mask_block_n, bfla_mask_stride_b, - bfla_mask_stride_h, bfla_mask_stride_q, bfla_mask_stride_k, + launcher_flash_attention_forward_paged_splitkv_impl< + 256, false, false>( + q, k_cache, v_cache, out_fp16, softmax_lse, split_tmp_out, + split_tmp_row_max, split_tmp_row_sum, block_table, seq_lens, softmax_scale, is_causal, k_scale, v_scale, window_size_left, - window_size_right, stream); - } else { - if (use_low_smem_output_stride_268) { - if (use_sw_pipeline_qk) { - if (use_sw_pipeline_pv) { - launcher_flash_attention_forward_paged_impl< - D, KV_DTYPE, true, false, false, false, true, true, true>( - Q, K_cache, V_cache, Out, softmax_lse, block_table, - seq_lens, bfla_mask_ptr, bfla_mask_block_n, - bfla_mask_stride_b, bfla_mask_stride_h, - bfla_mask_stride_q, bfla_mask_stride_k, softmax_scale, - is_causal, k_scale, v_scale, window_size_left, - window_size_right, stream); - } else { - launcher_flash_attention_forward_paged_impl< - D, KV_DTYPE, true, false, false, false, true, true, - false>(Q, K_cache, V_cache, Out, softmax_lse, block_table, - seq_lens, bfla_mask_ptr, bfla_mask_block_n, - bfla_mask_stride_b, bfla_mask_stride_h, - bfla_mask_stride_q, bfla_mask_stride_k, - softmax_scale, is_causal, k_scale, v_scale, - window_size_left, window_size_right, stream); - } - } else if (use_sw_pipeline_pv) { - launcher_flash_attention_forward_paged_impl< - D, KV_DTYPE, true, false, false, false, true, false, true>( - Q, K_cache, V_cache, Out, softmax_lse, block_table, - seq_lens, bfla_mask_ptr, bfla_mask_block_n, - bfla_mask_stride_b, bfla_mask_stride_h, bfla_mask_stride_q, - bfla_mask_stride_k, softmax_scale, is_causal, k_scale, - v_scale, window_size_left, window_size_right, stream); - } else { - launcher_flash_attention_forward_paged_impl< - D, KV_DTYPE, true, false, false, false, true>( - Q, K_cache, V_cache, Out, softmax_lse, block_table, - seq_lens, bfla_mask_ptr, bfla_mask_block_n, - bfla_mask_stride_b, bfla_mask_stride_h, bfla_mask_stride_q, - bfla_mask_stride_k, softmax_scale, is_causal, k_scale, - v_scale, window_size_left, window_size_right, stream); - } - } else { - launcher_flash_attention_forward_paged_impl( - Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, - bfla_mask_ptr, bfla_mask_block_n, bfla_mask_stride_b, - bfla_mask_stride_h, bfla_mask_stride_q, bfla_mask_stride_k, - softmax_scale, is_causal, k_scale, v_scale, window_size_left, - window_size_right, stream); - } - } + window_size_right, split_kv_tokens, max_seq_len_hint, stream); } - } - } else { - launcher_flash_attention_forward_paged_impl( - Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, - bfla_mask_ptr, bfla_mask_block_n, bfla_mask_stride_b, - bfla_mask_stride_h, bfla_mask_stride_q, bfla_mask_stride_k, - softmax_scale, is_causal, k_scale, v_scale, window_size_left, - window_size_right, stream); - } - } else { - launcher_flash_attention_forward_paged_impl( - Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, - bfla_mask_ptr, bfla_mask_block_n, bfla_mask_stride_b, - bfla_mask_stride_h, bfla_mask_stride_q, bfla_mask_stride_k, - softmax_scale, is_causal, k_scale, v_scale, window_size_left, - window_size_right, stream); - } + } + + return out_fp16; } -at::Tensor flash_attention_prefill_paged_splitkv( - const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, - std::optional& out_, const at::Tensor& block_table, - const at::Tensor& seq_lens, const float softmax_scale, - const std::string& kv_cache_dtype, const float k_scale, const float v_scale, - const bool is_causal, const int window_size_left, - const int window_size_right, const int split_kv_tokens, - const int max_seq_len_hint) { - TORCH_CHECK(q.dtype() == torch::kFloat16, "q must be fp16"); - const int kv_dtype_code = kv_cache_dtype_code_from_string(kv_cache_dtype); - TORCH_CHECK( - kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP16 || - kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP8_E5M2, - "split-KV paged prefill supports fp16 and fp8_e5m2 KV cache only"); - if (kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP16) { - TORCH_CHECK(k_cache.dtype() == torch::kFloat16, - "fp16 split-KV requires fp16 k_cache"); - TORCH_CHECK(v_cache.dtype() == torch::kFloat16, - "fp16 split-KV requires fp16 v_cache"); - } else { - TORCH_CHECK(k_cache.dtype() == torch::kUInt8, - "fp8_e5m2 split-KV requires uint8 k_cache"); - TORCH_CHECK(v_cache.dtype() == torch::kUInt8, - "fp8_e5m2 split-KV requires uint8 v_cache"); - TORCH_CHECK(k_scale > 0.f && v_scale > 0.f, - "fp8_e5m2 split-KV requires positive K/V scales"); - } - TORCH_CHECK(q.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda(), - "Tensors must be on CUDA"); - TORCH_CHECK(block_table.is_cuda() && seq_lens.is_cuda(), - "block_table and seq_lens must be CUDA tensors"); - TORCH_CHECK( - q.stride(-1) == 1 && k_cache.stride(-1) == 1 && v_cache.stride(-1) == 1, - "Last dim must be contiguous"); - TORCH_CHECK(k_cache.stride(1) % 8 == 0 && k_cache.stride(2) % 8 == 0 && - v_cache.stride(1) % 8 == 0 && v_cache.stride(2) % 8 == 0, - "Paged KV strides must be divisible by 8 half elements"); - - 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 num_kv_heads = k_cache.size(2); - const int page_block_size = k_cache.size(1); - - TORCH_CHECK(D == 256, "split-KV paged prefill supports D=256 only"); - TORCH_CHECK(H % num_kv_heads == 0, - "num_heads must be divisible by num_kv_heads"); - TORCH_CHECK(M > 1, "split-KV paged prefill requires prefill M > 1"); - TORCH_CHECK(page_block_size >= 16 && (page_block_size % 16) == 0, - "split-KV paged prefill requires page block size multiple of 16"); - TORCH_CHECK(window_size_left >= -1 && window_size_right >= -1, - "window sizes must be >= -1"); - - at::Tensor out_fp16 = out_.has_value() ? out_.value() : torch::zeros_like(q); - auto softmax_lse = - 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."); - - const bool fp16_kv = kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP16; - const bool use_low_smem_contig_fast = - fp16_kv && (page_block_size == 16 || - env_flag_enabled("VLLM_FLASH_V100_PREFILL_CONTIG_FAST")); - const bool use_low_smem_scalar_qk = - fp16_kv && env_flag_enabled("VLLM_FLASH_V100_PREFILL_D256_SCALAR_QK"); - const int block_n = - fp16_kv ? (use_low_smem_scalar_qk ? BLOCK_N_256_LOW_SMEM_SCALAR_QK - : BLOCK_N_256_LOW_SMEM) - : KernelConfig<256, false>::BLOCK_N; - const int split_tokens_rounded = std::max(split_kv_tokens, block_n); - const int split_kv_tiles = - std::max(1, (split_tokens_rounded + block_n - 1) / block_n); - const int max_hint = std::max(max_seq_len_hint, 1); - const int max_kv_tiles = std::max(1, (max_hint + block_n - 1) / block_n); - const int num_partitions = - std::max(1, (max_kv_tiles + split_kv_tiles - 1) / split_kv_tiles); - - if (num_partitions <= 1) { - if (fp16_kv) { - launcher_flash_attention_forward_paged<256, - flash_v100::KV_CACHE_DTYPE_FP16>( - q, k_cache, v_cache, out_fp16, softmax_lse, block_table, seq_lens, - softmax_scale, is_causal, k_scale, v_scale, window_size_left, - window_size_right, stream); +at::Tensor flash_attention_prefill_paged( + const at::Tensor& q, + const at::Tensor& k_cache, + const at::Tensor& v_cache, + std::optional& out_, + const at::Tensor& block_table, + const at::Tensor& seq_lens, + const float softmax_scale, + const std::string& kv_cache_dtype, + const float k_scale, + const float v_scale, + const bool is_causal, + const int window_size_left, + const int window_size_right +) { + TORCH_CHECK(q.dtype() == torch::kFloat16, "q must be fp16"); + const int kv_dtype_code = kv_cache_dtype_code_from_string(kv_cache_dtype); + TORCH_CHECK(kv_dtype_code >= 0, "Unsupported kv_cache_dtype: ", kv_cache_dtype); + if (kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP16) { + TORCH_CHECK(k_cache.dtype() == torch::kFloat16, "k_cache must be fp16"); + TORCH_CHECK(v_cache.dtype() == torch::kFloat16, "v_cache must be fp16"); } else { - launcher_flash_attention_forward_paged< - 256, flash_v100::KV_CACHE_DTYPE_FP8_E5M2>( - q, k_cache, v_cache, out_fp16, softmax_lse, block_table, seq_lens, - softmax_scale, is_causal, k_scale, v_scale, window_size_left, - window_size_right, stream); + TORCH_CHECK(k_cache.dtype() == torch::kUInt8, + "fp8 k_cache must be stored as uint8"); + TORCH_CHECK(v_cache.dtype() == torch::kUInt8, + "fp8 v_cache must be stored as uint8"); + TORCH_CHECK(k_scale > 0.f && v_scale > 0.f, + "fp8 k/v scales must be positive"); } - return out_fp16; - } - - auto tmp_options = torch::dtype(torch::kFloat32).device(q.device()); - auto split_tmp_out = torch::empty({B, H, num_partitions, M, D}, tmp_options); - auto split_tmp_row_max = torch::empty({B, H, num_partitions, M}, tmp_options); - auto split_tmp_row_sum = torch::empty({B, H, num_partitions, M}, tmp_options); - - if (!fp16_kv) { - launcher_flash_attention_forward_paged_splitkv_impl< - 256, flash_v100::KV_CACHE_DTYPE_FP8_E5M2, false, false, false>( - q, k_cache, v_cache, out_fp16, softmax_lse, split_tmp_out, - split_tmp_row_max, split_tmp_row_sum, block_table, seq_lens, - softmax_scale, is_causal, k_scale, v_scale, window_size_left, - window_size_right, split_kv_tokens, max_seq_len_hint, stream); - } else if (use_low_smem_contig_fast) { - if (use_low_smem_scalar_qk) { - launcher_flash_attention_forward_paged_splitkv_impl< - 256, flash_v100::KV_CACHE_DTYPE_FP16, true, true, true>( - q, k_cache, v_cache, out_fp16, softmax_lse, split_tmp_out, - split_tmp_row_max, split_tmp_row_sum, block_table, seq_lens, - softmax_scale, is_causal, k_scale, v_scale, window_size_left, - window_size_right, split_kv_tokens, max_seq_len_hint, stream); - } else { - launcher_flash_attention_forward_paged_splitkv_impl< - 256, flash_v100::KV_CACHE_DTYPE_FP16, true, true, false>( - q, k_cache, v_cache, out_fp16, softmax_lse, split_tmp_out, - split_tmp_row_max, split_tmp_row_sum, block_table, seq_lens, - softmax_scale, is_causal, k_scale, v_scale, window_size_left, - window_size_right, split_kv_tokens, max_seq_len_hint, stream); - } - } else { - if (use_low_smem_scalar_qk) { - launcher_flash_attention_forward_paged_splitkv_impl< - 256, flash_v100::KV_CACHE_DTYPE_FP16, true, false, true>( - q, k_cache, v_cache, out_fp16, softmax_lse, split_tmp_out, - split_tmp_row_max, split_tmp_row_sum, block_table, seq_lens, - softmax_scale, is_causal, k_scale, v_scale, window_size_left, - window_size_right, split_kv_tokens, max_seq_len_hint, stream); - } else { - launcher_flash_attention_forward_paged_splitkv_impl< - 256, flash_v100::KV_CACHE_DTYPE_FP16, true, false, false>( - q, k_cache, v_cache, out_fp16, softmax_lse, split_tmp_out, - split_tmp_row_max, split_tmp_row_sum, block_table, seq_lens, - softmax_scale, is_causal, k_scale, v_scale, window_size_left, - window_size_right, split_kv_tokens, max_seq_len_hint, stream); + TORCH_CHECK(q.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda(), + "Tensors must be on CUDA"); + TORCH_CHECK(block_table.is_cuda() && seq_lens.is_cuda(), + "block_table and seq_lens must be CUDA tensors"); + TORCH_CHECK(q.stride(-1) == 1 && k_cache.stride(-1) == 1 && + v_cache.stride(-1) == 1, + "Last dim must be contiguous"); + TORCH_CHECK(k_cache.stride(1) % 8 == 0 && k_cache.stride(2) % 8 == 0 && + v_cache.stride(1) % 8 == 0 && v_cache.stride(2) % 8 == 0, + "Paged KV strides must be divisible by 8 half elements"); + + 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 num_kv_heads = k_cache.size(2); + + TORCH_CHECK(D <= 256 && D % 8 == 0 && D % 2 == 0, + "D must be even, <=256, multiple of 8"); + TORCH_CHECK(H % num_kv_heads == 0, + "num_heads must be divisible by num_kv_heads"); + TORCH_CHECK(window_size_left >= -1 && window_size_right >= -1, + "window sizes must be >= -1"); + + at::Tensor out_fp16 = out_.has_value() ? out_.value() : torch::zeros_like(q); + auto softmax_lse = 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."); + + #define LAUNCH_PAGED_TYPED(HDIM, KV_DTYPE_CODE) \ + launcher_flash_attention_forward_paged( \ + q, k_cache, v_cache, out_fp16, softmax_lse, block_table, seq_lens, \ + softmax_scale, is_causal, k_scale, v_scale, window_size_left, \ + window_size_right, stream) + + #define LAUNCH_PAGED_BY_KV(HDIM) \ + do { \ + switch (kv_dtype_code) { \ + case flash_v100::KV_CACHE_DTYPE_FP16: \ + LAUNCH_PAGED_TYPED(HDIM, flash_v100::KV_CACHE_DTYPE_FP16); \ + break; \ + case flash_v100::KV_CACHE_DTYPE_FP8_E4M3: \ + LAUNCH_PAGED_TYPED(HDIM, flash_v100::KV_CACHE_DTYPE_FP8_E4M3); \ + break; \ + case flash_v100::KV_CACHE_DTYPE_FP8_E5M2: \ + LAUNCH_PAGED_TYPED(HDIM, flash_v100::KV_CACHE_DTYPE_FP8_E5M2); \ + break; \ + default: \ + TORCH_CHECK(false, "Unsupported kv_cache_dtype: ", kv_cache_dtype); \ + } \ + } while (0) + + switch (D) { + case 16: + LAUNCH_PAGED_BY_KV(16); + break; + case 32: + LAUNCH_PAGED_BY_KV(32); + break; + case 64: + LAUNCH_PAGED_BY_KV(64); + break; + case 128: + LAUNCH_PAGED_BY_KV(128); + break; + case 256: + LAUNCH_PAGED_BY_KV(256); + break; + default: + TORCH_CHECK(false, "Unsupported D: ", D); } - } - return out_fp16; -} + #undef LAUNCH_PAGED_BY_KV + #undef LAUNCH_PAGED_TYPED -at::Tensor flash_attention_prefill_paged( - const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, - std::optional& out_, const at::Tensor& block_table, - const at::Tensor& seq_lens, const float softmax_scale, - const std::string& kv_cache_dtype, const float k_scale, const float v_scale, - const bool is_causal, const int window_size_left, - const int window_size_right, const std::optional& anchor_lens, - const int64_t anchored_window) { - TORCH_CHECK(q.dtype() == torch::kFloat16, "q must be fp16"); - const int kv_dtype_code = kv_cache_dtype_code_from_string(kv_cache_dtype); - TORCH_CHECK(kv_dtype_code >= 0, - "Unsupported kv_cache_dtype: ", kv_cache_dtype); - if (kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP16) { - TORCH_CHECK(k_cache.dtype() == torch::kFloat16, "k_cache must be fp16"); - TORCH_CHECK(v_cache.dtype() == torch::kFloat16, "v_cache must be fp16"); - } else { - TORCH_CHECK(k_cache.dtype() == torch::kUInt8, - "fp8 k_cache must be stored as uint8"); - TORCH_CHECK(v_cache.dtype() == torch::kUInt8, - "fp8 v_cache must be stored as uint8"); - TORCH_CHECK(k_scale > 0.f && v_scale > 0.f, - "fp8 k/v scales must be positive"); - } - TORCH_CHECK(q.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda(), - "Tensors must be on CUDA"); - TORCH_CHECK(block_table.is_cuda() && seq_lens.is_cuda(), - "block_table and seq_lens must be CUDA tensors"); - TORCH_CHECK( - q.stride(-1) == 1 && k_cache.stride(-1) == 1 && v_cache.stride(-1) == 1, - "Last dim must be contiguous"); - TORCH_CHECK(k_cache.stride(1) % 8 == 0 && k_cache.stride(2) % 8 == 0 && - v_cache.stride(1) % 8 == 0 && v_cache.stride(2) % 8 == 0, - "Paged KV strides must be divisible by 8 half elements"); - - 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 num_kv_heads = k_cache.size(2); - - TORCH_CHECK(D <= 256 && D % 8 == 0 && D % 2 == 0, - "D must be even, <=256, multiple of 8"); - TORCH_CHECK(H % num_kv_heads == 0, - "num_heads must be divisible by num_kv_heads"); - TORCH_CHECK(window_size_left >= -1 && window_size_right >= -1, - "window sizes must be >= -1"); - - TORCH_CHECK(anchored_window >= 0 && anchored_window <= INT_MAX, - "anchored_window must fit in a non-negative int32"); - TORCH_CHECK((anchored_window > 0) == anchor_lens.has_value(), - "anchor_lens and a positive anchored_window must be provided " - "together"); - const bool use_anchored = anchor_lens.has_value(); - const int* anchor_lens_ptr = nullptr; - if (use_anchored) { - const at::Tensor& anchors = anchor_lens.value(); - TORCH_CHECK(anchors.is_cuda(), "anchor_lens must be on CUDA"); - TORCH_CHECK(anchors.device() == q.device(), - "anchor_lens must be on the q device"); - TORCH_CHECK(anchors.dtype() == torch::kInt32, "anchor_lens must be int32"); - TORCH_CHECK(anchors.dim() == 1 && anchors.size(0) == B, - "anchor_lens must have shape [B]"); - TORCH_CHECK(anchors.is_contiguous(), "anchor_lens must be contiguous"); - TORCH_CHECK(kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP16, - "anchored decode window requires an fp16 KV cache"); - TORCH_CHECK(is_causal, "anchored decode window requires causal attention"); - TORCH_CHECK(window_size_left == -1 && window_size_right == -1, - "anchored decode window cannot be combined with " - "sliding-window attention"); - anchor_lens_ptr = anchors.data_ptr(); - } - - at::Tensor out_fp16 = out_.has_value() ? out_.value() : torch::zeros_like(q); - auto softmax_lse = - 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."); - -#define LAUNCH_PAGED_TYPED(HDIM, KV_DTYPE_CODE) \ - launcher_flash_attention_forward_paged( \ - q, k_cache, v_cache, out_fp16, softmax_lse, block_table, seq_lens, \ - softmax_scale, is_causal, k_scale, v_scale, window_size_left, \ - window_size_right, stream, nullptr, 0, 0, 0, 0, 0, anchor_lens_ptr, \ - static_cast(anchored_window)) - -#define LAUNCH_PAGED_BY_KV(HDIM) \ - do { \ - switch (kv_dtype_code) { \ - case flash_v100::KV_CACHE_DTYPE_FP16: \ - LAUNCH_PAGED_TYPED(HDIM, flash_v100::KV_CACHE_DTYPE_FP16); \ - break; \ - case flash_v100::KV_CACHE_DTYPE_FP8_E4M3: \ - LAUNCH_PAGED_TYPED(HDIM, flash_v100::KV_CACHE_DTYPE_FP8_E4M3); \ - break; \ - case flash_v100::KV_CACHE_DTYPE_FP8_E5M2: \ - LAUNCH_PAGED_TYPED(HDIM, flash_v100::KV_CACHE_DTYPE_FP8_E5M2); \ - break; \ - default: \ - TORCH_CHECK(false, "Unsupported kv_cache_dtype: ", kv_cache_dtype); \ - } \ - } while (0) - - switch (D) { - case 16: - LAUNCH_PAGED_BY_KV(16); - break; - case 32: - LAUNCH_PAGED_BY_KV(32); - break; - case 64: - LAUNCH_PAGED_BY_KV(64); - break; - case 128: - LAUNCH_PAGED_BY_KV(128); - break; - case 256: - LAUNCH_PAGED_BY_KV(256); - break; - default: - TORCH_CHECK(false, "Unsupported D: ", D); - } - -#undef LAUNCH_PAGED_BY_KV -#undef LAUNCH_PAGED_TYPED - - return out_fp16; + return out_fp16; } std::vector flash_attention_prefill_paged_d256_bm32_allp_pair_scratch( - const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, - std::optional& out_, std::optional& softmax_lse_, - const at::Tensor& block_table, const at::Tensor& seq_lens, - const float softmax_scale) { - TORCH_CHECK(q.dim() == 4, "q must have shape [B, H, M, D]"); - TORCH_CHECK(k_cache.dim() == 4 && v_cache.dim() == 4, - "k_cache and v_cache must have shape [blocks, page, Hkv, D]"); - TORCH_CHECK(block_table.dim() == 2, - "block_table must have shape [B, max_num_blocks]"); - TORCH_CHECK(seq_lens.dim() == 1, "seq_lens must have shape [B]"); - TORCH_CHECK(q.dtype() == torch::kFloat16, "q must be fp16"); - TORCH_CHECK(k_cache.dtype() == torch::kFloat16, "k_cache must be fp16"); - TORCH_CHECK(v_cache.dtype() == torch::kFloat16, "v_cache must be fp16"); - 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.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda(), - "q, k_cache, and v_cache must be CUDA tensors"); - TORCH_CHECK(block_table.is_cuda() && seq_lens.is_cuda(), - "block_table and seq_lens must be CUDA tensors"); - TORCH_CHECK( - q.device() == k_cache.device() && q.device() == v_cache.device() && - q.device() == block_table.device() && q.device() == seq_lens.device(), - "all input tensors must be on the q device"); - TORCH_CHECK(q.is_contiguous(), "q must be contiguous [B, H, M, D]"); - TORCH_CHECK(block_table.is_contiguous(), "block_table must be contiguous"); - TORCH_CHECK(seq_lens.is_contiguous(), "seq_lens must be contiguous"); - TORCH_CHECK(k_cache.sizes() == v_cache.sizes(), - "k_cache and v_cache must have the same shape"); - TORCH_CHECK(k_cache.stride(-1) == 1 && v_cache.stride(-1) == 1, - "K/V head dimensions must be contiguous"); - TORCH_CHECK(k_cache.stride(1) % 8 == 0 && k_cache.stride(2) % 8 == 0 && - v_cache.stride(1) % 8 == 0 && v_cache.stride(2) % 8 == 0, - "K/V strides must be divisible by 8 half elements"); - - 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 num_kv_heads = k_cache.size(2); - - TORCH_CHECK(B > 0 && H > 0, "q batch and head dimensions must be positive"); - TORCH_CHECK(M >= D256_BM32_PHASE_BLOCK_M, - "fixed BM32 prefill requires M >= ", D256_BM32_PHASE_BLOCK_M); - TORCH_CHECK(D == D256_BM32_PHASE_D, - "fixed BM32 prefill requires D=", D256_BM32_PHASE_D); - TORCH_CHECK(k_cache.size(0) > 0 && num_kv_heads > 0, - "K/V cache block and head dimensions must be positive"); - TORCH_CHECK(k_cache.size(1) == D256_BM32_PHASE_PAGE_BLOCK_SIZE, - "fixed BM32 prefill requires page size ", - D256_BM32_PHASE_PAGE_BLOCK_SIZE); - TORCH_CHECK(k_cache.size(3) == D256_BM32_PHASE_D, - "K/V cache head dimension must be D=", D256_BM32_PHASE_D); - TORCH_CHECK(H % num_kv_heads == 0, - "num_q_heads must be divisible by num_kv_heads"); - TORCH_CHECK(block_table.size(0) == B && seq_lens.size(0) == B, - "block_table and seq_lens batch dimensions must equal q batch"); - TORCH_CHECK(block_table.size(1) > 0, - "block_table must provide at least one page index"); - - if (out_.has_value()) { - const at::Tensor& out = out_.value(); - TORCH_CHECK(out.is_cuda() && out.device() == q.device(), - "out must be a CUDA tensor on the q device"); - TORCH_CHECK(out.dtype() == torch::kFloat16, "out must be fp16"); - TORCH_CHECK(out.sizes() == q.sizes(), "out must have shape [B, H, M, D]"); - TORCH_CHECK(out.is_contiguous(), "out must be contiguous [B, H, M, D]"); - } - if (softmax_lse_.has_value()) { - const at::Tensor& softmax_lse = softmax_lse_.value(); - TORCH_CHECK(softmax_lse.is_cuda() && softmax_lse.device() == q.device(), - "softmax_lse must be a CUDA tensor on the q device"); - TORCH_CHECK(softmax_lse.dtype() == torch::kFloat32, - "softmax_lse must be fp32"); - TORCH_CHECK(softmax_lse.dim() == 3 && softmax_lse.size(0) == B && - softmax_lse.size(1) == H && softmax_lse.size(2) == M, - "softmax_lse must have shape [B, H, M]"); - TORCH_CHECK(softmax_lse.is_contiguous(), - "softmax_lse must be contiguous [B, H, M]"); - } - - c10::cuda::CUDAGuard device_guard(q.device()); - auto stream = at::cuda::getCurrentCUDAStream().stream(); - cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; - const cudaError_t capture_error = - cudaStreamIsCapturing(stream, &capture_status); - TORCH_CHECK(capture_error == cudaSuccess, "cudaStreamIsCapturing failed: ", - cudaGetErrorString(capture_error)); - TORCH_CHECK(capture_status == cudaStreamCaptureStatusNone, - "fixed BM32 prefill CUDA graph capture is not validated"); - - auto props = at::cuda::getCurrentDeviceProperties(); - TORCH_CHECK(props->major == 7 && props->minor == 0, - "fixed BM32 prefill supports only SM70 GPUs"); - - at::Tensor out_fp16 = out_.has_value() ? out_.value() : torch::empty_like(q); - at::Tensor softmax_lse = - softmax_lse_.has_value() - ? softmax_lse_.value() - : torch::empty({B, H, M}, - torch::dtype(torch::kFloat32).device(q.device())); - - launch_flash_attention_forward_paged_d256_bm32_phase_kernel( - q, k_cache, v_cache, out_fp16, softmax_lse, block_table, seq_lens, - softmax_scale, stream); - - return {out_fp16, softmax_lse}; + const at::Tensor& q, + const at::Tensor& k_cache, + const at::Tensor& v_cache, + std::optional& out_, + std::optional& softmax_lse_, + const at::Tensor& block_table, + const at::Tensor& seq_lens, + const float softmax_scale +) { + TORCH_CHECK(q.dim() == 4, "q must have shape [B, H, M, D]"); + TORCH_CHECK(k_cache.dim() == 4 && v_cache.dim() == 4, + "k_cache and v_cache must have shape [blocks, page, Hkv, D]"); + TORCH_CHECK(block_table.dim() == 2, + "block_table must have shape [B, max_num_blocks]"); + TORCH_CHECK(seq_lens.dim() == 1, "seq_lens must have shape [B]"); + TORCH_CHECK(q.dtype() == torch::kFloat16, "q must be fp16"); + TORCH_CHECK(k_cache.dtype() == torch::kFloat16, + "k_cache must be fp16"); + TORCH_CHECK(v_cache.dtype() == torch::kFloat16, + "v_cache must be fp16"); + 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.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda(), + "q, k_cache, and v_cache must be CUDA tensors"); + TORCH_CHECK(block_table.is_cuda() && seq_lens.is_cuda(), + "block_table and seq_lens must be CUDA tensors"); + TORCH_CHECK( + q.device() == k_cache.device() && q.device() == v_cache.device() + && q.device() == block_table.device() && q.device() == seq_lens.device(), + "all input tensors must be on the q device"); + TORCH_CHECK(q.is_contiguous(), "q must be contiguous [B, H, M, D]"); + TORCH_CHECK(block_table.is_contiguous(), "block_table must be contiguous"); + TORCH_CHECK(seq_lens.is_contiguous(), "seq_lens must be contiguous"); + TORCH_CHECK(k_cache.sizes() == v_cache.sizes(), + "k_cache and v_cache must have the same shape"); + TORCH_CHECK(k_cache.stride(-1) == 1 && v_cache.stride(-1) == 1, + "K/V head dimensions must be contiguous"); + TORCH_CHECK(k_cache.stride(1) % 8 == 0 && k_cache.stride(2) % 8 == 0 + && v_cache.stride(1) % 8 == 0 + && v_cache.stride(2) % 8 == 0, + "K/V strides must be divisible by 8 half elements"); + + 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 num_kv_heads = k_cache.size(2); + + TORCH_CHECK(B > 0 && H > 0, "q batch and head dimensions must be positive"); + TORCH_CHECK(M >= D256_BM32_PHASE_BLOCK_M, + "fixed BM32 prefill requires M >= ", + D256_BM32_PHASE_BLOCK_M); + TORCH_CHECK(D == D256_BM32_PHASE_D, + "fixed BM32 prefill requires D=", D256_BM32_PHASE_D); + TORCH_CHECK(k_cache.size(0) > 0 && num_kv_heads > 0, + "K/V cache block and head dimensions must be positive"); + TORCH_CHECK(k_cache.size(1) == D256_BM32_PHASE_PAGE_BLOCK_SIZE, + "fixed BM32 prefill requires page size ", + D256_BM32_PHASE_PAGE_BLOCK_SIZE); + TORCH_CHECK(k_cache.size(3) == D256_BM32_PHASE_D, + "K/V cache head dimension must be D=", + D256_BM32_PHASE_D); + TORCH_CHECK(H % num_kv_heads == 0, + "num_q_heads must be divisible by num_kv_heads"); + TORCH_CHECK(block_table.size(0) == B && seq_lens.size(0) == B, + "block_table and seq_lens batch dimensions must equal q batch"); + TORCH_CHECK(block_table.size(1) > 0, + "block_table must provide at least one page index"); + + if (out_.has_value()) { + const at::Tensor& out = out_.value(); + TORCH_CHECK(out.is_cuda() && out.device() == q.device(), + "out must be a CUDA tensor on the q device"); + TORCH_CHECK(out.dtype() == torch::kFloat16, "out must be fp16"); + TORCH_CHECK(out.sizes() == q.sizes(), "out must have shape [B, H, M, D]"); + TORCH_CHECK(out.is_contiguous(), "out must be contiguous [B, H, M, D]"); + } + if (softmax_lse_.has_value()) { + const at::Tensor& softmax_lse = softmax_lse_.value(); + TORCH_CHECK(softmax_lse.is_cuda() && softmax_lse.device() == q.device(), + "softmax_lse must be a CUDA tensor on the q device"); + TORCH_CHECK(softmax_lse.dtype() == torch::kFloat32, + "softmax_lse must be fp32"); + TORCH_CHECK(softmax_lse.dim() == 3 && softmax_lse.size(0) == B + && softmax_lse.size(1) == H && softmax_lse.size(2) == M, + "softmax_lse must have shape [B, H, M]"); + TORCH_CHECK(softmax_lse.is_contiguous(), + "softmax_lse must be contiguous [B, H, M]"); + } + + c10::cuda::CUDAGuard device_guard(q.device()); + auto stream = at::cuda::getCurrentCUDAStream().stream(); + cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; + const cudaError_t capture_error = cudaStreamIsCapturing(stream, &capture_status); + TORCH_CHECK(capture_error == cudaSuccess, "cudaStreamIsCapturing failed: ", + cudaGetErrorString(capture_error)); + TORCH_CHECK(capture_status == cudaStreamCaptureStatusNone, + "fixed BM32 prefill CUDA graph capture is not validated"); + + auto props = at::cuda::getCurrentDeviceProperties(); + TORCH_CHECK(props->major == 7 && props->minor == 0, + "fixed BM32 prefill supports only SM70 GPUs"); + + at::Tensor out_fp16 = out_.has_value() ? out_.value() : torch::empty_like(q); + at::Tensor softmax_lse = + softmax_lse_.has_value() + ? softmax_lse_.value() + : torch::empty({B, H, M}, + torch::dtype(torch::kFloat32).device(q.device())); + + launch_flash_attention_forward_paged_d256_bm32_phase_kernel( + q, k_cache, v_cache, out_fp16, softmax_lse, block_table, seq_lens, + softmax_scale, stream); + + return {out_fp16, softmax_lse}; } std::vector flash_attention_prefill_paged_d256_bm32_allp_pair_scratch_splitkv3( - const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, - std::optional& out_, std::optional& softmax_lse_, - at::Tensor& split_tmp_out, at::Tensor& split_tmp_row_max, - at::Tensor& split_tmp_row_sum, const at::Tensor& block_table, - const int64_t actual_n, const float softmax_scale) { - TORCH_CHECK(q.dim() == 4, "q must have shape [B, H, M, D]"); - TORCH_CHECK(k_cache.dim() == 4 && v_cache.dim() == 4, - "k_cache and v_cache must have shape [blocks, page, Hkv, D]"); - TORCH_CHECK(block_table.dim() == 2, - "block_table must have shape [B, max_num_blocks]"); - TORCH_CHECK(q.dtype() == torch::kFloat16, "q must be fp16"); - TORCH_CHECK(k_cache.dtype() == torch::kFloat16, "k_cache must be fp16"); - TORCH_CHECK(v_cache.dtype() == torch::kFloat16, "v_cache must be fp16"); - TORCH_CHECK(block_table.dtype() == torch::kInt32, - "block_table must be int32"); - TORCH_CHECK(q.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda(), - "q, k_cache, and v_cache must be CUDA tensors"); - TORCH_CHECK(block_table.is_cuda(), "block_table must be a CUDA tensor"); - TORCH_CHECK(q.device() == k_cache.device() && - q.device() == v_cache.device() && - q.device() == block_table.device(), - "all input tensors must be on the q device"); - TORCH_CHECK(q.is_contiguous(), "q must be contiguous [B, H, M, D]"); - TORCH_CHECK(block_table.is_contiguous(), "block_table must be contiguous"); - TORCH_CHECK(k_cache.sizes() == v_cache.sizes(), - "k_cache and v_cache must have the same shape"); - TORCH_CHECK(k_cache.stride(-1) == 1 && v_cache.stride(-1) == 1, - "K/V head dimensions must be contiguous"); - TORCH_CHECK(k_cache.stride(1) % 8 == 0 && k_cache.stride(2) % 8 == 0 && - v_cache.stride(1) % 8 == 0 && v_cache.stride(2) % 8 == 0, - "K/V strides must be divisible by 8 half elements"); - - 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 num_kv_heads = k_cache.size(2); - - TORCH_CHECK(B == 1, "fixed BM32 split-KV host-length path requires B=1"); - TORCH_CHECK(H > 0, "q head dimension must be positive"); - TORCH_CHECK( - M >= D256_BM32_PHASE_BLOCK_M, - "fixed BM32 split-KV prefill requires M >= ", D256_BM32_PHASE_BLOCK_M); - TORCH_CHECK(D == D256_BM32_PHASE_D, - "fixed BM32 split-KV prefill requires D=", D256_BM32_PHASE_D); - TORCH_CHECK(k_cache.size(0) > 0 && num_kv_heads > 0, - "K/V cache block and head dimensions must be positive"); - TORCH_CHECK(k_cache.size(1) == D256_BM32_PHASE_PAGE_BLOCK_SIZE, - "fixed BM32 split-KV prefill requires page size ", - D256_BM32_PHASE_PAGE_BLOCK_SIZE); - TORCH_CHECK(k_cache.size(3) == D256_BM32_PHASE_D, - "K/V cache head dimension must be D=", D256_BM32_PHASE_D); - TORCH_CHECK(H % num_kv_heads == 0, - "num_q_heads must be divisible by num_kv_heads"); - TORCH_CHECK(block_table.size(0) == B, - "block_table batch dimension must equal q batch"); - TORCH_CHECK(actual_n >= M && actual_n <= 2147483647LL, - "actual_n must satisfy M <= actual_n <= INT_MAX"); - const int64_t required_pages = - (actual_n + D256_BM32_PHASE_PAGE_BLOCK_SIZE - 1) / - D256_BM32_PHASE_PAGE_BLOCK_SIZE; - TORCH_CHECK(block_table.size(1) >= required_pages, - "block_table does not cover actual_n host length"); - - if (out_.has_value()) { - const at::Tensor& out = out_.value(); - TORCH_CHECK(out.is_cuda() && out.device() == q.device(), - "out must be a CUDA tensor on the q device"); - TORCH_CHECK(out.dtype() == torch::kFloat16, "out must be fp16"); - TORCH_CHECK(out.sizes() == q.sizes(), "out must have shape [B, H, M, D]"); - TORCH_CHECK(out.is_contiguous(), "out must be contiguous [B, H, M, D]"); - } - if (softmax_lse_.has_value()) { - const at::Tensor& softmax_lse = softmax_lse_.value(); - TORCH_CHECK(softmax_lse.is_cuda() && softmax_lse.device() == q.device(), - "softmax_lse must be a CUDA tensor on the q device"); - TORCH_CHECK(softmax_lse.dtype() == torch::kFloat32, - "softmax_lse must be fp32"); - TORCH_CHECK(softmax_lse.dim() == 3 && softmax_lse.size(0) == B && - softmax_lse.size(1) == H && softmax_lse.size(2) == M, - "softmax_lse must have shape [B, H, M]"); - TORCH_CHECK(softmax_lse.is_contiguous(), - "softmax_lse must be contiguous [B, H, M]"); - } - - const std::vector expected_partial_out = { - B, H, D256_BM32_PHASE_SPLIT_PARTS, M, D}; - const std::vector expected_partial_rows = { - B, H, D256_BM32_PHASE_SPLIT_PARTS, M}; - for (const at::Tensor* workspace : - {&split_tmp_out, &split_tmp_row_max, &split_tmp_row_sum}) { - TORCH_CHECK(workspace->is_cuda() && workspace->device() == q.device(), - "split-KV workspaces must be CUDA tensors on the q device"); - TORCH_CHECK(workspace->dtype() == torch::kFloat32, - "split-KV workspaces must be fp32"); - TORCH_CHECK(workspace->is_contiguous(), - "split-KV workspaces must be contiguous"); - } - TORCH_CHECK(split_tmp_out.sizes() == expected_partial_out, - "split_tmp_out must have shape [B, H, 3, M, D]"); - TORCH_CHECK(split_tmp_row_max.sizes() == expected_partial_rows, - "split_tmp_row_max must have shape [B, H, 3, M]"); - TORCH_CHECK(split_tmp_row_sum.sizes() == expected_partial_rows, - "split_tmp_row_sum must have shape [B, H, 3, M]"); - - c10::cuda::CUDAGuard device_guard(q.device()); - auto stream = at::cuda::getCurrentCUDAStream().stream(); - cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; - const cudaError_t capture_error = - cudaStreamIsCapturing(stream, &capture_status); - TORCH_CHECK(capture_error == cudaSuccess, "cudaStreamIsCapturing failed: ", - cudaGetErrorString(capture_error)); - TORCH_CHECK(capture_status == cudaStreamCaptureStatusNone, - "fixed BM32 split-KV CUDA graph capture is not validated"); - - auto props = at::cuda::getCurrentDeviceProperties(); - TORCH_CHECK(props->major == 7 && props->minor == 0, - "fixed BM32 split-KV prefill supports only SM70 GPUs"); - - at::Tensor out_fp16 = out_.has_value() ? out_.value() : torch::empty_like(q); - at::Tensor softmax_lse = - softmax_lse_.has_value() - ? softmax_lse_.value() - : torch::empty({B, H, M}, - torch::dtype(torch::kFloat32).device(q.device())); - launch_flash_attention_forward_paged_d256_bm32_splitkv3_kernel( - q, k_cache, v_cache, out_fp16, softmax_lse, split_tmp_out, - split_tmp_row_max, split_tmp_row_sum, block_table, - static_cast(actual_n), softmax_scale, stream); - - return {out_fp16, softmax_lse}; + const at::Tensor& q, + const at::Tensor& k_cache, + const at::Tensor& v_cache, + std::optional& out_, + std::optional& softmax_lse_, + at::Tensor& split_tmp_out, + at::Tensor& split_tmp_row_max, + at::Tensor& split_tmp_row_sum, + const at::Tensor& block_table, + const int64_t actual_n, + const float softmax_scale +) { + TORCH_CHECK(q.dim() == 4, "q must have shape [B, H, M, D]"); + TORCH_CHECK(k_cache.dim() == 4 && v_cache.dim() == 4, + "k_cache and v_cache must have shape [blocks, page, Hkv, D]"); + TORCH_CHECK(block_table.dim() == 2, + "block_table must have shape [B, max_num_blocks]"); + TORCH_CHECK(q.dtype() == torch::kFloat16, "q must be fp16"); + TORCH_CHECK(k_cache.dtype() == torch::kFloat16, + "k_cache must be fp16"); + TORCH_CHECK(v_cache.dtype() == torch::kFloat16, + "v_cache must be fp16"); + TORCH_CHECK(block_table.dtype() == torch::kInt32, + "block_table must be int32"); + TORCH_CHECK(q.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda(), + "q, k_cache, and v_cache must be CUDA tensors"); + TORCH_CHECK(block_table.is_cuda(), "block_table must be a CUDA tensor"); + TORCH_CHECK( + q.device() == k_cache.device() && q.device() == v_cache.device() + && q.device() == block_table.device(), + "all input tensors must be on the q device"); + TORCH_CHECK(q.is_contiguous(), "q must be contiguous [B, H, M, D]"); + TORCH_CHECK(block_table.is_contiguous(), "block_table must be contiguous"); + TORCH_CHECK(k_cache.sizes() == v_cache.sizes(), + "k_cache and v_cache must have the same shape"); + TORCH_CHECK(k_cache.stride(-1) == 1 && v_cache.stride(-1) == 1, + "K/V head dimensions must be contiguous"); + TORCH_CHECK(k_cache.stride(1) % 8 == 0 && k_cache.stride(2) % 8 == 0 + && v_cache.stride(1) % 8 == 0 + && v_cache.stride(2) % 8 == 0, + "K/V strides must be divisible by 8 half elements"); + + 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 num_kv_heads = k_cache.size(2); + + TORCH_CHECK(B == 1, + "fixed BM32 split-KV host-length path requires B=1"); + TORCH_CHECK(H > 0, "q head dimension must be positive"); + TORCH_CHECK(M >= D256_BM32_PHASE_BLOCK_M, + "fixed BM32 split-KV prefill requires M >= ", + D256_BM32_PHASE_BLOCK_M); + TORCH_CHECK(D == D256_BM32_PHASE_D, + "fixed BM32 split-KV prefill requires D=", + D256_BM32_PHASE_D); + TORCH_CHECK(k_cache.size(0) > 0 && num_kv_heads > 0, + "K/V cache block and head dimensions must be positive"); + TORCH_CHECK(k_cache.size(1) == D256_BM32_PHASE_PAGE_BLOCK_SIZE, + "fixed BM32 split-KV prefill requires page size ", + D256_BM32_PHASE_PAGE_BLOCK_SIZE); + TORCH_CHECK(k_cache.size(3) == D256_BM32_PHASE_D, + "K/V cache head dimension must be D=", + D256_BM32_PHASE_D); + TORCH_CHECK(H % num_kv_heads == 0, + "num_q_heads must be divisible by num_kv_heads"); + TORCH_CHECK(block_table.size(0) == B, + "block_table batch dimension must equal q batch"); + TORCH_CHECK(actual_n >= M && actual_n <= 2147483647LL, + "actual_n must satisfy M <= actual_n <= INT_MAX"); + const int64_t required_pages = + (actual_n + D256_BM32_PHASE_PAGE_BLOCK_SIZE - 1) + / D256_BM32_PHASE_PAGE_BLOCK_SIZE; + TORCH_CHECK(block_table.size(1) >= required_pages, + "block_table does not cover actual_n host length"); + + if (out_.has_value()) { + const at::Tensor& out = out_.value(); + TORCH_CHECK(out.is_cuda() && out.device() == q.device(), + "out must be a CUDA tensor on the q device"); + TORCH_CHECK(out.dtype() == torch::kFloat16, "out must be fp16"); + TORCH_CHECK(out.sizes() == q.sizes(), "out must have shape [B, H, M, D]"); + TORCH_CHECK(out.is_contiguous(), "out must be contiguous [B, H, M, D]"); + } + if (softmax_lse_.has_value()) { + const at::Tensor& softmax_lse = softmax_lse_.value(); + TORCH_CHECK(softmax_lse.is_cuda() && softmax_lse.device() == q.device(), + "softmax_lse must be a CUDA tensor on the q device"); + TORCH_CHECK(softmax_lse.dtype() == torch::kFloat32, + "softmax_lse must be fp32"); + TORCH_CHECK(softmax_lse.dim() == 3 && softmax_lse.size(0) == B + && softmax_lse.size(1) == H && softmax_lse.size(2) == M, + "softmax_lse must have shape [B, H, M]"); + TORCH_CHECK(softmax_lse.is_contiguous(), + "softmax_lse must be contiguous [B, H, M]"); + } + + const std::vector expected_partial_out = { + B, H, D256_BM32_PHASE_SPLIT_PARTS, M, D}; + const std::vector expected_partial_rows = { + B, H, D256_BM32_PHASE_SPLIT_PARTS, M}; + for (const at::Tensor* workspace : { + &split_tmp_out, &split_tmp_row_max, &split_tmp_row_sum}) { + TORCH_CHECK(workspace->is_cuda() && workspace->device() == q.device(), + "split-KV workspaces must be CUDA tensors on the q device"); + TORCH_CHECK(workspace->dtype() == torch::kFloat32, + "split-KV workspaces must be fp32"); + TORCH_CHECK(workspace->is_contiguous(), + "split-KV workspaces must be contiguous"); + } + TORCH_CHECK(split_tmp_out.sizes() == expected_partial_out, + "split_tmp_out must have shape [B, H, 3, M, D]"); + TORCH_CHECK(split_tmp_row_max.sizes() == expected_partial_rows, + "split_tmp_row_max must have shape [B, H, 3, M]"); + TORCH_CHECK(split_tmp_row_sum.sizes() == expected_partial_rows, + "split_tmp_row_sum must have shape [B, H, 3, M]"); + + c10::cuda::CUDAGuard device_guard(q.device()); + auto stream = at::cuda::getCurrentCUDAStream().stream(); + cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; + const cudaError_t capture_error = cudaStreamIsCapturing(stream, &capture_status); + TORCH_CHECK(capture_error == cudaSuccess, "cudaStreamIsCapturing failed: ", + cudaGetErrorString(capture_error)); + TORCH_CHECK(capture_status == cudaStreamCaptureStatusNone, + "fixed BM32 split-KV CUDA graph capture is not validated"); + + auto props = at::cuda::getCurrentDeviceProperties(); + TORCH_CHECK(props->major == 7 && props->minor == 0, + "fixed BM32 split-KV prefill supports only SM70 GPUs"); + + at::Tensor out_fp16 = out_.has_value() ? out_.value() : torch::empty_like(q); + at::Tensor softmax_lse = + softmax_lse_.has_value() + ? softmax_lse_.value() + : torch::empty({B, H, M}, + torch::dtype(torch::kFloat32).device(q.device())); + launch_flash_attention_forward_paged_d256_bm32_splitkv3_kernel( + q, k_cache, v_cache, out_fp16, softmax_lse, + split_tmp_out, split_tmp_row_max, split_tmp_row_sum, + block_table, static_cast(actual_n), + softmax_scale, stream); + + return {out_fp16, softmax_lse}; } at::Tensor flash_attention_prefill_paged_bfla( - const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, - std::optional& out_, const at::Tensor& block_table, - const at::Tensor& seq_lens, const at::Tensor& bfla_block_mask, - const int bfla_mask_block_n, const float softmax_scale, - const std::string& kv_cache_dtype, const float k_scale, const float v_scale, - const bool is_causal, const int window_size_left, - const int window_size_right) { - TORCH_CHECK(q.dtype() == torch::kFloat16, "q must be fp16"); - const int kv_dtype_code = kv_cache_dtype_code_from_string(kv_cache_dtype); - TORCH_CHECK(kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP16, - "BFLA paged prefill supports fp16 KV cache only"); - TORCH_CHECK(k_cache.dtype() == torch::kFloat16, "k_cache must be fp16"); - TORCH_CHECK(v_cache.dtype() == torch::kFloat16, "v_cache must be fp16"); - TORCH_CHECK(q.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda(), - "Tensors must be on CUDA"); - TORCH_CHECK(block_table.is_cuda() && seq_lens.is_cuda(), - "block_table and seq_lens must be CUDA tensors"); - TORCH_CHECK(bfla_block_mask.is_cuda(), - "bfla_block_mask must be a CUDA tensor"); - TORCH_CHECK(bfla_block_mask.dtype() == torch::kInt32, - "bfla_block_mask must be int32"); - TORCH_CHECK(bfla_block_mask.dim() == 4, - "bfla_block_mask must have shape [B, Hkv, q_tiles, kv_tiles]"); - TORCH_CHECK( - q.stride(-1) == 1 && k_cache.stride(-1) == 1 && v_cache.stride(-1) == 1, - "Last dim must be contiguous"); - TORCH_CHECK(k_cache.stride(1) % 8 == 0 && k_cache.stride(2) % 8 == 0 && - v_cache.stride(1) % 8 == 0 && v_cache.stride(2) % 8 == 0, - "Paged KV strides must be divisible by 8 half elements"); - - 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 num_kv_heads = k_cache.size(2); - const int page_block_size = k_cache.size(1); - - TORCH_CHECK(D == 256, "BFLA paged prefill supports D=256 only"); - TORCH_CHECK(H % num_kv_heads == 0, - "num_heads must be divisible by num_kv_heads"); - TORCH_CHECK(M > 1, "BFLA paged prefill requires prefill M > 1"); - TORCH_CHECK(page_block_size >= 16 && (page_block_size % 16) == 0, - "BFLA paged prefill requires page block size multiple of 16"); - TORCH_CHECK(is_causal, "BFLA paged prefill supports causal attention only"); - TORCH_CHECK(window_size_left == -1 && window_size_right == -1, - "BFLA paged prefill does not support sliding window"); - TORCH_CHECK(bfla_mask_block_n > 0, "bfla_mask_block_n must be positive"); - TORCH_CHECK(B <= bfla_block_mask.size(0), - "bfla_block_mask batch dimension must cover q"); - TORCH_CHECK(num_kv_heads <= bfla_block_mask.size(1), - "bfla_block_mask head dimension must cover KV heads"); - - at::Tensor out_fp16 = out_.has_value() ? out_.value() : torch::zeros_like(q); - auto softmax_lse = - 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."); - - launcher_flash_attention_forward_paged<256, flash_v100::KV_CACHE_DTYPE_FP16>( - q, k_cache, v_cache, out_fp16, softmax_lse, block_table, seq_lens, - softmax_scale, is_causal, k_scale, v_scale, window_size_left, - window_size_right, stream, bfla_block_mask.data_ptr(), - bfla_mask_block_n, bfla_block_mask.stride(0), bfla_block_mask.stride(1), - bfla_block_mask.stride(2), bfla_block_mask.stride(3)); - - return out_fp16; + const at::Tensor& q, + const at::Tensor& k_cache, + const at::Tensor& v_cache, + std::optional& out_, + const at::Tensor& block_table, + const at::Tensor& seq_lens, + const at::Tensor& bfla_block_mask, + const int bfla_mask_block_n, + const float softmax_scale, + const std::string& kv_cache_dtype, + const float k_scale, + const float v_scale, + const bool is_causal, + const int window_size_left, + const int window_size_right +) { + TORCH_CHECK(q.dtype() == torch::kFloat16, "q must be fp16"); + const int kv_dtype_code = kv_cache_dtype_code_from_string(kv_cache_dtype); + TORCH_CHECK(kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP16, + "BFLA paged prefill supports fp16 KV cache only"); + TORCH_CHECK(k_cache.dtype() == torch::kFloat16, "k_cache must be fp16"); + TORCH_CHECK(v_cache.dtype() == torch::kFloat16, "v_cache must be fp16"); + TORCH_CHECK(q.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda(), + "Tensors must be on CUDA"); + TORCH_CHECK(block_table.is_cuda() && seq_lens.is_cuda(), + "block_table and seq_lens must be CUDA tensors"); + TORCH_CHECK(bfla_block_mask.is_cuda(), + "bfla_block_mask must be a CUDA tensor"); + TORCH_CHECK(bfla_block_mask.dtype() == torch::kInt32, + "bfla_block_mask must be int32"); + TORCH_CHECK(bfla_block_mask.dim() == 4, + "bfla_block_mask must have shape [B, Hkv, q_tiles, kv_tiles]"); + TORCH_CHECK(q.stride(-1) == 1 && k_cache.stride(-1) == 1 && + v_cache.stride(-1) == 1, + "Last dim must be contiguous"); + TORCH_CHECK(k_cache.stride(1) % 8 == 0 && k_cache.stride(2) % 8 == 0 && + v_cache.stride(1) % 8 == 0 && v_cache.stride(2) % 8 == 0, + "Paged KV strides must be divisible by 8 half elements"); + + 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 num_kv_heads = k_cache.size(2); + const int page_block_size = k_cache.size(1); + + TORCH_CHECK(D == 256, "BFLA paged prefill supports D=256 only"); + TORCH_CHECK(H % num_kv_heads == 0, + "num_heads must be divisible by num_kv_heads"); + TORCH_CHECK(M > 1, "BFLA paged prefill requires prefill M > 1"); + TORCH_CHECK(page_block_size >= 16 && (page_block_size % 16) == 0, + "BFLA paged prefill requires page block size multiple of 16"); + TORCH_CHECK(is_causal, "BFLA paged prefill supports causal attention only"); + TORCH_CHECK(window_size_left == -1 && window_size_right == -1, + "BFLA paged prefill does not support sliding window"); + TORCH_CHECK(bfla_mask_block_n > 0, + "bfla_mask_block_n must be positive"); + TORCH_CHECK(B <= bfla_block_mask.size(0), + "bfla_block_mask batch dimension must cover q"); + TORCH_CHECK(num_kv_heads <= bfla_block_mask.size(1), + "bfla_block_mask head dimension must cover KV heads"); + + at::Tensor out_fp16 = out_.has_value() ? out_.value() : torch::zeros_like(q); + auto softmax_lse = 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."); + + launcher_flash_attention_forward_paged<256, flash_v100::KV_CACHE_DTYPE_FP16>( + q, k_cache, v_cache, out_fp16, softmax_lse, block_table, seq_lens, + softmax_scale, is_causal, k_scale, v_scale, window_size_left, + window_size_right, stream, bfla_block_mask.data_ptr(), + bfla_mask_block_n, bfla_block_mask.stride(0), + bfla_block_mask.stride(1), bfla_block_mask.stride(2), + bfla_block_mask.stride(3)); + + return out_fp16; } at::Tensor flash_attention_decode_paged_wmma( - const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, - std::optional& out_, const at::Tensor& block_table, - const at::Tensor& seq_lens, const float softmax_scale, - const std::string& kv_cache_dtype, const float k_scale, - const float v_scale) { - TORCH_CHECK(q.dtype() == torch::kFloat16, "q must be fp16"); - const int kv_dtype_code = kv_cache_dtype_code_from_string(kv_cache_dtype); - TORCH_CHECK(kv_dtype_code >= 0, - "Unsupported kv_cache_dtype: ", kv_cache_dtype); - if (kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP16) { - TORCH_CHECK(k_cache.dtype() == torch::kFloat16, "k_cache must be fp16"); - TORCH_CHECK(v_cache.dtype() == torch::kFloat16, "v_cache must be fp16"); - } else { - TORCH_CHECK(k_cache.dtype() == torch::kUInt8, - "fp8 k_cache must be stored as uint8"); - TORCH_CHECK(v_cache.dtype() == torch::kUInt8, - "fp8 v_cache must be stored as uint8"); - TORCH_CHECK(k_scale > 0.f && v_scale > 0.f, - "fp8 k/v scales must be positive"); - } - TORCH_CHECK(q.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda(), - "Tensors must be on CUDA"); - TORCH_CHECK(block_table.is_cuda() && seq_lens.is_cuda(), - "block_table and seq_lens must be CUDA tensors"); - TORCH_CHECK(q.dim() == 3, "q must have shape [B, H, D]"); - TORCH_CHECK(block_table.dim() == 2, - "block_table must have shape [B, max_num_blocks]"); - TORCH_CHECK(seq_lens.dim() == 1, "seq_lens must have shape [B]"); - TORCH_CHECK( - q.stride(-1) == 1 && k_cache.stride(-1) == 1 && v_cache.stride(-1) == 1, - "Last dim must be contiguous"); - TORCH_CHECK(k_cache.stride(1) % 8 == 0 && k_cache.stride(2) % 8 == 0 && - v_cache.stride(1) % 8 == 0 && v_cache.stride(2) % 8 == 0, - "Paged KV strides must be divisible by 8 half elements"); - - const int B = q.size(0); - const int H = q.size(1); - const int D = q.size(2); - const int num_kv_heads = k_cache.size(2); - - TORCH_CHECK(B <= block_table.size(0), "block_table batch size must cover q"); - TORCH_CHECK(B <= seq_lens.size(0), "seq_lens batch size must cover q"); - TORCH_CHECK(D <= 256 && D % 8 == 0 && D % 2 == 0, - "D must be even, <=256, multiple of 8"); - TORCH_CHECK(H % num_kv_heads == 0, - "num_heads must be divisible by num_kv_heads"); - - at::Tensor out_fp16 = out_.has_value() ? out_.value() : torch::zeros_like(q); - TORCH_CHECK(out_fp16.is_cuda(), "out must be on CUDA"); - TORCH_CHECK(out_fp16.dtype() == torch::kFloat16, "out must be fp16"); - TORCH_CHECK(out_fp16.sizes() == q.sizes(), "out must have same shape as q"); - TORCH_CHECK(out_fp16.stride(-1) == 1, "out last dim must be contiguous"); - - at::Tensor q_m1 = q.unsqueeze(2); - at::Tensor out_m1 = out_fp16.unsqueeze(2); - auto softmax_lse = - torch::zeros({B, H, 1}, 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."); - -#define LAUNCH_DECODE_WMMA_TYPED(HDIM, KV_DTYPE_CODE) \ - launcher_flash_attention_forward_paged( \ - q_m1, k_cache, v_cache, out_m1, softmax_lse, block_table, seq_lens, \ - softmax_scale, true, k_scale, v_scale, -1, -1, stream) - -#define LAUNCH_DECODE_WMMA_BY_KV(HDIM) \ - do { \ - switch (kv_dtype_code) { \ - case flash_v100::KV_CACHE_DTYPE_FP16: \ - LAUNCH_DECODE_WMMA_TYPED(HDIM, flash_v100::KV_CACHE_DTYPE_FP16); \ - break; \ - case flash_v100::KV_CACHE_DTYPE_FP8_E4M3: \ - LAUNCH_DECODE_WMMA_TYPED(HDIM, flash_v100::KV_CACHE_DTYPE_FP8_E4M3); \ - break; \ - case flash_v100::KV_CACHE_DTYPE_FP8_E5M2: \ - LAUNCH_DECODE_WMMA_TYPED(HDIM, flash_v100::KV_CACHE_DTYPE_FP8_E5M2); \ - break; \ - default: \ - TORCH_CHECK(false, "Unsupported kv_cache_dtype: ", kv_cache_dtype); \ - } \ - } while (0) - - switch (D) { - case 16: - LAUNCH_DECODE_WMMA_BY_KV(16); - break; - case 32: - LAUNCH_DECODE_WMMA_BY_KV(32); - break; - case 64: - LAUNCH_DECODE_WMMA_BY_KV(64); - break; - case 128: - LAUNCH_DECODE_WMMA_BY_KV(128); - break; - case 256: - LAUNCH_DECODE_WMMA_BY_KV(256); - break; - default: - TORCH_CHECK(false, "Unsupported D: ", D); - } - -#undef LAUNCH_DECODE_WMMA_BY_KV -#undef LAUNCH_DECODE_WMMA_TYPED - - C10_CUDA_KERNEL_LAUNCH_CHECK(); - return out_fp16; + const at::Tensor& q, + const at::Tensor& k_cache, + const at::Tensor& v_cache, + std::optional& out_, + const at::Tensor& block_table, + const at::Tensor& seq_lens, + const float softmax_scale, + const std::string& kv_cache_dtype, + const float k_scale, + const float v_scale +) { + TORCH_CHECK(q.dtype() == torch::kFloat16, "q must be fp16"); + const int kv_dtype_code = kv_cache_dtype_code_from_string(kv_cache_dtype); + TORCH_CHECK(kv_dtype_code >= 0, "Unsupported kv_cache_dtype: ", kv_cache_dtype); + if (kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP16) { + TORCH_CHECK(k_cache.dtype() == torch::kFloat16, "k_cache must be fp16"); + TORCH_CHECK(v_cache.dtype() == torch::kFloat16, "v_cache must be fp16"); + } else { + TORCH_CHECK(k_cache.dtype() == torch::kUInt8, + "fp8 k_cache must be stored as uint8"); + TORCH_CHECK(v_cache.dtype() == torch::kUInt8, + "fp8 v_cache must be stored as uint8"); + TORCH_CHECK(k_scale > 0.f && v_scale > 0.f, + "fp8 k/v scales must be positive"); + } + TORCH_CHECK(q.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda(), + "Tensors must be on CUDA"); + TORCH_CHECK(block_table.is_cuda() && seq_lens.is_cuda(), + "block_table and seq_lens must be CUDA tensors"); + TORCH_CHECK(q.dim() == 3, "q must have shape [B, H, D]"); + TORCH_CHECK(block_table.dim() == 2, + "block_table must have shape [B, max_num_blocks]"); + TORCH_CHECK(seq_lens.dim() == 1, "seq_lens must have shape [B]"); + TORCH_CHECK(q.stride(-1) == 1 && k_cache.stride(-1) == 1 && + v_cache.stride(-1) == 1, + "Last dim must be contiguous"); + TORCH_CHECK(k_cache.stride(1) % 8 == 0 && k_cache.stride(2) % 8 == 0 && + v_cache.stride(1) % 8 == 0 && v_cache.stride(2) % 8 == 0, + "Paged KV strides must be divisible by 8 half elements"); + + const int B = q.size(0); + const int H = q.size(1); + const int D = q.size(2); + const int num_kv_heads = k_cache.size(2); + + TORCH_CHECK(B <= block_table.size(0), "block_table batch size must cover q"); + TORCH_CHECK(B <= seq_lens.size(0), "seq_lens batch size must cover q"); + TORCH_CHECK(D <= 256 && D % 8 == 0 && D % 2 == 0, + "D must be even, <=256, multiple of 8"); + TORCH_CHECK(H % num_kv_heads == 0, + "num_heads must be divisible by num_kv_heads"); + + at::Tensor out_fp16 = out_.has_value() ? out_.value() : torch::zeros_like(q); + TORCH_CHECK(out_fp16.is_cuda(), "out must be on CUDA"); + TORCH_CHECK(out_fp16.dtype() == torch::kFloat16, "out must be fp16"); + TORCH_CHECK(out_fp16.sizes() == q.sizes(), "out must have same shape as q"); + TORCH_CHECK(out_fp16.stride(-1) == 1, "out last dim must be contiguous"); + + at::Tensor q_m1 = q.unsqueeze(2); + at::Tensor out_m1 = out_fp16.unsqueeze(2); + auto softmax_lse = torch::zeros({B, H, 1}, + 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."); + + #define LAUNCH_DECODE_WMMA_TYPED(HDIM, KV_DTYPE_CODE) \ + launcher_flash_attention_forward_paged( \ + q_m1, k_cache, v_cache, out_m1, softmax_lse, block_table, seq_lens, \ + softmax_scale, true, k_scale, v_scale, -1, -1, stream) + + #define LAUNCH_DECODE_WMMA_BY_KV(HDIM) \ + do { \ + switch (kv_dtype_code) { \ + case flash_v100::KV_CACHE_DTYPE_FP16: \ + LAUNCH_DECODE_WMMA_TYPED(HDIM, flash_v100::KV_CACHE_DTYPE_FP16); \ + break; \ + case flash_v100::KV_CACHE_DTYPE_FP8_E4M3: \ + LAUNCH_DECODE_WMMA_TYPED(HDIM, flash_v100::KV_CACHE_DTYPE_FP8_E4M3); \ + break; \ + case flash_v100::KV_CACHE_DTYPE_FP8_E5M2: \ + LAUNCH_DECODE_WMMA_TYPED(HDIM, flash_v100::KV_CACHE_DTYPE_FP8_E5M2); \ + break; \ + default: \ + TORCH_CHECK(false, "Unsupported kv_cache_dtype: ", kv_cache_dtype); \ + } \ + } while (0) + + switch (D) { + case 16: + LAUNCH_DECODE_WMMA_BY_KV(16); + break; + case 32: + LAUNCH_DECODE_WMMA_BY_KV(32); + break; + case 64: + LAUNCH_DECODE_WMMA_BY_KV(64); + break; + case 128: + LAUNCH_DECODE_WMMA_BY_KV(128); + break; + case 256: + LAUNCH_DECODE_WMMA_BY_KV(256); + break; + default: + TORCH_CHECK(false, "Unsupported D: ", D); + } + + #undef LAUNCH_DECODE_WMMA_BY_KV + #undef LAUNCH_DECODE_WMMA_TYPED + + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out_fp16; } 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 e9453037b5..bbcb8f8397 100644 --- a/flash-attention-v100/kernel/fused_mha_forward_paged_full.cu +++ b/flash-attention-v100/kernel/fused_mha_forward_paged_full.cu @@ -16,621 +16,589 @@ 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); - - 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); +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) { - if (start_col >= start_row + valid_q_rows) { - continue; - } + 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 uint4* k_vec = reinterpret_cast(k_ptr + start_col * D); - uint4* sK_vec = reinterpret_cast(sK); + 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; + } -#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 uint4* q_vec = reinterpret_cast(q_ptr); + uint4* sQ_vec = reinterpret_cast(sQ); - 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; + #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 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; + 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; } } - } - store_matrix_sync(sS + tile_m * S_STRIDE + tile_n, acc_frag, S_STRIDE, - mem_row_major); - } - __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))); + 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; } - } - -#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(); + + 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) { - 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(); + + 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; + } + } } - } + __syncthreads(); - if (tail_col >= 0) { - sP_row_h[tail_col] = tail_value; - } + const uint4* v_vec = reinterpret_cast(v_ptr + start_col * D); + uint4* sV_vec = reinterpret_cast(sV); -#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(); + #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 uint4* v_vec = reinterpret_cast(v_ptr + start_col * D); - uint4* sV_vec = reinterpret_cast(sV); + 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(); -#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 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; - 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 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 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 int tile_m_idx = global_tile_idx / num_tiles_n_pv; + const int tile_d_idx = global_tile_idx % num_tiles_n_pv; - 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 = tile_m_idx * WMMA_M; + const int tile_d = tile_d_idx * WMMA_N; - const int tile_m_idx = global_tile_idx / num_tiles_n_pv; - const int tile_d_idx = global_tile_idx % num_tiles_n_pv; + if (tile_m >= valid_q_rows) continue; - const int tile_m = tile_m_idx * WMMA_M; - const int tile_d = tile_d_idx * WMMA_N; + fragment a_frag; + fragment b_frag; + fragment acc_frag; - if (tile_m >= valid_q_rows) continue; + load_matrix_sync(acc_frag, sO + tile_m * O_STRIDE + tile_d, O_STRIDE, mem_row_major); - fragment a_frag; - fragment b_frag; - fragment acc_frag; + #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; - load_matrix_sync(acc_frag, sO + tile_m * O_STRIDE + tile_d, O_STRIDE, - mem_row_major); + 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(); + } -#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; + 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(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); + if (tid < valid_q_rows) { + const float sum = fmaxf(sRowSum[tid], 1e-24f); + softmax_lse_ptr[tid] = sRowMax[tid] + logf(sum); } - __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 bb548ecc41..18c8274d2e 100644 --- a/flash-attention-v100/kernel/paged_kv_utils.cuh +++ b/flash-attention-v100/kernel/paged_kv_utils.cuh @@ -4,58 +4,78 @@ #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) { - const int virtual_page_idx = token_idx / page_block_size; - const int page_offset = token_idx % page_block_size; +__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 physical_block_idx = block_table[virtual_page_idx]; + const int virtual_page_idx = token_idx / page_block_size; + const int page_offset = token_idx % page_block_size; - 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 int physical_block_idx = block_table[virtual_page_idx]; - return physical_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; } -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; +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; + const int kv_head_offset = kv_head_idx * head_dim; - uint4* smem_vec = reinterpret_cast(smem_dst); + uint4* smem_vec = reinterpret_cast(smem_dst); -#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; + #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); + 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; + if (token_offset < num_tokens && vec_col < d_stride_uint4) { + const int global_token_idx = start_token_idx + token_offset; - 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); + 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 + ); - const uint4* kv_vec = - reinterpret_cast(kv_cache + base_offset); - kv_val = __ldg(&kv_vec[vec_col]); - } + 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; - } + 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 77a94b13e7..b5f321a28a 100644 --- a/flash-attention-v100/kernel/paged_to_contiguous.cu +++ b/flash-attention-v100/kernel/paged_to_contiguous.cu @@ -4,178 +4,215 @@ #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 aeda1b57cb..c4901cb6d7 100644 --- a/flash-attention-v100/kernel/paged_to_contiguous_fixed.cu +++ b/flash-attention-v100/kernel/paged_to_contiguous_fixed.cu @@ -4,85 +4,102 @@ #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 9845ab19fb..0017127d71 100644 --- a/flash-attention-v100/kernel/paged_to_contiguous_old.cu +++ b/flash-attention-v100/kernel/paged_to_contiguous_old.cu @@ -4,75 +4,85 @@ #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/requirements/common.txt b/requirements/common.txt index 27ee2a5918..0384ca2f96 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -3,31 +3,31 @@ cachetools psutil sentencepiece # Required for LLaMA tokenizer. numpy -requests >= 2.34.2 +requests >= 2.26.0 tqdm blake3 py-cpuinfo -transformers >= 5.15.1 +transformers >= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0 tokenizers >= 0.21.1 # Required for fast incremental detokenization. -safetensors >= 0.8.0 # MXFP4/MXFP6 dtype support (F8_E8M0, F4) added in 0.6.0: https://github.com/huggingface/safetensors/pull/611 +safetensors >= 0.6.2 # 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.141.1, < 0.142.0 # Required by FastAPI's form models in the OpenAI API server's audio transcriptions endpoint. -aiohttp >= 3.14.3 +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 openai >= 2.0.0 # For Responses API with reasoning content -pydantic >= 2.13.4, < 2.14.0 -prometheus_client >= 0.26.0 +pydantic >= 2.12.0, < 2.13.0 +prometheus_client >= 0.18.0 pillow # Required for image processing prometheus-fastapi-instrumentator >= 7.0.0, < 8.0.0 -tiktoken >= 0.14.0 # Required for DBRX tokenizer +tiktoken >= 0.6.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.3.1 +lark == 1.2.2 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.16.0 -filelock >= 3.32.3 # need to contain https://github.com/tox-dev/filelock/pull/317 +typing_extensions >= 4.10 +filelock >= 3.16.1 # 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.44.0 -opentelemetry-api >= 1.44.0 +opentelemetry-sdk >= 1.27.0 +opentelemetry-api >= 1.27.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 e22cde62d2..5ec338af73 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.67.0; platform_machine != "s390x" # Required for N-gram speculative decoding +numba == 0.65.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 52161d05f1..8d92a144c1 100644 --- a/requirements/cuda.txt +++ b/requirements/cuda.txt @@ -1,7 +1,7 @@ # Common dependencies -r common.txt -numba == 0.67.0 # Required for N-gram speculative decoding +numba == 0.65.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.3.3 +fastsafetensors >= 0.2.2 # QuACK and Cutlass DSL for FA4 (cute-DSL implementation) -nvidia-cutlass-dsl[cu13]==4.7.0 +nvidia-cutlass-dsl[cu13]==4.5.2 quack-kernels>=0.3.3 # Tokenspeed_MLA for faster mla with spec decode -tokenspeed-mla==0.2.5 +tokenspeed-mla==0.1.2 # Humming kernels for quantization gemm humming-kernels[cu13]==0.1.2 diff --git a/requirements/docs.txt b/requirements/docs.txt index ffe58c6366..5331a3b79b 100644 --- a/requirements/docs.txt +++ b/requirements/docs.txt @@ -6,7 +6,7 @@ babel==2.18.0 # via # mkdocs-git-revision-date-localized-plugin # mkdocs-material -backrefs==8.0 +backrefs==7.0 # via mkdocs-material bracex==2.6 # via wcmatch @@ -30,7 +30,7 @@ gitdb==4.0.12 # via gitpython gitpython==3.1.50 # via mkdocs-git-revision-date-localized-plugin -griffelib==2.2.0 +griffelib==2.0.2 # 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.3 +packaging==26.2 # 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.16.4 +ruff==0.15.12 # via -r requirements/docs.in selectolax==0.4.8 # via mkdocs-glightbox diff --git a/requirements/rocm.txt b/requirements/rocm.txt index b99e5820ea..0520f4ca1e 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.83.0 -grpcio-reflection==1.83.0 +grpcio==1.78.0 +grpcio-reflection==1.78.0 -numba == 0.67.0 # Required for N-gram speculative decoding +numba == 0.65.0 # Required for N-gram speculative decoding # Dependencies for AMD GPUs datasets peft pytest-asyncio -tensorizer==2.12.1 +tensorizer==2.10.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.16.1 -conch-triton-kernels==1.3 +runai-model-streamer[s3,gcs,azure]==0.15.7 +conch-triton-kernels==1.2.1 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 f9e2114a48..e0981c6fdc 100644 --- a/requirements/test/cuda.in +++ b/requirements/test/cuda.in @@ -1,6 +1,6 @@ # testing pytest -tensorizer==2.12.1 +tensorizer==2.10.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.20.0 # required for phi-4-mm test +peft>=0.18.1 # 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.15.1 +transformers==5.5.3 tokenizers==0.22.2 schemathesis>=3.39.15 # Required for openai schema test. # quantization -bitsandbytes==0.50.1 +bitsandbytes==0.49.2 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.83.0 -grpcio-reflection==1.83.0 +grpcio==1.78.0 +grpcio-reflection==1.78.0 -arctic-inference == 0.2.0; platform_machine == "x86_64" # Required for suffix decoding test -numba == 0.67.0 # Required for N-gram speculative decoding +arctic-inference == 0.1.1; platform_machine == "x86_64" # Required for suffix decoding test +numba == 0.65.0 # Required for N-gram speculative decoding numpy -runai-model-streamer[s3,gcs,azure]==0.16.1 -fastsafetensors>=0.3.3; platform_machine == "x86_64" +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 instanttensor>=0.1.5; platform_machine == "x86_64" -pydantic>=2.13.4 # 2.11 leads to error on python 3.13 +pydantic>=2.12 # 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 e6a621d712..831de114dd 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.14.3 +aiohttp==3.13.3 # via # -c requirements/common.txt # aiohttp-cors @@ -31,7 +31,7 @@ anyio==4.6.2.post1 # via # httpx # starlette -arctic-inference==0.2.0 +arctic-inference==0.1.1 # 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.50.1 +bitsandbytes==0.49.2 # via -r requirements/test/cuda.in -black==26.5.1 +black==24.10.0 # 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==6.0.0.post1 +chardet==5.2.0 # 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.141.1 +fastapi==0.128.0 # 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.3.3 +fastsafetensors==0.2.2 # via # -c requirements/cuda.txt # -r requirements/test/cuda.in -filelock==3.32.4 +filelock==3.16.1 # 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.34.0 +google-api-core==2.24.2 # 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.75.2 +googleapis-common-protos==1.70.0 # 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.83.0 +grpcio==1.78.0 # via # -r requirements/test/cuda.in # grpcio-reflection # ray -grpcio-reflection==1.83.0 +grpcio-reflection==1.78.0 # via -r requirements/test/cuda.in h11==0.14.0 # via @@ -323,6 +323,8 @@ 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 @@ -378,7 +380,7 @@ libnacl==2.1.0 # via tensorizer librosa==0.10.2.post1 # via -r requirements/test/cuda.in -llvmlite==0.49.0 +llvmlite==0.47.0 # via numba lm-eval==0.4.12 # via -r requirements/test/cuda.in @@ -398,7 +400,7 @@ markupsafe==3.0.1 # werkzeug matplotlib==3.9.2 # via -r requirements/test/cuda.in -mbstrdecoder==1.1.5 +mbstrdecoder==1.1.3 # via # dataproperty # pytablewriter @@ -443,7 +445,7 @@ nltk==3.9.1 # via rouge-score num2words==0.5.14 # via -r requirements/test/cuda.in -numba==0.67.0 +numba==0.65.0 # via # -c requirements/cuda.txt # -r requirements/test/cuda.in @@ -486,7 +488,6 @@ numpy==2.2.6 # scikit-learn # scipy # segmentation-models-pytorch - # sentence-transformers # soxr # statsmodels # tensorizer @@ -552,28 +553,28 @@ opencv-python-headless==4.13.0.90 # mistral-common openpyxl==3.1.5 # via -r requirements/test/cuda.in -opentelemetry-api==1.44.0 +opentelemetry-api==1.35.0 # via # -c requirements/common.txt # opentelemetry-exporter-prometheus # opentelemetry-sdk # opentelemetry-semantic-conventions -opentelemetry-exporter-prometheus==0.65b0 +opentelemetry-exporter-prometheus==0.56b0 # via ray -opentelemetry-proto==1.44.0 +opentelemetry-proto==1.35.0 # via ray -opentelemetry-sdk==1.44.0 +opentelemetry-sdk==1.35.0 # via # -c requirements/common.txt # opentelemetry-exporter-prometheus # ray -opentelemetry-semantic-conventions==0.65b0 +opentelemetry-semantic-conventions==0.56b0 # via opentelemetry-sdk optuna==3.6.1 # via genai-perf orjson==3.11.5 # via genai-perf -packaging==26.3 +packaging==24.2 # via # accelerate # bitsandbytes @@ -603,19 +604,19 @@ pandas==2.2.3 # fastparquet # genai-perf # statsmodels -pathspec==1.1.1 +pathspec==0.12.1 # via black pathvalidate==3.2.1 # via pytablewriter patsy==1.0.1 # via statsmodels -peft==0.20.0 +peft==0.18.1 # 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==12.3.0 +pillow==10.4.0 # via # genai-perf # imagehash @@ -647,7 +648,7 @@ portalocker==2.10.1 # via sacrebleu pqdm==0.2.0 # via -r requirements/test/cuda.in -prometheus-client==0.26.0 +prometheus-client==0.22.0 # via # -c requirements/common.txt # opentelemetry-exporter-prometheus @@ -656,9 +657,9 @@ propcache==0.2.0 # via # aiohttp # yarl -proto-plus==1.28.4 +proto-plus==1.26.1 # via google-api-core -protobuf==7.36.0 +protobuf==6.33.6 # via # -c requirements/common.txt # google-api-core @@ -693,7 +694,7 @@ pycparser==2.22 # via cffi pycryptodomex==3.22.0 # via blobfile -pydantic==2.13.4 +pydantic==2.12.0 # via # -c requirements/common.txt # -r requirements/test/cuda.in @@ -706,7 +707,7 @@ pydantic==2.13.4 # openai-harmony # pydantic-extra-types # ray -pydantic-core==2.46.4 +pydantic-core==2.41.1 # via pydantic pydantic-extra-types==2.10.5 # via mistral-common @@ -761,8 +762,6 @@ 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 @@ -804,7 +803,7 @@ regex==2026.2.28 # sacrebleu # tiktoken # transformers -requests==2.34.2 +requests==2.32.3 # via # -c requirements/common.txt # azure-core @@ -845,19 +844,19 @@ rpds-py==0.20.1 # referencing rsa==4.9.1 # via google-auth -runai-model-streamer==0.16.1 +runai-model-streamer==0.15.7 # via -r requirements/test/cuda.in -runai-model-streamer-azure==0.16.1 +runai-model-streamer-azure==0.15.7 # via runai-model-streamer -runai-model-streamer-gcs==0.16.1 +runai-model-streamer-gcs==0.15.7 # via runai-model-streamer -runai-model-streamer-s3==0.16.1 +runai-model-streamer-s3==0.15.7 # via runai-model-streamer s3transfer==0.10.3 # via boto3 sacrebleu==2.4.3 # via lm-eval -safetensors==0.8.0 +safetensors==0.7.0 # via # -c requirements/common.txt # accelerate @@ -891,7 +890,7 @@ scipy==1.13.1 # vocos segmentation-models-pytorch==0.5.0 # via -r requirements/test/cuda.in -sentence-transformers==6.0.0 +sentence-transformers==5.2.0 # via # -r requirements/test/cuda.in # mteb @@ -964,7 +963,7 @@ tenacity==9.1.2 # gpt-oss # lm-eval # plotly -tensorizer==2.12.1 +tensorizer==2.10.1 # via -r requirements/test/cuda.in termcolor==3.1.0 # via gpt-oss @@ -972,7 +971,7 @@ threadpoolctl==3.5.0 # via scikit-learn tifffile==2025.3.30 # via scikit-image -tiktoken==0.14.0 +tiktoken==0.12.0 # via # -c requirements/common.txt # gpt-oss @@ -987,7 +986,6 @@ tokenizers==0.22.2 # via # -c requirements/common.txt # -r requirements/test/cuda.in - # sentence-transformers # transformers tomli==2.2.1 # via schemathesis @@ -1041,7 +1039,7 @@ tqdm==4.67.3 # segmentation-models-pytorch # sentence-transformers # transformers -transformers==5.15.1 +transformers==5.5.3 # via # -c requirements/common.txt # -r requirements/test/cuda.in @@ -1068,10 +1066,9 @@ typer==0.15.2 # transformers types-python-dateutil==2.9.0.20241206 # via arrow -typing-extensions==4.16.0 +typing-extensions==4.15.0 # via # -c requirements/common.txt - # aiohttp # aiosignal # albumentations # alembic @@ -1100,9 +1097,7 @@ typing-extensions==4.16.0 # typer # typing-inspection typing-inspection==0.4.2 - # via - # fastapi - # pydantic + # via pydantic tzdata==2024.2 # via pandas uri-template==1.3.0 @@ -1141,3 +1136,5 @@ 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 3910081624..97e0658fb1 100644 --- a/requirements/test/rocm.in +++ b/requirements/test/rocm.in @@ -2,7 +2,7 @@ # testing pytest -tensorizer==2.12.1 +tensorizer==2.10.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.20.0 # required for phi-4-mm test +peft>=0.15.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.15.1 +transformers==5.5.3 tokenizers==0.22.2 schemathesis>=3.39.15 # Required for openai schema test # quantization -bitsandbytes==0.50.1 +bitsandbytes==0.49.2 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.83.0 -grpcio-reflection==1.83.0 +grpcio==1.78.0 +grpcio-reflection==1.78.0 -arctic-inference==0.2.0 # Required for suffix decoding test -numba==0.67.0 # Required for N-gram speculative decoding +arctic-inference==0.1.1 # Required for suffix decoding test +numba==0.65.0 # Required for N-gram speculative decoding numpy -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 +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 instanttensor>=0.1.5 -pydantic>=2.13.4 # 2.11 leads to error on python 3.13 +pydantic>=2.12 # 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 5a40ff1c4a..c39f268709 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.14.3 +aiohttp==3.13.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.2.0 +arctic-inference==0.1.1 # 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.50.1 +bitsandbytes==0.49.2 # via -r requirements/test/rocm.in black==26.3.1 # via datamodel-code-generator @@ -226,25 +226,23 @@ et-xmlfile==2.0.0 # via openpyxl evaluate==0.4.6 # via lm-eval -fastapi==0.141.1 +fastapi==0.135.2 # via # -c requirements/common.txt # -r requirements/test/../common.txt # gpt-oss # model-hosting-container-standards -fastapi-cli==0.0.32 +fastapi-cli==0.0.24 # via fastapi fastapi-cloud-cli==0.16.1 # via fastapi-cli fastar==0.10.0 - # via - # fastapi - # fastapi-cloud-cli + # via fastapi-cloud-cli fastparquet==2026.3.0 # via genai-perf -fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@daf6cc3fd803ecbf85865062c860122bc42ffefa +fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@65d80088fca7a8f567fba30415fbcc80f7d2259c # via -r requirements/test/rocm.in -filelock==3.32.4 +filelock==3.25.2 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -282,7 +280,7 @@ gguf==0.18.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt -google-api-core==2.34.0 +google-api-core==2.30.0 # via # google-cloud-core # google-cloud-storage @@ -303,7 +301,7 @@ google-crc32c==1.8.0 # google-resumable-media google-resumable-media==2.8.0 # via google-cloud-storage -googleapis-common-protos==1.75.2 +googleapis-common-protos==1.73.0 # via # google-api-core # opentelemetry-exporter-otlp-proto-grpc @@ -314,14 +312,14 @@ graphql-core==3.2.8 # via hypothesis-graphql greenlet==3.3.2 # via sqlalchemy -grpcio==1.83.0 +grpcio==1.78.0 # via # -c requirements/rocm.txt # -r requirements/test/rocm.in # grpcio-reflection # opentelemetry-exporter-otlp-proto-grpc # ray -grpcio-reflection==1.83.0 +grpcio-reflection==1.78.0 # via # -c requirements/rocm.txt # -r requirements/test/rocm.in @@ -399,6 +397,8 @@ 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.3.1 +lark==1.2.2 # 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.49.0 +llvmlite==0.47.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.67.0 +numba==0.65.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.44.0 +opentelemetry-api==1.40.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -656,27 +656,27 @@ opentelemetry-api==1.44.0 # opentelemetry-exporter-prometheus # opentelemetry-sdk # opentelemetry-semantic-conventions -opentelemetry-exporter-otlp==1.44.0 +opentelemetry-exporter-otlp==1.40.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt -opentelemetry-exporter-otlp-proto-common==1.44.0 +opentelemetry-exporter-otlp-proto-common==1.40.0 # via # opentelemetry-exporter-otlp-proto-grpc # opentelemetry-exporter-otlp-proto-http -opentelemetry-exporter-otlp-proto-grpc==1.44.0 +opentelemetry-exporter-otlp-proto-grpc==1.40.0 # via opentelemetry-exporter-otlp -opentelemetry-exporter-otlp-proto-http==1.44.0 +opentelemetry-exporter-otlp-proto-http==1.40.0 # via opentelemetry-exporter-otlp -opentelemetry-exporter-prometheus==0.65b0 +opentelemetry-exporter-prometheus==0.61b0 # via ray -opentelemetry-proto==1.44.0 +opentelemetry-proto==1.40.0 # via # opentelemetry-exporter-otlp-proto-common # opentelemetry-exporter-otlp-proto-grpc # opentelemetry-exporter-otlp-proto-http # ray -opentelemetry-sdk==1.44.0 +opentelemetry-sdk==1.40.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -685,7 +685,7 @@ opentelemetry-sdk==1.44.0 # opentelemetry-exporter-prometheus # opentelemetry-semantic-conventions-ai # ray -opentelemetry-semantic-conventions==0.65b0 +opentelemetry-semantic-conventions==0.61b0 # 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.20.0 +peft==0.18.1 # 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.26.0 +prometheus-client==0.24.1 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -800,9 +800,9 @@ propcache==0.4.1 # via # aiohttp # yarl -proto-plus==1.28.4 +proto-plus==1.27.1 # via google-api-core -protobuf==7.36.0 +protobuf==6.33.6 # 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.13.4 +pydantic==2.12.5 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -865,7 +865,7 @@ pydantic==2.13.4 # pydantic-settings # ray # xgrammar -pydantic-core==2.46.4 +pydantic-core==2.41.5 # via pydantic pydantic-extra-types==2.11.1 # via @@ -992,7 +992,7 @@ regex==2026.2.28 # sacrebleu # tiktoken # transformers -requests==2.34.2 +requests==2.32.5 # 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.16.1 +runai-model-streamer==0.15.7 # via # -c requirements/rocm.txt # -r requirements/test/rocm.in -runai-model-streamer-azure==0.16.1 +runai-model-streamer-azure==0.15.7 # via runai-model-streamer -runai-model-streamer-gcs==0.16.1 +runai-model-streamer-gcs==0.15.7 # via runai-model-streamer -runai-model-streamer-s3==0.16.1 +runai-model-streamer-s3==0.15.7 # via runai-model-streamer s3transfer==0.16.0 # via boto3 sacrebleu==2.6.0 # via lm-eval -safetensors==0.8.0 +safetensors==0.7.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.12.1 +tensorizer==2.10.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.14.0 +tiktoken==0.12.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.15.1 +transformers==5.5.3 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -1272,11 +1272,10 @@ typer==0.24.1 # huggingface-hub # perceptron # transformers -typing-extensions==4.16.0 +typing-extensions==4.15.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt - # aiohttp # aiosignal # albumentations # alembic @@ -1379,6 +1378,8 @@ 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 ae93c647bb..5581d0a079 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.14.3 +aiohttp==3.13.4 # 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.141.1 +fastapi==0.135.2 # via # -c requirements/common.txt # gpt-oss -filelock==3.32.4 +filelock==3.25.2 # 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.49.0 +llvmlite==0.47.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.67.0 +numba==0.65.0 # via # -c requirements/xpu.txt # librosa @@ -428,7 +428,7 @@ pycparser==3.0 # via cffi pycryptodomex==3.23.0 # via blobfile -pydantic==2.13.4 +pydantic==2.12.5 # via # -c requirements/common.txt # albumentations @@ -438,7 +438,7 @@ pydantic==2.13.4 # mteb # openai-harmony # pydantic-extra-types -pydantic-core==2.46.4 +pydantic-core==2.41.5 # via pydantic pydantic-extra-types==2.11.1 # via mistral-common @@ -509,7 +509,7 @@ regex==2026.3.32 # sacrebleu # tiktoken # transformers -requests==2.34.2 +requests==2.33.1 # 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.8.0 +safetensors==0.7.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.14.0 +tiktoken==0.12.0 # via # -c requirements/common.txt # gpt-oss @@ -668,7 +668,7 @@ tqdm==4.67.3 # pqdm # sentence-transformers # transformers -transformers==5.15.1 +transformers==5.5.3 # via # -c requirements/common.txt # sentence-transformers @@ -683,10 +683,9 @@ typer==0.24.1 # via # huggingface-hub # transformers -typing-extensions==4.16.0 +typing-extensions==4.15.0 # via # -c requirements/common.txt - # aiohttp # aiosignal # albumentations # anyio diff --git a/requirements/tpu.txt b/requirements/tpu.txt index ab8fa029d8..539f2320ba 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.27.0 +tpu-inference==0.20.0 diff --git a/requirements/xpu.txt b/requirements/xpu.txt index b7d9b5ce71..897e2080da 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.67.0 # Required for N-gram speculative decoding +numba == 0.65.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 a9c6a90ab7..6ca12d8387 100644 --- a/setup.py +++ b/setup.py @@ -891,9 +891,6 @@ def extract_precompiled_and_patch_package( flash_qla_sm70_ext_regex = re.compile( r"flash_qla/ops/gated_delta_rule/chunk/sm70/[^/]+\.so" ) - sm70_sampler_ext_regex = re.compile( - r"vllm/_sm70_sampler_C(?:\.[^/]+)?\.so$" - ) file_members = [] for member in wheel.filelist: if member.filename in exact_members: @@ -913,7 +910,6 @@ def extract_precompiled_and_patch_package( or deep_gemm_regex.match(member.filename) or flash_attn_v100_ext_regex.match(member.filename) or flash_qla_sm70_ext_regex.match(member.filename) - or sm70_sampler_ext_regex.match(member.filename) ): file_members.append(member) @@ -1239,8 +1235,6 @@ def _read_requirements(filename: str) -> list[str]: ext_modules.append(CMakeExtension(name="vllm._rocm_C")) if _is_cuda(): - if _cuda_arch_contains(7, 0): - ext_modules.append(CMakeExtension(name="vllm._sm70_sampler_C")) build_sm70_fa2 = _cuda_arch_contains(7, 0) and not _cuda_arch_at_least(8, 0) if _cuda_arch_at_least(8, 0) or build_sm70_fa2: ext_modules.append(CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa2_C")) @@ -1392,10 +1386,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.12.1"], - "fastsafetensors": ["fastsafetensors >= 0.3.3"], + "tensorizer": ["tensorizer==2.10.1"], + "fastsafetensors": ["fastsafetensors >= 0.2.2"], "instanttensor": ["instanttensor >= 0.1.5"], - "runai": ["runai-model-streamer[s3,gcs,azure] >= 0.16.1"], + "runai": ["runai-model-streamer[s3,gcs,azure] >= 0.15.7"], "audio": [ "av", "scipy", @@ -1408,7 +1402,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.4.0"], + "helion": ["helion==1.0.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_dsv4_quality_tools.py b/tests/benchmarks/test_dsv4_quality_tools.py deleted file mode 100644 index 09697faec9..0000000000 --- a/tests/benchmarks/test_dsv4_quality_tools.py +++ /dev/null @@ -1,141 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from copy import deepcopy - -from benchmarks.benchmark_dsv4_gsm8k_api import _last_integer -from benchmarks.compare_dsv4_quality_results import ( - _compare_gsm8k, - _compare_humaneval, - _compare_needle, -) - - -def test_last_integer_normalizes_signed_integral_decimal() -> None: - assert _last_integer("work... #### -1,024") == -1024 - assert _last_integer("answer: 12.0") == 12 - assert _last_integer("answer: 12.5") is None - assert _last_integer("no numeric answer") is None - - -def _gsm8k_artifact(predictions: list[int | None]) -> dict: - expected = [13, 40] - rows = [ - { - "index": index, - "question": f"question-{index}", - "expected": answer, - "predicted": prediction, - "correct": prediction is not None and prediction == answer, - "invalid": prediction is None, - } - for index, (answer, prediction) in enumerate(zip(expected, predictions)) - ] - return { - "contract": { - "questions": 2, - "few_shot": 5, - "temperature": 0.0, - "top_p": 1.0, - "seed": 42, - "max_tokens": 256, - "strictly_sequential": True, - "train_selection": "first_n", - "test_selection": "first_n", - "prompt_format": "gsm8k_question_answer_v1", - "answer_normalization": "last_signed_integral_decimal_v1", - "stop_sequences": ["Question"], - }, - "input_manifest": { - "train": {"sha256": "train-hash"}, - "test": {"sha256": "test-hash"}, - }, - "correct": sum(bool(row["correct"]) for row in rows), - "invalid": sum(bool(row["invalid"]) for row in rows), - "rows": rows, - } - - -def test_gsm8k_gate_allows_balanced_answer_flips_without_greedy_identity() -> None: - reference = _gsm8k_artifact([13, 42]) - candidate = _gsm8k_artifact([12, 40]) - - result = _compare_gsm8k(reference, candidate) - - assert result["passed"] - assert result["correct_delta"] == 0 - assert result["regressions"] == [0] - assert result["improvements"] == [1] - assert result["exact_prediction_matches"] == 0 - - -def test_gsm8k_gate_rejects_aggregate_drop_contract_drift_and_bad_summary() -> None: - reference = _gsm8k_artifact([13, 40]) - candidate = _gsm8k_artifact([12, 40]) - assert not _compare_gsm8k(reference, candidate)["passed"] - - candidate = deepcopy(reference) - candidate["contract"]["seed"] = 7 - assert not _compare_gsm8k(reference, candidate)["passed"] - - candidate = deepcopy(reference) - candidate["correct"] = 0 - assert not _compare_gsm8k(reference, candidate)["passed"] - - -def test_humaneval_gate_uses_aggregate_quality_and_reports_flips() -> None: - reference = { - "human_eval": { - "passed": 1, - "records": [ - {"task_id": "a", "passed": True, "response": "a0"}, - {"task_id": "b", "passed": False, "response": "b0"}, - ], - } - } - candidate = { - "human_eval": { - "passed": 1, - "records": [ - {"task_id": "a", "passed": False, "response": "a1"}, - {"task_id": "b", "passed": True, "response": "b1"}, - ], - } - } - - result = _compare_humaneval(reference, candidate) - - assert result["passed"] - assert result["regressions"] == ["a"] - assert result["improvements"] == ["b"] - - -def test_needle_gate_uses_aggregate_hits_and_reports_flips() -> None: - reference = [ - { - "sample_id": "a", - "target_tokens": 100, - "depth": 0.5, - "code": "x", - "hit": True, - "hit_anywhere": True, - }, - { - "sample_id": "b", - "target_tokens": 100, - "depth": 0.5, - "code": "y", - "hit": False, - "hit_anywhere": False, - }, - ] - candidate = [ - {**reference[0], "hit": False, "hit_anywhere": False}, - {**reference[1], "hit": True, "hit_anywhere": True}, - ] - - result = _compare_needle(reference, candidate) - - assert result["passed"] - assert result["final_hit_regressions"] == ["a"] - assert result["anywhere_hit_regressions"] == ["a"] diff --git a/tests/config/test_prefix_anchored_swa.py b/tests/config/test_prefix_anchored_swa.py deleted file mode 100644 index 6a801406fc..0000000000 --- a/tests/config/test_prefix_anchored_swa.py +++ /dev/null @@ -1,158 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from types import SimpleNamespace - -import pytest -import torch - -import vllm.platforms as platforms -from vllm.config.attention import AttentionConfig -from vllm.config.vllm import apply_prefix_anchored_swa_constraints -from vllm.model_executor.layers.attention.attention import Attention -from vllm.v1.attention.backend import AttentionType -from vllm.v1.attention.backends.registry import AttentionBackendEnum -from vllm.v1.kv_cache_interface import PrefixAnchoredSWASpec - - -def _config(window: int = 128) -> SimpleNamespace: - return SimpleNamespace( - attention_config=AttentionConfig( - prefix_anchored_decode_window=window, - ), - cache_config=SimpleNamespace( - cache_dtype="auto", - enable_prefix_caching=True, - kv_offloading_size=None, - ), - model_config=SimpleNamespace(dtype=torch.float16, use_mla=False), - parallel_config=SimpleNamespace( - decode_context_parallel_size=1, - prefill_context_parallel_size=1, - ), - kv_transfer_config=None, - speculative_config=None, - ) - - -@pytest.fixture -def sm70_platform(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr( - platforms, - "current_platform", - SimpleNamespace( - is_cuda=lambda: True, - is_device_capability=lambda capability: capability == (7, 0), - ), - ) - - -def test_engine_contract_is_model_identity_independent(sm70_platform: None): - cfg = _config() - - apply_prefix_anchored_swa_constraints(cfg) - - assert cfg.attention_config.backend == AttentionBackendEnum.FLASH_ATTN_V100 - assert not cfg.cache_config.enable_prefix_caching - - -@pytest.mark.parametrize("window", [0, -1]) -def test_engine_option_rejects_nonpositive_window(window: int): - with pytest.raises(ValueError): - AttentionConfig(prefix_anchored_decode_window=window) - - -def test_engine_contract_rejects_non_fp16_kv(sm70_platform: None): - cfg = _config() - cfg.cache_config.cache_dtype = "fp8_e5m2" - - with pytest.raises(ValueError, match="requires an fp16 KV cache"): - apply_prefix_anchored_swa_constraints(cfg) - - -@pytest.mark.parametrize( - ("path", "value", "message"), - [ - ("model_config.use_mla", True, "MLA attention"), - ("attention_config.use_non_causal", True, "causal decoder"), - ("parallel_config.decode_context_parallel_size", 2, "context parallelism"), - ("cache_config.kv_offloading_size", 1.0, "KV-cache offloading"), - ( - "kv_transfer_config", - SimpleNamespace(kv_connector="ExampleConnector"), - "KV connectors", - ), - ("speculative_config", object(), "speculative decoding"), - ( - "attention_config.backend", - AttentionBackendEnum.TRITON_ATTN, - "requires FLASH_ATTN_V100", - ), - ], -) -def test_engine_contract_rejects_unsupported_runtime_combinations( - sm70_platform: None, - path: str, - value: object, - message: str, -): - cfg = _config() - target = cfg - parts = path.split(".") - for part in parts[:-1]: - target = getattr(target, part) - setattr(target, parts[-1], value) - - with pytest.raises(ValueError, match=message): - apply_prefix_anchored_swa_constraints(cfg) - - -def test_engine_contract_rejects_non_sm70(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr( - platforms, - "current_platform", - SimpleNamespace( - is_cuda=lambda: True, - is_device_capability=lambda capability: False, - ), - ) - - with pytest.raises(ValueError, match="SM70 GPU"): - apply_prefix_anchored_swa_constraints(_config()) - - -def test_base_attention_emits_prefix_anchored_spec_without_model_subclass(): - layer = object.__new__(Attention) - object.__setattr__(layer, "attn_type", AttentionType.DECODER) - object.__setattr__( - layer, - "attn_backend", - SimpleNamespace(get_name=lambda: "FLASH_ATTN_V100"), - ) - object.__setattr__(layer, "kv_cache_dtype", "auto") - object.__setattr__(layer, "kv_cache_torch_dtype", torch.float16) - object.__setattr__(layer, "sliding_window", None) - object.__setattr__(layer, "num_kv_heads", 4) - object.__setattr__(layer, "head_size", 128) - object.__setattr__(layer, "head_size_v", 128) - object.__setattr__(layer, "has_sink", False) - vllm_config = SimpleNamespace( - cache_config=SimpleNamespace(block_size=16), - attention_config=SimpleNamespace(prefix_anchored_decode_window=256), - ) - - spec = Attention.get_kv_cache_spec(layer, vllm_config) - - assert isinstance(spec, PrefixAnchoredSWASpec) - assert spec.decode_sliding_window == 256 - - -def test_prefix_anchored_spec_rejects_nonpositive_window(): - with pytest.raises(ValueError, match="greater than zero"): - PrefixAnchoredSWASpec( - block_size=16, - num_kv_heads=1, - head_size=64, - dtype=torch.float16, - decode_sliding_window=0, - ) diff --git a/tests/conftest.py b/tests/conftest.py index 882138f911..3eaebc38bc 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -468,7 +468,7 @@ def _init( model_name, revision=revision, device=self.device, - model_kwargs=model_kwargs, + automodel_args=model_kwargs, trust_remote_code=trust_remote_code, ) else: diff --git a/tests/distributed/test_custom_all_reduce_allocator.py b/tests/distributed/test_custom_all_reduce_allocator.py deleted file mode 100644 index d455012fc9..0000000000 --- a/tests/distributed/test_custom_all_reduce_allocator.py +++ /dev/null @@ -1,224 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import pytest -import regex as re - -from vllm.distributed.device_communicators import custom_all_reduce - -pytestmark = pytest.mark.skip_global_cleanup - - -def _patch_allocator_settings(monkeypatch: pytest.MonkeyPatch) -> list[str]: - settings: list[str] = [] - monkeypatch.setattr(custom_all_reduce.current_platform, "is_cuda", lambda: True) - monkeypatch.setattr( - custom_all_reduce.torch.cuda.memory, - "_set_allocator_settings", - settings.append, - ) - return settings - - -def test_cuda_ipc_capture_temporarily_disables_expandable_segments( - monkeypatch: pytest.MonkeyPatch, -) -> None: - allocator_conf = ( - "max_split_size_mb:512,garbage_collection_threshold:0.8," - "expandable_segments:True" - ) - monkeypatch.setenv( - "PYTORCH_CUDA_ALLOC_CONF", - allocator_conf, - ) - settings = _patch_allocator_settings(monkeypatch) - - with custom_all_reduce._disable_expandable_segments_for_cuda_ipc(True): - assert settings == [allocator_conf.replace("True", "False")] - - assert settings == [allocator_conf.replace("True", "False"), allocator_conf] - - -def test_cuda_ipc_capture_preserves_allocator_state( - monkeypatch: pytest.MonkeyPatch, -) -> None: - allocator_conf = ( - "max_split_size_mb:512,roundup_power2_divisions:4," - "garbage_collection_threshold:0.8,expandable_segments:True" - ) - monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", allocator_conf) - monkeypatch.setattr(custom_all_reduce.current_platform, "is_cuda", lambda: True) - allocator_state: dict[str, int | float | bool] = { - "max_split_size_mb": 512, - "roundup_power2_divisions": 4, - "garbage_collection_threshold": 0.8, - "expandable_segments": True, - } - - def set_allocator_settings(conf: str) -> None: - # Model PyTorch's parseArgs behavior: these options reset whenever - # they are omitted from a runtime settings update. - allocator_state.update( - max_split_size_mb=-1, - roundup_power2_divisions=0, - garbage_collection_threshold=0.0, - ) - for key in allocator_state: - match = re.search(rf"(?:^|,){key}:([^,]+)", conf) - if match is None: - continue - value = match.group(1) - if key == "garbage_collection_threshold": - allocator_state[key] = float(value) - elif key == "expandable_segments": - allocator_state[key] = value == "True" - else: - allocator_state[key] = int(value) - - monkeypatch.setattr( - custom_all_reduce.torch.cuda.memory, - "_set_allocator_settings", - set_allocator_settings, - ) - - with custom_all_reduce._disable_expandable_segments_for_cuda_ipc(True): - assert allocator_state == { - "max_split_size_mb": 512, - "roundup_power2_divisions": 4, - "garbage_collection_threshold": 0.8, - "expandable_segments": False, - } - - assert allocator_state == { - "max_split_size_mb": 512, - "roundup_power2_divisions": 4, - "garbage_collection_threshold": 0.8, - "expandable_segments": True, - } - - -def test_cuda_ipc_capture_supports_unified_allocator_conf( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.delenv("PYTORCH_CUDA_ALLOC_CONF", raising=False) - allocator_conf = "max_split_size_mb:512,expandable_segments:True" - monkeypatch.setenv("PYTORCH_ALLOC_CONF", allocator_conf) - settings = _patch_allocator_settings(monkeypatch) - - with custom_all_reduce._disable_expandable_segments_for_cuda_ipc(True): - pass - - assert settings == [allocator_conf.replace("True", "False"), allocator_conf] - - -@pytest.mark.parametrize("legacy_conf", ["", "expandable_segments:False"]) -def test_legacy_allocator_conf_takes_precedence_over_unified_conf( - monkeypatch: pytest.MonkeyPatch, - legacy_conf: str, -) -> None: - monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", legacy_conf) - monkeypatch.setenv("PYTORCH_ALLOC_CONF", "expandable_segments:True") - settings = _patch_allocator_settings(monkeypatch) - - with custom_all_reduce._disable_expandable_segments_for_cuda_ipc(True): - pass - - assert settings == [] - - -def test_enabled_legacy_allocator_conf_takes_precedence( - monkeypatch: pytest.MonkeyPatch, -) -> None: - legacy_conf = "max_split_size_mb:512,expandable_segments:True" - monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", legacy_conf) - monkeypatch.setenv("PYTORCH_ALLOC_CONF", "expandable_segments:False") - settings = _patch_allocator_settings(monkeypatch) - - with custom_all_reduce._disable_expandable_segments_for_cuda_ipc(True): - pass - - assert settings == [legacy_conf.replace("True", "False"), legacy_conf] - - -def test_cuda_ipc_capture_preserves_bracketed_allocator_syntax( - monkeypatch: pytest.MonkeyPatch, -) -> None: - allocator_conf = ( - "roundup_power2_divisions:[256:1,512:2,>:4], expandable_segments : True" - ) - monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", allocator_conf) - settings = _patch_allocator_settings(monkeypatch) - - with custom_all_reduce._disable_expandable_segments_for_cuda_ipc(True): - pass - - assert settings == [ - allocator_conf.replace("True", "False"), - allocator_conf, - ] - - -def test_cuda_ipc_capture_restores_allocator_after_error( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") - settings = _patch_allocator_settings(monkeypatch) - - with ( - pytest.raises(RuntimeError, match="capture failed"), - custom_all_reduce._disable_expandable_segments_for_cuda_ipc(True), - ): - raise RuntimeError("capture failed") - - assert settings == ["expandable_segments:False", "expandable_segments:True"] - - -def test_custom_allreduce_registers_graph_buffers_before_allocator_restore( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") - settings = _patch_allocator_settings(monkeypatch) - communicator = object.__new__(custom_all_reduce.CustomAllreduce) - communicator.disabled = False - communicator._IS_CAPTURING = False - communicator._ptr = None - registration_state: list[tuple[list[str], bool]] = [] - - def register_graph_buffers() -> None: - registration_state.append((list(settings), communicator._IS_CAPTURING)) - - monkeypatch.setattr( - communicator, - "register_graph_buffers", - register_graph_buffers, - ) - - with communicator.capture(): - assert communicator._IS_CAPTURING - assert settings == ["expandable_segments:False"] - - assert registration_state == [(["expandable_segments:False"], False)] - assert settings == ["expandable_segments:False", "expandable_segments:True"] - - -@pytest.mark.parametrize( - ("active", "allocator_conf"), - [ - (False, "expandable_segments:True"), - (True, "expandable_segments:False"), - (True, ""), - ], -) -def test_cuda_ipc_capture_allocator_guard_is_otherwise_a_noop( - monkeypatch: pytest.MonkeyPatch, - active: bool, - allocator_conf: str, -) -> None: - monkeypatch.delenv("PYTORCH_ALLOC_CONF", raising=False) - monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", allocator_conf) - settings = _patch_allocator_settings(monkeypatch) - - with custom_all_reduce._disable_expandable_segments_for_cuda_ipc(active): - pass - - assert settings == [] diff --git a/tests/distributed/test_static_tensor_transfer.py b/tests/distributed/test_static_tensor_transfer.py deleted file mode 100644 index 4a1890f5dc..0000000000 --- a/tests/distributed/test_static_tensor_transfer.py +++ /dev/null @@ -1,63 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from typing import Any - -import pytest -import torch - -from vllm.distributed.parallel_state import GroupCoordinator - - -class _DummyWork: - def wait(self) -> None: - pass - - -def _group(rank_in_group: int, world_size: int = 2) -> GroupCoordinator: - group = GroupCoordinator.__new__(GroupCoordinator) - group.world_size = world_size - group.rank_in_group = rank_in_group - group.ranks = list(range(world_size)) - group.use_cpu_custom_send_recv = False - group.device_group = None - group.cpu_group = None - return group - - -def test_static_tensor_dict_transfer_uses_caller_owned_buffers( - monkeypatch: pytest.MonkeyPatch, -) -> None: - calls: list[tuple[str, int, torch.Tensor]] = [] - - def fake_isend(tensor: torch.Tensor, *, dst: int, group: Any) -> _DummyWork: - del group - calls.append(("send", dst, tensor)) - return _DummyWork() - - def fake_irecv(tensor: torch.Tensor, *, src: int, group: Any) -> _DummyWork: - del group - tensor.fill_(7) - calls.append(("recv", src, tensor)) - return _DummyWork() - - monkeypatch.setattr(torch.distributed, "isend", fake_isend) - monkeypatch.setattr(torch.distributed, "irecv", fake_irecv) - - sent = torch.arange(8, dtype=torch.float32) - send_handles = _group(0).isend_tensor_dict_static({"hidden_states": sent}) - - received = torch.empty_like(sent) - recv_handles = _group(1).irecv_tensor_dict_static({"hidden_states": received}) - - assert len(send_handles) == 1 - assert len(recv_handles) == 1 - assert calls == [("send", 1, sent), ("recv", 0, received)] - torch.testing.assert_close(received, torch.full_like(received, 7)) - - -def test_static_tensor_dict_transfer_rejects_non_tensors() -> None: - with pytest.raises(TypeError, match="only accepts tensors"): - _group(0).isend_tensor_dict_static( # type: ignore[arg-type] - {"hidden_states": object()} - ) diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 541f3d5d42..9f2d8819ef 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -237,35 +237,19 @@ def _fake_qwen_hybrid_model_config(): ) -def _apply_sm70_defaults( - monkeypatch, - *, - env=None, - speculative_config=None, - enable_prefix_caching=None, - max_num_seqs=None, -): +def _apply_sm70_defaults(monkeypatch, *, env=None, speculative_config=None): for key in ( "VLLM_1CAT_ENABLE_SM70_MTP_DEFAULTS", "VLLM_1CAT_ENABLE_QWEN35_MTP_DEFAULTS", "VLLM_1CAT_DISABLE_SM70_MTP_DEFAULTS", "VLLM_1CAT_DISABLE_QWEN35_MTP_DEFAULTS", - "VLLM_SM70_MTP_SPLIT_DRAFT_CUDAGRAPHS", ): monkeypatch.delenv(key, raising=False) for key, value in (env or {}).items(): monkeypatch.setenv(key, value) monkeypatch.setattr("vllm.engine.arg_utils.current_platform", _FakeSM70Platform()) - monkeypatch.setattr( - "vllm.engine.arg_utils.get_model_path", lambda model, revision: model - ) - args = EngineArgs( - model="dummy", - tensor_parallel_size=4, - enable_prefix_caching=enable_prefix_caching, - max_num_seqs=max_num_seqs, - ) + args = EngineArgs(model="dummy", tensor_parallel_size=4) args.speculative_config = speculative_config args._maybe_apply_sm70_mtp_defaults( UsageContext.OPENAI_API_SERVER, @@ -313,39 +297,6 @@ def test_sm70_mtp_defaults_require_env_opt_in(monkeypatch): ] -def test_sm70_mtp_split_cudagraphs_are_opt_in(monkeypatch): - args = _apply_sm70_defaults( - monkeypatch, - env={ - "VLLM_1CAT_ENABLE_SM70_MTP_DEFAULTS": "1", - "VLLM_SM70_MTP_SPLIT_DRAFT_CUDAGRAPHS": "1", - }, - ) - - assert args.compilation_config.cudagraph_capture_sizes == [5, 10, 20] - - -def test_sm70_mtp_split_cudagraphs_cover_production_batches(monkeypatch): - args = _apply_sm70_defaults( - monkeypatch, - env={ - "VLLM_1CAT_ENABLE_SM70_MTP_DEFAULTS": "1", - "VLLM_SM70_MTP_SPLIT_DRAFT_CUDAGRAPHS": "1", - }, - max_num_seqs=16, - ) - - assert args.compilation_config.cudagraph_capture_sizes == [ - 5, - 10, - 20, - 30, - 40, - 60, - 80, - ] - - def test_sm70_mtp_defaults_treat_zero_env_as_disabled(monkeypatch): args = _apply_sm70_defaults( monkeypatch, @@ -388,37 +339,6 @@ def test_sm70_explicit_mtp_still_gets_safe_defaults(monkeypatch): assert args.max_num_seqs == 4 -def test_sm70_explicit_dflash_preserves_probabilistic_default(monkeypatch): - args = _apply_sm70_defaults( - monkeypatch, - speculative_config={"method": "dflash", "num_speculative_tokens": 7}, - ) - - assert args.speculative_config == { - "method": "dflash", - "num_speculative_tokens": 7, - "draft_sample_method": "probabilistic", - } - assert args.enable_prefix_caching is True - assert args.mamba_cache_mode == "align" - - -def test_sm70_explicit_dflash_accepts_explicit_greedy(monkeypatch): - args = _apply_sm70_defaults( - monkeypatch, - speculative_config={ - "method": "dflash", - "num_speculative_tokens": 7, - "draft_sample_method": "greedy", - }, - enable_prefix_caching=True, - ) - - assert args.speculative_config["draft_sample_method"] == "greedy" - assert args.enable_prefix_caching is True - assert args.mamba_cache_mode == "align" - - @pytest.mark.parametrize( ("cli_args", "expected"), [ diff --git a/tests/kernels/attention/test_sm70_flash_v100_anchored_swa.py b/tests/kernels/attention/test_sm70_flash_v100_anchored_swa.py deleted file mode 100644 index d93cbc878e..0000000000 --- a/tests/kernels/attention/test_sm70_flash_v100_anchored_swa.py +++ /dev/null @@ -1,288 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""SM70 Flash-V100 anchored decode-window mask kernel tests. - -Validates the masked (second) kernel instantiations of the paged decode and -paged prefill kernels against a plain torch reference: -``causal AND (kv < anchor_len OR q_abs - kv < window)``, element-wise. - -Requires a CUDA device and the flash_attn_v100 extension built with the -anchored kernel arguments; skipped otherwise. -""" - -from collections.abc import Callable -from typing import Any - -import pytest -import torch - -cuda_available = torch.cuda.is_available() -FlashOp = Callable[..., Any] -flash_ops: tuple[FlashOp, FlashOp] | None = None -if cuda_available: - try: - from flash_attn_v100 import ( - flash_attn_decode_paged, - flash_attn_prefill_paged, - ) - - flash_ops = (flash_attn_decode_paged, flash_attn_prefill_paged) - except ImportError: - flash_ops = None - -pytestmark = pytest.mark.skipif( - not cuda_available or flash_ops is None, - reason="requires CUDA and the flash_attn_v100 extension", -) - - -def _require_flash_ops() -> tuple[FlashOp, FlashOp]: - assert flash_ops is not None - return flash_ops - - -def _anchored_mask( - q_positions: torch.Tensor, - kv_len: int, - anchor_len: int, - window: int, - device: str, -) -> torch.Tensor: - """Boolean keep-mask [num_q, kv_len]; True = attend.""" - kv = torch.arange(kv_len, device=device)[None, :] - q = q_positions[:, None] - causal = kv <= q - in_prefix = kv < anchor_len - in_window = (q - kv) < window - return causal & (in_prefix | in_window) - - -def _ref_attention( - q: torch.Tensor, # [num_q, H, D] - k: torch.Tensor, # [kv_len, H_kv, D] - v: torch.Tensor, - mask: torch.Tensor, # [num_q, kv_len] - scale: float, -) -> torch.Tensor: - group = q.shape[1] // k.shape[1] - k_exp = k.repeat_interleave(group, dim=1) - v_exp = v.repeat_interleave(group, dim=1) - scores = torch.einsum("qhd,khd->hqk", q.float(), k_exp.float()) * scale - scores = scores.masked_fill(~mask[None], float("-inf")) - probs = torch.softmax(scores, dim=-1) - return torch.einsum("hqk,khd->qhd", probs, v_exp.float()) - - -def _build_paged_kv( - kv_len: int, - block_size: int, - num_kv_heads: int, - head_dim: int, - device: str, - dtype: torch.dtype, - seed: int, -): - torch.manual_seed(seed) - num_blocks = (kv_len + block_size - 1) // block_size + 2 - k_cache = torch.randn( - num_blocks, block_size, num_kv_heads, head_dim, device=device, dtype=dtype - ) - v_cache = torch.randn_like(k_cache) - block_table = torch.arange(num_blocks, dtype=torch.int32, device=device).unsqueeze( - 0 - ) - k_lin = k_cache.reshape(-1, num_kv_heads, head_dim)[:kv_len] - v_lin = v_cache.reshape(-1, num_kv_heads, head_dim)[:kv_len] - return k_cache, v_cache, block_table, k_lin, v_lin - - -@pytest.mark.parametrize( - "head_dim,num_heads,num_kv_heads", - [(256, 8, 4), (128, 8, 8), (128, 8, 4), (64, 4, 4)], -) -# kv_len=1100 spans three 512-token partitions with the middle partition -# entirely inside the masked gap (exercises the neutral-stats path). -@pytest.mark.parametrize( - "anchor,window,kv_len", [(32, 16, 96), (16, 8, 40), (32, 16, 1100)] -) -def test_decode_paged_anchored_matches_reference( - head_dim, num_heads, num_kv_heads, anchor, window, kv_len -): - device = "cuda" - dtype = torch.float16 - block_size = 16 - scale = head_dim**-0.5 - - k_cache, v_cache, block_table, k_lin, v_lin = _build_paged_kv( - kv_len, block_size, num_kv_heads, head_dim, device, dtype, seed=0 - ) - q = torch.randn(1, num_heads, head_dim, device=device, dtype=dtype) - seq_lens = torch.tensor([kv_len], dtype=torch.int32, device=device) - anchor_lens = torch.tensor([anchor], dtype=torch.int32, device=device) - - out = torch.empty_like(q) - decode_op, _ = _require_flash_ops() - decode_op( - q, - k_cache, - v_cache, - block_table, - seq_lens, - softmax_scale=scale, - out=out, - anchor_lens=anchor_lens, - anchored_window=window, - ) - - q_pos = torch.tensor([kv_len - 1], device=device) - mask = _anchored_mask(q_pos, kv_len, anchor, window, device) - assert not mask[0].all(), "test must exercise a non-trivial gap" - ref = _ref_attention(q[0].unsqueeze(0), k_lin, v_lin, mask, scale) - torch.testing.assert_close(out[0].float(), ref[0], atol=2e-2, rtol=2e-2) - - -def test_decode_paged_anchored_off_path_unchanged(): - """Without anchor arguments the decode op matches its own baseline.""" - device = "cuda" - dtype = torch.float16 - head_dim, num_heads, num_kv_heads = 128, 8, 8 - kv_len, block_size = 96, 16 - scale = head_dim**-0.5 - - k_cache, v_cache, block_table, k_lin, v_lin = _build_paged_kv( - kv_len, block_size, num_kv_heads, head_dim, device, dtype, seed=1 - ) - q = torch.randn(1, num_heads, head_dim, device=device, dtype=dtype) - seq_lens = torch.tensor([kv_len], dtype=torch.int32, device=device) - - out = torch.empty_like(q) - decode_op, _ = _require_flash_ops() - decode_op(q, k_cache, v_cache, block_table, seq_lens, softmax_scale=scale, out=out) - q_pos = torch.tensor([kv_len - 1], device=device) - causal = _anchored_mask(q_pos, kv_len, kv_len, kv_len, device) - assert causal[0].all() - ref = _ref_attention(q[0].unsqueeze(0), k_lin, v_lin, causal, scale) - torch.testing.assert_close(out[0].float(), ref[0], atol=2e-2, rtol=2e-2) - - -def test_decode_paged_uses_per_request_anchor_lengths(): - device = "cuda" - dtype = torch.float16 - head_dim, num_heads, num_kv_heads = 128, 8, 4 - kv_len, block_size, window = 96, 16, 16 - anchors = (16, 32) - scale = head_dim**-0.5 - - k_cache, v_cache, block_table, k_lin, v_lin = _build_paged_kv( - kv_len, block_size, num_kv_heads, head_dim, device, dtype, seed=3 - ) - block_table = block_table.repeat(2, 1) - q = torch.randn(2, num_heads, head_dim, device=device, dtype=dtype) - seq_lens = torch.full((2,), kv_len, dtype=torch.int32, device=device) - anchor_lens = torch.tensor(anchors, dtype=torch.int32, device=device) - - decode_op, _ = _require_flash_ops() - out = torch.empty_like(q) - decode_op( - q, - k_cache, - v_cache, - block_table, - seq_lens, - softmax_scale=scale, - out=out, - anchor_lens=anchor_lens, - anchored_window=window, - ) - - q_pos = torch.tensor([kv_len - 1], device=device) - for batch_idx, anchor in enumerate(anchors): - mask = _anchored_mask(q_pos, kv_len, anchor, window, device) - ref = _ref_attention(q[batch_idx].unsqueeze(0), k_lin, v_lin, mask, scale) - torch.testing.assert_close(out[batch_idx].float(), ref[0], atol=2e-2, rtol=2e-2) - - -@pytest.mark.parametrize( - "head_dim,num_heads,num_kv_heads", - [(256, 8, 4), (128, 8, 4)], -) -@pytest.mark.parametrize( - "anchor,window,kv_len,q_len", [(32, 16, 96, 8), (16, 8, 48, 48)] -) -def test_prefill_paged_anchored_matches_reference( - head_dim, num_heads, num_kv_heads, anchor, window, kv_len, q_len -): - """Masked paged prefill vs torch reference. - - Covers both the chunked-continuation shape (q_len < kv_len, queries in - the decode region) and the full-sequence shape (q_len == kv_len, where - prompt-region queries must stay purely causal). - """ - device = "cuda" - dtype = torch.float16 - block_size = 16 - scale = head_dim**-0.5 - - k_cache, v_cache, block_table, k_lin, v_lin = _build_paged_kv( - kv_len, block_size, num_kv_heads, head_dim, device, dtype, seed=2 - ) - q = torch.randn(1, q_len, num_heads, head_dim, device=device, dtype=dtype) - seq_lens = torch.tensor([kv_len], dtype=torch.int32, device=device) - anchor_lens = torch.tensor([anchor], dtype=torch.int32, device=device) - - _, prefill_op = _require_flash_ops() - out = prefill_op( - q, - k_cache, - v_cache, - block_table, - seq_lens, - softmax_scale=scale, - causal=True, - anchor_lens=anchor_lens, - anchored_window=window, - ) - - # Query absolute positions: the last q_len positions of the sequence. - q_pos = torch.arange(kv_len - q_len, kv_len, device=device) - mask = _anchored_mask(q_pos, kv_len, anchor, window, device) - ref = _ref_attention(q[0], k_lin, v_lin, mask, scale) - torch.testing.assert_close(out[0].float(), ref.float(), atol=2e-2, rtol=2e-2) - - -def test_prefill_paged_uses_per_request_anchor_lengths(): - device = "cuda" - dtype = torch.float16 - head_dim, num_heads, num_kv_heads = 128, 8, 4 - kv_len, q_len, block_size, window = 96, 8, 16, 16 - anchors = (16, 32) - scale = head_dim**-0.5 - - k_cache, v_cache, block_table, k_lin, v_lin = _build_paged_kv( - kv_len, block_size, num_kv_heads, head_dim, device, dtype, seed=4 - ) - block_table = block_table.repeat(2, 1) - q = torch.randn(2, q_len, num_heads, head_dim, device=device, dtype=dtype) - seq_lens = torch.full((2,), kv_len, dtype=torch.int32, device=device) - anchor_lens = torch.tensor(anchors, dtype=torch.int32, device=device) - - _, prefill_op = _require_flash_ops() - out = prefill_op( - q, - k_cache, - v_cache, - block_table, - seq_lens, - softmax_scale=scale, - causal=True, - anchor_lens=anchor_lens, - anchored_window=window, - ) - - q_pos = torch.arange(kv_len - q_len, kv_len, device=device) - for batch_idx, anchor in enumerate(anchors): - mask = _anchored_mask(q_pos, kv_len, anchor, window, device) - ref = _ref_attention(q[batch_idx], k_lin, v_lin, mask, scale) - torch.testing.assert_close( - out[batch_idx].float(), ref.float(), atol=2e-2, rtol=2e-2 - ) diff --git a/tests/kernels/attention/test_sm70_flash_v100_decode_planner.py b/tests/kernels/attention/test_sm70_flash_v100_decode_planner.py index 9f7f609587..d168663a14 100644 --- a/tests/kernels/attention/test_sm70_flash_v100_decode_planner.py +++ b/tests/kernels/attention/test_sm70_flash_v100_decode_planner.py @@ -188,70 +188,6 @@ def test_static_decode_short_workspace_can_preserve_partition_boundaries( assert plan.workspace_num_partitions == 4 -@pytest.mark.parametrize( - ( - "seq_len", - "max_seq_len_hint", - "workspace_seq_capacity", - "routing_enabled", - "capture_active", - "expected_hint", - ), - ( - (4096, 4096, 8191, True, False, 4096), - (16001, 16001, 262144, True, False, 16001), - (4096, None, 8191, True, False, 8191), - (1, 1, 262144, True, True, 262144), - (4096, 4096, 8191, False, False, 0), - ), -) -def test_xqa_batch_context_route_prefers_live_context( - monkeypatch, - seq_len, - max_seq_len_hint, - workspace_seq_capacity, - routing_enabled, - capture_active, - expected_hint, -) -> None: - _clear_decode_caches() - captured_hints: list[int] = [] - monkeypatch.setattr( - flash_attn_v100, - "_cuda_graph_capture_active", - lambda: capture_active, - ) - - def fake_xqa(*args): - captured_hints.append(args[-1]) - return torch.empty_like(args[0]) - - monkeypatch.setattr( - flash_attn_v100.flash_attn_v100_cuda, - "decode_paged_xqa_fwd", - fake_xqa, - ) - q = torch.empty((4, 6, 256), dtype=torch.float16) - k_cache = torch.empty((2048, 16, 1, 256), dtype=torch.uint8) - v_cache = torch.empty_like(k_cache) - block_table = torch.zeros((4, 2048), dtype=torch.int32) - seq_lens = torch.full((4,), seq_len, dtype=torch.int32) - - flash_attn_v100.flash_attn_decode_paged_xqa( - q, - k_cache, - v_cache, - block_table, - seq_lens, - kv_cache_dtype="fp8_e5m2", - max_seq_len_hint=max_seq_len_hint, - workspace_seq_capacity_hint=workspace_seq_capacity, - batch_context_routing=routing_enabled, - ) - - assert captured_hints == [expected_hint] - - @torch.inference_mode() def test_stale_active_num_partitions_does_not_truncate_decode( monkeypatch, diff --git a/tests/kernels/attention/test_sm70_flash_v100_grouped_verify.py b/tests/kernels/attention/test_sm70_flash_v100_grouped_verify.py deleted file mode 100644 index 84da54590b..0000000000 --- a/tests/kernels/attention/test_sm70_flash_v100_grouped_verify.py +++ /dev/null @@ -1,245 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from __future__ import annotations - -import math - -import pytest -import torch - -_EXPANDED_BASELINE_MEAN_ERROR = { - (16, 8, 120): 5.159488409844926e-6, - (16, 8, 121): 8.621820597909391e-6, - (16, 3, 253): 6.551763817697065e-6, - (16, 4, 252): 5.888879059057217e-6, - (16, 5, 507): 4.691817139246268e-6, - (16, 6, 506): 4.563758920994587e-6, - (16, 8, 1016): 3.1533952551399125e-6, -} - - -def _require_grouped_verify(): - if not torch.cuda.is_available(): - pytest.skip("CUDA is required") - if torch.cuda.get_device_capability() != (7, 0): - pytest.skip("grouped Flash-V100 verification is SM70-only") - flash_attn_v100 = pytest.importorskip("flash_attn_v100") - if not hasattr(flash_attn_v100, "flash_attn_grouped_verify_paged"): - pytest.skip("Flash-V100 extension lacks grouped verification") - return flash_attn_v100 - - -def _make_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, page_size, 1, 256) - key_source = torch.randn(source_shape, dtype=torch.float16, device="cuda").mul_( - 0.25 - ) - value_source = torch.randn_like(key_source).mul_(0.25) - key_cache = key_source.to(torch.float8_e5m2).view(torch.uint8) - value_cache = value_source.to(torch.float8_e5m2).view(torch.uint8) - 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, - value_cache: torch.Tensor, - block_table: torch.Tensor, - seq_len: int, - *, - k_scale: float, - v_scale: float, -) -> torch.Tensor: - query_len = query.shape[0] - prefix_len = seq_len - query_len - logical_pages = math.ceil(seq_len / key_cache.shape[1]) - physical_pages = block_table[0, :logical_pages].to(dtype=torch.long) - key = ( - key_cache.view(torch.float8_e5m2)[physical_pages] - .flatten(0, 1)[:seq_len, 0] - .float() - .mul_(k_scale) - ) - value = ( - value_cache.view(torch.float8_e5m2)[physical_pages] - .flatten(0, 1)[:seq_len, 0] - .float() - .mul_(v_scale) - ) - rows = [] - for token_idx in range(query_len): - visible = prefix_len + token_idx + 1 - scores = torch.einsum( - "hd,nd->hn", query[token_idx].float(), key[:visible] - ).mul_(256**-0.5) - probabilities = torch.softmax(scores, dim=-1) - rows.append(torch.einsum("hn,nd->hd", probabilities, value[:visible])) - return torch.stack(rows).half() - - -@pytest.mark.parametrize( - ("page_size", "query_len", "prefix_len", "strict_fp32_gate"), - [ - (16, 1, 127, True), - (16, 8, 120, False), - (16, 8, 121, False), - (16, 3, 253, False), - (16, 4, 252, False), - (16, 5, 507, False), - (16, 6, 506, False), - (16, 8, 1016, False), - (16, 3, 1025, True), - (784, 5, 2049, True), - (1648, 8, 4097, True), - (3296, 8, 8193, True), - ], -) -@torch.inference_mode() -def test_grouped_verify_matches_fp32_reference_with_random_pages( - page_size: int, - query_len: int, - prefix_len: int, - strict_fp32_gate: bool, -) -> None: - flash_attn_v100 = _require_grouped_verify() - torch.manual_seed(20260824 + page_size + query_len) - query, key_cache, value_cache, block_table, seq_lens = _make_case( - page_size=page_size, - query_len=query_len, - prefix_len=prefix_len, - ) - k_scale = 0.5 - v_scale = 2.0 - expected = _reference( - query, - key_cache, - value_cache, - block_table, - int(seq_lens.item()), - k_scale=k_scale, - v_scale=v_scale, - ) - two_pass = flash_attn_v100.flash_attn_grouped_verify_paged( - query, - key_cache, - value_cache, - block_table, - seq_lens, - k_scale=k_scale, - v_scale=v_scale, - one_pass=False, - ).clone() - one_pass = flash_attn_v100.flash_attn_grouped_verify_paged( - query, - key_cache, - value_cache, - block_table, - seq_lens, - k_scale=k_scale, - v_scale=v_scale, - one_pass=True, - ).clone() - 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, - ) - accepted_xqa = flash_attn_v100.flash_attn_decode_paged_xqa( - query, - key_cache, - value_cache, - decode_block_table, - decode_seq_lens, - kv_cache_dtype="fp8_e5m2", - k_scale=k_scale, - v_scale=v_scale, - max_seq_len_hint=prefix_len + query_len, - workspace_seq_capacity_hint=prefix_len + query_len, - ) - torch.accelerator.synchronize() - - for actual in (two_pass, one_pass): - difference = actual.float().sub(expected.float()).abs() - if strict_fp32_gate: - assert difference.max().item() <= 6.2e-5 - assert difference.mean().item() <= 6.0e-6 - one_pass_error = one_pass.float().sub(expected.float()).abs() - accepted_error = accepted_xqa.float().sub(expected.float()).abs() - assert one_pass_error.max().item() <= accepted_error.max().item() - if strict_fp32_gate: - mean_error_limit = accepted_error.mean().item() - else: - mean_error_limit = _EXPANDED_BASELINE_MEAN_ERROR[ - (page_size, query_len, prefix_len) - ] - assert one_pass_error.mean().item() <= mean_error_limit + 5.0e-7 - torch.testing.assert_close(one_pass, two_pass, atol=6.2e-5, rtol=2.0e-3) - - -@torch.inference_mode() -def test_grouped_verify_cuda_graph_replay_tracks_runtime_seq_len() -> 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, - prefix_len=4097, - ) - output = torch.empty_like(query) - flash_attn_v100.flash_attn_grouped_verify_paged( - query, - key_cache, - value_cache, - block_table, - seq_lens, - out=output, - one_pass=True, - ) - torch.accelerator.synchronize() - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - flash_attn_v100.flash_attn_grouped_verify_paged( - query, - key_cache, - value_cache, - block_table, - seq_lens, - out=output, - one_pass=True, - ) - - for prefix_len in (127, 2049, 4097): - seq_lens.fill_(prefix_len + query.shape[0]) - graph.replay() - torch.accelerator.synchronize() - expected = _reference( - query, - key_cache, - value_cache, - block_table, - prefix_len + query.shape[0], - k_scale=1.0, - v_scale=1.0, - ) - difference = output.float().sub(expected.float()).abs() - assert difference.max().item() <= 6.2e-5 - assert difference.mean().item() <= 6.0e-6 diff --git a/tests/kernels/attention/test_sm70_flash_v100_smallq_metadata.py b/tests/kernels/attention/test_sm70_flash_v100_smallq_metadata.py deleted file mode 100644 index 3912f0f614..0000000000 --- a/tests/kernels/attention/test_sm70_flash_v100_smallq_metadata.py +++ /dev/null @@ -1,232 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import pytest -import torch - -from vllm.v1.attention.backends.flash_attn_v100 import ( - _sm70_prepare_grouped_smallq_decode_metadata, - _sm70_prepare_smallq_decode_metadata, -) - - -def _reference_smallq_metadata( - block_table: torch.Tensor, - seq_lens: torch.Tensor, - query_start_loc: torch.Tensor, - *, - num_query_tokens: int, - real_num_query_tokens: int, -) -> tuple[torch.Tensor, torch.Tensor]: - query_lens = query_start_loc[1:] - query_start_loc[:-1] - repeat_query_lens = query_lens.clone() - repeat_query_lens[-1] += num_query_tokens - real_num_query_tokens - effective_seq_lens = torch.maximum(seq_lens, query_lens) - decode_block_table = torch.repeat_interleave( - block_table.clamp_min(0), - repeat_query_lens, - dim=0, - output_size=num_query_tokens, - ).contiguous() - seq_lens_rep = torch.repeat_interleave( - effective_seq_lens, - repeat_query_lens, - output_size=num_query_tokens, - ) - query_lens_rep = torch.repeat_interleave( - query_lens, - repeat_query_lens, - output_size=num_query_tokens, - ) - start_locs_rep = torch.repeat_interleave( - query_start_loc[:-1], - repeat_query_lens, - output_size=num_query_tokens, - ) - token_indices = torch.arange( - num_query_tokens, dtype=seq_lens.dtype, device=seq_lens.device - ) - decode_seq_lens = ( - seq_lens_rep - query_lens_rep + token_indices - start_locs_rep + 1 - ).contiguous() - if real_num_query_tokens < num_query_tokens: - padding_mask = token_indices >= real_num_query_tokens - decode_seq_lens = decode_seq_lens.masked_fill(padding_mask, 0) - decode_block_table = decode_block_table.masked_fill(padding_mask[:, None], 0) - return decode_block_table, decode_seq_lens - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -@pytest.mark.parametrize( - ("query_lens", "seq_lens", "num_query_tokens"), - [ - ([8], [128], 8), - ([3, 2, 0], [8, 7, 0], 6), - ([1, 4, 2, 0], [1, 10, 8, 0], 8), - ], -) -def test_sm70_fused_smallq_metadata_matches_repeat_interleave( - query_lens: list[int], - seq_lens: list[int], - num_query_tokens: int, -): - if torch.cuda.get_device_capability() != (7, 0): - pytest.skip("SM70 CUDA device required") - - device = torch.device("cuda") - num_reqs = len(query_lens) - query_start_loc = torch.tensor( - [0, *torch.tensor(query_lens).cumsum(0).tolist()], - dtype=torch.int32, - device=device, - ) - real_num_query_tokens = sum(query_lens) - seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int32, device=device) - block_table = torch.arange( - num_reqs * 7, - dtype=torch.int32, - device=device, - ).view(num_reqs, 7) - block_table[:, 0] = -1 - out_block_table = torch.empty( - (num_query_tokens, block_table.shape[1]), - dtype=torch.int32, - device=device, - ) - out_seq_lens = torch.empty(num_query_tokens, dtype=torch.int32, device=device) - out_query_start_loc = torch.empty(num_reqs + 1, dtype=torch.int32, device=device) - - _sm70_prepare_smallq_decode_metadata( - out_block_table, - out_seq_lens, - out_query_start_loc, - block_table, - seq_lens_tensor, - query_start_loc, - num_reqs=num_reqs, - num_query_tokens=num_query_tokens, - real_num_query_tokens=real_num_query_tokens, - ) - expected_block_table, expected_seq_lens = _reference_smallq_metadata( - block_table, - seq_lens_tensor, - query_start_loc, - num_query_tokens=num_query_tokens, - real_num_query_tokens=real_num_query_tokens, - ) - torch.accelerator.synchronize() - - assert torch.equal(out_block_table, expected_block_table) - assert torch.equal(out_seq_lens, expected_seq_lens) - assert torch.equal(out_query_start_loc, query_start_loc) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -def test_sm70_grouped_smallq_metadata_matches_individual_launches(): - if torch.cuda.get_device_capability() != (7, 0): - pytest.skip("SM70 CUDA device required") - - device = torch.device("cuda") - num_groups = 5 - num_reqs = 3 - num_query_tokens = 8 - real_num_query_tokens = 6 - block_cols_by_group = [3, 5, 4, 5, 3] - query_start_loc = torch.tensor([0, 3, 5, 6], dtype=torch.int32, device=device) - seq_lens = torch.tensor([19, 13, 9], dtype=torch.int32, device=device) - block_tables = [ - ( - torch.arange( - num_reqs * block_cols, - dtype=torch.int32, - device=device, - ) - .view(num_reqs, block_cols) - .add_(group_id * 1000) - ) - for group_id, block_cols in enumerate(block_cols_by_group) - ] - for table in block_tables: - table[:, 0] = -1 - - grouped_tables = [ - torch.empty((num_query_tokens, block_cols), dtype=torch.int32, device=device) - for block_cols in block_cols_by_group - ] - grouped_seq_lens = [ - torch.empty(num_query_tokens, dtype=torch.int32, device=device) - for _ in range(num_groups) - ] - grouped_query_locs = [ - torch.empty(num_reqs + 1, dtype=torch.int32, device=device) - for _ in range(num_groups) - ] - descriptor = _sm70_prepare_grouped_smallq_decode_metadata( - grouped_tables, - grouped_seq_lens, - grouped_query_locs, - block_tables, - seq_lens, - query_start_loc, - num_reqs=num_reqs, - num_query_tokens=num_query_tokens, - real_num_query_tokens=real_num_query_tokens, - ) - - reference_tables = [] - reference_seq_lens = [] - for table in block_tables: - out_table = torch.empty( - (num_query_tokens, table.shape[1]), dtype=torch.int32, device=device - ) - out_seq = torch.empty_like(grouped_seq_lens[0]) - out_query = torch.empty_like(grouped_query_locs[0]) - _sm70_prepare_smallq_decode_metadata( - out_table, - out_seq, - out_query, - table, - seq_lens, - query_start_loc, - num_reqs=num_reqs, - num_query_tokens=num_query_tokens, - real_num_query_tokens=real_num_query_tokens, - ) - reference_tables.append(out_table) - reference_seq_lens.append(out_seq) - assert torch.equal(out_query, query_start_loc) - torch.accelerator.synchronize() - - for grouped, reference in zip(grouped_tables, reference_tables): - assert torch.equal(grouped, reference) - for grouped, reference in zip(grouped_seq_lens, reference_seq_lens): - assert torch.equal(grouped, reference) - for grouped in grouped_query_locs: - assert torch.equal(grouped, query_start_loc) - - # Reuse the persistent pointer tables and prove that no stale group output - # survives a replay with different sequence metadata. - seq_lens.add_(7) - replay_descriptor = _sm70_prepare_grouped_smallq_decode_metadata( - grouped_tables, - grouped_seq_lens, - grouped_query_locs, - block_tables, - seq_lens, - query_start_loc, - num_reqs=num_reqs, - num_query_tokens=num_query_tokens, - real_num_query_tokens=real_num_query_tokens, - descriptor=descriptor, - ) - torch.accelerator.synchronize() - assert replay_descriptor is descriptor - _, expected_seq_lens = _reference_smallq_metadata( - block_tables[0], - seq_lens, - query_start_loc, - num_query_tokens=num_query_tokens, - real_num_query_tokens=real_num_query_tokens, - ) - for grouped in grouped_seq_lens: - assert torch.equal(grouped, expected_seq_lens) diff --git a/tests/kernels/attention/test_sm70_flash_v100_splitkv3.py b/tests/kernels/attention/test_sm70_flash_v100_splitkv3.py index 6e63720ed1..4d829aaeb9 100644 --- a/tests/kernels/attention/test_sm70_flash_v100_splitkv3.py +++ b/tests/kernels/attention/test_sm70_flash_v100_splitkv3.py @@ -37,14 +37,18 @@ 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( @@ -198,4 +202,6 @@ 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/attention/test_sm70_flash_v100_varlen_layout.py b/tests/kernels/attention/test_sm70_flash_v100_varlen_layout.py index df09e884c8..83d0cd761b 100644 --- a/tests/kernels/attention/test_sm70_flash_v100_varlen_layout.py +++ b/tests/kernels/attention/test_sm70_flash_v100_varlen_layout.py @@ -624,333 +624,3 @@ def test_sm70_flash_v100_fp8_e5m2_xqa_decode_matches_scalar( ).float() assert torch.all(delta <= torch.maximum(one_ulp, torch.full_like(delta, 1e-4))) assert float(delta.mean()) <= 1e-5 - - -@pytest.mark.parametrize(("k_scale", "v_scale"), ((1.0, 1.0), (0.04, 0.25))) -@pytest.mark.parametrize("partition_size", (64, 128, 256)) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -@torch.inference_mode() -def test_sm70_flash_v100_fp8_e4m3_g6_xqa_decode_matches_scalar( - k_scale, - v_scale, - partition_size, -): - 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(20260823) - seq_len = 2049 - block_size = 1568 - head_dim = 256 - cache_shape = (3, block_size, 1, head_dim) - key_cache = ( - torch.randn(cache_shape, dtype=torch.float16, device="cuda") - .to(torch.float8_e4m3fn) - .view(torch.uint8) - ) - value_cache = ( - torch.randn(cache_shape, dtype=torch.float16, device="cuda") - .to(torch.float8_e4m3fn) - .view(torch.uint8) - ) - query = torch.randn(1, 6, head_dim, dtype=torch.float16, device="cuda") - block_table = torch.tensor([[2, 0]], dtype=torch.int32, device="cuda") - seq_lens = torch.tensor([seq_len], dtype=torch.int32, device="cuda") - kwargs = { - "kv_cache_dtype": "fp8_e4m3", - "k_scale": k_scale, - "v_scale": v_scale, - "max_seq_len_hint": seq_len, - } - - expected = flash_attn_v100.flash_attn_decode_paged( - query, - key_cache, - value_cache, - block_table, - seq_lens, - partition_size_hint=256, - **kwargs, - ) - actual = flash_attn_v100.flash_attn_decode_paged_xqa( - query, - key_cache, - value_cache, - block_table, - seq_lens, - partition_size_hint=partition_size, - **kwargs, - ) - - torch.accelerator.synchronize() - delta = (actual.float() - expected.float()).abs() - positive = torch.full_like(expected, torch.inf) - negative = torch.full_like(expected, -torch.inf) - one_ulp = torch.maximum( - (torch.nextafter(expected, positive) - expected).abs(), - (torch.nextafter(expected, negative) - expected).abs(), - ).float() - assert torch.all(delta <= torch.maximum(one_ulp, torch.full_like(delta, 1e-4))) - assert float(delta.mean()) <= 1e-5 - - -@pytest.mark.parametrize("batch_size", (2, 16)) -@pytest.mark.parametrize("partition_size", (64, 128, 256)) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -@torch.inference_mode() -def test_sm70_flash_v100_fp8_e4m3_batched_g6_xqa_matches_scalar( - monkeypatch, - batch_size, - partition_size, -): - if torch.cuda.get_device_capability() != (7, 0): - pytest.skip("FlashAttention-V100 regression is SM70/V100 only") - - monkeypatch.setenv("VLLM_FLASH_V100_E4M3_BATCH_XQA", "1") - flash_attn_v100 = pytest.importorskip("flash_attn_v100") - torch.manual_seed(20260824) - block_size = 1568 - head_dim = 256 - blocks_per_seq = 2 - num_blocks = batch_size * blocks_per_seq - cache_shape = (num_blocks, block_size, 1, head_dim) - key_cache = ( - torch.randn(cache_shape, dtype=torch.float16, device="cuda") - .to(torch.float8_e4m3fn) - .view(torch.uint8) - ) - value_cache = ( - torch.randn(cache_shape, dtype=torch.float16, device="cuda") - .to(torch.float8_e4m3fn) - .view(torch.uint8) - ) - query = torch.randn( - batch_size, - 6, - head_dim, - dtype=torch.float16, - device="cuda", - ) - block_table = ( - torch.arange(num_blocks, dtype=torch.int32, device="cuda") - .view(batch_size, blocks_per_seq) - .flip(1) - ) - seq_lens = 1025 + (torch.arange(batch_size, device="cuda") * 61) % 900 - seq_lens = seq_lens.to(torch.int32) - kwargs = { - "kv_cache_dtype": "fp8_e4m3", - "k_scale": 0.04, - "v_scale": 0.25, - "max_seq_len_hint": int(seq_lens.max().item()), - } - - expected = flash_attn_v100.flash_attn_decode_paged( - query, - key_cache, - value_cache, - block_table, - seq_lens, - partition_size_hint=256, - **kwargs, - ) - actual = flash_attn_v100.flash_attn_decode_paged_xqa( - query, - key_cache, - value_cache, - block_table, - seq_lens, - partition_size_hint=partition_size, - **kwargs, - ) - - torch.accelerator.synchronize() - delta = (actual.float() - expected.float()).abs() - positive = torch.full_like(expected, torch.inf) - negative = torch.full_like(expected, -torch.inf) - one_ulp = torch.maximum( - (torch.nextafter(expected, positive) - expected).abs(), - (torch.nextafter(expected, negative) - expected).abs(), - ).float() - assert torch.all(delta <= torch.maximum(one_ulp, torch.full_like(delta, 1e-4))) - assert float(delta.mean()) <= 1e-5 - - -@pytest.mark.parametrize("batch_size", (8, 16)) -@pytest.mark.parametrize("partition_size", (64, 128, 256)) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -@torch.inference_mode() -def test_sm70_flash_v100_fp8_e4m3_batched_long_context_route_matches_scalar( - monkeypatch, - batch_size, - partition_size, -): - if torch.cuda.get_device_capability() != (7, 0): - pytest.skip("FlashAttention-V100 regression is SM70/V100 only") - - monkeypatch.setenv("VLLM_FLASH_V100_E4M3_BATCH_XQA", "1") - monkeypatch.delenv("VLLM_FLASH_V100_E4M3_BATCH_XQA_OPTIMIZED", raising=False) - monkeypatch.setenv("VLLM_FLASH_V100_E4M3_PAGE800_FASTPATH_TRACE", "0") - flash_attn_v100 = pytest.importorskip("flash_attn_v100") - torch.manual_seed(20260824 + batch_size) - block_size = 800 - head_dim = 256 - seq_lens_cpu = 15360 + torch.arange(batch_size, dtype=torch.int32) * 113 - max_seq_len = int(seq_lens_cpu.max().item()) - blocks_per_seq = (max_seq_len + block_size - 1) // block_size - num_blocks = batch_size * blocks_per_seq - cache_shape = (num_blocks, 2, block_size, 1, head_dim) - kv_cache = ( - torch.randn(cache_shape, dtype=torch.float16, device="cuda") - .to(torch.float8_e4m3fn) - .view(torch.uint8) - ) - key_cache, value_cache = kv_cache.unbind(1) - query = torch.randn( - batch_size, - 6, - head_dim, - dtype=torch.float16, - device="cuda", - ) - block_table = ( - torch.arange(num_blocks, dtype=torch.int32, device="cuda") - .view(batch_size, blocks_per_seq) - .flip(1) - ) - seq_lens = seq_lens_cpu.to(device="cuda") - kwargs = { - "kv_cache_dtype": "fp8_e4m3", - "k_scale": 0.04, - "v_scale": 0.25, - "max_seq_len_hint": max_seq_len, - "workspace_seq_capacity_hint": 32768, - } - - expected = flash_attn_v100.flash_attn_decode_paged( - query, - key_cache, - value_cache, - block_table, - seq_lens, - partition_size_hint=256, - **kwargs, - ) - monkeypatch.setenv("VLLM_FLASH_V100_E4M3_PAGE800_FASTPATH", "0") - fallback = flash_attn_v100.flash_attn_decode_paged_xqa( - query, - key_cache, - value_cache, - block_table, - seq_lens, - partition_size_hint=partition_size, - batch_context_routing=True, - **kwargs, - ) - monkeypatch.setenv("VLLM_FLASH_V100_E4M3_PAGE800_FASTPATH", "1") - actual = flash_attn_v100.flash_attn_decode_paged_xqa( - query, - key_cache, - value_cache, - block_table, - seq_lens, - partition_size_hint=partition_size, - batch_context_routing=True, - **kwargs, - ) - - torch.accelerator.synchronize() - assert torch.equal(actual, fallback) - delta = (actual.float() - expected.float()).abs() - positive = torch.full_like(expected, torch.inf) - negative = torch.full_like(expected, -torch.inf) - one_ulp = torch.maximum( - (torch.nextafter(expected, positive) - expected).abs(), - (torch.nextafter(expected, negative) - expected).abs(), - ).float() - assert torch.all(delta <= torch.maximum(one_ulp, torch.full_like(delta, 1e-4))) - assert float(delta.mean()) <= 1e-5 - - -@pytest.mark.parametrize( - ("seq_len", "expected_partition"), - ( - (8191, 64), - (8192, 256), - (16383, 256), - (16384, 512), - (32767, 512), - (32768, 896), - (49151, 896), - (49152, 1664), - ), -) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -@torch.inference_mode() -def test_sm70_flash_v100_fp8_e4m3_g6_routes_wave_partitions_exactly( - monkeypatch, - seq_len, - expected_partition, -): - 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") - monkeypatch.delenv("VLLM_FLASH_V100_DECODE_PARTITION_SIZE", raising=False) - monkeypatch.setenv("VLLM_FLASH_V100_XQA_E4M3_G6_P256_BEGIN", "8192") - monkeypatch.setenv("VLLM_FLASH_V100_XQA_E4M3_G6_P512_BEGIN", "16384") - monkeypatch.setenv("VLLM_FLASH_V100_XQA_E4M3_G6_P896_BEGIN", "32768") - monkeypatch.setenv("VLLM_FLASH_V100_XQA_E4M3_G6_P1664_BEGIN", "49152") - monkeypatch.setenv("VLLM_FLASH_V100_XQA_E4M3_G6_WAVE_PARTITIONS", "1") - monkeypatch.setenv("VLLM_FLASH_V100_XQA_E4M3_G6_MERGED_WAVE_LAUNCH", "1") - - torch.manual_seed(20260824) - block_size = 1568 - head_dim = 256 - num_blocks = (seq_len + block_size - 1) // block_size - cache_shape = (num_blocks, 2, block_size, 1, head_dim) - kv_cache = ( - torch.randn(cache_shape, dtype=torch.float16, device="cuda") - .to(torch.float8_e4m3fn) - .view(torch.uint8) - ) - key_cache, value_cache = kv_cache.unbind(1) - assert not key_cache.is_contiguous() - assert not value_cache.is_contiguous() - query = torch.randn(1, 6, head_dim, dtype=torch.float16, device="cuda") - block_table = torch.arange(num_blocks, dtype=torch.int32, device="cuda").view(1, -1) - seq_lens = torch.tensor([seq_len], dtype=torch.int32, device="cuda") - kwargs = { - "kv_cache_dtype": "fp8_e4m3", - "k_scale": 0.04, - "v_scale": 0.25, - "max_seq_len_hint": seq_len, - "workspace_seq_capacity_hint": 262144, - } - - monkeypatch.setenv("VLLM_FLASH_V100_XQA_E4M3_G6_P64_P256_AUTO", "0") - expected = flash_attn_v100.flash_attn_decode_paged_xqa( - query, - key_cache, - value_cache, - block_table, - seq_lens, - partition_size_hint=expected_partition, - **kwargs, - ).clone() - - monkeypatch.setenv("VLLM_FLASH_V100_XQA_E4M3_G6_P64_P256_AUTO", "1") - actual = flash_attn_v100.flash_attn_decode_paged_xqa( - query, - key_cache, - value_cache, - block_table, - seq_lens, - partition_size_hint=64, - batch_context_routing=True, - **kwargs, - ) - - torch.accelerator.synchronize() - assert torch.equal(actual, expected) diff --git a/tests/kernels/core/test_batched_weight_rms_norm.py b/tests/kernels/core/test_batched_weight_rms_norm.py deleted file mode 100644 index 790367eaad..0000000000 --- a/tests/kernels/core/test_batched_weight_rms_norm.py +++ /dev/null @@ -1,63 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for the batched-weight RMS norm kernel. - -The outermost input dimension selects the corresponding weight row. DFlash -uses this for its fused per-layer K normalization. -""" - -import pytest -import torch - -from vllm import _custom_ops as ops -from vllm.platforms import current_platform -from vllm.utils.torch_utils import set_random_seed - -pytestmark = pytest.mark.skipif( - not current_platform.is_cuda_alike(), - reason="rms_norm requires a CUDA/ROCm device", -) - - -@pytest.mark.parametrize( - "shape", - [ - (28, 17, 128), - (1, 5, 2, 128), - (28, 13, 8, 128), - (6, 3, 4, 769), - ], -) -@pytest.mark.parametrize("dtype", [torch.half, torch.bfloat16, torch.float]) -@pytest.mark.parametrize("seed", [42]) -@torch.inference_mode() -def test_rms_norm_matches_loop( - shape: tuple[int, ...], dtype: torch.dtype, seed: int -) -> None: - set_random_seed(seed) - torch.set_default_device("cuda") - - num_rows, hidden = shape[0], shape[-1] - eps = 1e-6 - x = torch.randn(*shape, dtype=dtype) * 0.1 - weight = torch.randn(num_rows, hidden, dtype=dtype) * 0.1 + 1.0 - - out_ref = torch.empty_like(x) - for i in range(x.shape[0]): - ops.rms_norm(out_ref[i], x[i], weight[i], eps) - - out = torch.empty_like(x) - ops.rms_norm(out, x, weight, eps) - torch.testing.assert_close(out, out_ref, atol=0, rtol=0) - - -@torch.inference_mode() -def test_rms_norm_validates_shapes() -> None: - torch.set_default_device("cuda") - - x = torch.randn(4, 8, 128, dtype=torch.float) - out = torch.empty_like(x) - with pytest.raises(RuntimeError): - ops.rms_norm(out, x, torch.randn(3, 128), 1e-6) - with pytest.raises(RuntimeError): - ops.rms_norm(out, x, torch.randn(4, 64), 1e-6) diff --git a/tests/kernels/core/test_sm70_dflash2_gemma_rms.py b/tests/kernels/core/test_sm70_dflash2_gemma_rms.py deleted file mode 100644 index 87c0eda3da..0000000000 --- a/tests/kernels/core/test_sm70_dflash2_gemma_rms.py +++ /dev/null @@ -1,121 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import pytest -import torch - -from vllm.model_executor.layers.layernorm import ( - GemmaRMSNorm, - _sm70_dflash2_gemma_fused_add_rms_norm, - _use_sm70_dflash2_gemma_fused_add_rms, -) - - -@pytest.mark.parametrize("num_tokens", [1, 8, 32, 137, 512]) -@pytest.mark.parametrize("weight_dtype", [torch.float16, torch.bfloat16, torch.float32]) -def test_sm70_dflash2_gemma_fused_add_rms_is_within_one_fp16_ulp( - num_tokens: int, - weight_dtype: torch.dtype, -): - if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0): - pytest.skip("SM70 CUDA device required") - - torch.manual_seed(20260823) - x = torch.randn((num_tokens, 5120), dtype=torch.float16, device="cuda").mul_(0.125) - residual = torch.randn((num_tokens, 5120), dtype=torch.float32, device="cuda").mul_( - 0.125 - ) - weight = torch.randn(5120, dtype=torch.float32, device="cuda").mul_(0.05) - weight = weight.to(weight_dtype) - - expected_normalized, expected_residual = GemmaRMSNorm._forward_static_with_residual( - weight, 1e-6, x, residual - ) - actual_normalized, actual_residual = _sm70_dflash2_gemma_fused_add_rms_norm( - x, residual, weight, 1e-6 - ) - - lower = torch.nextafter( - expected_normalized, - torch.full_like(expected_normalized, -float("inf")), - ) - upper = torch.nextafter( - expected_normalized, - torch.full_like(expected_normalized, float("inf")), - ) - assert torch.all((actual_normalized >= lower) & (actual_normalized <= upper)) - assert torch.equal(actual_residual, expected_residual) - - -def test_sm70_dflash2_gemma_fused_add_rms_graph_replay_reads_current_inputs(): - if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0): - pytest.skip("SM70 CUDA device required") - - torch.manual_seed(20260823) - x = torch.randn((8, 5120), dtype=torch.float16, device="cuda").mul_(0.125) - residual = torch.randn((8, 5120), dtype=torch.float32, device="cuda").mul_(0.125) - weight = torch.randn(5120, dtype=torch.float32, device="cuda").mul_(0.05) - _sm70_dflash2_gemma_fused_add_rms_norm(x, residual, weight, 1e-6) - torch.accelerator.synchronize() - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - actual_normalized, actual_residual = _sm70_dflash2_gemma_fused_add_rms_norm( - x, residual, weight, 1e-6 - ) - - x.copy_(torch.randn_like(x).mul_(0.125)) - residual.copy_(torch.randn_like(residual).mul_(0.125)) - weight.copy_(torch.randn_like(weight).mul_(0.05)) - graph.replay() - torch.accelerator.synchronize() - - expected_normalized, expected_residual = GemmaRMSNorm._forward_static_with_residual( - weight, 1e-6, x, residual - ) - lower = torch.nextafter( - expected_normalized, - torch.full_like(expected_normalized, -float("inf")), - ) - upper = torch.nextafter( - expected_normalized, - torch.full_like(expected_normalized, float("inf")), - ) - assert torch.all((actual_normalized >= lower) & (actual_normalized <= upper)) - assert torch.equal(actual_residual, expected_residual) - - -def test_sm70_dflash2_gemma_fused_add_rms_allows_aot_warmup_shape(monkeypatch): - if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0): - pytest.skip("SM70 CUDA device required") - - monkeypatch.setenv("VLLM_SM70_DFLASH2_FUSED_GEMMA_RMS", "1") - monkeypatch.setenv("VLLM_SM70_FLASH_V100_0DOT3_COMPILE_GRAPH", "1") - x = torch.empty((512, 5120), dtype=torch.float16, device="cuda") - residual = torch.empty_like(x, dtype=torch.float32) - weight = torch.empty(5120, dtype=torch.float16, device="cuda") - - assert _use_sm70_dflash2_gemma_fused_add_rms(x, residual, weight) - - -def test_sm70_dflash2_gemma_fused_predicate_is_fullgraph_safe(monkeypatch): - if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0): - pytest.skip("SM70 CUDA device required") - - monkeypatch.setenv("VLLM_SM70_DFLASH2_FUSED_GEMMA_RMS", "1") - monkeypatch.setenv("VLLM_SM70_FLASH_V100_0DOT3_COMPILE_GRAPH", "1") - x = torch.zeros((8, 5120), dtype=torch.float16, device="cuda") - residual = torch.empty_like(x, dtype=torch.float32) - weight = torch.empty(5120, dtype=torch.float16, device="cuda") - - @torch.compile(backend="eager", fullgraph=True) - def apply_predicate( - values: torch.Tensor, - residual_values: torch.Tensor, - norm_weight: torch.Tensor, - ) -> torch.Tensor: - if _use_sm70_dflash2_gemma_fused_add_rms(values, residual_values, norm_weight): - return values + 1 - return values - - torch.testing.assert_close(apply_predicate(x, residual, weight), x + 1) diff --git a/tests/kernels/mamba/test_mrv2_mamba_align.py b/tests/kernels/mamba/test_mrv2_mamba_align.py deleted file mode 100644 index f6bb99f974..0000000000 --- a/tests/kernels/mamba/test_mrv2_mamba_align.py +++ /dev/null @@ -1,196 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from types import SimpleNamespace - -import pytest -import torch - -from vllm.platforms import current_platform -from vllm.triton_utils import triton -from vllm.v1.worker.gpu.mamba_align import ( - preprocess_mamba_align_fused_kernel, - run_mamba_align_postprocess, - run_mamba_align_precopy, -) - -pytestmark = pytest.mark.skipif( - not current_platform.is_cuda(), reason="MRV2 Mamba align needs CUDA/Triton" -) - - -def _make_context( - conv: torch.Tensor, - temporal: torch.Tensor, - block_table: torch.Tensor, - *, - block_size: int = 4, - max_num_reqs: int | None = None, -) -> SimpleNamespace: - states = (conv, temporal) - device = block_table.device - return SimpleNamespace( - is_initialized=True, - num_layers=1, - num_state_types=2, - block_table_ptrs=torch.tensor( - [block_table.data_ptr()], dtype=torch.int64, device=device - ), - block_table_stride_req=block_table.stride(0), - state_base_addrs=torch.tensor( - [state.data_ptr() for state in states], - dtype=torch.int64, - device=device, - ), - state_block_strides=torch.tensor( - [state.stride(0) * state.element_size() for state in states], - dtype=torch.int64, - device=device, - ), - state_elem_sizes=torch.tensor( - [state.element_size() for state in states], - dtype=torch.int32, - device=device, - ), - state_inner_sizes=torch.tensor( - [conv.stride(1), temporal[0].numel()], - dtype=torch.int64, - device=device, - ), - state_conv_widths=torch.tensor( - [conv.size(1), 0], dtype=torch.int32, device=device - ), - state_group_indices=torch.zeros(2, dtype=torch.int32, device=device), - block_size=block_size, - num_accepted_tokens_out=torch.empty( - max_num_reqs or block_table.shape[0], - dtype=torch.int32, - device=device, - ), - ) - - -def test_mrv2_align_prefix_seed_uses_resolved_mamba_block_size() -> None: - device = torch.device("cuda") - # Batch order is intentionally different from persistent request-slot order. - idx_mapping = torch.tensor([1, 0], dtype=torch.int32, device=device) - state_idx = torch.full((2,), -2, dtype=torch.int32, device=device) - num_computed = torch.tensor([8, 0], dtype=torch.int32, device=device) - query_start_loc = torch.tensor([0, 1, 2], dtype=torch.int32, device=device) - num_accepted = torch.ones(2, dtype=torch.int32, device=device) - src_col = torch.empty(2, dtype=torch.int32, device=device) - token_bias = torch.empty(2, dtype=torch.int32, device=device) - - preprocess_mamba_align_fused_kernel[(1,)]( - idx_mapping, - state_idx, - num_computed, - query_start_loc, - num_accepted, - src_col, - token_bias, - 2, - BLOCK_SIZE=triton.next_power_of_2(2), - MAMBA_BLOCK_SIZE=8, - ) - torch.accelerator.synchronize() - - # req0 resumes after one complete Mamba block and crosses into column 1; - # req1 is fresh and has no source state before entering column 0. - assert src_col.cpu().tolist() == [0, -1] - assert state_idx.cpu().tolist() == [1, 0] - assert token_bias.cpu().tolist() == [0, 0] - - -@pytest.mark.parametrize("token_bias", [0, 1, 2]) -def test_mrv2_align_precopy_matches_v1_sd_semantics(token_bias: int) -> None: - torch.manual_seed(42) - device = torch.device("cuda") - num_reqs, num_cols = 4, 6 - conv_width, conv_dim = 6, 128 - num_blocks = num_reqs * num_cols + 1 - block_table = torch.arange(1, num_blocks, dtype=torch.int32, device=device).reshape( - num_reqs, num_cols - ) - conv = torch.randn( - num_blocks, conv_width, conv_dim, dtype=torch.float16, device=device - ) - temporal = torch.randn(num_blocks, 4, 16, 16, dtype=torch.float32, device=device) - conv_before = conv.clone() - temporal_before = temporal.clone() - - # State arrays use persistent request-slot order; block-table rows use - # current batch order. The non-identity mapping proves the two are not mixed. - src_col = torch.tensor([-1, 1, 1, 2], dtype=torch.int32, device=device) - dst_col = torch.tensor([0, 1, 2, 0], dtype=torch.int32, device=device) - bias = torch.tensor( - [0, 0, token_bias, token_bias], dtype=torch.int32, device=device - ) - idx_mapping = torch.tensor([2, 0, 3, 1], dtype=torch.int32, device=device) - - run_mamba_align_precopy( - _make_context(conv, temporal, block_table), - num_reqs, - dst_col, - src_col, - bias, - idx_mapping, - ) - torch.accelerator.synchronize() - - conv_ref = conv_before.clone() - temporal_ref = temporal_before.clone() - for batch_idx, req_idx in enumerate(idx_mapping.cpu().tolist()): - src = int(src_col[req_idx]) - dst = int(dst_col[req_idx]) - accepted_bias = int(bias[req_idx]) - if src < 0 or src == dst: - continue - src_block = int(block_table[batch_idx, src]) - dst_block = int(block_table[batch_idx, dst]) - temporal_block = int(block_table[batch_idx, src + accepted_bias]) - conv_ref[dst_block, : conv_width - accepted_bias] = conv_before[ - src_block, accepted_bias: - ] - temporal_ref[dst_block] = temporal_before[temporal_block] - - torch.testing.assert_close(conv, conv_ref, rtol=0, atol=0) - torch.testing.assert_close(temporal, temporal_ref, rtol=0, atol=0) - - -def test_mrv2_align_postprocess_same_block_shift_is_exact() -> None: - torch.manual_seed(43) - device = torch.device("cuda") - conv_width, conv_dim = 6, 128 - block_table = torch.tensor([[1, 2, 3, 4, 5]], dtype=torch.int32, device=device) - conv = torch.randn(6, conv_width, conv_dim, dtype=torch.float16, device=device) - temporal = torch.randn(6, 4, 16, 16, dtype=torch.float32, device=device) - conv_before = conv.clone() - temporal_before = temporal.clone() - - # accepted=3 and new_computed=8 means running=6, aligned=8, bias=2. - # src_col == dst_col == 1 exercises the overlapping conv memmove. - accepted = torch.tensor([3, 1, 1, 1], dtype=torch.int32, device=device) - state_idx = torch.tensor([1, 0, 0, 0], dtype=torch.int32, device=device) - new_computed = torch.tensor([8, 0, 0, 0], dtype=torch.int32, device=device) - idx_mapping = torch.tensor([0], dtype=torch.int32, device=device) - run_mamba_align_postprocess( - _make_context( - conv, temporal, block_table, block_size=4, max_num_reqs=accepted.numel() - ), - 1, - accepted, - state_idx, - new_computed, - idx_mapping, - ) - torch.accelerator.synchronize() - - dst_block = int(block_table[0, 1]) - temporal_src_block = int(block_table[0, 3]) - conv_ref = conv_before.clone() - temporal_ref = temporal_before.clone() - conv_ref[dst_block, : conv_width - 2] = conv_before[dst_block, 2:] - temporal_ref[dst_block] = temporal_before[temporal_src_block] - torch.testing.assert_close(conv, conv_ref, rtol=0, atol=0) - torch.testing.assert_close(temporal, temporal_ref, rtol=0, atol=0) - assert int(accepted[0]) == 1 diff --git a/tests/kernels/moe/test_moe_permute_unpermute.py b/tests/kernels/moe/test_moe_permute_unpermute.py index dc596f8795..662da0fb62 100644 --- a/tests/kernels/moe/test_moe_permute_unpermute.py +++ b/tests/kernels/moe/test_moe_permute_unpermute.py @@ -463,7 +463,6 @@ def test_sm70_single_token_weighted_reduce_matches_moe_unpermute(): ] cases = [ ("topk2_h256_contiguous", 2, 256, 0), - ("topk6_h4096_contiguous", 6, 4096, 0), ("topk8_h2048_contiguous", 8, 2048, 0), ("topk8_h7168_padded_stride", 8, 7168, 64), ] @@ -474,7 +473,9 @@ 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 ) @@ -482,8 +483,12 @@ 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/moe/test_sm70_unquantized_moe_config.py b/tests/kernels/moe/test_sm70_unquantized_moe_config.py deleted file mode 100644 index ed024de854..0000000000 --- a/tests/kernels/moe/test_sm70_unquantized_moe_config.py +++ /dev/null @@ -1,93 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import pytest - -import vllm.envs as envs -import vllm.model_executor.layers.fused_moe.fused_moe as fused_moe_module -from vllm.model_executor.layers.fused_moe.fused_moe import ( - _get_sm70_mtp_moe_decode_config, - force_sm70_mtp_moe_legacy_config, -) - - -@pytest.fixture(autouse=True) -def _enable_tuned_mtp_config(monkeypatch): - monkeypatch.setattr(envs, "VLLM_SM70_MTP_MOE_TUNED_CONFIG", True) - - -def test_mtp_sm70_decode_config_keeps_legacy_tile_at_m1(): - assert _get_sm70_mtp_moe_decode_config(1, 256, 128, 2048, 8) is None - - -@pytest.mark.parametrize("m", range(2, 17)) -def test_mtp_sm70_decode_config_uses_exact_local_tile(m): - config = _get_sm70_mtp_moe_decode_config(m, 256, 128, 2048, 8) - - assert config is not None - assert config["BLOCK_SIZE_M"] == 8 - assert config["BLOCK_SIZE_N"] == 128 - assert config["BLOCK_SIZE_K"] == 32 - - -@pytest.mark.parametrize( - "shape", - [ - (17, 256, 128, 2048, 8), - (2, 128, 128, 2048, 8), - (2, 256, 512, 2048, 8), - (2, 256, 128, 4096, 8), - (2, 256, 128, 2048, 4), - ], -) -def test_mtp_sm70_decode_config_is_shape_bounded(shape): - assert _get_sm70_mtp_moe_decode_config(*shape) is None - - -def test_mtp_sm70_decode_config_can_be_disabled(monkeypatch): - monkeypatch.setattr(envs, "VLLM_SM70_MTP_MOE_TUNED_CONFIG", False) - - assert _get_sm70_mtp_moe_decode_config(2, 256, 128, 2048, 8) is None - - -def test_mtp_sm70_decode_config_can_be_forced_to_legacy_for_warmup(): - with force_sm70_mtp_moe_legacy_config(): - config = _get_sm70_mtp_moe_decode_config(2, 256, 128, 2048, 8) - - assert config is None - assert _get_sm70_mtp_moe_decode_config(2, 256, 128, 2048, 8) is not None - - -class _FakeSM70Platform: - @staticmethod - def is_cuda(): - return True - - @staticmethod - def is_rocm(): - return False - - @staticmethod - def has_device_capability(capability): - return capability == 70 - - -def test_mtp_sm70_decode_config_is_selected_when_opted_in(monkeypatch): - monkeypatch.setattr(fused_moe_module, "current_platform", _FakeSM70Platform()) - - config = fused_moe_module.get_default_config(2, 256, 128, 2048, 8, None) - - assert config["BLOCK_SIZE_M"] == 8 - assert config["BLOCK_SIZE_N"] == 128 - assert config["BLOCK_SIZE_K"] == 32 - - -def test_mtp_sm70_decode_config_rolls_back_to_0dot3(monkeypatch): - monkeypatch.setattr(fused_moe_module, "current_platform", _FakeSM70Platform()) - monkeypatch.setattr(envs, "VLLM_SM70_MTP_MOE_TUNED_CONFIG", False) - - config = fused_moe_module.get_default_config(1, 256, 128, 2048, 8, None) - - assert config["BLOCK_SIZE_N"] == 32 - assert config["BLOCK_SIZE_K"] == 64 - assert "num_warps" not in config diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index dd7e475549..c6daab2d86 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -12,7 +12,6 @@ """ import math -from types import SimpleNamespace import pytest import torch @@ -25,11 +24,6 @@ from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import ( _fused_kv_compress_norm_rope_insert_indexer_attn, _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn, - compress_norm_rope_store_triton, -) -from vllm.models.deepseek_v4.common.ops.save_partial_states import ( - save_partial_states_to_ring, - stage_partial_states_from_ring, ) from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4 @@ -613,12 +607,6 @@ def test_fused_kv_insert_indexer(num_tokens: int, kv_block_size: int, use_fp4: b kv_cache, slot_mapping, kv_block_size, - state_cache, - state_cache.stride(0), - state_cache, - state_cache.stride(0), - rms_weight, - rms_weight.stride(0), HEAD_SIZE=HEAD_DIM, TRITON_BLOCK_SIZE=HEAD_DIM, STATE_WIDTH=coff * HEAD_DIM, @@ -630,10 +618,6 @@ def test_fused_kv_insert_indexer(num_tokens: int, kv_block_size: int, use_fp4: b TOKEN_STRIDE=TOKEN_STRIDE, SCALE_DIM=SCALE_DIM, KV_BLOCK_STRIDE=kv_cache.stride(0), - USE_PRIVATE_STATE=False, - USE_DENSE_PRIVATE_STATE=False, - RING_SIZE=1, - PRIVATE_HISTORY_SIZE=1, num_warps=1, ) @@ -683,212 +667,3 @@ def test_fused_kv_insert_indexer(num_tokens: int, kv_block_size: int, use_fp4: b assert torch.equal(actual_scale, scale[i : i + 1]), ( f"token {i}: scale {actual_scale.item()} != {scale[i].item()}" ) - - -@pytest.mark.parametrize( - ("head_dim", "compress_ratio", "use_fp4_cache"), - [(512, 4, False), (512, 128, False), (128, 4, False)], -) -def test_private_compressor_ring_matches_paged_state_bytes( - head_dim: int, compress_ratio: int, use_fp4_cache: bool -): - """The SM70 private-state gather must preserve final cache bytes exactly.""" - torch.manual_seed(19) - device = "cuda" - overlap = compress_ratio == 4 - coff = 1 + overlap - state_width = coff * head_dim - state_block_size = 4 if compress_ratio == 4 else 8 - chunk_start = 12 if compress_ratio == 4 else 120 - num_tokens = 12 if compress_ratio == 4 else 136 - total_tokens = chunk_start + num_tokens - ring_size = coff * compress_ratio + 8 - - all_kv = torch.randn(total_tokens, state_width, dtype=torch.float32, device=device) - all_score = torch.randn_like(all_kv) - ape = torch.randn(compress_ratio, state_width, dtype=torch.float32, device=device) - absolute_positions = torch.arange(total_tokens, device=device) - stored_score = all_score + ape[absolute_positions % compress_ratio] - dense_state = torch.cat((all_kv, stored_score), dim=-1) - - num_state_blocks = math.ceil(total_tokens / state_block_size) - state_cache = torch.zeros( - num_state_blocks, - state_block_size, - 2 * state_width, - dtype=torch.float32, - device=device, - ) - state_cache.view(-1, 2 * state_width)[:total_tokens].copy_(dense_state) - block_table = torch.arange( - num_state_blocks, dtype=torch.int32, device=device - ).unsqueeze(0) - - state_ring = torch.empty( - 1, ring_size, 2 * state_width, dtype=torch.float32, device=device - ) - history_start = max(0, chunk_start - ring_size) - history_positions = torch.arange(history_start, chunk_start, device=device) - state_ring[0, history_positions % ring_size] = dense_state[history_positions] - - positions = torch.arange( - chunk_start, total_tokens, dtype=torch.int64, device=device - ) - valid_slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device) - token_to_req = torch.zeros(num_tokens, dtype=torch.int32, device=device) - compressed_slot_mapping = torch.where( - (positions + 1) % compress_ratio == 0, - positions // compress_ratio, - -1, - ) - kv_block_size = 16 - num_compressed_tokens = math.ceil(total_tokens / compress_ratio) - num_kv_blocks = math.ceil(num_compressed_tokens / kv_block_size) + 1 - - if head_dim == 512: - token_stride, scale_dim, quant_block = 576, 8, 64 - elif use_fp4_cache: - token_stride, scale_dim, quant_block = 64, 4, 32 - else: - token_stride, scale_dim, quant_block = 128, 4, 128 - bytes_per_token = token_stride + scale_dim - paged_output = torch.zeros( - num_kv_blocks, - kv_block_size, - bytes_per_token, - dtype=torch.uint8, - device=device, - ) - private_output = torch.zeros_like(paged_output) - k_cache_metadata = SimpleNamespace(slot_mapping=compressed_slot_mapping) - rms_weight = torch.randn(head_dim, dtype=torch.bfloat16, device=device) - cos_sin_cache = torch.randn(total_tokens + 1, 64, device=device) - current_kv = all_kv[chunk_start:] - current_score = all_score[chunk_start:] - - common_kwargs = dict( - current_kv=current_kv, - current_score=current_score, - ape=ape, - num_actual=num_tokens, - token_to_req_indices=token_to_req, - positions=positions, - slot_mapping=valid_slot_mapping, - block_table=block_table, - block_size=state_block_size, - state_width=state_width, - cos_sin_cache=cos_sin_cache, - k_cache_metadata=k_cache_metadata, - pdl_kwargs={}, - head_dim=head_dim, - rope_head_dim=64, - compress_ratio=compress_ratio, - overlap=overlap, - use_fp4_cache=use_fp4_cache, - rms_norm_weight=rms_weight, - rms_norm_eps=1e-6, - quant_block=quant_block, - token_stride=token_stride, - scale_dim=scale_dim, - ) - compress_norm_rope_store_triton( - state_cache=state_cache, - kv_cache=paged_output, - use_private_state=False, - **common_kwargs, - ) - compress_norm_rope_store_triton( - state_cache=state_ring, - kv_cache=private_output, - use_private_state=True, - ring_size=ring_size, - **common_kwargs, - ) - - assert torch.equal(private_output, paged_output) - - if compress_ratio == 128: - dense_state = torch.empty( - compress_ratio + num_tokens, - 2 * state_width, - dtype=torch.float32, - device=device, - ) - stage_partial_states_from_ring( - kv=current_kv, - score=current_score, - ape=ape, - positions=positions, - state_ring=state_ring, - dense_state=dense_state, - slot_mapping=valid_slot_mapping, - state_width=state_width, - history_size=compress_ratio, - compress_ratio=compress_ratio, - ) - dense_output = torch.zeros_like(paged_output) - compress_norm_rope_store_triton( - state_cache=dense_state.unsqueeze(0), - kv_cache=dense_output, - use_dense_private_state=True, - private_history_size=compress_ratio, - **common_kwargs, - ) - assert torch.equal(dense_output, paged_output) - - -@pytest.mark.parametrize( - ("head_dim", "compress_ratio"), [(512, 4), (512, 128), (128, 4)] -) -def test_private_compressor_ring_keeps_only_chunk_tail( - head_dim: int, compress_ratio: int -): - overlap = compress_ratio == 4 - state_width = (1 + overlap) * head_dim - ring_size = (1 + overlap) * compress_ratio + 8 - chunk_start = 4096 - num_tokens = 8192 - positions = torch.arange( - chunk_start, - chunk_start + num_tokens, - dtype=torch.int64, - device="cuda", - ) - kv = torch.randn(num_tokens, state_width, dtype=torch.float32, device="cuda") - score = torch.randn_like(kv) - ape = torch.randn(compress_ratio, state_width, dtype=torch.float32, device="cuda") - state_ring = torch.full( - (1, ring_size, 2 * state_width), - float("nan"), - dtype=torch.float32, - device="cuda", - ) - slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device="cuda") - token_to_req = torch.zeros(num_tokens, dtype=torch.int32, device="cuda") - seq_lens = torch.tensor( - [chunk_start + num_tokens], dtype=torch.int32, device="cuda" - ) - - save_partial_states_to_ring( - kv=kv, - score=score, - ape=ape, - positions=positions, - state_ring=state_ring, - slot_mapping=slot_mapping, - token_to_req_indices=token_to_req, - seq_lens=seq_lens, - state_width=state_width, - compress_ratio=compress_ratio, - ) - - tail_positions = positions[-ring_size:] - actual = state_ring[0, tail_positions % ring_size] - expected = torch.cat( - ( - kv[-ring_size:], - score[-ring_size:] + ape[tail_positions % compress_ratio], - ), - dim=-1, - ) - assert torch.equal(actual, expected) diff --git a/tests/kernels/test_deepseek_v4_sm70_indexer.py b/tests/kernels/test_deepseek_v4_sm70_indexer.py index 2d30d4bf6a..ab98194cff 100644 --- a/tests/kernels/test_deepseek_v4_sm70_indexer.py +++ b/tests/kernels/test_deepseek_v4_sm70_indexer.py @@ -4,322 +4,33 @@ import pytest import torch -from vllm.models.deepseek_v4.sm70 import indexer as sm70_indexer -from vllm.models.deepseek_v4.sm70.indexer import ( - sm70_indexer_decode_logits, - sm70_indexer_prefill_logits, -) +from vllm.models.deepseek_v4.sm70.indexer import sm70_indexer_prefill_logits + -requires_sm70 = pytest.mark.skipif( +@pytest.mark.skipif( not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0), reason="requires NVIDIA V100/SM70", ) - -_HEAD_DIM = 128 -_BLOCK_SIZE = 64 - -# E4M3FN bit patterns with exact values. Building keys from a palette keeps the -# reference bit-exact and needs no native FP8 conversion, which SM70 lacks. -_FP8_BITS = torch.tensor([0x00, 0x30, 0x38, 0x3C, 0x40, 0xB0, 0xB8, 0xC0]) -_FP8_VALUES = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, -0.5, -1.0, -2.0]) - - -def _random_fp8_keys(num_keys: int, generator: torch.Generator): - """Return (uint8 bit patterns, exact fp32 values) for `num_keys` keys.""" - picks = torch.randint(len(_FP8_BITS), (num_keys, _HEAD_DIM), generator=generator) - return _FP8_BITS[picks].to(torch.uint8).cuda(), _FP8_VALUES[picks].cuda() - - -def _reference_index_logits( - q: torch.Tensor, weights: torch.Tensor, k: torch.Tensor -) -> torch.Tensor: - """I[t, s] = sum_h weights[t, h] * relu(q[t, h] . k[s]). - - The model's actual indexer scoring function (reference `Indexer.forward`, - inference/model.py:427; the in-repo `fp8_mqa_logits` reference at - v1/attention/ops/rocm_aiter_mla_sparse.py:510 applies the same relu). It - does NOT factor to (sum_h weights[t, h] * q[t, h]) . k[s] - the relu sits - between the weighting and the head sum. - """ - per_head = torch.einsum("mhd,nd->mhn", q.float(), k.float()) - return torch.einsum("mhn,mh->mn", torch.relu(per_head), weights.float()) - - -def _reference_factored_logits( - q: torch.Tensor, weights: torch.Tensor, k: torch.Tensor -) -> torch.Tensor: - """The relu-free form, kept reachable via VLLM_SM70_INDEXER_RELU=0.""" - weighted_q = torch.einsum("thd,th->td", q.float(), weights.float()) - return weighted_q @ k.float().T - - -def _build_paged_index_cache( - value_bits: torch.Tensor, scales: torch.Tensor -) -> torch.Tensor: - """Pack keys into the block-major layout the C++ cache kernels write. - - `indexer_k_quant_and_cache` / `cp_gather_indexer_k_quant_cache` - (csrc/libtorch_stable/cache_kernels.cu) store, per block, all `block_size` - rows of 128 FP8 values first and only then that block's FP32 scales: - - value: block_base + pos_in_block * 128 + d - scale: block_base + block_size * 128 + pos_in_block * 4 - - It is NOT a per-token [128 values][4 scale] record, which is why only - token 0 of each block used to read back correctly. - """ - num_blocks, block_size, head_dim = value_bits.shape - assert scales.shape == (num_blocks, block_size) - value_bytes = value_bits.reshape(num_blocks, block_size * head_dim) - scale_bytes = scales.contiguous().view(torch.uint8) - return ( - torch.cat([value_bytes, scale_bytes], dim=1) - .reshape(num_blocks, block_size, head_dim + 4) - .contiguous() - ) - - -def _make_queries(rows: int, num_heads: int): - q = ( - torch.randn((rows, num_heads, _HEAD_DIM), device="cuda", dtype=torch.float16) - * 0.25 - ) - weights = torch.randn((rows, num_heads), device="cuda", dtype=torch.float32) - return q, weights - - -@requires_sm70 -@pytest.mark.parametrize("use_cublas", [True, False]) -def test_prefill_matches_the_reference_scoring_function(monkeypatch, use_cublas): - torch.manual_seed(20260819) - generator = torch.Generator().manual_seed(20260819) - num_queries, num_heads, num_keys = 5, 8, 37 - monkeypatch.setattr(sm70_indexer, "_PREFILL_CUBLAS", use_cublas) - - q, weights = _make_queries(num_queries, num_heads) - bits, values = _random_fp8_keys(num_keys, generator) - scales = torch.linspace(0.5, 2.0, num_keys, device="cuda", dtype=torch.float32) - k = values * scales[:, None] - - actual = sm70_indexer_prefill_logits( - q, bits.view(torch.float8_e4m3fn), scales, weights - ) - expected = _reference_index_logits(q, weights, k) - torch.testing.assert_close(actual, expected, rtol=1e-2, atol=1e-2) - - # The relu is not a rounding detail: without it the scoring is a different - # function, which is what broke long-context top-k selection. - factored = _reference_factored_logits(q, weights, k) - assert not torch.allclose(expected, factored, rtol=1e-2, atol=1e-2) - - -@requires_sm70 -def test_prefill_relu_off_restores_the_factored_form(monkeypatch): - """VLLM_SM70_INDEXER_RELU=0 keeps the old scoring reachable for A/B.""" - torch.manual_seed(20260819) - generator = torch.Generator().manual_seed(20260819) - num_queries, num_heads, num_keys = 4, 8, 23 - monkeypatch.setattr(sm70_indexer, "_RELU_LOGITS", False) - - q, weights = _make_queries(num_queries, num_heads) - bits, values = _random_fp8_keys(num_keys, generator) - scales = torch.linspace(0.5, 2.0, num_keys, device="cuda", dtype=torch.float32) - - actual = sm70_indexer_prefill_logits( - q, bits.view(torch.float8_e4m3fn), scales, weights - ) - expected = _reference_factored_logits(q, weights, values * scales[:, None]) - torch.testing.assert_close(actual, expected, rtol=1e-2, atol=1e-2) - - -@requires_sm70 -@pytest.mark.parametrize( - "relu,fused", - [ - pytest.param(True, True, id="relu"), - pytest.param(False, True, id="factored-fused"), - pytest.param(False, False, id="factored-gather"), - ], -) -def test_decode_reads_the_block_major_paged_cache(monkeypatch, relu, fused): - """Every decode path must address values and scales the way the cache is - written. Reading a token as a [128 values][4 scale] record lands on the - wrong bytes for every token but the first of each block, and reinterprets - four FP8 value bytes as the FP32 scale. - """ - torch.manual_seed(20260819) - generator = torch.Generator().manual_seed(20260819) - monkeypatch.setattr(sm70_indexer, "_RELU_LOGITS", relu) - monkeypatch.setattr(sm70_indexer, "_FUSED_DECODE_LOGITS", fused) - - num_heads = 8 - seq_lens_list = [1, _BLOCK_SIZE + 1, 3 * _BLOCK_SIZE - 5] - rows = len(seq_lens_list) - max_seq_len = max(seq_lens_list) - blocks_per_row = (max_seq_len + _BLOCK_SIZE - 1) // _BLOCK_SIZE - num_blocks = rows * blocks_per_row + 2 # +2 so no row can start at block 0 - - bits, values = _random_fp8_keys(num_blocks * _BLOCK_SIZE, generator) - value_bits = bits.reshape(num_blocks, _BLOCK_SIZE, _HEAD_DIM) - block_values = values.reshape(num_blocks, _BLOCK_SIZE, _HEAD_DIM) - # Distinct per-token scales inside every block: with one scale per block - # the layout bug is invisible. - scales = torch.rand((num_blocks, _BLOCK_SIZE), generator=generator).cuda() + 0.5 - cache = _build_paged_index_cache(value_bits, scales) - - # Shuffled, non-identity mapping so a wrong block stride cannot pass. - perm = torch.randperm(num_blocks, generator=generator)[: rows * blocks_per_row] - block_table = perm.to(torch.int32).cuda().reshape(rows, blocks_per_row) - seq_lens = torch.tensor(seq_lens_list, device="cuda", dtype=torch.int32) - q, weights = _make_queries(rows, num_heads) - - actual = sm70_indexer_decode_logits( - q, cache, weights, seq_lens, block_table, max_seq_len - ) - - reference = _reference_index_logits if relu else _reference_factored_logits - for row, seq_len in enumerate(seq_lens_list): - positions = torch.arange(seq_len, device="cuda") - physical = block_table[row][positions // _BLOCK_SIZE].long() - pos_in_block = positions % _BLOCK_SIZE - k = ( - block_values[physical, pos_in_block] - * scales[physical, pos_in_block][:, None] - ) - expected = reference(q[row : row + 1], weights[row : row + 1], k) - torch.testing.assert_close( - actual[row, :seq_len].unsqueeze(0), expected, rtol=1e-2, atol=1e-2 - ) - - -@requires_sm70 -def test_decode_cublas_keeps_the_fp32_topk_set_and_masks_graph_tail(monkeypatch): - torch.manual_seed(20260824) - generator = torch.Generator().manual_seed(20260824) - num_rows, num_heads = 8, 64 - live_seq_len = 1017 - graph_width = 2048 - blocks = (live_seq_len + _BLOCK_SIZE - 1) // _BLOCK_SIZE - - bits, _ = _random_fp8_keys(blocks * _BLOCK_SIZE, generator) - value_bits = bits.reshape(blocks, _BLOCK_SIZE, _HEAD_DIM) - scales = torch.rand((blocks, _BLOCK_SIZE), generator=generator).cuda() + 0.5 - cache = _build_paged_index_cache(value_bits, scales) - block_table = torch.full( - (1, graph_width // _BLOCK_SIZE), - -1, - dtype=torch.int32, - device="cuda", - ) - block_table[0, :blocks] = torch.randperm(blocks, generator=generator).to( - device="cuda", dtype=torch.int32 - ) - seq_lens = torch.arange( - live_seq_len - num_rows + 1, - live_seq_len + 1, - dtype=torch.int32, - device="cuda", - ).view(1, num_rows) - q, weights = _make_queries(num_rows, num_heads) - - monkeypatch.setattr(sm70_indexer, "_DECODE_CUBLAS", False) - baseline = sm70_indexer_decode_logits( - q, cache, weights, seq_lens, block_table, graph_width - ) - - static_workspace = ( - torch.empty((graph_width, _HEAD_DIM), dtype=torch.float16, device="cuda"), - torch.empty((graph_width,), dtype=torch.float32, device="cuda"), - torch.empty( - (num_rows * num_heads, graph_width), - dtype=torch.float32, - device="cuda", - ), - ) - - class StaticWorkspace: - @staticmethod - def get_simultaneous(*specs): - assert len(specs) == len(static_workspace) - return static_workspace - - monkeypatch.setattr(sm70_indexer, "_DECODE_CUBLAS", True) - monkeypatch.setattr(sm70_indexer, "_DECODE_CUBLAS_MIN_KEYS", 1) - monkeypatch.setattr( - sm70_indexer, "current_workspace_manager", lambda: StaticWorkspace() - ) - - def candidate_call(): - return sm70_indexer_decode_logits( - q, cache, weights, seq_lens, block_table, graph_width - ) - - capture_stream = torch.cuda.Stream() - capture_stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(capture_stream): - for _ in range(3): - candidate_call() - capture_stream.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph, stream=capture_stream): - candidate = candidate_call() - graph.replay() - torch.accelerator.synchronize() - - for row, row_len in enumerate(seq_lens.reshape(-1).tolist()): - expected_topk = torch.topk(baseline[row, :row_len], 512).indices.sort().values - actual_topk = torch.topk(candidate[row, :row_len], 512).indices.sort().values - assert torch.equal(actual_topk, expected_topk) - - # The graph bucket has no allocated blocks after the live request tail. - # Reaching this assertion proves the gather masked those -1 entries. - assert torch.isfinite(candidate[:, :live_seq_len]).all() - - # Replay the same graph with a shorter live request. This exercises the - # device-derived bound instead of merely validating the capture-time row. - seq_lens.sub_(127) - monkeypatch.setattr(sm70_indexer, "_DECODE_CUBLAS", False) - replay_baseline = sm70_indexer_decode_logits( - q, cache, weights, seq_lens, block_table, graph_width - ) - graph.replay() - torch.accelerator.synchronize() - for row, row_len in enumerate(seq_lens.reshape(-1).tolist()): - expected_topk = ( - torch.topk(replay_baseline[row, :row_len], 512).indices.sort().values - ) - actual_topk = torch.topk(candidate[row, :row_len], 512).indices.sort().values - assert torch.equal(actual_topk, expected_topk) - - -@requires_sm70 -def test_sm70_indexer_decode_valid_logits_do_not_depend_on_workspace_width() -> None: - torch.manual_seed(20260804) - num_queries = 8 +def test_sm70_indexer_prefill_dequantizes_fp8_storage_as_uint8() -> None: + torch.manual_seed(20260803) + num_queries = 3 num_heads = 4 num_keys = 7 head_dim = 128 - cache_block_size = 64 q = torch.randn( (num_queries, num_heads, head_dim), device="cuda", dtype=torch.float16 ) - weights = torch.randn((num_queries, num_heads), device="cuda", dtype=torch.float16) - cache = torch.zeros( - (1, cache_block_size, head_dim + 4), device="cuda", dtype=torch.uint8 - ) - cache[..., :head_dim].fill_(0x38) - cache[..., head_dim:].view(torch.float32).fill_(1.0) - seq_lens = torch.full((1, num_queries), num_keys, device="cuda", dtype=torch.int32) - block_table = torch.zeros((1, 1), device="cuda", dtype=torch.int32) + weights = torch.randn((num_queries, num_heads), device="cuda", dtype=torch.float32) + # E4M3FN bit pattern 0x38 is exactly 1.0. Constructing the tensor through + # uint8 also avoids requiring native FP8 conversion support on SM70. + k_bits = torch.full((num_keys, head_dim), 0x38, device="cuda", dtype=torch.uint8) + k_quant = k_bits.view(torch.float8_e4m3fn) + scales = torch.linspace(0.5, 2.0, num_keys, device="cuda", dtype=torch.float32) - narrow = sm70_indexer_decode_logits( - q, cache, weights, seq_lens, block_table, max_seq_len=16 - ) - wide = sm70_indexer_decode_logits( - q, cache, weights, seq_lens, block_table, max_seq_len=64 - ) + actual = sm70_indexer_prefill_logits(q, k_quant, scales, weights) - assert narrow.shape == (num_queries, 16) - assert wide.shape == (num_queries, 64) - torch.testing.assert_close(narrow[:, :num_keys], wide[:, :num_keys], rtol=0, atol=0) + weighted_q = torch.sum(q.float() * weights[:, :, None], dim=1).half() + k_reference = scales[:, None].expand(num_keys, head_dim).half() + expected = torch.mm(weighted_q, k_reference.T, out_dtype=torch.float32) + torch.testing.assert_close(actual, expected, rtol=2e-3, atol=2e-3) 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 3cb8dc13ba..8dff20c955 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 @@ -25,7 +25,10 @@ def _make_cos_sin_cache(max_position: int, device: str) -> torch.Tensor: inv_freq = 1.0 / ( 10000.0 - ** (torch.arange(0, ROPE_DIM, 2, dtype=torch.float32, device=device) / ROPE_DIM) + ** ( + torch.arange(0, ROPE_DIM, 2, dtype=torch.float32, device=device) + / ROPE_DIM + ) ) positions = torch.arange(max_position, dtype=torch.float32, device=device) frequencies = torch.outer(positions, inv_freq) @@ -117,45 +120,3 @@ def test_sm70_kv_rope_insert_is_repeatable_across_physical_blocks(): torch.testing.assert_close(gathered, gathered_outputs[0], rtol=0, atol=0) for q_out in q_outputs[1:]: torch.testing.assert_close(q_out, q_outputs[0], rtol=0, atol=0) - - -def test_sm70_tp4_fused_qnorm_kv_insert_matches_legacy( - monkeypatch: pytest.MonkeyPatch, -): - torch.manual_seed(20260826) - device = "cuda" - num_heads = 16 - block_size = 256 - block_stride = round_up(block_size * HEAD_BYTES, TOKEN_DATA_BYTES) - storage_shape = (3 * block_stride,) - q_input = torch.randn(1, num_heads, HEAD_DIM, dtype=torch.float16, device=device) - kv = torch.randn(1, HEAD_DIM, dtype=torch.float16, device=device) - positions = torch.tensor([1024], dtype=torch.int64, device=device) - slot_mapping = torch.tensor([block_size + 3], dtype=torch.int64, device=device) - cos_sin_cache = _make_cos_sin_cache(2048, device) - - def run(enabled: bool) -> tuple[torch.Tensor, torch.Tensor]: - monkeypatch.setenv("VLLM_SM70_DSV4_QNORM_KV_FUSED_TP4", "1" if enabled else "0") - q = q_input.clone() - storage = torch.full(storage_shape, 0xA5, dtype=torch.uint8, device=device) - cache = storage.as_strided( - (3, block_size, HEAD_BYTES), - (block_stride, HEAD_BYTES, 1), - ) - sm70_qnorm_rope_kv_fp8_insert( - q, - kv, - cache, - slot_mapping, - positions, - cos_sin_cache, - eps=1e-6, - block_size=block_size, - ) - torch.accelerator.synchronize() - return q, storage - - legacy_q, legacy_storage = run(False) - fused_q, fused_storage = run(True) - torch.testing.assert_close(fused_q, legacy_q, rtol=0, atol=0) - torch.testing.assert_close(fused_storage, legacy_storage, rtol=0, atol=0) diff --git a/tests/kernels/test_fused_recurrent_packed_decode.py b/tests/kernels/test_fused_recurrent_packed_decode.py index 10e6e6af42..ce9a687606 100644 --- a/tests/kernels/test_fused_recurrent_packed_decode.py +++ b/tests/kernels/test_fused_recurrent_packed_decode.py @@ -151,7 +151,9 @@ 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) @@ -159,7 +161,9 @@ 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( @@ -179,7 +183,9 @@ 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_fused_sigmoid_gating_delta_rule.py b/tests/kernels/test_fused_sigmoid_gating_delta_rule.py index 2ac073b410..2b03e83c30 100644 --- a/tests/kernels/test_fused_sigmoid_gating_delta_rule.py +++ b/tests/kernels/test_fused_sigmoid_gating_delta_rule.py @@ -7,12 +7,7 @@ from vllm.model_executor.layers.fla.ops import ( fused_recurrent_gated_delta_rule, - fused_recurrent_gated_delta_rule_packed_decode, fused_sigmoid_gating_delta_rule_update, - fused_sigmoid_gating_delta_rule_update_mixed_qkv_out, -) -from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( - fused_gdn_gating, ) from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed @@ -20,149 +15,6 @@ DEVICE = current_platform.device_type -def _bitwise_diff_stats( - actual: torch.Tensor, - expected: torch.Tensor, -) -> tuple[int, float, float]: - diff = actual.float() - expected.float() - return ( - torch.count_nonzero(actual != expected).item(), - diff.abs().max().item(), - diff.abs().mean().item(), - ) - - -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_split_verify_matches_fused_decode_beta_semantics( - dtype: torch.dtype, -) -> None: - torch.set_default_device(DEVICE) - if dtype is torch.bfloat16 and not torch.cuda.is_bf16_supported(): - pytest.skip("BF16 requires a supported accelerator") - set_random_seed(0) - - # Qwen3.8-27B TP4 per-rank head geometry (16 QK heads, 48 V heads). - num_qk_heads = 4 - num_v_heads = 12 - head_dim = 128 - scale = head_dim**-0.5 - # Make the decay exactly zero in both kernels so this regression isolates - # beta materialization instead of also measuring unrelated fused/split - # decay rounding. softplus(-100) rounds to zero in the kernel. - A_log = torch.zeros(num_v_heads, dtype=torch.float32) - dt_bias = torch.zeros(num_v_heads, dtype=dtype) - a = torch.full((1, num_v_heads), -100, dtype=dtype) - # sigmoid(-0.5) is not exactly representable in FP16 or BF16, exposing - # whether the split verifier keeps the fused decode FP32 beta semantics. - b = torch.full((1, num_v_heads), -0.5, dtype=dtype) - query = torch.rand(1, 1, num_qk_heads, head_dim, dtype=dtype) - key = torch.rand_like(query) - value = torch.rand(1, 1, num_v_heads, head_dim, dtype=dtype) - # A zero state also removes decay and dot-product reduction-order effects - # from this one-step comparison. The state transition then differs only - # through the materialized beta value. - initial_state = torch.zeros( - 1, - num_v_heads, - head_dim, - head_dim, - dtype=torch.float32, - ) - state_indices = torch.zeros(1, dtype=torch.int32) - cu_seqlens = torch.tensor([0, 1], dtype=torch.int32) - - fused_state = initial_state.clone() - fused_output, _ = fused_sigmoid_gating_delta_rule_update( - A_log=A_log, - a=a, - b=b, - dt_bias=dt_bias, - q=query, - k=key, - v=value, - scale=scale, - initial_state=fused_state, - inplace_final_state=True, - cu_seqlens=cu_seqlens, - ssm_state_indices=state_indices, - use_qk_l2norm_in_kernel=False, - ) - - packed_state = initial_state.clone() - packed_output = torch.empty_like(fused_output) - mixed_qkv = torch.cat( - [ - query.squeeze(0).reshape(1, -1), - key.squeeze(0).reshape(1, -1), - value.squeeze(0).reshape(1, -1), - ], - dim=-1, - ).contiguous() - fused_recurrent_gated_delta_rule_packed_decode( - mixed_qkv=mixed_qkv, - a=a, - b=b, - A_log=A_log, - dt_bias=dt_bias, - scale=scale, - initial_state=packed_state, - out=packed_output, - ssm_state_indices=state_indices, - use_qk_l2norm_in_kernel=False, - ) - - def run_split( - beta_dtype: torch.dtype | None, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - g, beta = fused_gdn_gating( - A_log, - a, - b, - dt_bias, - beta_dtype=beta_dtype, - ) - state = initial_state.clone() - output, _ = fused_recurrent_gated_delta_rule( - q=query, - k=key, - v=value, - g=g, - beta=beta, - scale=scale, - initial_state=state, - inplace_final_state=True, - ssm_state_indices=state_indices, - cu_seqlens=cu_seqlens, - use_qk_l2norm_in_kernel=False, - ) - return output, state, beta - - fixed_output, fixed_state, fixed_beta = run_split(torch.float32) - legacy_output, legacy_state, legacy_beta = run_split(None) - - assert fixed_beta.dtype is torch.float32 - assert legacy_beta.dtype is dtype - fixed_output_diff = _bitwise_diff_stats(fixed_output, fused_output) - fixed_state_diff = _bitwise_diff_stats(fixed_state, fused_state) - legacy_output_diff = _bitwise_diff_stats(legacy_output, fused_output) - legacy_state_diff = _bitwise_diff_stats(legacy_state, fused_state) - assert torch.equal(packed_output, fused_output), _bitwise_diff_stats( - packed_output, fused_output - ) - assert torch.equal(packed_state, fused_state), _bitwise_diff_stats( - packed_state, fused_state - ) - assert torch.equal(fixed_output, fused_output), fixed_output_diff - assert torch.equal(fixed_state, fused_state), ( - "fixed", - fixed_state_diff, - "legacy", - legacy_state_diff, - ) - assert legacy_state_diff[0] > 0, legacy_state_diff - assert legacy_output_diff[0] > 0, legacy_output_diff - - @pytest.mark.parametrize("tp_size", [1]) @pytest.mark.parametrize("num_reqs", [1, 2, 4]) @pytest.mark.parametrize("num_k_heads", [16]) @@ -342,176 +194,3 @@ def test_fused_sigmoid_gating_delta_rule_update_spec( torch.testing.assert_close( last_recurrent_state, last_recurrent_state_ref, atol=1e-2, rtol=1e-2 ) - - -@pytest.mark.parametrize("num_reqs", [1, 2]) -@pytest.mark.parametrize("num_speculative_tokens", [3, 7]) -@pytest.mark.parametrize("state_dtype", [torch.float16, torch.float32]) -def test_dflash2_packed_verify_matches_split_fp16_contract( - num_reqs: int, - num_speculative_tokens: int, - state_dtype: torch.dtype, -) -> None: - if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0): - pytest.skip("the DFlash2 verifier fast path is SM70-only") - - torch.set_default_device(DEVICE) - set_random_seed(11) - dtype = torch.float16 - # Production GDN verifier shape used by the recorded TP4 parity evidence. - num_k_heads, num_v_heads = 4, 12 - head_k_dim = head_v_dim = 128 - tokens_per_req = num_speculative_tokens + 1 - num_tokens = num_reqs * tokens_per_req - qkv_width = 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim - - mixed_qkv = torch.randn(num_tokens, qkv_width, dtype=dtype) - query, key, value = torch.split( - mixed_qkv, - [ - num_k_heads * head_k_dim, - num_k_heads * head_k_dim, - num_v_heads * head_v_dim, - ], - dim=-1, - ) - query = query.contiguous().view(1, num_tokens, num_k_heads, head_k_dim) - key = key.contiguous().view(1, num_tokens, num_k_heads, head_k_dim) - value = value.contiguous().view(1, num_tokens, num_v_heads, head_v_dim) - a = torch.randn(num_tokens, num_v_heads, dtype=dtype) - b = torch.randn(num_tokens, num_v_heads, dtype=dtype) - A_log = torch.randn(num_v_heads, dtype=torch.float32) - dt_bias = torch.randn(num_v_heads, dtype=dtype) - state_indices = torch.arange(num_tokens, dtype=torch.int32).view( - num_reqs, tokens_per_req - ) - num_accepted_tokens = torch.randint( - 1, tokens_per_req + 1, (num_reqs,), dtype=torch.int32 - ) - cu_seqlens = torch.arange(0, num_tokens + 1, tokens_per_req, dtype=torch.int32) - initial_state = torch.randn( - num_tokens + 3, - num_v_heads, - head_v_dim, - head_k_dim, - dtype=state_dtype, - ) - - g, beta = fused_gdn_gating(A_log, a, b, dt_bias) - reference_state = initial_state.clone() - reference_out, _ = fused_recurrent_gated_delta_rule( - q=query, - k=key, - v=value, - g=g, - beta=beta, - initial_state=reference_state, - inplace_final_state=True, - cu_seqlens=cu_seqlens, - ssm_state_indices=state_indices, - num_accepted_tokens=num_accepted_tokens, - use_qk_l2norm_in_kernel=True, - ) - - fused_state = initial_state.clone() - fused_out = torch.empty(num_tokens, 1, num_v_heads, head_v_dim, dtype=dtype) - fused_sigmoid_gating_delta_rule_update_mixed_qkv_out( - A_log=A_log, - a=a, - b=b, - dt_bias=dt_bias, - mixed_qkv=mixed_qkv, - num_q_heads=num_k_heads, - num_v_heads=num_v_heads, - head_k_dim=head_k_dim, - head_v_dim=head_v_dim, - out=fused_out, - initial_state=fused_state, - cu_seqlens=cu_seqlens, - ssm_state_indices=state_indices, - num_accepted_tokens=num_accepted_tokens, - use_qk_l2norm_in_kernel=True, - precomputed_g=g, - precomputed_beta=beta, - match_recurrent_schedule=True, - match_recurrent_numerics=True, - ) - torch.accelerator.synchronize() - - torch.testing.assert_close(fused_out.transpose(0, 1), reference_out, rtol=0, atol=0) - torch.testing.assert_close(fused_state, reference_state, rtol=0, atol=0) - - -def test_dflash2_quantized_state_reload_matches_repeated_decode() -> None: - if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0): - pytest.skip("the DFlash2 verifier fast path is SM70-only") - - torch.set_default_device(DEVICE) - set_random_seed(12) - dtype = torch.float16 - num_tokens = 4 - num_k_heads, num_v_heads = 2, 4 - head_k_dim = head_v_dim = 16 - qkv_width = 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim - mixed_qkv = torch.randn(num_tokens, qkv_width, dtype=dtype) - a = torch.randn(num_tokens, num_v_heads, dtype=dtype) - b = torch.randn(num_tokens, num_v_heads, dtype=dtype) - A_log = torch.randn(num_v_heads, dtype=torch.float32) - dt_bias = torch.randn(num_v_heads, dtype=dtype) - initial_state = torch.randn(2, num_v_heads, head_v_dim, head_k_dim, dtype=dtype) - shared_state_indices = torch.zeros((1, num_tokens), dtype=torch.int32) - accepted = torch.ones(1, dtype=torch.int32) - - fused_state = initial_state.clone() - fused_out = torch.empty(num_tokens, 1, num_v_heads, head_v_dim, dtype=dtype) - fused_sigmoid_gating_delta_rule_update_mixed_qkv_out( - A_log=A_log, - a=a, - b=b, - dt_bias=dt_bias, - mixed_qkv=mixed_qkv, - num_q_heads=num_k_heads, - num_v_heads=num_v_heads, - head_k_dim=head_k_dim, - head_v_dim=head_v_dim, - out=fused_out, - initial_state=fused_state, - cu_seqlens=torch.tensor([0, num_tokens], dtype=torch.int32), - ssm_state_indices=shared_state_indices, - num_accepted_tokens=accepted, - quantize_beta=True, - quantize_state_each_step=True, - match_recurrent_schedule=True, - use_qk_l2norm_in_kernel=True, - ) - - sequential_state = initial_state.clone() - sequential_outputs = [] - one_token_cu_seqlens = torch.tensor([0, 1], dtype=torch.int32) - one_state_index = torch.zeros(1, dtype=torch.int32) - for token_idx in range(num_tokens): - token_out = torch.empty(1, 1, num_v_heads, head_v_dim, dtype=dtype) - fused_sigmoid_gating_delta_rule_update_mixed_qkv_out( - A_log=A_log, - a=a[token_idx : token_idx + 1], - b=b[token_idx : token_idx + 1], - dt_bias=dt_bias, - mixed_qkv=mixed_qkv[token_idx : token_idx + 1], - num_q_heads=num_k_heads, - num_v_heads=num_v_heads, - head_k_dim=head_k_dim, - head_v_dim=head_v_dim, - out=token_out, - initial_state=sequential_state, - cu_seqlens=one_token_cu_seqlens, - ssm_state_indices=one_state_index, - num_accepted_tokens=accepted, - quantize_beta=True, - match_recurrent_schedule=True, - use_qk_l2norm_in_kernel=True, - ) - sequential_outputs.append(token_out) - torch.accelerator.synchronize() - - assert torch.equal(fused_out, torch.cat(sequential_outputs, dim=0)) - assert torch.equal(fused_state, sequential_state) diff --git a/tests/kernels/test_mhc_sm70_fp16.py b/tests/kernels/test_mhc_sm70_fp16.py index a0086f77e7..75b5ae7242 100644 --- a/tests/kernels/test_mhc_sm70_fp16.py +++ b/tests/kernels/test_mhc_sm70_fp16.py @@ -154,10 +154,7 @@ def test_mhc_sm70_fp16_block_m_prenorm_keeps_rows_independent() -> None: ), reason="NVIDIA V100/SM70 and TileLang required", ) -@pytest.mark.parametrize("num_tokens", [1, 4, 7, 8]) -def test_mhc_sm70_fp32_stage_matches_fused_decode_bitwise( - num_tokens: int, -) -> None: +def test_mhc_sm70_fp32_stage_matches_fused_decode_bitwise() -> None: from vllm.model_executor.kernels.mhc.tilelang_kernels import ( mhc_fused_tilelang, sm70_mhc_dot_from_fp32_stage_tilelang, @@ -168,26 +165,21 @@ def test_mhc_sm70_fp32_stage_matches_fused_decode_bitwise( hidden_size = 4096 hc_mult = 4 hc_out = 24 - tile_n = 2 if num_tokens < 8 else 3 - n_splits = 8 if num_tokens < 8 else 4 + n_splits = 8 device = mhc_tilelang.current_platform.device_type - x = torch.randn((num_tokens, hidden_size), device=device, dtype=torch.float16) + x = torch.randn((1, hidden_size), device=device, dtype=torch.float16) residual = torch.randn( - (num_tokens, hc_mult, hidden_size), device=device, dtype=torch.float16 - ) - post_mix = torch.randn((num_tokens, hc_mult), device=device, dtype=torch.float32) - comb_mix = torch.randn( - (num_tokens, hc_mult, hc_mult), device=device, dtype=torch.float32 + (1, hc_mult, hidden_size), device=device, dtype=torch.float16 ) + post_mix = torch.randn((1, hc_mult), device=device, dtype=torch.float32) + comb_mix = torch.randn((1, hc_mult, hc_mult), device=device, dtype=torch.float32) fn = torch.randn((hc_out, hc_mult, hidden_size), device=device, dtype=torch.float32) baseline_gemm = torch.empty( - (n_splits, num_tokens, hc_out), device=device, dtype=torch.float32 - ) - baseline_sqrsum = torch.empty( - (n_splits, num_tokens), device=device, dtype=torch.float32 + (n_splits, 1, hc_out), device=device, dtype=torch.float32 ) + baseline_sqrsum = torch.empty((n_splits, 1), device=device, dtype=torch.float32) baseline_residual = torch.empty_like(residual) candidate_gemm = torch.empty_like(baseline_gemm) candidate_sqrsum = torch.empty_like(baseline_sqrsum) @@ -206,7 +198,7 @@ def test_mhc_sm70_fp32_stage_matches_fused_decode_bitwise( hc_mult, hidden_size, hc_out, - tile_n=tile_n, + tile_n=2, n_splits=n_splits, use_fp16=True, ) @@ -229,10 +221,10 @@ def test_mhc_sm70_fp32_stage_matches_fused_decode_bitwise( hidden_size, hc_mult, hc_out, - tile_n=tile_n, + tile_n=2, n_splits=n_splits, ) - torch.accelerator.synchronize() + torch.cuda.synchronize() torch.testing.assert_close(candidate_residual, baseline_residual, rtol=0, atol=0) torch.testing.assert_close(candidate_gemm, baseline_gemm, rtol=0, atol=0) diff --git a/tests/kernels/test_sm70_deepseek_v4_fp16_gemv.py b/tests/kernels/test_sm70_deepseek_v4_fp16_gemv.py index f3ed2c5669..b5297317c0 100644 --- a/tests/kernels/test_sm70_deepseek_v4_fp16_gemv.py +++ b/tests/kernels/test_sm70_deepseek_v4_fp16_gemv.py @@ -8,9 +8,6 @@ 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, ) @@ -23,79 +20,13 @@ def reset_env_cache() -> Iterator[None]: envs.disable_envs_cache() -def test_sm70_dsv4_fp13_gemv_is_default_on_with_rollback(monkeypatch) -> None: +def test_sm70_dsv4_fp16_gemv_is_default_off(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", @@ -139,54 +70,3 @@ 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/kernels/test_sm70_qwen35_gdn_split.py b/tests/kernels/test_sm70_qwen35_gdn_split.py deleted file mode 100644 index e8a65c0002..0000000000 --- a/tests/kernels/test_sm70_qwen35_gdn_split.py +++ /dev/null @@ -1,127 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import pytest -import torch - -from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( - _sm70_pack_qwen_gdn_qkv, -) -from vllm.model_executor.models.qwen3_5 import ( - _sm70_materialize_qwen35_gdn_splits, -) - - -@pytest.mark.parametrize("num_rows", [1, 8, 32]) -def test_qwen35_gdn_split_materialization_is_bitwise_exact(num_rows: int): - if not torch.cuda.is_available(): - pytest.skip("CUDA is required for the Qwen3.5 GDN split kernel") - - torch.manual_seed(7) - qkv_size = 2560 - z_size = 1536 - ba_size = 12 - - # Slice wider allocations so the source row strides are deliberately - # larger than their logical widths, matching the projection-view contract. - qkvz_storage = torch.randn( - (num_rows, qkv_size + z_size + 7), - dtype=torch.float16, - device="cuda", - ) - ba_storage = torch.randn( - (num_rows, 2 * ba_size + 5), - dtype=torch.float16, - device="cuda", - ) - mixed_qkvz = qkvz_storage[:, : qkv_size + z_size] - mixed_ba = ba_storage[:, : 2 * ba_size] - - actual_z, actual_b, actual_a = _sm70_materialize_qwen35_gdn_splits( - mixed_qkvz, - mixed_ba, - qkv_size, - z_size, - ba_size, - ) - - expected_z = mixed_qkvz[:, qkv_size : qkv_size + z_size].contiguous() - expected_b = mixed_ba[:, :ba_size].contiguous() - expected_a = mixed_ba[:, ba_size : 2 * ba_size].contiguous() - assert torch.equal(actual_z, expected_z) - assert torch.equal(actual_b, expected_b) - assert torch.equal(actual_a, expected_a) - - -def test_qwen35_gdn_split_graph_replay_reads_current_projection_values(): - if not torch.cuda.is_available(): - pytest.skip("CUDA is required for the Qwen3.5 GDN split kernel") - - num_rows, qkv_size, z_size, ba_size = 8, 2560, 1536, 12 - mixed_qkvz = torch.randn( - (num_rows, qkv_size + z_size), dtype=torch.float16, device="cuda" - ) - mixed_ba = torch.randn((num_rows, 2 * ba_size), dtype=torch.float16, device="cuda") - _sm70_materialize_qwen35_gdn_splits(mixed_qkvz, mixed_ba, qkv_size, z_size, ba_size) - torch.accelerator.synchronize() - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - actual_z, actual_b, actual_a = _sm70_materialize_qwen35_gdn_splits( - mixed_qkvz, mixed_ba, qkv_size, z_size, ba_size - ) - - mixed_qkvz.copy_(torch.randn_like(mixed_qkvz)) - mixed_ba.copy_(torch.randn_like(mixed_ba)) - graph.replay() - torch.accelerator.synchronize() - - assert torch.equal(actual_z, mixed_qkvz[:, qkv_size:].contiguous()) - assert torch.equal(actual_b, mixed_ba[:, :ba_size].contiguous()) - assert torch.equal(actual_a, mixed_ba[:, ba_size:].contiguous()) - - -@pytest.mark.parametrize("num_rows", [1, 8]) -def test_qwen35_gdn_qkv_pack_is_bitwise_exact(num_rows: int): - if not torch.cuda.is_available(): - pytest.skip("CUDA is required for the Qwen3.5 GDN Q/K/V pack kernel") - - torch.manual_seed(11) - q_dim, k_dim, v_dim = 512, 512, 1536 - storage = torch.randn( - (num_rows, q_dim + k_dim + v_dim + 13), - dtype=torch.float16, - device="cuda", - ) - mixed_qkv = storage[:, : q_dim + k_dim + v_dim] - actual = _sm70_pack_qwen_gdn_qkv(mixed_qkv, q_dim, k_dim, v_dim) - - q, k, v = torch.split(mixed_qkv, [q_dim, k_dim, v_dim], dim=-1) - expected = torch.cat([q.reshape(-1), k.reshape(-1), v.reshape(-1)]) - assert torch.equal(actual, expected) - - -def test_qwen35_gdn_qkv_pack_graph_replay_reads_current_values(): - if not torch.cuda.is_available(): - pytest.skip("CUDA is required for the Qwen3.5 GDN Q/K/V pack kernel") - - num_rows, q_dim, k_dim, v_dim = 8, 512, 512, 1536 - mixed_qkv = torch.randn( - (num_rows, q_dim + k_dim + v_dim), - dtype=torch.float16, - device="cuda", - ) - _sm70_pack_qwen_gdn_qkv(mixed_qkv, q_dim, k_dim, v_dim) - torch.accelerator.synchronize() - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - actual = _sm70_pack_qwen_gdn_qkv(mixed_qkv, q_dim, k_dim, v_dim) - - mixed_qkv.copy_(torch.randn_like(mixed_qkv)) - graph.replay() - torch.accelerator.synchronize() - - q, k, v = torch.split(mixed_qkv, [q_dim, k_dim, v_dim], dim=-1) - expected = torch.cat([q.reshape(-1), k.reshape(-1), v.reshape(-1)]) - assert torch.equal(actual, expected) diff --git a/tests/model_executor/test_deepseek_v4_swa_rope.py b/tests/model_executor/test_deepseek_v4_swa_rope.py deleted file mode 100644 index e494d16cfc..0000000000 --- a/tests/model_executor/test_deepseek_v4_swa_rope.py +++ /dev/null @@ -1,137 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -"""YaRN must be off on DeepSeek-V4's sliding-window-only layers. - -For compress_ratio == 0 layers the reference passes original_seq_len=0 to -precompute_freqs_cis (inference/model.py:481-486), which skips the correction -range entirely and leaves plain rope at the base rope_theta. Swapping only -rope_theta and leaving rope_type="deepseek_yarn" with factor=16 keeps the -interpolation on for those layers. -""" - -import math - -import pytest -import torch - -from vllm.model_executor.layers.rotary_embedding import get_rope -from vllm.platforms import current_platform - -# DeepSeek-V4-Flash config.json. -_HEAD_SIZE = 512 -_ROPE_DIM = 64 -_MAX_POSITION = 262144 # served context; the model passes max_position_embeddings -_ORIGINAL_MAX_POSITION = 65536 -_YARN_FACTOR = 16 -_ROPE_THETA = 10000 # SWA-only layers; compressing layers use 160000 -_SWA_WINDOW = 128 - - -@pytest.fixture(autouse=True) -def _default_device(): - """DeepseekV4ScalingRotaryEmbedding builds its position grid on - current_platform.device_type but its inv_freq on the default device, so it - only assembles under the device context the model loader already sets.""" - with torch.device(current_platform.device_type): - yield - - -def _rope_parameters(factor: int, original_max_position: int) -> dict: - """What DeepseekV4Attention hands get_rope for a compress_ratio<=1 layer.""" - return { - "rope_type": "deepseek_yarn", - "beta_fast": 32, - "beta_slow": 1, - "factor": factor, - "original_max_position_embeddings": original_max_position, - "rope_theta": _ROPE_THETA, - "mscale": 0, - "mscale_all_dim": 0, - "is_deepseek_v4": True, - "rope_dim": _ROPE_DIM, - } - - -def _get_rope(factor: int, original_max_position: int): - return get_rope( - _HEAD_SIZE, - max_position=_MAX_POSITION, - rope_parameters=_rope_parameters(factor, original_max_position), - is_neox_style=False, - ) - - -def _plain_inv_freq() -> torch.Tensor: - exponents = torch.arange(0, _ROPE_DIM, 2, dtype=torch.float32) / _ROPE_DIM - return 1.0 / (_ROPE_THETA**exponents) - - -def test_factor_one_is_exactly_plain_rope(default_vllm_config) -> None: - """factor=1 collapses YaRN to an identity, so no class swap is needed. - - DeepseekV4ScalingRotaryEmbedding is what the fused kernels expect: it - rotates the *last* rotary_dim and keeps cos/sin in fp32. Neither is true of - the base RotaryEmbedding, so "turn YaRN off" has to mean factor=1, not - rope_type="default". - """ - swa = _get_rope(1, _MAX_POSITION) - - assert swa.mscale == 1.0 - # inv_freq_interpolation == inv_freq_extrapolation at factor=1, so the ramp - # blend `a * (1 - m) + a * m` is algebraically `a`; in fp32 it lands within - # a couple of ulp on the two pairs where the mask is neither 0 nor 1. - torch.testing.assert_close( - swa._compute_inv_freq(swa.scaling_factor), - _plain_inv_freq(), - rtol=1e-6, - atol=0, - ) - # Cache length is original_max_position * factor, so original_max_position - # has to carry the full range once factor is 1. - assert swa.cos_sin_cache.shape[0] == _MAX_POSITION - - -def test_yarn_config_still_interpolates_the_low_frequency_pairs( - default_vllm_config, -) -> None: - """The unfixed config is a real divergence, not a no-op.""" - yarn = _get_rope(_YARN_FACTOR, _ORIGINAL_MAX_POSITION) - - yarn_inv_freq = yarn._compute_inv_freq(yarn.scaling_factor) - plain = _plain_inv_freq() - assert not torch.allclose(yarn_inv_freq, plain) - - # Only the low-frequency pairs move: high-frequency pairs sit above the - # correction range and extrapolate unchanged. - moved = (yarn_inv_freq != plain).nonzero().flatten().tolist() - assert moved and min(moved) >= _ROPE_DIM // 4 - - # Accumulated rotation error across one 128-token SWA window. - worst_radians = ((plain - yarn_inv_freq).abs() * (_SWA_WINDOW - 1)).max().item() - assert 0.0 < math.degrees(worst_radians) < 5.0 - - -def test_shared_rope_parameters_dict_must_not_be_mutated(default_vllm_config) -> None: - """A shared rope_parameters dict must not be mutated in place. - - DeepseekV4Attention reads config.rope_parameters itself, and every later - layer reads it back, so setting factor=1 in place would disable YaRN - model-wide. - """ - shared = _rope_parameters(_YARN_FACTOR, _ORIGINAL_MAX_POSITION) - before = dict(shared) - - swa_parameters = dict(shared) - swa_parameters["factor"] = 1 - swa_parameters["original_max_position_embeddings"] = _MAX_POSITION - get_rope( - _HEAD_SIZE, - max_position=_MAX_POSITION, - rope_parameters=swa_parameters, - is_neox_style=False, - ) - - assert shared == before - compressing = _get_rope(_YARN_FACTOR, _ORIGINAL_MAX_POSITION) - assert compressing.scaling_factor == _YARN_FACTOR diff --git a/tests/model_executor/test_qwen3_5_quantization.py b/tests/model_executor/test_qwen3_5_quantization.py index 74914bd765..ae7e7a79c3 100644 --- a/tests/model_executor/test_qwen3_5_quantization.py +++ b/tests/model_executor/test_qwen3_5_quantization.py @@ -1,102 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from types import SimpleNamespace from unittest.mock import Mock, patch -import pytest -import torch - class _QuantConfig: def __init__(self) -> None: self.ignore: list[str] = [] self.config: dict[str, object] = {} - self.quantized_layers: dict[str, dict[str, str]] = {} - - -@pytest.mark.parametrize( - ("split_enabled", "norm_enabled"), - [ - (False, False), - (False, True), - ], -) -def test_sm70_gdn_qpn8_ba_split_is_disabled_without_split_flag( - monkeypatch, - split_enabled: bool, - norm_enabled: bool, -) -> None: - from vllm.model_executor.layers.mamba.gdn import qwen_gdn_linear_attn as gdn - - monkeypatch.setattr(gdn.envs, "VLLM_SM70_GDN_QPN8_BA_SPLIT", split_enabled) - monkeypatch.setattr(gdn.envs, "VLLM_SM70_GDN_RMSNORM_ONEPASS", norm_enabled) - - assert not gdn._sm70_gdn_qpn8_ba_split_enabled() - - -def test_sm70_gdn_qpn8_ba_split_rejects_unpaired_route(monkeypatch) -> None: - from vllm.model_executor.layers.mamba.gdn import qwen_gdn_linear_attn as gdn - - monkeypatch.setattr(gdn.envs, "VLLM_SM70_GDN_QPN8_BA_SPLIT", True) - monkeypatch.setattr(gdn.envs, "VLLM_SM70_GDN_RMSNORM_ONEPASS", False) - - with pytest.raises(RuntimeError, match="requires the accepted"): - gdn._sm70_gdn_qpn8_ba_split_enabled() - - -def test_sm70_gdn_qpn8_ba_split_requires_source_built_ops(monkeypatch) -> None: - from vllm.model_executor.layers.mamba.gdn import qwen_gdn_linear_attn as gdn - - monkeypatch.setattr(gdn.envs, "VLLM_SM70_GDN_QPN8_BA_SPLIT", True) - monkeypatch.setattr(gdn.envs, "VLLM_SM70_GDN_RMSNORM_ONEPASS", True) - monkeypatch.setattr( - gdn, - "_missing_sm70_gdn_qpn8_ba_ops", - lambda: ["fp8_qpn8_dispatch_ba_split_sm70_out"], - ) - - with pytest.raises(RuntimeError, match="requires the source-built"): - gdn._sm70_gdn_qpn8_ba_split_enabled() - - -def test_sm70_gdn_qpn8_ba_split_accepts_complete_contract(monkeypatch) -> None: - from vllm.model_executor.layers.mamba.gdn import qwen_gdn_linear_attn as gdn - - monkeypatch.setattr(gdn.envs, "VLLM_SM70_GDN_QPN8_BA_SPLIT", True) - monkeypatch.setattr(gdn.envs, "VLLM_SM70_GDN_RMSNORM_ONEPASS", True) - monkeypatch.setattr(gdn, "_missing_sm70_gdn_qpn8_ba_ops", lambda: []) - - assert gdn._sm70_gdn_qpn8_ba_split_enabled() - - -def test_sm70_gdn_qpn8_ba_weight_contract_is_layout_based() -> 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_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"), - bias=None, - ) - layer = SimpleNamespace( - enable_sm70_gdn_qpn8_ba_split=True, - in_proj_qkvz=qkvz, - in_proj_ba=ba, - ) - - assert gdn._sm70_gdn_qpn8_ba_weight_contract(layer) - - 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") - assert not gdn._sm70_gdn_qpn8_ba_weight_contract(layer) def test_qwen3_5_split_gdn_detects_compressed_tensors_ignore(): @@ -130,53 +41,6 @@ def test_qwen3_5_split_gdn_detects_compressed_tensors_config_ignore(): assert _uses_split_gdn_input_projections(quant_config) -def test_qwen3_5_split_gdn_detects_modelopt_mixed_unquantized_ba(): - from vllm.model_executor.models.qwen3_5 import ( - _uses_split_gdn_input_projections, - ) - - quant_config = _QuantConfig() - prefix = "model.language_model.layers.0.linear_attn" - quant_config.quantized_layers = { - f"{prefix}.in_proj_qkv": {"quant_algo": "FP8"}, - f"{prefix}.in_proj_z": {"quant_algo": "FP8"}, - } - - assert _uses_split_gdn_input_projections(quant_config) - - -def test_qwen3_5_split_gdn_keeps_same_precision_projections_fused(): - from vllm.model_executor.models.qwen3_5 import ( - _uses_split_gdn_input_projections, - ) - - quant_config = _QuantConfig() - prefix = "model.language_model.layers.0.linear_attn" - quant_config.quantized_layers = { - f"{prefix}.{name}": {"quant_algo": "FP8"} - for name in ("in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a") - } - - assert not _uses_split_gdn_input_projections(quant_config) - - -def test_qwen3_5_split_gdn_detects_different_modelopt_algorithms(): - from vllm.model_executor.models.qwen3_5 import ( - _uses_split_gdn_input_projections, - ) - - quant_config = _QuantConfig() - prefix = "model.language_model.layers.0.linear_attn" - quant_config.quantized_layers = { - f"{prefix}.in_proj_qkv": {"quant_algo": "FP8"}, - f"{prefix}.in_proj_z": {"quant_algo": "FP8"}, - f"{prefix}.in_proj_b": {"quant_algo": "W4A16_NVFP4"}, - f"{prefix}.in_proj_a": {"quant_algo": "W4A16_NVFP4"}, - } - - assert _uses_split_gdn_input_projections(quant_config) - - def test_qwen3_5_lm_head_receives_quant_config(): from vllm.model_executor.models.qwen3_5 import Qwen3_5ForCausalLMBase diff --git a/tests/model_executor/test_sm70_fp8_qpn8_pp2_tp4.py b/tests/model_executor/test_sm70_fp8_qpn8_pp2_tp4.py deleted file mode 100644 index 17580656d8..0000000000 --- a/tests/model_executor/test_sm70_fp8_qpn8_pp2_tp4.py +++ /dev/null @@ -1,229 +0,0 @@ -# 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 import envs -from vllm.model_executor.layers.quantization import fp8 - - -def _layer( - suffix: str, - tp_size: int, - k_dim: int, - n_dim: int, - *, - output_partition_sizes: list[int] | None = None, -) -> SimpleNamespace: - return SimpleNamespace( - prefix=f"arbitrary.engine.graph.{suffix}", - tp_size=tp_size, - input_size_per_partition=k_dim, - output_size_per_partition=n_dim, - output_partition_sizes=output_partition_sizes, - weight_block_size=[128, 128], - weight=torch.empty((n_dim, k_dim), device="meta"), - ) - - -def test_pp2_tp4_qpn8_is_default_on_with_explicit_rollback(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 fp8._sm70_fp8_qpn8_pp2_tp4_enabled() - - monkeypatch.setenv("VLLM_SM70_FP8_QPN8", "0") - envs.disable_envs_cache() - assert not fp8._sm70_fp8_qpn8_pp2_tp4_enabled() - - monkeypatch.delenv("VLLM_SM70_FP8_QPN8") - monkeypatch.setenv("VLLM_SM70_FP8_QPN8_PP2_TP4", "0") - envs.disable_envs_cache() - assert not fp8._sm70_fp8_qpn8_pp2_tp4_enabled() - finally: - envs.disable_envs_cache() - - -@pytest.mark.parametrize( - ("layer", "gated_silu", "expected"), - [ - (_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( - layer: SimpleNamespace, - gated_silu: bool, - expected: tuple[int, int, bool], -) -> None: - assert fp8._sm70_fp8_qpn8_pp2_tp4_config(layer, gated_silu=gated_silu) == expected - - -def test_pp2_tp4_qpn8_rejects_wrong_tensor_and_concurrency_roles() -> None: - wrong_tp = _layer("wq_b", 8, 1024, 8192) - assert fp8._sm70_fp8_qpn8_pp2_tp4_config(wrong_tp, gated_silu=False) is None - - concurrent_indexer = _layer("wq_b", 1, 1024, 8192) - assert ( - fp8._sm70_fp8_qpn8_pp2_tp4_config(concurrent_indexer, gated_silu=False) is None - ) - - excluded_gate = _layer( - "gate_up_proj", - 4, - 4096, - 1024, - output_partition_sizes=[512, 512], - ) - assert fp8._sm70_fp8_qpn8_pp2_tp4_config(excluded_gate, gated_silu=False) is None - assert fp8._sm70_fp8_qpn8_pp2_tp4_config(excluded_gate, gated_silu=True) is None - - wrong_layout = _layer("wo_b", 4, 2048, 4096) - wrong_layout.weight_block_size = [64, 128] - assert fp8._sm70_fp8_qpn8_pp2_tp4_config(wrong_layout, gated_silu=False) is None - - -def test_pp2_tp4_qpn8_runtime_and_workspace_contract() -> None: - 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, - ) - with patch.object(fp8, "get_current_vllm_config", return_value=config): - assert fp8._is_sm70_fp8_qpn8_pp2_tp4_runtime_contract() - config.speculative_config = object() - assert not fp8._is_sm70_fp8_qpn8_pp2_tp4_runtime_contract() - config.speculative_config = None - config.parallel_config.pipeline_parallel_size = 1 - assert not fp8._is_sm70_fp8_qpn8_pp2_tp4_runtime_contract() - config.parallel_config.pipeline_parallel_size = 2 - config.scheduler_config.max_num_seqs = 2 - assert not fp8._is_sm70_fp8_qpn8_pp2_tp4_runtime_contract() - config.scheduler_config.max_num_seqs = 1 - config.parallel_config.enable_dbo = True - assert not fp8._is_sm70_fp8_qpn8_pp2_tp4_runtime_contract() - config.parallel_config.enable_dbo = False - config.parallel_config.ubatch_size = 2 - assert not fp8._is_sm70_fp8_qpn8_pp2_tp4_runtime_contract() - - assert fp8._SM70_FP8_QPN8_PP2_TP4_WORKSPACE_ELEMENTS * 2 == 16 * 1024 * 1024 - - -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) == (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 = 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") - - calls: list[tuple[tuple[int, ...], int, int, bool, bool]] = [] - - def fake_dispatch( - out: torch.Tensor, - workspace_ptr: int, - input_: torch.Tensor, - codes: torch.Tensor, - scales: torch.Tensor, - split_k: int, - nacc: int, - prefetch: bool, - gated_silu: bool, - ) -> None: - del codes, scales - calls.append( - (tuple(input_.shape), workspace_ptr, split_k, prefetch, gated_silu) - ) - out.fill_(len(calls)) - 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, 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: - monkeypatch.delenv("VLLM_SM70_FP8_QPN8", raising=False) - monkeypatch.delenv("VLLM_SM70_FP8_QPN8_PP2_TP4", raising=False) - envs.disable_envs_cache() - layer = _layer("fused_wqa_wkv", 1, 4096, 1536) - layer.orig_dtype = torch.float16 - layer.is_bmm = False - layer.weight_scale_inv = torch.empty((12, 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, 1536), dtype=torch.uint8, device="meta"), - torch.empty((256, 48), 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 diff --git a/tests/model_executor/test_uva_tower_offload.py b/tests/model_executor/test_uva_tower_offload.py deleted file mode 100644 index 91e3d5de93..0000000000 --- a/tests/model_executor/test_uva_tower_offload.py +++ /dev/null @@ -1,89 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from collections.abc import Generator -from types import SimpleNamespace - -import torch.nn as nn - -from vllm.model_executor.models.interfaces import SupportsMultiModal -from vllm.model_executor.offloader import ( - BaseOffloader, - PrefetchOffloader, - get_offloader, - set_offloader, -) -from vllm.model_executor.offloader.uva import UVAOffloader - - -class _RecordingOffloader(BaseOffloader): - def __init__(self) -> None: - self.component_calls: list[tuple[nn.Module, str]] = [] - - def wrap_modules( - self, modules_generator: Generator[nn.Module, None, None] - ) -> list[nn.Module]: - raise AssertionError("tower registration must not reuse the layer wrapper") - - def wrap_module( - self, module: nn.Module, *, parameter_prefix: str = "" - ) -> nn.Module: - self.component_calls.append((module, parameter_prefix)) - return module - - -class _MultiModalModel(nn.Module): - _mark_tower_model = SupportsMultiModal._mark_tower_model - - -def _selector_only_offloader(selectors: set[str]) -> UVAOffloader: - offloader = UVAOffloader.__new__(UVAOffloader) - offloader.cpu_offload_params = selectors - return offloader - - -def test_uva_tower_selector_uses_qualified_parameter_path() -> None: - offloader = _selector_only_offloader({"visual"}) - - assert offloader._is_parameter_selected("patch_embed.proj.weight", "visual") - assert not offloader._is_parameter_selected("patch_embed.proj.weight", "audio") - assert not offloader._is_parameter_selected("patch_embed.proj.weight", "visualizer") - assert _selector_only_offloader(set())._is_parameter_selected( - "patch_embed.proj.weight", "visual" - ) - - -def test_uva_multi_segment_selector_remains_segment_exact() -> None: - offloader = _selector_only_offloader({"visual.patch_embed"}) - - assert offloader._is_parameter_selected("patch_embed.proj.weight", "visual") - assert not offloader._is_parameter_selected("patch_embedding.proj.weight", "visual") - - -def test_tower_registration_uses_single_component_wrapper() -> None: - model = _MultiModalModel() - offloader = _RecordingOffloader() - previous_offloader = get_offloader() - config = SimpleNamespace( - model_config=SimpleNamespace( - multimodal_config=SimpleNamespace(get_limit_per_prompt=lambda _modality: 1) - ) - ) - - set_offloader(offloader) - try: - with model._mark_tower_model(config, "image"): - model.visual = nn.Linear(4, 4) - finally: - set_offloader(previous_offloader) - - assert model._tower_model_names == ["visual"] - assert offloader.component_calls == [(model.visual, "visual")] - - -def test_prefetch_inherits_noop_component_wrapper() -> None: - tower = nn.Linear(4, 4) - offloader = PrefetchOffloader.__new__(PrefetchOffloader) - - assert PrefetchOffloader.wrap_module is BaseOffloader.wrap_module - assert offloader.wrap_module(tower, parameter_prefix="visual") is tower diff --git a/tests/models/multimodal/processing/test_gemma4.py b/tests/models/multimodal/processing/test_gemma4.py index 00475a5925..a355501fdd 100644 --- a/tests/models/multimodal/processing/test_gemma4.py +++ b/tests/models/multimodal/processing/test_gemma4.py @@ -10,7 +10,6 @@ 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 @@ -23,51 +22,6 @@ 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/test_deepseek_v4_pipeline.py b/tests/models/test_deepseek_v4_pipeline.py deleted file mode 100644 index 2920cd5c2c..0000000000 --- a/tests/models/test_deepseek_v4_pipeline.py +++ /dev/null @@ -1,50 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from types import SimpleNamespace - -import pytest -import torch - -from vllm.distributed.utils import get_pp_indices -from vllm.models.deepseek_v4.nvidia.model import DeepseekV4Model - - -@pytest.mark.parametrize( - ("pp_size", "expected_partitions"), - [ - (2, [(0, 22), (22, 43)]), - (4, [(0, 11), (11, 22), (22, 33), (33, 43)]), - ], -) -def test_deepseek_v4_pipeline_partitions_keep_dspark_targets_on_last_stage( - monkeypatch: pytest.MonkeyPatch, - pp_size: int, - expected_partitions: list[tuple[int, int]], -) -> None: - monkeypatch.delenv("VLLM_PP_LAYER_PARTITION", raising=False) - - partitions = [get_pp_indices(43, rank, pp_size) for rank in range(pp_size)] - - assert partitions == expected_partitions - last_start, last_end = partitions[-1] - assert all(last_start <= layer_id < last_end for layer_id in (40, 41, 42)) - - -def test_deepseek_v4_pipeline_intermediate_preserves_mhc_streams() -> None: - model = SimpleNamespace( - hc_mult=4, - config=SimpleNamespace(hidden_size=4096), - ) - - intermediate = DeepseekV4Model.make_empty_intermediate_tensors( - model, - batch_size=3, - dtype=torch.float16, - device=torch.device("cpu"), - ) - - assert set(intermediate.tensors) == {"hidden_states"} - assert intermediate["hidden_states"].shape == (3, 4, 4096) - assert intermediate["hidden_states"].dtype == torch.float16 - assert torch.count_nonzero(intermediate["hidden_states"]) == 0 diff --git a/tests/models/test_deepseek_v4_sm70_routes.py b/tests/models/test_deepseek_v4_sm70_routes.py index a983b2b206..0c5ea28f90 100644 --- a/tests/models/test_deepseek_v4_sm70_routes.py +++ b/tests/models/test_deepseek_v4_sm70_routes.py @@ -101,14 +101,6 @@ def test_sm70_sparse_qk_dsplit_uses_graph_workspace(): kwargs = qk_dsplit.call_args.kwargs assert kwargs["partial_qk"].shape == (1, 8, 8, 8, 16) assert kwargs["partial_probs"].shape == (1, 8, 8, 16) - assert kwargs["stage1_block_h"] == 8 - - -def test_sm70_sparse_qk_dsplit_uses_one_tp4_head_group(): - from vllm.models.deepseek_v4.sm70.sparse import _qk_dsplit_block_h - - assert _qk_dsplit_block_h(16) == 16 - assert _qk_dsplit_block_h(8) == 8 def test_sm75_does_not_select_sm70_impl(): @@ -205,52 +197,6 @@ def make_metadata(starts: list[int]) -> CommonAttentionMetadata: assert _get_c128_boundary(make_metadata([127, 10])) is True -def test_sm70_private_compressor_state_requires_a_contiguous_single_request(): - from vllm.models.deepseek_v4 import compressor - - platform = MagicMock() - platform.is_cuda.return_value = True - platform.is_device_capability.return_value = True - - config = MagicMock() - config.scheduler_config.max_num_seqs = 1 - config.cache_config.enable_prefix_caching = False - config.parallel_config.pipeline_parallel_size = 1 - config.kv_transfer_config = None - config.speculative_config.parallel_drafting = False - - with ( - patch.object(compressor, "current_platform", platform), - patch.object( - compressor.envs, - "VLLM_SM70_DSV4_PRIVATE_COMPRESSOR_STATE", - False, - ), - ): - assert not compressor._can_use_sm70_private_compressor_state(config) - - with ( - patch.object(compressor, "current_platform", platform), - patch.object( - compressor.envs, - "VLLM_SM70_DSV4_PRIVATE_COMPRESSOR_STATE", - True, - ), - ): - assert compressor._can_use_sm70_private_compressor_state(config) - - config.speculative_config.parallel_drafting = True - assert not compressor._can_use_sm70_private_compressor_state(config) - - config.speculative_config.parallel_drafting = False - config.scheduler_config.max_num_seqs = 2 - assert not compressor._can_use_sm70_private_compressor_state(config) - - config.scheduler_config.max_num_seqs = 1 - config.cache_config.enable_prefix_caching = True - assert not compressor._can_use_sm70_private_compressor_state(config) - - def test_v4_c128_metadata_keeps_graph_stable_row_stride(): from vllm.v1.attention.backends.mla import flashmla_sparse diff --git a/tests/quantization/test_nvfp4_mxfp4_gating.py b/tests/quantization/test_nvfp4_mxfp4_gating.py index d3f0ed641c..6f03dc5ade 100644 --- a/tests/quantization/test_nvfp4_mxfp4_gating.py +++ b/tests/quantization/test_nvfp4_mxfp4_gating.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from types import SimpleNamespace - import torch from torch.nn.parameter import Parameter @@ -20,9 +18,6 @@ from vllm.model_executor.layers.quantization.compressed_tensors.schemes import ( compressed_tensors_w4a4_mxfp4 as mxfp4_scheme, ) -from vllm.model_executor.layers.quantization.compressed_tensors.schemes import ( - compressed_tensors_w4a4_nvfp4 as nvfp4_scheme, -) from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_w4a4_mxfp4 import ( # noqa: E501 CompressedTensorsW4A4Mxfp4, ) @@ -98,92 +93,6 @@ def test_flashinfer_nvfp4_backends_reject_sm70(): assert "sm_100" in cudnn_reason -def _make_sm70_tp4_nvfp4_gate_up_layer() -> torch.nn.Module: - layer = torch.nn.Module() - layer.tp_size = 4 - layer.prefix = "model.language_model.layers.0.mlp.gate_up_proj" - layer.input_size_per_partition = 5120 - layer.output_size_per_partition = 8704 - layer.logical_widths = [4352, 4352] - return layer - - -def _make_sm70_tp4_nvfp4_down_layer() -> torch.nn.Module: - layer = torch.nn.Module() - layer.tp_size = 4 - layer.prefix = "model.language_model.layers.0.mlp.down_proj" - layer.input_size_per_partition = 4352 - layer.output_size_per_partition = 5120 - return layer - - -def test_sm70_tp4_nvfp4_gate_up_match_is_shape_exact(): - layer = _make_sm70_tp4_nvfp4_gate_up_layer() - - assert nvfp4_scheme._is_sm70_tp4_nvfp4_gate_up(layer) - - layer.tp_size = 2 - assert not nvfp4_scheme._is_sm70_tp4_nvfp4_gate_up(layer) - layer.tp_size = 4 - layer.output_size_per_partition = 8703 - assert not nvfp4_scheme._is_sm70_tp4_nvfp4_gate_up(layer) - layer.output_size_per_partition = 8704 - layer.prefix = "model.language_model.layers.0.mlp.down_proj" - assert not nvfp4_scheme._is_sm70_tp4_nvfp4_gate_up(layer) - - -def test_sm70_tp4_nvfp4_down_match_is_shape_exact(): - layer = _make_sm70_tp4_nvfp4_down_layer() - - assert nvfp4_scheme._is_sm70_tp4_nvfp4_down(layer) - - layer.input_size_per_partition = 4351 - assert not nvfp4_scheme._is_sm70_tp4_nvfp4_down(layer) - layer.input_size_per_partition = 4352 - layer.output_size_per_partition = 5119 - assert not nvfp4_scheme._is_sm70_tp4_nvfp4_down(layer) - - -def test_sm70_nvfp4_qpn4_runtime_contract_is_single_sequence_no_mtp( - monkeypatch, -): - config = SimpleNamespace( - scheduler_config=SimpleNamespace(max_num_seqs=1), - speculative_config=None, - ) - monkeypatch.setattr( - nvfp4_scheme, - "get_current_vllm_config", - lambda: config, - ) - - assert nvfp4_scheme._is_sm70_nvfp4_qpn4_runtime_contract() - - config.scheduler_config.max_num_seqs = 2 - assert not nvfp4_scheme._is_sm70_nvfp4_qpn4_runtime_contract() - config.scheduler_config.max_num_seqs = 1 - config.speculative_config = object() - assert not nvfp4_scheme._is_sm70_nvfp4_qpn4_runtime_contract() - - -def test_nvfp4_scheme_delegates_fused_silu_only_for_prepared_layer(monkeypatch): - scheme = CompressedTensorsW4A4Fp4() - layer = torch.nn.Module() - x = torch.ones((1, 4), dtype=torch.float16) - expected = torch.full((1, 2), 7.0, dtype=torch.float16) - - monkeypatch.setattr(nvfp4_scheme.sm70_tm, "has_prepared_linear", lambda _: True) - monkeypatch.setattr( - nvfp4_scheme.sm70_tm, - "apply_prepared_fused_silu_and_mul", - lambda prepared_layer, prepared_x: expected - if prepared_layer is layer and prepared_x is x - else None, - ) - - assert scheme.apply_fused_silu_and_mul(layer, x) is expected - - def _make_mxfp4_layer() -> torch.nn.Module: layer = torch.nn.Module() layer.weight_packed = Parameter( diff --git a/tests/quantization/test_sm70_compressed_tensors_fastpaths.py b/tests/quantization/test_sm70_compressed_tensors_fastpaths.py deleted file mode 100644 index 2512a5864c..0000000000 --- a/tests/quantization/test_sm70_compressed_tensors_fastpaths.py +++ /dev/null @@ -1,293 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from types import SimpleNamespace - -import torch -from compressed_tensors.quantization import ( - QuantizationArgs, - QuantizationStrategy, - QuantizationType, -) - -import vllm.envs as envs -from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501 - CompressedTensorsLinearMethod, -) -from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_w8a16_fp8 import ( # noqa: E501 - CompressedTensorsW8A16Fp8, - _sm70_channel_fp8_qpn8_config, -) - - -def _make_channel_fp8_scheme(monkeypatch, *, static_input: bool = False): - monkeypatch.setattr( - "vllm.model_executor.layers.quantization.compressed_tensors.schemes." - "compressed_tensors_w8a16_fp8.get_current_vllm_config", - lambda: SimpleNamespace(model_config=SimpleNamespace(dtype=torch.float16)), - ) - monkeypatch.setattr( - "vllm.model_executor.layers.quantization.compressed_tensors.schemes." - "compressed_tensors_w8a16_fp8.sm70_tm.is_exact_sm70_cuda_platform", - lambda: True, - ) - monkeypatch.setattr( - "vllm.model_executor.layers.quantization.compressed_tensors.schemes." - "compressed_tensors_w8a16_fp8.sm70_tm.use_turbomind", - lambda enabled: enabled, - ) - weight_quant = QuantizationArgs( - num_bits=8, - type=QuantizationType.FLOAT, - symmetric=True, - strategy=QuantizationStrategy.CHANNEL, - dynamic=False, - ) - return CompressedTensorsW8A16Fp8(weight_quant, static_input) - - -def test_compressed_tensors_channel_fp8_selects_turbomind_on_sm70(monkeypatch): - monkeypatch.setenv("VLLM_SM70_FP8_TURBOMIND", "1") - envs.disable_envs_cache() - try: - assert _make_channel_fp8_scheme(monkeypatch).use_sm70_fp8_turbomind - assert not _make_channel_fp8_scheme( - monkeypatch, static_input=True - ).use_sm70_fp8_turbomind - finally: - envs.disable_envs_cache() - - -def test_compressed_tensors_channel_fp8_prepares_and_dispatches_turbomind( - monkeypatch, -): - scheme = object.__new__(CompressedTensorsW8A16Fp8) - scheme.use_sm70_fp8_turbomind = True - layer = torch.nn.Module() - layer.register_parameter( - "weight", - torch.nn.Parameter( - torch.empty((6, 4), dtype=torch.float8_e4m3fn), requires_grad=False - ), - ) - layer.register_parameter( - "weight_scale", - torch.nn.Parameter(torch.ones((6, 1)), requires_grad=False), - ) - layer.orig_dtype = torch.float16 - layer.output_size_per_partition = 6 - prepare_calls = [] - - monkeypatch.setattr(torch.ops._C, "fp8_sm70_prepare", object(), raising=False) - - def fake_prepare(weight, scales, group_size, gated_silu): - prepare_calls.append( - (tuple(weight.shape), tuple(scales.shape), group_size, gated_silu) - ) - return ( - torch.empty((4, 6), dtype=torch.uint8), - torch.ones((1, 6), dtype=torch.float16), - torch.tensor([4, 6]), - ) - - module = ( - "vllm.model_executor.layers.quantization.compressed_tensors.schemes." - "compressed_tensors_w8a16_fp8" - ) - monkeypatch.setattr(f"{module}.sm70_ops.fp8_sm70_prepare", fake_prepare) - scheme.process_weights_after_loading(layer) - - assert prepare_calls == [((6, 4), (6, 1), 128, False)] - assert layer.sm70_fp8_turbomind - assert layer.sm70_fp8_channel_scale - assert tuple(layer.weight.shape) == (4, 6) - assert tuple(layer.weight_scale_inv.shape) == (1, 6) - assert not hasattr(layer, "weight_scale") - - gemm_calls = [] - - def fake_gemm(out, x, weight, scales, group_size, k_ld, q_ld, gated_silu): - gemm_calls.append( - ( - tuple(out.shape), - tuple(x.shape), - tuple(weight.shape), - tuple(scales.shape), - group_size, - k_ld, - q_ld, - gated_silu, - ) - ) - out.fill_(2) - - monkeypatch.setattr(f"{module}.sm70_ops.fp8_gemm_sm70_out", fake_gemm) - output = scheme.apply_weights( - layer, - torch.ones((2, 4), dtype=torch.float16), - torch.ones(6, dtype=torch.float16), - ) - - assert torch.equal(output, torch.full((2, 6), 3, dtype=torch.float16)) - assert gemm_calls == [((2, 6), (2, 4), (4, 6), (1, 6), 128, 4, 6, False)] - - -def test_compressed_tensors_channel_fp8_qpn8_shape_gate_is_exact(): - layer = SimpleNamespace( - tp_size=4, - prefix="model.language_model.layers.0.linear_attn.in_proj_qkvz", - weight=SimpleNamespace(shape=(4096, 5120)), - ) - assert _sm70_channel_fp8_qpn8_config(layer) == (16, 2, False) - - layer.prefix = "model.language_model.layers.3.self_attn.qkv_proj" - layer.weight = SimpleNamespace(shape=(3584, 5120)) - assert _sm70_channel_fp8_qpn8_config(layer) == (16, 2, False) - - layer.prefix = "model.language_model.layers.3.self_attn.o_proj" - layer.weight = SimpleNamespace(shape=(5120, 1536)) - assert _sm70_channel_fp8_qpn8_config(layer) == (12, 2, False) - - layer.prefix = "model.language_model.lm_head" - layer.weight = SimpleNamespace(shape=(62080, 5120)) - assert _sm70_channel_fp8_qpn8_config(layer) is None - - layer.tp_size = 2 - layer.prefix = "model.language_model.layers.3.self_attn.o_proj" - layer.weight = SimpleNamespace(shape=(5120, 1536)) - assert _sm70_channel_fp8_qpn8_config(layer) is None - - -def test_compressed_tensors_channel_fp8_qpn8_prepares_and_dispatches(monkeypatch): - scheme = object.__new__(CompressedTensorsW8A16Fp8) - scheme.use_sm70_fp8_turbomind = True - scheme.use_sm70_fp8_qpn8 = True - layer = torch.nn.Module() - layer.register_parameter( - "weight", - torch.nn.Parameter( - torch.empty((6, 4), dtype=torch.float8_e4m3fn), requires_grad=False - ), - ) - layer.register_parameter( - "weight_scale", - torch.nn.Parameter(torch.ones((6, 1)), requires_grad=False), - ) - layer.orig_dtype = torch.float16 - layer.output_size_per_partition = 6 - layer.prefix = "model.language_model.layers.63.mlp.gate_up_proj" - workspace = torch.empty(1, dtype=torch.float16) - prepare_calls = [] - - module = ( - "vllm.model_executor.layers.quantization.compressed_tensors.schemes." - "compressed_tensors_w8a16_fp8" - ) - monkeypatch.setattr(torch.ops._C, "fp8_sm70_prepare", object(), raising=False) - monkeypatch.setattr( - f"{module}._sm70_channel_fp8_qpn8_config", lambda layer: (16, 2, False) - ) - monkeypatch.setattr(f"{module}._is_sm70_fp8_qpn8_runtime_contract", lambda: True) - monkeypatch.setattr(f"{module}._missing_sm70_fp8_qpn8_ops", lambda: []) - monkeypatch.setattr( - f"{module}._get_sm70_fp8_prefill_exact_dense_workspace", - lambda weight: workspace, - ) - - def fake_prepare(weight, scales): - prepare_calls.append((tuple(weight.shape), tuple(scales.shape))) - return ( - torch.empty((4, 6), dtype=torch.uint8), - torch.ones((1, 6), dtype=torch.float16), - ) - - monkeypatch.setattr(f"{module}.sm70_ops.fp8_qpn8_prepare_sm70", fake_prepare) - scheme.process_weights_after_loading(layer) - - assert prepare_calls == [((6, 4), (6, 1))] - assert layer.sm70_fp8_turbomind - assert layer.sm70_fp8_qpn8 - assert layer.sm70_fp8_channel_scale - assert layer.sm70_fp8_qpn8_split_k == 16 - assert layer.sm70_fp8_qpn8_nacc == 2 - assert not layer.sm70_fp8_qpn8_prefetch - assert layer.sm70_fp8_gated_silu - assert layer.sm70_fp8_qpn8_gated_split_k == 8 - assert layer.sm70_fp8_qpn8_gated_nacc == 2 - assert not layer.sm70_fp8_qpn8_gated_prefetch - assert tuple(layer.weight.shape) == (4, 6) - assert tuple(layer.weight_scale_inv.shape) == (1, 6) - assert not hasattr(layer, "weight_scale") - - dispatch_calls = [] - - def fake_dispatch( - out, workspace_ptr, x, weight, scales, split_k, nacc, prefetch, gated_silu - ): - dispatch_calls.append( - ( - tuple(out.shape), - tuple(x.shape), - workspace_ptr, - tuple(weight.shape), - tuple(scales.shape), - split_k, - nacc, - prefetch, - gated_silu, - ) - ) - out.fill_(2) - - monkeypatch.setattr(f"{module}.sm70_ops.fp8_qpn8_dispatch_sm70_out", fake_dispatch) - output = scheme.apply_weights( - layer, - torch.ones((2, 4), dtype=torch.float16), - torch.ones(6, dtype=torch.float16), - ) - - assert torch.equal(output, torch.full((2, 6), 3, dtype=torch.float16)) - fused_output = scheme.apply_fused_silu_and_mul( - layer, torch.ones((2, 4), dtype=torch.float16) - ) - assert fused_output is not None - assert torch.equal(fused_output, torch.full((2, 3), 2, dtype=torch.float16)) - assert dispatch_calls == [ - ( - (2, 6), - (2, 4), - workspace.data_ptr(), - (4, 6), - (1, 6), - 16, - 2, - False, - False, - ), - ( - (2, 3), - (2, 4), - workspace.data_ptr(), - (4, 6), - (1, 6), - 8, - 2, - False, - True, - ), - ] - - -def test_compressed_tensors_linear_method_delegates_fused_apply(): - expected = torch.ones((1, 3), dtype=torch.float16) - layer = SimpleNamespace( - scheme=SimpleNamespace( - apply_fused_silu_and_mul=lambda layer, x: expected, - ) - ) - method = object.__new__(CompressedTensorsLinearMethod) - - assert ( - method.apply_fused_silu_and_mul(layer, torch.ones((1, 4), dtype=torch.float16)) - is expected - ) diff --git a/tests/quantization/test_sm70_fp8_prefill_exact_dense.py b/tests/quantization/test_sm70_fp8_prefill_exact_dense.py index 166f871954..800e658dc7 100644 --- a/tests/quantization/test_sm70_fp8_prefill_exact_dense.py +++ b/tests/quantization/test_sm70_fp8_prefill_exact_dense.py @@ -4,344 +4,18 @@ from types import SimpleNamespace import torch -from compressed_tensors.quantization import ( - QuantizationArgs, - QuantizationStrategy, - QuantizationType, -) import vllm.envs as envs -from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501 - CompressedTensorsLinearMethod, -) -from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_w8a16_fp8 import ( # noqa: E501 - CompressedTensorsW8A16Fp8, - _sm70_channel_fp8_qpn8_config, - _sm70_fp8_qpn8_enabled, -) from vllm.model_executor.layers.quantization.fp8 import ( - _SM70_FP8_EXACT_8K_PREFILL_M, _SM70_FP8_PREFILL_DENSE_MIN_M, _SM70_FP8_PREFILL_DENSE_WORKSPACE_BYTES, Fp8LinearMethod, _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, ) -def _make_channel_fp8_scheme( - monkeypatch, - *, - static_input: bool = False, - enable_qpn8_by_default: bool = False, -): - monkeypatch.setattr( - "vllm.model_executor.layers.quantization.compressed_tensors.schemes." - "compressed_tensors_w8a16_fp8.get_current_vllm_config", - lambda: SimpleNamespace(model_config=SimpleNamespace(dtype=torch.float16)), - ) - monkeypatch.setattr( - "vllm.model_executor.layers.quantization.compressed_tensors.schemes." - "compressed_tensors_w8a16_fp8.sm70_tm.is_exact_sm70_cuda_platform", - lambda: True, - ) - monkeypatch.setattr( - "vllm.model_executor.layers.quantization.compressed_tensors.schemes." - "compressed_tensors_w8a16_fp8.sm70_tm.use_turbomind", - lambda enabled: enabled, - ) - weight_quant = QuantizationArgs( - num_bits=8, - type=QuantizationType.FLOAT, - symmetric=True, - strategy=QuantizationStrategy.CHANNEL, - dynamic=False, - ) - return CompressedTensorsW8A16Fp8( - weight_quant, - static_input, - enable_sm70_qpn8_by_default=enable_qpn8_by_default, - ) - - -def test_compressed_tensors_channel_fp8_selects_turbomind_on_sm70(monkeypatch): - monkeypatch.setenv("VLLM_SM70_FP8_TURBOMIND", "1") - envs.disable_envs_cache() - try: - assert _make_channel_fp8_scheme(monkeypatch).use_sm70_fp8_turbomind - assert not _make_channel_fp8_scheme( - monkeypatch, static_input=True - ).use_sm70_fp8_turbomind - finally: - envs.disable_envs_cache() - - -def test_channel_fp8_qpn8_uses_mixed_nvfp4_default_only(monkeypatch): - monkeypatch.delenv("VLLM_SM70_FP8_QPN8", raising=False) - envs.disable_envs_cache() - try: - assert not _make_channel_fp8_scheme(monkeypatch).use_sm70_fp8_qpn8 - assert _make_channel_fp8_scheme( - monkeypatch, - enable_qpn8_by_default=True, - ).use_sm70_fp8_qpn8 - - monkeypatch.setenv("VLLM_SM70_FP8_QPN8", "0") - envs.disable_envs_cache() - assert not _make_channel_fp8_scheme( - monkeypatch, - enable_qpn8_by_default=True, - ).use_sm70_fp8_qpn8 - finally: - envs.disable_envs_cache() - - -def test_compressed_tensors_channel_fp8_prepares_and_dispatches_turbomind( - monkeypatch, -): - scheme = object.__new__(CompressedTensorsW8A16Fp8) - scheme.use_sm70_fp8_turbomind = True - layer = torch.nn.Module() - layer.register_parameter( - "weight", - torch.nn.Parameter( - torch.empty((6, 4), dtype=torch.float8_e4m3fn), requires_grad=False - ), - ) - layer.register_parameter( - "weight_scale", - torch.nn.Parameter(torch.ones((6, 1)), requires_grad=False), - ) - layer.orig_dtype = torch.float16 - layer.output_size_per_partition = 6 - prepare_calls = [] - - monkeypatch.setattr(torch.ops._C, "fp8_sm70_prepare", object(), raising=False) - - def fake_prepare(weight, scales, group_size, gated_silu): - prepare_calls.append( - (tuple(weight.shape), tuple(scales.shape), group_size, gated_silu) - ) - return ( - torch.empty((4, 6), dtype=torch.uint8), - torch.ones((1, 6), dtype=torch.float16), - torch.tensor([4, 6]), - ) - - monkeypatch.setattr( - "vllm.model_executor.layers.quantization.compressed_tensors.schemes." - "compressed_tensors_w8a16_fp8.sm70_ops.fp8_sm70_prepare", - fake_prepare, - ) - scheme.process_weights_after_loading(layer) - - assert prepare_calls == [((6, 4), (6, 1), 128, False)] - assert layer.sm70_fp8_turbomind - assert layer.sm70_fp8_channel_scale - assert tuple(layer.weight.shape) == (4, 6) - assert tuple(layer.weight_scale_inv.shape) == (1, 6) - assert not hasattr(layer, "weight_scale") - - gemm_calls = [] - - def fake_gemm(out, x, weight, scales, group_size, k_ld, q_ld, gated_silu): - gemm_calls.append( - ( - tuple(out.shape), - tuple(x.shape), - tuple(weight.shape), - tuple(scales.shape), - group_size, - k_ld, - q_ld, - gated_silu, - ) - ) - out.fill_(2) - - monkeypatch.setattr( - "vllm.model_executor.layers.quantization.compressed_tensors.schemes." - "compressed_tensors_w8a16_fp8.sm70_ops.fp8_gemm_sm70_out", - fake_gemm, - ) - output = scheme.apply_weights( - layer, - torch.ones((2, 4), dtype=torch.float16), - torch.ones(6, dtype=torch.float16), - ) - - assert torch.equal(output, torch.full((2, 6), 3, dtype=torch.float16)) - assert gemm_calls == [((2, 6), (2, 4), (4, 6), (1, 6), 128, 4, 6, False)] - - -def test_compressed_tensors_channel_fp8_qpn8_shape_gate_is_exact(): - layer = SimpleNamespace( - tp_size=4, - prefix="model.language_model.layers.0.linear_attn.in_proj_qkvz", - weight=SimpleNamespace(shape=(4096, 5120)), - ) - assert _sm70_channel_fp8_qpn8_config(layer) == (16, 2, False) - - layer.prefix = "model.language_model.layers.3.self_attn.qkv_proj" - layer.weight = SimpleNamespace(shape=(3584, 5120)) - assert _sm70_channel_fp8_qpn8_config(layer) == (16, 2, False) - - layer.prefix = "model.language_model.layers.3.self_attn.o_proj" - layer.weight = SimpleNamespace(shape=(5120, 1536)) - assert _sm70_channel_fp8_qpn8_config(layer) == (12, 2, False) - - layer.prefix = "model.language_model.lm_head" - layer.weight = SimpleNamespace(shape=(62080, 5120)) - assert _sm70_channel_fp8_qpn8_config(layer) is None - - layer.tp_size = 2 - layer.prefix = "model.language_model.layers.3.self_attn.o_proj" - layer.weight = SimpleNamespace(shape=(5120, 1536)) - assert _sm70_channel_fp8_qpn8_config(layer) is None - - -def test_compressed_tensors_channel_fp8_qpn8_prepares_and_dispatches(monkeypatch): - scheme = object.__new__(CompressedTensorsW8A16Fp8) - scheme.use_sm70_fp8_turbomind = True - scheme.use_sm70_fp8_qpn8 = True - layer = torch.nn.Module() - layer.register_parameter( - "weight", - torch.nn.Parameter( - torch.empty((6, 4), dtype=torch.float8_e4m3fn), requires_grad=False - ), - ) - layer.register_parameter( - "weight_scale", - torch.nn.Parameter(torch.ones((6, 1)), requires_grad=False), - ) - layer.orig_dtype = torch.float16 - layer.output_size_per_partition = 6 - layer.prefix = "model.language_model.layers.63.mlp.gate_up_proj" - workspace = torch.empty(1, dtype=torch.float16) - prepare_calls = [] - - module = ( - "vllm.model_executor.layers.quantization.compressed_tensors.schemes." - "compressed_tensors_w8a16_fp8" - ) - monkeypatch.setattr(torch.ops._C, "fp8_sm70_prepare", object(), raising=False) - monkeypatch.setattr( - f"{module}._sm70_channel_fp8_qpn8_config", lambda layer: (16, 2, False) - ) - monkeypatch.setattr(f"{module}._is_sm70_fp8_qpn8_runtime_contract", lambda: True) - monkeypatch.setattr(f"{module}._missing_sm70_fp8_qpn8_ops", lambda: []) - monkeypatch.setattr( - f"{module}._get_sm70_fp8_prefill_exact_dense_workspace", - lambda weight: workspace, - ) - - def fake_prepare(weight, scales): - prepare_calls.append((tuple(weight.shape), tuple(scales.shape))) - return ( - torch.empty((4, 6), dtype=torch.uint8), - torch.ones((1, 6), dtype=torch.float16), - ) - - monkeypatch.setattr(f"{module}.sm70_ops.fp8_qpn8_prepare_sm70", fake_prepare) - scheme.process_weights_after_loading(layer) - - assert prepare_calls == [((6, 4), (6, 1))] - assert layer.sm70_fp8_turbomind - assert layer.sm70_fp8_qpn8 - assert layer.sm70_fp8_channel_scale - assert layer.sm70_fp8_qpn8_split_k == 16 - assert layer.sm70_fp8_qpn8_nacc == 2 - assert not layer.sm70_fp8_qpn8_prefetch - assert layer.sm70_fp8_gated_silu - assert layer.sm70_fp8_qpn8_gated_split_k == 8 - assert layer.sm70_fp8_qpn8_gated_nacc == 2 - assert not layer.sm70_fp8_qpn8_gated_prefetch - assert tuple(layer.weight.shape) == (4, 6) - assert tuple(layer.weight_scale_inv.shape) == (1, 6) - assert not hasattr(layer, "weight_scale") - - dispatch_calls = [] - - def fake_dispatch( - out, workspace_ptr, x, weight, scales, split_k, nacc, prefetch, gated_silu - ): - dispatch_calls.append( - ( - tuple(out.shape), - tuple(x.shape), - workspace_ptr, - tuple(weight.shape), - tuple(scales.shape), - split_k, - nacc, - prefetch, - gated_silu, - ) - ) - out.fill_(2) - - monkeypatch.setattr(f"{module}.sm70_ops.fp8_qpn8_dispatch_sm70_out", fake_dispatch) - output = scheme.apply_weights( - layer, - torch.ones((2, 4), dtype=torch.float16), - torch.ones(6, dtype=torch.float16), - ) - - assert torch.equal(output, torch.full((2, 6), 3, dtype=torch.float16)) - fused_output = scheme.apply_fused_silu_and_mul( - layer, torch.ones((2, 4), dtype=torch.float16) - ) - assert fused_output is not None - assert torch.equal(fused_output, torch.full((2, 3), 2, dtype=torch.float16)) - assert dispatch_calls == [ - ( - (2, 6), - (2, 4), - workspace.data_ptr(), - (4, 6), - (1, 6), - 16, - 2, - False, - False, - ), - ( - (2, 3), - (2, 4), - workspace.data_ptr(), - (4, 6), - (1, 6), - 8, - 2, - False, - True, - ), - ] - - -def test_compressed_tensors_linear_method_delegates_fused_apply(): - expected = torch.ones((1, 3), dtype=torch.float16) - layer = SimpleNamespace( - scheme=SimpleNamespace( - apply_fused_silu_and_mul=lambda layer, x: expected, - ) - ) - method = object.__new__(CompressedTensorsLinearMethod) - - assert ( - method.apply_fused_silu_and_mul(layer, torch.ones((1, 4), dtype=torch.float16)) - is expected - ) - - def test_fp8_prefill_exact_dense_is_default_on(monkeypatch): monkeypatch.delenv("VLLM_SM70_FP8_PREFILL_EXACT_DENSE", raising=False) envs.disable_envs_cache() @@ -351,49 +25,9 @@ def test_fp8_prefill_exact_dense_is_default_on(monkeypatch): envs.disable_envs_cache() -def test_fp8_exact_8k_prefill_controls_are_generic_defaults(monkeypatch): - for name in ( - "VLLM_SM70_FP8_PREFILL_FAST_SELECTOR", - "VLLM_SM70_FP8_PREFILL_PRESCALED", - "VLLM_SM70_FP8_PREFILL_CUTLASS", - "VLLM_SM70_FP8_PREFILL_VISIBLE_DENSE_MM", - ): - monkeypatch.delenv(name, raising=False) - envs.disable_envs_cache() - try: - assert envs.VLLM_SM70_FP8_PREFILL_FAST_SELECTOR - assert envs.VLLM_SM70_FP8_PREFILL_PRESCALED - assert envs.VLLM_SM70_FP8_PREFILL_CUTLASS - assert not envs.VLLM_SM70_FP8_PREFILL_VISIBLE_DENSE_MM - finally: - envs.disable_envs_cache() - - -def test_fp8_qpn8_is_opt_in_except_for_mixed_nvfp4(monkeypatch): - monkeypatch.delenv("VLLM_SM70_FP8_QPN8", raising=False) - monkeypatch.delenv("VLLM_SM70_FP8_QPN8_LIBRARY", raising=False) - envs.disable_envs_cache() - try: - assert not envs.VLLM_SM70_FP8_QPN8 - assert envs.VLLM_SM70_FP8_QPN8_LIBRARY is None - assert not _sm70_fp8_qpn8_enabled(False) - assert _sm70_fp8_qpn8_enabled(True) - monkeypatch.setenv("VLLM_SM70_FP8_QPN8", "0") - envs.disable_envs_cache() - assert not envs.VLLM_SM70_FP8_QPN8 - assert not _sm70_fp8_qpn8_enabled(True) - monkeypatch.setenv("VLLM_SM70_FP8_QPN8", "1") - envs.disable_envs_cache() - assert envs.VLLM_SM70_FP8_QPN8 - assert _sm70_fp8_qpn8_enabled(False) - finally: - envs.disable_envs_cache() - - def test_fp8_prefill_exact_dense_shape_gate_is_narrow(): layer = SimpleNamespace( tp_size=4, - weight_block_size=[128, 128], prefix="model.language_model.layers.1.mlp.gate_up_proj", weight=SimpleNamespace(shape=(5120, 8704)), ) @@ -405,89 +39,13 @@ def test_fp8_prefill_exact_dense_shape_gate_is_narrow(): layer.tp_size = 4 layer.prefix = "model.language_model.layers.1.self_attn.qkv_proj" layer.weight = SimpleNamespace(shape=(5120, 3584)) - assert _is_sm70_fp8_prefill_exact_dense_layer(layer) + assert not _is_sm70_fp8_prefill_exact_dense_layer(layer) layer.prefix = "model.language_model.layers.1.linear_attn.in_proj_qkvz" layer.weight = SimpleNamespace(shape=(5120, 4096)) - assert _is_sm70_fp8_prefill_exact_dense_layer(layer) + assert not _is_sm70_fp8_prefill_exact_dense_layer(layer) layer.prefix = "model.language_model.layers.1.mlp.gate_up_proj" layer.weight = SimpleNamespace(shape=(5120, 8192)) assert not _is_sm70_fp8_prefill_exact_dense_layer(layer) - layer.weight_block_size = [64, 128] - layer.weight = SimpleNamespace(shape=(5120, 8704)) - assert not _is_sm70_fp8_prefill_exact_dense_layer(layer) - - -def test_fp8_exact_8k_prefill_gate_uses_operator_contract_only(): - layer = SimpleNamespace( - tp_size=4, - weight_block_size=[128, 128], - prefix="engine.layers.1.linear_attn.in_proj_qkvz", - weight=SimpleNamespace(shape=(5120, 4096)), - ) - - assert _is_sm70_fp8_exact_8k_prefill_layer(layer) - layer.prefix = "engine.layers.3.self_attn.qkv_proj" - layer.weight = SimpleNamespace(shape=(5120, 3584)) - assert _is_sm70_fp8_exact_8k_prefill_layer(layer) - - layer.tp_size = 2 - assert not _is_sm70_fp8_exact_8k_prefill_layer(layer) - layer.tp_size = 4 - layer.weight_block_size = [64, 128] - assert not _is_sm70_fp8_exact_8k_prefill_layer(layer) - layer.weight_block_size = [128, 128] - layer.weight = SimpleNamespace(shape=(5120, 4096)) - assert not _is_sm70_fp8_exact_8k_prefill_layer(layer) - - -def test_fp8_qpn8_shape_gate_uses_checkpoint_native_layout(): - layer = SimpleNamespace( - tp_size=4, - weight_block_size=[128, 128], - prefix="model.language_model.layers.1.mlp.gate_up_proj", - input_size_per_partition=5120, - output_size_per_partition=8704, - weight=SimpleNamespace(shape=(8704, 5120)), - ) - - assert _is_sm70_fp8_qpn8_layer(layer) - - layer.tp_size = 2 - assert not _is_sm70_fp8_qpn8_layer(layer) - layer.tp_size = 4 - layer.prefix = "model.language_model.layers.1.linear_attn.in_proj_qkvz" - layer.output_size_per_partition = 4096 - layer.weight = SimpleNamespace(shape=(4096, 5120)) - assert _is_sm70_fp8_qpn8_layer(layer) - layer.prefix = "model.language_model.layers.3.self_attn.qkv_proj" - layer.output_size_per_partition = 3584 - layer.weight = SimpleNamespace(shape=(3584, 5120)) - assert _is_sm70_fp8_qpn8_layer(layer) - layer.output_size_per_partition = 4096 - assert not _is_sm70_fp8_qpn8_layer(layer) - layer.output_size_per_partition = 3584 - layer.weight_block_size = [64, 128] - assert not _is_sm70_fp8_qpn8_layer(layer) - - -def test_fp8_qpn8_runtime_gate_uses_engine_contract_only(monkeypatch): - vllm_config = SimpleNamespace( - scheduler_config=SimpleNamespace(max_num_seqs=8), - speculative_config=None, - ) - monkeypatch.setattr( - "vllm.model_executor.layers.quantization.fp8.get_current_vllm_config", - lambda: vllm_config, - ) - - assert _is_sm70_fp8_qpn8_runtime_contract() - vllm_config.scheduler_config.max_num_seqs = 16 - assert not _is_sm70_fp8_qpn8_runtime_contract() - vllm_config.scheduler_config.max_num_seqs = 8 - vllm_config.speculative_config = object() - assert not _is_sm70_fp8_qpn8_runtime_contract() - monkeypatch.setenv("VLLM_SM70_FP8_QPN8", "1") - assert _is_sm70_fp8_qpn8_runtime_contract() def test_fp8_prefill_exact_dense_workspace_is_bounded(): @@ -517,51 +75,6 @@ def fake_empty(shape, *, dtype, device): _sm70_fp8_prefill_dense_workspaces.clear() -def test_fp8_prefill_visible_dense_mm_is_long_prefill_only(monkeypatch): - workspace = torch.empty(24, dtype=torch.float16) - weight = torch.empty((4, 6), dtype=torch.uint8) - scales = torch.empty((1, 6), dtype=torch.float16) - calls = [] - - def fake_dequantize(out, qweight, factors, group_size): - calls.append((qweight, factors, group_size)) - out.fill_(1) - - monkeypatch.setenv("VLLM_SM70_FP8_PREFILL_VISIBLE_DENSE_MM", "1") - monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 0) - monkeypatch.setattr( - "vllm.model_executor.layers.quantization.fp8.sm70_ops.fp8_sm70_dequantize_out", - fake_dequantize, - ) - envs.disable_envs_cache() - _sm70_fp8_prefill_dense_workspaces[(0, torch.float16)] = workspace - try: - long_prefill = _sm70_fp8_prefill_visible_dense_mm( - torch.ones((_SM70_FP8_PREFILL_DENSE_MIN_M, 4), dtype=torch.float16), - weight, - scales, - workspace.data_ptr(), - gated_silu=False, - min_prefill_m=_SM70_FP8_PREFILL_DENSE_MIN_M, - ) - tail = _sm70_fp8_prefill_visible_dense_mm( - torch.ones((1, 4), dtype=torch.float16), - weight, - scales, - workspace.data_ptr(), - gated_silu=False, - min_prefill_m=_SM70_FP8_PREFILL_DENSE_MIN_M, - ) - finally: - _sm70_fp8_prefill_dense_workspaces.clear() - envs.disable_envs_cache() - - assert long_prefill is not None - assert long_prefill.shape == (_SM70_FP8_PREFILL_DENSE_MIN_M, 6) - assert tail is None - assert calls == [(weight, scales, 128)] - - def test_fp8_prefill_dispatch_reaches_runtime_op_for_small_and_large_m(monkeypatch): calls = [] @@ -612,245 +125,3 @@ def fake_dispatch( False, ), ] - - -def test_fp8_prefill_prescaled_scales_only_reach_exact_8k_route(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_prefill_prescaled_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_prefill_prescaled_scales=prescaled_scales, - sm70_fp8_k_ld=4, - sm70_fp8_q_ld=6, - ) - method = SimpleNamespace() - - monkeypatch.setenv("VLLM_SM70_FP8_PREFILL_FAST_SELECTOR", "1") - monkeypatch.setenv("VLLM_SM70_FP8_PREFILL_PRESCALED", "1") - envs.disable_envs_cache() - try: - Fp8LinearMethod.apply(method, layer, torch.empty((1, 4), dtype=torch.float16)) - Fp8LinearMethod.apply( - method, - layer, - torch.empty((_SM70_FP8_EXACT_8K_PREFILL_M, 4), dtype=torch.float16), - ) - monkeypatch.setenv("VLLM_SM70_FP8_PREFILL_PRESCALED", "0") - envs.disable_envs_cache() - Fp8LinearMethod.apply( - method, - layer, - torch.empty((_SM70_FP8_EXACT_8K_PREFILL_M, 4), dtype=torch.float16), - ) - finally: - envs.disable_envs_cache() - - assert [(route, m) for route, m, _ in calls] == [ - ("default", 1), - ("prescaled", _SM70_FP8_EXACT_8K_PREFILL_M), - ("default", _SM70_FP8_EXACT_8K_PREFILL_M), - ] - assert calls[0][2] is normal_scales - assert calls[1][2] is prescaled_scales - 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 = [] - - def fake_dispatch( - out, workspace, input, codes, scales, split_k, nacc, prefetch, gated_silu - ): - calls.append( - ("dispatch", input.shape[0], workspace, split_k, nacc, prefetch, gated_silu) - ) - out.zero_() - - monkeypatch.setattr( - "vllm.model_executor.layers.quantization.fp8.sm70_ops." - "fp8_qpn8_dispatch_sm70_out", - fake_dispatch, - ) - layer = SimpleNamespace( - sm70_fp8_turbomind=True, - sm70_fp8_qpn8=True, - output_size_per_partition=6, - weight=torch.empty((4, 6), dtype=torch.uint8), - weight_scale_inv=torch.empty((1, 1), dtype=torch.float16), - sm70_fp8_qpn8_split_k=16, - sm70_fp8_qpn8_nacc=1, - sm70_fp8_qpn8_prefetch=False, - sm70_fp8_prefill_exact_dense_workspace_ptr=42, - ) - method = SimpleNamespace() - - for m in (1, 9): - output = Fp8LinearMethod.apply( - method, layer, torch.empty((m, 4), dtype=torch.float16) - ) - assert output.shape == (m, 6) - - assert calls == [ - ("dispatch", 1, 42, 16, 1, False, False), - ("dispatch", 9, 42, 16, 1, False, False), - ] - - -def test_fp8_qpn8_fused_gate_dispatches_without_intermediate(monkeypatch): - calls = [] - - def fake_dispatch( - out, workspace, input, codes, scales, split_k, nacc, prefetch, gated_silu - ): - calls.append( - ("dispatch", input.shape[0], workspace, split_k, nacc, prefetch, gated_silu) - ) - out.zero_() - - monkeypatch.setattr( - "vllm.model_executor.layers.quantization.fp8.sm70_ops." - "fp8_qpn8_dispatch_sm70_out", - fake_dispatch, - ) - layer = SimpleNamespace( - sm70_fp8_qpn8=True, - sm70_fp8_gated_silu=True, - output_size_per_partition=12, - weight=torch.empty((4, 12), dtype=torch.uint8), - weight_scale_inv=torch.empty((1, 1), dtype=torch.float16), - sm70_fp8_qpn8_gated_split_k=8, - sm70_fp8_qpn8_gated_nacc=2, - sm70_fp8_qpn8_gated_prefetch=True, - sm70_fp8_prefill_exact_dense_workspace_ptr=42, - ) - method = SimpleNamespace() - - for m in (8, 16): - output = Fp8LinearMethod.apply_fused_silu_and_mul( - method, layer, torch.empty((m, 4), dtype=torch.float16) - ) - assert output is not None - assert output.shape == (m, 6) - - assert calls == [ - ("dispatch", 8, 42, 8, 2, True, True), - ("dispatch", 16, 42, 8, 2, True, True), - ] diff --git a/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py b/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py deleted file mode 100644 index 065b447cc9..0000000000 --- a/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py +++ /dev/null @@ -1,192 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Narrow SM70 admission and shape gates for ModelOpt mixed NVFP4.""" - -from types import SimpleNamespace -from unittest.mock import patch - -import pytest -import torch - -from vllm import envs -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 ( - ModelOptFp8Config, - ModelOptMixedPrecisionConfig, - ModelOptNvFp4Config, -) -from vllm.model_executor.layers.quantization.nvfp4_sm70_moe import ( - ModelOptNvFp4SM70MoEMethod, - _prepare_compact_slot_groups, - _validate_weight_layout, - validate_nvfp4_sm70_moe_contract, -) - - -def _mixed_config() -> ModelOptMixedPrecisionConfig: - fp8 = ModelOptFp8Config("FP8", True, None, []) - nvfp4 = ModelOptNvFp4Config( - quant_method="NVFP4", - is_checkpoint_nvfp4_serialized=True, - ) - w4a16 = ModelOptNvFp4Config( - quant_method="W4A16_NVFP4", - is_checkpoint_nvfp4_serialized=True, - ) - return ModelOptMixedPrecisionConfig( - kv_cache_quant_method=None, - exclude_modules=[], - quantized_layers={"model.layers.0.mlp.experts": {"quant_algo": "W4A16_NVFP4"}}, - fp8_config=fp8, - nvfp4_config=nvfp4, - w4a16_nvfp4_config=w4a16, - ) - - -def _moe_contract(**overrides): - values = { - "num_experts": 256, - "experts_per_token": 8, - "hidden_dim": 2048, - "intermediate_size_per_partition": 128, - "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), - patch.object(sm70_tm, "use_turbomind", side_effect=[True, True]), - ): - assert ModelOptMixedPrecisionConfig.get_min_capability() == 70 - - with ( - patch.object(sm70_tm, "is_exact_sm70_cuda_platform", return_value=True), - patch.object(sm70_tm, "use_turbomind", side_effect=[True, False]), - ): - assert ModelOptMixedPrecisionConfig.get_min_capability() == 89 - - with patch.object(sm70_tm, "is_exact_sm70_cuda_platform", return_value=False): - assert ModelOptMixedPrecisionConfig.get_min_capability() == 89 - - -def test_nvfp4_grouped_prefill_defaults_on_and_can_be_disabled(monkeypatch): - monkeypatch.delenv("VLLM_SM70_NVFP4_MOE_GROUPED_PREFILL", raising=False) - assert envs.VLLM_SM70_NVFP4_MOE_GROUPED_PREFILL - - monkeypatch.setenv("VLLM_SM70_NVFP4_MOE_GROUPED_PREFILL", "0") - assert not envs.VLLM_SM70_NVFP4_MOE_GROUPED_PREFILL - - -@pytest.mark.parametrize( - ("field", "value"), - [ - ("num_experts", 128), - ("experts_per_token", 6), - ("hidden_dim", 4096), - ("intermediate_size_per_partition", 96), - ("tp_size", 2), - ], -) -def test_nvfp4_moe_contract_rejects_non_qwen36_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_rejects_shape_consistent_unvalidated_tp8(): - with pytest.raises(NotImplementedError, match="tensor parallel"): - validate_nvfp4_sm70_moe_contract( - _moe_contract(tp_size=8, intermediate_size_per_partition=64) - ) - - -def _meta_layer() -> SimpleNamespace: - experts, hidden, intermediate = 256, 2048, 128 - return SimpleNamespace( - local_num_experts=experts, - moe_config=_moe_contract(), - w13_weight=torch.empty( - experts, 2 * intermediate, hidden // 2, dtype=torch.uint8, device="meta" - ), - w13_weight_scale=torch.empty( - experts, - 2 * intermediate, - hidden // 16, - dtype=torch.float8_e4m3fn, - device="meta", - ), - w13_weight_scale_2=torch.empty(experts, 2, device="meta"), - w2_weight=torch.empty( - experts, hidden, intermediate // 2, dtype=torch.uint8, device="meta" - ), - w2_weight_scale=torch.empty( - experts, - hidden, - intermediate // 16, - dtype=torch.float8_e4m3fn, - device="meta", - ), - w2_weight_scale_2=torch.empty(experts, device="meta"), - ) - - -def test_nvfp4_moe_weight_layout_is_explicit(): - layer = _meta_layer() - _validate_weight_layout(layer) - - layer.w2_weight = torch.empty(256, 2048, 32, dtype=torch.uint8, device="meta") - with pytest.raises(ValueError, match="w2_weight"): - _validate_weight_layout(layer) - - -def test_nvfp4_sm70_moe_owns_routing_without_generic_modular_wrapper(): - method = ModelOptNvFp4SM70MoEMethod( - quant_config=ModelOptNvFp4Config( - quant_method="W4A16_NVFP4", - is_checkpoint_nvfp4_serialized=True, - ), - moe_config=_moe_contract(), - ) - - assert method.maybe_make_prepare_finalize() is None - - -@pytest.mark.skipif( - 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)) -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 - ) % 256 - compact_offsets = torch.empty(total_slots + 1, dtype=torch.int32, device="cuda") - active_expert_ids = torch.empty(total_slots, dtype=torch.int32, device="cuda") - - _prepare_compact_slot_groups(sorted_expert_ids, compact_offsets, active_expert_ids) - - assert torch.equal( - compact_offsets.cpu(), - torch.arange(total_slots + 1, dtype=torch.int32, device="cpu"), - ) - assert torch.equal(active_expert_ids.cpu(), sorted_expert_ids.cpu()) - - -def test_mixed_w4a16_moe_requires_turbomind_on_sm70(): - config = _mixed_config() - - class FakeRoutedExperts: - moe_config = _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=False), - pytest.raises(NotImplementedError, match="TurboMind"), - ): - config.get_quant_method(FakeRoutedExperts(), "model.layers.0.mlp.experts") diff --git a/tests/quantization/test_sm70_modelopt_nvfp4.py b/tests/quantization/test_sm70_modelopt_nvfp4.py deleted file mode 100644 index 98a7db2be3..0000000000 --- a/tests/quantization/test_sm70_modelopt_nvfp4.py +++ /dev/null @@ -1,255 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Exact-SM70 ModelOpt NVFP4 admission: capability, provenance, routing. - -These tests exist so PR 114/166 cannot recur: -- the class-wide gate is not blindly lowered to 70 -- scale convention is ModelOpt provenance (amax/2688), never max(|scale|) < 1 -- TurboMind prepare is exact SM70 only -""" - -from unittest.mock import MagicMock, patch - -import torch -from torch.nn.parameter import Parameter - -from vllm.config import KernelConfig, VllmConfig, set_current_vllm_config -from vllm.model_executor.layers.quantization import sm70_turbomind as sm70_tm -from vllm.model_executor.layers.quantization.modelopt import ( - ModelOptNvFp4Config, - ModelOptNvFp4LinearMethod, - ModelOptNvFp4W4A16LinearMethod, -) - - -def _modelopt_dequant_oracle( - packed_nibble: float, block_scale: float, weight_global_scale: float -) -> float: - """Independent ModelOpt NVFP4 dequant: nibble * block * (amax/2688).""" - return packed_nibble * block_scale * weight_global_scale - - -def _ct_dequant_oracle( - packed_nibble: float, block_scale: float, disk_global: float -) -> float: - """Independent CT NVFP4 dequant: nibble * block / disk_global.""" - return packed_nibble * block_scale / disk_global - - -def _rejected_magnitude_heuristic_dequant( - packed_nibble: float, block_scale: float, disk_global: float -) -> float: - """PR 166 heuristic: max(|block|) < 1 => drop the global factor.""" - global_scale = 1.0 if abs(block_scale) < 1 else 1.0 / disk_global - return packed_nibble * block_scale * global_scale - - -def test_modelopt_nvfp4_min_capability_requires_proven_backend(monkeypatch): - monkeypatch.delenv("VLLM_USE_NVFP4_CT_EMULATIONS", raising=False) - monkeypatch.delenv("VLLM_NVFP4_GEMM_BACKEND", raising=False) - monkeypatch.delenv("VLLM_SM70_QUANT_BACKEND", raising=False) - - with set_current_vllm_config( - VllmConfig(kernel_config=KernelConfig(linear_backend="auto")) - ): - assert ModelOptNvFp4Config.get_min_capability() == 70 - - monkeypatch.setenv("VLLM_SM70_NVFP4_TURBOMIND", "0") - with set_current_vllm_config( - VllmConfig(kernel_config=KernelConfig(linear_backend="auto")) - ): - assert ModelOptNvFp4Config.get_min_capability() == 75 - - monkeypatch.delenv("VLLM_SM70_NVFP4_TURBOMIND", raising=False) - monkeypatch.setenv("VLLM_SM70_QUANT_BACKEND", "marlin") - with set_current_vllm_config( - VllmConfig(kernel_config=KernelConfig(linear_backend="auto")) - ): - assert ModelOptNvFp4Config.get_min_capability() == 70 - - monkeypatch.setenv("VLLM_SM70_QUANT_BACKEND", "turbomind") - with set_current_vllm_config( - VllmConfig(kernel_config=KernelConfig(linear_backend="auto")) - ): - assert ModelOptNvFp4Config.get_min_capability() == 70 - - monkeypatch.setenv("VLLM_SM70_NVFP4_TURBOMIND", "0") - monkeypatch.setenv("VLLM_SM70_QUANT_BACKEND", "auto") - with set_current_vllm_config( - VllmConfig(kernel_config=KernelConfig(linear_backend="emulation")) - ): - assert ModelOptNvFp4Config.get_min_capability() == 70 - - -def test_modelopt_scale_provenance_is_not_the_pr166_heuristic(): - # yangzhuxinyzx 100x counterexample: valid CT encoding, FP8 block 0.5, - # disk global 100, nibble 2. Correct CT dequant is 0.01. The rejected - # magnitude rule rewrites it to 1.0. - nibble, block, disk_global = 2.0, 0.5, 100.0 - assert _ct_dequant_oracle(nibble, block, disk_global) == 0.01 - assert _rejected_magnitude_heuristic_dequant(nibble, block, disk_global) == 1.0 - - # ModelOpt provenance is explicit: global is already amax/2688. - # Same numeric tensors must keep the on-disk global (no rewrite). - modelopt_global = 100.0 - assert _modelopt_dequant_oracle(nibble, block, modelopt_global) == 100.0 - assert abs(block) < 1 - assert _modelopt_dequant_oracle(nibble, block, modelopt_global) != 1.0 - - -def _w4a16_layer(global_scale: float, block_scale: float) -> torch.nn.Module: - layer = torch.nn.Module() - layer.weight = Parameter( - torch.zeros((16, 8), dtype=torch.uint8), - requires_grad=False, - ) - layer.weight_scale = Parameter( - torch.full((16, 1), block_scale, dtype=torch.float8_e4m3fn), - requires_grad=False, - ) - layer.weight_scale_2 = Parameter( - torch.tensor([global_scale], dtype=torch.float32), - requires_grad=False, - ) - layer.input_scale = Parameter( - torch.tensor([1.0], dtype=torch.float32), - requires_grad=False, - ) - return layer - - -def test_w4a16_keeps_modelopt_global_scale_without_reciprocation_or_heuristic(): - layer = _w4a16_layer(global_scale=100.0, block_scale=0.5) - fake_kernel = MagicMock() - - with patch( - "vllm.model_executor.layers.quantization.modelopt.MarlinNvFp4LinearKernel", - return_value=fake_kernel, - ): - method = ModelOptNvFp4W4A16LinearMethod( - ModelOptNvFp4Config( - quant_method="W4A16_NVFP4", - is_checkpoint_nvfp4_serialized=True, - ) - ) - - with ( - patch.object(sm70_tm, "should_prepare_turbomind", return_value=False), - patch.object( - sm70_tm, - "prepare_nvfp4_linear", - side_effect=AssertionError("prepare must not run"), - ), - ): - method.process_weights_after_loading(layer) - - # Provenance: keep on-disk ModelOpt global. Do not invert. Do not - # force 1.0 because max(|weight_scale|) < 1. - assert torch.equal(layer.weight_global_scale.cpu(), torch.tensor(100.0)) - assert not hasattr(layer, "weight_scale_2") - fake_kernel.process_weights_after_loading.assert_called_once_with(layer) - - -def test_w4a16_turbomind_prepare_only_on_exact_sm70(): - layer = _w4a16_layer(global_scale=0.25, block_scale=1.0) - fake_kernel = MagicMock() - - with patch( - "vllm.model_executor.layers.quantization.modelopt.MarlinNvFp4LinearKernel", - return_value=fake_kernel, - ): - method = ModelOptNvFp4W4A16LinearMethod( - ModelOptNvFp4Config( - quant_method="W4A16_NVFP4", - is_checkpoint_nvfp4_serialized=True, - ) - ) - - def fake_prepare(prepared: torch.nn.Module) -> None: - setattr(prepared, sm70_tm.STATE_ATTR, object()) - - with ( - patch.object(sm70_tm, "should_prepare_turbomind", return_value=True), - patch.object(sm70_tm, "prepare_nvfp4_linear", side_effect=fake_prepare), - ): - method.process_weights_after_loading(layer) - - assert sm70_tm.has_prepared_linear(layer) - assert layer.weight.numel() == 0 - assert layer.weight_scale.numel() == 0 - fake_kernel.process_weights_after_loading.assert_not_called() - - -def test_w4a4_turbomind_prepare_skipped_off_sm70(): - layer = torch.nn.Module() - layer.input_scale = Parameter( - torch.ones(1, dtype=torch.float32), - requires_grad=False, - ) - layer.weight_scale_2 = Parameter( - torch.tensor([0.25], dtype=torch.float32), - requires_grad=False, - ) - layer.weight = Parameter( - torch.zeros((16, 8), dtype=torch.uint8), - requires_grad=False, - ) - layer.weight_scale = Parameter( - torch.ones((16, 1), dtype=torch.float8_e4m3fn), - requires_grad=False, - ) - fake_kernel = MagicMock() - - with patch( - "vllm.model_executor.layers.quantization.modelopt.init_nvfp4_linear_kernel", - return_value=fake_kernel, - ): - method = ModelOptNvFp4LinearMethod( - ModelOptNvFp4Config( - quant_method="NVFP4", - is_checkpoint_nvfp4_serialized=True, - ) - ) - - with ( - patch.object(sm70_tm, "should_prepare_turbomind", return_value=False), - patch.object( - sm70_tm, - "prepare_nvfp4_linear", - side_effect=AssertionError("prepare must not run off SM70"), - ), - ): - method.process_weights_after_loading(layer) - - fake_kernel.process_weights_after_loading.assert_called_once_with(layer) - assert not sm70_tm.has_prepared_linear(layer) - - -def test_apply_uses_prepared_turbomind_path(): - layer = torch.nn.Module() - fake_kernel = MagicMock() - fake_out = torch.zeros((1, 4), dtype=torch.float16) - - with patch( - "vllm.model_executor.layers.quantization.modelopt.MarlinNvFp4LinearKernel", - return_value=fake_kernel, - ): - method = ModelOptNvFp4W4A16LinearMethod( - ModelOptNvFp4Config( - quant_method="W4A16_NVFP4", - is_checkpoint_nvfp4_serialized=True, - ) - ) - - x = torch.ones((1, 16), dtype=torch.float16) - with ( - patch.object(sm70_tm, "has_prepared_linear", return_value=True), - patch.object( - sm70_tm, "apply_prepared_linear", return_value=fake_out - ) as apply_tm, - ): - out = method.apply(layer, x) - - assert out is fake_out - apply_tm.assert_called_once() - fake_kernel.apply_weights.assert_not_called() diff --git a/tests/quantization/test_sm70_modelopt_nvfp4_kernel_oracle.py b/tests/quantization/test_sm70_modelopt_nvfp4_kernel_oracle.py deleted file mode 100644 index ddffa37532..0000000000 --- a/tests/quantization/test_sm70_modelopt_nvfp4_kernel_oracle.py +++ /dev/null @@ -1,148 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Exact-SM70 ModelOpt NVFP4: real kernel path vs an independent fp32 reference. - -Runs only on a Volta (7,0) CUDA device. Builds a synthetic linear in the exact -on-disk ModelOpt layout (``weight`` uint8 E2M1 pairs, ``weight_scale`` fp8 e4m3 -per 16-group, ``weight_scale_2`` fp32 = amax/(6*448)), loads it through the real -ModelOpt W4A4 / W4A16 quant methods (so the SM70 TurboMind prepare/apply route is -exercised end to end), and compares against a pure-torch dequant + fp32 GEMM that -shares no code with vLLM's scale handling. A convention error (reciprocated or -dropped global, PR 166 class) shows up as an output-norm ratio of ~2688x or -~1/2688x and fails loudly; kernel arithmetic error shows up as low cosine. -""" - -import os - -import pytest -import torch - -from vllm.config import VllmConfig, set_current_vllm_config -from vllm.distributed import ( - ensure_model_parallel_initialized, - init_distributed_environment, -) -from vllm.model_executor.layers.linear import ReplicatedLinear -from vllm.model_executor.layers.quantization import sm70_turbomind as sm70_tm -from vllm.model_executor.layers.quantization.modelopt import ModelOptNvFp4Config - -E2M1 = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32) - - -def _is_exact_sm70() -> bool: - return torch.cuda.is_available() and torch.cuda.get_device_capability(0) == (7, 0) - - -pytestmark = pytest.mark.skipif( - not _is_exact_sm70(), reason="requires an exact SM70 (Volta) CUDA device" -) - - -def _decode_e2m1(packed: torch.Tensor) -> torch.Tensor: - lo = packed & 0x0F - hi = (packed >> 4) & 0x0F - nib = torch.stack([lo, hi], dim=-1).reshape(packed.shape[0], -1).long() - return E2M1[nib & 0x7] * torch.where((nib & 0x8) != 0, -1.0, 1.0) - - -def _make_modelopt_layer(n: int, k: int, seed: int): - """Synthetic ModelOpt NVFP4 tensors: nibbles, fp8 block scales, amax global.""" - g = torch.Generator().manual_seed(seed) - packed = torch.randint(0, 256, (n, k // 2), generator=g, dtype=torch.uint8) - amax = 0.35 - weight_scale_2 = torch.tensor(amax / (6.0 * 448.0), dtype=torch.float32) - # block scales in fp8: values so that block*global lands in a realistic weight range - block = torch.rand(n, k // 16, generator=g) * 400.0 + 20.0 - weight_scale = block.to(torch.float8_e4m3fn) - w_ref = ( - _decode_e2m1(packed) - * weight_scale.to(torch.float32).repeat_interleave(16, dim=1) - * weight_scale_2 - ) - return packed, weight_scale, weight_scale_2, w_ref - - -def _build_layer(quant_method: str, n: int, k: int, tensors, device): - packed, weight_scale, weight_scale_2, _ = tensors - os.environ.setdefault("MASTER_ADDR", "127.0.0.1") - os.environ.setdefault("MASTER_PORT", "29581") - with set_current_vllm_config(VllmConfig()): - try: - init_distributed_environment( - world_size=1, - rank=0, - local_rank=0, - distributed_init_method="env://", - backend="gloo", - ) - ensure_model_parallel_initialized(1, 1) - except Exception: - pass # already initialised by an earlier test in this process - qc = ModelOptNvFp4Config( - is_checkpoint_nvfp4_serialized=True, - kv_cache_quant_algo=None, - exclude_modules=[], - group_size=16, - quant_method=quant_method, - ) - lin = ReplicatedLinear( - k, - n, - bias=False, - quant_config=qc, - prefix="oracle", - params_dtype=torch.float16, - ).to(device) - for pname, param in lin.named_parameters(): - src = { - "weight": packed, - "weight_scale": weight_scale, - "weight_scale_2": weight_scale_2.reshape(1), - "input_scale": torch.tensor([1.0]), - }.get(pname) - if src is None: - continue - loader = getattr(param, "weight_loader", None) - if loader is not None: - loader( - param, src if param.dtype == torch.uint8 else src.to(param.dtype) - ) - else: - param.data.copy_(src.reshape(param.shape)) - lin.quant_method.process_weights_after_loading(lin) - return lin - - -@pytest.mark.parametrize("quant_method", ["NVFP4", "W4A16_NVFP4"]) -@pytest.mark.parametrize("n,k", [(512, 1024), (1024, 512)]) -def test_sm70_modelopt_nvfp4_kernel_matches_independent_reference( - monkeypatch, quant_method, n, k -): - monkeypatch.setenv("VLLM_SM70_NVFP4_TURBOMIND", "1") - monkeypatch.delenv("VLLM_SM70_QUANT_BACKEND", raising=False) - assert ModelOptNvFp4Config.get_min_capability() == 70 - - device = torch.device("cuda:0") - tensors = _make_modelopt_layer(n, k, seed=n + k) - w_ref = tensors[3] - lin = _build_layer(quant_method, n, k, tensors, device) - assert sm70_tm.has_prepared_linear(lin), "expected the exact-SM70 TurboMind route" - - torch.manual_seed(0) - x = (torch.randn(64, k) * 0.5).to(torch.float16) - y_ref = x.float() @ w_ref.T # independent fp32 reference - y_floor = (x @ w_ref.to(torch.float16).T).float() # fp16 accumulation floor - with torch.no_grad(): - y_dut, _ = lin(x.to(device)) - y_dut = y_dut.float().cpu() - - ratio = float(y_dut.norm() / y_ref.norm()) - cos = float( - torch.nn.functional.cosine_similarity(y_dut.flatten(), y_ref.flatten(), dim=0) - ) - rel = float((y_dut - y_ref).norm() / y_ref.norm()) - rel_floor = float((y_floor - y_ref).norm() / y_ref.norm()) - # A reciprocated/dropped global would put ratio at ~2688 or ~1/2688. - assert abs(ratio - 1.0) < 0.02, f"output scale ratio {ratio} (convention error?)" - assert cos > 0.999, f"cosine {cos} (kernel arithmetic error?)" - assert rel < max(5.0 * rel_floor, 5e-3), f"rel err {rel} vs fp16 floor {rel_floor}" diff --git a/tests/quantization/test_sm70_mxfp4_moe.py b/tests/quantization/test_sm70_mxfp4_moe.py index 9e3e9fa18f..ebb2e1785d 100644 --- a/tests/quantization/test_sm70_mxfp4_moe.py +++ b/tests/quantization/test_sm70_mxfp4_moe.py @@ -6,7 +6,6 @@ import pytest import torch -from vllm import _sm70_ops as sm70_ops from vllm import envs from vllm.model_executor.layers.fused_moe import ( FusedMoEConfig, @@ -22,7 +21,6 @@ from vllm.model_executor.layers.quantization.mxfp4_sm70_moe import ( Mxfp4SM70MoEMethod, _compact_mxfp4_active_experts, - _select_mxfp4_direct_order_offsets, _select_mxfp4_stage_dispatch, validate_mxfp4_sm70_moe_contract, validate_mxfp4_sm70_moe_weight_layout, @@ -211,21 +209,6 @@ def test_mxfp4_sm70_b1_dispatch_selects_six_runtime_experts(monkeypatch): assert count == 6 -def test_mxfp4_sm70_direct_order_ignores_compacted_offsets(): - buffers = { - # 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 - ), - "slot_expert_offsets": torch.arange(7, dtype=torch.int32), - } - - offsets = _select_mxfp4_direct_order_offsets(buffers) - - assert offsets is buffers["slot_expert_offsets"] - torch.testing.assert_close(offsets, torch.arange(7, dtype=torch.int32)) - - def test_mxfp4_sm70_b1_dispatch_rejects_incompatible_permute_fastpath( monkeypatch, ): @@ -325,9 +308,6 @@ def test_mxfp4_sm70_m8_grouped_dispatch_keeps_one_row_slots(monkeypatch): } monkeypatch.setattr(mxfp4_moe, "_mxfp4_active_expert_max_tokens", lambda: 8) monkeypatch.setattr(mxfp4_moe, "_mxfp4_grouped_m8_enabled", lambda: True) - monkeypatch.setattr( - mxfp4_moe, "_mxfp4_grouped_m8_expert_rows_enabled", lambda: False - ) offsets, expert_ids, count = _select_mxfp4_stage_dispatch( buffers, @@ -341,72 +321,6 @@ def test_mxfp4_sm70_m8_grouped_dispatch_keeps_one_row_slots(monkeypatch): assert count == 48 -def test_mxfp4_sm70_m8_expert_grouped_dispatch_keeps_real_segments(monkeypatch): - buffers = { - "compact_expert_offsets": torch.arange(49, dtype=torch.int32) * 2, - "slot_expert_offsets": torch.arange(49, dtype=torch.int32), - "permuted_experts_id": torch.arange(48, dtype=torch.int32).flip(0), - "active_expert_ids": torch.arange(48, dtype=torch.int32), - "expert_offsets": torch.arange(257, dtype=torch.int32), - "dense_expert_ids": torch.arange(256, dtype=torch.int32), - } - monkeypatch.setattr(mxfp4_moe, "_mxfp4_active_expert_max_tokens", lambda: 8) - monkeypatch.setattr(mxfp4_moe, "_mxfp4_grouped_m8_enabled", lambda: True) - monkeypatch.setattr( - mxfp4_moe, "_mxfp4_grouped_m8_expert_rows_enabled", lambda: True - ) - - offsets, expert_ids, count = _select_mxfp4_stage_dispatch( - buffers, - num_tokens=8, - num_experts=256, - fully_replicated_experts=True, - ) - - assert offsets is buffers["compact_expert_offsets"] - assert expert_ids is buffers["active_expert_ids"] - assert count == 48 - - -@pytest.mark.parametrize("expert_rows", [False, True]) -def test_mxfp4_sm70_m5_grouped_dispatch(expert_rows, monkeypatch): - buffers = { - "compact_expert_offsets": torch.arange(31, dtype=torch.int32) * 2, - "slot_expert_offsets": torch.arange(31, dtype=torch.int32), - "permuted_experts_id": torch.arange(30, dtype=torch.int32).flip(0), - "active_expert_ids": torch.arange(30, dtype=torch.int32), - "expert_offsets": torch.arange(257, dtype=torch.int32), - "dense_expert_ids": torch.arange(256, dtype=torch.int32), - } - monkeypatch.setattr(mxfp4_moe, "_mxfp4_active_expert_max_tokens", lambda: 8) - monkeypatch.setattr(mxfp4_moe, "_mxfp4_grouped_m8_enabled", lambda: False) - monkeypatch.setattr(mxfp4_moe, "_mxfp4_grouped_verifier_enabled", lambda: True) - monkeypatch.setattr( - mxfp4_moe, - "_mxfp4_grouped_m8_expert_rows_enabled", - lambda: expert_rows, - ) - - offsets, expert_ids, count = _select_mxfp4_stage_dispatch( - buffers, - num_tokens=5, - num_experts=256, - fully_replicated_experts=True, - ) - - expected_offsets = ( - buffers["compact_expert_offsets"] - if expert_rows - else buffers["slot_expert_offsets"] - ) - expected_ids = ( - buffers["active_expert_ids"] if expert_rows else buffers["permuted_experts_id"] - ) - assert offsets is expected_offsets - assert expert_ids is expected_ids - assert count == 30 - - @pytest.mark.skipif( not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0), reason="requires NVIDIA V100/SM70", @@ -449,117 +363,6 @@ def test_mxfp4_sm70_active_expert_compaction_replays_dynamic_routes(): ) -@pytest.mark.parametrize(("k", "n"), [(4096, 512), (4096, 1024), (512, 4096)]) -@pytest.mark.skipif( - not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0), - reason="requires NVIDIA V100/SM70", -) -def test_mxfp4_sm70_tp4_compact_shapes_match_per_expert( - monkeypatch: pytest.MonkeyPatch, - k: int, - n: int, -): - num_experts = 6 - group_size = 32 - torch.manual_seed(20260825) - packed = torch.randint(0, 16, (k, n), dtype=torch.uint8, device="cuda") - scales = torch.full((k // group_size, n), 127, dtype=torch.uint8, device="cuda") - weight, scale, meta = sm70_ops.mxfp4_sm70_prepare(packed, scales, group_size) - weights = weight.unsqueeze(0).repeat(num_experts, *([1] * weight.ndim)) - expert_scales = scale.unsqueeze(0).repeat(num_experts, *([1] * scale.ndim)) - ptrs_w, ptrs_s = torch.ops._C.awq_moe_build_strided_ptrs( - weights, - expert_scales, - int(meta[0].item()), - int(meta[1].item()), - num_experts, - ) - x = torch.randn((1, k), dtype=torch.float16, device="cuda") * 0.01 - grouped_input = x.expand(num_experts, -1).contiguous() - offsets = torch.arange(num_experts + 1, dtype=torch.int32, device="cuda") - expert_ids = torch.arange(num_experts, dtype=torch.int32, device="cuda") - reference = torch.empty((num_experts, n), dtype=torch.float16, device="cuda") - actual = torch.empty_like(reference) - - monkeypatch.setenv("VLLM_SM70_MXFP4_MOE_COMPACT_GROUPED_DECODE", "0") - torch.ops._C.mxfp4_moe_dense_stage_sm70_out( - reference, - grouped_input, - offsets, - expert_ids, - ptrs_w, - ptrs_s, - num_experts, - k, - n, - group_size, - ) - monkeypatch.setenv("VLLM_SM70_MXFP4_MOE_COMPACT_GROUPED_DECODE", "1") - torch.ops._C.mxfp4_moe_dense_stage_sm70_out( - actual, - grouped_input, - offsets, - expert_ids, - ptrs_w, - ptrs_s, - num_experts, - k, - n, - group_size, - ) - torch.accelerator.synchronize() - torch.testing.assert_close(actual, reference, rtol=0.0, atol=0.0) - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - torch.ops._C.mxfp4_moe_dense_stage_sm70_out( - actual, - grouped_input, - offsets, - expert_ids, - ptrs_w, - ptrs_s, - num_experts, - k, - n, - group_size, - ) - graph.replay() - torch.accelerator.synchronize() - torch.testing.assert_close(actual, reference, rtol=0.0, atol=0.0) - - if k == 4096: - direct = torch.empty_like(actual) - legacy_direct = torch.empty_like(actual) - direct_input = torch.empty_like(grouped_input) - direct_offsets = torch.empty_like(offsets) - inverse_indices = torch.empty(num_experts, dtype=torch.int32, device="cuda") - direct_expert_ids = torch.empty_like(expert_ids) - for broadcast, output in ((False, legacy_direct), (True, direct)): - monkeypatch.setenv( - "VLLM_SM70_MXFP4_MOE_BROADCAST_INPUT_DECODE", - "1" if broadcast else "0", - ) - torch.ops._C.mxfp4_moe_single_token_prepare_w13_sm70_out( - output, - direct_input, - x, - expert_ids.view(1, num_experts), - ptrs_w, - ptrs_s, - direct_offsets, - inverse_indices, - direct_expert_ids, - k, - n, - group_size, - k, - ) - torch.accelerator.synchronize() - torch.testing.assert_close(legacy_direct, reference, rtol=0.0, atol=0.0) - torch.testing.assert_close(direct, reference, rtol=0.0, atol=0.0) - - def test_mxfp4_sm70_post_load_reads_bias_from_method_config(monkeypatch): for op_name in ( "mxfp4_sm70_prepare", diff --git a/tests/quantization/test_sm70_nvfp4_qpn2.py b/tests/quantization/test_sm70_nvfp4_qpn2.py deleted file mode 100644 index 84f8174158..0000000000 --- a/tests/quantization/test_sm70_nvfp4_qpn2.py +++ /dev/null @@ -1,143 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from types import SimpleNamespace - -import torch -from torch.nn.parameter import Parameter - -import vllm.envs as envs -from vllm.model_executor.layers.quantization import sm70_turbomind as sm70_tm -from vllm.model_executor.layers.quantization.compressed_tensors.schemes import ( - compressed_tensors_w4a4_nvfp4 as nvfp4_scheme, -) -from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_w4a4_nvfp4 import ( # noqa: E501 - CompressedTensorsW4A4Fp4, -) - - -def test_nvfp4_qpn2_is_default_off_with_explicit_on(monkeypatch): - monkeypatch.delenv("VLLM_SM70_NVFP4_QPN2", raising=False) - envs.disable_envs_cache() - try: - assert not envs.VLLM_SM70_NVFP4_QPN2 - monkeypatch.setenv("VLLM_SM70_NVFP4_QPN2", "1") - envs.disable_envs_cache() - assert envs.VLLM_SM70_NVFP4_QPN2 - finally: - envs.disable_envs_cache() - - -def test_nvfp4_qpn2_shape_gate_is_exact_tp4(): - layer = SimpleNamespace( - tp_size=4, - prefix="model.language_model.layers.0.mlp.gate_up_proj", - input_size_per_partition=5120, - output_size_per_partition=8704, - weight=SimpleNamespace(shape=(8704, 2560)), - ) - assert nvfp4_scheme._is_qpn2_layer(layer) - - layer.tp_size = 2 - assert not nvfp4_scheme._is_qpn2_layer(layer) - layer.tp_size = 4 - layer.prefix = "model.language_model.layers.0.self_attn.qkv_proj" - assert not nvfp4_scheme._is_qpn2_layer(layer) - layer.prefix = "model.language_model.layers.0.mlp.down_proj" - layer.input_size_per_partition = 4352 - layer.output_size_per_partition = 5120 - layer.weight = SimpleNamespace(shape=(5120, 2176)) - assert nvfp4_scheme._is_qpn2_layer(layer) - layer.input_size_per_partition = 4096 - assert not nvfp4_scheme._is_qpn2_layer(layer) - - -def _make_small_layer() -> torch.nn.Module: - layer = torch.nn.Module() - layer.prefix = "model.language_model.layers.0.mlp.gate_up_proj" - layer.tp_size = 4 - layer.input_size_per_partition = 64 - layer.output_size_per_partition = 64 - layer.weight_packed = Parameter( - torch.zeros((64, 32), dtype=torch.uint8), requires_grad=False - ) - layer.weight_scale = Parameter( - torch.ones((64, 4), dtype=torch.float8_e4m3fn), requires_grad=False - ) - layer.weight_global_scale = Parameter( - torch.tensor([2.0, 1.0], dtype=torch.float32), requires_grad=False - ) - layer.input_global_scale = Parameter( - torch.tensor([4.0, 4.0], dtype=torch.float32), requires_grad=False - ) - return layer - - -def test_nvfp4_qpn2_prepare_and_dispatch_contract(monkeypatch): - monkeypatch.setenv("VLLM_SM70_NVFP4_QPN2", "1") - envs.disable_envs_cache() - layer = _make_small_layer() - calls = [] - - monkeypatch.setattr(nvfp4_scheme.sm70_tm, "use_turbomind", lambda enabled: True) - scheme = CompressedTensorsW4A4Fp4() - monkeypatch.setattr( - nvfp4_scheme.sm70_tm, - "should_prepare_turbomind", - lambda tensor, enabled: enabled, - ) - monkeypatch.setattr(nvfp4_scheme, "_is_qpn2_layer", lambda layer: True) - monkeypatch.setattr(nvfp4_scheme, "_missing_qpn2_ops", lambda: []) - 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( - nvfp4_scheme.sm70_ops, - "nvfp4_qpn2_prepare_sm70", - lambda weight, scales: ( - torch.empty_like(weight), - torch.empty(scales.shape, dtype=torch.uint8), - ), - ) - - def fake_prepare(prepared_layer, *, interleave_gated_silu=False): - assert not interleave_gated_silu - state = sm70_tm.SM70TurboMindLinearState( - weight=torch.empty((1,), dtype=torch.int32), - scales=torch.empty((1,), dtype=torch.float16), - group_size=16, - k_ld=64, - q_ld=64, - output_size=64, - op_kind="nvfp4", - ) - setattr(prepared_layer, sm70_tm.STATE_ATTR, state) - - monkeypatch.setattr(nvfp4_scheme.sm70_tm, "prepare_nvfp4_linear", fake_prepare) - - def fake_dispatch(*args): - calls.append(args) - args[0].fill_(3) - - monkeypatch.setattr( - nvfp4_scheme.sm70_ops, "nvfp4_qpn2_dispatch_sm70_out", fake_dispatch - ) - - 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.weight.numel() == 0 - assert layer.weight_scale.numel() == 0 - - x = torch.ones((8, 64), dtype=torch.float16) - raw = scheme.apply_weights(layer, x) - fused = scheme.apply_fused_silu_and_mul(layer, x) - assert raw.shape == (8, 64) - assert fused is not None and fused.shape == (8, 32) - assert torch.equal(raw, torch.full_like(raw, 3)) - assert torch.equal(fused, torch.full_like(fused, 3)) - assert calls[0][-1] is False - assert calls[1][-1] is True - finally: - envs.disable_envs_cache() diff --git a/tests/quantization/test_sm70_turbomind_adapter.py b/tests/quantization/test_sm70_turbomind_adapter.py index c04d961caf..af022393fe 100644 --- a/tests/quantization/test_sm70_turbomind_adapter.py +++ b/tests/quantization/test_sm70_turbomind_adapter.py @@ -91,224 +91,3 @@ def test_symmetric_int4_zero_points_are_eight(): assert zeros.dtype == torch.float16 assert zeros.tolist() == [[8, 8, 8], [8, 8, 8]] - - -def _make_nvfp4_state(tm, *, gated_silu: bool): - return tm.SM70TurboMindLinearState( - weight=torch.empty((3, 1), dtype=torch.uint8), - scales=torch.empty((1, 4), dtype=torch.float16), - group_size=16, - k_ld=3, - q_ld=4, - output_size=4, - op_kind="nvfp4", - gated_silu=gated_silu, - ) - - -def _make_nvfp4_qpn4_state(tm, *, gated_silu: bool): - output_size = 4 - return tm.SM70TurboMindLinearState( - weight=torch.empty((3, output_size // 2), dtype=torch.uint8), - scales=torch.empty( - (1, output_size), - dtype=torch.uint8 if gated_silu else torch.float16, - ), - group_size=16, - k_ld=0, - q_ld=0, - output_size=output_size, - op_kind="nvfp4_qpn4", - gated_silu=gated_silu, - dense_weight_ptr=1234, - global_scale=0.25 if gated_silu else 0.0, - use_scale_code=gated_silu, - ) - - -def test_nvfp4_gated_layout_is_deinterleaved_for_unfused_apply(monkeypatch): - tm = _load_adapter() - layer = torch.nn.Module() - setattr(layer, tm.STATE_ATTR, _make_nvfp4_state(tm, gated_silu=True)) - - from vllm import _sm70_ops as sm70_ops - - calls = [] - - def fake_gemm( - out, - x, - weight, - scales, - group_size, - k_ld, - q_ld, - gated_silu=False, - ): - del x, weight, scales, group_size, k_ld, q_ld - calls.append(gated_silu) - out.copy_(torch.tensor([[1.0, 10.0, 2.0, 20.0]], dtype=out.dtype)) - - monkeypatch.setattr(sm70_ops, "nvfp4_gemm_sm70_out", fake_gemm) - - output = tm.apply_prepared_linear( - layer, - torch.ones((1, 3), dtype=torch.float16), - bias=None, - ) - - assert calls == [False] - assert output.tolist() == [[1.0, 2.0, 10.0, 20.0]] - - -def test_nvfp4_gated_layout_uses_fused_silu_epilogue(monkeypatch): - tm = _load_adapter() - layer = torch.nn.Module() - setattr(layer, tm.STATE_ATTR, _make_nvfp4_state(tm, gated_silu=True)) - - from vllm import _sm70_ops as sm70_ops - - calls = [] - - def fake_gemm( - out, - x, - weight, - scales, - group_size, - k_ld, - q_ld, - gated_silu=False, - ): - del weight, scales - calls.append( - { - "out_shape": tuple(out.shape), - "x_shape": tuple(x.shape), - "group_size": group_size, - "k_ld": k_ld, - "q_ld": q_ld, - "gated_silu": gated_silu, - } - ) - out.fill_(3.0) - - monkeypatch.setattr(sm70_ops, "nvfp4_gemm_sm70_out", fake_gemm) - - output = tm.apply_prepared_fused_silu_and_mul( - layer, - torch.ones((2, 1, 3), dtype=torch.float16), - ) - - assert output is not None - assert output.shape == (2, 1, 2) - assert output.tolist() == [[[3.0, 3.0]], [[3.0, 3.0]]] - assert calls == [ - { - "out_shape": (2, 2), - "x_shape": (2, 3), - "group_size": 16, - "k_ld": 3, - "q_ld": 4, - "gated_silu": True, - } - ] - - -def test_nvfp4_fused_silu_rejects_non_gated_state(): - tm = _load_adapter() - layer = torch.nn.Module() - setattr(layer, tm.STATE_ATTR, _make_nvfp4_state(tm, gated_silu=False)) - - output = tm.apply_prepared_fused_silu_and_mul( - layer, - torch.ones((1, 3), dtype=torch.float16), - ) - - assert output is None - - -def test_nvfp4_qpn4_regular_apply_uses_opaque_dynamic_m_dispatch(monkeypatch): - tm = _load_adapter() - layer = torch.nn.Module() - setattr(layer, tm.STATE_ATTR, _make_nvfp4_qpn4_state(tm, gated_silu=False)) - - from vllm import _sm70_ops as sm70_ops - - calls = [] - - def fake_dispatch( - out, - dense_weight_ptr, - x, - weight, - scales, - global_scale, - use_scale_code, - gated_silu, - ): - del weight, scales - calls.append( - ( - tuple(out.shape), - dense_weight_ptr, - tuple(x.shape), - global_scale, - use_scale_code, - gated_silu, - ) - ) - out.fill_(2.0) - - monkeypatch.setattr(sm70_ops, "nvfp4_qpn4_dispatch_sm70_out", fake_dispatch) - output = tm.apply_prepared_linear( - layer, - torch.ones((2, 3), dtype=torch.float16), - bias=None, - ) - - assert output.tolist() == [[2.0] * 4, [2.0] * 4] - assert calls == [((2, 4), 1234, (2, 3), 0.0, False, False)] - - -def test_nvfp4_qpn4_fused_gate_uses_scale_code_dispatch(monkeypatch): - tm = _load_adapter() - layer = torch.nn.Module() - setattr(layer, tm.STATE_ATTR, _make_nvfp4_qpn4_state(tm, gated_silu=True)) - - from vllm import _sm70_ops as sm70_ops - - calls = [] - - def fake_dispatch( - out, - dense_weight_ptr, - x, - weight, - scales, - global_scale, - use_scale_code, - gated_silu, - ): - del weight, scales - calls.append( - ( - tuple(out.shape), - dense_weight_ptr, - tuple(x.shape), - global_scale, - use_scale_code, - gated_silu, - ) - ) - out.fill_(4.0) - - monkeypatch.setattr(sm70_ops, "nvfp4_qpn4_dispatch_sm70_out", fake_dispatch) - output = tm.apply_prepared_fused_silu_and_mul( - layer, - torch.ones((1, 3), dtype=torch.float16), - ) - - assert output is not None - assert output.tolist() == [[4.0, 4.0]] - assert calls == [((1, 2), 1234, (1, 3), 0.25, True, True)] diff --git a/tests/quantization/test_sm70_warmup.py b/tests/quantization/test_sm70_warmup.py index e6e12b922a..781c01d7e7 100644 --- a/tests/quantization/test_sm70_warmup.py +++ b/tests/quantization/test_sm70_warmup.py @@ -6,8 +6,6 @@ import torch from torch import nn -import vllm.envs as envs -from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod from vllm.model_executor.warmup import awq_sm70_warmup as warmup @@ -28,14 +26,6 @@ def _grouped_fp8_layer() -> nn.Module: return layer -def test_fp8_grouped_bmm_decode_defaults_on_with_rollback(monkeypatch): - monkeypatch.delenv("VLLM_SM70_FP8_GROUPED_BMM_DECODE", raising=False) - assert envs.VLLM_SM70_FP8_GROUPED_BMM_DECODE - - monkeypatch.setenv("VLLM_SM70_FP8_GROUPED_BMM_DECODE", "0") - assert not envs.VLLM_SM70_FP8_GROUPED_BMM_DECODE - - def test_fp8_warmup_discovers_grouped_bmm_by_per_group_shape(): layer = _grouped_fp8_layer() model = nn.Sequential(layer) @@ -45,14 +35,6 @@ def test_fp8_warmup_discovers_grouped_bmm_by_per_group_shape(): assert discovered == [(layer, False)] -def test_fp8_warmup_skips_static_qpn8_dispatch(): - layer = _grouped_fp8_layer() - layer.sm70_fp8_qpn8 = True - model = nn.Sequential(layer) - - assert list(warmup._iter_unique_fp8_dense_layers(model)) == [] - - def test_fp8_warmup_matches_grouped_bmm_runtime_slice(monkeypatch): layer = _grouped_fp8_layer() calls = [] @@ -86,344 +68,6 @@ def record_call(out, x, weight, scales, group_size, k_ld, q_ld, gated_silu): assert all(not call.gated_silu for call in calls) -def test_fp8_warmup_includes_one_launch_grouped_decode(monkeypatch): - layer = _grouped_fp8_layer() - layer.sm70_fp8_bmm_groups = 2 - layer.sm70_fp8_bmm_grouped_decode = True - layer.weight = nn.Parameter( - torch.empty((2, 128, 64), dtype=torch.uint8), requires_grad=False - ) - layer.weight_scale_inv = nn.Parameter( - torch.empty((2, 1, 64), dtype=torch.float16), requires_grad=False - ) - layer.sm70_fp8_bmm_grouped_offsets = torch.arange(3, dtype=torch.int32) - layer.sm70_fp8_bmm_grouped_ptrs_w = torch.empty(2, dtype=torch.int64) - layer.sm70_fp8_bmm_grouped_ptrs_s = torch.empty(2, dtype=torch.int64) - dense_calls = [] - grouped_calls = [] - monkeypatch.setattr(torch.ops._C, "fp8_gemm_sm70_out_meta", object(), raising=False) - monkeypatch.setattr( - warmup.sm70_ops, - "fp8_gemm_sm70_out", - lambda *args: dense_calls.append(args), - ) - monkeypatch.setattr( - warmup.sm70_ops, - "fp8_moe_gemm_sm70_per_expert_dispatch_out", - lambda *args: grouped_calls.append(args), - ) - - count = warmup._warmup_fp8_dense_layers([(layer, False)], [1, 4]) - - assert count == 3 - assert len(dense_calls) == 2 - assert len(grouped_calls) == 1 - call = grouped_calls[0] - assert tuple(call[0].shape) == (2, 64) - assert tuple(call[1].shape) == (2, 128) - assert call[2] is layer.sm70_fp8_bmm_grouped_offsets - assert call[3] is layer.sm70_fp8_bmm_grouped_ptrs_w - assert call[4] is layer.sm70_fp8_bmm_grouped_ptrs_s - assert call[5:] == (2, 128, 64, 128, False) - - -def test_fp8_grouped_bmm_decode_uses_one_dispatch(monkeypatch): - layer = _grouped_fp8_layer() - layer.sm70_fp8_bmm_groups = 2 - layer.sm70_fp8_bmm_grouped_decode = True - layer.sm70_fp8_bmm_grouped_offsets = torch.arange(3, dtype=torch.int32) - layer.sm70_fp8_bmm_grouped_ptrs_w = torch.empty(2, dtype=torch.int64) - layer.sm70_fp8_bmm_grouped_ptrs_s = torch.empty(2, dtype=torch.int64) - layer.output_size_per_partition = 128 - calls = [] - - def grouped(out, x, *args): - calls.append((out, x, args)) - out.copy_(torch.arange(out.numel(), dtype=out.dtype).reshape_as(out)) - - def fail_dense(*args): - raise AssertionError("dense fallback used") - - monkeypatch.setattr( - warmup.sm70_ops, - "fp8_moe_gemm_sm70_per_expert_dispatch_out", - grouped, - ) - monkeypatch.setattr( - warmup.sm70_ops, - "fp8_gemm_sm70_out", - fail_dense, - ) - x = torch.empty((1, 2, 128), dtype=torch.float16) - - out = object.__new__(Fp8LinearMethod).apply(layer, x) - - assert tuple(out.shape) == (1, 2, 64) - assert len(calls) == 1 - grouped_out, grouped_x, args = calls[0] - assert tuple(grouped_out.shape) == (2, 64) - assert tuple(grouped_x.shape) == (2, 128) - assert args[0] is layer.sm70_fp8_bmm_grouped_offsets - assert args[1] is layer.sm70_fp8_bmm_grouped_ptrs_w - assert args[2] is layer.sm70_fp8_bmm_grouped_ptrs_s - assert args[3:] == (2, 128, 64, 128, False) - torch.testing.assert_close( - out, - torch.arange(128, dtype=torch.float16).reshape(1, 2, 64), - ) - - -def test_fp8_grouped_bmm_decode_retains_multirow_fallback(monkeypatch): - layer = _grouped_fp8_layer() - layer.sm70_fp8_bmm_groups = 2 - layer.sm70_fp8_bmm_grouped_decode = True - layer.sm70_fp8_bmm_output_size = 64 - dense_calls = [] - - def dense(out, x, *args): - dense_calls.append((tuple(out.shape), tuple(x.shape), args)) - out.zero_() - - def fail_grouped(*args): - raise AssertionError("grouped decode must remain batch-one only") - - monkeypatch.setattr(warmup.sm70_ops, "fp8_gemm_sm70_out", dense) - monkeypatch.setattr( - warmup.sm70_ops, - "fp8_moe_gemm_sm70_per_expert_dispatch_out", - fail_grouped, - ) - x = torch.empty((2, 2, 128), dtype=torch.float16) - - out = object.__new__(Fp8LinearMethod).apply(layer, x) - - assert tuple(out.shape) == (2, 2, 64) - assert [call[:2] for call in dense_calls] == [ - ((2, 64), (2, 128)), - ((2, 64), (2, 128)), - ] - - -def test_fp8_warmup_supports_modelopt_turbomind_layout(monkeypatch): - layer = nn.Module() - layer.sm70_modelopt_fp8_turbomind = True - layer.sm70_modelopt_fp8_k_ld = 128 - layer.sm70_modelopt_fp8_q_ld = 64 - layer.output_size_per_partition = 64 - layer.weight = nn.Parameter( - torch.empty((128, 64), dtype=torch.uint8), requires_grad=False - ) - layer.weight_scale = nn.Parameter( - torch.empty((1, 64), dtype=torch.float16), requires_grad=False - ) - model = nn.Sequential(layer) - calls = [] - monkeypatch.setattr(torch.ops._C, "fp8_gemm_sm70_out_meta", object(), raising=False) - monkeypatch.setattr( - warmup.sm70_ops, - "fp8_gemm_sm70_out", - lambda out, x, weight, scales, group_size, k_ld, q_ld, gated_silu: calls.append( - (out.shape, x.shape, scales, group_size, k_ld, q_ld) - ), - ) - - discovered = list(warmup._iter_unique_fp8_dense_layers(model)) - count = warmup._warmup_fp8_dense_layers(discovered, [1, 4]) - - assert discovered == [(layer, False)] - assert count == 2 - assert [tuple(call[0]) for call in calls] == [(1, 64), (4, 64)] - assert all(call[2] is layer.weight_scale for call in calls) - assert all(call[3:] == (128, 128, 64) for call in calls) - - -def _nvfp4_moe_layer() -> nn.Module: - layer = nn.Module() - layer.sm70_nvfp4_moe = True - layer.moe_config = SimpleNamespace(experts_per_token=8) - layer.swiglu_limit = None - layer.sm70_nvfp4_num_experts = 256 - layer.sm70_nvfp4_w13_k_dim = 2048 - layer.sm70_nvfp4_w13_n_dim = 256 - layer.sm70_nvfp4_w2_k_dim = 128 - layer.sm70_nvfp4_w2_n_dim = 2048 - layer.sm70_nvfp4_group_size = 16 - layer.sm70_nvfp4_graph_safe_max_tokens = 18 - layer.sm70_nvfp4_compact_grouped_max_tokens = 10 - layer.w13_tm_weight = nn.Parameter( - torch.empty((1, 1), dtype=torch.uint8), requires_grad=False - ) - layer.w13_strided_ptrs_w = torch.empty(1, dtype=torch.uint8) - layer.w13_strided_ptrs_s = torch.empty(1, dtype=torch.uint8) - layer.w2_strided_ptrs_w = torch.empty(1, dtype=torch.uint8) - layer.w2_strided_ptrs_s = torch.empty(1, dtype=torch.uint8) - layer._nvfp4_sm70_dense_expert_ids = torch.arange(256, dtype=torch.int32) - return layer - - -def test_nvfp4_moe_warmup_discovers_and_uses_compact_decode_shapes(monkeypatch): - layer = _nvfp4_moe_layer() - model = nn.Sequential(layer) - calls = [] - monkeypatch.setattr( - torch.ops._C, "nvfp4_moe_dense_stage_sm70_out", object(), raising=False - ) - - def record_call( - out, - x, - offsets, - expert_ids, - ptrs_w, - ptrs_s, - num_experts, - k, - n, - group_size, - ): - calls.append( - ( - tuple(out.shape), - tuple(x.shape), - offsets.tolist(), - expert_ids.tolist(), - num_experts, - k, - n, - group_size, - ) - ) - - monkeypatch.setattr(warmup.sm70_ops, "nvfp4_moe_dense_stage_sm70_out", record_call) - monkeypatch.setattr( - warmup, - "_prepare_compact_slot_groups", - lambda sorted_ids, offsets, active_ids: ( - offsets.copy_(torch.arange(offsets.numel(), dtype=torch.int32)), - active_ids.copy_(sorted_ids), - ), - ) - monkeypatch.setattr( - torch.ops._C, - "silu_and_mul", - lambda out, gate_up: out.zero_(), - raising=False, - ) - - discovered = list(warmup._iter_unique_nvfp4_moe_layers(model)) - count = warmup._warmup_nvfp4_moe_decode_layers(discovered, [1, 2]) - - assert discovered == [layer] - assert count == 4 - assert [call[4] for call in calls] == [8, 8, 16, 16] - assert calls[0][2] == list(range(9)) - assert calls[0][3] == list(range(8)) - assert calls[2][2] == list(range(17)) - assert calls[2][3] == list(range(16)) - - -def test_nvfp4_moe_warmup_includes_opted_in_cuda_graph_shapes(monkeypatch): - monkeypatch.setattr(warmup.envs, "VLLM_SM70_NVFP4_MOE_TUNE_MAX_TOKENS", 640) - worker = SimpleNamespace( - vllm_config=SimpleNamespace( - compilation_config=SimpleNamespace( - cudagraph_capture_sizes=[1, 2, 4, 5, 8, 9, 18, 20, 40, 60, 80, 81] - ) - ) - ) - - assert warmup._get_nvfp4_moe_token_counts( - worker, [_nvfp4_moe_layer()], [1, 2, 4, 8] - ) == [*range(1, 11), 18, 20, 40, 60, 80] - - -def test_nvfp4_moe_warmup_uses_slot_compact_through_b10(monkeypatch): - layer = _nvfp4_moe_layer() - calls = [] - monkeypatch.setattr( - torch.ops._C, "nvfp4_moe_dense_stage_sm70_out", object(), raising=False - ) - monkeypatch.setattr( - warmup.sm70_ops, - "nvfp4_moe_dense_stage_sm70_out", - lambda *args: calls.append(args), - ) - monkeypatch.setattr( - warmup, - "_prepare_compact_slot_groups", - lambda sorted_ids, offsets, active_ids: ( - offsets.copy_(torch.arange(offsets.numel(), dtype=torch.int32)), - active_ids.copy_(sorted_ids), - ), - ) - monkeypatch.setattr( - torch.ops._C, - "silu_and_mul", - lambda out, gate_up: out.zero_(), - raising=False, - ) - - assert warmup._warmup_nvfp4_moe_decode_layers([layer], [9, 10]) == 4 - assert [call[6] for call in calls] == [72, 72, 80, 80] - assert all(call[2].numel() == call[6] + 1 for call in calls) - assert all(call[3].numel() == call[6] for call in calls) - assert calls[0][2].tolist() == list(range(73)) - assert calls[2][2].tolist() == list(range(81)) - - -def test_nvfp4_moe_warmup_uses_full_expert_groups_above_compact_b10(monkeypatch): - layer = _nvfp4_moe_layer() - calls = [] - monkeypatch.setattr( - torch.ops._C, "nvfp4_moe_dense_stage_sm70_out", object(), raising=False - ) - monkeypatch.setattr( - warmup.sm70_ops, - "nvfp4_moe_dense_stage_sm70_out", - lambda *args: calls.append(args), - ) - monkeypatch.setattr( - torch.ops._C, - "silu_and_mul", - lambda out, gate_up: out.zero_(), - raising=False, - ) - - assert warmup._warmup_nvfp4_moe_decode_layers([layer], [11]) == 2 - assert [call[6] for call in calls] == [256, 256] - assert all(call[2].numel() == 257 for call in calls) - assert all(call[3].numel() == 256 for call in calls) - assert calls[0][2][:89].tolist() == list(range(89)) - assert calls[0][2][89:].tolist() == [88] * (257 - 89) - - -def test_nvfp4_moe_warmup_supports_mtp4_c16_width(monkeypatch): - layer = _nvfp4_moe_layer() - calls = [] - monkeypatch.setattr( - torch.ops._C, "nvfp4_moe_dense_stage_sm70_out", object(), raising=False - ) - monkeypatch.setattr( - warmup.sm70_ops, - "nvfp4_moe_dense_stage_sm70_out", - lambda *args: calls.append(args), - ) - monkeypatch.setattr( - torch.ops._C, - "silu_and_mul", - lambda out, gate_up: out.zero_(), - raising=False, - ) - - assert warmup._warmup_nvfp4_moe_decode_layers([layer], [80]) == 2 - assert [tuple(call[0].shape) for call in calls] == [(640, 256), (640, 2048)] - assert [call[6] for call in calls] == [256, 256] - offsets = calls[0][2] - assert offsets.numel() == 257 - assert offsets[:4].tolist() == [0, 3, 6, 9] - assert offsets[-1].item() == 640 - - def test_fp8_coordinated_warmup_leader_broadcasts_rank0_lut(monkeypatch): import vllm.distributed.parallel_state as parallel_state diff --git a/tests/test_envs.py b/tests/test_envs.py index ffc79b49e4..db50deb241 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -121,60 +121,6 @@ def test_precompiled_install_flags_are_orthogonal() -> None: assert environment_variables["VLLM_USE_PRECOMPILED_RUST"]() is False -def test_sm70_concurrency_tuning_envs( - monkeypatch: pytest.MonkeyPatch, -) -> None: - names = ( - "VLLM_SM70_TP4_MTP_AR_BLOCK_TUNING", - "VLLM_SM70_TOPK_TOPP_8_WARPS", - "VLLM_SM70_MTP_MOE_TUNED_CONFIG", - ) - for name in names: - monkeypatch.delenv(name, raising=False) - assert environment_variables[name]() is False - monkeypatch.setenv(name, "1") - assert environment_variables[name]() is True - - name = "VLLM_SM70_TOPK_TOPP_B8_B16_8_WARPS" - monkeypatch.delenv(name, raising=False) - assert environment_variables[name]() is True - monkeypatch.setenv(name, "0") - assert environment_variables[name]() is False - - name = "VLLM_SM70_DSV4_QNORM_KV_FUSED_TP4" - monkeypatch.delenv(name, raising=False) - assert environment_variables[name]() is True - 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") - assert environment_variables["VLLM_SM70_NVFP4_MOE_TUNE_MAX_TOKENS"]() == 640 - - def test_flash_v100_g6_sawtooth_pipeline_envs( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -204,20 +150,6 @@ def test_flash_v100_g6_sawtooth_pipeline_envs( assert environment_variables[name]() == expected_int + 1 -def test_flash_v100_e4m3_page800_fastpath_envs( - monkeypatch: pytest.MonkeyPatch, -) -> None: - bool_defaults = { - "VLLM_FLASH_V100_E4M3_PAGE800_FASTPATH": True, - "VLLM_FLASH_V100_E4M3_PAGE800_FASTPATH_TRACE": False, - } - for name, expected_bool in bool_defaults.items(): - monkeypatch.delenv(name, raising=False) - assert environment_variables[name]() is expected_bool - monkeypatch.setenv(name, "0" if expected_bool else "1") - assert environment_variables[name]() is not expected_bool - - 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 9167493688..ca2a60d4ce 100644 --- a/tests/tool_parsers/test_qwen3coder_tool_parser.py +++ b/tests/tool_parsers/test_qwen3coder_tool_parser.py @@ -1384,7 +1384,7 @@ def test_streaming_multi_param_single_chunk(qwen3_tool_parser, qwen3_tokenizer): @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/test_deepseek_ocr_resolution_modes.py b/tests/transformers_utils/test_deepseek_ocr_resolution_modes.py deleted file mode 100644 index b7f8888ecd..0000000000 --- a/tests/transformers_utils/test_deepseek_ocr_resolution_modes.py +++ /dev/null @@ -1,163 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""DeepSeek-OCR resolution modes: per-request ``image_mode`` selection, -crop parametrization, and — critically — the producer/counter consistency -contract (`tokenize_with_images` versus -`count_image_tokens`), whose violation is the "N multimodal tokens vs M -placeholders" crash class.""" - -import pytest -from transformers.processing_utils import ProcessorMixin - -from vllm.transformers_utils.processors.deepseek_ocr import ( - BASE_SIZE, - CROP_MODE, - IMAGE_SIZE, - RESOLUTION_MODES, - DeepseekOCRProcessor, - count_image_tokens_for, -) - -# --------------------------------------------------------------------------- -# Pure-function tests (no tokenizer). -# --------------------------------------------------------------------------- - - -def test_mode_table_matches_official_config(): - # Verbatim from the model repo's config.py (captured in the issue refs). - assert RESOLUTION_MODES == { - "tiny": {"base_size": 512, "image_size": 512, "crop_mode": False}, - "small": {"base_size": 640, "image_size": 640, "crop_mode": False}, - "base": {"base_size": 1024, "image_size": 1024, "crop_mode": False}, - "large": {"base_size": 1280, "image_size": 1280, "crop_mode": False}, - "gundam": {"base_size": 1024, "image_size": 640, "crop_mode": True}, - } - # The module defaults ARE the gundam mode (byte-compat anchor). - assert (BASE_SIZE, IMAGE_SIZE, CROP_MODE) == (1024, 640, True) - - -def _count(mode: str, width: int, height: int, **kw) -> int: - cfg = RESOLUTION_MODES[mode] - return count_image_tokens_for( - image_width=width, - image_height=height, - base_size=cfg["base_size"], - image_size=cfg["image_size"], - cropping=cfg["crop_mode"], - **kw, - ) - - -def test_no_crop_modes_are_size_independent(): - for mode in ("tiny", "small", "base", "large"): - assert _count(mode, 100, 100) == _count(mode, 4000, 3000) - - -def test_gundam_small_image_equals_no_crop_arithmetic(): - # <= image_size on both dims: never cropped, regardless of crop_mode. - assert _count("gundam", 640, 480) == count_image_tokens_for( - image_width=640, - image_height=480, - base_size=1024, - image_size=640, - cropping=False, - ) - - -def test_gundam_large_image_counts_tiles(): - assert _count("gundam", 1700, 2200) > _count("gundam", 100, 100) - - -def test_max_crops_bounds_the_tile_count(): - big = _count("gundam", 4000, 4000, max_crops=9) - small = _count("gundam", 4000, 4000, max_crops=2) - assert big > small - - -# --------------------------------------------------------------------------- -# Processor tests with a minimal tokenizer (CPU-only, no model files). -# --------------------------------------------------------------------------- - - -class _TokenizerStub: - bos_token_id = 1 - eos_token_id = 2 - pad_token_id = 0 - pad_token = None - padding_side = "right" - vocab = {"": 99} - - def add_special_tokens(self, tokens): - self.pad_token = tokens["pad_token"] - - def encode(self, text, add_special_tokens=False): - del add_special_tokens - return [10] if text else [] - - def decode(self, tokens, **kwargs): - del kwargs - return " ".join(map(str, tokens)) - - -@pytest.fixture -def tokenizer(monkeypatch): - monkeypatch.setattr(ProcessorMixin, "__init__", lambda *args, **kwargs: None) - - return _TokenizerStub() - - -def _images(sizes): - from PIL import Image - - return [Image.new("RGB", s, color="white") for s in sizes] - - -def test_default_construction_is_gundam_byte_compat(tokenizer): - default = DeepseekOCRProcessor(tokenizer=tokenizer) - explicit = DeepseekOCRProcessor(tokenizer=tokenizer, image_mode="gundam") - for attr in ("base_size", "image_size", "crop_mode", "min_crops", "max_crops"): - assert getattr(default, attr) == getattr(explicit, attr) - assert default.crop_mode is True and default.image_size == 640 - - -def test_invalid_mode_and_crops_raise(tokenizer): - with pytest.raises(ValueError, match="image_mode"): - DeepseekOCRProcessor(tokenizer=tokenizer, image_mode="giant") - with pytest.raises(ValueError, match="crop bounds"): - DeepseekOCRProcessor(tokenizer=tokenizer, min_crops=5, max_crops=2) - with pytest.raises(ValueError, match="crop bounds"): - DeepseekOCRProcessor(tokenizer=tokenizer, min_crops=2, max_crops=10) - - -def test_image_mode_authoritative_over_trio(tokenizer): - # The processing info always forwards the constant trio; a named mode - # must win over it. - proc = DeepseekOCRProcessor( - tokenizer=tokenizer, - image_size=640, - base_size=1024, - crop_mode=True, - image_mode="base", - ) - assert (proc.base_size, proc.image_size, proc.crop_mode) == (1024, 1024, False) - - -@pytest.mark.parametrize("mode", [None, "base", "large", "tiny", "gundam"]) -@pytest.mark.parametrize("size", [(100, 100), (1200, 800), (1000, 2200)]) -def test_producer_counter_consistency(tokenizer, mode, size): - kwargs = {} if mode is None else {"image_mode": mode} - proc = DeepseekOCRProcessor(tokenizer=tokenizer, **kwargs) - out = proc(prompt="\nFree OCR. ", images=_images([size])) - produced = int(out["num_image_tokens"][0]) - counted = proc.count_image_tokens(image_width=size[0], image_height=size[1]) - assert produced == counted - - -def test_multi_image_producer_counter_consistency(tokenizer): - proc = DeepseekOCRProcessor(tokenizer=tokenizer) - sizes = [(1200, 800), (1000, 2200)] - out = proc(prompt="ab", images=_images(sizes)) - for i, size in enumerate(sizes): - produced = int(out["num_image_tokens"][i]) - counted = proc.count_image_tokens(image_width=size[0], image_height=size[1]) - assert produced == counted diff --git a/tests/v1/attention/test_gdn_metadata_builder.py b/tests/v1/attention/test_gdn_metadata_builder.py index 76626418e1..e43f1c0fa7 100644 --- a/tests/v1/attention/test_gdn_metadata_builder.py +++ b/tests/v1/attention/test_gdn_metadata_builder.py @@ -12,7 +12,6 @@ from pathlib import Path from types import SimpleNamespace -import numpy as np import pytest import torch @@ -37,82 +36,14 @@ build_gdn_spec_decode_state_contract, gdn_spec_metadata_tensors, get_registered_gdn_spec_metadata_tensors, - prepare_dflash2_gdn_group_metadata, ) from vllm.v1.attention.backends.utils import PAD_SLOT_ID from vllm.v1.kv_cache_interface import MambaSpec -from vllm.v1.worker.gpu.attn_utils import compute_common_gdn_attn_metadata -from vllm.v1.worker.gpu.model_states import mamba_hybrid -from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState BLOCK_SIZE = 16 DEVICE = torch.device("cpu") -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) - monkeypatch.delenv("VLLM_SM70_DUMP_GDN_GRAPH_DIR", raising=False) - - def fail_debug_work(*_args, **_kwargs): - raise AssertionError("disabled debug hook performed debug work") - - cuda_module = vars(qwen_gdn.torch)["cuda"] - monkeypatch.setattr(cuda_module, "is_current_stream_capturing", fail_debug_work) - monkeypatch.setattr(qwen_gdn, "_sm70_gdn_graph_buffer_copy", fail_debug_work) - qwen_gdn._sm70_dump_gdn_core_tensor( - "input_qkv", - "layer.0", - torch.zeros(1), - ) - - -def test_mamba_hybrid_passes_seq_lens_cpu_upper_bound(monkeypatch): - """Flash-V100 metadata must not lazily copy seq_lens back from the GPU.""" - - captured_kwargs = None - - def capture_build_attn_metadata(**kwargs): - nonlocal captured_kwargs - captured_kwargs = kwargs - return {} - - monkeypatch.setattr( - mamba_hybrid, - "build_attn_metadata", - capture_build_attn_metadata, - ) - state = object.__new__(MambaHybridModelState) - state.max_model_len = 4096 - seq_lens_cpu_upper_bound = torch.tensor([130], dtype=torch.int32) - input_batch = SimpleNamespace( - num_reqs=1, - num_tokens=8, - num_reqs_after_padding=1, - num_tokens_after_padding=8, - query_start_loc_np=np.array([0, 8], dtype=np.int32), - num_scheduled_tokens=np.array([8], dtype=np.int32), - is_prefilling_np=np.array([False]), - query_start_loc=torch.tensor([0, 8], dtype=torch.int32), - seq_lens=torch.tensor([123], dtype=torch.int32), - seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, - dcp_local_seq_lens=None, - prefix_anchor_lens=None, - ) - state.prepare_attn( - input_batch, - CUDAGraphMode.NONE, - block_tables=(), - slot_mappings=torch.empty((0, 8), dtype=torch.int32), - attn_groups=[], - kv_cache_config=SimpleNamespace(kv_cache_groups=[]), - for_capture=True, - ) - - assert captured_kwargs is not None - assert captured_kwargs["seq_lens_cpu_upper_bound"] is seq_lens_cpu_upper_bound - - @pytest.fixture def local_gdn_model(tmp_path: Path) -> str: model_dir = tmp_path / "gdn-test-model" @@ -241,7 +172,6 @@ def _create_gdn_builder( use_full_cuda_graph: bool = False, mamba_cache_mode: str = "none", max_cudagraph_capture_size: int = 4, - device: torch.device = DEVICE, ) -> GDNAttentionMetadataBuilder: """Create a GDNAttentionMetadataBuilder with minimal config.""" vllm_config = create_vllm_config(model_name=model_name, block_size=BLOCK_SIZE) @@ -251,11 +181,6 @@ def _create_gdn_builder( vllm_config.compilation_config.max_cudagraph_capture_size = ( max_cudagraph_capture_size ) - else: - # Keep the helper's default contract explicit. The repository-wide - # compilation default may select FULL graphs on CUDA platforms, which - # would pad otherwise eager metadata and make these tests route-dependent. - vllm_config.compilation_config.cudagraph_mode = CUDAGraphMode.NONE if num_speculative_tokens > 0: vllm_config.speculative_config = SpeculativeConfig( method="ngram", @@ -271,7 +196,7 @@ def _create_gdn_builder( kv_cache_spec=mamba_spec, layer_names=["layer.0"], vllm_config=vllm_config, - device=device, + device=DEVICE, ) @@ -287,73 +212,20 @@ def _build( builder: GDNAttentionMetadataBuilder, batch_spec: BatchSpec, num_decode_draft_tokens: list[int] | None = None, - use_common_metadata: bool = False, ) -> GDNAttentionMetadata: """Build GDN attention metadata, optionally with spec-decode kwargs.""" common = create_common_attn_metadata(batch_spec, BLOCK_SIZE, DEVICE) kwargs: dict = {} if num_decode_draft_tokens is not None: - num_decode_draft_tokens_cpu = torch.tensor( + kwargs["num_decode_draft_tokens_cpu"] = torch.tensor( num_decode_draft_tokens, dtype=torch.int32, device="cpu" ) - num_accepted_tokens = torch.ones( + kwargs["num_accepted_tokens"] = torch.ones( batch_spec.batch_size, dtype=torch.int32, device=DEVICE ) - kwargs["num_decode_draft_tokens_cpu"] = num_decode_draft_tokens_cpu - kwargs["num_accepted_tokens"] = num_accepted_tokens - if use_common_metadata: - kwargs["common_gdn_metadata"] = compute_common_gdn_attn_metadata( - num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, - query_start_loc=common.query_start_loc, - query_start_loc_cpu=common.query_start_loc_cpu, - num_spec_state_tokens=builder.num_spec_state_tokens, - legacy_mixed_decode_routing=( - qwen_gdn.envs.VLLM_SM70_MTP_LEGACY_GDN_MIXED_DECODE_ROUTING - ), - ) return builder.build(common_prefix_len=0, common_attn_metadata=common, **kwargs) -def _assert_gdn_metadata_equal( - actual: GDNAttentionMetadata, - expected: GDNAttentionMetadata, -) -> None: - scalar_fields = ( - "num_prefills", - "num_prefill_tokens", - "num_decodes", - "num_decode_tokens", - "num_spec_decodes", - "num_spec_decode_tokens", - "num_actual_tokens", - ) - for field in scalar_fields: - assert getattr(actual, field) == getattr(expected, field), field - - tensor_fields = ( - "has_initial_state", - "spec_query_start_loc", - "non_spec_query_start_loc", - "spec_state_indices_tensor", - "non_spec_state_indices_tensor", - "spec_sequence_masks", - "spec_token_indx", - "non_spec_token_indx", - "num_accepted_tokens", - "spec_state_slot_selectors", - "chunk_indices", - "chunk_offsets", - "batch_ptr", - "token_chunk_offset_ptr", - ) - for field in tensor_fields: - actual_tensor = getattr(actual, field) - expected_tensor = getattr(expected, field) - assert (actual_tensor is None) == (expected_tensor is None), field - if actual_tensor is not None: - assert torch.equal(actual_tensor, expected_tensor), field - - def _effective_spec_initial_state_slots( meta: GDNAttentionMetadata, ) -> list[int]: @@ -388,51 +260,6 @@ def test_gdn_build_classification( assert meta.num_spec_decodes == test_case.expected_num_spec_decodes -@pytest.mark.parametrize( - "test_case", - [ - case - for case in GDN_BUILD_TEST_CASES.values() - if case.num_decode_draft_tokens is not None - ], - ids=[ - name - for name, case in GDN_BUILD_TEST_CASES.items() - if case.num_decode_draft_tokens is not None - ], -) -def test_common_gdn_metadata_matches_per_group_builder( - test_case: GDNBuildTestCase, - local_gdn_model: str, -): - generic_builder = _create_gdn_builder( - local_gdn_model, test_case.num_speculative_tokens - ) - common_builder = _create_gdn_builder( - local_gdn_model, test_case.num_speculative_tokens - ) - batch = BatchSpec( - seq_lens=test_case.seq_lens, - query_lens=test_case.query_lens, - ) - - torch.manual_seed(0) - expected = _build( - generic_builder, - batch, - test_case.num_decode_draft_tokens, - ) - torch.manual_seed(0) - actual = _build( - common_builder, - batch, - test_case.num_decode_draft_tokens, - use_common_metadata=True, - ) - - _assert_gdn_metadata_equal(actual, expected) - - def test_mixed_decode_stays_decode_fastpath(local_gdn_model): """A non-spec query_len=1 row should stay on the decode fastpath even when another row is doing speculative verification.""" @@ -510,419 +337,6 @@ def test_full_cuda_graph_spec_replay_tail_uses_pad_slot(local_gdn_model): ] -def test_common_gdn_metadata_matches_full_graph_padding(local_gdn_model): - generic_builder = _create_gdn_builder( - local_gdn_model, - num_speculative_tokens=2, - use_full_cuda_graph=True, - ) - common_builder = _create_gdn_builder( - local_gdn_model, - num_speculative_tokens=2, - use_full_cuda_graph=True, - ) - batch = BatchSpec(seq_lens=[4097], query_lens=[3]) - common = create_common_attn_metadata(batch, BLOCK_SIZE, DEVICE).replace( - block_table_tensor=torch.tensor( - [[10, 11, 12, 13, 14]], - dtype=torch.int32, - device=DEVICE, - ), - num_actual_tokens=4, - ) - num_accepted_tokens = torch.tensor([3], dtype=torch.int32, device=DEVICE) - num_decode_draft_tokens_cpu = torch.tensor([2], dtype=torch.int32, device="cpu") - - expected = generic_builder.build( - common_prefix_len=0, - common_attn_metadata=common, - num_accepted_tokens=num_accepted_tokens, - num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, - ) - common_metadata = compute_common_gdn_attn_metadata( - num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, - query_start_loc=common.query_start_loc, - query_start_loc_cpu=common.query_start_loc_cpu, - num_spec_state_tokens=common_builder.num_spec_state_tokens, - legacy_mixed_decode_routing=( - qwen_gdn.envs.VLLM_SM70_MTP_LEGACY_GDN_MIXED_DECODE_ROUTING - ), - ) - assert common_metadata is not None - actual = common_builder.build( - common_prefix_len=0, - common_attn_metadata=common, - num_accepted_tokens=num_accepted_tokens, - num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, - common_gdn_metadata=common_metadata, - ) - - _assert_gdn_metadata_equal(actual, expected) - - -def test_prepared_dflash2_metadata_belongs_to_exact_builder(local_gdn_model): - builder = _create_gdn_builder( - local_gdn_model, - num_speculative_tokens=2, - use_full_cuda_graph=True, - ) - batch = BatchSpec(seq_lens=[4097], query_lens=[3]) - common = create_common_attn_metadata(batch, BLOCK_SIZE, DEVICE).replace( - block_table_tensor=torch.tensor( - [[10, 11, 12, 13]], dtype=torch.int32, device=DEVICE - ), - num_actual_tokens=4, - ) - num_accepted_tokens = torch.tensor([2], dtype=torch.int32, device=DEVICE) - num_decode_draft_tokens_cpu = torch.tensor([2], dtype=torch.int32, device="cpu") - common_metadata = compute_common_gdn_attn_metadata( - num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, - query_start_loc=common.query_start_loc, - query_start_loc_cpu=common.query_start_loc_cpu, - num_spec_state_tokens=builder.num_spec_state_tokens, - legacy_mixed_decode_routing=False, - ) - assert common_metadata is not None - prepared = builder.build( - common_prefix_len=0, - common_attn_metadata=common, - num_accepted_tokens=num_accepted_tokens, - num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, - common_gdn_metadata=common_metadata, - ) - - actual = builder.build( - common_prefix_len=0, - common_attn_metadata=common, - num_accepted_tokens=num_accepted_tokens, - num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, - common_gdn_metadata=common_metadata, - prepared_dflash2_metadata=prepared, - ) - assert actual is prepared - cached_contract = prepared._prepared_spec_metadata_tensors - assert cached_contract is not None - - replayed = builder.build( - common_prefix_len=0, - common_attn_metadata=common, - num_accepted_tokens=num_accepted_tokens, - num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, - common_gdn_metadata=common_metadata, - prepared_dflash2_metadata=prepared, - ) - assert replayed is prepared - assert prepared._prepared_spec_metadata_tensors is cached_contract - - other_builder = _create_gdn_builder( - local_gdn_model, - num_speculative_tokens=2, - use_full_cuda_graph=True, - ) - with pytest.raises(ValueError, match="does not belong"): - other_builder.build( - common_prefix_len=0, - common_attn_metadata=common, - num_accepted_tokens=num_accepted_tokens, - num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, - common_gdn_metadata=common_metadata, - prepared_dflash2_metadata=prepared, - ) - - -def test_dflash2_fused_gdn_group_metadata_replay( - monkeypatch, - local_gdn_model, -): - if not torch.cuda.is_available(): - pytest.skip("CUDA is required for the fused GDN metadata kernel") - device = torch.device("cuda") - if torch.cuda.get_device_capability(device) != (7, 0): - pytest.skip("the fused GDN metadata kernel is SM70-only") - - # The accepted DFlash2 verifier keeps the legacy spec-core operator off. - # Fused metadata must therefore own its graph-stable buffers independently. - monkeypatch.setenv("VLLM_SM70_QWEN_GDN_SPEC_CORE_OP", "0") - monkeypatch.setenv("VLLM_SM70_DFLASH2_FUSED_GDN_METADATA", "1") - monkeypatch.setenv("VLLM_SM70_DFLASH2_GDN_METADATA_SHADOW", "1") - qwen_gdn.envs.disable_envs_cache() - - builders = [ - _create_gdn_builder( - local_gdn_model, - num_speculative_tokens=7, - use_full_cuda_graph=True, - max_cudagraph_capture_size=16, - device=device, - ) - for _ in range(2) - ] - width = builders[0].num_spec_state_tokens + 1 - assert width == 8 - query_start_loc = torch.tensor([0, 8, 16], dtype=torch.int32, device=device) - query_start_loc_cpu = query_start_loc.cpu() - common_metadata = compute_common_gdn_attn_metadata( - num_decode_draft_tokens_cpu=torch.tensor([7, 7], dtype=torch.int32), - query_start_loc=query_start_loc, - query_start_loc_cpu=query_start_loc_cpu, - num_spec_state_tokens=7, - legacy_mixed_decode_routing=False, - ) - assert common_metadata is not None - - block_tables = ( - torch.arange(10, 10 + 2 * width, dtype=torch.int32, device=device).view( - 2, width - ), - torch.arange(100, 100 + 2 * width, dtype=torch.int32, device=device).view( - 2, width - ), - ) - accepted = torch.tensor([3, 5], dtype=torch.int32, device=device) - for builder in builders: - builder.spec_state_indices_tensor.fill_(777) - - first_result = prepare_dflash2_gdn_group_metadata( - builders_by_group=[(0, builders[0]), (1, builders[1])], - block_tables=block_tables, - common_gdn_metadata=common_metadata, - num_accepted_tokens=accepted, - num_actual_tokens=16, - descriptor=None, - ) - assert first_result is not None - prepared, descriptor = first_result - torch.accelerator.synchronize(device) - - for group_id, builder in enumerate(builders): - state = prepared[id(builder)].spec_state_indices_tensor - assert state is not None - torch.testing.assert_close(state[:2], block_tables[group_id], rtol=0, atol=0) - torch.testing.assert_close( - state[2:], torch.full_like(state[2:], PAD_SLOT_ID), rtol=0, atol=0 - ) - first = prepared[id(builders[0])] - assert first.spec_query_start_loc is not None - assert first.spec_query_start_loc.tolist() == [0, 8, 16] + [16] * 14 - assert first.num_accepted_tokens is not None - assert first.num_accepted_tokens.tolist() == [3, 5] + [1] * 14 - assert first.spec_state_slot_selectors is not None - assert first.spec_state_slot_selectors.tolist() == [3, 5] + [1] * 14 - assert first.spec_sequence_masks is not None - assert first.spec_sequence_masks.tolist() == [True, True] + [False] * 14 - - # Poison every persistent output and replay with new source values. This - # catches stale graph-tail data and accidental reuse of one group's table. - for builder in builders: - builder.spec_state_indices_tensor.fill_(888) - common_buffers = builders[0]._ddtree_fast_common_buffers - assert common_buffers is not None - common_buffers.spec_sequence_masks.fill_(True) - common_buffers.spec_query_start_loc.fill_(-1) - common_buffers.num_accepted_tokens.fill_(-1) - common_buffers.spec_state_slot_selectors.fill_(-1) - block_tables[0].add_(1000) - block_tables[1].add_(2000) - accepted.copy_(torch.tensor([7, 2], dtype=torch.int32, device=device)) - - second_result = prepare_dflash2_gdn_group_metadata( - builders_by_group=[(0, builders[0]), (1, builders[1])], - block_tables=block_tables, - common_gdn_metadata=common_metadata, - num_accepted_tokens=accepted, - num_actual_tokens=16, - descriptor=descriptor, - ) - assert second_result is not None - replayed, replay_descriptor = second_result - assert replay_descriptor is descriptor - assert replayed is prepared - torch.accelerator.synchronize(device) - for group_id, builder in enumerate(builders): - state = replayed[id(builder)].spec_state_indices_tensor - assert state is not None - torch.testing.assert_close(state[:2], block_tables[group_id], rtol=0, atol=0) - torch.testing.assert_close( - state[2:], torch.full_like(state[2:], PAD_SLOT_ID), rtol=0, atol=0 - ) - assert replayed[id(builders[0])].num_accepted_tokens is not None - assert replayed[id(builders[0])].num_accepted_tokens.tolist() == [7, 2] + [1] * 14 - - # The pointer tables remain valid when the graph batch shape changes, but - # the graph-buffer views do not. Ensure the cache key rebuilds their - # lengths instead of returning the steady-state B2 metadata. - single_common = compute_common_gdn_attn_metadata( - num_decode_draft_tokens_cpu=torch.tensor([7], dtype=torch.int32), - query_start_loc=torch.tensor([0, 8], dtype=torch.int32, device=device), - query_start_loc_cpu=torch.tensor([0, 8], dtype=torch.int32), - num_spec_state_tokens=7, - legacy_mixed_decode_routing=False, - ) - assert single_common is not None - single_result = prepare_dflash2_gdn_group_metadata( - builders_by_group=[(0, builders[0]), (1, builders[1])], - block_tables=(block_tables[0][:1], block_tables[1][:1]), - common_gdn_metadata=single_common, - num_accepted_tokens=accepted[:1], - num_actual_tokens=8, - descriptor=descriptor, - ) - assert single_result is not None - single_prepared, single_descriptor = single_result - assert single_descriptor is descriptor - assert single_prepared is not replayed - torch.accelerator.synchronize(device) - single = single_prepared[id(builders[0])] - assert single.spec_state_indices_tensor is not None - assert single.spec_state_indices_tensor.shape == (8, 8) - assert single.spec_query_start_loc is not None - assert single.spec_query_start_loc.tolist() == [0, 8] + [8] * 7 - assert single.num_accepted_tokens is not None - assert single.num_accepted_tokens.tolist() == [7] + [1] * 7 - - -def test_dflash2_fused_gdn_group_metadata_align_replay( - monkeypatch, - local_gdn_model, -): - if not torch.cuda.is_available(): - pytest.skip("CUDA is required for the fused GDN metadata kernel") - device = torch.device("cuda") - if torch.cuda.get_device_capability(device) != (7, 0): - pytest.skip("the fused GDN metadata kernel is SM70-only") - - monkeypatch.setenv("VLLM_SM70_QWEN_GDN_SPEC_CORE_OP", "0") - monkeypatch.setenv("VLLM_SM70_DFLASH2_FUSED_GDN_METADATA", "1") - monkeypatch.setenv("VLLM_SM70_DFLASH2_GDN_METADATA_SHADOW", "1") - qwen_gdn.envs.disable_envs_cache() - - builders = [ - _create_gdn_builder( - local_gdn_model, - num_speculative_tokens=7, - use_full_cuda_graph=True, - mamba_cache_mode="align", - max_cudagraph_capture_size=16, - device=device, - ) - for _ in range(2) - ] - width = builders[0].num_spec_state_tokens + 1 - query_start_loc = torch.tensor([0, 8, 16], dtype=torch.int32, device=device) - common_metadata = compute_common_gdn_attn_metadata( - num_decode_draft_tokens_cpu=torch.tensor([7, 7], dtype=torch.int32), - query_start_loc=query_start_loc, - query_start_loc_cpu=query_start_loc.cpu(), - num_spec_state_tokens=7, - legacy_mixed_decode_routing=False, - ) - assert common_metadata is not None - - table_columns = 32 - block_tables = ( - torch.arange( - 10, - 10 + 2 * table_columns, - dtype=torch.int32, - device=device, - ).view(2, table_columns), - torch.arange( - 100, - 100 + 2 * table_columns, - dtype=torch.int32, - device=device, - ).view(2, table_columns), - ) - # Batch rows deliberately map to non-adjacent persistent request slots. - req_index_mapping = torch.tensor([2, 0], dtype=torch.int32, device=device) - state_start_indices = torch.tensor([3, 12, 7, 18], dtype=torch.int32, device=device) - accepted = torch.tensor([3, 5], dtype=torch.int32, device=device) - - # Align metadata is incomplete without both persistent-state inputs. - assert ( - prepare_dflash2_gdn_group_metadata( - builders_by_group=[(0, builders[0]), (1, builders[1])], - block_tables=block_tables, - common_gdn_metadata=common_metadata, - num_accepted_tokens=accepted, - num_actual_tokens=16, - descriptor=None, - state_start_indices=state_start_indices, - ) - is None - ) - - first_result = prepare_dflash2_gdn_group_metadata( - builders_by_group=[(0, builders[0]), (1, builders[1])], - block_tables=block_tables, - common_gdn_metadata=common_metadata, - num_accepted_tokens=accepted, - num_actual_tokens=16, - descriptor=None, - state_start_indices=state_start_indices, - req_index_mapping=req_index_mapping, - ) - assert first_result is not None - prepared, descriptor = first_result - torch.accelerator.synchronize(device) - - starts = [7, 3] - for group_id, builder in enumerate(builders): - state = prepared[id(builder)].spec_state_indices_tensor - assert state is not None - expected = torch.stack( - [ - block_tables[group_id][row, start : start + width] - for row, start in enumerate(starts) - ] - ) - torch.testing.assert_close(state[:2], expected, rtol=0, atol=0) - torch.testing.assert_close( - state[2:], torch.full_like(state[2:], PAD_SLOT_ID), rtol=0, atol=0 - ) - - # Replay with the same persistent pointers but a different request mapping - # and state columns. CUDA graph replay must observe the live values. - for builder in builders: - builder.spec_state_indices_tensor.fill_(888) - # The model runner allocates this short mapping per batch. Its pointer is - # intentionally allowed to change without rebuilding the group descriptor. - req_index_mapping = torch.tensor([1, 3], dtype=torch.int32, device=device) - state_start_indices.copy_( - torch.tensor([2, 9, 5, 20], dtype=torch.int32, device=device) - ) - second_result = prepare_dflash2_gdn_group_metadata( - builders_by_group=[(0, builders[0]), (1, builders[1])], - block_tables=block_tables, - common_gdn_metadata=common_metadata, - num_accepted_tokens=accepted, - num_actual_tokens=16, - descriptor=descriptor, - state_start_indices=state_start_indices, - req_index_mapping=req_index_mapping, - ) - assert second_result is not None - replayed, replay_descriptor = second_result - assert replay_descriptor is descriptor - assert replayed is prepared - torch.accelerator.synchronize(device) - - starts = [9, 20] - for group_id, builder in enumerate(builders): - state = replayed[id(builder)].spec_state_indices_tensor - assert state is not None - expected = torch.stack( - [ - block_tables[group_id][row, start : start + width] - for row, start in enumerate(starts) - ] - ) - torch.testing.assert_close(state[:2], expected, rtol=0, atol=0) - torch.testing.assert_close( - state[2:], torch.full_like(state[2:], PAD_SLOT_ID), rtol=0, atol=0 - ) - - def test_full_cuda_graph_capture_single_token_decode_is_not_spec(local_gdn_model): builder = _create_gdn_builder( local_gdn_model, @@ -1054,7 +468,9 @@ def test_spec_state_slot_selector_can_differ_from_accepted_count(local_gdn_model common_prefix_len=0, common_attn_metadata=common, num_accepted_tokens=torch.tensor([2], dtype=torch.int32, device=DEVICE), - spec_state_slot_selectors=torch.tensor([4], dtype=torch.int32, device=DEVICE), + spec_state_slot_selectors=torch.tensor( + [4], dtype=torch.int32, device=DEVICE + ), num_decode_draft_tokens_cpu=torch.tensor([4], dtype=torch.int32, device="cpu"), ) @@ -1333,7 +749,9 @@ def test_gdn_state_contract_debug_assert_rejects_pad_accepted_slot( spec_sequence_masks_cpu=torch.tensor( [True], dtype=torch.bool, device="cpu" ), - num_accepted_tokens=torch.tensor([5], dtype=torch.int32, device=DEVICE), + num_accepted_tokens=torch.tensor( + [5], dtype=torch.int32, device=DEVICE + ), current_state_block_ids=None, is_mamba_cache_all=False, ) @@ -1360,7 +778,9 @@ def test_spec_commit_pure_decode_consumes_padded_graph_rows_without_metadata_pat num_spec_decodes=2, num_spec_decode_tokens=10, num_actual_tokens=10, - spec_query_start_loc=torch.tensor([0, 5, 10], dtype=torch.int32, device=DEVICE), + spec_query_start_loc=torch.tensor( + [0, 5, 10], dtype=torch.int32, device=DEVICE + ), spec_state_indices_tensor=torch.tensor( [[10, 11, 12, 13, 14], [20, 21, 22, 23, 24]], dtype=torch.int32, @@ -1401,6 +821,7 @@ def test_spec_commit_pure_decode_consumes_padded_graph_rows_without_metadata_pat seen: dict[str, object] = {} class FakeLayer: + def __init__(self) -> None: self.conv1d = torch.nn.Identity() self.conv1d.weight = torch.empty((8, 1, 1), device=DEVICE) @@ -1441,11 +862,10 @@ def fake_conv_update( seen["conv"] = True return x + 1 - def fake_gating(a_log, a_tensor, b_tensor, dt_bias, *, beta_dtype=None): + def fake_gating(a_log, a_tensor, b_tensor, dt_bias): del a_log, dt_bias assert a_tensor is a assert b_tensor is b - assert beta_dtype is torch.float32 seen["gating"] = True return a_tensor + 2, b_tensor + 3 @@ -1603,7 +1023,9 @@ def test_sm70_qwen_gdn_blocks_003_spec_core_only_for_deep_native_mtp( expected, ): monkeypatch.setenv("VLLM_SM70_QWEN_GDN_003_SPEC_CORE_OP", "1") - monkeypatch.delenv("VLLM_SM70_QWEN_GDN_003_SPEC_ALLOW_DEEP_MTP", raising=False) + monkeypatch.delenv( + "VLLM_SM70_QWEN_GDN_003_SPEC_ALLOW_DEEP_MTP", raising=False + ) _clear_envs_cache() vllm_config = SimpleNamespace( @@ -1614,7 +1036,9 @@ def test_sm70_qwen_gdn_blocks_003_spec_core_only_for_deep_native_mtp( ) assert ( - qwen_gdn._sm70_qwen_gdn_block_003_spec_for_deep_native_mtp(vllm_config) + qwen_gdn._sm70_qwen_gdn_block_003_spec_for_deep_native_mtp( + vllm_config + ) is expected ) @@ -1633,7 +1057,9 @@ def test_sm70_qwen_gdn_003_spec_core_deep_mtp_override_is_diagnostic( ) ) - assert not qwen_gdn._sm70_qwen_gdn_block_003_spec_for_deep_native_mtp(vllm_config) + assert not qwen_gdn._sm70_qwen_gdn_block_003_spec_for_deep_native_mtp( + vllm_config + ) def test_sm70_qwen_gdn_003_spec_core_route_has_priority( diff --git a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py index 7c7d535f08..159bb8af3f 100644 --- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py +++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py @@ -63,9 +63,6 @@ def test_indexer_builder_deepseek_v4_compressed_slot_mapping_uses_storage_block_ md = builder.build(common_prefix_len=0, common_attn_metadata=common) - assert md.max_seq_len == 280 - assert md.compressed_max_seq_len == 70 - # The compressed slot_mapping retains the original uncompressed size (40). # Only every compress_ratio-th position gets a valid slot; the rest are -1. assert md.slot_mapping.numel() == 40 diff --git a/tests/v1/attention/test_prefix_anchored_swa_mask.py b/tests/v1/attention/test_prefix_anchored_swa_mask.py deleted file mode 100644 index 0e6a744e3a..0000000000 --- a/tests/v1/attention/test_prefix_anchored_swa_mask.py +++ /dev/null @@ -1,181 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for the prefix-anchored sliding-window attention mask. - -Semantics under test: keys inside the prompt/prefix (``kv < anchor_len``) -stay globally visible; generated keys are additionally restricted to a -sliding window (``q_abs - kv < window``); everything is AND-ed with the -causal mask. - -The CPU tests validate the mask formula (the same one the masked kernels -implement) against plain causal attention with a torch reference. -""" - -from types import SimpleNamespace -from typing import Any, cast - -import pytest -import torch - -from vllm.v1.attention.backend import AttentionType -from vllm.v1.attention.backends.flash_attn_v100 import ( - FlashAttnV100Impl, - FlashAttnV100MetadataBuilder, -) - - -def anchored_swa_mask( - seq_len: int, anchor_len: int, window: int, device: str = "cpu" -) -> torch.Tensor: - """Boolean keep-mask [seq_len, seq_len]; True = attend. - - Mirrors the kernel formula: - ``causal AND (kv < anchor_len OR q_abs - kv < window)``. - """ - q = torch.arange(seq_len, device=device)[:, None] - kv = torch.arange(seq_len, device=device)[None, :] - causal = kv <= q - in_prefix = kv < anchor_len - in_window = (q - kv) < window - return causal & (in_prefix | in_window) - - -def causal_mask(seq_len: int, device: str = "cpu") -> torch.Tensor: - q = torch.arange(seq_len, device=device)[:, None] - kv = torch.arange(seq_len, device=device)[None, :] - return kv <= q - - -def ref_attention( - q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor -) -> torch.Tensor: - """Plain fp32 reference attention. q/k/v: [heads, seq, dim].""" - scale = q.shape[-1] ** -0.5 - scores = torch.einsum("hqd,hkd->hqk", q.float(), k.float()) * scale - scores = scores.masked_fill(~mask[None], float("-inf")) - probs = torch.softmax(scores, dim=-1) - return torch.einsum("hqk,hkd->hqd", probs, v.float()) - - -def _flash_impl(window: int | None) -> FlashAttnV100Impl: - impl = object.__new__(FlashAttnV100Impl) - object.__setattr__(impl, "prefix_anchored_decode_window", window) - object.__setattr__(impl, "attn_type", AttentionType.DECODER) - object.__setattr__(impl, "kv_cache_dtype", "auto") - return impl - - -@pytest.mark.cpu_test -def test_backend_default_off_does_not_inspect_metadata(): - class NoMetadataAccess: - def __getattr__(self, name: str): - raise AssertionError(f"default-off path inspected {name}") - - metadata = cast(Any, NoMetadataAccess()) - assert _flash_impl(None)._anchored_swa_params(metadata) == (None, 0) - - -@pytest.mark.cpu_test -def test_backend_enabled_fails_closed_without_matching_metadata(): - with pytest.raises(RuntimeError, match="metadata does not match"): - _flash_impl(128)._anchored_swa_params(cast(Any, object())) - - -@pytest.mark.cpu_test -def test_metadata_builder_attaches_prefix_lengths_for_eager_path(): - builder = object.__new__(FlashAttnV100MetadataBuilder) - object.__setattr__(builder, "decode_sliding_window", 128) - object.__setattr__(builder, "device", torch.device("cpu")) - object.__setattr__( - builder, - "persistent_prefix_anchor_lens", - torch.empty(4, dtype=torch.int32), - ) - metadata = SimpleNamespace() - common = SimpleNamespace( - prefix_anchor_lens=torch.tensor([7, 11], dtype=torch.int64), - num_reqs=2, - ) - - builder._attach_prefix_anchored_metadata( - cast(Any, metadata), - cast(Any, common), - ) - - assert metadata.decode_sliding_window == 128 - assert metadata.prefix_anchor_lens.dtype == torch.int32 - assert metadata.prefix_anchor_lens.tolist() == [7, 11] - - -@pytest.mark.cpu_test -def test_metadata_builder_fails_closed_without_prefix_lengths(): - builder = object.__new__(FlashAttnV100MetadataBuilder) - object.__setattr__(builder, "decode_sliding_window", 128) - object.__setattr__(builder, "device", torch.device("cpu")) - object.__setattr__( - builder, - "persistent_prefix_anchor_lens", - torch.empty(4, dtype=torch.int32), - ) - - with pytest.raises(RuntimeError, match="requires per-request prefix lengths"): - builder._attach_prefix_anchored_metadata( - cast(Any, SimpleNamespace()), - cast(Any, SimpleNamespace(prefix_anchor_lens=None, num_reqs=2)), - ) - - -@pytest.mark.cpu_test -def test_mask_equals_causal_while_within_window(): - # While the number of generated tokens is <= window, no key falls in the - # gap, so the mask must be exactly the causal mask. - anchor, window = 6, 8 - for seq_len in range(1, anchor + window + 1): - assert torch.equal( - anchored_swa_mask(seq_len, anchor, window), - causal_mask(seq_len), - ), f"seq_len={seq_len}" - - -@pytest.mark.cpu_test -def test_mask_prunes_gap_beyond_window(): - anchor, window = 6, 8 - seq_len = anchor + window + 3 # 3 generated keys fall out of the window - mask = anchored_swa_mask(seq_len, anchor, window) - causal = causal_mask(seq_len) - assert not torch.equal(mask, causal) - # The last query keeps: the full prefix and the window of recent keys... - last = mask[-1] - assert last[:anchor].all(), "prefix must stay globally visible" - assert last[seq_len - window :].all(), "recent window must be visible" - # ...and drops exactly the gap keys in between. - gap = last[anchor : seq_len - window] - assert not gap.any(), "gap keys must be masked" - # Prefix queries (inside the prompt) are never affected. - assert torch.equal(mask[: anchor + window], causal[: anchor + window]) - - -@pytest.mark.cpu_test -def test_attention_output_equivalence_and_divergence(): - torch.manual_seed(0) - heads, dim = 2, 32 - anchor, window = 6, 8 - - # Within the window: identical outputs. - seq_len = anchor + window - q, k, v = (torch.randn(heads, seq_len, dim) for _ in range(3)) - out_anchored = ref_attention(q, k, v, anchored_swa_mask(seq_len, anchor, window)) - out_causal = ref_attention(q, k, v, causal_mask(seq_len)) - torch.testing.assert_close(out_anchored, out_causal) - - # Beyond the window: the late queries must diverge from full attention - # (negative control proving the mask actually removes information). - seq_len = anchor + 2 * window - q, k, v = (torch.randn(heads, seq_len, dim) for _ in range(3)) - out_anchored = ref_attention(q, k, v, anchored_swa_mask(seq_len, anchor, window)) - out_causal = ref_attention(q, k, v, causal_mask(seq_len)) - assert not torch.allclose(out_anchored[:, -1], out_causal[:, -1]) - # Queries that see no gap still match exactly. - torch.testing.assert_close( - out_anchored[:, : anchor + window], out_causal[:, : anchor + window] - ) diff --git a/tests/v1/attention/test_sm70_flash_v100_policy.py b/tests/v1/attention/test_sm70_flash_v100_policy.py index fb2e195796..f16ce3800c 100644 --- a/tests/v1/attention/test_sm70_flash_v100_policy.py +++ b/tests/v1/attention/test_sm70_flash_v100_policy.py @@ -237,35 +237,6 @@ def test_sm70_fa2_d256_prefill_env_is_default_on(monkeypatch): assert envs.VLLM_FLASH_V100_FA2_D256_PREFILL is False -def test_sm70_d256_gqa_architecture_env_is_default_on(monkeypatch): - import vllm.envs as envs - - name = "VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL" - monkeypatch.delenv(name, raising=False) - envs.disable_envs_cache() - assert envs.VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL is True - - monkeypatch.setenv(name, "0") - envs.disable_envs_cache() - assert envs.VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL is False - - -def test_sm70_e4m3_batch_xqa_env_contract(monkeypatch): - import vllm.envs as envs - - monkeypatch.delenv("VLLM_FLASH_V100_E4M3_BATCH_XQA", raising=False) - monkeypatch.delenv("VLLM_FLASH_V100_E4M3_BATCH_XQA_OPTIMIZED", raising=False) - envs.disable_envs_cache() - assert envs.VLLM_FLASH_V100_E4M3_BATCH_XQA is True - assert envs.VLLM_FLASH_V100_E4M3_BATCH_XQA_OPTIMIZED is True - - monkeypatch.setenv("VLLM_FLASH_V100_E4M3_BATCH_XQA", "0") - monkeypatch.setenv("VLLM_FLASH_V100_E4M3_BATCH_XQA_OPTIMIZED", "0") - envs.disable_envs_cache() - assert envs.VLLM_FLASH_V100_E4M3_BATCH_XQA is False - assert envs.VLLM_FLASH_V100_E4M3_BATCH_XQA_OPTIMIZED is False - - def test_sm70_e5m2_decode_fast_route_envs_are_default_on(monkeypatch): import vllm.envs as envs @@ -274,7 +245,6 @@ def test_sm70_e5m2_decode_fast_route_envs_are_default_on(monkeypatch): "VLLM_FLASH_V100_XQA_E5M2_G6_SPLIT_REDUCE", "VLLM_FLASH_V100_XQA_E5M2_PARTITION_PAGE_IDS", "VLLM_FLASH_V100_XQA_E5M2_PAIR_LOAD", - "VLLM_FLASH_V100_XQA_E5M2_BATCH_WIDE_LOAD", ) for name in names: monkeypatch.delenv(name, raising=False) @@ -284,7 +254,6 @@ def test_sm70_e5m2_decode_fast_route_envs_are_default_on(monkeypatch): assert envs.VLLM_FLASH_V100_XQA_E5M2_G6_SPLIT_REDUCE is True assert envs.VLLM_FLASH_V100_XQA_E5M2_PARTITION_PAGE_IDS is True assert envs.VLLM_FLASH_V100_XQA_E5M2_PAIR_LOAD is True - assert envs.VLLM_FLASH_V100_XQA_E5M2_BATCH_WIDE_LOAD is True assert envs.VLLM_FLASH_V100_XQA_E5M2_P1024_BEGIN == 61633 for name in names: @@ -294,7 +263,6 @@ def test_sm70_e5m2_decode_fast_route_envs_are_default_on(monkeypatch): assert envs.VLLM_FLASH_V100_XQA_E5M2_G6_SPLIT_REDUCE is False assert envs.VLLM_FLASH_V100_XQA_E5M2_PARTITION_PAGE_IDS is False assert envs.VLLM_FLASH_V100_XQA_E5M2_PAIR_LOAD is False - assert envs.VLLM_FLASH_V100_XQA_E5M2_BATCH_WIDE_LOAD is False monkeypatch.setenv("VLLM_FLASH_V100_XQA_E5M2_P1024_BEGIN", "49152") envs.disable_envs_cache() @@ -408,15 +376,10 @@ def test_sm70_prefill_dense_splitkv3_env_is_default_on(monkeypatch): monkeypatch.delenv("VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3", raising=False) monkeypatch.delenv("VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3_MIN_KV", raising=False) - monkeypatch.delenv( - "VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3_Q8000_EXPERIMENTAL", - raising=False, - ) envs.disable_envs_cache() assert envs.VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3 is True assert envs.VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3_MIN_KV == 32768 - assert envs.VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3_Q8000_EXPERIMENTAL is False monkeypatch.setenv("VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3", "0") envs.disable_envs_cache() @@ -617,7 +580,6 @@ def test_prefix_prefill_prioritizes_gathered_exact_dense_over_paged( impl.head_size = 256 impl.scale = 0.0625 impl.sliding_window = None - impl.prefix_anchored_decode_window = None impl.kv_cache_dtype = "auto" impl.use_flash_v100_decode = False impl.use_decode_paged_prefill = False @@ -738,36 +700,6 @@ def test_sm70_splitd_d256_loader_requires_exact_ops(monkeypatch): assert flash_v100._get_sm70_splitd_d256_ops() == (dense, paged, splitkv3) -def test_sm70_d256_gqa_architecture_loader_is_optional(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, - ) - - architecture = object() - fake_ops = SimpleNamespace( - _vllm_fa2_C=SimpleNamespace( - sm70_d256_gqa_architecture_fwd=architecture, - ) - ) - monkeypatch.setattr(flash_v100, "torch", SimpleNamespace(ops=fake_ops)) - monkeypatch.setattr( - flash_v100, - "_sm70_d256_gqa_architecture_op_checked", - False, - ) - monkeypatch.setattr(flash_v100, "_sm70_d256_gqa_architecture_op", None) - - assert flash_v100._get_sm70_d256_gqa_architecture_op() is architecture - - def test_prefill_dense_splitkv3_workspace_reuses_exact_shape(monkeypatch): import vllm.v1.attention.backends.flash_attn_v100 as flash_v100 @@ -811,8 +743,6 @@ def test_prefill_dense_splitkv3_policy_is_exact_shape_bounded(monkeypatch): key = torch.empty((1, 65536, 1, 256), dtype=torch.float16, device="meta") key_32k = torch.empty((1, 32768, 1, 256), dtype=torch.float16, device="meta") key_16k = torch.empty((1, 16384, 1, 256), dtype=torch.float16, device="meta") - query_8k = torch.empty((1, 8000, 6, 256), dtype=torch.float16, device="meta") - query_8192 = torch.empty((1, 8192, 6, 256), dtype=torch.float16, device="meta") assert flash_v100._should_use_prefill_dense_splitkv3( query, @@ -842,168 +772,6 @@ def test_prefill_dense_splitkv3_policy_is_exact_shape_bounded(monkeypatch): max_seqlen_k=65536, splitkv3_op=object(), ) - assert not flash_v100._should_use_prefill_dense_splitkv3( - query_8k, - key, - max_seqlen_q=8000, - max_seqlen_k=65536, - splitkv3_op=object(), - ) - - monkeypatch.setenv("VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3_Q8000_EXPERIMENTAL", "1") - envs.disable_envs_cache() - assert flash_v100._should_use_prefill_dense_splitkv3( - query_8k, - key, - max_seqlen_q=8000, - max_seqlen_k=65536, - splitkv3_op=object(), - ) - assert not flash_v100._should_use_prefill_dense_splitkv3( - query_8192, - key, - max_seqlen_q=8192, - max_seqlen_k=65536, - splitkv3_op=object(), - ) - - -def test_prefill_d256_gqa_architecture_policy_is_shape_family_bounded(monkeypatch): - import vllm.envs as envs - import vllm.v1.attention.backends.flash_attn_v100 as flash_v100 - - name = "VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL" - query = torch.empty((1, 8000, 6, 256), dtype=torch.float16, device="meta") - key = torch.empty((1, 128000, 1, 256), dtype=torch.float16, device="meta") - value = torch.empty_like(key) - - monkeypatch.setenv(name, "0") - envs.disable_envs_cache() - assert not flash_v100._should_use_prefill_d256_gqa_architecture( - query, - key, - value, - max_seqlen_q=8000, - max_seqlen_k=128000, - softmax_scale=0.0625, - architecture_op=object(), - ) - - monkeypatch.delenv(name, raising=False) - envs.disable_envs_cache() - for kv_len in range(40000, 128001, 8000): - family_key = torch.empty( - (1, kv_len, 1, 256), dtype=torch.float16, device="meta" - ) - assert flash_v100._should_use_prefill_d256_gqa_architecture( - query, - family_key, - torch.empty_like(family_key), - max_seqlen_q=8000, - max_seqlen_k=kv_len, - softmax_scale=0.0625, - architecture_op=object(), - ) - assert not flash_v100._should_use_prefill_d256_gqa_architecture( - query[:, :7999], - key, - value, - max_seqlen_q=7999, - max_seqlen_k=128000, - softmax_scale=0.0625, - architecture_op=object(), - ) - assert not flash_v100._should_use_prefill_d256_gqa_architecture( - query, - key[:, :127999], - value[:, :127999], - max_seqlen_q=8000, - max_seqlen_k=127999, - softmax_scale=0.0625, - architecture_op=object(), - ) - assert not flash_v100._should_use_prefill_d256_gqa_architecture( - query, - key[:, :44000], - value[:, :44000], - max_seqlen_q=8000, - max_seqlen_k=44000, - softmax_scale=0.0625, - architecture_op=object(), - ) - assert not flash_v100._should_use_prefill_d256_gqa_architecture( - query, - key, - value, - max_seqlen_q=8000, - max_seqlen_k=128000, - softmax_scale=1.0, - architecture_op=object(), - ) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -def test_prefill_d256_gqa_architecture_oom_uses_dense_fallback(monkeypatch): - import vllm.envs as envs - import vllm.v1.attention.backends.flash_attn_v100 as flash_v100 - - if torch.cuda.get_device_capability() != (7, 0): - pytest.skip("SM70/V100 is required") - - name = "VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL" - monkeypatch.delenv(name, raising=False) - envs.disable_envs_cache() - monkeypatch.setattr(flash_v100, "_is_cuda_graph_capturing", lambda _: False) - - dense_calls = 0 - - def dense_op(query, key, value, out, softmax_scale, causal): - nonlocal dense_calls - dense_calls += 1 - out.fill_(3) - return out - - def architecture_oom(*args, **kwargs): - raise torch.OutOfMemoryError("expected architecture workspace OOM") - - monkeypatch.setattr( - flash_v100, - "_get_sm70_splitd_d256_ops", - lambda: (dense_op, object(), None), - ) - monkeypatch.setattr( - flash_v100, - "_get_sm70_d256_gqa_architecture_op", - lambda: architecture_oom, - ) - monkeypatch.setattr(flash_v100, "_warned_prefill_d256_gqa_architecture_oom", False) - - query = torch.zeros((1, 8000, 6, 256), dtype=torch.float16, device="cuda") - key = torch.zeros((1, 40000, 1, 256), dtype=torch.float16, device="cuda") - value = torch.zeros_like(key) - out = torch.zeros_like(query) - cu_seqlens_q = torch.tensor([0, 8000], dtype=torch.int32, device="cuda") - cu_seqlens_k = torch.tensor([0, 40000], dtype=torch.int32, device="cuda") - - result = flash_v100._try_sm70_fa2_d256_prefill( - query, - key, - value, - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_k=cu_seqlens_k, - max_seqlen_q=8000, - max_seqlen_k=40000, - softmax_scale=0.0625, - causal=True, - window_size=(-1, -1), - out=out, - ) - - assert result is not None - assert result.data_ptr() == out.data_ptr() - assert dense_calls == 1 - assert torch.all(out == 3) - assert flash_v100._warned_prefill_d256_gqa_architecture_oom is True def test_dense_prefill_restores_uniform_batch_view_for_exact_splitd(monkeypatch): @@ -1528,27 +1296,6 @@ def test_sm70_nomtp_cudagraph_capture_sizes_cover_concurrency( assert _sm70_nomtp_cudagraph_capture_sizes(max_num_seqs) == expected -@pytest.mark.parametrize( - ("max_num_seqs", "expected"), - [ - (1, [5]), - (2, [5, 10]), - (4, [5, 10, 20]), - (6, [5, 10, 20, 30]), - (12, [5, 10, 20, 30, 40, 60]), - (16, [5, 10, 20, 30, 40, 60, 80]), - (32, [5, 10, 20, 30, 40, 60, 80]), - ], -) -def test_sm70_mtp_cudagraph_capture_sizes_cover_production_concurrency( - max_num_seqs: int, - expected: list[int], -): - from vllm.config.vllm import _sm70_mtp_cudagraph_capture_sizes - - assert _sm70_mtp_cudagraph_capture_sizes(max_num_seqs, 5) == expected - - def test_flash_v100_decode_query_does_not_attach_smallq_metadata( monkeypatch, local_flash_v100_model, @@ -1655,89 +1402,6 @@ def fake_decode( assert torch.all(output == 1) -def test_flash_v100_dflash2_grouped_verify_uses_original_request_metadata(): - from vllm.v1.attention.backends.flash_attn_v100 import FlashAttnV100Impl - - impl = FlashAttnV100Impl( - num_heads=6, - head_size=256, - scale=1.0, - num_kv_heads=1, - alibi_slopes=None, - sliding_window=None, - kv_cache_dtype="fp8_e5m2", - ) - impl.use_dflash2_grouped_verify = True - captured: dict[str, object] = {} - - def grouped_verify( - query, - key_cache, - value_cache, - block_table, - seq_lens, - **kwargs, - ): - captured["block_table"] = block_table - captured["seq_lens"] = seq_lens - captured["one_pass"] = kwargs["one_pass"] - kwargs["out"].fill_(1) - - impl.flash_attn_grouped_verify_paged = grouped_verify - impl.flash_attn_decode_paged = lambda *args, **kwargs: pytest.fail( - "exact DFlash2 grouped route fell through to independent rows" - ) - 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, - causal=True, - is_dflash_selector_target=True, - max_model_len=32768, - query_start_loc=torch.tensor([0, 8], 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), - ) - layer = SimpleNamespace(_k_scale_float=0.5, _v_scale_float=2.0) - query = torch.zeros((8, 6, 256), dtype=torch.float16) - output = torch.zeros_like(query) - key_cache = torch.zeros((2, 3296, 1, 256), dtype=torch.uint8) - value_cache = torch.zeros_like(key_cache) - - attn_metadata.max_model_len = 8192 - assert not impl._dflash2_grouped_verify_allowed( - query, - key_cache, - value_cache, - attn_metadata, - num_query_tokens=8, - ) - attn_metadata.max_model_len = 32768 - result = impl._flash_v100_small_query_prefill_as_decode( - layer, - query, - key_cache, - value_cache, - attn_metadata, - output, - attn_metadata.query_start_loc, - original_seq_lens, - ) - - assert result is output - captured_block_table = captured["block_table"] - captured_seq_lens = captured["seq_lens"] - assert isinstance(captured_block_table, torch.Tensor) - assert isinstance(captured_seq_lens, torch.Tensor) - assert captured_block_table.data_ptr() == original_block_table.data_ptr() - assert captured_seq_lens.data_ptr() == original_seq_lens.data_ptr() - assert captured["one_pass"] is True - assert torch.all(output == 1) - - def test_flash_v100_decode_forwards_shape_hints(): from vllm.v1.attention.backends.flash_attn_v100 import FlashAttnV100Impl @@ -1852,126 +1516,6 @@ def fail_scalar(*args, **kwargs): assert torch.all(output == 1) -def test_flash_v100_decode_uses_xqa_for_e4m3_g6_d256(monkeypatch): - from vllm.v1.attention.backends.flash_attn_v100 import FlashAttnV100Impl - - monkeypatch.delenv("VLLM_FLASH_V100_DECODE_USE_XQA", raising=False) - monkeypatch.delenv("VLLM_FLASH_V100_DECODE_PARTITION_SIZE", raising=False) - - impl = FlashAttnV100Impl( - num_heads=6, - head_size=256, - scale=1.0, - num_kv_heads=1, - alibi_slopes=None, - sliding_window=None, - kv_cache_dtype="fp8_e4m3", - ) - calls: list[tuple[str, int | None, str | None]] = [] - - def hit_xqa(*args, **kwargs): - calls.append( - ( - "xqa", - kwargs.get("partition_size_hint"), - kwargs.get("kv_cache_dtype"), - ) - ) - kwargs["out"].fill_(1) - - def fail_scalar(*args, **kwargs): - raise AssertionError("E4M3 G6/D256 decode should select XQA") - - impl.flash_attn_decode_paged_xqa = hit_xqa # type: ignore[method-assign] - impl.flash_attn_decode_paged = fail_scalar # type: ignore[method-assign] - attn_metadata = SimpleNamespace( - num_actual_tokens=1, - block_table=torch.tensor([[0]], dtype=torch.int32), - seq_lens=torch.tensor([1025], dtype=torch.int32), - flash_v100_decode_max_seq_len_hint=1025, - flash_v100_decode_workspace_seq_capacity_hint=262144, - flash_v100_decode_active_num_partitions=torch.tensor([5], dtype=torch.int32), - ) - layer = SimpleNamespace(_k_scale_float=0.04, _v_scale_float=0.25) - query = torch.zeros((1, 6, 256), dtype=torch.float16) - output = torch.zeros_like(query) - kv_cache = torch.zeros((2, 2, 1568, 1, 256), dtype=torch.uint8) - - result = impl._flash_v100_decode( - layer, - query, - query, - query, - kv_cache, - attn_metadata, - output, - ) - - assert result is output - assert calls == [("xqa", 64, "fp8_e4m3")] - assert torch.all(output == 1) - - -@pytest.mark.parametrize( - ("batch_size", "enabled", "expected_route"), - ((2, False, "scalar"), (2, True, "xqa"), (16, True, "xqa"), (17, True, "scalar")), -) -def test_flash_v100_e4m3_batched_xqa_is_exactly_gated( - monkeypatch, batch_size, enabled, expected_route -): - from vllm.v1.attention.backends.flash_attn_v100 import FlashAttnV100Impl - - monkeypatch.setenv("VLLM_FLASH_V100_E4M3_BATCH_XQA", "1" if enabled else "0") - monkeypatch.delenv("VLLM_FLASH_V100_DECODE_PARTITION_SIZE", raising=False) - impl = FlashAttnV100Impl( - num_heads=6, - head_size=256, - scale=1.0, - num_kv_heads=1, - alibi_slopes=None, - sliding_window=None, - kv_cache_dtype="fp8_e4m3", - ) - calls: list[str] = [] - - def hit_xqa(*args, **kwargs): - calls.append("xqa") - kwargs["out"].fill_(1) - - def hit_scalar(*args, **kwargs): - calls.append("scalar") - kwargs["out"].fill_(1) - - impl.flash_attn_decode_paged_xqa = hit_xqa # type: ignore[method-assign] - impl.flash_attn_decode_paged = hit_scalar # type: ignore[method-assign] - attn_metadata = SimpleNamespace( - num_actual_tokens=batch_size, - block_table=torch.zeros((batch_size, 2), dtype=torch.int32), - seq_lens=torch.full((batch_size,), 16384, dtype=torch.int32), - flash_v100_decode_max_seq_len_hint=16384, - flash_v100_decode_workspace_seq_capacity_hint=32768, - flash_v100_decode_active_num_partitions=torch.tensor([64], dtype=torch.int32), - ) - layer = SimpleNamespace(_k_scale_float=0.04, _v_scale_float=0.25) - query = torch.zeros((batch_size, 6, 256), dtype=torch.float16) - output = torch.zeros_like(query) - kv_cache = torch.zeros((2, 2, 1568, 1, 256), dtype=torch.uint8) - - result = impl._flash_v100_decode( - layer, - query, - query, - query, - kv_cache, - attn_metadata, - output, - ) - - assert result is output - assert calls == [expected_route] - assert torch.all(output == 1) - - @pytest.mark.parametrize( ("num_heads", "seq_len", "expected_route"), ( @@ -2179,113 +1723,6 @@ def test_flash_v100_fp8_xqa_graph_capture_uses_static_context_hint( assert mod._decode_fp8_xqa_allowed(metadata, torch.empty(1)) is expected -@pytest.mark.parametrize( - ("routing_enabled", "graph_variant", "expected"), - ( - (True, None, True), - (True, 0, False), - (True, -1, True), - (False, None, False), - (False, -1, False), - ), -) -def test_flash_v100_batch_context_routing_isolated_by_graph_variant( - routing_enabled, - graph_variant, - expected, -): - from vllm.v1.attention.backends import flash_attn_v100 as mod - - assert ( - mod._batch_context_routing_for_graph_variant( - routing_enabled, - graph_variant, - ) - is expected - ) - - -@pytest.mark.parametrize( - ("cache_dtype", "e4m3_enabled", "expected"), - ( - ("fp8_e5m2", False, True), - ("fp8_e4m3", True, True), - ("fp8", True, True), - ("fp8_e4m3", False, False), - ("auto", True, False), - ), -) -def test_flash_v100_batch_context_routing_accepts_exact_fp8_xqa_formats( - monkeypatch, - cache_dtype, - e4m3_enabled, - expected, -): - from vllm.v1.attention.backends import flash_attn_v100 as mod - - monkeypatch.setenv("VLLM_FLASH_V100_E4M3_BATCH_XQA", "1" if e4m3_enabled else "0") - assert mod._batch_context_routing_cache_dtype_supported(cache_dtype) is expected - - -@pytest.mark.parametrize("routing_enabled", (True, False)) -def test_flash_v100_fp8_xqa_full_capacity_graph_preserves_baseline_xqa( - monkeypatch, - routing_enabled, -): - from vllm.v1.attention.backends.flash_attn_v100 import FlashAttnV100Impl - - impl = FlashAttnV100Impl( - num_heads=6, - head_size=256, - scale=1.0, - num_kv_heads=1, - alibi_slopes=None, - sliding_window=None, - kv_cache_dtype="fp8_e5m2", - ) - calls: list[tuple[str, bool]] = [] - - def hit_xqa(*args, **kwargs): - calls.append(("xqa", kwargs.get("batch_context_routing", False))) - kwargs["out"].fill_(1) - - def hit_scalar(*args, **kwargs): - calls.append(("scalar", False)) - kwargs["out"].fill_(1) - - impl.flash_attn_decode_paged_xqa = hit_xqa # type: ignore[method-assign] - impl.flash_attn_decode_paged = hit_scalar # type: ignore[method-assign] - attn_metadata = SimpleNamespace( - num_actual_tokens=4, - block_table=torch.zeros((4, 256), dtype=torch.int32), - seq_lens=torch.full((4,), 1, dtype=torch.int32), - flash_v100_cudagraph_capture=True, - flash_v100_batch_context_routing=routing_enabled, - flash_v100_decode_max_seq_len_hint=1, - flash_v100_static_decode_seq_hint=262144, - flash_v100_decode_workspace_seq_capacity_hint=262144, - flash_v100_decode_active_num_partitions=torch.tensor([16], dtype=torch.int32), - ) - layer = SimpleNamespace(_k_scale_float=1.0, _v_scale_float=1.0) - query = torch.zeros((4, 6, 256), dtype=torch.float16) - output = torch.zeros_like(query) - kv_cache = torch.zeros((2, 2, 16, 1, 256), dtype=torch.uint8) - - result = impl._flash_v100_decode( - layer, - query, - query, - query, - kv_cache, - attn_metadata, - output, - ) - - assert result is output - assert calls == [("xqa", routing_enabled)] - assert torch.all(output == 1) - - def test_flash_v100_mtp5_dual_cta_partition_policy(monkeypatch): from vllm.v1.attention.backends import flash_attn_v100 as mod @@ -2642,7 +2079,7 @@ def test_flash_v100_decode_e5m2_keeps_small_pages_on_default_planner(monkeypatch ) -def test_flash_v100_decode_e4m3_uses_exact_p64_partition_hint(monkeypatch): +def test_flash_v100_decode_e4m3_does_not_use_e5m2_partition_hint(monkeypatch): from vllm.v1.attention.backends.flash_attn_v100 import ( _g6_aligned_page_partition_size_hint, ) @@ -2654,7 +2091,7 @@ def test_flash_v100_decode_e4m3_uses_exact_p64_partition_hint(monkeypatch): assert ( _g6_aligned_page_partition_size_hint(query, key_cache, key_cache, "fp8_e4m3") - == 64 + is None ) diff --git a/tests/v1/core/test_async_scheduler.py b/tests/v1/core/test_async_scheduler.py index e4ea646ed9..eaf4ef82ad 100644 --- a/tests/v1/core/test_async_scheduler.py +++ b/tests/v1/core/test_async_scheduler.py @@ -306,9 +306,6 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): request.status = RequestStatus.RUNNING request.num_computed_tokens = request.num_tokens request.num_output_placeholders = 1 - # This fixture calls update_from_output without schedule(); mirror the - # scheduled token that would normally be recorded by _update_after_schedule. - request.num_in_flight_tokens = 1 scheduler.perf_metrics = None scheduler.connector = None diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 1ec844fead..94e9f6f4c1 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -3,7 +3,6 @@ import hashlib import importlib from collections.abc import Callable -from types import SimpleNamespace from typing import Any import pytest @@ -183,37 +182,6 @@ def new_mamba_spec( ) -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( - num_blocks=1, - kv_cache_tensors=[], - kv_cache_groups=[ - KVCacheGroupSpec(["attention"], new_kv_cache_spec(block_size=16)), - KVCacheGroupSpec( - ["mamba"], - new_mamba_spec(block_size=1648, mamba_cache_mode="align"), - ), - ], - ) - vllm_config = SimpleNamespace( - cache_config=SimpleNamespace( - block_size=16, - enable_prefix_caching=True, - hash_block_size=None, - ), - parallel_config=SimpleNamespace( - decode_context_parallel_size=1, - prefill_context_parallel_size=1, - ), - kv_transfer_config=None, - ) - - assert kv_cache_utils.resolve_kv_cache_block_sizes( - kv_cache_config, vllm_config - ) == (1648, 16) - - @pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor]) def test_none_hash(monkeypatch, hash_fn): import vllm.v1.core.kv_cache_utils @@ -1394,8 +1362,6 @@ def test_get_max_concurrency_for_kv_cache_config(): enable_chunked_prefill=True, max_model_len=model_config.max_model_len, is_encoder_decoder=model_config.is_encoder_decoder, - # Pin to sync: SWA per-request bounds grow with overlapping batches. - async_scheduling=False, ) vllm_config = VllmConfig( @@ -1932,60 +1898,6 @@ def test_get_kv_cache_spec_kind_prefers_specific_attention_subclasses(): ) -def test_deepseek_v4_tuple_width_minimizes_physical_pool_pages(): - full_spec = MLAAttentionSpec( - block_size=256, - num_kv_heads=1, - head_size=512, - dtype=torch.uint8, - cache_dtype_str="fp8_ds_mla", - compress_ratio=4, - alignment=576, - model_version="deepseek_v4", - ) - swa_spec = SlidingWindowMLASpec( - block_size=64, - num_kv_heads=1, - head_size=512, - dtype=torch.uint8, - sliding_window=128, - cache_dtype_str="fp8_ds_mla", - alignment=576, - model_version="deepseek_v4", - ) - grouped_specs = [ - UniformTypeKVCacheSpecs( - block_size=256, - kv_cache_specs={f"full_{i}": full_spec for i in range(21)}, - ), - UniformTypeKVCacheSpecs( - block_size=64, - kv_cache_specs={f"swa_{i}": swa_spec for i in range(46)}, - ), - ] - vllm_config = SimpleNamespace( - model_config=SimpleNamespace(max_model_len=1_048_576), - scheduler_config=SimpleNamespace(max_num_batched_tokens=8192), - parallel_config=SimpleNamespace( - decode_context_parallel_size=1, - prefill_context_parallel_size=1, - ), - ) - - assert ( - kv_cache_utils._select_deepseek_v4_tuple_width( - vllm_config, # type: ignore[arg-type] - grouped_specs, - ) - == 21 - ) - groups = kv_cache_utils._get_kv_cache_groups_uniform_groups( - vllm_config, # type: ignore[arg-type] - grouped_specs, - ) - assert [len(group.layer_names) for group in groups] == [21, 16, 15, 15] - - def test_get_kv_cache_spec_kind_unwraps_uniform_type_specs(): uniform_mla_spec = UniformTypeKVCacheSpecs( block_size=16, diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index 5b1d53cd8a..546412b1d2 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -1009,98 +1009,6 @@ def test_prefill_hybrid_model_mamba_align(): manager.free(req0) -def test_mamba_align_offload_handoff_tracks_relocated_mtp_boundary(): - """Expose the committed state block, not its stale logical position. - - A large second prefill chunk relocates the four MTP scratch blocks. The - first scratch block moves from logical position 1 to position 5 and becomes - the state after token 24. An append-only connector mirror would still see - that physical block at position 1 and could save the token-24 state under - the token-8 key. - """ - block_size = 4 - kv_cache_config = KVCacheConfig( - num_blocks=64, - kv_cache_tensors=[], - kv_cache_groups=[ - KVCacheGroupSpec( - ["full"], - FullAttentionSpec( - block_size=block_size, - num_kv_heads=1, - head_size=1, - dtype=torch.float32, - ), - ), - KVCacheGroupSpec( - ["mamba"], - MambaSpec( - block_size=block_size, - shapes=((1, 1),), - dtypes=(torch.float32,), - mamba_cache_mode="align", - num_speculative_blocks=4, - ), - ), - ], - ) - manager = KVCacheManager( - kv_cache_config, - max_model_len=8192, - enable_caching=True, - hash_block_size=block_size, - ) - request = make_request("0", [0] * 24, block_size, sha256) - computed_blocks, num_computed_tokens = manager.get_computed_blocks(request) - - assert ( - manager.allocate_slots( - request, - num_new_tokens=4, - num_new_computed_tokens=num_computed_tokens, - new_computed_blocks=computed_blocks, - ) - is not None - ) - mamba_manager = manager.coordinator.single_type_managers[1] - mamba_blocks = mamba_manager.req_to_blocks[request.request_id] - first_boundary_block_id = mamba_blocks[0].block_id - first_scratch_block = mamba_blocks[1] - - assert manager.take_boundary_state_offloads() == { - request.request_id: [(1, first_boundary_block_id, 4)] - } - - request.num_computed_tokens = 4 - assert manager.allocate_slots(request, num_new_tokens=20) is not None - mamba_blocks = mamba_manager.req_to_blocks[request.request_id] - - assert mamba_blocks[1].is_null - assert mamba_blocks[5] is first_scratch_block - assert manager.take_boundary_state_offloads() == { - request.request_id: [(1, first_scratch_block.block_id, 24)] - } - assert manager.take_boundary_state_offloads() == {} - - manager.free(request) - - # A request freed before the scheduler drains its offer must not leave a - # block ID that can be recycled and later misidentified as its boundary. - freed_request = make_request("freed", [1] * 4, block_size, sha256) - computed_blocks, num_computed_tokens = manager.get_computed_blocks(freed_request) - assert ( - manager.allocate_slots( - freed_request, - num_new_tokens=4, - num_new_computed_tokens=num_computed_tokens, - new_computed_blocks=computed_blocks, - ) - is not None - ) - manager.free(freed_request) - assert manager.take_boundary_state_offloads() == {} - - def test_prefill_plp(): """Test prefill with APC and some prompt logprobs (plp) requests. @@ -2876,8 +2784,7 @@ def test_can_fit_full_sequence_swa_cap_admits_long_prompt(): manager = KVCacheManager( config, max_model_len=max_model_len, - # Single (sync) batch in flight, so in-flight tokens == batched tokens. - max_in_flight_tokens=max_num_batched_tokens, + max_num_batched_tokens=max_num_batched_tokens, enable_caching=True, hash_block_size=block_size, ) @@ -2933,8 +2840,7 @@ def test_can_fit_full_sequence_full_attention_still_gates_oversized(): manager = KVCacheManager( config, max_model_len=max_model_len, - # Single (sync) batch in flight, so in-flight tokens == batched tokens. - max_in_flight_tokens=max_num_batched_tokens, + max_num_batched_tokens=max_num_batched_tokens, enable_caching=True, hash_block_size=block_size, ) diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index cfd91b83da..db33b4c6df 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -23,7 +23,6 @@ from vllm.sampling_params import SamplingParams, StructuredOutputsParams from vllm.utils.hashing import sha256 from vllm.v1.core.encoder_cache_manager import EncoderCacheManager -from vllm.v1.core.kv_cache_manager import KVCacheManager from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput from vllm.v1.core.sched.scheduler import Scheduler @@ -42,46 +41,6 @@ pytestmark = pytest.mark.cpu_test -@pytest.mark.parametrize( - "num_computed_tokens,expected_attention_ids", - [ - (0, []), - (1, [10]), - (16, [10]), - (17, [10, 11]), - (64, [10, 11, 12, 13]), - ], -) -def test_block_ids_for_computed_tokens_exclude_attention_lookahead( - num_computed_tokens: int, - expected_attention_ids: list[int], -): - """Only attention blocks beyond the computed-token frontier are clipped. - - Non-attention state groups retain their allocator-defined block table. - """ - manager = object.__new__(KVCacheManager) - attention_spec = FullAttentionSpec( - block_size=16, - num_kv_heads=1, - head_size=1, - dtype=torch.float32, - ) - manager.kv_cache_config = Mock( - kv_cache_groups=[ - Mock(kv_cache_spec=attention_spec), - Mock(kv_cache_spec=object()), - ] - ) - manager.get_block_ids = Mock(return_value=([10, 11, 12, 13], [20, 21, 22])) - - block_ids = manager.get_block_ids_for_computed_tokens( - "request", num_computed_tokens - ) - - assert block_ids == (expected_attention_ids, [20, 21, 22]) - - def test_add_requests(): scheduler = create_scheduler() requests = create_requests(num_requests=10) diff --git a/tests/v1/core/test_single_type_kv_cache_manager.py b/tests/v1/core/test_single_type_kv_cache_manager.py index 3af8c830cb..f59830dcd7 100644 --- a/tests/v1/core/test_single_type_kv_cache_manager.py +++ b/tests/v1/core/test_single_type_kv_cache_manager.py @@ -14,14 +14,9 @@ ) from vllm.v1.core.single_type_kv_cache_manager import ( ChunkedLocalAttentionManager, - PrefixAnchoredSWAManager, SlidingWindowManager, ) -from vllm.v1.kv_cache_interface import ( - ChunkedLocalAttentionSpec, - PrefixAnchoredSWASpec, - SlidingWindowSpec, -) +from vllm.v1.kv_cache_interface import ChunkedLocalAttentionSpec, SlidingWindowSpec pytestmark = pytest.mark.cpu_test @@ -486,103 +481,3 @@ def test_predictor_matches_allocator_blocks_calculation_with_admission_cap(): f"but allocator pulled {len(new_blocks)}" ) total_computed = num_tokens - - -def test_prefix_anchored_swa_remove_skipped_blocks_gap_range(): - block_size = 4 - spec = PrefixAnchoredSWASpec( - block_size=block_size, - num_kv_heads=1, - head_size=1, - dtype=torch.float32, - decode_sliding_window=8, - ) - block_pool = BlockPool( - num_gpu_blocks=2000, enable_caching=True, hash_block_size=block_size - ) - manager = PrefixAnchoredSWAManager( - spec, - block_pool=block_pool, - enable_caching=True, - kv_cache_group_id=0, - ) - - null_block_id = block_pool.null_block.block_id - original_block_ids = list(range(1000, 1010)) - block_table = [KVCacheBlock(id_) for id_ in original_block_ids] - manager.req_to_blocks["test"] = block_table - - prefix_len = 16 - - # Without num_prompt_tokens, gap blocks are not evicted (full-attention - # default path). - manager.remove_skipped_blocks("test", 28) - assert [b.block_id for b in block_table] == original_block_ids - - # Gap = block 4 only (tokens [16, 20) fall in the gap). - manager.remove_skipped_blocks("test", 28, num_prompt_tokens=prefix_len) - expected = original_block_ids.copy() - expected[4] = null_block_id - assert [b.block_id for b in block_table] == expected - - # Window moves: blocks 5 and 6 also enter the gap; block 4 is already - # null and stays null. - manager.remove_skipped_blocks("test", 36, num_prompt_tokens=prefix_len) - expected[5] = null_block_id - expected[6] = null_block_id - assert [b.block_id for b in block_table] == expected - - # Blocks holding the prefix and the live decode window are never touched. - assert [b.block_id for b in block_table[:4]] == original_block_ids[:4] - assert [b.block_id for b in block_table[7:]] == original_block_ids[7:] - - -def test_prefix_anchored_swa_spec_merge(): - def make_spec(window: int) -> PrefixAnchoredSWASpec: - return PrefixAnchoredSWASpec( - block_size=16, - num_kv_heads=2, - head_size=64, - dtype=torch.float16, - decode_sliding_window=window, - ) - - merged = PrefixAnchoredSWASpec.merge([make_spec(128), make_spec(128)]) - assert isinstance(merged, PrefixAnchoredSWASpec) - assert merged.decode_sliding_window == 128 - assert merged.block_size == 16 - assert merged.num_kv_heads == 2 - assert merged.head_size == 64 - assert merged.head_size_v == 64 - assert merged.dtype == torch.float16 - - with pytest.raises(AssertionError): - PrefixAnchoredSWASpec.merge([make_spec(128), make_spec(64)]) - - -def test_prefix_anchored_swa_manager_registered(): - from vllm.v1.core.single_type_kv_cache_manager import ( - get_manager_for_kv_cache_spec, - spec_manager_map, - ) - - assert spec_manager_map[PrefixAnchoredSWASpec] is PrefixAnchoredSWAManager - - spec = PrefixAnchoredSWASpec( - block_size=16, - num_kv_heads=1, - head_size=64, - dtype=torch.float16, - decode_sliding_window=128, - ) - block_pool = BlockPool(num_gpu_blocks=100, enable_caching=False, hash_block_size=16) - manager = get_manager_for_kv_cache_spec( - spec, - max_in_flight_tokens=1024, - max_model_len=4096, - block_pool=block_pool, - enable_caching=False, - kv_cache_group_id=0, - ) - assert isinstance(manager, PrefixAnchoredSWAManager) - assert manager.decode_sliding_window == 128 diff --git a/tests/v1/core/test_swa_inflight_window_free.py b/tests/v1/core/test_swa_inflight_window_free.py deleted file mode 100644 index 3c66fd533b..0000000000 --- a/tests/v1/core/test_swa_inflight_window_free.py +++ /dev/null @@ -1,351 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Out-of-window block frees vs in-flight GPU steps. - -With async scheduling / PP, `num_computed_tokens` optimistically includes -tokens of unprocessed steps whose attention windows still read the blocks -just below the optimistic boundary (and rejected spec tokens can roll it -back), so `allocate_slots` frees on the processed-token basis: -`num_computed_tokens - num_in_flight_tokens`. -""" - -from types import SimpleNamespace - -import pytest -import torch - -from vllm import envs -from vllm.config import VllmConfig -from vllm.v1.core.sched.output import SchedulerOutput -from vllm.v1.kv_cache_interface import ( - ChunkedLocalAttentionSpec, - PrefixAnchoredSWASpec, - SlidingWindowSpec, -) -from vllm.v1.outputs import ModelRunnerOutput - -from .utils import create_requests, create_scheduler, mock_kv - -NUM_PROMPT_TOKENS = 100 -BLOCK_SIZE = 16 -SLIDING_WINDOW = 16 -# Tokens 0..84 are outside the window of the next token to compute -# (100 - 16 + 1 = 85), i.e. 5 full blocks. -NUM_OUT_OF_WINDOW_BLOCKS = 85 // BLOCK_SIZE -# Chunked-local skips whole chunks left of the current one: -# (100 // 32) * 32 = 96 settled tokens -> 6 full blocks. -CHUNK_SIZE = 32 -NUM_OUT_OF_CHUNK_BLOCKS = (NUM_PROMPT_TOKENS // CHUNK_SIZE) * CHUNK_SIZE // BLOCK_SIZE - - -@pytest.mark.parametrize( - ("pp_size", "async_scheduling", "queue_depth", "expected_batches"), - [ - (1, False, 0, 1), - (1, True, 0, 2), - (1, True, 4, 4), - (3, False, 0, 3), - (3, True, 4, 3), - ], -) -def test_max_in_flight_tokens_matches_executor_concurrency( - monkeypatch, - pp_size: int, - async_scheduling: bool, - queue_depth: int, - expected_batches: int, -): - monkeypatch.setattr(envs, "VLLM_SM70_ASYNC_SCHEDULING_QUEUE_DEPTH", queue_depth) - config = object.__new__(VllmConfig) - config.parallel_config = SimpleNamespace(pipeline_parallel_size=pp_size) - config.scheduler_config = SimpleNamespace( - async_scheduling=async_scheduling, - max_num_batched_tokens=1024, - ) - - assert config.max_concurrent_batches == expected_batches - assert config.max_in_flight_tokens == expected_batches * 1024 - - -def _make_model_runner_output( - scheduler_output: SchedulerOutput, - token_id: int = 0, -) -> ModelRunnerOutput: - req_ids = list(scheduler_output.num_scheduled_tokens.keys()) - return ModelRunnerOutput( - req_ids=req_ids, - req_id_to_index={req_id: i for i, req_id in enumerate(req_ids)}, - sampled_token_ids=[[token_id] for _ in req_ids], - logprobs=None, - prompt_logprobs_dict={}, - pooler_output=[], - ) - - -def _create_swa_scheduler(async_scheduling: bool): - return create_scheduler( - block_size=BLOCK_SIZE, - async_scheduling=async_scheduling, - kv_cache_spec=SlidingWindowSpec( - block_size=BLOCK_SIZE, - num_kv_heads=1, - head_size=1, - dtype=torch.float32, - sliding_window=SLIDING_WINDOW, - ), - ) - - -def _create_chunked_scheduler(async_scheduling: bool): - return create_scheduler( - block_size=BLOCK_SIZE, - async_scheduling=async_scheduling, - kv_cache_spec=ChunkedLocalAttentionSpec( - block_size=BLOCK_SIZE, - num_kv_heads=1, - head_size=1, - dtype=torch.float32, - attention_chunk_size=CHUNK_SIZE, - ), - ) - - -def _create_anchored_scheduler(async_scheduling: bool): - return create_scheduler( - block_size=BLOCK_SIZE, - async_scheduling=async_scheduling, - enable_prefix_caching=False, - kv_cache_spec=PrefixAnchoredSWASpec( - block_size=BLOCK_SIZE, - num_kv_heads=1, - head_size=1, - dtype=torch.float32, - decode_sliding_window=SLIDING_WINDOW, - ), - ) - - -def _num_null_blocks(scheduler, request_id: str) -> int: - manager = scheduler.kv_cache_manager.coordinator.single_type_managers[0] - null_block = manager._null_block - return sum(1 for b in manager.req_to_blocks[request_id] if b is null_block) - - -def test_num_in_flight_tokens_accounting(): - scheduler = create_scheduler(async_scheduling=True) - request = create_requests(num_requests=1, num_tokens=NUM_PROMPT_TOKENS)[0] - scheduler.add_request(request) - - out0 = scheduler.schedule() - assert request.num_in_flight_tokens == NUM_PROMPT_TOKENS - # Async: decode scheduled before the prefill output is processed. - out1 = scheduler.schedule() - assert request.num_in_flight_tokens == NUM_PROMPT_TOKENS + 1 - - scheduler.update_from_output(out0, _make_model_runner_output(out0)) - assert request.num_in_flight_tokens == 1 - scheduler.update_from_output(out1, _make_model_runner_output(out1)) - assert request.num_in_flight_tokens == 0 - - -def test_swa_free_waits_for_in_flight_step(): - """Async: out-of-window blocks stay allocated until the step that still - reads them has been processed.""" - scheduler = _create_swa_scheduler(async_scheduling=True) - request = create_requests( - num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE - )[0] - scheduler.add_request(request) - req_id = request.request_id - block_pool = scheduler.kv_cache_manager.block_pool - - out0 = scheduler.schedule() # prefill, in flight from here on - free_after_prefill = block_pool.get_num_free_blocks() - - # Decode scheduled while the prefill still reads the out-of-window blocks: - # they must not be freed yet. - out1 = scheduler.schedule() - assert _num_null_blocks(scheduler, req_id) == 0 - assert block_pool.get_num_free_blocks() == free_after_prefill - - # Prefill output processed; the next allocate frees the out-of-window - # blocks. - scheduler.update_from_output(out0, _make_model_runner_output(out0)) - scheduler.schedule() - assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_WINDOW_BLOCKS - assert ( - block_pool.get_num_free_blocks() - == free_after_prefill + NUM_OUT_OF_WINDOW_BLOCKS - ) - # Not double-freed on the following steps. - scheduler.update_from_output(out1, _make_model_runner_output(out1)) - scheduler.schedule() - assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_WINDOW_BLOCKS - - -def test_swa_free_immediate_when_sync(): - """Sync: no in-flight step at schedule time, frees happen at the first - decode allocation as before.""" - scheduler = _create_swa_scheduler(async_scheduling=False) - request = create_requests( - num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE - )[0] - scheduler.add_request(request) - req_id = request.request_id - - out0 = scheduler.schedule() - scheduler.update_from_output(out0, _make_model_runner_output(out0)) - assert request.num_in_flight_tokens == 0 - - scheduler.schedule() - assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_WINDOW_BLOCKS - - -def test_anchored_swa_gap_free_waits_for_in_flight_step(): - """Prefix-anchored SWA (windowed OCR decode): gap blocks between the - prefill tail and the decode window are freed on the processed-token - basis under async scheduling — an in-flight step whose attention still - reads the gap must hold the blocks.""" - scheduler = _create_anchored_scheduler(async_scheduling=True) - request = create_requests( - num_requests=1, - num_tokens=NUM_PROMPT_TOKENS, - block_size=BLOCK_SIZE, - max_tokens=200, - ignore_eos=True, - )[0] - scheduler.add_request(request) - req_id = request.request_id - - # Prefill, then decode until a full gap block exists on the PROCESSED - # basis: gap = [ceil(P/bs), (T - W) // bs) needs T >= 144 for block 7 - # (P=100 -> first_gap_block 7; (144-16)//16 = 8). - out_prev = scheduler.schedule() # prefill (in flight) - scheduler.update_from_output(out_prev, _make_model_runner_output(out_prev)) - for _ in range(44): # T: 100 -> 144, settled step by step - out = scheduler.schedule() - scheduler.update_from_output(out, _make_model_runner_output(out)) - assert request.num_in_flight_tokens == 0 - assert request.num_computed_tokens >= 144 - - # Async over-scheduling: the settled-basis allocate on this schedule may - # evict gap blocks for the SETTLED T. - out_a = scheduler.schedule() - null_settled = _num_null_blocks(scheduler, req_id) - assert null_settled >= 1 # gap eviction engaged at all - - # Second in-flight step: the optimistic boundary advances by 1 token but - # the processed basis does not — no additional frees may happen. - out_b = scheduler.schedule() - assert _num_null_blocks(scheduler, req_id) == null_settled - - scheduler.update_from_output(out_a, _make_model_runner_output(out_a)) - scheduler.update_from_output(out_b, _make_model_runner_output(out_b)) - - -def test_swa_admission_cap_accounts_for_overlapping_batches(): - spec = SlidingWindowSpec( - block_size=16, - num_kv_heads=1, - head_size=1, - dtype=torch.float32, - sliding_window=1024, - ) - base = spec.max_admission_blocks_per_request( - max_in_flight_tokens=1024, max_model_len=16384 - ) - # (1024 - 1 + 1024) tokens -> 128 blocks, +1 for window misalignment. - assert base == 129 - overlapped = spec.max_admission_blocks_per_request( - max_in_flight_tokens=2 * 1024, max_model_len=16384 - ) - # One extra in-flight chunk is held back: (1024 - 1 + 2 * 1024) tokens. - assert overlapped == 193 - - -def test_chunked_local_free_waits_for_in_flight_step(): - """Chunked-local attention frees whole chunks left of the current one, and - is exposed to the same load-WAR: with async scheduling those chunks must - stay allocated until the in-flight step that still reads them settles.""" - scheduler = _create_chunked_scheduler(async_scheduling=True) - request = create_requests( - num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE - )[0] - scheduler.add_request(request) - req_id = request.request_id - block_pool = scheduler.kv_cache_manager.block_pool - - out0 = scheduler.schedule() # prefill, in flight from here on - free_after_prefill = block_pool.get_num_free_blocks() - - # Decode scheduled while the prefill still reads the out-of-chunk blocks. - out1 = scheduler.schedule() - assert _num_null_blocks(scheduler, req_id) == 0 - assert block_pool.get_num_free_blocks() == free_after_prefill - - # Prefill output processed; the next allocate frees the out-of-chunk blocks. - scheduler.update_from_output(out0, _make_model_runner_output(out0)) - scheduler.schedule() - assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_CHUNK_BLOCKS - assert ( - block_pool.get_num_free_blocks() == free_after_prefill + NUM_OUT_OF_CHUNK_BLOCKS - ) - # Not double-freed on the following steps. - scheduler.update_from_output(out1, _make_model_runner_output(out1)) - scheduler.schedule() - assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_CHUNK_BLOCKS - - -def test_chunked_local_admission_cap_accounts_for_overlapping_batches(): - spec = ChunkedLocalAttentionSpec( - block_size=16, - num_kv_heads=1, - head_size=1, - dtype=torch.float32, - attention_chunk_size=1024, - ) - base = spec.max_admission_blocks_per_request( - max_in_flight_tokens=1024, max_model_len=16384 - ) - # (1024 + 1024) tokens -> 128 blocks. - assert base == 128 - overlapped = spec.max_admission_blocks_per_request( - max_in_flight_tokens=2 * 1024, max_model_len=16384 - ) - # One extra in-flight chunk is held back: (1024 + 2 * 1024) tokens. - assert overlapped == 192 - - -def test_connector_finish_frees_on_settled_basis(): - """The out-of-window prune done at request finish, before the block table - is handed to a KV connector, must use the same processed-token basis.""" - scheduler = create_scheduler( - block_size=BLOCK_SIZE, - async_scheduling=True, - use_kv_connector=mock_kv(matched_tokens=0, is_async=False), - kv_cache_spec=SlidingWindowSpec( - block_size=BLOCK_SIZE, - num_kv_heads=1, - head_size=1, - dtype=torch.float32, - sliding_window=SLIDING_WINDOW, - ), - ) - request = create_requests( - num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE - )[0] - scheduler.add_request(request) - req_id = request.request_id - - out0 = scheduler.schedule() # prefill, in flight from here on - scheduler.schedule() # decode over-scheduled: num_computed_tokens optimistic - - # Finishing now (connector store) must NOT prune out-of-window blocks the - # in-flight prefill still reads. - scheduler._connector_finished(request) - assert _num_null_blocks(scheduler, req_id) == 0 - - # Once the in-flight step settles, the same prune releases them. - scheduler.update_from_output(out0, _make_model_runner_output(out0)) - scheduler._connector_finished(request) - assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_WINDOW_BLOCKS diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py index 472c1876f4..2d9834d2e3 100644 --- a/tests/v1/core/utils.py +++ b/tests/v1/core/utils.py @@ -28,7 +28,6 @@ FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, - KVCacheSpec, ) from vllm.v1.request import Request from vllm.v1.structured_output import StructuredOutputManager @@ -58,7 +57,6 @@ def create_scheduler( pipeline_parallel_size: int = 1, use_ec_connector: bool = False, ec_role: str | None = None, - kv_cache_spec: KVCacheSpec | None = None, ) -> Scheduler | AsyncScheduler: """Create scheduler under test. @@ -146,17 +144,20 @@ def create_scheduler( speculative_config=speculative_config, ec_transfer_config=ec_transfer_config, ) - if kv_cache_spec is None: - kv_cache_spec = FullAttentionSpec( - block_size=block_size, - num_kv_heads=1, - head_size=1, - dtype=torch.float32, - ) kv_cache_config = KVCacheConfig( num_blocks=num_blocks, # A large number of blocks to hold all requests kv_cache_tensors=[], - kv_cache_groups=[KVCacheGroupSpec(["layer"], kv_cache_spec)], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer"], + FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ) + ], ) cache_config.num_gpu_blocks = num_blocks scheduler_cls = AsyncScheduler if async_scheduling else Scheduler diff --git a/tests/v1/cudagraph/test_cudagraph_dispatch.py b/tests/v1/cudagraph/test_cudagraph_dispatch.py index e78c865a8c..0c959f6b71 100644 --- a/tests/v1/cudagraph/test_cudagraph_dispatch.py +++ b/tests/v1/cudagraph/test_cudagraph_dispatch.py @@ -19,12 +19,7 @@ VllmConfig, ) from vllm.config.lora import LoRAConfig -from vllm.forward_context import ( - CUDAGRAPH_VARIANT_DEFAULT, - CUDAGRAPH_VARIANT_LONG_CONTEXT, - BatchDescriptor, - set_forward_context, -) +from vllm.forward_context import BatchDescriptor, set_forward_context from vllm.platforms import current_platform from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher @@ -378,9 +373,7 @@ def test_dsv4_decode_context_bucket_is_derived_on_sm70(self, monkeypatch): cudagraph_capture_sizes=[1, 2], ) config = _create_vllm_config(comp_config, max_num_seqs=2) - # Route by the compressed-index operator contract, not architecture or - # checkpoint identity. - config.model_config.architectures = ["UnrelatedForCausalLM"] + config.model_config.architectures = ["DeepseekV4ForCausalLM"] config.model_config.max_model_len = 4096 config.model_config.hf_config.index_topk = 512 config.model_config.hf_config.compress_ratios = [128, 4, 128] @@ -408,119 +401,6 @@ def test_dsv4_decode_context_bucket_is_derived_on_sm70(self, monkeypatch): disabled = CudagraphDispatcher(config) assert not disabled.sm70_dsv4_decode_context_buckets - def test_dsv4_context_bucket_requires_explicit_override_for_mtp_on_sm70( - self, monkeypatch - ): - monkeypatch.delenv("VLLM_SM70_DSV4_DECODE_CONTEXT_BUCKETS", raising=False) - monkeypatch.delenv("VLLM_SM70_MTP_CONTEXT_BUCKETS", raising=False) - comp_config = CompilationConfig( - cudagraph_mode="FULL_DECODE_ONLY", - mode=CompilationMode.NONE, - cudagraph_capture_sizes=[8, 16], - ) - config = _create_vllm_config(comp_config, max_num_seqs=2) - config.speculative_config = MagicMock(num_speculative_tokens=7) - config.model_config.architectures = ["DeepseekV4ForCausalLM"] - config.model_config.max_model_len = 4096 - config.model_config.hf_config.index_topk = 512 - config.model_config.hf_config.compress_ratios = [128, 4, 128] - - with ( - patch.object(current_platform, "is_cuda", return_value=True), - patch.object(current_platform, "is_device_capability", return_value=True), - ): - default_dispatcher = CudagraphDispatcher(config) - assert not default_dispatcher.has_attention_context_buckets - - monkeypatch.setenv("VLLM_SM70_DSV4_DECODE_CONTEXT_BUCKETS", "2048") - dispatcher = CudagraphDispatcher(config) - dispatcher.initialize_cudagraph_keys( - cudagraph_mode=comp_config.cudagraph_mode, - uniform_decode_query_len=8, - ) - - bounded = BatchDescriptor( - num_tokens=8, - num_reqs=1, - uniform=True, - attention_context_bucket=2048, - ) - assert dispatcher.has_attention_context_buckets - assert bounded in dispatcher.cudagraph_keys[CUDAGraphMode.FULL] - - mode, desc = dispatcher.dispatch( - num_tokens=8, - uniform_decode=True, - attention_context_len=1024, - ) - assert mode == CUDAGraphMode.FULL - assert desc == bounded - - monkeypatch.setenv("VLLM_SM70_MTP_CONTEXT_BUCKETS", "") - explicitly_disabled = CudagraphDispatcher(config) - explicitly_disabled.initialize_cudagraph_keys( - cudagraph_mode=comp_config.cudagraph_mode, - uniform_decode_query_len=8, - ) - assert not explicitly_disabled.has_attention_context_buckets - - def test_dsv4_long_context_buckets_bound_mtp_graph_width(self, monkeypatch): - monkeypatch.setenv( - "VLLM_SM70_DSV4_DECODE_CONTEXT_BUCKETS", - "2048,4096,16384,65536,131072", - ) - monkeypatch.delenv("VLLM_SM70_MTP_CONTEXT_BUCKETS", raising=False) - comp_config = CompilationConfig( - cudagraph_mode="FULL_DECODE_ONLY", - mode=CompilationMode.NONE, - cudagraph_capture_sizes=[8, 16], - ) - config = _create_vllm_config(comp_config, max_num_seqs=2) - config.speculative_config = MagicMock(num_speculative_tokens=7) - config.model_config.architectures = ["DeepseekV4ForCausalLM"] - config.model_config.max_model_len = 262144 - config.model_config.hf_config.index_topk = 512 - config.model_config.hf_config.compress_ratios = [128, 4, 128] - - with ( - patch.object(current_platform, "is_cuda", return_value=True), - patch.object(current_platform, "is_device_capability", return_value=True), - ): - dispatcher = CudagraphDispatcher(config) - dispatcher.initialize_cudagraph_keys( - cudagraph_mode=comp_config.cudagraph_mode, - uniform_decode_query_len=8, - ) - - assert dispatcher.sm70_dsv4_decode_context_buckets == ( - 2048, - 4096, - 16384, - 65536, - 131072, - ) - expected_buckets = { - 3000: 4096, - 5000: 16384, - 20000: 65536, - 100000: 131072, - } - for context_len, expected_bucket in expected_buckets.items(): - mode, descriptor = dispatcher.dispatch( - num_tokens=8, - uniform_decode=True, - attention_context_len=context_len, - ) - assert mode == CUDAGraphMode.FULL - assert descriptor.attention_context_bucket == expected_bucket - - _, unbounded = dispatcher.dispatch( - num_tokens=8, - uniform_decode=True, - attention_context_len=200000, - ) - assert unbounded.attention_context_bucket is None - def test_fp8_e5m2_decode_context_bucket_is_shape_derived_on_sm70(self, monkeypatch): monkeypatch.delenv("VLLM_SM70_FP8_KV_DECODE_CONTEXT_BUCKETS", raising=False) comp_config = CompilationConfig( @@ -542,8 +422,6 @@ def test_fp8_e5m2_decode_context_bucket_is_shape_derived_on_sm70(self, monkeypat patch("vllm.v1.cudagraph_dispatcher.envs") as mock_envs, ): mock_envs.VLLM_SM70_FLASH_ATTN_V100 = True - mock_envs.VLLM_FLASH_V100_XQA_BATCH_CONTEXT_ROUTING = False - mock_envs.VLLM_FLASH_V100_DECODE_PARTITION_SIZE = None dispatcher = CudagraphDispatcher(config) dispatcher.initialize_cudagraph_keys( cudagraph_mode=comp_config.cudagraph_mode, @@ -575,312 +453,6 @@ def test_fp8_e5m2_decode_context_bucket_is_shape_derived_on_sm70(self, monkeypat assert mode == CUDAGraphMode.FULL assert desc == BatchDescriptor(num_tokens=1, num_reqs=1, uniform=True) - @pytest.mark.parametrize( - ("num_tokens", "context_len", "expected_variant"), - ( - (4, 8191, CUDAGRAPH_VARIANT_DEFAULT), - (4, 12288, CUDAGRAPH_VARIANT_DEFAULT), - (4, 16383, CUDAGRAPH_VARIANT_DEFAULT), - (4, 16384, CUDAGRAPH_VARIANT_LONG_CONTEXT), - (8, 4095, CUDAGRAPH_VARIANT_DEFAULT), - (8, 12288, CUDAGRAPH_VARIANT_DEFAULT), - (8, 16383, CUDAGRAPH_VARIANT_DEFAULT), - (8, 16384, CUDAGRAPH_VARIANT_LONG_CONTEXT), - (16, 1024, CUDAGRAPH_VARIANT_DEFAULT), - (16, 16001, CUDAGRAPH_VARIANT_DEFAULT), - (16, 16383, CUDAGRAPH_VARIANT_DEFAULT), - (16, 16384, CUDAGRAPH_VARIANT_LONG_CONTEXT), - ), - ) - def test_fp8_e5m2_batch_context_graph_routing_on_sm70( - self, - monkeypatch, - num_tokens, - context_len, - expected_variant, - ): - monkeypatch.delenv("VLLM_SM70_FP8_KV_DECODE_CONTEXT_BUCKETS", raising=False) - comp_config = CompilationConfig( - cudagraph_mode="FULL_DECODE_ONLY", - mode=CompilationMode.NONE, - cudagraph_capture_sizes=[1, 2, 4, 8, 16], - ) - config = _create_vllm_config(comp_config, max_num_seqs=16) - config.cache_config.cache_dtype = "fp8_e5m2" - config.attention_config.backend = None - config.model_config.max_model_len = 262144 - config.model_config.hf_text_config.num_attention_heads = 24 - config.model_config.hf_text_config.num_key_value_heads = 4 - config.model_config.hf_text_config.head_dim = 256 - - with ( - patch.object(current_platform, "is_cuda", return_value=True), - patch.object(current_platform, "is_device_capability", return_value=True), - patch("vllm.v1.cudagraph_dispatcher.envs") as mock_envs, - ): - mock_envs.VLLM_SM70_FLASH_ATTN_V100 = True - mock_envs.VLLM_FLASH_V100_XQA_BATCH_CONTEXT_ROUTING = True - mock_envs.VLLM_FLASH_V100_DECODE_PARTITION_SIZE = None - dispatcher = CudagraphDispatcher(config) - dispatcher.initialize_cudagraph_keys( - cudagraph_mode=comp_config.cudagraph_mode, - uniform_decode_query_len=1, - ) - - mode, desc = dispatcher.dispatch( - num_tokens=num_tokens, - uniform_decode=True, - attention_context_len=context_len, - ) - - assert dispatcher.sm70_fp8_kv_batch_context_routing - assert dispatcher.has_attention_context_specialization - assert mode == CUDAGraphMode.FULL - assert desc == BatchDescriptor( - num_tokens=num_tokens, - num_reqs=num_tokens, - uniform=True, - graph_variant=expected_variant, - ) - assert ( - BatchDescriptor( - num_tokens=num_tokens, - num_reqs=num_tokens, - uniform=True, - graph_variant=CUDAGRAPH_VARIANT_LONG_CONTEXT, - ) - in dispatcher.cudagraph_keys[CUDAGraphMode.FULL] - ) - - def test_fp8_e5m2_long_context_variants_capture_after_baseline(self, monkeypatch): - comp_config = CompilationConfig( - cudagraph_mode="FULL_DECODE_ONLY", - mode=CompilationMode.NONE, - cudagraph_capture_sizes=[1, 2, 4, 8, 16], - ) - config = _create_vllm_config(comp_config, max_num_seqs=16) - config.cache_config.cache_dtype = "fp8_e5m2" - config.attention_config.backend = None - config.model_config.max_model_len = 262144 - config.model_config.hf_text_config.num_attention_heads = 24 - config.model_config.hf_text_config.num_key_value_heads = 4 - config.model_config.hf_text_config.head_dim = 256 - - with ( - patch.object(current_platform, "is_cuda", return_value=True), - patch.object(current_platform, "is_device_capability", return_value=True), - patch("vllm.v1.cudagraph_dispatcher.envs") as mock_envs, - ): - mock_envs.VLLM_SM70_FLASH_ATTN_V100 = True - mock_envs.VLLM_FLASH_V100_XQA_BATCH_CONTEXT_ROUTING = True - mock_envs.VLLM_FLASH_V100_DECODE_PARTITION_SIZE = None - dispatcher = CudagraphDispatcher(config) - dispatcher.initialize_cudagraph_keys( - cudagraph_mode=comp_config.cudagraph_mode, - uniform_decode_query_len=1, - ) - - full = dict(dispatcher.get_capture_descs())[CUDAGraphMode.FULL] - variants = [desc.graph_variant for desc in full] - first_long = variants.index(CUDAGRAPH_VARIANT_LONG_CONTEXT) - assert all( - variant == CUDAGRAPH_VARIANT_DEFAULT for variant in variants[:first_long] - ) - assert variants[first_long:] == [CUDAGRAPH_VARIANT_LONG_CONTEXT] * 3 - - @pytest.mark.parametrize( - ("e4m3_enabled", "expected_routing"), - ((True, True), (False, False)), - ) - def test_fp8_e4m3_batch_context_graph_routing_on_sm70( - self, - monkeypatch, - e4m3_enabled, - expected_routing, - ): - comp_config = CompilationConfig( - cudagraph_mode="FULL_DECODE_ONLY", - mode=CompilationMode.NONE, - cudagraph_capture_sizes=[1, 2, 4, 8, 16], - ) - config = _create_vllm_config(comp_config, max_num_seqs=16) - config.cache_config.cache_dtype = "fp8_e4m3" - config.attention_config.backend = None - config.model_config.max_model_len = 32768 - config.model_config.hf_text_config.num_attention_heads = 24 - config.model_config.hf_text_config.num_key_value_heads = 4 - config.model_config.hf_text_config.head_dim = 256 - - with ( - patch.object(current_platform, "is_cuda", return_value=True), - patch.object(current_platform, "is_device_capability", return_value=True), - patch("vllm.v1.cudagraph_dispatcher.envs") as mock_envs, - ): - mock_envs.VLLM_SM70_FLASH_ATTN_V100 = True - mock_envs.VLLM_FLASH_V100_XQA_BATCH_CONTEXT_ROUTING = True - mock_envs.VLLM_FLASH_V100_DECODE_PARTITION_SIZE = None - mock_envs.VLLM_FLASH_V100_E4M3_BATCH_XQA = e4m3_enabled - dispatcher = CudagraphDispatcher(config) - dispatcher.initialize_cudagraph_keys( - cudagraph_mode=comp_config.cudagraph_mode, - uniform_decode_query_len=1, - ) - - mode, desc = dispatcher.dispatch( - num_tokens=8, - uniform_decode=True, - attention_context_len=16384, - ) - - assert dispatcher.sm70_fp8_kv_batch_context_routing is expected_routing - assert mode == CUDAGraphMode.FULL - assert desc == BatchDescriptor( - num_tokens=8, - num_reqs=8, - uniform=True, - graph_variant=( - CUDAGRAPH_VARIANT_LONG_CONTEXT - if expected_routing - else CUDAGRAPH_VARIANT_DEFAULT - ), - ) - - @pytest.mark.parametrize( - ("auto_enabled", "wave_enabled", "expected_variant"), - ( - (True, True, CUDAGRAPH_VARIANT_LONG_CONTEXT), - (False, True, CUDAGRAPH_VARIANT_DEFAULT), - (True, False, CUDAGRAPH_VARIANT_DEFAULT), - ), - ) - def test_fp8_e4m3_b1_wave_graph_routing_on_sm70( - self, - monkeypatch, - auto_enabled, - wave_enabled, - expected_variant, - ): - monkeypatch.delenv("VLLM_SM70_FP8_KV_DECODE_CONTEXT_BUCKETS", raising=False) - monkeypatch.setenv( - "VLLM_FLASH_V100_XQA_E4M3_G6_P64_P256_AUTO", - "1" if auto_enabled else "0", - ) - monkeypatch.setenv( - "VLLM_FLASH_V100_XQA_E4M3_G6_WAVE_PARTITIONS", - "1" if wave_enabled else "0", - ) - monkeypatch.setenv("VLLM_FLASH_V100_XQA_E4M3_G6_P512_BEGIN", "49152") - comp_config = CompilationConfig( - cudagraph_mode="FULL_DECODE_ONLY", - mode=CompilationMode.NONE, - cudagraph_capture_sizes=[1, 2, 4, 8, 16], - ) - config = _create_vllm_config(comp_config, max_num_seqs=16) - config.cache_config.cache_dtype = "fp8_e4m3" - config.attention_config.backend = None - config.model_config.max_model_len = 262144 - config.model_config.hf_text_config.num_attention_heads = 24 - config.model_config.hf_text_config.num_key_value_heads = 4 - config.model_config.hf_text_config.head_dim = 256 - - with ( - patch.object(current_platform, "is_cuda", return_value=True), - patch.object(current_platform, "is_device_capability", return_value=True), - patch("vllm.v1.cudagraph_dispatcher.envs") as mock_envs, - ): - mock_envs.VLLM_SM70_FLASH_ATTN_V100 = True - mock_envs.VLLM_FLASH_V100_XQA_BATCH_CONTEXT_ROUTING = True - mock_envs.VLLM_FLASH_V100_DECODE_PARTITION_SIZE = None - mock_envs.VLLM_FLASH_V100_E4M3_BATCH_XQA = True - dispatcher = CudagraphDispatcher(config) - dispatcher.initialize_cudagraph_keys( - cudagraph_mode=comp_config.cudagraph_mode, - uniform_decode_query_len=1, - ) - - mode, below = dispatcher.dispatch( - num_tokens=1, - uniform_decode=True, - attention_context_len=49151, - ) - assert mode == CUDAGraphMode.FULL - assert below.graph_variant == CUDAGRAPH_VARIANT_DEFAULT - - mode, at_boundary = dispatcher.dispatch( - num_tokens=1, - uniform_decode=True, - attention_context_len=49152, - ) - assert mode == CUDAGraphMode.FULL - assert at_boundary.graph_variant == expected_variant - assert dispatcher.sm70_e4m3_b1_wave_context_routing is ( - expected_variant == CUDAGRAPH_VARIANT_LONG_CONTEXT - ) - - def test_fp8_e4m3_batch_context_graph_routing_rejects_flashinfer(self, monkeypatch): - comp_config = CompilationConfig( - cudagraph_mode="FULL_DECODE_ONLY", - mode=CompilationMode.NONE, - cudagraph_capture_sizes=[1, 2, 4, 8, 16], - ) - config = _create_vllm_config(comp_config, max_num_seqs=16) - config.cache_config.cache_dtype = "fp8_e4m3" - config.attention_config.backend = "FLASHINFER_SM70" - config.model_config.max_model_len = 32768 - config.model_config.hf_text_config.num_attention_heads = 24 - config.model_config.hf_text_config.num_key_value_heads = 4 - config.model_config.hf_text_config.head_dim = 256 - - with ( - patch.object(current_platform, "is_cuda", return_value=True), - patch.object(current_platform, "is_device_capability", return_value=True), - patch("vllm.v1.cudagraph_dispatcher.envs") as mock_envs, - ): - mock_envs.VLLM_SM70_FLASH_ATTN_V100 = True - mock_envs.VLLM_FLASH_V100_XQA_BATCH_CONTEXT_ROUTING = True - mock_envs.VLLM_FLASH_V100_DECODE_PARTITION_SIZE = None - mock_envs.VLLM_FLASH_V100_E4M3_BATCH_XQA = True - dispatcher = CudagraphDispatcher(config) - - assert not dispatcher.sm70_fp8_kv_batch_context_routing - - def test_fp8_e5m2_batch_context_graph_routing_rollback(self, monkeypatch): - comp_config = CompilationConfig( - cudagraph_mode="FULL_DECODE_ONLY", - mode=CompilationMode.NONE, - cudagraph_capture_sizes=[1, 2, 4, 8, 16], - ) - config = _create_vllm_config(comp_config, max_num_seqs=16) - config.cache_config.cache_dtype = "fp8_e5m2" - config.model_config.max_model_len = 262144 - config.model_config.hf_text_config.num_attention_heads = 24 - config.model_config.hf_text_config.num_key_value_heads = 4 - config.model_config.hf_text_config.head_dim = 256 - - with ( - patch.object(current_platform, "is_cuda", return_value=True), - patch.object(current_platform, "is_device_capability", return_value=True), - patch("vllm.v1.cudagraph_dispatcher.envs") as mock_envs, - ): - mock_envs.VLLM_SM70_FLASH_ATTN_V100 = True - mock_envs.VLLM_FLASH_V100_XQA_BATCH_CONTEXT_ROUTING = False - mock_envs.VLLM_FLASH_V100_DECODE_PARTITION_SIZE = None - dispatcher = CudagraphDispatcher(config) - dispatcher.initialize_cudagraph_keys( - cudagraph_mode=comp_config.cudagraph_mode, - uniform_decode_query_len=1, - ) - - mode, desc = dispatcher.dispatch( - num_tokens=8, - uniform_decode=True, - attention_context_len=4096, - ) - - assert not dispatcher.sm70_fp8_kv_batch_context_routing - assert mode == CUDAGraphMode.FULL - assert desc == BatchDescriptor(num_tokens=8, num_reqs=8, uniform=True) - @pytest.mark.parametrize("cache_dtype", ("auto", "fp8_e4m3")) def test_fp8_e5m2_decode_context_bucket_does_not_match_other_cache_dtypes( self, monkeypatch, cache_dtype @@ -918,35 +490,6 @@ def test_fp8_e5m2_decode_context_bucket_explicit_empty_disables(self, monkeypatc assert not dispatcher.sm70_fp8_kv_decode_context_buckets -def test_negative_graph_variant_uses_an_isolated_pool(): - wrapper = object.__new__(CUDAGraphWrapper) - wrapper.graph_pool = object() - wrapper.concrete_cudagraph_entries = {} - wrapper.semantic_variant_graph_pools = {} - default_desc = BatchDescriptor( - num_tokens=4, - graph_variant=CUDAGRAPH_VARIANT_DEFAULT, - ) - long_context_desc = replace( - default_desc, - graph_variant=CUDAGRAPH_VARIANT_LONG_CONTEXT, - ) - isolated_pool = object() - - with patch.object( - current_platform, - "graph_pool_handle", - return_value=isolated_pool, - ) as graph_pool_handle: - assert wrapper._graph_pool_for_descriptor(default_desc) is wrapper.graph_pool - assert wrapper._graph_pool_for_descriptor(long_context_desc) is isolated_pool - assert wrapper._graph_pool_for_descriptor(long_context_desc) is isolated_pool - - graph_pool_handle.assert_called_once_with() - wrapper.clear_graphs() - assert not wrapper.semantic_variant_graph_pools - - @pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda") class TestCUDAGraphWrapper: def setup_method(self): diff --git a/tests/v1/cudagraph/test_sm70_mtp_split_cudagraphs.py b/tests/v1/cudagraph/test_sm70_mtp_split_cudagraphs.py deleted file mode 100644 index 23c58085e8..0000000000 --- a/tests/v1/cudagraph/test_sm70_mtp_split_cudagraphs.py +++ /dev/null @@ -1,96 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from types import SimpleNamespace -from unittest.mock import MagicMock - -import pytest -import torch - -from vllm.config import CompilationConfig, CUDAGraphMode -from vllm.v1.worker.gpu.cudagraph_utils import CudaGraphManager - - -def _make_config(max_num_seqs: int, verifier_sizes: list[int]): - speculative_config = SimpleNamespace( - method="mtp", - num_speculative_state_tokens=lambda: 4, - ) - return SimpleNamespace( - scheduler_config=SimpleNamespace(max_num_seqs=max_num_seqs), - parallel_config=SimpleNamespace( - tensor_parallel_size=4, - data_parallel_size=1, - ), - compilation_config=CompilationConfig( - cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE, - cudagraph_capture_sizes=verifier_sizes.copy(), - max_cudagraph_capture_size=max(verifier_sizes), - ), - speculative_config=speculative_config, - ) - - -@pytest.mark.parametrize( - ("max_num_seqs", "request_sizes", "verifier_sizes"), - [ - (16, [1, 2, 4, 6, 8, 12, 16], [5, 10, 20, 30, 40, 60, 80]), - (32, [1, 2, 4, 6, 8, 12, 16, 32], [5, 10, 20, 30, 40, 60, 80, 160]), - ], -) -def test_split_managers_keep_exact_full_graph_shapes( - monkeypatch, - max_num_seqs: int, - request_sizes: list[int], - verifier_sizes: list[int], -): - monkeypatch.setenv("VLLM_SM70_MTP_SPLIT_DRAFT_CUDAGRAPHS", "1") - monkeypatch.setattr( - "vllm.v1.worker.gpu.cudagraph_utils.current_platform.is_cuda", - lambda: True, - ) - monkeypatch.setattr( - "vllm.v1.worker.gpu.cudagraph_utils.current_platform.is_device_capability", - lambda capability: capability == (7, 0), - ) - 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 = _make_config(max_num_seqs, verifier_sizes) - shared_sizes = config.compilation_config.cudagraph_capture_sizes.copy() - shared_max_size = config.compilation_config.max_cudagraph_capture_size - target = CudaGraphManager( - config, - torch.device("cpu"), - CUDAGraphMode.FULL_DECODE_ONLY, - decode_query_len=5, - ) - draft = CudaGraphManager( - config, - torch.device("cpu"), - CUDAGraphMode.FULL_DECODE_ONLY, - decode_query_len=1, - ) - - assert target._capture_sizes == verifier_sizes - assert draft._capture_sizes == request_sizes - assert config.compilation_config.cudagraph_capture_sizes == shared_sizes - assert config.compilation_config.max_cudagraph_capture_size == shared_max_size - - target._graphs_captured = True - draft._graphs_captured = True - for batch_size in request_sizes: - target_desc = target.dispatch(batch_size, batch_size * 5, 5) - draft_desc = draft.dispatch(batch_size, batch_size, 1) - assert target_desc.cg_mode == CUDAGraphMode.FULL - assert target_desc.num_tokens == batch_size * 5 - assert target_desc.num_reqs == batch_size - assert draft_desc.cg_mode == CUDAGraphMode.FULL - assert draft_desc.num_tokens == batch_size - assert draft_desc.num_reqs == batch_size diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index f377f244b2..5cc7065c1a 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Iterable -from types import SimpleNamespace from unittest.mock import MagicMock import pytest @@ -14,21 +13,15 @@ from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID from vllm.distributed.kv_events import BlockRemoved, BlockStored from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import ( - GroupOffloadConfig, OffloadingConnectorScheduler, - RequestOffloadState, - SchedulerOffloadConfig, ) from vllm.v1.core.kv_cache_utils import BlockHash from vllm.v1.kv_cache_interface import ( FullAttentionSpec, - KVCacheConfig, KVCacheGroupSpec, - MambaSpec, SlidingWindowSpec, ) from vllm.v1.kv_offload.base import ( - GPULoadStoreSpec, OffloadingEvent, OffloadingManager, OffloadPolicy, @@ -813,192 +806,6 @@ def test_request_level_policy_stores_all_blocks(request_runner, async_scheduling # --------------------------------------------------------------------------- -def _make_exact_boundary_scheduler() -> OffloadingConnectorScheduler: - scheduler = object.__new__(OffloadingConnectorScheduler) - scheduler.config = SchedulerOffloadConfig( - kv_group_configs=( - GroupOffloadConfig( - group_idx=0, - gpu_block_size=16, - offloaded_block_size=16, - hash_block_size_factor=4, - sliding_window_size_in_blocks=None, - ), - GroupOffloadConfig( - group_idx=1, - gpu_block_size=16, - offloaded_block_size=16, - hash_block_size_factor=4, - sliding_window_size_in_blocks=1, - requires_exact_boundary_source=True, - ), - ), - block_size_factor=1, - num_workers=1, - offload_prompt_only=False, - ) - scheduler.manager = MagicMock(spec=OffloadingManager) - scheduler.manager.prepare_store.side_effect = lambda keys, req_context: ( - generate_store_output(keys) - ) - request = MagicMock() - request.request_id = "req" - request.kv_transfer_params = None - request.block_hashes = [BlockHash(f"h{i}".encode()) for i in range(4)] - request.num_computed_tokens = 0 - request.num_prompt_tokens = 16 - request.num_tokens = 16 - request.is_finished.return_value = False - req_context = ReqContext(req_id=request.request_id) - req_status = RequestOffloadState( - config=scheduler.config, - req=request, - req_context=req_context, - offloading_context=RequestOffloadingContext(), - ) - req_status.update_offload_keys() - req_status.group_states[0].block_ids[:] = [11] - # This is the connector's stale append-only view of the Mamba table. - req_status.group_states[1].block_ids[:] = [21] - - scheduler._req_status = {request.request_id: req_status} - scheduler._current_batch_load_jobs = {} - scheduler._current_batch_jobs_to_flush = set() - scheduler._blocks_being_loaded = set() - scheduler._job_counter = 0 - scheduler._stale_job_threshold = 0 - scheduler._jobs = {} - scheduler._block_id_to_pending_jobs = {} - return scheduler - - -def test_mamba_align_group_requires_exact_boundary_source(): - block_size = 16 - kv_cache_config = KVCacheConfig( - num_blocks=32, - kv_cache_tensors=[], - kv_cache_groups=[ - KVCacheGroupSpec( - ["full"], - FullAttentionSpec( - block_size=block_size, - num_kv_heads=1, - head_size=1, - dtype=torch.float32, - ), - ), - KVCacheGroupSpec( - ["mamba"], - MambaSpec( - block_size=block_size, - shapes=((1, 1),), - dtypes=(torch.float32,), - mamba_cache_mode="align", - ), - ), - ], - ) - spec = SimpleNamespace( - vllm_config=SimpleNamespace(parallel_config=SimpleNamespace(world_size=4)), - kv_cache_config=kv_cache_config, - gpu_block_size=(block_size, block_size), - block_size_factor=1, - hash_block_size=block_size, - offload_prompt_only=False, - ) - - config = SchedulerOffloadConfig.from_spec(spec) - - assert not config.kv_group_configs[0].requires_exact_boundary_source - assert config.kv_group_configs[1].requires_exact_boundary_source - - -def _empty_store_output(**overrides): - output = SimpleNamespace( - scheduled_new_reqs=[], - scheduled_cached_reqs=SimpleNamespace( - req_ids=[], new_block_ids=[], resumed_req_ids=set() - ), - num_scheduled_tokens={}, - finished_req_ids=set(), - preempted_req_ids=set(), - boundary_state_offloads=None, - ) - for name, value in overrides.items(): - setattr(output, name, value) - return output - - -def test_mamba_boundary_store_uses_exact_source_not_positional_mirror(): - scheduler = _make_exact_boundary_scheduler() - output = _empty_store_output( - num_scheduled_tokens={"req": 16}, - boundary_state_offloads={"req": [(1, 99, 16)]}, - ) - - jobs = scheduler._build_boundary_state_store_jobs(output) - - assert len(jobs) == 1 - [job_id] = jobs - src_spec, _ = jobs[job_id].transfer_spec - assert isinstance(src_spec, GPULoadStoreSpec) - assert src_spec.block_ids.tolist() == [99] - assert src_spec.block_ids.tolist() != [21] - assert src_spec.group_sizes == [0, 1] - assert src_spec.block_indices == [0, 0] - assert scheduler._block_id_to_pending_jobs == {99: {job_id}} - - -def test_normal_store_excludes_positional_mamba_align_source(): - scheduler = _make_exact_boundary_scheduler() - output = _empty_store_output( - scheduled_new_reqs=[SimpleNamespace(req_id="req", block_ids=None)], - num_scheduled_tokens={"req": 16}, - ) - - jobs = scheduler._build_store_jobs(output) - - assert len(jobs) == 1 - [job] = jobs.values() - src_spec, _ = job.transfer_spec - assert isinstance(src_spec, GPULoadStoreSpec) - assert src_spec.block_ids.tolist() == [11] - assert src_spec.group_sizes == [1, 0] - assert 21 not in src_spec.block_ids - - -def test_mamba_boundary_store_flushes_before_source_block_reuse(): - scheduler = _make_exact_boundary_scheduler() - first = _empty_store_output( - num_scheduled_tokens={"req": 16}, - boundary_state_offloads={"req": [(1, 99, 16)]}, - ) - first_meta = scheduler.build_connector_meta(first) - [job_id] = first_meta.store_jobs - - second = _empty_store_output( - scheduled_cached_reqs=SimpleNamespace( - req_ids=["req"], - new_block_ids=[([12], [99])], - resumed_req_ids=set(), - ), - num_scheduled_tokens={"req": 0}, - ) - second_meta = scheduler.build_connector_meta(second) - - assert second_meta.jobs_to_flush == {job_id} - - -def test_pending_job_cleanup_deduplicates_block_ids(): - """MTP lookahead may alias a physical block within one completed job.""" - scheduler = object.__new__(OffloadingConnectorScheduler) - scheduler._block_id_to_pending_jobs = {7: {42}} - - scheduler._remove_pending_job(42, [7, 7]) - - assert scheduler._block_id_to_pending_jobs == {} - - def test_loads_do_not_populate_fence_index(request_runner): """Loads don't populate _block_id_to_pending_jobs (protected by delay_free_blocks while in flight).""" @@ -1153,8 +960,8 @@ def setup(r, max_offload_tokens): token_ids=[0] * offloaded_block_size * 3, kv_transfer_params={"max_offload_tokens": max_offload_tokens}, ) - r.manager.prepare_store.side_effect = lambda keys, req_context: ( - generate_store_output(keys) + r.manager.prepare_store.side_effect = ( + lambda keys, req_context: generate_store_output(keys) ) # With sync scheduling, the connector flushes completed stores when the @@ -1252,8 +1059,8 @@ def test_offload_prompt_only(request_runner, async_scheduling: bool): extra_config_overrides={"offload_prompt_only": True}, ) - runner.manager.prepare_store.side_effect = lambda keys, req_context: ( - generate_store_output(keys) + runner.manager.prepare_store.side_effect = ( + lambda keys, req_context: generate_store_output(keys) ) runner.new_request(token_ids=[0] * offloaded_block_size * num_prompt_blocks) diff --git a/tests/v1/kv_connector/unit/test_config.py b/tests/v1/kv_connector/unit/test_config.py index 05189639c8..33c9abd09e 100644 --- a/tests/v1/kv_connector/unit/test_config.py +++ b/tests/v1/kv_connector/unit/test_config.py @@ -104,16 +104,6 @@ def test_kv_connector_rejects_expandable_segments(monkeypatch, kv_connector): _build_config(kv_connector=kv_connector) -def test_native_offloading_allows_expandable_segments(monkeypatch): - """Native CPU offload resolves GPU virtual addresses for each CUDA copy - and explicitly opts into the VMM-safe transfer contract.""" - monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") - - cfg = _build_config(kv_connector="OffloadingConnector") - - assert cfg._connector_supports_vmm_safe_transfers() - - def test_kv_connector_allows_expandable_segments_with_sleep_mode(monkeypatch): """Sleep mode routes KV allocations through CuMemAllocator's pool, which auto-disables expandable_segments (see #40812).""" diff --git a/tests/v1/sample/test_sampler.py b/tests/v1/sample/test_sampler.py index f8c0a5bc65..c67199fa40 100644 --- a/tests/v1/sample/test_sampler.py +++ b/tests/v1/sample/test_sampler.py @@ -5,7 +5,6 @@ import pytest import torch -import vllm.envs as envs from tests.v1.sample.utils import create_allowed_token_ids from vllm.platforms import current_platform from vllm.utils.platform_utils import is_pin_memory_available @@ -26,30 +25,6 @@ MAX_NUM_PROMPT_TOKENS = 64 -def test_sm70_compact_topk20_metadata_gate(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(envs, "VLLM_SM70_COMPACT_TOPK20_SAMPLER", True) - metadata = _create_default_sampling_metadata( - num_output_tokens=0, - batch_size=1, - vocab_size=VOCAB_SIZE, - device=torch.device("cpu"), - ) - metadata.all_greedy = False - metadata.all_random = True - metadata.max_num_logprobs = None - metadata.top_k_cpu = (20,) - metadata.top_p_cpu = (0.95,) - metadata.temperature_cpu = (1.0,) - - assert Sampler._sm70_compact_topk20_metadata_supported(metadata, "raw_logprobs") - metadata.top_k_cpu = (19,) - assert not Sampler._sm70_compact_topk20_metadata_supported(metadata, "raw_logprobs") - metadata.top_k_cpu = (20,) - assert not Sampler._sm70_compact_topk20_metadata_supported( - metadata, "processed_logits" - ) - - def _create_fake_logits(batch_size: int, vocab_size: int) -> torch.Tensor: fake_logits = torch.full((batch_size, vocab_size), 1e-2, dtype=torch.float) return fake_logits diff --git a/tests/v1/sample/test_topk_topp_sampler.py b/tests/v1/sample/test_topk_topp_sampler.py index 3619d75dae..2fc26bfe93 100644 --- a/tests/v1/sample/test_topk_topp_sampler.py +++ b/tests/v1/sample/test_topk_topp_sampler.py @@ -54,69 +54,6 @@ def reset_default_device(): torch.set_default_device(original_device) -@pytest.mark.skipif(not HAS_TRITON, reason="Triton is not available") -def test_sm70_topk_topp_8_warp_policy(monkeypatch: pytest.MonkeyPatch): - from vllm.v1.sample.ops import topk_topp_triton - - monkeypatch.delenv("VLLM_SM70_TOPK_TOPP_B8_B16_8_WARPS", raising=False) - monkeypatch.delenv("VLLM_SM70_TOPK_TOPP_8_WARPS", raising=False) - monkeypatch.setattr(topk_topp_triton.current_platform, "is_cuda", lambda: True) - monkeypatch.setattr( - topk_topp_triton.current_platform, - "is_device_capability", - lambda capability: capability == (7, 0), - ) - - use_8_warps = topk_topp_triton._use_sm70_topk_topp_8_warps - cuda = torch.device("cuda:0") - assert use_8_warps(cuda, 8, 248_320, True, True) - assert use_8_warps(cuda, 16, 248_320, True, True) - assert not use_8_warps(cuda, 5, 248_320, True, True) - - monkeypatch.setenv("VLLM_SM70_TOPK_TOPP_8_WARPS", "1") - assert use_8_warps(cuda, 5, 248_320, True, True) - monkeypatch.setenv("VLLM_SM70_TOPK_TOPP_B8_B16_8_WARPS", "0") - assert not use_8_warps(cuda, 8, 248_320, True, True) - - assert not use_8_warps(cuda, 16, 128_000, True, True) - assert not use_8_warps(cuda, 16, 248_320, False, True) - assert not use_8_warps(torch.device("cpu"), 16, 248_320, True, True) - - -@pytest.mark.parametrize("batch_size", (8, 16)) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -def test_sm70_topk_topp_b8_b16_8_warps_matches_rollback( - monkeypatch: pytest.MonkeyPatch, - batch_size: int, -): - if torch.cuda.get_device_capability() != (7, 0): - pytest.skip("The eight-warp production schedule is SM70-only") - - import vllm.envs as envs - from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton - - generator = Generator(device="cuda").manual_seed(20260824 + batch_size) - logits = torch.randn( - (batch_size, 248_320), - generator=generator, - device="cuda", - dtype=torch.float32, - ) - k = torch.full((batch_size,), 20, device="cuda", dtype=torch.int32) - p = torch.full((batch_size,), 0.95, device="cuda", dtype=torch.float32) - - monkeypatch.setenv("VLLM_SM70_TOPK_TOPP_B8_B16_8_WARPS", "0") - envs.disable_envs_cache() - rollback = apply_top_k_top_p_triton(logits.clone(), k, p) - - monkeypatch.setenv("VLLM_SM70_TOPK_TOPP_B8_B16_8_WARPS", "1") - envs.disable_envs_cache() - candidate = apply_top_k_top_p_triton(logits.clone(), k, p) - - torch.accelerator.synchronize() - assert torch.equal(candidate, rollback) - - def test_flashinfer_sampler_default_falls_back_when_import_fails(monkeypatch): from vllm.v1.sample.ops import topk_topp_sampler as sampler_mod diff --git a/tests/v1/spec_decode/test_ddtree_flex_attention.py b/tests/v1/spec_decode/test_ddtree_flex_attention.py index e7a90e71e1..4728babacb 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 2236e4c699..5ddfa45fb4 100644 --- a/tests/v1/spec_decode/test_ddtree_gpu_runner_positions.py +++ b/tests/v1/spec_decode/test_ddtree_gpu_runner_positions.py @@ -304,7 +304,9 @@ 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)] @@ -315,7 +317,9 @@ 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) ) @@ -385,7 +389,9 @@ 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 6188ec103e..3d58b3822c 100644 --- a/tests/v1/spec_decode/test_ddtree_metadata.py +++ b/tests/v1/spec_decode/test_ddtree_metadata.py @@ -34,8 +34,7 @@ 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 deleted file mode 100644 index 98cda4fe8c..0000000000 --- a/tests/v1/spec_decode/test_dflash2.py +++ /dev/null @@ -1,1310 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from types import SimpleNamespace -from unittest.mock import Mock - -import numpy as np -import pytest -import torch - -import vllm.model_executor.layers.logits_processor as logits_processor_module -import vllm.model_executor.models.qwen3_dflash as dflash_model -import vllm.model_executor.models.qwen3_dflash2 as dflash2_model -import vllm.v1.attention.backends.flash_attn_v100 as flash_v100 -import vllm.v1.worker.gpu.attn_utils as attn_utils -import vllm.v1.worker.gpu.spec_decode.dflash.speculator as dflash_speculator -from vllm import envs -from vllm.config.speculative import SpeculativeConfig -from vllm.model_executor.layers.logits_processor import LogitsProcessor -from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( - _is_dflash2_spec_config, -) -from vllm.model_executor.layers.vocab_parallel_embedding import ( - UnquantizedEmbeddingMethod, - VocabParallelEmbedding, - _sm70_dflash2_dense_order_topk, - _sm70_dflash2_rerank_output_buffers, -) -from vllm.model_executor.models.dflash_sm70 import ( - DFLASH_SM70_GATE_UP_INPUT_SCALE, - DFLASH_SM70_WIDE_OUTPUT_SCALE, - DFlashSM70RMSNorm, - dflash_scale_output_sm70, - dflash_silu_and_mul_sm70, -) -from vllm.model_executor.models.qwen3_dflash import ( - DFlashQwen3ForCausalLM, - DFlashQwen3Model, - _dflash_layer_causal, -) -from vllm.model_executor.models.qwen3_dflash2 import ( - DFlash2Qwen3ForCausalLM, - DFlash2Qwen3Model, - _grouped_conv, - _score_edges, -) -from vllm.v1.attention.backends.flash_attn_v100 import FlashAttnV100Impl -from vllm.v1.core.kv_cache_utils import unify_kv_cache_spec_page_size -from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec -from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample -from vllm.v1.worker.gpu.spec_decode import init_speculator -from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator -from vllm.v1.worker.gpu.spec_decode.dflash2.sparse_rejection import ( - _parse_alignment_steps, - _supports_sparse_sampling_contract, -) -from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import ( - DFlash2Speculator, - _requires_sm70_tail, - _selector_walk_kernel, -) - - -def test_dflash2_gdn_fastpaths_are_default_off(monkeypatch): - names = ( - "VLLM_FLASH_V100_DFLASH2_GROUPED_VERIFY", - "VLLM_SM70_DFLASH2_QPN8_RERANK", - "VLLM_SM70_DFLASH2_VERIFY_FASTPATH", - "VLLM_SM70_DFLASH2_FUSED_GDN_METADATA", - "VLLM_SM70_DFLASH2_GDN_METADATA_SHADOW", - "VLLM_SM70_DFLASH2_FUSED_GDN_VERIFY", - "VLLM_SM70_DFLASH2_FUSED_GDN_NORM", - "VLLM_SM70_DFLASH2_FUSED_GDN_SPLIT", - "VLLM_SM70_DFLASH2_FUSED_SMALLQ_METADATA", - "VLLM_SM70_DFLASH2_GROUPED_SMALLQ_METADATA", - "VLLM_SM70_DFLASH2_FUSED_QKV_PACK", - "VLLM_SM70_DFLASH2_FUSED_GEMMA_RMS", - "VLLM_SM70_DFLASH2_SPARSE_TARGET_REJECTION", - "VLLM_SM70_DFLASH2_SHARDED_CONTEXT_FC", - ) - for name in names: - monkeypatch.delenv(name, raising=False) - envs.disable_envs_cache() - try: - assert not any(getattr(envs, name) for name in names) - 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() - try: - assert envs.VLLM_SM70_TP4_PUSH_ALLREDUCE - monkeypatch.setenv("VLLM_SM70_TP4_PUSH_ALLREDUCE", "0") - envs.disable_envs_cache() - assert not envs.VLLM_SM70_TP4_PUSH_ALLREDUCE - finally: - envs.disable_envs_cache() - - -def _bare_dflash2_model() -> DFlash2Qwen3Model: - model = DFlash2Qwen3Model.__new__(DFlash2Qwen3Model) - torch.nn.Module.__init__(model) - model.quant_config = None - return model - - -def _fake_rms_norm( - output: torch.Tensor, - input_: torch.Tensor, - weight: torch.Tensor, - epsilon: float, -) -> None: - normalized = input_.float() * torch.rsqrt( - input_.float().pow(2).mean(dim=-1, keepdim=True) + epsilon - ) - if weight.ndim == 2: - weight = weight.view(weight.shape[0], *([1] * (input_.ndim - 2)), -1) - output.copy_((normalized * weight.float()).to(output.dtype)) - - -def _bare_context_k_model() -> DFlashQwen3Model: - model = DFlashQwen3Model.__new__(DFlashQwen3Model) - torch.nn.Module.__init__(model) - model._rms_norm_eps = 1e-6 - model._k_norm_weights = torch.tensor( - [ - [1.0, 1.0, 1.0, 1.0], - [0.5, 1.5, 0.75, 1.25], - [1.5, 0.5, 1.25, 0.75], - ], - dtype=torch.float32, - ) - model._batched_k_norm_runtime_verified = None - return model - - -def test_context_k_norm_trusts_verified_batched_runtime(monkeypatch): - model = _bare_context_k_model() - all_k = torch.randn(3, 2, 1, 4) - calls = [] - - def grouped_runtime(output, input_, weight, epsilon): - calls.append(weight.ndim) - _fake_rms_norm(output, input_, weight, epsilon) - - monkeypatch.setattr(dflash_model.ops, "rms_norm", grouped_runtime) - first = model._normalize_context_k(all_k) - assert model._batched_k_norm_runtime_verified is True - assert calls == [2, 1, 1, 1, 2] - - calls.clear() - second = model._normalize_context_k(all_k) - assert calls == [2] - assert torch.equal(second, first) - - -def test_context_k_norm_falls_back_for_stale_stable_binary(monkeypatch): - model = _bare_context_k_model() - # An all-zero profiling input would make the stale grouped output appear - # correct unless the capability probe uses its own nonzero fixture. - all_k = torch.zeros(3, 2, 1, 4) - calls = [] - - def stale_runtime(output, input_, weight, epsilon): - calls.append(weight.ndim) - # Historical stable-ABI binaries silently used row zero for a 2-D - # weight tensor. One-dimensional calls remained correct. - effective_weight = weight[0] if weight.ndim == 2 else weight - _fake_rms_norm(output, input_, effective_weight, epsilon) - - monkeypatch.setattr(dflash_model.ops, "rms_norm", stale_runtime) - actual = model._normalize_context_k(all_k) - expected = torch.empty_like(all_k) - for layer_idx in range(all_k.shape[0]): - _fake_rms_norm( - expected[layer_idx], - all_k[layer_idx], - model._k_norm_weights[layer_idx], - model._rms_norm_eps, - ) - - assert model._batched_k_norm_runtime_verified is False - assert calls == [2, 1, 1, 1, 1, 1, 1] - assert torch.equal(actual, expected) - - calls.clear() - repeated = model._normalize_context_k(all_k) - assert calls == [1, 1, 1] - assert torch.equal(repeated, expected) - - -def test_sm70_tp4_shards_only_compatible_dflash2_context_projection(monkeypatch): - model = _bare_dflash2_model() - config = SimpleNamespace( - parallel_config=SimpleNamespace(tensor_parallel_size=4), - model_config=SimpleNamespace(dtype=torch.float16), - ) - created = {} - - def fake_column_parallel(**kwargs): - created.update(kwargs) - return SimpleNamespace() - - monkeypatch.setattr( - dflash2_model.envs, - "VLLM_SM70_DFLASH2_SHARDED_CONTEXT_FC", - True, - ) - monkeypatch.setattr(dflash2_model.current_platform, "is_cuda", lambda: True) - monkeypatch.setattr( - dflash2_model.current_platform, - "is_device_capability", - lambda capability: capability == 70, - ) - monkeypatch.setattr(dflash2_model, "ColumnParallelLinear", fake_column_parallel) - - projection = model._make_context_projection( - vllm_config=config, - input_size=25600, - output_size=5120, - prefix="model.fc", - ) - - assert created["gather_output"] is True - assert created["input_size"] == 25600 - assert created["output_size"] == 5120 - assert projection._sm70_f16_force_enable is True - assert projection._sm70_f16_max_m == 64 - - -@pytest.mark.parametrize( - ("enabled", "tp_size", "input_size", "output_size"), - [ - (False, 4, 25600, 5120), - (True, 2, 25600, 5120), - (True, 4, 5120, 5120), - (True, 4, 25600, 4096), - ], -) -def test_sharded_context_projection_falls_back_outside_exact_contract( - monkeypatch, enabled, tp_size, input_size, output_size -): - model = _bare_dflash2_model() - config = SimpleNamespace( - parallel_config=SimpleNamespace(tensor_parallel_size=tp_size), - model_config=SimpleNamespace(dtype=torch.float16), - ) - sentinel = object() - monkeypatch.setattr( - dflash2_model.envs, - "VLLM_SM70_DFLASH2_SHARDED_CONTEXT_FC", - enabled, - ) - monkeypatch.setattr(dflash2_model.current_platform, "is_cuda", lambda: True) - monkeypatch.setattr( - dflash2_model.current_platform, - "is_device_capability", - lambda capability: capability == 70, - ) - monkeypatch.setattr( - DFlashQwen3Model, - "_make_context_projection", - lambda *_args, **_kwargs: sentinel, - ) - - projection = model._make_context_projection( - vllm_config=config, - input_size=input_size, - output_size=output_size, - prefix="model.fc", - ) - - assert projection is sentinel - - -@pytest.mark.parametrize("block_size", [4, 6, 8]) -def test_grouped_conv_matches_reference(block_size: int): - torch.manual_seed(0) - batch, taps, num_groups, group_size = 3, 3, 4, 2 - hidden = torch.randn(batch * block_size, num_groups * group_size) - delta = torch.randn(batch * block_size, taps, num_groups) - base = torch.randn(taps, num_groups * group_size) - - actual = _grouped_conv( - hidden, delta, base, block_size, num_groups, group_size, taps - ) - hidden_blocks = hidden.view(batch, block_size, num_groups, group_size) - expected = torch.zeros_like(hidden_blocks) - base = base.view(taps, num_groups, group_size) - delta = delta.view(batch, block_size, taps, num_groups) - for position in range(block_size): - for tap in range(min(taps, position + 1)): - expected[:, position] += ( - base[tap] + delta[:, position, tap, :, None] - ) * hidden_blocks[:, position - tap] - - torch.testing.assert_close(actual, expected.flatten(0, 1).flatten(-2)) - - -def test_selector_edges_match_sequential_reference(): - torch.manual_seed(1) - batch, steps, top_k, rank = 2, 4, 3, 5 - vocab = 17 - predecessors = torch.randn(vocab, rank) - successors = torch.randn(vocab, rank) - candidate_ids = torch.randint(vocab, (batch, steps, top_k)) - unary = torch.randn(batch, steps, top_k) - hidden = torch.randn(batch, steps, rank) - anchors = torch.randint(vocab, (batch,)) - - actual = _score_edges( - predecessors, - successors, - candidate_ids, - unary, - hidden, - anchors, - top_k, - ) - expected = torch.empty_like(actual) - for step in range(steps): - pred = ( - anchors[:, None].expand(-1, top_k) - if step == 0 - else candidate_ids[:, step - 1] - ) - expected[:, step] = unary[:, step, None] + torch.einsum( - "bpr,bcr->bpc", - predecessors[pred] * hidden[:, step, None], - successors[candidate_ids[:, step]], - ) - - torch.testing.assert_close(actual, expected) - - -def _stub_base(monkeypatch: pytest.MonkeyPatch, draft_logits): - def init_base(self, _vllm_config, device): - self.draft_model_config = SimpleNamespace( - hf_config=SimpleNamespace(dflash_config={"selector_top_k": 16}) - ) - self.max_num_reqs = 2 - self.num_query_per_req = 8 - self.num_speculative_steps = 7 - self.vocab_size = 31 - self.draft_tokens = torch.empty((2, 7), dtype=torch.int64, device=device) - self.draft_logits = draft_logits - - monkeypatch.setattr(DFlashSpeculator, "__init__", init_base) - - -def test_selector_leaves_greedy_without_proposal_logits(monkeypatch): - _stub_base(monkeypatch, None) - speculator = DFlash2Speculator(None, torch.device("cpu")) - assert speculator.draft_logits is None - - -def test_selector_default_path_does_not_allocate_sparse_score_cache(monkeypatch): - allocated = torch.zeros((2, 7, 31), dtype=torch.float32) - _stub_base(monkeypatch, allocated) - monkeypatch.setattr(envs, "VLLM_SM70_DFLASH2_SPARSE_TARGET_REJECTION", False) - monkeypatch.setattr(envs, "VLLM_SPEC_DUMP_ALIGNMENT", False) - speculator = DFlash2Speculator(None, torch.device("cpu")) - assert speculator.draft_logits is allocated - assert torch.isneginf(speculator.draft_logits).all() - assert speculator.get_sparse_draft_logits() is None - assert speculator.get_selector_alignment_shadow() is None - - -def test_selector_opt_in_allocates_sparse_score_cache(monkeypatch): - allocated = torch.zeros((2, 7, 31), dtype=torch.float32) - _stub_base(monkeypatch, allocated) - monkeypatch.setattr(envs, "VLLM_SM70_DFLASH2_SPARSE_TARGET_REJECTION", True) - speculator = DFlash2Speculator(None, torch.device("cpu")) - sparse_logits = speculator.get_sparse_draft_logits() - assert sparse_logits is not None - candidate_ids, candidate_scores = sparse_logits - assert candidate_ids.shape == (2, 7, 16) - assert candidate_scores.shape == (2, 7, 16) - assert candidate_scores.dtype is torch.float32 - - -def test_selector_alignment_shadow_is_explicit_and_keeps_full_lattice(monkeypatch): - allocated = torch.zeros((2, 7, 31), dtype=torch.float32) - _stub_base(monkeypatch, allocated) - monkeypatch.setattr(envs, "VLLM_SM70_DFLASH2_SPARSE_TARGET_REJECTION", True) - monkeypatch.setattr(envs, "VLLM_SPEC_DUMP_ALIGNMENT", True) - - speculator = DFlash2Speculator(None, torch.device("cpu")) - shadow = speculator.get_selector_alignment_shadow() - - assert shadow is not None - candidate_ids, unary_logits, lattice_scores = shadow - assert candidate_ids.shape == (2, 7, 16) - assert unary_logits.shape == (2, 7, 16) - assert unary_logits.dtype is torch.float32 - assert lattice_scores.shape == (2, 7, 16, 16) - assert lattice_scores.dtype is torch.float32 - - -@pytest.mark.parametrize( - ("raw", "expected"), - [ - (None, None), - ("", None), - ("1,3-5", {1, 3, 4, 5}), - ("5-3", set()), - ("bad", set()), - ("-1", set()), - ], -) -def test_selector_alignment_step_filter(raw, expected): - assert _parse_alignment_steps(raw) == expected - - -def test_selector_uses_checkpoint_top16_and_fp32_proposal_cache(monkeypatch): - _stub_base(monkeypatch, None) - speculator = DFlash2Speculator(None, torch.device("cpu")) - dtype, fill = DFlash2Speculator.draft_logits_spec(None, None) - assert speculator.selector_top_k == 16 - assert dtype is torch.float32 - assert fill == float("-inf") - - -def _sparse_sampling_contract_fixture(): - idx = np.array([0], dtype=np.int32) - sampling_states = SimpleNamespace( - temperature=SimpleNamespace(np=np.array([1.0], dtype=np.float32)), - top_k=SimpleNamespace(np=np.array([20], dtype=np.int32)), - top_p=SimpleNamespace(np=np.array([0.95], dtype=np.float32)), - min_p=SimpleNamespace(np=np.array([0.0], dtype=np.float32)), - max_num_logprobs=Mock(return_value=-1), - ) - sampler = SimpleNamespace( - sampling_states=sampling_states, - penalties_state=SimpleNamespace(use_penalty=np.array([False])), - logit_bias_state=SimpleNamespace(use_logit_bias=np.array([False])), - bad_words_state=SimpleNamespace( - num_bad_words=SimpleNamespace(np=np.array([0], dtype=np.int32)) - ), - logprob_token_ids_state=SimpleNamespace(max_num_token_ids=Mock(return_value=0)), - compute_nans=False, - ) - rejection_sampler = SimpleNamespace( - rejection_sample_method="standard", - sampler=sampler, - ) - input_batch = SimpleNamespace( - num_reqs=1, - is_prefilling_np=np.array([False]), - idx_mapping_np=idx, - ) - return rejection_sampler, input_batch - - -def test_sparse_target_rejection_accepts_official_sampling_contract(): - rejection_sampler, input_batch = _sparse_sampling_contract_fixture() - assert _supports_sparse_sampling_contract(rejection_sampler, input_batch) - - -@pytest.mark.parametrize( - ("field", "value"), - [ - ("temperature", 0.0), - ("top_k", 16), - ("top_p", 0.0), - ("min_p", 0.05), - ("penalty", True), - ("logit_bias", True), - ("bad_words", 1), - ("logprobs", 1), - ("custom_logprobs", 1), - ("compute_nans", True), - ], -) -def test_sparse_target_rejection_falls_back_for_unsupported_sampling( - field: str, - value: float | int | bool, -): - rejection_sampler, input_batch = _sparse_sampling_contract_fixture() - sampler = rejection_sampler.sampler - if field in {"temperature", "top_k", "top_p", "min_p"}: - getattr(sampler.sampling_states, field).np[0] = value - elif field == "penalty": - sampler.penalties_state.use_penalty[0] = value - elif field == "logit_bias": - sampler.logit_bias_state.use_logit_bias[0] = value - elif field == "bad_words": - sampler.bad_words_state.num_bad_words.np[0] = value - elif field == "logprobs": - sampler.sampling_states.max_num_logprobs.return_value = value - elif field == "custom_logprobs": - sampler.logprob_token_ids_state.max_num_token_ids.return_value = value - else: - sampler.compute_nans = value - - assert not _supports_sparse_sampling_contract(rejection_sampler, input_batch) - - -def test_probabilistic_selector_caches_temperature_applied_scores(): - if not torch.cuda.is_available(): - pytest.skip("CUDA is required for the DFlash2 selector kernel") - - device = torch.device("cuda") - num_steps, top_k = 2, 4 - scores = torch.tensor( - [ - [ - [ - [0.0, 0.5, 1.0, 1.5], - [2.0, 2.5, 3.0, 3.5], - [4.0, 4.5, 5.0, 5.5], - [6.0, 6.5, 7.0, 7.5], - ], - [ - [8.0, 8.5, 9.0, 9.5], - [10.0, 10.5, 11.0, 11.5], - [12.0, 12.5, 13.0, 13.5], - [14.0, 14.5, 15.0, 15.5], - ], - ] - ], - dtype=torch.float32, - device=device, - ) - candidates = torch.arange(top_k, dtype=torch.int64, device=device).repeat( - 1, num_steps, 1 - ) - sample_pos = torch.tensor([10, 11], dtype=torch.int64, device=device) - req_state = torch.zeros(num_steps, dtype=torch.int32, device=device) - temperature = torch.tensor([0.5], dtype=torch.float32, device=device) - seeds = torch.tensor([123], dtype=torch.int64, device=device) - tokens = torch.full((num_steps,), -1, dtype=torch.int64, device=device) - realized = torch.full( - (1, num_steps, top_k), - float("nan"), - dtype=torch.float32, - device=device, - ) - path_state = torch.empty(1, dtype=torch.int32, device=device) - - _selector_walk_kernel[(1,)]( - scores, - candidates, - sample_pos, - req_state, - temperature, - seeds, - tokens, - realized, - path_state, - num_steps=num_steps, - walk_steps=num_steps, - top_k=top_k, - BLOCK_K=top_k, - SAMPLE_PROBABILISTIC=True, - USE_FP64=False, - num_warps=1, - ) - - first_index = int(tokens[0].item()) - torch.testing.assert_close(realized[0, 0], scores[0, 0, 0] / temperature[0]) - torch.testing.assert_close( - realized[0, 1], scores[0, 1, first_index] / temperature[0] - ) - - -def test_probabilistic_cache_respects_column_stride(): - if not torch.cuda.is_available(): - pytest.skip("CUDA is required for the Gumbel kernel") - - device = torch.device("cuda") - vocab_size = 17 - padded_vocab_size = 23 - logits = torch.arange(vocab_size, dtype=torch.float32, device=device)[None] - storage = torch.full( - (1, 2, padded_vocab_size), - float("nan"), - dtype=torch.float32, - device=device, - ) - cache = storage[:, :, :vocab_size] - - gumbel_sample( - logits, - expanded_idx_mapping=torch.tensor([0], dtype=torch.int32, device=device), - temperature=torch.tensor([1.0], dtype=torch.float32, device=device), - seed=torch.tensor([123], dtype=torch.int64, device=device), - pos=torch.tensor([7], dtype=torch.int64, device=device), - apply_temperature=True, - output_processed_logits=cache, - output_processed_logits_col=torch.tensor(1, device=device), - ) - - assert torch.isnan(cache[0, 0]).all() - torch.testing.assert_close(cache[0, 1], logits[0]) - assert torch.isnan(storage[:, :, vocab_size:]).all() - - -def test_probabilistic_cache_keeps_ids_and_scores_in_request_slot_order(monkeypatch): - if not torch.cuda.is_available(): - pytest.skip("CUDA is required for the DFlash2 cache kernel") - - device = torch.device("cuda") - dense_cache = torch.zeros((2, 7, 31), dtype=torch.float32, device=device) - _stub_base(monkeypatch, dense_cache) - monkeypatch.setattr(envs, "VLLM_SM70_DFLASH2_SPARSE_TARGET_REJECTION", True) - speculator = DFlash2Speculator(None, device) - speculator.sample_idx_mapping = torch.tensor( - [1] * 7 + [0] * 7, - dtype=torch.int32, - device=device, - ) - candidate_ids = torch.stack( - ( - torch.arange(16, dtype=torch.int64, device=device).repeat(7, 1), - torch.arange(15, 31, dtype=torch.int64, device=device).repeat(7, 1), - ) - ) - selector_scores = torch.arange( - 2 * 7 * 16, - dtype=torch.float32, - device=device, - ).view(2, 7, 16) - speculator._selector_scores.copy_(selector_scores) - - speculator._cache_draft_logits(candidate_ids, num_sample=14) - sparse_logits = speculator.get_sparse_draft_logits() - assert sparse_logits is not None - cached_ids, cached_scores = sparse_logits - - assert torch.equal(cached_ids[1], candidate_ids[0]) - assert torch.equal(cached_ids[0], candidate_ids[1]) - assert torch.equal(cached_scores[1], selector_scores[0]) - assert torch.equal(cached_scores[0], selector_scores[1]) - assert torch.equal(dense_cache[1].gather(1, candidate_ids[0]), selector_scores[0]) - assert torch.equal(dense_cache[0].gather(1, candidate_ids[1]), selector_scores[1]) - - -def test_dflash2_selector_contract_dispatches_to_mrv2(monkeypatch): - monkeypatch.setattr(DFlash2Speculator, "__init__", lambda self, *_args: None) - config = SimpleNamespace( - speculative_config=SimpleNamespace( - method="dflash", - draft_model_config=SimpleNamespace( - hf_config=SimpleNamespace(dflash_config={"selector_top_k": 16}) - ), - ) - ) - assert isinstance(init_speculator(config, torch.device("cpu")), DFlash2Speculator) - - -def test_dflash_without_selector_stays_on_official_mrv2_speculator(monkeypatch): - monkeypatch.setattr(DFlashSpeculator, "__init__", lambda self, *_args: None) - config = SimpleNamespace( - speculative_config=SimpleNamespace( - method="dflash", - draft_model_config=SimpleNamespace( - hf_config=SimpleNamespace(dflash_config={}) - ), - ) - ) - assert isinstance(init_speculator(config, torch.device("cpu")), DFlashSpeculator) - - -@pytest.mark.parametrize( - ("method", "selector_top_k", "expected"), - [ - ("dflash", 16, True), - ("dflash", 0, False), - ("dflash_ddtree", 16, False), - ("mtp", 16, False), - ("eagle3", 16, False), - ], -) -def test_fused_gdn_verify_config_uses_selector_engine_contract( - method: str, - selector_top_k: int, - expected: bool, -): - config = SimpleNamespace( - speculative_config=SimpleNamespace( - method=method, - draft_model_config=SimpleNamespace( - hf_config=SimpleNamespace( - dflash_config={"selector_top_k": selector_top_k} - ) - ), - ) - ) - - assert _is_dflash2_spec_config(config) is expected - - -@pytest.mark.parametrize( - ("method", "expected"), - [ - ("dflash", False), - ("dflash_ddtree", True), - ("eagle3", True), - ("mtp", True), - ], -) -def test_only_mrv2_dflash_skips_eagle_prefix_block_drop(method, expected): - config = SimpleNamespace( - method=method, - use_eagle=lambda: True, - use_dflash=lambda: method == "dflash", - ) - assert SpeculativeConfig.use_eagle_kv_cache(config) is expected - - -@pytest.mark.parametrize("method", ["eagle3", "mtp"]) -def test_non_dflash_speculators_keep_eagle_dispatch(monkeypatch, method): - from vllm.v1.worker.gpu.spec_decode.eagle.speculator import EagleSpeculator - - monkeypatch.setattr(EagleSpeculator, "__init__", lambda self, *_args: None) - config = SimpleNamespace( - speculative_config=SimpleNamespace( - method=method, - use_eagle=lambda: True, - ) - ) - assert isinstance(init_speculator(config, torch.device("cpu")), EagleSpeculator) - - -def test_top_level_noncausal_override_wins_over_sliding_layer_default(): - config = SimpleNamespace( - is_causal=False, - dflash_config={}, - layer_types=["sliding_attention"], - ) - assert _dflash_layer_causal(config, 0) is False - - -def test_aux_hidden_states_follow_loaded_draft_projection_dtype(): - fc = torch.nn.Linear(10, 2, bias=False, dtype=torch.float16) - fc.input_size = 10 - model = SimpleNamespace(use_aux_hidden_state=True, fc=fc) - outer = SimpleNamespace(model=model) - hidden_states = torch.randn(3, 10, dtype=torch.float32) - - output = DFlashQwen3ForCausalLM.combine_hidden_states(outer, hidden_states) - - assert output.dtype is torch.float16 - - -@pytest.mark.parametrize("mamba_page_size_padded", [None, 16 * 512]) -def test_fp16_draft_cache_grows_padded_fp8_hybrid_pages( - mamba_page_size_padded: int | None, -): - block_size = 16 - target_page_size = 16 * 512 - specs = { - "target.attn": FullAttentionSpec( - block_size=block_size, - num_kv_heads=1, - head_size=256, - dtype=torch.float8_e5m2, - ), - "target.mamba": MambaSpec( - block_size=block_size, - shapes=((target_page_size,),), - dtypes=(torch.uint8,), - page_size_padded=mamba_page_size_padded, - ), - "draft.attn": FullAttentionSpec( - block_size=block_size, - num_kv_heads=2, - head_size=128, - dtype=torch.float16, - ), - } - - unified = unify_kv_cache_spec_page_size(specs) - - expected_page_size = specs["draft.attn"].page_size_bytes - assert {spec.page_size_bytes for spec in unified.values()} == {expected_page_size} - assert unified["target.attn"].block_size == 2 * block_size - assert unified["target.mamba"].block_size == 2 * block_size - assert unified["target.mamba"].page_size_padded == expected_page_size - - -def test_flashinfer_topk_is_capability_gated_on_sm70(monkeypatch): - dflash2_model._flashinfer_topk.cache_clear() - monkeypatch.setattr(dflash2_model.current_platform, "is_cuda", lambda: True) - monkeypatch.setattr( - dflash2_model.current_platform, - "has_device_capability", - lambda capability: capability <= 70, - ) - monkeypatch.setattr(dflash2_model, "has_flashinfer", lambda: True) - assert dflash2_model._flashinfer_topk() is None - dflash2_model._flashinfer_topk.cache_clear() - - -def test_target_topk_uses_reranked_local_candidates_without_dense_logits(monkeypatch): - dense_apply = Mock(side_effect=AssertionError("dense logits must not run")) - values = torch.linspace(5.0, -5.0, 20, dtype=torch.float16).reshape(1, 20) - ids = torch.arange(100, 120, dtype=torch.int64).reshape(1, 20) - lm_head = SimpleNamespace( - quant_method=SimpleNamespace(apply=dense_apply), - weight=torch.empty((62080, 1), dtype=torch.float16), - maybe_get_sm70_dflash2_top20=lambda hidden, top_k, bias: (values, ids), - ) - processor = LogitsProcessor(vocab_size=248320, scale=0.5, soft_cap=3.0) - monkeypatch.setattr( - logits_processor_module, - "get_tensor_model_parallel_world_size", - lambda: 1, - ) - - actual_ids, actual_values = processor.get_topk_tokens_and_logits( - lm_head, - torch.empty((1, 1), dtype=torch.float16), - 20, - ) - - expected_values = (torch.tanh(values / 3.0) * 3.0 * 0.5).float() - assert torch.equal(actual_ids, ids) - torch.testing.assert_close(actual_values, expected_values) - dense_apply.assert_not_called() - - -def test_draft_candidates_use_reranked_lm_head_without_dense_logits(monkeypatch): - dense_apply = Mock(side_effect=AssertionError("dense logits must not run")) - values = torch.linspace(4.0, -4.0, 16, dtype=torch.float16).reshape(1, 16) - ids = torch.arange(200, 216, dtype=torch.int64).reshape(1, 16) - candidate_path = Mock(return_value=(values, ids)) - lm_head = SimpleNamespace( - quant_method=UnquantizedEmbeddingMethod(), - maybe_get_sm70_dflash2_top20=candidate_path, - ) - lm_head.quant_method.apply = dense_apply - model = SimpleNamespace(candidate_selector=SimpleNamespace(top_k=16)) - dflash2 = SimpleNamespace( - lm_head=lm_head, - model=model, - output_multiplier=1.0, - final_logit_softcapping=None, - ) - hidden_states = torch.empty((1, 8), dtype=torch.float16) - monkeypatch.setattr( - dflash2_model, - "get_tensor_model_parallel_world_size", - lambda: 1, - ) - - actual_ids, actual_values = DFlash2Qwen3ForCausalLM.compute_candidates( - dflash2, hidden_states - ) - - candidate_path.assert_called_once_with(hidden_states, 16) - dense_apply.assert_not_called() - assert torch.equal(actual_ids, ids) - assert torch.equal(actual_values, values.float()) - - -def test_lm_head_candidate_interface_falls_back_when_rerank_is_disabled(monkeypatch): - monkeypatch.setattr(envs, "VLLM_SM70_DFLASH2_QPN8_RERANK", False) - layer = SimpleNamespace() - - assert ( - VocabParallelEmbedding.maybe_get_sm70_dflash2_top20( - layer, - torch.empty((1, 8)), - 20, - ) - is None - ) - - -@pytest.mark.parametrize("selector_k", [16, 20]) -@pytest.mark.parametrize("num_rows", [1, 7, 8]) -def test_qpn8_rerank_output_buffers_are_contiguous(selector_k, num_rows): - layer = SimpleNamespace() - for top_k in (16, 20): - setattr( - layer, - f"_sm70_dflash2_rerank_values_{top_k}", - torch.empty((8, top_k), dtype=torch.float16), - ) - setattr( - layer, - f"_sm70_dflash2_rerank_positions_{top_k}", - torch.empty((8, top_k), dtype=torch.int64), - ) - setattr( - layer, - f"_sm70_dflash2_rerank_ids_{top_k}", - torch.empty((8, top_k), dtype=torch.int64), - ) - - buffers = _sm70_dflash2_rerank_output_buffers(layer, num_rows, selector_k) - - assert all(buffer.shape == (num_rows, selector_k) for buffer in buffers) - assert all(buffer.is_contiguous() for buffer in buffers) - - -def test_qpn8_rerank_restores_dense_vocab_tie_order(): - vocab_start = 96 - candidate_ids = torch.tensor( - [[9, 1, 7, 3, 12, 4], [18, 2, 15, 5, 11, 8]], dtype=torch.int64 - ) - candidate_logits = torch.tensor( - [[5, 5, 5, 5, 4, 3], [7, 7, 7, 7, 6, 5]], dtype=torch.float16 - ) - sparse_logits = torch.empty((2, 24), dtype=torch.float16) - actual_values = torch.empty((2, 4), dtype=torch.float16) - actual_ids = torch.empty((2, 4), dtype=torch.int64) - - _sm70_dflash2_dense_order_topk( - sparse_logits, - candidate_ids, - candidate_logits, - actual_values, - actual_ids, - 4, - vocab_start, - ) - - reference = torch.full_like(sparse_logits, -float("inf")) - reference.scatter_(1, candidate_ids, candidate_logits) - expected_values, expected_ids = torch.topk(reference, 4, dim=-1, sorted=True) - assert torch.equal(actual_values, expected_values) - assert torch.equal(actual_ids, expected_ids + vocab_start) - - -@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]) -def test_sm70_f16_rerank_topk_matches_composite_key_contract( - selector_k: int, num_rows: int -): - if torch.cuda.get_device_capability() != (7, 0): - pytest.skip("the exact rerank top-k kernel is SM70-only") - required_ops = ( - "sm70_f16_rerank_keys_out", - "sm70_f16_rerank_topk_out", - ) - if any(not hasattr(torch.ops._C, name) for name in required_ops): - pytest.skip("the exact rerank top-k operators are not built") - - from vllm import _sm70_ops - - generator = torch.Generator().manual_seed(20260824) - local_vocab = 62080 - candidate_ids = torch.stack( - [torch.randperm(local_vocab, generator=generator)[:64] for _ in range(num_rows)] - ).to(device="cuda", dtype=torch.int64) - torch.manual_seed(20260824) - candidate_logits = torch.randn( - (num_rows, 64), dtype=torch.float16, device="cuda" - ).clamp_(-3, 3) - # Force more equal maxima than either requested top-k so ID-ascending tie - # precedence is part of every comparison, not just an incidental edge. - candidate_logits[:, :24] = 4 - if num_rows > 1: - candidate_logits[1, 24:28] = torch.tensor( - [-float("inf"), -0.0, 0.0, float("inf")], - dtype=torch.float16, - device="cuda", - ) - - actual_values = torch.empty( - (num_rows, selector_k), dtype=torch.float16, device="cuda" - ) - actual_ids = torch.empty((num_rows, selector_k), dtype=torch.int64, device="cuda") - vocab_start = 186240 - _sm70_ops.sm70_f16_rerank_topk_out( - actual_values, - actual_ids, - candidate_logits, - candidate_ids, - vocab_start, - ) - - keys = torch.empty_like(candidate_ids) - _sm70_ops.sm70_f16_rerank_keys_out(keys, candidate_logits, candidate_ids) - _, positions = torch.topk(keys, selector_k, dim=-1, sorted=True) - expected_values = candidate_logits.gather(1, positions) - expected_ids = candidate_ids.gather(1, positions).add_(vocab_start) - torch.accelerator.synchronize() - - assert torch.equal( - actual_values.view(torch.int16), expected_values.view(torch.int16) - ) - assert torch.equal(actual_ids, expected_ids) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -@pytest.mark.parametrize("num_rows", [1, 7, 8]) -def test_sm70_dflash2_exact_rerank_matches_gathered_bmm(num_rows): - if torch.cuda.get_device_capability() != (7, 0): - pytest.skip("the exact rerank kernel is SM70-only") - if not hasattr(torch.ops._C, "sm70_f16_indexed_rerank_out"): - pytest.skip("the exact TurboMind rerank op is not built") - - from vllm import _sm70_ops - - torch.manual_seed(20260824) - device = torch.device("cuda") - candidates, vocab, hidden_size = 64, 256, 5120 - hidden = torch.randn((num_rows, hidden_size), dtype=torch.float16, device=device) - weight = torch.randn((vocab, hidden_size), dtype=torch.float16, device=device) - candidate_ids = torch.randint( - 0, vocab, (num_rows, candidates), dtype=torch.int64, device=device - ) - actual = torch.empty((num_rows, candidates), dtype=torch.float16, device=device) - selected_raw = torch.empty( - (num_rows * candidates, hidden_size), dtype=torch.float16, device=device - ) - selected_packed = torch.empty_like(selected_raw) - expanded = torch.empty( - (num_rows, num_rows * candidates), dtype=torch.float16, device=device - ) - partials = torch.empty( - (num_rows, num_rows * candidates), dtype=torch.float32, device=device - ) - barriers = torch.zeros(64, dtype=torch.int32, device=device) - - _sm70_ops.sm70_f16_indexed_rerank_out( - actual, - hidden, - weight, - candidate_ids, - selected_raw, - selected_packed, - expanded, - partials, - barriers, - 128, - 10, - ) - gathered = weight.index_select(0, candidate_ids.reshape(-1)) - expected = torch.bmm( - gathered.view(num_rows, candidates, hidden_size), - hidden.unsqueeze(-1), - ).squeeze(-1) - - torch.testing.assert_close(actual, expected, atol=0.125, rtol=0.002) - - -@pytest.mark.parametrize( - ("capability", "num_steps", "expected"), - [((7, 0), 7, True), ((7, 0), 1, False), ((8, 0), 7, False)], -) -def test_selector_tail_split_is_sm70_only(monkeypatch, capability, num_steps, expected): - monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: capability) - assert _requires_sm70_tail(torch.device("cuda:0"), num_steps) is expected - - -def test_noncausal_draft_cannot_enter_flash_v100_small_query_fast_path(): - impl = SimpleNamespace( - use_flash_v100_decode=True, - smallq_decode_max_query_len=8, - smallq_decode_max_model_len=4096, - ) - metadata = SimpleNamespace( - causal=False, - query_start_loc=torch.tensor([0, 8], dtype=torch.int32), - max_model_len=128, - ) - assert not FlashAttnV100Impl._small_query_decode_enabled(impl, metadata) - - -def test_dflash_attention_builders_receive_the_draft_model_config(monkeypatch): - def fake_replace(config, **updates): - return SimpleNamespace(source=config, **updates) - - monkeypatch.setattr(dflash_speculator, "replace", fake_replace) - target_model_config = object() - draft_model_config = object() - attention_config = object() - speculator = SimpleNamespace( - vllm_config=SimpleNamespace( - model_config=target_model_config, - attention_config=attention_config, - ), - draft_model_config=draft_model_config, - requires_non_causal=True, - ) - - config = DFlashSpeculator.attn_vllm_config.fget(speculator) - - assert config.model_config is draft_model_config - assert config.attention_config.source is attention_config - assert config.attention_config.use_non_causal is True - - -@pytest.mark.parametrize( - ("mask", "expected"), - [ - (np.array([True], dtype=np.bool_), True), - (np.array([True, True], dtype=np.bool_), True), - (np.array([True, False], dtype=np.bool_), False), - (np.array([False], dtype=np.bool_), False), - (None, False), - ], -) -def test_dflash_context_only_prefill_requires_every_live_request(mask, expected): - batch = SimpleNamespace( - num_reqs=1 if mask is None else mask.size, - is_incomplete_prefilling_np=mask, - ) - assert dflash_speculator._is_context_only_prefill(batch) is expected - - -def test_dflash_intermediate_prefill_materializes_context_without_query(monkeypatch): - prepare_inputs = Mock() - monkeypatch.setattr(dflash_speculator, "prepare_dflash_inputs", prepare_inputs) - - speculator = DFlashSpeculator.__new__(DFlashSpeculator) - speculator.max_model_len = 256 - speculator.num_query_per_req = 8 - speculator.num_speculative_steps = 7 - speculator.max_num_reqs = 1 - speculator.max_num_tokens = 8 - speculator.hidden_states = torch.zeros(8, 3) - speculator.draft_tokens = torch.zeros(1, 7, dtype=torch.int64) - speculator.input_buffers = object() - speculator.context_positions = torch.zeros(8, dtype=torch.int64) - speculator.sample_indices = torch.zeros(7, dtype=torch.int64) - speculator.sample_pos = torch.zeros(7, dtype=torch.int64) - speculator.sample_idx_mapping = torch.zeros(7, dtype=torch.int32) - speculator.temperature = torch.ones(1) - speculator.seeds = torch.zeros(1, dtype=torch.int64) - speculator.parallel_drafting_token_id = 0 - speculator.sample_from_anchor = False - speculator.draft_kv_cache_group_id = 0 - speculator.draft_kv_cache_group_ids = [0] - speculator.block_tables = SimpleNamespace( - slot_mappings=[torch.zeros(8, dtype=torch.int64)], - input_block_tables=[torch.zeros(1, 8, dtype=torch.int32)], - kernel_block_sizes=[1], - cp_rank=0, - cp_size=1, - cp_interleave=1, - ) - speculator._context_slot_mappings = torch.zeros(1, 8, dtype=torch.int64) - speculator._layer_group_idx = None - speculator._context_only_prefill_logged = False - speculator._prepare_ngram_assist = Mock( - side_effect=AssertionError("ngram lookup must not run mid-prefill") - ) - speculator.model = SimpleNamespace( - precompute_and_store_context_kv=Mock(), - ) - - input_batch = SimpleNamespace( - num_reqs=1, - num_tokens=4, - seq_lens_cpu_upper_bound=torch.tensor([4], dtype=torch.int32), - is_incomplete_prefilling_np=np.array([True], dtype=np.bool_), - ) - output = speculator.propose( - input_batch=input_batch, - attn_metadata={}, - slot_mappings={}, - last_hidden_states=torch.arange(12, dtype=torch.float32).view(4, 3), - aux_hidden_states=None, - num_sampled=torch.zeros(1, dtype=torch.int32), - num_rejected=torch.zeros(1, dtype=torch.int32), - last_sampled=torch.zeros(1, dtype=torch.int64), - next_prefill_tokens=torch.zeros(1, dtype=torch.int64), - temperature=torch.ones(1), - seeds=torch.zeros(1, dtype=torch.int64), - ) - - prepare_inputs.assert_called_once() - speculator.model.precompute_and_store_context_kv.assert_called_once() - torch.testing.assert_close( - speculator.hidden_states[:4], - torch.arange(12, dtype=torch.float32).view(4, 3), - ) - assert output.tolist() == [[-1] * 7] - - -def test_noncausal_dflash_capture_binds_paged_prefix_attention(monkeypatch): - monkeypatch.setattr(flash_v100, "_is_cuda_graph_capturing", lambda _query: True) - output = torch.empty(8, 1, 1) - paged_prefix = Mock(return_value=output) - impl = SimpleNamespace( - _supports_flash_v100_path=lambda: True, - _layer_debug_info=lambda _layer: { - "layer_name": "draft", - "is_dflash_draft_attn": True, - }, - use_triton_prefill=False, - use_decode_scalar_paged=True, - use_decode_paged_prefill=False, - use_flash_v100_prefill_paged=True, - _small_query_decode_enabled=lambda _metadata: False, - _flash_v100_prefill_with_prefix=paged_prefix, - ) - metadata = SimpleNamespace( - max_query_len=8, - max_seq_len=1024, - num_actual_tokens=8, - causal=False, - query_start_loc=torch.tensor([0, 8], dtype=torch.int32), - query_start_loc_cpu=torch.tensor([0, 8], dtype=torch.int32), - # Capture-time CPU metadata looks like no-prefix prefill, while the - # persistent device metadata is updated before replay. - seq_lens=torch.tensor([17], dtype=torch.int32), - seq_lens_cpu=torch.tensor([8], dtype=torch.int32), - block_table=torch.tensor([[0]], dtype=torch.int32), - ) - query = torch.empty(8, 1, 1) - layer = SimpleNamespace(is_dflash_draft_attn=True) - - result = FlashAttnV100Impl.forward( - impl, - layer, - query, - query, - query, - torch.empty(1), - metadata, - output, - ) - - assert result is output - paged_prefix.assert_called_once() - - -def test_draft_attention_causality_is_resolved_per_kv_group(monkeypatch): - observed_causality = [] - - class CapturedCommonAttentionMetadata: - def __init__(self, **kwargs): - observed_causality.append(kwargs["causal"]) - - monkeypatch.setattr( - attn_utils, - "CommonAttentionMetadata", - CapturedCommonAttentionMetadata, - ) - kv_cache_config = SimpleNamespace( - kv_cache_groups=[SimpleNamespace(), SimpleNamespace()] - ) - attn_utils.build_attn_metadata( - attn_groups=[[], []], - num_reqs=1, - num_tokens=8, - query_start_loc_gpu=torch.tensor([0, 8], dtype=torch.int32), - query_start_loc_cpu=torch.tensor([0, 8], dtype=torch.int32), - max_query_len=8, - seq_lens=torch.tensor([8], dtype=torch.int32), - max_seq_len=8, - block_tables=[ - torch.zeros((1, 1), dtype=torch.int32), - torch.zeros((1, 1), dtype=torch.int32), - ], - slot_mappings=torch.zeros((2, 8), dtype=torch.int64), - kv_cache_config=kv_cache_config, - causal={0: False, 1: True}, - ) - - assert observed_causality == [False, True] - - -def test_sm70_rmsnorm_keeps_bf16_residual_in_fp32(): - norm = DFlashSM70RMSNorm(8, 1e-6, torch.float16) - norm.weight.data.copy_(torch.linspace(0.5, 1.5, 8, dtype=torch.float16)) - x = torch.full((2, 8), 300.0, dtype=torch.float16) - residual = torch.full((2, 8), 70000.0, dtype=torch.float32) - - output, residual_output = norm(x, residual) - expected_residual = ( - (x.float() * DFLASH_SM70_WIDE_OUTPUT_SCALE + residual) - .to(torch.bfloat16) - .float() - ) - - assert residual_output.dtype is torch.float32 - assert torch.isfinite(output).all() - assert residual_output.max() > torch.finfo(torch.float16).max - torch.testing.assert_close(residual_output, expected_residual, rtol=0, atol=0) - - -def test_sm70_swiglu_uses_power_of_two_row_scale(): - gate_up = ( - torch.tensor( - [ - [2000.0, -1500.0, 1000.0, 800.0, 1800.0, 900.0, -700.0, 600.0], - [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], - ], - dtype=torch.float16, - ) - / DFLASH_SM70_GATE_UP_INPUT_SCALE - ) - transported, row_scales = dflash_silu_and_mul_sm70(gate_up) - - assert torch.all(row_scales >= 1) - torch.testing.assert_close( - torch.log2(row_scales), - torch.log2(row_scales).round(), - rtol=0, - atol=0, - ) - assert torch.isfinite(transported).all() - - down = torch.ones((2, 8), dtype=torch.float16) - restored = dflash_scale_output_sm70(down, row_scales) - expected = ( - down.float() * (row_scales[:, None] / DFLASH_SM70_WIDE_OUTPUT_SCALE) - ).half() - torch.testing.assert_close(restored, expected, rtol=0, atol=0) diff --git a/tests/v1/spec_decode/test_dflash2_ngram_assist.py b/tests/v1/spec_decode/test_dflash2_ngram_assist.py deleted file mode 100644 index cbf44d08f4..0000000000 --- a/tests/v1/spec_decode/test_dflash2_ngram_assist.py +++ /dev/null @@ -1,386 +0,0 @@ -# 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 - -from vllm.config.speculative import SpeculativeConfig -from vllm.v1.spec_decode.ngram_proposer import ( - _find_longest_matched_ngram_and_propose_tokens, -) -from vllm.v1.worker.gpu.spec_decode.dflash2.ngram_assist import ( - DFlash2NgramAssist, - find_split_ngram_proposal, -) -from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import ( - DFlash2Speculator, - _apply_ngram_draft_kernel, -) -from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import ( - dflash2_sparse_topk_rejection_sample, - rejection_sample, -) - - -@pytest.mark.parametrize("split", [0, 1, 4, 7, 10]) -@pytest.mark.parametrize(("min_ngram", "max_ngram"), [(1, 1), (2, 3), (3, 5)]) -def test_split_lookup_matches_standalone_ngram( - split: int, min_ngram: int, max_ngram: int -) -> None: - tokens = np.array([1, 2, 3, 4, 1, 2, 3, 5, 1, 2], dtype=np.int32) - expected = _find_longest_matched_ngram_and_propose_tokens( - tokens, - min_ngram, - max_ngram, - max_model_len=64, - k=3, - ) - prefix = tokens[:split] - suffix = tokens[split:].astype(np.int64) - actual = find_split_ngram_proposal( - prefix, - len(prefix), - suffix, - len(suffix), - min_ngram, - max_ngram, - max_model_len=64, - num_draft_tokens=3, - ) - np.testing.assert_array_equal(actual, expected) - - -def test_batch_assist_reports_only_full_width_hits() -> None: - assist = DFlash2NgramAssist(2, 3, num_draft_tokens=2, max_model_len=64) - token_ids = np.zeros((3, 16), dtype=np.int32) - token_ids[2, :6] = [1, 2, 3, 4, 1, 2] - token_ids[0, :4] = [8, 9, 10, 11] - sampled = np.array([[3, -1], [12, -1]], dtype=np.int64) - sampled_lens = np.array([1, 1], dtype=np.int32) - output = np.zeros((2, 2), dtype=np.int64) - output_lens = np.zeros(2, dtype=np.int32) - - hits = assist.propose( - token_ids, - req_state_indices=np.array([2, 0], dtype=np.int32), - prior_lengths=np.array([6, 4], dtype=np.int32), - sampled_token_ids=sampled, - num_sampled_tokens=sampled_lens, - eligible=np.array([True, True]), - output_tokens=output, - output_lengths=output_lens, - ) - - assert hits == 1 - np.testing.assert_array_equal(output[0], [4, 1]) - np.testing.assert_array_equal(output_lens, [2, 0]) - assert assist.num_eligible == 2 - assert assist.num_full_hits == 1 - - -class _CopyEvent: - def __init__(self) -> None: - self.synchronized = False - - def synchronize(self) -> None: - self.synchronized = True - - -def _host_only_speculator() -> DFlash2Speculator: - speculator = object.__new__(DFlash2Speculator) - speculator._ngram_assist = DFlash2NgramAssist( - 2, 3, num_draft_tokens=2, max_model_len=64 - ) - speculator._ngram_num_hits = 0 - speculator._ngram_rounds = 0 - speculator._ngram_skipped_rounds = 0 - speculator.num_speculative_steps = 2 - speculator._ngram_tokens_cpu_tensor = torch.zeros((2, 2), dtype=torch.int64) - speculator._ngram_lengths_cpu_tensor = torch.zeros(2, dtype=torch.int32) - speculator._ngram_tokens_cpu = speculator._ngram_tokens_cpu_tensor.numpy() - speculator._ngram_lengths_cpu = speculator._ngram_lengths_cpu_tensor.numpy() - speculator._ngram_tokens = torch.zeros((2, 2), dtype=torch.int64) - speculator._ngram_lengths = torch.zeros(2, dtype=torch.int32) - return speculator - - -def test_host_token_state_is_requested_only_for_enabled_assist() -> None: - speculator = object.__new__(DFlash2Speculator) - speculator._ngram_assist = None - assert not speculator.requires_host_token_state - - speculator._ngram_assist = object() # type: ignore[assignment] - assert speculator.requires_host_token_state - - -def test_prepare_assist_uses_request_slots_and_skips_only_all_hit() -> None: - speculator = _host_only_speculator() - token_ids = np.zeros((3, 16), dtype=np.int32) - token_ids[2, :6] = [1, 2, 3, 4, 1, 2] - token_ids[0, :6] = [7, 8, 9, 10, 7, 8] - batch = SimpleNamespace( - has_structured_output_reqs=False, - num_reqs=2, - num_draft_tokens_per_req=np.array([2, 2], dtype=np.int32), - seq_lens_cpu_upper_bound=torch.tensor([8, 8], dtype=torch.int32), - is_prefilling_np=np.array([False, False]), - idx_mapping_np=np.array([2, 0], dtype=np.int32), - ) - event = _CopyEvent() - - skip = speculator._prepare_ngram_assist( - batch, - event, - sampled_token_ids_cpu=np.array([[3], [9]], dtype=np.int64), - num_sampled_tokens_cpu=np.array([1, 1], dtype=np.int32), - all_token_ids_cpu=token_ids, - ) - - assert event.synchronized - assert skip - assert speculator._ngram_num_hits == 2 - assert torch.equal( - speculator._ngram_tokens, - torch.tensor([[4, 1], [10, 7]], dtype=torch.int64), - ) - - -def test_prepare_assist_keeps_dflash_draft_for_mixed_hits() -> None: - speculator = _host_only_speculator() - token_ids = np.zeros((3, 16), dtype=np.int32) - token_ids[2, :6] = [1, 2, 3, 4, 1, 2] - token_ids[0, :6] = [7, 8, 9, 10, 11, 12] - batch = SimpleNamespace( - has_structured_output_reqs=False, - num_reqs=2, - num_draft_tokens_per_req=np.array([2, 2], dtype=np.int32), - seq_lens_cpu_upper_bound=torch.tensor([8, 8], dtype=torch.int32), - is_prefilling_np=np.array([False, False]), - idx_mapping_np=np.array([2, 0], dtype=np.int32), - ) - event = _CopyEvent() - - skip = speculator._prepare_ngram_assist( - batch, - event, - sampled_token_ids_cpu=np.array([[3], [13]], dtype=np.int64), - num_sampled_tokens_cpu=np.array([1, 1], dtype=np.int32), - all_token_ids_cpu=token_ids, - ) - - assert event.synchronized - assert not skip - assert speculator._ngram_assist is not None - assert speculator._ngram_assist.num_full_hits == 1 - assert speculator._ngram_num_hits == 0 - assert torch.count_nonzero(speculator._ngram_tokens) == 0 - - -def test_prepare_assist_bypasses_structured_output_without_sync() -> None: - speculator = _host_only_speculator() - batch = SimpleNamespace(has_structured_output_reqs=True) - event = _CopyEvent() - - assert not speculator._prepare_ngram_assist( - batch, - event, - sampled_token_ids_cpu=np.zeros((1, 1), dtype=np.int64), - num_sampled_tokens_cpu=np.ones(1, dtype=np.int32), - all_token_ids_cpu=np.zeros((1, 8), dtype=np.int32), - ) - assert not event.synchronized - - -def test_ngram_assist_rejects_non_dflash_method() -> None: - with pytest.raises(ValueError, match="only supported with method='dflash'"): - SpeculativeConfig( - method="ngram", - num_speculative_tokens=2, - ngram_assist=True, - ) - - -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") - - device = torch.device("cuda") - num_reqs, num_steps, top_k, vocab_size = 2, 3, 4, 32 - ngram_tokens = torch.tensor( - [[11, 12, 13], [21, 22, 23]], dtype=torch.int64, device=device - ) - ngram_lengths = torch.tensor([3, 0], dtype=torch.int32, device=device) - # Batch row 0 belongs to request-state slot 1. - sample_req_state = torch.tensor( - [1, 1, 1, 0, 0, 0], dtype=torch.int32, device=device - ) - draft_tokens = torch.full( - (num_reqs, num_steps), 7, dtype=torch.int64, device=device - ) - cached_ids = ( - torch.arange( - num_reqs * num_steps * top_k, dtype=torch.int64, device=device - ).view(num_reqs, num_steps, top_k) - % vocab_size - ) - cached_scores = torch.randn( - num_reqs, num_steps, top_k, dtype=torch.float32, device=device - ) - draft_logits = torch.full( - (num_reqs, num_steps, vocab_size), - -float("inf"), - dtype=torch.float32, - device=device, - ) - draft_logits.scatter_(2, cached_ids, cached_scores) - miss_ids = cached_ids[0].clone() - miss_scores = cached_scores[0].clone() - miss_logits = draft_logits[0].clone() - - _apply_ngram_draft_kernel[(num_reqs * num_steps,)]( - ngram_tokens, - ngram_lengths, - sample_req_state, - draft_tokens, - draft_tokens.stride(0), - cached_ids, - cached_scores, - cached_ids.stride(0), - cached_ids.stride(1), - draft_logits, - draft_logits.stride(0), - draft_logits.stride(1), - num_steps=num_steps, - top_k=top_k, - BLOCK_K=top_k, - CACHE_DRAFT_LOGITS=True, - CACHE_SCORES=True, - num_warps=1, - ) - - assert torch.equal(draft_tokens[0], ngram_tokens[0]) - assert torch.equal(draft_tokens[1], torch.full((num_steps,), 7, device=device)) - assert torch.equal(cached_ids[0], miss_ids) - assert torch.equal(cached_scores[0], miss_scores) - assert torch.equal(draft_logits[0], miss_logits) - assert torch.equal(cached_ids[1, :, 0], ngram_tokens[0]) - assert torch.equal(cached_scores[1, :, 0], torch.zeros(num_steps, device=device)) - assert torch.isneginf(cached_scores[1, :, 1:]).all() - assert torch.equal( - draft_logits[1].gather(1, ngram_tokens[0, :, None]).squeeze(1), - torch.zeros(num_steps, device=device), - ) - - -@pytest.mark.parametrize("top_p", [1.0, 0.95]) -def test_ngram_one_hot_sparse_rejection_matches_dense(top_p: float) -> None: - if not torch.cuda.is_available(): - pytest.skip("CUDA is required for rejection sampling") - - from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p - - torch.manual_seed(20260825) - device = torch.device("cuda") - num_reqs, num_steps = 16, 3 - target_top_k, draft_top_k, vocab_size = 20, 16, 97 - rows_per_req = num_steps + 1 - num_logits = num_reqs * rows_per_req - - target_ids = torch.stack( - [ - torch.randperm(vocab_size, device=device)[:target_top_k] - for _ in range(num_logits) - ] - ) - target_logits = ( - torch.randn(num_logits, target_top_k, dtype=torch.float32, device=device) - .sort(dim=-1, descending=True) - .values - ) - target_dense = torch.full( - (num_logits, vocab_size), - -float("inf"), - dtype=torch.float32, - device=device, - ) - target_dense.scatter_(1, target_ids, target_logits) - target_dense = apply_top_k_top_p( - target_dense, - torch.full((num_logits,), target_top_k, dtype=torch.int32, device=device), - torch.full((num_logits,), top_p, dtype=torch.float32, device=device), - ) - - proposals = torch.randint( - 0, vocab_size, (num_reqs, num_steps), dtype=torch.int64, device=device - ) - draft_ids = torch.zeros( - num_reqs, num_steps, draft_top_k, dtype=torch.int64, device=device - ) - draft_ids[:, :, 0] = proposals - draft_logits = torch.full( - (num_reqs, num_steps, draft_top_k), - -float("inf"), - dtype=torch.float32, - device=device, - ) - draft_logits[:, :, 0] = 0.0 - draft_dense = torch.full( - (num_reqs, num_steps, vocab_size), - -float("inf"), - dtype=torch.float32, - device=device, - ) - draft_dense.scatter_(2, proposals[:, :, None], 0.0) - - draft_sampled_2d = torch.zeros( - num_reqs, rows_per_req, dtype=torch.int64, device=device - ) - draft_sampled_2d[:, 1:] = proposals - draft_sampled = draft_sampled_2d.flatten() - cu_num_logits = ( - torch.arange(num_reqs + 1, dtype=torch.int32, device=device) * rows_per_req - ) - idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device) - expanded_idx_mapping = idx_mapping.repeat_interleave(rows_per_req) - expanded_local_pos = torch.arange( - rows_per_req, dtype=torch.int32, device=device - ).repeat(num_reqs) - positions = torch.arange(num_logits, dtype=torch.int64, device=device) + 4096 - temperature = torch.ones(num_reqs, dtype=torch.float32, device=device) - seeds = torch.arange(100, 100 + num_reqs, dtype=torch.int64, device=device) - - dense_sampled, dense_lengths = rejection_sample( - target_dense, - draft_dense, - draft_sampled, - cu_num_logits, - positions, - idx_mapping, - expanded_idx_mapping, - expanded_local_pos, - temperature, - seeds, - num_steps, - ) - sparse_sampled, sparse_lengths = dflash2_sparse_topk_rejection_sample( - target_ids, - target_logits, - draft_ids, - draft_logits, - draft_sampled, - cu_num_logits, - positions, - idx_mapping, - temperature, - torch.full((num_reqs,), top_p, dtype=torch.float32, device=device), - seeds, - num_steps, - ) - - assert torch.equal(sparse_lengths, dense_lengths) - steps = torch.arange(rows_per_req, device=device).unsqueeze(0) - valid = steps < dense_lengths[:, None] - assert torch.equal(sparse_sampled[valid], dense_sampled[valid]) diff --git a/tests/v1/spec_decode/test_dflash2_selector_analysis.py b/tests/v1/spec_decode/test_dflash2_selector_analysis.py deleted file mode 100644 index 3316e78b6c..0000000000 --- a/tests/v1/spec_decode/test_dflash2_selector_analysis.py +++ /dev/null @@ -1,102 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import importlib.util -import sys -from pathlib import Path - -import torch - -_ANALYZER_PATH = ( - Path(__file__).parents[3] / "benchmarks/analyze_dflash2_selector_alignment.py" -) -_SPEC = importlib.util.spec_from_file_location( - "analyze_dflash2_selector_alignment", _ANALYZER_PATH -) -assert _SPEC is not None and _SPEC.loader is not None -_ANALYZER = importlib.util.module_from_spec(_SPEC) -sys.modules[_SPEC.name] = _ANALYZER -_SPEC.loader.exec_module(_ANALYZER) - - -def _record(): - candidate_ids = torch.tensor([[10, 11], [20, 21]]) - lattice = torch.tensor( - [ - [[0.0, 0.0], [0.0, 0.0]], - [[0.0, 4.0], [0.0, -4.0]], - ], - dtype=torch.float64, - ) - return _ANALYZER.AlignmentRecord( - path=Path("synthetic.pt"), - step=1, - temperature=1.0, - top_p=1.0, - target_topk_ids=candidate_ids, - target_topk_logits=torch.zeros((2, 2), dtype=torch.float64), - candidate_ids=candidate_ids, - realized_logits=torch.stack((lattice[0, 0], lattice[1, 0])), - unary_logits=torch.zeros((2, 2), dtype=torch.float64), - lattice_scores=lattice, - draft_sampled=torch.tensor([1, 10, 20]), - num_sampled=2, - ) - - -def test_compact_top_p_keeps_the_crossing_token(): - probs = _ANALYZER._compact_probs(torch.log(torch.tensor([0.6, 0.3, 0.1])), 0.7) - torch.testing.assert_close( - probs, torch.tensor([2 / 3, 1 / 3, 0.0], dtype=torch.float64) - ) - - -def test_beta_zero_reconstructs_realized_selector_rows(): - record = _record() - current = _ANALYZER._proposal_probs( - record, - _ANALYZER.ProposalConfig(name="current", use_cached_logits=True), - ) - reconstructed = _ANALYZER._proposal_probs( - record, - _ANALYZER.ProposalConfig(name="beta-zero"), - ) - - for actual, expected in zip(reconstructed, current): - torch.testing.assert_close(actual, expected) - - -def test_future_message_can_prefer_a_better_supported_branch(): - record = _record() - local = _ANALYZER._proposal_probs( - record, - _ANALYZER.ProposalConfig(name="local"), - )[0] - global_chain = _ANALYZER._proposal_probs( - record, - _ANALYZER.ProposalConfig(name="global", future_beta=1.0), - )[0] - - assert local[0] == local[1] - assert global_chain[0] > global_chain[1] - - -def test_greedy_mixture_is_normalized_and_moves_exact_mass(): - record = _record() - baseline = _ANALYZER._proposal_probs( - record, - _ANALYZER.ProposalConfig(name="baseline", use_cached_logits=True), - )[1] - mixed = _ANALYZER._proposal_probs( - record, - _ANALYZER.ProposalConfig( - name="mixture", - greedy_mix=0.25, - use_cached_logits=True, - ), - )[1] - - expected = baseline * 0.75 - expected[torch.argmax(baseline)] += 0.25 - torch.testing.assert_close(mixed, expected) - torch.testing.assert_close(mixed.sum(), torch.tensor(1.0, dtype=torch.float64)) diff --git a/tests/v1/spec_decode/test_dflash_mrv2_config.py b/tests/v1/spec_decode/test_dflash_mrv2_config.py deleted file mode 100644 index 7a0292d9d2..0000000000 --- a/tests/v1/spec_decode/test_dflash_mrv2_config.py +++ /dev/null @@ -1,199 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from dataclasses import dataclass -from types import SimpleNamespace - -import pytest -import torch - -from vllm.config.speculative import SpeculativeConfig -from vllm.config.vllm import VllmConfig -from vllm.model_executor.models.qwen3_dflash import DFlashQwen3ForCausalLM -from vllm.v1.spec_decode.dflash import DFlashProposer -from vllm.v1.spec_decode.llm_base_proposer import SpecDecodeBaseProposer -from vllm.v1.worker.gpu.attn_utils import build_attn_metadata -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 - - -def _config(method: str, num_speculative_tokens: int = 7) -> SpeculativeConfig: - config = object.__new__(SpeculativeConfig) - config.method = method - config.num_speculative_tokens = num_speculative_tokens - config.parallel_drafting = True - return config - - -def test_dflash_and_ddtree_routes_are_disjoint() -> None: - mrv2 = _config("dflash") - ddtree = _config("dflash_ddtree") - - assert mrv2.use_dflash() - assert not mrv2.use_dflash_ddtree() - assert mrv2.use_dflash_family() - assert not ddtree.use_dflash() - assert ddtree.use_dflash_ddtree() - assert ddtree.use_dflash_family() - - -def test_mrv2_dflash_reserves_all_mask_slots() -> None: - assert _config("dflash").max_num_new_slots_for_drafting == 7 - # The retained DDTree path keeps its existing flat-parallel slot contract. - assert _config("dflash_ddtree").max_num_new_slots_for_drafting == 6 - - -def test_dflash_forces_v2_and_rejects_explicit_v1(monkeypatch) -> None: - fake = SimpleNamespace( - speculative_config=SimpleNamespace(use_dflash=lambda: True), - ) - monkeypatch.setattr("vllm.config.vllm.envs.VLLM_USE_V2_MODEL_RUNNER", None) - assert VllmConfig.use_v2_model_runner.fget(fake) - - monkeypatch.setattr("vllm.config.vllm.envs.VLLM_USE_V2_MODEL_RUNNER", False) - with pytest.raises(ValueError, match="implemented only by Model Runner V2"): - VllmConfig.use_v2_model_runner.fget(fake) - - -def test_dflash_target_layers_use_boundary_indices() -> None: - spec_config = SimpleNamespace( - draft_model_config=SimpleNamespace( - hf_config=SimpleNamespace(target_layer_ids=[5, 19, 33, 47, 61]) - ) - ) - assert get_eagle3_aux_layers_from_config(spec_config) == (6, 20, 34, 48, 62) - - -def test_parallel_drafting_token_id_prefers_dflash_config() -> None: - config = SimpleNamespace( - dflash_config={"mask_token_id": 248070}, - pard_token=1, - ) - assert get_parallel_drafting_token_id(config) == 248070 - - -def test_ddtree_topk_adapter_uses_official_logits_processor() -> None: - hidden_states = torch.randn(2, 4) - expected_ids = torch.tensor([[3, 2], [1, 0]]) - expected_logprobs = torch.tensor([[-0.1, -0.2], [-0.3, -0.4]]) - - class Processor: - def get_topk_tokens_and_logprobs(self, head, states, top_k): - assert head == "head" - assert states is hidden_states - assert top_k == 2 - return expected_ids, expected_logprobs - - model = SimpleNamespace( - draft_id_to_target_id=None, - logits_processor=Processor(), - lm_head="head", - ) - actual = DFlashQwen3ForCausalLM.get_topk_tokens_and_logprobs( - model, hidden_states, 2 - ) - assert actual is not None - assert actual[0] is expected_ids - assert actual[1] is expected_logprobs - - -def test_dflash_aux_projection_matches_fc_weight_dtype() -> None: - class FC: - input_size = 4 - weight = torch.empty((4, 4), dtype=torch.float16) - - def __call__(self, hidden_states: torch.Tensor) -> torch.Tensor: - assert hidden_states.dtype == self.weight.dtype - return hidden_states - - model = SimpleNamespace( - model=SimpleNamespace(use_aux_hidden_state=True, fc=FC()), - ) - hidden_states = torch.ones((2, 4), dtype=torch.float32) - - output = DFlashQwen3ForCausalLM.combine_hidden_states(model, hidden_states) - - assert output.dtype == torch.float16 - - -def test_draft_kv_dtype_is_public_and_defaults_to_inherit() -> None: - assert SpeculativeConfig.kv_cache_dtype is None - - -def test_ddtree_draft_config_combines_sm70_kv_and_non_causal(monkeypatch) -> None: - @dataclass(frozen=True) - class CacheConfig: - cache_dtype: str - - @dataclass(frozen=True) - class AttentionConfig: - use_non_causal: bool - - @dataclass(frozen=True) - class DraftConfig: - cache_config: CacheConfig - attention_config: AttentionConfig - - base = DraftConfig( - cache_config=CacheConfig(cache_dtype="fp8_e5m2"), - attention_config=AttentionConfig(use_non_causal=False), - ) - proposer = object.__new__(DFlashProposer) - proposer.speculative_config = SimpleNamespace(kv_cache_dtype=None) - proposer.dflash_causal = False - - monkeypatch.setattr( - SpecDecodeBaseProposer, - "_create_draft_vllm_config", - lambda _self: base, - ) - monkeypatch.setattr( - "vllm.v1.spec_decode.dflash.current_platform", - SimpleNamespace( - is_cuda=lambda: True, - is_device_capability=lambda capability: capability == 70, - ), - ) - - draft = DFlashProposer._create_draft_vllm_config(proposer) - - assert draft.cache_config.cache_dtype == "auto" - assert draft.attention_config.use_non_causal is True - - -def test_draft_attention_metadata_resolves_causal_per_kv_group() -> None: - class Builder: - def build(self, common_prefix_len, common_attn_metadata, **_kwargs): - assert common_prefix_len == 0 - return common_attn_metadata - - class Group: - def __init__(self, layer_name: str): - self.layer_names = [layer_name] - self.builder = Builder() - - def get_metadata_builder(self, _index: int): - return self.builder - - attn_metadata = build_attn_metadata( - attn_groups=[[Group("layer.0")], [Group("layer.1")]], - num_reqs=1, - num_tokens=1, - query_start_loc_gpu=torch.tensor([0, 1], dtype=torch.int32), - query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.int32), - max_query_len=1, - seq_lens=torch.tensor([1], dtype=torch.int32), - max_seq_len=1, - block_tables=[ - torch.zeros((1, 1), dtype=torch.int32), - torch.zeros((1, 1), dtype=torch.int32), - ], - slot_mappings=torch.zeros((2, 1), dtype=torch.int64), - kv_cache_config=SimpleNamespace(kv_cache_groups=[object(), object()]), - causal={0: False, 1: True}, - ) - - assert attn_metadata["layer.0"].causal is False - assert attn_metadata["layer.1"].causal is True diff --git a/tests/v1/spec_decode/test_dspark.py b/tests/v1/spec_decode/test_dspark.py index 3f3251c7b9..aaddd48e98 100644 --- a/tests/v1/spec_decode/test_dspark.py +++ b/tests/v1/spec_decode/test_dspark.py @@ -1,100 +1,20 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from pathlib import Path from types import MethodType, SimpleNamespace import pytest import torch import torch.nn as nn -from vllm.config import ModelConfig, ParallelConfig, SpeculativeConfig from vllm.models.deepseek_v4.nvidia.dspark import ( DSparkDeepseekV4ForCausalLM, DSparkDeepseekV4Model, DSparkMarkovHead, ) -from vllm.platforms import current_platform -from vllm.transformers_utils.configs.deepseek_v4 import DeepseekV4Config from vllm.v1.spec_decode.dspark import DSparkProposer -def _make_dspark_speculative_config( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, - *, - pp_size: int, - tp_size: int, -) -> tuple[ParallelConfig, SpeculativeConfig]: - monkeypatch.setattr(current_platform, "device_count", lambda: 8) - DeepseekV4Config( - architectures=["DeepseekV4ForCausalLM"], - hidden_size=4096, - num_hidden_layers=43, - num_attention_heads=64, - num_key_value_heads=1, - vocab_size=129280, - num_nextn_predict_layers=1, - dspark_target_layer_ids=[40, 41, 42], - dspark_block_size=5, - dspark_noise_token_id=128799, - dspark_markov_rank=256, - ).save_pretrained(tmp_path) - target_model_config = ModelConfig( - model=str(tmp_path), - tokenizer=str(tmp_path), - runner="generate", - dtype=torch.float16, - max_model_len=2048, - ) - target_parallel_config = ParallelConfig( - pipeline_parallel_size=pp_size, - tensor_parallel_size=tp_size, - ) - speculative_config = SpeculativeConfig( - target_model_config=target_model_config, - target_parallel_config=target_parallel_config, - method="dspark", - num_speculative_tokens=7, - ) - return target_parallel_config, speculative_config - - -@pytest.mark.parametrize(("pp_size", "tp_size"), [(2, 4), (4, 2)]) -def test_dspark_draft_is_local_to_last_pipeline_stage( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, - pp_size: int, - tp_size: int, -) -> None: - target_parallel_config, speculative_config = _make_dspark_speculative_config( - tmp_path, - monkeypatch, - pp_size=pp_size, - tp_size=tp_size, - ) - - assert target_parallel_config.pipeline_parallel_size == pp_size - assert speculative_config.draft_parallel_config.pipeline_parallel_size == 1 - assert speculative_config.draft_parallel_config.tensor_parallel_size == tp_size - assert speculative_config.draft_model_config.architectures == ["DSparkDraftModel"] - - -def test_dspark_rejects_target_layers_outside_final_pipeline_stage( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setenv("VLLM_PP_LAYER_PARTITION", "42,1") - - with pytest.raises(ValueError, match="final target pipeline stage"): - _make_dspark_speculative_config( - tmp_path, - monkeypatch, - pp_size=2, - tp_size=4, - ) - - def test_deepseek_v4_dspark_checkpoint_name_mapping() -> None: remap = DSparkDeepseekV4ForCausalLM._remap_dspark_name @@ -107,9 +27,7 @@ def test_deepseek_v4_dspark_checkpoint_name_mapping() -> None: assert remap("mtp.2.markov_head.markov_w1.weight") == ( "model.markov_head.markov_w1.weight" ) - assert remap("mtp.2.confidence_head.proj.weight") == ( - "model.confidence_head.proj.weight" - ) + assert remap("mtp.2.confidence_head.proj.weight") is None class _FakeDSparkModel: @@ -129,13 +47,6 @@ def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor: bias.scatter_(1, ((previous + 1) % self.vocab_size).unsqueeze(1), 10.0) return bias - @staticmethod - def confidence_logits( - hidden_states: torch.Tensor, - markov_embeds: torch.Tensor, - ) -> torch.Tensor: - return hidden_states.sum(dim=-1) + markov_embeds.sum(dim=-1) - @staticmethod def map_draft_to_target(token_ids: torch.Tensor) -> torch.Tensor: return token_ids @@ -150,8 +61,6 @@ def test_dspark_markov_sampling_is_sequential( proposer.model = _FakeDSparkModel() proposer._enable_probabilistic_draft_probs = False proposer._static_draft_vocab = None - proposer.collect_confidence_logits = False - proposer.confidence_scheduling_enabled = False proposer.input_ids = torch.zeros(2 * num_speculative_tokens, dtype=torch.int32) proposer.input_ids[0] = 1 proposer.input_ids[num_speculative_tokens] = 4 @@ -174,8 +83,6 @@ def test_dspark_markov_sampling_is_sequential( for anchor in (1, 4) ] assert output.view(2, num_speculative_tokens).tolist() == expected - assert proposer._last_confidence_logits is None - assert proposer._last_verification_lengths is None @pytest.mark.parametrize("num_speculative_tokens", [5, 7]) @@ -187,8 +94,6 @@ def test_dspark_probabilistic_sampling_returns_sequential_probs( proposer.model = _FakeDSparkModel() proposer._enable_probabilistic_draft_probs = True proposer._static_draft_vocab = None - proposer.collect_confidence_logits = True - proposer.confidence_scheduling_enabled = True proposer.input_ids = torch.zeros(2 * num_speculative_tokens, dtype=torch.int32) proposer.input_ids[0] = 1 proposer.input_ids[num_speculative_tokens] = 4 @@ -227,37 +132,6 @@ def sample_from_logits( for anchor in (1, 4) ] assert output.view(2, num_speculative_tokens).tolist() == expected - assert proposer._last_confidence_logits.shape == (2, num_speculative_tokens) - assert proposer._last_verification_lengths.shape == (2,) - confidence_logits = proposer.take_last_confidence_logits() - verification_lengths = proposer.take_last_verification_lengths() - assert confidence_logits is not None - assert verification_lengths is not None - assert proposer.take_last_confidence_logits() is None - assert proposer.take_last_verification_lengths() is None - - -def test_dspark_static_verification_cap_skips_confidence_projection() -> None: - proposer = object.__new__(DSparkProposer) - proposer.num_speculative_tokens = 5 - proposer.model = _FakeDSparkModel() - proposer._enable_probabilistic_draft_probs = False - proposer._static_draft_vocab = None - proposer.collect_confidence_logits = False - proposer.confidence_scheduling_enabled = True - proposer._max_verification_tokens = 2 - proposer.input_ids = torch.zeros(5, dtype=torch.int32) - proposer._anchor_indices = torch.tensor([0], dtype=torch.int64) - - output, draft_probs = proposer._sample_draft_tokens( - torch.zeros(5, 2), - sampling_metadata=None, # type: ignore[arg-type] - ) - - assert output.shape == (5,) - assert draft_probs is None - assert proposer._last_confidence_logits is None - assert proposer._last_verification_lengths.tolist() == [2] def test_dspark_replicated_linears_return_tensors() -> None: @@ -280,42 +154,3 @@ def test_dspark_replicated_linears_return_tensors() -> None: markov_head.markov_w2 = nn.Identity() markov_embed = torch.randn(3, 4) assert torch.equal(markov_head.bias(markov_embed), markov_embed) - - -def test_dspark_confidence_projection_rejects_missing_weight() -> None: - draft_model = object.__new__(DSparkDeepseekV4ForCausalLM) - nn.Module.__init__(draft_model) - draft_model._confidence_head_loaded = False - - with pytest.raises(RuntimeError, match="did not provide confidence_head"): - draft_model.confidence_logits(torch.empty(1, 1), torch.empty(1, 1)) - - -def test_dspark_confidence_scheduler_keeps_only_confident_prefix() -> None: - proposer = object.__new__(DSparkProposer) - proposer._confidence_temperatures = torch.tensor([[1.0, 2.0, 1.0, 1.0, 1.0]]) - proposer._confidence_threshold = 0.7 - proposer._max_verification_tokens = 4 - logits = torch.tensor( - [ - [2.0, 2.0, 0.0, 4.0, 4.0], - [2.0, 4.0, 4.0, 4.0, 4.0], - [0.0, 4.0, 4.0, 4.0, 4.0], - ] - ) - - lengths = proposer._select_verification_lengths(logits) - - assert lengths.dtype == torch.int32 - assert lengths.tolist() == [2, 4, 0] - - -def test_dspark_confidence_scheduler_zero_threshold_uses_cap() -> None: - proposer = object.__new__(DSparkProposer) - proposer._confidence_temperatures = torch.ones((1, 5)) - proposer._confidence_threshold = 0.0 - proposer._max_verification_tokens = 2 - - lengths = proposer._select_verification_lengths(torch.randn(3, 5)) - - assert lengths.tolist() == [2, 2, 2] diff --git a/tests/v1/spec_decode/test_eagle.py b/tests/v1/spec_decode/test_eagle.py index 4de8315c06..c13de6d4f7 100644 --- a/tests/v1/spec_decode/test_eagle.py +++ b/tests/v1/spec_decode/test_eagle.py @@ -63,7 +63,7 @@ def _create_proposer( elif method == "draft_model": target_model_dir = model_dir draft_model_dir = ar_draft_model_dir - elif method == "dflash_ddtree": + elif method == "dflash": target_model_dir = dflash_target_dir draft_model_dir = dflash_dir else: @@ -73,7 +73,7 @@ def _create_proposer( model=target_model_dir, runner="generate", max_model_len=100, - trust_remote_code=(method == "dflash_ddtree"), + trust_remote_code=(method == "dflash"), ) speculative_config = SpeculativeConfig( @@ -105,7 +105,7 @@ def _create_proposer( attention_config=AttentionConfig(backend=attention_backend), ) - if method == "dflash_ddtree": + if method == "dflash": proposer = DFlashProposer(vllm_config=vllm_config, device=device) elif "eagle" in method: proposer = EagleProposer(vllm_config=vllm_config, device=device) @@ -1132,10 +1132,7 @@ def test_set_inputs_first_pass_dflash(): device = torch.device(current_platform.device_type) num_speculative_tokens = 3 - proposer = _create_proposer("dflash_ddtree", num_speculative_tokens) - # Production initialization assigns the draft KV cache group before this - # method runs. This isolated input-preparation test supplies that contract. - proposer.kv_cache_gid = 0 + proposer = _create_proposer("dflash", num_speculative_tokens) mask_token_id = proposer.parallel_drafting_token_id # Setup batch with 3 requests diff --git a/tests/v1/spec_decode/test_mtp.py b/tests/v1/spec_decode/test_mtp.py index ec3fca1d0c..7c478f81d8 100644 --- a/tests/v1/spec_decode/test_mtp.py +++ b/tests/v1/spec_decode/test_mtp.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from types import SimpleNamespace from unittest import mock import pytest @@ -13,7 +12,6 @@ create_standard_kv_cache_spec, try_get_attention_backend, ) -from vllm import envs from vllm.config import ( CacheConfig, DeviceConfig, @@ -28,121 +26,11 @@ from vllm.platforms import current_platform from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.spec_decode.eagle import EagleProposer -from vllm.v1.spec_decode.llm_base_proposer import ( - SpecDecodeBaseProposer, - _sm70_mtp_hotpath_warmup_batch_sizes, - _sm70_mtp_moe_warmup_sizes, -) mimo_7b_dir = "XiaomiMiMo/MiMo-7B-Base" DEVICE_TYPE = current_platform.device_type -def test_sm70_mtp_moe_warmup_sizes(): - assert _sm70_mtp_moe_warmup_sizes(256, 8, 8192) == (9, 33, 257) - assert _sm70_mtp_moe_warmup_sizes(256, 8, 32) == (9,) - assert _sm70_mtp_moe_warmup_sizes(256, 8, 8192, 16) == ( - *range(1, 17), - 33, - 257, - ) - assert _sm70_mtp_moe_warmup_sizes(0, 8, 8192) == () - - -def test_sm70_mtp_hotpath_warmup_batch_sizes(): - assert _sm70_mtp_hotpath_warmup_batch_sizes(1) == (1,) - assert _sm70_mtp_hotpath_warmup_batch_sizes(4) == (4,) - assert _sm70_mtp_hotpath_warmup_batch_sizes(4, include_alternate=True) == (1, 4) - assert _sm70_mtp_hotpath_warmup_batch_sizes(64, include_alternate=True) == (1, 4) - - -def test_sm70_mtp_concurrency_warmup_is_default_off(monkeypatch): - monkeypatch.delenv("VLLM_SM70_MTP_CONCURRENCY_WARMUP", raising=False) - envs.disable_envs_cache() - try: - assert not envs.VLLM_SM70_MTP_CONCURRENCY_WARMUP - finally: - envs.disable_envs_cache() - - -def test_sm70_mtp_moe_warmup_runs_each_shape_once(): - dummy_run = mock.Mock() - draft_model_config = SimpleNamespace( - is_moe=True, - hf_text_config=SimpleNamespace( - num_experts_per_tok=8, - moe_intermediate_size=512, - ), - get_num_experts=lambda: 256, - get_hidden_size=lambda: 4096, - ) - proposer = SimpleNamespace( - method="mtp", - device=torch.device("cuda"), - draft_model_config=draft_model_config, - vllm_config=SimpleNamespace( - parallel_config=SimpleNamespace(tensor_parallel_size=4) - ), - max_num_tokens=8192, - dummy_run=dummy_run, - _sm70_mtp_moe_warmed=False, - ) - - with ( - mock.patch.object(current_platform, "is_device_capability", return_value=True), - mock.patch.object(torch.accelerator, "synchronize"), - ): - assert SpecDecodeBaseProposer.warmup_sm70_mtp_moe_kernels(proposer) == ( - "mtp_draft_moe", - ) - - assert [call.args[0] for call in dummy_run.call_args_list] == [9, 33, 257] - assert all(call.kwargs["spec_step_idx"] == 0 for call in dummy_run.call_args_list) - - -def test_sm70_mtp_moe_warmup_covers_opted_in_exact_tiles(): - dummy_run = mock.Mock() - draft_model_config = SimpleNamespace( - is_moe=True, - hf_text_config=SimpleNamespace( - num_experts_per_tok=8, - moe_intermediate_size=512, - ), - get_num_experts=lambda: 256, - get_hidden_size=lambda: 2048, - ) - proposer = SimpleNamespace( - method="mtp", - device=torch.device("cuda"), - draft_model_config=draft_model_config, - vllm_config=SimpleNamespace( - parallel_config=SimpleNamespace(tensor_parallel_size=4) - ), - max_num_tokens=8192, - dummy_run=dummy_run, - _sm70_mtp_moe_warmed=False, - ) - - with ( - mock.patch.object(current_platform, "is_device_capability", return_value=True), - mock.patch.object(torch.accelerator, "synchronize"), - mock.patch.object(envs, "VLLM_SM70_MTP_MOE_TUNED_CONFIG", True), - ): - assert SpecDecodeBaseProposer.warmup_sm70_mtp_moe_kernels(proposer) == ( - "mtp_draft_moe", - ) - - assert [call.args[0] for call in dummy_run.call_args_list] == [ - *range(1, 17), - 33, - 257, - 1, - 2, - 9, - 10, - ] - - def _create_mtp_proposer(num_speculative_tokens: int) -> EagleProposer: """Create an MTP proposer with unified model configuration.""" model_config = ModelConfig( diff --git a/tests/v1/spec_decode/test_mtp_draft_prob_alignment.py b/tests/v1/spec_decode/test_mtp_draft_prob_alignment.py index 26316d35d8..66ddb3429c 100644 --- a/tests/v1/spec_decode/test_mtp_draft_prob_alignment.py +++ b/tests/v1/spec_decode/test_mtp_draft_prob_alignment.py @@ -4,10 +4,7 @@ import pytest import torch -from vllm.v1.spec_decode.draft_prob_alignment import ( - get_aligned_draft_probs, - get_aligned_draft_scalar_values, -) +from vllm.v1.spec_decode.draft_prob_alignment import get_aligned_draft_probs from vllm.v1.spec_decode.metadata import SpecDecodeMetadata @@ -118,27 +115,3 @@ def test_get_spec_decode_draft_probs_rejects_cached_token_mismatch(): draft_prob_token_ids=torch.tensor([[1, 9]], dtype=torch.int32), spec_decode_metadata=_metadata([[1, 2]]), ) - - -def test_get_aligned_draft_scalar_values_binds_prefixes_by_req_id(): - values = torch.tensor([[0.1, 0.2, 0.3], [1.1, 1.2, 1.3]]) - actual = get_aligned_draft_scalar_values( - req_ids=["b", "a"], - values=values, - value_req_ids=["a", "b"], - value_token_ids=torch.tensor([[10, 11, 12], [20, 21, 22]], dtype=torch.int32), - spec_decode_metadata=_metadata([[20], [10, 11]]), - ) - - assert torch.equal(actual, torch.tensor([1.1, 0.1, 0.2])) - - -def test_get_aligned_draft_scalar_values_rejects_token_mismatch(): - with pytest.raises(RuntimeError, match="do not match verifier draft tokens"): - get_aligned_draft_scalar_values( - req_ids=["a"], - values=torch.tensor([[0.1, 0.2]]), - value_req_ids=["a"], - value_token_ids=[[1, 9]], - spec_decode_metadata=_metadata([[1, 2]]), - ) diff --git a/tests/v1/spec_decode/test_rejection_sampler_utils.py b/tests/v1/spec_decode/test_rejection_sampler_utils.py index 83bb10599e..28f6044de8 100644 --- a/tests/v1/spec_decode/test_rejection_sampler_utils.py +++ b/tests/v1/spec_decode/test_rejection_sampler_utils.py @@ -7,7 +7,6 @@ import torch from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import ( - dflash2_sparse_topk_rejection_sample, rejection_sample, ) @@ -130,130 +129,6 @@ def _assert_distribution_match( ) -@pytest.mark.parametrize("num_speculative_steps", [3, 7]) -@pytest.mark.parametrize("top_p", [1.0, 0.95]) -def test_dflash2_sparse_topk_matches_dense_rejection( - num_speculative_steps: int, - top_p: float, -): - """Compact p/q support must preserve the dense DFlash2 decision path.""" - torch.manual_seed(20260823) - device = torch.device("cuda") - num_reqs = 4 - target_top_k = 20 - draft_top_k = 16 - rows_per_req = num_speculative_steps + 1 - num_logits = num_reqs * rows_per_req - - # Unique logits avoid making the reference depend on unspecified top-k - # ordering for equal FP16 values. Adversarial tie behavior is a separate - # candidate-set policy gate rather than a rejection-kernel property. - target_topk_ids = torch.stack( - [ - torch.randperm(VOCAB_SIZE, device=device)[:target_top_k] - for _ in range(num_logits) - ] - ) - target_topk_logits = ( - torch.randn(num_logits, target_top_k, device=device) - .sort(dim=-1, descending=True) - .values - + torch.arange(target_top_k, device=device).flip(0) * 1e-3 - ) - target_dense = torch.full( - (num_logits, VOCAB_SIZE), - -float("inf"), - dtype=torch.float32, - device=device, - ) - target_dense.scatter_(1, target_topk_ids, target_topk_logits) - - top_k_tensor = torch.full( - (num_logits,), target_top_k, dtype=torch.int32, device=device - ) - 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) - - draft_topk_ids = torch.stack( - [ - torch.stack( - [ - torch.randperm(VOCAB_SIZE, device=device)[:draft_top_k] - for _ in range(num_speculative_steps) - ] - ) - for _ in range(num_reqs) - ] - ) - draft_topk_logits = torch.randn( - num_reqs, - num_speculative_steps, - draft_top_k, - dtype=torch.float32, - device=device, - ) - draft_dense = torch.full( - (num_reqs, num_speculative_steps, VOCAB_SIZE), - -float("inf"), - dtype=torch.float32, - device=device, - ) - draft_dense.scatter_(2, draft_topk_ids, draft_topk_logits) - - draft_sampled_2d = torch.zeros( - num_reqs, rows_per_req, dtype=torch.int64, device=device - ) - draft_sampled_2d[:, 1:] = draft_topk_ids[:, :, 0] - draft_sampled = draft_sampled_2d.reshape(-1) - cu_num_logits = ( - torch.arange(num_reqs + 1, dtype=torch.int32, device=device) * rows_per_req - ) - idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device) - expanded_idx_mapping = idx_mapping.repeat_interleave(rows_per_req) - expanded_local_pos = torch.arange( - 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) - 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) - - dense_sampled, dense_num_sampled = rejection_sample( - processed_target, - draft_dense, - draft_sampled, - cu_num_logits, - pos, - idx_mapping, - expanded_idx_mapping, - expanded_local_pos, - temperature, - seeds, - num_speculative_steps, - ) - sparse_sampled, sparse_num_sampled = dflash2_sparse_topk_rejection_sample( - target_topk_ids, - target_topk_logits, - draft_topk_ids, - draft_topk_logits, - draft_sampled, - cu_num_logits, - pos, - idx_mapping, - temperature, - top_p_per_req, - seeds, - num_speculative_steps, - ) - - assert torch.equal(sparse_num_sampled, dense_num_sampled) - steps = torch.arange(rows_per_req, device=device).unsqueeze(0) - valid = steps < dense_num_sampled.unsqueeze(1) - assert torch.equal(sparse_sampled[valid], dense_sampled[valid]) - - @pytest.mark.parametrize( "num_speculative_steps,temperature", [ diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index b2c7d06de1..8d7826116e 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -26,7 +26,6 @@ from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2 from vllm.platforms import current_platform from vllm.sampling_params import SamplingParams -from vllm.sequence import IntermediateTensors from vllm.utils.mem_constants import GiB_bytes from vllm.utils.system_utils import update_environment_variables from vllm.utils.torch_utils import set_random_seed @@ -52,27 +51,6 @@ DEVICE_TYPE = current_platform.device_type -def test_unwrap_pipeline_intermediate_hidden_states_uses_semantic_key(): - hidden_states = torch.zeros((2, 3)) - residual = torch.ones_like(hidden_states) - output = IntermediateTensors({"residual": residual, "hidden_states": hidden_states}) - - assert ( - gpu_model_runner_module._unwrap_pipeline_intermediate_hidden_states(output) - is hidden_states - ) - assert ( - gpu_model_runner_module._unwrap_pipeline_intermediate_hidden_states( - hidden_states - ) - is hidden_states - ) - with pytest.raises(RuntimeError, match="must contain a 'hidden_states' tensor"): - gpu_model_runner_module._unwrap_pipeline_intermediate_hidden_states( - IntermediateTensors({"residual": residual}) - ) - - def initialize_kv_cache(runner: GPUModelRunner): """ Only perform necessary steps in GPUModelRunner.initialize_kv_cache() @@ -817,256 +795,6 @@ def test_sample_passes_reordered_draft_probs_to_rejection_sampler(): assert torch.equal(passed_draft_probs, expected_draft_probs) -def test_dspark_verification_lengths_trim_scheduler_drafts(monkeypatch): - runner = object.__new__(GPUModelRunner) - runner._draft_token_ids = torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]) - runner._draft_token_req_ids = ["a", "b"] - runner.draft_token_ids_event = object() - runner.draft_token_ids_cpu = runner._draft_token_ids.clone() - runner.dspark_confidence_scheduling = True - runner._dspark_verification_lengths_cpu = torch.tensor([2, 0], dtype=torch.int32) - monkeypatch.setattr( - "vllm.v1.worker.gpu_model_runner.sm70_trace_event_sync", - lambda *_args, **_kwargs: None, - ) - - token_ids, req_ids = runner._get_draft_token_ids_cpu() - - assert req_ids == ["a", "b"] - assert token_ids == [[1, 2], []] - - -@pytest.mark.skip_global_cleanup -def test_sync_mamba_accepted_token_state_tracks_request_ownership(): - runner = object.__new__(GPUModelRunner) - move_state = object() - resumed_state = object() - replaced_old_state = object() - replaced_new_state = object() - runner._mamba_accepted_token_state_rows = { - "old": (0, object()), - "resumed": (1, resumed_state), - "move": (2, move_state), - "same-id": (3, replaced_old_state), - } - runner.num_accepted_tokens = SimpleNamespace( - np=np.array([5, 3, 4, 5, 0], dtype=np.int32) - ) - runner.spec_state_slot_selectors = SimpleNamespace( - np=np.array([5, 4, 2, 5, 0], dtype=np.int32) - ) - runner.input_batch = SimpleNamespace( - req_ids=["new", "move", "resumed", "same-id"], - num_accepted_tokens_cpu=np.full(5, 9, dtype=np.int32), - spec_num_accepted_tokens_cpu=np.full(5, 9, dtype=np.int32), - ) - runner.requests = { - "new": object(), - "move": move_state, - "resumed": resumed_state, - "same-id": replaced_new_state, - } - scheduler_output = SimpleNamespace( - scheduled_new_reqs=[SimpleNamespace(req_id="new")], - scheduled_cached_reqs=SimpleNamespace(resumed_req_ids={"resumed"}), - ) - - GPUModelRunner._sync_mamba_accepted_token_state( - runner, scheduler_output, num_reqs=4 - ) - - # Only the continuing request keeps its prior accepted-token state. A new - # request, a resumed request, and a replacement using the same request ID - # all start with zero speculative bias (the encoded value 1). - expected_counts = np.array([1, 4, 1, 1], dtype=np.int32) - expected_selectors = np.array([1, 2, 1, 1], dtype=np.int32) - np.testing.assert_array_equal(runner.num_accepted_tokens.np[:4], expected_counts) - np.testing.assert_array_equal( - runner.spec_state_slot_selectors.np[:4], expected_selectors - ) - np.testing.assert_array_equal( - runner.input_batch.num_accepted_tokens_cpu[:4], expected_counts - ) - np.testing.assert_array_equal( - runner.input_batch.spec_num_accepted_tokens_cpu[:4], expected_selectors - ) - - -@pytest.mark.skip_global_cleanup -def test_mamba_align_postprocess_uses_runner_owned_d2h_buffers(monkeypatch): - runner = object.__new__(GPUModelRunner) - req_state = object() - runner.speculative_config = object() - runner.model_config = SimpleNamespace(is_hybrid=True) - runner.cache_config = SimpleNamespace(mamba_cache_mode="align") - runner.num_accepted_tokens = SimpleNamespace( - cpu=torch.zeros(1, dtype=torch.int32), - gpu=torch.zeros(1, dtype=torch.int32), - ) - runner.spec_state_slot_selectors = SimpleNamespace( - cpu=torch.zeros(1, dtype=torch.int32), - gpu=torch.zeros(1, dtype=torch.int32), - ) - runner.input_batch = SimpleNamespace( - req_ids=["req"], - num_accepted_tokens_cpu_tensor=torch.full((1,), 9, dtype=torch.int32), - spec_num_accepted_tokens_cpu_tensor=torch.full((1,), 9, dtype=torch.int32), - ) - runner.requests = {"req": req_state} - runner._get_mamba_bufs = Mock( - return_value=SimpleNamespace(postprocess_align=object()) - ) - runner.kv_cache_config = object() - runner.compilation_config = SimpleNamespace(static_forward_context=object()) - runner.model = SimpleNamespace(get_mamba_state_copy_func=Mock(return_value=())) - runner.mamba_state_idx = {} - runner.num_accepted_tokens_event = Mock() - postprocess = Mock() - monkeypatch.setattr( - gpu_model_runner_module.envs, - "VLLM_MAMBA_ALIGN_CPU_POSTPROCESS", - False, - ) - monkeypatch.setattr( - gpu_model_runner_module.mamba_utils, - "stage_postprocess_inputs_to_gpu", - Mock(), - ) - monkeypatch.setattr( - gpu_model_runner_module.mamba_utils, - "postprocess_mamba_align_gpu", - postprocess, - ) - - GPUModelRunner._update_states_after_model_execute( - runner, - output_token_ids=torch.tensor([[10, 11, -1]], dtype=torch.int32), - scheduler_output=SimpleNamespace(scheduled_ddtree_payloads={}), - ) - - kwargs = postprocess.call_args.kwargs - assert kwargs["num_accepted_tokens_cpu_tensor"] is runner.num_accepted_tokens.cpu - assert ( - kwargs["spec_num_accepted_tokens_cpu_tensor"] - is runner.spec_state_slot_selectors.cpu - ) - assert ( - kwargs["num_accepted_tokens_cpu_tensor"] - is not runner.input_batch.num_accepted_tokens_cpu_tensor - ) - assert ( - kwargs["spec_num_accepted_tokens_cpu_tensor"] - is not runner.input_batch.spec_num_accepted_tokens_cpu_tensor - ) - # The GPU path must not leave any asynchronous DMA targeting mutable - # InputBatch rows; both accepted-state arrays are remapped only after the - # event is synchronized in _prepare_inputs(). - assert torch.equal( - runner.input_batch.num_accepted_tokens_cpu_tensor, - torch.full((1,), 9, dtype=torch.int32), - ) - assert torch.equal( - runner.input_batch.spec_num_accepted_tokens_cpu_tensor, - torch.full((1,), 9, dtype=torch.int32), - ) - - -@pytest.mark.skip_global_cleanup -def test_mamba_align_cpu_postprocess_mirrors_runner_owned_snapshot(monkeypatch): - runner = object.__new__(GPUModelRunner) - req_state = object() - runner.speculative_config = object() - runner.model_config = SimpleNamespace(is_hybrid=True) - runner.cache_config = SimpleNamespace(mamba_cache_mode="align") - runner.num_accepted_tokens = SimpleNamespace( - cpu=torch.zeros(1, dtype=torch.int32), - gpu=torch.zeros(1, dtype=torch.int32), - ) - runner.spec_state_slot_selectors = SimpleNamespace( - cpu=torch.zeros(1, dtype=torch.int32), - gpu=torch.zeros(1, dtype=torch.int32), - ) - runner.input_batch = SimpleNamespace( - req_ids=["req"], - num_accepted_tokens_cpu_tensor=torch.zeros(1, dtype=torch.int32), - spec_num_accepted_tokens_cpu_tensor=torch.zeros(1, dtype=torch.int32), - ) - runner.requests = {"req": req_state} - runner._get_mamba_bufs = Mock(return_value=SimpleNamespace(preprocess=object())) - runner.kv_cache_config = object() - runner.compilation_config = SimpleNamespace(static_forward_context=object()) - runner.model = SimpleNamespace(get_mamba_state_copy_func=Mock(return_value=())) - runner.mamba_state_idx = {} - runner.num_accepted_tokens_event = Mock() - - def cpu_postprocess(*args, **kwargs): - runner.input_batch.num_accepted_tokens_cpu_tensor.fill_(3) - runner.input_batch.spec_num_accepted_tokens_cpu_tensor.fill_(2) - - monkeypatch.setattr( - gpu_model_runner_module.envs, - "VLLM_MAMBA_ALIGN_CPU_POSTPROCESS", - True, - ) - monkeypatch.setattr( - gpu_model_runner_module.mamba_utils, - "postprocess_mamba", - cpu_postprocess, - ) - - GPUModelRunner._update_states_after_model_execute( - runner, - output_token_ids=torch.tensor([[10, 11, -1]], dtype=torch.int32), - scheduler_output=SimpleNamespace(scheduled_ddtree_payloads={}), - ) - - assert runner.num_accepted_tokens.cpu.tolist() == [3] - assert runner.spec_state_slot_selectors.cpu.tolist() == [2] - - -@pytest.mark.skip_global_cleanup -def test_mamba_all_mode_uses_runner_owned_d2h_buffers(monkeypatch): - runner = object.__new__(GPUModelRunner) - req_state = object() - runner.speculative_config = object() - runner.model_config = SimpleNamespace(is_hybrid=True) - runner.cache_config = SimpleNamespace(mamba_cache_mode="all") - runner.num_spec_tokens = 4 - runner.num_accepted_tokens = SimpleNamespace( - cpu=torch.zeros(1, dtype=torch.int32), - gpu=torch.zeros(1, dtype=torch.int32), - ) - runner.spec_state_slot_selectors = SimpleNamespace( - cpu=torch.zeros(1, dtype=torch.int32), - gpu=torch.zeros(1, dtype=torch.int32), - ) - runner.input_batch = SimpleNamespace( - req_ids=["req"], - num_accepted_tokens_cpu_tensor=torch.full((1,), 9, dtype=torch.int32), - spec_num_accepted_tokens_cpu_tensor=torch.full((1,), 9, dtype=torch.int32), - ) - runner.requests = {"req": req_state} - runner.kv_cache_config = object() - runner.mamba_state_idx = {} - runner.num_accepted_tokens_event = Mock() - monkeypatch.setattr( - gpu_model_runner_module.mamba_utils, - "postprocess_mamba_all", - Mock(), - ) - - GPUModelRunner._update_states_after_model_execute( - runner, - output_token_ids=torch.tensor([[10, 11, -1]], dtype=torch.int32), - scheduler_output=SimpleNamespace(scheduled_ddtree_payloads={}), - ) - - assert runner.num_accepted_tokens.cpu.tolist() == [2] - assert runner.spec_state_slot_selectors.cpu.tolist() == [2] - assert runner.input_batch.num_accepted_tokens_cpu_tensor.tolist() == [9] - assert runner.input_batch.spec_num_accepted_tokens_cpu_tensor.tolist() == [9] - - def test_init_kv_cache_with_kv_sharing_invalid_target_layer_order(default_vllm_config): torch.set_default_dtype(torch.float16) layer_0 = "model.layers.0.self_attn.attn" diff --git a/tests/v1/worker/test_gpu_model_runner_v2_kv_zero.py b/tests/v1/worker/test_gpu_model_runner_v2_kv_zero.py deleted file mode 100644 index 19e7c4fe03..0000000000 --- a/tests/v1/worker/test_gpu_model_runner_v2_kv_zero.py +++ /dev/null @@ -1,87 +0,0 @@ -# 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.v1.worker.gpu import model_runner as mrv2 - - -class FakeKVBlockZeroer: - def __init__(self, device: torch.device, pin_memory: bool): - self.device = device - self.pin_memory = pin_memory - self.init_kwargs: dict[str, Any] | None = None - - def init_meta(self, **kwargs: Any) -> None: - kwargs["attn_groups_iter"] = list(kwargs["attn_groups_iter"]) - self.init_kwargs = kwargs - - -def test_v2_initializes_kv_zeroer_from_all_attention_groups(monkeypatch) -> None: - monkeypatch.setattr(mrv2, "KVBlockZeroer", FakeKVBlockZeroer) - monkeypatch.setattr(mrv2, "is_pin_memory_available", lambda: False) - - groups = [[object()], [object(), object()]] - static_forward_context = object() - runner = mrv2.GPUModelRunner.__new__(mrv2.GPUModelRunner) - runner.device = torch.device("cpu") - runner.attn_groups = groups - runner._kernel_block_sizes = [16, 32] - runner.cache_config = SimpleNamespace(cache_dtype="auto") - runner.compilation_config = SimpleNamespace( - static_forward_context=static_forward_context - ) - - runner._init_kv_zero_meta() - - zeroer = runner._kv_block_zeroer - assert isinstance(zeroer, FakeKVBlockZeroer) - assert zeroer.device == runner.device - assert zeroer.pin_memory is False - assert zeroer.init_kwargs == { - "attn_groups_iter": [*groups[0], *groups[1]], - "kernel_block_sizes": [16, 32], - "cache_dtype": "auto", - "runner_only_attn_layers": set(), - "static_forward_context": static_forward_context, - } - - -def test_v2_clears_new_cache_blocks_before_zero_token_return() -> None: - events: list[object] = [] - empty_output = object() - - def no_forward(_: Any) -> object: - events.append("no_forward") - return empty_output - - runner = mrv2.GPUModelRunner.__new__(mrv2.GPUModelRunner) - runner.finish_requests = lambda _: events.append("finish") - runner.free_states = lambda _: events.append("free") - runner.add_requests = lambda _: events.append("add") - runner.update_requests = lambda _: events.append("update") - runner.block_tables = SimpleNamespace( - apply_staged_writes=lambda: events.append("apply") - ) - runner._zero_block_ids = lambda block_ids: events.append(("zero", block_ids)) - runner.kv_connector = SimpleNamespace(no_forward=no_forward) - scheduler_output = SimpleNamespace( - total_num_scheduled_tokens=0, - new_block_ids_to_zero=[3, 7], - ) - - output = runner.execute_model(scheduler_output) - - assert output is empty_output - assert events == [ - "finish", - "free", - "add", - "update", - "apply", - ("zero", [3, 7]), - "no_forward", - ] diff --git a/tests/v1/worker/test_gpu_warmup_blocks.py b/tests/v1/worker/test_gpu_warmup_blocks.py deleted file mode 100644 index b4fcb8bd0a..0000000000 --- a/tests/v1/worker/test_gpu_warmup_blocks.py +++ /dev/null @@ -1,98 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""MRV2 warmup must reserve the same speculative KV tail as the scheduler.""" - -import pytest -import torch - -from vllm.config.speculative import SpeculativeConfig -from vllm.config.vllm import VllmConfig -from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec -from vllm.v1.worker.gpu.warmup import _reserved_block_count - -BLOCK_SIZE = 16 -MAX_MODEL_LEN = 1024 -NUM_SPEC_TOKENS = 7 - - -def _speculative_config(method: str) -> SpeculativeConfig: - config = object.__new__(SpeculativeConfig) - object.__setattr__(config, "method", method) - object.__setattr__(config, "num_speculative_tokens", NUM_SPEC_TOKENS) - object.__setattr__(config, "ddtree_disable_tree_verify", False) - object.__setattr__(config, "ddtree_budget", 24) - return config - - -class _Config: - speculative_config: SpeculativeConfig | None = None - diffusion_config = None - num_speculative_tokens = VllmConfig.num_speculative_tokens - num_lookahead_tokens = VllmConfig.num_lookahead_tokens - - -@pytest.mark.parametrize( - ("method", "expected"), - [ - ("dflash", NUM_SPEC_TOKENS + 1), - ("dflash_ddtree", 25), - ("eagle3", NUM_SPEC_TOKENS), - ("mtp", NUM_SPEC_TOKENS), - ("dspark", NUM_SPEC_TOKENS), - ("draft_model", NUM_SPEC_TOKENS), - ("ngram", 0), - ], -) -def test_num_lookahead_tokens_per_method(method: str, expected: int) -> None: - config = _Config() - config.speculative_config = _speculative_config(method) - - assert config.num_lookahead_tokens == expected - - -def test_num_lookahead_tokens_without_speculation() -> None: - assert _Config().num_lookahead_tokens == 0 - - -def _full_attention_spec() -> FullAttentionSpec: - return FullAttentionSpec( - block_size=BLOCK_SIZE, - num_kv_heads=1, - head_size=1, - dtype=torch.float16, - ) - - -def _mamba_spec(mode: str) -> MambaSpec: - return MambaSpec( - block_size=BLOCK_SIZE, - shapes=((1,),), - dtypes=(torch.float16,), - mamba_cache_mode=mode, - num_speculative_blocks=NUM_SPEC_TOKENS, - ) - - -@pytest.mark.parametrize( - ("spec", "expected"), - [ - (_full_attention_spec(), 2), - (_mamba_spec("align"), 8), - (_mamba_spec("none"), 9), - (_mamba_spec("all"), 9), - ], -) -def test_reserved_block_count_matches_speculative_tail(spec, expected) -> None: - # 14 scheduled tokens plus DFlash's eight lookahead positions crosses the - # second attention block. Align-mode Mamba excludes lookahead from its - # token range but always retains seven speculative state blocks. - assert ( - _reserved_block_count( - 14, - spec, - num_lookahead_tokens=NUM_SPEC_TOKENS + 1, - max_model_len=MAX_MODEL_LEN, - max_encoder_len=0, - ) - == expected - ) diff --git a/tests/v1/worker/test_gpu_worker_static_pp.py b/tests/v1/worker/test_gpu_worker_static_pp.py deleted file mode 100644 index b7e9e8da3a..0000000000 --- a/tests/v1/worker/test_gpu_worker_static_pp.py +++ /dev/null @@ -1,153 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from types import SimpleNamespace -from unittest.mock import patch - -import torch - -import vllm.envs as envs -from vllm.config import CUDAGraphMode -from vllm.sequence import IntermediateTensors -from vllm.v1.worker import gpu_worker -from vllm.v1.worker.gpu_model_runner import GPUModelRunner -from vllm.v1.worker.gpu_worker import Worker - - -def _worker() -> SimpleNamespace: - model_config = SimpleNamespace( - dtype=torch.float16, - ) - return SimpleNamespace( - model_runner=SimpleNamespace( - model=SimpleNamespace( - make_empty_intermediate_tensors=lambda batch_size, dtype, device: ( - IntermediateTensors( - { - "hidden_states": torch.empty( - (batch_size, 4, 4096), dtype=dtype, device=device - ) - } - ) - ) - ) - ), - vllm_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, - compilation_config=SimpleNamespace( - cudagraph_mode=CUDAGraphMode.FULL, - pass_config=SimpleNamespace(enable_sp=False), - ), - model_config=model_config, - ), - ) - - -def test_static_pp_admission_uses_engine_and_tensor_geometry(monkeypatch) -> None: - monkeypatch.setenv("VLLM_SM70_PP_STATIC_HIDDEN_TRANSFER", "1") - envs.disable_envs_cache() - worker = _worker() - try: - with ( - patch.object(gpu_worker.current_platform, "is_cuda", return_value=True), - patch.object( - gpu_worker.current_platform, - "is_device_capability", - return_value=True, - ), - ): - assert Worker._use_sm70_static_pp_hidden_transfer(worker, 1) - assert not Worker._use_sm70_static_pp_hidden_transfer(worker, 2) - - wrong_schema = _worker() - wrong_schema.model_runner.model.make_empty_intermediate_tensors = ( - lambda batch_size, dtype, device: IntermediateTensors( - { - "hidden_states": torch.empty( - (batch_size, 4096), dtype=dtype, device=device - ) - } - ) - ) - assert not Worker._use_sm70_static_pp_hidden_transfer(wrong_schema, 1) - finally: - envs.disable_envs_cache() - - -def test_static_pp_admission_rejects_unsafe_concurrency(monkeypatch) -> None: - monkeypatch.setenv("VLLM_SM70_PP_STATIC_HIDDEN_TRANSFER", "1") - envs.disable_envs_cache() - try: - for field, value in ( - ("max_num_seqs", 2), - ("pipeline_parallel_size", 1), - ("tensor_parallel_size", 8), - ("enable_dbo", True), - ("ubatch_size", 2), - ): - worker = _worker() - target = ( - worker.vllm_config.scheduler_config - if field == "max_num_seqs" - else worker.vllm_config.parallel_config - ) - setattr(target, field, value) - with ( - patch.object(gpu_worker.current_platform, "is_cuda", return_value=True), - patch.object( - gpu_worker.current_platform, - "is_device_capability", - return_value=True, - ), - ): - assert not Worker._use_sm70_static_pp_hidden_transfer(worker, 1) - finally: - envs.disable_envs_cache() - - -def test_static_pp_schema_is_exact() -> None: - hidden_states = SimpleNamespace( - shape=(1, 4, 4096), - dtype=torch.float16, - is_cuda=True, - is_contiguous=lambda: True, - ) - assert Worker._is_static_pp_hidden_tensor_dict( - {"hidden_states": hidden_states}, - 1, # type: ignore[dict-item] - ) - hidden_states.shape = (1, 4096) - assert not Worker._is_static_pp_hidden_tensor_dict( - {"hidden_states": hidden_states}, - 1, # type: ignore[dict-item] - ) - - -def test_static_pp_receive_buffer_skips_self_copy() -> None: - hidden_states = torch.ones((4, 8), dtype=torch.float16) - runner = SimpleNamespace( - intermediate_tensors=IntermediateTensors({"hidden_states": hidden_states}), - vllm_config=SimpleNamespace( - parallel_config=SimpleNamespace(tensor_parallel_size=4) - ), - ) - source = IntermediateTensors({"hidden_states": hidden_states}) - version = hidden_states._version - - with patch( - "vllm.v1.worker.gpu_model_runner.is_residual_scattered_for_sp", - return_value=False, - ): - result = GPUModelRunner.sync_and_gather_intermediate_tensors( - runner, 1, source, True - ) - - assert hidden_states._version == version - assert result["hidden_states"].data_ptr() == hidden_states.data_ptr() diff --git a/tests/v1/worker/test_kv_block_zeroer.py b/tests/v1/worker/test_kv_block_zeroer.py deleted file mode 100644 index 332111b7f4..0000000000 --- a/tests/v1/worker/test_kv_block_zeroer.py +++ /dev/null @@ -1,161 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""KVBlockZeroer must zero the right blocks under a block-major interleaved KV -layout (n_segs segments one cell apart, so the block stride is -n_segs * page_size_el). Build such a layout directly and check that zeroing -block 1 clears only block 1's cells; the old single-page_size_el stride zeros -the wrong cells, so this fails before the fix and passes after. -""" - -from types import SimpleNamespace - -import pytest -import torch - -from vllm.v1.kv_cache_interface import MambaSpec -from vllm.v1.worker.utils import KVBlockZeroer, _infer_segment_block_strides - -CELL_ELS = 8 # int32 elements zeroed per (block, segment) -CELL_BYTES = CELL_ELS * 4 -N_LAYERS = 3 # -> n_segs = 3 interleaved segments -N_BLOCKS = 4 - - -@pytest.mark.parametrize( - ("seg_addrs", "page_size_el", "expected"), - [ - ([0x1000], 8, [8]), - ([0x1000, 0x2000], 8, [8, 8]), - ([0x1000, 0x1020, 0x1040], 8, [24, 24, 24]), - # Two independent two-segment interleaved pools. Input order differs - # from address order to prove strides stay bound to their segments. - ([0x3020, 0x1000, 0x3000, 0x1020], 8, [16, 16, 16, 16]), - ], -) -def test_infer_segment_block_strides( - seg_addrs: list[int], - page_size_el: int, - expected: list[int], -) -> None: - assert _infer_segment_block_strides(seg_addrs, page_size_el) == expected - - -def test_infer_segment_block_strides_rejects_duplicate_addresses() -> None: - with pytest.raises(ValueError, match="must be unique"): - _infer_segment_block_strides([0x1000, 0x1000], 8) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -def test_zero_block_ids_block_major_interleaved() -> None: - device = torch.device("cuda") - - # Block-major interleaved buffer: cell k holds (block B, layer i) where - # k == B * N_LAYERS + i. Fill with a sentinel so any write is visible. - flat = torch.ones(N_BLOCKS * N_LAYERS * CELL_ELS, dtype=torch.int32, device=device) - - # One mamba spec whose page is exactly one cell; page_size_padded lets us - # pin page_size_bytes without depending on real state shapes. - spec = MambaSpec( - block_size=16, - shapes=((1,),), - dtypes=(torch.int32,), - page_size_padded=CELL_BYTES, - ) - layer_names = [f"mamba.{i}" for i in range(N_LAYERS)] - group = SimpleNamespace( - kv_cache_spec=spec, - kv_cache_group_id=0, - layer_names=layer_names, - backend=None, # unused on the mamba path - ) - # Layer i's state view starts one cell into the buffer, so the segment base - # addresses are exactly one cell apart -> the interleaved layout. - static_forward_context = { - name: SimpleNamespace(kv_cache=[flat[i * CELL_ELS :]]) - for i, name in enumerate(layer_names) - } - - zeroer = KVBlockZeroer(device, pin_memory=False) - zeroer.init_meta( - [group], - kernel_block_sizes=[spec.block_size], - cache_dtype="auto", - runner_only_attn_layers=set(), - static_forward_context=static_forward_context, - ) - - # Zero block 1 (block 0 is correct under both old and new strides, so it - # does not discriminate; block 1 does). - target = 1 - zeroer.zero_block_ids([target]) - torch.accelerator.synchronize() - - cells = flat.view(N_BLOCKS * N_LAYERS, CELL_ELS) - zeroed = {int(k) for k in range(cells.shape[0]) if bool((cells[k] == 0).all())} - expected = {target * N_LAYERS + i for i in range(N_LAYERS)} - assert zeroed == expected, ( - f"expected only block {target}'s cells {sorted(expected)} to be " - f"zeroed, got {sorted(zeroed)}" - ) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -def test_zero_block_ids_multiple_interleaved_pools() -> None: - device = torch.device("cuda") - layers_per_pool = 2 - buffers = [ - torch.ones( - N_BLOCKS * layers_per_pool * CELL_ELS, - dtype=torch.int32, - device=device, - ) - for _ in range(2) - ] - spec = MambaSpec( - block_size=16, - shapes=((1,),), - dtypes=(torch.int32,), - page_size_padded=CELL_BYTES, - ) - groups = [] - static_forward_context = {} - for pool_index, buffer in enumerate(buffers): - layer_names = [ - f"mamba.{pool_index}.{layer_index}" - for layer_index in range(layers_per_pool) - ] - groups.append( - SimpleNamespace( - kv_cache_spec=spec, - kv_cache_group_id=pool_index, - layer_names=layer_names, - backend=None, - ) - ) - static_forward_context.update( - { - name: SimpleNamespace(kv_cache=[buffer[index * CELL_ELS :]]) - for index, name in enumerate(layer_names) - } - ) - - zeroer = KVBlockZeroer(device, pin_memory=False) - zeroer.init_meta( - groups, - kernel_block_sizes=[spec.block_size] * len(groups), - cache_dtype="auto", - runner_only_attn_layers=set(), - static_forward_context=static_forward_context, - ) - - target = 1 - zeroer.zero_block_ids([target]) - torch.accelerator.synchronize() - - expected = {target * layers_per_pool + index for index in range(layers_per_pool)} - for buffer in buffers: - cells = buffer.view(N_BLOCKS * layers_per_pool, CELL_ELS) - zeroed = { - index for index in range(cells.shape[0]) if bool((cells[index] == 0).all()) - } - assert zeroed == expected diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index fc2ecea746..0a64067c1a 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -543,32 +543,6 @@ def test_warmup_batch_memcpy_kernel(): assert worker_mamba_utils.warmup_batch_memcpy_kernel(torch.device("cuda:0")) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_warmup_fused_postprocess_does_not_modify_state(): - device = torch.device("cuda:0") - cfg = _TestConfig() - layer_names = ["layer_0"] - kv_cache_config = _make_kv_cache_config(cfg, layer_names) - gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device) - _, _, conv_state, temporal_state, _, forward_context = _make_dual_layer_state( - cfg, device - ) - block_table = torch.arange( - cfg.max_num_reqs * cfg.num_blocks, - dtype=torch.int32, - device=device, - ).reshape(cfg.max_num_reqs, cfg.num_blocks) - gpu_ctx.initialize_from_forward_context( - kv_cache_config, forward_context, _COPY_FUNCS, [block_table] - ) - conv_before = conv_state.clone() - temporal_before = temporal_state.clone() - - assert gpu_ctx.warmup_fused_postprocess() - assert torch.equal(conv_state, conv_before) - assert torch.equal(temporal_state, temporal_before) - - @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") class TestPostprocessMambaFusedKernel: """Tests for postprocess_mamba_fused_kernel comparing GPU vs CPU paths.""" @@ -715,7 +689,9 @@ def test_mtp4_boundary_rollover_counts_1_to_5(self, device, test_config): torch.manual_seed(20260617) accepted = [1, 2, 3, 4, 5] - req_ids = [f"same_{i}" for i in accepted] + [f"cross_{i}" for i in accepted] + req_ids = [f"same_{i}" for i in accepted] + [ + f"cross_{i}" for i in accepted + ] num_accepted_tokens = accepted + accepted # scheduled=5 and draft=4 gives running_state = computed + 1. num_scheduled_tokens = {req_id: 5 for req_id in req_ids} @@ -766,7 +742,9 @@ def test_mtp4_boundary_rollover_counts_1_to_5(self, device, test_config): torch.accelerator.synchronize() gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device) - block_table_gpu = torch.zeros(len(req_ids), 8, dtype=torch.int32, device=device) + block_table_gpu = torch.zeros( + len(req_ids), 8, dtype=torch.int32, device=device + ) for i, block_ids in enumerate(block_ids_per_req): block_table_gpu[i, : len(block_ids)] = torch.tensor( block_ids, dtype=torch.int32, device=device diff --git a/tools/pre_commit/check_forbidden_imports.py b/tools/pre_commit/check_forbidden_imports.py index b71d64a677..365b2f5bb7 100644 --- a/tools/pre_commit/check_forbidden_imports.py +++ b/tools/pre_commit/check_forbidden_imports.py @@ -80,12 +80,7 @@ class ForbiddenImport: allowed_pattern=re.compile( "from vllm.triton_utils import (triton|tl|tl, triton)" ), - 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", - }, + allowed_files={"vllm/triton_utils/importing.py"}, ), } diff --git a/tools/pre_commit/check_torch_cuda.py b/tools/pre_commit/check_torch_cuda.py index 2511f241ba..bfbb36ffbf 100644 --- a/tools/pre_commit/check_torch_cuda.py +++ b/tools/pre_commit/check_torch_cuda.py @@ -19,17 +19,7 @@ "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..ec62b2101e 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -3032,32 +3032,6 @@ def sm70_tp4_all_reduce_gemma_rms_norm( ) -def sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather( - fa: int, - inp: torch.Tensor, - residual: torch.Tensor, - weight: torch.Tensor, - normalized_out: torch.Tensor, - residual_out: torch.Tensor, - reg_input_buffer: int, - reg_output_buffer: int, - reg_buffer_sz_bytes: int, - epsilon: float, -) -> None: - torch.ops._C_custom_ar.sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather( - fa, - inp, - residual, - weight, - normalized_out, - residual_out, - reg_input_buffer, - reg_output_buffer, - reg_buffer_sz_bytes, - epsilon, - ) - - def all_reduce_sum2( fa: int, inp_a: torch.Tensor, @@ -3149,18 +3123,10 @@ def meta_size() -> int: return torch.ops._C_custom_ar.meta_size() -def sm70_tp4_push_allreduce_buffer_size() -> int: - return torch.ops._C_custom_ar.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) -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) - - 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) diff --git a/vllm/_sm70_ops.py b/vllm/_sm70_ops.py index 5870e325e4..0bcd044544 100644 --- a/vllm/_sm70_ops.py +++ b/vllm/_sm70_ops.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import os -from pathlib import Path from typing import TYPE_CHECKING import torch @@ -11,44 +9,6 @@ current_platform.import_kernels() - -def _maybe_load_fp8_qpn8_library() -> None: - """Load an explicitly selected source-built QPN8 extension. - - Production builds register these operators in ``vllm._C``. This opt-in - path lets source experiments add only the QPN8 operators to an otherwise - compatible installed build, including in spawned TP workers. - """ - library_path = os.getenv("VLLM_SM70_FP8_QPN8_LIBRARY") - if library_path is None: - return - generic_override = os.getenv("VLLM_SM70_FP8_QPN8") - specific_override = os.getenv("VLLM_SM70_FP8_QPN8_PP2_TP4") - generic_enabled = generic_override == "1" - default_route_enabled = generic_override != "0" and specific_override != "0" - if generic_enabled or default_route_enabled: - torch.ops.load_library(library_path) - - -_maybe_load_fp8_qpn8_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"): - return - - library_path = os.getenv("VLLM_SM70_SAMPLER_LIBRARY") - if library_path is None: - bundled = sorted(Path(__file__).resolve().parent.glob("_sm70_sampler_C*.so")) - if bundled: - library_path = str(bundled[-1]) - if library_path is not None: - torch.ops.load_library(library_path) - - -_maybe_load_sm70_sampler_library() - if TYPE_CHECKING: def register_fake(fn): @@ -461,647 +421,6 @@ def _fp8_gemm_sm70_out_fake( return None -def fp8_gemm_sm70_prefill_prescaled_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_prefill_prescaled_out")( - out, input, qweight, prescaled_factors, group_size, k_ld, q_ld - ) - - -if hasattr(torch.ops._C, "fp8_gemm_sm70_prefill_prescaled_out"): - - @register_fake("_C::fp8_gemm_sm70_prefill_prescaled_out") - def _fp8_gemm_sm70_prefill_prescaled_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_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) - - -if hasattr(torch.ops._C, "fp8_qpn8_prepare_sm70"): - - @register_fake("_C::fp8_qpn8_prepare_sm70") - def _fp8_qpn8_prepare_sm70_fake( - qweight: torch.Tensor, - scales: torch.Tensor, - ) -> list[torch.Tensor]: - n = qweight.size(0) - k = qweight.size(1) - codes = torch.empty((k, n), dtype=torch.uint8, device=qweight.device) - scale_shape = (1, n) if scales.shape == (n, 1) else (k // 128, n // 32) - group_scales = torch.empty( - scale_shape, dtype=torch.float16, device=scales.device - ) - return [codes, group_scales] - - -def fp8_qpn8_dequantize_sm70_out( - out: torch.Tensor, - codes: torch.Tensor, - 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) - - -if hasattr(torch.ops._C, "fp8_qpn8_dequantize_sm70_out"): - - @register_fake("_C::fp8_qpn8_dequantize_sm70_out") - def _fp8_qpn8_dequantize_sm70_out_fake( - out: torch.Tensor, - codes: torch.Tensor, - group_scales: torch.Tensor, - ) -> None: - return None - - -def fp8_qpn8_prefill_sm70_out( - out: torch.Tensor, - dense_weight_ptr: int, - input: torch.Tensor, - codes: torch.Tensor, - group_scales: torch.Tensor, - gated_silu: bool, -) -> None: - """Dequantize QPN8 into bounded workspace and run a large-M FP16 GEMM.""" - _op("fp8_qpn8_prefill_sm70_out")( - out, - dense_weight_ptr, - input, - codes, - group_scales, - gated_silu, - ) - - -if hasattr(torch.ops._C, "fp8_qpn8_prefill_sm70_out"): - - @register_fake("_C::fp8_qpn8_prefill_sm70_out") - def _fp8_qpn8_prefill_sm70_out_fake( - out: torch.Tensor, - dense_weight_ptr: int, - input: torch.Tensor, - codes: torch.Tensor, - group_scales: torch.Tensor, - gated_silu: bool, - ) -> None: - return None - - -def fp8_qpn8_dispatch_sm70_out( - out: torch.Tensor, - dense_weight_ptr: int, - input: torch.Tensor, - codes: torch.Tensor, - group_scales: torch.Tensor, - split_k: int, - accumulator_chains: int, - prefetch_codes: bool, - gated_silu: bool, -) -> None: - """Runtime-dispatch dynamic M without specializing a Python branch.""" - _op("fp8_qpn8_dispatch_sm70_out")( - out, - dense_weight_ptr, - input, - codes, - group_scales, - split_k, - accumulator_chains, - prefetch_codes, - gated_silu, - ) - - -if hasattr(torch.ops._C, "fp8_qpn8_dispatch_sm70_out"): - - @register_fake("_C::fp8_qpn8_dispatch_sm70_out") - def _fp8_qpn8_dispatch_sm70_out_fake( - out: torch.Tensor, - dense_weight_ptr: int, - input: torch.Tensor, - codes: torch.Tensor, - group_scales: torch.Tensor, - split_k: int, - accumulator_chains: int, - prefetch_codes: bool, - gated_silu: bool, - ) -> None: - return None - - -def fp8_qpn8_gemm_sm70_out( - out: torch.Tensor, - input: torch.Tensor, - codes: torch.Tensor, - group_scales: torch.Tensor, - split_k: int, - accumulator_chains: int, - fast_decoder: bool, - prefetch_codes: bool = False, -) -> None: - """Run the SM70 QPN8 FP8 GEMM into ``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")( - out, - input, - codes, - group_scales, - split_k, - accumulator_chains, - fast_decoder, - prefetch_codes, - ) - - -if hasattr(torch.ops._C, "fp8_qpn8_gemm_sm70_out"): - - @register_fake("_C::fp8_qpn8_gemm_sm70_out") - def _fp8_qpn8_gemm_sm70_out_fake( - out: torch.Tensor, - input: torch.Tensor, - codes: torch.Tensor, - group_scales: torch.Tensor, - split_k: int, - accumulator_chains: int, - fast_decoder: bool, - prefetch_codes: bool, - ) -> None: - return None - - -def fp8_qpn8_gemm_ba_split_sm70_out( - qkv_out: torch.Tensor, - z_out: torch.Tensor, - b_out: torch.Tensor, - a_out: torch.Tensor, - input: torch.Tensor, - codes: torch.Tensor, - group_scales: torch.Tensor, - 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")( - qkv_out, - z_out, - b_out, - a_out, - input, - codes, - group_scales, - ba_weight, - ) - - -if hasattr(torch.ops._C, "fp8_qpn8_gemm_ba_split_sm70_out"): - - @register_fake("_C::fp8_qpn8_gemm_ba_split_sm70_out") - def _fp8_qpn8_gemm_ba_split_sm70_out_fake( - qkv_out: torch.Tensor, - z_out: torch.Tensor, - b_out: torch.Tensor, - a_out: torch.Tensor, - input: torch.Tensor, - codes: torch.Tensor, - group_scales: torch.Tensor, - ba_weight: torch.Tensor, - ) -> None: - return None - - -def fp8_qpn8_dispatch_ba_split_sm70_out( - qkv_out: torch.Tensor, - z_out: torch.Tensor, - b_out: torch.Tensor, - a_out: torch.Tensor, - qkvz_staging: torch.Tensor, - ba_staging: torch.Tensor, - dense_weight_ptr: int, - input: torch.Tensor, - codes: torch.Tensor, - group_scales: torch.Tensor, - 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")( - qkv_out, - z_out, - b_out, - a_out, - qkvz_staging, - ba_staging, - dense_weight_ptr, - input, - codes, - group_scales, - ba_weight, - ) - - -if hasattr(torch.ops._C, "fp8_qpn8_dispatch_ba_split_sm70_out"): - - @register_fake("_C::fp8_qpn8_dispatch_ba_split_sm70_out") - def _fp8_qpn8_dispatch_ba_split_sm70_out_fake( - qkv_out: torch.Tensor, - z_out: torch.Tensor, - b_out: torch.Tensor, - a_out: torch.Tensor, - qkvz_staging: torch.Tensor, - ba_staging: torch.Tensor, - dense_weight_ptr: int, - input: torch.Tensor, - codes: torch.Tensor, - group_scales: torch.Tensor, - ba_weight: torch.Tensor, - ) -> None: - return None - - -def fp8_qpn8_gated_pair_sm70_out( - out: torch.Tensor, - input: torch.Tensor, - codes: torch.Tensor, - group_scales: torch.Tensor, - split_k: int, - accumulator_chains: int, - fast_decoder: bool, - prefetch_codes: bool = False, -) -> None: - """Run the single-kernel paired-tile QPN8 gated SiLU experiment.""" - _op("fp8_qpn8_gated_pair_sm70_out")( - out, - input, - codes, - group_scales, - split_k, - accumulator_chains, - fast_decoder, - prefetch_codes, - ) - - -if hasattr(torch.ops._C, "fp8_qpn8_gated_pair_sm70_out"): - - @register_fake("_C::fp8_qpn8_gated_pair_sm70_out") - def _fp8_qpn8_gated_pair_sm70_out_fake( - out: torch.Tensor, - input: torch.Tensor, - codes: torch.Tensor, - group_scales: torch.Tensor, - split_k: int, - accumulator_chains: int, - fast_decoder: bool, - prefetch_codes: bool, - ) -> None: - return None - - -def nvfp4_qpn4_prepare_sm70( - qweight: torch.Tensor, - scales: torch.Tensor, -) -> list[torch.Tensor]: - """Pack unpacked FP4 weights and FP16 group scales into QPN4 layout.""" - return _op("nvfp4_qpn4_prepare_sm70")(qweight, scales) - - -if hasattr(torch.ops._C, "nvfp4_qpn4_prepare_sm70"): - - @register_fake("_C::nvfp4_qpn4_prepare_sm70") - def _nvfp4_qpn4_prepare_sm70_fake( - qweight: torch.Tensor, - scales: torch.Tensor, - ) -> list[torch.Tensor]: - k, n = qweight.shape - codes = torch.empty((k, n // 2), dtype=torch.uint8, device=qweight.device) - packed_scales = torch.empty_like(scales) - return [codes, packed_scales] - - -def nvfp4_qpn4_prepare_scale_code_sm70( - qweight: torch.Tensor, - scale_codes: torch.Tensor, -) -> list[torch.Tensor]: - """Pack unpacked FP4 weights and checkpoint E4M3 group-scale codes.""" - return _op("nvfp4_qpn4_prepare_scale_code_sm70")(qweight, scale_codes) - - -if hasattr(torch.ops._C, "nvfp4_qpn4_prepare_scale_code_sm70"): - - @register_fake("_C::nvfp4_qpn4_prepare_scale_code_sm70") - def _nvfp4_qpn4_prepare_scale_code_sm70_fake( - qweight: torch.Tensor, - scale_codes: torch.Tensor, - ) -> list[torch.Tensor]: - k, n = qweight.shape - codes = torch.empty((k, n // 2), dtype=torch.uint8, device=qweight.device) - packed_scale_codes = torch.empty( - (k // 16, n), - dtype=torch.uint8, - device=scale_codes.device, - ) - return [codes, packed_scale_codes] - - -def nvfp4_qpn4_dequantize_sm70_out( - out: torch.Tensor, - codes: torch.Tensor, - scales: torch.Tensor, - global_scale: float, - use_scale_code: bool, -) -> None: - _op("nvfp4_qpn4_dequantize_sm70_out")( - out, codes, scales, global_scale, use_scale_code - ) - - -if hasattr(torch.ops._C, "nvfp4_qpn4_dequantize_sm70_out"): - - @register_fake("_C::nvfp4_qpn4_dequantize_sm70_out") - def _nvfp4_qpn4_dequantize_sm70_out_fake( - out: torch.Tensor, - codes: torch.Tensor, - scales: torch.Tensor, - global_scale: float, - use_scale_code: bool, - ) -> None: - return None - - -def nvfp4_qpn4_prefill_sm70_out( - out: torch.Tensor, - dense_weight_ptr: int, - input: torch.Tensor, - codes: torch.Tensor, - scales: torch.Tensor, - global_scale: float, - use_scale_code: bool, - gated_silu: bool, -) -> None: - _op("nvfp4_qpn4_prefill_sm70_out")( - out, - dense_weight_ptr, - input, - codes, - scales, - global_scale, - use_scale_code, - gated_silu, - ) - - -if hasattr(torch.ops._C, "nvfp4_qpn4_prefill_sm70_out"): - - @register_fake("_C::nvfp4_qpn4_prefill_sm70_out") - def _nvfp4_qpn4_prefill_sm70_out_fake( - out: torch.Tensor, - dense_weight_ptr: int, - input: torch.Tensor, - codes: torch.Tensor, - scales: torch.Tensor, - global_scale: float, - use_scale_code: bool, - gated_silu: bool, - ) -> None: - return None - - -def nvfp4_qpn4_dispatch_sm70_out( - out: torch.Tensor, - dense_weight_ptr: int, - input: torch.Tensor, - codes: torch.Tensor, - scales: torch.Tensor, - global_scale: float, - use_scale_code: bool, - gated_silu: bool, -) -> None: - """Dispatch exact M=1 decode or bounded-workspace large-M prefill.""" - _op("nvfp4_qpn4_dispatch_sm70_out")( - out, - dense_weight_ptr, - input, - codes, - scales, - global_scale, - use_scale_code, - gated_silu, - ) - - -if hasattr(torch.ops._C, "nvfp4_qpn4_dispatch_sm70_out"): - - @register_fake("_C::nvfp4_qpn4_dispatch_sm70_out") - def _nvfp4_qpn4_dispatch_sm70_out_fake( - out: torch.Tensor, - dense_weight_ptr: int, - input: torch.Tensor, - codes: torch.Tensor, - scales: torch.Tensor, - global_scale: float, - use_scale_code: bool, - gated_silu: bool, - ) -> None: - return None - - -def nvfp4_qpn2_prepare_sm70( - weight_packed: torch.Tensor, - weight_scale: torch.Tensor, -) -> list[torch.Tensor]: - """Pack checkpoint-native NVFP4 tensors into the QPN2 layout.""" - return _op("nvfp4_qpn2_prepare_sm70")(weight_packed, weight_scale) - - -if hasattr(torch.ops._C, "nvfp4_qpn2_prepare_sm70"): - - @register_fake("_C::nvfp4_qpn2_prepare_sm70") - def _nvfp4_qpn2_prepare_sm70_fake( - weight_packed: torch.Tensor, - weight_scale: torch.Tensor, - ) -> list[torch.Tensor]: - codes = torch.empty_like(weight_packed) - scales = torch.empty( - weight_scale.shape, dtype=torch.uint8, device=weight_scale.device - ) - return [codes, scales] - - -def nvfp4_qpn2_gemm_sm70_out( - out: torch.Tensor, - input: torch.Tensor, - codes: torch.Tensor, - scales: torch.Tensor, - global_scale: float, - split_k: int, - accumulator_chains: int, -) -> None: - _op("nvfp4_qpn2_gemm_sm70_out")( - out, - input, - codes, - scales, - global_scale, - split_k, - accumulator_chains, - ) - - -if hasattr(torch.ops._C, "nvfp4_qpn2_gemm_sm70_out"): - - @register_fake("_C::nvfp4_qpn2_gemm_sm70_out") - def _nvfp4_qpn2_gemm_sm70_out_fake( - out: torch.Tensor, - input: torch.Tensor, - codes: torch.Tensor, - scales: torch.Tensor, - global_scale: float, - split_k: int, - accumulator_chains: int, - ) -> None: - return None - - -def nvfp4_qpn2_gated_sm70_out( - out: torch.Tensor, - input: torch.Tensor, - codes: torch.Tensor, - scales: torch.Tensor, - global_scale: float, - split_k: int, - accumulator_chains: int, -) -> None: - _op("nvfp4_qpn2_gated_sm70_out")( - out, - input, - codes, - scales, - global_scale, - split_k, - accumulator_chains, - ) - - -if hasattr(torch.ops._C, "nvfp4_qpn2_gated_sm70_out"): - - @register_fake("_C::nvfp4_qpn2_gated_sm70_out") - def _nvfp4_qpn2_gated_sm70_out_fake( - out: torch.Tensor, - input: torch.Tensor, - codes: torch.Tensor, - scales: torch.Tensor, - global_scale: float, - split_k: int, - accumulator_chains: int, - ) -> None: - return None - - -def nvfp4_qpn2_dispatch_sm70_out( - out: torch.Tensor, - input: torch.Tensor, - codes: torch.Tensor, - scales: torch.Tensor, - global_scale: float, - split_k: int, - accumulator_chains: int, - tm_weight: torch.Tensor, - tm_scales: torch.Tensor, - tm_group_size: int, - tm_k_ld: int, - tm_q_ld: int, - gated_silu: bool, -) -> None: - """Select QPN2 for M<=8 and TurboMind for larger dynamic M.""" - _op("nvfp4_qpn2_dispatch_sm70_out")( - out, - input, - codes, - scales, - global_scale, - split_k, - accumulator_chains, - tm_weight, - tm_scales, - tm_group_size, - tm_k_ld, - tm_q_ld, - gated_silu, - ) - - -if hasattr(torch.ops._C, "nvfp4_qpn2_dispatch_sm70_out"): - - @register_fake("_C::nvfp4_qpn2_dispatch_sm70_out") - def _nvfp4_qpn2_dispatch_sm70_out_fake( - out: torch.Tensor, - input: torch.Tensor, - codes: torch.Tensor, - scales: torch.Tensor, - global_scale: float, - split_k: int, - accumulator_chains: int, - tm_weight: torch.Tensor, - tm_scales: torch.Tensor, - tm_group_size: int, - tm_k_ld: int, - tm_q_ld: int, - gated_silu: bool, - ) -> None: - return None - - def fp8_gemm_sm70_prefill_dispatch_out( out: torch.Tensor, dense_weight_ptr: int, @@ -1222,50 +541,6 @@ def _mxfp4_moe_dense_stage_sm70_out_fake( return None -def nvfp4_moe_dense_stage_sm70_out( - out: torch.Tensor, - input: torch.Tensor, - expert_offsets: torch.Tensor, - dense_expert_ids: torch.Tensor, - ptrs_w: torch.Tensor, - ptrs_s: torch.Tensor, - num_experts: int, - k: int, - n: int, - group_size: int, -) -> None: - _op("nvfp4_moe_dense_stage_sm70_out")( - out, - input, - expert_offsets, - dense_expert_ids, - ptrs_w, - ptrs_s, - num_experts, - k, - n, - group_size, - ) - - -if hasattr(torch.ops._C, "nvfp4_moe_dense_stage_sm70_out"): - - @register_fake("_C::nvfp4_moe_dense_stage_sm70_out") - def _nvfp4_moe_dense_stage_sm70_out_fake( - out: torch.Tensor, - input: torch.Tensor, - expert_offsets: torch.Tensor, - dense_expert_ids: torch.Tensor, - ptrs_w: torch.Tensor, - ptrs_s: torch.Tensor, - num_experts: int, - k: int, - n: int, - group_size: int, - ) -> None: - return None - - def mxfp4_moe_single_token_prepare_w13_sm70_out( gate_up: torch.Tensor, compact_input: torch.Tensor, @@ -1492,145 +767,6 @@ def _sm70_f16_gemm_out_fake( return None -def sm70_f16_indexed_rerank_out( - out: torch.Tensor, - input: torch.Tensor, - weight: torch.Tensor, - candidate_ids: torch.Tensor, - selected_raw: torch.Tensor, - selected_packed: torch.Tensor, - expanded: torch.Tensor, - partials: torch.Tensor, - barriers: torch.Tensor, - cta_n: int, - split_k: int, -) -> None: - _op("sm70_f16_indexed_rerank_out")( - out, - input, - weight, - candidate_ids, - selected_raw, - selected_packed, - expanded, - partials, - barriers, - cta_n, - split_k, - ) - - -if hasattr(torch.ops._C, "sm70_f16_indexed_rerank_out"): - - @register_fake("_C::sm70_f16_indexed_rerank_out") - def _sm70_f16_indexed_rerank_out_fake( - out: torch.Tensor, - input: torch.Tensor, - weight: torch.Tensor, - candidate_ids: torch.Tensor, - selected_raw: torch.Tensor, - selected_packed: torch.Tensor, - expanded: torch.Tensor, - partials: torch.Tensor, - barriers: torch.Tensor, - cta_n: int, - split_k: int, - ) -> None: - return None - - -def sm70_f16_indexed_rerank_packed_out( - out: torch.Tensor, - input: torch.Tensor, - packed_weight: torch.Tensor, - candidate_ids: torch.Tensor, - selected_packed: torch.Tensor, - expanded: torch.Tensor, - partials: torch.Tensor, - barriers: torch.Tensor, - cta_n: int, - split_k: int, -) -> None: - _op("sm70_f16_indexed_rerank_packed_out")( - out, - input, - packed_weight, - candidate_ids, - selected_packed, - expanded, - partials, - barriers, - cta_n, - split_k, - ) - - -if hasattr(torch.ops._C, "sm70_f16_indexed_rerank_packed_out"): - - @register_fake("_C::sm70_f16_indexed_rerank_packed_out") - def _sm70_f16_indexed_rerank_packed_out_fake( - out: torch.Tensor, - input: torch.Tensor, - packed_weight: torch.Tensor, - candidate_ids: torch.Tensor, - selected_packed: torch.Tensor, - expanded: torch.Tensor, - partials: torch.Tensor, - barriers: torch.Tensor, - cta_n: int, - split_k: int, - ) -> None: - return None - - -def sm70_f16_rerank_keys_out( - keys: torch.Tensor, - logits: torch.Tensor, - candidate_ids: torch.Tensor, -) -> None: - _op("sm70_f16_rerank_keys_out")(keys, logits, candidate_ids) - - -if hasattr(torch.ops._C, "sm70_f16_rerank_keys_out"): - - @register_fake("_C::sm70_f16_rerank_keys_out") - def _sm70_f16_rerank_keys_out_fake( - keys: torch.Tensor, - logits: torch.Tensor, - candidate_ids: torch.Tensor, - ) -> None: - return None - - -def sm70_f16_rerank_topk_out( - values_out: torch.Tensor, - ids_out: torch.Tensor, - logits: torch.Tensor, - candidate_ids: torch.Tensor, - vocab_start_index: int, -) -> None: - _op("sm70_f16_rerank_topk_out")( - values_out, - ids_out, - logits, - candidate_ids, - vocab_start_index, - ) - - -if hasattr(torch.ops._C, "sm70_f16_rerank_topk_out"): - - @register_fake("_C::sm70_f16_rerank_topk_out") - def _sm70_f16_rerank_topk_out_fake( - values_out: torch.Tensor, - ids_out: torch.Tensor, - logits: torch.Tensor, - candidate_ids: torch.Tensor, - vocab_start_index: int, - ) -> None: - return None - - def sm70_f16_lm_head_top1_out( values_out: torch.Tensor, indices_out: torch.Tensor, @@ -1803,114 +939,6 @@ def _sm70_sample_packed_top20_out_fake( return None -def sm70_sample_sorted_top20_philox_out( - sampled_token_out: torch.Tensor, - sparse_ids_out: torch.Tensor, - sparse_probs_out: torch.Tensor, - top_values: torch.Tensor, - top_indices: torch.Tensor, - generator: torch.Generator | None, - vocab_size: int, - top_p: float, -) -> None: - _op("sm70_sample_sorted_top20_philox_out")( - sampled_token_out, - sparse_ids_out, - sparse_probs_out, - top_values, - top_indices, - generator, - vocab_size, - top_p, - ) - - -if hasattr(torch.ops._C, "sm70_sample_sorted_top20_philox_out"): - - @register_fake("_C::sm70_sample_sorted_top20_philox_out") - def _sm70_sample_sorted_top20_philox_out_fake( - sampled_token_out: torch.Tensor, - sparse_ids_out: torch.Tensor, - sparse_probs_out: torch.Tensor, - top_values: torch.Tensor, - top_indices: torch.Tensor, - generator: torch.Generator | None, - vocab_size: int, - top_p: float, - ) -> None: - return None - - -def sm70_sample_sorted_top20_philox_token_out( - sampled_token_out: torch.Tensor, - top_values: torch.Tensor, - top_indices: torch.Tensor, - generator: torch.Generator | None, - vocab_size: int, - top_p: float, -) -> None: - _op("sm70_sample_sorted_top20_philox_token_out")( - sampled_token_out, - top_values, - top_indices, - generator, - vocab_size, - top_p, - ) - - -if hasattr(torch.ops._C, "sm70_sample_sorted_top20_philox_token_out"): - - @register_fake("_C::sm70_sample_sorted_top20_philox_token_out") - def _sm70_sample_sorted_top20_philox_token_out_fake( - sampled_token_out: torch.Tensor, - top_values: torch.Tensor, - top_indices: torch.Tensor, - generator: torch.Generator | None, - vocab_size: int, - top_p: float, - ) -> None: - return None - - -def sm70_sample_chunked_top20_philox_token_out( - sampled_token_out: torch.Tensor, - global_values: torch.Tensor, - local_indices: torch.Tensor, - global_positions: torch.Tensor, - generator: torch.Generator | None, - vocab_size: int, - top_p: float, - chunk_size: int, -) -> None: - _op("sm70_sample_chunked_top20_philox_token_out")( - sampled_token_out, - global_values, - local_indices, - global_positions, - generator, - vocab_size, - top_p, - chunk_size, - ) - - -if hasattr(torch.ops._C, "sm70_sample_chunked_top20_philox_token_out"): - - @register_fake("_C::sm70_sample_chunked_top20_philox_token_out") - def _sm70_sample_chunked_top20_philox_token_out_fake( - sampled_token_out: torch.Tensor, - global_values: torch.Tensor, - local_indices: torch.Tensor, - global_positions: torch.Tensor, - generator: torch.Generator | None, - vocab_size: int, - top_p: float, - chunk_size: int, - ) -> None: - return None - - def sm70_dynamic_draft_vocab_update_tail_out( lru_token_ids: torch.Tensor, local_tail_token_ids: torch.Tensor, diff --git a/vllm/compilation/cuda_graph.py b/vllm/compilation/cuda_graph.py index 5e80444996..5664cb7361 100644 --- a/vllm/compilation/cuda_graph.py +++ b/vllm/compilation/cuda_graph.py @@ -192,7 +192,8 @@ def __init__( self.first_run_finished = False self.is_debugging_mode = envs.VLLM_LOGGING_LEVEL == "DEBUG" self.check_input_addresses = ( - self.is_debugging_mode or envs.VLLM_CUDAGRAPH_INPUT_ADDR_DEBUG + self.is_debugging_mode + or envs.VLLM_CUDAGRAPH_INPUT_ADDR_DEBUG ) self._runnable_str = str(runnable) if self.is_debugging_mode else None @@ -210,7 +211,6 @@ def __init__( # the entries for different batch descriptors that we need to capture # cudagraphs for. self.concrete_cudagraph_entries: dict[BatchDescriptor, CUDAGraphEntry] = {} - self.semantic_variant_graph_pools: dict[int, Any] = {} CUDAGraphWrapper._all_instances.add(self) @@ -264,20 +264,6 @@ def _collect_tensor_debug( def clear_graphs(self) -> None: self.concrete_cudagraph_entries.clear() - self.semantic_variant_graph_pools.clear() - - def _graph_pool_for_descriptor(self, batch_descriptor: BatchDescriptor) -> Any: - if batch_descriptor.graph_variant >= 0: - return self.graph_pool - # Negative variants are mutually exclusive semantic paths for an - # otherwise identical graph shape. They can replay in arbitrary order, - # so do not let a later capture reuse the default graph's intermediates. - graph_variant = batch_descriptor.graph_variant - if graph_variant not in self.semantic_variant_graph_pools: - self.semantic_variant_graph_pools[graph_variant] = ( - current_platform.graph_pool_handle() - ) - return self.semantic_variant_graph_pools[graph_variant] def __call__(self, *args: Any, **kwargs: Any) -> Any | None: if not is_forward_context_available(): @@ -349,9 +335,8 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any | None: ) ) - graph_pool = self._graph_pool_for_descriptor(entry.batch_descriptor) - if graph_pool is not None: - set_graph_pool_id(graph_pool) + if self.graph_pool is not None: + set_graph_pool_id(self.graph_pool) else: set_graph_pool_id(current_platform.graph_pool_handle()) @@ -362,7 +347,7 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any | None: # mind-exploding: carefully manage the reference and memory. with torch.cuda.graph( cudagraph, - pool=graph_pool, + pool=self.graph_pool, stream=current_stream(), ): # `output` is managed by pytorch's cudagraph pool diff --git a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py index 7726c72238..aa0f68fa85 100644 --- a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py +++ b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py @@ -442,89 +442,6 @@ def _sm70_ar_gemma_rms_match(match: pm.Match) -> bool: return False -def _sm70_tp4_fusion_diagnostic(graph: fx.Graph) -> str: - """Summarize real post-grad users when the fail-closed pattern misses.""" - all_reduces: list[fx.Node] = [] - direct_users: dict[str, int] = {} - for node in graph.nodes: - if node.target != _ALL_REDUCE_OP: - continue - value = node.meta.get("val") - if value is None or value.dim() != 2 or value.shape[-1] != 5120: - continue - all_reduces.append(node) - for user in node.users: - target = str(user.target) - direct_users[target] = direct_users.get(target, 0) + 1 - - traces: list[str] = [] - for root in all_reduces: - queue: list[tuple[fx.Node, int]] = [(root, 0)] - seen = {root} - trace_nodes: list[tuple[fx.Node, int]] = [] - while queue and len(trace_nodes) < 32: - node, depth = queue.pop(0) - trace_nodes.append((node, depth)) - if depth == 10: - continue - for user in node.users: - if user not in seen: - seen.add(user) - queue.append((user, depth + 1)) - targets = " ".join(str(node.target) for node, _ in trace_nodes) - if "mean" not in targets and "rsqrt" not in targets and "rms" not in targets: - continue - lines = [] - for node, depth in trace_nodes: - value = node.meta.get("val") - shape = getattr(value, "shape", None) - dtype = getattr(value, "dtype", None) - lines.append( - f"{depth}:{node.name}:{node.target}:shape={shape}:dtype={dtype}:" - f"args={node.args}:kwargs={node.kwargs}" - ) - traces.append("\n".join(lines)) - if len(traces) == 4: - break - - return ( - f"all_reduce_f16_5120={len(all_reduces)}; " - f"direct_users={direct_users}; candidate_traces=\n" + "\n---\n".join(traces) - ) - - -def _sm70_tp4_expected_graph_matches(graph: fx.Graph) -> int: - """Count supported FP32-residual boundaries in one post-grad graph.""" - count = 0 - for node in graph.nodes: - if ( - node.target - != torch.ops.vllm.sm70_gemma_long_prefill_fused_add_rms_norm.default - ): - continue - reduced = node.args[0] - residual = node.args[1] - if not isinstance(reduced, fx.Node) or reduced.target != _ALL_REDUCE_OP: - continue - if not isinstance(residual, fx.Node): - continue - residual_value = residual.meta.get("val") - if residual_value is None or residual_value.dtype != torch.float32: - continue - input_node = reduced.args[0] - if not isinstance(input_node, fx.Node): - continue - value = input_node.meta.get("val") - if ( - value is not None - and value.dtype == torch.float16 - and value.dim() == 2 - and value.shape[-1] == 5120 - ): - count += 1 - return count - - class Sm70AllReduceGemmaRMSNormPattern(BasePattern): def __init__( self, @@ -596,73 +513,6 @@ def replacement( ) -class Sm70Tp4LongPrefillFusedNormPattern(BasePattern): - """Fuse TP4 AR plus the accepted opaque mixed-dtype Gemma norm.""" - - def __init__( - self, - epsilon: float, - dtype: torch.dtype, - device: str | None, - ) -> None: - super().__init__(dtype, device) - self.epsilon = epsilon - self.group_name = get_tp_group().unique_name - - def get_inputs(self) -> list[torch.Tensor]: - return [ - self.empty_f32(5, 5120), - self.empty(5, 5120), - self.empty(5120), - ] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - residual: torch.Tensor, - input: torch.Tensor, - weight: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - reduced = tensor_model_parallel_all_reduce(input) - return torch.ops.vllm.sm70_gemma_long_prefill_fused_add_rms_norm( - reduced, - residual, - weight, - self.epsilon, - ) - - def replacement( - residual: torch.Tensor, - input: torch.Tensor, - weight: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - return torch.ops.vllm.sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather( - input, - residual, - weight, - self.epsilon, - group_name=self.group_name, - ) - - pm.register_replacement( - pattern, - replacement, - self.get_inputs(), - pm.fwd_only, - pm_pass, - extra_check=_sm70_ar_gemma_rms_match, - ) - - first_return_only = lambda fn: lambda a, b, c: fn(a, b, c)[0] - pm.register_replacement( - first_return_only(pattern), # type: ignore[no-untyped-call] - first_return_only(replacement), # type: ignore[no-untyped-call] - self.get_inputs(), - pm.fwd_only, - pm_pass, - extra_check=_sm70_ar_gemma_rms_match, - ) - - class AllReduceFusedRMSNormStaticQuantFP8Pattern(BasePattern): """ This pattern replaces the allreduce + rms norm (without residual) @@ -1024,7 +874,6 @@ def __init__(self, config: VllmConfig) -> None: self.hidden_dim = config.model_config.get_hidden_size() self.group = get_tp_group().device_group rank = get_tensor_model_parallel_rank() - self.rank = rank sm70_common = ( self.hidden_dim == 5120 and self.model_dtype == torch.float16 @@ -1033,20 +882,7 @@ def __init__(self, config: VllmConfig) -> None: self.sm70_tp2_mode = ( envs.VLLM_SM70_TP2_AR_GEMMA_RMS_FUSION and self.tp_size == 2 and sm70_common ) - self.sm70_tp4_long_mode = ( - envs.VLLM_SM70_TP4_LONG_PREFILL_FUSED_NORM - and self.tp_size == 4 - and sm70_common - and config.parallel_config.pipeline_parallel_size == 1 - and config.speculative_config is None - ) - self.sm70_tp4_model_expected_patterns = ( - 2 * config.model_config.get_num_layers(config.parallel_config) - 1 - if self.sm70_tp4_long_mode - else 0 - ) - self.sm70_tp4_total_matches = 0 - self.sm70_mode = self.sm70_tp2_mode or self.sm70_tp4_long_mode + self.sm70_mode = self.sm70_tp2_mode if self.sm70_mode: device_comm = get_tp_group().device_communicator ca_comm = ( @@ -1058,25 +894,6 @@ def __init__(self, config: VllmConfig) -> None: self.tp_size, ) return - if self.sm70_tp4_long_mode and ca_comm.long_prefill_output_ptrs is None: - logger.warning_once( - "SM70 TP4 long-prefill collective-norm buffers are unavailable." - ) - return - if self.sm70_tp4_long_mode and not hasattr( - torch.ops.vllm, - "sm70_gemma_long_prefill_fused_add_rms_norm", - ): - logger.warning_once( - "SM70 long-prefill Gemma norm op is unavailable for TP4 fusion." - ) - return - if self.sm70_tp4_long_mode and not hasattr( - torch.ops.vllm, - "sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather", - ): - logger.warning_once("SM70 TP4 fused collective-norm op is unavailable.") - return self.max_token_num = config.scheduler_config.max_num_batched_tokens self.register_sm70_patterns() self.dump_patterns(config, self.patterns) @@ -1151,14 +968,6 @@ def __init__(self, config: VllmConfig) -> None: @enable_fake_mode def register_sm70_patterns(self) -> None: - if self.sm70_tp4_long_mode: - Sm70Tp4LongPrefillFusedNormPattern( - 1e-6, - self.model_dtype, - self.device, - ).register(self.patterns) - self.disabled = False - return for residual_dtype in (torch.float16, torch.float32): Sm70AllReduceGemmaRMSNormPattern( 1e-6, @@ -1220,12 +1029,6 @@ def is_applicable_for_range(self, compile_range: Range) -> bool: if self.disabled: logger.warning_once("AllReduce fusion pass is disabled.") return False - if getattr(self, "sm70_tp4_long_mode", False): - return bool( - compile_range.is_single_size() - and 8000 <= compile_range.start <= min(8192, self.max_token_num) - and compile_range.start % self.tp_size == 0 - ) return bool(compile_range.end <= self.max_token_num) @VllmInductorPass.time_and_log @@ -1234,36 +1037,7 @@ def __call__(self, graph: fx.Graph) -> None: logger.debug("AllReduceFusionPass disabled") return - expected_graph_matches = ( - _sm70_tp4_expected_graph_matches(graph) if self.sm70_tp4_long_mode else 0 - ) self.matched_count = self.patterns.apply(graph) - if self.sm70_tp4_long_mode and self.matched_count != expected_graph_matches: - if self.rank == 0: - logger.error( - "SM70 TP4 fusion post-grad diagnostic:\n%s", - _sm70_tp4_fusion_diagnostic(graph), - ) - raise RuntimeError( - "SM70 TP4 long-prefill fused collective-norm requires a " - "complete residual chain in every piecewise graph: expected " - f"{expected_graph_matches} matches, got " - f"{self.matched_count}." - ) - if self.sm70_tp4_long_mode: - self.sm70_tp4_total_matches += self.matched_count - if self.sm70_tp4_total_matches > self.sm70_tp4_model_expected_patterns: - raise RuntimeError( - "SM70 TP4 long-prefill fused collective-norm exceeded the " - "model residual-chain count: expected at most " - f"{self.sm70_tp4_model_expected_patterns}, got " - f"{self.sm70_tp4_total_matches}." - ) - if self.sm70_tp4_total_matches == self.sm70_tp4_model_expected_patterns: - logger.info_once( - "SM70 TP4 fused the complete %d-boundary model residual chain.", - self.sm70_tp4_total_matches, - ) logger.debug("Replaced %s patterns", self.matched_count) if getattr(self, "sm70_mode", False): logger.info( diff --git a/vllm/compilation/passes/fusion/matcher_utils.py b/vllm/compilation/passes/fusion/matcher_utils.py index 97ce957f00..08dfbd4937 100644 --- a/vllm/compilation/passes/fusion/matcher_utils.py +++ b/vllm/compilation/passes/fusion/matcher_utils.py @@ -44,7 +44,8 @@ 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 9fef59b00b..01f5d5a210 100644 --- a/vllm/compilation/passes/fusion/qk_norm_rope_fusion.py +++ b/vllm/compilation/passes/fusion/qk_norm_rope_fusion.py @@ -198,7 +198,8 @@ 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 d5de88ba2c..b2c465b1ff 100644 --- a/vllm/compilation/passes/fusion/rms_quant_fusion.py +++ b/vllm/compilation/passes/fusion/rms_quant_fusion.py @@ -124,9 +124,12 @@ 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/fusion/sequence_parallelism.py b/vllm/compilation/passes/fusion/sequence_parallelism.py index be8bff058b..8d0f40e2c7 100644 --- a/vllm/compilation/passes/fusion/sequence_parallelism.py +++ b/vllm/compilation/passes/fusion/sequence_parallelism.py @@ -256,71 +256,6 @@ def replacement( ) -class MiddleAllReduceSm70GemmaLongPrefillPattern(_SequenceParallelPatternHelper): - """Keep the accepted mixed-dtype SM70 norm while applying SP.""" - - def __init__( - self, - epsilon: float, - dtype: torch.dtype, - device: str | None, - weight_dtype: torch.dtype, - ) -> None: - super().__init__(epsilon, dtype, device) - self.weight_dtype = weight_dtype - - def get_inputs(self) -> list[torch.Tensor]: - residual = torch.empty([4, 4], device=self.device, dtype=torch.float32) - mm_1 = torch.empty([4, 4], device=self.device, dtype=self.dtype) - weight = torch.empty([4], device=self.device, dtype=self.weight_dtype) - return [residual, mm_1, weight] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - residual: torch.Tensor, - mm_1: torch.Tensor, - weight: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - all_reduce = self._all_reduce(mm_1) - return torch.ops.vllm.sm70_gemma_long_prefill_fused_add_rms_norm( - all_reduce, - residual, - weight, - self.epsilon, - ) - - def replacement( - residual: torch.Tensor, - mm_1: torch.Tensor, - weight: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - reduce_scatter = self._reduce_scatter(mm_1) - local_len = reduce_scatter.size(0) - residual = residual[ - self.tp_rank * local_len : self.tp_rank * local_len + local_len, ... - ] - normalized, residual_out = ( - torch.ops.vllm.sm70_gemma_long_prefill_fused_add_rms_norm( - reduce_scatter, - residual, - weight, - self.epsilon, - ) - ) - return self._all_gather(normalized), residual_out - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - pm.register_replacement( - get_first_out_wrapper(pattern), - get_first_out_wrapper(replacement), - self.get_inputs(), - pm.fwd_only, - pm_pass, - ) - - class FirstAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper): def __init__( self, @@ -648,16 +583,6 @@ def __init__(self, config: VllmConfig) -> None: MiddleAllReduceRMSNormPattern( epsilon, self.model_dtype, self.device ).register(self.patterns) - if epsilon == 1e-6 and hasattr( - torch.ops.vllm, - "sm70_gemma_long_prefill_fused_add_rms_norm", - ): - MiddleAllReduceSm70GemmaLongPrefillPattern( - epsilon, - self.model_dtype, - self.device, - self.model_dtype, - ).register(self.patterns) self.dump_patterns(config, self.patterns) diff --git a/vllm/compilation/passes/utility/fix_functionalization.py b/vllm/compilation/passes/utility/fix_functionalization.py index 3e8c3af1c0..3f955d53cf 100644 --- a/vllm/compilation/passes/utility/fix_functionalization.py +++ b/vllm/compilation/passes/utility/fix_functionalization.py @@ -78,7 +78,9 @@ 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/attention.py b/vllm/config/attention.py index 487cafe7fb..52ce9f102a 100644 --- a/vllm/config/attention.py +++ b/vllm/config/attention.py @@ -3,7 +3,7 @@ from typing import Any, Literal -from pydantic import Field, field_validator +from pydantic import field_validator from vllm.config.utils import config from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum @@ -53,12 +53,6 @@ class AttentionConfig: use_non_causal: bool = False """Whether to use non-causal (bidirectional) attention.""" - prefix_anchored_decode_window: int | None = Field(default=None, gt=0, le=2**31 - 1) - """Bound generated-token KV to this sliding window while keeping the - request prefix globally visible. This is an explicit, model-agnostic - engine option. It is currently supported only by the SM70 - ``FLASH_ATTN_V100`` fp16 causal path. ``None`` disables it.""" - flex_attn_block_m: int | None = None """Triton kernel BLOCK_M tile size for flex attention. Must be a power of 2 >= 16. If None and VLLM_BATCH_INVARIANT=1, diff --git a/vllm/config/offload.py b/vllm/config/offload.py index 75aa758536..ad65e8acf3 100644 --- a/vllm/config/offload.py +++ b/vllm/config/offload.py @@ -40,9 +40,6 @@ class UVAOffloadConfig: - For parameter name "mlp.experts.w2_weight": - "experts" or "experts.w2_weight" will match. - "expert" or "w2" will NOT match (must be exact segments). - - A tower component name such as "visual" selects that complete - multimodal tower, including parameters whose names are relative to - the tower module. This allows distinguishing parameters like "w2_weight" and "w2_weight_scale". """ diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index c74aeb882e..27b1de7f73 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -2,14 +2,12 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import copy -import math from typing import TYPE_CHECKING, Any, Literal, get_args from pydantic import Field, SkipValidation, field_validator, model_validator from typing_extensions import Self from vllm.config import LoadConfig -from vllm.config.cache import CacheDType from vllm.config.kernel import MoEBackend from vllm.config.model import ModelConfig from vllm.config.parallel import ParallelConfig @@ -54,8 +52,7 @@ "gemma4_mtp", ] NgramGPUTypes = Literal["ngram_gpu"] -DFlashModelTypes = Literal["dflash"] -DFlashDDTreeModelTypes = Literal["dflash_ddtree"] +DFlashModelTypes = Literal["dflash", "dflash_ddtree"] DSparkModelTypes = Literal["dspark"] DDTreeBuildMode = Literal["best_first", "root_leaf", "spine_leaf"] EagleModelTypes = Literal[ @@ -64,7 +61,6 @@ "extract_hidden_states", MTPModelTypes, DFlashModelTypes, - DFlashDDTreeModelTypes, DSparkModelTypes, ] SpeculativeMethod = Literal[ @@ -123,10 +119,6 @@ class SpeculativeConfig: """Attention backend to use for the draft model. When `None`, the backend is automatically selected. Useful when the drafter requires a different attention backend (e.g. DFlash needs a non-causal-capable backend like FLASH_ATTN).""" - kv_cache_dtype: CacheDType | None = None - """KV cache dtype for the draft model. When ``None``, the draft inherits - the target cache dtype, except for SM70 DFlash with a quantized target, - which is resolved to ``auto`` during draft model loading.""" max_model_len: int | None = Field(default=None, ge=1) """The maximum model length of the draft model. Used when testing the ability to skip speculation for some sequences.""" @@ -158,11 +150,6 @@ class SpeculativeConfig: prompt_lookup_min: int | None = Field(default=None, ge=1) """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.""" # Alternative drafting strategies parallel_drafting: bool = False @@ -299,20 +286,6 @@ def _resolve_synthetic_acceptance_rates( during rejection sampling. This comes at the cost of additional GPU memory usage.""" - dspark_confidence_threshold: float = Field(default=0.0, ge=0.0, le=1.0) - """Conditional-acceptance threshold for DSpark prefix scheduling. A value - of zero keeps every proposed position. Positive values retain only the - leading positions whose calibrated confidence is at least the threshold.""" - - dspark_confidence_temperatures: list[float] | None = None - """Optional positive per-position temperatures applied to DSpark confidence - logits before prefix scheduling. The list must match num_speculative_tokens.""" - - dspark_max_verification_tokens: int | None = Field(default=None, ge=0) - """Optional cap on DSpark draft tokens submitted to target verification. - DSpark still generates the checkpoint's complete block; only a prefix is - scheduled, so values below the checkpoint block size remain lossless.""" - def compute_hash(self) -> str: """ WARNING: Whenever a new field is added to this config, @@ -334,7 +307,7 @@ def compute_hash(self) -> str: "eagle3", "extract_hidden_states", ) - or self.use_dflash_family() + or self.use_dflash() or self.use_dspark() ) factors.append(uses_aux_hidden_states) @@ -670,22 +643,6 @@ def __post_init__(self): if self.method in ("ngram", "[ngram]"): self.method = "ngram" - if self.ngram_assist: - if self.prompt_lookup_min is None and self.prompt_lookup_max is None: - self.prompt_lookup_min = 5 - self.prompt_lookup_max = 5 - elif self.prompt_lookup_min is None: - self.prompt_lookup_min = self.prompt_lookup_max - elif self.prompt_lookup_max is None: - self.prompt_lookup_max = self.prompt_lookup_min - assert self.prompt_lookup_min is not None - assert self.prompt_lookup_max is not None - if self.prompt_lookup_min > self.prompt_lookup_max: - raise ValueError( - f"prompt_lookup_min={self.prompt_lookup_min} must " - f"be <= prompt_lookup_max={self.prompt_lookup_max}" - ) - if self.method in ("ngram", "ngram_gpu"): # Set default values if not provided if self.prompt_lookup_min is None and self.prompt_lookup_max is None: @@ -762,9 +719,8 @@ def __post_init__(self): self.draft_parallel_config = self.target_parallel_config else: - if not self.ngram_assist: - self.prompt_lookup_max = 0 - self.prompt_lookup_min = 0 + self.prompt_lookup_max = 0 + self.prompt_lookup_min = 0 if self.model is not None: self.draft_model_config = ModelConfig( @@ -792,7 +748,7 @@ def __post_init__(self): # Automatically detect the method if ( self.method in ("eagle", "eagle3") - or self.use_dflash_family() + or self.use_dflash() or self.use_dspark() ): pass @@ -832,7 +788,7 @@ def __post_init__(self): ) # Replace hf_config for EAGLE draft_model - if self.method in ("eagle", "eagle3") or self.use_dflash_family(): + if self.method in ("eagle", "eagle3") or self.use_dflash(): from vllm.transformers_utils.configs.eagle import EAGLEConfig from vllm.transformers_utils.configs.speculators import ( SpeculatorsConfig, @@ -844,9 +800,7 @@ def __post_init__(self): ): pass else: - eagle_method = ( - "dflash" if self.use_dflash_family() else self.method - ) + eagle_method = "dflash" if self.use_dflash() else self.method eagle_config = EAGLEConfig( self.draft_model_config.hf_config, method=eagle_method, @@ -861,7 +815,7 @@ def __post_init__(self): draft_hf_config.architectures = ["DSparkDraftModel"] self.update_arch_() - if self.use_dflash_family() or self.use_dspark(): + if self.use_dflash() or self.use_dspark(): self.parallel_drafting = True if self.num_speculative_tokens is not None and hasattr( @@ -910,32 +864,6 @@ def __post_init__(self): f"{self.num_speculative_tokens}. Smaller values " "produce incorrect output." ) - if ( - self.dspark_max_verification_tokens is not None - and self.dspark_max_verification_tokens - > self.num_speculative_tokens - ): - raise ValueError( - "dspark_max_verification_tokens must be <= " - f"num_speculative_tokens ({self.num_speculative_tokens}); " - f"got {self.dspark_max_verification_tokens}." - ) - temperatures = self.dspark_confidence_temperatures - if temperatures is not None: - if len(temperatures) != self.num_speculative_tokens: - raise ValueError( - "dspark_confidence_temperatures must have length " - f"{self.num_speculative_tokens}; got {temperatures}." - ) - if not all( - math.isfinite(value) and value > 0.0 - for value in temperatures - ): - raise ValueError( - "dspark_confidence_temperatures entries must be " - f"finite and positive; got {temperatures}." - ) - self._verify_dspark_final_stage_ownership() if self.use_dflash_ddtree() and self.ddtree_budget is None: self.ddtree_budget = self.num_speculative_tokens @@ -958,48 +886,16 @@ def __post_init__(self): self.draft_parallel_config = ( SpeculativeConfig.create_draft_parallel_config( - self.target_parallel_config, - self.draft_tensor_parallel_size, - # The runner constructs the complete DSpark drafter on - # the target model's final PP rank. It is local to that - # stage rather than partitioned across target PP ranks. - speculative_draft_pipeline_parallel_size=( - 1 if self.use_dspark() else None - ), + self.target_parallel_config, self.draft_tensor_parallel_size ) ) return self - def _verify_dspark_final_stage_ownership(self) -> None: - """Validate the layer contract for a final-stage-local drafter.""" - from vllm.distributed.utils import get_pp_indices - - pp_size = self.target_parallel_config.pipeline_parallel_size - num_layers = self.target_model_config.get_total_num_hidden_layers() - last_start, last_end = get_pp_indices(num_layers, pp_size - 1, pp_size) - layer_ids = tuple( - getattr( - self.draft_model_config.hf_config, - "dspark_target_layer_ids", - (), - ) - ) - outside_last_stage = [ - layer_id for layer_id in layer_ids if not last_start <= layer_id < last_end - ] - if not layer_ids or outside_last_stage: - raise ValueError( - "DSpark's final-stage-local drafter requires every " - "dspark_target_layer_id to belong to the final target pipeline " - f"stage [{last_start}, {last_end}); got {layer_ids}. Adjust " - "VLLM_PP_LAYER_PARTITION or the DSpark layer contract." - ) - 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.2.0`." + "Install via `pip install arctic-inference==0.1.1`." ) if self.num_speculative_tokens is None: # Suffix decoding decides the actual number of speculative tokens @@ -1133,19 +1029,13 @@ def update_arch_(self): def create_draft_parallel_config( target_parallel_config: ParallelConfig, speculative_draft_tensor_parallel_size: int, - speculative_draft_pipeline_parallel_size: int | None = None, ) -> ParallelConfig: """Create a parallel config for use by the draft worker. - This is mostly a copy of the target parallel config, except the draft - tensor size and, for final-stage-local drafters, pipeline size. + This is mostly a copy of the target parallel config, except the tp_size. """ draft_parallel_config = ParallelConfig( - pipeline_parallel_size=( - target_parallel_config.pipeline_parallel_size - if speculative_draft_pipeline_parallel_size is None - else speculative_draft_pipeline_parallel_size - ), + pipeline_parallel_size=target_parallel_config.pipeline_parallel_size, tensor_parallel_size=speculative_draft_tensor_parallel_size, distributed_executor_backend=target_parallel_config.distributed_executor_backend, max_parallel_loading_workers=target_parallel_config.max_parallel_loading_workers, @@ -1185,13 +1075,6 @@ def _verify_args(self) -> Self: "Expected num_speculative_tokens to be greater " f"than zero ({self.num_speculative_tokens})." ) - if self.ngram_assist: - if not self.use_dflash(): - raise ValueError("ngram_assist is only supported with method='dflash'.") - draft_hf_config = getattr(self.draft_model_config, "hf_config", None) - dflash_config = getattr(draft_hf_config, "dflash_config", None) or {} - if int(dflash_config.get("selector_top_k", 0) or 0) <= 0: - raise ValueError("ngram_assist requires DFlash2 selector capability.") if self.use_dflash_ddtree(): if self.ddtree_budget is None: self.ddtree_budget = self.num_speculative_tokens @@ -1252,12 +1135,6 @@ def max_num_new_slots_for_drafting(self) -> int: Calculate the maximum number of new slots that might be added to the batch when drafting. """ - if self.use_dflash(): - # MRV2 DFlash uses one bonus query followed by K mask queries. The - # scheduler already owns the target query slot, so K more slots - # must be available to the drafter. - return self.num_speculative_tokens - slots_per_req = 0 # for serial non-draft-model methods, no change needed if self.parallel_drafting: # For parallel drafting, we need one new slot per 'masked' token @@ -1287,28 +1164,16 @@ def use_step3p5_mtp(self) -> bool: def use_eagle(self) -> bool: return ( self.method in ("eagle", "eagle3", "mtp") - or self.use_dflash_family() + or self.use_dflash() or self.use_dspark() ) - def use_eagle_kv_cache(self) -> bool: - """Whether prefix hits need Eagle's final target-block recompute.""" - # MRV2 DFlash caches its projected draft context KV alongside the - # target KV. It only needs target hidden states for the uncached suffix, - # so dropping an additional target block is both unnecessary and can - # eliminate all prefix hits when hybrid page unification makes blocks - # large. Keep DDTree and DSpark on their established V1 semantics. - return self.use_eagle() and not self.use_dflash() - def use_dflash(self) -> bool: - return self.method == "dflash" + return self.method in get_args(DFlashModelTypes) def use_dflash_ddtree(self) -> bool: return self.method == "dflash_ddtree" - def use_dflash_family(self) -> bool: - return self.use_dflash() or self.use_dflash_ddtree() - def use_dspark(self) -> bool: return self.method == "dspark" diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 961b979c12..ea67803abb 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -68,7 +68,6 @@ DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset({"Qwen3ForCausalLM"}) _SM70_NOMTP_CUDAGRAPH_CAPTURE_SIZES = (1, 2, 4, 8, 16) -_SM70_MTP_CUDAGRAPH_REQUEST_SIZES = (1, 2, 4, 6, 8, 12, 16) def _sm70_nomtp_cudagraph_capture_sizes(max_num_seqs: int) -> list[int]: @@ -80,19 +79,6 @@ def _sm70_nomtp_cudagraph_capture_sizes(max_num_seqs: int) -> list[int]: return sorted(capture_sizes) -def _sm70_mtp_cudagraph_capture_sizes( - max_num_seqs: int, - decode_query_len: int, -) -> list[int]: - """Return exact SM70 MTP verifier token shapes for production requests.""" - max_graph_reqs = min(max(int(max_num_seqs), 1), 16) - request_sizes = { - size for size in _SM70_MTP_CUDAGRAPH_REQUEST_SIZES if size <= max_graph_reqs - } - request_sizes.add(max_graph_reqs) - return [decode_query_len * size for size in sorted(request_sizes)] - - class OptimizationLevel(IntEnum): """Optimization level enum.""" @@ -230,81 +216,6 @@ def enable_mla_dual_rms_norm_fusion(cfg: "VllmConfig") -> bool: return rocm_aiter_ops.is_enabled() and check_aiter_fused_qk_rmsnorm() -def apply_prefix_anchored_swa_constraints(cfg: "VllmConfig") -> None: - """Validate the model-agnostic prefix-anchored SWA engine contract.""" - window = cfg.attention_config.prefix_anchored_decode_window - if window is None: - return - - from vllm.platforms import current_platform - from vllm.v1.attention.backends.registry import AttentionBackendEnum - - if not ( - current_platform.is_cuda() and current_platform.is_device_capability((7, 0)) - ): - raise ValueError("prefix_anchored_decode_window requires an NVIDIA SM70 GPU") - - model_config = cfg.model_config - if model_config is None: - raise ValueError("prefix_anchored_decode_window requires a model config") - if model_config.use_mla: - raise ValueError("prefix_anchored_decode_window does not support MLA attention") - if cfg.attention_config.use_non_causal: - raise ValueError( - "prefix_anchored_decode_window requires causal decoder attention" - ) - - cache_dtype = cfg.cache_config.cache_dtype - effective_cache_dtype = model_config.dtype if cache_dtype == "auto" else cache_dtype - if effective_cache_dtype not in (torch.float16, "float16"): - raise ValueError( - "prefix_anchored_decode_window requires an fp16 KV cache; " - f"got cache_dtype={cache_dtype!r} and model dtype={model_config.dtype}" - ) - - parallel_config = cfg.parallel_config - if ( - parallel_config.decode_context_parallel_size != 1 - or parallel_config.prefill_context_parallel_size != 1 - ): - raise ValueError( - "prefix_anchored_decode_window does not support decode or prefill " - "context parallelism" - ) - if cfg.cache_config.kv_offloading_size is not None: - raise ValueError( - "prefix_anchored_decode_window does not support KV-cache offloading" - ) - if ( - cfg.kv_transfer_config is not None - and cfg.kv_transfer_config.kv_connector is not None - ): - raise ValueError("prefix_anchored_decode_window does not support KV connectors") - if cfg.speculative_config is not None: - raise ValueError( - "prefix_anchored_decode_window does not yet support speculative decoding" - ) - - backend = cfg.attention_config.backend - if backend is None: - cfg.attention_config.backend = AttentionBackendEnum.FLASH_ATTN_V100 - logger.info( - "Prefix-anchored sliding-window attention: auto-selecting the " - "FLASH_ATTN_V100 backend." - ) - elif backend != AttentionBackendEnum.FLASH_ATTN_V100: - raise ValueError( - "prefix_anchored_decode_window requires FLASH_ATTN_V100; " - f"got {backend.name}" - ) - if cfg.cache_config.enable_prefix_caching: - cfg.cache_config.enable_prefix_caching = False - logger.info( - "Prefix-anchored sliding-window attention: disabling prefix " - "caching (windowed decode KV is not reusable across requests)." - ) - - OPTIMIZATION_LEVEL_00 = { "compilation_config": { "pass_config": { @@ -600,28 +511,6 @@ def compute_hash(self) -> str: ] return hash_str - @property - def max_concurrent_batches(self) -> int: - # Scheduler-side upper bound mirroring the executors' - # `max_concurrent_batches` (multiproc: queue depth under async - # scheduling, else the PP pipeline depth). - pp_size = self.parallel_config.pipeline_parallel_size - if pp_size <= 1 and self.scheduler_config.async_scheduling: - return max(2, envs.VLLM_SM70_ASYNC_SCHEDULING_QUEUE_DEPTH) - return pp_size - - @property - def max_in_flight_tokens(self) -> int: - # Upper bound on tokens that are scheduled but not yet settled (freed): - # every concurrent batch may hold up to a full `max_num_batched_tokens`. - # Recycling-aware KV cache specs (sliding-window, chunked-local, - # prefix-anchored SWA) reserve for this because out-of-window blocks - # are freed on the processed-token basis, so in-flight steps - # transiently keep their blocks. - return ( - self.max_concurrent_batches * self.scheduler_config.max_num_batched_tokens - ) - @property def num_speculative_tokens(self) -> int: if ( @@ -631,39 +520,12 @@ def num_speculative_tokens(self) -> int: return self.speculative_config.num_speculative_tokens return 0 - @property - def num_lookahead_tokens(self) -> int: - """KV positions reserved beyond the target model's scheduled query.""" - speculative_config = self.speculative_config - if speculative_config is None: - return 0 - if speculative_config.use_dflash_family(): - # DFlash has one anchor query plus the checkpoint-owned draft state. - return speculative_config.num_speculative_state_tokens() + 1 - if speculative_config.use_eagle() or speculative_config.uses_draft_model(): - return self.num_speculative_tokens - return 0 - @property def use_v2_model_runner(self) -> bool: use_v2_model_runner = envs.VLLM_USE_V2_MODEL_RUNNER - is_mrv2_dflash = ( - self.speculative_config is not None and self.speculative_config.use_dflash() - ) if use_v2_model_runner is not None: - if is_mrv2_dflash and not use_v2_model_runner: - raise ValueError( - "method='dflash' is implemented only by Model Runner V2. " - "Use method='dflash_ddtree' for the retained V1 DDTree route." - ) return use_v2_model_runner - # The old V1 flat-DFlash route is intentionally removed. Force every - # method='dflash' checkpoint through the MRV2 speculator and fail in - # _validate_v2_model_runner if another configured feature is unsupported. - if is_mrv2_dflash: - return True - if not self._is_default_v2_model_runner_model(): return False @@ -964,8 +826,8 @@ def _verify_kv_transfer_compat(self) -> None: # producing RDMA failures like IBV_WC_REM_ACCESS_ERR / # NIXL_ERR_REMOTE_DISCONNECT at the first inter-node KV transfer. # We can't enumerate every in-tree and out-of-tree connector that - # pins memory, so we conservatively reject the combination unless the - # connector explicitly opts into the VMM-safe transfer contract. + # pins memory, so we conservatively reject the combination whenever + # any KV connector is configured. # # CuMem allocator is exempt: CuMemAllocator.use_memory_pool toggles # expandable_segments off around its pool (see #40812), so the KV @@ -977,14 +839,6 @@ def _verify_kv_transfer_compat(self) -> None: return if self.model_config is not None and (self.model_config.enable_cumem_allocator): return - if self._connector_supports_vmm_safe_transfers(): - logger.warning_once( - "Allowing KV connector %s with " - "PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True because " - "it implements SupportsVmmSafeTransfers.", - self.kv_transfer_config.kv_connector, - ) - return raise ValueError( f"KV connector {self.kv_transfer_config.kv_connector} is " @@ -999,25 +853,6 @@ def _verify_kv_transfer_compat(self) -> None: "expandable_segments is automatically disabled)." ) - def _connector_supports_vmm_safe_transfers(self) -> bool: - """Resolve the connector and check its explicit VMM-safety marker.""" - from vllm.distributed.kv_transfer.kv_connector.factory import ( - KVConnectorFactory, - ) - from vllm.distributed.kv_transfer.kv_connector.v1 import ( - supports_vmm_safe_transfers, - ) - - kv_transfer_config = self.kv_transfer_config - if kv_transfer_config is None: - return False - - try: - connector_cls = KVConnectorFactory.get_connector_class(kv_transfer_config) - return supports_vmm_safe_transfers(connector_cls) - except (ValueError, AttributeError, ImportError, TypeError): - return False - def __post_init__(self): """Verify configs are valid & consistent with each other.""" @@ -1062,8 +897,6 @@ def __post_init__(self): if self.lora_config is not None: self.lora_config.verify_with_model_config(self.model_config) - apply_prefix_anchored_swa_constraints(self) - if ( self.mamba_config.enable_stochastic_rounding and self.cache_config.mamba_ssm_cache_dtype != "float16" @@ -1474,6 +1307,7 @@ def __post_init__(self): env_name, env_value, ) + sm70_flash_0dot3_compile_graph = envs.VLLM_SM70_FLASH_V100_0DOT3_COMPILE_GRAPH sm70_flash_no_compile_graph = ( envs.VLLM_SM70_FLASH_V100_DECODE_GRAPH_NO_COMPILE @@ -1507,6 +1341,11 @@ def __post_init__(self): self.speculative_config is not None and self.speculative_config.num_speculative_tokens ): + cudagraph_capture_sizes = ( + [1, 2, 4, 8, 9, 18] + if self.parallel_config.tensor_parallel_size >= 4 + else [1, 2, 4, 8, 9] + ) decode_query_len = ( self.speculative_config.num_speculative_state_tokens() + 1 ) @@ -1524,47 +1363,26 @@ def __post_init__(self): smallq_env, decode_query_len, ) - if ( - envs.VLLM_SM70_MTP_SPLIT_DRAFT_CUDAGRAPHS - and self.speculative_config.method == "mtp" - ): - cudagraph_capture_sizes = _sm70_mtp_cudagraph_capture_sizes( - self.scheduler_config.max_num_seqs, - decode_query_len, - ) - logger.info_once( - "Using split SM70 MTP verifier cudagraph token " - "shapes %s for Flash-V100 compile graph.", - tuple(cudagraph_capture_sizes), - ) - else: - cudagraph_capture_sizes = ( - [1, 2, 4, 8, 9, 18] - if self.parallel_config.tensor_parallel_size >= 4 - else [1, 2, 4, 8, 9] - ) - max_graph_reqs = ( - 4 - if self.parallel_config.tensor_parallel_size >= 4 - else 1 - ) - max_graph_reqs = min( - max(int(self.scheduler_config.max_num_seqs), 1), - max_graph_reqs, - ) - cudagraph_capture_sizes = sorted( - set(cudagraph_capture_sizes) - | { - decode_query_len * num_reqs - for num_reqs in range(1, max_graph_reqs + 1) - } - ) - logger.info_once( - "Using SM70 speculative verifier cudagraph shapes " - "%sx1..%s for Flash-V100 compile graph.", - decode_query_len, - max_graph_reqs, - ) + max_graph_reqs = ( + 4 if self.parallel_config.tensor_parallel_size >= 4 else 1 + ) + max_graph_reqs = min( + max(int(self.scheduler_config.max_num_seqs), 1), + max_graph_reqs, + ) + cudagraph_capture_sizes = sorted( + set(cudagraph_capture_sizes) + | { + decode_query_len * num_reqs + for num_reqs in range(1, max_graph_reqs + 1) + } + ) + logger.info_once( + "Using SM70 speculative verifier cudagraph shapes " + "%sx1..%s for Flash-V100 compile graph.", + decode_query_len, + max_graph_reqs, + ) elif cudagraph_capture_sizes != [1, 2]: logger.info_once( "Using SM70 no-MTP decode cudagraph request shapes %s.", @@ -1871,22 +1689,6 @@ def enable_quant_fp8_custom_op_for_blocked_weights() -> bool: CUDAGraphMode.FULL_DECODE_ONLY ) - # Prefix-anchored sliding-window attention: full-graph capture - # bakes attention metadata that carries no per-request anchor - # lengths, so the anchored decode-window mask cannot run inside a - # full cudagraph. Piecewise graphs keep attention outside capture. - if ( - self.attention_config.prefix_anchored_decode_window is not None - and self.compilation_config.cudagraph_mode.has_full_cudagraphs() - ): - logger.info( - "Prefix-anchored sliding-window attention does not " - "support full cudagraphs. Overriding cudagraph_mode " - "from %s to PIECEWISE.", - self.compilation_config.cudagraph_mode.name, - ) - self.compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE - # Check if KV connector requires PIECEWISE mode for CUDA graphs if ( self.kv_transfer_config is not None @@ -2653,6 +2455,13 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]: model_config = self.model_config speculative_config = self.speculative_config + if ( + model_config is not None + and model_config.has_inner_state + and self.cache_config.mamba_cache_mode == "align" + ): + unsupported.append("hybrid/mamba models with align cache mode") + if self.parallel_config.prefill_context_parallel_size > 1: unsupported.append("prefill context parallelism") @@ -2669,19 +2478,11 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]: # TODO: ngram / ngram_gpu are not supported by the v2 model runner yet if speculative_config.method in ("ngram", "ngram_gpu"): unsupported.append("ngram/ngram_gpu speculative decoding") - elif speculative_config.method not in ( - "eagle", - "eagle3", - "mtp", - "dflash", - ): + elif speculative_config.method not in ("eagle", "eagle3", "mtp"): unsupported.append(f"speculative method '{speculative_config.method}'") # V2 EagleSpeculator does not support parallel_drafting (required by PEagle) - if ( - speculative_config.parallel_drafting - and speculative_config.method != "dflash" - ): + if speculative_config.parallel_drafting: unsupported.append("parallel drafting for speculative decoding") if ( @@ -2689,11 +2490,6 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]: and self.parallel_config.pipeline_parallel_size > 1 ): unsupported.append("EAGLE3 with pipeline parallelism") - if ( - speculative_config.method == "dflash" - and self.parallel_config.pipeline_parallel_size > 1 - ): - unsupported.append("DFlash with pipeline parallelism") if self.parallel_config.enable_dbo: unsupported.append("dual batch overlap") @@ -2796,6 +2592,10 @@ def validate_block_size(self) -> None: "to schedule a multiple of block_size tokens even if they are " "in the middle of a mm input" ) + # TODO: support align mamba cache mode for model runner v2 + assert not envs.VLLM_USE_V2_MODEL_RUNNER, ( + "Model Runner V2 has not yet supported mamba_cache_mode='align'. " + ) @model_validator(mode="after") def validate_nvfp4_kv_cache_with_mla(self) -> "VllmConfig": diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index fc46b2419e..34b999e5d5 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -21,10 +21,9 @@ from .base_device_communicator import DeviceCommunicatorBase logger = init_logger(__name__) -_SM70_TP4_LONG_PREFILL_BUFFER_BYTES = 8192 * 5120 * 2 -_SEEN_TP_ALLREDUCE_PATHS: set[tuple[str, str, tuple[int, ...], torch.dtype, int]] = ( - set() -) +_SEEN_TP_ALLREDUCE_PATHS: set[ + tuple[str, str, tuple[int, ...], torch.dtype, int] +] = set() def _trace_all_reduce_path( @@ -83,7 +82,6 @@ def __init__( use_top1_custom_ar = False use_sm70_awq_mlp_down_tile_ar = False use_sm70_awq_mlp_down_tile_overlap = False - use_sm70_tp4_long_prefill_fused_norm = False use_torch_symm_mem = False use_flashinfer_allreduce = False else: @@ -95,19 +93,15 @@ def __init__( use_sm70_awq_mlp_down_tile_overlap = ( envs.VLLM_SM70_AWQ_MLP_DOWN_TILE_OVERLAP ) - use_sm70_tp4_long_prefill_fused_norm = ( - envs.VLLM_SM70_TP4_LONG_PREFILL_FUSED_NORM - and self.world_size == 4 - and current_platform.is_device_capability(70) - ) use_torch_symm_mem = envs.VLLM_ALLREDUCE_USE_SYMM_MEM use_flashinfer_allreduce = envs.VLLM_ALLREDUCE_USE_FLASHINFER self.use_custom_allreduce = use_custom_allreduce self.use_top1_custom_ar = use_top1_custom_ar self.use_sm70_awq_mlp_down_tile_ar = use_sm70_awq_mlp_down_tile_ar - self.use_sm70_awq_mlp_down_tile_overlap = use_sm70_awq_mlp_down_tile_overlap - self.use_sm70_tp4_long_prefill_fused_norm = use_sm70_tp4_long_prefill_fused_norm + self.use_sm70_awq_mlp_down_tile_overlap = ( + use_sm70_awq_mlp_down_tile_overlap + ) self.use_torch_symm_mem = use_torch_symm_mem self.use_flashinfer_allreduce = use_flashinfer_allreduce @@ -155,7 +149,6 @@ def __init__( or use_top1_custom_ar or use_sm70_awq_mlp_down_tile_ar or use_sm70_awq_mlp_down_tile_overlap - or use_sm70_tp4_long_prefill_fused_norm ) and self.world_size > 1: # Initialize a custom fast all-reduce implementation. # `use_top1_custom_ar` deliberately only provisions the IPC/signaling @@ -168,12 +161,6 @@ def __init__( symm_mem_enabled=( self.symm_mem_comm is not None and not self.symm_mem_comm.disabled ), - max_size=( - _SM70_TP4_LONG_PREFILL_BUFFER_BYTES - if use_sm70_tp4_long_prefill_fused_norm - else 8192 * 1024 - ), - long_prefill_fusion_enabled=use_sm70_tp4_long_prefill_fused_norm, ) if current_platform.is_rocm(): @@ -448,31 +435,6 @@ def sm70_tp2_all_reduce_gemma_rms_norm( input_, residual, weight, epsilon ) - def sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather( - self, - input_: torch.Tensor, - residual: torch.Tensor, - weight: torch.Tensor, - epsilon: float, - ) -> tuple[torch.Tensor, torch.Tensor]: - ca_comm = self.ca_comm - if ( - not self.use_sm70_tp4_long_prefill_fused_norm - or ca_comm is None - or ca_comm.disabled - or not ca_comm.can_sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather( - input_, residual, weight - ) - ): - raise RuntimeError( - "SM70 TP4 long-prefill fused collective-norm route was " - "compiled for an unsupported runtime shape or communicator." - ) - _trace_all_reduce_path(self, "sm70_tp4_fused_rs_norm_ag", input_) - return ca_comm.sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather( - input_, residual, weight, epsilon - ) - def sm70_awq_mlp_down_tile_all_reduce(self, input_): if not self.use_sm70_awq_mlp_down_tile_ar: return None @@ -506,14 +468,18 @@ def sm70_awq_mlp_down_tile_gemm_reduce( k_ld, q_ld, tile_numel=envs.VLLM_SM70_AWQ_MLP_DOWN_TILE_OVERLAP_TILE_NUMEL, - reducer_blocks=(envs.VLLM_SM70_AWQ_MLP_DOWN_TILE_OVERLAP_REDUCER_BLOCKS), + reducer_blocks=( + envs.VLLM_SM70_AWQ_MLP_DOWN_TILE_OVERLAP_REDUCER_BLOCKS + ), kernel_reducer_blocks=( envs.VLLM_SM70_AWQ_MLP_DOWN_TILE_OVERLAP_KERNEL_REDUCER_BLOCKS ), overlap=envs.VLLM_SM70_AWQ_MLP_DOWN_TILE_OVERLAP_SIDE_STREAM, ) if out is not None: - _trace_all_reduce_path(self, "sm70_awq_mlp_down_tile_overlap", input_) + _trace_all_reduce_path( + self, "sm70_awq_mlp_down_tile_overlap", input_ + ) return out def reduce_scatter(self, input_: torch.Tensor, dim: int = -1): diff --git a/vllm/distributed/device_communicators/custom_all_reduce.py b/vllm/distributed/device_communicators/custom_all_reduce.py index 8d55ceeaff..a06dfebc26 100644 --- a/vllm/distributed/device_communicators/custom_all_reduce.py +++ b/vllm/distributed/device_communicators/custom_all_reduce.py @@ -1,11 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import os from contextlib import contextmanager from typing import cast -import regex as re import torch import torch.distributed as dist from torch.distributed import ProcessGroup @@ -30,49 +28,6 @@ logger = init_logger(__name__) _SM70_TP8_HIERARCHICAL_ELEMENTS = 4096 -_EXPANDABLE_SEGMENTS_TRUE_PATTERN = re.compile( - r"((?:^|,)\s*expandable_segments\s*:\s*)True(?=\s*(?:,|$))" -) - - -def _get_effective_allocator_conf() -> str: - """Return the allocator config selected by PyTorch on CUDA. - - PyTorch 2.10 prefers the legacy CUDA-specific variable whenever it is - present, even when its value is empty, and otherwise falls back to the - unified accelerator variable. - """ - if "PYTORCH_CUDA_ALLOC_CONF" in os.environ: - return os.environ["PYTORCH_CUDA_ALLOC_CONF"] - return os.environ.get("PYTORCH_ALLOC_CONF", "") - - -@contextmanager -def _disable_expandable_segments_for_cuda_ipc(active: bool): - """Keep CUDA-IPC graph buffers on cudaMalloc-backed allocations. - - ``cudaIpcGetMemHandle`` cannot export CUDA VMM allocations created when - ``expandable_segments:True`` is active through either supported allocator - environment variable. Custom all-reduce exports graph buffers through that - legacy IPC API, so disable expandable segments for the complete capture - and registration window while preserving every other allocator option. - """ - allocator_conf = _get_effective_allocator_conf() - disabled_conf, replacements = _EXPANDABLE_SEGMENTS_TRUE_PATTERN.subn( - r"\1False", allocator_conf - ) - should_disable = active and current_platform.is_cuda() and replacements > 0 - if should_disable: - logger.info_once( - "Temporarily disabling expandable_segments during custom allreduce " - "CUDA graph capture for CUDA IPC compatibility." - ) - torch.cuda.memory._set_allocator_settings(disabled_conf) - try: - yield - finally: - if should_disable: - torch.cuda.memory._set_allocator_settings(allocator_conf) def _sm70_tp8_hierarchical_peer_ranks(rank: int) -> tuple[int, ...]: @@ -113,7 +68,6 @@ def __init__( device: int | str | torch.device, max_size=8192 * 1024, symm_mem_enabled=False, - long_prefill_fusion_enabled=False, ) -> None: """ Args: @@ -127,8 +81,6 @@ def __init__( """ self._IS_CAPTURING = False self.disabled = True - self.long_prefill_output_ptrs: list[int] | None = None - self.sm70_tp4_push_buffer_ptrs: list[int] | None = None if not custom_ar: # disable because of missing custom allreduce library @@ -181,7 +133,6 @@ def __init__( if ( current_platform.is_cuda() and symm_mem_enabled - and not long_prefill_fusion_enabled and device_capability is not None ): device_capability_str = device_capability.as_version_str() @@ -300,12 +251,6 @@ def __init__( 8 * 1024 * 1024, dtype=torch.uint8, device=self.device ) self.max_size = max_size - # Provisioning the long-prefill fusion buffers must not widen ordinary - # custom-AR dispatch beyond its established 8-MiB policy. The fused - # predicate below uses the full allocation capacity explicitly. - self.dispatch_max_size = ( - min(max_size, 8192 * 1024) if long_prefill_fusion_enabled else max_size - ) self.rank = rank self.world_size = world_size self.fully_connected = fully_connected @@ -314,33 +259,6 @@ def __init__( self.meta_ptrs, self.rank_data, rank, self.fully_connected ) ops.register_buffer(self._ptr, self.buffer_ptrs) - if long_prefill_fusion_enabled: - self.long_prefill_output_ptrs = self.create_shared_buffer( - max_size, - group=group, - ) - ops.register_buffer(self._ptr, self.long_prefill_output_ptrs) - if ( - envs.VLLM_SM70_TP4_PUSH_ALLREDUCE - and world_size == 4 - and fully_connected - and current_platform.is_cuda() - and device_capability is not None - and device_capability.major == 7 - and device_capability.minor == 0 - ): - push_buffer_size = ops.sm70_tp4_push_allreduce_buffer_size() - self.sm70_tp4_push_buffer_ptrs = self.create_shared_buffer( - push_buffer_size, - group=group, - ) - ops.register_sm70_tp4_push_allreduce_buffer( - self._ptr, self.sm70_tp4_push_buffer_ptrs - ) - logger.info( - "SM70 TP4 SGLang-style push all-reduce enabled for the " - "FP16 80-KiB verifier and 8-KiB decode payloads." - ) @contextmanager def capture(self): @@ -349,14 +267,13 @@ def capture(self): `register_graph_buffers` call at the end of the context. It records all the buffer addresses used in the CUDA graph. """ - with _disable_expandable_segments_for_cuda_ipc(not self.disabled): - try: - self._IS_CAPTURING = True - yield - finally: - self._IS_CAPTURING = False - if not self.disabled: - self.register_graph_buffers() + try: + self._IS_CAPTURING = True + yield + finally: + self._IS_CAPTURING = False + if not self.disabled: + self.register_graph_buffers() def register_graph_buffers(self): handle, offset = ops.get_graph_buffer_ipc_meta(self._ptr) @@ -394,7 +311,7 @@ def should_custom_ar(self, inp: torch.Tensor): # for 4 or more non NVLink-capable GPUs, custom allreduce provides # little performance improvement over NCCL. if self.world_size == 2 or self.fully_connected: - return inp_size < self.dispatch_max_size + return inp_size < self.max_size return False def all_reduce( @@ -536,72 +453,6 @@ def can_sm70_tp4_all_reduce_gemma_rms_norm( and inp.device == residual.device == weight.device ) - def sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather( - self, - inp: torch.Tensor, - residual: torch.Tensor, - weight: torch.Tensor, - epsilon: float, - *, - normalized_out: torch.Tensor | None = None, - residual_out: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Benchmark-only long-prefill fused RS + Gemma RMSNorm + AG.""" - if not self.can_sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather( - inp, residual, weight - ): - raise RuntimeError("SM70 TP4 long-prefill fused norm is unavailable") - assert self.long_prefill_output_ptrs is not None - if normalized_out is None: - normalized_out = torch.empty_like(inp) - if residual_out is None: - residual_out = torch.empty_like(residual, dtype=torch.float32) - graph_registered = ( - self._IS_CAPTURING and torch.cuda.is_current_stream_capturing() - ) - ops.sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather( - self._ptr, - inp, - residual, - weight, - normalized_out, - residual_out, - 0 if graph_registered else self.buffer_ptrs[self.rank], - 0 if graph_registered else self.long_prefill_output_ptrs[self.rank], - self.max_size, - epsilon, - ) - return normalized_out, residual_out - - def can_sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather( - self, - inp: torch.Tensor, - residual: torch.Tensor, - weight: torch.Tensor, - ) -> bool: - return ( - not self.disabled - and self.long_prefill_output_ptrs is not None - and self.world_size == 4 - and self.fully_connected - and inp.is_cuda - and inp.dtype == torch.float16 - and inp.ndim == 2 - and inp.shape[0] % self.world_size == 0 - and 1 <= inp.shape[0] // self.world_size <= 2048 - and inp.shape[1] == 5120 - and inp.numel() * inp.element_size() <= self.max_size - and residual.shape == inp.shape - and residual.dtype == torch.float32 - and weight.ndim == 1 - and weight.numel() == 5120 - and weight.dtype in (torch.float16, torch.float32) - and inp.is_contiguous() - and residual.is_contiguous() - and weight.is_contiguous() - and inp.device == residual.device == weight.device - ) - def top1_argmax( self, input_pair: torch.Tensor, @@ -867,15 +718,6 @@ def close(self): self._ptr = 0 self.free_shared_buffer(self.meta_ptrs, rank=self.rank) self.free_shared_buffer(self.buffer_ptrs, rank=self.rank) - if self.long_prefill_output_ptrs is not None: - self.free_shared_buffer( - self.long_prefill_output_ptrs, - rank=self.rank, - ) - self.long_prefill_output_ptrs = None - if self.sm70_tp4_push_buffer_ptrs is not None: - self.free_shared_buffer(self.sm70_tp4_push_buffer_ptrs, rank=self.rank) - self.sm70_tp4_push_buffer_ptrs = None def __del__(self): self.close() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/__init__.py b/vllm/distributed/kv_transfer/kv_connector/v1/__init__.py index 96ab863d02..47329207f4 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/__init__.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/__init__.py @@ -4,9 +4,7 @@ KVConnectorBase_V1, KVConnectorRole, SupportsHMA, - SupportsVmmSafeTransfers, supports_hma, - supports_vmm_safe_transfers, ) from vllm.distributed.kv_transfer.kv_connector.v1.decode_bench_connector import ( # noqa: E501 DecodeBenchConnector, @@ -17,7 +15,5 @@ "KVConnectorBase_V1", "supports_hma", "SupportsHMA", - "SupportsVmmSafeTransfers", - "supports_vmm_safe_transfers", "DecodeBenchConnector", ] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/base.py index bf98b69872..fb5658da88 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/base.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/base.py @@ -121,21 +121,6 @@ def supports_hma(connector: Any) -> bool: return isinstance(connector, SupportsHMA) -class SupportsVmmSafeTransfers(ABC): # noqa: B024 - """Marker for connectors safe under CUDA VMM remapping. - - Implement this marker only when every KV transfer dereferences GPU - virtual addresses at transfer time (for example, cuMemcpy/cudaMemcpy), - and the connector never registers KV device memory with RDMA, GDS, - CUDA IPC, or another mechanism that depends on stable physical pages. - """ - - -def supports_vmm_safe_transfers(connector_cls: type) -> bool: - """Return whether a connector class explicitly declares VMM safety.""" - return issubclass(connector_cls, SupportsVmmSafeTransfers) - - class KVConnectorRole(enum.Enum): # Connector running in the scheduler process SCHEDULER = 0 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index d0f2124b22..6560bcf9de 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -15,7 +15,7 @@ TransferJob, ) from vllm.logger import init_logger -from vllm.utils.math_utils import cdiv, round_down +from vllm.utils.math_utils import cdiv from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_cache_interface import ( @@ -73,10 +73,6 @@ class GroupOffloadConfig(NamedTuple): # than the MLA full-attention group). # None for full-attention groups or when the optimization doesn't apply. alignment_block_count: int | None = None - # Mamba ``align`` tables mutate existing logical positions. Such groups - # must be stored from exact core-provided boundary block IDs instead of the - # connector's append-only positional mirror. - requires_exact_boundary_source: bool = False def get_sliding_window_size_in_blocks( @@ -94,24 +90,6 @@ def get_sliding_window_size_in_blocks( return None -def resolve_mamba_align_size(spec: "OffloadingSpec") -> int | None: - """Scan all KV cache groups in *spec* and return the single mamba alignment - size, or None if no group requires mamba alignment. - - For MambaSpec groups in "align" cache mode the hit window must be rounded - down to a multiple of the offloaded block size. Asserts that all such - groups agree on the same value. - """ - mamba_align_size: int | None = None - for idx, gpu_block_size in enumerate(spec.gpu_block_size): - kv_spec = spec.kv_cache_config.kv_cache_groups[idx].kv_cache_spec - if isinstance(kv_spec, MambaSpec) and kv_spec.mamba_cache_mode == "align": - offload_block_size = gpu_block_size * spec.block_size_factor - assert mamba_align_size is None or mamba_align_size == offload_block_size - mamba_align_size = offload_block_size - return mamba_align_size - - class SchedulerOffloadConfig(NamedTuple): kv_group_configs: tuple[GroupOffloadConfig, ...] block_size_factor: int @@ -155,12 +133,6 @@ def _alignment_block_count( return None return per_segment - def _requires_exact_boundary_source(group_idx: int) -> bool: - kv_spec = spec.kv_cache_config.kv_cache_groups[group_idx].kv_cache_spec - return ( - isinstance(kv_spec, MambaSpec) and kv_spec.mamba_cache_mode == "align" - ) - return cls( num_workers=spec.vllm_config.parallel_config.world_size, kv_group_configs=tuple( @@ -181,9 +153,6 @@ def _requires_exact_boundary_source(group_idx: int) -> bool: alignment_block_count=_alignment_block_count( gpu_block_size * spec.block_size_factor, sw ), - requires_exact_boundary_source=( - _requires_exact_boundary_source(idx) - ), ) for idx, gpu_block_size in enumerate(spec.gpu_block_size) ), @@ -313,7 +282,6 @@ def _sliding_window_sort_key(i: int) -> int: # used by _lookup self._sliding_window_groups: tuple[int, ...] = tuple(sliding_window_groups) self._lookup_groups = tuple(full_attention_groups) + self._sliding_window_groups - self._mamba_align_size: int | None = resolve_mamba_align_size(spec) self._req_status: dict[ReqId, RequestOffloadState] = {} self._current_batch_load_jobs: dict[int, TransferJob] = {} @@ -344,10 +312,7 @@ def _generate_job_id(self) -> int: return job_id def _remove_pending_job(self, job_id: int, block_ids: list[int] | None) -> None: - # Registration is idempotent because each block maps to a set of job IDs. - # MTP lookahead may report the same physical block ID more than once, so - # completion must use the same set semantics instead of removing twice. - for bid in set(block_ids or ()): + for bid in block_ids or (): pending = self._block_id_to_pending_jobs[bid] pending.remove(job_id) if not pending: @@ -433,11 +398,6 @@ def _lookup(self, req_status: RequestOffloadState) -> int | None: # for sliding window attention, we must reduce by 1 to make sure # we still have a hit after reduction max_hit_size_tokens -= 1 - if self._mamba_align_size is not None: - # Constrain hit-window to the mamba block size. - max_hit_size_tokens = round_down( - max_hit_size_tokens, self._mamba_align_size - ) num_hit_tokens: int = 0 defer_lookup = False lookup_groups = self._lookup_groups @@ -767,8 +727,6 @@ def _build_store_jobs( for group_config, group_state in zip( self.config.kv_group_configs, req_status.group_states ): - if group_config.requires_exact_boundary_source: - continue num_blocks = num_offloadable_tokens // group_config.offloaded_block_size start_block_idx = group_state.next_stored_block_idx if num_blocks <= start_block_idx: @@ -837,10 +795,6 @@ def _build_store_jobs( for group_config, group_state in zip( self.config.kv_group_configs, req_status.group_states ): - if group_config.requires_exact_boundary_source: - group_sizes.append(0) - block_indices.append(0) - continue is_sliding_window = ( group_config.sliding_window_size_in_blocks is not None ) @@ -918,99 +872,6 @@ def _build_store_jobs( return store_jobs - def _build_boundary_state_store_jobs( - self, - scheduler_output: SchedulerOutput, - ) -> dict[int, TransferJob]: - """Build stores from exact committed Mamba boundary-state blocks.""" - handoffs = scheduler_output.boundary_state_offloads - if not handoffs: - return {} - - store_jobs: dict[int, TransferJob] = {} - num_groups = len(self.config.kv_group_configs) - for req_id, entries in handoffs.items(): - req_status = self._req_status.get(req_id) - if req_status is None: - continue - req_status.update_offload_keys() - req = req_status.req - num_scheduled_tokens = scheduler_output.num_scheduled_tokens.get(req_id, 0) - num_offloadable_tokens = min( - req.num_computed_tokens + num_scheduled_tokens, - req.num_tokens, - ) - if req_status.max_offload_tokens is not None: - num_offloadable_tokens = min( - num_offloadable_tokens, req_status.max_offload_tokens - ) - if self.config.offload_prompt_only: - num_offloadable_tokens = min( - num_offloadable_tokens, req.num_prompt_tokens - ) - - for group_idx, block_id, boundary_tokens in entries: - group_config = self.config.kv_group_configs[group_idx] - if not group_config.requires_exact_boundary_source: - continue - if ( - block_id == 0 - or boundary_tokens > num_offloadable_tokens - or boundary_tokens % group_config.offloaded_block_size != 0 - ): - continue - - offload_block_idx = ( - boundary_tokens // group_config.offloaded_block_size - 1 - ) - group_state = req_status.group_states[group_idx] - if offload_block_idx >= len(group_state.offload_keys): - continue - offload_key = group_state.offload_keys[offload_block_idx] - store_output = self.manager.prepare_store( - [offload_key], req_status.req_context - ) - if store_output is None: - logger.warning( - "Request %s: cannot store Mamba boundary at %d tokens", - req_id, - boundary_tokens, - ) - continue - keys_to_store = set(store_output.keys_to_store) - if not keys_to_store: - continue - assert keys_to_store == {offload_key} - - job_id = self._generate_job_id() - req_status.transfer_jobs.add(job_id) - self._block_id_to_pending_jobs.setdefault(block_id, set()).add(job_id) - self._jobs[job_id] = TransferJobStatus( - req_id=req_id, - pending_count=self.config.num_workers, - keys=keys_to_store, - is_store=True, - sliding_window_block_ids=[block_id], - ) - - group_sizes = [0] * num_groups - group_sizes[group_idx] = 1 - block_indices = [0] * num_groups - block_indices[group_idx] = ( - boundary_tokens // group_config.gpu_block_size - 1 - ) - src_spec = GPULoadStoreSpec( - [block_id], - group_sizes=group_sizes, - block_indices=block_indices, - ) - store_jobs[job_id] = TransferJob( - req_id=req_id, - transfer_spec=(src_spec, store_output.store_spec), - ) - - return store_jobs - def build_connector_meta( self, scheduler_output: SchedulerOutput ) -> KVConnectorMetadata: @@ -1030,11 +891,9 @@ def build_connector_meta( ): self._current_batch_jobs_to_flush.update(self._jobs.keys()) - boundary_store_jobs = self._build_boundary_state_store_jobs(scheduler_output) - normal_store_jobs = self._build_store_jobs(scheduler_output) meta = OffloadingConnectorMetadata( load_jobs=self._current_batch_load_jobs, - store_jobs=boundary_store_jobs | normal_store_jobs, + store_jobs=self._build_store_jobs(scheduler_output), jobs_to_flush=self._current_batch_jobs_to_flush, ) self._current_batch_load_jobs = {} diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py index e1be60bcd4..20888c71f8 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py @@ -11,7 +11,6 @@ KVConnectorBase_V1, KVConnectorRole, SupportsHMA, - SupportsVmmSafeTransfers, ) from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( @@ -44,15 +43,7 @@ from vllm.v1.request import Request -class OffloadingConnector(KVConnectorBase_V1, SupportsHMA, SupportsVmmSafeTransfers): - """Native CPU KV offload connector. - - Transfers resolve the current tensor virtual address for every operation - and use CUDA memcpy. The connector does not register or export GPU KV - memory through RDMA, GDS, or CUDA IPC, so it is safe under CUDA VMM - remapping. - """ - +class OffloadingConnector(KVConnectorBase_V1, SupportsHMA): @property def prefer_cross_layer_blocks(self) -> bool: return True diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 8449f958cb..a567922199 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -182,33 +182,6 @@ def sm70_tp2_all_reduce_gemma_rms_norm_fake( return torch.empty_like(tensor), torch.empty_like(residual, dtype=torch.float32) -def sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather( - tensor: torch.Tensor, - residual: torch.Tensor, - weight: torch.Tensor, - epsilon: float, - group_name: str, -) -> tuple[torch.Tensor, torch.Tensor]: - assert group_name in _groups, f"Group {group_name} is not found." - group = _groups[group_name]() - if group is None: - raise ValueError(f"Group {group_name} is destroyed.") - return group._sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather_out_place( - tensor, residual, weight, epsilon - ) - - -def sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather_fake( - tensor: torch.Tensor, - residual: torch.Tensor, - weight: torch.Tensor, - epsilon: float, - group_name: str, -) -> tuple[torch.Tensor, torch.Tensor]: - del weight, epsilon, group_name - return torch.empty_like(tensor), torch.empty_like(residual, dtype=torch.float32) - - def sm70_awq_mlp_down_tile_all_reduce( tensor: torch.Tensor, group_name: str ) -> torch.Tensor: @@ -398,12 +371,6 @@ def patched_fused_scaled_matmul_reduce_scatter( fake_impl=sm70_tp2_all_reduce_gemma_rms_norm_fake, ) -direct_register_custom_op( - op_name="sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather", - op_func=sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather, - fake_impl=sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather_fake, -) - direct_register_custom_op( op_name="sm70_awq_mlp_down_tile_all_reduce", op_func=sm70_awq_mlp_down_tile_all_reduce, @@ -728,27 +695,9 @@ def _sm70_tp2_all_reduce_gemma_rms_norm_out_place( raise RuntimeError("Device communicator lacks SM70 fused AR RMSNorm") return fused_op(input_, residual, weight, epsilon) - def _sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather_out_place( - self, - input_: torch.Tensor, - residual: torch.Tensor, - weight: torch.Tensor, - epsilon: float, - ) -> tuple[torch.Tensor, torch.Tensor]: - if self.device_communicator is None: - raise ValueError("No device communicator found") - fused_op = getattr( - self.device_communicator, - "sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather", - None, - ) - if fused_op is None: - raise RuntimeError( - "Device communicator lacks SM70 TP4 fused collective-norm" - ) - return fused_op(input_, residual, weight, epsilon) - - def sm70_awq_mlp_down_tile_all_reduce(self, input_: torch.Tensor) -> torch.Tensor: + def sm70_awq_mlp_down_tile_all_reduce( + self, input_: torch.Tensor + ) -> torch.Tensor: if self.world_size == 1: return input_ @@ -1223,68 +1172,6 @@ def isend_tensor_dict( return handles - def isend_tensor_dict_static( - self, - tensor_dict: dict[str, torch.Tensor], - dst: int | None = None, - ) -> list[Handle]: - """Send tensors without metadata when both peers know the schema. - - The receiver must provide the same ordered keys, shapes, dtypes, and - devices to :meth:`irecv_tensor_dict_static`. Callers must fall back to - :meth:`isend_tensor_dict` whenever that contract can change. - """ - if self.world_size <= 1: - return [] - if self.use_cpu_custom_send_recv: - raise RuntimeError("static tensor transfer requires a process group") - - if dst is None: - dst = (self.rank_in_group + 1) % self.world_size - assert 0 <= dst < self.world_size, f"Invalid dst rank ({dst})" - - handles: list[Handle] = [] - for tensor in tensor_dict.values(): - if not isinstance(tensor, torch.Tensor): - raise TypeError("static tensor transfer only accepts tensors") - if tensor.numel() == 0: - continue - comm_group = self.cpu_group if tensor.is_cpu else self.device_group - handle = torch.distributed.isend( - tensor, dst=self.ranks[dst], group=comm_group - ) - if tensor.is_cuda: - tensor.record_stream(torch.cuda.current_stream(tensor.device)) - handles.append(handle) - return handles - - def irecv_tensor_dict_static( - self, - tensor_dict: dict[str, torch.Tensor], - src: int | None = None, - ) -> list[Handle]: - """Receive metadata-free tensors into caller-owned buffers.""" - if self.world_size <= 1: - return [] - if self.use_cpu_custom_send_recv: - raise RuntimeError("static tensor transfer requires a process group") - - if src is None: - src = (self.rank_in_group - 1) % self.world_size - assert 0 <= src < self.world_size, f"Invalid src rank ({src})" - - handles: list[Handle] = [] - for tensor in tensor_dict.values(): - if not isinstance(tensor, torch.Tensor): - raise TypeError("static tensor transfer only accepts tensors") - if tensor.numel() == 0: - continue - comm_group = self.cpu_group if tensor.is_cpu else self.device_group - handles.append( - torch.distributed.irecv(tensor, src=self.ranks[src], group=comm_group) - ) - return handles - def recv_tensor_dict( self, src: int | None = None, diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 31850497d3..121d2a30f8 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -93,11 +93,7 @@ ) from vllm.config.scheduler import SchedulerPolicy from vllm.config.utils import get_field -from vllm.config.vllm import ( - OptimizationLevel, - PerformanceMode, - _sm70_mtp_cudagraph_capture_sizes, -) +from vllm.config.vllm import OptimizationLevel, PerformanceMode from vllm.logger import init_logger, suppress_logging from vllm.platforms import CpuArchEnum, current_platform from vllm.plugins import load_general_plugins @@ -1769,11 +1765,10 @@ def _maybe_apply_sm70_mtp_defaults( has_linear_attention = any( layer_type == "linear_attention" for layer_type in layer_types ) + if is_server and self.enable_prefix_caching is None and has_linear_attention: self.enable_prefix_caching = True - profile_updates.append( - f"enable_prefix_caching={self.enable_prefix_caching}" - ) + profile_updates.append("enable_prefix_caching=True") if ( self.enable_prefix_caching and has_linear_attention @@ -1819,9 +1814,12 @@ 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. + # For SM70 native MTP, official Qwen sampling is non-greedy + # (temperature=1.0/top_p=0.95/top_k=20). Probabilistic draft + # sampling provides draft_probs for standard rejection sampling and + # is the validated high-acceptance path. Greedy requests still take + # the argmax fast path inside the proposer when sampling metadata is + # all-greedy, so this default does not penalize deterministic runs. draft_sample_method = "probabilistic" self.speculative_config["draft_sample_method"] = draft_sample_method profile_updates.append( @@ -1862,6 +1860,11 @@ def _maybe_apply_sm70_mtp_defaults( and self.compilation_config.cudagraph_capture_sizes is None and self.compilation_config.max_cudagraph_capture_size is None ): + cudagraph_capture_sizes = ( + [1, 2, 4, 8, 9, 18] + if self.tensor_parallel_size >= 4 + else [1, 2, 4, 8, 9] + ) num_speculative_tokens = int( self.speculative_config.get("num_speculative_tokens") or 0 ) @@ -1878,22 +1881,9 @@ def _maybe_apply_sm70_mtp_defaults( num_speculative_state_tokens = max( num_speculative_tokens, ddtree_budget ) - decode_query_len = num_speculative_state_tokens + 1 - max_num_seqs = max(int(self.max_num_seqs or 1), 1) - if envs.VLLM_SM70_MTP_SPLIT_DRAFT_CUDAGRAPHS and spec_method == "mtp": - cudagraph_capture_sizes = _sm70_mtp_cudagraph_capture_sizes( - max_num_seqs, - decode_query_len, - ) - profile_updates.append( - f"mtp_split_verifier_cudagraph_shapes={cudagraph_capture_sizes}" - ) - else: - cudagraph_capture_sizes = ( - [1, 2, 4, 8, 9, 18] - if self.tensor_parallel_size >= 4 - else [1, 2, 4, 8, 9] - ) + if num_speculative_state_tokens > 0: + decode_query_len = num_speculative_state_tokens + 1 + max_num_seqs = max(int(self.max_num_seqs or 1), 1) cudagraph_capture_sizes = sorted( set(cudagraph_capture_sizes) | { diff --git a/vllm/envs.py b/vllm/envs.py index 2d3e04aa45..499eb2f577 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -136,9 +136,6 @@ VLLM_SM70_AWQ_TP2_FAST_SELECTOR: bool = True VLLM_SM70_AWQ_TP2_FAST_TARGETS: str | None = None VLLM_SM70_TP2_AR_GEMMA_RMS_FUSION: bool = False - VLLM_SM70_TP4_LONG_PREFILL_FUSED_NORM: bool = False - VLLM_SM70_TP4_LONG_FUSED_NORM_THREADS: int = 512 - VLLM_SM70_TP4_LONG_FUSED_NORM_BLOCKS: int = 80 VLLM_SM70_AWQ_MLP_ENGINE: bool = False VLLM_SM70_AWQ_PREFILL_EXACT_DENSE: bool = True VLLM_SM70_AWQ_MLP_DOWN_TILE_AR: bool = False @@ -156,23 +153,11 @@ VLLM_SM70_FP8_COORDINATED_TUNING: bool = True VLLM_SM70_FP8_REUSE_IMPORTED_CACHE: bool = False VLLM_SM70_FP8_SAFE_FAST_SELECTOR: bool = False - 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_FP8_QPN8_LIBRARY: str | None = None - VLLM_SM70_SAMPLER_LIBRARY: str | None = None - VLLM_SM70_FP8_PREFILL_VISIBLE_DENSE_MM: bool = False - VLLM_SM70_NVFP4_QPN2: bool = False 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_AWQ_REUSE_IMPORTED_CACHE: bool = False VLLM_SM70_AWQ_WARMUP: bool = True VLLM_SM70_AWQ_WARMUP_MAX_M: int = 16 @@ -184,50 +169,22 @@ 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 VLLM_SM70_AWQ_MOE_TUNE_MAX_TOKENS: int = 128 - VLLM_SM70_NVFP4_MOE_TUNE_MAX_TOKENS: int = 128 VLLM_SM70_ENABLE_DENSE_F16_FASTPATH: bool = False VLLM_SM70_ENABLE_LM_HEAD_FASTPATH: bool = False VLLM_SM70_LM_HEAD_TOP1: bool = True VLLM_SM70_LM_HEAD_TOP1_TC: bool = False - 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_VERIFY_FASTPATH: bool = False - VLLM_SM70_DFLASH2_FUSED_GDN_METADATA: bool = False - VLLM_SM70_DFLASH2_GDN_METADATA_SHADOW: bool = False - VLLM_SM70_DFLASH2_GDN_SYNC_ASSERT: bool = False - VLLM_SM70_DFLASH2_FUSED_GDN_VERIFY: bool = False - VLLM_SM70_DFLASH2_FUSED_GDN_NORM: bool = False - VLLM_SM70_DFLASH2_FUSED_GDN_SPLIT: bool = False - VLLM_SM70_DFLASH2_FUSED_SMALLQ_METADATA: bool = False - VLLM_SM70_DFLASH2_GROUPED_SMALLQ_METADATA: bool = False - VLLM_SM70_DFLASH2_FUSED_QKV_PACK: bool = False - VLLM_SM70_DFLASH2_FUSED_GEMMA_RMS: bool = False - 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_TOP1_CUSTOM_AR: bool = False VLLM_SM70_GREEDY_TOKEN_FASTPATH: bool = True VLLM_SM70_GREEDY_TOKEN_FASTPATH_TRACE: bool = False - VLLM_SM70_COMPACT_TOPK20_SAMPLER: bool = False - VLLM_SM70_CHUNKED_TOPK20_CHUNKS: int = 0 - VLLM_SM70_TP_LOCAL_TOPK20_SAMPLER: bool = False - VLLM_SM70_TOPK_TOPP_8_WARPS: bool = False - VLLM_SM70_TOPK_TOPP_B8_B16_8_WARPS: bool = True VLLM_SM70_ASYNC_SCHEDULING_QUEUE_DEPTH: int = 0 VLLM_SM70_ASYNC_STAGED_INPUT_PREP: bool = False VLLM_SM70_ASYNC_CPU_TRACE: bool = False VLLM_SM70_ASYNC_CPU_TRACE_EVERY: int = 16 VLLM_TP_ALLREDUCE_TRACE: bool = False VLLM_CUSTOM_ALLREDUCE_BLOCK_LIMIT: int | None = None - VLLM_SM70_TP4_MTP_AR_BLOCK_TUNING: bool = False VLLM_SM70_TP4_M5_AR_THREADS: int | None = None - VLLM_SM70_TP4_SMALL_AR_PACK32: bool = False VLLM_SM70_F16_DENSE_ALLOWLIST: str | None = None VLLM_SM70_MOE_DENSE_ALLOWLIST: str | None = None VLLM_SM70_F16_DENSE_MAX_M: int = 64 @@ -241,21 +198,13 @@ VLLM_SM70_FP8_TURBOMIND: bool = True VLLM_SM70_FP8_DENSE_GATED_SILU: bool = True VLLM_SM70_NVFP4_TURBOMIND: bool = True - VLLM_SM70_NVFP4_MOE_GROUPED_PREFILL: bool = True - VLLM_SM70_NVFP4_DENSE_GATED_SILU: bool = True - VLLM_SM70_NVFP4_QPN4: bool = True - VLLM_SM70_NVFP4_QPN4_DOWN_SCALE_CODE: bool = False 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 = True + VLLM_SM70_MXFP4_MOE_COMPACT_GROUPED_DECODE: bool = False 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 = True - VLLM_SM70_MXFP4_MOE_DIRECT_ORDER_DECODE: bool = True - VLLM_SM70_MXFP4_MOE_BROADCAST_INPUT_DECODE: bool = True + VLLM_SM70_MXFP4_MOE_DIRECT_TOP6_DECODE: bool = False VLLM_SM70_DSV4_SPARSE_MLA_SPLITK_SWA: bool = False VLLM_SM70_DSV4_SPARSE_MLA_SPLITK_C4: bool = False VLLM_SM70_DSV4_SPARSE_MLA_SPLITK_C128: bool = False @@ -288,12 +237,8 @@ VLLM_SM70_DECODE_EVENT_TRACE_EVERY: int = 16 VLLM_SM70_MTP_PROFILE: bool = False VLLM_SM70_MTP_PROFILE_INTERVAL: int = 16 - VLLM_SM70_MTP_SPLIT_DRAFT_CUDAGRAPHS: bool = False - VLLM_SM70_MTP_CONCURRENCY_WARMUP: bool = False VLLM_SM70_MTP_CONTEXT_BUCKETS: str | None = None VLLM_SM70_DSV4_DECODE_CONTEXT_BUCKETS: str | None = None - VLLM_SM70_DSV4_PRIVATE_COMPRESSOR_STATE: bool = False - VLLM_SM70_FP8_KV_DECODE_CONTEXT_BUCKETS: str | None = None VLLM_SM70_MTP_CONTEXT_BUCKET_PARTITION_SIZE: str | None = None VLLM_SM70_REJECTION_PROFILE: bool = False VLLM_SM70_REJECTION_PROFILE_INTERVAL: int = 20 @@ -353,8 +298,6 @@ VLLM_FLASH_V100_PREFILL_GATHER_DENSE_MIN_KV: int = 8192 VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3: bool = True VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3_MIN_KV: int = 32768 - VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3_Q8000_EXPERIMENTAL: bool = False - VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL: bool = True VLLM_FLASH_V100_PREFILL_SPLIT_KV: bool = False VLLM_FLASH_V100_PREFILL_SPLIT_KV_TOKENS: int = 32768 VLLM_FLASH_V100_PREFILL_SPLIT_KV_MIN_Q: int = 1 @@ -392,8 +335,6 @@ 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_MIN_MODEL_LEN: int = 32768 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 @@ -403,20 +344,11 @@ VLLM_FLASH_V100_XQA_G6_QK_PIPELINE: bool = True VLLM_FLASH_V100_XQA_G6_QK_PIPELINE_WARPS: int = 8 VLLM_FLASH_V100_XQA_G6_QK_PIPELINE_TRACE: bool = False - VLLM_FLASH_V100_XQA_G6_DUAL_CTA: bool = False - VLLM_FLASH_V100_XQA_SPLIT_REDUCE: bool = False - VLLM_FLASH_V100_E4M3_BATCH_XQA: bool = True - VLLM_FLASH_V100_E4M3_BATCH_XQA_OPTIMIZED: bool = True - VLLM_FLASH_V100_E4M3_PAGE800_FASTPATH: bool = True - VLLM_FLASH_V100_E4M3_PAGE800_FASTPATH_TRACE: bool = False - VLLM_FLASH_V100_XQA_BATCH_CONTEXT_ROUTING: bool = True - VLLM_FLASH_V100_XQA_BATCH_CONTEXT_ROUTING_TRACE: bool = False VLLM_FLASH_V100_XQA_E5M2_G6_DUAL_CTA: bool = True VLLM_FLASH_V100_XQA_E5M2_G6_SPLIT_REDUCE: bool = True VLLM_FLASH_V100_XQA_E5M2_P1024_BEGIN: int = 61633 VLLM_FLASH_V100_XQA_E5M2_PARTITION_PAGE_IDS: bool = True VLLM_FLASH_V100_XQA_E5M2_PAIR_LOAD: bool = True - VLLM_FLASH_V100_XQA_E5M2_BATCH_WIDE_LOAD: bool = True VLLM_FLASH_V100_XQA_E5M2_G6_DUAL_CTA_TRACE: bool = False VLLM_FLASH_V100_TRACE_DECODE_ACTIVE: bool = False VLLM_FLASH_V100_DECODE_USE_SCALAR_PAGED: bool = True @@ -511,8 +443,6 @@ VLLM_SM70_QWEN_GDN_DISABLE_INPUT_CORE_OP: bool = False VLLM_SM70_QWEN_GDN_INPUT_PROJECTION_OP: bool = False VLLM_SM70_QWEN_GDN_OUTPUT_PROJECTION_OP: bool = False - VLLM_SM70_GDN_QPN8_BA_SPLIT: bool = False - VLLM_SM70_GDN_RMSNORM_ONEPASS: bool = False VLLM_SM70_GEMMA_RMS_NORM_EAGER: bool = False VLLM_SM70_GEMMA_RMS_NORM_COMPILE_NATIVE: bool = False VLLM_SM70_GEMMA_LONG_PREFILL_FUSED: bool = True @@ -530,7 +460,6 @@ VLLM_QWEN3NEXT_ENABLE_SHARED_MOE_OVERLAP: bool = False VLLM_SM70_DISABLE_QWEN3NEXT_SHARED_MOE_OVERLAP: bool = False VLLM_SM70_UNQUANTIZED_MOE_0DOT3_CONFIG: bool = True - VLLM_SM70_MTP_MOE_TUNED_CONFIG: bool = False VLLM_SM70_DENSE_CUDAGRAPH_CAPTURE: bool = False VLLM_SM70_USE_BREAKABLE_CUDAGRAPH: bool = False VLLM_SM70_FLASH_V100_0DOT3_COMPILE_GRAPH: bool = False @@ -1655,17 +1584,6 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_TP2_AR_GEMMA_RMS_FUSION": lambda: bool( int(os.getenv("VLLM_SM70_TP2_AR_GEMMA_RMS_FUSION", "0")) ), - # Default-off exact-8K TP4 experiment: fuse the row-parallel reduction, - # mixed-FP32-residual Gemma RMSNorm, and normalized-output all-gather. - "VLLM_SM70_TP4_LONG_PREFILL_FUSED_NORM": lambda: bool( - int(os.getenv("VLLM_SM70_TP4_LONG_PREFILL_FUSED_NORM", "0")) - ), - "VLLM_SM70_TP4_LONG_FUSED_NORM_THREADS": lambda: int( - os.getenv("VLLM_SM70_TP4_LONG_FUSED_NORM_THREADS", "512") - ), - "VLLM_SM70_TP4_LONG_FUSED_NORM_BLOCKS": lambda: int( - os.getenv("VLLM_SM70_TP4_LONG_FUSED_NORM_BLOCKS", "80") - ), # Experimental TileRT-inspired dense MLP lane for SM70 AWQ decode. The # first stage fuses gate_up_proj + SiluAndMul through the TurboMind GEMM # epilogue for M=1/TP2 while keeping the existing down_proj/reduce path. @@ -1683,31 +1601,6 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_FP8_PREFILL_EXACT_DENSE": lambda: bool( int(os.getenv("VLLM_SM70_FP8_PREFILL_EXACT_DENSE", "1")) ), - # Memory-neutral QPN8 layout for shape- and runtime-gated TP4 block-FP8 - # dense projections. Pure-FP8 checkpoints must opt in; - # 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_FP8_QPN8_PP2_TP4": lambda: bool( - int(os.getenv("VLLM_SM70_FP8_QPN8_PP2_TP4", "1")) - ), - # 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_FP8_PREFILL_CUTLASS": lambda: bool( - int(os.getenv("VLLM_SM70_FP8_PREFILL_CUTLASS", "1")) - ), - "VLLM_SM70_FP8_PREFILL_VISIBLE_DENSE_MM": lambda: bool( - int(os.getenv("VLLM_SM70_FP8_PREFILL_VISIBLE_DENSE_MM", "0")) - ), - # 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"))), # 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. @@ -1759,18 +1652,6 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_FP8_SAFE_FAST_SELECTOR": lambda: bool( int(os.getenv("VLLM_SM70_FP8_SAFE_FAST_SELECTOR", "0")) ), - "VLLM_SM70_FP8_GROUPED_BMM_DECODE": lambda: bool( - int(os.getenv("VLLM_SM70_FP8_GROUPED_BMM_DECODE", "1")) - ), - "VLLM_SM70_FP8_PREFILL_FAST_SELECTOR": lambda: bool( - int(os.getenv("VLLM_SM70_FP8_PREFILL_FAST_SELECTOR", "1")) - ), - "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")) ), @@ -1783,9 +1664,6 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_NVFP4_TUNE_SMALL_SHAPES": lambda: bool( int(os.getenv("VLLM_SM70_NVFP4_TUNE_SMALL_SHAPES", "1")) ), - "VLLM_SM70_NVFP4_QWEN38_TP4_M1_FAST_SELECTOR": lambda: bool( - int(os.getenv("VLLM_SM70_NVFP4_QWEN38_TP4_M1_FAST_SELECTOR", "1")) - ), "VLLM_SM70_AWQ_REUSE_IMPORTED_CACHE": lambda: bool( int(os.getenv("VLLM_SM70_AWQ_REUSE_IMPORTED_CACHE", "0")) ), @@ -1818,26 +1696,12 @@ 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")) ), - "VLLM_SM70_DSV4_QNORM_KV_FUSED_TP4": lambda: bool( - int(os.getenv("VLLM_SM70_DSV4_QNORM_KV_FUSED_TP4", "1")) - ), - # Skip PP metadata and TP reconstruction only for the exact, replicated - # SM70 B1 hidden-state schema validated by the worker on both stages. - "VLLM_SM70_PP_STATIC_HIDDEN_TRANSFER": lambda: bool( - int(os.getenv("VLLM_SM70_PP_STATIC_HIDDEN_TRANSFER", "1")) - ), "VLLM_SM70_AWQ_MOE_TUNE_MAX_TOKENS": lambda: int( os.getenv("VLLM_SM70_AWQ_MOE_TUNE_MAX_TOKENS", "128") ), - "VLLM_SM70_NVFP4_MOE_TUNE_MAX_TOKENS": lambda: int( - os.getenv("VLLM_SM70_NVFP4_MOE_TUNE_MAX_TOKENS", "128") - ), # Experimental unquantized FP16 SM70 TurboMind fast paths. Keep default-off: # these must pass the numeric policy and model-level token gate before # becoming part of the default V100 route. @@ -1873,110 +1737,6 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_LM_HEAD_TOP1_TC": lambda: bool( int(os.getenv("VLLM_SM70_LM_HEAD_TOP1_TC", "0")) ), - # Candidate-only QPN8 LM head. QPN8 selects a conservative local top-64 - # support, then directly re-evaluates the original FP16 rows and restores - # dense-vocabulary top-k tie order. Keep default-off until real-hidden, - # quality, and complete Graph/Nsight gates all pass. - "VLLM_SM70_DFLASH2_QPN8_RERANK": lambda: bool( - int(os.getenv("VLLM_SM70_DFLASH2_QPN8_RERANK", "0")) - ), - # Audit-only eager mode: execute QPN8+rerank, compare it with the dense - # local top-k, and return the dense result so the baseline trajectory is - # unchanged. This intentionally synchronizes for diagnostics. - "VLLM_SM70_DFLASH2_QPN8_RERANK_SHADOW": lambda: bool( - int(os.getenv("VLLM_SM70_DFLASH2_QPN8_RERANK_SHADOW", "0")) - ), - # Preserve the production full-vocabulary torch.topk tie contract by - # default. The candidate-order experiment avoids a host-blocking SM70 - # multi-block top-k, but may be enabled only for paired quality tests. - "VLLM_SM70_DFLASH2_QPN8_DENSE_ORDER": lambda: bool( - int(os.getenv("VLLM_SM70_DFLASH2_QPN8_DENSE_ORDER", "1")) - ), - # 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( - int(os.getenv("VLLM_SM70_DFLASH2_VERIFY_FASTPATH", "0")) - ), - # Build all selector-based DFlash target GDN state-index metadata with one - # pointer-table Triton launch. Keep separate from the shared-classification - # gate until the fixed-trajectory and mixed-batch Graph checks pass. - "VLLM_SM70_DFLASH2_FUSED_GDN_METADATA": lambda: bool( - int(os.getenv("VLLM_SM70_DFLASH2_FUSED_GDN_METADATA", "0")) - ), - # Debug-only oracle: materialize the legacy advanced-indexing contract and - # compare it with the fused persistent buffers before graph replay. - "VLLM_SM70_DFLASH2_GDN_METADATA_SHADOW": lambda: bool( - int(os.getenv("VLLM_SM70_DFLASH2_GDN_METADATA_SHADOW", "0")) - ), - # Debug-only device-value assertion. The metadata producers and FULL graph - # replay share the current CUDA stream, so production ordering does not - # require a host-blocking Tensor.item(). The no-sync route preserves the - # accepted token/acceptance trajectory; opt in only when auditing metadata. - "VLLM_SM70_DFLASH2_GDN_SYNC_ASSERT": lambda: bool( - int(os.getenv("VLLM_SM70_DFLASH2_GDN_SYNC_ASSERT", "0")) - ), - # Selector-based DFlash packed GDN verification kernel. This is deliberately - # independent from the shared-metadata umbrella gate so each optimization - # can be paired against the unchanged verifier in isolation. - "VLLM_SM70_DFLASH2_FUSED_GDN_VERIFY": lambda: bool( - int(os.getenv("VLLM_SM70_DFLASH2_FUSED_GDN_VERIFY", "0")) - ), - # Route compatible target GDN output gates through the existing one-pass - # CUDA RMSNormGated implementation. This remains an explicit opt-in. - "VLLM_SM70_DFLASH2_FUSED_GDN_NORM": lambda: bool( - int(os.getenv("VLLM_SM70_DFLASH2_FUSED_GDN_NORM", "0")) - ), - # Fuse compatible nonzero-offset GDN z/b/a materialization into one - # copy kernel. This must stay separate from the plain-view path because - # nonzero-offset views are unsafe under the SM70 compile/full-graph route. - "VLLM_SM70_DFLASH2_FUSED_GDN_SPLIT": lambda: bool( - int(os.getenv("VLLM_SM70_DFLASH2_FUSED_GDN_SPLIT", "0")) - ), - # Build Flash-V100 small-query verifier rows directly in their persistent - # graph buffers. This replaces four repeat_interleave scans per KV group. - # The matched TP4 trace is token/acceptance exact and cuts the synchronized - # DFlash2 draft-to-target interval from 5.720 ms to 1.911 ms on V100. - "VLLM_SM70_DFLASH2_FUSED_SMALLQ_METADATA": lambda: bool( - int(os.getenv("VLLM_SM70_DFLASH2_FUSED_SMALLQ_METADATA", "0")) - ), - # Collapse five compatible target small-query metadata launches into one - # heterogeneous-width pointer-table kernel. The paired - # TP4 node trace is token/acceptance exact and reduces synchronized D2T. - "VLLM_SM70_DFLASH2_GROUPED_SMALLQ_METADATA": lambda: bool( - int(os.getenv("VLLM_SM70_DFLASH2_GROUPED_SMALLQ_METADATA", "0")) - ), - # Copy the post-convolution Q/K/V row slices into the recurrent kernel's - # packed contiguous layout with one bitwise Triton launch. This replaces - # three reshape copies plus one cat in the exact block-eight verifier. - "VLLM_SM70_DFLASH2_FUSED_QKV_PACK": lambda: bool( - int(os.getenv("VLLM_SM70_DFLASH2_FUSED_QKV_PACK", "0")) - ), - # Fuse the FP16 projection + FP32 residual + Gemma RMSNorm suffix used by - # small DFlash2 verifier graphs. Default-off pending numeric/quality gates. - "VLLM_SM70_DFLASH2_FUSED_GEMMA_RMS": lambda: bool( - int(os.getenv("VLLM_SM70_DFLASH2_FUSED_GEMMA_RMS", "0")) - ), - # Avoid materializing/gathering full-vocabulary target logits when the - # DFlash2 proposal and target sampling distributions both have compact - # top-k support. Default-off until paired output/acceptance and end-to-end - # V100 gates pass; unsupported sampling features fall back to dense logits. - "VLLM_SM70_DFLASH2_SPARSE_TARGET_REJECTION": lambda: bool( - int(os.getenv("VLLM_SM70_DFLASH2_SPARSE_TARGET_REJECTION", "0")) - ), - # Compute the compatible 25600->5120 target-hidden projection as four output - # shards, then all-gather only the 80-KiB block-eight result. This - # remains opt-in until the acceptance/quality gate follows the positive - # production-weight TP4 microbenchmark. - "VLLM_SM70_DFLASH2_SHARDED_CONTEXT_FC": lambda: bool( - int(os.getenv("VLLM_SM70_DFLASH2_SHARDED_CONTEXT_FC", "0")) - ), - # Default-on SGLang-style push collective for the validated FP16 80-KiB - # verifier and 8-KiB decode payloads on fully-connected SM70 TP4 CUDA - # Graphs. Other devices, topologies, sizes, and eager calls retain the - # ordinary pull path; explicit 0 is the rollback. - "VLLM_SM70_TP4_PUSH_ALLREDUCE": lambda: bool( - int(os.getenv("VLLM_SM70_TP4_PUSH_ALLREDUCE", "1")) - ), # 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 @@ -1988,16 +1748,6 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_GREEDY_TOKEN_FASTPATH_TRACE": lambda: bool( int(os.getenv("VLLM_SM70_GREEDY_TOKEN_FASTPATH_TRACE", "0")) ), - # Opt-in V100 launch for the exact validated combined top-k/top-p shapes. - "VLLM_SM70_TOPK_TOPP_8_WARPS": lambda: bool( - int(os.getenv("VLLM_SM70_TOPK_TOPP_8_WARPS", "0")) - ), - # The exact B8/B16, 248320-column sampler contract uses one logits row per - # request. Eight warps preserves Qrita's masking math while improving SM70 - # reduction parallelism. Set to 0 to restore Triton's launch heuristic. - "VLLM_SM70_TOPK_TOPP_B8_B16_8_WARPS": lambda: bool( - int(os.getenv("VLLM_SM70_TOPK_TOPP_B8_B16_8_WARPS", "1")) - ), # Diagnostic SM70 async scheduling depth override. Default 0 preserves # upstream behavior. Values >2 let no-PP async scheduling enqueue more real # decode steps before collecting CPU-visible outputs; this changes queueing @@ -2025,18 +1775,11 @@ def _resolve_rust_frontend_path() -> str | None: if "VLLM_CUSTOM_ALLREDUCE_BLOCK_LIMIT" in os.environ else None ), - # Opt in exact MTP4 verifier payloads on fully-connected SM70 TP4. - "VLLM_SM70_TP4_MTP_AR_BLOCK_TUNING": lambda: bool( - int(os.getenv("VLLM_SM70_TP4_MTP_AR_BLOCK_TUNING", "0")) - ), "VLLM_SM70_TP4_M5_AR_THREADS": lambda: ( int(os.environ["VLLM_SM70_TP4_M5_AR_THREADS"]) if "VLLM_SM70_TP4_M5_AR_THREADS" in os.environ else None ), - "VLLM_SM70_TP4_SMALL_AR_PACK32": lambda: bool( - int(os.getenv("VLLM_SM70_TP4_SMALL_AR_PACK32", "0")) - ), # Optional custom allreduce for the tiny per-rank top1 pair. Keep default # off until the communicator path has same-criterion model evidence. "VLLM_SM70_TOP1_CUSTOM_AR": lambda: bool( @@ -2089,34 +1832,6 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_NVFP4_TURBOMIND": lambda: bool( int(os.getenv("VLLM_SM70_NVFP4_TURBOMIND", "1")) ), - # Dispatch all 256 routed experts in one grouped TurboMind call for the - # exact Qwen3.6-35B-A3B TP1/2/4 NVFP4 prefill shapes. B1-B8 decode keeps - # the compact active-expert route; larger graph shapes use full groups. - "VLLM_SM70_NVFP4_MOE_GROUPED_PREFILL": lambda: bool( - int(os.getenv("VLLM_SM70_NVFP4_MOE_GROUPED_PREFILL", "1")) - ), - "VLLM_SM70_NVFP4_DENSE_GATED_SILU": lambda: bool( - int(os.getenv("VLLM_SM70_NVFP4_DENSE_GATED_SILU", "1")) - ), - # Memory-neutral native QPN4 M=1 decode for the accepted Qwen3.8-27B - # NVFP4 TP4 no-MTP contract. Larger-M prefill dequantizes into one bounded - # shared FP16 workspace. Set to 0 to retain the TurboMind NVFP4 route. - "VLLM_SM70_NVFP4_QPN4": lambda: bool(int(os.getenv("VLLM_SM70_NVFP4_QPN4", "1"))), - "VLLM_SM70_NVFP4_QPN4_DOWN_SCALE_CODE": lambda: bool( - int(os.getenv("VLLM_SM70_NVFP4_QPN4_DOWN_SCALE_CODE", "0")) - ), - # Exact no-MTP random sampler for the Qwen3.8 TP4 batch-one contract. - # The base flag enables exact compact sampling after full-logits gather. - # The separate TP-local flag opts into rank-local top-20 reduction. - "VLLM_SM70_COMPACT_TOPK20_SAMPLER": lambda: bool( - int(os.getenv("VLLM_SM70_COMPACT_TOPK20_SAMPLER", "0")) - ), - "VLLM_SM70_CHUNKED_TOPK20_CHUNKS": lambda: int( - os.getenv("VLLM_SM70_CHUNKED_TOPK20_CHUNKS", "0") - ), - "VLLM_SM70_TP_LOCAL_TOPK20_SAMPLER": lambda: bool( - int(os.getenv("VLLM_SM70_TP_LOCAL_TOPK20_SAMPLER", "0")) - ), "VLLM_SM70_MXFP4_TURBOMIND": lambda: bool( int(os.getenv("VLLM_SM70_MXFP4_TURBOMIND", "1")) ), @@ -2185,24 +1900,13 @@ 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", "1")) + int(os.getenv("VLLM_SM70_MXFP4_MOE_COMPACT_GROUPED_DECODE", "0")) ), # Experimental verifier-M8 grouped dispatch. This collapses the fixed 48 # active-expert stage calls into one TurboMind grouped launch. "VLLM_SM70_MXFP4_MOE_GROUPED_M8": lambda: bool( int(os.getenv("VLLM_SM70_MXFP4_MOE_GROUPED_M8", "0")) ), - # Extend the one-launch grouped verifier route from M=8 to M=2..M=8. - # Kept independent until each width passes its CUDA Graph and model gates. - "VLLM_SM70_MXFP4_MOE_GROUPED_VERIFIER": lambda: bool( - int(os.getenv("VLLM_SM70_MXFP4_MOE_GROUPED_VERIFIER", "0")) - ), - # Group verifier slots routed to the same expert into one multi-row group. - # This changes the MXFP4 reduction tactic, so it stays behind an - # independent acceptance/quality gate. - "VLLM_SM70_MXFP4_MOE_GROUPED_M8_EXPERT_ROWS": lambda: bool( - int(os.getenv("VLLM_SM70_MXFP4_MOE_GROUPED_M8_EXPERT_ROWS", "0")) - ), # Use deterministic W13/W2 tactics for the fixed 48 one-row # verifier groups instead of relying on capture-time autotune stability. "VLLM_SM70_MXFP4_MOE_GROUPED_M8_FAST_SELECTOR": lambda: bool( @@ -2211,17 +1915,7 @@ 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", "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. 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", "1")) - ), - "VLLM_SM70_MXFP4_MOE_BROADCAST_INPUT_DECODE": lambda: bool( - int(os.getenv("VLLM_SM70_MXFP4_MOE_BROADCAST_INPUT_DECODE", "1")) + int(os.getenv("VLLM_SM70_MXFP4_MOE_DIRECT_TOP6_DECODE", "0")) ), # FP8 caller for the generic SM70 TurboMind active-source-group compact # decode path. The backend scheduler keeps source expert group semantics @@ -2311,26 +2005,10 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_MTP_PROFILE_INTERVAL": lambda: max( 1, int(os.getenv("VLLM_SM70_MTP_PROFILE_INTERVAL", "16")) ), - "VLLM_SM70_MTP_SPLIT_DRAFT_CUDAGRAPHS": lambda: bool( - int(os.getenv("VLLM_SM70_MTP_SPLIT_DRAFT_CUDAGRAPHS", "0")) - ), - # Compile alternate single/concurrent MTP helper signatures at startup. - # Default-off until matched cold-start and steady-state evidence is complete. - "VLLM_SM70_MTP_CONCURRENCY_WARMUP": lambda: bool( - int(os.getenv("VLLM_SM70_MTP_CONCURRENCY_WARMUP", "0")) - ), "VLLM_SM70_MTP_CONTEXT_BUCKETS": lambda: os.getenv("VLLM_SM70_MTP_CONTEXT_BUCKETS"), "VLLM_SM70_DSV4_DECODE_CONTEXT_BUCKETS": lambda: os.getenv( "VLLM_SM70_DSV4_DECODE_CONTEXT_BUCKETS" ), - # The private ring changes compressor-state ownership and lifetime. Keep it - # opt-in until its long-context quality and capacity gates are complete. - "VLLM_SM70_DSV4_PRIVATE_COMPRESSOR_STATE": lambda: bool( - int(os.getenv("VLLM_SM70_DSV4_PRIVATE_COMPRESSOR_STATE", "0")) - ), - "VLLM_SM70_FP8_KV_DECODE_CONTEXT_BUCKETS": lambda: os.getenv( - "VLLM_SM70_FP8_KV_DECODE_CONTEXT_BUCKETS" - ), "VLLM_SM70_MTP_CONTEXT_BUCKET_PARTITION_SIZE": lambda: os.getenv( "VLLM_SM70_MTP_CONTEXT_BUCKET_PARTITION_SIZE" ), @@ -2491,22 +2169,6 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3_MIN_KV": lambda: int( os.getenv("VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3_MIN_KV", "32768") ), - "VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3_Q8000_EXPERIMENTAL": lambda: bool( - int( - os.getenv( - "VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3_Q8000_EXPERIMENTAL", - "0", - ) - ) - ), - "VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL": lambda: bool( - int( - os.getenv( - "VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL", - "1", - ) - ) - ), "VLLM_FLASH_V100_PREFILL_SPLIT_KV": lambda: bool( int(os.getenv("VLLM_FLASH_V100_PREFILL_SPLIT_KV", "0")) ), @@ -2629,12 +2291,6 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_FLASH_V100_DECODE_USE_XQA": lambda: bool( 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")) - ), - "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_DECODE_XQA_Q4_MIN_SEQ_LEN": lambda: int( os.getenv("VLLM_FLASH_V100_DECODE_XQA_Q4_MIN_SEQ_LEN", "32768") ), @@ -2671,32 +2327,6 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_FLASH_V100_XQA_G6_QK_PIPELINE_TRACE": lambda: bool( int(os.getenv("VLLM_FLASH_V100_XQA_G6_QK_PIPELINE_TRACE", "0")) ), - "VLLM_FLASH_V100_XQA_G6_DUAL_CTA": lambda: bool( - int(os.getenv("VLLM_FLASH_V100_XQA_G6_DUAL_CTA", "0")) - ), - "VLLM_FLASH_V100_XQA_SPLIT_REDUCE": lambda: bool( - int(os.getenv("VLLM_FLASH_V100_XQA_SPLIT_REDUCE", "0")) - ), - # Exact SM70 G6/D256 E4M3 XQA for B2-B16. Set to 0 to restore the scalar - # paged decoder for matched control runs or emergency rollback. - "VLLM_FLASH_V100_E4M3_BATCH_XQA": lambda: bool( - int(os.getenv("VLLM_FLASH_V100_E4M3_BATCH_XQA", "1")) - ), - "VLLM_FLASH_V100_E4M3_BATCH_XQA_OPTIMIZED": lambda: bool( - int(os.getenv("VLLM_FLASH_V100_E4M3_BATCH_XQA_OPTIMIZED", "1")) - ), - "VLLM_FLASH_V100_E4M3_PAGE800_FASTPATH": lambda: bool( - int(os.getenv("VLLM_FLASH_V100_E4M3_PAGE800_FASTPATH", "1")) - ), - "VLLM_FLASH_V100_E4M3_PAGE800_FASTPATH_TRACE": lambda: bool( - int(os.getenv("VLLM_FLASH_V100_E4M3_PAGE800_FASTPATH_TRACE", "0")) - ), - "VLLM_FLASH_V100_XQA_BATCH_CONTEXT_ROUTING": lambda: bool( - int(os.getenv("VLLM_FLASH_V100_XQA_BATCH_CONTEXT_ROUTING", "1")) - ), - "VLLM_FLASH_V100_XQA_BATCH_CONTEXT_ROUTING_TRACE": lambda: bool( - int(os.getenv("VLLM_FLASH_V100_XQA_BATCH_CONTEXT_ROUTING_TRACE", "0")) - ), "VLLM_FLASH_V100_XQA_E5M2_G6_DUAL_CTA": lambda: bool( int(os.getenv("VLLM_FLASH_V100_XQA_E5M2_G6_DUAL_CTA", "1")) ), @@ -2712,9 +2342,6 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_FLASH_V100_XQA_E5M2_PAIR_LOAD": lambda: bool( int(os.getenv("VLLM_FLASH_V100_XQA_E5M2_PAIR_LOAD", "1")) ), - "VLLM_FLASH_V100_XQA_E5M2_BATCH_WIDE_LOAD": lambda: bool( - int(os.getenv("VLLM_FLASH_V100_XQA_E5M2_BATCH_WIDE_LOAD", "1")) - ), "VLLM_FLASH_V100_XQA_E5M2_G6_DUAL_CTA_TRACE": lambda: bool( int(os.getenv("VLLM_FLASH_V100_XQA_E5M2_G6_DUAL_CTA_TRACE", "0")) ), @@ -3048,14 +2675,6 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_QWEN_GDN_OUTPUT_PROJECTION_OP": lambda: bool( int(os.getenv("VLLM_SM70_QWEN_GDN_OUTPUT_PROJECTION_OP", "0")) ), - # Exact-shape SM70 GDN decode fusions. The measured workload is evidence, - # not a model-identity selector; runtime admission uses operator contracts. - "VLLM_SM70_GDN_QPN8_BA_SPLIT": lambda: bool( - int(os.getenv("VLLM_SM70_GDN_QPN8_BA_SPLIT", "0")) - ), - "VLLM_SM70_GDN_RMSNORM_ONEPASS": lambda: bool( - int(os.getenv("VLLM_SM70_GDN_RMSNORM_ONEPASS", "0")) - ), # Diagnostic-only: keep Qwen3.5/Gemma RMSNorm arithmetic behind an opaque # custom-op boundary under the SM70 compile/FULL graph lane. "VLLM_SM70_GEMMA_RMS_NORM_EAGER": lambda: bool( @@ -3129,11 +2748,6 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_UNQUANTIZED_MOE_0DOT3_CONFIG": lambda: bool( int(os.getenv("VLLM_SM70_UNQUANTIZED_MOE_0DOT3_CONFIG", "1")) ), - # Opt-in decode tiles for the exact E256/N128/K2048/top-k8 SM70 contract. - # Larger or unmatched token shapes retain the 0.0.3 config. - "VLLM_SM70_MTP_MOE_TUNED_CONFIG": lambda: bool( - int(os.getenv("VLLM_SM70_MTP_MOE_TUNED_CONFIG", "0")) - ), # Legacy SM70 CUDA-graph capture-size tuning from 0.0.3. Default-off # because dense capture can increase startup/compile cost; when enabled on # SM70 it restores the old fine-grained small-batch capture list. diff --git a/vllm/forward_context.py b/vllm/forward_context.py index 154a85dd03..bc42ef9104 100644 --- a/vllm/forward_context.py +++ b/vllm/forward_context.py @@ -19,9 +19,6 @@ logger = init_logger(__name__) -CUDAGRAPH_VARIANT_DEFAULT = 0 -CUDAGRAPH_VARIANT_LONG_CONTEXT = -1 - track_batchsize: bool = envs.VLLM_LOG_BATCHSIZE_INTERVAL >= 0 last_logging_time: float = 0 forward_start_time: float = 0 @@ -59,13 +56,11 @@ class BatchDescriptor: (like fused_moe_lora) whose grid size depends on num_active_loras to be properly captured. """ - graph_variant: int = CUDAGRAPH_VARIANT_DEFAULT + graph_variant: int = 0 """ Optional semantic specialization for CUDA graphs that share the same batch shape but execute different model branches. The default keeps existing graph - keys unchanged. Negative values are reserved for target-model attention - specializations; MTP drafter uses non-negative values to separate per-step - layer selection. + keys unchanged. MTP drafter uses this to separate per-step layer selection. """ attention_context_bucket: int | None = None """ diff --git a/vllm/model_executor/kernels/mhc/tilelang.py b/vllm/model_executor/kernels/mhc/tilelang.py index e591018d4f..de5f02d8ce 100644 --- a/vllm/model_executor/kernels/mhc/tilelang.py +++ b/vllm/model_executor/kernels/mhc/tilelang.py @@ -626,14 +626,12 @@ def mhc_fused_post_pre_tilelang( and use_fp16 and capability is not None and capability.to_int() == 70 - and 1 <= num_tokens <= 8 + and num_tokens == 1 and hidden_size == 4096 and hc_mult == 4 and hc_mult3 == 24 - and ( - (num_tokens < 8 and tile_n == 2 and n_splits == 8) - or (num_tokens == 8 and tile_n == 3 and n_splits == 4) - ) + and tile_n == 2 + and n_splits == 8 ) gemm_out_mul = torch.empty( diff --git a/vllm/model_executor/kernels/mhc/tilelang_kernels.py b/vllm/model_executor/kernels/mhc/tilelang_kernels.py index bf4a94e556..777804072f 100644 --- a/vllm/model_executor/kernels/mhc/tilelang_kernels.py +++ b/vllm/model_executor/kernels/mhc/tilelang_kernels.py @@ -539,7 +539,7 @@ def sm70_mhc_post_fp32_stage_tilelang( n_thr: int = 256, n_splits: int = 8, ) -> tilelang.JITKernel: - """Materialize a small-token post mapping without rounding the dot input.""" + """Materialize the M=1 post mapping without rounding the dot input.""" num_tokens = T.dynamic("num_tokens") h_per_split = hidden // n_splits h_iters = h_per_split // n_thr @@ -608,7 +608,7 @@ def sm70_mhc_dot_from_fp32_stage_tilelang( tile_n: int = 2, n_splits: int = 8, ) -> tilelang.JITKernel: - """Run small-token dot accumulation from the exact FP32 post stage.""" + """Run the M=1 dot accumulation from the exact FP32 post stage.""" num_tokens = T.dynamic("num_tokens") h_per_split = hidden // n_splits h_iters = h_per_split // n_thr diff --git a/vllm/model_executor/layers/attention/attention.py b/vllm/model_executor/layers/attention/attention.py index bfedffcffd..9a85a75ddf 100644 --- a/vllm/model_executor/layers/attention/attention.py +++ b/vllm/model_executor/layers/attention/attention.py @@ -41,8 +41,6 @@ from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheSpec, - KVQuantMode, - PrefixAnchoredSWASpec, SlidingWindowSpec, get_kv_quant_mode, ) @@ -374,42 +372,6 @@ def __init__( f"but got {self.attn_backend.get_name()}." ) - prefix_anchored_decode_window = getattr( - vllm_config.attention_config, - "prefix_anchored_decode_window", - None, - ) - if ( - prefix_anchored_decode_window is not None - and attn_type == AttentionType.DECODER - ): - if self.attn_backend.get_name() != "FLASH_ATTN_V100": - raise ValueError( - "prefix-anchored SWA requires standard full decoder " - "attention on FLASH_ATTN_V100" - ) - if self.sliding_window is not None: - raise ValueError( - "prefix-anchored SWA cannot be combined with per-layer " - "sliding-window attention" - ) - if ( - get_kv_quant_mode(self.kv_cache_dtype) != KVQuantMode.NONE - or self.kv_cache_torch_dtype != torch.float16 - ): - raise ValueError("prefix-anchored SWA requires an fp16 KV cache") - if self.head_size_v != self.head_size: - raise ValueError( - "prefix-anchored SWA requires identical key/value head sizes" - ) - if self.has_sink: - raise ValueError( - "prefix-anchored SWA cannot be combined with attention sinks" - ) - extra_impl_args["prefix_anchored_decode_window"] = ( - prefix_anchored_decode_window - ) - if self.attn_backend.get_name() == "FLEX_ATTENTION": block_m = vllm_config.attention_config.flex_attn_block_m block_n = vllm_config.attention_config.flex_attn_block_n @@ -639,45 +601,6 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: # Should not be called for enc-dec or encoder-only attention. assert self.attn_type == AttentionType.DECODER quant_mode = get_kv_quant_mode(self.kv_cache_dtype) - anchored_window = getattr( - vllm_config.attention_config, - "prefix_anchored_decode_window", - None, - ) - if anchored_window is not None: - if self.attn_backend.get_name() != "FLASH_ATTN_V100": - raise ValueError( - "prefix-anchored SWA requires the FLASH_ATTN_V100 backend" - ) - if self.sliding_window is not None: - raise ValueError( - "prefix-anchored SWA cannot be combined with per-layer " - "sliding-window attention" - ) - if ( - quant_mode != KVQuantMode.NONE - or self.kv_cache_torch_dtype != torch.float16 - ): - raise ValueError("prefix-anchored SWA requires an fp16 KV cache") - if self.kv_cache_dtype.startswith("turboquant_"): - raise ValueError("prefix-anchored SWA does not support TurboQuant KV") - if self.head_size_v != self.head_size: - raise ValueError( - "prefix-anchored SWA requires identical key/value head sizes" - ) - if self.has_sink: - raise ValueError( - "prefix-anchored SWA cannot be combined with attention sinks" - ) - return PrefixAnchoredSWASpec( - block_size=block_size, - num_kv_heads=self.num_kv_heads, - head_size=self.head_size, - head_size_v=self.head_size_v, - dtype=self.kv_cache_torch_dtype, - kv_quant_mode=quant_mode, - decode_sliding_window=anchored_window, - ) if self.sliding_window is not None: assert not vllm_config.model_config.use_mla, ( "MLA is not supported for slidingwindow" 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 f86223b07c..d0735721da 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,9 +45,7 @@ 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 fecfa3a4e7..87708d2675 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,7 +76,8 @@ 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 35491a71e0..ed1d1a1ef4 100644 --- a/vllm/model_executor/layers/fla/ops/fused_recurrent.py +++ b/vllm/model_executor/layers/fla/ops/fused_recurrent.py @@ -56,7 +56,9 @@ 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") diff --git a/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py b/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py index 49f1e3aad1..83f599b62b 100644 --- a/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py +++ b/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py @@ -13,8 +13,6 @@ from vllm.triton_utils import tl, triton -from .op import exp - def _parse_positive_int_env(name: str) -> int | None: value = os.getenv(name) @@ -85,37 +83,6 @@ def _select_fused_sigmoid_schedule( return BV, num_warps, num_stages -def _select_fused_sigmoid_launch( - V: int, - N: int, - HV: int, - T: int, - device: torch.device, - *, - match_recurrent_schedule: bool, -) -> tuple[int, int, int]: - if not _use_sm70_fused_sigmoid_schedule(device): - return min(triton.next_power_of_2(V), 32), 4, 3 - if not match_recurrent_schedule: - return _select_fused_sigmoid_schedule(V, N, HV, device) - - # The unfused verifier uses this launch geometry. Keeping it for the first - # fused rollout removes schedule-induced reduction-order changes from the - # correctness comparison; schedule tuning is a separate performance gate. - from .fused_recurrent import ( - _select_sm70_bv, - _select_sm70_num_stages, - _select_sm70_num_warps, - ) - - BV = _select_sm70_bv(V, N, HV, device) - return ( - BV, - _select_sm70_num_warps(BV, N, HV), - _select_sm70_num_stages(T), - ) - - @triton.heuristics( { "USE_INITIAL_STATE": lambda args: args["h0"] is not None, @@ -173,10 +140,6 @@ def fused_sigmoid_gating_delta_rule_update_kernel( IS_DDTREE: tl.constexpr, IS_KDA: tl.constexpr, MIXED_QKV: tl.constexpr, - PRECOMPUTED_GATING: tl.constexpr, - QUANTIZE_BETA: tl.constexpr, - QUANTIZE_STATE_EACH_STEP: tl.constexpr, - MATCH_RECURRENT_NUMERICS: 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 @@ -274,44 +237,25 @@ def fused_sigmoid_gating_delta_rule_update_kernel( b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) - if PRECOMPUTED_GATING: - b_g = tl.load(p_a).to(tl.float32) - b_beta = tl.load(p_b).to(tl.float32) - else: - b_b = tl.load(p_b).to(tl.float32) + b_b = tl.load(p_b).to(tl.float32) - # If the model is loaded in fp16, without the .float() here, A - # might be -inf. - x = tl.load(p_a).to(tl.float32) + tl.load(p_dt_bias).to(tl.float32) - softplus_x = tl.where( - beta * x <= threshold, (1 / beta) * tl.log(1 + tl.exp(beta * x)), x - ) - b_g = -tl.exp(tl.load(p_A_log).to(tl.float32)) * softplus_x + # If the model is loaded in fp16, without the .float() here, A might be -inf + x = tl.load(p_a).to(tl.float32) + tl.load(p_dt_bias).to(tl.float32) + softplus_x = tl.where( + beta * x <= threshold, (1 / beta) * tl.log(1 + tl.exp(beta * x)), x + ) + b_g = -tl.exp(tl.load(p_A_log).to(tl.float32)) * softplus_x - # compute beta_output = sigmoid(b) - b_beta = tl.sigmoid(b_b.to(tl.float32)) - if QUANTIZE_BETA: - # The split GDN verifier materializes beta in b.dtype before - # the recurrent kernel consumes it. Preserve that boundary. - b_beta = b_beta.to(p_b.dtype.element_ty).to(tl.float32) + # compute beta_output = sigmoid(b) + b_beta = tl.sigmoid(b_b.to(tl.float32)) if USE_QK_L2NORM_IN_KERNEL: - if MATCH_RECURRENT_NUMERICS: - # Preserve the split recurrent verifier's operation order. - # Multiplication by rsqrt is mathematically equivalent but can - # move a probabilistic rejection decision by one FP16 ULP. - b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6) - b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6) - else: - b_q = b_q * (tl.rsqrt(tl.sum(b_q * b_q) + 1e-6)) - b_k = b_k * (tl.rsqrt(tl.sum(b_k * b_k) + 1e-6)) + b_q = b_q * (tl.rsqrt(tl.sum(b_q * b_q) + 1e-6)) + b_k = b_k * (tl.rsqrt(tl.sum(b_k * b_k) + 1e-6)) b_q = b_q * scale # [BV, BK] if not IS_KDA: - if MATCH_RECURRENT_NUMERICS: - b_h *= exp(b_g) - else: - b_h *= tl.exp(b_g) + b_h *= tl.exp(b_g) else: b_h *= tl.exp(b_g[None, :]) # [BV] @@ -333,17 +277,10 @@ def fused_sigmoid_gating_delta_rule_update_kernel( p_ht = ht + final_state_idx * stride_final_state_token p_ht = p_ht + i_hv * V * K + o_v[:, None] * K + o_k[None, :] tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) - if QUANTIZE_STATE_EACH_STEP: - # Target verification handles several linear tokens in one - # invocation. Reloading the just-persisted state reproduces - # the low-precision store/load boundary of repeated decode. - b_h = tl.load(p_ht, mask=mask_h, other=0).to(tl.float32) else: p_ht = ht + (bos + i_t) * stride_final_state_token p_ht = p_ht + i_hv * V * K + o_v[:, None] * K + o_k[None, :] tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) - if QUANTIZE_STATE_EACH_STEP: - b_h = tl.load(p_ht, mask=mask_h, other=0).to(tl.float32) if MIXED_QKV: p_q += QKV_STRIDE @@ -473,10 +410,6 @@ def fused_sigmoid_gating_delta_rule_update( USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, IS_KDA=is_kda, MIXED_QKV=False, - PRECOMPUTED_GATING=False, - QUANTIZE_BETA=False, - QUANTIZE_STATE_EACH_STEP=False, - MATCH_RECURRENT_NUMERICS=False, num_warps=num_warps, num_stages=num_stages, ) @@ -608,10 +541,6 @@ def fused_sigmoid_gating_delta_rule_update_mixed_qkv( USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, IS_KDA=False, MIXED_QKV=True, - PRECOMPUTED_GATING=False, - QUANTIZE_BETA=False, - QUANTIZE_STATE_EACH_STEP=False, - MATCH_RECURRENT_NUMERICS=False, num_warps=num_warps, num_stages=num_stages, ) @@ -635,21 +564,9 @@ def fused_sigmoid_gating_delta_rule_update_mixed_qkv_out( initial_state: torch.Tensor | None = None, cu_seqlens: torch.Tensor | None = None, ssm_state_indices: torch.Tensor | None = None, - num_accepted_tokens: torch.Tensor | None = None, use_qk_l2norm_in_kernel: bool = False, - quantize_beta: bool = False, - quantize_state_each_step: bool = False, - match_recurrent_schedule: bool = False, - match_recurrent_numerics: bool = False, - precomputed_g: torch.Tensor | None = None, - precomputed_beta: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - """Mixed-QKV update that writes into a caller-provided output buffer. - - ``precomputed_g`` and ``precomputed_beta`` retain the split verifier's - exact gating materialization while still removing the packed-QKV rearrange. - Omitting both computes gating inside the recurrent kernel. - """ + """Mixed-QKV fused update that writes into a caller-provided decode output.""" if mixed_qkv.ndim != 2: raise ValueError("mixed_qkv must have shape [T, qkv_hidden].") if not mixed_qkv.is_contiguous(): @@ -658,10 +575,6 @@ def fused_sigmoid_gating_delta_rule_update_mixed_qkv_out( raise ValueError("cu_seqlens is required for mixed_qkv_out.") if ssm_state_indices is None: raise ValueError("ssm_state_indices is required for mixed_qkv_out.") - if (precomputed_g is None) != (precomputed_beta is None): - raise ValueError( - "precomputed_g and precomputed_beta must be provided together." - ) T = mixed_qkv.shape[0] N = cu_seqlens.numel() - 1 @@ -684,28 +597,6 @@ def fused_sigmoid_gating_delta_rule_update_mixed_qkv_out( else: assert scale > 0, "scale must be positive" - use_precomputed_gating = precomputed_g is not None - if use_precomputed_gating: - assert precomputed_g is not None - assert precomputed_beta is not None - if precomputed_g.numel() != T * HV: - raise ValueError( - f"precomputed_g must contain {T * HV} values, " - f"got {precomputed_g.numel()}." - ) - if precomputed_beta.numel() != T * HV: - raise ValueError( - f"precomputed_beta must contain {T * HV} values, " - f"got {precomputed_beta.numel()}." - ) - if quantize_beta: - raise ValueError("precomputed_beta is already materialized in its dtype.") - kernel_a = precomputed_g.contiguous().view(T, HV) - kernel_b = precomputed_beta.contiguous().view(T, HV) - else: - kernel_a = a.contiguous() - kernel_b = b.contiguous() - assert initial_state is not None final_state = initial_state stride_init_state_token = initial_state.stride(0) @@ -717,21 +608,18 @@ def fused_sigmoid_gating_delta_rule_update_mixed_qkv_out( stride_indices_seq, stride_indices_tok = ssm_state_indices.stride() BK = triton.next_power_of_2(K) - BV, num_warps, num_stages = _select_fused_sigmoid_launch( - V, - N, - HV, - T, - mixed_qkv.device, - match_recurrent_schedule=match_recurrent_schedule, + BV, num_warps, num_stages = ( + _select_fused_sigmoid_schedule(V, N, HV, mixed_qkv.device) + if _use_sm70_fused_sigmoid_schedule(mixed_qkv.device) + else (min(triton.next_power_of_2(V), 32), 4, 3) ) NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) assert NK == 1, "NK > 1 is not supported yet" fused_sigmoid_gating_delta_rule_update_kernel[(NK, NV, N * HV)]( A_log=A_log, - a=kernel_a, - b=kernel_b, + a=a.contiguous(), + b=b.contiguous(), dt_bias=dt_bias, beta=beta, threshold=threshold, @@ -744,7 +632,7 @@ def fused_sigmoid_gating_delta_rule_update_mixed_qkv_out( ht=final_state, cu_seqlens=cu_seqlens, ssm_state_indices=ssm_state_indices, - num_accepted_tokens=num_accepted_tokens, + num_accepted_tokens=None, ddtree_parent_ids=None, scale=scale, N=N, @@ -770,10 +658,6 @@ def fused_sigmoid_gating_delta_rule_update_mixed_qkv_out( USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, IS_KDA=False, MIXED_QKV=True, - PRECOMPUTED_GATING=use_precomputed_gating, - QUANTIZE_BETA=quantize_beta, - QUANTIZE_STATE_EACH_STEP=quantize_state_each_step, - MATCH_RECURRENT_NUMERICS=match_recurrent_numerics, num_warps=num_warps, num_stages=num_stages, ) diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py index 922309a5dc..fb94ed46a1 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe.py @@ -5,8 +5,6 @@ import functools import json import os -from collections.abc import Iterator -from contextlib import contextmanager from typing import Any import torch @@ -36,20 +34,6 @@ logger = init_logger(__name__) -_force_sm70_mtp_moe_legacy_config = False - - -@contextmanager -def force_sm70_mtp_moe_legacy_config() -> Iterator[None]: - """Temporarily select the legacy SM70 MoE tile during startup warmup.""" - global _force_sm70_mtp_moe_legacy_config - previous = _force_sm70_mtp_moe_legacy_config - _force_sm70_mtp_moe_legacy_config = True - try: - yield - finally: - _force_sm70_mtp_moe_legacy_config = previous - @triton.jit def write_zeros_to_output( @@ -1216,32 +1200,6 @@ def should_moe_wna16_use_cuda( ) -def _get_sm70_mtp_moe_decode_config( - M: int, - E: int, - N: int, - K: int, - topk: int, -) -> dict[str, int] | None: - """Return the graph-tuned exact-shape SM70 MTP tile for M2-M16.""" - if _force_sm70_mtp_moe_legacy_config or not envs.VLLM_SM70_MTP_MOE_TUNED_CONFIG: - return None - if (E, N, K, topk) != (256, 128, 2048, 8) or not 2 <= M <= 16: - return None - - # TP4 shards the checkpoint-global I512 expert width to I128 per rank. - # M2-M16 all select this tile in the exact local-shape graph oracle. - return { - "BLOCK_SIZE_M": 8, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 32, - "GROUP_SIZE_M": 1, - "SPLIT_K": 1, - "num_warps": 4, - "num_stages": 3, - } - - def get_default_config( M: int, E: int, @@ -1304,11 +1262,7 @@ def get_default_config( # upstream's larger small-batch tiles are poor for single-token SM70 # decode, which is the FP8-MoE-dequant fallback route used by the # old 35B-FP8 baseline. - sm70_mtp_config = _get_sm70_mtp_moe_decode_config(M, E, N, K, topk) - if sm70_mtp_config is not None: - config = sm70_mtp_config - logger.info_once(f"Using tuned exact-shape SM70 MTP MoE config: {config}") - elif M <= E: + if M <= E: config = { "BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 32, @@ -1324,7 +1278,9 @@ def get_default_config( "GROUP_SIZE_M": 8, "SPLIT_K": 1, } - logger.info_once(f"Using SM70 0.0.3 unquantized MoE default config: {config}") + logger.info_once( + f"Using SM70 0.0.3 unquantized MoE default config: {config}" + ) else: # General defaults for bf16/fp16 and fp8 per-tensor. # Tile sizes scale with batch: small batches are memory-bound 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 4b485202ad..414bc561d1 100644 --- a/vllm/model_executor/layers/fused_moe/router/gate_linear.py +++ b/vllm/model_executor/layers/fused_moe/router/gate_linear.py @@ -57,7 +57,6 @@ 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 @@ -96,19 +95,12 @@ def forward( ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]: import vllm._custom_ops as ops - if ( - envs.VLLM_SM70_DSV4_FP16_GEMV or envs.VLLM_SM70_DSV4_FP13_GEMV - ) and self.out_dtype is not None: + if envs.VLLM_SM70_DSV4_FP16_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, - getattr(self, "_sm70_dsv4_fp13_weight", None), - ) + output = maybe_sm70_dsv4_fp16_gemv(x, self.weight, self.out_dtype) if output is not None: return output, None diff --git a/vllm/model_executor/layers/layernorm.py b/vllm/model_executor/layers/layernorm.py index 0106f989bc..0a72c8a363 100644 --- a/vllm/model_executor/layers/layernorm.py +++ b/vllm/model_executor/layers/layernorm.py @@ -14,100 +14,11 @@ from vllm.model_executor.custom_op import CustomOp from vllm.model_executor.layers.batch_invariant import rms_norm_batch_invariant from vllm.platforms import current_platform -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_dflash2_gemma_fused_add_rms_kernel( - x, - residual, - weight, - normalized_out, - residual_out, - hidden_size: tl.constexpr, - BLOCK_SIZE: tl.constexpr, - epsilon, -): - row = tl.program_id(0) - cols = tl.arange(0, BLOCK_SIZE) - mask = cols < hidden_size - values = tl.load(x + row * hidden_size + cols, mask=mask, other=0.0).to(tl.float32) - values += tl.load(residual + row * hidden_size + cols, mask=mask, other=0.0).to( - tl.float32 - ) - tl.store(residual_out + row * hidden_size + cols, values, mask=mask) - - variance = tl.sum(tl.where(mask, values * values, 0.0), axis=0) - inverse_rms = tl.rsqrt(variance / hidden_size + epsilon) - gemma_weight = tl.load(weight + cols, mask=mask, other=0.0).to(tl.float32) + 1.0 - tl.store( - normalized_out + row * hidden_size + cols, - values * inverse_rms * gemma_weight, - mask=mask, - ) - - -def _sm70_dflash2_gemma_fused_add_rms_norm( - x: torch.Tensor, - residual: torch.Tensor, - weight: torch.Tensor, - variance_epsilon: float, - *, - num_warps: int = 8, -) -> tuple[torch.Tensor, torch.Tensor]: - normalized_out = torch.empty_like(x) - residual_out = torch.empty_like(residual) - _sm70_dflash2_gemma_fused_add_rms_kernel[(x.shape[0],)]( - x, - residual, - weight, - normalized_out, - residual_out, - hidden_size=x.shape[1], - BLOCK_SIZE=triton.next_power_of_2(x.shape[1]), - epsilon=variance_epsilon, - num_warps=num_warps, - num_stages=1, - ) - return normalized_out, residual_out - - -def _use_sm70_dflash2_gemma_fused_add_rms( - x: torch.Tensor, - residual: torch.Tensor | None, - weight: torch.Tensor, -) -> bool: - # Keep the dynamic token dimension out of this Python predicate. AOT traces - # the target once at a large warmup shape; a decode-only row bound would be - # constant-folded there and would leave the M=8 replay on the decomposed path. - return bool( - envs.VLLM_SM70_DFLASH2_FUSED_GEMMA_RMS - and envs.VLLM_SM70_FLASH_V100_0DOT3_COMPILE_GRAPH - # Keep the cached platform query out of the AOT fullgraph. The helper - # below is explicitly constant-foldable by Dynamo. - and _sm70_gemma_long_prefill_available() - and residual is not None - and x.is_cuda - and x.dtype == torch.float16 - and residual.dtype == torch.float32 - and weight.dtype in (torch.float16, torch.bfloat16, torch.float32) - and x.ndim == 2 - and x.shape[0] > 0 - and x.shape[1] == 5120 - and residual.shape == x.shape - and residual.device == x.device - and weight.device == x.device - and weight.ndim == 1 - and weight.numel() == 5120 - and x.is_contiguous() - and residual.is_contiguous() - and weight.is_contiguous() - ) - - @torch.compiler.assume_constant_result def _sm70_gemma_long_prefill_available() -> bool: return current_platform.is_device_capability(70) @@ -427,14 +338,6 @@ def forward_native( residual: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """PyTorch-native implementation equivalent to forward().""" - if _use_sm70_dflash2_gemma_fused_add_rms(x, residual, self.weight): - assert residual is not None - return _sm70_dflash2_gemma_fused_add_rms_norm( - x, - residual, - self.weight, - self.variance_epsilon, - ) if self._use_sm70_long_prefill_fused(x, residual): assert residual is not None if not torch.compiler.is_compiling(): @@ -481,14 +384,6 @@ def forward_cuda( x: torch.Tensor, residual: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - if _use_sm70_dflash2_gemma_fused_add_rms(x, residual, self.weight): - assert residual is not None - return _sm70_dflash2_gemma_fused_add_rms_norm( - x, - residual, - self.weight, - self.variance_epsilon, - ) if self._use_sm70_long_prefill_fused(x, residual): assert residual is not None if not torch.compiler.is_compiling(): diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index 68f99ee3eb..3632f22c47 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -64,9 +64,6 @@ def _maybe_sm70_dense_forward( return None if not hasattr(torch.ops._C, "sm70_f16_gemm"): return None - max_m = getattr(layer, "_sm70_f16_max_m", None) - if max_m is not None and x.numel() // x.shape[-1] > max_m: - return None if envs.VLLM_SM70_F16_DENSE_DEBUG or envs.VLLM_QWEN3_NEXT_SM70_TRACE: rows = x.numel() // x.shape[-1] @@ -369,15 +366,6 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: if not current_platform.is_cuda_alike(): 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 @@ -1805,10 +1793,14 @@ def forward( if output is None: output_parallel = _maybe_sm70_dense_forward(self, input_parallel, bias_) if output_parallel is None: - output_parallel = self.quant_method.apply(self, input_parallel, bias_) + output_parallel = self.quant_method.apply( + self, input_parallel, bias_ + ) if self.reduce_results and self.tp_size > 1: - output = _maybe_sm70_awq_mlp_down_tile_all_reduce(self, output_parallel) + output = _maybe_sm70_awq_mlp_down_tile_all_reduce( + self, output_parallel + ) if output is None: output = tensor_model_parallel_all_reduce(output_parallel) else: diff --git a/vllm/model_executor/layers/logits_processor.py b/vllm/model_executor/layers/logits_processor.py index 8bf4ef2a28..1b49aa4f6e 100644 --- a/vllm/model_executor/layers/logits_processor.py +++ b/vllm/model_executor/layers/logits_processor.py @@ -151,9 +151,8 @@ def _cuda_graph_capture_active() -> bool: return False -def _maybe_sync_top1_all_gather( - module: torch.nn.Module, local_pair: torch.Tensor -) -> None: +def _maybe_sync_top1_all_gather(module: torch.nn.Module, + local_pair: torch.Tensor) -> None: raw_steps = envs.VLLM_SM70_SYNC_TOP1_ALLGATHER_STEPS if not raw_steps or not local_pair.is_cuda: return @@ -257,8 +256,8 @@ def _get_logits( pre_stream_sync_ms = 0.0 pre_device_extra_sync_ms = 0.0 if profile_enabled: - pre_stream_sync_ms, pre_device_extra_sync_ms = _cuda_sync_breakdown_ms( - hidden_states + pre_stream_sync_ms, pre_device_extra_sync_ms = ( + _cuda_sync_breakdown_ms(hidden_states) ) pre_sync_ms = pre_stream_sync_ms + pre_device_extra_sync_ms @@ -316,7 +315,9 @@ def get_top_tokens( ) tp_size = get_tensor_model_parallel_world_size() - local_top1 = lm_head.maybe_get_sm70_lm_head_top1(hidden_states, embedding_bias) + local_top1 = lm_head.maybe_get_sm70_lm_head_top1( + hidden_states, embedding_bias + ) if local_top1 is None: logits = lm_head.quant_method.apply( lm_head, hidden_states, bias=embedding_bias @@ -395,12 +396,14 @@ def get_topk_tokens_and_logprobs( pre_stream_sync_ms = 0.0 pre_device_extra_sync_ms = 0.0 if profile_enabled: - pre_stream_sync_ms, pre_device_extra_sync_ms = _cuda_sync_breakdown_ms( - hidden_states + pre_stream_sync_ms, pre_device_extra_sync_ms = ( + _cuda_sync_breakdown_ms(hidden_states) ) pre_sync_ms = pre_stream_sync_ms + pre_device_extra_sync_ms stage_start = _cuda_stage_start(profile_enabled) - logits = lm_head.quant_method.apply(lm_head, hidden_states, bias=embedding_bias) + logits = lm_head.quant_method.apply( + lm_head, hidden_states, bias=embedding_bias + ) local_lm_head_ms = _cuda_stage_ms(profile_enabled, stage_start) stage_start = _cuda_stage_start(profile_enabled) @@ -527,66 +530,43 @@ def get_topk_tokens_and_logits( pre_stream_sync_ms = 0.0 pre_device_extra_sync_ms = 0.0 if profile_enabled: - pre_stream_sync_ms, pre_device_extra_sync_ms = _cuda_sync_breakdown_ms( - hidden_states + pre_stream_sync_ms, pre_device_extra_sync_ms = ( + _cuda_sync_breakdown_ms(hidden_states) ) pre_sync_ms = pre_stream_sync_ms + pre_device_extra_sync_ms stage_start = _cuda_stage_start(profile_enabled) - local_candidates = lm_head.maybe_get_sm70_dflash2_top20( - hidden_states, - top_k, - embedding_bias, + logits = lm_head.quant_method.apply( + lm_head, hidden_states, bias=embedding_bias ) - if local_candidates is None: - logits = lm_head.quant_method.apply( - lm_head, hidden_states, bias=embedding_bias - ) - local_lm_head_ms = _cuda_stage_ms(profile_enabled, stage_start) + local_lm_head_ms = _cuda_stage_ms(profile_enabled, stage_start) - stage_start = _cuda_stage_start(profile_enabled) - if self.soft_cap is not None: - logits = torch.tanh(logits / self.soft_cap) * self.soft_cap - if self.scale != 1.0: - logits = logits * self.scale + stage_start = _cuda_stage_start(profile_enabled) + if self.soft_cap is not None: + logits = torch.tanh(logits / self.soft_cap) * self.soft_cap + if self.scale != 1.0: + logits = logits * self.scale - num_pad = lm_head.shard_indices.num_org_vocab_padding - if num_pad > 0: - logits[..., -num_pad:] = -float("inf") + num_pad = lm_head.shard_indices.num_org_vocab_padding + if num_pad > 0: + logits[..., -num_pad:] = -float("inf") - logits_float = logits.float() - postprocess_ms = _cuda_stage_ms(profile_enabled, stage_start) + logits_float = logits.float() + postprocess_ms = _cuda_stage_ms(profile_enabled, stage_start) - local_k = min(top_k, logits_float.shape[-1]) - stage_start = _cuda_stage_start(profile_enabled) - local_vals, local_indices = torch.topk( - logits_float, - k=local_k, - dim=-1, - ) - local_topk_ms = _cuda_stage_ms(profile_enabled, stage_start) + local_k = min(top_k, logits_float.shape[-1]) + stage_start = _cuda_stage_start(profile_enabled) + local_vals, local_indices = torch.topk( + logits_float, + k=local_k, + dim=-1, + ) + local_topk_ms = _cuda_stage_ms(profile_enabled, stage_start) - stage_start = _cuda_stage_start(profile_enabled) - vocab_start = lm_head.shard_indices.org_vocab_start_index - local_global_indices = local_indices + vocab_start - local_indices_ms = _cuda_stage_ms(profile_enabled, stage_start) - local_vocab_size = int(logits_float.shape[-1]) - else: - local_vals, local_global_indices = local_candidates - local_lm_head_ms = _cuda_stage_ms(profile_enabled, stage_start) - stage_start = _cuda_stage_start(profile_enabled) - # Positive scale and tanh soft-capping are monotonic, so applying - # them after the exact local top-k preserves the dense ordering. - if self.soft_cap is not None: - local_vals = torch.tanh(local_vals / self.soft_cap) * self.soft_cap - if self.scale != 1.0: - local_vals = local_vals * self.scale - local_vals = local_vals.float() - postprocess_ms = _cuda_stage_ms(profile_enabled, stage_start) - local_k = int(local_vals.shape[-1]) - local_topk_ms = 0.0 - local_indices_ms = 0.0 - local_vocab_size = int(lm_head.weight.shape[0]) + stage_start = _cuda_stage_start(profile_enabled) + vocab_start = lm_head.shard_indices.org_vocab_start_index + local_global_indices = local_indices + vocab_start + local_indices_ms = _cuda_stage_ms(profile_enabled, stage_start) tp_size = get_tensor_model_parallel_world_size() gather_vals_ms = 0.0 @@ -635,7 +615,7 @@ def get_topk_tokens_and_logits( gather_indices_ms, merge_topk_ms, int(hidden_states.shape[0]), - local_vocab_size, + int(logits_float.shape[-1]), top_k, local_k, tp_size, 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..c28dd507f6 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 @@ -12,7 +12,6 @@ from einops import rearrange from torch import nn -from vllm import _sm70_ops as sm70_ops from vllm import envs from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig, get_current_vllm_config @@ -75,7 +74,6 @@ get_registered_gdn_spec_metadata_tensors, ) from vllm.v1.attention.backends.utils import compute_causal_conv1d_metadata -from vllm.v1.worker.gpu.spec_decode import uses_dflash_selector_engine # Optional ROCm AITER Triton kernels for the GDN decode fast-path. # Availability is checked centrally via rocm_aiter_ops; the actual function @@ -105,80 +103,6 @@ _DFLASH_DDTREE_PATH_PROBE_REPORTS = 0 -@triton.jit -def _sm70_pack_qwen_gdn_qkv_kernel( - mixed_qkv, - packed, - input_row_stride: tl.int64, - num_rows: tl.constexpr, - q_dim: tl.constexpr, - k_dim: tl.constexpr, - v_dim: tl.constexpr, - BLOCK_Q: tl.constexpr, - BLOCK_K: tl.constexpr, - BLOCK_V: tl.constexpr, -): - """Bitwise row-slice copy into [all Q][all K][all V] storage.""" - row = tl.program_id(0) - input_row = mixed_qkv + row * input_row_stride - - q_cols = tl.arange(0, BLOCK_Q) - q_mask = q_cols < q_dim - q = tl.load(input_row + q_cols, mask=q_mask) - tl.store(packed + row * q_dim + q_cols, q, mask=q_mask) - - k_cols = tl.arange(0, BLOCK_K) - k_mask = k_cols < k_dim - k = tl.load(input_row + q_dim + k_cols, mask=k_mask) - q_numel = num_rows * q_dim - tl.store(packed + q_numel + row * k_dim + k_cols, k, mask=k_mask) - - v_cols = tl.arange(0, BLOCK_V) - v_mask = v_cols < v_dim - v = tl.load(input_row + q_dim + k_dim + v_cols, mask=v_mask) - qk_numel = q_numel + num_rows * k_dim - tl.store(packed + qk_numel + row * v_dim + v_cols, v, mask=v_mask) - - -def _sm70_pack_qwen_gdn_qkv( - mixed_qkv: torch.Tensor, - q_dim: int, - k_dim: int, - v_dim: int, -) -> torch.Tensor: - """Materialize the recurrent Q/K/V layout with one copy launch.""" - if mixed_qkv.ndim != 2: - raise ValueError("mixed_qkv must be rank two") - if mixed_qkv.stride(1) != 1: - raise ValueError("mixed_qkv must be contiguous by row") - if min(q_dim, k_dim, v_dim) <= 0: - raise ValueError("Q/K/V widths must be positive") - if mixed_qkv.shape[1] < q_dim + k_dim + v_dim: - raise ValueError("mixed_qkv is narrower than the Q/K/V widths") - - num_rows = mixed_qkv.shape[0] - packed = torch.empty( - num_rows * (q_dim + k_dim + v_dim), - dtype=mixed_qkv.dtype, - device=mixed_qkv.device, - ) - _sm70_pack_qwen_gdn_qkv_kernel[(num_rows,)]( - mixed_qkv, - packed, - mixed_qkv.stride(0), - num_rows=num_rows, - q_dim=q_dim, - k_dim=k_dim, - v_dim=v_dim, - BLOCK_Q=triton.next_power_of_2(q_dim), - BLOCK_K=triton.next_power_of_2(k_dim), - BLOCK_V=triton.next_power_of_2(v_dim), - num_warps=8, - num_stages=1, - ) - return packed - - def _ddtree_parent_ids_require_branch( parent_ids: torch.Tensor | None, num_tree_tokens_cpu: torch.Tensor | None, @@ -503,20 +427,12 @@ def _sm70_dump_gdn_core_tensor( tensor: torch.Tensor, source: str = "core", ) -> None: - dump_dir = os.getenv("VLLM_SM70_DUMP_GDN_CORE_DIR") - graph_dump = os.getenv("VLLM_SM70_DUMP_GDN_GRAPH_BUFFERS") == "1" and bool( - os.getenv("VLLM_SM70_DUMP_GDN_GRAPH_DIR") - ) - # Keep the production path free of CUDA runtime queries. This hook runs - # several times per GDN layer, while core dumps are normally disabled. - if not dump_dir and not graph_dump: - return - if graph_dump: - _sm70_gdn_graph_buffer_copy(label, layer_name, tensor, source) - if not dump_dir or not tensor.is_cuda: - return + _sm70_gdn_graph_buffer_copy(label, layer_name, tensor, source) if torch.cuda.is_current_stream_capturing(): return + dump_dir = os.getenv("VLLM_SM70_DUMP_GDN_CORE_DIR") + if not dump_dir: + return raw_layer_ids = os.getenv("VLLM_SM70_DUMP_GDN_CORE_LAYER_IDS") if raw_layer_ids: layer_idx = _sm70_gdn_layer_idx(layer_name) @@ -674,10 +590,6 @@ def _sm70_qwen_gdn_spec_method(vllm_config: object) -> str | None: return getattr(spec_config, "method", None) -def _is_dflash2_spec_config(vllm_config: object) -> bool: - return uses_dflash_selector_engine(vllm_config) - - def _sm70_qwen_gdn_full_forward_enabled( layer_name: LayerNameType, *, @@ -1338,60 +1250,6 @@ def _sm70_compile_graph_slice_dim( return tensor.index_select(dim, indices) -def _sm70_gdn_rmsnorm_onepass_enabled() -> bool: - return envs.VLLM_SM70_GDN_RMSNORM_ONEPASS - - -_SM70_GDN_QPN8_BA_REQUIRED_OPS = ( - "fp8_qpn8_gemm_ba_split_sm70_out", - "fp8_qpn8_dispatch_ba_split_sm70_out", -) - - -def _missing_sm70_gdn_qpn8_ba_ops() -> list[str]: - return [ - name - for name in _SM70_GDN_QPN8_BA_REQUIRED_OPS - if not hasattr(torch.ops._C, name) - ] - - -def _sm70_gdn_qpn8_ba_split_enabled() -> bool: - if envs.VLLM_SM70_GDN_QPN8_BA_SPLIT and not envs.VLLM_SM70_GDN_RMSNORM_ONEPASS: - raise RuntimeError( - "VLLM_SM70_GDN_QPN8_BA_SPLIT=1 requires the accepted " - "VLLM_SM70_GDN_RMSNORM_ONEPASS=1 pair." - ) - if envs.VLLM_SM70_GDN_QPN8_BA_SPLIT: - missing_ops = _missing_sm70_gdn_qpn8_ba_ops() - if missing_ops: - raise RuntimeError( - "VLLM_SM70_GDN_QPN8_BA_SPLIT=1 requires the source-built " - f"SM70 GDN extension; missing ops: {missing_ops}." - ) - return envs.VLLM_SM70_GDN_QPN8_BA_SPLIT - - -@triton.jit -def _sm70_gdn_rmsnorm_gated_onepass_kernel( - x_ptr, - z_ptr, - weight_ptr, - out_ptr, - D: tl.constexpr, - EPS: tl.constexpr, -): - row = tl.program_id(0) - offsets = tl.arange(0, D) - x = tl.load(x_ptr + row * D + offsets).to(tl.float32) - z = tl.load(z_ptr + row * D + offsets).to(tl.float32) - weight = tl.load(weight_ptr + offsets).to(tl.float32) - sumsq = tl.sum(x * x, axis=0) - rstd = tl.rsqrt(sumsq / D + EPS) - gated = z * tl.sigmoid(z) - tl.store(out_ptr + row * D + offsets, x * rstd * weight * gated) - - def _sm70_qwen_gdn_rmsnorm_gated_impl( x: torch.Tensor, z: torch.Tensor, @@ -1401,39 +1259,6 @@ def _sm70_qwen_gdn_rmsnorm_gated_impl( norm_before_gate: bool, activation: str, ) -> torch.Tensor: - if ( - _sm70_gdn_rmsnorm_onepass_enabled() - and current_platform.is_device_capability(70) - and x.is_cuda - and x.dtype == torch.float16 - and x.shape == (12, 128) - and x.is_contiguous() - and z.is_cuda - and z.device == x.device - and z.dtype == x.dtype - and z.shape == x.shape - and z.is_contiguous() - and weight.is_cuda - and weight.device == x.device - and weight.dtype == x.dtype - and weight.shape == (128,) - and weight.is_contiguous() - and group_size <= 0 - and norm_before_gate - and activation in ("silu", "swish") - ): - out = torch.empty_like(x) - _sm70_gdn_rmsnorm_gated_onepass_kernel[(12,)]( - x, - z, - weight, - out, - D=128, - EPS=eps, - num_warps=2, - ) - return out - from vllm.model_executor.layers.fla.ops.layernorm_guard import rmsnorm_fn return rmsnorm_fn( @@ -2370,54 +2195,12 @@ def __init__( 1, 1 + int(vllm_config.speculative_config.num_speculative_state_tokens()), ) - self.enable_sm70_gdn_rmsnorm_onepass = ( - _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.tp_size == 4 - and self.num_v_heads == 48 - and self.head_v_dim == 128 - and not self.gqa_interleaved_layout - ) - 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.tp_size == 4 - and self.num_v_heads == 48 - and self.head_v_dim == 128 - and not self.gqa_interleaved_layout - and not self.disable_tp_for_ba_proj - and self.in_proj_ba is not None - and _sm70_gdn_qpn8_ba_split_enabled() - ) self.enable_packed_recurrent_decode = ( envs.VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE ) self.enable_sm70_fused_sigmoid_mixed_qkv = ( envs.VLLM_SM70_FUSED_SIGMOID_MIXED_QKV ) - self.enable_sm70_dflash2_fused_gdn_verify = bool( - envs.VLLM_SM70_DFLASH2_FUSED_GDN_VERIFY - and current_platform.is_device_capability(70) - and _is_dflash2_spec_config(vllm_config) - ) - self.enable_sm70_dflash2_fused_gdn_norm = bool( - envs.VLLM_SM70_DFLASH2_FUSED_GDN_NORM - and current_platform.is_device_capability(70) - and _is_dflash2_spec_config(vllm_config) - ) - self.enable_sm70_dflash2_fused_gdn_split = bool( - envs.VLLM_SM70_DFLASH2_FUSED_GDN_SPLIT - and current_platform.is_device_capability(70) - and _is_dflash2_spec_config(vllm_config) - ) - self.enable_sm70_dflash2_fused_qkv_pack = bool( - envs.VLLM_SM70_DFLASH2_FUSED_QKV_PACK - and current_platform.is_device_capability(70) - and _is_dflash2_spec_config(vllm_config) - ) self.compare_sm70_fused_sigmoid_mixed_qkv = ( envs.VLLM_SM70_FUSED_SIGMOID_MIXED_QKV_COMPARE ) @@ -2506,26 +2289,6 @@ def __init__( "SM70 FlashQLA GDN decode route enabled.", scope="local", ) - if self.enable_sm70_dflash2_fused_gdn_verify: - logger.info_once( - "SM70 DFlash2 packed GDN target-verification route enabled.", - scope="local", - ) - if self.enable_sm70_dflash2_fused_gdn_norm: - logger.info_once( - "SM70 DFlash2 one-pass GDN output RMSNorm route enabled.", - scope="local", - ) - if self.enable_sm70_dflash2_fused_gdn_split: - logger.info_once( - "SM70 DFlash2 fused GDN z/b/a materialization route enabled.", - scope="local", - ) - if self.enable_sm70_dflash2_fused_qkv_pack: - logger.info_once( - "SM70 DFlash2 fused post-convolution Q/K/V packing enabled.", - scope="local", - ) if envs.is_set("VLLM_QWEN3_NEXT_FUSED_SIGMOID_GATING"): logger.info_once( "VLLM_QWEN3_NEXT_FUSED_SIGMOID_GATING is upstream-split " @@ -2885,19 +2648,11 @@ def rearrange_mixed_qkv(self, mixed_qkv): k_dim = self.key_dim // self.tp_size v_dim = self.value_dim // self.tp_size - if ( - self.enable_sm70_dflash2_fused_qkv_pack - and seq_len == 8 - and mixed_qkv.is_cuda - and mixed_qkv.dtype == torch.float16 - and mixed_qkv.stride(1) == 1 - ): - fused = _sm70_pack_qwen_gdn_qkv(mixed_qkv, q_dim, k_dim, v_dim) - else: - query, key, value = torch.split(mixed_qkv, [q_dim, k_dim, v_dim], dim=-1) - fused = torch.cat( - [query.reshape(-1), key.reshape(-1), value.reshape(-1)], dim=0 - ) + query, key, value = torch.split(mixed_qkv, [q_dim, k_dim, v_dim], dim=-1) + + fused = torch.cat( + [query.reshape(-1), key.reshape(-1), value.reshape(-1)], dim=0 + ) q_size = seq_len * q_dim k_size = seq_len * k_dim @@ -2990,7 +2745,6 @@ def _probe_ddtree_path_replay_states( a_spec.index_select(0, path_token_idx), b_spec.index_select(0, path_token_idx), self.dt_bias, - beta_dtype=torch.float32, ) local_ssm_state = initial_ssm_state[req_row : req_row + 1].clone() _, root_final_state = fused_recurrent_gated_delta_rule( @@ -3087,7 +2841,6 @@ def _probe_ddtree_path_replay_states( a_spec.index_select(0, path_token_idx), b_spec.index_select(0, path_token_idx), self.dt_bias, - beta_dtype=torch.float32, ) local_ssm_state = initial_ssm_state[req_row : req_row + 1].clone() _, path_final_state = fused_recurrent_gated_delta_rule( @@ -3381,7 +3134,6 @@ def _forward_ddtree_gdn_pure_spec( a_spec, b_spec, self.dt_bias, - beta_dtype=torch.float32, ) core_attn_out_linear, _ = fused_recurrent_gated_delta_rule( q=query_linear, @@ -3497,7 +3249,6 @@ def _forward_ddtree_gdn_pure_spec( a_spec.index_select(0, token_idx), b_spec.index_select(0, token_idx), self.dt_bias, - beta_dtype=torch.float32, ) cu_seqlens = torch.arange( input_state_idx.shape[0] + 1, @@ -3585,7 +3336,6 @@ def _forward_ddtree_gdn_pure_spec( a_spec.index_select(0, token_idx), b_spec.index_select(0, token_idx), self.dt_bias, - beta_dtype=torch.float32, ) cu_seqlens = torch.arange( conv_state_indices.shape[0] + 1, @@ -3810,7 +3560,7 @@ def _forward_core_decode_flashqla( def forward( self, hidden_states: torch.Tensor, - output: torch.Tensor | None, + output: torch.Tensor, ): if self.maybe_sm70_qwen_gdn_full_forward: layer_name = _encode_layer_name(self.prefix) @@ -3832,14 +3582,14 @@ def forward( layer_name, ) return - return self._forward_method(hidden_states, output) + self._forward_method(hidden_states, output) def _full_forward( self, hidden_states: torch.Tensor, output: torch.Tensor, ): - return self._forward_method(hidden_states, output) + self._forward_method(hidden_states, output) def _compute_output_projection( self, @@ -3862,44 +3612,7 @@ def _compute_output_projection( ) z = _sm70_dump_gdn_projection_tensor("proj_z", layer_name, z) profile_start = _sm70_gdn_prefill_profile_start() - use_sm70_onepass_norm = ( - self.enable_sm70_gdn_rmsnorm_onepass - and num_tokens == 1 - and core_attn_out.dtype == torch.float16 - and core_attn_out.shape == (12, 128) - and core_attn_out.is_contiguous() - and z.dtype == core_attn_out.dtype - and z.shape == core_attn_out.shape - and z.is_contiguous() - and self.norm.weight.dtype == core_attn_out.dtype - and self.norm.weight.shape == (128,) - and self.norm.weight.is_contiguous() - and self.norm.group_size is None - and self.norm.norm_before_gate - and self.norm.activation in ("silu", "swish") - ) - if self.enable_sm70_dflash2_fused_gdn_norm: - # RMSNormGated is decomposed into a nine-kernel native chain when - # it is reached from the opaque Qwen GDN op under the default - # inductor custom-op policy. Its canonical CUDA implementation is - # a one-pass Triton kernel with the same FP32 accumulation and - # FP16 output contract, so call it directly on this isolated path. - core_attn_out = self.norm.forward_cuda(core_attn_out, z) - elif use_sm70_onepass_norm: - _log_runtime_route_once( - "SM70 GDN 12x128 RMSNormGated one-pass route enabled." - ) - core_attn_out = torch.ops.vllm.sm70_qwen_gdn_rmsnorm_gated( - core_attn_out, - z, - self.norm.weight, - self.norm.eps, - -1, - True, - self.norm.activation, - ) - else: - core_attn_out = self.norm(core_attn_out, z) + core_attn_out = self.norm(core_attn_out, z) _sm70_gdn_prefill_profile_end( layer_name, "projection_norm", @@ -3932,19 +3645,18 @@ def _output_projection( self, core_attn_out: torch.Tensor, z: torch.Tensor, - output: torch.Tensor | None, + output: torch.Tensor, num_tokens: int, ) -> torch.Tensor: layer_name = _encode_layer_name(self.prefix) proj_out = self._compute_output_projection(core_attn_out, z, num_tokens) - if output is not None: - output[:num_tokens] = proj_out - _sm70_gdn_graph_buffer_copy( - "proj_output_after_write", - layer_name, - output[:num_tokens], - "proj", - ) + output[:num_tokens] = proj_out + _sm70_gdn_graph_buffer_copy( + "proj_output_after_write", + layer_name, + output[:num_tokens], + "proj", + ) return proj_out def forward_hip( @@ -4003,7 +3715,7 @@ def forward_hip( def forward_cuda( self, hidden_states: torch.Tensor, - output: torch.Tensor | None, + output: torch.Tensor, ): """ Forward pass with three parts: @@ -4041,8 +3753,8 @@ def forward_cuda( layer_name, ) else: - return self._output_projection(core_attn_out, z, output, num_tokens) - return None + self._output_projection(core_attn_out, z, output, num_tokens) + return if envs.VLLM_SM70_QWEN_GDN_INPUT_PROJECTION_OP: mixed_qkv = torch.empty( @@ -4093,8 +3805,8 @@ def forward_cuda( layer_name, ) else: - return self._output_projection(core_attn_out, z, output, num_tokens) - return None + self._output_projection(core_attn_out, z, output, num_tokens) + return # ============================================================ # Part 1: Input Projection @@ -4192,8 +3904,7 @@ def forward_cuda( layer_name, ) else: - return self._output_projection(core_attn_out, z, output, num_tokens) - return None + self._output_projection(core_attn_out, z, output, num_tokens) def forward_xpu( self, @@ -5102,85 +4813,6 @@ def _forward_core_rocm( kv_cache=kv_cache, ) - def _can_use_dflash2_packed_gdn_verify( - self, - *, - mixed_qkv: torch.Tensor | None, - a: torch.Tensor, - b: torch.Tensor, - core_attn_out: torch.Tensor, - ssm_state: torch.Tensor, - attn_metadata: GDNAttentionMetadata, - ) -> bool: - return bool( - self.enable_sm70_dflash2_fused_gdn_verify - and mixed_qkv is not None - and attn_metadata.spec_sequence_masks is not None - and attn_metadata.num_spec_decodes > 0 - and attn_metadata.num_prefills == 0 - and attn_metadata.num_decodes == 0 - and attn_metadata.ddtree_parent_ids is None - and attn_metadata.spec_query_start_loc is not None - and attn_metadata.spec_state_indices_tensor is not None - and attn_metadata.spec_state_slot_selectors is not None - and mixed_qkv.is_cuda - and mixed_qkv.dtype == torch.float16 - and a.dtype == torch.float16 - and b.dtype == torch.float16 - # The supported verifier contract keeps recurrent state in FP32; - # an explicit FP16 cache override is also supported. - and ssm_state.dtype in (torch.float16, torch.float32) - and mixed_qkv.is_contiguous() - and a.is_contiguous() - and b.is_contiguous() - and core_attn_out.is_contiguous() - ) - - def _forward_dflash2_packed_gdn_verify( - self, - *, - mixed_qkv: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - core_attn_out: torch.Tensor, - ssm_state: torch.Tensor, - spec_query_start_loc: torch.Tensor, - spec_state_indices_tensor: torch.Tensor, - spec_state_slot_selectors: torch.Tensor, - num_spec_decodes: int, - ) -> torch.Tensor: - num_tokens = mixed_qkv.shape[0] - out = core_attn_out[:num_tokens].unsqueeze(1) - g, beta = fused_gdn_gating(self.A_log, a, b, self.dt_bias) - fused_sigmoid_gating_delta_rule_update_mixed_qkv_out( - A_log=self.A_log, - a=a, - b=b, - dt_bias=self.dt_bias, - mixed_qkv=mixed_qkv, - num_q_heads=self.num_k_heads // self.tp_size, - num_v_heads=self.num_v_heads // self.tp_size, - head_k_dim=self.head_k_dim, - head_v_dim=self.head_v_dim, - scale=self.head_k_dim**-0.5, - initial_state=ssm_state, - out=out, - cu_seqlens=spec_query_start_loc[: num_spec_decodes + 1], - ssm_state_indices=spec_state_indices_tensor[:num_spec_decodes], - num_accepted_tokens=spec_state_slot_selectors[:num_spec_decodes], - use_qk_l2norm_in_kernel=True, - # Keep gating materialization and launch geometry bitwise aligned - # with the current verifier. This stage removes packed-QKV - # rearrangement and the final output copy without changing the - # target recurrence. Fully fused gating stays separately testable. - precomputed_g=g, - precomputed_beta=beta, - quantize_state_each_step=False, - match_recurrent_schedule=True, - match_recurrent_numerics=True, - ) - return out.transpose(0, 1) - def _forward_core( self, mixed_qkv: torch.Tensor, @@ -5533,18 +5165,7 @@ def _forward_core( core_attn_out[:num_non_spec_tokens] = core_attn_out_non_spec.squeeze(0) return - use_dflash2_packed_gdn_verify = self._can_use_dflash2_packed_gdn_verify( - mixed_qkv=mixed_qkv_spec, - a=a, - b=b, - core_attn_out=core_attn_out, - ssm_state=ssm_state, - attn_metadata=attn_metadata, - ) - if use_dflash2_packed_gdn_verify: - query_spec, key_spec, value_spec = None, None, None - else: - query_spec, key_spec, value_spec = self.rearrange_mixed_qkv(mixed_qkv_spec) + query_spec, key_spec, value_spec = self.rearrange_mixed_qkv(mixed_qkv_spec) prefill_l2norm_in_kernel = False use_original_flashqla_prefill = _sm70_flashqla_original_prefill_enabled() prefill_gate_is_exp = False @@ -5659,27 +5280,7 @@ def _forward_core( assert spec_token_indx is not None a_spec = a.index_select(0, spec_token_indx) b_spec = b.index_select(0, spec_token_indx) - if use_dflash2_packed_gdn_verify: - assert mixed_qkv_spec is not None - assert spec_query_start_loc is not None - assert spec_state_indices_tensor is not None - assert spec_state_slot_selectors is not None - core_attn_out_spec = self._forward_dflash2_packed_gdn_verify( - mixed_qkv=mixed_qkv_spec, - a=a_spec, - b=b_spec, - core_attn_out=core_attn_out, - ssm_state=ssm_state, - spec_query_start_loc=spec_query_start_loc, - spec_state_indices_tensor=spec_state_indices_tensor, - spec_state_slot_selectors=spec_state_slot_selectors, - num_spec_decodes=attn_metadata.num_spec_decodes, - ) - last_recurrent_state = ssm_state - _log_runtime_route_once( - "SM70 DFlash2 packed GDN target-verification route hit." - ) - elif ddtree_tree_gdn_pure_spec: + if ddtree_tree_gdn_pure_spec: assert mixed_qkv_spec is not None assert spec_query_start_loc is not None assert spec_state_indices_tensor is not None @@ -5712,15 +5313,8 @@ def _forward_core( ) last_recurrent_state = ssm_state else: - assert query_spec is not None - assert key_spec is not None - assert value_spec is not None g_spec, beta_spec = fused_gdn_gating( - self.A_log, - a_spec, - b_spec, - self.dt_bias, - beta_dtype=torch.float32, + self.A_log, a_spec, b_spec, self.dt_bias ) core_attn_out_spec, last_recurrent_state = ( fused_recurrent_gated_delta_rule( @@ -5921,8 +5515,7 @@ def _forward_core( merged_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec) core_attn_out[:num_actual_tokens] = merged_out.squeeze(0) elif spec_sequence_masks is not None: - if not use_dflash2_packed_gdn_verify: - core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0) + core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0) elif core_attn_out_non_spec is not None and not ( attn_metadata.num_prefills > 0 and spec_sequence_masks is None @@ -6705,41 +6298,12 @@ def _patch_fallback_metadata(name: str, value: object) -> None: max_query_len=spec_state_indices_tensor.size(-1), validate_data=False, ) - if isinstance(self, QwenGatedDeltaNetAttention) and ( - self._can_use_dflash2_packed_gdn_verify( - mixed_qkv=mixed_qkv_spec, - a=a, - b=b, - core_attn_out=core_attn_out, - ssm_state=ssm_state, - attn_metadata=attn_metadata, - ) - ): - self._forward_dflash2_packed_gdn_verify( - mixed_qkv=mixed_qkv_spec, - a=a, - b=b, - core_attn_out=core_attn_out, - ssm_state=ssm_state, - spec_query_start_loc=spec_query_start_loc, - spec_state_indices_tensor=spec_state_indices_tensor, - spec_state_slot_selectors=spec_state_slot_selectors, - num_spec_decodes=attn_metadata.num_spec_decodes, - ) - _log_runtime_route_once( - "SM70 DFlash2 packed GDN target-verification route hit." - ) - _sm70_dump_gdn_core_tensor( - "core_out", layer_name, core_attn_out, metadata_source - ) - return core_attn_out query_spec, key_spec, value_spec = self.rearrange_mixed_qkv(mixed_qkv_spec) g_spec, beta_spec = fused_gdn_gating( self.A_log, a, b, self.dt_bias, - beta_dtype=torch.float32, ) core_attn_out_spec, _ = fused_recurrent_gated_delta_rule( q=query_spec, @@ -7118,77 +6682,6 @@ def qwen_gdn_output_projection_fake( """Fake implementation for torch.compile.""" -def _sm70_gdn_qpn8_ba_weight_contract( - self: "QwenGatedDeltaNetAttention", -) -> bool: - """Check the loaded operator layout without using model identity.""" - qkvz = self.in_proj_qkvz - ba = self.in_proj_ba - 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 - return bool( - self.enable_sm70_gdn_qpn8_ba_split - and ba is not None - and getattr(qkvz, "sm70_fp8_qpn8", False) - 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 qkvz_weight.is_contiguous() - and isinstance(qkvz_scales, torch.Tensor) - and qkvz_scales.dtype == torch.float16 - and qkvz_scales.shape == (1, 4096) - 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.is_contiguous() - and qkvz_weight.device == qkvz_scales.device == ba_weight.device - and getattr(qkvz, "bias", None) is None - and getattr(ba, "bias", None) is None - ) - - -def _sm70_gdn_qpn8_ba_dispatch_eligible( - self: "QwenGatedDeltaNetAttention", - hidden_states: torch.Tensor, - layer_name: LayerNameType, -) -> bool: - qkvz = self.in_proj_qkvz - ba = self.in_proj_ba - return bool( - _sm70_gdn_qpn8_ba_weight_contract(self) - and ba is not None - and not _sm70_gdn_projection_dump_requested(layer_name) - and hasattr(torch.ops._C, "fp8_qpn8_dispatch_ba_split_sm70_out") - and hasattr(torch.ops._C, "fp8_qpn8_gemm_ba_split_sm70_out") - 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.is_contiguous() - and hidden_states.device == qkvz.weight.device == ba.weight.device - ) - - -def _sm70_gdn_qpn8_ba_split_eligible( - self: "QwenGatedDeltaNetAttention", - hidden_states: torch.Tensor, - z_out: torch.Tensor, - layer_name: LayerNameType, -) -> bool: - return ( - _sm70_gdn_qpn8_ba_dispatch_eligible(self, hidden_states, layer_name) - and hidden_states.shape == (1, 5120) - and z_out.is_cuda - and z_out.device == hidden_states.device - and z_out.dtype == hidden_states.dtype - and z_out.shape == (1, 12, 128) - and z_out.is_contiguous() - ) - - def qwen_gdn_input_projection_core( hidden_states: torch.Tensor, z_out: torch.Tensor, @@ -7207,75 +6700,46 @@ def qwen_gdn_input_projection_core( layer_name, hidden_states, ) - if _sm70_gdn_qpn8_ba_split_eligible( - self, - hidden_states, - z_out, + mixed_qkvz, _ = self.in_proj_qkvz(hidden_states) + ba, _ = self.in_proj_ba(hidden_states) + mixed_qkvz = _sm70_dump_gdn_projection_tensor( + "input_core_in_proj_qkvz", layer_name, - ): - mixed_qkv = torch.empty( - (1, 2560), dtype=hidden_states.dtype, device=hidden_states.device - ) - b = torch.empty((1, 12), dtype=hidden_states.dtype, device=hidden_states.device) - a = torch.empty_like(b) - sm70_ops.fp8_qpn8_gemm_ba_split_sm70_out( - mixed_qkv, - z_out, - b, - a, - hidden_states, - self.in_proj_qkvz.weight, - self.in_proj_qkvz.weight_scale_inv, - self.in_proj_ba.weight, + mixed_qkvz, + ) + ba = _sm70_dump_gdn_projection_tensor("input_core_in_proj_ba", layer_name, ba) + + if self.gqa_interleaved_layout: + query, key, value, z, b, a = self.fix_query_key_value_ordering( + mixed_qkvz, + ba, ) - _log_runtime_route_once( - "SM70 GDN QPN8 K5120/N4096 plus FP16 b/a N24 split route enabled." + query, key, value = map( + lambda x: rearrange(x, "l p d -> l (p d)"), (query, key, value) ) - z = z_out + mixed_qkv = torch.cat((query, key, value), dim=-1) else: - mixed_qkvz, _ = self.in_proj_qkvz(hidden_states) - ba, _ = self.in_proj_ba(hidden_states) - mixed_qkvz = _sm70_dump_gdn_projection_tensor( - "input_core_in_proj_qkvz", - layer_name, - mixed_qkvz, - ) - ba = _sm70_dump_gdn_projection_tensor("input_core_in_proj_ba", layer_name, ba) - - if self.gqa_interleaved_layout: - query, key, value, z, b, a = self.fix_query_key_value_ordering( - mixed_qkvz, - ba, - ) - query, key, value = map( - lambda x: rearrange(x, "l p d -> l (p d)"), - (query, key, value), - ) - mixed_qkv = torch.cat((query, key, value), dim=-1) - else: - qkv_size = (self.key_dim * 2 + self.value_dim) // self.tp_size - z_size = self.value_dim // self.tp_size - mixed_qkv = mixed_qkvz[..., :qkv_size] - if envs.VLLM_SM70_GDN_MIXED_QKV_CONTIGUOUS: - mixed_qkv = mixed_qkv.contiguous() - z = _sm70_compile_graph_slice_dim(mixed_qkvz, -1, qkv_size, z_size) - z = z.reshape(z.size(0), -1, self.head_v_dim) - ba_size = ba.shape[-1] // 2 - b = ba[..., :ba_size] - a = _sm70_compile_graph_slice_dim(ba, -1, ba_size, ba_size) - if self.disable_tp_for_ba_proj and self.tp_size > 1: - ba_chunk = self.num_v_heads // self.tp_size - ba_start = self.tp_rank * ba_chunk - b = b[:, ba_start : ba_start + ba_chunk] - a = a[:, ba_start : ba_start + ba_chunk] - b = b.contiguous() - a = a.contiguous() + qkv_size = (self.key_dim * 2 + self.value_dim) // self.tp_size + z_size = self.value_dim // self.tp_size + mixed_qkv = mixed_qkvz[..., :qkv_size] + if envs.VLLM_SM70_GDN_MIXED_QKV_CONTIGUOUS: + mixed_qkv = mixed_qkv.contiguous() + z = _sm70_compile_graph_slice_dim(mixed_qkvz, -1, qkv_size, z_size) + z = z.reshape(z.size(0), -1, self.head_v_dim) + ba_size = ba.shape[-1] // 2 + b = ba[..., :ba_size] + a = _sm70_compile_graph_slice_dim(ba, -1, ba_size, ba_size) + if self.disable_tp_for_ba_proj and self.tp_size > 1: + ba_chunk = self.num_v_heads // self.tp_size + ba_start = self.tp_rank * ba_chunk + b = b[:, ba_start : ba_start + ba_chunk] + a = a[:, ba_start : ba_start + ba_chunk] + b = b.contiguous() + a = a.contiguous() - if envs.VLLM_SM70_GDN_Z_CONTIGUOUS and current_platform.is_device_capability( - 70 - ): - z = z.contiguous() - z_out.copy_(z) + if envs.VLLM_SM70_GDN_Z_CONTIGUOUS and current_platform.is_device_capability(70): + z = z.contiguous() + z_out.copy_(z) _qwen_gdn_run_recurrent_core( self, @@ -7320,29 +6784,7 @@ def qwen_gdn_input_projection( layer_name, hidden_states, ) - if _sm70_gdn_qpn8_ba_split_eligible( - self, - hidden_states, - z_out, - layer_name, - ): - sm70_ops.fp8_qpn8_gemm_ba_split_sm70_out( - mixed_qkv_out, - z_out, - b_out, - a_out, - hidden_states, - self.in_proj_qkvz.weight, - self.in_proj_qkvz.weight_scale_inv, - self.in_proj_ba.weight, - ) - _log_runtime_route_once( - "SM70 GDN QPN8 K5120/N4096 plus FP16 b/a N24 split route enabled." - ) - return - mixed_qkvz, _ = self.in_proj_qkvz(hidden_states) - assert self.in_proj_ba is not None ba, _ = self.in_proj_ba(hidden_states) if self.gqa_interleaved_layout: @@ -7542,25 +6984,18 @@ def fused_gdn_gating( dt_bias: torch.Tensor, beta: float = 1.0, threshold: float = 20.0, - *, - beta_dtype: torch.dtype | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """ Fused computation of g and beta for Gated Delta Net. g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias) - beta_output = b.sigmoid(), materialized in beta_dtype when provided. - - Speculative recurrent updates request FP32 beta to match the fused packed - decode transition. Other callers retain the historical b.dtype output. + beta_output = b.sigmoid() TODO maybe use torch.compile to replace this triton kernel """ batch, num_heads = a.shape seq_len = 1 grid = (batch, seq_len, triton.cdiv(num_heads, 8)) g = torch.empty(1, batch, num_heads, dtype=torch.float32, device=a.device) - if beta_dtype is None: - beta_dtype = b.dtype - beta_output = torch.empty(1, batch, num_heads, dtype=beta_dtype, device=b.device) + beta_output = torch.empty(1, batch, num_heads, dtype=b.dtype, device=b.device) fused_gdn_gating_kernel[grid]( g, beta_output, diff --git a/vllm/model_executor/layers/quantization/auto_gptq.py b/vllm/model_executor/layers/quantization/auto_gptq.py index 9b481a0241..ce0ed71b49 100644 --- a/vllm/model_executor/layers/quantization/auto_gptq.py +++ b/vllm/model_executor/layers/quantization/auto_gptq.py @@ -480,11 +480,15 @@ 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( @@ -494,7 +498,9 @@ 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 860daeb8d6..2766740c06 100644 --- a/vllm/model_executor/layers/quantization/awq_marlin.py +++ b/vllm/model_executor/layers/quantization/awq_marlin.py @@ -610,7 +610,9 @@ 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/compressed_tensors.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py index 7791dc2f1a..406546d90c 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py @@ -380,12 +380,6 @@ def _is_nvfp4_format(quant_args: QuantizationArgs): and is_symmetric ) - def _contains_nvfp4_weights(self) -> bool: - return any( - scheme is not None and self._is_nvfp4_format(scheme.get("weights")) - for scheme in self.target_scheme_map.values() - ) - @staticmethod def _is_mxfp4(quant_args: QuantizationArgs) -> bool: if quant_args is None: @@ -682,7 +676,6 @@ def _get_scheme_from_parts( return CompressedTensorsW8A16Fp8( weight_quant=weight_quant, is_static_input_scheme=not input_quant.dynamic, - enable_sm70_qpn8_by_default=self._contains_nvfp4_weights(), ) # note: input_quant can be None @@ -691,7 +684,6 @@ def _get_scheme_from_parts( return CompressedTensorsW8A16Fp8( weight_quant=weight_quant, is_static_input_scheme=is_static_input_scheme, - enable_sm70_qpn8_by_default=self._contains_nvfp4_weights(), ) if self._is_static_tensor_w8a8(weight_quant, input_quant): @@ -862,19 +854,6 @@ def apply( raise ValueError("A scheme must be defined for each layer") return scheme.apply_weights(layer, x, bias=bias) - def apply_fused_silu_and_mul( - self, - layer: torch.nn.Module, - x: torch.Tensor, - ) -> torch.Tensor | None: - scheme = layer.scheme - if scheme is None: - raise ValueError("A scheme must be defined for each layer") - fused_apply = getattr(scheme, "apply_fused_silu_and_mul", None) - if fused_apply is None: - return None - return fused_apply(layer, x) - class CompressedTensorsKVCacheMethod(BaseKVCacheMethod): """ 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 5f1fc7e1b9..1fa4f8632e 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,7 +97,8 @@ 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 f9d6f0b9b5..07e252758f 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,11 +108,15 @@ 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..e740664b4e 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 @@ -5,9 +5,7 @@ import torch from torch.nn.parameter import Parameter -from vllm import _sm70_ops as sm70_ops from vllm import envs -from vllm.config import get_current_vllm_config from vllm.logger import init_logger from vllm.model_executor.kernels.linear import init_nvfp4_linear_kernel from vllm.model_executor.layers.quantization import sm70_turbomind as sm70_tm @@ -23,96 +21,14 @@ logger = init_logger(__name__) -def _is_sm70_tp4_nvfp4_gate_up(layer: torch.nn.Module) -> bool: - return bool( - getattr(layer, "tp_size", 1) == 4 - and getattr(layer, "prefix", "").rsplit(".", 1)[-1] == "gate_up_proj" - and getattr(layer, "input_size_per_partition", 0) == 5120 - and getattr(layer, "output_size_per_partition", 0) == 8704 - and getattr(layer, "logical_widths", None) == [4352, 4352] - ) - - -def _is_sm70_tp4_nvfp4_down(layer: torch.nn.Module) -> bool: - return bool( - getattr(layer, "tp_size", 1) == 4 - and getattr(layer, "prefix", "").rsplit(".", 1)[-1] == "down_proj" - and getattr(layer, "input_size_per_partition", 0) == 4352 - and getattr(layer, "output_size_per_partition", 0) == 5120 - ) - - -def _is_sm70_nvfp4_qpn4_runtime_contract() -> bool: - """Admit only the measured single-sequence, no-MTP decode contract.""" - vllm_config = get_current_vllm_config() - scheduler_config = getattr(vllm_config, "scheduler_config", None) - max_num_seqs = int(getattr(scheduler_config, "max_num_seqs", 1)) - speculative_config = getattr(vllm_config, "speculative_config", None) - return max_num_seqs == 1 and speculative_config is None - - -_SM70_NVFP4_QPN4_REQUIRED_OPS = ( - "nvfp4_qpn4_prepare_sm70", - "nvfp4_qpn4_prepare_scale_code_sm70", - "nvfp4_qpn4_dequantize_sm70_out", - "nvfp4_qpn4_prefill_sm70_out", - "nvfp4_qpn4_dispatch_sm70_out", -) - - -def _missing_sm70_nvfp4_qpn4_ops() -> list[str]: - return [ - name - for name in _SM70_NVFP4_QPN4_REQUIRED_OPS - if not hasattr(torch.ops._C, name) - ] - - __all__ = ["CompressedTensorsW4A4Fp4"] -_SM70_NVFP4_QPN2_CONFIGS = { - # (K, N, fused gated-SiLU): (split-K, independent accumulator chains) - (5120, 8704, False): (8, 2), - (5120, 8704, True): (8, 2), - (4352, 5120, False): (16, 2), -} -_SM70_NVFP4_QPN2_SHAPES = { - # Checkpoint-native packed tensors are [N, K/2]. - "gate_up_proj": (8704, 2560), - "down_proj": (5120, 2176), -} -_SM70_NVFP4_QPN2_REQUIRED_OPS = ( - "nvfp4_qpn2_prepare_sm70", - "nvfp4_qpn2_gemm_sm70_out", - "nvfp4_qpn2_gated_sm70_out", - "nvfp4_qpn2_dispatch_sm70_out", -) - - -def _is_qpn2_layer(layer: torch.nn.Module) -> bool: - if getattr(layer, "tp_size", 1) != 4: - return False - suffix = getattr(layer, "prefix", "").rsplit(".", 1)[-1] - expected = _SM70_NVFP4_QPN2_SHAPES.get(suffix) - if expected is None or tuple(layer.weight.shape) != expected: - return False - expected_n, expected_packed_k = expected - return bool( - getattr(layer, "input_size_per_partition", 0) == expected_packed_k * 2 - and getattr(layer, "output_size_per_partition", 0) == expected_n - ) - - -def _missing_qpn2_ops() -> list[str]: - return [ - name - for name in _SM70_NVFP4_QPN2_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": + if ( + envs.VLLM_USE_NVFP4_CT_EMULATIONS + or envs.VLLM_NVFP4_GEMM_BACKEND == "emulation" + ): return True from vllm.config import get_current_vllm_config_or_none @@ -237,106 +153,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: logger.info_once( "SM70 compressed-tensors NVFP4 TurboMind W4A16 dense path enabled." ) - is_qpn4_gate = _is_sm70_tp4_nvfp4_gate_up(layer) - is_qpn4_down = _is_sm70_tp4_nvfp4_down(layer) - qpn4_model_layer = envs.VLLM_SM70_NVFP4_QPN4 and ( - is_qpn4_down or (is_qpn4_gate and envs.VLLM_SM70_NVFP4_DENSE_GATED_SILU) - ) - qpn4_runtime = ( - _is_sm70_nvfp4_qpn4_runtime_contract() if qpn4_model_layer else False - ) - if qpn4_model_layer and not qpn4_runtime: - logger.info_once( - "The SM70 NVFP4 QPN4 route retains TurboMind unless the " - "runtime contract is max_num_seqs=1 with no MTP." - ) - if qpn4_model_layer and qpn4_runtime: - missing_ops = _missing_sm70_nvfp4_qpn4_ops() - if missing_ops: - logger.warning_once( - "The automatic SM70 NVFP4 QPN4 route is unavailable " - "in the loaded vllm._C; retaining TurboMind. Missing " - f"ops: {missing_ops}." - ) - workspace = ( - None - if missing_ops - else sm70_tm.get_nvfp4_qpn4_dense_workspace(layer.weight) - ) - if not missing_ops and workspace is not None: - sm70_tm.prepare_nvfp4_qpn4_linear( - layer, - workspace, - gated_silu=is_qpn4_gate, - ) - layer.weight = Parameter( - torch.empty(0, dtype=torch.uint8, device=layer.weight.device), - requires_grad=False, - ) - layer.weight_scale = Parameter( - torch.empty( - 0, - dtype=torch.float8_e4m3fn, - device=layer.weight_scale.device, - ), - requires_grad=False, - ) - logger.info_once( - "Memory-neutral SM70 NVFP4 QPN4 M=1 decode " - "path enabled with bounded FP16 prefill workspace." - ) - return - if not missing_ops: - logger.warning_once( - "Insufficient memory for the bounded SM70 NVFP4 QPN4 " - "prefill workspace; retaining TurboMind." - ) - use_qpn2 = bool(envs.VLLM_SM70_NVFP4_QPN2 and _is_qpn2_layer(layer)) - if use_qpn2: - missing_ops = _missing_qpn2_ops() - if missing_ops: - logger.warning_once( - "The requested SM70 NVFP4 QPN2 route is unavailable; " - f"retaining TurboMind. Missing ops: {missing_ops}." - ) - use_qpn2 = False - if use_qpn2: - qpn2_codes, qpn2_scales = sm70_ops.nvfp4_qpn2_prepare_sm70( - layer.weight.data, layer.weight_scale.data - ) - qpn2_global_scale = float(layer.weight_global_scale.item()) - - use_gated_silu = bool( - envs.VLLM_SM70_NVFP4_DENSE_GATED_SILU and is_qpn4_gate and not use_qpn2 - ) - sm70_tm.prepare_nvfp4_linear( - layer, - interleave_gated_silu=use_gated_silu, - ) - if use_qpn2: - suffix = layer.prefix.rsplit(".", 1)[-1] - k = layer.input_size_per_partition - n = layer.output_size_per_partition - split_k, nacc = _SM70_NVFP4_QPN2_CONFIGS[(k, n, False)] - layer.register_buffer( - "sm70_nvfp4_qpn2_codes", qpn2_codes, persistent=False - ) - layer.register_buffer( - "sm70_nvfp4_qpn2_scales", qpn2_scales, persistent=False - ) - layer.sm70_nvfp4_qpn2 = True - layer.sm70_nvfp4_qpn2_global_scale = qpn2_global_scale - layer.sm70_nvfp4_qpn2_split_k = split_k - layer.sm70_nvfp4_qpn2_nacc = nacc - layer.sm70_nvfp4_qpn2_gated_silu = suffix == "gate_up_proj" - logger.info_once( - "SM70 NVFP4 QPN2 M<=8 route enabled for a compatible " - "TP4 projection contract." - ) - elif use_gated_silu: - logger.info_once( - "SM70 NVFP4 TurboMind gated-SiLU single-layout path enabled." - ) + sm70_tm.prepare_nvfp4_linear(layer) layer.weight = Parameter( torch.empty(0, dtype=torch.uint8, device=layer.weight.device), requires_grad=False, @@ -363,68 +180,6 @@ def apply_weights( x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: - if getattr(layer, "sm70_nvfp4_qpn2", False): - return self._apply_qpn2(layer, x, bias, gated_silu=False) if sm70_tm.has_prepared_linear(layer): return sm70_tm.apply_prepared_linear(layer, x, bias) return self._fallback_kernel().apply_weights(layer=layer, x=x, bias=bias) - - def apply_fused_silu_and_mul( - self, - layer: torch.nn.Module, - x: torch.Tensor, - ) -> torch.Tensor | None: - if getattr(layer, "sm70_nvfp4_qpn2_gated_silu", False): - return self._apply_qpn2(layer, x, None, gated_silu=True) - if not sm70_tm.has_prepared_linear(layer): - return None - return sm70_tm.apply_prepared_fused_silu_and_mul(layer, x) - - @staticmethod - def _apply_qpn2( - layer: torch.nn.Module, - x: torch.Tensor, - bias: torch.Tensor | None, - *, - gated_silu: bool, - ) -> torch.Tensor: - if x.dtype != torch.float16: - raise RuntimeError( - f"SM70 NVFP4 QPN2 requires float16 activations, got {x.dtype}." - ) - x_2d = x.reshape(-1, x.shape[-1]) - if x_2d.stride(-1) != 1: - x_2d = x_2d.contiguous() - output_size = layer.output_size_per_partition - if gated_silu: - output_size //= 2 - out_2d = torch.empty( - (x_2d.shape[0], output_size), dtype=x.dtype, device=x.device - ) - if x_2d.shape[0] == 0: - 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) - if gated_silu: - split_k, nacc = _SM70_NVFP4_QPN2_CONFIGS[ - (x_2d.shape[1], output_size * 2, True) - ] - sm70_ops.nvfp4_qpn2_dispatch_sm70_out( - out_2d, - x_2d, - layer.sm70_nvfp4_qpn2_codes, - layer.sm70_nvfp4_qpn2_scales, - float(layer.sm70_nvfp4_qpn2_global_scale), - split_k, - nacc, - state.weight, - state.scales, - state.group_size, - state.k_ld, - state.q_ld, - gated_silu, - ) - if bias is not None: - out_2d.add_(bias) - return out_2d.reshape(*x.shape[:-1], output_size) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py index 2b7ca24ff1..61f0f5b923 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py @@ -1,20 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import os from collections.abc import Callable import torch from compressed_tensors.quantization import QuantizationArgs, QuantizationStrategy -from vllm import _sm70_ops as sm70_ops -from vllm import envs from vllm.config import get_current_vllm_config -from vllm.logger import init_logger from vllm.model_executor.kernels.linear import ( init_wfp8_a16_linear_kernel, ) -from vllm.model_executor.layers.quantization import sm70_turbomind as sm70_tm from vllm.model_executor.layers.quantization.compressed_tensors.schemes import ( CompressedTensorsScheme, ) @@ -22,11 +17,6 @@ STRATEGY_TO_PARAMETER_TYPE, STRATEGY_TO_WEIGHT_QUANT_KEY, ) -from vllm.model_executor.layers.quantization.fp8 import ( - _get_sm70_fp8_prefill_exact_dense_workspace, - _is_sm70_fp8_qpn8_runtime_contract, - _missing_sm70_fp8_qpn8_ops, -) from vllm.model_executor.layers.quantization.utils.fp8_utils import ( create_fp8_scale_parameter, create_fp8_weight_parameter, @@ -44,53 +34,9 @@ __all__ = ["CompressedTensorsW8A16Fp8"] -logger = init_logger(__name__) - -_SM70_CHANNEL_FP8_QPN8_SHAPES = { - "in_proj_qkvz": (4096, 5120), - "qkv_proj": (3584, 5120), - "out_proj": (5120, 1536), - "o_proj": (5120, 1536), - "gate_up_proj": (8704, 5120), - "down_proj": (5120, 4352), -} -_SM70_CHANNEL_FP8_QPN8_CONFIGS = { - # (K, N): (split-K, accumulator chains, prefetch codes) - (5120, 4096): (16, 2, False), - (5120, 3584): (16, 2, False), - (1536, 5120): (12, 2, False), - (5120, 8704): (16, 2, False), - (4352, 5120): (16, 2, False), -} -_SM70_CHANNEL_FP8_QPN8_GATED_CONFIG = (8, 2, False) - - -def _sm70_fp8_qpn8_enabled(enable_by_default: bool) -> bool: - """Resolve QPN8 while preserving the validated mixed-NVFP4 default.""" - if os.getenv("VLLM_SM70_FP8_QPN8") is None: - return enable_by_default - return envs.VLLM_SM70_FP8_QPN8 - - -def _sm70_channel_fp8_qpn8_config( - layer: torch.nn.Module, -) -> tuple[int, int, bool] | None: - if getattr(layer, "tp_size", 1) != 4: - return None - suffix = getattr(layer, "prefix", "").rsplit(".", 1)[-1] - if tuple(layer.weight.shape) != _SM70_CHANNEL_FP8_QPN8_SHAPES.get(suffix): - return None - n_dim, k_dim = (int(dim) for dim in layer.weight.shape) - return _SM70_CHANNEL_FP8_QPN8_CONFIGS.get((k_dim, n_dim)) - class CompressedTensorsW8A16Fp8(CompressedTensorsScheme): - def __init__( - self, - weight_quant: QuantizationArgs, - is_static_input_scheme: bool, - enable_sm70_qpn8_by_default: bool = False, - ): + def __init__(self, weight_quant: QuantizationArgs, is_static_input_scheme: bool): self.weight_quant = weight_quant self.strategy = weight_quant.strategy self.out_dtype = torch.get_default_dtype() @@ -102,15 +48,6 @@ def __init__( self.activation_quant_key = ( kFp8StaticTensorSym if is_static_input_scheme else kFp8DynamicTensorSym ) - self.use_sm70_fp8_turbomind = ( - self.strategy == QuantizationStrategy.CHANNEL - and not self.is_static_input_scheme - and sm70_tm.is_exact_sm70_cuda_platform() - and sm70_tm.use_turbomind(envs.VLLM_SM70_FP8_TURBOMIND) - ) - self.use_sm70_fp8_qpn8 = self.use_sm70_fp8_turbomind and _sm70_fp8_qpn8_enabled( - enable_sm70_qpn8_by_default - ) @classmethod def get_min_capability(cls) -> int: @@ -171,9 +108,6 @@ def create_weights( ) layer.register_parameter("input_scale", input_scale) - if self.use_sm70_fp8_turbomind: - return - self.linear_kernel = init_wfp8_a16_linear_kernel( weight_quant_key=self.weight_quant_key, activation_quant_key=self.activation_quant_key, @@ -183,112 +117,6 @@ def create_weights( ) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - if self.use_sm70_fp8_turbomind: - if layer.orig_dtype != torch.float16: - raise RuntimeError( - "SM70 TurboMind channel-FP8 requires fp16 model weights, " - f"got {layer.orig_dtype}." - ) - if not hasattr(torch.ops._C, "fp8_sm70_prepare"): - raise RuntimeError( - "VLLM_SM70_FP8_TURBOMIND=1 requires a build with CUDA " - "arch 7.0 and the SM70 TurboMind extension." - ) - weight_scale = layer.weight_scale.to(torch.float32).contiguous() - expected_scale_shape = (layer.weight.shape[0], 1) - if tuple(weight_scale.shape) != expected_scale_shape: - raise RuntimeError( - "SM70 TurboMind channel-FP8 expected weight scales with " - f"shape {expected_scale_shape}, got {tuple(weight_scale.shape)}." - ) - qpn8_config = ( - _sm70_channel_fp8_qpn8_config(layer) - if getattr(self, "use_sm70_fp8_qpn8", False) - else None - ) - qpn8_concurrency = ( - _is_sm70_fp8_qpn8_runtime_contract() - if qpn8_config is not None - else False - ) - if qpn8_config is not None and not qpn8_concurrency: - logger.info_once( - "The SM70 channel-FP8 QPN8 route retains TurboMind when " - "MTP is enabled or max_num_seqs exceeds 8." - ) - if qpn8_config is not None and qpn8_concurrency: - missing_ops = _missing_sm70_fp8_qpn8_ops() - if missing_ops and os.getenv("VLLM_SM70_FP8_QPN8") is not None: - raise RuntimeError( - "VLLM_SM70_FP8_QPN8=1 requires the source-built SM70 " - f"QPN8 extension; missing ops: {missing_ops}." - ) - workspace = ( - None - if missing_ops - else _get_sm70_fp8_prefill_exact_dense_workspace(layer.weight) - ) - if not missing_ops and workspace is not None: - qpn8_codes, qpn8_scales = sm70_ops.fp8_qpn8_prepare_sm70( - layer.weight, weight_scale - ) - replace_parameter(layer, "weight", qpn8_codes) - del layer._parameters["weight_scale"] - replace_parameter(layer, "weight_scale_inv", qpn8_scales) - split_k, nacc, prefetch = qpn8_config - layer.input_scale = None - layer.sm70_fp8_turbomind = True - layer.sm70_fp8_qpn8 = True - layer.sm70_fp8_channel_scale = True - layer.sm70_fp8_qpn8_split_k = split_k - layer.sm70_fp8_qpn8_nacc = nacc - layer.sm70_fp8_qpn8_prefetch = prefetch - layer.sm70_fp8_prefill_exact_dense_workspace_ptr = ( - workspace.data_ptr() - ) - if getattr(layer, "prefix", "").rsplit(".", 1)[-1] == ( - "gate_up_proj" - ): - gated_split_k, gated_nacc, gated_prefetch = ( - _SM70_CHANNEL_FP8_QPN8_GATED_CONFIG - ) - layer.sm70_fp8_gated_silu = True - layer.sm70_fp8_qpn8_gated_split_k = gated_split_k - layer.sm70_fp8_qpn8_gated_nacc = gated_nacc - layer.sm70_fp8_qpn8_gated_prefetch = gated_prefetch - logger.info_once( - "Memory-neutral SM70 channel-FP8 QPN8 path enabled " - "for accepted Qwen3.8-27B TP4 dense shapes." - ) - return - if missing_ops: - logger.warning_once( - "The SM70 channel-FP8 QPN8 operators are unavailable; " - "retaining the TurboMind layout." - ) - else: - logger.warning_once( - "Insufficient memory for the SM70 channel-FP8 QPN8 " - "prefill workspace; retaining the TurboMind layout." - ) - tm_weight, tm_scales, meta = sm70_ops.fp8_sm70_prepare( - layer.weight, weight_scale, 128, False - ) - replace_parameter(layer, "weight", tm_weight) - del layer._parameters["weight_scale"] - replace_parameter(layer, "weight_scale_inv", tm_scales) - layer.input_scale = None - layer.sm70_fp8_turbomind = True - layer.sm70_fp8_channel_scale = True - 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()) - logger.info_once( - "SM70 compressed-tensors channel-FP8 TurboMind W8A16 dense " - "path enabled." - ) - return - if self.strategy == QuantizationStrategy.BLOCK: assert self.is_static_input_scheme is False # MarlinFP8ScaledMMLinearKernel uses "weight_scale_inv" for block @@ -317,83 +145,4 @@ def apply_weights( x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: - if getattr(layer, "sm70_fp8_turbomind", False): - if x.dtype != torch.float16: - raise RuntimeError( - "SM70 TurboMind channel-FP8 requires float16 activations, " - f"got {x.dtype}." - ) - out_shape = (*x.shape[:-1], layer.output_size_per_partition) - x_2d = x.reshape(-1, x.shape[-1]) - if x_2d.stride(-1) != 1: - x_2d = x_2d.contiguous() - out_2d = torch.empty( - (x_2d.shape[0], layer.output_size_per_partition), - device=x.device, - dtype=x.dtype, - ) - if x_2d.shape[0] == 0: - return out_2d.reshape(out_shape) - if getattr(layer, "sm70_fp8_qpn8", False): - sm70_ops.fp8_qpn8_dispatch_sm70_out( - out_2d, - int(layer.sm70_fp8_prefill_exact_dense_workspace_ptr), - x_2d, - layer.weight, - layer.weight_scale_inv, - int(layer.sm70_fp8_qpn8_split_k), - int(layer.sm70_fp8_qpn8_nacc), - bool(layer.sm70_fp8_qpn8_prefetch), - False, - ) - else: - sm70_ops.fp8_gemm_sm70_out( - out_2d, - x_2d, - layer.weight, - layer.weight_scale_inv, - 128, - layer.sm70_fp8_k_ld, - layer.sm70_fp8_q_ld, - False, - ) - if bias is not None: - out_2d.add_(bias) - return out_2d.reshape(out_shape) return self.linear_kernel.apply_weights(layer, x, bias) - - def apply_fused_silu_and_mul( - self, - layer: torch.nn.Module, - x: torch.Tensor, - ) -> torch.Tensor | None: - if not getattr(layer, "sm70_fp8_qpn8", False): - return None - if not getattr(layer, "sm70_fp8_gated_silu", False): - return None - if x.dtype != torch.float16: - raise RuntimeError( - "SM70 channel-FP8 QPN8 gated-SiLU requires float16 " - f"activations, got {x.dtype}." - ) - x_2d = x.reshape(-1, x.shape[-1]) - if x_2d.stride(-1) != 1: - x_2d = x_2d.contiguous() - out_features = layer.output_size_per_partition // 2 - out_2d = torch.empty( - (x_2d.shape[0], out_features), device=x.device, dtype=x.dtype - ) - if x_2d.shape[0] == 0: - return out_2d.reshape(*x.shape[:-1], out_features) - sm70_ops.fp8_qpn8_dispatch_sm70_out( - out_2d, - int(layer.sm70_fp8_prefill_exact_dense_workspace_ptr), - x_2d, - layer.weight, - layer.weight_scale_inv, - int(layer.sm70_fp8_qpn8_gated_split_k), - int(layer.sm70_fp8_qpn8_gated_nacc), - bool(layer.sm70_fp8_qpn8_gated_prefetch), - True, - ) - return out_2d.reshape(*x.shape[:-1], out_features) diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 8f8c975d85..25eb4e363f 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import os from typing import TYPE_CHECKING, Any import torch @@ -101,11 +100,6 @@ logger = init_logger(__name__) _SM70_FP8_PREFILL_DENSE_MIN_M = 3920 -_SM70_FP8_EXACT_8K_PREFILL_M = 8000 -_SM70_FP8_EXACT_8K_PREFILL_SHAPES = { - "in_proj_qkvz": (5120, 4096), - "qkv_proj": (5120, 3584), -} _SM70_FP8_PREFILL_DENSE_SHAPES = { "gate_up_proj": (5120, 8704), "down_proj": (4352, 5120), @@ -118,237 +112,20 @@ _SM70_FP8_PREFILL_DENSE_WORKSPACE_BYTES = ( _SM70_FP8_PREFILL_DENSE_WORKSPACE_ELEMENTS * torch.float16.itemsize ) -_SM70_FP8_QPN8_CONFIGS = { - # (K, N, fused gated-SiLU): (split-K, accumulator chains, prefetch codes) - (4352, 5120, False): (16, 1, False), - (1536, 5120, False): (16, 1, False), - (5120, 4096, False): (16, 2, False), - (5120, 3584, False): (16, 2, False), - (5120, 8704, False): (16, 1, True), - (5120, 8704, True): (8, 2, True), -} -_SM70_FP8_QPN8_EXTRA_SHAPES = { - "in_proj_qkvz": (5120, 4096), - "qkv_proj": (5120, 3584), -} -_SM70_FP8_QPN8_PP2_TP4_CONFIGS = { - # Real-weight speed winners remain numerically bounded for every admitted - # projection. The shared-expert gate/up role is deliberately absent below - # after its fused activation matched only 53.5% of TurboMind FP16 elements. - (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. - "fused_wqa_wkv": {(1, 4096, 1536)}, - "wq_b": {(4, 1024, 8192)}, - "wo_b": {(4, 2048, 4096)}, - "down_proj": {(4, 512, 4096)}, -} -_SM70_FP8_QPN8_PP2_TP4_WORKSPACE_ELEMENTS = max( - k * n for k, n, _ in _SM70_FP8_QPN8_PP2_TP4_CONFIGS -) -_SM70_FP8_QPN8_REQUIRED_OPS = ( - "fp8_qpn8_prepare_sm70", - "fp8_qpn8_dequantize_sm70_out", - "fp8_qpn8_prefill_sm70_out", - "fp8_qpn8_dispatch_sm70_out", - "fp8_qpn8_gemm_sm70_out", - "fp8_qpn8_gated_pair_sm70_out", -) -_SM70_FP8_QPN8_MAX_NUM_SEQS = 8 # Layers retain only data_ptr(), so this cache owns each allocation's lifetime. _sm70_fp8_prefill_dense_workspaces: dict[tuple[int, torch.dtype], torch.Tensor] = {} -_sm70_fp8_qpn8_pp2_tp4_workspaces: dict[tuple[int, torch.dtype], torch.Tensor] = {} - - -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 - if getattr(layer, "weight_block_size", None) != [128, 128]: - return False - suffix = getattr(layer, "prefix", "").rsplit(".", 1)[-1] - expected = _SM70_FP8_EXACT_8K_PREFILL_SHAPES.get(suffix) - return expected is not None and tuple(layer.weight.shape) == expected def _is_sm70_fp8_prefill_exact_dense_layer(layer: torch.nn.Module) -> bool: if getattr(layer, "tp_size", 1) != 4: return False - if getattr(layer, "weight_block_size", None) != [128, 128]: - return False suffix = getattr(layer, "prefix", "").rsplit(".", 1)[-1] expected = _SM70_FP8_PREFILL_DENSE_SHAPES.get(suffix) - if expected is None: - expected = _SM70_FP8_EXACT_8K_PREFILL_SHAPES.get(suffix) if expected is None: return False return tuple(layer.weight.shape) == expected -def _is_sm70_fp8_qpn8_layer(layer: torch.nn.Module) -> bool: - """Admit only shapes with an accepted bounded-workspace prefill route.""" - if getattr(layer, "tp_size", 1) != 4: - return False - if getattr(layer, "weight_block_size", None) != [128, 128]: - return False - suffix = getattr(layer, "prefix", "").rsplit(".", 1)[-1] - expected_kn = _SM70_FP8_PREFILL_DENSE_SHAPES.get(suffix) - if expected_kn is None: - expected_kn = _SM70_FP8_QPN8_EXTRA_SHAPES.get(suffix) - if expected_kn is None: - return False - # Checkpoint-native block-FP8 weights are [N, K]; the shared prefill - # workspace and QPN8 replacement parameter are [K, N]. - if tuple(reversed(layer.weight.shape)) != expected_kn: - return False - expected_k, expected_n = expected_kn - return bool( - getattr(layer, "input_size_per_partition", 0) == expected_k - and getattr(layer, "output_size_per_partition", 0) == expected_n - ) - - -def _is_sm70_fp8_qpn8_runtime_contract() -> bool: - """Admit only the no-MTP M range with a measured native decode kernel.""" - vllm_config = get_current_vllm_config() - scheduler_config = getattr(vllm_config, "scheduler_config", None) - max_num_seqs = int(getattr(scheduler_config, "max_num_seqs", 1)) - speculative_config = getattr(vllm_config, "speculative_config", None) - if max_num_seqs > _SM70_FP8_QPN8_MAX_NUM_SEQS: - return False - if speculative_config is None: - return True - # Speculative workloads require an explicit operator opt-in; admission is - # still based only on bounded concurrency and the layer tensor contract. - return os.getenv("VLLM_SM70_FP8_QPN8") is not None - - -def _sm70_fp8_qpn8_pp2_tp4_enabled() -> bool: - """Resolve the validated default while retaining explicit rollback.""" - generic_override = os.getenv("VLLM_SM70_FP8_QPN8") - if generic_override is not None and not envs.VLLM_SM70_FP8_QPN8: - return False - specific_override = os.getenv("VLLM_SM70_FP8_QPN8_PP2_TP4") - if specific_override is not None: - return envs.VLLM_SM70_FP8_QPN8_PP2_TP4 - if generic_override is not None: - return envs.VLLM_SM70_FP8_QPN8 - return envs.VLLM_SM70_FP8_QPN8_PP2_TP4 - - -def _is_sm70_fp8_qpn8_pp2_tp4_runtime_contract() -> bool: - """Require the measured serialized PP2 x TP4 single-request route.""" - 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 _sm70_fp8_qpn8_pp2_tp4_config( - layer: torch.nn.Module, *, gated_silu: bool -) -> tuple[int, int, bool] | None: - """Select by operator/tensor contract, never model or checkpoint identity.""" - if getattr(layer, "weight_block_size", None) != [128, 128]: - return None - suffix = getattr(layer, "prefix", "").rsplit(".", 1)[-1] - accepted = _SM70_FP8_QPN8_PP2_TP4_SHAPES.get(suffix) - if accepted is None: - return None - k_dim = int(getattr(layer, "input_size_per_partition", 0)) - n_dim = int(getattr(layer, "output_size_per_partition", 0)) - layer_contract = (int(getattr(layer, "tp_size", 1)), k_dim, n_dim) - if layer_contract not in accepted: - return None - if tuple(reversed(layer.weight.shape)) != (k_dim, n_dim): - return None - if gated_silu: - output_partitions = getattr(layer, "output_partition_sizes", None) - if ( - suffix != "gate_up_proj" - or not isinstance(output_partitions, list) - or output_partitions != [n_dim // 2, n_dim // 2] - ): - return None - return _SM70_FP8_QPN8_PP2_TP4_CONFIGS.get((k_dim, n_dim, gated_silu)) - - -def _sm70_fp8_qpn8_pp2_tp4_bmm_config( - layer: torch.nn.Module, -) -> tuple[int, int, bool] | None: - if ( - getattr(layer, "prefix", "").rsplit(".", 1)[-1] != "wo_a" - or getattr(layer, "weight_block_size", None) != [128, 128] - or int(getattr(layer, "tp_size", 1)) != 4 - or int(getattr(layer, "bmm_batch_size", 0)) != 2 - or int(getattr(layer, "input_size_per_partition", 0)) != 4096 - or int(getattr(layer, "output_size_per_partition", 0)) != 2048 - or tuple(layer.weight.shape) != (2048, 4096) - ): - return None - return _SM70_FP8_QPN8_PP2_TP4_CONFIGS[(4096, 1024, False)] - - -def _missing_sm70_fp8_qpn8_ops() -> list[str]: - return [ - name for name in _SM70_FP8_QPN8_REQUIRED_OPS if not hasattr(torch.ops._C, name) - ] - - def _get_sm70_fp8_prefill_exact_dense_workspace( weight: torch.Tensor, ) -> torch.Tensor | None: @@ -375,78 +152,6 @@ def _get_sm70_fp8_prefill_exact_dense_workspace( return workspace -def _get_sm70_fp8_qpn8_pp2_tp4_workspace( - weight: torch.Tensor, -) -> torch.Tensor | None: - """Allocate one bounded FP16 prefill fallback per device.""" - device_index = weight.device.index - if device_index is None: - device_index = torch.accelerator.current_device_index() - cache_key = (device_index, torch.float16) - workspace = _sm70_fp8_qpn8_pp2_tp4_workspaces.get(cache_key) - if workspace is not None: - return workspace - try: - workspace = torch.empty( - (_SM70_FP8_QPN8_PP2_TP4_WORKSPACE_ELEMENTS,), - dtype=torch.float16, - device=weight.device, - ) - except torch.OutOfMemoryError: - logger.warning_once( - "Insufficient memory for the bounded SM70 PP2 x TP4 QPN8 " - "prefill workspace; retaining TurboMind FP8." - ) - return None - _sm70_fp8_qpn8_pp2_tp4_workspaces[cache_key] = workspace - return workspace - - -def _sm70_fp8_prefill_visible_dense_mm( - input: torch.Tensor, - weight: torch.Tensor, - scales: torch.Tensor, - dense_weight_ptr: int | None, - *, - gated_silu: bool, - min_prefill_m: int, -) -> torch.Tensor | None: - """Expose the long-prefill dense MM to AsyncTP pattern matching. - - This diagnostic route intentionally keeps the accepted dequantization and - FP16 MM arithmetic while moving ``aten.mm`` out of the opaque C++ wrapper. - """ - if not envs.VLLM_SM70_FP8_PREFILL_VISIBLE_DENSE_MM: - return None - if dense_weight_ptr is None: - return None - if input.dtype != torch.float16 or input.shape[0] < min_prefill_m: - return None - - device_index = input.device.index - if device_index is None: - device_index = torch.accelerator.current_device_index() - workspace = _sm70_fp8_prefill_dense_workspaces.get((device_index, torch.float16)) - if workspace is None: - return None - if not torch.compiler.is_compiling() and workspace.data_ptr() != dense_weight_ptr: - return None - - dense_weight = workspace.narrow(0, 0, weight.numel()).view(weight.shape) - sm70_ops.fp8_sm70_dequantize_out(dense_weight, weight, scales, 128) - dense_out = torch.mm(input, dense_weight) - if not gated_silu: - return dense_out - - out = torch.empty( - (input.shape[0], dense_out.shape[1] // 2), - dtype=input.dtype, - device=input.device, - ) - sm70_ops.silu_and_mul_interleaved(out, dense_out) - return out - - class Fp8Config(QuantizationConfig): """Config class for FP8.""" @@ -829,15 +534,6 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: if weight_scale_inv.dtype != torch.float32: weight_scale_inv = weight_scale_inv.to(torch.float32) if getattr(layer, "is_bmm", False): - qpn8_bmm_config = ( - _sm70_fp8_qpn8_pp2_tp4_bmm_config(layer) - if _sm70_fp8_qpn8_pp2_tp4_enabled() - else None - ) - qpn8_bmm_runtime = bool( - qpn8_bmm_config is not None - and _is_sm70_fp8_qpn8_pp2_tp4_runtime_contract() - ) group_count = int(getattr(layer, "bmm_batch_size", 0)) if group_count <= 0 or weight.shape[0] % group_count != 0: raise RuntimeError( @@ -861,73 +557,6 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: f"{weight_scale_inv.shape[0]}." ) - if qpn8_bmm_config is not None and not qpn8_bmm_runtime: - logger.info_once( - "Grouped SM70 QPN8 retains TurboMind outside the " - "serialized PP2 x TP4 single-request contract." - ) - if qpn8_bmm_runtime: - missing_ops = _missing_sm70_fp8_qpn8_ops() - explicitly_enabled = ( - os.getenv("VLLM_SM70_FP8_QPN8") == "1" - or os.getenv("VLLM_SM70_FP8_QPN8_PP2_TP4") == "1" - ) - if missing_ops and explicitly_enabled: - raise RuntimeError( - "The explicitly enabled SM70 PP2 x TP4 QPN8 route " - f"requires source-built operators; missing: {missing_ops}." - ) - if missing_ops: - logger.warning_once( - "The default SM70 PP2 x TP4 QPN8 route is unavailable " - "in the loaded vllm._C; retaining TurboMind FP8." - ) - workspace = ( - None - if missing_ops - else _get_sm70_fp8_qpn8_pp2_tp4_workspace(weight) - ) - if workspace is not None: - qpn8_weights = [] - qpn8_scales = [] - for group_idx in range(group_count): - row_start = group_idx * rows_per_group - scale_start = group_idx * scale_rows_per_group - qpn8_weight, qpn8_scale = sm70_ops.fp8_qpn8_prepare_sm70( - weight[ - row_start : row_start + rows_per_group - ].contiguous(), - weight_scale_inv[ - scale_start : scale_start + scale_rows_per_group - ].contiguous(), - ) - qpn8_weights.append(qpn8_weight) - qpn8_scales.append(qpn8_scale) - replace_parameter(layer, "weight", torch.stack(qpn8_weights)) - replace_parameter( - layer, "weight_scale_inv", torch.stack(qpn8_scales) - ) - assert qpn8_bmm_config is not None - split_k, nacc, prefetch = qpn8_bmm_config - layer.input_scale = None - layer.sm70_fp8_turbomind = True - layer.sm70_fp8_qpn8 = True - layer.sm70_fp8_qpn8_bmm = True - layer.sm70_fp8_bmm = True - layer.sm70_fp8_bmm_groups = group_count - layer.sm70_fp8_bmm_output_size = rows_per_group - layer.sm70_fp8_qpn8_split_k = split_k - layer.sm70_fp8_qpn8_nacc = nacc - layer.sm70_fp8_qpn8_prefetch = prefetch - layer.sm70_fp8_prefill_exact_dense_workspace_ptr = ( - workspace.data_ptr() - ) - logger.info_once( - "Default SM70 grouped QPN8 enabled for the validated " - "serialized PP2 x TP4 tensor contract." - ) - return - prepared_weights = [] prepared_scales = [] metas = [] @@ -967,182 +596,12 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: layer.register_buffer("sm70_fp8_meta", first_meta, persistent=False) layer.sm70_fp8_k_ld = int(first_meta[0].item()) layer.sm70_fp8_q_ld = int(first_meta[1].item()) - if ( - envs.VLLM_SM70_FP8_GROUPED_BMM_DECODE - and group_count == 2 - and rows_per_group == 1024 - and int(layer.weight.shape[1]) == 4096 - and hasattr( - torch.ops._C, - "fp8_moe_gemm_sm70_per_expert_dispatch_out", - ) - and hasattr(torch.ops._C, "awq_moe_build_strided_ptrs") - ): - ptrs_w, ptrs_s = sm70_ops.awq_moe_build_strided_ptrs( - layer.weight, - layer.weight_scale_inv, - layer.sm70_fp8_k_ld, - layer.sm70_fp8_q_ld, - group_count, - ) - layer.register_buffer( - "sm70_fp8_bmm_grouped_ptrs_w", ptrs_w, persistent=False - ) - layer.register_buffer( - "sm70_fp8_bmm_grouped_ptrs_s", ptrs_s, persistent=False - ) - layer.register_buffer( - "sm70_fp8_bmm_grouped_offsets", - torch.arange( - group_count + 1, - dtype=torch.int32, - device=layer.weight.device, - ), - persistent=False, - ) - layer.sm70_fp8_bmm_grouped_decode = True - logger.info_once( - "SM70 FP8 one-launch grouped-BMM decode path enabled." - ) logger.info_once( "SM70 FP8 TurboMind grouped-BMM path enabled for DeepSeek V4." ) return is_gated_silu_layer = self._is_sm70_gated_silu_layer(layer) use_gated_silu = is_gated_silu_layer and envs.VLLM_SM70_FP8_DENSE_GATED_SILU - generic_qpn8_candidate = ( - envs.VLLM_SM70_FP8_QPN8 and _is_sm70_fp8_qpn8_layer(layer) - ) - pp2_tp4_qpn8_config = ( - _sm70_fp8_qpn8_pp2_tp4_config(layer, gated_silu=False) - if _sm70_fp8_qpn8_pp2_tp4_enabled() - else None - ) - pp2_tp4_qpn8_candidate = pp2_tp4_qpn8_config is not None - qpn8_candidate_layer = generic_qpn8_candidate or pp2_tp4_qpn8_candidate - qpn8_runtime = bool( - ( - pp2_tp4_qpn8_candidate - and _is_sm70_fp8_qpn8_pp2_tp4_runtime_contract() - ) - or (generic_qpn8_candidate and _is_sm70_fp8_qpn8_runtime_contract()) - ) - if qpn8_candidate_layer and not qpn8_runtime: - logger.info_once( - "The SM70 FP8 QPN8 route retains TurboMind unless its " - "bounded-concurrency runtime contract is explicit." - ) - if qpn8_candidate_layer and qpn8_runtime: - pp2_tp4_gated_config = None - if pp2_tp4_qpn8_candidate and use_gated_silu: - pp2_tp4_gated_config = _sm70_fp8_qpn8_pp2_tp4_config( - layer, gated_silu=True - ) - if pp2_tp4_gated_config is None: - raise RuntimeError( - "The SM70 PP2 x TP4 QPN8 gate/up layer violated " - "its fused-SiLU tensor contract." - ) - missing_ops = _missing_sm70_fp8_qpn8_ops() - if missing_ops: - explicitly_enabled = ( - os.getenv("VLLM_SM70_FP8_QPN8") == "1" - or os.getenv("VLLM_SM70_FP8_QPN8_PP2_TP4") == "1" - ) - if explicitly_enabled: - raise RuntimeError( - "The explicitly enabled SM70 QPN8 route requires " - f"source-built operators; missing: {missing_ops}." - ) - logger.warning_once( - "The default SM70 FP8 QPN8 route is unavailable in " - "the loaded vllm._C; retaining the TurboMind layout." - ) - - if missing_ops: - workspace = None - elif pp2_tp4_qpn8_candidate: - workspace = _get_sm70_fp8_qpn8_pp2_tp4_workspace(weight) - else: - workspace = _get_sm70_fp8_prefill_exact_dense_workspace(weight) - if not missing_ops and workspace is not None: - qpn8_codes, qpn8_scales = sm70_ops.fp8_qpn8_prepare_sm70( - weight, weight_scale_inv - ) - k_dim, n_dim = (int(dim) for dim in qpn8_codes.shape) - if pp2_tp4_qpn8_candidate: - assert pp2_tp4_qpn8_config is not None - split_k, nacc, prefetch = pp2_tp4_qpn8_config - else: - split_k, nacc, prefetch = _SM70_FP8_QPN8_CONFIGS[ - (k_dim, n_dim, False) - ] - replace_parameter(layer, "weight", qpn8_codes) - replace_parameter(layer, "weight_scale_inv", qpn8_scales) - layer.input_scale = None - layer.sm70_fp8_turbomind = True - layer.sm70_fp8_qpn8 = True - layer.sm70_fp8_qpn8_split_k = split_k - layer.sm70_fp8_qpn8_nacc = nacc - layer.sm70_fp8_qpn8_prefetch = prefetch - layer.sm70_fp8_prefill_exact_dense_workspace_ptr = ( - workspace.data_ptr() - ) - if use_gated_silu: - if pp2_tp4_qpn8_candidate: - assert pp2_tp4_gated_config is not None - gated_split_k, gated_nacc, gated_prefetch = ( - pp2_tp4_gated_config - ) - else: - gated_split_k, gated_nacc, gated_prefetch = ( - _SM70_FP8_QPN8_CONFIGS[(k_dim, n_dim, True)] - ) - layer.sm70_fp8_gated_silu = True - layer.sm70_fp8_gated_silu_primary = True - layer.sm70_fp8_qpn8_gated_split_k = gated_split_k - layer.sm70_fp8_qpn8_gated_nacc = gated_nacc - layer.sm70_fp8_qpn8_gated_prefetch = gated_prefetch - if pp2_tp4_qpn8_candidate: - logger.info_once( - "Default SM70 QPN8 enabled for the validated " - "serialized PP2 x TP4 operator contract." - ) - else: - logger.info_once( - "Memory-neutral SM70 FP8 QPN8 path enabled for " - "accepted TP4 block-FP8 operator shapes." - ) - return - if not missing_ops: - logger.warning_once( - "Insufficient memory for the SM70 FP8 QPN8 prefill " - "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, @@ -1170,66 +629,16 @@ 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 - and hasattr(torch.ops._C, "fp8_gemm_sm70_prefill_prescaled_out") - and _is_sm70_fp8_exact_8k_prefill_layer(layer) - ): - layer.register_buffer( - "sm70_fp8_prefill_prescaled_scales", - tm_scales.mul(256), - persistent=False, - ) - logger.info_once( - "SM70 block-FP8 exact-8K pre-scaled projection path enabled." - ) if ( envs.VLLM_SM70_FP8_PREFILL_EXACT_DENSE and hasattr(torch.ops._C, "fp8_gemm_sm70_prefill_dispatch_out") and _is_sm70_fp8_prefill_exact_dense_layer(layer) ): - is_exact_8k_projection = _is_sm70_fp8_exact_8k_prefill_layer(layer) workspace = _get_sm70_fp8_prefill_exact_dense_workspace(tm_weight) if workspace is not None: layer.sm70_fp8_prefill_exact_dense_workspace_ptr = ( workspace.data_ptr() ) - layer.sm70_fp8_prefill_exact_dense_min_m = ( - _SM70_FP8_EXACT_8K_PREFILL_M - if is_exact_8k_projection - else _SM70_FP8_PREFILL_DENSE_MIN_M - ) logger.info_once( "SM70 FP8 exact-dense prefill path enabled with a bounded " "85 MiB workspace." @@ -1349,77 +758,6 @@ def apply( bias: torch.Tensor | None = None, ) -> torch.Tensor: if getattr(layer, "sm70_fp8_turbomind", False): - if getattr(layer, "sm70_fp8_qpn8", False): - if x.dtype != torch.float16: - raise RuntimeError( - "SM70 FP8 QPN8 currently requires float16 activations, " - f"got {x.dtype}." - ) - if getattr(layer, "sm70_fp8_qpn8_bmm", False): - group_count = int(layer.sm70_fp8_bmm_groups) - output_size = int(layer.sm70_fp8_bmm_output_size) - if x.ndim < 2 or x.shape[-2] != group_count: - raise RuntimeError( - "SM70 grouped QPN8 input must end in [groups, K], got " - f"{tuple(x.shape)} for groups={group_count}." - ) - x_grouped = x.reshape(-1, group_count, x.shape[-1]) - x_by_group = x_grouped.transpose(0, 1).contiguous() - out_by_group = torch.empty( - (group_count, x_grouped.shape[0], output_size), - device=x.device, - dtype=x.dtype, - ) - if x_grouped.shape[0] == 0: - return out_by_group.transpose(0, 1).reshape( - *x.shape[:-2], group_count, output_size - ) - for group_idx in range(group_count): - sm70_ops.fp8_qpn8_dispatch_sm70_out( - out_by_group[group_idx], - int(layer.sm70_fp8_prefill_exact_dense_workspace_ptr), - x_by_group[group_idx], - layer.weight[group_idx], - layer.weight_scale_inv[group_idx], - int(layer.sm70_fp8_qpn8_split_k), - int(layer.sm70_fp8_qpn8_nacc), - bool(layer.sm70_fp8_qpn8_prefetch), - False, - ) - out = out_by_group.transpose(0, 1).reshape( - *x.shape[:-2], group_count, output_size - ) - if bias is not None: - out.add_(bias.view(group_count, output_size)) - return out - - out_shape = (*x.shape[:-1], layer.output_size_per_partition) - x_2d = x.reshape(-1, x.shape[-1]) - if x_2d.stride(-1) != 1: - x_2d = x_2d.contiguous() - out_2d = torch.empty( - (x_2d.shape[0], layer.output_size_per_partition), - device=x.device, - dtype=x.dtype, - ) - if x_2d.shape[0] == 0: - return out_2d.reshape(out_shape) - sm70_ops.fp8_qpn8_dispatch_sm70_out( - out_2d, - int(layer.sm70_fp8_prefill_exact_dense_workspace_ptr), - x_2d, - layer.weight, - layer.weight_scale_inv, - int(layer.sm70_fp8_qpn8_split_k), - int(layer.sm70_fp8_qpn8_nacc), - bool(layer.sm70_fp8_qpn8_prefetch), - False, - ) - out = out_2d.reshape(out_shape) - if bias is not None: - out.add_(bias) - return out - if getattr(layer, "sm70_fp8_bmm", False): group_count = int(layer.sm70_fp8_bmm_groups) output_size = int(layer.sm70_fp8_bmm_output_size) @@ -1435,34 +773,17 @@ def apply( device=x.device, dtype=x.dtype, ) - if ( - getattr(layer, "sm70_fp8_bmm_grouped_decode", False) - and x_grouped.shape[0] == 1 - ): - sm70_ops.fp8_moe_gemm_sm70_per_expert_dispatch_out( - out_by_group.reshape(group_count, output_size), - x_by_group.reshape(group_count, x.shape[-1]), - layer.sm70_fp8_bmm_grouped_offsets, - layer.sm70_fp8_bmm_grouped_ptrs_w, - layer.sm70_fp8_bmm_grouped_ptrs_s, - group_count, - x.shape[-1], - output_size, + for group_idx in range(group_count): + sm70_ops.fp8_gemm_sm70_out( + out_by_group[group_idx], + x_by_group[group_idx], + layer.weight[group_idx], + layer.weight_scale_inv[group_idx], 128, + layer.sm70_fp8_k_ld, + layer.sm70_fp8_q_ld, False, ) - else: - for group_idx in range(group_count): - sm70_ops.fp8_gemm_sm70_out( - out_by_group[group_idx], - x_by_group[group_idx], - layer.weight[group_idx], - layer.weight_scale_inv[group_idx], - 128, - layer.sm70_fp8_k_ld, - layer.sm70_fp8_q_ld, - False, - ) out = out_by_group.transpose(0, 1).reshape( *x.shape[:-2], group_count, output_size ) @@ -1482,42 +803,7 @@ def apply( prefill_workspace_ptr = getattr( layer, "sm70_fp8_prefill_exact_dense_workspace_ptr", None ) - 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", - _SM70_FP8_PREFILL_DENSE_MIN_M, - ) - visible_dense_out = _sm70_fp8_prefill_visible_dense_mm( - x_2d, - layer.weight, - layer.weight_scale_inv, - prefill_workspace_ptr, - gated_silu=False, - min_prefill_m=prefill_min_m, - ) - 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: + if prefill_workspace_ptr is not None and x_2d.dtype == torch.float16: sm70_ops.fp8_gemm_sm70_prefill_dispatch_out( out_2d, prefill_workspace_ptr, @@ -1528,22 +814,7 @@ def apply( layer.sm70_fp8_k_ld, layer.sm70_fp8_q_ld, False, - prefill_min_m, - ) - elif ( - prefill_prescaled_scales is not None - and envs.VLLM_SM70_FP8_PREFILL_FAST_SELECTOR - and envs.VLLM_SM70_FP8_PREFILL_PRESCALED - and x_2d.shape[0] == _SM70_FP8_EXACT_8K_PREFILL_M - ): - sm70_ops.fp8_gemm_sm70_prefill_prescaled_out( - out_2d, - x_2d, - layer.weight, - prefill_prescaled_scales, - 128, - layer.sm70_fp8_k_ld, - layer.sm70_fp8_q_ld, + _SM70_FP8_PREFILL_DENSE_MIN_M, ) else: sm70_ops.fp8_gemm_sm70_out( @@ -1618,36 +889,6 @@ def apply_fused_silu_and_mul( layer: torch.nn.Module, x: torch.Tensor, ) -> torch.Tensor | None: - if getattr(layer, "sm70_fp8_qpn8", False): - if not getattr(layer, "sm70_fp8_gated_silu", False): - return None - if x.dtype != torch.float16: - raise RuntimeError( - "SM70 FP8 QPN8 gated-SiLU requires float16 activations, " - f"got {x.dtype}." - ) - x_2d = x.reshape(-1, x.shape[-1]) - if x_2d.stride(-1) != 1: - x_2d = x_2d.contiguous() - out_features = layer.output_size_per_partition // 2 - out_2d = torch.empty( - (x_2d.shape[0], out_features), device=x.device, dtype=x.dtype - ) - if x_2d.shape[0] == 0: - return out_2d.reshape(*x.shape[:-1], out_features) - sm70_ops.fp8_qpn8_dispatch_sm70_out( - out_2d, - int(layer.sm70_fp8_prefill_exact_dense_workspace_ptr), - x_2d, - layer.weight, - layer.weight_scale_inv, - int(layer.sm70_fp8_qpn8_gated_split_k), - int(layer.sm70_fp8_qpn8_gated_nacc), - bool(layer.sm70_fp8_qpn8_gated_prefetch), - True, - ) - return out_2d.reshape(*x.shape[:-1], out_features) - if not getattr(layer, "sm70_fp8_gated_silu", False): return None if not getattr(layer, "sm70_fp8_turbomind", False): @@ -1676,25 +917,6 @@ def apply_fused_silu_and_mul( layer, "sm70_fp8_prefill_exact_dense_workspace_ptr", None ) if prefill_workspace_ptr is not None and x_2d.dtype == torch.float16: - visible_dense_out = _sm70_fp8_prefill_visible_dense_mm( - x_2d, - weight, - scales, - prefill_workspace_ptr, - gated_silu=True, - min_prefill_m=getattr( - layer, - "sm70_fp8_prefill_exact_dense_min_m", - _SM70_FP8_PREFILL_DENSE_MIN_M, - ), - ) - if visible_dense_out is not None: - return visible_dense_out.reshape(*x.shape[:-1], out_features) - min_prefill_m = getattr( - layer, - "sm70_fp8_prefill_exact_dense_min_m", - _SM70_FP8_PREFILL_DENSE_MIN_M, - ) sm70_ops.fp8_gemm_sm70_prefill_dispatch_out( out_2d, prefill_workspace_ptr, @@ -1705,7 +927,7 @@ def apply_fused_silu_and_mul( k_ld, q_ld, True, - min_prefill_m, + _SM70_FP8_PREFILL_DENSE_MIN_M, ) return out_2d.reshape(*x.shape[:-1], out_features) sm70_ops.fp8_gemm_sm70_out( diff --git a/vllm/model_executor/layers/quantization/input_quant_fp8.py b/vllm/model_executor/layers/quantization/input_quant_fp8.py index a6cbe801e6..c32904339a 100644 --- a/vllm/model_executor/layers/quantization/input_quant_fp8.py +++ b/vllm/model_executor/layers/quantization/input_quant_fp8.py @@ -227,7 +227,9 @@ 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..eabaf62be7 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -8,7 +8,6 @@ from torch.nn.parameter import Parameter import vllm.model_executor.layers.fused_moe.modular_kernel as mk -from vllm import envs from vllm.config import get_current_vllm_config from vllm.logger import init_logger from vllm.model_executor.kernels.linear import ( @@ -49,7 +48,6 @@ UnquantizedLinearMethod, ) from vllm.model_executor.layers.quantization import QuantizationMethods -from vllm.model_executor.layers.quantization import sm70_turbomind as sm70_tm from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, QuantizeMethodBase, @@ -453,10 +451,6 @@ def __init__(self, quant_config: ModelOptFp8Config) -> None: self.quant_config = quant_config self.out_dtype = torch.get_default_dtype() self.input_dtype = get_current_vllm_config().model_config.dtype - self.use_sm70_fp8_turbomind = ( - sm70_tm.is_exact_sm70_cuda_platform() - and sm70_tm.use_turbomind(envs.VLLM_SM70_FP8_TURBOMIND) - ) def create_weights( self, @@ -506,9 +500,6 @@ def create_weights( scale[:] = torch.finfo(torch.float32).min layer.register_parameter("input_scale", scale) - if self.use_sm70_fp8_turbomind: - return - self.fp8_linear = init_fp8_linear_kernel( activation_quant_key=kFp8StaticTensorSym, weight_quant_key=kFp8StaticTensorSym, @@ -525,40 +516,6 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: max_w_scale, weight = requantize_with_max_scale( layer.weight, layer.weight_scale, layer.logical_widths ) - - if self.use_sm70_fp8_turbomind: - if not hasattr(torch.ops._C, "fp8_sm70_prepare"): - raise RuntimeError( - "ModelOpt FP8 on SM70 requires the TurboMind extension " - "with fp8_sm70_prepare." - ) - if self.input_dtype != torch.float16: - raise RuntimeError( - "ModelOpt FP8 TurboMind on SM70 requires FP16 activations, " - f"got {self.input_dtype}." - ) - - from vllm import _sm70_ops as sm70_ops - - block_size = 128 - out_blocks = (weight.shape[0] + block_size - 1) // block_size - in_blocks = (weight.shape[1] + block_size - 1) // block_size - block_scales = ( - max_w_scale.reshape(1, 1).expand(out_blocks, in_blocks).contiguous() - ) - tm_weight, tm_scales, meta = sm70_ops.fp8_sm70_prepare( - weight, block_scales, block_size - ) - replace_parameter(layer, "weight", tm_weight) - replace_parameter(layer, "weight_scale", tm_scales) - layer.input_scale = None - layer.sm70_modelopt_fp8_turbomind = True - layer.sm70_modelopt_fp8_group_size = block_size - layer.sm70_modelopt_fp8_k_ld = int(meta[0].item()) - layer.sm70_modelopt_fp8_q_ld = int(meta[1].item()) - logger.info_once("SM70 ModelOpt FP8 TurboMind W8A16 dense path enabled.") - return - layer.weight = Parameter(weight.t(), requires_grad=False) layer.weight_scale = Parameter(max_w_scale, requires_grad=False) layer.input_scale = Parameter(layer.input_scale.max(), requires_grad=False) @@ -570,31 +527,6 @@ def apply( x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: - if getattr(layer, "sm70_modelopt_fp8_turbomind", False): - if x.dtype != torch.float16: - raise TypeError( - "ModelOpt FP8 TurboMind on SM70 requires FP16 activations." - ) - from vllm import _sm70_ops as sm70_ops - - x_2d = x.reshape(-1, x.shape[-1]) - out = torch.empty( - (x_2d.shape[0], layer.output_size_per_partition), - dtype=x.dtype, - device=x.device, - ) - sm70_ops.fp8_gemm_sm70_out( - out, - x_2d, - layer.weight, - layer.weight_scale, - layer.sm70_modelopt_fp8_group_size, - layer.sm70_modelopt_fp8_k_ld, - layer.sm70_modelopt_fp8_q_ld, - ) - if bias is not None: - out.add_(bias) - return out.reshape(*x.shape[:-1], layer.output_size_per_partition) return self.fp8_linear.apply_weights(layer, x, bias) @@ -1070,51 +1002,6 @@ def apply( ModelOptFp8Config.KVCacheMethodCls = ModelOptKVCacheMethod -def _explicit_nvfp4_emulation_requested() -> bool: - """True when an explicit NVFP4 software-emulation backend is selected. - - Matches compressed-tensors W4A4 NVFP4: emulation is a proven software - path, so SM70 may load. Cutlass/FlashInfer are not. - """ - if envs.VLLM_USE_NVFP4_CT_EMULATIONS or envs.VLLM_NVFP4_GEMM_BACKEND == "emulation": - return True - - from vllm.config import get_current_vllm_config_or_none - - vllm_config = get_current_vllm_config_or_none() - return ( - vllm_config is not None - and vllm_config.kernel_config.linear_backend == "emulation" - ) - - -def _try_prepare_sm70_modelopt_nvfp4(layer: torch.nn.Module) -> bool: - """Prepare exact-SM70 TurboMind NVFP4 weights for a ModelOpt linear. - - ModelOpt stores ``weight_scale_2`` / ``weight_global_scale`` as - ``amax / (6 * 448)`` (the Marlin multiplier). TurboMind combine is - ``block * global``. Do not infer convention from scale magnitude. - """ - if not sm70_tm.should_prepare_turbomind( - layer.weight, envs.VLLM_SM70_NVFP4_TURBOMIND - ): - return False - logger.info_once( - "SM70 ModelOpt NVFP4 TurboMind dense path enabled " - "(weight-only; activations remain half)." - ) - sm70_tm.prepare_nvfp4_linear(layer) - layer.weight = Parameter( - torch.empty(0, dtype=torch.uint8, device=layer.weight.device), - requires_grad=False, - ) - layer.weight_scale = Parameter( - torch.empty(0, dtype=torch.float8_e4m3fn, device=layer.weight_scale.device), - requires_grad=False, - ) - return True - - class ModelOptNvFp4Config(ModelOptQuantConfigBase): """Config class for ModelOpt FP4.""" @@ -1163,18 +1050,6 @@ def get_supported_act_dtypes(self) -> list[torch.dtype]: @classmethod def get_min_capability(cls) -> int: - # Do not unconditionally lower the class-wide gate to 70. W4A4 - # (quant_algo=NVFP4) and W4A16_NVFP4 share this config; SM70 is - # admitted only when a proven software backend is selected. - # Exact-device routing still happens later via - # sm70_tm.should_prepare_turbomind (capability == (7, 0)). - if ( - sm70_tm.use_turbomind(envs.VLLM_SM70_NVFP4_TURBOMIND) - or sm70_tm.forces_marlin() - ): - return 70 - if _explicit_nvfp4_emulation_requested(): - return 70 return 75 @classmethod @@ -1345,9 +1220,6 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: (1.0 / layer.input_global_scale).to(torch.float32), requires_grad=False ) - if _try_prepare_sm70_modelopt_nvfp4(layer): - return - # Convert layer to NVFP4 linear kernel format self.kernel.process_weights_after_loading(layer) @@ -1357,8 +1229,6 @@ def apply( x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: - if sm70_tm.has_prepared_linear(layer): - return sm70_tm.apply_prepared_linear(layer, x, bias) return self.kernel.apply_weights(layer=layer, x=x, bias=bias) @@ -1490,16 +1360,13 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # Rename weight_scale_2 -> weight_global_scale. NO reciprocation: # ModelOpt already stores amax/2688, which is exactly what Marlin - # and SM70 TurboMind consume. This is checkpoint-format provenance, - # not a magnitude heuristic (see PR 166 rejection). + # consumes via nvfp4_marlin_process_global_scale (called inside the + # Marlin adapter's process_weights_after_loading). layer.weight_global_scale = Parameter( layer.weight_scale_2.max().to(torch.float32), requires_grad=False ) del layer.weight_scale_2 - if _try_prepare_sm70_modelopt_nvfp4(layer): - return - self.kernel.process_weights_after_loading(layer) def apply( @@ -1508,8 +1375,6 @@ def apply( x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: - if sm70_tm.has_prepared_linear(layer): - return sm70_tm.apply_prepared_linear(layer, x, bias) return self.kernel.apply_weights(layer=layer, x=x, bias=bias) @@ -2352,12 +2217,6 @@ def get_supported_act_dtypes(self) -> list[torch.dtype]: @classmethod def get_min_capability(cls) -> int: - if ( - sm70_tm.is_exact_sm70_cuda_platform() - and sm70_tm.use_turbomind(envs.VLLM_SM70_FP8_TURBOMIND) - and sm70_tm.use_turbomind(envs.VLLM_SM70_NVFP4_TURBOMIND) - ): - return 70 return 89 @classmethod @@ -2553,20 +2412,6 @@ def get_quant_method( moe_config=layer.moe_config, ) if quant_algo == "W4A16_NVFP4": - if sm70_tm.is_exact_sm70_cuda_platform(): - if not sm70_tm.should_use_nvfp4_moe_turbomind(): - raise NotImplementedError( - "ModelOpt W4A16 NVFP4 MoE on SM70 requires the " - "TurboMind backend." - ) - from vllm.model_executor.layers.quantization.nvfp4_sm70_moe import ( - ModelOptNvFp4SM70MoEMethod, - ) - - return ModelOptNvFp4SM70MoEMethod( - quant_config=self.w4a16_nvfp4_config, - moe_config=layer.moe_config, - ) return ModelOptNvFp4FusedMoE( quant_config=self.w4a16_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 dc39f8426a..6e255a5150 100644 --- a/vllm/model_executor/layers/quantization/mxfp4_sm70_moe.py +++ b/vllm/model_executor/layers/quantization/mxfp4_sm70_moe.py @@ -64,27 +64,6 @@ def _mxfp4_grouped_m8_enabled() -> bool: return bool(envs.VLLM_SM70_MXFP4_MOE_GROUPED_M8) -def _mxfp4_grouped_verifier_enabled() -> bool: - return bool(envs.VLLM_SM70_MXFP4_MOE_GROUPED_VERIFIER) - - -def _mxfp4_grouped_verifier_for_tokens(num_tokens: int) -> bool: - return bool( - (num_tokens == _GRAPH_SAFE_MAX_TOKENS and _mxfp4_grouped_m8_enabled()) - or ( - 1 < num_tokens <= _GRAPH_SAFE_MAX_TOKENS - and _mxfp4_grouped_verifier_enabled() - ) - ) - - -def _mxfp4_grouped_m8_expert_rows_enabled() -> bool: - return bool( - (_mxfp4_grouped_m8_enabled() or _mxfp4_grouped_verifier_enabled()) - and envs.VLLM_SM70_MXFP4_MOE_GROUPED_M8_EXPERT_ROWS - ) - - @triton.jit def _compact_sorted_experts_kernel( sorted_expert_ids_ptr, @@ -161,13 +140,7 @@ def _select_mxfp4_stage_dispatch( # tail entries as zero-row experts, avoiding a host readback of the # dynamic unique-expert count. graph_expert_slots = num_tokens * _DEEPSEEK_V4_FLASH_TOP_K - if _mxfp4_grouped_verifier_for_tokens(num_tokens): - if _mxfp4_grouped_m8_expert_rows_enabled(): - return ( - buffers["compact_expert_offsets"], - buffers["active_expert_ids"], - graph_expert_slots, - ) + if num_tokens == _GRAPH_SAFE_MAX_TOKENS and _mxfp4_grouped_m8_enabled(): return ( buffers["slot_expert_offsets"], buffers["permuted_experts_id"], @@ -185,18 +158,6 @@ def _select_mxfp4_stage_dispatch( return buffers["expert_offsets"], buffers["dense_expert_ids"], num_experts -def _select_mxfp4_direct_order_offsets( - buffers: dict[str, torch.Tensor], -) -> torch.Tensor: - """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 one row per route, - so it must use the separately maintained slot offsets instead. - """ - return buffers["slot_expert_offsets"] - - def validate_mxfp4_sm70_moe_contract( *, global_num_experts: int, @@ -355,11 +316,6 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: ) if envs.VLLM_SM70_MXFP4_MOE_DIRECT_TOP6_DECODE: required_ops += ("mxfp4_moe_single_token_prepare_w13_sm70_out",) - if ( - envs.VLLM_SM70_MXFP4_MOE_DIRECT_TOP6_DECODE - and envs.VLLM_SM70_MXFP4_MOE_DIRECT_ORDER_DECODE - ): - required_ops += ("awq_moe_single_token_weighted_reduce_out",) missing_ops = [name for name in required_ops if not hasattr(torch.ops._C, name)] if missing_ops: raise RuntimeError( @@ -743,49 +699,6 @@ def apply( and layer.expert_map is None and layer.local_num_experts == layer.global_num_experts ) - direct_order = ( - direct_top6 - and envs.VLLM_SM70_MXFP4_MOE_DIRECT_ORDER_DECODE - and topk_ids.dtype == torch.int32 - and topk_ids.is_contiguous() - ) - if direct_order: - route_ids = topk_ids.view(-1) - route_offsets = _select_mxfp4_direct_order_offsets(buffers) - sm70_ops.mxfp4_moe_dense_stage_sm70_out( - buffers["gate_up"], - x, - route_offsets, - route_ids, - layer.w13_strided_ptrs_w, - layer.w13_strided_ptrs_s, - _DEEPSEEK_V4_FLASH_TOP_K, - layer.sm70_mxfp4_w13_k_dim, - layer.sm70_mxfp4_w13_n_dim, - layer.sm70_mxfp4_group_size, - ) - self._apply_swiglu(layer, buffers["intermediate"], buffers["gate_up"]) - sm70_ops.mxfp4_moe_dense_stage_sm70_out( - buffers["sorted_output"], - buffers["intermediate"], - route_offsets, - route_ids, - layer.w2_strided_ptrs_w, - layer.w2_strided_ptrs_s, - _DEEPSEEK_V4_FLASH_TOP_K, - layer.sm70_mxfp4_w2_k_dim, - layer.sm70_mxfp4_w2_n_dim, - layer.sm70_mxfp4_group_size, - ) - 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, - ) - return output if direct_top6: sm70_ops.mxfp4_moe_single_token_prepare_w13_sm70_out( buffers["gate_up"], @@ -854,8 +767,7 @@ def apply( num_tokens > 1 and num_tokens <= _mxfp4_active_expert_max_tokens() and not ( - _mxfp4_grouped_verifier_for_tokens(num_tokens) - and not _mxfp4_grouped_m8_expert_rows_enabled() + num_tokens == _GRAPH_SAFE_MAX_TOKENS and _mxfp4_grouped_m8_enabled() ) and layer.expert_map is None and layer.local_num_experts == layer.global_num_experts diff --git a/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py b/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py deleted file mode 100644 index 7eb31d9ce3..0000000000 --- a/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py +++ /dev/null @@ -1,635 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Native SM70 TurboMind NVFP4 MoE for Qwen3.6-35B-A3B. - -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 -load time, repacks both tensors for TurboMind, and never materializes an FP16 -expert-weight copy. -""" - -from __future__ import annotations - -from typing import Final - -import torch -from torch.nn import Parameter - -from vllm import _sm70_ops as sm70_ops -from vllm.logger import init_logger -from vllm.model_executor.layers.fused_moe import ( - FusedMoEConfig, - FusedMoEMethodBase, - FusedMoEQuantConfig, - MoEActivation, - RoutedExperts, - SharedExperts, -) -from vllm.model_executor.layers.quantization.modelopt import ( - ModelOptNvFp4Config, - ModelOptNvFp4FusedMoE, -) -from vllm.model_executor.layers.quantization.sm70_turbomind import ( - NVFP4_GROUP_SIZE, - is_exact_sm70_cuda, - unpack_mxfp4_weight, -) -from vllm.triton_utils import tl, triton - -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) -_GRAPH_SAFE_MAX_TOKENS: Final = 18 -_COMPACT_GROUPED_MAX_TOKENS: Final = 10 - - -@triton.jit -def _prepare_compact_slot_groups_kernel( - sorted_expert_ids_ptr, - compact_offsets_ptr, - active_expert_ids_ptr, - TOTAL_SLOTS: tl.constexpr, - BLOCK: tl.constexpr, -): - offsets = tl.arange(0, BLOCK) - valid = offsets < TOTAL_SLOTS - expert_ids = tl.load( - sorted_expert_ids_ptr + offsets, - mask=valid, - other=-1, - ) - tl.store( - compact_offsets_ptr + offsets, - offsets, - mask=offsets <= TOTAL_SLOTS, - ) - tl.store( - active_expert_ids_ptr + offsets, - expert_ids, - mask=valid, - ) - - -def _prepare_compact_slot_groups( - sorted_expert_ids: torch.Tensor, - compact_offsets: torch.Tensor, - active_expert_ids: torch.Tensor, -) -> None: - total_slots = sorted_expert_ids.numel() - max_slots = _COMPACT_GROUPED_MAX_TOKENS * _QWEN36_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) - # TurboMind's compact grouped dispatch forces one row per group. Keep each - # routed slot independent even when adjacent slots select the same expert; - # coalescing duplicate expert IDs would make the forced one-row scheduler - # silently skip or miscompute the additional rows. - _prepare_compact_slot_groups_kernel[(1,)]( - sorted_expert_ids, - compact_offsets, - active_expert_ids, - TOTAL_SLOTS=total_slots, - BLOCK=block, - num_warps=1, - ) - - -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: - raise NotImplementedError( - "SM70 TurboMind NVFP4 MoE currently supports tensor parallel " - f"sizes {_QWEN36_SUPPORTED_TP_SIZES}, got {moe.tp_size}." - ) - local_intermediate = moe.intermediate_size_per_partition - if local_intermediate <= 0 or local_intermediate % NVFP4_GROUP_SIZE: - raise NotImplementedError( - "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: - 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}." - ) - if moe.moe_parallel_config.use_all2all_kernels: - raise NotImplementedError( - "SM70 TurboMind NVFP4 MoE does not support DP+EP all-to-all." - ) - - -def _validate_weight_layout(layer: RoutedExperts) -> None: - local_experts = int(layer.local_num_experts) - hidden = int(layer.moe_config.hidden_dim) - intermediate = int(layer.moe_config.intermediate_size_per_partition) - expected = { - "w13_weight": (local_experts, 2 * intermediate, hidden // 2), - "w13_weight_scale": ( - local_experts, - 2 * intermediate, - hidden // NVFP4_GROUP_SIZE, - ), - "w13_weight_scale_2": (local_experts, 2), - "w2_weight": (local_experts, hidden, intermediate // 2), - "w2_weight_scale": ( - local_experts, - hidden, - intermediate // NVFP4_GROUP_SIZE, - ), - "w2_weight_scale_2": (local_experts,), - } - tensors = {name: getattr(layer, name) for name in expected} - for name, shape in expected.items(): - if tuple(tensors[name].shape) != shape: - raise ValueError( - f"SM70 NVFP4 MoE layout mismatch for {name}: " - f"expected {shape}, got {tuple(tensors[name].shape)}." - ) - if layer.w13_weight.dtype != torch.uint8 or layer.w2_weight.dtype != torch.uint8: - raise TypeError("SM70 NVFP4 MoE requires packed uint8 expert weights.") - if ( - layer.w13_weight_scale.dtype != torch.float8_e4m3fn - or layer.w2_weight_scale.dtype != torch.float8_e4m3fn - ): - raise TypeError("SM70 NVFP4 MoE requires FP8 E4M3 block scales.") - - -class ModelOptNvFp4SM70MoEMethod(ModelOptNvFp4FusedMoE): - """Qwen3.6 ModelOpt W4A16_NVFP4 experts on native TurboMind SM70.""" - - def __init__( - self, - quant_config: ModelOptNvFp4Config, - moe_config: FusedMoEConfig, - ) -> None: - FusedMoEMethodBase.__init__(self, moe_config) - if quant_config.quant_method != "W4A16_NVFP4": - raise NotImplementedError( - "SM70 TurboMind ModelOpt NVFP4 MoE currently requires " - "W4A16_NVFP4 checkpoint weights." - ) - self.quant_config = quant_config - self.use_a16 = True - self.use_global_sf = False - validate_nvfp4_sm70_moe_contract(moe_config) - - @property - def supports_eplb(self) -> bool: - return False - - def maybe_make_prepare_finalize( - self, - routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, - ) -> None: - # This method owns routing, TurboMind expert GEMMs, and unpermutation. - # Do not wrap it in the generic ModelOpt modular-kernel path. - del routing_tables - return None - - def process_weights_after_loading(self, layer: RoutedExperts) -> None: - required_ops = ( - "nvfp4_sm70_prepare", - "nvfp4_moe_dense_stage_sm70_out", - "awq_moe_build_strided_ptrs", - ) - missing = [name for name in required_ops if not hasattr(torch.ops._C, name)] - if missing: - raise RuntimeError( - "Qwen3.6 NVFP4 MoE on SM70 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." - ) - if self.moe.has_bias: - raise NotImplementedError("SM70 NVFP4 MoE does not support expert bias.") - if layer.activation != MoEActivation.SILU: - raise NotImplementedError( - "SM70 NVFP4 MoE currently supports SwiGLU/SILU only." - ) - if layer.apply_router_weight_on_input: - raise NotImplementedError( - "SM70 NVFP4 MoE does not support router weights on input." - ) - if layer.expert_map is not None: - raise NotImplementedError( - "SM70 NVFP4 MoE currently requires fully replicated experts." - ) - if layer.local_num_experts != layer.global_num_experts: - raise NotImplementedError( - "SM70 NVFP4 MoE currently requires local and global experts to match." - ) - - validate_nvfp4_sm70_moe_contract(layer.moe_config) - _validate_weight_layout(layer) - num_experts = int(layer.local_num_experts) - hidden = int(layer.moe_config.hidden_dim) - intermediate = int(layer.moe_config.intermediate_size_per_partition) - - w13_tm_weights: list[torch.Tensor] = [] - w13_tm_scales: list[torch.Tensor] = [] - w13_meta: list[torch.Tensor] = [] - w2_tm_weights: list[torch.Tensor] = [] - w2_tm_scales: list[torch.Tensor] = [] - w2_meta: list[torch.Tensor] = [] - for expert_id in range(num_experts): - w13_packed = unpack_mxfp4_weight(layer.w13_weight[expert_id].data) - w13_scales = layer.w13_weight_scale[expert_id].float().clone() - w13_global = layer.w13_weight_scale_2[expert_id].float() - w13_scales[:intermediate].mul_(w13_global[0]) - w13_scales[intermediate:].mul_(w13_global[1]) - prepared_w13 = sm70_ops.nvfp4_sm70_prepare( - w13_packed, - w13_scales.half().t().contiguous(), - NVFP4_GROUP_SIZE, - ) - w13_tm_weights.append(prepared_w13[0]) - w13_tm_scales.append(prepared_w13[1]) - w13_meta.append(prepared_w13[2]) - - w2_packed = unpack_mxfp4_weight(layer.w2_weight[expert_id].data) - w2_scales = ( - layer.w2_weight_scale[expert_id].float() - * layer.w2_weight_scale_2[expert_id].float() - ) - prepared_w2 = sm70_ops.nvfp4_sm70_prepare( - w2_packed, - w2_scales.half().t().contiguous(), - NVFP4_GROUP_SIZE, - ) - w2_tm_weights.append(prepared_w2[0]) - w2_tm_scales.append(prepared_w2[1]) - w2_meta.append(prepared_w2[2]) - - layer.w13_tm_weight = Parameter( - torch.stack(w13_tm_weights), requires_grad=False - ) - layer.w13_tm_scales = Parameter(torch.stack(w13_tm_scales), requires_grad=False) - layer.w2_tm_weight = Parameter(torch.stack(w2_tm_weights), requires_grad=False) - layer.w2_tm_scales = Parameter(torch.stack(w2_tm_scales), requires_grad=False) - - w13_k_ld = int(w13_meta[0][0].item()) - 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()) - w13_ptrs = sm70_ops.awq_moe_build_strided_ptrs( - layer.w13_tm_weight, - layer.w13_tm_scales, - w13_k_ld, - w13_q_ld, - num_experts, - ) - w2_ptrs = sm70_ops.awq_moe_build_strided_ptrs( - layer.w2_tm_weight, - layer.w2_tm_scales, - w2_k_ld, - w2_q_ld, - num_experts, - ) - layer.w13_strided_ptrs_w = Parameter(w13_ptrs[0], requires_grad=False) - layer.w13_strided_ptrs_s = Parameter(w13_ptrs[1], requires_grad=False) - layer.w2_strided_ptrs_w = Parameter(w2_ptrs[0], requires_grad=False) - layer.w2_strided_ptrs_s = Parameter(w2_ptrs[1], requires_grad=False) - - layer.sm70_nvfp4_moe = True - layer.sm70_nvfp4_num_experts = num_experts - layer.sm70_nvfp4_hidden_size = hidden - layer.sm70_nvfp4_intermediate_size = intermediate - layer.sm70_nvfp4_w13_k_dim = hidden - layer.sm70_nvfp4_w13_n_dim = 2 * intermediate - layer.sm70_nvfp4_w2_k_dim = intermediate - layer.sm70_nvfp4_w2_n_dim = hidden - layer.sm70_nvfp4_group_size = NVFP4_GROUP_SIZE - layer.sm70_nvfp4_graph_safe_max_tokens = _GRAPH_SAFE_MAX_TOKENS - layer.sm70_nvfp4_compact_grouped_max_tokens = _COMPACT_GROUPED_MAX_TOKENS - self._allocate_graph_safe_decode_buffers(layer) - - del layer.w13_weight - del layer.w13_weight_scale - del layer.w13_weight_scale_2 - del layer.w13_input_scale - del layer.w2_weight - del layer.w2_weight_scale - 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).", - num_experts, - _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 - experts = int(layer.sm70_nvfp4_num_experts) - hidden = int(layer.sm70_nvfp4_hidden_size) - intermediate = int(layer.sm70_nvfp4_intermediate_size) - - layer._nvfp4_sm70_output = torch.empty( - _GRAPH_SAFE_MAX_TOKENS, hidden, dtype=torch.float16, device=device - ) - layer._nvfp4_sm70_permuted_input = torch.empty( - max_slots, hidden, dtype=torch.float16, device=device - ) - layer._nvfp4_sm70_gate_up = torch.empty( - max_slots, 2 * intermediate, dtype=torch.float16, device=device - ) - layer._nvfp4_sm70_intermediate = torch.empty( - max_slots, intermediate, dtype=torch.float16, device=device - ) - layer._nvfp4_sm70_sorted_output = torch.empty( - max_slots, hidden, dtype=torch.float16, device=device - ) - layer._nvfp4_sm70_expert_offsets = torch.empty( - experts + 1, dtype=torch.int32, device=device - ) - layer._nvfp4_sm70_expert_offsets64 = torch.empty( - experts + 1, dtype=torch.int64, device=device - ) - layer._nvfp4_sm70_inv_permuted_idx = torch.empty( - _GRAPH_SAFE_MAX_TOKENS, - _QWEN36_TOP_K, - dtype=torch.int32, - device=device, - ) - layer._nvfp4_sm70_topk_ids = torch.empty( - _GRAPH_SAFE_MAX_TOKENS, - _QWEN36_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) - layer._nvfp4_sm70_permuted_idx = torch.empty( - max_slots, dtype=torch.int32, device=device - ) - layer._nvfp4_sm70_permuted_experts_id = torch.empty( - max_slots, dtype=torch.int32, device=device - ) - layer._nvfp4_sm70_sorted_row_idx = torch.empty( - max_slots, dtype=torch.int32, device=device - ) - layer._nvfp4_sm70_topk_ids_for_sort = torch.empty( - max_slots, dtype=torch.int32, device=device - ) - workspace_size = torch.ops._moe_C.moe_permute_sort_workspace_size( - max_slots, layer.global_num_experts - ) - layer._nvfp4_sm70_sort_workspace = torch.empty( - workspace_size, dtype=torch.int8, device=device - ) - layer._nvfp4_sm70_dense_expert_ids = torch.arange( - experts, dtype=torch.int32, device=device - ) - layer._nvfp4_sm70_compact_offsets = torch.empty( - max_slots + 1, dtype=torch.int32, device=device - ) - layer._nvfp4_sm70_active_expert_ids = torch.empty( - max_slots, dtype=torch.int32, device=device - ) - - @staticmethod - def _persistent_buffers( - layer: RoutedExperts, num_tokens: int - ) -> dict[str, torch.Tensor]: - slots = num_tokens * _QWEN36_TOP_K - return { - "output": layer._nvfp4_sm70_output[:num_tokens], - "permuted_input": layer._nvfp4_sm70_permuted_input[:slots], - "gate_up": layer._nvfp4_sm70_gate_up[:slots], - "intermediate": layer._nvfp4_sm70_intermediate[:slots], - "sorted_output": layer._nvfp4_sm70_sorted_output[:slots], - "expert_offsets": layer._nvfp4_sm70_expert_offsets, - "expert_offsets64": layer._nvfp4_sm70_expert_offsets64, - "inv_permuted_idx": layer._nvfp4_sm70_inv_permuted_idx[:num_tokens], - "topk_ids": layer._nvfp4_sm70_topk_ids[:num_tokens], - "token_expert_indices": ( - layer._nvfp4_sm70_token_expert_indices[:num_tokens] - ), - "permuted_idx": layer._nvfp4_sm70_permuted_idx[:slots], - "sort_workspace": layer._nvfp4_sm70_sort_workspace, - "permuted_experts_id": layer._nvfp4_sm70_permuted_experts_id[:slots], - "sorted_row_idx": layer._nvfp4_sm70_sorted_row_idx[:slots], - "topk_ids_for_sort": layer._nvfp4_sm70_topk_ids_for_sort[:slots], - "dense_expert_ids": layer._nvfp4_sm70_dense_expert_ids, - "compact_offsets": layer._nvfp4_sm70_compact_offsets[: slots + 1], - "active_expert_ids": layer._nvfp4_sm70_active_expert_ids[:slots], - } - - @staticmethod - def _eager_buffers( - layer: RoutedExperts, num_tokens: int - ) -> dict[str, torch.Tensor]: - device = layer.w13_tm_weight.device - slots = num_tokens * _QWEN36_TOP_K - experts = int(layer.sm70_nvfp4_num_experts) - hidden = int(layer.sm70_nvfp4_hidden_size) - intermediate = int(layer.sm70_nvfp4_intermediate_size) - workspace_size = torch.ops._moe_C.moe_permute_sort_workspace_size( - slots, layer.global_num_experts - ) - return { - "output": torch.empty( - num_tokens, hidden, dtype=torch.float16, device=device - ), - "permuted_input": torch.empty( - slots, hidden, dtype=torch.float16, device=device - ), - "gate_up": torch.empty( - slots, 2 * intermediate, dtype=torch.float16, device=device - ), - "intermediate": torch.empty( - slots, intermediate, dtype=torch.float16, device=device - ), - "sorted_output": torch.empty( - slots, hidden, dtype=torch.float16, device=device - ), - "expert_offsets": torch.empty( - experts + 1, dtype=torch.int32, device=device - ), - "expert_offsets64": torch.empty( - experts + 1, dtype=torch.int64, device=device - ), - "inv_permuted_idx": torch.empty( - num_tokens, _QWEN36_TOP_K, dtype=torch.int32, device=device - ), - "topk_ids": torch.empty( - num_tokens, _QWEN36_TOP_K, dtype=torch.int32, device=device - ), - "token_expert_indices": torch.arange( - slots, dtype=torch.int32, device=device - ).view(num_tokens, _QWEN36_TOP_K), - "permuted_idx": torch.empty(slots, dtype=torch.int32, device=device), - "sort_workspace": torch.empty( - workspace_size, dtype=torch.int8, device=device - ), - "permuted_experts_id": torch.empty(slots, dtype=torch.int32, device=device), - "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), - "active_expert_ids": torch.empty(slots, dtype=torch.int32, device=device), - } - - def _get_buffers( - self, layer: RoutedExperts, num_tokens: int - ) -> dict[str, torch.Tensor]: - if 0 < num_tokens <= _GRAPH_SAFE_MAX_TOKENS: - return self._persistent_buffers(layer, num_tokens) - return self._eager_buffers(layer, num_tokens) - - @staticmethod - def _apply_swiglu( - layer: RoutedExperts, out: torch.Tensor, gate_up: torch.Tensor - ) -> None: - if layer.swiglu_limit is None: - torch.ops._C.silu_and_mul(out, gate_up) - else: - torch.ops._C.silu_and_mul_with_clamp( - out, gate_up, float(layer.swiglu_limit) - ) - - def apply( - self, - layer: RoutedExperts, - x: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - shared_experts: SharedExperts | None, - shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor: - del shared_experts, shared_experts_input - if not x.is_cuda or x.dtype != torch.float16 or x.ndim != 2: - 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: - raise ValueError( - "SM70 NVFP4 MoE activation hidden size mismatch: expected " - f"{_QWEN36_HIDDEN_SIZE}, got {x.shape[1]}." - ) - 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: - raise TypeError("SM70 NVFP4 MoE requires float32 top-k weights.") - - num_tokens = x.shape[0] - if num_tokens == 0: - return x.new_empty((0, _QWEN36_HIDDEN_SIZE)) - 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) - - if num_tokens <= _COMPACT_GROUPED_MAX_TOKENS: - _prepare_compact_slot_groups( - buffers["permuted_experts_id"], - buffers["compact_offsets"], - buffers["active_expert_ids"], - ) - stage_offsets = buffers["compact_offsets"] - stage_expert_ids = buffers["active_expert_ids"] - stage_experts = slots - else: - stage_offsets = buffers["expert_offsets"] - stage_expert_ids = buffers["dense_expert_ids"] - stage_experts = int(layer.sm70_nvfp4_num_experts) - - sm70_ops.nvfp4_moe_dense_stage_sm70_out( - buffers["gate_up"], - buffers["permuted_input"], - stage_offsets, - stage_expert_ids, - layer.w13_strided_ptrs_w, - layer.w13_strided_ptrs_s, - stage_experts, - layer.sm70_nvfp4_w13_k_dim, - layer.sm70_nvfp4_w13_n_dim, - layer.sm70_nvfp4_group_size, - ) - self._apply_swiglu(layer, buffers["intermediate"], buffers["gate_up"]) - sm70_ops.nvfp4_moe_dense_stage_sm70_out( - buffers["sorted_output"], - buffers["intermediate"], - stage_offsets, - stage_expert_ids, - layer.w2_strided_ptrs_w, - layer.w2_strided_ptrs_s, - stage_experts, - layer.sm70_nvfp4_w2_k_dim, - 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, - ) - return output - - def apply_monolithic( - self, - layer: RoutedExperts, - x: torch.Tensor, - router_logits: torch.Tensor, - input_ids: torch.Tensor | None = None, - ) -> torch.Tensor: - del layer, x, router_logits, input_ids - raise NotImplementedError("SM70 NVFP4 MoE is not a monolithic route.") - - def get_fused_moe_quant_config( # type: ignore[override] - self, layer: RoutedExperts - ) -> FusedMoEQuantConfig | None: - del layer - return None diff --git a/vllm/model_executor/layers/quantization/sm70_turbomind.py b/vllm/model_executor/layers/quantization/sm70_turbomind.py index a2482692d7..d948a94a7f 100644 --- a/vllm/model_executor/layers/quantization/sm70_turbomind.py +++ b/vllm/model_executor/layers/quantization/sm70_turbomind.py @@ -14,7 +14,6 @@ COMPRESSED_UINT4_GROUP_SIZES = (32, 128) MXFP4_GROUP_SIZE = 32 NVFP4_GROUP_SIZE = 16 -NVFP4_QPN4_DENSE_WORKSPACE_ELEMENTS = 5120 * 8704 STATE_ATTR = "_sm70_turbomind_linear" SM70QuantBackend = Literal["auto", "marlin", "turbomind"] @@ -27,15 +26,7 @@ class SM70TurboMindLinearState: k_ld: int q_ld: int output_size: int - op_kind: Literal["uint4", "mxfp4", "nvfp4", "nvfp4_qpn4"] - gated_silu: bool = False - dense_weight_ptr: int = 0 - global_scale: float = 0.0 - use_scale_code: bool = False - - -# States retain only data_ptr(), so this cache owns the bounded allocation. -_nvfp4_qpn4_dense_workspaces: dict[tuple[int, torch.dtype], torch.Tensor] = {} + op_kind: Literal["uint4", "mxfp4", "nvfp4"] def quant_backend() -> SM70QuantBackend: @@ -73,13 +64,6 @@ def should_use_mxfp4_moe_turbomind() -> bool: ) -def should_use_nvfp4_moe_turbomind() -> bool: - """Select the native NVFP4 MoE path only on exact SM70.""" - return is_exact_sm70_cuda_platform() and use_turbomind( - envs.VLLM_SM70_NVFP4_TURBOMIND - ) - - def should_prepare_turbomind( tensor: torch.Tensor, default_enabled: bool, @@ -150,27 +134,19 @@ def _store_state( layer: torch.nn.Module, weight: torch.Tensor, scales: torch.Tensor, - meta: torch.Tensor | None, + meta: torch.Tensor, group_size: int, output_size: int, - op_kind: Literal["uint4", "mxfp4", "nvfp4", "nvfp4_qpn4"], - gated_silu: bool = False, - dense_weight_ptr: int = 0, - global_scale: float = 0.0, - use_scale_code: bool = False, + op_kind: Literal["uint4", "mxfp4", "nvfp4"], ) -> None: state = SM70TurboMindLinearState( weight=weight, scales=scales, group_size=group_size, - k_ld=0 if meta is None else int(meta[0]), - q_ld=0 if meta is None else int(meta[1]), + k_ld=int(meta[0]), + q_ld=int(meta[1]), output_size=output_size, op_kind=op_kind, - gated_silu=gated_silu, - dense_weight_ptr=dense_weight_ptr, - global_scale=global_scale, - use_scale_code=use_scale_code, ) setattr(layer, STATE_ATTR, state) @@ -308,71 +284,6 @@ def prepare_nvfp4_linear( NVFP4_GROUP_SIZE, qweight.size(1), "nvfp4", - interleave_gated_silu, - ) - - -def get_nvfp4_qpn4_dense_workspace(weight: torch.Tensor) -> torch.Tensor | None: - device_index = weight.device.index - if device_index is None: - device_index = torch.accelerator.current_device_index() - cache_key = (device_index, torch.float16) - workspace = _nvfp4_qpn4_dense_workspaces.get(cache_key) - if workspace is not None: - return workspace - try: - workspace = torch.empty( - (NVFP4_QPN4_DENSE_WORKSPACE_ELEMENTS,), - dtype=torch.float16, - device=weight.device, - ) - except torch.OutOfMemoryError: - return None - _nvfp4_qpn4_dense_workspaces[cache_key] = workspace - return workspace - - -def prepare_nvfp4_qpn4_linear( - layer: torch.nn.Module, - workspace: torch.Tensor, - gated_silu: bool, -) -> None: - """Replace one accepted NVFP4 dense weight with memory-neutral QPN4.""" - from vllm import _sm70_ops as sm70_ops - - qweight = unpack_mxfp4_weight(layer.weight.data) - use_scale_code = gated_silu or envs.VLLM_SM70_NVFP4_QPN4_DOWN_SCALE_CODE - global_scale = 0.0 - if use_scale_code: - global_scale = float(layer.weight_global_scale.detach().float().item()) - raw_scale_codes = layer.weight_scale.data.t().contiguous() - packed_weight, packed_scales = sm70_ops.nvfp4_qpn4_prepare_scale_code_sm70( - qweight, raw_scale_codes - ) - else: - fp16_scales = ( - ( - layer.weight_scale.data.t().to(torch.float32) - * layer.weight_global_scale.to(torch.float32) - ) - .to(torch.float16) - .contiguous() - ) - packed_weight, packed_scales = sm70_ops.nvfp4_qpn4_prepare_sm70( - qweight, fp16_scales - ) - _store_state( - layer, - packed_weight, - packed_scales, - None, - NVFP4_GROUP_SIZE, - qweight.size(1), - "nvfp4_qpn4", - gated_silu=gated_silu, - dense_weight_ptr=workspace.data_ptr(), - global_scale=global_scale, - use_scale_code=use_scale_code, ) @@ -421,88 +332,8 @@ def apply_prepared_linear( state.k_ld, state.q_ld, ) - elif state.op_kind == "nvfp4_qpn4": - if reshaped_x.dtype != torch.float16: - raise RuntimeError( - f"SM70 NVFP4 QPN4 requires float16 activations, got {reshaped_x.dtype}." - ) - if reshaped_x.stride(-1) != 1: - reshaped_x = reshaped_x.contiguous() - sm70_ops.nvfp4_qpn4_dispatch_sm70_out( - out, - state.dense_weight_ptr, - reshaped_x, - state.weight, - state.scales, - state.global_scale, - state.use_scale_code, - False, - ) else: raise AssertionError(f"unknown SM70 TurboMind op kind: {state.op_kind}") - if state.gated_silu and state.op_kind == "nvfp4": - out_features = state.output_size // 2 - out = ( - out.reshape(reshaped_x.shape[0], out_features, 2) - .transpose(1, 2) - .reshape(reshaped_x.shape[0], state.output_size) - ) if bias is not None: out.add_(bias) return out.reshape(out_shape) - - -def apply_prepared_fused_silu_and_mul( - layer: torch.nn.Module, - x: torch.Tensor, -) -> torch.Tensor | None: - state = getattr(layer, STATE_ATTR, None) - if ( - state is None - or state.op_kind not in ("nvfp4", "nvfp4_qpn4") - or not state.gated_silu - ): - return None - if x.dtype != torch.float16: - raise RuntimeError( - "SM70 TurboMind NVFP4 gated-SiLU requires float16 activations, " - f"got {x.dtype}." - ) - - reshaped_x = x.reshape(-1, x.shape[-1]) - if reshaped_x.stride(-1) != 1: - reshaped_x = reshaped_x.contiguous() - out_features = state.output_size // 2 - out = torch.empty( - (reshaped_x.shape[0], out_features), - dtype=x.dtype, - device=x.device, - ) - if reshaped_x.shape[0] == 0: - return out.reshape(*x.shape[:-1], out_features) - - from vllm import _sm70_ops as sm70_ops - - if state.op_kind == "nvfp4_qpn4": - sm70_ops.nvfp4_qpn4_dispatch_sm70_out( - out, - state.dense_weight_ptr, - reshaped_x, - state.weight, - state.scales, - state.global_scale, - state.use_scale_code, - True, - ) - else: - sm70_ops.nvfp4_gemm_sm70_out( - out, - reshaped_x, - state.weight, - state.scales, - state.group_size, - state.k_ld, - state.q_ld, - True, - ) - return out.reshape(*x.shape[:-1], out_features) 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 68da508afc..f9152a52e7 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py @@ -28,7 +28,10 @@ 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/w8a8_utils.py b/vllm/model_executor/layers/quantization/utils/w8a8_utils.py index 7a08f380d6..40a350f4d8 100644 --- a/vllm/model_executor/layers/quantization/utils/w8a8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/w8a8_utils.py @@ -22,7 +22,8 @@ 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: @@ -32,7 +33,8 @@ 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: @@ -42,7 +44,8 @@ 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.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 89f7f2d457..d3b10092a6 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -355,7 +355,7 @@ def sparse_attn_indexer( padded_weights.reshape(num_padded_tokens, -1), seq_lens, decode_metadata.block_table, - attn_metadata_narrowed.compressed_max_seq_len, + attn_metadata_narrowed.max_seq_len, ) elif current_platform.is_xpu(): if padded_q_scale is not None: @@ -397,7 +397,7 @@ def sparse_attn_indexer( topk_indices, topk_workspace, topk_tokens, - logits.shape[1], + attn_metadata_narrowed.max_seq_len, ) else: ops.top_k_per_row_decode( diff --git a/vllm/model_executor/layers/vocab_parallel_embedding.py b/vllm/model_executor/layers/vocab_parallel_embedding.py index 9f69084c62..73b9e24167 100644 --- a/vllm/model_executor/layers/vocab_parallel_embedding.py +++ b/vllm/model_executor/layers/vocab_parallel_embedding.py @@ -35,14 +35,6 @@ DEFAULT_VOCAB_PADDING_SIZE = 64 logger = init_logger(__name__) -_SM70_DFLASH2_QPN8_LM_HEAD_SHAPE = (62080, 5120) -_SM70_DFLASH2_QPN8_MAX_ROWS = 8 -_SM70_DFLASH2_QPN8_CANDIDATES = 64 -_SM70_DFLASH2_QPN8_SPLIT_K = 8 -_SM70_DFLASH2_QPN8_ACCUMULATOR_CHAINS = 1 -_SM70_DFLASH2_RERANK_CTA_N = 128 -_SM70_DFLASH2_RERANK_SPLIT_K = 10 - def _sm70_env_bool(name: str, default: bool) -> bool: value = os.getenv(name) @@ -52,16 +44,8 @@ def _sm70_env_bool(name: str, default: bool) -> bool: def _sm70_lm_head_top1_default() -> bool: - return not _sm70_env_bool("VLLM_SM70_FLASH_V100_0DOT3_COMPILE_GRAPH", False) - - -def _sm70_dflash2_qpn8_rerank_enabled() -> bool: - return envs.VLLM_SM70_DFLASH2_QPN8_RERANK - - -def _sm70_dflash2_qpn8_rerank_requested() -> bool: - return ( - _sm70_dflash2_qpn8_rerank_enabled() or envs.VLLM_SM70_DFLASH2_QPN8_RERANK_SHADOW + return not _sm70_env_bool( + "VLLM_SM70_FLASH_V100_0DOT3_COMPILE_GRAPH", False ) @@ -81,7 +65,6 @@ def _is_sm70_lm_head_fastpath_eligible(layer: torch.nn.Module) -> bool: _sm70_env_bool("VLLM_SM70_ENABLE_LM_HEAD_FASTPATH", False) or _sm70_env_bool("VLLM_SM70_LM_HEAD_TOP1", _sm70_lm_head_top1_default()) or _sm70_env_bool("VLLM_SM70_LM_HEAD_TOP1_TC", False) - or _sm70_dflash2_qpn8_rerank_requested() ): _trace_sm70_lm_head_skip("disabled") return False @@ -111,229 +94,6 @@ def _is_sm70_lm_head_fastpath_eligible(layer: torch.nn.Module) -> bool: return True -def _is_sm70_dflash2_qpn8_rerank_eligible(layer: torch.nn.Module) -> bool: - if not _sm70_dflash2_qpn8_rerank_requested(): - return False - if tuple(layer.weight.shape) != _SM70_DFLASH2_QPN8_LM_HEAD_SHAPE: - logger.warning_once( - "SM70 DFlash2 QPN8 rerank requires LM-head shape %s; got %s. " - "Using the dense LM head.", - _SM70_DFLASH2_QPN8_LM_HEAD_SHAPE, - tuple(layer.weight.shape), - ) - return False - if getattr(layer, "tp_size", 1) != 4: - logger.warning_once( - "SM70 DFlash2 QPN8 rerank requires TP4; using the dense LM head." - ) - return False - if layer.shard_indices.num_org_vocab_padding != 0: - logger.warning_once( - "SM70 DFlash2 QPN8 rerank requires an unpadded local vocabulary; " - "using the dense LM head." - ) - return False - required_ops = ( - "fp8_qpn8_prepare_sm70", - "fp8_qpn8_gemm_sm70_out", - "sm70_f16_indexed_rerank_packed_out", - "sm70_f16_rerank_keys_out", - "sm70_f16_rerank_topk_out", - ) - missing = [name for name in required_ops if not hasattr(torch.ops._C, name)] - if missing: - logger.warning_once( - "SM70 DFlash2 QPN8 rerank operators are unavailable (%s); using " - "the dense LM head.", - ", ".join(missing), - ) - return False - return True - - -@torch.inference_mode() -def _prepare_sm70_dflash2_qpn8_rerank(layer: torch.nn.Module) -> bool: - if not _is_sm70_dflash2_qpn8_rerank_eligible(layer): - return False - - weight = layer.weight - rows, hidden = weight.shape - qweight = torch.empty_like(weight, dtype=torch.float8_e4m3fn) - channel_scales = torch.empty((rows, 1), dtype=torch.float32, device=weight.device) - # Quantize in bounded row chunks so startup never materializes another - # full-size FP32 LM head. The original FP16 parameter remains the oracle - # used by the exact candidate rerank. - chunk_rows = 4096 - for begin in range(0, rows, chunk_rows): - end = min(begin + chunk_rows, rows) - weight_f32 = weight[begin:end].float() - scales = weight_f32.abs().amax(dim=1, keepdim=True).div_(448.0) - scales.clamp_(min=torch.finfo(torch.float32).tiny) - channel_scales[begin:end].copy_(scales) - qweight[begin:end].copy_( - weight_f32.div_(scales).clamp_(-448.0, 448.0).to(torch.float8_e4m3fn) - ) - - codes, packed_scales = sm70_ops.fp8_qpn8_prepare_sm70(qweight, channel_scales) - del qweight, channel_scales, weight_f32, scales - torch.accelerator.empty_cache() - - device = weight.device - max_rows = _SM70_DFLASH2_QPN8_MAX_ROWS - candidates = _SM70_DFLASH2_QPN8_CANDIDATES - layer.register_buffer("_sm70_dflash2_qpn8_codes", codes, persistent=False) - layer.register_buffer("_sm70_dflash2_qpn8_scales", packed_scales, persistent=False) - layer.register_buffer( - "_sm70_dflash2_qpn8_logits", - torch.empty((max_rows, rows), dtype=torch.float16, device=device), - persistent=False, - ) - layer.register_buffer( - "_sm70_dflash2_qpn8_values", - torch.empty((max_rows, candidates), dtype=torch.float16, device=device), - persistent=False, - ) - layer.register_buffer( - "_sm70_dflash2_qpn8_ids", - torch.empty((max_rows, candidates), dtype=torch.int64, device=device), - persistent=False, - ) - layer.register_buffer( - "_sm70_dflash2_rerank_logits", - torch.empty((max_rows, candidates), dtype=torch.float16, device=device), - persistent=False, - ) - selected_rows = max_rows * candidates - layer.register_buffer( - "_sm70_dflash2_rerank_selected_raw", - torch.empty((selected_rows, hidden), dtype=torch.float16, device=device), - persistent=False, - ) - layer.register_buffer( - "_sm70_dflash2_rerank_selected_packed", - torch.empty((selected_rows, hidden), dtype=torch.float16, device=device), - persistent=False, - ) - layer.register_buffer( - "_sm70_dflash2_rerank_expanded", - torch.empty((max_rows, selected_rows), dtype=torch.float16, device=device), - persistent=False, - ) - layer.register_buffer( - "_sm70_dflash2_rerank_partials", - torch.empty((max_rows, selected_rows), dtype=torch.float32, device=device), - persistent=False, - ) - layer.register_buffer( - "_sm70_dflash2_rerank_barriers", - torch.zeros(64, dtype=torch.int32, device=device), - persistent=False, - ) - layer.register_buffer( - "_sm70_dflash2_rerank_dense_logits", - torch.empty((max_rows, rows), dtype=torch.float16, device=device), - persistent=False, - ) - layer.register_buffer( - "_sm70_dflash2_rerank_keys", - torch.empty((max_rows, candidates), dtype=torch.int64, device=device), - persistent=False, - ) - # Keep distinct top-16 and top-20 outputs. Slicing the columns of one - # [max_rows, 20] allocation for top-16 leaves a row stride of 20 and makes - # the result non-contiguous. The TP all-gather requires contiguous inputs, - # and inserting a runtime contiguous() copy would add work to both graphs. - for selector_k in (16, 20): - layer.register_buffer( - f"_sm70_dflash2_rerank_values_{selector_k}", - torch.empty((max_rows, selector_k), dtype=torch.float16, device=device), - persistent=False, - ) - layer.register_buffer( - f"_sm70_dflash2_rerank_positions_{selector_k}", - torch.empty((max_rows, selector_k), dtype=torch.int64, device=device), - persistent=False, - ) - layer.register_buffer( - f"_sm70_dflash2_rerank_ids_{selector_k}", - torch.empty((max_rows, selector_k), dtype=torch.int64, device=device), - persistent=False, - ) - layer.register_buffer( - f"_sm70_dflash2_rerank_key_values_{selector_k}", - torch.empty((max_rows, selector_k), dtype=torch.int64, device=device), - persistent=False, - ) - layer._sm70_dflash2_qpn8_rerank_prepared = True - logger.info_once( - "SM70 DFlash2 QPN8 top-64 plus exact TurboMind FP16 rerank layout prepared." - ) - return True - - -def _sm70_dflash2_rerank_output_buffers( - layer: torch.nn.Module, - num_rows: int, - selector_k: int, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Select graph-stable, contiguous rerank outputs for top-16 or top-20.""" - if selector_k == 16: - values = layer._sm70_dflash2_rerank_values_16[:num_rows] - positions = layer._sm70_dflash2_rerank_positions_16[:num_rows] - ids = layer._sm70_dflash2_rerank_ids_16[:num_rows] - elif selector_k == 20: - values = layer._sm70_dflash2_rerank_values_20[:num_rows] - positions = layer._sm70_dflash2_rerank_positions_20[:num_rows] - ids = layer._sm70_dflash2_rerank_ids_20[:num_rows] - else: - raise ValueError(f"Unsupported DFlash2 rerank top-k: {selector_k}") - return values, positions, ids - - -def _sm70_dflash2_dense_order_topk( - sparse_logits: torch.Tensor, - candidate_ids: torch.Tensor, - candidate_logits: torch.Tensor, - values: torch.Tensor, - ids: torch.Tensor, - selector_k: int, - vocab_start_index: int, -) -> None: - """Restore dense-vocabulary tie order after sparse candidate reranking.""" - sparse_logits.fill_(-float("inf")) - sparse_logits.scatter_(1, candidate_ids, candidate_logits) - torch.topk( - sparse_logits, - selector_k, - dim=-1, - sorted=True, - out=(values, ids), - ) - ids.add_(vocab_start_index) - - -def _sm70_dflash2_candidate_order_topk( - candidate_ids: torch.Tensor, - candidate_logits: torch.Tensor, - values: torch.Tensor, - ids: torch.Tensor, - selector_k: int, - vocab_start_index: int, -) -> None: - """Select exact values with original-vocabulary tie precedence.""" - if values.size(1) != selector_k: - raise ValueError( - f"DFlash2 rerank output width {values.size(1)} != top-k {selector_k}" - ) - sm70_ops.sm70_f16_rerank_topk_out( - values, - ids, - candidate_logits, - candidate_ids, - vocab_start_index, - ) - - def maybe_prepare_sm70_lm_head_top1(layer: torch.nn.Module) -> bool: if getattr(layer, "_sm70_f16_prepared", False): return True @@ -343,9 +103,7 @@ def maybe_prepare_sm70_lm_head_top1(layer: torch.nn.Module) -> bool: layer._sm70_f16_tm_weight = prepared[0] layer._sm70_f16_k_ld = int(prepared[1][0].item()) layer._sm70_f16_prepared = True - if _prepare_sm70_dflash2_qpn8_rerank(layer): - pass - elif envs.VLLM_SM70_ENABLE_LM_HEAD_FASTPATH: + if envs.VLLM_SM70_ENABLE_LM_HEAD_FASTPATH: logger.info_once("SM70 dense fp16 fast path enabled for LM head.") else: logger.info_once("SM70 LM head top1 layout prepared.") @@ -406,7 +164,8 @@ def _maybe_sm70_lm_head_top1( or hasattr(torch.ops._C, "sm70_f16_lm_head_top1_tc_out") ): logger.warning_once( - "SM70 LM head top1 requested, but no top1 op is available; falling back." + "SM70 LM head top1 requested, but no top1 op is available; " + "falling back." ) return None if x.dtype != torch.float16 or not x.is_cuda: @@ -434,7 +193,9 @@ def _maybe_sm70_lm_head_top1( values = torch.empty((x_2d.size(0),), dtype=torch.float32, device=x_2d.device) indices = torch.empty((x_2d.size(0),), dtype=torch.int64, device=x_2d.device) - if lm_head_top1_tc and hasattr(torch.ops._C, "sm70_f16_lm_head_top1_tc_out"): + if lm_head_top1_tc and hasattr( + torch.ops._C, "sm70_f16_lm_head_top1_tc_out" + ): tm_weight = getattr(layer, "_sm70_f16_tm_weight", None) k_ld = getattr(layer, "_sm70_f16_k_ld", None) if tm_weight is not None and k_ld is not None: @@ -455,7 +216,9 @@ def _maybe_sm70_lm_head_top1( if num_rows != 1: return None - if not lm_head_top1 or not hasattr(torch.ops._C, "sm70_f16_lm_head_top1_out"): + if not lm_head_top1 or not hasattr( + torch.ops._C, "sm70_f16_lm_head_top1_out" + ): return None sm70_ops.sm70_f16_lm_head_top1_out( @@ -471,146 +234,6 @@ def _maybe_sm70_lm_head_top1( return values.reshape(*x.shape[:-1]), indices.reshape(*x.shape[:-1]) -def _maybe_sm70_dflash2_qpn8_rerank( - layer: torch.nn.Module, - x: torch.Tensor, - selector_k: int, - bias: torch.Tensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor] | None: - """Return shard-local top-k after QPN8 support search and FP16 rerank.""" - if not _sm70_dflash2_qpn8_rerank_requested(): - return None - if not getattr(layer, "_sm70_dflash2_qpn8_rerank_prepared", False): - return None - if selector_k not in (16, 20) or bias is not None: - return None - if x.dtype != torch.float16 or not x.is_cuda: - return None - if torch.cuda.get_device_capability(x.device) != (7, 0): - return None - - x_2d = x.reshape(-1, x.shape[-1]) - num_rows = x_2d.size(0) - if not 1 <= num_rows <= _SM70_DFLASH2_QPN8_MAX_ROWS: - return None - if not x_2d.is_contiguous(): - x_2d = x_2d.contiguous() - - qpn8_logits = layer._sm70_dflash2_qpn8_logits[:num_rows] - qpn8_values = layer._sm70_dflash2_qpn8_values[:num_rows] - qpn8_ids = layer._sm70_dflash2_qpn8_ids[:num_rows] - sm70_ops.fp8_qpn8_gemm_sm70_out( - qpn8_logits, - x_2d, - layer._sm70_dflash2_qpn8_codes, - layer._sm70_dflash2_qpn8_scales, - _SM70_DFLASH2_QPN8_SPLIT_K, - _SM70_DFLASH2_QPN8_ACCUMULATOR_CHAINS, - False, - False, - ) - # The exact FP16 rerank is permutation-invariant over this approximate - # support, so skip the unnecessary 64-element result sort. This keeps the - # official PyTorch multiblock selector while avoiding its final bitonic - # kernel and leaves candidate quality unchanged. - torch.topk( - qpn8_logits, - _SM70_DFLASH2_QPN8_CANDIDATES, - dim=-1, - sorted=False, - out=(qpn8_values, qpn8_ids), - ) - - sm70_ops.sm70_f16_indexed_rerank_packed_out( - layer._sm70_dflash2_rerank_logits[:num_rows], - x_2d, - layer._sm70_f16_tm_weight, - qpn8_ids, - layer._sm70_dflash2_rerank_selected_packed, - layer._sm70_dflash2_rerank_expanded, - layer._sm70_dflash2_rerank_partials, - layer._sm70_dflash2_rerank_barriers, - _SM70_DFLASH2_RERANK_CTA_N, - _SM70_DFLASH2_RERANK_SPLIT_K, - ) - rerank_logits = layer._sm70_dflash2_rerank_logits[:num_rows] - values, _positions, ids = _sm70_dflash2_rerank_output_buffers( - layer, num_rows, selector_k - ) - if envs.VLLM_SM70_DFLASH2_QPN8_DENSE_ORDER: - _sm70_dflash2_dense_order_topk( - layer._sm70_dflash2_rerank_dense_logits[:num_rows], - qpn8_ids, - rerank_logits, - values, - ids, - selector_k, - layer.shard_indices.org_vocab_start_index, - ) - else: - _sm70_dflash2_candidate_order_topk( - qpn8_ids, - rerank_logits, - values, - ids, - selector_k, - layer.shard_indices.org_vocab_start_index, - ) - - if envs.VLLM_SM70_DFLASH2_QPN8_RERANK_SHADOW: - if torch.compiler.is_compiling() or torch.cuda.is_current_stream_capturing(): - raise RuntimeError( - "VLLM_SM70_DFLASH2_QPN8_RERANK_SHADOW is eager-only; disable " - "CUDA Graph/torch.compile for the real-hidden coverage audit." - ) - dense_logits = layer.quant_method.apply(layer, x_2d, bias=None) - dense_values, dense_ids = torch.topk( - dense_logits, - selector_k, - dim=-1, - sorted=True, - ) - support_matches = (dense_ids[:, :, None] == qpn8_ids[:, None, :]).any(dim=-1) - dense_global_ids = dense_ids + layer.shard_indices.org_vocab_start_index - rerank_matches = (dense_global_ids[:, :, None] == ids[:, None, :]).any(dim=-1) - call = int(getattr(layer, "_sm70_dflash2_qpn8_shadow_calls", 0)) + 1 - layer._sm70_dflash2_qpn8_shadow_calls = call - logger.info( - "SM70_DFLASH2_QPN8_SHADOW call=%d rows=%d top_k=%d " - "support_missing=%d exact_set_rows=%d top1_match_rows=%d " - "ordered_value_max_abs=%.6g", - call, - num_rows, - selector_k, - int((~support_matches).sum().item()), - int(rerank_matches.all(dim=-1).sum().item()), - int((dense_global_ids[:, 0] == ids[:, 0]).sum().item()), - float((dense_values.float() - values.float()).abs().max().item()), - ) - output_shape = (*x.shape[:-1], selector_k) - return dense_values.reshape(output_shape), dense_global_ids.reshape( - output_shape - ) - - 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, - ) - output_shape = (*x.shape[:-1], selector_k) - return values.reshape(output_shape), ids.reshape(output_shape) - - -def _maybe_sm70_dflash2_lm_head_top20( - layer: torch.nn.Module, - x: torch.Tensor, - selector_k: int, - bias: torch.Tensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor] | None: - """Return selector candidates from the opt-in QPN8 plus FP16 rerank.""" - return _maybe_sm70_dflash2_qpn8_rerank(layer, x, selector_k, bias) - - class UnquantizedEmbeddingMethod(QuantizeMethodBase): """Unquantized method for embeddings.""" @@ -1089,14 +712,6 @@ def maybe_get_sm70_lm_head_top1( ) -> tuple[torch.Tensor, torch.Tensor] | None: return _maybe_sm70_lm_head_top1(self, hidden_states, bias) - def maybe_get_sm70_dflash2_top20( - self, - hidden_states: torch.Tensor, - selector_k: int, - bias: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor] | None: - return _maybe_sm70_dflash2_lm_head_top20(self, hidden_states, selector_k, bias) - def extra_repr(self) -> str: s = f"num_embeddings={self.num_embeddings_per_partition}" s += f", embedding_dim={self.embedding_dim}" diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index de9786c777..d48ec981af 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -361,7 +361,9 @@ 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 " diff --git a/vllm/model_executor/models/deepseek_ocr.py b/vllm/model_executor/models/deepseek_ocr.py index 78ad964a3c..2575d3dcd4 100644 --- a/vllm/model_executor/models/deepseek_ocr.py +++ b/vllm/model_executor/models/deepseek_ocr.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Inference-only Deepseek-OCR model compatible with HuggingFace weights.""" +import math from collections.abc import Iterable, Mapping, Sequence from typing import Annotated, Literal @@ -52,7 +53,7 @@ BASE_SIZE, CROP_MODE, DeepseekOCRProcessor, - count_image_tokens_for, + count_tiles, ) from vllm.utils.tensor_schema import TensorSchema, TensorShape from vllm.v1.sample.logits_processor import ( @@ -189,15 +190,12 @@ def get_hf_config(self): return self.ctx.get_hf_config(DeepseekVLV2Config) def get_hf_processor(self, **kwargs: object): - v1_processor_config: dict[str, object] = dict(strategy="v1") - if "image_mode" not in kwargs: - # Only inject the constant trio when no named mode is requested; - # the processor treats ``image_mode`` as authoritative. - v1_processor_config.update( - image_size=IMAGE_SIZE, - base_size=BASE_SIZE, - crop_mode=CROP_MODE, - ) + v1_processor_config = dict( + image_size=IMAGE_SIZE, + base_size=BASE_SIZE, + crop_mode=CROP_MODE, + strategy="v1", + ) return self.ctx.get_hf_processor( DeepseekOCRProcessor, @@ -210,13 +208,35 @@ def get_supported_mm_limits(self) -> Mapping[str, int | None]: def get_num_image_tokens( self, *, image_width: int, image_height: int, cropping: bool = True ) -> int: - return count_image_tokens_for( - image_width=image_width, - image_height=image_height, - base_size=BASE_SIZE, - image_size=IMAGE_SIZE, - cropping=cropping and CROP_MODE, - ) + image_size = IMAGE_SIZE + base_size = BASE_SIZE + patch_size = 16 + downsample_ratio = 4 + + if CROP_MODE: + if image_width <= 640 and image_height <= 640: + crop_ratio = [1, 1] + else: + # find the closest aspect ratio to the target + crop_ratio = count_tiles( + image_width, image_height, image_size=IMAGE_SIZE + ) + + num_width_tiles, num_height_tiles = crop_ratio + else: + num_width_tiles = num_height_tiles = 1 + + h = w = math.ceil((base_size // patch_size) / downsample_ratio) + + h2 = w2 = math.ceil((image_size // patch_size) / downsample_ratio) + + global_views_tokens = h * (w + 1) + if num_width_tiles > 1 or num_height_tiles > 1: + local_views_tokens = (num_height_tiles * h2) * (num_width_tiles * w2 + 1) + else: + local_views_tokens = 0 + + return global_views_tokens + local_views_tokens + 1 def get_image_size_with_most_features(self) -> ImageSize: if IMAGE_SIZE == 1024 and BASE_SIZE == 1280: @@ -314,12 +334,10 @@ def get_replacement_deepseek_vl2(item_idx: int): else: size = images.get_image_size(item_idx) - # Count with the SAME processor instance the request's - # mm kwargs construct (mode/crops) — the producer/counter - # pair must never diverge. - num_image_tokens = hf_processor.count_image_tokens( + num_image_tokens = self.info.get_num_image_tokens( image_width=size.width, image_height=size.height, + cropping=CROP_MODE, ) return [image_token_id] * num_image_tokens diff --git a/vllm/model_executor/models/dflash_sm70.py b/vllm/model_executor/models/dflash_sm70.py deleted file mode 100644 index 039ea71a27..0000000000 --- a/vllm/model_executor/models/dflash_sm70.py +++ /dev/null @@ -1,300 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Range-preserving BF16 arithmetic for DFlash2 drafts on SM70. - -The checkpoint was trained with BF16 activations, but Volta executes its dense -GEMMs in FP16. These kernels retain FP16 Tensor Core transport while preserving -BF16 residual range and rounding at the model's trained rounding points. - -Adapted from haohervchb/sglang-V100@5526ef1c6a82. -""" - -from __future__ import annotations - -import torch -from torch import nn - -from vllm.triton_utils import tl, triton - -DFLASH_SM70_GATE_UP_INPUT_SCALE = 32.0 -DFLASH_SM70_WIDE_OUTPUT_SCALE = 256.0 - - -@triton.jit -def _round_fp32_to_bf16_rne(value): - """Round FP32 to BF16 RNE and keep the result represented as FP32.""" - bits = value.to(tl.uint32, bitcast=True) - rounded_bits = (bits + 0x7FFF + ((bits >> 16) & 1)) & 0xFFFF0000 - return rounded_bits.to(tl.float32, bitcast=True) - - -@triton.jit -def _dflash_silu_and_mul_sm70_kernel( - output_ptr, - output_scales_ptr, - gate_up_ptr, - n_cols: tl.constexpr, - gate_up_scale: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - row = tl.program_id(0).to(tl.int64) - offsets = tl.arange(0, BLOCK_SIZE) - mask = offsets < n_cols - input_offsets = row * 2 * n_cols + offsets - output_offsets = row * n_cols + offsets - - gate = tl.load(gate_up_ptr + input_offsets, mask=mask, other=0.0).to(tl.float32) - up = tl.load(gate_up_ptr + input_offsets + n_cols, mask=mask, other=0.0).to( - tl.float32 - ) - gate = _round_fp32_to_bf16_rne(gate * gate_up_scale) - up = _round_fp32_to_bf16_rne(up * gate_up_scale) - activated_gate = _round_fp32_to_bf16_rne(gate * tl.sigmoid(gate)) - output = _round_fp32_to_bf16_rne(activated_gate * up) - - # A per-row power-of-two divisor retains exact transport scaling. - max_abs = tl.max(tl.abs(output), axis=0) - output_scale = tl.maximum(max_abs / 32752.0, 1.0) - output_scale = tl.exp2(tl.ceil(tl.log2(output_scale))) - tl.store(output_scales_ptr + row, output_scale) - tl.store(output_ptr + output_offsets, output / output_scale, mask=mask) - - -@triton.jit -def _dflash_scale_output_sm70_kernel( - output_ptr, - row_scales_ptr, - n_cols: tl.constexpr, - residual_scale: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - row = tl.program_id(0).to(tl.int64) - offsets = tl.arange(0, BLOCK_SIZE) - mask = offsets < n_cols - row_offsets = row * n_cols + offsets - value = tl.load(output_ptr + row_offsets, mask=mask, other=0.0).to(tl.float32) - row_scale = tl.load(row_scales_ptr + row).to(tl.float32) - tl.store( - output_ptr + row_offsets, - value * (row_scale / residual_scale), - mask=mask, - ) - - -@triton.jit -def _dflash_rmsnorm_sm70_kernel( - output_ptr, - residual_output_ptr, - x_ptr, - residual_ptr, - weight_ptr, - n_cols: tl.constexpr, - eps: tl.constexpr, - residual_scale: tl.constexpr, - HAS_RESIDUAL: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - row = tl.program_id(0).to(tl.int64) - offsets = tl.arange(0, BLOCK_SIZE) - mask = offsets < n_cols - row_offsets = row * n_cols + offsets - - value = tl.load(x_ptr + row_offsets, mask=mask, other=0.0).to(tl.float32) - if HAS_RESIDUAL: - residual = tl.load(residual_ptr + row_offsets, mask=mask, other=0.0).to( - tl.float32 - ) - value = _round_fp32_to_bf16_rne(value * residual_scale + residual) - tl.store(residual_output_ptr + row_offsets, value, mask=mask) - - variance = tl.sum(value * value, axis=0) / n_cols - normalized = value * tl.rsqrt(variance + eps) - weight = tl.load(weight_ptr + offsets, mask=mask, other=0.0).to(tl.float32) - weight = _round_fp32_to_bf16_rne(weight) - output = _round_fp32_to_bf16_rne(normalized * weight) - tl.store(output_ptr + row_offsets, output, mask=mask) - - -def _num_warps(n_cols: int) -> int: - return max(min(triton.next_power_of_2(triton.cdiv(n_cols, 256)), 16), 4) - - -def _silu_and_mul_reference( - gate_up: torch.Tensor, gate_up_scale: float -) -> tuple[torch.Tensor, torch.Tensor]: - gate, up = gate_up.float().chunk(2, dim=-1) - gate = (gate * gate_up_scale).to(torch.bfloat16).float() - up = (up * gate_up_scale).to(torch.bfloat16).float() - activated = (gate * torch.sigmoid(gate)).to(torch.bfloat16).float() - output = (activated * up).to(torch.bfloat16).float() - max_abs = output.abs().amax(dim=-1) - row_scales = torch.maximum(max_abs / 32752.0, torch.ones_like(max_abs)) - row_scales = torch.pow(2.0, torch.ceil(torch.log2(row_scales))) - return (output / row_scales.unsqueeze(-1)).to(gate_up.dtype), row_scales - - -def dflash_silu_and_mul_sm70( - gate_up: torch.Tensor, - gate_up_scale: float = DFLASH_SM70_GATE_UP_INPUT_SCALE, -) -> tuple[torch.Tensor, torch.Tensor]: - if gate_up.ndim < 2 or gate_up.shape[-1] % 2: - raise ValueError( - "DFlash SM70 SwiGLU expects [..., 2 * intermediate_size], " - f"got {tuple(gate_up.shape)}." - ) - if gate_up_scale <= 0: - raise ValueError(f"gate_up_scale must be positive, got {gate_up_scale}.") - if not gate_up.is_cuda: - return _silu_and_mul_reference(gate_up, gate_up_scale) - if gate_up.dtype != torch.float16: - raise ValueError( - f"DFlash SM70 SwiGLU requires FP16 CUDA transport, got {gate_up.dtype}." - ) - if not gate_up.is_contiguous(): - gate_up = gate_up.contiguous() - - n_cols = gate_up.shape[-1] // 2 - n_rows = gate_up.numel() // (2 * n_cols) - output = torch.empty( - (*gate_up.shape[:-1], n_cols), - dtype=gate_up.dtype, - device=gate_up.device, - ) - output_scales = torch.empty(n_rows, dtype=torch.float32, device=gate_up.device) - block_size = triton.next_power_of_2(n_cols) - _dflash_silu_and_mul_sm70_kernel[(n_rows,)]( - output, - output_scales, - gate_up, - n_cols=n_cols, - gate_up_scale=gate_up_scale, - BLOCK_SIZE=block_size, - num_warps=_num_warps(n_cols), - ) - return output, output_scales - - -def dflash_scale_output_sm70( - output: torch.Tensor, - row_scales: torch.Tensor, - residual_scale: float = DFLASH_SM70_WIDE_OUTPUT_SCALE, -) -> torch.Tensor: - if output.ndim < 2: - raise ValueError(f"DFlash SM70 output must have ndim >= 2, got {output.ndim}.") - if residual_scale <= 0: - raise ValueError(f"residual_scale must be positive, got {residual_scale}.") - n_cols = output.shape[-1] - n_rows = output.numel() // n_cols - if row_scales.numel() != n_rows: - raise ValueError( - f"DFlash SM70 row-scale mismatch: rows={n_rows}, " - f"scales={row_scales.numel()}." - ) - if not output.is_cuda: - return (output.float() * (row_scales.reshape(-1, 1) / residual_scale)).to( - output.dtype - ) - if output.dtype != torch.float16 or not output.is_contiguous(): - raise ValueError("DFlash SM70 W2 scaling requires contiguous CUDA FP16 output.") - block_size = triton.next_power_of_2(n_cols) - _dflash_scale_output_sm70_kernel[(n_rows,)]( - output, - row_scales, - n_cols=n_cols, - residual_scale=residual_scale, - BLOCK_SIZE=block_size, - num_warps=_num_warps(n_cols), - ) - return output - - -class DFlashSM70RMSNorm(nn.Module): - """RMSNorm with FP32 residual transport and explicit BF16 RNE.""" - - def __init__( - self, - hidden_size: int, - eps: float, - dtype: torch.dtype, - residual_scale: float = DFLASH_SM70_WIDE_OUTPUT_SCALE, - ) -> None: - super().__init__() - self.hidden_size = hidden_size - self.variance_epsilon = eps - self.residual_scale = residual_scale - self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype)) - - def _reference( - self, x: torch.Tensor, residual: torch.Tensor | None - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - orig_dtype = x.dtype - value = x.float() - residual_output = None - if residual is not None: - value = ( - (value * self.residual_scale + residual.float()) - .to(torch.bfloat16) - .float() - ) - residual_output = value - variance = value.square().mean(dim=-1, keepdim=True) - normalized = value * torch.rsqrt(variance + self.variance_epsilon) - weight = self.weight.float().to(torch.bfloat16).float() - output = (normalized * weight).to(torch.bfloat16).to(orig_dtype) - if residual_output is None: - return output - return output, residual_output - - def forward( - self, x: torch.Tensor, residual: torch.Tensor | None = None - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - if not x.is_cuda: - return self._reference(x, residual) - if x.dtype != torch.float16: - raise ValueError( - f"DFlash SM70 RMSNorm requires FP16 CUDA transport, got {x.dtype}." - ) - if not x.is_contiguous(): - x = x.contiguous() - if residual is not None and not residual.is_contiguous(): - residual = residual.contiguous() - - n_rows = x.numel() // self.hidden_size - output = torch.empty_like(x) - has_residual = residual is not None - residual_output = ( - torch.empty_like(x, dtype=torch.float32) if has_residual else output - ) - residual_ptr = residual if residual is not None else x - block_size = triton.next_power_of_2(self.hidden_size) - _dflash_rmsnorm_sm70_kernel[(n_rows,)]( - output, - residual_output, - x, - residual_ptr, - self.weight, - n_cols=self.hidden_size, - eps=self.variance_epsilon, - residual_scale=self.residual_scale, - HAS_RESIDUAL=has_residual, - BLOCK_SIZE=block_size, - num_warps=_num_warps(self.hidden_size), - ) - if has_residual: - return output, residual_output - return output - - -class DFlashSM70MLP(nn.Module): - """Wrap the official Qwen MLP with overflow-safe BF16-equivalent math.""" - - def __init__(self, mlp: nn.Module) -> None: - super().__init__() - self.gate_up_proj = mlp.gate_up_proj - self.down_proj = mlp.down_proj - - def forward(self, x: torch.Tensor) -> torch.Tensor: - gate_up, _ = self.gate_up_proj(x / DFLASH_SM70_GATE_UP_INPUT_SCALE) - x, row_scales = dflash_silu_and_mul_sm70(gate_up) - x, _ = self.down_proj(x) - return dflash_scale_output_sm70(x, row_scales) diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py index 5fc0ba512a..b546040b74 100644 --- a/vllm/model_executor/models/gemma4_mm.py +++ b/vllm/model_executor/models/gemma4_mm.py @@ -102,19 +102,6 @@ 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 # --------------------------------------------------------------------------- @@ -961,7 +948,6 @@ 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 @@ -969,27 +955,9 @@ 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, **tower_init_kwargs - ) + self.vision_tower = AutoModel.from_config(config=config.vision_config) self.embed_vision = Gemma4MultimodalEmbedder( config.vision_config, config.text_config ) @@ -997,9 +965,7 @@ 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, **tower_init_kwargs - ) + self.audio_tower = AutoModel.from_config(config=config.audio_config) # 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 @@ -1228,16 +1194,9 @@ 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=attention_mask, + attention_mask=~pad_tensor, pixel_position_ids=pp_tensor, ) hidden_states = encoder_outputs.last_hidden_state @@ -1343,14 +1302,9 @@ 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=attention_mask, + attention_mask=~pad_chunk, 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..46b3269f8b 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -265,8 +265,6 @@ def _mark_tower_model( If `targets` is set, instead include descendants that are an instance of `targets`, even if they aren't direct children. """ - from vllm.model_executor.offloader import get_offloader - from .utils import StageMissingLayer, collect_children, no_init_weights if isinstance(modalities, str): @@ -292,14 +290,6 @@ def _mark_tower_model( yield self._tower_model_names = children_names - offloader = get_offloader() - for name in children_names: - if not name: - continue - tower = self.get_submodule(name) - wrapped = offloader.wrap_module(tower, parameter_prefix=name) - if wrapped is not tower: - self.set_submodule(name, wrapped) @contextmanager def _mark_composite_model( diff --git a/vllm/model_executor/models/qwen2_moe.py b/vllm/model_executor/models/qwen2_moe.py index ecfbf4f064..f27c4f3037 100644 --- a/vllm/model_executor/models/qwen2_moe.py +++ b/vllm/model_executor/models/qwen2_moe.py @@ -149,7 +149,9 @@ 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) diff --git a/vllm/model_executor/models/qwen3_5.py b/vllm/model_executor/models/qwen3_5.py index dbc891fee5..b91c77dbc0 100644 --- a/vllm/model_executor/models/qwen3_5.py +++ b/vllm/model_executor/models/qwen3_5.py @@ -48,7 +48,6 @@ _resolve_qwen_gdn_kv_cache_args, _sm70_compile_graph_slice_dim, _sm70_dump_gdn_projection_tensor, - _sm70_gdn_qpn8_ba_dispatch_eligible, ) from vllm.model_executor.layers.mamba.mamba_utils import ( MambaStateCopyFunc, @@ -75,7 +74,6 @@ Qwen3_5MoeConfig, Qwen3_5MoeTextConfig, ) -from vllm.triton_utils import tl, triton from vllm.utils.torch_utils import ( _encode_layer_name, ) @@ -118,98 +116,6 @@ logger = init_logger(__name__) - -@triton.jit -def _sm70_qwen35_gdn_split_kernel( - mixed_qkvz, - mixed_ba, - output, - stride_qkvz_row: tl.int64, - stride_ba_row: tl.int64, - num_rows: tl.constexpr, - qkv_size: tl.constexpr, - z_size: tl.constexpr, - ba_size: tl.constexpr, - BLOCK_Z: tl.constexpr, - BLOCK_BA: tl.constexpr, -): - row = tl.program_id(0) - - z_cols = tl.arange(0, BLOCK_Z) - z_mask = z_cols < z_size - z = tl.load( - mixed_qkvz + row * stride_qkvz_row + qkv_size + z_cols, - mask=z_mask, - ) - tl.store(output + row * z_size + z_cols, z, mask=z_mask) - - ba_cols = tl.arange(0, BLOCK_BA) - ba_mask = ba_cols < ba_size - b = tl.load(mixed_ba + row * stride_ba_row + ba_cols, mask=ba_mask) - a = tl.load( - mixed_ba + row * stride_ba_row + ba_size + ba_cols, - mask=ba_mask, - ) - z_numel = num_rows * z_size - ba_numel = num_rows * ba_size - tl.store(output + z_numel + row * ba_size + ba_cols, b, mask=ba_mask) - tl.store( - output + z_numel + ba_numel + row * ba_size + ba_cols, - a, - mask=ba_mask, - ) - - -def _sm70_materialize_qwen35_gdn_splits( - mixed_qkvz: torch.Tensor, - mixed_ba: torch.Tensor, - qkv_size: int, - z_size: int, - ba_size: int, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Materialize Qwen3.5 z/b/a tails with one bitwise copy kernel.""" - if mixed_qkvz.ndim != 2 or mixed_ba.ndim != 2: - raise ValueError("mixed_qkvz and mixed_ba must both be rank two") - if mixed_qkvz.shape[0] != mixed_ba.shape[0]: - raise ValueError("mixed_qkvz and mixed_ba row counts must match") - if mixed_qkvz.dtype != mixed_ba.dtype: - raise ValueError("mixed_qkvz and mixed_ba dtypes must match") - if mixed_qkvz.stride(1) != 1 or mixed_ba.stride(1) != 1: - raise ValueError("mixed_qkvz and mixed_ba must be contiguous by row") - if mixed_qkvz.shape[1] < qkv_size + z_size: - raise ValueError("mixed_qkvz is narrower than qkv_size + z_size") - if mixed_ba.shape[1] < 2 * ba_size: - raise ValueError("mixed_ba is narrower than 2 * ba_size") - - num_rows = mixed_qkvz.shape[0] - z_numel = num_rows * z_size - ba_numel = num_rows * ba_size - packed = torch.empty( - z_numel + 2 * ba_numel, - dtype=mixed_qkvz.dtype, - device=mixed_qkvz.device, - ) - _sm70_qwen35_gdn_split_kernel[(num_rows,)]( - mixed_qkvz, - mixed_ba, - packed, - mixed_qkvz.stride(0), - mixed_ba.stride(0), - num_rows=num_rows, - qkv_size=qkv_size, - z_size=z_size, - ba_size=ba_size, - BLOCK_Z=triton.next_power_of_2(z_size), - BLOCK_BA=triton.next_power_of_2(ba_size), - num_warps=8, - num_stages=1, - ) - z = packed[:z_numel].view(num_rows, z_size) - b = packed[z_numel : z_numel + ba_numel].view(num_rows, ba_size) - a = packed[z_numel + ba_numel :].view(num_rows, ba_size) - return z, b, a - - def _parse_sm70_moe_dense_allowlist() -> set[str] | None: raw = envs.VLLM_SM70_MOE_DENSE_ALLOWLIST if raw is None: @@ -240,63 +146,15 @@ def add_ignored_modules(value: object) -> None: raw_config = getattr(quant_config, "config", None) if isinstance(raw_config, dict): add_ignored_modules(raw_config.get("ignore")) - if any( + if not ignored_modules: + return False + return any( module_name == "linear_attn" or module_name.endswith(".linear_attn") or ("linear_attn.in_proj_a" in module_name) or ("linear_attn.in_proj_b" in module_name) for module_name in ignored_modules - ): - return True - - # ModelOpt mixed checkpoints describe quantized modules positively in - # ``quantized_layers`` instead of listing the unquantized b/a projections - # in an ignore list. A fused qkv/z/b/a Linear would then allocate the b/a - # rows in the qkv/z quant format and leave their scale slots uninitialized. - # Compare algorithms per GDN layer and split whenever b/a is absent or - # differs from the quantized qkv/z group. - quantized_layers = getattr(quant_config, "quantized_layers", None) - if not isinstance(quantized_layers, dict): - return False - - per_layer: dict[str, dict[str, str]] = {} - projection_names = { - "in_proj_qkv", - "in_proj_z", - "in_proj_b", - "in_proj_a", - } - marker = ".linear_attn." - for module_name, layer_info in quantized_layers.items(): - if marker not in module_name or not isinstance(layer_info, dict): - continue - layer_prefix, projection_name = module_name.rsplit(marker, 1) - if projection_name not in projection_names: - continue - quant_algo = str(layer_info.get("quant_algo", "")).upper() - if quant_algo: - per_layer.setdefault(layer_prefix, {})[projection_name] = quant_algo - - for projection_algos in per_layer.values(): - qkvz_algos = { - projection_algos[name] - for name in ("in_proj_qkv", "in_proj_z") - if name in projection_algos - } - if not qkvz_algos: - continue - ba_algos = { - projection_algos[name] - for name in ("in_proj_b", "in_proj_a") - if name in projection_algos - } - has_ba_pair = all( - name in projection_algos for name in ("in_proj_b", "in_proj_a") - ) - if not has_ba_pair or ba_algos != qkvz_algos: - return True - - return False + ) def _get_default_sm70_dense_force_suffixes(tp_size: int) -> set[str]: @@ -390,7 +248,7 @@ def create_ba_proj( def forward_cuda( self, hidden_states: torch.Tensor, - output: torch.Tensor | None, + output: torch.Tensor, ): num_tokens = hidden_states.size(0) layer_name = _encode_layer_name(self.prefix) @@ -402,52 +260,7 @@ def forward_cuda( "gdn_hidden_states", layer_name, hidden_states ) - if _sm70_gdn_qpn8_ba_dispatch_eligible( - self, - hidden_states, - layer_name, - ): - mixed_qkv = torch.empty( - (num_tokens, qkv_size), - dtype=hidden_states.dtype, - device=hidden_states.device, - ) - z = torch.empty( - (num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim), - dtype=hidden_states.dtype, - device=hidden_states.device, - ) - b = torch.empty( - (num_tokens, ba_size), - dtype=hidden_states.dtype, - device=hidden_states.device, - ) - a = torch.empty_like(b) - qkvz_staging = torch.empty( - (num_tokens, qkv_size + z_size), - dtype=hidden_states.dtype, - device=hidden_states.device, - ) - ba_staging = torch.empty( - (num_tokens, ba_size * 2), - dtype=hidden_states.dtype, - device=hidden_states.device, - ) - assert self.in_proj_ba is not None - torch.ops._C.fp8_qpn8_dispatch_ba_split_sm70_out( - mixed_qkv, - z, - b, - a, - qkvz_staging, - ba_staging, - int(self.in_proj_qkvz.sm70_fp8_prefill_exact_dense_workspace_ptr), - hidden_states, - self.in_proj_qkvz.weight, - self.in_proj_qkvz.weight_scale_inv, - self.in_proj_ba.weight, - ) - elif self.use_split_input_projections: + if self.use_split_input_projections: mixed_qkvz, _ = self.in_proj_qkvz(hidden_states) assert self.in_proj_ba is not None ba, _ = self.in_proj_ba(hidden_states) @@ -456,18 +269,9 @@ def forward_cuda( ) ba = _sm70_dump_gdn_projection_tensor("in_proj_ba", layer_name, ba) mixed_qkv = mixed_qkvz[..., :qkv_size] - if self.enable_sm70_dflash2_fused_gdn_split: - z, b, a = _sm70_materialize_qwen35_gdn_splits( - mixed_qkvz, - ba, - qkv_size, - z_size, - ba_size, - ) - else: - z = _sm70_compile_graph_slice_dim(mixed_qkvz, -1, qkv_size, z_size) - b = ba[..., :ba_size] - a = _sm70_compile_graph_slice_dim(ba, -1, ba_size, ba_size) + z = _sm70_compile_graph_slice_dim(mixed_qkvz, -1, qkv_size, z_size) + b = ba[..., :ba_size] + a = _sm70_compile_graph_slice_dim(ba, -1, ba_size, ba_size) else: mixed_qkvzba, _ = self.in_proj_qkvz(hidden_states) mixed_qkvzba = _sm70_dump_gdn_projection_tensor( @@ -477,9 +281,15 @@ def forward_cuda( ba_start = z_start + z_size a_start = ba_start + ba_size mixed_qkv = mixed_qkvzba[..., :qkv_size] - z = _sm70_compile_graph_slice_dim(mixed_qkvzba, -1, z_start, z_size) - b = _sm70_compile_graph_slice_dim(mixed_qkvzba, -1, ba_start, ba_size) - a = _sm70_compile_graph_slice_dim(mixed_qkvzba, -1, a_start, ba_size) + z = _sm70_compile_graph_slice_dim( + mixed_qkvzba, -1, z_start, z_size + ) + b = _sm70_compile_graph_slice_dim( + mixed_qkvzba, -1, ba_start, ba_size + ) + a = _sm70_compile_graph_slice_dim( + mixed_qkvzba, -1, a_start, ba_size + ) mixed_qkv = _sm70_dump_gdn_projection_tensor( "split_mixed_qkv", layer_name, mixed_qkv @@ -510,7 +320,7 @@ def forward_cuda( conv_state_cache=conv_state_cache, ssm_state_cache=ssm_state_cache, ) - return self._output_projection(core_attn_out, z, output, num_tokens) + self._output_projection(core_attn_out, z, output, num_tokens) class Qwen3_5DecoderLayer(Qwen3NextDecoderLayer): diff --git a/vllm/model_executor/models/qwen3_5_mtp.py b/vllm/model_executor/models/qwen3_5_mtp.py index 88f6ba5ca1..dea3dab7d8 100644 --- a/vllm/model_executor/models/qwen3_5_mtp.py +++ b/vllm/model_executor/models/qwen3_5_mtp.py @@ -462,7 +462,8 @@ 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_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index 554d0459a8..63ddb51248 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import io +import os from collections.abc import Iterable import torch @@ -10,12 +10,10 @@ from transformers import Qwen3Config from vllm import _custom_ops as ops +from vllm import envs from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config -from vllm.distributed import ( - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, -) +from vllm.distributed import get_tensor_model_parallel_world_size from vllm.logger import init_logger from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.layernorm import RMSNorm @@ -37,17 +35,12 @@ ) from vllm.multimodal.inputs import NestedTensors from vllm.transformers_utils.config import set_default_rope_theta -from vllm.transformers_utils.repo_utils import get_hf_file_bytes from vllm.v1.attention.backend import AttentionType -from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( - get_eagle3_aux_layers_from_config, -) from .qwen2 import Qwen2MLP as Qwen3MLP from .qwen3 import Qwen3ForCausalLM from .utils import ( AutoWeightsLoader, - WeightsMapper, get_draft_quant_config, maybe_prefix, process_eagle_weight, @@ -56,121 +49,43 @@ logger = init_logger(__name__) -_SLIDING_ATTENTION = "sliding_attention" +def _get_dflash_per_layer_sliding_window( + config: Qwen3Config, + layer_idx: int, +) -> int | None: + layer_types = getattr(config, "layer_types", None) + if layer_types is None: + return None + layer_type = layer_types[layer_idx] + if layer_type == "sliding_attention": + return getattr(config, "sliding_window", None) + if layer_type == "full_attention": + return None + raise ValueError(f"Invalid DFlash layer_type {layer_type}") -def _dflash_layer_causal(config: Qwen3Config, layer_idx: int) -> bool: - """Resolve explicit causality before falling back to legacy layer defaults.""" - is_causal = getattr(config, "is_causal", None) - if is_causal is not None: - return bool(is_causal) - override = (getattr(config, "dflash_config", None) or {}).get("causal") - if override is not None: - return bool(override) - layer_types = getattr(config, "layer_types", None) - return bool(layer_types) and layer_types[layer_idx] == _SLIDING_ATTENTION - - -def dflash_has_any_non_causal(config: Qwen3Config) -> bool: - """Whether the draft needs a non-causal-capable backend, resolved from config - (config mirror of the model's ``get_draft_attn_causal``, usable pre-build).""" - return not all( - _dflash_layer_causal(config, i) for i in range(config.num_hidden_layers) - ) - - -def dflash_target_rope_is_neox_style(target_model: nn.Module) -> bool | None: - """The target's RoPE layout, from its first attention layer. - - A DFlash head must rotate Q/K the way the target it was distilled against - does, and a mismatch is silent — acceptance collapses but nothing errors and - the output stays correct. Draft checkpoints do not carry this, so take it - from the target. None if the target uses no RoPE. - """ - language_model = ( - target_model.get_language_model() - if hasattr(target_model, "get_language_model") - else target_model - ) - for module in language_model.modules(): - style = getattr(module, "is_neox_style", None) - if isinstance(style, bool): - return style - return None - - -def _get_dflash_fc_input_size(vllm_config: VllmConfig) -> int: - spec_config = vllm_config.speculative_config - config = spec_config.draft_model_config.hf_config - aux_layers = get_eagle3_aux_layers_from_config(spec_config) - num_features_to_use = len(aux_layers) if aux_layers else config.num_hidden_layers - target_hidden_size = ( - getattr(config, "target_hidden_size", None) or config.hidden_size - ) - return target_hidden_size * num_features_to_use - - -def _resolve_layer_attention( - config: Qwen3Config, layer_idx: int -) -> tuple[int | None, bool]: - """Resolve ``(sliding_window, causal)`` for one DFlash draft layer. - - +----------------------+-------------------------+--------------------------------+ - | Config | ``layer_type`` | *``causal`` | - +======================+=========================+================================+ - | ``layer_types`` | SWA if ``use_swa`` | True if ``layer_types[i]=SWA`` | - | | else ``layer_types[i]`` | else False | - +----------------------+-------------------------+--------------------------------+ - | ``layer_types=None`` | SWA | False | - | + ``use_swa=True`` | | | - +----------------------+-------------------------+--------------------------------+ - | ``layer_types=None`` | Full | False | - | + ``use_swa=False`` | | | - +----------------------+-------------------------+--------------------------------+ - * If ``dflash_config.causal`` is set, its value overrides ``causal`` for all layers. - - This is to support a varied ecosystem of checkpoints, including: - - XiaomiMiMo/MiMo-V2.5-Pro-FP4-DFlash (sets "use_swa", assumes non-causal) - - z-lab/gemma-4-31B-it-DFlash (has mixed layer types, assumes causal only for SWA) - - z-lab/Qwen3.5-9B-DFlash ("standard" DFlash, all full attn, assumes non-causal) - """ - dflash_config = getattr(config, "dflash_config", None) or {} - layer_types = getattr(config, "layer_types", None) - use_swa = dflash_config.get("use_swa", False) +def _debug_sync_context_kv_enabled() -> bool: + return envs.VLLM_DFLASH_SYNC_CONTEXT_KV - any_sliding = False - if layer_types is not None: - num_sliding = sum(lt == _SLIDING_ATTENTION for lt in layer_types) - any_sliding = num_sliding > 0 - # Mixed sliding/full attention needs multiple KV groups (V2 runner only). - if ( - 0 < num_sliding < len(layer_types) - and not get_current_vllm_config().use_v2_model_runner - ): - raise NotImplementedError( - "DFlash drafters with mixed sliding/full attention require " - "the V2 model runner; relaunch with " - "VLLM_USE_V2_MODEL_RUNNER=1." - ) - # ``use_swa`` forces SWA on every layer, even an all-full ``layer_types``. - if layer_types is None or (use_swa and not any_sliding): - is_sliding = use_swa - else: - is_sliding = layer_types[layer_idx] == _SLIDING_ATTENTION +def _skip_context_kv_precompute_enabled() -> bool: + return envs.VLLM_DFLASH_SKIP_CONTEXT_KV_PRECOMPUTE - sliding_window = None - if is_sliding: - sliding_window = dflash_config.get( - "swa_window_size", getattr(config, "sliding_window", None) - ) - if sliding_window is None: - raise ValueError( - "DFlash sliding attention requires a window size configured in " - "dflash_config.swa_window_size or the top-level sliding_window." - ) - return sliding_window, _dflash_layer_causal(config, layer_idx) +def _debug_context_kv_enabled() -> bool: + return envs.VLLM_DFLASH_DEBUG_CONTEXT_KV + + +def _dump_draft_layer_hiddens_enabled() -> bool: + return envs.VLLM_DFLASH_DUMP_LAYER_HIDDENS + + +def _dump_draft_layer0_components_enabled() -> bool: + return envs.VLLM_DFLASH_DUMP_LAYER0_COMPONENTS + + +def _dump_draft_attn_components_enabled() -> bool: + return envs.VLLM_DFLASH_DUMP_ATTN_COMPONENTS class DFlashQwen3Attention(nn.Module): @@ -190,18 +105,16 @@ def __init__( head_dim: int | None = None, rms_norm_eps: float = 1e-06, attention_bias: bool = False, - add_swa_attention_sink_bias: bool = False, - sliding_window: int | None = None, - causal: bool = False, - is_neox_style: bool = True, cache_config: CacheConfig | None = None, quant_config: QuantizationConfig | None = None, prefix: str = "", attn_type: str = AttentionType.DECODER, + per_layer_sliding_window: int | None = None, ) -> None: super().__init__() self.layer_name = prefix self.hidden_size = hidden_size + self.per_layer_sliding_window = per_layer_sliding_window tp_size = get_tensor_model_parallel_world_size() self.total_num_heads = num_heads assert self.total_num_heads % tp_size == 0 @@ -237,17 +150,8 @@ def __init__( self.rotary_emb = get_rope( self.head_dim, max_position=max_position, - is_neox_style=is_neox_style, rope_parameters=rope_parameters, ) - - self.attention_sink_bias = ( - torch.nn.Parameter(torch.empty(self.num_heads), requires_grad=False) - if add_swa_attention_sink_bias - else None - ) - - self.sliding_window = sliding_window self.attn = Attention( self.num_heads, self.head_dim, @@ -255,17 +159,14 @@ def __init__( num_kv_heads=self.num_kv_heads, cache_config=cache_config, quant_config=quant_config, - per_layer_sliding_window=sliding_window, prefix=f"{prefix}.attn", attn_type=attn_type, - sinks=self.attention_sink_bias, + per_layer_sliding_window=per_layer_sliding_window, ) - # The in-tree FlashAttention-V100 backend uses this marker to select - # its DFlash-compatible non-causal paged-prefill fallback. self.attn.is_dflash_draft_attn = True - self.causal = causal self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) + self._attn_components_dumped = False def forward( self, @@ -276,7 +177,7 @@ def forward( with the context K/V from the target model's hidden states. This forward op computes attention for the query tokens only. See also: precompute_and_store_context_kv""" - qkv, _ = self.qkv_proj(hidden_states) + qkv = F.linear(hidden_states, self.qkv_proj.weight, self.qkv_proj.bias) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) # Per-head RMSNorm @@ -291,67 +192,64 @@ def forward( q, k = self.rotary_emb(positions, q, k) attn_output = self.attn(q, k, v) - output_input_scale = getattr(self, "output_input_scale", 1.0) - if output_input_scale != 1.0: - attn_output = attn_output / output_input_scale output, _ = self.o_proj(attn_output) + if ( + _dump_draft_attn_components_enabled() + and not self._attn_components_dumped + and positions.shape[-1] == 17 + and int(positions.max().item()) > 0 + and ".layers.32.self_attn" in self.layer_name + ): + dump_path = f"/tmp/dflash_draft_attn_components_pid{os.getpid()}.pt" + torch.save( + { + "positions": positions.detach().cpu(), + "q": q.detach().to(torch.float16).cpu(), + "k": k.detach().to(torch.float16).cpu(), + "v": v.detach().to(torch.float16).cpu(), + "attn_output": attn_output.detach().to(torch.float16).cpu(), + "o_proj_output": output.detach().to(torch.float16).cpu(), + }, + dump_path, + ) + self._attn_components_dumped = True + logger.warning("Saved DFlash draft attn components to %s", dump_path) return output class DFlashQwen3DecoderLayer(nn.Module): - attention_cls = DFlashQwen3Attention - def __init__( self, vllm_config: VllmConfig, *, config: Qwen3Config, - layer_idx: int, cache_config: CacheConfig | None = None, quant_config: QuantizationConfig | None = None, + layer_idx: int, prefix: str = "", ) -> None: super().__init__() self.hidden_size = config.hidden_size set_default_rope_theta(config, default_theta=1000000) attn_type = AttentionType.DECODER - - # DFlash drafts store the sink-bias flag inside dflash_config; fall back - # to the top-level attribute used by other (e.g. MiMo) configs. - dflash_config = getattr(config, "dflash_config", None) or {} - add_swa_attention_sink_bias = dflash_config.get( - "attention_sink_bias", - getattr(config, "add_swa_attention_sink_bias", False), + per_layer_sliding_window = _get_dflash_per_layer_sliding_window( + config, layer_idx ) - # Resolve this layer's attention mode (full vs sliding window, causal vs - # non-causal) from the draft config. - sliding_window, causal = _resolve_layer_attention(config, layer_idx) - - # RoPE layout, copied off the target at load time by the draft loader - # (see `dflash_target_rope_is_neox_style`). Checkpoints do not carry it: - # a head distilled from an interleaved-RoPE target must rotate the way - # that target does, or every drafted Q/K is wrong and acceptance - # collapses with no error raised. - is_neox_style = getattr(config, "is_neox_style", True) - - self.self_attn = self.attention_cls( + self.self_attn = DFlashQwen3Attention( hidden_size=self.hidden_size, num_heads=config.num_attention_heads, max_position=config.max_position_embeddings, num_kv_heads=config.num_key_value_heads, rms_norm_eps=config.rms_norm_eps, attention_bias=getattr(config, "attention_bias", False), - add_swa_attention_sink_bias=add_swa_attention_sink_bias, - sliding_window=sliding_window, - causal=causal, - is_neox_style=is_neox_style, head_dim=getattr(config, "head_dim", None), cache_config=cache_config, quant_config=quant_config, rope_parameters=config.rope_parameters, prefix=f"{prefix}.self_attn", attn_type=attn_type, + per_layer_sliding_window=per_layer_sliding_window, ) self.mlp = Qwen3MLP( hidden_size=self.hidden_size, @@ -364,6 +262,7 @@ def __init__( self.post_attention_layernorm = RMSNorm( config.hidden_size, eps=config.rms_norm_eps ) + self._layer0_components_dumped = False def forward( self, @@ -376,50 +275,49 @@ def forward( else: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) + ln_out = hidden_states hidden_states = self.self_attn( positions=positions, hidden_states=hidden_states, ) + attn_out = hidden_states hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + post_attn_ln_out = hidden_states hidden_states = self.mlp(hidden_states) + if ( + _dump_draft_layer0_components_enabled() + and not self._layer0_components_dumped + and positions.shape[-1] == 17 + and int(positions.max().item()) > 0 + and ".layers.32.self_attn" in self.self_attn.layer_name + ): + mlp_gate_up, _ = self.mlp.gate_up_proj(post_attn_ln_out) + mlp_act_out = self.mlp.act_fn(mlp_gate_up) + dump_path = f"/tmp/dflash_draft_layer0_components_pid{os.getpid()}.pt" + torch.save( + { + "positions": positions.detach().cpu(), + "input_layernorm_out": ln_out.detach().to(torch.float16).cpu(), + "self_attn_out": attn_out.detach().to(torch.float16).cpu(), + "post_attention_layernorm_out": post_attn_ln_out.detach() + .to(torch.float16) + .cpu(), + "mlp_gate_up": mlp_gate_up.detach().to(torch.float16).cpu(), + "mlp_act_out": mlp_act_out.detach().to(torch.float16).cpu(), + "mlp_out": hidden_states.detach().to(torch.float16).cpu(), + "residual": residual.detach().to(torch.float16).cpu(), + }, + dump_path, + ) + self._layer0_components_dumped = True + logger.warning("Saved DFlash draft layer0 components to %s", dump_path) return hidden_states, residual @support_torch_compile class DFlashQwen3Model(nn.Module): - decoder_layer_cls = DFlashQwen3DecoderLayer - - hf_to_vllm_mapper = WeightsMapper( - orig_to_new_substr={ - "midlayer.": "layers.0.", - # Muse-Glimmer-30B-assistant names the aux-hidden-state encoder - # `encoder.fc` / `encoder.output_norm_enc`; this head calls them - # `fc` / `hidden_norm`. Same tensors and shapes, different names. - "encoder.output_norm_enc.": "hidden_norm.", - "encoder.fc.": "fc.", - }, - ) - - def _make_context_projection( - self, - *, - vllm_config: VllmConfig, - input_size: int, - output_size: int, - prefix: str, - ) -> nn.Module: - return ReplicatedLinear( - input_size=input_size, - output_size=output_size, - bias=False, - params_dtype=vllm_config.model_config.dtype, - quant_config=self.quant_config, - prefix=prefix, - return_bias=False, - ) - def __init__( self, *, @@ -448,37 +346,37 @@ def __init__( prefix=maybe_prefix(prefix, "embed_tokens"), ) - # Masked query slots are fed to the draft as `mask_token_id`. Most DFlash - # checkpoints will have the mask embedding in the vocabulary embedding table - # at that slot id. Some checkpoints (XiaomiMiMo/MiMo-V2.5-Pro-FP4-DFlash) ship - # with a separate mask embedding tensor to use instead. When present, we load it - # and substitute it for embed_tokens[mask_token_id] when computing embeddings. - self.mask_token_id = drafter_config.get("mask_token_id") - self.mask_embedding = nn.Parameter( - torch.zeros(self.config.hidden_size, dtype=vllm_config.model_config.dtype), - requires_grad=False, - ) - self.has_separate_mask_embedding = False - self.layers = nn.ModuleList( [ - self.decoder_layer_cls( + DFlashQwen3DecoderLayer( current_vllm_config, config=self.config, - layer_idx=layer_idx, cache_config=current_vllm_config.cache_config, quant_config=self.quant_config, + layer_idx=layer_idx, prefix=maybe_prefix(prefix, f"layers.{layer_idx + start_layer_id}"), ) for layer_idx in range(self.config.num_hidden_layers) ] ) if self.use_aux_hidden_state: - self.fc = self._make_context_projection( - vllm_config=vllm_config, - input_size=_get_dflash_fc_input_size(vllm_config), + num_features_to_use = self.config.num_hidden_layers + if "target_layer_ids" in drafter_config: + num_features_to_use = len(drafter_config["target_layer_ids"]) + elif "layer_ids" in drafter_config: + num_features_to_use = len(drafter_config["layer_ids"]) + if hasattr(self.config, "target_hidden_size"): + fc_input_size = self.config.target_hidden_size * num_features_to_use + else: + fc_input_size = self.config.hidden_size * num_features_to_use + self.fc = ReplicatedLinear( + input_size=fc_input_size, output_size=self.config.hidden_size, + bias=False, + params_dtype=vllm_config.model_config.dtype, + quant_config=self.quant_config, prefix=maybe_prefix(prefix, "fc"), + return_bias=False, ) self.hidden_norm = RMSNorm( self.config.hidden_size, @@ -488,43 +386,11 @@ def __init__( self.config.hidden_size, eps=self.config.rms_norm_eps, ) + self._debug_context_kv_dumped = False + self._draft_layer_hiddens_dumped = False def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - embeds = self.embed_tokens(input_ids) - if self.has_separate_mask_embedding and self.mask_token_id is not None: - # Replace masked slots with the dedicated mask embedding. - is_mask = (input_ids == self.mask_token_id).unsqueeze(-1) - embeds = torch.where(is_mask, self.mask_embedding.to(embeds.dtype), embeds) - return embeds - - def _build_context_kv_buffers( - self, - layers_attn: list[nn.Module], - has_bias: bool, - ) -> None: - self._hidden_norm_weight = self.hidden_norm.weight.data - - # KV projection weights: [num_layers * 2 * kv_size, hidden_size] - kv_weights = [a.qkv_proj.weight[a.q_size :] for a in layers_attn] - self._fused_kv_weight = torch.cat(kv_weights, dim=0) - if has_bias: - kv_biases = [a.qkv_proj.bias[a.q_size :] for a in layers_attn] - self._fused_kv_bias: torch.Tensor | None = torch.cat(kv_biases, dim=0) - else: - self._fused_kv_bias = None - - # K-norm weights stacked into one contiguous [num_layers, head_dim] - # tensor so the per-layer K-norm runs as a single grouped kernel. - self._k_norm_weights = torch.stack( - [a.k_norm.weight.data for a in layers_attn], dim=0 - ).contiguous() - # The grouped call below relies on the stable-ABI RMSNorm extension - # selecting one weight row per outer (layer) index. Older binaries - # silently accepted the 2-D tensor but applied row zero to every - # layer, which severely lowers DFlash acceptance without affecting - # target-only output quality. Verify the loaded binary once before - # trusting its grouped result. - self._batched_k_norm_runtime_verified: bool | None = None + return self.embed_tokens(input_ids) def _build_fused_kv_buffers(self) -> None: """Build fused weight buffers for precompute_and_store_context_kv. @@ -538,7 +404,19 @@ def _build_fused_kv_buffers(self) -> None: attn0 = layers_attn[0] has_bias = attn0.qkv_proj.bias is not None - self._build_context_kv_buffers(layers_attn, has_bias) + self._hidden_norm_weight = self.hidden_norm.weight.data + + # KV projection weights: [num_layers * 2 * kv_size, hidden_size] + kv_weights = [a.qkv_proj.weight[a.q_size :] for a in layers_attn] + self._fused_kv_weight = torch.cat(kv_weights, dim=0) + if has_bias: + kv_biases = [a.qkv_proj.bias[a.q_size :] for a in layers_attn] + self._fused_kv_bias: torch.Tensor | None = torch.cat(kv_biases, dim=0) + else: + self._fused_kv_bias = None + + # K-norm weights: list of [head_dim] tensors, one per layer. + self._k_norm_weights = [a.k_norm.weight.data for a in layers_attn] # RoPE parameters self._rope_head_size = attn0.rotary_emb.head_size @@ -569,112 +447,11 @@ def _build_fused_kv_buffers(self) -> None: # References to inner Attention layers for direct cache writes self._attn_layers = [layer.self_attn.attn for layer in self.layers] - def _project_context_kv( - self, - context_states: torch.Tensor, - num_ctx: int, - num_layers: int, - num_kv_heads: int, - head_dim: int, - ) -> tuple[torch.Tensor, torch.Tensor]: - # --- Fused KV projection (one GEMM for all layers) --- - normed_context_states = self._normalize_context_states(context_states) - all_kv_flat = F.linear( - normed_context_states, self._fused_kv_weight, self._fused_kv_bias - ) - # Single contiguous copy that separates K/V and transposes to - # layer-major layout. Result: [2, L, num_ctx, nkv, hd] contiguous. - # Indexing dim-0 gives contiguous [L, num_ctx, nkv, hd] for K and V. - all_kv = ( - all_kv_flat.view(num_ctx, num_layers, 2, num_kv_heads, head_dim) - .permute(2, 1, 0, 3, 4) - .contiguous() - ) - all_k = all_kv[0] # [L, num_ctx, nkv, hd], contiguous - all_v = all_kv[1] # [L, num_ctx, nkv, hd], contiguous - return all_k, all_v - - def _normalize_context_states(self, context_states: torch.Tensor) -> torch.Tensor: - normed_context_states = torch.empty_like(context_states) - ops.rms_norm( - normed_context_states, - context_states, - self._hidden_norm_weight, - self._rms_norm_eps, - ) - return normed_context_states - - def _normalize_context_k_per_layer(self, all_k: torch.Tensor) -> torch.Tensor: - """Reference path that works with pre-batched-weight RMSNorm binaries.""" - all_k_normed = torch.empty_like(all_k) - for layer_idx in range(all_k.shape[0]): - ops.rms_norm( - all_k_normed[layer_idx], - all_k[layer_idx], - self._k_norm_weights[layer_idx], - self._rms_norm_eps, - ) - return all_k_normed - - def _verify_batched_context_k_norm_runtime( - self, dtype: torch.dtype, device: torch.device - ) -> bool: - """Check the loaded stable extension with nonzero per-layer inputs.""" - num_layers, hidden_size = self._k_norm_weights.shape - probe = torch.ones( - (num_layers, 1, 1, hidden_size), - dtype=dtype, - device=device, - ) - grouped = torch.empty_like(probe) - ops.rms_norm( - grouped, - probe, - self._k_norm_weights, - self._rms_norm_eps, - ) - per_layer = self._normalize_context_k_per_layer(probe) - return torch.equal(grouped, per_layer) - - def _normalize_context_k(self, all_k: torch.Tensor) -> torch.Tensor: - # --- Grouped RMSNorm K across all layers ([L, num_ctx, nkv, hd]) --- - # The weight is selected per layer by the outermost (layer) index. - runtime_verified = getattr(self, "_batched_k_norm_runtime_verified", None) - if runtime_verified is None: - # This one-time comparison is normally consumed by model profiling - # or warmup. Use synthetic nonzero inputs so an all-zero dummy run - # cannot falsely approve a stale binary. The equality check - # deliberately synchronizes once: silently using the wrong K - # weights is a much larger failure than a cold startup fence. - runtime_verified = self._verify_batched_context_k_norm_runtime( - all_k.dtype, all_k.device - ) - self._batched_k_norm_runtime_verified = runtime_verified - if not runtime_verified: - logger.warning_once( - "The loaded stable-ABI RMSNorm extension does not preserve " - "DFlash per-layer batched weights. Falling back to the " - "bitwise per-layer path; rebuild the vLLM stable extension " - "to restore the grouped K-normalization fast path." - ) - - if not runtime_verified: - return self._normalize_context_k_per_layer(all_k) - - grouped = torch.empty_like(all_k) - ops.rms_norm( - grouped, - all_k, - self._k_norm_weights, - self._rms_norm_eps, - ) - return grouped - def precompute_and_store_context_kv( self, context_states: torch.Tensor, context_positions: torch.Tensor, - context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None, + context_slot_mapping: torch.Tensor | dict[str, torch.Tensor] | None = None, ) -> None: """Precompute K/V for context states write them into each layer's KV cache. @@ -688,6 +465,10 @@ def precompute_and_store_context_kv( When context_slot_mapping is None (e.g. during dummy_run) only the computation runs, and no K/V is written to cache. """ + if _skip_context_kv_precompute_enabled(): + logger.warning("Skipping DFlash context KV precompute for debugging.") + return + if not hasattr(self, "_num_attn_layers"): logger.warning_once( "DFlash buffer initialization was skipped. If dummy weights are not " @@ -701,8 +482,35 @@ def precompute_and_store_context_kv( hd = self._head_dim nkv = self._num_kv_heads - all_k, all_v = self._project_context_kv(context_states, num_ctx, L, nkv, hd) - all_k_normed = self._normalize_context_k(all_k) + # --- Fused KV projection (one GEMM for all layers) --- + normed_context_states = torch.empty_like(context_states) + ops.rms_norm( + normed_context_states, + context_states, + self._hidden_norm_weight, + self._rms_norm_eps, + ) + all_kv_flat = F.linear( + normed_context_states, self._fused_kv_weight, self._fused_kv_bias + ) + # Single contiguous copy that separates K/V and transposes to + # layer-major layout. Result: [2, L, num_ctx, nkv, hd] contiguous. + # Indexing dim-0 gives contiguous [L, num_ctx, nkv, hd] for K and V. + all_kv = ( + all_kv_flat.view(num_ctx, L, 2, nkv, hd).permute(2, 1, 0, 3, 4).contiguous() + ) + all_k = all_kv[0] # [L, num_ctx, nkv, hd], contiguous + all_v = all_kv[1] # [L, num_ctx, nkv, hd], contiguous + + # --- Per-layer RMSNorm K (3D: [num_ctx, nkv, hd] per layer) --- + all_k_normed = torch.empty_like(all_k) + for i in range(L): + ops.rms_norm( + all_k_normed[i], + all_k[i], + self._k_norm_weights[i], + self._rms_norm_eps, + ) # --- Fused RoPE across all layers --- # View as [L * num_ctx, kv] so RoPE sees one big batch (no copy). @@ -726,22 +534,145 @@ def precompute_and_store_context_kv( # --- Per-layer cache insert --- all_k_final = all_k_flat.view(L, num_ctx, nkv, hd) - per_layer = isinstance(context_slot_mapping, (list, tuple)) + debug_direct_refs: dict[ + int, list[tuple[str, int, torch.Tensor, torch.Tensor]] + ] = {} + if _debug_context_kv_enabled() and not self._debug_context_kv_dumped: + sample_tokens = min(8, num_ctx) + debug_ranges: list[tuple[str, int, int]] = [("head", 0, sample_tokens)] + if num_ctx > sample_tokens: + tail_start = max(0, num_ctx - sample_tokens) + if tail_start > 0: + debug_ranges.append(("tail", tail_start, num_ctx)) + + debug_cos_sin_cache = self._rope_cos_sin_cache + for layer_idx in sorted({0, L // 2, L - 1}): + attn_layer = self.layers[layer_idx].self_attn + layer_refs: list[tuple[str, int, torch.Tensor, torch.Tensor]] = [] + for sample_name, sample_start, sample_end in debug_ranges: + sample_len = sample_end - sample_start + direct_kv = F.linear( + normed_context_states[sample_start:sample_end], + attn_layer.qkv_proj.weight[attn_layer.q_size :], + ( + None + if attn_layer.qkv_proj.bias is None + else attn_layer.qkv_proj.bias[attn_layer.q_size :] + ), + ) + direct_k, direct_v = direct_kv.split([kv, kv], dim=-1) + direct_k = direct_k.view(sample_len, nkv, hd).contiguous() + direct_v = direct_v.view(sample_len, nkv, hd).contiguous() + direct_k_norm = torch.empty_like(direct_k) + ops.rms_norm( + direct_k_norm, + direct_k, + attn_layer.k_norm.weight.data, + self._rms_norm_eps, + ) + direct_k_flat = direct_k_norm.view(sample_len, kv) + if debug_cos_sin_cache.dtype != direct_k_flat.dtype: + debug_cos_sin_cache = debug_cos_sin_cache.to( + dtype=direct_k_flat.dtype + ) + ops.rotary_embedding( + context_positions[sample_start:sample_end], + direct_k_flat, + None, + self._rope_head_size, + debug_cos_sin_cache, + self._rope_is_neox, + ) + direct_k = direct_k_flat.view(sample_len, nkv, hd) + layer_refs.append((sample_name, sample_start, direct_k, direct_v)) + + fused_k_diff = ( + all_k_final[layer_idx, sample_start:sample_end] - direct_k + ).abs() + fused_v_diff = ( + all_v[layer_idx, sample_start:sample_end] - direct_v + ).abs() + logger.warning( + "DFlash context KV fused-vs-direct layer=%d sample=%s " + "start=%d tokens=%d k_max=%.6f k_mean=%.6f " + "v_max=%.6f v_mean=%.6f", + layer_idx, + sample_name, + sample_start, + sample_len, + float(fused_k_diff.max().item()), + float(fused_k_diff.mean().item()), + float(fused_v_diff.max().item()), + float(fused_v_diff.mean().item()), + ) + debug_direct_refs[layer_idx] = layer_refs + for i in range(L): - slot_mapping = ( - context_slot_mapping[i] if per_layer else context_slot_mapping - ) - if slot_mapping is None: - continue # dummy run: skip cache ops attn = self._attn_layers[i] + layer_context_slot_mapping = context_slot_mapping + if isinstance(context_slot_mapping, dict): + layer_context_slot_mapping = context_slot_mapping[attn.layer_name] + assert layer_context_slot_mapping is not None kv_cache = attn.kv_cache + if isinstance(kv_cache, list): + kv_cache = kv_cache[0] + if _debug_sync_context_kv_enabled(): + logger.warning( + "DFlash context KV update layer=%s kv_cache_shape=%s " + "slot_min=%s slot_max=%s num_ctx=%d", + attn.layer_name, + tuple(kv_cache.shape), + int(layer_context_slot_mapping.min().item()), + int(layer_context_slot_mapping.max().item()), + num_ctx, + ) attn.impl.do_kv_cache_update( attn, all_k_final[i], all_v[i], kv_cache, - slot_mapping, + layer_context_slot_mapping, ) + if ( + i in debug_direct_refs + and _debug_context_kv_enabled() + and not self._debug_context_kv_dumped + ): + block_size = kv_cache.shape[2] + for ( + sample_name, + sample_start, + direct_k, + direct_v, + ) in debug_direct_refs[i]: + sample_len = direct_k.shape[0] + sample_end = sample_start + sample_len + slots = layer_context_slot_mapping[sample_start:sample_end].to( + torch.long + ) + block_idx = torch.div(slots, block_size, rounding_mode="floor") + block_off = torch.remainder(slots, block_size) + cache_k = kv_cache[block_idx, 0, block_off] + cache_v = kv_cache[block_idx, 1, block_off] + cache_k_diff = (cache_k - direct_k).abs() + cache_v_diff = (cache_v - direct_v).abs() + logger.warning( + "DFlash context KV cache-readback layer=%d sample=%s " + "start=%d tokens=%d k_max=%.6f k_mean=%.6f " + "v_max=%.6f v_mean=%.6f", + i, + sample_name, + sample_start, + sample_len, + float(cache_k_diff.max().item()), + float(cache_k_diff.mean().item()), + float(cache_v_diff.max().item()), + float(cache_v_diff.mean().item()), + ) + if i == max(debug_direct_refs): + self._debug_context_kv_dumped = True + if _debug_sync_context_kv_enabled(): + torch.accelerator.synchronize(all_k_final.device) def forward( self, @@ -755,34 +686,44 @@ def forward( hidden_states = input_embeds residual = None + layer_hidden_states: list[torch.Tensor] | None = None + should_dump_layer_hiddens = ( + _dump_draft_layer_hiddens_enabled() + and not self._draft_layer_hiddens_dumped + and positions.shape[-1] == 17 + and int(positions.max().item()) > 0 + ) + if should_dump_layer_hiddens: + layer_hidden_states = [] for layer in self.layers: hidden_states, residual = layer( positions=positions, hidden_states=hidden_states, residual=residual, ) + if layer_hidden_states is not None: + layer_hidden_states.append( + hidden_states.detach().to(torch.float16).cpu() + ) hidden_states, _ = self.norm(hidden_states, residual) + if layer_hidden_states is not None: + dump_path = f"/tmp/dflash_draft_layers_pid{os.getpid()}.pt" + torch.save( + { + "positions": positions.detach().cpu(), + "input_embeds": input_embeds.detach().to(torch.float16).cpu(), + "layer_hidden_states": layer_hidden_states, + "final_hidden_states": hidden_states.detach() + .to(torch.float16) + .cpu(), + }, + dump_path, + ) + self._draft_layer_hiddens_dumped = True + logger.warning("Saved DFlash draft layer hiddens to %s", dump_path) return hidden_states - def _preprocess( - self, weights: Iterable[tuple[str, torch.Tensor]] - ) -> Iterable[tuple[str, torch.Tensor]]: - tp_size = get_tensor_model_parallel_world_size() - tp_rank = get_tensor_model_parallel_rank() - for name, loaded_weight in weights: - if "attention_sink_bias" in name: - # Sink bias is per-head; shard it across TP ranks like the - # attention heads themselves. - heads_per_rank = loaded_weight.shape[0] // tp_size - loaded_weight = loaded_weight.narrow( - 0, tp_rank * heads_per_rank, heads_per_rank - ) - yield name, loaded_weight - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - # onecat's current AutoWeightsLoader predates stacked mapper metadata, - # so preserve the upstream naming contract while passing shard ids - # directly to the existing linear weight loaders. stacked_params_mapping = [ (".qkv_proj", ".q_proj", "q"), (".qkv_proj", ".k_proj", "k"), @@ -792,11 +733,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ] params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() - for name, loaded_weight in self._preprocess(weights): - mapped_name = self.hf_to_vllm_mapper._map_name(name) - if mapped_name is None: - continue - name = mapped_name + for name, loaded_weight in weights: + if "midlayer." in name: + name = name.replace("midlayer.", "layers.0.") if self.quant_config is not None and ( scale_name := self.quant_config.get_cache_scale(name) ): @@ -817,7 +756,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: continue name = name.replace(weight_name, param_name) param = params_dict[name] - param.weight_loader(param, loaded_weight, shard_id) + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) break else: param = params_dict[name] @@ -828,18 +768,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: class DFlashQwen3ForCausalLM(Qwen3ForCausalLM): - model_cls = DFlashQwen3Model - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): nn.Module.__init__(self) - self.draft_model_config = vllm_config.speculative_config.draft_model_config - self.config = self.draft_model_config.hf_config + self.config = vllm_config.speculative_config.draft_model_config.hf_config if getattr(self.config, "draft_vocab_size", None) is None: self.config.draft_vocab_size = getattr(self.config, "vocab_size", None) target_layer_num = vllm_config.model_config.get_num_layers( vllm_config.parallel_config ) - self.model = self.model_cls( + self.model = DFlashQwen3Model( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"), start_layer_id=target_layer_num, @@ -879,14 +816,6 @@ def forward( ) -> torch.Tensor: return self.model(input_ids, positions, inputs_embeds) - def get_draft_kv_cache_layer_names(self) -> list[str]: - return [layer.self_attn.attn.layer_name for layer in self.model.layers] - - def get_draft_attn_causal(self) -> list[bool]: - """Per-layer attention causality, aligned with - get_draft_kv_cache_layer_names.""" - return [layer.self_attn.causal for layer in self.model.layers] - def compute_logits( self, hidden_states: torch.Tensor, @@ -909,12 +838,6 @@ def get_topk_tokens_and_logprobs( hidden_states: torch.Tensor, top_k: int, ) -> tuple[torch.Tensor, torch.Tensor] | None: - """Compatibility adapter for the retained V1 DDTree proposer. - - MRV2 DFlash does not use this method. Keeping the adapter on the - official model avoids retaining the removed experimental flat-DFlash - implementation solely for DDTree payload construction. - """ if self.draft_id_to_target_id is not None: return None return self.logits_processor.get_topk_tokens_and_logprobs( @@ -927,7 +850,7 @@ def precompute_and_store_context_kv( self, context_states: torch.Tensor, context_positions: torch.Tensor, - context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None, + context_slot_mapping: torch.Tensor | dict[str, torch.Tensor] | None = None, ) -> None: """Precompute projected + RoPE'd K/V and write to cache.""" self.model.precompute_and_store_context_kv( @@ -943,17 +866,9 @@ def combine_hidden_states( needs_squeeze = hidden_states.dim() == 1 if needs_squeeze: hidden_states = hidden_states.unsqueeze(0) - expected = self.model.fc.input_size - if hidden_states.shape[-1] != expected: - raise ValueError( - f"DFlash drafter expects {expected} concatenated aux hidden " - f"features but received {hidden_states.shape[-1]}. This usually " - "means the draft model's target_layer_ids reference layers that " - "do not exist in the target model (incompatible draft/target pair)." - ) - # AWQ target paths can surface fp32 auxiliary states even when the - # draft projection is initialized in fp16. Match the FC weight dtype - # before dispatching the linear kernel. + # Keep the auxiliary hidden-state projection on the same dtype as the + # DFlash FC weights. AWQ target paths may surface fp32 hidden states + # even when the draft projection is initialized in fp16. fc_weight = getattr(self.model.fc, "weight", None) if fc_weight is not None and hidden_states.dtype != fc_weight.dtype: hidden_states = hidden_states.to(dtype=fc_weight.dtype) @@ -968,9 +883,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): includes_embed_tokens = False for name, loaded_weight in weights: assert "mask_hidden" not in name, ( - "DFlash embeds masked slots via mask_token_id (optionally " - "overridden by a mask_embedding.pt file); it should not ship a " - "mask_hidden weight." + "DFlash should use mask_token_id to embed the padding hidden state" ) if "t2d" in name: continue @@ -984,13 +897,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): model_weights[name] = loaded_weight process_eagle_weight(self, name) - # Route the separately-trained mask embedding (if shipped) through the - # standard weight loader alongside the rest of the draft weights. - mask_embedding = self._read_mask_embedding() - if mask_embedding is not None: - model_weights["model.mask_embedding"] = mask_embedding - self.model.has_separate_mask_embedding = True - skip_substrs = [] if not includes_draft_id_mapping: skip_substrs.append("draft_id_to_target_id") @@ -998,8 +904,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): skip_substrs.append("embed_tokens") if not self.model.use_aux_hidden_state: skip_substrs.append("fc.") - if not self.model.has_separate_mask_embedding: - skip_substrs.append("mask_embedding") loader = AutoWeightsLoader( self, skip_prefixes=None, @@ -1007,42 +911,3 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): ) loader.load_weights(model_weights.items()) self.model._build_fused_kv_buffers() - - def _read_mask_embedding(self) -> torch.Tensor | None: - """Checks for an override mask embedding in `mask_embedding.pt` and returns it. - - Some checkpoints ship a separately-trained mask embedding for the mask token, - which we use to overwrite the embedding for `mask_token_id`. This helper - checks for the file, loads the pytorch tensor, and returns the embedding to use. - - Returns None if the override file is not present. - """ - mask_token_id = self.model.mask_token_id - if mask_token_id is None: - return None - - MASK_EMBEDDING_FILENAME = "mask_embedding.pt" - data = get_hf_file_bytes( - MASK_EMBEDDING_FILENAME, - self.draft_model_config.model, - self.draft_model_config.revision, - ) - if data is None: - return None - - state = torch.load(io.BytesIO(data), weights_only=True) - if isinstance(state, dict): - if state.get("mask_token_id", mask_token_id) != mask_token_id: - raise ValueError( - f"{MASK_EMBEDDING_FILENAME} mask_token_id does not match " - f"dflash_config.mask_token_id ({mask_token_id}). " - f"Got {state.get('mask_token_id')}." - ) - state = state["embedding"] - - logger.info( - "Loaded DFlash mask embedding for mask_token_id %s from %s", - mask_token_id, - MASK_EMBEDDING_FILENAME, - ) - return state.reshape(-1) diff --git a/vllm/model_executor/models/qwen3_dflash2.py b/vllm/model_executor/models/qwen3_dflash2.py deleted file mode 100644 index af6560385c..0000000000 --- a/vllm/model_executor/models/qwen3_dflash2.py +++ /dev/null @@ -1,484 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from collections.abc import Callable -from functools import cache - -import torch -import torch.nn.functional as F -from torch import nn - -import vllm.envs as envs -from vllm.compilation.backends import set_model_tag -from vllm.compilation.decorators import support_torch_compile -from vllm.config import CacheConfig, VllmConfig -from vllm.distributed import ( - get_tensor_model_parallel_world_size, - tensor_model_parallel_all_gather, -) -from vllm.logger import init_logger -from vllm.model_executor.layers.linear import ( - ColumnParallelLinear, - ReplicatedLinear, - UnquantizedLinearMethod, -) -from vllm.model_executor.layers.quantization.base_config import QuantizationConfig -from vllm.model_executor.layers.vocab_parallel_embedding import ( - UnquantizedEmbeddingMethod, -) -from vllm.platforms import current_platform -from vllm.utils.flashinfer import has_flashinfer - -from .dflash_sm70 import ( - DFLASH_SM70_WIDE_OUTPUT_SCALE, - DFlashSM70MLP, - DFlashSM70RMSNorm, -) -from .qwen3_dflash import ( - DFlashQwen3DecoderLayer, - DFlashQwen3ForCausalLM, - DFlashQwen3Model, -) -from .utils import maybe_prefix - -logger = init_logger(__name__) - - -def _use_sm70_bf16_emulation(config) -> bool: - declared_dtype = getattr(config, "dtype", None) - if declared_dtype is None: - declared_dtype = getattr(config, "torch_dtype", None) - is_bf16 = declared_dtype is torch.bfloat16 or str(declared_dtype).lower() in { - "bfloat16", - "bf16", - "torch.bfloat16", - } - return ( - is_bf16 - and current_platform.is_cuda() - and current_platform.is_device_capability(70) - ) - - -@cache -def _flashinfer_topk() -> Callable[..., tuple[torch.Tensor, torch.Tensor]] | None: - """Return FlashInfer radix top-k only where the installed kernel can run. - - Presence is not a sufficient capability check: FlashInfer can be installed - in the V100 environment while its radix top-k has no SM70 implementation. - """ - if not current_platform.is_cuda(): - return None - if not current_platform.has_device_capability(80): - logger.info_once( - "DFlash2 disables FlashInfer top-k below SM80; using torch.topk." - ) - return None - if not has_flashinfer(): - logger.info_once( - "FlashInfer is unavailable; the DFlash2 selector uses torch.topk." - ) - return None - from flashinfer import top_k - - return top_k - - -def _topk(scores: torch.Tensor, k: int) -> tuple[torch.Tensor, torch.Tensor]: - impl = _flashinfer_topk() - if impl is None or not scores.is_cuda: - return torch.topk(scores, k, dim=-1) - return impl(scores, k, sorted=True, deterministic=True) - - -def _grouped_conv( - hidden_states: torch.Tensor, - delta: torch.Tensor, - base: torch.Tensor, - block_size: int, - num_groups: int, - group_size: int, - taps: int, -) -> torch.Tensor: - blocks = hidden_states.unflatten(-1, (num_groups, group_size)) - coefficients = base.view(1, taps, num_groups, group_size) + delta.unsqueeze(-1) - output = coefficients[:, 0] * blocks - position = torch.arange(hidden_states.shape[0], device=hidden_states.device) - if block_size & (block_size - 1) == 0: - position = position & (block_size - 1) - else: - position = position % block_size - for tap in range(1, taps): - shifted = F.pad(blocks[:-tap], (0, 0, 0, 0, tap, 0)) - output += coefficients[:, tap] * shifted * (position >= tap).view(-1, 1, 1) - return output.flatten(-2) - - -class DFlashGroupedConv(nn.Module): - def __init__( - self, - hidden_size: int, - taps: int, - group_size: int, - block_size: int, - params_dtype: torch.dtype, - prefix: str, - ) -> None: - super().__init__() - if hidden_size % group_size: - raise ValueError( - f"conv_group_size={group_size} must divide hidden_size={hidden_size}." - ) - self.block_size = block_size - self.taps = taps - self.group_size = group_size - self.num_groups = hidden_size // group_size - self.base_kernel = nn.Parameter( - torch.empty(2, taps, hidden_size, dtype=params_dtype), - requires_grad=False, - ) - self.kernel_projection = ReplicatedLinear( - hidden_size, - 2 * taps * self.num_groups, - bias=False, - params_dtype=params_dtype, - quant_config=None, - prefix=maybe_prefix(prefix, "kernel_projection"), - return_bias=False, - ) - - def _convolve( - self, hidden_states: torch.Tensor, delta: torch.Tensor, side: int - ) -> torch.Tensor: - return _grouped_conv( - hidden_states, - delta, - self.base_kernel[side], - self.block_size, - self.num_groups, - self.group_size, - self.taps, - ) - - def prepare(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - coefficients = self.kernel_projection(hidden_states).reshape( - hidden_states.shape[0], 2, self.taps, self.num_groups - ) - return self._convolve(hidden_states, coefficients[:, 0], 0), coefficients[:, 1] - - def finish( - self, hidden_states: torch.Tensor, coefficients: torch.Tensor - ) -> torch.Tensor: - return self._convolve(hidden_states, coefficients, 1) - - -class DFlash2Qwen3DecoderLayer(DFlashQwen3DecoderLayer): - def __init__( - self, - vllm_config: VllmConfig, - *, - config, - layer_idx: int, - cache_config: CacheConfig | None = None, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ) -> None: - super().__init__( - vllm_config, - config=config, - layer_idx=layer_idx, - cache_config=cache_config, - quant_config=quant_config, - prefix=prefix, - ) - self.use_sm70_bf16_emulation = _use_sm70_bf16_emulation(config) - if self.use_sm70_bf16_emulation: - runtime_dtype = vllm_config.model_config.dtype - if runtime_dtype != torch.float16: - raise ValueError( - "BF16 DFlash2 on SM70 requires FP16 runtime transport; " - f"got runtime dtype {runtime_dtype}." - ) - self.self_attn.output_input_scale = DFLASH_SM70_WIDE_OUTPUT_SCALE - self.mlp = DFlashSM70MLP(self.mlp) - self.input_layernorm = DFlashSM70RMSNorm( - config.hidden_size, - config.rms_norm_eps, - runtime_dtype, - ) - self.post_attention_layernorm = DFlashSM70RMSNorm( - config.hidden_size, - config.rms_norm_eps, - runtime_dtype, - ) - draft_config = config.dflash_config - speculative_config = vllm_config.speculative_config - assert speculative_config is not None - conv_args = dict( - 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, - params_dtype=vllm_config.model_config.dtype, - ) - self.attention_conv = DFlashGroupedConv( - **conv_args, prefix=maybe_prefix(prefix, "attention_conv") - ) - self.mlp_conv = DFlashGroupedConv( - **conv_args, prefix=maybe_prefix(prefix, "mlp_conv") - ) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - residual: torch.Tensor | None, - ) -> tuple[torch.Tensor, torch.Tensor]: - if residual is None: - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - else: - hidden_states, residual = self.input_layernorm(hidden_states, residual) - - hidden_states, coefficients = self.attention_conv.prepare(hidden_states) - hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states) - hidden_states = self.attention_conv.finish(hidden_states, coefficients) - - hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) - hidden_states, coefficients = self.mlp_conv.prepare(hidden_states) - hidden_states = self.mlp(hidden_states) - hidden_states = self.mlp_conv.finish(hidden_states, coefficients) - return hidden_states, residual - - -def _score_edges( - predecessor_table: torch.Tensor, - successor_table: torch.Tensor, - candidate_ids: torch.Tensor, - unary_logits: torch.Tensor, - hidden: torch.Tensor, - anchor_token_ids: torch.Tensor, - top_k: int, -) -> torch.Tensor: - successors = successor_table[candidate_ids] - predecessor_ids = torch.cat( - ( - anchor_token_ids[:, None, None].expand(-1, 1, top_k), - candidate_ids[:, :-1], - ), - dim=1, - ) - predecessors = predecessor_table[predecessor_ids] - return unary_logits[:, :, None] + torch.einsum( - "blpr,blcr->blpc", predecessors * hidden[:, :, None], successors - ) - - -@support_torch_compile -class CandidateSelector(nn.Module): - def __init__( - self, - hidden_size: int, - vocab_size: int, - rank: int, - top_k: int, - params_dtype: torch.dtype, - prefix: str, - ) -> None: - super().__init__() - self.top_k = top_k - self.predecessor_codebook = nn.Parameter( - torch.empty(vocab_size, rank, dtype=params_dtype), requires_grad=False - ) - self.successor_codebook = nn.Parameter( - torch.empty(vocab_size, rank, dtype=params_dtype), requires_grad=False - ) - self.hidden_projection = ReplicatedLinear( - hidden_size, - rank, - bias=False, - params_dtype=params_dtype, - quant_config=None, - prefix=maybe_prefix(prefix, "hidden_projection"), - return_bias=False, - ) - - def forward( - self, - candidate_ids: torch.Tensor, - unary_logits: torch.Tensor, - hidden_states: torch.Tensor, - anchor_token_ids: torch.Tensor, - ) -> torch.Tensor: - hidden = self.hidden_projection(hidden_states) - return _score_edges( - self.predecessor_codebook, - self.successor_codebook, - candidate_ids, - unary_logits, - hidden, - anchor_token_ids, - self.top_k, - ) - - -class DFlash2Qwen3Model(DFlashQwen3Model): - decoder_layer_cls = DFlash2Qwen3DecoderLayer - - def _make_context_projection( - self, - *, - vllm_config: VllmConfig, - input_size: int, - output_size: int, - prefix: str, - ) -> nn.Module: - use_sharded_fc = ( - envs.VLLM_SM70_DFLASH2_SHARDED_CONTEXT_FC - and self.quant_config is None - and current_platform.is_cuda() - and current_platform.is_device_capability(70) - and vllm_config.parallel_config.tensor_parallel_size == 4 - and input_size == 25600 - and output_size == 5120 - ) - if not use_sharded_fc: - return super()._make_context_projection( - vllm_config=vllm_config, - input_size=input_size, - output_size=output_size, - prefix=prefix, - ) - - projection = ColumnParallelLinear( - input_size=input_size, - output_size=output_size, - bias=False, - gather_output=True, - params_dtype=vllm_config.model_config.dtype, - quant_config=None, - prefix=prefix, - return_bias=False, - ) - projection._sm70_f16_force_enable = True - # TurboMind wins at the 8-row decode shape, while cuBLAS is faster for - # 128+ context rows. Keep the exact small-M path and let ordinary - # F.linear handle chunked-prefill projection shapes. - projection._sm70_f16_max_m = 64 - logger.info_once( - "Using TP4 output-sharded DFlash2 target-hidden projection on SM70 " - "(25600->5120 plus compact all-gather)." - ) - return projection - - def __init__( - self, - *, - vllm_config: VllmConfig, - start_layer_id: int = 0, - prefix: str = "", - ) -> None: - super().__init__( - vllm_config=vllm_config, - start_layer_id=start_layer_id, - prefix=prefix, - ) - draft_config = self.config.dflash_config - self.use_sm70_bf16_emulation = _use_sm70_bf16_emulation(self.config) - if self.use_sm70_bf16_emulation: - runtime_dtype = vllm_config.model_config.dtype - self.hidden_norm = DFlashSM70RMSNorm( - self.config.hidden_size, - self.config.rms_norm_eps, - runtime_dtype, - residual_scale=1.0, - ) - self.norm = DFlashSM70RMSNorm( - self.config.hidden_size, - self.config.rms_norm_eps, - runtime_dtype, - ) - logger.info_once( - "Using range-preserving BF16 DFlash2 arithmetic on SM70 " - "(FP32 residual, BF16 RNE, scaled FP16 projections)." - ) - self.input_embedding_scale = float( - draft_config.get("input_embedding_scale", 1.0) - ) - # The selector is compiled separately from the draft backbone. A unique - # model tag prevents the two incompatible input signatures from sharing - # a persistent compile-cache namespace. - with set_model_tag("dflash2_candidate_selector"): - self.candidate_selector = CandidateSelector( - hidden_size=self.config.hidden_size, - vocab_size=self.config.vocab_size, - rank=int(draft_config["selector_rank"]), - top_k=int(draft_config["selector_top_k"]), - params_dtype=vllm_config.model_config.dtype, - prefix=maybe_prefix(prefix, "candidate_selector"), - ) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return super().embed_input_ids(input_ids) * self.input_embedding_scale - - def _normalize_context_states(self, context_states: torch.Tensor) -> torch.Tensor: - if self.use_sm70_bf16_emulation: - return self.hidden_norm(context_states) - return super()._normalize_context_states(context_states) - - -class DFlash2Qwen3ForCausalLM(DFlashQwen3ForCausalLM): - model_cls = DFlash2Qwen3Model - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: - super().__init__(vllm_config=vllm_config, prefix=prefix) - draft_config = self.config.dflash_config - self.output_multiplier = float(draft_config.get("output_multiplier", 1.0)) - softcap = float(draft_config.get("final_logit_softcapping") or 0.0) - self.final_logit_softcapping = softcap if softcap > 0 else None - - def compute_candidates( - self, hidden_states: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor]: - if not isinstance( - self.lm_head.quant_method, - (UnquantizedEmbeddingMethod, UnquantizedLinearMethod), - ): - raise ValueError( - "DFlash2 requires an unquantized target LM head for candidate TopK; " - f"got {type(self.lm_head.quant_method).__name__}." - ) - - selector = self.model.candidate_selector - local_candidates = self.lm_head.maybe_get_sm70_dflash2_top20( - hidden_states, selector.top_k - ) - if local_candidates is None: - logits = self.lm_head.quant_method.apply( - self.lm_head, hidden_states, bias=None - ) - num_pad = self.lm_head.shard_indices.num_org_vocab_padding - if num_pad > 0: - logits[..., -num_pad:] = -float("inf") - values, ids = _topk(logits, selector.top_k) - ids = ids.to(torch.int64) + self.lm_head.shard_indices.org_vocab_start_index - else: - values, ids = local_candidates - - if get_tensor_model_parallel_world_size() > 1: - values = tensor_model_parallel_all_gather(values, dim=-1) - ids = tensor_model_parallel_all_gather(ids, dim=-1) - - if values.shape[-1] > selector.top_k: - values, selected = _topk(values, selector.top_k) - ids = ids.gather(-1, selected) - - values = values.float() * self.output_multiplier - if self.final_logit_softcapping is not None: - cap = self.final_logit_softcapping - values = torch.tanh(values / cap) * cap - return ids, values - - -EntryClass = DFlash2Qwen3ForCausalLM diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index b54454adad..cf4034724a 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/vllm/model_executor/models/qwen3_next.py @@ -549,7 +549,7 @@ def __init__( def forward( self, positions: torch.Tensor, - output: torch.Tensor | None, + output: torch.Tensor, hidden_states: torch.Tensor, ): qkv, _ = self.qkv_proj(hidden_states) @@ -623,16 +623,13 @@ def forward( attn_output, ) - projected_output, _ = self.o_proj(attn_output) - if output is not None: - output[:] = projected_output + output[:], _ = self.o_proj(attn_output) _sm70_dump_qwen_layer_tensor( "full_attn_o_proj_out", self.layer_idx, "full_attention", - projected_output, + output, ) - return projected_output class Qwen3NextDecoderLayer(nn.Module): @@ -747,33 +744,21 @@ def forward( hidden_states, ) - use_direct_attention_output = ( - envs.VLLM_SM70_TP4_LONG_PREFILL_FUSED_NORM and torch.compiler.is_compiling() - ) - self_attention_output = ( - None if use_direct_attention_output else torch.empty_like(hidden_states) - ) + self_attention_output = torch.empty_like(hidden_states) if self.layer_type == "linear_attention": - projected_attention_output = self.linear_attn( + self.linear_attn( hidden_states=hidden_states, output=self_attention_output, ) elif self.layer_type == "full_attention": - projected_attention_output = self.self_attn( + self.self_attn( hidden_states=hidden_states, output=self_attention_output, positions=positions, ) else: raise ValueError("Invalid layer_type") - if use_direct_attention_output: - if projected_attention_output is None: - raise RuntimeError( - "SM70 TP4 fused prefill requires a direct attention output" - ) - hidden_states = projected_attention_output - else: - hidden_states = self_attention_output + hidden_states = self_attention_output hidden_states = _sm70_dump_qwen_layer_tensor( "attn_out", self.layer_idx, diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 5723dbe285..3a7a7b4ba0 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -598,7 +598,6 @@ "EagleLlama4ForCausalLM": ("llama4_eagle", "EagleLlama4ForCausalLM"), "EagleMiniCPMForCausalLM": ("minicpm_eagle", "EagleMiniCPMForCausalLM"), "DFlashDraftModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"), - "DFlash2DraftModel": ("qwen3_dflash2", "DFlash2Qwen3ForCausalLM"), "DSparkDraftModel": ( "vllm.models.deepseek_v4", "DSparkDeepseekV4ForCausalLM", diff --git a/vllm/model_executor/offloader/base.py b/vllm/model_executor/offloader/base.py index b29b57d34f..ceff60cd4c 100644 --- a/vllm/model_executor/offloader/base.py +++ b/vllm/model_executor/offloader/base.py @@ -76,18 +76,6 @@ def post_init(self): """ return - def wrap_module( - self, module: nn.Module, *, parameter_prefix: str = "" - ) -> nn.Module: - """Optionally offload one non-layer component. - - Layer-prefetch backends intentionally inherit this no-op. UVA - overrides it so model components created outside ``make_layers`` can - participate without invoking a backend's multi-layer wrapper twice. - """ - del parameter_prefix - return module - def sync_prev_onload(self) -> None: # noqa: B027 """Sync previous onload operations. Override in subclasses.""" pass diff --git a/vllm/model_executor/offloader/uva.py b/vllm/model_executor/offloader/uva.py index 67218d5dd6..51eb1a14fc 100644 --- a/vllm/model_executor/offloader/uva.py +++ b/vllm/model_executor/offloader/uva.py @@ -61,30 +61,7 @@ def wrap_modules( ) return modules - def wrap_module( - self, module: nn.Module, *, parameter_prefix: str = "" - ) -> nn.Module: - """Apply UVA offloading to one component with its qualified prefix.""" - previous_offload_bytes = self.cpu_offload_bytes - module = self._maybe_offload_to_cpu(module, parameter_prefix=parameter_prefix) - if self.cpu_offload_bytes > previous_offload_bytes: - logger.info( - "Total CPU offloaded parameters: %s", - format_gib(self.cpu_offload_bytes), - ) - return module - - def _is_parameter_selected(self, name: str, parameter_prefix: str = "") -> bool: - if not self.cpu_offload_params: - return True - qualified_name = f"{parameter_prefix}.{name}" if parameter_prefix else name - return any( - f".{param}." in f".{qualified_name}." for param in self.cpu_offload_params - ) - - def _maybe_offload_to_cpu( - self, module: nn.Module, *, parameter_prefix: str = "" - ) -> nn.Module: + def _maybe_offload_to_cpu(self, module: nn.Module) -> nn.Module: """Offload module parameters to CPU using UVA if budget allows.""" if (params := next(module.parameters(), None)) is None: return module @@ -106,17 +83,16 @@ def _maybe_offload_to_cpu( # one module might have some parameters offloaded and some not break - # Add dots to require exact path-segment matches. A qualified - # prefix lets selectors such as ``visual`` target an entire tower - # even though ``module.named_parameters()`` returns relative names. - if not self._is_parameter_selected(name, parameter_prefix): - continue - - # A component may contain layers already visited by make_layers(). - # Do not copy or account for the same UVA parameter twice when the - # enclosing component is registered afterwards. - if getattr(p, "_vllm_is_uva_offloaded", False): - continue + if self.cpu_offload_params: + # Check if parameter belongs to the offloading set + # Add dots here to ensure we match full segments only + # e.g., "experts.w2_weight" matches "mlp.experts.w2_weight" + # but not "mlp.experts.w2_weight_scale" + should_offload = any( + f".{param}." in f".{name}." for param in self.cpu_offload_params + ) + if not should_offload: + continue cpu_data = p.data.to(device="cpu") if self.pin_memory: diff --git a/vllm/model_executor/warmup/awq_sm70_warmup.py b/vllm/model_executor/warmup/awq_sm70_warmup.py index 9bc11324d9..e12589a2c7 100644 --- a/vllm/model_executor/warmup/awq_sm70_warmup.py +++ b/vllm/model_executor/warmup/awq_sm70_warmup.py @@ -15,9 +15,6 @@ from vllm import _sm70_ops as sm70_ops from vllm.logger import init_logger from vllm.model_executor.layers.quantization import sm70_turbomind as sm70_tm -from vllm.model_executor.layers.quantization.nvfp4_sm70_moe import ( - _prepare_compact_slot_groups, -) if TYPE_CHECKING: from vllm.v1.worker.gpu_worker import Worker @@ -47,7 +44,6 @@ def _lut_cache_disabled_for_dynamic_quant_dispatch( has_fp8_dense: bool, fp4_kinds: set[str], has_mxfp4_moe: bool = False, - has_nvfp4_moe: bool = False, ) -> bool: awq_no_preserve = ( has_awq_dense @@ -63,9 +59,7 @@ def _lut_cache_disabled_for_dynamic_quant_dispatch( mxfp4_dynamic = ( "mxfp4" in fp4_kinds or has_mxfp4_moe ) and envs.VLLM_SM70_MXFP4_TUNE_SMALL_SHAPES - nvfp4_dynamic = ( - "nvfp4" in fp4_kinds or has_nvfp4_moe - ) and envs.VLLM_SM70_NVFP4_TUNE_SMALL_SHAPES + nvfp4_dynamic = "nvfp4" in fp4_kinds and envs.VLLM_SM70_NVFP4_TUNE_SMALL_SHAPES return awq_no_preserve or fp8_dynamic or mxfp4_dynamic or nvfp4_dynamic @@ -265,15 +259,7 @@ def _iter_unique_fp8_dense_layers( ) -> Iterable[tuple[torch.nn.Module, bool]]: seen: set[tuple[int, int, bool]] = set() for layer in model.modules(): - if not ( - getattr(layer, "sm70_fp8_turbomind", False) - or getattr(layer, "sm70_modelopt_fp8_turbomind", False) - ): - continue - # QPN8 has a static source-selected dispatch and does not populate the - # TurboMind LUT warmed below. CUDA graph capture exercises its admitted - # M<=8 kernels separately. - if getattr(layer, "sm70_fp8_qpn8", False): + if not getattr(layer, "sm70_fp8_turbomind", False): continue if getattr(layer, "sm70_fp8_bmm", False): @@ -307,20 +293,14 @@ def _iter_unique_fp8_dense_layers( def _iter_unique_fp4_dense_layers(model: torch.nn.Module) -> Iterable[Any]: - seen: set[tuple[str, int, int, int, bool]] = set() + seen: set[tuple[str, int, int, int]] = set() for layer in model.modules(): state = getattr(layer, sm70_tm.STATE_ATTR, None) if state is None or state.op_kind not in ("mxfp4", "nvfp4"): continue k_dim = int(state.weight.shape[0]) n_dim = int(state.output_size) - key = ( - state.op_kind, - k_dim, - n_dim, - int(state.group_size), - bool(state.gated_silu), - ) + key = (state.op_kind, k_dim, n_dim, int(state.group_size)) if key in seen: continue seen.add(key) @@ -407,45 +387,6 @@ def _iter_unique_mxfp4_moe_layers( yield layer -def _is_nvfp4_moe_sm70_layer(layer: torch.nn.Module) -> bool: - return getattr(layer, "sm70_nvfp4_moe", False) and all( - hasattr(layer, name) - for name in ( - "w13_strided_ptrs_w", - "w13_strided_ptrs_s", - "w2_strided_ptrs_w", - "w2_strided_ptrs_s", - "sm70_nvfp4_num_experts", - "sm70_nvfp4_w13_k_dim", - "sm70_nvfp4_w13_n_dim", - "sm70_nvfp4_w2_k_dim", - "sm70_nvfp4_w2_n_dim", - "sm70_nvfp4_group_size", - "_nvfp4_sm70_dense_expert_ids", - ) - ) - - -def _iter_unique_nvfp4_moe_layers( - model: torch.nn.Module, -) -> Iterable[torch.nn.Module]: - seen: set[tuple[int, int, int, int, int]] = set() - for layer in model.modules(): - if not _is_nvfp4_moe_sm70_layer(layer): - continue - key = ( - int(layer.sm70_nvfp4_w13_k_dim), - int(layer.sm70_nvfp4_w13_n_dim), - int(layer.sm70_nvfp4_w2_k_dim), - int(layer.sm70_nvfp4_w2_n_dim), - int(layer.sm70_nvfp4_num_experts), - ) - if key in seen: - continue - seen.add(key) - yield layer - - def _warmup_dense_layers( dense_layers: list[torch.nn.Module], m_values: list[int], @@ -482,7 +423,6 @@ def _warmup_fp8_dense_layers( calls = 0 for layer, gated_silu in dense_layers: - is_modelopt = getattr(layer, "sm70_modelopt_fp8_turbomind", False) if getattr(layer, "sm70_fp8_bmm", False): assert not gated_silu weight = layer.weight[0] @@ -504,14 +444,9 @@ def _warmup_fp8_dense_layers( n_dim = int(weight.shape[1]) // 2 else: weight = layer.weight - if is_modelopt: - scales = layer.weight_scale - k_ld = int(layer.sm70_modelopt_fp8_k_ld) - q_ld = int(layer.sm70_modelopt_fp8_q_ld) - else: - scales = layer.weight_scale_inv - k_ld = int(layer.sm70_fp8_k_ld) - q_ld = int(layer.sm70_fp8_q_ld) + scales = layer.weight_scale_inv + k_ld = int(layer.sm70_fp8_k_ld) + q_ld = int(layer.sm70_fp8_q_ld) n_dim = int(layer.output_size_per_partition) device = weight.device @@ -523,27 +458,6 @@ def _warmup_fp8_dense_layers( out, x, weight, scales, 128, k_ld, q_ld, gated_silu ) calls += 1 - if getattr(layer, "sm70_fp8_bmm_grouped_decode", False): - group_count = int(layer.sm70_fp8_bmm_groups) - grouped_x = torch.empty( - (group_count, k_dim), dtype=torch.float16, device=device - ) - grouped_out = torch.empty( - (group_count, n_dim), dtype=torch.float16, device=device - ) - sm70_ops.fp8_moe_gemm_sm70_per_expert_dispatch_out( - grouped_out, - grouped_x, - layer.sm70_fp8_bmm_grouped_offsets, - layer.sm70_fp8_bmm_grouped_ptrs_w, - layer.sm70_fp8_bmm_grouped_ptrs_s, - group_count, - k_dim, - n_dim, - 128, - False, - ) - calls += 1 return calls @@ -555,8 +469,7 @@ def _warmup_fp4_dense_layers( for state in dense_layers: device = state.weight.device k_dim = int(state.weight.shape[0]) - gated_silu = bool(state.gated_silu) - n_dim = int(state.output_size) // (2 if gated_silu else 1) + n_dim = int(state.output_size) for m_dim in m_values: x = torch.empty((m_dim, k_dim), dtype=torch.float16, device=device) out = torch.empty((m_dim, n_dim), dtype=torch.float16, device=device) @@ -571,7 +484,7 @@ def _warmup_fp4_dense_layers( int(state.group_size), int(state.k_ld), int(state.q_ld), - gated_silu, + False, ) elif state.op_kind == "nvfp4": if not hasattr(torch.ops._C, "nvfp4_gemm_sm70_out"): @@ -584,7 +497,7 @@ def _warmup_fp4_dense_layers( int(state.group_size), int(state.k_ld), int(state.q_ld), - gated_silu, + False, ) else: continue @@ -731,133 +644,6 @@ def _warmup_mxfp4_moe_b1_layers(moe_layers: list[torch.nn.Module]) -> int: return calls -def _warmup_nvfp4_moe_decode_layers( - moe_layers: list[torch.nn.Module], token_counts: list[int] -) -> int: - """Tune compact and full-group NVFP4 graph shapes.""" - if not hasattr(torch.ops._C, "nvfp4_moe_dense_stage_sm70_out"): - return 0 - - calls = 0 - for layer in moe_layers: - device = layer.w13_tm_weight.device - top_k = int(layer.moe_config.experts_per_token) - max_experts = int(layer.sm70_nvfp4_num_experts) - compact_max_tokens = int( - getattr(layer, "sm70_nvfp4_compact_grouped_max_tokens", 8) - ) - for num_tokens in token_counts: - total_slots = num_tokens * top_k - if num_tokens <= compact_max_tokens: - stage_experts = total_slots - sorted_expert_ids = layer._nvfp4_sm70_dense_expert_ids[:total_slots] - expert_offsets = torch.empty( - total_slots + 1, dtype=torch.int32, device=device - ) - active_expert_ids = torch.empty( - total_slots, dtype=torch.int32, device=device - ) - _prepare_compact_slot_groups( - sorted_expert_ids, - expert_offsets, - active_expert_ids, - ) - else: - stage_experts = max_experts - expert_offsets = _build_balanced_offsets( - total_slots, max_experts, device - ) - active_expert_ids = layer._nvfp4_sm70_dense_expert_ids - permuted_input = torch.empty( - total_slots, - int(layer.sm70_nvfp4_w13_k_dim), - dtype=torch.float16, - device=device, - ) - gate_up = torch.empty( - total_slots, - int(layer.sm70_nvfp4_w13_n_dim), - dtype=torch.float16, - device=device, - ) - intermediate = torch.empty( - total_slots, - int(layer.sm70_nvfp4_w2_k_dim), - dtype=torch.float16, - device=device, - ) - sorted_output = torch.empty( - total_slots, - int(layer.sm70_nvfp4_w2_n_dim), - dtype=torch.float16, - device=device, - ) - sm70_ops.nvfp4_moe_dense_stage_sm70_out( - gate_up, - permuted_input, - expert_offsets, - active_expert_ids, - layer.w13_strided_ptrs_w, - layer.w13_strided_ptrs_s, - stage_experts, - int(layer.sm70_nvfp4_w13_k_dim), - int(layer.sm70_nvfp4_w13_n_dim), - int(layer.sm70_nvfp4_group_size), - ) - if layer.swiglu_limit is None: - torch.ops._C.silu_and_mul(intermediate, gate_up) - else: - torch.ops._C.silu_and_mul_with_clamp( - intermediate, gate_up, float(layer.swiglu_limit) - ) - sm70_ops.nvfp4_moe_dense_stage_sm70_out( - sorted_output, - intermediate, - expert_offsets, - active_expert_ids, - layer.w2_strided_ptrs_w, - layer.w2_strided_ptrs_s, - stage_experts, - int(layer.sm70_nvfp4_w2_k_dim), - int(layer.sm70_nvfp4_w2_n_dim), - int(layer.sm70_nvfp4_group_size), - ) - calls += 2 - return calls - - -def _get_nvfp4_moe_token_counts( - worker: Worker, - moe_layers: list[torch.nn.Module], - base_token_counts: list[int], -) -> list[int]: - """Include graph-safe and tuned NVFP4 CUDA-graph token widths.""" - if not moe_layers: - return base_token_counts - graph_safe_max_tokens = max( - int(getattr(layer, "sm70_nvfp4_graph_safe_max_tokens", 8)) - for layer in moe_layers - ) - compact_max_tokens = max( - int(getattr(layer, "sm70_nvfp4_compact_grouped_max_tokens", 8)) - for layer in moe_layers - ) - max_top_k = max(int(layer.moe_config.experts_per_token) for layer in moe_layers) - tuned_max_tokens = max(0, int(envs.VLLM_SM70_NVFP4_MOE_TUNE_MAX_TOKENS)) // max( - max_top_k, 1 - ) - max_tokens = max(graph_safe_max_tokens, tuned_max_tokens) - compilation_config = getattr( - getattr(worker, "vllm_config", None), "compilation_config", None - ) - capture_sizes = getattr(compilation_config, "cudagraph_capture_sizes", None) or [] - return sorted( - set(base_token_counts) - | set(range(1, compact_max_tokens + 1)) - | {int(size) for size in capture_sizes if 0 < int(size) <= max_tokens} - ) - - def _warmup_moe_single_token_layers(moe_layers: list[torch.nn.Module]) -> int: if not ( hasattr(torch.ops._C, "awq_moe_single_token_dense_w13_sm70_out") @@ -1011,14 +797,12 @@ def sm70_awq_warmup(worker: Worker) -> None: fp4_dense_layers = list(_iter_unique_fp4_dense_layers(model)) moe_layers = list(_iter_unique_moe_layers(model)) mxfp4_moe_layers = list(_iter_unique_mxfp4_moe_layers(model)) - nvfp4_moe_layers = list(_iter_unique_nvfp4_moe_layers(model)) if not ( dense_layers or fp8_dense_layers or fp4_dense_layers or moe_layers or mxfp4_moe_layers - or nvfp4_moe_layers ): return @@ -1028,7 +812,6 @@ def sm70_awq_warmup(worker: Worker) -> None: bool(fp8_dense_layers), fp4_kinds, bool(mxfp4_moe_layers), - bool(nvfp4_moe_layers), ) imported_records = _load_lut_cache(device, skip_import=skip_lut_cache) if imported_records > 0: @@ -1040,26 +823,20 @@ def sm70_awq_warmup(worker: Worker) -> None: m_values = _get_decode_m_values(worker) moe_token_counts = _get_moe_token_counts(worker) - nvfp4_moe_token_counts = _get_nvfp4_moe_token_counts( - worker, nvfp4_moe_layers, moe_token_counts - ) lut_path = _resolve_lut_path(device) logger.info( "Warming up SM70 TurboMind accepted routes (%d AWQ dense layer " "shapes, %d FP8 dense layer shapes, %d FP4 dense layer shapes, " - "%d MoE layer shapes, %d MXFP4 MoE B1 layer shapes, " - "%d NVFP4 MoE layer shapes, dense_m=%s, " - "moe_tokens=%s, nvfp4_moe_tokens=%s, lut_path=%s).", + "%d MoE layer shapes, %d MXFP4 MoE B1 layer shapes, dense_m=%s, " + "moe_tokens=%s, lut_path=%s).", len(dense_layers), len(fp8_dense_layers), len(fp4_dense_layers), len(moe_layers), len(mxfp4_moe_layers), - len(nvfp4_moe_layers), m_values, moe_token_counts, - nvfp4_moe_token_counts, lut_path, ) with torch.inference_mode(): @@ -1073,22 +850,17 @@ def sm70_awq_warmup(worker: Worker) -> None: moe_stage_calls = _warmup_moe_dense_stage_layers(moe_layers, moe_token_counts) single_token_calls = _warmup_moe_single_token_layers(moe_layers) mxfp4_moe_stage_calls = _warmup_mxfp4_moe_b1_layers(mxfp4_moe_layers) - nvfp4_moe_stage_calls = _warmup_nvfp4_moe_decode_layers( - nvfp4_moe_layers, nvfp4_moe_token_counts - ) torch.accelerator.synchronize(device) logger.info( "SM70 TurboMind warmup finished (%d AWQ dense calls, %d FP8 dense " "calls, %d FP4 dense calls, %d MoE stage calls, " - "%d single-token active-expert calls, %d MXFP4 MoE B1 stage calls, " - "%d NVFP4 MoE stage calls).", + "%d single-token active-expert calls, %d MXFP4 MoE B1 stage calls).", dense_calls, fp8_dense_calls, fp4_dense_calls, moe_stage_calls, single_token_calls, mxfp4_moe_stage_calls, - nvfp4_moe_stage_calls, ) exported_records = _save_lut_cache(device, skip_export=skip_lut_cache) if exported_records > 0: diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index bbdd5b1053..0920e6db4d 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -308,7 +308,6 @@ def __init__( rotate=True, prefix=f"{prefix}.compressor", k_cache_prefix=self.mla_attn.prefix, - validity_cache_prefix=self.swa_cache_layer.prefix, ) def forward( @@ -416,11 +415,6 @@ 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 @@ -440,11 +434,6 @@ 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 @@ -457,11 +446,6 @@ 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 @@ -866,7 +850,6 @@ 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" @@ -940,7 +923,7 @@ def forward( attn_metadata = get_forward_context().attn_metadata if isinstance(attn_metadata, dict): indexer_metadata = cast(Any, attn_metadata[self.k_cache.prefix]) - if indexer_metadata.compressed_max_seq_len <= self.topk_tokens: + if indexer_metadata.max_seq_len // self.compress_ratio <= self.topk_tokens: compressor(compressed_kv_score, positions, rotary_emb) assert self.topk_indices_buffer is not None num_tokens = ( diff --git a/vllm/models/deepseek_v4/common/ops/__init__.py b/vllm/models/deepseek_v4/common/ops/__init__.py index ae26189162..dc6f3c608d 100644 --- a/vllm/models/deepseek_v4/common/ops/__init__.py +++ b/vllm/models/deepseek_v4/common/ops/__init__.py @@ -11,11 +11,7 @@ from .fused_inv_rope_fp8_quant import fused_inv_rope_fp8_quant from .fused_mtp_input_rmsnorm import fused_mtp_input_rmsnorm, mtp_shared_head_rmsnorm from .fused_qk_rmsnorm import fused_q_kv_rmsnorm -from .save_partial_states import ( - save_partial_states, - save_partial_states_to_ring, - stage_partial_states_from_ring, -) +from .save_partial_states import save_partial_states __all__ = [ "MXFP4_BLOCK_SIZE", @@ -29,6 +25,4 @@ "mtp_shared_head_rmsnorm", "quantize_and_insert_k_cache", "save_partial_states", - "save_partial_states_to_ring", - "stage_partial_states_from_ring", ] diff --git a/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py b/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py index 9cec99251f..f5e0886972 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py +++ b/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py @@ -30,148 +30,8 @@ from .fused_indexer_q import _fp32x2_to_fp4x2 -@triton.jit -def _gather_compressor_window( - state_cache_ptr, - state_cache_stride0, - state_cache_stride1, - current_kv_ptr, - current_kv_stride, - current_score_ptr, - current_score_stride, - ape_ptr, - ape_stride, - block_table_ptr, - block_table_stride, - positions_ptr, - position, - token_idx, - req_idx, - block_size, - block, - mask, - HEAD_SIZE: tl.constexpr, - STATE_WIDTH: tl.constexpr, - COMPRESS_RATIO: tl.constexpr, - OVERLAP: tl.constexpr, - USE_PRIVATE_STATE: tl.constexpr, - USE_DENSE_PRIVATE_STATE: tl.constexpr, - RING_SIZE: tl.constexpr, - PRIVATE_HISTORY_SIZE: tl.constexpr, -): - """Gather a compression window from paged state or current input + ring.""" - tokens = tl.arange(0, (1 + OVERLAP) * COMPRESS_RATIO) - pos = position - (1 + OVERLAP) * COMPRESS_RATIO + 1 + tokens - mask_pos = pos >= 0 - head_offset = (tokens >= COMPRESS_RATIO).to(tl.int32) * HEAD_SIZE - - if USE_DENSE_PRIVATE_STATE: - current_start = position - token_idx - dense_rows = pos - current_start + PRIVATE_HISTORY_SIZE - row_base = ( - state_cache_ptr - + req_idx.to(tl.int64) * state_cache_stride0 - + dense_rows.to(tl.int64) * state_cache_stride1 - + head_offset - ) - combined_mask = mask_pos[:, None] & mask[None, :] - score = tl.load( - row_base[:, None] + STATE_WIDTH + block[None, :], - mask=combined_mask, - other=float("-inf"), - ) - kv = tl.load( - row_base[:, None] + block[None, :], - mask=combined_mask, - other=0.0, - ) - return kv, score - - if USE_PRIVATE_STATE: - # This route is gated to one request with contiguous positions. - current_start = position - token_idx - current_idx = pos - current_start - from_current = mask_pos & (current_idx >= 0) - from_history = mask_pos & (current_idx < 0) - current_mask = from_current[:, None] & mask[None, :] - history_mask = from_history[:, None] & mask[None, :] - offsets = head_offset[:, None] + block[None, :] - - current_kv = tl.load( - current_kv_ptr - + current_idx[:, None].to(tl.int64) * current_kv_stride - + offsets, - mask=current_mask, - other=0.0, - ).to(tl.float32) - current_score = tl.load( - current_score_ptr - + current_idx[:, None].to(tl.int64) * current_score_stride - + offsets, - mask=current_mask, - other=0.0, - ).to(tl.float32) - ape = tl.load( - ape_ptr - + (pos % COMPRESS_RATIO)[:, None].to(tl.int64) * ape_stride - + offsets, - mask=current_mask, - other=0.0, - ).to(tl.float32) - - safe_pos = tl.maximum(pos, 0) - ring_base = ( - state_cache_ptr - + req_idx.to(tl.int64) * state_cache_stride0 - + (safe_pos % RING_SIZE).to(tl.int64) * state_cache_stride1 - + head_offset - ) - history_kv = tl.load( - ring_base[:, None] + block[None, :], - mask=history_mask, - other=0.0, - ) - history_score = tl.load( - ring_base[:, None] + STATE_WIDTH + block[None, :], - mask=history_mask, - other=float("-inf"), - ) - kv = tl.where(current_mask, current_kv, history_kv) - score = tl.where(current_mask, current_score + ape, history_score) - return kv, score - - block_indices = pos // block_size - block_numbers = tl.load( - block_table_ptr + req_idx * block_table_stride + block_indices, - mask=mask_pos, - other=0, - ) - block_offsets = pos % block_size - row_base = ( - state_cache_ptr - + block_numbers.to(tl.int64) * state_cache_stride0 - + block_offsets * state_cache_stride1 - + head_offset - ) - combined_mask = mask_pos[:, None] & mask[None, :] - score = tl.load( - row_base[:, None] + STATE_WIDTH + block[None, :], - mask=combined_mask, - other=float("-inf"), - ) - kv = tl.load( - row_base[:, None] + block[None, :], - mask=combined_mask, - other=0.0, - ) - return kv, score - - def compress_norm_rope_store_triton( state_cache: torch.Tensor, - current_kv: torch.Tensor, - current_score: torch.Tensor, - ape: torch.Tensor, num_actual: int, token_to_req_indices: torch.Tensor, positions: torch.Tensor, @@ -193,10 +53,6 @@ def compress_norm_rope_store_triton( quant_block: int, token_stride: int, scale_dim: int, - use_private_state: bool = False, - use_dense_private_state: bool = False, - ring_size: int = 1, - private_history_size: int = 1, ) -> None: """Shared triton launcher for the fused compress+norm+RoPE+insert path. @@ -235,12 +91,6 @@ def compress_norm_rope_store_triton( kv_cache, k_cache_metadata.slot_mapping, kv_cache.shape[1], # paged KV cache block size (tokens per block) - current_kv, - current_kv.stride(0), - current_score, - current_score.stride(0), - ape, - ape.stride(0), # constexprs HEAD_SIZE=head_dim, TRITON_BLOCK_SIZE=triton.next_power_of_2(head_dim), @@ -256,10 +106,6 @@ def compress_norm_rope_store_triton( USE_SOFTWARE_FP8=( current_platform.is_cuda() and current_platform.is_device_capability((7, 0)) ), - USE_PRIVATE_STATE=use_private_state, - USE_DENSE_PRIVATE_STATE=use_dense_private_state, - RING_SIZE=ring_size, - PRIVATE_HISTORY_SIZE=private_history_size, num_warps=num_warps, **pdl_kwargs, ) @@ -291,13 +137,6 @@ def _fused_kv_compress_norm_rope_insert_sparse_attn( k_cache_ptr, kv_slot_mapping_ptr, kv_cache_block_size, - # current chunk (private-state path) - current_kv_ptr, - current_kv_stride, - current_score_ptr, - current_score_stride, - ape_ptr, - ape_stride, # ── constexprs ── HEAD_SIZE: tl.constexpr, TRITON_BLOCK_SIZE: tl.constexpr, @@ -311,10 +150,6 @@ def _fused_kv_compress_norm_rope_insert_sparse_attn( SCALE_DIM: tl.constexpr, # 8 for DeepseekV4 (7 real + 1 pad) KV_BLOCK_STRIDE: tl.constexpr, USE_SOFTWARE_FP8: tl.constexpr, - USE_PRIVATE_STATE: tl.constexpr, - USE_DENSE_PRIVATE_STATE: tl.constexpr, - RING_SIZE: tl.constexpr, - PRIVATE_HISTORY_SIZE: tl.constexpr, ): """Fused compress → RMSNorm → FP8 quant (nope) → RoPE → bf16 store (rope). @@ -337,40 +172,48 @@ def _fused_kv_compress_norm_rope_insert_sparse_attn( req_idx = tl.load(token_to_req_indices_ptr + token_idx) # ── Gather state cache entries ──────────────────────────────────── + start = position - (1 + OVERLAP) * COMPRESS_RATIO + 1 + tokens = tl.arange(0, (1 + OVERLAP) * COMPRESS_RATIO) + pos = start + tokens + mask_pos = pos >= 0 + + block_indices = pos // block_size + block_numbers = tl.load( + block_table_ptr + req_idx * block_table_stride + block_indices, + mask=mask_pos, + other=0, + ) + block_offsets = pos % block_size + head_offset = (tokens >= COMPRESS_RATIO).to(tl.int32) * HEAD_SIZE + block = tl.arange(0, TRITON_BLOCK_SIZE) mask = block < HEAD_SIZE - kv, score = _gather_compressor_window( - state_cache_ptr, - state_cache_stride0, - state_cache_stride1, - current_kv_ptr, - current_kv_stride, - current_score_ptr, - current_score_stride, - ape_ptr, - ape_stride, - block_table_ptr, - block_table_stride, - positions_ptr, - position, - token_idx, - req_idx, - block_size, - block, - mask, - HEAD_SIZE, - STATE_WIDTH, - COMPRESS_RATIO, - OVERLAP, - USE_PRIVATE_STATE, - USE_DENSE_PRIVATE_STATE, - RING_SIZE, - PRIVATE_HISTORY_SIZE, + block_numbers_i64 = block_numbers.to(tl.int64) + + # Precomputed row base shared by score and kv loads + row_base = ( + state_cache_ptr + + block_numbers_i64 * state_cache_stride0 + + block_offsets * state_cache_stride1 + + head_offset ) + combined_mask = mask_pos[:, None] & mask[None, :] + # ── Softmax + weighted sum ─────────────────────────────────────── + score = tl.load( + row_base[:, None] + STATE_WIDTH + block[None, :], + mask=combined_mask, + other=float("-inf"), + ) score = tl.softmax(score, dim=0) + kv = tl.load( + row_base[:, None] + block[None, :], + mask=combined_mask, + other=0.0, + ) + compressed_kv = tl.sum(kv * score, axis=0) # [TRITON_BLOCK_SIZE] fp32 # ── RMSNorm (fp32 throughout) ────────────────────────────────────── @@ -488,13 +331,6 @@ def _fused_kv_compress_norm_rope_insert_indexer_attn( k_cache_ptr, kv_slot_mapping_ptr, kv_cache_block_size, - # current chunk (private-state path) - current_kv_ptr, - current_kv_stride, - current_score_ptr, - current_score_stride, - ape_ptr, - ape_stride, # ── constexprs ── HEAD_SIZE: tl.constexpr, TRITON_BLOCK_SIZE: tl.constexpr, @@ -508,10 +344,6 @@ def _fused_kv_compress_norm_rope_insert_indexer_attn( SCALE_DIM: tl.constexpr, # 4 for indexer (1 float32) KV_BLOCK_STRIDE: tl.constexpr, USE_SOFTWARE_FP8: tl.constexpr, - USE_PRIVATE_STATE: tl.constexpr, - USE_DENSE_PRIVATE_STATE: tl.constexpr, - RING_SIZE: tl.constexpr, - PRIVATE_HISTORY_SIZE: tl.constexpr, ): """Fused compress → RMSNorm → RoPE → FP8 quant → store. @@ -538,38 +370,46 @@ def _fused_kv_compress_norm_rope_insert_indexer_attn( req_idx = tl.load(token_to_req_indices_ptr + token_idx) # ── Gather state cache entries ──────────────────────────────────── + start = position - (1 + OVERLAP) * COMPRESS_RATIO + 1 + tokens = tl.arange(0, (1 + OVERLAP) * COMPRESS_RATIO) + pos = start + tokens + mask_pos = pos >= 0 + + block_indices = pos // block_size + block_numbers = tl.load( + block_table_ptr + req_idx * block_table_stride + block_indices, + mask=mask_pos, + other=0, + ) + block_offsets = pos % block_size + head_offset = (tokens >= COMPRESS_RATIO).to(tl.int32) * HEAD_SIZE + block = tl.arange(0, TRITON_BLOCK_SIZE) mask = block < HEAD_SIZE - kv, score = _gather_compressor_window( - state_cache_ptr, - state_cache_stride0, - state_cache_stride1, - current_kv_ptr, - current_kv_stride, - current_score_ptr, - current_score_stride, - ape_ptr, - ape_stride, - block_table_ptr, - block_table_stride, - positions_ptr, - position, - token_idx, - req_idx, - block_size, - block, - mask, - HEAD_SIZE, - STATE_WIDTH, - COMPRESS_RATIO, - OVERLAP, - USE_PRIVATE_STATE, - USE_DENSE_PRIVATE_STATE, - RING_SIZE, - PRIVATE_HISTORY_SIZE, + block_numbers_i64 = block_numbers.to(tl.int64) + + row_base = ( + state_cache_ptr + + block_numbers_i64 * state_cache_stride0 + + block_offsets * state_cache_stride1 + + head_offset + ) + + combined_mask = mask_pos[:, None] & mask[None, :] + + score = tl.load( + row_base[:, None] + STATE_WIDTH + block[None, :], + mask=combined_mask, + other=float("-inf"), ) score = tl.softmax(score, dim=0) + kv = tl.load( + row_base[:, None] + block[None, :], + mask=combined_mask, + other=0.0, + ) + compressed_kv = tl.sum(kv * score, axis=0) # [TRITON_BLOCK_SIZE] fp32 # ── RMSNorm (fp32 throughout) ────────────────────────────────────── @@ -672,13 +512,6 @@ def _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn( k_cache_ptr, kv_slot_mapping_ptr, kv_cache_block_size, - # current chunk (private-state path) - current_kv_ptr, - current_kv_stride, - current_score_ptr, - current_score_stride, - ape_ptr, - ape_stride, # ── constexprs ── HEAD_SIZE: tl.constexpr, TRITON_BLOCK_SIZE: tl.constexpr, @@ -692,10 +525,6 @@ def _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn( SCALE_DIM: tl.constexpr, # HEAD_SIZE // QUANT_BLOCK = 4 ue8m0 bytes/token KV_BLOCK_STRIDE: tl.constexpr, USE_SOFTWARE_FP8: tl.constexpr, - USE_PRIVATE_STATE: tl.constexpr, - USE_DENSE_PRIVATE_STATE: tl.constexpr, - RING_SIZE: tl.constexpr, - PRIVATE_HISTORY_SIZE: tl.constexpr, ): """Fused compress → RMSNorm → RoPE → MXFP4 quant → store. @@ -724,38 +553,46 @@ def _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn( req_idx = tl.load(token_to_req_indices_ptr + token_idx) # ── Gather state cache entries ──────────────────────────────────── + start = position - (1 + OVERLAP) * COMPRESS_RATIO + 1 + tokens = tl.arange(0, (1 + OVERLAP) * COMPRESS_RATIO) + pos = start + tokens + mask_pos = pos >= 0 + + block_indices = pos // block_size + block_numbers = tl.load( + block_table_ptr + req_idx * block_table_stride + block_indices, + mask=mask_pos, + other=0, + ) + block_offsets = pos % block_size + head_offset = (tokens >= COMPRESS_RATIO).to(tl.int32) * HEAD_SIZE + block = tl.arange(0, TRITON_BLOCK_SIZE) mask = block < HEAD_SIZE - kv, score = _gather_compressor_window( - state_cache_ptr, - state_cache_stride0, - state_cache_stride1, - current_kv_ptr, - current_kv_stride, - current_score_ptr, - current_score_stride, - ape_ptr, - ape_stride, - block_table_ptr, - block_table_stride, - positions_ptr, - position, - token_idx, - req_idx, - block_size, - block, - mask, - HEAD_SIZE, - STATE_WIDTH, - COMPRESS_RATIO, - OVERLAP, - USE_PRIVATE_STATE, - USE_DENSE_PRIVATE_STATE, - RING_SIZE, - PRIVATE_HISTORY_SIZE, + block_numbers_i64 = block_numbers.to(tl.int64) + + row_base = ( + state_cache_ptr + + block_numbers_i64 * state_cache_stride0 + + block_offsets * state_cache_stride1 + + head_offset + ) + + combined_mask = mask_pos[:, None] & mask[None, :] + + score = tl.load( + row_base[:, None] + STATE_WIDTH + block[None, :], + mask=combined_mask, + other=float("-inf"), ) score = tl.softmax(score, dim=0) + kv = tl.load( + row_base[:, None] + block[None, :], + mask=combined_mask, + other=0.0, + ) + compressed_kv = tl.sum(kv * score, axis=0) # [TRITON_BLOCK_SIZE] fp32 # ── RMSNorm (fp32 throughout) ────────────────────────────────────── diff --git a/vllm/models/deepseek_v4/common/ops/save_partial_states.py b/vllm/models/deepseek_v4/common/ops/save_partial_states.py index 07b40cfcea..e3d7d38f45 100644 --- a/vllm/models/deepseek_v4/common/ops/save_partial_states.py +++ b/vllm/models/deepseek_v4/common/ops/save_partial_states.py @@ -45,93 +45,6 @@ def save_partial_states( ) -def save_partial_states_to_ring( - kv: torch.Tensor, - score: torch.Tensor, - ape: torch.Tensor, - positions: torch.Tensor, - state_ring: torch.Tensor, - slot_mapping: torch.Tensor, - token_to_req_indices: torch.Tensor, - seq_lens: torch.Tensor, - state_width: int, - compress_ratio: int, - pdl_kwargs: dict | None = None, -) -> None: - """Keep only the compressor history that can cross a forward boundary.""" - num_slots = slot_mapping.shape[0] - head_size = kv.shape[-1] - _save_partial_states_to_ring_kernel[(num_slots,)]( - kv, - kv.stride(0), - score, - score.stride(0), - ape, - ape.stride(0), - positions, - state_ring, - state_ring.stride(0), - state_ring.stride(1), - slot_mapping, - token_to_req_indices, - seq_lens, - HEAD_SIZE=head_size, - TRITON_BLOCK_SIZE=triton.next_power_of_2(head_size), - STATE_WIDTH=state_width, - RING_SIZE=state_ring.shape[1], - COMPRESS_RATIO=compress_ratio, - **(pdl_kwargs or {}), - ) - - -def stage_partial_states_from_ring( - kv: torch.Tensor, - score: torch.Tensor, - ape: torch.Tensor, - positions: torch.Tensor, - state_ring: torch.Tensor, - dense_state: torch.Tensor, - slot_mapping: torch.Tensor, - state_width: int, - history_size: int, - compress_ratio: int, - pdl_kwargs: dict | None = None, -) -> None: - """Stage private history and the current chunk into one dense state view.""" - num_slots = slot_mapping.shape[0] - copy_block_size = triton.next_power_of_2(2 * state_width) - _copy_ring_history_to_dense_kernel[(history_size,)]( - positions, - state_ring, - state_ring.stride(1), - dense_state, - dense_state.stride(0), - STATE_DIM=2 * state_width, - HISTORY_SIZE=history_size, - RING_SIZE=state_ring.shape[1], - BLOCK_SIZE=copy_block_size, - **(pdl_kwargs or {}), - ) - _save_current_states_to_dense_kernel[(num_slots,)]( - kv, - kv.stride(0), - score, - score.stride(0), - ape, - ape.stride(0), - positions, - dense_state, - dense_state.stride(0), - slot_mapping, - HEAD_SIZE=kv.shape[-1], - BLOCK_SIZE=triton.next_power_of_2(kv.shape[-1]), - STATE_WIDTH=state_width, - HISTORY_SIZE=history_size, - COMPRESS_RATIO=compress_ratio, - **(pdl_kwargs or {}), - ) - - @triton.jit def _save_partial_states_kernel( kv_ptr, @@ -186,123 +99,3 @@ def _save_partial_states_kernel( score + ape, mask=mask, ) - - -@triton.jit -def _save_partial_states_to_ring_kernel( - kv_ptr, - kv_stride, - score_ptr, - score_stride, - ape_ptr, - ape_stride, - positions_ptr, - state_ring_ptr, - state_ring_stride0, - state_ring_stride1, - slot_mapping_ptr, - token_to_req_indices_ptr, - seq_lens_ptr, - HEAD_SIZE: tl.constexpr, - TRITON_BLOCK_SIZE: tl.constexpr, - STATE_WIDTH: tl.constexpr, - RING_SIZE: tl.constexpr, - COMPRESS_RATIO: tl.constexpr, -): - token_idx = tl.program_id(0) - if tl.load(slot_mapping_ptr + token_idx) < 0: - return - - req_idx = tl.load(token_to_req_indices_ptr + token_idx) - position = tl.load(positions_ptr + token_idx) - seq_len = tl.load(seq_lens_ptr + req_idx) - if position < seq_len - RING_SIZE: - return - - ring_row = position % RING_SIZE - base_ptr = ( - state_ring_ptr - + req_idx.to(tl.int64) * state_ring_stride0 - + ring_row.to(tl.int64) * state_ring_stride1 - ) - block = tl.arange(0, TRITON_BLOCK_SIZE) - mask = block < HEAD_SIZE - - kv = tl.load(kv_ptr + token_idx * kv_stride + block, mask=mask).to(tl.float32) - tl.store(base_ptr + block, kv, mask=mask) - - ape_row = position % COMPRESS_RATIO - ape = tl.load(ape_ptr + ape_row * ape_stride + block, mask=mask).to(tl.float32) - score = tl.load(score_ptr + token_idx * score_stride + block, mask=mask).to( - tl.float32 - ) - tl.store(base_ptr + STATE_WIDTH + block, score + ape, mask=mask) - - -@triton.jit -def _copy_ring_history_to_dense_kernel( - positions_ptr, - state_ring_ptr, - state_ring_stride1, - dense_state_ptr, - dense_state_stride0, - STATE_DIM: tl.constexpr, - HISTORY_SIZE: tl.constexpr, - RING_SIZE: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - history_idx = tl.program_id(0) - current_start = tl.load(positions_ptr) - position = current_start - HISTORY_SIZE + history_idx - if position < 0: - return - - block = tl.arange(0, BLOCK_SIZE) - mask = block < STATE_DIM - ring_row = position % RING_SIZE - values = tl.load( - state_ring_ptr + ring_row.to(tl.int64) * state_ring_stride1 + block, - mask=mask, - ) - tl.store( - dense_state_ptr + history_idx * dense_state_stride0 + block, - values, - mask=mask, - ) - - -@triton.jit -def _save_current_states_to_dense_kernel( - kv_ptr, - kv_stride, - score_ptr, - score_stride, - ape_ptr, - ape_stride, - positions_ptr, - dense_state_ptr, - dense_state_stride0, - slot_mapping_ptr, - HEAD_SIZE: tl.constexpr, - BLOCK_SIZE: tl.constexpr, - STATE_WIDTH: tl.constexpr, - HISTORY_SIZE: tl.constexpr, - COMPRESS_RATIO: tl.constexpr, -): - token_idx = tl.program_id(0) - if tl.load(slot_mapping_ptr + token_idx) < 0: - return - - block = tl.arange(0, BLOCK_SIZE) - mask = block < HEAD_SIZE - row_ptr = dense_state_ptr + (HISTORY_SIZE + token_idx) * dense_state_stride0 - kv = tl.load(kv_ptr + token_idx * kv_stride + block, mask=mask).to(tl.float32) - tl.store(row_ptr + block, kv, mask=mask) - - position = tl.load(positions_ptr + token_idx) - ape_row = position % COMPRESS_RATIO - ape = tl.load(ape_ptr + ape_row * ape_stride + block, mask=mask).to(tl.float32) - score = tl.load(score_ptr + token_idx * score_stride + block, mask=mask).to( - tl.float32 - ) - tl.store(row_ptr + STATE_WIDTH + block, score + ape, mask=mask) diff --git a/vllm/models/deepseek_v4/compressor.py b/vllm/models/deepseek_v4/compressor.py index f9d8d43e4c..2738efc293 100644 --- a/vllm/models/deepseek_v4/compressor.py +++ b/vllm/models/deepseek_v4/compressor.py @@ -1,17 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Callable from dataclasses import dataclass from typing import Any, ClassVar, cast import torch from torch import nn -from vllm import envs from vllm.config import CUDAGraphMode, VllmConfig, get_current_vllm_config from vllm.forward_context import get_forward_context -from vllm.logger import init_logger from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import MergedColumnParallelLinear @@ -21,8 +18,6 @@ from vllm.models.deepseek_v4.common.ops.fused_indexer_q import MXFP4_BLOCK_SIZE from vllm.models.deepseek_v4.common.ops.save_partial_states import ( save_partial_states, - save_partial_states_to_ring, - stage_partial_states_from_ring, ) from vllm.platforms import current_platform from vllm.v1.attention.backend import ( @@ -37,28 +32,6 @@ MLAAttentionSpec, SlidingWindowMLASpec, ) -from vllm.v1.worker.workspace import current_workspace_manager - -logger = init_logger(__name__) - - -def _can_use_sm70_private_compressor_state(vllm_config: VllmConfig) -> bool: - kv_transfer_config = vllm_config.kv_transfer_config - spec_config = vllm_config.speculative_config - return ( - envs.VLLM_SM70_DSV4_PRIVATE_COMPRESSOR_STATE - and current_platform.is_cuda() - and current_platform.is_device_capability((7, 0)) - and vllm_config.scheduler_config.max_num_seqs == 1 - and not vllm_config.cache_config.enable_prefix_caching - and vllm_config.parallel_config.pipeline_parallel_size == 1 - # The private gather assumes one contiguous token chain. Parallel - # drafting uses multiple branches and must retain the paged fallback. - and (spec_config is None or not spec_config.parallel_drafting) - and ( - kv_transfer_config is None or not kv_transfer_config.is_kv_transfer_instance - ) - ) def _get_c128_boundary(metadata: CommonAttentionMetadata) -> bool | None: @@ -168,27 +141,12 @@ def __init__( dtype: torch.dtype, compress_ratio: int, prefix: str, - private_ring_size: int = 0, - max_num_reqs: int = 1, ): super().__init__() self.state_dim = state_dim self.dtype = dtype self.prefix = prefix self.kv_cache = torch.tensor([]) - self.private_ring_size = private_ring_size - if private_ring_size > 0: - self.register_buffer( - "private_state", - torch.empty( - (max_num_reqs, private_ring_size, state_dim), - dtype=dtype, - device=current_platform.device_type, - ), - persistent=False, - ) - else: - self.private_state = None compilation_config = get_current_vllm_config().compilation_config if prefix in compilation_config.static_forward_context: raise ValueError(f"Duplicate layer name: {prefix}") @@ -212,9 +170,7 @@ def __init__( else: raise ValueError(f"Invalid compress ratio: {compress_ratio}") - def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: - if self.private_state is not None: - return None + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: return SlidingWindowMLASpec( # only has one vector instead of K + V block_size=self.block_size, num_kv_heads=1, @@ -250,7 +206,6 @@ def __init__( rotate: bool = False, prefix: str = "", k_cache_prefix="", - validity_cache_prefix="", use_fp4_cache: bool = False, ): super().__init__() @@ -260,7 +215,6 @@ def __init__( self.rotate = rotate self.prefix = prefix self.k_cache_prefix = k_cache_prefix - self.validity_cache_prefix = validity_cache_prefix self.use_fp4_cache = use_fp4_cache config = vllm_config.model_config.hf_config @@ -273,17 +227,6 @@ def __init__( self.overlap = compress_ratio == 4 self.coff = 1 + self.overlap - self.use_private_state = _can_use_sm70_private_compressor_state(vllm_config) - num_speculative_tokens = ( - vllm_config.speculative_config.num_speculative_tokens - if vllm_config.speculative_config is not None - else 0 - ) - private_ring_size = ( - self.coff * compress_ratio + max(1, num_speculative_tokens + 1) - if self.use_private_state - else 0 - ) state_dtype = torch.float32 self.ape = nn.Parameter( @@ -304,7 +247,6 @@ 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( @@ -312,15 +254,7 @@ def __init__( dtype=state_dtype, compress_ratio=compress_ratio, prefix=f"{prefix}.state_cache", - private_ring_size=private_ring_size, - max_num_reqs=self.max_num_reqs, ) - if self.use_private_state: - logger.info_once( - "Using SM70 private DeepSeek V4 compressor state rings " - "(max_num_seqs=1, prefix cache disabled); paged compressor " - "state allocation is disabled." - ) # Save reference to static_forward_context for forward-time KV cache lookup. # get_current_vllm_config() is only available during __init__, not forward. @@ -369,33 +303,19 @@ def forward( if not isinstance(attn_metadata, dict): return - seq_lens = None - if self.use_private_state: - assert self.validity_cache_prefix, ( - "Private compressor state requires an SWA metadata source." - ) - validity_metadata = cast(Any, attn_metadata[self.validity_cache_prefix]) - token_to_req_indices = validity_metadata.token_to_req_indices - slot_mapping = validity_metadata.slot_mapping - block_table = validity_metadata.block_table - block_size = validity_metadata.block_size - seq_lens = validity_metadata.seq_lens - state_cache = self.state_cache.private_state - assert state_cache is not None - assert token_to_req_indices is not None and seq_lens is not None - state_width = self.state_cache.state_dim // 2 - state_metadata = None - else: - state_metadata = cast( - CompressorMetadata, attn_metadata[self.state_cache.prefix] - ) - token_to_req_indices = state_metadata.token_to_req_indices - slot_mapping = state_metadata.slot_mapping - block_table = state_metadata.block_table - block_size = state_metadata.block_size - state_cache = self.state_cache.kv_cache - state_width = state_cache.shape[-1] // 2 + state_metadata = cast( + CompressorMetadata, attn_metadata[self.state_cache.prefix] + ) + token_to_req_indices = state_metadata.token_to_req_indices + slot_mapping = state_metadata.slot_mapping num_actual = slot_mapping.shape[0] + block_table = state_metadata.block_table + block_size = state_metadata.block_size + + # [num_blocks, block_size, kv_dim+score_dim], where kv_dim == score_dim + state_cache = self.state_cache.kv_cache + # kv_state stored in first half, score_state stored in second half + state_width = state_cache.shape[-1] // 2 pdl_kwargs = ( {} if current_platform.is_rocm() or current_platform.is_xpu() @@ -408,48 +328,21 @@ def forward( # GEMM; state_cache from this kernel) but neither emits/waits on PDL # grid dependency primitives, so launch_pdl=True caused a # read-after-write race and non-deterministic output. - if not self.use_private_state: - save_partial_states( - kv=kv, - score=score, - ape=self.ape, - positions=positions, - state_cache=state_cache, - slot_mapping=slot_mapping, - block_size=block_size, - state_width=state_width, - compress_ratio=self.compress_ratio, - pdl_kwargs=pdl_kwargs, - ) - - use_dense_private_state = self.use_private_state and self.compress_ratio == 128 - compression_state_cache = state_cache - private_history_size = self.coff * self.compress_ratio - if use_dense_private_state: - (dense_state,) = current_workspace_manager().get_simultaneous( - ( - (private_history_size + num_actual, 2 * state_width), - torch.float32, - ) - ) - stage_partial_states_from_ring( - kv=kv, - score=score, - ape=self.ape, - positions=positions, - state_ring=state_cache, - dense_state=dense_state, - slot_mapping=slot_mapping, - state_width=state_width, - history_size=private_history_size, - compress_ratio=self.compress_ratio, - pdl_kwargs=pdl_kwargs, - ) - compression_state_cache = dense_state.unsqueeze(0) + save_partial_states( + kv=kv, + score=score, + ape=self.ape, + positions=positions, + state_cache=state_cache, + slot_mapping=slot_mapping, + block_size=block_size, + state_width=state_width, + compress_ratio=self.compress_ratio, + pdl_kwargs=pdl_kwargs, + ) if ( current_platform.is_cuda() - and not self.use_private_state and self.head_dim == 512 and self.compress_ratio == 128 and forward_context.cudagraph_runtime_mode != CUDAGraphMode.FULL @@ -468,7 +361,6 @@ def forward( k_cache_metadata = cast(Any, attn_metadata[self.k_cache_prefix]) kv_cache = self._static_forward_context[self.k_cache_prefix].kv_cache - compress_norm_rope_store_fn: Callable[..., None] if current_platform.is_cuda() and not current_platform.is_device_capability( (7, 0) ): @@ -490,8 +382,8 @@ def forward( # Always use a triton kernel. compress_norm_rope_store_fn = compress_norm_rope_store_triton - compress_kwargs = dict( - state_cache=compression_state_cache, + compress_norm_rope_store_fn( + state_cache=state_cache, num_actual=num_actual, token_to_req_indices=token_to_req_indices, positions=positions, @@ -514,31 +406,3 @@ def forward( token_stride=self._token_stride, scale_dim=self._scale_dim, ) - if compress_norm_rope_store_fn is compress_norm_rope_store_triton: - compress_kwargs.update( - current_kv=kv, - current_score=score, - ape=self.ape, - use_private_state=self.use_private_state - and not use_dense_private_state, - use_dense_private_state=use_dense_private_state, - ring_size=self.state_cache.private_ring_size or 1, - private_history_size=private_history_size, - ) - compress_norm_rope_store_fn(**compress_kwargs) - - if self.use_private_state: - assert seq_lens is not None - save_partial_states_to_ring( - kv=kv, - score=score, - ape=self.ape, - positions=positions, - state_ring=state_cache, - slot_mapping=slot_mapping, - token_to_req_indices=token_to_req_indices, - seq_lens=seq_lens, - state_width=state_width, - compress_ratio=self.compress_ratio, - pdl_kwargs=pdl_kwargs, - ) diff --git a/vllm/models/deepseek_v4/nvidia/dspark.py b/vllm/models/deepseek_v4/nvidia/dspark.py index 84f711947c..da3319124a 100644 --- a/vllm/models/deepseek_v4/nvidia/dspark.py +++ b/vllm/models/deepseek_v4/nvidia/dspark.py @@ -15,7 +15,6 @@ import torch import torch.nn as nn -from vllm import envs from vllm.config import VllmConfig, get_current_vllm_config from vllm.distributed import ( get_tensor_model_parallel_rank, @@ -72,40 +71,8 @@ def bias(self, markov_embed: torch.Tensor) -> torch.Tensor: return self.markov_w2(markov_embed) -class DSparkConfidenceHead(nn.Module): - """Replicated FP32 conditional-acceptance predictor.""" - - def __init__(self, input_size: int, prefix: str) -> None: - super().__init__() - # The checkpoint tensor is BF16, but the official implementation keeps - # this projection and its input in FP32 for calibration accuracy. - self.proj = ReplicatedLinear( - input_size, - 1, - bias=False, - params_dtype=torch.float32, - quant_config=None, - return_bias=False, - prefix=maybe_prefix(prefix, "proj"), - ) - - def forward( - self, - hidden_states: torch.Tensor, - markov_embeds: torch.Tensor, - ) -> torch.Tensor: - features = torch.cat((hidden_states.float(), markov_embeds.float()), dim=-1) - return self.proj(features).squeeze(-1) - - class DSparkDeepseekV4Model(nn.Module): - def __init__( - self, - *, - vllm_config: VllmConfig, - prefix: str = "", - use_confidence_head: bool = False, - ) -> None: + 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 @@ -180,14 +147,6 @@ def __init__( config.dspark_markov_rank, prefix=maybe_prefix(prefix, "markov_head"), ) - self.confidence_head = ( - DSparkConfidenceHead( - config.hidden_size + config.dspark_markov_rank, - prefix=maybe_prefix(prefix, "confidence_head"), - ) - if use_confidence_head - else None - ) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) @@ -304,16 +263,8 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: assert vllm_config.speculative_config is not None self.draft_model_config = vllm_config.speculative_config.draft_model_config self.config = self.draft_model_config.hf_config - speculative_config = vllm_config.speculative_config - self._confidence_head_required = bool( - speculative_config.dspark_confidence_threshold > 0.0 - or envs.VLLM_SPEC_DUMP_ALIGNMENT - ) - self._confidence_head_loaded = False self.model = DSparkDeepseekV4Model( - vllm_config=vllm_config, - prefix=maybe_prefix(prefix, "model"), - use_confidence_head=self._confidence_head_required, + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) self.lm_head = ParallelLMHead( self.config.vocab_size, @@ -360,21 +311,6 @@ def markov_embed(self, token_ids: torch.Tensor) -> torch.Tensor: def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor: return self.model.markov_head.bias(markov_embed) - def confidence_logits( - self, - hidden_states: torch.Tensor, - markov_embeds: torch.Tensor, - ) -> torch.Tensor: - if not self._confidence_head_loaded: - raise RuntimeError( - "DSpark confidence logits were requested, but the draft " - "checkpoint did not provide confidence_head.proj.weight." - ) - confidence_head = self.model.confidence_head - if confidence_head is None: - raise RuntimeError("DSpark confidence head was not initialized.") - return confidence_head(hidden_states, markov_embeds) - def map_draft_to_target(self, draft_ids: torch.Tensor) -> torch.Tensor: return draft_ids @@ -416,11 +352,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: name = self._remap_dspark_name(checkpoint_name) if name is None: continue - # Ordinary serving does not instantiate the optional diagnostic / - # scheduler head, so ignore its checkpoint tensor without adding - # an unused parameter to the runtime model. - if name.startswith("model.confidence_head.") and name not in params_dict: - continue if name.endswith(".scale"): suffix = ( expert_scale_suffix @@ -485,13 +416,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: for layer in self.model.layers: layer.ffn.finalize_mega_moe_weights() - confidence_weight_name = "model.confidence_head.proj.weight" - self._confidence_head_loaded = confidence_weight_name in loaded_params - if self._confidence_head_required and not self._confidence_head_loaded: - raise ValueError( - "DSpark confidence collection requires the draft checkpoint " - f"parameter {confidence_weight_name!r}." - ) logger.info_once("DSpark draft model loaded: %d params", len(loaded_params)) return loaded_params @@ -502,13 +426,14 @@ def _remap_dspark_name(name: str) -> str | None: return None stage = int(match.group(1)) rest = match.group(2) + if rest.startswith("confidence_head."): + return None head_prefixes = ( "norm.", "hc_head_fn", "hc_head_base", "hc_head_scale", "markov_head.", - "confidence_head.", ) if rest.startswith(("main_proj.", "main_norm.")) or rest.startswith( head_prefixes diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 750b3179ea..fb165aa730 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -718,27 +718,6 @@ def __init__( rope_parameters["mscale_all_dim"] = 0 # Disable mscale rope_parameters["is_deepseek_v4"] = True rope_parameters["rope_dim"] = self.rope_head_dim - if self.compress_ratio <= 1: - # The sliding-window-only layers get base rope_theta above, but the - # reference also turns YaRN itself off for them: it passes - # original_seq_len=0, and precompute_freqs_cis then skips the - # correction range entirely (inference/model.py:481-486, :227). - # Swapping only theta leaves the interpolation on, which shifts the - # low-frequency pairs by up to ~2 degrees across the 128-token - # window. factor=1 makes the interpolation an exact identity - # (inv_freq_interpolation == inv_freq_extrapolation) and mscale - # collapses to 1.0, so this is plain rope without swapping the - # class -- DeepseekV4ScalingRotaryEmbedding still rotates the last - # rotary_dim and keeps the cache in fp32, which the fused kernels - # require. Copy first: the dict is config.rope_parameters itself - # and every later layer reads it back. - rope_parameters = dict(rope_parameters) - rope_parameters["factor"] = 1 - # Cache length is original_max_position * factor, so it has to - # carry the full position range now that factor is 1. - rope_parameters["original_max_position_embeddings"] = ( - self.max_position_embeddings - ) self.rotary_emb = get_rope( self.head_dim, max_position=self.max_position_embeddings, @@ -801,10 +780,6 @@ def __init__( quant_config=quant_config, prefix=prefix, ) - if self.indexer is not None: - self.indexer.compressor.validity_cache_prefix = ( - self.mla_attn.swa_cache_layer.prefix - ) def forward( self, diff --git a/vllm/models/deepseek_v4/sm70/gemv.py b/vllm/models/deepseek_v4/sm70/gemv.py index 20378173e7..d1374fbb76 100644 --- a/vllm/models/deepseek_v4/sm70/gemv.py +++ b/vllm/models/deepseek_v4/sm70/gemv.py @@ -16,11 +16,6 @@ _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 @@ -41,240 +36,15 @@ def _sm70_dsv4_fp16_gemv_kernel( tl.store(out_ptr + row, acc) -@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( +def can_use_sm70_dsv4_fp16_gemv( x: torch.Tensor, weight: torch.Tensor, output_dtype: torch.dtype, ) -> bool: return ( - current_platform.is_cuda() + envs.VLLM_SM70_DSV4_FP16_GEMV + and 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) @@ -289,57 +59,11 @@ def _has_sm70_dsv4_gemv_contract( ) -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/indexer.py b/vllm/models/deepseek_v4/sm70/indexer.py index 291f150cc3..369dc2e147 100644 --- a/vllm/models/deepseek_v4/sm70/indexer.py +++ b/vllm/models/deepseek_v4/sm70/indexer.py @@ -3,100 +3,16 @@ """DeepSeek V4 C4 indexer fallback using FP16 HMMA on SM70.""" -import os - import torch from vllm.models.deepseek_v4.common.ops.fp8_software import ( fp8_e4m3fn_bits_to_fp32, ) -from vllm.platforms import current_platform from vllm.triton_utils import tl, triton -from vllm.v1.worker.workspace import current_workspace_manager _INDEX_HEAD_DIM = 128 _INDEX_CACHE_BYTES = _INDEX_HEAD_DIM + 4 -# Fuse dequant+scoring into one paged kernel instead of materialising -# [rows, max_seq_len, 128] fp16 and running a dense bmm. Set to 0 to get the -# old gather+bmm path back for an A/B. Only affects decode; prefill scores a -# contiguous chunk whose width is the real chunk length already. -_FUSED_DECODE_LOGITS = os.getenv("VLLM_SM70_INDEXER_FUSED_LOGITS", "1") == "1" - -# Score with the model's actual indexer function, -# I[t, s] = sum_h weights[t, h] * relu(q[t, h] . k[s]), -# instead of the cheaper (sum_h weights[t, h] * q[t, h]) . k[s]. Those are -# equal only without the relu, and the relu is not optional: it is what the -# checkpoint was trained with (reference `Indexer.forward`, -# inference/model.py:427, and the in-repo `fp8_mqa_logits` reference at -# v1/attention/ops/rocm_aiter_mla_sparse.py:510 -- every non-SM70 backend -# applies it inside its logits kernel). Dropping it costs 64x fewer flops and -# ranks the wrong keys, which only shows up once top-k is actually selective. -# Set 0 to restore the factored form for an A/B. -_RELU_LOGITS = os.getenv("VLLM_SM70_INDEXER_RELU", "1") == "1" - -# Key-axis splits are chosen so the launch is roughly this many blocks, which -# keeps all 80 SMs busy at long context without paying for thousands of -# no-op programs at short context. -_LOGITS_TARGET_BLOCKS = 512 -_LOGITS_BLOCK_N = 32 -# The relu path multiplies q against k with an MMA instead of a broadcast -# multiply, so the key tile has to be at least one 16x16x16 HMMA fragment -# wide; 64 keeps the [64 heads, 128] x [128, N] shape efficient. -_RELU_BLOCK_N = 64 -_RELU_BLOCK_M = 32 - -# Route prefill scoring through cuBLAS instead of the Triton MMA. Set 0 for an -# A/B against the Triton kernel above (which stays as the reference shape and -# is what the decode path still uses, since decode gathers from a paged cache). -_PREFILL_CUBLAS = os.getenv("VLLM_SM70_INDEXER_PREFILL_CUBLAS", "1") == "1" -# Cap (MiB) on the [tokens*heads, key_tile] fp16 score tile. At the default -# 2048-token prefill chunk and 64 heads this is 128 KiB per key, so the tile -# is what keeps a long-context chunk from asking for gigabytes in one go. -_PREFILL_TILE_MB = int(os.getenv("VLLM_SM70_INDEXER_PREFILL_TILE_MB", "192")) -_EPILOGUE_BLOCK_K = 256 -_EPILOGUE_BLOCK_H = 8 - -# At long context the paged Triton decode kernel rescans the same FP8 MQA keys -# once for every verifier row and head. For a bounded one-request verifier -# shape, gather the keys once and let cuBLAS score all rows/heads together. -# FP32 output is deliberate: an FP16 score intermediate produced rare top-k -# set changes at 16K/64K compressed tokens. Short and generic shapes retain the -# fused paged kernel, whose launch overhead is lower there. -_DECODE_CUBLAS = os.getenv("VLLM_SM70_INDEXER_DECODE_CUBLAS", "0") == "1" -_DECODE_CUBLAS_MIN_KEYS = int( - os.getenv("VLLM_SM70_INDEXER_DECODE_CUBLAS_MIN_KEYS", "1024") -) -_DECODE_CUBLAS_MAX_ROWS = 8 - - -@triton.jit -def _index_cache_ptrs( - cache_ptr, - physical_block, - pos_in_block, - block_stride, - cache_block_size, - head_dim: tl.constexpr, -): - """Value and scale pointers for one token of the paged indexer cache. - - The layout is block-major, NOT token-interleaved: within a block come all - `cache_block_size` rows of FP8 values, and only then the FP32 scales, one - per token (csrc/libtorch_stable/cache_kernels.cu, the - `indexer_k_quant_and_cache` store and the `cp_gather_indexer_k_quant_cache` - load both address it this way). Reading a token as - `[128 values][4 scale]` at `pos * cache.stride(1)` therefore lands on the - wrong bytes for every token but the first of each block, and the "scale" - it picks up is really four FP8 value bytes. - """ - base = cache_ptr + physical_block.to(tl.int64) * block_stride - value_ptr = base + pos_in_block.to(tl.int64) * head_dim - scale_ptr = (base + cache_block_size * head_dim + pos_in_block * 4).to( - tl.pointer_type(tl.float32) - ) - return value_ptr, scale_ptr - @triton.jit def _weighted_query_kernel( @@ -163,6 +79,7 @@ def _dequant_paged_index_k_kernel( seq_lens_ptr, out_ptr, cache_stride0, + cache_stride1, block_table_stride0, out_stride0, out_stride1, @@ -183,13 +100,10 @@ def _dequant_paged_index_k_kernel( mask=valid, other=0, ) - token_ptr, scale_ptr = _index_cache_ptrs( - cache_ptr, - physical_block, - pos_in_block, - cache_stride0, - cache_block_size, - head_dim, + token_ptr = ( + cache_ptr + + physical_block.to(tl.int64) * cache_stride0 + + pos_in_block * cache_stride1 ) packed = tl.load( token_ptr[:, None] + dim_offsets[None, :], @@ -197,6 +111,7 @@ def _dequant_paged_index_k_kernel( other=0, ) fp8 = fp8_e4m3fn_bits_to_fp32(packed) + scale_ptr = (token_ptr + head_dim).to(tl.pointer_type(tl.float32)) scale = tl.load(scale_ptr, mask=valid, other=0.0).to(tl.float32) dequant = fp8 * scale[:, None] tl.store( @@ -209,285 +124,6 @@ def _dequant_paged_index_k_kernel( ) -@triton.jit -def _gather_paged_index_k_unscaled_kernel( - cache_ptr, - block_table_ptr, - seq_lens_ptr, - out_k_ptr, - out_scale_ptr, - cache_stride0, - out_k_stride0, - cache_block_size, - max_seq_len, - num_rows, - head_dim: tl.constexpr, - MAX_ROWS: tl.constexpr, - BLOCK_N: tl.constexpr, -): - """Gather one request's live paged keys once for all verifier rows. - - ``max_seq_len`` is a CUDA-graph bucket and can be much wider than the - currently allocated request. Derive the live bound on device from the - row lengths so graph replay never dereferences a ``-1`` block-table tail. - Rows outside the live bound are zeroed to keep the workspace deterministic; - the downstream top-k still applies each row's own sequence length. - """ - key_offsets = tl.program_id(0) * BLOCK_N + tl.arange(0, BLOCK_N) - dim_offsets = tl.arange(0, head_dim) - row_offsets = tl.arange(0, MAX_ROWS) - row_lens = tl.load( - seq_lens_ptr + row_offsets, - mask=row_offsets < num_rows, - other=0, - ) - live_seq_len = tl.max(row_lens, axis=0) - in_bucket = key_offsets < max_seq_len - valid = in_bucket & (key_offsets < live_seq_len) - - block_in_seq = key_offsets // cache_block_size - pos_in_block = key_offsets % cache_block_size - physical_block = tl.load( - block_table_ptr + block_in_seq, - mask=valid, - other=0, - ) - token_ptr, scale_ptr = _index_cache_ptrs( - cache_ptr, - physical_block, - pos_in_block, - cache_stride0, - cache_block_size, - head_dim, - ) - packed = tl.load( - token_ptr[:, None] + dim_offsets[None, :], - mask=valid[:, None], - other=0, - ) - scales = tl.load(scale_ptr, mask=valid, other=0.0).to(tl.float32) - tl.store( - out_k_ptr + key_offsets[:, None] * out_k_stride0 + dim_offsets[None, :], - fp8_e4m3fn_bits_to_fp32(packed).to(tl.float16), - mask=in_bucket[:, None], - ) - tl.store(out_scale_ptr + key_offsets, scales, mask=in_bucket) - - -@triton.jit -def _paged_index_logits_kernel( - q_ptr, - cache_ptr, - block_table_ptr, - seq_lens_ptr, - out_ptr, - q_stride0, - cache_stride0, - block_table_stride0, - out_stride0, - cache_block_size, - head_dim: tl.constexpr, - BLOCK_N: tl.constexpr, - NUM_SPLITS: tl.constexpr, -): - """Index logits straight out of the paged FP8 cache, no gather buffer. - - `_dequant_paged_index_k_kernel` + bmm has a grid and a GEMM shape derived - from `max_seq_len`, and under a captured decode CUDA Graph that is frozen - at `max_model_len` forever. Here the grid is static -- (rows, NUM_SPLITS), - sized from `max_seq_len` on the host -- and only each program's trip count - follows the live `seq_len`, so one captured graph serves any context while - a short context costs almost nothing. - """ - row_idx = tl.program_id(0) - split_id = tl.program_id(1) - dim_offsets = tl.arange(0, head_dim) - q = tl.load(q_ptr + row_idx * q_stride0 + dim_offsets).to(tl.float32) - - seq_len = tl.load(seq_lens_ptr + row_idx) - chunks = (seq_len + BLOCK_N - 1) // BLOCK_N - chunks_per_split = (chunks + NUM_SPLITS - 1) // NUM_SPLITS - lo = split_id * chunks_per_split * BLOCK_N - hi = tl.minimum(lo + chunks_per_split * BLOCK_N, seq_len) - - for start in range(lo, hi, BLOCK_N): - key_offsets = start + tl.arange(0, BLOCK_N) - valid = key_offsets < seq_len - block_in_seq = key_offsets // cache_block_size - pos_in_block = key_offsets % cache_block_size - physical_block = tl.load( - block_table_ptr + row_idx * block_table_stride0 + block_in_seq, - mask=valid, - other=0, - ) - token_ptr, scale_ptr = _index_cache_ptrs( - cache_ptr, - physical_block, - pos_in_block, - cache_stride0, - cache_block_size, - head_dim, - ) - packed = tl.load( - token_ptr[:, None] + dim_offsets[None, :], - mask=valid[:, None], - other=0, - ) - fp8 = fp8_e4m3fn_bits_to_fp32(packed) - scale = tl.load(scale_ptr, mask=valid, other=0.0).to(tl.float32) - # Per-token scale factors out of the dot product. - logits = tl.sum(fp8 * q[None, :], axis=1) * scale - tl.store(out_ptr + row_idx * out_stride0 + key_offsets, logits, mask=valid) - - -@triton.jit -def _contiguous_index_logits_relu_kernel( - q_ptr, - w_ptr, - k_ptr, - k_scale_ptr, - out_ptr, - q_stride0, - q_stride1, - w_stride0, - k_stride0, - out_stride0, - num_q, - num_k, - num_heads, - head_dim: tl.constexpr, - BLOCK_M: tl.constexpr, - BLOCK_N: tl.constexpr, -): - """Prefill index scores, per head, with the relu inside the head sum. - - Reference shape for the scoring, and the A/B fallback for the cuBLAS path - (`VLLM_SM70_INDEXER_PREFILL_CUBLAS=0`). It is not the fast path: Triton's - `tl.dot` does not reach Volta's tensor cores, so this tops out near 3.3 - TFLOP/s no matter how the tile is shaped -- folding heads into the M - dimension to get a [256, 128] x [128, 64] MMA instead of a per-head - [32, 128] x [128, 64] one measured 3.34 vs 3.39 TFLOP/s, i.e. nothing. - cuBLAS does the same shape at 54. - - The key tile is loaded once and reused by every head (the indexer is MQA). - FP8 keys are widened to FP16 *unscaled* -- e4m3 maxes out at 448, which - fp16 holds exactly -- and the per-key scale is applied to the finished - accumulator. That keeps `k * scale` out of fp16 entirely, where a large - ue8m0 scale could otherwise overflow. - """ - pid_m = tl.program_id(0) - pid_n = tl.program_id(1) - offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - dim_offsets = tl.arange(0, head_dim) - m_valid = offs_m < num_q - n_valid = offs_n < num_k - - packed = tl.load( - k_ptr + offs_n[:, None] * k_stride0 + dim_offsets[None, :], - mask=n_valid[:, None], - other=0, - ) - k_t = tl.trans(fp8_e4m3fn_bits_to_fp32(packed).to(tl.float16)) - scale = tl.load(k_scale_ptr + offs_n, mask=n_valid, other=0.0).to(tl.float32) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - for head in range(num_heads): - q = tl.load( - q_ptr - + offs_m[:, None] * q_stride0 - + head * q_stride1 - + dim_offsets[None, :], - mask=m_valid[:, None], - other=0.0, - ) - w = tl.load(w_ptr + offs_m * w_stride0 + head, mask=m_valid, other=0.0) - acc += tl.maximum(tl.dot(q, k_t), 0.0) * w[:, None].to(tl.float32) - - tl.store( - out_ptr + offs_m[:, None] * out_stride0 + offs_n[None, :], - acc * scale[None, :], - mask=m_valid[:, None] & n_valid[None, :], - ) - - -@triton.jit -def _paged_index_logits_relu_kernel( - q_ptr, - w_ptr, - cache_ptr, - block_table_ptr, - seq_lens_ptr, - out_ptr, - q_stride0, - q_stride1, - w_stride0, - cache_stride0, - block_table_stride0, - out_stride0, - cache_block_size, - head_dim: tl.constexpr, - NUM_HEADS: tl.constexpr, - BLOCK_N: tl.constexpr, - NUM_SPLITS: tl.constexpr, -): - """Decode counterpart of `_contiguous_index_logits_relu_kernel`. - - Same static-grid property as `_paged_index_logits_kernel`: the grid is - (rows, NUM_SPLITS) and only the trip count follows the live `seq_len`, so - one captured CUDA graph serves any context length. - """ - row_idx = tl.program_id(0) - split_id = tl.program_id(1) - head_offsets = tl.arange(0, NUM_HEADS) - dim_offsets = tl.arange(0, head_dim) - q = tl.load( - q_ptr - + row_idx * q_stride0 - + head_offsets[:, None] * q_stride1 - + dim_offsets[None, :] - ) - w = tl.load(w_ptr + row_idx * w_stride0 + head_offsets).to(tl.float32) - - seq_len = tl.load(seq_lens_ptr + row_idx) - chunks = (seq_len + BLOCK_N - 1) // BLOCK_N - chunks_per_split = (chunks + NUM_SPLITS - 1) // NUM_SPLITS - lo = split_id * chunks_per_split * BLOCK_N - hi = tl.minimum(lo + chunks_per_split * BLOCK_N, seq_len) - - for start in range(lo, hi, BLOCK_N): - key_offsets = start + tl.arange(0, BLOCK_N) - valid = key_offsets < seq_len - block_in_seq = key_offsets // cache_block_size - pos_in_block = key_offsets % cache_block_size - physical_block = tl.load( - block_table_ptr + row_idx * block_table_stride0 + block_in_seq, - mask=valid, - other=0, - ) - token_ptr, scale_ptr = _index_cache_ptrs( - cache_ptr, - physical_block, - pos_in_block, - cache_stride0, - cache_block_size, - head_dim, - ) - packed = tl.load( - token_ptr[:, None] + dim_offsets[None, :], - mask=valid[:, None], - other=0, - ) - k_t = tl.trans(fp8_e4m3fn_bits_to_fp32(packed).to(tl.float16)) - scale = tl.load(scale_ptr, mask=valid, other=0.0).to(tl.float32) - # relu(q_h . k) is per head; only then does the weighted head sum - # collapse it to one logit. The per-key scale is positive, so pulling - # it out past the relu and the sum is exact. - scores = tl.maximum(tl.dot(q, k_t), 0.0) * w[:, None] - logits = tl.sum(scores, axis=0) * scale - tl.store(out_ptr + row_idx * out_stride0 + key_offsets, logits, mask=valid) - - def _combine_index_queries(q: torch.Tensor, weights: torch.Tensor) -> torch.Tensor: assert q.dtype == torch.float16 and q.ndim == 3 assert q.shape[-1] == _INDEX_HEAD_DIM @@ -510,134 +146,6 @@ def _combine_index_queries(q: torch.Tensor, weights: torch.Tensor) -> torch.Tens return out -@triton.jit -def _relu_weight_headsum_kernel( - s_ptr, - w_ptr, - scale_ptr, - out_ptr, - s_stride0, - s_stride1, - w_stride0, - out_stride0, - num_k, - num_heads: tl.constexpr, - BLOCK_K: tl.constexpr, - BLOCK_H: tl.constexpr, -): - """relu -> weight -> head sum -> key scale, in a single pass over scores. - - Epilogue for the cuBLAS prefill path. Doing it as torch ops instead costs - three extra round trips through the [tokens, heads, keys] score tile, which - is the largest tensor in the whole indexer; fusing them measured 13 -> 26 - TFLOP/s end to end. - """ - token = tl.program_id(0) - k_base = tl.program_id(1) * BLOCK_K - k_offsets = k_base + tl.arange(0, BLOCK_K) - k_valid = k_offsets < num_k - - acc = tl.zeros((BLOCK_K,), dtype=tl.float32) - for head in range(0, num_heads, BLOCK_H): - h_offsets = head + tl.arange(0, BLOCK_H) - s = tl.load( - s_ptr - + token * s_stride0 - + h_offsets[:, None] * s_stride1 - + k_offsets[None, :], - mask=k_valid[None, :], - other=0.0, - ).to(tl.float32) - w = tl.load(w_ptr + token * w_stride0 + h_offsets).to(tl.float32) - acc += tl.sum(tl.maximum(s, 0.0) * w[:, None], axis=0) - - scale = tl.load(scale_ptr + k_offsets, mask=k_valid, other=0.0).to(tl.float32) - tl.store(out_ptr + token * out_stride0 + k_offsets, acc * scale, mask=k_valid) - - -@triton.jit -def _dequant_index_k_unscaled_kernel( - k_ptr, - out_ptr, - k_stride0, - out_stride0, - head_dim: tl.constexpr, -): - """Widen FP8 index keys to FP16 without applying the per-key scale. - - e4m3 tops out at 448 and the scale is folded into the epilogue instead, so - the GEMM never sees `k * scale` -- a large ue8m0 scale would otherwise - overflow fp16 on the way in. - """ - row_idx = tl.program_id(0) - offsets = tl.arange(0, head_dim) - values = tl.load( - k_ptr.to(tl.pointer_type(tl.uint8)) + row_idx * k_stride0 + offsets - ) - tl.store( - out_ptr + row_idx * out_stride0 + offsets, - fp8_e4m3fn_bits_to_fp32(values).to(tl.float16), - ) - - -def _prefill_logits_cublas(q, k_quant, k_scales, weights): - """Prefill index scores through cuBLAS, with the relu inside the head sum. - - Triton's `tl.dot` does not reach Volta's tensor cores: on the exact shape - this needs it measures 3.3 TFLOP/s, against 54 TFLOP/s for cuBLAS fp16 and - 13 TFLOP/s for plain fp32 CUDA cores. The relu blocks the one-GEMM - factorisation, but it does not block cuBLAS -- stacking heads into the M - dimension turns the scoring into a single [tokens*heads, 128] x [128, keys] - GEMM whose output is reduced over heads afterwards. Keys are tiled so that - intermediate stays bounded; it is the largest allocation in the indexer. - """ - num_q, num_heads, head_dim = q.shape - num_k = k_quant.shape[0] - out = torch.empty((num_q, num_k), dtype=torch.float32, device=q.device) - - k_fp16 = torch.empty((num_k, head_dim), dtype=torch.float16, device=q.device) - _dequant_index_k_unscaled_kernel[(num_k,)]( - k_quant.view(torch.uint8), - k_fp16, - k_quant.stride(0), - k_fp16.stride(0), - head_dim=head_dim, - num_warps=4, - ) - - q2 = q.reshape(num_q * num_heads, head_dim) - if not q2.is_contiguous(): - q2 = q2.contiguous() - - rows = num_q * num_heads - budget = max(1, _PREFILL_TILE_MB * 2**20 // max(1, rows * 2)) - k_tile = max(_RELU_BLOCK_N, min(num_k, budget // _RELU_BLOCK_N * _RELU_BLOCK_N)) - - for start in range(0, num_k, k_tile): - stop = min(start + k_tile, num_k) - width = stop - start - # fp16 out with fp32 accumulate: the epilogue reduces in fp32 and only - # the top-k *set* is consumed downstream, so the 5e-4 rounding here is - # far below the gap between neighbouring index scores. - scores = torch.mm(q2, k_fp16[start:stop].t()) - _relu_weight_headsum_kernel[(num_q, triton.cdiv(width, _EPILOGUE_BLOCK_K))]( - scores, - weights, - k_scales[start:stop], - out[:, start:stop], - scores.stride(0) * num_heads, - scores.stride(0), - weights.stride(0), - out.stride(0), - width, - num_heads=num_heads, - BLOCK_K=_EPILOGUE_BLOCK_K, - BLOCK_H=_EPILOGUE_BLOCK_H, - num_warps=4, - ) - return out - - def sm70_indexer_prefill_logits( q: torch.Tensor, k_quant: torch.Tensor, @@ -650,37 +158,6 @@ def sm70_indexer_prefill_logits( k_scales = k_scale_storage.view(torch.float32).reshape(-1) assert k_scales.shape[0] == k_quant.shape[0] - if _RELU_LOGITS and _PREFILL_CUBLAS: - return _prefill_logits_cublas(q, k_quant, k_scales, weights) - - if _RELU_LOGITS: - num_q, num_heads, _ = q.shape - num_k = k_quant.shape[0] - out = torch.empty((num_q, num_k), dtype=torch.float32, device=q.device) - block_m, block_n = _RELU_BLOCK_M, _RELU_BLOCK_N - _contiguous_index_logits_relu_kernel[ - (triton.cdiv(num_q, block_m), triton.cdiv(num_k, block_n)) - ]( - q, - weights, - k_quant.view(torch.uint8), - k_scales, - out, - q.stride(0), - q.stride(1), - weights.stride(0), - k_quant.stride(0), - out.stride(0), - num_q, - num_k, - num_heads, - head_dim=_INDEX_HEAD_DIM, - BLOCK_M=block_m, - BLOCK_N=block_n, - num_warps=8, - ) - return out - weighted_q = _combine_index_queries(q, weights) k_fp16 = torch.empty(k_quant.shape, dtype=torch.float16, device=k_quant.device) _dequant_contiguous_index_k_kernel[(k_quant.shape[0],)]( @@ -695,68 +172,6 @@ def sm70_indexer_prefill_logits( return torch.mm(weighted_q, k_fp16.t(), out_dtype=torch.float32) -def _decode_logits_cublas( - q: torch.Tensor, - cache: torch.Tensor, - weights: torch.Tensor, - flat_lens: torch.Tensor, - block_table: torch.Tensor, - max_seq_len: int, -) -> torch.Tensor: - """Score a one-request decode by sharing one paged gather across rows.""" - total_rows, num_heads, head_dim = q.shape - workspace = current_workspace_manager() - gathered_k, k_scales, scores = workspace.get_simultaneous( - ((max_seq_len, head_dim), torch.float16), - ((max_seq_len,), torch.float32), - ((total_rows * num_heads, max_seq_len), torch.float32), - ) - block_n = 16 - _gather_paged_index_k_unscaled_kernel[(triton.cdiv(max_seq_len, block_n),)]( - cache, - block_table, - flat_lens, - gathered_k, - k_scales, - cache.stride(0), - gathered_k.stride(0), - cache.shape[1], - max_seq_len, - total_rows, - head_dim=head_dim, - MAX_ROWS=_DECODE_CUBLAS_MAX_ROWS, - BLOCK_N=block_n, - num_warps=4, - ) - - q2 = q.reshape(total_rows * num_heads, head_dim) - torch.mm( - q2, - gathered_k.t(), - out=scores, - out_dtype=torch.float32, - ) - out = torch.empty((total_rows, max_seq_len), dtype=torch.float32, device=q.device) - _relu_weight_headsum_kernel[ - (total_rows, triton.cdiv(max_seq_len, _EPILOGUE_BLOCK_K)) - ]( - scores, - weights, - k_scales, - out, - scores.stride(0) * num_heads, - scores.stride(0), - weights.stride(0), - out.stride(0), - max_seq_len, - num_heads=num_heads, - BLOCK_K=_EPILOGUE_BLOCK_K, - BLOCK_H=_EPILOGUE_BLOCK_H, - num_warps=4, - ) - return out - - def sm70_indexer_decode_logits( q: torch.Tensor, cache: torch.Tensor, @@ -768,142 +183,33 @@ def sm70_indexer_decode_logits( """Gather paged FP8 index keys and compute batched decode scores.""" assert cache.dtype == torch.uint8 and cache.ndim == 3 assert cache.shape[-1] >= _INDEX_CACHE_BYTES - # The relu path scores per head, so it needs q as [rows, heads, dim]; the - # factored path collapses the head axis up front. - weighted_q = q if _RELU_LOGITS else _combine_index_queries(q, weights) + weighted_q = _combine_index_queries(q, weights) - single_request = block_table.shape[0] == 1 - native_rows = seq_lens.ndim == 2 - if native_rows: + if seq_lens.ndim == 2: next_n = seq_lens.shape[1] flat_lens = seq_lens.reshape(-1).to(torch.int32) + block_table = block_table.repeat_interleave(next_n, dim=0) else: flat_lens = seq_lens.reshape(-1).to(torch.int32) assert flat_lens.shape[0] == weighted_q.shape[0] - - max_seq_len = max(1, int(max_seq_len)) - total_rows = weighted_q.shape[0] - - if ( - _RELU_LOGITS - and _DECODE_CUBLAS - and current_platform.is_cuda() - and current_platform.is_device_capability((7, 0)) - and max_seq_len >= _DECODE_CUBLAS_MIN_KEYS - and single_request - and (native_rows or total_rows == 1) - and 0 < total_rows <= _DECODE_CUBLAS_MAX_ROWS - and q.ndim == 3 - and q.dtype == torch.float16 - and q.shape[2] == _INDEX_HEAD_DIM - and q.shape[1] > 0 - and q.shape[1] % _EPILOGUE_BLOCK_H == 0 - and weights.ndim == 2 - and weights.shape == q.shape[:2] - and weights.dtype == torch.float32 - and cache.ndim == 3 - and cache.dtype == torch.uint8 - and cache.shape[1] > 0 - and cache.shape[2] == _INDEX_CACHE_BYTES - and cache.stride(1) == _INDEX_CACHE_BYTES - and cache.stride(2) == 1 - and block_table.ndim == 2 - and block_table.dtype == torch.int32 - and flat_lens.dtype == torch.int32 - and triton.cdiv(max_seq_len, cache.shape[1]) <= block_table.shape[1] - and q.device - == cache.device - == weights.device - == flat_lens.device - == block_table.device - and q.is_contiguous() - and weights.is_contiguous() - and cache.is_contiguous() - and flat_lens.is_contiguous() - and block_table.is_contiguous() - ): - return _decode_logits_cublas( - q, - cache, - weights, - flat_lens, - block_table, - max_seq_len, - ) - - if native_rows: - block_table = block_table.repeat_interleave(next_n, dim=0) assert block_table.shape[0] == weighted_q.shape[0] - if _RELU_LOGITS: - out = torch.empty( - (total_rows, max_seq_len), dtype=torch.float32, device=q.device - ) - max_chunks = triton.cdiv(max_seq_len, _RELU_BLOCK_N) - splits = max(1, min(max_chunks, _LOGITS_TARGET_BLOCKS // total_rows)) - _paged_index_logits_relu_kernel[(total_rows, splits)]( - q, - weights, - cache, - block_table, - flat_lens, - out, - q.stride(0), - q.stride(1), - weights.stride(0), - cache.stride(0), - block_table.stride(0), - out.stride(0), - cache.shape[1], - head_dim=_INDEX_HEAD_DIM, - NUM_HEADS=q.shape[1], - BLOCK_N=_RELU_BLOCK_N, - NUM_SPLITS=splits, - num_warps=8, - ) - # Positions >= seq_len are left uninitialised; every consumer masks by - # seq_lens (see the note on the factored path below). - return out - - if _FUSED_DECODE_LOGITS: - out = torch.empty( - (total_rows, max_seq_len), dtype=torch.float32, device=q.device - ) - max_chunks = triton.cdiv(max_seq_len, _LOGITS_BLOCK_N) - splits = max(1, min(max_chunks, _LOGITS_TARGET_BLOCKS // total_rows)) - _paged_index_logits_kernel[(total_rows, splits)]( - weighted_q, - cache, - block_table, - flat_lens, - out, - weighted_q.stride(0), - cache.stride(0), - block_table.stride(0), - out.stride(0), - cache.shape[1], - head_dim=_INDEX_HEAD_DIM, - BLOCK_N=_LOGITS_BLOCK_N, - NUM_SPLITS=splits, - num_warps=4, - ) - # Positions >= seq_len are left uninitialised, exactly as the - # single-chunk gather path leaves them; every consumer masks by - # seq_lens. - return out - - block_n = 16 + max_seq_len = max(1, int(max_seq_len)) gathered_k = torch.empty( - (total_rows, max_seq_len, _INDEX_HEAD_DIM), + (weighted_q.shape[0], max_seq_len, _INDEX_HEAD_DIM), dtype=torch.float16, device=q.device, ) - _dequant_paged_index_k_kernel[(total_rows, triton.cdiv(max_seq_len, block_n))]( + block_n = 16 + _dequant_paged_index_k_kernel[ + (weighted_q.shape[0], triton.cdiv(max_seq_len, block_n)) + ]( cache, block_table, flat_lens, gathered_k, cache.stride(0), + cache.stride(1), block_table.stride(0), gathered_k.stride(0), gathered_k.stride(1), 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 d520cf833f..ac88ab5d3d 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 @@ -5,13 +5,9 @@ import torch -import vllm.envs as envs from vllm.models.deepseek_v4.common.ops.cache_utils import ( quantize_and_insert_k_cache, ) -from vllm.models.deepseek_v4.common.ops.fp8_software import ( - fp32_to_fp8_e4m3fn_bits, -) from vllm.triton_utils import tl, triton _HEAD_DIM = 512 @@ -97,111 +93,6 @@ def _sm70_qnorm_rope_kernel( ) -@triton.jit -def _sm70_qnorm_rope_parallel_kv_insert_kernel( - q_ptr, - kv_ptr, - cache_ptr, - slot_mapping_ptr, - positions_ptr, - cos_sin_cache_ptr, - cache_stride0, - eps: tl.constexpr, - num_tokens, - num_heads: tl.constexpr, - cache_block_size: tl.constexpr, - HEAD_DIM: tl.constexpr, - ROPE_DIM: tl.constexpr, - NOPE_DIM: tl.constexpr, - HALF_ROPE: tl.constexpr, - TOKEN_DATA_SIZE: tl.constexpr, -): - """Fuse TP4 Q transform with eight exact packed-KV writer CTAs.""" - token_idx = tl.program_id(0) - worker_idx = tl.program_id(1) - if token_idx >= num_tokens: - return - - if worker_idx < num_heads: - q_position = tl.load(positions_ptr + token_idx).to(tl.int64) - q_pair_idx = tl.arange(0, HALF_ROPE) - q_cos = tl.load(cos_sin_cache_ptr + q_position * ROPE_DIM + q_pair_idx).to( - tl.float32 - ) - q_sin = tl.load( - cos_sin_cache_ptr + q_position * ROPE_DIM + HALF_ROPE + q_pair_idx - ).to(tl.float32) - q_base = q_ptr + token_idx * num_heads * HEAD_DIM + worker_idx * HEAD_DIM - q_offsets = tl.arange(0, HEAD_DIM) - q_values = tl.load(q_base + q_offsets).to(tl.float32) - q_rrms = tl.rsqrt(tl.sum(q_values * q_values, axis=0) / HEAD_DIM + eps) - q_values *= q_rrms - tl.store( - q_base + q_offsets, - q_values.to(q_ptr.type.element_ty), - mask=q_offsets < NOPE_DIM, - ) - - q_even_offsets = NOPE_DIM + q_pair_idx * 2 - q_odd_offsets = q_even_offsets + 1 - q_even = tl.load(q_base + q_even_offsets).to(tl.float32) * q_rrms - q_odd = tl.load(q_base + q_odd_offsets).to(tl.float32) * q_rrms - tl.store( - q_base + q_even_offsets, - (q_even * q_cos - q_odd * q_sin).to(q_ptr.type.element_ty), - ) - tl.store( - q_base + q_odd_offsets, - (q_even * q_sin + q_odd * q_cos).to(q_ptr.type.element_ty), - ) - else: - part_idx = worker_idx - num_heads - slot_idx = tl.load(slot_mapping_ptr + token_idx) - if slot_idx < 0: - return - - block_idx = slot_idx // cache_block_size - pos_in_block = slot_idx % cache_block_size - cache_block = cache_ptr + block_idx.to(tl.int64) * cache_stride0 - token_data = cache_block + pos_in_block * TOKEN_DATA_SIZE - token_scales = ( - cache_block + cache_block_size * TOKEN_DATA_SIZE + pos_in_block * 8 - ) - kv_row = kv_ptr + token_idx * HEAD_DIM - kv_offsets = tl.arange(0, 64) - if part_idx < 7: - input_offsets = part_idx * 64 + kv_offsets - kv_values = tl.load(kv_row + input_offsets) - block_max = tl.maximum(tl.max(tl.abs(kv_values), axis=0), 1.0e-4) - exponent = tl.ceil(tl.log2(block_max / 448.0)) - scale = tl.exp2(exponent) - scaled = tl.clamp(kv_values / scale, -448.0, 448.0) - packed = fp32_to_fp8_e4m3fn_bits(scaled.to(tl.float32)) - tl.store(token_data + input_offsets, packed) - encoded = tl.maximum(tl.minimum(exponent + 127.0, 255.0), 0.0) - tl.store(token_scales + part_idx, encoded.to(tl.uint8)) - else: - tl.store(token_scales + 7, tl.zeros((), dtype=tl.uint8)) - kv_position = tl.load(positions_ptr + token_idx).to(tl.int64) - kv_pair_idx = tl.arange(0, HALF_ROPE) - kv_cos = tl.load( - cos_sin_cache_ptr + kv_position * ROPE_DIM + kv_pair_idx - ).to(tl.float32) - kv_sin = tl.load( - cos_sin_cache_ptr + kv_position * ROPE_DIM + HALF_ROPE + kv_pair_idx - ).to(tl.float32) - kv_even = tl.load(kv_row + NOPE_DIM + kv_pair_idx * 2).to(tl.float32) - kv_odd = tl.load(kv_row + NOPE_DIM + kv_pair_idx * 2 + 1).to(tl.float32) - rotated_even = (kv_even * kv_cos - kv_odd * kv_sin).to(tl.float16) - rotated_odd = (kv_even * kv_sin + kv_odd * kv_cos).to(tl.float16) - bf16_out = (token_data + NOPE_DIM).to(tl.pointer_type(tl.bfloat16)) - tl.store(bf16_out + kv_pair_idx * 2, rotated_even.to(tl.bfloat16)) - tl.store( - bf16_out + kv_pair_idx * 2 + 1, - rotated_odd.to(tl.bfloat16), - ) - - def sm70_qnorm_rope_kv_fp8_insert( q: torch.Tensor, kv: torch.Tensor, @@ -221,29 +112,6 @@ def sm70_qnorm_rope_kv_fp8_insert( assert q.is_contiguous() and kv.is_contiguous() num_tokens, num_heads, _ = q.shape - if envs.VLLM_SM70_DSV4_QNORM_KV_FUSED_TP4 and num_tokens == 1 and num_heads == 16: - cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1) - _sm70_qnorm_rope_parallel_kv_insert_kernel[(num_tokens, num_heads + 8)]( - q, - kv, - cache_2d, - slot_mapping, - positions, - cos_sin_cache, - cache_2d.stride(0), - eps, - num_tokens, - num_heads=num_heads, - cache_block_size=block_size, - HEAD_DIM=_HEAD_DIM, - ROPE_DIM=_ROPE_DIM, - NOPE_DIM=_NOPE_DIM, - HALF_ROPE=_HALF_ROPE, - TOKEN_DATA_SIZE=576, - num_warps=4, - ) - return q - kv_roped = torch.empty_like(kv) _sm70_qnorm_rope_kernel[(num_tokens, num_heads + 1)]( q, diff --git a/vllm/models/deepseek_v4/sm70/sparse.py b/vllm/models/deepseek_v4/sm70/sparse.py index 1eec99a606..f507372031 100644 --- a/vllm/models/deepseek_v4/sm70/sparse.py +++ b/vllm/models/deepseek_v4/sm70/sparse.py @@ -40,11 +40,6 @@ logger = init_logger(__name__) -def _qk_dsplit_block_h(num_heads: int) -> int: - """Use one CTA head group for the exact TP4 decode shape.""" - return 16 if num_heads == 16 else 8 - - class DeepseekV4SM70SparseBackend(DeepseekV4FlashMLASparseBackend): supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16] @@ -263,7 +258,6 @@ def _forward_decode( sm70_sparse_attention_paged_fp8_splitk_qk_dsplit( partial_qk=partial_qk, partial_probs=partial_probs, - stage1_block_h=_qk_dsplit_block_h(q.shape[1]), **common_kwargs, ) else: diff --git a/vllm/tool_parsers/qwen3coder_tool_parser.py b/vllm/tool_parsers/qwen3coder_tool_parser.py index 7a3a63ad53..2b21f8f6ae 100644 --- a/vllm/tool_parsers/qwen3coder_tool_parser.py +++ b/vllm/tool_parsers/qwen3coder_tool_parser.py @@ -358,7 +358,10 @@ def extract_tool_calls_streaming( if pos != -1 ] tool_start = min(tool_start_candidates) if tool_start_candidates else -1 - if self.tool_call_start_token_id in delta_token_ids or tool_start != -1: + if ( + self.tool_call_start_token_id in delta_token_ids + or tool_start != -1 + ): self.is_tool_call_started = True # Return any content before the tool call if tool_start > len(previous_text): diff --git a/vllm/transformers_utils/processors/deepseek_ocr.py b/vllm/transformers_utils/processors/deepseek_ocr.py index 62972ad242..68a2b1aaaa 100644 --- a/vllm/transformers_utils/processors/deepseek_ocr.py +++ b/vllm/transformers_utils/processors/deepseek_ocr.py @@ -3,7 +3,7 @@ # adapted from https://github.com/deepseek-ai/DeepSeek-OCR/blob/main/DeepSeek-OCR-master/DeepSeek-OCR-vllm/process/image_process.py # and https://github.com/deepseek-ai/DeepSeek-OCR-2/blob/main/DeepSeek-OCR2-master/DeepSeek-OCR2-vllm/process/image_process.py import math -from typing import Literal, TypedDict +from typing import Literal import torch import torchvision.transforms as T @@ -11,30 +11,18 @@ from transformers import BatchFeature, LlamaTokenizerFast from transformers.processing_utils import ProcessorMixin - -# Official resolution modes, verbatim from the model's config.py: -# https://github.com/deepseek-ai/DeepSeek-OCR/blob/8cf003d38821fa1b19c73da3bd1b0dc262ea8136/DeepSeek-OCR-master/DeepSeek-OCR-vllm/config.py#L1-L6 -# Selectable per request via the ``image_mode`` mm_processor_kwarg; the -# default stays the Gundam configuration below (byte-compatible with the -# previous hardcoded constants). -class ResolutionMode(TypedDict): - base_size: int - image_size: int - crop_mode: bool - - -RESOLUTION_MODES: dict[str, ResolutionMode] = { - "tiny": {"base_size": 512, "image_size": 512, "crop_mode": False}, - "small": {"base_size": 640, "image_size": 640, "crop_mode": False}, - "base": {"base_size": 1024, "image_size": 1024, "crop_mode": False}, - "large": {"base_size": 1280, "image_size": 1280, "crop_mode": False}, - "gundam": {"base_size": 1024, "image_size": 640, "crop_mode": True}, -} - +# TODO(Isotr0py): change modes for variants +# see: https://github.com/deepseek-ai/DeepSeek-OCR/blob/8cf003d38821fa1b19c73da3bd1b0dc262ea8136/DeepSeek-OCR-master/DeepSeek-OCR-vllm/config.py#L1-L6 +# Tiny: base_size = 512, image_size = 512, crop_mode = False +# Small: base_size = 640, image_size = 640, crop_mode = False +# Base: base_size = 1024, image_size = 1024, crop_mode = False +# Large: base_size = 1280, image_size = 1280, crop_mode = False +# Gundam: base_size = 1024, image_size = 640, crop_mode = True BASE_SIZE = 1024 IMAGE_SIZE = 640 CROP_MODE = True +# TODO(Isotr0py): Expose as mm_kwargs MIN_CROPS = 2 MAX_CROPS = 6 # max:9; If your GPU memory is small, it is recommended to set it to 6. @@ -129,53 +117,6 @@ def dynamic_preprocess( return processed_images, target_aspect_ratio -def count_image_tokens_for( - *, - image_width: int, - image_height: int, - base_size: int, - image_size: int, - cropping: bool, - min_crops: int = MIN_CROPS, - max_crops: int = MAX_CROPS, - strategy: Literal["v1", "v2"] = "v1", - patch_size: int = 16, - downsample_ratio: int = 4, -) -> int: - """Placeholder-token count mirroring ``tokenize_with_images``.""" - if image_width <= image_size and image_height <= image_size: - crop_ratio = (1, 1) - elif cropping: - crop_ratio = count_tiles( - image_width, - image_height, - min_num=min_crops, - max_num=max_crops, - image_size=image_size, - ) - else: - crop_ratio = (1, 1) - num_width_tiles, num_height_tiles = crop_ratio - - num_queries = math.ceil((image_size // patch_size) / downsample_ratio) - num_queries_base = math.ceil((base_size // patch_size) / downsample_ratio) - - num_tokens_base = ( - (num_queries_base * (num_queries_base + 1)) - if strategy == "v1" - else num_queries_base * num_queries_base - ) - total = num_tokens_base + 1 - if num_width_tiles > 1 or num_height_tiles > 1: - num_tokens_per_row = ( - num_queries * num_width_tiles + 1 - if strategy == "v1" - else num_queries * num_width_tiles - ) - total += num_tokens_per_row * (num_queries * num_height_tiles) - return total - - class ImageTransform: def __init__( self, @@ -219,63 +160,22 @@ def __init__( ignore_id: int = -100, image_size: int = IMAGE_SIZE, base_size: int = BASE_SIZE, - crop_mode: bool = CROP_MODE, - image_mode: str | None = None, - min_crops: int = MIN_CROPS, - max_crops: int = MAX_CROPS, strategy: Literal["v1", "v2"] = "v1", **kwargs, ): - if image_mode is not None: - # The named mode is authoritative over any individually passed - # base_size/image_size/crop_mode (the vLLM processing info always - # forwards the defaults, so a raise-on-conflict would never let a - # mode through). - if not isinstance(image_mode, str): - raise ValueError( - f"image_mode must be a string, got {type(image_mode).__name__}." - ) - try: - mode = RESOLUTION_MODES[image_mode] - except KeyError: - raise ValueError( - f"Unknown image_mode {image_mode!r}; valid modes: " - f"{sorted(RESOLUTION_MODES)}" - ) from None - base_size = mode["base_size"] - image_size = mode["image_size"] - crop_mode = mode["crop_mode"] - if not ( - type(min_crops) is int - and type(max_crops) is int - and 1 <= min_crops <= max_crops <= 9 - ): - raise ValueError( - f"Invalid crop bounds: min_crops={min_crops!r}, " - f"max_crops={max_crops!r} (need 1 <= min <= max <= 9)." - ) - if type(crop_mode) is not bool: - raise ValueError(f"crop_mode must be a bool, got {crop_mode!r}.") - if patch_size <= 0 or downsample_ratio <= 0: - raise ValueError("patch_size and downsample_ratio must be positive.") self.image_size = image_size self.base_size = base_size - self.crop_mode = crop_mode - self.image_mode = image_mode - self.min_crops = min_crops - self.max_crops = max_crops # image token calculation strategy for # Deepseek-OCR and Deepseek-OCR-2 self.strategy = strategy - if strategy not in ("v1", "v2"): - raise ValueError("Only 'v1' and 'v2' strategies are supported.") + assert strategy in ["v1", "v2"], "Only 'v1' and 'v2' strategies are supported." - self.patch_size = patch_size + self.patch_size = 16 self.image_mean = image_mean self.image_std = image_std self.normalize = normalize - self.downsample_ratio = downsample_ratio + self.downsample_ratio = 4 self.image_transform = ImageTransform( mean=image_mean, std=image_std, normalize=normalize @@ -303,28 +203,6 @@ def __init__( **kwargs, ) - def count_image_tokens(self, *, image_width: int, image_height: int) -> int: - """Number of placeholder tokens for one image. - - This mirrors ``tokenize_with_images`` exactly (same instance - parameters and same crop arithmetic) so the prompt-update counting - can never diverge from what processing - produces — a mismatch here is the "N multimodal tokens vs M - placeholders" crash class. - """ - return count_image_tokens_for( - image_width=image_width, - image_height=image_height, - base_size=self.base_size, - image_size=self.image_size, - cropping=self.crop_mode, - min_crops=self.min_crops, - max_crops=self.max_crops, - strategy=self.strategy, - patch_size=self.patch_size, - downsample_ratio=self.downsample_ratio, - ) - @property def bos_id(self): return self.tokenizer.bos_token_id @@ -352,7 +230,7 @@ def process_one( self, prompt: str, images: list[Image.Image], - crop_mode: bool | None = None, + crop_mode: bool = CROP_MODE, ): """ @@ -375,8 +253,6 @@ def process_one( ) sft_format = prompt - if crop_mode is None: - crop_mode = self.crop_mode ( input_ids, @@ -412,7 +288,7 @@ def __call__( *, prompt: str, images: list[Image.Image], - crop_mode: bool | None = None, + crop_mode: bool = CROP_MODE, **kwargs, ): prepare = self.process_one( @@ -456,10 +332,7 @@ def tokenize_with_images( crop_ratio = [1, 1] elif cropping: images_crop_raw, crop_ratio = dynamic_preprocess( - image, - min_num=self.min_crops, - max_num=self.max_crops, - image_size=self.image_size, + image, image_size=self.image_size ) else: crop_ratio = [1, 1] diff --git a/vllm/utils/cache.py b/vllm/utils/cache.py index 5f7647dce6..4338983f90 100644 --- a/vllm/utils/cache.py +++ b/vllm/utils/cache.py @@ -118,8 +118,10 @@ def stat(self, *, delta: bool = False) -> CacheInfo: return info def touch(self, key: _K) -> None: - if key in self: + try: self._LRUCache__order.move_to_end(key) # type: ignore + except KeyError: + self._LRUCache__order[key] = None # type: ignore @overload def get(self, key: _K, /) -> _V | None: ... diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 4425fae48b..b1cb946d33 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -388,9 +388,6 @@ class CommonAttentionMetadata: causal: bool = True - cudagraph_graph_variant: int | None = None - """Optional semantic variant for CUDA graph capture dummy runs.""" - # Needed by FastPrefillAttentionBuilder logits_indices_padded: torch.Tensor | None = None num_logits_indices: int | None = None @@ -419,14 +416,6 @@ class CommonAttentionMetadata: decode rows (assumes every draft was accepted). Not safe for kernels that need exact per-row context lengths on decode rows.""" - prefix_anchor_lens: torch.Tensor | None = None - """(batch_size,) per-request prefix length (prompt token count) for - prefix-anchored sliding-window attention. Tokens with logical index below - this stay globally visible; later (generated) tokens additionally see a - fixed sliding window. None disables the mechanism. The attention backend - copies this into its own persistent buffer and reads the window size from - the cache group's ``PrefixAnchoredSWASpec``.""" - # WARNING: Deprecated fields. Will be removed in a future release (v0.15.0) _seq_lens_cpu: torch.Tensor | None = None _num_computed_tokens_cpu: torch.Tensor | None = None @@ -527,7 +516,6 @@ def unpadded( block_table_tensor=self.block_table_tensor[:num_actual_reqs], slot_mapping=self.slot_mapping[:num_actual_tokens], causal=self.causal, - cudagraph_graph_variant=self.cudagraph_graph_variant, logits_indices_padded=self.logits_indices_padded, num_logits_indices=self.num_logits_indices, encoder_seq_lens=maybe_slice_reqs(self.encoder_seq_lens), @@ -535,7 +523,6 @@ def unpadded( dcp_local_seq_lens=maybe_slice_reqs(self.dcp_local_seq_lens), dcp_local_seq_lens_cpu=maybe_slice_reqs(self.dcp_local_seq_lens_cpu), is_prefilling=maybe_slice_reqs(self.is_prefilling), - prefix_anchor_lens=maybe_slice_reqs(self.prefix_anchor_lens), ) diff --git a/vllm/v1/attention/backends/ddtree_branch_triton.py b/vllm/v1/attention/backends/ddtree_branch_triton.py index 22a2c9f046..21e38930bf 100644 --- a/vllm/v1/attention/backends/ddtree_branch_triton.py +++ b/vllm/v1/attention/backends/ddtree_branch_triton.py @@ -77,7 +77,10 @@ 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) @@ -98,10 +101,18 @@ 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 @@ -115,7 +126,9 @@ 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 c58258410f..3117c01660 100644 --- a/vllm/v1/attention/backends/fa_utils.py +++ b/vllm/v1/attention/backends/fa_utils.py @@ -17,7 +17,6 @@ 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..1ecbbdb308 100644 --- a/vllm/v1/attention/backends/flash_attn_v100.py +++ b/vllm/v1/attention/backends/flash_attn_v100.py @@ -16,17 +16,14 @@ import os import time from collections.abc import Callable -from dataclasses import dataclass from functools import partial from typing import cast import torch import vllm.envs as envs -from vllm.forward_context import CUDAGRAPH_VARIANT_LONG_CONTEXT from vllm.logger import init_logger from vllm.platforms import current_platform -from vllm.triton_utils import tl, triton from vllm.v1.attention.backend import AttentionCGSupport, AttentionType from vllm.v1.attention.backends.triton_attn import ( TritonAttentionBackend, @@ -34,396 +31,10 @@ TritonAttentionMetadata, TritonAttentionMetadataBuilder, ) -from vllm.v1.kv_cache_interface import PrefixAnchoredSWASpec -from vllm.v1.worker.gpu.spec_decode import uses_dflash_selector_engine logger = init_logger(__name__) -@triton.jit -def _sm70_prepare_smallq_decode_metadata_kernel( - out_block_table_ptr, - out_seq_lens_ptr, - out_query_start_loc_ptr, - block_table_ptr, - seq_lens_ptr, - query_start_loc_ptr, - block_table_stride, - out_block_table_stride, - num_reqs, - num_query_tokens, - real_num_query_tokens, - block_cols, - REQ_BLOCK: tl.constexpr, - BLOCK_COLS: tl.constexpr, -): - token_idx = tl.program_id(0) - - # Find the request that owns this flattened query token. Equal trailing - # boundaries are CUDA-graph padding; clamping maps them to the final padded - # request, matching repeat_query_lens[-1] += padding_tokens. - req_offsets = tl.arange(0, REQ_BLOCK) - req_mask = req_offsets < num_reqs - query_ends = tl.load( - query_start_loc_ptr + req_offsets + 1, - mask=req_mask, - other=0x7FFFFFFF, - ) - req_idx = tl.sum((token_idx >= query_ends).to(tl.int32), axis=0) - req_idx = tl.minimum(req_idx, num_reqs - 1) - - query_start = tl.load(query_start_loc_ptr + req_idx) - query_end = tl.load(query_start_loc_ptr + req_idx + 1) - query_len = query_end - query_start - seq_len = tl.load(seq_lens_ptr + req_idx) - effective_seq_len = tl.maximum(seq_len, query_len) - decode_seq_len = effective_seq_len - query_len + token_idx - query_start + 1 - is_padding = token_idx >= real_num_query_tokens - tl.store( - out_seq_lens_ptr + token_idx, - tl.where(is_padding, 0, decode_seq_len), - ) - - cols = tl.arange(0, BLOCK_COLS) - col_mask = cols < block_cols - block_ids = tl.load( - block_table_ptr + req_idx * block_table_stride + cols, - mask=col_mask, - other=0, - ) - block_ids = tl.maximum(block_ids, 0) - block_ids = tl.where(is_padding, 0, block_ids) - tl.store( - out_block_table_ptr + token_idx * out_block_table_stride + cols, - block_ids, - mask=col_mask, - ) - - # The same launch also refreshes the graph-stable query boundaries. - if token_idx == 0: - boundary_offsets = tl.arange(0, REQ_BLOCK) - boundary_mask = boundary_offsets < num_reqs + 1 - boundaries = tl.load( - query_start_loc_ptr + boundary_offsets, - mask=boundary_mask, - other=0, - ) - tl.store( - out_query_start_loc_ptr + boundary_offsets, - boundaries, - mask=boundary_mask, - ) - - -def _sm70_prepare_smallq_decode_metadata( - out_block_table: torch.Tensor, - out_seq_lens: torch.Tensor, - out_query_start_loc: torch.Tensor, - block_table: torch.Tensor, - seq_lens: torch.Tensor, - query_start_loc: torch.Tensor, - *, - num_reqs: int, - num_query_tokens: int, - real_num_query_tokens: int, -) -> None: - """Materialize persistent Flash-V100 verifier metadata in one launch.""" - if num_reqs <= 0 or num_query_tokens <= 0: - raise ValueError("small-query metadata requires positive request/token counts") - block_cols = int(block_table.shape[1]) - if block_cols <= 0: - raise ValueError("small-query block table must have at least one column") - if out_block_table.shape[0] < num_query_tokens: - raise ValueError("small-query output block table is too small") - if out_seq_lens.numel() < num_query_tokens: - raise ValueError("small-query output sequence buffer is too small") - if out_query_start_loc.numel() < num_reqs + 1: - raise ValueError("small-query output boundary buffer is too small") - - _sm70_prepare_smallq_decode_metadata_kernel[(num_query_tokens,)]( - out_block_table, - out_seq_lens, - out_query_start_loc, - block_table, - seq_lens, - query_start_loc, - block_table.stride(0), - out_block_table.stride(0), - num_reqs, - num_query_tokens, - real_num_query_tokens, - block_cols, - REQ_BLOCK=triton.next_power_of_2(num_reqs + 1), - BLOCK_COLS=triton.next_power_of_2(block_cols), - num_warps=1, - ) - - -@triton.jit -def _load_sm70_smallq_i32_ptr(ptrs, group_id): - ptr = tl.load(ptrs + group_id) - return tl.cast(ptr, tl.pointer_type(tl.int32)) - - -@triton.jit -def _sm70_prepare_grouped_smallq_decode_metadata_kernel( - out_block_table_ptrs, - out_seq_lens_ptrs, - out_query_start_loc_ptrs, - block_table_ptrs, - block_table_strides, - out_block_table_strides, - block_col_counts, - seq_lens_ptr, - query_start_loc_ptr, - num_reqs, - real_num_query_tokens, - REQ_BLOCK: tl.constexpr, - BLOCK_COLS: tl.constexpr, -): - """Materialize every target full-attention group's verifier metadata.""" - group_id = tl.program_id(0) - token_idx = tl.program_id(1) - out_block_table = _load_sm70_smallq_i32_ptr(out_block_table_ptrs, group_id) - out_seq_lens = _load_sm70_smallq_i32_ptr(out_seq_lens_ptrs, group_id) - out_query_start_loc = _load_sm70_smallq_i32_ptr(out_query_start_loc_ptrs, group_id) - block_table = _load_sm70_smallq_i32_ptr(block_table_ptrs, group_id) - block_table_stride = tl.load(block_table_strides + group_id) - out_block_table_stride = tl.load(out_block_table_strides + group_id) - block_cols = tl.load(block_col_counts + group_id) - - req_offsets = tl.arange(0, REQ_BLOCK) - req_mask = req_offsets < num_reqs - query_ends = tl.load( - query_start_loc_ptr + req_offsets + 1, - mask=req_mask, - other=0x7FFFFFFF, - ) - req_idx = tl.sum((token_idx >= query_ends).to(tl.int32), axis=0) - req_idx = tl.minimum(req_idx, num_reqs - 1) - - query_start = tl.load(query_start_loc_ptr + req_idx) - query_end = tl.load(query_start_loc_ptr + req_idx + 1) - query_len = query_end - query_start - seq_len = tl.load(seq_lens_ptr + req_idx) - effective_seq_len = tl.maximum(seq_len, query_len) - decode_seq_len = effective_seq_len - query_len + token_idx - query_start + 1 - is_padding = token_idx >= real_num_query_tokens - tl.store( - out_seq_lens + token_idx, - tl.where(is_padding, 0, decode_seq_len), - ) - - cols = tl.arange(0, BLOCK_COLS) - col_mask = cols < block_cols - block_ids = tl.load( - block_table + req_idx * block_table_stride + cols, - mask=col_mask, - other=0, - ) - block_ids = tl.maximum(block_ids, 0) - block_ids = tl.where(is_padding, 0, block_ids) - tl.store( - out_block_table + token_idx * out_block_table_stride + cols, - block_ids, - mask=col_mask, - ) - - if token_idx == 0: - boundary_offsets = tl.arange(0, REQ_BLOCK) - boundary_mask = boundary_offsets < num_reqs + 1 - boundaries = tl.load( - query_start_loc_ptr + boundary_offsets, - mask=boundary_mask, - other=0, - ) - tl.store( - out_query_start_loc + boundary_offsets, - boundaries, - mask=boundary_mask, - ) - - -@dataclass -class DFlash2SmallQGroupDescriptor: - """Persistent pointer tables for grouped Flash-V100 verifier metadata.""" - - key: tuple[object, ...] - block_table_ptrs: torch.Tensor - out_block_table_ptrs: torch.Tensor - out_seq_lens_ptrs: torch.Tensor - out_query_start_loc_ptrs: torch.Tensor - block_table_strides: torch.Tensor - out_block_table_strides: torch.Tensor - block_col_counts: torch.Tensor - - -@dataclass(frozen=True) -class DFlash2SmallQPreparedMetadata: - """Graph-stable buffers already refreshed by the grouped launch.""" - - builder_id: int - num_reqs: int - num_query_tokens: int - max_seq_len_hint: int - workspace_seq_capacity_hint: int - partition_size_hint: int | None = None - - -def _sm70_prepare_grouped_smallq_decode_metadata( - out_block_tables: list[torch.Tensor], - out_seq_lens: list[torch.Tensor], - out_query_start_locs: list[torch.Tensor], - block_tables: list[torch.Tensor], - seq_lens: torch.Tensor, - query_start_loc: torch.Tensor, - *, - num_reqs: int, - num_query_tokens: int, - real_num_query_tokens: int, - descriptor: DFlash2SmallQGroupDescriptor | None = None, -) -> DFlash2SmallQGroupDescriptor: - """Refresh N full-attention cache groups in one Triton launch.""" - num_groups = len(block_tables) - if num_groups <= 0: - raise ValueError("grouped small-query metadata requires at least one group") - if not ( - len(out_block_tables) - == len(out_seq_lens) - == len(out_query_start_locs) - == num_groups - ): - raise ValueError("grouped small-query metadata lists must have equal lengths") - if num_reqs <= 0 or num_query_tokens <= 0: - raise ValueError("grouped small-query metadata requires positive sizes") - if not 0 <= real_num_query_tokens <= num_query_tokens: - raise ValueError("real query token count exceeds grouped launch size") - - block_col_counts = [int(table.shape[1]) for table in block_tables] - if any(cols <= 0 for cols in block_col_counts): - raise ValueError("grouped small-query block tables cannot be empty") - max_block_cols = max(block_col_counts) - device = block_tables[0].device - if ( - seq_lens.device != device - or query_start_loc.device != device - or seq_lens.dtype != torch.int32 - or query_start_loc.dtype != torch.int32 - or not seq_lens.is_contiguous() - or not query_start_loc.is_contiguous() - or seq_lens.numel() < num_reqs - or query_start_loc.numel() < num_reqs + 1 - ): - raise ValueError("grouped small-query input metadata contract mismatch") - key: tuple[object, ...] = ( - device.type, - device.index, - tuple(block_col_counts), - tuple(table.data_ptr() for table in block_tables), - tuple(table.data_ptr() for table in out_block_tables), - tuple(tensor.data_ptr() for tensor in out_seq_lens), - tuple(tensor.data_ptr() for tensor in out_query_start_locs), - tuple(table.stride(0) for table in block_tables), - tuple(table.stride(0) for table in out_block_tables), - ) - if descriptor is None or descriptor.key != key: - for group, ( - block_table, - out_block_table, - out_seq_len, - out_query_start, - ) in enumerate( - zip( - block_tables, - out_block_tables, - out_seq_lens, - out_query_start_locs, - strict=True, - ) - ): - block_cols = block_col_counts[group] - if ( - block_table.device != device - or out_block_table.device != device - or out_seq_len.device != device - or out_query_start.device != device - or block_table.dtype != torch.int32 - or out_block_table.dtype != torch.int32 - or out_seq_len.dtype != torch.int32 - or out_query_start.dtype != torch.int32 - or block_table.ndim != 2 - or out_block_table.ndim != 2 - or block_table.shape[0] < num_reqs - or out_block_table.shape[0] < num_query_tokens - or out_block_table.shape[1] < block_cols - or block_table.stride(1) != 1 - or out_block_table.stride(1) != 1 - or not out_seq_len.is_contiguous() - or not out_query_start.is_contiguous() - or out_seq_len.numel() < num_query_tokens - or out_query_start.numel() < num_reqs + 1 - ): - raise ValueError( - f"grouped small-query metadata contract mismatch for group {group}" - ) - descriptor = DFlash2SmallQGroupDescriptor( - key=key, - block_table_ptrs=torch.tensor( - [table.data_ptr() for table in block_tables], - dtype=torch.uint64, - device=device, - ), - out_block_table_ptrs=torch.tensor( - [table.data_ptr() for table in out_block_tables], - dtype=torch.uint64, - device=device, - ), - out_seq_lens_ptrs=torch.tensor( - [tensor.data_ptr() for tensor in out_seq_lens], - dtype=torch.uint64, - device=device, - ), - out_query_start_loc_ptrs=torch.tensor( - [tensor.data_ptr() for tensor in out_query_start_locs], - dtype=torch.uint64, - device=device, - ), - block_table_strides=torch.tensor( - [table.stride(0) for table in block_tables], - dtype=torch.int64, - device=device, - ), - out_block_table_strides=torch.tensor( - [table.stride(0) for table in out_block_tables], - dtype=torch.int64, - device=device, - ), - block_col_counts=torch.tensor( - block_col_counts, - dtype=torch.int32, - device=device, - ), - ) - - _sm70_prepare_grouped_smallq_decode_metadata_kernel[(num_groups, num_query_tokens)]( - descriptor.out_block_table_ptrs, - descriptor.out_seq_lens_ptrs, - descriptor.out_query_start_loc_ptrs, - descriptor.block_table_ptrs, - descriptor.block_table_strides, - descriptor.out_block_table_strides, - descriptor.block_col_counts, - seq_lens, - query_start_loc, - num_reqs, - real_num_query_tokens, - REQ_BLOCK=triton.next_power_of_2(num_reqs + 1), - BLOCK_COLS=triton.next_power_of_2(max_block_cols), - num_warps=1, - ) - return descriptor - - class FlashAttnV100Metadata(TritonAttentionMetadata): """Static view of Flash-V100 fields attached to Triton metadata.""" @@ -432,10 +43,7 @@ class FlashAttnV100Metadata(TritonAttentionMetadata): causal: bool max_model_len: int flash_v100_cudagraph_capture: bool - flash_v100_batch_context_routing: bool flash_v100_contig_dense_cache: dict[tuple[int, int, int, int, int], int] - prefix_anchor_lens: torch.Tensor | None = None - decode_sliding_window: int | None = None flash_v100_decode_max_seq_len_hint: int | None flash_v100_decode_workspace_seq_capacity_hint: int | None flash_v100_static_decode_seq_hint: int | None @@ -451,7 +59,6 @@ class FlashAttnV100Metadata(TritonAttentionMetadata): smallq_decode_max_seq_len_hint: int | None smallq_decode_workspace_seq_capacity_hint: int | None smallq_decode_partition_size_hint: int | None - is_dflash_selector_target: bool def _as_flash_v100_metadata( @@ -462,25 +69,6 @@ def _as_flash_v100_metadata( return cast(FlashAttnV100Metadata, attn_metadata) -def _batch_context_routing_for_graph_variant( - routing_enabled: bool, - graph_variant: int | None, -) -> bool: - if not routing_enabled: - return False - if graph_variant is None: - # Eager execution can route directly from the live batch and context. - return True - return graph_variant == CUDAGRAPH_VARIANT_LONG_CONTEXT - - -def _batch_context_routing_cache_dtype_supported(cache_dtype: str | None) -> bool: - """Admit the exact FP8 XQA formats implemented by Flash-V100.""" - return cache_dtype == "fp8_e5m2" or ( - cache_dtype in ("fp8", "fp8_e4m3") and envs.VLLM_FLASH_V100_E4M3_BATCH_XQA - ) - - def _sm70_profile_trace(message: str, *args: object) -> None: if envs.VLLM_SM70_PROFILE_TRACE: if args: @@ -494,16 +82,12 @@ def _sm70_profile_trace(message: str, *args: object) -> None: _flash_attn_decode_paged = None _flash_attn_decode_paged_xqa = None _flash_attn_decode_paged_wmma = None -_flash_attn_grouped_verify_paged = None -_flash_attn_grouped_verify_checked = False _flash_attn_prefill_paged = None _flash_attn_prefill_paged_bhmd = None _flash_attn_prefill_paged_bfla = None _flash_attn_prefill_paged_splitkv = None _sm70_splitd_d256_ops = None _sm70_splitd_d256_ops_checked = False -_sm70_d256_gqa_architecture_op = None -_sm70_d256_gqa_architecture_op_checked = False _sm70_fa2_cu_seqlens_cache: dict[ tuple[int, int, int, int], tuple[torch.Tensor, torch.Tensor] ] = {} @@ -517,7 +101,6 @@ def _sm70_profile_trace(message: str, *args: object) -> None: _warned_decode_strict_fallback = False _warned_prefill_gather_oom = False _warned_prefill_dense_splitkv3_oom = False -_warned_prefill_d256_gqa_architecture_oom = False _logged_prefill_flash = False _logged_prefill_prefix_flash = False _logged_prefill_prefix_contig_dense = False @@ -526,11 +109,8 @@ def _sm70_profile_trace(message: str, *args: object) -> None: _logged_prefill_paged_cache = False _logged_prefill_smallq_decode = False _logged_prefill_smallq_decode_xqa = False -_logged_prefill_smallq_grouped_verify = False -_logged_prefill_smallq_grouped_verify_gate = False _logged_prefill_fa2_d256 = False _logged_prefill_dense_splitkv3 = False -_logged_prefill_d256_gqa_architecture = False _logged_prefill_triton_safe = False _logged_decode_flash = False _logged_decode_dense_reference = False @@ -736,7 +316,7 @@ def _g6_aligned_page_partition_size_hint( ): return None if ( - kv_cache_dtype in ("auto", "float16", "bfloat16") + kv_cache_dtype in ("auto", "bfloat16") and key_cache.dtype == torch.float16 and key_cache.shape[1] == 784 ): @@ -744,12 +324,6 @@ def _g6_aligned_page_partition_size_hint( # selects between them from device seq_lens. Plan the p256 workspace # envelope once. return 256 - if kv_cache_dtype in ("fp8", "fp8_e4m3") and key_cache.dtype == torch.uint8: - # Qwen3.8 TP4 has G6/D256 attention and checkpoint-provided E4M3 - # scales. The p64 route wins the accepted 1K-2K operator sweep while - # preserving the p256 scalar route's E4M3 conversion within one fp16 - # output ULP. - return 64 if kv_cache_dtype == "fp8_e5m2" and key_cache.dtype == torch.uint8: # Plan the largest p256 workspace once. The extension selects p256 or # p1024 from device seq_lens, so CUDA graph replay keeps one @@ -917,15 +491,6 @@ def _decode_xqa_allowed_for_q_per_kv( return int(seq_hint) >= _decode_xqa_q4_min_seq_len() -def _e4m3_batch_xqa_allowed(query: torch.Tensor) -> bool: - """Gate the exact SM70 E4M3 G6 batched XQA route.""" - return ( - envs.VLLM_FLASH_V100_E4M3_BATCH_XQA - and 1 < query.shape[0] <= 16 - and query.shape[1:] == (6, 256) - ) - - def _same_storage(left: torch.Tensor, right: torch.Tensor) -> bool: return left.untyped_storage().data_ptr() == right.untyped_storage().data_ptr() @@ -1220,12 +785,7 @@ def _get_flash_ops(): _flash_attn_prefill_paged_splitkv = flash_attn_prefill_paged_splitkv except ImportError: _flash_attn_prefill_paged_splitkv = None - except ImportError as exc: - logger.warning_once( - "Flash-V100 Python operators could not be imported (%s: %s).", - type(exc).__name__, - exc, - ) + except ImportError: _flash_attn_func = None _flash_attn_bhmd_func = None _flash_attn_decode_paged = None @@ -1248,23 +808,6 @@ def _get_flash_ops(): ) -def _get_flash_grouped_verify_op(): - """Load the optional exact SM70 DFlash2 grouped verifier.""" - global _flash_attn_grouped_verify_paged - global _flash_attn_grouped_verify_checked - if _flash_attn_grouped_verify_checked: - return _flash_attn_grouped_verify_paged - - _flash_attn_grouped_verify_checked = True - try: - from flash_attn_v100 import flash_attn_grouped_verify_paged - - _flash_attn_grouped_verify_paged = flash_attn_grouped_verify_paged - except ImportError: - _flash_attn_grouped_verify_paged = None - return _flash_attn_grouped_verify_paged - - def _get_sm70_splitd_d256_ops(): """Load the exact SM70 Split-D dense and paged prefill operators.""" global _sm70_splitd_d256_ops @@ -1299,35 +842,6 @@ def _get_sm70_splitd_d256_ops(): return _sm70_splitd_d256_ops -def _get_sm70_d256_gqa_architecture_op(): - """Load the optional SM70 GQA long-prefill architecture operator.""" - global _sm70_d256_gqa_architecture_op - global _sm70_d256_gqa_architecture_op_checked - if _sm70_d256_gqa_architecture_op_checked: - return _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 - - _sm70_d256_gqa_architecture_op = getattr( - torch.ops._vllm_fa2_C, - "sm70_d256_gqa_architecture_fwd", - None, - ) - except (AttributeError, ImportError, RuntimeError) as exc: - _sm70_d256_gqa_architecture_op = None - if envs.VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL: - logger.warning_once( - "SM70 D256 GQA architecture operator is unavailable " - "(%s: %s); using the exact dense prefill kernel.", - type(exc).__name__, - exc, - ) - return _sm70_d256_gqa_architecture_op - - def _uniform_cu_seqlens( tensor: torch.Tensor, *, @@ -1423,60 +937,18 @@ def _should_use_prefill_dense_splitkv3( return ( envs.VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3 and splitkv3_op is not None - and ( - query.shape == (1, 4096, 6, 256) - or ( - envs.VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3_Q8000_EXPERIMENTAL - and query.shape == (1, 8000, 6, 256) - ) - ) + and query.shape == (1, 4096, 6, 256) and key.ndim == 4 and key.shape[0] == 1 and key.shape[1] == max_seqlen_k and key.shape[2:] == (1, 256) - and max_seqlen_q == query.shape[1] + and max_seqlen_q == 4096 and max_seqlen_k >= envs.VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3_MIN_KV and max_seqlen_k > max_seqlen_q and not _is_cuda_graph_capturing(query) ) -def _should_use_prefill_d256_gqa_architecture( - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - *, - max_seqlen_q: int, - max_seqlen_k: int, - softmax_scale: float, - architecture_op: Callable[..., torch.Tensor] | None, -) -> bool: - """Gate the measured Q8000/KV40K..128K/Hq6/Hkv1/D256 family.""" - return ( - envs.VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL - and architecture_op is not None - and query.shape == (1, 8000, 6, 256) - and key.ndim == 4 - and key.shape[0] == 1 - and key.shape[2:] == (1, 256) - and value.shape == key.shape - and max_seqlen_q == 8000 - and max_seqlen_k == key.shape[1] - and 40000 <= max_seqlen_k <= 128000 - and max_seqlen_k % 8000 == 0 - and query.dtype == torch.float16 - and key.dtype == query.dtype - and value.dtype == query.dtype - and query.device == key.device - and query.device == value.device - and query.is_contiguous() - and key.is_contiguous() - and value.is_contiguous() - and abs(softmax_scale - 0.0625) <= 1.0e-8 - and not _is_cuda_graph_capturing(query) - ) - - def _try_sm70_fa2_d256_prefill( query: torch.Tensor, key: torch.Tensor, @@ -1493,9 +965,6 @@ def _try_sm70_fa2_d256_prefill( seqused_k: torch.Tensor | None = None, block_table: torch.Tensor | None = None, ) -> torch.Tensor | None: - global _logged_prefill_d256_gqa_architecture - global _warned_prefill_d256_gqa_architecture_oom - int32_max = torch.iinfo(torch.int32).max if not envs.VLLM_FLASH_V100_FA2_D256_PREFILL: return None @@ -1596,47 +1065,7 @@ def _try_sm70_fa2_d256_prefill( ) if splitd_eligible: splitd_out = out if out is not None else torch.empty_like(query) - architecture_op = ( - _get_sm70_d256_gqa_architecture_op() - if envs.VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL - else None - ) - if _should_use_prefill_d256_gqa_architecture( - query, - key, - value, - max_seqlen_q=max_seqlen_q, - max_seqlen_k=max_seqlen_k, - softmax_scale=softmax_scale, - architecture_op=architecture_op, - ): - assert architecture_op is not None - try: - splitd_result = architecture_op( - query, - key, - value, - splitd_out, - softmax_scale, - True, - ) - except torch.OutOfMemoryError: - if not _warned_prefill_d256_gqa_architecture_oom: - logger.warning( - "Insufficient memory for the default-on " - "SM70 D256 GQA long-prefill architecture; " - "falling back to the exact dense kernel." - ) - _warned_prefill_d256_gqa_architecture_oom = True - if splitd_result is not None: - if not _logged_prefill_d256_gqa_architecture: - logger.info( - "FLASH_ATTN_V100 SM70 D256 GQA " - "8K-by-40K..128K architecture route active." - ) - _logged_prefill_d256_gqa_architecture = True - _record_route("prefill_dense_d256_gqa_arch_long") - if splitd_result is None and _should_use_prefill_dense_splitkv3( + if _should_use_prefill_dense_splitkv3( query, key, max_seqlen_q=max_seqlen_q, @@ -3032,76 +2461,13 @@ class FlashAttnV100MetadataBuilder(TritonAttentionMetadataBuilder): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) spec_config = getattr(self.vllm_config, "speculative_config", None) - cache_config = getattr(self.vllm_config, "cache_config", None) - model_config = self.vllm_config.model_config - hf_text_config = getattr(model_config, "hf_text_config", None) - num_attention_heads = getattr(hf_text_config, "num_attention_heads", None) - num_key_value_heads = getattr(hf_text_config, "num_key_value_heads", None) - head_dim = getattr(hf_text_config, "head_dim", None) - batch_context_shape_supported = ( - isinstance(num_attention_heads, int) - and isinstance(num_key_value_heads, int) - and num_key_value_heads > 0 - and num_attention_heads == 6 * num_key_value_heads - and head_dim == 256 - ) self._is_speculative_draft_model = ( spec_config is not None and getattr(spec_config, "draft_model_config", None) is self.vllm_config.model_config ) - self._batch_context_routing_enabled = ( - envs.VLLM_FLASH_V100_XQA_BATCH_CONTEXT_ROUTING - and envs.VLLM_FLASH_V100_DECODE_PARTITION_SIZE is None - and spec_config is None - and _batch_context_routing_cache_dtype_supported( - getattr(cache_config, "cache_dtype", None) - ) - and batch_context_shape_supported - ) - self._is_dflash_draft_model = self._is_speculative_draft_model and ( - getattr(spec_config, "method", None) == "dflash" - ) - use_dflash = bool( - spec_config is not None - and callable(getattr(spec_config, "use_dflash", None)) - and spec_config.use_dflash() - ) - selector_engine = uses_dflash_selector_engine(self.vllm_config) - 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 selector_engine - ) - self._use_sm70_dflash2_fused_smallq_metadata = bool( - envs.VLLM_SM70_DFLASH2_FUSED_SMALLQ_METADATA - and self.device.type == "cuda" - and current_platform.is_device_capability(70) - and use_dflash - and selector_engine - ) - if self._use_sm70_dflash2_fused_smallq_metadata: - logger.info_once( - "SM70 DFlash2 fused Flash-V100 small-query metadata active." - ) self._draft_block_table: torch.Tensor | None = None self._draft_seq_lens: torch.Tensor | None = None - # Prefix-anchored SWA: persistent per-request prompt-length buffer so - # the device address stays stable across steps. - kv_cache_spec = self.kv_cache_spec - self.decode_sliding_window = ( - kv_cache_spec.decode_sliding_window - if isinstance(kv_cache_spec, PrefixAnchoredSWASpec) - else None - ) - self.persistent_prefix_anchor_lens: torch.Tensor | None = None - if self.decode_sliding_window is not None: - self.persistent_prefix_anchor_lens = torch.empty( - self.vllm_config.scheduler_config.max_num_seqs, - dtype=torch.int32, - device=self.device, - ) self._draft_query_start_loc: torch.Tensor | None = None self._flash_draft_buffer_shape: tuple[int, int] | None = None self._smallq_decode_block_table: torch.Tensor | None = None @@ -3111,40 +2477,6 @@ def __init__(self, *args, **kwargs): self._smallq_buffer_shape: tuple[int, int, int] | None = None self._decode_active_num_partitions: torch.Tensor | None = None - def _attach_prefix_anchored_metadata( - self, - attn_metadata: TritonAttentionMetadata, - common_attn_metadata, - ) -> None: - window = self.decode_sliding_window - if window is None: - return - - prefix_anchor_lens = common_attn_metadata.prefix_anchor_lens - if prefix_anchor_lens is None: - raise RuntimeError( - "prefix-anchored SWA requires per-request prefix lengths" - ) - assert self.persistent_prefix_anchor_lens is not None - anchor_reqs = common_attn_metadata.num_reqs - if prefix_anchor_lens.ndim != 1 or prefix_anchor_lens.numel() < anchor_reqs: - raise RuntimeError( - "prefix-anchored SWA prefix lengths must have shape [num_reqs]" - ) - if anchor_reqs > self.persistent_prefix_anchor_lens.numel(): - raise RuntimeError( - "prefix-anchored SWA request count exceeds metadata capacity" - ) - - prefix_anchor_lens = prefix_anchor_lens.to( - device=self.device, dtype=torch.int32, non_blocking=True - ) - persistent_anchor_lens = self.persistent_prefix_anchor_lens[:anchor_reqs] - persistent_anchor_lens.copy_(prefix_anchor_lens[:anchor_reqs]) - flash_metadata = _as_flash_v100_metadata(attn_metadata) - flash_metadata.prefix_anchor_lens = persistent_anchor_lens - flash_metadata.decode_sliding_window = window - def _attach_common_flash_metadata( self, attn_metadata: TritonAttentionMetadata, @@ -3170,15 +2502,8 @@ def _attach_common_flash_metadata( else common_attn_metadata.seq_lens_cpu ) flash_metadata.causal = common_attn_metadata.causal - flash_metadata.is_dflash_selector_target = self._is_dflash_selector_target flash_metadata.max_model_len = self.vllm_config.model_config.max_model_len flash_metadata.flash_v100_cudagraph_capture = False - flash_metadata.flash_v100_batch_context_routing = ( - _batch_context_routing_for_graph_variant( - self._batch_context_routing_enabled, - getattr(common_attn_metadata, "cudagraph_graph_variant", None), - ) - ) def _attach_ddtree_metadata( self, @@ -3568,45 +2893,6 @@ def _clear_smallq_decode_metadata( flash_metadata.smallq_decode_workspace_seq_capacity_hint = None flash_metadata.smallq_decode_partition_size_hint = None - def _attach_prepared_dflash2_smallq_metadata( - self, - attn_metadata: TritonAttentionMetadata, - prepared: DFlash2SmallQPreparedMetadata, - ) -> None: - """Attach buffers refreshed by the cross-cache-group launch.""" - if prepared.builder_id != id(self): - raise ValueError("grouped small-query metadata belongs to another builder") - if ( - self._smallq_decode_block_table is None - or self._smallq_decode_seq_lens is None - or self._smallq_query_start_loc is None - or self._smallq_buffer_shape is None - ): - raise RuntimeError("grouped small-query metadata has no persistent buffers") - token_capacity, req_capacity, _ = self._smallq_buffer_shape - if ( - prepared.num_query_tokens > token_capacity - or prepared.num_reqs > req_capacity - ): - raise RuntimeError("grouped small-query metadata exceeds captured capacity") - - self._clear_smallq_decode_metadata(attn_metadata) - flash_metadata = _as_flash_v100_metadata(attn_metadata) - flash_metadata.smallq_decode_block_table = self._smallq_decode_block_table[ - : prepared.num_query_tokens - ] - flash_metadata.smallq_decode_seq_lens = self._smallq_decode_seq_lens[ - : prepared.num_query_tokens - ] - flash_metadata.smallq_query_start_loc = self._smallq_query_start_loc[ - : prepared.num_reqs + 1 - ] - flash_metadata.smallq_decode_max_seq_len_hint = prepared.max_seq_len_hint - flash_metadata.smallq_decode_workspace_seq_capacity_hint = ( - prepared.workspace_seq_capacity_hint - ) - flash_metadata.smallq_decode_partition_size_hint = prepared.partition_size_hint - def _update_smallq_decode_metadata( self, attn_metadata: TritonAttentionMetadata, @@ -3704,110 +2990,92 @@ def _update_smallq_decode_metadata( profile_stage_t0 = time.perf_counter() if profile_enabled else 0.0 query_start_loc = attn_metadata.query_start_loc[: num_reqs + 1] + query_lens = query_start_loc[1:] - query_start_loc[:-1] + real_query_lens = query_lens real_num_query_tokens = int(query_start_loc_cpu[-1].item()) if real_num_query_tokens > num_query_tokens: return + + repeat_query_lens = query_lens padding_tokens = num_query_tokens - real_num_query_tokens + if padding_tokens > 0: + repeat_query_lens = query_lens.clone() + repeat_query_lens[-1] += padding_tokens + seq_lens = attn_metadata.seq_lens[:num_reqs] + effective_seq_lens = torch.maximum( + seq_lens, + real_query_lens.to(dtype=seq_lens.dtype), + ) + block_table = block_table.clamp_min(0) prep_ms = ( (time.perf_counter() - profile_stage_t0) * 1000.0 if profile_enabled else 0.0 ) profile_stage_t0 = time.perf_counter() if profile_enabled else 0.0 - if self._use_sm70_dflash2_fused_smallq_metadata: - _sm70_prepare_smallq_decode_metadata( - self._smallq_decode_block_table, - self._smallq_decode_seq_lens, - self._smallq_query_start_loc, - block_table, - seq_lens, - query_start_loc, - num_reqs=num_reqs, - num_query_tokens=num_query_tokens, - real_num_query_tokens=real_num_query_tokens, - ) - expand_ms = ( - (time.perf_counter() - profile_stage_t0) * 1000.0 - if profile_enabled - else 0.0 - ) - copy_ms = 0.0 - else: - query_lens = query_start_loc[1:] - query_start_loc[:-1] - real_query_lens = query_lens - repeat_query_lens = query_lens - if padding_tokens > 0: - repeat_query_lens = query_lens.clone() - repeat_query_lens[-1] += padding_tokens - - effective_seq_lens = torch.maximum( - seq_lens, - real_query_lens.to(dtype=seq_lens.dtype), - ) - clamped_block_table = block_table.clamp_min(0) - decode_block_table = torch.repeat_interleave( - clamped_block_table, - repeat_query_lens, - dim=0, - output_size=num_query_tokens, + decode_block_table = torch.repeat_interleave( + block_table, + repeat_query_lens, + dim=0, + output_size=num_query_tokens, + ).contiguous() + seq_lens_rep = torch.repeat_interleave( + effective_seq_lens, + repeat_query_lens, + output_size=num_query_tokens, + ) + query_lens_rep = torch.repeat_interleave( + real_query_lens.to(dtype=seq_lens.dtype), + repeat_query_lens, + output_size=num_query_tokens, + ) + start_locs_rep = torch.repeat_interleave( + query_start_loc[:-1].to(dtype=seq_lens.dtype), + repeat_query_lens, + output_size=num_query_tokens, + ) + token_indices = self._smallq_token_indices[:num_query_tokens].to( + dtype=seq_lens.dtype + ) + offsets = token_indices - start_locs_rep + 1 + decode_seq_lens = (seq_lens_rep - query_lens_rep + offsets).contiguous() + if padding_tokens > 0: + padding_mask = token_indices >= real_num_query_tokens + decode_seq_lens = torch.where( + padding_mask, + torch.zeros_like(decode_seq_lens), + decode_seq_lens, ).contiguous() - seq_lens_rep = torch.repeat_interleave( - effective_seq_lens, - repeat_query_lens, - output_size=num_query_tokens, - ) - query_lens_rep = torch.repeat_interleave( - real_query_lens.to(dtype=seq_lens.dtype), - repeat_query_lens, - output_size=num_query_tokens, - ) - start_locs_rep = torch.repeat_interleave( - query_start_loc[:-1].to(dtype=seq_lens.dtype), - repeat_query_lens, - output_size=num_query_tokens, - ) - token_indices = self._smallq_token_indices[:num_query_tokens].to( - dtype=seq_lens.dtype - ) - offsets = token_indices - start_locs_rep + 1 - decode_seq_lens = (seq_lens_rep - query_lens_rep + offsets).contiguous() - if padding_tokens > 0: - padding_mask = token_indices >= real_num_query_tokens - decode_seq_lens = torch.where( - padding_mask, - torch.zeros_like(decode_seq_lens), - decode_seq_lens, - ).contiguous() - decode_block_table = torch.where( - padding_mask[:, None], - torch.zeros_like(decode_block_table), - decode_block_table, - ).contiguous() - expand_ms = ( - (time.perf_counter() - profile_stage_t0) * 1000.0 - if profile_enabled - else 0.0 - ) - - profile_stage_t0 = time.perf_counter() if profile_enabled else 0.0 - self._smallq_decode_block_table[:num_query_tokens].copy_( + decode_block_table = torch.where( + padding_mask[:, None], + torch.zeros_like(decode_block_table), decode_block_table, - non_blocking=True, - ) - self._smallq_decode_seq_lens[:num_query_tokens].copy_( - decode_seq_lens, - non_blocking=True, - ) - self._smallq_query_start_loc[: num_reqs + 1].copy_( - query_start_loc, - non_blocking=True, - ) - copy_ms = ( - (time.perf_counter() - profile_stage_t0) * 1000.0 - if profile_enabled - else 0.0 - ) + ).contiguous() + expand_ms = ( + (time.perf_counter() - profile_stage_t0) * 1000.0 + if profile_enabled + else 0.0 + ) + + profile_stage_t0 = time.perf_counter() if profile_enabled else 0.0 + self._smallq_decode_block_table[:num_query_tokens].copy_( + decode_block_table, + non_blocking=True, + ) + self._smallq_decode_seq_lens[:num_query_tokens].copy_( + decode_seq_lens, + non_blocking=True, + ) + self._smallq_query_start_loc[: num_reqs + 1].copy_( + query_start_loc, + non_blocking=True, + ) + copy_ms = ( + (time.perf_counter() - profile_stage_t0) * 1000.0 + if profile_enabled + else 0.0 + ) profile_stage_t0 = time.perf_counter() if profile_enabled else 0.0 flash_metadata.smallq_decode_block_table = self._smallq_decode_block_table[ @@ -3850,7 +3118,7 @@ def _update_smallq_decode_metadata( "total_ms=%.3f clear_ms=%.3f guard_ms=%.3f ensure_ms=%.3f " "prep_ms=%.3f expand_ms=%.3f copy_ms=%.3f hint_ms=%.3f " "num_reqs=%d num_query_tokens=%d real_query_tokens=%d " - "padding_tokens=%d block_cols=%d fused=%s", + "padding_tokens=%d block_cols=%d", (time.perf_counter() - profile_t0) * 1000.0, clear_ms, guard_ms, @@ -3864,7 +3132,6 @@ def _update_smallq_decode_metadata( real_num_query_tokens, padding_tokens, int(block_table.shape[1]), - self._use_sm70_dflash2_fused_smallq_metadata, ) if _draft_graph_debug_enabled(): _graph_metadata_debug_log( @@ -3908,7 +3175,6 @@ def build_for_cudagraph_capture(self, common_attn_metadata): attn_metadata = super().build_for_cudagraph_capture(common_attn_metadata) self._attach_common_flash_metadata(attn_metadata, common_attn_metadata) flash_metadata = _as_flash_v100_metadata(attn_metadata) - self._attach_prefix_anchored_metadata(attn_metadata, common_attn_metadata) flash_metadata.seq_lens_cpu = capture_seq_lens_cpu # The Triton builder shortens capture seq_lens to 1 so full graph @@ -3937,13 +3203,12 @@ def build_for_cudagraph_capture(self, common_attn_metadata): workspace_seq_capacity_cap=workspace_seq_capacity_cap, partition_size_hint=partition_size_hint, ) - if max_query_len == 1 or self._is_dflash_draft_model: + else: # PIECEWISE graph replay captures the q=1 decode kernel arguments # during metadata warmup. Runtime drafting updates the persistent - # draft metadata buffers, so capture must bind the graph to the - # same buffers instead of transient dummy capture tensors. DFlash - # parallel drafting has q=K+1; its non-causal paged-prefill graph - # consumes the same dynamic block table and sequence lengths. + # draft metadata buffers in build_for_drafting(), so capture must + # bind the graph to the same buffers instead of transient dummy + # capture tensors. self._stabilize_draft_graph_metadata( attn_metadata, common_attn_metadata, @@ -3970,13 +3235,11 @@ def build( fast_build: bool = False, ddtree_parent_ids: torch.Tensor | None = None, ddtree_num_tree_tokens_cpu: torch.Tensor | None = None, - prepared_dflash2_smallq_metadata: (DFlash2SmallQPreparedMetadata | None) = None, ): attn_metadata = super().build( common_prefix_len, common_attn_metadata, fast_build ) self._attach_common_flash_metadata(attn_metadata, common_attn_metadata) - self._attach_prefix_anchored_metadata(attn_metadata, common_attn_metadata) self._attach_ddtree_metadata( attn_metadata, ddtree_parent_ids=ddtree_parent_ids, @@ -3989,14 +3252,13 @@ def build( and getattr(attn_metadata, "max_query_len", 1) > 1 and bool(torch.any(ddtree_num_tree_tokens_cpu[:num_reqs] > 0).item()) ) - if self._flash_draft_buffer_shape is not None and ( + if ( getattr(attn_metadata, "max_query_len", 1) == 1 - or self._is_dflash_draft_model + and self._flash_draft_buffer_shape is not None ): # FULL graph capture binds q=1 decode to these persistent buffers. - # DFlash binds its q=K+1 paged-prefill graph to the same buffers. - # Refresh them on every runtime step so replay sees the current - # request's block table and sequence metadata. + # Refresh them on every runtime decode step so replay sees the + # current request's block table and sequence metadata. self._stabilize_draft_graph_metadata( attn_metadata, common_attn_metadata, @@ -4012,20 +3274,13 @@ def build( # max_model_len=262144 / S=4224 / ps=256). The cap keeps launch equal # to the runtime coverage; the interface floors it at effective # max_seq_len (_get_decode_plan:165-169), so it can never under-cover. - if prepared_dflash2_smallq_metadata is not None: - self._attach_prepared_dflash2_smallq_metadata( - attn_metadata, - prepared_dflash2_smallq_metadata, - ) - else: - self._update_smallq_decode_metadata( - attn_metadata, - common_attn_metadata, - workspace_seq_capacity_cap=( - int(getattr(common_attn_metadata, "max_seq_len", 0) or 0) - or None - ), - ) + self._update_smallq_decode_metadata( + attn_metadata, + common_attn_metadata, + workspace_seq_capacity_cap=( + int(getattr(common_attn_metadata, "max_seq_len", 0) or 0) or None + ), + ) self._attach_decode_shape_hints(attn_metadata, common_attn_metadata) self._update_decode_active_num_partitions(attn_metadata, stage="build") self._debug_draft_metadata("build", attn_metadata, common_attn_metadata) @@ -4126,177 +3381,10 @@ def build_for_drafting(self, common_attn_metadata, draft_index: int): return attn_metadata -def prepare_dflash2_smallq_group_metadata( - *, - builders_by_group: list[tuple[int, FlashAttnV100MetadataBuilder]], - block_tables: tuple[torch.Tensor, ...], - seq_lens: torch.Tensor, - query_start_loc: torch.Tensor, - query_start_loc_cpu: torch.Tensor, - num_reqs: int, - num_query_tokens: int, - max_seq_len_hint: int, - workspace_seq_capacity_cap: int | None, - descriptor: DFlash2SmallQGroupDescriptor | None, -) -> ( - tuple[ - dict[int, DFlash2SmallQPreparedMetadata], - DFlash2SmallQGroupDescriptor, - ] - | None -): - """Refresh all pure-DFlash2 target full-attention metadata in one launch.""" - - def fallback(reason: str) -> None: - logger.info_once( - "DFlash2 grouped small-query metadata fell back to per-group launches: %s", - reason, - ) - - if ( - not envs.VLLM_SM70_DFLASH2_FUSED_SMALLQ_METADATA - or not envs.VLLM_SM70_DFLASH2_GROUPED_SMALLQ_METADATA - or not builders_by_group - or num_reqs <= 0 - or num_query_tokens <= 1 - or max_seq_len_hint <= 0 - ): - fallback("route-or-shape guard") - return None - if ( - seq_lens.device.type != "cuda" - or seq_lens.dtype != torch.int32 - or query_start_loc.device != seq_lens.device - or query_start_loc.dtype != torch.int32 - or query_start_loc_cpu.device.type != "cpu" - or query_start_loc_cpu.dtype != torch.int32 - or seq_lens.numel() < num_reqs - or query_start_loc.numel() < num_reqs + 1 - or query_start_loc_cpu.numel() < num_reqs + 1 - ): - fallback("common tensor contract") - return None - - real_num_query_tokens = int(query_start_loc_cpu[num_reqs].item()) - if not 0 < real_num_query_tokens <= num_query_tokens: - fallback("real query-token count") - return None - - input_tables: list[torch.Tensor] = [] - output_tables: list[torch.Tensor] = [] - output_seq_lens: list[torch.Tensor] = [] - output_query_start_locs: list[torch.Tensor] = [] - builder_ids: list[int] = [] - workspace_hints: list[int] = [] - seen_builders: set[int] = set() - for group_id, builder in builders_by_group: - builder_id = id(builder) - if builder_id in seen_builders: - continue - seen_builders.add(builder_id) - if ( - not builder._use_sm70_dflash2_fused_smallq_metadata - or group_id < 0 - or group_id >= len(block_tables) - or builder._smallq_decode_block_table is None - or builder._smallq_decode_seq_lens is None - or builder._smallq_query_start_loc is None - or builder._smallq_buffer_shape is None - ): - fallback("builder route or persistent buffers") - return None - - input_table = block_tables[group_id] - output_table = builder._smallq_decode_block_table - output_seq = builder._smallq_decode_seq_lens - output_query = builder._smallq_query_start_loc - token_capacity, req_capacity, builder_block_cols = builder._smallq_buffer_shape - input_block_cols = int(input_table.shape[1]) - if ( - input_table.device != seq_lens.device - or input_table.dtype != torch.int32 - or input_table.ndim != 2 - or input_table.shape[0] < num_reqs - or not input_table.is_contiguous() - or input_block_cols <= 0 - or builder_block_cols != input_block_cols - or token_capacity < num_query_tokens - or req_capacity < num_reqs - or output_table.dtype != torch.int32 - or output_seq.dtype != torch.int32 - or output_query.dtype != torch.int32 - or not output_table.is_contiguous() - or not output_seq.is_contiguous() - or not output_query.is_contiguous() - ): - fallback( - "cache-group tensor contract " - f"(group={group_id}, input_shape={tuple(input_table.shape)}, " - f"input_dtype={input_table.dtype}, input_device={input_table.device}, " - f"input_stride={input_table.stride()}, " - f"input_contiguous={input_table.is_contiguous()}, " - f"buffer_shape={builder._smallq_buffer_shape}, " - f"output_dtype={output_table.dtype}, " - f"output_stride={output_table.stride()}, " - f"output_contiguous={output_table.is_contiguous()}, " - f"seq_dtype={output_seq.dtype}, " - f"seq_contiguous={output_seq.is_contiguous()}, " - f"query_dtype={output_query.dtype}, " - f"query_contiguous={output_query.is_contiguous()}, " - f"num_reqs={num_reqs}, num_query_tokens={num_query_tokens}, " - f"input_block_cols={input_block_cols}, " - f"builder_block_cols={builder_block_cols})" - ) - return None - - raw_seq_capacity = input_block_cols * int(builder.block_size) - if workspace_seq_capacity_cap is not None: - raw_seq_capacity = min( - raw_seq_capacity, - max(max_seq_len_hint, int(workspace_seq_capacity_cap)), - ) - input_tables.append(input_table) - output_tables.append(output_table) - output_seq_lens.append(output_seq) - output_query_start_locs.append(output_query) - builder_ids.append(builder_id) - workspace_hints.append(raw_seq_capacity) - - if not input_tables: - fallback("no distinct full-attention builders") - return None - descriptor = _sm70_prepare_grouped_smallq_decode_metadata( - output_tables, - output_seq_lens, - output_query_start_locs, - input_tables, - seq_lens[:num_reqs], - query_start_loc[: num_reqs + 1], - num_reqs=num_reqs, - num_query_tokens=num_query_tokens, - real_num_query_tokens=real_num_query_tokens, - descriptor=descriptor, - ) - prepared = { - builder_id: DFlash2SmallQPreparedMetadata( - builder_id=builder_id, - num_reqs=num_reqs, - num_query_tokens=num_query_tokens, - max_seq_len_hint=max_seq_len_hint, - workspace_seq_capacity_hint=workspace_hint, - ) - for builder_id, workspace_hint in zip(builder_ids, workspace_hints) - } - return prepared, descriptor - - class FlashAttnV100Impl(TritonAttentionImpl): """Flash Attention V100 implementation with explicit fallback policy.""" def __init__(self, *args, **kwargs): - self.prefix_anchored_decode_window = kwargs.pop( - "prefix_anchored_decode_window", None - ) super().__init__(*args, **kwargs) self.kv_cache_dtype = _normalize_flash_v100_kv_cache_dtype(self.kv_cache_dtype) _log_kv_dtype_contract(self.kv_cache_dtype) @@ -4311,7 +3399,6 @@ def __init__(self, *args, **kwargs): self.flash_attn_prefill_paged_bfla, self.flash_attn_prefill_paged_splitkv, ) = _get_flash_ops() - self.flash_attn_grouped_verify_paged = _get_flash_grouped_verify_op() self.fp8_e5m2_paged_kv_to_fp16 = _get_fp8_e5m2_paged_kv_bridge_op() # V100 FA2 kernels consume fp16 Q. FP8 KV cache support is implemented # as storage compression only, with K/V dequantized inside FA2 kernels. @@ -4326,16 +3413,10 @@ def __init__(self, *args, **kwargs): "workspace_seq_capacity_hint", "active_num_partitions", "partition_size_hint", - "anchor_lens", - "anchored_window", ) if self.flash_attn_decode_paged is not None and _callable_accepts_keyword(self.flash_attn_decode_paged, name) } - self._flash_prefill_paged_supports_anchor = ( - self.flash_attn_prefill_paged is not None - and _callable_accepts_keyword(self.flash_attn_prefill_paged, "anchor_lens") - ) paged_prefill_enable = os.getenv("VLLM_FLASH_V100_ENABLE_PAGED_PREFILL") paged_prefill_disable = ( os.getenv("VLLM_FLASH_V100_DISABLE_PAGED_PREFILL", "0") == "1" @@ -4430,18 +3511,6 @@ def __init__(self, *args, **kwargs): self.use_decode_xqa and os.getenv("VLLM_FLASH_V100_SMALLQ_DECODE_USE_XQA", "1") == "1" ) - self.use_dflash2_grouped_verify = ( - self.flash_attn_grouped_verify_paged is not None - and envs.VLLM_FLASH_V100_DFLASH2_GROUPED_VERIFY - and current_platform.is_device_capability(70) - ) - self.dflash2_grouped_verify_min_model_len = ( - envs.VLLM_FLASH_V100_DFLASH2_GROUPED_VERIFY_MIN_MODEL_LEN - ) - if self.dflash2_grouped_verify_min_model_len < 1: - raise ValueError( - "VLLM_FLASH_V100_DFLASH2_GROUPED_VERIFY_MIN_MODEL_LEN must be positive" - ) decode_scalar_paged_env = os.getenv("VLLM_FLASH_V100_DECODE_USE_SCALAR_PAGED") self.use_decode_scalar_paged = decode_scalar_paged_env != "0" self.compare_bhmd_out_dir = os.getenv("VLLM_FLASH_V100_COMPARE_BHMD_OUT_DIR") @@ -4467,52 +3536,6 @@ def __init__(self, *args, **kwargs): self._decode_cache_len = 0 self._decode_cache_capacity = 0 - if self.prefix_anchored_decode_window is not None: - if ( - self.prefix_anchored_decode_window <= 0 - or self.attn_type != AttentionType.DECODER - or self.kv_cache_dtype != "auto" - ): - raise ValueError( - "prefix-anchored SWA requires a positive window, causal " - "decoder attention, and an fp16 KV cache" - ) - if self.use_triton_prefill: - raise ValueError( - "prefix-anchored SWA cannot use the Triton prefill fallback" - ) - if ( - not self.use_flash_v100_decode - or not self.use_decode_scalar_paged - or not {"anchor_lens", "anchored_window"} - <= self._flash_decode_paged_kwargs - ): - raise RuntimeError( - "prefix-anchored SWA requires the masked scalar paged " - "decode extension" - ) - if ( - not self.use_flash_v100_prefill_paged - or not self._flash_prefill_paged_supports_anchor - ): - raise RuntimeError( - "prefix-anchored SWA requires the masked paged prefill extension" - ) - - # Select the only two routes that carry the anchored mask once at - # construction time. The default-off hot path therefore retains - # its existing route predicates without extra metadata parsing. - self.smallq_decode_max_query_len = 0 - self.use_decode_paged_prefill = False - self.use_decode_dense_cache = False - self.use_decode_dense_reference = False - self.use_decode_xqa = False - self.use_smallq_decode_xqa = False - self.use_flash_v100_prefill_splitkv = False - self.use_flash_v100_prefill_bfla = False - self.use_flash_v100_prefill_contig_dense = False - self.use_flash_v100_prefill_gather_dense = False - def _reset_decode_cache(self) -> None: self._decode_cache_k = None self._decode_cache_v = None @@ -5085,8 +4108,6 @@ def _call_flash_attn_decode_paged( workspace_seq_capacity_hint: int | None = None, active_num_partitions: int | None = None, partition_size_hint: int | None = None, - anchor_lens: torch.Tensor | None = None, - anchored_window: int = 0, ) -> None: kwargs: dict[str, object] = { "softmax_scale": softmax_scale, @@ -5102,15 +4123,6 @@ def _call_flash_attn_decode_paged( "FLASH_ATTN_V100 decode op does not support sliding-window " "attention with this extension build." ) - if anchor_lens is not None and anchored_window > 0: - if "anchor_lens" not in self._flash_decode_paged_kwargs: - raise RuntimeError( - "FLASH_ATTN_V100 decode op does not support the anchored " - "decode-window mask with this extension build; rebuild " - "flash_attn_v100." - ) - kwargs["anchor_lens"] = anchor_lens - kwargs["anchored_window"] = anchored_window optional_kwargs = { "max_seq_len_hint": max_seq_len_hint, "workspace_seq_capacity_hint": workspace_seq_capacity_hint, @@ -5129,122 +4141,6 @@ def _call_flash_attn_decode_paged( **kwargs, ) - def _dflash2_grouped_verify_allowed( - self, - query: torch.Tensor, - key_cache: torch.Tensor, - value_cache: torch.Tensor, - attn_metadata: TritonAttentionMetadata, - *, - num_query_tokens: int, - ) -> bool: - """Gate the verifier on its hardware and tensor-layout contract.""" - global _logged_prefill_smallq_grouped_verify_gate - block_table = getattr(attn_metadata, "block_table", None) - seq_lens = getattr(attn_metadata, "seq_lens", None) - allowed = bool( - self.use_dflash2_grouped_verify - and self.flash_attn_grouped_verify_paged is not None - and getattr(attn_metadata, "is_dflash_selector_target", False) - and getattr(attn_metadata, "max_model_len", 0) - >= 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 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) - and tuple(key_cache.shape[2:]) == (1, 256) - and tuple(value_cache.shape) == tuple(key_cache.shape) - and key_cache.dtype == torch.uint8 - and value_cache.dtype == torch.uint8 - and key_cache.stride(-1) == 1 - and value_cache.stride(-1) == 1 - and self.kv_cache_dtype == "fp8_e5m2" - and block_table is not None - and block_table.ndim == 2 - and block_table.shape[0] == 1 - and block_table.device == query.device - and block_table.dtype == torch.int32 - and block_table.is_contiguous() - and seq_lens is not None - and seq_lens.ndim == 1 - and seq_lens.shape[0] == 1 - and seq_lens.device == query.device - and seq_lens.dtype == torch.int32 - and seq_lens.is_contiguous() - ) - if ( - self.use_dflash2_grouped_verify - and not allowed - and not _logged_prefill_smallq_grouped_verify_gate - ): - logger.info( - "FLASH_ATTN_V100 DFlash2 grouped verifier gate rejected: " - "op=%s marker=%s max_model_len=%s min_model_len=%s " - "causal=%s window=%s actual=%d q=%s/%s " - "k=%s/%s v=%s/%s kv_dtype=%s block_table=%s/%s " - "seq_lens=%s/%s.", - self.flash_attn_grouped_verify_paged is not None, - getattr(attn_metadata, "is_dflash_selector_target", False), - getattr(attn_metadata, "max_model_len", None), - self.dflash2_grouped_verify_min_model_len, - getattr(attn_metadata, "causal", True), - self._flash_v100_window_size(causal=True), - num_query_tokens, - tuple(query.shape), - query.dtype, - tuple(key_cache.shape), - key_cache.dtype, - tuple(value_cache.shape), - value_cache.dtype, - self.kv_cache_dtype, - None if block_table is None else tuple(block_table.shape), - None if block_table is None else block_table.dtype, - None if seq_lens is None else tuple(seq_lens.shape), - None if seq_lens is None else seq_lens.dtype, - ) - _logged_prefill_smallq_grouped_verify_gate = True - return allowed - - def _call_dflash2_grouped_verify( - self, - layer: torch.nn.Module, - query: torch.Tensor, - key_cache: torch.Tensor, - value_cache: torch.Tensor, - attn_metadata: TritonAttentionMetadata, - *, - out: torch.Tensor, - ) -> None: - global _logged_prefill_smallq_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)." - ) - _logged_prefill_smallq_grouped_verify = True - self.flash_attn_grouped_verify_paged( - query, - key_cache, - value_cache, - attn_metadata.block_table[:1], - attn_metadata.seq_lens[:1], - softmax_scale=self.scale, - out=out, - kv_cache_dtype=self.kv_cache_dtype, - k_scale=float(layer._k_scale_float), - v_scale=float(layer._v_scale_float), - one_pass=True, - ) - _log_fp8_kv_cache_route("decode", self.kv_cache_dtype, "dflash2_grouped_verify") - _record_route("prefill_smallq_dflash2_grouped_verify") - def _smallq_decode_xqa_allowed( self, query: torch.Tensor, @@ -5275,7 +4171,7 @@ def _smallq_decode_xqa_allowed( return False fp16_kv = ( - self.kv_cache_dtype in ("auto", "float16", "bfloat16") + self.kv_cache_dtype in ("auto", "bfloat16") and key_cache.dtype == torch.float16 and value_cache.dtype == torch.float16 ) @@ -5392,34 +4288,6 @@ def _call_flash_attn_smallq_decode_paged( ) _record_route("prefill_smallq_decode_scalar") - def _anchored_swa_params( - self, - attn_metadata: TritonAttentionMetadata, - ) -> tuple[torch.Tensor | None, int]: - """Anchored decode-window mask parameters, when active. - - Returns ``(prefix_anchor_lens, decode_sliding_window)`` when this - decoder cache group carries the engine's prefix-anchored spec and - per-request prompt lengths; otherwise ``(None, 0)``. - """ - window = self.prefix_anchored_decode_window - if window is None: - return None, 0 - - metadata_window = getattr(attn_metadata, "decode_sliding_window", None) - anchor_lens = getattr(attn_metadata, "prefix_anchor_lens", None) - if ( - self.attn_type != AttentionType.DECODER - or self.kv_cache_dtype != "auto" - or metadata_window != window - or anchor_lens is None - ): - raise RuntimeError( - "FLASH_ATTN_V100 prefix-anchored SWA metadata does not match " - "the enabled decoder-layer contract" - ) - return anchor_lens, int(window) - def _small_query_decode_enabled( self, attn_metadata: TritonAttentionMetadata, @@ -5430,6 +4298,7 @@ def _small_query_decode_enabled( or self.smallq_decode_max_query_len <= 0 ): return False + query_start_loc_cpu = getattr(attn_metadata, "query_start_loc_cpu", None) query_start_loc = ( query_start_loc_cpu @@ -5488,14 +4357,7 @@ def forward( "Flash op is unavailable or the attention features/KV cache dtype " "are unsupported. Select TRITON_ATTN for a full Triton route, or " "set VLLM_FLASH_V100_ALLOW_TRITON_FALLBACK=1 for explicit " - "diagnostic fallback. " - f"Details: layer={layer_info.get('layer_name')!r}, " - f"flash_ops_available={self.use_flash_v100}, " - f"attn_type={self.attn_type!r}, " - f"has_alibi={self.alibi_slopes is not None}, " - f"logits_soft_cap={self.logits_soft_cap!r}, " - f"has_sinks={self.sinks is not None}, " - f"kv_cache_dtype={self.kv_cache_dtype!r}." + "diagnostic fallback." ) if not (self.allow_triton_fallback or is_dflash_draft_attn): raise RuntimeError(message) @@ -5626,26 +4488,6 @@ def forward( # look like no-prefix prefill, while replayed MTP verification # is a uniform small-query decode over an existing KV prefix. # Capture the same small-query kernel branch that replay needs. - is_dflash_non_causal = bool( - getattr(layer, "is_dflash_draft_attn", False) - ) and not bool(getattr(attn_metadata, "causal", True)) - if is_dflash_non_causal: - # DFlash pre-inserts target context K/V before replay. Its - # dummy capture has seq_len == query_len and would - # otherwise freeze the no-prefix dense branch into the - # graph. Bind directly to the non-causal paged-prefix - # kernel; runtime updates its persistent sequence and - # block-table buffers before every replay. - _record_route("prefill_capture_dflash_noncausal_paged") - return self._flash_v100_prefill_with_prefix( - layer, - query, - key, - value, - kv_cache, - attn_metadata, - output, - ) smallq_decode = self._small_query_decode_enabled(attn_metadata) if smallq_decode: if _draft_graph_debug_enabled(): @@ -6479,10 +5321,6 @@ def _flash_v100_decode( ) -> torch.Tensor: """Decode path using Flash V100 directly over paged KV cache.""" window_size = self._flash_v100_window_size(causal=True) - if self.prefix_anchored_decode_window is None: - anchor_lens, anchored_window = None, 0 - else: - anchor_lens, anchored_window = self._anchored_swa_params(attn_metadata) num_actual_tokens = attn_metadata.num_actual_tokens query = query[:num_actual_tokens] out_view = output[:num_actual_tokens] @@ -6497,21 +5335,13 @@ def _flash_v100_decode( else 0 ) xqa_kv_supported = ( - ( - self.kv_cache_dtype in ("auto", "float16", "bfloat16") - and key_cache.dtype == torch.float16 - and value_cache.dtype == torch.float16 - ) - or ( - self.kv_cache_dtype == "fp8_e5m2" - and key_cache.dtype == torch.uint8 - and value_cache.dtype == torch.uint8 - ) - or ( - self.kv_cache_dtype in ("fp8", "fp8_e4m3") - and key_cache.dtype == torch.uint8 - and value_cache.dtype == torch.uint8 - ) + self.kv_cache_dtype in ("auto", "bfloat16") + and key_cache.dtype == torch.float16 + and value_cache.dtype == torch.float16 + ) or ( + self.kv_cache_dtype == "fp8_e5m2" + and key_cache.dtype == torch.uint8 + and value_cache.dtype == torch.uint8 ) # FP8 G4 XQA had no end-to-end gain on 35B-A3B TP4 and has no accepted @@ -6525,13 +5355,6 @@ def _flash_v100_decode( and key_cache.shape[2] > 0 and query.shape[1] % key_cache.shape[2] == 0 and _decode_xqa_allowed_for_q_per_kv(q_per_kv, attn_metadata) - and ( - self.kv_cache_dtype not in ("fp8", "fp8_e4m3") - or ( - q_per_kv == 6 - and (query.shape[0] == 1 or _e4m3_batch_xqa_allowed(query)) - ) - ) and ( self.kv_cache_dtype != "fp8_e5m2" or (q_per_kv != 4 and _decode_fp8_xqa_allowed(attn_metadata, query)) @@ -6585,13 +5408,6 @@ def _flash_v100_decode( None, ), partition_size_hint=partition_size_hint, - batch_context_routing=bool( - getattr( - attn_metadata, - "flash_v100_batch_context_routing", - False, - ) - ), ) _record_route("decode_xqa_paged") return output @@ -6632,8 +5448,6 @@ def _flash_v100_decode( "flash_v100_decode_active_num_partitions", None, ), - anchor_lens=anchor_lens, - anchored_window=anchored_window, ) _record_route("decode_scalar_paged") return output @@ -6981,25 +5795,6 @@ def _flash_v100_small_query_prefill_as_decode( int(query.shape[0]), int(output.shape[0]), ) - if self.use_dflash2_grouped_verify and self._dflash2_grouped_verify_allowed( - query, - key_cache, - value_cache, - attn_metadata, - num_query_tokens=num_query_tokens, - ): - query = query[:num_query_tokens] - out_view = output[:num_query_tokens] - self._call_dflash2_grouped_verify( - layer, - query, - key_cache, - value_cache, - attn_metadata, - out=out_view, - ) - return output - persistent_decode_block_table = getattr( attn_metadata, "smallq_decode_block_table", @@ -7521,25 +6316,6 @@ def _flash_v100_prefill_with_prefix( global _logged_prefill_compare, _logged_prefill_smallq_decode causal = getattr(attn_metadata, "causal", True) window_size = self._flash_v100_window_size(causal) - if self.prefix_anchored_decode_window is None: - anchor_lens, anchored_window = None, 0 - else: - anchor_lens, anchored_window = self._anchored_swa_params(attn_metadata) - if anchor_lens is not None: - # Fail closed: with the anchored decode-window mask active the - # KV cache manager evicts gap blocks, so running any unmasked - # prefill route would silently produce wrong output. - if not self.use_flash_v100_prefill_paged: - raise RuntimeError( - "FLASH_ATTN_V100 anchored decode-window mask requires " - "the paged prefill kernel; it is disabled or unavailable." - ) - if not self._flash_prefill_paged_supports_anchor: - raise RuntimeError( - "FLASH_ATTN_V100 prefill op does not support the " - "anchored decode-window mask with this extension build; " - "rebuild flash_attn_v100." - ) num_actual_tokens = attn_metadata.num_actual_tokens query = query[:num_actual_tokens] out_view = output[:num_actual_tokens] @@ -7575,11 +6351,6 @@ def _flash_v100_prefill_with_prefix( attn_metadata, query_start_loc, ): - if anchor_lens is not None: - raise RuntimeError( - "FLASH_ATTN_V100 anchored decode-window mask does not " - "support ddtree drafting metadata." - ) return self._flash_v100_ddtree_small_query_prefill_dense( layer, query, @@ -7595,7 +6366,6 @@ def _flash_v100_prefill_with_prefix( if ( causal - and anchor_lens is None and self.use_flash_v100_decode and self.smallq_decode_max_query_len > 0 and max_query_len <= self.smallq_decode_max_query_len @@ -7635,36 +6405,6 @@ def _flash_v100_prefill_with_prefix( q_len = end - start seq_len = int(seq_lens[i].item()) q_seq = query[start:end].unsqueeze(0) - if anchor_lens is not None: - # Anchored decode-window mask: single masked paged - # prefill route; every unmasked fast path is bypassed. - _record_route("prefill_prefix_paged_anchored") - out_seq = self._run_prefill_paged_call( - route="prefill_prefix_paged_anchored", - q_len=q_len, - seq_len=seq_len, - heads_q=query.shape[1], - heads_kv=num_kv_heads, - head_dim=head_dim, - block_size=block_size, - fn=lambda q_seq=q_seq, i=i: self.flash_attn_prefill_paged( # type: ignore[misc] - q_seq, - key_cache, - value_cache, - attn_metadata.block_table[i : i + 1], - attn_metadata.seq_lens[i : i + 1], - softmax_scale=self.scale, - kv_cache_dtype=self.kv_cache_dtype, - k_scale=float(layer._k_scale_float), - v_scale=float(layer._v_scale_float), - causal=causal, - window_size=window_size, - anchor_lens=anchor_lens[i : i + 1], - anchored_window=anchored_window, - ), - ) - out_view[start:end].copy_(out_seq.squeeze(0)) - continue bfla_block_mask = None use_bfla = self._should_use_prefill_bfla( q_len=q_len, diff --git a/vllm/v1/attention/backends/flashinfer_sm70.py b/vllm/v1/attention/backends/flashinfer_sm70.py index 42762bb5e8..182234612a 100644 --- a/vllm/v1/attention/backends/flashinfer_sm70.py +++ b/vllm/v1/attention/backends/flashinfer_sm70.py @@ -15,7 +15,6 @@ from vllm.platforms.interface import DeviceCapability from vllm.v1.attention.backend import AttentionType from vllm.v1.attention.backends.flash_attn_v100 import ( - DFlash2SmallQPreparedMetadata, FlashAttnV100Backend, FlashAttnV100Impl, FlashAttnV100MetadataBuilder, @@ -263,7 +262,6 @@ def build( fast_build: bool = False, ddtree_parent_ids=None, ddtree_num_tree_tokens_cpu=None, - prepared_dflash2_smallq_metadata: DFlash2SmallQPreparedMetadata | None = None, ): attn_metadata = super().build( common_prefix_len, @@ -271,7 +269,6 @@ def build( fast_build, ddtree_parent_ids, ddtree_num_tree_tokens_cpu, - prepared_dflash2_smallq_metadata, ) return self._attach_planner_decision( attn_metadata, diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 33f496826c..b52a06422c 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -27,7 +27,6 @@ split_decodes_and_prefills, ) from vllm.v1.kv_cache_interface import AttentionSpec, MambaSpec -from vllm.v1.worker.gpu.attn_utils import CommonGDNSpecMetadata logger = init_logger(__name__) @@ -60,18 +59,6 @@ class _GDNDdTreeFastCommonBuffers: token_index_initialized_size: int = 0 -@dataclass -class DFlash2GDNGroupDescriptor: - """Persistent pointer tables for one target runner's GDN cache groups.""" - - key: tuple[object, ...] - block_table_ptrs: torch.Tensor - state_output_ptrs: torch.Tensor - block_table_strides: torch.Tensor - prepared_key: tuple[int, int, int, int] | None = None - prepared_metadata: dict[int, "GDNAttentionMetadata"] | None = None - - _GDN_DDTREE_FAST_COMMON_BUFFERS: dict[ tuple[str, int | None, int, int], _GDNDdTreeFastCommonBuffers ] = {} @@ -97,9 +84,8 @@ def _get_ddtree_gdn_fast_common_buffers( dtype=torch.bool, device=device, ) - spec_token_capacity = decode_cudagraph_max_bs * width - spec_token_indx = torch.arange( - spec_token_capacity, + spec_token_indx = torch.empty( + (decode_cudagraph_max_bs * width,), dtype=torch.int32, device=device, ) @@ -130,7 +116,6 @@ def _get_ddtree_gdn_fast_common_buffers( spec_query_start_loc=spec_query_start_loc, num_accepted_tokens=num_accepted_tokens, spec_state_slot_selectors=spec_state_slot_selectors, - token_index_initialized_size=spec_token_capacity, ) _GDN_DDTREE_FAST_COMMON_BUFFERS[key] = buffers return buffers @@ -164,90 +149,6 @@ def _dflash_ddtree_gdn_fast_build_triton_enabled() -> bool: ).strip().lower() not in ("0", "false", "no", "off", "") -@triton.jit -def _load_gdn_i32_ptr(ptr_to_ptr): - ptr = tl.load(ptr_to_ptr) - ptr = tl.cast(ptr, tl.pointer_type(tl.int32)) - return tl.multiple_of(ptr, 16) - - -@triton.jit -def _dflash2_gdn_group_metadata_kernel( - block_table_ptrs, - state_output_ptrs, - block_table_strides, - state_start_indices, - req_index_mapping, - spec_query_start_loc_src, - num_accepted_src, - state_selector_src, - spec_sequence_masks_out, - spec_query_start_loc_out, - num_accepted_out, - state_selector_out, - num_spec_decodes, - batch_size, - WIDTH: tl.constexpr, - PAD_ID: tl.constexpr, - BLOCK: tl.constexpr, - USE_STATE_START: tl.constexpr, -): - """Write every GDN group's state IDs and the shared graph metadata.""" - group_id = tl.program_id(0) - block_table = _load_gdn_i32_ptr(block_table_ptrs + group_id) - state_output = _load_gdn_i32_ptr(state_output_ptrs + group_id) - block_table_stride = tl.load(block_table_strides + group_id) - - offsets = tl.arange(0, BLOCK) - rows = offsets // WIDTH - columns = offsets % WIDTH - output_mask = offsets < batch_size * WIDTH - live_state_mask = output_mask & (rows < num_spec_decodes) - state_columns = columns - if USE_STATE_START: - req_indices = tl.load( - req_index_mapping + rows, - mask=live_state_mask, - other=0, - ) - live_state_mask &= req_indices >= 0 - state_starts = tl.load( - state_start_indices + req_indices, - mask=live_state_mask, - other=-1, - ) - state_columns = columns + state_starts - live_state_mask &= (state_starts >= 0) & (state_columns < block_table_stride) - state_ids = tl.load( - block_table + rows * block_table_stride + state_columns, - mask=live_state_mask, - other=PAD_ID, - ) - tl.store(state_output + offsets, state_ids, mask=output_mask) - - # All groups share these buffers. Restrict the stores to group zero so the - # launch has one writer without introducing a second metadata kernel. - is_common_writer = group_id == 0 - row_mask = (offsets < batch_size) & is_common_writer - live_row_mask = row_mask & (offsets < num_spec_decodes) - accepted = tl.load(num_accepted_src + offsets, mask=live_row_mask, other=1) - selectors = tl.load(state_selector_src + offsets, mask=live_row_mask, other=1) - tl.store( - spec_sequence_masks_out + offsets, offsets < num_spec_decodes, mask=row_mask - ) - tl.store(num_accepted_out + offsets, accepted, mask=row_mask) - tl.store(state_selector_out + offsets, selectors, mask=row_mask) - - query_mask = (offsets < batch_size + 1) & is_common_writer - query_src_offsets = tl.minimum(offsets, num_spec_decodes) - query_offsets = tl.load( - spec_query_start_loc_src + query_src_offsets, - mask=query_mask, - other=0, - ) - tl.store(spec_query_start_loc_out + offsets, query_offsets, mask=query_mask) - - @triton.jit def _ddtree_gdn_fast_metadata_kernel( state_src, @@ -452,12 +353,6 @@ class GDNAttentionMetadata: batch_ptr: torch.Tensor | None = None token_chunk_offset_ptr: torch.Tensor | None = None - # Pure DFlash2 full-graph replay updates these persistent tensors in place. - # Cache the tuple consumed by the layer registry so the prepared fast path - # does not allocate two empty CUDA tensors and rebuild the same tuple for - # every GDN cache group on every verification round. - _prepared_spec_metadata_tensors: GDN_SPEC_METADATA_TENSORS | None = None - @dataclass class GDNSpecDecodeStateContract: @@ -821,10 +716,7 @@ def __init__( self._ddtree_fast_tail_key: tuple[int, int, int] | None = None if ( self.use_spec_decode - and ( - envs.VLLM_SM70_QWEN_GDN_SPEC_CORE_OP - or envs.VLLM_SM70_DFLASH2_FUSED_GDN_METADATA - ) + and envs.VLLM_SM70_QWEN_GDN_SPEC_CORE_OP and _dflash_ddtree_gdn_shared_common_enabled() ): self._ddtree_fast_common_buffers = _get_ddtree_gdn_fast_common_buffers( @@ -1291,8 +1183,6 @@ def build( # type: ignore[override] spec_state_slot_selectors: torch.Tensor | None = None, num_decode_draft_tokens_cpu: torch.Tensor | None = None, spec_sequence_masks_cpu: torch.Tensor | None = None, - common_gdn_metadata: CommonGDNSpecMetadata | None = None, - prepared_dflash2_metadata: GDNAttentionMetadata | None = None, current_state_block_ids: torch.Tensor | None = None, ddtree_parent_ids: torch.Tensor | None = None, ddtree_num_tree_tokens_cpu: torch.Tensor | None = None, @@ -1309,37 +1199,6 @@ def build( # type: ignore[override] query_start_loc = m.query_start_loc query_start_loc_cpu = m.query_start_loc_cpu - if prepared_dflash2_metadata is not None: - if ( - for_cudagraph_capture - or common_gdn_metadata is None - or ddtree_parent_ids is not None - or current_state_block_ids is not None - ): - raise ValueError( - "Prepared DFlash2 GDN metadata is valid only for MRV2 " - "runtime graph replay" - ) - prepared_state = prepared_dflash2_metadata.spec_state_indices_tensor - if ( - prepared_state is None - or prepared_state.data_ptr() - != self.spec_state_indices_tensor.data_ptr() - ): - raise ValueError( - "Prepared DFlash2 GDN metadata does not belong to this builder" - ) - prepared_tensors = prepared_dflash2_metadata._prepared_spec_metadata_tensors - if prepared_tensors is None: - prepared_tensors = gdn_spec_metadata_tensors( - prepared_dflash2_metadata, - query_start_loc.device, - ) - prepared_dflash2_metadata._prepared_spec_metadata_tensors = ( - prepared_tensors - ) - register_gdn_spec_metadata_tensors(self.layer_names, prepared_tensors) - return prepared_dflash2_metadata if fast_build and _dflash_ddtree_gdn_fast_build_enabled(): fast_metadata = self._build_fast_pure_ddtree_full_graph( common_attn_metadata=common_attn_metadata, @@ -1368,12 +1227,6 @@ def build( # type: ignore[override] is_mamba_cache_all = self.vllm_config.cache_config.mamba_cache_mode == "all" num_reqs = query_start_loc_cpu.numel() - 1 - if common_gdn_metadata is not None: - if ddtree_parent_ids is not None or current_state_block_ids is not None: - raise ValueError( - "Common MRV2 GDN metadata cannot be used with DDTree state" - ) - spec_sequence_masks_cpu = common_gdn_metadata.spec_sequence_masks_cpu if spec_sequence_masks_cpu is not None: assert spec_sequence_masks_cpu.dtype == torch.bool assert spec_sequence_masks_cpu.ndim == 1 @@ -1410,9 +1263,6 @@ def build( # type: ignore[override] if not self.use_spec_decode: spec_sequence_masks = None num_spec_decodes = 0 - elif common_gdn_metadata is not None: - spec_sequence_masks = common_gdn_metadata.spec_sequence_masks - num_spec_decodes = common_gdn_metadata.num_spec_decodes else: if spec_sequence_masks_cpu is None: if num_decode_draft_tokens_cpu is None: @@ -1623,11 +1473,9 @@ def build( # type: ignore[override] else: # query_start_loc may be padded for CUDA graph replay. The CPU # metadata is authoritative for the live request count here. - query_lens = None - if common_gdn_metadata is None: - query_lens = query_lens_cpu.to( - query_start_loc.device, non_blocking=True - ) + query_lens = query_lens_cpu.to( + query_start_loc.device, non_blocking=True + ) profile_state_contract_t0 = ( time.perf_counter() if metadata_profile else 0.0 ) @@ -1647,142 +1495,113 @@ def build( # type: ignore[override] time.perf_counter() - profile_state_contract_t0 ) * 1000.0 - if common_gdn_metadata is not None: - num_prefills = common_gdn_metadata.num_prefills - num_prefill_tokens = common_gdn_metadata.num_prefill_tokens - num_decodes = common_gdn_metadata.num_decodes - num_decode_tokens = common_gdn_metadata.num_decode_tokens - num_spec_decode_tokens = common_gdn_metadata.num_spec_decode_tokens - spec_token_indx = common_gdn_metadata.spec_token_indx - non_spec_token_indx = common_gdn_metadata.non_spec_token_indx - spec_query_start_loc = common_gdn_metadata.spec_query_start_loc - non_spec_query_start_loc = ( - common_gdn_metadata.non_spec_query_start_loc + if envs.VLLM_SM70_MTP_LEGACY_GDN_MIXED_DECODE_ROUTING: + # 0.0.3 kept ordinary query_len==1 rows on the decode path even + # when another row was running speculative verification. This + # is an A/B guard for MTP-only recurrent-state corruption. + num_decodes = (non_spec_query_lens_cpu == 1).sum().item() + num_prefills = ( + non_spec_query_lens_cpu.size(0) - num_decodes - num_zero_len ) - non_spec_query_start_loc_cpu = ( - common_gdn_metadata.non_spec_query_start_loc_cpu - ) - spec_state_indices_tensor = state_contract.spec_state_indices_tensor - non_spec_state_indices_tensor = ( - None - if num_prefills == 0 and num_decodes == 0 - else state_contract.non_spec_state_indices_tensor + num_decode_tokens = num_decodes + num_prefill_tokens = ( + non_spec_query_lens_cpu.sum().item() - num_decode_tokens ) else: - if envs.VLLM_SM70_MTP_LEGACY_GDN_MIXED_DECODE_ROUTING: - # 0.0.3 kept ordinary query_len==1 rows on the decode - # path even when another row was running speculative - # verification. - num_decodes = (non_spec_query_lens_cpu == 1).sum().item() - num_prefills = ( - non_spec_query_lens_cpu.size(0) - num_decodes - num_zero_len - ) - num_decode_tokens = num_decodes - num_prefill_tokens = ( - non_spec_query_lens_cpu.sum().item() - num_decode_tokens - ) - else: - # Mixed non-spec rows use the prefill path so their GDN - # state metadata stays separate from verification rows. - num_decodes = 0 - num_prefills = non_spec_query_lens_cpu.size(0) - num_zero_len - num_decode_tokens = 0 - num_prefill_tokens = non_spec_query_lens_cpu.sum().item() - num_spec_decode_tokens = ( - query_lens_cpu.sum().item() - - num_prefill_tokens - - num_decode_tokens - ) + # When active spec decodes are present, route non-spec requests + # through the prefill path so mixed batches keep separate GDN + # state metadata for spec and non-spec tokens. + num_decodes = 0 + num_prefills = non_spec_query_lens_cpu.size(0) - num_zero_len + num_decode_tokens = 0 + num_prefill_tokens = non_spec_query_lens_cpu.sum().item() + num_spec_decode_tokens = ( + query_lens_cpu.sum().item() - num_prefill_tokens - num_decode_tokens + ) - if num_prefills == 0 and num_decodes == 0: - spec_token_size = min( - num_spec_decodes * (self.num_spec_state_tokens + 1), - query_start_loc_cpu[-1].item(), - ) - spec_token_indx = torch.arange( - spec_token_size, - 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 = ( - state_contract.spec_state_indices_tensor - ) - non_spec_state_indices_tensor = None - # Padded sequences are always at the back. - 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: - assert query_lens is not None - spec_token_masks = torch.repeat_interleave( - spec_sequence_masks, - query_lens, - output_size=query_start_loc_cpu[-1].item(), + if num_prefills == 0 and num_decodes == 0: + spec_token_size = min( + num_spec_decodes * (self.num_spec_state_tokens + 1), + query_start_loc_cpu[-1].item(), + ) + spec_token_indx = torch.arange( + spec_token_size, + 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 = state_contract.spec_state_indices_tensor + if for_cudagraph_capture: + spec_state_indices_tensor = torch.full_like( + spec_state_indices_tensor, PAD_SLOT_ID ) - index = torch.argsort(spec_token_masks, stable=True) - num_non_spec_tokens = num_prefill_tokens + num_decode_tokens - non_spec_token_indx = index[:num_non_spec_tokens] - spec_token_indx = index[num_non_spec_tokens:] + non_spec_state_indices_tensor = None + # Padded sequences are always at the back, so the first + # num_spec_decodes + 1 entries of query_start_loc already + # contain the correct cumulative token counts. + 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: + spec_token_masks = torch.repeat_interleave( + spec_sequence_masks, + query_lens, + output_size=query_start_loc_cpu[-1].item(), + ) + index = torch.argsort(spec_token_masks, stable=True) + num_non_spec_tokens = num_prefill_tokens + num_decode_tokens + non_spec_token_indx = index[:num_non_spec_tokens] + spec_token_indx = index[num_non_spec_tokens:] - spec_state_indices_tensor = ( - state_contract.spec_state_indices_tensor + spec_state_indices_tensor = state_contract.spec_state_indices_tensor + non_spec_state_indices_tensor = ( + state_contract.non_spec_state_indices_tensor + ) + if for_cudagraph_capture: + spec_state_indices_tensor = torch.full_like( + spec_state_indices_tensor, PAD_SLOT_ID ) - non_spec_state_indices_tensor = ( - state_contract.non_spec_state_indices_tensor + non_spec_state_indices_tensor = torch.full_like( + non_spec_state_indices_tensor, PAD_SLOT_ID ) - if for_cudagraph_capture: - spec_state_indices_tensor = torch.full_like( - spec_state_indices_tensor, PAD_SLOT_ID - ) - non_spec_state_indices_tensor = torch.full_like( - non_spec_state_indices_tensor, PAD_SLOT_ID - ) - spec_query_start_loc = torch.zeros( - num_spec_decodes + 1, - dtype=torch.int32, - device=query_start_loc.device, - ) - torch.cumsum( - query_lens[spec_sequence_masks], - dim=0, - out=spec_query_start_loc[1:], - ) - non_spec_query_start_loc = torch.zeros( - query_lens.size(0) - num_spec_decodes + 1, - dtype=torch.int32, - device=query_start_loc.device, - ) - torch.cumsum( - query_lens[~spec_sequence_masks], - dim=0, - out=non_spec_query_start_loc[1:], - ) - non_spec_query_start_loc_cpu = torch.zeros( - query_lens_cpu.size(0) - num_spec_decodes + 1, - dtype=torch.int32, - device="cpu", - ) - torch.cumsum( - query_lens_cpu[~spec_sequence_masks_cpu], - dim=0, - out=non_spec_query_start_loc_cpu[1:], - ) + spec_query_start_loc = torch.zeros( + num_spec_decodes + 1, + dtype=torch.int32, + device=query_start_loc.device, + ) + torch.cumsum( + query_lens[spec_sequence_masks], + dim=0, + out=spec_query_start_loc[1:], + ) + non_spec_query_start_loc = torch.zeros( + query_lens.size(0) - num_spec_decodes + 1, + dtype=torch.int32, + device=query_start_loc.device, + ) + torch.cumsum( + query_lens[~spec_sequence_masks], + dim=0, + out=non_spec_query_start_loc[1:], + ) + non_spec_query_start_loc_cpu = torch.zeros( + query_lens_cpu.size(0) - num_spec_decodes + 1, + dtype=torch.int32, + device="cpu", + ) + torch.cumsum( + query_lens_cpu[~spec_sequence_masks_cpu], + dim=0, + out=non_spec_query_start_loc_cpu[1:], + ) num_accepted_tokens = state_contract.num_accepted_tokens spec_state_slot_selectors = state_contract.spec_state_slot_selectors assert spec_query_start_loc is not None - if common_gdn_metadata is None: - assert spec_query_start_loc[-1].item() == num_spec_decode_tokens - else: - assert ( - common_gdn_metadata.num_spec_decode_tokens == num_spec_decode_tokens - ) + assert spec_query_start_loc[-1].item() == num_spec_decode_tokens assert spec_state_indices_tensor is not None assert spec_state_indices_tensor.shape[0] == num_spec_decodes @@ -2176,319 +1995,3 @@ def build_for_cudagraph_capture( spec_sequence_masks_cpu=spec_sequence_masks_cpu, for_cudagraph_capture=True, ) - - -def prepare_dflash2_gdn_group_metadata( - *, - builders_by_group: list[tuple[int, GDNAttentionMetadataBuilder]], - block_tables: tuple[torch.Tensor, ...], - common_gdn_metadata: CommonGDNSpecMetadata, - num_accepted_tokens: torch.Tensor, - num_actual_tokens: int, - descriptor: DFlash2GDNGroupDescriptor | None, - state_start_indices: torch.Tensor | None = None, - req_index_mapping: torch.Tensor | None = None, -) -> ( - tuple[ - dict[int, GDNAttentionMetadata], - DFlash2GDNGroupDescriptor, - ] - | None -): - """Prepare all pure-MRV2 DFlash2 GDN graph metadata in one launch. - - ``mamba_cache_mode=none`` reads the first speculative state columns. - ``mamba_cache_mode=align`` supplies the authoritative, post-precopy state - column for each live request. DFlash2 batches keep live speculative rows at - the front and CUDA-graph padding at the back, so one pointer-table kernel can - perform the same state selection and tail fill without ten independent - gather/copy pipelines. - """ - if not envs.VLLM_SM70_DFLASH2_FUSED_GDN_METADATA: - return None - if not builders_by_group or num_actual_tokens <= 0: - return None - if num_accepted_tokens.device.type != "cuda": - return None - if num_accepted_tokens.dtype != torch.int32 or num_accepted_tokens.ndim != 1: - return None - - use_state_start = state_start_indices is not None - if use_state_start != (req_index_mapping is not None): - return None - if use_state_start: - assert state_start_indices is not None - assert req_index_mapping is not None - if ( - state_start_indices.device != num_accepted_tokens.device - or state_start_indices.dtype != torch.int32 - or state_start_indices.ndim != 1 - or req_index_mapping.device != num_accepted_tokens.device - or req_index_mapping.dtype != torch.int32 - or req_index_mapping.ndim != 1 - ): - return None - - spec_mask_cpu = common_gdn_metadata.spec_sequence_masks_cpu - num_spec_decodes = common_gdn_metadata.num_spec_decodes - if ( - common_gdn_metadata.num_prefills != 0 - or common_gdn_metadata.num_decodes != 0 - or num_spec_decodes <= 0 - or num_spec_decodes > num_actual_tokens - or num_spec_decodes > spec_mask_cpu.numel() - or common_gdn_metadata.num_spec_decode_tokens > num_actual_tokens - or ( - req_index_mapping is not None - and req_index_mapping.numel() < num_spec_decodes - ) - ): - return None - if not bool(torch.all(spec_mask_cpu[:num_spec_decodes]).item()): - return None - if bool(torch.any(spec_mask_cpu[num_spec_decodes:]).item()): - return None - - query_start_loc = common_gdn_metadata.spec_query_start_loc - if ( - query_start_loc.device != num_accepted_tokens.device - or query_start_loc.dtype != torch.int32 - or query_start_loc.ndim != 1 - or query_start_loc.numel() != num_spec_decodes + 1 - ): - return None - if num_accepted_tokens.numel() < num_spec_decodes: - return None - - first_builder = builders_by_group[0][1] - mamba_cache_mode = first_builder.vllm_config.cache_config.mamba_cache_mode - if mamba_cache_mode not in ("none", "align"): - return None - if use_state_start != (mamba_cache_mode == "align"): - return None - width = first_builder.num_spec_state_tokens + 1 - common_buffers = first_builder._ddtree_fast_common_buffers - if common_buffers is None: - return None - if ( - num_actual_tokens > first_builder.decode_cudagraph_max_bs - or common_gdn_metadata.num_spec_decode_tokens - > first_builder.decode_cudagraph_max_bs - ): - return None - - input_tables: list[torch.Tensor] = [] - output_states: list[torch.Tensor] = [] - builder_ids: list[int] = [] - group_ids: list[int] = [] - seen_builders: set[int] = set() - for group_id, builder in builders_by_group: - builder_id = id(builder) - if builder_id in seen_builders: - continue - seen_builders.add(builder_id) - if ( - builder.num_spec_state_tokens + 1 != width - or not builder.use_full_cuda_graph - or builder.vllm_config.cache_config.mamba_cache_mode != mamba_cache_mode - or builder.decode_cudagraph_max_bs < num_actual_tokens - or builder._ddtree_fast_common_buffers is not common_buffers - or group_id < 0 - or group_id >= len(block_tables) - ): - return None - block_table = block_tables[group_id] - state_output = builder.spec_state_indices_tensor - if ( - block_table.device != num_accepted_tokens.device - or block_table.dtype != torch.int32 - or block_table.ndim != 2 - or block_table.shape[0] < num_spec_decodes - or block_table.shape[1] < width - or not block_table.is_contiguous() - or state_output.device != num_accepted_tokens.device - or state_output.dtype != torch.int32 - or state_output.ndim != 2 - or state_output.shape[0] < num_actual_tokens - or state_output.shape[1] != width - or not state_output.is_contiguous() - ): - return None - input_tables.append(block_table) - output_states.append(state_output) - builder_ids.append(builder_id) - group_ids.append(group_id) - - if not input_tables: - return None - spec_token_size = common_gdn_metadata.spec_token_indx.numel() - if ( - spec_token_size > common_buffers.spec_token_indx.numel() - or num_actual_tokens > common_buffers.spec_sequence_masks.numel() - or num_actual_tokens + 1 > common_buffers.spec_query_start_loc.numel() - ): - return None - if spec_token_size > common_buffers.token_index_initialized_size: - return None - - descriptor_key: tuple[object, ...] = ( - num_accepted_tokens.device.type, - num_accepted_tokens.device.index, - width, - tuple(group_ids), - tuple(builder_ids), - tuple(table.data_ptr() for table in input_tables), - tuple(state.data_ptr() for state in output_states), - tuple(table.stride(0) for table in input_tables), - use_state_start, - common_buffers.spec_sequence_masks.data_ptr(), - common_buffers.spec_token_indx.data_ptr(), - common_buffers.non_spec_token_indx.data_ptr(), - common_buffers.spec_query_start_loc.data_ptr(), - common_buffers.num_accepted_tokens.data_ptr(), - common_buffers.spec_state_slot_selectors.data_ptr(), - ) - if descriptor is None or descriptor.key != descriptor_key: - descriptor = DFlash2GDNGroupDescriptor( - key=descriptor_key, - block_table_ptrs=torch.tensor( - [table.data_ptr() for table in input_tables], - dtype=torch.uint64, - device=num_accepted_tokens.device, - ), - state_output_ptrs=torch.tensor( - [state.data_ptr() for state in output_states], - dtype=torch.uint64, - device=num_accepted_tokens.device, - ), - block_table_strides=torch.tensor( - [table.stride(0) for table in input_tables], - dtype=torch.int64, - device=num_accepted_tokens.device, - ), - ) - - block = triton.next_power_of_2( - max(num_actual_tokens * width, num_actual_tokens + 1) - ) - _dflash2_gdn_group_metadata_kernel[(len(input_tables),)]( - descriptor.block_table_ptrs, - descriptor.state_output_ptrs, - descriptor.block_table_strides, - num_accepted_tokens if state_start_indices is None else state_start_indices, - num_accepted_tokens if req_index_mapping is None else req_index_mapping, - query_start_loc, - num_accepted_tokens, - num_accepted_tokens, - common_buffers.spec_sequence_masks, - common_buffers.spec_query_start_loc, - common_buffers.num_accepted_tokens, - common_buffers.spec_state_slot_selectors, - num_spec_decodes, - num_actual_tokens, - WIDTH=width, - PAD_ID=PAD_SLOT_ID, - BLOCK=block, - USE_STATE_START=use_state_start, - num_warps=1, - ) - common_buffers.initialized_key = ( - num_actual_tokens, - common_gdn_metadata.num_spec_decode_tokens, - width, - ) - - prepared_key = ( - num_spec_decodes, - num_actual_tokens, - common_gdn_metadata.num_spec_decode_tokens, - spec_token_size, - ) - prepared = descriptor.prepared_metadata - if descriptor.prepared_key != prepared_key or prepared is None: - prepared = {} - for builder_id, state_output in zip(builder_ids, output_states): - prepared[builder_id] = GDNAttentionMetadata( - num_prefills=0, - num_prefill_tokens=0, - num_decodes=0, - num_decode_tokens=0, - num_spec_decodes=num_spec_decodes, - num_spec_decode_tokens=common_gdn_metadata.num_spec_decode_tokens, - num_actual_tokens=num_actual_tokens, - has_initial_state=None, - chunk_indices=None, - chunk_offsets=None, - spec_query_start_loc=common_buffers.spec_query_start_loc[ - : num_actual_tokens + 1 - ], - non_spec_query_start_loc=None, - spec_state_indices_tensor=state_output[:num_actual_tokens], - non_spec_state_indices_tensor=None, - spec_sequence_masks=common_buffers.spec_sequence_masks[ - :num_actual_tokens - ], - spec_token_indx=common_buffers.spec_token_indx[:spec_token_size], - non_spec_token_indx=common_buffers.non_spec_token_indx[:0], - num_accepted_tokens=common_buffers.num_accepted_tokens[ - :num_actual_tokens - ], - spec_state_slot_selectors=common_buffers.spec_state_slot_selectors[ - :num_actual_tokens - ], - ddtree_parent_ids=None, - ddtree_num_tree_tokens_cpu=None, - nums_dict=None, - batch_ptr=None, - token_chunk_offset_ptr=None, - ) - descriptor.prepared_key = prepared_key - descriptor.prepared_metadata = prepared - - assert prepared is not None - if envs.VLLM_SM70_DFLASH2_GDN_METADATA_SHADOW: - spec_mask = common_gdn_metadata.spec_sequence_masks - expected_accepted = num_accepted_tokens[spec_mask] - for group_id, builder in builders_by_group: - actual = prepared[id(builder)] - actual_state = actual.spec_state_indices_tensor - assert actual_state is not None - source_table = block_tables[group_id][spec_mask] - if use_state_start: - assert state_start_indices is not None - assert req_index_mapping is not None - req_indices = req_index_mapping[:num_spec_decodes].to(torch.long) - starts = state_start_indices[req_indices].to(torch.long) - columns = starts[:, None] + torch.arange( - width, - dtype=torch.long, - device=source_table.device, - ) - expected_state = torch.gather(source_table, 1, columns) - else: - expected_state = source_table[:, :width] - torch.testing.assert_close( - actual_state[:num_spec_decodes], - expected_state, - rtol=0, - atol=0, - ) - torch.testing.assert_close( - actual_state[num_spec_decodes:], - torch.full_like( - actual_state[num_spec_decodes:], - PAD_SLOT_ID, - ), - rtol=0, - atol=0, - ) - torch.testing.assert_close( - common_buffers.num_accepted_tokens[:num_spec_decodes], - expected_accepted, - rtol=0, - atol=0, - ) - if torch.any(common_buffers.spec_sequence_masks[num_spec_decodes:]).item(): - raise AssertionError("DFlash2 fused GDN metadata left a live padded row") - - return prepared, descriptor diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 2171f74647..2870ec9a15 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -203,11 +203,6 @@ class DeepseekV32IndexerMetadata: # hacky way to access the data now, need to be in chunked meta seq_lens: torch.Tensor max_seq_len: int - # Decode seq_lens and the indexer KV cache use compressed positions for - # DeepSeek V4. Keep the matching host bound explicit: CUDA Graph capture - # needs a static logits width, but using the uncompressed attention bound - # launches 4x/128x too much work for C4/C128 layers. - compressed_max_seq_len: int slot_mapping: torch.Tensor # New for MLA (compared to FlashAttention) @@ -632,9 +627,6 @@ def build( attn_metadata = DeepseekV32IndexerMetadata( seq_lens=common_attn_metadata.seq_lens, max_seq_len=common_attn_metadata.max_seq_len, - compressed_max_seq_len=max( - 1, common_attn_metadata.max_seq_len // self.compress_ratio - ), slot_mapping=compressed_slot_mapping, num_decodes=num_decodes, num_decode_tokens=num_decode_tokens, diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index 9688a42b64..a96c0dd115 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -42,8 +42,12 @@ 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" ) diff --git a/vllm/v1/attention/ops/triton_prefill_attention.py b/vllm/v1/attention/ops/triton_prefill_attention.py index 8450ba1eda..4e28e95e95 100644 --- a/vllm/v1/attention/ops/triton_prefill_attention.py +++ b/vllm/v1/attention/ops/triton_prefill_attention.py @@ -184,7 +184,11 @@ 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 0d59a03535..ef5b256eb6 100644 --- a/vllm/v1/attention/ops/triton_unified_attention.py +++ b/vllm/v1/attention/ops/triton_unified_attention.py @@ -558,7 +558,9 @@ 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: @@ -918,8 +920,12 @@ 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..c5e8953745 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -34,7 +34,7 @@ def __init__( self, kv_cache_config: KVCacheConfig, max_model_len: int, - max_in_flight_tokens: int, + max_num_batched_tokens: int, use_eagle: bool, enable_caching: bool, enable_kv_cache_events: bool, @@ -66,7 +66,7 @@ def __init__( self.single_type_managers = tuple( get_manager_for_kv_cache_spec( kv_cache_spec=kv_cache_group.kv_cache_spec, - max_in_flight_tokens=max_in_flight_tokens, + max_num_batched_tokens=max_num_batched_tokens, max_model_len=max_model_len, block_pool=self.block_pool, enable_caching=enable_caching, @@ -236,10 +236,7 @@ def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]: ] def remove_skipped_blocks( - self, - request_id: str, - processed_computed_tokens: int, - num_prompt_tokens: int | None = None, + self, request_id: str, total_computed_tokens: int ) -> None: """ Remove the blocks that are no longer needed from `blocks` and replace @@ -247,16 +244,11 @@ def remove_skipped_blocks( Args: request_id: The request ID. - processed_computed_tokens: Computed-token prefix length covering - fully processed and committed tokens only (safe to free). - num_prompt_tokens: Optional prompt length. Prefix-anchored SWA - managers use this to free gap blocks between the prefill tail - and the decode window; other manager types ignore it. + total_computed_tokens: The total number of computed tokens, including + local computed tokens and external computed tokens. """ for manager in self.single_type_managers: - manager.remove_skipped_blocks( - request_id, processed_computed_tokens, num_prompt_tokens - ) + manager.remove_skipped_blocks(request_id, total_computed_tokens) def get_blocks(self, request_id: str) -> tuple[list[KVCacheBlock], ...]: """ @@ -293,7 +285,7 @@ def __init__( self, kv_cache_config: KVCacheConfig, max_model_len: int, - max_in_flight_tokens: int, + max_num_batched_tokens: int, use_eagle: bool, enable_kv_cache_events: bool, dcp_world_size: int, @@ -304,7 +296,7 @@ def __init__( super().__init__( kv_cache_config, max_model_len, - max_in_flight_tokens, + max_num_batched_tokens, use_eagle, False, enable_kv_cache_events, @@ -340,7 +332,7 @@ def __init__( self, kv_cache_config: KVCacheConfig, max_model_len: int, - max_in_flight_tokens: int, + max_num_batched_tokens: int, use_eagle: bool, enable_caching: bool, enable_kv_cache_events: bool, @@ -352,7 +344,7 @@ def __init__( super().__init__( kv_cache_config, max_model_len, - max_in_flight_tokens, + max_num_batched_tokens, use_eagle, enable_caching, enable_kv_cache_events, @@ -407,7 +399,7 @@ def __init__( self, kv_cache_config: KVCacheConfig, max_model_len: int, - max_in_flight_tokens: int, + max_num_batched_tokens: int, use_eagle: bool, enable_caching: bool, enable_kv_cache_events: bool, @@ -419,7 +411,7 @@ def __init__( super().__init__( kv_cache_config, max_model_len, - max_in_flight_tokens, + max_num_batched_tokens, use_eagle, enable_caching, enable_kv_cache_events, @@ -618,7 +610,7 @@ def _get_block_hashes(kv_cache_spec: KVCacheSpec) -> BlockHashList: def get_kv_cache_coordinator( kv_cache_config: KVCacheConfig, max_model_len: int, - max_in_flight_tokens: int, + max_num_batched_tokens: int, use_eagle: bool, enable_caching: bool, enable_kv_cache_events: bool, @@ -631,7 +623,7 @@ def get_kv_cache_coordinator( return KVCacheCoordinatorNoPrefixCache( kv_cache_config, max_model_len, - max_in_flight_tokens, + max_num_batched_tokens, use_eagle, enable_kv_cache_events, dcp_world_size=dcp_world_size, @@ -643,7 +635,7 @@ def get_kv_cache_coordinator( return UnitaryKVCacheCoordinator( kv_cache_config, max_model_len, - max_in_flight_tokens, + max_num_batched_tokens, use_eagle, enable_caching, enable_kv_cache_events, @@ -655,7 +647,7 @@ def get_kv_cache_coordinator( return HybridKVCacheCoordinator( kv_cache_config, max_model_len, - max_in_flight_tokens, + max_num_batched_tokens, use_eagle, enable_caching, enable_kv_cache_events, diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 947d8cb0ad..9359d8843a 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -8,14 +8,10 @@ from vllm.distributed.kv_events import BlockStored, KVCacheEvent from vllm.logger import init_logger -from vllm.utils.math_utils import cdiv from vllm.v1.core.kv_cache_coordinator import get_kv_cache_coordinator from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector from vllm.v1.core.kv_cache_utils import KVCacheBlock from vllm.v1.kv_cache_interface import ( - AttentionSpec, - CrossAttentionSpec, - EncoderOnlyAttentionSpec, KVCacheConfig, get_kv_cache_spec_kind, get_kv_cache_spec_sliding_window, @@ -117,7 +113,7 @@ def __init__( kv_cache_config: KVCacheConfig, max_model_len: int, hash_block_size: int, - max_in_flight_tokens: int | None = None, + max_num_batched_tokens: int | None = None, enable_caching: bool = True, use_eagle: bool = False, log_stats: bool = False, @@ -130,8 +126,8 @@ def __init__( # When unset, fall back to `max_model_len` so the recycling-aware cap # collapses to the prior (uncapped) admission behavior. The scheduler # always supplies the real value at runtime. - if max_in_flight_tokens is None: - max_in_flight_tokens = max_model_len + if max_num_batched_tokens is None: + max_num_batched_tokens = max_model_len self.enable_caching = enable_caching self.use_eagle = use_eagle @@ -145,7 +141,7 @@ def __init__( self.coordinator = get_kv_cache_coordinator( kv_cache_config=kv_cache_config, max_model_len=self.max_model_len, - max_in_flight_tokens=max_in_flight_tokens, + max_num_batched_tokens=max_num_batched_tokens, use_eagle=self.use_eagle, enable_caching=self.enable_caching, enable_kv_cache_events=enable_kv_cache_events, @@ -374,13 +370,8 @@ def allocate_slots( # insufficient free blocks. # Should call this function before allocating new blocks to reduce # the number of evicted blocks. - # Free on the processed-token basis: in-flight steps' attention windows - # still read blocks below the optimistic boundary, and rejected spec - # tokens can roll it back. self.coordinator.remove_skipped_blocks( - request.request_id, - max(0, total_computed_tokens - request.num_in_flight_tokens), - num_prompt_tokens=request.num_prompt_tokens, + request.request_id, total_computed_tokens ) num_blocks_to_allocate = self.coordinator.get_num_blocks_to_allocate( @@ -446,24 +437,17 @@ def free(self, request: Request) -> None: self.coordinator.free(request.request_id) def remove_skipped_blocks( - self, - request_id: str, - processed_computed_tokens: int, - num_prompt_tokens: int | None = None, + self, request_id: str, total_computed_tokens: int ) -> None: """Remove the blocks that are no longer needed from `blocks` and replace the removed blocks with null_block. Args: request_id: The request ID. - processed_computed_tokens: Computed-token prefix length covering - fully processed and committed tokens only (safe to free). - num_prompt_tokens: Optional prompt length for prefix-anchored SWA - gap eviction. + total_computed_tokens: The total number of computed tokens, including + local computed tokens and external computed tokens. """ - self.coordinator.remove_skipped_blocks( - request_id, processed_computed_tokens, num_prompt_tokens - ) + self.coordinator.remove_skipped_blocks(request_id, total_computed_tokens) def evict_blocks(self, block_ids: set[int]) -> None: """evict blocks from the prefix cache by their block IDs. @@ -557,26 +541,6 @@ def get_block_ids(self, request_id: str) -> tuple[list[int], ...]: """Get the block ids of a request.""" return self.get_blocks(request_id).get_block_ids() - def get_block_ids_for_computed_tokens( - self, - request_id: str, - num_computed_tokens: int, - ) -> tuple[list[int], ...]: - """Get block ids covering the request's computed tokens.""" - block_ids = self.get_block_ids(request_id) - clipped_block_ids: list[list[int]] = [] - for group, ids in zip(self.kv_cache_config.kv_cache_groups, block_ids): - spec = group.kv_cache_spec - if not isinstance(spec, AttentionSpec) or isinstance( - spec, (CrossAttentionSpec, EncoderOnlyAttentionSpec) - ): - clipped_block_ids.append(ids) - continue - - num_valid_blocks = cdiv(num_computed_tokens, spec.block_size) - clipped_block_ids.append(ids[:num_valid_blocks]) - return tuple(clipped_block_ids) - def cache_blocks(self, request: Request, num_computed_tokens: int) -> None: """Cache the blocks for the request, if enabled. @@ -601,28 +565,6 @@ def take_new_block_ids(self) -> list[int]: ids.extend(mgr.take_new_block_ids()) return ids - def take_boundary_state_offloads( - self, - ) -> dict[str, list[tuple[int, int, int]]]: - """Drain exact Mamba ``align`` state blocks for KV connectors. - - Returns ``{request_id: [(group_id, block_id, boundary_tokens), ...]}``. - These scheduler-local handoffs avoid reconstructing mutable Mamba block - tables from a connector's append-only block-ID mirror. - """ - offloads: dict[str, list[tuple[int, int, int]]] = {} - for mgr in self.coordinator.single_type_managers: - for ( - request_id, - group_id, - block, - boundary_tokens, - ) in mgr.take_pending_boundary_state_offloads(): - offloads.setdefault(request_id, []).append( - (group_id, block.block_id, boundary_tokens) - ) - return offloads - def new_step_starts(self) -> None: """Called when a new step is started.""" self.coordinator.new_step_starts() diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 9f1f447340..7f3a5e4fdf 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -611,12 +611,12 @@ def resolve_kv_cache_block_sizes( if not (cache_config.enable_prefix_caching or connector_enabled): return scheduler_block_size, scheduler_block_size - # Mamba groups outside align mode break divisibility; back off to the - # scheduler block size. Read the mode from the resolved group spec because - # its block size may have been updated independently of cache_config. + # Mamba groups with block_size != cache_config.block_size + # (mamba_cache_mode != "align") break divisibility; back off to the + # scheduler block size. if any( isinstance(g.kv_cache_spec, MambaSpec) - and g.kv_cache_spec.mamba_cache_mode != "align" + and g.kv_cache_spec.block_size != cache_config.block_size for g in groups ): return scheduler_block_size, scheduler_block_size @@ -1043,23 +1043,8 @@ def unify_kv_cache_spec_page_size( ) ratio = max_page_size // layer_page_size new_block_size = layer_spec.block_size * ratio - replace_args = {"block_size": new_block_size} - # A padded page does not grow when only block_size changes. This - # happens for hybrid Mamba targets when a higher-precision draft - # cache has a larger page than the target cache. Keep the logical - # block-size adjustment and grow the physical padding with it. - if ( - isinstance(layer_spec, MambaSpec) - or getattr(layer_spec, "page_size_padded", None) is not None - ): - replace_args["page_size_padded"] = max_page_size - new_spec = replace(layer_spec, **replace_args) - assert new_spec.page_size_bytes == max_page_size, ( - f"Failed to unify KV page for {layer_name}: " - f"spec={layer_spec!r}, old_page={layer_page_size}, " - f"target_page={max_page_size}, replacement={replace_args!r}, " - f"new_page={new_spec.page_size_bytes}" - ) + new_spec = replace(layer_spec, block_size=new_block_size) + assert new_spec.page_size_bytes == max_page_size new_kv_cache_spec[layer_name] = new_spec return new_kv_cache_spec @@ -1470,36 +1455,42 @@ def group_and_unify_kv_cache_specs( return [mla_uniform_spec, *swa_uniform_specs] -def _select_deepseek_v4_tuple_width( - vllm_config: VllmConfig, - grouped_specs: list[UniformTypeKVCacheSpecs], -) -> int: - """Pick the DeepSeek V4 tuple width with the smallest physical pool. - - Padding layer counts alone is insufficient: full/compressed MLA may need - thousands more pages than a short sliding-window group at long context. - Widening every physical block to save one SWA subgroup can therefore cost - substantially more memory than it saves. - """ - tuple_counts = [spec.get_num_layer_tuples() for spec in grouped_specs] - page_counts = [spec.max_memory_usage_pages(vllm_config) for spec in grouped_specs] - min_width = tuple_counts[0] - max_width = max(tuple_counts) - - def pool_pages(width: int) -> int: - return width * sum( - cdiv(tuple_count, width) * page_count - for tuple_count, page_count in zip(tuple_counts, page_counts) - ) +def _approximate_gcd(values: Sequence[int], *, lower_bound: int | None = None) -> int: + """Pick a chunk size that minimizes total upward padding. - return min( - range(min_width, max_width + 1), - key=lambda width: (pool_pages(width), width), - ) + Each x is rounded up to a multiple of d: + + x -> ceil(x / d) * d + + Total padding is: + + pad(d) = sum_i (ceil(x_i / d) * d - x_i) + + We brute-force d in [lower_bound, max(values)] (fine for small lists / small + maxima) and return the d with minimum padding. Ties prefer larger d. + """ + if not values: + raise ValueError("values must be non-empty") + if any(x <= 0 for x in values): + raise ValueError(f"values must be positive, got: {list(values)!r}") + + min_d = max(1, lower_bound if lower_bound is not None else 1) + max_d = max(values) + if min_d > max_d: + return min_d + + best_d = min_d + best_pad: int | None = None + for d in range(min_d, max_d + 1): + pad = sum((d - (x % d)) % d for x in values) + if best_pad is None or pad < best_pad or (pad == best_pad and d > best_d): + best_pad = pad + best_d = d + + return best_d def _get_kv_cache_groups_uniform_groups( - vllm_config: VllmConfig, grouped_specs: list[UniformTypeKVCacheSpecs], ) -> list[KVCacheGroupSpec]: """ @@ -1530,8 +1521,10 @@ def _get_kv_cache_groups_uniform_groups( num_layer_tuples_per_group: list[int] = [ g_spec.get_num_layer_tuples() for g_spec in grouped_specs ] - # Choose `num_layer_tuples` by physical bytes, not layer-count padding. - num_layer_tuples = _select_deepseek_v4_tuple_width(vllm_config, grouped_specs) + # Choose `num_layer_tuples` to minimize total padding across groups. + num_layer_tuples = _approximate_gcd( + num_layer_tuples_per_group, lower_bound=num_layer_tuples_per_group[0] + ) # Round up to the nearest multiple of `num_layer_tuples` (i.e., padding) num_layer_tuples_per_group = [ round_up(x, num_layer_tuples) for x in num_layer_tuples_per_group @@ -1658,9 +1651,7 @@ def get_kv_cache_groups( # yet some layers are full attention while others are sliding window # attention in different sizes. Need to group layers into multiple # UniformTypeKVCacheSpecs. - kv_cache_groups = _get_kv_cache_groups_uniform_groups( - vllm_config, grouped_specs - ) + kv_cache_groups = _get_kv_cache_groups_uniform_groups(grouped_specs) _annotate_eagle_groups_deepseek_v4(vllm_config, kv_cache_spec, kv_cache_groups) return kv_cache_groups diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py index 88da600d24..c153d1103c 100644 --- a/vllm/v1/core/sched/output.py +++ b/vllm/v1/core/sched/output.py @@ -246,11 +246,6 @@ class SchedulerOutput: # preventing stale NaN/data from corrupting attention or SSM computation. new_block_ids_to_zero: list[int] | None = None - # Exact Mamba ``align`` boundary-state sources offered to a KV connector: - # {request_id: [(group_id, block_id, boundary_tokens), ...]}. - # Scheduler-local; cleared after connector metadata is built. - boundary_state_offloads: dict[str, list[tuple[int, int, int]]] | None = None - @classmethod def make_empty(cls) -> "SchedulerOutput": return cls( diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 86b5a1aeff..7fa7757ea7 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -247,10 +247,25 @@ def __init__( ) self.use_eagle = False - self.num_spec_tokens = vllm_config.num_speculative_tokens - self.num_lookahead_tokens = vllm_config.num_lookahead_tokens - if speculative_config and speculative_config.use_eagle_kv_cache(): - self.use_eagle = True + self.num_spec_tokens = self.num_lookahead_tokens = 0 + if speculative_config: + self.num_spec_tokens = speculative_config.num_speculative_tokens + if speculative_config.use_eagle(): + self.use_eagle = True + self.num_lookahead_tokens = self.num_spec_tokens + if speculative_config.uses_draft_model(): + self.num_lookahead_tokens = self.num_spec_tokens + if speculative_config.use_dflash(): + # DFlash requires an extra lookahead slot since it uses in-fill-style + # decoding instead of standard next-token sampling, so it has a query + # for the last sampled token plus queries for each draft token. + self.num_lookahead_tokens = ( + speculative_config.num_speculative_state_tokens() + 1 + ) + if speculative_config.use_dspark(): + # The anchor itself is the first prediction position, with no + # separate bonus query, so DSpark needs exactly N slots. + self.num_lookahead_tokens = self.num_spec_tokens # Create the KV cache manager. if hash_block_size is None: @@ -258,7 +273,7 @@ def __init__( self.kv_cache_manager = KVCacheManager( kv_cache_config=kv_cache_config, max_model_len=self.max_model_len, - max_in_flight_tokens=vllm_config.max_in_flight_tokens, + max_num_batched_tokens=self.scheduler_config.max_num_batched_tokens, enable_caching=self.cache_config.enable_prefix_caching, use_eagle=self.use_eagle, log_stats=self.log_stats, @@ -400,6 +415,9 @@ def _mamba_block_aligned_split( num_new_local_computed_tokens: int = 0, num_external_computed_tokens: int = 0, ) -> int: + assert num_external_computed_tokens == 0, ( + "External KV connector is not verified yet" + ) num_computed_tokens = ( request.num_computed_tokens + num_new_local_computed_tokens @@ -905,8 +923,7 @@ def schedule(self) -> SchedulerOutput: # The request cannot be scheduled. break - # Skip block alignment when setting up async receive (no local work). - if self.need_mamba_block_aligned_split and not load_kv_async: + if self.need_mamba_block_aligned_split: num_new_tokens = self._mamba_block_aligned_split( request, num_new_tokens, @@ -1103,9 +1120,6 @@ def schedule(self) -> SchedulerOutput: if self.needs_kv_cache_zeroing else None ) - # A Mamba ``align`` table is not append-only. Drain exact committed - # boundary-state block IDs before the connector builds store jobs. - boundary_state_offloads = self.kv_cache_manager.take_boundary_state_offloads() scheduler_output = SchedulerOutput( scheduled_new_reqs=new_reqs_data, @@ -1124,9 +1138,6 @@ def schedule(self) -> SchedulerOutput: finished_req_ids=self.finished_req_ids, free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(), new_block_ids_to_zero=new_block_ids_to_zero, - boundary_state_offloads=( - boundary_state_offloads if self.connector is not None else None - ), ) # NOTE(Kuntai): this function is designed for multiple purposes: @@ -1137,9 +1148,6 @@ def schedule(self) -> SchedulerOutput: meta = self._build_kv_connector_meta(self.connector, scheduler_output) scheduler_output.kv_connector_metadata = meta - # Exact source IDs are scheduler-side connector inputs only. - scheduler_output.boundary_state_offloads = None - # Build the connector meta for ECConnector if self.ec_connector is not None: ec_meta: ECConnectorMetadata = self.ec_connector.build_connector_meta( @@ -1192,7 +1200,6 @@ def _update_after_schedule(self, scheduler_output: SchedulerOutput) -> None: for req_id, num_scheduled_token in num_scheduled_tokens.items(): request = self.requests[req_id] request.num_computed_tokens += num_scheduled_token - request.num_in_flight_tokens += num_scheduled_token request.is_prefill_chunk = request.num_computed_tokens < ( request.num_tokens + request.num_output_placeholders ) @@ -1577,12 +1584,10 @@ def update_from_output( stopped_preempted_reqs: set[Request] = set() for req_id, num_tokens_scheduled in num_scheduled_tokens.items(): assert num_tokens_scheduled > 0 - request = self.requests.get(req_id) - if request is not None: - request.num_in_flight_tokens -= num_tokens_scheduled if failed_kv_load_req_ids and req_id in failed_kv_load_req_ids: # skip failed or rescheduled requests from KV load failure continue + request = self.requests.get(req_id) if request is None or request.is_finished(): # The request is already finished. This can happen if the # request is aborted while the model is executing it (e.g., @@ -2386,19 +2391,13 @@ def _connector_finished( return False, None # Free any out-of-window prefix blocks before we hand the block table to - # the connector, on the processed-token basis (see `allocate_slots`). + # the connector. self.kv_cache_manager.remove_skipped_blocks( request_id=request.request_id, - processed_computed_tokens=max( - 0, request.num_computed_tokens - request.num_in_flight_tokens - ), - num_prompt_tokens=request.num_prompt_tokens, + total_computed_tokens=request.num_computed_tokens, ) - block_ids = self.kv_cache_manager.get_block_ids_for_computed_tokens( - request_id=request.request_id, - num_computed_tokens=request.num_computed_tokens, - ) + block_ids = self.kv_cache_manager.get_block_ids(request.request_id) if not isinstance(self.connector, SupportsHMA): # NOTE(Kuntai): We should deprecate this code path after we enforce diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index b3980c592f..cd000dc849 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -20,7 +20,6 @@ KVCacheSpec, MambaSpec, MLAAttentionSpec, - PrefixAnchoredSWASpec, SinkFullAttentionSpec, SlidingWindowMLASpec, SlidingWindowSpec, @@ -69,14 +68,6 @@ def __init__( self._max_admission_blocks_per_request = max_admission_blocks_per_request self.new_block_ids: list[int] = [] - # Mamba ``align`` block tables are sparse and mutable: committed - # boundary states can be followed by speculative scratch blocks that - # are relocated in place. A connector therefore cannot recover an - # exact boundary-state source from an append-only block-ID mirror. - self._pending_boundary_state_offloads: list[ - tuple[str, int, KVCacheBlock, int] - ] = [] - # Mapping from request ID to blocks to track the blocks allocated # for each request, so that we can free the blocks when the request # is finished. @@ -91,18 +82,6 @@ def __init__( self.kv_cache_group_id = kv_cache_group_id self._null_block = block_pool.null_block - def take_pending_boundary_state_offloads( - self, - ) -> list[tuple[str, int, KVCacheBlock, int]]: - """Drain exact Mamba ``align`` boundary-state handoffs. - - Entries are ``(request_id, group_id, block, boundary_tokens)``. - Non-Mamba managers leave the list empty. - """ - pending = self._pending_boundary_state_offloads - self._pending_boundary_state_offloads = [] - return pending - @classmethod def _get_num_evictable_blocks(cls, blocks: Sequence[KVCacheBlock]): return sum(blk.ref_cnt == 0 and not blk.is_null for blk in blocks) @@ -446,41 +425,8 @@ def find_longest_cache_hit( raise NotImplementedError - def _remove_blocks_in_range( - self, - request_id: str, - first_block: int, - last_block: int, - ) -> None: - """Free blocks in ``[first_block, last_block)`` and replace with null_block. - - Iterates backward so newly-evictable tail blocks are reached even after - earlier blocks in the range were nulled in a prior call. - """ - if request_id not in self.req_to_blocks: - return - if first_block >= last_block: - return - blocks = self.req_to_blocks[request_id] - last_block = min(last_block, len(blocks)) - - freed: list[KVCacheBlock] = [] - for i in range(last_block - 1, first_block - 1, -1): - if blocks[i] == self._null_block: - # If the block is already a null block, the blocks before it - # should also have been set to null blocks by the previous calls - # to this function. - break - freed.append(blocks[i]) - blocks[i] = self._null_block - if freed: - self.block_pool.free_blocks(freed) - def remove_skipped_blocks( - self, - request_id: str, - processed_computed_tokens: int, - num_prompt_tokens: int | None = None, + self, request_id: str, total_computed_tokens: int ) -> None: """ Remove and free the blocks that are no longer needed for attention computation. @@ -491,27 +437,36 @@ def remove_skipped_blocks( Args: request_id: The request ID. - processed_computed_tokens: Computed-token prefix length covering - fully processed and committed tokens only (safe to free). - num_prompt_tokens: Optional prompt length for attention types (e.g. - prefix-anchored SWA) that evict a middle gap rather than a head - prefix. Ignored by the default implementation. + total_computed_tokens: The total number of computed tokens, including + local computed tokens and external computed tokens. """ - del num_prompt_tokens # Remove the blocks that will be skipped during attention computation. - num_skipped_tokens = self.get_num_skipped_tokens(processed_computed_tokens) + num_skipped_tokens = self.get_num_skipped_tokens(total_computed_tokens) if num_skipped_tokens <= 0: # This indicates that ALL tokens are inside attention window. # Thus we do not need to free any blocks outside attention window. # A typical case is full attention that we never free any token # before the request is finished. return + blocks = self.req_to_blocks[request_id] num_skipped_blocks = num_skipped_tokens // self.block_size # `num_skipped_tokens` may include tokens that haven't been allocated yet # (e.g., when the attention window moves into the external computed tokens - # range); `_remove_blocks_in_range` caps to the number of blocks that - # currently exist for this request. - self._remove_blocks_in_range(request_id, 0, num_skipped_blocks) + # range), so we must cap to the number of blocks that currently exist for + # this request. + num_skipped_blocks = min(num_skipped_blocks, len(blocks)) + removed_blocks: list[KVCacheBlock] = [] + # Because the block starts from index 0, the num_skipped_block-th block + # corresponds to index num_skipped_blocks - 1. + for i in range(num_skipped_blocks - 1, -1, -1): + if blocks[i] == self._null_block: + # If the block is already a null block, the blocks before it + # should also have been set to null blocks by the previous calls + # to this function. + break + removed_blocks.append(blocks[i]) + blocks[i] = self._null_block + self.block_pool.free_blocks(removed_blocks) def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: """ @@ -592,53 +547,6 @@ def get_num_common_prefix_blocks(self, running_request_id: str) -> int: return num_common_blocks -class PrefixAnchoredSWAManager(FullAttentionManager): - """KV cache manager for prefix-anchored sliding-window attention. - - When ``num_prompt_tokens`` is supplied to ``remove_skipped_blocks``, frees - gap blocks between the prefill tail and the current decode window. This - bounds per-request KV memory at O(prefix_len + decode_sliding_window) - instead of growing linearly with decode length. - """ - - def __init__(self, kv_cache_spec: PrefixAnchoredSWASpec, **kwargs) -> None: - super().__init__(kv_cache_spec, **kwargs) - self.decode_sliding_window: int = kv_cache_spec.decode_sliding_window - - def remove_skipped_blocks( - self, - request_id: str, - processed_computed_tokens: int, - num_prompt_tokens: int | None = None, - ) -> None: - """Free gap blocks that are no longer needed for attention. - - Gap = blocks entirely within - [ceil(prefix_len / block_size) * block_size, - max(prefix_len, processed_computed_tokens - decode_sliding_window)) - - Freed blocks are replaced with null_block in req_to_blocks so the - block_table handed to the attention backend stays valid (null_block - KV is all-zero; the backend mask marks gap positions as non-visible - so they never contribute to the output). - """ - if num_prompt_tokens is None: - super().remove_skipped_blocks( - request_id, processed_computed_tokens, num_prompt_tokens - ) - return - - bs = self.block_size - # First block fully after the prefill boundary. - first_gap_block = cdiv(num_prompt_tokens, bs) - # Decode window start position; blocks before this are evictable. - window_start = max( - num_prompt_tokens, processed_computed_tokens - self.decode_sliding_window - ) - last_gap_block = window_start // bs # exclusive upper bound - self._remove_blocks_in_range(request_id, first_gap_block, last_gap_block) - - class SlidingWindowManager(SingleTypeKVCacheManager): def __init__(self, kv_cache_spec: SlidingWindowSpec, **kwargs) -> None: super().__init__(kv_cache_spec, **kwargs) @@ -1002,17 +910,17 @@ def find_longest_cache_hit( return computed_blocks - def remove_skipped_blocks( - self, - request_id: str, - processed_computed_tokens: int, - num_prompt_tokens: int | None = None, - ) -> None: + def remove_skipped_blocks(self, request_id: str, num_computed_tokens: int) -> None: assert isinstance(self.kv_cache_spec, MambaSpec) - super().remove_skipped_blocks( - request_id, processed_computed_tokens, num_prompt_tokens - ) + # NOTE (tdoublep) with async scheduling, the num_computed_tokens can contain + # draft tokens from the previous step that may or may not be rejected later. + # This can make us think we are further ahead in the sequence than we actually + # are, so let's assume that all tokens are rejected so we don't free blocks + # that we might actually need. + num_computed_tokens = max(0, num_computed_tokens - self.num_speculative_blocks) + + super().remove_skipped_blocks(request_id, num_computed_tokens) if self.mamba_cache_mode == "align": # `last_state_block_idx` refers to the block index allocated two steps ago. # The block allocated in the previous step is used to copy Mamba states @@ -1025,7 +933,7 @@ def remove_skipped_blocks( if ( last_state_block_idx is not None and last_state_block_idx - < cdiv(processed_computed_tokens, self.block_size) - 1 + < cdiv(num_computed_tokens, self.block_size) - 1 ): blocks = self.req_to_blocks[request_id] if blocks[last_state_block_idx] != self._null_block: @@ -1187,14 +1095,6 @@ def free(self, request_id: str) -> None: if self.mamba_cache_mode == "align": self._allocated_block_reqs.discard(request_id) self.last_state_block_idx.pop(request_id, None) - # A hand-off is valid only while its source still belongs to this - # request. Drop offers that have not reached the connector before - # returning the request's blocks to the pool. - self._pending_boundary_state_offloads = [ - entry - for entry in self._pending_boundary_state_offloads - if entry[0] != request_id - ] super().free(request_id) def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: @@ -1215,22 +1115,13 @@ def cache_blocks( super().cache_blocks(request, num_tokens, alignment_tokens=alignment_tokens) num_cached_blocks_after = self.num_cached_block.get(request.request_id, 0) if num_cached_blocks_after > num_cached_blocks_before: - blocks = self.req_to_blocks[request.request_id] - for block_idx in range(num_cached_blocks_before, num_cached_blocks_after): - block = blocks[block_idx] + for block in self.req_to_blocks[request.request_id][ + num_cached_blocks_before:num_cached_blocks_after + ]: if block.is_null: continue assert block.block_hash is not None self.cached_blocks_this_step.add(block.block_hash) - if self.mamba_cache_mode == "align": - self._pending_boundary_state_offloads.append( - ( - request.request_id, - self.kv_cache_group_id, - block, - (block_idx + 1) * self.block_size, - ) - ) def new_step_starts(self) -> None: self.cached_blocks_this_step.clear() @@ -1319,7 +1210,6 @@ def __init__( TQFullAttentionSpec: FullAttentionManager, MLAAttentionSpec: FullAttentionManager, HiddenStateCacheSpec: FullAttentionManager, - PrefixAnchoredSWASpec: PrefixAnchoredSWAManager, SlidingWindowSpec: SlidingWindowManager, SlidingWindowMLASpec: SlidingWindowManager, ChunkedLocalAttentionSpec: ChunkedLocalAttentionManager, @@ -1331,7 +1221,7 @@ def __init__( def get_manager_for_kv_cache_spec( kv_cache_spec: KVCacheSpec, - max_in_flight_tokens: int, + max_num_batched_tokens: int, max_model_len: int, **kwargs, ) -> SingleTypeKVCacheManager: @@ -1342,7 +1232,7 @@ def get_manager_for_kv_cache_spec( if isinstance(kv_cache_spec, (SlidingWindowSpec, ChunkedLocalAttentionSpec)): kwargs["max_admission_blocks_per_request"] = ( kv_cache_spec.max_admission_blocks_per_request( - max_in_flight_tokens=max_in_flight_tokens, + max_num_batched_tokens=max_num_batched_tokens, max_model_len=max_model_len, ) ) diff --git a/vllm/v1/cudagraph_dispatcher.py b/vllm/v1/cudagraph_dispatcher.py index b50f362594..862c3f2855 100644 --- a/vllm/v1/cudagraph_dispatcher.py +++ b/vllm/v1/cudagraph_dispatcher.py @@ -7,21 +7,13 @@ from vllm import envs from vllm.config import CUDAGraphMode, VllmConfig -from vllm.forward_context import ( - CUDAGRAPH_VARIANT_DEFAULT, - CUDAGRAPH_VARIANT_LONG_CONTEXT, - BatchDescriptor, -) +from vllm.forward_context import BatchDescriptor from vllm.logger import init_logger from vllm.lora.utils import get_captured_lora_counts from vllm.platforms import current_platform logger = init_logger(__name__) -_SM70_FP8_KV_BATCH_CONTEXT_SIZES = frozenset((4, 8, 16)) -_SM70_FP8_KV_LONG_CONTEXT_MIN_SEQ_LEN = 16384 -_SM70_E4M3_B1_WAVE_LONG_CONTEXT_MIN_SEQ_LEN = 49152 - def _get_sm70_context_buckets(env_name: str) -> tuple[int, ...]: """Parse optional one-request attention graph context buckets.""" @@ -47,13 +39,15 @@ def _get_sm70_dsv4_decode_context_buckets( env_name = "VLLM_SM70_DSV4_DECODE_CONTEXT_BUCKETS" if env_name in os.environ: return _get_sm70_context_buckets(env_name) - # Speculative graphs cover multiple query rows and require a separate - # end-to-end long-context gate. They remain available through the explicit - # environment override above. if vllm_config.speculative_config is not None: return () model_config = vllm_config.model_config + architectures = getattr(model_config, "architectures", ()) + if not isinstance(architectures, (list, tuple)) or ( + "DeepseekV4ForCausalLM" not in architectures + ): + return () if not ( current_platform.is_cuda() and current_platform.is_device_capability((7, 0)) ): @@ -72,88 +66,64 @@ def _get_sm70_dsv4_decode_context_buckets( if not positive_ratios: return () - short_bucket = topk_tokens * min(positive_ratios) - if short_bucket <= 0 or model_config.max_model_len <= short_bucket: + bucket = topk_tokens * min(positive_ratios) + if bucket <= 0 or model_config.max_model_len <= bucket: return () - # Keep the graph width close enough to the live context that the C4 - # indexer does not score the model's entire maximum length immediately - # after crossing the short-context bypass. The sparse progression limits - # startup graph count while covering the first long-indexer bucket and the - # 16K/64K/128K service points used by the long-context quality gates. - bucket_multipliers = (1, 2, 8, 32, 64) - buckets = tuple( - short_bucket * multiplier - for multiplier in bucket_multipliers - if short_bucket * multiplier < model_config.max_model_len - ) logger.info_once( - "Auto-enabling SM70 compressed-index decode CUDA graph context buckets " - "%s. Set %s explicitly to override or to an empty value to disable.", - buckets, + "Auto-enabling the SM70 DeepSeek V4 short-context decode CUDA graph " + "at %d tokens. Set %s explicitly to override or to an empty value " + "to disable.", + bucket, env_name, ) - return buckets + return (bucket,) -def _is_sm70_fp8_kv_decode_shape( +def _get_sm70_fp8_kv_decode_context_buckets( vllm_config: VllmConfig, - *, - q_per_kv_values: tuple[int, ...], - cache_dtypes: tuple[str, ...] = ("fp8_e5m2",), - backend_names: tuple[str, ...] = ( - "FLASH_ATTN_V100", - "FLASHINFER_SM70", - ), -) -> bool: +) -> tuple[int, ...]: + """Select a scalar-only short graph for the SM70 E5M2 D256 route.""" + env_name = "VLLM_SM70_FP8_KV_DECODE_CONTEXT_BUCKETS" + if env_name in os.environ: + return _get_sm70_context_buckets(env_name) if vllm_config.speculative_config is not None: - return False + return () if not ( current_platform.is_cuda() and current_platform.is_device_capability((7, 0)) ): - return False + return () if not envs.VLLM_SM70_FLASH_ATTN_V100: - return False + return () cache_config = getattr(vllm_config, "cache_config", None) - if getattr(cache_config, "cache_dtype", None) not in cache_dtypes: - return False + if getattr(cache_config, "cache_dtype", None) != "fp8_e5m2": + return () attention_config = getattr(vllm_config, "attention_config", None) attention_backend = getattr(attention_config, "backend", None) attention_backend_name = getattr(attention_backend, "name", attention_backend) - if attention_backend is not None and attention_backend_name not in backend_names: - return False + if attention_backend is not None and attention_backend_name not in ( + "FLASH_ATTN_V100", + "FLASHINFER_SM70", + ): + return () model_config = vllm_config.model_config hf_text_config = getattr(model_config, "hf_text_config", None) num_attention_heads = getattr(hf_text_config, "num_attention_heads", None) num_key_value_heads = getattr(hf_text_config, "num_key_value_heads", None) head_dim = getattr(hf_text_config, "head_dim", None) - return bool( + if not ( isinstance(num_attention_heads, int) and isinstance(num_key_value_heads, int) and num_key_value_heads > 0 and num_attention_heads % num_key_value_heads == 0 - and num_attention_heads // num_key_value_heads in q_per_kv_values + and num_attention_heads // num_key_value_heads in (6, 8) and head_dim == 256 - ) - - -def _get_sm70_fp8_kv_decode_context_buckets( - vllm_config: VllmConfig, -) -> tuple[int, ...]: - """Select a scalar-only B1 graph for the SM70 E5M2 D256 route.""" - env_name = "VLLM_SM70_FP8_KV_DECODE_CONTEXT_BUCKETS" - if env_name in os.environ: - return _get_sm70_context_buckets(env_name) - if not _is_sm70_fp8_kv_decode_shape( - vllm_config, - q_per_kv_values=(6, 8), ): return () bucket = 8192 - model_config = vllm_config.model_config max_model_len = getattr(model_config, "max_model_len", None) if not isinstance(max_model_len, int) or max_model_len <= bucket: return () @@ -168,35 +138,6 @@ def _get_sm70_fp8_kv_decode_context_buckets( return (bucket,) -def _sm70_fp8_kv_batch_context_routing_enabled( - vllm_config: VllmConfig, -) -> bool: - if ( - not envs.VLLM_FLASH_V100_XQA_BATCH_CONTEXT_ROUTING - or envs.VLLM_FLASH_V100_DECODE_PARTITION_SIZE is not None - ): - return False - cache_dtype = getattr(vllm_config.cache_config, "cache_dtype", None) - cache_dtypes: tuple[str, ...] - backend_names: tuple[str, ...] - if cache_dtype in ("fp8", "fp8_e4m3"): - if not envs.VLLM_FLASH_V100_E4M3_BATCH_XQA: - return False - cache_dtypes = ("fp8", "fp8_e4m3") - # E4M3 batch-context variants exist to select Flash-V100 XQA launch - # routes. An explicit FlashInfer backend cannot consume that metadata. - backend_names = ("FLASH_ATTN_V100",) - else: - cache_dtypes = ("fp8_e5m2",) - backend_names = ("FLASH_ATTN_V100", "FLASHINFER_SM70") - return _is_sm70_fp8_kv_decode_shape( - vllm_config, - q_per_kv_values=(6,), - cache_dtypes=cache_dtypes, - backend_names=backend_names, - ) - - class CudagraphDispatcher: """ Runtime cudagraph dispatcher to dispatch keys for multiple set of @@ -233,27 +174,7 @@ def __init__(self, vllm_config: VllmConfig): self.sm70_fp8_kv_decode_context_buckets = ( _get_sm70_fp8_kv_decode_context_buckets(vllm_config) ) - self.sm70_fp8_kv_batch_context_routing = ( - _sm70_fp8_kv_batch_context_routing_enabled(vllm_config) - ) - cache_dtype = getattr(vllm_config.cache_config, "cache_dtype", None) - self.sm70_e4m3_b1_wave_context_routing = bool( - self.sm70_fp8_kv_batch_context_routing - and cache_dtype in ("fp8", "fp8_e4m3") - and os.getenv("VLLM_FLASH_V100_XQA_E4M3_G6_P64_P256_AUTO", "0") == "1" - and os.getenv("VLLM_FLASH_V100_XQA_E4M3_G6_WAVE_PARTITIONS", "0") == "1" - ) - self.sm70_e4m3_b1_wave_context_min_seq_len = max( - 1, - int( - os.getenv( - "VLLM_FLASH_V100_XQA_E4M3_G6_P512_BEGIN", - str(_SM70_E4M3_B1_WAVE_LONG_CONTEXT_MIN_SEQ_LEN), - ) - ), - ) self._logged_sm70_context_bucket = False - self._logged_sm70_batch_context_variant = False # Dict to store valid cudagraph dispatching keys. self.cudagraph_keys: dict[CUDAGraphMode, set[BatchDescriptor]] = { @@ -388,9 +309,7 @@ def add_cudagraph_key( def _active_context_buckets(self) -> tuple[int, ...]: if self.uniform_decode_query_len > 1: - if "VLLM_SM70_MTP_CONTEXT_BUCKETS" in os.environ: - return self.sm70_mtp_context_buckets - return self.sm70_dsv4_decode_context_buckets + return self.sm70_mtp_context_buckets return tuple( sorted( set(self.sm70_dsv4_decode_context_buckets) @@ -398,102 +317,39 @@ def _active_context_buckets(self) -> tuple[int, ...]: ) ) - def _is_batch_context_variant_descriptor( - self, - batch_descriptor: BatchDescriptor, - ) -> bool: - return bool( - self.sm70_fp8_kv_batch_context_routing - and batch_descriptor.uniform - and batch_descriptor.num_reqs is not None - and batch_descriptor.num_tokens == batch_descriptor.num_reqs - and ( - batch_descriptor.num_reqs in _SM70_FP8_KV_BATCH_CONTEXT_SIZES - or ( - self.sm70_e4m3_b1_wave_context_routing - and batch_descriptor.num_reqs == 1 - ) - ) - ) - - def _batch_context_variant_min_seq_len( - self, - batch_descriptor: BatchDescriptor, - ) -> int: - if self.sm70_e4m3_b1_wave_context_routing and batch_descriptor.num_reqs == 1: - return self.sm70_e4m3_b1_wave_context_min_seq_len - return _SM70_FP8_KV_LONG_CONTEXT_MIN_SEQ_LEN - - def _context_buckets_for_descriptor( - self, - batch_descriptor: BatchDescriptor, - ) -> tuple[int, ...]: - if not batch_descriptor.uniform or batch_descriptor.num_reqs is None: - return () - - if self.uniform_decode_query_len > 1: - if batch_descriptor.num_tokens == self.uniform_decode_query_len: - if "VLLM_SM70_MTP_CONTEXT_BUCKETS" in os.environ: - return self.sm70_mtp_context_buckets - return self.sm70_dsv4_decode_context_buckets - return () - - buckets: set[int] = set() - if batch_descriptor.num_tokens == 1: - buckets.update(self.sm70_dsv4_decode_context_buckets) - buckets.update(self.sm70_fp8_kv_decode_context_buckets) - - return tuple(sorted(buckets)) - @property def has_attention_context_buckets(self) -> bool: return bool(self._active_context_buckets()) - @property - def has_attention_context_specialization(self) -> bool: - return bool( - self.has_attention_context_buckets or self.sm70_fp8_kv_batch_context_routing - ) - def _add_context_bucket_keys(self, batch_descriptor: BatchDescriptor) -> None: - """Add attention-specialized graphs for this exact descriptor.""" - context_buckets = self._context_buckets_for_descriptor(batch_descriptor) + """Add bounded graphs only for a single uniform-decode request.""" + context_buckets = self._active_context_buckets() + if ( + not context_buckets + or not batch_descriptor.uniform + or batch_descriptor.num_tokens != self.uniform_decode_query_len + ): + return + for bucket in context_buckets: self.add_cudagraph_key( CUDAGraphMode.FULL, replace(batch_descriptor, attention_context_bucket=bucket), ) - if self._is_batch_context_variant_descriptor(batch_descriptor): - self.add_cudagraph_key( - CUDAGraphMode.FULL, - replace( - batch_descriptor, - graph_variant=CUDAGRAPH_VARIANT_LONG_CONTEXT, - ), - ) def _get_context_bucket_descriptor( self, batch_descriptor: BatchDescriptor, attention_context_len: int | None, ) -> BatchDescriptor: - """Return the graph specialization selected by the active context.""" - if ( - attention_context_len is not None - and attention_context_len - >= self._batch_context_variant_min_seq_len(batch_descriptor) - and self._is_batch_context_variant_descriptor(batch_descriptor) - ): - return replace( - batch_descriptor, - graph_variant=CUDAGRAPH_VARIANT_LONG_CONTEXT, - ) - - context_buckets = self._context_buckets_for_descriptor(batch_descriptor) + """Return a bounded graph key only when its capacity is sufficient.""" + context_buckets = self._active_context_buckets() if ( attention_context_len is None or attention_context_len <= 0 or not context_buckets + or not batch_descriptor.uniform + or batch_descriptor.num_tokens != self.uniform_decode_query_len ): return batch_descriptor @@ -681,17 +537,6 @@ def dispatch( batch_desc_to_check, ) self._logged_sm70_context_bucket = True - if ( - batch_desc_to_check.graph_variant == CUDAGRAPH_VARIANT_LONG_CONTEXT - and not self._logged_sm70_batch_context_variant - ): - logger.info( - "Using SM70 long-context XQA CUDA graph: " - "context=%d descriptor=%s", - attention_context_len, - batch_desc_to_check, - ) - self._logged_sm70_batch_context_variant = True return CUDAGraphMode.FULL, batch_desc_to_check # A bounded graph might be absent because capture-size policy # excluded it. Preserve the existing full-context graph then. @@ -730,32 +575,16 @@ def get_capture_descs(self) -> list[tuple[CUDAGraphMode, list[BatchDescriptor]]] if descs: # Capture the generic full-context graph before smaller # context-specialized graphs so its allocations seed the pool. - if self.sm70_fp8_kv_batch_context_routing: - # Preserve the rollback graph capture order before adding - # long-context variants. This keeps baseline graph memory - # allocation and warmup behavior reproducible. - descs.sort( - key=lambda d: ( - d.graph_variant == CUDAGRAPH_VARIANT_DEFAULT, - d.num_tokens, - d.attention_context_bucket - if d.attention_context_bucket is not None - else float("inf"), - d.num_active_loras, - ), - reverse=True, - ) - else: - descs.sort( - key=lambda d: ( - d.num_tokens, - d.attention_context_bucket - if d.attention_context_bucket is not None - else float("inf"), - d.num_active_loras, - ), - reverse=True, - ) + descs.sort( + key=lambda d: ( + d.num_tokens, + d.attention_context_bucket + if d.attention_context_bucket is not None + else float("inf"), + d.num_active_loras, + ), + reverse=True, + ) result.append((mode, descs)) return result diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index c1f11a69ba..4e094eb6a6 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -461,7 +461,9 @@ 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() @@ -476,7 +478,9 @@ 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: @@ -493,7 +497,9 @@ 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. @@ -605,7 +611,9 @@ 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..31ee89bc72 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -403,58 +403,12 @@ class HiddenStateCacheSpec(MLAAttentionSpec): pass -@dataclass(frozen=True, kw_only=True) -class PrefixAnchoredSWASpec(FullAttentionSpec): - """KV cache spec for prefix-anchored sliding-window attention. - - Prefill (prompt/prefix) tokens are always globally visible. - Only the last ``decode_sliding_window`` generated tokens are kept in the - KV cache; gap blocks (between the prefill tail and the current decode - window) are evicted during each decode step to bound per-request memory - at O(prefix_blocks + window_blocks). - """ - - decode_sliding_window: int - - def __post_init__(self) -> None: - super().__post_init__() - if self.decode_sliding_window <= 0: - raise ValueError("decode_sliding_window must be greater than zero") - - @classmethod - def merge(cls, specs: list[Self]) -> Self: - assert all(isinstance(spec, PrefixAnchoredSWASpec) for spec in specs), ( - "All attention layers in the same KV cache group must be " - "PrefixAnchoredSWASpec." - ) - windows = {spec.decode_sliding_window for spec in specs} - assert len(windows) == 1, ( - "All prefix-anchored SWA layers must share the same " - f"decode_sliding_window, got {windows}" - ) - # Delegate common field merging to the parent, then reattach the - # decode window. - base = FullAttentionSpec.merge(specs) # type: ignore[arg-type] - return cls( - block_size=base.block_size, - num_kv_heads=base.num_kv_heads, - head_size=base.head_size, - head_size_v=base.head_size_v, - dtype=base.dtype, - kv_quant_mode=base.kv_quant_mode, - page_size_padded=base.page_size_padded, - sliding_window=base.sliding_window, - attention_chunk_size=base.attention_chunk_size, - decode_sliding_window=windows.pop(), - ) - - @dataclass(frozen=True, kw_only=True) class ChunkedLocalAttentionSpec(AttentionSpec): attention_chunk_size: int def max_admission_blocks_per_request( - self, max_in_flight_tokens: int, max_model_len: int + self, max_num_batched_tokens: int, max_model_len: int ) -> int: """Per-request admission cap, in blocks. @@ -464,15 +418,15 @@ def max_admission_blocks_per_request( """ # During chunked prefill, we hold KV for at most one chunk window. num_tokens = min( - self.attention_chunk_size + max_in_flight_tokens, max_model_len + self.attention_chunk_size + max_num_batched_tokens, max_model_len ) return cdiv(num_tokens, self.block_size) def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: max_model_len = vllm_config.model_config.max_model_len - max_in_flight_tokens = vllm_config.max_in_flight_tokens + max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens max_blocks = self.max_admission_blocks_per_request( - max_in_flight_tokens=max_in_flight_tokens, max_model_len=max_model_len + max_num_batched_tokens=max_num_batched_tokens, max_model_len=max_model_len ) return max_blocks * self.page_size_bytes @@ -507,7 +461,7 @@ def real_page_size_bytes(self) -> int: ) def max_admission_blocks_per_request( - self, max_in_flight_tokens: int, max_model_len: int + self, max_num_batched_tokens: int, max_model_len: int ) -> int: """Per-request admission cap, in blocks. @@ -520,7 +474,9 @@ def max_admission_blocks_per_request( # During chunked prefill, we hold KV for the last `sliding_window-1` # computed tokens plus the newly scheduled tokens, and never more # than `max_model_len`. - num_tokens = min(self.sliding_window - 1 + max_in_flight_tokens, max_model_len) + num_tokens = min( + self.sliding_window - 1 + max_num_batched_tokens, max_model_len + ) # +1 because the sliding window may not start from the beginning of # the block. E.g. block size 4 and num_token 4 needs two blocks # [XXCD][EF] to store the 6-token window [CDEF]. @@ -531,9 +487,9 @@ def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: "DCP not support sliding window." ) max_model_len = vllm_config.model_config.max_model_len - max_in_flight_tokens = vllm_config.max_in_flight_tokens + max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens max_blocks = self.max_admission_blocks_per_request( - max_in_flight_tokens=max_in_flight_tokens, max_model_len=max_model_len + max_num_batched_tokens=max_num_batched_tokens, max_model_len=max_model_len ) return max_blocks * self.page_size_bytes diff --git a/vllm/v1/request.py b/vllm/v1/request.py index 095ee6e975..26cc82fc4a 100644 --- a/vllm/v1/request.py +++ b/vllm/v1/request.py @@ -141,11 +141,6 @@ def __init__( self.num_output_placeholders = 0 self.async_tokens_to_discard = 0 - # Tokens of steps whose output is not yet processed (async scheduling - # and PP run ahead of the GPU); `num_computed_tokens` counts them - # optimistically. - self.num_in_flight_tokens = 0 - self.spec_token_ids: list[int] = [] self.num_computed_tokens = 0 self.cache_salt: str | None = cache_salt diff --git a/vllm/v1/sample/ops/topk_topp_triton.py b/vllm/v1/sample/ops/topk_topp_triton.py index 55965bed40..4b17831dfc 100644 --- a/vllm/v1/sample/ops/topk_topp_triton.py +++ b/vllm/v1/sample/ops/topk_topp_triton.py @@ -11,8 +11,6 @@ import torch -from vllm import envs -from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.utils.math_utils import next_power_of_2 from vllm.utils.platform_utils import num_compute_units @@ -20,29 +18,6 @@ _TRITON_TABLE_CACHE: dict[tuple[torch.device], tuple[torch.Tensor, torch.Tensor]] = {} _TRITON_BUFFER_CACHE: dict[tuple[torch.device, torch.dtype, int], torch.Tensor] = {} - -def _use_sm70_topk_topp_8_warps( - device: torch.device, - batch_size: int, - vocab_size: int, - topk_enabled: bool, - topp_enabled: bool, -) -> bool: - if ( - device.type != "cuda" - or not current_platform.is_cuda() - or not current_platform.is_device_capability((7, 0)) - or vocab_size != 248_320 - or not topk_enabled - or not topp_enabled - ): - return False - - if batch_size in (8, 16) and envs.VLLM_SM70_TOPK_TOPP_B8_B16_8_WARPS: - return True - return envs.VLLM_SM70_TOPK_TOPP_8_WARPS and batch_size in (5, 10, 20, 40, 60, 80) - - # fmt: off _NORMAL_CDF_TO_SIGMA_TABLE = [ 3.656, 3.650, 3.650, 3.650, 3.626, 3.626, 3.626, 3.514, 3.514, 3.503, @@ -1069,19 +1044,6 @@ def apply_top_k_top_p_triton( else: block_size, block_size_trunc = 8192, 4096 - launch_kwargs: dict[str, int] = {} - if _use_sm70_topk_topp_8_warps( - logits.device, - batch_size, - vocab_size, - topk_enabled, - topp_enabled, - ): - # Eight warps preserve the tile and masking algorithm while using the - # otherwise idle lanes in its reduction-heavy passes. B8/B16 no-MTP - # rows are default-on; the measured MTP verifier rows remain opt-in. - launch_kwargs["num_warps"] = 8 - _topk_topp_kernel[(NUM_PROGRAMS,)]( logits, logits.stride(0), @@ -1097,7 +1059,6 @@ def apply_top_k_top_p_triton( BLOCK_SIZE_TRUNC=block_size_trunc, TOPK_ENABLED=topk_enabled, TOPP_ENABLED=topp_enabled, - **launch_kwargs, ) return logits diff --git a/vllm/v1/sample/rejection_sampler.py b/vllm/v1/sample/rejection_sampler.py index 7933896f7c..ba05ef5db5 100644 --- a/vllm/v1/sample/rejection_sampler.py +++ b/vllm/v1/sample/rejection_sampler.py @@ -270,7 +270,6 @@ def forward( # [num_tokens + batch_size, vocab_size] logits: torch.Tensor, sampling_metadata: SamplingMetadata, - draft_confidence_logits: torch.Tensor | None = None, ) -> SamplerOutput: """ Args: @@ -332,7 +331,6 @@ def forward( bonus_logits_indices=bonus_logits_indices, profile_events=profile_events, profile_total_start=profile_total_start, - draft_confidence_logits=draft_confidence_logits, ) # When indexing with a tensor (bonus_logits_indices), PyTorch # creates a new tensor with separate storage from the original @@ -428,7 +426,6 @@ def forward( bonus_token_ids=bonus_token_ids, output_token_ids=output_token_ids, sampling_metadata=sampling_metadata, - draft_confidence_logits=draft_confidence_logits, ) logprobs_tensors = None @@ -458,7 +455,6 @@ def _forward_combined_bonus( bonus_logits_indices: torch.Tensor, profile_events: list[tuple[str, torch.cuda.Event, torch.cuda.Event]] | None, profile_total_start: torch.cuda.Event | None, - draft_confidence_logits: torch.Tensor | None, ) -> SamplerOutput: profile_constraints_start = _rejection_profile_start(profile_events) sampled_logits = logits.to(torch.float32) @@ -514,7 +510,6 @@ def _forward_combined_bonus( bonus_token_ids=bonus_token_ids, output_token_ids=output_token_ids, sampling_metadata=sampling_metadata, - draft_confidence_logits=draft_confidence_logits, ) return SamplerOutput( sampled_token_ids=output_token_ids, @@ -579,7 +574,6 @@ def _maybe_dump_alignment( bonus_token_ids: torch.Tensor, output_token_ids: torch.Tensor, sampling_metadata: SamplingMetadata, - draft_confidence_logits: torch.Tensor | None, ) -> None: global _SPEC_ALIGNMENT_DUMP_COUNT, _SPEC_ALIGNMENT_STEP_COUNTER _SPEC_ALIGNMENT_STEP_COUNTER += 1 @@ -657,9 +651,6 @@ def _maybe_dump_alignment( "sampling_output_token_ids_tail": [ list(ids[-32:]) for ids in sampling_metadata.output_token_ids ], - "sampling_output_lengths": [ - len(ids) for ids in sampling_metadata.output_token_ids - ], "sampling_spec_token_ids": [ list(ids) for ids in (sampling_metadata.spec_token_ids or []) ], @@ -681,23 +672,9 @@ def _maybe_dump_alignment( .cpu() ) payload["recovered_residual_mass"] = residual_mass.detach().cpu() - payload["conditional_acceptance_targets"] = ( - torch.minimum(target_probs, draft_probs).sum(dim=-1).detach().cpu() - ) draft_topk = torch.topk(draft_probs, k=k, dim=-1) payload["draft_topk_ids"] = draft_topk.indices.detach().cpu() payload["draft_topk_values"] = draft_topk.values.detach().cpu() - if draft_confidence_logits is not None: - if draft_confidence_logits.shape != (target_logits.shape[0],): - raise RuntimeError( - "Aligned DSpark confidence logits must have one value per " - "draft token: " - f"confidence={tuple(draft_confidence_logits.shape)}, " - f"target={target_logits.shape[0]}." - ) - payload["draft_confidence_logits"] = ( - draft_confidence_logits.detach().float().cpu() - ) _SPEC_ALIGNMENT_DUMP_COUNT += 1 dump_dir = os.getenv("VLLM_SPEC_DUMP_ALIGNMENT_DIR", "/tmp") os.makedirs(dump_dir, exist_ok=True) diff --git a/vllm/v1/sample/sampler.py b/vllm/v1/sample/sampler.py index 8d960ed297..9f565ce458 100644 --- a/vllm/v1/sample/sampler.py +++ b/vllm/v1/sample/sampler.py @@ -8,9 +8,7 @@ import torch import torch.nn as nn -from vllm import envs from vllm.config.model import LogprobsMode -from vllm.logger import init_logger from vllm.utils.platform_utils import is_pin_memory_available from vllm.v1.outputs import LogprobsTensors, SamplerOutput from vllm.v1.sample.metadata import SamplingMetadata @@ -21,7 +19,6 @@ _SAMPLING_EPS = 1e-5 _SM70_LOGITS_DUMP_COUNTER = 0 -logger = init_logger(__name__) def _sm70_cuda_graph_capture_active() -> bool: @@ -80,7 +77,10 @@ def tensor_cpu(value: torch.Tensor | None) -> torch.Tensor | None: ], }, path - / (f"sampler_logits_pid{os.getpid()}_step{_SM70_LOGITS_DUMP_COUNTER:04d}.pt"), + / ( + f"sampler_logits_pid{os.getpid()}" + f"_step{_SM70_LOGITS_DUMP_COUNTER:04d}.pt" + ), ) @@ -303,217 +303,6 @@ def apply_temperature( def greedy_sample(logits: torch.Tensor) -> torch.Tensor: return logits.argmax(dim=-1).view(-1) - @staticmethod - def _try_sm70_compact_topk20_sample( - logits: torch.Tensor, - sampling_metadata: SamplingMetadata, - logprobs_mode: LogprobsMode, - ) -> torch.Tensor | None: - if not Sampler._sm70_compact_topk20_metadata_supported( - sampling_metadata, logprobs_mode - ): - return None - if ( - logits.dtype != torch.float32 - or not logits.is_cuda - or logits.shape != (1, 248320) - or torch.cuda.get_device_capability(logits.device) != (7, 0) - ): - return None - - chunks = envs.VLLM_SM70_CHUNKED_TOPK20_CHUNKS - if ( - chunks > 0 - and 248320 % chunks == 0 - and hasattr(torch.ops._C, "sm70_sample_chunked_top20_philox_token_out") - and len(sampling_metadata.generators) <= 1 - ): - from vllm import _sm70_ops as sm70_ops - - chunk_size = 248320 // chunks - local_values, local_indices = torch.topk( - logits.view(chunks, chunk_size), - k=20, - dim=-1, - largest=True, - sorted=False, - ) - global_values, global_positions = torch.topk( - local_values.view(1, -1), - k=20, - dim=-1, - largest=True, - sorted=False, - ) - sampled = torch.empty((1,), dtype=torch.int64, device=logits.device) - generator = next(iter(sampling_metadata.generators.values()), None) - sm70_ops.sm70_sample_chunked_top20_philox_token_out( - sampled, - global_values, - local_indices, - global_positions, - generator, - 248320, - 0.95, - chunk_size, - ) - logger.info_once( - "SM70 exact %d-way chunked top-k20 Philox sampler path enabled.", - chunks, - ) - return sampled - - # The sparse-Philox kernel canonicalizes the selected pairs itself, so - # skip torch.topk's final ordering pass only when that exact path will - # consume the result. Keep sorted output for the legacy fallback. - use_sparse_philox = ( - hasattr(torch.ops._C, "sm70_sample_sorted_top20_philox_token_out") - or hasattr(torch.ops._C, "sm70_sample_sorted_top20_philox_out") - ) and len(sampling_metadata.generators) <= 1 - top_values, top_indices = torch.topk( - logits, - k=20, - dim=-1, - largest=True, - sorted=not use_sparse_philox, - ) - return Sampler._sm70_sample_top20_pairs( - top_values, top_indices, sampling_metadata - ) - - @staticmethod - def _sm70_compact_topk20_metadata_supported( - sampling_metadata: SamplingMetadata, - logprobs_mode: LogprobsMode, - ) -> bool: - if not envs.VLLM_SM70_COMPACT_TOPK20_SAMPLER: - return False - if logprobs_mode in ("processed_logits", "processed_logprobs"): - return False - if sampling_metadata.max_num_logprobs is not None: - return False - if sampling_metadata.logprob_token_ids: - return False - if not sampling_metadata.all_random: - return False - if sampling_metadata.top_k_cpu != (20,): - return False - top_p_cpu = sampling_metadata.top_p_cpu - temperature_cpu = sampling_metadata.temperature_cpu - return not ( - top_p_cpu is None - or len(top_p_cpu) != 1 - or abs(top_p_cpu[0] - 0.95) > 1e-6 - or temperature_cpu is None - or len(temperature_cpu) != 1 - or abs(temperature_cpu[0] - 1.0) > 1e-6 - ) - - @staticmethod - def _sm70_sample_top20_pairs( - top_values: torch.Tensor, - top_indices: torch.Tensor, - sampling_metadata: SamplingMetadata, - ) -> torch.Tensor | None: - if ( - top_values.dtype != torch.float32 - or top_indices.dtype != torch.int64 - or not top_values.is_cuda - or not top_indices.is_cuda - or top_values.shape != (1, 20) - or top_indices.shape != (1, 20) - or top_values.device != top_indices.device - or torch.cuda.get_device_capability(top_values.device) != (7, 0) - ): - return None - - from vllm import _sm70_ops as sm70_ops - - has_token_philox = hasattr( - torch.ops._C, "sm70_sample_sorted_top20_philox_token_out" - ) - has_metadata_philox = hasattr( - torch.ops._C, "sm70_sample_sorted_top20_philox_out" - ) - if not ( - has_token_philox - or has_metadata_philox - or hasattr(torch.ops._C, "sm70_sample_packed_top20_out") - ): - return None - - sampled = torch.empty((1,), dtype=torch.int64, device=top_values.device) - generators = sampling_metadata.generators - if has_token_philox and len(generators) <= 1: - generator = next(iter(generators.values()), None) - sm70_ops.sm70_sample_sorted_top20_philox_token_out( - sampled, - top_values, - top_indices, - generator, - 248320, - 0.95, - ) - logger.info_once( - "SM70 exact token-only Philox compact top-k20 sampler path enabled." - ) - return sampled - - sparse_ids = torch.empty((20,), dtype=torch.int64, device=top_values.device) - sparse_probs = torch.empty((20,), dtype=torch.float32, device=top_values.device) - if has_metadata_philox and len(generators) <= 1: - generator = next(iter(generators.values()), None) - sm70_ops.sm70_sample_sorted_top20_philox_out( - sampled, - sparse_ids, - sparse_probs, - top_values, - top_indices, - generator, - 248320, - 0.95, - ) - logger.info_once( - "SM70 exact sparse-Philox compact top-k20 sampler path enabled." - ) - return sampled - - top_indices_float = top_indices[0].to(torch.float32) - pairs = torch.stack( - (top_values[0], top_indices_float, top_indices_float), dim=-1 - ) - exponentials = torch.empty( - (1, 248320), dtype=torch.float32, device=top_values.device - ) - if len(generators) != 1: - exponentials.exponential_() - for row, generator in generators.items(): - exponentials[row].exponential_(generator=generator) - - sm70_ops.sm70_sample_packed_top20_out( - sampled, - sparse_ids, - sparse_probs, - pairs, - exponentials, - 0.95, - ) - logger.info_once("SM70 exact compact top-k20 random sampler path enabled.") - return sampled - - def sm70_compact_topk20_pairs_sample( - self, - top_values: torch.Tensor, - top_indices: torch.Tensor, - sampling_metadata: SamplingMetadata, - ) -> torch.Tensor | None: - """Sample exact global top-20 pairs while preserving native RNG use.""" - if not self._sm70_compact_topk20_metadata_supported( - sampling_metadata, self.logprobs_mode - ): - return None - return self._sm70_sample_top20_pairs(top_values, top_indices, sampling_metadata) - def sample( self, logits: torch.Tensor, @@ -556,12 +345,6 @@ def sample( for processor in sampling_metadata.logitsprocs.argmax_invariant: logits = processor.apply(logits) - compact_sampled = self._try_sm70_compact_topk20_sample( - logits, sampling_metadata, logprobs_mode - ) - if compact_sampled is not None: - return compact_sampled, None - # Apply top_k and/or top_p. random_sampled, processed_logprobs = self.topk_topp_sampler( logits, diff --git a/vllm/v1/sample/thinking_budget_state.py b/vllm/v1/sample/thinking_budget_state.py index 0f7ab58187..4dd764766d 100644 --- a/vllm/v1/sample/thinking_budget_state.py +++ b/vllm/v1/sample/thinking_budget_state.py @@ -153,7 +153,9 @@ 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"). @@ -161,7 +163,9 @@ 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/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index 6cdd5e5b41..24b6a178ce 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -119,7 +119,9 @@ def __init__( self.cpu_coordinator: KVCacheCoordinator = get_kv_cache_coordinator( kv_cache_config=self.cpu_kv_cache_config, max_model_len=vllm_config.model_config.max_model_len, - max_in_flight_tokens=vllm_config.max_in_flight_tokens, + max_num_batched_tokens=( + vllm_config.scheduler_config.max_num_batched_tokens + ), use_eagle=False, enable_caching=True, enable_kv_cache_events=self.enable_kv_cache_events, diff --git a/vllm/v1/spec_decode/ddtree_metadata.py b/vllm/v1/spec_decode/ddtree_metadata.py index e8dc694676..0cf0cf2d07 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,7 +98,8 @@ 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( @@ -114,7 +115,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 @@ -125,9 +126,10 @@ 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) @@ -155,8 +157,7 @@ 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] = [] @@ -171,18 +172,15 @@ 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))) @@ -194,20 +192,24 @@ 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), ) @@ -224,8 +226,7 @@ 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} @@ -251,8 +252,7 @@ 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/dflash.py b/vllm/v1/spec_decode/dflash.py index ffabcdc576..f90928b8b7 100644 --- a/vllm/v1/spec_decode/dflash.py +++ b/vllm/v1/spec_decode/dflash.py @@ -55,32 +55,6 @@ def _ddtree_worker_profile_enabled() -> bool: class DFlashProposer(SpecDecodeBaseProposer): - @override - def _create_draft_vllm_config(self) -> VllmConfig: - base = super()._create_draft_vllm_config() - if ( - self.speculative_config.kv_cache_dtype is None - and str(base.cache_config.cache_dtype).startswith("fp8") - and current_platform.is_cuda() - and current_platform.is_device_capability(70) - ): - logger.info_once( - "Using FP16 draft KV cache for SM70 DFlash DDTree while the " - "target uses %s.", - base.cache_config.cache_dtype, - ) - base = replace( - base, - cache_config=replace(base.cache_config, cache_dtype="auto"), - ) - return replace( - base, - attention_config=replace( - base.attention_config, - use_non_causal=not self.dflash_causal, - ), - ) - def __init__( self, vllm_config: VllmConfig, @@ -89,7 +63,7 @@ def __init__( ): assert vllm_config.speculative_config is not None assert ( - vllm_config.speculative_config.use_dflash_ddtree() + vllm_config.speculative_config.use_dflash() or vllm_config.speculative_config.use_dspark() ) super().__init__( @@ -349,6 +323,17 @@ def warmup_sm70_dflash_hotpath_kernels(self) -> tuple[str, ...]: "dflash_expand_inputs", ) + @override + def _create_draft_vllm_config(self) -> VllmConfig: + base = super()._create_draft_vllm_config() + return replace( + base, + attention_config=replace( + base.attention_config, + use_non_causal=not self.dflash_causal, + ), + ) + @override def _warn_if_multimodal(self): # Override to allow multimodal inputs since DFlash supports Qwen3.5 models diff --git a/vllm/v1/spec_decode/draft_prob_alignment.py b/vllm/v1/spec_decode/draft_prob_alignment.py index 21e6dd8b64..8a0ad46c7e 100644 --- a/vllm/v1/spec_decode/draft_prob_alignment.py +++ b/vllm/v1/spec_decode/draft_prob_alignment.py @@ -160,108 +160,3 @@ def get_aligned_draft_probs( if not draft_probs_rows: return None return torch.cat(draft_probs_rows, dim=0).contiguous() - - -def get_aligned_draft_scalar_values( - *, - req_ids: Sequence[str], - values: torch.Tensor | None, - value_req_ids: Sequence[str] | None, - value_token_ids: DraftProbTokenIds | None, - spec_decode_metadata: SpecDecodeMetadata, -) -> torch.Tensor | None: - """Align a cached per-draft scalar tensor to verifier request order.""" - if values is None or value_req_ids is None: - return None - if values.ndim != 2: - raise RuntimeError( - "Cached draft scalar values must have shape " - f"[num_requests, max_draft_tokens], got {tuple(values.shape)}." - ) - if values.shape[0] != len(value_req_ids): - raise RuntimeError( - "Cached draft scalar row count does not match request ids: " - f"rows={values.shape[0]}, req_ids={len(value_req_ids)}." - ) - if value_token_ids is None: - raise RuntimeError( - "Cached draft scalar values are missing their draft token ids; " - "cannot verify alignment." - ) - if isinstance(value_token_ids, torch.Tensor): - if value_token_ids.ndim != 2 or value_token_ids.shape[0] != len(value_req_ids): - raise RuntimeError( - "Cached draft scalar token ids must have one 2-D row per " - f"request, got shape={tuple(value_token_ids.shape)}, " - f"req_ids={len(value_req_ids)}." - ) - elif len(value_token_ids) != len(value_req_ids): - raise RuntimeError( - "Cached draft scalar token row count does not match request ids: " - f"rows={len(value_token_ids)}, req_ids={len(value_req_ids)}." - ) - if len(req_ids) != len(spec_decode_metadata.num_draft_tokens): - raise RuntimeError( - "Spec decode draft metadata batch size does not match current " - f"requests: metadata={len(spec_decode_metadata.num_draft_tokens)}, " - f"req_ids={len(req_ids)}." - ) - expected_num_draft_tokens = sum(spec_decode_metadata.num_draft_tokens) - if spec_decode_metadata.draft_token_ids.numel() != expected_num_draft_tokens: - raise RuntimeError( - "Spec decode draft token count does not match per-request counts: " - f"draft_token_ids={spec_decode_metadata.draft_token_ids.numel()}, " - f"num_draft_tokens={expected_num_draft_tokens}." - ) - - row_by_req_id = {req_id: idx for idx, req_id in enumerate(value_req_ids)} - aligned_rows: list[torch.Tensor] = [] - draft_token_offset = 0 - for req_id, num_draft in zip(req_ids, spec_decode_metadata.num_draft_tokens): - if num_draft == 0: - continue - row_idx = row_by_req_id.get(req_id) - if row_idx is None: - raise RuntimeError( - f"Missing cached draft scalar values for request {req_id}; " - "cannot verify alignment." - ) - if values.shape[1] < num_draft: - raise RuntimeError( - "Cached draft scalar values do not have enough positions for " - f"request {req_id}: available={values.shape[1]}, needed={num_draft}." - ) - expected_tokens = spec_decode_metadata.draft_token_ids[ - draft_token_offset : draft_token_offset + num_draft - ] - if isinstance(value_token_ids, torch.Tensor): - if value_token_ids.shape[1] < num_draft: - raise RuntimeError( - "Cached draft scalar token ids do not have enough positions " - f"for request {req_id}: available={value_token_ids.shape[1]}, " - f"needed={num_draft}." - ) - cached_tokens = value_token_ids[row_idx, :num_draft].to( - device=expected_tokens.device, - dtype=expected_tokens.dtype, - ) - else: - cached_token_row = value_token_ids[row_idx] - if len(cached_token_row) < num_draft: - raise RuntimeError( - "Cached draft scalar token ids do not have enough positions " - f"for request {req_id}: available={len(cached_token_row)}, " - f"needed={num_draft}." - ) - cached_tokens = torch.tensor( - cached_token_row[:num_draft], - device=expected_tokens.device, - dtype=expected_tokens.dtype, - ) - _assert_cached_tokens_match(cached_tokens, expected_tokens, req_id) - aligned_rows.append(values[row_idx, :num_draft]) - draft_token_offset += num_draft - - if not aligned_rows: - return None - return torch.cat(aligned_rows, dim=0).contiguous() diff --git a/vllm/v1/spec_decode/dspark.py b/vllm/v1/spec_decode/dspark.py index cb013f86e4..93d8fb2abb 100644 --- a/vllm/v1/spec_decode/dspark.py +++ b/vllm/v1/spec_decode/dspark.py @@ -5,7 +5,6 @@ import torch from typing_extensions import override -from vllm import envs from vllm.config import VllmConfig from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.spec_decode.dflash import DFlashProposer @@ -28,76 +27,6 @@ def __init__( torch.arange(self.max_batch_size, dtype=torch.int64, device=device) * self.num_query_per_req ) - self._confidence_threshold = self.speculative_config.dspark_confidence_threshold - self._confidence_temperatures: torch.Tensor | None = None - if self._confidence_threshold > 0.0: - temperatures = self.speculative_config.dspark_confidence_temperatures - if temperatures is None: - temperatures = [1.0] * self.num_speculative_tokens - self._confidence_temperatures = torch.tensor( - temperatures, - dtype=torch.float32, - device=device, - ).view(1, -1) - configured_cap = self.speculative_config.dspark_max_verification_tokens - self._max_verification_tokens = ( - self.num_speculative_tokens if configured_cap is None else configured_cap - ) - self.confidence_scheduling_enabled = ( - self._confidence_threshold > 0.0 - or self._max_verification_tokens < self.num_speculative_tokens - ) - # A fixed cap needs verification lengths but not a confidence - # projection. Collect logits only when the calibrated threshold or the - # explicit alignment diagnostic consumes them. - self.collect_confidence_logits = ( - self._confidence_threshold > 0.0 or envs.VLLM_SPEC_DUMP_ALIGNMENT - ) - self._last_confidence_logits: torch.Tensor | None = None - self._last_verification_lengths: torch.Tensor | None = None - - def take_last_confidence_logits(self) -> torch.Tensor | None: - confidence_logits = self._last_confidence_logits - self._last_confidence_logits = None - return confidence_logits - - def take_last_verification_lengths(self) -> torch.Tensor | None: - verification_lengths = self._last_verification_lengths - self._last_verification_lengths = None - return verification_lengths - - def _select_verification_lengths( - self, - confidence_logits: torch.Tensor, - ) -> torch.Tensor: - temperatures = getattr(self, "_confidence_temperatures", None) - if temperatures is None: - temperatures = torch.ones( - (1, confidence_logits.shape[1]), - dtype=torch.float32, - device=confidence_logits.device, - ) - threshold = float(getattr(self, "_confidence_threshold", 0.0)) - cap = int( - getattr( - self, - "_max_verification_tokens", - confidence_logits.shape[1], - ) - ) - if threshold <= 0.0: - return torch.full( - (confidence_logits.shape[0],), - cap, - dtype=torch.int32, - device=confidence_logits.device, - ) - - conditional_probs = torch.sigmoid(confidence_logits.float() / temperatures) - confident_prefix = ( - (conditional_probs >= threshold).to(torch.int32).cumprod(dim=1) - ) - return confident_prefix.sum(dim=1).clamp_max_(cap).to(torch.int32) @override def _sample_draft_tokens( @@ -125,19 +54,9 @@ def _sample_draft_tokens( # contents across asynchronous scheduling and CUDA Graph replay. prev = self.input_ids[self._anchor_indices[:batch_size]].to(torch.long) draft_tokens: list[torch.Tensor] = [] - self._last_confidence_logits = None - self._last_verification_lengths = None - collect_confidence_logits = bool( - getattr(self, "collect_confidence_logits", False) - ) - markov_embeds: list[torch.Tensor] | None = ( - [] if collect_confidence_logits else None - ) draft_probs: list[torch.Tensor] | None = None for step in range(self.num_speculative_tokens): markov_embed = self.model.markov_embed(prev) - if markov_embeds is not None: - markov_embeds.append(markov_embed) step_logits = base_logits[:, step] + self.model.markov_bias(markov_embed) sampled, step_probs = self._sample_from_logits( step_logits, @@ -150,26 +69,6 @@ def _sample_draft_tokens( draft_probs = [] draft_probs.append(step_probs) - if markov_embeds is not None: - block_hidden_states = hidden_states.view( - batch_size, self.num_speculative_tokens, -1 - ) - self._last_confidence_logits = self.model.confidence_logits( - block_hidden_states, - torch.stack(markov_embeds, dim=1), - ) - if bool(getattr(self, "confidence_scheduling_enabled", False)): - if self._last_confidence_logits is None: - self._last_verification_lengths = torch.full( - (batch_size,), - int(self._max_verification_tokens), - dtype=torch.int32, - device=hidden_states.device, - ) - else: - self._last_verification_lengths = self._select_verification_lengths( - self._last_confidence_logits - ) flat_tokens = torch.stack(draft_tokens, dim=1).reshape(-1) if draft_probs is None: return flat_tokens, None diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index 98bed3ac9d..85811c4c75 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -87,47 +87,6 @@ def _sm70_mtp_profile_interval() -> int: return envs.VLLM_SM70_MTP_PROFILE_INTERVAL -def _sm70_mtp_moe_warmup_sizes( - num_experts: int, - top_k: int, - max_num_tokens: int, - decode_tile_max_tokens: int = 0, -) -> tuple[int, ...]: - """Pick one token count for each SM70 Triton MoE dispatch region.""" - if num_experts <= 0 or top_k <= 0 or max_num_tokens <= 0: - return () - - def prefer_unaligned(size: int) -> int: - # Triton otherwise specializes another kernel when token_count * top_k - # first loses 16-byte divisibility. This cannot happen when top_k is - # itself a multiple of 16. - if top_k % 16: - while size * top_k % 16 == 0: - size += 1 - return size - - sizes = { - # First sorted-assignment shape after the naive small-token path. - prefer_unaligned(num_experts // (top_k * 4) + 1), - # First shape whose sorted-token allocation is not block aligned. - prefer_unaligned((num_experts + top_k - 1) // top_k), - # First shape using the large-M SM70 0.0.3 MoE tile. - prefer_unaligned(num_experts + 1), - } - sizes.update(range(1, min(decode_tile_max_tokens, max_num_tokens) + 1)) - return tuple(sorted(size for size in sizes if size <= max_num_tokens)) - - -def _sm70_mtp_hotpath_warmup_batch_sizes( - max_batch_size: int, include_alternate: bool = False -) -> tuple[int, ...]: - """Preserve the primary warmup and optionally cover its alternate shape.""" - concurrent_batch_size = max(1, min(int(max_batch_size), 4)) - if include_alternate: - return tuple(sorted({1, concurrent_batch_size})) - return (concurrent_batch_size,) - - def _is_dflash_method(method: str | None) -> bool: return method in ("dflash", "dflash_ddtree", "dspark") @@ -876,125 +835,6 @@ def _sm70_mtp_profile_report( summary, ) - def _warmup_sm70_mtp_hotpath_batch( - self, batch_size: int, vocab_size: int, warm_draft_topk: bool - ) -> None: - valid_sampled_tokens_count = None - for num_sampled_tokens in sorted({1, self.num_speculative_tokens + 1}): - sampled_token_ids = torch.zeros( - (batch_size, num_sampled_tokens), - dtype=torch.int32, - device=self.device, - ) - discard_request_mask = torch.zeros( - batch_size, dtype=torch.bool, device=self.device - ) - backup_next_token_ids = torch.zeros( - batch_size, dtype=torch.int32, device=self.device - ) - next_token_ids = torch.empty( - batch_size, dtype=torch.int32, device=self.device - ) - next_valid_count = torch.empty_like(next_token_ids) - eagle_prepare_next_token_padded_kernel[(batch_size,)]( - sampled_token_ids, - discard_request_mask, - backup_next_token_ids, - next_token_ids, - next_valid_count, - vocab_size, - num_sampled_tokens, - batch_size, - sampled_token_ids.stride(0), - BLOCK_SIZE_TOKENS=next_power_of_2(num_sampled_tokens), - ) - if num_sampled_tokens == self.num_speculative_tokens + 1: - valid_sampled_tokens_count = next_valid_count - assert valid_sampled_tokens_count is not None - - next_token_ids = torch.empty(batch_size, dtype=torch.int32, device=self.device) - cu_num_draft_tokens = ( - torch.arange( - 1, - batch_size + 1, - dtype=torch.int32, - device=self.device, - ) - * self.num_speculative_tokens - ) - query_start_loc = torch.arange( - batch_size + 1, dtype=torch.int32, device=self.device - ) - token_indices_to_sample = torch.empty_like(next_token_ids) - num_rejected_tokens_gpu = torch.empty_like(next_token_ids) - eagle_prepare_inputs_padded_kernel[(batch_size,)]( - cu_num_draft_tokens, - valid_sampled_tokens_count, - query_start_loc, - token_indices_to_sample, - num_rejected_tokens_gpu, - batch_size, - ) - - for dtype, replace_from, replace_to in ( - (torch.float32, 0, 1), # temperature - (torch.float32, 0, 0), # top_p - (torch.int32, 0, 0), # top_k - ): - rejection_expand_input = torch.ones( - batch_size, dtype=dtype, device=self.device - ) - rejection_expand_output = torch.empty( - batch_size * self.num_speculative_tokens, - dtype=dtype, - device=self.device, - ) - rejection_expand_kernel[(batch_size,)]( - rejection_expand_output, - rejection_expand_input, - cu_num_draft_tokens, - replace_from, - replace_to, - MAX_NUM_TOKENS=MAX_SPEC_LEN, - ) - - n_blocks_per_req = max( - 1, (self.max_model_len + self.block_size - 1) // self.block_size - ) - positions = torch.zeros(batch_size, dtype=torch.int64, device=self.device) - block_table = torch.zeros( - (batch_size, n_blocks_per_req), - dtype=torch.int32, - device=self.device, - ) - seq_lens = torch.ones(batch_size, dtype=torch.int32, device=self.device) - out_positions = torch.empty_like(positions) - out_slot_mapping = torch.empty( - batch_size, dtype=torch.int64, device=self.device - ) - eagle_step_update_slot_mapping_and_metadata( - positions_1d=positions, - block_table_tensor=block_table, - seq_lens=seq_lens, - block_size=self.block_size, - max_model_len=self.max_model_len, - out_clamped_positions=out_positions, - out_slot_mapping=out_slot_mapping, - input_batch_size=batch_size, - ) - - if warm_draft_topk: - draft_logits = torch.zeros( - (batch_size, vocab_size), dtype=torch.float32, device=self.device - ) - draft_top_k = torch.full( - (batch_size,), - min(20, vocab_size), - dtype=torch.int32, - device=self.device, - ) - apply_top_k_top_p(draft_logits, draft_top_k, None) - def warmup_sm70_mtp_hotpath_kernels(self) -> tuple[str, ...]: """Warm MTP helper kernels that otherwise JIT on the first request.""" if ( @@ -1010,101 +850,126 @@ def warmup_sm70_mtp_hotpath_kernels(self) -> tuple[str, ...]: return () try: + batch_size = max(1, min(self.max_batch_size, 4)) vocab_size = max(2, self.draft_model_config.get_vocab_size()) - expanded_warmup = envs.VLLM_SM70_MTP_CONCURRENCY_WARMUP - for batch_size in _sm70_mtp_hotpath_warmup_batch_sizes( - self.max_batch_size, include_alternate=expanded_warmup - ): - self._warmup_sm70_mtp_hotpath_batch( - batch_size, vocab_size, warm_draft_topk=expanded_warmup + + valid_sampled_tokens_count = None + for num_sampled_tokens in sorted({1, self.num_speculative_tokens + 1}): + sampled_token_ids = torch.zeros( + (batch_size, num_sampled_tokens), + dtype=torch.int32, + device=self.device, ) - torch.accelerator.synchronize() - except Exception as err: # pragma: no cover - best-effort warmup - logger.warning_once("SM70 MTP hotpath warmup skipped: %s", err) - return () + discard_request_mask = torch.zeros( + batch_size, dtype=torch.bool, device=self.device + ) + backup_next_token_ids = torch.zeros( + batch_size, dtype=torch.int32, device=self.device + ) + next_token_ids = torch.empty( + batch_size, dtype=torch.int32, device=self.device + ) + next_valid_count = torch.empty_like(next_token_ids) + eagle_prepare_next_token_padded_kernel[(batch_size,)]( + sampled_token_ids, + discard_request_mask, + backup_next_token_ids, + next_token_ids, + next_valid_count, + vocab_size, + num_sampled_tokens, + batch_size, + sampled_token_ids.stride(0), + BLOCK_SIZE_TOKENS=next_power_of_2(num_sampled_tokens), + ) + if num_sampled_tokens == self.num_speculative_tokens + 1: + valid_sampled_tokens_count = next_valid_count + assert valid_sampled_tokens_count is not None - warmed_kernels = [ - "mtp_prepare_next_token", - "mtp_prepare_inputs", - "mtp_rejection_expand", - "mtp_step_slot_mapping", - ] - if envs.VLLM_SM70_MTP_CONCURRENCY_WARMUP: - warmed_kernels.append("mtp_draft_topk") - return tuple(warmed_kernels) + next_token_ids = torch.empty( + batch_size, dtype=torch.int32, device=self.device + ) - def warmup_sm70_mtp_moe_kernels(self) -> tuple[str, ...]: - """Warm the finite Triton-MoE shape specializations used by MTP.""" - if ( - self.method != "mtp" - or self.device.type != "cuda" - or not current_platform.is_device_capability(70) - or not envs.VLLM_SM70_UNQUANTIZED_MOE_0DOT3_CONFIG - or not self.draft_model_config.is_moe - ): - return () - if getattr(self, "_sm70_mtp_moe_warmed", False): - return () - self._sm70_mtp_moe_warmed = True - - text_config = self.draft_model_config.hf_text_config - top_k = int(getattr(text_config, "num_experts_per_tok", 0)) - num_experts = self.draft_model_config.get_num_experts() - tp_size = self.vllm_config.parallel_config.tensor_parallel_size - use_mtp_decode_tiles = ( - envs.VLLM_SM70_MTP_MOE_TUNED_CONFIG - and num_experts == 256 - and top_k == 8 - and self.draft_model_config.get_hidden_size() == 2048 - and int(getattr(text_config, "moe_intermediate_size", 0)) == 512 - and tp_size == 4 - ) - sizes = _sm70_mtp_moe_warmup_sizes( - num_experts, - top_k, - self.max_num_tokens, - decode_tile_max_tokens=16 if use_mtp_decode_tiles else 0, - ) - legacy_sizes = (1, 2, 9, 10) if use_mtp_decode_tiles else () - if not sizes: - return () + cu_num_draft_tokens = ( + torch.arange( + 1, + batch_size + 1, + dtype=torch.int32, + device=self.device, + ) + * self.num_speculative_tokens + ) + query_start_loc = torch.arange( + batch_size + 1, dtype=torch.int32, device=self.device + ) + token_indices_to_sample = torch.empty_like(next_token_ids) + num_rejected_tokens_gpu = torch.empty_like(next_token_ids) + eagle_prepare_inputs_padded_kernel[(batch_size,)]( + cu_num_draft_tokens, + valid_sampled_tokens_count, + query_start_loc, + token_indices_to_sample, + num_rejected_tokens_gpu, + batch_size, + ) - try: - for num_tokens in sizes: - self.dummy_run( - num_tokens, - use_cudagraphs=False, - spec_step_idx=0, + for dtype, replace_from, replace_to in ( + (torch.float32, 0, 1), # temperature + (torch.float32, 0, 0), # top_p + (torch.int32, 0, 0), # top_k + ): + rejection_expand_input = torch.ones( + batch_size, dtype=dtype, device=self.device ) - if legacy_sizes: - # The tuned drafter tile replaces legacy graph - # specializations that used to be compiled incidentally. - # Other unquantized MoE calls can reuse those legacy Triton - # signatures during serving. Cover both token divisibility - # cases for the naive (M1/M2) and aligned (M9/M10) assignment - # paths instead of allowing a first-request JIT spike. - from vllm.model_executor.layers.fused_moe.fused_moe import ( - force_sm70_mtp_moe_legacy_config, + rejection_expand_output = torch.empty( + batch_size * self.num_speculative_tokens, + dtype=dtype, + device=self.device, + ) + rejection_expand_kernel[(batch_size,)]( + rejection_expand_output, + rejection_expand_input, + cu_num_draft_tokens, + replace_from, + replace_to, + MAX_NUM_TOKENS=MAX_SPEC_LEN, ) - with force_sm70_mtp_moe_legacy_config(): - for num_tokens in legacy_sizes: - self.dummy_run( - num_tokens, - use_cudagraphs=False, - spec_step_idx=0, - ) + n_blocks_per_req = max( + 1, (self.max_model_len + self.block_size - 1) // self.block_size + ) + positions = torch.zeros(batch_size, dtype=torch.int64, device=self.device) + block_table = torch.zeros( + (batch_size, n_blocks_per_req), + dtype=torch.int32, + device=self.device, + ) + seq_lens = torch.ones(batch_size, dtype=torch.int32, device=self.device) + out_positions = torch.empty_like(positions) + out_slot_mapping = torch.empty( + batch_size, dtype=torch.int64, device=self.device + ) + eagle_step_update_slot_mapping_and_metadata( + positions_1d=positions, + block_table_tensor=block_table, + seq_lens=seq_lens, + block_size=self.block_size, + max_model_len=self.max_model_len, + out_clamped_positions=out_positions, + out_slot_mapping=out_slot_mapping, + input_batch_size=batch_size, + ) torch.accelerator.synchronize() except Exception as err: # pragma: no cover - best-effort warmup - logger.warning_once("SM70 MTP MoE warmup skipped: %s", err) + logger.warning_once("SM70 MTP hotpath warmup skipped: %s", err) return () - logger.info_once( - "SM70 MTP MoE warmup finished: token_counts=%s, legacy_token_counts=%s", - sizes, - legacy_sizes, + return ( + "mtp_prepare_next_token", + "mtp_prepare_inputs", + "mtp_rejection_expand", + "mtp_step_slot_mapping", ) - return ("mtp_draft_moe",) def take_last_draft_probs(self) -> torch.Tensor | None: return self._last_draft_probs @@ -2104,15 +1969,6 @@ def _create_draft_vllm_config(self) -> VllmConfig: ), ) - if spec_cfg.kv_cache_dtype is not None: - base = replace( - base, - cache_config=replace( - base.cache_config, - cache_dtype=spec_cfg.kv_cache_dtype, - ), - ) - return base def _get_model(self) -> nn.Module: diff --git a/vllm/v1/worker/block_table.py b/vllm/v1/worker/block_table.py index 8329a2c196..616fdd15bd 100644 --- a/vllm/v1/worker/block_table.py +++ b/vllm/v1/worker/block_table.py @@ -104,7 +104,9 @@ 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..c2ba12437f 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Iterable, Sequence from dataclasses import dataclass from typing import Any, cast @@ -372,159 +372,6 @@ def build_slot_mappings_by_layer( return slot_mappings_by_layer -@dataclass(frozen=True) -class CommonGDNSpecMetadata: - """Batch-level GDN speculative metadata shared by every cache group.""" - - num_prefills: int - num_prefill_tokens: int - num_decodes: int - num_decode_tokens: int - num_spec_decodes: int - num_spec_decode_tokens: int - spec_query_start_loc: torch.Tensor - non_spec_query_start_loc: torch.Tensor | None - non_spec_query_start_loc_cpu: torch.Tensor | None - spec_sequence_masks_cpu: torch.Tensor - spec_sequence_masks: torch.Tensor - spec_token_indx: torch.Tensor - non_spec_token_indx: torch.Tensor - - -def compute_common_gdn_attn_metadata( - *, - num_decode_draft_tokens_cpu: torch.Tensor | None, - query_start_loc: torch.Tensor, - query_start_loc_cpu: torch.Tensor, - num_spec_state_tokens: int, - legacy_mixed_decode_routing: bool, -) -> CommonGDNSpecMetadata | None: - """Compute GDN speculative batch metadata once instead of per group. - - State block IDs remain group-specific and are deliberately excluded. This - is the dependency boundary needed to share request classification, token - indices, and query offsets without changing recurrent-state ownership. - """ - if num_spec_state_tokens <= 0 or num_decode_draft_tokens_cpu is None: - return None - - num_reqs = query_start_loc_cpu.numel() - 1 - if num_decode_draft_tokens_cpu.ndim != 1: - raise ValueError("num_decode_draft_tokens_cpu must be one-dimensional") - if num_decode_draft_tokens_cpu.numel() != num_reqs: - raise ValueError("num_decode_draft_tokens_cpu must align with query_start_loc") - - spec_sequence_masks_cpu = num_decode_draft_tokens_cpu >= 0 - num_spec_decodes = int(spec_sequence_masks_cpu.sum().item()) - if num_spec_decodes == 0: - return None - num_spec_draft_tokens = int( - num_decode_draft_tokens_cpu[spec_sequence_masks_cpu].sum().item() - ) - if num_spec_draft_tokens == 0: - return None - - spec_sequence_masks = spec_sequence_masks_cpu.to( - query_start_loc.device, non_blocking=True - ) - query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] - non_spec_query_lens_cpu = query_lens_cpu[~spec_sequence_masks_cpu] - num_zero_len = int((non_spec_query_lens_cpu == 0).sum().item()) - - if legacy_mixed_decode_routing: - num_decodes = int((non_spec_query_lens_cpu == 1).sum().item()) - num_prefills = non_spec_query_lens_cpu.numel() - num_decodes - num_zero_len - num_decode_tokens = num_decodes - num_prefill_tokens = ( - int(non_spec_query_lens_cpu.sum().item()) - num_decode_tokens - ) - else: - num_decodes = 0 - num_prefills = non_spec_query_lens_cpu.numel() - num_zero_len - num_decode_tokens = 0 - num_prefill_tokens = int(non_spec_query_lens_cpu.sum().item()) - - num_spec_decode_tokens = ( - int(query_lens_cpu.sum().item()) - num_prefill_tokens - num_decode_tokens - ) - if num_prefills == 0 and num_decodes == 0: - spec_token_size = min( - num_spec_decodes * (num_spec_state_tokens + 1), - int(query_start_loc_cpu[-1].item()), - ) - spec_token_indx = torch.arange( - spec_token_size, - dtype=torch.int32, - device=query_start_loc.device, - ) - non_spec_token_indx = torch.empty( - 0, - dtype=torch.int32, - device=query_start_loc.device, - ) - 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: - query_lens = query_lens_cpu.to(query_start_loc.device, non_blocking=True) - spec_token_masks = torch.repeat_interleave( - spec_sequence_masks, - query_lens, - output_size=int(query_start_loc_cpu[-1].item()), - ) - index = torch.argsort(spec_token_masks, stable=True) - num_non_spec_tokens = num_prefill_tokens + num_decode_tokens - non_spec_token_indx = index[:num_non_spec_tokens] - spec_token_indx = index[num_non_spec_tokens:] - - spec_query_start_loc = torch.zeros( - num_spec_decodes + 1, - dtype=torch.int32, - device=query_start_loc.device, - ) - torch.cumsum( - query_lens[spec_sequence_masks], - dim=0, - out=spec_query_start_loc[1:], - ) - non_spec_query_start_loc = torch.zeros( - query_lens.numel() - num_spec_decodes + 1, - dtype=torch.int32, - device=query_start_loc.device, - ) - torch.cumsum( - query_lens[~spec_sequence_masks], - dim=0, - out=non_spec_query_start_loc[1:], - ) - non_spec_query_start_loc_cpu = torch.zeros( - query_lens_cpu.numel() - num_spec_decodes + 1, - dtype=torch.int32, - device="cpu", - ) - torch.cumsum( - query_lens_cpu[~spec_sequence_masks_cpu], - dim=0, - out=non_spec_query_start_loc_cpu[1:], - ) - - return CommonGDNSpecMetadata( - num_prefills=num_prefills, - num_prefill_tokens=num_prefill_tokens, - num_decodes=num_decodes, - num_decode_tokens=num_decode_tokens, - num_spec_decodes=num_spec_decodes, - num_spec_decode_tokens=num_spec_decode_tokens, - spec_query_start_loc=spec_query_start_loc, - non_spec_query_start_loc=non_spec_query_start_loc, - non_spec_query_start_loc_cpu=non_spec_query_start_loc_cpu, - spec_sequence_masks_cpu=spec_sequence_masks_cpu, - spec_sequence_masks=spec_sequence_masks, - spec_token_indx=spec_token_indx, - non_spec_token_indx=non_spec_token_indx, - ) - - def build_attn_metadata( attn_groups: list[list[AttentionGroup]], num_reqs: int, @@ -542,8 +389,6 @@ def build_attn_metadata( positions: torch.Tensor | None = None, model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None, for_cudagraph_capture: bool = False, - causal: bool | torch.Tensor | Mapping[int, bool] = True, - prefix_anchor_lens: torch.Tensor | None = None, ) -> dict[str, Any]: seq_lens = seq_lens[:num_reqs] if dcp_local_seq_lens is not None: @@ -556,12 +401,6 @@ def build_attn_metadata( for i in range(num_kv_cache_groups): block_table = block_tables[i] slot_mapping = slot_mappings[i] - # Hybrid drafters can assign different causality to each KV group. - # A Mapping must be resolved here; passing the dict itself into the - # backend metadata makes it truthy and silently selects causal kernels. - group_causal = ( - causal if isinstance(causal, (bool, torch.Tensor)) else causal.get(i, True) - ) common_attn_metadata_extra_kwargs = ( model_specific_attn_metadata.get_extra_common_attn_kwargs(i, num_reqs) @@ -579,10 +418,9 @@ def build_attn_metadata( max_query_len=max_query_len, block_table_tensor=block_table, slot_mapping=slot_mapping, - causal=group_causal, + causal=True, dcp_local_seq_lens=dcp_local_seq_lens, positions=positions, - prefix_anchor_lens=prefix_anchor_lens, **common_attn_metadata_extra_kwargs, ) diff --git a/vllm/v1/worker/gpu/buffer_utils.py b/vllm/v1/worker/gpu/buffer_utils.py index 2dd6bf92b8..5963790a77 100644 --- a/vllm/v1/worker/gpu/buffer_utils.py +++ b/vllm/v1/worker/gpu/buffer_utils.py @@ -137,12 +137,6 @@ def __init__( self.write_starts = new_buffer(self.num_rows, dtype=torch.int32) self.write_cu_lens = new_buffer(self.num_rows, dtype=torch.int32) - def get_cpu_view(self) -> torch.Tensor: - """Return the host view when this tensor is backed by mapped memory.""" - if not hasattr(self, "_uva_buf"): - raise RuntimeError("This StagedWriteTensor is not UVA-backed") - return self._uva_buf.cpu - def stage_write( self, index: int, start: int, x: Iterable[int] | Iterable[float] ) -> None: diff --git a/vllm/v1/worker/gpu/cp_utils.py b/vllm/v1/worker/gpu/cp_utils.py index 08a1a258ed..6dd8fd3474 100644 --- a/vllm/v1/worker/gpu/cp_utils.py +++ b/vllm/v1/worker/gpu/cp_utils.py @@ -59,24 +59,3 @@ def _dcp_local_seq_lens_kernel( # For [num_reqs, max_num_reqs), pad with 0 local_seq_lens = tl.where(block < num_reqs, local_seq_lens, 0) tl.store(out_ptr + block, local_seq_lens, mask=block < max_num_reqs) - - -@triton.jit -def cp_local_slot( - positions, - block_numbers, - block_size, - cp_rank, - CP_SIZE: tl.constexpr, - CP_INTERLEAVE: tl.constexpr, - PAD_ID: tl.constexpr, -): - """Return rank-local KV slots, or PAD_ID for non-local positions.""" - block_offsets = positions % (block_size * CP_SIZE) - if CP_SIZE == 1: - return block_numbers * block_size + block_offsets - is_local = block_offsets // CP_INTERLEAVE % CP_SIZE == cp_rank - rounds = block_offsets // (CP_INTERLEAVE * CP_SIZE) - remainder = block_offsets % CP_INTERLEAVE - local_offsets = rounds * CP_INTERLEAVE + remainder - return tl.where(is_local, block_numbers * block_size + local_offsets, PAD_ID) diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index a0953c44c5..4ecbb9f97d 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -34,36 +34,6 @@ logger = init_logger(__name__) -def _use_split_sm70_mtp_cudagraphs(vllm_config: VllmConfig) -> bool: - speculative_config = vllm_config.speculative_config - return bool( - envs.VLLM_SM70_MTP_SPLIT_DRAFT_CUDAGRAPHS - and speculative_config is not None - and speculative_config.method == "mtp" - and current_platform.is_cuda() - and current_platform.is_device_capability((7, 0)) - ) - - -def _split_sm70_mtp_manager_capture_sizes( - vllm_config: VllmConfig, - decode_query_len: int, -) -> list[int]: - capture_sizes = sorted(vllm_config.compilation_config.cudagraph_capture_sizes or []) - if decode_query_len != 1: - return capture_sizes - - speculative_config = vllm_config.speculative_config - assert speculative_config is not None - verifier_query_len = speculative_config.num_speculative_state_tokens() + 1 - return [ - size // verifier_query_len - for size in capture_sizes - if size % verifier_query_len == 0 - and size // verifier_query_len <= vllm_config.scheduler_config.max_num_seqs - ] - - class CapturedAttentionState(NamedTuple): attn_metadata: dict[str, Any] | None slot_mappings: dict[str, torch.Tensor] @@ -141,32 +111,15 @@ def __init__( self._graphs_captured = False self._candidates: list[list[BatchExecutionDescriptor]] = [] self._capture_descs: dict[CUDAGraphMode, list[BatchExecutionDescriptor]] = {} - if _use_split_sm70_mtp_cudagraphs(vllm_config): - self._capture_sizes = _split_sm70_mtp_manager_capture_sizes( - vllm_config, - decode_query_len, - ) - logger.info_once( - "Using split SM70 MTP CUDA graph manager shapes: " - "decode_query_len=%d capture_sizes=%s shared_verifier_shapes=%s.", - decode_query_len, - tuple(self._capture_sizes), - tuple(self.compilation_config.cudagraph_capture_sizes or ()), - ) - else: - # Legacy MRV1 behavior. The rollback path deliberately preserves - # the shared in-place normalization used before split managers. - self.compilation_config.adjust_cudagraph_sizes_for_spec_decode( - self.decode_query_len, self.tp_size - ) - self._capture_sizes = list( - self.compilation_config.cudagraph_capture_sizes or [] - ) + # adjust the cudagraph sizes to be a multiple of the uniform decode query length + self.compilation_config.adjust_cudagraph_sizes_for_spec_decode( + self.decode_query_len, self.tp_size + ) self._init_candidates() def _init_candidates(self) -> None: """Build priority-ordered candidate lists for each token count.""" - capture_sizes = self._capture_sizes + capture_sizes = self.compilation_config.cudagraph_capture_sizes if not (self.cudagraph_mode and capture_sizes): return diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index 1e117e499d..b253d7d8c0 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -84,14 +84,6 @@ class InputBatch: # Whether any requests in batch use structured output. has_structured_output_reqs: bool - # [num_reqs_after_padding] per-request prompt length for prefix-anchored - # sliding-window attention (optional). - prefix_anchor_lens: torch.Tensor | None = None - # CPU mask for requests whose scheduled chunk does not finish prefill. - # Model-specific runners may use this to omit work whose outputs cannot be - # consumed until a later chunk. None is equivalent to all False. - is_incomplete_prefilling_np: np.ndarray | None = None - @classmethod def make_dummy( cls, diff --git a/vllm/v1/worker/gpu/mamba_align.py b/vllm/v1/worker/gpu/mamba_align.py deleted file mode 100644 index 52a03e83c3..0000000000 --- a/vllm/v1/worker/gpu/mamba_align.py +++ /dev/null @@ -1,319 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""MRV2 GPU state migration for hybrid Mamba ``align`` prefix caching. - -This is the narrow Model Runner V2 portion of upstream vLLM PR #42406. It is -kept separate from :mod:`vllm.v1.worker.mamba_utils` because that module also -owns this tree's V1 DDTree state-selection compatibility path. -""" - -import torch - -from vllm.triton_utils import tl, triton -from vllm.v1.worker.mamba_utils import MambaSpecDecodeGPUContext - -# Temporal states are much larger than conv states. Splitting their byte range -# keeps block-boundary copies from serializing on one CTA at small batch sizes. -_TEMPORAL_TILES = 16 - - -@triton.jit -def _copy_mamba_state_block( - state_idx, - block_table_row, - src_col, - dst_col, - token_bias, - block_table_ptrs_ptr, - block_table_stride_req, - state_base_addrs_ptr, - state_block_strides_ptr, - state_elem_sizes_ptr, - state_inner_sizes_ptr, - state_conv_widths_ptr, - state_group_indices_ptr, - tile_idx, - COPY_BLOCK_SIZE: tl.constexpr, - TEMPORAL_TILES: tl.constexpr, -): - """Copy one SD-layout conv or temporal state using V1 align semantics.""" - state_base_addr = tl.load(state_base_addrs_ptr + state_idx) - state_block_stride = tl.load(state_block_strides_ptr + state_idx) - state_elem_size = tl.load(state_elem_sizes_ptr + state_idx).to(tl.int64) - state_inner_size = tl.load(state_inner_sizes_ptr + state_idx).to(tl.int64) - conv_width = tl.load(state_conv_widths_ptr + state_idx) - - group_idx = tl.load(state_group_indices_ptr + state_idx).to(tl.int64) - group_base_addr = tl.load(block_table_ptrs_ptr + group_idx) - block_table = group_base_addr.to(tl.pointer_type(tl.int32)) - block_table += block_table_row * block_table_stride_req - - dst_block_id = tl.load(block_table + dst_col).to(tl.int64) - dst_addr = state_base_addr + dst_block_id * state_block_stride - is_conv_state = conv_width > 0 - - if is_conv_state: - # Conv states are small and may be shifted left within the same block. - # One CTA owns the copy so a block-local barrier can make each chunk - # memmove-safe before lower addresses overwrite overlapping input. - if tile_idx > 0: - return - src_block_id = tl.load(block_table + src_col).to(tl.int64) - token_bytes = state_inner_size * state_elem_size - src_addr = ( - state_base_addr - + src_block_id * state_block_stride - + token_bias.to(tl.int64) * token_bytes - ) - copy_size = (conv_width - token_bias).to(tl.int64) * token_bytes - tile_start = copy_size * 0 - tile_end = copy_size - is_left_overlap = (dst_addr < src_addr) & (dst_addr + copy_size > src_addr) - else: - # Temporal state chooses the accepted speculative column. - src_block_id = tl.load(block_table + src_col + token_bias).to(tl.int64) - src_addr = state_base_addr + src_block_id * state_block_stride - copy_size = state_inner_size * state_elem_size - bytes_per_tile = (copy_size + TEMPORAL_TILES - 1) // TEMPORAL_TILES - tile_start = tile_idx.to(tl.int64) * bytes_per_tile - tile_end = tl.minimum(tile_start + bytes_per_tile, copy_size) - is_left_overlap = False - - offsets = tl.arange(0, COPY_BLOCK_SIZE) - for chunk_start in range(tile_start, tile_end, COPY_BLOCK_SIZE): - byte_offsets = chunk_start + offsets - mask = byte_offsets < tile_end - src = (src_addr + byte_offsets).to(tl.pointer_type(tl.uint8)) - dst = (dst_addr + byte_offsets).to(tl.pointer_type(tl.uint8)) - data = tl.load(src, mask=mask) - if is_left_overlap: - tl.debug_barrier() - tl.store(dst, data, mask=mask) - - -@triton.jit -def preprocess_mamba_align_fused_kernel( - idx_mapping_ptr, - state_idx_ptr, - num_computed_tokens_ptr, - query_start_loc_ptr, - num_accepted_tokens_ptr, - src_col_ptr, - token_bias_ptr, - num_reqs, - BLOCK_SIZE: tl.constexpr, - MAMBA_BLOCK_SIZE: tl.constexpr, -): - """Advance running state columns and emit any required pre-copy inputs.""" - rows = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - mask = rows < num_reqs - req_indices = tl.load(idx_mapping_ptr + rows, mask=mask, other=0) - - state_idx = tl.load(state_idx_ptr + req_indices, mask=mask, other=-1) - num_computed = tl.load(num_computed_tokens_ptr + req_indices, mask=mask, other=0) - # A new request is seeded here, after the resolved MambaSpec is available. - # cache_config.block_size can be smaller than the actual Mamba block when - # an FP16 DFlash cache shares the hybrid page pool, so add_request cannot - # derive the state column from the global cache block size. - seeded_state_idx = (num_computed + MAMBA_BLOCK_SIZE - 1) // MAMBA_BLOCK_SIZE - 1 - state_idx = tl.where(state_idx == -2, seeded_state_idx, state_idx) - num_accepted = tl.load(num_accepted_tokens_ptr + req_indices, mask=mask, other=1) - token_bias = tl.maximum(num_accepted - 1, 0) - tl.store(src_col_ptr + req_indices, state_idx, mask=mask) - tl.store(token_bias_ptr + req_indices, token_bias, mask=mask) - - query_start = tl.load(query_start_loc_ptr + rows, mask=mask, other=0) - query_end = tl.load(query_start_loc_ptr + rows + 1, mask=mask, other=0) - computed_after = num_computed + query_end - query_start - new_state_idx = (computed_after + MAMBA_BLOCK_SIZE - 1) // MAMBA_BLOCK_SIZE - 1 - tl.store(state_idx_ptr + req_indices, new_state_idx, mask=mask) - crossed_boundary = (state_idx >= 0) & (state_idx != new_state_idx) - tl.store( - num_accepted_tokens_ptr + req_indices, - 1, - mask=mask & crossed_boundary, - ) - - -@triton.jit -def _precopy_mamba_align_kernel( - state_idx_ptr, - src_col_ptr, - token_bias_ptr, - block_table_ptrs_ptr, - block_table_stride_req, - state_base_addrs_ptr, - state_block_strides_ptr, - state_elem_sizes_ptr, - state_inner_sizes_ptr, - state_conv_widths_ptr, - state_group_indices_ptr, - idx_mapping_ptr, - num_reqs, - COPY_BLOCK_SIZE: tl.constexpr, - TEMPORAL_TILES: tl.constexpr, -): - batch_idx = tl.program_id(0) - state_idx = tl.program_id(1) - tile_idx = tl.program_id(2) - if batch_idx >= num_reqs: - return - req_idx = tl.load(idx_mapping_ptr + batch_idx) - if req_idx < 0: - return - - src_col = tl.load(src_col_ptr + req_idx) - dst_col = tl.load(state_idx_ptr + req_idx) - if src_col < 0 or src_col == dst_col: - return - token_bias = tl.load(token_bias_ptr + req_idx) - _copy_mamba_state_block( - state_idx, - batch_idx, - src_col, - dst_col, - token_bias, - block_table_ptrs_ptr, - block_table_stride_req, - state_base_addrs_ptr, - state_block_strides_ptr, - state_elem_sizes_ptr, - state_inner_sizes_ptr, - state_conv_widths_ptr, - state_group_indices_ptr, - tile_idx, - COPY_BLOCK_SIZE, - TEMPORAL_TILES, - ) - - -@triton.jit -def _postprocess_mamba_align_kernel( - num_accepted_snapshot_ptr, - num_accepted_out_ptr, - state_idx_ptr, - new_num_computed_tokens_ptr, - block_table_ptrs_ptr, - block_table_stride_req, - state_base_addrs_ptr, - state_block_strides_ptr, - state_elem_sizes_ptr, - state_inner_sizes_ptr, - state_conv_widths_ptr, - state_group_indices_ptr, - idx_mapping_ptr, - num_reqs, - MAMBA_BLOCK_SIZE: tl.constexpr, - COPY_BLOCK_SIZE: tl.constexpr, - TEMPORAL_TILES: tl.constexpr, -): - batch_idx = tl.program_id(0) - state_meta_idx = tl.program_id(1) - tile_idx = tl.program_id(2) - if batch_idx >= num_reqs: - return - req_idx = tl.load(idx_mapping_ptr + batch_idx) - if req_idx < 0: - return - - num_accepted = tl.load(num_accepted_snapshot_ptr + req_idx) - src_col = tl.load(state_idx_ptr + req_idx) - new_num_computed = tl.load(new_num_computed_tokens_ptr + req_idx) - num_tokens_running_state = new_num_computed - num_accepted + 1 - aligned_new_computed = (new_num_computed // MAMBA_BLOCK_SIZE) * MAMBA_BLOCK_SIZE - if aligned_new_computed < num_tokens_running_state: - return - - token_bias = aligned_new_computed - num_tokens_running_state - dst_col = aligned_new_computed // MAMBA_BLOCK_SIZE - 1 - if src_col == dst_col and state_meta_idx == 0 and tile_idx == 0: - tl.store(num_accepted_out_ptr + req_idx, 1) - if src_col == dst_col and token_bias == 0: - return - - _copy_mamba_state_block( - state_meta_idx, - batch_idx, - src_col, - dst_col, - token_bias, - block_table_ptrs_ptr, - block_table_stride_req, - state_base_addrs_ptr, - state_block_strides_ptr, - state_elem_sizes_ptr, - state_inner_sizes_ptr, - state_conv_widths_ptr, - state_group_indices_ptr, - tile_idx, - COPY_BLOCK_SIZE, - TEMPORAL_TILES, - ) - - -def run_mamba_align_precopy( - ctx: MambaSpecDecodeGPUContext, - num_reqs: int, - state_idx: torch.Tensor, - src_col: torch.Tensor, - token_bias: torch.Tensor, - idx_mapping: torch.Tensor, -) -> 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) - _precopy_mamba_align_kernel[grid]( - state_idx, - src_col, - token_bias, - ctx.block_table_ptrs, - ctx.block_table_stride_req, - ctx.state_base_addrs, - ctx.state_block_strides, - ctx.state_elem_sizes, - ctx.state_inner_sizes, - ctx.state_conv_widths, - ctx.state_group_indices, - idx_mapping, - num_reqs, - COPY_BLOCK_SIZE=1024, - TEMPORAL_TILES=_TEMPORAL_TILES, - ) - - -def run_mamba_align_postprocess( - ctx: MambaSpecDecodeGPUContext, - num_reqs: int, - num_accepted_tokens: torch.Tensor, - state_idx: torch.Tensor, - new_num_computed_tokens: torch.Tensor, - idx_mapping: torch.Tensor, -) -> None: - if num_reqs == 0 or not ctx.is_initialized: - return - # Snapshot the decision array because one program may reset a request's - # 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) - _postprocess_mamba_align_kernel[grid]( - snapshot, - num_accepted_tokens, - state_idx, - new_num_computed_tokens, - ctx.block_table_ptrs, - ctx.block_table_stride_req, - ctx.state_base_addrs, - ctx.state_block_strides, - ctx.state_elem_sizes, - ctx.state_inner_sizes, - ctx.state_conv_widths, - ctx.state_group_indices, - idx_mapping, - num_reqs, - MAMBA_BLOCK_SIZE=ctx.block_size, - COPY_BLOCK_SIZE=1024, - TEMPORAL_TILES=_TEMPORAL_TILES, - ) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 47d1da107e..5cba66e5c9 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -47,7 +47,6 @@ from vllm.tasks import SupportedTask from vllm.utils.math_utils import cdiv from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib -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 @@ -96,19 +95,14 @@ from vllm.v1.worker.gpu.sample.sampler import Sampler from vllm.v1.worker.gpu.shutdown import free_before_shutdown from vllm.v1.worker.gpu.spec_decode import init_speculator -from vllm.v1.worker.gpu.spec_decode.dflash2.sparse_rejection import ( - try_dflash2_sparse_target_rejection, -) from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( set_eagle3_aux_hidden_state_layers, ) from vllm.v1.worker.gpu.spec_decode.rejection_sampler import RejectionSampler -from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator from vllm.v1.worker.gpu.spec_decode.utils import DraftTokensHandler from vllm.v1.worker.gpu.states import RequestState from vllm.v1.worker.gpu.structured_outputs import StructuredOutputsWorker from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin -from vllm.v1.worker.utils import KVBlockZeroer logger = init_logger(__name__) @@ -181,14 +175,11 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): if self.is_last_pp_rank: self.speculator = init_speculator(self.vllm_config, self.device) - if self.speculative_config.method in ("eagle3", "dflash"): - # EAGLE3 and DFlash may require auxiliary target hidden states. + if self.speculative_config.method == "eagle3": + # EAGLE3 may require auxiliary hidden states from target model outputs. self.use_aux_hidden_state_outputs = True if self.use_pp: - raise ValueError( - f"{self.speculative_config.method} with pipeline parallel " - "is not supported." - ) + raise ValueError("EAGLE3 with pipeline parallel is not supported.") # Draft tokens propagation - for spec-dec + struct outputs. self.draft_tokens_handler = DraftTokensHandler(self.device) @@ -212,20 +203,6 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): max_num_tokens=self.max_num_tokens, device=self.device, ) - # Prefix-anchored SWA: persistent GPU buffer for per-request prompt - # lengths (stable device address across steps). - self.prefix_anchor_lens_buffer: torch.Tensor | None = None - if ( - getattr( - self.vllm_config.attention_config, - "prefix_anchored_decode_window", - None, - ) - is not None - ): - self.prefix_anchor_lens_buffer = torch.zeros( - self.max_num_reqs, dtype=torch.int32, device=self.device - ) self.sampler: Sampler | None = None self.rejection_sampler: RejectionSampler | None = None @@ -419,7 +396,6 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: self.attn_groups, attn_cg_support, kernel_block_sizes = init_attn_backend( self.kv_cache_config, self.vllm_config, self.device ) - self._kernel_block_sizes = kernel_block_sizes self.block_tables = BlockTables( block_sizes=block_sizes, max_num_reqs=self.max_num_reqs, @@ -448,21 +424,11 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: cudagraph_mode, decode_query_len=self.decode_query_len, ) - check_attention_cp_compatibility(self.vllm_config) - if isinstance(self.speculator, DraftModelSpeculator): - self.speculator.set_attn( - self.model_state, - self.kv_cache_config, - self.block_tables, - self.input_buffers, - self.attn_groups, - ) - # DFlash sizes its graph mode from the selected draft attention - # backend, which is available only after set_attn(). - self.speculator.init_cudagraph_manager(cudagraph_mode) - elif self.speculator is not None: - # Preserve the existing Eagle initialization order. + if self.speculator is not None: self.speculator.init_cudagraph_manager(cudagraph_mode) + + check_attention_cp_compatibility(self.vllm_config) + if self.speculator is not None: # HACK(woosuk) self.speculator.set_attn( self.model_state, self.kv_cache_config, self.block_tables @@ -481,24 +447,6 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: ) self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict) - def _init_kv_zero_meta(self) -> None: - """Precompute metadata used to clear newly allocated cache blocks.""" - self._kv_block_zeroer = KVBlockZeroer( - self.device, pin_memory=is_pin_memory_available() - ) - self._kv_block_zeroer.init_meta( - attn_groups_iter=(group for groups in self.attn_groups for group in groups), - kernel_block_sizes=self._kernel_block_sizes, - cache_dtype=self.cache_config.cache_dtype, - runner_only_attn_layers=set(), - static_forward_context=self.compilation_config.static_forward_context, - ) - - def _zero_block_ids(self, block_ids: list[int]) -> None: - """Clear cache blocks before their first use after allocation.""" - if hasattr(self, "_kv_block_zeroer"): - self._kv_block_zeroer.zero_block_ids(block_ids) - @torch.inference_mode() @step_eplb_after(is_dummy=True) def _dummy_run( @@ -702,10 +650,7 @@ def capture_model(self) -> int: use_aux_hidden_state_outputs=self.use_aux_hidden_state_outputs, ) if self.speculator is not None: - if isinstance(self.speculator, DraftModelSpeculator): - self.speculator.capture() - else: - self.speculator.capture(captured_attn_states) + self.speculator.capture(captured_attn_states) end_time = time.perf_counter() end_free_gpu_memory = torch.cuda.mem_get_info()[0] @@ -875,13 +820,6 @@ def prepare_inputs( query_start_loc_np = query_start_loc_np[: num_reqs_padded + 1] query_start_loc = self.input_buffers.query_start_loc[: num_reqs_padded + 1] is_prefilling_np = self.req_states.is_prefilling(idx_mapping_np) - computed_prefill_lens = self.req_states.num_computed_prefill_tokens[ - idx_mapping_np - ] - prefill_lens = self.req_states.prefill_len.np[idx_mapping_np] - is_incomplete_prefilling_np = is_prefilling_np & ( - computed_prefill_lens + num_scheduled_tokens < prefill_lens - ) # Get prefill tokens if any. if np.any(is_prefilling_np): @@ -940,16 +878,6 @@ def prepare_inputs( out=seq_lens_cpu_upper_bound_np[:num_reqs], ) seq_lens_cpu_upper_bound = torch.from_numpy(seq_lens_cpu_upper_bound_np) - - prefix_anchor_lens = None - if self.prefix_anchor_lens_buffer is not None: - prefix_anchor_lens = self.prefix_anchor_lens_buffer[:num_reqs_padded] - prefix_anchor_lens[:num_reqs] = self.req_states.prompt_len.gpu[ - idx_mapping[:num_reqs] - ] - if num_reqs_padded > num_reqs: - prefix_anchor_lens[num_reqs:].zero_() - return InputBatch( req_ids=req_ids, num_reqs=num_reqs, @@ -975,8 +903,6 @@ def prepare_inputs( cu_num_logits=cu_num_logits, cu_num_logits_np=cu_num_logits_np, has_structured_output_reqs=scheduler_output.has_structured_output_requests, - prefix_anchor_lens=prefix_anchor_lens, - is_incomplete_prefilling_np=is_incomplete_prefilling_np, ) def prepare_attn( @@ -1013,44 +939,31 @@ def sample( grammar_output: GrammarOutput | None, ) -> tuple[SamplerOutput, torch.Tensor, torch.Tensor]: sample_hidden_states = hidden_states[input_batch.logits_indices] - sampler_output = None - if input_batch.num_draft_tokens > 0: + logits = self.model.compute_logits(sample_hidden_states) + if grammar_output is not None: + # Apply grammar bitmask to the logits in-place. + assert self.structured_outputs_worker is not None + self.structured_outputs_worker.apply_grammar_bitmask( + logits, + input_batch, + grammar_output.structured_output_request_ids, + grammar_output.grammar_bitmask, + ) + + if input_batch.num_draft_tokens == 0: + # No draft tokens (common case). + assert self.sampler is not None + sampler_output = self.sampler(logits, input_batch) + else: + # Rejection sampling for spec decoding. assert self.rejection_sampler is not None assert self.speculator is not None - sampler_output = try_dflash2_sparse_target_rejection( - self.model, - self.speculator, - self.rejection_sampler, - sample_hidden_states, + sampler_output = self.rejection_sampler( + logits, input_batch, - grammar_output, + # Draft logits are needed for probabilistic rejection sampling. + self.speculator.draft_logits, ) - if sampler_output is None: - logits = self.model.compute_logits(sample_hidden_states) - if grammar_output is not None: - # Apply grammar bitmask to the logits in-place. - assert self.structured_outputs_worker is not None - self.structured_outputs_worker.apply_grammar_bitmask( - logits, - input_batch, - grammar_output.structured_output_request_ids, - grammar_output.grammar_bitmask, - ) - - if input_batch.num_draft_tokens == 0: - # No draft tokens (common case). - assert self.sampler is not None - sampler_output = self.sampler(logits, input_batch) - else: - # Rejection sampling for spec decoding. - assert self.rejection_sampler is not None - assert self.speculator is not None - sampler_output = self.rejection_sampler( - logits, - input_batch, - # Draft logits are needed for probabilistic rejection sampling. - self.speculator.draft_logits, - ) # Get the number of sampled and rejected tokens. # For chunked prefills, num_sampled and num_rejected are both 0. @@ -1089,11 +1002,7 @@ def postprocess( self.req_states.total_len.gpu, ) - self.model_state.postprocess_state( - input_batch, - num_sampled, - self.req_states.num_computed_tokens.gpu, - ) + self.model_state.postprocess_state(input_batch, num_sampled) @torch.inference_mode() def execute_model( @@ -1111,8 +1020,6 @@ def execute_model( self.add_requests(scheduler_output) self.update_requests(scheduler_output) self.block_tables.apply_staged_writes() - if scheduler_output.new_block_ids_to_zero: - self._zero_block_ids(scheduler_output.new_block_ids_to_zero) if scheduler_output.total_num_scheduled_tokens == 0: # No need to run the model. empty_output = self.kv_connector.no_forward(scheduler_output) @@ -1151,14 +1058,6 @@ def execute_model( # Prepare all the inputs and copy to the input buffers. input_batch = self.prepare_inputs(scheduler_output, batch_desc) block_tables, slot_mappings = self.prepare_attn(input_batch) - # Hybrid Mamba align-mode prefix caching migrates recurrent state - # across block boundaries before attention metadata consumes it. - self.model_state.preprocess_state( - input_batch, - block_tables, - self.kv_cache_config, - self.req_states.num_computed_tokens.gpu, - ) if self.lora_config: # Activate LoRA adapters. @@ -1420,9 +1319,6 @@ def sample_tokens( if self.speculator is not None: assert self.sampler is not None - requires_host_token_state = bool( - getattr(self.speculator, "requires_host_token_state", False) - ) # Let the target override the hidden state fed to the drafter # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The # target returns a persistent buffer sized at max_num_batched_tokens; @@ -1444,24 +1340,6 @@ def sample_tokens( self.sampler.sampling_states.temperature.gpu, self.sampler.sampling_states.seeds.gpu, mm_inputs=mm_inputs, - output_copy_event=( - async_output.copy_event if requires_host_token_state else None - ), - sampled_token_ids_cpu=( - async_output.sampled_token_ids - if requires_host_token_state - else None - ), - num_sampled_tokens_cpu=( - async_output.num_sampled_tokens_np - if requires_host_token_state - else None - ), - all_token_ids_cpu=( - self.req_states.all_token_ids.get_cpu_view().numpy() - if requires_host_token_state - else None - ), ) self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens self.draft_tokens_handler.set_draft_tokens(input_batch, draft_tokens) diff --git a/vllm/v1/worker/gpu/model_states/default.py b/vllm/v1/worker/gpu/model_states/default.py index f6b62dfcd5..7f7955a58a 100644 --- a/vllm/v1/worker/gpu/model_states/default.py +++ b/vllm/v1/worker/gpu/model_states/default.py @@ -195,6 +195,5 @@ def prepare_attn( dcp_local_seq_lens=input_batch.dcp_local_seq_lens, positions=input_batch.positions, for_cudagraph_capture=for_capture, - prefix_anchor_lens=input_batch.prefix_anchor_lens, ) return attn_metadata diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py index 0e268af17b..721e5c2013 100644 --- a/vllm/v1/worker/gpu/model_states/interface.py +++ b/vllm/v1/worker/gpu/model_states/interface.py @@ -56,21 +56,10 @@ def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: def apply_staged_writes(self) -> None: return None - def preprocess_state( - self, - input_batch: InputBatch, - block_tables: tuple[torch.Tensor, ...], - kv_cache_config: KVCacheConfig, - num_computed_tokens: torch.Tensor, - ) -> None: - """Hook for model-specific state movement before a real forward.""" - return None - def postprocess_state( self, input_batch: InputBatch, num_sampled: torch.Tensor, - num_computed_tokens: torch.Tensor | None = None, ) -> None: return None diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 140cf97413..93115fdf64 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -7,64 +7,24 @@ import torch import torch.nn as nn -from vllm import envs from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode -from vllm.logger import init_logger -from vllm.model_executor.layers.mamba.mamba_utils import ( - get_conv_copy_spec, - is_conv_state_dim_first, -) -from vllm.platforms import current_platform -from vllm.triton_utils import triton -from vllm.utils.platform_utils import is_pin_memory_available -from vllm.v1.attention.backends.flash_attn_v100 import ( - DFlash2SmallQGroupDescriptor, - DFlash2SmallQPreparedMetadata, - FlashAttnV100MetadataBuilder, - prepare_dflash2_smallq_group_metadata, -) -from vllm.v1.attention.backends.gdn_attn import ( - DFlash2GDNGroupDescriptor, - GDNAttentionMetadata, - GDNAttentionMetadataBuilder, - prepare_dflash2_gdn_group_metadata, -) +from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadataBuilder -from vllm.v1.core.sched.output import NewRequestData -from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec -from vllm.v1.utils import CpuGpuBuffer -from vllm.v1.worker.gpu.attn_utils import ( - CommonGDNSpecMetadata, - build_attn_metadata, - compute_common_gdn_attn_metadata, -) +from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.worker.gpu.attn_utils import build_attn_metadata from vllm.v1.worker.gpu.input_batch import InputBatch -from vllm.v1.worker.gpu.mamba_align import ( - preprocess_mamba_align_fused_kernel, - run_mamba_align_postprocess, - run_mamba_align_precopy, -) from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache 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.utils import AttentionGroup -logger = init_logger(__name__) - @dataclass class MambaHybridAttnMetadata(ModelSpecificAttnMetadata): is_prefilling: torch.Tensor num_accepted_tokens: torch.Tensor | None = None num_decode_draft_tokens_cpu: torch.Tensor | None = None - common_gdn_metadata: CommonGDNSpecMetadata | None = None - prepared_dflash2_gdn_metadata: dict[int, GDNAttentionMetadata] | None = None - prepared_dflash2_smallq_metadata: ( - dict[int, DFlash2SmallQPreparedMetadata] | None - ) = None def get_extra_common_attn_kwargs( self, @@ -78,22 +38,12 @@ def get_extra_attn_kwargs( attn_metadata_builder: Any, num_reqs: int, ) -> dict[str, Any]: - if isinstance(attn_metadata_builder, FlashAttnV100MetadataBuilder): - return { - "prepared_dflash2_smallq_metadata": ( - None - if self.prepared_dflash2_smallq_metadata is None - else self.prepared_dflash2_smallq_metadata.get( - id(attn_metadata_builder) - ) - ) - } if not isinstance( attn_metadata_builder, (Mamba2AttentionMetadataBuilder, GDNAttentionMetadataBuilder), ): return {} - kwargs = { + return { "num_accepted_tokens": None if self.num_accepted_tokens is None else self.num_accepted_tokens[:num_reqs], @@ -101,13 +51,6 @@ def get_extra_attn_kwargs( if self.num_decode_draft_tokens_cpu is None else self.num_decode_draft_tokens_cpu[:num_reqs], } - if isinstance(attn_metadata_builder, GDNAttentionMetadataBuilder): - kwargs["common_gdn_metadata"] = self.common_gdn_metadata - if self.prepared_dflash2_gdn_metadata is not None: - kwargs["prepared_dflash2_metadata"] = ( - self.prepared_dflash2_gdn_metadata.get(id(attn_metadata_builder)) - ) - return kwargs class MambaHybridModelState(DefaultModelState): @@ -121,187 +64,9 @@ def __init__( device: torch.device, ) -> None: super().__init__(vllm_config, model, encoder_cache, device) - self.cache_config = vllm_config.cache_config self.num_accepted_tokens_gpu = torch.ones( self.max_num_reqs, dtype=torch.int32, device=self.device ) - self._use_dflash2_common_gdn_metadata = bool( - envs.VLLM_SM70_DFLASH2_VERIFY_FASTPATH - and uses_dflash_selector_engine(vllm_config) - ) - if self._use_dflash2_common_gdn_metadata: - logger.info_once("DFlash2 shared GDN batch metadata fast path enabled.") - self._use_dflash2_fused_gdn_metadata = bool( - self._use_dflash2_common_gdn_metadata - and envs.VLLM_SM70_DFLASH2_FUSED_GDN_METADATA - and self.cache_config.mamba_cache_mode in ("none", "align") - and device.type == "cuda" - and current_platform.is_device_capability(70) - ) - self._dflash2_gdn_builders: ( - list[tuple[int, GDNAttentionMetadataBuilder]] | None - ) = None - self._dflash2_gdn_group_descriptor: DFlash2GDNGroupDescriptor | None = None - self._dflash2_fused_gdn_metadata_logged = False - self._use_dflash2_grouped_smallq_metadata = bool( - self._use_dflash2_common_gdn_metadata - and envs.VLLM_SM70_DFLASH2_FUSED_SMALLQ_METADATA - and envs.VLLM_SM70_DFLASH2_GROUPED_SMALLQ_METADATA - and device.type == "cuda" - and current_platform.is_device_capability(70) - ) - self._dflash2_smallq_builders: ( - list[tuple[int, FlashAttnV100MetadataBuilder]] | None - ) = None - self._dflash2_smallq_group_descriptor: DFlash2SmallQGroupDescriptor | None = ( - None - ) - self._dflash2_grouped_smallq_metadata_logged = False - self._align_mode = self.cache_config.mamba_cache_mode == "align" - if self._align_mode: - self._mamba_state_idx_gpu = torch.full( - (self.max_num_reqs,), -2, dtype=torch.int32, device=self.device - ) - self._mamba_src_col_gpu = torch.full( - (self.max_num_reqs,), -1, dtype=torch.int32, device=self.device - ) - self._mamba_token_bias_gpu = torch.zeros( - self.max_num_reqs, dtype=torch.int32, device=self.device - ) - self._mamba_ctx: MambaSpecDecodeGPUContext | None = None - self._mamba_group_ids: list[int] = [] - self._mamba_spec: MambaSpec | None = None - - def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: - super().add_request(req_index, new_req_data) - # A recycled request slot may contain the previous request's accepted - # count. The neutral value is required in both align and non-align mode. - self.num_accepted_tokens_gpu[req_index].fill_(1) - if self._align_mode: - # Delay state-column seeding until preprocess_state has the - # resolved MambaSpec block size. -2 is distinct from -1 (fresh - # request with no computed state). - self._mamba_state_idx_gpu[req_index].fill_(-2) - - 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] - return self._mamba_group_ids, self._mamba_spec - - def _ensure_align_ctx( - self, - kv_cache_config: KVCacheConfig, - mamba_group_ids: list[int], - block_tables: tuple[torch.Tensor, ...], - ) -> MambaSpecDecodeGPUContext: - if self._mamba_ctx is None: - copy_funcs = self.model.get_mamba_state_copy_func() - # 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(): - raise ValueError( - "MRV2 align prefix caching with speculative decoding requires " - "the default SD Mamba conv-state layout on this V100 path." - ) - self._mamba_ctx = MambaSpecDecodeGPUContext.create( - max_num_reqs=self.max_num_reqs, - kv_cache_config=kv_cache_config, - num_state_types=len(copy_funcs), - device=self.device, - make_buffer=lambda n, dtype: CpuGpuBuffer( - n, - dtype=dtype, - device=self.device, - pin_memory=is_pin_memory_available(), - ), - ) - ctx = self._mamba_ctx - if not ctx.is_initialized: - 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(), - [block_tables[group_id] for group_id in mamba_group_ids], - ) - return ctx - - def preprocess_state( - self, - input_batch: InputBatch, - block_tables: tuple[torch.Tensor, ...], - kv_cache_config: KVCacheConfig, - num_computed_tokens: torch.Tensor, - ) -> None: - """Move align-mode recurrent state before a real MRV2 forward.""" - if not self._align_mode or input_batch.num_reqs == 0: - return - mamba_group_ids, mamba_spec = self._get_mamba_group_info(kv_cache_config) - ctx = self._ensure_align_ctx(kv_cache_config, mamba_group_ids, block_tables) - - block = 256 - preprocess_mamba_align_fused_kernel[ - (triton.cdiv(input_batch.num_reqs, block),) - ]( - input_batch.idx_mapping, - self._mamba_state_idx_gpu, - num_computed_tokens, - input_batch.query_start_loc, - self.num_accepted_tokens_gpu, - self._mamba_src_col_gpu, - self._mamba_token_bias_gpu, - input_batch.num_reqs, - BLOCK_SIZE=block, - MAMBA_BLOCK_SIZE=mamba_spec.block_size, - ) - run_mamba_align_precopy( - ctx, - input_batch.num_reqs, - self._mamba_state_idx_gpu, - self._mamba_src_col_gpu, - self._mamba_token_bias_gpu, - input_batch.idx_mapping, - ) - - def _get_dflash2_gdn_builders( - self, - attn_groups: list[list[AttentionGroup]], - ) -> list[tuple[int, GDNAttentionMetadataBuilder]]: - if self._dflash2_gdn_builders is None: - builders: list[tuple[int, GDNAttentionMetadataBuilder]] = [] - for kv_cache_group_id, groups in enumerate(attn_groups): - for group in groups: - builder = group.get_metadata_builder(0) - if isinstance(builder, GDNAttentionMetadataBuilder): - builders.append((kv_cache_group_id, builder)) - self._dflash2_gdn_builders = builders - return self._dflash2_gdn_builders - - def _get_dflash2_smallq_builders( - self, - attn_groups: list[list[AttentionGroup]], - ) -> list[tuple[int, FlashAttnV100MetadataBuilder]]: - if self._dflash2_smallq_builders is None: - builders: list[tuple[int, FlashAttnV100MetadataBuilder]] = [] - for kv_cache_group_id, groups in enumerate(attn_groups): - for group in groups: - builder = group.get_metadata_builder(0) - if isinstance(builder, FlashAttnV100MetadataBuilder): - builders.append((kv_cache_group_id, builder)) - self._dflash2_smallq_builders = builders - return self._dflash2_smallq_builders def prepare_attn( self, @@ -331,9 +96,6 @@ def prepare_attn( # compute them during actual (non-capture) forward execution. num_accepted_tokens = None num_decode_draft_tokens_cpu = None - common_gdn_metadata = None - prepared_dflash2_gdn_metadata = None - prepared_dflash2_smallq_metadata = None if not for_capture: num_accepted_tokens = self.num_accepted_tokens_gpu.new_ones(num_reqs) num_accepted_tokens[: input_batch.num_reqs] = self.num_accepted_tokens_gpu[ @@ -353,93 +115,13 @@ def prepare_attn( -1, ) num_decode_draft_tokens_cpu = torch.from_numpy(num_decode_draft_tokens_np) - if self._use_dflash2_common_gdn_metadata: - speculative_config = self.vllm_config.speculative_config - assert speculative_config is not None - common_gdn_metadata = compute_common_gdn_attn_metadata( - num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, - query_start_loc=input_batch.query_start_loc, - query_start_loc_cpu=query_start_loc_cpu, - num_spec_state_tokens=( - speculative_config.num_speculative_state_tokens() - ), - legacy_mixed_decode_routing=( - envs.VLLM_SM70_MTP_LEGACY_GDN_MIXED_DECODE_ROUTING - ), - ) - - if ( - self._use_dflash2_fused_gdn_metadata - and cudagraph_mode == CUDAGraphMode.FULL - and common_gdn_metadata is not None - ): - prepared_result = prepare_dflash2_gdn_group_metadata( - builders_by_group=self._get_dflash2_gdn_builders(attn_groups), - block_tables=block_tables, - common_gdn_metadata=common_gdn_metadata, - num_accepted_tokens=num_accepted_tokens, - num_actual_tokens=num_tokens, - descriptor=self._dflash2_gdn_group_descriptor, - state_start_indices=( - self._mamba_state_idx_gpu if self._align_mode else None - ), - req_index_mapping=( - input_batch.idx_mapping if self._align_mode else None - ), - ) - if prepared_result is not None: - ( - prepared_dflash2_gdn_metadata, - self._dflash2_gdn_group_descriptor, - ) = prepared_result - if not self._dflash2_fused_gdn_metadata_logged: - logger.info( - "DFlash2 fused GDN metadata active for %d cache groups.", - len(prepared_dflash2_gdn_metadata), - ) - self._dflash2_fused_gdn_metadata_logged = True - - if ( - self._use_dflash2_grouped_smallq_metadata - and cudagraph_mode == CUDAGraphMode.FULL - and common_gdn_metadata is not None - ): - seq_lens_cpu = input_batch.seq_lens_cpu_upper_bound[:num_reqs] - max_seq_len_hint = int(seq_lens_cpu.max().item()) - prepared_smallq_result = prepare_dflash2_smallq_group_metadata( - builders_by_group=self._get_dflash2_smallq_builders(attn_groups), - block_tables=block_tables, - seq_lens=input_batch.seq_lens, - query_start_loc=input_batch.query_start_loc, - query_start_loc_cpu=query_start_loc_cpu, - num_reqs=num_reqs, - num_query_tokens=num_tokens, - max_seq_len_hint=max_seq_len_hint, - workspace_seq_capacity_cap=self.max_model_len, - descriptor=self._dflash2_smallq_group_descriptor, - ) - if prepared_smallq_result is not None: - ( - prepared_dflash2_smallq_metadata, - self._dflash2_smallq_group_descriptor, - ) = prepared_smallq_result - if not self._dflash2_grouped_smallq_metadata_logged: - logger.info( - "DFlash2 grouped small-query metadata active for %d " - "full-attention cache groups.", - len(prepared_dflash2_smallq_metadata), - ) - self._dflash2_grouped_smallq_metadata_logged = True mamba_attn_metadata = MambaHybridAttnMetadata( is_prefilling=is_prefilling, num_accepted_tokens=num_accepted_tokens, num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, - common_gdn_metadata=common_gdn_metadata, - prepared_dflash2_gdn_metadata=prepared_dflash2_gdn_metadata, - prepared_dflash2_smallq_metadata=prepared_dflash2_smallq_metadata, ) - attn_metadata = build_attn_metadata( + return build_attn_metadata( attn_groups=attn_groups, num_reqs=num_reqs, num_tokens=num_tokens, @@ -451,49 +133,18 @@ def prepare_attn( block_tables=block_tables, slot_mappings=slot_mappings, kv_cache_config=kv_cache_config, - seq_lens_cpu_upper_bound=input_batch.seq_lens_cpu_upper_bound, dcp_local_seq_lens=input_batch.dcp_local_seq_lens, model_specific_attn_metadata=mamba_attn_metadata, for_cudagraph_capture=for_capture, - prefix_anchor_lens=input_batch.prefix_anchor_lens, ) - if common_gdn_metadata is not None: - assert common_gdn_metadata.spec_query_start_loc.numel() == ( - common_gdn_metadata.num_spec_decodes + 1 - ) - # Metadata kernels/copies and FULL CUDA Graph replay are enqueued on - # the same current stream, which already supplies the required - # device ordering. Retain the old value check as an opt-in debug - # fence while the no-sync route is quality- and trace-gated. - if envs.VLLM_SM70_DFLASH2_GDN_SYNC_ASSERT: - assert ( - common_gdn_metadata.spec_query_start_loc[-1].item() - == common_gdn_metadata.num_spec_decode_tokens - ) - return attn_metadata def postprocess_state( self, input_batch: InputBatch, num_sampled: torch.Tensor, - num_computed_tokens: torch.Tensor | None = None, ) -> None: # Chunked prefill does not sample a token, so num_sampled can be 0. # Mamba treats num_accepted_tokens=1 as the neutral non-spec value. self.num_accepted_tokens_gpu[input_batch.idx_mapping] = torch.clamp( num_sampled, min=1 ) - if ( - self._align_mode - and num_computed_tokens is not None - and self._mamba_ctx is not None - and input_batch.num_reqs > 0 - ): - run_mamba_align_postprocess( - self._mamba_ctx, - input_batch.num_reqs, - self.num_accepted_tokens_gpu, - self._mamba_state_idx_gpu, - num_computed_tokens, - input_batch.idx_mapping, - ) diff --git a/vllm/v1/worker/gpu/model_states/whisper.py b/vllm/v1/worker/gpu/model_states/whisper.py index 191b930a2e..0ef3cadc87 100644 --- a/vllm/v1/worker/gpu/model_states/whisper.py +++ b/vllm/v1/worker/gpu/model_states/whisper.py @@ -161,7 +161,6 @@ def prepare_attn( dcp_local_seq_lens=input_batch.dcp_local_seq_lens, model_specific_attn_metadata=whisper_attn_metadata, for_cudagraph_capture=for_capture, - prefix_anchor_lens=input_batch.prefix_anchor_lens, ) return attn_metadata diff --git a/vllm/v1/worker/gpu/sample/gumbel.py b/vllm/v1/worker/gpu/sample/gumbel.py index 77fb4a16f2..aaa49283d3 100644 --- a/vllm/v1/worker/gpu/sample/gumbel.py +++ b/vllm/v1/worker/gpu/sample/gumbel.py @@ -2,16 +2,18 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch -from vllm.triton_utils import HAS_TRITON, tl, tldevice, triton +from vllm.triton_utils import HAS_TRITON, tl, triton -# Smallest positive value produced by Triton's fp32 `tl.rand`. Used to clamp -# zero draws before the flipped Gumbel transform below. +# Smallest positive normal fp32 value. Used to clamp the uniform draw so that +# `log(u)` cannot produce -inf (and thus `-log(-log(u))` stays finite). # # Triton requires globals accessed from `@triton.jit` functions to be wrapped # in `tl.constexpr(...)`. We can only do that when Triton is actually # available — on the CPU worker path `tl` is a placeholder whose `constexpr` # attribute is `None`, and `tl.constexpr(...)` would crash at import time. -_TL_RAND_MIN = tl.constexpr(4.6566127342e-10) if HAS_TRITON else 4.6566127342e-10 +_FP32_TINY = ( + tl.constexpr(float.fromhex("0x1p-126")) if HAS_TRITON else float.fromhex("0x1p-126") +) @triton.jit @@ -73,46 +75,6 @@ def tl_rand64(seed, offset, includes_zero: tl.constexpr): return u -@triton.jit -def tl_rand32(seed, offset, includes_zero: tl.constexpr): - u = tl.rand(seed, offset) - if not includes_zero: - u = tl.maximum(u, _TL_RAND_MIN) - return u - - -@triton.jit -def gumbel_noised_argmax( - logits, - keys, - mask, - seed, - pos, - temp, - USE_FP64: tl.constexpr, - APPLY_TEMPERATURE: tl.constexpr = True, -): - """Argmax under token-keyed Gumbel noise, or plain argmax at temp zero.""" - if temp != 0.0 and APPLY_TEMPERATURE: - logits = logits / temp - - if USE_FP64: - logits = logits.to(tl.float64) - if temp != 0.0: - gumbel_seed = tl.randint(seed, pos) - if USE_FP64: - u = tl_rand64(gumbel_seed, keys, includes_zero=False) - gumbel_noise = -tl.log(-tl.log(u)) - else: - u = tl_rand32(gumbel_seed, keys, includes_zero=False) - # Flip the transform so the deciding tail lies near zero, where - # fp32 uniforms retain resolution. `log1p` avoids cancellation. - gumbel_noise = -tl.log(-tldevice.log1p(-u)) - logits = tl.where(mask, logits + gumbel_noise, float("-inf")) - - return tl.max(logits, axis=0, return_indices=True) - - @triton.jit def gumbel_block_argmax( logits, @@ -125,24 +87,21 @@ def gumbel_block_argmax( pos_ptr, processed_logits_ptr, processed_logits_stride, - processed_logits_col_stride, processed_logits_col_ptr, vocab_size, APPLY_TEMPERATURE: tl.constexpr, USE_FP64: tl.constexpr, ): - req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx).to(tl.int64) - is_valid_req = req_state_idx >= 0 - temp = tl.load(temp_ptr + req_state_idx, mask=is_valid_req, other=0.0).to( - tl.float32 - ) + req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx) + temp = tl.load(temp_ptr + req_state_idx).to(tl.float32) if temp != 0.0 and APPLY_TEMPERATURE: - # Preserve the established MRV2 contract: proposal logits handed to - # rejection sampling have already had request temperature applied. + # Apply temperature. + # NOTE(woosuk): Match the behavior of _temperature_kernel. + # E.g., if the kernel uses tl.div_rn, we should use tl.div_rn here too. logits = logits / temp if processed_logits_ptr is not None: - # Store the temperature-applied proposal distribution. + # Store the temperature-applied logits. if processed_logits_col_ptr is not None: col = tl.load(processed_logits_col_ptr) else: @@ -150,24 +109,34 @@ def gumbel_block_argmax( tl.store( processed_logits_ptr + req_state_idx * processed_logits_stride - + col * processed_logits_col_stride + + col * vocab_size + block, logits, - mask=mask & is_valid_req, + mask=mask, ) - seed = tl.load(seeds_ptr + req_state_idx, mask=is_valid_req, other=0) - pos = tl.load(pos_ptr + token_idx) - return gumbel_noised_argmax( - logits, - block, - mask & is_valid_req, - seed, - pos, - temp, - USE_FP64=USE_FP64, - APPLY_TEMPERATURE=False, - ) + # fp32 is the default reduction dtype; fp64 is ~1/32–1/64x the throughput + # on H100/Ada/Blackwell and empirically indistinguishable for Gumbel-max. + if USE_FP64: + logits = logits.to(tl.float64) + if temp != 0.0: + # Calculate the seed for gumbel noise. + seed = tl.load(seeds_ptr + req_state_idx) + pos = tl.load(pos_ptr + token_idx) + gumbel_seed = tl.randint(seed, pos) + + if USE_FP64: + u = tl_rand64(gumbel_seed, block, includes_zero=False) + else: + u = tl.rand(gumbel_seed, block) + u = tl.maximum(u, _FP32_TINY) + gumbel_noise = -tl.log(-tl.log(u)) + + # Apply gumbel noise. + logits = tl.where(mask, logits + gumbel_noise, float("-inf")) + + value, idx = tl.max(logits, axis=0, return_indices=True) + return value, idx @triton.jit @@ -178,7 +147,6 @@ def _gumbel_sample_kernel( local_max_stride, processed_logits_ptr, processed_logits_stride, - processed_logits_col_stride, processed_logits_col_ptr, logits_ptr, logits_stride, @@ -213,7 +181,6 @@ def _gumbel_sample_kernel( pos_ptr, processed_logits_ptr, processed_logits_stride, - processed_logits_col_stride, processed_logits_col_ptr, vocab_size, APPLY_TEMPERATURE=APPLY_TEMPERATURE, @@ -248,7 +215,6 @@ def gumbel_sample( local_max.stride(0), output_processed_logits, output_processed_logits.stride(0) if output_processed_logits is not None else 0, - output_processed_logits.stride(1) if output_processed_logits is not None else 0, output_processed_logits_col, logits, logits.stride(0), diff --git a/vllm/v1/worker/gpu/spec_decode/__init__.py b/vllm/v1/worker/gpu/spec_decode/__init__.py index ee25697099..536b7526bd 100644 --- a/vllm/v1/worker/gpu/spec_decode/__init__.py +++ b/vllm/v1/worker/gpu/spec_decode/__init__.py @@ -1,51 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Mapping - import torch from vllm.config import VllmConfig -def uses_dflash_selector_engine(vllm_config: object) -> bool: - """Return whether the configured DFlash algorithm exposes selector inputs. - - Dispatch is based on the inference algorithm's configuration contract, not - a model, architecture, or checkpoint identity. - """ - speculative_config = getattr(vllm_config, "speculative_config", None) - if ( - speculative_config is None - or getattr(speculative_config, "method", None) != "dflash" - ): - return False - 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 {} - if not isinstance(dflash_config, Mapping): - return False - try: - return int(dflash_config.get("selector_top_k", 0) or 0) > 0 - except (TypeError, ValueError): - return False - - def init_speculator(vllm_config: VllmConfig, device: torch.device): speculative_config = vllm_config.speculative_config assert speculative_config is not None - if speculative_config.method == "dflash": - draft_config = speculative_config.draft_model_config - if draft_config is None: - raise ValueError("method='dflash' requires a draft model config") - if uses_dflash_selector_engine(vllm_config): - from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import ( - DFlash2Speculator, - ) - - return DFlash2Speculator(vllm_config, device) - from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator - - return DFlashSpeculator(vllm_config, device) if speculative_config.use_eagle(): from vllm.v1.worker.gpu.spec_decode.eagle.speculator import EagleSpeculator diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/__init__.py b/vllm/v1/worker/gpu/spec_decode/dflash/__init__.py deleted file mode 100644 index 208f01a7cb..0000000000 --- a/vllm/v1/worker/gpu/spec_decode/dflash/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py b/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py deleted file mode 100644 index cfbfc5aeb9..0000000000 --- a/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py +++ /dev/null @@ -1,125 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Callable, Mapping - -import torch - -from vllm.config.compilation import CUDAGraphMode -from vllm.v1.kv_cache_interface import KVCacheConfig -from vllm.v1.worker.gpu.attn_utils import ( - build_attn_metadata, - build_slot_mappings_by_layer, -) -from vllm.v1.worker.gpu.block_table import BlockTables -from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens -from vllm.v1.worker.gpu.cudagraph_utils import ( - BatchExecutionDescriptor, - CapturedAttentionState, - CudaGraphManager, -) -from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers -from vllm.v1.worker.utils import AttentionGroup - - -def _prepare_dflash_inputs_to_capture( - num_reqs: int, - num_tokens: int, - input_buffers: InputBuffers, - block_tables: BlockTables, - attn_groups: list[list[AttentionGroup]], - kv_cache_config: KVCacheConfig, - max_model_len: int, - skip_attn: bool, - causal: bool | Mapping[int, bool], -) -> CapturedAttentionState: - input_batch = InputBatch.make_dummy(num_reqs, num_tokens, input_buffers) - input_block_tables = block_tables.get_dummy_block_tables(num_reqs) - slot_mappings = block_tables.get_dummy_slot_mappings(num_tokens) - slot_mappings_by_layer = build_slot_mappings_by_layer( - slot_mappings, kv_cache_config - ) - - attn_metadata = None - if not skip_attn: - query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np) - dcp_local_seq_lens = None - if block_tables.cp_size > 1: - prepare_dcp_local_seq_lens( - input_buffers.dcp_local_seq_lens, - input_buffers.seq_lens, - num_reqs, - block_tables.cp_size, - block_tables.cp_rank, - block_tables.cp_interleave, - ) - dcp_local_seq_lens = input_buffers.dcp_local_seq_lens - attn_metadata = build_attn_metadata( - attn_groups=attn_groups, - num_reqs=num_reqs, - num_tokens=num_tokens, - query_start_loc_gpu=input_batch.query_start_loc, - query_start_loc_cpu=query_start_loc_cpu, - max_query_len=num_tokens // num_reqs, - seq_lens=input_batch.seq_lens, - dcp_local_seq_lens=dcp_local_seq_lens, - max_seq_len=max_model_len, - block_tables=input_block_tables, - slot_mappings=slot_mappings, - kv_cache_config=kv_cache_config, - for_cudagraph_capture=True, - causal=causal, - ) - return CapturedAttentionState(attn_metadata, slot_mappings_by_layer) - - -class DFlashCudaGraphManager(CudaGraphManager): - """DFlash CudaGraphManager for the parallel-drafting query forward, - building its own attention metadata from scratch.""" - - def capture( - self, - forward_fn: Callable, - input_buffers: InputBuffers, - block_tables: BlockTables, - attn_groups: list[list[AttentionGroup]], - kv_cache_config: KVCacheConfig, - max_model_len: int, - causal: bool | Mapping[int, bool], - progress_bar_desc: str = "Capturing CUDA graphs", - ) -> None: - def create_forward_fn( - desc: BatchExecutionDescriptor, - ) -> tuple[Callable[[CUDAGraphMode], None], CapturedAttentionState]: - num_tokens = desc.num_tokens - num_reqs = desc.num_reqs or min(num_tokens, self.max_num_reqs) - num_tokens_across_dp = ( - torch.full((self.dp_size,), num_tokens, dtype=torch.int32, device="cpu") - if self.dp_size > 1 - else None - ) - attn_state = _prepare_dflash_inputs_to_capture( - num_reqs, - num_tokens, - input_buffers, - block_tables, - attn_groups, - kv_cache_config, - max_model_len, - skip_attn=(desc.cg_mode == CUDAGraphMode.PIECEWISE), - causal=causal, - ) - attn_metadata, slot_mappings = attn_state - - return ( - lambda cg_mode: forward_fn( - num_reqs, - num_tokens, - attn_metadata, - slot_mappings, - num_tokens_across_dp, - cg_mode, - ), - attn_state, - ) - - super().capture(create_forward_fn, progress_bar_desc) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py deleted file mode 100644 index 9d08b77ee5..0000000000 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ /dev/null @@ -1,829 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Mapping -from typing import Any - -import numpy as np -import torch -import torch.nn as nn - -from vllm.config import VllmConfig, replace -from vllm.config.compilation import CUDAGraphMode -from vllm.forward_context import BatchDescriptor, set_forward_context -from vllm.logger import init_logger -from vllm.triton_utils import tl, triton -from vllm.v1.attention.backend import AttentionCGSupport -from vllm.v1.attention.backends.utils import PAD_SLOT_ID -from vllm.v1.kv_cache_interface import KVCacheConfig -from vllm.v1.utils import record_function_or_nullcontext -from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer -from vllm.v1.worker.gpu.block_table import BlockTables -from vllm.v1.worker.gpu.cp_utils import cp_local_slot, prepare_dcp_local_seq_lens -from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp -from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers -from vllm.v1.worker.gpu.model_states.interface import ModelState -from vllm.v1.worker.gpu.spec_decode.dflash.cudagraph import DFlashCudaGraphManager -from vllm.v1.worker.gpu.spec_decode.dflash.utils import load_dflash_model -from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator -from vllm.v1.worker.gpu.spec_decode.utils import get_parallel_drafting_token_id -from vllm.v1.worker.utils import AttentionGroup - -logger = init_logger(__name__) - - -def _is_context_only_prefill(input_batch: InputBatch) -> bool: - incomplete_prefill = input_batch.is_incomplete_prefilling_np - return bool( - incomplete_prefill is not None - and incomplete_prefill.size == input_batch.num_reqs - and np.all(incomplete_prefill) - ) - - -class DFlashSpeculator(DraftModelSpeculator): - _speculator_name = "DFlash" # For logging, so we can share methods with subclasses - - def __init__(self, vllm_config: VllmConfig, device: torch.device): - super().__init__(vllm_config, device) - - self.hidden_states = torch.zeros( - self.max_num_tokens, self.hidden_size, dtype=self.dtype, device=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 - - self.parallel_drafting_token_id = get_parallel_drafting_token_id( - self.draft_model_config.hf_config - ) - - from vllm.model_executor.models.qwen3_dflash import dflash_has_any_non_causal - - self.requires_non_causal = dflash_has_any_non_causal( - self.draft_model_config.hf_config - ) - - # Whether the anchor query position is itself a prediction. DFlash default uses - # the anchor as the bonus token (only mask tokens predict); DSpark samples from - # the anchor and the N-1 mask token positions. See _prepare_dflash_inputs_kernel - dflash_config = ( - getattr(self.draft_model_config.hf_config, "dflash_config", None) or {} - ) - if dflash_config.get("sample_from_anchor", False): - raise ValueError( - "sample_from_anchor=True is not supported for DFlash. " - "DFlash uses a fixed 1+N query layout where the anchor " - "is the bonus token." - ) - self.sample_from_anchor = False - - # Context positions for the K/V precompute. Populated by - # prepare_dflash_inputs, and processed by the model's - # precompute_and_store_context_kv method. NOT captured by CUDA graphs. - self.context_positions = torch.zeros( - self.max_num_tokens, dtype=torch.int64, device=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 - self.sample_indices = torch.zeros( - max_num_sampled_tokens, dtype=torch.int64, device=device - ) - self.sample_pos = torch.zeros( - max_num_sampled_tokens, dtype=torch.int64, device=device - ) - # -1 marks an inert sampling row. CUDA graph capture can execute the - # full buffer before a real batch has populated it, so zero would make - # every padding row scatter into request slot 0. - self.sample_idx_mapping = torch.full( - (max_num_sampled_tokens,), -1, dtype=torch.int32, device=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 - ).repeat(self.max_num_reqs) - - self.query_cudagraph_manager: DFlashCudaGraphManager | None = None - self.draft_kv_cache_group_id: int = -1 - self._context_only_prefill_logged = False - - @property - def attn_vllm_config(self) -> VllmConfig: - # The draft's attention differs from the target's in causality. - return replace( - self.vllm_config, - model_config=self.draft_model_config, - attention_config=replace( - self.vllm_config.attention_config, - use_non_causal=self.requires_non_causal, - ), - ) - - def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: - wants_full = cudagraph_mode.decode_mode() == CUDAGraphMode.FULL - supports_full = ( - self.attn_cg_support.min_cg_support.value - >= AttentionCGSupport.UNIFORM_BATCH.value - ) - if wants_full and not supports_full: - logger.warning( - "%s draft attention (%s) does not support full CUDA graphs; " - "running the draft eagerly.", - self._speculator_name, - self.attn_cg_support.min_cg_attn_backend, - ) - # PIECEWISE cudagraphs are not supported for dflash. - if wants_full and supports_full: - cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY - else: - cudagraph_mode = CUDAGraphMode.NONE - - self.query_cudagraph_manager = DFlashCudaGraphManager( - self.vllm_config, - self.device, - cudagraph_mode, - decode_query_len=self.num_query_per_req, - ) - - def capture(self) -> None: - logger.info("Capturing model for %s speculator...", self._speculator_name) - # Padded sample rows must not scatter into a live request during capture. - self.sample_indices.zero_() - self.sample_pos.zero_() - self.sample_idx_mapping.fill_(-1) - assert self.query_cudagraph_manager is not None - self.query_cudagraph_manager.capture( - self._generate_draft, - self.input_buffers, - self.block_tables, - self.attn_groups, - self.kv_cache_config, - self.max_model_len, - causal=self._group_causal, - progress_bar_desc=f"Capturing {self._speculator_name.lower()} CUDA graphs", - ) - - def load_draft_model( - self, - target_model: nn.Module, - target_attn_layer_names: set[str], - ) -> nn.Module: - return load_dflash_model(target_model, self.vllm_config) - - def set_attn( - self, - model_state: ModelState, - kv_cache_config: KVCacheConfig, - block_tables: BlockTables, - target_input_buffers: InputBuffers, - target_attn_groups: list[list[AttentionGroup]], - ) -> None: - super().set_attn( - model_state, - kv_cache_config, - block_tables, - target_input_buffers, - target_attn_groups, - ) - - self.draft_kv_cache_group_ids = [ - gid for gid, g in enumerate(self.attn_groups) if g - ] - assert self.draft_kv_cache_group_ids, "No draft attention groups found." - self.draft_kv_cache_group_id = self.draft_kv_cache_group_ids[0] - - # Per-group context slot buffers for the precompute (one row per group). - self._context_slot_mappings = torch.zeros( - len(self.draft_kv_cache_group_ids), - self.max_num_tokens, - dtype=torch.int64, - device=self.device, - ) - - # Map each draft decoder layer to the index (within draft_kv_cache_group_ids) - # of the kv-cache group its cache belongs to. Models that share a single group - # leave this as None and share one context slot mapping. - self._layer_group_idx: list[int] | None = None - # Per-KV-group causal, falling back to whether the drafter is all-causal. - self._group_causal: dict[int, bool] | bool = not self.requires_non_causal - if hasattr(self.model, "get_draft_kv_cache_layer_names"): - layer_names = self.model.get_draft_kv_cache_layer_names() - name_to_gid = { - ln: gid - for gid, group in enumerate(kv_cache_config.kv_cache_groups) - for ln in group.layer_names - } - gid_to_idx = {gid: i for i, gid in enumerate(self.draft_kv_cache_group_ids)} - self._layer_group_idx = [ - gid_to_idx[name_to_gid[name]] for name in layer_names - ] - if hasattr(self.model, "get_draft_attn_causal"): - self._group_causal = { - name_to_gid[name]: layer_causal - for name, layer_causal in zip( - layer_names, self.model.get_draft_attn_causal() - ) - } - - @torch.inference_mode() - def _run_model( - self, - num_tokens: int, - attn_metadata: dict[str, Any] | None, - slot_mappings: dict[str, torch.Tensor] | None, - num_tokens_across_dp: torch.Tensor | None, - cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE, - ) -> torch.Tensor: - batch_descriptor = BatchDescriptor(num_tokens=num_tokens) - with set_forward_context( - attn_metadata, - self.vllm_config, - num_tokens=num_tokens, - cudagraph_runtime_mode=cudagraph_runtime_mode, - num_tokens_across_dp=num_tokens_across_dp, - slot_mapping=slot_mappings, - batch_descriptor=batch_descriptor, - ): - last_hidden_states = self.model( - input_ids=self.input_buffers.input_ids[:num_tokens], - positions=self.input_buffers.positions[:num_tokens], - inputs_embeds=None, - ) - return last_hidden_states - - def _generate_draft( - self, - num_reqs: int, - num_tokens_padded: int, - attn_metadata: dict[str, Any] | None, - slot_mappings: dict[str, torch.Tensor] | None, - num_tokens_across_dp: torch.Tensor | None, - cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE, - ) -> None: - last_hidden_states = self._run_model( - num_tokens_padded, - attn_metadata, - slot_mappings, - num_tokens_across_dp, - cudagraph_runtime_mode, - ) - num_sample = num_reqs * self.num_speculative_steps - 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. - draft_tokens = self.sample_draft( - sample_hidden_states, - self.sample_pos[:num_sample] - 2, - self.sample_idx_mapping[:num_sample], - self.temperature, - self.seeds, - self.sample_col[:num_sample], - self.draft_logits, - ) - self.draft_tokens[:num_reqs] = draft_tokens.view( - num_reqs, self.num_speculative_steps - ) - - def _build_draft_attn_metadata( - self, - num_reqs: int, - num_reqs_padded: int, - num_tokens_padded: int, - seq_lens_cpu_upper_bound: torch.Tensor, - step: int, - num_query_per_req: int | None = None, - causal: bool | Mapping[int, bool] = False, - query_start_loc_np: np.ndarray | None = None, - dcp_local_seq_lens: torch.Tensor | None = None, - ) -> dict[str, Any] | None: - if not self.draft_attn_layer_names: - return None - assert num_query_per_req is None # Omitted for DFlash, read from self instead - if dcp_local_seq_lens is None and self.block_tables.cp_size > 1: - prepare_dcp_local_seq_lens( - self.input_buffers.dcp_local_seq_lens, - self.input_buffers.seq_lens, - num_reqs, - self.block_tables.cp_size, - self.block_tables.cp_rank, - self.block_tables.cp_interleave, - ) - dcp_local_seq_lens = self.input_buffers.dcp_local_seq_lens - return super()._build_draft_attn_metadata( - num_reqs, - num_reqs_padded, - num_tokens_padded, - seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, - step=step, - num_query_per_req=self.num_query_per_req, - causal=causal, - query_start_loc_np=query_start_loc_np, - dcp_local_seq_lens=dcp_local_seq_lens, - ) - - def _prepare_ngram_assist( - self, - input_batch: InputBatch, - output_copy_event: torch.cuda.Event | None, - sampled_token_ids_cpu: np.ndarray | None, - num_sampled_tokens_cpu: np.ndarray | None, - all_token_ids_cpu: np.ndarray | None, - ) -> bool: - """Prepare an optional draftless proposal; return whether it is complete.""" - return False - - def _apply_ngram_assist(self, num_reqs: int) -> None: - """Override model proposals for ngram-hit rows, if configured.""" - - @torch.inference_mode() - def propose( - self, - input_batch: InputBatch, - attn_metadata: dict[str, Any], - slot_mappings: dict[str, torch.Tensor], - # [num_tokens, hidden_size] - last_hidden_states: torch.Tensor, - # num_layers x [num_tokens, hidden_size] - aux_hidden_states: list[torch.Tensor] | None, - # [num_reqs] - num_sampled: torch.Tensor, - # [num_reqs] - num_rejected: torch.Tensor, - # [max_num_reqs] - last_sampled: torch.Tensor, - # [max_num_reqs] - next_prefill_tokens: torch.Tensor, - # [max_num_reqs] - temperature: torch.Tensor, - # [max_num_reqs] - seeds: torch.Tensor, - num_tokens_across_dp: torch.Tensor | None = None, - dummy_run: bool = False, - skip_attn_for_dummy_run: bool = False, - mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, - is_profile: bool = False, - output_copy_event: torch.cuda.Event | None = None, - sampled_token_ids_cpu: np.ndarray | None = None, - num_sampled_tokens_cpu: np.ndarray | None = None, - all_token_ids_cpu: np.ndarray | None = None, - ) -> torch.Tensor: - num_reqs = input_batch.num_reqs - num_target_tokens = input_batch.num_tokens - num_query_tokens = num_reqs * self.num_query_per_req - max_seq_len = input_batch.seq_lens_cpu_upper_bound[:num_reqs].max().item() - self.draft_max_seq_len = min( - max_seq_len + self.num_query_per_req, self.max_model_len - ) - - # NOTE: To avoid CPU-GPU synchronization without CPU knowing the - # number of rejected tokens, we maintain the size of input_ids and - # hidden_states the same as the target model's. This means, we pad each - # request's query length to include any rejected positions. - if aux_hidden_states: - with record_function_or_nullcontext("dflash: concatenate target hidden"): - combined_target_hidden = torch.cat(aux_hidden_states, dim=-1) - with record_function_or_nullcontext("dflash: project target hidden"): - hidden_states = self.model.combine_hidden_states(combined_target_hidden) - else: - hidden_states = last_hidden_states - with record_function_or_nullcontext("dflash: stage target hidden"): - self.hidden_states[:num_target_tokens].copy_( - hidden_states[:num_target_tokens] - ) - - if dummy_run and skip_attn_for_dummy_run: - # Memory profiling path: block_tables / kv_cache_config are not initialized. - # Since DFlash needs to build its own attention metadata, we must skip the - # preparation in this path and run a minimal forward pass. - self.model.precompute_and_store_context_kv( - self.hidden_states[:num_target_tokens], - self.context_positions[:num_target_tokens], - ) - self._generate_draft( - num_reqs, - num_query_tokens, - attn_metadata=None, - slot_mappings=None, - num_tokens_across_dp=num_tokens_across_dp, - cudagraph_runtime_mode=CUDAGraphMode.NONE, - ) - return self.draft_tokens[:num_reqs] - - # The query slot mapping is written into the shared BlockTables slot_mappings. - # That buffer's address is what the captured CUDA graph reads from at replay. - assert self.draft_kv_cache_group_id >= 0 - # Support multiple draft KV cache groups by preparing inputs once for each - with record_function_or_nullcontext("dflash: prepare inputs"): - for i, gid in enumerate(self.draft_kv_cache_group_ids): - prepare_dflash_inputs( - self.input_buffers, - self.block_tables.slot_mappings[gid], - self.context_positions, - self._context_slot_mappings[i], - self.sample_indices, - self.sample_pos, - self.sample_idx_mapping, - self.temperature, - self.seeds, - input_batch, - num_sampled, - num_rejected, - last_sampled, - next_prefill_tokens, - temperature, - seeds, - self.block_tables.input_block_tables[gid], - self.block_tables.kernel_block_sizes[gid], - self.block_tables.cp_rank, - self.block_tables.cp_size, - self.block_tables.cp_interleave, - self.parallel_drafting_token_id, - self.num_query_per_req, - self.num_speculative_steps, - self.max_num_reqs, - self.max_num_tokens, - self.max_model_len, - self.sample_from_anchor, - ) - - # Pre-insert context K/V into the cache. Runs eagerly outside the captured graph - # because the context shape varies per step. During dummy runs the block tables - # are placeholders, so we skip the cache write to avoid clobbering real entries. - # Each layer uses the context slots of its own kv-cache group. - if dummy_run: - context_slots: torch.Tensor | list[torch.Tensor | None] | None = None - elif self._layer_group_idx is not None: - context_slots = [ - self._context_slot_mappings[gidx][:num_target_tokens] - for gidx in self._layer_group_idx - ] - else: - context_slots = self._context_slot_mappings[0][:num_target_tokens] - with record_function_or_nullcontext("dflash: materialize context kv"): - self.model.precompute_and_store_context_kv( - self.hidden_states[:num_target_tokens], - self.context_positions[:num_target_tokens], - context_slots, - ) - - if not dummy_run and _is_context_only_prefill(input_batch): - # Intermediate chunked-prefill steps only need to materialize the - # target hidden states into draft KV. Their proposed tokens cannot - # be consumed before the final prompt chunk, which will run the - # normal query graph and replace this sentinel output. - self.draft_tokens[:num_reqs].fill_(-1) - if not self._context_only_prefill_logged: - logger.info( - "%s context-only path is skipping unused draft queries " - "for intermediate prefill chunks.", - self._speculator_name, - ) - self._context_only_prefill_logged = True - return self.draft_tokens[:num_reqs] - - if self._prepare_ngram_assist( - input_batch, - output_copy_event, - sampled_token_ids_cpu, - num_sampled_tokens_cpu, - all_token_ids_cpu, - ): - self._apply_ngram_assist(num_reqs) - return self.draft_tokens[:num_reqs] - - with record_function_or_nullcontext("dflash: query and selector"): - # Every DFlash step has exactly num_query_per_req tokens, so we can - # use FULL CUDA graphs. - batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp( - self.query_cudagraph_manager, - num_reqs, - num_query_tokens, - uniform_token_count=self.num_query_per_req, - dp_size=self.dp_size, - dp_rank=self.dp_rank, - need_eager=is_profile, - ) - - num_reqs_padded = batch_desc.num_reqs or num_reqs - num_tokens_padded = batch_desc.num_tokens - - # Rebuild the draft attention metadata even when replaying the FULL - # graph so that any attention metadata builder state is updated. - draft_attn_metadata = self._build_draft_attn_metadata( - num_reqs=num_reqs, - num_reqs_padded=num_reqs_padded, - num_tokens_padded=num_tokens_padded, - seq_lens_cpu_upper_bound=input_batch.seq_lens_cpu_upper_bound, - step=self.num_query_per_req, - causal=self._group_causal, - ) - draft_slot_mappings_by_layer = build_slot_mappings_by_layer( - self.block_tables.slot_mappings[:, :num_tokens_padded], - self.kv_cache_config, - ) - - if batch_desc.cg_mode == CUDAGraphMode.FULL: - assert self.query_cudagraph_manager is not None - self.query_cudagraph_manager.run_fullgraph(batch_desc) - else: - self._generate_draft( - num_reqs, - num_tokens_padded, - draft_attn_metadata, - draft_slot_mappings_by_layer, - num_tokens_across_dp=num_tokens_across_dp, - cudagraph_runtime_mode=batch_desc.cg_mode, - ) - - self._apply_ngram_assist(num_reqs) - - return self.draft_tokens[:num_reqs] - - -@triton.jit -def _prepare_dflash_inputs_kernel( - # Outputs - out_input_ids_ptr, - out_query_positions_ptr, - out_query_start_loc_ptr, - out_seq_lens_ptr, - out_query_slot_mapping_ptr, - out_context_positions_ptr, - out_context_slot_mapping_ptr, - out_sample_indices_ptr, - out_sample_pos_ptr, - out_sample_idx_mapping_ptr, - out_temperature_ptr, - out_seeds_ptr, - # Inputs from target batch - target_positions_ptr, - target_query_start_loc_ptr, - idx_mapping_ptr, - last_sampled_ptr, - next_prefill_tokens_ptr, - num_sampled_ptr, - num_rejected_ptr, - # Sampling params - temperature_ptr, - seeds_ptr, - # Block table for slot mapping lookup. - block_table_ptr, - block_table_stride, - # Scalars - parallel_drafting_token_id, - block_size, - num_query_per_req, - num_speculative_steps, - max_num_reqs, - max_num_tokens, - max_model_len, - cp_rank, - SAMPLE_FROM_ANCHOR: tl.constexpr, - PAD_SLOT_ID: tl.constexpr, - CP_SIZE: tl.constexpr, - CP_INTERLEAVE: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - req_idx = tl.program_id(0) - block_idx = tl.program_id(1) - num_reqs = tl.num_programs(0) - req_state_idx = tl.load(idx_mapping_ptr + req_idx) - - ctx_start = tl.load(target_query_start_loc_ptr + req_idx) - ctx_end = tl.load(target_query_start_loc_ptr + req_idx + 1) - num_ctx = ctx_end - ctx_start - - num_rejected = tl.load(num_rejected_ptr + req_idx) - valid_ctx_end = ctx_end - num_rejected - num_valid_ctx = valid_ctx_end - ctx_start - - num_sampled = tl.load(num_sampled_ptr + req_idx) - if num_sampled > 0: - bonus_token = tl.load(last_sampled_ptr + req_state_idx).to(tl.int32) - else: - # Chunked prefilling: splice in the next prefill token. - bonus_token = tl.load(next_prefill_tokens_ptr + req_state_idx).to(tl.int32) - - last_valid_pos = tl.load(target_positions_ptr + valid_ctx_end - 1) - query_base = req_idx * num_query_per_req - - j = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - is_ctx = j < num_ctx - is_valid_ctx = j < num_valid_ctx - is_query = (j >= num_valid_ctx) & (j < num_valid_ctx + num_query_per_req) - query_off = j - num_valid_ctx - - # --- Context positions / slots --- - ctx_pos_idx = ctx_start + tl.where(is_ctx, j, 0) - ctx_pos = tl.load(target_positions_ptr + ctx_pos_idx, mask=is_valid_ctx, other=0) - ctx_block_num = ctx_pos // (block_size * CP_SIZE) - ctx_block_num = tl.minimum(ctx_block_num, block_table_stride - 1) - ctx_block_id = tl.load( - block_table_ptr + req_idx * block_table_stride + ctx_block_num, - mask=is_valid_ctx, - other=0, - ).to(tl.int64) - # Block 0 is the null block. Old sliding-window context positions can map - # to it after eviction; rejected suffix rows are invalid context as well. - # Neither kind of row may write draft KV into physical block 0. - ctx_resident = is_valid_ctx & (ctx_block_id != 0) - local_ctx_slot = cp_local_slot( - ctx_pos, ctx_block_id, block_size, cp_rank, CP_SIZE, CP_INTERLEAVE, PAD_SLOT_ID - ) - ctx_slot = tl.where( - ctx_resident, - local_ctx_slot, - PAD_SLOT_ID, - ) - # Stored over the full [0, num_ctx) span while the loads above are masked to - # [0, num_valid_ctx): the rejected suffix rows in between get position 0 and - # PAD_SLOT_ID. That is intentional — those rows write no KV and their - # positions are never consumed, but the span must stay fully initialized so - # a replayed graph cannot observe a stale value from an earlier batch. - tl.store(out_context_positions_ptr + ctx_start + j, ctx_pos, mask=is_ctx) - tl.store(out_context_slot_mapping_ptr + ctx_start + j, ctx_slot, mask=is_ctx) - - # --- Query positions / input_ids / slots --- - query_pos = last_valid_pos + 1 + query_off - query_idx = query_base + query_off - is_bonus = is_query & (query_off == 0) - input_id = tl.where(is_bonus, bonus_token, parallel_drafting_token_id) - - q_block_num = query_pos // (block_size * CP_SIZE) - q_block_num = tl.minimum(q_block_num, block_table_stride - 1) - q_block_id = tl.load( - block_table_ptr + req_idx * block_table_stride + q_block_num, - mask=is_query, - other=0, - ).to(tl.int64) - # A null block is never a writable cache slot. This can occur when a - # sliding-window block table contains evicted/global padding entries. - q_resident = is_query & (q_block_id != 0) - local_q_slot = cp_local_slot( - query_pos, - q_block_id, - block_size, - cp_rank, - CP_SIZE, - CP_INTERLEAVE, - PAD_SLOT_ID, - ) - q_slot = tl.where( - q_resident, - local_q_slot, - PAD_SLOT_ID, - ) - - tl.store(out_input_ids_ptr + query_idx, input_id, mask=is_query) - clamped_query_pos = tl.minimum(query_pos, max_model_len - 1) - tl.store(out_query_positions_ptr + query_idx, clamped_query_pos, mask=is_query) - tl.store(out_query_slot_mapping_ptr + query_idx, q_slot, mask=is_query) - - # --- Sample indices / positions / idx_mapping --- - # When SAMPLE_FROM_ANCHOR (DSpark), so we sample at EVERY query position - # and each position k predicts the NEXT token (sampled position = query_pos + 1). - # Otherwise (DFlash default) the anchor is the bonus token and only the mask tokens - # at offsets > 0 are sampled from, each AT its own position. - sample_off = 0 if SAMPLE_FROM_ANCHOR else 1 - is_sample = is_query & (query_off >= sample_off) - sample_idx = req_idx * num_speculative_steps + (query_off - sample_off) - sample_pos = query_pos + 1 if SAMPLE_FROM_ANCHOR else query_pos - tl.store(out_sample_indices_ptr + sample_idx, query_idx, mask=is_sample) - tl.store(out_sample_pos_ptr + sample_idx, sample_pos, mask=is_sample) - tl.store(out_sample_idx_mapping_ptr + sample_idx, req_state_idx, mask=is_sample) - - if block_idx == 0: - tl.store(out_query_start_loc_ptr + req_idx, query_base) - # seq_lens is the absolute sequence length the draft attention - # reads up to (context + query), not just the count of accepted - # tokens this step. - tl.store( - out_seq_lens_ptr + req_idx, - tl.minimum(last_valid_pos + 1 + num_query_per_req, max_model_len), - ) - # Copy sampling state. - tl.store( - out_temperature_ptr + req_state_idx, - tl.load(temperature_ptr + req_state_idx), - ) - tl.store(out_seeds_ptr + req_state_idx, tl.load(seeds_ptr + req_state_idx)) - if req_idx == num_reqs - 1: - # Pad per-request buffers to max_num_reqs for CUDA graph safety. - last_query_end = num_reqs * num_query_per_req - for i in range(num_reqs, max_num_reqs + 1, BLOCK_SIZE): - block = i + tl.arange(0, BLOCK_SIZE) - mask = block < max_num_reqs + 1 - tl.store(out_query_start_loc_ptr + block, last_query_end, mask=mask) - for i in range(num_reqs, max_num_reqs, BLOCK_SIZE): - block = i + tl.arange(0, BLOCK_SIZE) - mask = block < max_num_reqs - tl.store(out_seq_lens_ptr + block, 0, mask=mask) - # Padded sample slots point at query index 0 (a valid row in - # last_hidden_states) so CG replay never reads OOB. Padded - # sample idx mappings point to -1, which is ignored during - # sampling to prevent writing stale values to draft logits. - pad_start = num_reqs * num_speculative_steps - pad_end = max_num_reqs * num_speculative_steps - for i in range(pad_start, pad_end, BLOCK_SIZE): - block = i + tl.arange(0, BLOCK_SIZE) - mask = block < pad_end - tl.store(out_sample_indices_ptr + block, 0, mask=mask) - tl.store(out_sample_pos_ptr + block, 0, mask=mask) - tl.store(out_sample_idx_mapping_ptr + block, -1, mask=mask) - # Pad query slot mappings past num_query_tokens with PAD so the - # captured CG sees PAD slots (no K/V write) for replay sizes - # larger than the current request count. - q_pad_start = num_reqs * num_query_per_req - for i in range(q_pad_start, max_num_tokens, BLOCK_SIZE): - block = i + tl.arange(0, BLOCK_SIZE) - mask = block < max_num_tokens - tl.store(out_query_slot_mapping_ptr + block, PAD_SLOT_ID, mask=mask) - - -def prepare_dflash_inputs( - input_buffers: InputBuffers, - query_slot_mapping: torch.Tensor, - context_positions: torch.Tensor, - context_slot_mapping: torch.Tensor, - sample_indices: torch.Tensor, - sample_pos: torch.Tensor, - sample_idx_mapping: torch.Tensor, - temperature: torch.Tensor, - seeds: torch.Tensor, - input_batch: InputBatch, - # [num_reqs] - num_sampled: torch.Tensor, - # [num_reqs] - num_rejected: torch.Tensor, - # [max_num_reqs] - last_sampled: torch.Tensor, - # [max_num_reqs] - next_prefill_tokens: torch.Tensor, - # [max_num_reqs] - input_temperature: torch.Tensor, - # [max_num_reqs] - input_seeds: torch.Tensor, - # [max_num_reqs, max_num_blocks] - block_table: torch.Tensor, - block_size: int, - cp_rank: int, - cp_size: int, - cp_interleave: int, - parallel_drafting_token_id: int, - num_query_per_req: int, - num_speculative_steps: int, - max_num_reqs: int, - max_num_tokens: int, - max_model_len: int, - sample_from_anchor: bool = False, -) -> None: - num_reqs = input_batch.num_reqs - assert num_reqs > 0 - # Cover the longest possible per-request span (ctx + query). Use the max - # per-request query length, not the total token count across the batch. - max_target_query_len = int(input_batch.num_scheduled_tokens.max()) - max_tokens_per_req = max_target_query_len + num_query_per_req - BLOCK_SIZE = min(256, triton.next_power_of_2(max(1, max_tokens_per_req))) - num_blocks = triton.cdiv(max_tokens_per_req, BLOCK_SIZE) - _prepare_dflash_inputs_kernel[(num_reqs, num_blocks)]( - input_buffers.input_ids, - input_buffers.positions, - input_buffers.query_start_loc, - input_buffers.seq_lens, - query_slot_mapping, - context_positions, - context_slot_mapping, - sample_indices, - sample_pos, - sample_idx_mapping, - temperature, - seeds, - input_batch.positions, - input_batch.query_start_loc, - input_batch.idx_mapping, - last_sampled, - next_prefill_tokens, - num_sampled, - num_rejected, - input_temperature, - input_seeds, - block_table, - block_table.stride(0), - parallel_drafting_token_id, - block_size, - num_query_per_req, - num_speculative_steps, - max_num_reqs, - max_num_tokens, - max_model_len, - cp_rank, - SAMPLE_FROM_ANCHOR=sample_from_anchor, - PAD_SLOT_ID=PAD_SLOT_ID, - CP_SIZE=cp_size, - CP_INTERLEAVE=cp_interleave, - BLOCK_SIZE=BLOCK_SIZE, - ) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/utils.py b/vllm/v1/worker/gpu/spec_decode/dflash/utils.py deleted file mode 100644 index f714b3db30..0000000000 --- a/vllm/v1/worker/gpu/spec_decode/dflash/utils.py +++ /dev/null @@ -1,101 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import torch.nn as nn - -from vllm.config import VllmConfig, replace -from vllm.distributed.parallel_state import get_pp_group -from vllm.logger import init_logger -from vllm.model_executor.model_loader import get_model -from vllm.platforms import current_platform -from vllm.v1.worker.gpu.spec_decode.eagle.utils import ( - _should_share, - get_target_lm_head, -) - -logger = init_logger(__name__) - - -def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module: - from vllm.compilation.backends import set_model_tag - from vllm.model_executor.models.qwen3_dflash import ( - dflash_has_any_non_causal, - dflash_target_rope_is_neox_style, - ) - - speculative_config = vllm_config.speculative_config - assert speculative_config is not None - draft_model_config = speculative_config.draft_model_config - draft_cache_dtype = speculative_config.kv_cache_dtype - if ( - draft_cache_dtype is None - and str(vllm_config.cache_config.cache_dtype).startswith("fp8") - and current_platform.is_cuda() - and current_platform.is_device_capability(70) - ): - draft_cache_dtype = "auto" - logger.info_once( - "Using FP16 draft KV cache for SM70 DFlash while the target uses %s.", - vllm_config.cache_config.cache_dtype, - ) - # The drafter must rotate Q/K the way its target does. Take that from the - # built target before super() constructs the draft. - is_neox_style = dflash_target_rope_is_neox_style(target_model) - if is_neox_style is not None: - draft_model_config.hf_config.is_neox_style = is_neox_style - # Select an attention backend that supports the drafter's attention: mixing - # a non-causal layer onto a causal-only backend would fail. - draft_vllm_config = replace( - vllm_config, - attention_config=replace( - vllm_config.attention_config, - use_non_causal=dflash_has_any_non_causal(draft_model_config.hf_config), - backend=speculative_config.attention_backend, - ), - cache_config=( - replace( - vllm_config.cache_config, - cache_dtype=draft_cache_dtype, - ) - if draft_cache_dtype is not None - else vllm_config.cache_config - ), - ) - with set_model_tag("dflash_head"): - dflash_model = get_model( - vllm_config=draft_vllm_config, model_config=draft_model_config - ) - - target_language_model = ( - target_model.get_language_model() - if hasattr(target_model, "get_language_model") - else target_model - ) - # MuseGlimmerForCausalLM marks its inner MuseGlimmerModel as the language - # model, so get_language_model() already returns the inner module and has - # no .model of its own. - target_inner = getattr(target_language_model, "model", target_language_model) - draft_inner = dflash_model.model - - # Skip embedding sharing under PP — each rank owns its own embedding. - if get_pp_group().world_size == 1: - target_embed = getattr(target_inner, "embed_tokens", None) or getattr( - target_inner, "embedding", None - ) - draft_embed = getattr(draft_inner, "embed_tokens", None) - if target_embed is not None and _should_share( - dflash_model, "has_own_embed_tokens", draft_embed, target_embed - ): - if draft_embed is not None: - del draft_inner.embed_tokens - draft_inner.embed_tokens = target_embed - - target_lm_head = get_target_lm_head(target_model, target_language_model) - draft_lm_head = getattr(dflash_model, "lm_head", None) - if target_lm_head is not None and _should_share( - dflash_model, "has_own_lm_head", draft_lm_head, target_lm_head - ): - if draft_lm_head is not None: - del dflash_model.lm_head - dflash_model.lm_head = target_lm_head - - return dflash_model diff --git a/vllm/v1/worker/gpu/spec_decode/dflash2/__init__.py b/vllm/v1/worker/gpu/spec_decode/dflash2/__init__.py deleted file mode 100644 index 208f01a7cb..0000000000 --- a/vllm/v1/worker/gpu/spec_decode/dflash2/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/v1/worker/gpu/spec_decode/dflash2/ngram_assist.py b/vllm/v1/worker/gpu/spec_decode/dflash2/ngram_assist.py deleted file mode 100644 index fe89c55512..0000000000 --- a/vllm/v1/worker/gpu/spec_decode/dflash2/ngram_assist.py +++ /dev/null @@ -1,162 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Host prompt-ngram lookup for the MRV2 DFlash2 assistant.""" - -from __future__ import annotations - -import time - -import numpy as np -from numba import njit - - -@njit(inline="always") -def _split_token( - prefix: np.ndarray, - prefix_len: int, - suffix: np.ndarray, - index: int, -) -> int: - if index < prefix_len: - return int(prefix[index]) - return int(suffix[index - prefix_len]) - - -@njit(cache=True) -def find_split_ngram_proposal( - prefix: np.ndarray, - prefix_len: int, - suffix: np.ndarray, - suffix_len: int, - min_ngram: int, - max_ngram: int, - max_model_len: int, - num_draft_tokens: int, -) -> np.ndarray: - """Run the prompt-lookup KMP over ``prefix + suffix`` without copying it.""" - total_tokens = prefix_len + suffix_len - if total_tokens < min_ngram: - return np.empty(0, dtype=np.int32) - - num_draft_tokens = min(num_draft_tokens, max_model_len - total_tokens) - if num_draft_tokens <= 0: - return np.empty(0, dtype=np.int32) - - # This is the same reverse-KMP policy as the standalone vLLM ngram - # proposer: prefer the longest suffix match and the earliest occurrence. - lps = np.zeros(max_ngram, dtype=np.int32) - longest_ngram = 0 - position = 0 - prev_lps = 0 - i = 1 - while i < total_tokens: - prefix_token = _split_token( - prefix, - prefix_len, - suffix, - total_tokens - 1 - prev_lps, - ) - current_token = _split_token( - prefix, - prefix_len, - suffix, - total_tokens - 1 - i, - ) - if prefix_token == current_token: - prev_lps += 1 - if prev_lps >= longest_ngram: - longest_ngram = prev_lps - position = i - if i < max_ngram: - lps[i] = prev_lps - if prev_lps == max_ngram: - prev_lps = lps[max_ngram - 1] - i += 1 - elif prev_lps != 0: - prev_lps = lps[prev_lps - 1] - else: - i += 1 - - if longest_ngram < min_ngram: - return np.empty(0, dtype=np.int32) - - start = total_tokens - 1 - position + longest_ngram - proposal_len = min(num_draft_tokens, total_tokens - start) - proposal = np.empty(proposal_len, dtype=np.int32) - for j in range(proposal_len): - proposal[j] = _split_token(prefix, prefix_len, suffix, start + j) - return proposal - - -class DFlash2NgramAssist: - """Batch adapter and counters around the split prompt lookup.""" - - def __init__( - self, - min_ngram: int, - max_ngram: int, - num_draft_tokens: int, - max_model_len: int, - ) -> None: - self.min_ngram = min_ngram - self.max_ngram = max_ngram - self.num_draft_tokens = num_draft_tokens - self.max_model_len = max_model_len - self.num_eligible = 0 - self.num_full_hits = 0 - self.lookup_seconds = 0.0 - - # Compile the small split-input kernel at startup, not on the first - # user-visible decode token. - find_split_ngram_proposal( - np.zeros(1, dtype=np.int32), - 1, - np.zeros(1, dtype=np.int64), - 0, - min_ngram, - max_ngram, - max_model_len, - num_draft_tokens, - ) - - def propose( - self, - token_ids_cpu: np.ndarray, - req_state_indices: np.ndarray, - prior_lengths: np.ndarray, - sampled_token_ids: np.ndarray, - num_sampled_tokens: np.ndarray, - eligible: np.ndarray, - output_tokens: np.ndarray, - output_lengths: np.ndarray, - ) -> int: - """Fill fixed-width outputs and return the number of full-width hits.""" - num_reqs = len(req_state_indices) - output_tokens[:num_reqs].fill(0) - output_lengths[:num_reqs].fill(0) - full_hits = 0 - started = time.perf_counter() - for batch_idx in range(num_reqs): - if not eligible[batch_idx]: - continue - self.num_eligible += 1 - req_state_idx = int(req_state_indices[batch_idx]) - proposal = find_split_ngram_proposal( - token_ids_cpu[req_state_idx], - int(prior_lengths[batch_idx]), - sampled_token_ids[batch_idx], - int(num_sampled_tokens[batch_idx]), - self.min_ngram, - self.max_ngram, - self.max_model_len, - self.num_draft_tokens, - ) - proposal_len = len(proposal) - output_lengths[batch_idx] = proposal_len - if proposal_len: - output_tokens[batch_idx, :proposal_len] = proposal - if proposal_len == self.num_draft_tokens: - full_hits += 1 - self.num_full_hits += full_hits - self.lookup_seconds += time.perf_counter() - started - return full_hits diff --git a/vllm/v1/worker/gpu/spec_decode/dflash2/sparse_rejection.py b/vllm/v1/worker/gpu/spec_decode/dflash2/sparse_rejection.py deleted file mode 100644 index 01b42e3968..0000000000 --- a/vllm/v1/worker/gpu/spec_decode/dflash2/sparse_rejection.py +++ /dev/null @@ -1,261 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Strictly gated compact target rejection for DFlash2 on SM70.""" - -from __future__ import annotations - -import os -from typing import TYPE_CHECKING, Any - -import numpy as np -import torch - -import vllm.envs as envs -from vllm.logger import init_logger -from vllm.v1.worker.gpu.sample.output import SamplerOutput -from vllm.v1.worker.gpu.sample.states import NO_LOGPROBS -from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import DFlash2Speculator -from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import ( - dflash2_sparse_topk_rejection_sample, -) - -if TYPE_CHECKING: - from vllm.v1.core.sched.output import GrammarOutput - from vllm.v1.worker.gpu.input_batch import InputBatch - from vllm.v1.worker.gpu.spec_decode.rejection_sampler import RejectionSampler - -logger = init_logger(__name__) - -_TARGET_TOP_K = 20 -_SELECTOR_ALIGNMENT_DUMP_COUNT = 0 -_SELECTOR_ALIGNMENT_STEP = 0 - - -def _parse_alignment_steps(raw_steps: str | None) -> set[int] | None: - if not raw_steps: - return None - steps: set[int] = set() - try: - for item in raw_steps.split(","): - item = item.strip() - if not item: - continue - if "-" in item: - start_text, end_text = item.split("-", 1) - start = int(start_text) - end = int(end_text) - if start < 0 or end < start: - return set() - steps.update(range(start, end + 1)) - else: - step = int(item) - if step < 0: - return set() - steps.add(step) - except ValueError: - return set() - return steps - - -def _safe_dump_tag(raw_tag: str) -> str: - return "".join(ch if ch.isalnum() or ch in "._-" else "_" for ch in raw_tag) - - -def _diagnostic_rank() -> int: - return int(os.getenv("RANK", os.getenv("LOCAL_RANK", "0"))) - - -def _maybe_dump_selector_alignment( - *, - speculator: DFlash2Speculator, - rejection_sampler: RejectionSampler, - input_batch: InputBatch, - target_topk_ids: torch.Tensor, - target_topk_logits: torch.Tensor, - draft_topk_ids: torch.Tensor, - draft_topk_logits: torch.Tensor, - draft_sampled: torch.Tensor, - pos: torch.Tensor, - sampled: torch.Tensor, - num_sampled: torch.Tensor, -) -> None: - """Dump one exact B1 selector/target alignment record when requested.""" - if not envs.VLLM_SPEC_DUMP_ALIGNMENT: - return - if _diagnostic_rank() != 0: - return - - global _SELECTOR_ALIGNMENT_DUMP_COUNT, _SELECTOR_ALIGNMENT_STEP - _SELECTOR_ALIGNMENT_STEP += 1 - if _SELECTOR_ALIGNMENT_DUMP_COUNT >= envs.VLLM_SPEC_DUMP_ALIGNMENT_LIMIT: - return - selected_steps = _parse_alignment_steps(envs.VLLM_SPEC_DUMP_ALIGNMENT_STEPS) - if selected_steps is not None and _SELECTOR_ALIGNMENT_STEP not in selected_steps: - return - - shadow = speculator.get_selector_alignment_shadow() - if shadow is None: - return - shadow_ids, unary_logits, lattice_scores = shadow - req_state = int(input_batch.idx_mapping_np[0]) - packed_row = 0 - sampling_states = rejection_sampler.sampler.sampling_states - - with torch.no_grad(): - draft_sampled_cpu = draft_sampled.detach().cpu() - if bool(torch.all(draft_sampled_cpu == 0).item()): - return - payload = { - "format": "dflash2_selector_alignment_v1", - "rank": _diagnostic_rank(), - "step": _SELECTOR_ALIGNMENT_STEP, - "request_state": req_state, - "selector_top_k": speculator.selector_top_k, - "num_speculative_steps": speculator.num_speculative_steps, - "target_topk_ids": target_topk_ids.detach().cpu(), - "target_topk_logits": target_topk_logits.detach().float().cpu(), - "draft_candidate_ids": draft_topk_ids[req_state].detach().cpu(), - "draft_realized_logits": ( - draft_topk_logits[req_state].detach().float().cpu() - ), - "selector_candidate_ids": shadow_ids[packed_row].detach().cpu(), - "selector_unary_logits": (unary_logits[packed_row].detach().float().cpu()), - "selector_lattice_scores": ( - lattice_scores[packed_row].detach().float().cpu() - ), - "draft_sampled": draft_sampled_cpu, - "positions": pos.detach().cpu(), - "cu_num_logits": input_batch.cu_num_logits.detach().cpu(), - "idx_mapping": input_batch.idx_mapping.detach().cpu(), - "temperature": float(sampling_states.temperature.np[req_state]), - "top_p": float(sampling_states.top_p.np[req_state]), - "top_k": int(sampling_states.top_k.np[req_state]), - "sampled_token_ids": sampled.detach().cpu(), - "num_sampled": num_sampled.detach().cpu(), - } - _SELECTOR_ALIGNMENT_DUMP_COUNT += 1 - dump_dir = os.getenv("VLLM_SPEC_DUMP_ALIGNMENT_DIR", "/tmp") - os.makedirs(dump_dir, exist_ok=True) - tag = _safe_dump_tag(os.getenv("VLLM_SPEC_DUMP_ALIGNMENT_TAG", "")) - tag_part = f"{tag}_" if tag else "" - dump_path = os.path.join( - dump_dir, - f"spec_alignment_dflash2_selector_{tag_part}pid{os.getpid()}_" - f"step{_SELECTOR_ALIGNMENT_STEP:06d}_" - f"{_SELECTOR_ALIGNMENT_DUMP_COUNT}.pt", - ) - torch.save(payload, dump_path) - logger.warning("Dumped DFlash2 selector alignment diagnostics to %s", dump_path) - - -def _supports_sparse_sampling_contract( - rejection_sampler: RejectionSampler, - input_batch: InputBatch, -) -> bool: - """Whether compact logits preserve every requested sampling transform.""" - if rejection_sampler.rejection_sample_method != "standard": - return False - # Start with the single-request path used by the latency target. The - # kernel supports batches, but mixed-request graph validation is a - # separate promotion gate. - if input_batch.num_reqs != 1 or np.any(input_batch.is_prefilling_np): - return False - - sampler = rejection_sampler.sampler - idx = input_batch.idx_mapping_np - states = sampler.sampling_states - temperatures = states.temperature.np[idx] - top_k = states.top_k.np[idx] - top_p = states.top_p.np[idx] - if np.any(temperatures <= 0.0): - return False - if np.any(top_k != _TARGET_TOP_K): - return False - if np.any((top_p <= 0.0) | (top_p > 1.0)): - return False - if np.any(states.min_p.np[idx] != 0.0): - return False - - if np.any(sampler.penalties_state.use_penalty[idx]): - return False - if np.any(sampler.logit_bias_state.use_logit_bias[idx]): - return False - if np.any(sampler.bad_words_state.num_bad_words.np[idx] != 0): - return False - if states.max_num_logprobs(idx) != NO_LOGPROBS: - return False - if sampler.logprob_token_ids_state.max_num_token_ids(idx) != 0: - return False - return not sampler.compute_nans - - -def try_dflash2_sparse_target_rejection( - model: Any, - speculator: Any, - rejection_sampler: RejectionSampler, - sample_hidden_states: torch.Tensor, - input_batch: InputBatch, - grammar_output: GrammarOutput | None, -) -> SamplerOutput | None: - """Sample from compact target/draft supports, or return ``None`` safely.""" - if not envs.VLLM_SM70_DFLASH2_SPARSE_TARGET_REJECTION: - return None - if not isinstance(speculator, DFlash2Speculator): - return None - if grammar_output is not None or input_batch.has_structured_output_reqs: - return None - if sample_hidden_states.device.type != "cuda": - return None - if torch.cuda.get_device_capability(sample_hidden_states.device) != (7, 0): - return None - if not hasattr(model, "get_topk_tokens_and_logits"): - return None - if not _supports_sparse_sampling_contract(rejection_sampler, input_batch): - return None - - sparse_draft_logits = speculator.get_sparse_draft_logits() - if sparse_draft_logits is None: - return None - draft_topk_ids, draft_topk_logits = sparse_draft_logits - target_topk_ids, target_topk_logits = model.get_topk_tokens_and_logits( - sample_hidden_states, - _TARGET_TOP_K, - ) - draft_sampled = input_batch.input_ids[input_batch.logits_indices] - pos = input_batch.positions[input_batch.logits_indices] - sampled, num_sampled = dflash2_sparse_topk_rejection_sample( - target_topk_ids, - target_topk_logits, - draft_topk_ids, - draft_topk_logits, - draft_sampled, - input_batch.cu_num_logits, - pos, - input_batch.idx_mapping, - rejection_sampler.sampler.sampling_states.temperature.gpu, - rejection_sampler.sampler.sampling_states.top_p.gpu, - rejection_sampler.sampler.sampling_states.seeds.gpu, - rejection_sampler.num_speculative_steps, - use_fp64=rejection_sampler.sampler.use_fp64_gumbel, - ) - if envs.VLLM_SPEC_DUMP_ALIGNMENT: - _maybe_dump_selector_alignment( - speculator=speculator, - rejection_sampler=rejection_sampler, - input_batch=input_batch, - target_topk_ids=target_topk_ids, - target_topk_logits=target_topk_logits, - draft_topk_ids=draft_topk_ids, - draft_topk_logits=draft_topk_logits, - draft_sampled=draft_sampled, - pos=pos, - sampled=sampled, - num_sampled=num_sampled, - ) - logger.info_once("Using SM70 DFlash2 compact target top-k rejection sampling.") - return SamplerOutput( - sampled_token_ids=sampled, - logprobs_tensors=None, - num_nans=None, - num_sampled=num_sampled, - ) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py deleted file mode 100644 index 64f987238a..0000000000 --- a/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py +++ /dev/null @@ -1,621 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from typing import Any - -import numpy as np -import torch - -from vllm import envs -from vllm.config import VllmConfig -from vllm.config.compilation import CUDAGraphMode -from vllm.logger import init_logger -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.ngram_assist import DFlash2NgramAssist - -logger = init_logger(__name__) - - -def _requires_sm70_tail(device: torch.device, num_steps: int) -> bool: - """Whether the final dependent selector slot needs its own kernel.""" - return ( - device.type == "cuda" - and num_steps > 1 - and torch.cuda.get_device_capability(device) == (7, 0) - ) - - -@triton.jit -def _selector_walk_kernel( - scores_ptr, - candidate_ptr, - sample_pos_ptr, - req_state_ptr, - temperature_ptr, - seeds_ptr, - tokens_ptr, - realized_scores_ptr, - path_state_ptr, - num_steps: tl.constexpr, - walk_steps: tl.constexpr, - top_k: tl.constexpr, - BLOCK_K: tl.constexpr, - SAMPLE_PROBABILISTIC: tl.constexpr, - USE_FP64: tl.constexpr, -): - row = tl.program_id(0) - offsets = tl.arange(0, BLOCK_K) - mask = offsets < top_k - req_state = tl.load(req_state_ptr + row * num_steps) - valid = req_state >= 0 - temperature = tl.load(temperature_ptr + req_state, mask=valid, other=0.0) - seed = tl.load(seeds_ptr + req_state, mask=valid, other=0) - previous = 0 - for step in range(walk_steps): - flat = row * num_steps + step - score_base = (flat * top_k + previous) * top_k - scores = tl.load( - scores_ptr + score_base + offsets, - mask=mask & valid, - other=float("-inf"), - ).to(tl.float64 if USE_FP64 else tl.float32) - if SAMPLE_PROBABILISTIC and temperature != 0.0: - # Cache the exact temperature-applied proposal scores expected by - # the shared rejection sampler. This keeps Eagle/MTP's established - # contract unchanged while matching the DFlash2 selector draw. - scores = scores / temperature - candidate_base = flat * top_k - candidates = tl.load( - candidate_ptr + candidate_base + offsets, - mask=mask & valid, - other=0, - ) - - # Candidate token IDs key the noise, matching the target sampler. - position = tl.load(sample_pos_ptr + flat) - 1 - _, index = gumbel_noised_argmax( - scores, - candidates, - mask & valid, - seed, - position, - temperature if SAMPLE_PROBABILISTIC else 0.0, - USE_FP64=USE_FP64, - APPLY_TEMPERATURE=False, - ) - - tl.store( - realized_scores_ptr + candidate_base + offsets, - scores, - mask=mask & valid, - ) - token = tl.load(candidate_ptr + candidate_base + index, mask=valid, other=0) - tl.store(tokens_ptr + flat, token, mask=valid) - previous = index - - if walk_steps < num_steps: - tl.store(path_state_ptr + row, previous, mask=valid) - - -@triton.jit -def _selector_walk_tail_kernel( - scores_ptr, - candidate_ptr, - sample_pos_ptr, - req_state_ptr, - temperature_ptr, - seeds_ptr, - tokens_ptr, - realized_scores_ptr, - path_state_ptr, - num_steps: tl.constexpr, - top_k: tl.constexpr, - BLOCK_K: tl.constexpr, - SAMPLE_PROBABILISTIC: tl.constexpr, - USE_FP64: tl.constexpr, -): - """Write the final dependent slot separately on SM70. - - Triton can drop the seventh store of a fully unrolled selector walk during - CUDA Graph replay on Volta. The first six slots remain fused; this tiny tail - consumes their persistent path state and guarantees the seventh write. - """ - row = tl.program_id(0) - offsets = tl.arange(0, BLOCK_K) - mask = offsets < top_k - req_state = tl.load(req_state_ptr + row * num_steps) - valid = req_state >= 0 - temperature = tl.load(temperature_ptr + req_state, mask=valid, other=0.0) - seed = tl.load(seeds_ptr + req_state, mask=valid, other=0) - previous = tl.load(path_state_ptr + row, mask=valid, other=0) - step: tl.constexpr = num_steps - 1 - flat = row * num_steps + step - score_base = (flat * top_k + previous) * top_k - scores = tl.load( - scores_ptr + score_base + offsets, - mask=mask & valid, - other=float("-inf"), - ).to(tl.float64 if USE_FP64 else tl.float32) - if SAMPLE_PROBABILISTIC and temperature != 0.0: - scores = scores / temperature - candidate_base = flat * top_k - candidates = tl.load( - candidate_ptr + candidate_base + offsets, - mask=mask & valid, - other=0, - ) - position = tl.load(sample_pos_ptr + flat) - 1 - _, index = gumbel_noised_argmax( - scores, - candidates, - mask & valid, - seed, - position, - temperature if SAMPLE_PROBABILISTIC else 0.0, - USE_FP64=USE_FP64, - APPLY_TEMPERATURE=False, - ) - tl.store( - realized_scores_ptr + candidate_base + offsets, - scores, - mask=mask & valid, - ) - token = tl.load(candidate_ptr + candidate_base + index, mask=valid, other=0) - tl.store(tokens_ptr + flat, token, mask=valid) - - -@triton.jit -def _cache_draft_logits_kernel( - draft_logits_ptr, - cached_candidate_ptr, - cached_score_ptr, - candidate_ptr, - scores_ptr, - req_state_ptr, - 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_state = tl.load(req_state_ptr + flat) - step = flat % num_steps - 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 - old_token_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_token_ids, -float("inf"), mask=mask) - token_ids = tl.load(candidate_ptr + candidate_base + offsets, mask=mask) - scores = tl.load(scores_ptr + candidate_base + offsets, mask=mask) - tl.store(logits_base + token_ids, scores, mask=mask) - tl.store(cached_candidate_ptr + cache_base + offsets, token_ids, mask=mask) - if CACHE_SCORES: - tl.store(cached_score_ptr + cache_base + offsets, scores, mask=mask) - - -@triton.jit -def _apply_ngram_draft_kernel( - ngram_tokens_ptr, - ngram_lengths_ptr, - sample_req_state_ptr, - draft_tokens_ptr, - draft_tokens_stride, - cached_candidate_ptr, - cached_score_ptr, - cache_stride_0, - cache_stride_1, - draft_logits_ptr, - draft_logits_stride_0, - draft_logits_stride_1, - num_steps: tl.constexpr, - top_k: tl.constexpr, - BLOCK_K: tl.constexpr, - CACHE_DRAFT_LOGITS: tl.constexpr, - CACHE_SCORES: tl.constexpr, -): - flat = tl.program_id(0) - batch_idx = flat // num_steps - step = flat % num_steps - req_state = tl.load(sample_req_state_ptr + flat) - valid = (req_state >= 0) & (tl.load(ngram_lengths_ptr + batch_idx) == num_steps) - token = tl.load(ngram_tokens_ptr + flat, mask=valid, other=0).to(tl.int64) - tl.store( - draft_tokens_ptr + batch_idx * draft_tokens_stride + step, - token, - mask=valid, - ) - - if CACHE_DRAFT_LOGITS: - offsets = tl.arange(0, BLOCK_K) - topk_mask = valid & (offsets < top_k) - cache_base = ( - cached_candidate_ptr + req_state * cache_stride_0 + step * cache_stride_1 - ) - old_ids = tl.load(cache_base + offsets, mask=topk_mask, other=0) - 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=topk_mask) - - is_proposal = offsets == 0 - new_ids = tl.where(is_proposal, token, 0) - new_scores = tl.where(is_proposal, 0.0, -float("inf")) - tl.store(cache_base + offsets, new_ids, mask=topk_mask) - if CACHE_SCORES: - score_base = ( - cached_score_ptr + req_state * cache_stride_0 + step * cache_stride_1 - ) - tl.store(score_base + offsets, new_scores, mask=topk_mask) - tl.store(logits_base + token, 0.0, mask=valid) - - -class DFlash2Speculator(DFlashSpeculator): - _speculator_name = "DFlash2" - - def __init__(self, vllm_config: VllmConfig, device: torch.device): - super().__init__(vllm_config, device) - draft_config = self.draft_model_config.hf_config.dflash_config - self.selector_top_k = int(draft_config["selector_top_k"]) - self._anchor_indices = ( - torch.arange(self.max_num_reqs, dtype=torch.int64, device=device) - * self.num_query_per_req - ) - self._selector_scores = torch.empty( - self.max_num_reqs, - self.num_speculative_steps, - 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._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._selector_path_state = torch.empty( - self.max_num_reqs, dtype=torch.int32, device=device - ) - self._alignment_candidate_ids: torch.Tensor | None = None - self._alignment_unary_logits: torch.Tensor | None = None - self._alignment_lattice_scores: torch.Tensor | None = None - if ( - envs.VLLM_SPEC_DUMP_ALIGNMENT - and envs.VLLM_SM70_DFLASH2_SPARSE_TARGET_REJECTION - ): - packed_shape = ( - self.max_num_reqs, - self.num_speculative_steps, - self.selector_top_k, - ) - self._alignment_candidate_ids = torch.empty( - packed_shape, dtype=torch.int64, device=device - ) - self._alignment_unary_logits = torch.empty( - packed_shape, dtype=torch.float32, device=device - ) - self._alignment_lattice_scores = torch.empty( - (*packed_shape, self.selector_top_k), - dtype=torch.float32, - device=device, - ) - self._use_sm70_tail = _requires_sm70_tail(device, self.num_speculative_steps) - if self.draft_logits is not None: - # The cache kernel writes only K columns; all other vocabulary - # columns must remain impossible. - self.draft_logits.fill_(-float("inf")) - - self._ngram_assist: DFlash2NgramAssist | None = None - self._ngram_num_hits = 0 - 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 - ): - min_ngram = speculative_config.prompt_lookup_min - max_ngram = speculative_config.prompt_lookup_max - assert min_ngram is not None and max_ngram is not None - self._ngram_assist = DFlash2NgramAssist( - min_ngram=min_ngram, - max_ngram=max_ngram, - num_draft_tokens=self.num_speculative_steps, - max_model_len=self.max_model_len, - ) - self._ngram_tokens_cpu_tensor = torch.zeros( - self.max_num_reqs, - self.num_speculative_steps, - dtype=torch.int64, - device="cpu", - pin_memory=True, - ) - self._ngram_lengths_cpu_tensor = torch.zeros( - self.max_num_reqs, - dtype=torch.int32, - device="cpu", - pin_memory=True, - ) - self._ngram_tokens_cpu = self._ngram_tokens_cpu_tensor.numpy() - self._ngram_lengths_cpu = self._ngram_lengths_cpu_tensor.numpy() - self._ngram_tokens = torch.zeros( - self.max_num_reqs, - self.num_speculative_steps, - dtype=torch.int64, - device=device, - ) - self._ngram_lengths = torch.zeros( - self.max_num_reqs, dtype=torch.int32, device=device - ) - logger.info( - "Enabled DFlash2 ngram assist with prompt lookup [%d, %d] " - "and draft width %d.", - min_ngram, - max_ngram, - self.num_speculative_steps, - ) - - @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 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. - return torch.float32, -float("inf") - - def _sample_path( - self, - candidate_ids: torch.Tensor, - scores: torch.Tensor, - num_reqs: int, - ) -> None: - # The SM70 tail must consume the exact same packed lattice as the - # prefix walk. Keep one persistent view instead of creating temporary - # contiguous inputs only for the first launch. - 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 - ) - _selector_walk_kernel[(num_reqs,)]( - scores, - candidate_ids, - self.sample_pos, - self.sample_idx_mapping, - self.temperature, - self.seeds, - self.draft_tokens, - self._selector_scores, - self._selector_path_state, - num_steps=self.num_speculative_steps, - walk_steps=walk_steps, - top_k=self.selector_top_k, - BLOCK_K=block_k, - SAMPLE_PROBABILISTIC=self.draft_logits is not None, - USE_FP64=self.use_fp64_gumbel, - num_warps=1, - ) - if self._use_sm70_tail: - _selector_walk_tail_kernel[(num_reqs,)]( - scores, - candidate_ids, - self.sample_pos, - self.sample_idx_mapping, - self.temperature, - self.seeds, - self.draft_tokens, - self._selector_scores, - self._selector_path_state, - num_steps=self.num_speculative_steps, - top_k=self.selector_top_k, - BLOCK_K=block_k, - SAMPLE_PROBABILISTIC=self.draft_logits is not None, - USE_FP64=self.use_fp64_gumbel, - num_warps=1, - ) - - def _cache_draft_logits(self, candidate_ids: torch.Tensor, num_sample: int) -> None: - draft_logits = self.draft_logits - assert draft_logits is not None - cached_scores = self._cached_candidate_scores - block_k = triton.next_power_of_2(self.selector_top_k) - _cache_draft_logits_kernel[(num_sample,)]( - draft_logits, - self._cached_candidate_ids, - self._selector_scores if cached_scores is None else cached_scores, - candidate_ids, - self._selector_scores, - self.sample_idx_mapping, - 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 get_sparse_draft_logits( - self, - ) -> tuple[torch.Tensor, torch.Tensor] | None: - """Return request-slot proposal candidates for sparse rejection.""" - if self.draft_logits is None or self._cached_candidate_scores is None: - return None - return self._cached_candidate_ids, self._cached_candidate_scores - - def get_selector_alignment_shadow( - self, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None: - """Return packed selector tensors for an explicitly enabled diagnostic.""" - if self._alignment_candidate_ids is None: - return None - assert self._alignment_unary_logits is not None - assert self._alignment_lattice_scores is not None - return ( - self._alignment_candidate_ids, - self._alignment_unary_logits, - self._alignment_lattice_scores, - ) - - def _prepare_ngram_assist( - self, - input_batch, - output_copy_event: torch.cuda.Event | None, - sampled_token_ids_cpu: np.ndarray | None, - num_sampled_tokens_cpu: np.ndarray | None, - all_token_ids_cpu: np.ndarray | None, - ) -> bool: - assist = self._ngram_assist - self._ngram_num_hits = 0 - if ( - assist is None - or input_batch.has_structured_output_reqs - or output_copy_event is None - or sampled_token_ids_cpu is None - or num_sampled_tokens_cpu is None - or all_token_ids_cpu is None - ): - return False - - # The copy stream only depends on target sampling. The main stream can - # materialize DFlash context K/V while the host waits here, so lookup - # does not serialize the context projection. - output_copy_event.synchronize() - num_reqs = input_batch.num_reqs - num_draft_tokens = input_batch.num_draft_tokens_per_req - if num_draft_tokens is None: - num_draft_tokens = np.zeros(num_reqs, dtype=np.int32) - prior_lengths = ( - input_batch.seq_lens_cpu_upper_bound[:num_reqs].numpy() - num_draft_tokens - ) - eligible = (~input_batch.is_prefilling_np[:num_reqs]) & ( - num_sampled_tokens_cpu[:num_reqs] > 0 - ) - full_hits = assist.propose( - all_token_ids_cpu, - input_batch.idx_mapping_np, - prior_lengths, - sampled_token_ids_cpu, - num_sampled_tokens_cpu, - eligible, - self._ngram_tokens_cpu, - self._ngram_lengths_cpu, - ) - self._ngram_rounds += 1 - skip_query = num_reqs > 0 and full_hits == num_reqs - self._ngram_num_hits = full_hits if skip_query else 0 - if skip_query: - self._ngram_tokens[:num_reqs].copy_( - self._ngram_tokens_cpu_tensor[:num_reqs], non_blocking=True - ) - self._ngram_lengths[:num_reqs].copy_( - self._ngram_lengths_cpu_tensor[:num_reqs], non_blocking=True - ) - - self._ngram_skipped_rounds += int(skip_query) - if ( - envs.VLLM_DFLASH_PROFILE - and self._ngram_rounds % envs.VLLM_DFLASH_PROFILE_LOG_INTERVAL == 0 - ): - eligible_count = max(assist.num_eligible, 1) - logger.info( - "DFLASH2_NGRAM_PROFILE rounds=%d eligible=%d full_hits=%d " - "hit_rate=%.4f skipped_query_rounds=%d lookup_avg_ms=%.4f", - self._ngram_rounds, - assist.num_eligible, - assist.num_full_hits, - assist.num_full_hits / eligible_count, - self._ngram_skipped_rounds, - assist.lookup_seconds * 1000.0 / self._ngram_rounds, - ) - return skip_query - - def _apply_ngram_assist(self, num_reqs: int) -> None: - if self._ngram_assist is None or self._ngram_num_hits == 0: - return - draft_logits = self.draft_logits - cached_scores = self._cached_candidate_scores - block_k = triton.next_power_of_2(self.selector_top_k) - _apply_ngram_draft_kernel[(num_reqs * self.num_speculative_steps,)]( - self._ngram_tokens, - self._ngram_lengths, - self.sample_idx_mapping, - self.draft_tokens, - self.draft_tokens.stride(0), - self._cached_candidate_ids, - self._selector_scores if cached_scores is None else cached_scores, - self._cached_candidate_ids.stride(0), - self._cached_candidate_ids.stride(1), - self._selector_scores if draft_logits is None else draft_logits, - 0 if draft_logits is None else draft_logits.stride(0), - 0 if draft_logits is None else draft_logits.stride(1), - num_steps=self.num_speculative_steps, - top_k=self.selector_top_k, - BLOCK_K=block_k, - CACHE_DRAFT_LOGITS=draft_logits is not None, - CACHE_SCORES=cached_scores is not None, - num_warps=1, - ) - - def _generate_draft( - self, - num_reqs: int, - num_tokens_padded: int, - attn_metadata: dict[str, Any] | None, - slot_mappings: dict[str, torch.Tensor] | None, - num_tokens_across_dp: torch.Tensor | None, - cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE, - ) -> None: - last_hidden_states = self._run_model( - num_tokens_padded, - attn_metadata, - slot_mappings, - num_tokens_across_dp, - cudagraph_runtime_mode, - ) - num_sample = num_reqs * self.num_speculative_steps - hidden_states = last_hidden_states[self.sample_indices[:num_sample]].view( - num_reqs, self.num_speculative_steps, -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 - ) - unary_logits = unary_logits.view_as(candidate_ids) - anchor_token_ids = self.input_buffers.input_ids[self._anchor_indices[:num_reqs]] - scores = self.model.model.candidate_selector( - candidate_ids, - unary_logits, - hidden_states, - anchor_token_ids, - ) - if self._alignment_candidate_ids is not None: - assert self._alignment_unary_logits is not None - assert self._alignment_lattice_scores is not None - self._alignment_candidate_ids[:num_reqs].copy_(candidate_ids) - 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) - if self.draft_logits is not None: - self._cache_draft_logits(candidate_ids, num_sample) diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py b/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py index 8c8d56ba31..d805c88582 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py @@ -38,17 +38,9 @@ def get_eagle3_aux_layers_from_config( if not (spec_config and spec_config.draft_model_config): return None hf_config = spec_config.draft_model_config.hf_config - layer_ids = getattr(hf_config, "eagle_aux_hidden_state_layer_ids", None) - if not layer_ids: - dflash_config = getattr(hf_config, "dflash_config", None) - if dflash_config and isinstance(dflash_config, dict): - # DFlash checkpoint indices identify decoder layers, while target - # hidden-state capture uses layer boundaries. - layer_ids = [i + 1 for i in (dflash_config.get("target_layer_ids") or [])] - if not layer_ids: - target_layer_ids = getattr(hf_config, "target_layer_ids", None) - if target_layer_ids: - layer_ids = [i + 1 for i in target_layer_ids] + if not hasattr(hf_config, "eagle_aux_hidden_state_layer_ids"): + return None + layer_ids = hf_config.eagle_aux_hidden_state_layer_ids if layer_ids and isinstance(layer_ids, (list, tuple)): return tuple(layer_ids) return None diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py index 8376a51558..6ae3fe793b 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import Any -import numpy as np import torch import torch.nn as nn @@ -482,10 +481,6 @@ def propose( skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, is_profile: bool = False, - output_copy_event: torch.cuda.Event | None = None, - sampled_token_ids_cpu: np.ndarray | None = None, - num_sampled_tokens_cpu: np.ndarray | None = None, - all_token_ids_cpu: np.ndarray | None = None, ) -> torch.Tensor: num_tokens = input_batch.num_tokens_after_padding num_reqs = input_batch.num_reqs diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py index 6826a234f0..fcbfc5569e 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py @@ -24,13 +24,6 @@ def _should_share(eagle: nn.Module, flag: str, draft, target) -> bool: return torch.equal(w, target.weight) -def get_target_lm_head(target_model: nn.Module, target_language_model: nn.Module): - """Return the language-model head for plain and conditional targets.""" - return getattr(target_language_model, "lm_head", None) or getattr( - target_model, "lm_head", None - ) - - def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module: from vllm.compilation.backends import set_model_tag diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py index da9cca92d3..0cfbdf4182 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py @@ -2,12 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch -from vllm.triton_utils import tl, tldevice, triton -from vllm.v1.worker.gpu.sample.gumbel import ( - gumbel_block_argmax, - gumbel_noised_argmax, - tl_rand32, -) +from vllm.triton_utils import tl, triton +from vllm.v1.worker.gpu.sample.gumbel import gumbel_block_argmax, tl_rand64 @triton.jit @@ -247,7 +243,7 @@ def _rejection_kernel( if SYNTHETIC_MODE: pos = tl.load(pos_ptr + logit_idx) - u = tl_rand32(seed, pos, includes_zero=False) + u = tl_rand64(seed, pos, includes_zero=False) rate = tl.load(synthetic_conditional_rates_ptr + i) accepted &= u < rate else: @@ -271,7 +267,7 @@ def _rejection_kernel( ) target_log_prob = target_logit - target_lse pos = tl.load(pos_ptr + logit_idx) - u = tl_rand32(seed, pos, includes_zero=False) + u = tl_rand64(seed, pos, includes_zero=False) if HAS_DRAFT_LOGITS: draft_logit = tl.load( draft_logits_ptr @@ -392,7 +388,7 @@ def _resample_kernel( ratio = tl.exp(draft_log_probs - target_log_probs) residual_logits = tl.where( ratio < 1.0, - target_log_probs + tldevice.log1p(-ratio), + target_log_probs + tl.log(1 - ratio), float("-inf"), ).to(tl.float32) else: @@ -417,7 +413,6 @@ def _resample_kernel( pos_ptr, None, # processed_logits_ptr 0, # processed_logits_stride - 0, # processed_logits_col_stride None, # processed_logits_col_ptr vocab_size, APPLY_TEMPERATURE=False, @@ -495,327 +490,6 @@ def _insert_resampled_kernel( ) -@triton.jit -def _dflash2_compact_target_row( - target_ids_ptr, - target_logits_ptr, - target_stride, - row, - temperature, - top_p, - TOP_K: tl.constexpr, - BLOCK_K: tl.constexpr, -): - offsets = tl.arange(0, BLOCK_K) - mask = offsets < TOP_K - token_ids = tl.load( - target_ids_ptr + row * target_stride + offsets, - mask=mask, - other=0, - ).to(tl.int64) - logits = tl.load( - target_logits_ptr + row * target_stride + offsets, - mask=mask, - other=float("-inf"), - ).to(tl.float32) - if temperature != 1.0: - logits = logits / temperature - - # get_topk_tokens_and_logits returns scores in descending order. Apply - # nucleus filtering only within that exact top-k support, matching the - # target's top-k-first sampling contract without materializing V columns. - row_max = tl.max(logits, axis=0) - unnormalized = tl.where(mask, tl.exp(logits - row_max), 0.0) - normalizer = tl.sum(unnormalized, axis=0) - probs = unnormalized / normalizer - cumulative_before = tl.cumsum(probs, axis=0) - probs - keep = mask & ((top_p >= 1.0) | (cumulative_before < top_p)) - logits = tl.where(keep, logits, float("-inf")) - - kept_max = tl.max(logits, axis=0) - kept_sumexp = tl.sum(tl.where(keep, tl.exp(logits - kept_max), 0.0), axis=0) - logsumexp = kept_max + tl.log(kept_sumexp) - return token_ids, logits, keep, logsumexp - - -@triton.jit -def _dflash2_sparse_topk_rejection_kernel( - # [num_reqs, num_speculative_steps + 1] - sampled_ptr, - sampled_stride, - # [num_reqs] - num_sampled_ptr, - # [num_logits, target_top_k] - target_ids_ptr, - target_logits_ptr, - target_stride, - # [max_num_reqs, num_speculative_steps, draft_top_k] - draft_ids_ptr, - draft_logits_ptr, - draft_stride_0, - draft_stride_1, - # [num_logits] - draft_sampled_ptr, - # [num_reqs + 1] - cu_num_logits_ptr, - # [num_reqs] - idx_mapping_ptr, - # [max_num_reqs] - temperature_ptr, - top_p_ptr, - seed_ptr, - # [num_logits] - pos_ptr, - num_speculative_steps: tl.constexpr, - TARGET_TOP_K: tl.constexpr, - TARGET_BLOCK_K: tl.constexpr, - DRAFT_TOP_K: tl.constexpr, - DRAFT_BLOCK_K: tl.constexpr, - USE_FP64: tl.constexpr, -): - """DFlash2 chain rejection on compact target/draft supports. - - The target distribution is zero outside its top-k/top-p support and the - DFlash2 proposal is zero outside selector top-k. Therefore both the - acceptance ratio and ``relu(p - q)`` recovery distribution are exact on - these compact sets; scanning or storing a full-vocabulary row is unnecessary. - """ - req_idx = tl.program_id(0) - req_state_idx = tl.load(idx_mapping_ptr + req_idx) - start_idx = tl.load(cu_num_logits_ptr + req_idx) - end_idx = tl.load(cu_num_logits_ptr + req_idx + 1) - num_tokens = end_idx - start_idx - temperature = tl.load(temperature_ptr + req_state_idx).to(tl.float32) - top_p = tl.load(top_p_ptr + req_state_idx).to(tl.float32) - seed = tl.load(seed_ptr + req_state_idx) - - target_offsets = tl.arange(0, TARGET_BLOCK_K) - target_mask = target_offsets < TARGET_TOP_K - draft_offsets = tl.arange(0, DRAFT_BLOCK_K) - draft_mask = draft_offsets < DRAFT_TOP_K - - accepted = True - accepted_steps = tl.zeros((), dtype=tl.int32) - for step in range(num_speculative_steps): - if accepted and step < num_tokens - 1: - row = start_idx + step - target_ids, target_logits, target_keep, target_lse = ( - _dflash2_compact_target_row( - target_ids_ptr, - target_logits_ptr, - target_stride, - row, - temperature, - top_p, - TOP_K=TARGET_TOP_K, - BLOCK_K=TARGET_BLOCK_K, - ) - ) - proposed = tl.load(draft_sampled_ptr + row + 1).to(tl.int64) - target_proposed_logit = tl.max( - tl.where( - target_keep & (target_ids == proposed), - target_logits, - float("-inf"), - ), - axis=0, - ) - - draft_base = ( - draft_logits_ptr - + req_state_idx * draft_stride_0 - + step * draft_stride_1 - ) - draft_id_base = ( - draft_ids_ptr + req_state_idx * draft_stride_0 + step * draft_stride_1 - ) - draft_ids = tl.load( - draft_id_base + draft_offsets, - mask=draft_mask, - other=0, - ).to(tl.int64) - draft_logits = tl.load( - draft_base + draft_offsets, - mask=draft_mask, - other=float("-inf"), - ).to(tl.float32) - draft_max, draft_sumexp = _compute_block_max_and_sumexp(draft_logits) - draft_lse = draft_max + tl.log(draft_sumexp) - draft_proposed_logit = tl.max( - tl.where( - draft_mask & (draft_ids == proposed), - draft_logits, - float("-inf"), - ), - axis=0, - ) - - pos = tl.load(pos_ptr + row) - uniform = tl_rand32(seed, pos, includes_zero=False) - accepted &= (target_proposed_logit - target_lse) > ( - tl.log(uniform) + draft_proposed_logit - draft_lse - ) - tl.store( - sampled_ptr + req_idx * sampled_stride + step, - proposed, - ) - accepted_steps += accepted.to(tl.int32) - - resample_step = accepted_steps - resample_row = start_idx + resample_step - target_ids, target_logits, target_keep, target_lse = _dflash2_compact_target_row( - target_ids_ptr, - target_logits_ptr, - target_stride, - resample_row, - temperature, - top_p, - TOP_K=TARGET_TOP_K, - BLOCK_K=TARGET_BLOCK_K, - ) - is_bonus = resample_row == end_idx - 1 - residual_logits = target_logits - if not is_bonus: - draft_base = ( - draft_logits_ptr - + req_state_idx * draft_stride_0 - + resample_step * draft_stride_1 - ) - draft_id_base = ( - draft_ids_ptr - + req_state_idx * draft_stride_0 - + resample_step * draft_stride_1 - ) - draft_ids = tl.load( - draft_id_base + draft_offsets, - mask=draft_mask, - other=0, - ).to(tl.int64) - draft_logits = tl.load( - draft_base + draft_offsets, - mask=draft_mask, - other=float("-inf"), - ).to(tl.float32) - draft_max, draft_sumexp = _compute_block_max_and_sumexp(draft_logits) - draft_lse = draft_max + tl.log(draft_sumexp) - - id_match = target_ids[:, None] == draft_ids[None, :] - draft_logits_at_target = tl.max( - tl.where( - id_match & target_mask[:, None] & draft_mask[None, :], - draft_logits[None, :], - float("-inf"), - ), - axis=1, - ) - target_log_probs = target_logits - target_lse - draft_log_probs = draft_logits_at_target - draft_lse - ratio = tl.exp(draft_log_probs - target_log_probs) - residual_logits = tl.where( - target_keep & (ratio < 1.0), - target_log_probs + tldevice.log1p(-ratio), - float("-inf"), - ) - - pos = tl.load(pos_ptr + resample_row) - _, selected_offset = gumbel_noised_argmax( - residual_logits, - target_ids, - target_keep, - seed, - pos, - temperature, - USE_FP64=USE_FP64, - APPLY_TEMPERATURE=False, - ) - selected_token = tl.load( - target_ids_ptr + resample_row * target_stride + selected_offset - ) - tl.store( - sampled_ptr + req_idx * sampled_stride + resample_step, - selected_token, - ) - tl.store(num_sampled_ptr + req_idx, resample_step + 1) - - -def dflash2_sparse_topk_rejection_sample( - target_topk_ids: torch.Tensor, - target_topk_logits: torch.Tensor, - draft_topk_ids: torch.Tensor, - draft_topk_logits: torch.Tensor, - draft_sampled: torch.Tensor, - cu_num_logits: torch.Tensor, - pos: torch.Tensor, - idx_mapping: torch.Tensor, - temperature: torch.Tensor, - top_p: torch.Tensor, - seed: torch.Tensor, - num_speculative_steps: int, - use_fp64: bool = False, -) -> tuple[torch.Tensor, torch.Tensor]: - """Run exact DFlash2 rejection on compact top-k distributions. - - Callers must enforce the no-penalty/no-logit-bias contract and provide - target candidates after a global top-k merge. Draft scores must be the - temperature-applied scores used to draw the corresponding proposal. - """ - if target_topk_ids.ndim != 2 or target_topk_logits.ndim != 2: - raise ValueError("target top-k tensors must be two-dimensional") - if target_topk_ids.shape != target_topk_logits.shape: - raise ValueError("target top-k ids/logits shapes must match") - if draft_topk_ids.ndim != 3 or draft_topk_logits.ndim != 3: - raise ValueError("draft top-k tensors must be three-dimensional") - if draft_topk_ids.shape != draft_topk_logits.shape: - raise ValueError("draft top-k ids/logits shapes must match") - if draft_topk_ids.shape[1] != num_speculative_steps: - raise ValueError("draft top-k step count does not match configuration") - if target_topk_ids.shape[0] != draft_sampled.shape[0]: - raise ValueError("target rows and sampled-token rows must match") - if not 0 < target_topk_ids.shape[1] <= 64: - raise ValueError("target top-k width must be in [1, 64]") - if not 0 < draft_topk_ids.shape[2] <= 64: - raise ValueError("draft top-k width must be in [1, 64]") - - target_topk_ids = target_topk_ids.contiguous() - target_topk_logits = target_topk_logits.contiguous() - draft_topk_ids = draft_topk_ids.contiguous() - draft_topk_logits = draft_topk_logits.contiguous() - num_reqs = cu_num_logits.shape[0] - 1 - sampled = draft_sampled.new_empty( - num_reqs, - num_speculative_steps + 1, - dtype=torch.int64, - ) - num_sampled = sampled.new_empty(num_reqs, dtype=torch.int32) - _dflash2_sparse_topk_rejection_kernel[(num_reqs,)]( - sampled, - sampled.stride(0), - num_sampled, - target_topk_ids, - target_topk_logits, - target_topk_ids.stride(0), - draft_topk_ids, - draft_topk_logits, - draft_topk_ids.stride(0), - draft_topk_ids.stride(1), - draft_sampled, - cu_num_logits, - idx_mapping, - temperature, - top_p, - seed, - pos, - num_speculative_steps=num_speculative_steps, - TARGET_TOP_K=target_topk_ids.shape[1], - TARGET_BLOCK_K=triton.next_power_of_2(target_topk_ids.shape[1]), - DRAFT_TOP_K=draft_topk_ids.shape[2], - DRAFT_BLOCK_K=triton.next_power_of_2(draft_topk_ids.shape[2]), - USE_FP64=use_fp64, - num_warps=1, - ) - return sampled, num_sampled - - def rejection_sample( # [num_logits, V] target_logits: torch.Tensor, diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py deleted file mode 100644 index 0775cfdca6..0000000000 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ /dev/null @@ -1,371 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from abc import ABC, abstractmethod -from collections.abc import Mapping -from typing import Any - -import numpy as np -import torch -import torch.nn as nn - -from vllm.config import VllmConfig, get_layers_from_vllm_config -from vllm.config.compilation import CUDAGraphMode -from vllm.logger import init_logger -from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase -from vllm.model_executor.models import supports_multimodal -from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.v1.kv_cache_interface import KVCacheConfig -from vllm.v1.worker.gpu.attn_utils import ( - build_attn_metadata, - init_attn_backend, -) -from vllm.v1.worker.gpu.block_table import BlockTables -from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers -from vllm.v1.worker.gpu.model_states.interface import ModelState -from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample -from vllm.v1.worker.utils import AttentionGroup - -logger = init_logger(__name__) - - -class BaseSpeculator(ABC): - @abstractmethod - def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: - pass - - @abstractmethod - def capture(self) -> None: - pass - - @abstractmethod - def propose( - self, - input_batch: InputBatch, - attn_metadata: dict[str, Any], - slot_mappings: dict[str, torch.Tensor], - # [num_tokens, hidden_size] - last_hidden_states: torch.Tensor, - # num_layers x [num_tokens, hidden_size] - aux_hidden_states: list[torch.Tensor] | None, - # [num_reqs] - num_sampled: torch.Tensor, - # [num_reqs] - num_rejected: torch.Tensor, - # [max_num_reqs] - last_sampled: torch.Tensor, - # [num_prefill_lookahead, max_num_reqs] - next_prefill_tokens: torch.Tensor, - # [max_num_reqs] - temperature: torch.Tensor, - # [max_num_reqs] - seeds: torch.Tensor, - num_tokens_across_dp: torch.Tensor | None = None, - dummy_run: bool = False, - skip_attn_for_dummy_run: bool = False, - mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, - is_profile: bool = False, - output_copy_event: torch.cuda.Event | None = None, - sampled_token_ids_cpu: np.ndarray | None = None, - num_sampled_tokens_cpu: np.ndarray | None = None, - all_token_ids_cpu: np.ndarray | None = None, - ) -> torch.Tensor: - pass - - -class DraftModelSpeculator(BaseSpeculator): - def __init__(self, vllm_config: VllmConfig, device: torch.device): - self.vllm_config = vllm_config - self.device = device - - assert vllm_config.speculative_config is not None - self.speculative_config = vllm_config.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 - - self.scheduler_config = vllm_config.scheduler_config - self.max_num_reqs = self.scheduler_config.max_num_seqs - self.max_num_tokens = self.scheduler_config.max_num_batched_tokens - self.max_model_len = vllm_config.model_config.max_model_len - self.draft_max_seq_len = self.max_model_len - # 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() - # Widen for HC-multiplexed residuals (e.g. DeepSeek V4 feeds the MTP - # draft the target's pre-hc_head (T, hc_mult * hidden_size) residual). - # Non-HC models default to hc_mult=1 and are unaffected. - hc_mult = getattr(self.draft_model_config.hf_config, "hc_mult", 1) - self.hidden_size = self.hidden_size * hc_mult - 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 - ) - - # DP configuration - self.dp_size = vllm_config.parallel_config.data_parallel_size - self.dp_rank = vllm_config.parallel_config.data_parallel_rank - - self.input_buffers = InputBuffers( - max_num_reqs=self.max_num_reqs, - max_num_tokens=self.max_num_tokens, - device=device, - ) - self.idx_mapping = torch.zeros( - self.max_num_reqs, dtype=torch.int32, device=device - ) - self.temperature = torch.zeros( - self.max_num_reqs, dtype=torch.float32, device=device - ) - self.seeds = torch.zeros(self.max_num_reqs, dtype=torch.int64, device=device) - self.draft_tokens = torch.zeros( - self.max_num_reqs, - self.num_speculative_steps, - dtype=torch.int64, - device=device, - ) - self.arange = torch.arange( - self.max_num_reqs + 1, dtype=torch.int32, device="cpu" - ) - - self.draft_logits: torch.Tensor | None = None - if self.speculative_config.draft_sample_method == "probabilistic": - # Temperature-applied logits, cached from the previous decode step. - dtype, fill = self.draft_logits_spec(vllm_config) - self.draft_logits = torch.full( - ( - self.max_num_reqs, - self.num_speculative_steps, - self.vocab_size, - ), - fill, - dtype=dtype, - device=device, - ) - - self.supports_mm_inputs = False - - @abstractmethod - def load_draft_model( - self, - target_model: nn.Module, - target_attn_layer_names: set[str], - ) -> nn.Module: - pass - - def load_model(self, target_model: nn.Module) -> None: - target_attn_layer_names = set( - get_layers_from_vllm_config( - self.vllm_config, - AttentionLayerBase, # type: ignore[type-abstract] - ).keys() - ) - - self.model = self.load_draft_model(target_model, target_attn_layer_names) - self._validate_local_argmax_reduction() - - all_attn_layers = set[str]( - get_layers_from_vllm_config( - self.vllm_config, - AttentionLayerBase, # type: ignore[type-abstract] - ).keys() - ) - self.draft_attn_layer_names = all_attn_layers - target_attn_layer_names - - target_supports_mm = MULTIMODAL_REGISTRY.supports_multimodal_inputs( - self.vllm_config.model_config - ) - draft_supports_mm = supports_multimodal(self.model) - self.supports_mm_inputs = target_supports_mm and draft_supports_mm - if target_supports_mm and not draft_supports_mm: - logger.warning_once( - "Draft model %s does not support external multimodal embeddings. " - "Embeddings from the target model will not be passed to the " - "drafter; using text-only draft inputs instead.", - type(self.model).__name__, - ) - - @property - def attn_vllm_config(self) -> VllmConfig: - """Config for the draft's attention metadata builders. Overridden by - speculators whose attention mode differs from the target's.""" - return self.vllm_config - - def set_attn( - self, - model_state: ModelState, - kv_cache_config: KVCacheConfig, - block_tables: BlockTables, - target_input_buffers: InputBuffers, - target_attn_groups: list[list[AttentionGroup]], - ) -> None: - self.model_state = model_state - self.kv_cache_config = kv_cache_config - self.attn_groups, self.attn_cg_support, _ = init_attn_backend( - kv_cache_config, - self.attn_vllm_config, - self.device, - active_layer_names=self.draft_attn_layer_names, - ) - self.block_tables = block_tables - # The target model runner's buffers and attention groups. Draft - # prefill reuses the target model's attention metadata, so its - # cudagraph capture must build dummy metadata through the same - # builders and buffers. - self.target_input_buffers = target_input_buffers - self.target_attn_groups = target_attn_groups - - def _build_draft_attn_metadata( - self, - num_reqs: int, - num_reqs_padded: int, - num_tokens_padded: int, - seq_lens_cpu_upper_bound: torch.Tensor, - step: int, - num_query_per_req: int = 1, - causal: bool | Mapping[int, bool] = True, - query_start_loc_np: np.ndarray | None = None, - dcp_local_seq_lens: torch.Tensor | None = None, - ) -> dict[str, Any] | None: - if query_start_loc_np is not None: - # Non-uniform query layout (e.g. multi-module MTP's mixed - # prefill/decode queries); num_query_per_req is ignored. - query_start_loc_cpu = torch.empty(num_reqs_padded + 1, dtype=torch.int32) - query_start_loc_cpu[: num_reqs + 1] = torch.from_numpy( - query_start_loc_np[: num_reqs + 1] - ) - query_start_loc_cpu[num_reqs:] = query_start_loc_cpu[num_reqs] - max_query_len = int( - (query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]).max() - ) - else: - # Uniform query: query_start_loc[i] = min(i, num_reqs) * num_query_per_req. - # Clamp keeps the series non-decreasing past num_reqs, which some - # attention backends require. - query_start_loc_cpu = ( - torch.clamp(self.arange[: num_reqs_padded + 1], max=num_reqs) - * num_query_per_req - ) - max_query_len = num_query_per_req - block_tables = [ - x[:num_reqs_padded] for x in self.block_tables.input_block_tables - ] - slot_mappings = self.block_tables.slot_mappings[:, :num_tokens_padded] - draft_seq_lens_cpu_upper_bound = torch.zeros( - num_reqs_padded, dtype=torch.int32, device="cpu" - ) - torch.add( - seq_lens_cpu_upper_bound[:num_reqs], - step, - out=draft_seq_lens_cpu_upper_bound[:num_reqs], - ) - draft_seq_lens_cpu_upper_bound[:num_reqs].clamp_(max=self.max_model_len) - attn_metadata = build_attn_metadata( - attn_groups=self.attn_groups, - num_reqs=num_reqs_padded, - num_tokens=num_tokens_padded, - query_start_loc_gpu=self.input_buffers.query_start_loc[ - : num_reqs_padded + 1 - ], - query_start_loc_cpu=query_start_loc_cpu, - max_query_len=max_query_len, - seq_lens=self.input_buffers.seq_lens[:num_reqs_padded], - dcp_local_seq_lens=( - None - if dcp_local_seq_lens is None - else dcp_local_seq_lens[:num_reqs_padded] - ), - max_seq_len=self.draft_max_seq_len, - block_tables=block_tables, - slot_mappings=slot_mappings, - kv_cache_config=self.kv_cache_config, - causal=causal, - seq_lens_cpu_upper_bound=draft_seq_lens_cpu_upper_bound, - ) - return attn_metadata - - def draft_logits_spec(self, vllm_config: VllmConfig) -> tuple[torch.dtype, float]: - """Dtype and initial value for the cached proposal distribution. - - Most speculators overwrite every vocabulary column and use the target - head dtype. DFlash2 overrides this because it writes only K candidates - and its selector/rejection path must share FP32 scores exactly. - """ - return vllm_config.model_config.head_dtype, 0.0 - - 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 _greedy_sample_draft(self, hidden_states: torch.Tensor) -> torch.Tensor: - 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 sample_draft( - self, - hidden_states: torch.Tensor, - positions: torch.Tensor, - idx_mapping: torch.Tensor, - temperature: torch.Tensor, - seeds: 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( - logits, - idx_mapping, - temperature, - seeds, - positions + 1, - apply_temperature=True, - output_processed_logits=draft_logits, - output_processed_logits_col=draft_step, - use_fp64=self.use_fp64_gumbel, - ) - return self._greedy_sample_draft(hidden_states) - - def _copy_request_inputs( - self, - num_reqs: int, - # [num_reqs] - idx_mapping: torch.Tensor, - # [max_num_reqs] - temperature: torch.Tensor, - # [max_num_reqs] - seeds: torch.Tensor, - ) -> None: - # Copy temperature, seeds, and idx mapping to the pre-allocated buffers. - # NOTE(woosuk): For draft sampling, we only consider the temperature - # and ignore the other sampling parameters such as top_k and top_p, - # for simplicity and performance. - # While this may slightly degrade the acceptance rate, it does not - # affect the output distribution after rejection sampling. - self.temperature.copy_(temperature) - self.seeds.copy_(seeds) - self.idx_mapping[:num_reqs].copy_(idx_mapping) - # idx_mapping for CG padded requests points to -1, which is ignored - # during sampling to prevent writing stale values to draft logits. - self.idx_mapping[num_reqs:].fill_(-1) diff --git a/vllm/v1/worker/gpu/spec_decode/utils.py b/vllm/v1/worker/gpu/spec_decode/utils.py index e721e85b35..e1fa21aeb8 100644 --- a/vllm/v1/worker/gpu/spec_decode/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/utils.py @@ -45,18 +45,3 @@ def get_draft_tokens(self) -> DraftTokenIds | None: # This case only happens when async scheduling is disabled. draft_token_ids = [[-1] * self.num_draft_tokens for _ in self.req_ids] return DraftTokenIds(self.req_ids, draft_token_ids) - - -def get_parallel_drafting_token_id(hf_config) -> int: - """Resolve the mask token used by a parallel-drafting checkpoint.""" - dflash_config = getattr(hf_config, "dflash_config", None) or {} - if "mask_token_id" in dflash_config: - return int(dflash_config["mask_token_id"]) - if hasattr(hf_config, "pard_token"): - return int(hf_config.pard_token) - if hasattr(hf_config, "ptd_token_id"): - return int(hf_config.ptd_token_id) - raise ValueError( - "Model config must specify dflash_config.mask_token_id, pard_token, " - "or ptd_token_id for parallel drafting." - ) diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index 1cdcf567d6..83d87c74a4 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -15,46 +15,10 @@ NewRequestData, SchedulerOutput, ) -from vllm.v1.kv_cache_interface import CrossAttentionSpec, KVCacheSpec, MambaSpec from vllm.v1.request import Request from vllm.v1.worker.gpu.model_runner import GPUModelRunner -def _reserved_block_count( - num_tokens: int, - kv_cache_spec: KVCacheSpec, - *, - num_lookahead_tokens: int, - max_model_len: int, - max_encoder_len: int, -) -> int: - """Match the scheduler's block reservation in hand-built warmup batches.""" - if isinstance(kv_cache_spec, CrossAttentionSpec): - return cdiv(max_encoder_len, kv_cache_spec.block_size) - num_speculative_blocks = 0 - if isinstance(kv_cache_spec, MambaSpec): - num_speculative_blocks = kv_cache_spec.num_speculative_blocks - if kv_cache_spec.mamba_cache_mode == "align": - return cdiv(num_tokens, kv_cache_spec.block_size) + num_speculative_blocks - num_tokens = min(num_tokens + num_lookahead_tokens, max_model_len) - return cdiv(num_tokens, kv_cache_spec.block_size) + num_speculative_blocks - - -def _warmup_block_counter( - model_runner: GPUModelRunner, -) -> Callable[[int, KVCacheSpec], int]: - def block_count(num_tokens: int, kv_cache_spec: KVCacheSpec) -> int: - return _reserved_block_count( - num_tokens, - kv_cache_spec, - num_lookahead_tokens=model_runner.vllm_config.num_lookahead_tokens, - max_model_len=model_runner.max_model_len, - max_encoder_len=getattr(model_runner.model_state, "max_encoder_len", 0), - ) - - return block_count - - @torch.inference_mode() def warmup_kernels( model_runner: GPUModelRunner, @@ -82,10 +46,9 @@ def warmup_kernels( num_kv_cache_groups = len(kv_cache_groups) # Compute per-request block counts for each KV cache group. - block_count = _warmup_block_counter(model_runner) - kv_cache_specs = [group.kv_cache_spec for group in kv_cache_groups] - prefill_block_counts = [block_count(prompt_len, spec) for spec in kv_cache_specs] - decode_block_counts = [block_count(decode_len, spec) for spec in kv_cache_specs] + group_block_sizes = [g.kv_cache_spec.block_size for g in kv_cache_groups] + prefill_block_counts = [cdiv(prompt_len, bs) for bs in group_block_sizes] + decode_block_counts = [cdiv(decode_len, bs) for bs in group_block_sizes] decode_block_deltas = [ d - p for d, p in zip(decode_block_counts, prefill_block_counts) ] diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index c150cf791e..e6dbb57163 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -204,7 +204,6 @@ from vllm.v1.spec_decode.draft_prob_alignment import ( clone_draft_prob_token_ids, get_aligned_draft_probs, - get_aligned_draft_scalar_values, ) from vllm.v1.spec_decode.dspark import DSparkProposer from vllm.v1.spec_decode.eagle import EagleProposer @@ -271,19 +270,6 @@ _SM70_DECODE_EVENT_TRACE_CONFIG_LOGGED = False -def _unwrap_pipeline_intermediate_hidden_states( - output: torch.Tensor | IntermediateTensors, -) -> torch.Tensor: - if not isinstance(output, IntermediateTensors): - return output - try: - return output.tensors["hidden_states"] - except KeyError as exc: - raise RuntimeError( - "Pipeline intermediate output must contain a 'hidden_states' tensor." - ) from exc - - def _dflash_ddtree_debug_enabled() -> bool: return os.getenv("VLLM_DFLASH_DDTREE_DEBUG", "0") == "1" @@ -405,10 +391,7 @@ def _dflash_ddtree_target_forward_nvtx_enabled() -> bool: def _dflash_ddtree_target_forward_profiler_step() -> int: - raw = os.getenv( - "VLLM_SM70_SPEC_TARGET_FORWARD_PROFILER_STEP", - os.getenv("VLLM_DFLASH_DDTREE_TARGET_FORWARD_PROFILER_STEP", "0"), - ) + raw = os.getenv("VLLM_DFLASH_DDTREE_TARGET_FORWARD_PROFILER_STEP", "0") try: return max(0, int(raw)) except ValueError: @@ -1158,19 +1141,13 @@ def _dflash_ddtree_target_forward_profile_scope( if ( not use_spec_decode or self.speculative_config is None - or not ( - self.speculative_config.use_dflash_ddtree() - or self.speculative_config.use_dspark() - ) + or not self.speculative_config.use_dflash_ddtree() or self.device.type != "cuda" ): yield return - nvtx_enabled = bool( - _dflash_ddtree_target_forward_nvtx_enabled() - or os.getenv("VLLM_SM70_SPEC_TARGET_FORWARD_NVTX", "0") == "1" - ) + nvtx_enabled = _dflash_ddtree_target_forward_nvtx_enabled() profiler_step = _dflash_ddtree_target_forward_profiler_step() if not nvtx_enabled and profiler_step <= 0: yield @@ -1178,13 +1155,9 @@ def _dflash_ddtree_target_forward_profile_scope( step = getattr(self, "_dflash_ddtree_target_forward_profile_step", 0) + 1 self._dflash_ddtree_target_forward_profile_step = step - label_prefix = ( - "ddtree_target_forward" - if self.speculative_config.use_dflash_ddtree() - else "dspark_target_forward" - ) label = ( - f"{label_prefix}:step={step}:tokens={num_tokens}:reqs={num_reqs}" + "ddtree_target_forward" + f":step={step}:tokens={num_tokens}:reqs={num_reqs}" f":mode={cudagraph_mode.name}:pid={os.getpid()}" ) profile_this_step = profiler_step > 0 and step == profiler_step @@ -1400,23 +1373,6 @@ def __init__( # Async scheduling self.use_async_scheduling = bool(self.scheduler_config.async_scheduling) - self.dspark_confidence_scheduling = bool( - self.speculative_config is not None - and self.speculative_config.use_dspark() - and ( - self.speculative_config.dspark_confidence_threshold > 0.0 - or ( - self.speculative_config.dspark_max_verification_tokens is not None - and self.speculative_config.dspark_max_verification_tokens - < self.num_spec_tokens - ) - ) - ) - if self.dspark_confidence_scheduling and self.use_async_scheduling: - raise ValueError( - "DSpark confidence prefix scheduling currently requires " - "synchronous scheduling." - ) self._sm70_async_worker_execute_trace_step = 0 self._sm70_async_worker_sample_trace_step = 0 self._sm70_async_worker_input_prep_trace_step = 0 @@ -1543,7 +1499,7 @@ def __init__( elif self.speculative_config.use_dspark(): self.drafter = DSparkProposer(self.vllm_config, self.device, self) self.use_aux_hidden_state_outputs = True - elif self.speculative_config.use_dflash_ddtree(): + elif self.speculative_config.use_dflash(): self.drafter = DFlashProposer(self.vllm_config, self.device, self) self.use_aux_hidden_state_outputs = ( self.drafter.eagle3_use_aux_hidden_state @@ -1811,18 +1767,6 @@ def __init__( self._draft_probs: torch.Tensor | None = None self._draft_prob_req_ids: list[str] | None = None self._draft_prob_token_ids: list[list[int]] | torch.Tensor | None = None - self._draft_confidence_logits: torch.Tensor | None = None - self._draft_confidence_req_ids: list[str] | None = None - self._draft_confidence_token_ids: list[list[int]] | torch.Tensor | None = None - self._dspark_verification_lengths: torch.Tensor | None = None - self._dspark_verification_lengths_cpu: torch.Tensor | None = None - if self.dspark_confidence_scheduling: - self._dspark_verification_lengths_cpu = torch.empty( - self.max_num_reqs, - dtype=torch.int32, - device="cpu", - pin_memory=self.pin_memory, - ) self._dflash_ddtree_payloads: tuple[DDTreeDraftPayload, ...] | None = None self._ddtree_parent_metadata: DDTreeParentMetadata | None = None self._ddtree_accepted_rows_cpu_sidecar: list[list[int]] | None = None @@ -1864,12 +1808,6 @@ def __init__( self.valid_sampled_token_count_cpu: torch.Tensor | None = None self.draft_token_ids_cpu: torch.Tensor | None = None self.num_accepted_tokens_event: torch.Event | None = None - # Row ownership for the runner-owned accepted-token D2H snapshots. - # Keep the request object as well as its ID so abort-and-resubmit with - # the same ID is still treated as a new request. - self._mamba_accepted_token_state_rows: dict[ - str, tuple[int, CachedRequestState] - ] = {} if self.num_spec_tokens: self.draft_token_ids_event = torch.Event() self.num_accepted_tokens_event = torch.Event() @@ -2188,33 +2126,10 @@ def _warmup_sm70_aux_kernels(self) -> None: and mamba_utils.warmup_batch_memcpy_kernel(self.device) ): warmed.append("mamba_batch_memcpy") - if ( - self.cache_config.mamba_cache_mode == "align" - and self.speculative_config is not None - and self.model_config.is_hybrid - ): - mamba_bufs = self._get_mamba_bufs() - postprocess_ctx = mamba_bufs.postprocess_align - assert postprocess_ctx is not None - if not postprocess_ctx.is_initialized: - postprocess_ctx.initialize_from_forward_context( - self.kv_cache_config, - self.compilation_config.static_forward_context, - self.model.get_mamba_state_copy_func(), - [ - self.input_batch.block_table[group_id].get_device_tensor(1) - for group_id in postprocess_ctx.mamba_group_ids - ], - ) - if postprocess_ctx.warmup_fused_postprocess(): - warmed.append("mamba_spec_postprocess") drafter = getattr(self, "drafter", None) mtp_warmup = getattr(drafter, "warmup_sm70_mtp_hotpath_kernels", None) if mtp_warmup is not None: warmed.extend(mtp_warmup()) - mtp_moe_warmup = getattr(drafter, "warmup_sm70_mtp_moe_kernels", None) - if mtp_moe_warmup is not None: - warmed.extend(mtp_moe_warmup()) dflash_warmup = getattr(drafter, "warmup_sm70_dflash_hotpath_kernels", None) if dflash_warmup is not None: warmed.extend(dflash_warmup()) @@ -2674,10 +2589,6 @@ def _update_states_after_model_execute( # Count only the contiguous accepted prefix. Values after the first -1 # are rejected/padding slots and may contain stale token ids. num_reqs = output_token_ids.size(0) - self._mamba_accepted_token_state_rows = { - req_id: (i, self.requests[req_id]) - for i, req_id in enumerate(self.input_batch.req_ids[:num_reqs]) - } profile_count_t0 = time.perf_counter() if profile_enabled else 0.0 sidecar_values = None if ( @@ -2804,16 +2715,19 @@ def _update_states_after_model_execute( else None ), ) - # CPU postprocess updates InputBatch in place. Mirror its final - # values into the runner-owned snapshot used by the next step. - self.num_accepted_tokens.cpu[:num_reqs].copy_( - self.input_batch.num_accepted_tokens_cpu_tensor[:num_reqs] - ) - self.spec_state_slot_selectors.cpu[:num_reqs].copy_( - self.input_batch.spec_num_accepted_tokens_cpu_tensor[:num_reqs] - ) else: profile_mamba_stage_t0 = time.perf_counter() if profile_enabled else 0.0 + if spec_state_slot_selectors_cpu is not None: + self.input_batch.spec_num_accepted_tokens_cpu_tensor[ + :num_reqs + ].copy_(spec_state_slot_selectors_cpu) + else: + self.input_batch.spec_num_accepted_tokens_cpu_tensor[ + :num_reqs + ].copy_( + self.spec_state_slot_selectors.gpu[:num_reqs], + non_blocking=True, + ) assert mamba_bufs.postprocess_align is not None mamba_utils.stage_postprocess_inputs_to_gpu( mamba_bufs.postprocess_align, @@ -2835,9 +2749,11 @@ def _update_states_after_model_execute( num_reqs=num_reqs, num_accepted_tokens_gpu=self.num_accepted_tokens.gpu, spec_state_slot_selectors_gpu=(self.spec_state_slot_selectors.gpu), - num_accepted_tokens_cpu_tensor=self.num_accepted_tokens.cpu, + num_accepted_tokens_cpu_tensor=( + self.input_batch.num_accepted_tokens_cpu_tensor + ), spec_num_accepted_tokens_cpu_tensor=( - self.spec_state_slot_selectors.cpu + self.input_batch.spec_num_accepted_tokens_cpu_tensor ), input_batch=self.input_batch, kv_cache_config=self.kv_cache_config, @@ -2870,10 +2786,10 @@ def _update_states_after_model_execute( spec_state_slot_selectors_cpu ) else: - self.num_accepted_tokens.cpu[:num_reqs].copy_( + self.input_batch.num_accepted_tokens_cpu_tensor[:num_reqs].copy_( self.num_accepted_tokens.gpu[:num_reqs], non_blocking=True ) - self.spec_state_slot_selectors.cpu[:num_reqs].copy_( + self.input_batch.spec_num_accepted_tokens_cpu_tensor[:num_reqs].copy_( self.spec_state_slot_selectors.gpu[:num_reqs], non_blocking=True, ) @@ -4742,45 +4658,6 @@ def _get_encoder_seq_lens( return encoder_seq_lens, encoder_seq_lens_cpu - def _sync_mamba_accepted_token_state( - self, - scheduler_output: "SchedulerOutput", - num_reqs: int, - ) -> None: - """Remap the previous step's accepted-token state by request. - - The GPU postprocess copies into runner-owned CPU buffers. InputBatch - rows can be removed, reused, or condensed before this synchronization, - so row-for-row writeback is not safe even with synchronous scheduling. - """ - previous_counts = self.num_accepted_tokens.np.copy() - previous_selectors = self.spec_state_slot_selectors.np.copy() - previous_rows = self._mamba_accepted_token_state_rows - reset_req_ids = set(scheduler_output.scheduled_cached_reqs.resumed_req_ids) - reset_req_ids.update( - req_data.req_id for req_data in scheduler_output.scheduled_new_reqs - ) - - current_counts = self.num_accepted_tokens.np[:num_reqs] - current_selectors = self.spec_state_slot_selectors.np[:num_reqs] - for current_idx, req_id in enumerate(self.input_batch.req_ids[:num_reqs]): - previous = previous_rows.get(req_id) - if ( - previous is not None - and req_id not in reset_req_ids - and self.requests.get(req_id) is previous[1] - ): - previous_idx = previous[0] - current_counts[current_idx] = previous_counts[previous_idx] - current_selectors[current_idx] = previous_selectors[previous_idx] - else: - # A new/restored request has no accepted speculative prefix. - current_counts[current_idx] = 1 - current_selectors[current_idx] = 1 - - self.input_batch.num_accepted_tokens_cpu[:num_reqs] = current_counts - self.input_batch.spec_num_accepted_tokens_cpu[:num_reqs] = current_selectors - def _prepare_inputs( self, scheduler_output: "SchedulerOutput", @@ -4992,7 +4869,32 @@ def _prepare_inputs( self.num_accepted_tokens_event, "GPUModelRunner.num_accepted_tokens_event.synchronize", ) - self._sync_mamba_accepted_token_state(scheduler_output, num_reqs) + # Async mode: condense() reordered indices, use prev_positions mapping + if self.use_async_scheduling and prev_req_id_to_index: + prev_idx = self.prev_positions.np[:num_reqs] + new_mask = prev_idx < 0 + prev_idx_or_zero = np.where(new_mask, 0, prev_idx) + align_counts = self.input_batch.num_accepted_tokens_cpu[ + prev_idx_or_zero + ].copy() + align_counts[new_mask] = 1 + self.input_batch.num_accepted_tokens_cpu[:num_reqs] = align_counts + spec_counts = self.input_batch.spec_num_accepted_tokens_cpu[ + prev_idx_or_zero + ] + spec_counts = spec_counts.copy() + spec_counts[new_mask] = 1 + self.input_batch.spec_num_accepted_tokens_cpu[:num_reqs] = spec_counts + self.num_accepted_tokens.np[:num_reqs] = align_counts + self.spec_state_slot_selectors.np[:num_reqs] = spec_counts + else: + # Non-async mode: use values directly + self.num_accepted_tokens.np[:num_reqs] = ( + self.input_batch.num_accepted_tokens_cpu[:num_reqs] + ) + self.spec_state_slot_selectors.np[:num_reqs] = ( + self.input_batch.spec_num_accepted_tokens_cpu[:num_reqs] + ) self.num_accepted_tokens.np[num_reqs:].fill(1) self.spec_state_slot_selectors.np[num_reqs:].fill(1) self._copy_buffer_to_gpu(self.num_accepted_tokens) @@ -5251,7 +5153,6 @@ def _build_attention_metadata( slot_mappings: dict[int, torch.Tensor] | None = None, ddtree_parent_metadata: DDTreeParentMetadata | None = None, cudagraph_capture_max_seq_len: int | None = None, - cudagraph_graph_variant: int | None = None, ) -> tuple[PerLayerAttnMetadata, CommonAttentionMetadata | None]: """ :return: tuple[attn_metadata, spec_decode_common_attn_metadata] @@ -5356,20 +5257,6 @@ def _get_block_table(kv_cache_gid: int): seq_lens_cpu = None num_computed_tokens_cpu = None - # Prefix-anchored SWA: pass per-request prompt lengths so the - # attention backend can keep the prefix globally visible. The backend - # owns the persistent device buffer. - prefix_anchor_lens = None - if ( - getattr( - self.vllm_config.attention_config, - "prefix_anchored_decode_window", - None, - ) - is not None - ): - prefix_anchor_lens = num_prompt_tokens_cpu - cm_base = CommonAttentionMetadata( query_start_loc=self.query_start_loc.gpu[: num_reqs_padded + 1], query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs_padded + 1], @@ -5384,10 +5271,8 @@ def _get_block_table(kv_cache_gid: int): block_table_tensor=block_table_gid_0, slot_mapping=slot_mapping_gid_0, causal=True, - cudagraph_graph_variant=cudagraph_graph_variant, is_prefilling=is_prefilling, positions=self.positions[:num_tokens_padded], - prefix_anchor_lens=prefix_anchor_lens, ) current_mamba_state_block_ids_by_gid: dict[int, torch.Tensor] = {} @@ -6465,14 +6350,9 @@ def sync_and_gather_intermediate_tensors( local_len = num_tokens // tp v = get_tp_group().all_gather(v[:local_len], dim=0) - destination = self.intermediate_tensors[k][:num_tokens] - source = v[:num_tokens] - if ( - destination.data_ptr() != source.data_ptr() - or destination.shape != source.shape - or destination.stride() != source.stride() - ): - destination.copy_(source, non_blocking=True) + self.intermediate_tensors[k][:num_tokens].copy_( + v[:num_tokens], non_blocking=True + ) return IntermediateTensors( {k: v[:num_tokens] for k, v in self.intermediate_tensors.items()} @@ -7008,26 +6888,6 @@ def _sample( _maybe_dump_sm70_qwen_layer_graph_buffers("pre_sample") _maybe_sync_sm70_sample_tensors(sample_hidden_states) if spec_decode_metadata is None: - if logits is None and self._can_use_sm70_compact_topk20_tokens( - scheduler_output, spec_decode_metadata - ): - topk_token_ids, topk_logits = self.model.get_topk_tokens_and_logits( - sample_hidden_states, - 20, - ) - compact_sampled = self.sampler.sm70_compact_topk20_pairs_sample( - topk_logits, - topk_token_ids, - sampling_metadata, - ) - if compact_sampled is not None: - logger.info_once( - "SM70 TP-local exact top-k20 random sampling path enabled." - ) - return SamplerOutput( - sampled_token_ids=compact_sampled.to(torch.int32).unsqueeze(-1), - logprobs_tensors=None, - ) sampler_output = self._sample_greedy_token_fastpath( logits, sampling_metadata, @@ -7460,9 +7320,6 @@ def _sample( ) draft_probs = self._get_spec_decode_draft_probs(spec_decode_metadata) - draft_confidence_logits = self._get_spec_decode_confidence_logits( - spec_decode_metadata - ) if ( self.speculative_config is not None and self.speculative_config.method == "mtp" @@ -7482,7 +7339,6 @@ def _sample( draft_probs, logits, sampling_metadata, - draft_confidence_logits=draft_confidence_logits, ) target_candidate_ids = self.rejection_sampler.take_last_target_candidate_ids() if target_candidate_ids is not None and hasattr( @@ -7582,88 +7438,6 @@ def _can_use_greedy_token_fastpath( self._trace_greedy_token_fastpath("enabled") return True - def _can_use_sm70_compact_topk20_tokens( - self, - scheduler_output: "SchedulerOutput", - spec_decode_metadata: SpecDecodeMetadata | None, - ) -> bool: - if not envs.VLLM_SM70_COMPACT_TOPK20_SAMPLER: - return False - if not envs.VLLM_SM70_TP_LOCAL_TOPK20_SAMPLER: - return False - if spec_decode_metadata is not None: - return self._reject_sm70_compact_topk20_tokens("spec_decode") - if self.is_pooling_model or self.broadcast_pp_output: - return self._reject_sm70_compact_topk20_tokens("pooling_or_pp_broadcast") - if scheduler_output.has_structured_output_requests: - return self._reject_sm70_compact_topk20_tokens("structured_output") - if self.device.type != "cuda": - return self._reject_sm70_compact_topk20_tokens("non_cuda") - if torch.cuda.get_device_capability(self.device) != (7, 0): - return self._reject_sm70_compact_topk20_tokens("non_sm70") - if self.model_config.get_vocab_size() != 248320: - return self._reject_sm70_compact_topk20_tokens("vocab_size") - if not hasattr(self.model, "get_topk_tokens_and_logits"): - return self._reject_sm70_compact_topk20_tokens("missing_model_topk") - if self.input_batch.num_reqs != 1: - return self._reject_sm70_compact_topk20_tokens("batch_size") - - sampling_metadata = self.input_batch.sampling_metadata - if not sampling_metadata.all_random or sampling_metadata.all_greedy: - return self._reject_sm70_compact_topk20_tokens("sampling_mode") - if sampling_metadata.max_num_logprobs is not None: - return self._reject_sm70_compact_topk20_tokens("logprobs") - if sampling_metadata.logprob_token_ids: - return self._reject_sm70_compact_topk20_tokens("logprob_token_ids") - if not sampling_metadata.no_penalties: - return self._reject_sm70_compact_topk20_tokens("penalties") - if sampling_metadata.allowed_token_ids_mask is not None: - return self._reject_sm70_compact_topk20_tokens("allowed_token_ids") - if sampling_metadata.bad_words_token_ids: - return self._reject_sm70_compact_topk20_tokens("bad_words") - holder = sampling_metadata.thinking_budget_state_holder - if holder is not None and holder.has_tracked_requests(): - return self._reject_sm70_compact_topk20_tokens("thinking_budget") - if not self._non_argmax_logits_processors_inactive(sampling_metadata): - return self._reject_sm70_compact_topk20_tokens("logits_processor") - if not self._sm70_argmax_logits_processors_inactive(sampling_metadata): - return self._reject_sm70_compact_topk20_tokens( - "argmax_invariant_logits_processor" - ) - if sampling_metadata.top_k_cpu != (20,): - return self._reject_sm70_compact_topk20_tokens("top_k") - - top_p_cpu = sampling_metadata.top_p_cpu - if top_p_cpu is None or len(top_p_cpu) != 1 or abs(top_p_cpu[0] - 0.95) > 1e-6: - return self._reject_sm70_compact_topk20_tokens("top_p") - temperature_cpu = sampling_metadata.temperature_cpu - if ( - temperature_cpu is None - or len(temperature_cpu) != 1 - or abs(temperature_cpu[0] - 1.0) > 1e-6 - ): - return self._reject_sm70_compact_topk20_tokens("temperature") - return True - - @staticmethod - def _reject_sm70_compact_topk20_tokens(reason: str) -> bool: - logger.info_once("SM70 TP-local top-k20 route rejected: %s", reason) - return False - - @staticmethod - def _sm70_argmax_logits_processors_inactive( - sampling_metadata: SamplingMetadata, - ) -> bool: - from vllm.v1.sample.logits_processor.builtin import MinPLogitsProcessor - - for processor in sampling_metadata.logitsprocs.argmax_invariant: - if isinstance(processor, MinPLogitsProcessor): - if processor.min_p_count: - return False - continue - return False - return True - def _trace_greedy_token_fastpath(self, reason: str) -> None: if not self.sm70_greedy_token_fastpath_trace: return @@ -8012,7 +7786,7 @@ def _determine_batch_execution_and_padding( attention_context_len is None and uniform_decode and num_reqs > 0 - and self.cudagraph_dispatcher.has_attention_context_specialization + and self.cudagraph_dispatcher.has_attention_context_buckets ): attention_context_len = int( self.optimistic_seq_lens_cpu[:num_reqs].max().item() @@ -8664,9 +8438,6 @@ def execute_model( self._can_use_greedy_token_fastpath( scheduler_output, spec_decode_metadata ) - or self._can_use_sm70_compact_topk20_tokens( - scheduler_output, spec_decode_metadata - ) or self._can_use_ddtree_greedy_top_tokens( scheduler_output, spec_decode_metadata ) @@ -9021,10 +8792,6 @@ def sample_tokens( self._draft_probs = None self._draft_prob_req_ids = None self._draft_prob_token_ids = None - self._draft_confidence_logits = None - self._draft_confidence_req_ids = None - self._draft_confidence_token_ids = None - self._dspark_verification_lengths = None self._dflash_ddtree_payloads = None self._ddtree_parent_metadata = None self._ddtree_accepted_rows_cpu_sidecar = None @@ -9168,10 +8935,6 @@ def propose_draft_token_ids(sampled_token_ids): self._draft_probs = None self._draft_prob_req_ids = None self._draft_prob_token_ids = None - self._draft_confidence_logits = None - self._draft_confidence_req_ids = None - self._draft_confidence_token_ids = None - self._dspark_verification_lengths = None self._copy_draft_token_ids_to_cpu(scheduler_output) with record_function_or_nullcontext("gpu_model_runner: bookkeep"): @@ -9476,15 +9239,6 @@ def _copy_draft_token_ids_to_cpu( self.draft_token_ids_cpu[:num_reqs].copy_( draft_token_ids, non_blocking=True ) - if self.dspark_confidence_scheduling: - lengths = self._dspark_verification_lengths - lengths_cpu = self._dspark_verification_lengths_cpu - if lengths is None or lengths_cpu is None: - raise RuntimeError( - "DSpark confidence scheduling is enabled but the " - "proposer did not return verification lengths." - ) - lengths_cpu[:num_reqs].copy_(lengths[:num_reqs], non_blocking=True) else: # No copy needed, just zero-out cpu tensor. self.draft_token_ids_cpu[:num_reqs] = 0 @@ -9505,16 +9259,7 @@ def _get_draft_token_ids_cpu(self) -> tuple[list[list[int]], list[str]]: self.draft_token_ids_event, "GPUModelRunner.draft_token_ids_event.synchronize", ) - draft_token_ids = self.draft_token_ids_cpu[: len(req_ids)].tolist() - if self.dspark_confidence_scheduling: - lengths_cpu = self._dspark_verification_lengths_cpu - assert lengths_cpu is not None - lengths = lengths_cpu[: len(req_ids)].tolist() - draft_token_ids = [ - token_ids[: max(0, min(int(length), len(token_ids)))] - for token_ids, length in zip(draft_token_ids, lengths, strict=True) - ] - return draft_token_ids, req_ids + return self.draft_token_ids_cpu[: len(req_ids)].tolist(), req_ids def _copy_valid_sampled_token_count( self, next_token_ids: torch.Tensor, valid_sampled_tokens_count: torch.Tensor @@ -9564,22 +9309,6 @@ def _get_spec_decode_draft_probs( spec_decode_metadata=spec_decode_metadata, ) - def _get_spec_decode_confidence_logits( - self, spec_decode_metadata: SpecDecodeMetadata - ) -> torch.Tensor | None: - # Alignment has a small but non-zero launch/copy cost. Keep it out of - # ordinary serving until confidence scheduling is explicitly enabled; - # alignment dumps are the calibration path used before that gate. - if not envs.VLLM_SPEC_DUMP_ALIGNMENT: - return None - return get_aligned_draft_scalar_values( - req_ids=self.input_batch.req_ids, - values=self._draft_confidence_logits, - value_req_ids=self._draft_confidence_req_ids, - value_token_ids=self._draft_confidence_token_ids, - spec_decode_metadata=spec_decode_metadata, - ) - def propose_draft_token_ids( self, scheduler_output: "SchedulerOutput", @@ -9598,10 +9327,6 @@ def propose_draft_token_ids( self._draft_probs = None self._draft_prob_req_ids = None self._draft_prob_token_ids = None - self._draft_confidence_logits = None - self._draft_confidence_req_ids = None - self._draft_confidence_token_ids = None - self._dspark_verification_lengths = None self._dflash_ddtree_payloads = None self._ddtree_parent_metadata = None self._ddtree_accepted_rows_cpu_sidecar = None @@ -9726,7 +9451,7 @@ def propose_draft_token_ids( elif ( spec_config.use_eagle() - or spec_config.use_dflash_ddtree() + or spec_config.use_dflash() or spec_config.uses_draft_model() ): assert isinstance( @@ -9918,18 +9643,6 @@ def propose_draft_token_ids( self._draft_prob_token_ids = clone_draft_prob_token_ids( draft_token_ids ) - if hasattr(self.drafter, "take_last_confidence_logits"): - confidence_logits = self.drafter.take_last_confidence_logits() - if confidence_logits is not None: - self._draft_confidence_logits = confidence_logits - self._draft_confidence_req_ids = self.input_batch.req_ids.copy() - self._draft_confidence_token_ids = clone_draft_prob_token_ids( - draft_token_ids - ) - if hasattr(self.drafter, "take_last_verification_lengths"): - verification_lengths = self.drafter.take_last_verification_lengths() - if verification_lengths is not None: - self._dspark_verification_lengths = verification_lengths if hasattr(self.drafter, "take_last_ddtree_payloads"): self._dflash_ddtree_payloads = self.drafter.take_last_ddtree_payloads() first_payload = ( @@ -10788,11 +10501,6 @@ def _dummy_run( if is_graph_capturing else None ), - cudagraph_graph_variant=( - batch_desc.graph_variant - if batch_descriptor_override is not None - else None - ), ) _sm70_profile_trace( "_dummy_run attention metadata built num_tokens=%s " @@ -10907,15 +10615,6 @@ def _dummy_run( else: hidden_states = outputs - if isinstance(hidden_states, IntermediateTensors): - assert not get_pp_group().is_last_rank - # Non-last PP ranks return intermediate tensors instead of - # logits-bearing hidden states. Dummy-run callers only need a - # tensor to preserve the profiling/capture return contract. - hidden_states = _unwrap_pipeline_intermediate_hidden_states( - hidden_states - ) - if self.speculative_config and ( self.speculative_config.use_eagle() or self.speculative_config.uses_draft_model() @@ -11102,21 +10801,6 @@ def _dummy_sampler_run( logits, dummy_metadata, ) - # The mixed-sampling warmup above passes a non-null is_greedy - # tensor. Pure greedy serving passes None and otherwise triggers a - # separate Triton compile during the first decode verifier step. - greedy_metadata = replace( - dummy_metadata, - temperature=torch.zeros_like(dummy_metadata.temperature), - all_greedy=True, - all_random=False, - ) - self.rejection_sampler( - dummy_spec_decode_metadata, - None, - logits, - greedy_metadata, - ) return sampler_output def _dummy_pooler_run_task( diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index c131a019b1..6ca3c687f7 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -220,9 +220,9 @@ def _scoped_allocator_max_split(self, max_split_size_mb: int): yield return - set_allocator_settings = getattr( - torch._C, "_accelerator_setAllocatorSettings", None - ) + set_allocator_settings = getattr(torch._C, + "_accelerator_setAllocatorSettings", + None) if set_allocator_settings is None: yield return @@ -788,92 +788,6 @@ def sample_tokens( ) -> ModelRunnerOutput | AsyncModelRunnerOutput: return self.model_runner.sample_tokens(grammar_output) - def _use_sm70_static_pp_hidden_transfer(self, num_tokens: int) -> bool: - """Whether this step has the exact metadata-free PP tensor contract.""" - if not envs.VLLM_SM70_PP_STATIC_HIDDEN_TRANSFER or num_tokens != 1: - return False - cached = getattr(self, "_sm70_static_pp_hidden_contract", None) - if cached is not None: - return bool(cached) - - config = self.vllm_config - parallel_config = config.parallel_config - compilation_config = config.compilation_config - model_config = config.model_config - enabled = not ( - parallel_config.pipeline_parallel_size != 2 - or parallel_config.tensor_parallel_size != 4 - or parallel_config.enable_dbo - or parallel_config.ubatch_size > 1 - or config.scheduler_config.max_num_seqs != 1 - or getattr(config, "speculative_config", None) is not None - or compilation_config.cudagraph_mode == CUDAGraphMode.NONE - or compilation_config.pass_config.enable_sp - or model_config.dtype != torch.float16 - ) - - schema_matches = False - if enabled: - try: - schema = self.model_runner.model.make_empty_intermediate_tensors( - batch_size=1, - dtype=model_config.dtype, - device=torch.device("meta"), - ).tensors - hidden_states = schema.get("hidden_states") - schema_matches = bool( - tuple(schema) == ("hidden_states",) - and isinstance(hidden_states, torch.Tensor) - and hidden_states.shape == (1, 4, 4096) - and hidden_states.dtype == torch.float16 - and hidden_states.is_contiguous() - ) - except (AttributeError, RuntimeError, TypeError): - # This is an optional fast path; unsupported model-provided - # schema construction retains the metadata protocol. - schema_matches = False - enabled = bool( - enabled - and schema_matches - and current_platform.is_cuda() - and current_platform.is_device_capability((7, 0)) - ) - self._sm70_static_pp_hidden_contract = enabled - return enabled - - @staticmethod - def _is_static_pp_hidden_tensor_dict( - tensor_dict: dict[str, torch.Tensor], num_tokens: int - ) -> bool: - if tuple(tensor_dict) != ("hidden_states",): - return False - hidden_states = tensor_dict["hidden_states"] - return bool( - hidden_states.shape == (num_tokens, 4, 4096) - and hidden_states.dtype == torch.float16 - and hidden_states.is_cuda - and hidden_states.is_contiguous() - ) - - def _static_pp_hidden_recv_buffers( - self, num_tokens: int - ) -> dict[str, torch.Tensor] | None: - if not self._use_sm70_static_pp_hidden_transfer(num_tokens): - return None - buffers = self.model_runner.intermediate_tensors - if buffers is None or tuple(buffers.tensors) != ("hidden_states",): - raise RuntimeError( - "static PP hidden transfer requires one persistent " - "hidden_states input buffer" - ) - hidden_states = buffers.tensors["hidden_states"][:num_tokens] - tensor_dict = {"hidden_states": hidden_states} - if not self._is_static_pp_hidden_tensor_dict(tensor_dict, num_tokens): - raise RuntimeError( - "static PP hidden transfer input buffer violates its fixed schema" - ) - return tensor_dict - @torch.inference_mode() def execute_model( self, scheduler_output: "SchedulerOutput" @@ -921,20 +835,12 @@ def execute_model( } if forward_pass and not get_pp_group().is_first_rank: - tensor_dict = self._static_pp_hidden_recv_buffers(num_scheduled_tokens) - if tensor_dict is not None: - comm_handles = get_pp_group().irecv_tensor_dict_static(tensor_dict) - comm_postprocess: list[Callable[[], None]] = [] - logger.info_once( - "Default SM70 metadata-free PP hidden transfer enabled." - ) - else: - tensor_dict, comm_handles, comm_postprocess = ( - get_pp_group().irecv_tensor_dict( - all_gather_group=get_tp_group(), - all_gather_tensors=all_gather_tensors, - ) + tensor_dict, comm_handles, comm_postprocess = ( + get_pp_group().irecv_tensor_dict( + all_gather_group=get_tp_group(), + all_gather_tensors=all_gather_tensors, ) + ) assert tensor_dict is not None intermediate_tensors = AsyncIntermediateTensors( tensor_dict, @@ -964,23 +870,12 @@ def execute_model( and not get_pp_group().is_last_rank ) - # The exact SM70 B1 path sends the replicated tensor directly into - # the next stage's persistent graph-input buffer without CPU metadata. - if self._use_sm70_static_pp_hidden_transfer(num_scheduled_tokens): - if not self._is_static_pp_hidden_tensor_dict( - output.tensors, num_scheduled_tokens - ): - raise RuntimeError( - "static PP hidden transfer output violates its fixed schema" - ) - self._pp_send_work = get_pp_group().isend_tensor_dict_static(output.tensors) - logger.info_once("Default SM70 metadata-free PP hidden transfer enabled.") - else: - self._pp_send_work = get_pp_group().isend_tensor_dict( - output.tensors, - all_gather_group=get_tp_group(), - all_gather_tensors=all_gather_tensors, - ) + # launch non-blocking send of intermediate tensors + self._pp_send_work = get_pp_group().isend_tensor_dict( + output.tensors, + all_gather_group=get_tp_group(), + all_gather_tensors=all_gather_tensors, + ) return None diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 7ffb0c28f9..24f7536887 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -659,63 +659,6 @@ def run_fused_postprocess( HAS_DDTREE_ACCEPTED_NODES=has_ddtree, ) - def warmup_fused_postprocess(self) -> bool: - """Compile MTP postprocess signatures without copying model state.""" - if not self.is_initialized: - 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]) - warmup_sizes.add(max(1, min(max_num_reqs, 4))) - for num_reqs in sorted(warmup_sizes): - num_accepted_tokens = torch.ones(num_reqs, dtype=torch.int32, device=device) - spec_state_slot_selectors = torch.ones( - num_reqs, dtype=torch.int32, device=device - ) - mamba_state_idx = torch.zeros(num_reqs, dtype=torch.int32, device=device) - num_scheduled_tokens = torch.ones( - num_reqs, dtype=torch.int32, device=device - ) - # Put every request one token past a block boundary. The kernel - # exits before dereferencing state or block-table pointers. - num_computed_tokens = torch.full( - (num_reqs,), self.block_size, dtype=torch.int32, device=device - ) - num_draft_tokens = torch.zeros(num_reqs, dtype=torch.int32, device=device) - ddtree_accepted_node_indices = torch.empty( - (1, 1), dtype=torch.int32, device=device - ) - - postprocess_mamba_fused_kernel[(num_reqs, total_states)]( - num_accepted_tokens, - spec_state_slot_selectors, - ddtree_accepted_node_indices, - ddtree_accepted_node_indices.stride(0), - mamba_state_idx, - num_scheduled_tokens, - num_computed_tokens, - num_draft_tokens, - self.block_table_ptrs, - self.block_table_stride_req, - self.state_base_addrs, - self.state_block_strides, - self.state_elem_sizes, - self.state_inner_sizes, - self.state_conv_widths, - self.state_group_indices, - self.num_accepted_tokens_out, - self.spec_state_slot_selectors_out, - num_reqs, - block_size=self.block_size, - COPY_BLOCK_SIZE=1024, - HAS_DDTREE_ACCEPTED_NODES=False, - ) - torch.accelerator.synchronize() - return True - @dataclasses.dataclass class MambaBuffers: diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index 6b634557a5..fb02fdc935 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -37,48 +37,9 @@ logger = init_logger(__name__) -def _infer_segment_block_strides( - seg_addrs: list[int], - page_size_el: int, -) -> list[int]: - """Infer logical block strides from block-major segment address runs. - - Dense segments advance by one page. An interleaved pool exposes one segment - base per cell, exactly ``page_size_el * 4`` bytes apart, and advances to the - next logical block after every segment in that address run. Treat separate - runs independently so multiple pools do not disable one another. - """ - if page_size_el <= 0: - raise ValueError(f"page_size_el must be positive, got {page_size_el}.") - if len(set(seg_addrs)) != len(seg_addrs): - raise ValueError("Segment addresses must be unique.") - - block_strides = [page_size_el] * len(seg_addrs) - if len(seg_addrs) < 2: - return block_strides - - cell_bytes = page_size_el * 4 - ordered = sorted((addr, index) for index, addr in enumerate(seg_addrs)) - run_start = 0 - for run_end in range(1, len(ordered) + 1): - if ( - run_end < len(ordered) - and ordered[run_end][0] - ordered[run_end - 1][0] == cell_bytes - ): - continue - run = ordered[run_start:run_end] - if len(run) > 1: - run_stride = page_size_el * len(run) - for _, original_index in run: - block_strides[original_index] = run_stride - run_start = run_end - return block_strides - - @triton.jit(do_not_specialize=["n_blocks"]) def _zero_kv_blocks_kernel( seg_addrs_ptr, - seg_block_strides_ptr, block_ids_ptr, n_blocks, N_SEGS: tl.constexpr, @@ -94,9 +55,6 @@ def _zero_kv_blocks_kernel( seg_addrs_ptr holds absolute byte addresses (int64) for each segment, allowing segments to live in different CUDA allocations. - seg_block_strides_ptr holds the per-segment block-to-block stride in - 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). """ @@ -111,11 +69,9 @@ def _zero_kv_blocks_kernel( 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_id.to(tl.int64) * PAGE_SIZE_EL + chunk_index.to(tl.int64) * BLOCK_SIZE ) cols = tl.arange(0, BLOCK_SIZE).to(tl.int64) tl.store(ptr + offset + cols, tl.zeros([BLOCK_SIZE], dtype=tl.int32)) @@ -132,7 +88,7 @@ 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, int, int, int] | None = None self._id_cap: int = 0 self._ids_pinned: torch.Tensor | None = None self._ids_gpu: torch.Tensor | None = None @@ -199,7 +155,8 @@ def init_meta( 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}" + f"Non-uniform page sizes: {page_size_el} vs " + f"{cur_page_el}" ) block_stride_bytes = cur_bytes @@ -230,7 +187,8 @@ def init_meta( 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}" + f"Non-uniform page sizes: {page_size_el} vs " + f"{cur_page_el}" ) seg_addrs.append(dp) @@ -239,18 +197,6 @@ def init_meta( return blk_size = min(largest_power_of_2_divisor(page_size_el), 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) - seg_block_strides = torch.tensor( - _infer_segment_block_strides(seg_addrs, page_size_el), - dtype=torch.int64, - device=self.device, - ) - self._id_cap = 8192 self._ids_pinned = torch.empty( self._id_cap, @@ -260,17 +206,16 @@ def init_meta( self._ids_gpu = torch.empty(self._id_cap, dtype=torch.int64, device=self.device) self._meta = ( torch.tensor(seg_addrs, dtype=torch.uint64, device=self.device), - seg_block_strides, page_size_el, blk_size, - n_segs, + len(seg_addrs), ) 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, page_size_el, blk_size, n_segs = self._meta n_blocks = len(block_ids) if n_blocks > self._id_cap: self._id_cap = n_blocks * 2 @@ -289,7 +234,6 @@ def zero_block_ids(self, block_ids: list[int]) -> None: grid = (n_blocks * n_segs * (page_size_el // blk_size),) _zero_kv_blocks_kernel[grid]( seg_addrs, - seg_block_strides, idx, n_blocks, N_SEGS=n_segs, From 99d67d74c63c563460d5b82fa2a9d2b82f673206 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Wed, 26 Aug 2026 09:56:06 +0000 Subject: [PATCH 020/105] Bump webcolors from 24.11.1 to 25.10.0 Bumps [webcolors](https://github.com/ubernostrum/webcolors) from 24.11.1 to 25.10.0. - [Changelog](https://github.com/ubernostrum/webcolors/blob/trunk/docs/changelog.rst) - [Commits](https://github.com/ubernostrum/webcolors/compare/24.11.1...25.10.0) --- updated-dependencies: - dependency-name: webcolors dependency-version: 25.10.0 dependency-type: direct:production update-type: version-update:semver-major ... Signed-off-by: dependabot[bot] --- requirements/test/cuda.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index 831de114dd..4222ef0f8c 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -1120,7 +1120,7 @@ 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 From 5ac5a29e588445dce369ec21c01190142205f62d Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Wed, 26 Aug 2026 09:56:13 +0000 Subject: [PATCH 021/105] Bump pandas from 2.2.3 to 2.3.3 Bumps [pandas](https://github.com/pandas-dev/pandas) from 2.2.3 to 2.3.3. - [Release notes](https://github.com/pandas-dev/pandas/releases) - [Commits](https://github.com/pandas-dev/pandas/compare/v2.2.3...v2.3.3) --- updated-dependencies: - dependency-name: pandas dependency-version: 2.3.3 dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] --- requirements/test/cuda.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index 831de114dd..10f12e60db 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -597,7 +597,7 @@ packaging==24.2 # statsmodels # transformers # typepy -pandas==2.2.3 +pandas==2.3.3 # via # datasets # evaluate From 698d7b987b1915d5a22d72b29e48419f02ae5d26 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Wed, 26 Aug 2026 09:56:23 +0000 Subject: [PATCH 022/105] Bump wrapt from 1.17.2 to 2.3.0 Bumps [wrapt](https://github.com/GrahamDumpleton/wrapt) from 1.17.2 to 2.3.0. - [Release notes](https://github.com/GrahamDumpleton/wrapt/releases) - [Changelog](https://github.com/GrahamDumpleton/wrapt/blob/develop/docs/changes.rst) - [Commits](https://github.com/GrahamDumpleton/wrapt/compare/1.17.2...2.3.0) --- updated-dependencies: - dependency-name: wrapt dependency-version: 2.3.0 dependency-type: direct:production update-type: version-update:semver-major ... Signed-off-by: dependabot[bot] --- requirements/test/cuda.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index 831de114dd..a90b545ab2 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -1126,7 +1126,7 @@ 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 From 720153de24dbf4ebb234897d4872cc1cb32ea544 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 26 Aug 2026 18:15:59 +0800 Subject: [PATCH 023/105] [Kernel][SM70] Default exact DFlash2 long-context verifier (#9) Source-audited against private main b06f54cee3. Includes strict route/shape admission, bitwise-equivalent long-context kernel optimizations, static-gate fixes, default-on promotion, and explicit rollback. Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../benchmark_flash_v100_dflash2_swa.py | 151 +++++++++++ docs/design/sm70_dflash2_long_verify_decay.md | 244 +++++++++++++++++- .../kernel/flash_decode_paged.cu | 186 +++++++++---- .../kernel/fused_mha_forward_paged.cu | 6 + .../test_sm70_flash_v100_dflash2_swa.py | 76 ++++++ .../test_sm70_flash_v100_grouped_verify.py | 72 ++++++ tests/test_envs.py | 13 + tests/v1/spec_decode/test_dflash2.py | 14 +- vllm/envs.py | 12 +- 9 files changed, 714 insertions(+), 60 deletions(-) create mode 100644 benchmarks/kernels/benchmark_flash_v100_dflash2_swa.py create mode 100644 tests/kernels/attention/test_sm70_flash_v100_dflash2_swa.py 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/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/flash-attention-v100/kernel/flash_decode_paged.cu b/flash-attention-v100/kernel/flash_decode_paged.cu index b6417092d4..7d704a9aaa 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; @@ -1777,6 +1788,11 @@ constexpr int kGroupedVerifyHeadsPerCta = 6; constexpr int kGroupedVerifyRows = kGroupedVerifyMaxQ * kGroupedVerifyHeadsPerCta; 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. @@ -1801,16 +1817,17 @@ constexpr int kGroupedVerifyOutputTilesPerWarp = 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, @@ -1853,13 +1870,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 +1884,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 +1905,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, @@ -1940,8 +1959,31 @@ __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; @@ -1952,11 +1994,11 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ const int local_head = row % kGroupedVerifyHeadsPerCta; 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 +2006,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 +2027,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) { @@ -2010,7 +2057,7 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ 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 +2085,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) { @@ -2067,10 +2118,13 @@ __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 { @@ -2085,7 +2139,7 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ 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 +2151,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 +2172,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 +2199,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]); } @@ -3516,11 +3571,13 @@ at::Tensor flash_attention_grouped_verify_paged( const dim3 partial_grid(1, kGroupedVerifySplits, 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(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>; \ + 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 +3589,8 @@ 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< \ + 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,15 +3604,40 @@ 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(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(TWO_PASS, 1648, SINGLE_QUERY, true, true); \ + } else if (stage_page_ids) { \ + LAUNCH_GROUPED_VERIFY_PARTIAL(TWO_PASS, 3296, SINGLE_QUERY, true, true); \ + } else if (fixed_interleaved_layout && page_size == 1648) { \ + LAUNCH_GROUPED_VERIFY_PARTIAL(TWO_PASS, 1648, SINGLE_QUERY, true, \ + false); \ + } else if (fixed_interleaved_layout) { \ + LAUNCH_GROUPED_VERIFY_PARTIAL(TWO_PASS, 3296, SINGLE_QUERY, true, \ + false); \ + } else if (page_size == 1648) { \ + LAUNCH_GROUPED_VERIFY_PARTIAL(TWO_PASS, 1648, SINGLE_QUERY, false, \ + false); \ + } else if (page_size == 3296) { \ + LAUNCH_GROUPED_VERIFY_PARTIAL(TWO_PASS, 3296, SINGLE_QUERY, false, \ + false); \ + } else { \ + LAUNCH_GROUPED_VERIFY_PARTIAL(TWO_PASS, 0, SINGLE_QUERY, false, false); \ + } \ } while (0) const bool single_query = q.size(0) == 1; 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/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..c7fab4fdb3 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, @@ -243,3 +266,52 @@ 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]) +@torch.inference_mode() +def test_grouped_verify_fixed_interleaved_is_bitwise( + page_size: 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=8, + 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/test_envs.py b/tests/test_envs.py index ffc79b49e4..983359d14f 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -218,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/v1/spec_decode/test_dflash2.py b/tests/v1/spec_decode/test_dflash2.py index 98cda4fe8c..07d1de30ec 100644 --- a/tests/v1/spec_decode/test_dflash2.py +++ b/tests/v1/spec_decode/test_dflash2.py @@ -63,7 +63,6 @@ def test_dflash2_gdn_fastpaths_are_default_off(monkeypatch): names = ( - "VLLM_FLASH_V100_DFLASH2_GROUPED_VERIFY", "VLLM_SM70_DFLASH2_QPN8_RERANK", "VLLM_SM70_DFLASH2_VERIFY_FASTPATH", "VLLM_SM70_DFLASH2_FUSED_GDN_METADATA", @@ -87,6 +86,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() diff --git a/vllm/envs.py b/vllm/envs.py index 2d3e04aa45..fce8b2f611 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -392,8 +392,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 @@ -2630,11 +2632,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") ), From a0336cfbb1b21a586c26d548c829dc0e8464aea9 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 26 Aug 2026 18:22:45 +0800 Subject: [PATCH 024/105] [Kernel][SM70] Default stable D256 GQA prefill through 256K (#8) Extends the quality-valid max-aware D256 GQA architecture to KV256K, removes the rejected raw-logit half2 experiment, preserves explicit rollback, and retains the recorded exact output hash with 19.5% prompt-throughput gain. Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../sm70_flash_attn_d256_gqa_arch.patch | 6 +- docs/design/sm70_qwen38_fp8_prefill_decay.md | 129 ++++++++++++++++++ .../attention/test_sm70_flash_v100_policy.py | 18 ++- vllm/v1/attention/backends/flash_attn_v100.py | 6 +- 4 files changed, 149 insertions(+), 10 deletions(-) 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/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/tests/v1/attention/test_sm70_flash_v100_policy.py b/tests/v1/attention/test_sm70_flash_v100_policy.py index fb2e195796..419c4d847c 100644 --- a/tests/v1/attention/test_sm70_flash_v100_policy.py +++ b/tests/v1/attention/test_sm70_flash_v100_policy.py @@ -891,7 +891,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 +904,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 +934,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(), ) diff --git a/vllm/v1/attention/backends/flash_attn_v100.py b/vllm/v1/attention/backends/flash_attn_v100.py index e6a6f15e65..e42b3df8ae 100644 --- a/vllm/v1/attention/backends/flash_attn_v100.py +++ b/vllm/v1/attention/backends/flash_attn_v100.py @@ -1451,7 +1451,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 +1462,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 +1632,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") From d63e9490f65f9e01f6649053c1ab72922034b931 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 26 Aug 2026 18:42:44 +0800 Subject: [PATCH 025/105] [Kernel][SM70] Default exact PP2 TP4 decode optimizations (#10) Default the audited non-fused shared-gate QPN8 and exact TP4 six-route MXFP4 QPN M1 paths with explicit rollback and old-extension fallback. Preserve the recorded PP2 x TP4 endpoint, rejected-path evidence, and long-context design scope. Tests: Ruff 0.16.4; clang-format 21; git diff --check; 39 passed, 4 skipped policy tests; recorded matching-extension suite 41 passed. Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- CMakeLists.txt | 1 + csrc/ops.h | 4 + csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu | 203 ++++++++++++++++++ csrc/torch_bindings.cpp | 7 + docs/design/sm70_v100_migration_control.md | 128 +++++++++++ .../test_sm70_fp8_qpn8_pp2_tp4.py | 99 +++++++++ tests/quantization/test_sm70_mxfp4_moe.py | 56 +++++ vllm/_sm70_ops.py | 27 +++ vllm/envs.py | 13 ++ .../model_executor/layers/quantization/fp8.py | 48 ++++- .../layers/quantization/mxfp4_sm70_moe.py | 79 +++++++ 11 files changed, 659 insertions(+), 6 deletions(-) create mode 100644 csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu diff --git a/CMakeLists.txt b/CMakeLists.txt index d07300fdff..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" diff --git a/csrc/ops.h b/csrc/ops.h index e2fd08b575..67170c205d 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -485,6 +485,10 @@ 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_dense_stage_sm70_out(torch::Tensor out, torch::Tensor input, torch::Tensor expert_offsets, torch::Tensor dense_expert_ids, 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..5d65a8267f --- /dev/null +++ b/csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu @@ -0,0 +1,203 @@ +// 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 +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); +} + +} // 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(); +} diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index 38278ee057..4621db264c 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -644,6 +644,13 @@ 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_dense_stage_sm70_out(" "Tensor(a!) out, Tensor input, Tensor expert_offsets, " diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index 7bb276c394..dba85d9161 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -43750,3 +43750,131 @@ Interpretation: occurs once after loading; capture and per-token execution perform no packing. 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 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. This is the new accepted speed endpoint, but it misses + the 100-token/s target by `27.628 token/s`. 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`, not endpoint + TPOT. Full endpoint, three-repeat pure decode and GSM8K-64 remain pidfd + queued 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 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 now defaults on, + retains the external production activation, and does not inspect model or + checkpoint identity. `VLLM_SM70_FP8_QPN8_PP2_TP4_SHARED_GATE=0` is the + rollback. A matched endpoint remains useful for attribution, but the + operator result is not presented as an endpoint-speed measurement. +- 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-100tps-20260826`; public upstream is + fetch-only for this work. 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 17580656d8..9fa5f8ecc3 100644 --- a/tests/model_executor/test_sm70_fp8_qpn8_pp2_tp4.py +++ b/tests/model_executor/test_sm70_fp8_qpn8_pp2_tp4.py @@ -90,6 +90,39 @@ 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_is_default_with_rollback(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) == ( + 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( @@ -227,3 +260,69 @@ def test_pp2_tp4_qpn8_default_prepares_matching_layer(monkeypatch) -> None: 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.delenv("VLLM_SM70_FP8_QPN8_PP2_TP4_SHARED_GATE", raising=False) + monkeypatch.delenv("VLLM_SM70_FP8_QPN8", raising=False) + monkeypatch.delenv("VLLM_SM70_FP8_QPN8_PP2_TP4", 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/quantization/test_sm70_mxfp4_moe.py b/tests/quantization/test_sm70_mxfp4_moe.py index 9e3e9fa18f..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, @@ -226,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/vllm/_sm70_ops.py b/vllm/_sm70_ops.py index 5870e325e4..70490b3c37 100644 --- a/vllm/_sm70_ops.py +++ b/vllm/_sm70_ops.py @@ -1222,6 +1222,33 @@ 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_dense_stage_sm70_out( out: torch.Tensor, input: torch.Tensor, diff --git a/vllm/envs.py b/vllm/envs.py index fce8b2f611..c09e3d89c7 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -166,6 +166,7 @@ VLLM_SM70_FP8_PREFILL_EXACT_DENSE: bool = True VLLM_SM70_FP8_QPN8: bool = False VLLM_SM70_FP8_QPN8_PP2_TP4: bool = True + VLLM_SM70_FP8_QPN8_PP2_TP4_SHARED_GATE: bool = True VLLM_SM70_FP8_QPN8_LIBRARY: str | None = None VLLM_SM70_SAMPLER_LIBRARY: str | None = None VLLM_SM70_FP8_PREFILL_VISIBLE_DENSE_MM: bool = False @@ -255,6 +256,7 @@ VLLM_SM70_MXFP4_MOE_GROUPED_M8_FAST_SELECTOR: bool = True 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 @@ -1697,6 +1699,12 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_FP8_QPN8_PP2_TP4": lambda: bool( int(os.getenv("VLLM_SM70_FP8_QPN8_PP2_TP4", "1")) ), + # Default non-fused QPN8 route for the exact PP2 x TP4 shared-expert + # gate/up tensor. The model-level clamp-SwiGLU remains external. Set to 0 + # to retain the TurboMind path. + "VLLM_SM70_FP8_QPN8_PP2_TP4_SHARED_GATE": lambda: bool( + int(os.getenv("VLLM_SM70_FP8_QPN8_PP2_TP4_SHARED_GATE", "1")) + ), # 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), @@ -2222,6 +2230,11 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_MXFP4_MOE_DIRECT_ORDER_DECODE": lambda: bool( 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")) ), diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 8f8c975d85..53a871e1b9 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -133,8 +133,8 @@ } _SM70_FP8_QPN8_PP2_TP4_CONFIGS = { # Real-weight speed winners remain numerically bounded for every admitted - # projection. The shared-expert gate/up role is deliberately absent below - # after its fused activation matched only 53.5% of TurboMind FP16 elements. + # 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), @@ -142,14 +142,15 @@ (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 @@ -299,6 +300,22 @@ 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.""" + 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 _sm70_fp8_qpn8_pp2_tp4_config( layer: torch.nn.Module, *, gated_silu: bool ) -> tuple[int, int, bool] | None: @@ -306,6 +323,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 @@ -1027,6 +1050,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 " @@ -1048,6 +1078,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( @@ -1108,6 +1139,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 " diff --git a/vllm/model_executor/layers/quantization/mxfp4_sm70_moe.py b/vllm/model_executor/layers/quantization/mxfp4_sm70_moe.py index dc39f8426a..cd254bb2a5 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: @@ -197,6 +222,27 @@ def _select_mxfp4_direct_order_offsets( 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 " @@ -749,6 +796,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) From fe9953db576ffc2bc113d07a67a655a24ba7c98e Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 26 Aug 2026 18:47:35 +0800 Subject: [PATCH 026/105] [Doc][SM70] Clarify no-copy endpoint attribution Record the route-hit 73.645-token combined endpoint and passing quality gate, while identifying the FP8 route change that prevents assigning its delta to MXFP4. Continue follow-up work from the merged private main. Co-authored-by: OpenAI Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_v100_migration_control.md | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index dba85d9161..f2866524be 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -43796,9 +43796,14 @@ Interpretation: 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`, not endpoint - TPOT. Full endpoint, three-repeat pure decode and GSM8K-64 remain pidfd - queued under + 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. This is a valid + combined endpoint but 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 @@ -43811,7 +43816,9 @@ Interpretation: 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. + operator projection is relabeled as a full-model speed result. A matching + merged-source AOT build and one-variable full-model A/B remain follow-up + evidence rather than merge prerequisites. - 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 @@ -43876,5 +43883,5 @@ Interpretation: 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-100tps-20260826`; public upstream is + `agent/private-v100-dsv4-pp2tp4-followup-20260826`; public upstream is fetch-only for this work. From e7bf082f99a3484b42d8ed3765aec4f904d8b5e2 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 26 Aug 2026 19:08:12 +0800 Subject: [PATCH 027/105] [Doc][SM70] Record merged defaults AOT validation Co-authored-by: OpenAI Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_v100_migration_control.md | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index f2866524be..7a3bfd753c 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -43816,9 +43816,19 @@ Interpretation: 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. A matching - merged-source AOT build and one-variable full-model A/B remain follow-up - evidence rather than merge prerequisites. + 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 From 3fffc42058e88797c462c7d78c2515198ee56006 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 26 Aug 2026 19:20:24 +0800 Subject: [PATCH 028/105] [Bugfix][SM70] Gate shared QPN8 on quality Keep the numerically sensitive shared-expert gate/up QPN8 route opt-in after its matched GSM8K endpoint regressed one item. Co-authored-by: OpenAI Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_v100_migration_control.md | 19 ++++++++++++++----- .../test_sm70_fp8_qpn8_pp2_tp4.py | 8 ++++++-- vllm/envs.py | 10 +++++----- 3 files changed, 25 insertions(+), 12 deletions(-) diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index 7a3bfd753c..0ee7164c8d 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -43847,11 +43847,20 @@ Interpretation: 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 now defaults on, - retains the external production activation, and does not inspect model or - checkpoint identity. `VLLM_SM70_FP8_QPN8_PP2_TP4_SHARED_GATE=0` is the - rollback. A matched endpoint remains useful for attribution, but the - operator result is not presented as an endpoint-speed measurement. + 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. - 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, 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 9fa5f8ecc3..6fc93b62bd 100644 --- a/tests/model_executor/test_sm70_fp8_qpn8_pp2_tp4.py +++ b/tests/model_executor/test_sm70_fp8_qpn8_pp2_tp4.py @@ -90,7 +90,7 @@ 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_is_default_with_rollback(monkeypatch) -> None: +def test_pp2_tp4_qpn8_shared_gate_requires_explicit_opt_in(monkeypatch) -> None: layer = _layer( "gate_up_proj", 4, @@ -104,6 +104,10 @@ def test_pp2_tp4_qpn8_shared_gate_is_default_with_rollback(monkeypatch) -> None: 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, @@ -263,7 +267,7 @@ def test_pp2_tp4_qpn8_default_prepares_matching_layer(monkeypatch) -> None: def test_pp2_tp4_qpn8_shared_gate_retains_external_activation(monkeypatch) -> None: - monkeypatch.delenv("VLLM_SM70_FP8_QPN8_PP2_TP4_SHARED_GATE", raising=False) + monkeypatch.setenv("VLLM_SM70_FP8_QPN8_PP2_TP4_SHARED_GATE", "1") monkeypatch.delenv("VLLM_SM70_FP8_QPN8", raising=False) monkeypatch.delenv("VLLM_SM70_FP8_QPN8_PP2_TP4", raising=False) envs.disable_envs_cache() diff --git a/vllm/envs.py b/vllm/envs.py index c09e3d89c7..4363243d65 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -166,7 +166,7 @@ VLLM_SM70_FP8_PREFILL_EXACT_DENSE: bool = True VLLM_SM70_FP8_QPN8: bool = False VLLM_SM70_FP8_QPN8_PP2_TP4: bool = True - VLLM_SM70_FP8_QPN8_PP2_TP4_SHARED_GATE: bool = True + 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_FP8_PREFILL_VISIBLE_DENSE_MM: bool = False @@ -1699,11 +1699,11 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_FP8_QPN8_PP2_TP4": lambda: bool( int(os.getenv("VLLM_SM70_FP8_QPN8_PP2_TP4", "1")) ), - # Default non-fused QPN8 route for the exact PP2 x TP4 shared-expert - # gate/up tensor. The model-level clamp-SwiGLU remains external. Set to 0 - # to retain the TurboMind path. + # 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", "1")) + 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. From 6d82fe5666e528f326859ca79f87ce555c1e4285 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 26 Aug 2026 19:29:30 +0800 Subject: [PATCH 029/105] [Kernel][SM70] Add exact shared-gate prescaled route Keep the route opt-in until its matched full-model speed and quality gate closes. Co-authored-by: OpenAI Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_v100_migration_control.md | 10 +++ .../test_sm70_fp8_qpn8_pp2_tp4.py | 79 +++++++++++++++++++ .../test_sm70_fp8_prefill_exact_dense.py | 25 ++++++ tests/test_envs.py | 6 ++ vllm/envs.py | 6 ++ .../model_executor/layers/quantization/fp8.py | 55 +++++++++---- 6 files changed, 165 insertions(+), 16 deletions(-) diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index 0ee7164c8d..0c3f0aa908 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -43861,6 +43861,16 @@ Interpretation: 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 TurboMind gated-interleaved layout, accumulation + path, layout restore and external clamp-SwiGLU unchanged. It only folds the + exact E4M3 exponent-bias factor into finite, reversibly prescaled FP16 block + scales. On the real layer-0 TP4-rank-0 shared gate/up tensor, all 64 dynamic + patterns are bitwise both 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`). Evidence is under + `/data/models/v100-dsv4-0731-pp2tp4-shared-gate-prescaled-exact-screen-20260826-r1/`. + `VLLM_SM70_FP8_PRESCALED_M1_SHARED_GATE=1` remains an isolated opt-in until + its matched full-model speed and dataset-quality pair closes. - 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, 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 6fc93b62bd..a6aeed52de 100644 --- a/tests/model_executor/test_sm70_fp8_qpn8_pp2_tp4.py +++ b/tests/model_executor/test_sm70_fp8_qpn8_pp2_tp4.py @@ -330,3 +330,82 @@ def test_pp2_tp4_qpn8_shared_gate_retains_external_activation(monkeypatch) -> No assert layer.sm70_fp8_qpn8_nacc == 2 assert not layer.sm70_fp8_qpn8_prefetch assert not getattr(layer, "sm70_fp8_gated_silu", False) + + +def test_pp2_tp4_shared_gate_prescaled_keeps_turbomind_layout(monkeypatch) -> None: + monkeypatch.setenv("VLLM_SM70_FP8_QPN8_PP2_TP4_SHARED_GATE", "0") + monkeypatch.setenv("VLLM_SM70_FP8_PRESCALED_M1_DECODE", "1") + monkeypatch.setenv("VLLM_SM70_FP8_PRESCALED_M1_SHARED_GATE", "1") + monkeypatch.setenv("VLLM_SM70_FP8_DENSE_GATED_SILU", "1") + monkeypatch.setenv("VLLM_SM70_FP8_PREFILL_FAST_SELECTOR", "0") + monkeypatch.setenv("VLLM_SM70_FP8_PREFILL_EXACT_DENSE", "0") + envs.disable_envs_cache() + + layer = torch.nn.Module() + layer.prefix = "arbitrary.layers.7.mlp.shared_experts.gate_up_proj" + layer.tp_size = 4 + layer.input_size_per_partition = 4096 + layer.output_size_per_partition = 1024 + layer.output_partition_sizes = [512, 512] + layer.weight_block_size = [128, 128] + layer.weight = torch.empty((1024, 4096), device="meta") + layer.weight_scale_inv = torch.empty((8, 32), device="meta") + layer.orig_dtype = torch.float16 + layer.is_bmm = False + + method = fp8.Fp8LinearMethod.__new__(fp8.Fp8LinearMethod) + method.use_marlin = False + method.use_sm70_fp8_turbomind = True + method.weight_block_size = [128, 128] + method.is_scale_e8m0 = True + 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, + ) + tm_scales = torch.full((32, 1024), 2**-12, dtype=torch.float16) + prepare_calls: list[bool] = [] + + def fake_prepare(weight, scales, group_size, gated_silu): + del weight, scales + assert group_size == 128 + prepare_calls.append(gated_silu) + return ( + torch.empty((4096, 1024), dtype=torch.uint8, device="meta"), + tm_scales, + torch.tensor([4096, 1024]), + ) + + monkeypatch.setattr( + torch.ops._C, "fp8_gemm_sm70_prescaled_m1_out", object(), raising=False + ) + try: + with ( + patch.object(fp8, "get_current_vllm_config", return_value=config), + patch.object( + fp8, + "process_fp8_weight_block_strategy", + side_effect=lambda weight, scales: (weight, scales), + ), + patch.object(fp8.sm70_ops, "fp8_sm70_prepare", side_effect=fake_prepare), + patch.object(fp8, "replace_parameter"), + ): + method.process_weights_after_loading(layer) + finally: + envs.disable_envs_cache() + + assert prepare_calls == [True] + assert layer.sm70_fp8_turbomind + assert not getattr(layer, "sm70_fp8_qpn8", False) + assert layer.sm70_fp8_gated_silu_primary + torch.testing.assert_close( + layer.sm70_fp8_decode_prescaled_scales, + tm_scales * 256, + rtol=0, + atol=0, + ) diff --git a/tests/quantization/test_sm70_fp8_prefill_exact_dense.py b/tests/quantization/test_sm70_fp8_prefill_exact_dense.py index 166f871954..74cd6d3c09 100644 --- a/tests/quantization/test_sm70_fp8_prefill_exact_dense.py +++ b/tests/quantization/test_sm70_fp8_prefill_exact_dense.py @@ -698,6 +698,31 @@ def test_fp8_prescaled_m1_decode_admission_is_exact(): assert not _is_sm70_fp8_prescaled_m1_decode_layer(layer) +def test_fp8_prescaled_m1_shared_gate_requires_opt_in(monkeypatch): + layer = SimpleNamespace( + prefix="model.layers.2.mlp.shared_experts.gate_up_proj", + tp_size=4, + weight_block_size=[128, 128], + input_size_per_partition=4096, + output_size_per_partition=1024, + output_partition_sizes=[512, 512], + weight=torch.empty((1024, 4096), device="meta"), + ) + monkeypatch.delenv("VLLM_SM70_FP8_PRESCALED_M1_SHARED_GATE", raising=False) + envs.disable_envs_cache() + try: + assert not _is_sm70_fp8_prescaled_m1_decode_layer(layer) + + monkeypatch.setenv("VLLM_SM70_FP8_PRESCALED_M1_SHARED_GATE", "1") + envs.disable_envs_cache() + assert _is_sm70_fp8_prescaled_m1_decode_layer(layer) + + layer.prefix = "model.layers.2.mlp.gate_up_proj" + assert not _is_sm70_fp8_prescaled_m1_decode_layer(layer) + finally: + envs.disable_envs_cache() + + 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) diff --git a/tests/test_envs.py b/tests/test_envs.py index 983359d14f..66a7741f0a 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -153,6 +153,12 @@ def test_sm70_concurrency_tuning_envs( monkeypatch.setenv(name, "0") assert environment_variables[name]() is False + name = "VLLM_SM70_FP8_PRESCALED_M1_SHARED_GATE" + monkeypatch.delenv(name, raising=False) + assert environment_variables[name]() is False + monkeypatch.setenv(name, "1") + assert environment_variables[name]() is True + name = "VLLM_SM70_MXFP4_MOE_BROADCAST_INPUT_DECODE" monkeypatch.delenv(name, raising=False) assert environment_variables[name]() is True diff --git a/vllm/envs.py b/vllm/envs.py index 4363243d65..35e06e8eac 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -160,6 +160,7 @@ 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_PRESCALED_M1_SHARED_GATE: bool = False VLLM_SM70_FP8_PREFILL_CUTLASS: bool = True VLLM_SM70_FP8_PRESERVE_DEFAULT_SPLITS: bool = True VLLM_SM70_FP8_PRESERVE_DEFAULT_SPLITS_ONLY: bool = False @@ -1781,6 +1782,11 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_FP8_PRESCALED_M1_DECODE": lambda: bool( int(os.getenv("VLLM_SM70_FP8_PRESCALED_M1_DECODE", "1")) ), + # Exact TurboMind-prescaled M=1 route for the PP2 x TP4 shared-expert + # gate/up tensor. This remains opt-in until its full-model quality gate. + "VLLM_SM70_FP8_PRESCALED_M1_SHARED_GATE": lambda: bool( + int(os.getenv("VLLM_SM70_FP8_PRESCALED_M1_SHARED_GATE", "0")) + ), "VLLM_SM70_FP8_PRESERVE_DEFAULT_SPLITS": lambda: bool( int(os.getenv("VLLM_SM70_FP8_PRESERVE_DEFAULT_SPLITS", "1")) ), diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 53a871e1b9..77db404e2c 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -169,9 +169,23 @@ _sm70_fp8_qpn8_pp2_tp4_workspaces: dict[tuple[int, torch.dtype], torch.Tensor] = {} -def _is_sm70_fp8_prescaled_m1_decode_layer(layer: torch.nn.Module) -> bool: - """Admit only the measured replicated fused WQA/WKV tensor contract.""" +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 measured replicated or explicitly enabled TP4 tensors.""" + fused_wqa_wkv = 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] @@ -179,6 +193,13 @@ def _is_sm70_fp8_prescaled_m1_decode_layer(layer: torch.nn.Module) -> bool: and int(getattr(layer, "output_size_per_partition", 0)) == 1536 and tuple(layer.weight.shape) == (1536, 4096) ) + return bool( + fused_wqa_wkv + or ( + envs.VLLM_SM70_FP8_PRESCALED_M1_SHARED_GATE + and _is_sm70_fp8_pp2_tp4_shared_gate_layer(layer) + ) + ) def _is_sm70_fp8_prescaled_m1_decode_runtime_contract() -> bool: @@ -304,16 +325,7 @@ 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.""" - 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) - ) + return _is_sm70_fp8_pp2_tp4_shared_gate_layer(layer) def _sm70_fp8_qpn8_pp2_tp4_config( @@ -1157,8 +1169,13 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: ) prescaled_decode_requested = envs.VLLM_SM70_FP8_PRESCALED_M1_DECODE - prescaled_decode_explicit = ( + prescaled_shared_gate_layer = _is_sm70_fp8_pp2_tp4_shared_gate_layer(layer) + prescaled_decode_explicit = bool( os.getenv("VLLM_SM70_FP8_PRESCALED_M1_DECODE") == "1" + or ( + prescaled_shared_gate_layer + and os.getenv("VLLM_SM70_FP8_PRESCALED_M1_SHARED_GATE") == "1" + ) ) prescaled_decode_layer = _is_sm70_fp8_prescaled_m1_decode_layer(layer) prescaled_decode_runtime = bool( @@ -1233,9 +1250,15 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: prescaled_scales, persistent=False, ) - logger.info_once( - "Exact SM70 fused-WQA/WKV prescaled decode path enabled." - ) + if prescaled_shared_gate_layer: + logger.info_once( + "Exact SM70 shared-expert gate/up prescaled decode " + "path enabled with external activation retained." + ) + else: + 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 From 59adf8d89a37cf9d10d4884380e036992d09a96f Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 26 Aug 2026 20:48:14 +0800 Subject: [PATCH 030/105] [Bugfix][SM70] Bound MXFP4 post-load memory Release per-expert conversion temporaries before pointer table construction and build strided pointer tables directly in torch-managed storage. This avoids unchecked cudaMallocAsync failures when V100 model loading exhausts driver headroom. Co-authored-by: OpenAI Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- csrc/sm70_turbomind/ops/awq_sm70_gemm.cu | 50 ++++++++++--------- .../layers/quantization/mxfp4_sm70_moe.py | 21 +++++--- 2 files changed, 41 insertions(+), 30 deletions(-) diff --git a/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu b/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu index d555a8be7d..f53e612cf7 100644 --- a/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu +++ b/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu @@ -5826,6 +5826,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 @@ -5844,12 +5861,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 = @@ -5857,30 +5868,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}; } diff --git a/vllm/model_executor/layers/quantization/mxfp4_sm70_moe.py b/vllm/model_executor/layers/quantization/mxfp4_sm70_moe.py index cd254bb2a5..56e2cdf789 100644 --- a/vllm/model_executor/layers/quantization/mxfp4_sm70_moe.py +++ b/vllm/model_executor/layers/quantization/mxfp4_sm70_moe.py @@ -489,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, @@ -519,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, " From a3cc771526dd19b07e230f0762e463add6092d67 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 26 Aug 2026 22:04:50 +0800 Subject: [PATCH 031/105] [Model] Start Qwen3.8 Flash Next SM70 adaptation Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_qwen38_flash_next_nvfp4.md | 90 +++++++++++++++++++++ 1 file changed, 90 insertions(+) create mode 100644 docs/design/sm70_qwen38_flash_next_nvfp4.md 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..b76c30c842 --- /dev/null +++ b/docs/design/sm70_qwen38_flash_next_nvfp4.md @@ -0,0 +1,90 @@ +# Qwen3.8 Flash Next NVFP4 on SM70 + +## Status and ownership + +- Status: bring-up in progress; no route, quality, memory, or speed claim yet. +- Integration line: `private/main`. +- Base SHA: `d63e9490f65f9e01f6649053c1ab72922034b931`. +- Model: `RadixArk/Qwen3.8-Flash-Next-NVFP4` at revision + `7b719225242aacd3dbd3f9407468c2ee9a9d2594`. +- Model download: `/data/models/RadixArk/Qwen3.8-Flash-Next-NVFP4`. +- 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. +- 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. + +## 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. + +## Acceptance gates + +1. Static route: Transformers config, model registry, multimodal processor, + QSA/GDN/GR/PLE modules, and ModelOpt NVFP4 mapping load without importing an + Ampere-only backend. +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 target is at + least 100 emitted tokens/s (at most 10 ms/token) 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. From d6f3515a8c6d668d54bc41f37f8c35fb6273b9c4 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 26 Aug 2026 23:30:35 +0800 Subject: [PATCH 032/105] [Core] Guard DFlash2 selector ordering Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/v1/spec_decode/test_dflash2.py | 19 +++++++++++++++ vllm/envs.py | 7 ++++++ .../layers/vocab_parallel_embedding.py | 23 +++++++++++++++++-- 3 files changed, 47 insertions(+), 2 deletions(-) diff --git a/tests/v1/spec_decode/test_dflash2.py b/tests/v1/spec_decode/test_dflash2.py index 07d1de30ec..6fa8aa81ff 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, @@ -64,6 +65,7 @@ def test_dflash2_gdn_fastpaths_are_default_off(monkeypatch): names = ( "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", @@ -939,6 +941,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]) diff --git a/vllm/envs.py b/vllm/envs.py index c09e3d89c7..887121a739 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -198,6 +198,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 @@ -1902,6 +1903,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( diff --git a/vllm/model_executor/layers/vocab_parallel_embedding.py b/vllm/model_executor/layers/vocab_parallel_embedding.py index 9f69084c62..def76aaaf6 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) From 8165861da11bb80d81bdbd18a59c5f73a7921ffe Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 26 Aug 2026 23:30:42 +0800 Subject: [PATCH 033/105] [Bugfix] Repair structured spec decode state Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../test_dflash2_structured_output.py | 148 +++++++++++++++ .../test_reasoning_structured_output.py | 55 +++++- vllm/v1/core/sched/scheduler.py | 18 +- vllm/v1/structured_output/__init__.py | 178 +++++++++++++----- vllm/v1/structured_output/backend_xgrammar.py | 18 +- vllm/v1/structured_output/request.py | 3 + 6 files changed, 355 insertions(+), 65 deletions(-) create mode 100644 tests/v1/spec_decode/test_dflash2_structured_output.py 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/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/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/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. From fbbdcab7fc455eb7eb45b24a757d0e06b418fd0d Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 26 Aug 2026 23:40:27 +0800 Subject: [PATCH 034/105] [Bugfix][SM70] Default quality-regressing routes off Remove the failed shared-gate prescaled integration, default PP2 TP4 QPN8 and packed FP13 off, and record the strict matched quality evidence. Preserve explicit opt-in compatibility for the existing diagnostic routes. Co-authored-by: OpenAI Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_v100_migration_control.md | 86 ++++++++------ .../test_sm70_deepseek_v4_fp16_gemv.py | 8 +- .../test_sm70_fp8_qpn8_pp2_tp4.py | 109 +++++------------- .../test_sm70_fp8_prefill_exact_dense.py | 25 ---- tests/test_envs.py | 6 - vllm/_sm70_ops.py | 4 +- vllm/envs.py | 22 ++-- .../model_executor/layers/quantization/fp8.py | 34 ++---- 8 files changed, 98 insertions(+), 196 deletions(-) diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index 0c3f0aa908..1e96efdb4e 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -43523,30 +43523,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. @@ -43555,10 +43554,10 @@ 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. +- 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%` @@ -43593,11 +43592,9 @@ Interpretation: `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. That single sample remains visible, but - it is not a greedy-identity gate and does not show that choosing the lowest - operator error improves aggregate task quality. The accepted repair - therefore preserves every numerically bounded speed winner and falls back - only the demonstrated gate/up outlier. Full artifacts are under + 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/`. @@ -43743,13 +43740,14 @@ Interpretation: 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. -- Production defaults on through `VLLM_SM70_DSV4_FP13_GEMV=1`; explicit `=0` - is the rollback. Admission checks SM70, CUDA device equality, dtype, shape, - layout, output contract, and the actual weight bits. It never reads a model - name or checkpoint identity, and an incompatible tensor retains FP16. Packing - occurs once after loading; capture and per-token execution perform no packing. - The retained buffers add 549.66 MiB without pipeline parallelism, or at most - 282.34 MiB per rank for the audited PP2 partition. +- FP13 now defaults off because these operator and overlap screens do not + constitute a matched model-level dataset gate. Explicit + `VLLM_SM70_DSV4_FP13_GEMV=1` remains available for an isolated future A/B. + 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 would 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 @@ -43762,6 +43760,15 @@ Interpretation: 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 current quality-safe source-`59adf8d89a` control explicitly disables all + FP8 QPN8 routes and packed FP13 while retaining 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 @@ -43780,8 +43787,8 @@ Interpretation: 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. This is the new accepted speed endpoint, but it misses - the 100-token/s target by `27.628 token/s`. Evidence is under + `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 @@ -43799,8 +43806,8 @@ Interpretation: 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. This is a valid - combined endpoint but not an isolated no-copy A/B: relative to the + 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 @@ -43869,8 +43876,15 @@ Interpretation: replay stably. Same-GPU A/B/B/A timing moves `62.298` to `20.562 us/layer` (`3.03x`). Evidence is under `/data/models/v100-dsv4-0731-pp2tp4-shared-gate-prescaled-exact-screen-20260826-r1/`. - `VLLM_SM70_FP8_PRESCALED_M1_SHARED_GATE=1` remains an isolated opt-in until - its matched full-model speed and dataset-quality pair closes. + 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, diff --git a/tests/kernels/test_sm70_deepseek_v4_fp16_gemv.py b/tests/kernels/test_sm70_deepseek_v4_fp16_gemv.py index f3ed2c5669..ae42aad9f3 100644 --- a/tests/kernels/test_sm70_deepseek_v4_fp16_gemv.py +++ b/tests/kernels/test_sm70_deepseek_v4_fp16_gemv.py @@ -23,15 +23,15 @@ def reset_env_cache() -> Iterator[None]: envs.disable_envs_cache() -def test_sm70_dsv4_fp13_gemv_is_default_on_with_rollback(monkeypatch) -> None: +def test_sm70_dsv4_fp13_gemv_is_default_off_with_opt_in(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 + assert not envs.VLLM_SM70_DSV4_FP13_GEMV - monkeypatch.setenv("VLLM_SM70_DSV4_FP13_GEMV", "0") + monkeypatch.setenv("VLLM_SM70_DSV4_FP13_GEMV", "1") envs.disable_envs_cache() - assert not envs.VLLM_SM70_DSV4_FP13_GEMV + assert envs.VLLM_SM70_DSV4_FP13_GEMV def test_sm70_dsv4_fp13_enables_fp16_fallback(monkeypatch) -> None: 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 a6aeed52de..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,6 +53,24 @@ 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"), [ @@ -209,9 +231,9 @@ def fake_dispatch( 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 @@ -267,9 +289,9 @@ def test_pp2_tp4_qpn8_default_prepares_matching_layer(monkeypatch) -> None: 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) - monkeypatch.delenv("VLLM_SM70_FP8_QPN8_PP2_TP4", raising=False) envs.disable_envs_cache() layer = _layer( "gate_up_proj", @@ -330,82 +352,3 @@ def test_pp2_tp4_qpn8_shared_gate_retains_external_activation(monkeypatch) -> No assert layer.sm70_fp8_qpn8_nacc == 2 assert not layer.sm70_fp8_qpn8_prefetch assert not getattr(layer, "sm70_fp8_gated_silu", False) - - -def test_pp2_tp4_shared_gate_prescaled_keeps_turbomind_layout(monkeypatch) -> None: - monkeypatch.setenv("VLLM_SM70_FP8_QPN8_PP2_TP4_SHARED_GATE", "0") - monkeypatch.setenv("VLLM_SM70_FP8_PRESCALED_M1_DECODE", "1") - monkeypatch.setenv("VLLM_SM70_FP8_PRESCALED_M1_SHARED_GATE", "1") - monkeypatch.setenv("VLLM_SM70_FP8_DENSE_GATED_SILU", "1") - monkeypatch.setenv("VLLM_SM70_FP8_PREFILL_FAST_SELECTOR", "0") - monkeypatch.setenv("VLLM_SM70_FP8_PREFILL_EXACT_DENSE", "0") - envs.disable_envs_cache() - - layer = torch.nn.Module() - layer.prefix = "arbitrary.layers.7.mlp.shared_experts.gate_up_proj" - layer.tp_size = 4 - layer.input_size_per_partition = 4096 - layer.output_size_per_partition = 1024 - layer.output_partition_sizes = [512, 512] - layer.weight_block_size = [128, 128] - layer.weight = torch.empty((1024, 4096), device="meta") - layer.weight_scale_inv = torch.empty((8, 32), device="meta") - layer.orig_dtype = torch.float16 - layer.is_bmm = False - - method = fp8.Fp8LinearMethod.__new__(fp8.Fp8LinearMethod) - method.use_marlin = False - method.use_sm70_fp8_turbomind = True - method.weight_block_size = [128, 128] - method.is_scale_e8m0 = True - 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, - ) - tm_scales = torch.full((32, 1024), 2**-12, dtype=torch.float16) - prepare_calls: list[bool] = [] - - def fake_prepare(weight, scales, group_size, gated_silu): - del weight, scales - assert group_size == 128 - prepare_calls.append(gated_silu) - return ( - torch.empty((4096, 1024), dtype=torch.uint8, device="meta"), - tm_scales, - torch.tensor([4096, 1024]), - ) - - monkeypatch.setattr( - torch.ops._C, "fp8_gemm_sm70_prescaled_m1_out", object(), raising=False - ) - try: - with ( - patch.object(fp8, "get_current_vllm_config", return_value=config), - patch.object( - fp8, - "process_fp8_weight_block_strategy", - side_effect=lambda weight, scales: (weight, scales), - ), - patch.object(fp8.sm70_ops, "fp8_sm70_prepare", side_effect=fake_prepare), - patch.object(fp8, "replace_parameter"), - ): - method.process_weights_after_loading(layer) - finally: - envs.disable_envs_cache() - - assert prepare_calls == [True] - assert layer.sm70_fp8_turbomind - assert not getattr(layer, "sm70_fp8_qpn8", False) - assert layer.sm70_fp8_gated_silu_primary - torch.testing.assert_close( - layer.sm70_fp8_decode_prescaled_scales, - tm_scales * 256, - rtol=0, - atol=0, - ) diff --git a/tests/quantization/test_sm70_fp8_prefill_exact_dense.py b/tests/quantization/test_sm70_fp8_prefill_exact_dense.py index 74cd6d3c09..166f871954 100644 --- a/tests/quantization/test_sm70_fp8_prefill_exact_dense.py +++ b/tests/quantization/test_sm70_fp8_prefill_exact_dense.py @@ -698,31 +698,6 @@ def test_fp8_prescaled_m1_decode_admission_is_exact(): assert not _is_sm70_fp8_prescaled_m1_decode_layer(layer) -def test_fp8_prescaled_m1_shared_gate_requires_opt_in(monkeypatch): - layer = SimpleNamespace( - prefix="model.layers.2.mlp.shared_experts.gate_up_proj", - tp_size=4, - weight_block_size=[128, 128], - input_size_per_partition=4096, - output_size_per_partition=1024, - output_partition_sizes=[512, 512], - weight=torch.empty((1024, 4096), device="meta"), - ) - monkeypatch.delenv("VLLM_SM70_FP8_PRESCALED_M1_SHARED_GATE", raising=False) - envs.disable_envs_cache() - try: - assert not _is_sm70_fp8_prescaled_m1_decode_layer(layer) - - monkeypatch.setenv("VLLM_SM70_FP8_PRESCALED_M1_SHARED_GATE", "1") - envs.disable_envs_cache() - assert _is_sm70_fp8_prescaled_m1_decode_layer(layer) - - layer.prefix = "model.layers.2.mlp.gate_up_proj" - assert not _is_sm70_fp8_prescaled_m1_decode_layer(layer) - finally: - envs.disable_envs_cache() - - 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) diff --git a/tests/test_envs.py b/tests/test_envs.py index 66a7741f0a..983359d14f 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -153,12 +153,6 @@ def test_sm70_concurrency_tuning_envs( monkeypatch.setenv(name, "0") assert environment_variables[name]() is False - name = "VLLM_SM70_FP8_PRESCALED_M1_SHARED_GATE" - monkeypatch.delenv(name, raising=False) - assert environment_variables[name]() is False - monkeypatch.setenv(name, "1") - assert environment_variables[name]() is True - name = "VLLM_SM70_MXFP4_MOE_BROADCAST_INPUT_DECODE" monkeypatch.delenv(name, raising=False) assert environment_variables[name]() is True diff --git a/vllm/_sm70_ops.py b/vllm/_sm70_ops.py index 70490b3c37..a00b80e614 100644 --- a/vllm/_sm70_ops.py +++ b/vllm/_sm70_ops.py @@ -25,8 +25,8 @@ def _maybe_load_fp8_qpn8_library() -> None: generic_override = os.getenv("VLLM_SM70_FP8_QPN8") specific_override = os.getenv("VLLM_SM70_FP8_QPN8_PP2_TP4") 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" + if generic_enabled or specific_enabled: torch.ops.load_library(library_path) diff --git a/vllm/envs.py b/vllm/envs.py index 35e06e8eac..43a506020e 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -160,13 +160,12 @@ 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_PRESCALED_M1_SHARED_GATE: bool = False 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_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 @@ -186,7 +185,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_FP13_GEMV: bool = False 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 @@ -1693,12 +1692,12 @@ 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. + # 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 @@ -1782,11 +1781,6 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_FP8_PRESCALED_M1_DECODE": lambda: bool( int(os.getenv("VLLM_SM70_FP8_PRESCALED_M1_DECODE", "1")) ), - # Exact TurboMind-prescaled M=1 route for the PP2 x TP4 shared-expert - # gate/up tensor. This remains opt-in until its full-model quality gate. - "VLLM_SM70_FP8_PRESCALED_M1_SHARED_GATE": lambda: bool( - int(os.getenv("VLLM_SM70_FP8_PRESCALED_M1_SHARED_GATE", "0")) - ), "VLLM_SM70_FP8_PRESERVE_DEFAULT_SPLITS": lambda: bool( int(os.getenv("VLLM_SM70_FP8_PRESERVE_DEFAULT_SPLITS", "1")) ), @@ -1835,7 +1829,7 @@ def _resolve_rust_frontend_path() -> str | None: 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")) + int(os.getenv("VLLM_SM70_DSV4_FP13_GEMV", "0")) ), "VLLM_SM70_DSV4_MHC_FP32_STAGE": lambda: bool( int(os.getenv("VLLM_SM70_DSV4_MHC_FP32_STAGE", "1")) diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 77db404e2c..f438c291da 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -184,8 +184,8 @@ def _is_sm70_fp8_pp2_tp4_shared_gate_layer(layer: torch.nn.Module) -> bool: def _is_sm70_fp8_prescaled_m1_decode_layer(layer: torch.nn.Module) -> bool: - """Admit only measured replicated or explicitly enabled TP4 tensors.""" - fused_wqa_wkv = 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] @@ -193,13 +193,6 @@ def _is_sm70_fp8_prescaled_m1_decode_layer(layer: torch.nn.Module) -> bool: and int(getattr(layer, "output_size_per_partition", 0)) == 1536 and tuple(layer.weight.shape) == (1536, 4096) ) - return bool( - fused_wqa_wkv - or ( - envs.VLLM_SM70_FP8_PRESCALED_M1_SHARED_GATE - and _is_sm70_fp8_pp2_tp4_shared_gate_layer(layer) - ) - ) def _is_sm70_fp8_prescaled_m1_decode_runtime_contract() -> bool: @@ -914,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 = ( @@ -1098,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." ) @@ -1169,13 +1162,8 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: ) prescaled_decode_requested = envs.VLLM_SM70_FP8_PRESCALED_M1_DECODE - prescaled_shared_gate_layer = _is_sm70_fp8_pp2_tp4_shared_gate_layer(layer) - prescaled_decode_explicit = bool( + prescaled_decode_explicit = ( os.getenv("VLLM_SM70_FP8_PRESCALED_M1_DECODE") == "1" - or ( - prescaled_shared_gate_layer - and os.getenv("VLLM_SM70_FP8_PRESCALED_M1_SHARED_GATE") == "1" - ) ) prescaled_decode_layer = _is_sm70_fp8_prescaled_m1_decode_layer(layer) prescaled_decode_runtime = bool( @@ -1250,15 +1238,9 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: prescaled_scales, persistent=False, ) - if prescaled_shared_gate_layer: - logger.info_once( - "Exact SM70 shared-expert gate/up prescaled decode " - "path enabled with external activation retained." - ) - else: - logger.info_once( - "Exact SM70 fused-WQA/WKV prescaled decode path enabled." - ) + 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 From 9852ce8a29bbda3e4e940880e3fd9cfd96bc96ec Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 26 Aug 2026 23:45:15 +0800 Subject: [PATCH 035/105] [Bench] Add DFlash2 multi-seed quality gates Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- benchmarks/benchmark_sm70_dflash2_gsm8k.py | 159 ++++++++-- .../compare_sm70_dflash2_quality_matrix.py | 280 ++++++++++++++++++ .../test_benchmark_sm70_dflash2_gsm8k.py | 51 ++++ ...est_compare_sm70_dflash2_quality_matrix.py | 45 +++ .../test_rejection_sampler_utils.py | 31 +- 5 files changed, 527 insertions(+), 39 deletions(-) create mode 100644 benchmarks/compare_sm70_dflash2_quality_matrix.py create mode 100644 tests/benchmarks/test_benchmark_sm70_dflash2_gsm8k.py create mode 100644 tests/benchmarks/test_compare_sm70_dflash2_quality_matrix.py 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/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/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/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, From 22a25e877fff7db150543f020a46c9f74ad634f6 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 00:01:27 +0800 Subject: [PATCH 036/105] [Doc][SM70] Record shared prescale quality root Correct the earlier same-accumulation assumption: the rejected shared-gate route selects a dedicated split-K-7 prescaled kernel whose reduction order does not match the control. Co-authored-by: OpenAI Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_v100_migration_control.md | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index 1e96efdb4e..b960d02e7a 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -43868,13 +43868,17 @@ Interpretation: 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 TurboMind gated-interleaved layout, accumulation - path, layout restore and external clamp-SwiGLU unchanged. It only folds the - exact E4M3 exponent-bias factor into finite, reversibly prescaled FP16 block - scales. On the real layer-0 TP4-rank-0 shared gate/up tensor, all 64 dynamic - patterns are bitwise both 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`). Evidence is under +- 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`, From ee4ac48a479c3dbd458d5f7c09a59f39fd271d82 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 00:55:56 +0800 Subject: [PATCH 037/105] [Doc] Record DFlash2 quality audit Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_dflash2_quality_audit.md | 316 +++++++++++++++++++++ docs/design/sm70_v100_migration_control.md | 62 ++++ 2 files changed, 378 insertions(+) create mode 100644 docs/design/sm70_dflash2_quality_audit.md 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_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index dba85d9161..5e7ef24715 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -43878,3 +43878,65 @@ Interpretation: PR for this continuation use the private remote and the isolated branch `agent/private-v100-dsv4-pp2tp4-100tps-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. From a90419ba0b11b0f72c2e6ac9b4a138792c1a79db Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 01:06:40 +0800 Subject: [PATCH 038/105] [Model] Add Qwen3.8 Flash Next SM70 V2 route Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_qwen38_flash_next_nvfp4.md | 97 +- .../test_weights_mapper_stacked.py | 76 + tests/models/qwen4_exp/__init__.py | 2 + tests/models/qwen4_exp/test_config.py | 271 ++++ tests/models/qwen4_exp/test_ple.py | 478 ++++++ tests/models/qwen4_exp/test_weight_loading.py | 131 ++ .../test_sm70_modelopt_mixed_nvfp4.py | 43 +- tests/v1/core/test_qwen4_exp_kv_cache.py | 198 +++ tests/v1/worker/test_gpu_warmup_blocks.py | 23 +- tests/v1/worker/test_qwen4_exp_ngram.py | 80 + tests/v1/worker/test_qwen4_exp_v2.py | 134 ++ vllm/config/compilation.py | 2 + vllm/config/vllm.py | 19 +- vllm/model_executor/layers/linear.py | 24 + vllm/model_executor/layers/mamba/abstract.py | 5 + .../layers/mamba/mamba_utils.py | 8 +- .../layers/quantization/modelopt.py | 35 + .../layers/quantization/nvfp4_sm70_moe.py | 116 +- .../layers/vocab_parallel_embedding.py | 18 +- vllm/model_executor/model_loader/utils.py | 3 +- vllm/model_executor/models/config.py | 73 + vllm/model_executor/models/interfaces.py | 14 +- vllm/model_executor/models/registry.py | 8 + vllm/model_executor/models/utils.py | 91 +- vllm/models/qwen4_exp/__init__.py | 61 + vllm/models/qwen4_exp/common/__init__.py | 17 + .../qwen4_exp/common/hyperconnection.py | 250 ++++ vllm/models/qwen4_exp/common/ple.py | 75 + vllm/models/qwen4_exp/common/qsa_cache.py | 824 +++++++++++ vllm/models/qwen4_exp/config.py | 252 ++++ vllm/models/qwen4_exp/nvidia/__init__.py | 2 + .../qwen4_exp/nvidia/hyperconnection.py | 205 +++ vllm/models/qwen4_exp/nvidia/indexer_qsa.py | 355 +++++ .../qwen4_exp/nvidia/low_latency_gemm.py | 185 +++ vllm/models/qwen4_exp/nvidia/model.py | 1070 ++++++++++++++ vllm/models/qwen4_exp/nvidia/model_state.py | 143 ++ vllm/models/qwen4_exp/nvidia/mtp.py | 460 ++++++ vllm/models/qwen4_exp/nvidia/ops/__init__.py | 3 + vllm/models/qwen4_exp/nvidia/ops/hc.py | 487 ++++++ vllm/models/qwen4_exp/nvidia/ops/qsa.py | 1116 ++++++++++++++ .../qwen4_exp/nvidia/ops/qsa_pre_indexer.py | 516 +++++++ vllm/models/qwen4_exp/nvidia/ple_layer.py | 1310 +++++++++++++++++ vllm/models/qwen4_exp/nvidia/qsa.py | 518 +++++++ vllm/transformers_utils/config.py | 2 + vllm/transformers_utils/configs/__init__.py | 6 + vllm/transformers_utils/configs/qwen4_exp.py | 15 + vllm/v1/attention/backends/short_conv_attn.py | 529 ++++++- vllm/v1/core/kv_cache_utils.py | 371 ++++- vllm/v1/core/single_type_kv_cache_manager.py | 94 ++ vllm/v1/kv_cache_interface.py | 26 + vllm/v1/worker/gpu/attn_utils.py | 4 +- vllm/v1/worker/gpu/block_table.py | 16 +- vllm/v1/worker/gpu/model_runner.py | 27 +- .../worker/gpu/model_states/mamba_hybrid.py | 9 +- vllm/v1/worker/gpu/warmup.py | 14 +- vllm/v1/worker/gpu_model_runner.py | 162 +- 56 files changed, 10965 insertions(+), 108 deletions(-) create mode 100644 tests/model_executor/test_weights_mapper_stacked.py create mode 100644 tests/models/qwen4_exp/__init__.py create mode 100644 tests/models/qwen4_exp/test_config.py create mode 100644 tests/models/qwen4_exp/test_ple.py create mode 100644 tests/models/qwen4_exp/test_weight_loading.py create mode 100644 tests/v1/core/test_qwen4_exp_kv_cache.py create mode 100644 tests/v1/worker/test_qwen4_exp_ngram.py create mode 100644 tests/v1/worker/test_qwen4_exp_v2.py create mode 100644 vllm/models/qwen4_exp/__init__.py create mode 100644 vllm/models/qwen4_exp/common/__init__.py create mode 100644 vllm/models/qwen4_exp/common/hyperconnection.py create mode 100644 vllm/models/qwen4_exp/common/ple.py create mode 100644 vllm/models/qwen4_exp/common/qsa_cache.py create mode 100644 vllm/models/qwen4_exp/config.py create mode 100644 vllm/models/qwen4_exp/nvidia/__init__.py create mode 100644 vllm/models/qwen4_exp/nvidia/hyperconnection.py create mode 100644 vllm/models/qwen4_exp/nvidia/indexer_qsa.py create mode 100644 vllm/models/qwen4_exp/nvidia/low_latency_gemm.py create mode 100644 vllm/models/qwen4_exp/nvidia/model.py create mode 100644 vllm/models/qwen4_exp/nvidia/model_state.py create mode 100644 vllm/models/qwen4_exp/nvidia/mtp.py create mode 100644 vllm/models/qwen4_exp/nvidia/ops/__init__.py create mode 100644 vllm/models/qwen4_exp/nvidia/ops/hc.py create mode 100644 vllm/models/qwen4_exp/nvidia/ops/qsa.py create mode 100644 vllm/models/qwen4_exp/nvidia/ops/qsa_pre_indexer.py create mode 100644 vllm/models/qwen4_exp/nvidia/ple_layer.py create mode 100644 vllm/models/qwen4_exp/nvidia/qsa.py create mode 100644 vllm/transformers_utils/configs/qwen4_exp.py diff --git a/docs/design/sm70_qwen38_flash_next_nvfp4.md b/docs/design/sm70_qwen38_flash_next_nvfp4.md index b76c30c842..fcf159bddc 100644 --- a/docs/design/sm70_qwen38_flash_next_nvfp4.md +++ b/docs/design/sm70_qwen38_flash_next_nvfp4.md @@ -2,12 +2,16 @@ ## Status and ownership -- Status: bring-up in progress; no route, quality, memory, or speed claim yet. +- Status: source bring-up implemented; CPU/configuration gates pass. Full-model + load, output quality, measured memory, and speed are not claimed yet. - Integration line: `private/main`. - 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. - 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). @@ -22,6 +26,10 @@ 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, @@ -33,6 +41,12 @@ The first correctness route deliberately excludes speculative decoding. - 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: disabled for the first route. The fixed QSA ring manager does + not yet implement reusable prefix blocks. ## Architecture facts that affect the port @@ -60,11 +74,41 @@ 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. + ## Acceptance gates -1. Static route: Transformers config, model registry, multimodal processor, +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. + 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, @@ -88,3 +132,50 @@ 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 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. +- 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 and full TP4 load + remains pending. +- 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/tests/model_executor/test_weights_mapper_stacked.py b/tests/model_executor/test_weights_mapper_stacked.py new file mode 100644 index 0000000000..07f45a55dd --- /dev/null +++ b/tests/model_executor/test_weights_mapper_stacked.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.model_executor.layers.linear import MergedColumnParallelLinear +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_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), + ] 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..f7abe35682 --- /dev/null +++ b/tests/models/qwen4_exp/test_config.py @@ -0,0 +1,271 @@ +# 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_initial_sm70_route_accepts_v2_runner( + 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=False), + 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_rejects_speculative_decode( + 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"), + ) + + with pytest.raises(NotImplementedError, match="initial SM70 V2 route"): + 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_ple.py b/tests/models/qwen4_exp/test_ple.py new file mode 100644 index 0000000000..fcb081596e --- /dev/null +++ b/tests/models/qwen4_exp/test_ple.py @@ -0,0 +1,478 @@ +# 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._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 + ) + 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(), + ) + expected = torch.arange(64, dtype=torch.float32).reshape(8, 8) % 16 + layer.weight.data.copy_(expected.to(torch.float8_e4m3fn)) + + accelerator_weight = layer.get_accelerator_weight(torch.device("cuda")) + output = F.embedding( + torch.tensor([0, 7], dtype=torch.int64, device="cuda"), accelerator_weight + ) + torch.cuda.synchronize() + + assert output.dtype == torch.float8_e4m3fn + 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, +) -> 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=torch.bfloat16, + 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=torch.bfloat16)) + return layer + + +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_weight_loading.py b/tests/models/qwen4_exp/test_weight_loading.py new file mode 100644 index 0000000000..81df1c5a9c --- /dev/null +++ b/tests/models/qwen4_exp/test_weight_loading.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm.models.qwen4_exp.nvidia.model import ( + Qwen4ExpForConditionalGeneration, + Qwen4ExpModel, + _remap_qsa_cache_scale_name, +) + + +@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" + ) + + +@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_modelopt_mixed_nvfp4.py b/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py index 065b447cc9..35982e851f 100644 --- a/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py +++ b/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py @@ -57,6 +57,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 +105,16 @@ 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_nvfp4_moe_contract_rejects_shape_consistent_unvalidated_tp8(): with pytest.raises(NotImplementedError, match="tensor parallel"): validate_nvfp4_sm70_moe_contract( @@ -160,7 +177,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 @@ -190,3 +207,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/v1/core/test_qwen4_exp_kv_cache.py b/tests/v1/core/test_qwen4_exp_kv_cache.py new file mode 100644 index 0000000000..1763bb1dda --- /dev/null +++ b/tests/v1/core/test_qwen4_exp_kv_cache.py @@ -0,0 +1,198 @@ +# 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) + + +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/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_qwen4_exp_ngram.py b/tests/v1/worker/test_qwen4_exp_ngram.py new file mode 100644 index 0000000000..c3904e1b65 --- /dev/null +++ b/tests/v1/worker/test_qwen4_exp_ngram.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +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 = {} + + 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..13e9ee95bb --- /dev/null +++ b/tests/v1/worker/test_qwen4_exp_v2.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +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 + + +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.cuda.synchronize() + + assert slot_mappings[0].tolist() == [-1, -1] + assert slot_mappings[1].tolist() == [ + 12 * 262144 + 153797, + 12 * 262144 + 165757, + ] diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index 2b0833e545..3a11face66 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", diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 961b979c12..a24c2ac59b 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -66,7 +66,13 @@ logger = init_logger(__name__) -DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset({"Qwen3ForCausalLM"}) +DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset( + { + "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 +699,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 diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index 68f99ee3eb..47477033f3 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 @@ -1215,6 +1216,29 @@ 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, self) + else: + param = getattr(self, name, self) + if param is None and name == "bias": + continue + 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. 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/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/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/nvfp4_sm70_moe.py b/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py index 7eb31d9ce3..ccc73f908c 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 Qwen 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 @@ -38,13 +38,15 @@ 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 +} +_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) @triton.jit @@ -80,7 +82,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 +102,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 +113,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 +170,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 +178,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 @@ -215,13 +210,11 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: missing = [name for name in required_ops if not hasattr(torch.ops._C, name)] 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 +309,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 +328,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 +370,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 ) @@ -417,7 +415,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 +444,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 +475,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 @@ -528,13 +527,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,11 +546,11 @@ 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 + slots = num_tokens * top_k topk_ids_i32 = buffers["topk_ids"] topk_ids_i32.copy_(topk_ids, non_blocking=True) buffers["permuted_idx"].fill_(slots) @@ -557,7 +561,7 @@ def apply( layer.expert_map, layer.global_num_experts, layer.local_num_experts, - _QWEN36_TOP_K, + top_k, buffers["permuted_input"], buffers["expert_offsets64"], buffers["inv_permuted_idx"], @@ -613,7 +617,7 @@ def apply( topk_weights, buffers["inv_permuted_idx"], buffers["expert_offsets64"], - _QWEN36_TOP_K, + top_k, output, ) return output diff --git a/vllm/model_executor/layers/vocab_parallel_embedding.py b/vllm/model_executor/layers/vocab_parallel_embedding.py index 9f69084c62..7bec6af393 100644 --- a/vllm/model_executor/layers/vocab_parallel_embedding.py +++ b/vllm/model_executor/layers/vocab_parallel_embedding.py @@ -824,6 +824,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 +855,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 +1081,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 de9786c777..7674b2c539 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -568,6 +568,77 @@ 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") + + if vllm_config.cache_config.enable_prefix_caching: + raise NotImplementedError( + "Qwen4Exp prefix caching is not enabled in the initial SM70 " + "route; disable it while the QSA ring cache is in use" + ) + + 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: + raise NotImplementedError( + "Qwen4Exp speculative decoding is disabled for the initial " + "SM70 V2 route. The checkpoint's PLE n-gram embedding remains " + "enabled; disable --speculative-config until the native MTP " + "follow-up is quality-gated." + ) + + +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 SnowflakeGteNewModelConfig(VerifyAndUpdateConfig): @staticmethod def verify_and_update_model_config(model_config: "ModelConfig") -> None: @@ -628,6 +699,8 @@ def verify_and_update_model_config(model_config: "ModelConfig") -> None: "Qwen3VLForSequenceClassification": Qwen3VLForSequenceClassificationConfig, "Qwen3_5ForConditionalGeneration": Qwen3_5ForConditionalGenerationConfig, "Qwen3_5MoeForConditionalGeneration": Qwen3_5ForConditionalGenerationConfig, + "Qwen4ExpForCausalLM": Qwen4ExpForCausalLMConfig, + "Qwen4ExpForConditionalGeneration": Qwen4ExpForConditionalGenerationConfig, "VoyageQwen3BidirectionalEmbedModel": VoyageQwen3BidirectionalEmbedModelConfig, "XLMRobertaModel": JinaRobertaModelConfig, } 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/registry.py b/vllm/model_executor/models/registry.py index 5723dbe285..e053e2def8 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -117,6 +117,10 @@ "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"), @@ -562,6 +566,10 @@ "qwen3_5", "Qwen3_5MoeForConditionalGeneration", ), + "Qwen4ExpForConditionalGeneration": ( + "vllm.models.qwen4_exp", + "Qwen4ExpForConditionalGeneration", + ), "RForConditionalGeneration": ("rvl", "RForConditionalGeneration"), "SkyworkR1VChatModel": ("skyworkr1v", "SkyworkR1VChatModel"), "SmolVLMForConditionalGeneration": ("smolvlm", "SmolVLMForConditionalGeneration"), diff --git a/vllm/model_executor/models/utils.py b/vllm/model_executor/models/utils.py index 095d0e363d..302e71dbd7 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 [ @@ -354,6 +377,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/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..8ca503d214 --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/hyperconnection.py @@ -0,0 +1,205 @@ +# 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.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 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, + ) + + if self.use_combine: + # produce injection logits for combine + 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) # [M, D] + block_input = hc_gate_mix(xn, gate, self.hc_count) + + 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, + ) + + if self.use_combine: + # produce injection logits for combine + 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) # [M, D] + block_input = hc_gate_mix(xn, gate, self.hc_count) + + 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..5f32b8ee6e --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -0,0 +1,1070 @@ +# 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.compilation.decorators import support_torch_compile +from vllm.config import VllmConfig +from vllm.distributed import get_pp_group +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) + + +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": + attn_out = self.linear_attn(hidden_states=block_input) + elif self.layer_type == "full_attention": + attn_out = self.self_attn( + hidden_states=block_input, + positions=positions, + ) + 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 = ( + 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]]: + 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]]: + 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_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 = ( + 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 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( + cls, + vllm_config: VllmConfig, + ) -> tuple[tuple[int, int], tuple[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..87599d99f8 --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/mtp.py @@ -0,0 +1,460 @@ +# 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 _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 = ( + 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( + 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 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(), + ) + return loader.load_weights(remap_weight_names()) + + +__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..a3487b916e --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/ops/hc.py @@ -0,0 +1,487 @@ +# 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) + BLOCK_SIZE = 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(), + ) + 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..b1fda40d18 --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py @@ -0,0 +1,1116 @@ +# 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 torch + +from vllm.platforms import current_platform +from vllm.triton_utils import HAS_TRITON, tl, triton + +_LOGITS_WORKSPACE_BYTES = 128 * 1024 * 1024 +_TOPK_WORKSPACE_BYTES = 1024 * 1024 + + +@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 _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 + BLOCK_N = 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 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 + + columns = page_table.shape[1] * k_cache.shape[1] + block_topk = token_topk // compress_ratio + rows_per_chunk = max(1, _LOGITS_WORKSPACE_BYTES // max(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) + 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, 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] + 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 + + 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_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..b4549add56 --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py @@ -0,0 +1,1310 @@ +# 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.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__) + + +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, params_dtype + 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, + scale_dtype=torch.bfloat16, + ) + 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: + 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, + ) + + host_weight = ModelWeightParameter( + data=torch.empty( + tuple(self.weight.shape), + dtype=self.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=torch.bfloat16, + ) + self._accelerator_weight_views: dict[int, torch.Tensor] = {} + 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.cuda.current_device() 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.cuda.device(device_index): + view = get_accelerator_view_from_cpu_tensor(self.weight) + self._accelerator_weight_views[device_index] = view + return view + + def prepare_accelerator_weight(self) -> None: + self.get_accelerator_weight(torch.device("cuda", torch.cuda.current_device())) + + +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 = query_start_loc.numel() - 1 + num_tokens = input_ids.shape[0] + if num_tokens > self.positions_buffer.numel(): + raise ValueError( + f"PLE received {num_tokens} tokens, but its workspace supports " + f"at most {self.positions_buffer.numel()}" + ) + if num_reqs > self.padded_buffer.shape[0]: + raise ValueError( + f"PLE received {num_reqs} requests, but its workspace supports " + f"at most {self.padded_buffer.shape[0]}" + ) + + 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[:num_reqs].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 + + +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..0bab984eac --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/qsa.py @@ -0,0 +1,518 @@ +# 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] + key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-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, + 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) + output, _ = self.o_proj(flat_output) + return 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/transformers_utils/config.py b/vllm/transformers_utils/config.py index f940739a96..23d833ac67 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -131,6 +131,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..700d9cda73 100644 --- a/vllm/transformers_utils/configs/__init__.py +++ b/vllm/transformers_utils/configs/__init__.py @@ -75,6 +75,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", @@ -143,6 +146,9 @@ "QianfanOCRVisionConfig", "Qwen3ASRConfig", "Qwen3NextConfig", + "Qwen4ExpConfig", + "Qwen4ExpTextConfig", + "Qwen4ExpVisionConfig", "Qwen3_5Config", "Qwen3_5TextConfig", "Qwen3_5MoeConfig", 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/v1/attention/backends/short_conv_attn.py b/vllm/v1/attention/backends/short_conv_attn.py index 9c85ec5efb..3e888e69b7 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 MambaSpec class ShortConvAttentionBackend(AttentionBackend): @@ -32,3 +47,513 @@ 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: MambaSpec, + 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, 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/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 9f1f447340..256443f819 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,6 +21,7 @@ from vllm.utils.torch_utils import get_dtype_size from vllm.v1.kv_cache_interface import ( ChunkedLocalAttentionSpec, + CircularBufferSpec, FullAttentionSpec, HiddenStateCacheSpec, KVCacheConfig, @@ -880,18 +881,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 +913,8 @@ 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 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 ): @@ -1248,6 +1247,39 @@ def _get_kv_cache_config_deepseek_v4( return num_blocks, kv_cache_tensors +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 +1324,10 @@ def get_kv_cache_config_from_groups( ) for layer_name in kv_cache_groups[0].layer_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 +1506,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 +1996,8 @@ def get_kv_cache_groups( ) _annotate_eagle_groups_deepseek_v4(vllm_config, kv_cache_spec, kv_cache_groups) return kv_cache_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 +2105,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 diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index b3980c592f..69f4717792 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -14,6 +14,7 @@ ) from vllm.v1.kv_cache_interface import ( ChunkedLocalAttentionSpec, + CircularBufferSpec, CrossAttentionSpec, FullAttentionSpec, HiddenStateCacheSpec, @@ -592,6 +593,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. @@ -1326,6 +1419,7 @@ def __init__( MambaSpec: MambaManager, CrossAttentionSpec: CrossAttentionManager, SinkFullAttentionSpec: SinkFullAttentionManager, + CircularBufferSpec: CircularBufferManager, } diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 83f32dfccf..d55f5a9e41 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -538,6 +538,26 @@ 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 + + @dataclass(frozen=True, kw_only=True) class SlidingWindowMLASpec(SlidingWindowSpec): """Sliding window attention with MLA cache format.""" @@ -611,6 +631,8 @@ class MambaSpec(KVCacheSpec): 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: @@ -749,6 +771,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/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/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 47d1da107e..b85b6d8e9b 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -50,7 +50,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 @@ -395,18 +400,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 +448,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 diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 140cf97413..2c55f6ac35 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 @@ -90,7 +93,11 @@ def get_extra_attn_kwargs( } if not isinstance( attn_metadata_builder, - (Mamba2AttentionMetadataBuilder, GDNAttentionMetadataBuilder), + ( + Mamba2AttentionMetadataBuilder, + GDNAttentionMetadataBuilder, + PleShortConvAttentionMetadataBuilder, + ), ): return {} kwargs = { 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..53efe8b579 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, @@ -1380,6 +1381,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 @@ -1713,6 +1727,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 +5108,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 @@ -6593,10 +6628,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 +6783,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 +8667,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 = ( @@ -10832,6 +10963,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: @@ -12049,9 +12190,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 From 6b23eea07b01f7ebccf93991503ae33d706a9cee Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 01:10:37 +0800 Subject: [PATCH 039/105] [Doc] Record verified ModelScope snapshot Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_qwen38_flash_next_nvfp4.md | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/docs/design/sm70_qwen38_flash_next_nvfp4.md b/docs/design/sm70_qwen38_flash_next_nvfp4.md index fcf159bddc..54191ae5c1 100644 --- a/docs/design/sm70_qwen38_flash_next_nvfp4.md +++ b/docs/design/sm70_qwen38_flash_next_nvfp4.md @@ -2,8 +2,9 @@ ## Status and ownership -- Status: source bring-up implemented; CPU/configuration gates pass. Full-model - load, output quality, measured memory, and speed are not claimed yet. +- Status: source bring-up implemented; CPU/configuration gates pass and the + local ModelScope snapshot is fully verified. Full-model load, output quality, + measured memory, and speed are not claimed yet. - Integration line: `private/main`. - Base SHA: `d63e9490f65f9e01f6649053c1ab72922034b931`. - Model: `RadixArk/Qwen3.8-Flash-Next-NVFP4` at revision @@ -11,7 +12,9 @@ - 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. + 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). @@ -135,6 +138,10 @@ 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. From faa1269df9836705d2b30cf959b031479f3cc6f3 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 01:46:36 +0800 Subject: [PATCH 040/105] [Bugfix][SM70] Restore D256 prefill sidecar loading Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- benchmarks/benchmark_sm70_decode.py | 41 +++++++--- docs/design/sm70_dflash2_prefill_closure.md | 79 +++++++++++++++++++ .../benchmarks/test_benchmark_sm70_decode.py | 27 +++++++ .../attention/test_sm70_flash_v100_policy.py | 29 +++++++ vllm/v1/attention/backends/flash_attn_v100.py | 28 +++++-- 5 files changed, 187 insertions(+), 17 deletions(-) create mode 100644 docs/design/sm70_dflash2_prefill_closure.md diff --git a/benchmarks/benchmark_sm70_decode.py b/benchmarks/benchmark_sm70_decode.py index bf6d364c19..29377def49 100644 --- a/benchmarks/benchmark_sm70_decode.py +++ b/benchmarks/benchmark_sm70_decode.py @@ -52,20 +52,37 @@ 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) + 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 import_exc: + library_path = os.getenv("VLLM_SM70_FA2_D256_LIBRARY") + if library_path is None: + return { + "available": False, + "error": f"{type(import_exc).__name__}: {import_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}, + } + 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 + 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 diff --git a/docs/design/sm70_dflash2_prefill_closure.md b/docs/design/sm70_dflash2_prefill_closure.md new file mode 100644 index 0000000000..7b2dd7460c --- /dev/null +++ b/docs/design/sm70_dflash2_prefill_closure.md @@ -0,0 +1,79 @@ +# 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 | +| Private PR #8/#13 | Qwen3.8-27B-FP8, TP4, input 261888, chunk 8192 with 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. 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. + +## 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. 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, to test the already validated Q8000 + architecture with DFlash2 and NVFP4; +3. only if chunk 4096 remains materially slower, profile a quality-exact Q4K + architecture generalization instead of weakening attention math. + +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. + diff --git a/tests/benchmarks/test_benchmark_sm70_decode.py b/tests/benchmarks/test_benchmark_sm70_decode.py index b9fe6ebb78..fa95dd5280 100644 --- a/tests/benchmarks/test_benchmark_sm70_decode.py +++ b/tests/benchmarks/test_benchmark_sm70_decode.py @@ -46,3 +46,30 @@ 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 fail_import(_name): + raise ImportError("source overlay shadows extension") + + 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", fail_import) + 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"] diff --git a/tests/v1/attention/test_sm70_flash_v100_policy.py b/tests/v1/attention/test_sm70_flash_v100_policy.py index 419c4d847c..99be3305e1 100644 --- a/tests/v1/attention/test_sm70_flash_v100_policy.py +++ b/tests/v1/attention/test_sm70_flash_v100_policy.py @@ -738,6 +738,35 @@ 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 + + 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.setitem(sys.modules, "vllm.vllm_flash_attn", None) + 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 diff --git a/vllm/v1/attention/backends/flash_attn_v100.py b/vllm/v1/attention/backends/flash_attn_v100.py index e42b3df8ae..2c2365badc 100644 --- a/vllm/v1/attention/backends/flash_attn_v100.py +++ b/vllm/v1/attention/backends/flash_attn_v100.py @@ -1274,8 +1274,21 @@ 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 + try: + # Importing the interface loads the bundled FA2 torch library. + from vllm.vllm_flash_attn import flash_attn_interface # noqa: F401 + except ImportError: + # Source-overlay deployments can intentionally keep native + # extensions outside the checkout. Match the existing SM70 + # sidecar convention and load only an explicitly selected binary. + library_path = os.getenv("VLLM_SM70_FA2_D256_LIBRARY") + if library_path is None: + raise + 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 +1298,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 +1321,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, From d944992cd5c8818743f8ac92a188e39da18cfcb5 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 01:48:25 +0800 Subject: [PATCH 041/105] [Core] Register SM70 D256 sidecar environment Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- vllm/envs.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/vllm/envs.py b/vllm/envs.py index 887121a739..5e1f0a953c 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -169,6 +169,7 @@ VLLM_SM70_FP8_QPN8_PP2_TP4_SHARED_GATE: bool = True 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_MXFP4_TUNE_SMALL_SHAPES: bool = True @@ -1710,6 +1711,7 @@ def _resolve_rust_frontend_path() -> str | None: # 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")) ), From 3fc247d1e9fedbef0743e07d6c64df898bbb9f1f Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 01:49:09 +0800 Subject: [PATCH 042/105] [Bugfix] Attach Qwen4Exp expert mapping Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/models/qwen4_exp/test_weight_loading.py | 54 +++++++++++++++++++ vllm/models/qwen4_exp/nvidia/model.py | 14 +++++ 2 files changed, 68 insertions(+) diff --git a/tests/models/qwen4_exp/test_weight_loading.py b/tests/models/qwen4_exp/test_weight_loading.py index 81df1c5a9c..fb5d9903c7 100644 --- a/tests/models/qwen4_exp/test_weight_loading.py +++ b/tests/models/qwen4_exp/test_weight_loading.py @@ -2,10 +2,13 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest +from torch import nn +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, ) @@ -99,6 +102,57 @@ def test_outer_checkpoint_mapper_selects_language_model_only_paths() -> None: ) +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_", "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"), + ] + + @pytest.mark.parametrize( ("checkpoint_name", "model_name"), [ diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py index 5f32b8ee6e..08520d3be4 100644 --- a/vllm/models/qwen4_exp/nvidia/model.py +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -11,6 +11,9 @@ 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, @@ -194,6 +197,17 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: 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, + ) class Qwen4ExpDecoderLayer(nn.Module): From 657c5d391dcfc339713a2c88e1eed1e44ed44770 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 01:58:43 +0800 Subject: [PATCH 043/105] [Bench] Require real D256 operator capability Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- benchmarks/benchmark_sm70_decode.py | 27 +++++++++++-------- docs/design/sm70_dflash2_prefill_closure.md | 16 +++++++++++ .../benchmarks/test_benchmark_sm70_decode.py | 9 ++++--- 3 files changed, 37 insertions(+), 15 deletions(-) diff --git a/benchmarks/benchmark_sm70_decode.py b/benchmarks/benchmark_sm70_decode.py index 29377def49..f10eac8d3d 100644 --- a/benchmarks/benchmark_sm70_decode.py +++ b/benchmarks/benchmark_sm70_decode.py @@ -54,20 +54,16 @@ def _sm70_fa2_d256_prefill_status(torch: Any) -> dict[str, Any]: optional_ops = ("sm70_d256_splitd_n32_dense_splitkv3_fwd",) 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 import_exc: - library_path = os.getenv("VLLM_SM70_FA2_D256_LIBRARY") - if library_path is None: - return { - "available": False, - "error": f"{type(import_exc).__name__}: {import_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}, - } + 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: @@ -80,6 +76,15 @@ def _sm70_fa2_d256_prefill_status(torch: Any) -> dict[str, Any]: "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") diff --git a/docs/design/sm70_dflash2_prefill_closure.md b/docs/design/sm70_dflash2_prefill_closure.md index 7b2dd7460c..db345122af 100644 --- a/docs/design/sm70_dflash2_prefill_closure.md +++ b/docs/design/sm70_dflash2_prefill_closure.md @@ -77,3 +77,19 @@ 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/`. diff --git a/tests/benchmarks/test_benchmark_sm70_decode.py b/tests/benchmarks/test_benchmark_sm70_decode.py index fa95dd5280..929630db88 100644 --- a/tests/benchmarks/test_benchmark_sm70_decode.py +++ b/tests/benchmarks/test_benchmark_sm70_decode.py @@ -52,9 +52,6 @@ def test_sm70_fa2_d256_prefill_status_accepts_explicit_sidecar(monkeypatch): namespace = types.SimpleNamespace() loaded: list[str] = [] - def fail_import(_name): - raise ImportError("source overlay shadows extension") - def load_library(path: str) -> None: loaded.append(path) namespace.sm70_d256_splitd_n32_dense_fwd = object() @@ -65,7 +62,11 @@ def load_library(path: str) -> None: load_library=load_library, ) fake_torch = types.SimpleNamespace(ops=fake_ops) - monkeypatch.setattr(benchmark_sm70_decode.importlib, "import_module", fail_import) + 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) From 09f68f836116c57bdd19f62c74b38ade6899a226 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 01:59:55 +0800 Subject: [PATCH 044/105] [Core] Report checkpoint name on weight-load failure Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- vllm/model_executor/models/utils.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/models/utils.py b/vllm/model_executor/models/utils.py index 302e71dbd7..561bd419a4 100644 --- a/vllm/model_executor/models/utils.py +++ b/vllm/model_executor/models/utils.py @@ -243,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) From be6c6373516b152de77b3141874b5f00d70c1b6c Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 02:08:03 +0800 Subject: [PATCH 045/105] [Bugfix][SM70] Validate D256 sidecar capability Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_dflash2_prefill_closure.md | 28 ++++++++++----- .../attention/test_sm70_flash_v100_policy.py | 11 +++++- vllm/v1/attention/backends/flash_attn_v100.py | 34 ++++++++++++------- 3 files changed, 52 insertions(+), 21 deletions(-) diff --git a/docs/design/sm70_dflash2_prefill_closure.md b/docs/design/sm70_dflash2_prefill_closure.md index db345122af..ac8267af4f 100644 --- a/docs/design/sm70_dflash2_prefill_closure.md +++ b/docs/design/sm70_dflash2_prefill_closure.md @@ -20,12 +20,14 @@ The historical short- and long-prefill numbers use different contracts: |---|---|---:|---| | 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 | -| Private PR #8/#13 | Qwen3.8-27B-FP8, TP4, input 261888, chunk 8192 with Q8000 aligned chunks, target-only | 2438.89 prompt tok/s | Stable max-aware D256 architecture | +| 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. The removed -raw-logit half2 polynomial must not be restored. +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 @@ -58,18 +60,28 @@ 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. The next paired measurements -are therefore deliberately separated: +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. Reaching Q8000 with +the existing kernel requires the historical FP16 Mamba/SSM cache override; +that arm must pass the quality audit before promotion. 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, to test the already validated Q8000 - architecture with DFlash2 and NVFP4; -3. only if chunk 4096 remains materially slower, profile a quality-exact Q4K +2. chunk 8192 plus the same sidecar and FP32 SSM state, to measure Q6592; +3. chunk 8192 plus FP16 Mamba/SSM cache, to test the already validated Q8000 + architecture with DFlash2 and NVFP4, initially as a quality-gated candidate; +4. if the FP16 state fails quality, profile a quality-exact Q6592/Q8240 architecture generalization instead of weakening attention math. Each candidate must prove operator-route hits, preserve output validity, fit diff --git a/tests/v1/attention/test_sm70_flash_v100_policy.py b/tests/v1/attention/test_sm70_flash_v100_policy.py index 99be3305e1..74bca1faef 100644 --- a/tests/v1/attention/test_sm70_flash_v100_policy.py +++ b/tests/v1/attention/test_sm70_flash_v100_policy.py @@ -741,6 +741,16 @@ def test_sm70_splitd_d256_loader_requires_exact_ops(monkeypatch): 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] = [] @@ -754,7 +764,6 @@ def load_library(path: str) -> None: load_library=load_library, ) monkeypatch.setattr(flash_v100, "torch", SimpleNamespace(ops=fake_ops)) - monkeypatch.setitem(sys.modules, "vllm.vllm_flash_attn", None) 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) diff --git a/vllm/v1/attention/backends/flash_attn_v100.py b/vllm/v1/attention/backends/flash_attn_v100.py index 2c2365badc..8baa45dda6 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 @@ -1274,21 +1275,30 @@ def _get_sm70_splitd_d256_ops(): _sm70_splitd_d256_ops_checked = True try: - try: + 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 - except ImportError: - # Source-overlay deployments can intentionally keep native - # extensions outside the checkout. Match the existing SM70 - # sidecar convention and load only an explicitly selected binary. + + 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 None: - raise - torch.ops.load_library(library_path) - logger.info( - "Loaded external SM70 D256 prefill library from %s.", - library_path, - ) + 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 From 86b19124e53b8b45b745334075d8e62872b5fab1 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 02:08:17 +0800 Subject: [PATCH 046/105] [Bugfix] Preserve QKV shard metadata in auto loading Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../test_weights_mapper_stacked.py | 42 ++++++++++++++++++- vllm/model_executor/layers/linear.py | 24 +++++++++++ 2 files changed, 65 insertions(+), 1 deletion(-) diff --git a/tests/model_executor/test_weights_mapper_stacked.py b/tests/model_executor/test_weights_mapper_stacked.py index 07f45a55dd..2cbe9b13fd 100644 --- a/tests/model_executor/test_weights_mapper_stacked.py +++ b/tests/model_executor/test_weights_mapper_stacked.py @@ -3,7 +3,10 @@ import torch -from vllm.model_executor.layers.linear import MergedColumnParallelLinear +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, +) from vllm.model_executor.models.utils import WeightsMapper @@ -74,3 +77,40 @@ def weight_loader(param, loaded_weight, shard_id) -> None: (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/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index 47477033f3..1d1f726b45 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -1337,6 +1337,30 @@ 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, self) + else: + param = getattr(self, name, self) + if param is None and name == "bias": + continue + 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, From 78225e67fa22003f6864e61528bda1935734d36c Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 02:16:52 +0800 Subject: [PATCH 047/105] [Model] Pass 1Cat GDN output buffer in Qwen4Exp Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/models/qwen4_exp/test_model_forward.py | 63 ++++++++++++++++++++ vllm/models/qwen4_exp/nvidia/model.py | 22 ++++++- 2 files changed, 84 insertions(+), 1 deletion(-) create mode 100644 tests/models/qwen4_exp/test_model_forward.py 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/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py index 08520d3be4..7eddc1bb35 100644 --- a/vllm/models/qwen4_exp/nvidia/model.py +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -8,6 +8,7 @@ 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 @@ -349,7 +350,26 @@ def forward( hidden_states, block_input, injection = attn_hc.mix(hidden_states) if self.layer_type == "linear_attention": - attn_out = self.linear_attn(hidden_states=block_input) + 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, From d174f3ff0ed1df9d9dedeeb9bb073d47b5f8cffe Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 02:25:34 +0800 Subject: [PATCH 048/105] [Model] Honor custom V2 model state providers Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/v1/worker/test_model_state_init.py | 41 +++++++++++++++++++++ vllm/v1/worker/gpu/model_states/__init__.py | 5 +++ 2 files changed, 46 insertions(+) create mode 100644 tests/v1/worker/test_model_state_init.py 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..7535a70b6c --- /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 = {} + + 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/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 From 7fb606850bae15770716c81fe663b091bb67a180 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 02:34:12 +0800 Subject: [PATCH 049/105] [Bugfix] Zero heterogeneous hybrid KV pages Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/v1/worker/test_kv_block_zeroer.py | 47 ++++++++++ vllm/v1/worker/utils.py | 112 +++++++++++++----------- 2 files changed, 107 insertions(+), 52 deletions(-) diff --git a/tests/v1/worker/test_kv_block_zeroer.py b/tests/v1/worker/test_kv_block_zeroer.py index 332111b7f4..4e31e11bde 100644 --- a/tests/v1/worker/test_kv_block_zeroer.py +++ b/tests/v1/worker/test_kv_block_zeroer.py @@ -159,3 +159,50 @@ 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} diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index 6b634557a5..8f011bc02a 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -16,7 +16,6 @@ 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.mem_utils import MemorySnapshot, format_gib from vllm.v1.attention.backend import ( AttentionBackend, @@ -75,14 +74,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 +95,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 +129,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 +160,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 +193,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 +209,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 +217,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 +263,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 +273,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 +296,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, ) From dc047cdf32ee25cf65aec177625540cb5819dfae Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 02:39:07 +0800 Subject: [PATCH 050/105] [Kernel][SM70] Stage QPN2-packed NVFP4 prefill Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- benchmarks/benchmark_sm70_nvfp4_gemm_micro.py | 80 +++++++++++++++++-- docs/design/sm70_dflash2_prefill_closure.md | 59 ++++++++++++-- tests/quantization/test_sm70_nvfp4_qpn2.py | 38 +++++++++ vllm/envs.py | 11 +++ .../schemes/compressed_tensors_w4a4_nvfp4.py | 59 ++++++++++++++ 5 files changed, 231 insertions(+), 16 deletions(-) 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/docs/design/sm70_dflash2_prefill_closure.md b/docs/design/sm70_dflash2_prefill_closure.md index ac8267af4f..c2b844db61 100644 --- a/docs/design/sm70_dflash2_prefill_closure.md +++ b/docs/design/sm70_dflash2_prefill_closure.md @@ -72,17 +72,19 @@ 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. Reaching Q8000 with -the existing kernel requires the historical FP16 Mamba/SSM cache override; -that arm must pass the quality audit before promotion. The next paired -measurements are therefore deliberately separated: +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 test the already validated Q8000 - architecture with DFlash2 and NVFP4, initially as a quality-gated candidate; -4. if the FP16 state fails quality, profile a quality-exact Q6592/Q8240 - architecture generalization instead of weakening attention math. +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 @@ -105,3 +107,44 @@ 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. 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. diff --git a/tests/quantization/test_sm70_nvfp4_qpn2.py b/tests/quantization/test_sm70_nvfp4_qpn2.py index 84f8174158..13ebd6fd49 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,24 @@ 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.weight.numel() == 0 assert layer.weight_scale.numel() == 0 @@ -139,5 +168,14 @@ 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) finally: envs.disable_envs_cache() diff --git a/vllm/envs.py b/vllm/envs.py index 5e1f0a953c..f95b2ec292 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -172,6 +172,8 @@ 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 @@ -1721,6 +1723,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. 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..48b35f730e 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,21 @@ 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() + ) 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 +442,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_codes, + layer.sm70_nvfp4_qpn2_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) From 702722bab51731812e561011cadc139f8b8148d1 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 03:11:46 +0800 Subject: [PATCH 051/105] [Bugfix][SM70] Bridge QPN2 prefill layouts Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_dflash2_prefill_closure.md | 27 ++++++++++++++++--- tests/quantization/test_sm70_nvfp4_qpn2.py | 12 +++++++++ .../schemes/compressed_tensors_w4a4_nvfp4.py | 11 ++++++-- 3 files changed, 45 insertions(+), 5 deletions(-) diff --git a/docs/design/sm70_dflash2_prefill_closure.md b/docs/design/sm70_dflash2_prefill_closure.md index c2b844db61..1621b4c077 100644 --- a/docs/design/sm70_dflash2_prefill_closure.md +++ b/docs/design/sm70_dflash2_prefill_closure.md @@ -20,6 +20,7 @@ The historical short- and long-prefill numbers use different contracts: |---|---|---:|---| | 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 | @@ -145,6 +146,26 @@ weight dequantization on every call: 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. 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. +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/tests/quantization/test_sm70_nvfp4_qpn2.py b/tests/quantization/test_sm70_nvfp4_qpn2.py index 13ebd6fd49..c9084e7cd6 100644 --- a/tests/quantization/test_sm70_nvfp4_qpn2.py +++ b/tests/quantization/test_sm70_nvfp4_qpn2.py @@ -156,6 +156,16 @@ def fake_prefill(*args): 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 @@ -177,5 +187,7 @@ def fake_prefill(*args): 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/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 48b35f730e..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 @@ -362,6 +362,13 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: 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." @@ -453,8 +460,8 @@ def _apply_qpn2( out_2d, prefill_dense_weight_ptr, x_2d, - layer.sm70_nvfp4_qpn2_codes, - layer.sm70_nvfp4_qpn2_scales, + layer.sm70_nvfp4_qpn2_prefill_codes, + layer.sm70_nvfp4_qpn2_prefill_scales, float(layer.sm70_nvfp4_qpn2_global_scale), True, gated_silu, From f73cedb3ec20b864049dc2e411851fa1ef049b5e Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 05:03:56 +0800 Subject: [PATCH 052/105] [Bugfix] Fall back for unsupported custom AR dtypes Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../test_custom_all_reduce_dispatch.py | 35 +++++++++++++++++++ .../device_communicators/custom_all_reduce.py | 2 ++ 2 files changed, 37 insertions(+) create mode 100644 tests/distributed/test_custom_all_reduce_dispatch.py 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/vllm/distributed/device_communicators/custom_all_reduce.py b/vllm/distributed/device_communicators/custom_all_reduce.py index 8d55ceeaff..790922fd6f 100644 --- a/vllm/distributed/device_communicators/custom_all_reduce.py +++ b/vllm/distributed/device_communicators/custom_all_reduce.py @@ -380,6 +380,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: From 49eb4707d8533aa4249d49dbb6d9716bf0a7258d Mon Sep 17 00:00:00 2001 From: Leonccaa <166551845+Leonccaa@users.noreply.github.com> Date: Wed, 26 Aug 2026 13:38:59 -0700 Subject: [PATCH 053/105] perf(nvfp4): use Triton weight dequant on SM70-SM88 Decode ModelOpt E4M3FN scale bytes with integer and float32 Triton operations when native FP8 scale decoding is unavailable. Preserve signed zero in the manual branch so all finite scale encodings bit-match the reference across 2D/3D and FP16/BF16 outputs. Assisted-by: OpenAI Codex Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com> --- .../quantization/test_nvfp4_emulation.py | 68 +++++++++++++++++++ .../utils/nvfp4_emulation_utils.py | 68 +++++++++++++++++-- 2 files changed, 132 insertions(+), 4 deletions(-) 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/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 ) From d64150fb9bd60eaa3bbac0db074e06652c0de7bc Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 05:18:08 +0800 Subject: [PATCH 054/105] [Core] Invoke custom KV cache bind hooks Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/models/qwen4_exp/test_qsa_cache.py | 38 +++++++++++++++++++ .../layers/attention_layer_base.py | 6 +++ vllm/v1/worker/utils.py | 2 +- 3 files changed, 45 insertions(+), 1 deletion(-) create mode 100644 tests/models/qwen4_exp/test_qsa_cache.py 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..494e9b5ae7 --- /dev/null +++ b/tests/models/qwen4_exp/test_qsa_cache.py @@ -0,0 +1,38 @@ +# 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 QSAKeyStateCache +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 = {} + 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 = [] + + 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] 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/v1/worker/utils.py b/vllm/v1/worker/utils.py index 8f011bc02a..a9e1bd2501 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -609,7 +609,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( From b37cd74c94f471d15d146ea3d8bd8ac316e1891b Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 05:48:20 +0800 Subject: [PATCH 055/105] [Bugfix] Match QSA to local Flash KV layout Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/models/qwen4_exp/test_qsa_cache.py | 59 ++++++++++++++++++++++++ vllm/models/qwen4_exp/nvidia/qsa.py | 4 +- 2 files changed, 62 insertions(+), 1 deletion(-) diff --git a/tests/models/qwen4_exp/test_qsa_cache.py b/tests/models/qwen4_exp/test_qsa_cache.py index 494e9b5ae7..9a5ed1daca 100644 --- a/tests/models/qwen4_exp/test_qsa_cache.py +++ b/tests/models/qwen4_exp/test_qsa_cache.py @@ -6,6 +6,7 @@ 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 @@ -36,3 +37,61 @@ def test_bind_qsa_key_cache_builds_key_and_mrope_views() -> None: 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/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py index 0bab984eac..b6b68b45cb 100644 --- a/vllm/models/qwen4_exp/nvidia/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/qsa.py @@ -147,7 +147,9 @@ def forward_qsa( 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] - key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) + # 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 ( From 7349fe259df3a0eefc274626a75e02abc4b34ed1 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 06:20:56 +0800 Subject: [PATCH 056/105] [Bugfix] Keep PLE scales in model dtype Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/models/qwen4_exp/test_ple.py | 12 ++++++++++-- vllm/models/qwen4_exp/nvidia/ple_layer.py | 9 ++++++--- 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index fcb081596e..14d8206fe7 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -58,6 +58,7 @@ def test_pinned_host_ple_allocates_tp_shard_without_device_table( 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 == {} @@ -286,6 +287,7 @@ def test_ngram_embedding_loads_fp8_shards_and_global_scale() -> None: 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( @@ -304,16 +306,22 @@ def _make_fp8_embedding_layer( layer = embedding_module.VocabParallelEmbedding( 3, 2, - params_dtype=torch.bfloat16, + 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=torch.bfloat16)) + 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])) diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py index b4549add56..4f38d0dc61 100644 --- a/vllm/models/qwen4_exp/nvidia/ple_layer.py +++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py @@ -162,7 +162,7 @@ def create_weights( params_dtype: torch.dtype, **extra_weight_attrs, ) -> None: - del input_size, output_size, params_dtype + 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 @@ -175,7 +175,10 @@ def create_weights( input_size_per_partition, None, weight_loader, - scale_dtype=torch.bfloat16, + # 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) @@ -293,7 +296,7 @@ def __init__( self.embedding_dim, None, self.weight_loader, - scale_dtype=torch.bfloat16, + scale_dtype=params_dtype, ) self._accelerator_weight_views: dict[int, torch.Tensor] = {} logger.info( From c228412e47aff9862841ca152f79c751fed8b781 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 06:53:46 +0800 Subject: [PATCH 057/105] [Kernel] Decode offloaded PLE FP8 rows on SM70 Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/models/qwen4_exp/test_ple.py | 23 ++-- vllm/models/qwen4_exp/nvidia/ple_layer.py | 125 ++++++++++++++++++++++ 2 files changed, 141 insertions(+), 7 deletions(-) diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index 14d8206fe7..102d1045da 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -78,6 +78,9 @@ def test_pinned_host_ple_fp8_rows_are_gatherable_on_sm70( 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, @@ -86,16 +89,22 @@ def test_pinned_host_ple_fp8_rows_are_gatherable_on_sm70( prefix="model.layers.2.ple.ngram_embedding", quant_method=Qwen4ExpPLEFp8EmbeddingMethod(), ) - expected = torch.arange(64, dtype=torch.float32).reshape(8, 8) % 16 - layer.weight.data.copy_(expected.to(torch.float8_e4m3fn)) - - accelerator_weight = layer.get_accelerator_weight(torch.device("cuda")) - output = F.embedding( - torch.tensor([0, 7], dtype=torch.int64, device="cuda"), accelerator_weight + 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.cuda.synchronize() - assert output.dtype == torch.float8_e4m3fn + assert output.dtype == torch.float16 + expected = raw.view(torch.float8_e4m3fn).float() * 0.25 torch.testing.assert_close(output.float().cpu(), expected[[0, 7]]) diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py index 4f38d0dc61..0f071ba965 100644 --- a/vllm/models/qwen4_exp/nvidia/ple_layer.py +++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py @@ -44,6 +44,7 @@ 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 ( @@ -68,6 +69,59 @@ 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 @@ -191,6 +245,9 @@ def apply( 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) @@ -299,6 +356,8 @@ def __init__( 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()), @@ -321,11 +380,38 @@ def get_accelerator_weight(self, device: torch.device) -> torch.Tensor: with torch.cuda.device(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.cuda.current_device())) + def embedding_lookup(self, input_: torch.Tensor) -> torch.Tensor: + """Gather FP8 UVA rows and emit scaled model-dtype values.""" + + device_index = ( + torch.cuda.current_device() + 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__( @@ -1298,6 +1384,45 @@ def qwen4_exp_ple_short_conv_fake( 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, From de10e2ea4c2f667354c5559c60e9d3cf5ca2508e Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 07:02:15 +0800 Subject: [PATCH 058/105] [Bugfix] Preserve dynamic PLE request shapes Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- vllm/models/qwen4_exp/nvidia/ple_layer.py | 19 +++++-------------- 1 file changed, 5 insertions(+), 14 deletions(-) diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py index 0f071ba965..1ccad20a4b 100644 --- a/vllm/models/qwen4_exp/nvidia/ple_layer.py +++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py @@ -571,18 +571,11 @@ def forward( ) -> torch.Tensor: input_ids = input_ids.reshape(-1).long() query_start_loc = query_start_loc.long() - num_reqs = query_start_loc.numel() - 1 + num_reqs = ngram_context.shape[0] num_tokens = input_ids.shape[0] - if num_tokens > self.positions_buffer.numel(): - raise ValueError( - f"PLE received {num_tokens} tokens, but its workspace supports " - f"at most {self.positions_buffer.numel()}" - ) - if num_reqs > self.padded_buffer.shape[0]: - raise ValueError( - f"PLE received {num_reqs} requests, but its workspace supports " - f"at most {self.padded_buffer.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] @@ -593,9 +586,7 @@ def forward( 0, packed.shape[1] - 1 ) packed[request_indices, columns] = input_ids - ngram_context = ngram_context[:num_reqs].to( - device=input_ids.device, dtype=torch.long - ) + 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( From 383ff458dbc8af1dda79a1a29bf700019030f054 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 08:03:49 +0800 Subject: [PATCH 059/105] [Kernel][SM70] Bypass single-token NVFP4 MoE sorting Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../test_sm70_modelopt_mixed_nvfp4.py | 52 ++++++ .../layers/quantization/nvfp4_sm70_moe.py | 175 ++++++++++++++---- 2 files changed, 193 insertions(+), 34 deletions(-) diff --git a/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py b/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py index 35982e851f..469e8aadb9 100644 --- a/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py +++ b/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py @@ -19,6 +19,8 @@ from vllm.model_executor.layers.quantization.nvfp4_sm70_moe import ( ModelOptNvFp4SM70MoEMethod, _prepare_compact_slot_groups, + _prepare_single_token_slots, + _single_token_weighted_reduce, _validate_weight_layout, validate_nvfp4_sm70_moe_contract, ) @@ -194,6 +196,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.cuda.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.cuda.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() diff --git a/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py b/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py index ccc73f908c..938439ddb9 100644 --- a/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py +++ b/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py @@ -49,6 +49,92 @@ _MAX_SUPPORTED_TOP_K: Final = max(contract[3] for contract in _SUPPORTED_CONTRACTS) +@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 def _prepare_compact_slot_groups_kernel( sorted_expert_ids_ptr, @@ -404,7 +490,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( @@ -491,7 +577,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), } @@ -549,31 +637,45 @@ def apply( return x.new_empty((0, hidden)) buffers = self._get_buffers(layer, num_tokens) output = buffers["output"] - output.zero_() slots = num_tokens * 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, - 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) + direct_single_token = num_tokens == 1 + 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"], @@ -582,7 +684,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) @@ -612,14 +714,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"], - 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( From c57d412ea5f7fb2db4d08f2ea030ef4af64a03d0 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 08:09:11 +0800 Subject: [PATCH 060/105] [Core][SM70] Start Qwen3.8 Flash Next prefill campaign Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> From da7beb2c72fb6df44409aa05cf50630a9d553fdc Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 08:29:25 +0800 Subject: [PATCH 061/105] [Model][SM70] Enable Qwen4Exp native MTP Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/models/qwen4_exp/test_config.py | 29 ++- tests/v1/spec_decode/test_qwen4_exp.py | 245 +++++++++++++++++++++++ vllm/config/speculative.py | 27 +++ vllm/model_executor/models/config.py | 24 ++- vllm/model_executor/models/registry.py | 1 + vllm/v1/spec_decode/llm_base_proposer.py | 6 +- vllm/v1/spec_decode/qwen4_exp.py | 178 ++++++++++++++++ vllm/v1/worker/gpu_model_runner.py | 16 +- 8 files changed, 517 insertions(+), 9 deletions(-) create mode 100644 tests/v1/spec_decode/test_qwen4_exp.py create mode 100644 vllm/v1/spec_decode/qwen4_exp.py diff --git a/tests/models/qwen4_exp/test_config.py b/tests/models/qwen4_exp/test_config.py index f7abe35682..ba72c90eb8 100644 --- a/tests/models/qwen4_exp/test_config.py +++ b/tests/models/qwen4_exp/test_config.py @@ -102,7 +102,7 @@ def test_qwen4_exp_defaults_to_v2_even_when_quantized_moe( assert VllmConfig._is_default_v2_model_runner_model(vllm_config) -def test_initial_sm70_v2_route_rejects_speculative_decode( +def test_initial_sm70_v2_route_accepts_native_mtp( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr( @@ -124,7 +124,32 @@ def test_initial_sm70_v2_route_rejects_speculative_decode( speculative_config=SimpleNamespace(method="mtp"), ) - with pytest.raises(NotImplementedError, match="initial SM70 V2 route"): + 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) 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/vllm/config/speculative.py b/vllm/config/speculative.py index c74aeb882e..07507b9be4 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", @@ -516,6 +517,23 @@ 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"], + } + ) + if hf_config.model_type == "exaone_moe": hf_config.model_type = "exaone_moe_mtp" if hf_config.model_type == "exaone_moe_mtp": @@ -1284,6 +1302,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/model_executor/models/config.py b/vllm/model_executor/models/config.py index 7674b2c539..dc0444b74c 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -623,12 +623,14 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: _strip_qwen4_exp_mrope(model_config) spec_config = vllm_config.speculative_config - if spec_config is not None: + if spec_config is not None and spec_config.method not in { + "mtp", + "ngram", + "ngram_gpu", + }: raise NotImplementedError( - "Qwen4Exp speculative decoding is disabled for the initial " - "SM70 V2 route. The checkpoint's PLE n-gram embedding remains " - "enabled; disable --speculative-config until the native MTP " - "follow-up is quality-gated." + "Qwen4Exp speculative decoding supports only its native MTP " + "checkpoint and linear n-gram proposers" ) @@ -639,6 +641,17 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: _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: @@ -701,6 +714,7 @@ def verify_and_update_model_config(model_config: "ModelConfig") -> None: "Qwen3_5MoeForConditionalGeneration": Qwen3_5ForConditionalGenerationConfig, "Qwen4ExpForCausalLM": Qwen4ExpForCausalLMConfig, "Qwen4ExpForConditionalGeneration": Qwen4ExpForConditionalGenerationConfig, + "Qwen4ExpMTP": Qwen4ExpMTPConfig, "VoyageQwen3BidirectionalEmbedModel": VoyageQwen3BidirectionalEmbedModelConfig, "XLMRobertaModel": JinaRobertaModelConfig, } diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index e053e2def8..80beb6e8fe 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -640,6 +640,7 @@ "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/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/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 53efe8b579..8ada303e4d 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -219,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, @@ -1518,6 +1519,7 @@ def __init__( | ExtractHiddenStatesProposer | Gemma4Proposer | Step3p5MTPProposer + | Qwen4ExpMTPProposer ) if self.speculative_config.method == "custom_class": self.drafter = create_custom_proposer( # type: ignore[assignment] @@ -1554,6 +1556,10 @@ 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 @@ -5664,7 +5670,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 ) From aa5384d298d15ac1f0c29e8ab48eeecf305120f0 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 08:41:56 +0800 Subject: [PATCH 062/105] [Core][SM70] Enable Qwen4Exp prefix caching Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/models/qwen4_exp/test_config.py | 4 ++-- tests/v1/core/test_qwen4_exp_kv_cache.py | 16 ++++++++++++++++ vllm/model_executor/models/config.py | 6 ------ vllm/v1/core/kv_cache_coordinator.py | 6 ++++-- vllm/v1/kv_cache_interface.py | 12 ++++++++++++ 5 files changed, 34 insertions(+), 10 deletions(-) diff --git a/tests/models/qwen4_exp/test_config.py b/tests/models/qwen4_exp/test_config.py index ba72c90eb8..c85602a401 100644 --- a/tests/models/qwen4_exp/test_config.py +++ b/tests/models/qwen4_exp/test_config.py @@ -22,7 +22,7 @@ ) -def test_initial_sm70_route_accepts_v2_runner( +def test_sm70_v2_route_accepts_prefix_caching( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr( @@ -42,7 +42,7 @@ def test_initial_sm70_route_accepts_v2_runner( hf_text_config=text_config, multimodal_config=SimpleNamespace(language_model_only=True), ), - cache_config=SimpleNamespace(enable_prefix_caching=False), + cache_config=SimpleNamespace(enable_prefix_caching=True), parallel_config=SimpleNamespace(enable_dbo=False, ubatch_size=1), speculative_config=None, use_v2_model_runner=True, diff --git a/tests/v1/core/test_qwen4_exp_kv_cache.py b/tests/v1/core/test_qwen4_exp_kv_cache.py index 1763bb1dda..4b777dc2e5 100644 --- a/tests/v1/core/test_qwen4_exp_kv_cache.py +++ b/tests/v1/core/test_qwen4_exp_kv_cache.py @@ -127,6 +127,22 @@ def test_qwen4_exp_csa_linear_cache_layout() -> None: ) 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( diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index dc0444b74c..1cb631a8af 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -594,12 +594,6 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: if text_config.hc_count <= 1: raise ValueError("Qwen4Exp requires hc_count > 1") - if vllm_config.cache_config.enable_prefix_caching: - raise NotImplementedError( - "Qwen4Exp prefix caching is not enabled in the initial SM70 " - "route; disable it while the QSA ring cache is in use" - ) - 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 diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index 224b4d98d7..eb0eec34dc 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -451,6 +451,8 @@ def verify_and_split_kv_cache_groups(self) -> None: ] = [] for i, g in enumerate(self.kv_cache_config.kv_cache_groups): + if not g.kv_cache_spec.prefix_cacheable: + continue manager_cls = self.single_type_managers[i].__class__ spec = g.kv_cache_spec @@ -465,8 +467,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/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index d55f5a9e41..d6ae8e0fc2 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -100,6 +100,10 @@ class KVCacheSpec: # number of tokens in a block block_size: int + @property + def prefix_cacheable(self) -> bool: + return True + @property def page_size_bytes(self) -> int: """ @@ -557,6 +561,10 @@ 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): @@ -740,6 +748,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()) From 606527ef79d5eb3aeb7a300495b615069fa1b36c Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 08:42:05 +0800 Subject: [PATCH 063/105] [Perf][SM70] Accelerate Qwen3.8 Flash Next prefill Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- csrc/sm70_turbomind/ops/awq_sm70_gemm.cu | 5 ++- tests/models/qwen4_exp/test_qsa_ops.py | 16 ++++++++ vllm/models/qwen4_exp/nvidia/ops/qsa.py | 47 +++++++++++++++++------- 3 files changed, 53 insertions(+), 15 deletions(-) create mode 100644 tests/models/qwen4_exp/test_qsa_ops.py diff --git a/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu b/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu index d555a8be7d..88d1dd9262 100644 --- a/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu +++ b/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu @@ -8583,8 +8583,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/tests/models/qwen4_exp/test_qsa_ops.py b/tests/models/qwen4_exp/test_qsa_ops.py new file mode 100644 index 0000000000..0026f2d59e --- /dev/null +++ b/tests/models/qwen4_exp/test_qsa_ops.py @@ -0,0 +1,16 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.models.qwen4_exp.nvidia.ops.qsa import ( + _qsa_sparse_launch_profile, +) + + +def test_sm70_qsa_prefill_uses_four_warps(): + assert _qsa_sparse_launch_profile(512, 8, True) == (64, 4, 4) + assert _qsa_sparse_launch_profile(8192, 8, True) == (64, 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) diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py index b1fda40d18..2052c98c26 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py @@ -858,20 +858,11 @@ def qsa_sparse_paged_attention( 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] - 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 + block_n, target_splits, partial_warps = _qsa_sparse_launch_profile( + base_programs, + block_m, + current_platform.is_device_capability(70), + ) num_tiles = triton.cdiv(logical_indices.shape[1], block_n) # Avoid empty splits when the selection width is smaller than the profile. @@ -953,6 +944,34 @@ def qsa_sparse_paged_attention( 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 + # preserve the exact reduction order while restoring warp-level + # parallelism for both split and non-split prefill profiles. + partial_warps = 4 + return block_n, target_splits, partial_warps + + def qsa_store_cache_rows( cache: torch.Tensor, slot_mapping: torch.Tensor, From c37a88424179c3e8be15a24ba202d39c364ecaa9 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 11:27:03 +0800 Subject: [PATCH 064/105] [Model] Add GLM-5.3-Flash SM70 support Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- csrc/ops.h | 40 +- csrc/sm70_turbomind/ops/awq_sm70_gemm.cu | 551 +++++++ csrc/torch_bindings.cpp | 24 + docs/design/sm70_glm53_flash_nvfp4.md | 176 +++ docs/design/sm70_v100_migration_control.md | 68 + .../test_kpool_decode_update_batched.py | 452 ++++++ tests/kernels/test_mhc_sm70_fp16.py | 420 +++++- .../test_weights_mapper_stacked.py | 33 + tests/models/glm5next/test_pp.py | 76 + tests/models/glm5next/test_sm70_kda.py | 50 + tests/models/glm5next/test_sm70_sparse.py | 291 ++++ .../test_sm70_glm53_modelopt_nvfp4.py | 77 + .../processors/test_glm5next.py | 401 +++++ .../attention/test_kpool_tail_slot_mapping.py | 331 +++++ .../test_sparse_indexer_decode_seq_lens.py | 156 ++ tests/v1/core/test_glm5_kv_cache_layout.py | 98 ++ .../worker/test_glm5_compressed_kv_cache.py | 67 + tests/v1/worker/test_gpu_model_runner_pp.py | 30 + tests/v1/worker/test_kv_block_zeroer.py | 38 + vllm/_sm70_ops.py | 128 ++ vllm/config/compilation.py | 1 + vllm/config/vllm.py | 8 +- vllm/model_executor/kernels/mhc/tilelang.py | 286 +++- vllm/model_executor/kernels/mhc/torch.py | 56 +- vllm/model_executor/kernels/mhc/triton.py | 446 ++++++ .../layers/attention/mla_attention.py | 13 +- .../layers/fla/ops/fused_recurrent.py | 27 +- vllm/model_executor/layers/fla/ops/kda.py | 116 +- .../layers/mamba/ops/gather_initial_states.py | 80 + .../layers/mamba/ops/scatter_states.py | 64 + vllm/model_executor/layers/mhc.py | 59 + vllm/model_executor/layers/mla.py | 25 +- .../layers/quantization/modelopt.py | 31 + .../layers/quantization/nvfp4_sm70_moe.py | 640 ++++++++ .../layers/quantization/sm70_turbomind.py | 7 + .../layers/sparse_attn_indexer_kpool.py | 1054 ++++++++++++++ vllm/model_executor/models/registry.py | 6 + vllm/model_executor/models/utils.py | 34 +- vllm/models/common/__init__.py | 2 + vllm/models/common/ops/__init__.py | 7 + vllm/models/common/ops/fused_qk_rmsnorm.py | 94 ++ vllm/models/common/ops/sequence_parallel.py | 68 + .../models/deepseek_v4/sm70/sparse_kernels.py | 23 +- vllm/models/glm5next/__init__.py | 18 + vllm/models/glm5next/nvidia/attention.py | 619 ++++++++ vllm/models/glm5next/nvidia/kda.py | 675 +++++++++ vllm/models/glm5next/nvidia/model.py | 1286 +++++++++++++++++ vllm/models/glm5next/nvidia/mtp.py | 424 ++++++ vllm/models/glm5next/nvidia/multimodal.py | 733 ++++++++++ .../glm5next/nvidia/ops/fused_eh_norm.py | 77 + .../glm5next/nvidia/ops/kpool_compress.py | 1061 ++++++++++++++ vllm/models/glm5next/sm70/__init__.py | 2 + vllm/models/glm5next/sm70/fp8_kv.py | 381 +++++ vllm/models/glm5next/sm70/sparse.py | 282 ++++ vllm/platforms/cuda.py | 17 +- vllm/platforms/interface.py | 15 + vllm/transformers_utils/config.py | 3 + vllm/transformers_utils/configs/__init__.py | 6 + vllm/transformers_utils/configs/glm5_next.py | 384 +++++ .../model_arch_config_convertor.py | 2 + .../transformers_utils/processors/__init__.py | 2 + .../transformers_utils/processors/glm5next.py | 1008 +++++++++++++ vllm/utils/deep_gemm.py | 4 + vllm/v1/attention/backend.py | 1 + vllm/v1/attention/backends/mla/indexer.py | 269 +++- vllm/v1/attention/backends/registry.py | 9 +- vllm/v1/core/kv_cache_coordinator.py | 16 +- vllm/v1/core/kv_cache_utils.py | 276 ++++ vllm/v1/core/single_type_kv_cache_manager.py | 97 ++ vllm/v1/engine/core.py | 24 +- vllm/v1/kv_cache_interface.py | 52 +- vllm/v1/worker/block_table.py | 28 +- vllm/v1/worker/gpu_model_runner.py | 64 +- vllm/v1/worker/utils.py | 150 +- 74 files changed, 14436 insertions(+), 203 deletions(-) create mode 100644 docs/design/sm70_glm53_flash_nvfp4.md create mode 100644 tests/kernels/test_kpool_decode_update_batched.py create mode 100644 tests/model_executor/test_weights_mapper_stacked.py create mode 100644 tests/models/glm5next/test_pp.py create mode 100644 tests/models/glm5next/test_sm70_kda.py create mode 100644 tests/models/glm5next/test_sm70_sparse.py create mode 100644 tests/quantization/test_sm70_glm53_modelopt_nvfp4.py create mode 100644 tests/transformers_utils/processors/test_glm5next.py create mode 100644 tests/v1/attention/test_kpool_tail_slot_mapping.py create mode 100644 tests/v1/attention/test_sparse_indexer_decode_seq_lens.py create mode 100644 tests/v1/core/test_glm5_kv_cache_layout.py create mode 100644 tests/v1/worker/test_glm5_compressed_kv_cache.py create mode 100644 tests/v1/worker/test_gpu_model_runner_pp.py create mode 100644 tests/v1/worker/test_kv_block_zeroer.py create mode 100644 vllm/model_executor/layers/mamba/ops/gather_initial_states.py create mode 100644 vllm/model_executor/layers/mamba/ops/scatter_states.py create mode 100644 vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py create mode 100644 vllm/model_executor/layers/sparse_attn_indexer_kpool.py create mode 100644 vllm/models/common/__init__.py create mode 100644 vllm/models/common/ops/__init__.py create mode 100644 vllm/models/common/ops/fused_qk_rmsnorm.py create mode 100644 vllm/models/common/ops/sequence_parallel.py create mode 100644 vllm/models/glm5next/__init__.py create mode 100644 vllm/models/glm5next/nvidia/attention.py create mode 100644 vllm/models/glm5next/nvidia/kda.py create mode 100644 vllm/models/glm5next/nvidia/model.py create mode 100644 vllm/models/glm5next/nvidia/mtp.py create mode 100644 vllm/models/glm5next/nvidia/multimodal.py create mode 100644 vllm/models/glm5next/nvidia/ops/fused_eh_norm.py create mode 100644 vllm/models/glm5next/nvidia/ops/kpool_compress.py create mode 100644 vllm/models/glm5next/sm70/__init__.py create mode 100644 vllm/models/glm5next/sm70/fp8_kv.py create mode 100644 vllm/models/glm5next/sm70/sparse.py create mode 100644 vllm/transformers_utils/configs/glm5_next.py create mode 100644 vllm/transformers_utils/processors/glm5next.py diff --git a/csrc/ops.h b/csrc/ops.h index 2ee0b700f5..be8c335a62 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -237,6 +237,29 @@ void sm70_f16_gemm_out(torch::Tensor out, int64_t k_ld, bool 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); + +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, @@ -503,15 +526,18 @@ void fp8_moe_dense_stage_sm70_out(torch::Tensor out, int64_t n, int64_t group_size); -void mxfp4_moe_dense_stage_sm70_out(torch::Tensor out, - torch::Tensor input, +void mxfp4_moe_dense_stage_sm70_out(torch::Tensor out, torch::Tensor input, + torch::Tensor expert_offsets, + torch::Tensor dense_expert_ids, + torch::Tensor ptrs_w, torch::Tensor ptrs_s, + int64_t num_experts, int64_t k, int64_t n, + int64_t group_size); + +void nvfp4_moe_dense_stage_sm70_out(torch::Tensor out, torch::Tensor input, torch::Tensor expert_offsets, torch::Tensor dense_expert_ids, - torch::Tensor ptrs_w, - torch::Tensor ptrs_s, - int64_t num_experts, - int64_t k, - int64_t n, + torch::Tensor ptrs_w, torch::Tensor ptrs_s, + int64_t num_experts, int64_t k, int64_t n, int64_t group_size); void mxfp4_moe_single_token_prepare_w13_sm70_out( diff --git a/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu b/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu index cd6bb22a70..2c988d9eb4 100644 --- a/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu +++ b/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu @@ -911,6 +911,7 @@ enum class TuneKeyKind : int { kMxfp4Dense = 6, kNvfp4Dense = 7, kMxfp4Moe = 8, + kNvfp4Moe = 9, }; struct DenseTuneKey { @@ -1046,6 +1047,9 @@ turbomind::gemm::DispatchPolicy select_fp8_moe_dispatch_policy( turbomind::gemm::DispatchPolicy select_mxfp4_moe_dispatch_policy( int device, int total_tokens, int n, int k, int num_experts, int group_size, cudaStream_t stream); +turbomind::gemm::DispatchPolicy select_nvfp4_moe_dispatch_policy( + int device, int total_tokens, int n, int k, int num_experts, int group_size, + cudaStream_t stream); bool tune_small_shapes_enabled() { const char* raw = std::getenv("VLLM_SM70_AWQ_TUNE_SMALL_SHAPES"); @@ -1087,6 +1091,11 @@ bool nvfp4_tune_small_shapes_enabled() { return raw == nullptr || std::atoi(raw) != 0; } +bool nvfp4_moe_grouped_prefill_enabled() { + const char* raw = std::getenv("VLLM_SM70_NVFP4_MOE_GROUPED_PREFILL"); + return raw == nullptr || std::atoi(raw) != 0; +} + bool fp8_moe_single_token_per_expert_dispatch_enabled() { const char* raw = std::getenv("VLLM_SM70_FP8_MOE_SINGLE_TOKEN_PER_EXPERT_DISPATCH"); @@ -1362,6 +1371,14 @@ turbomind::gemm::DispatchPolicy select_mxfp4_moe_dispatch_policy( TuneKeyKind::kMxfp4Moe, mxfp4_tune_small_shapes_enabled()); } +turbomind::gemm::DispatchPolicy select_nvfp4_moe_dispatch_policy( + int device, int total_tokens, int n, int k, int num_experts, int group_size, + cudaStream_t stream) { + return select_moe_dispatch_policy_impl( + device, total_tokens, n, k, num_experts, group_size, stream, + TuneKeyKind::kNvfp4Moe, nvfp4_tune_small_shapes_enabled()); +} + WorkspaceHolder& get_workspace(int device, cudaStream_t stream) { thread_local int cached_device = -1; thread_local cudaStream_t cached_stream = nullptr; @@ -1740,6 +1757,294 @@ __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, @@ -4946,6 +5251,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_lm_head_top1_out(torch::Tensor values_out, torch::Tensor indices_out, torch::Tensor _in_feats, torch::Tensor _kernel, @@ -7582,6 +7907,232 @@ void mxfp4_moe_dense_stage_sm70_out(torch::Tensor out, torch::Tensor input, } } +void nvfp4_moe_gemm_sm70_out_impl( + torch::Tensor out, torch::Tensor sorted_input, torch::Tensor expert_offsets, + torch::Tensor strided_ptrs_w, torch::Tensor strided_ptrs_s, + int64_t num_experts, int64_t k, int64_t n, int64_t group_size, + torch::Tensor b_group_indices, bool compact_grouped_rows = false) { + TORCH_CHECK( + sorted_input.is_cuda() && sorted_input.scalar_type() == torch::kFloat16, + "nvfp4_moe_gemm_sm70: input must be CUDA float16."); + TORCH_CHECK( + expert_offsets.is_cuda() && + expert_offsets.scalar_type() == torch::kInt32 && + expert_offsets.is_contiguous(), + "nvfp4_moe_gemm_sm70: expert_offsets must be contiguous CUDA int32."); + TORCH_CHECK(strided_ptrs_w.is_cuda() && strided_ptrs_s.is_cuda(), + "nvfp4_moe_gemm_sm70: strided_ptrs must be CUDA."); + TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kFloat16, + "nvfp4_moe_gemm_sm70: output must be CUDA float16."); + TORCH_CHECK(num_experts > 0 && k > 0 && n > 0, + "nvfp4_moe_gemm_sm70: invalid dimensions."); + TORCH_CHECK(group_size == 16, + "nvfp4_moe_gemm_sm70: only group_size=16 is supported."); + TORCH_CHECK(k % group_size == 0, + "nvfp4_moe_gemm_sm70: k must be divisible by group_size."); + TORCH_CHECK(sorted_input.dim() == 2 && sorted_input.size(1) == k, + "nvfp4_moe_gemm_sm70: input shape mismatch."); + TORCH_CHECK(out.dim() == 2 && out.size(0) == sorted_input.size(0) && + out.size(1) == n && out.stride(1) == 1, + "nvfp4_moe_gemm_sm70: output must be contiguous [tokens, n]."); + TORCH_CHECK(expert_offsets.numel() >= num_experts + 1, + "nvfp4_moe_gemm_sm70: expert_offsets too small."); + TORCH_CHECK(b_group_indices.is_cuda() && + b_group_indices.scalar_type() == torch::kInt32 && + b_group_indices.is_contiguous() && + b_group_indices.numel() >= num_experts, + "nvfp4_moe_gemm_sm70: B group indices must be contiguous CUDA " + "int32."); + + const int64_t total_tokens = sorted_input.size(0); + if (total_tokens == 0) { + return; + } + + const at::cuda::OptionalCUDAGuard device_guard(device_of(sorted_input)); + const int device = sorted_input.get_device(); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + const auto fp4_converters = turbomind::gemm::GetConverters( + turbomind::kHalf, turbomind::kFloat4_e2m1, turbomind::kHalf, true, 70); + const auto fp8_converters = turbomind::gemm::GetConverters( + turbomind::kHalf, turbomind::kFloat8_e4m3, turbomind::kHalf, true, 70); + const auto* conv_w = fp4_converters[0]; + const auto* conv_s = fp8_converters[1]; + TORCH_CHECK(conv_w && conv_s, + "nvfp4_moe_gemm_sm70: no compatible TurboMind converters."); + + turbomind::gemm::MatrixLayout desc_A{ + turbomind::kHalf, + turbomind::gemm::kRowMajor, + static_cast(total_tokens), + static_cast(k), + static_cast(sorted_input.stride(0)), + }; + desc_A.num = static_cast(num_experts); + desc_A.offsets = expert_offsets.data_ptr(); + turbomind::gemm::MatrixLayout desc_U{}; + + const auto order_w = conv_w->order; + const bool is_a_w = turbomind::gemm::get_operand_tag(conv_w->pack) == + turbomind::gemm::OPERAND_A; + const bool is_b_w = !is_a_w; + turbomind::gemm::MatrixLayout weight_desc{ + turbomind::kHalf, + order_w, + static_cast(n), + static_cast(k), + order_w == turbomind::gemm::kRowMajor ? static_cast(k) + : static_cast(n), + }; + if (is_b_w) { + std::swap(weight_desc.rows, weight_desc.cols); + weight_desc.order = ~weight_desc.order; + } + turbomind::gemm::MatrixLayout desc_B = weight_desc; + desc_B.type = turbomind::kFloat4_e2m1; + desc_B.pack = conv_w->pack; + if (is_a_w) { + desc_B = turbomind::gemm::transpose(desc_B); + } + desc_B.ld = 0; + desc_B.num = static_cast(num_experts); + desc_B.group_idxs = b_group_indices.data_ptr(); + + const auto order_s = conv_s->order; + const bool is_a_s = turbomind::gemm::get_operand_tag(conv_s->pack) == + turbomind::gemm::OPERAND_U; + const bool is_b_s = !is_a_s; + const int64_t num_groups = k / group_size; + turbomind::gemm::MatrixLayout scale_desc{ + turbomind::kUint16, order_s, + static_cast(n), static_cast(num_groups), + static_cast(n), + }; + if (is_b_s) { + std::swap(scale_desc.rows, scale_desc.cols); + scale_desc.order = ~scale_desc.order; + } + turbomind::gemm::MatrixLayout desc_V = scale_desc; + desc_V.pack = conv_s->pack; + if (is_a_s) { + desc_V = turbomind::gemm::transpose(desc_V); + } + desc_V.ld = 0; + desc_V.num = static_cast(num_experts); + desc_V.group_idxs = b_group_indices.data_ptr(); + + turbomind::gemm::MatrixLayout desc_D{ + turbomind::kHalf, + turbomind::gemm::kRowMajor, + static_cast(total_tokens), + static_cast(n), + static_cast(out.stride(0)), + }; + desc_D.num = static_cast(num_experts); + desc_D.offsets = expert_offsets.data_ptr(); + + turbomind::gemm::Operation op{}; + op.dispatch = vllm::awq_sm70::select_nvfp4_moe_dispatch_policy( + device, static_cast(total_tokens), static_cast(n), + static_cast(k), static_cast(num_experts), + static_cast(group_size), stream); + op.epilogue = turbomind::gemm::Epilogue::kNone; + op.quant_a = {turbomind::gemm::QuantType::kNone, 0}; + op.quant_b = {turbomind::gemm::QuantType::kK, static_cast(group_size)}; + op.batch_dim = 0; + op.dispatch_num_override = compact_grouped_rows ? 1 : 0; + op.active_group_count = + compact_grouped_rows ? -static_cast(num_experts) : 0; + + auto& workspace_holder = vllm::awq_sm70::get_workspace(device, stream); + auto& gemm = vllm::awq_sm70::get_gemm(device); + const int ec = + gemm.Run(op, 1.f, sorted_input.data_ptr(), desc_A, nullptr, desc_U, + strided_ptrs_w.data_ptr(), desc_B, strided_ptrs_s.data_ptr(), + desc_V, 0.f, out.data_ptr(), desc_D, out.data_ptr(), desc_D, + workspace_holder.workspace, stream); + TORCH_CHECK(ec == 0, + "nvfp4_moe_gemm_sm70: TurboMind batched GEMM failed (ec=", ec, + ")."); +} + +void nvfp4_moe_dense_stage_sm70_out(torch::Tensor out, torch::Tensor input, + torch::Tensor expert_offsets, + torch::Tensor dense_expert_ids, + torch::Tensor ptrs_w, torch::Tensor ptrs_s, + int64_t num_experts, int64_t k, int64_t n, + int64_t group_size) { + TORCH_CHECK(input.is_cuda() && input.scalar_type() == torch::kFloat16, + "nvfp4_moe_dense_stage_sm70_out: input must be CUDA float16."); + TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kFloat16, + "nvfp4_moe_dense_stage_sm70_out: out must be CUDA float16."); + TORCH_CHECK(expert_offsets.is_cuda() && + expert_offsets.scalar_type() == torch::kInt32 && + expert_offsets.is_contiguous(), + "nvfp4_moe_dense_stage_sm70_out: expert_offsets must be " + "contiguous CUDA int32."); + TORCH_CHECK(dense_expert_ids.is_cuda() && + dense_expert_ids.scalar_type() == torch::kInt32 && + dense_expert_ids.is_contiguous(), + "nvfp4_moe_dense_stage_sm70_out: dense_expert_ids must be " + "contiguous CUDA int32."); + TORCH_CHECK(ptrs_w.is_cuda() && ptrs_s.is_cuda(), + "nvfp4_moe_dense_stage_sm70_out: ptr rows must be CUDA."); + TORCH_CHECK(num_experts > 0, + "nvfp4_moe_dense_stage_sm70_out: num_experts must be positive."); + TORCH_CHECK(group_size == 16, + "nvfp4_moe_dense_stage_sm70_out: only group_size=16 is " + "supported."); + TORCH_CHECK(input.dim() == 2 && input.size(1) == k, + "nvfp4_moe_dense_stage_sm70_out: input shape mismatch."); + TORCH_CHECK( + out.dim() == 2 && out.size(0) == input.size(0) && out.size(1) == n, + "nvfp4_moe_dense_stage_sm70_out: out shape mismatch."); + TORCH_CHECK(expert_offsets.numel() >= num_experts + 1, + "nvfp4_moe_dense_stage_sm70_out: expert_offsets too small."); + TORCH_CHECK(dense_expert_ids.numel() >= num_experts, + "nvfp4_moe_dense_stage_sm70_out: dense_expert_ids too small."); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); + static std::atomic logged_nvfp4_dense_stage{0u}; + maybe_log_sm70_moe_route_once( + logged_nvfp4_dense_stage, + "SM70 NVFP4 MoE CUDA-graph-safe TurboMind path enabled C++ op reached", + input, input.size(0), num_experts); + constexpr int kNvfp4MaxCompactGroups = 8 * 8; + const bool compact_decode_shape = + input.size(0) == num_experts && num_experts <= kNvfp4MaxCompactGroups; + if (compact_decode_shape) { + nvfp4_moe_gemm_sm70_out_impl(out, input, expert_offsets, ptrs_w, ptrs_s, + num_experts, k, n, group_size, + dense_expert_ids, true); + return; + } + const bool exact_qwen36_prefill_shape = + input.size(0) > kNvfp4MaxCompactGroups && num_experts == 256 && + ((k == 2048 && (n == 1024 || n == 512 || n == 256)) || + (n == 2048 && (k == 512 || k == 256 || k == 128))); + if (vllm::awq_sm70::nvfp4_moe_grouped_prefill_enabled() && + exact_qwen36_prefill_shape) { + static std::atomic logged_nvfp4_grouped_prefill{0u}; + maybe_log_sm70_moe_route_once( + logged_nvfp4_grouped_prefill, + "SM70 NVFP4 MoE grouped TurboMind prefill path enabled C++ op reached", + input, input.size(0), num_experts); + nvfp4_moe_gemm_sm70_out_impl(out, input, expert_offsets, ptrs_w, ptrs_s, + num_experts, k, n, group_size, + dense_expert_ids); + return; + } + for (int expert = 0; expert < static_cast(num_experts); ++expert) { + torch::Tensor offsets = expert_offsets.narrow(0, expert, 2); + torch::Tensor expert_idx = dense_expert_ids.narrow(0, expert, 1); + nvfp4_moe_gemm_sm70_out_impl(out, input, offsets, ptrs_w, ptrs_s, 1, k, n, + group_size, expert_idx); + } +} + void mxfp4_moe_single_token_prepare_w13_sm70_out( torch::Tensor gate_up, torch::Tensor compact_input, torch::Tensor x, torch::Tensor topk_ids, torch::Tensor w13_ptrs_w, torch::Tensor w13_ptrs_s, diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index bd5b2c45e2..14c9199353 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -295,6 +295,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_lm_head_top1_out(Tensor(a!) values_out, " "Tensor(b!) indices_out, Tensor _in_feats, Tensor _kernel, int k_ld, " @@ -500,6 +516,14 @@ 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( + "nvfp4_moe_dense_stage_sm70_out(" + "Tensor(a!) out, Tensor input, Tensor expert_offsets, " + "Tensor dense_expert_ids, Tensor ptrs_w, Tensor ptrs_s, " + "int num_experts, int k, int n, int group_size) -> ()"); + ops.impl("nvfp4_moe_dense_stage_sm70_out", torch::kCUDA, + &nvfp4_moe_dense_stage_sm70_out); + ops.def( "mxfp4_moe_single_token_prepare_w13_sm70_out(" "Tensor(a!) gate_up, Tensor(b!) compact_input, Tensor x, " 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_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index 74b8425639..67146fbd63 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -41671,3 +41671,71 @@ Interpretation: formatter changes from current main and keep large C/CUDA formatter work 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. 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 75b5ae7242..eace9cdb42 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.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run(*graph_outputs) + graph.replay() + torch.cuda.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/model_executor/test_weights_mapper_stacked.py b/tests/model_executor/test_weights_mapper_stacked.py new file mode 100644 index 0000000000..6d1d4e8b48 --- /dev/null +++ b/tests/model_executor/test_weights_mapper_stacked.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.model_executor.models.utils import WeightsMapper + + +def test_weights_mapper_preserves_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" 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..fa5e5ffecd --- /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.cuda.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.cuda.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..b2642bd647 --- /dev/null +++ b/tests/models/glm5next/test_sm70_sparse.py @@ -0,0 +1,291 @@ +# 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.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_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.cuda.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.cuda.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.cuda.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.cuda.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/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/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_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_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..2c86332b6f --- /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.participates_in_prefix_caching + ] + 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/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_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_kv_block_zeroer.py b/tests/v1/worker/test_kv_block_zeroer.py new file mode 100644 index 0000000000..1fbd1e0882 --- /dev/null +++ b/tests/v1/worker/test_kv_block_zeroer.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest +import torch + +from vllm.platforms import current_platform +from vllm.v1.worker.utils import KVBlockZeroer + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") +def test_zeroer_supports_heterogeneous_page_sizes() -> None: + device = torch.device(current_platform.device_type) + 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/vllm/_sm70_ops.py b/vllm/_sm70_ops.py index 0bcd044544..7cf35983e1 100644 --- a/vllm/_sm70_ops.py +++ b/vllm/_sm70_ops.py @@ -541,6 +541,50 @@ def _mxfp4_moe_dense_stage_sm70_out_fake( return None +def nvfp4_moe_dense_stage_sm70_out( + out: torch.Tensor, + input: torch.Tensor, + expert_offsets: torch.Tensor, + dense_expert_ids: torch.Tensor, + ptrs_w: torch.Tensor, + ptrs_s: torch.Tensor, + num_experts: int, + k: int, + n: int, + group_size: int, +) -> None: + _op("nvfp4_moe_dense_stage_sm70_out")( + out, + input, + expert_offsets, + dense_expert_ids, + ptrs_w, + ptrs_s, + num_experts, + k, + n, + group_size, + ) + + +if hasattr(torch.ops._C, "nvfp4_moe_dense_stage_sm70_out"): + + @register_fake("_C::nvfp4_moe_dense_stage_sm70_out") + def _nvfp4_moe_dense_stage_sm70_out_fake( + out: torch.Tensor, + input: torch.Tensor, + expert_offsets: torch.Tensor, + dense_expert_ids: torch.Tensor, + ptrs_w: torch.Tensor, + ptrs_s: torch.Tensor, + num_experts: int, + k: int, + n: int, + group_size: int, + ) -> None: + return None + + def mxfp4_moe_single_token_prepare_w13_sm70_out( gate_up: torch.Tensor, compact_input: torch.Tensor, @@ -767,6 +811,90 @@ 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_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_lm_head_top1_out( values_out: torch.Tensor, indices_out: torch.Tensor, diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index 2b0833e545..fa7b695d0e 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -775,6 +775,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/vllm.py b/vllm/config/vllm.py index ea67803abb..e67637bc3f 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -66,7 +66,13 @@ logger = init_logger(__name__) -DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset({"Qwen3ForCausalLM"}) +DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset( + { + "Glm5NextForCausalLM", + "Glm5NextForConditionalGeneration", + "Qwen3ForCausalLM", + } +) _SM70_NOMTP_CUDAGRAPH_CAPTURE_SIZES = (1, 2, 4, 8, 16) diff --git a/vllm/model_executor/kernels/mhc/tilelang.py b/vllm/model_executor/kernels/mhc/tilelang.py index de5f02d8ce..63dd0ca22c 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, @@ -633,6 +783,13 @@ def mhc_fused_post_pre_tilelang( and tile_n == 2 and n_splits == 8 ) + 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, @@ -716,19 +873,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( @@ -748,7 +928,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, @@ -908,6 +1118,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/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/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/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index eabaf62be7..f1f8d60622 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -8,6 +8,7 @@ from torch.nn.parameter import Parameter import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm import envs from vllm.config import get_current_vllm_config from vllm.logger import init_logger from vllm.model_executor.kernels.linear import ( @@ -48,6 +49,7 @@ UnquantizedLinearMethod, ) from vllm.model_executor.layers.quantization import QuantizationMethods +from vllm.model_executor.layers.quantization import sm70_turbomind as sm70_tm from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, QuantizeMethodBase, @@ -1048,8 +1050,37 @@ 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: + # SM70 is admitted only when a software backend is selected. Exact + # device and MoE-contract checks happen during quant-method dispatch. + if ( + sm70_tm.use_turbomind(envs.VLLM_SM70_NVFP4_TURBOMIND) + or sm70_tm.forces_marlin() + ): + return 70 return 75 @classmethod diff --git a/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py b/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py new file mode 100644 index 0000000000..d395652853 --- /dev/null +++ b/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py @@ -0,0 +1,640 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Native SM70 TurboMind NVFP4 MoE for validated expert shapes. + +The route keeps ModelOpt NVFP4 expert weights packed. It combines the +checkpoint's FP8 block scales with its explicit ModelOpt global scales once at +load time, repacks both tensors for TurboMind, and never materializes an FP16 +expert-weight copy. +""" + +from __future__ import annotations + +from typing import Final + +import torch +from torch.nn import Parameter + +from vllm import _sm70_ops as sm70_ops +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe import ( + FusedMoEConfig, + FusedMoEMethodBase, + FusedMoEQuantConfig, + MoEActivation, + RoutedExperts, + SharedExperts, +) +from vllm.model_executor.layers.quantization.modelopt import ( + ModelOptNvFp4Config, + ModelOptNvFp4FusedMoE, +) +from vllm.model_executor.layers.quantization.sm70_turbomind import ( + NVFP4_GROUP_SIZE, + is_exact_sm70_cuda, + unpack_mxfp4_weight, +) +from vllm.triton_utils import tl, triton + +logger = init_logger(__name__) + +_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) + + +@triton.jit +def _prepare_compact_slot_groups_kernel( + sorted_expert_ids_ptr, + compact_offsets_ptr, + active_expert_ids_ptr, + TOTAL_SLOTS: tl.constexpr, + BLOCK: tl.constexpr, +): + offsets = tl.arange(0, BLOCK) + valid = offsets < TOTAL_SLOTS + expert_ids = tl.load( + sorted_expert_ids_ptr + offsets, + mask=valid, + other=-1, + ) + tl.store( + compact_offsets_ptr + offsets, + offsets, + mask=offsets <= TOTAL_SLOTS, + ) + tl.store( + active_expert_ids_ptr + offsets, + expert_ids, + mask=valid, + ) + + +def _prepare_compact_slot_groups( + sorted_expert_ids: torch.Tensor, + compact_offsets: torch.Tensor, + active_expert_ids: torch.Tensor, +) -> None: + total_slots = sorted_expert_ids.numel() + 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) + # TurboMind's compact grouped dispatch forces one row per group. Keep each + # routed slot independent even when adjacent slots select the same expert; + # coalescing duplicate expert IDs would make the forced one-row scheduler + # silently skip or miscompute the additional rows. + _prepare_compact_slot_groups_kernel[(1,)]( + sorted_expert_ids, + compact_offsets, + active_expert_ids, + TOTAL_SLOTS=total_slots, + BLOCK=block, + num_warps=1, + ) + + +def validate_nvfp4_sm70_moe_contract(moe: FusedMoEConfig) -> None: + """Reject every topology outside the validated SM70 NVFP4 contract.""" + if moe.tp_size not in _SUPPORTED_TP_SIZES: + raise NotImplementedError( + "SM70 TurboMind NVFP4 MoE currently supports tensor parallel " + 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: + raise NotImplementedError( + "SM70 TurboMind NVFP4 MoE requires a positive local intermediate " + f"size divisible by {NVFP4_GROUP_SIZE}, got {local_intermediate}." + ) + 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 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( + "SM70 TurboMind NVFP4 MoE does not support DP+EP all-to-all." + ) + + +def _validate_weight_layout(layer: RoutedExperts) -> None: + local_experts = int(layer.local_num_experts) + hidden = int(layer.moe_config.hidden_dim) + intermediate = int(layer.moe_config.intermediate_size_per_partition) + expected = { + "w13_weight": (local_experts, 2 * intermediate, hidden // 2), + "w13_weight_scale": ( + local_experts, + 2 * intermediate, + hidden // NVFP4_GROUP_SIZE, + ), + "w13_weight_scale_2": (local_experts, 2), + "w2_weight": (local_experts, hidden, intermediate // 2), + "w2_weight_scale": ( + local_experts, + hidden, + intermediate // NVFP4_GROUP_SIZE, + ), + "w2_weight_scale_2": (local_experts,), + } + tensors = {name: getattr(layer, name) for name in expected} + for name, shape in expected.items(): + if tuple(tensors[name].shape) != shape: + raise ValueError( + f"SM70 NVFP4 MoE layout mismatch for {name}: " + f"expected {shape}, got {tuple(tensors[name].shape)}." + ) + if layer.w13_weight.dtype != torch.uint8 or layer.w2_weight.dtype != torch.uint8: + raise TypeError("SM70 NVFP4 MoE requires packed uint8 expert weights.") + if ( + layer.w13_weight_scale.dtype != torch.float8_e4m3fn + or layer.w2_weight_scale.dtype != torch.float8_e4m3fn + ): + raise TypeError("SM70 NVFP4 MoE requires FP8 E4M3 block scales.") + + +class ModelOptNvFp4SM70MoEMethod(ModelOptNvFp4FusedMoE): + """ModelOpt NVFP4 experts with FP16 activations on native SM70.""" + + def __init__( + self, + quant_config: ModelOptNvFp4Config, + moe_config: FusedMoEConfig, + ) -> None: + FusedMoEMethodBase.__init__(self, moe_config) + if quant_config.quant_method not in {"NVFP4", "W4A16_NVFP4"}: + raise NotImplementedError( + "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 + self.use_global_sf = False + validate_nvfp4_sm70_moe_contract(moe_config) + + @property + def supports_eplb(self) -> bool: + return False + + def maybe_make_prepare_finalize( + self, + routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, + ) -> None: + # This method owns routing, TurboMind expert GEMMs, and unpermutation. + # Do not wrap it in the generic ModelOpt modular-kernel path. + del routing_tables + return None + + def process_weights_after_loading(self, layer: RoutedExperts) -> None: + required_ops = ( + "nvfp4_sm70_prepare", + "nvfp4_moe_dense_stage_sm70_out", + "awq_moe_build_strided_ptrs", + ) + missing = [name for name in required_ops if not hasattr(torch.ops._C, name)] + if missing: + raise RuntimeError( + "SM70 NVFP4 MoE requires the TurboMind extension " + "with " + ", ".join(missing) + "." + ) + if not hasattr(torch.ops._moe_C, "moe_permute_with_scratch"): + 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: + raise NotImplementedError( + "SM70 NVFP4 MoE currently supports SwiGLU/SILU only." + ) + if layer.apply_router_weight_on_input: + raise NotImplementedError( + "SM70 NVFP4 MoE does not support router weights on input." + ) + if layer.expert_map is not None: + raise NotImplementedError( + "SM70 NVFP4 MoE currently requires fully replicated experts." + ) + if layer.local_num_experts != layer.global_num_experts: + raise NotImplementedError( + "SM70 NVFP4 MoE currently requires local and global experts to match." + ) + + validate_nvfp4_sm70_moe_contract(layer.moe_config) + _validate_weight_layout(layer) + num_experts = int(layer.local_num_experts) + hidden = int(layer.moe_config.hidden_dim) + intermediate = int(layer.moe_config.intermediate_size_per_partition) + + w13_tm_weights: list[torch.Tensor] = [] + w13_tm_scales: list[torch.Tensor] = [] + w13_meta: list[torch.Tensor] = [] + w2_tm_weights: list[torch.Tensor] = [] + w2_tm_scales: list[torch.Tensor] = [] + w2_meta: list[torch.Tensor] = [] + for expert_id in range(num_experts): + w13_packed = unpack_mxfp4_weight(layer.w13_weight[expert_id].data) + w13_scales = layer.w13_weight_scale[expert_id].float().clone() + w13_global = layer.w13_weight_scale_2[expert_id].float() + w13_scales[:intermediate].mul_(w13_global[0]) + w13_scales[intermediate:].mul_(w13_global[1]) + prepared_w13 = sm70_ops.nvfp4_sm70_prepare( + w13_packed, + w13_scales.half().t().contiguous(), + NVFP4_GROUP_SIZE, + ) + w13_tm_weights.append(prepared_w13[0]) + w13_tm_scales.append(prepared_w13[1]) + w13_meta.append(prepared_w13[2]) + + w2_packed = unpack_mxfp4_weight(layer.w2_weight[expert_id].data) + w2_scales = ( + layer.w2_weight_scale[expert_id].float() + * layer.w2_weight_scale_2[expert_id].float() + ) + prepared_w2 = sm70_ops.nvfp4_sm70_prepare( + w2_packed, + w2_scales.half().t().contiguous(), + NVFP4_GROUP_SIZE, + ) + w2_tm_weights.append(prepared_w2[0]) + w2_tm_scales.append(prepared_w2[1]) + w2_meta.append(prepared_w2[2]) + + layer.w13_tm_weight = Parameter( + torch.stack(w13_tm_weights), requires_grad=False + ) + layer.w13_tm_scales = Parameter(torch.stack(w13_tm_scales), requires_grad=False) + layer.w2_tm_weight = Parameter(torch.stack(w2_tm_weights), requires_grad=False) + layer.w2_tm_scales = Parameter(torch.stack(w2_tm_scales), requires_grad=False) + + w13_k_ld = int(w13_meta[0][0].item()) + 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()) + w13_ptrs = sm70_ops.awq_moe_build_strided_ptrs( + layer.w13_tm_weight, + layer.w13_tm_scales, + w13_k_ld, + w13_q_ld, + num_experts, + ) + w2_ptrs = sm70_ops.awq_moe_build_strided_ptrs( + layer.w2_tm_weight, + layer.w2_tm_scales, + w2_k_ld, + w2_q_ld, + num_experts, + ) + layer.w13_strided_ptrs_w = Parameter(w13_ptrs[0], requires_grad=False) + layer.w13_strided_ptrs_s = Parameter(w13_ptrs[1], requires_grad=False) + layer.w2_strided_ptrs_w = Parameter(w2_ptrs[0], requires_grad=False) + layer.w2_strided_ptrs_s = Parameter(w2_ptrs[1], requires_grad=False) + + layer.sm70_nvfp4_moe = True + 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 + layer.sm70_nvfp4_w2_n_dim = hidden + layer.sm70_nvfp4_group_size = NVFP4_GROUP_SIZE + layer.sm70_nvfp4_graph_safe_max_tokens = _GRAPH_SAFE_MAX_TOKENS + layer.sm70_nvfp4_compact_grouped_max_tokens = _COMPACT_GROUPED_MAX_TOKENS + self._allocate_graph_safe_decode_buffers(layer) + + del layer.w13_weight + del layer.w13_weight_scale + del layer.w13_weight_scale_2 + del layer.w13_input_scale + del layer.w2_weight + del layer.w2_weight_scale + del layer.w2_weight_scale_2 + del layer.w2_input_scale + logger.info_once( + "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 + 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) + + layer._nvfp4_sm70_output = torch.empty( + _GRAPH_SAFE_MAX_TOKENS, hidden, dtype=torch.float16, device=device + ) + layer._nvfp4_sm70_permuted_input = torch.empty( + max_slots, hidden, dtype=torch.float16, device=device + ) + layer._nvfp4_sm70_gate_up = torch.empty( + max_slots, 2 * intermediate, dtype=torch.float16, device=device + ) + layer._nvfp4_sm70_intermediate = torch.empty( + max_slots, intermediate, dtype=torch.float16, device=device + ) + layer._nvfp4_sm70_sorted_output = torch.empty( + max_slots, hidden, dtype=torch.float16, device=device + ) + layer._nvfp4_sm70_expert_offsets = torch.empty( + experts + 1, dtype=torch.int32, device=device + ) + layer._nvfp4_sm70_expert_offsets64 = torch.empty( + experts + 1, dtype=torch.int64, device=device + ) + layer._nvfp4_sm70_inv_permuted_idx = torch.empty( + _GRAPH_SAFE_MAX_TOKENS, + top_k, + dtype=torch.int32, + device=device, + ) + layer._nvfp4_sm70_topk_ids = torch.empty( + _GRAPH_SAFE_MAX_TOKENS, + 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, top_k) + layer._nvfp4_sm70_permuted_idx = torch.empty( + max_slots, dtype=torch.int32, device=device + ) + layer._nvfp4_sm70_permuted_experts_id = torch.empty( + max_slots, dtype=torch.int32, device=device + ) + layer._nvfp4_sm70_sorted_row_idx = torch.empty( + max_slots, dtype=torch.int32, device=device + ) + layer._nvfp4_sm70_topk_ids_for_sort = torch.empty( + max_slots, dtype=torch.int32, device=device + ) + workspace_size = torch.ops._moe_C.moe_permute_sort_workspace_size( + max_slots, layer.global_num_experts + ) + layer._nvfp4_sm70_sort_workspace = torch.empty( + workspace_size, dtype=torch.int8, device=device + ) + layer._nvfp4_sm70_dense_expert_ids = torch.arange( + experts, dtype=torch.int32, device=device + ) + layer._nvfp4_sm70_compact_offsets = torch.empty( + max_slots + 1, dtype=torch.int32, device=device + ) + layer._nvfp4_sm70_active_expert_ids = torch.empty( + max_slots, dtype=torch.int32, device=device + ) + + @staticmethod + def _persistent_buffers( + layer: RoutedExperts, num_tokens: int + ) -> dict[str, torch.Tensor]: + 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], + "gate_up": layer._nvfp4_sm70_gate_up[:slots], + "intermediate": layer._nvfp4_sm70_intermediate[:slots], + "sorted_output": layer._nvfp4_sm70_sorted_output[:slots], + "expert_offsets": layer._nvfp4_sm70_expert_offsets, + "expert_offsets64": layer._nvfp4_sm70_expert_offsets64, + "inv_permuted_idx": layer._nvfp4_sm70_inv_permuted_idx[:num_tokens], + "topk_ids": layer._nvfp4_sm70_topk_ids[:num_tokens], + "token_expert_indices": ( + layer._nvfp4_sm70_token_expert_indices[:num_tokens] + ), + "permuted_idx": layer._nvfp4_sm70_permuted_idx[:slots], + "sort_workspace": layer._nvfp4_sm70_sort_workspace, + "permuted_experts_id": layer._nvfp4_sm70_permuted_experts_id[:slots], + "sorted_row_idx": layer._nvfp4_sm70_sorted_row_idx[:slots], + "topk_ids_for_sort": layer._nvfp4_sm70_topk_ids_for_sort[:slots], + "dense_expert_ids": layer._nvfp4_sm70_dense_expert_ids, + "compact_offsets": layer._nvfp4_sm70_compact_offsets[: slots + 1], + "active_expert_ids": layer._nvfp4_sm70_active_expert_ids[:slots], + } + + @staticmethod + def _eager_buffers( + layer: RoutedExperts, num_tokens: int + ) -> dict[str, torch.Tensor]: + device = layer.w13_tm_weight.device + 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) + workspace_size = torch.ops._moe_C.moe_permute_sort_workspace_size( + slots, layer.global_num_experts + ) + return { + "output": torch.empty( + num_tokens, hidden, dtype=torch.float16, device=device + ), + "permuted_input": torch.empty( + slots, hidden, dtype=torch.float16, device=device + ), + "gate_up": torch.empty( + slots, 2 * intermediate, dtype=torch.float16, device=device + ), + "intermediate": torch.empty( + slots, intermediate, dtype=torch.float16, device=device + ), + "sorted_output": torch.empty( + slots, hidden, dtype=torch.float16, device=device + ), + "expert_offsets": torch.empty( + experts + 1, dtype=torch.int32, device=device + ), + "expert_offsets64": torch.empty( + experts + 1, dtype=torch.int64, device=device + ), + "inv_permuted_idx": torch.empty( + num_tokens, top_k, dtype=torch.int32, device=device + ), + "topk_ids": torch.empty( + num_tokens, top_k, dtype=torch.int32, device=device + ), + "token_expert_indices": torch.arange( + slots, dtype=torch.int32, device=device + ).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 + ), + "permuted_experts_id": torch.empty(slots, dtype=torch.int32, device=device), + "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), + "active_expert_ids": torch.empty(slots, dtype=torch.int32, device=device), + } + + def _get_buffers( + self, layer: RoutedExperts, num_tokens: int + ) -> dict[str, torch.Tensor]: + if 0 < num_tokens <= _GRAPH_SAFE_MAX_TOKENS: + return self._persistent_buffers(layer, num_tokens) + return self._eager_buffers(layer, num_tokens) + + @staticmethod + def _apply_swiglu( + layer: RoutedExperts, out: torch.Tensor, gate_up: torch.Tensor + ) -> None: + if layer.swiglu_limit is None: + torch.ops._C.silu_and_mul(out, gate_up) + else: + torch.ops._C.silu_and_mul_with_clamp( + out, gate_up, float(layer.swiglu_limit) + ) + + def apply( + self, + layer: RoutedExperts, + x: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + shared_experts: SharedExperts | None, + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor: + del shared_experts, shared_experts_input + if not x.is_cuda or x.dtype != torch.float16 or x.ndim != 2: + 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.") + 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"{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_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: + raise TypeError("SM70 NVFP4 MoE requires float32 top-k weights.") + + num_tokens = x.shape[0] + if num_tokens == 0: + return x.new_empty((0, hidden)) + buffers = self._get_buffers(layer, num_tokens) + output = buffers["output"] + output.zero_() + slots = num_tokens * 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, + 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: + _prepare_compact_slot_groups( + buffers["permuted_experts_id"], + buffers["compact_offsets"], + buffers["active_expert_ids"], + ) + stage_offsets = buffers["compact_offsets"] + stage_expert_ids = buffers["active_expert_ids"] + stage_experts = slots + else: + stage_offsets = buffers["expert_offsets"] + stage_expert_ids = buffers["dense_expert_ids"] + stage_experts = int(layer.sm70_nvfp4_num_experts) + + sm70_ops.nvfp4_moe_dense_stage_sm70_out( + buffers["gate_up"], + buffers["permuted_input"], + stage_offsets, + stage_expert_ids, + layer.w13_strided_ptrs_w, + layer.w13_strided_ptrs_s, + stage_experts, + layer.sm70_nvfp4_w13_k_dim, + layer.sm70_nvfp4_w13_n_dim, + layer.sm70_nvfp4_group_size, + ) + self._apply_swiglu(layer, buffers["intermediate"], buffers["gate_up"]) + sm70_ops.nvfp4_moe_dense_stage_sm70_out( + buffers["sorted_output"], + buffers["intermediate"], + stage_offsets, + stage_expert_ids, + layer.w2_strided_ptrs_w, + layer.w2_strided_ptrs_s, + stage_experts, + layer.sm70_nvfp4_w2_k_dim, + 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"], + top_k, + output, + ) + return output + + def apply_monolithic( + self, + layer: RoutedExperts, + x: torch.Tensor, + router_logits: torch.Tensor, + input_ids: torch.Tensor | None = None, + ) -> torch.Tensor: + del layer, x, router_logits, input_ids + raise NotImplementedError("SM70 NVFP4 MoE is not a monolithic route.") + + def get_fused_moe_quant_config( # type: ignore[override] + self, layer: RoutedExperts + ) -> FusedMoEQuantConfig | None: + del layer + return None diff --git a/vllm/model_executor/layers/quantization/sm70_turbomind.py b/vllm/model_executor/layers/quantization/sm70_turbomind.py index d948a94a7f..faf91ee4ee 100644 --- a/vllm/model_executor/layers/quantization/sm70_turbomind.py +++ b/vllm/model_executor/layers/quantization/sm70_turbomind.py @@ -64,6 +64,13 @@ def should_use_mxfp4_moe_turbomind() -> bool: ) +def should_use_nvfp4_moe_turbomind() -> bool: + """Select the native NVFP4 MoE path only on exact SM70.""" + return is_exact_sm70_cuda_platform() and use_turbomind( + envs.VLLM_SM70_NVFP4_TURBOMIND + ) + + def should_prepare_turbomind( tensor: torch.Tensor, default_enabled: bool, 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/models/registry.py b/vllm/model_executor/models/registry.py index 3a7a7b4ba0..bd58453da7 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -122,6 +122,7 @@ "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 +411,10 @@ "Glm4vForConditionalGeneration": ("glm4_1v", "Glm4vForConditionalGeneration"), "Glm4vMoeForConditionalGeneration": ("glm4_1v", "Glm4vMoeForConditionalGeneration"), "GlmOcrForConditionalGeneration": ("glm_ocr", "GlmOcrForConditionalGeneration"), + "Glm5NextForConditionalGeneration": ( + "vllm.models.glm5next", + "Glm5NextForConditionalGeneration", + ), "GraniteSpeechForConditionalGeneration": ( "granite_speech", "GraniteSpeechForConditionalGeneration", @@ -628,6 +633,7 @@ "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"), diff --git a/vllm/model_executor/models/utils.py b/vllm/model_executor/models/utils.py index 095d0e363d..2c8e74afb3 100644 --- a/vllm/model_executor/models/utils.py +++ b/vllm/model_executor/models/utils.py @@ -46,6 +46,9 @@ 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) @@ -53,11 +56,21 @@ def __or__(self, other: "WeightsMapper") -> "WeightsMapper": """Combine two `WeightsMapper`s by merging their mappings.""" return WeightsMapper( 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 +85,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 +105,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 [ 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/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..f30912ccf7 --- /dev/null +++ b/vllm/models/glm5next/__init__.py @@ -0,0 +1,18 @@ +# 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 SM90 and above 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..327014ae20 --- /dev/null +++ b/vllm/models/glm5next/sm70/sparse.py @@ -0,0 +1,282 @@ +# 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.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[str]] = [ + "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, "model_type", None) != "glm5_next_text": + return "GLM5_SM70_SPARSE is restricted to glm5_next_text" + 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 + self.topk_indices_buffer = indexer.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 = [ + ((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/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/transformers_utils/config.py b/vllm/transformers_utils/config.py index f940739a96..7d657b5fb6 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", diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py index 5998e61dfd..17ae16dc54 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", @@ -106,6 +109,9 @@ "FunAudioChatConfig", "FunAudioChatAudioEncoderConfig", "Granite4VisionConfig", + "Glm5NextConfig", + "Glm5NextTextConfig", + "Glm5NextVisionConfig", "HunYuanVLConfig", "HunYuanVLTextConfig", "HunYuanVLVisionConfig", 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/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 b1cb946d33..7d07bf5b00 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -557,6 +557,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/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 2870ec9a15..88d91230b7 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 @@ -216,6 +280,113 @@ 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, + 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. @@ -230,6 +401,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 @@ -242,7 +414,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. @@ -288,6 +460,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. @@ -304,14 +485,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, @@ -343,6 +525,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, @@ -491,12 +677,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, @@ -541,7 +742,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, @@ -549,7 +750,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: @@ -558,6 +771,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] ) @@ -566,6 +785,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 @@ -584,6 +811,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: @@ -622,6 +870,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/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index c5e8953745..74a17ab44f 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -425,9 +425,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.participates_in_prefix_caching + ] 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." @@ -443,6 +448,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.participates_in_prefix_caching: + continue manager_cls = self.single_type_managers[i].__class__ spec = g.kv_cache_spec diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 7f3a5e4fdf..4ef70359c1 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -23,6 +23,7 @@ ChunkedLocalAttentionSpec, FullAttentionSpec, HiddenStateCacheSpec, + KpoolTailSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheSpec, @@ -916,6 +917,9 @@ 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 all( isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) for g in kv_cache_groups ): @@ -1233,6 +1237,211 @@ 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_from_groups( vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec], @@ -1277,6 +1486,39 @@ 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 all( isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) for group in kv_cache_groups @@ -1654,6 +1896,10 @@ def get_kv_cache_groups( kv_cache_groups = _get_kv_cache_groups_uniform_groups(grouped_specs) _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 # Pull HiddenStateCacheSpec layers out before the general multi-group # path so they don't affect page-size unification or grouping. @@ -1787,6 +2033,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/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index cd000dc849..07f546df24 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -17,6 +17,7 @@ CrossAttentionSpec, FullAttentionSpec, HiddenStateCacheSpec, + KpoolTailSpec, KVCacheSpec, MambaSpec, MLAAttentionSpec, @@ -697,6 +698,101 @@ 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, total_computed_tokens: int + ) -> None: + del request_id, total_computed_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) @@ -1212,6 +1308,7 @@ def __init__( HiddenStateCacheSpec: FullAttentionManager, SlidingWindowSpec: SlidingWindowManager, SlidingWindowMLASpec: SlidingWindowManager, + KpoolTailSpec: KpoolTailManager, ChunkedLocalAttentionSpec: ChunkedLocalAttentionManager, MambaSpec: MambaManager, CrossAttentionSpec: CrossAttentionManager, diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 4e094eb6a6..af0692e916 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.participates_in_prefix_caching + ] 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 31ee89bc72..a1b1f0419e 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 participates_in_prefix_caching(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. @@ -559,6 +574,29 @@ 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 participates_in_prefix_caching(self) -> bool: + return False + + @dataclass(frozen=True) class MambaSpec(KVCacheSpec): shapes: tuple[tuple[int, ...], ...] @@ -569,11 +607,15 @@ class MambaSpec(KVCacheSpec): num_speculative_blocks: int = 0 @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 @@ -674,6 +716,12 @@ class UniformTypeKVCacheSpecs(KVCacheSpec): kv_cache_specs: dict[str, KVCacheSpec] + @property + def participates_in_prefix_caching(self) -> bool: + return all( + spec.participates_in_prefix_caching 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()) 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_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index e6dbb57163..2b7ad07f95 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -253,6 +253,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, ) @@ -671,6 +672,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: @@ -10209,7 +10222,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. @@ -10683,16 +10699,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( @@ -10961,11 +10980,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) @@ -11850,15 +11875,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/utils.py b/vllm/v1/worker/utils.py index fb02fdc935..f3fc90e180 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,17 +33,26 @@ MambaSpec, UniformTypeKVCacheSpecs, ) +from vllm.v1.worker.block_table import get_block_table_width logger = init_logger(__name__) -@triton.jit(do_not_specialize=["n_blocks"]) +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 + + +@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. @@ -58,23 +67,24 @@ def _zero_kv_blocks_kernel( Programs are mapped as (block_index, seg_index, chunk_index). """ - 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) + block_stride_el = tl.load(seg_block_strides_ptr + seg_index) + 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) ptr = tl.cast(seg_addr, tl.pointer_type(tl.int32)) - offset = ( - block_id.to(tl.int64) * PAGE_SIZE_EL + 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: @@ -88,7 +98,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, 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 @@ -117,9 +129,20 @@ 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_segments: dict[tuple[int, int], int] = {} seg_addrs: list[int] = [] - page_size_el: int | None = None + seg_block_strides: list[int] = [] + seg_page_sizes: list[int] = [] + + def add_segment(addr: int, block_stride_el: int, page_size_el: int) -> None: + key = (addr, block_stride_el) + if (idx := seen_segments.get(key)) is not None: + seg_page_sizes[idx] = max(seg_page_sizes[idx], page_size_el) + return + seen_segments[key] = len(seg_addrs) + seg_addrs.append(addr) + seg_block_strides.append(block_stride_el) + seg_page_sizes.append(page_size_el) for group in attn_groups_iter: spec = group.kv_cache_spec @@ -142,22 +165,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 " - f"{cur_page_el}" - ) block_stride_bytes = cur_bytes outer_dims = [ @@ -168,7 +181,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, cur_page_el) elif isinstance(spec, MambaSpec): if not isinstance(kv, (list, tuple)) or not kv: continue @@ -176,27 +189,18 @@ 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 " - f"{cur_page_el}" - ) - seg_addrs.append(dp) - - if not seg_addrs or page_size_el is None: + add_segment(dp, cur_page_el, 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) self._id_cap = 8192 self._ids_pinned = torch.empty( self._id_cap, @@ -206,7 +210,9 @@ def init_meta( self._ids_gpu = torch.empty(self._id_cap, dtype=torch.int64, device=self.device) self._meta = ( torch.tensor(seg_addrs, dtype=torch.uint64, device=self.device), - page_size_el, + torch.tensor(seg_block_strides, dtype=torch.int64, device=self.device), + torch.tensor(seg_page_sizes, dtype=torch.int64, device=self.device), + (max_page_size_el + blk_size - 1) // blk_size, blk_size, len(seg_addrs), ) @@ -215,7 +221,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, 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 @@ -231,13 +244,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, ) @@ -269,20 +281,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 From a21d15e88092a4678aff62605828df63551eb69c Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 11:37:09 +0800 Subject: [PATCH 065/105] [Spec Decode][SM70] Handle Qwen4Exp MTP feedback states Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/v1/worker/test_qwen4_exp_v2.py | 37 +++++++++++++++++++ .../gpu/spec_decode/eagle/speculator.py | 8 ++-- 2 files changed, 42 insertions(+), 3 deletions(-) diff --git a/tests/v1/worker/test_qwen4_exp_v2.py b/tests/v1/worker/test_qwen4_exp_v2.py index 13e9ee95bb..29ab4fe9bb 100644 --- a/tests/v1/worker/test_qwen4_exp_v2.py +++ b/tests/v1/worker/test_qwen4_exp_v2.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from contextlib import nullcontext from types import SimpleNamespace import pytest @@ -16,6 +17,42 @@ ) 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_qsa_circular_group_uses_one_block_and_custom_slot_mapping( diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py index 8376a51558..3c4609d357 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py @@ -222,11 +222,13 @@ 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 _sample_draft( From eb9547ee4c397d33586bdb8fd82cff6f6e1131bf Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 11:49:22 +0800 Subject: [PATCH 066/105] [Core] Align Qwen4Exp short-conv metadata Port the applicable runtime hunk from origin/pr-338 commit 4de4737ad3 so CUDA graph capture passes the required bool dtype to async_tensor_h2d and uses the current AttentionSpec interface. Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- vllm/v1/attention/backends/short_conv_attn.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/vllm/v1/attention/backends/short_conv_attn.py b/vllm/v1/attention/backends/short_conv_attn.py index 3e888e69b7..91e3e7bcbc 100644 --- a/vllm/v1/attention/backends/short_conv_attn.py +++ b/vllm/v1/attention/backends/short_conv_attn.py @@ -21,7 +21,7 @@ compute_causal_conv1d_metadata, mamba_get_block_table_tensor, ) -from vllm.v1.kv_cache_interface import MambaSpec +from vllm.v1.kv_cache_interface import AttentionSpec class ShortConvAttentionBackend(AttentionBackend): @@ -111,7 +111,7 @@ class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder): def __init__( self, - kv_cache_spec: MambaSpec, + kv_cache_spec: AttentionSpec, layer_names: list[str], vllm_config: VllmConfig, device: torch.device, @@ -284,7 +284,9 @@ def build( # type: ignore[override] spec_sequence_masks = spec_sequence_masks_cpu else: spec_sequence_masks = async_tensor_h2d( - spec_sequence_masks_cpu, device=query_start_loc.device + spec_sequence_masks_cpu, + dtype=torch.bool, + device=query_start_loc.device, ) # For causal_conv1d (non-spec prefill Triton kernel metadata). From 6097689ea3f4ef004c36ec8a66ed456f0fa5431f Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 11:58:32 +0800 Subject: [PATCH 067/105] [Bugfix] Stage Qwen4Exp spec masks from lists Convert the CPU bool tensor to the list contract required by async_tensor_h2d. This preserves the pinned non-blocking copy and fixes full CUDA graph capture after the PR 338 API alignment. Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- vllm/v1/attention/backends/short_conv_attn.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/v1/attention/backends/short_conv_attn.py b/vllm/v1/attention/backends/short_conv_attn.py index 91e3e7bcbc..efbbd1e075 100644 --- a/vllm/v1/attention/backends/short_conv_attn.py +++ b/vllm/v1/attention/backends/short_conv_attn.py @@ -284,7 +284,7 @@ def build( # type: ignore[override] spec_sequence_masks = spec_sequence_masks_cpu else: spec_sequence_masks = async_tensor_h2d( - spec_sequence_masks_cpu, + spec_sequence_masks_cpu.tolist(), dtype=torch.bool, device=query_start_loc.device, ) From 96ffb203327f8c2df7186931572e68d153a6548e Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 13:04:57 +0800 Subject: [PATCH 068/105] [Bugfix][SM70] Reuse prebuilt FlashQLA extensions Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../gated_delta_rule/chunk/sm70/fused_fwd.py | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) 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 From d9a39ea43477e092a44cd7e7b4e21d5a25cf6cb6 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 16:28:50 +0800 Subject: [PATCH 069/105] [Core] Support heterogeneous Mamba state groups Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/v1/worker/test_mamba_utils.py | 94 ++++++++++- vllm/v1/worker/gpu/mamba_align.py | 6 +- .../worker/gpu/model_states/mamba_hybrid.py | 35 ++-- vllm/v1/worker/mamba_utils.py | 155 ++++++++++++------ 4 files changed, 223 insertions(+), 67 deletions(-) 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/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_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 2c55f6ac35..f38703a49c 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -52,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__) @@ -194,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( @@ -212,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." @@ -225,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, @@ -236,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/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( From b6ded5b8d4b701d9cb71fbfb2c864b8038e342e7 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 17:00:20 +0800 Subject: [PATCH 070/105] [Doc] Record Qwen4Exp MTP4 bring-up Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_qwen38_flash_next_nvfp4.md | 75 +++++++++++++++++++-- 1 file changed, 68 insertions(+), 7 deletions(-) diff --git a/docs/design/sm70_qwen38_flash_next_nvfp4.md b/docs/design/sm70_qwen38_flash_next_nvfp4.md index 54191ae5c1..169425e151 100644 --- a/docs/design/sm70_qwen38_flash_next_nvfp4.md +++ b/docs/design/sm70_qwen38_flash_next_nvfp4.md @@ -2,9 +2,12 @@ ## Status and ownership -- Status: source bring-up implemented; CPU/configuration gates pass and the - local ModelScope snapshot is fully verified. Full-model load, output quality, - measured memory, and speed are not claimed yet. +- Status: source bring-up, native Qwen4Exp MTP, and prefix-cache configuration + are implemented; focused CPU/configuration gates pass and the local + ModelScope snapshot is fully verified. Native MTP4 now completes TP4 model + load, graph capture, warmup, and two 1024x256 requests. Acceptance, repeated + token equality, memory, and steady decode are measured below; matched no-MTP + token equality and verifier cost remain pending. - Integration line: `private/main`. - Base SHA: `d63e9490f65f9e01f6649053c1ab72922034b931`. - Model: `RadixArk/Qwen3.8-Flash-Next-NVFP4` at revision @@ -48,8 +51,18 @@ The first correctness route deliberately excludes speculative decoding. 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: disabled for the first route. The fixed QSA ring manager does - not yet implement reusable prefix blocks. +- 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 @@ -107,6 +120,39 @@ 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. +## 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. + ## Acceptance gates 1. Static route: Transformers config, model registry/processor registration, @@ -150,13 +196,28 @@ only after profiles identify them as measured decode bottlenecks. 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 and full TP4 load - remains pending. + 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, From 31bd61db9c06c739832b6f7657b3469dd35adbbe Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 17:33:15 +0800 Subject: [PATCH 071/105] [Doc] Record Qwen3.8 Flash Next prefill Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_v100_migration_control.md | 49 ++++++++++++++++++++++ 1 file changed, 49 insertions(+) diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index dba85d9161..d2d33ea06d 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -24,6 +24,55 @@ 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. +- 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 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`. From 2366d547e1ce1cff2ddd771e5dddf2ba8033e55c Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 19:18:55 +0800 Subject: [PATCH 072/105] [Perf][SM70] Narrow Qwen3.8 QSA prefill tiles Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_v100_migration_control.md | 19 ++++++++++++++++++- tests/models/qwen4_exp/test_qsa_ops.py | 7 ++++--- vllm/models/qwen4_exp/nvidia/ops/qsa.py | 7 +++++-- 3 files changed, 27 insertions(+), 6 deletions(-) diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index d2d33ea06d..9a5c442c3f 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -53,6 +53,21 @@ Goal: `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. - 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 @@ -66,7 +81,9 @@ Goal: - 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 corresponding PR evidence is + `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 diff --git a/tests/models/qwen4_exp/test_qsa_ops.py b/tests/models/qwen4_exp/test_qsa_ops.py index 0026f2d59e..11f4de01c7 100644 --- a/tests/models/qwen4_exp/test_qsa_ops.py +++ b/tests/models/qwen4_exp/test_qsa_ops.py @@ -6,9 +6,10 @@ ) -def test_sm70_qsa_prefill_uses_four_warps(): - assert _qsa_sparse_launch_profile(512, 8, True) == (64, 4, 4) - assert _qsa_sparse_launch_profile(8192, 8, True) == (64, 1, 4) +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(): diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py index 2052c98c26..5cd2f60d3a 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py @@ -966,9 +966,12 @@ def _qsa_sparse_launch_profile( 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 - # preserve the exact reduction order while restoring warp-level - # parallelism for both split and non-split prefill profiles. + # 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 From a271361cac0b225c453800784dec728f16a8594a Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 19:30:07 +0800 Subject: [PATCH 073/105] [Doc] Record Qwen3.8 prefill follow-up boundary Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_v100_migration_control.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index 9a5c442c3f..1425b153b6 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -68,6 +68,12 @@ Goal: 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. +- 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 From 70b63a1e4502627285cb436aa71ca459da5b14c4 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 19:32:06 +0800 Subject: [PATCH 074/105] [Spec Decode] Reduce Qwen4Exp MTP cost Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/engine/test_arg_utils.py | 4 +- tests/v1/worker/test_qwen4_exp_v2.py | 57 ++++++++++ vllm/config/speculative.py | 24 ++++ vllm/engine/arg_utils.py | 7 +- vllm/models/qwen4_exp/nvidia/mtp.py | 5 + .../gpu/spec_decode/eagle/speculator.py | 107 ++++++++++++++---- 6 files changed, 176 insertions(+), 28 deletions(-) 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/v1/worker/test_qwen4_exp_v2.py b/tests/v1/worker/test_qwen4_exp_v2.py index 29ab4fe9bb..1e6d9d5af2 100644 --- a/tests/v1/worker/test_qwen4_exp_v2.py +++ b/tests/v1/worker/test_qwen4_exp_v2.py @@ -55,6 +55,63 @@ def test_qwen4_exp_mtp_v2_unpacks_logits_and_feedback_hidden_states( 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: diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 07507b9be4..9d18c98185 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -139,6 +139,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 @@ -368,6 +371,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 @@ -531,6 +543,9 @@ def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig: "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 + ), } ) @@ -986,6 +1001,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: 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/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py index 87599d99f8..3945ebc663 100644 --- a/vllm/models/qwen4_exp/nvidia/mtp.py +++ b/vllm/models/qwen4_exp/nvidia/mtp.py @@ -442,6 +442,11 @@ def compute_logits( ) -> 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: diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py index 3c4609d357..07afbd9c04 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py @@ -65,6 +65,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 +161,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, @@ -231,15 +246,35 @@ def run_model( 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( @@ -253,8 +288,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, @@ -279,10 +332,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, @@ -364,9 +416,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, @@ -434,11 +485,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 @@ -447,15 +500,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( @@ -550,6 +607,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 @@ -567,6 +625,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. @@ -595,12 +654,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] From 5378eb48789e60347751636c92cb4c0c4915cab4 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 22:08:56 +0800 Subject: [PATCH 075/105] [Bugfix] Harden Qwen3 coder tool parsing Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../test_qwen3coder_tool_parser.py | 190 ++++++++ vllm/tool_parsers/qwen3coder_tool_parser.py | 424 ++++++++++++++---- 2 files changed, 529 insertions(+), 85 deletions(-) diff --git a/tests/tool_parsers/test_qwen3coder_tool_parser.py b/tests/tool_parsers/test_qwen3coder_tool_parser.py index 9167493688..2df97f2536 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,6 +1537,40 @@ 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", + } + + @pytest.mark.parametrize( "deltas, expected_content", [ diff --git a/vllm/tool_parsers/qwen3coder_tool_parser.py b/vllm/tool_parsers/qwen3coder_tool_parser.py index 7a3a63ad53..6ac6b53b5f 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 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, @@ -575,6 +857,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( @@ -591,39 +881,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=[ @@ -634,10 +892,6 @@ def extract_tool_calls_streaming( ] ) - self.in_function = False - self.json_closed = True - self.accumulated_params = {} - return result return None From f5a93e2de2f1940de0e498afb01db85b77982948 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 22:09:03 +0800 Subject: [PATCH 076/105] [Bench] Add tool protocol quality gates Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- benchmarks/benchmark_sm70_tool_datasets.py | 546 +++++++++ benchmarks/benchmark_sm70_tool_protocol.py | 1291 ++++++++++++++++++++ 2 files changed, 1837 insertions(+) create mode 100644 benchmarks/benchmark_sm70_tool_datasets.py create mode 100644 benchmarks/benchmark_sm70_tool_protocol.py diff --git a/benchmarks/benchmark_sm70_tool_datasets.py b/benchmarks/benchmark_sm70_tool_datasets.py new file mode 100644 index 0000000000..095cfb2702 --- /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 re +import time +from pathlib import Path +from typing import Any + +import pandas as pd +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..83bfb648a2 --- /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 re +import time +from pathlib import Path +from typing import Any + +import jsonschema +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() From 5abeaa617754090ae20dc225134e310f729f0d2f Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 22:16:54 +0800 Subject: [PATCH 077/105] [Bench] Record per-prompt speculative acceptance Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- benchmarks/benchmark_sm70_model_tokens.py | 80 ++++++++++++++++++- .../benchmarks/test_benchmark_sm70_decode.py | 32 +++++++- 2 files changed, 110 insertions(+), 2 deletions(-) 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/tests/benchmarks/test_benchmark_sm70_decode.py b/tests/benchmarks/test_benchmark_sm70_decode.py index b9fe6ebb78..0895a8ab17 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,33 @@ 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_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 From 58beb3433b01320d1efb6e9d63f8a7485ef4ee89 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 22:24:17 +0800 Subject: [PATCH 078/105] Revert "[Release] Restore public source snapshot to v1.3.0 (#329)" This reverts commit f4c3451a2b2205dead8b0a62c968f7303f1cc702. Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .buildkite/test-amd.yaml | 2 +- .buildkite/test_areas/kernels.yaml | 2 +- .pre-commit-config.yaml | 12 +- CMakeLists.txt | 42 + README.md | 65 +- .../analyze_dflash2_selector_alignment.py | 517 ++ benchmarks/analyze_dsv4_dspark_confidence.py | 288 + benchmarks/benchmark_dspark_api.py | 1 + benchmarks/benchmark_dspark_stream.py | 1 + .../benchmark_dsv4_dspark_long_context.py | 576 ++ benchmarks/benchmark_dsv4_gsm8k_api.py | 231 + benchmarks/benchmark_dsv4_quality_api.py | 635 ++ .../benchmark_sm70_cublas_gemm_algorithms.py | 251 + .../benchmark_sm70_custom_all_reduce.py | 70 +- benchmarks/benchmark_sm70_d256_exact_dense.py | 215 + benchmarks/benchmark_sm70_decode.py | 17 +- .../benchmark_sm70_dflash2_gdn_metadata.py | 373 + benchmarks/benchmark_sm70_dflash2_gsm8k.py | 484 + ...benchmark_sm70_dflash2_sparse_rejection.py | 243 + .../benchmark_sm70_flashinfer_fixed_entry.py | 20 +- ...hmark_sm70_flashqla_indexed_state_cache.py | 48 +- .../benchmark_sm70_fp8_fixed_vs_dynamic.py | 17 +- ...hmark_sm70_fp8_prefill_exact_dense_blas.py | 368 + .../benchmark_sm70_fp8_prefill_tactics.py | 486 + benchmarks/benchmark_sm70_fp8_qpn8.py | 729 ++ benchmarks/benchmark_sm70_gdn_exactness.py | 396 +- benchmarks/benchmark_sm70_hmma_c2a_mincomm.py | 3 +- .../benchmark_sm70_instruction_latency.py | 3 +- benchmarks/benchmark_sm70_lm_head_top1.py | 18 +- .../benchmark_sm70_long_quality_suffix.py | 6 +- benchmarks/benchmark_sm70_model_tokens.py | 113 +- ...benchmark_sm70_mxfp4_moe_active_experts.py | 232 +- ..._sm70_native_bm32_allp_crossblock_micro.py | 60 +- ...ark_sm70_native_bm32_allp_scratch_micro.py | 73 +- ..._sm70_native_bm32_full_phase_allp_micro.py | 64 +- ...hmark_sm70_native_bm32_full_phase_micro.py | 64 +- ...nchmark_sm70_native_bm32_pv_panel_micro.py | 17 +- .../benchmark_sm70_prefill_gemm_collective.py | 395 + benchmarks/benchmark_sm70_raw_hmma_probe.py | 5 +- .../benchmark_sm70_rejection_sampler_micro.py | 34 +- .../benchmark_sm70_sglang_push_chain.py | 162 + benchmarks/benchmark_sm70_source_inventory.py | 88 +- ...chmark_sm70_tp2_compact_verifier_logits.py | 10 +- ...enchmark_sm70_tp4_fused_norm_collective.py | 305 + benchmarks/benchmark_sm70_turbomind_parity.py | 33 +- .../benchmark_sm70_wmma_fragment_probe.py | 38 +- benchmarks/compare_dsv4_quality_results.py | 425 + benchmarks/compare_sm70_tensor_dumps.py | 156 +- .../kernels/benchmark_device_communicators.py | 74 +- benchmarks/kernels/benchmark_dspark_markov.py | 1 + .../kernels/benchmark_dspark_metadata.py | 1 + .../benchmark_dspark_sm70_attention.py | 1 + .../kernels/benchmark_flash_v100_fp8_kv.py | 1 + .../benchmark_sm70_fp8_graph_capture.py | 1 + .../kernels/benchmark_sm70_nvfp4_qpn2.py | 499 ++ .../profile_sm70_fp8_qpn8_gated_pair_ncu.py | 140 + benchmarks/profile_sm70_fp8_qpn8_ncu.py | 194 + benchmarks/verify_sm70_artifact.py | 168 +- cmake/external_projects/vllm_flash_attn.cmake | 6 + .../sm70_flash_attn_d256_gqa_arch.patch | 1673 ++++ .../sm70_flash_attn_d256_k_pingpong.patch | 90 + csrc/custom_all_reduce.cu | 125 +- csrc/custom_all_reduce.cuh | 775 +- csrc/libtorch_stable/layernorm_kernels.cu | 27 +- .../marlin_moe_wna16/sm70_marlin_fp8_gemm.cu | 197 +- .../moe/marlin_moe_wna16/sm70_marlin_gemm.cuh | 1517 ++-- .../sm70_marlin_moe_dispatch.cu | 98 +- .../sm70_marlin_mxfp4_gemm.cu | 128 +- .../sm70_marlin_nvfp4_gemm.cu | 128 +- .../marlin_moe_wna16/sm70_marlin_u4_gemm.cu | 189 +- .../marlin_moe_wna16/sm70_marlin_u4b8_gemm.cu | 167 +- .../marlin_moe_wna16/sm70_marlin_u8_gemm.cu | 211 +- .../sm70_marlin_u8b128_gemm.cu | 191 +- csrc/moe/moe_permute_unpermute_op.cu | 17 +- .../moe_permute_unpermute_kernel.cu | 27 +- .../moe_permute_unpermute_kernel.h | 11 +- csrc/ops.h | 675 +- .../marlin/sm70_awq_marlin_repack.cu | 32 +- .../marlin/sm70_gptq_marlin_repack.cu | 34 +- .../marlin/sm70_marlin_common.cuh | 600 +- .../marlin/sm70_marlin_dispatch.cu | 129 +- .../marlin/sm70_marlin_fp8_gemm.cu | 220 +- csrc/quantization/marlin/sm70_marlin_gemm.cuh | 172 +- .../marlin/sm70_marlin_iterator_utils.cuh | 34 +- csrc/quantization/marlin/sm70_marlin_mma.cuh | 45 +- .../marlin/sm70_marlin_mxfp4_gemm.cu | 212 +- .../marlin/sm70_marlin_nvfp4_gemm.cu | 210 +- .../marlin/sm70_marlin_splitk.cuh | 50 +- .../marlin/sm70_marlin_u4_gemm.cu | 280 +- .../marlin/sm70_marlin_u4b8_gemm.cu | 218 +- .../marlin/sm70_marlin_u8_gemm.cu | 251 +- .../marlin/sm70_marlin_u8b128_gemm.cu | 222 +- csrc/sm70_tile_runtime_signal.cuh | 19 +- .../kernels/attention/quantization.h | 1373 ++- .../kernels/gemm/arch/config_sm70_s884.h | 219 +- .../src/turbomind/kernels/gemm/gemm.cu | 1344 +-- .../turbomind/kernels/gemm/iterator_sm70.h | 406 +- .../kernels/gemm/kernel/sm70_884_4.cu | 74 +- .../kernels/gemm/kernel/sm70_884_8.cu | 57 +- .../src/turbomind/kernels/gemm/transform.h | 286 +- .../src/turbomind/kernels/gemm/types.h | 355 +- csrc/sm70_turbomind/ops/LICENSE.v100-skinny | 21 + csrc/sm70_turbomind/ops/awq_sm70_gemm.cu | 1156 ++- csrc/sm70_turbomind/ops/fp8_qpn8_sm70.cu | 1255 +++ csrc/sm70_turbomind/ops/nvfp4_qpn2_sm70.cu | 546 ++ csrc/sm70_turbomind/ops/nvfp4_qpn4_sm70.cu | 1029 +++ .../ops/qwen38_prefill_cutlass.cu | 130 + .../ops/qwen38_prefill_cutlass.cuh | 13 + .../ops/sm70_top20_philox_fragment.cpp | 52 + .../ops/sm70_top20_philox_sampler.cu | 289 + csrc/sm70_turbomind/ops/tm_registry_sm70.cu | 75 +- csrc/torch_bindings.cpp | 198 +- docs/assets/wechat-group-qr-7.png | Bin 0 -> 529331 bytes ...6_27b_awq_mtp_target_verifier_fastpaths.md | 36 +- .../sm70_deepseek_v4_dspark_n7_latest_main.md | 28 + ...ek_v4_dspark_verifier_20ms_long_context.md | 217 + docs/design/sm70_deepseek_v4_flash_bringup.md | 34 + docs/design/sm70_dflash2_acceptance_rd.md | 77 + docs/design/sm70_dflash2_long_verify_decay.md | 465 + docs/design/sm70_dflash2_ngram_hybrid.md | 130 + docs/design/sm70_dflash2_nvfp4_17ms.md | 220 + docs/design/sm70_dflash2_target_graph_20ms.md | 1229 +++ ...0_nvfp4_turbomind_operator_optimization.md | 10 +- docs/design/sm70_prefix_anchored_swa.md | 57 + docs/design/sm70_qwen38_130_acceptance.md | 26 +- docs/design/sm70_qwen38_fp8_prefill_5500.md | 142 + docs/design/sm70_qwen38_fp8_prefill_decay.md | 307 + docs/design/sm70_qwen38_nvfp4_decode.md | 711 ++ docs/design/sm70_qwen38_qpn8_decode.md | 245 + docs/design/sm70_tile_runtime_exploration.md | 20 +- docs/design/sm70_v100_migration_control.md | 2079 +++++ docs/features/speculative_decoding/README.md | 27 + docs/models/supported_models.md | 8 + .../flash_attn_v100/__init__.py | 2 + .../flash_attn_v100/flash_attn_interface.py | 153 +- flash-attention-v100/include/fused_mha.h | 309 +- flash-attention-v100/include/fused_mma.h | 626 +- .../kernel/contiguous_to_paged.cu | 123 +- .../kernel/flash_decode_paged.cu | 1860 +++- .../kernel/flash_decode_turboquant.cu | 365 +- .../kernel/flash_v100_traits.cuh | 227 +- flash-attention-v100/kernel/fp8_kv_utils.cuh | 38 +- flash-attention-v100/kernel/fused_mha_api.cpp | 66 +- .../kernel/fused_mha_backward.cu | 2407 ++--- .../kernel/fused_mha_forward.cu | 2030 +++-- .../kernel/fused_mha_forward_paged.cu | 7866 ++++++++--------- .../kernel/fused_mha_forward_paged_full.cu | 1082 +-- .../kernel/paged_kv_utils.cuh | 98 +- .../kernel/paged_to_contiguous.cu | 355 +- .../kernel/paged_to_contiguous_fixed.cu | 167 +- .../kernel/paged_to_contiguous_old.cu | 140 +- requirements/common.txt | 26 +- requirements/cpu.txt | 2 +- requirements/cuda.txt | 8 +- requirements/docs.txt | 8 +- requirements/rocm.txt | 12 +- requirements/test/cuda.in | 22 +- requirements/test/cuda.txt | 93 +- requirements/test/rocm.in | 22 +- requirements/test/rocm.txt | 87 +- requirements/test/xpu.txt | 25 +- requirements/tpu.txt | 2 +- requirements/xpu.txt | 2 +- setup.py | 14 +- tests/benchmarks/test_dsv4_quality_tools.py | 141 + tests/config/test_prefix_anchored_swa.py | 158 + tests/conftest.py | 2 +- .../test_custom_all_reduce_allocator.py | 224 + .../test_static_tensor_transfer.py | 63 + tests/engine/test_arg_utils.py | 84 +- .../test_sm70_flash_v100_anchored_swa.py | 288 + .../test_sm70_flash_v100_decode_planner.py | 64 + .../test_sm70_flash_v100_grouped_verify.py | 245 + .../test_sm70_flash_v100_smallq_metadata.py | 232 + .../test_sm70_flash_v100_splitkv3.py | 12 +- .../test_sm70_flash_v100_varlen_layout.py | 330 + .../core/test_batched_weight_rms_norm.py | 63 + .../core/test_sm70_dflash2_gemma_rms.py | 121 + tests/kernels/mamba/test_mrv2_mamba_align.py | 196 + .../kernels/moe/test_moe_permute_unpermute.py | 13 +- .../moe/test_sm70_unquantized_moe_config.py | 93 + tests/kernels/test_compressor_kv_cache.py | 225 + .../kernels/test_deepseek_v4_sm70_indexer.py | 323 +- ...t_deepseek_v4_sm70_qnorm_rope_kv_insert.py | 47 +- .../test_fused_recurrent_packed_decode.py | 12 +- .../test_fused_sigmoid_gating_delta_rule.py | 321 + tests/kernels/test_mhc_sm70_fp16.py | 30 +- .../test_sm70_deepseek_v4_fp16_gemv.py | 122 +- tests/kernels/test_sm70_qwen35_gdn_split.py | 127 + .../test_deepseek_v4_swa_rope.py | 137 + .../test_qwen3_5_quantization.py | 136 + .../test_sm70_fp8_qpn8_pp2_tp4.py | 229 + .../model_executor/test_uva_tower_offload.py | 89 + .../multimodal/processing/test_gemma4.py | 46 + tests/models/test_deepseek_v4_pipeline.py | 50 + tests/models/test_deepseek_v4_sm70_routes.py | 54 + tests/quantization/test_nvfp4_mxfp4_gating.py | 91 + .../test_sm70_compressed_tensors_fastpaths.py | 293 + .../test_sm70_fp8_prefill_exact_dense.py | 733 +- .../test_sm70_modelopt_mixed_nvfp4.py | 192 + .../quantization/test_sm70_modelopt_nvfp4.py | 255 + .../test_sm70_modelopt_nvfp4_kernel_oracle.py | 148 + tests/quantization/test_sm70_mxfp4_moe.py | 197 + tests/quantization/test_sm70_nvfp4_qpn2.py | 143 + .../test_sm70_turbomind_adapter.py | 221 + tests/quantization/test_sm70_warmup.py | 356 + tests/test_envs.py | 68 + .../test_qwen3coder_tool_parser.py | 2 +- .../test_deepseek_ocr_resolution_modes.py | 163 + .../v1/attention/test_gdn_metadata_builder.py | 620 +- .../test_indexer_deepseek_v4_slot_mapping.py | 3 + .../test_prefix_anchored_swa_mask.py | 181 + .../attention/test_sm70_flash_v100_policy.py | 567 +- tests/v1/core/test_async_scheduler.py | 3 + tests/v1/core/test_kv_cache_utils.py | 88 + tests/v1/core/test_prefix_caching.py | 98 +- tests/v1/core/test_scheduler.py | 41 + .../core/test_single_type_kv_cache_manager.py | 107 +- .../v1/core/test_swa_inflight_window_free.py | 351 + tests/v1/core/utils.py | 21 +- tests/v1/cudagraph/test_cudagraph_dispatch.py | 461 +- .../test_sm70_mtp_split_cudagraphs.py | 96 + .../offloading_connector/test_scheduler.py | 201 +- tests/v1/kv_connector/unit/test_config.py | 10 + tests/v1/sample/test_sampler.py | 25 + tests/v1/sample/test_topk_topp_sampler.py | 63 + .../spec_decode/test_ddtree_flex_attention.py | 10 +- .../test_ddtree_gpu_runner_positions.py | 12 +- tests/v1/spec_decode/test_ddtree_metadata.py | 3 +- tests/v1/spec_decode/test_dflash2.py | 1310 +++ .../spec_decode/test_dflash2_ngram_assist.py | 386 + .../test_dflash2_selector_analysis.py | 102 + .../v1/spec_decode/test_dflash_mrv2_config.py | 199 + tests/v1/spec_decode/test_dspark.py | 167 +- tests/v1/spec_decode/test_eagle.py | 11 +- tests/v1/spec_decode/test_mtp.py | 112 + .../test_mtp_draft_prob_alignment.py | 29 +- .../test_rejection_sampler_utils.py | 125 + tests/v1/worker/test_gpu_model_runner.py | 272 + .../test_gpu_model_runner_v2_kv_zero.py | 87 + tests/v1/worker/test_gpu_warmup_blocks.py | 98 + tests/v1/worker/test_gpu_worker_static_pp.py | 153 + tests/v1/worker/test_kv_block_zeroer.py | 161 + tests/v1/worker/test_mamba_utils.py | 34 +- tools/pre_commit/check_forbidden_imports.py | 7 +- tools/pre_commit/check_torch_cuda.py | 10 + vllm/_custom_ops.py | 34 + vllm/_sm70_ops.py | 972 ++ vllm/compilation/cuda_graph.py | 25 +- .../passes/fusion/allreduce_rms_fusion.py | 228 +- .../passes/fusion/matcher_utils.py | 3 +- .../passes/fusion/qk_norm_rope_fusion.py | 3 +- .../passes/fusion/rms_quant_fusion.py | 9 +- .../passes/fusion/sequence_parallelism.py | 75 + .../passes/utility/fix_functionalization.py | 4 +- vllm/config/attention.py | 8 +- vllm/config/offload.py | 3 + vllm/config/speculative.py | 163 +- vllm/config/vllm.py | 282 +- .../device_communicators/cuda_communicator.py | 58 +- .../device_communicators/custom_all_reduce.py | 174 +- .../kv_transfer/kv_connector/v1/__init__.py | 4 + .../kv_transfer/kv_connector/v1/base.py | 15 + .../kv_connector/v1/offloading/scheduler.py | 147 +- .../kv_connector/v1/offloading_connector.py | 11 +- vllm/distributed/parallel_state.py | 119 +- vllm/engine/arg_utils.py | 44 +- vllm/envs.py | 394 +- vllm/forward_context.py | 9 +- vllm/model_executor/kernels/mhc/tilelang.py | 8 +- .../kernels/mhc/tilelang_kernels.py | 4 +- .../layers/attention/attention.py | 77 + .../layers/fla/ops/chunk_scaled_dot_kkt.py | 4 +- .../fla/ops/fused_gdn_prefill_post_conv.py | 3 +- .../layers/fla/ops/fused_recurrent.py | 4 +- .../layers/fla/ops/fused_sigmoid_gating.py | 156 +- .../layers/fused_moe/fused_moe.py | 52 +- .../layers/fused_moe/router/gate_linear.py | 12 +- vllm/model_executor/layers/layernorm.py | 105 + vllm/model_executor/layers/linear.py | 20 +- .../model_executor/layers/logits_processor.py | 100 +- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 701 +- .../layers/quantization/auto_gptq.py | 12 +- .../layers/quantization/awq_marlin.py | 4 +- .../compressed_tensors/compressed_tensors.py | 21 + .../schemes/compressed_tensors_w4a16_nvfp4.py | 3 +- .../schemes/compressed_tensors_w4a4_mxfp4.py | 8 +- .../schemes/compressed_tensors_w4a4_nvfp4.py | 255 +- .../schemes/compressed_tensors_w8a16_fp8.py | 253 +- .../model_executor/layers/quantization/fp8.py | 800 +- .../layers/quantization/input_quant_fp8.py | 4 +- .../layers/quantization/modelopt.py | 159 +- .../layers/quantization/mxfp4_sm70_moe.py | 92 +- .../layers/quantization/nvfp4_sm70_moe.py | 635 ++ .../layers/quantization/sm70_turbomind.py | 179 +- .../quantization/utils/marlin_utils_fp4.py | 5 +- .../layers/quantization/utils/w8a8_utils.py | 9 +- .../layers/sparse_attn_indexer.py | 4 +- .../layers/vocab_parallel_embedding.py | 407 +- vllm/model_executor/models/config.py | 4 +- vllm/model_executor/models/deepseek_ocr.py | 60 +- vllm/model_executor/models/dflash_sm70.py | 300 + vllm/model_executor/models/gemma4_mm.py | 54 +- vllm/model_executor/models/interfaces.py | 10 + vllm/model_executor/models/qwen2_moe.py | 4 +- vllm/model_executor/models/qwen3_5.py | 228 +- vllm/model_executor/models/qwen3_5_mtp.py | 3 +- vllm/model_executor/models/qwen3_dflash.py | 801 +- vllm/model_executor/models/qwen3_dflash2.py | 484 + vllm/model_executor/models/qwen3_next.py | 29 +- vllm/model_executor/models/registry.py | 1 + vllm/model_executor/offloader/base.py | 12 + vllm/model_executor/offloader/uva.py | 46 +- vllm/model_executor/warmup/awq_sm70_warmup.py | 254 +- vllm/models/deepseek_v4/attention.py | 19 +- .../models/deepseek_v4/common/ops/__init__.py | 8 +- .../common/ops/fused_compress_quant_cache.py | 373 +- .../common/ops/save_partial_states.py | 207 + vllm/models/deepseek_v4/compressor.py | 190 +- vllm/models/deepseek_v4/nvidia/dspark.py | 83 +- vllm/models/deepseek_v4/nvidia/model.py | 25 + vllm/models/deepseek_v4/sm70/gemv.py | 282 +- vllm/models/deepseek_v4/sm70/indexer.py | 726 +- .../sm70/qnorm_rope_kv_fp8_insert.py | 132 + vllm/models/deepseek_v4/sm70/sparse.py | 6 + vllm/tool_parsers/qwen3coder_tool_parser.py | 5 +- .../processors/deepseek_ocr.py | 157 +- vllm/utils/cache.py | 4 +- vllm/v1/attention/backend.py | 13 + .../backends/ddtree_branch_triton.py | 21 +- vllm/v1/attention/backends/fa_utils.py | 1 + vllm/v1/attention/backends/flash_attn_v100.py | 1468 ++- vllm/v1/attention/backends/flashinfer_sm70.py | 3 + vllm/v1/attention/backends/gdn_attn.py | 697 +- vllm/v1/attention/backends/mla/indexer.py | 8 + vllm/v1/attention/backends/registry.py | 8 +- .../attention/ops/triton_prefill_attention.py | 6 +- .../attention/ops/triton_unified_attention.py | 12 +- vllm/v1/core/kv_cache_coordinator.py | 40 +- vllm/v1/core/kv_cache_manager.py | 76 +- vllm/v1/core/kv_cache_utils.py | 95 +- vllm/v1/core/sched/output.py | 5 + vllm/v1/core/sched/scheduler.py | 57 +- vllm/v1/core/single_type_kv_cache_manager.py | 180 +- vllm/v1/cudagraph_dispatcher.py | 281 +- vllm/v1/engine/core.py | 16 +- vllm/v1/kv_cache_interface.py | 64 +- vllm/v1/request.py | 5 + vllm/v1/sample/ops/topk_topp_triton.py | 39 + vllm/v1/sample/rejection_sampler.py | 23 + vllm/v1/sample/sampler.py | 225 +- vllm/v1/sample/thinking_budget_state.py | 8 +- vllm/v1/simple_kv_offload/manager.py | 4 +- vllm/v1/spec_decode/ddtree_metadata.py | 70 +- vllm/v1/spec_decode/dflash.py | 39 +- vllm/v1/spec_decode/draft_prob_alignment.py | 105 + vllm/v1/spec_decode/dspark.py | 101 + vllm/v1/spec_decode/llm_base_proposer.py | 358 +- vllm/v1/worker/block_table.py | 4 +- vllm/v1/worker/gpu/attn_utils.py | 166 +- vllm/v1/worker/gpu/buffer_utils.py | 6 + vllm/v1/worker/gpu/cp_utils.py | 21 + vllm/v1/worker/gpu/cudagraph_utils.py | 57 +- vllm/v1/worker/gpu/input_batch.py | 8 + vllm/v1/worker/gpu/mamba_align.py | 319 + vllm/v1/worker/gpu/model_runner.py | 182 +- vllm/v1/worker/gpu/model_states/default.py | 1 + vllm/v1/worker/gpu/model_states/interface.py | 11 + .../worker/gpu/model_states/mamba_hybrid.py | 359 +- vllm/v1/worker/gpu/model_states/whisper.py | 1 + vllm/v1/worker/gpu/sample/gumbel.py | 106 +- vllm/v1/worker/gpu/spec_decode/__init__.py | 38 + .../worker/gpu/spec_decode/dflash/__init__.py | 2 + .../gpu/spec_decode/dflash/cudagraph.py | 125 + .../gpu/spec_decode/dflash/speculator.py | 829 ++ .../v1/worker/gpu/spec_decode/dflash/utils.py | 101 + .../gpu/spec_decode/dflash2/__init__.py | 2 + .../gpu/spec_decode/dflash2/ngram_assist.py | 162 + .../spec_decode/dflash2/sparse_rejection.py | 261 + .../gpu/spec_decode/dflash2/speculator.py | 621 ++ .../gpu/spec_decode/eagle/eagle3_utils.py | 14 +- .../gpu/spec_decode/eagle/speculator.py | 5 + vllm/v1/worker/gpu/spec_decode/eagle/utils.py | 7 + .../spec_decode/rejection_sampler_utils.py | 336 +- vllm/v1/worker/gpu/spec_decode/speculator.py | 371 + vllm/v1/worker/gpu/spec_decode/utils.py | 15 + vllm/v1/worker/gpu/warmup.py | 43 +- vllm/v1/worker/gpu_model_runner.py | 426 +- vllm/v1/worker/gpu_worker.py | 133 +- vllm/v1/worker/mamba_utils.py | 57 + vllm/v1/worker/utils.py | 72 +- 391 files changed, 68435 insertions(+), 16840 deletions(-) create mode 100644 benchmarks/analyze_dflash2_selector_alignment.py create mode 100644 benchmarks/analyze_dsv4_dspark_confidence.py create mode 100644 benchmarks/benchmark_dsv4_dspark_long_context.py create mode 100644 benchmarks/benchmark_dsv4_gsm8k_api.py create mode 100644 benchmarks/benchmark_dsv4_quality_api.py create mode 100644 benchmarks/benchmark_sm70_cublas_gemm_algorithms.py create mode 100644 benchmarks/benchmark_sm70_d256_exact_dense.py create mode 100644 benchmarks/benchmark_sm70_dflash2_gdn_metadata.py create mode 100644 benchmarks/benchmark_sm70_dflash2_gsm8k.py create mode 100644 benchmarks/benchmark_sm70_dflash2_sparse_rejection.py create mode 100644 benchmarks/benchmark_sm70_fp8_prefill_exact_dense_blas.py create mode 100644 benchmarks/benchmark_sm70_fp8_prefill_tactics.py create mode 100644 benchmarks/benchmark_sm70_fp8_qpn8.py create mode 100644 benchmarks/benchmark_sm70_prefill_gemm_collective.py create mode 100644 benchmarks/benchmark_sm70_sglang_push_chain.py create mode 100644 benchmarks/benchmark_sm70_tp4_fused_norm_collective.py create mode 100644 benchmarks/compare_dsv4_quality_results.py create mode 100644 benchmarks/kernels/benchmark_sm70_nvfp4_qpn2.py create mode 100644 benchmarks/profile_sm70_fp8_qpn8_gated_pair_ncu.py create mode 100644 benchmarks/profile_sm70_fp8_qpn8_ncu.py create mode 100644 cmake/patches/sm70_flash_attn_d256_gqa_arch.patch create mode 100644 cmake/patches/sm70_flash_attn_d256_k_pingpong.patch create mode 100644 csrc/sm70_turbomind/ops/LICENSE.v100-skinny create mode 100644 csrc/sm70_turbomind/ops/fp8_qpn8_sm70.cu create mode 100644 csrc/sm70_turbomind/ops/nvfp4_qpn2_sm70.cu create mode 100644 csrc/sm70_turbomind/ops/nvfp4_qpn4_sm70.cu create mode 100644 csrc/sm70_turbomind/ops/qwen38_prefill_cutlass.cu create mode 100644 csrc/sm70_turbomind/ops/qwen38_prefill_cutlass.cuh create mode 100644 csrc/sm70_turbomind/ops/sm70_top20_philox_fragment.cpp create mode 100644 csrc/sm70_turbomind/ops/sm70_top20_philox_sampler.cu create mode 100644 docs/assets/wechat-group-qr-7.png create mode 100644 docs/design/sm70_deepseek_v4_dspark_verifier_20ms_long_context.md create mode 100644 docs/design/sm70_dflash2_acceptance_rd.md create mode 100644 docs/design/sm70_dflash2_long_verify_decay.md create mode 100644 docs/design/sm70_dflash2_ngram_hybrid.md create mode 100644 docs/design/sm70_dflash2_nvfp4_17ms.md create mode 100644 docs/design/sm70_dflash2_target_graph_20ms.md create mode 100644 docs/design/sm70_prefix_anchored_swa.md create mode 100644 docs/design/sm70_qwen38_fp8_prefill_5500.md create mode 100644 docs/design/sm70_qwen38_nvfp4_decode.md create mode 100644 docs/design/sm70_qwen38_qpn8_decode.md create mode 100644 tests/benchmarks/test_dsv4_quality_tools.py create mode 100644 tests/config/test_prefix_anchored_swa.py create mode 100644 tests/distributed/test_custom_all_reduce_allocator.py create mode 100644 tests/distributed/test_static_tensor_transfer.py create mode 100644 tests/kernels/attention/test_sm70_flash_v100_anchored_swa.py create mode 100644 tests/kernels/attention/test_sm70_flash_v100_grouped_verify.py create mode 100644 tests/kernels/attention/test_sm70_flash_v100_smallq_metadata.py create mode 100644 tests/kernels/core/test_batched_weight_rms_norm.py create mode 100644 tests/kernels/core/test_sm70_dflash2_gemma_rms.py create mode 100644 tests/kernels/mamba/test_mrv2_mamba_align.py create mode 100644 tests/kernels/moe/test_sm70_unquantized_moe_config.py create mode 100644 tests/kernels/test_sm70_qwen35_gdn_split.py create mode 100644 tests/model_executor/test_deepseek_v4_swa_rope.py create mode 100644 tests/model_executor/test_sm70_fp8_qpn8_pp2_tp4.py create mode 100644 tests/model_executor/test_uva_tower_offload.py create mode 100644 tests/models/test_deepseek_v4_pipeline.py create mode 100644 tests/quantization/test_sm70_compressed_tensors_fastpaths.py create mode 100644 tests/quantization/test_sm70_modelopt_mixed_nvfp4.py create mode 100644 tests/quantization/test_sm70_modelopt_nvfp4.py create mode 100644 tests/quantization/test_sm70_modelopt_nvfp4_kernel_oracle.py create mode 100644 tests/quantization/test_sm70_nvfp4_qpn2.py create mode 100644 tests/transformers_utils/test_deepseek_ocr_resolution_modes.py create mode 100644 tests/v1/attention/test_prefix_anchored_swa_mask.py create mode 100644 tests/v1/core/test_swa_inflight_window_free.py create mode 100644 tests/v1/cudagraph/test_sm70_mtp_split_cudagraphs.py create mode 100644 tests/v1/spec_decode/test_dflash2.py create mode 100644 tests/v1/spec_decode/test_dflash2_ngram_assist.py create mode 100644 tests/v1/spec_decode/test_dflash2_selector_analysis.py create mode 100644 tests/v1/spec_decode/test_dflash_mrv2_config.py create mode 100644 tests/v1/worker/test_gpu_model_runner_v2_kv_zero.py create mode 100644 tests/v1/worker/test_gpu_warmup_blocks.py create mode 100644 tests/v1/worker/test_gpu_worker_static_pp.py create mode 100644 tests/v1/worker/test_kv_block_zeroer.py create mode 100644 vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py create mode 100644 vllm/model_executor/models/dflash_sm70.py create mode 100644 vllm/model_executor/models/qwen3_dflash2.py create mode 100644 vllm/v1/worker/gpu/mamba_align.py create mode 100644 vllm/v1/worker/gpu/spec_decode/dflash/__init__.py create mode 100644 vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py create mode 100644 vllm/v1/worker/gpu/spec_decode/dflash/speculator.py create mode 100644 vllm/v1/worker/gpu/spec_decode/dflash/utils.py create mode 100644 vllm/v1/worker/gpu/spec_decode/dflash2/__init__.py create mode 100644 vllm/v1/worker/gpu/spec_decode/dflash2/ngram_assist.py create mode 100644 vllm/v1/worker/gpu/spec_decode/dflash2/sparse_rejection.py create mode 100644 vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py create mode 100644 vllm/v1/worker/gpu/spec_decode/speculator.py 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 7d863a8053..d07300fdff 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -396,6 +396,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") "${VLLM_SM70_TURBOMIND_LMDEPLOY_ROOT}/src/turbomind/kernels/gemm/kernel/sm70_884_8.cu" "${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/nvfp4_qpn4_sm70.cu" + "${VLLM_SM70_TURBOMIND_ROOT}/ops/qwen38_prefill_cutlass.cu" + "${VLLM_SM70_TURBOMIND_ROOT}/ops/nvfp4_qpn2_sm70.cu" "${VLLM_SM70_TURBOMIND_ROOT}/ops/awq_sm70_gemm.cu") set_gencode_flags_for_srcs( @@ -709,6 +713,44 @@ if(VLLM_REGISTER_BASIC_ACTIVATION_IN_C) target_compile_definitions(_C PRIVATE VLLM_REGISTER_BASIC_ACTIVATION_IN_C=1) endif() +# Keep the accepted SM70 chunked sampler in a separate wheel module so adding +# it does not perturb the quality- and speed-frozen primary vllm._C module. +if(VLLM_GPU_LANG STREQUAL "CUDA" AND SM70_TURBOMIND_ARCHS) + set(VLLM_SM70_SAMPLER_FRAGMENT_SRCS + "csrc/sm70_turbomind/ops/sm70_top20_philox_fragment.cpp" + "csrc/sm70_turbomind/ops/sm70_top20_philox_sampler.cu") + set_gencode_flags_for_srcs( + SRCS "${VLLM_SM70_SAMPLER_FRAGMENT_SRCS}" + CUDA_ARCHS "${SM70_TURBOMIND_ARCHS}") + define_extension_target( + _sm70_sampler_C + DESTINATION vllm + LANGUAGE CUDA + 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 + CXX_STANDARD_REQUIRED ON + CXX_EXTENSIONS OFF + CUDA_STANDARD 17 + CUDA_STANDARD_REQUIRED ON + CUDA_EXTENSIONS OFF) + target_compile_definitions(_sm70_sampler_C PRIVATE + TORCH_API_INCLUDE_EXTENSION_H + $<$:__CUDA_NO_HALF_OPERATORS__> + $<$:__CUDA_NO_HALF_CONVERSIONS__> + $<$:__CUDA_NO_BFLOAT16_CONVERSIONS__> + $<$:__CUDA_NO_HALF2_OPERATORS__>) + target_compile_options(_sm70_sampler_C PRIVATE + $<$:-UNDEBUG> + $<$:-UNDEBUG;-lineinfo>) + target_link_libraries(_sm70_sampler_C PRIVATE + "${TORCH_INSTALL_PREFIX}/lib/libtorch_python.so") +endif() + if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # # _C_stable_libtorch extension (ops registered via STABLE_TORCH_LIBRARY) diff --git a/README.md b/README.md index 59a93472e6..f305d6e580 100644 --- a/README.md +++ b/README.md @@ -1,26 +1,20 @@ # 1Cat-vLLM -> 一猫之下始终相信，V100 不该在今天的大模型浪潮中被轻易宣判“过时”。 +> 一猫之下始终相信，V100 不该在今天的大模型浪潮里被轻易宣判“过时”。 > -> 1Cat-vLLM 是一个专注于 **SM70 / Tesla V100** 的 vLLM 工程分支。我们 -> 不追求覆盖所有硬件，而是围绕 Volta 架构补齐并优化现代大模型推理的 -> 关键路径，包括量化推理、注意力后端、长上下文、MTP 投机解码、 -> CUDA Graph、分布式通信、运行时策略和部署工具链，让 Qwen、DeepSeek -> 等新模型在 V100 上从“能跑”走向“可部署、可验证、可持续优化”。 +> 1Cat-vLLM 是面向 **SM70 / Tesla V100** 的 vLLM 工程分支。项目围绕 +> AWQ、注意力后端、长上下文稳定性、MTP 投机解码、运行时默认值和部署 +> 路径做了成体系的优化，让更多现代模型场景在 V100 上真正变得可用、 +> 好用、能持续部署。 > -> 我们希望把来自真实 V100 环境的工程实现、性能数据和验证过程贡献给 -> 开源社区，也欢迎仍在使用 V100 的个人开发者、工作室和团队参与测试、 -> 复现、反馈和共建，让这批依然有价值的算力继续发挥作用。 +> 我们希望把一猫之下在 V100 上的工程经验、优化成果和验证过程贡献给 +> 开源社区，也欢迎继续使用 V100 的个人开发者、工作室和团队一起反馈、 +> 复现和改进。 -1Cat-vLLM is a vLLM engineering fork focused on **SM70 / Tesla V100**. -Rather than targeting every hardware platform, it fills and optimizes the -critical inference paths required by modern models on Volta, including -quantized inference, attention backends, long-context serving, MTP speculative -decoding, CUDA Graph execution, distributed communication, runtime policies, -and deployment tooling. - -The goal is to make models such as Qwen and DeepSeek not merely runnable on -V100, but deployable, reproducible, and continuously optimizable. +1Cat-vLLM is a **Tesla V100 / SM70** focused vLLM fork for serving modern +Qwen-class AWQ and experimental FP8 models on Volta GPUs. It integrates +TurboMind-derived SM70 kernels, a V100 FlashAttention path, runtime defaults +for long-context serving, and OpenAI-compatible API fixes for common clients. ## Project Focus @@ -31,12 +25,6 @@ V100, but deployable, reproducible, and continuously optimizable. - **V100 FlashAttention path**: `FLASH_ATTN_V100` decode and prefill backend for Volta GPUs, with SM70 compile-graph, guarded XQA decode, and D=256 paged-prefix low-smem fast paths enabled by default. -- **DeepSeek V4 Flash**: supports running the original, unmodified DeepSeek V4 - Flash weights across eight GPUs. -- **Day-0 Qwen3.8-27B support**: native Day-0 support for Qwen3.8-27B. -- **Quantization support**: FP8, NVFP4, MXFP4, AWQ, and GPTQ inference paths - are included for supported model and hardware combinations. Availability - and production status vary by checkpoint and runtime route. - **Long-context serving**: public profiles default to 256K context where the model and memory budget allow it. - **MTP serving**: Qwen3.6-class MTP speculative decoding remains available as @@ -54,7 +42,6 @@ V100, but deployable, reproducible, and continuously optimizable. - `tclf90/Qwen3.6-27B-AWQ` - `tclf90/Qwen3.6-35B-A3B-AWQ` - `tclf90/Qwen3.5-122B-A10B-AWQ` for larger 4-GPU setups -- `Qwen/Qwen3.8-27B` for the latest supported Qwen release The launch examples use local paths such as `/path/to/Qwen3.6-27B-AWQ`. Replace them with your local model path or a Hugging Face repository id. @@ -70,7 +57,6 @@ validation. | --- | --- | | 4 x Tesla V100 32 GB | Main public reference target | | 2 x Tesla V100 32 GB | Supported for selected 27B profiles with lower concurrency | -| 8 x Tesla V100 32 GB | DeepSeek V4 Flash original-checkpoint TP8 target | Typical model placement: @@ -79,10 +65,6 @@ Typical model placement: the long-context public default. - `Qwen3.6-35B-A3B-AWQ`: TP4 recommended. - `Qwen3.5-122B-A10B-AWQ`: TP4 supported for larger deployments. -- `DeepSeek-V4-Flash`: the original mixed MXFP4/FP8 checkpoint is supported - on eight 32 GB V100 GPUs with TP8. -- `Qwen3.8-27B`: supported natively from Day 0, with validated FP16 and FP8 - serving paths. Multimodal defaults: @@ -367,26 +349,17 @@ python -m pip install -e . --no-build-isolation 请使用微信扫描下方二维码加入群组： -![1Cat-vLLM 微信交流群二维码](docs/assets/wechat-group-qr-6.png) - -> 提示：微信群二维码通常 7 天内有效。若扫描失败或提示过期，请联系微信号：YM_isi。 - -## Roadmap - -1Cat-vLLM will continue to prioritize Tesla V100 / SM70, making modern LLM inference more usable, stable, and deployable on existing GPUs. - -- **Near term:** Complete Volta compatibility and validation for DeepSeek V4, GLM5, and other new model families, while improving long-context serving, quantization, multi-GPU parallelism, and runtime stability. -- **Mid term:** Expand the Volta model support matrix with a focus on models up to 300B parameters, backed by reproducible performance, quality, and release validation. -- **Long term:** Build on the Volta foundation and extend hardware support in stages: **Volta → Turing/Ampere → CDNA**, bringing sustainable LLM inference to a broader range of existing accelerators. +![1Cat-vLLM 微信交流群二维码](docs/assets/wechat-group-qr-7.png) +> 提示：微信群二维码通常 7 天内有效。若扫描失败或提示过期，请重新打开本页查看最新图片，或关注仓库更新。 ## Repository Notes -- This fork focuses on SM70 quantized inference, V100-oriented attention and - long-context tuning, model-specific runtime and deployment paths, and - continued MTP and DFlash research. -- Production status is route-specific. Use the documented public profiles and - validated release matrices as the source of truth. +- Upstream project: [vLLM](https://github.com/vllm-project/vllm) +- This fork focuses on SM70 AWQ support, V100-oriented attention/runtime + tuning, and experimental FP8/MTP/DFlash validation paths. +- Prebuilt wheels are the public installation path. +- Source builds are for development and kernel work. ## Acknowledgements diff --git a/benchmarks/analyze_dflash2_selector_alignment.py b/benchmarks/analyze_dflash2_selector_alignment.py new file mode 100644 index 0000000000..b332106835 --- /dev/null +++ b/benchmarks/analyze_dflash2_selector_alignment.py @@ -0,0 +1,517 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Audit and sweep exact DFlash2 selector proposals from alignment dumps. + +This tool consumes the B1 records emitted by the GPU-runner compact rejection +path when ``VLLM_SPEC_DUMP_ALIGNMENT=1``. It never changes target sampling. +Counterfactual results are one-step overlap proxies on the recorded prefixes; +only an end-to-end run can establish a new acceptance-length result. +""" + +from __future__ import annotations + +import argparse +import glob +import json +import math +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any + +import torch + + +@dataclass(frozen=True) +class ProposalConfig: + name: str + proposal_temperature_scale: float = 1.0 + proposal_top_p: float = 1.0 + unary_scale: float = 1.0 + edge_scale: float = 1.0 + future_beta: float = 0.0 + greedy_mix: float = 0.0 + use_cached_logits: bool = False + + +@dataclass +class AlignmentRecord: + path: Path + step: int + temperature: float + top_p: float + target_topk_ids: torch.Tensor + target_topk_logits: torch.Tensor + candidate_ids: torch.Tensor + realized_logits: torch.Tensor + unary_logits: torch.Tensor + lattice_scores: torch.Tensor + draft_sampled: torch.Tensor + num_sampled: int + + +def _compact_probs(logits: torch.Tensor, top_p: float) -> torch.Tensor: + """Apply the compact top-p contract used by the sparse rejection kernel.""" + logits = logits.to(torch.float64) + order = torch.argsort(logits, descending=True, stable=True) + sorted_probs = torch.softmax(logits[order], dim=-1) + cumulative_before = torch.cumsum(sorted_probs, dim=-1) - sorted_probs + keep_sorted = (top_p >= 1.0) | (cumulative_before < top_p) + kept = torch.zeros_like(sorted_probs) + kept[keep_sorted] = sorted_probs[keep_sorted] + kept /= kept.sum() + probs = torch.zeros_like(kept) + probs[order] = kept + return probs + + +def _distribution_overlap( + target_ids: torch.Tensor, + target_probs: torch.Tensor, + proposal_ids: torch.Tensor, + proposal_probs: torch.Tensor, +) -> float: + matches = target_ids[:, None] == proposal_ids[None, :] + proposal_on_target = torch.where( + matches, + proposal_probs[None, :], + torch.zeros((), dtype=proposal_probs.dtype), + ).sum(dim=1) + return float(torch.minimum(target_probs, proposal_on_target).sum()) + + +def _support_mass( + target_ids: torch.Tensor, + target_probs: torch.Tensor, + proposal_ids: torch.Tensor, +) -> float: + supported = (target_ids[:, None] == proposal_ids[None, :]).any(dim=1) + return float(target_probs[supported].sum()) + + +def _transformed_lattice( + record: AlignmentRecord, config: ProposalConfig +) -> torch.Tensor: + unary = record.unary_logits[:, None, :].to(torch.float64) + lattice = record.lattice_scores.to(torch.float64) + edge = lattice - unary + denominator = record.temperature * config.proposal_temperature_scale + return (config.unary_scale * unary + config.edge_scale * edge) / denominator + + +def _backward_messages(lattice: torch.Tensor) -> torch.Tensor: + """Compute normalized log-sum-exp future messages for a KxK chain.""" + num_steps, top_k, _ = lattice.shape + future = torch.zeros((num_steps, top_k), dtype=lattice.dtype) + for step in range(num_steps - 2, -1, -1): + values = lattice[step + 1] + future[step + 1][None, :] + future[step] = torch.logsumexp(values, dim=-1) + future[step] -= future[step].max() + return future + + +def _path_predecessors(record: AlignmentRecord) -> list[int]: + num_steps = record.candidate_ids.shape[0] + if record.draft_sampled.numel() < num_steps + 1: + raise ValueError(f"{record.path}: draft_sampled is missing the anchor row") + predecessors = [0] + for step in range(num_steps - 1): + proposed = int(record.draft_sampled[step + 1]) + matches = torch.nonzero(record.candidate_ids[step] == proposed).flatten() + if matches.numel() != 1: + raise ValueError( + f"{record.path}: proposed token {proposed} is not unique at step {step}" + ) + predecessors.append(int(matches[0])) + return predecessors + + +def _proposal_probs( + record: AlignmentRecord, + config: ProposalConfig, +) -> list[torch.Tensor]: + if config.use_cached_logits: + rows = [ + _compact_probs(row, config.proposal_top_p) for row in record.realized_logits + ] + else: + lattice = _transformed_lattice(record, config) + future = _backward_messages(lattice) + predecessors = _path_predecessors(record) + rows = [ + _compact_probs( + lattice[step, predecessor] + config.future_beta * future[step], + config.proposal_top_p, + ) + for step, predecessor in enumerate(predecessors) + ] + if config.greedy_mix == 0.0: + return rows + if not 0.0 <= config.greedy_mix <= 1.0: + raise ValueError("greedy_mix must be in [0, 1]") + mixed_rows = [] + for probs in rows: + mixed = probs * (1.0 - config.greedy_mix) + mixed[torch.argmax(probs)] += config.greedy_mix + mixed_rows.append(mixed) + return mixed_rows + + +def _target_rows(record: AlignmentRecord) -> list[torch.Tensor]: + rows = [] + for logits in record.target_topk_logits[: record.candidate_ids.shape[0]]: + rows.append(_compact_probs(logits / record.temperature, record.top_p)) + return rows + + +def _load_record(path: Path) -> AlignmentRecord: + payload = torch.load(path, map_location="cpu", weights_only=True) + if payload.get("format") != "dflash2_selector_alignment_v1": + raise ValueError(f"{path}: unsupported alignment format") + candidate_ids = payload["selector_candidate_ids"].to(torch.int64) + cached_ids = payload["draft_candidate_ids"].to(torch.int64) + if not torch.equal(candidate_ids, cached_ids): + raise ValueError(f"{path}: packed selector IDs do not match request-slot cache") + record = AlignmentRecord( + path=path, + step=int(payload["step"]), + temperature=float(payload["temperature"]), + top_p=float(payload["top_p"]), + target_topk_ids=payload["target_topk_ids"].to(torch.int64), + target_topk_logits=payload["target_topk_logits"].to(torch.float64), + candidate_ids=candidate_ids, + realized_logits=payload["draft_realized_logits"].to(torch.float64), + unary_logits=payload["selector_unary_logits"].to(torch.float64), + lattice_scores=payload["selector_lattice_scores"].to(torch.float64), + draft_sampled=payload["draft_sampled"].to(torch.int64).flatten(), + num_sampled=int(payload["num_sampled"].flatten()[0]), + ) + return record + + +def _mean(values: list[float]) -> float: + return float(sum(values) / len(values)) if values else math.nan + + +def _completion_proxy(per_position_overlap: list[float]) -> float: + survival = 1.0 + result = 1.0 + for overlap in per_position_overlap: + survival *= overlap + result += survival + return result + + +def _evaluate(records: list[AlignmentRecord], config: ProposalConfig) -> dict[str, Any]: + num_steps = records[0].candidate_ids.shape[0] + overlaps = [[] for _ in range(num_steps)] + for record in records: + target_rows = _target_rows(record) + proposal_rows = _proposal_probs(record, config) + for step, (target_probs, proposal_probs) in enumerate( + zip(target_rows, proposal_rows) + ): + overlaps[step].append( + _distribution_overlap( + record.target_topk_ids[step], + target_probs, + record.candidate_ids[step], + proposal_probs, + ) + ) + means = [_mean(values) for values in overlaps] + return { + "config": asdict(config), + "num_records": len(records), + "mean_overlap_by_position": means, + "completion_length_proxy": _completion_proxy(means), + } + + +def _sweep_configs() -> list[ProposalConfig]: + configs = [ProposalConfig(name="current", use_cached_logits=True)] + seen: set[tuple[float, float, float, float, float, float]] = set() + + def add( + family: str, + temperature: float = 1.0, + top_p: float = 1.0, + unary: float = 1.0, + edge: float = 1.0, + beta: float = 0.0, + greedy_mix: float = 0.0, + ) -> None: + key = (temperature, top_p, unary, edge, beta, greedy_mix) + if key in seen: + return + seen.add(key) + configs.append( + ProposalConfig( + name=( + f"{family}:t={temperature:g},p={top_p:g},u={unary:g}," + f"e={edge:g},b={beta:g},g={greedy_mix:g}" + ), + proposal_temperature_scale=temperature, + proposal_top_p=top_p, + unary_scale=unary, + edge_scale=edge, + future_beta=beta, + greedy_mix=greedy_mix, + ) + ) + + for temperature in (0.7, 0.8, 0.9, 0.95, 1.0, 1.05, 1.1, 1.2, 1.3): + add("temperature", temperature=temperature) + for top_p in (0.9, 0.95, 0.98, 0.99): + for temperature in (0.8, 0.9, 1.0): + add("nucleus", temperature=temperature, top_p=top_p) + for greedy_mix in (0.1, 0.2, 0.3, 0.4, 0.5): + for temperature in (0.8, 0.9, 1.0, 1.1): + add( + "greedy-mixture", + temperature=temperature, + greedy_mix=greedy_mix, + ) + for unary in (0.75, 1.0, 1.25): + for edge in (0.5, 0.75, 1.0, 1.25, 1.5): + for temperature in (0.85, 1.0, 1.15): + add( + "edge-calibration", + temperature=temperature, + unary=unary, + edge=edge, + ) + for beta in (0.25, 0.5, 0.75, 1.0): + for temperature in (0.8, 0.9, 1.0, 1.1): + for edge in (0.75, 1.0, 1.25): + add( + "future-message", + temperature=temperature, + edge=edge, + beta=beta, + ) + return configs + + +def _baseline_diagnostics(records: list[AlignmentRecord]) -> dict[str, Any]: + num_steps = records[0].candidate_ids.shape[0] + support = [[] for _ in range(num_steps)] + cached_lattice_max_abs: list[float] = [] + for record in records: + target_rows = _target_rows(record) + predecessors = _path_predecessors(record) + for step, target_probs in enumerate(target_rows): + support[step].append( + _support_mass( + record.target_topk_ids[step], + target_probs, + record.candidate_ids[step], + ) + ) + expected = ( + record.lattice_scores[step, predecessors[step]] / record.temperature + ) + cached_lattice_max_abs.append( + float((expected - record.realized_logits[step]).abs().max()) + ) + support_means = [_mean(values) for values in support] + return { + "observed_mean_completion_tokens_per_round": _mean( + [float(record.num_sampled) for record in records] + ), + "candidate_support_mass_by_position": support_means, + "candidate_support_completion_upper_proxy": _completion_proxy(support_means), + "cached_vs_lattice_max_abs": max(cached_lattice_max_abs, default=math.nan), + } + + +def summarize(pattern: str, top_n: int) -> dict[str, Any]: + paths = [Path(path) for path in sorted(glob.glob(pattern))] + if not paths: + raise ValueError(f"no selector alignment dumps matched {pattern!r}") + records = [_load_record(path) for path in paths] + tune = [record for record in records if record.step % 2 == 0] + holdout = [record for record in records if record.step % 2 == 1] + if not tune or not holdout: + tune = records + holdout = records + + configs = _sweep_configs() + tune_results = [_evaluate(tune, config) for config in configs] + holdout_by_name = { + result["config"]["name"]: result + for result in (_evaluate(holdout, config) for config in configs) + } + current_tune = tune_results[0]["completion_length_proxy"] + current_holdout = holdout_by_name["current"]["completion_length_proxy"] + position_policy = [] + position_tune_overlaps = [] + position_holdout_overlaps = [] + for position in range(records[0].candidate_ids.shape[0]): + best = max( + tune_results, + key=lambda result: result["mean_overlap_by_position"][position], + ) + holdout_best = holdout_by_name[best["config"]["name"]] + position_policy.append( + { + "position": position, + "config": best["config"], + "tune_overlap": best["mean_overlap_by_position"][position], + "holdout_overlap": holdout_best["mean_overlap_by_position"][position], + } + ) + position_tune_overlaps.append(best["mean_overlap_by_position"][position]) + position_holdout_overlaps.append( + holdout_best["mean_overlap_by_position"][position] + ) + ranked = sorted( + tune_results[1:], + key=lambda result: result["completion_length_proxy"], + reverse=True, + ) + leaders = [] + for result in ranked[:top_n]: + holdout_result = holdout_by_name[result["config"]["name"]] + leaders.append( + { + "config": result["config"], + "tune_completion_length_proxy": result["completion_length_proxy"], + "tune_delta": result["completion_length_proxy"] - current_tune, + "holdout_completion_length_proxy": holdout_result[ + "completion_length_proxy" + ], + "holdout_delta": ( + holdout_result["completion_length_proxy"] - current_holdout + ), + "tune_mean_overlap_by_position": result["mean_overlap_by_position"], + "holdout_mean_overlap_by_position": holdout_result[ + "mean_overlap_by_position" + ], + } + ) + + return { + "contract": { + "glob": pattern, + "num_records": len(records), + "num_tune_records": len(tune), + "num_holdout_records": len(holdout), + "split": "even/odd sampler step", + "counterfactual_scope": "one-step recorded-prefix overlap proxy", + "end_to_end_claim": False, + }, + "baseline": _baseline_diagnostics(records), + "current": { + "tune": tune_results[0], + "holdout": holdout_by_name["current"], + }, + "position_policy": { + "selection": "best tune overlap independently at each depth", + "rows": position_policy, + "tune_completion_length_proxy": _completion_proxy(position_tune_overlaps), + "tune_delta": (_completion_proxy(position_tune_overlaps) - current_tune), + "holdout_completion_length_proxy": _completion_proxy( + position_holdout_overlaps + ), + "holdout_delta": ( + _completion_proxy(position_holdout_overlaps) - current_holdout + ), + }, + "leaders": leaders, + } + + +def render_markdown(summary: dict[str, Any]) -> str: + contract = summary["contract"] + baseline = summary["baseline"] + current = summary["current"] + position_policy = summary["position_policy"] + lines = [ + "# DFlash2 Selector Alignment Sweep", + "", + ( + f"- Records: `{contract['num_records']}` " + f"(tune `{contract['num_tune_records']}`, " + f"holdout `{contract['num_holdout_records']}`)" + ), + ( + "- Counterfactual scope: one-step overlap on recorded prefixes; " + "this is not an end-to-end acceptance claim." + ), + ( + "- Observed completion tokens/round: " + f"`{baseline['observed_mean_completion_tokens_per_round']:.4f}`" + ), + ( + "- Current overlap proxy (tune/holdout): " + f"`{current['tune']['completion_length_proxy']:.4f}` / " + f"`{current['holdout']['completion_length_proxy']:.4f}`" + ), + ( + "- Fixed-top16 support upper proxy: " + f"`{baseline['candidate_support_completion_upper_proxy']:.4f}`" + ), + ( + "- Cached/lattice max absolute mismatch: " + f"`{baseline['cached_vs_lattice_max_abs']:.3e}`" + ), + ( + "- Position-wise held-out proxy/delta: " + f"`{position_policy['holdout_completion_length_proxy']:.4f}` / " + f"`{position_policy['holdout_delta']:+.4f}`" + ), + "", + "| candidate | tune proxy | tune delta | holdout proxy | holdout delta |", + "| --- | ---: | ---: | ---: | ---: |", + ] + for row in summary["leaders"]: + lines.append( + f"| `{row['config']['name']}` | " + f"{row['tune_completion_length_proxy']:.4f} | " + f"{row['tune_delta']:+.4f} | " + f"{row['holdout_completion_length_proxy']:.4f} | " + f"{row['holdout_delta']:+.4f} |" + ) + lines.extend( + [ + "", + "## Position-wise calibration policy", + "", + "| position | tune overlap | holdout overlap | configuration |", + "| ---: | ---: | ---: | --- |", + ] + ) + for row in position_policy["rows"]: + lines.append( + f"| {row['position']} | {row['tune_overlap']:.4f} | " + f"{row['holdout_overlap']:.4f} | `{row['config']['name']}` |" + ) + lines.append("") + return "\n".join(lines) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument( + "--glob", + default="/tmp/spec_alignment_dflash2_selector_*.pt", + help="Glob for GPU-runner DFlash2 selector alignment records.", + ) + parser.add_argument("--top", type=int, default=20) + parser.add_argument("--out-json", type=Path) + parser.add_argument("--out-md", type=Path) + args = parser.parse_args() + + summary = summarize(args.glob, args.top) + markdown = render_markdown(summary) + if args.out_json is not None: + args.out_json.parent.mkdir(parents=True, exist_ok=True) + args.out_json.write_text(json.dumps(summary, indent=2), encoding="utf-8") + if args.out_md is not None: + args.out_md.parent.mkdir(parents=True, exist_ok=True) + args.out_md.write_text(markdown, encoding="utf-8") + print(markdown) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/analyze_dsv4_dspark_confidence.py b/benchmarks/analyze_dsv4_dspark_confidence.py new file mode 100644 index 0000000000..0445483f4c --- /dev/null +++ b/benchmarks/analyze_dsv4_dspark_confidence.py @@ -0,0 +1,288 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Calibrate and summarize DeepSeek-V4 DSpark confidence-head dumps.""" + +from __future__ import annotations + +import argparse +import json +import math +from collections import defaultdict +from pathlib import Path +from typing import Any + +import torch + + +def _ece(probs: torch.Tensor, targets: torch.Tensor, bins: int = 15) -> float: + if probs.numel() == 0: + return math.nan + edges = torch.linspace(0.0, 1.0, bins + 1, dtype=torch.float64) + total = probs.numel() + value = 0.0 + for index in range(bins): + if index + 1 == bins: + mask = (probs >= edges[index]) & (probs <= edges[index + 1]) + else: + mask = (probs >= edges[index]) & (probs < edges[index + 1]) + count = int(mask.sum()) + if count: + gap = (probs[mask].mean() - targets[mask].mean()).abs() + value += float(gap) * count / total + return value + + +def _soft_nll(logits: torch.Tensor, targets: torch.Tensor) -> float: + if logits.numel() == 0: + return math.nan + return float(torch.nn.functional.binary_cross_entropy_with_logits(logits, targets)) + + +def _fit_temperature(logits: torch.Tensor, targets: torch.Tensor) -> float: + if logits.numel() < 2: + return 1.0 + # A dense one-dimensional search is deterministic, stable for fractional + # TV labels, and avoids adding scipy/sklearn to the benchmark environment. + temperatures = torch.logspace(-1.0, 1.0, 801, dtype=torch.float64) + scaled = logits.to(torch.float64).unsqueeze(0) / temperatures.unsqueeze(1) + losses = torch.nn.functional.binary_cross_entropy_with_logits( + scaled, + targets.to(torch.float64).expand_as(scaled), + reduction="none", + ).mean(dim=1) + return float(temperatures[int(losses.argmin())]) + + +def _binary_auc(scores: torch.Tensor, labels: torch.Tensor) -> float: + positives = scores[labels == 1] + negatives = scores[labels == 0] + if positives.numel() == 0 or negatives.numel() == 0: + return math.nan + # Proposal blocks are at most five positions in the production contract; + # this exact pairwise definition is small and handles ties correctly. + comparisons = positives[:, None] - negatives[None, :] + return float((comparisons.gt(0).float() + 0.5 * comparisons.eq(0).float()).mean()) + + +def _summary( + logits: torch.Tensor, + targets: torch.Tensor, + temperature: float, + observed_logits: torch.Tensor | None = None, + observed_labels: torch.Tensor | None = None, +) -> dict[str, Any]: + scaled_logits = logits / temperature + probs = scaled_logits.sigmoid() + result: dict[str, Any] = { + "count": int(logits.numel()), + "temperature": temperature, + "predicted_mean": float(probs.mean()) if probs.numel() else math.nan, + "target_mean": float(targets.mean()) if targets.numel() else math.nan, + "ece": _ece(probs, targets), + "brier": ( + float(torch.square(probs - targets).mean()) if probs.numel() else math.nan + ), + "soft_nll": _soft_nll(scaled_logits, targets), + } + if ( + observed_logits is not None + and observed_labels is not None + and observed_labels.numel() + ): + result["observed_count"] = int(observed_labels.numel()) + result["observed_acceptance"] = float(observed_labels.mean()) + result["observed_auc"] = _binary_auc( + torch.sigmoid(observed_logits / temperature), observed_labels + ) + return result + + +def _load_rows(input_dir: Path, rank: int) -> list[dict[str, Any]]: + rows: list[dict[str, Any]] = [] + for path in sorted(input_dir.glob("spec_alignment_*.pt")): + payload = torch.load(path, map_location="cpu", weights_only=True) + if int(payload.get("rank", -1)) != rank: + continue + required = ( + "draft_confidence_logits", + "conditional_acceptance_targets", + "num_draft_tokens", + "output_valid_counts", + ) + if not all(key in payload for key in required): + continue + payload["_path"] = str(path) + rows.append(payload) + return rows + + +def analyze(rows: list[dict[str, Any]], calibration_fraction: float) -> dict[str, Any]: + if not rows: + raise ValueError("no rank-matching confidence alignment dumps were found") + max_position = max(max(row["num_draft_tokens"], default=0) for row in rows) + by_position: dict[int, dict[str, list[torch.Tensor]]] = defaultdict( + lambda: defaultdict(list) + ) + prefix_by_position: dict[int, dict[str, list[torch.Tensor]]] = defaultdict( + lambda: defaultdict(list) + ) + + for row_index, row in enumerate(rows): + logits = row["draft_confidence_logits"].to(torch.float64) + targets = row["conditional_acceptance_targets"].to(torch.float64) + valid_counts = row["output_valid_counts"].to(torch.int64) + offset = 0 + for request_index, count_raw in enumerate(row["num_draft_tokens"]): + count = int(count_raw) + request_logits = logits[offset : offset + count] + request_targets = targets[offset : offset + count] + accepted = max(0, min(count, int(valid_counts[request_index]) - 1)) + split = ( + "calibration" + if row_index < len(rows) * calibration_fraction + else "evaluation" + ) + cumulative_pred_logits = request_logits + cumulative_targets = request_targets.cumprod(dim=0) + for position in range(count): + bucket = by_position[position] + bucket[f"{split}_logits"].append( + request_logits[position : position + 1] + ) + bucket[f"{split}_targets"].append( + request_targets[position : position + 1] + ) + if position <= accepted and position < count: + observed = torch.tensor( + [1.0 if position < accepted else 0.0], dtype=torch.float64 + ) + bucket[f"{split}_observed_logits"].append( + request_logits[position : position + 1] + ) + bucket[f"{split}_observed"].append(observed) + + prefix_bucket = prefix_by_position[position] + prefix_bucket[f"{split}_logits"].append( + cumulative_pred_logits[: position + 1] + ) + prefix_bucket[f"{split}_targets"].append( + cumulative_targets[position : position + 1] + ) + offset += count + + temperatures: list[float] = [] + position_rows: list[dict[str, Any]] = [] + for position in range(max_position): + bucket = by_position[position] + calibration_logits = torch.cat(bucket["calibration_logits"]) + calibration_targets = torch.cat(bucket["calibration_targets"]) + temperature = _fit_temperature(calibration_logits, calibration_targets) + temperatures.append(temperature) + evaluation_logits = torch.cat( + bucket["evaluation_logits"] or bucket["calibration_logits"] + ) + evaluation_targets = torch.cat( + bucket["evaluation_targets"] or bucket["calibration_targets"] + ) + observed_parts = bucket["evaluation_observed"] or bucket["calibration_observed"] + observed_logits_parts = ( + bucket["evaluation_observed_logits"] + or bucket["calibration_observed_logits"] + ) + observed = torch.cat(observed_parts) if observed_parts else None + observed_logits = ( + torch.cat(observed_logits_parts) if observed_logits_parts else None + ) + position_rows.append( + { + "position": position + 1, + "raw": _summary( + evaluation_logits, + evaluation_targets, + 1.0, + observed_logits, + observed, + ), + "calibrated": _summary( + evaluation_logits, + evaluation_targets, + temperature, + observed_logits, + observed, + ), + } + ) + + prefix_rows: list[dict[str, Any]] = [] + for position in range(max_position): + bucket = prefix_by_position[position] + logits_parts = bucket["evaluation_logits"] or bucket["calibration_logits"] + target_parts = bucket["evaluation_targets"] or bucket["calibration_targets"] + prefix_probs = torch.stack( + [ + torch.sigmoid(part / torch.tensor(temperatures[: part.numel()])).prod() + for part in logits_parts + ] + ) + prefix_targets = torch.cat(target_parts) + prefix_rows.append( + { + "position": position + 1, + "count": int(prefix_probs.numel()), + "predicted_mean": float(prefix_probs.mean()), + "target_mean": float(prefix_targets.mean()), + "ece": _ece(prefix_probs, prefix_targets), + "brier": float(torch.square(prefix_probs - prefix_targets).mean()), + } + ) + + return { + "schema_version": 1, + "dump_count": len(rows), + "rank": int(rows[0].get("rank", -1)), + "step_range": [int(rows[0]["step"]), int(rows[-1]["step"])], + "calibration_fraction": calibration_fraction, + "sequential_temperatures": temperatures, + "per_conditional_position": position_rows, + "per_prefix_position": prefix_rows, + "source_files": [row["_path"] for row in rows], + } + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--input-dir", type=Path, required=True) + parser.add_argument("--rank", type=int, default=0) + parser.add_argument("--calibration-fraction", type=float, default=0.7) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + if not 0.0 < args.calibration_fraction < 1.0: + parser.error("--calibration-fraction must be in (0, 1)") + return args + + +def main() -> None: + args = parse_args() + rows = _load_rows(args.input_dir, args.rank) + result = analyze(rows, args.calibration_fraction) + args.output.parent.mkdir(parents=True, exist_ok=True) + + def json_safe(value: Any) -> Any: + if isinstance(value, float) and not math.isfinite(value): + return None + if isinstance(value, dict): + return {key: json_safe(item) for key, item in value.items()} + if isinstance(value, list): + return [json_safe(item) for item in value] + return value + + args.output.write_text( + json.dumps(json_safe(result), ensure_ascii=False, indent=2) + "\n" + ) + print( + json.dumps({key: result[key] for key in ("dump_count", "rank", "step_range")}) + ) + + +if __name__ == "__main__": + main() 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_dsv4_dspark_long_context.py b/benchmarks/benchmark_dsv4_dspark_long_context.py new file mode 100644 index 0000000000..d0a1287046 --- /dev/null +++ b/benchmarks/benchmark_dsv4_dspark_long_context.py @@ -0,0 +1,576 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Measure DSpark verification and output health across long contexts. + +The benchmark targets an already-running OpenAI-compatible server. It records +stream timestamps together with vLLM speculative metrics so decode speed, +verification-round cost, and draft acceptance can be compared per context +bucket instead of being conflated with prefill latency. +""" + +from __future__ import annotations + +import argparse +import http.client +import json +import statistics +import time +from pathlib import Path +from typing import Any +from urllib.parse import urlparse + +import regex as re +from transformers import AutoTokenizer + +_PROFILE_INTERVAL_MARKER = "SM70 spec runner profile interval_avg_ms " +_PROFILE_FIELD_PATTERN = re.compile(r"([a-z_]+)=([-+]?\d+(?:\.\d+)?(?:[eE][-+]?\d+)?)") + + +def _parse_lengths(raw: str) -> list[int]: + values = sorted({int(value.strip()) for value in raw.split(",") if value.strip()}) + if not values or values[0] <= 0: + raise ValueError("--lengths must contain positive integers") + return values + + +def _template_token_ids(tokenizer: Any, messages: list[dict[str, str]]) -> list[int]: + try: + encoded = tokenizer.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=True, + enable_thinking=False, + ) + except (TypeError, ValueError): + # DeepSeek V4's custom tokenizer mode supplies the template in the + # server rather than tokenizer_config.json. The constant allowance is + # only for fitting below the target; server-reported usage remains the + # authoritative prompt length in every result row. + encoded = [0] * 32 + for message in messages: + encoded.extend( + tokenizer.encode(message["content"], add_special_tokens=False) + ) + if isinstance(encoded, dict): + encoded = encoded["input_ids"] + elif hasattr(encoded, "input_ids"): + encoded = encoded.input_ids + if encoded and isinstance(encoded[0], list): + encoded = encoded[0] + return list(encoded) + + +def _filler_unit(index: int) -> str: + return ( + f"Record {index:06d}: unrelated inventory values " + f"{index * 17 % 9973}, {index * 31 % 7919}, {index * 43 % 6151}. " + "This record is only deterministic long-context filler and contains " + "no answer to the final instruction.\n" + ) + + +def _build_messages(unit_count: int, marker: str) -> list[dict[str, str]]: + filler = "".join(_filler_unit(index) for index in range(unit_count)) + instruction = ( + "\nFINAL INSTRUCTION: Ignore the inventory records. Begin the response " + f"with the exact marker {marker} on its own line, then write a concise " + "technical paragraph explaining why bounded CUDA-graph context buckets " + "reduce long-context decode latency. Do not repeat the marker." + ) + return [ + { + "role": "system", + "content": ( + "Follow the final user instruction exactly. Long repeated context " + "must not cause repetition, garbling, or loss of the suffix." + ), + }, + {"role": "user", "content": filler + instruction}, + ] + + +def _fit_messages( + tokenizer: Any, + target_tokens: int, + marker: str, +) -> tuple[list[dict[str, str]], int, int]: + sample_count = 64 + sample_tokens = len( + _template_token_ids(tokenizer, _build_messages(sample_count, marker)) + ) + base_tokens = len(_template_token_ids(tokenizer, _build_messages(0, marker))) + tokens_per_unit = max(1.0, (sample_tokens - base_tokens) / sample_count) + low = 0 + high = max(1, int(target_tokens / tokens_per_unit * 1.2) + 32) + best_messages = _build_messages(0, marker) + best_tokens = len(_template_token_ids(tokenizer, best_messages)) + best_units = 0 + for _ in range(24): + if low > high: + break + middle = (low + high) // 2 + messages = _build_messages(middle, marker) + prompt_tokens = len(_template_token_ids(tokenizer, messages)) + if prompt_tokens <= target_tokens: + best_messages = messages + best_tokens = prompt_tokens + best_units = middle + low = middle + 1 + else: + high = middle - 1 + return best_messages, best_tokens, best_units + + +def _metric_total(text: str, name: str) -> float: + return sum( + float(line.rsplit(" ", 1)[-1]) + for line in text.splitlines() + if line.startswith(name + "{") or line.startswith(name + " ") + ) + + +def _metrics_snapshot(host: str, port: int, timeout: int) -> dict[str, Any]: + connection = http.client.HTTPConnection(host, port, timeout=timeout) + try: + connection.request("GET", "/metrics") + response = connection.getresponse() + raw = response.read() + finally: + connection.close() + if response.status != 200: + raise RuntimeError(f"metrics endpoint returned HTTP {response.status}") + text = raw.decode("utf-8", errors="replace") + positions: dict[int, float] = {} + prefix = "vllm:spec_decode_num_accepted_tokens_per_pos_total{" + for line in text.splitlines(): + if not line.startswith(prefix): + continue + match = re.search(r'position="(\d+)"', line) + if match: + position = int(match.group(1)) + positions[position] = positions.get(position, 0.0) + float( + line.rsplit(" ", 1)[-1] + ) + return { + "rounds": _metric_total(text, "vllm:spec_decode_num_drafts_total"), + "proposed": _metric_total(text, "vllm:spec_decode_num_draft_tokens_total"), + "accepted": _metric_total(text, "vllm:spec_decode_num_accepted_tokens_total"), + "positions": positions, + } + + +def _longest_same_token_run(token_ids: list[int]) -> int: + longest = 0 + current = 0 + previous: int | None = None + for token_id in token_ids: + if token_id == previous: + current += 1 + else: + previous = token_id + current = 1 + longest = max(longest, current) + return longest + + +def _profile_log_offset(profile_log: Path | None) -> int | None: + if profile_log is None: + return None + try: + return profile_log.stat().st_size + except FileNotFoundError: + return 0 + + +def _read_verifier_profile( + profile_log: Path | None, + offset: int | None, + verifier_rows: int, +) -> dict[str, Any]: + if profile_log is None or offset is None: + return {"available": False, "reason": "profile log was not requested"} + try: + with profile_log.open("rb") as profile_file: + profile_file.seek(offset) + text = profile_file.read().decode("utf-8", errors="replace") + except FileNotFoundError: + return {"available": False, "reason": "profile log does not exist"} + + intervals: list[dict[str, Any]] = [] + for line in text.splitlines(): + if _PROFILE_INTERVAL_MARKER not in line: + continue + fields = { + key: float(value) + for key, value in _PROFILE_FIELD_PATTERN.findall( + line.split(_PROFILE_INTERVAL_MARKER, 1)[1] + ) + } + if int(fields.get("num_tokens", -1)) != verifier_rows: + continue + required = ("target_forward", "target_logits", "target_rejection_sample") + if any(name not in fields for name in required): + continue + interval = dict(fields) + for name in ( + "calls", + "interval_calls", + "interval_spec_steps", + "num_tokens", + "num_reqs", + ): + if name in interval: + interval[name] = int(interval[name]) + interval["complete_verifier_ms"] = sum(fields[name] for name in required) + intervals.append(interval) + + # The first/last profiler window can straddle prefill or request teardown. + # Long sweeps run with a short profile interval and retain only complete + # interior windows when enough samples are available. + if len(intervals) >= 4: + steady_intervals = intervals[1:-1] + elif len(intervals) == 3: + steady_intervals = intervals[1:] + else: + steady_intervals = intervals + median_names = ( + "target_forward", + "target_logits", + "target_rejection_sample", + "complete_verifier_ms", + "draft_total", + ) + medians = { + name: statistics.median( + float(interval[name]) for interval in steady_intervals if name in interval + ) + for name in median_names + if any(name in interval for interval in steady_intervals) + } + return { + "available": bool(intervals), + "profile_log": str(profile_log), + "verifier_rows": verifier_rows, + "raw_interval_count": len(intervals), + "steady_interval_count": len(steady_intervals), + "medians_ms": medians, + "intervals": intervals, + } + + +def _stream_request( + *, + host: str, + port: int, + model: str, + messages: list[dict[str, str]], + output_tokens: int, + temperature: float, + top_p: float, + seed: int, + timeout: int, +) -> dict[str, Any]: + payload = { + "model": model, + "messages": messages, + "max_tokens": output_tokens, + "temperature": temperature, + "top_p": top_p, + "seed": seed, + "ignore_eos": True, + "stream": True, + "stream_options": {"include_usage": True}, + "chat_template_kwargs": {"enable_thinking": False}, + } + connection = http.client.HTTPConnection(host, port, timeout=timeout) + encoded = json.dumps(payload, ensure_ascii=False).encode("utf-8") + started_ns = time.perf_counter_ns() + first_token_ns: int | None = None + last_token_ns: int | None = None + content_parts: list[str] = [] + reasoning_parts: list[str] = [] + usage: dict[str, Any] = {} + finish_reason: str | None = None + try: + connection.request( + "POST", + "/v1/chat/completions", + body=encoded, + headers={"Content-Type": "application/json"}, + ) + response = connection.getresponse() + if response.status != 200: + raw = response.read() + raise RuntimeError( + f"request returned HTTP {response.status}: " + f"{raw.decode('utf-8', errors='replace')}" + ) + for raw_line in response: + arrival_ns = time.perf_counter_ns() + line = raw_line.decode("utf-8", errors="replace").strip() + if not line.startswith("data: "): + continue + data = line[6:] + if data == "[DONE]": + break + event = json.loads(data) + if event.get("usage"): + usage = event["usage"] + choices = event.get("choices") or [] + if not choices: + continue + choice = choices[0] + if choice.get("finish_reason") is not None: + finish_reason = choice["finish_reason"] + delta = choice.get("delta") or {} + content = delta.get("content") or "" + reasoning = delta.get("reasoning_content") or delta.get("reasoning") or "" + if content or reasoning: + first_token_ns = first_token_ns or arrival_ns + last_token_ns = arrival_ns + if content: + content_parts.append(content) + if reasoning: + reasoning_parts.append(reasoning) + finally: + finished_ns = time.perf_counter_ns() + connection.close() + + completion_tokens = int(usage.get("completion_tokens") or 0) + decode_seconds = ( + (last_token_ns - first_token_ns) / 1_000_000_000 + if first_token_ns is not None and last_token_ns is not None + else None + ) + return { + "request": payload, + "usage": usage, + "finish_reason": finish_reason, + "wall_seconds": (finished_ns - started_ns) / 1_000_000_000, + "ttft_seconds": ( + (first_token_ns - started_ns) / 1_000_000_000 + if first_token_ns is not None + else None + ), + "pure_decode_seconds": decode_seconds, + "pure_decode_tokens_per_second": ( + (completion_tokens - 1) / decode_seconds + if completion_tokens > 1 and decode_seconds + else None + ), + "content": "".join(content_parts), + "reasoning": "".join(reasoning_parts), + } + + +def _run_case( + *, + tokenizer: Any, + host: str, + port: int, + model: str, + target_tokens: int, + output_tokens: int, + temperature: float, + top_p: float, + seed: int, + timeout: int, + num_speculative_tokens: int, + profile_log: Path | None, +) -> dict[str, Any]: + marker = f"DSV4-CONTEXT-{target_tokens}-OK" + messages, local_prompt_tokens, unit_count = _fit_messages( + tokenizer, target_tokens, marker + ) + profile_offset = _profile_log_offset(profile_log) + before = _metrics_snapshot(host, port, timeout) + result = _stream_request( + host=host, + port=port, + model=model, + messages=messages, + output_tokens=output_tokens, + temperature=temperature, + top_p=top_p, + seed=seed, + timeout=timeout, + ) + after = _metrics_snapshot(host, port, timeout) + rounds = after["rounds"] - before["rounds"] + proposed = after["proposed"] - before["proposed"] + accepted = after["accepted"] - before["accepted"] + accepted_per_position = [ + after["positions"].get(position, 0.0) - before["positions"].get(position, 0.0) + for position in range(num_speculative_tokens) + ] + verifier_profile = ( + _read_verifier_profile( + profile_log, + profile_offset, + num_speculative_tokens + 1, + ) + if num_speculative_tokens + else {"available": False, "reason": "speculative decoding is disabled"} + ) + output_text = f"{result['reasoning']}\n{result['content']}" + output_ids = tokenizer.encode(output_text, add_special_tokens=False) + completion_tokens = int(result["usage"].get("completion_tokens") or 0) + result.update( + { + "target_prompt_tokens": target_tokens, + "local_prompt_tokens": local_prompt_tokens, + "unit_count": unit_count, + "marker": marker, + "output_health": { + "marker_in_content": marker in result["content"], + "replacement_characters": output_text.count("\ufffd"), + "unique_output_tokens": len(set(output_ids)), + "longest_same_token_run": _longest_same_token_run(output_ids), + }, + "speculative_metrics": { + "rounds": rounds, + "proposed_tokens": proposed, + "accepted_tokens": accepted, + "accepted_fraction_of_proposed": ( + accepted / proposed if proposed else None + ), + "mean_emitted_tokens_per_round": ( + 1.0 + accepted / rounds if rounds else None + ), + "round_wall_milliseconds": ( + result["pure_decode_seconds"] * 1000.0 / rounds + if result["pure_decode_seconds"] and rounds + else None + ), + "unconditional_acceptance_per_position": [ + value / rounds if rounds else None + for value in accepted_per_position + ], + }, + "verifier_profile": verifier_profile, + } + ) + result["passed"] = bool( + completion_tokens == output_tokens + and result["output_health"]["marker_in_content"] + and result["output_health"]["replacement_characters"] == 0 + and result["output_health"]["unique_output_tokens"] >= 8 + and result["output_health"]["longest_same_token_run"] < 16 + ) + return result + + +def _summary(records: list[dict[str, Any]]) -> dict[str, Any]: + speeds = [ + float(record["pure_decode_tokens_per_second"]) + for record in records + if record["pure_decode_tokens_per_second"] is not None + ] + round_costs = [ + float(record["speculative_metrics"]["round_wall_milliseconds"]) + for record in records + if record["speculative_metrics"]["round_wall_milliseconds"] is not None + ] + verifier_costs = [ + float(record["verifier_profile"]["medians_ms"]["complete_verifier_ms"]) + for record in records + if record["verifier_profile"].get("medians_ms", {}).get("complete_verifier_ms") + is not None + ] + return { + "cases": len(records), + "passed": all(record["passed"] for record in records), + "pure_decode_tps_min": min(speeds) if speeds else None, + "pure_decode_tps_median": statistics.median(speeds) if speeds else None, + "round_wall_ms_max": max(round_costs) if round_costs else None, + "complete_verifier_ms_median": ( + statistics.median(verifier_costs) if verifier_costs else None + ), + "complete_verifier_ms_max": max(verifier_costs) if verifier_costs else None, + } + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", default="http://127.0.0.1:8000") + parser.add_argument("--model", required=True) + parser.add_argument("--model-path", type=Path, required=True) + parser.add_argument("--lengths", default="1024,4096,16384,65536,131072,252000") + parser.add_argument("--output-tokens", type=int, default=128) + parser.add_argument("--temperature", type=float, default=1.0) + parser.add_argument("--top-p", type=float, default=1.0) + parser.add_argument("--seed", type=int, default=20260824) + parser.add_argument("--num-speculative-tokens", type=int, default=7) + parser.add_argument("--profile-log", type=Path) + parser.add_argument("--timeout", type=int, default=1800) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + + if not 0 <= args.num_speculative_tokens <= 7: + raise ValueError("--num-speculative-tokens must be in [0, 7]") + + parsed = urlparse(args.base_url) + if parsed.scheme != "http" or not parsed.hostname: + raise ValueError("--base-url must be an http URL") + host = parsed.hostname + port = parsed.port or 80 + tokenizer = AutoTokenizer.from_pretrained( + str(args.model_path), + local_files_only=True, + trust_remote_code=True, + ) + + args.out.parent.mkdir(parents=True, exist_ok=True) + records: list[dict[str, Any]] = [] + for target_tokens in _parse_lengths(args.lengths): + record = _run_case( + tokenizer=tokenizer, + host=host, + port=port, + model=args.model, + target_tokens=target_tokens, + output_tokens=args.output_tokens, + temperature=args.temperature, + top_p=args.top_p, + seed=args.seed, + timeout=args.timeout, + num_speculative_tokens=args.num_speculative_tokens, + profile_log=args.profile_log, + ) + records.append(record) + contract = { + key: str(value) if isinstance(value, Path) else value + for key, value in vars(args).items() + } + payload = { + "contract": contract, + "records": records, + "summary": _summary(records), + } + args.out.write_text( + json.dumps(payload, ensure_ascii=False, indent=2) + "\n", + encoding="utf-8", + ) + print( + json.dumps( + { + "target_prompt_tokens": target_tokens, + "usage": record["usage"], + "pure_decode_tokens_per_second": record[ + "pure_decode_tokens_per_second" + ], + "speculative_metrics": record["speculative_metrics"], + "verifier_profile": record["verifier_profile"], + "output_health": record["output_health"], + "passed": record["passed"], + }, + ensure_ascii=False, + ), + flush=True, + ) + + return 0 if all(record["passed"] for record in records) else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/benchmark_dsv4_gsm8k_api.py b/benchmarks/benchmark_dsv4_gsm8k_api.py new file mode 100644 index 0000000000..49f19a22b7 --- /dev/null +++ b/benchmarks/benchmark_dsv4_gsm8k_api.py @@ -0,0 +1,231 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Run a deterministic, sequential GSM8K gate against a vLLM API server.""" + +from __future__ import annotations + +import argparse +import hashlib +import http.client +import json +import time +from decimal import Decimal, InvalidOperation +from pathlib import Path +from typing import Any +from urllib.parse import urlparse + +import regex as re + +_NUMBER_RE = re.compile(r"(?"] + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as source: + for chunk in iter(lambda: source.read(1 << 20), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _load_jsonl(path: Path) -> list[dict[str, Any]]: + return [ + json.loads(line) + for line in path.read_text(encoding="utf-8").splitlines() + if line and not line.startswith("#") + ] + + +def _last_integer(text: str) -> int | None: + numbers = _NUMBER_RE.findall(text) + if not numbers: + return None + try: + value = Decimal(numbers[-1].replace(",", "")) + except InvalidOperation: + return None + if not value.is_finite() or value != value.to_integral_value(): + return None + return int(value) + + +def _post_completion( + *, + host: str, + port: int, + timeout: int, + payload: dict[str, Any], +) -> tuple[int, dict[str, Any], float]: + connection = http.client.HTTPConnection(host, port, timeout=timeout) + started = time.perf_counter() + try: + connection.request( + "POST", + "/v1/completions", + body=json.dumps(payload).encode("utf-8"), + headers={"Content-Type": "application/json"}, + ) + response = connection.getresponse() + body = response.read() + finally: + connection.close() + elapsed = time.perf_counter() - started + try: + decoded = json.loads(body) + except json.JSONDecodeError as exc: + raise RuntimeError( + f"completion returned non-JSON HTTP {response.status}: " + f"{body.decode('utf-8', errors='replace')}" + ) from exc + return response.status, decoded, elapsed + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", default="http://127.0.0.1:8000") + parser.add_argument("--model", required=True) + parser.add_argument("--train", type=Path, required=True) + parser.add_argument("--test", type=Path, required=True) + parser.add_argument("--out", type=Path, required=True) + parser.add_argument("--limit", type=int, default=64) + parser.add_argument("--few-shot", type=int, default=5) + parser.add_argument("--max-tokens", type=int, default=256) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--timeout", type=int, default=1200) + parser.add_argument("--min-correct", type=int, default=0) + parser.add_argument("--max-invalid", type=int, default=0) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + if args.limit <= 0 or args.few_shot < 0: + raise ValueError("--limit must be positive and --few-shot non-negative") + if args.min_correct < 0 or args.max_invalid < 0: + raise ValueError("quality thresholds must be non-negative") + + parsed = urlparse(args.base_url) + if parsed.scheme != "http" or not parsed.hostname: + raise ValueError("--base-url must be an http URL") + host = parsed.hostname + port = parsed.port or 80 + + train = _load_jsonl(args.train) + test = _load_jsonl(args.test)[: args.limit] + if len(train) < args.few_shot: + raise RuntimeError( + f"training set has {len(train)} rows, fewer than {args.few_shot} shots" + ) + if len(test) < args.limit: + raise RuntimeError(f"test set has only {len(test)} rows, need {args.limit}") + + few_shot = "".join( + f"Question: {row['question']}\nAnswer: {row['answer']}\n\n" + for row in train[: args.few_shot] + ) + rows: list[dict[str, Any]] = [] + started = time.perf_counter() + for index, row in enumerate(test): + prompt = few_shot + f"Question: {row['question']}\nAnswer:" + request = { + "model": args.model, + "prompt": prompt, + "temperature": 0.0, + "top_p": 1.0, + "max_tokens": args.max_tokens, + "stop": _STOP_SEQUENCES, + "seed": args.seed, + } + status, response, elapsed = _post_completion( + host=host, + port=port, + timeout=args.timeout, + payload=request, + ) + choices = response.get("choices") or [] + choice = choices[0] if choices else {} + output = choice.get("text") or "" + expected = _last_integer(row["answer"]) + if expected is None: + raise ValueError(f"GSM8K item {index} has no integral reference answer") + predicted = _last_integer(output) + rows.append( + { + "index": index, + "question": row["question"], + "expected": expected, + "predicted": predicted, + "correct": predicted is not None and predicted == expected, + "invalid": predicted is None, + "status": status, + "elapsed_seconds": elapsed, + "finish_reason": choice.get("finish_reason"), + "usage": response.get("usage"), + "output": output, + } + ) + if status != 200: + raise RuntimeError(f"GSM8K item {index} returned HTTP {status}: {response}") + + wall_seconds = time.perf_counter() - started + correct = sum(bool(row["correct"]) for row in rows) + invalid = sum(bool(row["invalid"]) for row in rows) + completion_tokens = sum( + int((row.get("usage") or {}).get("completion_tokens") or 0) for row in rows + ) + passed = correct >= args.min_correct and invalid <= args.max_invalid + result = { + "contract": { + "base_url": args.base_url, + "model": args.model, + "questions": args.limit, + "few_shot": args.few_shot, + "temperature": 0.0, + "top_p": 1.0, + "seed": args.seed, + "max_tokens": args.max_tokens, + "strictly_sequential": True, + "train_selection": "first_n", + "test_selection": "first_n", + "prompt_format": "gsm8k_question_answer_v1", + "answer_normalization": "last_signed_integral_decimal_v1", + "stop_sequences": _STOP_SEQUENCES, + "min_correct": args.min_correct, + "max_invalid": args.max_invalid, + }, + "input_manifest": { + "train": {"path": str(args.train), "sha256": _sha256(args.train)}, + "test": {"path": str(args.test), "sha256": _sha256(args.test)}, + }, + "correct": correct, + "accuracy": correct / len(rows), + "invalid": invalid, + "wall_seconds": wall_seconds, + "completion_tokens": completion_tokens, + "aggregate_output_tokens_per_second": ( + completion_tokens / wall_seconds if wall_seconds else None + ), + "passed": passed, + "rows": rows, + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text( + json.dumps(result, ensure_ascii=False, indent=2) + "\n", + encoding="utf-8", + ) + print( + json.dumps( + { + "correct": correct, + "samples": len(rows), + "invalid": invalid, + "passed": passed, + } + ), + flush=True, + ) + return 0 if passed else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/benchmark_dsv4_quality_api.py b/benchmarks/benchmark_dsv4_quality_api.py new file mode 100644 index 0000000000..5eba6b30ea --- /dev/null +++ b/benchmarks/benchmark_dsv4_quality_api.py @@ -0,0 +1,635 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Run paired HumanEval and LongBench-subset gates against a vLLM API.""" + +from __future__ import annotations + +import argparse +import ctypes +import functools +import hashlib +import http.client +import json +import os +import resource +import statistics +import subprocess +import sys +import tempfile +from pathlib import Path +from typing import Any +from urllib.parse import urlparse + +import regex as re +from transformers import AutoTokenizer + +NO_CHAT_DATASETS = {"trec", "triviaqa", "samsum", "lsht", "lcc", "repobench-p"} + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as source: + for chunk in iter(lambda: source.read(1 << 20), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _post_json( + *, + host: str, + port: int, + path: str, + payload: dict[str, Any], + timeout: int, +) -> dict[str, Any]: + connection = http.client.HTTPConnection(host, port, timeout=timeout) + try: + connection.request( + "POST", + path, + body=json.dumps(payload, ensure_ascii=False).encode("utf-8"), + headers={"Content-Type": "application/json"}, + ) + response = connection.getresponse() + raw = response.read() + finally: + connection.close() + if response.status != 200: + raise RuntimeError( + f"{path} returned HTTP {response.status}: " + f"{raw.decode('utf-8', errors='replace')}" + ) + return json.loads(raw) + + +def _chat( + *, + host: str, + port: int, + model: str, + content: str, + max_tokens: int, + timeout: int, +) -> tuple[str, dict[str, Any]]: + result = _post_json( + host=host, + port=port, + path="/v1/chat/completions", + timeout=timeout, + payload={ + "model": model, + "messages": [{"role": "user", "content": content}], + "max_tokens": max_tokens, + "temperature": 0.0, + "top_p": 1.0, + "seed": 20260824, + "chat_template_kwargs": {"enable_thinking": False}, + }, + ) + message = result["choices"][0]["message"] + return message.get("content") or "", result + + +def _completion( + *, + host: str, + port: int, + model: str, + prompt: str, + max_tokens: int, + timeout: int, +) -> tuple[str, dict[str, Any]]: + result = _post_json( + host=host, + port=port, + path="/v1/completions", + timeout=timeout, + payload={ + "model": model, + "prompt": prompt, + "max_tokens": max_tokens, + "temperature": 0.0, + "top_p": 1.0, + "seed": 20260824, + }, + ) + return result["choices"][0].get("text") or "", result + + +def _extract_python(text: str, entry_point: str) -> str: + blocks = re.findall(r"```(?:python)?\s*\n(.*?)```", text, flags=re.DOTALL) + for block in blocks: + if re.search(rf"\bdef\s+{re.escape(entry_point)}\s*\(", block): + return block.strip() + if blocks: + return max(blocks, key=len).strip() + return text.strip() + + +def _sandbox_preexec(sandbox_dir: str) -> None: + """Deny filesystem/network access before executing generated code.""" + libc = ctypes.CDLL(None, use_errno=True) + libc.syscall.restype = ctypes.c_long + libc.prctl.argtypes = [ + ctypes.c_int, + ctypes.c_ulong, + ctypes.c_ulong, + ctypes.c_ulong, + ctypes.c_ulong, + ] + libc.prctl.restype = ctypes.c_int + sys_create_ruleset = 444 + sys_add_rule = 445 + sys_restrict_self = 446 + abi = libc.syscall(sys_create_ruleset, 0, 0, 1) + if abi < 1: + raise OSError(ctypes.get_errno(), "Landlock is unavailable") + + class RulesetAttr(ctypes.Structure): + _fields_ = [ + ("handled_access_fs", ctypes.c_uint64), + ("handled_access_net", ctypes.c_uint64), + ] + + class PathBeneathAttr(ctypes.Structure): + _pack_ = 1 + _fields_ = [ + ("allowed_access", ctypes.c_uint64), + ("parent_fd", ctypes.c_int32), + ] + + handled_fs = (1 << 13) - 1 + if abi >= 2: + handled_fs |= 1 << 13 + if abi >= 3: + handled_fs |= 1 << 14 + ruleset_attr = RulesetAttr(handled_fs, 3 if abi >= 4 else 0) + ruleset_fd = libc.syscall( + sys_create_ruleset, + ctypes.byref(ruleset_attr), + ctypes.sizeof(ruleset_attr), + 0, + ) + if ruleset_fd < 0: + raise OSError(ctypes.get_errno(), "landlock_create_ruleset") + + read_only = (1 << 0) | (1 << 2) | (1 << 3) + allowed_paths = ( + ("/usr", read_only), + ("/lib", read_only), + ("/lib64", read_only), + ("/dev/null", (1 << 1) | (1 << 2)), + ("/dev/urandom", 1 << 2), + (sandbox_dir, handled_fs), + ) + for path, allowed_access in allowed_paths: + parent_fd = os.open(path, os.O_PATH | os.O_CLOEXEC) + path_attr = PathBeneathAttr(allowed_access, parent_fd) + if ( + libc.syscall( + sys_add_rule, + ruleset_fd, + 1, + ctypes.byref(path_attr), + 0, + ) + < 0 + ): + raise OSError(ctypes.get_errno(), f"landlock_add_rule({path})") + os.close(parent_fd) + + if libc.prctl(38, 1, 0, 0, 0) < 0: + raise OSError(ctypes.get_errno(), "PR_SET_NO_NEW_PRIVS") + if libc.syscall(sys_restrict_self, ruleset_fd, 0) < 0: + raise OSError(ctypes.get_errno(), "landlock_restrict_self") + os.close(ruleset_fd) + + seccomp = ctypes.CDLL("libseccomp.so.2", use_errno=True) + seccomp.seccomp_init.argtypes = [ctypes.c_uint32] + seccomp.seccomp_init.restype = ctypes.c_void_p + seccomp.seccomp_syscall_resolve_name.argtypes = [ctypes.c_char_p] + seccomp.seccomp_syscall_resolve_name.restype = ctypes.c_int + seccomp.seccomp_rule_add.argtypes = [ + ctypes.c_void_p, + ctypes.c_uint32, + ctypes.c_int, + ctypes.c_uint, + ] + seccomp.seccomp_rule_add.restype = ctypes.c_int + seccomp.seccomp_load.argtypes = [ctypes.c_void_p] + seccomp.seccomp_load.restype = ctypes.c_int + seccomp.seccomp_release.argtypes = [ctypes.c_void_p] + seccomp.seccomp_release.restype = None + context = seccomp.seccomp_init(0x7FFF0000) + if not context: + raise RuntimeError("seccomp_init failed") + deny_with_eperm = 0x00050000 | 1 + for name in ( + b"socket", + b"socketpair", + b"connect", + b"bind", + b"listen", + b"accept", + b"accept4", + b"sendto", + b"sendmsg", + b"recvfrom", + b"recvmsg", + b"kill", + b"tkill", + b"tgkill", + b"pidfd_open", + b"pidfd_getfd", + b"pidfd_send_signal", + b"ptrace", + b"process_vm_readv", + b"process_vm_writev", + b"kcmp", + b"fork", + b"vfork", + b"clone", + b"clone3", + b"unshare", + b"setns", + b"mount", + b"umount2", + b"pivot_root", + b"bpf", + b"perf_event_open", + b"userfaultfd", + b"keyctl", + b"add_key", + b"request_key", + b"open_by_handle_at", + b"name_to_handle_at", + ): + syscall_number = seccomp.seccomp_syscall_resolve_name(name) + if ( + syscall_number >= 0 + and seccomp.seccomp_rule_add(context, deny_with_eperm, syscall_number, 0) + != 0 + ): + raise RuntimeError(f"seccomp_rule_add failed for {name!r}") + if seccomp.seccomp_load(context) != 0: + raise RuntimeError("seccomp_load failed") + seccomp.seccomp_release(context) + + resource.setrlimit(resource.RLIMIT_CPU, (5, 5)) + resource.setrlimit(resource.RLIMIT_AS, (256 << 20, 256 << 20)) + resource.setrlimit(resource.RLIMIT_FSIZE, (1 << 20, 1 << 20)) + resource.setrlimit(resource.RLIMIT_NPROC, (16, 16)) + + +def _execute_humaneval(program: str, timeout: int = 8) -> dict[str, Any]: + with tempfile.TemporaryDirectory(prefix="dsv4-humaneval-") as sandbox_dir: + output_path = Path(sandbox_dir) / "output.log" + try: + with output_path.open("wb") as output_file: + completed = subprocess.run( + ["/usr/bin/python3", "-I", "-S", "-"], + input=program.encode("utf-8"), + stdout=output_file, + stderr=subprocess.STDOUT, + cwd=sandbox_dir, + env={"PATH": "/usr/bin", "HOME": sandbox_dir}, + preexec_fn=functools.partial(_sandbox_preexec, sandbox_dir), + timeout=timeout, + ) + output = output_path.read_bytes()[-2000:].decode("utf-8", errors="replace") + return { + "passed": completed.returncode == 0, + "returncode": completed.returncode, + "output": output, + } + except subprocess.TimeoutExpired as exc: + return { + "passed": False, + "returncode": None, + "output": f"timeout: {exc}", + } + + +def _run_humaneval( + *, + host: str, + port: int, + model: str, + tasks_path: Path, + limit: int, + timeout: int, +) -> dict[str, Any]: + tasks = json.loads(tasks_path.read_text(encoding="utf-8"))[:limit] + records: list[dict[str, Any]] = [] + for task in tasks: + prompt = ( + "Complete the Python function below. Return runnable Python code " + "only, without Markdown or explanation.\n\n" + task["prompt"] + ) + response, raw = _chat( + host=host, + port=port, + model=model, + content=prompt, + max_tokens=768, + timeout=timeout, + ) + code = _extract_python(response, task["entry_point"]) + if re.search(rf"\bdef\s+{re.escape(task['entry_point'])}\s*\(", code): + candidate = code + else: + candidate = task["prompt"] + code + program = candidate + "\n" + task["test"] + f"\ncheck({task['entry_point']})\n" + execution = _execute_humaneval(program) + records.append( + { + "task_id": task["task_id"], + "passed": execution["passed"], + "execution": execution, + "response": response, + "usage": raw.get("usage"), + } + ) + passed = sum(int(record["passed"]) for record in records) + return { + "samples": len(records), + "passed": passed, + "pass_at_1": passed / len(records) if records else None, + "records": records, + } + + +def _truncate_middle(tokenizer: Any, prompt: str, max_input_tokens: int) -> str: + token_ids = tokenizer.encode(prompt, add_special_tokens=False) + if len(token_ids) <= max_input_tokens: + return prompt + half = max_input_tokens // 2 + return tokenizer.decode( + token_ids[:half], skip_special_tokens=True + ) + tokenizer.decode(token_ids[-half:], skip_special_tokens=True) + + +def _select_longbench_rows( + rows: list[dict[str, Any]], limit: int +) -> list[dict[str, Any]]: + rows = [row for row in rows if int(row.get("length", 0)) >= 8000] or rows + if len(rows) <= limit: + return rows + if limit == 1: + return [rows[len(rows) // 2]] + last = len(rows) - 1 + return [rows[round(index * last / (limit - 1))] for index in range(limit)] + + +def _run_longbench( + *, + tokenizer: Any, + host: str, + port: int, + model: str, + data_dir: Path, + longbench_root: Path, + datasets: list[str], + limit: int, + max_input_tokens: int, + timeout: int, +) -> dict[str, Any]: + config_dir = longbench_root / "config" + prompt_formats = json.loads( + (config_dir / "dataset2prompt.json").read_text(encoding="utf-8") + ) + max_output_tokens = json.loads( + (config_dir / "dataset2maxlen.json").read_text(encoding="utf-8") + ) + sys.path.insert(0, str(longbench_root.resolve())) + try: + from eval import dataset2metric # type: ignore + except ModuleNotFoundError as exc: + raise RuntimeError( + "LongBench metric dependencies are missing; install rouge, jieba, " + "and fuzzywuzzy into PYTHONPATH" + ) from exc + + by_dataset: dict[str, Any] = {} + for dataset in datasets: + rows = [] + for source_index, line in enumerate( + (data_dir / f"{dataset}.jsonl").read_text(encoding="utf-8").splitlines() + ): + if not line.strip(): + continue + row = json.loads(line) + row["_source_index"] = source_index + rows.append(row) + selected = _select_longbench_rows(rows, limit) + records: list[dict[str, Any]] = [] + for row in selected: + prompt = prompt_formats[dataset].format(**row) + prompt = _truncate_middle(tokenizer, prompt, max_input_tokens) + if dataset in NO_CHAT_DATASETS: + prediction, raw = _completion( + host=host, + port=port, + model=model, + prompt=prompt, + max_tokens=int(max_output_tokens[dataset]), + timeout=timeout, + ) + else: + prediction, raw = _chat( + host=host, + port=port, + model=model, + content=prompt, + max_tokens=int(max_output_tokens[dataset]), + timeout=timeout, + ) + score = max( + float( + dataset2metric[dataset]( + prediction, + answer, + all_classes=row.get("all_classes", []), + ) + ) + for answer in row["answers"] + ) + records.append( + { + "source_index": row["_source_index"], + "length": row.get("length"), + "score": score, + "prediction": prediction, + "answers": row["answers"], + "usage": raw.get("usage"), + } + ) + scores = [float(record["score"]) for record in records] + by_dataset[dataset] = { + "samples": len(records), + "score": 100.0 * statistics.mean(scores) if scores else None, + "records": records, + } + dataset_scores = [ + float(result["score"]) + for result in by_dataset.values() + if result["score"] is not None + ] + return { + "datasets": by_dataset, + "average_score": statistics.mean(dataset_scores) if dataset_scores else None, + } + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", default="http://127.0.0.1:8000") + parser.add_argument("--model", required=True) + parser.add_argument("--model-path", type=Path, required=True) + parser.add_argument("--out", type=Path, required=True) + parser.add_argument( + "--humaneval-tasks", + type=Path, + default=Path( + "bench_results/ddtree_humaneval_sweep_20260625/humaneval_tasks.json" + ), + ) + parser.add_argument("--humaneval-limit", type=int, default=32) + parser.add_argument( + "--longbench-data-dir", + type=Path, + default=Path("benchmark-data/longbench/data"), + ) + parser.add_argument( + "--longbench-root", + type=Path, + default=Path("third_party/LongBench/LongBench"), + ) + parser.add_argument( + "--longbench-datasets", + default="hotpotqa,multifieldqa_zh,gov_report,lcc", + ) + parser.add_argument("--longbench-limit", type=int, default=4) + parser.add_argument("--longbench-max-input-tokens", type=int, default=32768) + parser.add_argument("--timeout", type=int, default=1200) + parser.add_argument("--min-humaneval-passes", type=int, default=0) + parser.add_argument("--min-longbench-average", type=float, default=0.0) + args = parser.parse_args() + + parsed = urlparse(args.base_url) + if parsed.scheme != "http" or not parsed.hostname: + raise ValueError("--base-url must be an http URL") + host = parsed.hostname + port = parsed.port or 80 + tokenizer = AutoTokenizer.from_pretrained( + str(args.model_path), + local_files_only=True, + trust_remote_code=True, + ) + + datasets = [ + value.strip() for value in args.longbench_datasets.split(",") if value.strip() + ] + input_manifest = { + "humaneval_tasks": { + "path": str(args.humaneval_tasks), + "sha256": _sha256(args.humaneval_tasks), + }, + "longbench_prompt_config": { + "path": str(args.longbench_root / "config/dataset2prompt.json"), + "sha256": _sha256(args.longbench_root / "config/dataset2prompt.json"), + }, + "longbench_maxlen_config": { + "path": str(args.longbench_root / "config/dataset2maxlen.json"), + "sha256": _sha256(args.longbench_root / "config/dataset2maxlen.json"), + }, + "longbench_datasets": { + dataset: { + "path": str(args.longbench_data_dir / f"{dataset}.jsonl"), + "sha256": _sha256(args.longbench_data_dir / f"{dataset}.jsonl"), + } + for dataset in datasets + }, + } + + human_eval = _run_humaneval( + host=host, + port=port, + model=args.model, + tasks_path=args.humaneval_tasks, + limit=args.humaneval_limit, + timeout=args.timeout, + ) + longbench = _run_longbench( + tokenizer=tokenizer, + host=host, + port=port, + model=args.model, + data_dir=args.longbench_data_dir, + longbench_root=args.longbench_root, + datasets=datasets, + limit=args.longbench_limit, + max_input_tokens=args.longbench_max_input_tokens, + timeout=args.timeout, + ) + passed = bool( + human_eval["passed"] >= args.min_humaneval_passes + and longbench["average_score"] is not None + and longbench["average_score"] >= args.min_longbench_average + ) + result = { + "contract": { + key: str(value) if isinstance(value, Path) else value + for key, value in vars(args).items() + }, + "evaluation_contract": { + "version": "dsv4_api_quality_v1", + "strictly_sequential": True, + "temperature": 0.0, + "top_p": 1.0, + "seed": 20260824, + "human_eval": { + "selection": "first_n", + "limit": args.humaneval_limit, + "endpoint": "chat_completions", + "enable_thinking": False, + "max_tokens": 768, + "execution": "landlock_seccomp_python_isolated_v1", + }, + "longbench": { + "datasets": datasets, + "selection": "length_ge_8000_evenly_spaced_v1", + "limit_per_dataset": args.longbench_limit, + "max_input_tokens": args.longbench_max_input_tokens, + "truncation": "middle_first_half_last_half_v1", + "endpoint_policy": "official_no_chat_dataset_set_v1", + "max_output_tokens": "dataset2maxlen_sha256_manifest", + }, + }, + "input_manifest": input_manifest, + "human_eval": human_eval, + "longbench": longbench, + "passed": passed, + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text( + json.dumps(result, ensure_ascii=False, indent=2) + "\n", + encoding="utf-8", + ) + print( + json.dumps( + { + "human_eval_passes": human_eval["passed"], + "human_eval_samples": human_eval["samples"], + "longbench_average": longbench["average_score"], + "passed": passed, + }, + ensure_ascii=False, + ), + flush=True, + ) + return 0 if passed else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/benchmark_sm70_cublas_gemm_algorithms.py b/benchmarks/benchmark_sm70_cublas_gemm_algorithms.py new file mode 100644 index 0000000000..f8c33ed8f1 --- /dev/null +++ b/benchmarks/benchmark_sm70_cublas_gemm_algorithms.py @@ -0,0 +1,251 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Screen legacy cuBLAS Tensor Core algorithms on exact SM70 QKV shapes.""" + +from __future__ import annotations + +import argparse +import ctypes +import json +import statistics +from collections.abc import Callable +from pathlib import Path +from typing import Any + +import torch + +_CUBLAS_STATUS_SUCCESS = 0 +_CUBLAS_OP_N = 0 +_CUDA_R_16F = 2 +_CUBLAS_COMPUTE_32F = 68 +_CUBLAS_TENSOR_OP_MATH = 1 +_CUBLAS_GEMM_DEFAULT_TENSOR_OP = 99 +_CUBLAS_GEMM_ALGO0_TENSOR_OP = 100 + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--json-out", type=Path, required=True) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--m", type=int, default=8000) + parser.add_argument("--k", type=int, default=5120) + parser.add_argument("--n", type=int, nargs="+", default=[4096, 3584]) + parser.add_argument("--warmup", type=int, default=3) + parser.add_argument("--iters", type=int, default=10) + parser.add_argument("--trials", type=int, default=5) + parser.add_argument("--seed", type=int, default=20260823) + return parser.parse_args() + + +def _event_trials( + launch: Callable[[], None], warmup: int, iters: int, trials: int +) -> dict[str, Any]: + for _ in range(warmup): + launch() + torch.accelerator.synchronize() + samples: list[float] = [] + for _ in range(trials): + start = torch.Event(enable_timing=True) + end = torch.Event(enable_timing=True) + start.record() + for _ in range(iters): + launch() + end.record() + end.synchronize() + samples.append(float(start.elapsed_time(end) / iters)) + return { + "median_ms": float(statistics.median(samples)), + "mean_ms": float(statistics.mean(samples)), + "min_ms": min(samples), + "max_ms": max(samples), + "samples_ms": samples, + } + + +class _Cublas: + def __init__(self, stream: int) -> None: + self.lib = ctypes.CDLL("libcublas.so.12") + self.handle = ctypes.c_void_p() + self.lib.cublasCreate_v2.argtypes = [ctypes.POINTER(ctypes.c_void_p)] + self.lib.cublasCreate_v2.restype = ctypes.c_int + self.lib.cublasDestroy_v2.argtypes = [ctypes.c_void_p] + self.lib.cublasDestroy_v2.restype = ctypes.c_int + self.lib.cublasSetStream_v2.argtypes = [ctypes.c_void_p, ctypes.c_void_p] + self.lib.cublasSetStream_v2.restype = ctypes.c_int + self.lib.cublasSetMathMode.argtypes = [ctypes.c_void_p, ctypes.c_int] + self.lib.cublasSetMathMode.restype = ctypes.c_int + self.lib.cublasGemmEx.argtypes = [ + ctypes.c_void_p, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_int, + ctypes.c_int, + ctypes.c_void_p, + ctypes.c_int, + ctypes.c_int, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ] + self.lib.cublasGemmEx.restype = ctypes.c_int + self._check(self.lib.cublasCreate_v2(ctypes.byref(self.handle)), "create") + self._check( + self.lib.cublasSetStream_v2(self.handle, ctypes.c_void_p(stream)), + "set stream", + ) + self._check( + self.lib.cublasSetMathMode(self.handle, _CUBLAS_TENSOR_OP_MATH), + "set math mode", + ) + + @staticmethod + def _check(status: int, operation: str) -> None: + if status != _CUBLAS_STATUS_SUCCESS: + raise RuntimeError(f"cuBLAS {operation} failed with status {status}.") + + def gemm( + self, + output: torch.Tensor, + inputs: torch.Tensor, + weight: torch.Tensor, + algo: int, + ) -> int: + m, k = inputs.shape + weight_k, n = weight.shape + if weight_k != k or output.shape != (m, n): + raise ValueError("Incompatible row-major GEMM shapes.") + alpha = ctypes.c_float(1.0) + beta = ctypes.c_float(0.0) + # Row-major C=MxN is column-major C^T=NxM. Compute + # C^T = weight^T(NxK) * inputs^T(KxM) without materializing transposes. + return int( + self.lib.cublasGemmEx( + self.handle, + _CUBLAS_OP_N, + _CUBLAS_OP_N, + n, + m, + k, + ctypes.byref(alpha), + ctypes.c_void_p(weight.data_ptr()), + _CUDA_R_16F, + n, + ctypes.c_void_p(inputs.data_ptr()), + _CUDA_R_16F, + k, + ctypes.byref(beta), + ctypes.c_void_p(output.data_ptr()), + _CUDA_R_16F, + n, + _CUBLAS_COMPUTE_32F, + algo, + ) + ) + + def close(self) -> None: + if self.handle.value is not None: + self._check(self.lib.cublasDestroy_v2(self.handle), "destroy") + self.handle = ctypes.c_void_p() + + +def _run_shape( + args: argparse.Namespace, device: torch.device, cublas: _Cublas, n: int +) -> dict[str, Any]: + generator = torch.Generator(device=device).manual_seed(args.seed + n) + inputs = torch.randn( + (args.m, args.k), generator=generator, device=device, dtype=torch.float16 + ).mul_(0.02) + weight = torch.randn( + (args.k, n), generator=generator, device=device, dtype=torch.float16 + ).mul_(0.02) + reference = torch.empty((args.m, n), device=device, dtype=torch.float16) + torch.mm(inputs, weight, out=reference) + torch.accelerator.synchronize() + + algorithms = [_CUBLAS_GEMM_DEFAULT_TENSOR_OP] + list( + range(_CUBLAS_GEMM_ALGO0_TENSOR_OP, _CUBLAS_GEMM_ALGO0_TENSOR_OP + 16) + ) + results: list[dict[str, Any]] = [] + for algo in algorithms: + output = torch.empty_like(reference) + + def launch(output: torch.Tensor = output, algo: int = algo) -> None: + status = cublas.gemm(output, inputs, weight, algo) + if status != _CUBLAS_STATUS_SUCCESS: + raise RuntimeError(f"cuBLAS algo {algo} returned status {status}.") + + try: + timing = _event_trials(launch, args.warmup, args.iters, args.trials) + launch() + torch.accelerator.synchronize() + difference = (output.float() - reference.float()).abs() + results.append( + { + "algo": algo, + "status": "success", + "timing": timing, + "useful_tflops": float( + 2 * args.m * n * args.k / (timing["median_ms"] * 1e-3) / 1e12 + ), + "max_abs_diff": float(difference.max().item()), + "mean_abs_diff": float(difference.mean().item()), + "equal_elements": int((output == reference).sum().item()), + "elements": output.numel(), + } + ) + except RuntimeError as error: + torch.accelerator.synchronize() + results.append({"algo": algo, "status": "rejected", "error": str(error)}) + return { + "m": args.m, + "n": n, + "k": args.k, + "algorithms": sorted( + results, + key=lambda item: item.get("timing", {}).get("median_ms", float("inf")), + ), + } + + +def main() -> int: + args = _parse_args() + device = torch.device(args.device) + if torch.cuda.get_device_capability(device) != (7, 0): + raise RuntimeError("This benchmark requires SM70/V100.") + torch.accelerator.set_device_index(device) + cublas = _Cublas(torch.cuda.current_stream(device).cuda_stream) + try: + shapes = [_run_shape(args, device, cublas, n) for n in args.n] + finally: + cublas.close() + payload = { + "environment": { + "device": torch.cuda.get_device_name(device), + "torch": torch.__version__, + "torch_cuda": torch.version.cuda, + "m": args.m, + "k": args.k, + "n": args.n, + "warmup": args.warmup, + "iters": args.iters, + "trials": args.trials, + "seed": args.seed, + }, + "shapes": shapes, + } + args.json_out.parent.mkdir(parents=True, exist_ok=True) + args.json_out.write_text(json.dumps(payload, indent=2) + "\n") + print(json.dumps(payload, indent=2)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/benchmark_sm70_custom_all_reduce.py b/benchmarks/benchmark_sm70_custom_all_reduce.py index 88d6398437..156484f987 100644 --- a/benchmarks/benchmark_sm70_custom_all_reduce.py +++ b/benchmarks/benchmark_sm70_custom_all_reduce.py @@ -51,15 +51,19 @@ def _make_input( elif pattern == "random_small": generator = torch.Generator(device="cuda") generator.manual_seed(seed + rank) - tensor = (torch.rand(size, device="cuda", generator=generator) - 0.5).to( - dtype - ) + tensor = (torch.rand(size, device="cuda", generator=generator) - 0.5).to(dtype) elif pattern == "model_like": generator = torch.Generator(device="cuda") generator.manual_seed(seed + rank) tensor = (torch.randn(size, device="cuda", generator=generator) * 0.03).to( dtype ) + elif pattern == "positive_zero": + tensor = torch.zeros(size, device="cuda", dtype=dtype) + elif pattern == "signed_zero": + tensor = torch.zeros(size, device="cuda", dtype=dtype) + if rank % 2: + tensor = -tensor else: raise ValueError(f"unsupported pattern: {pattern}") return tensor @@ -132,7 +136,13 @@ def _reference_custom_order( def _compare(out: torch.Tensor, ref: torch.Tensor) -> dict[str, Any]: diff = (out.float() - ref.float()).abs() mismatch = out != ref + bit_dtype = torch.int16 if out.element_size() == 2 else torch.int32 + out_bits = out.view(bit_dtype) + ref_bits = ref.view(bit_dtype) + bit_mismatch = out_bits != ref_bits + negative_zero = -32768 if out.element_size() == 2 else -2147483648 nonzero = int(mismatch.sum().item()) + bitwise_nonzero = int(bit_mismatch.sum().item()) first_mismatch: dict[str, Any] | None = None if nonzero: idx = int(torch.nonzero(mismatch, as_tuple=False)[0].item()) @@ -144,6 +154,10 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> dict[str, Any]: } return { "equal": bool(torch.equal(out, ref)), + "bitwise_equal": bitwise_nonzero == 0, + "bitwise_nonzero_count": bitwise_nonzero, + "out_negative_zero_count": int((out_bits == negative_zero).sum().item()), + "ref_negative_zero_count": int((ref_bits == negative_zero).sum().item()), "max_diff": float(diff.max().item()), "mean_diff": float(diff.mean().item()), "nonzero_count": nonzero, @@ -292,12 +306,15 @@ def _make_gemma_rms_norm_inputs( shared_generator = torch.Generator(device="cuda") shared_generator.manual_seed(seed + 1009) - residual = torch.randn( - (tokens, _GEMMA_RMS_NORM_HIDDEN_SIZE), - device="cuda", - dtype=torch.float32, - generator=shared_generator, - ) * 0.03 + residual = ( + torch.randn( + (tokens, _GEMMA_RMS_NORM_HIDDEN_SIZE), + device="cuda", + dtype=torch.float32, + generator=shared_generator, + ) + * 0.03 + ) residual = residual.to(residual_dtype) weight = ( torch.randn( @@ -570,6 +587,11 @@ def _check_graph_multi( for i in range(multi_count) ] refs = [_reference_all_reduce(inp) for inp in inputs] + num_bytes = size * torch.empty((), dtype=dtype).element_size() + algo = _custom_allreduce_algo(ca.world_size, ca.fully_connected, num_bytes) + custom_order_refs = [ + _reference_custom_order(inp, algo, ca.world_size) for inp in inputs + ] graph_timing: dict[str, float] | None = None if timing_iterations: outputs, graph_timing = _time_custom_ar_graph( @@ -578,6 +600,9 @@ def _check_graph_multi( else: outputs = _run_graph(ca, inputs, graph_replays) comparisons = [_compare(out, ref) for out, ref in zip(outputs, refs)] + custom_order_comparisons = [ + _compare(out, ref) for out, ref in zip(outputs, custom_order_refs) + ] max_diff = max(item["max_diff"] for item in comparisons) nonzero_count = sum(item["nonzero_count"] for item in comparisons) first_bad = next( @@ -591,14 +616,27 @@ def _check_graph_multi( return { "mode": "graph_multi", "size": size, - "bytes": size * torch.empty((), dtype=dtype).element_size(), + "bytes": num_bytes, "dtype": str(dtype).replace("torch.", ""), "pattern": pattern, + "custom_allreduce_algo": algo, "multi_count": multi_count, "all_equal": all(item["equal"] for item in comparisons), "max_diff": float(max_diff), "nonzero_count": int(nonzero_count), "first_bad": first_bad, + "custom_order_all_equal": all( + item["equal"] for item in custom_order_comparisons + ), + "custom_order_bitwise_all_equal": all( + item["bitwise_equal"] for item in custom_order_comparisons + ), + "custom_order_bitwise_nonzero_count": sum( + item["bitwise_nonzero_count"] for item in custom_order_comparisons + ), + "custom_order_max_diff": max( + item["max_diff"] for item in custom_order_comparisons + ), "graph_timing": graph_timing, } @@ -835,11 +873,11 @@ def main() -> None: if item["pattern"] not in require_exact_patterns: continue if item["mode"] == "graph_multi": - equal = item["all_equal"] + equal = item["custom_order_bitwise_all_equal"] comparison = item else: - equal = item["comparison"]["equal"] - comparison = item["comparison"] + comparison = item["custom_order_comparison"] + equal = comparison["bitwise_equal"] if not equal: required_failures.append( { @@ -857,9 +895,11 @@ def main() -> None: "custom_allreduce_disabled": False, "fully_connected": ca.fully_connected, "max_size_bytes": ca.max_size, - "sm70_tp4_m5_ar_threads": os.environ.get( - "VLLM_SM70_TP4_M5_AR_THREADS" + "sm70_tp4_push_buffer_registered": ( + ca.sm70_tp4_push_buffer_ptrs is not None ), + "sm70_tp4_m5_ar_threads": os.environ.get("VLLM_SM70_TP4_M5_AR_THREADS"), + "sm70_tp4_push_allreduce": os.environ.get("VLLM_SM70_TP4_PUSH_ALLREDUCE"), "required_exact_patterns": sorted(require_exact_patterns), "required_exact_passed": not required_failures, "required_failures": required_failures[:16], diff --git a/benchmarks/benchmark_sm70_d256_exact_dense.py b/benchmarks/benchmark_sm70_d256_exact_dense.py new file mode 100644 index 0000000000..6b32dbab23 --- /dev/null +++ b/benchmarks/benchmark_sm70_d256_exact_dense.py @@ -0,0 +1,215 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Benchmark one SM70 D256 exact-dense prefill extension. + +Run control and candidate shared libraries in separate processes because both +register the same Torch operator namespace. A saved control tensor provides +the elementwise numerical gate for a candidate process. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import statistics +from pathlib import Path +from typing import Any + +import torch + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--extension", type=Path, required=True) + parser.add_argument("--query-len", type=int, default=8192) + parser.add_argument("--kv-len", type=int, default=131072) + parser.add_argument("--heads-q", type=int, default=6) + parser.add_argument("--heads-kv", type=int, default=1) + parser.add_argument("--head-dim", type=int, default=256) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--iters", type=int, default=5) + parser.add_argument("--trials", type=int, default=7) + parser.add_argument("--seed", type=int, default=20260824) + parser.add_argument("--reference-out", type=Path) + parser.add_argument("--write-reference", action="store_true") + parser.add_argument("--json-out", type=Path, required=True) + return parser.parse_args() + + +def _digest_bytes(data: bytes) -> str: + return hashlib.sha256(data).hexdigest() + + +def _digest_tensor(tensor: torch.Tensor) -> str: + return _digest_bytes(tensor.contiguous().view(torch.uint8).numpy().tobytes()) + + +def _event_trials( + launch: Any, + *, + warmup: int, + iters: int, + trials: int, +) -> dict[str, Any]: + for _ in range(warmup): + launch() + torch.accelerator.synchronize() + + samples: list[float] = [] + for _ in range(trials): + start = torch.Event(enable_timing=True) + end = torch.Event(enable_timing=True) + start.record() + for _ in range(iters): + launch() + end.record() + end.synchronize() + samples.append(float(start.elapsed_time(end) / iters)) + return { + "median_ms": float(statistics.median(samples)), + "mean_ms": float(statistics.mean(samples)), + "min_ms": min(samples), + "max_ms": max(samples), + "samples_ms": samples, + } + + +def _causal_tflops( + *, + query_len: int, + kv_len: int, + heads_q: int, + head_dim: int, + elapsed_ms: float, +) -> float: + prefix_len = kv_len - query_len + mean_visible_keys = prefix_len + (query_len + 1) / 2.0 + flops = 4.0 * query_len * mean_visible_keys * heads_q * head_dim + return flops / (elapsed_ms * 1e9) + + +@torch.inference_mode() +def main() -> int: + args = _parse_args() + extension = args.extension.resolve() + if not extension.is_file(): + raise FileNotFoundError(extension) + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + if torch.cuda.get_device_capability() != (7, 0): + raise RuntimeError("This benchmark requires SM70/V100") + if args.head_dim != 256: + raise ValueError("The exact operator requires head_dim=256") + if args.query_len <= 0 or args.query_len > args.kv_len: + raise ValueError("Expected 0 < query_len <= kv_len") + if args.query_len % 64 or args.kv_len % 32: + raise ValueError("Query/KV lengths must be divisible by 64/32") + if args.heads_q % args.heads_kv: + raise ValueError("heads_q must be divisible by heads_kv") + if min(args.warmup, args.iters, args.trials) <= 0: + raise ValueError("warmup, iters, and trials must be positive") + if args.write_reference and args.reference_out is None: + raise ValueError("--write-reference requires --reference-out") + + torch.ops.load_library(str(extension)) + op = torch.ops._vllm_fa2_C.sm70_d256_splitd_n32_dense_fwd + + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(args.seed) + query = torch.randn( + (1, args.query_len, args.heads_q, args.head_dim), + dtype=torch.float16, + device=device, + generator=generator, + ) + key = torch.randn( + (1, args.kv_len, args.heads_kv, args.head_dim), + dtype=torch.float16, + device=device, + generator=generator, + ) + value = torch.randn( + (1, args.kv_len, args.heads_kv, args.head_dim), + dtype=torch.float16, + device=device, + generator=generator, + ) + output = torch.empty_like(query) + scale = args.head_dim**-0.5 + + def launch() -> None: + op(query, key, value, output, scale, True) + + timing = _event_trials( + launch, + warmup=args.warmup, + iters=args.iters, + trials=args.trials, + ) + launch() + torch.accelerator.synchronize() + output_cpu = output.cpu() + + comparison = None + if args.reference_out is not None and args.reference_out.exists(): + reference = torch.load( + args.reference_out, + map_location="cpu", + weights_only=True, + ) + if reference.shape != output_cpu.shape or reference.dtype != output_cpu.dtype: + raise RuntimeError("Reference shape or dtype does not match output") + difference = (reference.float() - output_cpu.float()).abs() + comparison = { + "reference_sha256": _digest_tensor(reference), + "max_abs_diff": float(difference.max().item()), + "mean_abs_diff": float(difference.mean().item()), + "equal_elements": int((reference == output_cpu).sum().item()), + "elements": output_cpu.numel(), + } + if args.write_reference: + assert args.reference_out is not None + args.reference_out.parent.mkdir(parents=True, exist_ok=True) + torch.save(output_cpu, args.reference_out) + + median_ms = float(timing["median_ms"]) + payload = { + "environment": { + "device": torch.cuda.get_device_name(), + "torch": torch.__version__, + "torch_cuda": torch.version.cuda, + "extension": str(extension), + "extension_sha256": _digest_bytes(extension.read_bytes()), + "query_len": args.query_len, + "kv_len": args.kv_len, + "heads_q": args.heads_q, + "heads_kv": args.heads_kv, + "head_dim": args.head_dim, + "dtype": "float16", + "causal": True, + "softmax_scale": scale, + "seed": args.seed, + "warmup": args.warmup, + "iters": args.iters, + "trials": args.trials, + }, + "timing": timing, + "causal_tflops": _causal_tflops( + query_len=args.query_len, + kv_len=args.kv_len, + heads_q=args.heads_q, + head_dim=args.head_dim, + elapsed_ms=median_ms, + ), + "output_sha256": _digest_tensor(output_cpu), + "comparison": comparison, + } + args.json_out.parent.mkdir(parents=True, exist_ok=True) + args.json_out.write_text(json.dumps(payload, indent=2) + "\n") + print(json.dumps(payload, indent=2), flush=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/benchmark_sm70_decode.py b/benchmarks/benchmark_sm70_decode.py index 1e50bbfa23..bf6d364c19 100644 --- a/benchmarks/benchmark_sm70_decode.py +++ b/benchmarks/benchmark_sm70_decode.py @@ -30,7 +30,13 @@ def _module_file(module_name: str) -> str | None: module = sys.modules.get(module_name) if module is not None: return getattr(module, "__file__", None) - spec = importlib.util.find_spec(module_name) + try: + spec = importlib.util.find_spec(module_name) + except (AttributeError, ImportError, ValueError): + # Optional extension parents can raise while find_spec imports them. + # Diagnostics must report an unavailable extension, not discard an + # otherwise successful benchmark at result-serialization time. + return None return spec.origin if spec is not None else None @@ -163,6 +169,7 @@ def _tracked_env() -> dict[str, str]: "VLLM_FLASH_", "FLASH_ATTN", "FLA_", + "PREFIX_", "CUDA_MODULE_LOADING", "TORCH_CUDA_ARCH_LIST", "TORCHINDUCTOR_", @@ -683,6 +690,10 @@ def _sm70_attention_policy(kv_cache_dtype: Any) -> dict[str, Any]: "VLLM_FLASH_V100_XQA_E5M2_PAIR_LOAD", True, ) + e5m2_batch_wide_load = _env_bool( + "VLLM_FLASH_V100_XQA_E5M2_BATCH_WIDE_LOAD", + True, + ) decode_dynamic_partitions = _env_bool( "VLLM_FLASH_V100_DECODE_DYNAMIC_PARTITIONS", True, @@ -798,6 +809,10 @@ def _sm70_attention_policy(kv_cache_dtype: Any) -> dict[str, Any]: "VLLM_FLASH_V100_XQA_E5M2_PAIR_LOAD" ), "e5m2_pair_load_effective": e5m2_pair_load, + "VLLM_FLASH_V100_XQA_E5M2_BATCH_WIDE_LOAD": os.environ.get( + "VLLM_FLASH_V100_XQA_E5M2_BATCH_WIDE_LOAD" + ), + "e5m2_batch_wide_load_effective": e5m2_batch_wide_load, "exact_mtp5_fp8_p1024_dual_cta_policy": ( smallq_decode_use_xqa and mtp5_xqa_dual_cta diff --git a/benchmarks/benchmark_sm70_dflash2_gdn_metadata.py b/benchmarks/benchmark_sm70_dflash2_gdn_metadata.py new file mode 100644 index 0000000000..7c138529f0 --- /dev/null +++ b/benchmarks/benchmark_sm70_dflash2_gdn_metadata.py @@ -0,0 +1,373 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Microbenchmark DFlash2 target GDN metadata construction on SM70.""" + +from __future__ import annotations + +import argparse +import json +import os +import statistics +import time +from dataclasses import dataclass +from types import SimpleNamespace +from typing import Any + +import torch + +from vllm import envs +from vllm.triton_utils import triton +from vllm.v1.attention.backends.gdn_attn import ( + DFlash2GDNGroupDescriptor, + _dflash2_gdn_group_metadata_kernel, + _get_ddtree_gdn_fast_common_buffers, + build_gdn_spec_decode_state_contract, + prepare_dflash2_gdn_group_metadata, +) +from vllm.v1.attention.backends.utils import PAD_SLOT_ID +from vllm.v1.worker.gpu.attn_utils import compute_common_gdn_attn_metadata + + +@dataclass +class _MicroBuilder: + num_spec_state_tokens: int + decode_cudagraph_max_bs: int + spec_state_indices_tensor: torch.Tensor + _ddtree_fast_common_buffers: Any + use_full_cuda_graph: bool = True + vllm_config: Any = None + + def __post_init__(self) -> None: + self.vllm_config = SimpleNamespace( + cache_config=SimpleNamespace(mamba_cache_mode="none") + ) + + +def _percentile(samples: list[float], quantile: float) -> float: + ordered = sorted(samples) + position = (len(ordered) - 1) * quantile + lower = int(position) + upper = min(lower + 1, len(ordered) - 1) + weight = position - lower + return ordered[lower] + (ordered[upper] - ordered[lower]) * weight + + +def _distribution(samples: list[float]) -> dict[str, float]: + return { + "mean": statistics.fmean(samples), + "p50": _percentile(samples, 0.50), + "p90": _percentile(samples, 0.90), + "p99": _percentile(samples, 0.99), + "min": min(samples), + "max": max(samples), + } + + +def _profile_cuda_launches(step: Any) -> int: + from torch.profiler import ProfilerActivity, profile + + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + step() + torch.accelerator.synchronize() + return sum( + 1 + for event in prof.events() + if str(getattr(event, "device_type", "")).lower().endswith("cuda") + ) + + +def _measure(step: Any, repeats: int) -> dict[str, dict[str, float]]: + wall_samples: list[float] = [] + gpu_samples: list[float] = [] + start = torch.Event(enable_timing=True) + end = torch.Event(enable_timing=True) + for _ in range(repeats): + torch.accelerator.synchronize() + wall_start = time.perf_counter_ns() + start.record() + step() + end.record() + end.synchronize() + wall_samples.append((time.perf_counter_ns() - wall_start) / 1e6) + gpu_samples.append(start.elapsed_time(end)) + return { + "synchronized_wall_ms": _distribution(wall_samples), + "cuda_event_ms": _distribution(gpu_samples), + } + + +def _run_case( + batch_size: int, + num_groups: int, + width: int, + warmup: int, + repeats: int, +) -> dict[str, Any]: + device = torch.device("cuda") + num_actual_tokens = batch_size * width + common_buffers = _get_ddtree_gdn_fast_common_buffers( + device, + num_actual_tokens, + width, + ) + builders = [ + _MicroBuilder( + num_spec_state_tokens=width - 1, + decode_cudagraph_max_bs=num_actual_tokens, + spec_state_indices_tensor=torch.empty( + (num_actual_tokens, width), dtype=torch.int32, device=device + ), + _ddtree_fast_common_buffers=common_buffers, + ) + for _ in range(num_groups) + ] + block_tables = tuple( + ( + torch.arange( + batch_size * width, + dtype=torch.int32, + device=device, + ).view(batch_size, width) + + group_id * 10_000 + ).contiguous() + for group_id in range(num_groups) + ) + query_start_loc = torch.arange( + 0, + num_actual_tokens + 1, + width, + dtype=torch.int32, + device=device, + ) + query_start_loc_cpu = query_start_loc.cpu() + draft_counts_cpu = torch.full((batch_size,), width - 1, dtype=torch.int32) + common_metadata = compute_common_gdn_attn_metadata( + num_decode_draft_tokens_cpu=draft_counts_cpu, + query_start_loc=query_start_loc, + query_start_loc_cpu=query_start_loc_cpu, + num_spec_state_tokens=width - 1, + legacy_mixed_decode_routing=False, + ) + assert common_metadata is not None + accepted_values = [3, 5, 7, 2] + accepted = torch.tensor( + [accepted_values[index % len(accepted_values)] for index in range(batch_size)], + dtype=torch.int32, + device=device, + ) + seq_lens = torch.full((batch_size,), 4097, dtype=torch.int32, device=device) + spec_mask_cpu = torch.ones(batch_size, dtype=torch.bool) + spec_mask = common_metadata.spec_sequence_masks + + def legacy_step() -> None: + for group_id, builder in enumerate(builders): + contract = build_gdn_spec_decode_state_contract( + block_table_tensor=block_tables[group_id], + seq_lens=seq_lens, + block_size=16, + num_spec=width - 1, + spec_sequence_masks_cpu=spec_mask_cpu, + num_accepted_tokens=accepted, + current_state_block_ids=None, + is_mamba_cache_all=False, + ) + builder.spec_state_indices_tensor[:batch_size].copy_( + contract.spec_state_indices_tensor, + non_blocking=True, + ) + builder.spec_state_indices_tensor[batch_size:].fill_(PAD_SLOT_ID) + + common_buffers.spec_sequence_masks[:batch_size].copy_( + spec_mask, non_blocking=True + ) + common_buffers.spec_sequence_masks[batch_size:].fill_(False) + common_buffers.spec_token_indx[:num_actual_tokens].copy_( + common_metadata.spec_token_indx, + non_blocking=True, + ) + common_buffers.spec_query_start_loc[: batch_size + 1].copy_( + common_metadata.spec_query_start_loc, + non_blocking=True, + ) + spec_num_query_tokens = common_metadata.spec_query_start_loc[-1] + common_buffers.spec_query_start_loc[batch_size + 1 :].fill_( + spec_num_query_tokens + ) + common_buffers.num_accepted_tokens[:batch_size].copy_( + contract.num_accepted_tokens, + non_blocking=True, + ) + common_buffers.num_accepted_tokens[batch_size:].fill_(1) + common_buffers.spec_state_slot_selectors[:batch_size].copy_( + contract.spec_state_slot_selectors, + non_blocking=True, + ) + common_buffers.spec_state_slot_selectors[batch_size:].fill_(1) + + descriptor: DFlash2GDNGroupDescriptor | None = None + + def fused_step() -> None: + nonlocal descriptor + result = prepare_dflash2_gdn_group_metadata( + builders_by_group=[ + (group_id, builder) for group_id, builder in enumerate(builders) + ], + block_tables=block_tables, + common_gdn_metadata=common_metadata, + num_accepted_tokens=accepted, + num_actual_tokens=num_actual_tokens, + descriptor=descriptor, + ) + if result is None: + raise RuntimeError("fused DFlash2 GDN metadata path was not eligible") + _, descriptor = result + + def fused_kernel_only_step() -> None: + if descriptor is None: + raise RuntimeError("fused descriptor has not been initialized") + block = triton.next_power_of_2( + max(num_actual_tokens * width, num_actual_tokens + 1) + ) + _dflash2_gdn_group_metadata_kernel[(num_groups,)]( + descriptor.block_table_ptrs, + descriptor.state_output_ptrs, + descriptor.block_table_strides, + common_metadata.spec_query_start_loc, + accepted, + accepted, + common_buffers.spec_sequence_masks, + common_buffers.spec_query_start_loc, + common_buffers.num_accepted_tokens, + common_buffers.spec_state_slot_selectors, + batch_size, + num_actual_tokens, + WIDTH=width, + PAD_ID=PAD_SLOT_ID, + BLOCK=block, + num_warps=1, + ) + + def fused_with_host_fence_step() -> None: + fused_step() + if ( + common_metadata.spec_query_start_loc[-1].item() + != common_metadata.num_spec_decode_tokens + ): + raise AssertionError("invalid speculative query span") + + legacy_step() + torch.accelerator.synchronize() + expected_states = [ + builder.spec_state_indices_tensor.clone() for builder in builders + ] + expected_common = ( + common_buffers.spec_sequence_masks.clone(), + common_buffers.spec_query_start_loc.clone(), + common_buffers.num_accepted_tokens.clone(), + common_buffers.spec_state_slot_selectors.clone(), + ) + for builder in builders: + builder.spec_state_indices_tensor.fill_(777) + common_buffers.spec_sequence_masks.fill_(True) + common_buffers.spec_query_start_loc.fill_(-1) + common_buffers.num_accepted_tokens.fill_(-1) + common_buffers.spec_state_slot_selectors.fill_(-1) + fused_step() + torch.accelerator.synchronize() + for actual_builder, expected_state in zip(builders, expected_states): + torch.testing.assert_close( + actual_builder.spec_state_indices_tensor, + expected_state, + rtol=0, + atol=0, + ) + for actual, expected in zip( + ( + common_buffers.spec_sequence_masks, + common_buffers.spec_query_start_loc, + common_buffers.num_accepted_tokens, + common_buffers.spec_state_slot_selectors, + ), + expected_common, + ): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + for _ in range(warmup): + legacy_step() + fused_step() + torch.accelerator.synchronize() + + legacy_launches = _profile_cuda_launches(legacy_step) + fused_launches = _profile_cuda_launches(fused_step) + fused_with_host_fence_launches = _profile_cuda_launches(fused_with_host_fence_step) + fused_kernel_only_launches = _profile_cuda_launches(fused_kernel_only_step) + legacy = _measure(legacy_step, repeats) + fused = _measure(fused_step, repeats) + fused_with_host_fence = _measure(fused_with_host_fence_step, repeats) + fused_kernel_only = _measure(fused_kernel_only_step, repeats) + return { + "batch_size": batch_size, + "num_actual_tokens": num_actual_tokens, + "num_groups": num_groups, + "state_width": width, + "correctness": "bitwise_equal_after_poison", + "legacy_cuda_launches": legacy_launches, + "fused_cuda_launches": fused_launches, + "fused_with_host_fence_cuda_launches": fused_with_host_fence_launches, + "fused_kernel_only_cuda_launches": fused_kernel_only_launches, + "legacy": legacy, + "fused": fused, + "fused_with_host_fence": fused_with_host_fence, + "fused_kernel_only": fused_kernel_only, + "host_fence_wall_cost_p50_ms": ( + fused_with_host_fence["synchronized_wall_ms"]["p50"] + - fused["synchronized_wall_ms"]["p50"] + ), + "fused_python_overhead_p50_ms": ( + fused["synchronized_wall_ms"]["p50"] + - fused_kernel_only["synchronized_wall_ms"]["p50"] + ), + "wall_speedup": ( + legacy["synchronized_wall_ms"]["p50"] / fused["synchronized_wall_ms"]["p50"] + ), + "gpu_speedup": (legacy["cuda_event_ms"]["p50"] / fused["cuda_event_ms"]["p50"]), + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--batch-sizes", type=int, nargs="+", default=[1, 2, 4]) + parser.add_argument("--num-groups", type=int, default=10) + parser.add_argument("--width", type=int, default=8) + parser.add_argument("--warmup", type=int, default=20) + 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 microbenchmark requires SM70") + os.environ["VLLM_SM70_DFLASH2_FUSED_GDN_METADATA"] = "1" + os.environ["VLLM_SM70_DFLASH2_GDN_METADATA_SHADOW"] = "0" + envs.disable_envs_cache() + + results = [ + _run_case( + batch_size=batch_size, + num_groups=args.num_groups, + width=args.width, + warmup=args.warmup, + repeats=args.repeats, + ) + for batch_size in args.batch_sizes + ] + print( + json.dumps( + {"device": torch.cuda.get_device_name(), "cases": results}, + indent=2, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/benchmark_sm70_dflash2_gsm8k.py b/benchmarks/benchmark_sm70_dflash2_gsm8k.py new file mode 100644 index 0000000000..ac92519225 --- /dev/null +++ b/benchmarks/benchmark_sm70_dflash2_gsm8k.py @@ -0,0 +1,484 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Measure target-only or DFlash2 on a fixed local quality subset.""" + +from __future__ import annotations + +import argparse +import ast +import hashlib +import json +import math +import random +import subprocess +import time +from collections import Counter +from pathlib import Path +from typing import Any + +import regex as re +from benchmark_sm70_decode import ( + _diff_spec_metrics, + _hash_ids, + _json_safe, + _module_file, + _module_realpath, + _request_metrics_dict, + _spec_metrics_snapshot, + _tracked_env, +) + +INVALID_ANSWER = -9_999_999 +GSM8K_PROMPT_SUFFIX = ( + "\nPlease reason step by step, and put your final answer within \\boxed{}." +) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--model", type=Path, required=True) + parser.add_argument("--draft-model", type=Path) + parser.add_argument("--dataset", type=Path, required=True) + parser.add_argument( + "--dataset-format", + choices=("gsm8k", "turns"), + default="gsm8k", + help=( + "Input schema. 'gsm8k' scores question/answer rows; 'turns' uses " + "the first preformatted turn and leaves task scoring to a separate " + "evaluator." + ), + ) + parser.add_argument("--out", type=Path, required=True) + parser.add_argument("--mode", choices=("target-only", "dflash"), required=True) + parser.add_argument("--num-questions", type=int, default=64) + parser.add_argument("--start-index", type=int, default=0) + parser.add_argument( + "--dataset-order", + choices=("sequential", "zlab-shuffle42"), + default="sequential", + ) + parser.add_argument("--max-tokens", type=int, default=256) + parser.add_argument("--warmup-tokens", type=int, default=32) + parser.add_argument("--temperature", type=float, default=1.0) + parser.add_argument("--tensor-parallel-size", type=int, default=4) + parser.add_argument("--max-model-len", type=int, default=2048) + 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("--gpu-memory-utilization", type=float, default=0.8) + parser.add_argument( + "--target-kv-cache-dtype", + choices=("auto", "fp8_e5m2"), + default="fp8_e5m2", + ) + parser.add_argument("--enforce-eager", action="store_true") + parser.add_argument( + "--draft-sample-method", + choices=("greedy", "probabilistic"), + default="greedy", + ) + parser.add_argument( + "--draft-attention-backend", + choices=("FLASH_ATTN_V100", "TRITON_ATTN"), + default="FLASH_ATTN_V100", + help="Draft-only attention backend; the target remains on FLASH_ATTN_V100.", + ) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument( + "--request-seed", + type=int, + default=0, + help="Sampling seed for every request; use -1 for server-style random seeds.", + ) + parser.add_argument( + "--cuda-profiler-capture", + action="store_true", + help="Wrap the measured generation in cudaProfilerStart/Stop for nsys.", + ) + return parser.parse_args() + + +def _sha256_file(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 _model_weight_files(model: Path) -> dict[str, str]: + index = model / "model.safetensors.index.json" + if not index.is_file(): + weight = model / "model.safetensors" + return {weight.name: str(weight.resolve())} + weight_map = json.loads(index.read_text())["weight_map"] + return { + name: str((model / name).resolve()) for name in sorted(set(weight_map.values())) + } + + +def _answer_value(text: str) -> int: + numbers = re.findall(r"\d+", text.replace(",", "")) + if not numbers: + return INVALID_ANSWER + try: + return int(ast.literal_eval(numbers[-1])) + except (SyntaxError, ValueError): + return INVALID_ANSWER + + +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 _distribution(values: list[float]) -> dict[str, float | int | None]: + return { + "count": len(values), + "mean": sum(values) / len(values) if values else None, + "p50": _percentile(values, 0.50), + "p90": _percentile(values, 0.90), + "p99": _percentile(values, 0.99), + "min": min(values) if values else None, + "max": max(values) if values else None, + } + + +def _load_rows(args: argparse.Namespace) -> list[tuple[int, dict[str, Any]]]: + rows = [ + json.loads(line) + for line in args.dataset.read_text(encoding="utf-8").splitlines() + if line + ] + indices = list(range(len(rows))) + if args.dataset_order == "zlab-shuffle42": + random.Random(42).shuffle(indices) + selected_indices = indices[args.start_index : args.start_index + args.num_questions] + if len(selected_indices) != args.num_questions: + raise ValueError( + f"Requested {args.num_questions} rows at index {args.start_index}, " + f"but the dataset provided {len(selected_indices)}." + ) + return [(index, rows[index]) for index in selected_indices] + + +def _prompt_content(row: dict[str, Any], dataset_format: str) -> str: + if dataset_format == "gsm8k": + return str(row["question"]) + GSM8K_PROMPT_SUFFIX + turns = row.get("turns") + if not isinstance(turns, list) or not turns or not isinstance(turns[0], str): + raise ValueError("turns dataset rows must contain a non-empty string list") + return turns[0] + + +def _summarize_requests(cases: list[dict[str, Any]]) -> dict[str, Any]: + metric_names = ( + "queued_time", + "first_token_latency", + "prefill_time", + "decode_time", + "steady_decode_tps", + "tpot_seconds", + ) + summary = {} + for metric_name in metric_names: + values = [ + float(case["request_metrics"][metric_name]) + for case in cases + if case["request_metrics"] is not None + and case["request_metrics"].get(metric_name) is not None + ] + summary[metric_name] = _distribution(values) + + prefill_tps = [ + case["prompt_tokens"] / case["request_metrics"]["prefill_time"] + for case in cases + if case["request_metrics"] is not None + and case["request_metrics"].get("prefill_time") + ] + summary["prefill_tokens_per_second"] = _distribution(prefill_tps) + summary["prompt_tokens"] = _distribution( + [float(case["prompt_tokens"]) for case in cases] + ) + summary["output_tokens"] = _distribution( + [float(case["output_tokens"]) for case in cases] + ) + return summary + + +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") + + import torch + import vllm._C as vllm_c + import vllm._C_stable_libtorch as vllm_c_stable + from transformers import AutoTokenizer + + import vllm + from vllm import LLM, SamplingParams + + rows = _load_rows(args) + tokenizer = AutoTokenizer.from_pretrained(str(args.model)) + prompt_contents = [_prompt_content(row, args.dataset_format) for _, row in rows] + prompts = [ + tokenizer.apply_chat_template( + [ + { + "role": "user", + "content": prompt_content, + } + ], + tokenize=False, + add_generation_prompt=True, + enable_thinking=True, + reasoning_effort="xhigh", + ) + for prompt_content in prompt_contents + ] + prompt_token_counts = [len(tokenizer.encode(prompt)) for prompt in prompts] + + speculative_config = None + if args.mode == "dflash": + speculative_config = { + "method": "dflash", + "model": str(args.draft_model), + "revision": "dedf8df68adfb1afeaf7b7480c0a0243108177b4", + "num_speculative_tokens": 7, + "kv_cache_dtype": "auto", + "attention_backend": args.draft_attention_backend, + "draft_sample_method": args.draft_sample_method, + "enforce_eager": args.enforce_eager, + } + + engine_kwargs = { + "model": str(args.model), + "tensor_parallel_size": args.tensor_parallel_size, + "dtype": "half", + "kv_cache_dtype": args.target_kv_cache_dtype, + "attention_backend": "FLASH_ATTN_V100", + "max_model_len": args.max_model_len, + "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, + "disable_log_stats": False, + "enforce_eager": args.enforce_eager, + "seed": args.seed, + "speculative_config": speculative_config, + } + if args.cuda_profiler_capture: + engine_kwargs["profiler_config"] = {"profiler": "cuda"} + + load_started = time.perf_counter() + llm = LLM(**engine_kwargs) + load_seconds = time.perf_counter() - load_started + + request_seed = None if args.request_seed == -1 else args.request_seed + warmup_sampling = SamplingParams( + temperature=args.temperature, + top_p=0.95, + top_k=20, + max_tokens=args.warmup_tokens, + seed=request_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) + ) + 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() + spec_after = _spec_metrics_snapshot(llm) + + cases = [] + for ( + dataset_index, + row, + ), prompt_content, prompt_tokens, output, request_spec in zip( + rows, prompt_contents, prompt_token_counts, outputs, request_spec_metrics + ): + result = output.outputs[0] + token_ids = list(result.token_ids) + if args.dataset_format == "gsm8k": + prediction = _answer_value(result.text) + expected = _answer_value(row["answer"]) + correct = prediction == expected + else: + prediction = None + expected = None + correct = None + cases.append( + { + "dataset_index": dataset_index, + "question": row.get("question"), + "prompt_content": prompt_content, + "expected_answer": expected, + "predicted_answer": prediction, + "correct": correct, + "prompt_tokens": prompt_tokens, + "output_tokens": len(token_ids), + "finish_reason": result.finish_reason, + "stop_reason": result.stop_reason, + "text": result.text, + "token_ids": token_ids, + "token_hash": _hash_ids(token_ids), + "spec_decode_metrics": request_spec, + "request_metrics": _request_metrics_dict( + output.metrics, + len(token_ids), + ), + } + ) + + total_output_tokens = sum(case["output_tokens"] for case in cases) + scored_cases = [case for case in cases if case["correct"] is not None] + correct = sum(bool(case["correct"]) for case in scored_cases) + invalid = sum(case["predicted_answer"] == INVALID_ANSWER for case in scored_cases) + per_request_acceptance_lengths = [ + float(case["spec_decode_metrics"]["acceptance_length"]) + for case in cases + if case["spec_decode_metrics"] is not None + and case["spec_decode_metrics"].get("acceptance_length") is not None + ] + per_request_completion_tokens_per_verification_step = [ + case["output_tokens"] / case["spec_decode_metrics"]["num_drafts"] + for case in cases + if case["spec_decode_metrics"] is not None + and case["spec_decode_metrics"].get("num_drafts", 0) > 0 + ] + c_extension = Path(vllm_c.__file__).resolve() + c_stable_extension = Path(vllm_c_stable.__file__).resolve() + payload = { + "contract": { + "source_sha": subprocess.check_output( + ["git", "rev-parse", "HEAD"], + cwd=Path(__file__).resolve().parents[1], + text=True, + ).strip(), + "mode": args.mode, + "dataset": str(args.dataset), + "dataset_sha256": _sha256_file(args.dataset), + "dataset_format": args.dataset_format, + "start_index": args.start_index, + "num_questions": args.num_questions, + "dataset_order": args.dataset_order, + "model": str(args.model), + "model_config_sha256": _sha256_file(args.model / "config.json"), + "model_index_sha256": ( + _sha256_file(args.model / "model.safetensors.index.json") + if (args.model / "model.safetensors.index.json").is_file() + else None + ), + "model_weight_files": _model_weight_files(args.model), + "model_weights_realpath": str((args.model / "model.safetensors").resolve()), + "draft_model": str(args.draft_model) if args.draft_model else None, + "graph": not args.enforce_eager, + "sequential": args.sequential, + "sampling": { + "temperature": args.temperature, + "top_p": 0.95, + "top_k": 20, + "max_tokens": args.max_tokens, + "seed": request_seed, + "ignore_eos": False, + "thinking": True, + "reasoning_effort": "xhigh", + "prompt_suffix": ( + GSM8K_PROMPT_SUFFIX if args.dataset_format == "gsm8k" else None + ), + }, + "engine_kwargs": engine_kwargs, + }, + "runtime": { + "vllm_version": getattr(vllm, "__version__", None), + "vllm_file": getattr(vllm, "__file__", None), + "torch_version": torch.__version__, + "torch_cuda": torch.version.cuda, + "cuda_device_count": torch.accelerator.device_count(), + "device_capabilities": [ + list(torch.cuda.get_device_capability(index)) + for index in range(torch.accelerator.device_count()) + ], + "c_extension": str(c_extension), + "c_extension_sha256": _sha256_file(c_extension), + "c_stable_extension": str(c_stable_extension), + "c_stable_extension_sha256": _sha256_file(c_stable_extension), + "flash_attn_v100_python": _module_file("flash_attn_v100"), + "flash_attn_v100_cuda": _module_realpath("flash_attn_v100_cuda"), + "tracked_env": _tracked_env(), + "load_seconds": load_seconds, + }, + "results": { + "elapsed_seconds": elapsed_seconds, + "total_output_tokens": total_output_tokens, + "aggregate_output_tokens_per_second": ( + total_output_tokens / elapsed_seconds + ), + "questions_per_second": len(cases) / elapsed_seconds, + "accuracy": correct / len(scored_cases) if scored_cases else None, + "invalid_answer_rate": ( + invalid / len(scored_cases) if scored_cases else None + ), + "finish_reasons": dict( + Counter(str(case["finish_reason"]) for case in cases) + ), + "request_metrics": _summarize_requests(cases), + "spec_decode_metrics": _diff_spec_metrics(spec_before, spec_after), + "per_request_acceptance_length": _distribution( + per_request_acceptance_lengths + ), + "per_request_completion_tokens_per_verification_step": _distribution( + per_request_completion_tokens_per_verification_step + ), + }, + "cases": cases, + } + payload = _json_safe(payload) + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") + printable = {"contract": payload["contract"], "results": payload["results"]} + print(json.dumps(printable, indent=2)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/benchmark_sm70_dflash2_sparse_rejection.py b/benchmarks/benchmark_sm70_dflash2_sparse_rejection.py new file mode 100644 index 0000000000..51cc7ecf35 --- /dev/null +++ b/benchmarks/benchmark_sm70_dflash2_sparse_rejection.py @@ -0,0 +1,243 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Benchmark compact DFlash2 top-k/top-p rejection on one SM70 GPU. + +This isolates target sampling after the TP merge. The separate TP4 compact +logit benchmark measures local top-k and candidate transport. +""" + +from __future__ import annotations + +import argparse +import json +import statistics +from collections.abc import Callable +from pathlib import Path +from typing import Any + +import torch + +from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p +from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import ( + dflash2_sparse_topk_rejection_sample, + rejection_sample, +) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--vocab-size", type=int, default=248320) + parser.add_argument("--num-speculative-steps", type=int, default=7) + parser.add_argument("--target-top-k", type=int, default=20) + parser.add_argument("--draft-top-k", type=int, default=16) + parser.add_argument("--top-p", type=float, default=0.95) + parser.add_argument("--warmup", type=int, default=50) + parser.add_argument("--iters", type=int, default=500) + parser.add_argument("--seed", type=int, default=20260823) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--json-out", type=Path) + return parser.parse_args() + + +def _time_cuda( + operation: Callable[[], Any], + device: torch.device, + warmup: int, + iters: int, +) -> dict[str, float]: + for _ in range(warmup): + operation() + torch.accelerator.synchronize(device) + + events = [ + (torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)) + for _ in range(iters) + ] + for start, end in events: + start.record() + operation() + end.record() + events[-1][1].synchronize() + samples = [float(start.elapsed_time(end)) for start, end in events] + ordered = sorted(samples) + return { + "mean_ms": statistics.fmean(samples), + "p50_ms": statistics.median(samples), + "p90_ms": ordered[int(0.9 * (iters - 1))], + "p99_ms": ordered[int(0.99 * (iters - 1))], + "min_ms": ordered[0], + "max_ms": ordered[-1], + } + + +def main() -> int: + args = _parse_args() + if not 0.0 < args.top_p <= 1.0: + raise ValueError("--top-p must be in (0, 1]") + if not 0 < args.target_top_k <= 64: + raise ValueError("--target-top-k must be in [1, 64]") + if not 0 < args.draft_top_k <= 64: + raise ValueError("--draft-top-k must be in [1, 64]") + + device = torch.device(args.device) + torch.accelerator.set_device_index(device) + capability = torch.cuda.get_device_capability(device) + if capability != (7, 0): + raise RuntimeError(f"Expected SM70, got sm_{capability[0]}{capability[1]}.") + torch.manual_seed(args.seed) + + num_steps = args.num_speculative_steps + num_logits = num_steps + 1 + raw_target = torch.randn( + num_logits, + args.vocab_size, + dtype=torch.float16, + device=device, + ) + target_topk_logits_fp16, target_topk_ids = torch.topk( + raw_target, + k=args.target_top_k, + dim=-1, + ) + target_topk_logits = target_topk_logits_fp16.float() + target_k = torch.full( + (num_logits,), + args.target_top_k, + dtype=torch.int32, + device=device, + ) + target_p_rows = torch.full( + (num_logits,), args.top_p, dtype=torch.float32, device=device + ) + processed_target = apply_top_k_top_p(raw_target.float(), target_k, target_p_rows) + + draft_topk_ids = target_topk_ids[:num_steps, : args.draft_top_k].view( + 1, num_steps, args.draft_top_k + ) + draft_topk_logits = ( + target_topk_logits[:num_steps, : args.draft_top_k] + + torch.randn( + num_steps, + args.draft_top_k, + dtype=torch.float32, + device=device, + ) + * 0.2 + ).view(1, num_steps, args.draft_top_k) + dense_draft = torch.full( + (1, num_steps, args.vocab_size), + -float("inf"), + dtype=torch.float32, + device=device, + ) + dense_draft.scatter_(2, draft_topk_ids, draft_topk_logits) + + draft_sampled = torch.zeros(num_logits, dtype=torch.int64, device=device) + draft_sampled[1:] = draft_topk_ids[0, :, 0] + cu_num_logits = torch.tensor([0, num_logits], dtype=torch.int32, device=device) + pos = torch.arange(num_logits, dtype=torch.int64, device=device) + 32768 + idx_mapping = torch.zeros(1, dtype=torch.int32, device=device) + expanded_idx_mapping = torch.zeros(num_logits, dtype=torch.int32, device=device) + expanded_local_pos = torch.arange(num_logits, dtype=torch.int32, device=device) + temperature = torch.ones(1, dtype=torch.float32, device=device) + top_p_per_req = torch.full((1,), args.top_p, dtype=torch.float32, device=device) + seeds = torch.tensor([args.seed], dtype=torch.int64, device=device) + + def dense_rejection() -> tuple[torch.Tensor, torch.Tensor]: + return rejection_sample( + processed_target, + dense_draft, + draft_sampled, + cu_num_logits, + pos, + idx_mapping, + expanded_idx_mapping, + expanded_local_pos, + temperature, + seeds, + num_steps, + ) + + def sparse_rejection() -> tuple[torch.Tensor, torch.Tensor]: + return dflash2_sparse_topk_rejection_sample( + target_topk_ids, + target_topk_logits, + draft_topk_ids, + draft_topk_logits, + draft_sampled, + cu_num_logits, + pos, + idx_mapping, + temperature, + top_p_per_req, + seeds, + num_steps, + ) + + def dense_topk_topp() -> torch.Tensor: + return apply_top_k_top_p(raw_target.float(), target_k, target_p_rows) + + def dense_finalize() -> tuple[torch.Tensor, torch.Tensor]: + target = dense_topk_topp() + return rejection_sample( + target, + dense_draft, + draft_sampled, + cu_num_logits, + pos, + idx_mapping, + expanded_idx_mapping, + expanded_local_pos, + temperature, + seeds, + num_steps, + ) + + dense_out, dense_count = dense_rejection() + sparse_out, sparse_count = sparse_rejection() + torch.accelerator.synchronize(device) + counts_equal = bool(torch.equal(dense_count, sparse_count)) + valid = torch.arange(num_logits, device=device) < dense_count[0] + tokens_equal = bool(torch.equal(dense_out[0, valid], sparse_out[0, valid])) + + result = { + "device": torch.cuda.get_device_name(device), + "device_capability": list(capability), + "shape": { + "num_logits": num_logits, + "vocab_size": args.vocab_size, + "target_top_k": args.target_top_k, + "draft_top_k": args.draft_top_k, + "top_p": args.top_p, + }, + "correctness": { + "num_sampled_equal": counts_equal, + "valid_tokens_equal": tokens_equal, + "dense_num_sampled": int(dense_count[0].item()), + "sparse_num_sampled": int(sparse_count[0].item()), + }, + "timings": { + "dense_topk_topp_only": _time_cuda( + dense_topk_topp, device, args.warmup, args.iters + ), + "dense_rejection_only": _time_cuda( + dense_rejection, device, args.warmup, args.iters + ), + "sparse_topk_topp_rejection": _time_cuda( + sparse_rejection, device, args.warmup, args.iters + ), + "dense_topk_topp_plus_rejection": _time_cuda( + dense_finalize, device, args.warmup, args.iters + ), + }, + } + text = json.dumps(result, indent=2, sort_keys=True) + print(text) + if args.json_out is not None: + args.json_out.parent.mkdir(parents=True, exist_ok=True) + args.json_out.write_text(text + "\n", encoding="utf-8") + return 0 if counts_equal and tokens_equal else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) 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_fp8_prefill_exact_dense_blas.py b/benchmarks/benchmark_sm70_fp8_prefill_exact_dense_blas.py new file mode 100644 index 0000000000..7f515a0438 --- /dev/null +++ b/benchmarks/benchmark_sm70_fp8_prefill_exact_dense_blas.py @@ -0,0 +1,368 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Benchmark the exact-dense SM70 FP8 prefill path with a selected BLAS backend. + +This loads the real TP-local Qwen3.8 gate-up and down-projection weights. Run +each BLAS backend in a fresh process so global ATen and cuBLAS caches cannot +contaminate the comparison. +""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib.util +import json +import math +import os +import statistics +import sys +from collections.abc import Callable +from pathlib import Path +from typing import Any + +import torch +from safetensors import safe_open + +import vllm + + +def _load_vllm_c_extension() -> None: + if "vllm._C" in sys.modules: + vllm._C = sys.modules["vllm._C"] + return + candidate = os.getenv("VLLM_BENCH_C_EXTENSION") + if candidate is None: + importlib.import_module("vllm._C") + return + spec = importlib.util.spec_from_file_location("vllm._C", candidate) + if spec is None or spec.loader is None: + raise RuntimeError(f"Cannot load candidate extension: {candidate}") + module = importlib.util.module_from_spec(spec) + sys.modules["vllm._C"] = module + vllm._C = module + spec.loader.exec_module(module) + + +_load_vllm_c_extension() + +from vllm import _sm70_ops as sm70_ops # noqa: E402 + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument( + "--model", + type=Path, + required=True, + ) + parser.add_argument("--json-out", type=Path, required=True) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--tp-size", type=int, default=4) + parser.add_argument("--tp-rank", type=int, default=0) + parser.add_argument( + "--cases", + nargs="+", + choices=("gate_up", "down", "attention_output"), + default=("gate_up", "down"), + ) + parser.add_argument("--m", type=int, default=8000) + parser.add_argument( + "--backend", + choices=("exact-dense", "resident-turbomind-f16"), + default="exact-dense", + ) + parser.add_argument( + "--gated-silu", + action="store_true", + help="Benchmark gate_up with the fused interleaved SiLU-and-mul output.", + ) + parser.add_argument( + "--blas-library", + choices=("default", "cublas", "cublaslt"), + required=True, + ) + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--iters", type=int, default=30) + parser.add_argument("--trials", type=int, default=7) + parser.add_argument("--seed", type=int, default=20260823) + parser.add_argument("--cuda-graph", action="store_true") + return parser.parse_args() + + +def _load_keys(model: Path, keys: list[str]) -> dict[str, torch.Tensor]: + index = json.loads((model / "model.safetensors.index.json").read_text()) + weight_map: dict[str, str] = index["weight_map"] + by_file: dict[str, list[str]] = {} + for key in keys: + by_file.setdefault(weight_map[key], []).append(key) + + result: dict[str, torch.Tensor] = {} + for filename, file_keys in by_file.items(): + with safe_open(model / filename, framework="pt", device="cpu") as handle: + for key in file_keys: + result[key] = handle.get_tensor(key) + return result + + +def _column_shard( + weight: torch.Tensor, + scales: torch.Tensor, + tp_size: int, + tp_rank: int, +) -> tuple[torch.Tensor, torch.Tensor]: + n = int(weight.shape[0]) + if n % tp_size: + raise ValueError(f"N={n} is not divisible by TP={tp_size}.") + begin = tp_rank * (n // tp_size) + end = begin + n // tp_size + return ( + weight.view(torch.uint8)[begin:end].contiguous(), + scales[begin // 128 : math.ceil(end / 128)].contiguous(), + ) + + +def _row_shard( + weight: torch.Tensor, + scales: torch.Tensor, + tp_size: int, + tp_rank: int, +) -> tuple[torch.Tensor, torch.Tensor]: + k = int(weight.shape[1]) + if k % tp_size: + raise ValueError(f"K={k} is not divisible by TP={tp_size}.") + begin = tp_rank * (k // tp_size) + end = begin + k // tp_size + return ( + weight.view(torch.uint8)[:, begin:end].contiguous(), + scales[:, begin // 128 : math.ceil(end / 128)].contiguous(), + ) + + +def _load_case( + model: Path, + case: str, + tp_size: int, + tp_rank: int, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor, list[str]]: + if case == "attention_output": + prefixes = ["model.language_model.layers.3.self_attn.o_proj"] + else: + root = "model.language_model.layers.1.mlp" + projections = ("gate_proj", "up_proj") if case == "gate_up" else ("down_proj",) + prefixes = [f"{root}.{projection}" for projection in projections] + keys = [ + key + for prefix in prefixes + for key in (f"{prefix}.weight", f"{prefix}.weight_scale_inv") + ] + loaded = _load_keys(model, keys) + + raw_parts: list[torch.Tensor] = [] + scale_parts: list[torch.Tensor] = [] + fp8_dtype: torch.dtype | None = None + for prefix in prefixes: + weight = loaded[f"{prefix}.weight"] + scales = loaded[f"{prefix}.weight_scale_inv"].float() + fp8_dtype = weight.dtype if fp8_dtype is None else fp8_dtype + if weight.dtype != fp8_dtype: + raise ValueError("Fused gate-up projections must use one FP8 dtype.") + shard = ( + _column_shard(weight, scales, tp_size, tp_rank) + if case == "gate_up" + else _row_shard(weight, scales, tp_size, tp_rank) + ) + raw_parts.append(shard[0]) + scale_parts.append(shard[1]) + + assert fp8_dtype is not None + raw = torch.cat(raw_parts, dim=0) + scales = torch.cat(scale_parts, dim=0) + return ( + raw.to(device).view(fp8_dtype).contiguous(), + scales.to(device).contiguous(), + prefixes, + ) + + +def _event_trials( + launch: Callable[[], None], warmup: int, iters: int, trials: int +) -> dict[str, Any]: + for _ in range(warmup): + launch() + torch.accelerator.synchronize() + samples: list[float] = [] + for _ in range(trials): + start = torch.Event(enable_timing=True) + end = torch.Event(enable_timing=True) + start.record() + for _ in range(iters): + launch() + end.record() + end.synchronize() + samples.append(float(start.elapsed_time(end) / iters)) + return { + "median_ms": float(statistics.median(samples)), + "mean_ms": float(statistics.mean(samples)), + "min_ms": min(samples), + "max_ms": max(samples), + "samples_ms": samples, + } + + +def _digest(tensor: torch.Tensor) -> str: + raw = tensor.detach().contiguous().cpu().view(torch.uint8).numpy().tobytes() + return hashlib.sha256(raw).hexdigest() + + +def _file_digest(path: Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def _run_case( + args: argparse.Namespace, case: str, device: torch.device +) -> dict[str, Any]: + qweight, scales, prefixes = _load_case( + args.model, case, args.tp_size, args.tp_rank, device + ) + n, k = (int(dim) for dim in qweight.shape) + expected = { + "gate_up": (8704, 5120), + "down": (5120, 4352), + "attention_output": (5120, 1536), + }[case] + if (n, k) != expected: + raise ValueError(f"Unexpected {case} TP-local shape N={n}, K={k}.") + + tm_weight, tm_scales, meta = sm70_ops.fp8_sm70_prepare(qweight, scales, 128) + k_ld, q_ld = (int(value) for value in meta.tolist()) + generator = torch.Generator(device=device).manual_seed(args.seed + n + k) + inputs = torch.randn( + (args.m, k), generator=generator, device=device, dtype=torch.float16 + ).mul_(0.1) + output = torch.empty((args.m, n), device=device, dtype=torch.float16) + if args.gated_silu and case == "gate_up": + output = torch.empty((args.m, n // 2), device=device, dtype=torch.float16) + resident_bytes = 0 + if args.backend == "resident-turbomind-f16": + dense_weight = torch.empty((k, n), device=device, dtype=torch.float16) + sm70_ops.fp8_sm70_dequantize_out(dense_weight, tm_weight, tm_scales, 128) + source_weight = dense_weight.t().contiguous() + f16_tm_weight, f16_meta = sm70_ops.sm70_f16_prepare(source_weight) + f16_k_ld = int(f16_meta[0].item()) + resident_bytes = f16_tm_weight.numel() * f16_tm_weight.element_size() + + def eager_launch() -> None: + sm70_ops.sm70_f16_gemm_out( + output, + inputs, + f16_tm_weight, + f16_k_ld, + args.gated_silu and case == "gate_up", + ) + + else: + workspace = torch.empty((k, n), device=device, dtype=torch.float16) + resident_bytes = workspace.numel() * workspace.element_size() + + def eager_launch() -> None: + sm70_ops.fp8_gemm_sm70_prefill_dispatch_out( + output, + workspace.data_ptr(), + inputs, + tm_weight, + tm_scales, + 128, + k_ld, + q_ld, + args.gated_silu and case == "gate_up", + 3920, + ) + + cutlass_env = "VLLM_SM70_FP8_PREFILL_CUTLASS" + selected_cutlass = os.environ.get(cutlass_env) + os.environ[cutlass_env] = "0" + eager_launch() + torch.accelerator.synchronize(device) + reference = output.clone() + if selected_cutlass is None: + os.environ.pop(cutlass_env) + else: + os.environ[cutlass_env] = selected_cutlass + + launch = eager_launch + graph: torch.cuda.CUDAGraph | None = None + if args.cuda_graph: + for _ in range(3): + eager_launch() + torch.accelerator.synchronize(device) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + eager_launch() + launch = graph.replay + + timing = _event_trials(launch, args.warmup, args.iters, args.trials) + launch() + torch.accelerator.synchronize(device) + output_float = output.float() + reference_float = reference.float() + return { + "case": case, + "prefixes": prefixes, + "m": args.m, + "n": n, + "k": k, + "timing": timing, + "useful_tflops": 2 * args.m * n * k / (timing["median_ms"] * 1e-3) / 1e12, + "output_sha256": _digest(output), + "reference_sha256": _digest(reference), + "output_bitwise_equal": bool(torch.equal(output, reference)), + "output_max_abs_diff": float( + torch.max(torch.abs(output_float - reference_float)).item() + ), + "output_finite": bool(torch.isfinite(output).all().item()), + "resident_bytes": resident_bytes, + } + + +def main() -> int: + args = _parse_args() + device = torch.device(args.device) + if torch.cuda.get_device_capability(device) != (7, 0): + raise RuntimeError("This benchmark requires an SM70/V100 GPU.") + torch.accelerator.set_device_index(device.index or 0) + if args.blas_library != "default": + torch.backends.cuda.preferred_blas_library(args.blas_library) + + results = [_run_case(args, case, device) for case in args.cases] + extension_path = Path(vllm._C.__file__).resolve() + payload = { + "environment": { + "model": str(args.model), + "device": torch.cuda.get_device_name(device), + "capability": list(torch.cuda.get_device_capability(device)), + "torch": torch.__version__, + "torch_cuda": torch.version.cuda, + "blas_library_requested": args.blas_library, + "blas_library_effective": str(torch.backends.cuda.preferred_blas_library()), + "gated_silu": args.gated_silu, + "backend": args.backend, + "cuda_graph": args.cuda_graph, + "cutlass": os.getenv("VLLM_SM70_FP8_PREFILL_CUTLASS", ""), + "vllm_c_extension": str(extension_path), + "vllm_c_extension_sha256": _file_digest(extension_path), + "tp_size": args.tp_size, + "tp_rank": args.tp_rank, + "seed": args.seed, + }, + "results": results, + } + args.json_out.parent.mkdir(parents=True, exist_ok=True) + args.json_out.write_text(json.dumps(payload, indent=2) + "\n") + print(json.dumps(payload, indent=2)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/benchmark_sm70_fp8_prefill_tactics.py b/benchmarks/benchmark_sm70_fp8_prefill_tactics.py new file mode 100644 index 0000000000..5015fba613 --- /dev/null +++ b/benchmarks/benchmark_sm70_fp8_prefill_tactics.py @@ -0,0 +1,486 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Benchmark TurboMind tactics on Qwen3.8 FP8 prefill projections. + +This diagnostic loads the real TP-local fused QKVZ and QKV weights used by +Qwen3.8-27B-FP8. Run each dispatch policy in a fresh process: TurboMind caches +the first exact GEMM descriptor, so switching from ``default`` to ``measure`` +inside one process would not actually tune the already-cached shape. +""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib.util +import json +import math +import os +import statistics +import sys +from collections.abc import Callable +from pathlib import Path +from typing import Any, Literal + +import torch +from safetensors import safe_open + +import vllm + + +def _load_vllm_c_extension() -> None: + if "vllm._C" in sys.modules: + vllm._C = sys.modules["vllm._C"] + return + candidate = os.getenv("VLLM_BENCH_C_EXTENSION") + if candidate is None: + importlib.import_module("vllm._C") + return + spec = importlib.util.spec_from_file_location("vllm._C", candidate) + if spec is None or spec.loader is None: + raise RuntimeError(f"Cannot load candidate extension: {candidate}") + module = importlib.util.module_from_spec(spec) + sys.modules["vllm._C"] = module + vllm._C = module + spec.loader.exec_module(module) + + +_load_vllm_c_extension() + +from vllm import _sm70_ops as sm70_ops # noqa: E402 + +Case = Literal["gdn_qkvz", "full_qkv"] +Policy = Literal[ + "default", + "measure", + "resident-f16", + "resident-f16-measure", + "exact-dispatch", + "exact-mm", + "resident-mm", +] + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument( + "--model", + type=Path, + required=True, + ) + parser.add_argument("--json-out", type=Path, required=True) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--tp-size", type=int, default=4) + parser.add_argument("--tp-rank", type=int, default=0) + parser.add_argument( + "--cases", + choices=("gdn_qkvz", "full_qkv"), + nargs="+", + default=["gdn_qkvz", "full_qkv"], + ) + parser.add_argument("--m", type=int, nargs="+", default=[8000]) + parser.add_argument( + "--policy", + choices=( + "default", + "measure", + "resident-f16", + "resident-f16-measure", + "exact-dispatch", + "exact-mm", + "resident-mm", + ), + required=True, + ) + parser.add_argument( + "--safe-fast", + choices=("off", "preserve-splits", "preserve-split-count"), + default="off", + help="Optionally constrain measured tactics to the default split contract.", + ) + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--iters", type=int, default=50) + parser.add_argument("--trials", type=int, default=7) + parser.add_argument("--seed", type=int, default=20260823) + parser.add_argument("--exact-8k-prescaled-prefill", action="store_true") + parser.add_argument("--cuda-graph", action="store_true") + return parser.parse_args() + + +def _load_keys(model: Path, keys: list[str]) -> dict[str, torch.Tensor]: + index = json.loads((model / "model.safetensors.index.json").read_text()) + weight_map: dict[str, str] = index["weight_map"] + by_file: dict[str, list[str]] = {} + for key in keys: + by_file.setdefault(weight_map[key], []).append(key) + + result: dict[str, torch.Tensor] = {} + for filename, file_keys in by_file.items(): + with safe_open(model / filename, framework="pt", device="cpu") as handle: + for key in file_keys: + result[key] = handle.get_tensor(key) + return result + + +def _column_shard( + weight: torch.Tensor, + scales: torch.Tensor, + tp_size: int, + tp_rank: int, +) -> tuple[torch.Tensor, torch.Tensor]: + n = int(weight.shape[0]) + if n % tp_size: + raise ValueError(f"N={n} is not divisible by TP={tp_size}.") + begin = tp_rank * (n // tp_size) + end = begin + n // tp_size + scale_begin = begin // 128 + scale_end = math.ceil(end / 128) + return ( + weight.view(torch.uint8)[begin:end].contiguous(), + scales[scale_begin:scale_end].contiguous(), + ) + + +def _load_case( + model: Path, + case: Case, + tp_size: int, + tp_rank: int, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor, list[str]]: + root = "model.language_model.layers" + if case == "gdn_qkvz": + prefixes = [ + f"{root}.1.linear_attn.in_proj_qkv", + f"{root}.1.linear_attn.in_proj_z", + ] + elif case == "full_qkv": + prefixes = [ + f"{root}.3.self_attn.q_proj", + f"{root}.3.self_attn.k_proj", + f"{root}.3.self_attn.v_proj", + ] + else: + raise ValueError(f"Unknown case: {case}") + + keys = [ + key + for prefix in prefixes + for key in (f"{prefix}.weight", f"{prefix}.weight_scale_inv") + ] + loaded = _load_keys(model, keys) + raw_parts: list[torch.Tensor] = [] + scale_parts: list[torch.Tensor] = [] + fp8_dtype: torch.dtype | None = None + for prefix in prefixes: + weight = loaded[f"{prefix}.weight"] + scales = loaded[f"{prefix}.weight_scale_inv"].float() + if fp8_dtype is None: + fp8_dtype = weight.dtype + elif fp8_dtype != weight.dtype: + raise ValueError("Fused projections must use one FP8 dtype.") + raw, scale = _column_shard(weight, scales, tp_size, tp_rank) + raw_parts.append(raw) + scale_parts.append(scale) + + assert fp8_dtype is not None + raw = torch.cat(raw_parts, dim=0) + scales = torch.cat(scale_parts, dim=0) + qweight = raw.to(device).view(fp8_dtype).contiguous() + return qweight, scales.to(device).contiguous(), prefixes + + +def _configure_policy(policy: Policy, safe_fast: str, max_m: int) -> None: + os.environ["VLLM_SM70_FP8_DENSE_TUNE_MAX_M"] = str(max_m) + os.environ["VLLM_SM70_FP8_TUNE_SMALL_SHAPES"] = "1" if policy == "measure" else "0" + os.environ["VLLM_SM70_F16_DENSE_TUNE_MAX_M"] = str(max_m) + os.environ["VLLM_SM70_AWQ_TUNE_SMALL_SHAPES"] = ( + "1" if policy == "resident-f16-measure" else "0" + ) + os.environ["VLLM_SM70_FP8_SAFE_FAST_SELECTOR"] = "0" if safe_fast == "off" else "1" + os.environ["VLLM_SM70_FP8_PRESERVE_DEFAULT_SPLITS"] = ( + "1" if safe_fast == "preserve-splits" else "0" + ) + os.environ["VLLM_SM70_FP8_PRESERVE_DEFAULT_SPLITS_ONLY"] = ( + "1" if safe_fast == "preserve-split-count" else "0" + ) + + +def _event_trials( + launch: Callable[[], None], warmup: int, iters: int, trials: int +) -> dict[str, float]: + for _ in range(warmup): + launch() + torch.accelerator.synchronize() + + samples_ms: list[float] = [] + for _ in range(trials): + start = torch.Event(enable_timing=True) + end = torch.Event(enable_timing=True) + start.record() + for _ in range(iters): + launch() + end.record() + end.synchronize() + samples_ms.append(float(start.elapsed_time(end) / iters)) + ordered = sorted(samples_ms) + return { + "median_ms": float(statistics.median(samples_ms)), + "min_ms": ordered[0], + "max_ms": ordered[-1], + } + + +def _digest(tensor: torch.Tensor) -> str: + raw = tensor.detach().contiguous().cpu().view(torch.uint8).numpy().tobytes() + return hashlib.sha256(raw).hexdigest() + + +def _file_digest(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1 << 20), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _run_case( + args: argparse.Namespace, case: Case, m: int, device: torch.device +) -> dict[str, Any]: + qweight, scales, prefixes = _load_case( + args.model, case, args.tp_size, args.tp_rank, device + ) + n, k = (int(dim) for dim in qweight.shape) + if ( + k != 5120 + or (case == "gdn_qkvz" and n != 4096) + or (case == "full_qkv" and n != 3584) + ): + raise ValueError(f"Unexpected {case} TP-local shape N={n}, K={k}.") + + tm_weight, tm_scales, meta = sm70_ops.fp8_sm70_prepare(qweight, scales, 128) + if args.exact_8k_prescaled_prefill: + tm_scales.mul_(256) + k_ld, q_ld = (int(value) for value in meta.tolist()) + generator = torch.Generator(device=device).manual_seed(args.seed + m) + inputs = torch.randn( + (m, k), generator=generator, device=device, dtype=torch.float16 + ).mul_(0.1) + output = torch.empty((m, n), device=device, dtype=torch.float16) + resident_bytes = 0 + dense_workspace_bytes = 0 + reference_output = None + if args.policy in ( + "resident-f16", + "resident-f16-measure", + "exact-dispatch", + "exact-mm", + "resident-mm", + ): + reference_output = torch.empty_like(output) + if args.exact_8k_prescaled_prefill: + sm70_ops.fp8_gemm_sm70_prefill_prescaled_out( + reference_output, + inputs, + tm_weight, + tm_scales, + 128, + k_ld, + q_ld, + ) + else: + sm70_ops.fp8_gemm_sm70_out( + reference_output, + inputs, + tm_weight, + tm_scales, + 128, + k_ld, + q_ld, + False, + ) + dense_weight = torch.empty((k, n), device=device, dtype=torch.float16) + dense_workspace_bytes = int(dense_weight.numel() * dense_weight.element_size()) + if args.policy == "exact-dispatch": + + def launch( + output: torch.Tensor = output, + inputs: torch.Tensor = inputs, + dense_weight: torch.Tensor = dense_weight, + tm_weight: torch.Tensor = tm_weight, + tm_scales: torch.Tensor = tm_scales, + ) -> None: + sm70_ops.fp8_gemm_sm70_prefill_dispatch_out( + output, + dense_weight.data_ptr(), + inputs, + tm_weight, + tm_scales, + 128, + k_ld, + q_ld, + False, + m, + ) + + elif args.policy == "exact-mm": + + def launch( + output: torch.Tensor = output, + inputs: torch.Tensor = inputs, + dense_weight: torch.Tensor = dense_weight, + tm_weight: torch.Tensor = tm_weight, + tm_scales: torch.Tensor = tm_scales, + ) -> None: + sm70_ops.fp8_sm70_dequantize_out( + dense_weight, tm_weight, tm_scales, 128 + ) + torch.mm(inputs, dense_weight, out=output) + + elif args.policy == "resident-mm": + sm70_ops.fp8_sm70_dequantize_out(dense_weight, tm_weight, tm_scales, 128) + resident_bytes = dense_workspace_bytes + + def launch( + output: torch.Tensor = output, + inputs: torch.Tensor = inputs, + dense_weight: torch.Tensor = dense_weight, + ) -> None: + torch.mm(inputs, dense_weight, out=output) + + else: + sm70_ops.fp8_sm70_dequantize_out(dense_weight, tm_weight, tm_scales, 128) + resident_bytes = dense_workspace_bytes + source_f16_weight = dense_weight.t().contiguous() + f16_tm_weight, f16_meta = sm70_ops.sm70_f16_prepare(source_f16_weight) + f16_k_ld = int(f16_meta[0].item()) + + def launch( + output: torch.Tensor = output, + inputs: torch.Tensor = inputs, + f16_tm_weight: torch.Tensor = f16_tm_weight, + ) -> None: + sm70_ops.sm70_f16_gemm_out( + output, inputs, f16_tm_weight, f16_k_ld, False + ) + + else: + if args.exact_8k_prescaled_prefill: + + def launch( + output: torch.Tensor = output, + inputs: torch.Tensor = inputs, + tm_weight: torch.Tensor = tm_weight, + tm_scales: torch.Tensor = tm_scales, + ) -> None: + sm70_ops.fp8_gemm_sm70_prefill_prescaled_out( + output, inputs, tm_weight, tm_scales, 128, k_ld, q_ld + ) + + else: + + def launch( + output: torch.Tensor = output, + inputs: torch.Tensor = inputs, + tm_weight: torch.Tensor = tm_weight, + tm_scales: torch.Tensor = tm_scales, + ) -> None: + sm70_ops.fp8_gemm_sm70_out( + output, inputs, tm_weight, tm_scales, 128, k_ld, q_ld, False + ) + + if args.cuda_graph: + launch() + torch.accelerator.synchronize(device) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + launch() + launch = graph.replay + + timing = _event_trials(launch, args.warmup, args.iters, args.trials) + launch() + torch.accelerator.synchronize(device) + seconds = timing["median_ms"] * 1e-3 + result = { + "case": case, + "prefixes": prefixes, + "m": m, + "n": n, + "k": k, + "tm_meta": {"k_ld": k_ld, "q_ld": q_ld}, + "timing": timing, + "useful_tflops": float(2 * m * n * k / seconds / 1e12), + "dense_workspace_bytes": dense_workspace_bytes, + "resident_f16_bytes": resident_bytes, + "output_sha256": _digest(output), + "output_finite": bool(torch.isfinite(output).all().item()), + } + if reference_output is not None: + difference = (output.float() - reference_output.float()).abs() + result["reference"] = { + "output_sha256": _digest(reference_output), + "equal": bool(torch.equal(output, reference_output)), + "max_abs_diff": float(difference.max().item()), + "mean_abs_diff": float(difference.mean().item()), + "equal_elements": int((output == reference_output).sum().item()), + "elements": output.numel(), + } + del qweight, scales, tm_weight, tm_scales, inputs, output + torch.accelerator.empty_cache() + return result + + +def main() -> int: + args = _parse_args() + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required.") + device = torch.device(args.device) + if torch.cuda.get_device_capability(device) != (7, 0): + raise RuntimeError("This benchmark requires an SM70/V100 GPU.") + if args.tp_size <= 0 or not 0 <= args.tp_rank < args.tp_size: + raise ValueError("Invalid TP size/rank.") + if any(m < 1 for m in args.m): + raise ValueError("Every M must be positive.") + if args.warmup < 1 or args.iters < 1 or args.trials < 1: + raise ValueError("Warmup, iterations, and trials must be positive.") + + torch.accelerator.set_device_index(device.index or 0) + _configure_policy(args.policy, args.safe_fast, max(args.m)) + results = [_run_case(args, case, m, device) for case in args.cases for m in args.m] + extension_path = Path(vllm._C.__file__) + payload = { + "environment": { + "model": str(args.model), + "vllm_c_extension": str(extension_path), + "vllm_c_extension_sha256": _file_digest(extension_path), + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "device": torch.cuda.get_device_name(device), + "capability": list(torch.cuda.get_device_capability(device)), + "torch": torch.__version__, + "torch_cuda": torch.version.cuda, + "tp_size": args.tp_size, + "tp_rank": args.tp_rank, + "policy": args.policy, + "safe_fast": args.safe_fast, + "warmup": args.warmup, + "iters": args.iters, + "trials": args.trials, + "seed": args.seed, + "exact_8k_prescaled_prefill": args.exact_8k_prescaled_prefill, + "cuda_graph": args.cuda_graph, + "tm_gemm_tune": os.environ.get("TM_GEMM_TUNE"), + "fast_targets": os.environ.get("VLLM_SM70_AWQ_TP2_FAST_TARGETS"), + "fp8_prefill_fast_selector": os.environ.get( + "VLLM_SM70_FP8_PREFILL_FAST_SELECTOR" + ), + }, + "results": results, + } + args.json_out.parent.mkdir(parents=True, exist_ok=True) + args.json_out.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8") + print(json.dumps(payload, indent=2)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/benchmark_sm70_fp8_qpn8.py b/benchmarks/benchmark_sm70_fp8_qpn8.py new file mode 100644 index 0000000000..29584b4082 --- /dev/null +++ b/benchmarks/benchmark_sm70_fp8_qpn8.py @@ -0,0 +1,729 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Race experimental QPN8 against TurboMind on real TP-local FP8 weights. + +The benchmark is deliberately operator-only: it does not change model +dispatch. It covers the exact Qwen3.8-27B-FP8 TP4 decode shapes at M=1, 2, 4, +and 8, and measures eager launches separately from CUDA Graph replay. + +The production gate/up GEMM has a fused SiLU epilogue. ``gate_up_raw`` is +therefore opt-in and diagnostic only; its raw-GEMM timing is not an end-to-end +replacement claim. +""" + +from __future__ import annotations + +import argparse +import gc +import json +import math +import statistics +from collections.abc import Callable +from pathlib import Path +from typing import Any + +import torch +import vllm._C # noqa: F401 +from safetensors import safe_open + +from vllm import _sm70_ops as sm70_ops + +_KORDER8 = [0, 2, 4, 6, 1, 3, 5, 7, 8, 10, 12, 14, 9, 11, 13, 15] +_DEFAULT_CASES = ("down", "gdn_in", "output", "full_qkv") +_ALL_CASES = (*_DEFAULT_CASES, "gate_up_raw", "gate_up_fused") + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument( + "--model", + type=Path, + default=Path("/home/ymzx/models/Qwen3.8-27B-FP8"), + ) + parser.add_argument("--json-out", type=Path, required=True) + parser.add_argument( + "--qpn8-library", + type=Path, + help="Optional standalone library built from fp8_qpn8_sm70.cu.", + ) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--tp-size", type=int, default=4) + parser.add_argument("--tp-rank", type=int, default=0) + parser.add_argument("--m", type=int, nargs="+", default=[1, 2, 4, 8]) + parser.add_argument( + "--cases", choices=_ALL_CASES, nargs="+", default=list(_DEFAULT_CASES) + ) + parser.add_argument("--split-k", type=int, nargs="+", default=[4, 8, 16, 32]) + parser.add_argument("--nacc", type=int, nargs="+", default=[1, 2]) + parser.add_argument( + "--decoder", choices=("scalar", "fast"), nargs="+", default=["scalar", "fast"] + ) + parser.add_argument( + "--prefetch", choices=("off", "on"), nargs="+", default=["off", "on"] + ) + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--iters", type=int, default=100) + parser.add_argument("--trials", type=int, default=5) + parser.add_argument( + "--cache-state", + choices=("warm", "cold"), + nargs="+", + default=["warm", "cold"], + ) + parser.add_argument("--cache-scrub-mib", type=int, default=64) + parser.add_argument("--seed", type=int, default=1234) + parser.add_argument("--relative-l2-limit", type=float, default=5e-3) + parser.add_argument("--cosine-limit", type=float, default=0.999) + return parser.parse_args() + + +def _load_keys(model: Path, keys: list[str]) -> dict[str, torch.Tensor]: + index = json.loads((model / "model.safetensors.index.json").read_text()) + weight_map: dict[str, str] = index["weight_map"] + by_file: dict[str, list[str]] = {} + for key in keys: + by_file.setdefault(weight_map[key], []).append(key) + + result: dict[str, torch.Tensor] = {} + for filename, file_keys in by_file.items(): + with safe_open(model / filename, framework="pt", device="cpu") as handle: + for key in file_keys: + result[key] = handle.get_tensor(key) + return result + + +def _weight_keys(prefix: str) -> tuple[str, str]: + return f"{prefix}.weight", f"{prefix}.weight_scale_inv" + + +def _column_shard_raw( + weight: torch.Tensor, + scales: torch.Tensor, + tp_size: int, + tp_rank: int, +) -> tuple[torch.Tensor, torch.Tensor]: + n = weight.shape[0] + if n % tp_size: + raise ValueError(f"N={n} is not divisible by TP={tp_size}") + begin = tp_rank * (n // tp_size) + end = begin + n // tp_size + scale_begin = begin // 128 + scale_end = math.ceil(end / 128) + raw = weight.view(torch.uint8)[begin:end].contiguous() + return raw, scales[scale_begin:scale_end].contiguous() + + +def _row_shard_raw( + weight: torch.Tensor, + scales: torch.Tensor, + tp_size: int, + tp_rank: int, +) -> tuple[torch.Tensor, torch.Tensor]: + k = weight.shape[1] + if k % tp_size: + raise ValueError(f"K={k} is not divisible by TP={tp_size}") + begin = tp_rank * (k // tp_size) + end = begin + k // tp_size + scale_begin = begin // 128 + scale_end = math.ceil(end / 128) + raw = weight.view(torch.uint8)[:, begin:end].contiguous() + return raw, scales[:, scale_begin:scale_end].contiguous() + + +def _load_case( + model: Path, + case: str, + tp_size: int, + tp_rank: int, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor, str]: + root = "model.language_model.layers" + if case == "down": + prefixes = [f"{root}.1.mlp.down_proj"] + shard = "row" + note = "production down_proj TP row shard" + elif case == "gdn_in": + prefixes = [ + f"{root}.1.linear_attn.in_proj_qkv", + f"{root}.1.linear_attn.in_proj_z", + ] + shard = "column" + note = "production GDN qkv+z TP column shards concatenated" + elif case == "output": + prefixes = [f"{root}.1.linear_attn.out_proj"] + shard = "row" + note = "production GDN/full-attention output TP row-shard shape" + elif case == "full_qkv": + prefixes = [ + f"{root}.3.self_attn.q_proj", + f"{root}.3.self_attn.k_proj", + f"{root}.3.self_attn.v_proj", + ] + shard = "column" + note = "production full-attention q+k+v TP column shards concatenated" + elif case in ("gate_up_raw", "gate_up_fused"): + prefixes = [ + f"{root}.1.mlp.gate_proj", + f"{root}.1.mlp.up_proj", + ] + shard = "column" + note = ( + "production gate+up GEMM with fused SiLU-and-multiply" + if case == "gate_up_fused" + else "diagnostic raw gate+up GEMM; production has fused SiLU epilogue" + ) + else: + raise ValueError(f"unknown case: {case}") + + keys = [key for prefix in prefixes for key in _weight_keys(prefix)] + loaded = _load_keys(model, keys) + raw_parts: list[torch.Tensor] = [] + scale_parts: list[torch.Tensor] = [] + fp8_dtype: torch.dtype | None = None + for prefix in prefixes: + weight_key, scale_key = _weight_keys(prefix) + weight = loaded[weight_key] + scales = loaded[scale_key].float() + if fp8_dtype is None: + fp8_dtype = weight.dtype + elif weight.dtype != fp8_dtype: + raise ValueError("concatenated weights have different dtypes") + if shard == "column": + raw, scale = _column_shard_raw(weight, scales, tp_size, tp_rank) + else: + raw, scale = _row_shard_raw(weight, scales, tp_size, tp_rank) + raw_parts.append(raw) + scale_parts.append(scale) + + raw = torch.cat(raw_parts, dim=0) if len(raw_parts) > 1 else raw_parts[0] + scales = torch.cat(scale_parts, dim=0) if len(scale_parts) > 1 else scale_parts[0] + assert fp8_dtype is not None + qweight = raw.to(device, non_blocking=False).view(fp8_dtype).contiguous() + return qweight, scales.to(device).contiguous(), note + + +def _qpn8_prepack(raw_weight: torch.Tensor) -> torch.Tensor: + """[N,K] uint8 -> [tile N/32][group K/16][lane 32][16B].""" + n, k = raw_weight.shape + if n % 32 or k % 16: + raise ValueError(f"QPN8 packing requires N%32=K%16=0, got N={n} K={k}") + device = raw_weight.device + tiles, groups = n // 32, k // 16 + lane = torch.arange(32, device=device) + col = ((lane >> 2) & 3) * 8 + (lane & 3) + ((lane & 16) > 0).long() * 4 + korder = torch.tensor(_KORDER8, device=device) + group = torch.arange(groups, device=device) + kidx = group[:, None] * 16 + korder[None, :] + packed = torch.empty((tiles, groups, 32, 16), dtype=torch.uint8, device=device) + # Bound the temporary expanded index tensors for the 44-MiB projections. + tiles_per_chunk = max(1, 36864 // max(groups, 1)) + for tile_begin in range(0, tiles, tiles_per_chunk): + tile_end = min(tile_begin + tiles_per_chunk, tiles) + tile_count = tile_end - tile_begin + ncol = ( + torch.arange(tile_begin, tile_end, device=device)[:, None] * 32 + + col[None, :] + ) + packed[tile_begin:tile_end] = raw_weight[ + ncol.view(tile_count, 1, 32, 1).expand(tile_count, groups, 32, 16), + kidx.view(1, groups, 1, 16).expand(tile_count, groups, 32, 16), + ] + return packed.view(-1).contiguous() + + +def _qpn8_group_scales(scales: torch.Tensor, n: int, k: int) -> torch.Tensor: + expected = (math.ceil(n / 128), math.ceil(k / 128)) + if tuple(scales.shape) != expected or n % 128 or k % 128: + raise ValueError( + f"QPN8 scales need exact 128x128 blocks, got {tuple(scales.shape)} " + f"for N={n}, K={k}" + ) + return ( + scales.t() + .reshape(k // 128, n // 128) + .repeat_interleave(4, dim=1) + .mul(256.0) + .half() + .contiguous() + ) + + +def _dequantized_weight(qweight: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + n, k = qweight.shape + expanded = scales.repeat_interleave(128, 0).repeat_interleave(128, 1) + return qweight.float().mul(expanded[:n, :k]) + + +def _error_stats(actual: torch.Tensor, reference: torch.Tensor) -> dict[str, Any]: + actual_f32 = actual.float() + reference_f32 = reference.float() + delta = actual_f32 - reference_f32 + ref_norm = torch.linalg.vector_norm(reference_f32).clamp_min(1e-12) + relative_l2 = torch.linalg.vector_norm(delta) / ref_norm + cosine = torch.nn.functional.cosine_similarity( + actual_f32.flatten(), reference_f32.flatten(), dim=0 + ) + return { + "finite": bool(torch.isfinite(actual).all().item()), + "max_abs": float(delta.abs().max().item()), + "mean_abs": float(delta.abs().mean().item()), + "relative_l2": float(relative_l2.item()), + "cosine": float(cosine.item()), + } + + +def _event_trials( + launch: Callable[[], None], + warmup: int, + iters: int, + trials: int, + cache_scrub: torch.Tensor | None, +) -> dict[str, float]: + for _ in range(warmup): + if cache_scrub is not None: + cache_scrub.add_(1) + launch() + torch.accelerator.synchronize() + samples_us: list[float] = [] + for _ in range(trials): + if cache_scrub is None: + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + launch() + end.record() + end.synchronize() + samples_us.append(float(start.elapsed_time(end) * 1000.0 / iters)) + else: + # Scrub more than V100's L2 immediately before each call, while + # the event pair measures only the tested operator. + starts = [torch.cuda.Event(enable_timing=True) for _ in range(iters)] + ends = [torch.cuda.Event(enable_timing=True) for _ in range(iters)] + for start, end in zip(starts, ends): + cache_scrub.add_(1) + start.record() + launch() + end.record() + ends[-1].synchronize() + samples_us.append( + float( + statistics.mean( + start.elapsed_time(end) * 1000.0 + for start, end in zip(starts, ends) + ) + ) + ) + ordered = sorted(samples_us) + return { + "median_us": float(statistics.median(samples_us)), + "min_us": ordered[0], + "p10_us": ordered[max(0, math.ceil(0.1 * len(ordered)) - 1)], + "p90_us": ordered[min(len(ordered) - 1, math.ceil(0.9 * len(ordered)) - 1)], + "max_us": ordered[-1], + } + + +def _benchmark_eager( + launch: Callable[[], None], + warmup: int, + iters: int, + trials: int, + cache_scrub: torch.Tensor | None, +) -> dict[str, Any]: + timing = _event_trials(launch, warmup, iters, trials, cache_scrub) + launch() + torch.accelerator.synchronize() + return {"timing": timing, "replay_max_abs": None} + + +def _benchmark_graph( + launch: Callable[[], None], + output: torch.Tensor, + warmup: int, + iters: int, + trials: int, + cache_scrub: torch.Tensor | None, +) -> dict[str, Any]: + capture_stream = torch.cuda.Stream(device=output.device) + capture_stream.wait_stream(torch.cuda.current_stream(output.device)) + with torch.cuda.stream(capture_stream): + for _ in range(max(warmup, 3)): + launch() + capture_stream.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=capture_stream): + launch() + torch.accelerator.synchronize(output.device) + + timing = _event_trials(graph.replay, warmup, iters, trials, cache_scrub) + graph.replay() + torch.accelerator.synchronize(output.device) + first = output.clone() + graph.replay() + torch.accelerator.synchronize(output.device) + replay_max_abs = float((output - first).abs().max().item()) + del first, graph, capture_stream + gc.collect() + return {"timing": timing, "replay_max_abs": replay_max_abs} + + +def _derived_metrics( + timing: dict[str, float], m: int, n: int, k: int, bytes_per_call: int +) -> dict[str, float]: + seconds = timing["median_us"] * 1e-6 + return { + "effective_gbps": float(bytes_per_call / seconds / 1e9), + "useful_tflops": float(2 * m * n * k / seconds / 1e12), + } + + +def _quality_pass( + stats: dict[str, Any], relative_l2_limit: float, cosine_limit: float +) -> bool: + return bool( + stats["finite"] + and stats["relative_l2"] <= relative_l2_limit + and stats["cosine"] >= cosine_limit + ) + + +def _run_case(args: argparse.Namespace, case: str) -> dict[str, Any]: + device = torch.device(args.device) + qweight, scales, note = _load_case( + args.model, case, args.tp_size, args.tp_rank, device + ) + n, k = (int(dim) for dim in qweight.shape) + gated_silu = case == "gate_up_fused" + output_n = n // 2 if gated_silu else n + raw = qweight.view(torch.uint8) + reference_codes = _qpn8_prepack(raw) + reference_group_scales = _qpn8_group_scales(scales, n, k) + if hasattr(torch.ops._C, "fp8_qpn8_prepare_sm70"): + codes, group_scales = sm70_ops.fp8_qpn8_prepare_sm70(qweight, scales) + prepare_matches_reference = bool( + torch.equal(codes.view(-1), reference_codes) + and torch.equal(group_scales, reference_group_scales) + ) + if not prepare_matches_reference: + raise RuntimeError("source QPN8 prepare disagrees with Python reference") + else: + codes = reference_codes + group_scales = reference_group_scales + prepare_matches_reference = None + tm_weight, tm_scales, meta = sm70_ops.fp8_sm70_prepare( + qweight, scales, 128, gated_silu + ) + k_ld, q_ld = (int(value) for value in meta.tolist()) + reference_weight = _dequantized_weight(qweight, scales) + cache_scrub = torch.empty( + args.cache_scrub_mib * 1024 * 1024, + dtype=torch.uint8, + device=device, + ) + torch.accelerator.synchronize(device) + + case_result: dict[str, Any] = { + "case": case, + "note": note, + "shape": {"n": n, "k": k, "output_n": output_n}, + "qweight_dtype": str(qweight.dtype), + "scale_dtype": str(scales.dtype), + "tm_meta": {"k_ld": k_ld, "q_ld": q_ld}, + "packed_bytes": int(codes.numel()), + "group_scale_shape": list(group_scales.shape), + "source_prepare_matches_reference": prepare_matches_reference, + "rows": [], + } + + valid_configs = [ + (split_k, nacc, decoder == "fast", prefetch == "on") + for split_k in args.split_k + for nacc in args.nacc + for decoder in args.decoder + for prefetch in args.prefetch + if (k // 16) % split_k == 0 + and not (prefetch == "on" and decoder != "fast") + and not (gated_silu and split_k == 32) + ] + for m in args.m: + torch.manual_seed(args.seed + m) + inputs = torch.randn((m, k), device=device, dtype=torch.float16).mul_(0.1) + reference_raw = inputs.float().matmul(reference_weight.t()) + if gated_silu: + gate, up = reference_raw.chunk(2, dim=1) + reference = (torch.nn.functional.silu(gate) * up).half() + else: + reference = reference_raw.half() + + tm_out = torch.empty((m, output_n), device=device, dtype=torch.float16) + + def launch_tm( + tm_out: torch.Tensor = tm_out, + inputs: torch.Tensor = inputs, + tm_weight: torch.Tensor = tm_weight, + tm_scales: torch.Tensor = tm_scales, + ) -> None: + sm70_ops.fp8_gemm_sm70_out( + tm_out, + inputs, + tm_weight, + tm_scales, + 128, + k_ld, + q_ld, + gated_silu, + ) + + launch_tm() + torch.accelerator.synchronize(device) + tm_quality = _error_stats(tm_out, reference) + tm_bytes = n * k + 2 * (k // 128) * n + 2 * m * (output_n + k) + for cache_state in args.cache_state: + scrub = cache_scrub if cache_state == "cold" else None + for mode, benchmark in ( + ("eager", _benchmark_eager), + ("graph", _benchmark_graph), + ): + measured = ( + benchmark( + launch_tm, + args.warmup, + args.iters, + args.trials, + scrub, + ) + if mode == "eager" + else benchmark( + launch_tm, + tm_out, + args.warmup, + args.iters, + args.trials, + scrub, + ) + ) + case_result["rows"].append( + { + "backend": "turbomind_current", + "mode": mode, + "cache_state": cache_state, + "m": m, + "config": None, + "quality": tm_quality, + "quality_pass": _quality_pass( + tm_quality, args.relative_l2_limit, args.cosine_limit + ), + **measured, + **_derived_metrics(measured["timing"], m, n, k, tm_bytes), + } + ) + + for split_k, nacc, fast_decoder, prefetch_codes in valid_configs: + qpn_out = torch.empty((m, output_n), device=device, dtype=torch.float16) + + def launch_qpn( + split_k: int = split_k, + nacc: int = nacc, + fast_decoder: bool = fast_decoder, + prefetch_codes: bool = prefetch_codes, + qpn_out: torch.Tensor = qpn_out, + inputs: torch.Tensor = inputs, + codes: torch.Tensor = codes, + group_scales: torch.Tensor = group_scales, + gated_silu: bool = gated_silu, + ) -> None: + if gated_silu: + sm70_ops.fp8_qpn8_gated_pair_sm70_out( + qpn_out, + inputs, + codes, + group_scales, + split_k, + nacc, + fast_decoder, + prefetch_codes, + ) + else: + sm70_ops.fp8_qpn8_gemm_sm70_out( + qpn_out, + inputs, + codes, + group_scales, + split_k, + nacc, + fast_decoder, + prefetch_codes, + ) + + launch_qpn() + torch.accelerator.synchronize(device) + quality = _error_stats(qpn_out, reference) + quality_pass = _quality_pass( + quality, args.relative_l2_limit, args.cosine_limit + ) + config = { + "split_k": split_k, + "nacc": nacc, + "decoder": "fast" if fast_decoder else "scalar", + "prefetch": prefetch_codes, + } + qpn_bytes = n * k + 2 * (k // 128) * (n // 32) + 2 * m * (n + k) + if gated_silu: + qpn_bytes -= 2 * m * (n - output_n) + for cache_state in args.cache_state: + scrub = cache_scrub if cache_state == "cold" else None + for mode, benchmark in ( + ("eager", _benchmark_eager), + ("graph", _benchmark_graph), + ): + measured = ( + benchmark( + launch_qpn, + args.warmup, + args.iters, + args.trials, + scrub, + ) + if mode == "eager" + else benchmark( + launch_qpn, + qpn_out, + args.warmup, + args.iters, + args.trials, + scrub, + ) + ) + case_result["rows"].append( + { + "backend": "qpn8_experimental", + "mode": mode, + "cache_state": cache_state, + "m": m, + "config": config, + "quality": quality, + "quality_pass": quality_pass, + **measured, + **_derived_metrics(measured["timing"], m, n, k, qpn_bytes), + } + ) + + del qpn_out + + del inputs, reference_raw, reference, tm_out + + del ( + reference_weight, + tm_weight, + tm_scales, + codes, + group_scales, + reference_codes, + reference_group_scales, + qweight, + scales, + cache_scrub, + ) + torch.accelerator.empty_cache() + gc.collect() + return case_result + + +def _summarize(cases: list[dict[str, Any]]) -> list[dict[str, Any]]: + summary: list[dict[str, Any]] = [] + for case in cases: + rows: list[dict[str, Any]] = case["rows"] + for m in sorted({row["m"] for row in rows}): + for cache_state in sorted({row["cache_state"] for row in rows}): + for mode in ("eager", "graph"): + current = next( + row + for row in rows + if row["m"] == m + and row["mode"] == mode + and row["cache_state"] == cache_state + and row["backend"] == "turbomind_current" + ) + candidates = [ + row + for row in rows + if row["m"] == m + and row["mode"] == mode + and row["cache_state"] == cache_state + and row["backend"] == "qpn8_experimental" + and row["quality_pass"] + and (row["replay_max_abs"] in (None, 0.0)) + ] + best = min(candidates, key=lambda row: row["timing"]["median_us"]) + current_us = current["timing"]["median_us"] + best_us = best["timing"]["median_us"] + summary.append( + { + "case": case["case"], + "shape": case["shape"], + "m": m, + "mode": mode, + "cache_state": cache_state, + "current_us": current_us, + "best_qpn8_us": best_us, + "speedup": current_us / best_us, + "best_config": best["config"], + "best_effective_gbps": best["effective_gbps"], + "best_useful_tflops": best["useful_tflops"], + "best_quality": best["quality"], + } + ) + return summary + + +def main() -> int: + args = _parse_args() + if args.qpn8_library is not None: + torch.ops.load_library(str(args.qpn8_library.resolve())) + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + device = torch.device(args.device) + if torch.cuda.get_device_capability(device) != (7, 0): + raise RuntimeError("This benchmark requires an SM70/V100 device") + if not hasattr(torch.ops._C, "fp8_qpn8_gemm_sm70_out"): + raise RuntimeError("Missing _C::fp8_qpn8_gemm_sm70_out; build this source tree") + if args.tp_size <= 0 or not 0 <= args.tp_rank < args.tp_size: + raise ValueError("invalid TP size/rank") + if any(m < 1 or m > 8 for m in args.m): + raise ValueError("QPN8 operator supports M=1..8") + if args.warmup < 1 or args.iters < 1 or args.trials < 1 or args.cache_scrub_mib < 1: + raise ValueError("warmup, iters, trials, and cache scrub size must be positive") + + torch.accelerator.set_device_index(device.index or 0) + cases = [_run_case(args, case) for case in args.cases] + summary = _summarize(cases) + payload = { + "environment": { + "model": str(args.model), + "device": torch.cuda.get_device_name(device), + "capability": list(torch.cuda.get_device_capability(device)), + "torch": torch.__version__, + "torch_cuda": torch.version.cuda, + "tp_size": args.tp_size, + "tp_rank": args.tp_rank, + "m": args.m, + "warmup": args.warmup, + "iters": args.iters, + "trials": args.trials, + "cache_state": args.cache_state, + "cache_scrub_mib": args.cache_scrub_mib, + "seed": args.seed, + }, + "summary": summary, + "cases": cases, + } + args.json_out.parent.mkdir(parents=True, exist_ok=True) + args.json_out.write_text(json.dumps(payload, indent=2), encoding="utf-8") + print(json.dumps(summary, indent=2)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/benchmark_sm70_gdn_exactness.py b/benchmarks/benchmark_sm70_gdn_exactness.py index df3acaf627..165b033254 100644 --- a/benchmarks/benchmark_sm70_gdn_exactness.py +++ b/benchmarks/benchmark_sm70_gdn_exactness.py @@ -165,12 +165,15 @@ def _make_inputs(args: argparse.Namespace) -> dict[str, torch.Tensor]: ) * args.input_scale ).to(dtype) - g = -torch.rand( - (args.batch, args.seqlen, args.v_heads), - device=device, - dtype=torch.float32, - generator=generator, - ) * args.gate_scale + g = ( + -torch.rand( + (args.batch, args.seqlen, args.v_heads), + device=device, + dtype=torch.float32, + generator=generator, + ) + * args.gate_scale + ) beta = torch.sigmoid( torch.randn( (args.batch, args.seqlen, args.v_heads), @@ -211,9 +214,7 @@ def _make_inputs(args: argparse.Namespace) -> dict[str, torch.Tensor]: "initial_state": initial_state.cpu(), } if args.cu_seqlens: - result["cu_seqlens"] = torch.tensor( - [0, args.seqlen], dtype=torch.int32 - ) + result["cu_seqlens"] = torch.tensor([0, args.seqlen], dtype=torch.int32) return result @@ -565,25 +566,17 @@ def _make_spec_decode_cudagraph_inputs( max_query_len = args.num_spec + 1 graph_rows = args.graph_rows or max_query_len if graph_rows < max_query_len: - raise ValueError( - f"--graph-rows must be >= num_spec + 1 ({max_query_len})" - ) + raise ValueError(f"--graph-rows must be >= num_spec + 1 ({max_query_len})") if args.state_slots < max_query_len: - raise ValueError( - f"--state-slots must be >= num_spec + 1 ({max_query_len})" - ) + raise ValueError(f"--state-slots must be >= num_spec + 1 ({max_query_len})") if args.fixed_query_len is not None and not ( 1 <= args.fixed_query_len <= max_query_len ): - raise ValueError( - f"--fixed-query-len must be in [1, {max_query_len}]" - ) + raise ValueError(f"--fixed-query-len must be in [1, {max_query_len}]") if args.fixed_accepted is not None and not ( 1 <= args.fixed_accepted <= max_query_len ): - raise ValueError( - f"--fixed-accepted must be in [1, {max_query_len}]" - ) + raise ValueError(f"--fixed-accepted must be in [1, {max_query_len}]") q_dim = args.k_heads * args.head_k_dim k_dim = args.k_heads * args.head_k_dim @@ -769,7 +762,27 @@ def _load_or_create_spec_decode_cudagraph_inputs( payload = _torch_load(Path(args.input_file)) if "inputs" not in payload: raise ValueError(f"{args.input_file} does not contain saved inputs") - return payload["inputs"] + inputs = payload["inputs"] + available_steps = int(inputs["mixed_qkv_seq"].shape[0]) + if args.steps > available_steps: + raise ValueError( + f"--steps={args.steps} exceeds the {available_steps} saved steps" + ) + if args.steps < available_steps: + inputs = dict(inputs) + for name in ( + "mixed_qkv_seq", + "a_seq", + "b_seq", + "query_start_loc_seq", + "state_indices_seq", + "num_accepted_tokens_seq", + "query_lens", + "z_seq", + ): + if name in inputs: + inputs[name] = inputs[name][: args.steps].contiguous() + return inputs return _make_spec_decode_cudagraph_inputs(args) @@ -852,13 +865,7 @@ def _add_spec_projection_inputs( def _load_or_create_spec_projection_inputs( args: argparse.Namespace, ) -> dict[str, torch.Tensor]: - if args.input_file: - payload = _torch_load(Path(args.input_file)) - if "inputs" not in payload: - raise ValueError(f"{args.input_file} does not contain saved inputs") - inputs = payload["inputs"] - else: - inputs = _make_spec_decode_cudagraph_inputs(args) + inputs = _load_or_create_spec_decode_cudagraph_inputs(args) return _add_spec_projection_inputs(args, inputs) @@ -906,12 +913,14 @@ def _make_sigmoid_gating_inputs(args: argparse.Namespace) -> dict[str, torch.Ten ) * args.gate_scale ) - inputs.update({ - "a": a.cpu(), - "b": b.cpu(), - "A_log": a_log.cpu(), - "dt_bias": dt_bias.cpu(), - }) + inputs.update( + { + "a": a.cpu(), + "b": b.cpu(), + "A_log": a_log.cpu(), + "dt_bias": dt_bias.cpu(), + } + ) if getattr(args, "decode_segments", False): if args.batch != 1: raise ValueError("decode-segments fixture requires --batch 1") @@ -937,9 +946,7 @@ def _make_sigmoid_gating_inputs(args: argparse.Namespace) -> dict[str, torch.Ten dtype=state_dtype, ) inputs["initial_state"] = initial_state.cpu() - inputs["cu_seqlens"] = torch.arange( - 0, args.seqlen + 1, dtype=torch.int32 - ) + inputs["cu_seqlens"] = torch.arange(0, args.seqlen + 1, dtype=torch.int32) inputs["ssm_state_indices"] = torch.arange( 1, args.seqlen + 1, dtype=torch.int32 ) @@ -1105,9 +1112,7 @@ def run_op() -> tuple[torch.Tensor, torch.Tensor | None]: outputs: dict[str, torch.Tensor | None] = { "out": out.detach().cpu(), - "final_state": final_state.detach().cpu() - if final_state is not None - else None, + "final_state": final_state.detach().cpu() if final_state is not None else None, } if args.save_intermediates: pipeline_outputs = _run_chunk_pipeline(cuda_inputs, cu_seqlens, scale) @@ -1161,12 +1166,17 @@ def run_op() -> tuple[torch.Tensor, torch.Tensor | None]: out_path = Path(args.out) out_path.parent.mkdir(parents=True, exist_ok=True) torch.save(payload, out_path) - print(json.dumps({ - "out": str(out_path), - "elapsed_s": elapsed_s, - "timing_summary": payload["timing_summary"], - "output_summaries": payload["output_summaries"], - }, indent=2)) + print( + json.dumps( + { + "out": str(out_path), + "elapsed_s": elapsed_s, + "timing_summary": payload["timing_summary"], + "output_summaries": payload["output_summaries"], + }, + indent=2, + ) + ) return 0 @@ -1269,12 +1279,17 @@ def run_op() -> tuple[torch.Tensor, torch.Tensor]: out_path = Path(args.out) out_path.parent.mkdir(parents=True, exist_ok=True) torch.save(payload, out_path) - print(json.dumps({ - "out": str(out_path), - "elapsed_s": elapsed_s, - "timing_summary": payload["timing_summary"], - "output_summaries": payload["output_summaries"], - }, indent=2)) + print( + json.dumps( + { + "out": str(out_path), + "elapsed_s": elapsed_s, + "timing_summary": payload["timing_summary"], + "output_summaries": payload["output_summaries"], + }, + indent=2, + ) + ) return 0 @@ -1385,12 +1400,17 @@ def run_op() -> tuple[torch.Tensor, torch.Tensor]: out_path = Path(args.out) out_path.parent.mkdir(parents=True, exist_ok=True) torch.save(payload, out_path) - print(json.dumps({ - "out": str(out_path), - "elapsed_s": elapsed_s, - "timing_summary": payload["timing_summary"], - "output_summaries": payload["output_summaries"], - }, indent=2)) + print( + json.dumps( + { + "out": str(out_path), + "elapsed_s": elapsed_s, + "timing_summary": payload["timing_summary"], + "output_summaries": payload["output_summaries"], + }, + indent=2, + ) + ) return 0 @@ -1512,12 +1532,17 @@ def run_op() -> tuple[torch.Tensor, torch.Tensor]: out_path = Path(args.out) out_path.parent.mkdir(parents=True, exist_ok=True) torch.save(payload, out_path) - print(json.dumps({ - "out": str(out_path), - "elapsed_s": elapsed_s, - "timing_summary": payload["timing_summary"], - "output_summaries": payload["output_summaries"], - }, indent=2)) + print( + json.dumps( + { + "out": str(out_path), + "elapsed_s": elapsed_s, + "timing_summary": payload["timing_summary"], + "output_summaries": payload["output_summaries"], + }, + indent=2, + ) + ) return 0 @@ -1851,12 +1876,17 @@ def run_op() -> tuple[torch.Tensor, torch.Tensor]: out_path = Path(args.out) out_path.parent.mkdir(parents=True, exist_ok=True) torch.save(payload, out_path) - print(json.dumps({ - "out": str(out_path), - "elapsed_s": elapsed_s, - "timing_summary": payload["timing_summary"], - "output_summaries": payload["output_summaries"], - }, indent=2)) + print( + json.dumps( + { + "out": str(out_path), + "elapsed_s": elapsed_s, + "timing_summary": payload["timing_summary"], + "output_summaries": payload["output_summaries"], + }, + indent=2, + ) + ) return 0 @@ -1947,9 +1977,7 @@ def run_op() -> tuple[torch.Tensor, torch.Tensor]: "metadata": _metadata(args), "config": { "op": "flashqla_sm70_gdn_decode_full_route", - "route_components": ( - "fused_mixed_qkv+gating+l2norm+global_state_update" - ), + "route_components": ("fused_mixed_qkv+gating+l2norm+global_state_update"), "batch": args.batch, "k_heads": args.k_heads, "v_heads": args.v_heads, @@ -1984,12 +2012,17 @@ def run_op() -> tuple[torch.Tensor, torch.Tensor]: out_path = Path(args.out) out_path.parent.mkdir(parents=True, exist_ok=True) torch.save(payload, out_path) - print(json.dumps({ - "out": str(out_path), - "elapsed_s": elapsed_s, - "timing_summary": payload["timing_summary"], - "output_summaries": payload["output_summaries"], - }, indent=2)) + print( + json.dumps( + { + "out": str(out_path), + "elapsed_s": elapsed_s, + "timing_summary": payload["timing_summary"], + "output_summaries": payload["output_summaries"], + }, + indent=2, + ) + ) return 0 @@ -2116,13 +2149,10 @@ def launch_one( } comparisons = { "out": _compare_tensor(outputs["eager_out"], outputs["graph_out"]), - "final_state": _compare_tensor( - outputs["eager_state"], outputs["graph_state"] - ), + "final_state": _compare_tensor(outputs["eager_state"], outputs["graph_state"]), } strict_ok = all( - item["torch_equal"] and item["max_diff"] == 0.0 - for item in comparisons.values() + item["torch_equal"] and item["max_diff"] == 0.0 for item in comparisons.values() ) saved_inputs = {name: tensor.cpu() for name, tensor in inputs.items()} payload = { @@ -2160,12 +2190,17 @@ def launch_one( out_path = Path(args.out) out_path.parent.mkdir(parents=True, exist_ok=True) torch.save(payload, out_path) - print(json.dumps({ - "out": str(out_path), - "strict_ok": strict_ok, - "elapsed_s": elapsed_s, - "comparisons": comparisons, - }, indent=2)) + print( + json.dumps( + { + "out": str(out_path), + "strict_ok": strict_ok, + "elapsed_s": elapsed_s, + "comparisons": comparisons, + }, + indent=2, + ) + ) return 0 if strict_ok else 1 @@ -2247,13 +2282,10 @@ def launch_one( } comparisons = { "out": _compare_tensor(outputs["eager_out"], outputs["graph_out"]), - "final_state": _compare_tensor( - outputs["eager_state"], outputs["graph_state"] - ), + "final_state": _compare_tensor(outputs["eager_state"], outputs["graph_state"]), } strict_ok = all( - item["torch_equal"] and item["max_diff"] == 0.0 - for item in comparisons.values() + item["torch_equal"] and item["max_diff"] == 0.0 for item in comparisons.values() ) saved_inputs = {name: tensor.cpu() for name, tensor in inputs.items()} payload = { @@ -2288,12 +2320,17 @@ def launch_one( out_path = Path(args.out) out_path.parent.mkdir(parents=True, exist_ok=True) torch.save(payload, out_path) - print(json.dumps({ - "out": str(out_path), - "strict_ok": strict_ok, - "elapsed_s": elapsed_s, - "comparisons": comparisons, - }, indent=2)) + print( + json.dumps( + { + "out": str(out_path), + "strict_ok": strict_ok, + "elapsed_s": elapsed_s, + "comparisons": comparisons, + }, + indent=2, + ) + ) return 0 if strict_ok else 1 @@ -2364,6 +2401,9 @@ def launch_one( a, b, cuda_inputs["dt_bias"], + beta_dtype=( + torch.float32 if args.beta_output_dtype == "float32" else b.dtype + ), ) core_out, _ = fused_recurrent_gated_delta_rule( q=query, @@ -2471,8 +2511,7 @@ def launch_one( query_lens = inputs["query_lens"].to(torch.int64) live_token_mask = ( - torch.arange(max_query_len, dtype=torch.int64)[None, :] - < query_lens[:, None] + torch.arange(max_query_len, dtype=torch.int64)[None, :] < query_lens[:, None] ) outputs: dict[str, torch.Tensor | None] = { "eager_live_out": eager_out.detach().cpu()[live_token_mask], @@ -2529,6 +2568,7 @@ def launch_one( "conv_width": args.conv_width, "dtype": args.dtype, "state_dtype": args.state_dtype, + "beta_output_dtype": args.beta_output_dtype, "seed": args.seed, "scale": scale, "silu": args.silu, @@ -2557,13 +2597,18 @@ def launch_one( out_path = Path(args.out) out_path.parent.mkdir(parents=True, exist_ok=True) torch.save(payload, out_path) - print(json.dumps({ - "out": str(out_path), - "strict_ok": strict_ok, - "elapsed_s": elapsed_s, - "config": payload["config"], - "comparisons": comparisons, - }, indent=2)) + print( + json.dumps( + { + "out": str(out_path), + "strict_ok": strict_ok, + "elapsed_s": elapsed_s, + "config": payload["config"], + "comparisons": comparisons, + }, + indent=2, + ) + ) return 0 if strict_ok else 1 @@ -2650,6 +2695,9 @@ def launch_split( a, b, cuda_inputs["dt_bias"], + beta_dtype=( + torch.float32 if args.beta_output_dtype == "float32" else b.dtype + ), ) core_out, _ = fused_recurrent_gated_delta_rule( q=query, @@ -2704,7 +2752,9 @@ def launch_mixed_qkv( ) out.copy_(core_out.squeeze(0)) - def run_eager_sequence(launch_fn: Any) -> tuple[ + def run_eager_sequence( + launch_fn: Any, + ) -> tuple[ torch.Tensor, torch.Tensor, torch.Tensor, @@ -2804,8 +2854,7 @@ def time_graph_sequence(launch_fn: Any) -> float: query_lens = inputs["query_lens"].to(torch.int64) live_token_mask = ( - torch.arange(max_query_len, dtype=torch.int64)[None, :] - < query_lens[:, None] + torch.arange(max_query_len, dtype=torch.int64)[None, :] < query_lens[:, None] ) outputs: dict[str, torch.Tensor] = { "split_live_out": split_out.detach().cpu()[live_token_mask], @@ -2871,6 +2920,7 @@ def time_graph_sequence(launch_fn: Any) -> float: "conv_width": args.conv_width, "dtype": args.dtype, "state_dtype": args.state_dtype, + "beta_output_dtype": args.beta_output_dtype, "seed": args.seed, "scale": scale, "silu": args.silu, @@ -2898,13 +2948,18 @@ def time_graph_sequence(launch_fn: Any) -> float: out_path = Path(args.out) out_path.parent.mkdir(parents=True, exist_ok=True) torch.save(payload, out_path) - print(json.dumps({ - "out": str(out_path), - "strict_ok": strict_ok, - "timing": timing, - "config": payload["config"], - "comparisons": comparisons, - }, indent=2)) + print( + json.dumps( + { + "out": str(out_path), + "strict_ok": strict_ok, + "timing": timing, + "config": payload["config"], + "comparisons": comparisons, + }, + indent=2, + ) + ) return 0 if strict_ok else 1 @@ -3169,8 +3224,7 @@ def call_spec_commit( query_lens = inputs["query_lens"].to(torch.int64) live_token_mask = ( - torch.arange(max_query_len, dtype=torch.int64)[None, :] - < query_lens[:, None] + torch.arange(max_query_len, dtype=torch.int64)[None, :] < query_lens[:, None] ) outputs: dict[str, torch.Tensor | None] = { "standard_live_out": standard_out.detach().cpu()[live_token_mask], @@ -3286,13 +3340,18 @@ def call_spec_commit( out_path = Path(args.out) out_path.parent.mkdir(parents=True, exist_ok=True) torch.save(payload, out_path) - print(json.dumps({ - "out": str(out_path), - "strict_ok": strict_ok, - "elapsed_s": elapsed_s, - "config": payload["config"], - "comparisons": comparisons, - }, indent=2)) + print( + json.dumps( + { + "out": str(out_path), + "strict_ok": strict_ok, + "elapsed_s": elapsed_s, + "config": payload["config"], + "comparisons": comparisons, + }, + indent=2, + ) + ) return 0 if strict_ok else 1 @@ -3796,8 +3855,7 @@ def compile_direct_step( query_lens = inputs["query_lens"].to(torch.int64) live_token_mask = ( - torch.arange(max_query_len, dtype=torch.int64)[None, :] - < query_lens[:, None] + torch.arange(max_query_len, dtype=torch.int64)[None, :] < query_lens[:, None] ) outputs: dict[str, torch.Tensor | None] = { "standard_live_output": standard_output.detach().cpu()[live_token_mask], @@ -3986,12 +4044,15 @@ def compile_direct_step( "direct_compile_ssm_state", ] ) - strict_ok = all( - comparisons[name] is not None - and comparisons[name]["torch_equal"] - and comparisons[name]["max_diff"] == 0.0 - for name in strict_keys - ) and compile_step_error is None + strict_ok = ( + all( + comparisons[name] is not None + and comparisons[name]["torch_equal"] + and comparisons[name]["max_diff"] == 0.0 + for name in strict_keys + ) + and compile_step_error is None + ) accepted_hist = torch.bincount( inputs["num_accepted_tokens_seq"][:, 0].to(torch.int64), minlength=max_query_len + 1, @@ -4050,13 +4111,18 @@ def compile_direct_step( out_path = Path(args.out) out_path.parent.mkdir(parents=True, exist_ok=True) torch.save(payload, out_path) - print(json.dumps({ - "out": str(out_path), - "strict_ok": strict_ok, - "elapsed_s": elapsed_s, - "config": payload["config"], - "comparisons": comparisons, - }, indent=2)) + print( + json.dumps( + { + "out": str(out_path), + "strict_ok": strict_ok, + "elapsed_s": elapsed_s, + "config": payload["config"], + "comparisons": comparisons, + }, + indent=2, + ) + ) return 0 if strict_ok else 1 @@ -4231,9 +4297,7 @@ def make_parser() -> argparse.ArgumentParser: sigmoid.add_argument("--v-heads", type=int, default=8) sigmoid.add_argument("--head-k-dim", type=int, default=128) sigmoid.add_argument("--head-v-dim", type=int, default=128) - sigmoid.add_argument( - "--dtype", choices=("float16", "bfloat16"), default="float16" - ) + sigmoid.add_argument("--dtype", choices=("float16", "bfloat16"), default="float16") sigmoid.add_argument( "--state-dtype", choices=("same", "float16", "bfloat16", "float32"), @@ -4294,9 +4358,7 @@ def make_parser() -> argparse.ArgumentParser: model_mixed = subparsers.add_parser("run-model-mixed-qkv-route") model_mixed.add_argument("--out", required=True) model_mixed.add_argument("--input-file") - model_mixed.add_argument( - "--case-name", default="gdn_model_mixed_qkv_route_smoke" - ) + model_mixed.add_argument("--case-name", default="gdn_model_mixed_qkv_route_smoke") model_mixed.add_argument("--batch", type=int, default=1) model_mixed.add_argument("--seqlen", type=int, default=16) model_mixed.add_argument("--k-heads", type=int, default=4) @@ -4414,16 +4476,12 @@ def make_parser() -> argparse.ArgumentParser: ) packed_cg.add_argument("--random-initial-state", action="store_true") packed_cg.add_argument("--capture-null-indices", action="store_true") - packed_cg.set_defaults( - func=run_packed_recurrent_cudagraph_case, l2norm_inputs=True - ) + packed_cg.set_defaults(func=run_packed_recurrent_cudagraph_case, l2norm_inputs=True) conv_cg = subparsers.add_parser("run-conv-update-cudagraph") conv_cg.add_argument("--out", required=True) conv_cg.add_argument("--input-file") - conv_cg.add_argument( - "--case-name", default="gdn_conv_update_cudagraph_exactness" - ) + conv_cg.add_argument("--case-name", default="gdn_conv_update_cudagraph_exactness") conv_cg.add_argument("--batch", type=int, default=2) conv_cg.add_argument("--seqlen", type=int, default=1024) conv_cg.add_argument("--k-heads", type=int, default=4) @@ -4448,9 +4506,7 @@ def make_parser() -> argparse.ArgumentParser: spec_cg = subparsers.add_parser("run-spec-decode-cudagraph") spec_cg.add_argument("--out", required=True) spec_cg.add_argument("--input-file") - spec_cg.add_argument( - "--case-name", default="gdn_spec_decode_cudagraph_exactness" - ) + spec_cg.add_argument("--case-name", default="gdn_spec_decode_cudagraph_exactness") spec_cg.add_argument("--steps", type=int, default=512) spec_cg.add_argument("--num-spec", type=int, default=4) spec_cg.add_argument("--graph-rows", type=int) @@ -4462,6 +4518,12 @@ def make_parser() -> argparse.ArgumentParser: spec_cg.add_argument("--head-v-dim", type=int, default=128) spec_cg.add_argument("--conv-width", type=int, default=4) spec_cg.add_argument("--dtype", choices=("float16", "bfloat16"), default="float16") + spec_cg.add_argument( + "--beta-output-dtype", + choices=("input", "float32"), + default="float32", + help="Materialize split-verifier beta in the input dtype or FP32.", + ) spec_cg.add_argument( "--state-dtype", choices=("same", "float16", "bfloat16", "float32"), @@ -4489,9 +4551,7 @@ def make_parser() -> argparse.ArgumentParser: spec_mixed = subparsers.add_parser("run-spec-mixed-qkv-candidate") spec_mixed.add_argument("--out", required=True) spec_mixed.add_argument("--input-file") - spec_mixed.add_argument( - "--case-name", default="gdn_spec_mixed_qkv_candidate" - ) + spec_mixed.add_argument("--case-name", default="gdn_spec_mixed_qkv_candidate") spec_mixed.add_argument("--steps", type=int, default=512) spec_mixed.add_argument("--num-spec", type=int, default=4) spec_mixed.add_argument("--graph-rows", type=int) @@ -4505,6 +4565,12 @@ def make_parser() -> argparse.ArgumentParser: spec_mixed.add_argument( "--dtype", choices=("float16", "bfloat16"), default="float16" ) + spec_mixed.add_argument( + "--beta-output-dtype", + choices=("input", "float32"), + default="float32", + help="Materialize split-verifier beta in the input dtype or FP32.", + ) spec_mixed.add_argument( "--state-dtype", choices=("same", "float16", "bfloat16", "float32"), @@ -4532,9 +4598,7 @@ def make_parser() -> argparse.ArgumentParser: spec_op = subparsers.add_parser("run-spec-commit-vs-standard") spec_op.add_argument("--out", required=True) spec_op.add_argument("--input-file") - spec_op.add_argument( - "--case-name", default="gdn_spec_commit_vs_standard_exactness" - ) + spec_op.add_argument("--case-name", default="gdn_spec_commit_vs_standard_exactness") spec_op.add_argument("--steps", type=int, default=512) spec_op.add_argument("--num-spec", type=int, default=4) spec_op.add_argument("--graph-rows", type=int) 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 ae26060075..2f0d580bcf 100644 --- a/benchmarks/benchmark_sm70_model_tokens.py +++ b/benchmarks/benchmark_sm70_model_tokens.py @@ -20,7 +20,9 @@ from typing import Any SOURCE_ROOT = Path(__file__).resolve().parents[1] -FLASH_V100_ROOT = SOURCE_ROOT / "flash-attention-v100" +FLASH_V100_ROOT = Path( + os.environ.get("SM70_FLASH_V100_ROOT", SOURCE_ROOT / "flash-attention-v100") +) for source_path in (FLASH_V100_ROOT, SOURCE_ROOT): source_path_str = str(source_path) if source_path_str not in sys.path: @@ -500,7 +502,10 @@ def _sm70_tune_policy() -> dict[str, Any]: } -def _sm70_turbomind_policy() -> dict[str, Any]: +def _sm70_turbomind_policy( + max_num_seqs: int | None = None, + speculative_enabled: bool = False, +) -> dict[str, Any]: awq_turbomind = _env_bool("VLLM_SM70_AWQ_TURBOMIND", True) awq_tune_raw = os.environ.get("VLLM_SM70_AWQ_TUNE_SMALL_SHAPES") awq_tune0_pinned = awq_tune_raw == "0" @@ -521,7 +526,13 @@ def _sm70_turbomind_policy() -> dict[str, Any]: "VLLM_SM70_FP8_PREFILL_EXACT_DENSE", True, ) + fp8_qpn8 = _env_bool("VLLM_SM70_FP8_QPN8", False) + fp8_prefill_visible_dense_mm = _env_bool( + "VLLM_SM70_FP8_PREFILL_VISIBLE_DENSE_MM", + False, + ) nvfp4_turbomind = _env_bool("VLLM_SM70_NVFP4_TURBOMIND", False) + nvfp4_moe_grouped_prefill = _env_bool("VLLM_SM70_NVFP4_MOE_GROUPED_PREFILL", True) mxfp4_turbomind = _env_bool("VLLM_SM70_MXFP4_TURBOMIND", False) fp8_dequant_fallback = _env_bool("VLLM_SM70_FP8_DEQUANT_FALLBACK", True) fp8_moe_dequant_fallback = _env_bool( @@ -532,6 +543,10 @@ def _sm70_turbomind_policy() -> dict[str, Any]: "VLLM_SM70_UNQUANTIZED_MOE_0DOT3_CONFIG", True, ) + qwen36_mtp_moe_tuned_config = _env_bool( + "VLLM_SM70_QWEN36_MTP_MOE_TUNED_CONFIG", + True, + ) awq_warmup = _env_bool("VLLM_SM70_AWQ_WARMUP", True) awq_moe_disable = _env_bool("VLLM_SM70_AWQ_MOE_DISABLE", False) awq_moe_batched = _env_bool("VLLM_SM70_AWQ_MOE_BATCHED_GEMM", True) @@ -577,8 +592,25 @@ def _sm70_turbomind_policy() -> dict[str, Any]: "VLLM_SM70_FP8_PREFILL_EXACT_DENSE" ), "fp8_prefill_exact_dense_effective": fp8_prefill_exact_dense, + "VLLM_SM70_FP8_QPN8": os.environ.get("VLLM_SM70_FP8_QPN8"), + "fp8_qpn8_effective": fp8_qpn8, + "fp8_qpn8_native_decode_max_m": 8, + "fp8_qpn8_configured_max_num_seqs": max_num_seqs, + "fp8_qpn8_concurrency_contract": (max_num_seqs is None or max_num_seqs <= 8), + "fp8_qpn8_no_mtp_contract": not speculative_enabled, + "fp8_qpn8_large_m_fallback": "bounded_fp16_weight_workspace", + "fp8_qpn8_large_m_gated_temporary": "M_x_8704_fp16", + "VLLM_SM70_FP8_QPN8_LIBRARY": os.environ.get("VLLM_SM70_FP8_QPN8_LIBRARY"), + "VLLM_SM70_FP8_PREFILL_VISIBLE_DENSE_MM": os.environ.get( + "VLLM_SM70_FP8_PREFILL_VISIBLE_DENSE_MM" + ), + "fp8_prefill_visible_dense_mm_effective": fp8_prefill_visible_dense_mm, "VLLM_SM70_NVFP4_TURBOMIND": os.environ.get("VLLM_SM70_NVFP4_TURBOMIND"), "nvfp4_turbomind_effective": nvfp4_turbomind, + "VLLM_SM70_NVFP4_MOE_GROUPED_PREFILL": os.environ.get( + "VLLM_SM70_NVFP4_MOE_GROUPED_PREFILL" + ), + "nvfp4_moe_grouped_prefill_effective": nvfp4_moe_grouped_prefill, "VLLM_SM70_MXFP4_TURBOMIND": os.environ.get("VLLM_SM70_MXFP4_TURBOMIND"), "mxfp4_turbomind_effective": mxfp4_turbomind, "VLLM_SM70_FP8_DEQUANT_FALLBACK": os.environ.get( @@ -593,6 +625,10 @@ def _sm70_turbomind_policy() -> dict[str, Any]: "VLLM_SM70_UNQUANTIZED_MOE_0DOT3_CONFIG" ), "unquantized_moe_0dot3_config_effective": (unquantized_moe_0dot3_config), + "VLLM_SM70_QWEN36_MTP_MOE_TUNED_CONFIG": os.environ.get( + "VLLM_SM70_QWEN36_MTP_MOE_TUNED_CONFIG" + ), + "qwen36_mtp_moe_tuned_config_effective": qwen36_mtp_moe_tuned_config, "VLLM_SM70_AWQ_WARMUP": os.environ.get("VLLM_SM70_AWQ_WARMUP"), "awq_warmup_effective": awq_warmup, "VLLM_SM70_AWQ_WARMUP_MAX_M": os.environ.get("VLLM_SM70_AWQ_WARMUP_MAX_M"), @@ -765,6 +801,15 @@ def _sm70_turbomind_policy() -> dict[str, Any]: and not fp8_moe_batched_w2_dispatch and not f16_dense ), + "accepted_qwen38_fp8_qpn8_default_policy": ( + fp8_turbomind + and fp8_dequant_fallback + and fp8_dense_gated_silu + and fp8_prefill_exact_dense + and fp8_qpn8 + and (max_num_seqs is None or max_num_seqs <= 8) + and not speculative_enabled + ), "accepted_awq_moe_default_policy": ( awq_turbomind and awq_moe_safe_default_selector @@ -810,7 +855,9 @@ def _sm70_turbomind_policy() -> dict[str, Any]: "route_hit_oracle": ( "Accepted dense route-hit requires logs such as " "`SM70 AWQ TurboMind dense path enabled`, " - "`SM70 FP8 TurboMind W8A16 dense path enabled`, and when warmup " + "`SM70 FP8 TurboMind W8A16 dense path enabled`, and for the " + "Qwen3.8-27B TP4 QPN8 default " + "`Memory-neutral SM70 FP8 QPN8 path enabled`; when warmup " "is relevant `SM70 AWQ warmup finished`. AWQ MoE production " "throughput evidence must keep the default fast route: batched " "GEMM enabled, legacy single-token compact enabled, and no " @@ -885,6 +932,10 @@ def _sm70_attention_policy(kv_cache_dtype: Any) -> dict[str, Any]: "VLLM_FLASH_V100_XQA_E5M2_PAIR_LOAD", True, ) + e5m2_batch_wide_load = _env_bool( + "VLLM_FLASH_V100_XQA_E5M2_BATCH_WIDE_LOAD", + True, + ) decode_fp8_xqa_min_seq_len = _env_int( "VLLM_FLASH_V100_DECODE_FP8_XQA_MIN_SEQ_LEN", 16384, @@ -996,6 +1047,10 @@ def _sm70_attention_policy(kv_cache_dtype: Any) -> dict[str, Any]: "VLLM_FLASH_V100_XQA_E5M2_PAIR_LOAD" ), "e5m2_pair_load_effective": e5m2_pair_load, + "VLLM_FLASH_V100_XQA_E5M2_BATCH_WIDE_LOAD": os.environ.get( + "VLLM_FLASH_V100_XQA_E5M2_BATCH_WIDE_LOAD" + ), + "e5m2_batch_wide_load_effective": e5m2_batch_wide_load, "VLLM_FLASH_V100_DECODE_FP8_XQA_MIN_SEQ_LEN": os.environ.get( "VLLM_FLASH_V100_DECODE_FP8_XQA_MIN_SEQ_LEN" ), @@ -1588,11 +1643,14 @@ def _dump(args: argparse.Namespace) -> int: ) prompts = _load_prompts(args) - if capture_decode_after_prefix_warmup and len(prompts) != 1: - raise ValueError( - "--cuda-profiler-capture-decode-after-prefix-warmup requires " - "exactly one prompt" - ) + if args.repeat_count < 1: + raise ValueError("--repeat-count must be positive") + if args.repeat_count > 1 and len(prompts) != 1: + raise ValueError("--repeat-count greater than one requires one prompt") + if args.repeat_count > 1 and ( + args.cuda_profiler_capture_generate or capture_decode_after_prefix_warmup + ): + raise ValueError("CUDA-profiler capture does not support repeated requests") sampling_params = SamplingParams( max_tokens=args.max_tokens, temperature=args.temperature, @@ -1654,21 +1712,30 @@ def _dump(args: argparse.Namespace) -> int: torch.accelerator.synchronize() torch.cuda.cudart().cudaProfilerStart() - start = time.perf_counter() + generate_seconds_by_repeat: list[float] = [] + outputs = [] try: - if args.sequential_prompts: - outputs = [] - for prompt in prompts: - outputs.extend(llm.generate([prompt], sampling_params)) - else: - outputs = llm.generate(prompts, sampling_params) + for repeat_index in range(args.repeat_count): + if ( + repeat_index + and args.reset_prefix_cache_between_repeats + and not llm.reset_prefix_cache() + ): + raise RuntimeError("Failed to reset the idle prefix cache") + generate_start = time.perf_counter() + if args.sequential_prompts: + for prompt in prompts: + outputs.extend(llm.generate([prompt], sampling_params)) + else: + outputs.extend(llm.generate(prompts, sampling_params)) + generate_seconds_by_repeat.append(time.perf_counter() - generate_start) finally: if args.cuda_profiler_capture_generate or capture_decode_after_prefix_warmup: import torch torch.accelerator.synchronize() torch.cuda.cudart().cudaProfilerStop() - generate_seconds = time.perf_counter() - start + generate_seconds = sum(generate_seconds_by_repeat) metrics_snapshot = _metric_snapshot(llm) import torch @@ -1747,7 +1814,10 @@ def _dump(args: argparse.Namespace) -> int: ], "env": _tracked_env(), "sm70_tune_policy": _sm70_tune_policy(), - "sm70_turbomind_policy": _sm70_turbomind_policy(), + "sm70_turbomind_policy": _sm70_turbomind_policy( + int(llm_kwargs["max_num_seqs"]), + speculative_enabled=llm_kwargs.get("speculative_config") is not None, + ), "sm70_attention_policy": _sm70_attention_policy( llm_kwargs.get("kv_cache_dtype") ), @@ -1778,6 +1848,9 @@ def _dump(args: argparse.Namespace) -> int: "sequential_prompts": args.sequential_prompts, "load_seconds": load_seconds, "generate_seconds": generate_seconds, + "generate_seconds_by_repeat": generate_seconds_by_repeat, + "repeat_count": args.repeat_count, + "reset_prefix_cache_between_repeats": (args.reset_prefix_cache_between_repeats), "total_output_tokens": total_output_tokens, "records": records, } @@ -2527,12 +2600,14 @@ def _parse_args() -> argparse.Namespace: "--cuda-profiler-capture-decode-after-prefix-warmup", action="store_true", help=( - "Prime one prompt with prefix caching while the CUDA profiler is " + "Prime the prompts with prefix caching while the CUDA profiler is " "off, then capture only the second matching llm.generate call. " - "Requires one non-sequential prompt and CUDA graphs." + "Requires non-sequential prompts and CUDA graphs." ), ) parser.add_argument("--max-tokens", type=int, default=16) + parser.add_argument("--repeat-count", type=int, default=1) + parser.add_argument("--reset-prefix-cache-between-repeats", action="store_true") parser.add_argument("--temperature", type=float, default=0.0) parser.add_argument("--top-p", type=float, default=1.0) parser.add_argument("--top-k", type=int, default=-1) diff --git a/benchmarks/benchmark_sm70_mxfp4_moe_active_experts.py b/benchmarks/benchmark_sm70_mxfp4_moe_active_experts.py index b58ba25399..b79a169e61 100644 --- a/benchmarks/benchmark_sm70_mxfp4_moe_active_experts.py +++ b/benchmarks/benchmark_sm70_mxfp4_moe_active_experts.py @@ -13,6 +13,7 @@ import json import os from dataclasses import dataclass +from pathlib import Path import torch @@ -399,6 +400,7 @@ def make_buffers() -> dict[str, torch.Tensor]: generic = make_buffers() direct = make_buffers() + direct_order = make_buffers() def generic_call() -> None: generic["output"].zero_() @@ -498,11 +500,61 @@ def direct_call() -> None: direct["output"], ) + def direct_order_call() -> None: + route_ids = topk_ids.view(-1) + sm70_ops.mxfp4_moe_dense_stage_sm70_out( + direct_order["gate_up"], + x, + direct_order["compact_offsets"], + route_ids, + w13_ptrs_w, + w13_ptrs_s, + top_k, + 4096, + 512, + 32, + ) + torch.ops._C.silu_and_mul_with_clamp( + direct_order["intermediate"], direct_order["gate_up"], 10.0 + ) + sm70_ops.mxfp4_moe_dense_stage_sm70_out( + direct_order["sorted_output"], + direct_order["intermediate"], + direct_order["compact_offsets"], + route_ids, + w2_ptrs_w, + w2_ptrs_s, + top_k, + 256, + 4096, + 32, + ) + sm70_ops.awq_moe_single_token_weighted_reduce_out( + direct_order["sorted_output"], + topk_weights, + direct_order["token_expert_indices"], + direct_order["output"], + top_k, + 4096, + ) + generic_call() direct_call() + direct_order_call() torch.accelerator.synchronize() initial_equal = torch.equal(generic["output"], direct["output"]) initial_max_abs = float((generic["output"] - direct["output"]).abs().max().item()) + direct_order_initial_equal = torch.equal(direct["output"], direct_order["output"]) + direct_order_initial_max_abs = float( + (direct["output"] - direct_order["output"]).abs().max().item() + ) + initial_sorted_src = torch.argsort(topk_ids.view(-1), stable=True) + direct_order_stage_parity = { + f"{name}_equal": torch.equal( + direct[name], direct_order[name].index_select(0, initial_sorted_src) + ) + for name in ("gate_up", "intermediate", "sorted_output") + } stage_parity = { "gate_up_equal": torch.equal(generic["gate_up"], direct["gate_up"]), "gate_up_max_abs": float( @@ -527,28 +579,52 @@ def direct_call() -> None: generic_graph = _capture(generic_call) direct_graph = _capture(direct_call) + direct_order_graph = _capture(direct_order_call) generic_ms = _time_graph(generic_graph, repeats) direct_ms = _time_graph(direct_graph, repeats) + direct_order_ms = _time_graph(direct_order_graph, repeats) topk_ids.copy_(torch.tensor([route_b], dtype=torch.int32, device=device)) generic_graph.replay() direct_graph.replay() + direct_order_graph.replay() torch.accelerator.synchronize() replay_equal = torch.equal(generic["output"], direct["output"]) replay_max_abs = float((generic["output"] - direct["output"]).abs().max().item()) + direct_order_replay_equal = torch.equal(direct["output"], direct_order["output"]) + direct_order_replay_max_abs = float( + (direct["output"] - direct_order["output"]).abs().max().item() + ) result = { "generic_graph_ms": generic_ms, "direct_graph_ms": direct_ms, + "direct_order_graph_ms": direct_order_ms, "speedup": generic_ms / direct_ms, "projected_savings_ms_per_token": (generic_ms - direct_ms) * 43, + "direct_order_speedup_vs_direct": direct_ms / direct_order_ms, + "direct_order_projected_savings_ms_per_token": (direct_ms - direct_order_ms) + * 43, "initial_bitwise_equal": initial_equal, "initial_max_abs": initial_max_abs, "dynamic_replay_bitwise_equal": replay_equal, "dynamic_replay_max_abs": replay_max_abs, + "direct_order_initial_bitwise_equal": direct_order_initial_equal, + "direct_order_initial_max_abs": direct_order_initial_max_abs, + "direct_order_dynamic_replay_bitwise_equal": direct_order_replay_equal, + "direct_order_dynamic_replay_max_abs": direct_order_replay_max_abs, "initial_stage_parity": stage_parity, + "direct_order_initial_sorted_stage_parity": direct_order_stage_parity, } - if not initial_equal or not replay_equal: + if not all( + ( + initial_equal, + replay_equal, + direct_order_initial_equal, + direct_order_replay_equal, + *direct_order_stage_parity.values(), + ) + ): raise RuntimeError(f"MXFP4 direct top-6 correctness gate failed: {result}") return result @@ -557,16 +633,19 @@ def benchmark_verifier_m8_pipeline( *, num_experts: int, top_k: int, + verifier_tokens: int, repeats: int, seed: int, route_case: str = "mixed", input_scale: float = 0.01, require_grouped_bitwise: bool = True, ) -> dict[str, object]: - """Compare dense-256, active-48 loop, and grouped verifier pipelines.""" - num_tokens = 8 + """Compare dense, active-loop, and grouped verifier pipelines.""" + num_tokens = int(verifier_tokens) if num_experts != 256 or top_k != 6: raise ValueError("The verifier benchmark requires 256 experts/top-k=6") + if not 2 <= num_tokens <= 8: + raise ValueError("The verifier benchmark requires M in [2, 8]") torch.manual_seed(seed) device = torch.device("cuda") @@ -670,11 +749,22 @@ def make_buffers() -> dict[str, torch.Tensor]: dense = make_buffers() active = make_buffers() grouped = make_buffers() + expert_grouped = make_buffers() def pipeline_call( - buffers: dict[str, torch.Tensor], *, active_only: bool, grouped_m8: bool + buffers: dict[str, torch.Tensor], + *, + active_only: bool, + grouped_verifier: bool, + expert_rows: bool = False, ) -> None: - os.environ["VLLM_SM70_MXFP4_MOE_GROUPED_M8"] = "1" if grouped_m8 else "0" + os.environ["VLLM_SM70_MXFP4_MOE_GROUPED_M8"] = "0" + os.environ["VLLM_SM70_MXFP4_MOE_GROUPED_VERIFIER"] = ( + "1" if grouped_verifier else "0" + ) + os.environ["VLLM_SM70_MXFP4_MOE_GROUPED_M8_EXPERT_ROWS"] = ( + "1" if expert_rows else "0" + ) buffers["output"].zero_() buffers["topk_ids_i32"].copy_(topk_ids) buffers["permuted_idx"].fill_(total_slots) @@ -698,8 +788,16 @@ def pipeline_call( buffers["expert_offsets"].copy_(buffers["expert_offsets64"]) if active_only: stage_offsets = buffers["compact_offsets"] - if grouped_m8: - stage_expert_ids = buffers["permuted_experts_id"] + if grouped_verifier: + if expert_rows: + _compact_mxfp4_active_experts( + buffers["permuted_experts_id"], + buffers["compact_offsets"], + buffers["active_expert_ids"], + ) + stage_expert_ids = buffers["active_expert_ids"] + else: + stage_expert_ids = buffers["permuted_experts_id"] else: _compact_mxfp4_active_experts( buffers["permuted_experts_id"], @@ -749,20 +847,32 @@ def pipeline_call( ) def dense_call() -> None: - pipeline_call(dense, active_only=False, grouped_m8=False) + pipeline_call(dense, active_only=False, grouped_verifier=False) def active_call() -> None: - pipeline_call(active, active_only=True, grouped_m8=False) + pipeline_call(active, active_only=True, grouped_verifier=False) def grouped_call() -> None: - pipeline_call(grouped, active_only=True, grouped_m8=True) + pipeline_call(grouped, active_only=True, grouped_verifier=True) + + def expert_grouped_call() -> None: + pipeline_call( + expert_grouped, + active_only=True, + grouped_verifier=True, + expert_rows=True, + ) dense_call() active_call() grouped_call() + expert_grouped_call() torch.accelerator.synchronize() initial_equal = torch.equal(dense["output"], active["output"]) grouped_initial_equal = torch.equal(dense["output"], grouped["output"]) + expert_grouped_initial_equal = torch.equal( + dense["output"], expert_grouped["output"] + ) stage_parity = { name: torch.equal(dense[name], active[name]) for name in ("gate_up", "intermediate", "sorted_output", "output") @@ -771,14 +881,26 @@ def grouped_call() -> None: name: torch.equal(dense[name], grouped[name]) for name in ("gate_up", "intermediate", "sorted_output", "output") } + expert_grouped_stage_parity = { + name: torch.equal(dense[name], expert_grouped[name]) + for name in ("gate_up", "intermediate", "sorted_output", "output") + } grouped_stage_max_abs = { name: float((dense[name] - grouped[name]).abs().max().item()) for name in ("gate_up", "intermediate", "sorted_output", "output") } + expert_grouped_stage_max_abs = { + name: float((dense[name] - expert_grouped[name]).abs().max().item()) + for name in ("gate_up", "intermediate", "sorted_output", "output") + } grouped_stages_finite = all( bool(torch.isfinite(grouped[name]).all().item()) for name in ("gate_up", "intermediate", "sorted_output", "output") ) + expert_grouped_stages_finite = all( + bool(torch.isfinite(expert_grouped[name]).all().item()) + for name in ("gate_up", "intermediate", "sorted_output", "output") + ) expected_active_ids = sorted({expert for row in route_a for expert in row}) actual_active_ids = active["active_expert_ids"][: len(expected_active_ids)] compact_ids_match = actual_active_ids.cpu().tolist() == expected_active_ids @@ -786,14 +908,17 @@ def grouped_call() -> None: dense_graph = _capture(dense_call) active_graph = _capture(active_call) grouped_graph = _capture(grouped_call) + expert_grouped_graph = _capture(expert_grouped_call) dense_ms = _time_graph(dense_graph, repeats) active_ms = _time_graph(active_graph, repeats) grouped_ms = _time_graph(grouped_graph, repeats) + expert_grouped_ms = _time_graph(expert_grouped_graph, repeats) topk_ids.copy_(torch.tensor(route_b, dtype=torch.int32, device=device)) dense_graph.replay() active_graph.replay() grouped_graph.replay() + expert_grouped_graph.replay() torch.accelerator.synchronize() replay_equal = torch.equal(dense["output"], active["output"]) replay_max_abs = float((dense["output"] - active["output"]).abs().max().item()) @@ -801,12 +926,21 @@ def grouped_call() -> None: grouped_replay_max_abs = float( (dense["output"] - grouped["output"]).abs().max().item() ) + expert_grouped_replay_equal = torch.equal(dense["output"], expert_grouped["output"]) + expert_grouped_replay_max_abs = float( + (dense["output"] - expert_grouped["output"]).abs().max().item() + ) route_b_output = grouped["output"].clone() + expert_grouped_route_b_output = expert_grouped["output"].clone() topk_ids.copy_(torch.tensor(route_a, dtype=torch.int32, device=device)) grouped_graph.replay() + expert_grouped_graph.replay() torch.accelerator.synchronize() route_changes_output = not torch.equal(grouped["output"], route_b_output) + expert_grouped_route_changes_output = not torch.equal( + expert_grouped["output"], expert_grouped_route_b_output + ) result = { "route_case": route_case, @@ -819,22 +953,37 @@ def grouped_call() -> None: "dense_graph_ms": dense_ms, "active_graph_ms": active_ms, "grouped_graph_ms": grouped_ms, + "expert_grouped_graph_ms": expert_grouped_ms, "speedup": dense_ms / active_ms, "grouped_vs_active_speedup": active_ms / grouped_ms, + "expert_grouped_vs_slot_grouped_speedup": grouped_ms / expert_grouped_ms, + "expert_grouped_projected_savings_ms_per_43_layers": ( + grouped_ms - expert_grouped_ms + ) + * 43, "grouped_projected_savings_ms_per_43_layers": (active_ms - grouped_ms) * 43, "projected_savings_ms_per_43_layers": (dense_ms - active_ms) * 43, "initial_bitwise_equal": initial_equal, "grouped_initial_bitwise_equal": grouped_initial_equal, + "expert_grouped_initial_bitwise_equal": expert_grouped_initial_equal, "initial_stage_bitwise_equal": stage_parity, "grouped_initial_stage_bitwise_equal": grouped_stage_parity, + "expert_grouped_initial_stage_bitwise_equal": expert_grouped_stage_parity, "grouped_initial_stage_max_abs": grouped_stage_max_abs, + "expert_grouped_initial_stage_max_abs": expert_grouped_stage_max_abs, "grouped_stages_finite": grouped_stages_finite, + "expert_grouped_stages_finite": expert_grouped_stages_finite, "compact_ids_match": compact_ids_match, "dynamic_replay_bitwise_equal": replay_equal, "dynamic_replay_max_abs": replay_max_abs, "grouped_dynamic_replay_bitwise_equal": grouped_replay_equal, "grouped_dynamic_replay_max_abs": grouped_replay_max_abs, + "expert_grouped_dynamic_replay_bitwise_equal": expert_grouped_replay_equal, + "expert_grouped_dynamic_replay_max_abs": expert_grouped_replay_max_abs, "dynamic_route_changes_output": route_changes_output, + "expert_grouped_dynamic_route_changes_output": ( + expert_grouped_route_changes_output + ), } base_gate_passed = ( initial_equal @@ -843,6 +992,8 @@ def grouped_call() -> None: and replay_equal and route_changes_output and grouped_stages_finite + and expert_grouped_stages_finite + and expert_grouped_route_changes_output ) grouped_bitwise_gate_passed = ( grouped_initial_equal @@ -854,7 +1005,7 @@ def grouped_call() -> None: if not base_gate_passed or ( require_grouped_bitwise and not grouped_bitwise_gate_passed ): - raise RuntimeError(f"MXFP4 verifier M8 correctness gate failed: {result}") + raise RuntimeError(f"MXFP4 verifier correctness gate failed: {result}") return result @@ -992,7 +1143,15 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--seed", type=int, default=17) parser.add_argument("--input-scale", type=float, default=0.01) parser.add_argument("--full-pipeline", action="store_true") + parser.add_argument("--json-out", type=Path) parser.add_argument("--verifier-m8-pipeline", action="store_true") + parser.add_argument( + "--verifier-tokens", + type=int, + choices=range(2, 9), + default=8, + help="Target verifier row count M for the grouped pipeline gate.", + ) parser.add_argument( "--allow-grouped-numeric-drift", action="store_true", @@ -1002,9 +1161,9 @@ def parse_args() -> argparse.Namespace: ) parser.add_argument( "--route-case", - choices=("mixed", "unique48", "hot6"), + choices=("mixed", "unique48", "hot6", "all"), default="mixed", - help="Verifier M8 expert-overlap distribution.", + help="Verifier expert-overlap distribution.", ) parser.add_argument( "--profile-active-once", @@ -1023,19 +1182,28 @@ def main() -> int: args = parse_args() _require_sm70() if args.verifier_m8_pipeline: - result = benchmark_verifier_m8_pipeline( - num_experts=args.num_experts, - top_k=args.top_k, - repeats=args.repeats, - seed=args.seed, - route_case=args.route_case, - input_scale=args.input_scale, - require_grouped_bitwise=not args.allow_grouped_numeric_drift, + route_cases = ( + ("mixed", "unique48", "hot6") + if args.route_case == "all" + else (args.route_case,) ) + result = { + route_case: benchmark_verifier_m8_pipeline( + num_experts=args.num_experts, + top_k=args.top_k, + verifier_tokens=args.verifier_tokens, + repeats=args.repeats, + seed=args.seed + index, + route_case=route_case, + input_scale=args.input_scale, + require_grouped_bitwise=not args.allow_grouped_numeric_drift, + ) + for index, route_case in enumerate(route_cases) + } print( json.dumps( { - "benchmark": "sm70_mxfp4_moe_verifier_m8_pipeline", + "benchmark": "sm70_mxfp4_moe_verifier_pipeline", "device": torch.cuda.get_device_name(), "result": result, }, @@ -1050,16 +1218,16 @@ def main() -> int: repeats=args.repeats, seed=args.seed, ) - print( - json.dumps( - { - "benchmark": "sm70_mxfp4_moe_direct_top6_pipeline", - "device": torch.cuda.get_device_name(), - "result": result, - }, - indent=2, - ) - ) + payload = { + "benchmark": "sm70_mxfp4_moe_direct_top6_pipeline", + "device": torch.cuda.get_device_name(), + "result": result, + } + rendered = json.dumps(payload, indent=2) + if args.json_out is not None: + args.json_out.parent.mkdir(parents=True, exist_ok=True) + args.json_out.write_text(rendered + "\n", encoding="utf-8") + print(rendered) return 0 if args.profile_active_once: if args.stage == "both": 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_prefill_gemm_collective.py b/benchmarks/benchmark_sm70_prefill_gemm_collective.py new file mode 100644 index 0000000000..807f761ca4 --- /dev/null +++ b/benchmarks/benchmark_sm70_prefill_gemm_collective.py @@ -0,0 +1,395 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Compare GEMM+all-reduce with pipelined GEMM+RS+AG on SM70. + +Launch this benchmark with torchrun. It models the exact Qwen3.8 TP4 +row-parallel down/output projection boundary without loading the model. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import statistics +from collections.abc import Callable +from pathlib import Path +from typing import Any + +import torch +import torch.distributed as dist +from torch.distributed._symmetric_memory import enable_symm_mem_for_group + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--m", type=int, default=8000) + parser.add_argument("--k", type=int, default=4352) + parser.add_argument("--n", type=int, default=5120) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--iters", type=int, default=10) + parser.add_argument("--trials", type=int, default=7) + parser.add_argument("--chunk-counts", type=int, nargs="+", default=[2, 4, 8]) + parser.add_argument("--chunk-axis", choices=("m", "n"), default="m") + parser.add_argument("--skip-symm-mem", action="store_true") + parser.add_argument("--include-gemma-norm", action="store_true") + parser.add_argument("--norm-epsilon", type=float, default=1e-6) + parser.add_argument("--seed", type=int, default=20260823) + parser.add_argument("--json-out", type=Path, required=True) + return parser.parse_args() + + +def _event_trials( + launch: Callable[[], None], warmup: int, iters: int, trials: int +) -> dict[str, Any]: + for _ in range(warmup): + launch() + torch.accelerator.synchronize() + + samples: list[float] = [] + for _ in range(trials): + dist.barrier() + start = torch.Event(enable_timing=True) + end = torch.Event(enable_timing=True) + start.record() + for _ in range(iters): + launch() + end.record() + end.synchronize() + samples.append(float(start.elapsed_time(end) / iters)) + return { + "median_ms": float(statistics.median(samples)), + "mean_ms": float(statistics.mean(samples)), + "min_ms": min(samples), + "max_ms": max(samples), + "samples_ms": samples, + } + + +def _digest(tensor: torch.Tensor) -> str: + raw = tensor.detach().contiguous().cpu().view(torch.uint8).numpy().tobytes() + return hashlib.sha256(raw).hexdigest() + + +def main() -> int: + args = _parse_args() + local_rank = int(os.environ["LOCAL_RANK"]) + torch.accelerator.set_device_index(local_rank) + dist.init_process_group(backend="nccl") + rank = dist.get_rank() + world_size = dist.get_world_size() + if world_size != 4: + raise RuntimeError("This benchmark requires TP4.") + if torch.cuda.get_device_capability(local_rank) != (7, 0): + raise RuntimeError("This benchmark requires SM70/V100 GPUs.") + if args.m % world_size: + raise ValueError("M must be divisible by the world size.") + chunk_extent = args.m if args.chunk_axis == "m" else args.n + if any(chunks <= 0 or chunk_extent % chunks for chunks in args.chunk_counts): + raise ValueError( + f"Every chunk count must be positive and divide {args.chunk_axis.upper()}." + ) + if args.include_gemma_norm: + import vllm.kernels # noqa: F401 + + if args.n != 5120: + raise ValueError("The SM70 Gemma long-prefill norm requires N=5120.") + if not hasattr(torch.ops._C, "sm70_gemma_long_prefill_fused_add_rms_norm"): + raise RuntimeError("The active extension lacks the SM70 Gemma norm op.") + + group_name = dist.group.WORLD.group_name + if not args.skip_symm_mem: + enable_symm_mem_for_group(group_name) + device = torch.device("cuda", local_rank) + generator = torch.Generator(device=device).manual_seed(args.seed + rank) + inputs = torch.randn( + (args.m, args.k), + generator=generator, + device=device, + dtype=torch.float16, + ).mul_(0.1) + weight = torch.randn( + (args.k, args.n), + generator=generator, + device=device, + dtype=torch.float16, + ).mul_(0.01) + baseline_out = torch.empty((args.m, args.n), device=device, dtype=torch.float16) + pipelined_out = torch.empty_like(baseline_out) + plain_sp_mm_out = torch.empty_like(baseline_out) + plain_sp_out = torch.empty_like(baseline_out) + plain_sp_shard = torch.empty( + (args.m // world_size, args.n), device=device, dtype=torch.float16 + ) + chunked_outputs = { + chunks: torch.empty_like(baseline_out) for chunks in args.chunk_counts + } + chunked_n_comm_outputs: dict[int, list[torch.Tensor]] = {} + if args.chunk_axis == "n": + chunked_n_comm_outputs = { + chunks: [ + torch.empty( + (args.m, args.n // chunks), + device=device, + dtype=torch.float16, + ) + for _ in range(chunks) + ] + for chunks in args.chunk_counts + } + residual = None + norm_weight = None + baseline_normalized = None + baseline_residual_out = None + pipelined_normalized_shard = None + pipelined_residual_shard = None + chunked_normalized: dict[int, torch.Tensor] = {} + chunked_residual_out: dict[int, torch.Tensor] = {} + if args.include_gemma_norm: + residual = torch.randn( + (args.m, args.n), + generator=torch.Generator(device=device).manual_seed(args.seed + 10000), + device=device, + dtype=torch.float32, + ).mul_(0.1) + norm_weight = torch.randn( + (args.n,), + generator=torch.Generator(device=device).manual_seed(args.seed + 20000), + device=device, + dtype=torch.bfloat16, + ).mul_(0.01) + baseline_normalized = torch.empty_like(baseline_out) + baseline_residual_out = torch.empty_like(residual) + pipelined_normalized_shard = torch.empty( + (args.m // world_size, args.n), device=device, dtype=torch.float16 + ) + pipelined_residual_shard = torch.empty( + (args.m // world_size, args.n), device=device, dtype=torch.float32 + ) + chunked_normalized = { + chunks: torch.empty_like(baseline_out) for chunks in args.chunk_counts + } + chunked_residual_out = { + chunks: torch.empty_like(residual) for chunks in args.chunk_counts + } + + def norm_out( + output: torch.Tensor, + residual_output: torch.Tensor, + input: torch.Tensor, + input_residual: torch.Tensor, + ) -> None: + assert norm_weight is not None + torch.ops._C.sm70_gemma_long_prefill_fused_add_rms_norm( + output, + residual_output, + input, + input_residual, + norm_weight, + args.norm_epsilon, + ) + + def baseline() -> None: + torch.mm(inputs, weight, out=baseline_out) + dist.all_reduce(baseline_out) + if args.include_gemma_norm: + assert baseline_normalized is not None + assert baseline_residual_out is not None + assert residual is not None + norm_out( + baseline_normalized, + baseline_residual_out, + baseline_out, + residual, + ) + + def pipelined() -> None: + shard = torch.ops.symm_mem.fused_matmul_reduce_scatter( + inputs, + weight, + "sum", + 0, + group_name, + ) + if args.include_gemma_norm: + assert residual is not None + assert pipelined_normalized_shard is not None + assert pipelined_residual_shard is not None + residual_shard = residual.narrow( + 0, rank * (args.m // world_size), args.m // world_size + ) + norm_out( + pipelined_normalized_shard, + pipelined_residual_shard, + shard, + residual_shard, + ) + shard = pipelined_normalized_shard + dist.all_gather_into_tensor(pipelined_out, shard) + + def plain_sp() -> None: + torch.mm(inputs, weight, out=plain_sp_mm_out) + dist.reduce_scatter_tensor(plain_sp_shard, plain_sp_mm_out) + shard = plain_sp_shard + if args.include_gemma_norm: + assert residual is not None + assert pipelined_normalized_shard is not None + assert pipelined_residual_shard is not None + residual_shard = residual.narrow( + 0, rank * (args.m // world_size), args.m // world_size + ) + norm_out( + pipelined_normalized_shard, + pipelined_residual_shard, + shard, + residual_shard, + ) + shard = pipelined_normalized_shard + dist.all_gather_into_tensor(plain_sp_out, shard) + + def make_chunked_all_reduce(chunks: int) -> Callable[[], None]: + if args.chunk_axis == "m": + lhs_chunks = inputs.chunk(chunks, dim=0) + rhs_chunks = (weight,) * chunks + output_chunks = chunked_outputs[chunks].chunk(chunks, dim=0) + else: + lhs_chunks = (inputs,) * chunks + rhs_chunks = weight.chunk(chunks, dim=1) + output_chunks = chunked_n_comm_outputs[chunks] + result_chunks = chunked_outputs[chunks].chunk( + chunks, dim=0 if args.chunk_axis == "m" else 1 + ) + + def launch() -> None: + works = [] + for lhs, rhs, output_chunk in zip(lhs_chunks, rhs_chunks, output_chunks): + torch.mm(lhs, rhs, out=output_chunk) + works.append(dist.all_reduce(output_chunk, async_op=True)) + for work, output_chunk, result_chunk in zip( + works, output_chunks, result_chunks + ): + work.wait() + if args.chunk_axis == "n": + result_chunk.copy_(output_chunk) + if args.include_gemma_norm: + assert residual is not None + norm_out( + chunked_normalized[chunks], + chunked_residual_out[chunks], + chunked_outputs[chunks], + residual, + ) + + return launch + + baseline_timing = _event_trials(baseline, args.warmup, args.iters, args.trials) + plain_sp_timing = _event_trials(plain_sp, args.warmup, args.iters, args.trials) + chunked_launches = { + chunks: make_chunked_all_reduce(chunks) for chunks in args.chunk_counts + } + chunked_timings = { + chunks: _event_trials(launch, args.warmup, args.iters, args.trials) + for chunks, launch in chunked_launches.items() + } + pipelined_timing = None + if not args.skip_symm_mem: + pipelined_timing = _event_trials( + pipelined, args.warmup, args.iters, args.trials + ) + baseline() + plain_sp() + for launch in chunked_launches.values(): + launch() + if not args.skip_symm_mem: + pipelined() + torch.accelerator.synchronize() + + if args.include_gemma_norm: + assert baseline_normalized is not None + baseline_result = baseline_normalized + else: + baseline_result = baseline_out + local_payload = { + "rank": rank, + "baseline": baseline_timing, + "baseline_sha256": _digest(baseline_result), + "plain_sp": None, + "pipelined": None, + "chunked": {}, + } + plain_sp_difference = (baseline_result.float() - plain_sp_out.float()).abs() + local_payload["plain_sp"] = { + "timing": plain_sp_timing, + "speedup": baseline_timing["median_ms"] / plain_sp_timing["median_ms"], + "max_abs_diff": float(plain_sp_difference.max().item()), + "mean_abs_diff": float(plain_sp_difference.mean().item()), + "equal_elements": int((baseline_result == plain_sp_out).sum().item()), + "elements": baseline_result.numel(), + "sha256": _digest(plain_sp_out), + } + if pipelined_timing is not None: + difference = (baseline_result.float() - pipelined_out.float()).abs() + local_payload["pipelined"] = { + "timing": pipelined_timing, + "speedup": baseline_timing["median_ms"] / pipelined_timing["median_ms"], + "max_abs_diff": float(difference.max().item()), + "mean_abs_diff": float(difference.mean().item()), + "equal_elements": int((baseline_result == pipelined_out).sum().item()), + "elements": baseline_result.numel(), + "sha256": _digest(pipelined_out), + } + for chunks, output in chunked_outputs.items(): + chunked_result = ( + chunked_normalized[chunks] if args.include_gemma_norm else output + ) + chunked_difference = (baseline_result.float() - chunked_result.float()).abs() + local_payload["chunked"][str(chunks)] = { + "timing": chunked_timings[chunks], + "speedup": baseline_timing["median_ms"] + / chunked_timings[chunks]["median_ms"], + "max_abs_diff": float(chunked_difference.max().item()), + "mean_abs_diff": float(chunked_difference.mean().item()), + "equal_elements": int((baseline_result == chunked_result).sum().item()), + "elements": baseline_result.numel(), + "sha256": _digest(chunked_result), + } + gathered: list[dict[str, Any] | None] = [None] * world_size + dist.all_gather_object(gathered, local_payload) + if rank == 0: + payload = { + "environment": { + "world_size": world_size, + "device": torch.cuda.get_device_name(local_rank), + "torch": torch.__version__, + "torch_cuda": torch.version.cuda, + "m": args.m, + "k": args.k, + "n": args.n, + "dtype": "float16", + "warmup": args.warmup, + "iters": args.iters, + "trials": args.trials, + "chunk_counts": args.chunk_counts, + "chunk_axis": args.chunk_axis, + "symm_mem": not args.skip_symm_mem, + "include_gemma_norm": args.include_gemma_norm, + "norm_epsilon": args.norm_epsilon, + "seed": args.seed, + "cuda_visible_devices": os.getenv("CUDA_VISIBLE_DEVICES"), + "nccl_algo": os.getenv("NCCL_ALGO"), + "nccl_proto": os.getenv("NCCL_PROTO"), + "nccl_min_nchannels": os.getenv("NCCL_MIN_NCHANNELS"), + "nccl_max_nchannels": os.getenv("NCCL_MAX_NCHANNELS"), + }, + "ranks": gathered, + } + args.json_out.parent.mkdir(parents=True, exist_ok=True) + args.json_out.write_text(json.dumps(payload, indent=2) + "\n") + print(json.dumps(payload, indent=2)) + dist.destroy_process_group() + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) 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_sglang_push_chain.py b/benchmarks/benchmark_sm70_sglang_push_chain.py new file mode 100644 index 0000000000..52b38ece3f --- /dev/null +++ b/benchmarks/benchmark_sm70_sglang_push_chain.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Benchmark SGLang's push all-reduce on a verifier-shaped graph chain. + +Run this script from a pinned SGLang checkout and environment. It deliberately +uses the public SGLang JIT communicator so that the comparison includes its +production graph-input registration and two-epoch push protocol. +""" + +from __future__ import annotations + +import argparse +import json +import os +from pathlib import Path + +import torch +import torch.distributed as dist + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--numel", type=int, default=8 * 5120) + parser.add_argument("--count", type=int, default=128) + parser.add_argument("--warmup", type=int, default=40) + parser.add_argument("--iters", type=int, default=200) + parser.add_argument("--max-push-blocks", type=int, default=80) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--json-out", type=Path, required=True) + return parser.parse_args() + + +def _max_rank_float(value: float) -> float: + result = torch.tensor([value], dtype=torch.float64) + dist.all_reduce(result, op=dist.ReduceOp.MAX) + return float(result.item()) + + +def main() -> None: + args = _parse_args() + rank = int(os.environ["RANK"]) + local_rank = int(os.environ["LOCAL_RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + if world_size != 4: + raise ValueError(f"this verifier microbenchmark requires TP4, got {world_size}") + + torch.accelerator.set_device_index(local_rank) + device = torch.device("cuda", local_rank) + dist.init_process_group(backend="gloo") + + import sglang.srt.distributed.parallel_state as parallel_state + from sglang.jit_kernel.all_reduce import ( + AllReduceAlgo, + _jit_custom_all_reduce_push_module, + ) + from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import ( + CustomAllReduceV2, + ) + + parallel_state._WORLD = parallel_state.init_world_group( + ranks=list(range(world_size)), + local_rank=local_rank, + backend="nccl", + ) + cpu_group = parallel_state._WORLD.cpu_group + nccl_group = parallel_state._WORLD.device_group + if nccl_group is None: + raise RuntimeError("SGLang did not create the NCCL device group") + + bytes_per_input = args.numel * torch.float16.itemsize + communicator = CustomAllReduceV2( + cpu_group, + device, + max_pull_size=bytes_per_input, + max_push_size=bytes_per_input, + max_push_blocks=args.max_push_blocks, + ) + if communicator.disabled: + raise RuntimeError("SGLang push all-reduce is disabled on this topology") + communicator.override_algo = AllReduceAlgo.ONE_SHOT_PUSH + _jit_custom_all_reduce_push_module(torch.float16, world_size) + + generator = torch.Generator(device=device) + generator.manual_seed(args.seed + rank) + source = ( + torch.rand( + (args.count, args.numel), + dtype=torch.float32, + device=device, + generator=generator, + ) + * 0.03 + + 0.01 + ).half() + work = source.clone() + + gathered = [torch.empty_like(source) for _ in range(world_size)] + dist.all_gather(gathered, source, group=nccl_group) + reference = gathered[0].float() + for peer in gathered[1:]: + reference.add_(peer.float()) + reference = reference.half() + + graph = torch.cuda.CUDAGraph() + with communicator.capture(), torch.cuda.graph(graph): + for index in range(args.count): + communicator.custom_all_reduce(work[index]) + torch.accelerator.synchronize() + dist.barrier() + + work.copy_(source) + graph.replay() + torch.accelerator.synchronize() + exact = bool(torch.equal(work, reference)) + difference = (work.float() - reference.float()).abs() + mismatch_count = int((work != reference).sum().item()) + + for _ in range(args.warmup): + graph.replay() + torch.accelerator.synchronize() + dist.barrier() + + start = torch.Event(enable_timing=True) + end = torch.Event(enable_timing=True) + start.record() + for _ in range(args.iters): + graph.replay() + end.record() + end.synchronize() + local_ms = start.elapsed_time(end) / args.iters + max_rank_ms = _max_rank_float(local_ms) + + if rank == 0: + result = { + "implementation": "sglang_v100_one_shot_push", + "world_size": world_size, + "torch_version": torch.__version__, + "cuda_version": torch.version.cuda, + "numel": args.numel, + "bytes_per_input": bytes_per_input, + "allreduces_per_graph_replay": args.count, + "max_push_blocks": args.max_push_blocks, + "warmup_replays": args.warmup, + "timed_replays": args.iters, + "max_rank_ms_per_graph_replay": max_rank_ms, + "max_rank_us_per_allreduce": max_rank_ms * 1000.0 / args.count, + "exact_rank_order_fp32_reference": exact, + "mismatch_count": mismatch_count, + "max_diff": float(difference.max().item()), + "mean_diff": float(difference.mean().item()), + } + args.json_out.parent.mkdir(parents=True, exist_ok=True) + args.json_out.write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(result, indent=2)) + + communicator.close() + dist.barrier() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() 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_tp2_compact_verifier_logits.py b/benchmarks/benchmark_sm70_tp2_compact_verifier_logits.py index d7560f8e69..16456da153 100644 --- a/benchmarks/benchmark_sm70_tp2_compact_verifier_logits.py +++ b/benchmarks/benchmark_sm70_tp2_compact_verifier_logits.py @@ -1,10 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Measure TP2 full-vocab versus compact top-k verifier-logit transport. +"""Measure TP full-vocab versus compact top-k verifier-logit transport. -Run on two identical GPUs without loading a model: +Run on two or more identical GPUs without loading a model: - CUDA_VISIBLE_DEVICES=0,1 torchrun --standalone --nproc_per_node=2 \ + CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --standalone --nproc_per_node=4 \ benchmarks/benchmark_sm70_tp2_compact_verifier_logits.py \ --json-out bench_results/.../tp2_compact_verifier_logits.json @@ -93,8 +93,8 @@ def main() -> None: local_rank = int(os.environ["LOCAL_RANK"]) rank = int(os.environ["RANK"]) world_size = int(os.environ["WORLD_SIZE"]) - if world_size != 2: - raise ValueError(f"This benchmark is TP2-only, got world_size={world_size}.") + if world_size < 2: + raise ValueError(f"This benchmark requires TP>1, got world_size={world_size}.") if args.vocab_size % world_size: raise ValueError("--vocab-size must be divisible by TP size.") if args.top_k <= 0 or args.top_k > args.vocab_size // world_size: diff --git a/benchmarks/benchmark_sm70_tp4_fused_norm_collective.py b/benchmarks/benchmark_sm70_tp4_fused_norm_collective.py new file mode 100644 index 0000000000..38e460417e --- /dev/null +++ b/benchmarks/benchmark_sm70_tp4_fused_norm_collective.py @@ -0,0 +1,305 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Benchmark SM70 fused reduce-scatter + Gemma RMSNorm + all-gather. + +Launch with torchrun on a fully connected TP4 V100 group. The candidate is a +benchmark-only implementation of the communication topology in NVIDIA's NCCL +Device API fused RMSNorm example, using vLLM's existing IPC buffers and +cross-device signal substrate so it can be measured with NCCL 2.27. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import statistics +from collections.abc import Callable +from pathlib import Path +from typing import Any + +import torch +import torch.distributed as dist + +import vllm.kernels # noqa: F401 +from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce +from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--m", type=int, default=8000) + parser.add_argument("--n", type=int, default=5120) + parser.add_argument("--threads", type=int, nargs="+", default=[256, 512, 1024]) + parser.add_argument("--blocks", type=int, nargs="+", default=[80]) + parser.add_argument("--warmup", type=int, default=3) + parser.add_argument("--iters", type=int, default=5) + parser.add_argument("--trials", type=int, default=5) + parser.add_argument("--epsilon", type=float, default=1e-6) + parser.add_argument("--seed", type=int, default=20260823) + parser.add_argument("--candidate-cuda-graph", action="store_true") + parser.add_argument("--json-out", type=Path, required=True) + return parser.parse_args() + + +def _event_trials( + launch: Callable[[], None], warmup: int, iters: int, trials: int +) -> dict[str, Any]: + for _ in range(warmup): + launch() + torch.accelerator.synchronize() + + samples: list[float] = [] + for _ in range(trials): + dist.barrier() + start = torch.Event(enable_timing=True) + end = torch.Event(enable_timing=True) + start.record() + for _ in range(iters): + launch() + end.record() + end.synchronize() + samples.append(float(start.elapsed_time(end) / iters)) + return { + "median_ms": float(statistics.median(samples)), + "mean_ms": float(statistics.mean(samples)), + "min_ms": min(samples), + "max_ms": max(samples), + "samples_ms": samples, + } + + +def _digest(tensor: torch.Tensor) -> str: + raw = tensor.detach().contiguous().cpu().view(torch.uint8).numpy().tobytes() + return hashlib.sha256(raw).hexdigest() + + +def main() -> int: + args = _parse_args() + local_rank = int(os.environ["LOCAL_RANK"]) + torch.accelerator.set_device_index(local_rank) + dist.init_process_group(backend="nccl") + cpu_group = dist.new_group(backend="gloo") + rank = dist.get_rank() + world_size = dist.get_world_size() + if world_size != 4: + raise RuntimeError("This benchmark requires TP4.") + if torch.cuda.get_device_capability(local_rank) != (7, 0): + raise RuntimeError("This benchmark requires SM70/V100 GPUs.") + if args.n != 5120 or args.m % world_size: + raise ValueError("The candidate requires N=5120 and TP-divisible M.") + if any(threads not in (256, 512, 1024) for threads in args.threads): + raise ValueError("--threads choices must be 256, 512, or 1024.") + if any(blocks < 1 or blocks > 512 for blocks in args.blocks): + raise ValueError("--blocks choices must be in [1, 512].") + variants = [(threads, blocks) for threads in args.threads for blocks in args.blocks] + + device = torch.device("cuda", local_rank) + generator = torch.Generator(device=device).manual_seed(args.seed + rank) + partial = torch.randn( + (args.m, args.n), + generator=generator, + device=device, + dtype=torch.float16, + ).mul_(0.1) + residual = torch.randn( + (args.m, args.n), + generator=torch.Generator(device=device).manual_seed(args.seed + 10000), + device=device, + dtype=torch.float32, + ).mul_(0.1) + weight = torch.randn( + (args.n,), + generator=torch.Generator(device=device).manual_seed(args.seed + 20000), + device=device, + dtype=torch.float16, + ).mul_(0.01) + + baseline_reduced = torch.empty_like(partial) + baseline_normalized = torch.empty_like(partial) + baseline_residual = torch.empty_like(residual) + candidate_normalized = {variant: torch.empty_like(partial) for variant in variants} + candidate_residual = { + variant: torch.empty_like(residual, dtype=torch.float32) for variant in variants + } + + max_size = partial.numel() * partial.element_size() + 1 + pynccl = PyNcclCommunicator(group=cpu_group, device=device) + if pynccl.disabled: + raise RuntimeError("PyNccl is unavailable for the production baseline.") + communicator = CustomAllreduce( + cpu_group, + device, + max_size=max_size, + long_prefill_fusion_enabled=True, + ) + if communicator.disabled: + raise RuntimeError("Custom all-reduce is unavailable for this TP4 group.") + + def norm_out( + output: torch.Tensor, + residual_output: torch.Tensor, + input: torch.Tensor, + input_residual: torch.Tensor, + ) -> None: + torch.ops._C.sm70_gemma_long_prefill_fused_add_rms_norm( + output, + residual_output, + input, + input_residual, + weight, + args.epsilon, + ) + + def baseline() -> None: + pynccl.all_reduce(partial, baseline_reduced) + norm_out( + baseline_normalized, + baseline_residual, + baseline_reduced, + residual, + ) + + def baseline_all_reduce() -> None: + pynccl.all_reduce(partial, baseline_reduced) + + def baseline_norm() -> None: + norm_out( + baseline_normalized, + baseline_residual, + baseline_reduced, + residual, + ) + + def d2d_copy() -> None: + baseline_reduced.copy_(partial) + + def make_candidate(threads: int, blocks: int) -> Callable[[], None]: + def launch() -> None: + os.environ["VLLM_SM70_TP4_LONG_FUSED_NORM_THREADS"] = str(threads) + os.environ["VLLM_SM70_TP4_LONG_FUSED_NORM_BLOCKS"] = str(blocks) + communicator.sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather( + partial, + residual, + weight, + args.epsilon, + normalized_out=candidate_normalized[(threads, blocks)], + residual_out=candidate_residual[(threads, blocks)], + ) + + return launch + + baseline_timing = _event_trials(baseline, args.warmup, args.iters, args.trials) + baseline_ar_timing = _event_trials( + baseline_all_reduce, args.warmup, args.iters, args.trials + ) + baseline_all_reduce() + baseline_norm_timing = _event_trials( + baseline_norm, args.warmup, args.iters, args.trials + ) + copy_timing = _event_trials(d2d_copy, args.warmup, args.iters, args.trials) + candidate_launches = {variant: make_candidate(*variant) for variant in variants} + candidate_graphs: list[torch.cuda.CUDAGraph] = [] + if args.candidate_cuda_graph: + graph_launches: dict[tuple[int, int], Callable[[], None]] = {} + for variant, launch in candidate_launches.items(): + launch() + torch.accelerator.synchronize() + dist.barrier() + graph = torch.cuda.CUDAGraph() + with communicator.capture(), torch.cuda.graph(graph): + launch() + candidate_graphs.append(graph) + graph_launches[variant] = graph.replay + candidate_launches = graph_launches + candidate_timings = { + variant: _event_trials(launch, args.warmup, args.iters, args.trials) + for variant, launch in candidate_launches.items() + } + + baseline() + for launch in candidate_launches.values(): + launch() + torch.accelerator.synchronize() + local_rows = args.m // world_size + residual_reference = baseline_residual.narrow(0, rank * local_rows, local_rows) + candidates: dict[str, Any] = {} + for variant in variants: + threads, blocks = variant + normalized_difference = ( + baseline_normalized.float() - candidate_normalized[variant].float() + ).abs() + candidate_residual_shard = candidate_residual[variant].narrow( + 0, rank * local_rows, local_rows + ) + residual_difference = (residual_reference - candidate_residual_shard).abs() + timing = candidate_timings[variant] + candidates[f"t{threads}_b{blocks}"] = { + "threads": threads, + "blocks": blocks, + "timing": timing, + "speedup": baseline_timing["median_ms"] / timing["median_ms"], + "normalized_max_abs_diff": float(normalized_difference.max().item()), + "normalized_mean_abs_diff": float(normalized_difference.mean().item()), + "normalized_equal_elements": int( + (baseline_normalized == candidate_normalized[variant]).sum().item() + ), + "normalized_elements": baseline_normalized.numel(), + "residual_max_abs_diff": float(residual_difference.max().item()), + "residual_mean_abs_diff": float(residual_difference.mean().item()), + "normalized_sha256": _digest(candidate_normalized[variant]), + "residual_shard_sha256": _digest(candidate_residual_shard), + } + + local_payload = { + "rank": rank, + "baseline": baseline_timing, + "baseline_all_reduce": baseline_ar_timing, + "baseline_norm": baseline_norm_timing, + "d2d_copy": copy_timing, + "baseline_normalized_sha256": _digest(baseline_normalized), + "baseline_residual_shard_sha256": _digest(residual_reference), + "candidates": candidates, + } + gathered: list[dict[str, Any] | None] = [None] * world_size + dist.all_gather_object(gathered, local_payload) + if rank == 0: + payload = { + "environment": { + "world_size": world_size, + "device": torch.cuda.get_device_name(local_rank), + "torch": torch.__version__, + "torch_cuda": torch.version.cuda, + "m": args.m, + "n": args.n, + "dtype": "float16", + "residual_dtype": "float32", + "threads": args.threads, + "blocks": args.blocks, + "candidate_cuda_graph": args.candidate_cuda_graph, + "warmup": args.warmup, + "iters": args.iters, + "trials": args.trials, + "epsilon": args.epsilon, + "seed": args.seed, + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "nccl_algo": os.environ.get("NCCL_ALGO"), + "nccl_proto": os.environ.get("NCCL_PROTO"), + }, + "ranks": gathered, + } + args.json_out.parent.mkdir(parents=True, exist_ok=True) + args.json_out.write_text(json.dumps(payload, indent=2) + "\n") + print(json.dumps(payload, indent=2), flush=True) + + communicator.close() + pynccl.destroy() + dist.barrier() + dist.destroy_process_group(cpu_group) + dist.destroy_process_group() + return 0 + + +if __name__ == "__main__": + raise SystemExit(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_dsv4_quality_results.py b/benchmarks/compare_dsv4_quality_results.py new file mode 100644 index 0000000000..073d8276de --- /dev/null +++ b/benchmarks/compare_dsv4_quality_results.py @@ -0,0 +1,425 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Compare matched no-speculation and DSpark quality artifacts.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + + +def _load_json(path: Path) -> dict[str, Any]: + return json.loads(path.read_text(encoding="utf-8")) + + +def _load_jsonl(path: Path) -> list[dict[str, Any]]: + return [ + json.loads(line) + for line in path.read_text(encoding="utf-8").splitlines() + if line.strip() + ] + + +def _manifest_hashes(value: Any, prefix: str = "") -> dict[str, str]: + if not isinstance(value, dict): + return {} + if isinstance(value.get("sha256"), str): + return {prefix: value["sha256"]} + hashes: dict[str, str] = {} + for key, child in value.items(): + child_prefix = f"{prefix}.{key}" if prefix else key + hashes.update(_manifest_hashes(child, child_prefix)) + return hashes + + +def _compare_humaneval( + reference: dict[str, Any], candidate: dict[str, Any] +) -> dict[str, Any]: + reference_rows = {row["task_id"]: row for row in reference["human_eval"]["records"]} + candidate_rows = {row["task_id"]: row for row in candidate["human_eval"]["records"]} + common = sorted(reference_rows.keys() & candidate_rows.keys()) + regressions = [ + task_id + for task_id in common + if reference_rows[task_id]["passed"] and not candidate_rows[task_id]["passed"] + ] + improvements = [ + task_id + for task_id in common + if not reference_rows[task_id]["passed"] and candidate_rows[task_id]["passed"] + ] + exact_responses = sum( + reference_rows[task_id].get("response") + == candidate_rows[task_id].get("response") + for task_id in common + ) + reference_passed = int(reference["human_eval"]["passed"]) + candidate_passed = int(candidate["human_eval"]["passed"]) + missing = sorted(reference_rows.keys() ^ candidate_rows.keys()) + return { + "reference_passed": reference_passed, + "candidate_passed": candidate_passed, + "samples": len(common), + "exact_response_matches": exact_responses, + "regressions": regressions, + "improvements": improvements, + "missing_or_extra_tasks": missing, + # Individual flips remain visible for diagnosis, but are not token- + # identity gates. Acceptance is based on matched inputs and aggregate + # task quality. + "passed": not missing and candidate_passed >= reference_passed, + } + + +def _compare_longbench( + reference: dict[str, Any], + candidate: dict[str, Any], + tolerance: float, +) -> dict[str, Any]: + reference_sets = reference["longbench"]["datasets"] + candidate_sets = candidate["longbench"]["datasets"] + missing = sorted(reference_sets.keys() ^ candidate_sets.keys()) + results: dict[str, Any] = {} + for dataset in sorted(reference_sets.keys() & candidate_sets.keys()): + reference_rows = { + int(row["source_index"]): row for row in reference_sets[dataset]["records"] + } + candidate_rows = { + int(row["source_index"]): row for row in candidate_sets[dataset]["records"] + } + common = sorted(reference_rows.keys() & candidate_rows.keys()) + row_mismatch = sorted(reference_rows.keys() ^ candidate_rows.keys()) + input_mismatch = [ + index + for index in common + if reference_rows[index].get("answers") + != candidate_rows[index].get("answers") + ] + regressions = [ + { + "source_index": index, + "reference_score": float(reference_rows[index]["score"]), + "candidate_score": float(candidate_rows[index]["score"]), + } + for index in common + if float(candidate_rows[index]["score"]) + tolerance + < float(reference_rows[index]["score"]) + ] + exact_predictions = sum( + reference_rows[index].get("prediction") + == candidate_rows[index].get("prediction") + for index in common + ) + reference_score = float(reference_sets[dataset]["score"]) + candidate_score = float(candidate_sets[dataset]["score"]) + results[dataset] = { + "reference_score": reference_score, + "candidate_score": candidate_score, + "samples": len(common), + "exact_prediction_matches": exact_predictions, + "regressions": regressions, + "missing_or_extra_rows": row_mismatch, + "input_mismatches": input_mismatch, + "passed": not row_mismatch + and not input_mismatch + and candidate_score + tolerance >= reference_score, + } + reference_average = float(reference["longbench"]["average_score"]) + candidate_average = float(candidate["longbench"]["average_score"]) + return { + "reference_average": reference_average, + "candidate_average": candidate_average, + "score_tolerance": tolerance, + "datasets": results, + "missing_or_extra_datasets": missing, + "passed": not missing + and all(result["passed"] for result in results.values()) + and candidate_average + tolerance >= reference_average, + } + + +def _compare_api_quality( + reference: dict[str, Any], + candidate: dict[str, Any], + tolerance: float, +) -> dict[str, Any]: + reference_hashes = _manifest_hashes(reference.get("input_manifest", {})) + candidate_hashes = _manifest_hashes(candidate.get("input_manifest", {})) + manifest_equal = bool(reference_hashes) and reference_hashes == candidate_hashes + reference_contract = reference.get("evaluation_contract") + candidate_contract = candidate.get("evaluation_contract") + contract_equal = ( + bool(reference_contract) and reference_contract == candidate_contract + ) + humaneval = _compare_humaneval(reference, candidate) + longbench = _compare_longbench(reference, candidate, tolerance) + return { + "input_hashes_equal": manifest_equal, + "reference_input_hashes": reference_hashes, + "candidate_input_hashes": candidate_hashes, + "evaluation_contract_equal": contract_equal, + "reference_evaluation_contract": reference_contract, + "candidate_evaluation_contract": candidate_contract, + "humaneval": humaneval, + "longbench": longbench, + "passed": manifest_equal + and contract_equal + and humaneval["passed"] + and longbench["passed"], + } + + +def _compare_gsm8k( + reference: dict[str, Any], candidate: dict[str, Any] +) -> dict[str, Any]: + reference_hashes = _manifest_hashes(reference.get("input_manifest", {})) + candidate_hashes = _manifest_hashes(candidate.get("input_manifest", {})) + manifest_equal = bool(reference_hashes) and reference_hashes == candidate_hashes + contract_keys = ( + "questions", + "few_shot", + "temperature", + "top_p", + "seed", + "max_tokens", + "strictly_sequential", + "train_selection", + "test_selection", + "prompt_format", + "answer_normalization", + "stop_sequences", + ) + reference_contract = { + key: reference.get("contract", {}).get(key) for key in contract_keys + } + candidate_contract = { + key: candidate.get("contract", {}).get(key) for key in contract_keys + } + contract_equal = reference_contract == candidate_contract + reference_rows = {int(row["index"]): row for row in reference["rows"]} + candidate_rows = {int(row["index"]): row for row in candidate["rows"]} + common = sorted(reference_rows.keys() & candidate_rows.keys()) + missing = sorted(reference_rows.keys() ^ candidate_rows.keys()) + expected_samples = reference_contract["questions"] + row_count_matches_contract = ( + isinstance(expected_samples, int) + and len(reference["rows"]) == expected_samples + and len(candidate["rows"]) == expected_samples + and len(reference_rows) == expected_samples + and len(candidate_rows) == expected_samples + ) + + def row_is_consistent(row: dict[str, Any]) -> bool: + expected = row.get("expected") + predicted = row.get("predicted") + return bool( + expected is not None + and bool(row.get("correct")) + == (predicted is not None and predicted == expected) + and bool(row.get("invalid")) == (predicted is None) + ) + + rows_self_consistent = all( + row_is_consistent(row) for row in [*reference["rows"], *candidate["rows"]] + ) + summaries_self_consistent = all( + int(result["correct"]) + == sum(bool(row.get("correct")) for row in result["rows"]) + and int(result["invalid"]) + == sum(bool(row.get("invalid")) for row in result["rows"]) + for result in (reference, candidate) + ) + input_mismatch = [ + index + for index in common + if ( + reference_rows[index].get("question"), + reference_rows[index].get("expected"), + ) + != ( + candidate_rows[index].get("question"), + candidate_rows[index].get("expected"), + ) + ] + regressions = [ + index + for index in common + if reference_rows[index]["correct"] and not candidate_rows[index]["correct"] + ] + improvements = [ + index + for index in common + if not reference_rows[index]["correct"] and candidate_rows[index]["correct"] + ] + exact_predictions = sum( + reference_rows[index].get("predicted") == candidate_rows[index].get("predicted") + for index in common + ) + return { + "input_hashes_equal": manifest_equal, + "reference_input_hashes": reference_hashes, + "candidate_input_hashes": candidate_hashes, + "evaluation_contract_equal": contract_equal, + "reference_evaluation_contract": reference_contract, + "candidate_evaluation_contract": candidate_contract, + "row_count_matches_contract": row_count_matches_contract, + "rows_self_consistent": rows_self_consistent, + "summaries_self_consistent": summaries_self_consistent, + "reference_correct": int(reference["correct"]), + "candidate_correct": int(candidate["correct"]), + "correct_delta": int(candidate["correct"]) - int(reference["correct"]), + "reference_invalid": int(reference["invalid"]), + "candidate_invalid": int(candidate["invalid"]), + "samples": len(common), + "exact_prediction_matches": exact_predictions, + "regressions": regressions, + "improvements": improvements, + "missing_or_extra_rows": missing, + "input_mismatches": input_mismatch, + # Report directional flips without requiring greedy/token identity. + # Matched aggregate correctness and validity are the quality gate. + "passed": manifest_equal + and contract_equal + and row_count_matches_contract + and rows_self_consistent + and summaries_self_consistent + and not missing + and not input_mismatch + and int(candidate["correct"]) >= int(reference["correct"]) + and int(candidate["invalid"]) <= int(reference["invalid"]), + } + + +def _compare_needle( + reference: list[dict[str, Any]], candidate: list[dict[str, Any]] +) -> dict[str, Any]: + reference_rows = {row["sample_id"]: row for row in reference} + candidate_rows = {row["sample_id"]: row for row in candidate} + common = sorted(reference_rows.keys() & candidate_rows.keys()) + missing = sorted(reference_rows.keys() ^ candidate_rows.keys()) + input_mismatch = [ + sample_id + for sample_id in common + if ( + reference_rows[sample_id].get("target_tokens"), + reference_rows[sample_id].get("depth"), + reference_rows[sample_id].get("code"), + ) + != ( + candidate_rows[sample_id].get("target_tokens"), + candidate_rows[sample_id].get("depth"), + candidate_rows[sample_id].get("code"), + ) + ] + final_regressions = [ + sample_id + for sample_id in common + if reference_rows[sample_id]["hit"] and not candidate_rows[sample_id]["hit"] + ] + anywhere_regressions = [ + sample_id + for sample_id in common + if reference_rows[sample_id]["hit_anywhere"] + and not candidate_rows[sample_id]["hit_anywhere"] + ] + reference_hits = sum(bool(row["hit"]) for row in reference_rows.values()) + candidate_hits = sum(bool(row["hit"]) for row in candidate_rows.values()) + reference_anywhere = sum( + bool(row["hit_anywhere"]) for row in reference_rows.values() + ) + candidate_anywhere = sum( + bool(row["hit_anywhere"]) for row in candidate_rows.values() + ) + return { + "reference_final_hits": reference_hits, + "candidate_final_hits": candidate_hits, + "reference_anywhere_hits": reference_anywhere, + "candidate_anywhere_hits": candidate_anywhere, + "samples": len(common), + "final_hit_regressions": final_regressions, + "anywhere_hit_regressions": anywhere_regressions, + "missing_or_extra_rows": missing, + "input_mismatches": input_mismatch, + "passed": not missing + and not input_mismatch + and candidate_hits >= reference_hits + and candidate_anywhere >= reference_anywhere, + } + + +def _paired_paths( + parser: argparse.ArgumentParser, + reference: Path | None, + candidate: Path | None, + label: str, +) -> bool: + if (reference is None) != (candidate is None): + parser.error(f"--reference-{label} and --candidate-{label} must be paired") + return reference is not None + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--reference-api", type=Path) + parser.add_argument("--candidate-api", type=Path) + parser.add_argument("--reference-gsm8k", type=Path) + parser.add_argument("--candidate-gsm8k", type=Path) + parser.add_argument("--reference-needle", type=Path) + parser.add_argument("--candidate-needle", type=Path) + parser.add_argument("--longbench-score-tolerance", type=float, default=0.0) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + if args.longbench_score_tolerance < 0: + parser.error("--longbench-score-tolerance must be non-negative") + + has_api = _paired_paths(parser, args.reference_api, args.candidate_api, "api") + has_gsm8k = _paired_paths( + parser, args.reference_gsm8k, args.candidate_gsm8k, "gsm8k" + ) + has_needle = _paired_paths( + parser, args.reference_needle, args.candidate_needle, "needle" + ) + if not (has_api or has_gsm8k or has_needle): + parser.error("at least one paired quality artifact must be provided") + + result: dict[str, Any] = { + "contract": { + key: str(value) if isinstance(value, Path) else value + for key, value in vars(args).items() + }, + } + gates: list[bool] = [] + if has_api: + assert args.reference_api is not None and args.candidate_api is not None + result["api_quality"] = _compare_api_quality( + _load_json(args.reference_api), + _load_json(args.candidate_api), + args.longbench_score_tolerance, + ) + gates.append(result["api_quality"]["passed"]) + if has_gsm8k: + assert args.reference_gsm8k is not None and args.candidate_gsm8k is not None + result["gsm8k"] = _compare_gsm8k( + _load_json(args.reference_gsm8k), _load_json(args.candidate_gsm8k) + ) + gates.append(result["gsm8k"]["passed"]) + if has_needle: + assert args.reference_needle is not None and args.candidate_needle is not None + result["needle"] = _compare_needle( + _load_jsonl(args.reference_needle), _load_jsonl(args.candidate_needle) + ) + gates.append(result["needle"]["passed"]) + result["passed"] = all(gates) + + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text( + json.dumps(result, ensure_ascii=False, indent=2) + "\n", encoding="utf-8" + ) + print(json.dumps(result, ensure_ascii=False), flush=True) + return 0 if result["passed"] else 1 + + +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_device_communicators.py b/benchmarks/kernels/benchmark_device_communicators.py index 24e22023b9..4afdf07723 100644 --- a/benchmarks/kernels/benchmark_device_communicators.py +++ b/benchmarks/kernels/benchmark_device_communicators.py @@ -49,9 +49,10 @@ # Default sequence lengths to benchmark DEFAULT_SEQUENCE_LENGTHS = [16, 64, 128, 512, 1024, 2048, 4096, 8192] -# Fixed hidden size and dtype for all benchmarks -HIDDEN_SIZE = 8192 -BENCHMARK_DTYPE = torch.bfloat16 +# Defaults preserve the original benchmark contract. CLI overrides make it +# possible to replay an exact model communication shape. +DEFAULT_HIDDEN_SIZE = 8192 +DEFAULT_BENCHMARK_DTYPE = torch.bfloat16 # CUDA graph settings CUDA_GRAPH_CAPTURE_CYCLES = 10 @@ -67,16 +68,20 @@ def __init__( device: torch.device, cpu_group: ProcessGroup, sequence_lengths: list[int], + hidden_size: int, + dtype: torch.dtype, ): self.rank = rank self.world_size = world_size self.device = device self.cpu_group = cpu_group + self.hidden_size = hidden_size + self.dtype = dtype # Calculate max_size_override based on largest sequence length max_seq_len = max(sequence_lengths) - max_tensor_elements = max_seq_len * HIDDEN_SIZE - self.max_size_override = max_tensor_elements * BENCHMARK_DTYPE.itemsize + 1 + max_tensor_elements = max_seq_len * self.hidden_size + self.max_size_override = max_tensor_elements * self.dtype.itemsize + 1 # Initialize communicators self.custom_allreduce = None @@ -337,7 +342,10 @@ def benchmark_allreduce_single( try: # Create test tensor (2D: sequence_length x hidden_size) tensor = torch.randn( - sequence_length, HIDDEN_SIZE, dtype=BENCHMARK_DTYPE, device=self.device + sequence_length, + self.hidden_size, + dtype=self.dtype, + device=self.device, ) if not should_use_fn(tensor): return None @@ -405,16 +413,17 @@ def _calculate_speedup_info(comm_results: dict[str, float]) -> str: def print_results( - results: dict[str, dict[str, float]], sequence_lengths: list[int], world_size: int + results: dict[str, dict[str, float]], + sequence_lengths: list[int], + world_size: int, + hidden_size: int, + dtype: torch.dtype, ): """Print benchmark results in a formatted table.""" print(f"\n{'=' * 130}") print("Device Communicator Benchmark Results") - print( - f"World Size: {world_size}, Data Type: {BENCHMARK_DTYPE}, " - f"Hidden Size: {HIDDEN_SIZE}" - ) + print(f"World Size: {world_size}, Data Type: {dtype}, Hidden Size: {hidden_size}") print(f"{'=' * 130}") # Get all communicator names @@ -436,15 +445,15 @@ def print_results( for seq_len in sequence_lengths: if seq_len in results: # Calculate tensor size in elements and bytes - tensor_elements = seq_len * HIDDEN_SIZE - tensor_bytes = tensor_elements * BENCHMARK_DTYPE.itemsize + tensor_elements = seq_len * hidden_size + tensor_bytes = tensor_elements * dtype.itemsize # Format tensor size (MB) tensor_size_mb = tensor_bytes / (1024 * 1024) tensor_size_str = f"{tensor_size_mb:.2f} MB" # Format tensor shape - tensor_shape = f"({seq_len}, {HIDDEN_SIZE})" + tensor_shape = f"({seq_len}, {hidden_size})" row = f"{tensor_shape:<20}{tensor_size_str:<15}" for comm in all_comms: @@ -483,9 +492,24 @@ def main(): "--num-trials", type=int, default=50, help="Number of benchmark trials" ) + parser.add_argument( + "--hidden-size", + type=int, + default=DEFAULT_HIDDEN_SIZE, + help="Hidden dimension of the all-reduce tensor.", + ) + + parser.add_argument( + "--dtype", + choices=("float16", "bfloat16"), + default="bfloat16", + help="All-reduce tensor dtype.", + ) + parser.add_argument("--output-json", type=str, help="Output results to JSON file") args = parser.parse_args() + benchmark_dtype = getattr(torch, args.dtype) # Initialize distributed if not dist.is_initialized(): @@ -506,7 +530,13 @@ def main(): # Initialize benchmark benchmark = CommunicatorBenchmark( - rank, world_size, device, cpu_group, args.sequence_lengths + rank, + world_size, + device, + cpu_group, + args.sequence_lengths, + args.hidden_size, + benchmark_dtype, ) # Run benchmarks @@ -518,7 +548,7 @@ def main(): "Benchmarking sequence length: %s (tensor shape: %s x %s)", seq_len, seq_len, - HIDDEN_SIZE, + args.hidden_size, ) results = benchmark.benchmark_allreduce( @@ -534,7 +564,13 @@ def main(): # Print results (only rank 0) if rank == 0: - print_results(all_results, args.sequence_lengths, world_size) + print_results( + all_results, + args.sequence_lengths, + world_size, + args.hidden_size, + benchmark_dtype, + ) # Save to JSON if requested if args.output_json: @@ -548,8 +584,8 @@ def main(): output_data = { "world_size": world_size, - "dtype": str(BENCHMARK_DTYPE), - "hidden_size": HIDDEN_SIZE, + "dtype": str(benchmark_dtype), + "hidden_size": args.hidden_size, "sequence_lengths": args.sequence_lengths, "num_warmup": args.num_warmup, "num_trials": args.num_trials, 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_fp8_kv.py b/benchmarks/kernels/benchmark_flash_v100_fp8_kv.py index 8de8b616c3..34d22ea3b5 100644 --- a/benchmarks/kernels/benchmark_flash_v100_fp8_kv.py +++ b/benchmarks/kernels/benchmark_flash_v100_fp8_kv.py @@ -43,6 +43,7 @@ "fp8_b1616_bridge", ) DECODE_VARIANTS = ( + "fp8_b16_scalar", "fp16_b784_xqa", "fp16_b784_scalar", "fp16_b1568_xqa", 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/kernels/benchmark_sm70_nvfp4_qpn2.py b/benchmarks/kernels/benchmark_sm70_nvfp4_qpn2.py new file mode 100644 index 0000000000..4393434cc6 --- /dev/null +++ b/benchmarks/kernels/benchmark_sm70_nvfp4_qpn2.py @@ -0,0 +1,499 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Race QPN2 against the current SM70 NVFP4 path on real Qwen3.8 weights. + +This is a deliberately narrow verifier microbenchmark. It loads the TP-local +gate/up and down projection shards from one native NVFP4 layer, prepares both +the current TurboMind layout and the QPN2 fragment layout, and measures an M=8 +CUDA-graph replay. Weight loading and preparation are outside the timed +region. + +QPN2 is compiled from an explicitly supplied source file so this benchmark can +evaluate a pinned external implementation before it is admitted to the vLLM +build. The result records the source digest and numerical delta against the +current production operator. +""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib.util +import json +import statistics +import sys +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch +from safetensors import safe_open +from torch.utils.cpp_extension import load + + +@dataclass(frozen=True) +class Projection: + name: str + packed: torch.Tensor + scales: torch.Tensor + inverse_global_scale: float + gated_silu: bool + calls_per_round: int + + +QPN2_CONFIGS = { + # (K, N): (split K, independent accumulator chains) + (5120, 8704): (8, 2), + (4352, 5120): (16, 2), +} + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as source: + for chunk in iter(lambda: source.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _load_projection_shards( + model: Path, + layer_index: int, + tp_rank: int, + tp_size: int, +) -> tuple[Projection, Projection]: + path = model / "model.safetensors" + prefix = f"model.language_model.layers.{layer_index}.mlp" + intermediate_size = 17408 + hidden_size = 5120 + if intermediate_size % tp_size or hidden_size % tp_size: + raise ValueError("model dimensions must divide the tensor-parallel size") + intermediate_per_rank = intermediate_size // tp_size + packed_k_per_rank = intermediate_per_rank // 2 + scale_k_per_rank = intermediate_per_rank // 16 + n_start = tp_rank * intermediate_per_rank + n_end = n_start + intermediate_per_rank + packed_k_start = tp_rank * packed_k_per_rank + packed_k_end = packed_k_start + packed_k_per_rank + scale_k_start = tp_rank * scale_k_per_rank + scale_k_end = scale_k_start + scale_k_per_rank + + with safe_open(path, framework="pt", device="cpu") as tensors: + gate_packed = tensors.get_slice(f"{prefix}.gate_proj.weight_packed")[ + n_start:n_end, : + ] + up_packed = tensors.get_slice(f"{prefix}.up_proj.weight_packed")[ + n_start:n_end, : + ] + gate_scales = tensors.get_slice(f"{prefix}.gate_proj.weight_scale")[ + n_start:n_end, : + ] + up_scales = tensors.get_slice(f"{prefix}.up_proj.weight_scale")[ + n_start:n_end, : + ] + gate_global = tensors.get_tensor(f"{prefix}.gate_proj.weight_global_scale") + up_global = tensors.get_tensor(f"{prefix}.up_proj.weight_global_scale") + + down_packed = tensors.get_slice(f"{prefix}.down_proj.weight_packed")[ + :, packed_k_start:packed_k_end + ] + down_scales = tensors.get_slice(f"{prefix}.down_proj.weight_scale")[ + :, scale_k_start:scale_k_end + ] + down_global = tensors.get_tensor(f"{prefix}.down_proj.weight_global_scale") + + gate_up = Projection( + name="gate_up_proj", + packed=torch.cat((gate_packed, up_packed), dim=0).contiguous(), + scales=torch.cat((gate_scales, up_scales), dim=0).contiguous(), + inverse_global_scale=float( + 1.0 / torch.cat((gate_global, up_global)).max().float() + ), + gated_silu=True, + calls_per_round=56, + ) + down = Projection( + name="down_proj", + packed=down_packed.contiguous(), + scales=down_scales.contiguous(), + inverse_global_scale=float(1.0 / down_global.max().float()), + gated_silu=False, + calls_per_round=56, + ) + return gate_up, down + + +def _unpack_nibbles(weight_packed: torch.Tensor) -> torch.Tensor: + """Convert checkpoint [N,K/2] bytes to TurboMind's [K,N] codes.""" + return ( + torch.stack((weight_packed & 15, weight_packed >> 4), dim=-1) + .flatten(start_dim=-2) + .t() + .contiguous() + ) + + +def _qpn2_prepack( + codes: torch.Tensor, scales: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """Byte-preserving permutation into v100-skinny's QPN2 layout.""" + n, k2 = codes.shape + k = k2 * 2 + if n % 32 or k % 64: + raise ValueError(f"QPN2 requires N%32=0 and K%64=0, got N={n}, K={k}") + device = codes.device + tiles, groups = n // 32, k // 16 + lane = torch.arange(32, device=device) + column = ((lane >> 2) & 3) * 8 + (lane & 3) + ((lane & 16) > 0).long() * 4 + k_order = torch.tensor( + [0, 2, 4, 6, 1, 3, 5, 7, 8, 10, 12, 14, 9, 11, 13, 15], + device=device, + ) + nibbles = torch.stack((codes & 0xF, codes >> 4), dim=-1).view(n, k) + group = torch.arange(groups, device=device) + k_index = group.view(groups, 1) * 16 + k_order.view(1, 16) + q_codes = torch.empty((tiles, groups, 32, 8), dtype=torch.uint8, device=device) + q_scales = torch.empty((tiles, groups, 32), dtype=torch.uint8, device=device) + # Bound the int64 gather intermediates for the gate/up matrix. + tile_chunk = max(1, 36864 // groups) + for tile_start in range(0, tiles, tile_chunk): + tile_end = min(tile_start + tile_chunk, tiles) + tile_count = tile_end - tile_start + n_column = torch.arange(tile_start, tile_end, device=device).view( + tile_count, 1 + ) * 32 + column.view(1, 32) + nibble_block = nibbles[ + n_column.view(tile_count, 1, 32, 1).expand(tile_count, groups, 32, 16), + k_index.view(1, groups, 1, 16).expand(tile_count, groups, 32, 16), + ] + q_codes[tile_start:tile_end] = nibble_block[..., 0::2] | ( + nibble_block[..., 1::2] << 4 + ) + q_scales[tile_start:tile_end] = scales[ + n_column.view(tile_count, 1, 32).expand(tile_count, groups, 32), + group.view(1, groups, 1).expand(tile_count, groups, 32), + ] + return q_codes.view(-1).contiguous(), q_scales.view(-1).contiguous() + + +def _capture(call: Callable[[], Any]) -> tuple[torch.cuda.CUDAGraph, Any]: + for _ in range(10): + call() + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = call() + torch.accelerator.synchronize() + return graph, output + + +def _time_graph( + graph: torch.cuda.CUDAGraph, warmup: int, iterations: int, trials: int +) -> dict[str, float | list[float]]: + for _ in range(warmup): + graph.replay() + torch.accelerator.synchronize() + samples = [] + for _ in range(trials): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iterations): + graph.replay() + end.record() + end.synchronize() + samples.append(float(start.elapsed_time(end)) * 1000.0 / iterations) + return { + "samples_us": samples, + "median_us": statistics.median(samples), + "min_us": min(samples), + "max_us": max(samples), + } + + +def _quality(actual: torch.Tensor, expected: torch.Tensor) -> dict[str, object]: + actual_f = actual.float() + expected_f = expected.float() + delta = actual_f - expected_f + denominator = torch.linalg.vector_norm(expected_f).clamp_min(1e-12) + return { + "finite": bool(torch.isfinite(actual).all()), + "exact": bool(torch.equal(actual, expected)), + "different_fraction": float((actual != expected).float().mean()), + "max_abs": float(delta.abs().max()), + "mean_abs": float(delta.abs().mean()), + "relative_l2": float(torch.linalg.vector_norm(delta) / denominator), + "cosine": float( + torch.nn.functional.cosine_similarity(actual_f, expected_f, dim=1) + .mean() + .item() + ), + } + + +def _fp32_reference( + projection: Projection, x: torch.Tensor, device: torch.device +) -> torch.Tensor: + """Evaluate the checkpoint's native E2M1 codes with FP32 accumulation.""" + packed = projection.packed.to(device) + nibbles = torch.stack((packed & 0xF, packed >> 4), dim=-1).flatten(start_dim=-2) + magnitudes = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], + dtype=torch.float32, + device=device, + ) + weights = magnitudes[(nibbles & 7).long()] + weights = torch.where((nibbles & 8) != 0, -weights, weights) + weights.mul_( + projection.scales.float() + .to(device) + .repeat_interleave(16, dim=1) + .mul_(projection.inverse_global_scale) + ) + raw = x.float().matmul(weights.t()) + if not projection.gated_silu: + return raw + gate, up = raw.chunk(2, dim=1) + return torch.nn.functional.silu(gate).mul_(up) + + +def _run_projection( + projection: Projection, + extension: Any | None, + production: bool, + m: int, + device: torch.device, + warmup: int, + iterations: int, + trials: int, +) -> dict[str, object]: + from vllm import _sm70_ops as sm70_ops + + packed = projection.packed.to(device) + scale_codes = projection.scales.view(torch.uint8).to(device) + qweight = _unpack_nibbles(projection.packed).to(device) + effective_scales = ( + projection.scales.t() + .float() + .mul(projection.inverse_global_scale) + .half() + .contiguous() + .to(device) + ) + k, n = qweight.shape + split_k, accumulator_chains = QPN2_CONFIGS[(k, n)] + x = torch.randn((m, k), dtype=torch.float16, device=device) * 0.1 + + tm_weight, tm_scales, meta = sm70_ops.nvfp4_sm70_prepare( + qweight, effective_scales, 16, False + ) + k_ld, q_ld = (int(value.item()) for value in meta[:2]) + tm_raw = torch.empty((m, n), dtype=torch.float16, device=device) + final_n = n // 2 if projection.gated_silu else n + tm_final = ( + torch.empty((m, final_n), dtype=torch.float16, device=device) + if projection.gated_silu + else tm_raw + ) + + def run_tm() -> torch.Tensor: + sm70_ops.nvfp4_gemm_sm70_out( + tm_raw, x, tm_weight, tm_scales, 16, k_ld, q_ld, False + ) + if projection.gated_silu: + torch.ops._C.silu_and_mul(tm_final, tm_raw) + return tm_final + + qpn_final = torch.empty((m, final_n), dtype=torch.float16, device=device) + if production: + qpn_codes, qpn_scales = sm70_ops.nvfp4_qpn2_prepare_sm70( + packed, projection.scales.to(device) + ) + + def run_qpn2() -> torch.Tensor: + if projection.gated_silu: + sm70_ops.nvfp4_qpn2_gated_sm70_out( + qpn_final, + x, + qpn_codes, + qpn_scales, + projection.inverse_global_scale, + split_k, + accumulator_chains, + ) + else: + sm70_ops.nvfp4_qpn2_gemm_sm70_out( + qpn_final, + x, + qpn_codes, + qpn_scales, + projection.inverse_global_scale, + split_k, + accumulator_chains, + ) + return qpn_final + + else: + if extension is None: + raise AssertionError("external QPN2 extension was not loaded") + qpn_codes, qpn_scales = _qpn2_prepack(packed, scale_codes) + + def run_qpn2() -> torch.Tensor: + qpn_raw = extension.gemm_qpn2( + x, + qpn_codes, + qpn_scales, + projection.inverse_global_scale, + n, + split_k, + accumulator_chains, + ) + if projection.gated_silu: + torch.ops._C.silu_and_mul(qpn_final, qpn_raw) + return qpn_final + return qpn_raw + + tm_graph, tm_output = _capture(run_tm) + qpn_graph, qpn_output = _capture(run_qpn2) + tm_timing = _time_graph(tm_graph, warmup, iterations, trials) + qpn_timing = _time_graph(qpn_graph, warmup, iterations, trials) + tm_graph.replay() + qpn_graph.replay() + torch.accelerator.synchronize() + reference = _fp32_reference(projection, x, device) + saved_us = float(tm_timing["median_us"]) - float(qpn_timing["median_us"]) + return { + "name": projection.name, + "m": m, + "n": n, + "k": k, + "gated_silu": projection.gated_silu, + "calls_per_round": projection.calls_per_round, + "qpn2_config": { + "split_k": split_k, + "accumulator_chains": accumulator_chains, + }, + "turbomind": tm_timing, + "qpn2": qpn_timing, + "speedup": float(tm_timing["median_us"]) / float(qpn_timing["median_us"]), + "saved_ms_per_round": saved_us * projection.calls_per_round / 1000.0, + "quality_vs_turbomind": _quality(qpn_output, tm_output), + "turbomind_quality_vs_fp32": _quality(tm_output, reference), + "qpn2_quality_vs_fp32": _quality(qpn_output, reference), + } + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument( + "--model", type=Path, default=Path("/home/ymzx/models/Qwen3.8-27B-NVFP4") + ) + parser.add_argument("--layer", type=int, default=55) + parser.add_argument("--tp-rank", type=int, default=0) + parser.add_argument("--tp-size", type=int, default=4) + parser.add_argument("--m", type=int, default=8) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--warmup", type=int, default=50) + parser.add_argument("--iterations", type=int, default=1000) + parser.add_argument("--trials", type=int, default=5) + parser.add_argument("--qpn2-source", type=Path) + parser.add_argument( + "--production-library", + type=Path, + help="Load a complete candidate vllm._C containing production QPN2 ops.", + ) + parser.add_argument("--extension-name", default="v100_skinny_qpn2_micro_v1") + parser.add_argument("--json-out", type=Path, required=True) + return parser.parse_args() + + +def _load_production_library(path: Path) -> None: + spec = importlib.util.spec_from_file_location("vllm._C", path) + if spec is None or spec.loader is None: + raise RuntimeError(f"cannot load production library: {path}") + module = importlib.util.module_from_spec(spec) + sys.modules["vllm._C"] = module + spec.loader.exec_module(module) + + +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, + ): + raise RuntimeError("benchmark requires an exact SM70 CUDA device") + if not 1 <= args.m <= 8: + raise ValueError("QPN2 benchmark supports M in [1, 8]") + production = args.production_library is not None + if production: + _load_production_library(args.production_library) + extension = None + else: + if args.qpn2_source is None: + raise ValueError("--qpn2-source is required without --production-library") + extension = load( + name=args.extension_name, + sources=[str(args.qpn2_source)], + extra_cuda_cflags=[ + "-O3", + "--use_fast_math", + "-lineinfo", + "-gencode=arch=compute_70,code=sm_70", + ], + verbose=False, + ) + torch.manual_seed(20260824) + projections = _load_projection_shards( + args.model, args.layer, args.tp_rank, args.tp_size + ) + rows = [ + _run_projection( + projection, + extension, + production, + args.m, + device, + args.warmup, + args.iterations, + args.trials, + ) + for projection in projections + ] + payload = { + "device": torch.cuda.get_device_name(device), + "capability": list(torch.cuda.get_device_capability(device)), + "torch": torch.__version__, + "model": str(args.model), + "model_config_sha256": _sha256_file(args.model / "config.json"), + "layer": args.layer, + "tp_rank": args.tp_rank, + "tp_size": args.tp_size, + "qpn2_source": str(args.qpn2_source) if args.qpn2_source else None, + "qpn2_source_sha256": ( + _sha256_file(args.qpn2_source) if args.qpn2_source else None + ), + "production_library": ( + str(args.production_library) if args.production_library else None + ), + "production_library_sha256": ( + _sha256_file(args.production_library) if args.production_library else None + ), + "warmup": args.warmup, + "iterations": args.iterations, + "trials": args.trials, + "rows": rows, + "projected_total_saved_ms_per_round": sum( + float(row["saved_ms_per_round"]) for row in rows + ), + } + args.json_out.parent.mkdir(parents=True, exist_ok=True) + args.json_out.write_text(json.dumps(payload, indent=2) + "\n") + print(json.dumps(payload, indent=2), flush=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/profile_sm70_fp8_qpn8_gated_pair_ncu.py b/benchmarks/profile_sm70_fp8_qpn8_gated_pair_ncu.py new file mode 100644 index 0000000000..3070b16236 --- /dev/null +++ b/benchmarks/profile_sm70_fp8_qpn8_gated_pair_ncu.py @@ -0,0 +1,140 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Paired NCU probe for current and QPN8 fused gate/up kernels.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import torch +import vllm._C # noqa: F401 +from benchmark_sm70_fp8_qpn8 import ( + _dequantized_weight, + _error_stats, + _load_case, + _qpn8_group_scales, + _qpn8_prepack, +) + +from vllm import _sm70_ops as sm70_ops + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument( + "--model", + type=Path, + default=Path("/home/ymzx/models/Qwen3.8-27B-FP8"), + ) + parser.add_argument("--qpn8-library", type=Path, required=True) + parser.add_argument("--out", type=Path, required=True) + parser.add_argument("--m", type=int, default=1) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--repeats", type=int, default=100) + return parser.parse_args() + + +def _time(launch, warmup: int, repeats: int) -> float: + for _ in range(warmup): + launch() + torch.accelerator.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(repeats): + launch() + end.record() + end.synchronize() + return float(start.elapsed_time(end) * 1000.0 / repeats) + + +def main() -> int: + args = _parse_args() + torch.ops.load_library(str(args.qpn8_library.resolve())) + device = torch.device("cuda:0") + torch.accelerator.set_device_index(device.index or 0) + if torch.cuda.get_device_capability(device) != (7, 0): + raise RuntimeError("This profiler requires SM70/V100") + + qweight, scales, _ = _load_case(args.model, "gate_up_fused", 4, 0, device) + n, k = (int(dim) for dim in qweight.shape) + hidden = n // 2 + codes = _qpn8_prepack(qweight.view(torch.uint8)) + group_scales = _qpn8_group_scales(scales, n, k) + tm_weight, tm_scales, meta = sm70_ops.fp8_sm70_prepare(qweight, scales, 128, True) + k_ld, q_ld = (int(value) for value in meta.tolist()) + inputs = torch.randn((args.m, k), device=device, dtype=torch.float16).mul_(0.1) + current_out = torch.empty((args.m, hidden), device=device, dtype=torch.float16) + qpn_out = torch.empty_like(current_out) + + def launch_current() -> None: + sm70_ops.fp8_gemm_sm70_out( + current_out, + inputs, + tm_weight, + tm_scales, + 128, + k_ld, + q_ld, + True, + ) + + def launch_qpn() -> None: + sm70_ops.fp8_qpn8_gated_pair_sm70_out( + qpn_out, + inputs, + codes, + group_scales, + 8, + 2, + True, + True, + ) + + current_us = _time(launch_current, args.warmup, args.repeats) + qpn_us = _time(launch_qpn, args.warmup, args.repeats) + reference_raw = inputs.float().matmul(_dequantized_weight(qweight, scales).t()) + gate, up = reference_raw.chunk(2, dim=1) + reference = (torch.nn.functional.silu(gate) * up).half() + quality = { + "current": _error_stats(current_out, reference), + "qpn8_pair": _error_stats(qpn_out, reference), + } + + torch.cuda.cudart().cudaProfilerStart() + torch.cuda.nvtx.range_push("current_gate_up_fused") + launch_current() + torch.cuda.nvtx.range_pop() + torch.cuda.nvtx.range_push("qpn8_gate_up_fused_pair") + launch_qpn() + torch.cuda.nvtx.range_pop() + torch.accelerator.synchronize(device) + torch.cuda.cudart().cudaProfilerStop() + + result = { + "device": torch.cuda.get_device_name(device), + "m": args.m, + "n": n, + "k": k, + "hidden": hidden, + "current_unprofiled_us": current_us, + "qpn8_pair_unprofiled_us": qpn_us, + "speedup": current_us / qpn_us, + "qpn8_config": { + "split_k": 8, + "nacc": 2, + "decoder": "fast", + "prefetch": True, + }, + "quality": quality, + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(result, indent=2), encoding="utf-8") + print(json.dumps(result, indent=2)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/profile_sm70_fp8_qpn8_ncu.py b/benchmarks/profile_sm70_fp8_qpn8_ncu.py new file mode 100644 index 0000000000..7bc5e02c55 --- /dev/null +++ b/benchmarks/profile_sm70_fp8_qpn8_ncu.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Profile only the real-shape QPN8 winners with Nsight Compute.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + +import torch +import vllm._C # noqa: F401 +from benchmark_sm70_fp8_qpn8 import ( + _load_case, + _qpn8_group_scales, + _qpn8_prepack, +) + +from vllm import _sm70_ops as sm70_ops + +_CONFIGS = { + "down": (16, 1, True), + "gdn_in": (32, 1, True), + "output": (16, 1, True), + "full_qkv": (16, 1, True), + "gate_up_raw": (16, 1, True), +} + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument( + "--model", + type=Path, + default=Path("/home/ymzx/models/Qwen3.8-27B-FP8"), + ) + parser.add_argument("--qpn8-library", type=Path, required=True) + parser.add_argument("--out", type=Path, required=True) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--m", type=int, default=1) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--repeats", type=int, default=100) + parser.add_argument("--paired-current", action="store_true") + return parser.parse_args() + + +def _time(launch: Any, warmup: int, repeats: int) -> float: + for _ in range(warmup): + launch() + torch.accelerator.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(repeats): + launch() + end.record() + end.synchronize() + return float(start.elapsed_time(end) * 1000.0 / repeats) + + +def main() -> int: + args = _parse_args() + torch.ops.load_library(str(args.qpn8_library.resolve())) + device = torch.device(args.device) + torch.accelerator.set_device_index(device.index or 0) + if torch.cuda.get_device_capability(device) != (7, 0): + raise RuntimeError("This profiler requires SM70/V100") + if args.m < 1 or args.m > 8: + raise ValueError("M must be in [1, 8]") + + prepared: list[dict[str, Any]] = [] + for case, (split_k, nacc, fast_decoder) in _CONFIGS.items(): + qweight, scales, note = _load_case(args.model, case, 4, 0, device) + n, k = (int(dim) for dim in qweight.shape) + codes = _qpn8_prepack(qweight.view(torch.uint8)) + group_scales = _qpn8_group_scales(scales, n, k) + tm_weight, tm_scales, meta = sm70_ops.fp8_sm70_prepare( + qweight, scales, 128, False + ) + k_ld, q_ld = (int(value) for value in meta.tolist()) + inputs = torch.randn((args.m, k), device=device, dtype=torch.float16).mul_(0.1) + qpn_output = torch.empty((args.m, n), device=device, dtype=torch.float16) + current_output = torch.empty_like(qpn_output) + + def launch_qpn( + output: torch.Tensor = qpn_output, + inputs: torch.Tensor = inputs, + codes: torch.Tensor = codes, + group_scales: torch.Tensor = group_scales, + split_k: int = split_k, + nacc: int = nacc, + fast_decoder: bool = fast_decoder, + ) -> None: + sm70_ops.fp8_qpn8_gemm_sm70_out( + output, + inputs, + codes, + group_scales, + split_k, + nacc, + fast_decoder, + False, + ) + + def launch_current( + output: torch.Tensor = current_output, + inputs: torch.Tensor = inputs, + tm_weight: torch.Tensor = tm_weight, + tm_scales: torch.Tensor = tm_scales, + k_ld: int = k_ld, + q_ld: int = q_ld, + ) -> None: + sm70_ops.fp8_gemm_sm70_out( + output, + inputs, + tm_weight, + tm_scales, + 128, + k_ld, + q_ld, + False, + ) + + prepared.append( + { + "case": case, + "note": note, + "m": args.m, + "n": n, + "k": k, + "split_k": split_k, + "nacc": nacc, + "decoder": "fast" if fast_decoder else "scalar", + "launch_qpn": launch_qpn, + "launch_current": launch_current, + "tensors": [ + qweight, + scales, + codes, + group_scales, + inputs, + qpn_output, + current_output, + tm_weight, + tm_scales, + ], + } + ) + + for item in prepared: + item["qpn_unprofiled_us"] = _time(item["launch_qpn"], args.warmup, args.repeats) + if args.paired_current: + item["current_unprofiled_us"] = _time( + item["launch_current"], args.warmup, args.repeats + ) + + torch.cuda.cudart().cudaProfilerStart() + torch.cuda.nvtx.range_push("qpn8_real_shape_winners") + for item in prepared: + if args.paired_current: + torch.cuda.nvtx.range_push(f"current_{item['case']}") + item["launch_current"]() + torch.cuda.nvtx.range_pop() + torch.cuda.nvtx.range_push(f"qpn8_{item['case']}") + item["launch_qpn"]() + torch.cuda.nvtx.range_pop() + torch.accelerator.synchronize(device) + torch.cuda.nvtx.range_pop() + torch.cuda.cudart().cudaProfilerStop() + + public = [ + { + key: value + for key, value in item.items() + if key not in ("launch_qpn", "launch_current", "tensors") + } + for item in prepared + ] + result = { + "device": torch.cuda.get_device_name(device), + "capability": list(torch.cuda.get_device_capability(device)), + "torch": torch.__version__, + "torch_cuda": torch.version.cuda, + "profiles": public, + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(result, indent=2), encoding="utf-8") + print(json.dumps(result, indent=2)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) 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/external_projects/vllm_flash_attn.cmake b/cmake/external_projects/vllm_flash_attn.cmake index 331556d232..8918d9e7b0 100644 --- a/cmake/external_projects/vllm_flash_attn.cmake +++ b/cmake/external_projects/vllm_flash_attn.cmake @@ -50,6 +50,12 @@ elseif(VLLM_FLASH_ATTN_SM70) COMMAND ${PATCH_EXECUTABLE} --batch --forward -p1 -l -i ${CMAKE_CURRENT_LIST_DIR}/../patches/sm70_flash_attn_d256_splitkv3.patch + COMMAND + ${PATCH_EXECUTABLE} --batch --forward -p1 -l + -i ${CMAKE_CURRENT_LIST_DIR}/../patches/sm70_flash_attn_d256_k_pingpong.patch + COMMAND + ${PATCH_EXECUTABLE} --batch --forward -p1 -l + -i ${CMAKE_CURRENT_LIST_DIR}/../patches/sm70_flash_attn_d256_gqa_arch.patch BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn ) else() diff --git a/cmake/patches/sm70_flash_attn_d256_gqa_arch.patch b/cmake/patches/sm70_flash_attn_d256_gqa_arch.patch new file mode 100644 index 0000000000..d6d65439a5 --- /dev/null +++ b/cmake/patches/sm70_flash_attn_d256_gqa_arch.patch @@ -0,0 +1,1673 @@ +diff --git a/CMakeLists.txt b/CMakeLists.txt +index acf2219..f5e2732 100644 +--- a/CMakeLists.txt ++++ b/CMakeLists.txt +@@ -84,6 +84,7 @@ string(REPLACE "-O2" "-O3" CMAKE_CUDA_FLAGS_RELWITHDEBINFO "${CMAKE_CUDA_FLAGS_R + + if (FA2_ENABLED) + set(FA2_GEN_SRCS ++ "csrc/flash_attn/src/flash_fwd_d256_gqa_arch_sm70.cu" + "csrc/flash_attn/src/flash_fwd_d256_splitd_sm70.cu" + "csrc/flash_attn/src/flash_fwd_hdim256_causal_prefill_sm70.cu" + "csrc/flash_attn/src/flash_fwd_split_hdim256_causal_prefill_sm70.cu") +@@ -114,7 +115,8 @@ if (FA2_ENABLED) + csrc/flash_attn + csrc/flash_attn/src + csrc/common +- csrc/cutlass/include) ++ csrc/cutlass/include ++ csrc/cutlass/examples/35_gemm_softmax) + + # custom definitions + target_compile_definitions(_vllm_fa2_C PRIVATE +@@ -125,5 +127,16 @@ if (FA2_ENABLED) + # FLASHATTENTION_DISABLE_UNEVEN_K + # FLASHATTENTION_DISABLE_LOCAL + FLASHATTENTION_DISABLE_PYBIND ++ PREFIX_TORCH_EXTENSION ++ PREFIX_QK_FULL_STATS ++ PREFIX_QK_SKIP_APPLY ++ QK_TB_M=128 ++ QK_TB_N=512 ++ QK_WARP_M=64 ++ QK_WARP_N=128 ++ PV_TB_M=64 ++ PV_TB_N=256 ++ PV_WARP_M=32 ++ PV_WARP_N=64 + ) + endif () +diff --git a/csrc/flash_attn/flash_api_torch_lib.cpp b/csrc/flash_attn/flash_api_torch_lib.cpp +index 4b7601a..f5b71a2 100644 +--- a/csrc/flash_attn/flash_api_torch_lib.cpp ++++ b/csrc/flash_attn/flash_api_torch_lib.cpp +@@ -33,6 +33,24 @@ at::Tensor sm70_d256_splitd_dense_splitkv3_fwd( + double softmax_scale, + bool causal); + ++at::Tensor sm70_d256_splitd_dense_state_fwd( ++ const at::Tensor &q, ++ const at::Tensor &k, ++ const at::Tensor &v, ++ at::Tensor &state_max, ++ at::Tensor &state_sum, ++ at::Tensor &out, ++ double softmax_scale, ++ bool causal); ++ ++at::Tensor sm70_d256_gqa_architecture_fwd( ++ const at::Tensor &q, ++ const at::Tensor &k, ++ const at::Tensor &v, ++ at::Tensor &out, ++ double softmax_scale, ++ bool causal); ++ + at::Tensor sm70_d256_splitd_paged_fwd( + const at::Tensor &q, + const at::Tensor &k, +@@ -98,6 +116,17 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { + ops.impl("sm70_d256_splitd_n32_dense_fwd", torch::kCUDA, + make_pytorch_shim(&sm70_d256_splitd_dense_fwd)); + ++ ops.def("sm70_d256_splitd_n32_dense_state_fwd(Tensor q, Tensor k, " ++ "Tensor v, Tensor(a!) state_max, Tensor(b!) state_sum, " ++ "Tensor(c!) out, float softmax_scale, bool causal) -> Tensor(c!)"); ++ ops.impl("sm70_d256_splitd_n32_dense_state_fwd", torch::kCUDA, ++ make_pytorch_shim(&sm70_d256_splitd_dense_state_fwd)); ++ ++ ops.def("sm70_d256_gqa_architecture_fwd(Tensor q, Tensor k, Tensor v, " ++ "Tensor(a!) out, float softmax_scale, bool causal) -> Tensor(a!)"); ++ ops.impl("sm70_d256_gqa_architecture_fwd", torch::kCUDA, ++ make_pytorch_shim(&sm70_d256_gqa_architecture_fwd)); ++ + ops.def("sm70_d256_splitd_n32_dense_splitkv3_fwd(Tensor q, Tensor k, " + "Tensor v, Tensor(a!) partial_out, Tensor(b!) partial_max, " + "Tensor(c!) partial_sum, Tensor(d!) out, float softmax_scale, " +diff --git a/csrc/flash_attn/src/flash_fwd_d256_gqa_arch_sm70.cu b/csrc/flash_attn/src/flash_fwd_d256_gqa_arch_sm70.cu +new file mode 100644 +index 0000000..197453f +--- /dev/null ++++ b/csrc/flash_attn/src/flash_fwd_d256_gqa_arch_sm70.cu +@@ -0,0 +1,1359 @@ ++// SPDX-License-Identifier: BSD-3-Clause ++ ++/*************************************************************************************************** ++ * End-to-end prefix architecture screen for V100/SM70. ++ * ++ * QK GEMM writes FP16 logits and row statistics. PV normalizes logits in its ++ * A-operand transform and writes one reusable FP16 block partial. A float ++ * online accumulator folds that partial into prefix softmax state before the ++ * next block, avoiding resident storage for every block partial. ++ **************************************************************************************************/ ++ ++#include ++#include ++#include ++ ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++ ++#if defined(PREFIX_TORCH_EXTENSION) ++#include ++#include ++#include ++#include ++#include "namespace_config.h" ++#endif ++ ++#include "cutlass/cutlass.h" ++#include "cutlass/device_kernel.h" ++#include "cutlass/epilogue/thread/linear_combination.h" ++#include "cutlass/gemm/kernel/default_gemm.h" ++#include "cutlass/gemm/kernel/gemm.h" ++#include "cutlass/gemm/threadblock/mma_pipelined.h" ++#include "cutlass/layout/matrix.h" ++#include "cutlass/numeric_conversion.h" ++#include "gemm_with_softmax.h" ++ ++namespace { ++ ++using Element = cutlass::half_t; ++ ++#ifndef QK_TB_M ++#define QK_TB_M 128 ++#endif ++#ifndef QK_TB_N ++#define QK_TB_N 128 ++#endif ++#ifndef QK_WARP_M ++#define QK_WARP_M 32 ++#endif ++#ifndef QK_WARP_N ++#define QK_WARP_N 64 ++#endif ++ ++using QKLayoutA = cutlass::layout::RowMajor; ++using QKLayoutB = cutlass::layout::ColumnMajor; ++using QKThreadblockShape = cutlass::gemm::GemmShape; ++using QKWarpShape = cutlass::gemm::GemmShape; ++using QKInstructionShape = cutlass::gemm::GemmShape<8, 8, 4>; ++using QKOutputOp = cutlass::epilogue::thread::LinearCombination< ++ Element, 8, Element, Element>; ++using QKGemm = cutlass::GemmSoftmax< ++ Element, ++ QKLayoutA, ++ Element, ++ QKLayoutB, ++ Element, ++ Element, ++ cutlass::arch::OpClassTensorOp, ++ cutlass::arch::Sm70, ++ QKThreadblockShape, ++ QKWarpShape, ++ QKInstructionShape, ++ QKOutputOp, ++ 2, ++ cutlass::MatrixShape<1, 1024>>; ++ ++#ifndef PV_TB_M ++#define PV_TB_M 64 ++#endif ++#ifndef PV_TB_N ++#define PV_TB_N 128 ++#endif ++#ifndef PV_WARP_M ++#define PV_WARP_M 32 ++#endif ++#ifndef PV_WARP_N ++#define PV_WARP_N 64 ++#endif ++ ++using PVLayout = cutlass::layout::RowMajor; ++using PVAccumulator = Element; ++using PVThreadblockShape = cutlass::gemm::GemmShape; ++using PVWarpShape = cutlass::gemm::GemmShape; ++using PVInstructionShape = cutlass::gemm::GemmShape<8, 8, 4>; ++using PVSwizzle = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>; ++using PVOutputOp = cutlass::epilogue::thread::LinearCombination< ++ Element, 8, PVAccumulator, float>; ++ ++constexpr int kPVAlignment = 8; ++ ++using PVDefaultKernel = typename cutlass::gemm::kernel::DefaultGemm< ++ Element, ++ PVLayout, ++ kPVAlignment, ++ Element, ++ PVLayout, ++ kPVAlignment, ++ Element, ++ PVLayout, ++ PVAccumulator, ++ cutlass::arch::OpClassTensorOp, ++ cutlass::arch::Sm70, ++ PVThreadblockShape, ++ PVWarpShape, ++ PVInstructionShape, ++ PVOutputOp, ++ PVSwizzle, ++ 2, ++ false, ++ cutlass::arch::OpMultiplyAdd>::GemmKernel; ++ ++using PVDefaultMma = typename PVDefaultKernel::Mma; ++using PVIteratorA = typename PVDefaultMma::IteratorA; ++using PVIteratorB = typename PVDefaultMma::IteratorB; ++using PVSmemIteratorA = typename PVDefaultMma::SmemIteratorA; ++using PVSmemIteratorB = typename PVDefaultMma::SmemIteratorB; ++ ++__device__ float const* g_row_max = nullptr; ++__device__ float const* g_row_inv_sum = nullptr; ++__device__ float* g_row_sum_out = nullptr; ++__device__ int g_rows = 0; ++ ++#if defined(PREFIX_QK_FULL_STATS) ++struct ExpRowSumTransformA { ++ using InputFragment = typename PVIteratorA::Fragment; ++ using OutputFragment = cutlass::Array< ++ typename PVSmemIteratorA::Element, InputFragment::kElements>; ++ using ThreadMap = typename PVIteratorA::ThreadMap; ++ static constexpr int kAccessesPerVector = ++ PVIteratorA::UnderlyingIterator::kAccessesPerVector; ++ static constexpr int kContiguousIterations = ++ ThreadMap::Iterations::kContiguous; ++ static constexpr int kStridedIterations = ThreadMap::Iterations::kStrided; ++ ++ float row_max[kStridedIterations]; ++ float row_inv_sum[kStridedIterations]; ++ ++ CUTLASS_DEVICE ++ ExpRowSumTransformA() { ++ auto thread_offset = ThreadMap::initial_offset(threadIdx.x); ++#pragma unroll ++ for (int s = 0; s < kStridedIterations; ++s) { ++ int row = blockIdx.x * PVThreadblockShape::kM ++ + thread_offset.strided() + s * ThreadMap::Delta::kStrided; ++ row_max[s] = row < g_rows ? g_row_max[row] : 0.0f; ++ row_inv_sum[s] = row < g_rows ? g_row_inv_sum[row] : 0.0f; ++ } ++ } ++ ++ CUTLASS_DEVICE ++ OutputFragment operator()(InputFragment const& input) { ++ OutputFragment output; ++ constexpr float kLog2E = 1.4426950408889634f; ++ constexpr int kElementsPerAccess = PVIteratorA::AccessType::kElements; ++ auto const* input_access = ++ reinterpret_cast(&input); ++ auto* output_access = ++ reinterpret_cast(&output); ++#pragma unroll ++ for (int s = 0; s < kStridedIterations; ++s) { ++#pragma unroll ++ for (int c = 0; c < kContiguousIterations; ++c) { ++#pragma unroll ++ for (int v = 0; v < kAccessesPerVector; ++v) { ++ int index = v + kAccessesPerVector * (c + s * kContiguousIterations); ++ typename PVIteratorA::AccessType transformed; ++#pragma unroll ++ for (int e = 0; e < kElementsPerAccess; ++e) { ++ float value = static_cast(input_access[index][e]); ++ transformed[e] = Element( ++ exp2f((value - row_max[s]) * kLog2E) * row_inv_sum[s]); ++ } ++ output_access[index] = transformed; ++ } ++ } ++ } ++ return output; ++ } ++}; ++#else ++struct ExpRowSumTransformA { ++ using InputFragment = typename PVIteratorA::Fragment; ++ using OutputFragment = cutlass::Array< ++ typename PVSmemIteratorA::Element, InputFragment::kElements>; ++ using ThreadMap = typename PVIteratorA::ThreadMap; ++ ++ static constexpr int kAccessesPerVector = ++ PVIteratorA::UnderlyingIterator::kAccessesPerVector; ++ static constexpr int kContiguousIterations = ++ ThreadMap::Iterations::kContiguous; ++ static constexpr int kStridedIterations = ThreadMap::Iterations::kStrided; ++ ++ float row_max[kStridedIterations]; ++ float row_sum[kStridedIterations]; ++ bool valid = true; ++ ++ CUTLASS_DEVICE ++ void set_valid(bool is_valid) { ++ valid = is_valid; ++ } ++ ++ CUTLASS_DEVICE ++ ExpRowSumTransformA() { ++ auto thread_offset = ThreadMap::initial_offset(threadIdx.x); ++#pragma unroll ++ for (int s = 0; s < kStridedIterations; ++s) { ++ int row = blockIdx.x * PVThreadblockShape::kM ++ + thread_offset.strided() + s * ThreadMap::Delta::kStrided; ++ row_max[s] = row < g_rows ? g_row_max[row] : 0.0f; ++ row_sum[s] = 0.0f; ++ } ++ } ++ ++ CUTLASS_DEVICE ++ OutputFragment operator()(InputFragment const& input) { ++ if (!valid) { ++ OutputFragment output; ++ output.clear(); ++ return output; ++ } ++ OutputFragment output; ++ constexpr float kLog2E = 1.4426950408889634f; ++ constexpr int kElementsPerAccess = PVIteratorA::AccessType::kElements; ++ auto const* input_access = ++ reinterpret_cast(&input); ++ auto* output_access = ++ reinterpret_cast(&output); ++#pragma unroll ++ for (int s = 0; s < kStridedIterations; ++s) { ++#pragma unroll ++ for (int c = 0; c < kContiguousIterations; ++c) { ++#pragma unroll ++ for (int v = 0; v < kAccessesPerVector; ++v) { ++ int index = v + kAccessesPerVector * (c + s * kContiguousIterations); ++ typename PVIteratorA::AccessType transformed; ++#pragma unroll ++ for (int e = 0; e < kElementsPerAccess; ++e) { ++ float value = static_cast(input_access[index][e]); ++ float weight = exp2f((value - row_max[s]) * kLog2E); ++ if (blockIdx.y == 0) { ++ row_sum[s] += weight; ++ } ++ transformed[e] = Element(weight); ++ } ++ output_access[index] = transformed; ++ } ++ } ++ } ++ return output; ++ } ++ ++ CUTLASS_DEVICE ++ void finalize() { ++ if (blockIdx.y != 0) { ++ return; ++ } ++ constexpr int kThreadsPerRow = ++ ThreadMap::Detail::WarpThreadArrangement::kContiguous; ++ auto thread_offset = ThreadMap::initial_offset(threadIdx.x); ++ int lane = threadIdx.x & 31; ++#pragma unroll ++ for (int s = 0; s < kStridedIterations; ++s) { ++ float sum = row_sum[s]; ++#pragma unroll ++ for (int offset = 1; offset < kThreadsPerRow; offset <<= 1) { ++ sum += __shfl_xor_sync(0xffffffffu, sum, offset); ++ } ++ int row = blockIdx.x * PVThreadblockShape::kM ++ + thread_offset.strided() + s * ThreadMap::Delta::kStrided; ++ if ((lane % kThreadsPerRow) == 0 && row < g_rows) { ++ g_row_sum_out[row] = sum; ++ } ++ } ++ } ++}; ++#endif ++ ++using PVTransformB = cutlass::NumericArrayConverter< ++ typename PVSmemIteratorB::Element, ++ typename PVIteratorB::Element, ++ PVIteratorB::Fragment::kElements>; ++using PVMma = cutlass::gemm::threadblock::MmaPipelined< ++ typename PVDefaultMma::Shape, ++ PVIteratorA, ++ PVSmemIteratorA, ++ PVIteratorB, ++ PVSmemIteratorB, ++ PVAccumulator, ++ PVLayout, ++ typename PVDefaultMma::Policy, ++ ExpRowSumTransformA, ++ PVTransformB>; ++using PVKernel = cutlass::gemm::kernel::Gemm< ++ PVMma, typename PVDefaultKernel::Epilogue, PVSwizzle, false>; ++ ++void check(cudaError_t result, char const* operation) { ++ if (result != cudaSuccess) { ++ std::cerr << operation << ": " << cudaGetErrorString(result) << "\n"; ++ std::exit(EXIT_FAILURE); ++ } ++} ++ ++void check(cutlass::Status status, char const* operation) { ++ if (status != cutlass::Status::kSuccess) { ++ std::cerr << operation << ": CUTLASS status " << int(status) << "\n"; ++ std::exit(EXIT_FAILURE); ++ } ++} ++ ++struct PVLauncher { ++ typename PVKernel::Params params; ++ dim3 grid; ++ dim3 block; ++ int smem_bytes; ++ ++ PVLauncher(Element* scores, Element* value, Element* output, int rows, int k) { ++ cutlass::gemm::GemmCoord problem(rows, 256, k); ++ PVSwizzle swizzle; ++ auto tiled_shape = swizzle.get_tiled_shape( ++ problem, ++ {PVThreadblockShape::kM, PVThreadblockShape::kN, ++ PVThreadblockShape::kK}, ++ 1); ++ params = typename PVKernel::Params( ++ problem, ++ tiled_shape, ++ {scores, PVLayout(k)}, ++ {value, PVLayout(256)}, ++ {output, PVLayout(256)}, ++ {output, PVLayout(256)}, ++ typename PVOutputOp::Params(1.0f, 0.0f), ++ nullptr); ++ grid = swizzle.get_grid_shape(tiled_shape); ++ block = dim3(PVKernel::kThreadCount, 1, 1); ++ smem_bytes = int(sizeof(typename PVKernel::SharedStorage)); ++ if (smem_bytes >= 48 * 1024) { ++ check(cudaFuncSetAttribute( ++ cutlass::Kernel, ++ cudaFuncAttributeMaxDynamicSharedMemorySize, ++ smem_bytes), ++ "set PV dynamic shared memory"); ++ } ++ } ++ ++ void launch(cudaStream_t stream) const { ++ cutlass::Kernel<<>>(params); ++ } ++}; ++ ++__global__ void prepare_weights( ++ float const* block_max, ++ float const* block_sum, ++ float* weights, ++ float* state_max, ++ float* state_sum, ++ int blocks, ++ int rows) { ++ int row = blockIdx.x * blockDim.x + threadIdx.x; ++ if (row >= rows) { ++ return; ++ } ++ float global_max = -CUDART_INF_F; ++ for (int block = 0; block < blocks; ++block) { ++ global_max = fmaxf(global_max, block_max[block * rows + row]); ++ } ++ float denominator = 0.0f; ++ for (int block = 0; block < blocks; ++block) { ++ float scale = exp2f( ++ (block_max[block * rows + row] - global_max) * ++ 1.4426950408889634f); ++#if defined(PREFIX_QK_FULL_STATS) ++ float scaled_sum = scale / block_sum[block * rows + row]; ++ weights[block * rows + row] = scaled_sum; ++ denominator += scaled_sum; ++#else ++ weights[block * rows + row] = scale; ++ denominator += scale * block_sum[block * rows + row]; ++#endif ++ } ++ float inverse = 1.0f / denominator; ++ for (int block = 0; block < blocks; ++block) { ++ weights[block * rows + row] *= inverse; ++ } ++ state_max[row] = global_max; ++ state_sum[row] = denominator; ++} ++ ++__global__ void merge_partials( ++ __half const* partials, ++ float const* weights, ++ __half* output, ++ int blocks, ++ int rows) { ++ int pair = blockIdx.x * blockDim.x + threadIdx.x; ++ constexpr int kPairsPerRow = 128; ++ int total_pairs = rows * kPairsPerRow; ++ if (pair >= total_pairs) { ++ return; ++ } ++ int row = pair / kPairsPerRow; ++ int column_pair = pair - row * kPairsPerRow; ++ float2 accumulator = make_float2(0.0f, 0.0f); ++ auto const* partials2 = reinterpret_cast<__half2 const*>(partials); ++ auto* output2 = reinterpret_cast<__half2*>(output); ++ for (int block = 0; block < blocks; ++block) { ++ int index = (block * rows + row) * kPairsPerRow + column_pair; ++ float2 value = __half22float2(partials2[index]); ++ float weight = weights[block * rows + row]; ++ accumulator.x = fmaf(weight, value.x, accumulator.x); ++ accumulator.y = fmaf(weight, value.y, accumulator.y); ++ } ++ output2[pair] = __floats2half2_rn(accumulator.x, accumulator.y); ++} ++ ++__global__ void update_prefix_accumulator( ++ __half const* block_output, ++ float const* block_max, ++ float const* block_inv_sum, ++ float* prefix_accumulator, ++ float* prefix_max, ++ float* prefix_sum, ++ int rows, ++ bool initialize) { ++ int row = blockIdx.x; ++ int d = threadIdx.x; ++ if (row >= rows || d >= 256) { ++ return; ++ } ++ __shared__ float scales[3]; ++ if (d == 0) { ++ constexpr float kLog2E = 1.4426950408889634f; ++ float next_max = block_max[row]; ++ float next_mass = 1.0f / block_inv_sum[row]; ++ if (initialize) { ++ scales[0] = 0.0f; ++ scales[1] = next_mass; ++ scales[2] = next_mass; ++ prefix_max[row] = next_max; ++ } else { ++ float old_max = prefix_max[row]; ++ float global_max = fmaxf(old_max, next_max); ++ float old_scale = exp2f((old_max - global_max) * kLog2E); ++ float next_scale = exp2f((next_max - global_max) * kLog2E); ++ scales[0] = old_scale; ++ scales[1] = next_mass * next_scale; ++ scales[2] = prefix_sum[row] * old_scale + scales[1]; ++ prefix_max[row] = global_max; ++ } ++ prefix_sum[row] = scales[2]; ++ } ++ __syncthreads(); ++ int64_t element = int64_t(row) * 256 + d; ++ float block_value = float(block_output[element]); ++ float old_value = initialize ? 0.0f : prefix_accumulator[element]; ++ prefix_accumulator[element] = ++ fmaf(old_value, scales[0], block_value * scales[1]); ++} ++ ++__global__ void prepare_prefix_update( ++ float const* block_max, ++ float const* block_inv_sum, ++ float* prefix_max, ++ float* prefix_sum, ++ float* old_scales, ++ float* block_scales, ++ int rows, ++ bool initialize) { ++ int row = blockIdx.x * blockDim.x + threadIdx.x; ++ if (row >= rows) { ++ return; ++ } ++ constexpr float kLog2E = 1.4426950408889634f; ++ float next_max = block_max[row]; ++ float next_mass = 1.0f / block_inv_sum[row]; ++ if (initialize) { ++ old_scales[row] = 0.0f; ++ block_scales[row] = next_mass; ++ prefix_max[row] = next_max; ++ prefix_sum[row] = next_mass; ++ return; ++ } ++ float old_max = prefix_max[row]; ++ float global_max = fmaxf(old_max, next_max); ++ float old_scale = exp2f((old_max - global_max) * kLog2E); ++ float block_scale = ++ next_mass * exp2f((next_max - global_max) * kLog2E); ++ old_scales[row] = old_scale; ++ block_scales[row] = block_scale; ++ prefix_max[row] = global_max; ++ prefix_sum[row] = prefix_sum[row] * old_scale + block_scale; ++} ++ ++__global__ void apply_prefix_update_half2( ++ __half const* block_output, ++ float* prefix_accumulator, ++ float const* old_scales, ++ float const* block_scales, ++ int rows, ++ bool initialize) { ++ int pair = blockIdx.x * blockDim.x + threadIdx.x; ++ constexpr int kPairsPerRow = 128; ++ if (pair >= rows * kPairsPerRow) { ++ return; ++ } ++ int row = pair / kPairsPerRow; ++ auto const* block2 = reinterpret_cast<__half2 const*>(block_output); ++ auto* accumulator2 = reinterpret_cast(prefix_accumulator); ++ float2 block_value = __half22float2(block2[pair]); ++ float2 old_value = initialize ++ ? make_float2(0.0f, 0.0f) ++ : accumulator2[pair]; ++ float old_scale = old_scales[row]; ++ float block_scale = block_scales[row]; ++ accumulator2[pair] = make_float2( ++ fmaf(old_value.x, old_scale, block_value.x * block_scale), ++ fmaf(old_value.y, old_scale, block_value.y * block_scale)); ++} ++ ++__global__ void finalize_prefix_accumulator( ++ float const* prefix_accumulator, ++ float const* prefix_sum, ++ __half* output, ++ int rows) { ++ int pair = blockIdx.x * blockDim.x + threadIdx.x; ++ constexpr int kPairsPerRow = 128; ++ if (pair >= rows * kPairsPerRow) { ++ return; ++ } ++ int row = pair / kPairsPerRow; ++ int64_t first = int64_t(pair) * 2; ++ float inverse = 1.0f / prefix_sum[row]; ++ output[first] = __float2half_rn(prefix_accumulator[first] * inverse); ++ output[first + 1] = ++ __float2half_rn(prefix_accumulator[first + 1] * inverse); ++} ++ ++__global__ void merge_prefix_accumulator_tail( ++ float const* prefix_accumulator, ++ float const* prefix_max, ++ float const* prefix_sum, ++ __half const* tail_output, ++ float const* tail_max, ++ float const* tail_sum, ++ __half* output, ++ int rows) { ++ int row = blockIdx.x; ++ int d = threadIdx.x; ++ __shared__ float masses[3]; ++ if (d == 0) { ++ constexpr float kLog2E = 1.4426950408889634f; ++ float global_max = fmaxf(prefix_max[row], tail_max[row]); ++ masses[0] = exp2f((prefix_max[row] - global_max) * kLog2E); ++ masses[1] = tail_sum[row] ++ * exp2f((tail_max[row] - global_max) * kLog2E); ++ masses[2] = 1.0f / (prefix_sum[row] * masses[0] + masses[1]); ++ } ++ __syncthreads(); ++ int64_t element = int64_t(row) * 256 + d; ++ float numerator = prefix_accumulator[element] * masses[0] ++ + float(tail_output[element]) * masses[1]; ++ output[element] = __float2half_rn(numerator * masses[2]); ++} ++ ++__global__ void merge_prefix_tail( ++ __half const* prefix_output, ++ float const* prefix_max, ++ float const* prefix_sum, ++ __half const* tail_output, ++ float const* tail_max, ++ float const* tail_sum, ++ __half* output, ++ int rows) { ++ int row = blockIdx.x; ++ int d = threadIdx.x; ++ __shared__ float masses[3]; ++ if (d == 0) { ++ float global_max = fmaxf(prefix_max[row], tail_max[row]); ++ float prefix_scale = exp2f( ++ (prefix_max[row] - global_max) * 1.4426950408889634f); ++ float tail_scale = exp2f( ++ (tail_max[row] - global_max) * 1.4426950408889634f); ++ masses[0] = prefix_sum[row] * prefix_scale; ++ masses[1] = tail_sum[row] * tail_scale; ++ masses[2] = 1.0f / (masses[0] + masses[1]); ++ } ++ __syncthreads(); ++ int64_t element = int64_t(row) * 256 + d; ++ float numerator = float(prefix_output[element]) * masses[0] ++ + float(tail_output[element]) * masses[1]; ++ output[element] = __float2half_rn(numerator * masses[2]); ++} ++ ++using DenseTailRaw = cudaError_t (*)( ++ const void*, ++ const void*, ++ const void*, ++ float*, ++ float*, ++ void*, ++ int, ++ int, ++ int, ++ int, ++ float, ++ cudaStream_t); ++ ++__global__ void fill_pattern(__half* data, size_t elements, uint32_t seed) { ++ size_t index = size_t(blockIdx.x) * blockDim.x + threadIdx.x; ++ if (index >= elements) { ++ return; ++ } ++ uint32_t value = uint32_t(index) ^ seed; ++ value ^= value >> 16; ++ value *= 0x7feb352du; ++ value ^= value >> 15; ++ value *= 0x846ca68bu; ++ value ^= value >> 16; ++ float uniform = float(value & 0xffffu) * (2.0f / 65535.0f) - 1.0f; ++ data[index] = __float2half_rn(uniform); ++} ++ ++struct BlockOperators { ++ int width; ++ QKGemm qk; ++ std::unique_ptr pv; ++}; ++ ++} // namespace ++ ++#if !defined(PREFIX_TORCH_EXTENSION) ++int main(int argc, char** argv) { ++ int block_n = argc > 1 ? std::atoi(argv[1]) : 2048; ++ int iterations = argc > 2 ? std::atoi(argv[2]) : 5; ++#if defined(PREFIX_FULL_ENDPOINT) ++ if (argc <= 3) { ++ std::cerr << "full endpoint requires the exact-tail extension path\n"; ++ return EXIT_FAILURE; ++ } ++#endif ++ constexpr int kRows = 48000; ++ constexpr int kHeadDim = 256; ++ constexpr int kPrefix = 120000; ++ constexpr int kTail = 8000; ++ constexpr int kTotalKV = kPrefix + kTail; ++#if defined(PREFIX_FULL_ENDPOINT) ++ constexpr int kStoredKV = kTotalKV; ++#else ++ constexpr int kStoredKV = kPrefix; ++#endif ++ int blocks = (kPrefix + block_n - 1) / block_n; ++ ++ Element* query = nullptr; ++ Element* key = nullptr; ++ Element* value = nullptr; ++ Element* scores = nullptr; ++ Element* partials = nullptr; ++ Element* output = nullptr; ++ float* qk_norm = nullptr; ++ float* qk_sum = nullptr; ++ float* prefix_accumulator = nullptr; ++ float* prefix_max = nullptr; ++ float* prefix_sum = nullptr; ++ float* old_scales = nullptr; ++ float* block_scales = nullptr; ++#if defined(PREFIX_FULL_ENDPOINT) ++ Element* tail_output = nullptr; ++ Element* final_output = nullptr; ++ Element* exact_output = nullptr; ++ float* tail_max = nullptr; ++ float* tail_sum = nullptr; ++#endif ++ ++ int qk_tiles_n = (block_n + QKThreadblockShape::kN - 1) / ++ QKThreadblockShape::kN; ++ size_t partial_elements = size_t(kRows) * kHeadDim; ++ check(cudaMalloc(&query, size_t(kRows) * kHeadDim * sizeof(Element)), "allocate Q"); ++ check(cudaMalloc(&key, size_t(kStoredKV) * kHeadDim * sizeof(Element)), "allocate K"); ++ check(cudaMalloc(&value, size_t(kStoredKV) * kHeadDim * sizeof(Element)), "allocate V"); ++ check(cudaMalloc(&scores, size_t(kRows) * block_n * sizeof(Element)), "allocate scores"); ++ check(cudaMalloc(&partials, partial_elements * sizeof(Element)), "allocate partials"); ++ check(cudaMalloc(&output, size_t(kRows) * kHeadDim * sizeof(Element)), "allocate output"); ++ check(cudaMalloc(&qk_norm, size_t(qk_tiles_n) * kRows * sizeof(float)), "allocate QK max"); ++ check(cudaMalloc(&qk_sum, size_t(qk_tiles_n) * kRows * sizeof(float)), "allocate QK sum"); ++ check(cudaMalloc( ++ &prefix_accumulator, ++ size_t(kRows) * kHeadDim * sizeof(float)), ++ "allocate prefix accumulator"); ++ check(cudaMalloc(&prefix_max, size_t(kRows) * sizeof(float)), "allocate prefix max"); ++ check(cudaMalloc(&prefix_sum, size_t(kRows) * sizeof(float)), "allocate prefix sum"); ++ check(cudaMalloc(&old_scales, size_t(kRows) * sizeof(float)), ++ "allocate old scales"); ++ check(cudaMalloc(&block_scales, size_t(kRows) * sizeof(float)), ++ "allocate block scales"); ++#if defined(PREFIX_FULL_ENDPOINT) ++ check(cudaMalloc(&tail_output, size_t(kRows) * kHeadDim * sizeof(Element)), ++ "allocate tail output"); ++ check(cudaMalloc(&final_output, size_t(kRows) * kHeadDim * sizeof(Element)), ++ "allocate final output"); ++ check(cudaMalloc(&exact_output, size_t(kRows) * kHeadDim * sizeof(Element)), ++ "allocate exact output"); ++ check(cudaMalloc(&tail_max, size_t(kRows) * sizeof(float)), "allocate tail max"); ++ check(cudaMalloc(&tail_sum, size_t(kRows) * sizeof(float)), "allocate tail sum"); ++#endif ++ size_t query_elements = size_t(kRows) * kHeadDim; ++ size_t kv_elements = size_t(kStoredKV) * kHeadDim; ++ fill_pattern<<<(query_elements + 255) / 256, 256>>>( ++ reinterpret_cast<__half*>(query), query_elements, 0x12345678u); ++ fill_pattern<<<(kv_elements + 255) / 256, 256>>>( ++ reinterpret_cast<__half*>(key), kv_elements, 0x9abcdef0u); ++ fill_pattern<<<(kv_elements + 255) / 256, 256>>>( ++ reinterpret_cast<__half*>(value), kv_elements, 0x31415926u); ++ check(cudaGetLastError(), "initialize deterministic inputs"); ++ ++#if defined(PREFIX_FULL_ENDPOINT) ++ void* tail_handle = dlopen(argv[3], RTLD_NOW | RTLD_LOCAL); ++ if (tail_handle == nullptr) { ++ std::cerr << "load exact-tail extension: " << dlerror() << "\n"; ++ return EXIT_FAILURE; ++ } ++ dlerror(); ++ auto tail_raw = reinterpret_cast( ++ dlsym(tail_handle, "onecat_sm70_d256_dense_state_raw")); ++ char const* symbol_error = dlerror(); ++ if (symbol_error != nullptr || tail_raw == nullptr) { ++ std::cerr << "load exact-tail entry point: " ++ << (symbol_error == nullptr ? "missing symbol" : symbol_error) ++ << "\n"; ++ return EXIT_FAILURE; ++ } ++#endif ++ ++ std::vector operators; ++ operators.reserve(blocks); ++ for (int block = 0; block < blocks; ++block) { ++ int begin = block * block_n; ++ int width = std::min(block_n, kPrefix - begin); ++ BlockOperators operation; ++ operation.width = width; ++ typename QKGemm::Arguments arguments( ++ {kRows, width, kHeadDim}, ++ 1, ++ {query, QKLayoutA(kHeadDim)}, ++ {key + size_t(begin) * kHeadDim, QKLayoutB(kHeadDim)}, ++ {scores, typename QKGemm::LayoutC(width)}, ++ {scores, typename QKGemm::LayoutC(width)}, ++ {Element(0.0625f), Element(0.0f)}, ++ {qk_norm, typename QKGemm::LayoutN(kRows)}, ++ {qk_sum, typename QKGemm::LayoutS(kRows)}, ++ {scores, typename QKGemm::LayoutSoft(width)}); ++ check(operation.qk.initialize(arguments), "initialize QK"); ++ operation.pv = std::make_unique( ++ scores, ++ value + size_t(begin) * kHeadDim, ++ partials, ++ kRows, ++ width); ++ operators.push_back(std::move(operation)); ++ } ++ ++ check(cudaMemcpyToSymbol(g_rows, &kRows, sizeof(kRows)), "set row count"); ++ float const* persistent_max_ptr = qk_norm; ++ check(cudaMemcpyToSymbol( ++ g_row_max, ++ &persistent_max_ptr, ++ sizeof(persistent_max_ptr)), ++ "set persistent PV max pointer"); ++#if defined(PREFIX_QK_FULL_STATS) ++ float const* persistent_inv_sum_ptr = qk_sum; ++ check(cudaMemcpyToSymbol( ++ g_row_inv_sum, ++ &persistent_inv_sum_ptr, ++ sizeof(persistent_inv_sum_ptr)), ++ "set persistent PV inverse-sum pointer"); ++#endif ++ cudaStream_t stream = nullptr; ++ auto launch = [&](cudaEvent_t blocks_done, ++ cudaEvent_t prefix_done, ++ cudaEvent_t tail_done) { ++ for (int block = 0; block < blocks; ++block) { ++ check(operators[block].qk(stream), "launch QK max"); ++#if !defined(PREFIX_QK_FULL_STATS) ++ float* sum_ptr = block_sum + size_t(block) * kRows; ++ check(cudaMemcpyToSymbolAsync( ++ g_row_sum_out, ++ &sum_ptr, ++ sizeof(sum_ptr), ++ 0, ++ cudaMemcpyHostToDevice, ++ stream), ++ "set PV sum pointer"); ++#endif ++ operators[block].pv->launch(stream); ++ prepare_prefix_update<<<(kRows + 255) / 256, 256, 0, stream>>>( ++ qk_norm, ++ qk_sum, ++ prefix_max, ++ prefix_sum, ++ old_scales, ++ block_scales, ++ kRows, ++ block == 0); ++ int pairs = kRows * kHeadDim / 2; ++ apply_prefix_update_half2<<<(pairs + 255) / 256, 256, 0, stream>>>( ++ reinterpret_cast<__half const*>(partials), ++ prefix_accumulator, ++ old_scales, ++ block_scales, ++ kRows, ++ block == 0); ++ } ++ if (blocks_done != nullptr) { ++ check(cudaEventRecord(blocks_done, stream), "record QK/PV blocks done"); ++ } ++#if !defined(PREFIX_FULL_ENDPOINT) ++ int pairs = kRows * kHeadDim / 2; ++ finalize_prefix_accumulator<<<(pairs + 255) / 256, 256, 0, stream>>>( ++ prefix_accumulator, ++ prefix_sum, ++ reinterpret_cast<__half*>(output), ++ kRows); ++#endif ++ if (prefix_done != nullptr) { ++ check(cudaEventRecord(prefix_done, stream), "record prefix merge done"); ++ } ++#if defined(PREFIX_FULL_ENDPOINT) ++ check(tail_raw( ++ query, ++ key + size_t(kPrefix) * kHeadDim, ++ value + size_t(kPrefix) * kHeadDim, ++ tail_max, ++ tail_sum, ++ tail_output, ++ kTail, ++ kTail, ++ 6, ++ 1, ++ 0.0625f, ++ stream), ++ "launch exact causal tail"); ++ if (tail_done != nullptr) { ++ check(cudaEventRecord(tail_done, stream), "record causal tail done"); ++ } ++ merge_prefix_accumulator_tail<<>>( ++ prefix_accumulator, ++ prefix_max, ++ prefix_sum, ++ reinterpret_cast<__half const*>(tail_output), ++ tail_max, ++ tail_sum, ++ reinterpret_cast<__half*>(final_output), ++ kRows); ++#endif ++ }; ++ ++ launch(nullptr, nullptr, nullptr); ++ check(cudaDeviceSynchronize(), "warmup synchronize"); ++ cudaEvent_t start; ++ cudaEvent_t stop; ++ check(cudaEventCreate(&start), "create start"); ++ check(cudaEventCreate(&stop), "create stop"); ++ check(cudaEventRecord(start), "record start"); ++ for (int iteration = 0; iteration < iterations; ++iteration) { ++ launch(nullptr, nullptr, nullptr); ++ } ++ check(cudaEventRecord(stop), "record stop"); ++ check(cudaEventSynchronize(stop), "synchronize stop"); ++ float elapsed_ms = 0.0f; ++ check(cudaEventElapsedTime(&elapsed_ms, start, stop), "elapsed time"); ++ float endpoint_ms = elapsed_ms / float(iterations); ++ constexpr double kFrozenUsefulTflop = 60.0 * 0.101581; ++#if defined(PREFIX_FULL_ENDPOINT) ++ cudaEvent_t phase_start; ++ cudaEvent_t blocks_done; ++ cudaEvent_t prefix_done; ++ cudaEvent_t tail_done; ++ cudaEvent_t phase_stop; ++ check(cudaEventCreate(&phase_start), "create phase start"); ++ check(cudaEventCreate(&blocks_done), "create blocks done"); ++ check(cudaEventCreate(&prefix_done), "create prefix done"); ++ check(cudaEventCreate(&tail_done), "create tail done"); ++ check(cudaEventCreate(&phase_stop), "create phase stop"); ++ check(cudaEventRecord(phase_start, stream), "record phase start"); ++ launch(blocks_done, prefix_done, tail_done); ++ check(cudaEventRecord(phase_stop, stream), "record phase stop"); ++ check(cudaEventSynchronize(phase_stop), "synchronize phase stop"); ++ float blocks_ms = 0.0f; ++ float prefix_merge_ms = 0.0f; ++ float tail_ms = 0.0f; ++ float final_merge_ms = 0.0f; ++ check(cudaEventElapsedTime(&blocks_ms, phase_start, blocks_done), ++ "measure QK/PV blocks"); ++ check(cudaEventElapsedTime(&prefix_merge_ms, blocks_done, prefix_done), ++ "measure prefix merge"); ++ check(cudaEventElapsedTime(&tail_ms, prefix_done, tail_done), ++ "measure causal tail"); ++ check(cudaEventElapsedTime(&final_merge_ms, tail_done, phase_stop), ++ "measure final merge"); ++ check(tail_raw( ++ query, ++ key, ++ value, ++ tail_max, ++ tail_sum, ++ exact_output, ++ kTail, ++ kTotalKV, ++ 6, ++ 1, ++ 0.0625f, ++ stream), ++ "launch monolithic exact reference"); ++ check(cudaDeviceSynchronize(), "synchronize exact reference"); ++ size_t output_elements = size_t(kRows) * kHeadDim; ++ std::vector<__half> candidate_host(output_elements); ++ std::vector<__half> reference_host(output_elements); ++ check(cudaMemcpy( ++ candidate_host.data(), ++ final_output, ++ output_elements * sizeof(__half), ++ cudaMemcpyDeviceToHost), ++ "copy candidate output"); ++ check(cudaMemcpy( ++ reference_host.data(), ++ exact_output, ++ output_elements * sizeof(__half), ++ cudaMemcpyDeviceToHost), ++ "copy exact output"); ++ double squared_error = 0.0; ++ double squared_reference = 0.0; ++ double absolute_error = 0.0; ++ float max_absolute_error = 0.0f; ++ size_t equal_elements = 0; ++ for (size_t index = 0; index < output_elements; ++index) { ++ float candidate_value = __half2float(candidate_host[index]); ++ float reference_value = __half2float(reference_host[index]); ++ float difference = candidate_value - reference_value; ++ float absolute = std::fabs(difference); ++ squared_error += double(difference) * difference; ++ squared_reference += double(reference_value) * reference_value; ++ absolute_error += absolute; ++ max_absolute_error = std::max(max_absolute_error, absolute); ++ equal_elements += candidate_value == reference_value; ++ } ++ double relative_l2 = std::sqrt(squared_error / squared_reference); ++ double mean_absolute_error = absolute_error / double(output_elements); ++ double measured_tflops = kFrozenUsefulTflop / (endpoint_ms * 1.0e-3); ++ std::cout << "{\n" ++ << " \"block_n\": " << block_n << ",\n" ++ << " \"blocks\": " << blocks << ",\n" ++ << " \"full_attention_ms\": " << endpoint_ms << ",\n" ++ << " \"useful_causal_tflops\": " << measured_tflops << ",\n" ++ << " \"prefix_tokens\": " << kPrefix << ",\n" ++ << " \"causal_tail_tokens\": " << kTail << ",\n" ++ << " \"phases_ms\": {\n" ++ << " \"prefix_qk_pv_blocks\": " << blocks_ms << ",\n" ++ << " \"prefix_state_and_partial_merge\": " ++ << prefix_merge_ms << ",\n" ++ << " \"exact_causal_tail\": " << tail_ms << ",\n" ++ << " \"prefix_tail_merge\": " << final_merge_ms << "\n" ++ << " },\n" ++ << " \"numerical\": {\n" ++ << " \"max_abs_diff\": " << max_absolute_error << ",\n" ++ << " \"mean_abs_diff\": " << mean_absolute_error << ",\n" ++ << " \"relative_l2\": " << relative_l2 << ",\n" ++ << " \"equal_elements\": " << equal_elements << ",\n" ++ << " \"elements\": " << output_elements << "\n" ++ << " },\n" ++ << " \"partial_storage_gib\": " ++ << double(partial_elements * sizeof(Element)) / double(1ull << 30) ++ << ",\n" ++ << " \"note\": \"measured serial prefix + exact tail state + final merge endpoint\"\n" ++ << "}\n"; ++ check(cudaEventDestroy(phase_start), "destroy phase start"); ++ check(cudaEventDestroy(blocks_done), "destroy blocks done"); ++ check(cudaEventDestroy(prefix_done), "destroy prefix done"); ++ check(cudaEventDestroy(tail_done), "destroy tail done"); ++ check(cudaEventDestroy(phase_stop), "destroy phase stop"); ++#else ++ double projected_full_ms = endpoint_ms + 4.83; ++ double projected_tflops = kFrozenUsefulTflop / (projected_full_ms * 1.0e-3); ++ std::cout << "{\n" ++ << " \"block_n\": " << block_n << ",\n" ++ << " \"blocks\": " << blocks << ",\n" ++ << " \"prefix_ms_including_merge\": " << endpoint_ms << ",\n" ++ << " \"accepted_tail_ms\": 4.83,\n" ++ << " \"projected_full_ms\": " << projected_full_ms << ",\n" ++ << " \"projected_useful_tflops\": " << projected_tflops << ",\n" ++ << " \"partial_storage_gib\": " ++ << double(partial_elements * sizeof(Element)) / double(1ull << 30) ++ << ",\n" ++ << " \"note\": \"prefix endpoint; prefix/tail state merge excluded\"\n" ++ << "}\n"; ++#endif ++ ++ check(cudaEventDestroy(start), "destroy start"); ++ check(cudaEventDestroy(stop), "destroy stop"); ++ check(cudaFree(query), "free Q"); ++ check(cudaFree(key), "free K"); ++ check(cudaFree(value), "free V"); ++ check(cudaFree(scores), "free scores"); ++ check(cudaFree(partials), "free partials"); ++ check(cudaFree(output), "free output"); ++ check(cudaFree(qk_norm), "free QK max"); ++ check(cudaFree(qk_sum), "free QK sum"); ++ check(cudaFree(prefix_accumulator), "free prefix accumulator"); ++ check(cudaFree(prefix_max), "free prefix max"); ++ check(cudaFree(prefix_sum), "free prefix sum"); ++ check(cudaFree(old_scales), "free old scales"); ++ check(cudaFree(block_scales), "free block scales"); ++#if defined(PREFIX_FULL_ENDPOINT) ++ check(cudaFree(tail_output), "free tail output"); ++ check(cudaFree(final_output), "free final output"); ++ check(cudaFree(exact_output), "free exact output"); ++ check(cudaFree(tail_max), "free tail max"); ++ check(cudaFree(tail_sum), "free tail sum"); ++ dlclose(tail_handle); ++#endif ++ return 0; ++} ++#else ++ ++extern "C" cudaError_t onecat_sm70_d256_dense_state_raw( ++ const void*, ++ const void*, ++ const void*, ++ float*, ++ float*, ++ void*, ++ int, ++ int, ++ int, ++ int, ++ float, ++ cudaStream_t); ++ ++namespace FLASH_NAMESPACE { ++ ++struct Sm70GqaScoreWorkspace { ++ at::Tensor scores; ++ cudaEvent_t completion = nullptr; ++ bool completion_recorded = false; ++ std::mutex launch_mutex; ++ ++ Sm70GqaScoreWorkspace(const at::Tensor& q, int rows, int block_n) ++ : scores(at::empty({rows, block_n}, q.options())) { ++ C10_CUDA_CHECK( ++ cudaEventCreateWithFlags(&completion, cudaEventDisableTiming)); ++ } ++ ++ ~Sm70GqaScoreWorkspace() { ++ if (completion != nullptr) { ++ cudaEventDestroy(completion); ++ } ++ } ++}; ++ ++using Sm70GqaScoreWorkspacePtr = std::shared_ptr; ++ ++std::mutex& sm70_gqa_score_cache_mutex() { ++ static std::mutex cache_mutex; ++ return cache_mutex; ++} ++ ++std::map& sm70_gqa_score_cache() { ++ static std::map cache; ++ return cache; ++} ++ ++Sm70GqaScoreWorkspacePtr get_sm70_gqa_score_workspace( ++ const at::Tensor& q, ++ int rows, ++ int block_n) { ++ std::lock_guard lock(sm70_gqa_score_cache_mutex()); ++ int device = q.get_device(); ++ auto& cache = sm70_gqa_score_cache(); ++ auto& workspace = cache[device]; ++ if (!workspace) { ++ workspace = ++ std::make_shared(q, rows, block_n); ++ } ++ return workspace; ++} ++ ++Sm70GqaScoreWorkspacePtr find_sm70_gqa_score_workspace(int device) { ++ std::lock_guard lock(sm70_gqa_score_cache_mutex()); ++ auto& cache = sm70_gqa_score_cache(); ++ auto found = cache.find(device); ++ return found == cache.end() ? nullptr : found->second; ++} ++ ++at::Tensor sm70_d256_gqa_architecture_fwd( ++ const at::Tensor& q, ++ const at::Tensor& k, ++ const at::Tensor& v, ++ at::Tensor& out, ++ double softmax_scale, ++ bool causal) { ++ constexpr int kQuery = 8000; ++ constexpr int kRows = 48000; ++ constexpr int kHeadDim = 256; ++ constexpr int kHeadsQ = 6; ++ constexpr int kHeadsKV = 1; ++ constexpr int kTail = 8000; ++ constexpr int kMinTotalKV = 40000; ++ constexpr int kMaxTotalKV = 128000; ++ constexpr int kTotalKVStep = 8000; ++ constexpr int kBlockN = 8192; ++ ++ TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda() && out.is_cuda(), ++ "SM70 GQA architecture requires CUDA tensors"); ++ TORCH_CHECK(q.scalar_type() == at::ScalarType::Half ++ && k.scalar_type() == q.scalar_type() ++ && v.scalar_type() == q.scalar_type() ++ && out.scalar_type() == q.scalar_type(), ++ "SM70 GQA architecture requires FP16 q, k, v, and out"); ++ TORCH_CHECK(q.sizes() == at::IntArrayRef({1, kQuery, kHeadsQ, kHeadDim}) ++ && k.dim() == 4 && k.size(0) == 1 ++ && k.size(2) == kHeadsKV && k.size(3) == kHeadDim ++ && v.sizes() == k.sizes() && out.sizes() == q.sizes(), ++ "SM70 GQA architecture only accepts the validated " ++ "Q8000/Hq6/Hkv1/D256 dense shape family"); ++ 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] " ++ "with an 8000-token step, got ", ++ total_kv); ++ const int prefix = total_kv - kTail; ++ const int blocks = (prefix + kBlockN - 1) / kBlockN; ++ TORCH_CHECK(q.is_contiguous() && k.is_contiguous() && v.is_contiguous() ++ && out.is_contiguous(), ++ "SM70 GQA architecture requires contiguous tensors"); ++ TORCH_CHECK(q.get_device() == k.get_device() ++ && q.get_device() == v.get_device() ++ && q.get_device() == out.get_device(), ++ "SM70 GQA architecture tensors must share one device"); ++ TORCH_CHECK(causal, "SM70 GQA architecture requires causal attention"); ++ TORCH_CHECK(std::abs(softmax_scale - 0.0625) < 1.0e-8, ++ "SM70 GQA architecture requires D256 softmax scale 1/16"); ++ ++ const at::cuda::OptionalCUDAGuard device_guard(q.device()); ++ const int qk_tiles_n = ++ (kBlockN + QKThreadblockShape::kN - 1) ++ / QKThreadblockShape::kN; ++ cudaStream_t stream = at::cuda::getCurrentCUDAStream(); ++ ++ constexpr size_t kScratchAlignment = 256; ++ size_t scratch_bytes = 0; ++ auto reserve_scratch = [&](size_t bytes) { ++ scratch_bytes = ++ (scratch_bytes + kScratchAlignment - 1) & ~(kScratchAlignment - 1); ++ size_t offset = scratch_bytes; ++ scratch_bytes += bytes; ++ return offset; ++ }; ++ size_t partial_offset = reserve_scratch( ++ size_t(kRows) * kHeadDim * sizeof(Element)); ++ size_t prefix_accumulator_offset = reserve_scratch( ++ size_t(kRows) * kHeadDim * sizeof(float)); ++ size_t tail_output_offset = reserve_scratch( ++ size_t(kRows) * kHeadDim * sizeof(Element)); ++ size_t qk_norm_offset = reserve_scratch( ++ size_t(qk_tiles_n) * kRows * sizeof(float)); ++ size_t qk_sum_offset = reserve_scratch( ++ size_t(qk_tiles_n) * kRows * sizeof(float)); ++ size_t prefix_max_offset = ++ reserve_scratch(size_t(kRows) * sizeof(float)); ++ size_t prefix_sum_offset = ++ reserve_scratch(size_t(kRows) * sizeof(float)); ++ size_t old_scale_offset = ++ reserve_scratch(size_t(kRows) * sizeof(float)); ++ size_t block_scale_offset = ++ reserve_scratch(size_t(kRows) * sizeof(float)); ++ size_t tail_max_offset = ++ reserve_scratch(size_t(kRows) * sizeof(float)); ++ size_t tail_sum_offset = ++ reserve_scratch(size_t(kRows) * sizeof(float)); ++ int device = q.get_device(); ++ auto score_workspace = find_sm70_gqa_score_workspace(device); ++ if (!score_workspace) { ++ constexpr size_t kRequiredPostWorkspaceHeadroom = 128 * 1024 * 1024; ++ constexpr size_t kScoreBytes = ++ size_t(kRows) * kBlockN * sizeof(Element); ++ size_t free_bytes = 0; ++ size_t total_bytes = 0; ++ C10_CUDA_CHECK(cudaMemGetInfo(&free_bytes, &total_bytes)); ++ size_t required_bytes = ++ kScoreBytes + scratch_bytes + kRequiredPostWorkspaceHeadroom; ++ TORCH_CHECK_WITH( ++ OutOfMemoryError, ++ free_bytes >= required_bytes, ++ "SM70 GQA architecture requires ", ++ required_bytes / (1024 * 1024), ++ " MiB free before its first workspace allocation, including " ++ "128 MiB of downstream headroom, but only ", ++ free_bytes / (1024 * 1024), ++ " MiB remains out of ", ++ total_bytes / (1024 * 1024), ++ " MiB"); ++ score_workspace = ++ get_sm70_gqa_score_workspace(q, kRows, kBlockN); ++ } ++ std::unique_lock launch_lock(score_workspace->launch_mutex); ++ if (score_workspace->completion_recorded) { ++ C10_CUDA_CHECK( ++ cudaStreamWaitEvent(stream, score_workspace->completion, 0)); ++ } ++ at::Tensor scratch = at::empty( ++ {static_cast(scratch_bytes)}, ++ q.options().dtype(at::ScalarType::Byte)); ++ auto* scratch_base = scratch.data_ptr(); ++ ++ auto* query = reinterpret_cast(q.data_ptr()); ++ auto* key = reinterpret_cast(k.data_ptr()); ++ auto* value = reinterpret_cast(v.data_ptr()); ++ auto* score_ptr = ++ reinterpret_cast(score_workspace->scores.data_ptr()); ++ auto* partial_ptr = ++ reinterpret_cast(scratch_base + partial_offset); ++ float* prefix_accumulator_ptr = ++ reinterpret_cast(scratch_base + prefix_accumulator_offset); ++ auto* tail_output_ptr = ++ reinterpret_cast(scratch_base + tail_output_offset); ++ auto* output_ptr = reinterpret_cast(out.data_ptr()); ++ float* qk_norm_ptr = ++ reinterpret_cast(scratch_base + qk_norm_offset); ++ float* qk_sum_ptr = ++ reinterpret_cast(scratch_base + qk_sum_offset); ++ float* prefix_max_ptr = ++ reinterpret_cast(scratch_base + prefix_max_offset); ++ float* prefix_sum_ptr = ++ reinterpret_cast(scratch_base + prefix_sum_offset); ++ float* old_scale_ptr = ++ reinterpret_cast(scratch_base + old_scale_offset); ++ float* block_scale_ptr = ++ reinterpret_cast(scratch_base + block_scale_offset); ++ float* tail_max_ptr = ++ reinterpret_cast(scratch_base + tail_max_offset); ++ float* tail_sum_ptr = ++ reinterpret_cast(scratch_base + tail_sum_offset); ++ ++ C10_CUDA_CHECK(onecat_sm70_d256_dense_state_raw( ++ query, ++ key + size_t(prefix) * kHeadDim, ++ value + size_t(prefix) * kHeadDim, ++ tail_max_ptr, ++ tail_sum_ptr, ++ tail_output_ptr, ++ kTail, ++ kTail, ++ kHeadsQ, ++ kHeadsKV, ++ static_cast(softmax_scale), ++ stream)); ++ C10_CUDA_CHECK(cudaMemcpyToSymbolAsync( ++ g_rows, ++ &kRows, ++ sizeof(kRows), ++ 0, ++ cudaMemcpyHostToDevice, ++ stream)); ++ float const* persistent_max_ptr = qk_norm_ptr; ++ C10_CUDA_CHECK(cudaMemcpyToSymbolAsync( ++ g_row_max, ++ &persistent_max_ptr, ++ sizeof(persistent_max_ptr), ++ 0, ++ cudaMemcpyHostToDevice, ++ stream)); ++ float const* persistent_inv_sum_ptr = qk_sum_ptr; ++ C10_CUDA_CHECK(cudaMemcpyToSymbolAsync( ++ g_row_inv_sum, ++ &persistent_inv_sum_ptr, ++ sizeof(persistent_inv_sum_ptr), ++ 0, ++ cudaMemcpyHostToDevice, ++ stream)); ++ ++ for (int block = 0; block < blocks; ++block) { ++ int begin = block * kBlockN; ++ int width = std::min(kBlockN, prefix - begin); ++ BlockOperators operation; ++ operation.width = width; ++ typename QKGemm::Arguments arguments( ++ {kRows, width, kHeadDim}, ++ 1, ++ {query, QKLayoutA(kHeadDim)}, ++ {key + size_t(begin) * kHeadDim, QKLayoutB(kHeadDim)}, ++ {score_ptr, typename QKGemm::LayoutC(width)}, ++ {score_ptr, typename QKGemm::LayoutC(width)}, ++ {Element(static_cast(softmax_scale)), Element(0.0f)}, ++ {qk_norm_ptr, typename QKGemm::LayoutN(kRows)}, ++ {qk_sum_ptr, typename QKGemm::LayoutS(kRows)}, ++ {score_ptr, typename QKGemm::LayoutSoft(width)}); ++ TORCH_CHECK(operation.qk.initialize(arguments) == cutlass::Status::kSuccess, ++ "initialize SM70 GQA QK block ", block, " failed"); ++ operation.pv = std::make_unique( ++ score_ptr, ++ value + size_t(begin) * kHeadDim, ++ partial_ptr, ++ kRows, ++ width); ++ TORCH_CHECK(operation.qk(stream) == cutlass::Status::kSuccess, ++ "launch SM70 GQA QK block ", block, " failed"); ++ operation.pv->launch(stream); ++ prepare_prefix_update<<<(kRows + 255) / 256, 256, 0, stream>>>( ++ qk_norm_ptr, ++ qk_sum_ptr, ++ prefix_max_ptr, ++ prefix_sum_ptr, ++ old_scale_ptr, ++ block_scale_ptr, ++ kRows, ++ block == 0); ++ constexpr int kPairs = kRows * kHeadDim / 2; ++ apply_prefix_update_half2<<<(kPairs + 255) / 256, 256, 0, stream>>>( ++ reinterpret_cast<__half const*>(partial_ptr), ++ prefix_accumulator_ptr, ++ old_scale_ptr, ++ block_scale_ptr, ++ kRows, ++ block == 0); ++ } ++ merge_prefix_accumulator_tail<<>>( ++ prefix_accumulator_ptr, ++ prefix_max_ptr, ++ prefix_sum_ptr, ++ reinterpret_cast<__half const*>(tail_output_ptr), ++ tail_max_ptr, ++ tail_sum_ptr, ++ reinterpret_cast<__half*>(output_ptr), ++ kRows); ++ C10_CUDA_KERNEL_LAUNCH_CHECK(); ++ C10_CUDA_CHECK(cudaEventRecord(score_workspace->completion, stream)); ++ score_workspace->completion_recorded = true; ++ return out; ++} ++ ++} // namespace FLASH_NAMESPACE ++ ++#endif +diff --git a/csrc/flash_attn/src/flash_fwd_d256_splitd_sm70.cu b/csrc/flash_attn/src/flash_fwd_d256_splitd_sm70.cu +index 7a0f626..1fd4008 100644 +--- a/csrc/flash_attn/src/flash_fwd_d256_splitd_sm70.cu ++++ b/csrc/flash_attn/src/flash_fwd_d256_splitd_sm70.cu +@@ -398,7 +398,11 @@ __device__ __forceinline__ int64_t paged_kv_thread_offset( + + (tid % kThreadsPerRow) * kElemsPerLoad; + } + +-template ++template < ++ typename Element, ++ bool PagedKV, ++ bool SplitKV3 = false, ++ bool StoreState = false> + __global__ __launch_bounds__(Sm70D256SplitDTraits::kNThreads, 1) + void sm70_d256_splitd_dense_kernel( + const Element *__restrict__ q, +@@ -859,6 +863,29 @@ void sm70_d256_splitd_dense_kernel( + } + } + } else { ++ if constexpr (StoreState) { ++ if ((lane & 0x0e) == 0) { ++#pragma unroll ++ for (int slot = 0; slot < Traits::kQkRowsPerThread; ++slot) { ++ const int row = FLASH_NAMESPACE::sm70_row_slot< ++ Traits::kQkWarpRows>(slot, lane); ++ const int query_row = query_row_base ++ + group_row_base ++ + n_warp * Traits::kQkWarpRows + row; ++ const int64_t state_row = ++ (static_cast(batch) * query_len + query_row) ++ * heads_q ++ + head_q; ++ // The standalone prefix path stores maxima after applying ++ // softmax_scale. Export the dense-tail state in the same ++ // natural-log coordinate so the two states can be merged ++ // directly with exp(prefix_max - global_max). ++ partial_max[state_row] = row_max[slot] ++ * (softmax_scale_log2 * float(M_LN2)); ++ partial_sum[state_row] = row_sum[slot]; ++ } ++ } ++ } + if ((lane & 0x0e) == 0) { + #pragma unroll + for (int slot = 0; slot < Traits::kQkRowsPerThread; ++slot) { +@@ -1042,6 +1069,154 @@ at::Tensor sm70_d256_splitd_dense_fwd( + return out; + } + ++at::Tensor sm70_d256_splitd_dense_state_fwd( ++ const at::Tensor &q, ++ const at::Tensor &k, ++ const at::Tensor &v, ++ at::Tensor &state_max, ++ at::Tensor &state_sum, ++ at::Tensor &out, ++ double softmax_scale, ++ bool causal) { ++ TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda() ++ && state_max.is_cuda() && state_sum.is_cuda() ++ && out.is_cuda(), ++ "dense-state inputs must be CUDA tensors"); ++ TORCH_CHECK(q.scalar_type() == at::ScalarType::Half ++ && k.scalar_type() == q.scalar_type() ++ && v.scalar_type() == q.scalar_type() ++ && out.scalar_type() == q.scalar_type(), ++ "dense-state requires FP16 q, k, v, and out"); ++ TORCH_CHECK(state_max.scalar_type() == at::ScalarType::Float ++ && state_sum.scalar_type() == at::ScalarType::Float, ++ "dense-state statistics must use FP32"); ++ TORCH_CHECK(q.dim() == 4 && k.dim() == 4 && v.dim() == 4 ++ && q.size(0) == k.size(0) && k.sizes() == v.sizes(), ++ "dense-state q, k, and v shapes are invalid"); ++ TORCH_CHECK(q.size(3) == Sm70D256SplitDTraits::kHeadDim ++ && k.size(3) == Sm70D256SplitDTraits::kHeadDim ++ && q.size(2) % k.size(2) == 0, ++ "dense-state requires D256 and valid GQA heads"); ++ TORCH_CHECK(q.size(1) <= k.size(1) ++ && q.size(1) % Sm70D256SplitDTraits::kBlockM == 0 ++ && k.size(1) % Sm70D256SplitDTraits::kBlockN == 0, ++ "dense-state lengths do not match the exact tiles"); ++ TORCH_CHECK(out.sizes() == q.sizes() && out.is_contiguous() ++ && state_max.is_contiguous() && state_sum.is_contiguous(), ++ "dense-state outputs must be contiguous"); ++ const int64_t rows = q.size(0) * q.size(1) * q.size(2); ++ TORCH_CHECK(state_max.numel() == rows && state_sum.numel() == rows, ++ "dense-state statistic shapes are invalid"); ++ TORCH_CHECK(causal, "dense-state requires causal attention"); ++ ++ const at::cuda::OptionalCUDAGuard device_guard(q.device()); ++ const dim3 block(Sm70D256SplitDTraits::kNThreads); ++ const dim3 grid( ++ q.size(1) / Sm70D256SplitDTraits::kBlockM, ++ q.size(0), ++ q.size(2)); ++ auto stream = at::cuda::getCurrentCUDAStream(); ++ auto kernel = sm70_d256_splitd_dense_kernel< ++ cutlass::half_t, false, false, true>; ++ C10_CUDA_CHECK(cudaFuncSetAttribute( ++ kernel, ++ cudaFuncAttributeMaxDynamicSharedMemorySize, ++ Sm70D256SplitDTraits::kSmemBytes)); ++ kernel<<>>( ++ reinterpret_cast(q.data_ptr()), ++ reinterpret_cast(k.data_ptr()), ++ reinterpret_cast(v.data_ptr()), ++ reinterpret_cast(out.data_ptr()), ++ static_cast(q.stride(0)), ++ static_cast(q.stride(1)), ++ static_cast(q.stride(2)), ++ static_cast(k.stride(0)), ++ static_cast(k.stride(1)), ++ static_cast(k.stride(2)), ++ static_cast(v.stride(0)), ++ static_cast(v.stride(1)), ++ static_cast(v.stride(2)), ++ q.size(1), ++ k.size(1), ++ q.size(2), ++ k.size(2), ++ static_cast(softmax_scale * M_LOG2E), ++ nullptr, ++ 0, ++ 0, ++ nullptr, ++ state_max.data_ptr(), ++ state_sum.data_ptr()); ++ C10_CUDA_KERNEL_LAUNCH_CHECK(); ++ return out; ++} ++ ++// Raw fixed-layout entry point for the task-local end-to-end architecture ++// harness. Keeping this in the same cubin guarantees that the measured tail ++// is the exact production candidate kernel rather than a reimplemented proxy. ++extern "C" cudaError_t onecat_sm70_d256_dense_state_raw( ++ const void *q, ++ const void *k, ++ const void *v, ++ float *state_max, ++ float *state_sum, ++ void *out, ++ int query_len, ++ int kv_len, ++ int heads_q, ++ int heads_kv, ++ float softmax_scale, ++ cudaStream_t stream) { ++ if (q == nullptr || k == nullptr || v == nullptr || state_max == nullptr ++ || state_sum == nullptr || out == nullptr || query_len <= 0 ++ || kv_len < query_len || heads_q <= 0 || heads_kv <= 0 ++ || heads_q % heads_kv != 0 ++ || query_len % Sm70D256SplitDTraits::kBlockM != 0 ++ || kv_len % Sm70D256SplitDTraits::kBlockN != 0) { ++ return cudaErrorInvalidValue; ++ } ++ const dim3 block(Sm70D256SplitDTraits::kNThreads); ++ const dim3 grid( ++ query_len / Sm70D256SplitDTraits::kBlockM, ++ 1, ++ heads_q); ++ auto kernel = sm70_d256_splitd_dense_kernel< ++ cutlass::half_t, false, false, true>; ++ cudaError_t result = cudaFuncSetAttribute( ++ kernel, ++ cudaFuncAttributeMaxDynamicSharedMemorySize, ++ Sm70D256SplitDTraits::kSmemBytes); ++ if (result != cudaSuccess) { ++ return result; ++ } ++ kernel<<>>( ++ static_cast(q), ++ static_cast(k), ++ static_cast(v), ++ static_cast(out), ++ query_len * heads_q * Sm70D256SplitDTraits::kHeadDim, ++ heads_q * Sm70D256SplitDTraits::kHeadDim, ++ Sm70D256SplitDTraits::kHeadDim, ++ kv_len * heads_kv * Sm70D256SplitDTraits::kHeadDim, ++ heads_kv * Sm70D256SplitDTraits::kHeadDim, ++ Sm70D256SplitDTraits::kHeadDim, ++ kv_len * heads_kv * Sm70D256SplitDTraits::kHeadDim, ++ heads_kv * Sm70D256SplitDTraits::kHeadDim, ++ Sm70D256SplitDTraits::kHeadDim, ++ query_len, ++ kv_len, ++ heads_q, ++ heads_kv, ++ softmax_scale * float(M_LOG2E), ++ nullptr, ++ 0, ++ 0, ++ nullptr, ++ state_max, ++ state_sum); ++ return cudaPeekAtLastError(); ++} ++ + at::Tensor sm70_d256_splitd_dense_splitkv3_fwd( + const at::Tensor &q, + const at::Tensor &k, +diff --git a/csrc/cutlass/examples/35_gemm_softmax/gemm_with_softmax.h b/csrc/cutlass/examples/35_gemm_softmax/gemm_with_softmax.h +index 31b2b769..2cb4ba74 100644 +--- a/csrc/cutlass/examples/35_gemm_softmax/gemm_with_softmax.h ++++ b/csrc/cutlass/examples/35_gemm_softmax/gemm_with_softmax.h +@@ -623,6 +623,7 @@ public: + return cutlass::Status::kErrorInternal; + } + ++#if !defined(PREFIX_QK_SKIP_APPLY) + // + // Launch the SoftmaxApplyKernel + // +@@ -646,6 +647,7 @@ public: + if (result != cudaSuccess) { + return cutlass::Status::kErrorInternal; + } ++#endif + + return cutlass::Status::kSuccess; + } diff --git a/cmake/patches/sm70_flash_attn_d256_k_pingpong.patch b/cmake/patches/sm70_flash_attn_d256_k_pingpong.patch new file mode 100644 index 0000000000..a215fcd9de --- /dev/null +++ b/cmake/patches/sm70_flash_attn_d256_k_pingpong.patch @@ -0,0 +1,90 @@ +--- a/csrc/flash_attn/src/flash_fwd_d256_splitd_sm70.cu ++++ b/csrc/flash_attn/src/flash_fwd_d256_splitd_sm70.cu +@@ -126,8 +126,16 @@ + Stride<_4, _1>>{})); + + static constexpr int kQElements = size(SmemLayoutQ{}); + static constexpr int kKVElements = size(SmemLayoutKV{}); + static_assert(size(SmemLayoutV{}) == kKVElements); ++ // The second K stage is live only before P is materialized, so it may ++ // alias the beginning of the later P region. Keep the K stages disjoint ++ // and on the same 128-byte bank phase without moving the V/P regions. ++ static constexpr int kKStageElements = 2240; ++ static_assert(kKStageElements >= cosize(SmemLayoutK{})); ++ static_assert( ++ kKStageElements + cosize(SmemLayoutK{}) ++ <= 2 * kKVElements + size(SmemLayoutP{})); + static constexpr int kPElements = size(SmemLayoutP{}); + static constexpr int kExchangeRows = kMmaGroups * kGroupRows; + static constexpr int kTensorSmemBytes = +@@ -521,11 +525,11 @@ + make_gmem_ptr(k + k_batch_offset + head_kv * k_head_stride), + make_shape(kv_len, Int{}), + make_stride(k_row_stride, _1{})); +- auto sK = make_tensor( ++ auto sKFirst = make_tensor( + make_smem_ptr(kv_smem_ptr), typename Traits::SmemLayoutK{}); +- auto tKsKRaw = gmem_k_thread.partition_D(sK); +- auto tKsK = reshape_kv_thread_tensor(tKsKRaw); +- auto tKrKNext = make_fragment_like(tKsK); ++ auto tKsKFirstRaw = gmem_k_thread.partition_D(sKFirst); ++ auto tKsKFirst = reshape_kv_thread_tensor(tKsKFirstRaw); ++ auto tKrKNext = make_fragment_like(tKsKFirst); + auto gKFirst = local_tile( + mK, + Shape, Int>{}, +@@ -543,7 +547,7 @@ + k_row_stride); + tKgKFirst.data() = mK.data() + k_thread_tile_base; + } +- copy_even_tile(gmem_k_copy, tKgKFirst, tKsK); ++ copy_even_tile(gmem_k_copy, tKgKFirst, tKsKFirst); + __syncthreads(); + + for (int n_block = n_block_max; n_block >= n_block_min; --n_block) { +@@ -558,6 +562,11 @@ + + #pragma unroll + for (int d_chunk = 0; d_chunk < Traits::kDChunks; ++d_chunk) { ++ auto sK = make_tensor( ++ make_smem_ptr( ++ kv_smem_ptr ++ + (d_chunk & 1) * Traits::kKStageElements), ++ typename Traits::SmemLayoutK{}); + if (d_chunk + 1 < Traits::kDChunks) { + auto gKNext = local_tile( + mK, +@@ -596,8 +605,15 @@ + acc_s, tSrQ, tSrK, tSsQ, tSsK, qk_tiled_mma, + smem_copy_q, smem_copy_k, smem_thread_q, smem_thread_k); + if (d_chunk + 1 < Traits::kDChunks) { +- __syncthreads(); +- cute::copy(tKrKNext, tKsK); ++ auto sKNext = make_tensor( ++ make_smem_ptr( ++ kv_smem_ptr ++ + ((d_chunk + 1) & 1) * Traits::kKStageElements), ++ typename Traits::SmemLayoutK{}); ++ auto tKsKNextRaw = gmem_k_thread.partition_D(sKNext); ++ auto tKsKNext = ++ reshape_kv_thread_tensor(tKsKNextRaw); ++ cute::copy(tKrKNext, tKsKNext); + __syncthreads(); + } + } +@@ -770,7 +786,14 @@ + } + } + if (n_block > n_block_min) { +- cute::copy(tKrKNext, tKsK); ++ auto sKNextBlock = make_tensor( ++ make_smem_ptr(kv_smem_ptr), ++ typename Traits::SmemLayoutK{}); ++ auto tKsKNextBlockRaw = ++ gmem_k_thread.partition_D(sKNextBlock); ++ auto tKsKNextBlock = ++ reshape_kv_thread_tensor(tKsKNextBlockRaw); ++ cute::copy(tKrKNext, tKsKNextBlock); + __syncthreads(); + } + } diff --git a/csrc/custom_all_reduce.cu b/csrc/custom_all_reduce.cu index c294e2acd4..e55f515ebb 100644 --- a/csrc/custom_all_reduce.cu +++ b/csrc/custom_all_reduce.cu @@ -161,12 +161,12 @@ void sm70_all_reduce_gemma_rms_norm_impl( "SM70 TP4 Gemma RMSNorm prototype residual must be float32."); } else { TORCH_CHECK(residual.scalar_type() == at::ScalarType::Half || - residual.scalar_type() == at::ScalarType::Float, + residual.scalar_type() == at::ScalarType::Float, "SM70 Gemma RMSNorm prototype residual must be float16 or " "float32."); } TORCH_CHECK(weight.scalar_type() == at::ScalarType::Half || - weight.scalar_type() == at::ScalarType::Float, + weight.scalar_type() == at::ScalarType::Float, "SM70 Gemma RMSNorm prototype weight must be float16 or " "float32."); TORCH_CHECK_EQ(inp.dim(), 2); @@ -176,8 +176,9 @@ void sm70_all_reduce_gemma_rms_norm_impl( TORCH_CHECK(normalized_out.sizes() == inp.sizes(), "SM70 Gemma RMSNorm prototype normalized_out shape must match " "inp."); - TORCH_CHECK(residual_out.sizes() == inp.sizes(), - "SM70 Gemma RMSNorm prototype residual_out shape must match inp."); + TORCH_CHECK( + residual_out.sizes() == inp.sizes(), + "SM70 Gemma RMSNorm prototype residual_out shape must match inp."); TORCH_CHECK_EQ(weight.dim(), 1); TORCH_CHECK_EQ(weight.numel(), kHiddenSize); TORCH_CHECK_EQ(residual.get_device(), inp.get_device()); @@ -277,6 +278,83 @@ void sm70_tp4_all_reduce_gemma_rms_norm( reg_buffer_sz_bytes, epsilon); } +void sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather( + fptr_t _fa, torch::Tensor& inp, torch::Tensor& residual, + torch::Tensor& weight, torch::Tensor& normalized_out, + torch::Tensor& residual_out, fptr_t _reg_input_buffer, + fptr_t _reg_output_buffer, int64_t reg_buffer_sz_bytes, double epsilon) { + constexpr int64_t kWorldSize = 4; + constexpr int64_t kHiddenSize = vllm::kSm70GemmaRmsNormHiddenSize; + auto fa = reinterpret_cast(_fa); + const at::cuda::OptionalCUDAGuard device_guard(device_of(inp)); + auto stream = c10::cuda::getCurrentCUDAStream().stream(); + + TORCH_CHECK_EQ(fa->world_size_, kWorldSize); + TORCH_CHECK(fa->fully_connected_); + TORCH_CHECK_EQ(inp.scalar_type(), at::ScalarType::Half); + TORCH_CHECK_EQ(residual.scalar_type(), at::ScalarType::Float); + TORCH_CHECK(weight.scalar_type() == at::ScalarType::Half || + weight.scalar_type() == at::ScalarType::Float); + TORCH_CHECK_EQ(normalized_out.scalar_type(), at::ScalarType::Half); + TORCH_CHECK_EQ(residual_out.scalar_type(), at::ScalarType::Float); + TORCH_CHECK_EQ(inp.dim(), 2); + TORCH_CHECK_EQ(inp.size(1), kHiddenSize); + TORCH_CHECK_EQ(inp.size(0) % kWorldSize, 0); + TORCH_CHECK_LE(inp.size(0) / kWorldSize, + vllm::kSm70LongPrefillMaxTokensPerRank); + TORCH_CHECK(residual.sizes() == inp.sizes()); + TORCH_CHECK(normalized_out.sizes() == inp.sizes()); + TORCH_CHECK(residual_out.sizes() == inp.sizes()); + TORCH_CHECK_EQ(weight.dim(), 1); + TORCH_CHECK_EQ(weight.numel(), kHiddenSize); + TORCH_CHECK_EQ(residual.get_device(), inp.get_device()); + TORCH_CHECK_EQ(weight.get_device(), inp.get_device()); + TORCH_CHECK_EQ(normalized_out.get_device(), inp.get_device()); + TORCH_CHECK_EQ(residual_out.get_device(), inp.get_device()); + TORCH_CHECK(inp.is_contiguous()); + TORCH_CHECK(residual.is_contiguous()); + TORCH_CHECK(weight.is_contiguous()); + TORCH_CHECK(normalized_out.is_contiguous()); + TORCH_CHECK(residual_out.is_contiguous()); + + auto* reg_input_buffer = reinterpret_cast(_reg_input_buffer); + auto* reg_output_buffer = reinterpret_cast(_reg_output_buffer); + const int64_t input_size = inp.numel() * inp.element_size(); + TORCH_CHECK_LE(input_size, reg_buffer_sz_bytes); + if (reg_input_buffer != nullptr) { + AT_CUDA_CHECK(cudaMemcpyAsync(reg_input_buffer, inp.data_ptr(), input_size, + cudaMemcpyDeviceToDevice, stream)); + } else { + reg_input_buffer = inp.data_ptr(); + } + if (reg_output_buffer == nullptr) { + reg_output_buffer = normalized_out.data_ptr(); + } + + auto* input_ptr = reinterpret_cast(reg_input_buffer); + auto* shared_output_ptr = reinterpret_cast(reg_output_buffer); + auto* residual_ptr = reinterpret_cast(residual.data_ptr()); + auto* residual_out_ptr = reinterpret_cast(residual_out.data_ptr()); + const int num_tokens = static_cast(inp.size(0)); + const float epsilon_f = static_cast(epsilon); + if (weight.scalar_type() == at::ScalarType::Float) { + fa->sm70_reduce_scatter_gemma_rms_norm_all_gather<4, float, float>( + stream, input_ptr, shared_output_ptr, residual_ptr, + reinterpret_cast(weight.data_ptr()), residual_out_ptr, + num_tokens, kHiddenSize, epsilon_f); + } else { + fa->sm70_reduce_scatter_gemma_rms_norm_all_gather<4, float, half>( + stream, input_ptr, shared_output_ptr, residual_ptr, + reinterpret_cast(weight.data_ptr()), residual_out_ptr, + num_tokens, kHiddenSize, epsilon_f); + } + if (_reg_output_buffer != 0) { + AT_CUDA_CHECK(cudaMemcpyAsync(normalized_out.data_ptr(), reg_output_buffer, + input_size, cudaMemcpyDeviceToDevice, + stream)); + } +} + void all_reduce_sum2(fptr_t _fa, torch::Tensor& inp_a, torch::Tensor& inp_b, torch::Tensor& out) { auto fa = reinterpret_cast(_fa); @@ -307,17 +385,17 @@ void all_reduce_sum2(fptr_t _fa, torch::Tensor& inp_a, torch::Tensor& inp_b, switch (out.scalar_type()) { case at::ScalarType::Float: { - fa->allreduce_sum2(stream, reinterpret_cast(inp_a.data_ptr()), - reinterpret_cast(inp_b.data_ptr()), - reinterpret_cast(out.data_ptr()), - out.numel()); + fa->allreduce_sum2( + stream, reinterpret_cast(inp_a.data_ptr()), + reinterpret_cast(inp_b.data_ptr()), + reinterpret_cast(out.data_ptr()), out.numel()); break; } case at::ScalarType::Half: { - fa->allreduce_sum2(stream, reinterpret_cast(inp_a.data_ptr()), - reinterpret_cast(inp_b.data_ptr()), - reinterpret_cast(out.data_ptr()), - out.numel()); + fa->allreduce_sum2( + stream, reinterpret_cast(inp_a.data_ptr()), + reinterpret_cast(inp_b.data_ptr()), + reinterpret_cast(out.data_ptr()), out.numel()); break; } #if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) @@ -363,8 +441,7 @@ void top1_argmax(fptr_t _fa, torch::Tensor& input_pair, torch::Tensor& output, } void tile_runtime_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out, - fptr_t _reg_buffer, - int64_t reg_buffer_sz_bytes, + fptr_t _reg_buffer, int64_t reg_buffer_sz_bytes, int64_t tile_numel, int64_t engine_blocks, int64_t compute_iters) { auto fa = reinterpret_cast(_fa); @@ -412,8 +489,7 @@ void tile_runtime_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out, void tile_runtime_all_reduce_engine(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out, fptr_t _reg_buffer, int64_t reg_buffer_sz_bytes, - int64_t tile_numel, - int64_t producer_blocks, + int64_t tile_numel, int64_t producer_blocks, int64_t reducer_blocks, int64_t compute_iters) { auto fa = reinterpret_cast(_fa); @@ -499,6 +575,10 @@ void dispose(fptr_t _fa) { int64_t meta_size() { return sizeof(vllm::Signal); } +int64_t sm70_tp4_push_allreduce_buffer_size() { + return vllm::kSm70Tp4PushAllreduceBufferBytes; +} + void register_buffer(fptr_t _fa, const std::vector& fake_ipc_ptrs) { auto fa = reinterpret_cast(_fa); TORCH_CHECK(fake_ipc_ptrs.size() == fa->world_size_); @@ -509,6 +589,19 @@ void register_buffer(fptr_t _fa, const std::vector& fake_ipc_ptrs) { fa->register_buffer(ipc_ptrs); } +void register_sm70_tp4_push_allreduce_buffer( + fptr_t _fa, const std::vector& fake_ipc_ptrs) { + auto fa = reinterpret_cast(_fa); + TORCH_CHECK_EQ(fake_ipc_ptrs.size(), + static_cast(vllm::kSm70Tp4PushAllreduceWorldSize)); + TORCH_CHECK_EQ(fake_ipc_ptrs.size(), static_cast(fa->world_size_)); + void* ipc_ptrs[vllm::kSm70Tp4PushAllreduceWorldSize]; + for (size_t peer = 0; peer < fake_ipc_ptrs.size(); ++peer) { + ipc_ptrs[peer] = reinterpret_cast(fake_ipc_ptrs[peer]); + } + fa->register_sm70_tp4_push_buffer(ipc_ptrs); +} + // Use vector to represent byte data for python binding compatibility. std::tuple, std::vector> get_graph_buffer_ipc_meta(fptr_t _fa) { diff --git a/csrc/custom_all_reduce.cuh b/csrc/custom_all_reduce.cuh index 511b7cc668..f98aef77f5 100644 --- a/csrc/custom_all_reduce.cuh +++ b/csrc/custom_all_reduce.cuh @@ -51,9 +51,38 @@ inline hipPointer_attribute rangeStartAddrAttr = constexpr size_t kSm70Tp2SmallAllreduceBytes = 40 * 1024; constexpr size_t kSm70Tp8HierarchicalAllreduceBytes = 4096 * sizeof(half); +constexpr size_t kSm70Tp4MtpVerifierBytesPerRequest = 5 * 2048 * sizeof(half); constexpr int kSm70Tp8CompletionSignalSlotBase = 2; constexpr int kSm70GemmaRmsNormHiddenSize = 5120; constexpr int kSm70GemmaRmsNormThreads = 1024; +constexpr int kSm70LongPrefillSignalBlocks = + sm70_tile_runtime::kMaxSignalBlocks; +constexpr int kSm70LongPrefillMaxTokensPerRank = 2048; +// DFlash2 verifies eight positions at once. These constants mirror SGLang's +// two-epoch one-shot push collective while keeping its storage separate from +// vLLM's pull-collective metadata and intermediate buffers. +constexpr int kSm70Tp4PushAllreduceWorldSize = 4; +constexpr int kSm70Tp4PushAllreduceBlocks = 80; +constexpr int kSm70Tp4PushAllreduceThreads = 128; +constexpr int kSm70Tp4PushAllreduceEpochs = 2; +constexpr uint16_t kSm70Tp4PushAllreduceSentinel = 0x7f7f; +constexpr int kSm70Tp4PushAllreduceSentinelByte = 0x7f; +constexpr size_t kSm70Tp4PushAllreduceBytes = + 8 * kSm70GemmaRmsNormHiddenSize * sizeof(half); +constexpr size_t kSm70Tp4PushAllreduce8KiBBytes = 4096 * sizeof(half); +constexpr size_t kSm70Tp4PushAllreduceSignalBytes = + ((kSm70Tp4PushAllreduceBlocks * sizeof(uint32_t) + 127) / 128) * 128; +constexpr size_t kSm70Tp4PushAllreduceBufferBytes = + kSm70Tp4PushAllreduceSignalBytes + kSm70Tp4PushAllreduceEpochs * + kSm70Tp4PushAllreduceWorldSize * + kSm70Tp4PushAllreduceBytes; + +inline int sm70_tp4_push_allreduce_blocks(size_t bytes) { + if (bytes == kSm70Tp4PushAllreduceBytes) { + return kSm70Tp4PushAllreduceBlocks; + } + return bytes == kSm70Tp4PushAllreduce8KiBBytes ? 4 : 0; +} inline int sm70_gemma_rms_norm_threads() { const char* raw = std::getenv("VLLM_SM70_TP2_AR_GEMMA_RMS_THREADS"); @@ -64,6 +93,20 @@ inline int sm70_gemma_rms_norm_threads() { : kSm70GemmaRmsNormThreads; } +inline int sm70_tp4_long_fused_norm_threads() { + const char* raw = std::getenv("VLLM_SM70_TP4_LONG_FUSED_NORM_THREADS"); + if (raw == nullptr) return 512; + const int threads = std::atoi(raw); + return threads == 256 || threads == 512 || threads == 1024 ? threads : 512; +} + +inline int sm70_tp4_long_fused_norm_blocks() { + const char* raw = std::getenv("VLLM_SM70_TP4_LONG_FUSED_NORM_BLOCKS"); + if (raw == nullptr) return 80; + const int blocks = std::atoi(raw); + return blocks >= 1 && blocks <= kSm70LongPrefillSignalBlocks ? blocks : 80; +} + inline bool custom_allreduce_current_device_is_sm70() { #ifndef USE_ROCM int device = 0; @@ -86,7 +129,7 @@ inline bool custom_allreduce_current_device_is_sm70() { } inline bool sm70_tp8_hierarchical_custom_ar_enabled(int world_size, - bool fully_connected) { + bool fully_connected) { const char* raw = std::getenv("VLLM_SM70_TP8_HIERARCHICAL_CUSTOM_AR"); return raw != nullptr && std::atoi(raw) != 0 && world_size == 8 && !fully_connected && custom_allreduce_current_device_is_sm70(); @@ -96,15 +139,30 @@ inline bool sm70_tp8_hierarchical_peer(int rank, int peer) { return rank / 4 == peer / 4 || rank + 4 == peer || peer + 4 == rank; } -inline int custom_allreduce_block_limit(int default_limit, - int world_size, - size_t bytes) { +inline int custom_allreduce_block_limit(int default_limit, int world_size, + bool fully_connected, size_t bytes, + bool tune_sm70_tp4_mtp) { const char* raw = std::getenv("VLLM_CUSTOM_ALLREDUCE_BLOCK_LIMIT"); if (raw == nullptr || raw[0] == '\0') { if (world_size == 2 && bytes <= kSm70Tp2SmallAllreduceBytes && custom_allreduce_current_device_is_sm70()) { return 1; } + const char* tune_raw = std::getenv("VLLM_SM70_TP4_MTP_AR_BLOCK_TUNING"); + const bool tuning_enabled = + tune_raw != nullptr && tune_raw[0] != '\0' && std::atoi(tune_raw) != 0; + if (tune_sm70_tp4_mtp && tuning_enabled && + default_limit == defaultBlockLimit && world_size == 4 && + fully_connected && custom_allreduce_current_device_is_sm70()) { + if (bytes == kSm70Tp4MtpVerifierBytesPerRequest) { + return 1; + } + if (bytes == 8 * kSm70Tp4MtpVerifierBytesPerRequest || + bytes == 12 * kSm70Tp4MtpVerifierBytesPerRequest || + bytes == 16 * kSm70Tp4MtpVerifierBytesPerRequest) { + return 8; + } + } return default_limit; } char* end = nullptr; @@ -117,8 +175,7 @@ inline int custom_allreduce_block_limit(int default_limit, return static_cast(parsed); } -inline int sm70_tp4_m5_allreduce_threads(int world_size, - bool fully_connected, +inline int sm70_tp4_m5_allreduce_threads(int world_size, bool fully_connected, size_t bytes) { const char* raw = std::getenv("VLLM_SM70_TP4_M5_AR_THREADS"); if (raw == nullptr || raw[0] == '\0') return 512; @@ -139,6 +196,15 @@ inline int sm70_tp4_m5_allreduce_threads(int world_size, return static_cast(parsed); } +inline bool sm70_tp4_small_allreduce_pack32(int world_size, + bool fully_connected, + size_t bytes) { + const char* raw = std::getenv("VLLM_SM70_TP4_SMALL_AR_PACK32"); + return raw != nullptr && std::atoi(raw) != 0 && world_size == 4 && + fully_connected && bytes == 5120 * sizeof(half) && + custom_allreduce_current_device_is_sm70(); +} + // Counter may overflow, but unsigned integer overflow is well-defined. using FlagType = sm70_tile_runtime::FlagType; using Signal = sm70_tile_runtime::Signal; @@ -277,8 +343,7 @@ static DINLINE FlagType ld_flag_volatile(FlagType* flag_addr) { } static DINLINE void st_flag_sys_visible(FlagType* flag_addr, FlagType flag) { - asm volatile("membar.sys; st.volatile.global.u32 [%1], %0;" - ::"r"(flag), + asm volatile("membar.sys; st.volatile.global.u32 [%1], %0;" ::"r"(flag), "l"(flag_addr) : "memory"); } @@ -292,11 +357,47 @@ static DINLINE FlagType ld_flag_sys_visible(FlagType* flag_addr) { return flag; } +static DINLINE void membar_sys() { asm volatile("membar.sys;" ::: "memory"); } + +// The TP4 pack32 route waits with volatile peer loads, then executes one +// system fence after the matching flag is visible. This preserves the input +// visibility contract without putting a system fence in every poll iteration. +template +DINLINE void sm70_pack32_barrier_at_start(const RankSignals& sg, + Signal* self_sg, int rank) { + uint32_t flag = self_sg->_flag[blockIdx.x] + 1; + if (threadIdx.x < ngpus) { + const unsigned int peer = threadIdx.x; + auto peer_counter_ptr = &sg.signals[peer]->start[blockIdx.x][rank]; + auto self_counter_ptr = &self_sg->start[blockIdx.x][peer]; + st_flag_sys_visible(peer_counter_ptr, flag); + while (ld_flag_volatile(self_counter_ptr) != flag); + membar_sys(); + } + __syncthreads(); + if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; +} + +template +DINLINE void sm70_pack32_barrier_at_end(const RankSignals& sg, Signal* self_sg, + int rank) { + __syncthreads(); + uint32_t flag = self_sg->_flag[blockIdx.x] + 1; + if (threadIdx.x < ngpus) { + const unsigned int peer = threadIdx.x; + auto peer_counter_ptr = &sg.signals[peer]->end[blockIdx.x][rank]; + auto self_counter_ptr = &self_sg->end[blockIdx.x][peer]; + st_flag_volatile(peer_counter_ptr, flag); + while (ld_flag_volatile(self_counter_ptr) != flag); + } + if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; +} + // This function is meant to be used as the first synchronization in the all -// reduce kernel. The all-reduce input is usually produced by kernels launched -// immediately before this kernel on each rank. Use a system memory fence around -// the peer-visible flag so faster upstream paths do not expose stale input to -// peer ranks. +// reduce kernel. Publish the peer flag after a system fence, poll it with a +// volatile load, then execute one system fence after the matching value is +// visible. The post-poll fence preserves input visibility without putting a +// system fence in every spin-loop iteration. template DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg, int rank) { @@ -307,7 +408,8 @@ DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg, // Write the expected counter value to peer and wait for correct value // from peer. st_flag_sys_visible(peer_counter_ptr, flag); - while (ld_flag_sys_visible(self_counter_ptr) != flag); + while (ld_flag_volatile(self_counter_ptr) != flag); + membar_sys(); } __syncthreads(); // use one thread to update flag @@ -343,6 +445,38 @@ DINLINE void barrier_at_end(const RankSignals& sg, Signal* self_sg, int rank) { #else +template +DINLINE void sm70_pack32_barrier_at_start(const RankSignals& sg, + Signal* self_sg, int rank) { + uint32_t flag = self_sg->_flag[blockIdx.x] + 1; + if (threadIdx.x < ngpus) { + const unsigned int peer = threadIdx.x; + __scoped_atomic_store_n(&sg.signals[peer]->start[blockIdx.x][rank], flag, + __ATOMIC_RELAXED, __MEMORY_SCOPE_SYSTEM); + while (__scoped_atomic_load_n(&self_sg->start[blockIdx.x][peer], + __ATOMIC_RELAXED, + __MEMORY_SCOPE_DEVICE) < flag); + } + __syncthreads(); + if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; +} + +template +DINLINE void sm70_pack32_barrier_at_end(const RankSignals& sg, Signal* self_sg, + int rank) { + __syncthreads(); + uint32_t flag = self_sg->_flag[blockIdx.x] + 1; + if (threadIdx.x < ngpus) { + const unsigned int peer = threadIdx.x; + __scoped_atomic_store_n(&sg.signals[peer]->end[blockIdx.x][rank], flag, + __ATOMIC_RELAXED, __MEMORY_SCOPE_SYSTEM); + while (__scoped_atomic_load_n(&self_sg->end[blockIdx.x][peer], + __ATOMIC_RELAXED, + __MEMORY_SCOPE_DEVICE) < flag); + } + if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; +} + template DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg, int rank) { @@ -396,6 +530,54 @@ DINLINE P packed_reduce(const P* ptrs[], int idx) { return downcast

(tmp); } +// Adapted from SGLang-V100's SM70-capable one-shot push collective at +// haohervchb/sglang-V100@845b9fdf7a7e. SGLang uses positive zero as its empty +// slot sentinel. This variant preinitializes slots with a reserved FP16 NaN +// payload instead, preserving every finite payload bit, including signed zero. +DINLINE void sm70_push_escape_sentinel(half& value) { + auto* bits = reinterpret_cast(&value); + if (*bits == kSm70Tp4PushAllreduceSentinel) *bits = 0x7e00u; +} + +DINLINE bool sm70_push_is_sentinel(const half& value) { + const auto* bits = reinterpret_cast(&value); + return *bits == kSm70Tp4PushAllreduceSentinel; +} + +template +DINLINE void sm70_push_load_volatile_16b(P& value, const void* address, + int offset) { + static_assert(alignof(P) == 16 && sizeof(P) == 16); + const auto* source = reinterpret_cast(address) + offset; + uint4 bits; + asm volatile("ld.volatile.global.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(bits.x), "=r"(bits.y), "=r"(bits.z), "=r"(bits.w) + : "l"(source)); + value = *reinterpret_cast(&bits); +} + +template +DINLINE void sm70_push_store_volatile_16b(const P& value, void* address, + int offset) { + static_assert(alignof(P) == 16 && sizeof(P) == 16); + const uint4 bits = *reinterpret_cast(&value); + auto* destination = reinterpret_cast(address) + offset; + asm volatile("st.volatile.global.v4.b32 [%4], {%0, %1, %2, %3};" + : + : "r"(bits.x), "r"(bits.y), "r"(bits.z), "r"(bits.w), + "l"(destination)); +} + +template +DINLINE P sm70_push_reduce(P (&values)[ngpus]) { + A accumulator = upcast(values[0]); +#pragma unroll + for (int rank = 1; rank < ngpus; ++rank) { + packed_assign_add(accumulator, upcast(values[rank])); + } + return downcast

(accumulator); +} + template DINLINE P packed_reduce_sum2(const P* ptrs_a[], const P* ptrs_b[], int idx) { P local = ptrs_a[0][idx]; @@ -428,17 +610,113 @@ __global__ void __launch_bounds__(512, 1) barrier_at_end(sg, self_sg, rank); } +template +__global__ void __launch_bounds__(512, 1) + sm70_cross_device_reduce_1stage_pack32(RankData* _dp, RankSignals sg, + Signal* self_sg, + half* __restrict__ result, int rank, + int size) { + using P = array_t; + using A = array_t; + auto dp = *_dp; + sm70_pack32_barrier_at_start(sg, self_sg, rank); + for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size; + idx += gridDim.x * blockDim.x) { + reinterpret_cast(result)[idx] = + packed_reduce((const P**)&dp.ptrs[0], idx); + } + sm70_pack32_barrier_at_end(sg, self_sg, rank); +} + +template +__global__ void __launch_bounds__(1024, 1) + sm70_cross_device_reduce_1stage_push(RankData push_buffers, + const half* __restrict__ input, + half* __restrict__ output, int rank, + int packed_size) { + static_assert(ngpus == kSm70Tp4PushAllreduceWorldSize); + using P = typename packed_t::P; + using A = typename packed_t::A; + + auto* local_storage = + const_cast(reinterpret_cast(push_buffers.ptrs[rank])); + auto* local_epochs = reinterpret_cast(local_storage); + const uint32_t epoch = local_epochs[blockIdx.x]; + constexpr int packed_stride = kSm70Tp4PushAllreduceBytes / sizeof(P); + const int epoch_offset = epoch * ngpus * packed_stride; + const int offset = blockIdx.x * blockDim.x + threadIdx.x; + + if (offset < packed_size) { + P value = reinterpret_cast(input)[offset]; +#pragma unroll + for (int element = 0; element < P::size; ++element) { + sm70_push_escape_sentinel(value.data[element]); + } + +#pragma unroll + for (int destination_rank = 0; destination_rank < ngpus; + ++destination_rank) { + auto* destination_base = const_cast( + reinterpret_cast(push_buffers.ptrs[destination_rank])); + void* destination = destination_base + kSm70Tp4PushAllreduceSignalBytes + + (epoch_offset + rank * packed_stride) * sizeof(P); + sm70_push_store_volatile_16b(value, destination, offset); + } + + P peer_values[ngpus]; + while (true) { + bool has_empty_slot = false; +#pragma unroll + for (int source_rank = 0; source_rank < ngpus; ++source_rank) { + const void* source = + local_storage + kSm70Tp4PushAllreduceSignalBytes + + (epoch_offset + source_rank * packed_stride) * sizeof(P); + sm70_push_load_volatile_16b(peer_values[source_rank], source, offset); +#pragma unroll + for (int element = 0; element < P::size; ++element) { + has_empty_slot |= + sm70_push_is_sentinel(peer_values[source_rank].data[element]); + } + } + if (!has_empty_slot) break; + } + + reinterpret_cast(output)[offset] = + sm70_push_reduce(peer_values); + + P empty; +#pragma unroll + for (int element = 0; element < P::size; ++element) { + *reinterpret_cast(&empty.data[element]) = + kSm70Tp4PushAllreduceSentinel; + } +#pragma unroll + for (int source_rank = 0; source_rank < ngpus; ++source_rank) { + void* source = local_storage + kSm70Tp4PushAllreduceSignalBytes + + (epoch_offset + source_rank * packed_stride) * sizeof(P); + sm70_push_store_volatile_16b(empty, source, offset); + } + } + + __syncthreads(); + if (threadIdx.x == 0) { + local_epochs[blockIdx.x] = (epoch + 1) % kSm70Tp4PushAllreduceEpochs; + } +} + // This prototype deliberately stays narrow: it mirrors the FP32 Gemma RMSNorm // reduction order for a [tokens, 5120] FP16 projection. Each CTA handles one // row, retains the normal all-reduce peer order, and applies only its local // residual after the peer reduction. template __global__ void __launch_bounds__(Threads, 1) - sm70_peer_reduce_gemma_rms_norm( - RankData* _dp, RankSignals sg, Signal* self_sg, - const ResidualT* __restrict__ residual, - const WeightT* __restrict__ weight, half* __restrict__ normalized_out, - float* __restrict__ residual_out, int rank, float epsilon) { + sm70_peer_reduce_gemma_rms_norm(RankData* _dp, RankSignals sg, + Signal* self_sg, + const ResidualT* __restrict__ residual, + const WeightT* __restrict__ weight, + half* __restrict__ normalized_out, + float* __restrict__ residual_out, int rank, + float epsilon) { using P = typename packed_t::P; using A = typename packed_t::A; constexpr int kPackedWidth = P::size; @@ -459,15 +737,14 @@ __global__ void __launch_bounds__(Threads, 1) for (int packed_idx = tid; packed_idx < packed_per_row; packed_idx += blockDim.x) { - const P reduced = - packed_reduce((const P**)&dp.ptrs[0], - row * packed_per_row + packed_idx); + const P reduced = packed_reduce( + (const P**)&dp.ptrs[0], row * packed_per_row + packed_idx); const int element_offset = row_offset + packed_idx * kPackedWidth; #pragma unroll for (int i = 0; i < kPackedWidth; ++i) { - const float value = __half2float(reduced.data[i]) + - sm70_gemma_rms_norm_to_float( - residual[element_offset + i]); + const float value = + __half2float(reduced.data[i]) + + sm70_gemma_rms_norm_to_float(residual[element_offset + i]); residual_values[packed_idx * kPackedWidth + i] = value; residual_out[element_offset + i] = value; } @@ -503,20 +780,140 @@ __global__ void __launch_bounds__(Threads, 1) const int column = element_offset + i; const float gemma_weight = sm70_gemma_rms_norm_to_float(weight[column]) + 1.0f; - normalized_out[row_offset + column] = __float2half_rn( - residual_values[column] * inverse_rms * gemma_weight); + normalized_out[row_offset + column] = + __float2half_rn(residual_values[column] * inverse_rms * gemma_weight); } } barrier_at_end(sg, self_sg, rank); } +// Benchmark-only SM70 counterpart of NCCL's fused LSA RMSNorm example: +// reduce-scatter token rows, apply the mixed-dtype Gemma RMSNorm locally, then +// all-gather by writing the owned normalized rows into every peer output. +// Each rank launches one CTA per owned token and therefore moves the same +// asymptotic peer traffic as ring reduce-scatter + all-gather, without +// materializing an intermediate all-reduce tensor. +template +__global__ void __launch_bounds__(Threads, 1) + sm70_peer_reduce_scatter_gemma_rms_norm_all_gather( + RankData* _input_dp, RankData* _output_dp, RankSignals sg, + Signal* self_sg, const ResidualT* __restrict__ residual, + const WeightT* __restrict__ weight, float* __restrict__ residual_out, + int rank, int num_tokens, float epsilon) { + static_assert(ngpus == 4); + using H4 = array_t; + constexpr int kVarianceVectorWidth = 4; + constexpr int kVectorsPerRow = + kSm70GemmaRmsNormHiddenSize / kVarianceVectorWidth; + constexpr int kVectorsPerThread = (kVectorsPerRow + Threads - 1) / Threads; + + __shared__ float inverse_rms; + // Preserve the accepted local GemmaNorm reduction topology. The valid-item + // count below limits participation to the threads in this launch. + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage reduce_store; + + const int tokens_per_rank = num_tokens / ngpus; + const int tid = threadIdx.x; + auto input_dp = *_input_dp; + auto output_dp = *_output_dp; + + barrier_at_start(sg, self_sg, rank); + + for (int local_row = blockIdx.x; local_row < tokens_per_rank; + local_row += gridDim.x) { + const int row = rank * tokens_per_rank + local_row; + const int vector_row_offset = row * kVectorsPerRow; + float4 row_values[kVectorsPerThread]; + float variance = 0.0f; +#pragma unroll + for (int iter = 0; iter < kVectorsPerThread; ++iter) { + const int vector_idx = tid + iter * Threads; + if (vector_idx >= kVectorsPerRow) continue; + const int global_vector_idx = vector_row_offset + vector_idx; + H4 reduced = + reinterpret_cast(input_dp.ptrs[0])[global_vector_idx]; +#pragma unroll + for (int peer = 1; peer < ngpus; ++peer) { + const H4 peer_value = + reinterpret_cast(input_dp.ptrs[peer])[global_vector_idx]; +#pragma unroll + for (int i = 0; i < kVarianceVectorWidth; ++i) { + // NCCL's f16 Sum specialization uses __hadd/__hadd2 at every ring + // step. Preserve that FP16 rounding contract instead of accumulating + // all four inputs in FP32 and rounding only once. + reduced.data[i] = __hadd(reduced.data[i], peer_value.data[i]); + } + } + const float4 residual_value = + reinterpret_cast(residual)[global_vector_idx]; + float4 value; + value.x = __half2float(reduced.data[0]) + residual_value.x; + value.y = __half2float(reduced.data[1]) + residual_value.y; + value.z = __half2float(reduced.data[2]) + residual_value.z; + value.w = __half2float(reduced.data[3]) + residual_value.w; + row_values[iter] = value; + // Keep the public [M, H] residual shape so the compiled model graph does + // not change shape. Only this rank's persistent token shard is valid; + // the next fused boundary reads the same shard and ignores other rows. + reinterpret_cast(residual_out)[global_vector_idx] = value; + // Match the accepted local GemmaNorm exactly: each thread accumulates + // vector indices tid, tid + Threads, ... in x/y/z/w order before CUB. + variance += value.x * value.x; + variance += value.y * value.y; + variance += value.z * value.z; + variance += value.w * value.w; + } + variance = + BlockReduce(reduce_store).Reduce(variance, CubAddOp{}, blockDim.x); + if (tid == 0) { + inverse_rms = rsqrtf(variance / kSm70GemmaRmsNormHiddenSize + epsilon); + } + __syncthreads(); + + const float scale = inverse_rms; +#pragma unroll + for (int iter = 0; iter < kVectorsPerThread; ++iter) { + const int vector_idx = tid + iter * Threads; + if (vector_idx >= kVectorsPerRow) continue; + const int column = vector_idx * kVarianceVectorWidth; + const float4 value = row_values[iter]; + H4 normalized_pack; + normalized_pack.data[0] = __float2half_rn( + value.x * scale * + (sm70_gemma_rms_norm_to_float(weight[column]) + 1.0f)); + normalized_pack.data[1] = __float2half_rn( + value.y * scale * + (sm70_gemma_rms_norm_to_float(weight[column + 1]) + 1.0f)); + normalized_pack.data[2] = __float2half_rn( + value.z * scale * + (sm70_gemma_rms_norm_to_float(weight[column + 2]) + 1.0f)); + normalized_pack.data[3] = __float2half_rn( + value.w * scale * + (sm70_gemma_rms_norm_to_float(weight[column + 3]) + 1.0f)); +#pragma unroll + for (int peer = 0; peer < ngpus; ++peer) { + auto* peer_output = + reinterpret_cast(const_cast(output_dp.ptrs[peer])); + peer_output[vector_row_offset + vector_idx] = normalized_pack; + } + } + // CUB's temporary storage is reused by the next persistent row. + __syncthreads(); + } + + // The normalized rows are written directly into peer IPC buffers. Unlike + // an ordinary final all-reduce barrier, this barrier must publish those + // remote stores before a peer starts its local output copy. + barrier_at_end(sg, self_sg, rank); +} + template __global__ void __launch_bounds__(256, 1) sm70_tile_runtime_reduce_kernel( - RankData* _dp, RankSignals sg, Signal* self_sg, - const T* __restrict__ input, T* __restrict__ staging, - T* __restrict__ result, int rank, int packed_size, int tile_packed_size, - int tile_count, int compute_iters) { + RankData* _dp, RankSignals sg, Signal* self_sg, const T* __restrict__ input, + T* __restrict__ staging, T* __restrict__ result, int rank, int packed_size, + int tile_packed_size, int tile_count, int compute_iters) { using P = typename packed_t::P; using A = typename packed_t::A; @@ -567,11 +964,10 @@ __global__ void __launch_bounds__(256, 1) sm70_tile_runtime_reduce_kernel( template __global__ void __launch_bounds__(256, 1) sm70_tile_runtime_engine_kernel( - RankData* _dp, RankSignals sg, Signal* self_sg, - const T* __restrict__ input, T* __restrict__ staging, - T* __restrict__ result, int rank, int packed_size, int tile_packed_size, - int tile_count, int producer_blocks, int reducer_blocks, - int compute_iters) { + RankData* _dp, RankSignals sg, Signal* self_sg, const T* __restrict__ input, + T* __restrict__ staging, T* __restrict__ result, int rank, int packed_size, + int tile_packed_size, int tile_count, int producer_blocks, + int reducer_blocks, int compute_iters) { using P = typename packed_t::P; using A = typename packed_t::A; @@ -634,10 +1030,11 @@ __global__ void __launch_bounds__(256, 1) sm70_tile_runtime_engine_kernel( template __global__ void __launch_bounds__(256, 1) - sm70_tile_runtime_wait_reduce_kernel( - RankData* _dp, RankSignals sg, Signal* self_sg, - T* __restrict__ result, int rank, int packed_size, - int tile_packed_size, int tile_count) { + sm70_tile_runtime_wait_reduce_kernel(RankData* _dp, RankSignals sg, + Signal* self_sg, + T* __restrict__ result, int rank, + int packed_size, int tile_packed_size, + int tile_count) { using P = typename packed_t::P; using A = typename packed_t::A; @@ -669,9 +1066,11 @@ __global__ void __launch_bounds__(256, 1) } template -__global__ void __launch_bounds__(512, 1) cross_device_reduce_sum2_1stage( - RankData* _dp_a, RankData* _dp_b, RankSignals sg, Signal* self_sg, - T* __restrict__ result, int rank, int size) { +__global__ void __launch_bounds__(512, 1) + cross_device_reduce_sum2_1stage(RankData* _dp_a, RankData* _dp_b, + RankSignals sg, Signal* self_sg, + T* __restrict__ result, int rank, + int size) { using P = typename packed_t::P; using A = typename packed_t::A; auto dp_a = *_dp_a; @@ -679,9 +1078,8 @@ __global__ void __launch_bounds__(512, 1) cross_device_reduce_sum2_1stage( barrier_at_start(sg, self_sg, rank); for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size; idx += gridDim.x * blockDim.x) { - ((P*)result)[idx] = - packed_reduce_sum2((const P**)&dp_a.ptrs[0], - (const P**)&dp_b.ptrs[0], idx); + ((P*)result)[idx] = packed_reduce_sum2( + (const P**)&dp_a.ptrs[0], (const P**)&dp_b.ptrs[0], idx); } barrier_at_end(sg, self_sg, rank); } @@ -691,9 +1089,8 @@ DINLINE P* get_tmp_buf(Signal* sg) { return (P*)(((Signal*)sg) + 1); } -DINLINE FlagType sm70_tp8_clique_barrier(const RankSignals& sg, - Signal* self_sg, int rank, - int signal_slot) { +DINLINE FlagType sm70_tp8_clique_barrier(const RankSignals& sg, Signal* self_sg, + int rank, int signal_slot) { const int tid = threadIdx.x; const int clique_base = rank < 4 ? 0 : 4; const FlagType flag = self_sg->_flag[0] + 1; @@ -707,9 +1104,10 @@ DINLINE FlagType sm70_tp8_clique_barrier(const RankSignals& sg, return flag; } -static __global__ void __launch_bounds__(512, 1) sm70_tp8_hierarchical_reduce( - RankData* _dp, RankSignals sg, Signal* self_sg, - half* __restrict__ result, int rank, int packed_size) { +static __global__ void __launch_bounds__(512, 1) + sm70_tp8_hierarchical_reduce(RankData* _dp, RankSignals sg, Signal* self_sg, + half* __restrict__ result, int rank, + int packed_size) { using P = typename packed_t::P; using A = typename packed_t::A; @@ -723,8 +1121,7 @@ static __global__ void __launch_bounds__(512, 1) sm70_tp8_hierarchical_reduce( const int partial_slot = (self_sg->_flag[0] >> 1) & 1; const FlagType clique_flag = sm70_tp8_clique_barrier(sg, self_sg, rank, partial_slot); - A* const self_partial = - get_tmp_buf(self_sg) + partial_slot * packed_size; + A* const self_partial = get_tmp_buf(self_sg) + partial_slot * packed_size; const A* const pair_partial = get_tmp_buf(sg.signals[pair_rank]) + partial_slot * packed_size; @@ -747,11 +1144,9 @@ static __global__ void __launch_bounds__(512, 1) sm70_tp8_hierarchical_reduce( const FlagType pair_flag = clique_flag + 1; if (tid < 4) { const int peer = clique_base + tid; - const int completion_slot = - kSm70Tp8CompletionSignalSlotBase + partial_slot; + const int completion_slot = kSm70Tp8CompletionSignalSlotBase + partial_slot; st_flag_volatile(&sg.signals[peer]->end[completion_slot][rank], pair_flag); - while (ld_flag_volatile(&self_sg->end[completion_slot][peer]) != - pair_flag); + while (ld_flag_volatile(&self_sg->end[completion_slot][peer]) != pair_flag); } else if (tid == 4) { st_flag_release(&sg.signals[pair_rank]->end[partial_slot][rank], pair_flag); while (ld_flag_acquire(&self_sg->end[partial_slot][pair_rank]) != @@ -814,9 +1209,11 @@ __global__ void __launch_bounds__(512, 1) } template -__global__ void __launch_bounds__(512, 1) cross_device_reduce_sum2_2stage( - RankData* _dp_a, RankData* _dp_b, RankSignals sg, Signal* self_sg, - T* __restrict__ result, int rank, int size) { +__global__ void __launch_bounds__(512, 1) + cross_device_reduce_sum2_2stage(RankData* _dp_a, RankData* _dp_b, + RankSignals sg, Signal* self_sg, + T* __restrict__ result, int rank, + int size) { int tid = blockIdx.x * blockDim.x + threadIdx.x; int stride = gridDim.x * blockDim.x; using P = typename packed_t::P; @@ -841,8 +1238,7 @@ __global__ void __launch_bounds__(512, 1) cross_device_reduce_sum2_2stage( // Stage 1 mirrors cross_device_reduce_2stage, but each rank first forms // its local input_a + input_b value before the cross-rank reduction. for (int idx = start + tid; idx < end; idx += stride) { - tmp_out[idx - start] = - packed_reduce_sum2(ptrs_a, ptrs_b, idx); + tmp_out[idx - start] = packed_reduce_sum2(ptrs_a, ptrs_b, idx); } barrier_at_end(sg, self_sg, rank); @@ -922,6 +1318,8 @@ class CustomAllreduce { std::vector graph_unreg_buffers_; // a map from IPC handles to opened IPC pointers std::map ipc_handles_; + RankData sm70_tp4_push_buffers_{}; + bool sm70_tp4_push_buffers_registered_ = false; /** * Signals are an array of ipc-enabled buffers from all ranks. @@ -1002,6 +1400,27 @@ class CustomAllreduce { buffers_[ptrs[rank_]] = d_data; } + void register_sm70_tp4_push_buffer(void** ptrs) { + if (world_size_ != kSm70Tp4PushAllreduceWorldSize || !fully_connected_ || + !custom_allreduce_current_device_is_sm70()) { + throw std::runtime_error( + "SM70 push all-reduce requires fully-connected TP4 on SM70."); + } + for (int peer = 0; peer < world_size_; ++peer) { + if (ptrs[peer] == nullptr) { + throw std::runtime_error( + "SM70 push all-reduce received a null peer buffer."); + } + sm70_tp4_push_buffers_.ptrs[peer] = ptrs[peer]; + } + auto* local_data = + static_cast(ptrs[rank_]) + kSm70Tp4PushAllreduceSignalBytes; + CUDACHECK(cudaMemset( + local_data, kSm70Tp4PushAllreduceSentinelByte, + kSm70Tp4PushAllreduceBufferBytes - kSm70Tp4PushAllreduceSignalBytes)); + sm70_tp4_push_buffers_registered_ = true; + } + RankData* rank_data_for_buffer(cudaStream_t stream, void* buffer, const char* op_name) { RankData* ptrs; @@ -1013,11 +1432,10 @@ class CustomAllreduce { } else { auto it = buffers_.find(buffer); if (it == buffers_.end()) { - throw std::runtime_error(std::string(op_name) + - " buffer address " + - std::to_string( - reinterpret_cast(buffer)) + - " is not registered!"); + throw std::runtime_error( + std::string(op_name) + " buffer address " + + std::to_string(reinterpret_cast(buffer)) + + " is not registered!"); } ptrs = it->second; } @@ -1077,7 +1495,8 @@ class CustomAllreduce { void allreduce(cudaStream_t stream, T* input, T* output, int size, int threads = 512, int block_limit = defaultBlockLimit) { block_limit = custom_allreduce_block_limit( - block_limit, world_size_, static_cast(size) * sizeof(T)); + block_limit, world_size_, fully_connected_, + static_cast(size) * sizeof(T), true); auto d = packed_t::P::size; if (size % d != 0) throw std::runtime_error( @@ -1108,6 +1527,21 @@ class CustomAllreduce { size /= d; auto bytes = size * sizeof(typename packed_t::P); if constexpr (std::is_same_v) { + // The push protocol amortizes peer polling across captured collective + // chains. A lone eager call stays on the ordinary registered-buffer + // pull path. + if (sm70_tp4_push_buffers_registered_ && + status == cudaStreamCaptureStatusActive && + world_size_ == kSm70Tp4PushAllreduceWorldSize && fully_connected_ && + custom_allreduce_current_device_is_sm70()) { + const int push_blocks = sm70_tp4_push_allreduce_blocks(bytes); + if (push_blocks > 0) { + sm70_cross_device_reduce_1stage_push + <<>>( + sm70_tp4_push_buffers_, input, output, rank_, size); + return; + } + } if (sm70_tp8_hierarchical_custom_ar_enabled(world_size_, fully_connected_) && bytes == kSm70Tp8HierarchicalAllreduceBytes) { @@ -1116,10 +1550,20 @@ class CustomAllreduce { return; } } - threads = sm70_tp4_m5_allreduce_threads(world_size_, fully_connected_, - bytes); + threads = + sm70_tp4_m5_allreduce_threads(world_size_, fully_connected_, bytes); int blocks = std::min(block_limit, (size + threads - 1) / threads); + if constexpr (std::is_same_v) { + if (blocks == 1 && threads == 512 && + sm70_tp4_small_allreduce_pack32(world_size_, fully_connected_, + bytes)) { + sm70_cross_device_reduce_1stage_pack32<4><<<1, 512, 0, stream>>>( + ptrs, sg_, self_sg_, output, rank_, bytes / 32); + return; + } + } + // Check environment variable once const char* env_algo = std::getenv("VLLM_CUSTOM_ALLREDUCE_ALGO"); bool force_1stage = false; @@ -1180,11 +1624,11 @@ class CustomAllreduce { template void sm70_allreduce_gemma_rms_norm(cudaStream_t stream, half* input, - const ResidualT* residual, - const WeightT* weight, - half* normalized_out, - float* residual_out, int num_tokens, - int hidden_size, float epsilon) { + const ResidualT* residual, + const WeightT* weight, + half* normalized_out, float* residual_out, + int num_tokens, int hidden_size, + float epsilon) { if (world_size_ != ngpus || !custom_allreduce_current_device_is_sm70()) { throw std::runtime_error("SM70 Gemma RMSNorm prototype requires TP" + std::to_string(ngpus) + " on an SM70 device."); @@ -1198,18 +1642,18 @@ class CustomAllreduce { if (num_tokens <= 0 || num_tokens > kMaxTokens) { throw std::runtime_error( "SM70 Gemma RMSNorm prototype supports tokens in [1, " + - std::to_string(kMaxTokens) + "]. Got " + - std::to_string(num_tokens) + "."); + std::to_string(kMaxTokens) + "]. Got " + std::to_string(num_tokens) + + "."); } - RankData* ptrs = rank_data_for_buffer( - stream, input, "SM70 Gemma RMSNorm prototype input"); + RankData* ptrs = rank_data_for_buffer(stream, input, + "SM70 Gemma RMSNorm prototype input"); if constexpr (ngpus == 4) { // Keep the same 1024-thread CUB reduction topology as vLLM rms_norm. // packed_reduce preserves the cross_device_reduce_1stage // rank order before the FP32 residual add. sm70_peer_reduce_gemma_rms_norm + ResidualT, WeightT> <<>>( ptrs, sg_, self_sg_, residual, weight, normalized_out, residual_out, rank_, epsilon); @@ -1217,11 +1661,11 @@ class CustomAllreduce { } const int threads = sm70_gemma_rms_norm_threads(); -#define VLLM_LAUNCH_SM70_GEMMA_RMS_NORM(THREADS) \ +#define VLLM_LAUNCH_SM70_GEMMA_RMS_NORM(THREADS) \ sm70_peer_reduce_gemma_rms_norm \ - <<>>( \ - ptrs, sg_, self_sg_, residual, weight, normalized_out, \ - residual_out, rank_, epsilon) + <<>>(ptrs, sg_, self_sg_, residual, \ + weight, normalized_out, \ + residual_out, rank_, epsilon) switch (threads) { case 256: VLLM_LAUNCH_SM70_GEMMA_RMS_NORM(256); @@ -1236,12 +1680,63 @@ class CustomAllreduce { #undef VLLM_LAUNCH_SM70_GEMMA_RMS_NORM } + template + void sm70_reduce_scatter_gemma_rms_norm_all_gather( + cudaStream_t stream, half* input, half* shared_output, + const ResidualT* residual, const WeightT* weight, float* residual_out, + int num_tokens, int hidden_size, float epsilon) { + if (world_size_ != ngpus || !fully_connected_ || + !custom_allreduce_current_device_is_sm70()) { + throw std::runtime_error( + "SM70 long-prefill fused norm requires fully connected TP" + + std::to_string(ngpus) + " on SM70."); + } + if (hidden_size != kSm70GemmaRmsNormHiddenSize || num_tokens <= 0 || + num_tokens % ngpus != 0) { + throw std::runtime_error( + "SM70 long-prefill fused norm requires a TP-divisible [M, 5120] " + "input."); + } + const int tokens_per_rank = num_tokens / ngpus; + if (tokens_per_rank > kSm70LongPrefillMaxTokensPerRank) { + throw std::runtime_error( + "SM70 long-prefill fused norm exceeds the token-shard capacity."); + } + + RankData* input_ptrs = + rank_data_for_buffer(stream, input, "long-prefill fused norm input"); + RankData* output_ptrs = rank_data_for_buffer( + stream, shared_output, "long-prefill fused norm output"); + const int threads = sm70_tp4_long_fused_norm_threads(); + const int blocks = + std::min(tokens_per_rank, sm70_tp4_long_fused_norm_blocks()); +#define VLLM_LAUNCH_SM70_TP4_LONG_FUSED_NORM(THREADS) \ + sm70_peer_reduce_scatter_gemma_rms_norm_all_gather \ + <<>>(input_ptrs, output_ptrs, sg_, self_sg_, \ + residual, weight, residual_out, rank_, \ + num_tokens, epsilon) + switch (threads) { + case 256: + VLLM_LAUNCH_SM70_TP4_LONG_FUSED_NORM(256); + break; + case 1024: + VLLM_LAUNCH_SM70_TP4_LONG_FUSED_NORM(1024); + break; + default: + VLLM_LAUNCH_SM70_TP4_LONG_FUSED_NORM(512); + break; + } +#undef VLLM_LAUNCH_SM70_TP4_LONG_FUSED_NORM + } + template void allreduce_sum2(cudaStream_t stream, T* input_a, T* input_b, T* output, int size, int threads = 512, int block_limit = defaultBlockLimit) { block_limit = custom_allreduce_block_limit( - block_limit, world_size_, static_cast(size) * sizeof(T)); + block_limit, world_size_, fully_connected_, + static_cast(size) * sizeof(T), false); auto d = packed_t::P::size; if (size % d != 0) throw std::runtime_error( @@ -1293,29 +1788,28 @@ class CustomAllreduce { } } -#define SUM2_KL(ngpus, name) \ - name<<>>(ptrs_a, ptrs_b, sg_, \ - self_sg_, output, rank_, \ - size); -#define SUM2_CASE(ngpus) \ - case ngpus: { \ - if (force_1stage) { \ - SUM2_KL(ngpus, cross_device_reduce_sum2_1stage); \ - } else if (force_2stage) { \ - SUM2_KL(ngpus, cross_device_reduce_sum2_2stage); \ - } else { \ - if (world_size_ == 2) { \ - SUM2_KL(ngpus, cross_device_reduce_sum2_1stage); \ - } else if (fully_connected_) { \ - if ((world_size_ <= 4 && bytes < 512 * 1024) || \ - (world_size_ <= 8 && bytes < 256 * 1024)) { \ +#define SUM2_KL(ngpus, name) \ + name<<>>( \ + ptrs_a, ptrs_b, sg_, self_sg_, output, rank_, size); +#define SUM2_CASE(ngpus) \ + case ngpus: { \ + if (force_1stage) { \ + SUM2_KL(ngpus, cross_device_reduce_sum2_1stage); \ + } else if (force_2stage) { \ + SUM2_KL(ngpus, cross_device_reduce_sum2_2stage); \ + } else { \ + if (world_size_ == 2) { \ + SUM2_KL(ngpus, cross_device_reduce_sum2_1stage); \ + } else if (fully_connected_) { \ + if ((world_size_ <= 4 && bytes < 512 * 1024) || \ + (world_size_ <= 8 && bytes < 256 * 1024)) { \ SUM2_KL(ngpus, cross_device_reduce_sum2_1stage); \ - } else { \ + } else { \ SUM2_KL(ngpus, cross_device_reduce_sum2_2stage); \ - } \ - } \ - } \ - break; \ + } \ + } \ + } \ + break; \ } switch (world_size_) { @@ -1355,12 +1849,12 @@ class CustomAllreduce { const int packed_size = size / pack; const int tile_packed_size = tile_numel / pack; - const int tile_count = (packed_size + tile_packed_size - 1) / tile_packed_size; + const int tile_count = + (packed_size + tile_packed_size - 1) / tile_packed_size; if (tile_count <= 0 || tile_count > kMaxBlocks) { throw std::runtime_error( "SM70 tile runtime prototype supports tile_count in [1, " + - std::to_string(kMaxBlocks) + "]. Got " + - std::to_string(tile_count)); + std::to_string(kMaxBlocks) + "]. Got " + std::to_string(tile_count)); } auto it = buffers_.find(staging); @@ -1377,14 +1871,12 @@ class CustomAllreduce { blocks = std::max(1, std::min(blocks, tile_count)); compute_iters = std::max(0, compute_iters); -#define TILE_RUNTIME_CASE(ngpus) \ - case ngpus: { \ - sm70_tile_runtime_reduce_kernel \ - <<>>(ptrs, sg_, self_sg_, input, staging, \ - output, rank_, packed_size, \ - tile_packed_size, tile_count, \ - compute_iters); \ - break; \ +#define TILE_RUNTIME_CASE(ngpus) \ + case ngpus: { \ + sm70_tile_runtime_reduce_kernel<<>>( \ + ptrs, sg_, self_sg_, input, staging, output, rank_, packed_size, \ + tile_packed_size, tile_count, compute_iters); \ + break; \ } switch (world_size_) { @@ -1424,8 +1916,7 @@ class CustomAllreduce { if (tile_count <= 0 || tile_count > kMaxBlocks) { throw std::runtime_error( "SM70 tile runtime engine supports tile_count in [1, " + - std::to_string(kMaxBlocks) + "]. Got " + - std::to_string(tile_count)); + std::to_string(kMaxBlocks) + "]. Got " + std::to_string(tile_count)); } auto it = buffers_.find(staging); @@ -1447,15 +1938,13 @@ class CustomAllreduce { const int threads = 256; const int blocks = producer_blocks + reducer_blocks; -#define TILE_RUNTIME_ENGINE_CASE(ngpus) \ - case ngpus: { \ - sm70_tile_runtime_engine_kernel \ - <<>>(ptrs, sg_, self_sg_, input, staging, \ - output, rank_, packed_size, \ - tile_packed_size, tile_count, \ - producer_blocks, reducer_blocks, \ - compute_iters); \ - break; \ +#define TILE_RUNTIME_ENGINE_CASE(ngpus) \ + case ngpus: { \ + sm70_tile_runtime_engine_kernel<<>>( \ + ptrs, sg_, self_sg_, input, staging, output, rank_, packed_size, \ + tile_packed_size, tile_count, producer_blocks, reducer_blocks, \ + compute_iters); \ + break; \ } switch (world_size_) { @@ -1469,8 +1958,7 @@ class CustomAllreduce { template void tile_runtime_wait_reduce(cudaStream_t stream, T* staging, T* output, - int size, int tile_numel, - int reducer_blocks) { + int size, int tile_numel, int reducer_blocks) { if (world_size_ != 2) { throw std::runtime_error( "SM70 tile runtime wait-reduce currently supports only TP2."); @@ -1494,8 +1982,7 @@ class CustomAllreduce { if (tile_count <= 0 || tile_count > kMaxBlocks) { throw std::runtime_error( "SM70 tile runtime wait-reduce supports tile_count in [1, " + - std::to_string(kMaxBlocks) + "]. Got " + - std::to_string(tile_count)); + std::to_string(kMaxBlocks) + "]. Got " + std::to_string(tile_count)); } RankData* ptrs; @@ -1521,13 +2008,13 @@ class CustomAllreduce { constexpr int threads = 256; -#define TILE_RUNTIME_WAIT_REDUCE_CASE(ngpus) \ - case ngpus: { \ - sm70_tile_runtime_wait_reduce_kernel \ - <<>>( \ - ptrs, sg_, self_sg_, output, rank_, packed_size, \ - tile_packed_size, tile_count); \ - break; \ +#define TILE_RUNTIME_WAIT_REDUCE_CASE(ngpus) \ + case ngpus: { \ + sm70_tile_runtime_wait_reduce_kernel \ + <<>>( \ + ptrs, sg_, self_sg_, output, rank_, packed_size, tile_packed_size, \ + tile_count); \ + break; \ } switch (world_size_) { @@ -1556,11 +2043,11 @@ class CustomAllreduce { ptrs = it->second; } -#define TOP1_CASE(ngpus) \ - case ngpus: { \ - cross_device_top1_argmax<<<1, 32, 0, stream>>>( \ - ptrs, sg_, self_sg_, output, rank_); \ - break; \ +#define TOP1_CASE(ngpus) \ + case ngpus: { \ + cross_device_top1_argmax \ + <<<1, 32, 0, stream>>>(ptrs, sg_, self_sg_, output, rank_); \ + break; \ } switch (world_size_) { diff --git a/csrc/libtorch_stable/layernorm_kernels.cu b/csrc/libtorch_stable/layernorm_kernels.cu index bc5b367b7d..ce6a524b36 100644 --- a/csrc/libtorch_stable/layernorm_kernels.cu +++ b/csrc/libtorch_stable/layernorm_kernels.cu @@ -26,20 +26,25 @@ __global__ void rms_norm_kernel( const int64_t input_stride_d4, // input.stride(-4) const int64_t input_shape_d2, // input.size(-2) const int64_t input_shape_d3, // input.size(-3) - const scalar_t* __restrict__ weight, // [hidden_size] + const scalar_t* __restrict__ weight, // [hidden_size] or + // [num_groups, hidden_size] + const int64_t weight_stride, // 0 or weight.stride(0) const float epsilon, const int num_tokens, const int hidden_size) { __shared__ float s_variance; float variance = 0.0f; const scalar_t* input_row; + const scalar_t* weight_row; if constexpr (NUM_DIMS == 2) { // 2D for layernorm normal case [batch_size, hidden] input_row = input + blockIdx.x * input_stride_d2; + weight_row = weight + blockIdx.x * weight_stride; } else if constexpr (NUM_DIMS == 3) { // 3D for q/k norm [batch_size, num_heads, head_size] int batch_idx = blockIdx.x / input_shape_d2; int head_idx = blockIdx.x % input_shape_d2; input_row = input + batch_idx * input_stride_d3 + head_idx * input_stride_d2; + weight_row = weight + batch_idx * weight_stride; } else if constexpr (NUM_DIMS == 4) { // 4D for transformers model_impl qk norm [batch, seq, head, head_dim] int batch_idx = blockIdx.x / (input_shape_d3 * input_shape_d2); @@ -48,6 +53,7 @@ __global__ void rms_norm_kernel( int head_idx = remaining % input_shape_d2; input_row = input + batch_idx * input_stride_d4 + seq_idx * input_stride_d3 + head_idx * input_stride_d2; + weight_row = weight + batch_idx * weight_stride; } auto vec_op = [&variance](const vec_n_t& vec) { @@ -75,7 +81,7 @@ __global__ void rms_norm_kernel( scalar_t* out_row = out + blockIdx.x * hidden_size; auto* v_in = reinterpret_cast*>(input_row); - auto* v_w = reinterpret_cast*>(weight); + auto* v_w = reinterpret_cast*>(weight_row); auto* v_out = reinterpret_cast*>(out_row); for (int i = threadIdx.x; i < hidden_size / VEC_SIZE; i += blockDim.x) { vec_n_t dst; @@ -291,7 +297,8 @@ __global__ void __launch_bounds__(kSm70GemmaLongPrefillThreads, 2) void rms_norm(torch::stable::Tensor& out, // [..., hidden_size] torch::stable::Tensor& input, // [..., hidden_size] - torch::stable::Tensor& weight, // [hidden_size] + torch::stable::Tensor& weight, // [hidden_size] or + // [num_groups, hidden_size] double epsilon) { STD_TORCH_CHECK(out.is_contiguous()); if (input.stride(-1) != 1) { @@ -299,6 +306,16 @@ void rms_norm(torch::stable::Tensor& out, // [..., hidden_size] } STD_TORCH_CHECK(input.stride(-1) == 1); STD_TORCH_CHECK(weight.is_contiguous()); + int64_t weight_stride = 0; + if (weight.dim() == 1) { + STD_TORCH_CHECK(weight.size(0) == input.size(-1)); + } else if (weight.dim() == 2) { + STD_TORCH_CHECK(weight.size(0) == input.size(0)); + STD_TORCH_CHECK(weight.size(-1) == input.size(-1)); + weight_stride = weight.stride(0); + } else { + STD_TORCH_CHECK(false, "rms_norm weight must be 1D or 2D"); + } int hidden_size = input.size(-1); @@ -330,8 +347,8 @@ void rms_norm(torch::stable::Tensor& out, // [..., hidden_size] out.mutable_data_ptr(), input.const_data_ptr(), input_stride_d2, input_stride_d3, input_stride_d4, input_shape_d2, - input_shape_d3, weight.const_data_ptr(), epsilon, - num_tokens, hidden_size); + input_shape_d3, weight.const_data_ptr(), + weight_stride, epsilon, num_tokens, hidden_size); }); }); }); 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 2ee0b700f5..e2fd08b575 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -116,8 +116,7 @@ std::vector fp8_sm70_prepare(torch::Tensor _kernel, int64_t group_size, bool interleave_gated_silu); -void fp8_sm70_dequantize_out(torch::Tensor out, - torch::Tensor _kernel, +void fp8_sm70_dequantize_out(torch::Tensor out, torch::Tensor _kernel, torch::Tensor _scaling_factors, int64_t group_size); @@ -133,131 +132,205 @@ std::vector nvfp4_sm70_prepare(torch::Tensor _kernel, std::vector sm70_f16_prepare(torch::Tensor _kernel); -torch::Tensor awq_gemm_sm70(torch::Tensor _in_feats, - torch::Tensor _kernel, - torch::Tensor _scaling_factors, - int64_t group_size, - int64_t k_ld, - int64_t q_ld); +torch::Tensor awq_gemm_sm70(torch::Tensor _in_feats, torch::Tensor _kernel, + torch::Tensor _scaling_factors, int64_t group_size, + int64_t k_ld, int64_t q_ld); torch::Tensor sm70_f16_gemm(torch::Tensor _in_feats, torch::Tensor _kernel); -void awq_gemm_sm70_out(torch::Tensor out, - torch::Tensor _in_feats, - torch::Tensor _kernel, - torch::Tensor _scaling_factors, - int64_t group_size, - int64_t k_ld, - int64_t q_ld, +void awq_gemm_sm70_out(torch::Tensor out, torch::Tensor _in_feats, + torch::Tensor _kernel, torch::Tensor _scaling_factors, + int64_t group_size, int64_t k_ld, int64_t q_ld, bool gated_silu); -void awq_gemm_sm70_out_tile_reduce(torch::Tensor out, - torch::Tensor staging, - torch::Tensor _in_feats, - torch::Tensor _kernel, - torch::Tensor _scaling_factors, - int64_t group_size, - int64_t k_ld, - int64_t q_ld, - int64_t fa_ptr, - int64_t tile_numel, - int64_t reducer_blocks, - int64_t kernel_reducer_blocks, - bool overlap); - -void fp8_gemm_sm70_out(torch::Tensor out, - torch::Tensor _in_feats, - torch::Tensor _kernel, - torch::Tensor _scaling_factors, - int64_t group_size, - int64_t k_ld, - int64_t q_ld, +void awq_gemm_sm70_out_tile_reduce( + torch::Tensor out, torch::Tensor staging, torch::Tensor _in_feats, + torch::Tensor _kernel, torch::Tensor _scaling_factors, int64_t group_size, + int64_t k_ld, int64_t q_ld, int64_t fa_ptr, int64_t tile_numel, + int64_t reducer_blocks, int64_t kernel_reducer_blocks, bool overlap); + +void fp8_gemm_sm70_out(torch::Tensor out, torch::Tensor _in_feats, + torch::Tensor _kernel, torch::Tensor _scaling_factors, + int64_t group_size, int64_t k_ld, int64_t q_ld, bool gated_silu); +std::vector fp8_qpn8_prepare_sm70(torch::Tensor qweight, + torch::Tensor scales); + +void fp8_qpn8_dequantize_sm70_out(torch::Tensor out, torch::Tensor codes, + torch::Tensor group_scales); + +void fp8_qpn8_prefill_sm70_out(torch::Tensor out, int64_t dense_weight_ptr, + torch::Tensor input, torch::Tensor codes, + torch::Tensor group_scales, bool gated_silu); + +void fp8_qpn8_dispatch_sm70_out(torch::Tensor out, int64_t dense_weight_ptr, + torch::Tensor input, torch::Tensor codes, + torch::Tensor group_scales, int64_t split_k, + int64_t accumulator_chains, bool prefetch_codes, + bool gated_silu); + +void fp8_qpn8_gemm_sm70_out(torch::Tensor out, torch::Tensor input, + torch::Tensor codes, torch::Tensor group_scales, + int64_t split_k, int64_t accumulator_chains, + bool fast_decoder, bool prefetch_codes); + +void fp8_qpn8_gemm_ba_split_sm70_out(torch::Tensor qkv_out, torch::Tensor z_out, + torch::Tensor b_out, torch::Tensor a_out, + torch::Tensor input, torch::Tensor codes, + torch::Tensor group_scales, + torch::Tensor ba_weight); + +void fp8_qpn8_dispatch_ba_split_sm70_out( + torch::Tensor qkv_out, torch::Tensor z_out, torch::Tensor b_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); + +void fp8_qpn8_gated_pair_sm70_out(torch::Tensor out, torch::Tensor input, + torch::Tensor codes, + torch::Tensor group_scales, int64_t split_k, + int64_t accumulator_chains, bool fast_decoder, + bool prefetch_codes); + +std::vector nvfp4_qpn4_prepare_sm70(torch::Tensor qweight, + torch::Tensor scales); + +std::vector nvfp4_qpn4_prepare_scale_code_sm70( + torch::Tensor qweight, torch::Tensor scale_codes); + +void nvfp4_qpn4_dequantize_sm70_out(torch::Tensor out, torch::Tensor codes, + torch::Tensor scales, double global_scale, + bool use_scale_code); + +void nvfp4_qpn4_prefill_sm70_out(torch::Tensor out, int64_t dense_weight_ptr, + torch::Tensor input, torch::Tensor codes, + torch::Tensor scales, double global_scale, + bool use_scale_code, bool gated_silu); + +void nvfp4_qpn4_dispatch_sm70_out(torch::Tensor out, int64_t dense_weight_ptr, + torch::Tensor input, torch::Tensor codes, + torch::Tensor scales, double global_scale, + bool use_scale_code, bool gated_silu); + +void fp8_gemm_sm70_prefill_prescaled_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); + +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); + +void nvfp4_qpn2_gemm_sm70_out(torch::Tensor out, torch::Tensor input, + torch::Tensor codes, torch::Tensor scales, + double global_scale, int64_t split_k, + int64_t accumulator_chains); + +void nvfp4_qpn2_gated_sm70_out(torch::Tensor out, torch::Tensor input, + torch::Tensor codes, torch::Tensor scales, + double global_scale, int64_t split_k, + int64_t accumulator_chains); + +void nvfp4_qpn2_dispatch_sm70_out(torch::Tensor out, torch::Tensor input, + torch::Tensor codes, torch::Tensor scales, + double global_scale, int64_t split_k, + int64_t accumulator_chains, + torch::Tensor tm_weight, + torch::Tensor tm_scales, + int64_t tm_group_size, int64_t tm_k_ld, + int64_t tm_q_ld, bool gated_silu); + 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, int64_t k_ld, int64_t q_ld, bool gated_silu, int64_t min_prefill_m); -void mxfp4_gemm_sm70_out(torch::Tensor out, - torch::Tensor _in_feats, - torch::Tensor _kernel, - torch::Tensor _scaling_factors, - int64_t group_size, - int64_t k_ld, - int64_t q_ld, +void mxfp4_gemm_sm70_out(torch::Tensor out, torch::Tensor _in_feats, + torch::Tensor _kernel, torch::Tensor _scaling_factors, + int64_t group_size, int64_t k_ld, int64_t q_ld, bool gated_silu); -void nvfp4_gemm_sm70_out(torch::Tensor out, - torch::Tensor _in_feats, - torch::Tensor _kernel, - torch::Tensor _scaling_factors, - int64_t group_size, - int64_t k_ld, - int64_t q_ld, +void nvfp4_gemm_sm70_out(torch::Tensor out, torch::Tensor _in_feats, + torch::Tensor _kernel, torch::Tensor _scaling_factors, + int64_t group_size, int64_t k_ld, int64_t q_ld, bool gated_silu); -void nvfp4_gemv_sm70_raw_out(torch::Tensor out, - torch::Tensor _in_feats, +void nvfp4_gemv_sm70_raw_out(torch::Tensor out, torch::Tensor _in_feats, torch::Tensor _kernel, torch::Tensor _scaling_factors, - torch::Tensor partials, - int64_t group_size, + torch::Tensor partials, int64_t group_size, int64_t split_k); -void nvfp4_gemv_sm70_warp_out(torch::Tensor out, - torch::Tensor _in_feats, +void nvfp4_gemv_sm70_warp_out(torch::Tensor out, torch::Tensor _in_feats, torch::Tensor _kernel, torch::Tensor _scaling_factors, int64_t group_size); -void nvfp4_gemv_sm70_h2_out(torch::Tensor out, - torch::Tensor _in_feats, +void nvfp4_gemv_sm70_h2_out(torch::Tensor out, torch::Tensor _in_feats, torch::Tensor _kernel, torch::Tensor _scaling_factors, - torch::Tensor partials, - int64_t group_size, + torch::Tensor partials, int64_t group_size, int64_t split_k); -void fp8_gemm_sm70_out_auto(torch::Tensor out, - torch::Tensor _in_feats, +void fp8_gemm_sm70_out_auto(torch::Tensor out, torch::Tensor _in_feats, torch::Tensor _kernel, torch::Tensor _scaling_factors); -void fp8_gemm_sm70_out_meta(torch::Tensor out, - torch::Tensor _in_feats, +void fp8_gemm_sm70_out_meta(torch::Tensor out, torch::Tensor _in_feats, torch::Tensor _kernel, - torch::Tensor _scaling_factors, - torch::Tensor _meta, + torch::Tensor _scaling_factors, torch::Tensor _meta, bool gated_silu); -void sm70_f16_gemm_out(torch::Tensor out, - torch::Tensor _in_feats, - torch::Tensor _kernel, - int64_t k_ld, - bool gated_silu); +void sm70_f16_gemm_out(torch::Tensor out, torch::Tensor _in_feats, + torch::Tensor _kernel, int64_t k_ld, bool gated_silu); + +void sm70_f16_indexed_rerank_out(torch::Tensor out, torch::Tensor _in_feats, + torch::Tensor _kernel, + torch::Tensor candidate_ids, + torch::Tensor selected_raw, + torch::Tensor selected_packed, + torch::Tensor expanded, torch::Tensor partials, + torch::Tensor barriers, int64_t cta_n, + int64_t split_k); + +void sm70_f16_indexed_rerank_packed_out( + torch::Tensor out, torch::Tensor _in_feats, torch::Tensor _packed_kernel, + torch::Tensor candidate_ids, torch::Tensor selected_packed, + torch::Tensor expanded, torch::Tensor partials, torch::Tensor barriers, + int64_t cta_n, int64_t split_k); + +void sm70_f16_rerank_keys_out(torch::Tensor keys, torch::Tensor logits, + torch::Tensor candidate_ids); + +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_f16_lm_head_top1_out(torch::Tensor values_out, torch::Tensor indices_out, - torch::Tensor _in_feats, - torch::Tensor _kernel, - int64_t k_ld, - int64_t vocab_start_index, + torch::Tensor _in_feats, torch::Tensor _kernel, + int64_t k_ld, int64_t vocab_start_index, int64_t num_vocab_padding); void sm70_f16_lm_head_top1_tc_out(torch::Tensor values_out, torch::Tensor indices_out, torch::Tensor _in_feats, - torch::Tensor _kernel, - int64_t k_ld, + torch::Tensor _kernel, int64_t k_ld, int64_t vocab_start_index, int64_t num_vocab_padding); void sm70_f16_lm_head_top20_tc_out(torch::Tensor values_out, torch::Tensor indices_out, torch::Tensor _in_feats, - torch::Tensor _kernel, - int64_t k_ld, + torch::Tensor _kernel, int64_t k_ld, int64_t vocab_start_index, int64_t num_vocab_padding); @@ -273,27 +346,20 @@ void sm70_sample_packed_top20_out(torch::Tensor sampled_token_out, torch::Tensor sparse_ids_out, torch::Tensor sparse_probs_out, torch::Tensor gathered_pairs, - torch::Tensor exponential, - double top_p); + torch::Tensor exponential, double top_p); void sm70_dynamic_draft_vocab_update_tail_out( - torch::Tensor lru_token_ids, - torch::Tensor local_tail_token_ids, - torch::Tensor source_row_indices, - torch::Tensor observed_output_ids, - torch::Tensor target_candidate_ids, - torch::Tensor base_token_mask, - int64_t full_vocab_size, - int64_t local_shard_start, + torch::Tensor lru_token_ids, torch::Tensor local_tail_token_ids, + torch::Tensor source_row_indices, torch::Tensor observed_output_ids, + torch::Tensor target_candidate_ids, torch::Tensor base_token_mask, + int64_t full_vocab_size, int64_t local_shard_start, int64_t local_shard_end); void sm70_dynamic_draft_vocab_refresh_tail_weight_out( - torch::Tensor local_tail_weight, - torch::Tensor source_weight, + torch::Tensor local_tail_weight, torch::Tensor source_weight, torch::Tensor source_row_indices); -void sm70_f16_gate_mul_out(torch::Tensor out, - torch::Tensor _in_feats, +void sm70_f16_gate_mul_out(torch::Tensor out, torch::Tensor _in_feats, torch::Tensor _gate_weight); int64_t sm70_gemm_import_cache(torch::Tensor device_hint, @@ -302,343 +368,188 @@ int64_t sm70_gemm_import_cache(torch::Tensor device_hint, int64_t sm70_gemm_export_cache(torch::Tensor device_hint, const std::string& path); -std::vector awq_moe_build_strided_ptrs( - torch::Tensor tm_weights, - torch::Tensor tm_scales, - int64_t k_ld, - int64_t q_ld, - int64_t num_experts); +std::vector awq_moe_build_strided_ptrs(torch::Tensor tm_weights, + torch::Tensor tm_scales, + int64_t k_ld, + int64_t q_ld, + int64_t num_experts); -void awq_moe_gemm_sm70_out(torch::Tensor out, - torch::Tensor sorted_input, +void awq_moe_gemm_sm70_out(torch::Tensor out, torch::Tensor sorted_input, torch::Tensor expert_offsets, torch::Tensor strided_ptrs_w, - torch::Tensor strided_ptrs_s, - int64_t num_experts, - int64_t k, - int64_t n, - int64_t group_size, + torch::Tensor strided_ptrs_s, int64_t num_experts, + int64_t k, int64_t n, int64_t group_size, bool gated_silu); -void awq_moe_gemm_sm70_per_expert_dispatch_out(torch::Tensor out, - torch::Tensor sorted_input, - torch::Tensor expert_offsets, - torch::Tensor strided_ptrs_w, - torch::Tensor strided_ptrs_s, - int64_t num_experts, - int64_t k, - int64_t n, - int64_t group_size, - bool gated_silu); - -void awq_moe_dense_stage_sm70_out(torch::Tensor out, - torch::Tensor input, +void awq_moe_gemm_sm70_per_expert_dispatch_out( + torch::Tensor out, torch::Tensor sorted_input, torch::Tensor expert_offsets, + torch::Tensor strided_ptrs_w, torch::Tensor strided_ptrs_s, + int64_t num_experts, int64_t k, int64_t n, int64_t group_size, + bool gated_silu); + +void awq_moe_dense_stage_sm70_out(torch::Tensor out, torch::Tensor input, torch::Tensor expert_offsets, torch::Tensor dense_expert_ids, - torch::Tensor ptrs_w, - torch::Tensor ptrs_s, - int64_t num_experts, - int64_t k, - int64_t n, + torch::Tensor ptrs_w, torch::Tensor ptrs_s, + int64_t num_experts, int64_t k, int64_t n, int64_t group_size); void awq_moe_active_dense_stage_sm70_out( - torch::Tensor out, - torch::Tensor input, - torch::Tensor permuted_experts_id, - torch::Tensor active_expert_offsets, - torch::Tensor active_expert_ids, - torch::Tensor ptrs_w, - torch::Tensor ptrs_s, - int64_t total_slots, - int64_t k, - int64_t n, - int64_t group_size); + torch::Tensor out, torch::Tensor input, torch::Tensor permuted_experts_id, + torch::Tensor active_expert_offsets, torch::Tensor active_expert_ids, + torch::Tensor ptrs_w, torch::Tensor ptrs_s, int64_t total_slots, int64_t k, + int64_t n, int64_t group_size); void awq_moe_single_token_dense_stage_sm70_out( - torch::Tensor out, - torch::Tensor input, - torch::Tensor expert_offsets, - torch::Tensor sorted_expert_ids, - torch::Tensor ptrs_w, - torch::Tensor ptrs_s, - int64_t top_k, - int64_t k, - int64_t n, - int64_t group_size); + torch::Tensor out, torch::Tensor input, torch::Tensor expert_offsets, + torch::Tensor sorted_expert_ids, torch::Tensor ptrs_w, torch::Tensor ptrs_s, + int64_t top_k, int64_t k, int64_t n, int64_t group_size); void awq_moe_single_token_indexed_dense_stage_sm70_out( - torch::Tensor out, - torch::Tensor input, - torch::Tensor expert_offsets, - torch::Tensor sorted_expert_ids, - torch::Tensor ptrs_w, - torch::Tensor ptrs_s, - int64_t top_k, - int64_t k, - int64_t n, - int64_t group_size); + torch::Tensor out, torch::Tensor input, torch::Tensor expert_offsets, + torch::Tensor sorted_expert_ids, torch::Tensor ptrs_w, torch::Tensor ptrs_s, + int64_t top_k, int64_t k, int64_t n, int64_t group_size); void awq_moe_single_token_dense_w13_sm70_out( - torch::Tensor gate_up, - torch::Tensor compact_input, - torch::Tensor x, - torch::Tensor topk_ids, - torch::Tensor w13_ptrs_w, - torch::Tensor w13_ptrs_s, - torch::Tensor expert_offsets, - torch::Tensor expert_offsets64, - torch::Tensor inv_permuted_idx, - torch::Tensor sorted_expert_ids, - int64_t w13_k, - int64_t w13_n, - int64_t group_size, + torch::Tensor gate_up, torch::Tensor compact_input, torch::Tensor x, + torch::Tensor topk_ids, torch::Tensor w13_ptrs_w, torch::Tensor w13_ptrs_s, + torch::Tensor expert_offsets, torch::Tensor expert_offsets64, + torch::Tensor inv_permuted_idx, torch::Tensor sorted_expert_ids, + int64_t w13_k, int64_t w13_n, int64_t group_size, int64_t hidden_logical_size); void awq_moe_single_token_indexed_dense_w13_sm70_out( - torch::Tensor gate_up, - torch::Tensor compact_input, - torch::Tensor x, - torch::Tensor topk_ids, - torch::Tensor w13_ptrs_w, - torch::Tensor w13_ptrs_s, - torch::Tensor expert_offsets, - torch::Tensor expert_offsets64, - torch::Tensor inv_permuted_idx, - torch::Tensor sorted_expert_ids, - int64_t w13_k, - int64_t w13_n, - int64_t group_size, + torch::Tensor gate_up, torch::Tensor compact_input, torch::Tensor x, + torch::Tensor topk_ids, torch::Tensor w13_ptrs_w, torch::Tensor w13_ptrs_s, + torch::Tensor expert_offsets, torch::Tensor expert_offsets64, + torch::Tensor inv_permuted_idx, torch::Tensor sorted_expert_ids, + int64_t w13_k, int64_t w13_n, int64_t group_size, int64_t hidden_logical_size); void awq_moe_single_token_compact_dense_w13_sm70_out( - torch::Tensor gate_up, - torch::Tensor compact_input, - torch::Tensor x, - torch::Tensor topk_ids, - torch::Tensor w13_ptrs_w, - torch::Tensor w13_ptrs_s, - torch::Tensor compact_w13_ptrs_w, - torch::Tensor compact_w13_ptrs_s, - torch::Tensor expert_offsets, - torch::Tensor expert_offsets64, - torch::Tensor inv_permuted_idx, - torch::Tensor sorted_expert_ids, - int64_t w13_k, - int64_t w13_n, - int64_t group_size, + torch::Tensor gate_up, torch::Tensor compact_input, torch::Tensor x, + torch::Tensor topk_ids, torch::Tensor w13_ptrs_w, torch::Tensor w13_ptrs_s, + torch::Tensor compact_w13_ptrs_w, torch::Tensor compact_w13_ptrs_s, + torch::Tensor expert_offsets, torch::Tensor expert_offsets64, + torch::Tensor inv_permuted_idx, torch::Tensor sorted_expert_ids, + int64_t w13_k, int64_t w13_n, int64_t group_size, int64_t hidden_logical_size); -void awq_moe_single_token_exact_layout_prepare(torch::Tensor topk_ids, - torch::Tensor x, - torch::Tensor compact_input, - torch::Tensor expert_offsets, - torch::Tensor expert_offsets64, - torch::Tensor inv_permuted_idx, - int64_t num_experts); +void awq_moe_single_token_exact_layout_prepare( + torch::Tensor topk_ids, torch::Tensor x, torch::Tensor compact_input, + torch::Tensor expert_offsets, torch::Tensor expert_offsets64, + torch::Tensor inv_permuted_idx, int64_t num_experts); void awq_moe_single_token_weighted_reduce_out(torch::Tensor sorted_output, torch::Tensor topk_weights, torch::Tensor inv_permuted_idx, - torch::Tensor out, - int64_t top_k, + torch::Tensor out, int64_t top_k, int64_t hidden_logical_size); void awq_moe_single_token_sm70_out( - torch::Tensor out, - torch::Tensor x, - torch::Tensor topk_weights, - torch::Tensor topk_ids, - torch::Tensor src_w13_ptrs_w_rows, - torch::Tensor src_w13_ptrs_s_rows, - torch::Tensor src_w2_ptrs_w_rows, - torch::Tensor src_w2_ptrs_s_rows, - torch::Tensor compact_input, - torch::Tensor intermediate, - torch::Tensor sorted_output, - torch::Tensor sorted_weights, - torch::Tensor dst_w13_ptrs_w_rows, - torch::Tensor dst_w13_ptrs_s_rows, - torch::Tensor dst_w2_ptrs_w_rows, - torch::Tensor dst_w2_ptrs_s_rows, - torch::Tensor expert_offsets, - torch::Tensor inv_permuted_idx, - int64_t w13_k, - int64_t w13_n, - int64_t w2_k, - int64_t w2_n, - int64_t group_size, - int64_t hidden_logical_size); - -void fp8_moe_gemm_sm70_out(torch::Tensor out, - torch::Tensor sorted_input, + torch::Tensor out, torch::Tensor x, torch::Tensor topk_weights, + torch::Tensor topk_ids, torch::Tensor src_w13_ptrs_w_rows, + torch::Tensor src_w13_ptrs_s_rows, torch::Tensor src_w2_ptrs_w_rows, + torch::Tensor src_w2_ptrs_s_rows, torch::Tensor compact_input, + torch::Tensor intermediate, torch::Tensor sorted_output, + torch::Tensor sorted_weights, torch::Tensor dst_w13_ptrs_w_rows, + torch::Tensor dst_w13_ptrs_s_rows, torch::Tensor dst_w2_ptrs_w_rows, + torch::Tensor dst_w2_ptrs_s_rows, torch::Tensor expert_offsets, + torch::Tensor inv_permuted_idx, int64_t w13_k, int64_t w13_n, int64_t w2_k, + int64_t w2_n, int64_t group_size, int64_t hidden_logical_size); + +void fp8_moe_gemm_sm70_out(torch::Tensor out, torch::Tensor sorted_input, torch::Tensor expert_offsets, torch::Tensor strided_ptrs_w, - torch::Tensor strided_ptrs_s, - int64_t num_experts, - int64_t k, - int64_t n, - int64_t group_size, + torch::Tensor strided_ptrs_s, int64_t num_experts, + int64_t k, int64_t n, int64_t group_size, bool gated_silu); -void fp8_moe_gemm_sm70_per_expert_dispatch_out(torch::Tensor out, - torch::Tensor sorted_input, - torch::Tensor expert_offsets, - torch::Tensor strided_ptrs_w, - torch::Tensor strided_ptrs_s, - int64_t num_experts, - int64_t k, - int64_t n, - int64_t group_size, - bool gated_silu); - -void fp8_moe_dense_stage_sm70_out(torch::Tensor out, - torch::Tensor input, +void fp8_moe_gemm_sm70_per_expert_dispatch_out( + torch::Tensor out, torch::Tensor sorted_input, torch::Tensor expert_offsets, + torch::Tensor strided_ptrs_w, torch::Tensor strided_ptrs_s, + int64_t num_experts, int64_t k, int64_t n, int64_t group_size, + bool gated_silu); + +void fp8_moe_dense_stage_sm70_out(torch::Tensor out, torch::Tensor input, torch::Tensor expert_offsets, torch::Tensor dense_expert_ids, - torch::Tensor ptrs_w, - torch::Tensor ptrs_s, - int64_t num_experts, - int64_t k, - int64_t n, + torch::Tensor ptrs_w, torch::Tensor ptrs_s, + int64_t num_experts, int64_t k, int64_t n, int64_t group_size); -void mxfp4_moe_dense_stage_sm70_out(torch::Tensor out, - torch::Tensor input, +void mxfp4_moe_dense_stage_sm70_out(torch::Tensor out, torch::Tensor input, + torch::Tensor expert_offsets, + torch::Tensor dense_expert_ids, + torch::Tensor ptrs_w, torch::Tensor ptrs_s, + int64_t num_experts, int64_t k, int64_t n, + int64_t group_size); + +void nvfp4_moe_dense_stage_sm70_out(torch::Tensor out, torch::Tensor input, torch::Tensor expert_offsets, torch::Tensor dense_expert_ids, - torch::Tensor ptrs_w, - torch::Tensor ptrs_s, - int64_t num_experts, - int64_t k, - int64_t n, + torch::Tensor ptrs_w, torch::Tensor ptrs_s, + int64_t num_experts, int64_t k, int64_t n, int64_t group_size); void mxfp4_moe_single_token_prepare_w13_sm70_out( - torch::Tensor gate_up, - torch::Tensor compact_input, - torch::Tensor x, - torch::Tensor topk_ids, - torch::Tensor w13_ptrs_w, - torch::Tensor w13_ptrs_s, - torch::Tensor expert_offsets, - torch::Tensor inv_permuted_idx, - torch::Tensor sorted_expert_ids, - int64_t w13_k, - int64_t w13_n, - int64_t group_size, - int64_t hidden_logical_size); + torch::Tensor gate_up, torch::Tensor compact_input, torch::Tensor x, + torch::Tensor topk_ids, torch::Tensor w13_ptrs_w, torch::Tensor w13_ptrs_s, + torch::Tensor expert_offsets, torch::Tensor inv_permuted_idx, + torch::Tensor sorted_expert_ids, int64_t w13_k, int64_t w13_n, + int64_t group_size, int64_t hidden_logical_size); void fp8_moe_single_token_dense_stage_sm70_out( - torch::Tensor out, - torch::Tensor input, - torch::Tensor expert_offsets, - torch::Tensor sorted_expert_ids, - torch::Tensor ptrs_w, - torch::Tensor ptrs_s, - int64_t top_k, - int64_t k, - int64_t n, - int64_t group_size); + torch::Tensor out, torch::Tensor input, torch::Tensor expert_offsets, + torch::Tensor sorted_expert_ids, torch::Tensor ptrs_w, torch::Tensor ptrs_s, + int64_t top_k, int64_t k, int64_t n, int64_t group_size); void fp8_moe_single_token_indexed_dense_stage_sm70_out( - torch::Tensor out, - torch::Tensor input, - torch::Tensor expert_offsets, - torch::Tensor sorted_expert_ids, - torch::Tensor ptrs_w, - torch::Tensor ptrs_s, - int64_t top_k, - int64_t k, - int64_t n, - int64_t group_size); + torch::Tensor out, torch::Tensor input, torch::Tensor expert_offsets, + torch::Tensor sorted_expert_ids, torch::Tensor ptrs_w, torch::Tensor ptrs_s, + int64_t top_k, int64_t k, int64_t n, int64_t group_size); void fp8_moe_single_token_dense_w13_sm70_out( - torch::Tensor gate_up, - torch::Tensor compact_input, - torch::Tensor x, - torch::Tensor topk_ids, - torch::Tensor w13_ptrs_w, - torch::Tensor w13_ptrs_s, - torch::Tensor expert_offsets, - torch::Tensor expert_offsets64, - torch::Tensor inv_permuted_idx, - torch::Tensor sorted_expert_ids, - int64_t w13_k, - int64_t w13_n, - int64_t group_size, + torch::Tensor gate_up, torch::Tensor compact_input, torch::Tensor x, + torch::Tensor topk_ids, torch::Tensor w13_ptrs_w, torch::Tensor w13_ptrs_s, + torch::Tensor expert_offsets, torch::Tensor expert_offsets64, + torch::Tensor inv_permuted_idx, torch::Tensor sorted_expert_ids, + int64_t w13_k, int64_t w13_n, int64_t group_size, int64_t hidden_logical_size); void fp8_moe_single_token_indexed_dense_w13_sm70_out( - torch::Tensor gate_up, - torch::Tensor compact_input, - torch::Tensor x, - torch::Tensor topk_ids, - torch::Tensor w13_ptrs_w, - torch::Tensor w13_ptrs_s, - torch::Tensor expert_offsets, - torch::Tensor expert_offsets64, - torch::Tensor inv_permuted_idx, - torch::Tensor sorted_expert_ids, - int64_t w13_k, - int64_t w13_n, - int64_t group_size, + torch::Tensor gate_up, torch::Tensor compact_input, torch::Tensor x, + torch::Tensor topk_ids, torch::Tensor w13_ptrs_w, torch::Tensor w13_ptrs_s, + torch::Tensor expert_offsets, torch::Tensor expert_offsets64, + torch::Tensor inv_permuted_idx, torch::Tensor sorted_expert_ids, + int64_t w13_k, int64_t w13_n, int64_t group_size, int64_t hidden_logical_size); void fp8_moe_single_token_compact_dense_w13_sm70_out( - torch::Tensor gate_up, - torch::Tensor compact_input, - torch::Tensor x, - torch::Tensor topk_ids, - torch::Tensor w13_ptrs_w, - torch::Tensor w13_ptrs_s, - torch::Tensor compact_w13_ptrs_w, - torch::Tensor compact_w13_ptrs_s, - torch::Tensor expert_offsets, - torch::Tensor expert_offsets64, - torch::Tensor inv_permuted_idx, - torch::Tensor sorted_expert_ids, - int64_t w13_k, - int64_t w13_n, - int64_t group_size, + torch::Tensor gate_up, torch::Tensor compact_input, torch::Tensor x, + torch::Tensor topk_ids, torch::Tensor w13_ptrs_w, torch::Tensor w13_ptrs_s, + torch::Tensor compact_w13_ptrs_w, torch::Tensor compact_w13_ptrs_s, + torch::Tensor expert_offsets, torch::Tensor expert_offsets64, + torch::Tensor inv_permuted_idx, torch::Tensor sorted_expert_ids, + int64_t w13_k, int64_t w13_n, int64_t group_size, int64_t hidden_logical_size); void fp8_moe_single_token_sm70_out( - torch::Tensor out, - torch::Tensor x, - torch::Tensor topk_weights, - torch::Tensor topk_ids, - torch::Tensor src_w13_ptrs_w_rows, - torch::Tensor src_w13_ptrs_s_rows, - torch::Tensor src_w2_ptrs_w_rows, - torch::Tensor src_w2_ptrs_s_rows, - torch::Tensor compact_input, - torch::Tensor gate_up, - torch::Tensor intermediate, - torch::Tensor sorted_output, - torch::Tensor sorted_weights, - torch::Tensor dst_w13_ptrs_w_rows, - torch::Tensor dst_w13_ptrs_s_rows, - torch::Tensor dst_w2_ptrs_w_rows, - torch::Tensor dst_w2_ptrs_s_rows, - torch::Tensor expert_offsets, - torch::Tensor inv_permuted_idx, - torch::Tensor sorted_expert_ids, - torch::Tensor broadcast_input_indices, - torch::Tensor w2_raw_weight, - torch::Tensor w2_raw_scale_inv, - int64_t w13_k, - int64_t w13_n, - int64_t w2_k, - int64_t w2_n, - int64_t group_size, - int64_t hidden_logical_size, - bool fused_gated_silu, - bool fused_weighted_reduce, - bool broadcast_input, - bool w2_direct_reduce, - bool indexed_expert_ptrs, - bool exact_per_route); + torch::Tensor out, torch::Tensor x, torch::Tensor topk_weights, + torch::Tensor topk_ids, torch::Tensor src_w13_ptrs_w_rows, + torch::Tensor src_w13_ptrs_s_rows, torch::Tensor src_w2_ptrs_w_rows, + torch::Tensor src_w2_ptrs_s_rows, torch::Tensor compact_input, + torch::Tensor gate_up, torch::Tensor intermediate, + torch::Tensor sorted_output, torch::Tensor sorted_weights, + torch::Tensor dst_w13_ptrs_w_rows, torch::Tensor dst_w13_ptrs_s_rows, + torch::Tensor dst_w2_ptrs_w_rows, torch::Tensor dst_w2_ptrs_s_rows, + torch::Tensor expert_offsets, torch::Tensor inv_permuted_idx, + torch::Tensor sorted_expert_ids, torch::Tensor broadcast_input_indices, + torch::Tensor w2_raw_weight, torch::Tensor w2_raw_scale_inv, int64_t w13_k, + int64_t w13_n, int64_t w2_k, int64_t w2_n, int64_t group_size, + int64_t hidden_logical_size, bool fused_gated_silu, + bool fused_weighted_reduce, bool broadcast_input, bool w2_direct_reduce, + bool indexed_expert_ptrs, bool exact_per_route); #endif void static_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input, @@ -664,27 +575,30 @@ void all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out, void sm70_tp2_all_reduce_gemma_rms_norm( fptr_t _fa, torch::Tensor& inp, torch::Tensor& residual, torch::Tensor& weight, torch::Tensor& normalized_out, - torch::Tensor& residual_out, fptr_t reg_buffer, - int64_t reg_buffer_sz_bytes, double epsilon); + torch::Tensor& residual_out, fptr_t reg_buffer, int64_t reg_buffer_sz_bytes, + double epsilon); void sm70_tp4_all_reduce_gemma_rms_norm( fptr_t _fa, torch::Tensor& inp, torch::Tensor& residual, torch::Tensor& weight, torch::Tensor& normalized_out, - torch::Tensor& residual_out, fptr_t reg_buffer, - int64_t reg_buffer_sz_bytes, double epsilon); + torch::Tensor& residual_out, fptr_t reg_buffer, int64_t reg_buffer_sz_bytes, + double epsilon); +void sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather( + fptr_t _fa, torch::Tensor& inp, torch::Tensor& residual, + torch::Tensor& weight, torch::Tensor& normalized_out, + torch::Tensor& residual_out, fptr_t reg_input_buffer, + fptr_t reg_output_buffer, int64_t reg_buffer_sz_bytes, double epsilon); void all_reduce_sum2(fptr_t _fa, torch::Tensor& inp_a, torch::Tensor& inp_b, torch::Tensor& out); void top1_argmax(fptr_t _fa, torch::Tensor& input_pair, torch::Tensor& output, fptr_t reg_buffer, int64_t reg_buffer_sz_bytes); void tile_runtime_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out, - fptr_t reg_buffer, - int64_t reg_buffer_sz_bytes, + fptr_t reg_buffer, int64_t reg_buffer_sz_bytes, int64_t tile_numel, int64_t engine_blocks, int64_t compute_iters); void tile_runtime_all_reduce_engine(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out, fptr_t reg_buffer, int64_t reg_buffer_sz_bytes, - int64_t tile_numel, - int64_t producer_blocks, + int64_t tile_numel, int64_t producer_blocks, int64_t reducer_blocks, int64_t compute_iters); void tile_runtime_wait_reduce(fptr_t _fa, torch::Tensor& staging, @@ -692,7 +606,10 @@ void tile_runtime_wait_reduce(fptr_t _fa, torch::Tensor& staging, int64_t reducer_blocks); void dispose(fptr_t _fa); int64_t meta_size(); +int64_t sm70_tp4_push_allreduce_buffer_size(); void register_buffer(fptr_t _fa, const std::vector& fake_ipc_ptrs); +void register_sm70_tp4_push_allreduce_buffer( + fptr_t _fa, const std::vector& fake_ipc_ptrs); std::tuple, std::vector> get_graph_buffer_ipc_meta(fptr_t _fa); void register_graph_buffers(fptr_t _fa, 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_tile_runtime_signal.cuh b/csrc/sm70_tile_runtime_signal.cuh index 401b9f99eb..061a8147c2 100644 --- a/csrc/sm70_tile_runtime_signal.cuh +++ b/csrc/sm70_tile_runtime_signal.cuh @@ -4,17 +4,19 @@ namespace vllm::sm70_tile_runtime { -// Keep the default all-reduce block limit unchanged, but allow the -// TileRT-style MLP down-proj path to publish one flag per TurboMind N tile. +// Keep the default TileRT/all-reduce block limit unchanged. The larger signal +// capacity is reserved for long-prefill collective-fusion experiments that use +// one cross-rank barrier per owned token. constexpr int kMaxBlocks = 64; +constexpr int kMaxSignalBlocks = 512; constexpr int kMaxRanks = 8; using FlagType = uint32_t; struct Signal { - alignas(128) FlagType start[kMaxBlocks][kMaxRanks]; - alignas(128) FlagType end[kMaxBlocks][kMaxRanks]; - alignas(128) FlagType _flag[kMaxBlocks]; + alignas(128) FlagType start[kMaxSignalBlocks][kMaxRanks]; + alignas(128) FlagType end[kMaxSignalBlocks][kMaxRanks]; + alignas(128) FlagType _flag[kMaxSignalBlocks]; }; struct __align__(16) RankData { @@ -28,14 +30,13 @@ struct __align__(16) RankSignals { #if !defined(USE_ROCM) static __device__ __forceinline__ void store_flag_sys_visible( FlagType* flag_addr, FlagType flag) { - asm volatile("membar.sys; st.volatile.global.u32 [%1], %0;" - ::"r"(flag), + asm volatile("membar.sys; st.volatile.global.u32 [%1], %0;" ::"r"(flag), "l"(flag_addr) : "memory"); } -static __device__ __forceinline__ FlagType load_flag_sys_visible( - FlagType* flag_addr) { +static __device__ __forceinline__ FlagType +load_flag_sys_visible(FlagType* flag_addr) { FlagType flag; asm volatile("ld.volatile.global.u32 %0, [%1]; membar.sys;" : "=r"(flag) diff --git a/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/attention/quantization.h b/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/attention/quantization.h index 3a3fd5e6ff..fbb63b8426 100644 --- a/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/attention/quantization.h +++ b/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/attention/quantization.h @@ -14,875 +14,858 @@ namespace turbomind { #define TM_ROUND_USE_CVT_RNI 1 -inline constexpr bool kFuseU4F16Dequant = false; +inline constexpr bool kFuseU4F16Dequant = false; inline constexpr bool kForceIntZeroPoint = false; -template -__device__ T Infinity() -{ - if constexpr (std::is_same_v) { - return __ushort_as_half((unsigned short)0x7C00U); - } +template +__device__ T Infinity() { + if constexpr (std::is_same_v) { + return __ushort_as_half((unsigned short)0x7C00U); + } #if __CUDA_ARCH__ >= 800 - if constexpr (std::is_same_v) { - return __ushort_as_bfloat16((unsigned short)0x7F80U); - } + if constexpr (std::is_same_v) { + return __ushort_as_bfloat16((unsigned short)0x7F80U); + } #endif - if constexpr (std::is_same_v) { - return __int_as_float(0x7f800000U); - } + if constexpr (std::is_same_v) { + return __int_as_float(0x7f800000U); + } - return T{}; + return T{}; } -template -__device__ constexpr T Max(T a, T b) -{ - if constexpr (std::is_same_v) { - return __hmax(a, b); - } +template +__device__ constexpr T Max(T a, T b) { + if constexpr (std::is_same_v) { + return __hmax(a, b); + } #if __CUDA_ARCH__ >= 800 - if constexpr (std::is_same_v) { - return __hmax(a, b); - } + if constexpr (std::is_same_v) { + return __hmax(a, b); + } #endif - if constexpr (std::is_same_v) { - return fmaxf(a, b); - } + if constexpr (std::is_same_v) { + return fmaxf(a, b); + } - if constexpr (std::is_same_v) { - return max(a, b); - } + if constexpr (std::is_same_v) { + return max(a, b); + } - return T{}; + return T{}; } -template -__device__ constexpr T Min(T a, T b) -{ - if constexpr (std::is_same_v) { - return __hmin(a, b); - } +template +__device__ constexpr T Min(T a, T b) { + if constexpr (std::is_same_v) { + return __hmin(a, b); + } #if __CUDA_ARCH__ >= 800 - if constexpr (std::is_same_v) { - return __hmin(a, b); - } + if constexpr (std::is_same_v) { + return __hmin(a, b); + } #endif - if constexpr (std::is_same_v) { - return fminf(a, b); - } + if constexpr (std::is_same_v) { + return fminf(a, b); + } - if constexpr (std::is_same_v) { - return min(a, b); - } + if constexpr (std::is_same_v) { + return min(a, b); + } - return T{}; + return T{}; } -template -inline __device__ Array cvt_f16x4_u8(const Array& src) -{ - static constexpr uint32_t f16_magic = 0x64000000; - // 01234567 01234567 - // SEEEEEMM MMMMMMMM - // 1MM XXXXXXXX - // (1 + x/2^10) * 2^(e-15) -> e-15=10 -> e=25=16+8+1 -> 01100100b -> 0x64 - Array dst; - dst[0] = __byte_perm((uint32_t&)src, f16_magic, 0x7170); - dst[1] = __byte_perm((uint32_t&)src, f16_magic, 0x7372); - if constexpr (norm) { - for (int i = 0; i < 4; ++i) { - ((Array&)dst)[i] -= __ushort_as_half(0x6400U); - } +template +inline __device__ Array cvt_f16x4_u8(const Array& src) { + static constexpr uint32_t f16_magic = 0x64000000; + // 01234567 01234567 + // SEEEEEMM MMMMMMMM + // 1MM XXXXXXXX + // (1 + x/2^10) * 2^(e-15) -> e-15=10 -> e=25=16+8+1 -> 01100100b -> 0x64 + Array dst; + dst[0] = __byte_perm((uint32_t&)src, f16_magic, 0x7170); + dst[1] = __byte_perm((uint32_t&)src, f16_magic, 0x7372); + if constexpr (norm) { + for (int i = 0; i < 4; ++i) { + ((Array&)dst)[i] -= __ushort_as_half(0x6400U); } - return (Array&)dst; + } + return (Array&)dst; } -template -inline __device__ Array cvt_f16x2x2_u8_trans(const Array& src) -{ - static constexpr uint32_t f16_magic = 0x64000000; - // 01234567 01234567 - // SEEEEEMM MMMMMMMM - // 1MM XXXXXXXX - // (1 + x/2^10) * 2^(e-15) -> e-15=10 -> e=25=16+8+1 -> 01100100b -> 0x64 - Array dst; - dst[0] = __byte_perm((uint32_t&)src, f16_magic, 0x7270); - dst[1] = __byte_perm((uint32_t&)src, f16_magic, 0x7371); - if constexpr (norm) { - for (int i = 0; i < 4; ++i) { - ((Array&)dst)[i] -= __ushort_as_half(0x6400U); - } +template +inline __device__ Array cvt_f16x2x2_u8_trans( + const Array& src) { + static constexpr uint32_t f16_magic = 0x64000000; + // 01234567 01234567 + // SEEEEEMM MMMMMMMM + // 1MM XXXXXXXX + // (1 + x/2^10) * 2^(e-15) -> e-15=10 -> e=25=16+8+1 -> 01100100b -> 0x64 + Array dst; + dst[0] = __byte_perm((uint32_t&)src, f16_magic, 0x7270); + dst[1] = __byte_perm((uint32_t&)src, f16_magic, 0x7371); + if constexpr (norm) { + for (int i = 0; i < 4; ++i) { + ((Array&)dst)[i] -= __ushort_as_half(0x6400U); } - return (Array&)dst; + } + return (Array&)dst; } -inline __device__ Array cvt_bf16x4_u8(const Array& src) -{ - // 01234567 01234567 01234567 01234567 - // SEEEEEEE EMMMMMMM MMMMMMMM MMMMMMMM - // 1MM... XXXXXXXX - // (1 + x/2^15) * 2^(e-127) -> e-127=15 -> e=142 -> 01000111 -> 0x47 - static constexpr uint32_t f32_magic = 0x47000000; // 32768 - - Array tmp; - tmp[0] = __byte_perm((uint32_t&)src, f32_magic, 0x7604); - tmp[1] = __byte_perm((uint32_t&)src, f32_magic, 0x7614); - tmp[2] = __byte_perm((uint32_t&)src, f32_magic, 0x7624); - tmp[3] = __byte_perm((uint32_t&)src, f32_magic, 0x7634); - - auto& vec = (Array&)tmp; - - Array dst; - PRAGMA_UNROLL - for (int i = 0; i < 4; ++i) { - dst[i] = __float2bfloat16(vec[i] - 32768.f); - } - return dst; +inline __device__ Array cvt_bf16x4_u8( + const Array& src) { + // 01234567 01234567 01234567 01234567 + // SEEEEEEE EMMMMMMM MMMMMMMM MMMMMMMM + // 1MM... XXXXXXXX + // (1 + x/2^15) * 2^(e-127) -> e-127=15 -> e=142 -> 01000111 -> 0x47 + static constexpr uint32_t f32_magic = 0x47000000; // 32768 + + Array tmp; + tmp[0] = __byte_perm((uint32_t&)src, f32_magic, 0x7604); + tmp[1] = __byte_perm((uint32_t&)src, f32_magic, 0x7614); + tmp[2] = __byte_perm((uint32_t&)src, f32_magic, 0x7624); + tmp[3] = __byte_perm((uint32_t&)src, f32_magic, 0x7634); + + auto& vec = (Array&)tmp; + + Array dst; + PRAGMA_UNROLL + for (int i = 0; i < 4; ++i) { + dst[i] = __float2bfloat16(vec[i] - 32768.f); + } + return dst; } -inline __device__ Array cvt_f32x4_u8(const Array& src) -{ - // 01234567 01234567 01234567 01234567 - // SEEEEEEE EMMMMMMM MMMMMMMM MMMMMMMM - // 1MM... XXXXXXXX - // (1 + x/2^15) * 2^(e-127) -> e-127=15 -> e=142 -> 01000111 -> 0x47 - static constexpr uint32_t f32_magic = 0x47000000; // 32768 - - Array tmp; - tmp[0] = __byte_perm((uint32_t&)src, f32_magic, 0x7604); - tmp[1] = __byte_perm((uint32_t&)src, f32_magic, 0x7614); - tmp[2] = __byte_perm((uint32_t&)src, f32_magic, 0x7624); - tmp[3] = __byte_perm((uint32_t&)src, f32_magic, 0x7634); - - auto& vec = (Array&)tmp; - PRAGMA_UNROLL - for (int i = 0; i < 4; ++i) { - vec[i] -= 32768.f; - } - return vec; +inline __device__ Array cvt_f32x4_u8(const Array& src) { + // 01234567 01234567 01234567 01234567 + // SEEEEEEE EMMMMMMM MMMMMMMM MMMMMMMM + // 1MM... XXXXXXXX + // (1 + x/2^15) * 2^(e-127) -> e-127=15 -> e=142 -> 01000111 -> 0x47 + static constexpr uint32_t f32_magic = 0x47000000; // 32768 + + Array tmp; + tmp[0] = __byte_perm((uint32_t&)src, f32_magic, 0x7604); + tmp[1] = __byte_perm((uint32_t&)src, f32_magic, 0x7614); + tmp[2] = __byte_perm((uint32_t&)src, f32_magic, 0x7624); + tmp[3] = __byte_perm((uint32_t&)src, f32_magic, 0x7634); + + auto& vec = (Array&)tmp; + PRAGMA_UNROLL + for (int i = 0; i < 4; ++i) { + vec[i] -= 32768.f; + } + return vec; } -template -inline __device__ Array cvt_bf16x8_u4(const Array& src) -{ +template +inline __device__ Array cvt_bf16x8_u4( + const Array& src) { #if __CUDA_ARCH__ >= 800 - // 01234567 01234567 - // SEEEEEEE EMMMMMMM - // 1...XXXX - // (1 + x/2^7) * 2^(e-127) -> e-127=7 -> e=134 -> 0100 0011 -> 0x43 - static constexpr uint32_t TEMPLATE = 0x43004300; // nv_bfloat162(128, 128) - static constexpr uint32_t MASK = 0x000f000f; - static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa; - - Array h; - - static_assert(sizeof(Array) == sizeof(Array)); - - uint32_t const& i4s = reinterpret_cast(src); - const uint32_t i4s_4 = i4s >> 4; - const uint32_t i4s_8 = i4s >> 8; - const uint32_t i4s_12 = i4s >> 12; - - asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[0]) : "r"(i4s), "n"(MASK), "n"(TEMPLATE), "n"(immLut)); - asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[1]) : "r"(i4s_4), "n"(MASK), "n"(TEMPLATE), "n"(immLut)); - asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[2]) : "r"(i4s_8), "n"(MASK), "n"(TEMPLATE), "n"(immLut)); - asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[3]) : "r"(i4s_12), "n"(MASK), "n"(TEMPLATE), "n"(immLut)); - - if constexpr (norm) { - auto result = reinterpret_cast(h.data()); - PRAGMA_UNROLL - for (int i = 0; i < 8; ++i) { - result[i] -= nv_bfloat16(128.f); - } + // 01234567 01234567 + // SEEEEEEE EMMMMMMM + // 1...XXXX + // (1 + x/2^7) * 2^(e-127) -> e-127=7 -> e=134 -> 0100 0011 -> 0x43 + static constexpr uint32_t TEMPLATE = 0x43004300; // nv_bfloat162(128, 128) + static constexpr uint32_t MASK = 0x000f000f; + static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa; + + Array h; + + static_assert(sizeof(Array) == sizeof(Array)); + + uint32_t const& i4s = reinterpret_cast(src); + const uint32_t i4s_4 = i4s >> 4; + const uint32_t i4s_8 = i4s >> 8; + const uint32_t i4s_12 = i4s >> 12; + + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[0]) + : "r"(i4s), "n"(MASK), "n"(TEMPLATE), "n"(immLut)); + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[1]) + : "r"(i4s_4), "n"(MASK), "n"(TEMPLATE), "n"(immLut)); + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[2]) + : "r"(i4s_8), "n"(MASK), "n"(TEMPLATE), "n"(immLut)); + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[3]) + : "r"(i4s_12), "n"(MASK), "n"(TEMPLATE), "n"(immLut)); + + if constexpr (norm) { + auto result = reinterpret_cast(h.data()); + PRAGMA_UNROLL + for (int i = 0; i < 8; ++i) { + result[i] -= nv_bfloat16(128.f); } - return (Array&)h; + } + return (Array&)h; #else - return {}; + return {}; #endif } #if TM_ROUND_USE_CVT_RNI -template -inline __device__ T round(float x) -{ - uint32_t y{}; - if constexpr (std::is_same_v) { - asm("cvt.rni.sat.u8.f32 %0, %1;\n" : "=r"(y) : "f"(x)); - } - else if constexpr (std::is_same_v) { - asm("cvt.rni.sat.u16.f32 %0, %1;\n" : "=r"(y) : "f"(x)); - } - else if constexpr (std::is_same_v) { - asm("cvt.rni.sat.u32.f32 %0, %1;\n" : "=r"(y) : "f"(x)); - } - else if constexpr (std::is_same_v) { - asm("cvt.rni.sat.s32.f32 %0, %1;\n" : "=r"(y) : "f"(x)); - } - else { - static_assert(!std::is_same_v, "not implemented"); - } - return y; +template +inline __device__ T round(float x) { + uint32_t y{}; + if constexpr (std::is_same_v) { + asm("cvt.rni.sat.u8.f32 %0, %1;\n" : "=r"(y) : "f"(x)); + } else if constexpr (std::is_same_v) { + asm("cvt.rni.sat.u16.f32 %0, %1;\n" : "=r"(y) : "f"(x)); + } else if constexpr (std::is_same_v) { + asm("cvt.rni.sat.u32.f32 %0, %1;\n" : "=r"(y) : "f"(x)); + } else if constexpr (std::is_same_v) { + asm("cvt.rni.sat.s32.f32 %0, %1;\n" : "=r"(y) : "f"(x)); + } else { + static_assert(!std::is_same_v, "not implemented"); + } + return y; } -template -inline __device__ T round(half x) -{ - uint32_t y{}; - if constexpr (std::is_same_v) { - asm("cvt.rni.sat.u8.f16 %0, %1;\n" : "=r"(y) : "h"((uint16_t&)x)); - } - else if constexpr (std::is_same_v) { - asm("cvt.rni.sat.u16.f16 %0, %1;\n" : "=r"(y) : "h"((uint16_t&)x)); - } - else if constexpr (std::is_same_v) { - asm("cvt.rni.sat.u32.f16 %0, %1;\n" : "=r"(y) : "h"((uint16_t&)x)); - } - else if constexpr (std::is_same_v) { - asm("cvt.rni.sat.s32.f16 %0, %1;\n" : "=r"(y) : "h"((uint16_t&)x)); - } - else { - static_assert(!std::is_same_v, "not implemented"); - } - return y; +template +inline __device__ T round(half x) { + uint32_t y{}; + if constexpr (std::is_same_v) { + asm("cvt.rni.sat.u8.f16 %0, %1;\n" : "=r"(y) : "h"((uint16_t&)x)); + } else if constexpr (std::is_same_v) { + asm("cvt.rni.sat.u16.f16 %0, %1;\n" : "=r"(y) : "h"((uint16_t&)x)); + } else if constexpr (std::is_same_v) { + asm("cvt.rni.sat.u32.f16 %0, %1;\n" : "=r"(y) : "h"((uint16_t&)x)); + } else if constexpr (std::is_same_v) { + asm("cvt.rni.sat.s32.f16 %0, %1;\n" : "=r"(y) : "h"((uint16_t&)x)); + } else { + static_assert(!std::is_same_v, "not implemented"); + } + return y; } #else -template -inline __device__ T round(float x) -{ - x += .5f; - - uint32_t y{}; - if constexpr (std::is_same_v) { - asm("cvt.rzi.sat.u8.f32 %0, %1;\n" : "=r"(y) : "f"(x)); - } - else if constexpr (std::is_same_v) { - asm("cvt.rzi.sat.u16.f32 %0, %1;\n" : "=r"(y) : "f"(x)); - } - else if constexpr (std::is_same_v) { - asm("cvt.rzi.sat.u32.f32 %0, %1;\n" : "=r"(y) : "f"(x)); - } - else { - static_assert(!std::is_same_v, "not implemented"); - } - return y; +template +inline __device__ T round(float x) { + x += .5f; + + uint32_t y{}; + if constexpr (std::is_same_v) { + asm("cvt.rzi.sat.u8.f32 %0, %1;\n" : "=r"(y) : "f"(x)); + } else if constexpr (std::is_same_v) { + asm("cvt.rzi.sat.u16.f32 %0, %1;\n" : "=r"(y) : "f"(x)); + } else if constexpr (std::is_same_v) { + asm("cvt.rzi.sat.u32.f32 %0, %1;\n" : "=r"(y) : "f"(x)); + } else { + static_assert(!std::is_same_v, "not implemented"); + } + return y; } -template -inline __device__ T round(half x) -{ - x += half(.5f); - - uint32_t y{}; - if constexpr (std::is_same_v) { - asm("cvt.rzi.sat.u8.f16 %0, %1;\n" : "=r"(y) : "h"((uint16_t&)x)); - } - else if constexpr (std::is_same_v) { - asm("cvt.rzi.sat.u16.f16 %0, %1;\n" : "=r"(y) : "h"((uint16_t&)x)); - } - else if constexpr (std::is_same_v) { - asm("cvt.rzi.sat.u32.f16 %0, %1;\n" : "=r"(y) : "h"((uint16_t&)x)); - } - else { - static_assert(!std::is_same_v, "not implemented"); - } - return y; +template +inline __device__ T round(half x) { + x += half(.5f); + + uint32_t y{}; + if constexpr (std::is_same_v) { + asm("cvt.rzi.sat.u8.f16 %0, %1;\n" : "=r"(y) : "h"((uint16_t&)x)); + } else if constexpr (std::is_same_v) { + asm("cvt.rzi.sat.u16.f16 %0, %1;\n" : "=r"(y) : "h"((uint16_t&)x)); + } else if constexpr (std::is_same_v) { + asm("cvt.rzi.sat.u32.f16 %0, %1;\n" : "=r"(y) : "h"((uint16_t&)x)); + } else { + static_assert(!std::is_same_v, "not implemented"); + } + return y; } #endif -template -inline __device__ To quant(Ti x, B n_bits) -{ - auto y = round(x); - if constexpr (n_bits < sizeof(To) * 8) { // saturate operation for sub-byte type - return min(y, To((1 << n_bits) - 1)); - } - else { - return y; - } +template +inline __device__ To quant(Ti x, B n_bits) { + auto y = round(x); + if constexpr (n_bits < + sizeof(To) * 8) { // saturate operation for sub-byte type + return min(y, To((1 << n_bits) - 1)); + } else { + return y; + } } -template -__device__ inline void warp_minmax(Array& stats, const Array& x) -{ +template +__device__ inline void warp_minmax(Array& stats, const Array& x) { + PRAGMA_UNROLL + for (int i = 0; i < C; ++i) { + stats[0] = Min(stats[0], x[i]); + stats[1] = Max(stats[1], x[i]); + } + if constexpr (sizeof(T) == 2) { PRAGMA_UNROLL - for (int i = 0; i < C; ++i) { - stats[0] = Min(stats[0], x[i]); - stats[1] = Max(stats[1], x[i]); - } - if constexpr (sizeof(T) == 2) { - PRAGMA_UNROLL - for (int mask = WarpThreadC / 2; mask > 0; mask /= 2) { - Array tmp; - (uint32_t&)tmp = __shfl_xor_sync(uint32_t(-1), (uint32_t&)stats, mask); - stats[0] = Min(stats[0], tmp[0]); - stats[1] = Max(stats[1], tmp[1]); - } + for (int mask = WarpThreadC / 2; mask > 0; mask /= 2) { + Array tmp; + (uint32_t&)tmp = __shfl_xor_sync(uint32_t(-1), (uint32_t&)stats, mask); + stats[0] = Min(stats[0], tmp[0]); + stats[1] = Max(stats[1], tmp[1]); } - else { - PRAGMA_UNROLL - for (int mask = WarpThreadC / 2; mask > 0; mask /= 2) { - stats[0] = Min(stats[0], __shfl_xor_sync(uint32_t(-1), stats[0], mask)); - stats[1] = Max(stats[1], __shfl_xor_sync(uint32_t(-1), stats[1], mask)); - } + } else { + PRAGMA_UNROLL + for (int mask = WarpThreadC / 2; mask > 0; mask /= 2) { + stats[0] = Min(stats[0], __shfl_xor_sync(uint32_t(-1), stats[0], mask)); + stats[1] = Max(stats[1], __shfl_xor_sync(uint32_t(-1), stats[1], mask)); } + } } -template -__device__ void warp_stats(Array (¶m)[S], const Array (&x)[S][C], B n_bits) -{ +template +__device__ void warp_stats(Array (¶m)[S], + const Array (&x)[S][C], B n_bits) { + PRAGMA_UNROLL + for (int s = 0; s < S; ++s) { + Array stats{Infinity(), -Infinity()}; PRAGMA_UNROLL - for (int s = 0; s < S; ++s) { - Array stats{Infinity(), -Infinity()}; - PRAGMA_UNROLL - for (int c = 0; c < C; ++c) { - warp_minmax(stats, x[s][c]); - } - const float inv_q_max = fdividef(1.f, float((1 << n_bits) - 1)); - const float scale = ((float)stats[1] - (float)stats[0]) * inv_q_max; - param[s][0] = (P)scale; - param[s][1] = (P)stats[0]; - - if constexpr (kForceIntZeroPoint) { + for (int c = 0; c < C; ++c) { + warp_minmax(stats, x[s][c]); + } + const float inv_q_max = fdividef(1.f, float((1 << n_bits) - 1)); + const float scale = ((float)stats[1] - (float)stats[0]) * inv_q_max; + param[s][0] = (P)scale; + param[s][1] = (P)stats[0]; + + if constexpr (kForceIntZeroPoint) { #if TM_ROUND_USE_CVT_RNI - // rintf -> cvt.rni.f32.f32 - param[s][1] = (P)(rintf((float)stats[0] / scale) * scale); + // rintf -> cvt.rni.f32.f32 + param[s][1] = (P)(rintf((float)stats[0] / scale) * scale); #else - // roundf -> cvt.rzi.f32.f32(x + 0.5) - param[s][1] = (P)(roundf((float)stats[0] / scale) * scale); + // roundf -> cvt.rzi.f32.f32(x + 0.5) + param[s][1] = (P)(roundf((float)stats[0] / scale) * scale); #endif - } } + } } -template -__device__ void -quantize(Array (&dst)[S][C], const Array (&src)[S][C], const Array (¶ms)[S], B n_bits) -{ +template +__device__ void quantize(Array (&dst)[S][C], + const Array (&src)[S][C], + const Array (¶ms)[S], B n_bits) { + PRAGMA_UNROLL + for (int s = 0; s < S; ++s) { + P inv_scale = (P)fdividef(1.f, (float)params[s][0]); + P zero = params[s][1]; PRAGMA_UNROLL - for (int s = 0; s < S; ++s) { - P inv_scale = (P)fdividef(1.f, (float)params[s][0]); - P zero = params[s][1]; - PRAGMA_UNROLL - for (int c = 0; c < C; ++c) { - PRAGMA_UNROLL - for (int i = 0; i < N; ++i) { - const auto v = ((P)src[s][c][i] - zero) * inv_scale; - dst[s][c][i] = quant(v, n_bits); - } - } - } + for (int c = 0; c < C; ++c) { + PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + const auto v = ((P)src[s][c][i] - zero) * inv_scale; + dst[s][c][i] = quant(v, n_bits); + } + } + } } ////////////////////////////////////////////////////////////////////////////////////////////////// -// generic case for floating point -> floating point / integer -> integer conversion -template +// generic case for floating point -> floating point / integer -> integer +// conversion +template struct ConvertKvCache { - __device__ __host__ ConvertKvCache(float, float) {} - template - __device__ static auto convert(const Array& vi) - { - Array vo; - PRAGMA_UNROLL - for (int i = 0; i < N; ++i) { - vo[i] = (To)vi[i]; - } - return vo; - } - template - inline __device__ auto operator()(const Array& vi) const -> Array - { - return convert(vi); - } + __device__ __host__ ConvertKvCache(float, float) {} + template + __device__ static auto convert(const Array& vi) { + Array vo; + PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + vo[i] = (To)vi[i]; + } + return vo; + } + template + inline __device__ auto operator()(const Array& vi) const + -> Array { + return convert(vi); + } }; // generic case for converting to same type, bypass -template +template struct ConvertKvCache { - __device__ __host__ ConvertKvCache(float, float) {} - template - __device__ static auto convert(const Array& v) - { - return v; - } - template - inline __device__ auto operator()(const Array& v) const -> Array - { - return convert(v); - } + __device__ __host__ ConvertKvCache(float, float) {} + template + __device__ static auto convert(const Array& v) { + return v; + } + template + inline __device__ auto operator()(const Array& v) const -> Array { + return convert(v); + } }; // floating point -> u8 -template +template struct ConvertKvCache { - T inv_scale_; - T zero_; - __device__ ConvertKvCache(T scale, T zero): zero_{zero} - { - // NVCC complains if we put this in the member init list - inv_scale_ = (T)fdividef(1.f, (float)scale); - } - - template - __device__ auto operator()(const Array& vi) const - { - Array vo; - PRAGMA_UNROLL - for (int i = 0; i < N; ++i) { - vo[i] = quant((vi[i] - zero_) * inv_scale_, std::integral_constant{}); - } - return vo; + T inv_scale_; + T zero_; + __device__ ConvertKvCache(T scale, T zero) : zero_{zero} { + // NVCC complains if we put this in the member init list + inv_scale_ = (T)fdividef(1.f, (float)scale); + } + + template + __device__ auto operator()(const Array& vi) const { + Array vo; + PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + vo[i] = quant((vi[i] - zero_) * inv_scale_, + std::integral_constant{}); } + return vo; + } }; -template +template struct ConvertKvCache { - T inv_scale_; - T zero_; - __device__ ConvertKvCache(T scale, T zero): zero_{zero} - { - // NVCC complains if we put this in the member init list - inv_scale_ = (T)fdividef(1.f, (float)scale); - } - - static __device__ Array pack(const Array& vi) - { - Array ui = (Array&)vi; - - ui[0] |= (ui[0] >> 12); - ui[1] |= (ui[1] >> 12); - - // 7 6 5 4 3 2 1 0 - // _7_67564_3_23120 - uint32_t uo = __byte_perm(ui[0], ui[1], 0x5140); - - return (Array&)uo; + T inv_scale_; + T zero_; + __device__ ConvertKvCache(T scale, T zero) : zero_{zero} { + // NVCC complains if we put this in the member init list + inv_scale_ = (T)fdividef(1.f, (float)scale); + } + + static __device__ Array pack(const Array& vi) { + Array ui = (Array&)vi; + + ui[0] |= (ui[0] >> 12); + ui[1] |= (ui[1] >> 12); + + // 7 6 5 4 3 2 1 0 + // _7_67564_3_23120 + uint32_t uo = __byte_perm(ui[0], ui[1], 0x5140); + + return (Array&)uo; + } + + /// TODO: try cvt.pack.sat.u4 + template + __device__ auto operator()(const Array& vi) const { + static_assert(N % 8 == 0); + Array tmp; + PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + tmp[i] = quant((vi[i] - zero_) * inv_scale_, + std::integral_constant{}); } - - /// TODO: try cvt.pack.sat.u4 - template - __device__ auto operator()(const Array& vi) const - { - static_assert(N % 8 == 0); - Array tmp; - PRAGMA_UNROLL - for (int i = 0; i < N; ++i) { - tmp[i] = quant((vi[i] - zero_) * inv_scale_, std::integral_constant{}); - } - Array vo; - PRAGMA_UNROLL - for (int i = 0; i < N; i += 8) { - (Array&)vo[i] = pack((Array&)tmp[i]); - } - return vo; + Array vo; + PRAGMA_UNROLL + for (int i = 0; i < N; i += 8) { + (Array&)vo[i] = pack((Array&)tmp[i]); } + return vo; + } }; -template<> +template <> struct ConvertKvCache { - - half scale_; - half zero_; - - __device__ ConvertKvCache(half scale, half zero) - { - scale_ = scale; - zero_ = zero; - } - - static __device__ Array cvt_f16x8_u4(const Array& vi) - { - Array result; - uint32_t* h = reinterpret_cast(&result); - uint32_t const& i4s = reinterpret_cast(vi); - static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa; - static constexpr uint32_t BOT_MASK = 0x000f000f; - static constexpr uint32_t TOP_MASK = 0x00f000f0; - static constexpr uint32_t MAGIC_NUM_0 = 0x64006400; // `1024` - static constexpr uint32_t MAGIC_NUM_1 = 0x54005400; // `64` - // const uint32_t top_i4s = i4s >> 8; - uint32_t top_i4s = __byte_perm(i4s, 0, 0x4321); - asm("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[0]) : "r"(i4s), "n"(BOT_MASK), "n"(MAGIC_NUM_0), "n"(immLut)); - asm("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[1]) : "r"(i4s), "n"(TOP_MASK), "n"(MAGIC_NUM_1), "n"(immLut)); - asm("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[2]) : "r"(top_i4s), "n"(BOT_MASK), "n"(MAGIC_NUM_0), "n"(immLut)); - asm("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[3]) : "r"(top_i4s), "n"(TOP_MASK), "n"(MAGIC_NUM_1), "n"(immLut)); - asm("sub.f16x2 %0, %1, %2;\n" : "=r"(h[0]) : "r"(h[0]), "r"(MAGIC_NUM_0)); - asm("sub.f16x2 %0, %1, %2;\n" : "=r"(h[1]) : "r"(h[1]), "r"(MAGIC_NUM_1)); - asm("sub.f16x2 %0, %1, %2;\n" : "=r"(h[2]) : "r"(h[2]), "r"(MAGIC_NUM_0)); - asm("sub.f16x2 %0, %1, %2;\n" : "=r"(h[3]) : "r"(h[3]), "r"(MAGIC_NUM_1)); - return result; - } - - static __device__ Array cvt_f16x8_u4_biased(const Array& vi) - { - Array result; - uint32_t* h = reinterpret_cast(&result); - uint32_t const& i4s = reinterpret_cast(vi); - static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa; - static constexpr uint32_t BOT_MASK = 0x000f000f; - static constexpr uint32_t TOP_MASK = 0x00f000f0; - static constexpr uint32_t MAGIC_NUM_1 = 0x54005400; // `64` - static constexpr uint32_t MAGIC_NUM_2 = MAGIC_NUM_1 >> 4; // `64` >> 4 - const uint32_t top_i4s = i4s >> 8; - // uint32_t top_i4s = __byte_perm(i4s, 0, 0x4321); - asm("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[0]) : "r"(i4s), "n"(BOT_MASK), "n"(MAGIC_NUM_2), "n"(immLut)); - asm("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[1]) : "r"(i4s), "n"(TOP_MASK), "n"(MAGIC_NUM_1), "n"(immLut)); - asm("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[2]) : "r"(top_i4s), "n"(BOT_MASK), "n"(MAGIC_NUM_2), "n"(immLut)); - asm("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[3]) : "r"(top_i4s), "n"(TOP_MASK), "n"(MAGIC_NUM_1), "n"(immLut)); - h[0] <<= 4; - h[2] <<= 4; - return result; - } - - template - __device__ static auto convert(const Array& vi) - { - Array vo; - PRAGMA_UNROLL - for (int i = 0; i < N; i += 8) { - auto& v = (Array&)vo[i]; - if constexpr (kFuseU4F16Dequant) { - v = cvt_f16x8_u4_biased((Array&)vi[i]); - } - else { - v = cvt_f16x8_u4((Array&)vi[i]); - } - } - return vo; - } - - template - __device__ auto operator()(const Array& vi) const - { - auto vo = convert(vi); - PRAGMA_UNROLL - for (int i = 0; i < N; ++i) { - vo[i] = vo[i] * scale_ + zero_; - } - return vo; + half scale_; + half zero_; + + __device__ ConvertKvCache(half scale, half zero) { + scale_ = scale; + zero_ = zero; + } + + static __device__ Array cvt_f16x8_u4(const Array& vi) { + Array result; + uint32_t* h = reinterpret_cast(&result); + uint32_t const& i4s = reinterpret_cast(vi); + static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa; + static constexpr uint32_t BOT_MASK = 0x000f000f; + static constexpr uint32_t TOP_MASK = 0x00f000f0; + static constexpr uint32_t MAGIC_NUM_0 = 0x64006400; // `1024` + static constexpr uint32_t MAGIC_NUM_1 = 0x54005400; // `64` + // const uint32_t top_i4s = i4s >> 8; + uint32_t top_i4s = __byte_perm(i4s, 0, 0x4321); + asm("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[0]) + : "r"(i4s), "n"(BOT_MASK), "n"(MAGIC_NUM_0), "n"(immLut)); + asm("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[1]) + : "r"(i4s), "n"(TOP_MASK), "n"(MAGIC_NUM_1), "n"(immLut)); + asm("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[2]) + : "r"(top_i4s), "n"(BOT_MASK), "n"(MAGIC_NUM_0), "n"(immLut)); + asm("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[3]) + : "r"(top_i4s), "n"(TOP_MASK), "n"(MAGIC_NUM_1), "n"(immLut)); + asm("sub.f16x2 %0, %1, %2;\n" : "=r"(h[0]) : "r"(h[0]), "r"(MAGIC_NUM_0)); + asm("sub.f16x2 %0, %1, %2;\n" : "=r"(h[1]) : "r"(h[1]), "r"(MAGIC_NUM_1)); + asm("sub.f16x2 %0, %1, %2;\n" : "=r"(h[2]) : "r"(h[2]), "r"(MAGIC_NUM_0)); + asm("sub.f16x2 %0, %1, %2;\n" : "=r"(h[3]) : "r"(h[3]), "r"(MAGIC_NUM_1)); + return result; + } + + static __device__ Array cvt_f16x8_u4_biased( + const Array& vi) { + Array result; + uint32_t* h = reinterpret_cast(&result); + uint32_t const& i4s = reinterpret_cast(vi); + static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa; + static constexpr uint32_t BOT_MASK = 0x000f000f; + static constexpr uint32_t TOP_MASK = 0x00f000f0; + static constexpr uint32_t MAGIC_NUM_1 = 0x54005400; // `64` + static constexpr uint32_t MAGIC_NUM_2 = MAGIC_NUM_1 >> 4; // `64` >> 4 + const uint32_t top_i4s = i4s >> 8; + // uint32_t top_i4s = __byte_perm(i4s, 0, 0x4321); + asm("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[0]) + : "r"(i4s), "n"(BOT_MASK), "n"(MAGIC_NUM_2), "n"(immLut)); + asm("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[1]) + : "r"(i4s), "n"(TOP_MASK), "n"(MAGIC_NUM_1), "n"(immLut)); + asm("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[2]) + : "r"(top_i4s), "n"(BOT_MASK), "n"(MAGIC_NUM_2), "n"(immLut)); + asm("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[3]) + : "r"(top_i4s), "n"(TOP_MASK), "n"(MAGIC_NUM_1), "n"(immLut)); + h[0] <<= 4; + h[2] <<= 4; + return result; + } + + template + __device__ static auto convert(const Array& vi) { + Array vo; + PRAGMA_UNROLL + for (int i = 0; i < N; i += 8) { + auto& v = (Array&)vo[i]; + if constexpr (kFuseU4F16Dequant) { + v = cvt_f16x8_u4_biased((Array&)vi[i]); + } else { + v = cvt_f16x8_u4((Array&)vi[i]); + } + } + return vo; + } + + template + __device__ auto operator()(const Array& vi) const { + auto vo = convert(vi); + PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + vo[i] = vo[i] * scale_ + zero_; } + return vo; + } }; -template<> +template <> struct ConvertKvCache { + nv_bfloat16 scale_; + nv_bfloat16 zero_; - nv_bfloat16 scale_; - nv_bfloat16 zero_; - - __device__ ConvertKvCache(nv_bfloat16 scale, nv_bfloat16 zero) - { - scale_ = scale; - zero_ = zero; - } + __device__ ConvertKvCache(nv_bfloat16 scale, nv_bfloat16 zero) { + scale_ = scale; + zero_ = zero; + } - template - __device__ static Array convert(const Array& vi) - { - Array vo{}; - PRAGMA_UNROLL - for (int i = 0; i < N; i += 8) { - auto& v = (Array&)vo[i]; - auto u = cvt_bf16x8_u4((Array&)vi[i]); - v = (Array&)u; - } - return vo; - } - - template - __device__ Array operator()(const Array& vi) const - { - auto vo = convert(vi); - PRAGMA_UNROLL - for (int i = 0; i < N; ++i) { - vo[i] = vo[i] * scale_ + zero_; - } - return (Array&)vo; + template + __device__ static Array convert(const Array& vi) { + Array vo{}; + PRAGMA_UNROLL + for (int i = 0; i < N; i += 8) { + auto& v = (Array&)vo[i]; + auto u = cvt_bf16x8_u4((Array&)vi[i]); + v = (Array&)u; + } + return vo; + } + + template + __device__ Array operator()( + const Array& vi) const { + auto vo = convert(vi); + PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + vo[i] = vo[i] * scale_ + zero_; } + return (Array&)vo; + } }; -template<> +template <> struct ConvertKvCache { - #if 1 - ConvertKvCache impl_; + ConvertKvCache impl_; - __device__ ConvertKvCache(float scale, float zero): impl_{scale, zero} {} + __device__ ConvertKvCache(float scale, float zero) : impl_{scale, zero} {} - template - __device__ auto operator()(const Array& vi) const - { - return cast(impl_(vi)); - } + template + __device__ auto operator()(const Array& vi) const { + return cast(impl_(vi)); + } #else - static __device__ Array cvt_f16x8_u4_biased(const Array& vi) - { - Array result; - uint32_t* h = reinterpret_cast(&result); - uint32_t const& i4s = reinterpret_cast(vi); - static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa; - static constexpr uint32_t BOT_MASK = 0x000f000f; - static constexpr uint32_t TOP_MASK = 0x00f000f0; - static constexpr uint32_t MAGIC_NUM_1 = 0x54005400; // `64` - static constexpr uint32_t MAGIC_NUM_2 = MAGIC_NUM_1 >> 4; // `64` >> 4 - const uint32_t top_i4s = i4s >> 8; - asm("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[0]) : "r"(i4s), "n"(BOT_MASK), "n"(MAGIC_NUM_2), "n"(immLut)); - asm("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[1]) : "r"(i4s), "n"(TOP_MASK), "n"(MAGIC_NUM_1), "n"(immLut)); - asm("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[2]) : "r"(top_i4s), "n"(BOT_MASK), "n"(MAGIC_NUM_2), "n"(immLut)); - asm("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[3]) : "r"(top_i4s), "n"(TOP_MASK), "n"(MAGIC_NUM_1), "n"(immLut)); - h[0] <<= 4; - h[2] <<= 4; - return result; - } - float scale_; - float zero_; - __device__ ConvertKvCache(float scale, float zero) - { - scale_ = scale; - zero_ = zero - scale * 64.f; - } - template - __device__ auto operator()(const Array& vi) const - { - auto vo = cast(cvt_f16x8_u4_biased(vi)); - using namespace ops; - return vo * scale_ + zero_; - } + static __device__ Array cvt_f16x8_u4_biased( + const Array& vi) { + Array result; + uint32_t* h = reinterpret_cast(&result); + uint32_t const& i4s = reinterpret_cast(vi); + static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa; + static constexpr uint32_t BOT_MASK = 0x000f000f; + static constexpr uint32_t TOP_MASK = 0x00f000f0; + static constexpr uint32_t MAGIC_NUM_1 = 0x54005400; // `64` + static constexpr uint32_t MAGIC_NUM_2 = MAGIC_NUM_1 >> 4; // `64` >> 4 + const uint32_t top_i4s = i4s >> 8; + asm("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[0]) + : "r"(i4s), "n"(BOT_MASK), "n"(MAGIC_NUM_2), "n"(immLut)); + asm("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[1]) + : "r"(i4s), "n"(TOP_MASK), "n"(MAGIC_NUM_1), "n"(immLut)); + asm("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[2]) + : "r"(top_i4s), "n"(BOT_MASK), "n"(MAGIC_NUM_2), "n"(immLut)); + asm("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[3]) + : "r"(top_i4s), "n"(TOP_MASK), "n"(MAGIC_NUM_1), "n"(immLut)); + h[0] <<= 4; + h[2] <<= 4; + return result; + } + float scale_; + float zero_; + __device__ ConvertKvCache(float scale, float zero) { + scale_ = scale; + zero_ = zero - scale * 64.f; + } + template + __device__ auto operator()(const Array& vi) const { + auto vo = cast(cvt_f16x8_u4_biased(vi)); + using namespace ops; + return vo * scale_ + zero_; + } #endif }; // u8 -> f32/f16/bf16 -template +template struct ConvertKvCache { - T scale_; - T zero_; - __device__ ConvertKvCache(T scale, T zero): scale_{scale}, zero_{zero} {} - - template - __device__ static auto convert(const Array& vi) - { - Array vo; - PRAGMA_UNROLL - for (int n = 0; n < N; n += 4) { - auto& ui = (const Array&)vi[n]; - auto& uo = (Array&)vo[n]; - - if constexpr (std::is_same_v) { - uo = cvt_f16x4_u8(ui); - } - else if constexpr (std::is_same_v) { - uo = cvt_f32x4_u8(ui); - } + T scale_; + T zero_; + __device__ ConvertKvCache(T scale, T zero) : scale_{scale}, zero_{zero} {} + + template + __device__ static auto convert(const Array& vi) { + Array vo; + PRAGMA_UNROLL + for (int n = 0; n < N; n += 4) { + auto& ui = (const Array&)vi[n]; + auto& uo = (Array&)vo[n]; + + if constexpr (std::is_same_v) { + uo = cvt_f16x4_u8(ui); + } else if constexpr (std::is_same_v) { + uo = cvt_f32x4_u8(ui); + } #if __CUDA_ARCH__ >= 800 - else if constexpr (std::is_same_v) { - uo = cvt_bf16x4_u8(ui); - } + else if constexpr (std::is_same_v) { + uo = cvt_bf16x4_u8(ui); + } #endif - } - return vo; } + return vo; + } - template - __device__ auto operator()(const Array& vi) const - { - auto vo = convert(vi); - PRAGMA_UNROLL - for (int i = 0; i < N; ++i) { - vo[i] = vo[i] * scale_ + zero_; - } - return vo; + template + __device__ auto operator()(const Array& vi) const { + auto vo = convert(vi); + PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + vo[i] = vo[i] * scale_ + zero_; } + return vo; + } }; -template +template struct ConvertKvCache { + __device__ static Array cvt_bf16x8_e2m1( + const Array& vi) { + const uint32_t& x = (const uint32_t&)vi; - __device__ static Array cvt_bf16x8_e2m1(const Array& vi) - { - const uint32_t& x = (const uint32_t&)vi; - - constexpr uint32_t S = 0x80008000U; - constexpr uint32_t EM = 0x01C001C0U; + constexpr uint32_t S = 0x80008000U; + constexpr uint32_t EM = 0x01C001C0U; - Array vo; + Array vo; - // clang-format off + // clang-format off vo[0] = (x << 12 & S) | (x << 6 & EM); vo[1] = (x << 8 & S) | (x << 2 & EM); vo[2] = (x << 4 & S) | (x >> 2 & EM); vo[3] = (x << 0 & S) | (x >> 6 & EM); - // clang-format on + // clang-format on - constexpr uint32_t e = (127U - 1U + 127U) << 7U; - constexpr uint32_t ee = e << 16U | e; + constexpr uint32_t e = (127U - 1U + 127U) << 7U; + constexpr uint32_t ee = e << 16U | e; - PRAGMA_UNROLL - for (int i = 0; i < 4; ++i) { + PRAGMA_UNROLL + for (int i = 0; i < 4; ++i) { #if TURBOMIND_ARCH_SM90 - asm("mul.rn.bf16x2 %0, %1, %2;" : "=r"(vo[i]) : "r"(vo[i]), "r"(ee)); + asm("mul.rn.bf16x2 %0, %1, %2;" : "=r"(vo[i]) : "r"(vo[i]), "r"(ee)); #else - asm("fma.rn.bf16x2 %0, %1, %2, %3;" : "=r"(vo[i]) : "r"(vo[i]), "r"(ee), "r"(0)); + asm("fma.rn.bf16x2 %0, %1, %2, %3;" + : "=r"(vo[i]) + : "r"(vo[i]), "r"(ee), "r"(0)); #endif - } - - return (Array&)vo; } - __device__ static Array cvt_f16x8_e2m1(const Array& vi) - { - const uint32_t& x = (const uint32_t&)vi; + return (Array&)vo; + } + + __device__ static Array cvt_f16x8_e2m1( + const Array& vi) { + const uint32_t& x = (const uint32_t&)vi; - constexpr uint32_t S = 0x80008000U; - constexpr uint32_t EM = 0x0E000E00U; + constexpr uint32_t S = 0x80008000U; + constexpr uint32_t EM = 0x0E000E00U; - Array vo; + Array vo; - // clang-format off + // clang-format off vo[0] = (x << 12 & S) | (x << 9 & EM); vo[1] = (x << 8 & S) | (x << 5 & EM); vo[2] = (x << 4 & S) | (x << 1 & EM); vo[3] = (x << 0 & S) | (x >> 3 & EM); - // clang-format on + // clang-format on - constexpr uint32_t e = (15U - 1U + 15U) << 10U; - constexpr uint32_t ee = e << 16U | e; - - PRAGMA_UNROLL - for (int i = 0; i < 4; ++i) { - asm volatile("mul.f16x2 %0, %1, %2;" : "=r"(vo[i]) : "r"(vo[i]), "r"(ee)); - } + constexpr uint32_t e = (15U - 1U + 15U) << 10U; + constexpr uint32_t ee = e << 16U | e; - return (Array&)vo; + PRAGMA_UNROLL + for (int i = 0; i < 4; ++i) { + asm volatile("mul.f16x2 %0, %1, %2;" : "=r"(vo[i]) : "r"(vo[i]), "r"(ee)); } - template - __device__ static auto convert(const Array& vi) - { - Array vo; - PRAGMA_UNROLL - for (int i = 0; i < N; i += 8) { - auto& v = (Array&)vo[i]; - if constexpr (std::is_same_v) { - v = cvt_bf16x8_e2m1((Array&)vi[i]); - } - else if constexpr (std::is_same_v) { - v = cvt_f16x8_e2m1((Array&)vi[i]); - } - else { - static_assert(N != N, "not implemented"); - } - } - return vo; - } + return (Array&)vo; + } + + template + __device__ static auto convert(const Array& vi) { + Array vo; + PRAGMA_UNROLL + for (int i = 0; i < N; i += 8) { + auto& v = (Array&)vo[i]; + if constexpr (std::is_same_v) { + v = cvt_bf16x8_e2m1((Array&)vi[i]); + } else if constexpr (std::is_same_v) { + v = cvt_f16x8_e2m1((Array&)vi[i]); + } else { + static_assert(N != N, "not implemented"); + } + } + return vo; + } }; -__device__ inline Array cvt_bf16x4_e4m3(const Array& vi) -{ - const uint32_t& x = (const uint32_t&)vi; +__device__ inline Array cvt_bf16x4_e4m3( + const Array& vi) { + const uint32_t& x = (const uint32_t&)vi; - // 0 7 C 0 - // SEEEEEEEEMMMMMMMSEEEEEEEEMMMMMMM - // SEEEEMMM SEEEEMMM - // SEEEEMMM SEEEEMMM + // 0 7 C 0 + // SEEEEEEEEMMMMMMMSEEEEEEEEMMMMMMM + // SEEEEMMM SEEEEMMM + // SEEEEMMM SEEEEMMM - constexpr uint32_t S = 0x80008000U; - constexpr uint32_t EM = 0x07F007F0U; + constexpr uint32_t S = 0x80008000U; + constexpr uint32_t EM = 0x07F007F0U; - Array vo; + Array vo; - vo[0] = (x << 8 & S) | (x << 4 & EM); - vo[1] = (x << 0 & S) | (x >> 4 & EM); + vo[0] = (x << 8 & S) | (x << 4 & EM); + vo[1] = (x << 0 & S) | (x >> 4 & EM); - constexpr uint32_t e = (127U - 7U + 127U) << 7U; - constexpr uint32_t ee = e << 16U | e; + constexpr uint32_t e = (127U - 7U + 127U) << 7U; + constexpr uint32_t ee = e << 16U | e; - PRAGMA_UNROLL - for (int i = 0; i < 2; ++i) { + PRAGMA_UNROLL + for (int i = 0; i < 2; ++i) { #if TURBOMIND_ARCH_SM90 - asm("mul.rn.bf16x2 %0, %1, %2;" : "=r"(vo[i]) : "r"(vo[i]), "r"(ee)); + asm("mul.rn.bf16x2 %0, %1, %2;" : "=r"(vo[i]) : "r"(vo[i]), "r"(ee)); #else - asm("fma.rn.bf16x2 %0, %1, %2, %3;" : "=r"(vo[i]) : "r"(vo[i]), "r"(ee), "r"(0)); + asm("fma.rn.bf16x2 %0, %1, %2, %3;" + : "=r"(vo[i]) + : "r"(vo[i]), "r"(ee), "r"(0)); #endif - } + } - return (Array&)vo; + return (Array&)vo; } -__device__ inline Array cvt_f16x4_e4m3(const Array& vi) -{ - const uint32_t& x = (const uint32_t&)vi; +template +__device__ inline Array cvt_f16x4_e4m3( + const Array& vi) { + const uint32_t& x = (const uint32_t&)vi; - // 3 F 8 0 - // SEEEEEMMMMMMMMMMSEEEEEMMMMMMMMMM - // SEEEEMMM SEEEEMMM - // SEEEEMMM SEEEEMMM + // 3 F 8 0 + // SEEEEEMMMMMMMMMMSEEEEEMMMMMMMMMM + // SEEEEMMM SEEEEMMM + // SEEEEMMM SEEEEMMM - constexpr uint32_t S = 0x80008000U; - constexpr uint32_t EM = 0x3F803F80U; + constexpr uint32_t S = 0x80008000U; + constexpr uint32_t EM = 0x3F803F80U; - Array vo; + Array vo; - vo[0] = (x << 8 & S) | (x << 7 & EM); - vo[1] = (x << 0 & S) | (x >> 1 & EM); + vo[0] = (x << 8 & S) | (x << 7 & EM); + vo[1] = (x << 0 & S) | (x >> 1 & EM); - constexpr uint32_t e = (15U - 7U + 15U) << 10U; - constexpr uint32_t ee = e << 16U | e; + constexpr uint32_t e = (15U - 7U + 15U) << 10U; + constexpr uint32_t ee = e << 16U | e; + if constexpr (ApplyExponentBias) { PRAGMA_UNROLL for (int i = 0; i < 2; ++i) { - asm("mul.rn.f16x2 %0, %1, %2;" : "=r"(vo[i]) : "r"(vo[i]), "r"(ee)); + asm("mul.rn.f16x2 %0, %1, %2;" : "=r"(vo[i]) : "r"(vo[i]), "r"(ee)); } + } - return (Array&)vo; + return (Array&)vo; } -template +template struct ConvertKvCache { - - template - __device__ static auto convert(const Array& vi) - { - static_assert(N % 4 == 0); - Array vo; - PRAGMA_UNROLL - for (int i = 0; i < N; i += 4) { - auto& v = (Array&)vo[i]; - if constexpr (std::is_same_v) { - v = cvt_bf16x4_e4m3((Array&)vi[i]); - } - else if constexpr (std::is_same_v) { - v = cvt_f16x4_e4m3((Array&)vi[i]); - } - else { - static_assert(N != N, "not implemented"); - } - } - return vo; - } + template + __device__ static auto convert(const Array& vi) { + static_assert(N % 4 == 0); + Array vo; + PRAGMA_UNROLL + for (int i = 0; i < N; i += 4) { + auto& v = (Array&)vo[i]; + if constexpr (std::is_same_v) { + v = cvt_bf16x4_e4m3((Array&)vi[i]); + } else if constexpr (std::is_same_v) { + v = cvt_f16x4_e4m3((Array&)vi[i]); + } else { + static_assert(N != N, "not implemented"); + } + } + return vo; + } }; -template -inline __device__ void StoreQuantParam(T* dst, Array src) -{ - Store(dst, src); +template +inline __device__ void StoreQuantParam(T* dst, Array src) { + Store(dst, src); } -template<> -inline __device__ void StoreQuantParam(half* dst, Array src) -{ - if constexpr (kFuseU4F16Dequant) { - src[1] = src[1] - src[0] * __ushort_as_half(0x5400); - } - Store(dst, src); +template <> +inline __device__ void StoreQuantParam(half* dst, + Array src) { + if constexpr (kFuseU4F16Dequant) { + src[1] = src[1] - src[0] * __ushort_as_half(0x5400); + } + Store(dst, src); } } // namespace turbomind diff --git a/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/arch/config_sm70_s884.h b/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/arch/config_sm70_s884.h index 5c7499df34..dc8e6c3011 100644 --- a/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/arch/config_sm70_s884.h +++ b/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/arch/config_sm70_s884.h @@ -19,117 +19,85 @@ namespace turbomind::gemm::sm70_s884 { -template +template struct Sm70_s884 { + static_assert(A::SmemCopyAtom::K == B::SmemCopyAtom::K); - static_assert(A::SmemCopyAtom::K == B::SmemCopyAtom::K); - - static constexpr int SMEM_M = A::SmemCopyAtom::M / A::SmemCopyAtom::kFragNum; - static constexpr int SMEM_N = B::SmemCopyAtom::M / B::SmemCopyAtom::kFragNum; - static constexpr int SMEM_K = A::SmemCopyAtom::K; - - static constexpr auto MODE_ = group_axis >= 0 ? Striding::kBlocked : Striding::kFlat; - - static constexpr auto MODE_A = group_axis == 0 ? Striding::kIndexed : MODE_; - static constexpr auto MODE_B = group_axis == 1 ? Striding::kIndexed : MODE_; - static constexpr auto MODE_C = MODE_; - - template - struct Type { - - // (TM, TN, TK) = R(MMA_Atom, SmemCopy_Atom) - using MMA_Atom = SM70_MMA_884; - - using Partition = Blocked; - using MMA_Map = MMA_Map; - - using MMA = Tiled_MMA_v2; - - using Mainloop = MainloopSm70, - TransformA, - U, - GroupSizeU, - B, - IteratorSm70, - TransformB, - V, - GroupSizeV, - Stages, - true, - GmemLookahead>; // FusePrefetch_ - - static constexpr int CHUNK_K = std::lcm(std::lcm(GroupSizeU, GroupSizeV), CTA_K); - - using Scheduler = SchedulerSm70; - - static constexpr int TILE_C_M = TILE_C_M_ == -1 ? CTA_M : TILE_C_M_; - static constexpr int TILE_C_N = TILE_C_N_ == -1 ? CTA_N : TILE_C_N_; - - using Epilogue = gemm::Epilogue_, - Operand_C, - MODE_C, - SplitK>; - - using Kernel = GemmUniversal; - }; + static constexpr int SMEM_M = A::SmemCopyAtom::M / A::SmemCopyAtom::kFragNum; + static constexpr int SMEM_N = B::SmemCopyAtom::M / B::SmemCopyAtom::kFragNum; + static constexpr int SMEM_K = A::SmemCopyAtom::K; + + static constexpr auto MODE_ = + group_axis >= 0 ? Striding::kBlocked : Striding::kFlat; + + static constexpr auto MODE_A = group_axis == 0 ? Striding::kIndexed : MODE_; + static constexpr auto MODE_B = group_axis == 1 ? Striding::kIndexed : MODE_; + static constexpr auto MODE_C = MODE_; + + template + struct Type { + // (TM, TN, TK) = R(MMA_Atom, SmemCopy_Atom) + using MMA_Atom = SM70_MMA_884; + + using Partition = Blocked; + using MMA_Map = + MMA_Map; + + using MMA = Tiled_MMA_v2; + + using IteratorA = + std::conditional_t, + IteratorSm70>; + using IteratorB = + std::conditional_t, + IteratorSm70>; + + using Mainloop = + MainloopSm70; // FusePrefetch_ + + static constexpr int CHUNK_K = + std::lcm(std::lcm(GroupSizeU, GroupSizeV), CTA_K); + + using Scheduler = SchedulerSm70; + + static constexpr int TILE_C_M = TILE_C_M_ == -1 ? CTA_M : TILE_C_M_; + static constexpr int TILE_C_N = TILE_C_N_ == -1 ? CTA_N : TILE_C_N_; + + using Epilogue = gemm::Epilogue_, + Operand_C, MODE_C, SplitK>; + + using Kernel = GemmUniversal; + }; }; -template -using Config_U4_d = Sm70_s884::Operand, - Transform_Default, - VoidOperand, - typename GetOperand::Operand, - Transform_HMMA_SIMT_B, - typename GetOperand::Operand, - kRowMajor, - half, - raster_order, - -1>; - -template -using Config_U4_d_A8x64Swizzle = Sm70_s884, - Transform_Default, - VoidOperand, - typename GetOperand::Operand, - Transform_HMMA_SIMT_B, - typename GetOperand::Operand, - kRowMajor, - half, - raster_order, - -1>; - -template +template +using Config_U4_d = Sm70_s884< + typename GetOperand::Operand, + Transform_Default, VoidOperand, + typename GetOperand::Operand, + Transform_HMMA_SIMT_B, + typename GetOperand::Operand, + kRowMajor, half, raster_order, -1>; + +template +using Config_U4_d_A8x64Swizzle = Sm70_s884< + Operand_A_Swizzle_8x64, Transform_Default, VoidOperand, + typename GetOperand::Operand, + Transform_HMMA_SIMT_B, + typename GetOperand::Operand, + kRowMajor, half, raster_order, -1>; + +template using Config_U4_g = Sm70_s884, // A Transform_Default, // tarnsform A VoidOperand, // U @@ -138,10 +106,9 @@ using Config_U4_g = Sm70_s884, // A Operand_V_Pack, // V kRowMajor, // order_C half, // Tc - raster_order, - 0>; + raster_order, 0>; -template +template using Config_MXF4 = Sm70_s884, // A Transform_Default, // tarnsform A VoidOperand, // U @@ -150,10 +117,9 @@ using Config_MXF4 = Sm70_s884, // A Operand_V_Pack, // V kRowMajor, // order_C half, // Tc - raster_order, - group_axis>; + raster_order, group_axis>; -template +template using Config_NVF4 = Sm70_s884, // A Transform_Default, // tarnsform A VoidOperand, // U @@ -162,10 +128,9 @@ using Config_NVF4 = Sm70_s884, // A Operand_V_Pack, // V kRowMajor, // order_C half, // Tc - raster_order, - group_axis>; + raster_order, group_axis>; -template +template using Config_E4M3 = Sm70_s884, // A Transform_Default, // tarnsform A VoidOperand, // U @@ -174,10 +139,21 @@ using Config_E4M3 = Sm70_s884, // A Operand_V_Pack, // V kRowMajor, // order_C half, // Tc - raster_order, - group_axis>; - -template + raster_order, group_axis>; + +template +using Config_E4M3_Prescaled = + Sm70_s884, // A + Transform_Default, // transform A + VoidOperand, // U + Operand_B_Pack, // B + Transform_HMMA_SIMT_B_PrescaledE4M3, + Operand_V_Pack, // pre-scaled V + kRowMajor, // order_C + half, // Tc + raster_order, group_axis>; + +template using Config_F16 = Sm70_s884, // A Transform_Default, // tarnsform A VoidOperand, // U @@ -186,7 +162,6 @@ using Config_F16 = Sm70_s884, // A VoidOperand, // V kRowMajor, // order_C half, // Tc - raster_order, - group_axis>; + raster_order, group_axis>; } // namespace turbomind::gemm::sm70_s884 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 82f5c511b4..4c776b38ab 100644 --- a/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/gemm.cu +++ b/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/gemm.cu @@ -25,653 +25,739 @@ namespace turbomind::gemm { -void ExportDispatchCache(std::ostream& os, const std::vector>& entries); +void ExportDispatchCache( + std::ostream& os, + const std::vector>& entries); -void ImportDispatchCache(std::istream& is, +void ImportDispatchCache(std::istream& is, std::vector>& entries, - const std::vector>& kernels); + const std::vector>& kernels); namespace { -template -std::vector ArgSort(size_t size, const Cmp& cmp) -{ - std::vector idxs(size); - std::iota(idxs.begin(), idxs.end(), 0); - std::stable_sort(idxs.begin(), idxs.end(), cmp); - return idxs; +template +std::vector ArgSort(size_t size, const Cmp& cmp) { + std::vector idxs(size); + std::iota(idxs.begin(), idxs.end(), 0); + std::stable_sort(idxs.begin(), idxs.end(), cmp); + return idxs; } -bool GemmTraceEnabled() -{ - const char* raw = std::getenv("TM_GEMM_TRACE"); - return raw && std::atoi(raw) != 0; +bool GemmTraceEnabled() { + const char* raw = std::getenv("TM_GEMM_TRACE"); + return raw && std::atoi(raw) != 0; } -int GemmTraceLimit() -{ - const char* raw = std::getenv("TM_GEMM_TRACE_LIMIT"); - return raw ? std::max(std::atoi(raw), 0) : 256; +int GemmTraceLimit() { + const char* raw = std::getenv("TM_GEMM_TRACE_LIMIT"); + return raw ? std::max(std::atoi(raw), 0) : 256; } -bool GemmTraceFilterAllows(const std::string& desc) -{ - const char* raw = std::getenv("TM_GEMM_TRACE_FILTER"); - return !raw || !*raw || desc.find(raw) != std::string::npos; +bool GemmTraceFilterAllows(const std::string& desc) { + const char* raw = std::getenv("TM_GEMM_TRACE_FILTER"); + return !raw || !*raw || desc.find(raw) != std::string::npos; } -bool Sm70AwqTp2FastSelectorEnabled() -{ - const char* raw = std::getenv("VLLM_SM70_AWQ_TP2_FAST_SELECTOR"); - return !raw || std::atoi(raw) != 0; +bool Sm70AwqTp2FastSelectorEnabled() { + const char* raw = std::getenv("VLLM_SM70_AWQ_TP2_FAST_SELECTOR"); + return !raw || std::atoi(raw) != 0; } // Keep the accepted QKVZ route enabled by default, while allowing an isolated // end-to-end A/B comparison without disabling the other exact small-M routes. -bool Sm70AwqTp4QkvCta64Enabled() -{ - const char* raw = std::getenv("VLLM_SM70_AWQ_TP4_QKV_CTA64"); - return !raw || std::atoi(raw) != 0; +bool Sm70AwqTp4QkvCta64Enabled() { + const char* raw = std::getenv("VLLM_SM70_AWQ_TP4_QKV_CTA64"); + return !raw || std::atoi(raw) != 0; } // Default-on A/B gate for the exact TP4 MTP verifier M=5 routes. -bool Sm70AwqMtpM5FastSelectorEnabled() -{ - const char* raw = std::getenv("VLLM_SM70_AWQ_MTP_M5_FAST_SELECTOR"); - return !raw || std::atoi(raw) != 0; +bool Sm70AwqMtpM5FastSelectorEnabled() { + const char* raw = std::getenv("VLLM_SM70_AWQ_MTP_M5_FAST_SELECTOR"); + return !raw || std::atoi(raw) != 0; } -bool Sm70Mxfp4MoeGroupedM8FastSelectorEnabled() -{ - const char* grouped = std::getenv("VLLM_SM70_MXFP4_MOE_GROUPED_M8"); - if (!grouped || std::atoi(grouped) == 0) { - return false; - } - const char* raw = std::getenv("VLLM_SM70_MXFP4_MOE_GROUPED_M8_FAST_SELECTOR"); - return !raw || std::atoi(raw) != 0; +bool Sm70Mxfp4MoeGroupedM8FastSelectorEnabled() { + const char* grouped = std::getenv("VLLM_SM70_MXFP4_MOE_GROUPED_M8"); + if (!grouped || std::atoi(grouped) == 0) { + return false; + } + const char* raw = std::getenv("VLLM_SM70_MXFP4_MOE_GROUPED_M8_FAST_SELECTOR"); + return !raw || std::atoi(raw) != 0; +} + +bool Sm70Nvfp4Qwen38Tp4M1FastSelectorEnabled() { + const char* raw = std::getenv("VLLM_SM70_NVFP4_QWEN38_TP4_M1_FAST_SELECTOR"); + return !raw || std::atoi(raw) != 0; +} + +// Default-on gate for exact block-FP8 8K prefill GEMM descriptors. +bool Sm70Fp8BlockPrefillFastSelectorEnabled() { + const char* raw = std::getenv("VLLM_SM70_FP8_PREFILL_FAST_SELECTOR"); + return !raw || std::atoi(raw) != 0; +} + +bool Sm70Fp8GroupedBmmDecodeEnabled() { + const char* raw = std::getenv("VLLM_SM70_FP8_GROUPED_BMM_DECODE"); + return !raw || std::atoi(raw) != 0; } struct Sm70AwqTp2FastTarget { - int n; - int k; - int cta_m; - int cta_n; - int cta_k; - int splits; - int swizzle; - bool require_mgroup; - std::string name_contains; + int n; + int k; + int cta_m; + int cta_n; + int cta_k; + int splits; + int swizzle; + bool require_mgroup; + std::string name_contains; }; -std::optional GetSm70AwqTp2EnvFastTarget(const GemmDesc& desc, - const std::string_view desc_str) -{ - const char* raw = std::getenv("VLLM_SM70_AWQ_TP2_FAST_TARGETS"); - if (!raw || !*raw) { - return std::nullopt; - } - const std::string targets(raw); - size_t begin = 0; - while (begin <= targets.size()) { - const size_t end = targets.find(';', begin); - const auto entry = targets.substr(begin, end == std::string::npos ? std::string::npos : end - begin); - const size_t sep = entry.find('|'); - if (sep != std::string::npos && entry.substr(0, sep) == desc_str) { - std::string spec = entry.substr(sep + 1); - std::string name_contains; - if (const size_t name_sep = spec.find('@'); name_sep != std::string::npos) { - name_contains = spec.substr(name_sep + 1); - spec.resize(name_sep); - } - std::replace(spec.begin(), spec.end(), 'x', ','); - std::replace(spec.begin(), spec.end(), ':', ','); - for (char& ch : spec) { - if (ch == ',') { - ch = ' '; - } - } - std::istringstream is(spec); - int cta_m{}; - int cta_n{}; - int cta_k{}; - int splits{}; - int swizzle{}; - int require_mgroup{}; - if (is >> cta_m >> cta_n >> cta_k >> splits >> swizzle >> require_mgroup) { - if (name_contains.empty()) { - is >> name_contains; - } - return Sm70AwqTp2FastTarget{desc.n, - desc.k, - cta_m, - cta_n, - cta_k, - splits, - swizzle, - require_mgroup != 0, - name_contains}; - } - if (GemmTraceEnabled() && GemmTraceFilterAllows(std::string(desc_str))) { - std::cerr << "[TM_GEMM_FAST_SELECTOR] desc=" << desc_str - << " stage=invalid_env_target entry=" << entry << std::endl; - } - return std::nullopt; - } - if (end == std::string::npos) { - break; - } - begin = end + 1; - } +std::optional GetSm70AwqTp2EnvFastTarget( + const GemmDesc& desc, const std::string_view desc_str) { + const char* raw = std::getenv("VLLM_SM70_AWQ_TP2_FAST_TARGETS"); + if (!raw || !*raw) { return std::nullopt; + } + const std::string targets(raw); + size_t begin = 0; + while (begin <= targets.size()) { + const size_t end = targets.find(';', begin); + const auto entry = targets.substr( + begin, end == std::string::npos ? std::string::npos : end - begin); + const size_t sep = entry.find('|'); + if (sep != std::string::npos && entry.substr(0, sep) == desc_str) { + std::string spec = entry.substr(sep + 1); + std::string name_contains; + if (const size_t name_sep = spec.find('@'); + name_sep != std::string::npos) { + name_contains = spec.substr(name_sep + 1); + spec.resize(name_sep); + } + std::replace(spec.begin(), spec.end(), 'x', ','); + std::replace(spec.begin(), spec.end(), ':', ','); + for (char& ch : spec) { + if (ch == ',') { + ch = ' '; + } + } + std::istringstream is(spec); + int cta_m{}; + int cta_n{}; + int cta_k{}; + int splits{}; + int swizzle{}; + int require_mgroup{}; + if (is >> cta_m >> cta_n >> cta_k >> splits >> swizzle >> + require_mgroup) { + if (name_contains.empty()) { + is >> name_contains; + } + return Sm70AwqTp2FastTarget{ + desc.n, desc.k, cta_m, cta_n, + cta_k, splits, swizzle, require_mgroup != 0, + name_contains}; + } + if (GemmTraceEnabled() && GemmTraceFilterAllows(std::string(desc_str))) { + std::cerr << "[TM_GEMM_FAST_SELECTOR] desc=" << desc_str + << " stage=invalid_env_target entry=" << entry << std::endl; + } + return std::nullopt; + } + if (end == std::string::npos) { + break; + } + begin = end + 1; + } + return std::nullopt; } -std::optional GetSm70AwqTp2FastTarget(const GemmDesc& desc) -{ - const bool awq_fast_selector_enabled = Sm70AwqTp2FastSelectorEnabled(); - const std::string desc_str = to_string(desc); - if (awq_fast_selector_enabled) { - if (auto target = GetSm70AwqTp2EnvFastTarget(desc, desc_str)) { - return target; - } - } - if (!awq_fast_selector_enabled) { - return std::nullopt; - } - if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_5x17408x5120_1") { - return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 64, 64, 1, 0, false, ""}; - } - if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_5x8192x5120_1") { - return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 64, 64, 2, 0, false, ""}; - } - if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_5x5120x3072_1") { - return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 64, 64, 3, 0, false, ""}; - } - if (Sm70AwqMtpM5FastSelectorEnabled()) { - // TP4 MTP verifier, M=5. Both routes are bitwise exact on all TP shards. - if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_5x8704x5120_1") { - return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 64, 64, 2, 4, false, ""}; - } - if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_5x4096x5120_1") { - return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 64, 64, 4, 4, false, ""}; - } - } - if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_1x17408x5120_1") { - return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 256, 64, 3, 3, false, ""}; - } - if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_1x8704x5120_1") { - return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 64, 64, 2, 4, false, ""}; - } - if (Sm70AwqTp4QkvCta64Enabled() - && desc_str == "sm70_f16_u4k128_f16_tnt_fff_1x4096x5120_1") { - return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 64, 64, 4, 4, false, ""}; - } - if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_1x5120x1536_1") { - return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 256, 64, 12, 4, false, "s884_1x4x1"}; - } - if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_1x5120x3072_1") { - return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 256, 64, 7, 0, false, ""}; - } +std::optional GetSm70AwqTp2FastTarget( + const GemmDesc& desc) { + const std::string desc_str = to_string(desc); + if (Sm70Nvfp4Qwen38Tp4M1FastSelectorEnabled()) { + // Qwen3.8-27B NVFP4 TP4 decode. The exact-MNK registry entries keep + // these small-N tactics out of prefill and unrelated model shapes. + if (desc_str == "sm70_f16_e2m1k16_f16_tnt_fff_1x8704x5120_1" || + desc_str == "sm70_f16_e2m1k16_f16_tnt_fff_1x5120x4352_1") { + return Sm70AwqTp2FastTarget{ + desc.n, desc.k, 8, 32, 64, 3, 4, true, "c8x32_a1x1x64_01"}; + } + } + if (Sm70Fp8BlockPrefillFastSelectorEnabled() && + (desc_str == "sm70_f16_e4m3k128_f16_tnt_fff_8000x4096x5120_1" || + desc_str == "sm70_f16_e4m3k128_f16_tnt_fff_8000x3584x5120_1")) { + return Sm70AwqTp2FastTarget{desc.n, desc.k, 64, 256, 16, 1, 3, true, ""}; + } + if (Sm70Fp8GroupedBmmDecodeEnabled() && + desc_str == "sm70_f16_e4m3k128_f16_tnt_bbb_2x1024x4096_1") { + // Match the accepted dense 1x1024x4096 accumulation tree while launching + // both independent WO-A groups through one blocked descriptor. + return Sm70AwqTp2FastTarget{ + desc.n, desc.k, 8, 128, 64, 7, 2, true, "c8x128_a1x1x64_01"}; + } + const bool awq_fast_selector_enabled = Sm70AwqTp2FastSelectorEnabled(); + if (awq_fast_selector_enabled) { + if (auto target = GetSm70AwqTp2EnvFastTarget(desc, desc_str)) { + return target; + } + } + const char* selector_rerank = std::getenv("VLLM_SM70_DFLASH2_QPN8_RERANK"); + const char* selector_shadow = + std::getenv("VLLM_SM70_DFLASH2_QPN8_RERANK_SHADOW"); + const bool exact_selector_rerank = + (selector_rerank && std::atoi(selector_rerank) != 0) || + (selector_shadow && std::atoi(selector_shadow) != 0); + if (exact_selector_rerank && desc.arch == 700 && desc.type_a == kHalf && + desc.type_b == kHalf && desc.type_c == kHalf && desc.m >= 1 && + desc.m <= 8 && desc.n == 62080 && desc.k == 5120 && desc.num == 1) { + return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 256, 64, + 10, 1, true, "s884_1x4x1"}; + } + if (!awq_fast_selector_enabled) { return std::nullopt; + } + if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_5x17408x5120_1") { + return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 64, 64, 1, 0, false, ""}; + } + if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_5x8192x5120_1") { + return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 64, 64, 2, 0, false, ""}; + } + if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_5x5120x3072_1") { + return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 64, 64, 3, 0, false, ""}; + } + if (Sm70AwqMtpM5FastSelectorEnabled()) { + // TP4 MTP verifier, M=5. Both routes are bitwise exact on all TP shards. + if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_5x8704x5120_1") { + return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 64, 64, 2, 4, false, ""}; + } + if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_5x4096x5120_1") { + return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 64, 64, 4, 4, false, ""}; + } + } + if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_1x17408x5120_1") { + return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 256, 64, 3, 3, false, ""}; + } + if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_1x8704x5120_1") { + return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 64, 64, 2, 4, false, ""}; + } + if (Sm70AwqTp4QkvCta64Enabled() && + desc_str == "sm70_f16_u4k128_f16_tnt_fff_1x4096x5120_1") { + return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 64, 64, 4, 4, false, ""}; + } + if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_1x5120x1536_1") { + return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 256, 64, + 12, 4, false, "s884_1x4x1"}; + } + if (desc_str == "sm70_f16_u4k128_f16_tnt_fff_1x5120x3072_1") { + return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 256, 64, 7, 0, false, ""}; + } + return std::nullopt; } -std::optional GetSm70Mxfp4MoeGroupedM8FastTarget(const GemmDesc& desc) -{ - if (!Sm70Mxfp4MoeGroupedM8FastSelectorEnabled()) { - return std::nullopt; - } - const std::string desc_str = to_string(desc); - if (desc_str == "sm70_f16_e2m1k32_f16_tnt_bbb_48x512x4096_1") { - return Sm70AwqTp2FastTarget{ - desc.n, desc.k, 8, 128, 64, 5, 0, true, "c8x128_a1x1x64_01"}; - } - if (desc_str == "sm70_f16_e2m1k32_f16_tnt_bbb_48x4096x256_1") { - return Sm70AwqTp2FastTarget{ - desc.n, desc.k, 16, 128, 32, 1, 0, true, "c16x128_a1x1x32_01"}; - } +std::optional GetSm70Mxfp4MoeGroupedM8FastTarget( + const GemmDesc& desc) { + if (!Sm70Mxfp4MoeGroupedM8FastSelectorEnabled()) { return std::nullopt; + } + const std::string desc_str = to_string(desc); + if (desc_str == "sm70_f16_e2m1k32_f16_tnt_bbb_48x512x4096_1") { + return Sm70AwqTp2FastTarget{ + desc.n, desc.k, 8, 128, 64, 5, 0, true, "c8x128_a1x1x64_01"}; + } + if (desc_str == "sm70_f16_e2m1k32_f16_tnt_bbb_48x4096x256_1") { + return Sm70AwqTp2FastTarget{ + desc.n, desc.k, 16, 128, 32, 1, 0, true, "c16x128_a1x1x32_01"}; + } + return std::nullopt; } -bool MatchesSm70AwqTp2FastKernel(const Kernel& kernel, const Sm70AwqTp2FastTarget& target) -{ - const int3 cta = kernel.cta_tile_size(); - if (cta.x != target.cta_m || cta.y != target.cta_n || cta.z != target.cta_k) { - return false; - } - const std::string name = kernel.name(); - const bool is_mgroup = name.find("mgroup") != std::string::npos; - if (is_mgroup != target.require_mgroup) { - return false; - } - return target.name_contains.empty() || name.find(target.name_contains) != std::string::npos; +std::optional GetSm70Fp8BlockPrefillPrescaledTarget( + const GemmDesc& desc) { + const std::string desc_str = to_string(desc); + if (desc_str == "sm70_f16_e4m3k128_f16_tnt_fff_8000x4096x5120_1" || + desc_str == "sm70_f16_e4m3k128_f16_tnt_fff_8000x3584x5120_1") { + 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; } -void MaybeTraceSm70AwqTp2FastSelector(const GemmDesc& desc, const char* stage, const LaunchSpec* spec = nullptr) -{ - if (!GemmTraceEnabled()) { - return; - } - const std::string desc_str = to_string(desc); - if (!GemmTraceFilterAllows(desc_str)) { - return; - } - std::cerr << "[TM_GEMM_FAST_SELECTOR] desc=" << desc_str << " stage=" << stage; - if (spec && spec->kernel) { - std::cerr << " kernel=" << spec->kernel->name() << " splits=" << spec->splits - << " swizzle=" << spec->swizzle; - } - std::cerr << std::endl; +bool MatchesSm70AwqTp2FastKernel(const Kernel& kernel, + const Sm70AwqTp2FastTarget& target) { + const int3 cta = kernel.cta_tile_size(); + if (cta.x != target.cta_m || cta.y != target.cta_n || cta.z != target.cta_k) { + return false; + } + const std::string name = kernel.name(); + const bool is_mgroup = name.find("mgroup") != std::string::npos; + if (is_mgroup != target.require_mgroup) { + return false; + } + return target.name_contains.empty() || + name.find(target.name_contains) != std::string::npos; } -std::optional SelectSm70AwqTp2FastSpec(Context& ctx, - const std::vector& specs, - const Sm70AwqTp2FastTarget& target, - size_t barriers_size, - size_t partials_size) -{ - for (const auto& spec : specs) { - if (!spec.kernel || !MatchesSm70AwqTp2FastKernel(*spec.kernel, target)) { - continue; - } - if (spec.splits != target.splits) { - continue; - } - auto selected = spec; - const auto& actual_desc = ctx.get_desc(*selected.kernel); - const int4 shape{actual_desc.m, actual_desc.n, actual_desc.k, actual_desc.num}; - if (target.swizzle > selected.kernel->GetMaxSwizzle(shape)) { - continue; - } - (void)barriers_size; - (void)partials_size; - selected.splits = target.splits; - selected.swizzle = target.swizzle; - MaybeTraceSm70AwqTp2FastSelector(ctx.desc(), "selected", &selected); - return selected; - } - MaybeTraceSm70AwqTp2FastSelector(ctx.desc(), "no_match"); - return std::nullopt; +void MaybeTraceSm70AwqTp2FastSelector(const GemmDesc& desc, const char* stage, + const LaunchSpec* spec = nullptr) { + if (!GemmTraceEnabled()) { + return; + } + const std::string desc_str = to_string(desc); + if (!GemmTraceFilterAllows(desc_str)) { + return; + } + std::cerr << "[TM_GEMM_FAST_SELECTOR] desc=" << desc_str + << " stage=" << stage; + if (spec && spec->kernel) { + std::cerr << " kernel=" << spec->kernel->name() + << " splits=" << spec->splits << " swizzle=" << spec->swizzle; + } + std::cerr << std::endl; } -const char* ToString(DispatchPolicy policy) -{ - if ((policy & DispatchPolicy::kPreserveDefaultSplits) - || (policy & DispatchPolicy::kPreserveDefaultSplitCount) - || (policy & DispatchPolicy::kMxfp4MoeGroupedM8Fast)) { - static thread_local std::string text; - auto base = - static_cast((int)policy & ~(int)DispatchPolicy::kPreserveDefaultSplits - & ~(int)DispatchPolicy::kPreserveDefaultSplitCount - & ~(int)DispatchPolicy::kMxfp4MoeGroupedM8Fast); - text = std::string(ToString(base)); - if (policy & DispatchPolicy::kPreserveDefaultSplits) { - text += "|preserve_default_splits"; - } - if (policy & DispatchPolicy::kPreserveDefaultSplitCount) { - text += "|preserve_default_split_count"; - } - if (policy & DispatchPolicy::kMxfp4MoeGroupedM8Fast) { - text += "|mxfp4_moe_grouped_m8_fast"; - } - return text.c_str(); - } - switch (policy) { - case DispatchPolicy::kDefault: - return "default"; - case DispatchPolicy::kMeasure: - return "measure"; - case DispatchPolicy::kReuse: - return "reuse"; - case DispatchPolicy::kAppend: - return "append"; - default: - return "unknown"; - } +std::optional SelectSm70AwqTp2FastSpec( + Context& ctx, const std::vector& specs, + const Sm70AwqTp2FastTarget& target, size_t barriers_size, + size_t partials_size) { + for (const auto& spec : specs) { + if (!spec.kernel || !MatchesSm70AwqTp2FastKernel(*spec.kernel, target)) { + continue; + } + if (spec.splits != target.splits) { + continue; + } + auto selected = spec; + const auto& actual_desc = ctx.get_desc(*selected.kernel); + const int4 shape{actual_desc.m, actual_desc.n, actual_desc.k, + actual_desc.num}; + if (target.swizzle > selected.kernel->GetMaxSwizzle(shape)) { + continue; + } + (void)barriers_size; + (void)partials_size; + selected.splits = target.splits; + selected.swizzle = target.swizzle; + MaybeTraceSm70AwqTp2FastSelector(ctx.desc(), "selected", &selected); + return selected; + } + MaybeTraceSm70AwqTp2FastSelector(ctx.desc(), "no_match"); + return std::nullopt; } -void MaybeTraceGemmDispatch(const GemmDesc& desc, DispatchPolicy policy, const LaunchSpec& spec, bool measured) -{ - if (!GemmTraceEnabled()) { - return; - } - const std::string desc_str = to_string(desc); - if (!GemmTraceFilterAllows(desc_str)) { - return; - } - static std::atomic logged{0}; - const int limit = GemmTraceLimit(); - if (limit > 0 && logged.fetch_add(1, std::memory_order_relaxed) >= limit) { - return; - } +const char* ToString(DispatchPolicy policy) { + if ((policy & DispatchPolicy::kPreserveDefaultSplits) || + (policy & DispatchPolicy::kPreserveDefaultSplitCount) || + (policy & DispatchPolicy::kMxfp4MoeGroupedM8Fast) || + (policy & DispatchPolicy::kSm70Fp8PrefillPrescaled)) { + static thread_local std::string text; + auto base = static_cast( + (int)policy & ~(int)DispatchPolicy::kPreserveDefaultSplits & + ~(int)DispatchPolicy::kPreserveDefaultSplitCount & + ~(int)DispatchPolicy::kMxfp4MoeGroupedM8Fast & + ~(int)DispatchPolicy::kSm70Fp8PrefillPrescaled); + text = std::string(ToString(base)); + if (policy & DispatchPolicy::kPreserveDefaultSplits) { + text += "|preserve_default_splits"; + } + if (policy & DispatchPolicy::kPreserveDefaultSplitCount) { + text += "|preserve_default_split_count"; + } + if (policy & DispatchPolicy::kMxfp4MoeGroupedM8Fast) { + text += "|mxfp4_moe_grouped_m8_fast"; + } + if (policy & DispatchPolicy::kSm70Fp8PrefillPrescaled) { + text += "|sm70_fp8_prefill_prescaled"; + } + return text.c_str(); + } + switch (policy) { + case DispatchPolicy::kDefault: + return "default"; + case DispatchPolicy::kMeasure: + return "measure"; + case DispatchPolicy::kReuse: + return "reuse"; + case DispatchPolicy::kAppend: + return "append"; + default: + return "unknown"; + } +} - std::cerr << "[TM_GEMM_TRACE] desc=" << desc_str << " policy=" << ToString(policy) << " measured=" - << (measured ? 1 : 0); - if (spec.kernel) { - const int3 cta = spec.kernel->cta_tile_size(); - const int3 mma = spec.kernel->warp_tile_size(); - std::cerr << " kernel=" << spec.kernel->name() << " splits=" << spec.splits << " swizzle=" << spec.swizzle - << " cta=" << cta.x << "x" << cta.y << "x" << cta.z << " mma=" << mma.x << "x" << mma.y - << "x" << mma.z << " stages=" << spec.kernel->stages() << " smem=" << spec.kernel->smem_size(); - } - else { - std::cerr << " kernel="; - } - std::cerr << std::endl; +void MaybeTraceGemmDispatch(const GemmDesc& desc, DispatchPolicy policy, + const LaunchSpec& spec, bool measured) { + if (!GemmTraceEnabled()) { + return; + } + const std::string desc_str = to_string(desc); + if (!GemmTraceFilterAllows(desc_str)) { + return; + } + static std::atomic logged{0}; + const int limit = GemmTraceLimit(); + if (limit > 0 && logged.fetch_add(1, std::memory_order_relaxed) >= limit) { + return; + } + + std::cerr << "[TM_GEMM_TRACE] desc=" << desc_str + << " policy=" << ToString(policy) + << " measured=" << (measured ? 1 : 0); + if (spec.kernel) { + const int3 cta = spec.kernel->cta_tile_size(); + const int3 mma = spec.kernel->warp_tile_size(); + std::cerr << " kernel=" << spec.kernel->name() << " splits=" << spec.splits + << " swizzle=" << spec.swizzle << " cta=" << cta.x << "x" << cta.y + << "x" << cta.z << " mma=" << mma.x << "x" << mma.y << "x" + << mma.z << " stages=" << spec.kernel->stages() + << " smem=" << spec.kernel->smem_size() + << " regs=" << spec.kernel->info().attr.numRegs + << " max_ctas=" << spec.kernel->info().max_active_ctas; + } else { + std::cerr << " kernel="; + } + std::cerr << std::endl; } } // namespace struct Gemm::Impl { - - Impl(): - props_{GetCudaDeviceProps()}, + Impl() + : props_{GetCudaDeviceProps()}, arch_{props_->major * 100 + props_->minor * 10}, registry_{props_}, - cache_{registry_.kernels()} - { - if (arch_ == 700) { - // V100 decode is dominated by many tiny GEMM/GEMV problems. A - // broader search space consistently finds better launch specs than - // the generic defaults for these SM70 workloads. - tuning_.max_splits = 16; - tuning_.max_waves = 32; - tuning_.swizzle = {0, 1, 2, 3, 4}; - tuning_.top_k = 0; - tuning_.clusters = 0; - tuning_.min_iter = 2; - tuning_.max_iter = 20; - tuning_.max_time = 2.f; - } - if (auto str = std::getenv("TM_GEMM_TUNE")) { - try { - ParseTuningParams(tuning_, str); - } - catch (...) { - std::cerr << "[Gemm2] Failed to parse `TM_GEMM_TUNE`, default value will be used.\n"; - tuning_ = {}; - } - } - if (std::getenv("TM_GEMM_WARN_CACHE_MISS")) { - warn_cache_miss_ = true; - } - measurer_.emplace(CreateStoppingCriterion(tuning_.min_iter, tuning_.max_iter, tuning_.max_time)); - } - - // find launch spec in dispatch cache, dispatch by heuristic on cache miss - LaunchSpec Dispatch(Context& ctx, DispatchPolicy policy, size_t barriers_size, size_t partials_size) - { - const auto& desc = ctx.desc(); - const auto is_feasible = [&](const LaunchSpec& spec) { - return spec.kernel && spec.kernel->is_feasible(ctx.get_desc(*spec.kernel)); - }; - if (policy & DispatchPolicy::kMxfp4MoeGroupedM8Fast) { - if (auto fast_target = GetSm70Mxfp4MoeGroupedM8FastTarget(desc)) { - auto specs = Find(ctx, barriers_size, partials_size, 0); - if (auto fast_spec = SelectSm70AwqTp2FastSpec( - ctx, specs, *fast_target, barriers_size, partials_size)) { - return *fast_spec; - } - return {}; - } - } - if (policy & DispatchPolicy::kReuse) { - if (auto spec = cache_.LowerBound(desc); spec && is_feasible(*spec)) { - return *spec; - } - if (warn_cache_miss_) { - std::cerr << "Failed to find a feasible kernel in the cache, will dispatch by heuristic: " - << to_string(ctx.desc()) << std::endl; - } - } - - if (auto spec = cache_.Find(desc); spec && is_feasible(*spec)) { - return *spec; - } - - const auto fast_target = GetSm70AwqTp2FastTarget(desc); - auto specs = Find(ctx, barriers_size, partials_size, fast_target ? 0 : 1); - if (!specs.empty()) { - auto selected = specs.front(); - if (fast_target) { - if (auto fast_spec = - SelectSm70AwqTp2FastSpec(ctx, specs, *fast_target, barriers_size, partials_size)) { - selected = *fast_spec; - } - } - cache_.Insert(desc, selected); - return selected; + cache_{registry_.kernels()}, + sm70_fp8_prefill_cache_{registry_.kernels()} { + if (arch_ == 700) { + // V100 decode is dominated by many tiny GEMM/GEMV problems. A + // broader search space consistently finds better launch specs than + // the generic defaults for these SM70 workloads. + tuning_.max_splits = 16; + tuning_.max_waves = 32; + tuning_.swizzle = {0, 1, 2, 3, 4}; + tuning_.top_k = 0; + tuning_.clusters = 0; + tuning_.min_iter = 2; + tuning_.max_iter = 20; + tuning_.max_time = 2.f; + } + if (auto str = std::getenv("TM_GEMM_TUNE")) { + try { + ParseTuningParams(tuning_, str); + } catch (...) { + std::cerr << "[Gemm2] Failed to parse `TM_GEMM_TUNE`, default value " + "will be used.\n"; + tuning_ = {}; + } + } + if (std::getenv("TM_GEMM_WARN_CACHE_MISS")) { + warn_cache_miss_ = true; + } + measurer_.emplace(CreateStoppingCriterion( + tuning_.min_iter, tuning_.max_iter, tuning_.max_time)); + } + + // find launch spec in dispatch cache, dispatch by heuristic on cache miss + 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) { + if (auto spec = sm70_fp8_prefill_cache_.Find(desc); + spec && is_feasible(*spec)) { + return *spec; + } + if (auto fast_target = GetSm70Fp8BlockPrefillPrescaledTarget(desc)) { + 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); + return *fast_spec; + } + } + return {}; + } + if (policy & DispatchPolicy::kMxfp4MoeGroupedM8Fast) { + if (auto fast_target = GetSm70Mxfp4MoeGroupedM8FastTarget(desc)) { + 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; } return {}; - } - - std::vector Find(Context& ctx, size_t barrier_size, size_t partials_size, int top_k) + } + } + if (policy & DispatchPolicy::kReuse) { + if (auto spec = cache_.LowerBound(desc); spec && is_feasible(*spec)) { + return *spec; + } + if (warn_cache_miss_) { + std::cerr << "Failed to find a feasible kernel in the cache, will " + "dispatch by heuristic: " + << to_string(ctx.desc()) << std::endl; + } + } + + if (auto spec = cache_.Find(desc); spec && is_feasible(*spec)) { + return *spec; + } + + const auto fast_target = GetSm70AwqTp2FastTarget(desc); + auto specs = + Find(ctx, barriers_size, partials_size, fast_target ? 0 : 1, false); + if (!specs.empty()) { + auto selected = specs.front(); + if (fast_target) { + if (auto fast_spec = SelectSm70AwqTp2FastSpec( + ctx, specs, *fast_target, barriers_size, partials_size)) { + selected = *fast_spec; + } + } + cache_.Insert(desc, selected); + return selected; + } + return {}; + } + + std::vector Find(Context& ctx, size_t barrier_size, + 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; { - std::vector feasible = ctx.Filter(registry_.kernels()); - - std::vector> clusters; - { - std::vector tmp; - tmp.reserve(feasible.size()); - for (const auto& k : feasible) { - LaunchSpec spec{k}; - tmp.push_back(spec); - } - clusters = Cluster(tmp, ClusteringParam{false, true}); - } - std::vector proxies; - proxies.reserve(clusters.size()); - - for (const auto& c : clusters) { - proxies.push_back(c.front().kernel); - } - - std::vector> specs; - - PopulateParam param{}; - param.max_splits = tuning_.max_splits; - param.max_waves = tuning_.max_waves; - param.swizzle = tuning_.swizzle.at(0); - param.barriers_size = barrier_size; - param.partials_size = partials_size; - - for (int cluster_id = 0; cluster_id < (int)proxies.size(); ++cluster_id) { - auto& kernel = *proxies[cluster_id]; - - auto tmp = ctx.Populate(kernel, param); - for (const auto& s : tmp) { - specs.emplace_back(cluster_id, s); - } - } - - // std::cerr << "#kernel: " << kernels.size() << ", #cluster: " << clusters.size() - // << ", #metric: " << metrics.size() << "\n"; - - int64_t mio_max = 0; - int64_t mma_max = 0; - for (const auto& [_, s] : specs) { - auto& [mio, mma] = s.estimated; - mio_max = std::max(mio_max, mio); - mma_max = std::max(mma_max, mma); - } - std::vector mio_ratio; - std::vector mma_ratio; - std::vector avg_ratio; - for (const auto& [_, s] : specs) { - auto& [mio, mma] = s.estimated; - mio_ratio.push_back((float)mio / mio_max); - mma_ratio.push_back((float)mma / mma_max); - avg_ratio.push_back(.5 * (mio_ratio.back() + mma_ratio.back())); - } - auto idxs = ArgSort(specs.size(), [&](int i, int j) { // - return avg_ratio[i] < avg_ratio[j]; - }); - - // for (const auto& i : idxs) { - // auto [cid, s, m] = metrics[i]; - // std::cout << clusters[cid].front().kernel->name() << " s" << s << " " << avg_ratio[i] << " " << - // mio_ratio[i] - // << " " << mma_ratio[i] << " " << m.mio_cost << " " << m.mma_cost << "\n"; - // } - - top_k = top_k > 0 ? std::min(idxs.size(), top_k) : (int)idxs.size(); - std::vector ret; - ret.reserve(top_k); - for (int i = 0; i < top_k; ++i) { - const auto& [cluster_id, spec] = specs[idxs[i]]; - // Apply `splits` to all kernels in the cluster - for (const auto& s : clusters[cluster_id]) { - auto tmp = spec; - tmp.kernel = s.kernel; - ret.push_back(tmp); - } - } - - return ret; + std::vector tmp; + tmp.reserve(feasible.size()); + for (const auto& k : feasible) { + LaunchSpec spec{k}; + tmp.push_back(spec); + } + clusters = Cluster(tmp, ClusteringParam{false, true}); + } + std::vector proxies; + proxies.reserve(clusters.size()); + + for (const auto& c : clusters) { + proxies.push_back(c.front().kernel); + } + + std::vector> specs; + + PopulateParam param{}; + param.max_splits = tuning_.max_splits; + param.max_waves = tuning_.max_waves; + param.swizzle = tuning_.swizzle.at(0); + param.barriers_size = barrier_size; + param.partials_size = partials_size; + + for (int cluster_id = 0; cluster_id < (int)proxies.size(); ++cluster_id) { + auto& kernel = *proxies[cluster_id]; + + auto tmp = ctx.Populate(kernel, param); + for (const auto& s : tmp) { + specs.emplace_back(cluster_id, s); + } + } + + // std::cerr << "#kernel: " << kernels.size() << ", #cluster: " << + // clusters.size() + // << ", #metric: " << metrics.size() << "\n"; + + int64_t mio_max = 0; + int64_t mma_max = 0; + for (const auto& [_, s] : specs) { + auto& [mio, mma] = s.estimated; + mio_max = std::max(mio_max, mio); + mma_max = std::max(mma_max, mma); + } + std::vector mio_ratio; + std::vector mma_ratio; + std::vector avg_ratio; + for (const auto& [_, s] : specs) { + auto& [mio, mma] = s.estimated; + mio_ratio.push_back((float)mio / mio_max); + mma_ratio.push_back((float)mma / mma_max); + avg_ratio.push_back(.5 * (mio_ratio.back() + mma_ratio.back())); } + auto idxs = ArgSort(specs.size(), [&](int i, int j) { // + return avg_ratio[i] < avg_ratio[j]; + }); - template - int Measure( - Context& ctx, size_t barriers_size, size_t partials_size, int top_k, LaunchFunc launch_func, cudaStream_t st) - { - // Early exit on exact match - if (cache_.Find(ctx.desc())) { - return 0; - } - // std::cerr << "GEMM: " << desc.m << "x" << desc.n << "x" << desc.k << "\n"; + // for (const auto& i : idxs) { + // auto [cid, s, m] = metrics[i]; + // std::cout << clusters[cid].front().kernel->name() << " s" << s << " " + // << avg_ratio[i] << " " << mio_ratio[i] + // << " " << mma_ratio[i] << " " << m.mio_cost << " " << + // m.mma_cost << "\n"; + // } - const auto tmp = Find(ctx, barriers_size, partials_size, tuning_.top_k); + top_k = top_k > 0 ? std::min(idxs.size(), top_k) : (int)idxs.size(); + std::vector ret; + ret.reserve(top_k); + for (int i = 0; i < top_k; ++i) { + const auto& [cluster_id, spec] = specs[idxs[i]]; + // Apply `splits` to all kernels in the cluster + for (const auto& s : clusters[cluster_id]) { + auto tmp = spec; + tmp.kernel = s.kernel; + ret.push_back(tmp); + } + } - std::vector specs; - for (const auto& spec : tmp) { - // populate swizzle parameters - const auto swis = ctx.Swizzle(spec, tuning_.swizzle); - specs.insert(specs.end(), swis.begin(), swis.end()); - } + return ret; + } - specs = Sampler{*measurer_, tuning_.clusters}.Run(specs, launch_func, st); + template + int Measure(Context& ctx, size_t barriers_size, size_t partials_size, + int top_k, LaunchFunc launch_func, cudaStream_t st) { + // Early exit on exact match + if (cache_.Find(ctx.desc())) { + return 0; + } + // std::cerr << "GEMM: " << desc.m << "x" << desc.n << "x" << desc.k << + // "\n"; - // for (const auto& s : specs) { - // std::cout << s.kernel->name() // - // << " swizzle=" << s.swizzle // - // << ", splits=" << s.splits // - // << ", measured=" << s.measured << "ms\n"; - // break; - // } + const auto tmp = + Find(ctx, barriers_size, partials_size, tuning_.top_k, false); - if (!specs.empty()) { - cache_.Insert(ctx.desc(), specs.front()); - } - else { - std::cerr << "No valid kernel found for the problem\n"; - return -1; - } - - return 0; + std::vector specs; + for (const auto& spec : tmp) { + // populate swizzle parameters + const auto swis = ctx.Swizzle(spec, tuning_.swizzle); + specs.insert(specs.end(), swis.begin(), swis.end()); } - /// TODO: move to cuda utils - static std::unique_ptr GetCudaDeviceProps() - { - auto props = std::make_unique(); - int device_id = -1; - cudaGetDevice(&device_id); - cudaGetDeviceProperties(props.get(), device_id); - return props; - } + specs = Sampler{*measurer_, tuning_.clusters}.Run(specs, launch_func, st); - std::shared_ptr props_; + // for (const auto& s : specs) { + // std::cout << s.kernel->name() // + // << " swizzle=" << s.swizzle // + // << ", splits=" << s.splits // + // << ", measured=" << s.measured << "ms\n"; + // break; + // } - int arch_; + if (!specs.empty()) { + cache_.Insert(ctx.desc(), specs.front()); + } else { + std::cerr << "No valid kernel found for the problem\n"; + return -1; + } + + return 0; + } - Registry registry_; + /// TODO: move to cuda utils + static std::unique_ptr GetCudaDeviceProps() { + auto props = std::make_unique(); + int device_id = -1; + cudaGetDevice(&device_id); + cudaGetDeviceProperties(props.get(), device_id); + return props; + } - TuningParams tuning_; + std::shared_ptr props_; - bool warn_cache_miss_{}; + int arch_; + + Registry registry_; + + TuningParams tuning_; + + bool warn_cache_miss_{}; + + std::optional measurer_; - std::optional measurer_; + DispatchCache cache_; - DispatchCache cache_; + DispatchCache sm70_fp8_prefill_cache_; - std::mutex dispatch_mutex_; + std::mutex dispatch_mutex_; }; // implementation of GEMM interfaces -Gemm::Gemm(): impl_{new Impl{}} {} +Gemm::Gemm() : impl_{new Impl{}} {} Gemm::~Gemm() = default; -int Gemm::Run(const Operation& operation, - float alpha, - const void* A, - const MatrixLayout& Adesc, - const void* U, - const MatrixLayout& Udesc, - const void* B, - const MatrixLayout& Bdesc, - const void* V, - const MatrixLayout& Vdesc, - float beta, - const void* C, - const MatrixLayout& Cdesc, - void* D, - const MatrixLayout& Ddesc, - const Workspace& workspace, - cudaStream_t stream) -{ - - Context context{*impl_->props_}; - - const auto desc = context.Init(operation, Adesc, Udesc, Bdesc, Vdesc, Cdesc, Ddesc); - - if (!desc) { - fprintf(stderr, "invalid argument.\n"); - TM_CHECK(0); - return -1; - } - - const auto launch = [=](LaunchSpec spec, cudaStream_t st) { - auto _workspace = workspace; - return spec.kernel->Launch(operation, - alpha, - A, - Adesc, - U, - Udesc, - B, - Bdesc, - V, - Vdesc, - beta, - C, - Cdesc, - D, - Ddesc, - spec.swizzle, - spec.splits, - _workspace, - st); - }; - - std::optional dispatch_context_storage; - Context* dispatch_context = &context; - if (operation.dispatch_num_override > 0 && operation.dispatch_num_override != context.desc().num) { - MatrixLayout dispatch_Adesc = Adesc; - MatrixLayout dispatch_Bdesc = Bdesc; - MatrixLayout dispatch_Ddesc = Ddesc; - dispatch_Adesc.num = operation.dispatch_num_override; - dispatch_Bdesc.num = operation.dispatch_num_override; - dispatch_Ddesc.num = operation.dispatch_num_override; - dispatch_context_storage.emplace(*impl_->props_); - if (dispatch_context_storage->Init( - operation, dispatch_Adesc, Udesc, dispatch_Bdesc, Vdesc, Cdesc, dispatch_Ddesc)) { - dispatch_context = &*dispatch_context_storage; - } - else { - dispatch_context_storage.reset(); - dispatch_context = &context; - } - } +int Gemm::Run(const Operation& operation, float alpha, const void* A, + const MatrixLayout& Adesc, const void* U, + const MatrixLayout& Udesc, const void* B, + const MatrixLayout& Bdesc, const void* V, + const MatrixLayout& Vdesc, float beta, const void* C, + const MatrixLayout& Cdesc, void* D, const MatrixLayout& Ddesc, + const Workspace& workspace, cudaStream_t stream) { + Context context{*impl_->props_}; + + const auto desc = + context.Init(operation, Adesc, Udesc, Bdesc, Vdesc, Cdesc, Ddesc); + + if (!desc) { + fprintf(stderr, "invalid argument.\n"); + TM_CHECK(0); + return -1; + } + + const auto launch = [=](LaunchSpec spec, cudaStream_t st) { + auto _workspace = workspace; + return spec.kernel->Launch(operation, alpha, A, Adesc, U, Udesc, B, Bdesc, + V, Vdesc, beta, C, Cdesc, D, Ddesc, spec.swizzle, + spec.splits, _workspace, st); + }; + + std::optional dispatch_context_storage; + Context* dispatch_context = &context; + if (operation.dispatch_num_override > 0 && + operation.dispatch_num_override != context.desc().num) { + MatrixLayout dispatch_Adesc = Adesc; + MatrixLayout dispatch_Bdesc = Bdesc; + MatrixLayout dispatch_Ddesc = Ddesc; + dispatch_Adesc.num = operation.dispatch_num_override; + dispatch_Bdesc.num = operation.dispatch_num_override; + dispatch_Ddesc.num = operation.dispatch_num_override; + dispatch_context_storage.emplace(*impl_->props_); + if (dispatch_context_storage->Init(operation, dispatch_Adesc, Udesc, + dispatch_Bdesc, Vdesc, Cdesc, + dispatch_Ddesc)) { + dispatch_context = &*dispatch_context_storage; + } else { + dispatch_context_storage.reset(); + dispatch_context = &context; + } + } #if 0 if (operation.reserved) { @@ -687,70 +773,70 @@ int Gemm::Run(const Operation& operation, } #endif - LaunchSpec spec{}; - - const bool measured = operation.dispatch & DispatchPolicy::kMeasure; - { - std::lock_guard lock(impl_->dispatch_mutex_); - if (measured) { - impl_->Measure(*dispatch_context, workspace.barriers_size, workspace.partials_size, 1, launch, stream); - } - - spec = impl_->Dispatch(*dispatch_context, operation.dispatch, workspace.barriers_size, workspace.partials_size); - const bool preserve_default_kernel = operation.dispatch & DispatchPolicy::kPreserveDefaultSplits; - const bool preserve_default_split_count = operation.dispatch & DispatchPolicy::kPreserveDefaultSplitCount; - if (spec.kernel && (preserve_default_kernel || preserve_default_split_count)) { - auto default_specs = impl_->Find(*dispatch_context, workspace.barriers_size, workspace.partials_size, 1); - if (!default_specs.empty()) { - const auto default_spec = default_specs.front(); - if (preserve_default_kernel) { - spec.kernel = default_spec.kernel; - } - spec.splits = default_spec.splits; - const auto& default_desc = dispatch_context->get_desc(*spec.kernel); - spec.swizzle = std::min( - spec.swizzle, - spec.kernel->GetMaxSwizzle({ - default_desc.m, - default_desc.n, - default_desc.k, - default_desc.num, - })); - } - } - } - if (spec.kernel && dispatch_context != &context) { - const auto& actual_desc = context.get_desc(*spec.kernel); - spec.swizzle = - std::min(spec.swizzle, spec.kernel->GetMaxSwizzle({actual_desc.m, actual_desc.n, actual_desc.k, actual_desc.num})); - } - MaybeTraceGemmDispatch(dispatch_context->desc(), operation.dispatch, spec, measured); - - if (spec.kernel) { - // std::cout << "[Gemm] dispatch: " << spec.kernel->name() // - // << " split_k=" << spec.splits // - // << " swizzle=" << spec.swizzle << std::endl; - return launch(spec, stream); - } - - TM_CHECK(0) << "No feasible kernel found for the problem: " << to_string(context.desc()); - - return -1; + LaunchSpec spec{}; + + const bool measured = operation.dispatch & DispatchPolicy::kMeasure; + { + std::lock_guard lock(impl_->dispatch_mutex_); + if (measured) { + impl_->Measure(*dispatch_context, workspace.barriers_size, + workspace.partials_size, 1, launch, stream); + } + + spec = impl_->Dispatch(*dispatch_context, operation.dispatch, + workspace.barriers_size, workspace.partials_size); + const bool preserve_default_kernel = + operation.dispatch & DispatchPolicy::kPreserveDefaultSplits; + const bool preserve_default_split_count = + operation.dispatch & DispatchPolicy::kPreserveDefaultSplitCount; + if (spec.kernel && + (preserve_default_kernel || preserve_default_split_count)) { + auto default_specs = + impl_->Find(*dispatch_context, workspace.barriers_size, + workspace.partials_size, 1, false); + if (!default_specs.empty()) { + const auto default_spec = default_specs.front(); + if (preserve_default_kernel) { + spec.kernel = default_spec.kernel; + } + spec.splits = default_spec.splits; + const auto& default_desc = dispatch_context->get_desc(*spec.kernel); + spec.swizzle = std::min(spec.swizzle, spec.kernel->GetMaxSwizzle({ + default_desc.m, + default_desc.n, + default_desc.k, + default_desc.num, + })); + } + } + } + if (spec.kernel && dispatch_context != &context) { + const auto& actual_desc = context.get_desc(*spec.kernel); + spec.swizzle = + std::min(spec.swizzle, + spec.kernel->GetMaxSwizzle({actual_desc.m, actual_desc.n, + actual_desc.k, actual_desc.num})); + } + MaybeTraceGemmDispatch(dispatch_context->desc(), operation.dispatch, spec, + measured); + + if (spec.kernel) { + // std::cout << "[Gemm] dispatch: " << spec.kernel->name() // + // << " split_k=" << spec.splits // + // << " swizzle=" << spec.swizzle << std::endl; + return launch(spec, stream); + } + + TM_CHECK(0) << "No feasible kernel found for the problem: " + << to_string(context.desc()); + + return -1; } -int Gemm::Export(std::ostream& os) -{ - return impl_->cache_.Export(os); -} +int Gemm::Export(std::ostream& os) { return impl_->cache_.Export(os); } -int Gemm::Import(std::istream& is) -{ - return impl_->cache_.Import(is); -} +int Gemm::Import(std::istream& is) { return impl_->cache_.Import(is); } -std::vector Gemm::GetTuningSeq() const -{ - return impl_->tuning_.seq; -} +std::vector Gemm::GetTuningSeq() const { return impl_->tuning_.seq; } } // namespace turbomind::gemm diff --git a/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/iterator_sm70.h b/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/iterator_sm70.h index fc52463377..d7f4067595 100644 --- a/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/iterator_sm70.h +++ b/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/iterator_sm70.h @@ -16,262 +16,256 @@ namespace turbomind::gemm { -template -inline __device__ void _Ld(Array& dst, const T* src) -{ - static_assert(sizeof(Array) <= sizeof(uint4)); - - if constexpr (sizeof(Array) == sizeof(uint4)) { - (uint4&)dst = __ldcs((const uint4*)src); - } - else if constexpr (sizeof(Array) == sizeof(uint2)) { - (uint2&)dst = __ldcs((const uint2*)src); - } - else if constexpr (sizeof(Array) == sizeof(uint)) { - (uint&)dst = __ldcs((const uint*)src); - } - else { - static_assert(!std::is_same_v); - } +template +inline __device__ void _Ld(Array& dst, const T* src) { + static_assert(sizeof(Array) <= sizeof(uint4)); + + if constexpr (sizeof(Array) == sizeof(uint4)) { + (uint4&)dst = __ldcs((const uint4*)src); + } else if constexpr (sizeof(Array) == sizeof(uint2)) { + (uint2&)dst = __ldcs((const uint2*)src); + } else if constexpr (sizeof(Array) == sizeof(uint)) { + (uint&)dst = __ldcs((const uint*)src); + } else { + static_assert(!std::is_same_v); + } } -template +template struct GmemIteratorSm70 { + using ThreadMap = Map; - using ThreadMap = Map; - - using AccessType = Array; - using Pointer = get_pointer_type; - - using Policy = Policy_; + using AccessType = Array; + using Pointer = get_pointer_type; - static constexpr int ITER_S = Map::kIterS; - static constexpr int ITER_C = Map::kIterC; + using Policy = Policy_; - static constexpr Striding kMode = mode; - static constexpr bool is_indexed = mode == Striding::kIndexed; + static constexpr int ITER_S = Map::kIterS; + static constexpr int ITER_C = Map::kIterC; - const char* src_data_; + static constexpr Striding kMode = mode; + static constexpr bool is_indexed = mode == Striding::kIndexed; - int src_offset_; - int dst_offset_; + const char* src_data_; - int offset_c_; - int offset_s_; + int src_offset_; + int dst_offset_; - int src_step_c_; - int src_step_s_; + int offset_c_; + int offset_s_; - int src_step_k_; + int src_step_c_; + int src_step_s_; - Predicate pred_; + int src_step_k_; - bool g_mask{true}; + Predicate + pred_; - SmemAccessor smem_data_; + bool g_mask{true}; - static constexpr int2 kMK0 = cs2mk(SmemLayout::C0, SmemLayout::S0); - static constexpr int kPeriodC = ceil_div(SmemLayout::C0, Map::kDeltaC); - static constexpr int kPeriodS = ceil_div(SmemLayout::S0, Map::kDeltaS); + SmemAccessor smem_data_; - int phases_[kPeriodS][kPeriodC]; + static constexpr int2 kMK0 = cs2mk(SmemLayout::C0, SmemLayout::S0); + static constexpr int kPeriodC = ceil_div(SmemLayout::C0, Map::kDeltaC); + static constexpr int kPeriodS = ceil_div(SmemLayout::S0, Map::kDeltaS); - const char* src_data_vec_[ITER_S]; + int phases_[kPeriodS][kPeriodC]; - using Fragments = AccessType[Map::kIterS][Map::kIterC]; + const char* src_data_vec_[ITER_S]; - __device__ static constexpr int2 pack(int2 mk) - { - return Packing_v2::apply(mk); - } + using Fragments = AccessType[Map::kIterS][Map::kIterC]; - __device__ static constexpr int2 to_cs(int2 mk) - { - return mk2cs(mk.x, mk.y); - } + __device__ static constexpr int2 pack(int2 mk) { + return Packing_v2::apply(mk); + } - __device__ GmemIteratorSm70(): smem_data_{Pointer{nullptr}} {}; + __device__ static constexpr int2 to_cs(int2 mk) { + return mk2cs(mk.x, mk.y); + } - __device__ GmemIteratorSm70(const MatrixData& mat, int2 offset, int2 extent): smem_data_{Pointer{(T*)nullptr}} - { - const int warp_id = threadIdx.x / WARP_SIZE; - const int lane_id = threadIdx.x % WARP_SIZE; + __device__ GmemIteratorSm70() : smem_data_{Pointer{nullptr}} {}; - const Pointer data{(T*)mat.ptr.ptr}; - const int ld = mat.ptr.stride; + __device__ GmemIteratorSm70(const MatrixData& mat, int2 offset, int2 extent) + : smem_data_{Pointer{(T*)nullptr}} { + const int warp_id = threadIdx.x / WARP_SIZE; + const int lane_id = threadIdx.x % WARP_SIZE; - const int2 offsets = Map::get_offset(warp_id, lane_id); + const Pointer data{(T*)mat.ptr.ptr}; + const int ld = mat.ptr.stride; - offset_c_ = offsets.x; - offset_s_ = offsets.y; + const int2 offsets = Map::get_offset(warp_id, lane_id); - // auto src_ptr = reinterpret_cast((T*)data); + offset_c_ = offsets.x; + offset_s_ = offsets.y; - if constexpr (pred_.is_active) { - extent = to_cs(pack(extent)); - PRAGMA_UNROLL - for (int s = 0; s < Map::kIterS; ++s) { - PRAGMA_UNROLL - for (int c = 0; c < Map::kIterC; ++c) { - int ss = offset_s_ + s * Map::kDeltaS; - int cc = offset_c_ + c * Map::kDeltaC; - if (ss < extent.y && cc < extent.x) { - pred_.set(s, c); - } - } - } - } + // auto src_ptr = reinterpret_cast((T*)data); + if constexpr (pred_.is_active) { + extent = to_cs(pack(extent)); + PRAGMA_UNROLL + for (int s = 0; s < Map::kIterS; ++s) { PRAGMA_UNROLL - for (int s = 0; s < kPeriodS; ++s) { - PRAGMA_UNROLL - for (int c = 0; c < kPeriodC; ++c) { - phases_[s][c] = SmemLayout::apply(offset_s_ + s * Map::kDeltaS, offset_c_ + c * Map::kDeltaC); - } + for (int c = 0; c < Map::kIterC; ++c) { + int ss = offset_s_ + s * Map::kDeltaS; + int cc = offset_c_ + c * Map::kDeltaC; + if (ss < extent.y && cc < extent.x) { + pred_.set(s, c); + } } + } + } - const int src_offset = is_indexed ? offsets.x : offsets.x + offsets.y * ld; + PRAGMA_UNROLL + for (int s = 0; s < kPeriodS; ++s) { + PRAGMA_UNROLL + for (int c = 0; c < kPeriodC; ++c) { + phases_[s][c] = SmemLayout::apply(offset_s_ + s * Map::kDeltaS, + offset_c_ + c * Map::kDeltaC); + } + } - src_offset_ = src_offset * bitsof / bitsof; + const int src_offset = is_indexed ? offsets.x : offsets.x + offsets.y * ld; - src_step_c_ = bitsof * Map::kDeltaC / bitsof; - src_step_s_ = bitsof * Map::kDeltaS * ld / bitsof; + src_offset_ = src_offset * bitsof / bitsof; - src_step_k_ = bitsof * cs2mk(Map::kDimC, Map::kDimS * ld).y / bitsof; + src_step_c_ = bitsof * Map::kDeltaC / bitsof; + src_step_s_ = bitsof * Map::kDeltaS * ld / bitsof; - // initialize for the first tile - if constexpr (is_indexed) { - const int2 cta_cs = to_cs(offset); - for (int s = 0; s < ITER_S; ++s) { - const int ss = cta_cs.y + offset_s_ + s * Map::kDeltaS; - const int idx = (mat.idxs && pred_(s, 0)) ? __ldg(mat.idxs + ss) : ss; - const auto tmp = data + cs2idx({cta_cs.x, idx}, ld); - src_data_vec_[s] = reinterpret_cast((T*)tmp) + src_offset_; - } - } - else { - auto src_data = data + cs2idx(to_cs(pack(offset)), ld); - src_data_ = reinterpret_cast((T*)src_data) + src_offset_; - } - } + src_step_k_ = + bitsof * cs2mk(Map::kDimC, Map::kDimS * ld).y / bitsof; - __device__ constexpr int _src_step_k() const - { - return src_step_k_; + // initialize for the first tile + if constexpr (is_indexed) { + const int2 cta_cs = to_cs(offset); + for (int s = 0; s < ITER_S; ++s) { + const int ss = cta_cs.y + offset_s_ + s * Map::kDeltaS; + const int idx = (mat.idxs && pred_(s, 0)) ? __ldg(mat.idxs + ss) : ss; + const auto tmp = data + cs2idx({cta_cs.x, idx}, ld); + src_data_vec_[s] = reinterpret_cast((T*)tmp) + src_offset_; + } + } else { + auto src_data = data + cs2idx(to_cs(pack(offset)), ld); + src_data_ = reinterpret_cast((T*)src_data) + src_offset_; } - - __device__ void ClearSmem(int pipe_iter = 0) - { - PRAGMA_UNROLL - for (int s = 0; s < Map::kIterS; ++s) { - PRAGMA_UNROLL - for (int c = 0; c < Map::kIterC; ++c) { - const int pred_s = offset_s_ + s * Map::kDeltaS < Map::kDimS; - const int pred_c = offset_c_ + c * Map::kDeltaC < Map::kDimC; - auto ptr = &smem_data_(offset_s_ + s * Map::kDeltaS, offset_c_ + c * Map::kDeltaC); - if ((Map::kAlignedC && Map::kAlignedS) || (pred_s && pred_c)) { - turbomind::Store(ptr, Array{}); - } - } + } + + __device__ constexpr int _src_step_k() const { return src_step_k_; } + + __device__ void ClearSmem(int pipe_iter = 0) { + PRAGMA_UNROLL + for (int s = 0; s < Map::kIterS; ++s) { + PRAGMA_UNROLL + for (int c = 0; c < Map::kIterC; ++c) { + const int pred_s = offset_s_ + s * Map::kDeltaS < Map::kDimS; + const int pred_c = offset_c_ + c * Map::kDeltaC < Map::kDimC; + auto ptr = &smem_data_(offset_s_ + s * Map::kDeltaS, + offset_c_ + c * Map::kDeltaC); + if ((Map::kAlignedC && Map::kAlignedS) || (pred_s && pred_c)) { + turbomind::Store(ptr, Array{}); } + } } + } - __device__ void Advance() - { - if constexpr (!is_indexed) { - if (!g_mask) { - src_data_ -= _src_step_k(); - } - } + __device__ void Advance() { + if constexpr (!is_indexed) { + if (!g_mask) { + src_data_ -= _src_step_k(); + } } - - __device__ void Copy(std::true_type, T* dst, const char* __restrict__ src, bool mask) - { - if (mask) { - AccessType frag; - if constexpr (Policy_::kEvictPolicy != EvictPolicy::kEvictNormal) { - _Ld(frag, (const T*)src); - } - else { - Ldg(frag, (const T*)src); - } - turbomind::Store(dst, frag); - } + } + + __device__ void Copy(std::true_type, T* dst, const char* __restrict__ src, + bool mask) { + if (mask) { + AccessType frag; + if constexpr (Policy_::kEvictPolicy != EvictPolicy::kEvictNormal) { + _Ld(frag, (const T*)src); + } else { + Ldg(frag, (const T*)src); + } + turbomind::Store(dst, frag); } - - __device__ void Fetch(Fragments& frags, bool tile_mask) - { - PRAGMA_UNROLL - for (int s = 0; s < Map::kIterS; ++s) { - - if constexpr (is_indexed) { - src_data_ = src_data_vec_[s]; - } - - PRAGMA_UNROLL - for (int c = 0; c < Map::kIterC; ++c) { - Copy2(frags[s][c], src_data_ + src_step_c_ * c, tile_mask && g_mask && pred_(s, c)); - } - - if constexpr (is_indexed) { - src_data_vec_[s] += _src_step_k(); - } - else { - src_data_ += src_step_s_; - if (s == Map::kIterS - 1) { - src_data_ -= src_step_s_ * Map::kIterS; - src_data_ += _src_step_k(); - } - } + } + + __device__ void Fetch(Fragments& frags, bool tile_mask) { + PRAGMA_UNROLL + for (int s = 0; s < Map::kIterS; ++s) { + if constexpr (is_indexed) { + src_data_ = src_data_vec_[s]; + } + + PRAGMA_UNROLL + for (int c = 0; c < Map::kIterC; ++c) { + Copy2(frags[s][c], src_data_ + src_step_c_ * c, + tile_mask && g_mask && pred_(s, c)); + } + + if constexpr (is_indexed) { + src_data_vec_[s] += _src_step_k(); + } else { + src_data_ += src_step_s_; + if (s == Map::kIterS - 1) { + src_data_ -= src_step_s_ * Map::kIterS; + src_data_ += _src_step_k(); } + } } - - __device__ void Store(Fragments& frags) - { - PRAGMA_UNROLL - for (int s = 0; s < Map::kIterS; ++s) { - PRAGMA_UNROLL - for (int c = 0; c < Map::kIterC; ++c) { - // auto dst = &smem_data_(offset_s_ + s * Map::kDeltaS, offset_c_ + c * Map::kDeltaC); - - const int i0 = SmemLayout::apply( // - s / kPeriodS * kPeriodS * Map::kDeltaS, - c / kPeriodC * kPeriodC * Map::kDeltaC); - const int i1 = phases_[s % kPeriodS][c % kPeriodC]; - auto dst = &smem_data_.ptr_[i0 + i1]; - - if (pred_(s, c)) { - turbomind::Store(dst, frags[s][c]); - } - } + } + + __device__ void Store(Fragments& frags) { + PRAGMA_UNROLL + for (int s = 0; s < Map::kIterS; ++s) { + PRAGMA_UNROLL + for (int c = 0; c < Map::kIterC; ++c) { + // auto dst = &smem_data_(offset_s_ + s * Map::kDeltaS, offset_c_ + c * + // Map::kDeltaC); + + const int i0 = SmemLayout::apply( // + s / kPeriodS * kPeriodS * Map::kDeltaS, + c / kPeriodC * kPeriodC * Map::kDeltaC); + const int i1 = phases_[s % kPeriodS][c % kPeriodC]; + auto dst = &smem_data_.ptr_[i0 + i1]; + + if (pred_(s, c)) { + turbomind::Store(dst, frags[s][c]); } + } } - - __device__ void Copy2(AccessType& frag, const char* __restrict__ src, bool mask) - { - if (mask) { - if constexpr (Policy_::kEvictPolicy != EvictPolicy::kEvictNormal) { - _Ld(frag, (const T*)src); - } - else { - Ldg(frag, (const T*)src); - } - } + } + + __device__ void Copy2(AccessType& frag, const char* __restrict__ src, + bool mask) { + if (mask) { + if constexpr (Policy_::kEvictPolicy != EvictPolicy::kEvictNormal) { + _Ld(frag, (const T*)src); + } else { + Ldg(frag, (const T*)src); + } } + } }; -template +template struct IteratorSm70 { - template - using Type = GmemIteratorSm70; + template + using Type = GmemIteratorSm70; +}; + +// Exact-shape kernels may promise that every CTA covers a complete M/N/K tile. +// Keep ThreadMap's intrinsic bounds (for example the one-row scale tile), but +// remove runtime edge predicates from otherwise aligned A/B tiles. +template +struct IteratorSm70FullTile { + template + using Type = GmemIteratorSm70; }; } // namespace turbomind::gemm diff --git a/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/kernel/sm70_884_4.cu b/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/kernel/sm70_884_4.cu index 01d4db3296..ed4dce6012 100644 --- a/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/kernel/sm70_884_4.cu +++ b/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/kernel/sm70_884_4.cu @@ -13,32 +13,30 @@ using D = cache_policy::Default; namespace { -// Keep shape-specific small-N tactics out of prefill and unrelated CUDA graph shapes. -template -class ExactMKernelImpl final: public KernelImpl { -public: - bool is_feasible(const GemmDesc& desc) const noexcept override - { - return desc.m == ExactM && KernelImpl::is_feasible(desc); - } +// Keep shape-specific small-N tactics out of prefill and unrelated CUDA graph +// shapes. +template +class ExactMKernelImpl final : public KernelImpl { + public: + bool is_feasible(const GemmDesc& desc) const noexcept override { + return desc.m == ExactM && KernelImpl::is_feasible(desc); + } }; -template -class ExactMnkKernelImpl final: public KernelImpl { -public: - bool is_feasible(const GemmDesc& desc) const noexcept override - { - return desc.m == ExactM && desc.n == ExactN && desc.k == ExactK - && KernelImpl::is_feasible(desc); - } +template +class ExactMnkKernelImpl final : public KernelImpl { + public: + bool is_feasible(const GemmDesc& desc) const noexcept override { + return desc.m == ExactM && desc.n == ExactN && desc.k == ExactK && + KernelImpl::is_feasible(desc); + } }; } // namespace -void Registry::sm70_884_4() -{ - if constexpr (1) { - // clang-format off +void Registry::sm70_884_4() { + if constexpr (1) { + // clang-format off using C = Config_U4_d; Add>(); Add>(); @@ -65,11 +63,11 @@ void Registry::sm70_884_4() Add(std::make_unique>()); Add(std::make_unique>()); - // clang-format on - } + // clang-format on + } - if constexpr (1) { - // clang-format off + if constexpr (1) { + // clang-format off // GroupSizeV=128 using C = Config_U4_g; Add>(); @@ -98,30 +96,38 @@ void Registry::sm70_884_4() Add>(); Add>(); Add>(); - // clang-format on - } + // clang-format on + } - if constexpr (1) { - // clang-format off + if constexpr (1) { + // clang-format off using C = Config_MXF4; Add>(); Add>(); Add>(); Add>(); Add>(); - // clang-format on - } + // clang-format on + } - if constexpr (1) { - // clang-format off + if constexpr (1) { + // clang-format off using C = Config_NVF4; Add>(); Add>(); Add>(); Add>(); Add>(); - // clang-format on - } + using C32K64L1 = C::Type<8, 32, 64, 1, 1, 1, D, S, 2, true, 1, 16>; + using C32K64L2 = C::Type<8, 32, 64, 1, 1, 1, D, S, 2, true, 1, 16, -1, -1, 2>; + using C32K128L1 = C::Type<8, 32, 128, 1, 1, 1, D, S, 2, true, 1, 16>; + using C32K128L2 = C::Type<8, 32, 128, 1, 1, 1, D, S, 2, true, 1, 16, -1, -1, 2>; + Add(std::make_unique>()); + Add(std::make_unique>()); + Add(std::make_unique>()); + Add(std::make_unique>()); + // clang-format on + } } } // namespace turbomind::gemm 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 7ef3acf80b..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 @@ -11,10 +11,48 @@ using namespace cache_policy; using S = cache_policy::Stream; using D = cache_policy::Default; -void Registry::sm70_884_8() -{ - if constexpr (1) { - // clang-format off +namespace { + +template +class Qwen38PrescaledFullTileKernelImpl final : public KernelImpl { + public: + Qwen38PrescaledFullTileKernelImpl() { + this->info_.name += "_sm70_fp8_pscale_full"; + } + + bool is_feasible(const GemmDesc& desc) const noexcept override { + return desc.m == 8000 && (desc.n == 4096 || desc.n == 3584) && + desc.k == 5120 && desc.num == 1 && + KernelImpl::is_feasible(desc); + } +}; + +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() { + if constexpr (1) { + // clang-format off using C = Config_E4M3; Add>(); Add>(); @@ -29,8 +67,15 @@ void Registry::sm70_884_8() Add>(); Add>(); Add>(); - // clang-format on - } + + 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 + } } } // namespace turbomind::gemm diff --git a/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/transform.h b/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/transform.h index 1cbfda551a..48aeaf4be5 100644 --- a/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/transform.h +++ b/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/transform.h @@ -13,147 +13,181 @@ namespace turbomind::gemm { struct Transform_Default { - template - __device__ static void apply(Array (&frag)[K][Mf], int k, Array (&data)[K][Md], S&, int div) - { - static_assert(Nf * Mf == Nd * Md); - static_assert(Nd % Nf == 0 && Mf % Md == 0); - static_assert(sizeof(frag) == sizeof(data)); - - // Alignment must be manually enforced for `reinterpret_cast` - auto& frag_k = reinterpret_cast(&)[Md]>(frag[k]); - auto& data_k = data[k]; - - PRAGMA_UNROLL - for (int i = 0; i < std::size(frag_k); ++i) { - frag_k[i] = data_k[i]; - } + template + __device__ static void apply(Array (&frag)[K][Mf], int k, + Array (&data)[K][Md], S&, int div) { + static_assert(Nf * Mf == Nd * Md); + static_assert(Nd % Nf == 0 && Mf % Md == 0); + static_assert(sizeof(frag) == sizeof(data)); + + // Alignment must be manually enforced for `reinterpret_cast` + auto& frag_k = reinterpret_cast(&)[Md]>(frag[k]); + auto& data_k = data[k]; + + PRAGMA_UNROLL + for (int i = 0; i < std::size(frag_k); ++i) { + frag_k[i] = data_k[i]; } + } }; -template +template struct Transform_HMMA_16816 { - template - __device__ static void - apply(Array (&frag)[K][Mf], int k, Array (&data)[K][Md], Array (&stat)[Ks][Ms], int div) - { - static_assert(Nf * Mf == Nd * Md); - static_assert(Nd % Nf == 0 && Mf % Md == 0); - static_assert(Nf * Mf == Ns * Ms * 4); - - auto& frag_k = reinterpret_cast(&)[Md]>(frag[k]); - auto& stat_k = reinterpret_cast(&)[Ns * Ms]>(stat[k / div]); - auto& data_k = data[k]; - + template + __device__ static void apply(Array (&frag)[K][Mf], int k, + Array (&data)[K][Md], + Array (&stat)[Ks][Ms], int div) { + static_assert(Nf * Mf == Nd * Md); + static_assert(Nd % Nf == 0 && Mf % Md == 0); + static_assert(Nf * Mf == Ns * Ms * 4); + + auto& frag_k = reinterpret_cast(&)[Md]>(frag[k]); + auto& stat_k = reinterpret_cast(&)[Ns * Ms]>(stat[k / div]); + auto& data_k = data[k]; + + PRAGMA_UNROLL + for (int m = 0; m < Md; ++m) { + auto tmp = ConvertKvCache::convert(data_k[m]); + static_assert(Nd % 8 == 0); + PRAGMA_UNROLL + for (int i = 0; i < Nd; i += 8) { PRAGMA_UNROLL - for (int m = 0; m < Md; ++m) { - auto tmp = ConvertKvCache::convert(data_k[m]); - static_assert(Nd % 8 == 0); - PRAGMA_UNROLL - for (int i = 0; i < Nd; i += 8) { - PRAGMA_UNROLL - for (int s = 0; s < 2; ++s) { - PRAGMA_UNROLL - for (int c = 0; c < 2; ++c) { - const int idx = (m * Nd + i) / 8 * 2 + s * StatStepS + c * StatStepC; - dequant((Array&)tmp[i + s * 4 + c * 2], stat_k[idx]); - } - } - } - - frag_k[m] = tmp; + for (int s = 0; s < 2; ++s) { + PRAGMA_UNROLL + for (int c = 0; c < 2; ++c) { + const int idx = + (m * Nd + i) / 8 * 2 + s * StatStepS + c * StatStepC; + dequant((Array&)tmp[i + s * 4 + c * 2], stat_k[idx]); + } } - } + } - template - __device__ static void dequant(Array& x, Array s) - { - Array& _s = (Array&)s; - x[0] = __hfma(x[0], _s[0], _s[1]); - x[1] = __hfma(x[1], _s[0], _s[1]); - } - - __device__ static void dequant(Array& x, Array s) - { - bfloat16_t s1 = __ushort_as_bfloat16((uint16_t)s[0] << 7); - x[0] = __hmul(x[0], s1); - x[1] = __hmul(x[1], s1); - } - - __device__ static void dequant(Array& x, Array s) - { - // half_t s1 = __ushort_as_half(((uint16_t)s[0] + 15 - 127) << 10); - // Adjusted in `AdjustUe8m0ScaleForHalf` - half_t s1 = __ushort_as_half((uint16_t)s[0] << 10); - x[0] = __hmul(x[0], s1); - x[1] = __hmul(x[1], s1); - } - - __device__ static void dequant(Array& x, Array s) - { - auto s1 = __ushort_as_bfloat16(s[0]); - x[0] = __hmul(x[0], s1); - x[1] = __hmul(x[1], s1); - } - - __device__ static void dequant(Array& x, Array s) - { - auto s1 = __ushort_as_half(s[0]); - x[0] = __hmul(x[0], s1); - x[1] = __hmul(x[1], s1); + frag_k[m] = tmp; } + } + + template + __device__ static void dequant(Array& x, Array s) { + Array& _s = (Array&)s; + x[0] = __hfma(x[0], _s[0], _s[1]); + x[1] = __hfma(x[1], _s[0], _s[1]); + } + + __device__ static void dequant(Array& x, Array s) { + bfloat16_t s1 = __ushort_as_bfloat16((uint16_t)s[0] << 7); + x[0] = __hmul(x[0], s1); + x[1] = __hmul(x[1], s1); + } + + __device__ static void dequant(Array& x, Array s) { + // half_t s1 = __ushort_as_half(((uint16_t)s[0] + 15 - 127) << 10); + // Adjusted in `AdjustUe8m0ScaleForHalf` + half_t s1 = __ushort_as_half((uint16_t)s[0] << 10); + x[0] = __hmul(x[0], s1); + x[1] = __hmul(x[1], s1); + } + + __device__ static void dequant(Array& x, + Array s) { + auto s1 = __ushort_as_bfloat16(s[0]); + x[0] = __hmul(x[0], s1); + x[1] = __hmul(x[1], s1); + } + + __device__ static void dequant(Array& x, Array s) { + auto s1 = __ushort_as_half(s[0]); + x[0] = __hmul(x[0], s1); + x[1] = __hmul(x[1], s1); + } }; // Used by SM70 MMA struct Transform_HMMA_SIMT_B { - template - __device__ static void - apply(Array (&frag)[K][Mf], int k, Array (&data)[K][Md], Array (&stat)[Ks][Ms], int div) - { - static_assert(Nf * Mf == Nd * Md); - static_assert(Nd % Nf == 0 && Mf % Md == 0); - - auto& frag_k = reinterpret_cast(&)[Md]>(frag[k]); - auto& stat_k = reinterpret_cast(&)[Ns * Ms]>(stat[k / div]); - auto& data_k = data[k]; - - // static_assert(Nf != Nf); - - PRAGMA_UNROLL - for (int m = 0; m < Md; ++m) { - auto tmp = ConvertKvCache::convert(data_k[m]); - PRAGMA_UNROLL - for (int i = 0; i < Nd; i += 2) { - dequant((Array&)tmp[i], stat_k[(m * Nd + i) / Nf]); - } - frag_k[m] = tmp; - } - } - - template - __device__ static void dequant(Array& x, Array s) - { - Array& _s = (Array&)s; - - x[0] = __hfma(x[0], _s[0], _s[1]); - x[1] = __hfma(x[1], _s[0], _s[1]); - } - - __device__ static void dequant(Array& x, Array s) - { - // half_t s1 = __ushort_as_half(((uint16_t)s[0] + 15 - 127) << 10); - // Adjusted in `AdjustUe8m0ScaleForHalf` - half_t s1 = __ushort_as_half((uint16_t)s[0] << 10); - x[0] = __hmul(x[0], s1); - x[1] = __hmul(x[1], s1); + template + __device__ static void apply(Array (&frag)[K][Mf], int k, + Array (&data)[K][Md], + Array (&stat)[Ks][Ms], int div) { + static_assert(Nf * Mf == Nd * Md); + static_assert(Nd % Nf == 0 && Mf % Md == 0); + + auto& frag_k = reinterpret_cast(&)[Md]>(frag[k]); + auto& stat_k = reinterpret_cast(&)[Ns * Ms]>(stat[k / div]); + auto& data_k = data[k]; + + // static_assert(Nf != Nf); + + PRAGMA_UNROLL + for (int m = 0; m < Md; ++m) { + auto tmp = ConvertKvCache::convert(data_k[m]); + PRAGMA_UNROLL + for (int i = 0; i < Nd; i += 2) { + dequant((Array&)tmp[i], stat_k[(m * Nd + i) / Nf]); + } + frag_k[m] = tmp; } + } + + template + __device__ static void dequant(Array& x, Array s) { + Array& _s = (Array&)s; + + x[0] = __hfma(x[0], _s[0], _s[1]); + x[1] = __hfma(x[1], _s[0], _s[1]); + } + + __device__ static void dequant(Array& x, Array s) { + // half_t s1 = __ushort_as_half(((uint16_t)s[0] + 15 - 127) << 10); + // Adjusted in `AdjustUe8m0ScaleForHalf` + half_t s1 = __ushort_as_half((uint16_t)s[0] << 10); + x[0] = __hmul(x[0], s1); + x[1] = __hmul(x[1], s1); + } + + __device__ static void dequant(Array& x, Array s) { + auto s1 = __ushort_as_half(s[0]); + x[0] = __hmul(x[0], s1); + x[1] = __hmul(x[1], s1); + } +}; - __device__ static void dequant(Array& x, Array s) - { - auto s1 = __ushort_as_half(s[0]); - x[0] = __hmul(x[0], s1); - x[1] = __hmul(x[1], s1); +// FP8 scales that have absorbed the E4M3 exponent-bias factor (256) during +// one-time weight preparation. +struct Transform_HMMA_SIMT_B_PrescaledE4M3 { + template + __device__ static void apply(Array (&frag)[K][Mf], int k, + Array (&data)[K][Md], + Array (&stat)[Ks][Ms], int div) { + static_assert(std::is_same_v); + static_assert(std::is_same_v); + static_assert(std::is_same_v); + static_assert(Nf * Mf == Nd * Md); + static_assert(Nd % Nf == 0 && Mf % Md == 0); + static_assert(Nd % 4 == 0); + + auto& frag_k = reinterpret_cast(&)[Md]>(frag[k]); + auto& stat_k = reinterpret_cast(&)[Ns * Ms]>(stat[k / div]); + auto& data_k = data[k]; + + PRAGMA_UNROLL + for (int m = 0; m < Md; ++m) { + Array tmp; + PRAGMA_UNROLL + for (int i = 0; i < Nd; i += 4) { + auto& src = reinterpret_cast&>(data_k[m][i]); + reinterpret_cast&>(tmp[i]) = cvt_f16x4_e4m3(src); + } + PRAGMA_UNROLL + for (int i = 0; i < Nd; i += 2) { + auto scale = __ushort_as_half(stat_k[(m * Nd + i) / Nf][0]); + tmp[i] = __hmul(tmp[i], scale); + tmp[i + 1] = __hmul(tmp[i + 1], scale); + } + frag_k[m] = tmp; } + } }; } // namespace turbomind::gemm diff --git a/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/types.h b/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/types.h index 5d084456d5..475f623299 100644 --- a/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/types.h +++ b/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/types.h @@ -7,269 +7,244 @@ #include #if ENABLE_BF16 -#include + #include #endif namespace turbomind::gemm { -enum class Order : int -{ - kColMajor = 0, - kRowMajor = 1, +enum class Order : int { + kColMajor = 0, + kRowMajor = 1, }; inline constexpr Order kColMajor = Order::kColMajor; inline constexpr Order kRowMajor = Order::kRowMajor; -constexpr Order operator~(Order a) -{ - return a == kColMajor ? kRowMajor : kColMajor; +constexpr Order operator~(Order a) { + return a == kColMajor ? kRowMajor : kColMajor; } -constexpr const char* to_string(Order order) -{ - switch (order) { - case kColMajor: - return "Col"; - case kRowMajor: - return "Row"; - } - return ""; +constexpr const char* to_string(Order order) { + switch (order) { + case kColMajor: + return "Col"; + case kRowMajor: + return "Row"; + } + return ""; } using Pack = uint32_t; -typedef enum MMA_Tag -{ - HMMA_16816 = 0x100, // sm80+ - HMMA_1688 = 0x200, // sm75 - HMMA_884 = 0x300, // sm70 - HMMA_SIMT = 0x400, // sm75- +typedef enum MMA_Tag { + HMMA_16816 = 0x100, // sm80+ + HMMA_1688 = 0x200, // sm75 + HMMA_884 = 0x300, // sm70 + HMMA_SIMT = 0x400, // sm75- } MMA_Tag; -typedef enum Op_Tag -{ - OPERAND_A = 0x010, - OPERAND_B = 0x020, - OPERAND_U = 0x030, - OPERAND_V = 0x040, - OPERAND_C = 0x050, - OPERAND_D = 0x060, +typedef enum Op_Tag { + OPERAND_A = 0x010, + OPERAND_B = 0x020, + OPERAND_U = 0x030, + OPERAND_V = 0x040, + OPERAND_C = 0x050, + OPERAND_D = 0x060, } Op_Tag; -constexpr MMA_Tag get_mma_tag(Pack pack) -{ - return static_cast(pack & 0xf00); +constexpr MMA_Tag get_mma_tag(Pack pack) { + return static_cast(pack & 0xf00); } -constexpr Op_Tag get_operand_tag(Pack pack) -{ - return static_cast(pack & 0x0f0); +constexpr Op_Tag get_operand_tag(Pack pack) { + return static_cast(pack & 0x0f0); } -constexpr int get_pack_num(Pack pack) -{ - return pack & 0x00f; -} +constexpr int get_pack_num(Pack pack) { return pack & 0x00f; } -enum class Striding : int -{ - kFlat, // [1111,2222,3333] - kRagged, // [11,2222222,333] [0 , 2 , 9 ] - kIndexed, // [xx xxxxxxx xxx], [01, 2345678, 9ab] - kBlocked, // [11][22222][333] +enum class Striding : int { + kFlat, // [1111,2222,3333] + kRagged, // [11,2222222,333] [0 , 2 , 9 ] + kIndexed, // [xx xxxxxxx xxx], [01, 2345678, 9ab] + kBlocked, // [11][22222][333] }; -inline const char* to_string(Striding striding) -{ - switch (striding) { - case Striding::kFlat: - return "f"; - case Striding::kRagged: - return "r"; - case Striding::kIndexed: - return "i"; - case Striding::kBlocked: - return "b"; - default: - return "unknown"; - } +inline const char* to_string(Striding striding) { + switch (striding) { + case Striding::kFlat: + return "f"; + case Striding::kRagged: + return "r"; + case Striding::kIndexed: + return "i"; + case Striding::kBlocked: + return "b"; + default: + return "unknown"; + } } -enum class QuantType : int -{ - kNone = 0, - kK = 1, - kM = 2, - kB = 3, - kDefault = kK, +enum class QuantType : int { + kNone = 0, + kK = 1, + kM = 2, + kB = 3, + kDefault = kK, }; -inline const char* to_string(QuantType q) -{ - switch (q) { - case QuantType::kNone: - return "none"; - case QuantType::kK: - return "k"; - case QuantType::kM: - return "m"; - case QuantType::kB: - return "b"; - default: - return "unknown"; - } +inline const char* to_string(QuantType q) { + switch (q) { + case QuantType::kNone: + return "none"; + case QuantType::kK: + return "k"; + case QuantType::kM: + return "m"; + case QuantType::kB: + return "b"; + default: + return "unknown"; + } } -enum class Epilogue : int -{ - kNone = 0, - kChannelCombination = 0x1, - kGatedSilu = 0x2, - kMoeWeightedReduce = 0x4, - kTileAllReduce = 0x8, +enum class Epilogue : int { + kNone = 0, + kChannelCombination = 0x1, + kGatedSilu = 0x2, + kMoeWeightedReduce = 0x4, + kTileAllReduce = 0x8, }; struct QuantDesc { - QuantType type; - int group_size; + QuantType type; + int group_size; - operator bool() const noexcept - { - return (int)type || group_size; - } + operator bool() const noexcept { return (int)type || group_size; } }; -inline std::string to_string(QuantDesc desc) -{ - if (desc) { - return to_string(desc.type) + std::to_string(desc.group_size); - } - else { - return to_string(desc.type); - } +inline std::string to_string(QuantDesc desc) { + if (desc) { + return to_string(desc.type) + std::to_string(desc.group_size); + } else { + return to_string(desc.type); + } } -enum class DispatchPolicy : int -{ - kDefault = 0, - kMeasure = 1, - kReuse = 2, - kAppend = 3, - kPreserveDefaultSplits = 4, - kPreserveDefaultSplitCount = 8, - kMxfp4MoeGroupedM8Fast = 16, +enum class DispatchPolicy : int { + kDefault = 0, + kMeasure = 1, + kReuse = 2, + kAppend = 3, + kPreserveDefaultSplits = 4, + kPreserveDefaultSplitCount = 8, + kMxfp4MoeGroupedM8Fast = 16, + kSm70Fp8PrefillPrescaled = 32, }; -constexpr bool operator&(const DispatchPolicy& a, const DispatchPolicy& b) -{ - return ((int)a & (int)b); +constexpr bool operator&(const DispatchPolicy& a, const DispatchPolicy& b) { + return ((int)a & (int)b); } -constexpr DispatchPolicy operator|(const DispatchPolicy& a, const DispatchPolicy& b) -{ - return static_cast((int)a | (int)b); +constexpr DispatchPolicy operator|(const DispatchPolicy& a, + const DispatchPolicy& b) { + return static_cast((int)a | (int)b); } class Kernel; class Context; struct Tape { - int ctas; - int max_num; - int max_ctas; - char* buffer; - int4* gemm_shapes; - int4* tiled_shapes; - int4* tile_offsets; - int2* iter_k_ranges; - int* tile_ids; + int ctas; + int max_num; + int max_ctas; + char* buffer; + int4* gemm_shapes; + int4* tiled_shapes; + int4* tile_offsets; + int2* iter_k_ranges; + int* tile_ids; }; struct Operation { - DispatchPolicy dispatch; - Epilogue epilogue; - QuantDesc quant_a; - QuantDesc quant_b; - int batch_dim; - int dispatch_num_override; - int active_group_count; - void* reserved; - void* tile_allreduce; + DispatchPolicy dispatch; + Epilogue epilogue; + QuantDesc quant_a; + QuantDesc quant_b; + int batch_dim; + int dispatch_num_override; + int active_group_count; + void* reserved; + void* tile_allreduce; }; struct MoeWeightedReduceParam { - void* out; - const float* sorted_weights; - const int* offsets; + void* out; + const float* sorted_weights; + const int* offsets; }; struct TileAllReduceParam { - void* signals[8]; - void* self_signal; - void* rank_data; - void* output; - int rank; - int world_size; - int tile_numel; - int output_numel; - int reducer_blocks; - int kernel_reducer_blocks; - int producer_grid_x; - int producer_grid_y; - int producer_grid_z; + void* signals[8]; + void* self_signal; + void* rank_data; + void* output; + int rank; + int world_size; + int tile_numel; + int output_numel; + int reducer_blocks; + int kernel_reducer_blocks; + int producer_grid_x; + int producer_grid_y; + int producer_grid_z; }; -inline Operation transpose(Operation o) -{ - std::swap(o.quant_a, o.quant_b); - o.batch_dim = 1 - o.batch_dim; - return o; +inline Operation transpose(Operation o) { + std::swap(o.quant_a, o.quant_b); + o.batch_dim = 1 - o.batch_dim; + return o; } struct MatrixLayout { - DataType type; - Order order; - int rows; - int cols; - int ld; - Pack pack; - int num; - int* offsets; - int* idxs; - int* group_idxs; + DataType type; + Order order; + int rows; + int cols; + int ld; + Pack pack; + int num; + int* offsets; + int* idxs; + int* group_idxs; }; -inline std::ostream& operator<<(std::ostream& os, const MatrixLayout& x) -{ - os << x.type << " " << to_string(x.order) << " " << x.rows << " " << x.cols << " " << x.num << " " << x.ld; - return os; +inline std::ostream& operator<<(std::ostream& os, const MatrixLayout& x) { + os << x.type << " " << to_string(x.order) << " " << x.rows << " " << x.cols + << " " << x.num << " " << x.ld; + return os; } -inline int64_t byte_size(const MatrixLayout& m) -{ - return byte_size(m.type, (int64_t)m.rows * m.cols); +inline int64_t byte_size(const MatrixLayout& m) { + return byte_size(m.type, (int64_t)m.rows * m.cols); } -inline Striding get_mode(const MatrixLayout& m) -{ - if (m.idxs) { - return Striding::kIndexed; - } - else if (m.ld == 0 || m.offsets) { - return Striding::kBlocked; - } - return Striding::kFlat; +inline Striding get_mode(const MatrixLayout& m) { + if (m.idxs) { + return Striding::kIndexed; + } else if (m.ld == 0 || m.offsets) { + return Striding::kBlocked; + } + return Striding::kFlat; } struct Workspace { - void* barriers; - size_t barriers_size; - void* partials; - size_t partials_size; - void* tensormaps; - size_t tensormaps_size; - int* flags; + void* barriers; + size_t barriers_size; + void* partials; + size_t partials_size; + void* tensormaps; + size_t tensormaps_size; + int* flags; }; } // namespace turbomind::gemm diff --git a/csrc/sm70_turbomind/ops/LICENSE.v100-skinny b/csrc/sm70_turbomind/ops/LICENSE.v100-skinny new file mode 100644 index 0000000000..2649c53a63 --- /dev/null +++ b/csrc/sm70_turbomind/ops/LICENSE.v100-skinny @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 v100-skinny contributors + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu b/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu index cd6bb22a70..d555a8be7d 100644 --- a/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu +++ b/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu @@ -37,6 +37,7 @@ #include "src/turbomind/kernels/gemm/types.h" #include "src/turbomind/kernels/gemm/utils.h" #include "custom_all_reduce.cuh" +#include "qwen38_prefill_cutlass.cuh" namespace turbomind { void unpack_awq_gemm(uint4_t* dst, const uint4_t* src, int rows, int cols, @@ -911,6 +912,7 @@ enum class TuneKeyKind : int { kMxfp4Dense = 6, kNvfp4Dense = 7, kMxfp4Moe = 8, + kNvfp4Moe = 9, }; struct DenseTuneKey { @@ -1046,6 +1048,9 @@ turbomind::gemm::DispatchPolicy select_fp8_moe_dispatch_policy( turbomind::gemm::DispatchPolicy select_mxfp4_moe_dispatch_policy( int device, int total_tokens, int n, int k, int num_experts, int group_size, cudaStream_t stream); +turbomind::gemm::DispatchPolicy select_nvfp4_moe_dispatch_policy( + int device, int total_tokens, int n, int k, int num_experts, int group_size, + cudaStream_t stream); bool tune_small_shapes_enabled() { const char* raw = std::getenv("VLLM_SM70_AWQ_TUNE_SMALL_SHAPES"); @@ -1069,7 +1074,12 @@ 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() { + const char* raw = std::getenv("VLLM_SM70_MXFP4_MOE_BROADCAST_INPUT_DECODE"); + return raw == nullptr || std::atoi(raw) != 0; } bool mxfp4_moe_grouped_m8_enabled() { @@ -1077,6 +1087,16 @@ bool mxfp4_moe_grouped_m8_enabled() { return raw != nullptr && std::atoi(raw) != 0; } +bool mxfp4_moe_grouped_verifier_enabled() { + const char* raw = std::getenv("VLLM_SM70_MXFP4_MOE_GROUPED_VERIFIER"); + return raw != nullptr && std::atoi(raw) != 0; +} + +bool mxfp4_moe_grouped_m8_expert_rows_enabled() { + const char* raw = std::getenv("VLLM_SM70_MXFP4_MOE_GROUPED_M8_EXPERT_ROWS"); + return raw != nullptr && std::atoi(raw) != 0; +} + bool mxfp4_moe_grouped_m8_fast_selector_enabled() { const char* raw = std::getenv("VLLM_SM70_MXFP4_MOE_GROUPED_M8_FAST_SELECTOR"); return raw == nullptr || std::atoi(raw) != 0; @@ -1087,6 +1107,16 @@ bool nvfp4_tune_small_shapes_enabled() { return raw == nullptr || std::atoi(raw) != 0; } +bool nvfp4_moe_grouped_prefill_enabled() { + const char* raw = std::getenv("VLLM_SM70_NVFP4_MOE_GROUPED_PREFILL"); + return raw == nullptr || std::atoi(raw) != 0; +} + +bool nvfp4_qwen38_tp4_m1_fast_selector_enabled() { + const char* raw = std::getenv("VLLM_SM70_NVFP4_QWEN38_TP4_M1_FAST_SELECTOR"); + return raw == nullptr || std::atoi(raw) != 0; +} + bool fp8_moe_single_token_per_expert_dispatch_enabled() { const char* raw = std::getenv("VLLM_SM70_FP8_MOE_SINGLE_TOKEN_PER_EXPERT_DISPATCH"); @@ -1103,6 +1133,11 @@ bool fp8_safe_fast_selector_enabled() { return raw != nullptr && std::atoi(raw) != 0; } +bool fp8_grouped_bmm_decode_enabled() { + const char* raw = std::getenv("VLLM_SM70_FP8_GROUPED_BMM_DECODE"); + return raw == nullptr || std::atoi(raw) != 0; +} + bool awq_reuse_imported_cache_enabled() { const char* raw = std::getenv("VLLM_SM70_AWQ_REUSE_IMPORTED_CACHE"); return raw != nullptr && std::atoi(raw) != 0; @@ -1214,6 +1249,11 @@ int moe_tune_max_tokens() { return raw ? std::max(std::atoi(raw), 0) : 128; } +int nvfp4_moe_tune_max_tokens() { + const char* raw = std::getenv("VLLM_SM70_NVFP4_MOE_TUNE_MAX_TOKENS"); + return raw ? std::max(std::atoi(raw), 0) : 128; +} + int sm70_f16_dense_max_m() { const char* raw = std::getenv("VLLM_SM70_F16_DENSE_MAX_M"); return raw ? std::max(std::atoi(raw), 0) : 64; @@ -1266,6 +1306,19 @@ turbomind::gemm::DispatchPolicy select_dense_dispatch_policy_impl( turbomind::gemm::DispatchPolicy select_dense_dispatch_policy( int device, int m, int n, int k, int group_size, cudaStream_t stream) { + const char* dflash2_rerank = std::getenv("VLLM_SM70_DFLASH2_QPN8_RERANK"); + const char* dflash2_shadow = + std::getenv("VLLM_SM70_DFLASH2_QPN8_RERANK_SHADOW"); + const bool exact_dflash2_rerank = + (dflash2_rerank && std::atoi(dflash2_rerank) != 0) || + (dflash2_shadow && std::atoi(dflash2_shadow) != 0); + if (exact_dflash2_rerank && m >= 1 && m <= 8 && n == 62080 && k == 5120 && + group_size == 0) { + // The sparse reranker reproduces this exact split-K contract. Do not let + // concurrent startup noise choose a numerically different LM-head spec on + // one TP rank. + return turbomind::gemm::DispatchPolicy::kDefault; + } return select_dense_dispatch_policy_impl( device, m, n, k, group_size, stream, TuneKeyKind::kGenericDense, tune_small_shapes_enabled(), false, generic_dense_tune_max_m()); @@ -1304,6 +1357,11 @@ turbomind::gemm::DispatchPolicy select_mxfp4_dense_dispatch_policy( turbomind::gemm::DispatchPolicy select_nvfp4_dense_dispatch_policy( int device, int m, int n, int k, int group_size, cudaStream_t stream) { + if (nvfp4_qwen38_tp4_m1_fast_selector_enabled() && m == 1 && + group_size == 16 && + ((n == 8704 && k == 5120) || (n == 5120 && k == 4352))) { + return turbomind::gemm::DispatchPolicy::kDefault; + } return select_dense_dispatch_policy_impl( device, m, n, k, group_size, stream, TuneKeyKind::kNvfp4Dense, nvfp4_tune_small_shapes_enabled(), false, nvfp4_dense_tune_max_m()); @@ -1311,8 +1369,11 @@ turbomind::gemm::DispatchPolicy select_nvfp4_dense_dispatch_policy( turbomind::gemm::DispatchPolicy select_moe_dispatch_policy_impl( int device, int total_tokens, int n, int k, int num_experts, int group_size, - cudaStream_t stream, TuneKeyKind kind, bool tune_enabled) { - if (!tune_enabled || total_tokens > moe_tune_max_tokens()) { + cudaStream_t stream, TuneKeyKind kind, bool tune_enabled, + int max_tune_tokens = -1) { + const int tune_limit = + max_tune_tokens >= 0 ? max_tune_tokens : moe_tune_max_tokens(); + if (!tune_enabled || total_tokens > tune_limit) { return turbomind::gemm::DispatchPolicy::kDefault; } @@ -1342,6 +1403,12 @@ turbomind::gemm::DispatchPolicy select_moe_dispatch_policy( turbomind::gemm::DispatchPolicy select_fp8_moe_dispatch_policy( int device, int total_tokens, int n, int k, int num_experts, int group_size, cudaStream_t stream) { + if (fp8_grouped_bmm_decode_enabled() && total_tokens == 2 && n == 1024 && + k == 4096 && num_experts == 2 && group_size == 128) { + // The matching dense WO-A projection uses the fixed launch spec selected + // in gemm.cu. Do not let measurement replace its accumulation tree. + return turbomind::gemm::DispatchPolicy::kDefault; + } return select_moe_dispatch_policy_impl( device, total_tokens, n, k, num_experts, group_size, stream, TuneKeyKind::kFp8Moe, fp8_tune_small_shapes_enabled()); @@ -1354,6 +1421,7 @@ turbomind::gemm::DispatchPolicy select_mxfp4_moe_dispatch_policy( total_tokens == 48 && num_experts == 48 && group_size == 32 && ((n == 512 && k == 4096) || (n == 4096 && k == 256)); if (mxfp4_moe_grouped_m8_enabled() && + !mxfp4_moe_grouped_m8_expert_rows_enabled() && mxfp4_moe_grouped_m8_fast_selector_enabled() && exact_grouped_m8) { return turbomind::gemm::DispatchPolicy::kMxfp4MoeGroupedM8Fast; } @@ -1362,6 +1430,15 @@ turbomind::gemm::DispatchPolicy select_mxfp4_moe_dispatch_policy( TuneKeyKind::kMxfp4Moe, mxfp4_tune_small_shapes_enabled()); } +turbomind::gemm::DispatchPolicy select_nvfp4_moe_dispatch_policy( + int device, int total_tokens, int n, int k, int num_experts, int group_size, + cudaStream_t stream) { + return select_moe_dispatch_policy_impl( + device, total_tokens, n, k, num_experts, group_size, stream, + TuneKeyKind::kNvfp4Moe, nvfp4_tune_small_shapes_enabled(), + nvfp4_moe_tune_max_tokens()); +} + WorkspaceHolder& get_workspace(int device, cudaStream_t stream) { thread_local int cached_device = -1; thread_local cudaStream_t cached_stream = nullptr; @@ -2384,8 +2461,7 @@ __global__ void fp8_sm70_dequantize_kernel( } const int words_per_n_tile = k / kFp8QuantValuesPerWord; - const int n_in_tile = - static_cast(word_index % kFp8OutputColumnsPerTile); + const int n_in_tile = static_cast(word_index % kFp8OutputColumnsPerTile); const int64_t tile_word = word_index / kFp8OutputColumnsPerTile; const int k_word = static_cast(tile_word % words_per_n_tile); const int n_tile = static_cast(tile_word / words_per_n_tile); @@ -2394,10 +2470,12 @@ __global__ void fp8_sm70_dequantize_kernel( const int group = k_base / group_size; const uint2 quant = reinterpret_cast(packed_weight)[word_index]; - const auto& quant_lo = reinterpret_cast< - const turbomind::Array&>(quant.x); - const auto& quant_hi = reinterpret_cast< - const turbomind::Array&>(quant.y); + const auto& quant_lo = + reinterpret_cast&>( + quant.x); + const auto& quant_hi = + reinterpret_cast&>( + quant.y); const auto half_lo = turbomind::cvt_f16x4_e4m3(quant_lo); const auto half_hi = turbomind::cvt_f16x4_e4m3(quant_hi); const __half scale = packed_scales[static_cast(group) * n + col]; @@ -2422,10 +2500,8 @@ __global__ void fp8_sm70_dequantize_kernel( __hmul(half_hi[3], scale); } -void fp8_sm70_dequantize_out(torch::Tensor output, - torch::Tensor packed_weight, - torch::Tensor packed_scales, - int64_t group_size) { +void fp8_sm70_dequantize_out(torch::Tensor output, torch::Tensor packed_weight, + torch::Tensor packed_scales, int64_t group_size) { TORCH_CHECK( output.is_cuda() && packed_weight.is_cuda() && packed_scales.is_cuda(), "SM70 FP8 prefill dequant expects CUDA tensors."); @@ -2445,8 +2521,7 @@ void fp8_sm70_dequantize_out(torch::Tensor output, const at::cuda::OptionalCUDAGuard device_guard(device_of(output)); const int64_t k = output.size(0); const int64_t n = output.size(1); - TORCH_CHECK(k % group_size == 0 && - k % kFp8QuantValuesPerWord == 0 && + TORCH_CHECK(k % group_size == 0 && k % kFp8QuantValuesPerWord == 0 && n % kFp8OutputColumnsPerTile == 0, "SM70 FP8 prefill dequant shape alignment mismatch."); TORCH_CHECK(packed_weight.numel() == k * n, @@ -2459,15 +2534,14 @@ void fp8_sm70_dequantize_out(torch::Tensor output, "SM70 FP8 prefill tensors must share a device."); const int64_t words = packed_weight.numel() / kFp8QuantValuesPerWord; - const int blocks = static_cast( - (words + kFp8PrefillDequantThreads - 1) / kFp8PrefillDequantThreads); + const int blocks = static_cast((words + kFp8PrefillDequantThreads - 1) / + kFp8PrefillDequantThreads); fp8_sm70_dequantize_kernel<<>>( reinterpret_cast<__half*>(output.data_ptr()), packed_weight.data_ptr(), reinterpret_cast(packed_scales.data_ptr()), - static_cast(k), static_cast(n), - static_cast(group_size)); + static_cast(k), static_cast(n), static_cast(group_size)); C10_CUDA_KERNEL_LAUNCH_CHECK(); } @@ -2802,10 +2876,13 @@ std::vector fp8_sm70_prepare(torch::Tensor qweight, const int64_t k = qweight.size(1); TORCH_CHECK(k % 8 == 0 && n % 8 == 0, "fp8_sm70_prepare: K and N must be multiples of 8."); - TORCH_CHECK(scales.size(0) == (n + group_size - 1) / group_size, - "fp8_sm70_prepare: output scale block mismatch."); - TORCH_CHECK(scales.size(1) == (k + group_size - 1) / group_size, - "fp8_sm70_prepare: input scale block mismatch."); + const bool channelwise_scales = scales.size(0) == n && scales.size(1) == 1; + const bool blockwise_scales = + scales.size(0) == (n + group_size - 1) / group_size && + scales.size(1) == (k + group_size - 1) / group_size; + TORCH_CHECK(channelwise_scales || blockwise_scales, + "fp8_sm70_prepare: expected [N, 1] channel scales or " + "[ceil(N/128), ceil(K/128)] block scales."); const auto converters = turbomind::gemm::GetConverters( turbomind::kHalf, turbomind::kFloat8_e4m3, turbomind::kHalf, true, 70); @@ -2869,12 +2946,21 @@ std::vector fp8_sm70_prepare(torch::Tensor qweight, const bool is_B_s = !is_A_s; const int64_t num_groups = (k + group_size - 1) / group_size; - auto group_scales = scales.transpose(0, 1) - .contiguous() - .to(torch::kFloat16) - .repeat_interleave(group_size, 1) - .slice(1, 0, n) - .contiguous(); + torch::Tensor group_scales; + if (channelwise_scales) { + group_scales = scales.transpose(0, 1) + .contiguous() + .to(torch::kFloat16) + .repeat({num_groups, 1}) + .contiguous(); + } else { + group_scales = scales.transpose(0, 1) + .contiguous() + .to(torch::kFloat16) + .repeat_interleave(group_size, 1) + .slice(1, 0, n) + .contiguous(); + } if (interleave_gated_silu) { group_scales = interleave_gated_silu_cols(group_scales); } @@ -3391,7 +3477,8 @@ void awq_gemm_sm70_out_tile_reduce( void fp8_gemm_sm70_out(torch::Tensor out, torch::Tensor in_feats, torch::Tensor tm_weight, torch::Tensor tm_scales, int64_t group_size, int64_t k_ld, int64_t q_ld, - bool gated_silu) { + bool gated_silu, + bool exact_8k_prefill_prescaled = false) { TORCH_CHECK(in_feats.is_cuda(), "fp8_gemm_sm70: input must be CUDA."); TORCH_CHECK(tm_weight.is_cuda(), "fp8_gemm_sm70: weight must be CUDA."); TORCH_CHECK(tm_scales.is_cuda(), "fp8_gemm_sm70: scales must be CUDA."); @@ -3422,6 +3509,19 @@ void fp8_gemm_sm70_out(torch::Tensor out, torch::Tensor in_feats, "fp8_gemm_sm70: output rows must match input rows."); TORCH_CHECK(out.stride(1) == 1, "fp8_gemm_sm70: output must be row-major contiguous."); + if (exact_8k_prefill_prescaled) { + 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, "fp8_gemm_sm70: gated_silu requires even output dim."); @@ -3502,6 +3602,10 @@ void fp8_gemm_sm70_out(torch::Tensor out, torch::Tensor in_feats, op.dispatch = select_fp8_dense_dispatch_policy( device, static_cast(m), static_cast(n), static_cast(k), static_cast(group_size), stream); + if (exact_8k_prefill_prescaled) { + op.dispatch = + op.dispatch | turbomind::gemm::DispatchPolicy::kSm70Fp8PrefillPrescaled; + } op.epilogue = gated_silu ? turbomind::gemm::Epilogue::kGatedSilu : turbomind::gemm::Epilogue::kNone; op.quant_a = {turbomind::gemm::QuantType::kNone, 0}; @@ -3518,6 +3622,13 @@ void fp8_gemm_sm70_out(torch::Tensor out, torch::Tensor in_feats, TORCH_CHECK(ec == 0, "fp8_gemm_sm70: TurboMind GEMM failed."); } +bool sm70_fp8_prefill_cutlass_gated_silu_enabled( + const torch::Tensor& in_feats, const torch::Tensor& dense_weight) { + const char* raw = std::getenv("VLLM_SM70_FP8_PREFILL_CUTLASS"); + return (raw == nullptr || std::atoi(raw) != 0) && in_feats.size(0) == 8000 && + in_feats.size(1) == 5120 && dense_weight.size(1) == 8704; +} + void fp8_gemm_sm70_prefill_dispatch_out( torch::Tensor out, int64_t dense_weight_ptr, torch::Tensor in_feats, torch::Tensor tm_weight, torch::Tensor tm_scales, int64_t group_size, @@ -3534,6 +3645,19 @@ void fp8_gemm_sm70_prefill_dispatch_out( reinterpret_cast(dense_weight_ptr), {in_feats.size(1), tm_weight.size(1)}, in_feats.options()); fp8_sm70_dequantize_out(dense_weight, tm_weight, tm_scales, group_size); + if (gated_silu && + sm70_fp8_prefill_cutlass_gated_silu_enabled(in_feats, dense_weight)) { + auto gate_up = + at::empty({in_feats.size(0), dense_weight.size(1)}, in_feats.options()); + if (sm70_fp8_prefill_cutlass_out(gate_up, in_feats, dense_weight, false)) { + sm70_silu_and_mul_interleaved_fp16_out(out, gate_up); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return; + } + } + if (sm70_fp8_prefill_cutlass_out(out, in_feats, dense_weight, gated_silu)) { + return; + } if (gated_silu) { auto gate_up = at::mm(in_feats, dense_weight); sm70_silu_and_mul_interleaved_fp16_out(out, gate_up); @@ -4182,6 +4306,574 @@ void sm70_f16_gemm_out(torch::Tensor out, torch::Tensor in_feats, TORCH_CHECK(ec == 0, "sm70_f16_gemm: TurboMind GEMM failed."); } +template +__global__ void sm70_f16_gather_candidate_rows_kernel( + half* selected, const half* weight, const int64_t* candidate_ids, + int selected_rows, int hidden_size, int64_t weight_stride) { + const int row = blockIdx.x; + if (row >= selected_rows) { + return; + } + const int64_t token = __ldg(candidate_ids + row); + const half* src = weight + token * weight_stride; + half* dst = selected + static_cast(row) * hidden_size; + constexpr int kVectorHalf = sizeof(uint4) / sizeof(half); + const int vectors = hidden_size / kVectorHalf; + for (int idx = threadIdx.x; idx < vectors; idx += Threads) { + reinterpret_cast(dst)[idx] = + __ldg(reinterpret_cast(src) + idx); + } +} + +template +__global__ void sm70_f16_gather_packed_candidate_rows_kernel( + half* selected_packed, const half* packed_weight, + const int64_t* candidate_ids, int selected_rows, int hidden_size) { + // The SM70 HMMA 8x8x4 B layout packs every logical [32, 8] tile as + // eight consecutive uint4 vectors. Copy those vectors directly from the + // already prepared full LM-head layout instead of gathering 5.2 MiB of raw + // FP16 rows and running the generic layout converter every decode step. + constexpr int kRowsPerPack = 32; + constexpr int kColsPerVector = sizeof(uint4) / sizeof(half); + constexpr int kPackedTileElements = kRowsPerPack * kColsPerVector; + const int k_tiles = hidden_size / kColsPerVector; + const int vector_count = selected_rows * k_tiles; + for (int index = blockIdx.x * Threads + threadIdx.x; index < vector_count; + index += gridDim.x * Threads) { + const int selected_row = index / k_tiles; + const int k_tile = index - selected_row * k_tiles; + const int64_t token = __ldg(candidate_ids + selected_row); + const int64_t source = (token / kRowsPerPack) * hidden_size * kRowsPerPack + + static_cast(k_tile) * kPackedTileElements + + ((token % kRowsPerPack) / 4) * 4 * kColsPerVector + + (token % 4) * kColsPerVector; + const int64_t destination = + static_cast(selected_row / kRowsPerPack) * hidden_size * + kRowsPerPack + + static_cast(k_tile) * kPackedTileElements + + ((selected_row % kRowsPerPack) / 4) * 4 * kColsPerVector + + (selected_row % 4) * kColsPerVector; + *reinterpret_cast(selected_packed + destination) = + __ldg(reinterpret_cast(packed_weight + source)); + } +} + +template +__global__ void sm70_f16_extract_candidate_rows_kernel( + half* out, const half* expanded, int rows, int candidates, + int64_t expanded_stride) { + const int idx = blockIdx.x * Threads + threadIdx.x; + if (idx >= rows * candidates) { + return; + } + const int row = idx / candidates; + const int col = idx - row * candidates; + out[idx] = expanded[static_cast(row) * expanded_stride + + row * candidates + col]; +} + +template +__global__ void sm70_f16_rerank_keys_kernel(int64_t* keys, const half* logits, + const int64_t* candidate_ids, + int count) { + for (int index = blockIdx.x * Threads + threadIdx.x; index < count; + index += gridDim.x * Threads) { + const uint16_t raw = + __ldg(reinterpret_cast(logits) + index); + const uint16_t ordered = raw & 0x8000U + ? static_cast(~raw) + : static_cast(raw ^ 0x8000U); + // int64 top-k over this key is exactly lexicographic by FP16 value + // descending, then original vocabulary ID ascending. Candidate IDs are + // unique and below 2^32, so the key is collision-free. + const uint64_t tie = + 0xffffffffULL - static_cast(__ldg(candidate_ids + index)); + keys[index] = + static_cast((static_cast(ordered) << 32) | tie); + } +} + +__global__ void sm70_f16_rerank_topk_kernel(half* values_out, int64_t* ids_out, + const half* logits, + const int64_t* candidate_ids, + int rows, int top_k, + int64_t vocab_start_index) { + constexpr int kCandidates = 64; + const int row = blockIdx.x; + if (row >= rows) { + return; + } + + // One warp owns one candidate row. The composite key is identical to + // sm70_f16_rerank_keys_kernel: FP16 value descending, then original local + // vocabulary ID ascending. Keeping all 64 keys in shared memory lets the + // warp emit the first 16/20 order statistics without launching the generic + // top-k, gather, and ID-offset kernels. + __shared__ uint64_t candidate_keys[kCandidates]; + const int lane = threadIdx.x; + const int base = row * kCandidates; +#pragma unroll + for (int offset = 0; offset < kCandidates; offset += WARP_SIZE) { + const int position = lane + offset; + const uint16_t raw = + __ldg(reinterpret_cast(logits) + base + position); + const uint16_t ordered = raw & 0x8000U + ? static_cast(~raw) + : static_cast(raw ^ 0x8000U); + const uint64_t tie = + 0xffffffffULL - + static_cast(__ldg(candidate_ids + base + position)); + candidate_keys[position] = (static_cast(ordered) << 32) | tie; + } + __syncwarp(); + + for (int rank = 0; rank < top_k; ++rank) { + uint64_t best_key = candidate_keys[lane]; + int best_position = lane; + const uint64_t second_key = candidate_keys[lane + WARP_SIZE]; + if (second_key > best_key) { + best_key = second_key; + best_position = lane + WARP_SIZE; + } +#pragma unroll + for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) { + const uint64_t other_key = + __shfl_down_sync(0xffffffffU, best_key, offset); + const int other_position = + __shfl_down_sync(0xffffffffU, best_position, offset); + if (other_key > best_key) { + best_key = other_key; + best_position = other_position; + } + } + if (lane == 0) { + const uint16_t ordered = static_cast(best_key >> 32); + const uint16_t raw = ordered & 0x8000U + ? static_cast(ordered ^ 0x8000U) + : static_cast(~ordered); + values_out[row * top_k + rank] = __ushort_as_half(raw); + const uint32_t local_id = 0xffffffffU - static_cast(best_key); + ids_out[row * top_k + rank] = + vocab_start_index + static_cast(local_id); + // A valid local vocabulary ID is far below UINT32_MAX, so its key is + // never zero. Zero is therefore an exact in-warp consumed sentinel. + candidate_keys[best_position] = 0; + } + __syncwarp(); + } +} + +template +void launch_sm70_f16_indexed_rerank_gemm( + torch::Tensor out, torch::Tensor in_feats, torch::Tensor selected_packed, + torch::Tensor expanded, torch::Tensor partials, torch::Tensor barriers, + turbomind::gemm::MatrixLayout desc_B, int split_k, cudaStream_t stream) { + constexpr int kCandidates = 64; + constexpr int kThreads = 256; + const int m = static_cast(in_feats.size(0)); + const int k = static_cast(in_feats.size(1)); + const int selected_rows = m * kCandidates; + + turbomind::gemm::MatrixLayout desc_A{ + turbomind::kHalf, + turbomind::gemm::kRowMajor, + m, + k, + static_cast(in_feats.stride(0)), + }; + turbomind::gemm::MatrixLayout desc_U{}; + turbomind::gemm::MatrixLayout desc_V{}; + turbomind::gemm::MatrixLayout desc_D{ + turbomind::kHalf, turbomind::gemm::kRowMajor, m, + selected_rows, static_cast(expanded.stride(0)), + }; + + using namespace turbomind::gemm; + using C = sm70_s884::Config_F16; + using D = turbomind::cache_policy::Default; + using S = turbomind::cache_policy::Stream; + using Gemm = typename C::template Type<8, CTA_N, 64, 1, 4, 1, D, S, 2, true, + 1, 1>::Kernel; + using Sched = typename Gemm::Scheduler; + GemmParam param{ + to_param(in_feats.data_ptr(), desc_A), + to_param(selected_packed.data_ptr(), transpose(desc_B)), + to_param(nullptr, desc_U), + to_param(nullptr, desc_V), + }; + EpilogueParam epi_param{}; + epi_param.c = to_param(expanded.data_ptr(), desc_D); + turbomind::gemm::MatrixLayout desc_P = desc_D; + desc_P.ld = static_cast(partials.stride(0)); + epi_param.partials = to_param(partials.data_ptr(), desc_P); + epi_param.locks = barriers.data_ptr(); + epi_param.combine_mat = + MatrixCombination_v3{to_param(nullptr, MatrixLayout{}), 1.f, 0.f}; + Sched sched{{m, selected_rows, k, 1}, 0, split_k}; + sched.offsets_ = nullptr; + + const auto grid = sched.get_grid_shape(); + const auto block = Gemm::Impl::WARPS * WARP_SIZE; + constexpr int dynamic_smem_size = + static_cast(sizeof(typename Gemm::SharedStorage)); + auto kernel = gemm_kernel; + if constexpr (dynamic_smem_size > (48 << 10)) { + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + dynamic_smem_size); + } + kernel<<>>(param, epi_param, sched); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + const int output_count = m * kCandidates; + sm70_f16_extract_candidate_rows_kernel + <<<(output_count + kThreads - 1) / kThreads, kThreads, 0, stream>>>( + reinterpret_cast(out.data_ptr()), + reinterpret_cast(expanded.data_ptr()), m, + kCandidates, expanded.stride(0)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +template +void launch_sm70_f16_indexed_rerank( + torch::Tensor out, torch::Tensor in_feats, torch::Tensor weight, + torch::Tensor candidate_ids, torch::Tensor selected_raw, + torch::Tensor selected_packed, torch::Tensor expanded, + torch::Tensor partials, torch::Tensor barriers, int split_k, + cudaStream_t stream) { + constexpr int kCandidates = 64; + constexpr int kThreads = 256; + const int m = static_cast(in_feats.size(0)); + const int k = static_cast(in_feats.size(1)); + const int selected_rows = m * kCandidates; + + sm70_f16_gather_candidate_rows_kernel + <<>>( + reinterpret_cast(selected_raw.data_ptr()), + reinterpret_cast(weight.data_ptr()), + candidate_ids.data_ptr(), selected_rows, k, + weight.stride(0)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + const auto converters = turbomind::gemm::GetConverters( + turbomind::kHalf, turbomind::kHalf, turbomind::kHalf, true, 70); + const auto* conv_w = converters[0]; + TORCH_CHECK(conv_w, + "sm70_f16_indexed_rerank_out: no compatible TurboMind " + "converter."); + const auto order_w = conv_w->order; + const bool is_A_w = turbomind::gemm::get_operand_tag(conv_w->pack) == + turbomind::gemm::OPERAND_A; + const bool is_B_w = !is_A_w; + turbomind::gemm::MatrixLayout w_desc{ + turbomind::kHalf, + order_w, + selected_rows, + k, + order_w == turbomind::gemm::kRowMajor ? k : selected_rows, + }; + if (is_B_w) { + std::swap(w_desc.rows, w_desc.cols); + w_desc.order = ~w_desc.order; + } + turbomind::gemm::MatrixLayout desc_B = w_desc; + desc_B.pack = conv_w->pack; + if (is_A_w) { + desc_B = turbomind::gemm::transpose(desc_B); + } + TORCH_CHECK(conv_w->Convert(selected_raw.data_ptr(), w_desc, + selected_packed.data_ptr(), desc_B, stream) == 0, + "sm70_f16_indexed_rerank_out: candidate pack failed."); + + launch_sm70_f16_indexed_rerank_gemm(out, in_feats, selected_packed, + expanded, partials, barriers, + desc_B, split_k, stream); +} + +template +void launch_sm70_f16_indexed_rerank_packed( + torch::Tensor out, torch::Tensor in_feats, torch::Tensor packed_weight, + torch::Tensor candidate_ids, torch::Tensor selected_packed, + torch::Tensor expanded, torch::Tensor partials, torch::Tensor barriers, + int split_k, cudaStream_t stream) { + constexpr int kCandidates = 64; + constexpr int kThreads = 256; + const int m = static_cast(in_feats.size(0)); + const int k = static_cast(in_feats.size(1)); + const int selected_rows = m * kCandidates; + const int vector_count = selected_rows * k / 8; + const int blocks = std::min(1024, (vector_count + kThreads - 1) / kThreads); + sm70_f16_gather_packed_candidate_rows_kernel + <<>>( + reinterpret_cast(selected_packed.data_ptr()), + reinterpret_cast(packed_weight.data_ptr()), + candidate_ids.data_ptr(), selected_rows, k); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + turbomind::gemm::MatrixLayout desc_B{ + turbomind::kHalf, turbomind::gemm::kColMajor, k, selected_rows, k * 32, + }; + desc_B.pack = static_cast( + turbomind::gemm::HMMA_884 | turbomind::gemm::OPERAND_B | 1); + launch_sm70_f16_indexed_rerank_gemm(out, in_feats, selected_packed, + expanded, partials, barriers, + desc_B, split_k, stream); +} + +void sm70_f16_indexed_rerank_out(torch::Tensor out, torch::Tensor in_feats, + torch::Tensor weight, + torch::Tensor candidate_ids, + torch::Tensor selected_raw, + torch::Tensor selected_packed, + torch::Tensor expanded, torch::Tensor partials, + torch::Tensor barriers, int64_t cta_n, + int64_t split_k) { + constexpr int64_t kMaxRows = 8; + constexpr int64_t kCandidates = 64; + TORCH_CHECK(out.is_cuda() && in_feats.is_cuda() && weight.is_cuda() && + candidate_ids.is_cuda() && selected_raw.is_cuda() && + selected_packed.is_cuda() && expanded.is_cuda() && + partials.is_cuda() && barriers.is_cuda(), + "sm70_f16_indexed_rerank_out: tensors must be CUDA."); + TORCH_CHECK(out.scalar_type() == torch::kFloat16 && + in_feats.scalar_type() == torch::kFloat16 && + weight.scalar_type() == torch::kFloat16 && + selected_raw.scalar_type() == torch::kFloat16 && + selected_packed.scalar_type() == torch::kFloat16 && + expanded.scalar_type() == torch::kFloat16 && + partials.scalar_type() == torch::kFloat32, + "sm70_f16_indexed_rerank_out: data tensors must be float16."); + TORCH_CHECK(candidate_ids.scalar_type() == torch::kInt64, + "sm70_f16_indexed_rerank_out: candidate IDs must be int64."); + TORCH_CHECK(barriers.scalar_type() == torch::kInt32, + "sm70_f16_indexed_rerank_out: barriers must be int32."); + TORCH_CHECK(in_feats.dim() == 2 && weight.dim() == 2 && out.dim() == 2 && + candidate_ids.dim() == 2 && selected_raw.dim() == 2 && + selected_packed.dim() == 2 && expanded.dim() == 2 && + partials.dim() == 2, + "sm70_f16_indexed_rerank_out: tensors must be 2D."); + const int64_t m = in_feats.size(0); + const int64_t k = in_feats.size(1); + const int64_t selected_rows = m * kCandidates; + TORCH_CHECK(m >= 1 && m <= kMaxRows, + "sm70_f16_indexed_rerank_out: M must be in [1, 8]."); + TORCH_CHECK(k == weight.size(1) && k % 64 == 0, + "sm70_f16_indexed_rerank_out: invalid K dimension."); + TORCH_CHECK(out.size(0) == m && out.size(1) == kCandidates && + candidate_ids.size(0) >= m && + candidate_ids.size(1) == kCandidates, + "sm70_f16_indexed_rerank_out: expected [M, 64] output and " + "candidate IDs."); + TORCH_CHECK(selected_raw.size(0) >= selected_rows && + selected_raw.size(1) == k && + selected_packed.sizes() == selected_raw.sizes(), + "sm70_f16_indexed_rerank_out: candidate workspaces must be " + "[>=M*64, K]."); + TORCH_CHECK(expanded.size(0) >= m && expanded.size(1) >= selected_rows, + "sm70_f16_indexed_rerank_out: expanded workspace is too " + "small."); + TORCH_CHECK(partials.size(0) >= m && partials.size(1) >= selected_rows && + barriers.numel() >= 64, + "sm70_f16_indexed_rerank_out: split-K workspaces are too " + "small."); + TORCH_CHECK(out.is_contiguous() && in_feats.is_contiguous() && + weight.is_contiguous() && candidate_ids.is_contiguous() && + selected_raw.is_contiguous() && + selected_packed.is_contiguous() && expanded.is_contiguous() && + partials.is_contiguous() && barriers.is_contiguous(), + "sm70_f16_indexed_rerank_out: tensors must be contiguous."); + const int device = in_feats.get_device(); + TORCH_CHECK(out.get_device() == device && weight.get_device() == device && + candidate_ids.get_device() == device && + selected_raw.get_device() == device && + selected_packed.get_device() == device && + expanded.get_device() == device && + partials.get_device() == device && + barriers.get_device() == device, + "sm70_f16_indexed_rerank_out: tensors must share one device."); + TORCH_CHECK(cta_n == 128 || cta_n == 256, + "sm70_f16_indexed_rerank_out: cta_n must be 128 or 256."); + TORCH_CHECK(split_k >= 1 && split_k <= 16, + "sm70_f16_indexed_rerank_out: split_k must be in [1, 16]."); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(in_feats)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + if (cta_n == 128) { + launch_sm70_f16_indexed_rerank<128>( + out, in_feats, weight, candidate_ids, selected_raw, selected_packed, + expanded, partials, barriers, static_cast(split_k), stream); + } else { + launch_sm70_f16_indexed_rerank<256>( + out, in_feats, weight, candidate_ids, selected_raw, selected_packed, + expanded, partials, barriers, static_cast(split_k), stream); + } +} + +void sm70_f16_indexed_rerank_packed_out( + torch::Tensor out, torch::Tensor in_feats, torch::Tensor packed_weight, + torch::Tensor candidate_ids, torch::Tensor selected_packed, + torch::Tensor expanded, torch::Tensor partials, torch::Tensor barriers, + int64_t cta_n, int64_t split_k) { + constexpr int64_t kMaxRows = 8; + constexpr int64_t kCandidates = 64; + TORCH_CHECK(out.is_cuda() && in_feats.is_cuda() && packed_weight.is_cuda() && + candidate_ids.is_cuda() && selected_packed.is_cuda() && + expanded.is_cuda() && partials.is_cuda() && + barriers.is_cuda(), + "sm70_f16_indexed_rerank_packed_out: tensors must be CUDA."); + TORCH_CHECK(out.scalar_type() == torch::kFloat16 && + in_feats.scalar_type() == torch::kFloat16 && + packed_weight.scalar_type() == torch::kFloat16 && + selected_packed.scalar_type() == torch::kFloat16 && + expanded.scalar_type() == torch::kFloat16 && + partials.scalar_type() == torch::kFloat32, + "sm70_f16_indexed_rerank_packed_out: data tensors must have " + "the expected FP16/FP32 dtypes."); + TORCH_CHECK(candidate_ids.scalar_type() == torch::kInt64 && + barriers.scalar_type() == torch::kInt32, + "sm70_f16_indexed_rerank_packed_out: IDs/barriers must be " + "int64/int32."); + TORCH_CHECK(in_feats.dim() == 2 && packed_weight.dim() == 2 && + out.dim() == 2 && candidate_ids.dim() == 2 && + selected_packed.dim() == 2 && expanded.dim() == 2 && + partials.dim() == 2, + "sm70_f16_indexed_rerank_packed_out: tensors must be 2D."); + const int64_t m = in_feats.size(0); + const int64_t k = in_feats.size(1); + const int64_t selected_rows = m * kCandidates; + TORCH_CHECK(m >= 1 && m <= kMaxRows, + "sm70_f16_indexed_rerank_packed_out: M must be in [1, 8]."); + TORCH_CHECK(k == packed_weight.size(1) && k % 64 == 0 && + packed_weight.size(0) % 32 == 0, + "sm70_f16_indexed_rerank_packed_out: invalid packed weight " + "shape."); + TORCH_CHECK(out.size(0) == m && out.size(1) == kCandidates && + candidate_ids.size(0) >= m && + candidate_ids.size(1) == kCandidates, + "sm70_f16_indexed_rerank_packed_out: expected [M, 64] output " + "and candidate IDs."); + TORCH_CHECK( + selected_packed.size(0) >= selected_rows && selected_packed.size(1) == k, + "sm70_f16_indexed_rerank_packed_out: packed workspace must be " + "[>=M*64, K]."); + TORCH_CHECK(expanded.size(0) >= m && expanded.size(1) >= selected_rows && + partials.size(0) >= m && partials.size(1) >= selected_rows && + barriers.numel() >= 64, + "sm70_f16_indexed_rerank_packed_out: split-K workspaces are " + "too small."); + TORCH_CHECK(out.is_contiguous() && in_feats.is_contiguous() && + packed_weight.is_contiguous() && + candidate_ids.is_contiguous() && + selected_packed.is_contiguous() && expanded.is_contiguous() && + partials.is_contiguous() && barriers.is_contiguous(), + "sm70_f16_indexed_rerank_packed_out: tensors must be " + "contiguous."); + const int device = in_feats.get_device(); + TORCH_CHECK( + out.get_device() == device && packed_weight.get_device() == device && + candidate_ids.get_device() == device && + selected_packed.get_device() == device && + expanded.get_device() == device && partials.get_device() == device && + barriers.get_device() == device, + "sm70_f16_indexed_rerank_packed_out: tensors must share one " + "device."); + TORCH_CHECK(cta_n == 128 || cta_n == 256, + "sm70_f16_indexed_rerank_packed_out: cta_n must be 128 or " + "256."); + TORCH_CHECK(split_k >= 1 && split_k <= 16, + "sm70_f16_indexed_rerank_packed_out: split_k must be in " + "[1, 16]."); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(in_feats)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + if (cta_n == 128) { + launch_sm70_f16_indexed_rerank_packed<128>( + out, in_feats, packed_weight, candidate_ids, selected_packed, expanded, + partials, barriers, static_cast(split_k), stream); + } else { + launch_sm70_f16_indexed_rerank_packed<256>( + out, in_feats, packed_weight, candidate_ids, selected_packed, expanded, + partials, barriers, static_cast(split_k), stream); + } +} + +void sm70_f16_rerank_keys_out(torch::Tensor keys, torch::Tensor logits, + torch::Tensor candidate_ids) { + constexpr int64_t kCandidates = 64; + constexpr int kThreads = 128; + TORCH_CHECK(keys.is_cuda() && logits.is_cuda() && candidate_ids.is_cuda(), + "sm70_f16_rerank_keys_out: tensors must be CUDA."); + TORCH_CHECK(keys.scalar_type() == torch::kInt64 && + candidate_ids.scalar_type() == torch::kInt64 && + logits.scalar_type() == torch::kFloat16, + "sm70_f16_rerank_keys_out: expected int64/FP16 tensors."); + TORCH_CHECK(keys.sizes() == logits.sizes() && + keys.sizes() == candidate_ids.sizes() && keys.dim() == 2 && + keys.size(1) == kCandidates, + "sm70_f16_rerank_keys_out: expected matching [M, 64] " + "tensors."); + TORCH_CHECK(keys.is_contiguous() && logits.is_contiguous() && + candidate_ids.is_contiguous(), + "sm70_f16_rerank_keys_out: tensors must be contiguous."); + const int device = logits.get_device(); + TORCH_CHECK( + keys.get_device() == device && candidate_ids.get_device() == device, + "sm70_f16_rerank_keys_out: tensors must share one device."); + const at::cuda::OptionalCUDAGuard device_guard(device_of(logits)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const int count = static_cast(keys.numel()); + sm70_f16_rerank_keys_kernel + <<<(count + kThreads - 1) / kThreads, kThreads, 0, stream>>>( + keys.data_ptr(), + reinterpret_cast(logits.data_ptr()), + candidate_ids.data_ptr(), count); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +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) { + constexpr int64_t kCandidates = 64; + constexpr int kThreads = 32; + TORCH_CHECK(values_out.is_cuda() && ids_out.is_cuda() && logits.is_cuda() && + candidate_ids.is_cuda(), + "sm70_f16_rerank_topk_out: tensors must be CUDA."); + TORCH_CHECK(values_out.scalar_type() == torch::kFloat16 && + ids_out.scalar_type() == torch::kInt64 && + logits.scalar_type() == torch::kFloat16 && + candidate_ids.scalar_type() == torch::kInt64, + "sm70_f16_rerank_topk_out: expected FP16/int64 tensors."); + TORCH_CHECK(logits.dim() == 2 && logits.size(1) == kCandidates && + candidate_ids.sizes() == logits.sizes(), + "sm70_f16_rerank_topk_out: expected matching [M, 64] " + "candidate tensors."); + TORCH_CHECK(values_out.dim() == 2 && ids_out.sizes() == values_out.sizes() && + values_out.size(0) == logits.size(0) && + (values_out.size(1) == 16 || values_out.size(1) == 20), + "sm70_f16_rerank_topk_out: outputs must be matching [M, 16] " + "or [M, 20] tensors."); + TORCH_CHECK(values_out.is_contiguous() && ids_out.is_contiguous() && + logits.is_contiguous() && candidate_ids.is_contiguous(), + "sm70_f16_rerank_topk_out: tensors must be contiguous."); + const int device = logits.get_device(); + TORCH_CHECK(values_out.get_device() == device && + ids_out.get_device() == device && + candidate_ids.get_device() == device, + "sm70_f16_rerank_topk_out: tensors must share one device."); + const at::cuda::OptionalCUDAGuard device_guard(device_of(logits)); + const int rows = static_cast(logits.size(0)); + if (rows == 0) { + return; + } + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + sm70_f16_rerank_topk_kernel<<>>( + reinterpret_cast(values_out.data_ptr()), + ids_out.data_ptr(), + reinterpret_cast(logits.data_ptr()), + candidate_ids.data_ptr(), rows, + static_cast(values_out.size(1)), vocab_start_index); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + void sm70_f16_lm_head_top1_out(torch::Tensor values_out, torch::Tensor indices_out, torch::Tensor in_feats, torch::Tensor weight, @@ -4875,6 +5567,27 @@ void fp8_gemm_sm70_out(torch::Tensor out, torch::Tensor _in_feats, group_size, k_ld, q_ld, gated_silu); } +void fp8_gemm_sm70_prefill_prescaled_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) { + 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_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, @@ -4946,6 +5659,41 @@ 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_f16_indexed_rerank_out(torch::Tensor out, torch::Tensor _in_feats, + torch::Tensor _kernel, + torch::Tensor candidate_ids, + torch::Tensor selected_raw, + torch::Tensor selected_packed, + torch::Tensor expanded, torch::Tensor partials, + torch::Tensor barriers, int64_t cta_n, + int64_t split_k) { + vllm::awq_sm70::sm70_f16_indexed_rerank_out( + out, _in_feats, _kernel, candidate_ids, selected_raw, selected_packed, + expanded, partials, barriers, cta_n, split_k); +} + +void sm70_f16_indexed_rerank_packed_out( + torch::Tensor out, torch::Tensor _in_feats, torch::Tensor _packed_kernel, + torch::Tensor candidate_ids, torch::Tensor selected_packed, + torch::Tensor expanded, torch::Tensor partials, torch::Tensor barriers, + int64_t cta_n, int64_t split_k) { + vllm::awq_sm70::sm70_f16_indexed_rerank_packed_out( + out, _in_feats, _packed_kernel, candidate_ids, selected_packed, expanded, + partials, barriers, cta_n, split_k); +} + +void sm70_f16_rerank_keys_out(torch::Tensor keys, torch::Tensor logits, + torch::Tensor candidate_ids) { + vllm::awq_sm70::sm70_f16_rerank_keys_out(keys, logits, candidate_ids); +} + +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) { + vllm::awq_sm70::sm70_f16_rerank_topk_out(values_out, ids_out, logits, + candidate_ids, vocab_start_index); +} + void sm70_f16_lm_head_top1_out(torch::Tensor values_out, torch::Tensor indices_out, torch::Tensor _in_feats, torch::Tensor _kernel, @@ -5908,17 +6656,29 @@ void awq_moe_single_token_weighted_reduce_out(torch::Tensor sorted_output, const int hidden_logical_size_i = static_cast(hidden_logical_size); const int sorted_output_row_stride = static_cast(sorted_output.stride(0)); - if (top_k == 8 && (hidden_logical_size_i % 2) == 0 && + if ((top_k == 6 || top_k == 8) && (hidden_logical_size_i % 2) == 0 && (sorted_output_row_stride % 2) == 0) { const int blocks = std::max( 1, ((hidden_logical_size_i >> 1) + kThreads - 1) / kThreads); - awq_moe_single_token_weighted_reduce_half2_kernel<8> - <<>>( - reinterpret_cast(sorted_output.data_ptr()), - topk_weights.data_ptr(), - inv_permuted_idx.data_ptr(), - reinterpret_cast<__half*>(out.data_ptr()), - hidden_logical_size_i, sorted_output_row_stride); + if (top_k == 6) { + awq_moe_single_token_weighted_reduce_half2_kernel<6> + <<>>( + reinterpret_cast( + sorted_output.data_ptr()), + topk_weights.data_ptr(), + inv_permuted_idx.data_ptr(), + reinterpret_cast<__half*>(out.data_ptr()), + hidden_logical_size_i, sorted_output_row_stride); + } else { + awq_moe_single_token_weighted_reduce_half2_kernel<8> + <<>>( + reinterpret_cast( + sorted_output.data_ptr()), + topk_weights.data_ptr(), + inv_permuted_idx.data_ptr(), + reinterpret_cast<__half*>(out.data_ptr()), + hidden_logical_size_i, sorted_output_row_stride); + } } else { const int blocks = std::max(1, (hidden_logical_size_i + kThreads - 1) / kThreads); @@ -7365,7 +8125,8 @@ void mxfp4_moe_gemm_sm70_out_impl( torch::Tensor out, torch::Tensor sorted_input, torch::Tensor expert_offsets, torch::Tensor strided_ptrs_w, torch::Tensor strided_ptrs_s, int64_t num_experts, int64_t k, int64_t n, int64_t group_size, - torch::Tensor b_group_indices, bool compact_grouped_rows = false) { + torch::Tensor b_group_indices, bool compact_grouped_rows = false, + bool broadcast_input_rows = false) { TORCH_CHECK( sorted_input.is_cuda() && sorted_input.scalar_type() == torch::kFloat16, "mxfp4_moe_gemm_sm70: input must be CUDA float16."); @@ -7386,9 +8147,16 @@ void mxfp4_moe_gemm_sm70_out_impl( "mxfp4_moe_gemm_sm70: k must be divisible by group_size."); TORCH_CHECK(sorted_input.dim() == 2 && sorted_input.size(1) == k, "mxfp4_moe_gemm_sm70: input shape mismatch."); - TORCH_CHECK(out.dim() == 2 && out.size(0) == sorted_input.size(0) && + const int64_t total_tokens = + broadcast_input_rows ? out.size(0) : sorted_input.size(0); + TORCH_CHECK(out.dim() == 2 && out.size(0) == total_tokens && out.size(1) == n && out.stride(1) == 1, "mxfp4_moe_gemm_sm70: output must be contiguous [tokens, n]."); + TORCH_CHECK(!broadcast_input_rows || + (sorted_input.size(0) == 1 && total_tokens == num_experts && + compact_grouped_rows), + "mxfp4_moe_gemm_sm70: broadcast input requires one physical " + "row and one logical row per compact group."); TORCH_CHECK(expert_offsets.numel() >= num_experts + 1, "mxfp4_moe_gemm_sm70: expert_offsets too small."); TORCH_CHECK(b_group_indices.is_cuda() && @@ -7398,7 +8166,6 @@ void mxfp4_moe_gemm_sm70_out_impl( "mxfp4_moe_gemm_sm70: B group indices must be contiguous CUDA " "int32."); - const int64_t total_tokens = sorted_input.size(0); if (total_tokens == 0) { return; } @@ -7422,7 +8189,8 @@ void mxfp4_moe_gemm_sm70_out_impl( static_cast(sorted_input.stride(0)), }; desc_A.num = static_cast(num_experts); - desc_A.offsets = expert_offsets.data_ptr(); + desc_A.offsets = + broadcast_input_rows ? nullptr : expert_offsets.data_ptr(); turbomind::gemm::MatrixLayout desc_U{}; const auto order_w = conv_w->order; @@ -7494,8 +8262,18 @@ void mxfp4_moe_gemm_sm70_out_impl( op.quant_b = {turbomind::gemm::QuantType::kK, static_cast(group_size)}; op.batch_dim = 0; op.dispatch_num_override = compact_grouped_rows ? 1 : 0; - op.active_group_count = - compact_grouped_rows ? -static_cast(num_experts) : 0; + const bool dynamic_expert_rows = + compact_grouped_rows && num_experts >= 12 && num_experts <= 48 && + num_experts % 6 == 0 && + vllm::awq_sm70::mxfp4_moe_grouped_m8_expert_rows_enabled(); + // The negative active-group contract is valid only when every group owns + // exactly one row (for example compact B1/top-6 and slot-grouped M8). Real + // expert segments can contain multiple rows and have graph-dynamic empty + // tails, so retain the single-group dispatch choice while letting the + // standard offsets scheduler discover their bounds on device. + op.active_group_count = compact_grouped_rows && !dynamic_expert_rows + ? -static_cast(num_experts) + : 0; auto& workspace_holder = vllm::awq_sm70::get_workspace(device, stream); auto& gemm = vllm::awq_sm70::get_gemm(device); @@ -7538,8 +8316,12 @@ void mxfp4_moe_dense_stage_sm70_out(torch::Tensor out, torch::Tensor input, "supported."); TORCH_CHECK(input.dim() == 2 && input.size(1) == k, "mxfp4_moe_dense_stage_sm70_out: input shape mismatch."); + 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(0) == input.size(0) && out.size(1) == n, + 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."); @@ -7552,20 +8334,28 @@ void mxfp4_moe_dense_stage_sm70_out(torch::Tensor out, torch::Tensor input, logged_mxfp4_dense_stage, "SM70 MXFP4 MoE CUDA-graph-safe dense-stage path enabled C++ op reached", input, input.size(0), num_experts); - const bool compact_decode_shape = - input.size(0) == num_experts && num_experts == 6 && - ((k == 4096 && n == 512) || (k == 256 && n == 4096)); - if (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); + const bool compact_decode_shape = input.size(0) == num_experts && + num_experts == 6 && + ((k == 4096 && (n == 512 || n == 1024)) || + (n == 4096 && (k == 256 || k == 512))); + 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); return; } const bool grouped_m8_shape = input.size(0) == 48 && num_experts == 48 && ((k == 4096 && n == 512) || (k == 256 && n == 4096)); - if (vllm::awq_sm70::mxfp4_moe_grouped_m8_enabled() && grouped_m8_shape) { + const bool grouped_verifier_shape = + input.size(0) == num_experts && num_experts >= 12 && num_experts <= 48 && + num_experts % 6 == 0 && + ((k == 4096 && n == 512) || (k == 256 && n == 4096)); + if ((vllm::awq_sm70::mxfp4_moe_grouped_m8_enabled() && grouped_m8_shape) || + (vllm::awq_sm70::mxfp4_moe_grouped_verifier_enabled() && + grouped_verifier_shape)) { // Keep every routed slot as an independent one-row group. Repeated experts // reuse the same weight pointer row, while the scheduler retains the exact // arithmetic contract already accepted for compact B1 decode. @@ -7582,6 +8372,237 @@ void mxfp4_moe_dense_stage_sm70_out(torch::Tensor out, torch::Tensor input, } } +void nvfp4_moe_gemm_sm70_out_impl( + torch::Tensor out, torch::Tensor sorted_input, torch::Tensor expert_offsets, + torch::Tensor strided_ptrs_w, torch::Tensor strided_ptrs_s, + int64_t num_experts, int64_t k, int64_t n, int64_t group_size, + torch::Tensor b_group_indices, bool compact_grouped_rows = false) { + TORCH_CHECK( + sorted_input.is_cuda() && sorted_input.scalar_type() == torch::kFloat16, + "nvfp4_moe_gemm_sm70: input must be CUDA float16."); + TORCH_CHECK( + expert_offsets.is_cuda() && + expert_offsets.scalar_type() == torch::kInt32 && + expert_offsets.is_contiguous(), + "nvfp4_moe_gemm_sm70: expert_offsets must be contiguous CUDA int32."); + TORCH_CHECK(strided_ptrs_w.is_cuda() && strided_ptrs_s.is_cuda(), + "nvfp4_moe_gemm_sm70: strided_ptrs must be CUDA."); + TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kFloat16, + "nvfp4_moe_gemm_sm70: output must be CUDA float16."); + TORCH_CHECK(num_experts > 0 && k > 0 && n > 0, + "nvfp4_moe_gemm_sm70: invalid dimensions."); + TORCH_CHECK(group_size == 16, + "nvfp4_moe_gemm_sm70: only group_size=16 is supported."); + TORCH_CHECK(k % group_size == 0, + "nvfp4_moe_gemm_sm70: k must be divisible by group_size."); + TORCH_CHECK(sorted_input.dim() == 2 && sorted_input.size(1) == k, + "nvfp4_moe_gemm_sm70: input shape mismatch."); + TORCH_CHECK(out.dim() == 2 && out.size(0) == sorted_input.size(0) && + out.size(1) == n && out.stride(1) == 1, + "nvfp4_moe_gemm_sm70: output must be contiguous [tokens, n]."); + TORCH_CHECK(expert_offsets.numel() >= num_experts + 1, + "nvfp4_moe_gemm_sm70: expert_offsets too small."); + TORCH_CHECK(b_group_indices.is_cuda() && + b_group_indices.scalar_type() == torch::kInt32 && + b_group_indices.is_contiguous() && + b_group_indices.numel() >= num_experts, + "nvfp4_moe_gemm_sm70: B group indices must be contiguous CUDA " + "int32."); + + const int64_t total_tokens = sorted_input.size(0); + if (total_tokens == 0) { + return; + } + + const at::cuda::OptionalCUDAGuard device_guard(device_of(sorted_input)); + const int device = sorted_input.get_device(); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + const auto fp4_converters = turbomind::gemm::GetConverters( + turbomind::kHalf, turbomind::kFloat4_e2m1, turbomind::kHalf, true, 70); + const auto fp8_converters = turbomind::gemm::GetConverters( + turbomind::kHalf, turbomind::kFloat8_e4m3, turbomind::kHalf, true, 70); + const auto* conv_w = fp4_converters[0]; + const auto* conv_s = fp8_converters[1]; + TORCH_CHECK(conv_w && conv_s, + "nvfp4_moe_gemm_sm70: no compatible TurboMind converters."); + + turbomind::gemm::MatrixLayout desc_A{ + turbomind::kHalf, + turbomind::gemm::kRowMajor, + static_cast(total_tokens), + static_cast(k), + static_cast(sorted_input.stride(0)), + }; + desc_A.num = static_cast(num_experts); + desc_A.offsets = expert_offsets.data_ptr(); + turbomind::gemm::MatrixLayout desc_U{}; + + const auto order_w = conv_w->order; + const bool is_a_w = turbomind::gemm::get_operand_tag(conv_w->pack) == + turbomind::gemm::OPERAND_A; + const bool is_b_w = !is_a_w; + turbomind::gemm::MatrixLayout weight_desc{ + turbomind::kHalf, + order_w, + static_cast(n), + static_cast(k), + order_w == turbomind::gemm::kRowMajor ? static_cast(k) + : static_cast(n), + }; + if (is_b_w) { + std::swap(weight_desc.rows, weight_desc.cols); + weight_desc.order = ~weight_desc.order; + } + turbomind::gemm::MatrixLayout desc_B = weight_desc; + desc_B.type = turbomind::kFloat4_e2m1; + desc_B.pack = conv_w->pack; + if (is_a_w) { + desc_B = turbomind::gemm::transpose(desc_B); + } + desc_B.ld = 0; + desc_B.num = static_cast(num_experts); + desc_B.group_idxs = b_group_indices.data_ptr(); + + const auto order_s = conv_s->order; + const bool is_a_s = turbomind::gemm::get_operand_tag(conv_s->pack) == + turbomind::gemm::OPERAND_U; + const bool is_b_s = !is_a_s; + const int64_t num_groups = k / group_size; + turbomind::gemm::MatrixLayout scale_desc{ + turbomind::kUint16, order_s, + static_cast(n), static_cast(num_groups), + static_cast(n), + }; + if (is_b_s) { + std::swap(scale_desc.rows, scale_desc.cols); + scale_desc.order = ~scale_desc.order; + } + turbomind::gemm::MatrixLayout desc_V = scale_desc; + desc_V.pack = conv_s->pack; + if (is_a_s) { + desc_V = turbomind::gemm::transpose(desc_V); + } + desc_V.ld = 0; + desc_V.num = static_cast(num_experts); + desc_V.group_idxs = b_group_indices.data_ptr(); + + turbomind::gemm::MatrixLayout desc_D{ + turbomind::kHalf, + turbomind::gemm::kRowMajor, + static_cast(total_tokens), + static_cast(n), + static_cast(out.stride(0)), + }; + desc_D.num = static_cast(num_experts); + desc_D.offsets = expert_offsets.data_ptr(); + + turbomind::gemm::Operation op{}; + op.dispatch = vllm::awq_sm70::select_nvfp4_moe_dispatch_policy( + device, static_cast(total_tokens), static_cast(n), + static_cast(k), static_cast(num_experts), + static_cast(group_size), stream); + op.epilogue = turbomind::gemm::Epilogue::kNone; + op.quant_a = {turbomind::gemm::QuantType::kNone, 0}; + op.quant_b = {turbomind::gemm::QuantType::kK, static_cast(group_size)}; + op.batch_dim = 0; + op.dispatch_num_override = compact_grouped_rows ? 1 : 0; + op.active_group_count = + compact_grouped_rows ? -static_cast(num_experts) : 0; + + auto& workspace_holder = vllm::awq_sm70::get_workspace(device, stream); + auto& gemm = vllm::awq_sm70::get_gemm(device); + const int ec = + gemm.Run(op, 1.f, sorted_input.data_ptr(), desc_A, nullptr, desc_U, + strided_ptrs_w.data_ptr(), desc_B, strided_ptrs_s.data_ptr(), + desc_V, 0.f, out.data_ptr(), desc_D, out.data_ptr(), desc_D, + workspace_holder.workspace, stream); + TORCH_CHECK(ec == 0, + "nvfp4_moe_gemm_sm70: TurboMind batched GEMM failed (ec=", ec, + ")."); +} + +void nvfp4_moe_dense_stage_sm70_out(torch::Tensor out, torch::Tensor input, + torch::Tensor expert_offsets, + torch::Tensor dense_expert_ids, + torch::Tensor ptrs_w, torch::Tensor ptrs_s, + int64_t num_experts, int64_t k, int64_t n, + int64_t group_size) { + TORCH_CHECK(input.is_cuda() && input.scalar_type() == torch::kFloat16, + "nvfp4_moe_dense_stage_sm70_out: input must be CUDA float16."); + TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kFloat16, + "nvfp4_moe_dense_stage_sm70_out: out must be CUDA float16."); + TORCH_CHECK(expert_offsets.is_cuda() && + expert_offsets.scalar_type() == torch::kInt32 && + expert_offsets.is_contiguous(), + "nvfp4_moe_dense_stage_sm70_out: expert_offsets must be " + "contiguous CUDA int32."); + TORCH_CHECK(dense_expert_ids.is_cuda() && + dense_expert_ids.scalar_type() == torch::kInt32 && + dense_expert_ids.is_contiguous(), + "nvfp4_moe_dense_stage_sm70_out: dense_expert_ids must be " + "contiguous CUDA int32."); + TORCH_CHECK(ptrs_w.is_cuda() && ptrs_s.is_cuda(), + "nvfp4_moe_dense_stage_sm70_out: ptr rows must be CUDA."); + TORCH_CHECK(num_experts > 0, + "nvfp4_moe_dense_stage_sm70_out: num_experts must be positive."); + TORCH_CHECK(group_size == 16, + "nvfp4_moe_dense_stage_sm70_out: only group_size=16 is " + "supported."); + TORCH_CHECK(input.dim() == 2 && input.size(1) == k, + "nvfp4_moe_dense_stage_sm70_out: input shape mismatch."); + TORCH_CHECK( + out.dim() == 2 && out.size(0) == input.size(0) && out.size(1) == n, + "nvfp4_moe_dense_stage_sm70_out: out shape mismatch."); + TORCH_CHECK(expert_offsets.numel() >= num_experts + 1, + "nvfp4_moe_dense_stage_sm70_out: expert_offsets too small."); + TORCH_CHECK(dense_expert_ids.numel() >= num_experts, + "nvfp4_moe_dense_stage_sm70_out: dense_expert_ids too small."); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); + static std::atomic logged_nvfp4_dense_stage{0u}; + maybe_log_sm70_moe_route_once( + logged_nvfp4_dense_stage, + "SM70 NVFP4 MoE CUDA-graph-safe TurboMind path enabled C++ op reached", + input, input.size(0), num_experts); + constexpr int kNvfp4LegacyCompactGroups = 8 * 8; + // Cover the validated top-k8 B10 verifier bound. Duplicate expert selections + // remain separate one-row groups, so this is a routed-slot bound rather + // than an active-expert bound. Keep the dense-grouped cutoff independent: + // a disabled compact route above 64 rows must not fall back per expert. + constexpr int kNvfp4MaxCompactGroups = 10 * 8; + const bool compact_decode_shape = + input.size(0) == num_experts && num_experts <= kNvfp4MaxCompactGroups; + if (compact_decode_shape) { + nvfp4_moe_gemm_sm70_out_impl(out, input, expert_offsets, ptrs_w, ptrs_s, + num_experts, k, n, group_size, + dense_expert_ids, true); + return; + } + const bool exact_qwen36_prefill_shape = + input.size(0) > kNvfp4LegacyCompactGroups && num_experts == 256 && + ((k == 2048 && (n == 1024 || n == 512 || n == 256)) || + (n == 2048 && (k == 512 || k == 256 || k == 128))); + if (vllm::awq_sm70::nvfp4_moe_grouped_prefill_enabled() && + exact_qwen36_prefill_shape) { + static std::atomic logged_nvfp4_grouped_prefill{0u}; + maybe_log_sm70_moe_route_once( + logged_nvfp4_grouped_prefill, + "SM70 NVFP4 MoE grouped TurboMind prefill path enabled C++ op reached", + input, input.size(0), num_experts); + nvfp4_moe_gemm_sm70_out_impl(out, input, expert_offsets, ptrs_w, ptrs_s, + num_experts, k, n, group_size, + dense_expert_ids); + return; + } + for (int expert = 0; expert < static_cast(num_experts); ++expert) { + torch::Tensor offsets = expert_offsets.narrow(0, expert, 2); + torch::Tensor expert_idx = dense_expert_ids.narrow(0, expert, 1); + nvfp4_moe_gemm_sm70_out_impl(out, input, offsets, ptrs_w, ptrs_s, 1, k, n, + group_size, expert_idx); + } +} + void mxfp4_moe_single_token_prepare_w13_sm70_out( torch::Tensor gate_up, torch::Tensor compact_input, torch::Tensor x, torch::Tensor topk_ids, torch::Tensor w13_ptrs_w, torch::Tensor w13_ptrs_s, @@ -7606,9 +8627,10 @@ void mxfp4_moe_single_token_prepare_w13_sm70_out( TORCH_CHECK(topk_ids.numel() == kTopK, "mxfp4_moe_single_token_prepare_w13_sm70_out: exact top-k=6 is " "required."); - TORCH_CHECK(group_size == 32 && w13_k == hidden_logical_size && w13_n == 512, + TORCH_CHECK(group_size == 32 && w13_k == hidden_logical_size && + (w13_n == 512 || w13_n == 1024), "mxfp4_moe_single_token_prepare_w13_sm70_out: unsupported " - "DeepSeek V4 W13 shape contract."); + "MXFP4 W13 shape contract."); TORCH_CHECK( compact_input.is_cuda() && compact_input.scalar_type() == torch::kFloat16 && @@ -7655,6 +8677,9 @@ void mxfp4_moe_single_token_prepare_w13_sm70_out( x.size(0), kTopK); constexpr int kThreads = 256; + const bool broadcast_input = + w13_n == 512 && + vllm::awq_sm70::mxfp4_moe_broadcast_input_decode_enabled(); const int src_row_stride = kPtrRowBytes; const int row_bytes = kPtrRowBytes; if (topk_ids.scalar_type() == torch::kInt32) { @@ -7665,8 +8690,10 @@ void mxfp4_moe_single_token_prepare_w13_sm70_out( w13_ptrs_w.data_ptr(), w13_ptrs_s.data_ptr(), w13_ptrs_w.data_ptr(), w13_ptrs_s.data_ptr(), w13_ptrs_w.data_ptr(), w13_ptrs_s.data_ptr(), - reinterpret_cast<__half*>(compact_input.data_ptr()), nullptr, - expert_offsets.data_ptr(), nullptr, + broadcast_input + ? nullptr + : reinterpret_cast<__half*>(compact_input.data_ptr()), + nullptr, expert_offsets.data_ptr(), nullptr, inv_permuted_idx.data_ptr(), sorted_expert_ids.data_ptr(), kTopK, static_cast(hidden_logical_size), src_row_stride, src_row_stride, @@ -7679,8 +8706,10 @@ void mxfp4_moe_single_token_prepare_w13_sm70_out( w13_ptrs_w.data_ptr(), w13_ptrs_s.data_ptr(), w13_ptrs_w.data_ptr(), w13_ptrs_s.data_ptr(), w13_ptrs_w.data_ptr(), w13_ptrs_s.data_ptr(), - reinterpret_cast<__half*>(compact_input.data_ptr()), nullptr, - expert_offsets.data_ptr(), nullptr, + broadcast_input + ? nullptr + : reinterpret_cast<__half*>(compact_input.data_ptr()), + nullptr, expert_offsets.data_ptr(), nullptr, inv_permuted_idx.data_ptr(), sorted_expert_ids.data_ptr(), kTopK, static_cast(hidden_logical_size), src_row_stride, src_row_stride, @@ -7688,9 +8717,10 @@ void mxfp4_moe_single_token_prepare_w13_sm70_out( } C10_CUDA_KERNEL_LAUNCH_CHECK(); - mxfp4_moe_gemm_sm70_out_impl(gate_up, compact_input, expert_offsets, - w13_ptrs_w, w13_ptrs_s, kTopK, w13_k, w13_n, - group_size, sorted_expert_ids, true); + mxfp4_moe_gemm_sm70_out_impl(gate_up, broadcast_input ? x : compact_input, + expert_offsets, w13_ptrs_w, w13_ptrs_s, kTopK, + w13_k, w13_n, group_size, sorted_expert_ids, + true, broadcast_input); } void fp8_moe_single_token_dense_stage_sm70_out( diff --git a/csrc/sm70_turbomind/ops/fp8_qpn8_sm70.cu b/csrc/sm70_turbomind/ops/fp8_qpn8_sm70.cu new file mode 100644 index 0000000000..2dd90101e8 --- /dev/null +++ b/csrc/sm70_turbomind/ops/fp8_qpn8_sm70.cu @@ -0,0 +1,1255 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// The QPN8 execution layout is derived from dnv2003/v100-skinny (MIT) and +// its block-scale adaptation in haohervchb/sglang-V100. Automatic operator +// dispatch is restricted to accepted SM70 tensor/layout contracts and can be +// disabled with VLLM_SM70_FP8_QPN8=0. +// See LICENSE.v100-skinny in this directory for the retained MIT notice. + +#include +#include + +#include +#include +#include +#include +#include + +#include +#include + +namespace { + +__device__ __forceinline__ void fp8x8_to_half2x4(uint2 q, half2 out[4]) { +#pragma unroll + for (int i = 0; i < 4; ++i) { + const unsigned b0 = (q.x >> (8 * i)) & 0xffu; + const unsigned b1 = (q.y >> (8 * i)) & 0xffu; + const unsigned h0 = ((b0 & 0x80u) << 8) | ((b0 & 0x7fu) << 7); + const unsigned h1 = ((b1 & 0x80u) << 8) | ((b1 & 0x7fu) << 7); + const unsigned packed = h0 | (h1 << 16); + out[i] = *reinterpret_cast(&packed); + } +} + +__device__ __forceinline__ void fp8x8_to_half2x4_fast(uint2 q, half2 out[4]) { + constexpr unsigned kSign = 0x80008000u; + constexpr unsigned kExponentMantissa = 0x3f803f80u; + unsigned permuted[4]; + permuted[0] = __byte_perm(q.x, q.y, 0x0400); + permuted[1] = __byte_perm(q.x, q.y, 0x0501); + permuted[2] = __byte_perm(q.x, q.y, 0x0602); + permuted[3] = __byte_perm(q.x, q.y, 0x0703); +#pragma unroll + for (int i = 0; i < 4; ++i) { + const unsigned value = + ((permuted[i] << 8) & kSign) | ((permuted[i] << 7) & kExponentMantissa); + out[i] = *reinterpret_cast(&value); + } +} + +constexpr int kQpn8PrepareThreads = 256; + +__device__ __forceinline__ int qpn8_col_from_lane(int lane) { + return ((lane >> 2) & 3) * 8 + (lane & 3) + ((lane & 16) ? 4 : 0); +} + +__device__ __forceinline__ int qpn8_lane_from_col(int col) { + return (col & 3) | (((col >> 3) & 3) << 2) | (((col >> 2) & 1) << 4); +} + +__device__ __forceinline__ int qpn8_physical_k(int logical_k) { + const int local = logical_k & 7; + return (logical_k & 8) + (local >> 1) + ((local & 1) << 2); +} + +__global__ void fp8_qpn8_prepack_sm70_kernel( + uint8_t* __restrict__ codes, const uint8_t* __restrict__ qweight, int n, + int k) { + const size_t index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const size_t numel = static_cast(n) * k; + if (index >= numel) { + return; + } + + const int row = static_cast(index / k); + const int logical_k = static_cast(index % k); + const int tile = row >> 5; + const int lane = qpn8_lane_from_col(row & 31); + const int group = logical_k >> 4; + const int physical_k = qpn8_physical_k(logical_k & 15); + const int groups_k16 = k >> 4; + const size_t packed_index = + (((static_cast(tile) * groups_k16 + group) * 32 + lane) * 16 + + physical_k); + codes[packed_index] = qweight[index]; +} + +__global__ void fp8_qpn8_scale_sm70_kernel(half* __restrict__ group_scales, + const float* __restrict__ scales, + int n_blocks, int k_blocks) { + const int n_tiles = n_blocks * 4; + const int index = blockIdx.x * blockDim.x + threadIdx.x; + const int numel = k_blocks * n_tiles; + if (index >= numel) { + return; + } + + const int k_block = index / n_tiles; + const int n_tile = index - k_block * n_tiles; + group_scales[index] = + __float2half(scales[(n_tile >> 2) * k_blocks + k_block] * 256.0f); +} + +__global__ void fp8_qpn8_channel_scale_sm70_kernel( + half* __restrict__ channel_scales, const float* __restrict__ scales, + int n) { + const int col = blockIdx.x * blockDim.x + threadIdx.x; + if (col < n) { + channel_scales[col] = __float2half(scales[col] * 256.0f); + } +} + +__global__ void fp8_qpn8_dequantize_sm70_kernel( + half* __restrict__ output, const uint8_t* __restrict__ codes, + const half* __restrict__ group_scales, int n, int k, bool channel_scales) { + const size_t word_index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const size_t word_count = static_cast(n) * k / 16; + if (word_index >= word_count) { + return; + } + + const int groups_k16 = k >> 4; + const int lane = static_cast(word_index & 31); + const size_t tile_word = word_index >> 5; + const int group = static_cast(tile_word % groups_k16); + const int tile = static_cast(tile_word / groups_k16); + const int col = tile * 32 + qpn8_col_from_lane(lane); + const int tiles_n32 = n >> 5; + const half scale = channel_scales + ? group_scales[col] + : group_scales[(group >> 3) * tiles_n32 + tile]; + const half2 scale2 = __halves2half2(scale, scale); + const uint4 packed = reinterpret_cast(codes)[word_index]; + 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 pair = 0; pair < 8; ++pair) { + const half2 value = __hmul2(weights[pair], scale2); + const int k_base = group * 16 + pair * 2; + output[static_cast(k_base) * n + col] = __low2half(value); + output[static_cast(k_base + 1) * n + col] = __high2half(value); + } +} + +__global__ void fp8_qpn8_silu_and_mul_sm70_kernel( + half* __restrict__ output, const half* __restrict__ gate_up, int rows, + int hidden) { + const int row = blockIdx.x; + if (row >= rows) { + return; + } + const half* row_input = gate_up + static_cast(row) * hidden * 2; + half* row_output = output + static_cast(row) * hidden; + for (int col = threadIdx.x; col < hidden; col += blockDim.x) { + const float gate = __half2float(row_input[col]); + const float silu = gate / (1.0f + __expf(-gate)); + row_output[col] = __hmul(__float2half(silu), row_input[hidden + col]); + } +} + +#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)) + +__global__ void fp8_qpn8_ba_split_copy_sm70_kernel( + const half* __restrict__ qkvz, const half* __restrict__ ba, + half* __restrict__ qkv, half* __restrict__ z, half* __restrict__ b, + half* __restrict__ a, int m) { + constexpr int kQkvN = 2560; + constexpr int kZN = 1536; + constexpr int kQkvzN = kQkvN + kZN; + constexpr int kBAN = 24; + constexpr int kOutputN = kQkvzN + kBAN; + const int index = blockIdx.x * blockDim.x + threadIdx.x; + if (index >= m * kOutputN) { + return; + } + const int row = index / kOutputN; + const int col = index - row * kOutputN; + if (col < kQkvN) { + qkv[static_cast(row) * kQkvN + col] = + qkvz[static_cast(row) * kQkvzN + col]; + } else if (col < kQkvzN) { + z[static_cast(row) * kZN + col - kQkvN] = + qkvz[static_cast(row) * kQkvzN + col]; + } else if (col < kQkvzN + kBAN / 2) { + b[static_cast(row) * (kBAN / 2) + col - kQkvzN] = + ba[static_cast(row) * kBAN + col - kQkvzN]; + } else { + a[static_cast(row) * (kBAN / 2) + col - kQkvzN - kBAN / 2] = + ba[static_cast(row) * kBAN + col - kQkvzN]; + } +} + +template +__global__ void fp8_qpn8_sm70_kernel( + const uint8_t* __restrict__ codes, const half* __restrict__ group_scales, + const half* __restrict__ input, half* __restrict__ output, + half* __restrict__ z_output, const half* __restrict__ ba_weight, + half* __restrict__ ba_output, half* __restrict__ b_output, + half* __restrict__ a_output, int ba_n, int qkv_n, int n, int k, int m, + bool channel_scales) { + __shared__ float partials[SplitK][M1Only ? 32 : 256]; + + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + const int tile = blockIdx.x; + if constexpr (FusedBA) { + const int qpn_tiles = n >> 5; + if (tile >= qpn_tiles) { + constexpr int kBARowsPerBlock = 2; + constexpr int kBAThreadsPerRow = 256; + const int ba_block = tile - qpn_tiles; + const int ba_group = threadIdx.x / kBAThreadsPerRow; + const int ba_thread = threadIdx.x % kBAThreadsPerRow; + const int ba_warp = ba_thread >> 5; + const int ba_row = ba_block * kBARowsPerBlock + ba_group; + float value = 0.0f; + const half2* input2 = reinterpret_cast(input); + const half2* weight2 = reinterpret_cast( + ba_weight + static_cast(ba_row) * k); + for (int pair = ba_thread; pair < k / 2; pair += kBAThreadsPerRow) { + const float2 x = __half22float2(__ldg(input2 + pair)); + const float2 weight = __half22float2(__ldg(weight2 + pair)); + value = fmaf(x.x, weight.x, value); + value = fmaf(x.y, weight.y, value); + } +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value += __shfl_down_sync(0xffffffffU, value, offset); + } + if (lane == 0) { + partials[ba_group][ba_warp] = value; + } + __syncthreads(); + if (ba_warp == 0) { + value = lane < 8 ? partials[ba_group][lane] : 0.0f; +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value += __shfl_down_sync(0xffffffffU, value, offset); + } + if (lane == 0 && ba_row < ba_n) { + if constexpr (SplitOutputs) { + if (ba_row < ba_n / 2) { + b_output[ba_row] = __float2half(value); + } else { + a_output[ba_row - ba_n / 2] = __float2half(value); + } + } else { + ba_output[ba_row] = __float2half(value); + } + } + } + return; + } + } + 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_warp = groups_k16 / SplitK; + const int group_begin = warp * groups_per_warp; + const int tiles_n32 = n >> 5; + const uint4* code_ptr = reinterpret_cast(codes) + + static_cast(tile) * groups_k16 * 32 + lane; + const half* scale_ptr = group_scales + tile; + + float accum[NAcc][8]; +#pragma unroll + for (int chain = 0; chain < NAcc; ++chain) { +#pragma unroll + for (int i = 0; i < 8; ++i) { + accum[chain][i] = 0.0f; + } + } + int loaded_scale_group = -1; + half loaded_scale = + channel_scales + ? __ldg(group_scales + tile * 32 + qpn8_col_from_lane(lane)) + : __float2half(0.0f); + uint4 prefetched = make_uint4(0, 0, 0, 0); + if constexpr (PrefetchCodes) { + prefetched = __ldcs(code_ptr + static_cast(group_begin) * 32); + } + +#pragma unroll 4 + for (int group = group_begin; group < group_begin + groups_per_warp; + ++group) { + const int scale_group = group >> 3; + if (!channel_scales && scale_group != loaded_scale_group) { + loaded_scale = + __ldg(scale_ptr + static_cast(scale_group) * tiles_n32); + loaded_scale_group = scale_group; + } + + const uint4 packed = + PrefetchCodes ? prefetched + : __ldcs(code_ptr + static_cast(group) * 32); + uint4 next = make_uint4(0, 0, 0, 0); + if constexpr (PrefetchCodes) { + if (group + 1 < group_begin + groups_per_warp) { + next = __ldcs(code_ptr + static_cast(group + 1) * 32); + } + } + half2 weights[8]; + if constexpr (FastDecoder) { + fp8x8_to_half2x4_fast(make_uint2(packed.x, packed.y), weights); + fp8x8_to_half2x4_fast(make_uint2(packed.z, packed.w), weights + 4); + } else { + fp8x8_to_half2x4(make_uint2(packed.x, packed.y), weights); + fp8x8_to_half2x4(make_uint2(packed.z, packed.w), weights + 4); + } + + const half2 scale2 = __halves2half2(loaded_scale, loaded_scale); +#pragma unroll + for (int i = 0; i < 8; ++i) { + weights[i] = __hmul2(weights[i], scale2); + } + + uint4 input01 = make_uint4(0, 0, 0, 0); + uint4 input23 = make_uint4(0, 0, 0, 0); + if (row < m) { + const half* input_row = input + static_cast(row) * k; + 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); + 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]); + if constexpr (PrefetchCodes) { + prefetched = next; + } + } + +#pragma unroll + for (int chain = 1; chain < NAcc; ++chain) { +#pragma unroll + for (int i = 0; i < 8; ++i) { + accum[0][i] += accum[chain][i]; + } + } + + if constexpr (M1Only) { + if ((lane & 17) == 0) { +#pragma unroll + for (int pair = 0; pair < 2; ++pair) { +#pragma unroll + for (int offset = 0; offset < 2; ++offset) { + const int i = pair * 4 + offset; + const int output_col = + offset | (((lane >> 1) & 1) << 1) | (pair << 2); + partials[warp][quadpair * 8 + output_col] = accum[0][i]; + } + } + } + } else { +#pragma unroll + for (int i = 0; i < 8; ++i) { + const int output_row = (i & 2) | ((lane & 16) ? 4 : 0) | (lane & 1); + const int output_col = + (i & 1) | (((lane >> 1) & 1) << 1) | ((i >> 2) << 2); + partials[warp][output_row * 32 + quadpair * 8 + output_col] = accum[0][i]; + } + } + __syncthreads(); + + constexpr int kOutputElements = M1Only ? 32 : 256; + for (int element = threadIdx.x; element < kOutputElements; + element += blockDim.x) { + float value = 0.0f; +#pragma unroll + for (int k_warp = 0; k_warp < SplitK; ++k_warp) { + value += partials[k_warp][element]; + } + if constexpr (M1Only) { + const int output_col = tile * 32 + element; + if constexpr (SplitOutputs) { + if (output_col < qkv_n) { + output[output_col] = __float2half(value); + } else { + z_output[output_col - qkv_n] = __float2half(value); + } + } else { + output[output_col] = __float2half(value); + } + } else { + const int output_row = element >> 5; + const int output_col = element & 31; + if (output_row < m) { + output[static_cast(output_row) * n + tile * 32 + output_col] = + __float2half(value); + } + } + } +} + +template +void launch_fp8_qpn8_sm70(const uint8_t* codes, const half* group_scales, + const half* input, half* output, int n, int k, int m, + bool channel_scales, cudaStream_t stream) { + fp8_qpn8_sm70_kernel + <<<(n / 32), (32 * SplitK), 0, stream>>>( + codes, group_scales, input, output, nullptr, nullptr, nullptr, + nullptr, nullptr, 0, n, n, k, m, channel_scales); +} + +void launch_fp8_qpn8_ba_split_sm70(const uint8_t* codes, + const half* group_scales, const half* input, + half* qkv_output, half* z_output, + const half* ba_weight, half* b_output, + half* a_output, int ba_n, int qkv_n, int n, + int k, bool channel_scales, + cudaStream_t stream) { + fp8_qpn8_sm70_kernel<16, 2, true, false, true, true, true> + <<<(n / 32 + (ba_n + 1) / 2), 512, 0, stream>>>( + codes, group_scales, input, qkv_output, z_output, ba_weight, nullptr, + b_output, a_output, ba_n, qkv_n, n, k, 1, channel_scales); +} + +template +__global__ void fp8_qpn8_gated_pair_sm70_kernel( + const uint8_t* __restrict__ codes, const half* __restrict__ group_scales, + const half* __restrict__ input, half* __restrict__ output, int hidden, + int k, int m, bool channel_scales) { + __shared__ float partials[2][SplitK][M1Only ? 32 : 256]; + + const int lane = threadIdx.x & 31; + const int warp_in_block = threadIdx.x >> 5; + const int projection = warp_in_block / SplitK; + const int warp = warp_in_block - projection * SplitK; + const int hidden_tiles = hidden >> 5; + const int tiles_n32 = hidden_tiles * 2; + const int tile = blockIdx.x + projection * hidden_tiles; + 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_warp = groups_k16 / SplitK; + const int group_begin = warp * groups_per_warp; + const uint4* code_ptr = reinterpret_cast(codes) + + static_cast(tile) * groups_k16 * 32 + lane; + const half* scale_ptr = group_scales + tile; + + float accum[NAcc][8]; +#pragma unroll + for (int chain = 0; chain < NAcc; ++chain) { +#pragma unroll + for (int i = 0; i < 8; ++i) { + accum[chain][i] = 0.0f; + } + } + int loaded_scale_group = -1; + half loaded_scale = + channel_scales + ? __ldg(group_scales + tile * 32 + qpn8_col_from_lane(lane)) + : __float2half(0.0f); + uint4 prefetched = make_uint4(0, 0, 0, 0); + if constexpr (PrefetchCodes) { + prefetched = __ldcs(code_ptr + static_cast(group_begin) * 32); + } + +#pragma unroll 4 + for (int group = group_begin; group < group_begin + groups_per_warp; + ++group) { + const int scale_group = group >> 3; + if (!channel_scales && scale_group != loaded_scale_group) { + loaded_scale = + __ldg(scale_ptr + static_cast(scale_group) * tiles_n32); + loaded_scale_group = scale_group; + } + + const uint4 packed = + PrefetchCodes ? prefetched + : __ldcs(code_ptr + static_cast(group) * 32); + uint4 next = make_uint4(0, 0, 0, 0); + if constexpr (PrefetchCodes) { + if (group + 1 < group_begin + groups_per_warp) { + next = __ldcs(code_ptr + static_cast(group + 1) * 32); + } + } + half2 weights[8]; + if constexpr (FastDecoder) { + fp8x8_to_half2x4_fast(make_uint2(packed.x, packed.y), weights); + fp8x8_to_half2x4_fast(make_uint2(packed.z, packed.w), weights + 4); + } else { + fp8x8_to_half2x4(make_uint2(packed.x, packed.y), weights); + fp8x8_to_half2x4(make_uint2(packed.z, packed.w), weights + 4); + } + const half2 scale2 = __halves2half2(loaded_scale, loaded_scale); +#pragma unroll + for (int i = 0; i < 8; ++i) { + weights[i] = __hmul2(weights[i], scale2); + } + + uint4 input01 = make_uint4(0, 0, 0, 0); + uint4 input23 = make_uint4(0, 0, 0, 0); + if (row < m) { + const half* input_row = input + static_cast(row) * k; + 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); + 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]); + if constexpr (PrefetchCodes) { + prefetched = next; + } + } + +#pragma unroll + for (int chain = 1; chain < NAcc; ++chain) { +#pragma unroll + for (int i = 0; i < 8; ++i) { + accum[0][i] += accum[chain][i]; + } + } + if constexpr (M1Only) { + if ((lane & 17) == 0) { +#pragma unroll + for (int pair = 0; pair < 2; ++pair) { +#pragma unroll + for (int offset = 0; offset < 2; ++offset) { + const int i = pair * 4 + offset; + const int output_col = + offset | (((lane >> 1) & 1) << 1) | (pair << 2); + partials[projection][warp][quadpair * 8 + output_col] = accum[0][i]; + } + } + } + } else { +#pragma unroll + for (int i = 0; i < 8; ++i) { + const int output_row = (i & 2) | ((lane & 16) ? 4 : 0) | (lane & 1); + const int output_col = + (i & 1) | (((lane >> 1) & 1) << 1) | ((i >> 2) << 2); + partials[projection][warp][output_row * 32 + quadpair * 8 + output_col] = + accum[0][i]; + } + } + __syncthreads(); + + constexpr int kOutputElements = M1Only ? 32 : 256; + for (int element = threadIdx.x; element < kOutputElements; + element += blockDim.x) { + float gate = 0.0f; + float up = 0.0f; +#pragma unroll + for (int k_warp = 0; k_warp < SplitK; ++k_warp) { + gate += partials[0][k_warp][element]; + up += partials[1][k_warp][element]; + } + if constexpr (M1Only) { + const float silu = gate / (1.0f + __expf(-gate)); + output[blockIdx.x * 32 + element] = __float2half(silu * up); + } else { + const int output_row = element >> 5; + const int output_col = element & 31; + if (output_row < m) { + const float silu = gate / (1.0f + __expf(-gate)); + output[static_cast(output_row) * hidden + blockIdx.x * 32 + + output_col] = __float2half(silu * up); + } + } + } +} + +template +void launch_fp8_qpn8_gated_pair_sm70(const uint8_t* codes, + const half* group_scales, + const half* input, half* output, + int hidden, int k, int m, + bool channel_scales, cudaStream_t stream) { + fp8_qpn8_gated_pair_sm70_kernel + <<<(hidden / 32), (64 * SplitK), 0, stream>>>( + codes, group_scales, input, output, hidden, k, m, channel_scales); +} + +} // namespace + +std::vector fp8_qpn8_prepare_sm70(torch::Tensor qweight, + torch::Tensor scales) { + TORCH_CHECK(qweight.is_cuda() && scales.is_cuda(), + "fp8_qpn8_prepare_sm70: tensors must be CUDA tensors"); + TORCH_CHECK(qweight.scalar_type() == at::ScalarType::Float8_e4m3fn, + "fp8_qpn8_prepare_sm70: weight must be float8_e4m3fn"); + TORCH_CHECK(scales.scalar_type() == torch::kFloat32, + "fp8_qpn8_prepare_sm70: scales must be float32"); + TORCH_CHECK(qweight.dim() == 2 && scales.dim() == 2, + "fp8_qpn8_prepare_sm70: tensors must be 2D"); + TORCH_CHECK(qweight.is_contiguous() && scales.is_contiguous(), + "fp8_qpn8_prepare_sm70: tensors must be contiguous"); + TORCH_CHECK(qweight.get_device() == scales.get_device(), + "fp8_qpn8_prepare_sm70: tensors must share one device"); + + 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 || block_scales, + "fp8_qpn8_prepare_sm70: expected channel scales [N, 1] or " + "block scales [N/128, K/128]"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(qweight)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + auto codes = torch::empty( + {k, n}, + torch::TensorOptions().device(qweight.device()).dtype(torch::kUInt8)); + auto group_scales = torch::empty( + {channel_scales ? 1 : k / 128, channel_scales ? n : n / 32}, + torch::TensorOptions().device(qweight.device()).dtype(torch::kFloat16)); + + const int64_t weight_numel = n * k; + const int weight_blocks = static_cast( + (weight_numel + kQpn8PrepareThreads - 1) / kQpn8PrepareThreads); + fp8_qpn8_prepack_sm70_kernel<<>>( + codes.data_ptr(), + reinterpret_cast(qweight.data_ptr()), static_cast(n), + static_cast(k)); + + const int64_t scale_numel = group_scales.numel(); + const int scale_blocks = static_cast( + (scale_numel + kQpn8PrepareThreads - 1) / kQpn8PrepareThreads); + if (channel_scales) { + fp8_qpn8_channel_scale_sm70_kernel<<>>( + reinterpret_cast(group_scales.data_ptr()), + scales.data_ptr(), static_cast(n)); + } else { + fp8_qpn8_scale_sm70_kernel<<>>( + reinterpret_cast(group_scales.data_ptr()), + scales.data_ptr(), static_cast(n / 128), + static_cast(k / 128)); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {codes, group_scales}; +} + +void fp8_qpn8_dequantize_sm70_out(torch::Tensor out, torch::Tensor codes, + torch::Tensor group_scales) { + TORCH_CHECK(out.is_cuda() && codes.is_cuda() && group_scales.is_cuda(), + "fp8_qpn8_dequantize_sm70_out: tensors must be CUDA tensors"); + TORCH_CHECK(out.scalar_type() == torch::kFloat16 && + codes.scalar_type() == torch::kUInt8 && + group_scales.scalar_type() == torch::kFloat16, + "fp8_qpn8_dequantize_sm70_out: dtype mismatch"); + TORCH_CHECK(out.dim() == 2 && codes.dim() == 2 && group_scales.dim() == 2, + "fp8_qpn8_dequantize_sm70_out: tensors must be 2D"); + TORCH_CHECK(out.is_contiguous() && codes.is_contiguous() && + group_scales.is_contiguous(), + "fp8_qpn8_dequantize_sm70_out: tensors must be contiguous"); + TORCH_CHECK(out.get_device() == codes.get_device() && + out.get_device() == group_scales.get_device(), + "fp8_qpn8_dequantize_sm70_out: tensors must share one device"); + + 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 || block_scales, + "fp8_qpn8_dequantize_sm70_out: scale shape mismatch"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(out)); + const int64_t word_count = k * n / 16; + const int blocks = static_cast((word_count + kQpn8PrepareThreads - 1) / + kQpn8PrepareThreads); + fp8_qpn8_dequantize_sm70_kernel<<>>( + reinterpret_cast(out.data_ptr()), + codes.data_ptr(), + reinterpret_cast(group_scales.data_ptr()), + static_cast(n), static_cast(k), channel_scales); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void fp8_qpn8_prefill_sm70_out(torch::Tensor out, int64_t dense_weight_ptr, + torch::Tensor input, torch::Tensor codes, + torch::Tensor group_scales, bool gated_silu) { + TORCH_CHECK(input.is_cuda() && out.is_cuda(), + "fp8_qpn8_prefill_sm70_out: input and output must be CUDA"); + TORCH_CHECK(input.scalar_type() == torch::kFloat16 && + out.scalar_type() == torch::kFloat16, + "fp8_qpn8_prefill_sm70_out: input and output must be float16"); + TORCH_CHECK(input.dim() == 2 && out.dim() == 2 && dense_weight_ptr != 0, + "fp8_qpn8_prefill_sm70_out: invalid input or workspace"); + TORCH_CHECK(input.is_contiguous() && out.is_contiguous(), + "fp8_qpn8_prefill_sm70_out: tensors must be contiguous"); + + const int64_t m = input.size(0); + const int64_t k = input.size(1); + const int64_t n = codes.size(1); + TORCH_CHECK(codes.dim() == 2 && codes.size(0) == k, + "fp8_qpn8_prefill_sm70_out: code shape mismatch"); + TORCH_CHECK(out.size(0) == m && out.size(1) == (gated_silu ? n / 2 : n), + "fp8_qpn8_prefill_sm70_out: output shape mismatch"); + + auto dense_weight = torch::from_blob( + reinterpret_cast(dense_weight_ptr), {k, n}, input.options()); + fp8_qpn8_dequantize_sm70_out(dense_weight, codes, group_scales); + if (!gated_silu) { + at::mm_out(out, input, dense_weight); + return; + } + + auto gate_up = at::mm(input, dense_weight); + constexpr int kThreads = 256; + fp8_qpn8_silu_and_mul_sm70_kernel<<(m), kThreads, 0, + at::cuda::getCurrentCUDAStream()>>>( + reinterpret_cast(out.data_ptr()), + reinterpret_cast(gate_up.data_ptr()), + static_cast(m), static_cast(n / 2)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void fp8_qpn8_gemm_sm70_out(torch::Tensor out, torch::Tensor input, + torch::Tensor codes, torch::Tensor group_scales, + int64_t split_k, int64_t accumulator_chains, + bool fast_decoder, bool prefetch_codes) { + TORCH_CHECK(input.is_cuda() && out.is_cuda() && codes.is_cuda() && + group_scales.is_cuda(), + "fp8_qpn8_gemm_sm70_out: tensors must be CUDA tensors"); + TORCH_CHECK(input.scalar_type() == torch::kFloat16 && + out.scalar_type() == torch::kFloat16, + "fp8_qpn8_gemm_sm70_out: input and output must be float16"); + TORCH_CHECK(codes.scalar_type() == torch::kUInt8, + "fp8_qpn8_gemm_sm70_out: codes must be uint8"); + TORCH_CHECK(group_scales.scalar_type() == torch::kFloat16, + "fp8_qpn8_gemm_sm70_out: group scales must be float16"); + TORCH_CHECK(input.dim() == 2 && out.dim() == 2 && group_scales.dim() == 2, + "fp8_qpn8_gemm_sm70_out: expected 2D tensors"); + TORCH_CHECK(input.is_contiguous() && out.is_contiguous() && + codes.is_contiguous() && group_scales.is_contiguous(), + "fp8_qpn8_gemm_sm70_out: tensors must be contiguous"); + TORCH_CHECK(input.get_device() == out.get_device() && + input.get_device() == codes.get_device() && + input.get_device() == group_scales.get_device(), + "fp8_qpn8_gemm_sm70_out: tensors must share one device"); + + const int64_t m = input.size(0); + const int64_t k = input.size(1); + const int64_t n = out.size(1); + TORCH_CHECK(m >= 1 && m <= 8, "fp8_qpn8_gemm_sm70_out: M must be in [1, 8]"); + 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 || block_scales, + "fp8_qpn8_gemm_sm70_out: scale shape mismatch"); + TORCH_CHECK(split_k == 4 || split_k == 8 || split_k == 12 || split_k == 16 || + split_k == 32, + "fp8_qpn8_gemm_sm70_out: unsupported split_k"); + TORCH_CHECK((k / 16) % split_k == 0, + "fp8_qpn8_gemm_sm70_out: K/16 must be divisible by split_k"); + TORCH_CHECK(accumulator_chains == 1 || accumulator_chains == 2, + "fp8_qpn8_gemm_sm70_out: accumulator_chains must be 1 or 2"); + TORCH_CHECK(!prefetch_codes || fast_decoder, + "fp8_qpn8_gemm_sm70_out: prefetch experiment requires the " + "fast decoder"); + const bool exact_shape_split = split_k == 12; + TORCH_CHECK(!exact_shape_split || + (accumulator_chains == 2 && fast_decoder && !prefetch_codes), + "fp8_qpn8_gemm_sm70_out: split_k 12 requires the " + "fast decoder, two accumulator chains, and no prefetch"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const auto* code_ptr = codes.data_ptr(); + const auto* scale_ptr = + reinterpret_cast(group_scales.data_ptr()); + const auto* input_ptr = + reinterpret_cast(input.data_ptr()); + auto* output_ptr = reinterpret_cast(out.data_ptr()); + + // Keep the M=1 reduction order restricted to the two tuned output/down + // projections. Extending it to the other split variants changed the frozen + // random-sampling token stream without an end-to-end decode win. + if (m == 1 && accumulator_chains == 2 && fast_decoder && !prefetch_codes && + (split_k == 12 || split_k == 16)) { + if (split_k == 12) { + launch_fp8_qpn8_sm70<12, 2, true, false, true>( + code_ptr, scale_ptr, input_ptr, output_ptr, static_cast(n), + static_cast(k), 1, channel_scales, stream); + } else { + launch_fp8_qpn8_sm70<16, 2, true, false, true>( + code_ptr, scale_ptr, input_ptr, output_ptr, static_cast(n), + static_cast(k), 1, channel_scales, stream); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return; + } + +#define VLLM_LAUNCH_QPN8(SPLIT, NACC, FAST, PREFETCH) \ + launch_fp8_qpn8_sm70( \ + code_ptr, scale_ptr, input_ptr, output_ptr, static_cast(n), \ + static_cast(k), static_cast(m), channel_scales, stream) + + if (prefetch_codes) { + if (split_k == 4 && accumulator_chains == 1) { + VLLM_LAUNCH_QPN8(4, 1, true, true); + } else if (split_k == 4 && accumulator_chains == 2) { + VLLM_LAUNCH_QPN8(4, 2, true, true); + } else if (split_k == 8 && accumulator_chains == 1) { + VLLM_LAUNCH_QPN8(8, 1, true, true); + } else if (split_k == 8 && accumulator_chains == 2) { + VLLM_LAUNCH_QPN8(8, 2, true, true); + } else if (split_k == 16 && accumulator_chains == 1) { + VLLM_LAUNCH_QPN8(16, 1, true, true); + } else if (split_k == 16 && accumulator_chains == 2) { + VLLM_LAUNCH_QPN8(16, 2, true, true); + } else if (split_k == 32 && accumulator_chains == 1) { + VLLM_LAUNCH_QPN8(32, 1, true, true); + } else { + VLLM_LAUNCH_QPN8(32, 2, true, true); + } + } else if (fast_decoder) { + if (split_k == 4 && accumulator_chains == 1) { + VLLM_LAUNCH_QPN8(4, 1, true, false); + } else if (split_k == 4 && accumulator_chains == 2) { + VLLM_LAUNCH_QPN8(4, 2, true, false); + } else if (split_k == 8 && accumulator_chains == 1) { + VLLM_LAUNCH_QPN8(8, 1, true, false); + } else if (split_k == 8 && accumulator_chains == 2) { + VLLM_LAUNCH_QPN8(8, 2, true, false); + } else if (split_k == 12) { + VLLM_LAUNCH_QPN8(12, 2, true, false); + } else if (split_k == 16 && accumulator_chains == 1) { + VLLM_LAUNCH_QPN8(16, 1, true, false); + } else if (split_k == 16 && accumulator_chains == 2) { + VLLM_LAUNCH_QPN8(16, 2, true, false); + } else if (split_k == 32 && accumulator_chains == 1) { + VLLM_LAUNCH_QPN8(32, 1, true, false); + } else { + VLLM_LAUNCH_QPN8(32, 2, true, false); + } + } else if (split_k == 4 && accumulator_chains == 1) { + VLLM_LAUNCH_QPN8(4, 1, false, false); + } else if (split_k == 4 && accumulator_chains == 2) { + VLLM_LAUNCH_QPN8(4, 2, false, false); + } else if (split_k == 8 && accumulator_chains == 1) { + VLLM_LAUNCH_QPN8(8, 1, false, false); + } else if (split_k == 8 && accumulator_chains == 2) { + VLLM_LAUNCH_QPN8(8, 2, false, false); + } else if (split_k == 16 && accumulator_chains == 1) { + VLLM_LAUNCH_QPN8(16, 1, false, false); + } else if (split_k == 16 && accumulator_chains == 2) { + VLLM_LAUNCH_QPN8(16, 2, false, false); + } else if (split_k == 32 && accumulator_chains == 1) { + VLLM_LAUNCH_QPN8(32, 1, false, false); + } else { + VLLM_LAUNCH_QPN8(32, 2, false, false); + } +#undef VLLM_LAUNCH_QPN8 + + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void fp8_qpn8_gemm_ba_split_sm70_out(torch::Tensor qkv_out, torch::Tensor z_out, + torch::Tensor b_out, torch::Tensor a_out, + torch::Tensor input, torch::Tensor codes, + torch::Tensor group_scales, + torch::Tensor ba_weight) { + TORCH_CHECK(qkv_out.is_cuda() && z_out.is_cuda() && b_out.is_cuda() && + a_out.is_cuda() && input.is_cuda() && codes.is_cuda() && + group_scales.is_cuda() && ba_weight.is_cuda(), + "fp8_qpn8_gemm_ba_split_sm70_out: tensors must be CUDA tensors"); + TORCH_CHECK(qkv_out.scalar_type() == torch::kFloat16 && + z_out.scalar_type() == torch::kFloat16 && + b_out.scalar_type() == torch::kFloat16 && + a_out.scalar_type() == torch::kFloat16 && + input.scalar_type() == torch::kFloat16 && + group_scales.scalar_type() == torch::kFloat16 && + ba_weight.scalar_type() == torch::kFloat16 && + codes.scalar_type() == torch::kUInt8, + "fp8_qpn8_gemm_ba_split_sm70_out: dtype mismatch"); + TORCH_CHECK(qkv_out.is_contiguous() && z_out.is_contiguous() && + b_out.is_contiguous() && a_out.is_contiguous() && + input.is_contiguous() && codes.is_contiguous() && + group_scales.is_contiguous() && ba_weight.is_contiguous(), + "fp8_qpn8_gemm_ba_split_sm70_out: tensors must be contiguous"); + TORCH_CHECK(qkv_out.get_device() == z_out.get_device() && + qkv_out.get_device() == b_out.get_device() && + qkv_out.get_device() == a_out.get_device() && + qkv_out.get_device() == input.get_device() && + qkv_out.get_device() == codes.get_device() && + qkv_out.get_device() == group_scales.get_device() && + 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; + 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(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]"); + 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"); + TORCH_CHECK(ba_weight.dim() == 2 && ba_weight.size(0) == ba_n && + ba_weight.size(1) == k, + "fp8_qpn8_gemm_ba_split_sm70_out: b/a weight shape mismatch"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + launch_fp8_qpn8_ba_split_sm70( + codes.data_ptr(), + reinterpret_cast(group_scales.data_ptr()), + reinterpret_cast(input.data_ptr()), + reinterpret_cast(qkv_out.data_ptr()), + reinterpret_cast(z_out.data_ptr()), + reinterpret_cast(ba_weight.data_ptr()), + reinterpret_cast(b_out.data_ptr()), + reinterpret_cast(a_out.data_ptr()), + static_cast(ba_n), static_cast(qkv_n), static_cast(n), + static_cast(k), true, stream); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void fp8_qpn8_dispatch_sm70_out(torch::Tensor out, int64_t dense_weight_ptr, + torch::Tensor input, torch::Tensor codes, + torch::Tensor group_scales, int64_t split_k, + int64_t accumulator_chains, bool prefetch_codes, + bool gated_silu); + +void fp8_qpn8_dispatch_ba_split_sm70_out( + torch::Tensor qkv_out, torch::Tensor z_out, torch::Tensor b_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; + 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, + "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() && + ba_staging.is_cuda() && input.is_cuda() && codes.is_cuda() && + group_scales.is_cuda() && ba_weight.is_cuda(), + "fp8_qpn8_dispatch_ba_split_sm70_out: tensors must be CUDA"); + TORCH_CHECK(qkv_out.scalar_type() == torch::kFloat16 && + z_out.scalar_type() == torch::kFloat16 && + b_out.scalar_type() == torch::kFloat16 && + a_out.scalar_type() == torch::kFloat16 && + qkvz_staging.scalar_type() == torch::kFloat16 && + ba_staging.scalar_type() == torch::kFloat16 && + input.scalar_type() == torch::kFloat16 && + group_scales.scalar_type() == torch::kFloat16 && + ba_weight.scalar_type() == torch::kFloat16 && + codes.scalar_type() == torch::kUInt8, + "fp8_qpn8_dispatch_ba_split_sm70_out: dtype mismatch"); + TORCH_CHECK( + qkv_out.is_contiguous() && z_out.is_contiguous() && + b_out.is_contiguous() && a_out.is_contiguous() && + qkvz_staging.is_contiguous() && ba_staging.is_contiguous() && + input.is_contiguous() && codes.is_contiguous() && + group_scales.is_contiguous() && ba_weight.is_contiguous(), + "fp8_qpn8_dispatch_ba_split_sm70_out: tensors must be contiguous"); + TORCH_CHECK( + qkv_out.numel() == m * qkv_n && z_out.numel() == m * z_n && + b_out.numel() == m * ba_n / 2 && a_out.numel() == m * ba_n / 2 && + qkvz_staging.numel() == m * n && ba_staging.numel() == m * ba_n, + "fp8_qpn8_dispatch_ba_split_sm70_out: output shape mismatch"); + TORCH_CHECK( + qkv_out.get_device() == z_out.get_device() && + qkv_out.get_device() == b_out.get_device() && + qkv_out.get_device() == a_out.get_device() && + qkv_out.get_device() == qkvz_staging.get_device() && + qkv_out.get_device() == ba_staging.get_device() && + qkv_out.get_device() == input.get_device() && + qkv_out.get_device() == codes.get_device() && + qkv_out.get_device() == group_scales.get_device() && + qkv_out.get_device() == ba_weight.get_device(), + "fp8_qpn8_dispatch_ba_split_sm70_out: tensors must share one device"); + TORCH_CHECK(codes.dim() == 2 && codes.size(0) == k && codes.size(1) == n, + "fp8_qpn8_dispatch_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_dispatch_ba_split_sm70_out: scale shape mismatch"); + TORCH_CHECK(ba_weight.dim() == 2 && ba_weight.size(0) == ba_n && + ba_weight.size(1) == k, + "fp8_qpn8_dispatch_ba_split_sm70_out: b/a weight mismatch"); + TORCH_CHECK(m <= 8 || dense_weight_ptr != 0, + "fp8_qpn8_dispatch_ba_split_sm70_out: large-M requires the " + "dense QPN8 workspace"); + + if (m == 1) { + 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 " + "enabled.\n"); + std::fflush(stderr); + }); + fp8_qpn8_gemm_ba_split_sm70_out(qkv_out, z_out, b_out, a_out, input, codes, + group_scales, ba_weight); + return; + } + + fp8_qpn8_dispatch_sm70_out(qkvz_staging, dense_weight_ptr, input, codes, + group_scales, 16, 2, false, false); + at::mm_out(ba_staging, input, ba_weight.t()); + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); + constexpr int threads = 256; + const int64_t elements = m * (n + ba_n); + const int blocks = static_cast((elements + threads - 1) / threads); + fp8_qpn8_ba_split_copy_sm70_kernel<<>>( + reinterpret_cast(qkvz_staging.data_ptr()), + reinterpret_cast(ba_staging.data_ptr()), + reinterpret_cast(qkv_out.data_ptr()), + reinterpret_cast(z_out.data_ptr()), + reinterpret_cast(b_out.data_ptr()), + reinterpret_cast(a_out.data_ptr()), static_cast(m)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void fp8_qpn8_gated_pair_sm70_out(torch::Tensor out, torch::Tensor input, + torch::Tensor codes, + torch::Tensor group_scales, int64_t split_k, + int64_t accumulator_chains, bool fast_decoder, + bool prefetch_codes) { + TORCH_CHECK(out.is_cuda() && input.is_cuda() && codes.is_cuda() && + group_scales.is_cuda(), + "fp8_qpn8_gated_pair_sm70_out: tensors must be CUDA tensors"); + TORCH_CHECK(out.scalar_type() == torch::kFloat16 && + input.scalar_type() == torch::kFloat16 && + group_scales.scalar_type() == torch::kFloat16 && + codes.scalar_type() == torch::kUInt8, + "fp8_qpn8_gated_pair_sm70_out: dtype mismatch"); + TORCH_CHECK(out.is_contiguous() && input.is_contiguous() && + codes.is_contiguous() && group_scales.is_contiguous(), + "fp8_qpn8_gated_pair_sm70_out: tensors must be contiguous"); + TORCH_CHECK(out.dim() == 2 && input.dim() == 2 && group_scales.dim() == 2, + "fp8_qpn8_gated_pair_sm70_out: expected 2D tensors"); + TORCH_CHECK(out.get_device() == input.get_device() && + out.get_device() == codes.get_device() && + out.get_device() == group_scales.get_device(), + "fp8_qpn8_gated_pair_sm70_out: tensors must share one device"); + + const int64_t m = input.size(0); + const int64_t k = input.size(1); + const int64_t hidden = out.size(1); + 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(channel_scales || block_scales, + "fp8_qpn8_gated_pair_sm70_out: scale shape mismatch"); + TORCH_CHECK(split_k == 4 || split_k == 8 || split_k == 16, + "fp8_qpn8_gated_pair_sm70_out: split_k must be 4, 8, or 16"); + TORCH_CHECK((k / 16) % split_k == 0, + "fp8_qpn8_gated_pair_sm70_out: invalid split_k for K"); + TORCH_CHECK( + accumulator_chains == 1 || accumulator_chains == 2, + "fp8_qpn8_gated_pair_sm70_out: accumulator_chains must be 1 or 2"); + TORCH_CHECK(!prefetch_codes || fast_decoder, + "fp8_qpn8_gated_pair_sm70_out: prefetch requires fast decoder"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const auto* code_ptr = codes.data_ptr(); + const auto* scale_ptr = + reinterpret_cast(group_scales.data_ptr()); + const auto* input_ptr = + reinterpret_cast(input.data_ptr()); + auto* output_ptr = reinterpret_cast(out.data_ptr()); + + // The frozen gate/up route is split-8. Preserve the generic reduction order + // for split-4/16 rather than instantiating unaccepted M=1 variants. + if (m == 1 && split_k == 8 && accumulator_chains == 2 && fast_decoder && + !prefetch_codes) { + launch_fp8_qpn8_gated_pair_sm70<8, 2, true, false, true>( + code_ptr, scale_ptr, input_ptr, output_ptr, static_cast(hidden), + static_cast(k), 1, channel_scales, stream); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return; + } + +#define VLLM_LAUNCH_QPN8_GATED_PAIR(SPLIT, NACC, FAST, PREFETCH) \ + launch_fp8_qpn8_gated_pair_sm70( \ + code_ptr, scale_ptr, input_ptr, output_ptr, static_cast(hidden), \ + static_cast(k), static_cast(m), channel_scales, stream) + + if (prefetch_codes) { + if (split_k == 4 && accumulator_chains == 1) { + VLLM_LAUNCH_QPN8_GATED_PAIR(4, 1, true, true); + } else if (split_k == 4 && accumulator_chains == 2) { + VLLM_LAUNCH_QPN8_GATED_PAIR(4, 2, true, true); + } else if (split_k == 8 && accumulator_chains == 1) { + VLLM_LAUNCH_QPN8_GATED_PAIR(8, 1, true, true); + } else if (split_k == 8 && accumulator_chains == 2) { + VLLM_LAUNCH_QPN8_GATED_PAIR(8, 2, true, true); + } else if (split_k == 16 && accumulator_chains == 1) { + VLLM_LAUNCH_QPN8_GATED_PAIR(16, 1, true, true); + } else { + VLLM_LAUNCH_QPN8_GATED_PAIR(16, 2, true, true); + } + } else if (fast_decoder) { + if (split_k == 4 && accumulator_chains == 1) { + VLLM_LAUNCH_QPN8_GATED_PAIR(4, 1, true, false); + } else if (split_k == 4 && accumulator_chains == 2) { + VLLM_LAUNCH_QPN8_GATED_PAIR(4, 2, true, false); + } else if (split_k == 8 && accumulator_chains == 1) { + VLLM_LAUNCH_QPN8_GATED_PAIR(8, 1, true, false); + } else if (split_k == 8 && accumulator_chains == 2) { + VLLM_LAUNCH_QPN8_GATED_PAIR(8, 2, true, false); + } else if (split_k == 16 && accumulator_chains == 1) { + VLLM_LAUNCH_QPN8_GATED_PAIR(16, 1, true, false); + } else { + VLLM_LAUNCH_QPN8_GATED_PAIR(16, 2, true, false); + } + } else if (split_k == 4 && accumulator_chains == 1) { + VLLM_LAUNCH_QPN8_GATED_PAIR(4, 1, false, false); + } else if (split_k == 4 && accumulator_chains == 2) { + VLLM_LAUNCH_QPN8_GATED_PAIR(4, 2, false, false); + } else if (split_k == 8 && accumulator_chains == 1) { + VLLM_LAUNCH_QPN8_GATED_PAIR(8, 1, false, false); + } else if (split_k == 8 && accumulator_chains == 2) { + VLLM_LAUNCH_QPN8_GATED_PAIR(8, 2, false, false); + } else if (split_k == 16 && accumulator_chains == 1) { + VLLM_LAUNCH_QPN8_GATED_PAIR(16, 1, false, false); + } else { + VLLM_LAUNCH_QPN8_GATED_PAIR(16, 2, false, false); + } +#undef VLLM_LAUNCH_QPN8_GATED_PAIR + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void fp8_qpn8_dispatch_sm70_out(torch::Tensor out, int64_t dense_weight_ptr, + torch::Tensor input, torch::Tensor codes, + torch::Tensor group_scales, int64_t split_k, + int64_t accumulator_chains, bool prefetch_codes, + bool gated_silu) { + // Keep the M decision inside the opaque operator. AOTInductor compiles one + // dynamic M=1..8192 range for this model, so a Python shape branch traced at + // M=1/2 would otherwise be incorrectly reused by large-M prefill. + if (input.size(0) <= 8) { + if (gated_silu) { + fp8_qpn8_gated_pair_sm70_out(out, input, codes, group_scales, split_k, + accumulator_chains, true, prefetch_codes); + } else { + fp8_qpn8_gemm_sm70_out(out, input, codes, group_scales, split_k, + accumulator_chains, true, prefetch_codes); + } + return; + } + fp8_qpn8_prefill_sm70_out(out, dense_weight_ptr, input, codes, group_scales, + gated_silu); +} + +#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. +TORCH_LIBRARY_FRAGMENT(_C, ops) { + 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( + "fp8_qpn8_dequantize_sm70_out(Tensor(a!) out, Tensor codes, " + "Tensor group_scales) -> ()"); + ops.impl("fp8_qpn8_dequantize_sm70_out", torch::kCUDA, + &fp8_qpn8_dequantize_sm70_out); + ops.def( + "fp8_qpn8_prefill_sm70_out(Tensor(a!) out, int dense_weight_ptr, " + "Tensor input, Tensor codes, Tensor group_scales, bool gated_silu) -> " + "()"); + ops.impl("fp8_qpn8_prefill_sm70_out", torch::kCUDA, + &fp8_qpn8_prefill_sm70_out); + ops.def( + "fp8_qpn8_dispatch_sm70_out(Tensor(a!) out, int dense_weight_ptr, " + "Tensor input, Tensor codes, Tensor group_scales, int split_k, " + "int accumulator_chains, bool prefetch_codes, bool gated_silu) -> ()"); + ops.impl("fp8_qpn8_dispatch_sm70_out", torch::kCUDA, + &fp8_qpn8_dispatch_sm70_out); + ops.def( + "fp8_qpn8_gemm_sm70_out(Tensor(a!) out, Tensor input, Tensor codes, " + "Tensor group_scales, int split_k, int accumulator_chains, " + "bool fast_decoder, bool prefetch_codes) -> ()"); + ops.impl("fp8_qpn8_gemm_sm70_out", torch::kCUDA, &fp8_qpn8_gemm_sm70_out); + ops.def( + "fp8_qpn8_gemm_ba_split_sm70_out(Tensor(a!) qkv_out, Tensor(b!) " + "z_out, Tensor(c!) b_out, Tensor(d!) a_out, Tensor input, Tensor codes, " + "Tensor group_scales, Tensor ba_weight) -> ()"); + ops.impl("fp8_qpn8_gemm_ba_split_sm70_out", torch::kCUDA, + &fp8_qpn8_gemm_ba_split_sm70_out); + ops.def( + "fp8_qpn8_dispatch_ba_split_sm70_out(Tensor(a!) qkv_out, Tensor(b!) " + "z_out, Tensor(c!) b_out, Tensor(d!) a_out, Tensor(e!) qkvz_staging, " + "Tensor(f!) ba_staging, int dense_weight_ptr, Tensor input, Tensor " + "codes, Tensor group_scales, Tensor ba_weight) -> ()"); + ops.impl("fp8_qpn8_dispatch_ba_split_sm70_out", torch::kCUDA, + &fp8_qpn8_dispatch_ba_split_sm70_out); + ops.def( + "fp8_qpn8_gated_pair_sm70_out(Tensor(a!) out, Tensor input, " + "Tensor codes, Tensor group_scales, int split_k, " + "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); +} +#endif diff --git a/csrc/sm70_turbomind/ops/nvfp4_qpn2_sm70.cu b/csrc/sm70_turbomind/ops/nvfp4_qpn2_sm70.cu new file mode 100644 index 0000000000..641949e4b4 --- /dev/null +++ b/csrc/sm70_turbomind/ops/nvfp4_qpn2_sm70.cu @@ -0,0 +1,546 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// The QPN2 execution layout is derived from dnv2003/v100-skinny (MIT). +// See LICENSE.v100-skinny in this directory for the retained MIT notice. + +#include +#include + +#include +#include +#include +#include +#include + +#ifndef VLLM_NVFP4_QPN2_STANDALONE +void silu_and_mul(torch::Tensor& out, torch::Tensor& input); + +void nvfp4_gemm_sm70_out(torch::Tensor out, torch::Tensor input, + torch::Tensor weight, torch::Tensor scales, + int64_t group_size, int64_t k_ld, int64_t q_ld, + bool gated_silu); +#endif + +namespace { + +constexpr int kPrepareThreads = 256; + +__device__ __forceinline__ int qpn2_col_from_lane(int lane) { + return ((lane >> 2) & 3) * 8 + (lane & 3) + ((lane & 16) ? 4 : 0); +} + +__device__ __forceinline__ int qpn2_logical_k(int physical_k) { + const int local = physical_k & 7; + return (physical_k & 8) + ((local & 3) << 1) + (local >> 2); +} + +__global__ void nvfp4_qpn2_prepack_codes_kernel( + uint8_t* __restrict__ output, const uint8_t* __restrict__ weight, int n, + int k) { + const size_t index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const size_t numel = static_cast(n) * k / 2; + if (index >= numel) { + return; + } + + const int physical_byte = static_cast(index & 7); + size_t outer = index >> 3; + const int lane = static_cast(outer & 31); + outer >>= 5; + const int groups_k16 = k >> 4; + const int group = static_cast(outer % groups_k16); + const int tile = static_cast(outer / groups_k16); + const int column = tile * 32 + qpn2_col_from_lane(lane); + const int logical_k0 = group * 16 + qpn2_logical_k(physical_byte * 2); + const int logical_k1 = group * 16 + qpn2_logical_k(physical_byte * 2 + 1); + const int k_bytes = k >> 1; + const uint8_t packed0 = + weight[static_cast(column) * k_bytes + (logical_k0 >> 1)]; + const uint8_t packed1 = + weight[static_cast(column) * k_bytes + (logical_k1 >> 1)]; + const uint8_t code0 = + static_cast((packed0 >> ((logical_k0 & 1) * 4)) & 0x0f); + const uint8_t code1 = + static_cast((packed1 >> ((logical_k1 & 1) * 4)) & 0x0f); + output[index] = static_cast(code0 | (code1 << 4)); +} + +__global__ void nvfp4_qpn2_prepack_scales_kernel( + uint8_t* __restrict__ output, const uint8_t* __restrict__ scales, int n, + int k) { + const size_t index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const size_t numel = static_cast(n) * k / 16; + if (index >= numel) { + return; + } + + const int lane = static_cast(index & 31); + size_t outer = index >> 5; + const int groups_k16 = k >> 4; + const int group = static_cast(outer % groups_k16); + const int tile = static_cast(outer / groups_k16); + const int column = tile * 32 + qpn2_col_from_lane(lane); + output[index] = scales[static_cast(column) * groups_k16 + group]; +} + +__device__ __forceinline__ half2 fp8e4m3_to_half2(uint8_t value) { + const unsigned short bits = + ((static_cast(value) & 0x80u) << 8) | + ((static_cast(value) & 0x7fu) << 7); + const half converted = + __hmul(__ushort_as_half(bits), __ushort_as_half(0x5c00)); + return __halves2half2(converted, converted); +} + +__device__ __forceinline__ void dequant_e2m1x8(unsigned packed, half2 scale, + half2 output[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) { + output[index] = __hmul2(*reinterpret_cast(&values[index]), scale); + } +} + +#define VLLM_SM70_QPN2_MMA(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 nvfp4_qpn2_sm70_kernel(const uint8_t* __restrict__ codes, + const uint8_t* __restrict__ group_scales, + const half* __restrict__ input, + half* __restrict__ output, int n, int k, + int m, float global_scale) { + __shared__ float partials[SplitK][256]; + + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + const int tile = blockIdx.x; + 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_warp = groups_k16 / SplitK; + const int group_begin = warp * groups_per_warp; + const uint2* code_ptr = reinterpret_cast(codes) + + static_cast(tile) * groups_k16 * 32 + lane; + const uint8_t* scale_ptr = + group_scales + static_cast(tile) * groups_k16 * 32 + lane; + const half2 global_scale2 = __float2half2_rn(global_scale * 16384.0f); + + 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_warp; + ++group) { + const uint2 packed = __ldcs(code_ptr + static_cast(group) * 32); + const half2 scale = __hmul2( + fp8e4m3_to_half2(__ldg(scale_ptr + static_cast(group) * 32)), + global_scale2); + half2 weights[8]; + dequant_e2m1x8(packed.x, scale, weights); + dequant_e2m1x8(packed.y, scale, weights + 4); + + uint4 input01 = make_uint4(0, 0, 0, 0); + uint4 input23 = make_uint4(0, 0, 0, 0); + if (row < m) { + const half* input_row = input + static_cast(row) * k; + 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); + const unsigned* b = reinterpret_cast(weights); + VLLM_SM70_QPN2_MMA(accum[0], a0[0], a0[1], b[0], b[1]); + VLLM_SM70_QPN2_MMA(accum[1 % NAcc], a0[2], a0[3], b[2], b[3]); + VLLM_SM70_QPN2_MMA(accum[2 % NAcc], a1[0], a1[1], b[4], b[5]); + VLLM_SM70_QPN2_MMA(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]; + } + } + +#pragma unroll + for (int index = 0; index < 8; ++index) { + const int output_row = (index & 2) | ((lane & 16) ? 4 : 0) | (lane & 1); + const int output_col = + (index & 1) | (((lane >> 1) & 1) << 1) | ((index >> 2) << 2); + partials[warp][output_row * 32 + quadpair * 8 + output_col] = + accum[0][index]; + } + __syncthreads(); + + for (int element = threadIdx.x; element < 256; element += blockDim.x) { + float value = 0.0f; +#pragma unroll + for (int k_warp = 0; k_warp < SplitK; ++k_warp) { + value += partials[k_warp][element]; + } + const int output_row = element >> 5; + const int output_col = element & 31; + if (output_row < m) { + output[static_cast(output_row) * n + tile * 32 + output_col] = + __float2half(value); + } + } +} + +template +__global__ void nvfp4_qpn2_gated_sm70_kernel( + const uint8_t* __restrict__ codes, const uint8_t* __restrict__ group_scales, + const half* __restrict__ input, half* __restrict__ output, int hidden, + int k, int m, float global_scale) { + __shared__ float partials[2][SplitK][256]; + + const int lane = threadIdx.x & 31; + const int warp_in_block = threadIdx.x >> 5; + const int projection = warp_in_block / SplitK; + const int warp = warp_in_block - projection * SplitK; + const int hidden_tiles = hidden >> 5; + const int tile = blockIdx.x + projection * hidden_tiles; + 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_warp = groups_k16 / SplitK; + const int group_begin = warp * groups_per_warp; + const uint2* code_ptr = reinterpret_cast(codes) + + static_cast(tile) * groups_k16 * 32 + lane; + const uint8_t* scale_ptr = + group_scales + static_cast(tile) * groups_k16 * 32 + lane; + const half2 global_scale2 = __float2half2_rn(global_scale * 16384.0f); + + 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_warp; + ++group) { + const uint2 packed = __ldcs(code_ptr + static_cast(group) * 32); + const half2 scale = __hmul2( + fp8e4m3_to_half2(__ldg(scale_ptr + static_cast(group) * 32)), + global_scale2); + half2 weights[8]; + dequant_e2m1x8(packed.x, scale, weights); + dequant_e2m1x8(packed.y, scale, weights + 4); + + uint4 input01 = make_uint4(0, 0, 0, 0); + uint4 input23 = make_uint4(0, 0, 0, 0); + if (row < m) { + const half* input_row = input + static_cast(row) * k; + 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); + const unsigned* b = reinterpret_cast(weights); + VLLM_SM70_QPN2_MMA(accum[0], a0[0], a0[1], b[0], b[1]); + VLLM_SM70_QPN2_MMA(accum[1 % NAcc], a0[2], a0[3], b[2], b[3]); + VLLM_SM70_QPN2_MMA(accum[2 % NAcc], a1[0], a1[1], b[4], b[5]); + VLLM_SM70_QPN2_MMA(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]; + } + } + +#pragma unroll + for (int index = 0; index < 8; ++index) { + const int output_row = (index & 2) | ((lane & 16) ? 4 : 0) | (lane & 1); + const int output_col = + (index & 1) | (((lane >> 1) & 1) << 1) | ((index >> 2) << 2); + partials[projection][warp][output_row * 32 + quadpair * 8 + output_col] = + accum[0][index]; + } + __syncthreads(); + + for (int element = threadIdx.x; element < 256; element += blockDim.x) { + float gate = 0.0f; + float up = 0.0f; +#pragma unroll + for (int k_warp = 0; k_warp < SplitK; ++k_warp) { + gate += partials[0][k_warp][element]; + up += partials[1][k_warp][element]; + } + const int output_row = element >> 5; + const int output_col = element & 31; + if (output_row < m) { + // Match the existing SM70 silu_and_mul contract: round both GEMM + // outputs to FP16 before the activation, round SiLU to FP16, then use + // FP16 multiplication. This makes fusion numerically equivalent to + // QPN2 GEMM followed by the current activation kernel. + const half gate_half = __float2half(gate); + const half up_half = __float2half(up); + const float gate_float = __half2float(gate_half); + const half silu = __float2half(gate_float / (1.0f + expf(-gate_float))); + output[static_cast(output_row) * hidden + blockIdx.x * 32 + + output_col] = __hmul(silu, up_half); + } + } +} + +template +void launch_qpn2(const uint8_t* codes, const uint8_t* scales, const half* input, + half* output, int n, int k, int m, float global_scale, + cudaStream_t stream) { + nvfp4_qpn2_sm70_kernel<<<(n / 32), (32 * SplitK), 0, stream>>>( + codes, scales, input, output, n, k, m, global_scale); +} + +template +void launch_qpn2_gated(const uint8_t* codes, const uint8_t* scales, + const half* input, half* output, int hidden, int k, + int m, float global_scale, cudaStream_t stream) { + nvfp4_qpn2_gated_sm70_kernel + <<<(hidden / 32), (64 * SplitK), 0, stream>>>( + codes, scales, input, output, hidden, k, m, global_scale); +} + +void check_qpn2_tensors(const torch::Tensor& out, const torch::Tensor& input, + const torch::Tensor& codes, const torch::Tensor& scales, + bool gated_silu) { + TORCH_CHECK( + out.is_cuda() && input.is_cuda() && codes.is_cuda() && scales.is_cuda(), + "NVFP4 QPN2 tensors must be CUDA tensors"); + TORCH_CHECK(out.scalar_type() == torch::kFloat16 && + input.scalar_type() == torch::kFloat16 && + codes.scalar_type() == torch::kUInt8 && + scales.scalar_type() == torch::kUInt8, + "NVFP4 QPN2 dtype mismatch"); + TORCH_CHECK(out.is_contiguous() && input.is_contiguous() && + codes.is_contiguous() && scales.is_contiguous(), + "NVFP4 QPN2 tensors must be contiguous"); + TORCH_CHECK(out.dim() == 2 && input.dim() == 2, + "NVFP4 QPN2 input and output must be matrices"); + TORCH_CHECK(out.get_device() == input.get_device() && + out.get_device() == codes.get_device() && + out.get_device() == scales.get_device(), + "NVFP4 QPN2 tensors must share one device"); + const int64_t m = input.size(0); + const int64_t k = input.size(1); + const int64_t n = gated_silu ? out.size(1) * 2 : out.size(1); + TORCH_CHECK(m >= 1 && m <= 8 && out.size(0) == m, + "NVFP4 QPN2 requires M in [1, 8]"); + TORCH_CHECK(k > 0 && k % 64 == 0 && n > 0 && n % 32 == 0, + "NVFP4 QPN2 shape alignment mismatch"); + TORCH_CHECK(codes.numel() == n * k / 2 && scales.numel() == n * k / 16, + "NVFP4 QPN2 packed tensor size mismatch"); +} + +} // namespace + +std::vector nvfp4_qpn2_prepare_sm70(torch::Tensor weight_packed, + torch::Tensor weight_scale) { + TORCH_CHECK(weight_packed.is_cuda() && weight_scale.is_cuda(), + "nvfp4_qpn2_prepare_sm70 expects CUDA tensors"); + TORCH_CHECK(weight_packed.scalar_type() == torch::kUInt8 && + weight_scale.scalar_type() == at::ScalarType::Float8_e4m3fn, + "nvfp4_qpn2_prepare_sm70 expects uint8 weights and E4M3 scales"); + TORCH_CHECK(weight_packed.dim() == 2 && weight_scale.dim() == 2 && + weight_packed.is_contiguous() && weight_scale.is_contiguous(), + "nvfp4_qpn2_prepare_sm70 expects contiguous matrices"); + TORCH_CHECK(weight_packed.get_device() == weight_scale.get_device(), + "nvfp4_qpn2_prepare_sm70 tensors must share one device"); + + const int64_t n = weight_packed.size(0); + const int64_t k = weight_packed.size(1) * 2; + TORCH_CHECK(n > 0 && n % 32 == 0 && k > 0 && k % 64 == 0, + "nvfp4_qpn2_prepare_sm70 shape alignment mismatch"); + TORCH_CHECK(weight_scale.size(0) == n && weight_scale.size(1) == k / 16, + "nvfp4_qpn2_prepare_sm70 scale shape mismatch"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(weight_packed)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + auto codes = torch::empty_like(weight_packed); + auto scales = torch::empty({n, k / 16}, torch::TensorOptions() + .device(weight_packed.device()) + .dtype(torch::kUInt8)); + + const int code_blocks = + static_cast((codes.numel() + kPrepareThreads - 1) / kPrepareThreads); + nvfp4_qpn2_prepack_codes_kernel<<>>( + codes.data_ptr(), weight_packed.data_ptr(), + static_cast(n), static_cast(k)); + const int scale_blocks = static_cast( + (scales.numel() + kPrepareThreads - 1) / kPrepareThreads); + nvfp4_qpn2_prepack_scales_kernel<<>>( + scales.data_ptr(), + reinterpret_cast(weight_scale.data_ptr()), + static_cast(n), static_cast(k)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {codes, scales}; +} + +void nvfp4_qpn2_gemm_sm70_out(torch::Tensor out, torch::Tensor input, + torch::Tensor codes, torch::Tensor scales, + double global_scale, int64_t split_k, + int64_t accumulator_chains) { + check_qpn2_tensors(out, input, codes, scales, false); + TORCH_CHECK(split_k == 8 || split_k == 16 || split_k == 32, + "NVFP4 QPN2 split_k must be 8, 16, or 32"); + TORCH_CHECK((input.size(1) / 16) % split_k == 0, + "NVFP4 QPN2 K/16 must be divisible by split_k"); + TORCH_CHECK(accumulator_chains == 1 || accumulator_chains == 2, + "NVFP4 QPN2 accumulator_chains must be 1 or 2"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const auto* code_ptr = codes.data_ptr(); + const auto* scale_ptr = scales.data_ptr(); + const auto* input_ptr = + reinterpret_cast(input.data_ptr()); + auto* output_ptr = reinterpret_cast(out.data_ptr()); + const int n = static_cast(out.size(1)); + const int k = static_cast(input.size(1)); + const int m = static_cast(input.size(0)); + +#define VLLM_LAUNCH_QPN2(SPLIT, NACC) \ + launch_qpn2(code_ptr, scale_ptr, input_ptr, output_ptr, n, k, \ + m, static_cast(global_scale), stream) + if (split_k == 8 && accumulator_chains == 1) { + VLLM_LAUNCH_QPN2(8, 1); + } else if (split_k == 8) { + VLLM_LAUNCH_QPN2(8, 2); + } else if (split_k == 16 && accumulator_chains == 1) { + VLLM_LAUNCH_QPN2(16, 1); + } else if (split_k == 16) { + VLLM_LAUNCH_QPN2(16, 2); + } else if (accumulator_chains == 1) { + VLLM_LAUNCH_QPN2(32, 1); + } else { + VLLM_LAUNCH_QPN2(32, 2); + } +#undef VLLM_LAUNCH_QPN2 + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void nvfp4_qpn2_gated_sm70_out(torch::Tensor out, torch::Tensor input, + torch::Tensor codes, torch::Tensor scales, + double global_scale, int64_t split_k, + int64_t accumulator_chains) { + check_qpn2_tensors(out, input, codes, scales, true); + TORCH_CHECK(split_k == 8 || split_k == 16, + "NVFP4 QPN2 gated split_k must be 8 or 16"); + TORCH_CHECK((input.size(1) / 16) % split_k == 0, + "NVFP4 QPN2 gated K/16 must be divisible by split_k"); + TORCH_CHECK(accumulator_chains == 1 || accumulator_chains == 2, + "NVFP4 QPN2 gated accumulator_chains must be 1 or 2"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const auto* code_ptr = codes.data_ptr(); + const auto* scale_ptr = scales.data_ptr(); + const auto* input_ptr = + reinterpret_cast(input.data_ptr()); + auto* output_ptr = reinterpret_cast(out.data_ptr()); + const int hidden = static_cast(out.size(1)); + const int k = static_cast(input.size(1)); + const int m = static_cast(input.size(0)); + +#define VLLM_LAUNCH_QPN2_GATED(SPLIT, NACC) \ + launch_qpn2_gated(code_ptr, scale_ptr, input_ptr, output_ptr, \ + hidden, k, m, \ + static_cast(global_scale), stream) + if (split_k == 8 && accumulator_chains == 1) { + VLLM_LAUNCH_QPN2_GATED(8, 1); + } else if (split_k == 8) { + VLLM_LAUNCH_QPN2_GATED(8, 2); + } else if (accumulator_chains == 1) { + VLLM_LAUNCH_QPN2_GATED(16, 1); + } else { + VLLM_LAUNCH_QPN2_GATED(16, 2); + } +#undef VLLM_LAUNCH_QPN2_GATED + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +#ifndef VLLM_NVFP4_QPN2_STANDALONE +void nvfp4_qpn2_dispatch_sm70_out(torch::Tensor out, torch::Tensor input, + torch::Tensor codes, torch::Tensor scales, + double global_scale, int64_t split_k, + int64_t accumulator_chains, + torch::Tensor tm_weight, + torch::Tensor tm_scales, + int64_t tm_group_size, int64_t tm_k_ld, + int64_t tm_q_ld, bool gated_silu) { + if (input.size(0) <= 8) { + if (gated_silu) { + nvfp4_qpn2_gated_sm70_out(out, input, codes, scales, global_scale, + split_k, accumulator_chains); + } else { + nvfp4_qpn2_gemm_sm70_out(out, input, codes, scales, global_scale, split_k, + accumulator_chains); + } + return; + } + + if (!gated_silu) { + nvfp4_gemm_sm70_out(out, input, tm_weight, tm_scales, tm_group_size, + tm_k_ld, tm_q_ld, false); + return; + } + auto gate_up = + torch::empty({input.size(0), out.size(1) * 2}, input.options()); + nvfp4_gemm_sm70_out(gate_up, input, tm_weight, tm_scales, tm_group_size, + tm_k_ld, tm_q_ld, false); + silu_and_mul(out, gate_up); +} +#endif + +#ifdef VLLM_NVFP4_QPN2_STANDALONE +// Compile the exact production kernels as a task-local operator-race library +// before paying for a complete vLLM rebuild. Production registers these +// operators centrally in torch_bindings.cpp. +TORCH_LIBRARY_FRAGMENT(_C, ops) { + ops.def( + "nvfp4_qpn2_prepare_sm70(Tensor weight_packed, Tensor weight_scale) -> " + "Tensor[]"); + ops.impl("nvfp4_qpn2_prepare_sm70", torch::kCUDA, &nvfp4_qpn2_prepare_sm70); + ops.def( + "nvfp4_qpn2_gemm_sm70_out(Tensor(a!) out, Tensor input, Tensor codes, " + "Tensor scales, float global_scale, int split_k, " + "int accumulator_chains) -> ()"); + ops.impl("nvfp4_qpn2_gemm_sm70_out", torch::kCUDA, &nvfp4_qpn2_gemm_sm70_out); + ops.def( + "nvfp4_qpn2_gated_sm70_out(Tensor(a!) out, Tensor input, Tensor codes, " + "Tensor scales, float global_scale, int split_k, " + "int accumulator_chains) -> ()"); + ops.impl("nvfp4_qpn2_gated_sm70_out", torch::kCUDA, + &nvfp4_qpn2_gated_sm70_out); +} +#endif diff --git a/csrc/sm70_turbomind/ops/nvfp4_qpn4_sm70.cu b/csrc/sm70_turbomind/ops/nvfp4_qpn4_sm70.cu new file mode 100644 index 0000000000..0395289403 --- /dev/null +++ b/csrc/sm70_turbomind/ops/nvfp4_qpn4_sm70.cu @@ -0,0 +1,1029 @@ +// SPDX-License-Identifier: Apache-2.0 +// Qwen3.8 NVFP4 M=1 decode kernels for NVIDIA Volta SM70. + +#include +#include + +#include +#include +#include +#include + +namespace { + +constexpr int kPrepareThreads = 256; +constexpr float kFp4Bias = 16384.0f; + +struct SplitHalfScale { + half hi; + half lo; +}; + +SplitHalfScale split_half_scale(float value) { + const half hi = __float2half_rn(value); + const half lo = __float2half_rn(value - __half2float(hi)); + return {hi, lo}; +} + +__device__ __forceinline__ int qpn_col_from_lane(int lane) { + return ((lane >> 2) & 3) * 8 + (lane & 3) + ((lane & 16) ? 4 : 0); +} + +__device__ __forceinline__ int logical_k_from_physical(int physical) { + const int local = physical & 7; + return (physical & 8) + ((local & 3) << 1) + (local >> 2); +} + +__global__ void nvfp4_qpn4_prepack_codes_kernel( + uint8_t* __restrict__ codes, const uint8_t* __restrict__ qweight, int k, + int n) { + const size_t index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const size_t packed_numel = static_cast(k) * n / 2; + if (index >= packed_numel) { + return; + } + + const int physical_byte = static_cast(index & 7); + size_t outer = index >> 3; + const int lane = static_cast(outer & 31); + outer >>= 5; + const int groups_k16 = k >> 4; + const int group = static_cast(outer % groups_k16); + const int tile = static_cast(outer / groups_k16); + const int col = tile * 32 + qpn_col_from_lane(lane); + const int physical_k0 = physical_byte * 2; + const int logical_k0 = logical_k_from_physical(physical_k0); + const int logical_k1 = logical_k_from_physical(physical_k0 + 1); + const uint8_t lo = + qweight[static_cast(group * 16 + logical_k0) * n + col] & 0x0fU; + const uint8_t hi = + qweight[static_cast(group * 16 + logical_k1) * n + col] & 0x0fU; + codes[index] = lo | (hi << 4); +} + +__global__ void nvfp4_qpn4_prepack_scales_kernel( + half* __restrict__ packed_scales, const half* __restrict__ scales, int k, + int n) { + const size_t index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const size_t numel = static_cast(k / 16) * n; + if (index >= numel) { + return; + } + const int lane = static_cast(index & 31); + size_t outer = index >> 5; + const int groups_k16 = k >> 4; + const int group = static_cast(outer % groups_k16); + const int tile = static_cast(outer / groups_k16); + const int col = tile * 32 + qpn_col_from_lane(lane); + packed_scales[index] = __hmul(scales[static_cast(group) * n + col], + __float2half_rn(kFp4Bias)); +} + +__global__ void nvfp4_qpn4_prepack_scale_codes_kernel( + uint8_t* __restrict__ packed_scale_codes, + const uint8_t* __restrict__ scale_codes, int k, int n) { + const size_t index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const size_t numel = static_cast(k / 16) * n; + if (index >= numel) { + return; + } + const int lane = static_cast(index & 31); + size_t outer = index >> 5; + const int groups_k16 = k >> 4; + const int group = static_cast(outer % groups_k16); + const int tile = static_cast(outer / groups_k16); + const int col = tile * 32 + qpn_col_from_lane(lane); + packed_scale_codes[index] = scale_codes[static_cast(group) * n + col]; +} + +__device__ __forceinline__ void fp4x8_to_half2x4(unsigned x, half2 out[4]) { + constexpr unsigned kSign = 0x80008000U; + constexpr unsigned kExponentMantissa = 0x0E000E00U; + unsigned values[4]; + values[0] = ((x << 12) & kSign) | ((x << 9) & kExponentMantissa); + values[1] = ((x << 8) & kSign) | ((x << 5) & kExponentMantissa); + values[2] = ((x << 4) & kSign) | ((x << 1) & kExponentMantissa); + values[3] = ((x << 0) & kSign) | ((x >> 3) & kExponentMantissa); +#pragma unroll + for (int i = 0; i < 4; ++i) { + out[i] = *reinterpret_cast(&values[i]); + } +} + +__device__ __forceinline__ void fp4x16_to_half2x8(uint2 packed, half2 out[8]) { + fp4x8_to_half2x4(packed.x, out); + fp4x8_to_half2x4(packed.y, out + 4); +} + +__device__ __forceinline__ half nvfp4_scale_code_to_half(uint8_t scale_code, + half scale_hi, + half scale_lo) { + // NVFP4 group scales are positive, normal E4M3FN values. Adding the bias + // delta while shifting the E4M3 payload produces the exact FP16 value. + const unsigned half_bits = + (static_cast(scale_code) << 7U) + 0x2000U; + const half raw_scale = *reinterpret_cast(&half_bits); + const half correction = __hmul(raw_scale, scale_lo); + return __hfma(raw_scale, scale_hi, correction); +} + +template +__global__ void nvfp4_qpn4_dequantize_sm70_kernel( + half* __restrict__ output, const uint8_t* __restrict__ codes, + const void* __restrict__ packed_scales, half scale_hi, half scale_lo, int n, + int k) { + const size_t word_index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const size_t word_count = static_cast(k) * n / 16; + if (word_index >= word_count) { + return; + } + + const int lane = static_cast(word_index & 31); + size_t outer = word_index >> 5; + const int groups_k16 = k >> 4; + const int group = static_cast(outer % groups_k16); + const int tile = static_cast(outer / groups_k16); + const int col = tile * 32 + qpn_col_from_lane(lane); + half scale; + if constexpr (UseScaleCode) { + scale = nvfp4_scale_code_to_half( + reinterpret_cast(packed_scales)[word_index], scale_hi, + scale_lo); + } else { + scale = reinterpret_cast(packed_scales)[word_index]; + } + half2 weights[8]; + fp4x16_to_half2x8(reinterpret_cast(codes)[word_index], weights); + const half2 scale2 = __halves2half2(scale, scale); +#pragma unroll + for (int pair = 0; pair < 8; ++pair) { + const half2 value = __hmul2(weights[pair], scale2); + // fp4x8_to_half2x4 pairs nibbles (0,4), (1,5), (2,6), (3,7) + // within each 8-value word rather than adjacent physical K values. + const int physical_k0 = (pair & 3) + ((pair & 4) ? 8 : 0); + const int logical_k0 = logical_k_from_physical(physical_k0); + const int logical_k1 = logical_k_from_physical(physical_k0 + 4); + output[static_cast(group * 16 + logical_k0) * n + col] = + __low2half(value); + output[static_cast(group * 16 + logical_k1) * n + col] = + __high2half(value); + } +} + +__global__ void nvfp4_qpn4_silu_and_mul_sm70_kernel( + half* __restrict__ output, const half* __restrict__ gate_up, int rows, + int hidden) { + const int row = blockIdx.x; + if (row >= rows) { + return; + } + const half* row_input = gate_up + static_cast(row) * hidden * 2; + half* row_output = output + static_cast(row) * hidden; + for (int col = threadIdx.x; col < hidden; col += blockDim.x) { + const float gate = __half2float(row_input[col]); + const float silu = gate / (1.0f + __expf(-gate)); + row_output[col] = __hmul(__float2half(silu), row_input[hidden + col]); + } +} + +#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 nvfp4_qpn4_sm70_kernel(const uint8_t* __restrict__ codes, + const void* __restrict__ packed_scales, + half scale_hi, half scale_lo, + const half* __restrict__ input, + half* __restrict__ output, int n, int k, + int m) { + __shared__ float partials[SplitK][32]; + + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + const int tile = blockIdx.x; + 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_warp = groups_k16 / SplitK; + const int group_begin = warp * groups_per_warp; + const uint2* code_ptr = reinterpret_cast(codes) + + static_cast(tile) * groups_k16 * 32 + lane; + const size_t scale_offset = + static_cast(tile) * groups_k16 * 32 + lane; + const half* scale_ptr = + reinterpret_cast(packed_scales) + scale_offset; + const uint8_t* scale_code_ptr = + reinterpret_cast(packed_scales) + scale_offset; + + float accum[NAcc][8]; +#pragma unroll + for (int chain = 0; chain < NAcc; ++chain) { +#pragma unroll + for (int i = 0; i < 8; ++i) { + accum[chain][i] = 0.0f; + } + } + + uint2 prefetched = make_uint2(0, 0); + if constexpr (PrefetchCodes) { + prefetched = __ldcs(code_ptr + static_cast(group_begin) * 32); + } + +#pragma unroll 4 + for (int group = group_begin; group < group_begin + groups_per_warp; + ++group) { + const uint2 packed = + PrefetchCodes ? prefetched + : __ldcs(code_ptr + static_cast(group) * 32); + uint2 next = make_uint2(0, 0); + if constexpr (PrefetchCodes) { + if (group + 1 < group_begin + groups_per_warp) { + next = __ldcs(code_ptr + static_cast(group + 1) * 32); + } + } + half2 weights[8]; + fp4x16_to_half2x8(packed, weights); + half scale; + if constexpr (UseScaleCode) { + const uint8_t scale_code = + __ldg(scale_code_ptr + static_cast(group) * 32); + scale = nvfp4_scale_code_to_half(scale_code, scale_hi, scale_lo); + } else { + scale = __ldg(scale_ptr + static_cast(group) * 32); + } + const half2 scale2 = __halves2half2(scale, scale); +#pragma unroll + for (int i = 0; i < 8; ++i) { + weights[i] = __hmul2(weights[i], scale2); + } + + uint4 input01 = make_uint4(0, 0, 0, 0); + uint4 input23 = make_uint4(0, 0, 0, 0); + if (row < m) { + const half* input_row = input + static_cast(row) * k; + 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); + 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]); + if constexpr (PrefetchCodes) { + prefetched = next; + } + } + +#pragma unroll + for (int chain = 1; chain < NAcc; ++chain) { +#pragma unroll + for (int i = 0; i < 8; ++i) { + accum[0][i] += accum[chain][i]; + } + } + if ((lane & 17) == 0) { +#pragma unroll + for (int pair = 0; pair < 2; ++pair) { +#pragma unroll + for (int offset = 0; offset < 2; ++offset) { + const int i = pair * 4 + offset; + const int output_col = offset | (((lane >> 1) & 1) << 1) | (pair << 2); + partials[warp][quadpair * 8 + output_col] = accum[0][i]; + } + } + } + __syncthreads(); + + for (int element = threadIdx.x; element < 32; element += blockDim.x) { + float value = 0.0f; +#pragma unroll + for (int k_warp = 0; k_warp < SplitK; ++k_warp) { + value += partials[k_warp][element]; + } + output[tile * 32 + element] = __float2half(value); + } +} + +template +__global__ void nvfp4_qpn4_gated_sm70_kernel( + const uint8_t* __restrict__ codes, const void* __restrict__ packed_scales, + half scale_hi, half scale_lo, const half* __restrict__ input, + half* __restrict__ output, int hidden, int k, int m) { + __shared__ float partials[2][SplitK][32]; + + const int lane = threadIdx.x & 31; + const int warp_in_block = threadIdx.x >> 5; + const int projection = warp_in_block / SplitK; + const int warp = warp_in_block - projection * SplitK; + const int hidden_tiles = hidden >> 5; + const int tile = blockIdx.x + projection * hidden_tiles; + 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_warp = groups_k16 / SplitK; + const int group_begin = warp * groups_per_warp; + const uint2* code_ptr = reinterpret_cast(codes) + + static_cast(tile) * groups_k16 * 32 + lane; + const size_t scale_offset = + static_cast(tile) * groups_k16 * 32 + lane; + const half* scale_ptr = + reinterpret_cast(packed_scales) + scale_offset; + const uint8_t* scale_code_ptr = + reinterpret_cast(packed_scales) + scale_offset; + + float accum[NAcc][8]; +#pragma unroll + for (int chain = 0; chain < NAcc; ++chain) { +#pragma unroll + for (int i = 0; i < 8; ++i) { + accum[chain][i] = 0.0f; + } + } + uint2 prefetched = make_uint2(0, 0); + if constexpr (PrefetchCodes) { + prefetched = __ldcs(code_ptr + static_cast(group_begin) * 32); + } + +#pragma unroll 4 + for (int group = group_begin; group < group_begin + groups_per_warp; + ++group) { + const uint2 packed = + PrefetchCodes ? prefetched + : __ldcs(code_ptr + static_cast(group) * 32); + uint2 next = make_uint2(0, 0); + if constexpr (PrefetchCodes) { + if (group + 1 < group_begin + groups_per_warp) { + next = __ldcs(code_ptr + static_cast(group + 1) * 32); + } + } + half2 weights[8]; + fp4x16_to_half2x8(packed, weights); + half scale; + if constexpr (UseScaleCode) { + const uint8_t scale_code = + __ldg(scale_code_ptr + static_cast(group) * 32); + scale = nvfp4_scale_code_to_half(scale_code, scale_hi, scale_lo); + } else { + scale = __ldg(scale_ptr + static_cast(group) * 32); + } + const half2 scale2 = __halves2half2(scale, scale); +#pragma unroll + for (int i = 0; i < 8; ++i) { + weights[i] = __hmul2(weights[i], scale2); + } + + uint4 input01 = make_uint4(0, 0, 0, 0); + uint4 input23 = make_uint4(0, 0, 0, 0); + if (row < m) { + const half* input_row = input + static_cast(row) * k; + 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); + 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]); + if constexpr (PrefetchCodes) { + prefetched = next; + } + } + +#pragma unroll + for (int chain = 1; chain < NAcc; ++chain) { +#pragma unroll + for (int i = 0; i < 8; ++i) { + accum[0][i] += accum[chain][i]; + } + } + if ((lane & 17) == 0) { +#pragma unroll + for (int pair = 0; pair < 2; ++pair) { +#pragma unroll + for (int offset = 0; offset < 2; ++offset) { + const int i = pair * 4 + offset; + const int output_col = offset | (((lane >> 1) & 1) << 1) | (pair << 2); + partials[projection][warp][quadpair * 8 + output_col] = accum[0][i]; + } + } + } + __syncthreads(); + + for (int element = threadIdx.x; element < 32; element += blockDim.x) { + float gate = 0.0f; + float up = 0.0f; +#pragma unroll + for (int k_warp = 0; k_warp < SplitK; ++k_warp) { + gate += partials[0][k_warp][element]; + up += partials[1][k_warp][element]; + } + const float silu = gate / (1.0f + __expf(-gate)); + output[blockIdx.x * 32 + element] = __float2half(silu * up); + } +} + +template +void launch_qpn4(const uint8_t* codes, const void* scales, half scale_hi, + half scale_lo, const half* input, half* output, int n, int k, + int m, cudaStream_t stream) { + nvfp4_qpn4_sm70_kernel + <<<(n / 32), (32 * SplitK), 0, stream>>>( + codes, scales, scale_hi, scale_lo, input, output, n, k, m); +} + +template +void launch_qpn4_gated(const uint8_t* codes, const void* scales, half scale_hi, + half scale_lo, const half* input, half* output, + int hidden, int k, int m, cudaStream_t stream) { + nvfp4_qpn4_gated_sm70_kernel + <<<(hidden / 32), (64 * SplitK), 0, stream>>>( + codes, scales, scale_hi, scale_lo, input, output, hidden, k, m); +} + +} // namespace + +std::vector nvfp4_qpn4_prepare_sm70(torch::Tensor qweight, + torch::Tensor scales) { + TORCH_CHECK(qweight.is_cuda() && scales.is_cuda(), + "nvfp4_qpn4_prepare_sm70: tensors must be CUDA"); + TORCH_CHECK(qweight.scalar_type() == torch::kUInt8 && + scales.scalar_type() == torch::kFloat16, + "nvfp4_qpn4_prepare_sm70: expected uint8 codes and FP16 scales"); + TORCH_CHECK(qweight.dim() == 2 && scales.dim() == 2 && + qweight.is_contiguous() && scales.is_contiguous(), + "nvfp4_qpn4_prepare_sm70: tensors must be contiguous matrices"); + TORCH_CHECK(qweight.get_device() == scales.get_device(), + "nvfp4_qpn4_prepare_sm70: tensors must share one device"); + const int64_t k = qweight.size(0); + const int64_t n = qweight.size(1); + TORCH_CHECK(k > 0 && k % 128 == 0 && n > 0 && n % 32 == 0, + "nvfp4_qpn4_prepare_sm70: shape alignment mismatch"); + TORCH_CHECK(scales.size(0) == k / 16 && scales.size(1) == n, + "nvfp4_qpn4_prepare_sm70: scale shape mismatch"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(qweight)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + auto packed_codes = torch::empty( + {k, n / 2}, + torch::TensorOptions().device(qweight.device()).dtype(torch::kUInt8)); + auto packed_scales = torch::empty_like(scales); + const int64_t code_numel = packed_codes.numel(); + const int code_blocks = + static_cast((code_numel + kPrepareThreads - 1) / kPrepareThreads); + nvfp4_qpn4_prepack_codes_kernel<<>>( + packed_codes.data_ptr(), qweight.data_ptr(), + static_cast(k), static_cast(n)); + const int64_t scale_numel = packed_scales.numel(); + const int scale_blocks = + static_cast((scale_numel + kPrepareThreads - 1) / kPrepareThreads); + nvfp4_qpn4_prepack_scales_kernel<<>>( + reinterpret_cast(packed_scales.data_ptr()), + reinterpret_cast(scales.data_ptr()), + static_cast(k), static_cast(n)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {packed_codes, packed_scales}; +} + +std::vector nvfp4_qpn4_prepare_scale_code_sm70( + torch::Tensor qweight, torch::Tensor scale_codes) { + TORCH_CHECK(qweight.is_cuda() && scale_codes.is_cuda(), + "nvfp4_qpn4_prepare_scale_code_sm70: tensors must be CUDA"); + TORCH_CHECK(qweight.scalar_type() == torch::kUInt8 && + scale_codes.scalar_type() == at::ScalarType::Float8_e4m3fn, + "nvfp4_qpn4_prepare_scale_code_sm70: expected uint8 weights and " + "float8_e4m3fn scale codes"); + TORCH_CHECK(qweight.dim() == 2 && scale_codes.dim() == 2 && + qweight.is_contiguous() && scale_codes.is_contiguous(), + "nvfp4_qpn4_prepare_scale_code_sm70: tensors must be contiguous " + "matrices"); + TORCH_CHECK(qweight.get_device() == scale_codes.get_device(), + "nvfp4_qpn4_prepare_scale_code_sm70: tensors must share one " + "device"); + const int64_t k = qweight.size(0); + const int64_t n = qweight.size(1); + TORCH_CHECK(k > 0 && k % 128 == 0 && n > 0 && n % 32 == 0, + "nvfp4_qpn4_prepare_scale_code_sm70: shape alignment mismatch"); + TORCH_CHECK(scale_codes.size(0) == k / 16 && scale_codes.size(1) == n, + "nvfp4_qpn4_prepare_scale_code_sm70: scale shape mismatch"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(qweight)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + auto packed_codes = torch::empty( + {k, n / 2}, + torch::TensorOptions().device(qweight.device()).dtype(torch::kUInt8)); + auto packed_scale_codes = torch::empty( + {k / 16, n}, + torch::TensorOptions().device(qweight.device()).dtype(torch::kUInt8)); + const int64_t code_numel = packed_codes.numel(); + const int code_blocks = + static_cast((code_numel + kPrepareThreads - 1) / kPrepareThreads); + nvfp4_qpn4_prepack_codes_kernel<<>>( + packed_codes.data_ptr(), qweight.data_ptr(), + static_cast(k), static_cast(n)); + const int64_t scale_numel = packed_scale_codes.numel(); + const int scale_blocks = + static_cast((scale_numel + kPrepareThreads - 1) / kPrepareThreads); + nvfp4_qpn4_prepack_scale_codes_kernel<<>>( + packed_scale_codes.data_ptr(), + reinterpret_cast(scale_codes.data_ptr()), + static_cast(k), static_cast(n)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {packed_codes, packed_scale_codes}; +} + +void nvfp4_qpn4_dequantize_sm70_out(torch::Tensor out, torch::Tensor codes, + torch::Tensor scales, double global_scale, + bool use_scale_code) { + TORCH_CHECK(out.is_cuda() && codes.is_cuda() && scales.is_cuda(), + "nvfp4_qpn4_dequantize_sm70_out: tensors must be CUDA"); + TORCH_CHECK(out.scalar_type() == torch::kFloat16 && + codes.scalar_type() == torch::kUInt8, + "nvfp4_qpn4_dequantize_sm70_out: output must be FP16 and codes " + "must be uint8"); + TORCH_CHECK((use_scale_code && scales.scalar_type() == torch::kUInt8) || + (!use_scale_code && scales.scalar_type() == torch::kFloat16), + "nvfp4_qpn4_dequantize_sm70_out: scale dtype mismatch"); + TORCH_CHECK(out.dim() == 2 && codes.dim() == 2 && scales.dim() == 2 && + out.is_contiguous() && codes.is_contiguous() && + scales.is_contiguous(), + "nvfp4_qpn4_dequantize_sm70_out: tensors must be contiguous " + "matrices"); + TORCH_CHECK(out.get_device() == codes.get_device() && + out.get_device() == scales.get_device(), + "nvfp4_qpn4_dequantize_sm70_out: tensors must share one device"); + const int64_t k = out.size(0); + const int64_t n = out.size(1); + TORCH_CHECK(k > 0 && k % 128 == 0 && n > 0 && n % 32 == 0, + "nvfp4_qpn4_dequantize_sm70_out: shape alignment mismatch"); + TORCH_CHECK(codes.numel() == k * n / 2 && scales.numel() == k * n / 16, + "nvfp4_qpn4_dequantize_sm70_out: packed tensor size mismatch"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(out)); + const int64_t word_count = k * n / 16; + const int blocks = + static_cast((word_count + kPrepareThreads - 1) / kPrepareThreads); + const SplitHalfScale split_scale = + split_half_scale(static_cast(global_scale) * kFp4Bias); + const half zero_scale = __float2half_rn(0.0f); + if (use_scale_code) { + nvfp4_qpn4_dequantize_sm70_kernel + <<>>( + reinterpret_cast(out.data_ptr()), + codes.data_ptr(), scales.data_ptr(), + split_scale.hi, split_scale.lo, static_cast(n), + static_cast(k)); + } else { + nvfp4_qpn4_dequantize_sm70_kernel + <<>>( + reinterpret_cast(out.data_ptr()), + codes.data_ptr(), scales.data_ptr(), zero_scale, + zero_scale, static_cast(n), static_cast(k)); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void nvfp4_qpn4_prefill_sm70_out(torch::Tensor out, int64_t dense_weight_ptr, + torch::Tensor input, torch::Tensor codes, + torch::Tensor scales, double global_scale, + bool use_scale_code, bool gated_silu) { + TORCH_CHECK(input.is_cuda() && out.is_cuda(), + "nvfp4_qpn4_prefill_sm70_out: input and output must be CUDA"); + TORCH_CHECK(input.scalar_type() == torch::kFloat16 && + out.scalar_type() == torch::kFloat16, + "nvfp4_qpn4_prefill_sm70_out: input and output must be FP16"); + TORCH_CHECK(input.dim() == 2 && out.dim() == 2 && dense_weight_ptr != 0 && + input.is_contiguous() && out.is_contiguous(), + "nvfp4_qpn4_prefill_sm70_out: invalid input or workspace"); + const int64_t m = input.size(0); + const int64_t k = input.size(1); + TORCH_CHECK(codes.dim() == 2 && codes.size(0) == k, + "nvfp4_qpn4_prefill_sm70_out: code shape mismatch"); + const int64_t n = codes.size(1) * 2; + TORCH_CHECK(out.size(0) == m && out.size(1) == (gated_silu ? n / 2 : n), + "nvfp4_qpn4_prefill_sm70_out: output shape mismatch"); + + auto dense_weight = torch::from_blob( + reinterpret_cast(dense_weight_ptr), {k, n}, input.options()); + nvfp4_qpn4_dequantize_sm70_out(dense_weight, codes, scales, global_scale, + use_scale_code); + if (!gated_silu) { + at::mm_out(out, input, dense_weight); + return; + } + + auto gate_up = at::mm(input, dense_weight); + constexpr int kThreads = 256; + nvfp4_qpn4_silu_and_mul_sm70_kernel<<(m), kThreads, 0, + at::cuda::getCurrentCUDAStream()>>>( + reinterpret_cast(out.data_ptr()), + reinterpret_cast(gate_up.data_ptr()), + static_cast(m), static_cast(n / 2)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void nvfp4_qpn4_gemm_sm70_out(torch::Tensor out, torch::Tensor input, + torch::Tensor codes, torch::Tensor scales, + int64_t split_k, int64_t accumulator_chains, + bool prefetch_codes) { + TORCH_CHECK( + out.is_cuda() && input.is_cuda() && codes.is_cuda() && scales.is_cuda(), + "nvfp4_qpn4_gemm_sm70_out: tensors must be CUDA"); + TORCH_CHECK(out.scalar_type() == torch::kFloat16 && + input.scalar_type() == torch::kFloat16 && + codes.scalar_type() == torch::kUInt8 && + scales.scalar_type() == torch::kFloat16, + "nvfp4_qpn4_gemm_sm70_out: dtype mismatch"); + TORCH_CHECK(out.is_contiguous() && input.is_contiguous() && + codes.is_contiguous() && scales.is_contiguous(), + "nvfp4_qpn4_gemm_sm70_out: tensors must be contiguous"); + const int64_t m = input.size(0); + const int64_t k = input.size(1); + const int64_t n = out.size(1); + TORCH_CHECK(m == 1 && out.size(0) == 1, + "nvfp4_qpn4_gemm_sm70_out: M must be 1"); + TORCH_CHECK(k > 0 && k % 128 == 0 && n > 0 && n % 32 == 0, + "nvfp4_qpn4_gemm_sm70_out: shape alignment mismatch"); + TORCH_CHECK(codes.numel() == k * n / 2 && scales.numel() == k * n / 16, + "nvfp4_qpn4_gemm_sm70_out: packed tensor size mismatch"); + TORCH_CHECK(split_k == 4 || split_k == 8 || split_k == 10 || split_k == 16 || + split_k == 17, + "nvfp4_qpn4_gemm_sm70_out: unsupported split_k"); + TORCH_CHECK((k / 16) % split_k == 0, + "nvfp4_qpn4_gemm_sm70_out: invalid split_k for K"); + TORCH_CHECK(accumulator_chains == 1 || accumulator_chains == 2, + "nvfp4_qpn4_gemm_sm70_out: accumulator chains must be 1 or 2"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const auto* code_ptr = codes.data_ptr(); + const auto* scale_ptr = + reinterpret_cast(scales.data_ptr()); + const auto* input_ptr = + reinterpret_cast(input.data_ptr()); + auto* output_ptr = reinterpret_cast(out.data_ptr()); + const half zero_scale = __float2half_rn(0.0f); + +#define LAUNCH_QPN4(SPLIT, NACC, PREFETCH) \ + launch_qpn4( \ + code_ptr, scale_ptr, zero_scale, zero_scale, input_ptr, output_ptr, \ + static_cast(n), static_cast(k), static_cast(m), stream) +#define DISPATCH_NACC_PREFETCH(SPLIT) \ + do { \ + if (accumulator_chains == 1 && !prefetch_codes) { \ + LAUNCH_QPN4(SPLIT, 1, false); \ + } else if (accumulator_chains == 2 && !prefetch_codes) { \ + LAUNCH_QPN4(SPLIT, 2, false); \ + } else if (accumulator_chains == 1) { \ + LAUNCH_QPN4(SPLIT, 1, true); \ + } else { \ + LAUNCH_QPN4(SPLIT, 2, true); \ + } \ + } while (0) + if (split_k == 4) { + DISPATCH_NACC_PREFETCH(4); + } else if (split_k == 8) { + DISPATCH_NACC_PREFETCH(8); + } else if (split_k == 10) { + DISPATCH_NACC_PREFETCH(10); + } else if (split_k == 16) { + DISPATCH_NACC_PREFETCH(16); + } else { + DISPATCH_NACC_PREFETCH(17); + } +#undef DISPATCH_NACC_PREFETCH +#undef LAUNCH_QPN4 + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void nvfp4_qpn4_gated_sm70_out(torch::Tensor out, torch::Tensor input, + torch::Tensor codes, torch::Tensor scales, + int64_t split_k, int64_t accumulator_chains, + bool prefetch_codes) { + TORCH_CHECK( + out.is_cuda() && input.is_cuda() && codes.is_cuda() && scales.is_cuda(), + "nvfp4_qpn4_gated_sm70_out: tensors must be CUDA"); + TORCH_CHECK(out.scalar_type() == torch::kFloat16 && + input.scalar_type() == torch::kFloat16 && + codes.scalar_type() == torch::kUInt8 && + scales.scalar_type() == torch::kFloat16, + "nvfp4_qpn4_gated_sm70_out: dtype mismatch"); + TORCH_CHECK(out.is_contiguous() && input.is_contiguous() && + codes.is_contiguous() && scales.is_contiguous(), + "nvfp4_qpn4_gated_sm70_out: tensors must be contiguous"); + const int64_t m = input.size(0); + const int64_t k = input.size(1); + const int64_t hidden = out.size(1); + const int64_t n = hidden * 2; + TORCH_CHECK(m == 1 && out.size(0) == 1, + "nvfp4_qpn4_gated_sm70_out: M must be 1"); + TORCH_CHECK(k > 0 && k % 128 == 0 && hidden > 0 && hidden % 32 == 0, + "nvfp4_qpn4_gated_sm70_out: shape alignment mismatch"); + TORCH_CHECK(codes.numel() == k * n / 2 && scales.numel() == k * n / 16, + "nvfp4_qpn4_gated_sm70_out: packed tensor size mismatch"); + TORCH_CHECK(split_k == 4 || split_k == 8 || split_k == 10 || split_k == 16, + "nvfp4_qpn4_gated_sm70_out: unsupported split_k"); + TORCH_CHECK((k / 16) % split_k == 0, + "nvfp4_qpn4_gated_sm70_out: invalid split_k for K"); + TORCH_CHECK(accumulator_chains == 1 || accumulator_chains == 2, + "nvfp4_qpn4_gated_sm70_out: accumulator chains must be 1 or 2"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const auto* code_ptr = codes.data_ptr(); + const auto* scale_ptr = + reinterpret_cast(scales.data_ptr()); + const auto* input_ptr = + reinterpret_cast(input.data_ptr()); + auto* output_ptr = reinterpret_cast(out.data_ptr()); + const half zero_scale = __float2half_rn(0.0f); + +#define LAUNCH_QPN4_GATED(SPLIT, NACC, PREFETCH) \ + launch_qpn4_gated( \ + code_ptr, scale_ptr, zero_scale, zero_scale, input_ptr, output_ptr, \ + static_cast(hidden), static_cast(k), static_cast(m), \ + stream) +#define DISPATCH_GATED_NACC_PREFETCH(SPLIT) \ + do { \ + if (accumulator_chains == 1 && !prefetch_codes) { \ + LAUNCH_QPN4_GATED(SPLIT, 1, false); \ + } else if (accumulator_chains == 2 && !prefetch_codes) { \ + LAUNCH_QPN4_GATED(SPLIT, 2, false); \ + } else if (accumulator_chains == 1) { \ + LAUNCH_QPN4_GATED(SPLIT, 1, true); \ + } else { \ + LAUNCH_QPN4_GATED(SPLIT, 2, true); \ + } \ + } while (0) + if (split_k == 4) { + DISPATCH_GATED_NACC_PREFETCH(4); + } else if (split_k == 8) { + DISPATCH_GATED_NACC_PREFETCH(8); + } else if (split_k == 10) { + DISPATCH_GATED_NACC_PREFETCH(10); + } else { + DISPATCH_GATED_NACC_PREFETCH(16); + } +#undef DISPATCH_GATED_NACC_PREFETCH +#undef LAUNCH_QPN4_GATED + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void nvfp4_qpn4_gemm_scale_code_sm70_out(torch::Tensor out, torch::Tensor input, + torch::Tensor codes, + torch::Tensor scale_codes, + double global_scale, int64_t split_k, + int64_t accumulator_chains, + bool prefetch_codes) { + TORCH_CHECK(out.is_cuda() && input.is_cuda() && codes.is_cuda() && + scale_codes.is_cuda(), + "nvfp4_qpn4_gemm_scale_code_sm70_out: tensors must be CUDA"); + TORCH_CHECK(out.scalar_type() == torch::kFloat16 && + input.scalar_type() == torch::kFloat16 && + codes.scalar_type() == torch::kUInt8 && + scale_codes.scalar_type() == torch::kUInt8, + "nvfp4_qpn4_gemm_scale_code_sm70_out: dtype mismatch"); + TORCH_CHECK(out.is_contiguous() && input.is_contiguous() && + codes.is_contiguous() && scale_codes.is_contiguous(), + "nvfp4_qpn4_gemm_scale_code_sm70_out: tensors must be " + "contiguous"); + TORCH_CHECK(out.get_device() == input.get_device() && + out.get_device() == codes.get_device() && + out.get_device() == scale_codes.get_device(), + "nvfp4_qpn4_gemm_scale_code_sm70_out: tensors must share one " + "device"); + const int64_t m = input.size(0); + const int64_t k = input.size(1); + const int64_t n = out.size(1); + TORCH_CHECK(m == 1 && out.size(0) == 1, + "nvfp4_qpn4_gemm_scale_code_sm70_out: M must be 1"); + TORCH_CHECK(k > 0 && k % 128 == 0 && n > 0 && n % 32 == 0, + "nvfp4_qpn4_gemm_scale_code_sm70_out: shape alignment " + "mismatch"); + TORCH_CHECK(codes.numel() == k * n / 2 && scale_codes.numel() == k * n / 16, + "nvfp4_qpn4_gemm_scale_code_sm70_out: packed tensor size " + "mismatch"); + TORCH_CHECK(split_k == 4 || split_k == 8 || split_k == 10 || split_k == 16 || + split_k == 17, + "nvfp4_qpn4_gemm_scale_code_sm70_out: unsupported split_k"); + TORCH_CHECK((k / 16) % split_k == 0, + "nvfp4_qpn4_gemm_scale_code_sm70_out: invalid split_k for K"); + TORCH_CHECK(accumulator_chains == 1 || accumulator_chains == 2, + "nvfp4_qpn4_gemm_scale_code_sm70_out: accumulator chains must " + "be 1 or 2"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const auto* code_ptr = codes.data_ptr(); + const auto* scale_code_ptr = scale_codes.data_ptr(); + const SplitHalfScale split_scale = + split_half_scale(static_cast(global_scale) * kFp4Bias); + const auto* input_ptr = + reinterpret_cast(input.data_ptr()); + auto* output_ptr = reinterpret_cast(out.data_ptr()); + +#define LAUNCH_QPN4_LUT(SPLIT, NACC, PREFETCH) \ + launch_qpn4( \ + code_ptr, scale_code_ptr, split_scale.hi, split_scale.lo, input_ptr, \ + output_ptr, static_cast(n), static_cast(k), \ + static_cast(m), stream) +#define DISPATCH_QPN4_LUT(SPLIT) \ + do { \ + if (accumulator_chains == 1 && !prefetch_codes) { \ + LAUNCH_QPN4_LUT(SPLIT, 1, false); \ + } else if (accumulator_chains == 2 && !prefetch_codes) { \ + LAUNCH_QPN4_LUT(SPLIT, 2, false); \ + } else if (accumulator_chains == 1) { \ + LAUNCH_QPN4_LUT(SPLIT, 1, true); \ + } else { \ + LAUNCH_QPN4_LUT(SPLIT, 2, true); \ + } \ + } while (0) + if (split_k == 4) { + DISPATCH_QPN4_LUT(4); + } else if (split_k == 8) { + DISPATCH_QPN4_LUT(8); + } else if (split_k == 10) { + DISPATCH_QPN4_LUT(10); + } else if (split_k == 16) { + DISPATCH_QPN4_LUT(16); + } else { + DISPATCH_QPN4_LUT(17); + } +#undef DISPATCH_QPN4_LUT +#undef LAUNCH_QPN4_LUT + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void nvfp4_qpn4_gated_scale_code_sm70_out( + torch::Tensor out, torch::Tensor input, torch::Tensor codes, + torch::Tensor scale_codes, double global_scale, int64_t split_k, + int64_t accumulator_chains, bool prefetch_codes) { + TORCH_CHECK(out.is_cuda() && input.is_cuda() && codes.is_cuda() && + scale_codes.is_cuda(), + "nvfp4_qpn4_gated_scale_code_sm70_out: tensors must be CUDA"); + TORCH_CHECK(out.scalar_type() == torch::kFloat16 && + input.scalar_type() == torch::kFloat16 && + codes.scalar_type() == torch::kUInt8 && + scale_codes.scalar_type() == torch::kUInt8, + "nvfp4_qpn4_gated_scale_code_sm70_out: dtype mismatch"); + TORCH_CHECK(out.is_contiguous() && input.is_contiguous() && + codes.is_contiguous() && scale_codes.is_contiguous(), + "nvfp4_qpn4_gated_scale_code_sm70_out: tensors must be " + "contiguous"); + TORCH_CHECK(out.get_device() == input.get_device() && + out.get_device() == codes.get_device() && + out.get_device() == scale_codes.get_device(), + "nvfp4_qpn4_gated_scale_code_sm70_out: tensors must share one " + "device"); + const int64_t m = input.size(0); + const int64_t k = input.size(1); + const int64_t hidden = out.size(1); + const int64_t n = hidden * 2; + TORCH_CHECK(m == 1 && out.size(0) == 1, + "nvfp4_qpn4_gated_scale_code_sm70_out: M must be 1"); + TORCH_CHECK(k > 0 && k % 128 == 0 && hidden > 0 && hidden % 32 == 0, + "nvfp4_qpn4_gated_scale_code_sm70_out: shape alignment " + "mismatch"); + TORCH_CHECK(codes.numel() == k * n / 2 && scale_codes.numel() == k * n / 16, + "nvfp4_qpn4_gated_scale_code_sm70_out: packed tensor size " + "mismatch"); + TORCH_CHECK(split_k == 4 || split_k == 8 || split_k == 10 || split_k == 16, + "nvfp4_qpn4_gated_scale_code_sm70_out: unsupported split_k"); + TORCH_CHECK((k / 16) % split_k == 0, + "nvfp4_qpn4_gated_scale_code_sm70_out: invalid split_k for K"); + TORCH_CHECK(accumulator_chains == 1 || accumulator_chains == 2, + "nvfp4_qpn4_gated_scale_code_sm70_out: accumulator chains " + "must be 1 or 2"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const auto* code_ptr = codes.data_ptr(); + const auto* scale_code_ptr = scale_codes.data_ptr(); + const SplitHalfScale split_scale = + split_half_scale(static_cast(global_scale) * kFp4Bias); + const auto* input_ptr = + reinterpret_cast(input.data_ptr()); + auto* output_ptr = reinterpret_cast(out.data_ptr()); + +#define LAUNCH_QPN4_GATED_LUT(SPLIT, NACC, PREFETCH) \ + launch_qpn4_gated( \ + code_ptr, scale_code_ptr, split_scale.hi, split_scale.lo, input_ptr, \ + output_ptr, static_cast(hidden), static_cast(k), \ + static_cast(m), stream) +#define DISPATCH_QPN4_GATED_LUT(SPLIT) \ + do { \ + if (accumulator_chains == 1 && !prefetch_codes) { \ + LAUNCH_QPN4_GATED_LUT(SPLIT, 1, false); \ + } else if (accumulator_chains == 2 && !prefetch_codes) { \ + LAUNCH_QPN4_GATED_LUT(SPLIT, 2, false); \ + } else if (accumulator_chains == 1) { \ + LAUNCH_QPN4_GATED_LUT(SPLIT, 1, true); \ + } else { \ + LAUNCH_QPN4_GATED_LUT(SPLIT, 2, true); \ + } \ + } while (0) + if (split_k == 4) { + DISPATCH_QPN4_GATED_LUT(4); + } else if (split_k == 8) { + DISPATCH_QPN4_GATED_LUT(8); + } else if (split_k == 10) { + DISPATCH_QPN4_GATED_LUT(10); + } else { + DISPATCH_QPN4_GATED_LUT(16); + } +#undef DISPATCH_QPN4_GATED_LUT +#undef LAUNCH_QPN4_GATED_LUT + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void nvfp4_qpn4_dispatch_sm70_out(torch::Tensor out, int64_t dense_weight_ptr, + torch::Tensor input, torch::Tensor codes, + torch::Tensor scales, double global_scale, + bool use_scale_code, bool gated_silu) { + // Keep the dynamic-M decision inside the opaque operator so one compiled + // graph can safely serve both M=1 decode and large-M prefill. + if (input.size(0) == 1) { + if (gated_silu) { + TORCH_CHECK(use_scale_code, + "QPN4 gated decode requires E4M3 scale codes"); + nvfp4_qpn4_gated_scale_code_sm70_out(out, input, codes, scales, + global_scale, 8, 2, false); + } else if (use_scale_code) { + nvfp4_qpn4_gemm_scale_code_sm70_out(out, input, codes, scales, + global_scale, 17, 1, false); + } else { + nvfp4_qpn4_gemm_sm70_out(out, input, codes, scales, 17, 1, false); + } + return; + } + nvfp4_qpn4_prefill_sm70_out(out, dense_weight_ptr, input, codes, scales, + global_scale, use_scale_code, gated_silu); +} + +#ifdef VLLM_QPN4_STANDALONE +TORCH_LIBRARY_FRAGMENT(_C, ops) { + ops.def("nvfp4_qpn4_prepare_sm70(Tensor qweight, Tensor scales) -> Tensor[]"); + ops.impl("nvfp4_qpn4_prepare_sm70", torch::kCUDA, &nvfp4_qpn4_prepare_sm70); + ops.def( + "nvfp4_qpn4_prepare_scale_code_sm70(Tensor qweight, Tensor scale_codes) " + "-> Tensor[]"); + ops.impl("nvfp4_qpn4_prepare_scale_code_sm70", torch::kCUDA, + &nvfp4_qpn4_prepare_scale_code_sm70); + ops.def( + "nvfp4_qpn4_gemm_sm70_out(Tensor(a!) out, Tensor input, Tensor codes, " + "Tensor scales, int split_k, int accumulator_chains, " + "bool prefetch_codes) -> ()"); + ops.impl("nvfp4_qpn4_gemm_sm70_out", torch::kCUDA, &nvfp4_qpn4_gemm_sm70_out); + ops.def( + "nvfp4_qpn4_gated_sm70_out(Tensor(a!) out, Tensor input, Tensor codes, " + "Tensor scales, int split_k, int accumulator_chains, " + "bool prefetch_codes) -> ()"); + ops.impl("nvfp4_qpn4_gated_sm70_out", torch::kCUDA, + &nvfp4_qpn4_gated_sm70_out); + ops.def( + "nvfp4_qpn4_gemm_scale_code_sm70_out(Tensor(a!) out, Tensor input, " + "Tensor codes, Tensor scale_codes, float global_scale, int split_k, " + "int accumulator_chains, bool prefetch_codes) -> ()"); + ops.impl("nvfp4_qpn4_gemm_scale_code_sm70_out", torch::kCUDA, + &nvfp4_qpn4_gemm_scale_code_sm70_out); + ops.def( + "nvfp4_qpn4_gated_scale_code_sm70_out(Tensor(a!) out, Tensor input, " + "Tensor codes, Tensor scale_codes, float global_scale, int split_k, " + "int accumulator_chains, bool prefetch_codes) -> ()"); + ops.impl("nvfp4_qpn4_gated_scale_code_sm70_out", torch::kCUDA, + &nvfp4_qpn4_gated_scale_code_sm70_out); + ops.def( + "nvfp4_qpn4_dequantize_sm70_out(Tensor(a!) out, Tensor codes, Tensor " + "scales, float global_scale, bool use_scale_code) -> ()"); + ops.impl("nvfp4_qpn4_dequantize_sm70_out", torch::kCUDA, + &nvfp4_qpn4_dequantize_sm70_out); + ops.def( + "nvfp4_qpn4_prefill_sm70_out(Tensor(a!) out, int dense_weight_ptr, " + "Tensor input, Tensor codes, Tensor scales, float global_scale, bool " + "use_scale_code, bool gated_silu) -> ()"); + ops.impl("nvfp4_qpn4_prefill_sm70_out", torch::kCUDA, + &nvfp4_qpn4_prefill_sm70_out); + ops.def( + "nvfp4_qpn4_dispatch_sm70_out(Tensor(a!) out, int dense_weight_ptr, " + "Tensor input, Tensor codes, Tensor scales, float global_scale, bool " + "use_scale_code, bool gated_silu) -> ()"); + ops.impl("nvfp4_qpn4_dispatch_sm70_out", torch::kCUDA, + &nvfp4_qpn4_dispatch_sm70_out); +} +#endif diff --git a/csrc/sm70_turbomind/ops/qwen38_prefill_cutlass.cu b/csrc/sm70_turbomind/ops/qwen38_prefill_cutlass.cu new file mode 100644 index 0000000000..bba83cf9b4 --- /dev/null +++ b/csrc/sm70_turbomind/ops/qwen38_prefill_cutlass.cu @@ -0,0 +1,130 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +#include "qwen38_prefill_cutlass.cuh" + +#include +#include + +#include +#include + +#include +#include +#include +#include + +namespace vllm::awq_sm70 { +namespace { + +// Matches CUTLASS's generated +// cutlass_tensorop_f16_s884gemm_f16_128x256_32x2_nn_align8 operation. The +// profiler reports its logical instruction tile as 16x16x4, while the SM70 +// DefaultGemmUniversal specialization is built from native m8n8k4 MMA ops. +using Sm70Fp8PrefillCutlassKernel = + typename cutlass::gemm::kernel::DefaultGemmUniversal< + cutlass::half_t, cutlass::layout::RowMajor, + cutlass::ComplexTransform::kNone, 8, cutlass::half_t, + cutlass::layout::RowMajor, cutlass::ComplexTransform::kNone, 8, + cutlass::half_t, cutlass::layout::RowMajor, float, + cutlass::arch::OpClassTensorOp, cutlass::arch::Sm70, + cutlass::gemm::GemmShape<128, 256, 32>, + cutlass::gemm::GemmShape<64, 64, 32>, cutlass::gemm::GemmShape<8, 8, 4>, + cutlass::epilogue::thread::LinearCombination, + cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<8>, 2, + cutlass::arch::OpMultiplyAdd>::GemmKernel; +using Sm70Fp8PrefillCutlassGemm = + cutlass::gemm::device::GemmUniversalAdapter; + +void maybe_log_sm70_fp8_prefill_cutlass_route(int64_t m, int64_t n, int64_t k) { + const char* raw = std::getenv("VLLM_SM70_PROFILE_TRACE"); + if (raw == nullptr || std::strcmp(raw, "1") != 0) { + return; + } + const unsigned bit = n == 8704 ? 1u + : k == 4352 ? 2u + : k == 1536 ? 4u + : n == 4096 ? 8u + : 16u; + static std::atomic logged_shapes{0}; + const unsigned previous = + logged_shapes.fetch_or(bit, std::memory_order_relaxed); + if ((previous & bit) == 0u) { + std::cerr << "SM70 block-FP8 exact-8K CUTLASS projection route M=" << m + << " N=" << n << " K=" << k << std::endl; + } +} + +} // namespace + +bool sm70_fp8_prefill_cutlass_out(torch::Tensor out, torch::Tensor in_feats, + torch::Tensor dense_weight, bool gated_silu) { + const char* raw = std::getenv("VLLM_SM70_FP8_PREFILL_CUTLASS"); + if ((raw != nullptr && std::atoi(raw) == 0) || gated_silu || + in_feats.dim() != 2 || dense_weight.dim() != 2) { + return false; + } + const int64_t k = in_feats.size(1); + const int64_t n = dense_weight.size(1); + const bool exact_qkv = k == 5120 && (n == 4096 || n == 3584); + const bool exact_gate_up = k == 5120 && n == 8704; + const bool exact_down = k == 4352 && n == 5120; + const bool exact_attention_output = k == 1536 && n == 5120; + if (in_feats.size(0) != 8000 || (!exact_qkv && !exact_gate_up && + !exact_down && !exact_attention_output)) { + return false; + } + + TORCH_CHECK(out.is_cuda() && in_feats.is_cuda() && dense_weight.is_cuda(), + "SM70 block-FP8 prefill CUTLASS GEMM expects CUDA tensors."); + TORCH_CHECK(out.scalar_type() == at::ScalarType::Half && + in_feats.scalar_type() == at::ScalarType::Half && + dense_weight.scalar_type() == at::ScalarType::Half, + "SM70 block-FP8 prefill CUTLASS GEMM expects float16 tensors."); + TORCH_CHECK( + out.is_contiguous() && in_feats.is_contiguous() && + dense_weight.is_contiguous(), + "SM70 block-FP8 prefill CUTLASS GEMM expects contiguous tensors."); + TORCH_CHECK(out.dim() == 2 && dense_weight.size(0) == k && + out.size(0) == in_feats.size(0) && out.size(1) == n, + "SM70 block-FP8 prefill CUTLASS GEMM shape mismatch."); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(in_feats)); + const int m = static_cast(in_feats.size(0)); + const int problem_n = static_cast(n); + const int problem_k = static_cast(k); + auto* dense_ptr = + reinterpret_cast(dense_weight.data_ptr()); + auto* input_ptr = + reinterpret_cast(in_feats.data_ptr()); + auto* output_ptr = + reinterpret_cast(out.data_ptr()); + maybe_log_sm70_fp8_prefill_cutlass_route(m, problem_n, problem_k); + typename Sm70Fp8PrefillCutlassGemm::Arguments arguments{ + cutlass::gemm::GemmUniversalMode::kGemm, + {problem_n, m, problem_k}, + 1, + {1.0f, 0.0f}, + dense_ptr, + input_ptr, + output_ptr, + output_ptr, + 0, + 0, + 0, + 0, + problem_n, + problem_k, + problem_n, + problem_n}; + Sm70Fp8PrefillCutlassGemm gemm; + const cutlass::Status status = + gemm(arguments, nullptr, at::cuda::getCurrentCUDAStream()); + TORCH_CHECK(status == cutlass::Status::kSuccess, + "SM70 block-FP8 prefill CUTLASS GEMM failed: ", + cutlassGetStatusString(status), "."); + return true; +} + +} // namespace vllm::awq_sm70 diff --git a/csrc/sm70_turbomind/ops/qwen38_prefill_cutlass.cuh b/csrc/sm70_turbomind/ops/qwen38_prefill_cutlass.cuh new file mode 100644 index 0000000000..23ce6a1f7f --- /dev/null +++ b/csrc/sm70_turbomind/ops/qwen38_prefill_cutlass.cuh @@ -0,0 +1,13 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +#pragma once + +#include + +namespace vllm::awq_sm70 { + +bool sm70_fp8_prefill_cutlass_out(torch::Tensor out, torch::Tensor in_feats, + torch::Tensor dense_weight, bool gated_silu); + +} // namespace vllm::awq_sm70 diff --git a/csrc/sm70_turbomind/ops/sm70_top20_philox_fragment.cpp b/csrc/sm70_turbomind/ops/sm70_top20_philox_fragment.cpp new file mode 100644 index 0000000000..bc8534fcdf --- /dev/null +++ b/csrc/sm70_turbomind/ops/sm70_top20_philox_fragment.cpp @@ -0,0 +1,52 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +#include +#include + +#include + +void sm70_sample_sorted_top20_philox_out( + torch::Tensor sampled_token_out, torch::Tensor sparse_ids_out, + torch::Tensor sparse_probs_out, torch::Tensor top_values, + torch::Tensor top_indices, const std::optional& generator, + int64_t vocab_size, double top_p); + +void sm70_sample_sorted_top20_philox_token_out( + torch::Tensor sampled_token_out, torch::Tensor top_values, + torch::Tensor top_indices, const std::optional& generator, + int64_t vocab_size, double top_p); + +void sm70_sample_chunked_top20_philox_token_out( + torch::Tensor sampled_token_out, torch::Tensor global_values, + torch::Tensor local_indices, torch::Tensor global_positions, + const std::optional& generator, int64_t vocab_size, + double top_p, int64_t chunk_size); + +// Keep these registrations in a wheel-safe sidecar. Adding the sampler must +// not relink or perturb the speed- and quality-frozen primary vllm._C module; +// Python loads this fragment only for the explicitly admitted SM70 route. +TORCH_LIBRARY_FRAGMENT(_C, ops) { + ops.def( + "sm70_sample_sorted_top20_philox_out(Tensor(a!) sampled_token_out, " + "Tensor(b!) sparse_ids_out, Tensor(c!) sparse_probs_out, " + "Tensor top_values, Tensor top_indices, Generator? generator, " + "int vocab_size, float top_p) -> ()"); + ops.impl("sm70_sample_sorted_top20_philox_out", torch::kCUDA, + &sm70_sample_sorted_top20_philox_out); + + ops.def( + "sm70_sample_sorted_top20_philox_token_out(Tensor(a!) " + "sampled_token_out, Tensor top_values, Tensor top_indices, " + "Generator? generator, int vocab_size, float top_p) -> ()"); + ops.impl("sm70_sample_sorted_top20_philox_token_out", torch::kCUDA, + &sm70_sample_sorted_top20_philox_token_out); + + ops.def( + "sm70_sample_chunked_top20_philox_token_out(Tensor(a!) " + "sampled_token_out, Tensor global_values, Tensor local_indices, " + "Tensor global_positions, Generator? generator, int vocab_size, " + "float top_p, int chunk_size) -> ()"); + ops.impl("sm70_sample_chunked_top20_philox_token_out", torch::kCUDA, + &sm70_sample_chunked_top20_philox_token_out); +} diff --git a/csrc/sm70_turbomind/ops/sm70_top20_philox_sampler.cu b/csrc/sm70_turbomind/ops/sm70_top20_philox_sampler.cu new file mode 100644 index 0000000000..a09c094b98 --- /dev/null +++ b/csrc/sm70_turbomind/ops/sm70_top20_philox_sampler.cu @@ -0,0 +1,289 @@ +// SPDX-License-Identifier: Apache-2.0 +// Exact sparse Philox sampling for the SM70 Qwen3.8 batch-one top-20 route. + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace { + +constexpr int kTopK = 20; +constexpr int kWarpSize = 32; +constexpr int kNativeBlockSize = 256; +constexpr int kNativeUnroll = 4; +constexpr int kMaxGeneratorOffsetsPerCurandCall = 4; + +__device__ __forceinline__ bool top1_better(float value, int token, + float best_value, int best_token) { + return value > best_value || (value == best_value && token < best_token); +} + +__device__ __forceinline__ float native_exponential_for_index( + int64_t index, int64_t total_threads, at::PhiloxCudaState philox_args) { + auto [seed, offset] = at::cuda::philox::unpack(philox_args); + const int64_t values_per_round = total_threads * kNativeUnroll; + const int64_t round = index / values_per_round; + const int64_t within_round = index - round * values_per_round; + const uint64_t subsequence = + static_cast(within_round % total_threads); + const int component = static_cast(within_round / total_threads); + + curandStatePhilox4_32_10_t state; + curand_init(seed, subsequence, offset, &state); + float4 random = make_float4(0.f, 0.f, 0.f, 0.f); + for (int64_t iteration = 0; iteration <= round; ++iteration) { + random = curand_uniform4(&state); + } + const float uniform = reinterpret_cast(&random)[component]; + return at::transformation::exponential(uniform, 1.f); +} + +template +__global__ void sm70_sample_sorted_top20_philox_kernel( + int64_t* sampled_token_out, int64_t* sparse_ids_out, + float* sparse_probs_out, const float* top_values, + const int64_t* top_indices, const int64_t* local_indices, + const int64_t* global_positions, int64_t local_candidate_count, + int chunk_size, int64_t vocab_size, int64_t total_threads, float top_p, + at::PhiloxCudaState philox_args) { + const int lane = threadIdx.x; + __shared__ float selected_values[kTopK]; + __shared__ int64_t selected_ids[kTopK]; + __shared__ float selected_random[kTopK]; + + int64_t lane_token64 = -1; + if (lane < kTopK) { + if (ChunkedIndices) { + const int64_t position = global_positions[lane]; + if (position >= 0 && position < local_candidate_count) { + const int64_t local_index = local_indices[position]; + if (local_index >= 0 && local_index < chunk_size) { + lane_token64 = (position / kTopK) * chunk_size + local_index; + } + } + } else { + lane_token64 = top_indices[lane]; + } + } + const bool lane_valid = + lane < kTopK && lane_token64 >= 0 && lane_token64 < vocab_size; + const int lane_token = lane_valid ? static_cast(lane_token64) : -1; + const float lane_raw_value = lane < kTopK ? top_values[lane] : -FLT_MAX; + const float lane_value = isfinite(lane_raw_value) ? lane_raw_value : -FLT_MAX; + const int valid_count = __popc(__ballot_sync(0xffffffff, lane_valid)); + int lane_rank = 0; +#pragma unroll + for (int source = 0; source < kTopK; ++source) { + const int other_token = __shfl_sync(0xffffffff, lane_token, source); + const float other_value = __shfl_sync(0xffffffff, lane_value, source); + if (lane_valid && other_token >= 0 && other_token < vocab_size && + top1_better(other_value, other_token, lane_value, lane_token)) { + ++lane_rank; + } + } + if (lane_valid) { + selected_values[lane_rank] = lane_value; + selected_ids[lane_rank] = static_cast(lane_token); + } else if (lane >= valid_count && lane < kTopK) { + // Preserve the prior fallback for malformed candidate ids. + selected_values[lane] = -FLT_MAX; + selected_ids[lane] = lane; + } + __syncwarp(); + + if (lane < kTopK) { + selected_random[lane] = native_exponential_for_index( + selected_ids[lane], total_threads, philox_args); + } + __syncwarp(); + + if (lane == 0) { + const float max_value = selected_values[0]; + float probability_sum = 0.f; + for (int index = 0; index < kTopK; ++index) { + const float probability = __expf(selected_values[index] - max_value); + selected_values[index] = probability; + probability_sum += probability; + } + + float cumulative_probability = 0.f; + float filtered_probability_sum = 0.f; + for (int index = 0; index < kTopK; ++index) { + const float probability = selected_values[index] / probability_sum; + const bool keep = index == 0 || cumulative_probability < top_p; + cumulative_probability += probability; + selected_values[index] = keep ? probability : 0.f; + filtered_probability_sum += selected_values[index]; + } + + float best_score = -FLT_MAX; + int64_t sampled_token = selected_ids[0]; + for (int index = 0; index < kTopK; ++index) { + const float probability = + selected_values[index] / filtered_probability_sum; + if (EmitMetadata) { + sparse_probs_out[index] = probability; + sparse_ids_out[index] = selected_ids[index]; + } + const float score = probability / selected_random[index]; + if (score > best_score) { + best_score = score; + sampled_token = selected_ids[index]; + } + } + sampled_token_out[0] = sampled_token; + } +} + +template +void launch_sm70_sample_sorted_top20_philox( + torch::Tensor sampled_token_out, int64_t* sparse_ids_out, + float* sparse_probs_out, torch::Tensor top_values, + const int64_t* top_indices, const int64_t* local_indices, + const int64_t* global_positions, int64_t local_candidate_count, + int chunk_size, const std::optional& generator, + int64_t vocab_size, double top_p) { + const auto* properties = at::cuda::getCurrentDeviceProperties(); + const int64_t blocks_per_sm = + properties->maxThreadsPerMultiProcessor / kNativeBlockSize; + const int64_t required_blocks = + (vocab_size + kNativeBlockSize - 1) / kNativeBlockSize; + const int64_t grid_blocks = std::min( + required_blocks, properties->multiProcessorCount * blocks_per_sm); + const int64_t total_threads = grid_blocks * kNativeBlockSize; + const uint64_t counter_offset = + ((vocab_size - 1) / (total_threads * kNativeUnroll) + 1) * + kMaxGeneratorOffsetsPerCurandCall; + + auto* cuda_generator = at::get_generator_or_default( + generator, at::cuda::detail::getDefaultCUDAGenerator()); + at::PhiloxCudaState philox_args; + { + std::lock_guard lock(cuda_generator->mutex_); + philox_args = cuda_generator->philox_cuda_state(counter_offset); + } + + sm70_sample_sorted_top20_philox_kernel + <<<1, kWarpSize, 0, at::cuda::getCurrentCUDAStream()>>>( + sampled_token_out.data_ptr(), sparse_ids_out, + sparse_probs_out, top_values.data_ptr(), top_indices, + local_indices, global_positions, local_candidate_count, chunk_size, + vocab_size, total_threads, static_cast(top_p), philox_args); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +} // namespace + +void sm70_sample_sorted_top20_philox_out( + torch::Tensor sampled_token_out, torch::Tensor sparse_ids_out, + torch::Tensor sparse_probs_out, torch::Tensor top_values, + torch::Tensor top_indices, const std::optional& generator, + int64_t vocab_size, double top_p) { + TORCH_CHECK(sampled_token_out.is_cuda() && sparse_ids_out.is_cuda() && + sparse_probs_out.is_cuda() && top_values.is_cuda() && + top_indices.is_cuda(), + "sm70_sample_sorted_top20_philox_out: tensors must be CUDA"); + TORCH_CHECK(sampled_token_out.scalar_type() == torch::kInt64 && + sparse_ids_out.scalar_type() == torch::kInt64 && + sparse_probs_out.scalar_type() == torch::kFloat32 && + top_values.scalar_type() == torch::kFloat32 && + top_indices.scalar_type() == torch::kInt64, + "sm70_sample_sorted_top20_philox_out: dtype mismatch"); + TORCH_CHECK(sampled_token_out.numel() == 1 && + sparse_ids_out.numel() == kTopK && + sparse_probs_out.numel() == kTopK && + top_values.numel() == kTopK && top_indices.numel() == kTopK, + "sm70_sample_sorted_top20_philox_out: shape mismatch"); + TORCH_CHECK( + sampled_token_out.is_contiguous() && sparse_ids_out.is_contiguous() && + sparse_probs_out.is_contiguous() && top_values.is_contiguous() && + top_indices.is_contiguous(), + "sm70_sample_sorted_top20_philox_out: tensors must be contiguous"); + TORCH_CHECK(vocab_size >= kTopK && top_p > 0.0 && top_p <= 1.0, + "sm70_sample_sorted_top20_philox_out: invalid parameters"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(top_values)); + launch_sm70_sample_sorted_top20_philox( + sampled_token_out, sparse_ids_out.data_ptr(), + sparse_probs_out.data_ptr(), top_values, + top_indices.data_ptr(), nullptr, nullptr, 0, 0, generator, + vocab_size, top_p); +} + +void sm70_sample_sorted_top20_philox_token_out( + torch::Tensor sampled_token_out, torch::Tensor top_values, + torch::Tensor top_indices, const std::optional& generator, + int64_t vocab_size, double top_p) { + TORCH_CHECK( + sampled_token_out.is_cuda() && top_values.is_cuda() && + top_indices.is_cuda(), + "sm70_sample_sorted_top20_philox_token_out: tensors must be CUDA"); + TORCH_CHECK(sampled_token_out.scalar_type() == torch::kInt64 && + top_values.scalar_type() == torch::kFloat32 && + top_indices.scalar_type() == torch::kInt64, + "sm70_sample_sorted_top20_philox_token_out: dtype mismatch"); + TORCH_CHECK(sampled_token_out.numel() == 1 && top_values.numel() == kTopK && + top_indices.numel() == kTopK, + "sm70_sample_sorted_top20_philox_token_out: shape mismatch"); + TORCH_CHECK( + sampled_token_out.is_contiguous() && top_values.is_contiguous() && + top_indices.is_contiguous(), + "sm70_sample_sorted_top20_philox_token_out: tensors must be contiguous"); + TORCH_CHECK(vocab_size >= kTopK && top_p > 0.0 && top_p <= 1.0, + "sm70_sample_sorted_top20_philox_token_out: invalid parameters"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(top_values)); + launch_sm70_sample_sorted_top20_philox( + sampled_token_out, nullptr, nullptr, top_values, + top_indices.data_ptr(), nullptr, nullptr, 0, 0, generator, + vocab_size, top_p); +} + +void sm70_sample_chunked_top20_philox_token_out( + torch::Tensor sampled_token_out, torch::Tensor global_values, + torch::Tensor local_indices, torch::Tensor global_positions, + const std::optional& generator, int64_t vocab_size, + double top_p, int64_t chunk_size) { + TORCH_CHECK( + sampled_token_out.is_cuda() && global_values.is_cuda() && + local_indices.is_cuda() && global_positions.is_cuda(), + "sm70_sample_chunked_top20_philox_token_out: tensors must be CUDA"); + TORCH_CHECK(sampled_token_out.scalar_type() == torch::kInt64 && + global_values.scalar_type() == torch::kFloat32 && + local_indices.scalar_type() == torch::kInt64 && + global_positions.scalar_type() == torch::kInt64, + "sm70_sample_chunked_top20_philox_token_out: dtype mismatch"); + TORCH_CHECK(sampled_token_out.numel() == 1 && + global_values.numel() == kTopK && local_indices.dim() == 2 && + local_indices.size(1) == kTopK && + global_positions.numel() == kTopK, + "sm70_sample_chunked_top20_philox_token_out: shape mismatch"); + TORCH_CHECK( + sampled_token_out.is_contiguous() && global_values.is_contiguous() && + local_indices.is_contiguous() && global_positions.is_contiguous(), + "sm70_sample_chunked_top20_philox_token_out: tensors must be contiguous"); + TORCH_CHECK(chunk_size > 0 && + local_indices.size(0) * chunk_size == vocab_size && + vocab_size >= kTopK && top_p > 0.0 && top_p <= 1.0, + "sm70_sample_chunked_top20_philox_token_out: invalid parameters"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(global_values)); + launch_sm70_sample_sorted_top20_philox( + sampled_token_out, nullptr, nullptr, global_values, nullptr, + local_indices.data_ptr(), global_positions.data_ptr(), + local_indices.numel(), static_cast(chunk_size), generator, + vocab_size, top_p); +} 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 bd5b2c45e2..38278ee057 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -165,13 +165,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "Tensor? qzeros_or_none, bool inplace) -> Tensor"); // conditionally compiled so impl registrations are in source file -#ifdef ENABLE_SM70_TURBOMIND + #ifdef ENABLE_SM70_TURBOMIND ops.def("silu_and_mul_interleaved(Tensor! result, Tensor input) -> ()"); - ops.impl("silu_and_mul_interleaved", torch::kCUDA, - &silu_and_mul_interleaved); + ops.impl("silu_and_mul_interleaved", torch::kCUDA, &silu_and_mul_interleaved); ops.def( - "awq_sm70_prepare(Tensor _kernel, Tensor _scaling_factors, Tensor _zeros, " + "awq_sm70_prepare(Tensor _kernel, Tensor _scaling_factors, Tensor " + "_zeros, " "int group_size, bool interleave_gated_silu) -> Tensor[]"); ops.impl("awq_sm70_prepare", torch::kCUDA, &awq_sm70_prepare); @@ -193,8 +193,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.def( "fp8_sm70_dequantize_out(Tensor(a!) out, Tensor _kernel, " "Tensor _scaling_factors, int group_size) -> ()"); - ops.impl("fp8_sm70_dequantize_out", torch::kCUDA, - &fp8_sm70_dequantize_out); + ops.impl("fp8_sm70_dequantize_out", torch::kCUDA, &fp8_sm70_dequantize_out); ops.def( "mxfp4_sm70_prepare(Tensor _kernel, Tensor _scaling_factors, " @@ -237,6 +236,128 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "bool gated_silu) -> ()"); ops.impl("fp8_gemm_sm70_out", torch::kCUDA, &fp8_gemm_sm70_out); + 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( + "fp8_qpn8_dequantize_sm70_out(Tensor(a!) out, Tensor codes, " + "Tensor group_scales) -> ()"); + ops.impl("fp8_qpn8_dequantize_sm70_out", torch::kCUDA, + &fp8_qpn8_dequantize_sm70_out); + + ops.def( + "fp8_qpn8_prefill_sm70_out(Tensor(a!) out, int dense_weight_ptr, " + "Tensor input, Tensor codes, Tensor group_scales, bool gated_silu) -> " + "()"); + ops.impl("fp8_qpn8_prefill_sm70_out", torch::kCUDA, + &fp8_qpn8_prefill_sm70_out); + + ops.def( + "fp8_qpn8_dispatch_sm70_out(Tensor(a!) out, int dense_weight_ptr, " + "Tensor input, Tensor codes, Tensor group_scales, int split_k, " + "int accumulator_chains, bool prefetch_codes, bool gated_silu) -> ()"); + ops.impl("fp8_qpn8_dispatch_sm70_out", torch::kCUDA, + &fp8_qpn8_dispatch_sm70_out); + + ops.def( + "fp8_qpn8_gemm_sm70_out(Tensor(a!) out, Tensor input, Tensor codes, " + "Tensor group_scales, int split_k, int accumulator_chains, " + "bool fast_decoder, bool prefetch_codes) -> ()"); + ops.impl("fp8_qpn8_gemm_sm70_out", torch::kCUDA, &fp8_qpn8_gemm_sm70_out); + + ops.def( + "fp8_qpn8_gemm_ba_split_sm70_out(Tensor(a!) qkv_out, Tensor(b!) " + "z_out, Tensor(c!) b_out, Tensor(d!) a_out, Tensor input, Tensor codes, " + "Tensor group_scales, Tensor ba_weight) -> ()"); + ops.impl("fp8_qpn8_gemm_ba_split_sm70_out", torch::kCUDA, + &fp8_qpn8_gemm_ba_split_sm70_out); + + ops.def( + "fp8_qpn8_dispatch_ba_split_sm70_out(Tensor(a!) qkv_out, Tensor(b!) " + "z_out, Tensor(c!) b_out, Tensor(d!) a_out, Tensor(e!) qkvz_staging, " + "Tensor(f!) ba_staging, int dense_weight_ptr, Tensor input, Tensor " + "codes, Tensor group_scales, Tensor ba_weight) -> ()"); + ops.impl("fp8_qpn8_dispatch_ba_split_sm70_out", torch::kCUDA, + &fp8_qpn8_dispatch_ba_split_sm70_out); + + ops.def( + "fp8_qpn8_gated_pair_sm70_out(Tensor(a!) out, Tensor input, " + "Tensor codes, Tensor group_scales, int split_k, " + "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("nvfp4_qpn4_prepare_sm70(Tensor qweight, Tensor scales) -> Tensor[]"); + ops.impl("nvfp4_qpn4_prepare_sm70", torch::kCUDA, &nvfp4_qpn4_prepare_sm70); + + ops.def( + "nvfp4_qpn4_prepare_scale_code_sm70(Tensor qweight, Tensor " + "scale_codes) -> Tensor[]"); + ops.impl("nvfp4_qpn4_prepare_scale_code_sm70", torch::kCUDA, + &nvfp4_qpn4_prepare_scale_code_sm70); + + ops.def( + "nvfp4_qpn4_dequantize_sm70_out(Tensor(a!) out, Tensor codes, Tensor " + "scales, float global_scale, bool use_scale_code) -> ()"); + ops.impl("nvfp4_qpn4_dequantize_sm70_out", torch::kCUDA, + &nvfp4_qpn4_dequantize_sm70_out); + + ops.def( + "nvfp4_qpn4_prefill_sm70_out(Tensor(a!) out, int dense_weight_ptr, " + "Tensor input, Tensor codes, Tensor scales, float global_scale, bool " + "use_scale_code, bool gated_silu) -> ()"); + ops.impl("nvfp4_qpn4_prefill_sm70_out", torch::kCUDA, + &nvfp4_qpn4_prefill_sm70_out); + + ops.def( + "nvfp4_qpn4_dispatch_sm70_out(Tensor(a!) out, int dense_weight_ptr, " + "Tensor input, Tensor codes, Tensor scales, float global_scale, bool " + "use_scale_code, bool gated_silu) -> ()"); + ops.impl("nvfp4_qpn4_dispatch_sm70_out", torch::kCUDA, + &nvfp4_qpn4_dispatch_sm70_out); + + ops.def( + "fp8_gemm_sm70_prefill_prescaled_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_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[]"); + ops.impl("nvfp4_qpn2_prepare_sm70", torch::kCUDA, &nvfp4_qpn2_prepare_sm70); + + ops.def( + "nvfp4_qpn2_gemm_sm70_out(Tensor(a!) out, Tensor input, Tensor codes, " + "Tensor scales, float global_scale, int split_k, " + "int accumulator_chains) -> ()"); + ops.impl("nvfp4_qpn2_gemm_sm70_out", torch::kCUDA, &nvfp4_qpn2_gemm_sm70_out); + + ops.def( + "nvfp4_qpn2_gated_sm70_out(Tensor(a!) out, Tensor input, Tensor codes, " + "Tensor scales, float global_scale, int split_k, " + "int accumulator_chains) -> ()"); + ops.impl("nvfp4_qpn2_gated_sm70_out", torch::kCUDA, + &nvfp4_qpn2_gated_sm70_out); + + ops.def( + "nvfp4_qpn2_dispatch_sm70_out(Tensor(a!) out, Tensor input, " + "Tensor codes, Tensor scales, float global_scale, int split_k, " + "int accumulator_chains, Tensor tm_weight, Tensor tm_scales, " + "int tm_group_size, int tm_k_ld, int tm_q_ld, bool gated_silu) -> ()"); + ops.impl("nvfp4_qpn2_dispatch_sm70_out", torch::kCUDA, + &nvfp4_qpn2_dispatch_sm70_out); + ops.def( "fp8_gemm_sm70_prefill_dispatch_out(Tensor(a!) out, " "int dense_weight_ptr, Tensor _in_feats, Tensor _kernel, " @@ -261,40 +382,63 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "nvfp4_gemv_sm70_raw_out(Tensor(a!) out, Tensor _in_feats, " "Tensor _kernel, Tensor _scaling_factors, Tensor(b!) partials, " "int group_size, int split_k) -> ()"); - ops.impl("nvfp4_gemv_sm70_raw_out", torch::kCUDA, - &nvfp4_gemv_sm70_raw_out); + ops.impl("nvfp4_gemv_sm70_raw_out", torch::kCUDA, &nvfp4_gemv_sm70_raw_out); ops.def( "nvfp4_gemv_sm70_warp_out(Tensor(a!) out, Tensor _in_feats, " "Tensor _kernel, Tensor _scaling_factors, int group_size) -> ()"); - ops.impl("nvfp4_gemv_sm70_warp_out", torch::kCUDA, - &nvfp4_gemv_sm70_warp_out); + ops.impl("nvfp4_gemv_sm70_warp_out", torch::kCUDA, &nvfp4_gemv_sm70_warp_out); ops.def( "nvfp4_gemv_sm70_h2_out(Tensor(a!) out, Tensor _in_feats, " "Tensor _kernel, Tensor _scaling_factors, Tensor(b!) partials, " "int group_size, int split_k) -> ()"); - ops.impl("nvfp4_gemv_sm70_h2_out", torch::kCUDA, - &nvfp4_gemv_sm70_h2_out); + ops.impl("nvfp4_gemv_sm70_h2_out", torch::kCUDA, &nvfp4_gemv_sm70_h2_out); ops.def( "fp8_gemm_sm70_out_auto(Tensor(a!) out, Tensor _in_feats, " "Tensor _kernel, Tensor _scaling_factors) -> ()"); - ops.impl("fp8_gemm_sm70_out_auto", torch::kCUDA, - &fp8_gemm_sm70_out_auto); + ops.impl("fp8_gemm_sm70_out_auto", torch::kCUDA, &fp8_gemm_sm70_out_auto); ops.def( "fp8_gemm_sm70_out_meta(Tensor(a!) out, Tensor _in_feats, " "Tensor _kernel, Tensor _scaling_factors, Tensor _meta, " "bool gated_silu) -> ()"); - ops.impl("fp8_gemm_sm70_out_meta", torch::kCUDA, - &fp8_gemm_sm70_out_meta); + ops.impl("fp8_gemm_sm70_out_meta", torch::kCUDA, &fp8_gemm_sm70_out_meta); ops.def( "sm70_f16_gemm_out(Tensor(a!) out, Tensor _in_feats, Tensor _kernel, " "int k_ld, bool gated_silu) -> ()"); ops.impl("sm70_f16_gemm_out", torch::kCUDA, &sm70_f16_gemm_out); + ops.def( + "sm70_f16_indexed_rerank_out(Tensor(a!) out, Tensor _in_feats, " + "Tensor _kernel, Tensor candidate_ids, Tensor(b!) selected_raw, " + "Tensor(c!) selected_packed, Tensor(d!) expanded, Tensor(e!) partials, " + "Tensor(f!) barriers, int cta_n, int split_k) -> ()"); + ops.impl("sm70_f16_indexed_rerank_out", torch::kCUDA, + &sm70_f16_indexed_rerank_out); + + ops.def( + "sm70_f16_indexed_rerank_packed_out(Tensor(a!) out, " + "Tensor _in_feats, Tensor _packed_kernel, Tensor candidate_ids, " + "Tensor(b!) selected_packed, Tensor(c!) expanded, " + "Tensor(d!) partials, Tensor(e!) barriers, int cta_n, int split_k) " + "-> ()"); + ops.impl("sm70_f16_indexed_rerank_packed_out", torch::kCUDA, + &sm70_f16_indexed_rerank_packed_out); + + ops.def( + "sm70_f16_rerank_keys_out(Tensor(a!) keys, Tensor logits, " + "Tensor candidate_ids) -> ()"); + ops.impl("sm70_f16_rerank_keys_out", torch::kCUDA, &sm70_f16_rerank_keys_out); + + ops.def( + "sm70_f16_rerank_topk_out(Tensor(a!) values_out, " + "Tensor(b!) ids_out, Tensor logits, Tensor candidate_ids, " + "int vocab_start_index) -> ()"); + ops.impl("sm70_f16_rerank_topk_out", torch::kCUDA, &sm70_f16_rerank_topk_out); + ops.def( "sm70_f16_lm_head_top1_out(Tensor(a!) values_out, " "Tensor(b!) indices_out, Tensor _in_feats, Tensor _kernel, int k_ld, " @@ -500,6 +644,14 @@ 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( + "nvfp4_moe_dense_stage_sm70_out(" + "Tensor(a!) out, Tensor input, Tensor expert_offsets, " + "Tensor dense_expert_ids, Tensor ptrs_w, Tensor ptrs_s, " + "int num_experts, int k, int n, int group_size) -> ()"); + ops.impl("nvfp4_moe_dense_stage_sm70_out", torch::kCUDA, + &nvfp4_moe_dense_stage_sm70_out); + ops.def( "mxfp4_moe_single_token_prepare_w13_sm70_out(" "Tensor(a!) gate_up, Tensor(b!) compact_input, Tensor x, " @@ -576,7 +728,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "bool exact_per_route) -> ()"); ops.impl("fp8_moe_single_token_sm70_out", torch::kCUDA, &fp8_moe_single_token_sm70_out); -#endif + #endif #endif @@ -662,6 +814,14 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) { "int reg_buffer, int reg_buffer_sz_bytes, float epsilon) -> ()"); custom_ar.impl("sm70_tp4_all_reduce_gemma_rms_norm", torch::kCUDA, &sm70_tp4_all_reduce_gemma_rms_norm); + custom_ar.def( + "sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather(int fa, Tensor inp, " + "Tensor residual, Tensor weight, Tensor! normalized_out, Tensor! " + "residual_out, int reg_input_buffer, int reg_output_buffer, int " + "reg_buffer_sz_bytes, float epsilon) -> ()"); + custom_ar.impl("sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather", + torch::kCUDA, + &sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather); custom_ar.def( "all_reduce_sum2(int fa, Tensor inp_a, Tensor inp_b, Tensor! out) -> ()"); custom_ar.impl("all_reduce_sum2", torch::kCUDA, &all_reduce_sum2); @@ -689,8 +849,12 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) { custom_ar.def("dispose", &dispose); custom_ar.def("meta_size", &meta_size); + custom_ar.def("sm70_tp4_push_allreduce_buffer_size", + &sm70_tp4_push_allreduce_buffer_size); custom_ar.def("register_buffer", ®ister_buffer); + custom_ar.def("register_sm70_tp4_push_allreduce_buffer", + ®ister_sm70_tp4_push_allreduce_buffer); custom_ar.def("get_graph_buffer_ipc_meta", &get_graph_buffer_ipc_meta); custom_ar.def("register_graph_buffers", ®ister_graph_buffers); diff --git a/docs/assets/wechat-group-qr-7.png b/docs/assets/wechat-group-qr-7.png new file mode 100644 index 0000000000000000000000000000000000000000..a1aaf540dc2978225cfd4efd9c42673c9e53a053 GIT binary patch literal 529331 zcmeFZg;N~O7xxPU4-yFO?(R;2V8Md>;_mK{00F|{1a}C!0fOtY1PvZ!ad&suOXPRo zs`vf}w;rmtcBW@wrn>u_{+`b{&FA;ZvgoKps4y@v=<;%J)nH)YpI~4R?~&o2Th^}% zfX`pBZfdfUFcrfjJI^O?Ep+59m6TwZo{y1XV8d--;D6up{2_Y&z`(r7g@t+X{0;lN zEf?;8SK*&>U;OVF@&5OXFYVf&fAbBD{96ePZ`ggXZUc#|Z??YfS)OmI+55Wb;Ss5*&HSnA zemMnT8JTAc6`96;fTy8=MUi_0hyCwxpXcMQ^2l(X{QqhCcQyFK?$pu$r`0p(!>-rW z)HLC{|NGMJl+*v+1pnXN|MkiLJt_X1C;!>b|M5Zpd5Hf!#D5;*e=G8T*6^P-{AUgS zS;Jom{-09%PbvMUl>Qg4{$~yUzp;ilhVXy+!$EMIN1&n# zyQk)RTj+s?lX_5=+d^|Y`d7yHw`xd`3+cizhD*p1&{tGT{m&bJFM!v82seLMR(K{T z?GWCNAU;4(H7+YM0U7(oFbML~hA-nuKiog-wz*>bO0#B~V`*l>bJeL8%=Q#*`s7GH zF4QL4y0__J_z3low@}Np?{NACmj##_o}=cJ{Yzzz5n;W=Ah2Dsf7ZckAl4gDu0tGW z7h_za2GwVWT=B7^xjhf2?dJQdx(BZ1E21%CAqV}dX3q*{Ncr`Vh|5BPVeeD(>HRKj zq%U|F&d@?qSUV&CmRRs<_(nkF?_c$jdvkGj>6H^g$cJ~ER@o+!_QK1&^;CcwmiKRy|=5a0Uq?3e;&|7P+5M;kSlq&rx}P#1?lg!tW=F+rlv%`lN3s;0Fp`b zY+2ZG=RmG$bC36Jsu=R}G%xy#wU&k8jvk92n#VT(yukhqP0CdO0c2{Wb6xvq*;yR}54z+x5;?(12sxHXz z(F^TjX&vfu)*SQX5u&*F(+Nq#{E$aDXGGm3IqLpA(x2=}<)hnz`B2JlmiIA8|M|H= z`j3}v&+oDPL;@#2nCwN2n^X&I!~y5YU8@^vSNFL`cY7szrPRDD5VUaAwUExU)y?J^ z*;ATd)3e`=NK-15I#9w}%mxxoE@{>q#E+NJk&oJq3Gy-!2p544sV{31Ui$KD2~AxK5ivMt3tD3OYk(l&66Hzo%acU?MKXWh=Q6##jTl&ryEkRp- z{<)#4XREH%O-=aIrAJgByyZL42;CTm*nsS38@`$=xkui$XtYg^G@m={G&o)AqS(iE z@5d99gggxyd7berUMX1*Xw{H9X;*or%OGp8Ch3>Vn;OQ+Fq$xAvux5ZU6f=bMe<JB71g8iJE7M_maC_-=;ukh^F=ach(;cDQ6VoVTo38>e!Ro7-8^Rwu7_7(>Y? z8SIts8qXbtzn-38b3)a{(&OS=i4^bI+KnS#NY1rz-itW07Pj^x#f>6Nn6t7-y zB%_V6V@G*^`{4C;aY{>N>c}1);~Q^Cu|(L=aa`LM^VhHno%(dj5u;PKtpOQx!Wq2k z|9lT0Ec9qjaWMW*Pk6nWHZRm*Dqyngf5j~?o5u`Xap%ooEz0mM)YIn`f+~Ax*`~f5 z%P6dt)EHx9qFePq_hTaxI6t9tAMIYnVKyqBsGz(q^I-9p%i5sZA2ffpTiDn1iDjtl z0)WYwIVH(8RK~h$QHq#SD9Sx!W|y&zb;fa0LLN<%;ij65(!!O{f(xFgCQ+S>e3&(1 zkhWl*6x@V`0&Yw)ij6^RT(1jYBwl%JG!jX?onK)cmPA%jB6Mks5qcWZm!z2 ztMFDS+I&N7F$%qnGEx%XEx)%d*lR zvN{JF?!yaeCwqR3T=@En$utl4tJjSXUu<>Oh-WE}GFgEGaAqTw`cl;t^l8I1cC9vc zhGejk$2@=5d`pw|FKB#2!3$Jsm;-zKB0opmQyN|wDIUqr+j#F)QXl~e(bsCOsD~Z>3u6M9A95~-@fKlt-P1Emf05*K^kSzE?zqC`Uu=EyKhHv@K@OyiE z$uDLY4dy`Kh%aP-iIBaEH7k{;6Y{5%zC`-k_pF6fbGd{v#aQ+8=Ey2HjG{KJhx`{g zi*!Gd`wcczf*ks%6e;yuu-DE$_c|Ddeq|iDuN|#GOORbH$DvWeb!VrQzrUS=w7IMy z`exTL>dQ1+ts5B_9CyNn_|Qr&RX+f-H`tJ_MDvLt^sUSvt_){RqGiZmq}ms{y097Cg-8l7^!8N^1D9gKrmeOt%=+{8O7WqCIpC|NP{!kL)QV_IYgE8_@Cg^CgT)+1r zPHL_1cw->NDDcce4Z@I9@VMC5C-3I`0i(-31MzL&iudL5N@tLuH2)A_*vQYUxh&Ai zUY8BZUiYi2cCvEgJa9TQ{!C&kbkVL*I(=nyGw1QV_~}OcX=~|WVQJK_T7%z~Kp7}w z&(tPz?XUR|v*(nJr0g_K0g-x;4148}%&8Xl7Td3jQ9`*2b&dWt{r=HW(D4tXtHXu4 zxj7&s#Qz4WZNzCADRFdpFgIuiZay#TPK|FnPE)olS#@RS0fS2e<1cMWmbe$cEiUem zYUW)H<5dh;9V>>jhhu&=;!8EW&A*JfoOv02OB>$&4bDP`waI&NxS>a&3^-mL(KNa;2MmTIcu^wk<CVhzin0#zIbwJz?R7vUFJ zVUlyRD=y^u5zF{g|xVLA?w4r(=tVQ)b>aKk&xY zn=@U}vi7(i(eGwX;ygoH!f2~g@Lel)t*y3R=I&p|1nfS`7pb&$dMSA$JI&v!@F)IV{+$0Mp& zWBgr4JJBdUx@1DqadB@VzxuDgvklG!IXHYmO_Oinsj<{C+H;w|YSg^F5cP{dKi(Y5 znC2$lqpcT19jD;JbhEB`E$aqaN#7kKJ}z32TFX+gB;Yg3#I=7gtelJJ)Ggf@JCH~r zs;hOw6D$r8^fibPKeR>-_?!0+u2Vb<2UjpxpcIgIadC=%Zi393^lrGb^=8AMn&JjJzC- z@{%^?_7o$aM#te83hY|Fx2iJuMr4warAa&><2NSwX3$;hP7bFFKQKWs5@>3p)tRDc z$5*s|Lg$1x(+lo-c<5AdBF+DJmExCYY$V2(yoT5EwRdf8W9Q#4iMF&d~NzFP_vH4stAW{pj5Tx%0OAkMW0H&5gmt-HvhP$ zMLM;A3c=SJPc55TyQ)0`ur!@j)X9s5GU(iejLi7KFe<&C_X^Vyx@j+~seKiT3Oj!X znXou9mGGHTM1xFmG^|^-L3NT>2svB(34WeNRJ5)A66-E)pF%l7Bv!malzeK#g)hUI z?<_wcId9B&XmsN`+>eX&VIb%6!@A%7^+}D@$h3nPNzD6Rcbo#f^#jiat>}Q+#_}F> z!%8yJ&VX z{_-u93pXB>-1m=}jGTUqd_;-XMtS)Bi2%woRJ^@tIxflB4Svhe%4MvFZle(s_LguJ z7MIwil?v7E*&N(mN0MDahqDQw0UDP0PV(~YGlj%{si#hDOj z$SX-VX;aoEt~M+csf|9|g}SzK?jKO756ty~qLaE;wLcupfp&O!10~*EnO=0G5E>1# zIA)?YDad{DF<%vNB+h#YFg*BK`%jArBbNk1Uc^-X)fr5%wd-_Q3z_B2>ibP5x-7Q3 zO`LVXjIW99>VDb|*oNQfYZsE%{pfD2^92!Mr0Ibu#|m2_hQ|DocCv=Ea>yr8! zUMjC_T8-=`p56K@%KHPV&Kl<&_Z#?<{3S1J!_-TITY|PS`2a)s1JdHmyPz9b{Y>Bu zAe;AeMxy+aTsFQ3ODpK~GSfnIzujlYD;+%G?Yp_pIcw96@Fv4vckz46tn;g;uy&@W z&|mqlEOtGV@vlbBrVG>PsC#Fl$AtDh-fs@fpt1T-5Pz4rnmj@8=TYb#z%MzTu{Y=v zl^^iUR3sbTvUc|mHEGM9eTdD5eeY z5aXiCUuQ zf2G72@EXDFZ-~739CIxZCn5Zq;1uVG+EbwkQURYr-_C{3`pUj33H4`Wt^6t|?>8ad zAt3mZ{pUk2_;I@v^Jr;STwvM>$^q5-UYBn-{S`Q|LQbzo_e%j>YgHu6bjoMPC+}`& zIA);H1qx}Ez3Y6JbGi@SWKEAt4fGj^Qx}raMVCm@qo$)ST*`i9P#)@5uB|%Z&rK=*7p(s84%{XZ6%8IlHIY#wiTfzD$(v_MRJDyT2JAs(X(}3 zjEjx4gu#h~O`aZWCt&chzf@}YqQ#G`tF*hzJ76Ld86S%)o%!K>dqn5aWtxjuRP;nr?#0e zT$5XWwJBA>s7BL4lz>FmZLd{lgIIA0`VOpWHv!@Y4pjqI0+)S?7i}!f&IxO=-lJwA*$E(!_=Y`eBz`@0HdYXb9!?TNEYyfCFB zrCXZ%hGo~#ebF@uWEKU6f%XMUT6O1+-Ge@_8&{?xMS5h&KXXv(R<>-_XFx5at%Sr426?PfJUUPZGR=1wc9vkJN^u8SC9vEo$bBR~1T!?crS z@`$^H-CSyG^~nHbA?(*fIck?1QuwG$0OA_Hk^hgW)Ahc@r~9YJSYfZjuj;ETrnb>< zJ5C8`PL%Gao($CaPq!%EltqfW`4Uf&1li}1S`nd zbudR_v`1UB61!plQftk(s+3LZrn@AdIKbJ}bsj@>yw9p+K9Jgd^!tLI=_$zJ7=5AZ zvG<_uA+UHUXV}VxCi!Xmm1abCm&d8)c=pgVLd$xI%XLaN?%9b1z^jXX5c~x!56op*GS1zpfQpb-e2<8gM)@3n&4w?DU(T7dL#OT}dym!rr(q$(dbYJo zpkO|P6OJ1#`#UrtBq^D|LNDD;p}Uc>7jM(fk}z2%lxhv_I?vc=ov~Il_B=*wogV?- zd4{OCa?RQWRX(`QH($~x|AvU#Y36#b`pK55$1ui3pR4XotC9h099QM$SDuirz~S&N z6-4`wn`md+MgbS68w14TK|h_y?eTJ9+!WpNu!_)Lchnob{zu7<JWx=xW8Hf(R(ha9zAR-Y?UyoVty4S%_e#}IcqG{>J5xw9Flp&=oJu}dTSXD52xTC z2(A`GW^0-2S|@duO`4R&K{H$i%w~2n_@v~AYavWIN~6i!IsSEnS}S(Vsy3CNAk9v-_ca` z@U6Ac)&<)M%w!*cK6M^wh^XN{tWg?Lz6#@%ebrF{oR(@CGrU*)6`8ft+UY1VAhqF| zhs{q*xy7iJ#x7`X2;!S&8kU^sjL6#kM0>Ms{HD9lyO|LWv{5_H(Yl(-C%x;DOgF+| zY4&|f!;=JH+1D?eW4)8q$|9g{c-T)zkD~fWr&JTgt-NJg%I0xhXj|SxMQ=lNBO6CHbJfL#gs477QHC~3VVU;YJIIZ}ee^H;g<+BuKC-^!9gOGucrWN@ zT6H@{S1ETjDrVuUZ^3S-+LqRi~XPo$^B-fXvR|%k?TrsY_DG7sC(d8ePk}dFJyIqb zUOqbMs65_v$INr~=B8rwMQEJ#Nj8-Awke?Y=>){0EC%)TZ1okMh<6vK9DHe|KnGcE z@`g0fY5&@Q+i>5!VR|W0vt;k|)yc*#cZyEsDVqh63MD#MUGkXrnNiZI(?1^Qh4qZ8 z@JW~clp(h#!+5@Vt0tt8?${Nr;EF<$y0J+5?Ry5(q|BqSt+|&|L4sTJ<@ahq!RaRp zP~_JeOz{5t`nqwAts~u@eQl}oSwD*bpgj3(#@lJ(8>N!AiBv_AdUcZ4N4FAxaQYzL z*^G7D+t^WUupVd`mZ-(mK1gOLKZHLu%@6kwVe>0ze#7X9B&;dnx*na7@wVKQ^L5k) z-{pRfh}Lw4U3^XoS-crvp>o#)cd(17Q9aisbNHA{jp>ePW7Wv@l)#5!%F0{pEw0L`Dqzua;5w1W3P9U0 z&^CAAsXDu^?#hG8`Yz)0xNqomU^bT#%D{4LRZ8H9@`(I6az1$6v6m+S@3%M$EM7BvRZTv2pBkbnKG1FFzX8lc-qCrL?t- zt!>C$Fbp4IjjCgj)z`GPUN>eM@+{bU@#)r$?wm5CaAm{67L#z`$|D6gybuN4psXES z{lx42QBJM5Dv32 zlPMn!pO~+qhqwMGF&_@K2SYnG$`|giH#}SRJ-l*zY#9lHn?(HZ-G%}aS(SU+U|rfW zjdy-r8KL8I)?e|!>p%gccez7QD!s!m#fmmLHz_X29jWrA1q|)p`7MrD&K&J8)CsCE zG5COZ4l(@;O)AM%cSN*uNj`PbE7CSJhLdX)j^%T-i^4)hY|+L@*(zHMLPa4mS!*>u zRWJzUs`0vsxw4B?_diWAory3xE2o8(*m0)#!ys@*b9?Ji49)j5?2NyU z-HG3|c!`xLkETy#;w@72nZL}Ww8*6_*N>Xbp&KXnA^)5?S#OSSz>G2KMp*LHAZv_A zqLJluN~KyC-NgDhnaki;y>8bLjt_5+(4zb!vD7@UIjuwfg{%+1Wlv5T?{B4lFNZ?2 z=cFh$6S&v6SnoK&AH6|0fiL?m-o=MByt~^sk)L>I(zMC@F~50!zwTBxB$L^P`RqLe zbU}g@VpFGD>7KDszjflzol~i;FbLigvRsY}j-Z;*Wst8JF5k+mG5m&iW8YIW=PVeF zw9RYv0lzzQt>WcyY;~GXrTR#dIEr*LW4gCX#WyPw6`2RuEOtH{crTq|;#B_>gkrn1 z*5*|ete|w498N2jfhHbiX!on)QK*X8&H6Kvtvz0!S(Al`4{39|Uz>9QT=B=!L%@O! z&<=B6bc6xi)(Dk)t*J9#-kf~85&ZG#&xWyvGoWs4X|FRHN0h+pb*x)HbMvO|^k$f^ zir8S0G0W;8Dg?$v<;Ad1+)&#?>_X_ZFy>st^caQCgy~A1MoSbWe@6eOipbwQtD+aV zy3&8v1;A^$TN;#9=YE{a&iC%X88~R?Fuj+ zGD0=jL*36$T-9`Yz05&ap)R|w?m}E_sUVP&)~fVTdqb~Vohb)9-YE90Gn8l9+pqM% ztYXk{4!;r^M?qaV7AwV}Es%%yFOzHxduEbE9F~8p2^vto>+&jkS(Nk?Lh|&aa=+=% zM55gvYBkoB03&w?IjPu~hg5M}tPAH3NZj?nz72-<={{i!tEOTbr%Q3*nvJ4nd&@ix zI~^xg2#K>rCkIms*?01AFUxt{txg~}JekAzlPwun<@Ut%ViZwF(5Yoi?3NGj7pH<~ zWdJ^F%)RN|2C|cYXbk!o)VF5nWX^W>X7FZ}x7e51pr$rk%MPShzy6IADapo1J+$nt zH}70|NDrmyeymYo%x%_Sv4SJ!a7Q&b@~j70uvTYUlMn^b{t%_7;k96d+3b%E9wr=; z{XL5_UTjF}aQFc9dPiZmKISsjW6V{e3a9-GJ%=7$@VWF&!P!mp2iaPWsIrNME(c=4=a=Ucogr;gr2EJ9954N$&(^)mf7ywProyi+s7vc^>Q)!xj2 zSlJ-OR}5975mE4yLmJGT?TDLsGJ2k0Uo7ephUzJ__vV0K?B6CLz;kw_&C=xEU!lbD zkw#$7*qaYd;M-`kh45dSHU6t|Y;h6YHE_Ew1T38YK~(q4F-()t(2U@_+%KqDWeI zl8BD0luvT9Vx($TReYy|B9E{Ivq$-)-Aa5}G``&!kEEiJ`_Oy S)DV-pB>9rAiH zcS9+G%eyZ%K%N&A(YK#4H*ENDHd;FB*|7^z!BJ^?L@*`fqW7CX_O|D280XH0h&!)1 z2}f9Gr_`nWRHrk*kbf;mvI`Su!+eKFtq?#9xi4Q=tL`ytD$Rh02)ugzQ@*%MHzQTr z=L$RRBxdV$>w|Wm&IV4tIU&H}y5Zj%A^MCpmtG`G|E6d6sVEc@^-Si5yP{~VCGpDT zpM`84c1zNBO#^?~`uooZX(_Uo&A6ti2zUu^8}R5^sbphrskj7b60^rT&I|k9ZghM5 zk8M{!H2#3r-E5B5Y+0AD5z}vXkI&F#lGW*V*%NIt&GOMrcoJ*(0XL!w3Tk$|0Vpq> zi$6IPp*D1|Z2|qk1cnCdxV**4l(KX1Yr1WPdTscI zTR0|%t!h$=M`B|zl^sdBM2h+d-d}wuilwjR&Kh~YKmuYbPCs`*3swHi{>hF+I_4FJ zzz5R1Sp|K^G{R2ipwVU8&gF&|bg5iBK=_Ycd|9pudTok#DgVaV_c@x}o4J`w27iPL94__<;k8hqA4arAr6CW5R3w+)mT^uWYn=@tUg<>^7DXK9)mYmtV=Ot9d2|g+J1K(KXSR zc{q)5$-+y;-12+czzi}+31c6ze#%{qLAqV9Kq!bI_wbsJbmHofAO z- zCN+P`g@l?vu9=Zx6CbCOlc?3v!16Yk<`Lkf1T z&qN$=o1moo7ngin_^BrtWnkZ$MFKO;EEVRJpSD^ownKV_ngfF0=DwM^Y;Y|Sl4TFP zx7SML>1FSEPMr!a7Qgbl62E`-)YYvTlYe-GE|cgYrmLhq0Gw%5-*Luk5T=004iMgz zHo*(%03a}=;VFYSB=#=rVHX)3o?}*vH~F(0mH3d!ZUS$B)>8AF zrh`34u}rsI8+blcP%t8Pw;wXs@TzqUi@OX%8~FTFQ}fa_)qdL5u5hu3X4ZO;s5bQF z-gL;2VfVaCOKHQnY^bNS=onl^+vha4vIO`nDhnTYEE}8bEYTCu8u>&E32gfi3Vi&mh{}kSxpZNu>^HF6-bs- z9wkO-;B^`wAi0fRZqm78I?Y-qFE6F8f8w>&4_^zCBE<>*e>+~|qb=2Y2f zJ5w=r@^lW&xnI|Y+(N$@f_UJ;B{%$xC}(H~+vM29U}%z(VApwq;h`%QjGsh6$0 zJ0JKBZX$WDzC;pJm%HrU%1pZG;{EZP*p#;VclwTu>QgJzK;kmC0IgBs?h-#~Y zQxpyP&kpj~N60)$@aMvdY5!As&!GH5{ckXNBuzm=OwF;{ZfT*p;2gtSS?y6jRvz}v zW)Xy4myG+Bcmh$DFz;%fY|q zYSTdq|#u`wn>dtg|0EvRy}{5P)s=E z5>q87rveXjro>;n5j@nDp=zONq|L@58EeO`#rt*Ep7e^Q9FLl1pd9Y=baR#CRmQ^I zP=D6&vZL83fJMLii$Y*JbUG3|oIY5q<$S>Yn5Av-b*`ee6Ul?|EGFhk{QCG2d`G@d zPKJho+3fei6G7bX@|(6z<=udXg)V4#${aYPld%IxbvTKlKrdjAX2t!q14jYOBKh3R+qC)Z@z=0E7^Guo(zjR z1lRN9_3P~(l&!)Wth|qydjhn?mp+4Fg4#}EUgl4{sUIX#mvCgxMh9xTyLt4#J*|xd z_6amC5r~g?0LT~mk7CH@TY~N;qxJb!b{Bq~mFvcOkZ*B6@}!U=rYS$2+QE6G+01BP ze`H%tH&MP!9^<_Q9lz+g{AlVgA6ngRCu%?xRp31FP;ZmcTcQCc|Cu4Bqd}pq|Kfkx zC*v8SIbomZ{xwA48%7%{$wcGUZYJ+NhI?h+%^CapAb`2G9*(tj*+nSalZby5;@$VO zyp&CQe*jCiSU{v9mKi%r{b@d8l_%B9&T-9RH_uD#Bsr6HDqNpWK5o8 z?oB1{igg)k-+NndzTISq^|PBX8=Ut3ocG)0(R+ zVXsuIk0$G9oEADU`VCG{*{$=_`xPlIbSJ|JA)&}K=zg$e{|%R;#-<`WCs_3*lz3<2 z4)>;?S(1nbVG((b@xA5U?!oDf!#c#Xzn;}m^3O2P9J4-a4yTo4i9e^$CDfJ62h8(w z-8L$h4W`;GSJiy)>F$y@swLTtzh1YF)k~@&60PS%!p|tM zJy0%o1%-c~ z=OhROmg9s^bHcKKH}F=ii`5QQvOm4}KI1%?ejBN*kjS1GEra$CT8M`{FM%T;tTg}T z;_K<;iR^{YlU%38+Jj!h(IEh6U(*&z)*gxnoyBG5EI~r@uu(!|Y`r?^KI3$?N6@{! zh?hyKBPx6b7GBK@D&h|~n2Nv|9Mb7b%Ld1VbjMkFoI=3+cR%G8F^B~|5phPnk9`n3 z0sDn-kM*$(4iPY+@GS1+x=2Zd!T8lx;k~BA#cEzxPv!Y!!+kh!SSfzJ{350!@P3)P zzs4aW&(rVV?&=6~ze#!4uU!>H3}QLOT~TPfI_W(Rtc|>j-U0<9HgVy13y}7@Si@0U z^=HyP?x+Ms(zDhB&2%k}e=hSdF^v%k)sT4nbdlpCBI@O6G3bi11i)Sb2w|kD%TFyA z;H`)gbDdebOgrUt;F!21HiMNiBHiC&tlN;h6k@t^_m#P!-x^^Q2lcSnF3_Rg`xEdv zSaXME@H?q(+6(0SCX|%Q>d}35?;zG|M6;nmL-cbo|C(QkdIq2DaLFeTh!bddugUwA z{JRzt7R4X)ca{jYkLu5Od`OEA_^}el=>E>3sZ_JxBf}`}6o35bl%KtAx!M zFB}}^%>S^_a!U^aUpJ>z@~r_M^fuOF#oG)O0g3=Imh8;R;j>xiUt%Rk;>O}UNapy_ z(wv{8`vtO!dM)@mTNMp>l7SxA)vIbI?z#@@tXgO6kn8(#6*B*uUU9BHqx(~j5>S{d@yJ ziJFH|B=U#-nw)5txTMN0pOd?#r?%*%n`7gJi1JatBs32pqMW>ysf_N?+_RPNhCW%( z>>M|m#-StRB-&ediFrGXCLe;?Xv-A1heNWx6q7_N6m5cLKd3$5YRU&-&8)4UwdH~5 zg}HOPAs93i90^2AJihUE>G`<}?FS=y(LLR-o(0`(X@N8NiNeEspO%~+_m=KYyu2y= zu7~hPasu413NVSwH-i(w^WB50Yh1r@J!)*j=pD@Y0N~Y^j7OcT?AIJ3aNXYX%iDLF z3@SR|vNrW8q1d@!umir5Q15WR8eEO1g-CwijRmAGA++MHakm5*2$XE*wCY6&ZQNp2|Z4$Be zAJl!JdA*=>*oh9LimufgrM@BHBJ_=Vf=;C3mltN58%X2S=n51?;V@%Xn%2Wr~pw?h#hZ)1&AX^ldONZdQNm4XkOvy_Q2 zEi8Rvd~P5Q4NqYf!baicsV8~)N0&$(KDYanm($&@ip4zL)7BqTi@xJ}&6?B|NysQ* z zgMe`Dv$`ez^o$om>Tc`>x)mR`R|*~}D2k<3l(5uk+ml7)6Q#(W;`=D7UKI1Vh_pqr z2bd%c5YhI>Y0>C${|g}k8J_n^VA5KN{`Q~Zrf~A1p%g4qjFytnFRq(vik3bFZKYnI zi@<=z?y8CECIfyf!7~qlgXUR*$zr+)I5mfjsm}8P;Pe#osOscwRulQq?(lZOUQ53s zu1WM;ux@-boj`06jcr96lb&wC;*kl0h1(fC$k+$>eZ%;+M+QDQcbH;?BxI@|q?-0ZEU!wIbi2{P^S(eH2i~>jHGe+g02klK z^_lLot+Vl(UH>0gW*AgvbL-3Q}dO6GgStxAUa&Pg1^vwAkc9!oY?8DkX-NM-)xB85hXMMI@kklCErcb-hK~DuLM*+7p#aCjj zvebTlA%T#q4g4kNwyS4{b$6YJ6i&X>;^BSX%^&Lx%)9A%TpX1+*sFwy=|`>lwXXxW z4O6DxtieT7xuALzCmgeq7c5(RFhgxMww32_avRfbaOepA2Vw$#CuLGnH~s)OZF5i(mHJoq*AdDEh4}Gy81u zhw~S<Q-gsSGyJkvwzzy@D%RNZKPI4;cL0$-}AJ6uQ(3XJ%aQ#*j553jCQ#)bj!f z#fMzqT%|y65%6X+_&Yn@bW{aJTZTIU3)b($NQa~1_QD9oi)Kf}pYCVQQmmEBIA)cd zK5Teaee2R%)Qi>&c@UjoPm^o26&-9MP#tX`uusE?UPEm$!fb%LxIdQ;ZtTn-ry2SV zA;LL+m!GnG)^V`C937*$b8&0uEOm`EH#-Ge){h~UKxHS7Yu#LZ1kQ%nZl5(;vJR!v z*;WJBI`!`|D+7=F!pQ^9Rh`&nlJ3Qqo|erj40_ag53B|<2{|o&KU-L+_}9f)2w=(W zruZe7=#NDY!QxQ=3Vipqk)h)@fY9`33TkR(>4pXiucUunh0VWIvfpzmxecG~h0Y_0c^0^42Cw$9W)5PxOFJ@`!GpaFMe|~5_G#~6m)r(lOJ9@ z@ZS1>&$l`oU*QnjR-;-dY&a9r$?f5m2KK(LzBBc2nH3Zguq7qW7pzgv)Hx)Xe>*7| z^#)&Fu)vV>lWXTO_77iQAs-*T#iCkoo-Eu%O0n{3%^^=44Zh2X{0Pzs4T(=wHL8%V zsBpy5d^eVXpK9~Wy*bXe2D+?M12^tspBq5(EuOz-2TwE?zOmRzuW|TWG-&H?-xmn? z?)!*|k-6*r9MUX^Is2}YxoCWJdCzz2+oYErcy@6;Gk#%?nkXZ-q!wXP z@+XY4tO2;r9j6zp0vl0m?!fl1PdZ^cS)`Bn=;-^rvEfnaFIvizPXFIRXzAy48=#Se zZ!^_9ww$2O6jcHDrJ`>SzkU9bj zod6M~Lh|P?@m^wkzn5G%$Rf#o!4Kr*Q>g&co`pO}ofCBD`yqa~+wep$)CS3B@We>h za*~glRpSkxVM(G>({Q4PMK{cC;?j-BwQiJ4ySS->G?#K!;w_H=>S5tO5 zE>X=sT3)9q@nVc)#fD zyZ|w?d~_g@9|C2e7A&NJgttv&ouJ-w~|0SKEZn zboeE|(y^sG$Fl|(f6I~2KE#o)x$w>+dljz)atnf_2YSC-nR{Tfn*hL9-Ma^~KCORu z+v{1Lbh%r@jBVKv7_sb3h1m_{oI9Vrka5tpcQxRze0iw3DMmLDD3-Py2AYc8iogLf59of)=TlKi1hL>WkSv8A% zGy-pMo(>oxsdec5i~e6cXO@CsctF3E&NU~t-0oEt5dn#{9w1LmjRNVRCvEI(<%?)v zr>oIP5$-I3W{CaU6*amE61-(!B-pyCk~2K{uxcWArfIpo2B&cN-LUS)adHY_l+0)0 zw=)iTn1ehNZz*W30#^t;g!HK5^m`i)>5jT0`pkD6>O1H@w-=fgKzoS>G0YHUkcmC; ztlS6eTyBKnLg+0C;LAk3k6{Mr8jc*9?B=C$GX$;%-j^y^YUr$G+J)JoTil&8x_Af_ zuoYNLG4H|kxT_6Ih4Xf#$t4(fE<=D#Pmt;Zd`NL#UPrsS4xHGuB26uONcmK(=pS|Cc@4v1&>a*#)t9sR1tEwZ|F$-A&&h1%SsKTg0 ziMO0~OLS~Xh1}vxt2TuCiLL9cxR2A+ap#YZ1`qo0TLn(QapO5kG#mnNrN+>*$&??^uFfsc3>iye6vUQ zs6j-C*?D$3Um+41 zzaNj&ls`Q>jzlZ~@OZ*%U<#cM)1^w`XV(Va0_2mqsTv@tMkK+!H@mgP)^(_ zSs#{*_;}JG8=6NB#11S3%G}Kwk1S6aiwv;?lOTDwR@=bTzHrBy(KVYh?GiTin-=WI zkSg~Yuv4YHg2y3X}(gzZ`+jHtj#a#x)+S#J|i6 zkA^Hzs>!5orj6J@ibIa7w(Ga3dn1l6X$D5N_5yf#@+K!>zMr4-PgImU$DBj0Z0Tx2 zKBi9e+eNUz%k@?Ic6CSp@y|jB$7g3~i<8ox&D=%s^MQ`%?GD`ryvS}=_h)~D8Jskq z#T3u&HGEbUcf`R2Mbs_&lAqZ4zXGvqX3up$^|z0#@=_kt7+^(V+{j^X)RHl*D_8(STHvC=Yt2Z(M7y~S)5Y( zTrG5t9~(WJ*z_`m4o~DUvOEyxOyR5+IN~Y_SaED;Pw?=MOS2g8%x{zWGQb8@<(nLW z=c1;;B)bB8!=ShdEh*VlZq5+MDpHmY2HPwkIiI|M#NfA1P*>hZJE}k%ZwuIPVaLN4 zrtKXow(!`>=+q#U{1e&SzQYA}pI38*;?bFt`iQUp54;)EN$y*YAee*1xEfmnkG5oC zRQD+}r&^5%mktLLvu4dQBq3)6X!B0#k&!kq8qHm&oomlKiZOg^*q)q~A@I0&c@BK; zz+hJ(CY`ABWUF>IUQ#LxA9744mFlmTq%=rq?FRlcQYoB6KK zX$2^cebt7t@_5=>ba$FCpKt2_`k8#nfqke-w7=bqZgatN|60GE=)KXca;(aObefZ7 zxUT(6zv5Hh^lV#hkA}cwp4a(W4tkhfWMmW0s$g+uU)SoTVQ-&=>tlCnGsgiF37Koq z{jo9prkDM7a@;J&Sk;C7YWMw>C4u{|$f*!h%MLRZ=bO5S3q*ezzpiSjtQw`y@`TTHSvXst>OP*hf$j zf{oJrL)?82YgN&Zk0oFJZoXOhMqJ&QP}kF-V71tLmyEbg#;6ae=fG*|@D2Tyoxo;+ zE*$ylqxO7wJ@+Z1UuTXb zYnC9Ym5Rl_Khp0P*<%WFC%2|lXWMqA_XwxP2-2+I8&lD%5&b4t_!!{F0%*xEet!^n zdmbS49!YLO!9)(ejK}Y+3oxBszqM^}G4W`Eg2q%IGP@r45j6E7D)JBv7iUan2oPTGX#AqT>A;370dm?~!Q9?LxWD_ITk3EvKxBL|HdJAqgZ?El>~FVLQDG!o zl@{i1En>=>MB12(mb-NzlhmlW=V}K@l!7MiT+7CeAN<9NE0p&ceG%sfX1J2ss1>Qh!_=oNatwK5aTRIg^NFr=^KmOP#;n0<|ys{1_SOM zDbI3m%Hb02?Q)>iQ^nMMKk52x^yT|~dd)w?%}d}1oP8G}vij?w>5TUIHbt=f08eet zmc`0kAmZm!d=U-L4Ye5@p$1A+rPdkQzB$orlT0}zBz<`!e2c1SJYb-ycEMVvK!LT` zSxw^*!u(~Tw)j5GwQx4+)5r|?9-$-P86tn!wTVn=fEd*sJm<=v>t&f7LVkn8t{q0a zz$i4eXNuDNiWtoRd|qBB810Fj0=8hX@9gm()&@x$F~Px=LktBIC&dr$HmMOvj?e2| z_x@s|!3_PlKjXPbyTZb`Sff}4TrXCw?fav0PqXD~JJ~?_?9`zjB|}_MoC`O=HFz6{%{w^)A4%F709Z zn(>_-%A)FFBd;ymBM*}rC z=JtH-L~W6dw$(F{W zL@TPTQPHg$oMW-KfX5Lob_?f0@&2%Lyld(D&INmQ!$H57l}cNxdE8_ z&%ye1a}Ne#k1y{hZHw{vgHC(c9=%!)pHO&8zC%b#8kUOVnHdw@ir-K3Mx=~9iSn_? z4vHzp(o{=elzvx5X>ABlC^uDRh z&FWg|F4Bk19M{MVp7-`8mx4^$`};bcuYHfh`=@uC;)vca-Y+MG0d~;Z zmL$ATNGpKe(s~WrF$|IT2knlBJ_eB)1G^tvT>ctp_Y=3X=AWiLp03=u$A&8)i%9L0;2MD&ZK<-V#d<&Z4;amc`~c0b#2eeF<|==!SP+B%Jc!9g*%ZLA=WTH8Fbi$Cr5JpUD^073Wm}>2r|cM*SyrL!o)P^ zbUfksL(G_?(H&0PBtj5nkDQF<#vEBl-jHz|j(Q`ucCINlN#CG7c%rKuQ$$OSYF{*q z>efFsUeu+ao zug|Lzp6HT^scB<$N%`2n=_GS&HGxky?~u2zvmjvLPd*&XO63KHTXzRvucZsgUPC$u z)jc^)ISRmgH7f~VBn=w|6p*&M$93NKVr16kyCd zkiglQAp^z0W^nrUzEQ0lommJsj7ylH>!-O%q~P%^mt*Th^5elLrpwa1fm{ZTv`IJs zN89UV`7)w)`2f78g+|4uG@|vY+19}^D6>1TSxVB&7JoDz&kBv$H6(O0q1G^2wp(1_ zjaV}Zc9xN8MnXbSLCmzIqun%%U&E_DG;|sbs;1{PxO6rSdcbTVK9~KEI_D+HeEF(1 zzEzL&s}JWr_-PV9J0rD+{TU>uIoR(o*6DcccsuqEe79lOuTK}u%OD#3ejs$Z*uEw4g?KSx zOs93m>Y|PMPZO)O*Em{+1#9g5jhJLmN%L}|NLb=N!w4!$3S9Ay^WD;1!#-57cPMXZ zmhCr1%pS)uYb0U5MH)V}I=745jf(5tg5cMoS2rvGnmG5__QzkJ&?c_Sb<=_C86h{M zdDivlrxKE9p&rQ80LOhWKvodA^+$g>ct1`BqOldZD_E1^I^W z8Z}B<(r~^q+ryYK0$HWm^z?rK{QycBErr}Djr_ql`BN4M`IApb^TG_mfx z!sWn3mCHXT0t+03T*N59ekc(VPvW*+$1@*j4|A@?u~lkCV#I7KL^~7cR}lKe z74cU6(41p}^Tpu#k-xn_hu=$xE`Rb3g5vSuav&sFkIcEjfEPf^u!87SZJl6@vE6$6 zoK1f|C(+)O2U~Fcu;_rbH_@!wE(_zuU9qcy3lYXULaJFySro5;mLkzf{`BF?PnUBp z4Q%8*ZA>AXC~1~paILfSI8iKxHy{cadAf}DeyGWMJ+XP)v$^muhrArg`>uIV{ZpId zZbUQd{civL1H<1;%TG#Fot#t7%Z9Q6tZOFptMtzDk~)|ee>S?GxJV1ICgW7gp3a`M zPiGcS#_#}GYXFWo0y{r9zB>XigJphHQ5i1JG8~v=JsDwhaah^TA?Z9H4*-Ul_|Qbd z*Sd9@~mw_NvoE9UFP(Emv z!Xa}XLDxE~C+P}H7Txiu@3H@?q{Z)j_2SaOKjPr2Oc`9m`?_*fF@)CcLy9sgblXrK5T73l7y# zemeq$e>GM9U;xv{oWOnbyH&7Y_<;pq4sgjyuc$rb|uX&V_R}#xk@2EL<2vyPiQDW_RJMQ3$HcIDhE}FBRO#CCq`e z@l7vA;wd)M6*_1$ToQHhdfQEs&*Cr7&^bH3Oql3$Ay^51rY}MFyJ!}}Jqt;Kn8YC` zW~PHsb^Jr$pwVabWf4_76@xU;7 ze1lb~)2~oV(gdtuR>eyR(fqXs>1$08J#eAayY`p0{u?8nA9J@ZE&HbyAvkuZ@f7EGNg7#lsM5ze4D@Fd3wCt8F)=ca<{LjBzEGVxE=|qyyGdsc za6Ih%2EkRtl!OXRtJMwtCrfoAxr#$Xytb3tvXA~<@kK=VBm~KW&`oc$;;bjbOX--M z*`2l{4g*iZjE0=~Jp|RdIqc4n9*=^wiXo~F;P{h(3;FutxTa8EMduv=!VtfHFmAvjK7igvJPj6k=Gz z`Bhi1jkv>M^Grb#8=Oi7aw1>e++56`QF0n z{Ez`XHFrfJPV37jow%^fG=Dv3Mt@|Ad$=7Jggp=9a!M^XZ+Jls30Wzxq9cAnu)y#+ z8@k<{b-VB6sc=zI1bU1loVAi3zo=5&afg&K{|D^3O8@;_*+iC5| zh=?eQoD|P?Ici0?@g?KWFe#&EYT3;9!)%pAB59+%A~nJ=!BC?hD{8QCDd zw%t62???2KW6`kDCx}BW6$|Yc{+Zmx*F9lLW_AjrLPIWSUOiZ2E1^@TGK zOq&VCE%7tErs-Mpx);L!>&vvxi(j(#Nd~>>=VJe*S^Y@1dco(12QK8ME!L}OI(Rzh z8M6ZEHAaViL$hwj13nwIKm+!=7hCl8u&Ho|t3)-+g;S9~3I~M~+XoaKDPm8gLeyFs z9i-+b_-Mziu|0Vp;BM4rG;ENH z>%;auwNQ5C@Bh)9&6;UgIE)G>$dE15YZas{$Sry*RQy{BHJzc}Bpi0le%e(I;)qzt zmHeY~)JX|1hMJS%?pzLG|DMp`@oU%(V=L8XeL?TKgn7W(z;eYRE^17ZuNK!JiZulS zLl_&}_dzCRy35j4t*zQWvDu=e#LiK>VrriElRZ!Z!1UAck^rDlZ8?Wu^=6PRGXd|@AIM%Ag^oYC$pgBhIwhYHRC zufz}*4|a*qN3=$5Js#?leMrmH#anz1NI+DC-5sTIe90x+Hucua7GTF3h4$F`d-SsG zbt_r|XQRzF>gYp;+kh|fsMbfbkZxzbNlI_~+ffr3lGZ9A(J1Jgk_?BTP4>-wc8(IV zkxu~jGhWi1O;6!x`LtK-$gi;KwEqGtn~;wSp95kj2>V~~Ed;hEWQSKN-}vU)r@eyo zPUA#FC`3&|Xc@_kxceC{UgS0qZ_1&vA=nVcEI?84eL_0hyRuPV5$Qce#WY(LwPRES zoJj=!xsS+yaL&WtdFZC81$RUuLg4xJJT_rm1$Rj};nJBkWz4u{QI>Fju3Ph`8nV0nZs>lsKykVYiAb+Aff~Mo%BXe!#^VWFL-z4SK4pkT zd!D2XxSYHE0HW z4U+$EGa(H4vhNJH9sSj$-7yw2hB3sX8MmE~BKN@4M z>7|0vE~k}aM*VSonIy$A|LbbFLh`Bi(*xAu+9dbbMY7l3*4E(lbP^^>yi-8i;HhKu zx1S8bS06%VhlewyMfr5GkYS%7pQA{+kVMdrEoPY*ir6O3Ks^)#H}*_2-fPNnuG|$n z;!Qw=a(aGz8y6p0AD>!!{LuZwAnJ=p7VtJ;c-a+J80TwY-=>XRKX(i19skqK!$l>v zh4T7xnq5~)^yi#|M%v?o+IH$n;colYCWoX$-nflHFCtTDmtAaYHT!)ok_)zH5pVJ1 z{3$mr89FRxEAPd0?~JIx66frKm>s?}rv#^kRe-VTd9BW*v^c-l>(OP}`^Y1;eU^^O zJ6^?_+uRYIOX~!_Jql8xeJ%pc*CrbGE^lfuT52!xOs)CSFb3pwzixux3)TuzGzLmR zUY8CX>)pFp&&&V4))$C%Pv+mK|67hCgKrV#i12J>Q|_}58B)Qp@s>LfWJC`8)mY~* zlIUE)pum!svSUuv)S3YRfEbg*bh-OWFC&sN*R1;I>c zbewU;>d1MDnqGs9ZjcT^lZ_=_!UkHG9ap=k%y(lCN{H;H-WGGS8V$asg({Krk@paf z=6F5M>5nDDIR75@piT{v3%+-Q_e2cK#*c#jV8To zeaqG`VNqZ|4h_8l)_9f_EIUIcUvNf2UD4H=VXHtIW+bc zo4U2K&87zgm04$Q@~(4&Q13;GZQ{=0V2*1`w7|j09@0jp{-*2Jg+!PWnUEc-p`T3V ziNcSflTI39hJon&6Y96QTHy%8;+QXw`HqaS>*snMC*lSx4UkK4w{_2wE?MlFNzK%6 zn89yP0~U41m%^_*7U!yOOVg!$_xD%YKZ= z9RqdudDW^BRCN-lxd2Mi1?@2>cki0hF-E9;47byrb{aazqzW-MFEsQPjm8GSGQrG7W?-i!~&m3V!ZQ!hqJ)>_huyU4M^ui zA|A98lsYi6Zm-XgAu{L>FG}mU6|t>n)j!Ynz&ytZcF+3B|@bA_GN66 zG1KOX`o(z`rKEA;ZnZ%Ap+}SMMz8gBS#)+gE=@Tg?C9db%Ab*R!y!)5j5}g)GLoqM zQUCUK5Hoy~n+MlL>d|3^l2Y(aQ_c6u3bZ0Yem11XDvQqohtfM=Pntimb8rBC)#*}00a{^r_^jkFxp^qj1Lwp}<+rp)5k z!(dJiA8mY-&+QWl>dQE?VXo>i;>y5M;Z(P}zhaKip zPDUwIgXRsouJw*CvD#b3VuR2OM5GT|J_EYU^5{lDFg9Gtxv{`x^M#Tu(O21^!!>Sx zI%F40qHCv~*OxJxj>5OMewrh@izeiCPig)2;We8#uclL*JNLjO{I)v40o6I`5(iZMY!+1lJ$KM7{fxW#^&PJn_8?7kBekyg{QZ5D*E zv5}}uS&AhaaguQqrYzqz8?OX1`-PV_^shB0M-)75bbVf6m_9V2Z?K4rO+U?HBJG*_ zZPKlr21DPI|b05r;)}6bY+Niw!y1!XVX|63uLC}lv@>p{E^01^)dwm|E z7~Jr(sv;8JzBX7XP(6E0T5k0}c#iPFd4q)epWy$)fw%j9g9#G&u~d)0CAvBuf|IeM z#7~iY$ve_a#lU0cWh-Bfzi-AQ%5u_EMyzCsHnX}hwNf8(RFe}Or^}4nKRw>=*?-}9 zj>h~_-8-3tH9e)M0%*4jPe>ILtrKms@-S)_{odw+c=Zg1fER!p)gAO|W)|NR%r`NC zp=nEP4*88?k!uE@zeBBpsEXqb5lh)ddS+)!*YCwk+iHAFuJqKIJBt{M!>dQ~Pk2N% zU7BP}^p*dzIMQJ>fanm`lqTG{R-Q#`}f;oBQAH)BOxS%sq{@?$W1VX8x@ zx{mgBri}4I6Q_<)_E#8qZ{y_?Wr%quZEt{14V3(u+oA0po+<;w^+7w%nTf}}HhAD+ zU@7TM;k8Y5qf^|t5 z32sDrXIELgD7?jdAXn%h*#YbPiGJO_2t(ztzHP|0_vpp^#PA-~eMVb{Wvyf7^)$1n zazmz#q;}n+hu-lKad!!t-ShNNX}ag?=#i&9qwD-eBltg%!iV^x9(sW*O-RZ<)pkiEZeqlWk(~?xS z7G7Vx=KU5WO1@pMZzc7dmUmYF5v3m&?#qftn)|UEC0|tQeq~wt9!l9OK@I0FJ6Pfo9t~<+8Am&6s0C zN!8bHw~3-x`7@=hCBdrC_vH39AyT8cI?&^n!$~dcr527*Z}zw#v8ln181WB3o4=Lq z&sYDd|I=$e{XN5AWVrvXY`BJy@?XT-RO)bpU_cvfNGLE#;Rolzi3HH7z0oao7tIA_ zrl<>Fr@AAR#fCB(&aXA4iWw@Y>9>a}2GpkCO|?)qxRWnolmooTw3nEe>f;Q{U4N`- zVURYD0mC^>(C@_YA?*u9@r@7{e^U(+$JRIQdq%z%3_{?n*IGQqtGZXIJ;Osf?g+zLVllQud>Rrl=f99K!0> zE;q@DU(hW=H;~a}FuWjtTbytV{%lByH-z%e5C0+z&?tRJ_q~zl;U5Ic^g7k)%{HmQ zv3-rZ+IS7d?+AQlUT_#SPHqgf{>7o&;Z=JJ%LSXdy9hde0X&m)F@t#ZGYxd=BVXM| z%Gm#ZSlKO|#XXYtB}T zLmAtp+XmXVx0i{r*-ZkYyT({FQUz25m{BcVO^`hYy#u@*S2AS*qT-V2;iw3Dqt-bP zPS{C1bW&5<=$}OMX9pF|tT2f`754!}%{Cv?$CAjD?{Y;x)i9u2HDQFq9pf?TA$%%F zqer|nCy2C9G_s&4vI}1ceAWl_E}TwOX*Mj6%|8yDUT~%if`Mwh!w7D zJ59+)SP5{FC12MywLNoh=saHXoG%Nmf8k}A9k*X}h9F!!=c;DeS=y>byKi}j{wyaq z7xbEDQD-nSpIQDd_5C3oV*X|LFtJ8M{>c%90C0Y5)cs%_H)jzII6CDs68VAOsg*pO z#C%bEp-+=O;6EKyk}k>0WuN+wP)Jj+w*1Bn$Nv7uX*grJU18WUo8Oy-lHE z{3@)P#=`ghXoXowUSNx3BgIQ0H;O-Lv22zk>HAeCKvOZM30Nn{pF0gc-P9j>xV0A@ zpExzm43fT~KuL6tsy%ksRHen#gY8O@!{W5LROpXl$rpD4tRo4uDkf`j5aJN+@r)g& z?9&nf_RoqyS+=uouRm}d7)&34&eA=I?NNvwWw;whtFW(i!VgR5e|A(zwTAx>8hOKf zgyDN(r*-&0fc1h9>XUptFV9t-K-w9qaLn(Zu#Vn36MJx4*I(OKhENmwwHYr$o9eZP!37lR)MvQf!D_?Og6(wEm+4)$DpEeW zNBYTZ>s^cGrn`v@O3aJ-S|n6n%8gdm!X{NQ zG}GE1rYHW?L;%I~e8m^-IP&Xs&V;Y&dfanN4Uf6#<#o4+u9iRsS6~~Q4^hR;Dbl;v zH0U}|Gm86B1={U)-qWBC(oKp#b?tYDzZ~|Otkv32Vv8ZFH2TV`EqHj+x6KDv>CGZ> zx3m(Nm$D5Gf!;~nmEKb{?IuZ5(hBAOmh&GP7t{y(IfoLj`6pj>LWOtIhV4;JqqjJa zP+?|$yhpSn(xy3ML;$KpNkP*>v%9uTRGEjmfWFMVQ*(Y18J~lu675MEV`5#uqxQx_ zyzX&fGxM!HaZ<6P8-QezTepVr)>_7$a}DP_9yQNV;NfCevYbA!!lE{^!wQ9p(kdZN z{B5gjT$P+Ptl=W_E`>_1X?AzqjExg*^n~6xOx@GEI~koF^ke$Gsi>X$po>XKd|j_2 z;*tT8YbbT8lwN~h|qiTc{<{&k4+icM(*Ngj$Th&`K;|z}w;8!qk%Wf;h%w2cA z7uc680VOO>W{72tH^*e%A~~!vIGFcM`J1VD@nLrsM_A+Fh?}5qy`QFCku~L)Ynd?J zuv#M~!$y6P1%4$xJbVKJyQ^Avj0jAjUwOVN-Dk?gJAH} znlbPlm|}vzonT_BYpP0)eipRP0ZHJKW+-x%HH76jYHjc0^O3z|-=U4w=oCEkRSlxJ z9}mX2jHGii-)h}S|Dl(4oQ(T=xpVKcNOfIuC8Bg~e!4iYg^Ob8par-q(o&vzr@{xU9l6O&458 zz&fNm@-Ca^bx~;j5_wF!jI!MPFv%cs$O@ zXZHty+zxa7>@deIO%Yuaf;cSzN{UX5;iY^>icI)jz``!yAyfWu}DPySxu`J`su z`%z@433G6tfm7f<2iZ3^Xpq=*Ti2#d+r(@nIYL&30(~;$c&P=wR1oV>!Y!Kl{@3qf zb2;@}@?-u+7ZY!9LP|rrabwguKAx7331N$d{3b7hi+$_wuj}jY56E&YvV{xt?a&x1 zrkd#5p8}sJc_rtPpiw^1<}~JbaVSktrF?~6tjmr?yIP~FjuX#Cu1PMF$A!xha-rZ) z0TPaGud%{EqOiy4WFt^exF(czNl17+aKlTl0KFmC*RcUBk2MJF6cIQo{Fo??i?wBJ~>4~X3r_U9}MEZ)FCpgbcR0Vo4 zuukA*6EO?;WVWhuaq5A1)?L-{uobTJc9mqpkTgQCa=dHLh5yBpVjSw>T|D}ID>s^d zDC|(+ZCl_?K;YeE5pE^48u>X`&Cf@L_kBw|TCcq0^=1ABRU@2!L!|69|0^F{F2TOR zoeu`h${em7&MvVl@a5v|V{8hAXctG|azpfi+tUqKi(eXZ>0!`^F{cA~ae!J@^|I^h zJcH?f^N-u8IyFG!1?f+zRSz_b$Rh0HEMIm zeOEURFAuGn7^BXSjcN6Q$aQM-F0##@Z!O`q*yHy-M`UMYrwB1`L|BpWr}cu9k@gFTdYkaR)4gfzW>ayxn(kJmgM>v5M>zY|}HbTq|?Ap~svlh!|x4GSIz4WAM zQx=@W+}~*;hn#fcD>fBLb~RFz3av2>^j!^G`U#w$Qr<&@L!N{koLvj{(@Us{`glOa zS6DaH5NxJSyDcB@rPWYKw8jr=PxYK$%cRn~EOPxwM|*Mp1(YcssP?&JFYy09@0LnS zQP*1BGlHz==f%Q7xKtvlsq#{jEuWCj#c=VR0-c$?2?a1572`pHhxPmg`M(O=Ir`tL z0bP#8_-~2-K?eI2wtpwWLij=;jjJBmel&>ASL9igAxT^vUdCO-V7559a2WnX*C>r-)ec+@AztN^Bh;sOWa6deB# z&+XeLqQ;$DRmo)VGGOBm+;)gs-0wvauiRjV;^eL$fHAIK6H@0ySX820Yl2O({3oEqP!2S<-uk#u%r)vSYPyv1~`rEo9K zS|WU(Y2Om*9KP`z23Scx?}(_h-QO$71yFn>6YWjC4e{(d0T)8o4kJ+_VZyuTyfti> z;p<)SG^Jd$0V9w%qNpMx80=DZmQ%*7wpvxt|9F7R52p4x(}$Ji9|6L%2>M$a5hEj$ zIaQ8pSvaQBFQ@ew7fOQNX;x$g^TEkuJmo&)Z6AUl7UbApPdhs&VVUH4r|1mmZsE<) zcj(ku_S17(+{0Incwb=k`7eVtMk|5?4g=88-s1uFsEyNAs8r4tHtB5V~g^KQaTLwhi%WPq!R{u2!8;r;x_? zy-;K+wxbCN>|Ai#Uv~x}d^Um&|CE>b?|X1@3`Wbx=Y*I@Cdbp6**Q0JY7lnpM5NX0 zWM+y~HL{*U%eu14$4?dnC|#{(Npgl5gzbm7G>daT0ye%))sXyv(h8;E4?6+`6;{O6 zFQsb+?2=jCKVZkvL1k5(me$iVzn*P%H*e7CN6IsFn59SHUp?J1LHz8nyEGorcT%QL z3NL>POz@f;-9J(L)bo|PK-P7vY?KI3LF507wmg4v08gIr#DCluu%9MXu)-Zq0>0*I zf2iWl^}Am(Ds>tGWWoIt-%`~>@FugxP**eQhvDC9`8G?4^5p_v-g`l}nD3N036mV+L~d3d`g7;!(C`=n$h!(6OV z5`bjzNiZ2WQjJQH`?(>u{fiHZOmf>LJGM|^l$awxZk+Mpi(y38_b8pMN6lY(_|{j8 zYr{!}gAePfM!RI6Z#Ob^r)PA|0woSom)*O?OOx9-6W!JG#rymySKzhK=BCr=9YVBgLP zigq~dd|zGlybEjdF!}YraM1o64%)h`r2h$rM(A%g8xrY+6njcK50%{%Y<4^Wqlqsq zT6}sks*dHGvl(edpHIlraIqAbP%q_GcZIe;B#n;)WCFhQ0rhtOGYZGIVB zK#e`6nE@mrPIQ^CDpOhz6#EJ|Zl@|DHkfQn*BuCKm{zD$kUhBe)tG|A2fl7`eIdQV zy)<)pVFvn6gqJ4LQ#xUaKN@q%teI(kMB|+kl*me!9YLP&?m;z(yUpxLHV8->Z$Ic} z_PiG4V!3Enki=8EdhYM7iP**Y8!Le3PmTFjv(}TNZ3xFK25zQ)(xnp-uajzsfjkH3TJHS;IUp1bdt0CaqH4o;UT!w3M#C}JiiVD)gxbcg_ zT(4LUCu;c?0taH2S9k)rkOd63s|g=>zpu_V&Xjr;&i+q$e*PaPr(rhA;V{>RpK5>Ul>(+^JSk{)vwA9JX7`vk?UvLqPt*@mv%#1wGH`DpU|vC0T(r*eS&$>oF-(#p6=Rv` z%53kzZkDN7#)0)1OyL`S;-=^6GHFz|=24}iWuygQEg=GwFIhF=ONol@5P`NHI(VSI zCD)G;zGJGL35r;2VrE_smvj^w{uD&PDaKty_8?u?iPk|KvVw9TyJ#NJpII-miAX znRQyeCHKCILbyQXS69x{S9h_?taKxxD$KPMvfW+qV*7?L&t~YEkUQgGf-f9=n`(95 z>Lbe0!UX;tlX&W0x6$&Ym9RsbA2BBG>zy2)*{5%E`4t)rm@gBWm5@p>yu2_9H>dUL zIea$r3=*oEc}^-OR^(z@fU(UNVa|P)-l7FW$yL@=@2icES)Y)D=_D_|f5f>sE)P&n zBRrB%uB--w*Ovqy^ovarVnIDzXc7Ms$~+XC8?t_mPEF^r!n;Kg#Ze@@UZwz6~eUUaIn zRu)2Ku!Nx5FFM#0U<(VFQ{Hz zn{fQuCUQ}uHSU?g$@8=x_APA?R~sW?6nFEcU7$j!UcrjSky>&!Gf&r_HHcaiSds#Z zripJUjXv-1(>~DSnmA2y=tV$!7=W1&+cAV!3*L3oMv0BO!o0%i8aOemTVASjR$dqH zA(PH(oK0g_{+}=P9~uavn4jsm1Q$BLaU%6}4{79sbxixO; z9#ysZy2Ofe+tYfE0Wmx@sD|gf>+Ht4y3cJJ8txE%qm|T1Eq-b41+3yMJtDEm*Kpom z@q+EtYB~!lND{vK&aSa@h1AD=Bp#n`DBC+TBt`}=wg6DbM}R~!?!l@#^o$mWZxg<& zENSq;2>3$%D=u9gj#%f*Lyv@5hv*>~Mjz*RTs9drrU*awVM|{#FVNY(+0_>~ENcF_Bxj5?bOe@!GZ_|Y zY|(*9vMart!z3`YsTY2^U%VD&+!A*+7XuwPvi}RV!>)8#e~274JzDC6)Q1~m%bl%s z;&zW=<;ep%)#P~7o`*+zf*cSu;r5~*6KPQ@5^s&u zvCOfa7s>h`8H5qy4eev+-lz{2KP>J4V8nS6pNmq)v4^-jBt>X-IHujI3T`rGSI`1} zR6Blb1k{es0!}v-S=djGEf4;8*$J006LjH-?Djb1CYxne=+G-xv7-QU1XtLeRcQPC zmbzxe6U9AfLfx=!M~MHAtZxd>Y+Is@I~}uQ+jcr;$F|Kswr%5&txh_&ZKq?~wmLU^ zpMCbZ-@UK%ajlwjj8QeJN}4>Rz83D*EUb=R8v0Os_wj_)5?R+h?xfX0h}^&y`^z#b z89utAIrop-_hYXlS=pSN%nJG&QnKII4mFbzvS%2U2ow~ta+?ec--HU~MkfyL_D7zN zmpM^1g2{XFxfFGcfwZp=7k!@mZf{SMfv=gw6;hKN%Q<3POnni3tK!$w7`wQkCVz~L zj<>IzEPfy56S*gd@MM6fxJRaA7cFmZ%giiFGu1N=l)aVdz215qit-w-onK_MuCQhm zp5~^c0=hJ8WMkxXdAUqK+C!UNK3#%c&x90pfbacW5~$Z{Zp%!X>`bkEulIm3gOq3r zV?e9x$M&ar30r13OeDehYP&pO^$lXFCoTt2a_2GQnz}#yk%Q{}TK&J;-ro*)amItp}dDd231l~}EB0cZi?wX+mj!daSh%OR2v~;biy2PX! z+^6#C)SkBjEpzP_12AXp+S@yxlQzy}r_m0Cf+bu7joLGMwhPuOITR<{k+4XmI`uY{ zSp_`Z6rxD-n9?T@@M*{RPz^;|qRNaxU9l1reCdsLWR3bypJ_-^F9h1(?)woPM23 z4yk9_!Nj(t8Jp&{9peE_4c@_Rj;}YX92Z%>&wD$@&aKUUHxomxgn z4hljZO+lR@cH9WzGW0AY(Q`__ns&MmBPzYD3UkM~nNir8-O?X}#^}c;8he?s!jVl+ zE6b23GX}k3kW>M>tsWT^yAvlnZ+oH&bZw0>=_pXE<6``8GZd~DX9NTU0Qwa5>^(^R z9XFX{WUo+vQAs-$o2IeI-H8{kn$qk0i*%HeHceexN$GjE5#U;v43j-m>)O0+yw@M( zGkWvaF0b=^?ylzro5FwP(Ptj!)te>L`oAHqMGVrEClIY)hjxuyvO6o>0y84Q9}HT! zF4_^K%``S-UX0T)?YT`{*~|6SqbOcD7?P5nzudZ!#*&h`Z_5plfUiCJ(H7~= zm*{i=TlMR&^3dL^$v_U_8L5uMmB*xLWq;+r6eM1fjUp-FQ?g@2F#^ZAOH;@WzuNEk!pBkT{NET zP<-CQabMe+SKD-1`F%*1fVca=WiJ}Sn6{7p!v3iM8eXX6HQ9swSny`A0G3s2o3 zsUfjg>TT@-J;;0N+2g16sqY`l5S)LL_|~tF#q0CbYOH^P{KFopIbWBNpcEED65^(I zdPL#)4p(|@@S4?h2H%5N%jAlEPm)yH5JL(|> z(I&INOY8=6p~3XGhOR*dI+Op+aj&sp4QAF2f|_!sy?vM{#>9?xp+_8hB$SLsqlh+f zW_K34@=|Hamdj-)$1x;&&*;9=PkwtppepaihtxlEfJCn` z>yjTzspwVsQF_XH+7}=_EK%}{J?2ES8_39xhsr&ucZpl{GLO#poVeJ|1#4TnzPo{+^i13;e%Q8b#vvGN{kZevC*&@ zri1%lmGM~wa~ncivhCd7c!fWIv@HKa_&e!zsHXAX&{zE{^l6%{tN(e;Gl)U35;K>= zyu)`}p$D_QCDVe2bY>YTDM@ZK3;7Il8HJ(b^nR7y6;fJM=!VQkph z@5~sc;+_f&)$;aSoMZZNNw_z-qOZ4HyHwb*Wq7aWL{&~ckIE9wXpPY!Txo-D<%SXX zl=PYxpW8wQ$~li`gfjBBs!yv!dknFxCRD}6M85M_)pm?$(QM6!z|NWAG*CpfeC~rY zYFvfiu}Os^og>3lSI}PT@We6OlD-(UuUm*V=Om$lYI*o-RvfB(K-}KTnA6UisJ|+q zG+`=!7DwlwEwy^LZq(5@r1g-An~e3kCijKo{hZi2P<{GX&B(Jtr4r-JSO+G*LIT;X zPde~+1dYl~YXsGen_r*hEkDylX*4yXXUPEl3<)CoW(9s3_TLn|H_2xT-kZeo;U9sI z3x=dS#5&+fm3s?B$I^itH2S8Mz4C$Qhds1T zSI-XwXVoY^kGtmZY`{X7j`nY?fJY>!hDy&lbE|KcO;+pbDL3EypC}3N3wMu$AS6WN zX?~_eLZ+aNlmfT;b!-%Z6nYbjaa1#Ia0|3~CnIe%x;VK)O>+QWDU+vSLu^zW(6o!% zG?xoE+JBGTx^x_!(mD1?SbD6_>r9i0Hf~s}mcuv8Tx*l70IKQ$>|{T@nN-Ptdgaq9 zsT8z}i>#;1itmPx<^1M2-T%xotR4lxDH~jz-ZN*?Zcn|12!oJyUgTKgcLDxHE=&AF z4V%P1h5rBLnc08<6*t-DX}yJ`3b=ZP8i(@;`^C&K}AHL!kmnYXghJ%drTKIQvbURP(M0*Hy7x}oRfRVLPZ zOME03e|>#5U^iMANK9=7nKV`npRsb(!MgP5<=K6Ylz_DoeJ%2j8-vookcKubT2gbS z5yg1M#%iY)j8fPBjpn3grw_3HI@T?8T{7K+^_+sxkeQ^Rj`TzDBwGF!Ucjm)-vNBFRX5{*%Mu~o zADDm9?bx1W)xUGJTYbWWdx}hw`Uc7U+G6{G1&!}WG7)9XxZfEG)@>K=O<`i?847jO zE63m}-rD*{e&P)&Rcb8(px$hnFFY*LogiRxv#MtNSF?R zq#e>bDx0VYKm*C8&uyX_l(vUpHt=+|0|{m|b3xwG@-k0UwNEKDZktY-#D`%|krzGlZEN2_ZFO`19A>PH-ew{mWjz${7J4;K7Lqry!mV@tOHjTs z6_s9f`)iY5hP=U~iN_fIoM0I-!k5CY#UEEJv(XbZ4l;| za%jGt8tbs%a{?#C&KbqWE&>gde%$N)I*|gbb(jrys z-m5`7H2OjVv}oO1Hsi1lVa;2rfu<-#OX!V4*UF;_?6C?!sVtXu8o&kFCVEc>b+h3V zd;Z<$4ai2Gh-S7$!1n^((WIGFEw3G%HtB{Al@inMlhd#*%IA>1tV z6;nf13pX5DVfj~T&-#ImpH-=2VOtMyAUCG}IJ3`jz&|5Imi*y)$#u3b#6nMWzLR&l z>y{&j2(zd08aHOt@k|pVj84O!#IUpB;z@iAoFegt|P%2U(auytZEHJIV zVWGUtPT3<-osf7O9auj+Soak<+m3Y*HKpfEs3Ut~(RNRdF7HXCyVuRzo1GJ921dsD zz(heu&fA^%YTiNRrCTNJTrB^eqtEFc$4FOn`VG%a3bhiz*Q)|823fgmFBOaLH470c z@qhxPo#qf4v1iJ$>zd$0)H!TO6~SIVu+U|+8fR+dQ%3Rc#hYCTcsJiWWU%Cci|dwo zuk$@F(OE9(a0c;vpTnv+xsosM;J$2l7(Fh6L&Z|XN?__^6un>_Gue6X!x*=%7UuSMf}S$@R2{I4t@`- z(m%0Xg|x*^j~mV^{z(nfkEqYc;gI={8BP=Mf#{IV$Q!Ulh5oZr{9c$~YxAu++NsH; zqB;Q6R$TiUC=MJHs!)064rwK^!8IiW;UM{FwAf($A7Bgjt65ucH7H$*6ryV7`#|c- zqX#J1V!>Y~EO8fP(9BeD;XS)-~`=<}B2zbIpBl5qKV1(MsD4cWaQZCKIWs zDQ$M;1zF>^+5Su*1+_BcYrD7ut0c{S<>L} z+@pwbalma$L|8)P6otWDryoDI$Dh+=ymuFeXN{yJ6kM@RN!4KJvU-YxKmLCSkU$n2=;xoJ`k@bmZz8d7hxacNI^!Zbbs16XEZuw)ev0hrZamt^g1$VM?XE=GkS30 zpGU#C4v6XIW72n84qA?HDeygsS)L3&Ka;A1L-ODb--7o;o+bq~)ttX=hp)sWpeZm% zC-C~G^>vp^PTn6}({jO$RUo9hw=M=xG0@8b)?1g2-^Vce{`^ruSkw|JxaQ+>ofI?#JdLQ`ni(u<6iu2DhfKU83Giqqx5YKe^&znA9 zLSeggswf|P;K0ZLijU``IRhS9zcl$e00@`p&i%QBb-i6#hy*GVxG zu^Fr|wjOA4@t{tN8{pNUJ;LKjSvzYl4D@PW1PxG1u94rfV86?#i>$@?Lhg^eoT?E< zK-sOtW3yYpHPxr#DM)y>OV{t4B`MpN$=9-UfC85}W~kyJpph9{l$BL1hR?dxlM z4E)H}HX&Iu=<|WGD)$|8X_2`B)aHG9$8r&oyS@>PfZb}!Audi%&Q@m}kINz*dc9j} z;t2+>h^*!bMt;aZ=|d6D%~|u+^a(;i7CuuB3(e;48DZ`LWp+Lh_(8t8F3)`ajli;k zf&yFn+kwc#qwX`WOv*F%v-EIWm^#X_l7~x>Y7FezFb7N-w@u9{>SE@Gh7n4@Pe#(LDJLbVv?W(~2Z}|A*&9--zKC81{iM!uDxFOB+pUgMc=z4xcG*(-$Ed14pbF~kvYc|cVt}& zS*UBDt!I<1S{NVQY!4}K{lu;+enyPScCo3AOWutCa+swR7akIizBZP9W8O+7VaK09 zqjONXYx}=kjQ?+m-k;`VxW+<>kRJ^bPdQ2`Lw1U2h9%FpJwK7%IJVKZBCD&EV}~lI z!fTEV&1xyHJ9Oqp)-#@sVIjpxXCk(e)62CrNOt)Im!2N7c9Ry66tuuXFPSF82c10} z`N93xG|p)!Mk@pBVRU-zKC+wV0HQ_@jJ z*OpgKGg2Xr^LvLLp(taVt%9ase!ZLt3{~n)mlR&G!;WP`(-P~HvMk4yr|9Q0T?b$e zif2sO=LUy3PB$lRRRB9_*?ttCRzB_^dLJOFPtPh*_)B?EUu)#W9|Dm<@>)1H^&z+) z4_uV&Y4p4-6~u?X8K2}Zog9{Et@OK^xk>aE@n))UW?4IYl>S?c`fR1l`Ab&kn5WbI zEl&6ZAaXg-5PU>Aw-w~^`UOjWDqi*KTd)MVOn>Hz<0KPMXTTXyRCWR?#!|s|^QS}5 zX2Lb31@paLBNkqu@+#l*ebxnGDq*tF$Q7X&c7=FH(9dd#pP5QyJWdU1`*!NtdX05% zMDd?VJg_20j>%n8kSnZG)&th#FC6UMYL=|W*d-75nD}(=^?11g@A8-VJRP9a za3?~x@|Sa><-?JUX>PV3G;m^MG0aX&rsPuKrBV2GhZ*bBNCFYtgckRQZGL&l@(PvA z)RjQQlM`LR)u<`do@#!so4I^L@nXE&W5y6N`881)s|NY3dxF%V8tB#p4SA;9x}Vi7UzZ6>z&}o zw~NrOod8SCFc@0bu{t*rO^b)hpt;=~?>jKRj(BLq*A`|1sB^`!`cI$KXPE{0FFOe| zIgasvPOi8icyro4#z->ZXQAD1ZFJ9}e-;$S5qn^6;piR&|C)Fhtfrk|SJNMk6DBA2 ze}Sq3O2`;|)m5mBn$zdfdv(SVzE&zlJ+7j5>aEaN=K^DKQG3(UGfD&X(Q%_{DlJQ$ zM}djJYIHKIpqw)^Gdm+V2kv!z3Y44Aek0v66UKgg>~V$JVMxCS^cr`2B7klbBXT_E zeu~=?TpymMIp+gU+?TM~9emnFjW8ZC=J@V|U{clRZPR7J&Y|M)FVPJOr3}yoL?%mT zSL{aJq1ifwWv3-C1vV12-}wGK8L{qd0&;en;68?YFjdS16q$u5j*J-VzEUH12 zUB5KP6-OEHQWygMl+1AT=P#a>;A4|+h9VN=zubgf6Ji^*eCF9djuy^}{v~QYf90fo z>{$N;QC`*10?w6_7pug8D`F=-alRtlgP>QHyor0wsTlo0dBU0@&Ot{u$z12R?+0De zLLp3;V;YTvpBBOFctw`%Y-%r2Kc->L{*D`XQ>O8F=5@=5HsYCzO^P&cJSGZLCr@X7 zOe{rT%(F}%WF`&K0S*7&4mbPu?-bU~BAFJL2ixagWOy<&N7T~TI-hPH889*u2O@BK zZGhx+Sc3WNHsTJtYXM)?f<54vbSwKq^MXu}Y*jN_Q(v;s<~k^iKGw_*fTsf}!ATiY zJADegITMJFh!2m5fJtv@S`$(*VnVV|#=5Il=j3*;#2I|uX4a!+dvD(Zj?w~N3gla3 zz{j=Tn~5S`-N*h|mD6SGMr-Msw(2G;7e6Xk^^^PJ!JEMFva?G3qQ>_udL2an3*N;n zVjUZ+^B{?}Cxy+=@nX9Ke+J&geMe&3EQTv9fGe`s_g#^H(9a9?sK);%Qyydg!m-#g zjlb(&Ti_zyEYb+FeHhn7hc2)^fmd$5P-X?);2jXu0T>XR0r}i72HK~=TC;sa-Jlq+bTZ5`QC1J!%ZNhh9##FaxYa-cCNAt(r5&$;yy7qX-ia_ zP+1@k(BcW79-x<8+#lC@RI~KEe4_rJ`%GC#`)MqEK5tKna(v5O&SNz@aXt?tT|rWke`Gi;cK;{y{>uFD2%Dk) z%?zHZk^HVWZ<>uVvax1&%Lt{RiVugTx2PFvsVkyDtI5NV`W_x0avPBFpHP-&LVouI zdvQ!JYRjD4tk{(K^6{nr5 zndMZ-odYX)W!}+5bj66*h>1vQdeQGG#JountLuXLPaN=So{43|2i%fV($P%6yo&Et z|H`106ra+|OekqS-kqBe<(qs}mThv-LReYu;mwGeoSfQ!%s{Ph2jWl4bv_YeQP{oG z3t7pqqtnmGAeji%KfJC3x?1V_~*(5iN9f zo&o>4y!fPGy4CdQyagL248}~JiHmuQ?nXTrN{OYmhvaBKq9Ae6=LT5{-}{`87^3p% z3k5k-&$k$e!JWZ+?O>J4c{Ti{i(r+1xNJ!G|8KoXORoBUE<&5En?Bt1&9a{c8lAc9 zu)aOug8AlPB7+Y=_uWn>0H!YvE0$ei0t7~6lEV#>KZsq}@9U*Vv_?qyJa^2Om#D(I z#dZhnna!{i3TiqIQQAm~538~sU+tDZuR2(d71?1Lv@#S%`WcMergq$vxNZCB1D7dG z(%TLVo9qn)Q)mxUP@w*@F}uH>S9}SF9k%VN`2)`Zi&BHiq+l-L0ne1x`~;p4WBAQH zGHoOpNOw$gW&!lG&i0kS&vw_z`-k8OJ6u+?wS=P4R-EIzYBWrvsUgTdP7ajW^wGCvi${j^aL8`f|#vWp_%2^Z3Z zj1kxjk^#}VjGM-h3mn1k3K2L0kOt`OE0wJ^!Ol4@b%oyxz-umN0OSC$22+?k64^4BLZ_?I!P>NuHL_F!F9_64K!@b`& zJH@UzY{lnDdV#_T?Wzv}#O>Qpzu9v=ynb!;Y`p$AB3B7NzqZYEZOwn>vmgkzSdQ`c z)B@-M#NNPG9()1;kv!wyx~Go65vWNAVT*#6nIY%VgkV8k1cOEMj{BkenwiF;&8q~s zM=2VR0x?6O)M`*jnBu8?z*+URJI`_I-?gw)ZKAJ(@v+x3Tp09JoNl(%jGb5J#dt{R zF{J7oeTg6ZAmA^5(U{P=g!>bPMnFTt@1a>Fgzvy>LQJP=LttJ(X3B$16i#4{zT=n^29-sPfOC(Rq<&K$NjIB zU@H0(jQBZvO*j8GTzqUn3mEa{t zSKx1jFr1Iznx1p>c52?6&|k9yXfDmmae5ewCRyfy%Wix9YnNWWTQle##vhV)QOrLO z4AxK1l~YTr4RXk+ENmqK?Hd(@fgLxnC9Ckyb|aGgK3>}N;S+8kQgXV00M0VYH&ayG z*jtM{B0pyOCK94<6FVei$32nowTwi=Udh)8!f~BwT>78bH3e#i?sckb>}M5J9P_54 zB(gi{ernF>r>B0CAsZE&!fW&3(#)hH_O3JOo-=odVVi<2+bn{;^6UX5>_;}`4mFxF zFW3WFs8Tdyd}7tEYPk&WKAYC5?Ef0Do>2c{Ruq0_O8>~$%qBXj8?m|p z*m?KXAK*ccJNsiKeOsmShR`TB>h#jH1(>E^V~@HyvvxdSnn<0#^=);a0V2rsJuv&B zD}dM}(xK5aSD8_oq!Z{Lql-J8b{t(xCg)%D_|E@4xVUwb4=C9W<7M=Kv4KxJ#7m0$ zU5`132YP?RU^~n&AGbzHlM}x!Q8n#xIaILi95g0;^~7Zi02^<)!Yj7@2_FdR55{tr zYY_c@wcyLA4-#ARK#WNhXYr0cYVYA1)$|q_j)q8O`hAjnoou@7aQ`8x55jUG!J!=#Ag%zn;?}s!xJ3xmRL zz>LyPS%~lm2!zw()XNO~)T(laDufK^v$9!l!q`F-828JnTpRY|)$Byoap6BC=)_K8 zR2hY~&sDC$+wsI>yFace+}&tBw%s*s=smB_`dug-F(&M%`gmilY@HgW2Qq%f(hF zOITc&wPMlOr}SMiN?I4A1uJH9$pz4ua>nK%ZJtQ8(4ZKzWrWUa8tFKGA8lH5j-fwe zS<(tTRe`HQIV-%JM+6FTRts0W>ImMr@2NIgx8O=;`kS`Fsx<^Ao%=JtXVLSxM#K~$ zf-8d&NOK|#JX%#3T9++*kxmo-FYa;(KI<#KK2bK1e}nNKVG_UE!zwaii(E-Y?ieZC zjrV6-DfAl{e8|8XKxiH>e+Sa3q z9TZpe3R*~W{c-*E>(_8ZKPF70ZuY(o%&w*HuF+}9w^)yEIePsga~i@;H07Tzf>yO= zTy!26I2xalq6@RCF!X2rr3Ib0iVpKIE5{9_&F4Ees7pxnrqcs1*w}PirJSAl3cD_- z9%E|ugJjbjwKMRXsS-l`Y$K#v>I5G%%sN4?Nzbh1my^+#ldC`LkJoeL_HNCt9}SDT zuWM#{Ek5ra{IHwehrgDa_kJ2&tnW+EJ@PW)ZGs{}AUfib$>7Dx9==!i&4yP@Ebr$3 ztjT(sh3CNJ>1xFYWz=pj!iEL)g@ABcXAp{KX%@ORFbmfEWqreav~c`l?1H^VvrEEZ7A(rA2x5#q*gwwY{L*V=su|f3VP)u*^ZX|kTe&Bq3BRBMV8`Njt zb!#;(otTipcPFcL?$m~f>NfuqpJJVI5TB_=hRL1PJ9jtBZX(laKBDbJ^jLSB$`Ovt zIF5LPGd&n@RJt@rpUHM#M69<>JO6^2$qX5i&j2xuF1e?iKy0WgxMYhhosYk&uNhKA zg6clSC6%GwCN_L3OaVDDQ@5yR;Or$yXy()9_2B` zm}957{_L9@!L3~?3s~K&{Tm7}U4JU+9ozW>+2)N$yu!zlFUY*;!qh0lM2Em&+#$jQ z#}a=TrNhevQmGq_5mFjdX)6t#@x|NpxT4_i4_EAm27*jPuMTZV{KebV)ip2Afu4d6 zev1;HDo)rI37M_>Tad6Ud{WqAQx|)#G+DNc1K~9ApPVb)mIkZ(+K^a7%cAIr_{gm z6GDgo z&H_Nzj(GUZU@#7j$rbKF4#sv2#sbT>ZmMmc<0%D>V&a#op)a`dgx$)B0Vpja8Q4@* zuShu`{o#}+Ev3rJ#}#to-v5)b)c?)%@xs9Odk!ChYC#}`lz4@qeidBR zQVnnRSlfoy4BlAO6pc4tvxK4aOG4OU2Zog|TOi|aW+tFotxFS(IKIRBBFs4k1((R- z8B1U_z60;p-kWjM@sNgQy_a%Vok$whox_Of4wsphf+FsmxPv5esry?WtD%$PlHKJBW9v>@}x zciI${&+wc{W2R>FV;$HO$di(g+T6e!SkfX#GRVT*csZ>KH6_w^m@&`T3>J(4 zgvXJ7>Z!iUYjZm@i?2U6ySp0J{&VTO(VDd(qkI_DO^GTPUN_V9 zFfK*u|ENa2&yEn3M;z=%TFtbOga_>3I`KH$Tkil!XZ)}?BQv3X3*}f|xKdMc={U$_@Fvm#o!;pjv&w}fvo(>YkOaayR zHa$#6<2)=Ot4C+UMIfiROJ1p-!0BGRq)IF;!|^`WmySUt$YcpoJIwwDCry_4;<318 z_>e;|J^JreM*ZuR%%@eg_ym{F!>f82Xn$=tQs9>$@NNOY>ge5S{4_VIn<_uzyez8` zB3*g-No2kjH-)|5(@RbM!4giOgM^k2*q*UW1uAa3S$%0gyc*YO55b6UXl2d8duz-v z^XDeF!%m&{8YV{ zfu~1!^zCQZUmTY z__nZ6jDR%uxk)t^;ju-TDhdp4H^BhOZcnp#8K&b(Uug;u2u3SDc-SiV6UJ#FwIA3% zuoe1rL*+BxvNk8!#1Ymy;8^8V*;Ruojo<$9=QZ)YHoSnn_g-h$#H^6A(*C^|!u{Y@ za5rn*j=}z?I((%^`!z(oi{gHH*5{;(HvGf#Sm0LrsMzQ5-&*m1aR^oM^V*}LEQGNA z4NG(RAkYv_HJZ5?gy5!I_~*SUq@=;^H{y0>5VI1U-zUrwJ3Z3TVIj3m%i%lf@Awp3 z=axvAhEG~1NhqIL@#ez8q>sc+1y&sU(L8}k5b--yCVh;34l+ox5)WsJBc=7EZs^`M zX@bUneT*mGhAt=X5tOgu`?f|d!|*@SxlLNj>CbawfaM5sOqkGQU((+r%+m%}aCI?m zhz`r%R(I!=z6FE!yL7hJMcb_5h>SCarMJZ($HuSzY&p7$Cx{?71c5_xk}MasW^Z4YwJn!S#M8kT+xxlwZIZJtP)ma2dvhD6@`ldWV3#Q9 z84_jXQ}qTRf36lc*i@xHPRre1+d6eV(kDT_6|@y|*4^#cl#z|}|4<&@axk&3sh^qm zEPs8yZZta>5Gf?sa0<|DsVlPu^fg5}@TV^TM$)f!@!dEMK%~rFx~oAW4szf~Vu=LD7{yebIe!cDGSK)5SVrn*xvQ~Q zaqMAN2I1`Oz4)6zuvT=Lbu73VcImcE+2VorbM4o0vy9q+?lgo!8=)+SgG?5FB`H5k zB()pIVfE|!ZJh{}xqP`qr=@#d?e&$GXG34_68b*&`9e$26kt}_aZ|-z$n#5`>O<2% zjETz50tWM4H$90x?ms6Wq~HcZKS;g0h{;?nG28+j1;|f`ILwC@ISm!Z%^2NO8Ei+3 z684b8Qb4kY)Wie~beW*hR#aDlrN~Bv8Mw~(MV$FoO)g;F3t0x+$@tTxnLuuSbH?-x zTZoZ4J|6=gKm9|O?Y9V!TPZP#$J$935BhW8q=50*C@I=)F44UUub1|A95&Rr+lj?~ z68nZJCJ|r04LXE+1Y#vJ-N^nN-;0Z2j zxS)QfYTGM@v7Sqrf%p^RzkJpdxAJ4*$^hG7*LH?_Td~_5BgrJIH0n zYSSe%3IZU?43nM6kpt!Nz1C|Qqk69#6f!Cgcsnn#YRAd zwQi)RuuI16NxRCQ!&|mL935db3k&}RdADz$AfF@3c53twYBhr~MwRI0BA8pZ;jOq0 zU%NrM2)PncCfhMYQ}qEzHD?DCaTuY^dqmrdvYU7l^o45nuq)it)aG$W-F6Ts91J0=uIc4w{q&P3J1X z(>t4ah!pDZ6s%w-nD~!J=MU``2S|py<3My(?hJEXY(vB@Le$`2m7KCdSRs9>F_Pjn z6->!=`?wZ^C>8k!;ij~w-u{_EpnG!GfaHIne@NjoR(DIyy#HaS{Y_giU+E~pMZ0k_ zmlswL;mAS7v+3NWK#lX^D@3rwh)duIi||C9AK=6?Obm@z5|xDgnclT~h!xS3XaR+7 z+OFGJrbr8QI40)Rle0>SzUuP+K&!8lK+ZIMRz4y-3+*?&ejgEL;cBs`B@JJO0YcI2z4@0&% zD=CZe>=u(=^zP#_;|0I_7=!t!n0T(8oRAI!!xYj3ehG zs6LEW?I-B?)V{17{T-=qYcwE(-MnM{B;DY#kkO1lW*E1hclY=|wiVsNA+)a-4(lPh z-R`QqaG>&Fza|-~e?%N>x>5q!JJSdvJoqsnEkW4D9{P$!J_oEJux_lc&e`fRKHZ_2 z9XL*Xr?@gAYSu8)F*p@WWiQzM0(ZaAb}29xIwW2fur+bT7;`9~&~z0rXfWB(t>D() z#?zz2SD#HFhIz2xiwg)vz#E#~ z$*iHSWepQfwi2<6J%Wa;b0unVl~tWxc4{;uiZ|i&_9%?h0{}Ver9{q-p-Tjbsx`Uo zC)Or6hs1U?{TC~(4%>$X1+W}OZ&t1DFuw{ysU&Fx=;G)^i;H)zHH+> zKBV6Ll;(KAs47u7wmvk!t&hme=@P8Fq@-W*fUcuDZTWZ~mF+3wH~>_e$LVz8dLdYM`wt5N}P3(V~9Cvd4%E19B7QQ3g^a|#w=8T2@| zJq}E@;V>&Aq(7i!B9E>il>@IR+K?L(0@3k%JaQ6|SvjvuQjOjxq`mi=&h`EbQnPRas)RO4 zPaC z$y&@`yz#P^7k?IndP+f*3s%WRTsY6#A>@sd{|sDjVErN>zfL+M|J?~Kp?-|hhjRaI ziU?lGhjd9p}E#b62}+Q8K51`e^W@ zB<_NIht;mwZi1WWJPDa}=-mCo4jD;3<95za0dHy{OH5LG;X24@#g5DN=)z8uwO}_M z*eI*H)Hrtur?y%##dKtAW|ge>GCC2B!vMd)zew)g>tiI4K8*L>(kyDMkd)8)DIhm> z3fiAGu}r&BdeU*#K_8e?xozg$mGe`=1IwWsv+p_Yd1hbQwIE9`oM6r2-{?n%O!>I* zF5JNXHz8A^2;zm_6hw3s*@iepQIkbg%?A7dvsLK|wNt9xA*GM$3(k1ONFq?E;c73( z@Hu5;q8h}trYR8nBSL$99bbpqz0F$MUv2E+5cZ0zO2WxZnzIoLfpN>p39iNn^ByRO zF{QCZ5H!MPcmh~!zZvU#8#6%Nee*g!U*&)5LEPx*5KlU?)gujmub*M-vS?jgHM&Y} zkO|qLCpke5n(M+GthP2L?^hW&k+nw2W^gJvn+VbN_Se;@$l8|pu@nDdpm+KcjW7zl z)pCSp{GbjDh5|3E`q$eqf(J&fyDV-MlxgXfL%kO7Y}CYBA!MqLyx5n*T^g_mCJZ;N z$uRTpCIk=BI2=pc(j~idX?)zxmxK%*cPM+Yg+Ot2#xn>F(Vx91qA$6HFuw*UK33)! z9=RB7w4Ca?{#IgtrB`B0KGV@db5kA}eD_@7 zv|StHwS@`lj?9EUyUrTEui&`sgr_Dz!*#B-i8+hI^OZv7x8Sxeg{dA}s2x|dcGHuq zBk$G^140g}_ucCL`A?l#!V+>ymTg{kVgCDEn$O7eqsepe*shn;yzXLU@L>X@DpOj# z?%2)saWt~ZXM0S~=TC4d^yB?xo9okdV-=rxb$JFuR|J>A2Gz3N#kXN9(-WjU4Ns5o?u^NI-&dGEbmG?&cAtSG_$lS{En+O7%9r#Dc9{h*wD4+h zD+juIcN(J)Gwla?QNo)O^QXzW+4He-+{#JrGPi5+#Qf{{E*SbLRKpH_6CJ{ap2vCQ znVrw$MAOxsgFxy`t4%bbpN=-VI(btv@ zQDl7QheX6sZ&Nu>Zng>7;)hOG<@`*M2;xDsq~o17D|rsIiTKb4CagIx0G=BKTRQhCP382 z1W0At}kS?PT9mN4;H~ zVXz|6&k1n*tKyd5?rV-7qaJZ(wqN%20%ZEKH&`7#f7xe`uF{t0A3sfI@yv}Nb7F9f zPSnzWt`FlU`gYgxWoR0pF{R&|h11O2WZN$c{OIY>aS=!K1V6Ckp}NDlYlU`ryaH9$ z9(U5S=Rj|ZXxA~M*Xss@v9%~I;r=yEHfr=h-Ud|k`XRxyf`-@1u=LwdIFqUv>QRO8gHO=?*Y;7)z1=l-agvZWTlU(f2kTHFXImaI>#p9j2{PGzzVRPCJDhv{ z@|n6dBz=HR#O2k8`yB$ zwhgtWg_lmWpBEY`GzvQY{y%)ZRahL_(glcw;O|+PJ$k z?(XjH)44b2o`2?Hp8D&4g8KHZUA5M#RgZD-fj{@th836R0=2sXs;MiCwNfN0OcErK zt`2R=w%hry;QFL$-8W_DZ$iIWkj*BlijR~aI`eW(1rM3`lbC!pm!pz8IKaNk-wVAd zO$hbw$PgQ!W(w#ml6SU1my>%br7VNrUEJz{u9Y;Dh+Lity@h6HyZI zqjMn0j?VAP$b&f|-G} zT2$7UwMi_Ay2z2MT2JPOUCR>9G&P4v67c7^>Lv%TXjz+9A|~6{_*rYQHUM-#vodCa^OQb+13qx37+$t@ zJ|DG3yV0QTFhP9_XWccVd)X^vDeW(nQ&6-d-Qa^YgBu*JVu$c^xRW6`G32c&N=Q%1 zY)8^*s!3|!pY8^fEjmxvD4Gcof9g)7I3ot*I(@A8KNtQ_e`cmU_Vl;)pvMGr-@sCW z?hkUJj@j{YW=ROZMdPTuHyIZlIcv{Ildl-5vGZ$>WMb>yj&J{n@7;B;E6BL~?BUZE z?L(pGre~m6t?0$WZ6o!dhdIh)=?vtA4RwrgijM`^8^frA(ZkzWGNP3%&4SD7$T1 zlRq#24{o9T1GmUfTK&Df7pQQ;BAu`>pYh=~P+lOB9tR|u_j>w^3H(I41sR^79Ei@L zUM0n_i^SX7@9K!8UDO}>> zMpr05D?(jm!E-#cWt;V*dkF9PKR?zyi0xRN;Swi#TlHb(2}bO=`V z&h)-(o{*~(H~feNnD6}7r{^(!#>3spu|J&TDIMDV>{>GjtUaA>@4v74my0R_Vb|EI zqf7s(jzBptuM`0ux70pq~f6In*u?}keW`6jm(@7 z-VAK59ED%1?ppvPKC?PLd#>*CkBA)X4$QP0wYD=la(KnY`zkMLBA9E}8Wt-H)KuSOeV^?-bx(j!=FR3}&{$+` z()8$C#}ku&)KuPe5<9wOwJ56XEuncb6G~jr2`YiQ#efaSziQI+ENs$o3{(D6*mr~d z?DJ}#@b{lef*{TYrE3%XIJ@v!J zc|t9Q%b>u7CcH07-8OpDS9ciQaZI$!XcwN&4X2$wd zmoDz=c0LuJ=g&4t%6divb+*LyOinJX`a zs-e=GaDPx(vCICiGBf$ZkK@uV;r!B7Q7LfcQC zIY(>12lP?fAAiXxc05SA*m6!NY?wd}ld!lw&fF&LelOqRl#D{&?p#IKiQ&Egd_et0 z3j-%Pel|>wgbKge`m8O7u-Y`cIhucxV@VTH!#~QXg_AHidyEDh(WLmZIC*|OedM%P zVL-Ynm3aW5Sh<;2=Y&bdkOuiblucxn%N%5)Jlrb#+9KHsYyHGOZ{(QWUYRO$kQA0Ulr%I=H zNas@usk(zKiU&%|!LVku<(6=%p4&0n&O-gy3mWf+DXN6&hTRhVU4jpLWwo-*IR&;K z;zFLE&)e?T-qs-FcCPfma3G~rl`c&JN`0DCxDNMbKK!~yK;^}opA{}(IVm@WnDg_v~z z4gN^Net^xX-AI|Kh4U~xIRO!t@R~Jf^61?q`AJa!deQ{@z;bej{}l_I>*fHbTpxmg zKPVs{biB5w#rkeacEHc}1ho1MkHFLgH5kDXs5mDHI|Nra5TxWW9+y3XJtN5^Y)IyU zKHj&lADYXhF;_tKFUqDdA3f>`L-Y-7vv3UlOqyVNt+EtMSCP(9FT{c@Y{FS)8UkFV zJg2?Rv`*zY?`wZ|RCnbcMrUi5^D0U#z03808U!s@;&{Y;ZN_^Kde9$nxwtuQtMitf zU~e3j<~2D{f7>BeR%sl1Q3rFp&*RldPIG9L{y!;2AM`LsS7t~4(L5D&7HWZyQ~ATu zKIj~fH8)q$BMlK9lLHnrk`LaF23d5%jJ&$Dtf?D{UO7pbtavuzGj!s*aRPXDoC=hdruo|02S)n4|{tBmB~?fN*?@t$rx0slgQzQ%s3w{@V*JC zZB~4kV_iYjUIp++*H15OGXXxacx|_nxvy73>fPCopqU{t>OoJPNf?3Q%JP>co~elJ z>dpDO0k)4$($e-jDC?Sl;?I2f<`pJ<9^saI(n&f$Ci}A)p{y@%=qjHzmqMNNZ&r1q z`-BYIt^d}6{xtT=f8-7&N%dU+JxA_R0TgNOYL{m%AzC^+T>(?wh;-IZFb{2mRdoXA zAS~%Fd$M|ce1A`MZ=O_Z`yCiE#La%)q<_v@*tSG(uSljM_GV~dXBXYh9otrW|M2i| z^YEd4R~c~YXknia(4}6X1Cq}1l8yIYZfK=K(WMma8t)MYnCgZdF3fbp!JGccv~4pLUNS zucr`b5^hSfw^op)&c^fvwYvUSe{lY(KjgWNP5yft)WH>*GAx~+X1ZX#MG<$AVUBVq z&S8%NOxouHa+yTU6wN*E&WyOWP^un7jVUy>qEm8ccab8US&(ZTSI=5A*p(MIGuCRP z_9+zR3b^I_yyNdUsd#K0UO(EPC-UyED8ipQx(aSm+vdVtsj!#jIaDX2d-(G(AWSQi z(i}X``Mv2A97E}^`(Dd)vYnru0f&vu%=Y_d<-s22(wE(T^!$`{++S9oTks)to2T+B zj6HmDW-s2mH>GN$TRy%tU>0l#)Yr!&xz*8Wxn0b;%yf0H^m*H^Wb9*guowCf=E^|t zA7EY@5m#qGjd!sVUBee>`Ngu<_DHfgQ{TU^I=W-~_2||q zp=Wix)r70jt8#IF1+;QRPRrN4*?|5eI=!0gqUiFn+S#E7{(WVQ=8LOaW#i8ki{i`B zYzk3ePR(ZjZht!4-arXi>zrIHz0NS}%J07zG0g?}pWx?(&HmLOdLLnenHzGtZoyHl zdVI|&z6n*Q{;)yQ)N$*b@Um=z4;StXlpOwHR)jY3x_00G#}xmE$Ln=wi|6f8V#kZn zr=FYjx7|P6BL3`Rz&zT1xZK;LL@c@oZZ=aNjHF8|-PDh_N<|zvmf-f^TAP?q93F*j zZ0Sf@FH&2znQ5qH%sxEzxq+MQK-;vAr7;u;lp1fY2$Z$QvuB5We96R9rlF-^VIH4} zGv@TT1_$$Ik`b)k7WVxCD%pg@ zNW5(R+gM#3eiddZ;!ZglFQ}QuA-0%7UJk=k%YZ64j<{Su#HkDSaP$ zzj$h19I77h^4E=EshHEXXJRJcEPlYrED_ZyHVKEs*??B7$9w5~sISLQ)|jp7W+_(* z8Yx6G9he-^k(3-jVe*EbmS7-q`k7EzpnxYK-jXsg^L=QmClEaM2^^RL_R&dOBKDEvWWSt9H}n48cTd%iJ2=oEDBQGv zVCPda zMq1Y~E9(I&V`(q#0`c@ACUkDtnhtHh?S^vCa|7|Qt?IQGZA=qv&RmUz0);-7e-vS+ z!-P-4G-M!@=Ks9_cSgLqKe(Ru89MTis_x&?o&8T<1qo&c|JE)0(%^TP*Hn8N7J~(= z2|{qj4Ci#gN!MMC0b~FMjp!{y5yBbjwLbU~zNk76b&cZMx=#u~v}lIO?852A9q0Uc z{Edy(g#)*i`}6e>FFp%iK_1`d^2MOGPm+2~OZUP1nZ`o)Dwy}@YZ;@a&&X%dQFn)W z*v4bQ940&h0v_8J8&u~4{xJjc4Q}^9VB|PhGLATWStq$GNl)?ybJMVr*Ad?A- zKR$hT=_6rb^9Wa<3_4u%GZ$>r%}@U!C_BwIJ2~(f@_FDINs*RL>TD}Lbx>s<@%w%> z&n`TRSbWbWN+nI9x1xdqX9K6c=OG30{^K|2=7t6dN$;kJ{?!f;c@mcT1weobE6YZC zSgs9#2^X)7Pao&MUvYZ7IMw(B_!jY|rLhb5c6kzy=N?Wxj7^vvM=s}5U5OK}JGMQI z>Y5SNUuMM58*Dny_2cu}k!dkjvf~7xn|O@hfQx5M_ct+gc9CYYUpm8$4&TQ}jPDYhr=_Yw7M(UR}UH-Ax9@VCNDEbxMxy&+D&_Ds~fS1xwlg0Z)nL^K z>f^zpxzcpX3tFL|y7)Wi4Ks4}yG4GXojx$Y+qn0eXemen*8DnXWr43XnI4TbU*ycr z+i&ll&pSG5+BtC}Kgs>X?)0^{yEnIjR{-)_aHJLxw4ItMr?Mzy&_>;?U78!)T-Tpb zzIfdBjoaFm0id*V3N6-l_d6jLbEjJg_WQoMGJ6rty-GaP4v7IFAbT8Vj-;5F(w4F3 z9T@iQbTw)*u`!pTmgOA9 z0wnvl(TocRZEeczv zh|?4bW4;ASg4}fCVQ}p@=w&c>*-;pC=$CcsXh-AfZR?VXX2LYbe%M^ZMFP7T zQ_$nWWBcjL(R-oq*Ztl*eA$=QGYQ9|>|b&ml_N?#s_Lbk4_5iy_&nf{$D86BOMcGv zop8BU6fkw@thB9Gg}3Bhx2Dr?8?aqA-zi2Z+#&frO=q4FMcwci*jxu=>ovZu$TG7U zG2-HES0{iMccIH0p=8D=0IeSZ|4u~+iGkXM%ug|Gf0Y*?5wHOccL>WSsbV1sB0pJ# z$J~y;xya{?IR4-#!_6WyoP{QKWJO~BEHAZT)sa5IB(2){WSB3&fa2%q#!}9>fq1%* zR$SWW^l@8}D!NmPABL@tBjoU1{6v?Rv>DgVGPSmqAtV^wALSr{a|#9MaAmy8uB~9l`lguZ?w`~V1pE*3F&9VZMSpUIIi1!k=EP;x!OD%X8K&}?1nh)$GKvCyD9RYJGjS&6crTD z*V>+AGhaT>bC;JJBEvoxUjNSay@>^Yb~?SmKlSZ-y9C_@%`&7I;>8}lQB!D&GQFDo z@1VM?{VC{Fd5+EgmN;^7odHUS2QA-mH%v862ceD|yNZS0A9`UE5xv2^uY?Eb3B-aI z`K_-Xde4Z}c8$W83U`&PCHT!qg$8)Pgcr|rr8oKt z?1rs0=FOJ}hTDZgs2~AJ4P(B2Y~?q0G{0u4hj1`bIxu$UGz}R z!tfU9v_6X{4X3kd#1zJ&wA<@5k})sU-)!n)0GZ4DZR*G>rC)_6u)06L3d|R;{qPFW zV7*KmT8-?B`>NuYqS?}@`<2#Bq0x~;6=!1us#)Y^@*97$yaUvK)@|Y0RsncmTrwWN zm86D%=6~9Z*l8(a@w74&Mb@*DE)8>uphND5iG81$nx;Y-R@|!|l@ID${61G5bJ?A* zdEU2qe3q&_D~nLLi#9_@0D7heuV)Yghx&QrzP0@OOeTX+4<7cZ(}U@v(A7`FQS9N6}3Jh<@|xM`XQApD2*rvy)B(UJTvxgA6a zA}u3~N|@d<&NC?nH!cfIBys`y?h^%K& z>Vgo_h-rc_6Qv1GJG4|L|A6nEJNmj6Xu!tIEF@UT$=!Wa=%aHs3vX471&L5o!S$!y z)qOCtLe$D<2l~H~-Us(Tf*=0ZvcJ01PAE(;>ot%4^mD)WXW>bv*C{5SyV^{Ol31xB z2eKd!YqEQi9b)3l@1AHj-dE8GK=ZUblgv}UmjS=l`!%nQ=avq5D-oX+6hWmH@3R=`H#Nu*3)!&bslOl{ZI{cWR3 zW>qH}!k40|SY)*ysS9IG+TtUYF;GTA{JhN_E%1wwA@qnH2L>j$OSi|V$5uI-eTKGt z$p>7cEf#ZhR@v6FF(w3~DLm9kEHHE{Q-<7#sUDg{O*WJHu>pHHFuGURm!-ETLtZOC zFfWrfztK-3D)>`hAF5G_99E^4>z&R)4@k!52sA^VMeP$9SA;Gjyk-#an>WF@J9nGw zd>JKf9S>4)webl~AKe|wqgAkL^Qv<@UoEFwX0~6`Z}N1B+{L^Gk{3}<$6CxV;nE|; zRz!4xc0-qxj6zh0WzGKkTAQhWUh75YTD`y6<#-VoNf`1)h|tR?GZCKKFyxqUOGA9O z7*#omWbzBg3f)X02iEo9X#n<~5gCXW&{9I;*FOE1;10q2QvFw2#G(%B*C)fFOMA)> zp^gas^e*)aLvoS#2V03BrWUA9EWaTbR!TAFhHwp&u=!L+vO1RltQ7@ie`Y0wTABx8 z*HOzV46)u$PNIS#?22~abgUzayYDPNLYmk z35P?KMPVJU&HXF(6F>d>%1W;+fBN`CP`JB|v2j(+)ZtB1AXdb7!JtboS1*Fj#EI@+ zpdRB)ykFdFZ-n;15~>e#>N#QN3Kgppid@?H%ycEP!!DUnqD>}X<)fS>cNT(B&taJY z-}JGFT zV~=qr|MB|ffjt+FrcYD`1Lr$>A=pLc+hHa|r_kGF^d%|8*BLg{tOM;o=@v#u`eww| zNWC}uxJsvXo|cXrlW{I#wF`Jz%G0~N6d;Gs;I=uR4{p+kYY6Q-N*}}&_@tU=vpK|-^N(1)|l#64jLlpcn4s2Sld@eUG3_E zd3QPq1v#g-1H0cqh9wHanATZaLmK^rUskt%xJ|s*3ED%XX!a0OS-G2|PsKqN!`zKc z{}&x(hRTrt!+1`QdH#2Pv{(h~4op1ZgD7}lP#H?#)#Hxv31o%WXBy8*94dxqU7lRo znRTJ}YazeqQ^5N}^!o!yNCWrMlVpHQATHY(eWrxP)AL19!HviR{>S5_Jtq+X!ApI1a*GTiL2YuFO;jKXQN9gU?a+WJsBUO zGN9!}^GdlRzI{QB;}FS)xHQCa_7|~hAWjg$J5hb1sUgW8Pfkt5U6!~+fi*!E!>flseq~wx-fTJfh+rBYGmBZgX#RHCn}0*Ai0ILr08egic@MnZkTFq zGM=4F%;6WE9Hjw^uUoaJ{jt!8@2q)iCHZo!7wtC&xrz5SYWaywAI9xX50YhDD=Wt&P^6SmpdIgU8xBZ`E zi@hsAdlgB;6qF3QsQFKuUG^>QdLc{cx&t22-=jFWTwNPd6nKlnRYWxghE`uM<5N9C zxp>FOTbV=^Go zEEpB(5P5*LkGFxSu$b+rOx)i4ZO0q-WsLd8t1k+TiPjYi%T#!p44Q7QsWikB+6LUs z>zP}}h^GD#uG+yaTF)0sf$I|QbrWC_eZXc>wm69#wGs~h;!@1(%X%dl-!iV!vj=2}-Sn!;bVOb&rgT`@(X`U+rdBzDDy z!j|U3UO4zRZnD8Mv`>&nJ#udFR2B-xcF56|C*8u=H2#j^WH?g_&SK1b?G{x_zl2=p z&c(XhB&Y(aWAqD=;7koQmQZTKIVYGn?H4phe%`?5CLVa2#ToQhz<9X#qQQFnA65Z`VXUoqg;vNE+cM1;W|A^>lES`n`k-j8jVb1j&(u{ucljPK}Tbe#!Q?6ibxfqp@`H zx5vLq1E&%E)2R&W=Li!#IN?&qY}YKhbQ`3w5buzy{vR;t%@2A2 zg))5S{>lY50^z_cNDwi~v&>?*#q8LzXvHRSu}<8H@;M8NAN*gX?VrZjt}_K1hpowy2$aZ7B!FI?2l6j8T0 z7TUI0>w*~5(dGJ=Ayh*=$Tn-X@$~ciuTiQPAooXnmPUvXbdtrGMQ^_%PnikM1TV;`}fjf7#ZGxhEZZCiE}!x~b2ZI~};4cca@ z9^{Xl|I%2AIUr^{f6zSWufBx89AY9$Lq637N<4CbP767hCanTvn8ysL`oZ(fC5-GZW1WaH2FeTBF~-R& zWP)8y*O`ijM|)j<96IN4l4{iQFEEzWdUweUq&K!ki0VS0yqWv%x7_U&g|1|;RFtNP zjcY*d81hHg1t|CVnhZ(oUVanJMp~jDS94t8*e_#Jhc-hpo0db;J6kkyY5dujs?jH5 zF4^&+^vkGe|BR{x4yK!D+wARa;Y3_8aks$9{WkI0`&j8ohgmj5M72^I>6XQeVg8CTpXLII?A6sYSvyx{a_yt*Z zqdsjoGgkvmb1uPl;Js{ue$EOhSvqMvu8C}u*s=8gIw9nPVrxQUvHeeM;~^aR9ZRNl zk07r$18!D}c2K>WU5@A0U&hw`y9_?%edabtu49jlt+L_iemZM8_dWMo?>T%0w5Ip_!r4@h *sqpr1pR?l%rQ;?LZRYJ@fqlI zUWh`Hoo(-`TovkMJc62;jcp1%wy_Ou#uo9&%dbFH{>_LiX;`H6V`oPAZ~IPY>~&qM za@Z>^zAq-Sz`dyq87_|G-+tkrd%lh=09_6eJ`KXcM;g|Nc-y@g1p(Oj>~BUCyS06> zlV03;^ZjunH02bkXhV1NOhlaBmJ>#b60UiuUz32-_{UQAO78P*PD(9}i{9)pwzd_0 zCu;@`R++(`a~rnFII_v)wf|xz<9VRCu??A~|HbDrRKYJoD@Y>UhvG%Ziu(BiXvEwe zToxe^!Biej*WXSP^?e`vBk^wnLyGq^$Xna}?%n+E`|<5INI!PLbngT)i1az%pY$e; zZXKTCM)U=R#m>cAqM=-KD~K846dGl7NJH!K~e!qDMC z-IVK*WkK74yGee0%v>VGAb>X9p0VRhc_FZ!j!)eZsGhktD;ZD}7Q|js=qNmHhHBPi zC0^Pc7biCGFwPTsV(ZSE5uVoLPj~{xBgFrXW3P?=PKbue*}@pG9nthCMsJ>B5>Q>? zi|~&S1cV)fuD`Se1S%iN>dq5B@nsGS0a3B#`nIc>V&c ze1EFkhFD!)-TlhUQqg`}ZufdqC}fgWs$PkP{0MR{*-pE7dlM9?dhsVp!-?6k$|@;j zv0LnSBggKuY^gw+`HgBTGkS>hX&R@LJ5H8;k8~zo%$I>4I>Jo6ZhX4$&U{Whzu)0( zt^ji*8w`HNcfEuZ1~ZyI>=IXmw~klr81tHL>ldI`YWZoNali771i$}2%Qj74| zgtRf`SFyu%b2}2yEWTZkIpu@nujJ$u8lnRI`K+WOX`a1Pek4pEPn*boHh@FA2#!z~ zd$5q74zR9}@F`zFs0^XP`L6Mlex(7SpntUDb<I_8+kv3nU+v0J(r{3{2Tj}?DY#n3D#BR6lbuJLCI1mKm=T()YWeBbi!AENh>W`qTTzz8lfwPx1~H!6ZF|8%EpC4h(9Yb zST0Hx4ZM;|()-bSgWwn^?fTSzyW{KwH)pg$0q8m>|6=?ynl0cwmMdoZF20>Z%&rmg z7r_cptcSpUQ3QRcHMsv}dK8tw?FX-!TcMX8lNVk_DI_K+GxY9wRQg{>Yz^dGvh zecwUJbl&%LegUJczFu74J?V5h(In@-k(Cl*4v{wfoB}0on}xP16Pqks3FhL!op@XF zLR;C|&d&Hfs}MA#EGp1$H14>O7M;DTK42pgA%TyRv^xC|6QraMejpln)e0WbPbZO! z9e>^Y_49$~a|}-|9gcG?_dO44^-l3!k7!Nn?ynUpys4y`@#H@24QaT>)FDX8?h|jd z$zNd3N4Eb`5Z|EKs)f>2Zg!6-w`8s&fWxqh_} z%*aX5Yn=&M7_}-v&`EeiO%p#CfbNo^Yd85!KCp`Cc5i!jUp7n)IEf%l)1W_w8h-7X zU7TBtuKhrQ=2jYAIE#$2h%?Wy&E?R~(KJnh^&Cl+bGZoP$mie-tn_{yS+dIMw8@HDGBP$(C#Ab@RgPRtTwBD zXeiYw6_LX(8+Qzf1vZ*gO11EFh^N>+wkds;9WM}YE5qSNf`WL5b!^L-pn=XkR~__A z-0|A=aRPIuI{e?))6JN;TpSz*KJn^K6=B0)7{4Nsxwg({`Scxb)Ha%6cwY$}FQU&~ zm>52OgZdxeSo{Y_rA;pT^*69ostULv%M{n&`C!+Nw_`^O6zT2iS5wHcgf0p8H?M5Vf)f3BtJu5H$EA|m8|ON zAFv!d)9?*kkV>!}mt6)?7XFaD;+%-uP2r1r_+%fvEgt3VUV-cjqwt>J$O)jJN9P>J)S+WvuE?NmCAG22&!o^pjCA`NsQ_?%YMC)^v@ zu++8$$rT{bY2ZFm0@=W%z3G>J;sk-AY2%$NhDwh6jyx68j}AW`?%eEL26fB!O|LNU#ev5BrD?fvqQWIy4 z6%4!_McXsjxm+8B{CupOd^A*h+rpTMKAt%@KUnYE4E;D^$vOjmGNvYSih9|SDWG(y zMd{)VZ6cGo!AD0%z)`rO{>p8Ypc2O+Sv=)>ugansnGrfCXOa1-yh3L!aIv#TKKkc zfZ}?Jtft{8rhWG2;*gj8bF!7Sx7cXtWT2oW>z+QaZZKu%XA1Tk9heP23!@Vp`C4=7 zkBfdd8e?5P40|3lo-@;3+kG}@UytLc*N@kZKpmV=)p45}J}rTA_zu7N!W^dX#j3CC1)A|D0UX2-UyeU+!taAObWV&;-OFzJ6Epe7&82 z-%)(~&d}Rg&OJ;RzKg&-|AGdSvK+m!w*FMU4v~}fxf2#EA92Q+DDATxwj$U@E4@T4 zDvpl2)9Civ5@iG@08p{o*zU5;BWFT)ea+O~%ON5skUaL0j*eg$D#Fmf9{B9F?s+{n zWX(Iq2bzV>V)A$My2twYw&6}}Whe*68nOFXS9A5-r*WS>sK%%Mv3(QQw26LWztn;{ z!%{ACX&)3CB-W7~Uc-88GMt3Y5n-}#BPae;44P2kZl3<_jlPiG46>wv(nUM}JC@9F zTans=gmbRSn@s4H7Fj{QRkPT~*;_`TfqgFC(jfZO42%IN40$igt<$3?yeV~xtW@%1 z_!)v7oZkm{xy_q$tn2X3^-QUSF8kBQHrf(6sJiQ;=b##M;Vv!{oS8*bHXHm9!uS;&(ex9rH~ z#sp^z(Z^F40IE&nXPo^y8M@p0|NT&{wPZ}~+SUq20ry|WvKX?7zCz~B(?*@*7!sqD?_0QjIz^`g=r+3QV0ZKmUsrb;nF>RWfltX9NSfoiJave0V- z6bdAMwW<;Rr$GbtXO*vP|Db{EC176hhU6Q=+i;%|F@Ck8V!Y5LA)3+D{cyU)Uw`;` zrW9y^dVT~g^Z;HqLdPkEt^@bc_YW?9RQVRxHX%&nnTKP4UBHNC`puLkC+t@&I~8caeG74oPat0c9(*#(`hZ1R*v%hH*50Tm zpMcyHaL0GvDpngG5B_n!JcWYL^tk2}KZU}uLs1XWQYj+-ciEcUDnJtoT6UVEKu93Gj!U+;lzcd1+ z`9({N_-$`o;-q*KL7aCN3RNi@Xc5YHmHMDD#NgX4)YZngL~$0>wj}Ijm}NPu3Pn;g z7zHXW)b9je8FN$#j~YHfXwdwKf9q^PqQ`coF>7288O_%HGCaA;JbW?PV51iC;KKR5 z$y`3*b#E_1-0|uMpU1@)6m^!eMMctxtvzwstTHZB4L)ir@#9K})V1a~d7vaABze)J z-@AO3kgp$KP9IFy&PV3}>S;vC=WRR$;<7t;9PHnSM*lGm%blrq`^(KHR)Bqg83`yh zB672z;j=IlErRyWaXf6!`}or(4`W0BwQc=bK&OiU76hesL_#O&HrC~l2@h70Pz=JM z2l~g*>dC;2#o|J4JZ=so@NGdee*(5uL=U56*3d@AM5Mk>}*P zEF;BdCvrm@J3AX3+p!!w61@P;X2d8gm&gSyj)-D^f-0ttGaX$<&KnI~erS#9_RztaIM})lxXmW55EdY8GNGM`Vmyy5~vNzxgYDpTJ;m8d5g*N2Wiq75Bu^ z11d`p1kaw*h9{YX?E)+!VoPl2n>EvNmL_Tksa?}p66)UBSLXl!;~C|=qxHR>gbrt6 zGk*eG!UUTo`YS8kf<%rmwq+KjPT7aXoUT7hv&-##jwOYjP!$YV@?<+{BP65zaHKvy zQae4OFv>h2t{2Bo&4Y$48MCR38x}2rk-W-g}Z7I+hl*)2{^9}cG{N<&}#KK|slf{>mrfQ9)@q*j+lvuB z#QO?l5=X+lESMj+qjokwHFhS#M6sxIe5$A;PR$;y4)?b45oc?D7Gx4$BAmP{e=!jl)==gCg5K>dHm5^KM;xt#{3?f zun8;Jzz1uK%VvF-iVDB|!z=n4pkM=z+&Z3hIoGEWvkhbo5TyDsLy6ItD5Qo21RMhd zRk}GR@Ennbe>?T9w3ZT%6gm}*C>G#w>S*>D!1ZR)hHLic3!dTB+5>rnETTIwn4EWj z_o9OrJP6U3X-R#iwo*APBcDIWzZMClANGb!a!aToZH;y^F($9owRAwq?@(h$6iQB`vlG8x>FUR zEN1F{ZHco1-s5pU**HTi{!evpfB1*0YR==#`dgGq(ttsUbbR>rw1wRPvk_3%iS`1i z*Z!~-j)G;9H8lB)pkssB1deESaMO$kHW-SoTju@F?!BK|7(4z3Du)^}xzVrVWtXXE zh^XfyWZy5}rw{g1IM>w9DKxEI=?Y1ViqFhhOz3E^lLsmI2;cn3?->n%n%bhLZL zRxv2B>BmMI6n@52plaQV`^R=(4bWC{lj_W;FVaBh4W*W9zcN|pzYwIZI=XO#hiI$+ zHg)yHQJY?L{xP3;-G9wvRt~=WAc!Wi~QA7bzCdA9!> z7p5C6?jn!S*%A?XEPo?|GNzf8^<)kb9?BDcT_(#Y}6I^5ZPx{RTGrxH{u6*Al ze%k~=4e~g!I3Gj6$g5EHslJ7~{bJ?tM_~uKS~PVTXv9gP9l*T#ydLZO?Z&HLyiTo6 z#?gOBGlylL8ug+qj(fig3~j<0p0-1GxQs!g2k_u1aT7&@Oj$V3)`O6<^$`X=b@9YH zssz9b9ZrdeiK+=!3JKyE3fS%pY6;e!&})~M`7U@0y)*^+lB#30quE7KgqJoAq>SBT z(hiG4nCd0Aq^$<3<_5V_xPD#pkU!Ohb2G54l2`zuU&%k-fQ&Y1hzi#=^Zm!2=M}3K zyY`N?Y5MnXKM}G(jvX(;EcXxaI;@@}R0nu|*3(|@^ex*}Pn9!$bNZ;2-;84=(lMP} zW}lG;vmqnCZB>%8r>y@U;-FLU2f{5aj5PhLG-SiL@b2qtQnZG~ z6ZoAlWz>XzpmExqoQvE$7w2Wky0W)M?_DUT&t{ zo6af^&+6_ZA1$JGfO7(JsXCX|B#iTu&n4PQrAVIgZ=bD@m+9J@L42!bca@B76jL7| zzi}eCFUYlC^1Y3(EH5wLi{*>G^i7lF&XXl6^W@z%2IEWhve0vV#snTvL%g{d@q~SE z3Kldp_i6xISW&T=`geLNt6E^G=r^jg_YRpdc6<^zzol^s$IuWgMZc>R4^Q~oyFNktKU15KKZuzWlGlGsutA1Jd6*~_{4@;5iAl}$ z69iIx>xVC;*S8~Ak019(Mm`GFf*o{MOb}W({65fObl&5-_e|?CNXv}q*?xV#H>-se zRJ3da-s}5b^l+`Y2jTzV_IhcW7aD=74b2g&HcUcXvyLw?bzz=4l|0*9OSe)tG+02* z=@&;}Z;(T_kZitb{Dxc=zKkH&osh>pe|AJV_Iv0;|e~NHPkt# z;OprA^RtUeu(JO%+o!+;KAYvi-Eb9-Z!gxJ66ShA>GF6IF{VUxz`4lsvBb=lCS`MI zZgox;V;F+z!67-eVFwP<4lFz24)tdDoalXMFJ#;Gk6ujzk9oPVZOkQu4U-uc&4~7r z$k23B0X=Sq=}Z`i7HMk^P+UtBz=Wb1!s1bK;L7w%RK?DWuU z=5A)qn}pgZUtY6usbfw@+&$D-yeMhCKwc@O3;d;HlOUq$Kvh5^rP@FI&3hmOu&>ZV z0RAyeQuqMXihug)<(K-%8HQk4*S4-;*+(KsLYA^=i|m@q#-RPX-qp+xV}WH47Mr~G zU!`=0QQ-)_oJVz65b`gCeD4xHYTA{e`p!p)w{d@?XGDaXKw#8sLW_sU9;e@;Qy(RJ z=^L0xMOF9^a{04+XT_j5$s$dev9P%i5bC)DD^fHS7dd&;haj;{^a+lHcRLe{mroai ztQ>$US96#g@jCwXdoYB&{_>)!J{Nwng`g0BFP}j4{zBUQm$BU_JxNBRx%l^x|A)P| z3X7{-qD2FN0Kp0F?i$?P-Q5%1-6gm~aBJKlcw@mKxI2xzdvM5Ye%X7U``y?3c3#)& zZ_U-UYL2QJHRqV^Pw!^v%Z>)tExoR6)+oYHNNFh%jG+8}_To_m;#l3j%^z8058<37T<9#;1 z_Cwe8DA<0kN_Agj^=a+59U-EJHVz`|3n>eQvFpW#@qXn}q!Wb+68nJ-hHe{I#>J+o zrt=1`BeBED@{~;;Z$>5tv;>405Wao40xzD1nog9!+8ZAo*}PwQ7_n+;@;yRDkZEx zPiDWzo{(W{2p~85+B-}rH&M8G$$h%4Es}u5d~8UTRWVtMKS=A48GaIt1wLR>8M@v1 zYDo=VGO!5d@v$C1P0SP+`q_!Qg6%2vmsv}qg~U5c5Y$mI|Gu)40o%Z3!KEX{R_JUO z<(J8=Qn8@rymW2PLUM-N(}a2aaf@DmZyK^vNq+VY*wk)2vF01|$K^?U1yceEzLOd{ zxZ~jjCQ!`lCY9;4O~N&ho25pS&vWPji$;gbZFYw4(-mmVyPL}D{(LsK*yYP*y zVz=#~18Y}8!#^ebxTUP+6(wl^S-GQsp_xKr`lJRp(tDINX@jC0am@mq&6zJM^Rv-T`& zH>DRgx2oxntiF!7fM7mwJGIiyy=EOmte@{=2vRaP32*cBsC_72!aIDJ+V?L!LVKTj zds|n6oVK9+=7Mh~%i@56l+UM zk7ofhr1uPr)%04 zL7mh4M33>)r%~{9Fm_o2ZePxCZ8>d0d-Xa|ex!?0g99l|=PPcn(EIv6Llkr~tmiGS zzUSV5e!KC4h*ujLvmd(xO(n6U&Hdo5>f>Y19lAco6y`h;Fk5y`Bvl)DR$xWbZuEX3 zx!Jtcap~1#xC9<((|C<^O%1BF0Lgoy-#fc8Fu!I zkIZ$j^NJghPWm7J(5FPo-ko8%o~f)Zl~=2)!tw2fUj|3bMul2&$1(uzWwezro*XB} zZEIa&doKAn*g*ntuF9r$rl}5IW_E7l)M{$HYxYy$3fj%`VFPhtf`Mst%!b#;!FYGFE>+k5=Yxdx&$ioa&L-0Z@l?J0 zh_>@K50|0k^RtI+Rml|0lw!t3_rUPW7tS@CrW?TZMCZm9JVWtWn!Rr9zHgYP+3rQ& zdlPP^79J*~Za_X*dQ-Glvn+JdMJE|Jd17w6Z|1ANyznEq(Pgd8dEg}W0W;!k^*GE0 zI+Cv5FWl0~>|DUQK$ew(aE*Im?KRr~uqx2`lltnb+^0sD&DvhglZ;M}7GFaz}HiL}_c+DETQz}L{s$0nEJ2i9BuQ+9Y;OZQIdNhg&_ zBPs~Si6|E!z&7V*O{AA4=+bv~eA?!I&h78@daxOCh|y{O)BjHo;ndlbZG9Lf)*kjx zQ@?%S!>y^*2+-$=5V%LsIzMD)ra|xzGKX||K+!%D#L4dd1I?9^LX}(b>pkT_=f)z6 ztWS!2S+F?r`^;$1`6o*iXo>4bd^p^38pPH&btP20x)Qc@fZk0-RljEZIXz=eh`@e> z?ke@M-Jmd=OMNb&(W>kDNhP+2*#J`a(esMidI!(XKg1@RSqr`(iGyt33En>{$6akN zQa^||Gpt9dSySTrOM5kQmx=PWYD($+pi&()pCjH`Ai+J{H$l&t4)pcz7&r}4XvAIT zRRX1kvsy6U>&g`0-pGNVt5|frVEJn)7V-k=8-Q_|$l+tiHM@-~w=>eJs?Gz;vAf6M zS=3mz1c+p?&g5ZcHh zvv#7|8cVFHn&8$3rlaR&lGr<*YI!Z+K7VR$>csR&E#60>$6xi|i0ATfG;<64Iuyky zFv0X-;B+|Kop~ApMDsM-bnw>wz1VX`Hh{!LpeNAkvDtBb?sWK-3w9njW?R1o>f&~r zO`ox?Jy?k)xUN5M@9cIQ( zYpQn8YOQ(GWE0H^gqTJ^2?(&6obD-SRt(zuSqEJ8uq=XpKIM!nZkn(K)z>qCWI*jz zKtpaBHJWx?FDuaOl{{!KHJ0$2PM%G$roflN+OL+z?LsK3%eE{gm+E!ZRd zj1QLsjr_H`-h+0Lai3AHB5s!VW3ezR1?#eh)?^yl$_I;JmEJ9Bcunw5lrsUG=#9=m z?f_y9@sWI?L^z754l4pnCS;=D-BN5 zxzsfW`GdA%Y98Jh?S%T=DDHaR+cBGU-g?vZo2jxQGI)V*pWe;PRABp`N7Jy0KV45u zZ3C-S13L7pj<0LJSo#+ooo$Y8R(|`>gq_80KiS$FxEkwD5~pO>_AzET4wmmqeTG_m zQJ2r}BJa#J_D-nR+RSsLzUy&^*#X}_)?6rMAo04N@%-t5XLvO5x-8m1<+nFBvosry z`+%r?*1$hfKjPQ25H;-Atw{}n9J7pPVtiP0maXagqn-WRuhMpX-h-yDqA}C;>6WZ- z)}-0hPSBFib(2E!49ML1nPE8Fz+(MwnH$8GrLT<+BI5DcSe(9E`Pr1PAIr!<$mip7 zC}pi}Jvg4?of$#*-tM|r#|86Cm=qU7T{HD9=4sRF(o8*ovXXubc=1|HMufJ`Xb)Ov zTLP*=3Gl9edxpPT+rH=)ox!RMKZ7&E$h15wSnwO8;ft`7R4t7CE4_J(*D0u_2X`qf zw2^;bfq)I?L~{yvCKv03Ocv|xmVBjFA{pm?5GAIU!!6=TS2!FxPAmgEy2EesAt@`h zEh*sO@eW3YvaLRk*G-hd;d3lOif)MiR#z{4*h}lglfVT4%ulaM?x&o&R(qfF_^)GV z>eX65A#tr-1R|G5jT%$c&rN0XUS~?cuNCt1J-%O>vEGML>w-c3y2(N49hY@a?UobhvN60 zkSz=D^*Ry)rJNBQDNN{_5j^#t`*EhSN-pO&O|T?hbsw*v`pmonr&hf?sga*v3QPxX zneeSJbk&_eM4dl(qU}fOq_xN2d9|TY&A7kAX`vx%o z=)&uBZ=#~2>6#>B7|v}N0AJ(v!Bi3+n5@bED6{JKAV<5_>*wpZ+K z13G*fj0#JAPGJ^zJ6WW=w;`SYYI+p8=lH*_bik(o&kfSj4FMUMOD*(>n$&D>aTCcPq_G+=pz4U$|>v4^(t&wD8%t`}Qo5_)+Of_7{NOL*JA6%aU^}_qqO) z+rsU$(#8Fcx*wA268}m3KY|0VX16ldKS@thcmS+qK(O58=T;OdJLMgE6p25>hDP*l} z?bz-Ptg}lm^4JJ!Zv!z2N{7Qiryn#u1d`b=j!%F#>1XJ=!;tXR8oc-2Dw^pAmKCp- z2ckf)ct{27jpm%VOm{r-jnnD31(NgL1lPQv<+>2Vf!d{8(e7Twok)D$T!EhhD_gf8 z4kt;`82pQ%+Xrkq2|1i(X9jp~Ph2{9eMgJp_`PjvJsMt!RY)~DOUVfhY^v#)#_5~* zf2o2%QDdpfAkNa%<)en|S^1&O#92DXht2&E@gB)AWI^PdY6DC>YsgDxFEso?9R+VG zr!8wO2I*^Ha2wZ2e$LXAz_6)yzrYJM-Qxup;^~dv<$n5xCdXhoQxp@MULRU|`uX2%rp%>?c zeesR;3&vputYbwR<^6%Zzf4`f4m*aAF1nuVWfqw)$L@*3VA_eaoKM4x#i@dH!`rMS zqZCSF8!c`ZMV^A^PJlqB1*R8f7)95(>PjR*xLE;cN)CR)u%`88*Ido1DMDHVY@p!U zF@M5&x;|x=S|6^(5ij_*J#5L-3)IMuOZMq?tpXihhfYx;|*pLvPI>Q zXwd8L+jH2T$@G8Q6gy6+>Lqrd$myS?MFAB8J}ttbu5SVIM-c_;@+H$dWIZUsY>-8( zxs2Q}Q4x1?Lzg9D(+>Hb91ql1uY#OcEaNvgIS{YS#Txs}R#$%+*|?;+&Z0X)8<3!> zzchunDrx+5b*h=d53hp>Bea8=g5=f>ldqgz?KgU^Zy>ki?BH?}^IVbZ>X*>mIfB@- zL#6D}rTi_DLV(>=r9=IY6;YecQ0z_mAtpmq0f^@aWz5dLEg~ecYTn|ejX5|j!p}VP z7ks-E-b>f~exrg!Cp_eQ@mBePW+?lzL{2qZqQjP!-fxHPFE_{T>G4K;KW7zGBm6Xx z2;b4pWxLziA74ZSRge)i;k(LIwk%6c97dBp){OO{FuhMh+M(~PXKiY(t=Hp$ z6R3l3ZfCsl;B4`#CR4Mat>JS!8&w?$=&VRbkRtV|4T$CP5?XbaUu*jm5@-Lx-f5Yx z9ab%Nr5grrj=EdI8O%2(?Qv-u}_)Mt7;$6w|nrTAkGy9Z% zYNo52vV)6VQm>3nl(3?dn9p|stoL@`W6WT0!zheNtfN3Img2`Xi1JT$vjJ9#?DWX+ zbi1(9*-tFY_RpsF@TGXi25WB2IX|UbGAw+;{b2Uwh90jYCOb8;NkMwrQ^{Q8+|Oax zHTVD&4=c~`hSZnX@72bhFYfn)^=lL?Wi4%`GXLQlo5I1IMH8$t>OWU(f&=a)@6Gen ze7=)KbzMnTRJGtU$UI`g&8y=Oy&!#I_>68P{RGNyaU=Y;K%!xzd~})S3_BExv9W_N zWfVo-?>t(|fJK+-4>~UH1?(4Xwr|ku4omEj8OmGe4iI2cGK3)EcuqA|jR0Vks-`(; z-ZMv0T`7VU+Fl>j)Cw$V*&;t()`dGPD)Whjs)9p+&TkYtG}+$3ub|Xu&}H-TIXW^2 z$19i|FuAWXIwu{mpDH@nM=skg@gwm^UruOTpDN{2`Cwf+z>&5nF`HFwfP5646~^V` zJu_uJGLBiT>(wGJ!pCmGGgCvte)34@7B}%@RO|A$+E-X#@!cR>9V_O>-fz9-E^R*w z{!UVIQNR!*-ns2qw(?J|s46M6@_nualM$)2pZ1Cym_%J69*4Bz1n9#|lomI==9Ko! z8ps{F{1m~RBz@y7rJA&Us;y&_uyZwA)7sYlk}>tYXb3pS_tULyaXLFFSjd4Hs2#&q zspHipjIhL(CLivHmJ;j=!MxHJ{5~hpFZ`9+l6#X?EE?`LG}=s@S1gq(e2Q0REy5Yp z&F@qw_3Qa$=q~<~VzGtg?!7%^8@_GcA_mZw=+2>8j}I8>51SQ-Z)-1K9qXryJe0vF zix=9<^iGz`Z7j;iefe6>YLuH%fVoUt@;_OCG??jFB{vc`_`4vl!3`w*B)wReu~C)K zkWO zEbY07>M@CY>Y6CbHVz6H{z~Uz_dy$?bgPl}rG-_Cz3?=WzP7I3DmY3 z(IZ*xq*wf@%)|mA<*`z0Bq+Vr`Lrq44D>Yb_Nn3A(%jp-Ifh;YS?}n-`~=z?KLMwe z@E1R!1Y0R#1X~dcV^2!t7)8OG{?M1$N&3rAV@|ib0tdfSgq6czWsltJfoAx@1DgX4 zXq=Wm-eM?obA#gb@Z~DEQ_~Do`z7;ibo@bUwj+q$XJa|6WOkWF(5Ihv6xyx#V@>^f zkL@AJ=h1Te{Pq5A)H4(YcXfrzkUh{8*;0iIQd81&Pm!Wb;^)r_?X1DX?j-wsPtri>F;kGvYmk5rLjXVnH}uib!RV=hKA|?JYQQFT)o#rnzH;w z?B=3ysX|eTTPNOsOr*q6_xpF&ymiHn-`<3uGFkO+OE_Sn99UO*#7LB9#yC-xeU4E4 z5|ct%a?C0dU^AV;`Y4@0uot(?m0~2nR~l*^X*V|_?FC@)rMx_s5Ld$DjP##Ks)26R zEED^83w|pE!D8d`kX}{UO_*_I3oIC2FXAi7PaIfi`ZMNpF2ZgDdM0>Qv$N@0gR| z|I{R^9P!WRu8;1@=yDp{#Wl{1;dmLJz(sm-Y9}^U?=WGGUK)>FLtF0=C3^xqPB3-V8vWdD5QRbmES}upkCN{2zOR+d&vm zR;T|j13hwUu0p z|7PtU1D&zJGk#(fNH_K$+y3kRZ|xBU4DfA5jI#WH8YayJBsFI3?qj zl${~eQH5p`($-xoL7P_+O$=^1vepKkWr*ONb`mkc$OWBl6=4C$H|&z<*cvIC`j z{p*e}@9MeQ*OG+t=qjX4U^S&tIM1zdMs5T^eFBLBY+|IdgOZ@$k? z=&JrN3jF5|_~$hkZ+Gc{R!jc@(f{K%?pxitt*rU~`D9&G@ZJAA@PF<4e@*>gTl`<= z`VUzAGg1GV$p3db{SS8i9{@56VqE)I3*di#^AGa;|2nl=sw~w07XX`(A;BBMnxZxi z4-Y}&c5wf+S-8SbN*mXmOKWB3{x`_{6Y{}m|JF^If%@@3apFIp9vlKi-_Rb3U8Bsa z92TG7%f-w=A-CZ-R%jkZuorFmpDflYDE*xn z{311K9-NXr73$EBr8Yt+jH^eb-~OU$<4Xz^E=luLL;q(oPjsAi?!xCV?#-7*G-1*%d9o#R}u#p zu>g0LV<$j%_UX|j_4W1jixXSkpM6a1qWxe!CruF07HOAG8L}_P9THk);6ILFWGvu(C6~E-R!yRI*vjOBl%E;!uf*pIVKwh|>|^+4m_@oD zGt{~~xK+p&9MqQ-m_HM3}lJo*o z4-rY7Ar}q(+(DH%bjcY0D66=-d_nEnZ_9WV0y@_6t-Pth&Dc%`S@6JdPa~c~=>uGi z%+?bp?2?jhvZXAv5ZMA;8Ou=168r^1C|KM)VP_IB~eoxwEN};q9I>X0}_)!H+->mBm!+;K}R1&Xzgk3&JWb&3nF^%Ci2;l^xF}~(l4}> zW2PLcdIrNKmnf9N*%7y>gnBzlGZ}H0CU-vEc^VYe?2L8NPQ_M`XXRAupjtxEhP=Jk zOx){1T*u3)J*Vk~%dM)wuByQKwLn~3Br#I7peX(2`JdNw@VD?Ek`}y9pZK{IN!M{B z+j&!7@k57+jc_0MkqK#C^2l z%dBD}PNvsV$tq4PWl}3b)6}!fU)%vRTp)ss3?lI_iil`&97W2x*56lc0;nuJ%nE2$ za7n;WdDJK}$7(3AhiuUg24p+r0_GGs=mngIaJ@Wf`Z20WJ!O`$&7o)IuWu(J z8X>;!{PWa_j!6U-fu3{hHg8vkYFBGCtFR~U0+8nijRs@0<4#vRIrnbn1&xj3Zh__b zw);;}U_l-ivo8f;836y+eSb`m&P_!7`_YN#Q7VDU2+#2j!mQ!9l9^TV4M_%SxK;vD z!b{xujPP5s&>yEWt0PUfBerE8SReRs1pCZQotoNKPgl?V?%LLVGo~&W`oFj?=%;$1 zQPip370#;~0M5Wm-iVNGty2#B3kZLgN>5%?b5+O&4(?8|?&9i(R9ciqik%#J-+tB% zzhn1A%Qy5sZLIeDUBRI3+}h*HC)#@$yGxWqFg{+mle?q!IQ&-y)kO^v9F6Tu%e0Tg zZug}DAEbq|6Z7krT=$%6&ZC;48|+u)uISNFLwHA=(t#`Pw)9JS{nLP-+Nh2&Ruq$0 zfbqWs3&H;y0lA;aqv~eC|7C&bX>q_w7wY0p~zi>f=&t?CH6D>F<*&<<&{HNYtTHN z0e47OODvcL*!O}@d9483{|{IvhzT%VA&if_FwX7BTydbF=O73;chWp);U8D*sjFJB ze>@2Uy_Q!$(vEccpdsJA$z!lrdye%IyP5<)`rJa@ z1!pz-Uyoya9_jz)Tqg(3C3{`dSlG+1DO1ZAsWHO#?K^^66l8?{Tv6u6%8YM%<*x^K zzkT-ppj+Fimnw65oiT68r#nF7nmfu-rxFh!Du3{2#n#zn3 zTsDB+X(=!Tzvj&OwsGt|M|5A`1SSW-I=<&%7rDFCa7EY_dj(6Q?#|B(euR&wHd9m7 z+Wi>_XS(?rpqjjg{fqa)4!9~K`m3B$a>MfsEITQ|w(H z70Y)DD?;;46=){Xiw*AH6#DWjF0?#M))mp6BhlTPzNOR7AMi9!?4J(fKQV50&8V&9 zNG{8J+d{>muyPG)x4J@Xe5xfNH_2mBTKMDPSBdX(Nmk7f zo~j3Z$-wVnH-gIZb3H}ebgdV7$=#5<$<><^QnHwa43npBqcaGC+qJ(b7z7l3Vy|&a z7Au^8_)vWvWf(O*Df8?31h|mSxsmpvo%-_tmgyVT?6?A}d%1HfL&-@fLj6WrYPJNFhRAi zT&YhU+FsPUDWu|;u1pigHjfy)iW~aifDyB~q;&(96kW=B9?BUNy@bZc!973pc`WgN zG)uP^TJNUfIg93fttEP5xy<`&ZokERTG9Df(Fm{hJ4L#9_mG*Qg*?r{oKSkPZwIyh z-6BQ7ecQyJ=LvIY4f6^^&*w|9$X2!<>AlP*b<+sz5xHsAJZNfiGOts4uOF=6Axe0K zayyNGa^y6xQNEiNCKNDEbcsromlS>h^$kZn*sFoT=++A3o6FByC%SrYSRO|Sl#eSn zr>-u?f^4FN1T0+P@TEv2oA*LE_pU?-(S09Jf<*LBF+#&xKa7^3q$0;mR_M1F!O48& zdl*to+1MOCZI(X2xw3^-`UvD_-P?rfYJOm3%URI>?1;n~BnP`=Dw4uAf30u8dL1x1 z6YsL7h5^TiS~|iuT~?w|9`(!c&*MU-!83TTT3)tyl5xSJE{3aRQNHAqUfHC+MylpRj&xQ$l;-<6(^D{Vj+D}NP9jbT_okJy} z&qxQ@{|X{;Aiq)*7q^U&N=fa#_G=#clt*-HOtkmO0w8&So$B#l4J;WR5Aev&YHszv4OQ@Gp<^RC=j?J5lY!6XzO>r z>*m3nKaV-iG>#IA8gvN`$8mJYbI}W^O(7<`w$KKz(j!V_eW*|96XDRKiXwN6wp4Vz zTZ-R*16ll$_7>;Cjfja`jI~{cz9b`hkV;d^gp6H(YG41fuWPxIHt{%3eS(jQR=7h? zS>vto`(e(xZUB`2ynjyO+5Wfv8j9nm6y6huk~HLeftKk>VRey|Y*z3c4X zz2ILVq%!u|H$gm}cGA0|aGWM_82>8aGv|pruHF&!?xx(?lIXWG(QYCJ5Yv?OXRYGl zFkLk|sPybYKK4y|IQ+fnc@Sw<`;i&;W$t3@D7ROIxR?w#_aXIRG;&ig@}byF*zqnd zznO-t>!%RI$0EZk;;cl8bAHB~L3EP%UPWCT4t)o$tHvSujv=U=7;c~iK3`Mh#&?`5 z;9|OddTb5+?0D|A7sidy5gK~D|F|>2bQ~gHde2-S3@!SQHa##8E7<7%O~UjoQL?lD zviyxc%xo^48d`)WPSn|Fi^?y5>4QMzG*Jz~DfS^8S3v@xa9Vq4A-y*&_XZ;4rsWEp zZ?rKR%hG?Y^695D*!=9f)-ZidP_L*nG`qDdC2Ol8%^7Pa6dZ$Rz{Z^CZr+l-Kv1X?@BC`k|=<3uH&YHXGvEw3ODY3WBdY?#OWc4 zx#Ea)9$>7L8m;_~9~CBQ6v`#8p>d?47J9weY`Cd-Y_!-eonUigp}XCM0c<2?yr` zh(bSmD}>u$((nBQ2NrYAc;y62K>0lvDbxHwpzaO|WVsBpdD2>uSy2gP0mtf@Q)-*m zaA@H^JR@94XsY81`>lkw8L|w?d>cZt+%Xx)-a7<=E5s?kCgiX+ikhv|pxaVPet9D5 zuL8uXJ=$g-ZgqfMzL|U2<$Q<){6@7Q#P9X8P}Hl*N5z{;GBc_ocg{5x zaAvTK{I<24$8QO7ZGX&@$J#Z(Tq3XSW7R1_6UVr31z~7{xSr+}8(PCO&kY;4mt0&X zGj=o-KUhnAKQ3G}sNR?%sf47c9365W1#-K6kNb;InGS3;X&2sQELcplWe7#l)Mvlz zUJXf0K+`H||D-*MZYrD8+=$jpR{fdLi*GEDiwY^|$#F>E7hJZqRP{4rqzf(PeOc~n zd6X)U{V{_74P7`E2S_VM6(5dw?DhRvDYCOAoYt)%J)F28V!@?Mn9){!KgOn&HC^r1)r;lp8~GYQ-9?d6zqDPLXWnjj8h_Kf>mDNPu3mZyKx5 zVCBW+vbe*I_kkaH?H8|nenx^wWGbxbin}>-|5LWFnP2VA#&_HX&DNqa! z-Gc7+N2oktzDujsd{=Y+3{&;0T=NTUOX6)3p@7)}!(lMNjE zF#y=OgonR)_EZr*rHTdUL)JLq19m|=6dmUvtwlXhLGvD-98Lqha(`iLDe4G$jtQhDRZg()H~H~d7>NEeciqDv?^Nj z^8$g^epw5BfbnkBf|B-8S>?`Z>ye>4T(_L3Dz!p2KyDfmgk_#IDvCjY5h}F*X1 zfh!g6w=FdLVywyEOCkesLN~5;J9c%skOiV8H)V>z`h5PA3Yym+)|RazIK5-8t}b__ z5FSZCRZ}^#W+B4R2zG=`F#nG5XdI2dw@g$}SYDCp7>Y?wuk3d6IN+b-;M`KUmxXdn zM03=+{4#i8wkM;jvqWbRh8^PRF1XVrovG3X1#iLV{kg=_%+SoJWd~&e&I4d?7;Cp z0J$L@r&qX8ZALCbeQ+X43ZtshlK$sgX#lJl_z6>f$el(l9&McQu1K_aJl8qWj{#kE z8_WQeVs(v*jd?%XeE1~kr3|&Yj0l0r-#1`1)>WFPsRYt^sEl1(p0?+IB~Yo6wMIf@ z(3xc0`sDO=@590uM&k~euQhW-vodt1nI(jFR?n+m`VfaRX{yta<7(e?x%beJe5H}Q zRTKJj*B>?Fx+)b_3XYi!dD!@!eIbP{ctqnaqD{qW6evI4z`zOB&gO- zQUU;Np21B9`S>>u*A(PIB_Ig%XGAY!LSGdYOLG|Eow@`Q9;=3L+(OZ2Xfo0zE%0Rs*)W^C6vD)%|(zU zLE=UsMU3s>*$S0!PsyMh*cxa^@~bP!>^mG1qf=&y_@HQi7A1~}Yex_nxGAOnB&u*( z>AXk6v^tp2x4#vFb#{Wz9gMJn#)VDyCQ9Ug-#nVGk;Xfn@QyUUZ7;*x2TOEzK%aCb zhdbNwx!(|7ciC$!>y5&^<6&yN&3%{)Gl@nFZ3?U{zDK}7Yar_v+O5Jm_3(C^OF~8$?s)et<<)M zOp?~it|sse(&4eeRiud7+F>*wz4+xk#*vAPtcP%_DS$OK_BES@B^Vw`{FySQOQ{Kk%!&v{9I6FgViqwNQTH z-~jNedcBAx8A-wtIm(7_>}JVVl0M3KO7g@X$hEc@5e|@QXl9uFEFh&F68`1NqH1P_ z%^|dGBTa_{Fx}HTE$q?22)%YQLpHJD3XkXh2_`4H`x3@4-j3>Q!Ym&ywbP|UB<1Ll zKTcbDx-Fh#7<^}Vd`yTC_T44J^!w(0z$y+AoHKzT=hSd{iE;pfA`IjTIXbuLkHg)A zFf+^aN39Pm@U??`WTlv&T*#-%bX;K&sVtq`uFI-S({^>g8DUL5Ci{8+knFs0v96K& zWDT8rc#$!&N-`#ADL&l_Q#T`|E`AUa1U^`%p$?dmj!7=XBf?2R%2V+TPeaXh+jRV6 zDw;BAkW8&!r;%nayr2e3KzT0Qn9)w0(H#nUF!EKIxbdJsR+jFg-Jo8+JDJqM=`RU= zDX*72p3fDR8ujdfvIQwIv|v`VT{1$_9;RGH2nhwDUht`2kZ!L8C>zEbPQRhlOHrDA z7#cu4g#;uRRYv8H@e~?=N>3rhp$0gTHM2kqeiN-}c2@qV=0G;}Lm#z+O*@!}Zf}h^k&b|+~c7tW) z4qaJ*O*vifkCnM@z6v35xQmd8NCF!-+&S0WPR!8^S4*+Zv+v<|W2h9pmuDL4!WRtS zZg5c@t`&l3Aq?!h`ta^_my_iO8?&f>|& zwX{I6_n;_2O^G4V&7vJ0*A}>f<4*YMWg$1}0gYMwxjW2`r`v?D%ex$+PtVQ+3;ON^ zuSY$J^b^GN6IV)3E0WBcFZ`g$Q&bXu=?CSi`NF_1p9y)3BK^`5XK(>~W1K1)ON$$V z4yQ8jP&y)F6!o^c71Mf>w~ZVB{T=HO;XMXm;^2G(l(h{Rl%%{iDTogANk5<%b!sM! z9g7o+4D_UARhSz?ZGo{HA_XvGn=3kVGaG0KOpWiEKT{BJ@x?BY4~_H0T~BQgo#%g4 z{M53?f=QhOCgCV~c1~0TG|-8&O|bJ*Y2?)(o#P2}-p&#&=Q`;xdIu>E_@ZljFehDD z6Lt(#ydJ)Es5GK~R9iPO2do*(TrMwk%x?wB8^+&=3B6~*AFI@%pLmU_U8+ zRG9827L^J39;C>ajmX;pD8ih)n3VwNe7KEw&DQT>kV8q_MLP+9?Bi*BeMG1W_|UL`;DT zb5TdPDDuwSZ8xu~ZHT1gl7Mi4xMrUA6-wj=x$rPcza7w}LG~Te7+6tcVQLB|5{)n$ zi#xzp%j7vEOHE1DNKbqm1}MBBhEx{6gf9#bagB4e zL50D}PF7iTwhiy!9*MI%h*+D=%OAMOg-73DFEl!{Y|bi$Nl7Di#3kh@QbYywJG@{* zZ|zK8>0{PcsIA)QZM`H`tRC*mMzcurp`pqUs%_I9YS;1SWU7q{B`4(1`XBA;XS;IB zN@al4o*&`_arf!Z5{^;kv|@zQwV_M6j&P?EM=&lq!^uQ2(#yHX_HOa{;1ky=~jUbU48dnt{{S)IRUREa5v&HK8UWx>frellGE{M5{hE`0>!_#Lv z!Ibs&oVBxEKw9Vrhm_AN5Ipg%P)5bp$HG!fPPa%vbwd!Ku^JHHSjRX1r-$iTq%(*@ zRdy+@1a8`j&?@b&l=^l)yKlnmdH$Ilzu%(t;U}C;WRO2D#m7R_1v+O;=f#SsBgJmUgY~9ENe7bbTIwV(8D_mJ|ua>y@ErZR5hES+$?-cIGAC$vHU+w7aR% z;hKzpT+us{6^k!LcShXzvk^NZ*0&t^8+ckn+|A%#K*J`jad5?y|WQs$>g~dHJDpFa9F&a+q zVR{N~h02-x9WfxSK-ebgYEIKoo(>|Kok#Rw!FP25M%!zaXoaOYbg%%bAj5Z!kYI1}bok-DNl`a_knyg5~~ z4`_jLS_$9pl?-UG!7{-zd_@Hu_KU)JbHca~l#rVGgGcRwntLB!RB6VIBJ{^+iz8t5 zC?~&9T-;#N)}o4X-{ge`zo*cJ$~xMLmy1y#vW$HPiFO1TII@90`@2~zOc|>P2Obqb z+Ak4lyc=z+lq~RDG!B{!P#A#jFWg9B!PAFQxrp|K^k6m}3*2aDBn2aL%c51$wHQ=0 zwP_Oz0kq`Tq&P!1(+y}0ZeDbpW7S{@zpN~u^mIeg4NDYf!Pq zRG`WZZmyBvOR1v_JPpU$ZK4pEkR?Wlw){qVYGq*eAnZQcsGMs+Vr$6IMorQaK{FxF z4u0pOwEV#`wb|%BMW`ph*lgr5QM%{f~;BYkw!mxc(!R^K7<&{DB+Fxsw1xi#NB4#$tbtjj7^nUeZwq2L zT4SB7iWIUOQn+Uav}_OfNJA=~5Zeu9QxM^@52t7;9_t_!C$}1HlkFr=YO@-VauUqu zI~yZl9qt?<2q1}%1%GPdjlsS!83G5}K!Xm6usJPYK{|^Pl3G8ZmM2m7*dWi7Lg!&- z$57R&K=O9Vz>@7ZTeW=Y^?Fegmo5~#YHbNGKo)L}a_gmf-%>i#;@CKH)_hLFbLi}-#C8o>Xh4Yd8APbYN*n{sj$_`iZ6dt zz8Z3(9JKDjL3{a<5Ek%3v9~}O{i@4F=pBxU>u=k^wyNAfyh@wGI(36TB#<~+jqt zaZFo`%t4LhY)yX0LuJz|j%eYR35v#Ob1GsK2$s9=WY2#TWY}lj&dz<(jdQB$kAa=g zH3r~O1@sAi*pF#Zw9KUrfE1d#t0eyY%JAh#^NEj4aa*_&|5b>gBap5~KI2ZM?_6g~l&JFukMQ`xoH%8U zAX$%15apF)e1x=*+<2D7#LzIw=bLdK5&=HUs0)_}Edjul{#Eq~6@dv&sMe5sdM^c2 z1tXH2(g;P>masBu7m6-h+4LtWkLW*my{Ov~Ki=zeT~BiDWVuf0{TMrtFz%llioBvfvqA(_=lBr@1JAkDkZ`>t9jCz6_Xg= zxvAT8MoURyqt-+K6Gagh>76o{I6yCaNJ+1D87Y`ZYd(VuUH}Qv&V!WX|mK+>M{)FNT>DP(dsf#lZ7Wrhf5|Kz7)CHX(&@j zW(P8r#D_ACRA`uqAzS1VlpLxOO+=>)gr|h9%U4}>m7Ow_ z^coH_{Q)b?!xoJy21fndi@|h|PYT;bzit8c%27+?ECF6m6+MFV)~-vcvMz{8&<#Up z(v`9|iRdPuMaT^Ac0fqGv!boImvwGBx*_)P;z>_>QUh^#%t2fDIBf0F`kEDu&v)5- zrIxw~2x22`h3FDs$DBoRAJR}t&k5TonK6VW$hZmf4SX;Oq-ZuiVoorWmhdjIhX~j* zpH?LLS{PK*jyQ+EQfbG={HwGO*a8St^t%Nbu}nCLfT-1WHKA!nb7wBoC~w@4KjTY; zL*_IbCttT~>ZH~{HazVekKTVD3*Pk*}E4-T9h|HLs!a9Z4nZv}BJWScD7 z8Yi&fbcu@#iBP4c4kl6(@kSD~avpndv~-{?Dd;6QJc+?2ojj?Vi?K{@n@m#{g6PJZ zN-)a7c1JADnK`$CKADKXB$Ej`yzqrDv{~`(Z+|;auf0s^oFDX{2ieD{9laXeb(TZWD?|tKFyJeF~qD32+>fFMSvbd2|ZR$ie&1v zDPkjTqqcwtEPO@<6yhD6ncb^yf6*=?easjfdbJI;0V3|4H>J=}V1 zk$_c<`jbET69$UJ#3vF4|Ll(ij)YA2J~YHOUXnDPq!te0;|>z%*bf3aU^cLZNH4t_ zn5xTH-Ef4b8&eVKDvE6cRMLq#wOu-SM65YuOL>M+37;HT$;2O91gYYZcjLC~u=R$m z-eSy&4pAn=0MVDYC$&7}75#$ZQIC3*H=kAK&PMeKdembPp{hVY!es^=IoQgiA&Zj0 zfu7w{&;h*zG21-{qz?pW&M28(q8vWjr}%uIA5fHJrq4Bd)aZdDp7aS5RjjdlUDhZH zVOyqskc?ewJBy83`Ki`3$sJzhsFhf&NY z9a-#BOTPDOs3FtJ>hqufe0>IJ^2tpPA@t{e{^xxtOLbB_w7UB0tJCEatvPgzgucnA zhD3B*Lzzob9O*@ACTVeRwmoBI-0+0ma~2!Ofv*-a7D<32{P~OYim@Upvk=%C2vqdD zHQKSPIEjGK-v$52CmF{%9>cmbH^4t}SG+I(q10IKQl=sy6fKcu6LzfG^Gw^72@F!e|X5!!PqE~ImB3d z5mV{G4}P$pQSkw4%(b|_XU*YNuX>evs#zc%9BeVc%De92#KB_ac6R7JDxP@w*&rBY?N(GODo8i7^Ugy zuA77^o#!a4Eszcxeks%|eCFBA2Z;E=4*z`LE{1nBHZZw2r7%MvEn3h|O#Q9<`GFtz zPGk{*+Xj5@P({&stV)2&SSrCx_XT3O{m>8nkjVe>kAM6(e&aVXRvhBi15Us7aD#yH z&8@0#M7YgCEiKXGH`UQNj5e6KRLh-}yVCuYcAL`i88*i%pG@K$*mfE^ zglDwvDj7yRsx!j)x{}9phF>1anYJiHtYSt3npiuvf-Q6u{#|v|RnK_FGiZxUQHDzD zIU+X-jH=P|Rbf6#2IAl$qq9L$D2@+qGz>%4BOO*y1(cj1fCfXKH}g|NVp?R{%aiD@ z|N5`@hA^5iQH5uD;~U@j%x6B+&R&6{PV2rfGUth8M;JY09Yg?7H=vUK_L#>! zMj{xlIYUMc=6KRAr9H8RIF~_BT3rXZZ+g?4_!nxk2pSvnD}33)gnjSx+!06K?scO&Fr= zi8Zxu`%5MSLjaP_Mu14bgue%=QkbOZ)((A1(F)GA3ff>7US`Q%kuLFsCp^KgYdreV zkEWqRgQ85^sh9)2+Nhty(;21B@Sd8h#X2+&^DwizvAr(S3M&CZ2hA%?gFrn5^>tB8XfYB-4N^wPMcd!vkE`ZL?j*wk`#;i9SLCFo} zL5U6g+R*?d-moMkh}F<8M|;HvRzplFVDm9Hm$>OIUU}`$`!epO4{X5doG4-%J#`hm z-TWRgSO6OG3iKCgf?Vj}Rwkfj-INMz z%a%J20-g|yNq8m>()i99@G-VJd;BTKSZ|c~)eRFz?!hI$!snPG*!t!9Xnc~VwpLfT zD~Y((I1;6ZfypDv!H^%t#3eos;*{o$vggZ;EN5wy&SQ@#9Pi7Cnmkd1A@TRlIn?cm zv3zdr%-9Ho=*A)#+WseQx9ngO<`q|50r`0AfOe~@LRK>F!3R4DKpXbqOechCx z^iqOM-2w<@)J=1k581FyP$urAI26gm&?Z1;y43y9XxuDHuQBJeFDcBZOwnLSkftG; z4r6%+Vag{#rC5QOwko>R^vq-|w!L&Lw{tg>`3@-Ab#JUQYOxo9x{y7t-sNK&aJ1US zHl_<coJS8^n(i;o9gOui8pp4tfAz^!}FGQI7Gtj*wZt;a) zyEz1k_JzU$$ogbQ5@NZJax3+Y3OD!|sFmXUz2E!2+C_^9ZZiN9?FhXyJ{m~>cTagF z2YErbOf+5A2B7sDIW!omB~zcL2_59ysB+CU*I3r(2Zh$omUb=gd+CF!y=9{Y!v&dK zBr;3=^KiOkS2metQ^u|m`fa}B%;I678?VUIGhNpp!PJ0 zvS^5g&RAuv1~Q=V0UG9{`*k`ov-{eBMk{-CpZw$}bL#{Jr6h8X+YbF*-}PP3e)hBV ze|)gjj_T6EA~rpliMf1fx~ZE>2|CM64wbw!oF`w|4v9>8vw#<8`>7?gO|MPoO=A zHr(@Po_cbnT&7E2B;J&-@@SmbG!fx|t*iLnNt}TgfIk?;F-}Nb{#C_zP@YCYQJ0%^ z+@5qe;Hk$msEF!-D z`@f&gfo=_kfrX`@BQ7+I^=JdD+*7JN@{x~x%2S>K&(NQSVt4Q=EQNq}%qoqKnFA0S zMm)iWj}++Mx&bQ^OH48+Fv3$%WD_Rn}V~e7^+1X#EL@PX-*}91Km_a93~<$-*mLK^AaM4rEOBhi>1-*W(LLTZ3O(<6 z&$A3};gp`-ja@U?iYh(L^lQ?bX3R-leQb0bP$hsk3$aKsDvAd4rF}5!qQ)HGAoLQi zCzAGH!-skXWY|JLMVHav4xQ&b=Q&0mM!AdK%f6}MH5Q`P%6*;Q8%S9pD8*c-H86_& zX)a#%nFdN)&yT8bP&ZXnCs7XMWjbRn!C!c$nVhr*uknHxyugrXs4+Zxd|>iG5CaX) z6jql+0~ZGo+jS~;5J8HfCmZostK&u`r3P>u9Ew;GY}cIGGMYm-mqB|pe|;KHnLL+A z2{a^;8uZbW{7#xI>d^`%LE1fP@O0}LR!$+X#Sp0IcZ)S>nREP&U*k9jW=wbRtZ{fU zxdAgpzK08OxC9x;z_apRA_TP%a}RM#N+m{NRVxI}7;@fxVcf|mS==Ybk6u1CQmR%w z@hEXmt0jOJ1)-P5)?;j2xvV6haoH)QM=yK+ZIgglJaNPc;e@`nUtV#=6(%|^A>+F# zBDxp?@LN@}K#{gHtI#Oq0zaXzgxM_xGXGGNMj0S%r(8xu4V1P^L(BBBmIDTTUn z*kvs=^ej~rCg~W`DH?l5tV0GplQWAmp4a%&j6}MnQWXi(8!?y0*4ZR1-&%9lhxMZL z5UlV-$vRUsYsh}&$x4V-2ouNq-uJ%9(064_Pj!2=Hq4 zmd@u)5`MMCE{uMQ>Et#Jn5Bt+37Q08^84-MQA;MB)Rw7<^Bh6zGx2$+h%4*zs!x~UtUWA#@t`2&xW>DOd*Sw+j90$7?hB6&n**b$5R5LyqO$x=7^1e~rpnpmVG zd8Xj&$8UY>TPbE#(U(NmmewEfh(~xzUX$p=G)z%Xtf9=MWFugpSVYame3@78F?vf?-bM)UZF z?jg8*ZXkY+H{yE?^v3rL_wso%F3KUzE0`E53Omlp!CRomFm5$WoaLyZh$9LAUQww7 z&$zlMgOy81_>65cAGEZ^O8}f{0o}~+ef+~*1U>#adU`CA+aQzUB!!y~I9#mR2IYlb zrZr{^zB4J)8(DGreFktV=@v-g4CKV6+6De;xZs}_p@4an_Ht;8_4*(NyZmrKW=p&r z6B>O(oMvN~$LC5Fk=xm-l3BGjTFv@e#I+Z zA*}Ul88&Ll$A__suyRhvfV`O^Rq3hQc=m7l=p(4XA?>9uweeSl2-CGlai?e|IXau} zhbJ18P-0^Qbt}#=ftwu`R>@#1EpQe3TCfswjOENyn$R=GI%`Ce={&K5oKnBiO z4qGErTj(Tg(ISjB>9;!PQriK?>yDo@{A!LC%%hVaw%+#i2N>G$lE&2yOzs6Ox}Q>b z_PW4d&*9w4$(^1)EUllXM;RhKcw(uJWaSV6XCYV;4;#w##s@y|frtSj0?c?K&GrsX z(h1M-%*`Sbo*5+ZOn$1kRD>sqhDvNfjc0CWoqS_pV@Rh-MVWeqw|idv;um|v)+G^w zSeZsDyP$=>LP|NIX9x~tA#+Gm6^#{O>&Aj-(JzbW1`3OZ5-Nhiu)vCbC+ayFcA}3< zq9ekX=FJW~{irVvG$9lLko8!Dm{*IfFhc`9=^i&2uf(e!;<$xn3Vp4kTD2B=1|n*p zH;AD}pwGJmqi#$Ax(VPtL^D(pXB@Ll#xk;WAthrYVp8x= zLfQ?YZZ6SarJ{Uk%22?8!0;QpTy{*-aA{xa$#xq*? zA5xI{Xo$|h5dv5zj7* zdM=&W7<*kNVU|OGRU91wX#@c+eIS^aa+#3fBOxDuOC|Q&by*>H=}2En`WDZ4Iw;dC zj#k$P!U2ZTRz>`hBrwgG%~YtNu1kGV=&0ktKMf%tflh8Gn|-Y;GLbcNHOtXncxo(y z#7->IhY{=HGVw&9huakD5+BQe`k^Y>2^oEy@eEv}IXyj^7$pm~DV1+#EO_&9$1Dwo zZ1IYkR}f7xsY-z$gI+_V;i$yY0B*Wzokv02)i6Ch|QCR>`bRBoSz zcfSLmA4io2Y&s9hAkf(PF5OfhnOLcp{3K;?_jb_SG)Xx$Q7$$+( zBLfkat-pTgj@A0S~*t)!af3Mvym<0YEbxkg2w^IA3~!*8p0)W*B%KRO#0K;fC1H?i@t-C8~+GsXt$;2AWcHj{xs1Y(pPg{Wd^ih*pQ02-Xtl8F!8 z9I9(UHxXqNX%uOjQ^F^3sfY|Q<&; z3UdjUcuF|H1v;EbRM%m6)E{`RUQ2=$vZ#tJtysjHbcSJ7=R#mhB2dxqmTcBLuTchc zj!?)dVG!q3q)o4l>@FQ5;WxhVjXVuM!P6jf@C~Bbcrs}Yj0;4-p?X!xoRVAOS$RPW zfvbqqP49oW0OV87g%c#_f+9n5DjhXD#c6!bgS6q?Vt=72> zOe4$tvG%Z~qnb79oTfxJzIwsT)E4G6p{FUAt&gT$7DY`EvwGu`C#IFibVXKBSpU>t z&Ovnm06+jqL_t&`efuOcEeuWB{Yo~gbzlTaANtUT#$-a}5MlDv!A$F?qsnx7DmZjt z!^bpB@Md+w49vt4H}{U0SD*H@rx8j?WSL*)UeF3rIMrYO`q!JAO3QS^GS&`H_aG^1 zor*r;b8ejzBe5}1Xj5UE@r}Kv`dm4l7H|<`|kr zrb`AQritnXZKE-Ia&up(MKNPCdO8r~frah&h!Lc(`VA!4GDg=d`$fjc1r)%)*+jlPvl=6)=Oh!*%T9>`cP zOzmD-smILTM2e5$O}R{tnIAH(;idRaRd@wyr6yZEp9iQMhRGsN6Tg@<50t4D4~;-D z>=KPbQCuVbDk@b)er5^rXJeZI0qF?hqA0_kxg1M;J7%(lbc#We&Ytf3VPx0_PYZEE zkt`T80qimSc2sy%q8SBthgQ=8oYQ0)#lf!38<@5!wx>Z1Iz7`;j2Ssp&5CGZRm6rs zOHyDr!wXyaz&FLQm(@>Jc%M^Uvq>%G@r^|rdk{j&E}MjHuvZri11$i0B*ZkAO3j#U zmc|gh>T-H>C#UC!q4Ojt@Wdyj(xG&TCk*LN6E}_sL@At2k$t*JBGc385GkrigrS6m zV)q>alnzf@`YW&aM^-9sr8CnQ8yuXib(_;$@0D-J(o8-+chyx_8Txb;f#O)(fd+hy zo^9twErXSpU_JFncftn&z_FFy0L-)$0@G`h=rf!~$@S6=suV+yQe9G{^*UX5J+T}F zO3~2S*}Wh^5dMmcDspsksw$PEibN`gp(55Lz5a|_FJ(SfLPxI{c!Y|Ny6gfbwT%gg zbKs-YAtA>(Qk)aKQ4Cvk4gN|g?7Zf1-+4ZWfbH6j^61yl^h;8DY0ONK==4W^PyoD>e!vrjXQn#muvsYYb ztzBrs{lAc~-g*xU(Tp(eItD#kaV< zhJJ39*U3ZcctP&aM7g=F!D*Q=@h(PxhaxVP*HHwTIGc_U2ZIAY$2sv;Tn_&6WSj~w zBGXyEDsgQ1xD0dtFGWhb3ru>{Z*_514}SU}J}-3!tWdB{V|&tezl^?cRw*E;Gl zJO=9C$Pi?bgho*e30czpoYD=Th+>_d0uKD;2;+quKG2#=caSGj9ya;eCpc3KB{HQB zin*KQm}0tT)r!LS3tB|%xPb^g&Cz#>TPeXbR%`Ke)c>Eod;8HXtI7r*5K*HS{t3H- zdiz~)Lt;qUt@3$6;wc(W!EQiQ6g&XZU_4=i*y1558iI0&t%-(2O-%Fxx=An^?mnZJ z+H`^;-l{R;k?&Xa(8I-g54);%)!w^SJ@aDKTyu>%4$l~C&N=2<>n)ygquF)hZg5G` z4CiGbxyf* zFFjT+UYc9Z3cQ6_jWT&1k8g(Hf01$ZA>gVo!RqHe_jQEMEsK>y) z%r(X$Kc6nYB}k@GywuwG&+){ z1vggkHMMEtlJ*Q@6gIHO1V;p}qIX1~lKM2Nq!9O9^>YyzV!hbwxw-*AQwFwFf^+_t z9yN5-xa2UbS}6iI2Z3h4H^=1X&-dS+jXJTrmi2nehX5SX`2Jg#anA(^)n z^MNjzflL@?JH#4od!sbFph2L~saeGAm|4UnJVRh+MWE=2USIWLtCIFZ&d?o=LJb{p zpoTenX4yj>?^ioLX$UT>a0_T4K*p)15|vzXWb*w|@4Ry9uM4-PBNtefB@*~T;d+sy zT_sDI^NKhxz2&~>62MkAMBBfY-^y~shiME`(S@12Hr|y!D2>=~64xikHzQ^-82*_b> zmFZ`4yn>q@AH+E5?v9wgPA%``ds=dvK*K|CQ5RzETP;G4$Kxr+G82FvA$t0x5onGc zUL`YeHI%8KkAEHIoZbk??djALcwM9rq35Z1EwVHo!E-PdVEn?HNA7nTVlIr68{n*$ z-)Fp)Z+p!_%;19t=fnV30vnAY=7q=$K13(qJZ0sfBf z_zn!E@t$y|BOE4e2pN3frjAl^6NzCrTHGNlS&G1=(1$ojEv#_Kuas$FpksUy;v+}U zxhnX%&ogRDsx#cC7@n@W@jc~}!ueM{j((lq*Y)`xp2SgLCBh{*;Q9XdzaMDmJ)Cq; z`Hx~UO#Hc(jYbJanU(a_k6L*{1Rqy}h6?_3s;-Ps?IwsfzxmB_UH^`V!moVgEA5a{ z+jkZuQh?WB2nwT?sMZ%aTmwvkSci$o0dIyf$4Eh@Bk!YhLu+O{M zrYNGl(q`V(@lcoh(oP>nv6uPL&BDd*YOPX$l6xW;shYY^f)7!$0oGsWp%U@}Y+w5;<-JuJLxU zRDmyqH2SEmKK5c24uTwW1X^50qeMj7SFSvcz={}8Of=%0OMVsHQzjqOxs}{I1bysd zA6q8si656h!4Mo251^b9O=qmW;w!!aCQg+QMB(inq2aa)bPN}kwwb&@?&odo@fnxo z174)`wyL}$P#_BqD)33&XWqd)34`qt7;>~O2v7nxbw@%EA&HDW*y?KlrqdSht<`-#(&oxazbkHkTPYF*BIHXP`X^+nhqQ8vFL5-XkJQ(X zU70Cx_S!S=$gcf!_=E)iLc=9-V&GCiKq8bN4ArQO70y9Ja3Wipg9GrHOpqJF-ir5< z{A8DWJ7fsbciAg-!NXWbJgFpBRoDfM?5FF<3YepZfcc%yT#3MBO$0QQaATJp1js~T z9VzDc!4H1W4Tv*7ZaupItiUIffdaOMeYUSuoj6fReM!JHuLBqjV1*%8;I)b4Ee41U zR`HAOO9$eHj;1S91&GWC<`3KdDxvYp5N;YV)M)C$7Ns}D?A+_CptR=j;hf3FB{Ujw zMsN1g;?z+I!SUJGvc}77=6>{+Un*gB%HTmcWYz=)i)Qw9I%{8!fJ02ac2jID;x#5% zMzfJCSrNx28DL-{;YwX1emo+u7L?)GVX+QVMDhqtWKuQ37FFXE3m{6;-wKkd0tt-jK zBqb-ZUsx#`;j9D`E_0K5NE;VJRh;;ESVDr~_{de*>aIomApl*CU+|}udUP8;tB%qM z+~fd!ghnrcr*g<&D8rDFVk5-#*Ck7Nv7NrTTQxeRwhHFRw6Anno>C_K(W6Y>La2@a z0fM`RTy~{VUu+$hz-r0O<>+->CHE92OIgdhyO`$8(U-#b>!KdEOUh+ksxJ@p%W!_f zNq4h@n}z-a$+9GPzHX3-nZLB>j_^&m*l>g7=dYIBmUW!HknFaSWCFoEBc?O&n4rO5 z8s z(32}IG(I3R!y=F|ddZ*ng*}kH;0-$MBzn?o()5uC9dwstd)+;cmyMr3ftJ$}q@X_n zvh4#8&-sFS9-?{@ji_i0w5HArhElWYgDrBaT+58~VTb zo4|Fc+dpuOts>W@^dYtR>uZGz00F)4gsmOUc-+GY})NqLsHZ$Bes| z{ysJ?tuSFTR)MW{Q@D9oqf;hljeZ|hJ(p;jk7=CFUshBs(wOB-mNGH(T%0dJOp3@9 z6$(7Fa1q;pr4oi(RII3mxn(h8D`QUu6Gbf2h1w|j$SLBuOkC-!9i?DqATn3Sk5NFE zNJ%dQ7`n3rYAzug1(p?8adtH-IjXxOqOdi$*swJBWXZyjEr=CcOIF@aAs@=fb%`3b zha9y{k^5Bl)LS&1XYv&X?iAGVJCWo#raa(9n7%gPW4JnA{6j}coq z?GkH=v1m~0l3veBe5|5X-Yid^~WL z`pUZii>Hr%y+HBNkA4&lkfnv{o-xunx((Z@NRn+Ti6sEifh~f{vsrQ}rbwP8lRTAw zK97}r+k;oadw3!#_3^rhar#!Wk7WL^UI#xG}?=)dc`zRN4?bn^-f%}GIqEz?k{A1dx6=_nfJ z_&UJ8xgm}e_Xv<30pcn?gpPHw#i<9*X#_iQ0xzfN$9Brtp@1{PH9CqJJkcvyfdc+I zav!Htm5Y_=Jo=|Hrwqs0V}IR=!JUG|08Xn5lI=hhZ!1XL>8!JCx?L$s@V z;t_j~`09Z zd22zsd1htr2?1GQq8)*}%R1Zo+q#xq1b+C+d%2!os%t@`Slh80AFqAlsfF!0cy8%? zLUp!`o>bD;F(q}!>t3M2j{*Sfic^V!qvgIB@Wq6sKZnW$JOQ>Z5bvW`$%ushi6@@m zFBpOgq;^R>C<8AVaDfPCM{dd*9seV^afGYDSG;wpIMwk`7a4_*Ke-rNsfD*g{mQpT zd57rPOGyDEAekoWq;vG7q_NSZ6Cc`9tk2x9>x^s0Z;X@PM6}J5D@BXfwLBb#r zQcy-uaFIwEG=Nf13_}tVL%rM3HYJOOc`l*H89hGg?lR_}5wQ@qxk;R>Dv93dzN{e7 z5j|PbFn9O6-~Dd%@-eQr(}79oiBMks>Q_TXC8FIjZWCUHAvI{OZ}iv8h~ZNL*{{(f zP8SmX@t8~GvLcGvqI&IXUn>C?KVFIQpMmWrk2#?sajEvysX??;mE5*4%;=<}#!!+) zD?D8y$dQZ1l|}{lS_CR$0TzksgktX1blq9U5k?a9u!=r7ARGF~+IwU#M+3pk?u8V5 zL>JdcYmzoDH(ZiOIogSqU-4}s~)R!bV?=O*Ii~#Pwh?<7^jR+(4#|`#N9y{ zLyZbDaD^B*Eg2ZGvRNSf;2;=aB_=&p_Octpbw>~;Q9F$|gh{|5k#sFl33JxzIK(F! zv=o6&Akgf06K<>#&mmycRyDpG!5vNNW*Wt?ZZ^%vyi;pvvyMqI6Nwh|rWsSRX~Qv7 z;LJlaLu8U5f*_69W=7E^a-UP=6JbVKUCk+#sI88(Vve`IH?5LI5TO*n zhr;TaIDI4;E4=me&Ue1k%VQ|@)g`Plb!m(hd*!w<2n40z#zzqIyv5N=s5K zRZ+&LX9SSr^i6kCu$qS_VU>ko>WG!oKsb35PetV9hpcSt$Z&l4_slcTSP@d2GKi@` zRi{KaO_*5{36(Qyc$_S3@u#6>7OgCgD3GP$ZeR3tMobxM1VfDpMaDA9jp_M}|21Nj z8P06+^0U_$L8NFzpSq`-_#n{2DC<1WpaK!;C|yzs|6P5y+VVAD^EDp(7(JjZVX#5t zbH_Kl;SB&`a>7YqSSTA(JUmARB1M895;+0_m>8?$6qImSGTg1lOyb?xpk1A(CHSL}QrZQx4Q(Tgl2GM( z;NXuK8JYy=^u5Ownz&6q{pnBp_?c69ir4b_E)Tlul;M;F_<)I$AzEzE*nwivO6?#J z7~@Jujrf3fMZ7%aai(mgR zIAapK;%>Upc9t=&>Wg@nLtBc#bwpr&i(W??s&f$-1Pt7!&kWq2g#6i`{h8t0bfv{| zFuyKq{xHfT%hJcBmBmUXdo!Pn|98CO9bTHs%#j(32r~r!N@ff6*lNekf*Z>oOVo1NKS3Z`;N#j<@&x$ zR-dI+Y-Bh}y`tw4xx1R&%$9LeDvn|RZZH=ZGGQ@eZJ}Kwx1JXzU9i$!(Kl=8@WD`{ zdc39{xml;GTb}wT&M*3Utx;5q5Cs|ftcOy^R#do@s9vs?&a3TS`h;Jk#2m_yxR;np z^ruXMnKDZZx56}%^g1dwStEaJcS3iVC#}R=WqY^Z{ed3D5HmWMB|7{1T8Mxg-l$8NC50e{c#kiJ71x7j# zEZ!U@cT=G%&)Qh@1`WkzypFGW)vLU)Faf}8$%JVY4{|ZhVPbc{;$a0&ly`(-puj2~ zH<}&Mb|j)fkAnh}jCkhyxQ42nXk%0i*wG-xwDki&@B@AeP)BZE4ZI@}e_eP}C?cbq zzD}`KVNvi$v-Ht)Z>0!aUj&-{USHo;^Zf`Iud~{=+ifN=IT_PcGIJQ&%>+jC*~eAS ztZy#C(I{`?M`LB+)a0bQiANo!DM3XGJ{q$SC($MCW_LwJK||uQMc+(j;TE#pBhi_1 zvhuLhh^o}n(X@vr&LkqBqz4G^eeZjjEX^uc&|Q7lDCM2> zlbIvdS9DQW7+{XGXc0ET05K5>fHCvJ6$}%-RFyb63M!%#|GR9V=hmZ=B|A;1p`N9i zdU{b$mg`^Cdf37iYkEt3PvN{)Z6S_7vY;dbv7je~4m8m&^>u?Jf6y24t`0Y$1sb07 z*xiGPK}pQW;@~O!1r9Ip6y$}UU+#q-@YMq#Kk6c8;@L}qE92x|R#IOh6;CA?sG*|l zC8l!%oBq^u6g-ZIwWshqvZeQ(uDI>Cga9ibh)kz|^1i(z^hhy^%4Ty9QqMg|RRw_w zA_`-xU;<%XY&|f|jj4`|jc!9%Rj5Z>z0_UB7BN8ro9|rG7&TU!=adv2)ZIrIg5sfv z9&*$xoYm1FAObfu5@SCXZ)y+;)MloagS@Nk#c96Qg3?BwA4&S?M?b2s{KsoVl3OpL zGXYx!dQlkLm>DS^3u__5F;|xr^}R)mvD#bftX6>!J+Tc3ved&rbxgDjt7Ux}FTE0w zg{WJJ8sG&KOMQfF98-skaTOVT)Qwc}W&%J*7-|$ey2k{RN>pW_3RL->;fnfP-8Cm6 zk+8E}z8?XDwb|W+GTE<{aVF`k|BUPgddm)Dx?+e= zhMZRCA9%ymTOh~Dguo-*ndOmsk? z9GUI#G{@kn7ol`e*O5fr^kssIj5li$W>g@-DxD^Lk&WY$lVX+jci>x z2!@ss))p3Rae)skJ#3f#UJM>gL>fQxniYlxs1yY&-r>9ORCh^TDoBBMz3W|8W0p1E zXA^h{LVAiTtSBZ19Mhfao#JolL+ymqWvyD03{U{WH@)diRH8-#elHpVIAKW*D@b<= z#ptOQCX^T60MXaGyMEHf&H~$AVwy+Xe$L7r5!lgq-`l=}3IdJH^kJLRST`#?m!|IO zO(^{0c_4)0Y)msOc<^hdz%RZgkbnSf+Xr497b?iOz@YTCU;DKP0@o#u4~Ot4*AsqW zKo-|i+jG}^L62LS^Fm1m!QHe)G0UCIbAyy5An z#gW=>Sh7Sio-si}sJmIfj1$0}QcOqBcy(m5b}5A0*{R}Qb_|(b0v@red-c?aGMDh7 zFnV{vj2P1^wnPVahup?%L52~TyrIT>-4ZiS&mH7}p?|QauR910HFOa(>ZZ97%2@T4 zn^>o4$0`*zu4r7H0!)$V?G>8K63Oyu13&^}n2bZG*Rm2dmMxcF=U+VCIEEUS@Cdav zGL7FUKzXsD5=tnsMHg1U9I-)3M>#M@Fv_Jpi`&GK0hta~3rz45KAjO@ZXDKYx9Ei-dWHxL4&wUvyyrbW?RF_v3)Dv+ zebkfWgv2hBwiJQuh(NR7>*%%Wd=P<5qAvf%U;Kp>(Ufo8*JvnLF=u$e(!7;P$P8qz zGC5@y@fnpClqLhKaF3o8Q=P@4StL`St7af|i%V05x)u|LFRU!%5SwXS&5TYB9T6y& z6uuooPgbPtwZT-Hd6drjn#uGf$PrtOu2PS6-iS6iSuL5p%(X1?p_)9+{O4-awz`cJ zLXkP{723u@07_uzrmW|(xb(A`{(mOeUE(`Fl303HnWSs_p0UC;Spkbi-|qS!o3`pX z&O*>5P>v=-6vg9W&+^pg>s(AXTG7EW#q&C`{h}AW$p2;0-SUX9^|DeyBS-}QumAe5 z3t&zcMMCt9>K!K&gfJ0!Z1=OP7G!})``(VE)5;|=OpT#o%ha=y^HfUyW*LWGFPuOA z@sEG^cYn7p44JG^7U5r>KuScfb-~qg%}PFL}R1DEgo`_dXgoX11%6jo5Mcrf>QtU$NRt)slOS zX_z>AjONLzc0umdbf2^&g%HM|7ShOIA{F%c?s4$R#V5+6bUKE>Bv!Z!Fig8Fi=Ofp zVPGRkJh}^c9M9&fwsmpKaxbsksHj*EBn@V%Vk1vSg*miJ$5Rte@>MS?yQI-EKXV zD}V&9#EdObQDlYVQVMvX&5ycJbr~&7DY&te2xMfrWDc^41AmQ}tFTx135jL*){DUP zMWEU5_4Qpf--CcD!GLbjZ)!IIm|YCtS=(7*+`fI=`okn)zDKFunm{UIe779=gQW2K2)!goS_su>5(+V%CX!pXZpnUW z91ZuVe%fUAjSe(4_h1q!nQN(`P=FvUC{10=^pcy058wRF-z*=HhkH(_FB2>-=@6k< z>3B$P+XsQ51TIi#(O_xIeHOC76dsUQ_d%WQhUsm|!cabrC0LSMN9_wvb7_ zia*XoLn%TJ^+mQMu*{-%u*EIC@DafF0^E`T$0HHD46@&t<~g>NB6Cpc;+~RqsAy4i zN!?XS<}{pd!V#qK^x)32+b*62LnPP%u}9|HzX0gevK~01Ld+NsMUn9;E`V5oy;8Wu zj2)lxj+ht2cY88=QCa;t1d3;)#YtH1ZPzi&GNa z5r>o{ckb)x5WBEd(J_tW;*d+2t81~TWDCg*Ke7%LNY$-p8FrU1>eujeZRKOgOF z3$mh~c7cRA#s;Nu?7d@jB+=XU8yg+ln3x?*Y}+;`m~i5m*mfqiZQFJ-$;7se)AL{F z#d)9i^ZE3w)nB@+YS*rNS5WjLovEGt7x+8g8s(*5gqnS!OF z7eil3VAc)$xpCoZM%CVf3t+2mq1PS{+@r@6u-EJl6%UP!BlPj_oykFsn$hJ}VMKgs zY7b7QWLItU3gRG!l01)uD|MjdqEXRq)HM}dw>9Q-;p~lmNF(it=`>9C!71uDI5bGr znDqt@r{(NNw`B(T2HWG2D#dfN7A?oAqv3n1FcV>{HT+g8UBZ+HMI>U+$8d*}&QQuV>p3BW< z2pNesGGJ}BlX>+)r1TMoEmBE{sk7WHkf$wwRAr%-jTDJBo9VBl%LU^3gIJ&+-c0)R z(J^&9`)}!lRExSEC4o!H1KaXi9DD|Gi@Z^vaA(_PeibJqZ5OR2$#V0cn+KCCsJh6h z_w=Zk89)f6k5$*QjTb@68*`yO@S$Qu0o>l!4GKkQfgktK(DU2EPD$3aJj<((Jj}Od zHKWjEu)LrH1+t<%Gc6R7?OtI*GD?S90In9xDV7<60d^0|gPB=c>g1ax5_~Kq`Dq%gc`#~={Rv~5(LoYF ziNDU@vV^RSEmwuBqE~+Ztivg1z6Y3I#3T5yc^T&gi5&xoW4pp(vv&MQT*-9Gwz$L@;|5iH zqbe5i;W^xU@yFx2=&U8PwfZm||LnHV5k6Pt93%7TPW`QB0`P)1v*`{9mh+{=>j4ztd&<~N$f-eNNr)|P(1(Pr`4a0 zk0{_27Ygwt$c);CpHp@LPrWSP1tb@(2-F|PI@2YT=2a<6qs6L5Nb;lAOi_8+t!l;Yn2Tvm_(E;6>>&9!|U_s1R z38YJVTx`Ay=TRb9gW&ccR>za3czEZ)+5Tb77a3@uTq7WM+W3Oojx$|i8n6BGTq|_! zu?PkNMx{u&qR=Az+^l+A80%z`W7_r5Q2 z9tex0pFbW}@GuYXft-dN)Ba$(KI}&xt&mDC^HnA=s0lSxn!M&}m?FgkG(TMyS$m7J zqK)$qiMB4)Br6(nW;GF$+z&x$_L0l?jmzVUGSyyRp&TCQh=&nZ{Mpi2TqzC+Jql-~ zI$3@dk!&?i&v0a?dUo`tTy0`$*c1wxbqVK-B7+>5lxEZ|={05B(LE_5O9IxPsCv-c zwgdB*dD_ZHLAf-2t6ty*slm}TCxzq=Ad8z@h)2&=RlDyf5JhRYHXTz?rYhGt{c8o$ zob^$E|3v*=y%oS=TUyP^OFf@Ju%ML;%1K+~L#_$PqBY#GF-;}CrJ?nd$C(df(qRDV zk5Pssm5$35@!u)j4@y$B17=p)X3dVZt1{MA!2Ish3|U#p-Z$sT4R5}KASnYbQkEVJ zY;zoK-@a0byDgIY;ulFkbKyDtM3rVnB)z5o4eHdCnqH(fVpitmmY4fSW{nu31f~Vc z$Lf^G10f#f4gm>R)3DDA1_6>Z9)*Mq*a&AgCpo_&tw&Rp7Ix-xdI>k6xWD%ahdDEH zxzn;!qNXgmoQ}(ISzC1$2cY{KjAwXnvD{LO-Wtwu*VYok7*ja^gVw zBhIZ4*s6$z&s9j+%qvEcuMv+eSuuACV-(=^1}z*Q8=%qKRs{3EJV=do9uze}j9R?m^TR3tNFRa1OY>6?a`qdUeYT)j2ZnI zic#Xk2cW!`3^{dAS-C7u9`_oE{YMV8NXQx^$WCJN+jgs z>k1}&v+U7OMMfuhCYF0jg>}_4*D)uP5EpCHhDoZcHz5^jsw zWy@%y9NaOY-5{w@?HyZ(yXD?#*Tw0&z(der$B!z@#IJo*BOuYnIS*@^7?Eo+C!0cs zqtNb=(9(^|A9a@r>z%Lns_9wZVZp5&8(wi_cEGhesxw{aa%cT})MZX0sBK}>3I}%Y zk4p1L7l0v*oeAT^R|s5aw#QW_Sr z=c;??@g?0DF1cYNP4CjTemQIS`+vqD0vqE4rlb4lCMqTjo+%kRGf$V8sahJr=@(tv zd5pOQ3EV2N$Kv2Xxde8vk62`gHa|Gs>DbPlbfzcCTbfocC=QK+E$LgSiH#a;r1_vJ z&^&XpNqa?=ctkeP3$!A2x(c1}j2MiK9GfDdh`ap0eXe1CC9v`z5P? zC7?%=t`9o-&kpj(yzIaUQ#@pjstO@lHEWMQ&(U@TdJKTpKha7YN3-9%`E|d^J)5M!Ud#4aj;j=+Yq|oOwD*sE$BS)EOoyuijG3%F|KI|J9u3dJZf23=-7?D`s zX`lZ^o#~sb=trY`4g|uC2Mw~cYn$}i+ll&x=MGX`1od*bC`B2H+P-=!~Ti z(BFZhO&|hv%CXLJ5`o7@BdJLZurs>!v$UOmV`B`p4pLEGuG+D3Ci7A(f^03EG@t8V zRd@Y;u>2s86TDyniIMm* zyg|jLMw{b@N)$e}zl3~k0>OP6+#7qJ5@>#}spgRLUbeV^26_=aK9Zo0y^otPRtp@2 zwtgCpJi*W@gUMcJ)3!Y09C|_q_x8?-Kacnw6RP{F_mZS@$M_AQ$HLA~{Sg2XeJx6ia=Jd$6iaN41jB+}oZRpTs?GE@ z^b{lF-y1CZJU)T7aG(}pF_W(+pQ24-kt|})cOe^z>i1E=J4!5+Q{7qb^kyoLgEJa~ ztvFnVHge;2PtruOHJ<=NnRz74$4wn3A*y5725t8OOz8fg*5*OJY*jNA9J(|kq8^~+ zS<+d}DEc%bPN3~Z28)kDU*M3_4B4Eu`MfqwkL-7D(;9ile=~~gG>bce?JuFrIMa$N z7av9(0ZsTaRwy}AGD*U9f4*NnD*eWcO1u{P69PFx^DA_k-ztOv3M%!H0ScGKXa}1@ zN56Lbc+|Lqg4&L$g^K+IV#YEL&vkpW>R-RYXH?bzJ%%FsIkXY?5X05!XZ9=p?cV}> z5VVM7IG&Rr5PoqGP49Z#(ZNVolzO1mleP_`AU`j&$)MrEq-g;~@wf`7a+2eR=+5td zbZ#8V&8Mi%qa@8oB2U7jjIUHkjWjDK{?ZC=F(T0dR-r$V7{78rd5p;`rVyUL?Elp< zYco3i@UNHG!qOLd+`;MQexp#>t!m+IQu6)a##K22v(p`tuw~Eba*rn_!81FvB0w%V z90{Wq>o$i@bKi@uONN5_I>D=go-s+$)RyJA~L%%4gN*WQLSKdIp!84j&Tn zlw6#m4oP+XJsQiq2Le_!WNVm=`%4Clvr}Wibu^!>5>yp51gvDOO|DN;M=c!g3-Tm2 zYW@!)7RQ&TF`a>lZ<=ezr-q-N)sU-B3Lg$k6*ZGu-`mDSB;8!uk#m{M&GVK2?nMdI zD;G7Voy^S7v)Y8V(p9rur^Z*7vxt4)sZFkXQW{1s#C|Ra$EN>+Rk31{W`3BZF-w`D zA20RFh`~-@Sy5O~%K)6^F7&;iJPCpfg<~7`@;#^xD9p--LW*hygvjN&2Y*E=_b00) zG;J?0kZp@!@K9ewb$o7B-ZB!g@6lszl_=KMVi{BAi1FMD$eiM$p zaopB5evl=FhUt}GICQzVf|=ni<6yV5HLqC2X-OH|o4pAR@JY8jp|vI$bc z70U0XB(J#0j#7^4#XBd?c3<*$T4iF%`kcS})~FF3q^9&0Z*1;tes&Hb;4P9+%MufGVMr%6Di_!Yc>$r$W<)Z|X_oP;x0fS9AcN)E+&FU83 znqP8YNXm_MSYSnoe~2q08dXKtF84HrVMbb4w?G89hQW!4pnZlRa$aPto;u=5|fgMpMw=pttN?f_}xB0=O+Ok{Q z!0Z@|oR76xtbZwiXv=x+K z8b}k(zBO5@y&gOU^DTM)W>|>qr=E} z4?Ua58>J6xA}bbg!}1~`C8R654j86{l6aDdS#wG%lErFqRf3ip9! zCO`x`IY!|jvTjcy2A>KgDlOP)6rJ9GKqs+AViSojo4BqA%Q81sXy_Q9!lXRzH%ojT znWq|XR>am>E1Dr z12NdmL8h`FBQU=XSjzFTIkoq#zNi_N@zAM>=jVjE9`FIFKT>^t1&)#mpBJ>SQ$+d+ zBSdT@Pt)bnbQoOtLLRdPZ{UJL8Q78ut2BvRjJ7Fk>astMvx_?15E`U>%?qxeLhnyO z*gpwUhw#Kjd+1LYy~-?MD6@0!o?#ITQ81lj9mP+Or=;4|b~X3H;+ARd!J*bU!xcAL zn^)o?zjRlV=fi!D1Uln@Ikz4uL31iL(lVaiJD@9M7_Ae!rq1bKwXiH0GU%K`IZBUBSKWbV$!g;69Hwq4-cY zOj086;e%}HUHXgQz#O&$y=(tRWoXu6&uVwV)Dkk0_XUKAZO;|8#RU&!vWp_N2YmW1 z@gEUkWT8JqhI6S0BGie#M0$o<7WWZ|DGI*3=%kX&rncY<1WVoqL2hLIIo!be2! zKc$pF!-#<3o<$~eVxUd=v;|;<#+G|(V2G){D{RcxPH@A`zDNrrH7EsMvCh%z$5)6@ zv`x93*g`XdSu5j{1#eMc(}6-%5i(UhW>(Q#1+GH-&y8^rcGuY@7bQ zoW^vwv}D@Je%SR139W5Hv^e7ZRV(;?Z9FIj6T;4iGe*5c4Xh=mk-NdzV3#JCIT0p# zLOiKoQLMvMiHTp-B4i3GqWuPydEMc~Z9x&loeOu_0`%soKiaqBRfZwSv!2MtI2q{6T3HMl`mA zpKG9(NfwH7`jnRRL48W?$e-X%ALr^wmXHwk zSC$su2;ldfcLgnDXj{_FC7IGx!j>O*;->j4sa zZzxn`WFIF2yiBFUMIkvOrARn}ZymI;m?z3AZX!to+cNzTnK0md9+UNwybuG{m~RWW z!O_)n>Uxv7POP*kbF%ik&=y-gGP${woPYX;3Hma2fL!?l+-`*V-Zx;n(s!*|^)^iCC3`4-LYG30x(O zfce4c5OB9v{Evvaob@y?tgkpfrqA(OB}*jj{-|6+Jet}UN*P+?W<}$jA_<2|Ysc+l z!p40yCN*l9=VYYonZ=2WOWifl2*I}|r>@OIl^e)S?&S*)Iwzw-b+q?FbB{$^>2C0e z<*2e3Hv+J5nF-9qn`ii~K`h(?8HZeK^pYYkLo1yG-nLQmflH@W!Oqg{s`K5Bn`#AB z`P&AL-v+qm%|`R0xt)cUCPU@{5%rZfCVG17+soUoRRr(lT@bf;LZEfe4^AdZ|SK>mgBbcngC?`J|*| zOSGhTGXDIZJf|yS9*sXGj@*|uhTO#Y<_=E z{$BfMT(Kk6(k$#-n{u$yki>vG$*~+ch8e_(OA||@!(%x0FOTkiAG)l=q0Mqsl<^W% zYBo49NobM-`ZJTRWU#IjEC@476^a3}dH$<5)xbaq;GecKVrW8kF$P|yLQ0)83~m4U za5z4W6DgWzC&J>>D%J}0FADi!VVDkOysbvFfvU2v%|v2jeO|f z;bb$i5%V->D&b7{`uGBWQoJ>aox=nq6tc$<=_WGgv~R+M{t{Wb6mrV3JkI&at6zE6 zzyH%iW{iQ+m#hnHcyR7Hk#M-U3?M(XM{1rP3T^rT{{RzbTaQM7c&o z2=R@4-@MIwOeO3r%)@~qS9tR1hvvQzI4-dcYT6Rnl0rb9tp`e!nr7v=st0+#T_ha9 znN#>6ao`!xjW1O;?@;yZ`03F-n}zkYp5XF1f>8tE7NuwoTFaFAhi5|RgCG>JxSheL zob-#IZNa(Q4*5}{T5wdEar9%>0bE?`l zEXa;h)dv|+!vm>V&e>7T(O8NP>t>M5z5&UOgSrvymlWylC_{sp#6f9h2k$q`pHA`E zDHBH<9uHz@7(fIlE0a|ySgdc;-ob(zI?1=EtBuGc!3q0e85L_hAn);H2(3*l*((No zmJZy{dK;&}IUOh&q+%7myr4d8q4mjE(!8uMiyE>469&OG$golnljIxZ_Ydc5#1lGh zp2X$6w0Ywuh>97?(R!Le;nCPR$?;h}+Gft+zbpD55C#L4F5!gc)*^p)jmT)~GY88B z0z$rGbRZ&oeEv?wyxN%QOBJ0D>VHCvP4}QPrqmC(g?Q2&lIdy`{_*QM5&g#EcW%X^W*nVVW#Dy;c4ikt z!kF(&!YYV2GP6592r|z=+cyvh!M{hI!Yd8kC5?L8DgxU~%|SQR^~fJUmiNwmWm!&@ zrMrP;_I(458osg-igBxfAD!O0WEiw#kvsUNBbZ=`zc~cT9F3~88O%l<&|uKLwC%N) z=$-|TyyeV}))*y(N?P5E3f_rI1e;O!f71+UILf=ipGpMpLO7r#S6W6V%3FEb2lT7+ zs{d6YH+1vCY8${D zRsJR&yN0_qJn4qNF~t2&4(=cV&uP_8e8^tB=)wH@0KN~@SF}5*{A|m5#V-kW!4MPI zZ&{bm#5uYDRveRu-Ej0QIr7I6 zhLb{bJvv8x;t%>aA?NGOA{S&2wr7>dEx9O~F;8qH-7Cp|?Vwkfe`EVdezis6&?uV*yZQ00g%xT-It`NudA|sIzCOQ{C*OOCq=K^^e}7VIq+(*v zw#dvmzA^eCZHhe~RYrlGnS6@j$m$PnD@|kSMBb1KpJ@1!?~d`%-XJKoVd5S5A$Z)braYhvt9!j3Un2`PrHikvs!q}?p7#<6F(?B$<;+Z;zj z;ZG+hzeYU*^}fRBNBd@avOY5$7&E%ut`J`M%gU$n`gEq^*Fg-K7~;J4zA1< zSFMTjVze4diB9flR8ytM4hZg2N$t#IlAoeZ{RkuxuDEojANDv%ZUl|7yJ7v_!4XwF-=`S}-wk}`9ZGpapQf`GwY9k89Ix+NSwJvqE#ER?TU zd;WgTX!TP7aDSm56m&0RUCyI&FWr(idyta<9@+4uZJCcZiHBS%b*Px-b3{_W zA4a#yCX$VXl1r&x%23f9nY8V;u3@QKZrt|o6pbq?)qP8U7Q~cUZk#&4gVFeSP7Z35 z)!d+#u<>*`@|0c-*74V(zOb*Z0rW7ecW?W!Z#1Tq#m=@+5HiGOzs>F;NuGScd#TcD zCNFie+O3ucNeezS*3`s+^oiVx$1=^dDMWEPn_3Tg} zT^&ZJ!UOX$9PU(>BPd+Lx*u*7Dg0+*^F&N$c<1)Rnc$T(#6w9nb0dYq1FEDzs^jotm^JDQP>4#a9CBnG^lY_B;8-Cs!UXBJAw#G*H zZOA5;yRW4GN+GwMTh4pi4k0;kl;iu|gpa-Q>7}DBm-2_G3X+@frEct1za8fB6XC_MYmu6+0x$x`thB-QZ8Pj%|JL(ebQ;3V=#^z`i!fJzaaTIN z$d4QjU}G}exG5FGfeq5dE6J;F7IznB));0=YSuB($?am+xFX0nUqC8o_M>v2`D^z>SQiF>T1zmRi3{_m%yXb`;VP+L<%EtzfI&bjhDFgq?k~R_Q zgkL?cFTCenXlo_PQ}C{Ig504QE;WsfqbyZ(VxL+K!`XtUKG7sxBtO0oo^3u}d-crR zB7R(2jqS^^c<#iHlP}Jf$9F~T;j4Xe*~wSSRYyB&(!=%XV&DlTVwR=CqGRRrv7|Dd zDHA_ii991@=NjEfjfQR%89EeYsTlLwomIF>C0UUgjgC|-Uh2!63nE?BCQSwrB%(2C zv{=sv_FYyViI;~FXIJq=~iC+07X zEu;}5^F;asUK?0!N7w3!CHV%);3J0O3QY3v&=AORAU@F%nRvK!xe*>PyN6XLX9%L^fGJ&df3Qj6?FL4af|U zM}W1%&I&y4^O`B#$2)FsB|bwQuw)~IDQiJ%=aJhsQbZ_6CmNP<5k{+7oz83Y7Ud%i zTliJ79<1mlV$Eker+r7iPKKRz8qa$I5z$ zSj*l%L1c@iAmH+E-KQHY^lQgo(z3E8dj>Lj3?*1in4cLuJg5u)o%2bI6#G`6vI##I zFNaHie&+~;tkbC@nn0sqcP$h0WqTW}dgI!qAQQG~z?W^+v@{6Q$a6wz}Bq!P33(gTPXTy1)Ug!Cd0pBzaN6 znVn@F4-g(Fd||f2+h+$)Z-e7JEg-Eaa{#V`a@7)#8Rgh;B$n^ghbM1gk?Ko9t6!Sz zh*NiI28TQ+5vb;&IleoDbk{^CGs-{z z2xD%Z--KBo@01`)p;>Z(=rayY`Mr!zq9n5{iZ}){`gMT^p?W|be|!ucZt+O?W0?g9 z=`a-IAcj13nqW34x&lcI#;@wa?AGa?1KRyTzjn-?a=P0@WOAYfCX`Jg)Y)1aOcKL? zywHavI>JYl=fa`ZWKNR`6p2Eda3-?F`O7m4di2vrcM`N`a6{zj*8?Q%4FT=Euwma5NT1 z6p%N+y21|D50FXmOBM@hnA9N%Aeh=U{TU(EyzW^-5=lsnOkMOy;`a+zg!zR9^zivZ zRaPeWuwpu=OrRwft9EEes4Sr|0qhRmFNAr+KvkA2I3{RPj$v%ozCzP!!`O?LoY`QA zt-HzOhl6BblMD_RZwew3RV@fvYas&hB^ZO-9BK_`a$q8kd{MR6@nswHm8Fz-WZKAl zzoK}FZ}kY7VJ_k5V*Mq0oHKMaU7)5y;8hq`k6`*Fcld?YBG|nC-q}3@6CQC^i0J*nb1$6c!Uc%#pd+-|spNlZ+q&X#fBf zx6xABp4D<OWENipe1q69Xag~$FN@+Zqk$x4i zifP!_V(zGOI`oDmX#F@gGeZ>(cesQYizv-M@7~QM zQrXKlXp+yV6<&3Xv&P!Or3HgC#Q+2-&N0>kX9bYUGMo-*Li&++$$Gz0@sFqJDi`9j zCAe+klLyS{wpmyA{a0TC&XJi-^)x!@^^d(F-eB`!**%&_bo+%{-vv`=W>u*|ky74w zT8TB%oT@(HrHhy`aBTg6CwQZVIqv7naWQMk0WEQ%=Uv%*ZU`|sJv-4mZIJ@N;coC- zi#Q=Zxxy^JA6L;=Dp~k4xCLt}vPaCjG|Q~(p++C%+i1C!^Q=id;(S%^}r} zX9FUot$X6e7!z-;Tq4DG7|?3c2q=TydLBGAUalHuj~|*)Y8bGu))$3i@YuGTDbZ`W zfWbMdAp>vsw*M)>7iaC2H6h^zB?3)#(9r{^u*DRL1hHM7om?3wkOW#!(1{2K zIq@|YONlZUM_t!-uh67=_L$Au+}%jsP7Pkl?F+y-tT@ID8wYNfAyDBUzLbB$iY1jP zLoq>v*Js9Kb&QUVMB5!GLw1g`snl7>5T(ruCe}aNYOph!YY63hrxu)O^g1Cpcsi2G zSd=LGTvr{LnSOU2TBTBr8(lPkUZXkg25y^H=Fg&jZZX{7D-ykfg}EgcbO~IIs|ZQL zSoB4^q;Dou{1jUtu+m}BX5hWS^v6=ZjBI1JSFyMB~Js!0CDkruD zFh-8~hZ*10fc~ME2Y0bEEko*9FbcL-@YCkQekPZTg>1%+XYFN(bZ;2bO5)!O_qpRb z@<^5|tdd`b4OdReeFndgXFKbQ^i;wRCRNA^BF8I7f|i)J)pFTEDV67`&Qh?6OoM~D zVftRjDQ4|Zar-4#uK*#fVet9mo0C|g>LaZAzVKIET(7@V*Iqr#sXvO_Vhg|Hf3*^^ z_iPPM`Z2s7qmyV(-k{m+5~dFqnxp$PeN{rFNQKX%&-NJL{V^RiT%6U!|8kV?)WrDx{6QQMG;0bs$#&?L;K zD$}Xf1%(|@({X;YOIw1a%qWpNh$(y1SG|;YIcVvm0C>}3cfaCbMY^_j9^wpDpW(IHPu*M>fpNEw~dy2IK}Sdd|3&(2vwxtZqew? zVm$W4$*eGe-mr3aaOkn2`gt|}?DxTkRp;~mxaEYb1HH#!yk@15ud&v~?WS+#R3M_; z=vN=1z^^^IE^rU(kMul2ND0SVQc#r*V%YRmQdXeR=D>7;a&Jx>{*=%;W5sQPHiKvG<%WAdpg z-3RsAUpEQaiD~0ios=0dCDT)T=(O@wKVaoVjp9Ze!4E>}1}L2Dm}>ecmME_z`?lZ` z4Co-lbKo*?Kx61b@rpa~EdkP4s6|H=y$IBC((ewdh|N#=^OUu8rOY@OI(Onlj;)cr zh3}hgR~G?xSlUqGJqLhTeA+A((CH zrfs=S3VZ|KTaWgY_Q=uIfkgT`HGb(`k5+vIelrW3t-h_$SXMvXYdlHA@OX0-sBj^) zm4C79tZAr)tzA08aEK4>%$ zT;P76?}5kETDk-gJ+6{HGq;e`>3v6VT)gI{K&6-Nn_O} zTlTgM;tleb1|i*ibiaA;uV?imE(U4Bk3tAp~9RuUl z5h|{>9(iz& z1CLwmAIs7kb4f*OiIM=?r%Jp>L6^|d^{?!%uJQveIKTbxrqQd~suu80R`Y+u&5jsV0lakcD>DBA zyG2a8^}I1CpstVXi1%2fwvCiN?&#FKFN3AUYA8}{-yjIws^&46r3EFVPB6<>Ofk6} zj~N$UcWB?OQYF{@doIC14+wGqxzfM;5>#T&zu*4zWkmquOLxuke{R_S&j!#JqcuS|55Az9B4}dx~R2BM@EMK{eu4AGA+8J{x2n=fi8FQXY=K*j@EyP z1Py2l0p$Nt-G6ba@d0r->^FxSYyM}{{#V6+zjMO;pD_h(HG2IQx~>=-t`6aC*$(0{+;{_{5e@BbiBXXd~EgFqek|L!FG|Jq5wvE3ZdaZdm| zommQw>ka=?h&y_JfQ;{<4X>cBwR^8XaG&psbMv?4SO3-wncaI4KVG{(3H}&91sNXQ z34Y!RcJIk{L1q5}+2P7k0T1^(z&%q#YXk8Bfq70NOasmwM9L@7i=Fy21wNis2N>ra zLQKnn!2xPc(&ooSfc5vdJ1u+;mWfgspj`V)L%SG8u#&9{Z^7sls$cZedhN1-JtZAb zIT=dK2oacp#w}j6z&UW-e%8_)jfLU}n(GL43T!(8zAg517F%N0YF;3JM;ZYqQ|f;Z z{MZ#Vc34a}m~yq-wzv-s7$H!(A*N4zu6?~qPBL-)B6sJq_$5|+&h(;W1EAoo-6+2)|c+I9mQ$grw1bJKVArr+w(4u9ifEmND;E= z*4=iRrhon*pFuwQb^&*84~x_qoRZQ!sv-jcW(h6-R&vQA+;Mh_F5~Z_w_zha(w{c( z3UBEt!4IYUvq8ViG3K;_G>$Bq6$BFJFoBK$%BErAjO!*Z5hNh?Z-@g%JYC8Qf^*u4kc8v zK`7iT;{;^FiC_^!oxoZZ0#`{%CP$ycn8pUWXgLY_R*Hl$7>BGCs^vLoQN`7)FQ*1! z`oJMGv~SIXgG@V6N&`{zsUJ|bolAIUnT`eIL|WD0m1q--Rm>RonHc#grFf~&RC=L4 zS5bdo;LmcaP?yH3YF790O5Ux{Q9YtLi-~t|KK^-f3#t*?Cz@L9nox)_?UQ_x6r*v? z4)?68f+PRqh}S@-u(?h#&xWPkaOR}fJl@H}6B!(sF!QEv01o4Lo&%Sy$@6u_jsI@yJf+u6k41 zh*@NrYbS@wYw3I|>0yxtq&%1;l`{otQ61~eG0G)7r67f@b$^_8|EWx!Z?Bk;cMDbg zAtaK#kg0*7ZY^*?SckJeme+2W0R9#FR$H!|ooph1V~jj%tjtkC-^9O z+-ieQ51$I-g75Q!nd=J0ByKZ|yBlQoc7lOmN9cj%^3Nc(;LWGAPO_=y%MG<#j^IGq zx%ieR8Azw-`CR^I1JBQ6?ltD~Dds<+3v1z2c=u#`~N{fq$PRwGs zRND8(ry3T|>j=ZQo0RT1ac0k1S%aJS>bJ4#1^aUdElv5fF%PZ}uHagudhUArnu1`t z871&c^NR|UxZW{q4P3wvPb;h{-bAH74{ioW*+5JpWkyM1EOp6uw~-DZ9BW?)CoFCaA(7uFB(Fva$IV7AqZCFek69q|T?I<) zHz{eSn^IY7T%@<>ZSRUb9L)un)9X4RAa`5I}SX* z4}F=qb=u@Z7*xOjZ+sQT!_AZ*QxGAY9081Qy+E=+Ytov z2)I~w(DAS(%lShb6HC_}Nce8N`RnuDZPVxRpW+ro-$Ip7yA+dEozf%dpA0?`o_tj6 zKCP&q!q&B|y&fU}*n-TZDMaB~LSNTiO`nR=?f)70 zf<9HCgp8d@u1(vYJZ*`ZH5+6Q&kKoZRO5Gx6@|$<-_q@xJT@J4t#X6Hy|Xi&uT)jZ zzM$;ey$HpP$QkyGZvFZ={3Vnq*J)WsuaSP8LB-8+6E=1sw4pZI9i?i6)daq15raoT z?37zvl@kQsK!9n3m*WQ=msydtNsx$1QmC$JBp2nk=~QX7)>aT96=Y$TOi-U&j9;7V z?GqbWwe9--{X*8_RsD+THMDo-M(W3Jt*IMgad$v%p*ZQmlF^Z-CM`c)NtX=6Dn!Wx z%N`(In=}HvJ-bqP)PPb<7iUL$rGU~ZORw^+^+mhTIA`3cd$mPc3JupWTizjRA`6FbtC#10VQLm%M|#aee}eni8=U?yzzgJmprEJt z8fo+ZLG76*MttAnXjY4kbm!k-l8`aeccM;+fUq;RmTNe$R-Q1PiOQ*t8kG-YgC`yj z(G-grY{M4~lkYXn)CT2MA$)kFAjljE#<(DnLL@)z+vJhpE-~v1?ChDU3OndcFH{|w z4!doh*A$;jbZ@xkX~anI>buo47=xg$P<6QE^kr!kQ&xeOjqMVjZ`VD5mEq^q%<@(# z-JFaaBtQv%ph;z(Ip=NJYUa~`=4~=M58whyX(~NrQhblA| z4-M%b0tJ2p+A5*e$I}j;h91ypTK1O%naXA-d{FAo7@DxVhhO!7Zo!Q5f@6n0OtR|s z%)d9i$GJZL2z!)RPdVNmU@(6VY>t6J;ABJXfiDXK)dvp}Cvv@O#?#eq z?mEhpnc%5c<3N!vu^4n*Suw-qZty!Sj)Efy-NuF<-LQ570YZz;oFsX>@=sxmLKB+Y z>n(o7w^GMMHPCk6hr!c4!QXdhaTaWZDd1*J@cgl7l%-&!DoUx)*fy5Qf6b-hU*=_L z(H0}k3-RSO7=;?b-D&c5AOe2Wece7yqixNVgJOZb8y|&X%eUjtX2>8uV_MT;t~U(U zlCycd0lfY;`8oWt^elt-hXo#P?ucI}T@8FUyoOC+mFIVkXJ_yKJ_<}r=+TKo;X`kAeb@2t`Oe!ljvuTp zLN2y00D@SsKlq7vtYQ-z*1c*zsz0XdW)LYS4#1vBPsw-{gKmvr| z$dKP0perhnZsbZ$+(=Kw7MW8nUGoK>$_f_{ByEF&09KEe)N z7`148cTpsk$!%578%W7wN0R+~;YkBVs@sR2rQd*!wo+G>S51L&Mt}TiO6(VBeM-pP zF4`L-1{1z}U`7|j;$!2L!s#c7kT1V_hZ1^!TpSKOQs{BsLmoXH15TJZ92CWiZ;nW zTNUrmsu-fT<~+nr@)nme^t`vq>vdD&rl5GArh+3+`__kX?v4{c*H|zf|G?Py-TX?J zHouNo>pSc`9s$~*ClCm*=-n5I9yMQ&tB7}sQi#`(qBZbvAa(krK(6yAAoF4apRS06 zs89ab#oSD`P(f~r+_!PhVsgV0&yOCxs`$H%HWTw=L6x>ynn7v^hJlGRFwsbO#TsGO z?FV0cD&WnC)<91kaJkmn0tF2n@I1PihMC^5EWW>28^G%I&KtJ^8m5hK zR%R?zi%e|y)v>0cs2|U$0a#nSDl`TL?j#`LS9-{&-GRB*IgxL7dcUOcPT8y;iG4zy zP(pngOi`j<^#kb#D_S!%(0ch^ERK8RdUFr|8}sPCxT{o9q8k?qvE_x0ew^T6*T<;u ze)+1eBFtlu8eV+ll$8_l9o)k=OzCsKpBZw(_Jml zYcEwT@`3?S5)4bf=;(hI&OPQUtKWC4`%)fL94YV<@BnmeeZU|zm1kN$pI#qJ&0`+T zs7+lJ==pf3A;SEoZPhV?>Suw2F#Y%qViir;(=N^U37bLbAZGJKImFdnQ!Dp0ns@xT?DD(ogDX<$U-upE_T{0RPTQs4+6plmBs&B4wwtJqH8G0d7Sffpq1VWZp?5-e>Z*HX@<5@6%1ny#ua5jz5=OZ*<4U$-W12&s! z(L+iD_;FD+498$P$AERB9vj}s zY1xWqqOLo8yoCqpKFTRz$6=9GD%v9R-odr&>g|h=vDaWASwf7_rd;UdC<_&oV)Uxs z%C40Xbj-oSJ%))N!-%@3;8;3(jCpOzdv#mF7DGI?d)B^-T~;4uq2Ul^qROQtI7FOf zDGjY86&IIK9h0HiUKv?d<929VHuupmOziqBF`-$_C{w!Xygo^4m#Fl>Iw-9E_`V9; z&ZSWtbpJwqB{93xF+tt@4c*W}gIE*8K5p(Jql@sq)z}CPN=g!7YtK&Xo5j+E)ih#Srwh*j7eDwNJclLYpBm6~ERS zxZbaIQ7%z~R`C#6$vWm5Zfw=C<4@Y{$6y(d$b z%8C6-hHqz;o6MB&pKbW@xaf)mg%zaYI&UAi4XcxGG8sl^p1y(ir2n*qkZN4-8-c~99w0HA;v$348=-_bd;ZnDDP>m zDB7v!h((KO*gKl6y;YWWn;f0y26HBxh0S7Wl9ylyHzI0G4udW{b>Q=&RX+cs96468 zC6fO==Vu%4x@N!Jbkci;%Ydb{AQcblc9r|Y@S9-a&!5$$*M1_ncCZ7yq}dk{ZTs=)RG%B3WcNK!-Fxo^#a{-<6!EzaVT7m zb8VmH1$(znRRmMd^oELhK63jW&n3*O8lv%>70hG5Nx1<#7F+P;(9wTu*rNpKz_ zA!fV9C!^7v<7wl%4ZR_eY7t1XsYx8qC4VxZY#ek?P5eG`y9#gSg!?DL*vAaxJ(W5# z$e)fS{(ECmasbOt(>QmDn#_Sd`M26M8-olJJeUmPi|>A(Op$@UqIDCVGW!<8YLqc_ z3t{fCEK|{Dh!*1@E#<+d*V(Nk7CuI-0BkjQ%ycq7`Eak)B8WKfJi87(EfWv$%D(DE~6H|N$ zH6R+SvX07^g-%>_#pWkAFu*nH(Sp)@G~H#U_4^j_o8{MepK>qbSqt;44x zIZjrtGwj}jeq#InoAd1i!PA5fbbG*q>+I{q_*%I*vF@?Xx5EiXw4!Wd2@5&)cYSyb zUe?oJti>Zc&td}`M-J97uA3Nco%>piz(im+8KBkFgY!1vV?$N84b^?nUI;Ci2{Dm! z*r(TlV)$}Tk}?l|U~jCWRG;7g_smm$0qaws2R37Tjv&Q|78k^e}Bpru&dCp==-R}vmUnAj(0#|l% zllCEGB{>eF5h_b@DSM=u90>X4#h>!er7lJ>4MmZ>KUI8f9bC1wM;7V8+{N2=9395= zw6l739tFW=NJh+EJ>}NPkOKCbV|R^BOS;u1iRbHi>Jg_JE#62Ukb#j75xDYSr>_^` zv&i3z+78%6dt!=t2sXnW*VcKC{wE9Iw~CXG=<9`Iq0FCXZG~}e@-o?}7Uz{NPwPKk zZlfvQDr?RUNx8X<_1TBn@{y~+2j2FX#*A_fe?_r3)x=(mG(Meytzxgbt0v_$ZuVqBP~=7U8W?Z7nPic3 zM!Z(mY6cyv?CroKReVCwkNe3 z4uKX|&IL+{tL-T^VMv@dM*Gwn0|ys@;8}?NTMn(0#U%*iOt~AWzZ`oOs{1H0)<*)v~0Ct1mWH@v>mYvqp@%HBRO+8@lug z;I^`B^sw|@4RH#0weZIn-;S(t225WMQFtVL)ei4@?h}7Lp@=q)rehkVk64ZmHhtD)S~FOVq}z&^ajGI5h+63kmCoHQ&ea-DRd44Wcl@Ld z6|nhQUHBF;^^c6@q;&*i=w?7(gbBuvvV=JkJm=&A5^R+`0x0%i6(f&}@8-K!?KB_z zQBS#*IR_NPHk!T2AlxC6QPwtEQFK*drc@vQqeJO#=8eG5JvzNpoTe0|r%Fn}{+fi!W!+D|a~|wsKZC*$YCgmMvL56Tv&s@k5NMV+$>7Y(L`hzh4w=ZG72qAOCh3Q z5s@tjKO9I(^haTtXf&Z%xxq$j`8L4YwglzK`7@^?(8ky)>sAsah73zYKgM^#9M@Oa z_A2tkgY@1tt0v)h67v&DuaDLA?g>?y^s#GwF{y5h#5$_gsPSIa4~yfm+$$xuwg*Xp z8iR(M8N90w>|=)lC1o`8vSkoY0!vp(5Vjs$<)k|ut-LpdoLICyhz#&xMQU52sw5%0d%6SYnvLkhv@B@I@$##5?q$v7>7~KQu`+7N_E~}HOTpFN z#JGws^kjTzjPld4p%aLFaHS_qV!f(mc9(M0*JYz9Z}pm)iC*+pDO0MF3prSRkEA#X z%v6p7d#!cF(K?7RXR*O+fY!ec7eN>~a?X{});~G) zNrj+xU!+A|fhD)@L~aKj2c^|!Uw-Z9oDxstDbvr|rfhm_VxYe|;m>G(Js>s_z4olWN}`?F4{U3e zcrY;Pg4R4)JKJI@n_>@L$Nm?_^lBQ@ef{0WI}kBevKQi!^6bHBE1l z9A)<=N?*H6Azw9QmUIiV6!Rr-Oh%D_JH?WL5khn=H3!XBNbz!d`{fW@v(WY3+b=8e zh!`=NgKf-d<$Nk!yQ8HLC^tjL5ttC<+n6Q!kf6b*YllPchwNVT(NdxuqWEfPV~uw} znOh)f{Qi%opf|W>K^mF01o;ol6$4n&S08&+Yf zPuUU}3|X(?ll5ASji)m~;-T)5)PqY9PtPL4=zqPh5F+z)mEba5?aP!E>fN`An6iOl z<5K~t0)i<1Unt(3u^*pakJUX~eleCN|1nyw#*<#JYSp;6Z~tk`%ix_Su?uo!&KxG< z_M93Im)-z_!+_ptM@-QVd>(P-lG$g%3U}jZ^3l<%S*)5eHAW@nRCm3*va=($@y95f zLIu+~FLuwdV#>=AIwf9hQ{N0ks+`Iv`y-$bbSr|h6JDR;dX>yHyby-qsk(+=)-zzj z_(uSzkWrBtB}oK_rOe&2dx#{OZ&ZJylsr#NHKVe;Z>bCk0bIv^n@3NZX4kFTk~|dM zj)tYVy`VkvnHXUdTcb#@^_o(+;$Ursm@u;lrCMJ~v{_v=l5sC_a?J{1`PD?*95E+= zm8>GI*DeSWw~!ww*roiSOb&ONRV2akN>L(ij5)w!Ey^;i}N=fZC zuO5PN&rIu5jF9Zd3<+t1`@&tdNhJVifECk`B41fRSjPPgZxPcM(>TN@vTF*_ zaJl-J#tCcD_8km2F22{Vw(!ssF@*#7GmMU!OgNpwH~YAW%4y*&-O$v*Si~2&yIiGA z>xG6nARkN$n+MEy7;oRLWl$29S4Qk~++-V>Xz^ z*P3R?G6_uX)N46yORf^O)#oNMgcfqx3b!CB33S`$p4MYIMe9)1r`6j1i`Y-q?4TNCMGk?y#(1x@uhehi1IWh-FDZ z(q#FgzRkQMYJ8HhAKFq!iIunZSmPu+1?vo=YD|l&+cV6A)7jtFE`~&_SHjaRRe3l_ zFqY0}1P&upNP293Gu%dij7IY@B2`57TE2YSj;2WM&>NG*JPJJ0wHDFW4$^|iiAC@3>MG!)deyb5@D6;m!_EWpK#J>Xz@%eeqw}p9`z2u$lz?k!6uG{>v;UD z6B-DIsDG}6L*`G)HzkDaYnkChoT0jo>--1}M=t6D;~1&4Y>I;A0ozd|Q>0b1gKeNN zz1>`@ewNJBlaHI4XknfXW^X>#B4(2l`Gk8rySal6cUT+*BHyH4Ef72XdTeYJhomGE zKymqNv6?mPGxGhRdkM{PSb}SYETFGUee(6pUbUXwyN4g-0m%fIX(*J>-%_cb393XT zuoI%M^}66w$%0A0UY+b*$=R}HNUZ@CO`aiUM7U4y3jVlR7E3az3~j$eqJ3~PQmv_a z>s;1%)q)?fZ9PQgKNTBXb&uQWc`I=5;sI#RTT|3LFHZ>d4o)q#s2ji?pOYfJ+8t{b z6s~M=T_W#Fc~&Px?58T!XyH(P>hr_4!ugp4#y%Q5)cgS8bN>|^lrJN|boH|d-Nc(E zv$JnkQ)=wYxqBZWXk5<9o;R1KtG^#N|8&`X6de*^GGRb=;nNHXrQU0kOVM|XWrB*F zH1sbS+hMWYXyg7p`};tCj({Vvl_^@;oV4Oz{dJr$QGsk`B%kZ^#wCCjGSgi)o78@N zh}D!FI!sS7VpDaw#%>EG$+N@+#TGMmG0`|o12BDeA3BG&Q&40e_rv8l4+?#5`2)ql z^t)3Z0AgdY7}|??OIa!XrwHCY>Fgiswt$mRnx~$xI~komK)L$v&vzV&a~TgFGdT-( zKuPnkJQw`*oQOz7_-e5eNWVR*>o`>v5ypr$_$z(0pjU~7>Vz3qg@u~rI{x{)|N6Vz zRi!6dZ{a%v{3~yf;SdNKHH3)Ln(3e*a0gA1Arhe;oY#l&w6Wm%DRBKC^SS=!RS07o zTnE4heR+2U94lFq@5)~Cu%dMTCJmFR#?!Bgs-*~gq-TzbscYIGHFIyMD0tdq`TpPq zp*ws8CMyA zY>x3*^uBX_=wW8R{xU!8_Ho?%gZ(C4q>9@wISSiy&RGhIr(Nk6yQa=BxM>kYNlhbt(6nv3hjS%)YZ1Z&872B_^fztcb>rRT25kI%$wJ({vJ)g%&>sOc@ZD*8< zjsw^11av=sPw{w@?x-XYJofUoW%e}yxpzk~3Avj}Stm%lIN&6uZpqU;Xt#AmYUeJN z@XB`p4+3+2Q%Nl0Dh}-LM88*b;oFAR52gPG!HkV*RcBZ1h9EJ9QNxL7td-pOjoz}a z9Vwe7->f_4ygBXsWBF9|u+y4p^p8Mjd8;a$(^CKb`PUc0%D@slUscX1!@b=i&G98T zmTRuIm48rAd``M4kti}l@T&mQ53e7Zgif79feiWa&tpZ|B7$&DwT;VkUlUWEYR9JB z>o)td2X}Q`2WRQUQ|oT|h)erIe@%k0l`65f>KcRR8*dJN zsI5dfIBR1aLj{umo+b9x&lPs_L0#;A6-&3>_8wx z>zB`h!@JPAgL#9ejLuPS2|5*?(RG0KrvYdi!=$fu6svfv$jEk`=2zlB?ltdc-ed?r zLMWFReW){l^MYKLFG|mHP+jgP8;vnygB-^bE>_W1S4rediTv=Go(EKavFVoC>GjV~A+aCfYlh19 z4w*t5mh3)L3NL~cp$RPGW>QTRc@?xQZZ^#uF6l+1xqgO>!u!+d$yUl~{XIe?+kTx?6u>VmSY59`SJ0X`6;rvU5 zfUCNuDURs~JKi#80Gh*GyVT`V+>n6Yw!eUTe_i|;J?SDpkpj^_Y@5gUroIoYaVBD- z`)Ry8kOY~E@7>?3Fx6NT1oBrKp8fqe(xH<=%O{7FGGv`!;pH5+(x)%}QJz84Qy(Fc zu%v#I#WSt2O9G7M4O%=ie^2BGON}bj7J6-I4Ktg{0nx~CE%ViArSy>Bt(kkT+x9+4 z#XOyP`eivlF9K)m>U^d|ho=cg!KoLe%BRMH*U?3lm$S~vUL=3r%UJ$oI0dpYULbwg z6PG4(vnICKDjLV2QeC#kzqFX()v*eSM7pX&Wcrk?(MKhGteiGHE+xd)u5U@REzGSV z+PJE)&$_7KvWn_G*6!z)nxAc3)rwq+iFL%T#&0miz}97A;(~8$$}5r}R)!jqN*~}b zdTmN!v0WN3ehe=56+}=AOiVkL@FZU3gtveY@`*zQ=EC2=CY!i?Ga>r22!g^k>d)b; z;mw$=gbL(0?lUOK+9!acQRHn;pKjb1<%Y3TOJ)DMYmEu@b}7@4GHMikguWrAFR5US z>&{FWUe3NF-#pB~^J6kKMT273CaIRrGE;X3b^Dy1GHJ47a|@>nvVAvh84o?HZG$-W zH{*9Vi#m}UJ&&1@eb%I})%*u6Qn?SPgL!R)=_-ei41CdI2Ylr6)T7Bduvw^^WDWNz z;oZHumBtaNp)y>2yW@)`DbwFR*et6^l8crcf>Sjbg*2xlu*F75h(r77Zv|d6us`{_PRsC3*P0=Auqa7mpML)cQb$;_NqiA}a%;JN&0ZeQZp6=RQbe zq>nsJ{e!tLEiYix#Bh7Dt0LxokIIJ}@<4j~X-BdhbnU{lu(C*8Q&&?qcNO&bzblNH zsU8}p2}+>M;$jX9f)+f>%a=YjBz5vPaaj0NbHyAJEwoj=6VbYyWL~(y#8sVC5APAq zLArj@{;({uu$O}tU45Anli#=>68V-~2gKZ4!Q8fN^fNFqcRL{|HB=h}G^kstB5h9V zUX-sQUPyuugiO9VhymZ=#H0?#=N1yWrUO$b8RAulXqa_IE0$$bG+m6UsEMmATVf)3 z5iV}WRoqT8n5978rL<*Ld$`%*GpiZIt-)=jBvx%pvh)@C1uiTYf;!v7lxliZ5|`gc z6JvCm=j8}}v;IakyWYzTiVErpe>@-kX#bwCgpqv=Eq+8e$2ksN*)h(aixhi?KcL(_ zt2&Y5eVh-?e$|gfLo=?~43S)M?UrD6g=G^bs(?^y@4XZyu{VkxGx$2SPyhDV8`Op4 zfF%zvKTkx&jk7;)Vp0^LjMf=A_n$t|WXVqQu#d9m2OVSqEawsn6GSiwIGK$PZ(BC> z9XUIqaw7YUb=*+uc^e`_ zkAkRKtA`4U6oB;XTiB}|x>2alMjQShMw~mDi@&x<_YPI(JsZCo{ZifFFL+aBft-AL zt_cYtT;yE(NaBwE@BzBfp&EW`Oed`3SH;NA6ND~M(%p*mI0d9)NrS{5^eD0cISQ>= zh?3YT_XK1!(Rrwy48M3{-d$sgM5OTZ*fk?`-0_|Wf={t;?rgEB#n}70M4T*7h}Q;A z=WQ~Wl|?9YfZZ?X2}4FaL$7>mgD2>%t~vW|dvU-pRe-tQq5&Nyx0wkX8w(JVlLv+) zusA(?r0Od8`*awv;T6-4ujqpUK6X(oL;skQ8&kvtz$u5>GN{eE{ro_`SpJ3)J=$T0#AZ4i%SE+U9EK6@e2>zfU_F?k# zq-4!GLiaiQ_Mu9L)yFTuA9r+F#jSQshB=Sbv+hS7_tf?P#&%1?_jlVdzb#A`Y^)S5mch)6%KPM>z1b>%AP42UCO38;m7Eq`fydHB^_ zRC?C*x4Q1@F+hx-x7tXYcpS6u0+O8m=?I`O_YTC3nOfXdvJvHUXTWaa@!tc)NMDYM z#GoC+s z%2QI4ClUBdI3S%n?Bh0~5yv|Z+WH%Rw}lxqG0}tKWYo6I<_y*B85TfC zY8WW~_QaW&S9YoiGo$@_ImV$Sx}WY=-)X4L)S8hd)Xc}nJa=iQLoG0-w&%%t&gIsv zT%*p9TK0qUksfTwG4_pVcu19S^$M7k6bz_&%mcprRgE>V2YW$to>wmCRqdV)NcnS^ zQ$`sw7J&et-)atFP?Hv3`5=Pd`n$6W6J&vGvq`_ib~&$qnV*YZO$fZdskQ~G3G=8#jzArnhB!YV)7x8G+8Vh49Zb#I7_ofBl-5$e=RZQRzP{{kW| zLgkn7bQ(%w48E7t%3}fvY#-_IzbEifER7WI4YK)-N(?SY$=bmWtcnSnlvRx1(YI@C z@fCN=o5Ll^E7JY%eB=|y^Oe2`%rOMHV(-$5BQ(iQU^S|IbYw6+XG~1;z*&n z>=X;r(&0NNw+M35&p);DgoFnzlU?3}(_oQi{xN)jEUJJbjx>uG?jSsnv@rNq1C5u5 zyl2>8p>j_jNpk~n_v9)dlzN|j_x<-+4u~iLf%R^hp^*XNf!jZP-pa#DC=7(_S#Oce zk(5HLhR#H@JMPHblTOA4;yjsCo5FPE=gYQ89S4-EvA`=z;fR;Pq*AFjr#z*5-(EqB z0R0{Yi0{&GcCL+eWnt{f?6bLjUx*g?H4h1 zKZ!sO8>xE-g^#vzP&Hn9?Qqb(Q?)-`j%Rc+hC`n9BcjC5!5Ob1?~H1Ku`H?U2TiK_ zSIeS4sXNqDy^dl}7xd^en{pg|n_owgIXXNl5uhHv%WM?p5u1YOxXb_6AQ7G_qSaE|G#h!8S3e|9YRZ(=Q#o1Y#}Y~OYX^UJd5U>c z0c#e)3DbbLfm3jJGra|5_0H*+rBV0mh#oVZxrFG}OMOH~8qnyMv~XsyK%I%= z>QN!x=j3dsnm=f!q{oFs@MVQ_M!Go_Vu9AM9mA;;%o|j5v{# z%m~y|#=#wo+`{{xyPvr3>O4GXHub_xSA%V2vRj`n z_4W39#JjZPK4n_l6jDXNBM9Cq?@NxfQs$`2!%w(jVzwdvLz76Nxj2!vi;9veXF-gepOwXjT3ZRmjKZpSQsmWzxTBxh?&nKPIQ*&LU zgdD2zG19aB5*^oj%f7-H@teXi=d~J&(yd@FgzR8NJYGPc&a6jE#l~}Myb!aA`-*b^ zNSQAo=xXbX8p}08)KIbO9>Ym_X1~i2NxBlR(our6TxR6-v60DCEj*Xp$s&^%-R&oI zmI`CJdZvdF{XawAF?V|GCD*U&Y{FISceK=BPaTD@3Xym-vV_Sv+K?l7&MERU=n){R z6FMeX+KGH|G-yF!Cn1wWN5*X&qqB3CUBJ~5#5|FW>4602NI(Y`K24c|-l(d&VM16h z1lOG?f0xrI-g0+3#Ic5nE)-VDregHhVcsjVk=_hdb~ZoWA-QHgjg~iObr-Ag@Na9` z@wu<$NdNlgxy>ikjTssMCy?}E6Q+JMu+F|#JI0;tf+@J3q8;Fzkk3XJS;xnk`>^%_ zPMtQb10SuzJdi_57EVNN+_*is9#KS;fJ2H1EE?Z4Gpq|p!bCrY<3w+uX8_ieQU(=~ z0};-R%=BOTP56cdnWKnTqNcrM;o)~Fo{R_)mVKkNYKSl>pRnCcELfMWdLc&F6LZw} zk;q1EpAE@K-&&2C>*}gKiL%*5&1E*@1h=AwJjI|*_lNoql~Y}EfpfPZ!^)lu@dY>H zV>`6y0hLl6rye#ip=NW9L#%8dzV-7|ae7nad2W@V(Bxa7lxvB+g-f056N%9{V4gLu zEe_yrn5R9Bk%qon_D@8TA22N*{B^Bvi&@+{1ljNcRgm~V2YZAC(OFncAUepWmTSDo1c6mFs1b=T-8?Ht?ckD> z6Exk*lqv3+Nczy$t-ccnNCWY>`3Ez(pE=$qRDDEt5D_>| zMzN4y8xO{W`pM5JmSr(9t_P7$yZQ$js41@UKd%9pFA2os{>`k|L9{HWY z|Nkg*#D$(#%zW}s)zGYje=7gWQ~mY-bFcrt^{zr;*S&=1)^XEn{J)FNHVvo$-$hu6 zu96#YCD@1bK;M`b-(mT#?mxd1Ob<1EIo}OlqT;qwY&`<_pb)sqnE!pO|9DDX47^W^ zXPg@g$8#`TPZYsP_K3O9=xl#8X5W)wM*LMUeOq39_dep~%}>D>lyU0j=MLq+V*JF= zzjUkp0$VfEnu6<3nJ}<(Pik!T2TyDJ(EIwL`?E}aTH#^RF&oF5qCIID;ac_yhZ%cp z!=KGnCw_V(kX@j6GX9kMUJF^<)Q;=^@9GVPYBs9mQ;|qWv+sTa68oOD@)Xqi(>>vh$hUlF;hw zo!QTq`I`5eYWmhH)4UUij&_E@_SCy8d0=d%7G4sgxa3-5 zuDi+(YQ_O~i(HN{I29NzmV5~hcl-P+@FB;e>mHE6P!=o$`2MGKNiwa{l=&N%I4)oV zQwzl=vxai-&FkYwq;k&0u!4`CYvVVkMPYp~Lu6lJ-f%vqdVW$Gg@4YZe{Kvw_I{~v z_hc=)O?L1so#Wm#3Ya*g4HwT_7=wIDFO(y?``xwoG8^GpSE9M}R`uQ?+dFDQtvS*$ zofASOyyk5tWXRs>^^CkVi2gTct20trulS>t@6UY${TjfnP;futbo)KP%7o`@8(^QR zHnRj^Arh<5l-Cl^lip<0k3ast)W7@fhT9gGvLd|ZTq6B)k)$Pc&$!qqtQXwSgEDJ+ zGXyvbNBgnG)tQK^v+IfkadLK>kOCD*DsshB16wE6csj2EPY6TyT=6*VDlbCmS$f$0 zxRoJXb1$J5Q6y>p`_U9-!&^H3aAk=R5* z1B2t@>7N(bfk<_zF_-Uns6KR`gXU=qWz_{;UXXmolK2+0ub<67yNQQ&2VBwe&kD9m zZk7Nh-UegQHZTQkehUJuN8bV41}c*qo`#L}qR5Wt&7zUh@jM`-27WVKTlJ5xYviL9 zJJ!)VfJFk2LP<+LK8TCJQ~bSe$q_m!{GE2`(hKmJ5I%i??<*E`v@1z>wN1og^>)C0 zdI)!Pz2+)5c8J|iN$x@|8X?t!32;9hx!c_X#L-y-P?NK4hF)pJoVfcDZbZ%{`MLFAyMlnOR#x^%(B4LvsS}liqS6)kjjHXEB$u%8^dNj@G~)5ik_ea= z%i&_$h?FsFxR(7DNkF9Lt_?zQV;IcslR$i33a)KZaejZFV=SB6>j4&1-r*6zSuu(d zoV3y#ld|>A|Fou|4_$F9++a{s;*QGoeqiFQSj#LN&Fl-j>E$tnw$U!+UA0H-I>yGteSeMoZ}g&UPhEOo%wj_nwe<=#)z3S-!%Y;ZkLXL`29 z!AzW|WEMJU6p%0RVAz%n#GBI*Bl zT_IiD69NDjU!r5saVqVP{~6y^FerHK^p??(kB9ciA*d~J_VCOpP3DzXfYbAXUibZ^ z$KHiLW9trW7|w9GbbtKV!a*k>=zploJ^42UylwsMI8UG;PJ}HxiXT9XmG1JF5F&+7Tx9v2|e8f9))$H9)P$q z|2~k@K2DVm4tA&IS|Tq^nb}2EgB8gxvi%euetA5bUidefQ$oa`{xPDbu0`Sc69NX zz)=v)C0=J_M-A`?vPioTxID2KWJUc26PTk-lhJpQU!(?Q?%~=dBzqO6K|3|>K z`QDSu+abJ-qqMnyl;47T&y=101Y6Bqu9Kq?z=A;9O)1OpAjjp86NsOxv53HAkbi~4 zR=IwfBR443qmttBa`or=tP+i1$pfGTdU`&4&r^BB_cLN_n20)7u_`a@J@avM6$=|% z9E`3$^;h5z-Tj71;q|@5`99u$lu<>-58)u^FR278}O_>#E8NOL-3|B`JpN> zeyG)JBVql5D9EtVthXGJn%rBlQpyYyX7|@?qx8_< zPQdNAQ0BQc6i3bNO}a-TASddPhecmtRG1(KW-L@!t>T{eX(I!YRd~-K-0bN6 zM<{rg`G3ACd8sjs_6~Frj0;b4{6DheBoy0c=|i$w~rO@?l+RG;1;w`TnPDk!rc@t`jYr2 zMck_xz)|F)fT~02U{?w_r~of!vI>J3wH`~qK!~{``dq!NX)ksEB?(*XNv%#!@HMM3~G+Z zo7IE8Rvj{LJVSofe)gXh%shqT#^32{AvMV`1m+VWzFh0@Nt;@|sY6xg14%^|^A(6~ znmzv1dLImMS`vRYe1hWn3v^YAJehlsUwIcxRmVw6vcp(1Q-j$@B$2Q8hAE<{X(v_` zN{7Q@Xl>uX?O`k1g*Jgf?#Y~C+k#OQsF?DX5j{*rvVu#y10FI`aZg%W%Fq}V{`j6* z%mv||3gyH&X+I_z-|WVNCD?*945)iu zkjL0Hmh!7wv(N)Z&x#`AG1yj)c90N8wrMLsB#&z3Sc?gIK08LoTk49DGyp%OX_W)X$PJVqAwE8e!#bY?>6ut;oH{O{T6hEjsvE|L-4Aa8Hba^9fX7O zwoV8GLJC3CZ1Jg*5iOVvgwp-#Lz_Hv28&Vc-K*B?q*Zysud)NFv9|8ZsMGp^?RGu! z^f4;26Xphl#pTA;84#|CNpMsURz6f;;etCZJ2zsuH|AoG_1o`u0eTNI;}L4oCWWXp ziq~xnnpsGGGi~N>28dWMu6u6CHVi1_tfTV4aMq@P!rY7n!d9KpQ)jd{z_5 zolq)izE%`1~Y8c|kD7(Au54x|fSF(d4htW+;@Vhx(H< z{^a~pBX6gN3^k?}kB8c#p_zYdc8tYxVz(KJ04erIfLYfTO-1NN*wzd%o6ufqBP9nK2wCEIw1sx+~O!nOlvyfffZ|mc{uLEE%#Q17W)8RM#CIv3%3hX zc{%$?f8=JQDy=OwH6}7VlrqLV9zNkv@%EEdlm!qyOKQT{OsvDfBGp-hYfH?MIN&P+k8y*(arSeX!&4Sho&-uef1qm# zJkFV>?!>7IeIjAl!*!6zh+uqmQj~79d2|&X9M9Dj_~UW&Rfgj?0t>!;#AN2}5I_*f z-FdBB$HXScB-@K@e-A}Ove0Z;biY3 zOpY#kBpE~=n*EvKz`phQsEl?Omu2;|L$W#FxXXw9vX3b1^?typsk2Fzgke{t`zbfy z=UdQjYb73wm4y0j$cZIFBm42);17Tjd+d5V^^r2w^=CXaQ?AsT_;dA6Qi`8T>-H8*@lhBa3%jglK|Fr9K%m!dw1bmmz#d;#BKoGlbbbeUDQwL;g(n zQ8^}bF-;*%W76o$@4zYo1`5gOCsigHAg&>KrxQS%BPaC-*rr^MkX$2HB9wl*)kY~2 zYF@V(+HH7`A$%o7_T{6kGQ019)chA)?-W?a_x*v!wv)!TbK*2M8e1o3V>h;K+qP}n zjomhCo2E%~C*R-wKi&Iy9&$2g&+I*GeYC)cEyxm~7?Hz_(zc;H2Va}eKDf~^+egX-K3-Q$ZJ zA(aBeCMVp9q;>TNXGMEZ@n9I1Z#c=9ntEwPiCN|v$yN6a!`}^W&kch=g14hew0@5q z_s2{MK_r`rk|6pm1fovOPaVo?9Ha1^DiSW`k-`BFWFWI@Mm}I{{Vv-*b5J^oL(r5J zR(CLTLs#w)*APU`sY!+B>eqn`C>xhF(F{eIMY)^WuEl{JXo_n_OO$-_L)VhsWvJ_id- zab~lU2Hm~^^rpHZzlK&L#sRSW$ zQeJ?mMQv5*aQ|8RT}dNFK+n zg^Cv;OtznDK$^)vlu*7a=4T!45=zvC!xp0_eHox^3!NoRu0y?NL=`6e$nM8Y??h(C2rl*u&Sr&c0Myi%Sg1iNl&}8gV#q$N{ zR$koIT(i&C+mJ?|N%#ehD=y`vOB9@mm8`X~_N_?apUaSa;r>f#EONQ;VSVwr=C9az zAB87xCqLTp>x@7$l`&b6;gz04MC0YupZj`Dz!ux7XDQ;@F$`^KD> zdP>ie&@buk(z;%Q>gS&bGYtW{Io(pvn~ac4BdBcHM~Xiq!9sxp3iEDDY+WZ<39!}~ zSV5FgNgB!ccVnBPLlDJ)Uevz(CfE3p@x^3PkH3$4!L9%hFE(_y&iIRE>h|eCH8Zcc zinCSuSl1FLPgf>{+DACUwa@lsKy8MQK<{WIuX0|f39GuXGHK%oo#%PxN5{xPQ&*I_ zfqfh;5tZwef~so4=E{n2>PXxT7$wyAs#p(WUT(P{*!Uy(c$siOyL_f5ADSWe#Bp$(e)~qITb1Oqxb5 z)QPjz(3vbg_k07P*X-d1nP%96Ve32{F5D?xFx&kCjSXEG`v5Gy5{?mwrwt%fO#4&Y8~UBdjAp zutuS5mcIM>&?I_R=NH4Cy%zb4QW}_LK+$H=NyW6WRR9@s2#%W|n$<#Lqo;fbsRGUc z;V)D>pJF2>Woz$kY<`BN+=@K48Ix2n{ZCj0l^qz;eA#%yls>B1gbIxcEi0$!S*^7$ zo&%hbSvH}fR!$XiOlI%F4_Wq>=JN({*07uGx50MKl0QgnSoFDqwBn3#{*EOdEn>BE zY!-TdF$-vkf=4kJujQXGHB!Fv?^~yJ*$snK&@zP%4uiAAY{77)HzgOEY+NL$kbqtv zpNI6p9?1+7uzT?GCUTAsX>0h81`q}cs}A@?AD#C#`v}$CT8@cjCzt7lB;2K!&Zif%;@E>F{QtJ=JQ{FUrq29};b3xGi)LEk`Q{ z3!bQl2teVClXWcRDr`SK`_s`_yN`)$?BOz(WoAPyBk-!Q=SR@VS)BpnlI z&O~SnN5jB;ol?)ADjy!XLQbe}G?p4ywx1{lcFmL#m6CNm7?NG)nU%bZX{FtXC866( z&4=rF{B)!-6;7BCpIjlzwVsh>q9gYjhT%(&I3BA%r!Y7O`l)lI6%>@*M?a%5ipO7m zyC@qY@`7k6M4e&Hra)_k@y)2v1oSN=q?FGZSJ~+IUmil~;S?SDym|gpRgG{8yf|16 zFECTb;3i=u@;MH=P7b|%U=x>-TNVY#D?0Af>FhF)oLL-8Y8{RhA6{Jyk;C3ASdDW1 z*VgaA5U?o}gd3j%Hy-XQv6GuL*K$D}IqG zC4Jb&l+R4MW*L!MjilzDcuj|IEKrMOY=VGzo*|McnjAX?1;?I#D-UdittNR6HPXxo zYo7EeU2!Y0EzsS!o`6Bw+H-RoOb{acX|5`cv(N)=VF9M1Xeb&;2oV=GhD)oFvM{~g167#A;TVKzY?c(oH@~iEv)0&^fni;UCKJ$SZZGxGc(kTXnnbq#>oX2y%3i-`JQ=QeuB5y z5=Db>caj5q4Mjk>cj++KRt6zBXA%<~!z{swnO~8wP0w21y=ak0{LeQlA4+v5RV%7{ zB1S|+!(1gBYU+baq`iC6eA+k3IBW_H-$Q=5#N3v0O{wyuW^6j4u`@s~DOIXr#;zmT z69h-|MLkRxr{4r*htz^CnL{L)GtNB`#5+#M4!1cH=JJooP7-x9qxV;E#Ei3nO3V+4 z97beUno@kCrlHhvP>)R%gIfeEt*J?4Q2(WZD=|o+k~OZFR)s|SKf?1;3b2Od8NJl} zdZ#{NA-6JjnXspRCDUHDHIwTWN{pJP+}EWxBsZRzPRB;7LDCK|C&%P+CKBm|DSgo5 zNfpT0%KB*HdcV=zfPgJK9BS+n!Dv)pHP;!%`5>I_fvip&l_h%RG}xl_bt;rm*_7(` zaj_jVL};whPPkcH#OM2Wpbp%k#kFXkxNl0yKARE;_hXe`d{X7_Rn_s*fjT*)7W<*# zwE4~yb)vvZP*pRMb)=rE;do&Pbq!dDa-fUKVk#9iwK>)lI&llTkE3A?F1L5_)(Hea zHfde5$k4>1lEJdB&lxi&u60TiD6^Xotc;IM*c@Yx?^V&SSO!L{BmC&~5iRupofnUk z>=)vozI4rR@25v6rnRMiDBmiamBR5lj~udvs$5kTOA@Ct&0tdTQYs#NnVlN9Fx(%l zm$aC3Pno8vljM6a6C$Bcl{5y&r%>yzXw`rCf2594$k;RSi9>PKO%b|_pzQdpc+vpI zsE?`Tu3l6ifM1ZuH)#-?0C6%DW>hbSaro`8k~0{enFCmM#1d8y60kcQhuMb0V6haG-MwQ+Jdu!<3LeDHAS^^2`s7tac-Af=C5@uOIm+4V0fGQZVI zAd{4fx`#;ExO?FIUlu?u*r|;>=f@kzPL%P_&bMOBM8_;;51Da*{^^$6 z!2puh`NX(N%lA(K;iC5y=Amj5a$OoIGMrr>%}DB5Kb<&vD^jE&=a7zREQ_JI z44Ip$b(ao|9xb?hbXQ;lA(tpm^O0G6cE+>5IhSHTp_R{_+fecjP2M7OQT)dyIyR2(=ayfz9k~^VdAnKLYiDBbx%-Itk(5Z4TRmm( z!lfcAJ(r~G;}L^JO-p3ChFvb04xc$UyD*qkDisdHEa?YU(vEp#G;1n5UlQ_UIdS#5 z18P2V!E-)pJ3-(W2!z1pi48{42&pg`35m>_%a-tu!h13NOX<}ZMX245(~8DX0)q`* z{Am)N)(<5UXOZ0sG;UHJU+Kd0eL3c<(aH4jcSQbB8Cf88?R|63ikUEqe6ot~SWA9? zvH6RUhZ!W(5$+ttu(4LjUyZ}M@pptO2A{GsbVgv^DXUWWa)EN7BuE<|+T~p=)1*@Y zloylem|}Fhc#=S3bpA!e3AVelJcCBzk)LTI%=G>ea8K>6RAl{Fdeoy=rhZ0*OClG- z1E};N5V05wiCqwi(`kg}dNG1{$Zmj~E|^Wtr;3+-LI#<~Y_=UP!LsgFAu+;1?XQLb z>5G8fE@HIsI7>ZRic7O7oC=CU5O9zwTy%xW**a9RzAR=stXO~@a>UYi z)P#V~F`>m45$CA6UkHtpVt&BFI?}Qm-Yj;1vVxUuSr_M}dTg~zGWh#`-czzmU}S!7 z?j{V@<|%pCy%52icwjNu6r|b#JaNJEYK4}Ow?$PBOgtnfqk1;X3`McMViXh0dp6}8 zKu{4>K#go`Y8$%0z5dR)n%YmswsPR)x3dwWFa?~5Pn#1W?NhWlvfK(|AO5Hx-wSTI z0Kq!-fX2iAjD+op>GWqQQV6i;kGQBB{XEAZQBnhuLdBz5KwQp&bi_40gg%DuYRt#aN}BiRVM(ux*`m4D}YF* zD0tR|CGPeAE%;aBMZYjuYl?L{Xc2T$U*2$&KT zl)X?8HWt_%3cmb0JO%{BFQeEhgD_y`jSA|N95^+464zM{M{Jn!KVv;MP3m3i)#ox0Y-_Tc~+3y-N~ z0?Y~E0c;28vr09DA|lv4=@)g3XgBYdN|fVQ0`#cYw z{{(}srw8ko#jwMNIMHy2-n;R75NXR-w-u!*NHaSw&;d?JA8aF#)Jftzib=2R@RKnI zxA^(=C1d>9&gQkIZw_H)=*8H$(+nQDg7VgLArG)7(L7}=!c11|Lb3`4> z>4qL#Lh*QUfYyQpTIXMl@vFQ}`<0YuJpGMfy%c7s>SRTC+5p!74zgBqphvL@>ie%Q zl@=|e?o{K|ZHvI-?4rpqP@`l;Oxd>eGS{w#vPq9s3a^MvgD){hM6y&-=zlq7^K}ZM zk5l(8qzhV^IiFOe_8sm8)WX}$;y_c$GHIxiDJ zX)hyW0;O<6i1eJi%q}+sEeGtAgQ zh#!-`#S^qFu2)7l#7+1Vr~0)>UXHmk{L=fWjFY4pjhD25Y&8mKOD>uypmjFOz+I>5 z`T&wPJFX7xo;qC@6Ci$}5Gaq?(G2>V60!DLeS@P;~sW_%j`la$~GRG=64iUZBz+68Sil)c0*NcW{Cd(R|Fv}WUP`xct8{~b+KICG zKrHE}?Fr@81H1f@b)7_%r6l3`l4LSh^kmXoQHBwxcH{-ufE@2>9*|W^G4p`OWW^uJ zstPB58yXu(8<4TfsGPK@K9Y%-@<`%P;!)u zRCKlr8Ev#=G$pp#lwI$*$`dV)R`^K+BB*jYs`_-9ynt-q6m85rDzsI`X4Z^3Z;U5y zgYbA*PQFSGHz?WfoTyTF%}&D+u3>Fhuf#~=7}&hw*kMC#nzoW+d;`) znxN~#JjB)!Ll3sFr@tU7){N967+;c6aB+(UyYxcuxApS z52-G%Kpk}U#n_FOx6(%_0!(0aZ2+3Da zA>Wkc24Q_di*+fTFGsJS}0Ty3GbIdU-aUgyC$dwz4kWL*4)&)vB}yhA37RY)5m zZf8?bX3|I~Bxwy)&lgi@cRfkWVhduyn6|kUJc+2j&`(QgZDr{CNxBzs5F5jZ{gNeH zrV{(jlBTgxaeUxe8NK*p)vgcl1qRt6T%B%2j9Px0WGjLptZDk3%%lthZ-t+CZ8t z1BF`?o?8@Zc{7*w?*^!bxA9W7V7k;yNRA4{o7s(6xfq@P!jy2-%t-FxAf5P?$d#hK zqG!O~PU4UCoL1ZjPY0&}WfuAD7sC@lZhY9mJhiK%Ja6;KBjc zT7O^ln#oLqlR2S3d^6YEpB4(VMhw-pup@W@d?n+(qajN_brJ|n3R4r&!XgYLfq z;rBI)V5eyPugmYRH8QUQ$l_b*#SJkIbg9i-F4Nil(|vFZOJMF~M+M(l9jfqQC^kX) z2~&4Abhlr^Y6A35wtbUkZ#a&Mx$f(_FBWva4_9mlO(RIo7)_%?Hf!DI0AKF|1$TsJ&;a3u6obkiWSRj;w6a{H zx?wAFj=RCe(TcCjuGP=QE*svqF^0Zj5gZ92vA+w(3G@BE%|6$Nh3*np!STBo&88^o$BLqsA zZ~5C*-LIA^LZ}t1-o)}1CUP8D%!D5SvF>@U>pT!bZ9N!hk3=ShANrGl4lopAPsgeY zZ6f@22i$%JKuGkjZ5HS-=Ne?-q2PsUt8_u;Jl?@f))hx9o5npnnes#VE~CHChLURg zb_-1Zt*8bzMuO<)&cpB&Wr3(rEkQE^OP-3W`G9@Oh-E0f1xB|P3eW^S8iftN;;ME< zo(EB=Z+eD=55pHA^+ogw1R0`~G%J?>4;a(o4X6`U5l(EVtuQ8vo&kb&=|6q~K*hEo>GU`B*UG8HN7$M+(g zoE*DTia#cXT;}OizqE-Nzr14CBC~c57S4dI`doJzEL~rJ;VOPAw{z`@t$qBjn*yph zBIx(dG*&|Neze>?QH_2HjfxAM0B=JKo#7Vd#~0Wg1=ItJ)Z6D5&;=?_{~k*pRFkq7#+~ zWNPmfRnm@tP?(mv13+2%=Y97>hHd>>{&|B>CzO}zNctz}LRP&7+20Bq7ty5}d#A7* zR)jgI@h#6x{`KY!9aY*U*%G^W7spWEIKtjcZSYc?aej~J>3L%y|ART$C@BDtwdtkN ztlMtU;;`JAU`E_3Fw%r{;pWYj|f*V@PbIQ=2L3PvIL(zX} z!XIGD`K@0A`hTtru|nZB`W5&0z6*$OB!x`FD)Y7`fBQ9DVKs2is9Wx4TY2l#6P9c) zV3P%<)19CU#hPIB&UKB|N$%gq|QqHF3_yR2|B7M@1GSm1yu<*w-#sG8sXv9A5*b^z=4SCQyxoqD>~Q z%SAN3N|!Lt&q~yaxzo;wZjf~)nBd<(y7qNTEl5RVRSmoJW8#`)kF5uRg-Xr#bbQZA zOoAvvv#ZBnhBkERy)3y5aE$?z>+yIL_UbAwN>zjEy;R#wnScR;@Vg>hHL zX$m)DTwLKp&Hl9Ue|Lr!H|Y7sQ|+*B_`TD*BPX|ORcx=z@SPf@o|Y@cKuxy3;t{%^ zG0(LP=Vbb!*q{H5Es+mCMNFX41;Wm+@Zzp&F22yo<+L#hf0;r!4dX8zqY`VE%QNe} z;m0jxNYw$qD5Es`aq1;f@Vf_#wbRC%BV$n1dJ?>-GxBMKILkvO)(YqcV(WWF`;A(T z{6Gu)23?|u6KGk|It4n)B$gkM#;;V>$?1z6)Cd`~v_!7HQSb{Pl{aD@VfN{x9HQVvI>IMP^n9V_+9pB6*e|i;M9T|r7sS1Mv-+W2 z;!sunaCrXlo`s4?3KlX9c{!w*oIV;Um^`wQ9*L1LZ!6f#A3t^{1X+I7_brJ0&2L%Q`d$CCXui73@^7uCrmNOk{fK5y z?S)hA2C+`)*&ScusIt~%DNH7=n%WlXd}J|RRTn{+Ibwa4GBI2vo)c$U$t`AH2o9Ii zC~~OcV>VJX{|bgR0yoUF${lJ(xTr)SMPoF=BoU1oC*7YEMTK+*8JQC-aI=-fbxZR{ z;c_l>Z4jG+LNP7$N^c7tdJH*S9(CbkX^If8A((EDP9bWqSRL#srzA4xa8^3=Sh=2z zJP!Hu12*s?7CYn?Nc@DRo_;bo%`sA!3Z1xfE2hXy$SIa>B2^T{!IO|J_TYxerA%+i zg@j~(>^$GQlCq2+2pyTGZ6=^Ao@KlemtrXLX$+nPN4fZ#j@#6Q@9Y!jsvQl5lcre= zg@d0#pU+WL8N$fw2l{8BONa2(jBvgw%Ino&UKv!~YECh)*lk=`27AdP3RHVSsf7nc zoZLwIzBs;G6x9!r2%HsCb!cAR!7EJ|=Rc!h2+O&m`TnN!#~B9Qu0zG7^y!y1u(w5) zVue#LhE}P^@sai@5|j}`BVi}yWz@kAt_l5En^IjtMk_XHxe_fU!#JZ^KC^mX&p+gg z`LB*5u6vXY2(yGxYV^9C6*SYfh+V0>9I29OTbeHNU9=#5hPEQ7V(KE9*_D>kA{za| zQv1a7G6HhepqTlitVW4MKI|!yv>uqS5p0ruS{`#%8S$yg_@bL~E)3?QYF=PT2+alP=mT8OhlCBT;c6pJ1G`1bfjUZ zX>1hM7>-qrz&ZA@;e@4N=zWisp|bGC@~1J)e;TpH3V}SU-y}=DN|$KMW*UF})l$xG zmRTtXT4sDr-Mzt_mLI|g>dj>i9Zy8rs z>CeN(#$VGJ_n+SXPx#=rq>1##Y8#>7L-qTQZb8t9qxf-D>L47Oy=%or?{mn`Z-=ve z7BzPw4+`g+*Jl1uVjDL}OF55|FL0R=vGolLpKl#M&eGhn&EOB0J6+9u`oc@c+C9?s zssJ4hvqX9y-Zl2L^4z$!d04IHq&m~?C`?Ecwrt?C7>8hvPA2prGjrHaxg_W0;;d-WPfe4sFML*d;%e-1!DMMT zTkSJJ{B${c@3?&*&ea976P?sW)WuDi0nPvKhOA>)ul>-=Njg&jw%)B3OtOt!?-l0v zf(hCSImIpNFfptN9rj@+;ltx;LiMlfpd7VWLy=O|&A?OycG?Nrp&+i3xlh2q z1nveujNZ>4mrc;z`<+TfaKrvieUz;1P<54@|BS{$bG`Bu77gWG<&!?$;Ias3l=Ka* zO(Tr!EJE$`F`N$sN zQjTO2=ZSjXKQUFjTa=r6H-nkVk2S&*oPI8jy3((9@|xxQ&JLys?Y^-M3?Fm2fQu28 zFp+z|M{-U)XYU{I5M)`I?vj101T>V|71~$~)%l+mIDSTkZ!<2ZS39`)XB~+)^BL2? zeD>AH56Y%`;6yBi;6Sd1fAO#B@LLQarQT4Kt_J~m##v|r%A@*wS6XzrPQx)qkHz}? z_vs$50oi@O_y4*#`0jPPUsG4PJPVD+Fk$HuE%W+J2=3nzE_jXGROCtHib*JzTt`^l zv+(W%d5@${gOJ>}EWt<>)rWu5eJp#vwF?4$i}8JO#tTp9r;wyg&MJDk4{aPvQaCbQO4#X_LjiyS7C z=QtUd!IJBJ;PX&+$?l|_a=wcRKBLKN3ng%ySP?K3! z`CY(-t_{C0mkP$f+^W1_^r3xeXr_15p==1hzchenWT=nfF$f8?>uDoOpe4SKHsq+g z`1gSwHS+ln*_h#3&_4(*4iZ*&4GY}f{DmvtCS1(JWN3Z!r$@-P59(z)=ZoXYe^Y1Z~CK?9bz3Hl)0 zO(<19uCZrW4HFqet#uuNvaC0iFIVbhzdI-RA>hxNyq3GhSR9>!*;uc@5Wqe&JUZc4 z)c-rg8oLpg|L9I%b??-^KbtwwP;!}!IjubayCVF%YXOKF3t&=Rq1KchDQ7cGKk=4t z2jYz0BZ>9_>s?`Ju)4tC1Ne_m=mAe7?_?Pb2B}*3s?-u2>}nVAw*vrvAopWc$ULP5 zJhICYM+;N17lJO5+4UYA3fM%=HHFSN$1Sr) zDi?J661x~gnpJ*}gkiV=$*0HWG~p;?)*cFS9oVO})~R-Z+W^Lr|1?TB?E_>@{+Mo$ zTKP;LUI_1ZC;o)n;5BCW-31@%j|46~FguFg<0J5+_5d&Ql1=LeD9Y)g2-0VOVDtb$ zomeg;^rlK)ouwgT5g9}zfv?E08aH}V5`zBcS$UHP*OA@1X9!S|R6UZFN^v-UQTxX)wxzA(<=CXZazW*KlO zxls%4S2gO&{8ACJnegr52&68P6zj=HyY!njohy0Y-GIq$`o8Z8r&IUTTCvT;S+jg5 z{I%^1)q_*|GP~EtZ$6&!@8qdEj|oRlEYt&lm08Lbkf*|2gvWp5KhHf5q?ySzED_*Q zw{>^_0ZjM9iRO`@``t3g|ZdRomW}BHEe=Ekf_wu ztuX;*iC^Ks2j7mzt|M-)yaP1x1;1J=fW7^rI0Zd7fg(oI4F8N#`aB&-k9{ZH#E5ZM z1MBt;wX89yn8Hrj=~UR+aWps*v2eQCgHj+kBoUDo4x`WrdkpDNDHwmc(p|hfEL1pb zTGtLv=ye0BsB>;gpV(ZfK>r<%YvP5P12tiOEjvJ+9xF-SrQOb{ahYKte@IP|;Wfn| zN2FBDV5V+1xk)`2e$t#Q)FXGi%8$fVHo5oRgBK;T*_(kuIf@(=XL{K@sG2XA8@tHB zHb&Nam*CRwD?|i;Ap}76*8TZmiWxnf=q@cc>KYDs^kW#W{tL zw+%zjvRqZI+Y@C-O~1me*heBMZZXNqqmTYU7-IOr4jcOIt)dRH1%(c!x7S>Cmagy< zc~C}TQU0?dbK+7bbgqvVf@Ym|5ayI_#gVDXVm6N@u2|%bvxHKkg)OCcbPljsJ)>3C z{{q_!gJfZdm}#$$&eeA5tpYI~<#89dDs&^BmFq}v$O2yOpmS-4SuJIhlsR?LEowIP zDuCj{8ljdIs_xm)@FnNxdqD&bnFU}FSy|>d5jgUovoJNc*VySzluLy!`Isf@08V?0 zicoZ^321JiDIZOo)Rg2Tg{9jFW!#lC54vv>u|=Fs)n<8puJDHq(S*+k)_C#-5 z*vN}y9nvkjt_=ZS)`=NlzjCn=#jih>Ot1!_N1LP}jSJmc7D=NfwuC_w2|i~8Df<{y zjmbMHd`f);rtYA9Yzq($2KgtEq6 zmRPuQM5gKPMp8&^uc%|A=*kB#Zcw)S=~;joi;l+p0v4&{g6};$dGIPV%8|}hY25`p ziD8-FVpp!QO24|C?$q7vb(73vUPG48O82{)a?Og@A1aRYx$WsQQge(h0=L~`CO+l! zBYcqc$mfkpUw9%g_j}`maf(Ka#Pg+Dt3RA2^_DBuBHd7;tYSuM^p0d3f@!msm54iN z^A{4W6Fnw3kFGekdAzKYiZrIGm%>2gq@{*6;j);O>;YafXTC|pULzq?s&uaMQQvIb z=6?;k7IKi*YdzlmoFhK)<(mpNlJ^`N<-G7 z5hYFtCreQWa9Q&|rq778@%o z{rkLmfEPpXNYlSXavGwlmPQ7z)|_o%5G#F*tDJYyF%}g}ky+6`iUA@)(bM!q9vR2|h>ecdX=+v)K&BBE`Ebf)QzVhM>valFm(C zFWWiIbO6sL8dfg`6{_F?OL{n}<%NZUiK9zUIxH%Yd#P6|0>t(qVHaUfp~lZb3ABY< zTjr^@yMZe;D8dJNoArN(M6;00)>2n?Mu01QTTDdPg&Cxkxbu-1n}y_H6(mZX&~$1I zzuRmU%aF~tgJa!ug;d-eU4~=|#;g(KpGFcZG2Ndqz+2U50hE)b;kKrWA;mO)FVCr7 zHTyRaR@|5DKX2#3QqBn+AGJxStk1hviF zQL8{csjghw>KsP+EwvdJ;{w%1L}aodVWy$P)Wgfa#f0Zd`Dn$WUF8BehU^I$nAiE0 z550uVIWTA=d@?ezOL8HWxn>%dlgMSI97JtVu&zREn6r^lcA;VZ&ACt71J7$f93a(4 zkR|056TEduU_l+HGr)NLHGBk8C`3Q9uv3Ihdz6URUg zEiz@~^F@hc<1mzrzB*|?^c-EF2RaY7$WR>LVM~5Kze+q+F)smqsn(G!lKj+^LqX|* zYB6*>Z*c!O{j`sDQ)JYEmn{<#Wj_IPg7!CrxX(D*p6%?`wl1)5O%tRFAY>}L zsL!gw5iDP{pwLif?_fC8TOhuMS4}0dRN0zNvMS0Vq9)*p{VRMq;~n4!WH5d{bA|-d zS=dljICKKBF0+DbBFn!fWki<2?Ku%0@_8M-rz|3EYd@)Jo5!k-;VbD1{F)0I$6ZH9 z?dV-8jEKZntvDRbA{P(~3ekk#B}>mfErJYHN8p+s8Vm6+?0-mqcft!`3S6RR(mx`n z4_{1tk}x@><3)O~76IgEU!TT#yOGU!C5k*d{97D^$p8n}D$<+_)`mrf;jdpE1BR#X z$<{FR?be^pbStA)nWp=>tb|t+D*udIH& z^KU%X;(tgWN(@;K3C@1JPfFmM^@VLgtE2fE{cwFf>$U5VDX1}U8^*>*=pwQDa#(>) z1xM2UXnIZ8@9il}%e?U(&jTpJeEojfroM4^!ah{W2uJf(DsjB>9QW0O{&%ot3Io+f zp77b+`^QEED6D4clEV}Yxvu=5wN~#r?x|4nOX4b5k|LyeX@TFHlzYf%!8PTdXqTh`f*#iW;6e{^j+K=axd;Ec?F`%6O zB6?tJLfFm}=fAVFAO5c+;BUv39LdjLKsxBZ_t>BQ9p-bb2ZJ7pXq_hB6_00PoOgM}+e}{aro1v%xgO<>8cU%?}71`gK9T^;&x^s=rDz zUuj5IsP<7|B2GNq>k87o)uQ~-+isOLd^TxmHl8IMMoNx9ih`xOQ*}KmFqi_#T@m#{ zg9IvwT{9ArE<0_AcCxoJ|@>}9)c@=B@R&l&1PP=5a`g!M@jp{9v`QNbNb2D2W+P3kCBe zP`T><;j@JC>y`G}4!)5Wg>LJ8bzQ5fxkXiG7j#C3sn`mM zTtt+EHiCN)B-|Vp%~*knQ8=8y5_ZQ=t^u!4H?{&t;!9RF>4ZqwVIzUtQl5X{`m8-h z`8Cq|0gwX(ZeYH?&_Y8YITL{+T`^j8k{S84_(_W51$JYDmzH?=J3b#s%oH*cUvj>u zd}* z=zciUJXyQikcCJ|JNp1Bq3f$2R#fyYQ51ht0*p9m>p3Lu^WXbxQhrq77$4Gnc7+(> zT*4B+)lX-XB`!Q~H`+wJ)7<)M+(YvB+{UMARLPpqK=hL_MdnR`*B@!>x|)>|arfrj zX073;G4|;{YsJ2_>4Xj5_=aIsDYr!wY#I5}>eyBH%S^RlD)=cu{@j%XMbwJ=)(WmY z7SPxBW8C^{6UP)q8FV2D;)Z`n#Dpvt+qWAH)sLgdbVYjPOpD>Yx_cg-aM54)AAiX< z_7qzRiB+Q|xOyiO99s$9JfrYBnLa#i#O2ot!&1zD5A~3$X4A1Zf7jXl9ACW#tyDkV z)zy+KI_MH!m`|ErVXFvnJTz1HWUkl>+@FD%bS2PwO0APMh0taX@apL z?Qy-nZ?GGBzzckMA$k!UY-j8t-4}4dAA=9M#BO^_J3Yl0MU>vVF4z2R9yyoIoQgC^ z%vbL$76qqHJU{HwseYuqkzwE>7vFg|VfY9?rC1h_I#)wa=48L>m?`Y>>ucsjP7Om% z6{=Ke3@Nyzb|L?J`Y4d@b0mtt$TTyajg!f&-AIa?4$5BPL=MWMki;&cv1-;|!I?NZbm_uzsC8A0Svpng*AU)42-LEc zlz>dzi)Cd)Q}3LvA$E46>aYAr(4CKMT)i|taD$MU8Z;hi!&3nr!`_p{L%C^>$+A(^ z7oA@I&px`(lX-yCXdRFa=q2c zV2IMw$gHqB=gnmz{ajyLzxWKq$Rz_wL;hn^C(B%G-4 ze_bjj4Bi{!S0d-ut|g$ovCl`SRhjyinjR-!=oL+o8dmK6`W7f9LpXmE6R`4K!y|`X z3{y0n=}M|e^K1-W-9)=1j!FQ%6KyF(y1EObAJ1%KH8K<7uweiToz=9ro z!wyqr6H8pXI%5FSDqQr8=4^dod6<@NglKx!l9$Wc&eOEX18)#VL}mvkSgG{CE1QJ~ zvQXx;JaFdv@4~bYWJ_0-uZ;8T>~|hwI!45dzN`{eE@exGm6*Yr(-cM zlx60${VvdoHc@Hp+ZEi~v`aHOHRj?5{<{)pDk$C=g>===>TCaXfa#b)r4}VP^(s4f zAWrwX%pM-syI8g6t|jEMICMXVos%Bg;Ir=%rK=+%Ch0|xRdJpr+@&Y=X@iMm&CU8K zT;Fd;)nY#yZf@wA4t8{wrtlM?~t zJ1mR;zo3E%CX5DW_OOj#)#6-N%Ac(>Y&N-n(VdZfb0&PBG?aVqFg+{Er|Gn;+doA2 zdrNABS<*Xd1RU?tqcTggM90fM4Ap@%o$1Oxc&=v{x|GjCo#86ZX zwoF#dHjcpTJg6zeQslS418@3U5|r}trut>u&oE~0AcOlkI-8CF$EqJdZV2o-toIvb zCl=>rcEBgZUxg%I*a*2XEDN}Ik=^BiwP$ZhM+u#1ngq!OrphywZ!S}<#F!zJs2nT8 zY@N=Vu^|df1$I2qrw#wE8iE7gf^w6YA)>bfAZXE8^O8xv2O{SrOsFY@IEbLd5}ZNQ zEWpoRTOGWu1hMD?4OxSubBsPY*kn>Hl+v;1i@M&i*~1UD{HB_UkXMDnmmgYQ>cg+7cy$ni*Lr*ijcA-Ji@r?;80Msw;eb zenR=3cT}_Fous2Te?JQnL^+|&JbRBxfb4he49cnsMezwK4Ly@8!mdxqO2cJ^-T1-KzCmr~dQ;|su6Q0$T3jVM7=5GE?j`eR2hAk@fH zAn)oFBqA1IO5qsIO&aMq+1-tMTuLIbzkAVojY|1DZ?G7o|4S?+d_k186m#dx%B-kgaFXN3YB+d0Ty&^eEO3BPKzwyUGKA4`Q*Y)e!t>lzdXO*JW4 z)~w=in1OvvIr!~gi%CYTZ+o&xrVR7|&cXu;EX69A5`4~T_*c---wl0+j1Bgnm=rcY z=tz5SRm1L%+Cjb?@=c3XgUd}p1OyVfSVxI)LSs-9i&BN3P+2mpok7W zVP_V`h3mhuOxhU1AO$-Ax&+L%*8sWfO)yjtmMR&Sd4ItLr^yuAS0}73z?8W>J8m++ zx@ZBsR~WhQUL8Rec^nB78<B`ae%h7}TlUm8|UQ20mcR^`+Q z;e(aHvdc*g;FL`&>0pUUtLU%jzb8wN8DCo0HJd?pCLA3=H>8B8Vbo7VPl;C`rRV0? z(S%CpTM$K@<8Z`1-R1pCThYzT92CK-9E%q#8uHQ&*nrdCF_sgca_oa*p~0Y`?yJVS zz9xqvaGk^pei+NhI3Xc&85>I*=&jKg5YqZ)9vxU<}vYNv&-->lkcM$eW&W2u9ghsz8(_|gdJ ziwupU_AW|Kwb0((t(>O;pe%ogT@A_ zmJ%$%e;myoh#G$Jd)k#Ga0<35+9b6}pF1JucEXkdbc;WqSc-N>z}hyDb-0&2budwEdiOPtrFNdoEq z&lKgrTs0ZgjwBq`h`&II3iwMdEaikT@y59}%g+gu9K<mNN*6?5fBj)vGN6pNo3u zj^1B@H@9_{^EX&Xh-`iX10lZYs4?Z*e@M9wXg>KXe<`|X7>bFbNd%%&niE&#{m=AhQS87-@6)Hz- zAg}YF`7HquW_Z|1qr>@fbz*h<+o0mR2fy^9?=m-K)0W5Yx3a?6K%>zSbPheY{-;HYVFuOWLz6s+aOkqTk;zv3FF>lW;1nnUDl}_K%vh_OgOv;WRXu4 zGQH+G9_lai9ydL1@vnKTGz5*_*|uMvdZBr%K~s*yRS8ZpP0oQAPDP0K+ex;WO+GRQ z<(aQw3U3N8YdL`7{pk>wMx8{nO~;EK)1ukkN4IB@D?7%F#HHuG_Ck);;!3N7<;N`) z!y`|tTaEMiwW|oYe^=#mv{KXC&;R#4^W!}PsJ>;`)R5G`TJ!zzT)ul3WU?d#Rvydc zhGpF(2Xas2U4&G254=Bbl?ye4vh>;1vk*4j&%TENv8h@^Y>HYrQAwE8`pM? zrdGZ`{pjfV_tk?`jSvTS9JxMC|gb}>j;H7yj;})aG=6I3N*2v*l|d0-CS_$t>7hkD3gppiZcwk zhz)g_0QQWsn8mqfEjf!uxvFDbkJ5%!k>s{_!G%iz09Hq)hQy`0HB|Jv**zdrjR0=` zQ5!C*f&@X=${scE(PFy{{TqN;Ud+hP1(!qS3TqSpL=D4jilEg-D*U_#x(oBXlXhS# z?AMX9tJJ21bF_PPL4^EKMJr1j;uwWOPY!+5Vj(kKx$pO-_FlwDTsPK8lK}8{H_aq= z43YANm#!i>o0Jit@hE8p+=JjE=^VwvHe9Lj-=)vEQi%&L{Q>R)#~EO|pYAVrSWH10 z7?ri;=I7&lbBLMaWWSF-RBo;9IOjZ2z`)dy9S@)y)7f`0DZxRnws+XBHH+Aqte}y? zBT^~GjWj+o3kh}Vucbdsg?%v z2Da;H$}hL=(lVyI7rW{!itv5&haP*n*zSy$!D##P$GI{9d#8m}q%Q->mX_i#SrOAb z_taOU|6>6_zT&9$%y><^@G<{nkqRtNo)Oo`r|Th+g~&Y4jYJg1Yq33sVaF^<5JDhH zt0AL^nuXKT3O~0IK|qpgSPQ|Yi{*}!sGj_!e&1NEHfcNhXTlKaK;cl&@P08j(2p$c z7T64l!R*HsZp+XW$*F3uJ_@|j@lO-d7}7;ILd^WSeU>Q`_qz+fANLr}LM?l8uHx59 zG@w{v2Mm4LpLDRDx|7ET)uc}iw*9X>50__aWEDc8z#joKo z!Hy^Wf>8-RQO5N3vI{$*%Z5MgDz}l~yh$A@x7OAcjGfmailhAc&{33=H%Jx+{TFXz znq6n=tfx1m1BHE9**vK~o(Y=4>c!b1g>}>L-sXQ{V79!#ZYsf%5F&Gq-A0-(vGO89 zpvr%#f+;af@fwhc(r+jV<<|?spy~cf2sTNu7AaM(pebUeZx_9&QA0&nfQkIqxwxvX3jdEE~LV0inDTE;0yU@6+c>s?pGu(02mxV#D zhKm#*BC~jzKm~9+&nM|%w~h!pfuzOgul)y{g#ld0C}x=jKW9*aBB;4(*t1|kjU(XK zwd)HvIWy&8($TUBF4$GS$`r{S3%N&#PBKD=XEJqkiu}3N{6>tsM_ay_Rb0JV!6t*O z)o_qRK)WG_RiGu)vH^ZR^An(7B3Xdx|xfi+Pv&U8Bt7H>Ppypm0H? zsD-}Z39Do%xF-&;>HnL$!&Q(`f6g^ka3|$mgFOjG=d0FIru(iegt~c;ZiVILP={J7 zUkqTGs_F@lUZ^Fsi(;$<4A~$n3d5PPNqJMpo9Gj_luN8+Z{%cLF^19c&F0Yv z3&k34wZ%yKTjl7EaJh_}$GrD*s!?3Qq(qJwpz0l?2iDf*>r!f@z};$(O@oT&&hDLK zsfG!N1gIlb+|Ebp4YXq;%5dmOl>R!$l7Nx4By}S#sco6tsky;1%AZbfB{KSEk-+83 zIu#Bot5Dtb#fS76of}2V+E`0jil9*D4Nh?Fg!4yPtg4TYAq-e$u8mD$C}%Q}V)T!r z2G|FrduJ^kp-5O@Q2ipPJ6dFowqWRRF{4EMC2I9w5*?|img6?e>e)vURaBNx z>BOVi+*!g%)+6(IjdqgD7gcESV_GE;t&$xjsd#;_;Z^0;6;jr<-7A!K)v<%0{IzD? zQc;cvl(ZgR^u{|1i|b@BE*pt>f657)FS8hX_&y5mtU6kDVTDD&F%MRmju1ARPu(NB z6kMY#nlIbaf`&ySA^I{W6{-+(3tJDJsoEs1GqV|uI!8irYl*Ki$2COdV>KRz@oX5p zvkaz9_f3O#+y{Tz@v35g6U6x2U8&qXG}fTi$Z1RQ$2Hhsg2sNO0wg4q>V7*Js!Xn| zYygC-G+X&e$+FL}3dvMuVV8|ykyQmQQAoz6PeqhPvyQ&-VN5;!+~6qVg*oZ&|A#FQ z$Yy_7BHnSxd?qafW$KtNzn;=Y)Mos!b8hs%Up*ybWemPUd zah_BhdWBl5y|`BTD&C(2@w3V{RC%Tk=M4aXaF0{}eeIr(<9p*nj`PJJ-c;cIDr22S zIY@)5#yXA>r?_auu-$Z%G&S#npf`6-Pojk zlU788Uxx`#?Hzb9l`Q#9BT5V>)_N`4I|(MR&`*{lnvDWJLvXi*G%B&ecD=^#xo;O1 zU-?uU3I@Aqu++0D=nZ>PTK%&{E%)Ngv|!sQwqbA_gf zD$@8pu37vudnm&vjb8)ei?2}ufZ4?YZG_GXHyLwNI`&@_E=D--1EYrpFVptZY&?UR zq^Hx#mr7ifrc>1(1N&rsd=f;MeAXro23gg#4dw9rquJ66NI7%y)d?zQQEq@U?C;0q ziMa0t%F-qImdP)*VX)j&$qwrE2n9~E?eiN$O7r|Up!U2&kP$`KS#Yg0Xgv< z=;B?$jI-3>G@>ytp7s|!XdBWE5@zucH;dGjw*nsL#{+O&Qg)v|HL|VMXMFC>R7}O4 z0$a%UEti$8lW}c+>3Q4eaY}=V<+=Jb)pR&zX8{hsK}f(=^qH}JZpVR}#z=hWB4T9v znOi6C+Fkbm;OW;ddfToEn4wjE6FLq)3#1@pl$sOmSatko+bp318efH?u*+vA9!HE= z+5pYiB(fLTcL)6AIrrmKIHrk-fpuPtLAyb+15ZJ^=w23--k}2+k^J<;fdR`eQ$fNi zXnJ_0Ii;fPqiM-MKE2nu?{})(lpSxtQS>&}qf_)Tg5(u^Bf1)Q!d_ybs_HUtNVFi)GqKHtnwvGFJA3 zQR*I69?E+mS=a&lKRaD6ubvf^5@fCOj%hH}F}>G!#j`=V{aySUC+ zuDl20A=0rE2>(kH4?-laU3ER;IRU|Qe!T&%T1VOt)p zOi@nq2&P#glXu-^c~Fe*p+?}g$}H5=(rdoa6EWQw!w(rNrH6+j+m==g0SQt9_cLXJV1tX=c3S9!Gy(MFUM%!45Yg@;$ zkMe4HavA^_O!&-@UfE$MghE7Usm*7c3m}L1Z zi1;FHQJiT-ji&%>yf|*DUcF6QC59y$$&0X6%6rlz%~C#2hP_s*|IFTgul6ifJslW> zI1>q-^@lc^IJC$gJPZ9Qa%kUXGAlx}E6~~O=eqYn2peA(hU}-}WQh2)0NV1U36}QW zs61(nS=zuRv%w|JWsQci*k8TLRKaT+nj%;kOFp+`KvdUO8<`aBLZwgp5hbHv8O9Vd z{aY9h(cf?M?=0uPU)VEcSZxG60tw)?(WuRTd*>|}k4DDOlJdkm>0GAB+v|&?G$KuD zi*(iSR{nxcMTb<-Xfl|-5yL|Rb7CB#vma@1t}uzJ0mphjpN+6ohvy00ucn`{yx?u~ zDw8EaaQnp~(?*8m&qLyZ_U;ylM=^i$-@ohdXKiKJpUdUbXhlQ0R->XZOrSkgJ1n6w zPvAbY;Ce{&r3d9MUGVuJ`&SjI7(z z;3%dMY0<9$2eNFnZ`4mutvK?5Ts?dF>NEulKrdt;sK1jP`k zQJ>c?$|51wB*imU@11+1UTfPNcL7(c;mIf=Pv(Q#G13V4>tRA~n?5?bVPKd9Do^em z$@?ghW6OUvETFzw<+Ff6k|1U9fG3Xeli~062ap+DG-8~lA7eZse=(MucSy}qMieA>THo6x(M+{=g7a%J z=Ka|)0s3|!0uGK~RA2(nU(Dp4Ls7w=5T$6;9x(|5%4Qabqw=N`fYe2RRMODOGt81{ zOhHIj~L-G*m9#>=-RAKm7CJjlZPyw{KUjTk1cXC zh9eiAYiVU+`Z`jt|L_C{=vODXJE(VK1W)l96fo*nA!8lkcA1+lnZdrVGR=naq^25M zv#JVU+u~+!R$VK6_<@3Cj?VW{K)8?P;kj|!h$W1USH3JOrxnM1r-VA6}!uZy>k*~+9{uHLiaa7N`7lS}T=AFDXcIJOV4!``*GX5^Fs=Irp` zth}>FqT145ka$5?|F%tGinPWK9d{Z=q?u|m3Ui6%?jT=Ed)EH$Y$%gOu>8BPkA z$q2(CjhWq(M+sv!xIdk5c<)#RHWJ;89IaNui>@d}y%KHIU|OU{hlBmoL&rCac{>WP z$&6v*g0+?^1WT+5i-U7tkG1Hb9N{EyMkrJ|n-MoW!@boZjmDIm3$Mif-kPA)W(2@! zM!<6l{JmwuW{K^Sth4@OXlA3ZU`(B3xF7<|_Fy5}wdvV;mk9OxE-9x}%=725csmzb z03^WC8+%lp_<{32on~@CWovblE6&9<)}-Pf@&>L>55b#*yaiZ%^mQCfNyPF?LRYsZ z#{nbE+B5(4x_0cGu0~_7V+fVy#!KqM3GtZ#N-aWaQ0*so$%<0 zbmzsi89Cj4Q&h392<6X7TXS`>nkl)`XZ6dpC(q;qdWDQW^6)Vm*t5|MTm1)$>_Rn) zNmHn_L_vzMbIkH2LaJtMLGcSrTyeQh?O-{Rs0V|0gf*Qm%F!Su7`dNjXr;Nz0uwRR z!S2i#%~7Lpx=(bfRIHI+Z6~&nx&4D|yh~sUBFq^=X#=69CEQZ$=8vfpxg{1}sG~G1 z0V%5W#V1@{CUOV#rUiFp%dAQaD<-;8B1WiJ>es(6u{9w;~OMN6+`@-9HVuW+!p;Nbux7 zQ?e^0WbYNROFWn5Pl6Sihl#uM@N9W`CPlQB=-z4bE%zvbWE9>ca0eqQ9# zO4fXt^k9tfeoAMI-J2q;KYczfE41&jU1yn;DGI$Vb}D$>`ZcYwFUi06QHCTecy5j) za?jG~Mi}1qcZ#O1k;`WC$a8XbH<`#P5Omz07NtS>#E{C+6wIdh;3u9HJC}&jzcML$ z%Tsymro}S9g|pcIWsM_|KrXyysm>U*MCbuwgYVN&5jW*4GHH?*7_6KW2}xhe7lt{) z4#Em?cWbBw>O{mDn|k|MEMUu>>CJ9J+OH|;R@Fh?tGEbfSbUcgRQ8p{1&vR*%?XXp)-i}1LnaI?*-Msh`{~ASE(x^m=qA4S%R$re)w9fDG2&Fv zx-x*0>JCS@6Ov38i#(u!q!nD9$DaL!n*_C*;=!eA&iVb2kxE62p*Esz*?h0Tm_h#| zW{qYJ;&jGMY3dmWl6`)ZF~YTD)BIqJlR$Ho4N&fLf1y$P>ys$yT^tM4nZXIEz+^E2 zlj}f88o7@&Ki?BdJJnRyeN=;9#@>|eLCAso*8#8CQ_0%YACTg2ZD$sv8W?kYRS-jE z7?HmY{E5Doj9jZS5}Jk(S0&(`{iXe)NKOQ`C-YsGI2!HeR1sW=bWWsVLK>0{yuEE( z2PjJis4s9kWjut%@3_X@-=#)enByt=DZgZA&^r`hxxj*&70^>+GVU532-z5vF)%h% z>XIj2%?=jJjc!4x6LUm2;gh7ymczp_5v!!z)4nnm*o0G+QWZRJxKfd~9KS!|UAK6t zf+?ebVcm&_d!^-ugdW9w3{eTMsBy%6hXbd{gOW?L%gMFZ$)L$$E!$WtXHw_?On?0= zgViBA*$kJJDew@skR`j!C;5&!EHe=;e*eeb+2J=vV}wuV;&XTez&?WU+q)RH-OqI7 z5a21b?}p;%05hfSlKu&7&z)h*uMRMBVM3LDr32N{cBC-B;(q{9P%j|6@<*ml;85%7 zS(`Oep1eN>ejli6z5+GLN5IFGDWQ0qFZwpt#i?*ZIzMJkuFE?wkGo_KP-!GWO(dp> zfAd`cKvcWKS;lL;7vP6fIhGuJyy9_Pqn|@#jKb@~gA=6an}v>mu`2f?@&VltQWls& zvOE0fKCvl$HI;g#AJ_oUrRZdcr_pGq#S20*`>%voj?hE??`gJ6)cjvBn^bcei*pEq zENbUy!u9rOg$a**uSPxMlvI%Opy>*&JR8VepE?1TJOE>!|Ys8Zqc2{D^CSgm=(w=eVUQ3|o$NnLu5r_qi-}IX@!=?B1r9&|{cy0Qpe_Z% z+cci2JeTYIw1bes2@O+`LUSk``EsG4#d(S zE6a*EFb`$;;J1gK`*dL;2v0W>b*Trvl;r?pivacmI3HWsNBaUNswyn7TX1Qi7GQK^ z8|bhc((pNvq>2J&j-X>b5H9+2LRyi?=L2V#nG^mG?2pgyf1d(_bX#vh%B*xG8|7*S zcJ|xf-yd}YNV>E|a=$ob?aY{yq}260`G33vhV!-!wyrbl`R+EL3jxg5PGD=voEA1? z@3ZDw9$t1hS%wlT_zq-@_E2>lx;P4RuBCSP6b+MXYqwAiO8_d~)@NPzH{iLdb}J~q zv6j-z9i|4VZ|S>CahT-jlLvpc*XsdYN?DR{sGt|sg5 zN8!NtFTf$t8wQtq&P-jnbVkF>0hnz|T|Q<9VhP zNjkr&1q!t;n4ss#>EC5*4*&}6axi|W*v5BXjd-yKQV?{2X90NU0T|PpKJCB;Wahi{=&q{JU)QYOZhYLNZG9v?2F-s ztp~}4+H?Ah?-5|{yG&o#R389m^7$SXd0x3VgY#v(3O7yxL(-0;JV#Tp>v;bLY;US zXMc-?t*6{DOd!@WvXJx0?qKNmy8q?aw3geio^KacQ7L5ar!gFaMGFNQgmr6cq9#9o zg!i;dP9ee*AcuhwP`jChl%xK&eH|r><2y`Jbyd0!-mdgQBh{sh-Pl((4V^U9i_}NT zbsEN%(*tTq6)>~f#qOAxv6g(-%0v@GpYmj9N@1+sN`dmo7KR6IT2u}iGR&b9&M#EkIb{(7$2>78 zUW@{e7FlSWc>1Y8nQtSr{?1r^1)mw5yfYxck?(+L7ZM=YVy8jGk+Q85o^d2EsWF^w|W7FglkwC?>{+_$_Jsn8+-HOQ6#?9d4u7;{h|$5?d&f^~-eU?;)9SLX?5&Qd7!I zAI_-4eJFM{+MlKW&+kVUdBe6*zVIG6zXAWI$&Err4Eqm;077Bu#&CU%5c>A0cz1yC_-ubC$cQ~f5L z;ggL($=cbm{ZNy1$M(MZcpFCowaV-`Z=xkoc`gwlv=3bBsWhr9yQx>2V=OX7snd+D z@?5$6hbimZ0Oh&79aI)V)UJ4I`}ICWOj8|8E|pa}S2;r%+nlX<>8cLuwWRRt*TGL7 z;0!Ecevk+ZugV=g2w@rl0N@}(NJ&-$EB-=NzK9l{j2FqDMUBIrMUNDonsbxQ(kgyJ zp9KkK)M6faFa!EO^UokfWi z>T1l$xMrr-S)^5upG8=~I~en*K~{G_B41Qadz$ew`;r5YN>qgeIP{hPZ-kEGf35}5+6 z3{o1**K7C`doi-b>04!Z%=#Vrq4qsEl&6kDuIcozNGwb{+)Ob))ocFogeH>Yz@&KA zs;#s-w*S3JNbQ6m{L#`C2rgzBjB8)A!7u-pWq%mt*!W*VS zELA5^m?I}@a&sn*{uZfoT$pX}=eYa}$DMkf z9+eIe<(*s^73bq}@ z=2$dCFlE zWzV)&wom2RYXL_X;wedN#bFi$hXy0$QldVGuIcq!760*iv7f{Af&AY|cC zdz4)Xt#L)D(nvbs#Q!LVt#LNoH=+@}A5%evGAEh1!u7Z#x@FF3MO0K(W>hgAKlu@q zkQ<52oHHNJ*V*RvW<4L~`0e1C8%V%a1rwA#Bk)HZ>;8VhGOkVbU;El5jmC!Bk^wR` zO^FNe>`0h-KI54LmBVeYq&sJ6l$@y$sfyVJwt=8=Fw&em?J%HL5Lj#A)Ta7y`9;;? zjh&_qW?1dhGM%&m%hE3ii*Qt}s)kZ56uQ#u2`KpDNgouDayd!l{_kZXRqUt~d=c1G z)CG0tiZ9X^7)x@7QLuP2A9YepA>4(y6$-372Jd^IQOEPvf9SUij#vh3BAj&<1nI#f zETK)`nu1}=%F*>2G8_vc*Nqmaec`+O_jn4mp>XGC7(+s7|8dFiz#XO&zh&eliDF6x z)NR*7^6T8AuP}w8q&5{Nl6d39Y@x*+yPt9;ODH>9i#32vBqr!Ifd}Ih1oVd(E{BPh zVmXa))5RODSB$UfPzobQ}A%QA<7K!edM$Qa9hlcXD= zmLBbtSCUI2rbt61Nm$Zm>SBhZW(nX~D%ldj2S?YZ8`x%*B=J5am=UZ3Ek@!ZWE9W; z2lLYP8Zj5|cpV%A-ZeQPbzo3cm)7b-?eW&|o@z>g#3XIXnARxMmN6g(t+!k)&Vl{T zMTHMG5eu9=QVHiv9CJ_{5gm6-1ZE^I)bzMCy=u~qfaOyP@YV^oxb zw(eoOzoK5)#=-EbH*2?pJlU=Ly@}V5X?5nDGWvvsn(t%~RPE5Tex~`zM|>r^*?(Z& z$uZs}@MYYannFtV#g|j`G&y$t5&-Yh@%CulJ-HeXj<&3Ji03>r66elP<`^<&&+8&~ikIh~X(LP!QLnMnVQjTg*G(&k+ z(wJ#b-8Bv|nI2D9YS4WE*c;BcYF|X}T^pwJQ%6=H7>nPRSj)vaRY&10*vK%mF@Ehp zu)GR=0nWhbP(UOUO>H~5cgDC|$KmlAr31n3!G#D6;#nH@P@*kW8le&Al>O}BVx?OB z)a~H}f9Mn&4oC`iUdMP9@JAhs`f43R0_%P<)#gdNb%+fHP;mh32n z-ENzy_yvqH5>k8@9JFq=fSf|dN{&2mtexC$>;o{}UVR^AA`6GOG@WuBBWAcgR*j&) ze7)(jc&-teEL0$%ED*H=o1>~+<0ukUtiVuFR>;pqebHO~w7|DL53l^_&Lw*UVD~q4 z)U3L)_9L=JQ2ze99AJ2#NcycGH*Q~?hQJT>o*$Pg7m@4GPCe-c3Rz7gAmOok5dX(Q z0$%k3l>H5nC~@{O85FP-`5)92GmVnhZ6o8sJIdbb)_x>DmNO(LIM_dr-P9#g3wWM? z<@Fma&E82T(8A~?5^&(?O6ojW?(0^O zSj8B^g=UNyRKIn4qhHL*rQU!FjC_Q#ECc^htTSJCg*Si_`U22}>cr#b78%meM-SX1 zd_OY*Ywo%(b|pB?n8evhHU%zU_XCi9^BmrV$;RQko|uP8Px`CKf(a_F@}Iuv<^CZJ z)F1XP z&}0``p}`d@03eTbm#8C7R-Vkz4-z1=oO?{U{WA&&cU9Rl5&6vJ5^~frX!CN|GlSYO zAqd_?DBhR<4p{h;^MvsO3F-4i@tTbuGJ#`64Qe>Z!)JUCGnhAs+tR5zWAzhRVd7IRv zlbWmtAyPq0GRq0MqFYQ;N4#p4!yMkP6ni5pw3;bj$0>zCS*^W?P_9nlIW%OuX2hDe zQyoU=cEHJgV%q2cf}mMAx&$PtET4cR4y`lBxqDoQ$60^6S~_4o;VK|nh^l} z42hOdy8H_@eJ)iT&qN#)Y0l#FI$L3P_&}Heg)i1m`fvW@%7NWK&+Cnt?mP>cv@j}^ zKHzSFBt|7wfy$$4>rM>)-=;IRqV>NhIE#1|!N4#)Aj1cm4E_wG(0(cZw>8`9pqtT#aBob`aP0lyP{V(onn&Yzh2w) zXs+HR;n<@Bt0lx34(d;~H!;fZE;%|Z8Q*Igbi=}kyvw|dkaYPV+>9r$$aN39nhz8r zspiQ?&bdQoAhzFT*mu_ng-tX1YSt10Nh zpE9R?aQyb(fia3KZ~+qqvZ#O$cb7>W+5X@;^O)kKQwnYBj;#A3(Ih%^j_3!{w$8hw_i9AJ1cqgYQ1D90QaP z{6@n6p%B(SdCdw#_hrqb&$5u>3$k}AG+awel{k9G#P0nu%ONw?FehNFg7^v-#sgOk zpd!=IKr4HiLwos1cgUwNZQuSz(IESiu<%G8ei04If2p)PT$7c`o;%6Y5HLBFp~^#T72ktS09`62mXBCC2DB zgEd=Ou8fd2Q16`ie&s?#kju?FQlqDbS>c*bA|VT~v;+pr#*aB_z2x6ji#L^B3-LxU zm~-S_$lc`R0$ifky^CCqzy402UPR+v@0a`STca)ZxBu)DQh~!TLj~LDsl^VN!}bpY z7u|SLb>sochclWM-{7$!)$2R~4Q8TFR`*%RF12QsAxDlM%JG?M41V`*=o>Sh$%{50 zZ%KoOK%%ZqWNak?Bgp#m9&~|Y`!1& z<1(?;mylS&Y0BZCaiF+# zwEMn4lk>lC^!lRdG$Xbkactll1aoeA!Vze6iTgYhENa>G zRD;`Cnhb1&A0%UJF2ay1?$rhNsNmcS5g3%+et$Ss`*hz62Q@;eS7=Q|U$kL)Y3<;( zxDjHktEg0TZF&Wf|C}*RCAzX!8%%i~vE`P>m!w-6J+@3X#CQ4cLt=qS_(tc8VTXe3 z{pK=-5+B(-h*^AcaVy6Q!AHzAWz!aEX~sA|(6E%-zL~OMOQva@s-Mm7V?o~v#RIc} z3wXTnxCQTVsv&_O5U)`oNa&*$%$m(!%Z}PFK<9`0&i5joOAb=%Y=<^2A18S`hiGlc zB?{Uy!8x1yBDGi~Tr%7#)#FMjC(}O8<|Iup_x7}$Sg>0@X^P9bf}4y>$4P3b?pVwE zhmpk|`D-Z8IW$`#yVQs>RR1L7NZNvhPGkfwg@n&DbU3C1R->}x^IA2ub`tgW@aoVc z%gnt@%0RJW(vClP^SQb3`9UZv1g#M~~vncKsIN;j52#+ONzlJti z$7yKBhn_4^{vNe--M+>Y;9n*;H;dlG;soylY;j z?Uaejg38-=moU!LdTpvL-NamkwvC`{O)|src_KSM-Kfg;xUI%d_d7WK*UF=M>}t#I zF`J#C%xP0(j^iuyTaSMC&tLQHQakMl1dF_M>;GkeT!Mrueplq_ap$!D=qhd z+{dyWuWp02%^{ou2S?tRilstCW7r_{*qmE=+Kd|&E8^pHx?KCir}%Jm?o;W<^}8+i zeAML#Q-oG`=?`#=%N0v>$`|dO!}=Is}r#w)FIO8AT?L& z=NZ^jOPkPVivuV5i-R%q3qVkZBMLaetXwRzj|`K<;}6#)@o;oVyhy@-{nN?To;lK% zbVOL;q_e*2Otw=|ez+a{fqVW%;bmN5J6XqGRsP%mkP4PClPq8l++0g)`rN)La$AQZ zXI|;l4#-R2mw1po74&8Xg1F(`;{x?Z)O(&q42U3CI@}(pjTHiI>^UK@fDDI&F0@Ou zKkbV8Lfh)vTIS5ZKRrEuHAysaLnR^#EPNwQ)5+%`$VY&vJ?Jz-NYyiD=3EnKIsIS2 z6!`W@<0qPQCOC9dEoqAN2H!I4_u~Y&YiuOLYiO=2YpS^9|3N|MHOq&23CA0cm3PM4 zOKe69RUOcp)s?Ri{|`|D1u&VWr_?wkJ7BoUELLiRvPvHQZh3tB@HGY^s?OG&Iu090 zyN;SI51JUDr%^{6Lmq7t`me`CZbL<5S$+5`OM9(F6mkeLg6WX-kflY)6>HeAsFaM$G*2%@kbOHTG|t#N7_XEUGE{hXsdBXe-A=XXK3IGE)e;d$^MQ1nEcPSBd4n z4$(TYSMFm$52eg9GDUuJCyV(e1?ZmH(n_9~T&prq{gXci#lx68RMVzy`zv-tjQA^9 zD{9H+30U2q6A#5AA6T@_H*1RV%@Zluy(H+6sm}*iYj(c2~(|S95@QO;{y~r$BDm?dz{)59p5#i)_5IZ@G9n5LPNv2 z<(N{bDuXab{fIY4T~8 z-I@sK0p>k^XXNQ*lpXKkMp>-cx0f5w?I)l%|5}u&;LH&hYDpCCn$#kZivsEy*p+s9 zWwl~X85A~8PO+U-z+)Wv!ODGIUEaW%BBKgQyNRBOfXzc*NOezZ zkAzsSF$x|{?YoDSA8#wtJG3^XRJd9g4S{{7!j~DY(!N&_a>g|iqO*saXfHqya1^V# zO0!dc`uMS9x8qkW75hN5EpheV+u(#Z$~Aw$mo98LRDv6-c3rq{*%REb0c zLr1~#g}M0ys(1zLI0o|pifT^&Y@%J7)rs}hVSDLX|)LOgjDS^RFGVB@+a=l zcGho>dMqD8r+UCG4*s}`9<;DtgQM3+U!^z1m=+r=Xo14tvqINn-} zRqHU&W#0ix`GNWz5=*k^P>f5x(pYtXs|Nis14NK#A%l%5Sm3aF+HoV=3*lQmMx zw#^U|1X{Li@QM4;jnQY!yT=)UUBDzNe{>AdARGj;H}}?P-5GsQ05#=i3>qbSj_P6A zCrh_#EJJ=5&>>d|GGp<+M67;3@tWq%3fsW)?Eqf5^ZaMrUi<)s@ka%?!JhVIcTg|& z2#CN?z6|n0{|&?Uir#+!vM;!$#%D>>;dt`mgP6NX`o?`xx~eGj>npPVcKMpmpDuvh zv0_u(g%uiMlO~8^VOR$Ot;a3N!y`k$4|1lohg}*w_<82S)^qFwww993vK)ycTsqlI zs+#hN4LNjl{HMx6oB(dJQ?%#vW+|NunNRzjujnt5I&;*g3APzNm;XoBUkAnYyicGo z5Fog_y9HPz5S-v{!QCB#y9al74<0nQI{`v)cXxM(d-n5vfA9U{)=*Hj3tNZXIWy06 z_tQPfGf!j<$A}F%FFabL<-1|UG<{1ZY;*(I7xLgk1T$yN>ZECBLP4rabG1<{d9t4Z z1gLX5m;yZm9N}9Gqv-CXYlAVbWU?ig(enj~HEXkGIXMhzCKbP&5AS*amhrxYJNxTW zWpX!I#ZnI1@?`LNcUL1IbJA;(O*tQnzmUu8)6hUk<8`}ErS2G|WGuFe(A7ZDWk$@Chl zD!GUG0Iq(LWr+c|e~vZo3SciC+2{=aw5VjIn@v6jv>aEp#T?T8ZqTU<`av~1B>>WE z?W#F#T-HK05wmheD4Fv)c%>x6{?Azb5YEtgK}$Hcr90RhX=Yc|1$y^Z ziGXnfuShYHq@dm~q!1$5l8M`UrI{Q<%&lnv$2O$ViYVG>5{H;9|4Z-F1hj> z-r6%05F+$c)>@0$SAfn&xNgpLFW+p2k(W5z_vtW{)VFzezh77?)oId~BDyJx2OnYa zl+dkU=#QdXri&>PH~k}e1OdkKw`h#v_U>tiL8ol6Pn{V~HhW)mE`GaI9M!}NSGGb&k+D&CONKoHB=VLwOHGrFYwNh{r80;|m4N>|dJFJKHHp5($kFfY~4^V&+_ z@)_zIS}bo~pM z<}5A}%BgZ!+t>^Ye|ukS8#iCB{q^>)Dwyjh7pB?9`U#e08DT07a}R=nZcmYhB}qD6 z4E=*-IQ7pWBVj5jNCZhNY?i#Rhy^HwQoLU`s1F@y9@BlZUez7vEvd_^`#EiUfgXbd zPW|IV)}!3_OK!RN1eCbcc=QI$8g=1K@HO)?6#QtPp5y>c4Jvg zl@R0^*K(IXpA%`hI(Nw5)JCpzrp|!utc4M0YHQ{L(r{zZ_Iw6)QyV3~z8I2@AH9td zI@+IylqTqB#M9A1E~meqQPazPZT(J5e9{IE;dDaQUEQt1(Mzg@k5YXv3y_)C)T?)-44-%^OOn;`;HbGDyM-DsZtMNw=>q<>BI7Gf4NNq81f2PE}ZNP%lbH$sN#VN^Y zdUHc%y|myH-^eSndAu$O%nEP_Dv4w619Y@UBrj8-*KgrmN8^4t7S5jcVb8mOH0&n& z1H%JfM$%a}Ixn)23;_RMIWjmYH@kTUY*B+e!oR#22O&Z^wV;^}mgD`1h@4JpTa9VPXom8m#d>@{ zDdz4~Y(26+MO>%+*58Q$YP%u#{VXuF?gX`{beF7a?pVNC;0@nE&^EK0!yaxiy(e7c+j5i&YNHqy4wse}u~B_POO%L#1|bHS zjBLfVb7hxB$t!Ver)Icd_Lrq~pd`~HVob8&yB#+75<^<53D?VtrjaqbIMUL$3NUZN z^#s?(dv(Z|dAQ|mB}r6vn-H+N-odj8{SJZ|jVDkmGBr8-k*!P=@6PF17&6Izn7EDh zOtSs?HdsqA-d-h;;aQPm1S4Ov_!s7%U7w^HOob}rI}T(sk3k9h`JE$&RzrADgDVnA zR~Av7fqbT-MH&op7kRys69bZjKDjF2!=P?O41+q$Ua&+u(2T>TLS7w`X-$F(U#%5J zzabkP_2aU}iX^m!s#)2LE1>(ly8#QBQEHIXnmzj}-v~NoIBdeo`kAeqFc*X%{EMaJ zTm>`e>m9QtdcJig=s9wFifkN-8SuX9bvQMNb;kIAeBx`X%(ECN9Q2H2&Bl?J&2qsc zyf=bd2xy3=w0LK&mr|D=7mN}Q9* z>LYTGYehCtueLq+l%5XcLrUCIP1rFKh+1~68$wdAEEsyU-BaW3U=ed9qJ`$R&Too9 za#yL*_ME$LBBQ3`p}WQ({XSk~OE!*7Q`$pa6AD_(j3BcGTQO`_J3d|P-k5zEMh>Wp z>UM?T8mX%ra2rht4|t3>E4}4Jvq!!ui=Hfz0z3Tvw=qn3g}uf0ZGJ8 zYiggl+qmJ`$gHjn!HyttR#u?ba2G6uRO9{{>^O;Gx;7hSzueZcLhdT<7U3;%@&NcE)OfU8AP7O?JN>0}>6fP(*whbR#;Ar`~k03xAH`F?w5>}Q7N+(UTmwB7Yd+3277eU~5O6UtO zqw$pnMVj%75i~{3xP#H`m4jKGnjnX4k^3Ed4}=h3P@mtggXV%)rS96AtUQNS5Z^Sb z$jg*1bBKqMktT`76?>-fuO}=d_fQX|k<4xk$OYHhfvgBQh+kpftBU@L+red_LciH< zy9-<`)-IqepmCAHX+)jgd15dAS@PJRSMd028CO#aR^>K9S>Xg&Vr*`RCW2)E8= zz~8=d%f*n zT1>1bzs8gRtEvaO36#^QS};T4JXm)SJD~OklpsrVX;`5~i1#~>%d;u(r*1OG$Gz6B zC%9_$Mo0X_u6bSUDUHoiKl-aa*8RE3HJA-2#ReLY^@DuKlN;Z>C9MHh$*Qp#5@Me{uqVO8`yW3vs@2%!OZSxL3uzSzU>y|;p# z$vRmM9R3!o zGvBVa^?dQyifAMc`-1R-ibyx!5N`mkMj^HPe9~EbYOxpmzAm6O zcZdd`YaC`W`&0qoCX`rtD{bCpX>6-~s8IO=NTb?~VZ}n1IrDt${@Ud(+42U|zu#4m z{zUD1%`4T<^8qq@6edw-FBcrJPVc8)YDV)L0;*Ok^}m2&0ZS)^c!^Ad2s`&aLf;(; z-+seC>R(^gO^l+Pk>dQ_+C7N`XRovN`9jrfSmmI4nxxh5S4zQ0m$Jf#d>gLVS1@HT~_=B8lS9+tg5!tGLjI zQKnej($>+UGASp=!E?^@;f)sOKhrC{%m2OAT2`6uEEhwb(kF7C%5ojUTz-;GeCXhtgd^1C zXY#`tiFIFn9CIr+-5$GT&wcy{Sn>CO+l1+9Bqb#=07Qy zXr_BcY6rhO#J)d_zpjjr=Grb-zYE|hz95Xp20UxaWwRu|j=!CZgDq5`K}@nFn-#mi zDQ+e)y1qjWZtRvDv(+PTJ!Y?2h(#LIIo>uUWfy_fwjMn)Ym4J(l;V@bWS7B5sigzk zYPvazZaP_IGQ%m9-N{9$4kSQd{reT@y><&p11e#5M15zrC8@*fW3@n)1e29b0N6!HA2eW3Y3$ z?H61ALr-K=ucMB+tPG54YuBIh1h(IZA+m6?+%FedcOH&pTJG9h>&r*u+`c!yy^SFl zmcWsmMxCC3&GESBmgX=GkSD9=a~~aq|B%yDC4kSy*$~9o@)WHbS}{*fEB=>t;YB3wsF`iLd8t!hBR_q8Dq}kpoG3wH?^(1wR7jWEC|=?Y(L)+0mAUYA@fJ z2hMRy?-ov8XhO)p(#>OiZck8sX$yGm5SK5gyc><>D}8z|%x_Wja)rhkOH^I*eSQx9XMaWiD~a z=b6V%14Ce`xkhkTD_lM=Dd4Zm-X3Ap*q_z6gFHC7eLm^oOaa>{{c#z8YPJ$DM}D+X zLMoTAhaz8)0^YArT2!BWZwr*l1cWMC2MjC<%C&?3R&Y59L9p@%h}g=u+zzo&EM9HF z5i1)AKy6gr!Q-1#Y`!;>RpLEk)`I?*yySP?9XxoJ&Lvs*#kFi=Nc z0}OjP@|tYl>axSNgjKX=E3}$$Tx+jU`4HPdbxaEpV5Q9Y>ZQCdn~H+5Hf6x`Z7G(3 z0I|*?R%BHjGkns~;DVIHOJp|Pus#L4OT4A~A2CxXPxa~&0(tEtTI@iL21_iAdJa{sVJg0&#IKDplb95?wMcS5rEio~hc zSCDI<;MfP?I(0Bne$-#Qpm3vOz@UCu3?n0ZrZ#<_IVQy@x9u z@lO#PC7o+7jo%+H&7pZ?Ay7kYAO?iO*NzK6AIc*+A$;n*PzAK?OWCFiPH6BLV}-&M zSYQ4BZ`gU+zljU8f9f{iin?YE#_|LeeO<%%Z>^(_Nen!7(H@7J!>Eg=wh60A^jG?b zdNtGG*-%X#UNCBCKuS^YC{{x$pL%-IQQKEYfg#H2K}D9zDmGHYV$UP#?6{G15thwP zvEq;9@xX;Od_VLO5>giMp%S7!;)3*$vg zld1GaWE47+Tg^K$#NuQj%F(Lt*j*&K5PXL*g=ES-rH-hT&t5GVup695h0+E2{M$*V6yv_S$}Pci11o_*k0PefrpMVoK4BP zwN-1A$rtD<6lgO)&F03DF$m6sy3Bq%(nYnYvs~#+OI}|>pAr}mEM%pb=>{pe9F|^R ztddjuK3I+^R-XX9M-G&Ic6YbzB?t;}$)TbUwNaWs#v+m< znzrCbEBxW0%8{Lqkl9W!_rmNy*KrY@cB^*aS7HsXU9Ux@ z$+;H(p|MtFL&TTj!fhuD3c~lZ$ua#YD&44;VH%J&6p-hVj$oqm(`R~9GCMZ-zf$7A zI_?&2q4UxE*sx7=4P4gALVJ$v6Io+gZZC*q@@u=gE;sDwA)6A)5k;V&-kovIu`b8g zP2fW~-K`i8#$-_@aMq?Vhr2*5aC0-(F;FEK4IzsCDP#Fzm%5Gy@Z}AQpD-zv^}p5G z>3MRSJ7W@HnIfZM`WzUbIh(Z8aZ-|5Syka^^j9;8Y8%R6N+2K)Ww0eEML1JYkZ$V}WMTcrKAh?-ngk(1E-R0A zs>ND{ad>Bm5IP%ib`mdNEwr5VPX!x)&}D;(qW7gG<;3r=P@*kUz#Dbh%pW&Jv^^a8 z{7>c@Z7e*#5U+P{$?vpNt2^E@pY!?P7fL>r2~!XB1&1s@idax@x26>}pHB(s0c}j3 zf8+>ZTB%Mc24ZOAAYr(3naOPwsx(mu>t87&e8I57Eyh_48K5)B40mGK3zN-t##)=& zSS%5d)uS1>rn)dK#P*(O38Jkpsu-^CM>iK7Wa|Cjg6RaDW1sj36WfX!$?Y2`4C*Y* z%aXrlx-83i-MgkdgyrjglfnyFk6D48iJ7KWj-p!MOD^N4t+#bar}}|}O~RgOa(D0b z?5@Yt&1H78D$j|2oy(_-aS-IsOj&$2077Kmc4F}Vx#Z@C!i9p=0h!n=+Yl2hH{tsu z&%#xb(^Jw>|DMraw)An9haO(aQjQ`mQ9hZoNWi=ItA^g${_MDfYH0M2k;^FkR7Q2B z#}9+e6RLn+GV|&gxb%R=_i!0j;J>=Zx5vD4WWb$dTg_yJ z3vhNBs;c+lGkhNNsGD)sIp>?8aT!zcw#2_AtW?p+lS%V=@kP8hvuk=t9dw50HFGv7 z-hv=ATlRc@*1N@Hpkzq5S3&#G>VJzBjWJ!2D3R5rmVptZZBB0S6~nIq#xXJBMfFCc zr&xi*Dj~U8eQxSdapoWABt!djL4u-8#K(MLI*bCDyuXX*D7qgV=$suoqU^8%;JYs79=_ahArv)L;RpqCfoQDYk#NTJnr~I+lCMdJoZTf0?U_L#=T|i!3AwCdQGUdZK|!k ze!P!thi>s&uGb;k+k{hk+G}ZklY3-S&F#rh&^d;($o%xtgQ3?impnF;#3j=3{4?kT z<#*-y7c1i$x3nF2-XM#Uh4ap7-VV$k=*UV2aH9fg9KMtjZq0qfOPEc@XaByEJuY|Z z%i&y0i%|%*1dh$HEQi8)x>s>x?qr_6CZyw)rj>od8&Cdw97miUq5l^{$(^9p*qmY1 z%SBQHeJiS=$*ZkNeL zSL-q)`;Svwqiyu9^v`)TOk(OI@s2dIO9!Qvn{@ zFL@)1z{=S*S6pzYa|9Fq^AhSyH|k5p_)`Vy`vt1R@*!}$-E{78R6ka~KY_GN zvOd{(FP07LEXlOQK7_?KgX4=zcvs63rSwkHWUOWnPY05Qu$^rkUv=q}w}Fta-7p7Q zo(tcj5??DiYWSaH<+0hqU1 zT?Rwn7CRa0tv>ON`TUBxekT7zU9O*Uw-2qU^=sKfS8bigZN7g?y#E5ZqfJdd%T2_}=ILnd zqGXhk9VV~u*Yh1`E&^nkmdTwn7-TCmUsW@5AJcO!t)Bc(7)x6I?}qL~UwB(Tf6Q5D z0(`{EF#=rEgTDuU*M5i5zkq2>7O1StwE3_~k|gUT$LpTZIJtg-yl#+c>1BI;%6PeU z{<^mQ5WQ|OXvce|c*J6Jzlf4dEN+2+yRW)_&A*;__R+ihV~r)~*nSIbYhk0>^X*#j z{hFKz|1v|}UH^4f|CD#CVi6P5EFvEKI(cK&2F+(OGb%xY4RxauJ~I+N1M1!bfY&W5 z+XKv;qASW{$FE04M;7~}U=}SWwv*|A*>-`a0>LNi&OK}I=C6_dk5FS*eLqwy>-tne zbMHNJX-W1Mz-@bGr1^a9-%qn^@{8Yq6yjGO=^`mZjX}w+gCtw+eGBQ-FF9_jCk+3yJTDC54IOuQ=IkU#9a9a zWW9qIUL3UMp=t{qL^K=j9AoIE;Dwl^g3%X+)0MNVM4h=`DRViU7-Q1uki zl7ufNVj*&Rno!?|slhLwP9e>d{-wijrhufu3YwrXhpmCsWo33>iw_Ow6x`D+_&P1v zAQJ4ckKFs9$N>6iQ*D|>D!Mk39zt~s*l)MVkCpqr1E2pi{lMY<>iX-d+u9LV71#dD z>H1q`qLU$)S(@iruN+nJKllq^jqu6t$W*kCAD+MTluY=C^#Nt%mQh95bicN@)UE7c zAdoC?99Q?oQA_u%o)ymiYrssruf*O>YK$!qqy28;>M4a%s;YYpr3adbl1IQC{`FVW zxL2(hyuKU~LKfpWuL>2TJ}`e+Pt5`X4_kq$a&L)4v#8#_g$+g|vyIiSjYJ{4mfJcq z6L+*bceQEy$Z(HdfaZ$L^_S80xXD&!cN6`-NnQh97pLfAVcqH$(Q&zsBLen18Q}$x z=VT;S7e&u_Qbn>3DeUGaLef`zIfJPaspKA38H29Jzy5a6r2$b2{O2+xn^F2~MVjk! zzpEE)6F|Q3nOxO#;2wr^#O*j*N^7U1)7EM!V@Cvpg$y!i{NmP3|LIi~yhWX2@dO01 zBbjv7^K#kN+GEQPhI1-I0yFsMIG8M5{b*RmRRn`K3*?D%LVv%nKaj6d=MGL`Z6YvB9pbSE04SnhH;}Q zq24)2$KQ85MAegK$YqjWjd1IFHnAH^f0HqA(8=&@D4`zb{ecTx;%o29od20YMuOK_ z0!7xN=NIxy%s+ZiB&VSmOrpqef@aR|K1?Ue%_Mm{(C+~Ci@4N7*%gxjixFO-E$2_^ z0ufH12?EHYFldSme{#efA*a)OQ}bMfeTVeEOp~kAmpu681Bb)brlLR9F3a;Ue5|SH z$SVR}OTz&*0FU<3>sD4T<1FtdLqbutMTE)~@8)UJP3IgriuL!>M>KPioY0-iRMZG~ z1O2r)8hLuR`+RDFqO&PX<2m{|1HHR_3f_kxuM^Puu1{Ae7a5dg>0g3;bIhxde@Qex zvdVFI%E}!FlUgoHR#AAmPg>Xly({TZ#ofIp`M19nn2OAlQWHK|f9Ca&DRgmq?*# zz-b{gZt+3=6I9CA$k1B@jGDccdNNJ!pU8^}+8gFf_=eWut{z^In)I4AD~*T$u+G(Z zt?MMiIIIAtp*auw%5alVI^V0gUT&9#W+XX1SIb?BhRBoO*9tJ?arbJkp!L* zp2gQm($Yi9Qj+YtTN*jz(l{S2O4OyP^B`bIFJAKgbbV~H|f=rCMmh5kld)5 z=y?v4K9DlkL%QnB82S}oR8t-8{?--_Jx*pKJy=u%DIpiyjrM&?$eoe!Cr4@tm1q0S zLjGE`tom7dw|(5bXfJvu^g!_!hA6eT?K2uJXj~^zzABCal@$W!K9h$~IHO-;rS_Oq zS&svk?lUserk7s^ge0Q)I+G?x#WvWX=!oJAbF6#Ap2E_nb;fJc9N1(Jqpg1AFPnQ9 z+2a>d&`x8pOjM4YH&b@m82b&c06J#m01+^%ATnh00u&&c+(cRxtjv|nbA;A`FjL?A zwq~}X-K+0CK~h>M4iycR%~VV*mEt9#dBlO|gSFx#t-D{BV}y_haMzgXmo()Xh_^XK zTve#t9(0ShMxvSb`8CQ5sFnM0Ke&XO=upA=?iyW_yLmKVjOztZ5_+7~ne<86>z-RLVe*(HU@&QZoJa(2Y^0`$O8$!zPfc6A z^b6Va$5eYDxREU{)65`V!cwd^G*UdXFGk6A^};tdq-8jwk*OwO6T=Wn5bD@_@5!*> zqOF<6WZ#KQqAKTKd$e_2BX-~+>Gm=4knOtHLDuqLd zwYOSP-Q_hkqGbeb0(%p^>%@dKhbnbCm6@`Wi^=pD1_k;?;QO)LMq&82pxVJ`#z?sI zzk~GS&D;=0BsOfMAG~HTUniq>;G)d>b((`h^3CuKtwbqyRVC!oeZxf9n}2ooVI`As zSkNUr4Z0h1AKDSF*t2$26VRTqm@zN z&I(WNJ30>{JBQ9yrm>SD!WQ)f4y6=4|FGxeCwy#TCFvQ8{<8M4tN*rJpVo_&4>jjX z1@h#^VPm!O~dDZLZZ1WM_Ge_J_%k0P$ zakYlK^gXdos0+I1mfUA<`RC6yEBCv82cf=0e#0?>u?^qWaP@Y-rtxX;Y`r~Izw;y& zl7;*bY8#m_iW5kB(MehVJ)WTiN!UtNf+=NlaK1j?ALNJd(Lsb1Cpxd(181XnBe35y z?r@-^H~=5B_B-Ga-465(#$Cg4uH4)Pxl2M2-Z1G<2e zpcI7>5k$P7duuiI3+}I%01LOXre%suICcmR=mG`qj0NX@4p_Fac5(^xDvmxm11*xP z>C+j1-L2b|TL-3MpW<4$yI1`de<9SG)ZBXAS~4 zB7!RJ0&zJ}&@YLCkBOYf@F!1gk4vSb91ZY8x@g=8OP*RG(GiQiI<85zYId0bIr$Li zYvY>Lc^H(0+Mv4K$ktc~&UkqxPYY*iZiUwRnog4oTJTnMc=@yrcDDAO2qZCi`u@lG zex2nDhi}tk?At-?(jb6zLQl_`Q{>7RSx}q9t#Yc^f=S!v55tbIv=4A4$~DXR0=H!f zl}4Ev;7|Plm>^Go?NojZn|ds8`wY5yhw2i8gP~q<;$wu%F zIdeximzj+ITM`!p{G6=a1Tx)kD~wX>qP8VQvcpyw+wIk<3pD3Q?I^R-mfZU(NhTl7 zBL=g#y=VO>6yts=y*rTJJDfES%kRW=#`O2H^Rzw(ZM^n-j^kK**17TZZWnxi1fYWR zTMs%3iSIlG;JNy4?ssQi;ecQw+HR1f{59$fIIXVx0Ba<#eNc)Ct9a%##64>dAW87| zi1?Oz1%{lKh2xagVEF0+4HD5s1k~Od^q)Bt&K)9Hl$9-I-T-xt4O|7q#h))2<;H|g zA{$ns8xnpu5^isvg)qjhzH_m=!UsOtw+*{3*!Nc{ar%$R`j_fF+M72hw6^CVeTUvs z=Ya9oGcZ#z7nXV!nz)}lW4|Pni;m>$R`BPha&wmbba8e|S@V#aUs{x=w|>@wW$N}t zn)8CyN@zIn%=G|3t>L($j@&8T+ephV+x+?Z_LX?;#f|*Lbyrhbhfb?KjMf6nu7B19 zk@YTH*{tPq*VYDhPnLiW8|YJ3!mWuy6jD^Sxl1}>hHX9=;P zEV$OgkO>(|0`-+``jvqBPwKtul7G+Cf7G@0XEDcVy?bft473>SPjy8_Ml>znY2}00*+ z*yDw3VKxw_0ZKJR&y5eHbUyi4Z2Q-<&oTaXXI!UI>t(Ws=U-UHPRQUC48Xt#GBxaK z4uRVql;XAIVF;50DfL6cMK<){)j~VnW*uy<-tXaN%?pU(H!c`qR5b6{QFDRobOCl% znK3QFOH9GZLAvMIp3TpbIZOh`&JWN-Vj?^g8wF_Vn^)f($?~y&uU}hZJ3-v_z2V=C z(n#o+!Wt(=GAjqCHlU$h+EDIYTJAjVQIhY-|JvpbN}ttXeIN3fiv;g-$-6GR`61}fIDu#B&dc}huD{_H^nB1A7o_S6en!c4Rj`Pa6Mqs8aETZAq{%;e^y ztwCCv8Kc>#uB(G?Mi`Uw?#kC|eNH{o^;F)))*o+q^_2>Ow1*PNGqAy$RvUjCj$*3K{v+@#DAaa$jwN;8=?S^=zY9c8^zNS*tEN#=Ey!K@w$N zM)%$jVdVG0dKcw#EHUD7&}AF>v}WNIu;KK-XvBj4Va%m=o@;rPCODLlYaBCAAci2I z@g0@50|{CaB2xMW=ua|M->>KBHPTaiaj;orMDgAI;j}+Q!ub%9;nld)J-%ZlaPU&- z*vbW!l%i^|4A$W7B`!K{!KFx;_-{En$|lm?JpG}=|LmA-dN-^+DBWc zZ#bN?UoyjMcxo-e$W~~Fn7TrXe4Q1-9{RakR$UpkVyK@k6VJ?!$A&+jf*ECchWA#mv2D`Zd zzOgc1SXtzCXa>&;OQ3P!wajqt2Spv@Lcb{Jd9hmEB?ON=YMsWCzP)gmZbK2K8Umpf z+TmX{gs{NC)j*QKBSZk6Rz@23FCHD9rF=^VX$|gz%DU4sFU9tYQ5m`7{~lF4DwKe!C4h5g@*qt2NeKAYHgQzw=dcP)EO z^mv)hsTP40*8Ip8+{x74=TTLjNkBfCkbC2WBXq7Rr6s(oWpLC7h9o?lT{Ud+1$lpx1Xa)BRRjJ(bAyC zevn7gvr_)WVqHVmu@LK})UUOd;(BZ+LU7pz@7at1-znn3UGpD~1uk0uzhfUZ;B@u8 zCH*f`{Qs!o|M&YcCSJ@XGQrvTAD8-{CW(;3zD(a}#dW!9b0>>j zee}6VZ3JK10MQiS5&<^x?yJG`-t}p~)1K3R4}c&1?+H2aZv1BBi%c6!cJS+f)lFnz zbyK+IS^}7v{lCFg4QKOl>VTMvBY>OfmV=D0f1)ab_lG-Chvy)yDE?b_pB-AKNj z?qKF^0h~>pJywpbJ#YyU23`PBJ>G-U^MAkmzn`XLHQvWk+77b z`ui_VQZ4f{2b5>ZNDTTB2a#<8!4HI1rH&CUnSeA2&t92F>g>7U2PTud4l1YqkxZ_B zUrQ+H%3SYt(GiImxsUqn9D6GKr9m{mhkwRYR|_NM2hSJu*=$+AhR567DSmy zSK}qcW%kDX*+>|o*H7cu?v+5M-D5^8^sj3pvybHBWDY&+*6ZF}Hnv@b_<`0?)SHo2VM3b(%HQ@M{gNoL+wf`U#}8s z{byH(|0r)fB$)Q;UJ2J?5d|d>r4i0n(8q5=1Ou9kj^g6y7h3Bc0POJU0@G@DVRjq|?7|`%!&;?`?sJ4$fIv3Zj>+b- z&oW&|=i9SC>VrR+{>eA`ZN~xT2KCEvV?MM5!xHPicb-3Y-f0tA&b5O>s{GoTym%8R zJhp$<-^X&q68E0p>YbaqQD{-v-KyCQl*tjUP6=7i%IsvKA8eeegj|COQoH9yyX6pn zq}cO^m~%KZNcNFitGZ4ezCaa_Rs_ISzjlG^x5D7FjgQ(Egu6u@P+b;K=Xc-F?>0%_ zHl6aIgvDtDQ_G{jN7j?*QhKm^Ecq3|TCo-?;@8mg)w3+rFQai1yaO$;ezEuNsJ&Ra zz8fEX^R#%jqN4g3boYnBv^EQ+WVVy)V8#_II6m(uU3Lq z2b6+Wl;`e)ylg}0^USA5Kjw81$?D+PqO2L8yLT(spDMvje2B>l&?zZ@W^s6c%5i0+ z_B-n;43=^`jG&mNhD|t%-nrl@jcbWeh}%0g#;kMNzKz{cbg*IlgKYc0wEWs`w)Eh_OT@dKK(U9}xIUYmTHUD!!<~P1#6b zrMeU+i|&b<0u}_vYT_LTl#Hk+1fW>GpMk%Mp{+qSRcNtnKCo zws;E!TYxi1_`SI&P?wKg=$AD^C^0Mto0>`bR^ga|NZ72a;}`t}cC3FU!aml7Q1diU zc6VSVb>3`Gdu|^CGS;{J2ghFS+hK0j6C?ao|II|x+wIL@Pt(}qtKUeIcjW8{!E}aD zR8+&}x1QhPOL!eu$<_6e%H}~ZB$ZZhW({hcs84yQPqm=uTEHUW`y7irJXg}U$LFqh zFDRJybKUx@C$Q9_`t7LNRgvBRhjAG=8Nd5Z%YD&Y(;(DeK^`q0xRU9*ks-ddm1eAR zg(XP-(!Pld=&auzFjjlcwgrH%dndtp9wrIS%Ed-eFibu-|4Uu&oH0>jfDP1GH_TQyEDnlCm}OX_*2%3 zoaocVMf#La`oy7s&Ed(byQ%aMYkLtCkLB=?963iPpmqLD-?fYpTUq~(GWRc1v}P8j zAa0a(e~Jk6@_TgDB($hxXQ{l3-Q;O==I`H$UsF9f9bD5M9K>JeC|$iDx&9bm0FKPI zp`jEz#H_4&*e`>=Cxc+8bQf6r3yZm#BPj+u?_U6dQJOP5E?aM;s4Fc=-m;w}(U~Uf z{3{=ZEY^hl}f$z4BFeg09gD}Co87cVV@l_z2!|7!rX|# zf)TLxA`2m~kWew1=Mk?1t!KCiwtxHUtv_IGQsvfBPM(kz5%luVvAF#HYJEVphDbbY zNitCUr|$yjBqvTc3zf})W9FL}^4jnNq#(iAA7>uEyN>hZcE25Wy_IV*w7;%lX^_c| z&ikIuKVgkyvF{pUuo`BCpk*3pV3t^LGmWY0*4eo{Z-L&78v1d={4KTUJGZc_C1D!- z@rGlPwz~ikM_;Pu zwE$Bp!1gtyyCDkb7N4F26vdcOjWlXPG6T6U!?`yt)wh43<`~w*M`I=sBS23-LA9Cm z_bG8|xSzyROG0nX7tkB8~eC}q8pPc zpAYhT9D!*$nn4*p+rJh#ziy8{F%h2Ma;x@Qm0?zsOP9t>w0?&D{6qQi=oIPTE}=#; z>#_143lK-WPo#*Imiu3Bqv&W;F^&mgq)CSJ=h{Hc6^Q z8-{K+<%O73`JE_8`|v;o$lD6 zuG)UvP11^|8dt$5uP;J({>S(k3+_1~I_9;AxR+OwS;9?4k_25~$Ia4S>KqS53?=i% zk{ZzjN&Jff-M{HSc+##GSi|0=nKo5m5jwT=$mEW}7?wL8L%t$Gg$ChX{suhS_HW)* z;+>fw*SRb25p-Rm=(Hd8^KOHy+a&k)+wKemE?P98jq2yt>g^8HdC|_J8vwy`f;gDP z57-{4vm!6XC_2WJB3b47?v#44htAdigYFUBboJfo64Re5nL^f=h)LDx%*p%08OSgF z_znYd4_MJOjxe=1go^lC!o$MfD!9!qU>rq=LH>`Mn)^kS!UO<2X0OAP;d4kxX7Dd( z*a}Ctg{cI_>IAmj(Ujb5uHBtU?mWu(F$PrK7_|`u_nia=XW(j+gU^0lh^1rP6dEBj zg4<*A^KH}`xL*4c&EWTvqCy*06(|HuFFX*6rx{%ihKFirqr1iSy@NwR`dp>5pHgAY z?tRbiO{f-w?sdH)p2!Gk_#Y2>P&0){1}b{GmJEWd@3YsYXHTXZip{4>6-R5O9N7QP zL{GmUggZb=A}gP_fVqftlB<>9Y6E7p`p*t$KyNV4T`1=maL@R62ams01@ko#Rii^Q zIx<0|@PK+M=+q+BI6A0BDZl>vY4iyq#Z?yvkH|WlwQf@sRIh3qvNE{IJpYty+MD}F@~D|oE^dM8x-z&iuq`7okHfwVU-4Ij!F ziD0P0;%eubE-Y5@M7_=OCvWm(h++ z2AiY%<=xTV&g_8slk$&;aliB7qc@-mN};1za6`6a>iRpv*;?(026fzR2mT-G&a$b^ zuwA>jOR(Y+9E!WUOK^9W;tmB`+}+(Bio08JcPQ=@C|aDp;o0+id-o66laHBUCPQ-F z*D1$ZZRQ)g3CWp{^0QR?B|*#9h($6D5o21NhQ^s_Q}1Kf?w0l={H!xE@ zqfkdWg%Yh>uaec2TEuUmeqIwMcrHJ3a$dSM0UtYIh;dGh6Ay-^F=I{C(CG*oKoN~Q zjqT^KMvNHou=ijAp&=4!!D;DM_6-6pm*Q+?>W<`Z6@@Banw=h4^83`a|FQt!6zUO! zL~W3wDN%IPo)S?#h9@OwI2wJy_m~83&al@??T*s8oI2A~+{odX>QD+2Mrso^8x;?o zer_Y*c8gz?lxrN!b~jsH5TAhjpV^T2IaU||pL!?WeyY@+ST%h={^+s(po3{y_`}J$ zxgybEk6e5fv0eBDztP3*G$O7pmb|*+@e`N4O4Mw*KAh^y7sM<^vkBwbnp&(fmXkfb zWAeY<`%zQ#;dSv*c^2J_2_@^#9kmSyO>xCwo2!C~V{cTD;rO3f%w50s206%ubtO#u z(FS@-g=?o{Rdt?Ec2}x2YfxQgzW1u-6Yf45M|TQ(g84CuqHIhSDu~VYFuF}R3fA%5 zy2~dTo(j!5^I5Ykij-F0x+*y*CsR)sq06dcCNaeNBx&ugm4xZfTV2IC?n^@Hz6^^G zpU2ZEHR_{ecc=vyEjcjyg;NRXN{Bb}WKB4s=%h<}57S;IYmK08 zRzqUUEzzGfau|)qsrZa4^ek+-8Iv7ES@QJ@3llQ4D7F=U05#N5gG#>(xzE{Z6KL>3bDArjjBe`$_l`bg`Y0#a zfgTL7%;Np*moxq97hD;b3g?INBTO__0R9yk?)`d3doaOT9s|e$)fnm0g0Z@`F_v}a zRJjMByu)W-jn-+K)FY}ansftgK?AOri4IIzR`ug^1`0aA$?(HjU}l{M zZLPR2;WzP}Uu59UyNd(cFU)qZk(f&x!lzX4Y_=xkwt8gIj$+jMkrA(k$tTU335UCb zQERegdWtKMPCS*Pzo|>-tRtJVO$JZrz7r+oyGD}pfJNBgPTObXzl=?pbFvFkWfZ^X zk-eehr)MW!Fvx~WXE5-r!qb~cD|uaHhX-arRqnSgkZVa7Y4R}@vfl)x@W9>r_9q2H zFE`I2CZ93bll3%&bODjjeayt@3EF(8N$g&8wjNh}*EM#gx{~MFArNM_Q@k;th8Y=P z17qBFm%@});)_74XOf5rPxrKw=)$TbsL)gZ2|wUAMiE&COhj{PELwHNiYrUr6B63FddmzEk#Jn zE|{1&Jec5Ddt{ZqX`zuwRJ)cyXB$XL-cn}AYzY7mGf3G~lpGFS93(o!Mc`+SsF|< zy421Ga{NlY2jo$b2eZq}QX{wx7`5rHk{Pj1PdjBv& zk8z^gsVIHgJkVybs9Hr3ncz#+B}>dvf^*Xq+y*huG8=8;VKi3JKsfll0kWNaU2)c7 z{C?Qd3E-?QC%BKI->9S_{{`DqY?BF`(SrRrWtv$gR+M8eaPIt_h6`11flV|*>+_~L zb#U4{_wQA4TjSh%s8K_QjEI`mLT5}6;nFOlZLrPv6_TUnUU#2(iCFRO?{$s5BYTBS zK5;oeH{@ZJ&NZ2v09TxXa~#BZ)YFG?Me3(CX8=7((cFv(;63g8Zk~;N!)!Q%`iQQW z(Y7e|G=yvq9EwR?ZR6oud+Cunl*_w!Qrs6+=vZLnv$*Kx2@^9kLVRH_76a(LarA}*HI03ZkB_`>od>|= z*Wg?~XF)UkbTl%_;##c58-E<0`fn9ZFDYwH;gqJ!&vS5hV38$Ie3&Y0K)Jgo2LTn4 z%|YhV(CVVR91d?i?eI7zJHAT5=rNlcVq6iG^V|JOf0~o(nK>n$1#dEI;o!=G;}Em za03G_1Tupobyx~r))*jafV$4%K!r$N&h{N$RJ^XU@ozOWc1Ub{sm?l=MHYO~ks}OR z!R$l){%M6`SfyJzoTmIkc*O&p-b9UE6drT(x={H{o=vN<<-s%4!8h;YT0NL*MS7VC z1AB;iYCE8?6cwUGB9*&;(zAf&(Y~qQKK)(<%6}YgaCCGRZ>Jl0Z0_xJVz{I?*@$cy zSR6>a=U#;BSWYslLv`r>-Bk zp8Y%!lI5qs7;zTBSyG~sOvWsq5witoYXLYN+u45{m&S|Cc zRr|n=o!$f9wH#QcbZa%j%^UyraboGuUZ4GY33z(Hmw%VljKML1Te@X_-(!9ziea)# zN#a}l2VUC9+}$z0pkJXZHnFqkqe*I&O>*0>8;Fn6?#xYxB$W`3yXq?xLy72w2 z`S1AidV8kSO=7OJyks-905YaR5cK2~<{W$~jOm8JI!I)od#i8zo+|XhGp$ zT5WEnXzx? zpIel%J^dp4;*Jk*+&OslO|1o6%N!29#ACFLS5=qG@V4E*7aLaTdn;m=Vv1{#Jr%A_ zl=X=B=(|wfrxdUXz}Y{aA6H-NwS9tfGRy5-+)&)qdi@;1%A}>V(T1}GUYURBVj{cm z!nl8Ixv`#Pd`ie*Ik&oH680--dwv+^x6uKNYlEi4hgkH*-&6@ruj>x<0;QmLv!?y~?# zM+)8p3oFGX!LjU+8R^R^m@Jd)OSkTv_WO{^z03hVx3J_l(_@R{3zw_hDs9|==8O*` za>(um_d8KO_rJ>pp$FeHwLIG#a-uG6>0C80Ululm|$ z4Am=nuZSJ(&GL-ukV`Qr7ortrjncurtJretiFxL9phl1%_@pWu8F=jODnXLV#SsSY zr7h;6apB&AwNETZqGuBfsVa_tcu}{?gx^kc`})%6C_fwPib*O+K0U8GzW#LlMWDux zU^f)nrxEZm{-7_lLB451(8eZx?1P^UoU=VX-45K8+nHBcI18-PeY~NpS(|NT$zI+* z9KT-eu=r=L$Be0{-dKGJuvZ*ijeS|BYSr|VVj2{d?A!oPAV{YzqYh&`+fhWMVfA)> zn^TUT6LmY{d=Dttt@+yGr+L%_Wype~863Rl2_ z8?oox=$hX%&0qkL$*}hb+uVLc1RdxDmEnRr~_TeJ0*s&*x}a$p>5>tgp~j$ zE=6s(dpfC_pS?fy_Wnbf>-yltoIq2+m5m}U5PuN|TwUNq-rN%#d#p}MlAzxYbt&I( zl6R!t0Vuw8DH>b=nskXRb=-QnIpcXSn#l!oiQl!q0RskG5VG z)1&WaaeT%azP;2{P4>(n?xTLUv9s&FBZ<8F(;PM$f0)KQNOqxChC1Z-da;zvq@ zvD@oLYBX0>@M_on<--$R7BWuLBNXST1xyzwsc~wN`y^JY=}W?_-k-JzOcD3ZXgI1; z^XUxW{ZT!RjACT0dl`)SQ)8Dxp`+Y&llMpcJI(5ea3P;oDDi?eZHH(Y=}Y z=7b$;!8Vqx8XE`imt{|;PsJ`rT+i~vqI$l$0ky~IF#UFaVk)mifO8W>rZo$_J5m2= zh|#y4FH55~>5|&7hWc5VSEti1V?v-Ue2LX6Fyn`*{slw20j>TlvhFk#uJj6xv*=7T z2>~ZkJw4N04Sleh^|RPN)|nNkO>0v!@y&CxNo0Nmwr4t)v^nx?rcq*P4C#IKPpTBX z>*WWTNCt>v{4Z5ERxvuOnAq*`74it4&?HFB@Zr^h`63eM^y3f8s$q{uE6Nu8El>#^ zmAR+^pCQ>INeO&71jgDpyVTWj#xWOKZ6FR~^9!vTbH-#qczl~qK_g=)2akb`#RRnIyo#p)QLv@ zgjwp5KX`S$>qb1hKSMb%vRy`Djg2xLZwUZy3jf9v9nDNXdQy&ONacp!uE?08&3l$l9M_e=sQSr?zYpWxL} z0N&({=}Ylt$#E=hIsbtKBd}egu^gSs%BqJmOLA9Z!bSQC&zHzo0aZHAM%Man8MJJ> z_;+)_QV@%BWw}K>ri(Z(WKDBxe1(1^(6gc-a5U3kS;V-7aix+N;H)@o4N}tw$CI#m zjbH3G*#0EyT*MkK%UyLqZTt;wNnmXGgJE43)~{@hQ(twpAJ* z&-ZCtTzt)Ay)-W2%G6Yg-U+1R=sn-Sin4$5cutO`t~`pxh5L-o+E|+wQrflO1THX; zH{CLeSRf z?EH%AuK$~!n6pUxXORrz#h+C+DigJ-4~weEDI4G zkpV@rNg%{QSqsjwBk|y;wNq?&2lW-7w6IAeY(woQlQRnfWXLo5JG?o@ky%szL|jmt zB2A%eDYTF zc=hMz1TL^&4czH3CtxZwu!LN7iE#uaLY0L}lYlhxIIf^&L>J zCc}9J;QNgLc37TX-9pn}4 z=h>z2H9*7lkQvy)emA{UJ-LHRGWT%!L=mEZ4#dF~^2g3plQ*aiJ5w0XsE=6NJ3RnoQ&^yM`le$se zN#ip@Jb!@ggKF#pDp?-+8NFA7S>v#TU$$e^>FH^_(1#I`8X<_5K!(d&Dh*+YeADnM zDdH%o*w3foxje{7vj@ zN=Gts889C=)EY8`*$DMG80@Vdf{&6^!)c?;HkQ;8=?Yf|O4Z;mTTaPS>U9fg4^wXMIe#fK7i_G%9spRrFp55k~;ZG9!Afx!|L6?%s^($M~P9LoIgP0zl z_~X?G_G-zx(YPy+&V+rmVyXi%h11{Ex4pL`xd_6^PS(EE>?}ld^GOL%8mUFCEFyyq%`!%J0u%5TlL7L~S5wvIY@ zl$?pPG73NymF~^;O|$?h!2x9@h`$=K5&1vP6hw9P+jM@lyLu)~f^#?D)!g#1=Js-! zeHBp_LkiqCObes(p1t}#$7BLdGjPyNUd0ughD0Rn#`F8oD)i}kMh9W^1)=Y6Fhh?* z=%pKI{nr%7A*n`l$A+AWzZ=iG1^3xz`_>O0DTHP0IcPlcBoX(tElcbr(y4G#*ttOJ zk=#>b7)|5=@#$(xkUI!KrgnvbT#AK`@%Um4humP)YNpr+YE^B)?>(v6S_i%QLGXcD z;9aSu{gc?3o=P=fnLcgEDO**y>W+0xf9uR$c1KqkcDc|(T9w5>_vc^2{pGGbN7mHv zkgLRO%+$>0+1-B34WWFC_jqCw=IVMHLS#r>sUcvBMvryKRBVL_J3nk4Ju~RSoW^O_ zUiO&$tZKk2p3bH3xo^1oRI&W-a6_o=J1S72a4l|2RzmRHiOFZJN%J6^-bRq$WI7E< zB(qn|*Igc^P@gYL&)lL}DjO)}%to&BfXj*#E8bS9A*|SK&5FAAcdnsip_YqH-R`+y z+%E0hDNI;IRv9b8%RKBTp^G^fIq1_q!%2nu&qYE$mF_EV{LKLhso^PA8JJXG@s}8W z*qB;?!@hMBJFd=N)1xAZD47Ia1P~c^b8cTO`6Vn|%$eGdNoo{*s*e@k&#B}&EMdpH6rD)*x@05K zjv9T<$#8~rHp9^&QUj*hAw1E?>PS4Jfx!Y!Pf>b)LHvQH^!u1!1GE*yW1qj$slgo@ zzWNX)AW*P;Z5drr{*I?OjNKPZ+4Fdxc)!C&Z$tzYR(Kfl1D_IKF>xIDB<`znU*e`D zmXCqORr(!6kdl>#8QIj3hm#^QGPNwmf318}zRCf%0^yC*vNpDRDNj`ZdkMUL%QJ$Y z&)GTVDgdu%tMv);v&2MAnE!-SRk$Qne;Vg}AseAiRXFUVwTU)%SWiR1g*H|$OLg;e zPQlW|UvZzP?41%1a}C4h?->|8+B<7Dyx3zdWBT*zJV>g(+A5QyK-?(&P#l0Tf%H&epk*WEVFqPcnt(qYjwTij4WYE9sdHlcEfHa!4|z%ikNg zcB&HuhivRdxsQ?9o)5yC;{1}i(Yx0_!&3uAy#NeWcPZF(5ph1IW-!d0-_HH>ye z0pPU>eqCYu#7ynJy?JX=*&Qoe7yvX6xp`mIQ6U6%bjh-_E=wjIwUoiZfK<4IhV|7! zpu*25F`8=T(ZIZ`D>c^f{0WlWH``Klm6Ak%!U&e$#9 zdg??cZmp${eqM+qz?%?yah^&h3+!5!he)NZui+{hc3i`TGc5z4^2qpYGtLb zplIRPO0!r>i9#g4vzY7e6-jE8FUCm>1lDb?V`oq)a6ggL;rScbQpQwfN7Kl{0l%1V zsAtPkvzw6FT3982vBb9vOw7x5o}y(`2wpt8<1w|$rj-_bZl_>&nZgcNDD-c>e;(}= z(vih`h|8@)S9@6vQ5{E|3VmF3>ZNE2Tj@2YV7_T*dbNrpwzrJAnddFAioP5e6VD z?bg^MWD~+LKr>=$QhX7w>N&TBq&99J5T!JhwdPe=buk@lZ^ng-(g7lx1!fO_b>|DF zWLb+yO@&2aT>q&5{Jqnf;ix`oOI251QPC%UG60P^bZ~VjX6?Gvn9Ao6ZO{|1y*_KAz*-+Ah@CXKOkNUQ%`srj`CKl48nS9O_ zR!Hv*8>e1#%y5cWppXEa9P`vV`N{UX7=D{bX|JUqHh{+y+`7L^ z{~?N!7E1O2yJXrJP0&%SQIQ5`&A@1_iUhq|3ipb0F^f3HKCQQKch*~px%jpWQ%-KQ!U$`+;c>djCKJ7z!1{(9(t zFChd#W-AA>4BHkFHtt!;nFUXxm2BGaXwbHf~irFs13Cb*PJ0|YoLZJ2mc=?-Aulaca*^x9h$54 znNn~i$hSPr-<~wgtZR1HDYW97h+adOXtcM}W@0J^M+f28d zjRU`VfJGUmFnB=b+6?5Ufv_KRBx4ISqv}@!n)0ms);)aJV3YS>k+RlobIA!d}V9DlqeNS>(Gw!RRy19Y)P24~Z1$SH;t#Dt_ntV_k&MaY7HjNrk_< zSLN~eg(v^#9IZ=0QJw@X+9vS`E+mIqxP!U=QFJlC9 ze@=0&l@qTU;q2alWKUO+%w*1WW^#WSJ_;XHQ!yOCjddR{2Rl(eCQngq5W^ix9WU&* z@(RQ9`$i-grk?*$Up2LWqQ{;MlT4uo!fu4`(!^qcYs0oLli1{hHk9Ut<0QEKgBGx{ z$KGCrmdvgcyxlOS;Ii)`?F{f9I2gd-2U# zTY`)C*j2xgpS1Op1vyHaj4(=XgKw!?EKbXVqD+bN#tz4nI#7Ntak2_Z^Cu;xJBO5v z!{(5|4u@7nhyoPx9w4cCKnq z-e2WvWKo&9K_xE=!qlZWU;aU2ToQ=@__?yb-zHwgu@l!JC%(_zKcv#ylrDzdSMJ2k z0ci@u!xqF;hA=!>Dg}{*?9b^IP=EpBG^)@Oov+9pv~L}D72<6YZ&J)35~(~gS|c6I z^*lzn-BIB0aw%OrS)+(b}Dr;4j*N@cs#Z^(3&<)7V6Q zU7DzRpwh6h7g5$d#W>|pK&C|qhYFOTKtQ0SUzJp60eYkZCc-f|QfSMKp;ZgXpvYB} zMI>xsivNhlP*)Tu&iwKGvmM>qpym2^O`qu@yYnc_{fNM|=^yWMrrD3?nJqtZKR1#F zNt1WhV4697KGCUG7}HUKGigFPMRMn&#@#Y_b}E=#^(x+WKf{?aoY_7d9PYcNF15Xm32AxFR8dKan|N}F8KjEK)tQ_2x%EK59?wCD=``;arX(zbJs7{I4U2{Epw>H7oG6dggxkV4)QXIlo(HaB517fH@x9evMOH?BijNEjOeg zY8sCbG`Y0qWsoDvopd9x&4;OanFI(YQGTH>gX3`u0?Byb2KD5mLtQ_dgZ1SuqK+n| zD?Jf*HcNfIv-pPZ?bIgz07We&OX#n$&r0;ov~x;K#9$5Bk1r7i>R~E~zYd~@3wIX1 znwK>7?7<)PNyr3S>+J&(vBDr|ldbc!^R6>Wv1ubg$Kz?o`)R_6NrW4attl(jx@Jiz zz}A(6Sm0GVij{^}q*Inn5cyT+mxmz`$gi{5|l9}ME18 zbZnoBX2c)#9>}p#XvzJ+tE+aU&(*W|ppijz;a{7GR(;@t&ER>(X!D>+6pyk8 zrze*{2e2|sgZhZw26RaBFQmvLpdhvnV~N!x4xZGPg^#pvFC8V@#h(>R1<~dMU1rTv zHjk*1sAlSt0U8LNYVx^s1uMC(fqV5b=7knb;YVV^Si4Ulu`P2_gc-{1#VYA|RN72@ z&d%a$BJz~zu*rXtl~a84kLyQBio){DpVk~a&>X*RY<&2ccOmHn=z#v)eh<7j$x)d= zH!UP8koOOC6BAo{7iuSZTYI?VAh4R5S?&BLgPr3x5|g+n_D5TilF-_81_B7X?P(gX z7kBUP08)rP)!s?6B#>h7_zdjQ>$8`T!ayEO=}KCKHGO!0+>hIKoe@;N_+qIuoP zD^g>NOU$8pgJ_ti-$>^oQ*s;SNX4tP zI_xo#K9D!8Pkx_~kH2NGE#5f~NaLTUh!ZfCYSVq-E251sFoGiWyPL};gyMRDOl#K* zcBK7YXM8^wu;2hs>;}*wR4FY9XfNf$o1Soeo|@DBQe!d3!dW^~8`6nMw^$bWFN3+l zwCHo1er|Z^==zV}@?wORCM|@5H=-3q{$ATJTQuJgT)aXQr@ChGGNfTNA)x;K!~cuK z!j*}K;njh0#WcWXjQaSuwfAiy?|l*6+iNn`1~Qn1s^i3-ob7Y>K26C%GFhOKeNmzL zV&mKKcAykpy6k&5C*6Br2hekrOQXXqkxy&6sk(7$y?=oD>Y$Q5%+wI9V67NDyM6 zYPy_$h{WT*y_$(f&Yj~7nk7q4P5;`OEOw<6k3AhSwRZ{+gHrX`+DiwtA~l1V8@Y8){6bLSf082?)~MD(xMdS8xe)| z%W-?$m#Mcf-gaV3DdzTrM6Ur8RAIPQGGcC2epY@yrLZS)39|NrH}L+o<)TIqVi}7 zo2fweSE#_kbG?@_Ot9{%D+8r8p7YCXLxJK9^JW~X;U!4yohL|865{6lh6LkdHhf13 zBngo+hyB*$Q9^4Ci5W3r0A*baIvrSAq~_3*#prR!mdt)B7^?w=X#0gRrod2juNGTA z4lPl&)fNjHTHpM#{DuBmE8G?iasCi$?yIOd88HDD%4i2%URi8j-B#y79v(vfAett`EKk#mf zFi7H!!0z@gN#1dyIihz>EWG=Y6PO}~H@atqTc6~eo<8)Mq@Hpra{k;rd5)5@2*6&Z zm^-|PTMLmi7De}iH5DU3GVbGsc&8=+*A@t<+tNd34332UT!@q}SE4Q!3Q!BoFQ)BJ z5(7rj`YdeDE9V@YAf33Trn?VIUrbncHG>~>y?1s~FvuB;;5VQqksDO^}=-#gau$@ON z{fGov7Z+s(=rBy_H>t%NO}pr8B_*GAzdk7?E)-6*bvc=L34uNT3i0FU@n)ZiV^8Pyaxxr(C!g zi3GDmY14y>JV3+mD%xx$If^P178oMNKA*Rgnas??z^V(6+pWUww0iq~YO2OQRO^u{ z6OBhKnOF=OF-K`cFeybr+w$#5vt z0xvtRvhFbCLz9DW7+d%={&RR0Eq24O)eBP{3;^2>* zbW2r8i78haH*;rw4e9r+>5W6gHgO5pWr~nM)549y`fO(qpF^Y)z00yim!fM%~W zDJeA^mEpygZNrYPoESsRp}enPdvD;WDmIYCBQ(KXhL|M$x?;_8kgg`pw0ltOxp;(2 z^@!rNL^3{^cxamlgC73;qv(MCol+Ry@`X%f7z&^QZ)n@xl0-Y;hC)y`EG%A^lc@zV zy-M&`d;rkUPR-Q|le3W5vHj<$4b{Q((?&s{mHW~2Z2NDnQl|N=3J1HHt|CQq9M(PY zThaZDdd~|bWLfE)7ROGLxh1Eh4O8CKJCamo<&QSX$mEFL$6@Bz-$nl-#^HiCG7D(} zjDNO(ZM)C#EdMx4JhZ-=h|I8JXy^l}o9+r!2;uIMi?`mMX_#0_Gbz!mq{+vXd!8Uz zT%(_#D;T_BpHs<#JhkHfEVYNq{VQJG_Kjs5T*}}Y`Ld4C zifYxW$xs&57yfy10kf#>p{+%5cMj8xF1J0iMmdP%?K~2D;>UVXp(P^5^g@XSHFHYU zt@7v%Gnh#XUBSj#qi=$a&EVtW9=2ENI?HeGXft(Z0GauYwCfH*PJT(gv5eC25LUrx zdG8xW40Ti{=&vLJceHKsP&m1HxacX^jC@qnx@AkagN@;ny_;A_Ix6Aj3#?>bp$gfy z!G1PU{KtMiD!%c^;>zIPKwhL@^A4O%K#)6{K9?9ok=|}R{Kx1wm*BX7f&5U&h<2MD z_PeGv*T%pq$NdC1w&2{%Vh6YPQ8G=n%)mk3e;f07;_Hz0v-JdZZ_+n!jF3JK;Ubok zl0WZPMWk;si951^8DwQ#eFMQX##2$>vJPx#Wi$c8`_b&LQFQ$Zo*R0PL}0Oa4}3Ui z9Lwh8+oopyoyB+tT$aO=c=gfwGr7xkgL9`fwQSQG3|6n5od?JCL>vyqL4_uDpB@c8 zw;z7a+JWXNCW-b}gV2RewmtncB+Bx9li=P5)ozYSSK2wmg*X`5%xyaNB0s-we>Ea; zNmMby+0C=Ljv|))?dEcgq396k0%Yhp`57(+^9cuvX{9~}AsRuz9cHn)^H0xLO&>}7 zv$D$L%Iq6Kd@e)}WCnUPQ!SP9Y3pSLP;=SklG-8wJ!Pi^ez-Mv_2hE8{{;;go$2Bn0sk=iUUtoLmwf=cq9T7U(V73W2 zuE=aCa;o;eHW>O?;PsPUlt_0cQ1Wu0m!(;g8^EY9Q~Eg0OIog?!#Wr>$}AA_&VVdsMq=y{{HV z?L$ad^JK&h{EEVNV0S!t3E0bQc9iKMRnOr%WyQlt5q)B264!C+1??4DrEgEc@EVNr&B(zD8Mm`#h%eUBdqPjZyj}%G&#LPOk8wv7 z8}YS7&ikG8DRvQ9vSy|E#TS}h(hSoeMxA@GB z++Y_U&i+N~w*R_K+2O=L{%gcJwRSMIa`2s#Pt$UAD=-ZXcOb(mH)Em;i%7y_GeD4{ z9OlSXZLDk0Kznkkb9vy#_s&DY7)QeRv?~37zBrOFT|qa0nVVy8engTiiClcWMH{&k zGbp!g(59DG1yiyYjCi$XE%7hxgYp8#$_ypb1QX-c#V=LxE}6)X1NZwDgb2*qmiR}8;i*txVB&=oMl6P@m}pLEf*?*t-aR$R>pKIm_I z#sCr3oH_aNZ#rhan(fr37k0&Gvu&B)*24;PgfgYQTt&H^6`>3J+C5$4MG+L2)H4AxNO07Xx!pq8G+dD^nnnK2l&I)_ELTp5K8Z=@+!ruMtX8Ti1*cKax?Vh}V>c$68m8p&S`xcDQz68S?Cixbl!`@?hwlqFmyy7!`6GDcJo6w9K7C^0knG6laPW9s=^sF_4HB;4r$P2=*>ol)4w!qM2q325Mi48DQ z2h{*6BB8bVWhelq&VJg>^LatKxU1N_VX_teNFa|tD}R=??FUs|xJ%UKf-sGN{*k)w zyczuqL6LfU+YUaS^K^wTD=VedCp8H(1`4Ps%m9twKbvUFoKS2nL_`FOy_yx}5(v6><4aqGuk)S3k3)LAFk5hOmNHT}4T58>}#OU#GQv3q7n6f2~(-Ucn_zEZ$EQ~;s zl@@AzwIWbx%>@0=BXjDjtOiJ|t`skCQ05p*Yv^xDWS=v$s3mA#yUIM$&VHhE*Y5`CntSQ#D`6d*!|lCY)~M#Y${Fay0* zTP3^fI3?zq>ynjZ5E2xvq22@|_-}n0Fz2K$L>0cPj04deJ^oQ2&FH`3-SY!^%qvpm zgc3as>NBX~9J^}=YYfF$ip?1aT@oHqAud-4Z<&^6BaSK)V&%L}Bxu*wAvkHTozWzD z4+9E;*ovgEJGA7O6g2D>=~D!7Wu;_Vrg?e!btOa0ssdx#w12Qpi(R1wH(Wd*5d{4` z*eIyAf0(-Bq~kt_D}6>Fm=7)rS)hA3KzDS%D(aK|y|67EhoEhJWy|MS2LE~x=fv^D z<$G(F%Q^TBt2E5@lR0O(I1x-UJtV?hyk;oV9ro+T?IR;1?4$-BDEuNk3~hHn`opHv zQh;T&n}o3;Lb>E3KE|91sW!F-N`ZPwVET=&sB%^5fWf!WGkZgl>Aq^?$Rs3!5_;6B zd73%8wxTsUE^=ErRVZvaqJd8$m39LF@aQD_9TJ@;tr0bcj%_f>+CT>S)XcA+ zgS>HfpF0Oros{>y^BA}Uq%>~3rYO(i*^+fxhK-CPxJCh@D`Nc$nnvvOw8{_ol4Yn$ zmEY8gNY@|}{wfz11#_W_|BbPqR^H3%K^12>Zm?gXU#M-xAQ*wv3 z4LK}eHCXmm1W0DD!Vvc7kH=8!HO{cAwkH`+q*PJR){M3LzA#9caA`rwhTZDBV9VgO zC0t?Ic?9)b=|^9ysqW64z>r}HsU1dyD*eMWAExUknph5&IkwTC+9#>JLfPQVd{7{%Xf$DglA*-WpZ5RMs#CyxaM=oj6d)!|*AxDf%A z&J_kn-T-CqlR$Jl-HOnmHq|F~K6Y@u!a@T@Of{O#Mll%P_rJXcKFV39>L-G3Vhp_h z?<*ZP!1;k6*|uo%|IXjYfev`Kmlc-R?hQ$w{_rnJdOhMIAG6G|zkVdJzSMSeiTfMP z)&4tuJQVMBP0!<-WIilBmd1TkIQ;P22oHwOyR}I_+|b}cWmo&>WxURoi`;Ca_+Q@e z&Ul0TY1oCC84{i~ZZN8_S6^im;S(rcS_-wcejpDA*sfB|h4r zAX|=N(p9rXXYkzO-ZdX1ZzbX8xnFP8u5CBeKZv8>ds6vYd`W?I1pK@4aq1hhk`S-9&No|86z|&I zHjQGmhR9!-mZD$R&#J0vAvj!gfhPWzmuJrhNW{aB{5Q93DODlEewIv__~bqFl|;ig z)?GR`-*{mKL*B`0118G5zuF!=a>GMWw9c(_2&7$enPB|=g9NU@2J02NjKdDhJ^nR_ z)n;qBt{UwC9tR%gPP>9@4@+wc0~)t=ug^^%=|v9%bXv!(FRIq^??DdW!gey{zSCOo z8-l4*vOKt^Bq`hPT?bx>6Q+lHmPLmKJsUb=hfMx?t-K)O4Z z?p#Xg5=4+iy1N9HF29r#g1(2}%=y00f7zFK9%~iO7Za3XJ`ep+hO8eaq@)4{^Er(rJi`Q{q##D8)hJ-R;FV$$X-d4 z{Jjc=BqlJ9m|+15k-s1IRk}{(3=mu#)U)7zeZKH`pe@}{mOpFy zDiv_IYJ*oYQd@BUa^e312CfG{ino2jIBYIKob=uNv<-;Jsem9#Ta9@~;fS{c8C+!9 zi;(~E5CL2CR1S%=P&_GeYQsLie@&AG+`9}JH0gkYyHB8lz8)G@-4L;Ye>#%yyXj)xK#RL(JAD%fmzKq zuve=E*Gk8;FcWy+#QO%`P|{f?%;or5Tba84;lD2odBMl~|1r#Kk%YZrP)ad{NQKJB zpt3;)0nO4=LZDaOc-$m<#q52Q>;lwq^d)O2P^M}DK!aD&T=Ed@fV-{tXWgd}%T5LN zE*3Whz_f8tzQy_&2}-t39rUWC;krDIhMyq-u2e{F2mEc=i1?3|hr_}%0f#7u4<}&O zTopQRGvAZrqzQt&5;f1&=S$#HO2zWp6t*%oArh#_ju_zCo$72w({$PKdVBnPzLP$0 z9+C(gl#vko&P4rpA&MsEkxbB+YvJ)ll5Vb8HJ(BF^NC1{s-VWHb8H(3k5CFZQbpPH zJXi~SeA$shu&G#?j+vPtquv^5ERvw%wew3ZC_y=bUF?%Xv1Oz7x7x9J#M}k``lfsn zdL>mcrcAwO3^1|taOSI%Wb`0J&c9Be!MHaXmq+9C38B=6f`;>6wyDH;3CEr_C;Ng* z$|V4l5iN5^*q=vFWn)r}KN_?%T>>&7*zqqF8lc`z{8XfyQRXs1K+c2l0yVH>^#i{q z@$ZJI>^c@#vP|V^7nehLMjl+i0jpvB)RE5B5Y;4BJX;1kk5~U$A^J%CqCHD8GoB{& z-m0S#6B`ee7zGFQDewK3pHcM%aLJW@YGmS&@;dT-)}3;Nea{59zxCOU5yl+8#P%=i zyds((P(ghEopaK*8OQMfp+|TKa)X>ZB^GFQkHf%xT-5VJZ#?q3S%4ICZtO zaHZ}owRANi^t2g|h`a}O`@YjcQ|*cH>~V(YrrB{nm~?Hp_dWyWXocuz8;8E}<1fG0 z&j5|t+Y(3gxbBNCD$C=SNlIg38QGF-L@cbiSo~-A-o5Lnll`nt5*Hn8qx0k*j~4qe z+`vjcdJ z!^3TEJ;(gv7an2ae}OgPL?Qlr%|}w5AVGYE{Bk?Y=&_r4AI^>cq}=kak_m}B5`ez? zw&v++LlY?*&a43AP0A@HVSKlQiRbS9EP|KJuGmgEl1792_e1hDG`C{^i^dQ97qe2+TGCAB9j+Q|Q(pV5W+lQ%WD%GK3CWy26 z_P6$;F-tuhR(4U1rBH};+0{u6p`bRCAwq- z|6d3($hT%&*c1XGNUbZH{eqz9qKnP}za?8rW2rH=a??Y&$wTj~_I$+s=al5_bJ}4B z<)kUkHy3(Wv|DmP)N4z}$rMZ_Gg3t<2jijYNK&SojoMeDbW;$cS+AF9t^062#$xcf z*uU)`TGL~ppEishXH#WS2x^T$AY93r^G)q5;g#^1xWiLk(#N%kr=YWs4)z3MMQbzr zLB!Ut$BE>zTUd>E_t`7cO%J-UtrX8!Me+aMlV;TPoa)t<%G55En~W^r<4|DG#t;^P z6KwJJ#v=I%v=Ti!yWK|3IQ!=qSj9^rnU6|WJGmq7hS3+j@tCl+dX(=oHLf7Oc$H}fS;iHmj3RQZ{)?YT$X$0lX;@&zB zZ3`MbTb3*&CcU|W=lAQe4(EG2$wVEiR2Q*cfkWwxVg87Nu0%_ zlan45en65+mS;f3dJVkD^m-n--}WIzn3}dy^!zdt8iFFvb^SI&CB|W&!+n6tv>vv& zJJn%wPS>Wa;t-&i^(@}z7%|-w_vOhGbDFj>JAQ(=SqWh_>GEDt$-mN{C6?4>vto-2 z{Gc;3<2Xk?RlnIumPojU98(Ie1%bF^XXo&l$XnZ4xU*a!0ZM-~eCHcH=d1*?6vMQg zlVH)uc{k&WsHoP?dReK6M4rQ6h&f199t=B7*?4HRV=t@+G^9BMm80}JCsPd#)x?bW zp`sGARPb?ZTM zLl5VHky`N;_C_2rAdKkl<6|_VTN|GO8Q*pdGLSXgWep%8uMvIThz zk&4c-{7|w2Ux-EoDOwey55UGi6~x+#i%3Ub7GSqbQAi*6#yCQ1sE_7hCYgI&C!71jXlN=nN=~iiOweXDT?hoGf8Vq%HY5N@v6ZV2N*o zjcMdjp%j%$F`TsCKHO$gxRNz^KA$3CGzP2qCl2j_vgHfs#?#S6K13YuWx4Fn_BJc{ z+pH1E<@**{4h8+hZjVYNlDfI3j|<1U*rn808(%6Jtd|K|rH-f>Vz}hJ5`>B6ovczJ ze<3NE=3DNs1A{O1D@x{%M}_H|^YF>ue@$`{9x=|LPw%MDA>z;^g>qizFk#nG7Wn

F3ZR;8@jVg6`P!el0R8#BIo4FDAxb*bACnoANxhA^HZrYj zX+=2gUtqq8;{vmWZ!zah!Xjs0LF1-uU9S5$Sv4A+8ihw*Ov)`+dyuzLQL%|>o zCO5eq88Hb?i%+qdOgJfh6LMqcH#mH|EIBgKqt?zmL6=yza#jnV_A7N8Ouu*sh%j{6 z$sd~uU*~v7t3yEy{N%+*tm2M72@Xs3mb8?$qi1C9_+cD_WNX$zdaH^Ah>`7SL}vW-4pzQoJ1s&-V-J zP9!|Ei=vLZ-x-ZV&*|z$A)BhXV&MALYTp4ay{zT=cJV!8-nmwO&qR0bY7N7US+LJFgH)LbB5uX_Ozj)!lT##5jC-p!tw_!k_tK~ zX;t${87EIJ=$Gw5RBTL^Kq~2U{jUxB2W_g)ip4YcG2gXXL=y!YJjjX+(K~pR9Z(AO z3_4WSs)Zp54A10t;649lOex5oR;l8%H$xDvwOlaFhwhS#)2Vj!}@e~wRdrGbvCU`yk8sha1D=F#$Yj1Oati` zY7(4UzRd|tH2&-ChsnC!nCm6 zM2RrfkwACedrKI;C{B8bL;6muR5({vW7;9U5~;iqX0CpxoEiP8_?t)`&&u+}K>_U-4UF*Y7r} zvd#Ey(pI(W?+V~y6mb44d^68ng=iHIj&}k8SpA3YUOdk919|3_>)Edn`=!j<M-)0z~$ zcf@&v2m0G{b!PZp99S_u`Qa`D4gKYxg-t|4>^^60v!40M*J>Qw(!3nxXc1P1R)Ihn z2l14wt$Pu0B^ zahP@D8r9TC2+j}J-={pb)^@*MyB9vt;UWb8P}&T;ec-3dT3y@E&j0m=-&>X1#IQ`^ z2{6asDJo-1{TN51s8V>9i_LnhOdHZpuhcvkX9-e8RRdAd)VaqG@ArY3zZq#gEWZ=O zOEs6!Z$^b?%lISWGsu<;3G|p zmcW#iFV=GZTAM>(o5H{06QwHVrx)BTP9n0q-kN(2le`GH7_}Vs>KlLKBt5+I(q*P^ z1py&r&TWpsDy-0QOOYw>Q0;|%mA=WJ@f5-pw^srdD{^n^#zF-!2(~;){q_qJdaGA{H^CQ zM3E)T_sD{yv1iOtWG655(T|FJz}&jEa&_mXP+NlNm_rIWupuqe@xzbSMm-^@rpVzr z+A1usK}op4ef6+p#_VwZmR#tgdJd=8cG9(}&|$mK%RN{PIY(-oVk3UQT5f>*7Dl8s zNs+Ex|MTy!=4ESp@=Cf4^**%;^ZCud;0qQXD{9UfxD}E~fbJyVCW>PPsgmJaDGB$Wo@nSIi88G{}dc2A& zYay=r)Cr7mSvg&^ou>P0Yb}1&8|5!g+DB_t*USVelf+g&PFnbOE=gtsEM3bnN13aU zzBJ#W6u*Nv2`Du9c$}kY-_4+FIe$#w{!qh~5et3VqwN}YyB7R| z@2{BL4nM81J^5NFRjiHgnm(za-VdEWgDx&89-R zF-P!%caQ?Rq>)4ic_Sg{o7DyJ1)T-YwiJ-QD*>L2j~E5sx+Sum`%x^io&D(^7MzC_ zTjn2W;=8{yu+DGt8ET*>qwcD6_rQK}WH@s(n+v(H3VZk~0EQ{8o?QKQ z`A6dMp{j}xD?o{&u-C*kWT}nTKI`**<9Vm}@2-B*zW||(TG=rGD5Fa!W)%1oqCfO7 zAZ$9B_-A7jIGz3dh?2N0wY*wX53VcEYcrvp?wt8$57X+fkK)`u5BMBC8DQhd zI@--h=pS8VgyCI@2^WSa!ss!~%{0WtwuFaA^3bvL-mVcX<^nLjQ}Yt0@!_VW4ySED z$4_7ok$UfnKXrJ@Dn#GhO7}wcc9N4n2dbePhf1tYyL^K1n^ZHCB0-nc8wwoSOB_aQ zNi)~fx?W6pH<}+?szohKWuHo`-~BDO^wZ}ivj)E#F|uASsb;qST>l+P{NVd70o?Ka z@dSF4C14OgRJGM=*|46yw-z(6P4eAWIBevg{4YSZ8-1O@pk3Y|mFGK!W^#W%u0a3r z-_eLpIK~)y=8~>e)qi<(lxZSgPBeu--FP?Uxqw_>J#$(w8c$5g6* zU}Sl*9IfS`nG)JjgID6hI%3eSk(Z0b-97LZ%^?fACc>BwqmSU}yZqKi2f>3Dd8_?Y zCr`<&LQo5=@W9()DQI2FB>_N`k4+TcrJK)q)sRejY*)56+cofTrzOH#o_9Jg)Xb?n zC)(rKEX0SF7daa<3x+m-XW1}=d@_>swL_-<4i-n0Q3iNUYx-G;&g7EUj|KWaYviH@ ze=5MhD-1LiQ4Qq2+Qe>C&!r5t=8C*;$+yNa1LoH_QE4ozFnc+>%8_0}ZuzQ$PUc6l zO}8gj3h6dmJ9^R-)PDf~ewQpT9p9DgElQii-K63ttaR~Nl5xtn z+q2cF^!4ITc76Z~nVH!V2m2RbWo(3&P4BrHzOu8DJj}DfXW^||P+=lf{%VxFja!tP z)nafK^W_mu`hM_pmo}EF+qVa@gTl96cq3E6Sj9~Z4FBQ>#uEnqV{eX__&`NLU@U{q zx(E#UaFetR9ZPz*%)H))PO&cH|NNW7u~FZbIH{K4U_IrihI4B&nFcrZtF16^gVg7> z4ZfKu;wOkQg?*mZb|&tpvmOw$wC=%1LBlHl47*l^nhv)7_ZwNueH)m0=d2QQUn&$0 zl5e=O>L?m+6KgL)(L{x@ov(O#LjWY`RaRHHn>s?eFzBI1_{)F&fAm}0c2&b_qS#Lc zvmye#>>*CHV*Uzr6~YEjY@-3$38aW40W_#Z$INmH!baG3&fXf~j{zVv$_Q=7uc=2; zx*YY0xkbruD`Q3NP1i0K7My&=2-gmxF#NbgndcWBLb&twG)2ASMMH{^ZplL62`mFN zIPC}g=a6ePA|HO*uj8v`kKk$itY{rpo2dYQw)(&(-2-W+{9NBE1X*QRG>_B8<(9_k z_@>omG`K)#-l0HPP2nh{-zla3aPOT@BQysU{f`NMJvzTyAbG*@j;&y`t**BHXjn;}`NR3o zDP$%-Nvls_1hP1t=$&gXI>vdplYhmmPI%d+G;Sd3)_#5sUdpBtug~gPh=3xO^|6?+ z0YL=z=Qx7Y$4EE+P{@422FJ~50+R>qKKN8MlUQ}Mm&b;EB4c!k#({wla#{PsnB}J^UxF8z0MneI<*6P=WweC)o!5t~M4-mPxtCAAo-O!FGj4xh`WIwoDt zsDiTS>}=j_vUfnFLxG;Vo?M}uv?LXoN+{+7r!l@0I7W%_3T)$Uw?p9ZS5aU(;6KTo z?2rt+-ChO~&K?M=Le2`lbe32qof&iPWVCl(rada3%gH{Coi|fCe&z>a!VwuO)_WaZ zB!n%liL-Zx^zP`;y@98BnU8}0)%!}(rRusX8Poiy`q$cdxF8N>N#_jp9{K=947QR= zhWX{SiX72WHs5trsGeu!v409&YHu{Cej8CJGa}%X`d2f*~lfq|>bL>}^k6((?|nMH*#T>#IFa z%|fZ>Z0MylY<|{V8Y3>zSAd)K4ZIc%C%%N&*$8%^4+|9R-2?%e zOalYK+|vQdS{9hzEs;^bh)5BUYY0G;`+T)hM^|Djmc2MqXr zRLwTpw3^BcGDd{Q;V9pVCA1;ZuTE|wcC85Ba^tl=)<>1#ied(uQil$k>+!K4z=q9y zBE-7P<>MbbdqcTFx@4EiKGw=~$2oQoLisg=0_zY+KWVW9;{>a4kT0u7Dcw{|vRpQu zm)Mfl`enLk1sHiwfLOTbB9@46^wNa5ad!5Qs;2wt-(qf+<1(um>R*=Sil|PN> z3XT#mR0Mih&-%!=G7Byay!}K^MIdg%*RKL03O9JxCn_;k>iMOhI>!Kq$`6%)9HJSC&zw?$>m)YO?#i;su9;<^*c}pV}Q1VHJ2`DoRf(9 zAktecs)iJUFZMF2Yv%Y>F>$5XJWq9m(pLsFg%Qb6XUL?z^}yq;ylk=nh4cdI^>{j~ zE49E}{u|u|Gi(-u1fy6-EMlaffcO-Z$rBb;km4nyU?p=!6|2tqVTpFL4Jk1|cY9GB zKmxuRnejsfu{TmMM$A##y_UQzQdpXqR50Rxu|X6=)OZ=GN=CR@j-uPOHWHyez1Ik% z-151XuiPVJVjBt@Fo@KIfVs~pBABfzG5!!@|ASCNN$eWU zM|5t1N0NxTHKE~HU{YOBbbeXzEog)K;m!Be>##X6(T73-;s-3Zmu8RaY&9VuZE8?~BkL8?J;{?xX?cT$!Kt9G3R9qSNy8q>IS0(le|B2LZ9 zR^w%our#xvySXKL=-5& zU>VE7m;|>}*M8H`e&Cn4BH%~`kTgHKDZ7TPU8^FDLWW-dK4isHw^s&5y9PC)lY`6~xJj^~LqSEW4frc2cb|m%V3A!YA7Y1!mV6Vq@Vp-*&yOejoGA&m3Z zX3-^Z>8mCd4(u$&tJH}C(aAo)0q$&iPUN~!Ig;ERf4sXi;0KG=!L`9M^!dErd5^8j z|EB`q^_v2D{7MId##s7{W_XBbSdNY3;kT;G6dhGLf$`R0lOR@k>p6|MN@RgTDU>WT zgbLh$du`>OM`8VmK$$K|MVZ6w0!(K9(YJ=SWHViC;3!hWUmo_Ic5H~L z23QX6AjgpHxlWoCNjGX)B8KtseW~pA6n5LpGbJgkd1-5M7?B1rD$ z(#bM{(BSda>tD@pD<`eKlb6Ll#8HzeD6KgV=)*~7^|K}DKFMtym*u2flcI{7d^GNxGG`q-I%-4TU%N&5_as`Gvrc{9iKTQn zeX_BQmehJBiN0#3Nhu95A!j>%U?E8w-4D=pc|1oqQkMd?@jV9Vh?I!AEL+(yKN4rP zaE}6|&zhM%D6xD?=PrsvRb3m&qmT=-A7%2!8a!4ZKu31~IQ`h^EP3_m6C*q zQ(E5?Zlkl+`I1 zYz{Ph+UjDVq0V?J}@npP4^BHS~;JaCrF@#V$%Qx})Wge5(axeTU zo};^350_`SX~f+XB|n=|i>$G4DpKh0lmGEis1ZLJg60%NQCsWj(wXQ`Qf%BGbI@>z zO+NUx>KmuEfz2>lznfhga)z&`&&&I8DV|s^=-}I4V9b59!O$pG`FKnfALn*3kwFhy zN3Olz5g3V$Lct~Yam`gN;T3t~mmn%RS5rFtVL*^(o7oyES((Abbir!2F6!qLD~O9}r(SWeiQJUloG!jqKCJ?FpbxdQm#` zq!14Np40Ye2j3*dZZvA+wg-&jB#dxWH`6f(=%#l3YI>!{ji5+&^eajO(TO$!+X|Os zP%%0`ZjP77NN*7lL6kTXM-fd_M^Vg2d%n5F5xs3;)o=YbFL8eYxIgw(~s&nmUt_7+7x=w_rQ?J>OSC1P0*dj<(cib4mZ$X8Az~! z4{P(HE=fIe&`IXH9qNlpOd|Vz#a8@Yhz#9wYEf{*qg6?_G3Y>uqM*r%w@(V#mlYb|NZpH{#GPpM$6LK*l`n?*;o=j7{j%Go`KB5ZkLFqTCaWKcny`Yxv4fkn_s~d5lOlU1g){(Ob>d=dBu;<;IQ|IjB!; z$Gt1bJhMIS`Bdo2TfV81^mt_86zeER!BPsW$A#o62&^9S^_pU9O#yc7dPntZ6G3^p z4vlw7rn&hIB3j%2nBBA)bgguD*&+Ae05Z67Z*fH8UGEHW7&GKo;@tb+h}!1GIQ8_u zOswfwZ)G7aPbat$J!ID^9ChF1-QNqRk{^5w6oh_L7>luq>~_|lo=u4Bv`u389IuT% z$LZlGc!#;5Mu{3U>Mx9~BsA7kjmIse4UTrAF7S(%H!B)@Rx4>sO&JQf6$)`SQT}g2 zUUnDC%bXd}s3+BQ*sqY?+k%z;#)um2mdyI>YJ_%$Ac?nU&@3|b#F7^|3k_X0 zDB0!T^qw6Af#4b{Pf&i2!Q9FW*TX^#XjE0AK8i{BrJ*AS1q+(JW$jutvAHkvP>HjL z6tVMbH>QFej1i9xtnzlt^g2af;p@D%lVd5>d1L=^SV_W5vw12|4T&bAt(L)g9koK( zR=3q*0Apid!b~2v*PI_&5*tbDu;Z7Qa>FMA{h)~o*8P^AD{;(D$D~tAOMQ$ zr{52%ltd7fx6TV9SRi-h<4E4;4vMoN>Z4jCvF`YTr6(5xxGm*Ezvbi23IBp>NsL zWem+VQ}!D^f+_9Ib;Zqy!Ys3@Ip?fB=;K24qPKvm1%@UnsaY#{lvTGNv&zuD`PRj% zM%EXsm^iY3_&wfuof%Zds|GFgWUyNthDRVCl-Pe4tT4VDrJ{kqrhY!Y6)5OBd(;>K zf>^Xqz1@tupp)`tuUOL3N@UR;PO)9PgHf-*}L87cgg~=uo3%hHvCHm z)8gqagHvvS>)sinVFfb-Phs&lF}768PClu?{lx(J1u99@xMe4{cfakJwdXS8Lq=?H zA~l8P*eYs#^vq$MUf;sTqBc)sa2lSmO&Pw zu{{;2sp%1VG@Tt+vZN;q+I?w3|Jtis=(==a|d5MX>Yh`hCh}!(t6V-(;61&uh;x2@< zrZYn6YO>^xm6AR#__AaOyDQVldncO~sd_g| zb}B8f@d)Ooy8Ia@w8#6+3ma$mjPL)Yh$tl_u0W*ZNj@2#q&*LaeZ;nIJwrI zE?Z}7;zFHgs$#WQAz;#gYzLl>3BAn!d%?=b;6lKCup_XKd$T?QEC&}ABtpoS!aL2LVrQFbO7 z^@7Ub9pL^nd7Vb-yoDm1UqO?}@{R7lnqwg`rHSKzU6wI|{n3M_$H!L7~&4X5~b1zlg zy?st8!YG*)!QVnke#4KwR1;qo+Bl>P;LdwWhl??NK)!gbjTU1>^~1iF zo=kA%b(|AwuW8l}mc=C^c9n)?kggmH9Xc9|tTylI=ZRpx(^NWI(OCf2QfHWy+Y>0h zB}+U!NvrbyN`LmFw%|x?ir8N`qFcA*SEB=$Nh_nF{y#V{JrJ0{(8WCm#65EOjU5A= z`^z|&{%tm!VgH_e-g!F>MT^ds{A6Wh;plM-8n@5MqAAXbWVOk{(*;X05-}#2&rj;m zaC)1uOk=(4!(cHJOIGdG#LQDwzny97C?k-95&_Nc0y_^c@vkFmL37O9r zSoTQT=YH*~Qgyy?e@8xPfYi#OeT=y&B#ZlOxU#=M|zIxEK9@ zemWF+q5kEJ798OyexN8`XaJ)kn2M8NQvEsL2|Gc`={wz5XqWqXDAW|VN4Z@~HIq0o zcCP!V06K9Ud`JELC-B3&S+z;Cs)S~d6R5o`o_blMbe&T6T>gmB48%p0iKJc>w?rAZ zcm}S(Ctn4OS~3c6NlO1Twx85sWI$BCY#}8tyT>;>@^xUtpU34jg8>e7f(iNtp8TlA z|6&>f3;8}VOET}5ci-c!5@|_X0eh;sYTFFH4}7*{`%e%gqraZ`2p z-}en2(q(sXtpUD4@1x@PQjf=#n`iCnq;L72 zJG;JUvU>|*Y>r9DM8A>^JU$tiu=vfN@8Migj+kK?V~I9=L{bVyBG44rr|?zW z1lUk}F}=M^0&wpIpm)kO_fAL}%l99ZdgwSBhATHszz%}IQeBP_ zS=A=xkIPuVF!$1fuxt8I@Q;N>Jz*pAY}DKgdb!6J(!=uw*js%+@%2evoT*n$TuU3@ z*H`|gpaqZCUs5@JV1)X8k#by5K~?*6`DXf%4mW9RGH71atYQ_&Cjsc;j>dBUA^yqB zH*DZ$KTGQ~+}7TqIbQH(kyp?{;g5ZCq65S^?~q>!zVjx?t;+fCRY1X^)UW=(X z%yE)=KI2yIp0CW8*>KCs^lFJ%;n^}8aI@`4L!*X`jJXF3evB01$byT{$o&A9Jjn6n zZ$nL4`Lbbzv$3Q^l<5b?Hx()ZV-kF0-Wf-um@zV&;8S#C5>A%igyWGfy%S2XOCRF2 z!v^Qg&->U$Ez(yTB+6Jwy@qe(JFsVV5?0({OU0)Fzwj%a&l+n4+rE8@+yhYlF>4cM zxJ3)b-_fV=M}Ok_z6m&cP_#4-zU#UyArbfsZ$YwQPDixq&B|3pS8 z1Zi9@UMF7ovFiZx7Sw|l`xj_ZOUm)SqhPWxLXFTy@nL!$x1fDZZ6+V#qMNf^|Cb7Kgx2_2(9^OL}Z>G1J^DvWJUkpvC)LC=2M_RcqOvbPT*@9ZW z+U+s*F2y74&xYowkR(<|q}K-n5yDR-8IXvWPn!T-r z1M7K=fC}3?f9qp}da|1H^h^@}g!}AYn47j4fs-{9)QVBr?+x7vW(5ktI1!8L(=Y3o z!v5WbZ98QUvZ-ljJN*lf8%oxN&&^~g)2I#sMIV8z%*AXBR-r9$*Sd2<6fQ_bDI+FU zN9V&cVA2XX1@?GZZli{m?ceclFTa-lce3L6^VTPv)wrH_)@RKwm)Z+P@U!+;tf!o| zBcQD;ME3uJuBK#FRz6U)^KK2iTel`g<+rF)qV@z3P=rE9ZcCuIU+&>abHK2V-WFY* zWoO2-;bd`JIhn2v!|XjO1}EKH9>Bdg^GBhosJ$i2_XuX=Z|?ro4#v>V`y?J{JsP?M zA4NtbVn5SvEElb$qE-`bp^4J{@UnyVq8%#{cn+56h*VS|l)gojdMz(7W-v|UaSdn-OXA)f~>N>xv9FRq`opV5ticiq%?-dtEHaj~qwF%2u!%L%m&u?+k z;!H}odzQN6#4l_=fa2$s~M;%-J)M0jM%oS0_s#IC8f0jF{C|JG5jup z6-}Qd2a(Cfv!P~E&Jst_a;}_sbo{U$G>yub;JV_PA(2)p=AQ|L%5+uTb8JbeV)#oq z%~?K-spyUL=*hj?2DTMO<-8xcP2$)O6AJs8SnQ&Q>qtu)eH7x&Sk(M8{0P0lFg}(sB7Van&;3y&q$w^wP|$#l!e%>(O%v~0+Sd-W%-&<88#VHC#z9T1S#?tW)`K2t`z%_20pZ+H=wFbkIPT_PIY%n z?J&T?toy~wtPgRlW#RLQe6F|Wo&&c9G1RNX^=VflLRES&oV-*#I}L_`Y>BdKPNB-1 z4!+zMTp^fDICdzAz3J4OafRIjT4Kfuz*-U@nEDg&KyWYmX^|l-I5v66G4{yYqyPx&gO&2#v?aBGIHRTB(gOcU3 zavYPh<8Y|5u)TH&;F!AI5Yd(A#-uxbHIiXz$w(L~nn&C;6-7exq;z%msuXo}_Pe_Lton-y$c%aa`GEFma(T0ZAh+K zjQ=&9e0cvRRs}OHXV|{w8>@^un+)=?c)okjLC%Xjir(4l0(>m2=UA>h97SkDr z@TBcCLk+mlW8We-zqAi{QujN6vs9Bp)Rl+{tqPMh`BX-x>aXVX`<&=JYsIe%U&IplPAQ|F zj;dZeX5vVBM#-F#l`smvOFh_L;%7+UYc(?WLMi=&EjyEZ4xG;)&~)f#IsN|J|2e+p zrA`V6eOxGggi5W$&CD>jj56GS4-z|y*v8|t-l~a7WI8TI&WaZ@U@Vmv&usQfWXo62 z|E8iTFyp0R7XoO6my%*-*cT966QW}OxnEjY=JtdI?0uyB=$l$Bgu-->Wh$XkP{950XV>@nz0h)% zPFwwrhJ2o<&id?h#Oj7mpE1pOeTiLv(uBjgc(?G)y&moYr3!OWI0p!hUFz z{v+`2FhWznmd}15QZOzbMI0NsfnVeN-`9W7|Go1ZMCm_9u?|`K*c5VI&Cri9{2yq5 zRE_uct?bKN5CH7aQ)BKPIx31FlSz=ek|b*uHyK`suwWLGT@R!5{J#spApCKhcALyL z;z*|7t(r*v-TFDDcJQ7fd#v#omEScUzN+NytUR!SG0gAbAeSSF{1?A)KK@Mhy8@id zFOPN5688B%Q_b`?h2_5<%H)4>bQNKbtV{{Mu2p&TqWam%f8zoAb zX_YO&LOh|gxaeC@wh%?RU4Htf!Y#ZlU+cI3^<iCxgHTpz~S${zu>^YK#0SFoD3C zOFG#vl6M_G=y}CYAqxINmHp?ZFK^92f7GwgI2`M%9_;Tp6KCalx>(V?33onEBr`D( zLfXS3kVNVg%Fj{%;~Sycwm`@OrEzV1c zM1ZFcz>N*YQoI-5qBoNQTeHE6o|$?|PVVCK^6sNPMzlsK$3NiR^k3fo+?^?#m3Rcf zUn=LXEduDgj4*8_N|4^LdUMRm-ZA-<-K3cKJ900*$}UIfwU5wu_e$;gZU+%>RJq+L4=m)#Fj zvd-7XUH7;LuLu2c&c1NvZC|ce0lLQT6jDU|fW@*SPu6;*TZjHADAJO$KO!&8W3zc^A7|qv#4W)cMqUP}r?zRk8)|*Xm!f5m4_tbMyM*WG0Zs0I zpg-svV4S7uOM=vqtHH|cO>EOE;(qJw70&m~L%3}>$=)byf`-fi=edGvKY(UEO0@>E zF{n^|sf+wWhQ&fT?Dl!r8Ftuy%JS@8m9^22zmyiNf@JgA})_4^?z6K$qtd~EPlf7++jmf6`g z_8i6K9)LrAex*AwCPWB1MF-Y#L4R_&qq}3}PjP0X?u@q(bsb)Ob}!_Xe5de#WW8ln z8(jcyjT5xE7I$}-;_gmx*WzBF5Zv9}-GW0)ad(PK@#4i_ptQZim36L->5C4Pq8yftHEKO4KucYkhaYNnd zF&tRGBQ!fSjeced*H#m%Ac!+IMu@S8c4mT$0w2kq^iBP3aET4@cU523C;*JhiA=Fo zlm_A18Zowk8))34shPK~(ynC%oh9mM#gUc(O^c0L@N4v)#0CqdGA2KICUxC7BrYHz zifxg)-S!`_kn=q7F4(-?2lc(CFMsrN#3IGNr!n#2Db#tdJ^*>hWPjh`GxU^RU_uw$zTRsI<)XvJlMTpoKi16b;Jh7*AcHH<~#)AUt@k~=y)%+p`Y zyvsR(R@e}SMX@_T-&p&p9H#O9dF4C5joBAqm5{&wNj>kYfJC(iGpTp>0elawPIT2S zJoX>))P9T`+p!%a;1_yHjVJO0bVfsL@RP07)iY+VMzr;T>@Z_#4i5xd|Al_+> zww%l__-j6Cw}}pYS=aPIA&VNj|D2x4p*Txi(v{1F3p^~LX8%NxVmV`?EwM%NBxLN*s7TReE?<2HW2w7w6bWs60#19Xi@|+AIj0d6S z!IS)V;wUr!g@&ziGsx{tY3&JcAX}hLHoOz_<)o4i+f2{v7B6`0Br-OB!RBJRw@%!; z-+I|XFWb`mqRbK2$ncj1J$7L}y8r)ze+uAXS;DvL>4tMH5GQj>#HW<teBuCvvZ{#u?Ai7AweqWR-tFs8P*h4a{dQ>pHc$nC4VJ2OK!W={~ z2l*()lJ}B$<=3-(xo&k47)(%TPjcjj$~emEGmLEdlQ|s6l!ixru}SxXoBJYSL-(A# zoua9g-*jTj+0fT=UKv`O4Z3ASr0;#vbfJh&(0q+0m45HzwXtrGHdwFV3XQ%GWH-)Mm;{JRqGMx1D%zgJaiaDh#mSHkrn7egAG{#b;Vg zM0L?xc_SQocuts~f9b+PR^d3I0%cq3=aZ%&B~Vf_7EFuxVA5_EEeU|Lvc&@^Kx^m| zIpxXX?Z?!~iXm-FMx8|@I##?^B7~tD5{HCtUF6RD|C<*Wkh0cgA}d3d z2LpCtnHB4lGcjb!m{N6=Ul?<3kUTFxS8{j27bRBg zY_wVC2_nA4s#Q?^&`Kzw)hf7-vYRaX^H}uvw8+wSuCXo#f9je)VM=>Cz~UIyC{adpvN6g+iYUnzF$Ec zkD|Z*i(Wqe7Potxs9JXbRUV@^4KFfE!Rh3pu#q<+qsbDxAA0hrdBR{e+jx>g8ovzB zNqy<1ui^6kqy!gIF=%BB&wP_k(3%ZVx0`B3=Z0F943lt0!<2aAmI21cPqi9dnPC6GGdQ-vc{TksSmc+zDGp z=`*}U8*n*PMo9cEnXH0cC4(s^X)CoSbf_1xo2$*~wDsTX4~Ej=Hg!-UQK%N{^)*g> z7e;2=c6QKNOF4^)hUGjI&H=zX~@v}$VIC~fm_Ms{|l|!vgarKV{~pbc40n0o zE3{D~yAC|M9ws#3yMRZ(xL2H62)-$)>A2$m^Q@8m{D~4Ge_Rg8iYtC#t?6N{f$Vm# z80NgMwn3pESqh4e-k#%sok3)ON5b2&o0mVJt{~@6Ag`q2PQOQQq`nav2iX68%>O!# zwKHtsY#aJwj35}H9|=QuBz+G8C&4LrEtmXSQ(yj{69}N1SNnDz;dQt4KmSuCmOb(= zl@3YM|8vk8AgcNNyYqAJi(cMHWItPY$lQy|n~lYLS@&%p9M-M2-vduTa5~4gc;CFu ztEDga9HogHi z)lL5@R!o$D>vkCvJ;dK;Yyk@Xv6iw ztxCbK$+16V@8MGkjol=JZv`qB0nm?=ugdKknC<#$?8DzT3kx7uVyyI;pwPrK?Y3eA zY1d8(h3u++Ql>Eyhc8`#Gd<>rgw!-g}~Vxox=sysTh>RQ6UyPyR8tE59)j zabS7+^Y=1xI(hQ2&$R*L!qLV}kB^V2!J4A!q>26s8);p<@WVhqg5EIr9+8 z`yFTjFf-*JRWsqU>N0X*7a-{`V|xuas1HU#a_LAk^|1m}`{uyuyn1-yb(QFl+EgiJ z#;pr&!5b(|s8@LXq_RvPdzP}&lhy6kJJ*9eKdio~MSlGI`=7UL7+t>M!|NH;;%}_} zmwk5M?l=t=uq_P-5`y?ajjHa+r0>cT0t44MQhd3+F-mUMeeBDGWst=6DOEkOWywQD zmDr_&9OExyKH1@oa3@dw;CG@0q1JMeHJtb~5S9E|sjf!WAGaeqO%s!@StpZYy>->a zv_I}lv+}dhPYg&=er2SD8RL~*46k$$ zOn~9Y&m`-#9)7&9`cUvgWcyUjrH()~P89auyk>{4Jvv9X4KJ9pbC?uS5T8nei+^9vTQQ*j zve|%}_=(+z-|G^srZ~-?H-Pxg=3^)_&7g;iroIP~Cu@tE^k1Ot5V#H`kR`u+v6)Fs9E&C?ztoV9t3T=`+jJPA+CP3>JtlV*cCts3*s?J6Vg&w55wRn$d-EnC@P6-=l&-|)Gbe5!Omlgc-`D6=+fYz+B zmG0`0nmKhpKK3310i0zTP_wMTE?03~UOx%4jLk$jBcAgcHT!#HwgWH(=ry?NF6Foh z-1Or`fqRSQB@oUMth(?KBtdkhLkZw;vuXjV{u9s;m?`{pMa(gQ)`6;M2;!UG(WP~x z$MtoCN`_d^S-IfaIgdtYS_Qp!e*P_x#{bRG zX4m)`LbwF>1lB932=S=x#{P!0Cl#L9Sn!GM!^@YH#-g>}1GMr7z!^nW=!|Nh*p@1b zKB_3EZjk=x>){}Y1kN|76`CAommCip)xC7pY(~~%7BiV7%qkO7I8usGQgNyEr_c@% zDS*{(Zm#_E>tDd%0AT)WO+M$oYV2y;nY>1Bm{Tw8Gn%=RyU!-IlI_63f8PNT%aOEI zw}sys99SVx;*awY3d5F~P@EnGrOvlPz^}#DKkosdR1*$zm#mlqPHXzUV4*+y0IRtP z+9t4C2|i!52X8k(!-9VSAj~!CYF(VXhY<}Xm_&3l0;NlZrX*xz6C98nsi1`0tjPWGP%`6d7eUNG zV{fV0w=Uq?WeC2aB_-S^ymk$~BX!#wAIw0P(#NJvy6pj$CxU)&{wOS^?4H-`rso$) z?zYNDMT?CmQbGTY-CZEtyPr?_q>E6Tu8I`LVcV!1qwQCsjH5XaE;Pq?N&VD6C9AKT zc%v1*Q5_C%!I1erttVC%Z6XJ#Pkmt!(bQRP&omN(<569jVe#?e*Y*NLktITl)EDOp zu`u^;Di|}lr*UXY(O$d79}cq@WTkagHE>}iv?k+;Xok^qjO$bxyzVqWco}SLKirhvwDsQ#a*wHH(Aa&-9z1ec$$hKUwi{SlROU zZUg-X-0lpGq7z_FBw&5EBSn^qM{h%0`|#M{z6K&Tm;aY)ky$_@zarRqjG*s@ndyF5 zD;ZHf*FTJ3b+KEA(R$d6pcp}8*@Go@D3=guBab6b4JAl*p+Q{9H=$!ee^k_pL6bM* z;l$*!9l`Y`7~dsj!tA&G8w;D}eue-d{HE;IzsZ5tPL6hQPg5xy(H78)6Hk z>x~wVvG6fNq`rX_Od~4T7|5A*@IN#85{RZ$M#vt<<%xKLNG*kmAM<XI~08M4nv6{9^t-bG%NOGc>5ecT_A*SUml;Zx#weD{l@;+U>vTBnVn z$#z>T*s_eZ9&}*s)w1YdiI?Vj@~pKMWI!_X{h5G5*s7~SH;c+KpCU6J2lIRXXI(OS zd<2RGHclP+QHYXU|6-1J>=5Pe{85nu!4y_`1PBV|6y{ENnxJC4ux@6OBw?5e625M! zEUi4RrG_biQZ0495f#ZE@3zi{EI~u*r+?9ZJKD%a`HBh8=6os-}B-c8eck zRFMY)|3v;|QvkYQ>ZvU?=-w?C3b;AwT*O+#1GO#vs%fJ6 z`3SX##H%eu3ZBLqtg9=s-W_-y_Y*DjR{K*cq}c3+YlN_&2=HEp<>SEqL~@ol6{T&b z7Uh;|a+WD4{xSmWb2c;az_3nqv8DxWux@NNek_ZpYD;ccp2Cx48Mg<^Y@?_ICt7Cg zhj3=iP%gNW>nZd7Uye3-QDy&W9k#G=zBAG2Z!q@E6;c(IP1ae%a}>YKDuCeXV7p9k z;=fKEJJUi|)BDP6Cu3ONRaqv;!^R)MMVT9^*(J;kVKex?_x|)$T)GfmM4pHh;^ePC z6%CoA@4HUzJJKRh3~!<%^ZxqZ;dv(7lNSXnf?KevnWaXX)r8%S!WtY6y-s_hiLV%V ziL-RqGb<;DhN--2k8_YR2Lfp(P=^Um^IL`1HoA2(S3aF38nahpol6PWh8oZr#eU`W zr=o8jHv9;V#j*SLv<;G0{)V~eq0**F_(iacw7#cANv46@30J<~N;m8_VI6Bt@a^L? z>JUdT!AxoX(jjGVIDKKy6Wqbw4VG0|5f^AhUnMa=8m~1zw_(Fyf>mXWU5f4W`{ECZdvJp#w(d1-oa>N9Oo3<1;W%c0OSt7P(%Sl( zcKq?+Xi3w3_LyD!Aj>{|JYL9#4=cvqpV@svr)oGl^mQ&K}al2%nL^X z7V{B#WJo4n83@w!>mz6;o2qKw8*ovo+^n=;pc)9UyH5Zg=ee}hnqnv}HG%TC$}*ev zA{IC~$SRzVE^P2yz7RexyO{}GNK7od?3v)Y4Qg2LMawlyRw;^o3Um|>Op8!%_9|Iw zsm4Z&J)u_JqLFx5sBDc#az3SP>uNw8OQ9f8t zG?MSzFw7p^Ew!Cn z`UT~4<5dx8Qguhrs3&^#<|dYs9SFnPQIZTR{nUkF!pVJ&S!sO^(o_ke2R{6OHZBy{ zFQW>o8V?sQ#nZ>&2?~;=6I8v(<z$i=C$uqkGR12Dy!q>_=>U_^!$%O2ODAVk6}E0;1GlyIR0&Q z{QKdh2lhXPRf6D^FX5Dwz6czw#W@&C%QJrU?!ytA=mK&$ipB~>*cND%hT%WH()(F@ z5uZltTQJ~sDw#r}6~YJ6E;G?m(WK(q;j|a)v@5(VizH~%YWUlE%ZMq?-(5&+wSoA1 zLBs7%^@zN1UqCQvUzl|MOQsB@mNP}!Fj0HD+`w9U`k$Xk*D``VZmjXqk*zmR2Abqx*uhSii|J)CY>{pcb zx!XF2>7S@5;}@cbxg6upmG%|Z827T2%BC&gjctu1lAZ+Ia* zZ}Xo9ql$y_SC4XkA5!E8)%^x>P5<6z{eI~r(pi4MEqV%{ZHkI+J#V!CH1O76e?tbZ zTQd{GhE9s?+$1&y9oK7_GzFWtsR))UDc2)NQ0UFD6*4y#V>z zgdMPA)D(2-B%HxX^xXlLTJ15EDH6b;{{#5aZ(wXms(2;j6x-Y-{TSlox9IgPRsHAw z(UaAqAl9jwXFS5fDf<&Lw1wsmXAEpIIvh_*`V93w9lva%bczz}-$viB=fHUXeHBR} z+Ov}0$8f_(*9y6rw8HyxEv8maoMlgH*Ch5Oj-XTJziFZgBU#pnlxTafhOGf*g$}ga zA(r43bAO<{@h2rIaQJ5JgEdr?Xh@*%2IJZGucDc(ncrT-+jtcNmi=&$GRb&ore!!~ zZa}3{F28ylWIO+azyl_B9}c$*0ji=u#BYXik=rFccAYAlntB8ET>)kiO}}_Xdv~OM zb6p=w`X_Yc)F%fhPXC2_tz#x5T4@AM*BWwP{H+0Vom`=k=(*wCqOjn_5a-8oK7dnD z7*Ok)csMaqPbhRO0LSP!c`slo{O@h!1I6nuvf6ewm{v?rw^IJ!XQv%;VC~Kas1+RA zINM78`(7pf!13>v!xViz$ad8;==EmoEu>uSf-xGzUK`mMp%@7gqq@EP2pyORDr1?jPFkksP8T`zGpmN+ z>C^4Yc3%clz`Sdy#`*#4S{rKzZad>NVdnzHnbj01cD!I0gxCuGU_DXfFWDgQ`_@9i zmty|7-k-y^A~dx_Y?qJCc5VsUti&3?AS2qN5SP;3iom2`nBGz|w#sRgMr#3HAHt~j z9txR1P0OsGJ2?+FUH*=ky+&Z*W+3pKt6k%F1cmIcWGLO%g4Mo;@lv!r<{$H`nRw>` zv#dPLpZZzN-U>c&(KlPQ4argx?V_nIAMTXq-WCJUNaF@dl@B)kgxp4*w@iYckJPG8 zFB?xAb2Dn3=@xdSD1gT8`Y?m-8K&)yjBvJX*6F1-f0;W8n!Wn_TMzC{fI976(spm(63hS0%%b&^CVa?y3jkZecn1EGw!2s(aM*$fNw%5&TrMkyLjYDVc` zy0>IwW@X&&I3MMF_h*RW?|gYm0^cjo#5KMz64-l}@au`JseG2p9CVmv%>TW&%@>v5 z0LZ-HV=1mLifFeY-uQVEe##`U#eqv zl`MUtFmw&@7jPpZ(M75J65y5=HohrEQ*0Iv>=9eTj{gQY8tz{KT5}o$&ilGl!7f%Y zawYmDzf%(y!e1XdMn}MIg|Z?y!yc=@>a9C{Aht?zUWlSM{{8Uel;gQr#Ry#a#y9Qh zpW_7g!f|nukZcC1X0&72j5?B4w7#l%)&+?zcuF<48^HTBr$mJrmm7YHc2_c42JSIq z*uJOw0V8jF72bXOe%9ZJha8u72T0`odyahh|J#L09_;Q9brY47o3ufeZ%AeC5t{&& zuLvJa46c>I-moaH;IhxbmEwoJQ2?uQ9C_H1c^PylkW9?~TN_s?9fGggr%>j9+-#?6 zPQf;K>XBu^2yH-tWy$&3m2=LM!7=?ufXy)AQYW8_)nt91G7yIQCvrDtKY5w0Slx*E z!H|8o*6^uNOG9LZSYA@zpl6y$9{e|m(=Qjg#KyXG?G?tO3#F+p*;PBZj%*T_=C$12 zIHO*XIls&w+j8jzLNVp^svnN_heGp-UEn5FT`#^0wzpt05UzL5oYoEBS6q0%tuwux zaELtQy>w3$qzgTiL?y`(m9RU1%RJ@T%ja5zhDO*%_c#VcwQQIKJq=Mj_Bs77r28G# zukb?4P^H26>=i+M7@T<^k3VQBvNu z;bxRI{ot8H`H(pIX|X>~d(`TI@1xnvsN|PqfOx$4fW>sQQJAy}vp@La z1T+p6SPzQ0QvPP3)mW zIM||*Qf~qUgd65;SWo|ECfAxBUL*VIrTWXA=j?nfwba&R@1f&vxA5>Q?>_mqzXeBh zX#u_%Cl1REo|qhZ*vcjiSh~FyC^fLdp6)+>;e4CR8ojQ;qg9x;z+uRNG&-(EH!HMG zr(9n!ho5P)u5$UV`&838%dWD|+EP=b6&>GAQKE_C?eliOt3|+)6ka%zcM>wWg%Pq> zS{ly9tWoKV3yJ{=dDkzdD;CAh1xKU|*BMc#-=7WBOOo!#vG;(?I9)<3*W5`Y!}cN+ zO`I`pzYmpIikLqA=QXj)uDt|W`%FP(FV^DSi*M(%OsVuMchoTVC8Rkfh480n+K{vG z=~pZyDbdkc%c(857y?)_*OkFZ!JbAHn}SH63|Sl$Y;j`jT=aunQQu`M8Xh7EGIWEu zo9TY;Wp!KKwRpo@OP)B$s&AMbU6l~N7AP5d@SD{NOvk~emOvt_u27kS&JgLJ!{-?g zK+P9jLDS7AmMM{hMHF2=-Bs8m(xOrP5MJdM z?bajXa++zKWX@3Od|Qt^8SbX;8@3huVZdlR;*ehgBeH5Tz#*aYpTT%wjXEh(LVXCa zeC{ao`DCv>Hg2WkT**kQJJW2qUK8^s+)mjPhnjljCuP=>Go%jRb(=0ILV)cA1mbQuitAEA9~`MIUcb47;D4jf-M6~vcoDf zN7C!@CYd-`$E4^g92SMmC;AOD(FQGjY{LJj5WZMkaBTEhLgZ5KC>6 zxeqjjY-W=UqNF-;5pnPHVwIlob5+~Scu!;c$HgUaVMxrW1f88{PSyt=jmJram|f5E zOxiZ=At=UxG79=d7-I~X6nNtB)TsC?+s|IZVf|SY;*n!zw3^2e3z8@*Y)z^2BRRh* zhLWA?Yqh#d>F$m5je3$!xa`|U)F|A>9Yf6I?K|ToRH!5hei-_W#IRfN*O3-0n_otU zBlDh|hFfMYWm&LG<@6|Rz3Rat?``!&MjP^4Brd-9d~vh{23m3b-}`r(B|G~`;H*H| zORXxlO%l%TnBOLXS$du_F@6&<8!FtYehEAiv@O!RFgjyown`<CLg7TCuq?FpE*M@#$ypNW++bnF}DNKyXDNf{Qk!06D?qNXRb zk+Nx4n)Fat?oMT*0iAR!3F$iahT><30pG$cqAc2pIFShMi?CQqmc`<>RQ$MAPKb(~ z%2Pcq%GvB(YS@j|mLLDJ^_HM#uk=d~P>K)r~PLBd5Y~(@lc#!!&-SA>&Hgl*RtFiK87k5zmvvwf8m-^=*>%H9t{5-@ozG-8g2Oj+2 zt%IZqX$6ui4a4+*2Cexw>OVuL&azc(#7elGtP%5FaIY(v?$QI`>}qs%!knH(Er$E* zm!P2FrIoNV(NQK@=GPwlj~}s-+DEgwMiRNV*#nqp@|j!c6@on*l66pGm@J*;C~q}s z)y)Sc6MVEpeCS5~`q28#+ior3G7wn=HfoHqO46{CqPc!M&8YkB8^(LV_6+kM>sL!C zqiAG2WKU8S5y~w(Di#%VH9AS5Qt-S&S9;n0@?z&z#Pv^Pl%2#fQ0<*}A`JT0sr*z9 z3Ct-ySs@UT@rb!`v&>&9or`dRAuxe!+0v>sC_V=#+Nw^lz?raFxRI$fx^USo$64_= z_gS-BYJ~N8=QD;HmO{hiU`mLhZ6n46OG{-<7<8~w3ozRqu)Eywyy&%tWe!l0=&b|));P@HI7v1M8nOr@5UZ$RJ7lH0E7NyE zScB&vMyhjCD3qxS^@p{d5Ojk@uCA+NkrVc2nL-$o;%+wQG!Trtl)jwB*@tK%0AYd@ z#AhpQ9!BFf^wK#{=Df#;a4gfmpEWXl-!OmHXzkJplisd!Gjg$=X@E^jjz}w&-tC2QqFzumZQ-Tq1W$83oglx?n%W2FqhVX4`Ed+Pup5p9otwfewiRJ1Y zD@dGx$dLPDGRv&zkWV|BH`B-mhbU(}%O;R4gtiatwHUvax`QH<8o`~I1LL8pAQn~| zknbPZP>PNKN#c}JA!VpknPN4Za+p(at2Q@)DT7ePFrrUAu~ya6M?WRFGzxAA_tkzF zmI`Ny4<|e>j??Leswhh$EW6#asPtH;6DbmP<0QTo|KYO5P~^v@f>4? zLzz&@OwsRSvjREJT6b}7_Lf1vbu6JmuA;>X-$#I$8h>6*<>nIdv=W1zr+gND4Sfv%??F4I1QLN^3k5^9WG!5phf=> z)H#5-d|yh$#3;4Fr_>E5o@eXj$8ug#6L0K*r|C2>$h2En z2SMq7aN<(-{{vZ)K+xDqCvBsd!lDi725lxAQV5(7d$l8x z$x>DljJTGjvMRZAZIfz^9Q$)%%gPhZpZmsyw5?ZUaouX{t5mN!2oVpxu^AB8)&MQ9?zV z+MXX%yD0Ol(>^t6RdZ+7Zah$~zDZOt7v|w=Z)=dqY@$kwGaK5DYgFkg(I@PmC`bI_X{?nPE(bdP!fAbL;L^zTCi$>Q=jDoS7Z4MzL?xgt>G{%x**eIdAjdHZ= zZx65Sw#q5RVS_#N(P$Yvl+r=Xpj9Zwxj<>!|G00*87`jLgdg9HW%Ocki%+EpBIVwT z!u9444Z)qIr4-R{}$VhF(7WFtj>ziY(roO+nT^ldwVtgbS3uY``h;wVFP#^GblxMVaN1T?$Z zBRz-iiZ2k$Do_jaZIR(k!|Mhqt8j1bO7+_0w69Cz%o9{IGi2zRvnS1In`|7_pQ2=cUpjKbvCieQv})wDACV)Q z@vx!geBaG5=*3xihBhQ=KK+F7!IVgUp4-NFPa!3)l=(d+ZAX`9h7O(407wMR5g84? zk06NM%bO8f85%7+hn97MTIjnflNusf>qFht!fg7aJ5{q_M~P;uLi7QRg6)d2Qj+0N zzamz|2gfne|8{v(mUIo56liZn9%RyXwdueNm*aTfAh*GYd}_eV3_Lx1M>#1WhtR{z zZrwGWA-d}FyqLx$mxcEFW$v9k8hyyf3n}|vGm`>bmqW_D$pl}8s=ei|un^|UpS~b7j0gZY^NJu0dM4Y*kVltdbgqOulX_Ic~6~!)j!|tz>aQuz<31yWr2a zY4UU1an91rwqDaK4$KCgYwoN{!RZdI%S#_nXnap^;`$ZJXgg|_erdLLa~E^bns0_Q zOK;V<%O8^5+&WDIH2@DmW8cX!m|fIbjUJH=0&zEv7kRd#`njhGow8L1HHmzp4X7Mm`^C7a9 zsmh4vxAbgIrY-x$l@%K?gJkHl;{Z7_NUeB`i6}u(wyW_Z^O*`O|kq z^_Ff`X?efyg`HD1-+d6A+mtaLT402^G5!IF6Ve8ulAzPI3SO>cRxDYUbRrX@YKPmE z!Vi9R^a8IAnYHV2$K%NswY%b-Iue#o6E&v&8k=DzGQk zI2(tv%6s?-WM*2MT7wJGq!jI(*NDJUwB`ShmAg>YyWc7o)I6V&PR+~{`-5vw_vci>O!?_#d$aGF1+o~{ zV6j}Z8MWgQMH>iT;kYv;yy22DH~eNn(2I)GU_XTxKNU3k8DNX-A6Ul@spuwKF#M6H z9mouVE&k;eB~tDR64|TZGt}Vq0uc;1C9C(JPC(=Ob4~s4Sti||Ej+D^X? z$KzQhbDe$!NGvC!cNw5uB5^ows5AnLo8&UScHD3-Gh72N0zcVXN#r!Q9#J$R#Cvwy zWr16AvzuBW??`opK4_{Cd`{fLsyX}`(o4KmIp$&6AZsu*eGwK>Q18aNzLG?& zCfmVvHjIv(6N#1XjhVn>Q2#komy~dRomc`VfD#<#{2qT_7RC~n|1bi)poeZ)vQ1q- z6YcEesH~HEhjk(HDH$(B?G`9mt~A=Kq;Q&XwqdUWM#@m!;?cvgr}z_BEc$*j@HHIK zMR@k`ThI75V)iPEJ+UstBBCZ|ae_=>ftJacTPBY;e?}%f&a|<|G~a?#15tIQQ}Pr)Vl|_FP@MXjmia zN&C_`8j7?i>;gJWcZiN+Iec<{@z@Ict!65r{(n==#B^me2rFaWrV}$`5O$ZzX4D4VwJLf|*!%Ba?JZ?k z^TIIwK?|{LKeJaf6KJ4%aQY&47p!}j<(*)26SqFy4^pu5d1U6RA{lpMkQxjqXwEuv zM;xRs$w2Tt(m!nWAXrv;$+q;MwB@vA^&(5@J|O!QO#>!FnZ4l5L#g^fF0g}aH_jLU zXF`tQ`KmK8Q%{bt2rrr5v2m~;nRDCv3ovpPBBJSh!%Y`v z?Ah#zxSa;08c@vaMzMb;_}%|QMZru?Cj>o47 zjx$eyE`p7Cu=00{SIR(EMB)s6p9UWt z*CwBCGc_(z&CltgpgD7Bb{kR2ET!4P2yz#Tlj4A47-73hXxZTsie*0iAQz$2By#=J zNKKfcWOxsO%5cs0fiOlb10$`hlUlXPQ9A@gW-mF$#OP-Yi`jzW7Wn?@Cydr9aK5vS#D-1RFtjp0-D%=3H(c7$OY z40rOvsdF@$4gGlJ>W27YSTYy7<5~@95(U zx2N4PXI&;3e6u|-o9!U@M0w@|7HPD7dzm%wWL{!eWG#NWk#Bn;@tWb8ETC!wLR-({ zGi$QHV6!w=PK2iG3u^??$fQ1I+OIm*Js99P`{Mkf3gh$o_-M(A3~MAzd8iNHFvQN- zs?l!6DZ#Pz+ea!+elW>5s%Bzr&MCCtK$A7WP~52@VnD|QIvZHvPl5dAa(;dN`C%?M zh3-=R!^^_k7ArLsSiOB`n!?{(KLfI7vZ_it=<^U+9KaWB=HCD_wHgou0g;?4bfr5c z{T&dA1MD|6#bnD64^=~e*{EnGN-_Jk#q&*WTmKTXqWc>h(8}RH!|jZF4cLBP0UHRc z3Cn+|gUX-Bm(~80!<9s2$eCbF@jJ4JyfWT9EP%rwHEZ> z_Gdr?Mm+`wKGMIl?U69`Nqo6IDauxaAEt=qN?U9zv&Z20^6tfor*fOrlh zQwbNkd>De|sP$;^*9_QdQ3NzmV-S?&dZ?Nq{pJrq1b8hp=dVUj$5#4N0_nRwEZ+F8 z7~mRxQ@~6!yRrY8@}3#Q*Egm@l#gCh_dQBnx{dDtybn|Y01Md&eOF#Vx@@CYj##dW z;1-iTL$dFe1G4|V{1meiN^2O7NVz{C`TlYSK&blgNH7L(hNhkm+7(aocL$bKpb|v1k zl;;db=xX0%`!R0MY;YqB%Z8MIAkfDJd}=2iymCaRZ*QaOL!eEtLC(^1do7O!)5ewu zpz<;LSXbnfwL_bBSB3p{%UVEeDr^!&`Sv{Hs%EAnPon(Tk`z?$-oS+0@Yj4|JiL;@ z8f49FT_S?pmL2hbSO8+E#}ynwHTLbwJw0_3^SF$ur$V_x_+|99 zxIu-xwOZrBBIPa5g_6R{S5JWZ`9>Q^uvS7PtN5*!_+(TZDtR#uG6w+&R$plaX-{;m zql~oq0ZsEYenAWD(i29gYHUcxUp80&sWI_KuI4x(K@-mj}Pe1@q3I(tMPMp4HE2PMSfM$3ZkZj|Q>@bU0w{ztH<%;!BK%;;I7`7NT zAcx)4wdsF7(={YxYvIS3mp*I}&e$n5pC3of=I07&7xnEteIuCN!T)+%544Vy-HZgc zZj0XnC5I6f(-Od7u)joFx!&&uvuPtzCG5Y!BKSd&Xd3n41l@S+(#7`^fQ+#Jq~DB= zs?JH5@y8JRP1_Azmr02;bGHvfe}RwB5pebrYOM!JLCEHFe7r2;4}trza*ClkVWxpM zBkw+~4dt8jP%|^3FQxu50LYrJr|rafh61~TSvWHn*6qrE+*b;m_6wtM zXdOgrBSwZUlyw5h@w2-Bq_*A`T`Hxsp}F>9rHkEe>zw#ze&x5lv^_9BSXz0dTIhD1 zCasbgHtHXQ@XIb>Tti4C@J*2TdHZTS=PG-9IWcRELb`Grwuw1V~M%gfJK3Bky zx_{#)`9f|#DWw%2^lIH&*s zYlo}(^@!bF07CVR0V#weGQ58hdzx&{dj+xz@c0j!^>sSDs4)+T`Wkpwe!ulRCJuTV z;Edi^8nY$BjcT%F8RrxP;c*i^#^#iAHw|J~INBiJs#U$2;6|NMpWb#ZX#yUyPY_TC zE3J9ZY17}crpCS(#~8m{*WKMZ#6;8b|0Z=db2xiPd8+p6rTa|2CAfg9!UbZMn>WgB?civYdQ%rUT81{2kPG>kl8Tune6P~b&r z@8y7vh`=%qBf%Q#glc(-PO)KPQJdvUKIBc*Ygfz9KB_p-tJ2$^w?yJ*D9rNL)>yj3MN;7G>G9oX0bRmH&(tHe~vU8aAlReIQ^tTMCo4P)uM3r#;n6 z-iTO|Vw6XPpiPbaL7aE%SoGHt!{!XMhC#uw5+{X*X6h*pL~ z`2?dq*OS|tJ2i#iS{UkRU#pX6!KqhpPOMBaHwGLfB0ZyZiGotzzCu(i3n#9l?-MJF z-!={T3}^l^-f$UCSAX<;v43BhVr2PX3K9L5Y$YU~cI^Q3isfz`m6_cIYPYgxHIGOZI#BKG<5y}UBT zl$`c~5JfIATJDKegB4w*b{RS;*zyn2s%0SY&M@Y>k;7<0eXc0bEhTpXJb|Dve5I6x zG^N)E1FCotgm6!0wUzok>uStVrXRcU=Havu0^fiv`QY~sa4T-cOYcMJ$aC;50@7i| zHTLkxhzzx(TbjaiuyYf~95!q?2X*Zr`ia6kQsxlCTJ96jG1zQHBg4DLT~4XlQbumr zi`3FSuoEkphLQ5Xu&T~_>A40WRx;_W=DPbf2u5-(CfI>ubK2z9YcsB5qI^v07YH)i zoc?84XWZ~Vs5;A_IKr)4<3WQH+}+(FxVyW%yF>8cE`z%dE(sD`0|a+>cMXB$YtF4( zbx!r)se-28X|}BOtPd*~x`JSG-F-MBb)J)vsDluqwem)$(ZUG*+$>Ke#3b}M#-@J+ z$jG4JYI=I!Tc)Vxs#Q%&`cgzv$D{1$+RpU_N4pESmZ&%kSv*}Q6R@hR9d<7Worv^L zSPE7V>+t@E-{jVmF@6-`dqT+*)s?N@jn3*r^=l9qW2bkb!4LunbP3lFca<8;7-UxF z5JB_hI&GrX4+SbzAir}#@9&32hFRDb zlcEVywXQ!Nr1_~>;z`K06?y8TsN$-L?u{5)(M3?sm3T*a`V$l@(jbS8N%r54V`Ejg!O_W8i zM@3QwhbCpdEPD8yS;)42InIUAWv4Bh%g}Gy)*P4)v%kM+lJzTw+KA!j1)cPG()E$S z-{%raQYIP3^k~Xkxy_kGz+kOTu&>%E2CZc$%(vK_n21c(K1zKY4^V57X?TQ#S04eJ zKWy%5@_p}^n~|cpy+Tt3uRnuH$4abrs$^VX2%VES%`N!&1a-EHZ_|1YZ|}Moe%MJF z71_@WH;v}7Dj*O78?WEX&2&)B7`-U&h&2 z*ug7H&!5PEH>33me82h?lw}I6S>L1~P6{xYx1D+L8r2GtoHbEG5MYo|)Kfne41m@Y zY^>r$YbP@7&^s8!oDB+sK&mI%?VMp{8`;+ngUsl)Ss{I*mAvgON617p} z$yOu0#S$#fan-|@3cx@)#rO&V&J#tA2dV~Pqmyu3X!IVT5I*PTf28LHxosV<5`H3X-V-?S?H z%4czW1&{tY&*_NBT(hf1jb+FElfLeGTj5igFbKTcv&H$3iymEPIj#*biddmHkm~xm z`!vY;_6=ut-&0));{kQng~7o9zt`fepX!7Fmg_@6ZQKO(FD6>(&qK`}Drx_)Cf_2gbYq6f((`~x?Mly*GdgX~tQAqvYoL4Qra_C~bB+Xmpvvr} z9TYI0_}wBAR-eVu%V+q5343k1Yivfd^Hv{nD;2{>kLS4b#x3odE_ zfb{ZfHTMSgc4`m+KG42-K2iRnHwO_d)D*UUMc-uy?u(zO51)qfr!wrf-wf@ z&SkA1M5WE-$@Tg~ra6Noc%(B$e-AsrPeca?m3h;_g-W+%R67IEp^n8uKe%pdYpI%@ z6Ajx2RJl6lIozt7{wDq8{bT!5e>X~8g|fkE?65DDW2N7wg#Tcer|&Ad%JFnMk-HS`>ITz*%HOaRiVW?6GL{o2BewS%=S#a*c7V7^DtXevI6YO6 z^Xn5{ns{E>Xtt0JVG6c&`@8J8+vo5s8%x${_C6E2^N|+k`MjlY|NCwxx7*;*MAjrb z$-0Zfj80o(qx>BnEYye5cl7?}n`q_bRNd!3kPzCZ<{{Upx;eMFW+^L6-amgEnp+aT z;Gq{)U^i-IH2+2;mhExa$6SJuLT+TQ@W|AJf#cEE3lgC!1gxUexn0&n#N$VwyB8GX zD@8J@gc=n~JtPtaqQpp}MK)RYEsrf#kD>`UjWcw&aH~ zILTuKfY^y~%@t!~S>%0!D*TE*VZt_TQf#{$i%dS?r5lvVFrlh}W^$h@sYH~}Z5y$| zz!n|6vDEJAxs@!%vSgd;41bTt%6b3j;WL+Zl$h_5#fE&^XO_V+5wC*Ty7k5F=dAYF z(?s116N>}l7*`AKYgpV@Q7dGHFj{jtrqzIy73!^k8{UZ)=lM!S*>DXtx6YY5YE4B7 zG1Nh!PmlKoi@a*AGg**LgWuE3(>6j!1I`9qJhDAYs1|l@;T<;MTOxeVoVGGHX5r@_ z=d%_ozSLz@vbDaH*O`Y|Uxzap@)^ceIoNf>>c8?_3gIY0@p?VNcLuZ{O~8qN0KuE< z&0PG*)d{#QENngzJ*vz|<)EN6v_?|r_7yJE7%KcB)+XJYGDjHn+wG?r8k1f`Ox!O2 za1Cr=l|D%?r823KFH4B6H%|rkaidAI%_qy$eKY~9X)i7834|j-7MYLi9jQ(x0D7ME51BX8ShV`};Ze_*^ z4H~Te@~8n;+}W40+I>hY_gJc6IDO?xf+f{M72(*y6@bEwW?(z3DxIdL(hOJZ!Q+bl z5@rmWt8kFAcmu(V6*0c)6<9g02>DFgJDVcyl8SEJYT;NN@Cb}FF^OtNTPOdr_Ru}iq~rTWdJYjK z%Cd&vPBFimWIf2X{OOvo4g`hIU-4ZxZ8rirobIrz5i0k0@vopq9fTAv$jTUFs9lL> zmY@KM7V_x}${p|vSMVUCjzrE8BNXv{L_NvjPHe5??ATxS+miwIbfsrEh28M6<6bO7 z1m4dzD8DQ~7v)-#9)fZ#g(XDl9Yk{$Ci3U-?qz$u%RvPOEJfZ-N(Z0gR-^5BW*yd6)(|9}U&>Ob z*THH{_3fh)&Mo6qZ*pgR8W%zBXyf0!iQyq8S`j&6dT!I5wcVr&5LH1bQSb7Mb}b5; zA-mvW5Y(6U2xBR%I&2suGjon|iE-7eR!g!-SrG>75v_izw$pH4zJmmHYgKI*?6&4h zc#%kjk9G4xoymqBb*q&P(WH?iuc=7BupD|8!KB{oh@rbCMRlSpi9~EsfS;0JA4dA~ zj$Y_G`krbQt~7o`7wOpG!ytNK{}p02M!GOXn}1~#{Z2X#`41=9x*R$2LbU0P z>EYO!iX6*0)Q@TkKN3oTD&Xmk6mF*tT#Igs)tsv$e&3aZ{uiI!iE zM9tDUhQ$jLI%M+KpVLmcbxvFr!raxgALpgzbaMsI{K#3fciKvsa z&zA3Of0IPz5T^l-O%d=~>rW$FgcdEOl7Oh-At#?%LFNJvSJx zL<=a@x|=mclDx=0ZKgzC2AfZHwyKFl<;&>365XL+m7yrwoWhu&ubD^flB0QR1XRhW zt@IW+Ls|e_!?+~rfV%hX>x@YjQH4+#96TQMXG3jA#`{67r?!ABQ@gf$yvX`+i|3;p z^)aX|+_CeoC5X}qKbEcdghC#A((7-S^Cl{tIOaY|rO1Y{Vy9e_bvfNFo5;L3*;eQ_ z1#eExidz^xkRj)kUrN2wim4O)wf9LP;p}&haSJB6Ty>D-L84Bs&!dKD?yIlp_Pp$Jn?krM z9#T3E+akme>=h*tUN`dGrE-*AmliM3w;eaOS3$6ktNX)MEU2}*LlqC+?3Yu}gmCU$ zAG|dQY8~4SAV--93!X`aG}}H_%AaOhsbWuh&=jW*1V#7;Y!GJgvcp=!NA;<2xk<+2 zjq~Dtf?_CNzDC%7&KQ}DHWl(5Vnpg2*B)D6&SSF*8*?HXO{xI$Lf8qUfVXg)>pQQ8 z;0j{R?*)gxQ*($}@XX^Q^xSue;o`caiAGaB%hW|LZw{3dNtJAw5w|*S|9UhnZ&<(! z(CR5*A8$-h4kE`HTmOZq>>|N{v5CII(xJr3)*}Zy<>p)PUSdlOu?UDQR86e{0N!nQ$TzC*kclJqZr>l9<(i8GAeZZuL0rgXI@9?vSOVUQlc*SRC+LCwJX8L&N zU61RTpz~$>krzBR(YJ0?!hjX-Nr>kX$izNbYwp@gwTZ&SeGD|3A(-(iJcIc-=%GQB zSgI>Vwtg{&*-@iI1-BT^6dZ>?u;^$2&!P*@2wR||wV;V?exV>c9f##Jqb8qiPK=N= zbtl&@-dmi(=crJ2LjZ*m(kcRUo=Z>*R_=48CR{%enRdZH=j4-H>Bk8gr{_A86rl%Y z3i+0on_zIp$jb{RkKzzDhkPeWAf>SW5VSBL@u3bL79O3iIfOwuKg!PFr;(W3fnP4~ z1xW@^pksPrjJauc8MAVfN3rBy5h1pNvm0a?H(2W3VAGWGZRuK|KJ`1UL?iCOOV%!m zqN&JHL1ET~N$F*=j_Ks6WYxyRBv6Xg4JD>dd0Iyd>}0C<%XQTcFtEN4~P-zT$3A9 zZGU(&_zCHEZX>8)A_ZLZ7o{)i_hvQt&*OWQ?Ms$wX)I z({fFy1<4TfaZgYe+^w(~SW{ZQQ>UJRuk~Dl2N*4UP|d=OQljV^>15qupONvMCmiu0 z!+1tYvk7XNC;N-{PRb&yi~p#S4fKQe^<3}u>;+cc}K z@D5B4VqB2pp@EK^Y@1{vk(8ks_Xh-Mi{x!q?-V**ao{kGB-y??LP4G-9nf_8qME%) z)N=9qFE+l`>bFE4Y%wxj&)u4Paif1;qt1$1@ull|;K%*!a_*p^yBbzHXVnIab%AVh zu=p#kHF^USrt%B7ak$Vbo7WQ=F09J(EtoRYn|RI()9W5GW##J(IUkwep$BY0!~Ac# zOS9nFL1q;VQ7Qf4{ zWX3x+-1r)J7IunxA#}ve6(d;<;tI8kvhMJ`7cjzMO4_~QsK!_k723~Kv(ZZOameYs z~HOpj<2Omdzf`Z!|=hDKr>#PBEvZekm$bJ$O<|_ zwwwpoNW#O&j`dRtsM;_zt~g%m)GGy(4*ck*c=3OqWz;H(ljrv05v`mNf^t&#kEw+|PLnYwdU zSn*mj_=j4MgsoZM#X2!qjaZ{Cx6!&U9GwXs6kWM67fZw6m&i2zk-@V0B1)rH&Qq{< zfw?LUz{D=H_mmvy=WAlEo7=|FMw$CI`w?4!c(6+#A`77NWsUZc=fVB~06j-PE3N+OXsV+E{Us)+IZ^NWmBi{G#f~Ts9VryQs;CO_+_XYSCfCk%heM#alwSd)J z`FsOU_;w2HfCMIJIx;R1Q4R{jK0=|p1e`aHPqllD)@E*Bdwj5snB(XQQaO4vN79^Vj>!` zy62;+m0G#zi}#9ASe9%;bj{^mT0orYVPa}THkg`? zf+T73wd+e?9C=M@v|!bbs#_GwxaMuq?sf ze5ZE}KYl?m#Bp2up_OY$o!v?q%A9j!lrllDsUHwtnZ{fp&7lc(05A8RyL1~a>>oJJ zH|RFGvsg6TI@>XEB7l;1?WdvdqCIL=XN8F$T>Qw$*-&**HXT(n0aw?d< zHjEkj?P!8yY8}8*m!m3u`~j#TVwv0bF|4Wv@e|+Y6GBA|CzNQpbL4z0<tIt{W0Wf|ZL0AwZqH7WZL%rlUE!=2WWyESkJslZjs>shy~&724IXT@5M$XqVr7HLNiR7*AqgHi4nx%MvO**P2U= z5_0VD-!Fl2xGXSODcJ5j8top=M@7JoE+7uPpUDbUQ2cRcB&6hq&k4=39ENVqn&NA~ zq7BD>0~>z~SJr30Xo*=ETykK)_U~T#pN4yVWD+<_uqIAFHH>R?Aog7)k;ONH^K>B9 z4fXqUVqrH<2BwEyE#l*6W_TTelqbr#DLcnijxx^upNin9TP1O}pN*WPb3S1t;d;scN zz(b93*(laD;}KhE9p)17inaQa?p&jl`dNqdz9-=dQ}r_hQ_TPU6Aqibq-k22lmo#Qxbd!6$|57 z^JPrREb(h7HKnkYc(?-vC6S*iTEswrav)$b*oG`pU-ZiH8lg1 z0+kF3X?Ak`I6Mcci4ArtRODh48}kT<=jdkl8vU+#c7|%fUSpAOpsT^i5TMGV(E~u% z4$9z~7boswDom_q02*cpgd|z|(`kkk>YhwEoVnSh9!17H{I>VpdCa8vyF7>3Cfgr& zINRsh6I;zna5J#vC+n4PAcsE2Xk+{0B!V|K4Ib!@i%d{Qb2KU*R+VaA)KsBW4nM*> z9QyiwZfVNEYi1Ho*J8_onRgVJ@N(O9Z<2s#Y~8&c;*hbY-6D6m+)o*+?Fe^W)&`Zd z#8-bGHZ{?#Xu zSxGApJ>!76&m?qen3m0&rC%&A36&-|6;&1GB6#v+R!gQLf@Id}m+vuXyu5Z*q#d1g zlxp$I@s7kDB35{Rv0b#aO?uIrQ$P$v&nl)CD#_?{QIntWDTYJWI7JPXEq|gHAmPf< zaF;PEZp)MKKD&sgf|nt<7D)#2D#CS@?pzH!?l2=1@Uzi>dkq|C@5;A?esbZCU0644 zY%Q#zYsazBgZUfQJy_pP65{!vh6r;1 za2H3UH#fg(-J-Uzc>*+|xSKUnD~^+&{nOq%XTh8GxUxn}2m2>1#fD*1HUDRe15(ns zX~Grq&RU|`!g}UZ(c3r+o)W|lhDsC@O+z|rUDCp@#Rw`IAhr2J{KiHGUyub$e<;|O zj1fr_0bzuR`lB~ngbLEdsKuFhoVrRSQ<`7hfh4K)*7xm9WcBv=zFVVH-sAa5-56vo zO2}|la_(m1Qn86vN^8Tx^AzojF|4kBm6WvW`SQD@lLsEn!QrHnmsVGgoH-Nj*kVmy zz_DeE#r`|7^U^6r#Fr%u4MLzB4iR2c6Gq2O3(??uNHV+1;*fGTmMExbki~u@@@#d7 z>Jo9vXzB{srJ*3I4qb_1ZazcPm>a>z;;698k&x<|CfaiSwGtKiFr96yVyG?7VBUvG zor-QpdC1E@Gc|I49NxliW#1Kt;Qf1JlORYa8~t772&c>aS!PhD?pkb~$+IP8z0)Y8 zbg#I=#$$VaA?g?6Y_pU|x{6|S-H+7{D!7XYf1B6*&J=@`DCsOKtFy_euv3EyPYC7} zy>@cOa|xw5SJ9M^FZq$y>%P)LOb@os$n6 zR?mn_zLVU?$bpUED}_OG^?U~oZ2^^G8-8nQP??b&!B>wSbTr+a)B6V|_|*C#`!8G* zlWHJpOfji_V*e+WIyMlK7$rU zuYKy8nRA*;0Wa;nIx@;Kh;Du1aA7Oi67t9^0;k9)xV<+uk|HvHGF)RYf2)lhaV&I} zAB;OBqP3}D2h7~!R?NMTNm*1EmTWe6;e7MsN-|CFw#HH(p`&Qj+)8zru}3#*ZqJAJ zAmA}ZVR@d`(f|Q)&H~|F0qZ!t_)mxb#H6jYv-dla`*n?S zj!zymH+g&LF}|=|MR6A-Il!JXBzRJq+6_oyXNP`2)PErg6^VhJiar1;Zy$igvDzB^ z%P!BdsTwYtu_igO_fz*Z>^};6%ZXA_`?>?OA~E-uqGD3HxG|B}8FSl*HKHZi1fWR+E&=Kiz~Q>YVKlG# zTky70@Ox>??vrv!9%UfK+xNYM5}Wz~$yYgb*`FXI_BQ|yy&)jUt7v2!5v~zJHt&@c z#yK8J2+p>$a^EA70P2v$H6EEcQ=Q)3{Y*v0X58n0etjoJUy12uAopnKudoFU($txM z+^f-{@)ut|zYNj{&wS48*A-u`JBW__S1<4|ul~b|THzFGn)8JCyyTn1%5JMdQ1{om zelBX09w4tN;KJV6T-qTl=toi79MH0SY4ExLky{v4+9o$THJ?Iie3&k@&GYS0j7WAO zeZF7VUXB6t#0gbs%cg44yI z=qqbX7Cq$S%wQU_%dx56DonsIo^@rsLKW7|k42pCfDVD+>y4o0uVaQROLUKSM%h`( z60MHk`p4=uF;MT`FHU(l(BHFt;P%SqDVwHIar{F^v4uo;V=lgW*6`XgN30woc|xf zX71^|$6%QnwP{SqE@v`H_>k2(@D1G+@!?jXUB3XVa~EuatpT%CpWfavQ_{boyF9G*Zp&B3mR!(s=nLf zMZ`BizRO5FG!Y5`NHtM^AH=~Y8V&`_t4jNhpc-ck8zEzBMqq`OJktkJ!a9?oMxd^%hoF=v_PnWZFgoa8fabo@2vlPvu?jydmlxHz67LMBQQZVmX1~%$j0V zgAblZrlH!%Q(ykxVOn#==*O37Xt&h(0Nm?zVF*(t&mn+59+uD=9*u&d`Lf)c|G>!d z{d9P5`5gOa5M~4tJEFz?z@;=Euh+LnYwm838fm^K;Y-)V*!@Cv1-SR->+JMV%EP~Q zJ>G&i+b;x+5X$$<0k0(^fla4XZvu)2j=^})q_Gn_04g;>k#dOh(k^2Dd_5H2X5D9r zZ902x{NF?pF{Nc7n?Tk&yPBI=FOPbbZC(<}1-5r2GUBOKG;SXrVA3%c7yo*{KO^eu z>hyi5Q-)g8UCb&jJA2gUhwtQX!yNPBinWLcjN@6*tgIAOFzW!sinw@W)yl`Bf+5lyev2e6nNqTe0MjnCfz{%1){HOw@&ah#x5yOZJhbvV~On6unG-ewOZbck702B z?1bn4@S4l2CEEAV*|tGkJGJc)4V8vhYk?m%M||G63i8hHu^}yp_cw1a!f#||G}QCH zWr{oku=3i5p6R#_(#PHVR4ii*MpfPs;zt%aH5ZvAPu1H!fV@`_cnK(`>0rff4IQ8h@xc*`yMT(#+AcknBO=Hf0jcf>dM zq3fJ~M0yk5^GooZU>m{3dXV|GtjfO5$`Qp5eoc9IR;b1>;IijtXg$1A{UI}!((?S= zQc}-Cf}os~5s4PvwpcZRYFBIHtC9oi^m5Tsfys>^?bulEPu7D;0k6Ghg9FmxdiZ~d zKmH`N)mC4A_$&Q)e6&HeYz(>-F{)BVRWw^Jl{{3Z1a6)(!GaYd#b%JZm=2C2I!OYy zN(%-XJ&9Z`jhs{l-IrEI73{iSzeUnB)oBb`$#0(WJ=*j>>$HDsyjpjk>R!xtul@0~ zA78!6bQV`Fd>zM|Eq59UH=^%dl3AyE;gvI(LSjHy#F6t0f7(^O&s^yOIot@I4IaT^_) z1|<-;GyX6~x+KfxX)y91I>C%A5IdSQ5RC1o)ipXgJv^!S97Zgv1_KEooM3`S@2$XZ z-#Yi?AC;JsOUU5yg8wYPP1ltntDC2jw6*{H`^T`g96{*&w-u>fU_>EpIFz$pc`0eY zt}fcFP*tP)f$p5ML+A>F1vJ?}P5{w$W8X8`x2IEA5Pz9Bs{MKlD8Nc?1eB&bAR;IVt8oe3D{)Vt}HG(VXS-JH3{R z4%#(xqAZPs+yx(TpQBAR3MO?**@#nRGZ19@dwfQM7kpoOg*e3VRn?07*e zsi{G_1^`YEe0`WjR;anZ=YR1}?a!H@-mh^!s8$R)H1I;?C(Us_;mIeQkpI5gV@1~9 zei=IC?Bj&cTxPCax9GIMBTJZ=5D2<(RXCcj6J$g>pxgY{1=59Cz4W$nwKrgbnu`h+ zQdomfrZBQ+vAwXe7sP6K6y1t;FH zcLaXKtqX~mi-|6aERZjg9WNnsexSv-tZrFOocMS}gfVAIwMpc7yzCnM3%!d_2 zk``>`FkT_Ah{hWGJkT?#T1p4sNGNbN=|IJ=*{-&zAvLVb5jEi+TyP&DcXMAJ^!&b# zQTL&dw;`kDzAA0|FdGf3caNNZUq6Vc*d-**Bqx#2+OW(bS$3vGWZ_xKU^X%Bih@*v zKBfqk&n7Vp9~pA(1r2k>LG6M}q03{Fj%%HMY8KPuj{*qM5alGYl{kiHI`@2DWC_Ne z)*R!kVW(sMkJAcIE%T-ptn%bVNKZc*hv67yU_20iywpf(LoOJzft4~jO)CGO>oG$p zkbiEJ`4QRi>cf41S$PhJFx4q`54CJZLXq$5=a zNPaL$8V!SG%e@dVY-*0NR6nXI7GEfW4tOKGwc#l*CSnmmc@>3wxll@D7{MshM=!<%1~ZbgZ!8!~+2TaqZ99*AG!E^@juaIBWRgim zG|uBn488O-Hp?hXvw@4$*cFPB_%D5!dnCh`akkW%LMU9agcGjJ+*4glUyaYqWC2cf zH>0CGtz2`IkK*N2P$u1-;TXMvA!_-d256ZDw`moxfQQinBn0}MpK>$`>pnTSq*%wk zJkQ-%=4XM-jY4D?BjD~RD4E?T(fM7K$wTQYhZ^fd5YQL?gnjr3nDHMyF4F8MFG}Uu z=b8sos!(L+#6KSI*{a|Ow5@XRG>E@ldO5&8}Y;+ZFBrGNio)3c-WP*RYxd1;w5_A4&+YOY-ff1PJ%jq)sXy@8XHmK;ybOJJ*LM_F7RvK^gAo6 zvJ9j@a|tA34NHz7b9{SPst*(+)aacrWt3`e3X!hrhFVTi>D38G9};`UQ0wm$X)BE3 zr$lOwJi}0cs^~@fe#-)tpx6GHF{BDCH<#1%cHumf>cxoP+*#DCR9?ayJYs*OO!3on z>@KWm_2)1gh#yj6%57v;_k*rDYoavZvK;h-DP%E-V7W?q;AVW`*?R>rqTRh}A$j{w z9}r-@xC;{}bartGmZ7EKmMu~hdue5lH4$xJNe`2W>7W_YtQo2a$zmY23{TXRY@4Zi z#GoMJOcWPzy9(PVvjrKHIN z`Y53j;E`>mh5|IwQ}11Z#c@SIsCqv`PNyXESRFl~HB>MIMFuDD?4Hn7F3qx10V?fOh$boF1^v{y0|aZwRE-SE>vxlV1xEs=YNmRJgYVoj8-G5ki@S!LuIO8 zhd?d#P+;F2{MDL|uB9+31v_%u_OYx`3+m~Z3LPP;ednZ;2gSx6B4x(RIc@K$DF;-x z)^=?Vgek9v&#^tvu?p_W|8RZ=a{rs$dR6?XQbs+Tv(zNV?jE4RD-7s^1z(O!6u3V4 zqV(3COt-fg(~r}-*(;lF`0T}-FPI4X-L5YHlc5Midk6VE`mhy2rZ8I|p=PV#dH=FE z)i3{!nTm~nS;+`-U<3iztu5%6l8Zu-P{|wgx&-ehZ zW1`#Hb6>{N`V06g3-r!@n~ZeK_s?+SkGnO0YkXML{{oDg{Kw>t$cNTymdQZza>;s{ zp7yskv?l(S?KDilQ78}?k)xJehHP0Y7X~oe#St-pA@h%By`Ls`S+X)wK=mH&!l@!X zEl%rVavrzF`-1aXvI;Sci1Z3RQNZSVeYLI+FvWVqW2)w&d<2%?OX^6Ln+ZtuqxWrq zD&cSazn6OBn;7kF%ixTmEzp8dF23-pjViqyC2cYJuuQG?)%!LL`Xg^J#ksVJX(N!5 ztj}x!SYZve&98@Q#O}miJ%yB?md#OGJSGJn`?#J z(A;FWYnZ>lcD-pxC~)oq2<97;G-SlZj)4W-C(wJkaYh%+-|l{zUn+ms zdfwIq0Q&6+v45Nn9CXKcITvb4a>hdXWl5;7#`O64B91V_< zMg>DF7YQ*Eb5SmP#YDYh+lrJP@Wj`8Wu>Vja_h4=szo;AT)vKPQ1R3X7Ax55ph;V7 z01rP2@13)20XJtb%vw3YkN=fRP>xlInb{=ynX-vKB5>nh+@0sea|iAgxJ1V0M?Cy7%u6_?pPByW+*0F05RRL zzl<0U8H^!5x7Ztr6~O#%U;*?adQOQ}bdum2+SVKT)didoN!*aFKWO9RE%jNzKe$@> z(rr4E@AG8xz6%Tsl3B^qnsi8kxMv*S1e*~~7TqwgI3(@${1%W~6(ibaTFprNN+7vy*& zz9VkRiF2aY6TRbZO!hy)s<(x`bw<)phy@=guLfh1HzVedQS}E0w`Rmvs3=0H$QxFy za>(j7CB{}hkOZL3x?HqsQ8MWwH$2-Q<65P~;BEq_fv?5Vg&{>-rV?-{9LW9*xaazg zt=L^s!sbv&C2(P2Z5FwK<)l%WQ6(0|(eYI0R(pGqT|uPJ#+Z1`fqyZ>Ee>Qjs~F_?TnyTlC6j zHI(!YJV)2pFi@H|U9e#c%8v{{Jkbh7WIR8S-6C=e_dt&PHJ_Xi?<9w#a&9qpE#Q0t zM||#6ZL+ZOXG%0!b2ls2EEp&a5yPbL&)YMvxu>q&qb)IMszqPY zv1_P_tI$u8t~^8u+3Hn?MWM?La!Ay3GH^mx;@?|5C&8SSh(wW;rqvSsvGdt)x8__d z#?WfNh^RakWXd$w_Py_Mnc%GsMKX|lYAIz(XmRy&`xyvSzXs#Lk&fci!j*(rT+l*` zEm9@W-rLdd=o+RIrpKw#bEfpc+x-sByqX`{R+Yod?Kmy3x6Atg6_E?$g4(|dYx!G- z1?2eSe3!eSM=Vi|RT3hJ*TQs%bH+=;N)tV+jG^Mr7wlw~=xUQJ(RrNi*b#Qawduwq zbjHnyfb>UtB{xBSD%6vKIX3;?`N3`36HgU~Q;CCg9OF4oQw*65hOypqy}?U6SrQ&Z z&Plz5FZ`E^z5~}B-D^&fC-Cje_{ry|y792iS2mwzP^9iPyhY4O_o*oK*agkq{2FuT zWk=1woa=hLwn4npI_B?__GZEC$z5r5NkmF;UqV&zs#-vxF78%E%wWD)3|;gI4z%%- z7B$7!7}Il!J*VuN6D?+%KAbm$H3L_rj~)t66DJK`A+;C-KFp&n8{Ht*Z7NE7wN}e^ zx%mZes)pUHJPS1j0avmNrWLKhYb6E{JlxOq27C8-*D$UFI<`bd)t?c$TDG*(Kj|;S znn&p@p!HyIzPvZwH&3d~F8Js&bSbT=^{(Yc^nCQ>rJRPp`O_T=i@lrwabo@ZyIEed zv?}y1TwX6ufh>H2&fMhZIr+z4dS)9NHTAH>3d?3FE6*~Cbal_@KXf~gtQhIV*wApC z8AjtD0fQLib+kUq8?J>zY~q5blD{t6+W1PFZaA5hX&i}4ykXcrJkBq~RtmrK1GoAZ zUy;@Jg*MijVmHB6JeJZ0=@WQqi&r9-HhUKx@<~U3)D?rtI&Ah072_mZFi8?Oevm>Kx&W$F9}VW6ppcrW&U`F;GF3ryo1Q@ z4}=zW@Jf4$HBGZ7h1-V**RY836>#$yGC?(CgtYz1OK0UKljnj;39?I(sK;0)@JTo& ziIo0Lan-4Px?s75BljRyC1~Ts?XL>BF0Ula5kQVoc44705SyD5{+33}{SlEcn%o^Y zj@YsXR>XhQ8DZn4z35m?BvpTS(+Oy%;Xl23XIUxss`oX1Ib`|yKK=jA{OuphB=pXX z4O0EmfJZ#NeC%nC(+dDmkxoA3d>-VE&}CiOhGmMeJId;~9PauOQT9vTBSZRo)3*CO zodYyjr?FqW8aGV7BtjvsWK%>c%y2E>Zu0l&jfV}G=+Qobbt1~bPe|b509#oxLN7|9 z{6$hHTnUZQ@omIs*uuv>ZspEO&^(S=JmQA>7t^=(UE^_N)Zu~SWX7IDS{udlT&Z3Y z4kE(0vvO}6AXswyT^EKS=7n`)@D|&h-6>d+%rYkQ98Du5Ht} z7#5_0@|OQN=hSn77Fxd==7=EqRvd>7-GWR>l()x>lC#r*FVfiO(>tQ3%P2+p2QXT5 zlI$k=-!bK&{!vO7xDg6&3YTzna&KmO2~#XP5<4|=X}l8Y=UCN*^v4sqDKl@`#|&F4 zFIV|O+F3<*h_MbS94;}r6I!N{>l0(`D)-X7Yd}9nM2cef*Peoc1{xKmR6FbSia{VY zFzr<`B9TI2H+DH@6OaXQR4ZB1p5EknZ|ZXXR3|FxnZW-SYk?aZsa^sF$E@{$=u0zl*)~H8AAR60S317<1@C`8xAIxp+dysDY-2X|d z&R}J{czE=$aL@hfQupFtGGF7?X<%POV$$_`JC2PEW%yGrG#9OsmYcYHVh4$KlS|F8 zSSY5$rT!FFv?*&7oC@LKbuY6xk{i%6@VIbGHJKs<(atE5pyeY>=(e9M?RPZ;rK+nb zaiiz+Tkcd-+mMPHv7!luZ=EqMFlmAq7cT^!_CX9>e;@6L!*h*HA^O`|N26xCAHWR# zVGmZLE^5rV#Y=e?RGGp5fFi+x;g5Y=Z;^Vs&?;cP8TwTx^lluS;y+^~*k*1raRco* zMGd_HdtW!P)7emnuDflGE8O_T$rmZKZC`p&f>r@eKtEMKk{M;$#54<@l9r*7+h5?$ zy?|DH>&w5lOZCBO59aYtKS#^GCn*X`Hcp$`GDyz91I!f_ThS@26%VV~%f=Th zC_x_apJ39E3%1U%h56Wb7OnsrDqYwe5!$(dd|t8M)il!H+v__`lpkKOJtm?IOK2bm zg3X!`ofGJx&a@czB1xyA)5E@?;cgm2Y0v}G6@v-pY^e1$ zdg@r>JKT(yXCqy3M~JB75k_f8rkM8X8E=@&8UObQl1n5g^QO!hd-ZB$`7{6^@j(0| zQT-rRh5I}ljaVco2(;PW8hV0`%>)`YLcZ|r9mTVqrEa~X15hpnnm>F2h3}572yGY* zUq$RGNKeO$eE|hZ9&zdz%x>OZgCIZ3c-2(fx_oc~L@ZzY$%rYd=^pj~7){@$x0xRx zw0+KK|9XjD@*Fj%_P^@~bA1o9f}}?IqdBW1I8id^zyh@w29j6JV_I&EGf4l(+Ia`V z8Mp6#clF-e>K4&E(Id)Y3DHY*tBW9dN%XpERwwEby+#lcEqd<~5uyaqqb4Hf@t!#| zzccfm|IX}xhFR9`^L+37x~|XFv29TJ8%@FvotGWh@vjIHN7P37a16Zb1nk`dsq|eq zHuW!haOmbTYXpdu;@DNwf;rreZMCId6=w8+PUQ1{`cRvtY++uxaimI1~O4d=+#%% zS3EPH#7$+LjyJQJJbt>p{%of=H@<2p5}6X$)6yLNH1m`NGCAh28L92a*yXr#XSKV( z-7#FoPs`BdQ2_}db_YiH>ik)*Jl3U9@RP@n_;k?e9L~DRJHL4n6Ln;9@Pd znO#rOa0MF_(To6Gd6V}3XY_-=$v@`|bkA+q_mZQyGzs$gb}3}7mR;6pH#vJWQ;M*s zO3rq2XLTQZEil%gujBf|5DXGr&g(H<1<*q6dOv44J81J`PGqAQ5Gkc~GJvO%#<-Zp zp7A9Ge!VpzIk~O>US7mr=~*HLch`|C=phx;Bd;>*a1xOb3lgFmddz`-VmC*U65LO& z6<=TypQnTA-3#xtmEmVJn-rL`Z(j^?#*wZ3!GJgUr?n8({a{YH^*J3$+@!Co#26V+E0yZ^;9Aegn=Uok zXrmP`WalhX0XKf~s^U(9L3YQ}SaJ2G{L-_$%>+lgC?$65 z#{-H?Uer|bs02jwq#f%i&6XF&uHL5`l@Mpgw)^Gsw@69?qnQ6l0be1XTWex_GpVuQ z-0PCEPm}BcKWjXmz2~z|jf6KcVEn-+WA>a_OxIif&dbBXt7fjjNmRsdn zECn4=16VVJxz)$^xjbgx7?iT8}XiC~kzLn3@N0mng8>#k#w zgQTojOp#>am*01i+-L*o6{>$|t?a+eC7y*xQfT^ar``!jT^y}lboo2!Vz^(UK3kZ328L%2uRMA)k8)yr=E z@sSMvaWMnMq%qWdZqRUSoP+qah$M%j(!_Inqt`3m7S%kWSQgEwz%eO>44D^Od`k*M zB26Vs4G$&|4^KPS+Nz!(%&PLDZ9wx3G*SLf*=S2ms&V`4E)192;qFmXKP8U9Cph_# z;tj|k?*Pct>rs{U8y&M+|32)oR?_vvI0CjmuP)yAFs&ixRh~gQ%CO_&|9&j2lJeKD zkEURN{*7ac7M|^&9-`1kz+9DH$UUVnHNR8$F!R>qpUnVs1qWd~O1QRylb`)GxSm^L zB&;e-R86@+sm4ggz)?j$pPt?n2O!4^z}-1m;KU1$|{~MQ=t| zTMU)pF|b;0aa&8QdlPVp-G1*)7GUQa?jkLrMwSd2mSIVk=Q;PNsRh8MjQAohk2E*G zJbOKMS?EAk`OgeI1FNt zV)fJ;z_9OzX1c%1P8&Mr2;bA`md4r87vuZAsVl0jfACgtHBBJRjx6b^r-y~3&^>OU z;XS0?H2Dv{fu&jgm(=HFRFt-}Sp{j0ZQMSsLWbxkf@a45rQov4@%1&ma_i2?@fZ z&47izX>YXAbWVcq{$%SE#+$GPdz(-C=e@@7(2$X*30(KvA1DrSs={R2B%VlG9J|Ex z<${ou_!AM_lRP$O$&%#qS|a8C5IL^YsO(&pA0{A(0-m3wg)M6DP_$o!>0PhAlWB*< zF9Pu={BB^g`FJJYM>X@jzdaP6{@kN{@~?XPBluTPS(}3fh*IBnj<3^l6<7EcABnxE zSdh&?Xd{IZy71Ip3?7_U2d}tPt;-62B+gEEC(PUCWZb0Afe3*S-jA+rm@E&rUlABJ z7kYM5b4gAozN1!qZ||@G>_{#VOQR~tXxKpWnk81^q9ZT8^F8&7_78Ern@7p#ho z>**Y|-eop#_@mhT2nmvVC^GgSrN?cA_$)<+6*v5xItE&x#99ZfVZQT>G~VD|9Q%iT z!YnQ)$p-lO3F)3-v^J%@pXa8qwmjEUbOn9`k!9;Rb~#_4jc{FR@qU#bqA3F`RCN%Q z6UCL0r^&^srZ6+pd@C$^I^DLnOd7JJwF?Rh6FLLz88f)f+Xse9Ek5KJka{cEoL5xj zu?R%+Squr!>5I~TxMwp3Mkd5k@t|#;ffV~_FYp1A$om2v`fzhDlW)5goeaS?tkWlY z9IQ+cOb){b0ryVoC!4&j$g=a(c#t^De> z({VY2pH~TpoHvW6;SRV}r_jm6KUrAYSrKMds+%xSrF@U@Y>xtpA~x60cZ!!0atAF#n)fbFW-O>l0IvwjEPH_nsi2KurVf5+#HTa3J4AMn>GUTZozqPACF&kjt3f}f%HF3H=l7d8Su8PgYVH>gI7imHDz zPwb{l0RD@xtkmi8IbcJV5EgOo&lfB#bX3r*r|28%TWYdp)|`@bSu}!`-fCD^(uf^5 z0lPCo4+po=n=)4~h>no8+Hk?NTB+*o3LReFs%Zwkj}}{yVnT}JD<*{6eT~jE#~x9O zipXLO`QdoNlZPn&G$)lx?xf7Q4FLz|>Eb1=>&b-r$xQo3fq zF5h?CZf}zJlaJafCU7r^XUBlP)XjSh0>csIK>vG43dwFITQokUoZVA40XL85zZ3I8 z1Y+~IqZ2Q0wz;$&2VXcw#%4^sl#hMtob95PzD)+m(*|@7@@FB{7VsJ(*smpQ~&&oX-BU#;?}N@sQmH0vXxE4ap_eI z$r5qpljzT^P$~U|F=+UfFN%;FWuNM2ooZtuIm2|qF-3@_!R1$JveCU4C;6jPShxjnAfQ~RyuXHrJal81+Y1fsTr(Br}S1Uf8xuV^yru?*bs!4&}Abb zf?gas`j0fs-V}2t$X-gWP4C!cI+4sdcx5n+ij_2|Lgvj_RZWg4vbHHls_ojF zZ(V;6SwAY=X524#%brf;!um4CC72#Lh7TAw9E*C(Vg}RuxAsuy+^;l9Rj=1&)n{f# zh%>w2K$7trqWEQ=cHB zRT-0T29+^l6WgM&_d&icQz*YBbtC2$5c6cxCM{zwqM;)`5FqF>2x=^{eVWx|N_&s3 z_csUbFAe&@{(8NIl4+a})m^C;11~s=i1KkGq6@~5+D?^2P_q}$B3E|HK|+lG=x#(Q z#&6MkA0=CzBYTe`^@Bm=mVY9RtN=dTrE0GJg3W+M-wVu3pwq0%#{P-B#E=_1x0=0N zhhGg5r{&wB5mpzijZoV(0dY}1+~Zly=S~+@q-WoPLIp&GxSZFaTigh2h;r7HO!*Af zOm7^B>-uf_^O2aEx+!vHX2zxqX3u*i_&vP=ohLli_eXKE@fjTMeVNJQQ5eVnn{712 zBhwuHt~FU0+uIn~g73sgob^0GKToi$I>R6+XZ`tmAgO2hK_Sm$d}z|^Z3(nzCP~S+ zvDoBKHOV0BDQRfAoHh^ZzeW4QEsP@A$91ct|7SE7sbvGE$C$`{8iynQlgDRg4)srHvy z%LGbq??MJuapmkhHS@^h^2x=+$QSf0!kH}*Ngy_x*H%+IcI*+%xMHLOZlq0k z#8egDn~rbk^Sk%#Lf0vq2z4UJ!;fwI9p0!UkU{o^Z1LV>laj9=EO9gdcpSfvTGbzO z-!`*`BX2hS+pUt&ztt;;zq1)}2FG_t@!(e0Bb0}RP;blyDV|VP&8RNYV@gpu_vKle z&ol*Vz7IMEc*RBD%;!6Pg|X$%miUi2(nD&cK&2K+@oc^YwCHe^kBCDwPZoa4jfL`V zdZ`u^`Px&AI6u2Nr)P+~ffHHRi_WY`v1ga=jZN7Ny>ibkfKQ8T-Pp`L6qQ@`Jtw3_ zMx@$b?pbs*NO6j5dd}>>`1@NcRgaJ?oAp&mBQqTzBxspULIfmb@&em=>;~#k2DWcq z2m_DG1QumsCr)dR&A~IhvR9Jn9_cUeT-zU~#k2k6V1E=Y-jJB6u=Y5UL1!FNCH^@v zKoN`^2R{|>qr?{CbdzyFj;R=|M~9Q0xNgLR(L!**Yt+FK#@cH=GkWzcdp#~{~OY*D5eR;CO~f$BENX{s#>1_WQ%ZPw^keh&C8d-P5DIukn{M z44KB8J#a>sIk&Pz9)Q~Ot5mb0fEOLSptG%zc(pAD@$Sf^NvAY5dbyO$`5R3rP?bDV zOy6zVS4qYDCT#8xHF3q;n_kR?VR$hlm^x`%&VBgC#?)$idPq4*xHGXL)q@kZVTj*G z>-4GL6^ntG-7K&uoSiex=kpg;nup)Dy>~hA6!%I*_h~(4Xsr5ua7qn&t1C3@=_4A@ z*2e}OwLwQh@PQXFaN~lJgpHuR1q})JC-bunOU6`@_Z2If&jf5K*HP=t)rd3$k-Td4 zy_p&;noSSc_TJ-_qS8Vp2UG6MYKg{me5Lmpy-?>B)V$G|*2V{ghgT^k@e$%;4 zTYr3rU2tPCZiLIkRsuVR0;KV|S=Y~{BWOBLnv*=kOjo0elhUQdpQ{sev4b-CihIhW>-!}<7<{< z=&XR@EPkGzrY~#6J!?o)k2{mcOh~a~p3*m>$abC1*9Pg%Li^rX)FQD*!0CZpmuBPV zwq1M)o}?vmBJ@eo7E9v`&6*+3Z4P1!pP3{ikN6X@HZP-t zF>SKOMr()R7TOKdVHYDgyD7rQ?=2ZwfYbvtP`;mkqIsrE%!%0SC~w~B>v0=bPHj31 zeWFC2Mf}<&?VM|O<)sy(ufT?zXOikjcKC3K?VpHVomf&K6s$Tp0v{&tFRF(5*$YS= zbEG-3zuZU;L-LAO_zlD)$RE5uzka6s?-%dqhyf|wGCcdu{Liyd!U5Q4a*Sm}Uk{p{ z?ym-Hl%^=kidUofy(RWzLR+8oLk)Ab2aLjo>f=#t{j>9lY1-Z#BOs}=STzTEDxLR^ ziUU7U*T|~fszXc)SKpwOokmYQ$#8JIQQUI=aCU%qlBveP7o}4OE@l=X&}B}p(cWaY zW1kmnrBj*gi8@Xupr&Vm1@9mR4rvq#;|rqe^eJ)%uo(085A<&yU9$6wccb-__$lF) zrtVkRnd7LGr#4p3edRz|4}=5%HwW0uRf3Xhn_kPYcLxn@hBeF@HH{^YXVf(CJ9?Om zzlsP#62$?V+I#6g7&A#-3XQI;e(;5faPab|)QV@{$+ATw`tKdO{IF?Ms8Rp#65uuf zHeJtjSE=JYij$B?L}JAHT=HQz7Fiii7E}c)A3>cS=pAe)>8=$As@eLA%jkYQll8%; zG)ao1k4tFqwk-@dbt~--=x~+)Lmn@@-1N?=@)@v&Z%Drk{^_Hw@(YNSr_i3Tzbg59 z_tm8Pl%*Gbbuk|MB+spZ(^PXh;)rnBLj*gde-ydG3@nU$f(+>eVvxL<159TLS=h ziG0NX`dC%5BGZiJ(wDBS_pLPE%jIujW+_UPkVf3dSUvGn;lere{q9W+z+<%klt4q_ z@CM7ZDQ|Rm4>wQCa>S)_LSw?PmnZ26SiZDkbH5l5iZn}n8fN+(5KTTjo*&w6K~Hcr zn*B@9>Qohz0NZ*rUJ};bZ89a;dg*rRJRXXvfdGMD5K` z;s<)(P(jNu_zd zYo93q4y%J2jmN%HUwsU|m=`vN9{1p7;GNHlMyq~(q&YFYaPs%)TyhIvMgW9Q#F;0x12#h z+DZQ4A)D$aF`GRw9@8=4>swA<97#t)Yuze;MVks3-7)=P$pl3#hF?!0YYfkLaWAOk ztQwFHR_DAsy+(M$Y-LY?bjtntj)7A9jh?EwsJmDD>O_ZlKtmYM8ZhB70e_cA1hKT( zWQFn4i>dDbhhh35abiU$ZDBT{O?q}G;ELU@$sp;JQupI-jWS5{`tJ1&K)e4OEKFiD z5&vVyl1neIJ|NMU$PUp=VTg5nM{$*4(rK6Xv^A&hY0Nuz-MgYRuj3p?4WE*{zts~l!1(QIz;<-dUNyU*n4Q*U^r^;Ge|I0* z1A*W{-YY0sM&OE|0Q98l_*#nkUUK&;&@CUh;j%<@s8Z#8(Q)Ty z60osr9>IL9@uOZK)QcO>nnQG9Z3rDdSv{pS>I2a}?SdCM>TsnKQAU-wsM!oqI9Dez z#MX({kx*!qq};Q4b`E4;6yB|PC$g3>ZM?SIU?``_`(A*94J>c1`~k0$^~F>ISqCaL>(&pdc1}s!AKk7iO<>Yj$ayPczoWJ%?3LD*iT6RLtaMP97-$p)!;`_bb)WYfoOb)gq?Cd zm!aS8j-K*1wYzw{wkUkTE$s2;CmA*7q$2jxAxl@(N)$Gm7S=D5Z)R^W$}rqPf9NOm zcQM>E8dKzJG}<6}7w3G41*B) zVb+mB@9ojr+lWo@^M%ysJ!^}Pa~E}>6i+h_9f(gxC!p z@6pf;l&`j&4Jdefdhwl(|JL&^QFx`W8;3nqMx?yIw7wC|q^09sR8;IU@_EEw$Ov*~ zrRf}WK<5<50)9#Le1@L59B{u9!undS(iLjyWj|ru+J+6Wnfl*k{82sMOY(1{eMg zV-@?tu4-nWfnzH@XVezVOta<>H-foQmr;4CC@(Aav)QFL_Gb?Nj z{I6B*-RsuK#shNZU2-`D!zYA$sn&(-jbuirhqgCdjh~*##SY&C};=UL! z2-=Tx*;GTz@NL;ea7Ig8@k@U@lRrY(O$HOq?+A+_)3$*Nz-w6yYKnWh;F%;86UtRz z4*j)jGW#%`v11g2lHdLknh&(vLut(YuNZ>8BYFypvv+NBsh6YDCh6e=rC+AU~6Y z>$!Yt9t4rsq8kwy$Ma?KK>p`oeEl4@rV+_0Uk6 zOY>ZR!5W-JWtwIei58hiLa~#h(0_5xno!E$?V1$?l}w`&H{z(GzL)4|`73#IcGWz& z$29qgQmt)?EV-kbGm#!~1>JN|0sL$W5%_1qwF{8JTp~dT6Jt@zo1#YJ8<4h&=x_}Y-+d!tLE>~t5=wD7&g|= zs0GK3mQl)CvVC4IoX29L_Xe4bfLb)lMSJ24O7QUad6O=af=Te@Sox*RTA~9GH%DJL zbpoc{^ymm6CCeKw%PS>O?05`z6;Qi`Q)egxEQh4;J^hLsCnp@n?#6wlM=@V&I!4V)B2oQ3EJ#l?^%Me-Q$h^bI< zkNDUb8JNg&i;wWYsLBfBC{7;ID2kI1lSt_pmWIINaWDoX zn9^#HI`)ZdgBFjL3Eg96WB|Po=Uf`D3Pg89Ap97=MU}1G*pre2xwB=o&mN~~nS=x5 zk&6_VR}&}bo-k(&YkUkXTFOz)VnNEYb#3omWvDE5olJ-@tK(m?6R@+JSSHAVQY^ex zV@$^e+gLCB+$l`G@5Dc@MThYEX)yLt)C3pLI@VEh9%k{g!p#ib#tnZUq< zvc~=yTxW8JWcizpx{3g+4c*Ra#+*aCha9p5A`httdP^sb9?Qf(oA2cyq=wYzEIx zJ_x&hW(s!{rD$AT_(j(Dq1<4v0|ec-TECE_5RQ>KuqvP(({UW6L|FsrlrGkZbSrl1 z!0@L>U6&I-B96JQyq4!xLBjM+1M`NPudh<%MC$(cg(@3nW9ToglmqPq3+u5BeOE*$ z_b~o)jP3v4V*HQGR{rrXdiieAh&hqdbc$`8qJuTX0aKP`zvd( zL@cC?8VqvV2b|0<(L51l_pdhrvtZ}x2!DG$M_ZKXt3Vs1IIsDu7hHV-{Y1F-m-v?n z0QL-+74M|_T2R(%;gJI8VmcywgiHNh`qUIbAu{sB zjly(-;;|IckK(Z$Rgk#w+Z-6e?EN{oOkaHug%chJsL}yD9u{^7E;A4oA*B(#->G+ zf8mpVjG=L0ws*?gipESOUVNw;-TS!Jau!KYeH*! z*6`H?Gvy$~UGy~CBuh%`1JeR8mYx2<4GWX1Jt4`9U-Isbp0zD@*>z+h*IwP*00{4M zifQWP+$y|&`V{E;_<8&1C?!)vqbTOwacZG~7ZQ%(AFILIRCnx{yxbJHzx?g12_4Q>pAY2AMHv)h2w^HHQ3m?8_^nTL z*88o({`eIg?-k!djb2o$B+u^(s_ ze@Ne1Bupj1fmQ{7Vf}Z)`i0%V9JR*K!eewAWof5-*_aA8Bimsq&?fEmx!s7I_U25Hwb)u7lEkJ+bzw^^k;U9@>Y<=JNJe!C8hb!_Jv;#PJ zQJYyybtHOwd^r(HHas}jAO0Ltts3Ym(7fMfmj5h>Z^sAM5SY7meW#XJL8~Y`Vjkc> z8#;czo%f8TRfK*<#N0Xv-(K*$3qiPx!w*g+j%|&&c@~TvPO|zj*0|l!&`N> z(I!pIuq8jcB`p^5Dt;_67&W^lF1a?_ftrnmrRLj(={Mr3opzbTSf+PIDuk5<9%^oPB($)bmmV~F8$G+~=~>gJrbDalXY zSr`R2TAI% z`0e3hWM`^Pa4%@!5@;YtM{gbMm(opTSrNl<-MYN=il8!#0?%L2L62d6!$VFb>^N7(6#D{cx%G%mfGE%Gm@pOp84kDna17kSYekC^4)*J|3 zg z#@jfd`L8;r43U%G;M6Yze*R@;xOn!FJ|nypQC_F^3E%01unCd)`u;>h^|sv$3=5*b zS=T;&ukO4}Hu`(qmuuQ5NMnE0eeUV<>Y<1erVI(eDlh;#W>g&()jOsBdslC2C-U(o zK{$6s4)09r!p?u#fFkclNAqjxv9+sWve3A0AGi7ep_!@)ST2g|_SCY?bd^N0bKz~7 z&UA%u{e+PA3+#cii~V9oW(EttX;dI%MMxVm_^csXTt2QCmTBcb4ISw!b`u7>Dj`A`Hfv)pDmhp44D(h(!g)(P&kik|HQMNMdN`M=P~9@it6b zeYsL%-{W5DbG$h>xX*1#51SWgDQkc{Cv^M=dSvoqSIqz!y{}IzuZ{caEj7GYqlS&D zTBPa_gp_fxJv@U^si43JEn%fNot`+Z2Xk{;QY^#0($3d7Hu^g6Euu9me#Fg~Qcj*7 z9q()@o8Y&JaCwFm6|~}ySz?QM3{M_c&q_|4!rBLMWdk4Dkq|{p`kG(Uz@9zwRdadeTlr zG*RtnAO|^ImrgmK+vRrO%ohf4l1t~hm*tG^X>i8B${yLG5JH$FejpmkSz^$9wkaZ( z$yxX+=Kj8OFOivuf#rB*Z$g1zHqst#>1#=E)|(hD$4Tt8CKM{Ht%wOqfFFJGeE13P zIlI-`jKQveh}*93u(4Df&N6>zDz+Vx{d%3Ftd4f7fZNi#f$kA7jbsL4^xeJyl+ z`G0XJ{wI~9+w*P>AT4h}D7bh_#;qw*LS9;r8V2 zNuK_k-q!tp{38FKe+}{YFxIKro9x&2!Y?oRmV}Ske?5=8UjKF|%u3q)$MVAsHG9y5 zo*%K`Us|Cra>-H)W&5MKKi2P5T%Z07x%PU1yjJM59Tx8+{zHGnk)&Vi40_R=W}8|w zawJ+(^UYVa=uuAR9vw~(@x`zB;{NI$`;?pE6wa$p8thet5h(Hpb@c3Czlp16e$&&g zS%OXMjIP;snVJ&)i+y<&yK}uGjU}h~=9@Q)VL*k`Kcm-hgQBne1=;Dp`G6hD`vK+@ z*>b!<(ICa9p<$$so_skqXzmHr_%4i-} zU}fj<77z&y_mGR_kNG#at0Xo`AWWEclQCsha zv)@Y%ikGecSyjOOqtPv(u|EJHuOPVynx2*%d7%)p2Z<%S+Fj;e3`0fZvL>fV5`a_<)|H_ycG(|91qCsg>sVt`Rb2ZSsm%K)Kl z6cyOy4??!E4OMsaftv8XWfY|`d3$K-X`0xnhLKHV*-<1vDqa$Qa3;jcc%t!imOO~ z91xh(_Ibtqj#~1D-*r;y+NkRd)J z>8|_gH7}ZftOg%Bv>MdzfT2SBaZmSd0s8bAfOI5hdYw52;g-nms8Bqwr2ZQCqX6B6 zS8AZ1UI6C}%T(QEJ_?A{`i*mXH*NuYD2`r{-fI4+5tBa;A2a>GNE0L23`x&) z|4rR^HD&H!K)3rAEJdCpE#0DV58@XobkI1+@(jq)Sf8e6rNUo`ec>?(KH0{X7dPdP z61gSWa;m^x0||U z1|}}yU;ybLc1dNRf@k<+ZYLD*{{#ZM9}$J_aFmmLFsSd5XN#&pfAY!b*Hz$1{bR4X zz8qZ%ZB}8G(M~!cjB@C913Hpd*akiT191>Gpc42GR<6RVUOWZH9996k=?8S2Q%iSt z_1~>rFVt-UUPx&JH>U0}j-qYerOyEu=mFqdxpfUpVk>}7QbCtqF8+f%EkB@o`^Vnp$A6)2)ZCRH5*qg=bQ)j= zG9HXJnkT@FnU(3U-E>O+rk`(+VQCjPoJa2-%8W5y7l$t_BNp@eZ9ytL`WVuGpax(R z-XOxVPUHYgDFE!C{3Sr*1i%lF#}>e+0FT)x^)Ycul+98zu%`gXaYsqA7LBc&0xU&A zf{_cL23l5{$>?6jSgfjZclMr~M_3XuRM5dSFoPBrQm1;cBMb3@y|J`r7MT6(qqW3Q zcfIy9%yn=Zl3rW2_Rn0VPl1G+*kZQWq`||E{lvSF|308d+Mh-)b?63tqA@ATQsWnH zCkwba8T|6lp-j^{CTAtPC-cNz_{Vl>t@||>H)ZJ`V?ZM(Ltwpq*M0`b??-2|x+WD% zV-xPyS`(W}hkuJd0+)+BZvE17HMRH8+IP1=kHF>xH0U)q!krcaupz8?zFQ_&k+E%o z*bpXGfm{QgUj{3iWGL6Xbyc&dxAZkAZ4 z@;c~LvyoA=+3adrzn=Zi(Pjho_!DtX35Hh>T$SA(MQqhLhE-^xKPi96$i>zbP@@Fr{}Y3~*Z0 z88o@OQXx1oUh{PU!0_NlRUniVHNn-uMd7Bd;22NWF#t`Jiq-Th zE685lY_RrS5>k*LfI1Ml`-<2LyFQW^2?a!)bi@SdIO^8l0MRuTAhBl`E7D2T7?%(I zL?9WG+;vR28|pfc`89y&LH08v`2BuU;zK(r_91$#wfRt>k@6BAllH6YC^V({u&NCa z%?-V2eJK5ZSOB(~S>YtvrW*KGDjsifg)Je91K&^ykEE72fnXO9qoz`QZg}NeeE*;Id7tCy00*srLX0 zq?v#-)8II6nm+0F6J1hll!rDa9)}3Gf!wJOah+9HfP&R3z@hlsBM_l|BU`{^BZX

re;%}CYsVSxNG1&Qf-7y2u?>!e|h+hbpIVY{_&FKK;C1vp-zcjx2Uqwm{mA*xE# zMmA_F67mdQU(0h#lwG+qM)?VUTN$sObf3oKFtA9O znL>jgQ5FQTp9?ZXS0wyZV)m8n1?M(MZ+T}j+gy-I=h3K`=+I)vX%_Jh#uXgX_cdJ+ zxqr)JTu&pXK~Z$ov>>#y^rM0$dw<0tNk1$5v-V!jq&zJ5#~;g0(cFD&MJs{c3q6Y$ zY^v+&vLK462^g6B)c%Kt!B*nJv>Nl|z?X!K>CW^~lxf8A z(>!WSG!!^G%TUa_)!jYG=UXDxZCAxbzF zWn;Aw-J4RLVihcUdvQ#8JE0D<43^qcVtXOmMUesbJ-(y{dFFfB{R71P5ADEC8lLy~ z!P!(BLV1u(xuOiN&Fj-r3z^z; zRiR3&Yfp*@k1fVy{AYacdQsX1$xIdrC5{DY7`P~(uikoYv8S6!@8nN5G ztX+-SWpsQ?Au>?2?;90sVg+v4Jx7!!)g3S_iNehOCA6%| z)5nf|F0(t-4m0+2^u~md6kLr^Mw=z_qO2FxVg!YiX|x9d43ZD&SyL#_s2NhdDMk$}k?FWpAHrwds;=T(5zo^_nw)%HZ~geupd*55g}SX1Pr@#*-IOfG;80{o=lKM2$mo>b{z45gujm=`CG2H>Eq$r0rh?N*Pa|Iu zC0*rd7}gvTZsvsqE{YE8n|oMF2%DZhF}_nd~LBJ&T( zA@477vXz93NdEat3unrb8ZE=0L32{=mEWmkN@DphhisQ%Vs9B^W8(q9b<;UccM*YE z{fsoDOWgD^hm!DL9!>)2pkM`xq~L0&G-Yo~@9g_Cdv3*>rUe(v zRJXuvbX{;1qP{00a2}iqtgmXHngqkvUK&JC_fu$EzB-(u2m<4>U6rFc%M3JzD8jVp zcIXP?ybN#NnHIe?ODSOR))WpwHhE4ksi4V~3b;(_l1y6c_x2$E!UTl#WY(4(2QQo?BKhr?9nE zg;o~klzI2<#rv@Y#uDDWVi8RxSTs^*4c}ti)OLz<^8fqxL(ooKJX56&XTpYmn-7 zJ*oBcdc|N}Nr}q0YBj%cV|Bk{NZ*-{Xxfm(`d+id}@>55)11^_vlZ7Flt}EI6?myGhjnm zsr8*KhW)Wpld%n)Hc%-3Nwknobe0@f8IhskXaoHUCWlE%!!G{i-X8$PB|KV`{(wzY zL(m~)e7;#Vl7tScA?pO|Z1k7mXDW zQUNu~HNbQC?q7aQBY3K|H10jB8KO?7E{#Fnfx%{6>XCKBeO9K&2M(+%QbbYw1cESe zse5j|&rM=@)RwfolYR+3vt-5v-13LntlW%|J#}<%u-ZzC8?2s>OUh~`{M-_A@~!@c zZzBb@FpOmm9Xo_r_EjmZ%@oIniI#Hwfo6Uc9!Tg8ySQU~d^p`=^}@-g^%^*|E)B&_ z`y0fk5-khr13}j*lOJA1&~++(6k&^|xrr7IdH4?6VRHe(Cx$UPpOj})?Bq@Q{n^XX z>JT<}u1_%b)5p`HgEiFqNl-GaKYts?IK|Nr@}J91MSWhyb9(UQ+8xd9nzLLjp6^X( zr4gS&rJ&MRsH?*vO>s!It$h~T3g3QXJk((_$KXOf(0tVF+$vA^?6! zoU2MddWS{e2Ma*OO)(6Ha5gdeu1W|Llj;%P!(X%g-XTozU8v%p=QLc6Cs5&TP{TlK zIkmcQl`7k{ySPXdwhdEjtF1Y>Cpj~qH7Hm%g}4tHz)QRoCq^|5(6P!ZQOg`W~UisVp64@9NhN3|cxh zQ3jT;gruxD=Wo&l{mK(Th<~3vZ2h6Sk>!3M zE5ecNH`MNd>lD-YFTO$DJ912AS`@*B{Y3qGAqn*C^o#H^%0Wve#x2dldp5Sx<8>~n z-^bO28Lyn|F4~~|e@V|mB2&2Ee$mX??m}0q=^tS7#3!WtqzM`gMG~~?{ufi<7#?RA zwL7ud*j8iPwms8e;xuZk#z|w_Y;4<3V_S`lwn>{adB5+R>-?KvGtc$xJ$tXUZn+mt z>rY@vaIAI%e>j#{NU@rKJAy*R1&Sy4-`@;kP%rSf_rA`zUwLULx@|wO;)(CB2CRx? zzBiXRd60n+81opLa%*liv=dwQ_xugS$@<;XwX>?DQ$R=$OFiD@+F*Q~RRe2!dL4(d zS-l*lGD@-3BP=6GmZS4$u8%!qb?}n?QR9kqv~8AE85yb}n$=uYvgVwU*ttAB819sC zQiX5p0W$C9W=eL4Tb;8JhO1M=ybu@Gu8(ciX^j0@)m-kg%~Y>$QZGtuK@Tx2w*OU> z9Lb@Sry*&H^{HmMaH{6}*v~4C-SF5uNAqRApk>;RtAJV9Wf&>Hz227wPSzj`^oTOk z)bxl7Ynp+P6fbc9S^YP^+UJq(^xdXE_|Cgi`(B7e-^7djJ@h8UU|!-xeM~R{{{{!V z$k}tnHF%P8ksbHX43|^=r4_Xd+-WI)&4?bd6gP=mu{TFWx_$5^5D^C@s@C@Wu{HYS zL;oNvUys_UJO|OTn{J$rkZ44ZHk(eY4byN*(_QjfKKZl@%vX-7?=iCyn-t0Gdm^kU zbs;8j=scPCA`eKWyQg71ufy-3B& z7_MS7>PjZd1!@%5iwW5QDLC~G(9|ZqJ`;p}{R~C%vedG$L~bWBMi$x`v%S(5$ntjP z#&iYI5N8PSr~t$%c8R9dW(Z8;N}d12(4nZYOwyi=mK9mz&|e8Y%fN=KEiLR%u|SUg z69Opx*kXdqnl@H0GbLQrLD09DPHNtk#B@H!RKuaHln^FiX*8)SzMPXvXM#)Wi0DaV zQ&i!CEJjHa?s6%LjH9+tAX$v$xUA3}8^xX)O$$3rFr-9I(XD43D)f{o(t*eX*Uw z#;)l#tG=2{#}@=Xv!f0fJM{WFII#OrMvipw#f4_D{hkv%r!}2&( zqd1h$oDTj$Q7NM+t$98w&QCojfNSZdF9kHokD7f~=7CCMyf`dGN`%Li6JL?IFU|}S z$A$xy)hA6Fh=sJS%>3uyP7nbB571$Y-@!c^Xxo%h=}-d$%~#DDZMwQbxUtfeCFbVe@vSifafp@M{9}DD9XhC6Pu=KI?=})|2 zNx+ga?tecQw3h7|KkQA(71aQSSf|~z$E_-8Oq4zVlR4cW5)V;>LmpK&P4A_F(HOH; z&Wj098p0nLhrk4kZXTw=lToMsvacmX0Rklj$vPw&$o>Rah(fHAu5xW*I_C5XW6rKA zB$8p{mZTq8G;G0^?9rTw?pKIGVOSipsYRcdisXBX%Q|B3QCyVmfiksPl0{O(rl@~! zXuTCM`O=02#A9Miz0*<4Ru$Z2Jo^?ONI6dheIC;*w2TWG9+Cyo=GB{^;I*)b*jd>j z6+n72Yb;u~EJU9VxpTzd2Ne_)mk@$|7*^d_OPW@A^7kej^Je7ypbFFSQq}%Nes8ZO zn+iQopQV-0CWg$()YWM&%L!*38Mnh4M>E)z&8gFXU(8s{{6V!v>Jb37IoQ}x)zFXN zLEfS%L5j5OOnTQq}kJGnn(E(>urN;#qJSd`R7hX-w=$gO1Ppwc0d*zj;Ep01mwcrSHgKtYBS%vJj@=c1g z7aK+8W7u{YaR@9bbm(xm&Q$H2;>-*>JTmRhSDN-l2TJSs&h8S)2uyV< zvlGUTL}NOyA=Ib}-wp+%3o?8`kZpovxmSH1t{Y)bHJ&ccI*+KQRG^5o{=~CUmfH1C zY%OYbds!5crRDo#Rl{PyboUYWgp;Ka9x@wsgwLCC!%piT#!prf z@lY`%tKnDvh6O0ze3*}uwA_zk--cBopj?<7G-l<%=vIPCM_yAI^W+i+z?dD~xK=8W zNRND)Pb_M~i0k86D?j8PR$>I6?D2rv1q4a|U#<05>9S)p;60H+fhzy-JIGcu0=GDs zc34?S0^(Bn0eWm#=ZaY|d4wW7PLVxr29h)b&M{i71!lfFoG(ei4lJp2VwaRCRo?DVT(o1}M6J63&2IL}M$?0b>uQ+=F8_hhJ!zv!vWUwmr?Bhi5Gx+b4-v1KI|m^LE9{6kK| zh8c?q&fu`FIH@YJ<}Dak2#Gf#+yizJh^-`PtioUuqOuTlD*gm0@`L5l zOuZthc)O<7-!#k5I9LqY@dk^6ZjHak_j!Kt4b=tU7>cu`^?z9`Ot`^uTLOL|@--=A za@I8m$U6rd;#)zrJTt!O)KO~v+w&oI3!jo_crf{_JyZEZyJ z^&`=m1uzsCzYK6fprV!U-`B#2o3ItqqVpg=k9UyVmPm3xLTV;82#=9KhSTg3x~(!a zex6@;x39!i<%ktph;Z{ku-4D zhl(LTfFd7BC|-RMe(hITT(KKblIfTnlTd-FCa!jM=u} zPMIHMe3GMzO)g7Ay{2nXqLDBYRuH74V+K%^0PGUkDW*?!fq;fR{9C(hoE_YeiAY?@ zWL5!CxJIf1$g#7|2B<6TYggzcJ3+`A3H@$3ccwWCKCo;!hLuutH%Es*MW5UMl^3VHX3%~bJSll^k;YOQd@ne%*hTYZ!o_8 z8f0Zv2+9H9(I1hFjva@|q;XtaK zQ!4G^OiEHz{K}5kPx;P}Tt5{iLbcg(bPB=1Sunr8hp~O_%14P-tRa)JK zC@@mXW5=Q%yGS^ZU}(%F_zVvk8WR607O5@NDXetrI#y_*T+(Hv%-b)ps)vY8i{I%2 zb!yM*)5UdLDvcDAE8utT?7AY$yKExp2l>y)GnrU3I=^2f4e(KO`OMB{*NN2nAFQ9q*Qk5_LmMNoc8Y4%~v5oQoYX3P6oL_OHnhuXrC+g~X=a7gD}{eX$^ zInfh?3Qk~BpzUbrZOb6H3D1v1qLxRzWR&KwAdSH4IZIe9s=nMO?KwaTgbPZ+`t=Nr z9J{A|j8RyYD1n0~b+LqFLK6o~5lw=rqkr!%&MzmF%Rjb{Byb#2nXQ>#!RnACPyqUz zYPLe&K(q37L5#;H8I;oY7_LI$o~&zw2sm~sq1sErgrz&};V zF+#LbP{NK=jw*6EiINe1?YT^2Bv20Z2bo;tJaNfS6Mq=rIG@b}2+OJDN)V^gK7<$4KP1ILvT;14Vj3eh8PG6o(p1s}Ntg#j#{yBlOaN*128r}#3JuVxQNSXN;PQS;38J060h3F2ZOyDTrRRL_Fj zZ3bkQA0sl$m~8)Lg#(M`P&-jTGfjYv@JV||@7;8}b+!NfEJj{0RHP&k;#0zeMXvbT zW`>KRJ01^wtkBzPb^J68;>of9mjeB#()fsJrEP>KS%e)|u3#Mr8b}1JU<)|frtZG5 zU^9x@ITLm|NtQE_P+vEFy59K;;b~-GUc?G_&$!hD_x&S>3QU@jq{36vxCAx~k*R-S zM z(3c-R+xL)kdDrJ5D(sl;)Z8dAjJr)(~V^Uq@s zvNkguYj~`4_mteuA4~`L^29nV7zZx%vO25{Rb z`JP#Tlgg?OhJ2GIA`R{;++`SFS3&wDUeb+xfex#4sjIJ+_@Av^c)lCAxJc1IC7UNq zys?}0D)9;h2TrePc<76{wWgh5*lrx4<84GRYY3ym6i1AI6pA|=3ma%ptG0@DDMT7Anc>FFWGSN%cOeJwj z7#oT|0SV$5Yay(oY1exE;d1HV(V4~5U7dhSyj`SzDXU^Z;5KDaefF=2{O?$28Xm%g zFI=udgxr!>2WEzy5vwDJGX6**an1wpKdRYxy{1y8$uW z$gkD}b)*kkXgag1g@mq%{XJEyc32rMo9l&$;{Vo#i(GW9bYE#qhHuGqGANLtD`vf$ z=e{RDH414c86&uuRmQ+k;U=nIiM%`7jz5HgjF(4`<`YAPwImZ|5xcZbkr&zk6c$3o zzt}j*aOeucR`gqwcYKjrG3BIcy`+kU%SYI-|Gi6+36L)4g}$Ttgaino>E3c@Y(l~& zS_an=ykZtjC?{MM7vb=YZ{EE|fjl-?tWtEpRJ2?MA@?GNS48`;E3Des5!+WUQyev+ z!?v(P+P>iDGLcK0E54XcrlI*777em*T8xxvQyAGt2ch^NYk1X{d5sFLDdTET-CB+` ztfX`CrLgf`22%k1WN!RQA#hwDUniGHy%*&sZhBKsQa6md=e$k#wvB-nt>{NqM#x+Q zFR|`7QwRV1EG_Rpr)8b+aK@{&1rom@F5y2p*^VNJ6o%Q*MriGgXCN|~kHAyKjCB}O zLTjYJ{KPU8k#tshXwv7)jpqnC;X>Fjw&Tw9WIB<|1jyYOk2ux)A|e6r?S4Nd+xKT& za4^OQg*ru(Q8TG&0vA<(OOx8KQ>Pmo!wR;o-k3seLALr`;-pFx*++H<={##kaS++P zmskAu75)r{YI}dH5QbrBe(G4NcCf&*#4yLWaNwgkz`c%MNwF)-^{%+=c-hCZ-IP8q zABK1wd{t8weGl@5k%}BIF30ffgSqk?6>^DJhSs1eqV=IOh_#_vX=7A_9rF8&03mze z`9h@ke8tnnpT>$qOtw!|J#5EWIU9oF75*IHt~Jbah_{6d>pU%%i7rL%G^891>|5(} z2z-@iSU-8+Mp_O=PQjN6!s;%8tp5+;%gFkc=O+P<&|S#ukXr2mD{%tXLV-XIBd?M zv#HtTOgNz5o>hmb+{1n%k`!U7V`vjW3I+q$f(o%zMqQR_n)re^eh}~s(=)-ZkES(E zDu4UCS_fh+`$tts;wzhQ1GRCr?awd(3T_3CKI3CPf%;mdt)Ztu2%}qkhi+EqoGW8U zYq+rckWx4$OOvU@%SUJ0n$RcP$P;mJcNl&z>?(#famaAWo$3Yw!1Q zBZB_X8%(e6nSex!yZS|##hFyS{ZV?3F$wX9VFL5|do_nfwe6Z*2+1*zg`XEdZM7@o)2L~u&g`B-{<@3szS5W4D1nMx4T0m4Piy>*~bhZiIYB!BS`9PW}Gjs;hz(IgYk!9}&DXDlp z>IM(#ICE{|gPnEQqK8w>VGuC8!{Cb+_B@@##!=Q8;TFaga;^f0M+Xv`T<8R(@Lx5h z#G}&d8vF2|hU@}(1 zDi`|%KP<9uTSxgQe}9U-d71UzgtB?^tn~GJmVzsvAQ9ffysI!ZsPU*LqnFV#_5EXo z&K+Oh{IjgWykr*2Sy}ot;N6i$oim^1@pHD2#{<)okp@_qqtfojFxU1`_O{Zb8yUb@ zJ+4w6)?N#thX&k^eRSlYX!uVIVZFgFTof5R-roxA7toj=I>3+{loZdDT%j}2qu$O_ z5hqX;4HIAa8>BERrD|YDTVOYjC&f{BUF3--0+wQ#rM0EHPz+89r&R1;tPd@L(HSl} z_%%^YCC^A_A`>@y&fti;Zxo~ka|!UbLh?4=dxlO3V?MN7h>>fBBv1q!YH8pLZl%6V zA?HY)hIb`{4#nhV{f50xh>X1rvs1OdLkLrd59JMZKzCw`s>lWMllV_R`DMqRfDrK&6#gKs#`rZM z5UM__*+IVyUnL6ZYzR4{PKr~!2>mX=pcgf0xY8#Uq1k(8h@ejYa=k@mxp9F_t{5FqB zV*H}-28pKf(LUF)-GE`)k7Datt(@Ghc`XwALv6&?zM|7Y>3BrYz#8KB^7PxS?>%YT zy>~DGYD_%&-~kpERiY^V8h+-4TRe8!Feg=>3a=;1gU4(?V@jsJ5*$uCSV0@Zyjc9V zxQ2r<26W*i$i4xaiA`ILiI(c-6{@c-XJQbcG><+)3N>gl;@ES-K$K|n32o(M3su|I zIBYRNAsOLf_vV%rUAof(9XhiH3j?|b!}A1fToUb0ysqVn&9$O6U?t^K#~4y9*%~JQ zfAf+HD(L$w4}6I%gC|$m^)`R2N=9?rFzTo!dNvM0oxBTHZ5*+F9hr5epNX87C3csc zSH#2>J)zB=13kbCxZ*ZMd!CzVNB;%8q(!XCl?uHj*Jx$Z;qxy=%bk%1C~(mcq-#>D zwCsO|52(m$TUI;$JcV@&qZ)Iiz2Gezbt80-*ylk5Xs6N4-UNi|9I+Ow<~(5yr5_s~ zTlQP&VBU!xEhipH1F{48RZO9$jq~nZEN~$;&7tIvrlu_8P5icC7lio)<=*qguLP3` z3crctsikR}JRT=1@Ym@L?j$FEPig)zY^VwMI0aq4JfVRu#I2mkEx6Z zj2Yh~8g-4E~j$P-8e!BlaQbuz8d z&Zx+h^K#W$Jeww`SA_oSNMINw6$Cv*@YVA zR87MZgxO)BXnFLT`I2icW<&8%)4221_E97LmZ%-^{uv^{82ym=0?B${x8xFYG{s;q8wc^E;7>}7S7NUNpUC&6Xl zMGvCgP7F-&FQDglXFR0XZ7T?GpZ@GnwBkVC(*(6@kptm+yekBXXNy6Rn?YDqNXmx9 z6%VTOqJRm1w-wK#sa z$xyiwt_z9za?E6hSE)L)ivrDOU+mDvusn79rgA8Y10maZ&yb{C#vKoRA4*+hS7smW zCjAwt&gOTj^e(FEbXil=Hbx>b>uxSsp|46NgAKK-kFcL@UgsG)TUQ)}`9frZ+dCm- z{$E|90qfr4x}9*6;ThXE7Z)HBSN2)7Xdi84In7?Z{j(nwUfM)Dlf!6Xg$Trc?@ov# zPfa?%%21`^yw>m1?$Ufk$Kc%D1u|;=QSP3575gR_aD-dy8Jg4BT6(`OasU{0x6`Fr zZhrsAAw3gI=J{xB6l!cgG^h|8M~Oda!wKyJ<+{e&kVpTOK%QU4(K|^+H11*n7My}T z^7&uy+eX)<8qWnt~x z^sSeQx175;i`gYb@a`Y~HwY!Q)8HebX=3;A`Z|G{NNJ)`L z%B^&tlNevr2mvA%7haLwE=n3}dOSn3TiW*VoLC{mjHl&NTPgNBD1l1xsH-S|A-0t9 z8A`>QI$dcbpBpBsBH8^%6S)QRoc+(_Yz5q#q>HHSm^taLh1-Gs^@sXOy`H-iy0@!z zw$!^Oj13&QI<<$aZnsc!1ze5{!|vcY#)d$I^Z6EQj9nn% z9@h`yg~~)v7Nbx9?$UhDu~9Ot@$v>@LaQtsILBM1t>G@@g9eWE1*XzVifXFwJOVCe ziG$x`oKX=Bj8uRkkrnhoWdo}O1)@D#nEZ^P!3_>LQxZBAXsl$W4Aa?51z}jSdu3j7MfNxb zM~Qq%6P~xR>CT;`88X2J2{CK!@WPYPQxHtcb;RcP>X|_yw@A&xSKTAXk?#7f-V{XC z*pawbP1r%Oi@jPt?@Wr8<{`*^g0xl+dU$%ub&M1EI%t+}-FEOa(9 z5ygDy+ClR*!4lH+Lb?@8Sa)CY+lb|1gI^?pQqcjSwXN0K%ytqM%Kay3`orVUdTVv@ zOK^3F3nL_^+L*2l3$xP-5W|P z&ZZ%Y>3fD2YRTs{M#e0*Skk9$h!}?wP}_5M62Y4Jt6d#N1eQ(3$JF^2AC>TQibzKA zm$E0;qa~7H6F3-TFf7%nK(5s&JBe<>`;vc|-Bv#4pVsV(yg8O1) z3j=c*pn2GpPBQHae=@S7RHw$TO*e$(L?{3zVcE@T$oP|XlYR)@QHp@ageK(_$EB5b z1t)2VDZNFE@Mwo%S+~C!jC#KA$;2Jx`WrdD>^8D@Tos4V(ZvQ~SpQmXcXW+!-!uDm zE#iJl@;kV7@80JQ5@3|$o9^a<8YznpT|r9Sv%XYA$K}=^4{uVOF_dktHT^H5lZkH( z0#)AwJhwhN2S&$?3T~*1`vz}2apNh~yBU7upL@gRyY=7)v|qe7mp7e&B^ zB%~R+N;___$Mfux8p7n3TBXRf*E!8Z(H7R+rq3A>ejAk@2+N&GC}s`lz&k=E#!HBL z;1ZTId7q`$9juwY6k1IPK+B0iS5v1XDYC)X(~&Dn=(6#Ac7YuQ#n^!p&sl4V4?3MEa**iY zzXv4g4ki+RfX@&*{Ldie6^W8Lf{cP&-)*)gKgPZo-VB=>_T`XCZZUN=QtxrwEQNNfPTnSHzs5P zh7p|y0A+PIO6T_mu22VAe{Ou35Z6ofW-1C>)e-kMlUShM=cvA150Jq<56Lg6hX={I zt9PM9LcDh`2ce5p?%Z`sMHJc1Qc?d@5tRYB4GF}H1Qc2rtmIc`Ub!NyAWQj6I0kQQOA{=y9S+HygE z+v5y(Fka!YF7o&&K%Qm#+lI^Jt(_Xz{a6$Zymm&8^BbevA>!AYf`iw4Uijm$Rv%i` z@us_D8CL$5G{7Q?IAHCPEE#O;fLp(2&fu~6fWF~cV9%2`9i7KkkCjK5{yuGcf5Z>$ zklyGSmgHgWNC4W07Fr4Awq&8(?8ts}J2n8JBDeO>h<(<`mdIIjTG3MpNL%zEn(H;Y zA2c_Tc)0z&wZ3U@^K>~)S#8QbAn&1={1t2wPo<~Tj=pXD6Wnhs>%W~KjtHG%ayx-u zg`$5r@_1$-Fsl;*HT$(chB>>H#(y=sDD-`#?Q^}5hQK6Y=O-T}Vw0EgEt!FyIZ%B! zY&uDtZ}aZfxbb1b-bJ$VwZ?|8_vP9$#Jx{HE?H{`-sn@OarZXZ5@VJ6?}gN z{EE{0E_Qaj{#-z4*;#2S9Tt@nwENcDZ=FG^jO7hRTqQ{psXI1s<-T`*U)TTHaoN^b z=^nlG>k&Wx&0wucVOiG-p&AO3{(`K39VN7Tv?y(fjCteAhzrc5ZQABS1X0&#zoMp? zrog5RdykF@fAwgV4m?7X4=3f#^c~`1u1E|5%tT-ePn;&T?ZmI2?*mz*)v1@x@!LxklGl@Woa78y-Xc$h_hbBx0WY| zB9%qJX6cvlmaKjb5LtOtybl+9HO*W9v-wkWvPk<06i?jUqz|H^ee1T1q-)ETxPVubx9q z-(6v*RMBd~qPgXu!x_5Si|})Ao9aGPO$u&2n?n&3tt5WKCuCu1>zC5-9NjFyUvrW( zfrG_Ys?+~Uy?R4{T-)3^`v{lcg6zeot&T${6HA0bJdHM>pufUkMr=pC4=Vi(rZ?~m z1f^ZVnk~;qwFV5f+9U*agk9g{;OQhJaCx8Nl9cs6+4Wj>?z<{5BAAk3@ijaN_-QaI z%K-r_?hj}a0K8#I+XzzgXwL-8KI^dWZLTcGefh$I0k9vO;N@|Nyi8$bvOW{m^{3II zU<{q7VHsT#rC*{jN3nDM(8#T;sN7}*Mn99sy*m-FK6`7UF#hab@?GiNUP^V_>G7k_ zaZTKeN(5o+HlKf{f|lo%q=_q{qjMP-n*V2>l>UqI)}!Y>=Uqa=_VxFqMH;H~f=)&# zO`UqM)dtgDz^uqlL$A2Asj`HN5k5J-d7Hd_mumqzOsmSLET?5G`cx1+SXpmP@DG}2 zw3jqd1DZS4ffAS))}c%(S01|%oCc{vj`{S~XlgwtxYv(B>u_a+E{XpbC|W-tCx8_M z*c|q!P(*o?=uPJrIV+TqpaeqV)j=u^C!yp50i08SvryD)SJ!%J?LjLplm$4#I*IZc z2Ol0P`v3zK?AUwBg{ZLAU4?Njm)JfIqW?vDqccUbMC43(pi5nZrh1~+pA-mr5BP$! z9;PpFEd!9B%f9|8^XNJt9oGL>!=@95J{D@TCQflH`2ZLIoI~UO4pKg4SeFJ)mAU*UVU_(A@&q*PA&jD zk=-1>n_XWo*^$e6iW?D2KcVlr(ayRU-L!Dmri2>l?|=}k2?Oo~xRW5|a>NZAEn8~} zI2E>GtC7h zZ#c%t>k620SNVsxt4abkvnCA&1 zAaz5_>?gqNJeUpAT)lVr=9!J2fqI{`_>M+{rb)ZQq3c3!!~W^EcBCK_4^~ieEgd{C z0F(=5^ucgnxsYb{i5a_(zq9%7q5KDiSLk(5fRAo#3AQ2V{;y!MU%~Gsl}Gy|(Ac%d zjODjiT0HYkxzGuZ0aKcM=UqzeMK9jVSJfIU`HIUCg;F?!ey~6cJJhc5M1Yi-rZlGj zKyD#_{N_xr_fLTmC^`NnTai?*=|*EfXek;;ADw!3$ysAPz#J>yEJ!a+A1z2{H(>eG z*)Y$p&IB3(ZM4^7Ov1#mTZLY8U029vzC}Z2Iq1pO((zekT|xyrjkovjMz7yH`vtGH zEtD~K@;o{_*~v5EnD2>Cd6V9lB-3AONSFW*k7Klyldoa=wR9ieJ4fe=n^Du!pWKDK zE5at*{^&>iNFo+_z`|4nzVkXIf)chezWws9DYy9jP4RaBOeaYQRy+;hI~|P-43e55 z-KkrL^?1AhWGoyzj5yV~WA40YnZ+m5x2~$93DE$~8a$OFyOUl9s_ZXM9DAqJ3-2Mt zc}ja^N1a-p%XUn@atpqT8>M zqFv#eK#bQxvDd-BFV7#WS2DSove6S0g^>TR7eLJT5Ak0e5kbz2=ZFeLIzhl^8GPW7 zq|%29`Ya)ocT{`RxX!)*%Gxjb88VpUhlvi0os<)je@uF*H(6d?q6d zOK6TgN{qu8M<}aQv0eLtczi(i%)jA=-msiG{>5Yi)yK_ij2C58HUw>)PqRf_bxZbl zRPV2-jjb#9UE_Eq|5Y0pX2mhxlmX2B=RC+;@H>?;eihiLh+tp`sibLnj*gK$+OkW_ zi06--Q+?yeEHTHXI3S7y5F2((L=JYolq`zdwL*r6o8C@}_%b#ZuEC~qg^=_)C&~%z zdW&M};_fSfpV{}VUte1D{w$W+(}Sb1d{+m_j?u>;1h-?!B$+!oPAib{u1s%>jn@~f zr1XL{+sx-ptM={38mp*%Qk#|hYjXk4xX&kA(rNHvLp>^P8?%L%&C{3#zD4BP73RuesuES;wnkg24vDeW1_y0SU3TF=mJ85GP!uIif0v zIh`^2K5%w>%goyZuO6PB_%Sqo?9bJbt1uK9KQXt(km)(>Y}_}{OO3&EauP@ z+sui%zG$(Ku+VilsM~ptcC_y%wh3dOwoe~3{fJP28(Oo;VujX&7Xe$zY>s(6TMdGl zcVh9g*G_+q?Q{MBC4{VobL|u^q(9~L-Hx|4mfJ7u*-av|pY$&|7WgbxaP(ONUfaRP z@Qr^l(utE!u?Se%arGOu9uGm)MngqS;0Wi%%7@Ho2qZy+7w86~0RWTreotO$pK)Un zE7FC67cTc8O++^H3g4OtRyaK~r2wbnee@+SL)^^5M~?lTSq-n-yt=ZMixkBgIT9-p zlOpBWO(PYXhf$$zA_bC7D)KW7c!5>tz!mV|#)S^Yf-r@+Nh9@*w-9$AQC+OUXxIx%wdv<@=xt zGEP7L6VE2VE8sv|7k>onBJHcbUX$nl=QQ3rB+b41$%X1S1k#jYRXGbyE_w1XRQP%~WL^T9zv3MSKo6~OLtEmlHV^LBQEkZqc^Mo{o8f1x`GG7H=C zRxz0cvs;c6pv299{+{W*ZSDNc`B`((1);B$Pd!_!rGS-`HkwVJn_rUh9uy(Uu(mzPDzkSxO_+~ zi!G!apdM>$`YgB#z->}Utx*Ird{8P{C~!!$!zH|K+q5Qu*867qD+@U}7%&;oiUjUG zqp*%#C$!2oB}-*0S*XBFcG$`GJ}XZ&Xu^snv`qQMU^3Lt&=^99_F1TZ z@n?b^&u>exu(WX$!3HTIdvtp+x(7~kZNnnukcO8YtF3s9lzw~4A^QeP{ipI#ln$7a z#`5cQ^e0LrcWKn}dHeM9NA&2Lsn&BTUdcBM@n_RRTq3zq)_bM~leGm=>Gza4Z4u+T@(h}-P9nj58f}S|9 zW3RMlBu>L8NeWzn+T-I|PwA!|XDZ%FcjL;;W_^frG#ua(34liA6W%i|M#f5g)nHP7 z6g8F;Lg|H}MM3?mL{T8Z0n@Dn$Dn`mLoAM=1&IYTT6aoxP@SUh1BVVKOn*x$=Wj;_ z0op`r>A_tnd)^do*kC^xj*^ie!3^~hYgZE@HqSZt^(E6tifZMkKAWMjk*X*_uE)4B zPag&eo}FvyJ`pXhr8T8HqMFMHyPJ?`jcpzWq=`0xJ+|ugaZ^Oopv)C{s>&_@TsG24gnL ze++>w7yVTNsI_{aXJRN{O7PV_3+Egr*4e3QuIP!>MHoh_W*O>Ds(kh~F>I#7%Rb1Fi<( z;jsf_g9$+9*otyaL>T(*&F#P?g0OzTe22UBP#Y`%c?gMY7ot23yl3fID0=N3`X}_F z)XOMa=~;Y*Jvgo(S%{w?r_VcIZQZKFNmcNah7HBj>Cg%e^*D$<^CN(uekdIT3rK8c zuN;L?{j|@iHaRj5BAqhD9oL`dd%m{2ZhQTFd`lfkFZI*8+a9Xp{n}dSrqI^*>EBNw z=rkl~{DhAKRd8BnQQwv$lm+FXM-U^1$q12>9*X-9A^}BxrFIZx$l6U9$z|kuGX&m=F4dc4+&xp)18tMV+V;2>u)sk@CLzI__W_W0n48XS`EpbK2(Wa>g>~6|Oz(7ssMiXZYvZe}{3qnBg zK-@#zi%9Bny|8Ek$FdrSpbEfXSdYomv8Q_+`?3kLBu?Lu1e)_f3xSBpI%{T9^ANGv z4L-sg+XcX@V{`X>^iTN=Tk0m)JguR3F9d*3+n)FV9)fo+K}0hPEf|F}XDMk4kr)Yx z+{BcRsz{&;fA>`2gf$b`NW?K`0*Aabf3e07hwG44{z=HifApX{ylKEu;Os12J?5GeZFz-6 zW!*r`nFNDH9-BA1-u_-+Z}|r^*q|YJCsP?r;BJ)6kten~I`PJHd7ip3A)FsI5pC&{GttHAof7&Wh9mf}|bZ@^^{W9UtYG+ZUwcrv;vEsSje0ufGFNzf= z(0?9(+_5VfYkNKynVhqaUjgHk=?e$;iS4kN>IpPN2l0(OjR3^{DSfL+OrbUnUzjI9 z!lrDc{vs2cAsyQ^CuE7?J=)wmBop=PJ5ebjw`zN@R*4cP#?fN&Sx>rpouV#3&vGe0 zU?rNAIsznBph5~P3V5a61p0DfIZr)7R1K3HY56sL|!!=%~E>e(_dbEIwN+p7+825*9cf!}vtytWcdQF9lAj#68@O2}zjGQYvsA{e~8&5O`!Dh|h_E`I=U zS7{db8H_rHeU?YGile6V{%g1JZmh`rW7tZw_L|`3a~CO31qX8yEejXS#qS_ed+705YhJ?1pKV30s)WlQc4B9M*YE zb7tm?I251qh`!Ny2ps&pHs}0RvldZ*^p|ex*BF>%9!^I30T9bSNY2ro63}cuAAtU= zv3V4}sC|T|X3r|GF!?6B$um*u-^VHpEsL89TgsBTx)#%2Mxv*Cvy@-oO$nTChw)zo z3i}JWD`PXCWMr6~GYR+n*rww~KK-81^M|+(@@3}!5q;gmcus6Nl%uF1hX^zY8&ATS{TL>*(7~kk1u0BsZhhzpxL@!8vP@Y#8sBBulk?`NRT!) z>)&Z+!k15Q9>ZYOLlx zpHwxobb#ZWTBco=qYMbnyN$NeA2vHQzxY2qon=5((blzT>5}da>F(xm=R7>ynul*>T%~~*$NiA`ZMEKHdGd8b&6bo7WjkeGw{q`KU zpm;eej?oraj1iZqj}HpvATD%dlzHd0FDLM4`Tw5sFFP$Hty0biFV-5#AoxB2bOl0~ z_zCId(B!ogZJYu@k@sKx#B52pcghCZTfu3Zlb)4=;QKT_{nnd)7$5P`6+IxqvA(;p zjo4~*`?~>O+o_DwSU&w)_^nfaoDB`V$}W)tZ)5#p3wc4YL5VuK%6t05KkyuH49Ut* z$_GXqe)%k# zBM5rc3{=J{1GB|`UkKg#{o-TGQ7FzoYo)Y~ZKm%GpJJcQzHa`lk|!(K=c$nr^1^@W zBIqN`qjs*8hP=ju&DFDj;fOCEj6C!HQehOKc=vC?Lj^!=@|Ilrq3+BYA+p8=MLW%9 zT*u=g+TOnAT9dzC;OTE+{gv($Ii~|uM zsCf1B9h6FyRf>#iKW(4EZ22z-^fuD#pQ3G^AbR4#&o^nQ-#`5Ua9?`kt?AmHO|CfZ z<&ZlldBU=U=d1cOnDzI6jYVDLBe-RHfRF5NJ;?8Uv>U-^&>HWosJw;n{_Q7za*Ab`_WR_2v7$ZAk4!M4I}^uKt=oOb;>1Cd@2ekd=woZ< z5yNxVcvJ+~L*xqtghWbSMncv*MkvX~1Uzu&IQojH=GZ?g66svSLKPMmu*IqE5?KgU z)i;)Nr$B$Mqcbl1>?5!O^E!-PdQG}=4i&Vag%R%q{{p|fBu1(yv*rrKXoUUYtv6Rt ztk`=ZBG_G*X8`m4E}2q*>Z)xf@#^!x=Q@4^9Bx$h`VD@609e%g7qg=xIh~cCsWVm#UK$)iRElbk-NG`uK^CWP#vFA2tPdO!|un zVT`VyVd40CkB18;pCFXykMG+_;fGQ#cQU*Tpyk)DTnjFCh3{Gdz+rD|Iwe1N1RDW% zV#KtCd(iH|s=;)43@bSX{N0h6s*c5W3MyqvgQzQZ@1KUOq}T9P-pVCC%&D#VguM&4 znb0(ZZ}_?+-|Wgkor;CiQ2!c`{y3w8ATt~0)(ut~Dkq2)S})|OGSf|%Hx_d4I;A_Q z)i5PR)ed-V@?}Tu@AqzJgjDWr&nw*44R_u3L#h258N;=&dDpbnkXqZmF>CW^NfQ{- z&4h^to0-aPvCM(?f+na5Tf^eUw)p*Nrc z-}59aoI(w26RWsu8_5p$_&;ZyA>~yTEBLCxx5emNicJu-e3c&Jb`2Hcx1!%Q8!=O+ z+UFBCGb-cM$k(H?yXp9GP-dKY<@N>NgC_!ufcirP7vIK68w0j-;x-nyY;F6@mBs-* zOa|Hqy3Ug%F8je51dJS`hU{As4z7!cuJ%RC9+d7R0hShv125&aldY;sZd=}DI*uj{ zB2#U~csMJ?444^+YjKoR;^MCF={!iEmh&eN#apo6<50p=U=Tq5sHOIN`C+Gu#9~nP z_6z>~B;S*HnlpINJqkq+NmvnC!-G&M+k&l|eqz2p0Tm6*t&CQO+J_RcYTG}}U#=IS zE2;pIj8bHd277f|EL1RH><-yR8kOYl0+Uwhw@E7phhkGAy$@UpuzGNCva3>}>MZnY zKjeR~x;}LT7%0V z3iFG| zt!6geBF93zWXR@N;*fEjp!sbltlYYL^BN59c|~eZmqe?Z#FLQvet%dcAx02O-J3v& zc*XXuXn0}5>qH+eNduY3FN6XU*)VCsSl`L1-52iLSSMl=bT`TM(u_gRD<3WabY_Rl zQD2D3fk+m+NricQnlauHXG{$O%IS&z#F-%tU^pEtZMF-T~v6Azc zQF42O5zHlfHKlQAM|`ZE!4N}L!r9yrTJQ%Z`!LV%m2wtlyUe(a73N~Df#karB_Uk6 z{>im=Bn3K5ANMYCgeY*^FrC5X0>bYIQO6IlIWrYl3zjG{SKFYHR7tYp+#;d;DwDvW z47-R7jdjkKkR<|MdjoaBd|>$h{5Hk*{TyS30`!R^0{1its}J};1bj&D&;`}IX=S#^ zZ_`@V!_>P`_L-3uU}ca52K$KjYf_WrOGG4Ui5i4zXv)Q@bk|s@&*N!c74`d=Lqp~_ z2#gZq8_1c4%x%kZFtv=A(IShBY9Rd?(8aq5{^k&WfuDY1A93Xf=(Tx7Gt; zC3m7l@_dL+;ir0gz)%aT^Hb2TYUqEnE#*DR2H>b(zZFOuBvzTOZEc_}A_m2R?Hc7* zSp5(oPfD9U^n0O^3V${=o5EG4WcP-Slt$rFLq+uw!IsD*NhC(jj z9&%fIHu~1@Xn2Y7Y!D0?s2~V+TSg^ru2-GHyjMgGUsU)Eee0P{>r=TQLu4usQIiZPSVP8BbOG!Va zi>28zQMRet(%7wi`UKa_%3jCO!)a?1ac?Wd7zDY1(vMN@bi@rIT-$`zkaiiJVwHyc zsyt80p(IT3j?Mr|poZUXBAVZ?1k-pkSZy>gxFu2VQx@6JayXH+^X~*}vx>fYs9m~g zflYH{gsb+$;D12F8+1jH`z%~*yPZ#P%J1UN)JUQ_S5a0AE}V12u-mZ6v)Az&)?#sC zOGv6`trZfX`GpcDIpf&|TE|72=UI98T>Io6;&CwC>ksL+Nh2}-f^T|^Z@SzaLjekT>?gGP%W5F6>c0r%qccDG`(V*%kjT4yOjQ_;& z+kq{*Gw8iH5N*2a?0xR^pOOh7&Vr8cAP72)hU`%d39bos7*oOXH>)~=zn0W3%5Y^o zhZK*(J6ZfKnX9MK-%)Qhr5?gw^{RxwUR27zzTtu-McTH}qw)Wd9pMgAoi-`2_MRk_ zHIhVT;J}2HHt@TS3PTC^+-dY_Vhi36le`m+?9H%G*d;dHs8rUcEd8doR%6zfZo!&? z3tb{$^^fM$*_q@xwnLH&7d9KrLPeU|%Z6e-UwS2Rk5l-0ff?U>)%40SRI1oe*{~RB zNp^cFgdaBPNE1eD0ytP}&TRpaN;=|c#Y@3LXO%gu`JigX zl)m9#eQ~tK0^&NhA0!vdJP{$cvK*Xb(arPd(b6le_)^=mN1Mrg$1b&$AfvgB1x<)X zLvTZpt=d4Z#-gOl8k3F2ZIOd(oG9U>I*0Ll7tJxT(oQa`?gSix*k9U@dVsCN92ZM5 zo*k*Pjv*9AEveSSv7P-CkP74aWN6?MayNoBx8;uz zOMbKK`dBl3lFKt=^|59L`ptogh64*-G)sAT60BGsGt_vqI%Lb>p?c-JhB|+`Wq@m{ ziS;@gGqE zyFh$9Wj2~PImI_2PmI3cHDvjt3HHo3-y$2$b;{-7A>7>aVUfM26V_5 z9w?(Um>s$J2a2Tr8Z+0z^CS5B)lB^tw0sONWOL|UHi1ymImU@$ezxJEIvLGqf{8Cj zupZo+1QIa_vamHxw+uUk;?egc$$re!|59(13W^In)j`pae=C!#VZuVV#=fh`Ew6kT zZIKNe%S3FZG$mN`mM6K!Zu6j?n&&q`#OmZr(z@rqANuIe)z7r2ds#~;>ub>ZaQnZ! ziDN{M_fbAduWLRN)jBOWdWA8+m$g2-CMth6bYu|TfyQXa+?^3z29vXO53ABt;)yFk zH7{)2=n;|W;isTL$Ffi8_JeCAc-@yk5JH|)oUbFAP#{=BH4nBeqQ0K7g&mL6pwTr` z66$8zUP8K{-pT~?*x`isiENbbB??R|_{Tw7X>zE>Ug+{rtM~#_n9FizvV}8xlt$mY zN&oH+G$or7m0C1nslT*Dntv(&%@Z}7Ch*VG(@M6C z!7kAhOsCyAW_cbpi8_XWlBXz(uphNLuiJN$v?iB~onT=rg%1^rK#(Mmt^T0_RXE5C zYw!S{Gv*+ErA6Rrm_3;NO+3hf^Cl!iMO(2yizE(0&T!FkIPc_vl=7hlL*vtl1v`jr zMd4+-lory2hv2DPX9@DHr&y%M!=X@^L9Z`KVr+XPnZ|o#7%!Ut)(`3T#j;B2pIal zL|ero*PlR8{4N;^kFc+=)v+%9h5Ia&L;Y(6F{BYI!x(HDFA|$Rk(fwQqF>bzj_+aj zHVhN#PT_)Yh&}|aO`Yiy713GE@6TfXPm}lp?;&ck-F7=HVyVE4^|gGRwNc8$##!=EPb%qj(>re*mfe;xrb)`ttBA5K?~OUhhi^2!e42{cfsS zy!5aBoYu2JkgTkvMwBlE1;9l4w!3EPv6E0fMvtM~LZrk#%rZlj@7x%kBGf~mm7(X; zm&Xl7vJf2945gRmL;Dg#3PVXYa?fB==n)Oe8P<1fl@b>(`HDJ(p9TU>w-gb@)?G5u zxQY0+=XsJmo_$ahbTov>B_j?8EEGaDamW+D_Y_H$Oi9RtsMXkwUa5Yt_cyg~V-c3i zG&wIf7ZUrVDTW6+x zFKTGeNZwO-^v);J@?-I!=NBe9M5HhLzH>1U$uy@Ub>osO@uv4#!M!D(0Qse1^gqSK zBh>T}u%}=KA+JS%?vWPMQPtsLbiz`LkW;ajvC(9t@W|8$)lV-gmXuo5juFd)jd|)X zLo5wBmuIHz)DY>od^qm~-KQR7>=TDjht7%QVXD`?korSI<%R@py-1UhRAbc6K&N-s zti8V+SHZll4a+jl#VxG-2r&Mp(5`qS3X*qwu`+&C;)@mi)GaPB5xB7~u&9CbnC$Bb z{pFJxym7h^*3^-V#;gawyNls}oxf&tGIVsRJ!FoaO}}!bgW9ZNk)s12TWMcCQ#89K zwCeUM>;EOJGbo!yu-G2=G8?uoN)l*kSef+22p+~-q0KDtDHH$} zc`$1-ZQZ!WUFeMWxaML9L$1I(>YbjQ zzoocmgeMlMW{c(QSX2k1Rj5VVWOp7h8DPvQD2UnM4(!k{ z@sylmogd}8zY7Y6z4PH{`wJ$$>V^X(Yd~8lRfzfq1AXobNFx| zPHGNZ87YNZOC<}ripc#E@%DAtUVVCKvX48KmW}|1xj3=OD%~V4FX+LZ_aRx5BAgaW z_wH~m{E%4?iR@s`@;rl@qJnBy@Zg6P*<#tVZBq=Q%IG#)*yS4u75u(BAZ7N6a0 z5)voGGDat^#65QsY%dutq9*l9)n`{uz2PO-cit93&dmBXM#9EIEiTe?&i?&adfZat z)a!*_Y5{(!F`D@Ehv_(W?j(&n-tQuX>AGj!RT|LSiFHPp(JXaVEFMG>qcC~RGGau7 zyJrk;I}6#KkP%mB@P?}c)6_f178qG;A4PO7sF2kYcvWycE4VyXjbLc~2xsc(B>5;J zT78V&?<+zeVIp}-HPj|&0Sz4AEVtNPhI+XiH6<)mN5u-rJ#XM3U{uTf4G&}-*Ehp@ z-+eJ|54i33MPVjXELD^izBg-unIP^GEk_l3NMjW#sUmjPy@yg^6ccBED?Ai9?A2sZ z1mp^oZ|vlti&0Qek>JqXFM>QZsSPkslu#fzSG4Z?I4F%VsAVY65#rc_MN_DN?SjaR zS!##uo$id445tAqv8Wp9ep`P4GXoaI2nu^otvS@zFMO)U5q(|3(0M!JKvdLm)EGE) zn8#t6fT49g9eMomIJ&7N#Uf1-8Ys3Nphw;F=8Bxq$QuxPlqB+NvoGcztz3+%&A-_A zD`=T|CH%a@<;#0q_zC>u zlHRZk7PpfzI!1Z;q=noZIM#tiEa}ow2E~mn%K0%jN>12wZ6u%@7*W0FmqTUT^ks-` z`#$!={Vw@2DX9M&N+quxZvsvs4pqy1wZQa|gewh{k9T@DpTE>%{)DbU{3yCh@NmzKE6>FBr~r$OpK{ zSfT%c^14Gq^N{Q#8j(T9p3}lqAS?V-#Rugt&lO3+;+w`jQw}IMyx`C|(=R`rf4n_H z1WX93HaKLo;7@;m!H!Zdr>#TbuYS#n;ah`H|3DfR*#~Xf@$iB?99!owg4Jft`BfmAgod&qJsvTcgtRssLR3zjWj%L zIi1T=ry@Tpl&0Ql!fxnaN0cAuxW;`J(qyZ*+}hjtIj2-gJF0LG_cqS(=!JMY-kzRW z>#kNB08r}<(>%Xl5+k%|1yt0~6k-$WG&K_wg>F{*%YZ*0GLdBu5E)+5K$j5d~PtQQOSU8>7 z2x!%b_fU1?WeQNEnqSVdR=dg!qS*KKskyd8&TUsXe{|IK??!F}>{CC*>E}&gvcW47 zA<5aWW|KTzZh{#{U>f~P89oUR`O9|?eg=RrD~vYMO1wshCDoM%&R!iTALZ6S<`}Hh z))AR3Fx#()NM!C&Yzboo+9`d8oq$QW14Tnahrv2sn9c%fT!xY@Y%KUphXk7L#D}nbb^~$E3QMv$0ap5k{!+2U~Y#H2$Y;^IYG%iIBs|(^s-lx#& zXpj?UDr#@wTyzcV(BNxV#+wB;74c=_&1!Exys|zc07LHrr%CbCQ_~$PRo7(lI;z;c zcLTvBRHF!UEyo&)_6Z2kTYW>K2vrl*TDkpw92I*oB4Pei8;Cv0nmSnU?GZumOvlNMLjT0`Kwz zgAvBQwWmjbGevGbRHI}*(xj9gOUOrfp@w_=VYxcsk|!qUJz*Ss8v|0kcx{sggIi?M z)y?)Ux@Vtp$k*qMb(PP5G_n^;`p;RGbgi+tsjaQafG}!<-$Ee9883g_Q?b{5v4)z( z=)ZtW^i_)NSHhyQP;9PQ+~g-x|3^e?Q=}&+|8oQX9B;14I3KPrFa5)eDiGGLYe@yg z$g1^tD+jp=*}@4dExwl>2~ zVs(u+fQ1QGXP}A1e8SL*ljgwq8>CwI&0zai+8c7>5?a>k%ik|X;XibuT?2OKa7{+2 zi&s4k-(ve<$sfgAK@2=|_l~zG9YP6s2E@YhgyWk32EB*-6fje4Nm9b^HlWNS{j-Uv z2Zz1tD)Hg6JOE{Oe#Et4cLZ8R5~b1KaF#DrNg)@qr#~Yyhm_CE_zK70El+PTSaFdV z*UQ4FA{T2YDQS|>7Bci-$qM--p)iko^s)y#1C|_qP8+o>{u7=*x+UM=APdjHV6NpP z^VkE`=Y6rsMbigQM2QQWghNg+vRujL_Cre)Y0$U*RE(rvLRgo|MDaMn&5OUspI?t1 z%;p*3hO!)%h^+bu8kxcQH(oq^F)1<60O^U7Fkq_mo-CKr81)Es7K?uf@xmTiRN9j3 z`M)@uz{ZlPiU@g9eoEWCX^bu)v&FMHsk%+FNUaakpiTFxD*d7Fx{9=fwqPN8#LTEU zI$)WXXO>efw=CKMG^79h0|~@?H$38EzJ`W5GG#$pt#DPusv7y|<&+i%*`gHmHV0A~`OK!|c=l90o^Cmv7Q?;4^#KW*oRPMy#nEQH!~3YNj9f>c(dOhd_MR6#4FZ@wX=kdSE^7@3q$#n>PfIQ z%Cvwv!*1>vO@ho(=3$9o7n7`LL+@O>CoHU!TT$=0lB1dOdE5xABGvkw zMwWPb2GcN(Em>fBwDDO69B*7prHn~O)fh5UELs{X!Gfa3RzH?_*Y%cDV>2CiBC_E* z-jVgZ6eP;f_(qISd3|Hl$>9^RaU~`7OP1TaT12+PM;kyLvFgCFxv41;Wq6clF;vQ1 zK@Ck23{D_8sGkC~w}y0T=W(O8Zu`BrGEo4A3PM{*RGG3fp>BF3F7s70e`yT4H>$Lu z(pY(yr3Lsd9n&(ouWp$b$y>2KK8Yp>P1BfXqBimjz5ac6@IN0$QN6q{`uk1YsabDY zn_tLbu#Cp{gIQ!{^)i`oGIu;Pq$ZhJW))Z27m9aLl8z_^N?{X#8)cAmF-P)2Ex7#1 zlES?bz#d(1P<}0GN~cQti${WL%|k7kXRu*aLyUipaVKekA)F+kTM|wYkp$+}=9S5` z_%uh0V@CVhI1tgv2Q7sq5Qjs^I4gGzGnNw*rq<&ByrJ zQ+DJdJJ+u`6k4dupLCCiY2joGUSD#Bv{=G|Bbj=o*Lo(CwGw~h+yDxqCLxV);1)eL z&RR8gTR#q4Rz(5*5sy_UUjfgNBCY9(&GF|8rlbWj7zCRzRsBiu~_bi6v#!5>>3(y`dV zI$qabQpv^Ni-@nYTk0fh#tp-0U;VVYRok!BW|7qv#vyoI60jXC!Na@PhMa+%yINot z#1DGrGu=FxG~Yc~U9K==;n?h(UCh;^CrHr9f4krC{br;xT=L>R#u4!ag3HwbWr+_8 zG+kz2AR72ULKx_PDC^2h34(>t29ZvV1zC*XJ{OJb6zUYXB!{5N5{(n>fhPKy499A^ zwQ0lk(xc9=&ky=Ol2M?<5B>C1?+lD?V0Fy%*s#e<^I9-t)a!G#V024%p-bctEtg!EoIINkoD7tTr~NG^&)vOxdyQTa_JWzQxsR==gk$|5(oiAgf?)6)}S z%1|gOP)?V{j*$yCTBgUU{2NKWe+4g2aF7r zM+&dMaPWP~nOCdw!c=PTs>$&BX?m@fmP5p?5J*^itptEtPap1xFF#(t;v)#m2(w&8 z=c&wmJ$rIuXqBDh3gYM2*-=`^BzK1b@k=n$1Iq)w~KGl%R5oK5)zDuaK(-8 zeAsX@#}c?11oo9&aM)x!_puVzk8KcUsnA_=pZ~dxB%eG1#mv+6>h9CEYbP3+ z)M@KfL{C z3e@Rih#OHAHjEuKHb#%Y1N%?BDc?QY3DFaSzKz`?e>1>Ed#^&=oZIoelqJv2T>zy_!{*dY%_>DG zRcvt^Ywm|TRq|%MXrJe@$+jzGWB0!7T_?+2w@LDXT>weD>zt(b+?Bp{ zjL84Yh2yNa-=-2r^jujYI6im~^5(wv72XuQUDdM)>6Ec_a6jZ#{l`|nr@0F$j$bR) z|D;ius6QPT?e0EPVN^RmOZms})7&bg?O=SiYuN7dm-z67Vd38fz7LxtuLCG-ciAtq zfaxyTEb}0s;!44FLg}-bI(fTEieJ+ua8zvOxbMMc8>)JIcbqObte1UkCd-)FQ8jQr zV!4n(C;*12NcVAbv;smS4Y&(Qe_t61PBtdQ}3ioJ==Hm9}Z+Gg3lK4 zE)=Z*c|%WK2es-mX2<+w#bQo`lQw^&QIgVS1agP77v4tj50*&H z>G31cP~7-ze94jzx2s8$Fs1G?l(=VE*1v1=7|a|l{Z5A{s-`>9Q$pd5eP#a%M+$zR zGZ>Vw?C$u-Y8RTQUM=pl`R9an`{!KiHfP_79^XqhY2jGi1M4Elx3@cUU$XuU-{?0tBl^)nPBRy5L1 z_K%k`BhNIFd1~_UbjraS_B&n3>@g ze;Uez=CcZx0`@}}aEiSjCRJle>LeDW#-^ssI=&VyYc9{pjU~h)Xm=pfW}4QtI>Ij4 zH4w~dR2+U~3>Pe*HBAaiypmEc*Ow@t|Y%N&$b%Z{l)=K&-p&6S9ED>0Xt) zX082YYgVji`WRc02Z21;7|XdhLpwH20Y33pB+Sxr6Lmr5aOzjIab1Kb_(sdWYr%0t z<_Cdt)@YzmTrzjX(v7B1ao_An_0yn|CYkh(No`t@^yj_<6hKg9Yw0GU?1lNep48NT zKPt=fY~i`q>3Dk#DoEkE_7psScY1L|`N{2_hWoznytYg%y$RkeAP{t)EyBY8_jU;S z0!ch42o)In1$WUq0jvY;ZT~8mAwX($UKP10Msaf^IfVLYMOG0+;Qx{*E6E z+l#AM8QnLekK^(_w{6vJqRjR0^3v8j7S52-AZ`L?I%ry^n zx}7afZ!cB`H?6;Vd%pgRu&Y#lidTAKkU!VAy4}bEd&}pVVBd>)cMs2D#qOKsV!xYi z>l_vvakDQ3q{DF|j&$49So~w%MWyF{|DEMQCU_YQ;yV&6Q5*0sbhWhcp{+Y!QZRjK zI+Hb7IICMgntE5i@!PU==3{u+1n%d_@V$g?uY_yvA_(S0_cJ91|U*A^hN zH4V{`7pq~a1PWo*wnU=ED{fX^X9CG0WRZcDm2k#6Mp|pjne0>lX18+@3^cSnVe934 zMkyoJRK3WoKp^SN9D$SO9Ssw#t!_OEj4_2~hQI zV1q7Nh}HE%JtG!)eu9P>nMm`n#l7CfI8%{IQRPLt#vRMv$}23mC2udLj~ z7&AmJiENP~E`OVU0Rb*Ciw1W$!EH&kHo%j(Lz0?M5H0#gL|9*%c)?{rt27@EeK|;c zo0D-;M*AC`5kiX4Y9#X!ITyKWeqqHt1}uU46!#~(|5 zICt@lTcO(ewPIlQPoGT|^S0Wbka%A>=Y~`HRk(+3h@36emf#1UMJH`Jbr>db#t0lm ztv?)9cQqlomKycHpy{;2f_^UA{nWlryI&CAX5Y#+HFwG^r^>*iGOpOqgPbm{X6O|_ z^wrdcrAYkfI8nV=S^dzLR0;W_iqo*Ggl8=IzQpogDeN=V7+k9j3lgB%w!rBg`J}QT zG<6_o4}`0q%5cQSg}lXkfW8dJ47zH`eZ|KD~#YCY7Rt@@RbN$7)O{Mwh)svioK3+;O*$vu9A$1LM4k&~dGHd5+N_>tnqooQ$hZjY6lXK3lu> zfKQ!X?IU^U0oDyYqYx207gss29AR!H^{UIN_PoBtNQpeXPVr|_Ry5XOCS@u&G;;Bw z*ln0`!+{&s2 zSPl5iSItWIPNj`X&8f}h&6Sf<8^WL)1JQp;8&>*hSi zo^4@*+3;)`9PbAzpYf`S>!r=5{%r!@yUY2y_;CN5Xp^o#XVq+- z7o7RL)AeP4!3pQ}y`tpO_j6veBEP@Ae~@0i+WQqP?cnb)B7@os{J%V(+k*t?GOZu$aAM5gCq!_ITlfv@whHm)(B z!P>$D1{_lbS9z#rEYs&8}PZI!!_QPlS z2@qKXe{3gv`Xf;NF8jgVtVyqVy6N=N@zUP<+;w!QvG?ABKtg?fu99v$;(=uC6pimY zWYOz!H~blAMf>yZy21A!uMe%EnIfNkFTGJrV5jf@yWAJ4=U(^cmBrxj&;#F?+ZeuU zgEgum9XD+NGu&e2LfpE0wTNn%DgUJso<~dm7qi)q@a)iLj)9xeQP8=8gjkBW9n7nM3oTum+JLD+^d@uzYiZhjFcZHLo!Vb*J z!8*0$mjD*<0+#F=s@|z>&+8Y^>W;5IHOu1qXoY@25MZez5dS5uWgih@zPVNTDjghoGoD3$SYZf(=RjsW4E$U~83VsPb z)UROtuD)ZW!0~U~H+SV?q%*XEa$L=FOo)&@%)m@#Pxknk3@O zS^IHbXXtl_I$k)Bh`{Oi_3wQ8U8B}f=`M}+?6s3pyhfS7r?-y&z8Nl!zm#a;#$}3h zHcnj#eba&oP_%YG0ZqIa2B$bEO*LDnWlYnwe~psA^po3fwkQ#?W@hXchLyS!Nn%@q zTxmCQsEWXm!6}eTNACZuR)8UV-Iw0`a)@UT0G4n1uu*V8K1HC=EPF--vQw^wetDZo{J{(_ z(GIZuFtzE}d08&@m+>%BI>ZCjG0hX8lqH>7%y)F+qb-vxl;+zeQHE$du?)S*})Ku%yQBD$_EUt)5_NgP>P1Fk?7sf z?8$5*lw{;-P7*3t*>GtAC8pfK7WUH~f2=V26N6rtc>_;thT2x&ecnI7-Y#N_hi){$ ze9(DPX5zW>(#pOSBqfEQIEUU6_lT@zi_2|A0Xo-f`rYENd7%BzcM5-%K)rLxs%EW_ z!&Z`uQ05l!L#r?j*7#1cw4SKKPD-3IZNaWDN&2iT{3Hz1AH}l(^A~ zP|$oqoIMM#jgYoxvmm73Z`|u)X8I(aCyPM4gH&j+G|v{AV8KxiQgkxfGY!uBR*GhX zgxh3oX4}ru?UFSA-JLnd;`f1=86^Y0Z$b?H?qp*k{;L|86tw+0@dgLoTHU2yER7li z_^xTqGd?4042S~$>~BdWrjw<0_ccx=yPtv2-C-PpD?cX1xC3pA_&lhLgd)&HoYLM% zS?sTy#FpEl6734@iKlJu@04<3T6B~gkk5tpe+cyF%PkTa-&}U|o`HK~d;y9Cke^$+@TA>)|cdVNzNoEC^ zzr;May%*g@O_g0~V2m~@vXxjmbPCXS+upfnf8TFlRKRxUtNh6}VG2L`C;<*7)JOHV zKcHIXb=#8qS`s&GJ;R?lJF(I1Qmc1A!uO47$pJrwd?KIKP3Myd;XWQsIBMwu=$hDCqNa09Y^;>o5nY}~ zvSH3ijpcWH%k4kyAitu*bK(`T-!nvX0~tkbQT4n!&>e_wOhBy>(r45v*_tpqy&TrZ z4rAqw!q3^5Yf+rBIGWwIeo9G9sHH~$8kveCMBA^W`*guU6X1a59iiixqR$w@)vTOk zez0Wu?zsJV^FGALrz2+#S%U>GWunt){Pa{lYRO7Ts_95MQwkAH4LZ&`OaMzJv zbQd&4t6y$5)_a@OJmUqGUV*&!TWGv_SLx8eWNi+Xt_zWPBOb@W$#DZ+QOjuDe29uG zi)tfe597VxEx}4-<)%qUGwp1wac4$NV1-*81^adt6*i}VM;@8m506cU{iO3gEdFCs zCh^x|j%bac@2}Fmueh8A-zS}mr=sgVfN>`ThozH~YL3N) zdNk-O$zQ^5TUTgrZhRW%&V8H`NYz8)R|-p$wWx8{HIa`I+Q5tx5*Kq~lw0^fYdmQxE$L94mE<))M)Lz#OAhC;`2 z8RW<-p&lEG(*6{7YEQKIT`iS5g*$!YNluI&3wipWIC)LBP)Vwmf$|~+J5)EW#x9P| z>hb3P(R5Z}aW!qT#@z|-?gV%D!EJD7Ah=5)K!D)R4DRl(0fM^)x8UvseSrkR9=`wD zXB^J8deN);6aV>oHlKzq@Z`7aoghad7-`wUn4p{_^Exrkhz5tTxX`+4wv;!(nV~(3T0tz$+XKGZKl{Qi#9MoBsoaIL4CP1g5gV%<>AlP|L2I z#+4Ey?vS2;+5k#}diIf!w$ui*r~J^TEXTfCUXT53Hy8|RZpG%%M3SJ>oQN@)qY|-H z=7G#bJr3iX14_^S%!b;JmV$cXcvIT`*FhDXN)wb0srdJqUMh|2JE)Y~i@^poj{7vNl(6^67H_b*qAp;} zSJ*68UWed(M%}-b#&0zgPw|$(DWLCp$V#i)|6w!`oh)D@Cs7c>$A$skm`x+@8eg`@ z!U5LkfNdgiqE3CmLX=@K%&$s}{(A@CH$A@9X0RnB8O9+g^VE$UBr6SZk;pVVCfN<( zB`eia>Ex`exC}kKwilrh|QP$it1Ye@%C8h=~3GAx?2)jK}Ap zj`66?L3$H1(Tj=tDYEGndc09rI3SLAbQ+u5rAWpgJ%}rm5G4Ne@3EK4{Hx!w7>|@$ zM@Q>ylD^a{*WcX9a8xYcZp8HY!BvMY*j- z^~2OP>@x9H;aYZQlQJ=LWoQN(>(IfQlgdjnX;T=dj!7~zF9jd+=(NZ#JE|{07sT&~fr8tA-lKwbP5H6;g5?>tBvVAXxR!=K1S(+H|;h0pZtQPaLu&T zgZ8wAf1Kwp+qZjCoZf!FU5HFwfRl1Sl0Q-~-nQTF5TmRUgh!D2SjR7{w z$W#APVL;(6Mn6ead_UpxC^pGE3HuC#l2JUZm1A+~zP>)J{QXKVQ>^(O#fToe1%mPb0nb(690{>b6qTX~$z60^v6Ze@%}v7gYB<;#a<&l00p-5|AwRNt-% z+1T&_q2CoW>b|EC7j2+bV_rl?)S28-9du?B@FT`ry+f|gBNEA>HW3p;6u);s4aHx_ zgx)%3z@5pR&#N51-@mBXus=kSR-(ctNGLs|Dms|If(PLdfbC_{9@PsZ3O)tmf==WKz6q>+x$a{?hsY>=;|;w*J- zN>wQCI}@!`jo z9D{39-0y&n>a-de;U4Q`UKA>((Q)0w~)wyx*rLUhe%_3=z`+ccb>5- z@UsI?EH^$I;NEay7R#-B6sM#d;Uw6*6e{EiIMVqN22jLTTjc02{^?8oUGG47D6gsQ zGl@T;6NR`$K1sFx@Nn4Xv|>4Ly!fY~>MzPE>CHSq!<>rqX7MO$7QuPw3OXe52G&Uy zo`>uCbfI@y)rD0om0%Sytys8-IdAv8z8mklpn{9nsIBOqzV)uH2hK0V@b#Sd_d<-` zmx=nAZps9Q(12oTx|s+Y3cnpM7XM@GI+^Q_IsK@YlJEy7k4~2oQMZ?im$T~nWqM_g zrKj)d9F0a-@feA2_rPC%)s$!6oF3$ZZyFVBf%Ri26#Zi4*@V=G z*aCGZZGrZ)Hm`HEH>w`m(g3epbYnn;7|^*Is4wf#g{jg59c6Qb^v&z!>XE|5hV=~- znO_u7c;v~e^29solPFQjRMS$^_3RPU<8-aJjFTTd-mHCv1nU5zBnw~ju7~h7rrJRe zbI3~9=s^&h2Nm+a`Scz>NA~dq`{4?mSjo3WmYFGQ<|M*_48x68XCuv^9eApYr4!Z( zWAZ20x}HcwSBgaTA83NrwC4nG39D~*9LP=YO2T7$c zTi^uAm!oHsiQq@jNH5k;wphM%ISQjdqvab@(H(+;ijJm8#TAk=3I^K!s#t)^m9_v(?EmMWH`*-7EmX3eNtJo-hXqhD9>}h9nJ1Y z>j_{Y+dKM3Bm0#LV%160MptsWwg*PT@i$(uKdpock0P2eb_OjIPL6CvFp=$S&_it; z5(m+oyBpLh%O`Cdzdl(07yHn;cc#98_63A5Fw+;@;ISJ{bd+CBK-wK)e$^B!35Pms$JA(u)z(?PSh zICusv3;7D3gO&KScqC$kbo%!4g^Q}-N*)P<&5w20-M;ng*|+Bg{3zt^@!hi88R{oT zx!;)0_1Y8&lC`*_qkc0)V!;E;osS()2;TH^dKgOo{JV?&mDy4y7i5m_xM00eK zH!v-@eL*$#9a)mL&80j|&+PDlGuS&L&wIjoxgQU-hw2nE;!HPcH0>eqGKLVw{~#O@ zXP^l`eb+9=Yku2Qk0M8O;vf6D-to4Vl7H%;Jff(0V!?K%dqJne!T;fV3d~rA%aWe& z9sOXjeh)W#2!@zz{GXJ?FI#t%uDOdp4sPko9ytoS1?&q*V`8`~?{a-UR%?I93+46l z2#%gj@PFfQ^@~jxcvkQuH;7Jcbwno|)Zq4SF7(j9f{w4@@)5zPNe`V~GRO+Q`?K(; zB1dZDIj`%#@1ZS_lA!T7(qtZVWbQ)0Tooy}v$O&Y#(5>5s;|RI^wIgsLmCCcc2Lba zMljO!DZS#!T5T=V`m!TN^5Cwg1nhv?5$kU0AHQ1v(0`-1txI7pkBR(Kao?bgDa$k& zixo0x?1Pzymkl1X;6>p*2%D-v_0#sWqTVZH6nifQInTVK-yLy<3J}X!A^&x4w+^=1 zZ^0lO5WHD8<2iGDQ!q)s27j#Ohw!+USgQDe^(3?1q#^vsk__s@^V1GXk;!gPPN+iQ zANXHt1?Z2iBrnNfK~H>)#EA6|`4=vkl(y|O#>yNSv&*V&Cx&@G3@Z=;;un*6d8|9N z#yxlTH`$$}#w%o9%xtisU}VTqCcvlgr6%wx3VvxN2nsE6TfLS7!6c$!=kSzkpmKp$ zMQu`U0U^jFfoCC}I19^t#9VV#)G!7~Wh(`b(3HS>)uYYFFPU%MmHQikp3?oxEzqh| zZt_~RU5ku7BU3Sw(&ae{%#XSLI)m4-sGC_=W~Q;#rI;nxVU(fj@T;kgSReCmbW;0k zl;zNLr^CYLNWmSKZRL2L!BQBilRcb6H5H$eH0x7mm>4fy%Ea7L&;ZLVRI?bZmeh#y*II9T+LYD$C$2K1v%Vd*$s;+J;I(3=;Q_wbWfgSfTg5Tj^B z)-NG5N(oNjW)E5>cgZV@$e%=mbb0H!A!2ET7D<7Mj>*j_%rF{Cskxj>7@DQ3JCqW^YqY@U5IPg7UkG^cC9kQois+*ZDt!hWx5CP2e@R^11)G+wfu%lFbtA z6iR#X-GZX=wv(6#DuxrB{$~91huZ~B-MFtO&y6q^T~)fm-eL4%o|(jb>hWfxIKWPL zXMC$B+Jc>RQSXNjxj99`uvWAC-?rS@Gv#S{#K_~t4LB0}OXCkqh~W{rtqs^L2GBHj zT5C%~g4<59!dPEJqv&Nkl6!Ew8UDrT&&7>%zQ%Aamp|+2p(-+#gHU)KI%Q8JZq-#j z`^&^||7XH(1q#KZW8E*&2UH=z(eC0$z)1Orf3i^@_@pPU)uBFT-lFrgu~miD`sQ<#k!^KP+bd7nFD@|L5e> zSiuBwpL=Hvp}su9qmzCGs}#o04&f_g1F5UZge8kQlrpUXir^_oxaLIFPfie{d7Bgx)n6)Bx}@V53{v^DTXAoN={^KQr*m`)7- zL3?}{cm&!gnGfOHw{2dOpQ0_BzsB->5s}sUcg=NMx_}%+b0;sKy@bb?B=! zS_edTYigK97|^XmGZE7&BJZ~{!j-kl03KmE*M>`_jqxG3u=`FpzkJ_hcuK6Kk3n7}C2){`mu7bfTOeQ)5m)JhM2dYUCc$73Rs~)Dw5m&?a4Fo!6RY$-5U7bzYcWD$ zPpm8zMJ)8Dy8q#N%I4Su_*cRtehTeO{tp)P;!gtR%Df2KjWgq@z)ZcT_;jMU&QV~v z8yu?VL9l3_k-eLtMJgqfnQWe&>7EiMr2tz=tEwer%S9LWl{_81!Hq)=cw*djB24?WE{0bWjb>Kg;O<~v{ZXbK6(u`MAvl`*&7Ia z_qs!YAeDbm_!RA}(G}q&m%A)V#2|^8TxfTo`LhoBXS7QJhwB}J5X;hOuv@jm`yK@M zz6YV}X8(+rlTt_y1DV(Z=2K%>86G<#;$oV6uB8-Mv>wi|DSMpMS7vp056XeK1% z_5pF=S5eLsPiRSz0Nb~!+=zmt3~Sv&-0{cVYVD#Jg5_KV;@cm;48;84Xcfb)(}Z#` z^lmzGn1YLAq#`sCqH!4}o|<8kHOImX)jaecW@gkf5?^}UL6mx9?nMD#c8U{3&Q6LGw|W9aEv@_) z64bfz_;MCmOxYF|M6y!9TPl`+YPI1GoXD9wsZg!=+@yy`#P-XcW||?u8al=}diu7{ z=U!?F3NuOkdu=~yi}8!=cYJtU`e1y(iM_gpT15A)aju~(R&a=5*!36L$qbsx!VU)J zQ|fYM3&j)kX;iY1p=%W~v0d}>1^g7v3dehNe9rvjO($O^&Fq`LC80qVxSTCo0IUP2 zl%j%g!(bQFMunlctS@7%+V}T&pWxlT0JLyY7C?ssl+I_^l*$QR=?TtNrN+>6+&;%A zku&nHsmBR!xM6KK`B&iJyo@Hg7IKHvG4JkIfWf&D8FXl9%>$FaLXFD#MC0KbL-llR zUPBay1V`;M2LGUCUt6x#GkmIO;%eHOA4YK*^;G03L@Z?ovLmAafYPqvc^HDn*sx`5 zp{Y82l7_#X$HaLEHOvkOv!G3P>t?&KjC@wJIk7D7ogn$v*7(xvYT^m_M zhHl&Q*oKECN0=~roB(vAn_e?9w4gCrN)M>ia@qB}=U-`C#G3qxiO{E<8ZVjiZZ@2C1%!9jBTV(9QqxC!jPf<|V=)pAy;7X4B5*xDO(eU;%`VVAU5#|8cTP3m!8&7S4)!7_F{I`X&hlXKXqh)uJ{VP z1YDyU#x|oERc6os_mM^k$<$YIJW$kGe_yEs$qsr)p-TOjQUP?KXa1HUI<2M6KtJ@_ z)Ho%NgNU1|&#kP=V4KRlhLp?~jx?P4DwuKl$74fs?%b_|32|fKX%Y^iNl4-^8Ie=G2lmqLl6nUc0)If!wK zeOQXflRtQL4yMxJz6KM_uvin@RwtvyWQI0a{zCwdE(DpD2wlI`F$Iyj{|~s8)PD~t;@7FrsmI&j?XN4Qz7O@2uf8@Aa6a;1(xjkqJagIC((*XB zu%0pI+JCf_-;Z+xtN9ke7xxeL+HyxI^EJ}LMH^=j99>tE_Zo#h7V zE=aUwNFi(k<3A3@1nXqzWckW*WziKE6>5wIwBP_lMxnH*^U3OfK@>Z5Zu5xkz{gE@ zC2fOgx@HSSnnC*6{!bQ;>3!qWaot|XBGVMWIi$w11Cn}aB_W4}j|jHINw)?EObjku zyR(q~U;s@q5>l~)b!cSfkZv-2WO;iX!C~kGkcH4S$@^nbMGSmcE_hsp%Fa6aG;Z)T zUSJ&R5V1IXpX^R>JCO?D!llvOX3|zFMT)5?d)Dai#IRNq`l&$Ik+OVQ*5jBl#=$R> zARSzSpTbo-Z^Tw%g!NiS)tpRZ>QiM$GlQ#Z!~nGk8rXD)mW9hjL~EB$sZ-ql#{4&Ahx~ZXu%~7-iAxBb?P&I;vI}G1=bNy^fUllZ zZfTz0{s27rG%qXPqHhfDFlTsotw9JAHie~L|Mku6ef z?%fAqmHXJhn(}^uL1?DEyQabPTi;4x+U=iMm|qrU(!Ec|d_7j)>(1?kuiqHT?Y^Bg zO>#n8wkVGe&)~7el90%qaog5}x}ay_F-iIJKYGvwlqu;@aOjimD;wxVn>>KTW?`h5 zm{<+b)D*&M;Kk#gIAF+1!1Y~3tPM-w{e@amI-}#T&SE&|F=gF_ZHtJdZRtUfa#VyP zMg}zu^&|%rhi*+Syjh$@v=CLsU7;iK8ra*n;t@NJc8g;>`Y<4*Y}5$F_P4hVjJi8W`78-^zHge2@0(^VuFIbl z0@!DTBZyWLpL@#I=ivQ}65@m~CV5MLl>N7aG1)NG9eP1J0Sy4SOanXE9T#H|K};maw>!YR2sgY!aS6!r&ync?_Vw2Xul{Dyl=5ir8%sA7;x$G^`5 zio1D;oPt70{K#!$DsOVkYZ9MeUVqa>dLU4S0q8 z`3M{dXUIK*s@1g|GMjP2i*1m0^+e5moMH*lJNMeMiV40MnK4P}#x z`!7Ity=Z}1NxpGR!nq#?BChdmH{cU!Meh7eYNN3vuMT;9m4h?Tr-{qx@9=EP6a$*& zu4}$`?lz-nhe!YVgBQiL1iC@FW>k=4tllDSZbxt`K>FyW%C1oV_F3*5+*f736~qL@ zkF5uRAtrZbhw!S(^QU~rSss@|nHTb{-n%&eV^-ls8MeWKF!{w28G6Ei| zB_!VN(LWJXFx061`S|9oTATyFG&(~E=2`D%HAgg8ot`a}G3HxMy)SE*L-44?G_PR* zYCk)uDOwpT&Am2hmhj#awB4G(yN;p9+I+k006aQs#>>a04?q3DVE|{99(}|1p#GbK z3w3gPE-9iXc@P1shy5qM8H$vhf!sO%L)7FV&andJ(M*B7lsY=(*3yAb-&-GZSs+!!6HlvlMDup zTscPVq67fA2;sJ#7Y{9frFsFGHI(~#ouTM>-KMN1dVt~oF!g7zB|^|Lg&HO~yEcOU zUUnp4!2c^D>=)4>hur%yS8|I3-smN|*7`{;)@EJS+b;i{*C)OxQ8(71ehRCJe$%|{ z6!>}z0!RPflmRZpTpT(rTaj^HSi@PxK>&R)RFi2<(D*NP^+3jNOa!o;9Fik~_XRr4 z`wXQ>koz4Z8v#W&TKvZLsL^pLSqGA9oeyfZE>hg!gh=Mbhhzv{J!-dc{^S=CcIA1@ z{Q3xsX<|Qy)1{ZGs`7%-&s(j4z!WX)U?UD|tdaR2)aC_Hn7jnQu^q^c!^PrBE#3W4{0zGKZ#OTv{O>3wO1o3 z1GgpXR~Pp}2@#~$F)FK1THc%9E2dvATTQh4%8?(ztJtn}b#p^0<;qT3*3Gfa11&XN8X>ir?DQhV~#BDd3(e{R|{7ZHO3N&a+i z{}tCK|BU(1#x9*aW9hk;HlKUMm^+I~ECWc*?grdPWTdBh9#hOp|6z@-%8L|sqt<)X zQPLa7-E5MjbCps;!z%u6tq|J{!}pr5!8zqItyvJ3;(h%y$I};99An!0M6xxt_Gl}n zg!4YK*ATlI5=NqN=Htyek2QKkcI`M!yJ+E!L7^v1aA(1C)2PF(F(==dA$5Yh&!+UE zq`v&uB?Du8rxw{~8A-v|2xynMr(GXc?I>nP-Bk&3Y291R7KfHvi$5Pgav8nHRmWuk zkZcYvGk1Vo!p*E_ua3H~{OwU&F5cVSnEHiwjcLs1TCoz)WGl1=KCtJ4Gkjr+IVHHm zd0l{w%#k|FwddJZv9D?m9snmM$NWAU=U_k&WiI_%y27Q~|5NBWakme{)cHeNVRDs6xP za-&x+Ph2rl>tfVdm!<+Fu&iL))cH2S;;u+k24&Rr5;NG3aP0Zt9;2wBcvz$6NWI$Z@ksQ3@@3ZI~mZ&g~eE)qh$5#EXVIC^NKt~rs3DG^1?vM zsRi$HBk8VnBKh9;hB0hh6!HQ$NI_$n3k(ja3r@z&?v8tD7kx*wn9B-+5ke-b}8jQcli{(|^A+xLMe zXrz!XHjRAa%YV)J&dhi?5*j1?g6i)Mb{~`%+L_wi__7X~1{ahuHvuG%DQFQXu)4FP z6O>*9$3WJhjP>7C|B*E19vSS5k$P=*J1>c{3t{;E+Vp~i{XTl53{LX(p%J2GF>L_j z(D|#vW5y>Q{Q;ZU=o$;%HEKWA12c))rdx|)hY>9f>{=Lhs~@J();HXzrOFNIg{{={ z4}4@8`sRVH1~t5LmZnxYZg_G+D32H23mM@=jeQxb_&Y}51&biOh6T-k524)2vTT>8 z6Y}J7mPqNdT%g2$N%1i}ME;e*Z1WK+>7Z-2;Xs07j{I-v^R@Jo7)K21zuYJiq|dsu z%dXTLKuJcQ{C_U5H$X?$CXNa+jCI%X!14@+w2(6%zC{(hg#rd^sAoz+sJlQ^yLw?< zsFnM=n{Covu-RlKL;2L)yC@)5(o97d>a`a`&K}pqWUa-pz(OW@)Tyxfe6c50)@eKnPMiq(&+=qfhP zs+h(o3-kpPY?85Cw$rHoMd))R<(^iU=Rox|Dj4ifFR?%YsFHieiX94=)E}K@@>BA0 zt3A(?y<(ic5+APL8W>C<`fsDOJTS0pbfA?9egI6Ia);$8vys^6L(h^K0!i0U%QEhH zjIFv1pEKHjP?_O-N-a=fr^Y8c;QdX&dq49roV2!Oe+a3GC%0kjP;3JHN&d_w9B=%24JG{TsuLsViz0)VluO$U#TuocbA zaAP);e1`3$G4B`i%ruIZ1yy?m-AHEjHW_v2iJ#2dZ^Zza7Xn^vr7A886H2k1oWHOC}Cq)td1&yk2x{9ecL2VnQpG2TCFn zBGAvtaebVW^>?V+?a-0;difB$%B`l{C?l}~by2BP>uWwCxGmbIQ9~kxqEe{)n8jF_ z@+*aY)(}c$>u5|LkVsnne!$m+;Y*}(u1S#AdO_FMH?GWQO{j`$GZBlWOG7=aL00#8P?J%ML{N%f=*W{5k zmVD!|*7Xlw|JfGPl0QvPD2jej=(p~^8_v9QZc}(7q>Y!NdEi_CHdiN6Z7)t<78D9{ zpv(7`vRfvgi6ZL%)E#`qkp=3;gNw&e*H!xpO3?2)`L@G^lhM3+2*dEAfnwKMWkpL| zR$C=}q5j(vaEnjv&Tjo})69(v)$^l^C|G0g`>O&)Tm zJXGX-19B?LmewW@gn8+?Z_6J$4pi<+@@So?SxWn3OxSKSnLSSB6GL!Yy=(Mr?C|kA zoE)p%g~P1*SQW~2T1kptd>u{{Moxfy{X{(ouq42YBa&$#6xpgj`{YomjZ!g>qf8%S z3y3iy)%w6tK-iFsagh)s=Fq7W`%ZM~?HDR0w1hFRcJa#0?eyO>{~FWe3DiQ_WP+=@nA zvfm@Up&_!bJ5R3KTF7v#5tL@7yPpuWm9^N#&3IPCLQhmecT=IGrc$y_AY_=2t2wn1 zkV_L%KV;1`xJ}A<_2G7;rrB4@@hA)nOY2-FjRe^!tXXLp$*^_%G+XX>q_DWUgO3N2 zJ1Dx7$_NtfZCiJ>KBW<*`Ue(7!EpEfI$P>po?^!e8~1EUT=+cglppwCB}i;9m}oGp zp!wd^3fE|WB$-^ohVyX+w|zTEwDrKff=As;@cKDP{@V3r#I<8RW&aMj?vq2~2VH+-DdWB@7jCf@ES|aJnf%b=Ku*%pK3jJ!T z=r$i8LK`z&o%>3%;}FBXFFB>$U&;&vkIO{lLjt^&gz9;M^5_W8M)$_>-3$#p3xUQm z7SFL;_N6x;7(dS%c6cD>o$bdV#@A45Jk4O>Fu6G689o1m=f^KUQk{B^^S(4GNk~9U zA}d98m}O0PSlI<$s1YHO5y%=mA+m8WL3o&w-N;wRvri=Q`?#2%MFK4I+$buQ1HYY> zBz*elWSNM^bkKCdO5!@?b-Fj?>Qg`N-J%g9R^VuK`Yc{kVp zj`!nLrO;_SO&_dp|MTgEh1eZN@TJ(;Ona{qT!%2>X{RB+cRy}H4}|X}XEW-qVyXl0 z0T?qV7|YE7NzH6j?mg1+UJ=ewTG;H4mw=hw!isaMDc4hUnmywgpyas$=@1qfu2TAF z5^yt|ZsHN`I_2i7{i8nMknL4uk&&ytJsh6fWuPQ+FqtS z9AyJ%z1{WEu}{DI40zDx$WJ{3a*?pKG)g#;!->vPf#S`+UPVSlxYWz!o^SO-PY*}u zJO~$GXaMj963}Biv@A6o+Fi&rWJl-pc^!Uz3vvU{Hm~DxGB^dkKc-v+*cBz)pgS8VDLh=AQrSy21(I!J{V-+sPGDuCeP%WTk)@ve;OaMNEOD1P zb6@1}y5AN)JV>_vq%c~db;{Vzlt(^oRM*x+ENgI>d)T})T5%`UIz z=2V)dApmJx|5(ZIXl6^xb=)-?-4`CixpS%2gB++af6!Q`Qe8;G!GAwq&?CCt{Ruz+ zX$V)5yii#}g|yA8qLf%$mc=Lv0lzHkl5p1}X$WU!FFmeQ-@ z0X8?va=N?K{(O*Wz0KJgj+BHr1QA(}ge>Rgl8S0)Jc&8%xOx42Dn)gD(Q`)J_p{iL zXTDPR%k3mDUQOi5Ed|PVqo|5g%^&i0j8n;!-(P>46U?MX=k(LO$he%g z&Cp+gWq2D|#*}-*bosoIWf6Gyod54M|GVzWzm2O*jO~{r!Q=IyL7!J_fj8ypK&iHS zCEsJ=_6K{?Uloc|9NuBwouXi=3U(70GIhIvVc<^ojG1EAh+Mtu1!*S@h6Q@|4KA5_ zz%_l}@=83V_4l$@v%~l`9D~mYXha3O?(1)b0>*>5OU&)pUmrUx zL0#bL!ERyw**@>%qhumA%%>@(yJ;1=mB}*k*2oq^59kfxRK?uk*J8Z4@#Pig>+rU! z$?Tm^wBNWGOun!{5xBT=oJ{0*z_IB0^TF8Ef9TbdDJ*kHPf(r&E!Xtc`NGGJ;}aTI z{k4aHhxZ&@ND2+#UZKt)J8_XI3KFHo@l~KtD}V4NHHtE1gySoi=fFges#(jND(Wbq zj^7ETh6RqdH<>3$9C{D1xTD^Q6l5_0V}mI#Xien*%fzetTbwaIVf%HB?JchS?>hYe zv}q5C0g(~C=`fc;Q3Z{qH!Xs0xzuxQl4pQYTS$yZCMU&l4uKEXf>fJ%7k44d)|jqk zQMAZsQ+;;|lQTNGeeD_B06F2S*&ZxUq|P3Yi7l_|U|iXt?>DTg}ur z>TpQpB`I9mL2z}oHUWO;I0eqWbp9S_X|>uMsK>j7IClNB@^PGBS56sfDU3Iw+o725i#IL!>%xQ;3@zELg>u}-N9<89=-hp7=Dj9xfj%n>p>8J;k##uFg zrKyQJA@S7IqRQ}isanIBGHja#msYt?pAYVW)hP1&^MIk(ev{Q^$oDCTHZH0L_a(l3L|XMP zFbhgiNP~S=!}bBes8lD5X-p`?xT?Ju(XViJ+K zV)QkSnAmgn)vR#Qi&_^5*E(ds^m_IEtf8SiJOmG#C&f`)GBSAl@V-*`6|IkOoDBL& zv1~-PW%cr#e`Ki1)EDTaw0RmcvixsCKG_lr^YNAgSg0Q-G@$jqzDm>MhH*`KYc=hF zIOTI&SGYg&_@(-vy)#>lc4W09=!UEiH8)U9o&eKF&anUUo6z42Unw^u0{w$?xKGo0 z{+ox$SYv_FTw7%tBoujaOQ0-cIAg4k(nP_Qc^m^`1fJ&X4SD+Wj6G+JSwLUrl?|NL z8ae=}>HwQ(99hV|rz`_>&UQB*>mgt~zO6#iRewG0AjP(pdv{RFniDg??LL*OJ1X7~Dbw`iD?f{9~{-xx=6pZtj2=frXrI;fa?iBaC~G?KB!x z!ospy)8RYex##R^gi%;(Hwr!A@#w1nqA!zzt9H>qU}|MTsmIW7sF~kCorr{k_%3EO z@$FF9S>c=?yDG%HAR{)m=}sWTJJF)Hfbj`HDNo20EOog@VnC7!gSR~c!F1zf7Qq$0 zm|$SDoYMRraeiuDo7lI@EI)}RKhstG+lZul;`Bs88-FNt|Z%pr-t`IOwA zhC76lC1eydz3b*S_SfPAlJU-{)d#|9@y{|JRS<3)DH%Ks@++gU7=Bj>#-jz6W7zGK zB+r@(ho}tRdy|sV4A3x!DG`>;ee6Rshs-9x`xpadEi7@|yqVj&I!0nokPL_r2Y|9> zq15@4^w?xw(6D+o)rWo&MxVJ$x1ccuN6w{~0N;cxyd1B{h`6jCRmaSs&cL}+$dPI^ z=^6F7%ph~-{9b@nW3@9Eu1A^qvWg`UXO9gFUfB)fWD3}U2NM@lMgWfDG@9&bUmmDO z8YLBG5b2G@R?CMV!1VqAh5wfN@iW1Bl%&~yRm@bV?C8;_3*iGemqO1|K+uRRnxTo0 z=P6f#MQ%B|lwIBdT`KH_pmL{t4j5K4~c0WsxcGoCQcC(j+x-|ghu^Uf(Nqgv^ z1;oBR_R##lDFTmRC!q;N$H93hDRxYTO63s)TwVS9xjtA1Hnu{IE)-J^0>|zrH;$pv z!62YM#3xa#%-HD8o#WK-kboDW`mKAQdU6V;;&W-OliZhDoFkg;C|s11Zm=kbasDLv zR|RRGu)?^eicsyuF~u(jXtsV>c87>}ac_8;_Lnejk_D?8ZOTEp86Ryy)s=?+$YrWI zhiz&1XJZ%Dc?~Io-S6o~ga{D4jH-&Nh$#AA0ONAX%pLA=W~^(s59H@UOyB&k#@@sMOUZR!k-qbMK&)KMl8NF~lNRq8>ulba{ywi32$7Ln)SQQ148` zD#`)gn0KRB>$WMQ{3YJqkDtjfhNH8P1ES!MQt}P`tmTkBS_P@vOysIUCRFTe6@W+> ztBq7o)r1u;u(eE%(Q-y#t(JrsGKDxkrb?n)Uxmgs)Swlp-_iGN?;=MEB0!36?>GymD+H|*D_7nS+Phh=X0QvEGU&)H6hWGkD zS2R`u#1F6aJdt=8P$86XQYcuHF<%?^DP|=hcTmqngFLQV_W|n|0TbHFUn%xFMt!}y z1tqK*Citp0IRb+)R9bwcX^ughN}z*TsK%yYo&<;#1^Ln@P&OB@_4LEhyw8N*LKh#P z$wXe&^BX@r5+nJ4Da2KGm^8H+o+V6@tCiehSFGEiiF6{S!qk#=h3=$q=0gt@$Mu4DBWL_ii&0$wZ0)TW0-d~#qNjjITs2jDe#3%{LvZ3HHMi^Vqc<-XcQQ$ za%jVr1C3q5sr-kd3_ZgWeU3)%28l^5epAAK!XsMXgn-h;eN>6*ZqkaBF8!V+9^S$> z9}0~fJw+|Tzo7Yr8-%A><`C(oOxF<}0g6oI>T42N>etXuiWcBFB&6tqLx!CJ>UG*n8>SKp>ZFe5u! z&F5ayPdCdYctCu9c+EAYQpic?@RU63ME8x=h8?M$T#??>KYgjo9QCT`7hA8< z6tRbHSf)4XC&p7EbU=tu8nM%Gwho*!qC7hv$qPn#JOwsH#_{ES7`21u(fwV-BjiBH zVv{NM)`qUuMEQUy2%To+AmJw3hP)zcOHHnfND+m|((RV4EX3<%o7Ik%#EVxCJdq_sSWnK8ciUKGKLz;wue~G^d)((A zLDJEUfFfFI$?BDnG{<-+w`6?nQtG%>^6`VY50XgusSvwaOVpmKRn9m!3)PQzfNy53 zyt2&w8MjBlhpvDGknrp}2$&%fOF)?UAA?(W9Pu?|?z0U>@COcU(+j@WyjFEzjFkLv z*?+0JH)L}-t7|x8oCCo8j|9S(jG<=9vKhkBwUq#>w_p+@?ZV3950n#9A0r zDjkJLBs${HeEFhotMkP}C)Y2(`KiFCij-wj&i-W)1N5#3*M)c+^qAyynFbhQNmp)T z@zrMnA>d4t9d~OJVkT>T;6HI~q|4xddhq+cf%m>|D0Qz{8Tu!$p2>|r#;g6wJYkY< zh7w8sExW$z-STA;*q#IQMxNc{T=CjA6X>XrOw_c~*?&VLCQrSw%d}zd1B*>9me^I^@gM508 z46RsVxFyi)5}@kn=$yT4dl%d56N$sriH`WBY%HtPB~*beh_$ya5c`dV(VB(74y}yc zP;ZOC3_+CYBcMmS1IKc%d@<$eo=e__%5=P@_+vIpn=$uozdBW$(WIs#ZRwc-z)U|z zc{$2BdbN)UjNIjrFyyCTpsF2TM&t92msDObC5ZbGb0Ioa2#K_=fz|BL&dIIstK*EM z!qrmFR8fp8loZg-z&M8)xm?X;$c{Q!wG;H%mEW?&Mq=x)W_}{gLT7xBNZAgr*-{lD zh4Zuoy_vUL&fyr)BVY;bz(ij|1=_cnw-=4*e!2X&hEHedU#6aTLP9@7-5$Pk@I`^rO}5%cC1EKoXIQ@APQ3%vNJ;THLts+*C+Hv`r{no3d;d6!wH!_Mb~&0U?+IGI z9NBfz6dRG@k;hS_*S_I_9^T}7o4H6FnlbymeZ}lcH#F$+8LHP{!h(ee7{zruO(oDJ zNe!YYNsH11fbJisS|)fykfHTy0X+}P#$Taw5|ZaF%l(EL=W8Iee(Of@>%HQ;3^%FK z#>YaX5Rky+U`<+RH%Y=42i24BI?wASC3`W)hNN)yoBv5kUWZ@@K!ZwO22hz0TApOJ z@pEuvtAi$X^41jxxZgVzrKvpv+cE>iXSA&z9US;4rs7x$~thc#P4U{MRzFK zJ#i?{rAd2~DG>(>NFA~XL|T6*_7SG5x2iSJ&vD04Lh+W53o%@|bWd{G#YUqis=hB! zb@#{pyvp0__M14hVgXFzsgH>TLLBlw>o@5P+Xg0yo9MsLq+boqmp=|BHk5U3{2xr4 zJF;|1Uf69_MfkBn_(j4DLzGap74He!`F*O-RhR6KAa@fbW`pN}GAfOsY+W7ZLmE!jx%SjTbd`k+hm-B>@Bs)5e zI)5Ji7J_V^tPtIXV_rgk33UGVW)k|%;E(%((k+XI--PgSPB9w^L)qdgN*wvDISIjV zzuYK!r_^!|s(2{QQSjXmEes)L7Ni$ithniSoaV%{9MR8Pii}nO#6t}!U z$}uk@N06Fj#$Hqap$!erdF+2AEcNhU)Xj;1_8MQZ-c)X&#+N_uY$B)N=|#qqw|~Fi z^lUm3|3Q`fw#Th6P)tg_Z^;Tq9s;2VYST&{hP3Q|nL{6}JPvA_d*zA>izBf8XTPZn zAAS)2_sg;VHUI57->hq+S*qn*Yw~_?3)2khW%$Pc8Y+a(qZGad1qC;#t7q3b|KYWY zAK#2pxw`iaMcp1WFq2nKRL-8z@9=e`%?xO8Y5>7SzZ_(0Sbl(d72sfPh_TW--MYH&RN`*JL~t<-bTG>J6w%?83ymMcsC zO)|3GdEMu%vK;Nd2r?Hgq`rsuU^ds{gKh^EXxAm@JI+th-(sy&IIe|gI5yPU!Lx)> zO17^ho=p5SQ*4N130c>?f#s-Ju%$wfd%l06GzI7QT8~9}5(@ghYC=mB%vFUMj%l}6 zo&TAdaH1`g9aZu1-J-5qtqJXZJ>t!$vCf(prB?RDg?f#?s8Xw2ixVYO7hI$raP5lB z3n&XEosZc}aw|G_gk&;D?RKg1g%~4p7MVm3w&@v=)n}vyzhgCEg04LC*{Bp!Vkb9` zA?hOIoVuywkuTRNJVi9+10Q&b|FfTgrgHk;eH+Uo>qARyzjKLvIh-#)SAU{ty;Qck z72X^I^zAW=hzG;t4_qw2nB&uFgcYO>kn{ReEQC=`RPy9b(7oevg{%<=k|M-t$_h%Y z_Z5#jQ>m#8gN8Ees&(d3ldxTUc^tS0XP^|EyY6ZHV)uWy>Af4dHazno?k*Q3XS(hZ z-{SKf7Hkj^VI!=xAY(hse2RYco6zr=lM%v|zh<-jt78H`!jVEBxwFsmMnLWwIR{lS zq=NhqCOc%&hlo4t{#Sg7YBbj=lkB6i`Z&l=IGX*Lg+*PN{JkV^Z20(jiO-e(d4E?78ByWH2b@$dhwOG zNhDxer`|j5@&xuNn4VUmI%OFB=$tRjU z_&eza^J0~M8=J69dIs_b7j?gkK|)63WDrpZ?|efpCo!Ce(Z{Z`INHJSzxhPanLYO? zIj>FzKLtss`-Z}cGk>frL@(kqa_hU+l^*t-bRp1rnaA2?Lgd}jzApo z0*-*SBO6E=;e54v$O{?jctarnUg+M!8g>fk_4O1d<{lLvDwn$hwU~8T5-O`t%4K_ajW^m-1)mU$&siG48*xmvfAEiaNakkW5oBSGIN5D+7Ev$PIk# zVeAEVu2o`mRw5>CGB{#vB|TiTwZ3bk5OncJJ4ZZQHgxNn_i_#75)Bwrw=F8#iri+h&tI zjnM}0o$qhG|7X_9Ju_#{g}pz!Eq#4cOmrp~>(Mr4OFB?756DGQLLw7jFy@!)JWfs)ypDxrfssxN-{f}%UuQ0a3X ziepv&0ahh=fp3K_+^CRorZsWKS_@%C_L`G<-u35Vy#VHUS;oDj(AaljT9V3|C1+vQ zCbqfO37c}3QJB1Xi_=M0)#_{2D9mqDj6%GcA^gHt0wKEMB}E(5)NMR{XUf%<4897j zq7=hRzkf1NM)%^Qsigdz8>AabkLVeB(zA!iU0Sb7nD|PnO$`6U_>eXIjY?anhVQkq znu2I0jyq2GOX|mo20>j;ouK{+IqtQiGWn|!3gmw>kwEp+)Mpv~1q%mQ7~*TMQhTD{kHFaynxU$Cx}fjd%b1p@$6HlK+~83s{Pk z_NYjklTyzv-Fh@Yo%bxhSu|}jCwEGJE@;4!1csw0{5;KeS$N! zT!Ajsg*2&t9d9tqMP<^W>f)&`Wd>I3g@Be6SL3b)jaP?cZZThfQvognd@+VtdBR?1 zr1HvlXW`|cEQYMm0x#do{@FIMIHO<`1!4N$-4btUXPQ5!do2c7X&AdAj~K~>E3=^3$e^Ffsy;i5rlQkCC4Can z$DQKwY$O$Qeo{>PWSs2C6BeZCbP?(pRgqq@K0_Ou%j&vZ-2WET2`wW9w06JaLuzQ+eoiJzf&t1 zSGJpAg)HY{rcn<_ecFVTuh!J%@#PIzOZS^1D8!&d}JA;;9+eVcVQDTq)UVLRf= zXYw0K^BQfRtj~(T(08NY@M+&wXYT*b=#y>Ei9?WDD(~eJ!q=9znT+$Z?M`LBmln^P z>TvW5CX&pl;(%v{j7-}_os|eJ2YZ45Be~yQG|Qq17{ulh0=}57$({wjTzJUa=Jeo< z?AiRB4yAo7H}`oJ^@OiU&WFc`MWwWbL|td)ximk(8_g7uBI_TTc!s8^cEqAmF8MDt zV5_f1y?N}o&ya~Q3(S+=<+mPqdmQw`hZQ8PSo&EYxl4pHoMtoAOj6sFqYqk`NIB7} z<@L$c06IMjtSf=!8jt$FGIsxHdeGhH@tg#XJ-Bd+t)Kqge`Kqc!UO~IuB9+MIm=XD z)&KL*sG4zq9>R~+)i7a1$8BosS(tCkfL5>;}C!G%@nIP)>`xY&iNDg{_npCMI^mIp}QhNh9c> zBL0f-2gsQuFQf}aqI|LPgrKzBNND$Cf(8pqZ=>Ix-E`$ymeOm+;6sR$E+yq2mzdQNQ6Tke=K871GwDNcdV#WB# z=h7*CVh2s?MJ^f#0D7H@)-69mK*kr>86ZF`zyPWxVspWul#uYRIr%>VQ}2W!s^L0h zjF!8Jv^~jr+o%clUz;>!}nmd_<$qe47G8;7M;uCpX5G--j7uoiQPhX#;OD)WqvmP? zw&&B&ZQD%Zz-j{^n{z)>=)&|6Il}}Ol>+MLWldr+F=c&77e;;Oan8m(v;;CpL*F*iEt?3!Sj94#s^S7o^aW| zpj>e~p|akp5B>zixdkr6gd3Ufi8?6p12$0n07~uxo5#f?S4&z^)h0^dAtM>`m*4#h zRbpHv5h#A9yu;75pfEd}i&J13=7`(l@g*f0BpaVpBl%%W80kge4+=yD{*xU zleKfK4G9%K@|BU`h^(^LsH8(fN;+O6ZK3Y2IgX#I6rqQr zP)MHQ2%RUlpXYTQ?#aaVbl?bmU*g2qKA4A*xvG9Oa-iKPo)t#NWkF-09!9osgv6aW zxecVr^IXT!wT8)XvTOzRQ8X~6*3p!hth5SYX$|6Y3p zf~r#}0V`hRhsYzL;No2VL@~^9hF;`ZZr!H9-MqW{>7hOFUwhl>*sP22acS>*wboPh z>`{RQ;2@Z%=dlk=r&%FxHrXO-hqU6+TevyT|KnGoqH=4;pUW=vPpLU=F&4Yff2Yz3 z`xO7zv*7I{w-%#!4yMmbmWr8vhz+d@_bUK{F?)Y~Y-c!R^=_-h>>QyAxz(%M6#_S+>7cg{; zJklowzrco;v3cf`iEv6FjJKD!KqlD!eHM(bS@7bb;C8jmWlMv6u4bjGN8>rG!_t>; z$E)xa5ZC!SXpK0bJcBHGHA$@<<}Hkk=FrVO0}xUD$7nMpi;mTU{snbDXI(GkQgQYO z-+oM0w$kMIFHdU6|eg|=EuE)ikyk@uHIEdxNR^0-}in|E^Pzv7OBy#^PzYFB|A5mcih zNoamd{-@v&QTNX-q1AE}()D#>pkSdWxoDgbT8F5i@gG14uFVyW_K#2Lvc~w;ZX-SJ zom6?K&C!Weow%@~^RE2b^xDEsGiYLKo(M(ZR=)yAB#dZH&~aZ@I-4vZ_PN$Nh0@9F(3{tO4cE5?^P=@@c0(_ull_1ofi zRgfx}qPURv{Pvgj22h0`SsQ_*CV%?@n*NW9+W%Egn&K zPO{R;6!b@5$$?CLzl(O0%1aK*EPAG z1mzJ(3kjcb*n+XzW* z$K}&9@XPDLn6ECO&>lu{2yuBKqh;}gg%u9G#ijX9AgD9YNOJRX0_TY2ij8t+eQ-Ct zX?=2+`to{N`L7SQKcI6M-ruSH1hY&G$=K)$PT40R(;UBOE33;DOSQK*!0kRX=NWOj zO`o09&hAzosf%;#@{F|aP9auy)2gaJOuS-EuCcFX$0(BVd*Bf!@o6R}z&sVYBr>`& z=@N35TYvDnc>pNFTjq8K#=?uSo@9)V`VN!fV!gXKu_e~3{{=k_jR5k-X5qOd+N2`g z%(1_^VfkJ4BkY$hO~Y%42%2B1yDaOrteELAqytdj*hquRG@D9R z0+-h}KAr~c<7otfHWh&;mJ~{^XETR=N$YrdQ{0o$G@)#AU-S26v1+T)@uvHVWH_GV zR<&GS^XK`tr`+ryi>{zlZSEsG%1tQ>_v)L;Is*g8$mnl}x^#47nXV~CoxiXD2>-h{ z@bLMwYx}>`{Ja+OuA?%sR69Zr$4u^QM>XE;Q#ENnI$Irat;Pf5|L$2L++@ zp!#uTs;qfpe;(&Y9!dA)jYXtH6V1o!Wgg-C`(t-mcLVVCVy>G0nW^(GHU0X) zRw5OL(7q7SYesz)`|RFy;4sGvVEx5J+ccUH#cD}1Ed3RH>{2x%bP3(Yi^FQr|_Pku^>_~^aUi}TTx*&`s>`6E` zcTY)2c{0T3S56sB+(^(rau@B;a&(2cGsiu^=vhV^kxU?bccEDBM(dDXUeEqG280+N zyiW;s9&#e6^Dxm*CqovWZk7H>x>|fM)3ThRDqeG)`puWJx7b8e(f&EoVp3_X(Jxm$ zpkxpTUb7-}{A4g$uahOe{5MJ*I-d@NL@Q6Uo+i^e^#18P)b>Cjc;2hb5vu z(wfbl7>AuR^;*wG2O&r5S*M+=CKz)LZ(zIJonbL2TuNTpM~FDFq0V*$tT3iaO=1U5d5`xi?SK+;8|$H*;Zr}O;{wJHr2$X_A`oZliOkFh z^JH!SMaV-JGJ~i#K(K^2^2wKV`UQJ}gS5_`_mSMmMl`i=@cM|jI!n8_ItNG1x5y7q z7rWnXY^0q>lKNCried*<8QYIU_ZmrLD)4e8fK~n}ZZ_}OdGuWM9XZ)z(T-u2ho;3d zywmTN@vSkIfP)?k=Eu4%bUi4N9l()Y|NYIGY4{G!pljd?8j5xotJNt+zpnC}xtLoniEMx%%hL6k%)8AN; zJhL2JYn_FiL+LX{c*RE;;=#2XEea+&o7HBpvTomP70+MUBGR3cwJL@g#MA zGXFqx-_*VFTC-kIR#NSUBq9$BN*t}iFx(ZKDGh_@7wG6}Rn%OD(1(ND6z>%r7{<)6 zf-B<7+s#hAf?Tt0MuZu?%Lp=mTjV`etVUQ~7{aVHzxeGA^Gf8Jb&QpjQpwR8pF$+k zL&Ozk7UeOAC=v_#i5#0QGp;?3?I)wAEtUY27L|Ly&63I}UDJt&2c#z5TLH1DWduJTOt zNs}ac+yUz+7Wipimm$Nf5ve56blG{45{QL$SSc?ZkVx`SXvM|7^ofB$AhLuT>e^Bg zM6XnGSiNC4;Bn)~3v}r%e6!ztZ6ajHGf=!wUlsvN*p;aG^*xC_bXwbL|-+&a? zu0F6!0^)Gsw02_x0ss>*w+9KjclUXHw-1>YQ2Y4)>3Mux9TCpJB^cp3p+9;F9&2pV z?xsYrP(2%Lt~Z!lGL9(Tv(E6nFem?57CsGq5NnNF;0fzow=l~qcoZd)ba;5eMqB#p zOS;2`t*=ql%-F6rzi}Gk-h3G8YAdPv-)L-Z18Ox{I&zv5XMXl;yjz0{KZ zmlmA^4CABKLIuwIe}C-O_7G_B)3O#x#8S@XV`{*tK_V_={Z1zvC?go~F5;$*jfvBt zwZ+hLcakcOZm0;NQlMdRC7^u-xKaXdT?BbF-^Gh{h!umlvwt{|#ivZV%Rg(cEnAOL>#vq9*>Kt5HE}f6Zdot?eMACp1o#0<(lv&(Y0B!n zKDE2K6kjep!bys27#c^p#YM=72{(v#LaaA`g81^xX~WWMh$EL>fxnC1a;~U(NTM-A zs68;zDDYU+;lxWwF6PFjFD|k!K@m9M>e@I1I5SwjJRgf>vgmQiK^p{E5DGAX_~n`V zx0jd1HJ0C?5evd5?bwASIypp+sk9dnm>LjJ4-| zZH7YZ!y$ZQJPYyCiqn9M!71i*t}lt6RkBYRzBQDp0Nu%7wU%!Ti9kOt1{0>kZP4qGAg;5qTyu9up^PSfN#dv@5;NKZiGpM+356 zxzRP6P0<7v45K7q;OLCKqI*c=pt=lsTh4xKrxnZH#H*^9a?OQK(o+m<;-Z$&fb1+Z zNJ*r4&M_JV-n}79q@Immf9puAVMXle5iKTk!%zl2Vt`K0mC7l$oS`#aNKuJ4lBS^D zu6y4z%Of|uP*Q|k4b#wW0D5~*B&58&uilQn|1)0BPaYVfy|k&X)_p! z3J`mnClO$LZFK~f!>(~fCwK$(O#ESt-xX%1h0bInzu22NN(qjj;TtymcZ1e zMUk47qb9R)2z3T0;EWGy$rSs3(#^xu z(~1Um;a|&lIglKAU`)~t2#``wBXzS-8pI;Vu&?fjNU8cIKi9^<$FCnHVFsP^?Bs5J z-SN*S<_^WNqL#7sDRJ~u2Zu;X$H;w5hF?zW!OE8M&$IojmwWVk3bU7~vji zE1>cipA7a4&#PjS5j%cQo`Lg+s&y#a`x0-cl;h? z4&Yn-Pd@K-zQ6qoo$9b+;ti_$n*U`rj4V1V@M_oT-)9CsXNGa5q9Np!SBhP~cd9se>s9OT6 zcloj4@yE!2-#O-PwLD0A?vl&C0-%Izl)?@wD}_a?lBchRO_g<^nTT4$9){#51WZfG!*-h>GLO6?NBV8|)}%$b@C#Mz2ef!b@`XF%KTY!8i76 z${&G?b)79N5FFPC_%WrKgQ-F7mfi6oJL4S;v$9^ul@0?F>C+wG1wo^In*<$)VudQ^ z2Z=bXZ~OjH;xPb{&g%QrXi}*5V#>buoEsrh!p-YZVdz&B1_-bd#&XNo-gT8Z*IrQ= z&{Eh?nj=r4CZdIk6E}OC(E6HK4b@w`<~U73Yut-6(Np0_SJMW;w*$^e82`Bk|7eQ^ zrB5C|Z{9T#F5c!L{0a%INQ4Jz^E+zlyMF)I{YM$;PmB8GSfB8`rXC<<;$vn&T;<~JbWP03IT_lG0{>7=Ij z6QDf%1q{k?NBEJb$`Oem80Fm(DbvoB3m9dnl8tP7L*?0kI2pvhhQLNDA&`YAt$zu} zbNGsaPPo*0i!lB8tv}Zf1`L&&d1P(YH0>W)Ft{n#Fp5$jj`xuq10But`^I zu9->c<7q0>=*4D3uYMVTeI-oMkYUx$EwwDrv2ZhyR!iw|%2eEF6;~ZHeMmaJz*b|F z5Omflo;%B|2``SoVwrm~hEVBwU5KYHB-u-Hcmnv)cJRY|1?27_=mDhuLn+`sD>&GK zM0Fr`&anMB7J+1R+%KI7?L;0S#7-IhSw%D@+-<+m_6b0?x zl&`9L_2LzNr}36WR7NQETsVosvN;T;crD3SM@Pw4;W-gj}4T&I^0S%o|2+a^m1@>^Q)D$5`|94_6i}ndoMMbl6;gS2B z&}{~dEGc-VLkOO@kx^27th5G7hYB)z2<&;>$XN}1`IS}TK&;?xH<(1kSqchfH^?AD zdY!t|6WCFXF@26ivFR;VHHtgf1G$;5rCe|Fipk9(-f>3nUJJ{;uwSBFdO24L4hRJ8 zCaLy(Lvll=;+*Zss_ezEaZS;pFRUb9f~~UOshSyXE*~$fA`UBCAUb0|({Pf1O<#!G zm7{B9lkRlQG-z{2H8!+d20QstO=77bM&uX=)jK#Ltd2I-|G3Oy2pLk+jv*;s`zLW81`iqV6g8hF-B``AB?+-`MYTLD zn2^f$7NsQ!uWHq~Gbd|#W~g#9N1O}5%#48YPkOY3Z2KDF-{Sk6|AYft2}~0!_T0ZgoUMVe_&SGR#6pk^iRS5j>b1wy5SpXR ze08#YD)s-iuuT-@0jfK;0Y=8ow#Im9{+Ryb^~tUq^-ct_hiplN9XcXz8cIzVr^~$i z8Lp?^IJa=HB9^O7b6-jUWjW(i9UK}hR5nARO?Qik@;HLQSUSOUG!?Y-p4Ooy!ZeFDr)X3Fjungd06lJ3U6MEEC05NGU|_ z5k-zD7M~SaSbYch@X~|D%yiobRvjsw!w-{6-w&k-TBSHl{to(bu*~BMUvYps;NWaH z=Cp~H#W#OZ4aBSy?rXF5i$XXZb^S&yzj{X5fsvbuc+bYN`Z<1V2q7FtwT7Os*A-8h4(l)rsPh*)`_#JWN)U^ zA~sP$Bk*Mn?bo#la@1LEJ+~Xey)6Twm@Kjv-q#CT7EO zen!C^o}MppZaSJWW*s3FY{Lg5nuXM-Tfuu%~#Dsw@Gs6Z~)Y_^RH{ z+Sl2r^n*cv6Z}0nVD#?0zGq={1P;**CZ7_$5tdEV^-L_Sa#b(ZsQFK)a%O8JaPgY> z_13pN{s=-+&*Y%MAya*z3)=JMWh<+|3F>w1fP^!k3*om^P-ikTg0Y1_MLkI!s!2ld zJG95M;a&UQ2OijKnre>e5E~ai#G!P+D}(l6_;O$gL}Spo zN2!GIT^j=KYtZ=*lk(cRzrs;H99RDVq~{E=*7&MbPe>Vm{`*U)^;940)N`-Rj(dBv zv6$n--MKGDlT_o~UDZDN7fRO$b9Q!VW$6d;CQF#ffg^f=9eY%jFUcY-vf$jF^jUaz5OR znGcSDO`B&C_76d*y60`a>4cxE+vM1mT}|D?nswn1QOHlDBpD9BDhhn;@Yn@uTV&MG zQRBYc5mbirm+0o(d$x%c8SOV-raZsZ39&|SV?#y&TQx3ywICvJ!zpjT;wf4lb>0Az z8C_O6$94RUO2ADj&|(t*q7_HDaYC@j8-w2OLw^6!$t1!CQP^}RcwXD`@&eMbm0c3! zon(BJx9~x|ezpT!?M=m^9M zPHz(_-(D`8I&TxBo1m}CK(1)vb-6LJO7j@69y!YBfJb8t5MK+da2C3<(%}>`xLm%c z*yn3~03RA@mGqUEl*+#*j5m!VtIY`BDysC2WOJuxSFBNG9#88$uc;8p_tb|@88oOe zz1e+}H_N+W#!xY+nEhd-(Q1dMDO=6ZD95R^JDH}V=$qV16^+~Z*LfG{_}oN!VWvzJ zOnjFrGBc_u;*~b5OL4*A&opEm_}Mw?ljcX=`=z&Se;t;pwLYZMSm@O}1DtBV0VYHP z)%2=Z(HYj44=caFbw|xDQeH^;yud~9I`kR?nxhfaysfIUzTO%NcD=Tk<6Gax1E~X6 z&#gzY7T#R_T0YRba{a+;>tLy=bw|#VwBP()5;}CsL=L1w8s8XJ(WpNGoCX)qX7LvN zWG$x8hexk-00##smjbv_m#l7yityMYXz&8b6!dFJ zB^vPY1ZX}}VU6@$c%Z>pcb_tGLyd@RU9Lx1**R#hx+1?r@Q7z4rZRu^&N#uH+G?Vn zq~Db8!B8dG-7Coy9O@wvcQPsNw9yXFj-}rPvW_c$sLv1c6e65>sPM0}0!0K0)B zG->q#p`Rc*4$6-l2#R%QXYr6Fo@C66Uhv!`|C~2t!=$2te36p&ZEws?@8oUuGB8!4 zMfU@A)Eid0j)$zBf{5V1bEPD_NTOsZ!Q#7_3Lo|g`Bg?VSEV1ofy{v+A_aXU?Hy_mY@tRILA?{{vVP2#U4M%D`rM;!i|6-J&n z8mIc%`iINk2A&_)oa9MdG;V#0ja73HPD)DUam>ns*OF{z;p2Y*>h-jY5GCv&k?(YK^U&Cu}3yemK${sSjJ8Pu*UwyoOpX}6T zW|ytnt9_ddVBB3OeEW`TUDo~gZmFO)(=0^$a+l+wit_zvI~2{GX-%>K0 z%k%I353}^ZzI^yy-8(;E_}FdhO_?B0hWO#xZJM7RsfPANGld*)Tvz(K1+S>y#OS&c zef6?;B4y=?$?<1^)7=#Gc^teIQuF_P;9iOa#yzRNi%6VG^{(}^o9mCEX{&5MuEcr9 zEVOw=b&!&0Ym=DMEAyFY`&?fn&`GJ2TK+!^K&Vt8xiNLx+z7`i)=oO}km<1o(dpOg zLx-gvi+bGr;V|)+J3U(1%FMt?Z5q_yB?8w6sL#{HP=716uGc;`ox06 z7y`GLYWcA#^lcCbJAv%ipeu@GD7B7Z&)~TGBBXpw>~S{c%4uAqeYj8Pieu1g#iTl} z-!UQ|kbY$rOeM`aMEH0h>R%`74qDyf^rD6N;D0Mn)77KPNO|X`TJ&j}Z-MC6nao#X zwVtn&F;;dFC>GVuLwyEXVo!U7h`F%poSq&hNSJ5aE1Vg|tqA7;C5}v3xWYLIGrH4K!KW$QMZrMB+_>4i z_t2{fU(b`KfKnhpLbTW|po4Wc#!)w$h#WM6_%xz+%^eL{N|O_X^p6!a+p1iZYxB!Z za{Z@mQSQpPz(UT7Cj!fmsM}Bj$w<(+`^bI+vkLdt~by;?Jc(I^`t zI*Gp;2|~fh&B83aSQl^bOhsyM)IG zaYK>=+SurQ&g?#oO=i#%5Q?@-_(c1wA>0KH-~E;|3Eijbgd=hM!vecef6vi=S$SIL zZqcNf<9f2dgn5+A%%QeX6VYR+h=(va7rY7KhtDoGRe)30=xiF$} z83nze?g-8ba?pM>*J&ght8BUOi+C3FmyjuY5ZsJkA2gE=GZz^i8H6>jkYIbpRDvuQ z0*Xy+wYsEI77OC0!ThRzx?>h^CFgpqwqrPR6CiMoe2|C*XJwSvboAGXj<_ zc*&olxFSM-h}LjB3#BPNCmxYwv1Z+D#OO zbIO=X^Em_}>uPZI@ANvOAk}c)>=!VgqYGNUc)d$z2t}$QDsKRTB-F{a)Sb@|2P8*s zM?gQ`zL_z#*7{pEujkm$jIM`&9J1kYqURDJX)w~M zB1XuEPHi?p1q&cT@}?nIN#QA!RoxWvzh&bzo{f%)v<>7Y#E91K7N(V&Smn%;_^YAu znK#K_;|L?_4|Gl>IiZZ4RDI4Li{`^fK~O{~Vgk-o7nW4$8zr)vQAcMEwT=W_^v94; zo=kRjT?NLXzrl=*B8Zogb#%+By%e#uCY@oiRr2Ns&tRi+oE|dAlR2(a*t`rr*A%tL zM4=%u5fuCnl9%553&wxMbdcgIW1B>=iQp9$I~|T!+XSn|Aq%4?n`Opm&L6jN05XJ?=f_JvR}rSZ{jzxrH;7uWE5h= zRXdS_E$~9pEtf9bz6!KZMRUt(0!osL0x&#%Pzu#UGg{IN8Lol%DC*jmg|^HBo}%gu z*T@aMxg{7eUN;%q_tA@OxEkCe6GAwU)i z%v+rdp3+7sT(Uz5GQ^@hs$%)%$Ax#oU=9NHEP4{J>U$!4ZjPg6-*5j?M`TVL<$xbN zWYNmv`q%*$gS&6qeuPD!oMm> zn5YjEY7P7gqI4+Y7IvthoKlF_+v8U6CxqPZK&a_};&fC~@5O}|g*oFI@-OPk;@!E2 zjr8t_r!Eg=`Gzk+aeG`AVnP=EEYSl9P7bd(q3tA zsQ{0ba(sy~SkTe-RNsz6sm-dveL|g3Un(8kF2;o+O{HR#o=6otpbo`T#PO;j-En_C z1OuxHjbOMIKB=3boHd15n^C4vRh=EB8B+_lP+HBPW~WY(F8sSOXD4l_OVzW&cQB?A z7h!Dsn`mK6XED0 zJiJQ>D$*r&38h{e$ed7@WEAQ|UWH9OBhb9ax^T$?#g|h?Ldb6uD_>g5Sv6t~!SW%XVO z;DiS)1u+>5{xb$kFsRT9Xp|JxqF1MJo0h{to@enj^xG$0rcm~jr%w+?GuY)qvA$ox zXEVlr0u8b%Bsb|dxXPavqemhN5XU06)sH$?g643)?$t8uK*&$>lZz#co5)kmGBoNJ zjej<8>FD>q(JZzFW(zJt6mj7`RGadcNyVSKjgvb2?Zvyq-K|Kc-&&Y`jF6X+0{;>4 zUR80^h*6^k+A~coq%_So|LY^Jl!7FIDq*QpjUZ4Nm4z|S`oq3xM*ba&`J{%U={D7g z1%m6ZyUfvKL=X{-iy8?+D2gJaY`f+q+Qq;ek`QCZo+A7xK~(exM<1%S5yX>&Ji&b9 zjAttWoLLBpMDaBVzgEAEi)|hWez@DP6)4|LLCCS$Mo>t5Iyn2-iGDVfKtUq)50%ol zFGe4Cy~Yi2iztsZ#b4fdiT0;T(h|rZcq=f;gA@5$IxnIv$`z0#$wl?WQ#FP7gIQxu z)l&NPZ<_3@V>FfT9xFfdfA;bj#xRn5u0qR{3QXOf^HOfiQ*6q+qhm;QE6J%Ur*`8p z9UVV2tIGBPdTnLezs#ePIt9WU#+li!yzC6Hrn>~AYNzypsx$mR(1mJw#|V{Bf!=8aXnJyGM0{@bS-gYQ3=|86Bfsr>Sm5DD>*At%kX6|{`uxQAy?iv z76HB#f>ecZgrrxRfT)O5)Wb3dcD;i(OlZt2MX%vtwb0DF{z!D1@V>0XsPml&3BlA` zKE{X=8Y+@I@nV*UtdhMA;ejQy3xy6!>ICkl;$O@bS6%YjDACxH2ug7w?ome!A2Pz@v1Dl9gsVH_LP zCX`qqanj@@>(=c-#vRGE{K@fAgE+~Z&(fsHT!27`o1Z=S!on3l=*JM&Qjh4W752P2 z=DdJ^42nBlcFKj%=!NcA;V_V~u*9TwZ~4g-H%H-&+r&RdiFbvZuyNHsUP~d+k-(iX z$kNVWUS-A15rg)#>5H5+?bGRjT@*8;o^e)c`X_=Bya^1en*gW8A;|2hQo}ww1MDIY z{yeLE9wp?Ji7XLTz7hy6F1riruAvrzi!R7K{53w)dAfkS`>F^`(ht6sHS?SnJ{cpB zkKitNa%R-IB{<_Cdei4*yobP;HI7x|J`Ea~Gzt_n&kT8=V%=ibQR=Wkt z>~m~Z$;#vh?RDGdye@lE0m=(NnntMME^Xdla{XhR29!7O&f=u-*R^O0SIB>`@IiK} z6yYn?S~VuUx{i%~aEH|tTgpi_D@h`62O@mwPpV`t>Cs*K|96nQMsWA&pIB)jzce%w zA}J6c(SSe(u7uk9u3p+F`Yn+t{laSZb@fw?sKa zNUy;(w@_NC5i#a~k5}7SmsVjzd#6DC2K!W$QTU2bxY^MCdp*oi$27&IsT*nyPt|(@SepEfQ`ZcDiPtQAGIH^c_=<7J6 zfCCdiSr9~DqKk(v@IxqUb5Ptm&Udp z#w}%GLHL(PT<-nN+XPJZ6bMMGE`GrA>;p=e3&=Hva7n2PBx@J-N_5 z!3=0*0bAU?K3q40##gALpHM`XCg4C^=7CU8Jw! zFacj!;6K2>gM}(@o^SH!*ZUhhx!)}7uQEmWu*5YLua2Jp0Yv0_5WDa?&RR7t6%kho zPkEan-*O&5ehI+s`v4+3;Q#rLmv;hJ)Bm36m~kKSToyb2!ovr1##?TCNypd`)As5N z)&f`8m>*-^VEgZKJDK~e7@vyb@j0hNE!t#wPf5=Geu;RE2MKx2{9LBb2hpV{lI(Y> zBzf$LuFp0A{}JX#fy2T`w}X^d=Su}X7DqrFwj7>LIxsg)6bOjShA+|G2gZMI1NyOt zoxpd;wzY`<9O%7+L=mUhZ=Qo<#qMPyMafPo&f9MvMx5+hEJnQp9BF7Y%P1L?PulDe z(k>l!GZ(AOz#L0B*VG49%*3m-qykzdo0?5vKZ-g_u;25>@lIxd!QluWolX-<#1jzx-327c$O^h*{ScHo$JZl7tS?1xCu1;s6t%)|d>nP>> zoyvO}PVU77QjWRQe^E_NHH@U+*jm6QnkWv*3h!(HlsXxY6HLjXk0 z6wB;jli^TmIZ~D9BGCb1Z<)Z7}%SG+Ub zY>FO_*mR9z&)ca_$g&HjIkQ2Kg2O!8-9OKizI^Dhyd#KgnMBAb5R4XL%KheHW$#SG zH7W(R8H$bgJ%NSIt!<*~X5g!sYlB)|BKjPJ4!{EX*whKdeRH--ig?!Du2`P@^qJuw zYWpLQ?Ap zmJWU$5a8t^Y%~wq$0y}3IdW1_p9A^F6TnG6+npMaex5-aX#q4|DPPEaFF3o3c$OLD z*H`P5T>CulQ&ohYTO0gYNs|a;i;kW@uud$k^QdR-9GigIa_6?-@nx*=C0Z?RpF#8i z56L_$5P@scwOnxOOjN?4&-ne`_*Om4DqDjWNi!+U0O?yU*n-1?vky10FzI$aG2pJ` z!WYF|UXofa8}uKS%)EP;dgy!8|6}f~y4w2UtsUIm-Q7L7YjCGnaHqHxcQ5Y6-K99R zxKrFI#jVhyEl^6&=6}X}cRs*7a+QnhjL6t)t>2vUnY>~xF;z-G)K!$xQ@}!zx%Aot z*q*cNY)yT*R30P=Il1NO6P>QJn?NSI189hxG9k8Pr$?x^Q^>#;bAC!8S$zA-i*ep% zKFeyxdHX7Yask?scC~bYRwBeRx(}gd`k6%H>T!umqGZUVPYv zoAiW(;x9966@I88t>CPF+g!#^?1~lA3ps4Xdb$C8CAq(6SQ7Wj?a8apyhqK|#sVGP z-ZJ9g;>py&bB|mxZxbcuXyzSI3AJ`=2r{6b)t{;)c4T1jGTv3R=^noSe!0XY61|CH zRd8^nn)J9od7;!pZA!x?$LAy<>Qqha$|xi)kFl>(r*u;KDO>Ybs#EpkKBk@03{I}Z z3cT^pRcZ}h!U0wFW5?9Ftb!)FDDM%bH@~v3q!Wi@qX&2Azed$-hIiq6>e#OAO4)-$ zX$Bkh)oeG?IuVDrrFv-Ec2cz>UZS@Wh9cQD%0>iPa=HXY1hmQX+u#zb&$HO4=5rqn z2Z{GrD+hg|oz(NWeKwmBbm+_$(L#l=3s*pn^ytvBQCWN|WN9x<5?2&qDt*x(+U)8e zZX`QnS_TIh2H9A;g#I4zhycfzz?&bxUq~B;)FDU|S3+XWUvF4vg{eYVX^#7uL(#{i za=?-uOMJ4;PvZU>J5aqg6rx+W-FgH($xs*>b!^LN!whKSkGx0vpT9a^*ZFu#`>Jm& zBAqQ}gjh<#%BHEDZ>8s$ykY*g+JGO}&zFphX zgI(Qve3`ZU<#CAV)>KcQ(uFZvzlT}E7E7ri70S8&iA^(oa(9~*fPFg16+zdC&ppzv zXO56AAU@TvJ1CFr{;sqR#kEG|Dox;6>z=hfvw=R&rf5f}fBi)!hD4ec8fP-Yg7w8k zVMW&iS~*w!usWl_sUp&~A00i(y7p;)D1+b*y?}%(+yR77*Qg4|*HBN_Ldf`Sj01$2 zcbcL{uOZtnus2CKY(=fniepll%g#7JXB=8|31XQajV4&O8+wbgX>Q4AVBpZoD9HK4 z@MW+W-!tEK%%2OZP*v!&%@Q#WhWps6<%xC1X`Pei-rfFKk_;RFMphX}bruIb2u#J3 zV>#kp6V`~Y#0MXSC{>A;l4-AJ`#C&p9ERj&6sf;4m~Krf0|!5ow*#W;S$~e41m-8H zHXy=z172QdHtIy!CF~>~>hZvrH7pmR6eQOygPoc3x6C+&XAL%Pa1W8;l5!LKe{NlVEjHcY}-xMufc(A zQ*XS+9e0YISRF>AS`f$dp$12V8-zrwpsseEsRh!OF`^W-*3_4yQ&8q?CFX5}C(-gg zB%~r1IUvOB1JHG4K8dx8#AgEV(i@`c{2gcKn`!=oCQbfvGz+LkhdqwEtXOWTJK%9C zRrH!oq2cduJkXXXD36bY`g>55VUZ`IMvyT0AW7)b`USmGlz^lb$cxao)?<)j#2!*nM(bMIRwlTdPJ1e=?232#b#+3`eCA zaf1d{Cs;OK>#dyxbxg8k#A(JGBALsMBw-q!u^5e%}io9LT|{L5tz8%H%~*wfkq5p!_Il z6BD9#$K@7Xoo<42M^?g54{+7?3N|hPM?LoVFnTx8p^}qz8%>i@ za!bp3iO&q&{3EXzcjV$bK)OErPDvllhsmAml-#a$rCr$JCLp~DOaDcyMS#AH7?YoT z4=>8JT@G!W)=C2M=PueV1oYY#M?NZ9$(SK;9?JEj6#p&xXj1OT5Aw51mASMS{RM%0I-ST`|PEall3UrvHA|ZwbhMYh@BaU&~vS!r{1nX!`%|g7GlpQ#k z3;XdLm&o}bjl7#prkfL_Z)V0|G)|UM!ZS~=4X2^Z`uMSH!EubJ4G&&9=A0B}s)g=W zR7SrOUpBWpHUnkGN%wzmMqgjg3a6Yd=-Eqf9<}8RE8c9lN|`FS!bd%i3W#YzI1>rc zN}a%AX$SCO1py|0jJz;mm7h4*W>c?CG4?|iDSn(vr8Xn4WiH0Cu3dYGdGE=NrnO#~ zo&PT8UfAJla4CZy_FmC_UWnDWK;yr;F2K!iL@^vskB@n+) zgK$a{s)E~Eysp~FcN`w0h5Ew0lnRdCVLorJeeyZZwF@YT4Fek2yl|=#6)FQafYaS8 zv*K1Sdk^Tj~p!;si*%WNdiqTDY1n^ zJdwTY;liXp<*rd%UQf7vk&L%PraN2D0F7k7wyu#BUyi}=^NITeO-$c)(%VCf`iFF6 zQ&1aq^|pu1b13z5oKotEAfk`_NG^+6@54Y&qHfYDi2&}e*xrR7t5O(?stbZTk}|R_ zII1jCb64zVqs{h}sPk$%5{n~nbZq;p#%~Fp#ZX9?sCTb)J2BID1P@(AbXFYOXB;3! z0MF+chD;w|0Ik4XnqIKR=x>x z36pHr%*3`NTVp1zYahQdJAp6$?H|W_Hw?D23n*QwX9CoE3hKrf&Rn{8iGH~d|2&;Q@Rp0jIhVT@{NmZ}GpynKj4_}oepI<-l7*;a^cClM= z^GoB-cK&;0|i#|zuLUs0@Yzyc2+&gmyDVcep_T$~d zn4;`)Z}h~I>-T};+ItMI3$J@z zD=yDpG7E_nhv^;){{oJVm6<-^m-OK;Fmev`HRcq!4WamtGGh>=mh%~vaSUpWHE3V~ zzZKZ7_DZOs?+~MfhPZFlt^@N?{hr|MSYDZ`P8Y)tZth?XT0=Ne%MPk8|JOZ|#eUKR zyKIU2-QGa#>NJXvy9ZKh)dt(h`_>H^JdVvP*dWV#02Wt;q$pI2cy%#`$Gr50j~O$n z$A88HR9<;kM-R}T;zK+Q4WHTL!JlH$u*oMlAP_mLbs zp+5XV{CVz~vaCinXpe_7Nq5vWah|3QU4XxXn`3xJRl&?(eUltNh8!J{@_W4(__bTu z0LH({va7L3R6z*#TYALOpS?T1`V%jg;AMGcFK!jo2EqH$Z-H@cPAyX512eRim-O~H zjNw)aHNgZW)7d0Fr9jnV%Yd;)3FEN3Y$TP)4)P%kD2PbgSY$C+drWg2*5X5q^^wN-oK5SXi+XNroQ`(Os@?p{*fC_F z3S>nBuvE73pb@g(9D}PV9h>!Nhk`du|3j}&Ouh26Z>~mY+;`T>;@`@_r4#r$WdSJ` zwrbUU<&O)QpkI*_lKQGwGw=%ui#D@Q&%z}#ovsvA^m=$%Ao$>mO&__BR-WiTMFM}Y z8Xc!*wvP!v*da}?xNM)U)@<=1{EJPHG++7XBQEK?FfFrT3MM$Ywpa_3S5Q_|k?ft4 zgk@}u&GuvDaFP^7=&#EhM&`L`ByM|5P69CGz*~ z*UU9k3c2v-+@{k%P{%QAWa$SY%$Xr=iV=NDxc-A1q1=;N^-QEEMwj-t+`r(=ale(O z7-@NFuO9RE<>mKhNzPKe^6Vy+M+r^-Hq1hQCH%9$48v{>upntyh<%a&Fh3Sxp_zmG z{QK|6O%_1tY0e6u&XilhQ}$C%%Dn2SG!s*fjiC+!Cu^I^1 zpHFaZSc5%K7``5ymrK>S4>4!3(s-OXbVcuc_D~*BkY>y$L?X^z=TPBh0g;` z=vjhzVJzsmfygK1YZAyzzmggk5D0O{mLY;wA5t`_x@f+54ZH*G<$K58ucjcMn+Mt8 zcv=lNAAi&khOuJT3vwq)N}Fz1Nf~Ig4WkaNT*FHOFOBm8OH7lQe!1@zDN!!kKg&D~ zkLB2Fl&1m_ke@BMI3#Q#&Tv12GVVbcJfQj}0ch@jf$9f?6^o(0<&KtP7RG3Aj zJJUpgS$kynbTb?j@Yk@Y=nvU%w8)5IG!mTPtofg)men{I!i}kQ3AH0IH%>jlhu9LI zNsvi0pVytsp*d7EBVwP;POIn^+Lm1&K#czmsXG?1i%FtnNw+K}Xe&aWH!Ut?>W_Z2 zrj`g{S5U2^FKS3Dth{7uO0J)dlerg%1T(j5bi6@VyE!OeMwxvpu~ZUV*aFtXO8b4Y z^;!46F0(-AbBj?1om2~LLTz{Q{KqZ%YtZlg^sC#tZbmedf?4PwXheoN)uG~fkm^~w z(>KmSnl`OWHLfzSL4rHVQ$AFpPL@BM&uxe`*3Lv(x+zX}=233FcfWYA1`WIW>NGBa zOzV7#1H8Q*Zn`!(W@dIN0$-oRr$2G6u$@1Ws=UY6AnM3Vg{cTd{%4oGy?Mur^MFs@ z(b%yXi9I#589UN=ClKkbjwf#zb~fzev0RzX76t@FERbsqnnr+o@ z2gt=!J~9zW0F(b3fe)!^jDOZLD8SQ|i;k>S2nI1Z#1PaDx!JIrL;(i9oXIkYx|@pRhy)XP7HmGM&>y-?Rt|SXlTE4;!%J1C zOPQLfB`jXLCSTA324~rG+w3AsSqwYGQewa3ftjo0`L*`>#hVddNRe@PR-k=2#$?9S zK|X|MHwXVrAs3(N3aC?fE7#vMI&yUi;N&InjNN#OvTKw4!jV*iXgL-z39atk!wz?eV<2F)6R&05%1jyy&cHl;!zLB%URI^c-Z; zZ5X>VaTJf;vdwll%NTYZ!|1ty_=`8mUqP;SL{c;Ig3Qsk=67QPs`qu zmZmaeww=uyh;+y4?eh=2hz?N)I%h(c!W5cNp1%G2YcBOM{&Z$D@aRjzr|@(M6q;+< z7U>hI^HFmz_XIZIYdfEx{*y2DQ6?d-7)VPda7(QD zVj8#>v1iD(W&whL;&~j_z6>|5k5OOo;^^4)jxq&VfaMC{D~(M9Buh1+ue$rHc^Bpt z%1Y|%&TX?D+a3@-@5PhAuW24JCIdcOF29F}K4BxF_~~bdL8f^5^PO8!^lZEk20aUJ z$x%hIbd>LbaZtTBiXB2^9vPz4hyt059C?(5*;P2vt5iOOTJatH7sqoNQ}hB*rOAqU zwn^T;(jx#p8$yE8@3dRmOG*K{Bt96=x4)!}tvZ5qA|m3NWuqfb==Nc>iSDc20iFjcw7-%!pT=+Klen)MqfVN7 z0caQROm+@f!YIc#VL*MPXAGsai9u{ zY*Dahs0kp^V@W<=(Y+VbmIpYqIJyhR2UEsx`!g}~@@nL$Lur%*ngKENwGN{;44*bq zb=60s5T?=P0fDRq^8Co9;k4~ z+M<4JGMjXFcua^4>P~=>O0@fRXtX72lW*SvMIzwRC?Zo`@A$qS zK+OYnoXG}t(Ej-J$90Zz#>&xVe=$Kym2qqbD3g69Wo-j26(WONfA?Frk;%au{1`&^ zP>v0F=C}QthQmuTZPp_p>ZAci8Q!Mfox)GwFj|~yC6Bt+ya2**E8@N7C`Th0j(;0` z+W_z@zCF)=69dd2$CD|8icj$FPVchSLfL$_-4(}Q@?Co-ExMkEFaO$Z5kYM_s=cPDq~ zrrx4gr@w&?Kt&`_-yc{NWqjEIFaw+YADF(DW(bTeXbkWbdC92F4r58OA6%gi=YD%@ z|B4R-4%Rc2U*y-{xL4lWN&eXY6hOL9``Tq*4&W2&4+ycU9%8$paO+^{qSY0^>D0{W z52eJ)Tw`d<=a_}N4(+S%tESaR2I5tsp8>h|mAw+dDf4q!R=|>S3{u2>4$1XS35L}s zWlZ+sKaYEb;Uq`J*ACLChe~8^zy17I7jD5{@cmr*<7=yGws~-OT920|?dhn|bxBsm zPGj$nb1T(x(}&tz{E^oH-3D6jKb*^iU(YW|Zv)Fn-x&DTt6aafsNBr?)l(ywRFEMc z-#`4jZwl#8rzu8|hNVVhXxcnHPoFR<;V@7{K2Wh?@tzlsGCZ z+&Y>6*_;*nrkd*>_>2?=`ejH6qtRM{yVOM!;n$*w`WEkXY5SJdT|<&hC7zO59{4Lk zkDQ=TfBQu8pFiWg^p&yux(ADd5%N{x;tyb|$os)%2hmGw4Ra7%@a}uW@G4Xx zvD%_i*|yjUh(%F4!%iEk*o|%@rSFHOV)^Fv4$Q1@m~Nbwdr+_7nB3$^a!B^#hrD1$ z1ZmaSyKOJYB$7XgcA(`SIUnb^fd#{-mq+(8=U;6Qf`9T&GE_l)SbK* zTe_HDJGNK@=CGDoIahm_C2l}ZI#g;_PKAIrB|zZL#yg2)VBllMsLHU2q@rC0@|AlM zyI-apc&&y`Re?2;w{L_1>W5kkov<-gW^e6=dzK|X()mlK#6UZQry6POhk$4OTY1o$ zWh^@MaV(ta!IQUcRdfT60Jl#pE>y!zedH_dEYWtD_*gePF4mtb{7kLyPY>pUs{(F|hsXuR7c)9krQu+$& zWnGJ)U9wq})iJsj6bp?1uJCH7*fW2uxQcCOHlJ95ZHEvs2s~I1Ef*gWS_fu!srf~C z8uuVbZF|@61wkMI|7iuQJg?unh+xTV!J|L*<6Mrlod2833;~;YX$9o)>@S}1-n{^c z=5P59Z8By~3Yg}P{;|3mxwy{q{;obGD$JNQ)Lu{E4mlKkE@KT$%xllZh*%pf8-c#- z8T!=shA2bnZp z{&KrzsDlE{v<0-FxK-Ty3fvlPj)#>4{2GG(qDwH1#VT(fkiOgZ6**>_t^ zT?V^3{wBCwb>E{iApJpYo_r#6g8F&Df+-_-<_Q)FQxVn);MEPb36Mz52CTamd^q~Y zu#S?tHrst7aC)C~2)`w#H2lhvgKM+$6m_TNTinCYD7HoQ5TK7g-x|q(4e2x zjqNo&pgr}N&KqHMKot+OKYViu7)Y$0<)mPeSp^VdMjQ^0Sv16Cy~Y@id&yCM)UHJ# zRs3(5P~_lM9=5A&C27URD&OVPHp_l4Kn&8EAd>6Pa!C2mmi{mQUIMGB;`*$PzIL)J z)@Hk|>C>iA2GrU~K&7?mXVlauN28gbgJq$v^FQ7SH2Mr;uY9)x&tzXs4AiLdDqoXF z#ImM^o)0Ni-ZZBM4hLXK`71N@V>GBrr>>H%++I#k`o9k+K|j~>!^{DyVmARgdF;T& zm0kRD`=}P?o?U@^G*=_jRn9Q{!bItM@_<3cXj-VQ!1>f;N=vqhM;w>Zt%6L?TV6(h z^j6iGYB9<|l00+ydN682N&w@!&pY+U`hWXy?Em~;_?Y-DQ5_O?&oWm*#D#kiiI{n&qY2l6 zFuX4sVxefr!9xPqnwzC}HdYB=WR~o4rIE$Jw*jA>qu*V$#B>_lUBDmzL|=`WyD)G0YF7SFQy89*aa@+&yU_F0z1^xT=`i-v!^kkPc#G28MH#^*^Jz& zfIwc;!X~wo?Vi9Z_z}RR!P)o?7mMC@C1X>f?$u%0 z00y>+rY3tTJqW}_epHkY1ce^Nzrrw+rQq@MWg%V(z&LnNNit#3qg0jYFUQ{Ihg0UO zx?Wl{c#HSY+4y=abMV{sHTTU${5{s@4>Xtp5%qLkPz^eeKK9An2UK6t#@ZC{(HgL{ znvmYAFg1y)h{Sl=N&vV7lOth{2Wz?JO`^p?+0y&Zo?xZUf{20MSgyXSAuiw`qt1Td z2wTAsU=t}+?3a(vyPY`dznwzxUbl%MIW@8>_JGAp#0RwhW?ag{p-c|>}7 zZB&LE+21au3KqE=%}ArxG~-2n-YseTyReh|Zc`q)AoJ_;ZOA)z z`5N{9{dl%dS5aKgK(*rg(%AuNwUPX6L6#?R7Ab1#OxMD%;G7^p$@&ENo2jpTkCaD* zSqC&j*;AXye_K{bAG2ZS)(|1aVUxy}zPh5b%>r1ye*~`zbTZ@V0%J(7I`MN{13@Gv zY_10HPm8J;9Rc1I632mXUwcGaiunaPV&}2XIoyGFX(|8mLp0t?RZRbO(F5r%UVK}( zzVTM7y67c;mAO*q{^0V#uAm(;97Ws*VxE@H956qQG?moM1)3U>@f8Xiqm&4opZo44 zQ_)3+4ZFZRGJz>LeOkVp`L9bG=Hbqy&_p0Cu10@fm?VQ$)7L+ARFAQ1bVu;h?}Qz& ztNc;%&|w|@1XIn7>+J9-C2Uu6*k^i!8B6&Wj!-aXYc5fX%wRP6mq5g^4Z=r4AUx&^ zm%zd*^my4hL#|@pZpzB~H|?lK+uqlsL?QQe=H4F3Nq0!i5Ubu?4@kV0N0uEP;?Qjl zGLL@YzTFJc(I!bBHrbK!&JAHx9z;xv&|P2UD<`3pfBm{;X-%H}N1eYwNmafc-Gy)@ z!4-6sXG!m{>z-M2BtXZYkbMqPaYg?@R<6r1%>-xz%)nCzOB30+EkUrG98F+Wk4AGV z$81ldJa5z1|Jwo(j>G**0HXGVQy8wo>LBft5XO<6SGL0tRMe?=nGMZM$ zb=P$8H0vG(h-?&9O8vBP^gDaB`$S6px{8Mq%4XC-@-v(JB1OUvx?xG@>WK0Eylt*Q z=(OvR1-NP(RjyVwvTr34VnY~B9#ty*0L?s8;|1uSG|em+4qXH~SF)3@+|OTL<(aoM zjV)~C&3W}p+KZCDV1nirCn8AhcUj}||HlF-$Qo^2(<+h-N^z}!Xz6@Fuaz*FbqgFR z%#Cm2UA_5Trf77GnC9A~1R{&_XijY;0zfWX9RGqX)q}3?Gn)c!uAFP`G@P-CYyx4A zO|tpiHu96SBJnG+z-iv8(aoED5*CC}D2oHuS`zyn_b7Y7Q7zzG&~IxA zS7SU91=EeO(d-L_YY#Z}0KP817|5}>idCJm3SxH1KLZ~v=+_SEPEtxY^Ji|-_-IUJ18?Jvc8c1R;lh4Crn;#qUUUlsa^qP!~Ia|Y2Y(<`-HnH~~2?e?g2ps)%af&g}5olM6tk-3IAl68q{5jfo(Rfv{FW$@CUz*lYSNE?fD1- z66+F;nuX6(l{pWfah6endasTP7`h~U7x+irPxl%Z5=?vAKa$7@e8{!U_8(_bm%`tN z%tlpfP7RE|mG8d!dioDH3Qt_VGQ~=%-U`+*V#Ylx_eDrgiV<0y?h1aK=?$jU^2;Oyj44GVXgv=N#u){j*E~!qSm%=>PvM>9oWvgr zBVSkViW1+C&`*lG*`_`l0GHIIzT$jI-s%ReZ`-`xRr8}8L zZrNjx(JAHtRE~ehyGtIoU*MH6kl<0!*FX}8zaw1m3(@{@VvOdHtLC6_$ozfaqd%te z%qZ*PbTS`#65AFs+18$G2m{(@Ilyh$F_DiGD};MSSWCM`q-+ZImG09Dj7!*_VbaD% z)N}cWqb2nY_a>p66LPUe9dB|~J%>3$ILPRTNMtl+M(uK5=q8utI7_|(t*VpWY{6Z= z&8Z0kgojkhGsyZmb|UjwwLQ#Aj!uRv+eQMvI-R2lumsFEvhAk{!>s?Led6c`7>asX z=nF6RV+Le$0*N048G9A_;|(PgopX)dDcP>pVGlcKUZZJF`(eL|=A8EX1$`QROSKk1 z)59sROor%H+|SC^oQ4T$vg*--Ihg6_J9crFd8n1O477Q1T8ryiwCU(A>ye`t2rz?6 zxD87br7RUvDP4(&x-c=FCuTD0mPp@%9#Ve$B|;wO;YZ)qsPj@V>##PMeyZST*DF9Z zVFbs);L^B`*5akE>zqZW3GE6@kl#MYfrupfy_2xC_H(TYIlot&)*TNcC2_*5qgBXe zsq&aSN?|5#SL^F{pknb1i{#x}v=a{na&T{9z`)JKG87;bItkN{7)?LGL7OoKCZS>r zcC50J9`eJJBx&3V#jWG>)+lkRa*?*UgUeRQEyx2yaOvp24Z`xmw2klTh?+(m1EN4*ieEy&%@ddVFc}499^1`&>QoKvhlEUz<1n6C% zX7K^YJ^RhG9IhyEvZ@&Lj7QZ!Qbwa7@b{<6#vUv>4><1V%4@3NFNj#MjAVs2U!+1C zNr^*x2&(za^3~WzV~`S~M)aQ(AF6V!*}vX=yFc}|f$o-?n1t$=%Mr@5=2*;dEgYk% zdBt+{7*@fcFW?Hgp6t}yl&F)RJtQ2?Oj9E*#wir9(f6Q2N+5UO5X>8giT@((zh5Kl z-B@HX-h*(w*nDYdHm&X(N`(-fRMoRLxm@IF5qs(9R~c9Q{f*c7q?25a8Vq}$H2oe` zPK%!3N_wNkIMC~V z!zwcc<~Fn}AwT@Qn8H6NQqh+AC5w1HDYx;9J#j3$xXO=xX3&-?4If=JluV`mN$eVh z=e@rfbT!4n)X0>$(`{@CWrSePBy|A^eIQHaK~9#Inr}@QZLs*vV{ANy(-yvCP2Q_3 zN)mnTsDKKGJ0sF7p-Q0<5$~(QdgCubB;2@!A;T4HM1_hWkP;WE)q!d$#vIe|bsSsk z%T@E@QOKr0SkYXy`BEiU(1ek@fb?MGY~`*sBkVY~A~8JC;h}kX$hVHhQ$Cj<504jC zQ8E3r_GozT0pl9HaCMB(UyNW^7nfbS7#ZfYD8GenM=Gn72u2-uXoN}OG1(nO-Q+Fe z1NbA0FjQi6c}pN(WPO~=_%=kP13Ww}B^5`p0=K;EQ-$rSd3l)jOGJhuUo)=bd0*;YtOfFvuP;U_W`!e$c-O;w^b9%& zX7{Nk&`2m?=ECmWYDAjRmAKgFTsk*bzRv^m9ZkKbcm=W_?8FCMk&57d<6CaGd` zFB6J2PHBcAL3trG6uukC$0V{~6$ofJv5hEbXlGn>uf9WpJS6IR6~vCRhK9WGA6#*h7!E@?^CGTe;pl%F#^(@}v}lH>5`|=mT)%c=#-l!y6}EA#BaZpN%&r z`27OU{f$!9;GkR!s>20^!uAVQsIAg;q5TNLGh zTeMqdGZ}*gl%96%(tc6#s%tO%IVAQ?LEts&AN=#%yVXc7vMsS@!2s&qwBxjXvCCRz zr2}VWj1q%|vhs^)Gz$9Nm2TF?n)kD~Y-)dfgq4 zMpnjTcc9OY_0hL_&d}h5-!d~5dd2X^i_3P5B^rFpmhv=sJ*rTy(RDTz0v}SctXs~9$lUkQiQdz1qy3h&5j~vge>Gimcaqsa zV2CZ>Pnjq_4%{(KTX$!96?{D(B0=KdBbM$NYfp(Q`Z#X|>8Uf`q98rr15p#%$VD>c zRST()P*YN2(j$ayUqvp^B1mvpm67GJi*WOO`{jkEECRdjWf6rE&O{=1U1IOLP&J($ zV-b67Av?s>TIp3mAf2N#Aa=tMB@a0hhArl&cU`jyc0h}5+86ot=7pI@6%^?*R=SXY zax#z-C1HX1eh!M}el3K}I8$}bfIF!|m&%@J_uiv!lmNMU77MvRc;owODgh7EfaIiA zV@QKQ0frEvP`DeXQ4nV(mq^zp`Me$KQwGM;Ysgp+$h1!9)aH!!`9p02q|VYuYB;l# zv^7*U^;%V%${8k={ahI&-Ol@-@jjR^1KdvWFU&T*=se`AUam_CQa{L@&uC+Z_Tsud z;M4ix8}7xSX8W9&8mT8MT(vFuY&AAcgy@Pze${aYo|Hd4&2ot=wG5Ffs%v4hwNSd( zPxE%Asot56O4j+?hWCTV-P~h$+(r~C&W4bZ>Y_6hs0>VHzMg{&KE|x5d}>VyV7Eb_ zoh5id0Cg~k^f47!#Maoo;|cZloH7*7u)3?6Zw1b{0S#I#2nsmgVZ#|`J2+i1~g5L-3+-D81Q3Txz5LMA+?>6=p znpPnPU9%K(k>mvvrIaP!XgUqkDTUjZvUczvZQ9Sc$;URhq@zHcq{MU9j)^i3)Ne&k z?522Y>3@Ya$KMW;%RhSf1$;%+@gPDdW-`(pDuA7w01-DTp<`#{WL(Rl51aMJh50+s zLb-V&<%R{!rbNzf7p9Baboq!jZ#cSBg7%{w{@2+Ccsy!D5?)9dIb zlLG;$R=A`` z`i_w(8dp*=0XJ%0yrTB_REL06={z+9?!@+im;ekCrVqt4N{kB}zX@;M%T~#$?DA&q z$)akev>P6V@e_%S)Vud-G9Ei%@n&FQ@9eqfm}iVK#%M@}-u1aa$9zUJ>tVYvWy6x} zz+L;A@H()(2p91jS`phl3_^KQyxSanrnz*8xa#R^Yv9=fPlIRod2X-#RSG?Kz4Kgz zyX7Bg0$r~MB@{7kmKKh%2gXu==!SSPZ%Q!z#bLRMRz$qA!znP^X^*;dh@E0lDr-hH zx?v-hKF8vpTvF~_Y__h~^X6s7$Hg-UjD;jKQFOZBVVG0y@>A;HIk}gGmpOu6<#PRb z1W<_((swio)7ufS-^Ix}fl8w4KNgz2aMq+)zA<0&iiHRyXM8j4)=jPcc02cM--^JJ8pKTrV zPX5{zKEjcYYfbv3^C>B$t>eluutI5B5Z$;#lf&pYV@1H>$e34-#0WewFbol};Yx@E z%3rPJDC~~!^NNbrZX2Hl@-8XVKZuyF1l7JXm%MoM&Xg^1t!MSq(5Ih0Sk7;bSlzG~ z2|D4K;~U4x#X5T?x*TB|L0LAtFusuDnUuHJn?z7`(~_vjb?8==$JwK$MY&cM$wBG8 z7I!d-Y;2GW9G5>!DQqo|IBZ&CHFMVA7{)(>&sGTaP@Ia@04ru3G$mdE_ z31jQRO@=zUvoZeI=C}C_nE$ddXT*v-qh()LajNMmpPW|J!1#1=A40ugB*ks0uWYGc zmE~!?7+bMn1%bormNFULce=Hd1aY+7Xee-G;;}y9igvFstXKy`LQ;);R%%)A%L4xP zHS;xV@mZIB!Uzpgg+=s@g8^+|YMaI(uut|GaAR=h;(KOXpI);4iJ$_-Po9*v>=Z|R zLkXjOuUH+UJJw{Z?vFG202@sg%HS_U9Xja{mJwo5Ink8|ZEL1*s6>zFp*?n7B%Byn zukldGsX+5f{bg&tb-_DUBJT)`uO%JVZ@kNN2|le&C>={RW~ja5@1(tdQeCFU@`G2A zC%lQJRNTs$!$*=y?cqg^7#}xiKY?z*8S)G)l_^qEfbLSIZ3|{Ci3Et*H4bQ@$VLJg zUm2C_XwfGFO!=B^cUvb10-pMvZ~6j+D@k%zdM@ReMX4`kRYyr}`|lQv?6lH?k(b-? z>W}7cR&hnA3(dA-jc|gHqlc5lIYmcg!n@nFoUe}!oY01L)%P;HXF^Aw9){-ET_#E* z*VrBvYTegt`b=JJz(%!X$s1V84~x6*l`htLQHJr;lm-tujr$%1yw4J2IQ2!Jvr>8{ ziUoS6OTQxjvcCT?5lxXlhHRnic~qBoLd%s?5H0s|x!L5?UKezeUX<~O(Ki05j}MxD zzaDiU(tGZ*4$sfBLNC|ie-;XrZd5Qkw;dX^YqKOH3CUpNx~UAH2>SEy+{z*X8U`Y; zCqoEbI3`i@Y8htXfZB#`V>WYRW1`4VFI+fx8Sk_<<3q$`{8&a}a5h@fnnzsSwZ26!?p*197`FjuNlz zYKdaw!%4ze++0jz*+>NkE2kyl%iwn%b@0ja-Lbn7^0BB zCzDf{;>(P`3Jr9vHqYBaG?inJ&SI~h&-Z*jl!uO7byX1yBHg3|73wouA&$8qd$MUf4F!w_) z+wj^1wA<{BB!Xg=zPmxP5oS7e0&Kh6@80}FmcuT@kXE%7rv0gJb@Ivs$lM}2_X!}c zOiE~9TBE*nnCOhx_^EO+6PhYe87#2&#V?KVB$A9VayGYdLz7g|D&7iS0M^uZQ*vu2 z9oxs`*JJI<}6~g|=O`Y1{#ON${__&rfxB!~Jv9*L%~r>sMZgDYRD`Z;f{&vp!tO!1HI z7kqP`P^Kcs;gh);QgDQx*MWQ@X)sQW0tAtRJnYr?km(W+izBKmxI`gJZ){R7lV$Ji zR4#^~JGLg6Di9we^QGxIF~Q}LZvrKR0F8S{31zAI3w}mo7-WV6X2=Z7^ni2?yr~&5 zFv^B4NU#UQ_g%I?!UglziF6otxZfQFq}IB>c$4PgnEiK?=FCTb-l_uX#IuvCR;Kg9 zta4!3(&g)H<`!MuTXmiRjj9EU)j#R*FMu(R5j&(2>OWI!aAaQxZx)()Ofsy#fz$}J z-)B<5>TZv9|L|`5_r5 z!j<$(&{sD9?7?;xf1vBg!M=P*z1pP?UXT_%2FY&N^9?@n-2o5-UhH1R0d!ftBn8v` zB)1=kFZBOweSTz}`5**e-JC$gL3Vz*+WG08>t&y7b72b0>5J{J5g zj?AbBD2n?&U}Cfh*7j;$nSNu`kP!bS^Qw?ErUi|wZ@v&KF^Dsnz})uYtIZgq0W1Z* z-0x)HzOoTlSr^7!cKW4E*S>m?7zty zsIxLLHtt^5bxwjZxl=OO^zP0yLN*D_WwN@bL>Wl})Axoce>L}RA92Ex>%`FWj+)v> zg>u8?4|r#~ReCIU;n&j!E|*byH@*8hZI25L^VOII>d#0BS=WNDNi$xH6NP#=W?%NG zUP?OA9gTNmMgkp1LJi2|tu+H#zw;kB=jnRP*3eCa`JfbwBd+%{U}^wd z!u6-N>8TculEWDL`|7K-^yYBOxx|Gf_zEdxK73Pqt#l&U+qd2*4Jf(}*{*C}5m{?& z3i zCG6X+_%|!rT%0#Pf(vPVOE)$H{9U43FpXS2r=_nbip+vH93t`riZnE?{?o>rr8PSj z$b=TZ|K*H}9Rt2Gy5=@y9-}`cR8KVH!Z+vhD-m>(KDocU#>JV{+fS|)yP~AH99!iNU#4WGT^>In+W%G^U6+Oryp~u^>ps6=kVs6Zn5}mPXfh)c&{Koos zg@uuH!#0N+x|H2_$iQ^(W4@bu{`2L>&tK-;53IxgL)Tdb#lZwj8h3Yhcb8zn-QC@S zySoGpu)qSrAy^3RzPP))2M-nq&Mj}x1*g(*?B{%% zNQj+j-+CD04<6#q!yH)o)yI$yD}C`u;r_!*xHDvxPoJ0PW5#` zHL+HTb&;5Bd06zmu3@Y;S#!Fk1&+xnD<5@j;|5%q-pFW;HTax>?~rU4!@a^a5>F8D zmqZ-gj{Ck6gRE0DCLe}h`c>kAt{S{V^n@SSB}eMuaLa6c6Ew@9{d4fO7_ea(;|Gp7 zy3?-l#~t*r;}VXDX!>rz?b>F$gS_4LMMkg20j}DS@EFnYmCbv+`ECLpUr5uO-I}zi zmeOdpAy40jwdLdG?52N?KaI6+lIHB6PB)D`RI|$#M-eG7yHH9_2cfqu!q|6Q_Hj0WoxBfd~LiDAm zGCpWXlTaECnY^}+quPrwt-39GQgrcHKr8tum&4d%9PX;A+KW}!*dX%EdImz1m_#{Y z%u-?<(QL>Wie5%>i09J6C>@oP_G?W$44VntIYN-KmU}D1oKfm^U9HF&2ftAz9(gWf z-VD%0_g;i6+3p-$k^st*o(EPHRCP)9iSTp}zwE~K1#}Q&dP7Wh1azf5nqoyJro;f* z2O@{ujN1Ph+A&8kFhEx<-SF2Evbq_RYm@LfV;cSD{a{kD67#PU+4KijpoOoBC@)%3 zo>beMbi8n1^6VbwqzX0^VI#Hy`V^mg-$LR^pqF2sTc{xHN+-=UfMWWFeG#FQX3{P% zcD852isf@{(kJ zhZ#fZ9`I)nu7b?N6s8e2SxG~^&*%*Xe|VM#iA<}HLr2?WR#=!6Gs!R@fbEo*;tSQw$k$+q^8Kl0_P7azp(1Lm%PiifgV3TAuTv4}W!`MFVonp= zU{op4B-%9L_}ql&u1;Kqio9;HPGXn?d~E-0;eog2s{bJ|hEBnn|s8BYa+!c4zKy(2{DLw0>ylB3uI?Zls4_>2(u8Rrz9u$L?wrKGQCI zhR*;*7t7K?_Dc>s@iEfwVPy%cB%dL~S`u-r`BW}__iTSNUz(n27dbu{o8vYaRn|hn z8ey8Rlis_6ePJ_c79$JtyNE zV^mkcQMB%%0NOaZ2^^opHXa?ICW@6mjZnbB$VqPI1jitS7E2D#`};1wBOo(Tt7W{> zeNX*TBeaQwTcfK&(uD6^4yq2yya58ULNhXC$nylJYWV>ec*6p3DS7U=q2fuHzfQqw z$x;$}g5@aPiX&j`8%QV{=8(z;ye>(hjAFyX~%gI2kh0WP`gDz@dzn z2>a4=GlNV+_Hj_N&OM^0{jO2N_F421#mK`?YdR6d=OjD) zyL?Dou?33pjHjVdsIrzh^HtdK6YFjsyn|hgCP`xC+U9fS{eS zL32wPoV1M|bLDE7k)lKWY5agn*?sRbgvVsC&~HK_n^~k#@!FHfFyp z0C(2tRCx}7;F0huhn$7 zoW$lybO<4m!rSc0SXi{Uhw0&#FegkVU>)Jr$^_4IlLIOXTg2hrY*HV0~^jgk6!7;ipvyRE& zE)=$1RMz+JVRUe(zxvL7U6mGk3se#&o$oXVcR7xTL?nf4hc^{|5emE!>Z$e&ha~7k z+KNC|5whk8au;6y{lwiZZNWG#(SpEb7n$Bq^jk-$?tw+??~u4%pgz#G$ftd10*7j| z`Z?0)MzT-Qjy3ZBn8r}3Mf;$)dEH>e+5TQ0sAP)Mm@?>=ck&GMR#YJ}I&P9N;C352$O!%azrayIHl`i`MOnqq56s&)n zK?kX*gerWf&tjTnlyvI@T)$ z3V!(E;IwASHSNLSsNnOVYgRQXSR#Pmy@CW<*}S%xD;XEfoJ09VVkn2Q>xvOcMF*`Q z&9g>~LE9+xS^CNEJbPR8?UNb=2c<4)Kg>x_xWzC6;xG)lUu4yEA#i9$OI9`K63PwX z^s86OIZEm(5~G|)Q*qGBENXnZKZmInm{F!AOE}v_7LRiu8m3%t`eRtNqL!1B)T@#v zDTd>*YtW!KW=cri7BM{>s{bUU*#RmC^FoGyKh*dH+>CR8B7Xx=P*{0mS?0(>1xf4Z zqGFQ~tSlVj*|ZGB!=4IINkdZyN}uT-D{(2BO5PUvIHM*nt>bNvABHzg)U$jsme9E; zW8OwzR`g+N;~2G+5VC1O`vYmQaQoXfG*J?(5iWMM8xK0K~ea>EGMUcBUqF#)^06Fq{n0^&*z1jQkNqmsc+}=myvg69<<8%{ zx$~`CC(G!_kC+C}T544ttSiAiM;`@9xE*25$P-u$^woHd0eizDf81Ky7puI#i_qTI zlg#IuIa)iW<$9Qa>(_doH!tsRg@pz??F8OmZeLX@mh7+;|A8$%%$rGp2h{a6B=4w6 z9E3sy_KnjGiTAc&UnH`(KXG~AlwSO;PtD4>_k(NC%}-ng2b4NCS*^>0RK}MkAq5~Y zS5+!X$X{p4U%%Hpeuu*$3%mgb2!bIl-EIJ^+0T5Zp5U}a_7b7GGmBK3O+~i01F(M8 zy`6cTf3XQCU$~avrfp{N@9oeWJs>k*JdgCW`5Iba{m#rzt7dN*Hk$da@Q@It`^DrB zHQ6-+$B3D?=k@-l=%IJS$XY18cyUS@nXQ?RGkVW1ad+<}F6vMB2QXo{&K|gnJc<$` zzoYoTA2iP|^3aPe3O1~GHpxqOf@Ol@XK+C3GRBA?XLaEovX*6JG1iOlT}$v;%Qv0T zAEGe+ww6fbZ~m!1-mkgakGZ-%Pm5guivIe~F_Y2j#itoz=FL_A)KWR^M}E9K=+nKs z#m$YLr(Xu6}J$lUD8x;Y3E(NP4H(1IMqaa&P7UO14@jo z`=*1md>nBv=BoTh|m|1 z-shswN^Pa63-;p&$e+k;tq}Ru3?mp*qL+%kk{JhWq~Sm=HFU5IG2|4IveXaGT1Jg% z-|>?v_>Kx}Z{06X9YvRjYyUP$%c+iz4IgB{J#kEym+6RL+eovD^OSQWJnWK33^iS# z;R9r5*r&z2jToIFvJEBLP#2?7A4170+?D=;q~~@MuK7?M*eer?e+wYbOg0JWHa{6m zc)iK!*+|+V^l2Opb9;0#(hFU!LcbPzv)u^tmjc$(o`o-% zO?~^? zRJ{_f*NI*GMPqd}Cxxy3n98O1L?aDbjQ>kgUj`{kI!mM=(Y!Ona1vJ1JV4AIA4=14 zZ@DSuqQJyt*;`Rm>UA~$RUh2MJ$gFz8$^M=rSPDxDM^jA8+s9Vn3TriBTYN(H3Vzk z@C>e?_YdCrZjmk==F<^GYi1a2ih8sl|D2K(_=A;{y^XJ4(yTc|!)E*KV*8ITl~Sm_ zt5F?he!e(eTLS@y^*&fidfI;7+~y7ZLh$M#_RoXAa~Jin?3nIz%lQ2rIHEJy1ozyN zen=9F)IL&aD{S|pm%3{PGoGLg?P|^|Kzx`RmIZ`!L7gjh;bTqKY-J}nS6eIP5OPv} zDW%1Eez z*j)~h z*ZiEMhXHI8)@=Q!Y7~{RsAteN9jMNDDR??*P>n)CNU$k#Ju>&)Fq`xrE42QN zHX1m@MprnmOxx^o9YNH>7A>&IwHZI_L9evv{M%^gn|I#dzQ54q51^FYbP{%Uy(=yi zYyVQ31|{cK4~J77T(s12D4mzm3UD^_y*5;%qV$9S(|{s5vXp9Ka>*L@(b3Z+C~^P3 zi3+DJZruyGM+oj<53RfohAZMo^e0?#ot7wX>5B=t zwMQ+v(e5ri976SouM_qhvbtcFD(6|&c0px(@*5bu80;8tKFw7?*LKNe1FUWI47bS0 z4J-;#^KL^~nl$lmx^;HmKGms?!bli%YdBq^gxL(M8vnG0OHoGr;gEf{IbikqqPMm| zN7g2rV)`UC?=MYufXPDt%_dA)BGxDLx0Q5B3h7W%>n_Jr zXpcS!k1#d3!$AA6`OBy3OLzBb1ar2R2>KHLFh@0wYx8dORN_hD-y+M35LTVfR~GS6 zN`B(Waa^JTCgh~gk*v$52p;j8&G^0{mK*8SdwhM!#4f>GV0H zVaa&$J$=Xx{jJ}=RlgMEP@@fE^&-GWyHFcFCN>DsQa*Lo7X{a1HK6%q<4yzg>tRGz zw4Pq&2lg)xVoXRTz?TU?9XM|6>o8*pHH#p`g9>UxToq#IDBUg5#xLXq3?uBSIQ5r zl5y(G?%Dk)guxdX7Vgl7W?twc1-Q4xvl2fo9g9g&It}!=c44wxb5?4ie+Cqp7cxW` z!g2|zP%IGRC6zb=`^XTAMNgdvMn$P*c6}MLmZv+ijzmst-Z%P4gngzd?nXjU0LM=jIMgbY{RZg&NDAths&FY8?ysIq7!M8esW zJI_D2Agc@SbR-hG9UVhAiuA9JL^ghYGUO~)lU^k|QKsYUAUU~#EM^P*CfB#~ef~mt zm@*&Y8G(!tW@V&^6P2e3wUO;+!o8O?Jg&GVX>-l>nPEb=oo5%F93WC@o0u8Zj3RGQhl@Gttyc!|6N+XH?I-3-{;6UM~t#Wcxu6(+{WZtrq z!cTHQ*>eC;{49Z8%4Vg`a4Fxh#RRN{O4gue_(X3zEy1VSzUCG-Fi-jQ(to*w21(o` zqL+|>Ab1y@qX!1UwQm6yLS6uO+A1?C61^KriDK8zWMwxaBf%CD0;ZJN2eM#q-X5Ap zsn40zg7#}@&3+ui<^?k(cV?Yk+TdQrzh0cn1wV6ariNsN{dhOP!ztu;rDyj8k121wZ_yyJT*tL?W2XJNH z9ZnH!M)T{#q_CJKTQig0GFgci(g__rTIb?Nn`9?Y&K0q7P|2z3H}YXH>VjyCt-~;3 zvLx2=L?cKL(|{;9@oLjsBD{Ll-wSJ)wGOj>tTliTffFYZ0K#BowVk1udAO0g15og; z?*vz!E0Jh=1%z-FXr$+FpQVKJ=0IZ~SG8~W4}Z6fQ;Pvs+08v)dGj@UjZH$R<@vLB z=H5{RP^m@BIYAH!0g=R2Glfy=c7wiT{_#{F=I?d`$%&K&4$?=vc89k)F+h>sASQrI;F%-Lo1+GgNK<;R zuo=5fwC>R54x`eLg-nd>WZ5>{sSt05s330dA7%!^N2 zBKU1^sK~6FO|fcCAvz=DsRYaBS(9}2k1Cq)@#7_AE){qw0+cw)bx}|?RhZ3g^rv_u zVmPx8%&S50T==<;jNvKNQc0{5KrOQ)PJ@lR1xeyirzD}us0yc1?(toVF&Dt-AXh3Y zYX1)^1}7_~{dw3S7w9mKKwORn>7cp(Sc8kDbl0YfpUcMqAk0V91j?BX?$LC}IzL14 zWd++>Z#uDt7b4TF!YP4rzeqgUnj@+g!CZ@V5ChmUYW*7c9>(kf_xqjH z>4#a+a8F9|o%jF!Rs%$!eoDL+h{NiRY^N#I5bzP@fk~{+p`0}Zs}b4c|Z#``(U$vS{+Z) z{ec-7XEeUCok-|>`7hnR&bc#YwH>;madc_n*Hq;$T6oAp8Kp`3oCJuwy>b{SK1?A) zRSVxDzk_h=bdcPmJpc0C9$qGbD#qyh75Q8L``h+=03ZvblEldzT|Fr6Q-Jr)P|4&W zajIJ%Rx5T7hZ7i(D@6HKw(Vh8t5aS!6&kzZxj~iExNh?$HK6Mu^IFyE z`dn^>e8Rxm*xp_cl%7skry~leGJ!8AnW4qi?O0pvei`o}fsAgtfG~d+RxG6^lz{3v zAMZYA?%rb_>0pH!FSFXy>#+q#CYKGyH|f-fyWI$al@5Zsu_AitePgvj*cDV~SLpUp zgm*grHpF!KG%sDn6eOR{P{`1A2XNQ;(nGTl3yuvlKOsr`z}V>z9;m5Kgk&mwY7T*n z$Sv_f0o;iAnO2NnhBwZ}>l>-FW;2zPLr~^I=ojEEYsBLSz)cS9pkSUml+S#thyVbi_q6u+VAVGZu0?8Xfgr*~Ndp(&&)sJe59F z5@Hg0*$Ad4aSe<(81KtPLsYT9&as!seNIe~J$Ep@(>A4z+oSSZtTfSPjMhgWQI~3M z-lmP+k#qVipudK8YBu%$DE2Pf>pCDn4Mq>Tk0UE4HDk-K7t8BTXc$2W)Xu*B#5%XI z2n6!|yE+01+-C|~ZGVWCTV+Du4^7{qXkOYQ^Z(qZvpGD>S?FQ*{WI$W+1or)^#@Ov zw$60o5Siae3S?;^a2&SR&C1c5xiA{k+wpzB4Wa#pGIlX^=~uIcCg5fY1j(%w7qDr>;t&X=reh(!&(UVC}7Ch+z|=!_!iBKkL!lq>GB6= zl8Jm zTP;|k;-JAMZVtqZ z(V=iWNDdalJRM^Yvtuo^U+8wOna=Nzg%1+96TFGi%|yv8Zm|R}f8XbuPIwE4$)~K9 zAk4fV^&(*Sm@%fNWaxqc&PmZr8k_BRV5#)RHss!xlo35InkDN4T>79Repe@Sebl;) zNRnBENVc50C+0WJc)#0&#fqi9~((JFl)9=C_ddYq7WC54Wn8O>2otGcsh@^o{Ek{n7ZqZ~m&F zWhHfR1NtA8$E#Ku#}OX3Oo`xLWD~P9?<<*k;MAY(8JrkpL=-UV@lNp~KkL4|1Z+22 z4Sp?q|C`UiSCE55X;uxj+|;0~Le(Dp8+_n;yL=Ds;2%mXt+WZj8t=%hkv1-%lO@M{r;LG5D9fPzo zWJybO>K`_JC&17=LB^d^)@Ku-_=uaqP~6zK}-hbH+leVughic zZ|5*K{bDcu?_fj!FQ|F0ez}K_*k0Q~t1)>VO^bz%Px9gP(h*r_)_xOuNg#mUz5CP0 zad`WwYI{xKw_HisGRD#5KYGTXx4eXlc@>5L&=|6yT=;$lk#>6letAUzz+xuky9v+6p zRHBCSaz94FQ?tu+v!~!5n#w~Jrr%LYyRM?D|6Holoo7a3%Hf3q7$4Q;s(tI(E{(eL zhaC-tpZ;|41#WpwCUYi=KR(&d6yN4$oYo!TQ6%wN{p2uo*n5#_rHLij4?2f$T+NimS4SORC<1I zm04W#x>nPKe#MB_sG~{wTyVLI9@5IYo;4YM(P@SLDKx{+cVnN-=c~|$_yV3I@Z=M> z*e;?fh0E|N!s#Hp7ZXH zQ8VeiO0=p3Ay=+r^=$2#SyX1r&&u0>mAx%ZNCcotI^5Dd$cxsPMzMNW8oRhn&;xg% zAH9|=wh3^FyqGTD+!A)sUqI+HeX5~x>jK2xz?RfhBsiO-Y*m%u8~TvBMF&jpgp^h} zD3@s~c=?rbK?vcSl%&5ZQVHH#XrN{zpRO}wr!(AGYHZnY61G>-V{!C$@GRsc8O^Ax zfN;^3aGS(W>iLL@g?#bTTrH03=_H>%us!d8vH;%f;mK4+S)9nCZQ@|b2yL@qRZ6bG zpj07>e-(N4;<31+GJZ0kvq5}6rVA1d0ry-XY znbT|gSnIN4C%U`YpUbO@`UV0n;Dt3(A8l9scI((zYw;&e&@F!mH>(oaYD(6J!23`~ z{>g|a8dlV zl%qc=NLPXdWDkp54h!c-MDd+*Hl|Wp-%iFW2qHAKuC%l-f9h4qRQ4sXPzkceVl86( z^uLq(|3v$Gbv1k53N;;m<|%!=C(>}-N#WY)^oZ&ehGILnr~VXW`T$#*a7%_9y&g%y z&Fg`{4%IWr>b(1K6ydDD8ROg+W`52dYtmL)P83z}Gq(g&1{hzji@>4vZuXhB>tjdK z$UTBdsF`Cse}0_=#>+ej%S_t;GQqXrIXi45y^JJba9|GFPU;h()QkeJqIDu!0wyLw z-EN^njOtAD9g3QDi#f;>GCz;f;(w|GfAtKk@UR&~K!JQej2K61a4LbAha}-0U~$zo z6Q%r=&N|X-hkA%nm_%P`t>V%&G2x3XY(x6{bHFr zF`-ebfaqGJHSDya+&GwDGwm#MONs*zUd@=WW_X)7h{D1;hXQymEeE{u40Pc%nndHy z#HEKwM4Od%*X`CzGLw9;7s4P$g8JQN*O{zC&hJ!aowiOA`<$OPp|+Fli#I+OVD8(U zqfla)0QZX=;C!G+vST6{fvBOJ6$oL}`WDHUP+W|0*G)8X3M1!_?ourYDUHR}*rxqj z?|j;GcC8or2r=>7F{2RXf^Nu*bDJ~%kmD|OtpMA}uVe|!(n?KsOoa%qO`hslhG7pFSuYqxI7_<1Fh{WCYya%yeR zi&NT}!g)#;_})iZdnFdGi#vrXN)L&!LIR<@W1x=3l9VHa-8CSQF{V5IQ8D@JaUf3% zWS%m}Dq6GAXtrdesqg@H=c<<7@&+23u7onoGaCRdK|;Zxa-N@$P5B}s!#2Sdle*YYw=kH4I*XEIfm<* zJ#9feGQ}=s0lp(oc3_*vstVe-Ex3#X9V3h{gJ_{LO48cI68p$<->x3N?D)b^P{Rh# z9sD`MaqmW8f&aGJCXH<-9C1vk%~Xw5$%k;~RuafX zr9l`$*~wW)Tf-j9^p_nk#XKScy&L(CouOFg9uhoz4!AHT|YCnXzoMS18$lm-)c#;j>{nMO#LHDINL#%qkHc#36+KV?N= zTdRDj=kV@^_4XUi&I+S=P-&y5g!XTkgg$eh|n&EfpUeY*oU zuED67l`>qi@dUlqpRu~<~pB9lG@?YNP*Owyu>4kvalo6%xU z%*AHIMjJ_;!*yfz#`G*@Vh)(02TH{3s;$CgbdUz-gM#&_6u-=VyGh>z-1l76N9XGX zzPMF>y$7M*2;1}Xb;li;YQ{D1WFcb`(19ZH=Fr&zSjOFjEikl<3&rEG)Vfwe;NGJR zWSV5nDuSq|akVAEg%dJ)5};L32dt~n_SO<<8W*(FPLSe5W$-MmnAb+|a2i)*xDWr5 znPsgLK!7?9yYMbPM8+03L4f_Ub9&qV0OByYq14aQ^&!$Ne`w?xh#q#9wzWC--Pb!) zq)Ln(fS7|^gdp=?UbBcE1@CxrEt^e^>B~7fQR;M_I_u?`UCPqB zip3_*#ImV)5WR_={Yf8By>rt|YSVSSa>D0G=Eee*-VR>9D_;zZPOa0dy%ct1%daTD z&vQZ$Hm5c-q4yB=&75BjhPl4G0xR?7(UJ#=&EcRCrx59lPitxAb=5X{L#Q@Ql^!S( zW)xvj56=Di1Sx3{vBF7#sX4AcyC79r8l+R6z?Lf4YAZ95*F0wGrafUTI2)jBW?B_e z6?~OQ<`S-w?&O?1L$R$m0$F8DC9k6nOY%Ky%OF;=C#y}Bt*qoQf{$naPWwEXqmv#y zP2JM?lCgt_@dY>CyT5#MLtS9_Sgu(U3GKy99Xz}7F>H|K3KWqWu_`*0 zRvwQ7DOOS=(wxAS+@Y}46f|KL=04Fna}wWgHzsJt^!F8sn(ZszlLF&Pu@W_^3y))$ z%^YZthZGu-3Vaw3o||!N%~*kG8DpDX>g0HytXrpfQ&~m67RU2d-uL)4<&Zl#tNbOlLwP9Y&wPlk_hpC@?gcjVSQcsMt{45!cM2h>-;2WI!W! zbyShqK~77L!UKkO>r}=QgDFRDo?1*c`hiHqpUhQOGf{%wAGc-K(;xiNn`I-Fa3YEx zC5srHPmjlbNJ!R(+nR|U@x(q5KDFD$;Hh=ztozT|3wh>zZLqf#tK&zv=_|czkJKX2 zdOeAuI_SdTnO0lOQIj1_qISIx-ZhdP@uD%Xz!RMI^^8lknt2y2x$J%%V)G=XzM(JL zv}B1)rs+gv!z=@rmolMs#yMCA&av%TZJ8CY{=wkM3cK-?KlD~(!_LA4^cEwe zPvtV`H3>zuUG?%X9cj_X*AdlR@&tA|M4Dm&%vXoP6qo}p_yUd2#Y#V{GE_WH6KMYhip@YGC;mQw%U~fCX3J!8JkY~AOA$<->iY`6xpeResqY3F;ESR~wpzSi@OsQmoHKtUoh6R`& zOdz65M!_?#O?CWHPa(+eKO6sZ6gAmd@J{WyX6G#+550ikNSzjiiEAtZIWZoA#KIL1 zD(mIj62>}NMoy8(A%7W)&IL)Le()}R#QbOEMC5%yRGjJM<@?*^vuRAJ9iiSQqBZz& z^x&(Ati2e0-C%vorF49l>;vdzi^l}#&%*_K|K96fFYTHRHl^CGAlVGt=H zImT8ts9Y6%Esyv(b@8^qaY~gTw8}D?LqXb*7nxcQ!P?{W?B8eO(8MLl^9kSMtiO{N}U}?_8m%-omA)|5#ql3NaW3?g6X)#yl1sY-F>8g zPes#=0#)EMe8epu60&}lq;NU!SDoO2jeV~;^ZDURPltP$Rg1ftWN{VzS>!>rAy0r; zr%1sa&-Yg>kBT3NCZU=PsClfRE`LRAgV)I9XZ6r$F3U0({?@hqEaLElzxh@fQsF_b zr<&*xL@0|s8-wv1HXo_aDP7a^AqE95_#w*}>=&7eMLB;{r9OiBi+k`r>;k1xq7y8Q z+!n_KL-zeN3u?0D5MRe>w#j~F1w)K@H}v3CvZ{XrHl$n)(-KV=dnoP(&Q%^CnaeMr zQu3ov^Rl3`WY{fLnMaWo{X*C#%?F6m)^hCH(dY4=m+=ofm;0cv-?~mC_wT=$NxqTW zJWOJdh9|ND5zgH8=0BATbkwXwLprIZ% z1l)1av*_B!U zg{{qv+}zLiqmZ|w4<_w*j}clkZEyT;L(kt-z7J_@QBQ4jB}T3$|EjKK=4uym3m8Al zG|%vmZqReRZzl_<$vX$0SWKTgk3^=Y1dF{5j#6Zub z{9U8jPyof<66G4R1KDcdi5SBkB$Z{(Q8&^ZCTu1`;1`mWgZb{&OXz?t`ni(Bh?AS0 zc9b}31FyuV8$487<3tHJR%<%kOm2+~20O@<<#Rv7!*FK4di`)Q9_~c^&X2`0XfU z>!sG%z`pGhp+D(WxvYpnu)rksc#JiaJQe|K6lcr8%#9z`-Fqce$e*KN70mA6xX_SWBJ_S8Kr9fPppG?-|@> zM|H`6?H8+jWYSac+cryzhD4D6n~8c>*dTMGjhOwl@C}0!1OpPV zNPmI-@dOnP@@&E7?YGMqmOK-2iuZU$zR$g;$6IicT*OfiNRM5>>k6(^eH=k3I*G=Y zF3=&zl)^zl3qs>U7y31>ep}L;IS_s}H@GX4<5mb-cQo~&q*jmb0g=8TP>E>c!7ryG zI^|(hjxk#k6WR*i%>=ks6M2J__4Oe0$syoFULqEzwCow`x#Q`j9cgOcV8 zdL+#Tn*?CJo_~(JAEWHR9P4>mMS4ScMx8Sp(`!7JrsC*ddGn2Dn@)iPOObOY`W?4j zVJ;u<#E(%X-7i~3bn|?+{KTG$#ZIB)aKylBUKcoD^lx-<8QuZc?=iH#GQ2m7WJ!Ag zQlnS0-(J-oxt%d~Z3gI7OgUHNk0Y|(*YP2eBe8u)aD3YB&9XFT*)K$-AwsA4A=mg9 zEeRD;tBh|8s;@X@cW^9g!OS$6e5Ln+qUh+XUPr1u>*v_91NX)i&!0daY!rD}s6y@^ z%;;#Kjq14v>$}gna${k@E{LJug5r7f#w_N2ye^_d6`rbU> z*WIX;Y0lpck?ukP0z5Cqm}sYBx7=d;+I=a7WR#T|CepatFMN9{mmi|V#5pm3#rz>_ z0*fH47P>LJ2w)r}p@y1F+6k@H&#yY>yblAWh zk!Lxo_FUi$GYIb!PG3Em3yxZ1@7r6Co*%UCHU5sZY#wvLC`&jRI2=qI-e{x67lpd} zG7+mN!f=yYrQycg0(LsG7gF`FeUyv`h*=cvdkSK6)vFs@y+}QNu#oNwvy?NR4y%g9 z!(7WWl{CMSJQ3V4hrBK;uUQYFoDW!Y9Tc@m|A^8b0eL#iHGKA+Kt<}QenUtmXGOn) zyfcHkJ{?XwgEq9&hP|BjoAOuyEGE7#wQ7j`!60fdA$SG1l!$0_ z(Zw_ffF-6gg)SCg-p0guf?C(ie_Hs~?fC6qr=#hmdq#)v-PT;N{f*`Og6Yk{UTDcD zi0+0X7yD*vH;7vJK6Wu;1kQb_6}iXyO4R&`6QZS4sLFH8Svcsf_n%yoWN{jv^j^{L1-=yeK$$o_hQaw?S5@Yci9n@U)Yu@8rH4*j^i*Xp9=o&A#)AcSZ)7jQxjvE7ID?u=5x_p2dQs-Y#Ejd zUq=6(M*Zewq*@M2%=AMvFfixS9a*vGjU)4@K^e0}9+sh81<$ z`Ye!PJ;<&US-Em+T|__=x2?IQ@%IeR zL`|1u(n>pfv{bQ&2r+@j!qx^@lMFGK!ys(w{@vN^9v!FZ{UVNR6zKi$)&KNWs;ZKH zyKX~eE^gqovu`DAoT&KRX#Dbb4dex?Y+eAI>pw<-xudh-G@v%5FmgTV9h6Y=Kq zj(K+B$S|55l)GF2JU8prYTesvKW4Lt1ot|5P)j1ImuA9iUvl~3ofbuYXSLW@Cu z%L?6^Ez2hAWLl6agBnnNLmoeKw?^)$v}uRX>X ziF17dc7hrUM*f^QlTbwKM(M!)R(Eyn8>g5(oCRkh5~(Jj0F~Ilj#NxR4>uU1NyG92 zaLutavWU>JE8^3D;w-hGRC8jXEw22O8f!S6GO=Lt&`RYq zgc96Bu0@%Z#(|dgCu>IIuy`J-HC#TWs~J*L;-IW+7T4XAbzSsHiO=6Bf4L3t6zYp+ z7L9YMbi0dWMzQDDV)8`|eML)Ciu;{lDsBljVc| zwUr`Nokh~+)PH|%@PU;#qMe~eFLPVv_d1KhsHG}l=IQ-ni@y8C??2QGd^DUYZmJm zG`P|wYhnREpL%8i%3sKSNhja~9*OqMsuE2M!kyI2lFgw|_8jHB*G&=$1|cbQ5B^ou zx=7RaLuHK>6iV%%NjIK0h;7kIxy<(3^vEeGLdZ6qd#veQF_ct9=3w&Cu2z{e>ewNa zt3is3TrD3p*f$>CnCfk4m*4i4Rlk>i11e#dyOsZ_mj1wupPCA5rE!WLg!% ze&P-@YQB+(QWc$G+6x!D4Z}eQ57#5zE8a4}bFqJWTu2ZhgxI%G20VawY{#@HjOwpv z7H)UVgr%t3o6}a=0PwV3jHAor#bBDzu-1(<7`6{N6#RBpN;hnQB=&7A_Y`U;mI^*` zPxQtuc3|xh9J7ZE*Db-Mo)4Ys?JmE~TZ-0f>BLYhU3bcA=nc!!R|laG_vxYEI3;0V zL8l9IYVU?@K=iP-xq}xKd9GTW_jId*8W!{Kdy|)V+{~Vw+5cE^Gl8mIHi5dx<_?7> z%;NE;ioB2xqt>Z4*SG{Kl&dF1J%cM=Ffi^gi2<`_JMVAx!!>eUxwBuKK<%U#nD8=mtAWDncoR)FJBmoYI)RPx0=3^$P zwpwdwHi7#J1YZO9sbty)-Nr0(ihrFPMiFS8qN0wvGpH84?6pNDRH2l`)H(@xAv>L( zCs3KIv_zr*J_v8xn>-g4ifpq?4kei1qz?==l`PI+Zy!g!7ERAlr~%Cu$+SXIDB7X+ z8hB0zz=BE*?V5t_lZboLbcA8`fcGWsZkR%ly<{z}DzRn$b`LX9m(WAdd+8YI6s}3> zl?EEwU2nAmZb) zQGBF_{;?&ok9J`{_*HqT;wqdl4>id-p66MdHp`&}WUTA&I*vcnB!&s%S^voQuz*NOLF?!&!E zXM!&E#r9{0T>DNFs2Glm6SiFxpQ;4Kih{EV=yYsOVjXHQYNBfjc0t##)VlgLiJ0)G z`9&1%c}a|OD!?K;X>_O=uSe^F)fnagDX6{L&%VvL-Au>O9mvBDSCqKWx~^k?&fvcd z$tpjK$vxwFz52h%I_KcX!miuLo|qHc#>BR5+Y{Ti?POxxp4fIWv7O0GJon_S_ul%x zx__m*tNWa;lhaQ>d#|;BYd&IKes4^FGi!h1zMn|l%W=-@jHU6ruFJaEoupt>deIoa z9*6@JNYG!$kgz{w&1mDcwBq8<7dAbe;+!A%+UVU5dUkZ)3eocT1S?dX(%xO;c03`J zR7rwq3wGndTwG?AWd%9WKo~icoFL%8s&D|J~*~jD> zrMExer)$`XP6njN5ttpRoaC@2-E^Ds6OUYv$=5XH8_T{IG;-)WDWwZp=>BvuGXmDY z!i={C6E+EL7wwL6$)a!K>R}y{kO93lq_t2fC}6c`ad%%fTqX77m$lmYF4ZZ@>)XhF#IkypJvW}LrOd|Xp zjNa<1+oElLa(9vl*0W7l>!)EWdSK^lo&Nn-oPfRX2bjs4WqEB~b~iQxPYY{{gBf`l za6wWL$6uJHgdRQyD#bbhbh|L1N{X8rM4P#{C2eLrA-MW$16I}ZA=_YESfnVO*z@Ge z+VK9^5b~TDA@@Z6@T);3*p!)e3Dk|cy8P=Sg)@+omLZ!9N0+fu8FF9xbVlGY2fC_+ z3c>u1$=Y;|Xt=wK8U%S7njExis5E!uE>K&pnByeJlM!fEj~<$NQ5hc0e$w*@kC^&f z93ToSouFVFf1xG)I%AkHYri+dt8`@GbHvPY0zu5c8A?qOl0fcnL8>J^S*$w_GVV+w zOEeH>K#xNW@h0Bwh(0zWUe)}&&kgrtF*e-{7{qM9$V{>mccZ{TlOHA1HZ5?C4*28p zp4C(4r6uxUkNl-UQ1h7fwXqOOfFegZWJL643V}lMk|+xV#}U3JY!2JetW}UEwRV^V zPv4>41Z^;DC$|Rgeo`sQ1$lUh^v#HBt_kji~(-JVfm7v9PO193Z zZ>}Kf=hRV1>;1l?3FEBI*$LzEtci6fZk_P`wS=6Qw{s-%T)K+r*^vgRg*Jq&;twENzUN<7q^b1okS62{w5x zZwecSI)n}M{g#xs8YnXRVsmEK!)v(QoS;*yeX>ER)fpG2;$-N2)Tsy7N9oj1MZS(` zm^RCl&(dE4A|X5$bYY0RG-Iw9BN+jcECp#{t^U2ke2$rY(l7#SJJr{Ziw0A+_Pvbm zUTWXA4MkkhA~u&FtMcMjW5g&-SIkkaVPONHX~-@Wp(!!6Uyp_(-XxlGG0NEcq7gS0 zt*wv!hB>vv3oHgpLckKXqTE&2@L+;0L9lH+$U)O!9;9xo9jsuP-``k|k)f{nx~4$OCf%r*XY-!zso~=VrWkAPN8DKBtJlPRk+cv&Et@T8C2OR+(_(s zTT?*v7a0S;ET+H;hqGbVGBxn4o9H+EGxmd);;nBE*B7D5}g4 z9XOnM$Jnz8KLxKf2gD4#vEK-tFZdKucxF*5Y3xvXIP(yLSQ6+1mbDZhxQMn889}Ty z$xvK9X;bE=8mICgR7iO9zB8@7>WsWkrJ z&1XWW(HUZG^HUa*s1LsO2=4wJuX_VrZ%@Fum)#Zy{|tLStrV2(Y1Z=CyggLQ>QvYy(kG6d( zI$INxB7(=9CdF_#N}p;?#Ba`ZQ;Av$YSj1>i@B@B?txYM zYI8jhAy-m;qoD)&cm3+?6_AO*h}!C~T&pPpPIzyJaxDHZK4mIB~1!oy^`&{r|x z%84Ar*hzHF4ma#ij;G`p!Y{i1YG-?45fMj|!6ZA;9Bx5Em=XfU!*lMu-2M#HDihn_ zsN#;e=idcC%Ra;F=SeINOClZfVV`wm~$R5;d=&|WLESdN|`ns_m^>-yfm|!Y06wk4=diAjV2qK&)oQ;RA-&`!+ zJQVU!0v}L$n<-3l%<2#CfchIgkNEZq-2C9SUFAN8mHt~0q--mOo#>B_D|iESqhxY_ zm-=t`!JayPN7Xp(gBE~5slONgG=Wi|HXURNnr*h`Qm6H&+Njh!Bt{czt5ch`ngkOG zo}hCFGqowJ9n&?AZt6|j#*Rh1#{I?4LeMGA&iFKU5Sj^cEZPIKAkM9Yr$KJLf?Vba zQ>JzkD%7X!7F4!^Va;)fS_S%&9W?uBT|cT`$E~^KN34z>S?G8!6bI&L(qx zm;r7uI%t!UB!8;4m1&x2*EomHgsjWe+sg$}f*~epQ-Vu*np#`TW(u;Q0VtDB zz_YWSG48>6i zyoSW%{!-)?%ZL)Vi&W@e(GsJ3-A{!xhH9fM*^GK<7FN-a)zwdz%XO=C zX8D@>W^P*c#$THC+tS^Q9Rt|}q6xgPe^^cJr#h-EmYQqrxXtTj*+RFdEdTqsB?m@W zal_)g24}ZO5;a>KgE=f(G8`Tl-^~+p6!+xjPHchxeG<4s6CTfHmmL3Os!c+#taYC-cG<`zdJ(2?jySGdh^Z6y5i(G-lbgp$ijQqa@qW7y)f@ zdOh$$EMW#+nf5Blwpc<1d7%M1%~T?|5^@<}U~Ab0R(Cmb-7N}Vr?j0HQusFLCl;aT zc_fo&>??;YQp(*`g%XC#c#i#Dd=zW~!)l&;$d8*|?tv3FiF!|YoXq>pk&bv$<&2g% zH95qPXi3Hcg778ZHZUe)T64hrIyZLdJT2^u2uf~b3`-N7GocxIG$_g@#*VaPipoT@ zB-O|TNU>?iElc;NAd+fLr^e+hdCM&{aSw!Cd!;ih)C;M3^gU}XT(~u!p3%V zYaEf8oLR|>0_qIR88@yR7UK$jvoNWNFd=H$*Nkq3?ztq#*kJROfS_=WUeid6)TG1XM#$izGsZ3_(Z7nMm6y3y>n78h+!MhT;i)Gt=dL(>`)<_3Pn ziDb=1Y0CuPaC-6X^3!a*-G}KQrt!>{5Mk6&D+BX_&egR>cXc6G6$ld+pRk!$%ohQC!?;jHBGw7s4SVcF0%kv3c& z)CV|J!FF7t2*MlX88?xaN(w1!#LPkVHy}e_NeykSpB}>x*xp+r5tZQ>5v$V02qrqg zMHX;RqKHh9iC|k;EDq<6hNLrJ(Dts6bMhFIjq{V$Cj8!=GhV>}p+SX&;>j>Aav!Nm zH_)y|S!9X!SI-_eDo2Qx39yfizLJP6 zQyX7kjMM6-o&08=v}}2T8g0cRa}UHQUwtA;b%6XKdOwMNYBwV1^Tj0qA~Gb+XzjDo zn*F~1DY8HdJPo*aZEj!sNrIWo^KDCd78~Q5Ui|sg`RB&C+%wGl4$jw}BP4kD8QeRf zYx)=lIT#2(Zo{t>x4(+6&V?1*eeoc(@||q+nY6Qa&R+5030CkuSILPIW5nsaI)@;^ z%i+-EukQiirVqE_0$aUIQUSm@k(5uOS7yXl9ds|WSG_*H$- zS&g4HQk+xBPL6-O{nBK0U3J)!e{LC zO%5Ss<+4%ctUo7Am>Q{JQ`xX@+aHJcfFIj#z-k=ZAE@pmO9wqAbNd%+$n>R9W?7<(UUHR`F9Q4V(e(@d55E;Qge(Dyit`jyre7F96qDAB|KZ7R44Iog@y3;z}qI-Y0}_0gjh%3}3NFid>~zjcPkv z>-@I%K#@EG!sU&a#J**dx$pBf$@W&_XtA(0hBP%rYM^-}*f_>Dw?U!kl-vqt;3 z0OIBSR6vIzFdtTHh~u%jxno{8=WJ#qsEMkpBX?i`Fq=C70n(1Nx=Vdj{pfz~-e29> zn9xvYA{2ZN-rz~tA|v-gUAC&?VCysPelwz(&U@o5=+LEEABa!hxzPWcln<^hls#_& zC~t^j9E2+}*SgTpdp!zXpUKmk3~S13RBXsRb@854CZGO|xvza;X8?DZqx{RMJiYA2 z3ozT|gIij!B(b)!q15tIpGJOn!DjRz=69`ok4=y=6v?OafIfS>WF@K9nMH6%nioJy zIEgY>I~*@DEaolOj_H7ht|}GdU$LBwY`ghKriFe|bUr65yv^!e+KTHOV{Lc=|9*X$ zt*)&F^s^Pp1f!?LV_bT2T1l6s&fX^-v{g%$OrltV3aS&|-Rud-b$tEwNjI;jrd_?8 z;Ys{j?fuijm1+?}_daa!7E0$CZkAkABnh^j%)cD!v9XwU93WRBN*w2(O+B_q^K(Zz zaxk1@lN7b=+{wuUxPet&!pviIUp$P~)){(^8s_RdeSvWthab4x<+B%mhA#^4e|>EF zXIIb-a#>X4Cn<_a^#X!q2k_lKRoXsb{`f>5ls2M6_Y-f&YDRs!k$Ng`+o`oRAh$>f zb0toe^PeK~Z@cGGi{2`s(c(I4myKSt@WXe~0jRA{g-^WaL(=fwYwyh~FH3X8EXyk@Ez*BVn%$--o7;dguHs~)79nNUr${^=Cgl$)vGw{ z(K)dj9TRw86gZO4U@go<__S?sga}^Sp6}ho`*%N_n`S@)wooiTqZ}`BSzv}FI&Xa$ z!PcI+Iw%)*x?w$9>=uzBlJKRF&HSHeAF64a@ZlY~#8W3UAp`hR$o zB)9s;u|2*N{2U3bpY90^4s4MU_N>omO_@>p)_WHVYZ5id5tnxNI1txUN{^0EyJY8r zC!JX-xwY}Ja~t-JEIxlaeah`9H* z5jR#0+G2yE+QnQ7F&MN8(b_%a_A(uPRaYp)DkiJ=I^tM{BKs2*u!aUbf)0MtkCEgz zhjO*~os*4Vtz6(t%MJ&obd#0-+yKkWp#&7dv^;%2g&x3TKr;thU?|xSUE#;b2vbI{ zOfeJJ{Sr$#X~OUT*N*>B{&h4#D<&Y-NXu*5MxS&lR}ojpfbUkQt$8d7E~jXt!byAT z+)gJn-R>BLlh2aGOi`%!?TGO&_SgL|r_F+mhvT^Jtiq(1)-@we3X&!~sZrr0)vyx# zvRW+DN-B3b#-X?bi)V<<=Hua<@R@-mIaU#qaK!~0t$?TA1Y`S5i65g|hqu@6ecr$7 zoSA1rJH(_s(o}HD5oZCUkpDRRL5479kOQv)ItGEezMw}R#Xn~JwCe1TPtzQy8i^Z+ zc=ydGQ`+K^>g&W1h+_XvohN~JB*OzC zY)4>zFW#8*w(B(8q0x@w2#6+SCY{`lMb~OO|Cu9CJsON-4z66+Iw{(4 zW!7oi5p9k?1mx4STJ8w$1ogIoEYC;ex*HAE)cjPM6pjMmK(OPbwwkBJONPj59}?8!MWShK9yiJF9p7DxjnA1bMPWyrx7O^=l1(o%6Wo#S7#$?qtb|WRjExo) zrG~PT)U1_MYd9#MqfdNjvRi^@vuZjP%g+$8SWAZ~82C}IwG)E06osR07Q9If1-YR9 zjG1eiVJ)S(G7Z!ZB0fDTOwDNE#D)ytV*-!rJHG(Q1*fkOlo5^P$+ohwWGY94Jjp^l zc*9LJPO&<~T0RB~RvL59W%b`KXzOjaJHeu=r}y%a(NLRc5lr6<9npb3OO5M4juI4j zx34Ab3Vk6@mlVsHB_O9FTdkZ#cB-j1^iSSOKrexfCkb!iBHgw)oT#&9nF6N{1sz$RVel{n~nAdxtwU} zN5Nxrhf{U(!e3%7_j>zxl_Km##);xW7G1<}ds3z7DATg9WQ~Y}Q!R%|%W-w(&gmca ztFzENXV-CI>|x91)$ql{h;zg-ny5`cjpiV)&MILCY@Nj7UPur@f~r7zbyh{?J9_iSWRyoj)MjZ>H8@SQ4c;siXb6Ycddn zZIaey4cez4I4+m9eX+J1R_`IHr>c-Gim#`Cr5`n^9GFf4#j3kt7D{Nigo$@hPqDr~ zMB(QyO{b0DigQ5h#L0n;+}ScGLAd{pPC{RbYjIXWd!DoG8Ll~EQ-dOlC6_b&Yh-~$ zmPD(rbL6ks*qI(Tn9f(Q_upM5u}|jL-%D_w)wEL)lCT>K6O-qZZ24_Bd>~|^k^)UB zL?3OvlDsV7sG(SuBh^UT)-|Ml(w#&!XT4*t+D>|8)i`?@NfOidCP%WIG7_+?Qs=!| zWxAPlTY9!iD@(#kd~PAJ}QDs5vG- z@}ziO+$52JZdHw;Cu4p>wgoUiV!uitB(t2c8NWP1@Y8c~loyzSpA;WrD#C(~tw9?9 zk>2XUr@^c8p+TKxDH3hDnfZo()EnE<`6&Mk_u(^Xxp*pQFiV4gqn09zwsTy}+~>Aux&- zgdI|R7KLuc!(+p^+*Of(M(72O98u-kO6Iv-fe|I)l z58LS`t+mtX{oP+7i-i2I$Rc8u$I6jm82S*8EI52J)5KlSQ7!3~(VvSbN_t7NgoxqF zUP7_-PrdW}9J7pYq8*q0Sbq0|BQ=XzUgek;UnAgRv4rxE#YpPql+bcMd=78k$F@l6 zVV_O5LE||DFvj0S#5-h|&XN&NN}g2g9Y9VB06yL);ex_U)A;WTFUJ`My$ylz@o`Yw zF+86zPY?tM7}4>w$%4&{WAD^nS&z$x(HDN=JgNI{?73wm$)3GZmMc$@kzWlwsHXaJp()eZeduq?z;q<@^o9MqQ4~L3JBgKw zo>1YT%Fz2{4716NWKBFb_wxaX}HPD}&U} zZ1sE60GyYP&>nUCTigOg3h*@ZL==Hb=6z<)pA+~dS2P68B#o3xfm9z#J~Hb_%rEnPZuadJkc}Pc zPw31X=sg$W^+@X3jRWRtgT+9Vah|8E!gw-|VQ8ePMrt6GVjL zyJ?185Ww31V7BsBbnTyqVD9NZT~)tC2aKU@o*>Xpt-&6)>WxZa#RfFS0j4*rDwYjO z{r=oK+PhJKc$>uWPB z_6LznhK4(mj6(eVXD)8aa76!8RzG?n(;}~85OM~+OVM>fxC-Q@8Tzu?nMiG63^^RW?E1}{&E$2i2B_dQ+AlJe$Q%Ofk`Qlj$x zAF=fN5+$+xnfGOb#MELIU&7|y;N_2a{*9SNwH|#Qi8t-78im_)7>BxGJg!SEG@1kP zqT6qY<|^KUu#{_6^Z00~GsSTPaX z1+{X)iVvd&r;baZuz&Ygdb%VlXH`>Z$#*2^`kdw7Z+j2p`-SGeNa)(w|K8NK{_=a3 zERDD&O3u7kAzOwMWE}Y9c6iB%C-OC8ExC;^b91tV(IvejJW251qgbqpVP`V%s`29R z?njWWorzJwt3kmx(eub(0;o~%;P|RGAO63Eky5v^qYw(B!>SrcqHpVD zx9b-9^aL5&RYwI)5jouVW29fCk^c-J6B6Cu(%PpPEb~z^>%((sjl(KRz{8gWoBiC* zgoGV0E~3JS_x2$pUi>mhqQ!?suLRmw`Gf_|kB80(h$Wri^dM;T!2|id{L;G>x*ulv z%c!GavTM&jHMKQuOjoP5OaO>fRQNzt0?$?R;&ENMvl|3A7%ocS9%CEA%ra0dRKq~{ zculWt;V|+fk>z7=BmZ9(z&4rW^Y~$6(iIPeH~LMr&)Anln{iFfqb99~l+bm0<`|n= z)NbFmjTW1S6^$y1%Z41%by}0jHs%hWWi^dGO3D(~gvq~268})6n^bN7B?*tKP@q@9MWaw~=IIhHIrO`TpJUmboM^Zc`Ycq9OUW^Jl)V zh6h(u0b4MMaYw0%n3#=fg@6{>Gi*QAgBR62c#2ihlu02r^dK{F;@0zyKkBCEujlz* zA2U~yqpmyXu~`0(tg})>GBPxJk?Ptc5%d&Wg>)3?R!`wQ{EB~%!yai7hmGNP#9_t` zbzRro%u^DM<#k2WtLN&aNLq%oNjtcNQ!J_su;(Ei5# zJT(x(WX&Oa5Q>kP@v*h>wEgAm#Hv#h8)B-aghMwq0z&e@9gPJ2h1p^|i?d!@0y8=0 z)J?BoYR~ZMQF{}2cO{IxHh5dc6^b>Lv6^(DYIQ%;_=#41Y&!WL-yc=rdOq={q=;+^%Dk+hAVog~`j?3*HbR9sX|rJn473)q zklWsN9DaTogi=>oAy>e$Ow8boCsK4@MJ2$<-v3HUn#se@vl}g0dy1nAgMGKOe!vk{ z0ncS%QQu01JJ~v9S8Z*#sd?X(6W4DZ{8@u=h^*slg$t_tSE$ld;!7gH3eH;fDxLhY%-c}$}%MlU{@M@7p@^oHW;a=!(Vlp8u;q;C6OlQ zjn=Kt53*X}iu4z*iqGk1mUC1jhZSkG9tfI8>hGFG>SN=DAS*@K(OX?~lLjgent&5y zkI*BOTy+sYHb&{Yd7X{g4pDpI;N0?IAq(hc4?Jatz)qq)2m;Cs1>CJ3G`t|ufgEWa znzsHgu2fj4(DDMKlw~QVFBNXIs{Q443ZU)sjTS14Tw{ZxRy&gotl}dxY{Tv;PEF@D z(e;F!L4;sS6;&XaRE0(I)VQZCaMHq5Y3)TBXV@FxznyMkLS4*s#Rw0zB0N+_d!94# z%+F&-ls?Ir%pD+_yPd}oGNVc4PDwhLQSMF2l$Sf#>+8-gj>S%aqU?R(Hln`AzukfhFicG)Oad7S zW^lSc6mzek$^s#WKXl+E*E%igvJv|MBBM!Szc9E+WlVDwBsA;=6CNPwP0_AHP#E3j zK_9iiXArVO<<=-fOq05_t<0v$5w*ZjOj?hMqS>l9CJj-w$@>apnr}2;haV&R^O6ia zwP+M;9cK7L7BQ`aOi!xFX+-vUvXV@@f=|3tE)?D|2Y%hZev7}|XZ#ohJ!IDw-}d6* z!#rM|!pE`@rB9ly>^t%b79aYms!OxYcnw2wVlpf1eu`d!T$0edUY!U2=2j`{^8wU% zB66eMuMg6$h1?{BI*MXQ>K$?Q)AuzF#sukF*Hhp|W=VaV!x{9{qiXxMJaKlTSzZwP z8j3@~ag≥v5@tQ3!4A_zh)rrk;cPVoV_f?&6)YeUK^F+{VEoLzSy3I0I33na5sc z0Y%ydJtUN0MWm`K^qR75&iAx#S|F0rGFEe?hODheT|`y!8f>}u44VqE5F%h<7bNFW z^eNb*l8&f&m&}B)6T3~@szd8Xxe*S8qaWJ8U=tlwrc0{Nqh$9E$8ZW0qlwht!z4|5 zD5jjZc&Y~cl$*c)8vjI<5JHlUfgxwtjJ3NgCWfL)M)vfG%Br%vPz>L|++WjpMpQyO zXcJ{7YZ7pTB7#sZjYE3bk;X1zgQJUptDJfGNcIwTdhW{uM*B7J*?bh^%Jf?==^V-_ zj0>=e;|InGN)-I}VXFhg%b7boQoyy}|9kF#ykw08`Lyg8y>UiY^8b3~j$@0^ixb`u{&ZaCEZIx>*eI_4uOMpZXaH8@omN*hOkT%~v?~hk9`W{FOcW zA42*5%cuURMgM?E0F>jm%>^ImW%hr84uwQ0wTFFQOY8n!o^ZHgv>yG#D^d*Pk~>R% zDYki65k+`z%M`V1SyqoRD0}?(-Ul$M__OMtp(8&+3%HA_bY?1f0mbc4{5Q|CcKV-7 zsa!bV>CNZDRv#!-7yzHjKZ5~XOg;;lG=~Ud6`FS!Gy|s~4rc~oq`E+MKD(FOy!+lC zy@G;+JgnyL*9UEj;Nm!XW+xWzO9=FyuM$vlP5Frl`J_R@mio|w}CNFi$bBFAKGlPz%7RB(J z+|MehVGG()v1QfweLU7thwcE{r#}>vc-4Y6W0Wh>EUkNyAs^97SE-+TSw|Ns*G#?RdTJ`cyre z!qs54^+%2eF5-BuEk#2>w+ugkq>9(0=wu4ZsIa()M{39+D@B@ z(*rRySmQ{{np)MNrA}03`>$PYkSq8C#>|Z6!Ym~?i#2UK+I-A{w$ca`tY;YvdMRpE z>X|IZ2R$);tSx1m{eI-UqKQnQ0{0Cq{84+t6{m^8tP69r78nocVm%e4^~%UkJ?u6s z)#y51x%h#PBt#dI0vsu;dhXkwmtivEJSk*ZSMTeLo$@vi9}&Si%liNbp4Hm|4G`Rqz$#_QcM+H!xY@qO`5;oHSN!6nBYnc2MS&q{LAZc3q zUYU~+iD^4`s>Mb*FxEsPz6f)uw%qW0y#fZu|M>wsv61iJ z;d+4hLgsJVE{njGJE{O>n=nW+iU=A#lZ?3Ty6OF7|5N&HD6K6aw%*6q5vhJW2>g2h zCY&P&*p;UaERVsUaW1U`cxRA9pb0z*LW{tVPDo^j#F)?kx9q zN969!V$7yI(!6;{x9@SZ8;j(10x{)hg9lN>UIQ5w(Bv0a4bEbDEVMODPS8;6x|(sX z3(9#9IVasfxTQCc(Fp3`!f~78_J}RZsAT^`FZYXwQ&nv<*og||Vv_KWsK;qMESW)l zI% zotpC~I2)1T61KRpn>PGFjBvq!(V__Qd7u1l8cG}~t>4z;$d z&1J0tJE239juSv!FMTU(Npr=`VycsvzTvQtK_PAo)wyB`_d)A1vIxKn{^F}=pScM~ z#!5^>)kpi3$0KKbR%QHy=wPCNGS8!k^%m%bdAwe=rGE*z9SqAR_xbC#9{`T%0i?c; zUYHt0vPwISj7fR5l)@{So|6v<0MpP7VoGHx4m#o#iB^3@Yurs}|MI-Atw^u&sxXQ> zH`$BUhIni`07l*m<122>!LWJOw99slUk8+UM8y_W!!Ul+< z8=%&YL(_d7`PQrSyLFJv*KR|eb1!7uZ8LT?EKGc7I`*U!9%FtO5}~#;LCcAy-L)$P z!;J-YO!u~Wz)qi)8xJvSue=E;N4Wr-9wqek#cO>y-{mL2l#>50lD4)Tte!h4clA6H zPJ_qr2uRhr00{v2pYmCTTlOGA-(bZIQu)M-U_W*sz++5@BTjUsevm;cdI8XV?Es?fSzaq9FwAx8?s7L6=}*)>(CL{s zLoSl#jyH_sI?m{}n2gKmKV7J5$#L7u4HBy2Cizx_T04^dg+Q$v2x?#+<&(l&ie5kM z1Hw3UTC7v|Fa_|~B0n=FqGLbrrHI3m2FayFk%G?Nyf}QfHKQ0p^j?Ueq2Lr^0*`k< zdFMe}(L}QL|I~Je?9uBTr!!ZonU9uH)Hrc-xFmz-SPJLzZ$l_KU%1Z;pl^kqWdw`$ zR87UAn3=J2vCINWZwpM3O*o9_qdkjJe?6iifo)pf6k9|I;nvSIoZ9;X?8gY@DGyzx zYI;Sd{bdcKHNdx_xW1+P1cWTdc`D<<2{<1DF(;jK3yGmxE5_Jx2PUmbGi$>aD^BW4 z^UB96SSK;F3xNhsx=rS>OJTxT%Mk+{VI5hz1X+@y_cEpbyom7|$Y4m73mFYqxJhlV zY8X^UZCdAXZ%sT;J&diTs^FA}>>&SmNf?Te*4xAP7U7?cxI<_|*a{^*^TxLsq|8>| zFI3u+4QpRFwd>p+8E*1aTmXy&f72h;uXbUU!!u3t8b_)$;!a6hKTHkO9^}bg&D%1~ z6@OeyB}-*CCfbopq#QE2;CDnewHkM>S#q~2;y*~EaBN(-=~k3l$L88ow+uY;5_H=- zK$6lo+mt|t^094<7BC?VQXS$#9vE}>=gG~sE2C!fn{{c8IRCTmfgbF0(3G>rS;;`qA-MXC<3kQTc#D7JZFEsL4K(~2sRsu1kU zOcwi2M5VLo90-HQVii{1^RO?eO$NddlK+zkX;l#^$d$`DF%G4vtq$M!dmM;-6~#gu z#Sc_uNL13$Dx*g~mdBV-nAtaDCT3%@^?qEknFVxEUHBck+pPII>~-y%$WRiG}q0IOZ7J;6P?8`hzM| z9%r6$^T4IfV@%4ZD2&}vfsck3rSdt;Eio`FK{W3mp}YO1CR`YJ6FIs_h>Ijwkx?l} z$)HFP3RmIu5fq*uVzlX2Wwf9`RF<*NCYv%N$ovl~w(ga?z9<5$-D2!B;0pNGO*m$yPCkNwWp?SNUFW;>QT(6O%P#Q!Ui6&ZO^EQTfO;N&MRz z=aBkcQEMWh(jJm3_0((FY10y3ViB>zp!(;PEc_6nXPk)Uu9>QZ+Q}%C$;@zER5Ye6 zoE-w7tgC*rC%=T}t$%S?tqxD0#=f4LLZNBYAc#kHHg}d{u?8je3W^(y436E0;uzwv zl@9lHf~e>61j#Y}kWJ~_oj>ePqk@dsow{~Vd_qM=3%mDj=_q(jPWS+ocNr~&97^KmpAd>el|RaZ*;0n`;!X!*&rvy*iG}KrBbDc6urDOBGHW4>qu0*k zl5oQNA{J@QpfX$TDL{7navfkV!hP1g~@Dp4RcjzC#0wY7uCvA9g3nd4~cNZ z4p8@r{Ej=IsCr0*#3HAN6tow-b{A40jXZz2Z`H}4!O(pT^n1}`bLlyVAu-%b8=>LT zT%o_Sr@xRIZ&Yop*hRgcqc}D*GMY1JQ*7)fyUAvI8&B z{V&ejrQ`J8Fjj5+Qg0|MV&4;n#J#5&q+ezpv_(vQXHMiHL~#cqB+NgTEcdpm*5|Cl z(Tm}&jxpvS*Pc!Yo(>*+tilO7PJ$G%D4$ovz#D+Uc{f^_q&n_J--+ckM;j;g=*ia$ zP9nmhc;Qv7qsIs7a7m*01*TZFL+r#>!F-&WbxY^Sh2Zkv6~rYU5sM=_px2rSj)p1} ziq=&v7|ZZvQiYm^D8@?8^9JYRm4fM}Un){7Gk#-*tL>j~OLLTCtgZB^i@+?5rYHhZbj zKtW%-)o_u9L~7X6s$8en^~XZsHFL@Ua7r%x&Q$WUfN>Fx_vAID*t%0!YQA2VlurJTF*5e;GE zf<}=OZML~i@U^?F>Je@J?v2Wjla_#ko6#Yul?l71` zp5aTdu&M|h%J_vD6ix4nz*5nBpN!vFQnCu>p70H@lIBK$7E?tcmpg5lpcqjgzUL+P zIsh3@->pfYD-7M%6&WO4qs6IkCVua4@9)xE zE|S6;sQwAoH%1lc(?#1{HSq6wl8~4}4HrO2Q7*ed9~(ovq|`5c}-|0Dl{YC5&ko4>E}xBWG$M_`VO>8A~PHD&-}qy*vCY z<{T19scU-Kkhvi0R7hR4ks1Qw0~?NnE0H$D3XBRx854rl(+heK@)}s5u)XjjiEH71 zo6`P+R+Xyk+XooPlH&(eDu1Bs;j>Ekyzv|0SVGFzU1%cELE!J9g2Hc45MdekFteiD zQspd*H9);^gFsJnO7B%mh3rQ5SEjs4VeByH*{Ci`rZZB5cqR~1k|s|-bb%Nx=7P#5 zm1Uf)q>NVA54|laK zrbxq~0=Y}pVDuc9v5m94H07kJf8M|T^rwDWK0So-G~|K^>*w{XSIMOsXx2?n<=JO8 z9>Eg&ev4P@lR2xS4%7Yp2_f<)^gvokI)!^iDDN5VYVE_dR`N>zzCOxr1bL9@+;%lb z_}#uK|2WF1d+t!t{pGCem(Rm=@HyuozV8*xC$fD%F!-kFZ0zqRq1__ix78WmmGj|n zxR=taP@;}LH2+`+(7B;M?^@DK>F7cu1@VceFwWYZku07}-3nc=vl72YX*zPUC$S*v zQWB2)VWm3||31*(J#cROHC}2*rd81jZsfI6O}_nhIOKg9t!QG`H?f?uortI+m&oQz zPLuc*V^Ts(VfZG#a%eTc{A%$Di^DAyihW|g>ihTM=Q}NovT9sR82te>2Mz6*^`%E> zPn0>@OIFbs1X_;Vk8>c5la$cf2(WvA4%h6Cmm(v_`CAsihc*&S7%*@gLc*N@2qE%EdQ+rxfp%LXk$c!I$Lm>97X~ z?lTiy)bfXmB1n8oj0qJ9$wtF(pB%o=FfGx>4D<&oQwmT=Ovr~%V$Y8_?uq?;N84?! z?|e}~xaRl7x1T+bPo5T=nK+1|N-0JeuaJZEU%|O}djgj^ArP&)G$E)t$0$J!1Yr#j zg_pk&Tfes%;@v-^2$ar|tES>^#e-F>2d% zIR68HAWO-xfXT+{2uu~nwm*p}@HG|jG`?YWb7Pdx4H$3np6@u8zB`o8@6B>!ED zGE%E7kc*N-I;O$fkh==#O{$m-NQ^;PxTY?l@L}C<>@_dSKk(;B2oIk6043*LJX>QA zHcs4DX;*`@58p*~lgkoqK>OpwcKck_uDjgRSt_&e@0sbk)USv=mneERF_`$V%!X)s zCi+`3q4Mf}zb9+IA*R_la@-OAeJ{8Jx@{VW_9XD%kQ#))#d+A-WH#E99YWO=aH#qA z{!B*y?&a!I-SX=(WYT=qq)ycEaoV&_=LBwa^%JkCKh5=adR=Z(oFg*G+Rkapf#Xio z?LJW-SB>&FNJhSyuLJrJHi%RMYKiO zXzBQ(CBi3%&q%sr?1YcT7OlwVh|dJVv4+ZRQsg-pw!+sy1B@)&F%jo>#c*@ z`olMEoS?06X7T&Rv*L>Ae12KV#h{Cht# zw9oK(hw_txo02^-P<dJ)V3>_-6CLf2be1GE% ziDVCPb7Rv{luaTPn{m=9%tih@EA&o5JS_3QUt-lfxT(7a;ypN+RB?uwq3A<8x_2x% zc0UqPf1*J&HJ^Y!<~HDM6f6Scbfqh0jq}}we%ilD@gJH&KgJiL>JRFVlxo#gPmpP3 zxN2W~i`^=c6>q!Wu4r9!InWf$8eO4nz-0UW{u1fa*L1QsPQ!1QQU2RkAIm3>lEbGs z8l;RQ-n%#K%`&=Y$=1#1o`2;UkQe%?qL~?ck9dg8W0+EdmuTG9m!cc#B8oIQc3mx> z|Fw%#+K_xOlc-R%OQ;3r`P@+;*K-B*`D}Ahg!ZWLEMpX}lB{>m>-Oeg*NhhEN8jCY zyFfM(6P>M-Mo5rsoKWj6=iv(8w6+tTsx4j$rv^H%GUs}m-tQ> z#}JF$BQ$mKPP=Y8My3op0rf#MP!gSUM`1MQKtWf#{dm+hA98G-K`$Zp{ugHeV{eN0 z0z2|ZiU|$+!_|r6P||!YIn7Xz4V<=*`-mDl=?u$F2^U+u=}EH+J=DY(XaC9DuBNr zdq}MYr1%YuaF&3(`hal=&$xTfuK+yaEQwIo?bnB2)~racsml1<6f19?9#pcKvvUuo zrAAb*k056UDiFt7?7_5rX8JShtMHfP`~g^%7nihY1E+o)OJadficibB|P(%7F*&MCcG@G=M|;%jzMKFk%* zF=A&dFtFStEvz+^!GK*3M?Zd;NIy@m_6@bY@r+bf90I~7H$%q8tr+JAJcwbuIUC&i zJXy-8(+A7tvHm$S%Em&`Z+_4OrMrWw_3OM)d54*`Q3{PiSDwQj^SFtp|Mve4mPYsy z&=P54jxMJ2lZK{ce=Wlon4VJq37TU^kZ^YJbL_K9lp2y|Z#A0HtB|>TP^|0nbJ$Du z>78^)TlZqEouSBJo^9(*nNm19oD=}DUwyOjl|TjQvWi`lV6h;cv9>@v=$ z^wVpZlndr~#-*e8FmB|-*%HF}6!ss6(j`T)H>tPvqMKUs>8*9;`Fs+7MpfxG+#bO! zvRiik9j|J_`ApAw?v1uoyyX3q>QrqxhB&eu?M&8*QV@RAH^uDjhIbjc4rx@I4+WXD zLHLy!s;5IdD~<`mTj#!zHrM>y!4-eTYrHYC@!65U=1gstN)mz@~Jh}XDF5>hH@ZtcVDsyTOK zZZc5CgSi(%hFgVe_~IV$<1YJQZKrk4=shdj7Y#y$$Zc?uM)A1><1ym+!;J2P!Lg$1Q-nkHL z1mVJO9iYHZbRQtC^OAeRx{zsBJZG{om?*}`BZM)Oe0#Rp2yR)sC?9HVpo|07N&zHQ zS)64xDOl}jUe@ucbBO6CDIN-y;p;Q=9`8vP4Ei**G3)1=3Jla|@$5a5StV61Tb zq!vDQ17Y9*MGw`PY0c)b+$vo~8#JnJm@+Ff&IC2OwRM*EZ&fX|e0|EqkMuBxvCccS zp;&1-G}%d2IfS0`%f|C8{Bn`VyuZv@vJ;db9U1V{ML`&R1@R4ISOO@!^>=6WJMecC6bE^+@l zf}_Z){G__ru(=|$;X(*j+wcMZ7p@$v|4WYjwC^lHrbt{$uJZX@IjCx(apQ#0SBd#& z#JZ8N!}1WsB3(v{YLHPR9se_TO35BUsX8F%s} z`eF~NocDFAFaa^21zvd%dzLBX+^3TCrE6&NCTk>h0OQY0Pf~)me%XyS@wc~hV3*8@ zdieR0&1XA42Qykhs#^puN?Dq_|cRQSyjZACw_+K(4Bk~(JH&@MJ|IXb?!4d~P z@e>0BZLMu=Y-9G%E)}9120XSNPUfdB@m~!(nJJyYG(_MF23DvIv$FMsQc(_{k|lxO zzGvkL{SiY3Qg<$8%XWifnKan-LS2~3XY8c9gO{b-N~Jj)FQ^I6D*qgupxS-YVOMb4 zcTBq5GySlY=Xr>2fg@1sqK(Vfqy%f?D2{q7or^PWh*QQTct2FBZh7ZCVt-DL&*MWN z&QldWZgej-Fso=wZc9Kkq4=AgcRGvyxlVDxB~3CximQ?EdiH$$sy@c848t^~Gf;$Y z`f5zuK4=1A{Bc`Pi7SidHwKw~Vma61*n>N);E<|(UghfXf0=2>VVj8)HB$|PvI>W@P$ zYX%|^+0sQSIG>k3OmN4QG~h1J*lf5oW^TaBC|t$pGg^6 zy$T>8i9*&3G0s5IO6&y9BVj)0K^hR%VF4@ld(2xV z(_gvTli}etHoAOOI5g4_K8**U{_b+^Mmyp9pexqY(~##c*0r3{>{rbZOCKvO7xZxWX`*jvBC! z%~gwcxz4HU{m3z5>SfVQiXhEx)%875)D^L<52`j=q61Iy)_B<37Oesq986|lW$818 z##0!uV&LLBOQ?n*6sgGHJwqf{kbJw`u_QnVll`d&zL^E=}^~ zx_OPvmI%BaNjIPo{x~y$n`i{BfD=ZFfEV2-MazJz7%5$%HESwwz@7Fp-=-^q4_TAF zI)#owgH>K9WB7hmmX#DFSi$?Bf%|R#UpOlc;-d$GS@?76Ur0DVE{m7*R_y%pk)jR- z_439Z`WU^W<2!}b{36D5Zqm4Ns+^u4jy8#l;ql+1Nrx_%8;`pQ>%`j{D2KkuP9GMW zwc$VF>Pk;wT&9W`C(^jk{)X_c*kTyxB$t#ZjIoyd54$@(bTfEXEWyX2!o>-Avc+px z6IP#(>^G)?u*I0FM>e@|Q{eNMLUx*n;nsxRvy>h&o+j^DM~@mN98<1Cg^#LkS}N_T zY!ODjwa62LnApJUXDtvpoWoWVa^(SS)f**dOjOXhbU@6#=w_|IZ{)jwUsy>!Cz9wL zbKCHIv?=YYZ7Il9uRxcEiqL)3zGd;fteYa;rYPa#-3G_>VrTOzS#VKb%Xzxg~b;qf14dgdZi8?jvVd2M4%GIqG*GNMU0S1RbT1J0Z|y9*oV(t~IqdL#1n zqSLleNsLw_dF6EdS&%yBnv`&>ZXu>3TV4jbeO);-y`3X>_qWm=+LMbngncAFn3YX- z#TrE+(g_*x5lzf)U^v)_l#lw2Xj7-_uUZ}S>7mMna=f>3Z@$GrG}QbE!mlL^G#%#! zvrJICLA zoYtG&d``qlf%@k@BhQ~PmM=?7ii5V2Ag6@6bPBG{9*$OMMp5*Jt?AtfRMZfg6`Z&2 z9l0st0nWfF+^*O9o^&?j1y-8nK!m!4w$5|dgM-`VdT?$7=rhxDb!W{)Dj3um>*oyncMGPgYo6QLgM9WIZaTAsh-}BpN z5EJq{(`|G^c1R?3I~FYCt!|`?Zjvy&OeNusjhBb5xggs3@Ts^`AEh?!XBdW1lSq6^|`w+8M zG&sDfYYk|-sgLGRg_)LBu1AqAUc>WyJr86PJYE*Yn?;3~%Fju$oy-4I<7HI^@ z=2n~bUdhgs^NSxFPjiBnX_wi`$Y(P5Wh+j*P)-r-&G}C8zPHYx|HX>)% zc+Xn?Ffxb>*z-V!-hVW1WdHQ3lQKdI35p`mn*bin0tS z{#>_vONci5E+X?0D)NPMZ5(b>nF^Y=ehOy(O#0^a}9#i8pX)T}F z;UpCbBC2E3nWDJYY?uLNTob(h_y+5;So>EVFL~w2;fg3ixF|Qd<&yM(A@Ua*b`anG z9vOxg-3*+ZbPQUmetMoJ@sKg3>|dx`1DiE!mUT%^rX!~^c+yb#MW z+d1#&CEoibLyUiU_b`(%YhX(RiH5Y<59MlaUd!uW)m!py~YY+U1^Y2*O^cext! zlB*y-`g6Pl6x@H7Jn=rX43n(SqWaM#ivc_r0cI$_)=a7H_+@$e=LL6cFTRoO_cl09 z?6LofjzF8b)RoR5>czoPH32b{UaOyY^ltGEd(%-SjpqMEpl`vUV~}p9nd4P-y-^Ue zjA|V8IrLYNjyB9g&7vcdT^W(oz3LFA_oB^CBl0zB60lOV5D$_AJw2<$uk5LMUWy$q zD>~ivO}JV|{6wc9vX(*h4RdjOQTBwXm;Vsg_5^z;{Jrxcs=hryeMiEiXyn2on2$Nt z;DC4N$uLn)Z!PbiwTAZwb4g+V&|V*lc`H~F{U7diE|PLDCS__yH_rk-V3l< zW1G33Jx$QzHe}lZ~ns5<+)u3j62;48VdfR4kmM>T|wFs(u*7 z5_<3%swjlgfJgoMJSk2!=QaYLn3lB5~08aPLV2iL{O;2JBll0<)WK?t~qtpgq?Zw}n}=R(n@ z&o4pn6kWqoN~^$DgkKr>lMKXN<&#dM%w#Io(~jKO0_neA zSxu3}z?AW&BO<}jtk*zfCC1g!5(k_w8hM#HkcaHAC`XIBZcl(sxsGU|B>uqm{qN`J ze}JV%wfS^Y?s50sCv`adS!Ao)F}3cJIA#Ln)VELPlM+8Xor(a1xZlcxcx;F@X~?|5 zW*nCEl6SH{%(6lk0@e%yltav|usB{rgau1{e#cGOp7X4_Fd$=%=|!haw*Ou9l|g*X z-m0#rI|qjYY9)m}2V$gdGI~l<94?9$pftq+%RtFHrApQoKz&XX^uy09>N_c zvf$33%T6~dxw>ckqJX{v*Pj71V4F?_M0=pHlY2)`F^`#@g3BxG`-PSQr7yq-D_hQu z($Ck=nm=hLNE>?*=1`_)9Pqqi+VAeybUxl@R;L|rY+seW31vofO(c{J+U%9ke*hm+ z0-;E4foXJ(9-@RrubZg7$uHgeVZr|)(pm{yfA}t4g2dKZI{o7epg}Vbxk;?7gNI#c z#Nc+qE`$wz_|g4b{SppV(mmw%`0hWx+*NOPlyal+)#&YxL=OR8nB|!Jz{TKVCNk!H z1g?w@PGWbRyEaxLbd;=s85XHr>K8dB5~ilOgm!yVO=UO@t>pq59LxxIghr>XYHSXB z{<+D-mgVuP-(X-pZ65p`Ja?VwwhHw2-6ecLz>=-HX!g?9IT{&H}Mq>62CSP;oyD}PYCyJ#8aCx@~`1ub5m!ha-OL(&8 zRm@Bxy8QdM04{@D<5}X1A{#K^(Flxjzl{Vj(a$3oupY$V>-({9WCNStO#c{^baMlg z168}F1t63;$B6Jpd^1n&xDWXJh0s~jc9p1PzxL}5+ySss$VR9G+Kv<JWZs}PKJawa%Qq>cSd)PwFQ|V6~~Pw zg!eljA0l3mqCTwsG#~AcE$N4$#0(A_d}V+CURo?plk$9`5SLy4LDUZ4sgsm`PT(}c zE?a$-fRb#4U~L56+A<<5y?b2z@b!^4`%?nNS5K*PMp(E}@DXf_EebS58k3jH8UWb8 zewqdxLI|1Y2M%0?iyy|A^1sDPu!YkO`2p|8x!>MBAHKoLJB_r(Y|7`oT&Hvur}6f> z{x(B9Mlf}~t@`0%lv*(UORKw{OWU6Roo~NX>Pq^S>VOWb}Q8Jgcn#Q6dM2po*GTSZsIgma|UxF zqon4bG~ax&e4eay99g!JVogq}R562opuE|rynJVF8<)c!W8j8yr2|$B%51!n5O<|K zlZ|(N(<=`Dr)S1(LLvyO#_Iltd8N$_^jQgB2m%QT^q|(=TM!9deELbWSn~DtMamAG z3dP!GR7j|x;44wZfZbixXR9ERYfDt87@Kc6B|gAtCk6eEct_Xx+vH}A)P&2+l}4E4zgFB_imLL) z7*RSj62sY5V78B11Z)`%U4E5fhzM8Po=S<#DxR>Ph>RdFh}H)AM_A`31qd3Z=~9H- zFj-$ub7DcCu>0UpAJSrm^F&%~L3eaxSwx{c0S8AjczuRh%L+Fjlqgzt9k|JinUb%k zbY4tWn0e?nT9clhBuhOpkyr}>c)rYiL;Xh{y{3k^L3~dw2n{8jwmdtGUuCH}s0%vf zkfgzN~efgECxsPO9rC}XlDu9tE%l=K#yJ&ui zQiuu+;>-POz3$EP{G z+yfI3=rGc8@L2Ob7#-BS(vVs!%BPeT%HnWJw3Hb|B(Y;EScuGdjmqf7FX&_;INbS}RXOZ|zUxdKqkyw_nY;0%Ck`8FL<5E3Y%cGv z=3rUH!Fo<+)iN!Mxi}9G%mpNhyToUyCs6AV&Ir9)FF01|aVc4O0s=k-J=5B2>TVr7 zH_8u3fh4elLgLGimDT9~RL1pOnvla7Dsn`3 z9!R2R`$c3>ICGtR`bCo}@=fjBW$t z@PsWa2h-8I6DIvZV?fmk!mLB*2GO9yffFV5b!Ox@t(k( z>+o*DHV(Qlw*p=Hn=U=_{%es5S`V4yphDn5 zYC_TuBqk)ZjPm0YOaCv;3dKS10q@+neKZZ1J!({c;b;=Y;f}px+$akU?_63Io@{-b zZ!sZNQ5elKh4h-KD{FUVnQUj~3-0l%8FAJ~gj*yQlk!n)swYQECi;}_Kx0enM0DX` z3)aZ^J9iHy&Sx8E*VTc(8XU*osT=m?L=HlO3GY}=6mFZ9lFZ>;G@p8mLpvp&x#~X_ z+yVdB1)#k_w3aSd%T<|QqjE^oLHZt9lV^|j?{ej5JgS|QZr7MxPI#lrTvWZyQqD%H zPui+_xADc1WwzPyYdoWF(v<>Unq1at%Zb_pZOyr3;)=2BG@*DrG6#UKib>If9Bwig zjbw2lWUM?NJx$lpy=cxGhyFCX1BvA4oiY{e{!r~`tbCTCwskO|tQtWwu_}HQaLvxmIE!l*|WXKMds(qu58hX)WUJ1DJtf?GgGas zCUhuFjCus1@bQd*Ay!WYBBEjpUf~-}GrVj|+WydX+@~X5+bL@f;yHR7H};29!B#CO zk2h;nF`Q+=AnEWJH#&xzq^z1B4JertuPxk9GwKi;Ob7MN3j@oTytpk}!uzjYAU{;nQfFwx%>;DGQxY8*49o2&eCH z{o84Xv9_dgil!?vG`yzbNEsVZ(-dDsRGQ^jTd+b!*v5pzE+6>Hdf!P;8A2s@PGOy_ zud%P~1j7Q1*RYgWA~cWS&+!P&=c(ZgZe~HsZA-zdO#vbxxz)))-a{5DQ(L@h7T4-v zDsUbdII_c0ImeicW6(`PVOP(kY+0Ib74x=rOb}d;RotMGaTMW!+2kXKy$V>?L^Kei3bt2m z>sl2<;EMw&=7ob2pA4c@$(5}8J%gl>^+xTgapR!&1{Ww(OXhGC0_wnwO zGrv;#DfApuulV;uax{d2u7wJF>{_LJ%`OkhdQfzASmrs@vhHn@8L?cICzeUm;N7KI zV^|Bn67e#&Bs62u(Wm~ZbJKi~Pu1S;6VxJ87e*A7;!~QIJ;;1<{-2!1%WY1m4#?e$ zo9e~-n6cP)yU6d$bhp*OuQdWKzdXSPny1u%4{yHNndoG95Pt4G3N~Ff8MUGh`@QY( z3me&cg>OB&UWRT;vP)!zT`2Nc&Qc6cGZ^1qZJ1& zu7*`vWQv+Z@FnDnDoa3Ui-p1z4v1&v08wZpIfX0y(T?n}RrBz{*B=~v4Rc?#op^j9 zp+oU3u0ju%8mY|z=v`eGUlYuIfYog7IYPH_m})C$ZPg20$tXUFg^(8Rari50{?QZq z%vasKG=gao<5Wic#)jr+j`8bh-V&OcJCHvd67aIPH`YllaNKhTzfx7?hTIDX_f&x(>?*= z_rCH>9*^siG)as1U7Vh*S`7AMF}R=4%M#uDZ{!8f%YOD~XJ%-;%}X0W9giYu#^>re zK)8lsz$@!WmxA`J&mGbv&Q``^Oe-tiha8jJC!axcJ9!arFPdyz-##2fVX5`uRFd$n z?xJJXONsRkiCGNp0g&w%E`3a14JKZ8$0e8TKWvWH@jOujH0?q6&K4ws$3A1(M&-@a z$iQvvTdRMM|R}#>?4L;YaZ+M#(XL^gENXz)BRJLX8Smp6ORU9_ z4OCvf%oA^-cIK*eNjwv(4AxLnO_h6Wgu@C*6|mzJ8TIRFbPaf1khp!H7$J7<35McP z$^c2ty$S+d0#H=}7J|OVDS)MBTKo{WIay|(hki7{2sP7t{@s>oIyl>7bdy1t@b+T@ z#p`i#In`k)7xR!+LXkn27^3}I#VOI=M92C&rlU`maH6nbyw#N|o5o*5h=p}hs*{)@DE|UC5`RMKQ?L$7Zp%CZ1^D z1cJI$7eS=GPa5KWt(Yh(nWDZAc4DIvFK&`Ac*q_GxpOj&u3n z&VvCU9pj)&#UxdgiHL7O$eu-jUu$5AZQ$b4=nSwDQ=JX4REq^sN^wY8mNH+%enkJA zCdP2NXd30vK=WCLNo7q!-~|_qO#o~@^)V;-V&`dVqkG%FU=xrX6xX+9n}^QpXu$o= zZRZ)EZmIET%;F|wKFw+MS=REek7mLMo8>4D`nn8rzMQ1+PQq`qj59ZoCQ7(2d)&FO z4s3QMLh;O%b2{;$3+tGxOV-+eiK`Ew?UYUT)QK7SR;y=Y{b`|2iTe-8fB(Z`ian{m`UM?& zP-%uTT0GncW~HZ_uIMH}(I1s@O%|>qu0~&A%`Y=Q_@EftO?nb)%*BcSE>^j6fP?Fcm|f|Cp1*9j0h^p zDaz!@#si|>X*Q*121nc*cx{^Pl&~JI=Niq(Y0v>}(AibzSLxFnrTDEv_yvhSrorVs zO1>-tk^;=E`QkkAD!;3AgU>u>SgJerm z`$+x2I8_Id;5z(8PvJO_IeR%a@T`+l9B(r&Vc+ zn8oSC7wq|fqG5eRg$**|#DafoMTXYSy{7FnH%ggjWW^|=EEffxYrlB^RadV|zgR|W zz4ip3!OdC3!t<$zcOG<5sxW3o*V8LEepNq^7R78xE;=bmOLK9r#9Lf;|NPg_QX6J= z{#qpb%|f)h?@Q1a^n(3q!GCa-l=r4(>5w_h;r`WXd=psM)&fM|&hPz}+tKio+arh1 z*u@WLSO}s5$pkn3^yV5p?w1LJw&!%1hZ_+($9o$uUF84!2YG=+j7OL&SO}gRD63>5 z$eJ<8zUTJ8Bt`mp#2&>3MDA0y#+0_Vffse*?cKHw@qn|f3olU+`2dt9yAi`T=}!Qp zQvm0fuu<|lcc?vH9Ad?thNboi;gY1e_ah5Fx`qVfb$?ZemCPKoa zah>L-XYJ-%WQT%CTKQ^Bl@SJgzH#$k6F-}OB_eo)kQYWlI*R4KD3hjSbJeDSb=!HPy-4 z8pXNMkH1Tztu;+aGxYy`bt=6bUF=6!WQrJl^RAF6sJ8wlnbcIKqBXxi)&~bH<@mVU z_SZ)=GEz^Kl05*_S~MTCxK5v2;{h~ohVsGN$Opn2<5KP8P-*_#lzNMuKe%;{<*h2` z;Czj?2sGD!p%6a`##%ip*^wP!UE{E;{nV9cMTezXJJx9`LJlV()cVt*Q~?t~Si$0X z1cfoGl88U`@U~3j53l{pg1s1L&26jdZxZav&i2OYX+LpT&F>{X6?8$Xp)c14l8g_l z>9bVM3^BBf;iH9;rv!*77=qMoe}PyC(dHm9?GAGLv5GoQA>!M~+{$pqoEdd*QSMp{8y2DL11|}~3%uQ1D#F|d^>$Hn@1jpe} zTw6)kL^WjUqx@c|^iG3@1@vY$*JLvxp^hT(Bpn6v8822mYiNT<@KVe$uI(fNOqNx% zAdneo?n}y^(Py)fx!g)Z6?(VKuH>4N>IS|ULLLKz0>exKw@E|QmRqScEF`x=tvU6K zchKE$wOg9@s=S@RU4DD~kk<=qQsBCZg5@!tWu?Ar6m%NY-LSK-SlCuOAFVIOA@)9R zG@ldIaE#;5Xn$T?4CT2AgI5)AJ&_^-R6T{gl11cI6UVP6jLWlXh!$gxY0NgCm-Nin{< zI|h4Maf&nEO7ThzYcvvdHSNYv$1Y*B;pSeZPHRGBW;=teIEbPQgk(1TDI+B9+QZZ? z0&0_+lbUw3r&A_a6%%%8NfD!cH`IVFq#MTN5nCl9 z`8Gx*N@ii|jG)w!M8p5gMd!)(7J)l5V2L>XLa&V3o-t3Ns92o9Pmo}Upa-936rM51 zyn;?mAPC$A%#nWbOfEeqtjTL^#tR~UlkwyZl1Lbm){;lHLx0Nym0$4UXgz`H`+GV% zM#k&&@}b@wN4PPJGQH!%n>}uOs`3+rJ$2Hzc;Y^M4|6%5`Mx5GFbfGk{j|{SIyJ~D z!mdQ1wbE9if@WNarG(w8|2Sjy%U=qCW?sOkFeJ3Lolug=V*J=vM1C+m%G23b&0I+^ zzl;;}B2Epw4G8{D@GD3B&0F4eJ5_cJ6qVWavjj!G!6Y2XklHaV2yQy`|9 z^fIAC$neZ?X~iPSzC_?>>=l)VBAO&zPcG>W5i#`f>G#Cs?j%h$@(=P=RYWomMKf56f>VF@ zCsO2_l_gxeH_}SoQNA?j*p~r6)!bOmkzt8LP&|5f$L8D1M5sMbbz7SkD)OJ*K@SdQ zl#*MBCmDjf^^%!On(W7;>SSK^;c$_lYHkhZ1hC{n)U4yjv7`k7dvYVQJ6{)1Zo{T{ z=ah06JxoXW?cZi>I(!WswyId3m^2~^vU{kBKYV&eWFO+eg0eP7Dh=A*tY~tqH9EU1 zQ(@&pB!qo5)GiU!D3nnlDT95?J{NWrHf(Xmrj!u{FHk(G9j$c>9TEA@UFq=Sd|ti4L;d}d2)BMm!r1f2!gO@Zfy!!w zy%i}%5Bh$J(#XJwxF$ow=;b_TS)Y$(#C(V7hps~nvg)vcA7otT)>PGa?=sf1Fbc=Z z-%s%9Ff!wop}piH+MqxzSkdDG5;$5)9Il3hmEKqTBiR#0Kn4!-jfDEFC?kL^KKm~v z6=pAU{ol4#?0?Ul^wK*TzD*Br{AuB^7RY%66}Ohpq~m4m)!bW9l`zUie|Ca8mbYf% zh|nASsK5v#{K4=hiwH9-<4%u@u$G`&-&amfhr&RstfNsfhE6HHR@ezGbw}44MHeyP zsV<|LSQL%deHTfcWVgyDy2HRH3X~>7z8Ol;9x_~&5|KvW>lz2%=SXfMF>Y1SKNACD z<1w;kE!G;$t)}?Yd1RxOl>)w@$6@Yg7oY(t0p;=(AV1wDPg!uRdC|D>sC$P+?qX8! z%xvPz6Qy*2q{uDFc;K-#1$|n z@Pxu_9i{z1&w-Lwl~vn!osTra$J}^qW&YL0)WgngXpXy-pzfSU49D}`S%jp)ZwE)b z*k0@Ujhh7IWTX}x#kP8LE&c`ZqA~+Y>6tG|W{y0zYvW-CoC?j1VBwM$Wv$xq2&FRS z1yRKF;z1r7UR1qHlIZ}7YgJ~!=#^VOn(i@ZZ3=aDnfc*~zD2j{NG-9BiaKqo0^2tq zA2#0BtZbNe)|~(+5-g4nn`ktI^@5Nig87hN?tNbXe;OLH6#)v;ie&l^5?ut!)`-Y- zAT>MzZbae4QEM}*j~HrB1L~iAdK;BhrAviUGKbh>)_8w7tYvK%{ zdI?smz!fGAGd=#o;Cm`4LIc|!h2R2S3kn3AD5meRQyLjb+NCN zm6Q>r(lpfD`Lt1eEdNhSxYdtyzL+o6d~@}+!Sa>$Cln9ZyGt3O!*R6u2ME+baKSnM z<1qxe&Yrr6ef@|yp5{_6vahUi5F;8uQm`=?h!%)C;|u7bei4sV>e2N%>$ z{=WZ)`11oYmHB@L5dkC3c&>Y2{|x?-I7yR95rx%I`tL4(#<4vV+;JU+);ACKz3jl^ zR{9`kwbjKz=cTc*Q)*UeJoi_(-1qK%bxOfovtL|ve~+agkm1LYDD0=i&ku>uiJOL< zOKgOavmXRhEicX%aqAD(s33&4|JbM9a9zQ8pVSd0QTvxjN1OFp&Nov3t{et+fIWxa)u%Nfm z^RCaT8bYpEP-)IAtGiB|#A&7u_S*+b6nv19)cC5{&wjpt4+JLcy75*42;aJP-2Z#6 ziUX`-agsWcy}$-MI)n3}wX&Iws!!T#XPW7w>xMY6SbpPZD<5#P_!$8naLbj-^&k-; ze&3n^ciTvOH8w+#qu@b>4+d}}9P^tYDP^bS%=?HAi9fr)$qcF4#Tllb93ho1|hjbfAk zMf~ya{LncOhs1M&CB-0(60b-HjiZeu;Vn^|)9I|!YxF`k&@<8HA|M`zI4fac$vL95 z)5P;A+7jULbTz&F`>*iVY1h@jhG_85KZU=Fy}zD4+7tmpwXh$D|LraY{cva;_#<*o z{c=dH>Mh?1U`It~Ky6<4?Z--FHc90vdfrXwGk4&Ecd_@164Q48F%a!zb;eOb>#jBV zdYk=1C^(2ywF3|&_fAvQek!aEuw1>K8zA0Ft)S zh;iYoI`HU{Zc85W0}>qG$t(MmPUX9%WuF%RUU2V-sQUrn0m6U@u6i(ux~M1>pzQlG z1W*a@aZ0)_PZId|dt{T0L7xaM{I>0$sbM5rsp7YaF1rz^R8#QGdH3OccPlDt-n8dU zHwxwA-n7vT!7yNUS-~WsG^J`u#%}$T|4nnHRmj$`w_f~lTI202&E-xRZzS8)`;DxZ z;2(xYP8kL2a5uh&xR!+sB6a9q&dJI+vBexSDoJvT(<-p)|HD^&lXRT=&YUwny;WNN zTzb7DMD0*kP5W5hbmk&SK}Lob(OJ%qIy-hio`qPqOc=Nj|M%EtO;Gz)ivfm4Yur|I z@r*cfJe>T%3hA~8WR|2K>A+zfA`Nx7#)?RelHB?0UfAkb+<<_fPb`;9E=nh+QJ-Vd zKR&It3%;zkA8vOyx~2k}3P(O~_!G zq6H!bTnaBI8ZSt0gjP-Dzv7lnVVoJl8`X%|!b4#uT4i-D_H3HbqZD7>N>#wi9#(ef zK{8{QDOwFv^x&gM=}z-u(PkhFZnhe3y)vP`B3~jArIlP~V4H%LD8HJ{8co0Y5AC4j|!8h z?y5qB;qpV8-&pZM18XQxW_NxO|*lLZHy^N)O?q3AoGY^_(Jd!t2Y7XMk%?R1s}d zu*zyS!ZYsZa4bDz4!Hx<$j^v>^_w=*9TahWbA2WnIjN5}JW{J0^b(~QJ`EoyMXnH3 z6aNap1-se`&$-Z#!2c&I;|DZrwAB;qRu}yER!ccs8cS7&Ars9<&l*h{FMbj!c)fnA zJ5Iq=U4ZriOu$vi5R)tHuf)H#pn{662`o``HK&tLV1PE0@o~?@@&^$-(eZ2zC3{U^ z;=*@3(XSt#CqLY6#Zg()sd4hyCW32cG$?uYju=YIwXUxrX8uf+m1*K9SztsCEV!mb zo%~_!?{J*xf#)&bT{Kp@B-?V|xgU0bNL|AvlLo(;G>Y3g zbqS(t&1I3C%Jf=ufH2S3oDgEuOA>J&#A ze9H0WVtLcllHW!vY6mR-QBRJf%}e;gKqqT~fsp#(ppL#uhc^)m9qU%VZV661EJ2QD z1#!ykPhW>592>RkJ^Vrk!{M8=h{9OYv7TroWB8!5LSP_J{n$kvFae;p>SvvK&vl7O4-j2;U)Ou8R)#3hH;P7|rjDw6- zo-ZPFl{GrSH(N=}8c=ZO6E5)76UeNIvNC`AwB!V)FKj64Y?EJcG30_@dO+gW;;T^i zBFI0T(YNSy`0Fh&V7emjx(+6zk?A#b1X)x*RC0w*8>{pN$Ej@C@FtMV*4dF9uu|Ql z23GS68~YU9W-rFTxn(7f#Fys~Za=N4(z+q}^d3ZFP>_y{;g(yn(8e^UoFz#>RY)(< zI@yXh{(2x#=m-60PkxMN*^5L?!7t&+#l+QjPA8HVGPNX`)QJ!h#cZWs|orP zJ#o>gk{g?O;@sU_dQ&&n<>3^ba~WJ0AUAQY=}1Gs)8%L9GBP%$6XjFND!R&3xn(eT zOZZ;LMprn0j#l2#8(n7HdfQ8&{KNgaLGa$0?*0g=8Pt&+@8j|iEe#Q^@XCeR$IXQ* zL1vH0Tr|b@!kZ?v)nANcM-?j5;GZX?120dP@Yb@|)7RN(b;3!t%pJ2{Gt8~KN)Pep zeFkAEI9^fRNW3Ki9_&OF;uYDW4H}qOJUMlnaite3Ptq*SAJ*IWCz;tmDOxTH*b-Pv0Bp`tYXvTJTO4w z_2-*j3>WW6I%Z6kfh9u}z1}XmD;~iOj^L7cNbSQa_#$&z1Sr@UcR$)v#N4*aG`=N^ zBn;YJx(XAK`Mi23cf7QS>KXCJ-T#vbDVSb+!G!Cfkb;jcH|>P))YhG0~8Iwa&Tajh)0?Qz+gxlLBLgy29|6 zR&Vo2uxdnNX0ZJJAck@bO5GKVe6{I`JMqU;LgoMB>Z}5yYWuDa!vI4`gET{TBPre8 z-QC?K-QC^YNT+~wNT+m33kVW&>$`alzJu=&j+kNZ>-xu9zqQu5vXyqsL=doK^Uf83 z6XVJ`gp9g?UAPpO616OhZqsaoX``6c2$ku^FWc?8=bEK(+3)VNTmp&>7_G9%f2v{C zOp{eTAoE_HGinweZ8u$BUC|lIT5>u_X5w0uR~Yg6S^R|4Jz-76q^iHow|aqx=Sj-! zn_)d3$u!qPjveg#n^h>q1@^)LIX?E5B|CeqA;a2dM+^l942!IOc)VvKVmPUV-=1PO z6jSidpI`E~!yQU0gyxora#5gNq3dVnH6gIpV>B-Gm zlKg7TLV22w>ZlRnWtW~qBZB&AX8{K%Ymw2WHTwM`rAUQZBi=jF8TNdMTz*RSJKd~% zXAO?OSHw{VET}QvOU-aRqDg)KU#wOZLCo<18iN=c`__VdVz67TB9T zhFDf0X=~zXgGQ@U^EUY~cJ=oqC=Ou(qbm=srv?oLxI&L6EIMf0+zVHuV3&iuGWp9& zXN0zJYpVAR$);i%zL@`UDfllQpH`3lAefa@sI-ih_G3@S)4?(%{<_9azgY+CH?|#o z*8r|`Z(>Qno*CUoMfQ^^lw_XchlXH$ z6CEFsQQ02AZY%a+?W{T8pMQ}|6~QneDJTwQo4sY2K@wNmSY)JAt-w!IS55zRX2csl zN$!EqN`b*1E(1zOUr!+iS8EkM2o z_g#eH9|Gw$&0HbuIjsATRiMV&@NeRmA27mxC; zUKA5*&7TH-n*_Z!Euj(3tXNK^Wi>+fZSHLnX1&sRG2VkTX&W?Kz(_i#`8vgrC; z)l?A^B2;z)iXut&nPo7Gt!8Mltnr;87scwwF;n8yt#1&DwjPZtNNOShBv;%6w^0UP zj-9Y*g+P;7!G=!cCZuS_l#odY^R@oO?y0``_5rzX2oI6WsuwdAL1 zj9pj>ubEO9wxFwOxG+@KF3UBb3#nZ%kqv-Tg$sD4QV`O#1qsg7~@d zsLh$tEjw6wWN;XdKR4l>@F84h1bQ39oZnMZe_^Bz=D%Sa83HrLuxMHLdqfu+vt<)_ zqvYYVNEXnNRu3^@Pg<$JZTM=OC(DXN8Wr!MS1aPL%=JG~Q_lqT^Zb=f=B4rj-6MG7 zE!pZ>Ab@I%D#wTaSIIbgM53Zi26PN`vkqu*#D{Z+mVTipa1aYyYC$b@_ zo{5@Fu4}Y*hf;m2F4d(pqp-N(rXB!MY8~JfP|5*+zii=;bT8uZK^Cxt6X=PsfV~8T z!ih3Ghvgpp9IeahFrzRke)fkEAqUW%T| zu6rvObKHTSUgYI*jFG_I;WkyIL{tDTf4kf|b4)q*mWjiCFrdVKH-kU6p2+ zr|b>ybQe72a8wlKm|8O9bfc4Pp1BRNs_G9k%vhx0 zI0(n=qyCBKJ&LPM^&O6>4f82@^zlr0cjCqquK&J5*lAoYy*TRfB$b7oP^LR!Pc8ja z9-SYW-T^tCW(g@La1KP4Rm$zow&2>)Bkxkw>G(5Z8)3^UvCqZRuahr?MMtB}y*F<# zLl8qQ%PY$FAhp?v5P4%d$~-+@_&kSL^@3co(<5;8^5F;}Dl6rQIf^~%s4NpI;gI7> zE;rGK25Z_jRLbD=UO>sqmZ$FWCE)PCyJ=S1%MCd$W&O^Br${y-cTryGg$%kC7cFru zuLv%gMKSePO=9gd)TI|`UR;e)EXICA8W*~pb9o~EZ$IqLvcd(?TaPYq*gvXuZzclm zWc*t@d9+|VFq_>kNjn17uJl_;WBb9y!{~!d~<@4rO7AL_TB&o2WFpE{#gh z<%t^tGQ&T~5bq>I>is&ZfYxPJgdd$oI{0xWnD;k|m^FB}f5?Ku#@8YZAACDX3=P&+ z!G>}D(endc{($m?{UbC}J%}O00rDSSqA~W+k+O@!=e>v?l|0oJlCm4uo7UL*x(VLtqvA9sX-kwUABVZK)#p?Kd&?U^8{#B z-XMHygQcd!)#{QBZ2t|SmOXq_wu8m}XD;cjK#UE|OD-=OlvPXZNuF~JV*}@F{Z`6j z5xz)ei5#%!&jPV(s!kINDT@`?OzD8Rs#F+6O2OEoPR4(gKsf%Lw>hUoGh1wic+>#JXP4&vp%^C`!o0hHxX|k z7{e&6^^-y+A%?qikCg5~QoWT#IU*#ce(^Y;IJ$`!A{atb;{`gASScUtqPnSQnNHez z1QLSmxNC{cr>U#jTnhuZq)ng2X)d~q!9GAHxjP_joF*U;!ue@V_D4q;F$2v-ko|K| z^bY#+MYgFW1(FDoW;X8t!%u09cfzkO3PzS#l3(M;LviG|MD)uK&aYU?j+9eNz-)?W zx1nZ!rTd80_hw4cG_?X^(k(WIR%Jw>JW$AP0K{Fr>wCTfB%X3l$&s2eEn=7Xd&7pJ zpuFVuwib%^t$S{sKQ1!%=x$eck6cVHO>Mo1?vrn5<%X&`FcV~ACj8RxkesuZ{?|6h z^>b`%qvU`cNys6FOkcDbFXa$*sslWu1Rue%*$nk7s3X0VXj9XS&{qLXtD^QOZi4!q zz39r{Uk^>6`mb_oDV+9x4Gk;nAJKQl){*0qd0~oE`|F-jn|-4m4H^GO_B52szvbiO$`PiyAt;N5Wrhn%=+IH#g z?f%P;v5bu-{{h(y-=sy-j%s#ffc28^;{hErCBes|A)+Wp%Q?Ou`JDjVzhmp{d%I>E zaqry}BlIv!=~tt`uqp*c)%4fidoTD|3Sn1}*vQSVM5Jduu&mxFM-3U%sJMn zx-~t>9y$3kU}@izT6oGG$RHIAlTSv14Hhi3e?;zY#hA+sl+`Ri_Z&pZJgYa2U3O*TCygxv8#y2Ic^(Hz@X!#Shkh zp@g3@DgJpXhVCH|rTWMf7-%zrsYa1LLlS7Q6tL7c6J_zM04?8umtEEOXwn| ztRo#Ubj{yOo61t0nl7)j8O0@RV62})m=yn8B?O^`L=Gm33?6KLe==iEk}i~9M;E+y zP5SlWat-DBXjenE>zwFc;Tuh}HFHVOS%F=z^3B6H^XG_Rq4MWRY_NtBLkwji?{G)* z=rZE3+!L3cr=r;J9ggZ$ig|f^Z0ClLk6E8WGZuDy;C?1^QYf-UbQy<`#|Um=HoBC`y>m=zO68`q}-X(91RcTtJ)Xj`J&5AaUyWr@z07=CH3vTc$$Xy5xO& zPm7|)O~=c@bJF)stJGRs)B&3t!hN!!q%CT!i2lq^RhRoDWDf7v#^*F+vg$f2$({xg z?}{=x_}hgq6J>0s6BPfx`LR|XxLJgB*+P?!6r?$V#&fAk$PFN4hq)%_dEuNcBBF@%Y5T-DEg zvZM>*3T2A35=QQNp1+ebjDqeVLjthS_&%5FYm-|FRr8Nto=tVF7%@Cj8;a$cr4M!( zls>+v=Y1!BxYwu(Wo4fGV)mwT*+-~}(B!XZ-*rEOY;tNVzsx23kFYJrW#c9HmvJ_6 zV_VAv8U0;+dT953K%M-xLF}E8lA6LS*+hcLIQ{%5Z}xXGA9-&&@wZv3YSc+*uAoYb zSA34gC+Ar*<&9p2sTE@a<*qB3%NPuLk~b~ith zp+`BOnVKBa+@}aiy0^=EducSKi?4#TNfs-cq4OBoxFNL&ela$a))M2~wR&nb7;y~! zc-0K|niu5=a@$gqPugVDu||iUi#`zcEOTPI*(DNlnRo^tvKe=Uaofpa_F42Q)D}to zY~;^SUk;p}+{5IBl4}8GVOiWZxNR<+(Ig!toG-iy$v*HXT^3D0t*rnpkVv|@Y5R=P zOQiJ+g-1;tHu$*TQOsdm>1u@Im_%_dW`1=!VQsCazxb*aP4tP<_dwuziWOA{CVFRfb0@bv+X|U5tEPtZ3nvjTB zm_n(dyc5q_?6l81^P}hAjtj_+b|8=AWELFWcLZH#8u&ho?0d+2ye`MmknxXWx_ic0 z(4#M%$vk#d4O2=9;XSJ$6D`ES$b{L=(_q6;_6f#+A(RfdkrFcW!~I}|0i!yf7Mt3N z!AN1YVIuKf8vibj(V6osT*cH&-R9~JN-*uT*;T?)?&=(vPyA^2yG|sNL2nE1E;i7< zBXwDR_tuEKHm`4h@!a`;ro!N z>>gpXD6|wA-*B0t|BiA784Z)Chk9E=TFM<))Ch%^gf?7Ssp=sf(u*VR_v~ZZQOtyg z&A!v$s1+;}nWZ=|u&FY^G{&DpA(r(`N`Vq)F0|1f_wQ48IQnF~xi*aloG%Ewcu1m- zuP}xU`oYzXoV;GM2`MDOm$2wLQv!>dA5vfQ6`sTFv~b>&{15-S8}F;9%jb*1-%G>G zp=^$?WNSGx$E0TTBm#^Ju_Q?221DFC6^3tbmljIM?<3pxWCy?hJ&aN&-I+SIPTXzi zF=zxb$H2Tj*_)qActrccW+i2`4eA?k*WI(_dB#Q|=n?aeeUmw?=Uzg37>Vyy2HBVW z!mk1n_rZ^5X60<80#pkoT!;%cgbmM1JzP6em6$}-4@ez#kYEK3wfaxuOfa2t#C>aI zs=~HX=$(rj79+xhwz7Wt&4x+Cuc~>F&&oKxaITJHe!?F{Mj$ZpG`GcF8YW zs6&-AZ^zYZmf!?@gjK?jjz5O>;dy`G&JHJV-{NxBR6Py>rDnwWhtcq|@te&pnk<3& zuX3o8=-=Wuvrn1_G_hun#pB2Rg_9{YVX?sHH+B>JcgbM!jXtV(Mc=z_9#^W$$(+CA zKs*sU^{OonCVCNubf#B2m3tN2%#b>P1$RwIEK$2KsBs_*XFI8+AwX7jvHvf4@)6SoB4A}a_A^T;>k0Jh}NgvUV!_1YvbwjO+Wvd*;on-NwO;!ML9Vow1!vl}(Y5noZv=4*i<^ z3Mu->6&(GYfBq|L+}QQH+T`~5rD3NKRLhUH82-d7p#^_?Rco)TI- z$_=|NTN7GdSuBMz1XEwEkilylVJ-?vo9I?qB7#IDMu9n7H(V7$LCdlj^HmX6^1Ud! z`)JQ5ge#Kh?4gy{uR4V{9j`~JI>Yo)n1A_P#6)_}t>_%FA`S`Rhw;QO5sb7r{)?KS zZrM)32>MS)p6 zqTc-9=l|cOvCR&#C}|(G!7E9J_arQe{bFvgg{Adh9-=GVw%YC>U$x%EQabzwIi|cU zm-E(wcmu;wvs&r5!F;+0wU2ZH6LiZHb?aux^6FlTE}sZ;QYpL+BmAL0}-O>5gce%EcBY;w={{Gpzt7873)fo-)FemVS(i zhi4DYnt3Uk^5B40qBRm@jJ_Esi0MPhkkbC~Lf(x#{?kQ6{89(?TRnJh&5MC*9yozP z3&*?buy$C|6?w!%C>-qm-9?6~P@6Y>`%Jhv)3d2d$2@*UV4UkSGiGgeurh9N zrOmtsOg77qztQoL?fe2}L97I`r4DC+YIYCw#ZZpba6$msM@r@VQ%zJUP+=5#FwS0* z=Ex4mjS!=mylz`(B981TD8@-SDE&~iG+u#h-_5i-s(+pSfRephK>Cug9aBKG**5AD zAFA^?<GGujOFq5H5oQgDPpNgnp>_Wd z6o>{;)7z-AC%qirjv?5$V+b)hFXlmfu+)sXrMjBAo)OO^KV7KZuc3%Xe>1fQtz654 zoYgU#-_qv5OBPP31wOf0YbomlaPF@@c8!1RGpKMW)q>E%Bl<>OyF>|G0+ENmZEIvs zxS3gm)Q`V9-(Ld$EM9TBDKgJMzN0FW=ie7<4Pd(=Bezm^M(b}K-~j@l{%wk4(*s^Fs3)OpxWHFfMB+V` zkev|6qQKu9pZ@*lYnjeDe>!rkq2fe;Hx$s84`h_0W0Avp$85G$t>s20hg%F6?6Y`x z%U{bs?f?nTS0r9{^23OXiGpBHkFQRPH9rI-5)a zC^~vCKgSV$PV_s@B;~h_u?3@DG&p(InJ0b3GeD`w7yru@1YFW=l%OyEfK6$JZgYeX zL_`LQ5`Rn;RNAiB{~6RB4k^g6Jw$AW{Rl45%45ma|Kr+~yU1LIQ62W%Om)77LW70w zOLi~d=d-)SU|~-;zh|Temc~Z_^d_Z7X6oxk!!&^;s&8P1sKO^eQ#I?o&X7-$P1D5o zhWhz^L@T^esIu!(v~5OpW+A!zedl*a=*!>7FGb#A@*>LIhJ6;i*IkE&0C%Idh)jM!;OsZB`(9w?KpE>3gYt(3QU=M@O&yb;Io zQhZwdP$BNXtc}^=VodC_F7|!u3jPd;I90pww&T!l*RYM-U^BY6#<1=Bv;DKtOl9-H zK2ljTB^JETRRnFb<~=JEx}QBR+E(Ghd;G+_N`WPY;zPaOqA2i z*q!Kgk}L^qy2`OjQqAQQvDXMl@)_u6k<~z_cYp@iEjQUY=V-Y0pJT!W@ zf4>z&D(7%$8peQMZ|*iwc<}A@r1nisz)Jx6y#GHKO0K!Av5yddpsB5kTjRn-w5qi#gk_T#X8nB-?MaH|+8 zSBJ;$#C#rcJ~$`M;FuESMq9c~7C!1;z|~Kwj6yKmKmbh8K%82w*X31S=*$IW$)P;`xR&-VY2#&V#a6I1e3M|7^ifI9KYvp+p zM;s|J##MDn%iEJEQbw++|E!FiJg%v_M`ek2UcktR9ZEgaNHuAz*ffJgA%(>cxIGz}%r1 z(G4RHLH)SW$1po4>soWb_uu`}M~=rkw&PxT@Owmo!5?>}Jvb%)zkPUf)rHDS&gKSE z7El{0!q&~a)Z5F0)8tEw9Vh;2$j_6Ljh}E6&fb^fSr;Ibpy^)T1tKqenT1rr{6vli z%=x#uXj<>n7*6ia-AhEy;(iG;yD?8e74TEvG`LZda&oK>XH+ z83N*kT~eL4H9?ibUK;u)F_t?K9d5BKg>*z1}h{W<#2SuN2zr|lS!+1E110M*u91*PHu5XU4~^v zYP(`0iA#)1z%PDQ5=qj;GI9h%y&VPLA4%8>wCOsA;LU4q&s&GhRrHWr(typ`sFC+t z#(E*kXm8l1xAf%o|MU#&`E_H~=2{EZQ_;zU{yon<4v8`HudDAQ;A0RWH)t?3yoXOX zr+uc3h^+m{ai_9MZyYeTWCM!t6!JIE{3XybrVLg7|t|E(0vTcR@Is(N+&#a$~LsRW=jS$ zv5Br_fjK!O2D9d-mtBl)77Mp%%Xd6xPIKPjFN2C*_ES-<@Gu}N` zP%JfK)KAjoZkr1WxX|C_7!hfno9PGrC3Gk`3dt>Me1Lc3wGW{PuNyf5(LS(?1CeDd z3O6Pc0^n7&wEP%DklJYTXD1sA%odKe0~CtklsQ%D+z7m~(~;3wqOv-q@Zf-7xDvwZ z;W_hcY2Pbqx7Ovt>22-)+=E3}`fkL&xKX)GOs6^RN&aya*q0~bq-*urQGob8l*g-n z;iUbjV1b;_CF)8mIw<^n=D!!jyj4>(+Zm7YF_reP%ZoA<(V9*6dT44bdWNL%o-iVB zSMH`_w|yWZ@~xPU+Ye}_j?&ozj0C|i$M$Z$(cQQe-wBZ z7OBnD=1|Y4>zF(g_*1mMBLO=E3Inyuda9+$2W>BUg$!1`t&YHlrg#n2e#QAIE&uYh zz7y`y70Gdo4}5RIXs`Eei6b?%FMjZJCczE_@4zj%To8OlsN zw%ECz+vCHJ5k(4(E3x)HB(Qlw_}*dD{Kwnv2TFl411__J$67v(_beru5gZ z1!sm667ICCa>m+fQ=whFHEh& zuH1^JGrPRMNFa6K`o~Y;a{tBYFaAiax{fvgWG%Ofua5RGIooEh-7gy2^1*rj1B!{$ zlCYAJfQiPY@2>#R!7=_bJaHK|+OJ%ix$52LN2JQ0xOW?#h#kzV*!(n%u}jf!BNlUt z`M-w?BYI0pJZK=RWw*&5jhT`aG$U6!oVR0+@lb8N8rRKpNuqs{M3ee`9Kg(*nQN?- zdHsV351Tmb+A1-&qEY&Jn)ENAJFFp7QeV3703T;etY{lIl#x|xu~9Sc-eq?70r5m* z_(Cmxd7d$D#2ExM1x7r9SN6ug4oi)UMU)BsNh}!efV$ead&3?l4MCd~$QCc8bw#d^&Xim+X2*h+R&&G1B3!fXk z`!9%gpAgc+5y*IAi;-K{_HB1BFeCBhNKHl1PZE#Z^HH$Fj$OkXuKdjJ@MjxKI<_r+zt$sKD{W(jQw~P5Go|`Zdd$x)8#ntxTPku9t zc&nRyo+IF<9%s1)7Mt11(Ic59uASwx81p}!OqLBYI5|uKrTlcB@>f#}!zo8Qqh9yP z)?u`oVKU{=t@(kcQC6$Zd-osa4$HB%w~dHHx`T^M5aAvE3w;EdPc1rM9}BgoAn7L6 z0LFK|jEOwYJt%44-Qnf4ad&z+qnE>H=?7aI3O)o*QfDPwFVFFg|9}TV#wTfZ-v+!O z5H`=F*Cxg~6GSJh+fFq0Yi;M3*Xb`a_`!nYm|g*wCcCTB5si z6ZIop=R@mrrGM*Mr$*nZ_xQr@yMfqeYtYxVTM`lyV{TEim}iEdxBGYyA> z8k0%_h)ulW2o&#L@4GkNmeuS(Y)!VYdcC^SHJ}}HLVtXMv;WvJ<#ry6tRL=tBZCxZ1udYa4(3l$HsJWZ^b}&Dqmn@ZP}iiBk4`Mx&E!;tlL( zG8PgOvSQOKUjM-4O$keXKzFUV0`{I1%+RIdb$!6WBXj!Bmo>78vHbcmK zo0Um31i4$!XUhvIMB*<*Ux1QvogBi{mFbM}Gml7s6QpaNskQm<@!#_8K+zd?nAWa_ zr|wv&vVOuUvvNAUT=*r~t3@*7{2)Rpe*np}o4pHhC^W z)q|h2q<(No=J_c-2EY)BMiy1L&_YE@@aDc@+bF@mQ;_yK!Er2~kCgI6Q3gfQV8Awv z4CfqAXhqJ@0iyhqp;)4=PsuRmD4+^2d92DjA5P1`1Fr5{MsU_Yn37yw@v-QqYDbvKNNXdPEwp&} zr8WXcLe0!-b!|%~D(?|rrmE6#@{g51uOAs%!i2qo_YzWzZmHkJ{p^&#=^I)Ij6{Y@ zKXJGtr61F--G%`8VlFx#cmf`AJ3gcx%>Tr*iWPC&90L;^V5$ffzC+^D1tY*7-Ud;w z1b1I<-&#@vOr2l7KjMS9S|SGxAaw7T{+*fCw6guliR~L{zSQl6=!Dm2jnt-QPfK7M z!N6w38r9UUunc6Ma~9>*5$-iG$dT8?M}$$v`4P#URl&Hmwan<78Cn#hMJ8?>k*mSi zCaB#}sI64UP5Wb+y}b0)T-SWf|NUXTvJBmH^M0{rVN&+a29(arubD5&NvnBfvFv5g zFss}Y!#o*Y3q)ohiUflhx$51tCx1u62%?p;1cP$x2!jUmQOL|yZKl&~Omg8AdKFv5 z#^^7&IB88e0+ep8+v>{jOu=aP?nJY(%;D?iDGI6TzfVT>ENCpI8l7VMI*IkFPAL z8UJ*G%0oRi+|o&aU>PWWySz__U$S_u4sh+L-A!*|FU> zoS5_M`&=A7Q7*)vySSgEz!P>_`-4!JQg*nMvIQiS4U?xqS_f!mi3BSQyh#qEZ!!~S zwkE@UQnRXl$r*v~SRtN2(?5FNkO%UMKv>RTu_Aqm)pW|eYC)N|0|tbAw-CP?o*y5p zv$mC4QbI*7fn1~zHB(;#lR#p@af~T$#(0MoQc!ZjGDMNzzGv{&BA$bSSiGyV**-WG zHIY$oBPhf~3_WTzQAP{Z_Qc2pN|Rr%4)IWmdhvShiC#09P!r2TC~cepE54@*i9mNa zLzfNJSMg!OHI_bDLXVC;xDyPeS+H{9o|~N6}2mGmrM(%iXpW2;#;JrIyn* zNoi%XM``lPBS}a&{8DS#4btstp_2}1UMghE1AcYXMXTwiNlr1YCLmN)ibuX6W97x7rOnB*!BE2RqE!x3{z7kQ$4#6e z;rc%nF=@>X1DC0suyjrhME*f4tgW7cNtOsK(}PSVolpGhKlW{P=Y;hN z^|BH6VaOft!dOmveB&KdiX=u8mnfvCEe*p9)7Mb_7U7k@fE;JYFl0~&$~j$wxo{*h zCVx_HF}}P#xX+pX(tA~s$=Hw+*z=| z$4R${@Nw5m>)|otXm*?QSQ5LJOXQ?y$C$S&IO5%Gl;B!qxu;I%1{ZmH9Q|6E9 z?L3_(x8cyxs?)a7Rm&kcC6)0u9M4>dEADfKGExP$aq-zR(9?Sr5=Gy?lsZ+l{{Qlu zoyZL=3oOP`f)laLuG-!%NE=Ao_D?CIl#@APh{Z~anPmZyF3As-O^HsT5FCYb1UR5< z>g>=z=BHP=A4$!XVjlhE*!L3C$3HiIRNG`!`#pW3mZCi+5)PYVs;o9yY~9DBm8XeB z>hfHYXcniE_y^1^{kVB3Iv;iX9?XD|*)k@la|63~Azr#jgC)~QA@md;HZnCQICG%i zjFDq#ARqKalY$8T#r8V$j9ueUl9NDsd$P%fu8f$yZQ!iXd<`RKW{FA`+|!)y!qQ;l zmxdh8G8j5Nt*IKn=OA~ouNt&a;8JL$^pEuQ^i^^aH~d=70Ug@uJ3|H5XNaq7{o^>!7D8x875 z%`S;Q{zRbtl{nS;^{pEVBL}n=L>A<^M8BZPmbPQmc0wzPegYUWq8%FEOxg?xGW(xD zeviDmu}o>)sAk+w=Rl?}+Qn841@zGzKjZFlDOjCH>&gZ-&2cKt-jVaERXDh5%Fp5i zH7`a;{`B=@jwFzYIctpc(rB@~Af1b8k!N~Y2)AWHl+gb2PteRJ&JateD}V;uAS3g% zgl&?D{{4xL$-7DP+mZ(75apR@30{Sl?i5(+9O7slDkKd{VaJNPZk$Zgz}8PP(z@d+ zx=jYDUmZL(QwW;=R(;$sAZ~Wi5jtwOGixfzNtC>M(lo_JOsSKNv7K~HG(w<`EHx#3 zZbY~~#?Ka%+AqiS;Qr=4x*B>?MV!p0GPUhMA>Je9E}KQwDsMnC)cUd**hN=SICO$F*ysL6rO){H3-s%+Wanm}Ffe zu|`U`UIw*~+Wj;gRO-tj%^{Qi&Sfwd0UDtAP^_fB87CG6R{Gg$ml27{8 zB0c+*t6SQKT{s7W#2r)>{^0O8Dm0Wd&S~x1_SR)<68f>*^AG+2I29eT8g(?uHo~P% z+1mbN8slBUF2=gY+iP7(^!ell3&g<5%$RS<;UCj+b9r;55^)2y)E%QHiYK{VkAJs5 z!eEIN!N5@A61xmVdwAY|TMwO6Steg60_NyWmYRQVOlZln1PPQ0{c?A6A?A|CKXqQ= zsH!NR6_|o{0DrmEkMs_}B)^p6X{Vm0F2*6ik~N=?NfV@b^CTdPJy0M#SW%Je4mU)M zylZ8?X6UV`Xutv6p=1p$R#cAH;-AFJQE7PG{10Zidp z@f>>18<*_q;YZ2_i=&izx@;T;WrcGDKWCt;g|C%iGFD7agl{J(!c5qS5(Y6E}e zgKhcY2sPiHm+9R>USBu=$IumdPL=lkGmH(j{STpA`>(auX8AHhY%OdSL0YXHbtk>-e%&+}D%P>sE3mwbu4)D5cVG zJ%|6v=AzzC0s-)ED-VCn39~fwRigV5QgeUDw!d){$)!UsHg?hQwO=Kfot=gCDdCAb zsWWmfo{yKN=)l%lU;7i+2p8FB{9!?0 z=UfN`W1yey)&{`20t@Fez~&8f(U))u>B~|sbUOg&;SpF0$0^9XQ-Bj#b`AYmGm;`r z>bn<$NVpVFB9=>l&+qy)7=AJo<);VQ+JEO=T<#V*6aKA49oao5=8e7fkMw{^arD^R zX1kn&=c+8b_!XGQahm4Tv!Xp6Ef^4Tv~fTE)&`cLU(xxFFecu^t0o=6?N@!gvK5`w zpgxmo8C=33+17c>VBqA@w=qe)T>X2m{O^8V9R1kzuEgqA4_&nQCU~w+`;kDgxgDvf zqfp?B1HK_v(hP`bSmeBjfec`cT_?1T_>T+hu>T^IBQA9P$&U&0$r*TyHhH}r@3o+q zuT+d)IRip1PL$qKu#c}3oVn~lnL3nj0<46{!JmZy{ex(>>O@)GL=wy-qnlG-f6n>u z`myiVH}$Wxy*JClR)LD!=5XRweyXmV+e~i#ijPKrtxPu`knk>z@P=SPaxP80BFGfl z>ZM44XpS3OkQ15Uo#lKO+3VgZeVaD1F%89*yv1o%MzzmG>q>ulaSHUQkwQY+J^fH3 zTMV|RmXe>FSe8gNFFB5-XJX~dVo;StY(~}B4xpEFh2LT@~{ST-0#Cm;iR@T|PIBiPSmb8sls zG2X+^AlNTNl}5nmO2*Lts=X~Gel#9svslGmXEqD)aNH{UPN~oP*|v#(Z$a_V{xq{N zi4oKTL$(_yH>VYl`s?MQ$y$sBQyd~Z6pOzf15>8#x0oh45(AYpa~e*r+r}GTPhcy} zgQ4R5vq*h?9Q_sB(clwcLUFuwS?sT)@~(XflRBz>Oomihay`q&l-$lhBmi^=9~Ww@ zEvOVG0&;cF`uBj<=)M#`hW}|fX~mNy`=<3Vo#>>C_!l%lQBK8)>D07+t(bfCt z@R7sj+qsUpYjWi@I?F;;$y1rP>j0geLjX9l@@4Gtv+p`awiG!^OkUHUySH(kK19N4 zYH3)U*pNvZ&&U4z-#6dy-fkT~bJ+WKa82wSFhWI^{_!=fzhiSpvw%4Sa=|E=Vk zN4>GhL}wZJ$XBGPKH^xrSO ze%aa)nV8ZXS9& zgl1gFqo(rscW0sh%7C2=07L+9ap8HggVRSU$SKLTqA$hRk!-fNTnPU&T-xs%sjD6~ zEf9M)xnM-1Ze{^@SzJ(9cqq8fA!59r{RZ7U;pZSRIa{W)fDb;l2+lBkN(^s53YPEr z<<1i4*Rm*?XLZ6RnE*rB?iws_C<__^VZxf5;si|me*HD~Wj~V0=eV~FQ|6Hu@8wOs z0;GJIyO_$M5(s5|r_Pt9nteO=0kn>#!Etl3jx=T02z#Nw-?2EgFN0~QIBQEZC^DzH zq9fX=&r}VU5jSWE>s^^yycjVd7q}%p{TjHnw*hJjD)m4{1B>Y5Ar%kxoO*opiIHf- z0Op0*<`u6f3iBh{or~3d-tH?+4zyRrjk6D^B*uOqX}TA@veAiS*!n_gm$Z`HI}LG&vAUkH+3%ojM=Dq!N<`Vzb|@?jRmCS^|5X% z*?!X;2+AmVwMs}&x4@SI?3kLT4!`gI{lJJP27-XZz1-=gr>R;ZyuZBIiCapSYEfCT zdu{e2SbKWQ+L{%+*Or&@d6I^9_D~WVA{U+LR2erxc_!+R9BAc~w`3lqpn}|l#3i+< zaEBbuwrRo9=2S~AQ+TuDE6nV8y$_k!mN*46lhw@?JfO7U;K4Mvu2`RK(INg8q}kiq zg_Qt;CWYk+3Y%3dNdq@I>l2N&$@9;J;ZQhGoPpS{DRxQH&^k6mR$;h-JDy%+PrJgB z^hn#zFw>{OWikAT2%{%ds849L_7bJhNwlO?)*|u8MEPP?T6dOP2~OlQ3T8|sU7>RP zvb*qrVeed3#sWsH78wO&KQ0gUs@@-|l)DG((KjBC2RJkp^VF-o&~&Sgn~FBdmf_|$ z+!VSk)ZXW!VhXtG87BV8qjhe*+OZfQI#kVRxs`8{1~y(yis}VGrk&7yZL50~U7q+& z7@Nj|3z3|LExag97s=wRsHKqeHtexQDst=U3709vH+x-EVE3a?$)e#`${&L*i!Zt@ z)!iNBfL^u{(Zp@<=iwC=d_`;BnR>5WDR9SJBa>=K#|+mA)^Zf2JP>PQ>KFd=!exgq-*Z*hD8=_M1$nS^;N6g0a0U(r#4`zx zM!2q7WI5VJD3Mu{Y?_EqQ#FV>~si0)T@fHIr@w55)Y0&fq32j}-R09pa&A8nXFLu2ad@zDuOa&?KwMZj zeHGPTz=v-jlJ|iyG3n>nSDRv$kmyyCf}z-d1Q9x{iFVG_9;BbUdIF4)iqMA@%Jat+ zctvHIEVwE;xyi9R)f*7#ZZnthMSzCJW}~|7AT^!dslLa zOGkYFN7h*e#TBhx8X9-k;I6@40yOUKu8m8u5FkKscXyZI5&}VjySoJsZh;_yFsJXG znVPTWCshrnyPC7l-tSuLSp_RRI-{Kt`$9?{qo>2xXh4Oj2Xp?xD;+*rxpP{nyuDIB z7@R4L7;!^_Lv}<>2@B7ZENa7Q|M~c&MI(0t2NSQ5zxwF7+`x4h#g;zGLh841DXE_n zcLmMZOp0}jo-H8*S*1tpE1T^TtzReJxcPC}0saGW98czE{H=D|k$kyqHEUab%Pg0} zv$D&8=fUTq{fe4+c;C~^LSq{AbL(Zlrwu1Q_%_Pmf6QGkemUXcy|$9Q<;!~E&)?(f ziN7`5&5{~QHZQ0|@YB6x|HdT*(zA=)v~+cDZZN1wRQ$dA-QP0hxAhWrsLQ4~ zy=qi%ix2H}N($$6@)-o>zN;NSmBKOyKMEKI7%>e@I=poA2kI=dm zrb-RS;fXa!?9FVC%G?Ti$R;@>o%T=DcxCtDYayEhN!m}5qh(vZCfq8X!6?oWXdv{* z&{~+VAw0{vn|F$wv)06mVVk|{j+V7qxhRUcyl>;Bw%>F0D!}Bt6MFh3o_nH?;?!k;Y>HgO0eLFj3w3*TQmLS_&>Yn15 zS=`XYT5j0+SvqgFC;}#oUQG}*m=#t?{ivMTpNbZ0EB^N+^jLzuyJi$AV}!92{&a`* zu@A@zB~=xT_96`R9;Cc_H%|x=l)KZVAStt1l9FQ5t&r_uxyHpOC2tLkk%$wWM`TBl zh)gJFWikbmpyXen3Z`<7DX56I%Uz4!V>gWOnl;@q@5bt-o+=Ta&!5qKWWV{TA35$+ zrtX7HE$I6B+~p6jB@e>IoRN(UsvNA#DuTp$jBZRySYEJ8b}qklrn;*!PwgkGwTbmR z(6cV*sHYj-*?Ma{XG~xPPsK^MQA4e~g=%o7S_JQ}NM+!__!KhI0jE<4tN37mO!rP~ zV`$+PM6Lrb)h^^4n_P6vuoC*ORXojH^j8Ik%~f2)1{^s&uHaAe5av43dvsD!HC$P7 zuDaH*DxbjS1USHCUI)IpYRwj|6Hi{y+9OmrHeS|{p?ZQyP#G@M^Pg|OP{+j z*7|3l(If%do6)5@*PGbQnCsTpbs6BA)vO)H-t%r&%w(k)DNkZXXtB{VoZ%b(v`6cC z=31q7hcqH`MJfy(s9WIyD07haI!&eCiIKr!ZrqEyf_51EwIMa*QX%pm#Ax@Q*|fsE z!^kqu7F-j0aQSW~kIURzr&*JqPI39^^>2W@IA7O=GC1fELUov759LWQ(f6WEp1v10 z8Sw}&NOQWHxVo-r;B$lDCsz1R@85I%>k|iPQzE2CbIw-w&UNg;*30xDX(7q*Cskr- zH%2_INqy%0D1c;~vlC3jO+H}=TsqWjN z(Kip>$9)dkte>p5mv4sQGWazOuna7!fnjz0x9TN=g4_VaML7<`3)f>QODnve#qOz)vw77*@<_Ls0Y{!STp~r+%(|chN z!U~*t#!cG$v$ut%y>T)!k2-oLcdoC0Tu&GYge*3rqB7ua@hyIRfebZ%x;QabGhE=x zc+pQ(`|d9bNuw}EP|DATQQzxf9|o*0hzZ2GnrT+s9_ZB8mj`s!uX1H8e5mI4Jhmst zP|UzqT{D81mc6jIrJk`nLN^9hZpqIoZRTGU9~+I5uc<5e;5!QOsIt-P$%F`jXu_I- zAbOcAj`v+Y=E@7ys$2GBF_W-WYdaFb;{H@2IQZ(yoe*q7qL|B|IYvzz_oDg(*zwz9 zS$-dh=7G0tS-*!Ci)y{j$x6|nhpOb&UP40_PC3kHV5TZFTRw1pgz?$Qq0R0+M8sp&qeg7P1GP3J6qx8{`K8WQo;fr2IhNa zmyGGAg_p4gKM$4+!Ggxh1(S*3r5fsne(b~pV0X$-3o=M%=ghq~LsF)`KSf|j)*P+P z#jO_oB2!5m4r9^s6=SrXuVA`n?tw(CDk?7oB-a2rh=0{ZT|a}DRZRS{D@2Mp&#A2_nbq*c6dR%8@^lAw2OVu4?Z)IHHRemW3A8q(EarvP>`}2NM3x<|i(Os);CK{@^phQTKB^jFHP!Gu7RKlcJ!R=zL{YtW zI6BeS&HPi~fqxJ+s1Y!Ij^BL*AD>Y2GE4cQ@vSKN&rfa?#^Nn=t~-Q`=t%|xGY&4n zJ%L|9KE&u7e1xc|RK}akFwui8q2(E6q}HX}!V}`#%$jQ!rW4H|dUfO)oUR4MLwK5~ ziZH@fUF1yeKUOXF{QaL^mOq`J3&4(~ETiG*WG2q0KfO*}`exn$GtlW_jCP`V|CBh@ zl1X<^`a;*@wnKu>Ib4|(+83#eF`%1QkP~Eo(N6^8rAh1ERA$N!0-oB_IA>l$ zdnA==NI@-{$PI;gfQ&4N%+_Qyq6`@h?_H9Rz$ww?WMtP)7yf$3LS7C>mfaDZ3}mIn z$bKy}iAhGXk|}vv!pjKahcJzQ`vp?h%obc)bdPy&!VXkp6Jn!Mxw9>tX*ud$ym-@UA%3_~85=#OTT z8ff%iRNgP2vrvs_W8inF?!Bwum58rXb?eXr>kT*cwfH3O#(Buq`$Y!JWcU-XQwdtG znH?n8n36h5;kC{)$$5U3**iTIlW?HA6JThuxWFZ&D$FQfz|;_agsTJQBlD)U`*Tl} zpR279^!Es6+O;!4Y~nN|yRFmyQ4xgQLVLQBYE5u@lnDCTU&XmS5A=1`X-V6bk1 zT$Nc~GA@?D@50fyGDrg*z#a~7ej(MO!ui< zSsO)kvur&eWGQBrlitqc98)giZSXlAl`X!n%P2Ctv+q{xVwnZg;(=y)g_Ss4UeUht zPV5l++F?lif?s>DqD)8L6vvE-=Hm)aJ-YBGNabxzji8s3k69Dd8XQ{m`dexT*Y81h zOJ&pza5xW>&NfKGv&sVS{AeHpYW|q7lhc5jS_ml75FpG4oy9|cg!)rLig7T4({NVK zI?bkQ^)Nl)l$^X@1vpzKTmQ!|7*PWGzs)DB4H_q8d{9ony=gFS@&EpF-YoEIPpw)Q zK4@*!4c^GlOBs_(Q_Axj|KnQ!&!x#Dp`(x%6AkG9q$rTb`}|BuKv-fT z&RXdD4)%-^kseyV&2{_%UGj7s0dODcU% z9O^|Uk!DrJLgVlQjce^@nxE&(V%PF*yyXqc#X-;UjSzK=UQo7#6rJomHV$S1J>27m zJgj`YMB(#t-9wQ{W7L9K-$A>t2(N3rSLq+?k_C9d{}R6svb*fa?)Mbos_JrV_jHr& z##Xr{3*nW2Q_-YCw|7s#@(ZK5;aM#McyF@>^6t5Z)k-FjxgBehV1^sB?DuMW%5f%V zxRr#M7z6DU0L}2PTb;`@tsVyaug*GKNJnmQoXkb2+l;^~EfM<(*`16W1!L$=(0dQ! zM!0J({SCMMOnu*+v|?P^43@2Dfo*A)Bu{r{ydG=PS*X>sMNQl_+QLEA1s}diYT%8M z3E{^IU67|?o$dAbO~b#fW|yd@hO*pU2OkpCPoEB3EB`fINMV)!c)tJidfz;zn##bt z7h07(BEh6~&Awwe)-4D_xW=2zDcJy`+0+W5@+=zQ0zp9Fk82>^;hfDgTZyG4mf;Uz zQ-TlfDA92Khu;WW6?vr#VL)H=)?g~VG=Czin?c62f>uwl%$eQsnWZw2=d&)9^?a71 z5cJ_mbX`1Ba;W^$T$LFLg0o;{J|p|#QTK{Yp=EKxawdm1^u8Zu1}Q%Em1-Y2$PVaK zz)0h}2$GaQBA1!RBu> z90;wk_Ygmo7AtjF$D#)x_tHBfn}Xi`1#sez<7O8Cf_$Gxd6g60i)8ooO@Vjr@F4 zw4jSoKe!0ZB6Tc=)BmU~0J?B==KZjZ(&5QkLmc|odcfIXW4JbAF?S1;#)c-mOF;h~ zl3dKi*2hcdGz7t1KKKc~P~te5Hfu!4*G+HQ+@eO`y)%G^Z5^7xFzPJqjt$?^;rSOj zqyG9$Ww%&;7`lV58#=m80#HhpGa3L*a)~sH`USwK*N5n2dT?wdD~~37rlI%`-J+gi z`m}@9J*u;R0?}Mhyz{eX>5jEDfTvwgVpJwL@cF(+4wd_&yKu$*9X{+bVVNM%p89+n zH`#^xR}P^{rZy(%T<8Eq<|snFuU>Gl+UgyPwJtXIj&Jr{U}^xSyopV0Ir}#vuFM6d z_UfvNB6wbhQ(WP$jy7*W_Rt%}Sk9d}07dv`N$lB43kq#oKWNC2 z1L32`vrLQVn{V+K8a^HKd4h<=54rw@boY4O?!EMwlt4|XM?yp71YcoR(A(p=_0Ymm zdJ+FEj~o_t;tT+cPUSb3^evM!=`+jeKFTFf?TUB=+;yK0I6HAOvD_2Fw-~%o+7T22 zNskoRGINS0VIa9avF&;14%3ehw;WqKm~@ccyoO+`DsP00C*5iHgFvc|r0 z4UH;%_?4G2-yuwU4c*cTOAGqEZp?$#9DBbSdtLs8j8iC_77&C4(ZdcWS6SfdzY@Ma zTGj4ys=1L>&>AL{?0EsY48UD8dq2t7jGPj6M|SB||55BJ7z;jf-c46qifY_pOshh1 zj-L0^<$d>coGT1ZM?)212kE*7iA{ChWxLmm>35`>g}Ka^hI^>bw4|eHV%L zKncQ*r(qaOpL-k~S1B4Y5?>chtsU;Y?xO7&6G5q zWw(9BhdL@q8c^X-rKKj8WKctAb?gi8z>oinj$5ik52yfto9I3 z^#n8^Od0))36oKd<9|9xR*InTJ3=$P0?C0$orFepc&4}}qDruNJKN>NPBgIFz$xZI z#VIx6vJ|-HgI%n$JqoB7 zIJGENgV*(n)?ZoYwEPaFo-05GzV=JKij))W12JZe!qeP%Fkhh(zN3ynruwpPdn7r8 z2|uiT@Qv>R5KRvohDaVKELtRT9YpV$Li;c5+n#XV>sSivPB~pAIbQ&S0%SH#44=OO z9n=lY)QM7zZ%}+bIes?WpHt${Sum6-OhTnqk_ug>q~VLJjg%Rm8W0h~Eacm&>*0wp zX!+uy|I-}WC(3eoCXBDIh^KdrhiSOFW8Wb*1Poe?7PoNZ$cn5`4U~;q0sWOI{euX& z)hb*et-K`6SV67`qsU|x;NPbQ6e@-L0tcT)qxow9Vt5u864CMT^zVqAxF(KTI7+~8 z-mb6X+1y>3yD+?(lLLeUO`XU^rGISquwJjQn$ZY4vG=2P5d{l?!d^)4*`1d0pU03i z^IM68%&$ll1%bjP&~oRpfU$D@31)>9DHZ$W9d^QCay|~f&!iS(u=#dxMb`f0dd?}( zK#z%cih<=HA>0z-b70D+86?qdSkq*1G9t4Dj|5`C^&rtH$YHfcvY!I+CljX~L{;>9 z01ituQ9}8M0Lc;pCLiZ|KlzWP``eBL3TCpRv@9jNgx-$%5g&C9AH(6S9A%!#5y)K< z5psNVxSbT7lYg>-xI)Z@YUKTRID#iM3hRUGEw;kVd*w4aBk^X3zmDibU)R5tr?s$- z#qm!XbZ^N&hRJ1~W%K3k?8p$Az|-({F5IA2XKRH6OSG+}I3VROg(;JD1DUncKL}gU zr)Zb1eDnxj@38107ZSNAHp!ZKnizV7SSr`#qgpHhiQ~gIGgDvu(7|j%Anzy7h3!!b z({I+w!Ys|JE33BRCuo;EtS{ir_Hd<576(hXP=%!QZR41 z?%pJN0rI|O%AoeWU4d>|%>!Yr_D9rms8;~U_1t*RYCWf_j3Fe?3e{SZ3kK%(<+B$g zHB+$W4-kJjY|6)8u-i)fbjVXMOJpYvw=p+3>n13EYB22wFevf}JdqbQ^~V$tC1SWZ z$w#R~6CyaT6i6AKkCellCv8XOZ#DHpoAk|)v1d#;0n0rBFUn70n10)v3e$6X>;j;Q^Pf-WyF%DW7pEXga}$^0}gMt0gy>1xchkb zF=PQ0|CB@Amt@P4+l8+vC;OwuHv}=VQpYrhysaCaaHPK}?YG?jVF6T{ywd8$)f0Xt z58JyVN@daGic$TpdY&R_l3`6MaV3^^MHP-avWGo`pk5_HqWK;yXdo?49%Zk(r=TjZ zkW}D*A#9JvLV3k^Ol`^~%s`-^h%2HK)#yeA!1pIxX05W7cu<@eq|$_*SNKvlUfFHf$Sd?~qqKa~yh(vZfEYc~ zy4M_zfjAfi5xvixEkzoi1m*LO{h|F_m8n(9^fJ~L%6|p#yE$F&3|p60KhE0~Y;N*V z9#ZffM)nBRI80QxE_*B?%l786JrIW5(ArkS4(mucD4<(zyuD{%8=MX4Jr3!;xw!WX zcIWcYl`Is=>eM?(njYx>(^?;}d1`8BcMxUICoO;#S;@He=IovDh5jy18j%WV-#T$U zh+oI+FXfJgSjV1j-M7sUuC*NO?M^kV%ewED5uEW)UJG_EJ9N&D=SwTR!WTcm8Izf| z*t2c$x{r7D#$b;ys$Sazw2E?5CjxBNU32&>z+hZDOKTmbQI`4i% zDwUGqs)kzLw|=aQQ2|~#Isd88Ev{>L1CG@qV@*Mb=^2-qmVnzh!p_HSleIiVNJ-D5 zK>MOnVb9VNbrFkaDiS@Q8xm_zg1;rd#0@#Z&1fHB`1_#cxaS$a`p;}4k?8wudZmVQ zZbM#=@^z+GAI0xGi+={ohXS<4H(#D%t=BR@WDE?;Z`NqoUhW_{_?-MjOW{0* zyIYO@!TzkAUVSF6*X*TJFOAa6ahGs%qU93^8$E?+?Sf(Ouq3r+8Ks<{3uDaS0Aw%s zMk==_nufDLH8W8nM>`LA3!l^}9y(1J`6v zK=XQIPS)7RR7u)xIMsIg`T=9j`38%Nw-?| z>f%^MHf(ZtkI9OqhE>f&8ZnkZC8|xtOJVoHTD!Xm!~qB`9p2XLKeh}b9R|p#tEddI z^yB#(&$QQ6$!8ZU6QicPna!)?3dmK}8QWZ*VLU?DoDiw3t?8!RLx_F{IbKOhXS$p) z6n~?_BJ0ANPef1%6Q05B_pz=KZKEV{9ZQaSU-*Rcz0z9y3YI2sPJkC8$ZLp&5=qwb#BnJa@yWbGKOTHx|>J%^QlG>uT~w(**f|4LE>)5zIS{j zWcFYjRt|TrC zn-&L%{XS1zwX9pr3@Jp>@X?=SZp6457yhdfkr5v-N3}}B0tp6w%%dhr!b*MMN0Zs zcL)F*@P9RI_V32etRFOvdIa!T6XQ$ zeXq&>aT_dF5Pds$2u$mv^-2sl6*UMPkuAk*PwJ*4YvXI{=_MyDUNHo7(8eG74pZqw zOvT(y8U<@Sl(S%u;64)4wGU`nn0V1}l6=_Q|J-TnsNETZH04Hh`8!cO<4tw5yHw{r zH8RVWH0Er72$e@fq-%K-VNd@3^HT<_d-5$xn7B&bRm-JJ@g`a#8ppU+7Wv${m81o| z#ON6s7<@`%x60&Lc&byY&W=v@7+%f6+BYzZ)kkdt^ECX9BSO+zc-tx~sj9FAGXX*t zXbRrqlOCJEWsoox4WPqyR@qg`QlZqnpZ*1>&@n&6p}=zrSJQ-Ds$ok%cm|mKTBqU( zMMvRBwv8GWy1u$T844Mr1(ot({#aPpEVTKhQaEJZ>p^YK!Y$(%LIlStmSB}$I#(d6 z{xONkd?I~-^P@J+BSCFC&GOtRNZ2;nQom7@0%q`1{h8Hr9-PUGrrFPU*1SaMmVykX zP~qa`O;O4W+IlxIrcKLvMcG}Mh;y~QZ`1DLEjt!BM+(vc6_ATPYKIS>#SZjWH#7&~ z%@;xLELIeD_EZUpvCA(IRpt_jqBdq1V_+i!ET8~!FI1yZT@xmaP$TB?k*Gaf6`7Jg z5c>%QZ>OR@i-q_GMo8?ZgN9^_sj7!83>vdkdGrQCZz4Oe-~f&PBR?;^?T4!zl~4`J z@CqjPmn%N{C;dc06i$e|DH%#&^&CF24Eh>1m$QT0>oM1|)ega&ejqEXTL0VOs@{LY zJ-)`+E3}{p4jv^VY(Y4hZ3%EQ=8;X z`5%WhYuD%68-+agUWb8Nsk~DDG}uV6&DTo<3+Z(VFVJ0s%tN3#eT@Cs!oq{MpK_=s z^C%R9HD>b~v@age9w|7N?YWAg{9#{2^(TN-VzwzCLb%HPB^nU~8oVgR-`K1;Uh#O) z#iH;06kSrI$pQ1)eblMoo8_#`?BoCY7w~^RCOYS!q6)X7_5>wuyE37Wp_wo6N+Fog zXLIcS>DegeX706jc!qn13aWoTHji4aExa6<{i9bt`X_8J#a!VQIK-2`^Em`{M#Gg) za{eD6K-XV|Yn_~IiWRaL80x0JhgBr(*M~dAf%I2dt_kq}HS@U{Gpi}~KmSO19VrJx zn7pPRdIl+Sp8gGR#j*%ZKE4y{0sYZPIwgN};5F5s5uDtk@;7K9DP6oNSudX*7TA8* zXT+%gIr$fn_8F|<&=GKey#rVZHdV)E6RHu~yS~NOQ1(Sgzdhq7aNM0Gb2h}14G74t zU?;QZ5B6J@{WTkK@t!ir^ldq7Lo=2OzRFXK=`Bs4oc1*FrvC3e2gF|)-}mYNytX+f zV;F@Jj&P^7x@z2DJJTryv;iqN+-BMG3G%hox3&n`3?F^<%zyeWC$B7s)IPJSS1`)F zijB9U-2fv8&jNNz$LY~zgzei2v>l^zz=WMZ6pI8FE$+DrTh+x00Q?>*bo&HtJx=+w z^#I_qRf_VqKu;aZ*uPi^taCuc0C3-S#n1 zG;r}C3L>_X97+wjH2FK1(XH(~w&2Uq>u-#& z<*j%>I}4$tvhD@Dfp3`F5F;nGk8ATzafp3w-Q90H4dnAK^gnuhm3BXa|MD9u)Z5H1 zQ+4eU)N}4T@I2RH3zebKyE@Jjv@!o#=Q#T1u+dXMvT?i1`KSoUZ3*Ytr) zw2mKw>g;j`IrtO*6o@V0HkPLHXR0%uJ5@d{jii1v`h0q0#rT+i|9#jhl|9CS4PDL6 z6gW_Qy|kHm_Ul_@B>jR!$R^{6%di)M_j6AOOQpQ0a^DKAo2mQ!y5XHtW*AUQ@W?+= zc3Ed}dTPUM!HD)Z^4TM9Z)U-B_dyfgPVt+TXWoZ}5{ku)*BJF&){1q(Tiz?^fCuZa z(ujCXI2_GR(p+*INGhKNVy(rpU7z8@3J4Yd1+2dA+He$)osu0H2{$f(>io&FHun~@ zvo%eN+O*t0(G!gf+uXB;ey8={aCw@7J>=G$(GKgS9JaUv{u>w z4qgA1=YbuqTeL8wWI~G3R`-b@p@qFkwY8KltNpqHQyT)6sftqxJ{E2R$BU$2upV_qEHgQuPy$ z;QCP1URZU3)JV3R#~#92yJ1~#*yYBL7SS&Sghj(l;@}@R-F&w6)53#i`*$Q{lq1Cl z6z5t?CPqhV2XzlN&TJ#;UxQ3iJ4$_5cxAP6^Dnis-Ctd{g&y#~YEUa@p>;u*k3#H~ zKME|2f943`3E}yvx64nh`K#Af=-z*S(1E*iImt#n^hIb!>XED_)Ckm&dl)#bq>T&z zP8|vc;ifW>kgX%?pk;Rid@)osuvL8l%SbjCjTQfq7AF!2|EySjVwLjrfP@FG1E8L$6ch+(;Eq!ANzu+k; z%wUr;#@16o@PUzvbDET9<}|L?D~8r)R5vV z$ET|ZU8+30-9w0BhUR1y2~(K3XdqHn1Ix}WiMqE*Bb2RNHk6w4_acLFQ#u~I<*up z>}+MeO5=e>%hiIBRTH66uE8EsA~O-}*TcJHFB4LC(zN>0u@2?CLj-a4rax?ZguOGu zBcUTksv*2)H&Xh!5}D3TaA;g5^ExVFG}ct%M;&X76<@j%2jm>D6kzhhdBN$4iz27^ z&AlVds~^}sAa+PU7%+oXBkJfTP@9go^~z12CQ=VhYB|b-NErFoD>(<9a8CFpBFoIM z^|1vTkjc8!*`0yA+x!=DJ4q)_$IiFo4ms~%>))RK$ww#B7NK${R3kZm`d(mas?DOT zAgNMct}zB<%vZ&1fKReWlB`pRt=mdmTl-WQbAZd^U0_v{+Ovno2edCMGg35#0xZ5s;~Eof&Vt-| zZ*$Tf44NyN=*Qy;2&P!Z2}aK&HmXRgkOKSKl+T@(y++r?bdm6A&Fg$lQ`425PnvqY z>_SV~7{N<$sSkbB#$_(iU4O$EFRyk|Pq}(n(LTRb02#Ohb6w0J7Or)0ZV!=)xb<3) zMVM)1;ZyciD9P}tNzif*Mvt`+q-{6PhoXj`=I&;4dqXfn8o-t`vwoFHtAY?h>10dr ze7;Mak^w5+Bk6A@T6TiK9V#0{WgAV$7LH+^_Jc$t}Upf>dpt3(j?0#g#XHBFW9{( zVsjghNSW?GdI_?A=OhAdR!mIL_{IJ9#31e{2@LF4fqQKq^?-YMfn%3PBz`{}-7u$) zFv{MMFjR-Gz8Ya!uj10T8K880STdE#+MOC1&PCq+bFu|Z0R zE)Aw^TU)73E3qp#knF<*gA4VoE_PiLcn!3xAOZAW>7^BqoB$ud9(zT~^P5(Sm*{O=5Q zUkWB{W3Ft)pFva}$C) zz2CTzzq{$iS8}FucyPgc$B}bKrr}q(HF}ttV2aV z>pU;$*oxxf!BsR*LKdk_R2Ll-TY3xDehQE(>XrmiGi+Z4>x6z8v_*zx5{eJeoGr&$ zr*KeTKlg2#c^mUgvGf{=L_K9Yk1m|+%uu%4Lc9;@hVS8yTBmhwkpWSfR-RtNhCl-G zXbWMH-k%+mj#h~~rEom-k6-J(X2kSILfC~MXqnL(h#$zpxV!szsi`Enndh3W&XT@Z zbNOPSJJem<{rG{8Ytjry0xL+jQC*(%ay`^(E?f)Iup_g8jTAC$(ZntSgJkXA&_3+q zzw=tRwIlgEw}8Dvi`Iw|<@OqSsci9+sd1{%;4iQUC*6RW@L}lDlzF?!?n)H!*wJc6 z*lzVX@`f|Ove0sPw_%l;$iOl4dlrPz?! zkjbp!`L`A^J!X`g&_HM)I~RD_vNGN(U(FDv4811Z_-)5|#WI#qLW{D5ttkT)FnXiW-jW&IV_GDiEmuq8o*x&WU` zdL@E5cW;~LF}1cqvyHi zRq|uSWIF7#NOgFY}r89nR#FpxPD38?d zuTV~BS`1UNt5V0yGgT^!+>@x!u%+=W!ruB=#$qc4!i10n36j*vC8 zhNQIgPAqtOxA&7-ig{D_Cz&=24jMn_4h>toX*I8ruub`^h!H-{aX`NoUikFQOUaxh z`h0|7gkUqmq(zA|zApMVGv=6vCLy|=+D^>kse%MXZ5WWn9MtmU7|!P78?ZfdGS5A^eR-rTNHVTSc6Kq_2o8O{l?B%L-A(MO7sAuRa{Bbd-K&daxrJ zd$aQC$j##c9_6v48$_(czLmI=3K>wShStm*h;>E1&8z0D%_cJ8> z;E7#7?72l7pbL3LzFQ|Po~zy#GkXALU%<9=1evY#=Vwt7=zU^m56=>N>kW!dU?x0? z9ZSkQI!VKfV=295p*SfC?W0bl6Q}Xh*Bwrr)#f?L1QDO~FDL3D+hYGn<%HCR3WweY zV}rStGxDB3Y8MmONh0N5e~`Maai{MV^F{pFnQTMS_&sD=b*{AR=uGFtSDz{$WOTXe z#1e50Lm45$$dtZ!#d^Ac#D%Fah6&k0T+|2ycmw%_A6$+1KKh(+u65vks^dAB!FF_t zANbA6xqYh|<(8}(d~##0`{ZF!%@amw!4vi_4o@tW_w(HVFTYGBk^aT46GQWEWB(Tg zEj-5ZER}p35!~5W(hxe*;DTcMVONm|bc!p}eK?HAF?VZg`JADj79blH(>Ge21mc-3 z5VVQD%j_(4fXetVf>P zQaN*ja*#X5&l`#*!QCV@L)q(?)#CG9+zCIcP>yS&H#8OfXa8NJl%Y4nkK7FUYJ}is zUeF2pqaMpDb){134dxyeq+BLJ;|<~+?;~oVw@n%?OJE;!Bc=8iq9g~ zDaCu0GWiEhx}p!%6P-mkTmSHAh{DaKi>X-U#U#cIeIEv6=Y9GkQxp-5S&V4vnHIna zBJG2##ZBS^O4s{!LY9*|PIx)1`;C{%$x|JZd2T^_V0gbi$sc+B9>e%ivC$QLy%`BE zxQ8YQ1{f;DaaBcVZ;BcV2VpVMFuzIck7L(jZ!7JxdvIe`VW+s-ugTDc&l`WaMcT{a5p+nrjf@0N18?$R+w&P6b~1MzmjaiF!b0cL)4A4W5xMZwBlIr+RP)uLYxpZgY{tr79b=0U&$)? zrs%ue#`N59P*K- zl};0l2eV6H!!X-|A~A#aP;F~(!qmpoD0yTowVWz=ca|(FP-wNe2SXm%dALF4VMo}r zi#cv+t@4@N*4ISg{Zb~EVZwZ>N_+H{$NxD=v4J}4C@`HEl+6G0a3_UdASnq~&!{Y` zgW`Ari%fkug}p9JmBvV=)B=s-A0^{mT!Ho<7;cdy!qXq7-q#pyj#=DcHy=fFp0qr0 zmr~r+e$V(eVdFd6-4rzs5G9&ol^t(2;OJ$E+=J1Kw`~PF4OCHhTV7wBUEb%!n4DQY z6Y4DH^lPeP5+?V@81K_SXQQw&wA*Mv@@=_esr7kGwGZAlN8{jMoQr44`@B=*ma(sS zI_W+u^L#iFnB*OYeD%N^$xqRRRbyRacJ8kG^m1x;qM0$4(cdWQedHeb@((c40V;vm z(%oIJx2L-g=$`|T?z|6M$)cBV$D5afVk~r@mQgC98&L~uZKl=Cz znwZdg`#b%?pr#sA$hXxBji|u8e#*r)kR=W1c=p-W&#IN@(W6fG( zTQ}c-YjG1fJAn{7Df0&I1>hZh^`5%g@mrIHvG59GHQp&rQKi97#tfsz8RELTK|P7x zlXmmzf3wwqz)2%CB0a9QL>2Z6iaKugkAI&Fcf+0RRC!Y2megYYC70FoIvOtcbM9fJ z%Dx!O&E1|}=i1`iPmkl{;}h`+4N|Fh2pqnW6~y{d{c7IdECnw!5%Od}Vz+`}A_F@y zG#p@<0AOHE?iZ@3o<|eg7F6#Z=a>s1*Ijn51ILZF=4_%37%y+zD5g>#Q`5g9glj2sWvsDmGCjC2AJK?o|kP`N<;44|J= zJ8gjO4e)rdj0_??&yP5Jo)+uJKyoi9gQQX=n)Df&PsdOiiwj^+71oAh9_+OV>;aN9 z{Np9YdmRdT@0hU(yrGAYOL&PC-KcLTI zAfwr8?cO}4yn9S}_wf0hx(%ut4G3@F<)z8>k4lTc3bean(nEH?k7``e0=~Je+x9mA zDpG?jyCurMGcTS70&_<4K{UE8J3&K98Q`cDjDM>IFC?ie0t4B3xQ68~!f(y0&=VL! z>RSCGL^ghSFKfLad#v7bh7o@P994kncWTd1EeUPfh2pjl$>0^U?C^`mFxv6Yf9F31 z**W;cn%R2QIne#Pq2X8L0=GUCSP(t)$b!M;o18Dd`)hv5r)<`b^i%S7&4j(Z1vujm zsamSAq$aRt0OH20Eqck2v=j4(I({bb%(R;i{RiZ?lFkABU4dL;N?x)Fo1d@kqFZqD z03nuEq(xS?i5r1i8A=uM6SN^PP#P9C+fpOCE$vNpEN5fG<*%iW8i%xz+=Py_Yyo{N zv}zjV>K(Ub94`QzcY6RS#p(rbkEnpAul5g7&jZ5c zl~4PY@%gzP5H5aQ`HfIwet#9oSF=px8TMm827#X>gQ#bKSFN4mDhO%gSk(uYfGcTb za;z^`cyf}tf;8>U8wdwJKmk@f^du*mwvV2$pSYT!ql+_aL8}dqc$0rI2MEDD7$BD} zTB4TqHn4z`fG<3zA2OMLvF8cF+mmBxTMX)?;)Dhyy?kvgPbHsm~^4c|D&e_fCVpjWcxIVsjDWNxCY%aRP7{kEHj;C5t+jl&S=0VQ^}N5!aF({!lFj<~yZBcqQ1`V7`O!V0X5tul%K)3} z4#(A1kM$?iG`R;}zPaIR#W)&~0kSdG%S5_9He*0lXuV6@>o~}Bv;JYC)L}9i>0kN2 z)?G|Br3GGs=#%_Ks5z>13-*D{@7GX9(IS~9mteGCR08F6@B@06As0Sw%B9@ zELtCk-pNFOZ8z(!R|VX)Tjr23EoH`|J+@TXXXs@G%STVNTanMV@`B7>pOyt`!C8jS z-aT#I_&Eki630L#Q?V9IiUGCznv6cpZYZur{jfzP6>iXK2_G_#gN!I1edd zH&@I4QWhIipoPVde_P70N>xOAK&QOH5l?@ND4vn;wLErX=2afT?_b1#zkwiM$+ zpXoiT7SShujMh5h>G6&$Z%iUiv#j_`u!gT7ecvW3DJzdi3!x=+C{7-1mOgR{2k2kh zL9{bs|BI@(42q-c!bQrlD-D_%dH%WZ9Ct zE1E+9mWqPEt4uuO{NEQ;4^Dw>VZ+%=)&-&dl91SpquIzjB0(*<2BX zd&UkAYXqCG^AWjyWJ2!USo%|$pC5|;r&gdpRxa9R0!GgbNhFHlt9|wOWbFtzLva|n zE>s8Tb<-(oJ6>|~>!$vz=wYB+{_NnVq&_Q;;C#_wpj>F~WWdRe`~a_5pc#a!iB?Iz zqQBA?%`ky(7^#VGLG+1f&x=QoymYS=ZE-oTYN$EC{&>o0@k61JC{@F8=**nh{BCB-t+V{Kq=u!{;J46bF(v?S6aAgEJB=najmb-l&j{%V^R z!A)##jtsYD$K)&*Xj$NDPk^jWwY?UORD4nF-Vm?ueWQu2@Xd#(C|I}{86yVa{WgYG zNQWV`X)zJ*$edwgZIxi#a!ZjYE2xyMWCr1v2FML#NQKcf>CGdq4;1W$`_pcRXTYMC z2}=J{`Zo_5Z)fx_EA!dPZx#zmNn-_-cjz<=^xPUSh{8*uh~hVdW0N+obxr1Q?(VCu z;;swex^8-?;fSOjHhQQ-I?51nDQ}A@N~*yudiDj_s+CH0g@ixDpg7;{Ycsx5JNP89 z=dgLCIwJ-{I>He#4Qq_`g2yA=r3&q_M_q+M!{ze=12bi5A}}yB+un{?!*q@Ku~#{i zF7ywxNAepFQy{b6sW>^TsZt3DiAi9fK$S__47KmPRJ6KXp9k`=$1c2dVgGyTsL&60 za&#$4TT^mUA;iI*e1XVqnd09}9*)>LFo(X_GEh2Wszhn%J?+u<#-jz7lQ6jJEHIEI z6WidfP~UrsL8XxSSG9cN@y9)?2RD!Qdu*R)Ba)pE96N+zI%vti$~BHM?Gxf~#(Iwq z9N@-`pvJ~Wmlp=XivM${psOSr2Z7DiX~Pr)OlR6-A!mV*nvL2!-v$>)vsC*$n!aP) zY{eo@S+xYF;8&Vm{l`d+6J4@FrsT$(2^tYj@emHKJ(otg-@0IQEoD3lMsBv$p@R4L zgZAyG@uF5EoF#Lj(N|p8Ps4CsXKgd-Qx3L5QYPmAfU+;qB@(ozwgYzo0s@ z=HNrg5R`4K=0R7sBCjKN41f+c24SWCu2=7(>IBes4E&<#`tHw$+ef0EUp6I`^?6IR z=hXg^C#8cG>)0wHIQ6*h0aKwzpsuT+2;(A}7Y{94fmc5+n7q2|`OBe>AeNRRO6MXK zS_hM%$AB6{M6N;0XeKR)Em9PB^%HkP_-$4gU^o51+Nfj!Z@@i!as4wWH~WfnscF)w zH~6vpm~l6nT?VQLNe|_O-_ZYEfW4Owd9KCnyc}Si`F?wRuF|;sxy1P2_m4oL9KpKo zHnRdrk;t?LW)umxHI5=%U08X_!&X|TsRVLL2NT*Nj{#>&%6E0Jxm9>VipmReuqHWJ zDgt5NK2uAX2q1%9{sNp=o&l*&FH%XSsymrRr+(jJ*WbH0Wu`H};CjOjf$}Q=f?^9( z_5@GuAO)KJ+LB(P+HLr@z(;u-jF&!d=Hm}EP+x#sI-OcczhgAs%sq(OfkUWv9@D6B zA(Q%rBb3WBJhR!jVU?qlVdho}ZUVmb^Iye2>Her5^p^#J7f)=9{;lLzK+_t;+teW5 zmcBe)+Koa>^--5~d4@ek9Fw?Oz2Uy4BFToI$^T=}ZjHR0q7hIWjem#YF#2Q0oWNu> z7|16b1X?6OS3wkQA8BmjpU13i%2RlxU!%l#Or^~o_l|FoaHkrF-=k5ffu;k4V^MnN zZ9nCcVE`7b@xL4JKQrid6_LI~x~&D|G1Rw*CMSp5zy@=WsLt<8nq8zXkm?6M=n2zW z(uUUnILCC*mz3ew4saMR9vDi>S>Sshk}2c90?wXAbV{mFG5EI6IhiwQQ%V*3##@uUp4v)y6|v(|i65vESP>~7$NwR<6uD(nbQy8?yG*eB~iA8L11ES$xa>flz@6#s!o z=SCCS@5xkk3u0aIhen@FWzqP&m=HBsosF9HG>`NZN#fqR>RJgfCLy!wcebAv+^#42 z3w#UCkWJDUa0KK^*ik@~3WB=?zpomRIM|G(Uo#gw2Jhk%bW7MiFD6IjIdNRTwwaks zEPkxL2ih7py85yh9Q@7Sc?f-J!po6JBcJld#BN8i(bOpl2MA4hDTd8gLZ4?u+l~1Z zKqjswB{B*>3Q0$IdRscVLX=xne_3&w6)79V_pHJbIeuST*V-gP-V@${2ac3dE43KYKd z&olmAqUUG9S-p^s+T6{rr^He}BbY(i%jnLNOb()3(GlIZQu{AdIT?9eo)kyE{~sc`Ti&S&|cQTlP(^TrA9}<@lz0= z9quHm#|vpS5&p6=O1GHrov|`ro|W3-icRZs++1-&>AU3Susy${sk{<_eR@QYdYSu` zL(8sgt8WY5_cNP$PRX%MIT_jQy9Ar_k5VRz&pg?f{pCu<o{EhtGWAUv58$XY<_lQod;e3oXHcHA}*$E0&rf zI{sDGKx2w(%3)W;MQq-JU^n0;e+`58v)lM>-dL4DH2Aqx_pr2ELdJ4Dirhi)p!b8n z9`^?C&4|xt&W0~O+E>2p`}5%BUB*S+YKL#qPzF~xfvx%(!5Docpj2DkZXhGf@bgU5 z6;lAH#(u^(Iq(>M7hqPhm8_DiiL%G>FM4Du@@1ZTk`=9{drEh|02KPMQ8 z@$ckps_ux+Pqt@2LPXi4AG3~_7mE}F@cssO>*59(UgGk-r3OHDM$xS$@?1U_)rATT z?`t=n?-72#c1eUsG5iZgVM2pHbR9l~fzShhGfDVFSUjLIz4qK?KK(fnS2hMd7F035 z^kA3j*@fs8&Oi?1v+N$OBCn~GpwUAYIDDD6FS9<98B;WbQd!q90+yz@jXgbg`qExl z9#aZXf1TptB~>g?s(fi`Qhs(oF{=M);~Z@|xS}L|iV(k}kUeRLDFjV~e#|hj4oAFv zHsl(#wZ*hh|2?;aSe->-Tg<7{1Q%>$=#3O!5Y5Wd|4v%t>o<}ub3|oI0zo8HtEix8 z&5{JLG?B7tU^G+>?+nM0N)m?gFSDY2)cfGcjc69iuyktJ(!E7Lc2SLSCfHfTosr#BJ#TWn;=)6mT4?d_%e%U?MK? zcN_0QGw^7HH!B9sz$@ao`%;;&O}du?-BWh3xadt0Ee(U!bVCf6zttVgZfLDE!-kW( zJs)Ay%+F!$l3+GF)Hmb6&z&Zm7dK*-Qw;$Gq^~y4t%+_E_HAlgl1_XQ6x>=|_$UV3qK7@JK`|bI436}d;2g$bGI}Co`Fo1oZtz1Tyejv;+3TB6@s1f9d4e177c4mb^798Yq5zY1Mst@#PS09khPBCGSGi!Tegl9-e1_@Dx z)eNd#q@j z&w4CZ+rXeC4$ri>tE9iXWupB?Ll8Fa35R?=P|wb0EVpL?=q4OW*SXa9miO$?H1VqJ z*?9&WqE0dT$P}h>y=3g`x-dr~oCjQSxnoeb7aYPv6i4x0L0>#8HBq{bMY=iI;@`eWYP#@Ds{)XY>z;_3yHqE4=W28l9@$pK|C;0#xgf&G>*YJ% zZ}n1KKtRj%E)7o#(~}NfJ}TYeyU-<@A_}xA@>BMVp|03zbw;nkcP`O<%h67&MjvI2 zO&fjbY!;h-3puv<>yivB;X7O5F;Y`y>Nq)$6te))Z542>JDM?5z#Zoh7MA=k`G<`) zVy^F!z&v*;FbxJZ{z>E&4V+mU6LAMXZG(3RYwzWJ8?5Pyx?d_>U;Xq^UA+(*T}oiPic zzj~Ub4mnX3Up!E~=*kNh`T)|OUqAYZ*ooC>1v00}o!h#{xOVsFf7>HA_1$&CAw;ZQ zYQ`2mn<{udPs^KuPT*LH^Cef&2&2Bd9u;lMorkkPH(fZ6=1t05VYCVvqVr&VfOb-e z{5 z6pliq4p*-f@q#<(8?jO)I!00k1`aecD&3*W;6>HPF#9ac&h z=RA=#F=*IAgi)lGV2pE!0b*LN0tyDnLm5@tf5+Z`Lk%RJu_?$JxE3|hV`wd7)B(FD zM9~>EI^IV@DLSWs(&b2hRgq=50l~y4uy=6J_SXEXKpDLZG5FAaXR8rGUN4>#{i$GP z+w04NSA}qdCi1n~H&2BibSX_#F95}Q3b=d-P!6gd<;~tRz4m^X5oI5K~%4wPwo&=>x5MP1qbaw-D$Kxg!k9!ZNZ6cn8_}S5L$D$b<%@HkRP>= zw62jERH_d7tbbkshh9t!O`XS3KXL3FZIAnAC>tIM`$3F^l~$05#M=lo`sAW0J@9#u${4-vi?ehG1`0c(*7cG$rV zpsgwiHz43ObY88(*naDhCX}oTS|{v|MlTm16X}-L>i>7xkqTc-&JcV_k27hHgnP)9 zHEE|Zbwd=NFBx&gnvk#hT8#h-yrK)6Gt|jG_Tem-)bAs%eo!}swzI#X3kw^?;K-*w zX8@=^L~Sto@~5eBqX_TK>JQT114BV9qJW` z%OD_r+f?gBA6^mpUN|aSGv4+W5S~4>Z?69pz3Bg0{8Fa;S&mpaA(DZ7*QOwva=5h3 z5dfHt5BU~V5vX;k+7bCv*6$5~N~TgZ9)is?Zt1GCk~?u|<&wSgSF> zKVhu8^$zRP#Zz&t>a;Hnuz0qwu1{t&#Fnvw z`j*M?4n(@^8+jBb0L6x&HgayzB;D1vzqeNG8Qs(&&k(Pp^VbjQIr2s~iJu;{;Dz!7 zmIQD^k^tA(pAG7-G(l|t*9*XP=fC@^C$T3lbOzE=uB*L%RdB>!1f*_3Q$VSJbh06W zurD`BKX^h)#8DKHAN#ZKP$9J}B7X1Izz{D>74>%hNi3H<&?efbH^+;&385NZE(T=$y z77qZ?ooLd?&c{pym67p{*Yw#!5o4kk8Wjw%{{;tJM?oe)XXZN+=wnWPs+|BaVHFZW zQ+r=*6(`p)A{l#IJg;f^4^X7dF;lV<(Bq5q81NS!WggdGnjIa>@{-K2qwC0ix$=^a z$guk@Za``fGCf%^L*OI1H0VbKO_BE5nn-#eK6L%ra1Re|vK~a#BL7bp_nP_|LXat- z$YuHvDwA@m=;(?CII%f+hB44`LJ$1l3xo1eBsO>Q1m*{`8_`HaDM9BUk#Zh)kS#{; zJ%de5?TmCm>@oa+C5u~zx)N_u6*vI8_BU}x!JvN|=*sUKz}5gll)FA103yax3+oAn zt+RP2mni);AW+ifH~@urpN@aHY^_wZJ-& zfTu;m0s-IdoJ|KNg_y+bXi7lxV^TTlFYilM9LODUBa#>D{1O_$r_c{Sr85hT9=?`)G(_Pj-HE0XU1c1X?Gpyf0olv*pEH0xX^! z;X(x0-5m9(f*V=kDj(S5fuIwT$)Cj{MwpPCX{CXbb0AiE$-WUgr(9^w1+p3IEw*XR zmY2X#GeB4zhSbgm&TS!{0}UW%)!)58YhRkCDaPqb5y<$mJ^0cMd5P7gS3}(777dmf zIV8JH5%Lw|O#sTtpQY_~nD#I6fk{Sz}4l9ZtF;zY-qYXVy#_&N@JbnZGMEKvOIgfY@7^_8&2x z6`kj`(VmE@Ke(F!6jc3Uc>M#w0s(GA-08d83hYvU0Lk~aj%Vad<`w4iixVtiMw zy)0JRNHmFf$;-k{wi6gZh)scib ztUt$4r1)5{{AV1c%7QPElA(IM_2c%8!Oj?7XZ-@EH`O`$?ajGSIryon*@}_ z>4Z;r;wxX_$@tfO^J#M?MYT&<6j3FQVD?obmW_3htw!f~NUh&S*PABi`#9K#3+2-l zV{`W~F>C^Z(O&%X44IMoRu_$!>*LVc;A9Z#sZdbN1d+|E^20ik9`GeS&!#X_WQLt- zF`@kucTUc^xCs64$r&Nk!XIr*76bj!hvW|HG&ClBcikmA@x^uDoH5|Hly#+aOARK zQ2sCrrUMiu>|Fi9C+g;=>H|_?{5N{|RueJHfwUQ(;0EE8!Xq! z`0|V9A46|oX=98dwqZ}<-`H73bq{gN_W{%eW@5*=D&5*rDwRZep*34>Jfo=N=6aYZ zy+@3JMCI^z=jI>hIfKFK7^zB;1uBEIte*-P)=#JsQ4!{biur1w9uZx8a* zM_TX=Civb)l-0BNMjr#=+1$qn-v^Bp43go*KHG!3Rz#55sBBRh)Z-Z4L1R%Jd6@kc zwBy%qpGuKxI(d)+2DYkETZfUVTnsH!-zeyb*)NcucaqzsOWR%DHB9NzGFS=1)y42p zqjc?e>UPiLaVxWJ*5@L&=Q2Tp&_sWDnp*#Zqyav*XP17rZWyG=j75Vf;M+FREWA{z zF)62U?xOu#M!xnkQVL(OQ9giBvF+ltg=s&Jhd6t(h^W zoObX_VnA`aKI61J-viGQ1rCORJ5THFNK*gFCfO`ew8si2rP-AV-93aJbE+TRZ5>?N zN@*S;hueR3idT8sIjNAn;Ji=&QzV=l=>UG^+}fVM2jp|r%ub$rLjBG=)L5LS!PiW~X*n#$<2QC0KDd2N%JA(K1bf&I$!S^-ou`7_`RAkl3S>LVKn_artNO4T)!$y4=Rc| zlgk*$^MV=Cor$E(F_Cm4HC6v|td$ghzMZ6fXPxyU(HCg0bKVfO+WO32iSNs=V7#ZJ zjOHGe5s@bw&Vkys<6i`&GGo^f;w^ETynQy_rSD?9f9pqdsNRAN*OS^v<7Y^B@!1%4 zAXToPBwNiQV;N!y-5FRFm-Z6K66^EOB-P#0=6u_!G2Z2W2;Xd0;7A{^M@hOI9QRlz zyp5n->pRjC2{U@U)%C`^U({sLSqb;gqyg$Q7Hs@*i2`bU*MFta8 zQ~LpgQxy|z)5$G+4lq>%({;pD?1m4RYBoopB={TdhsA014=^-hS6i4+e*u`V6G#?l zqxsw{O0v~SuylrnOU}%~<2v9x2t>cQWmL0PBF>r;=c5tx^+z#w1DrXjUJ{y2-73p1 zV>MEMTvrQ4(UXK%6_y+NJMA*_8DK62>h8JgCjqW@;lmL@l(gEtFd{mUOwP{$*;4PK z>ZuX`T2M32yPqA5&x~N!eQ{Io# zLL!SGl8R0dValkZf+2NlWQy+lkLj`$7Kr6x3k9yYne-X49gp6|*m6DjCnIs`KCFR9On>NHdhaLCP$#%(g$53;JijQHw4N8vb2h zMlHX;!Cro8@I(DbKLi=^O^DzHF;v|)tdB(%eL+< zojPw~6!T(>Y$99ybV+djfGG?5c1vs7reWeQM`(YXJ*?WuHei+bOD~+t4nWiUOxTi0 zLqRe=9iW!d(u|S}d6l#mA-HJQzn89(_K`=SE1w4w%U38h>xDDfa9jD^Yg`l9R#L{5 zaI%SB9S_g1R-=@y%!g6GaK{xZoLcVZuO*k@1nj;sVbw?rXb8r_;-I!d4r?KnGegx2q_`G(M#faDh2qMJ#>K&NnqLk%$)R8 zJ9kA{&7DGEj~oq}pMcy!@1N9pn~!a4^Xk7dYvENqyAKnpzi@OqpOvzKhHZfBNBhO! zPMk9UOlZ5@%r9-CLhnF&0bM%3|AMEGHld_x%sJc(VZ>J$o_}gSl2wCNVsJr55btCW zWe=T6*FE`9dbC6ehbmmG2W+nd7(o{V>*;ChH7}?3mQxO$tM>#p537x_md6i%XJQrP zqZ)U)VfanMrL56An?Rqh2}RuGOnU!b^{xQ3+@|4aJ0eX^kTn*+SfVIv)8eM6SmC@i zQ&n>M&UN1J!~Kw_8(gB`xpzYMArZnP%MON3#w@?{)L+)TBikC-1%w!Cq%R z@51=WgNMVWkd~N~KT0-iXbyruAIldl6P=@2pKMe<%jOxRFHaKw(YK6_TfpqrSYKg9 z^0wfES&}UtDJ>ad;T(HCuz9SH9EAC^Fv(I;u7`1%hgQXQtgCZeG2QbBY_4(Y87pD} z03vL|e9azBnJJCS6C>pfYu_MRk)E}gL;FGK8F0a}H1F$nB8%fh`r59valK@7X!5-? zQOWai&|x(M4&5a-r52FF>lkjq;ME@Eh^K*gE}#52eJFMXB-1IotUF0o5lgT}9cRrc zYY4Km(6UzMbmg~&4?>?uTa|ABhv6tr>Z6`nbCN&!HObl?l|)t(?=DFmeui5YQy*#8 z8XHpZPF@?-9KCCweFYf#(2~%8R1dF*!lc%cu5Yzm_N&J}ML2d!ka0+R^OhY=nlwrd z%Zd-%BOh=O|{L|}5ReWKwGO0x~yA@H+_MO#1^lHM#75oY`1uL~}v zmJaY6 z>(_7c#R*zuC1qC_<-n1@HYpFn>SVoU%4Fqqfw#L87BQP<{c3d}d5H*-I z%hnd|VaaG$4EI~c4=SL6g9bHGKo8dcx-_x*UG$3gl4IGHBo{w@x-M-b3f#0?4m3nu5 z6Mg1}MiKbK_oS>!BoG7@X0|$KfaNPLR2;Wghq}t-j1flDroJ%*%?&U6^Q}rHXR&_< zr{mACjUAM!;V~Dxr&-QlRDiQ^Rk_OQ!h939mYM94O(|DF3_bv;_Z7s^C{aC^m=BX! zL#@c?c?kbfLjw6H94=#A3tOYxN%*i!K-;XVI_Nl;05^|&M1GAR1g#F&hm3Pk%{<;x z`IEE3uNXZv_Idr(Gjx32c(_a>DEj;$dN7=i|24!ilz@SU=@h}Hv8~XH8Y40Ea4q}0 zge=Q{>!D5xfE!5cLhxSx10|$#^L5u%TDR~}Q!vUAalR75=h3SB?+lI4m}7eW6I2Y=a& z+CYqPUy<7P+0@kS3<0ZRo(dV2LH0%D4E-hHk7^{}YFUqRumm#cZ>)61o#)uC(5qrp zK}#VTSV?tg*EHpb0k}j8!|8;1D;T=M@HA0E?4y+4<_Ah(dKk%2tNtuOT!~(6(}KIO z8Nmp)t4og1_8cPA!$dRGmK#TJ~cGTD*) zjcDz7oeGLkq`hWXh3VJUSyxO>Drxexgh*O`C1`vh{84PnBUb6&RHhEMd9b}oe^MiL zDS6F?WsD*jN^tY4$o!F~JtDv+>{^t6hv9u&JqeCa9GbA=VX{XKTNt~YTf`J<{<$w> zHx$8#U?yV=Nh@ZXaHdkvgwD;kOWKL$X-Fl`!%)RF$_bm|_A?c1;BZKa)CWc5N-e8e zGf-Mu90BB$?qsXi|IAL7K0=g_oua>Z^Z>AHi5`aG-9Ef6e27tNp!yy9dv()7-e}{V zAgNN}IX-?5Pd(m!Rm~`h?+KA5vyp&&XuLG_Ro>{Ei1v(*-50%^FmV zNbxi6j$f~qA+xGk>==eU0XU+cCFM4h98IF0kgeU)Gw#!;wPAm&BK zD!#`hAKWIG$5Dzz?=kk3s;gfNKc|otSyxgl|H$>^S)>Rn9zSS#u&e`w^u1bq!>8XO z93IYi`YhMMt4eUen~EMk+&xC`7gZB0{Z@Tnp*}LfIQx3P0C=k{FJzighgV;BurQza zzOfN92!4|U-qs87MpO&rH)66}cA0Y5ej|CGym!5e$#S1if5}Nw%2tl*YX;F&Vy*=J z&`c=DO~r_&Qka$tpc{e1(j5fCJjwET?)rA~yV+ zrYUqp=>rEyd!_t%M0^I_I3frkL@`JqEAusR8Q9$8{pm*&@T%q|@pBNRZDQdDFZ0C1 zGtCOBEEEz#27R3-kYz0eGiw8nZ!BBB?y7j!OA(>&X#1gwXMg(D>WUBZmF(MhhmMF@ z^sIbLDLADk_OZ+k&fbM+ayu@ODNHAV+=s4=&q2CYgO?29nWwB3?vwcIeGoX?hWMDV;RT8B-W#Ks( zBd75P9ij#mkdHG3BbP|IC_Jb8!cBtWFxY)buHm<)jB?Tz7%0UY_~jVAP6dPNc&$rX zDUIVRJ$e%CMHB0qxEabU$odMd6xfO%d(31^s!WxTT$iNdkMsNMIidoSK*k?L&5A3A zytc*mB#Uvo{zeSWFaz-R87a(%%TF(Oy*sx5p~zc}AUX-wa(QeDIXIken9@zAGV1+D zSSmiqVo&UHMYEY}(+O=7b&|=7_q_9SnwqRBV0H0G21zto1A*Sm2v;EPtBHiCWd(PHI=|gTXp=78;v<`VVbs z?DIks`PDO;&bKJiKpw^Y%Gcuodt8yxveW3NbG@EB`_s05eo2ht*r}EGk!y662RZPq zUiSkZ;o{*Q^BoZ|E;P9EcnE`1#*|K~bl$a6?h>ZZf5xp$wJ4<(SGDVTgzKE|1wW-m ziZ|k0B@J+vdu-#-Hqg*TTun!t%i+#?a5bc(PM^J4d(Hdq6oEn9hKP{>af%=OVBE z)LPe?;qh_ai_*Qo<8bvbWZgg-3qcQRz8J=nzs6|!u=M@_)i;usRjmfD?i`~_qR^_* zfGV^)kkVEIZ~~*#^~8Z{P$cqzdx^~A2{t7abcp63hcCt9JjsjSOfl}g$*%}VR2s=2 z3aLg?hVz^C2(Xkw7Q;85}Jqg zjwD0l)0TK*jI*u$qc*We-1}I;Rs?JRn<;B>$?#yqLy3nT%Yk+7+m0{ z5Z%)$<@4`ETJ}Y0O!Ydv<#t>@H4bbd5->81*lwbmdS*%#D15^uClRzO8t0NF=ZZ{~ zG`4mahQ$EwssxEMgob~M;Es%s?L-23HFt)NvU|;Og^8J^c0bxV6wPt+v8OoIsfz~d3W7{f#QP+EsXI(8v2P($b>A*Te zL%866C2y#|A570yF>-q@J~gj)S}Y|fEfSJqW7BrpyUKer|Cq>xbhMQ-KUzbcg{u^5 zDEN`FBXs!;-6oI1dz0dKRM|%`E^JA!iRfxRJHLe|{&h$0u%#Za623GmAx*{{biM~Z z)6uUd%UYuF>(tXg_bXY%zPPd_*mMIJ8tOVl8>H{4XFpBuLEQgV8Ra~$Vh;G?SS_a7 zh^wO}bdvn$if;cgm9f>iqC-A|g-Ur>HwHgEsV)kqgatjs<3J4N(SH9zT*x(%*wD?oj;9a>G;X@KUmbgtK*jbfNH=V~lbA zt-!{p)~mIZ?)L~5k%*b9TCvjwn4>iJ^bA=H9Yk`*$PE+45nyo|Jttxc8{8*w$S8F;0CfNj+(Z>c=-ZzuF4>g3C;QiS{wy_2GJJ1u+) z9h1FCbgdXnwzB{1brCG$kOXaM%M$Uw0<@+jCW0#{j^*C!10UV zyh5!(mF_7a{`nsQG|%PA>rWY&)+H-S^0j|&G_aqRgY2{(Z!NjS}oCbc1_<2NNwe|50CR0^ z+hS*7Kh_&Sd}m(>B+D7{#5z=A*6-Hm%kl4peY}j7e7xLnq7S^(HJD#WN4^V2>UOKC z#6oqjgL(;zwPI0Qp*;Q>(rgN7<#Vvwqg;LS};V>h`Ol5X#21*?5?mt_BgWzi>P-VGJ+;~ z>$}_Qp`iRxoeVQfC93(D`>yIAlgR|$x)%fX{ z7&ru;d-D|5i0YT|tW_73JgBk?o@U4f`} zX}k{XQL+7-CW>CY;PX*qRSoMs z6c^6EpGX~jfYn_vJwqpYzn+FKr&V|uQ?1~h9kXQyJeMG)=Di`cYs>?FIXB;`++M_nsSmvuVX9>DE$VJ^l+jLpsKTGb(Er;u->a3L>pMpk01h=_f&AGY zN)C*?AyNv!P zS%}TboMOdTgj6TJlrtU{Oo*%=CO1`B$fs++!RvnwUYn*#)Ru4mz68wIWJJTrFIlE- zmFGr~AqAU2A;2UiQ6aztQ(z_|r6I9Mh`_Tz!%NhYkp)X&V8oNriR#SsvM{yF?e?o*yhJNa(oz;b)F>5Ef4UDqa$_13=7#m9{w&q&L(lnq!TU^PjXf59OG1PdNLv39L@y}Z3e1Tm=nRm6Hj_;Yzcb|ND>_l;_ zKk@1}l{uPC2~tAixXQObKW%yg#EJis%b6Wrq6+!jo8-IxsX}68&Iy1x=2WfFOQBz*D)ee(FHCW-3{$Tajq{rRA(4s$r53N%rqJ?xz}$D!aFnn!leqrMqn~ z_lYwAtW~ZelDeB7>ENN|yNlxoo(7><_g=SsPVCpo) zWvSKT5zB{#plOf=x;GY4Ht#73of%*sA68Ps>Q0xUJD2yqA7!hn=mG!$OSD?4f1>C0 zBtO$Wrn9(3@%rH-1%{991A9vkkUeJsaMrs0kusNjC9)97H&zx~N%inMO|yu1D+6Xf zQ2+qI%NmeB-cs5>VC|*4^hT7qDle6R6f`AkgfyW<;40pb8UVz-PLwkBkgakU4?yo&G_*|khPRvn`kI%@CH01F>#^#1u@zh{7I{DIU;db{uF+|i4_hMUTnhm3=vYNq@}m>99z#f?uSE??trJ1(dU!OmHzunH$VKHh8Y^1It#?&U4g?LBC)dlc+aZ(* z`&EYjhp4j*iX+_Cb&v#?;0^-}GPt`0x8UyX!68_1cXxLS?(Xgmfe%Gi z^f}0!8cv~N4iYh7EF1&DpyU|LUqTD%OkHi%c(OeVe^;ybAfrRBy|b3P5SDFjR+L|W zf72lGgpGv{>dZD=zUBeUWR0C`uOsD#>iO03#Bcgk{T*-~vaLHcx_g6=Q_Ma=nJ9I3TM;qR=hLoVuF?D9upX9UgMNnQXKg_&r%xL`#iIjp8T8Uo+ zk$ESYrCVu$O5(MJy}yr9agI0X@aw-X7(CD1dN#kRJoVf*9$OpHo>XPA4jYN%ZDy%3 zmOX_mNX5EZj+Ji!>8+M8%|Z9M0?i#8n}2-1g1kCDn=ic|^}tQ-^6Vs^vG}o0EpQ}F zG8cO6pc*%Bf#7y$Md6z>eMnyvXGvfQ7S~ z+Z7Q-^;f+-H|a>@ecJjDWxHz&w3%`(i*-$CaIG`6!L73O?*WJwbZC~aQwhH-2|x_t zKQ%&=wQHdpLB7P=&*BuQAEic9vLaGW zdC60PiT(FYJc&QX*yDH5d-USWT?)ytehn#`M8(S9n41f_TA-2^%0@g%62gN!J?bvt)5^cey1h=CV^FY_czwQ44`V^Hq114V zwjz&fBViNd3(^SLJ5ZQETD16QiNA2BnpSNSeGb#s2v@1~Qb{oKC8+lzeO>FJmndP6 zho>`xczzgzyB=Y(XO_3##LfkH3xkB92$7>OW065mq1|8m5wuh`1L}LyAX^mm5&*wzYQnSen%v(KA~vrq>>gU zkB#Fos=%#T zvI!)ywf{1bAK3W9W+!Kom27WZxlPLS7MNU+@&rYn9{g< z!07K3%Qk|7RMVfERR}`Hu&QZnJ;w{-nrv?C$>D{R(U6qHzha#k5e$nSG|JJA6#Z39 z+`?^m4Hy_zAY&mz7o>x%gXVO^Lg8TQ$ccV!Xj5Gzt)R*v3y#J**cfK$6UHM7`o$0x zXyDd|qa>q_lS99T^Y_UV+b?GNRz5>w0Vm^WilNNw`?8NF)}Rq@I%-8p^Pj4N9j1RG zdV)#F(%C*o#F&GWj6T6jTj$Ev@;Y*u7j%-{iMg?1sh=^;7nx2VNfw5~j7T+B+eQ3Q zvXG?vMAe?SMk+eviLIK0ze*2eWdeBnXwEW3Ppcz>2O3|VzY%kf38oU`ln5QO4T{0@ zo4|Nzxqn8g9piWj$Pl_>|3r_dJ4#`}k4Ow5f^G9ZjP;SwBIlV)D}at93ejSKRD z;A?O`oZ_g2)m})C?agb;KH5stVZ@x=B|%PT1s#&0Di06Ud$7>No&w_>%%@JvJsrFN z*P*_h)Z8~)*9Cvg2RA}!9d^EX+W=&tmz#xGfb}9lff=(fG2*#fp>^>)@84bWg@bF0 z8ydyKr2MZv8tboYJvoN@hHx^nRH%%pndOe*>A1d?r-)0%d-G2_xIKp@rrHQmQ`gWa zE0IE5u1p-14Wh!}j_&u{9Hp087_9$F)*1S>uc*hp(};(RCc=~DDqEwlr%7KgNeM0Q z7?fb$jawrbKZf_mtk2v&IYOD3xTwEo4A*(8J$s<))_856sJ}3ixNI0s+v7aic;k%y z!I>toSXYA%CH6n}LF=Hc9C*7o;{D{rW5_e>ZkrU`fx^z$H_M;8fQ>SN(eGc9nbZ8& zazfmDrER9z!!j`s%`$|2wqWE%vUMM$#$}w)@ z17zX0>{;S!?_lu&@L9A$*y$3cLaOO~=tb6}ERiHC5~O^WXG` z`I*`NDC|!>^Kbj4P@e!3x-OTLtz zcleho>IFzFIv;LunRmICw_86W=5cYek3cyK&~=tLR=L@i`&@JRE?`^KyLVWxoaOwn z*nXkp(-|b|&k9V%Qq4|z z_DdC>dv8^37ZC_|SS81p<{v=prG+iY7~F^mH@^J5h}18m2y4#ZSRJNhvooGeHTRIy zV_*LDS#Y=qpyLk-79+WMBu4x6j|QO&Ur>+dFVN2?&5oPx4NaZ1v6QgmN~ zV_sQ~2k1IJHB1_PWRT3jA?|xSM z#G>Q!4bUvA#IDGDw`yX-etEuOOu5!Hm6B6Wr_$IQtbkppIM^7r_$6JgplUN@_SIN~ zo{dl-Fg|<)V1i6^m=t;KNs>uGMC}*v^G1opuEj}NJ35<*x+N?XJc}}Vt z=+nr|z9t#2oD)c{wjg-JPeyOw*%7ZgJyc=Js`{lQN=+qI{;=V&H>&sJfEEdP7j>Kh zHZB?VW!*d2CHU(|hcc|)KD_>}S$Eq*eC_REs?U*P)=O}xgWiUR00G-)wnWO90j_{Er14KDrGDEc@aP~BTs=idNo?q|Hq}UZ z0Fa%NZmQ}d=-aApc1fnnT%^lN;6T{WYrkuYX1tvNQXNTJ3TJ*$CtOETcOqyKe@)6?tn8r5MIGW94Sj5BIvM3HLKWJp0M-FM!uYT5_X8 z?4`qcJ@~c?fBItf@9!=HhonQBx}x2p*-0vbppUbM6_n{U0g~pwB^Q#y*6BcRg2W zE3N)XUt|{(hLmh}_lM_|UNR-=Rc8_Xnt}LtiKA?yq4M?@=ImpDiDWM8X4md*O?@Qo zdA1NDR(EZ(kwLBC>ytoUZ8)*hHT@g>vSi zyg7T3&0rOy1i}Y?*Q7#%T22-S5tYZR`8kP;_SM)1A`0W|s44Jp?l3@0=Z?tW1=FGE zJARmitWf{LxuTMdl2Uf6Q4pR%FS_%yDC>T@&#fh;-pAnRAM#IH-j zQz$q~BNr4k3B?%jRO0tBIGXAligqsYQzF{vWF`_4;4Vf1*|}q|v@p}uBJVT~E^FuT zK?IeMx!AYgL#B=i2yB^8^H8|qXpL+^Fn4#}cV9j-Ku_*OxiUb{l?6!S80mAev^b-M z7D@HhO=9)cgP*my=Mxw=~&xv(GQe{%J!aL0piNnB@97?G%69tWd zDOmi6@Myv8wU%Kxl9hAuB4dB%F$3DN8Ds}8 z`1`Fg4;gMMq6)B<=rQ4_&IiN2M(aEA_)pjTdX(Z{79Zn`D6?9L9{4f-BCqgKr9zfX zcr%DG%OiT@L5X}IxR6jCO-y!)5z25+FJ3QU34) zDeEX09^&aVM?D#RRfb9dR&BrR;j-$XlDkM8AQrmCYDC0#QTb(jBmY#+E6cGnV&qy+Ktf^%({ zx0%qi7W(O0;Nz*}swEmyI0{|xEuj)yQ2&NlcKy3c>D?AF+~3@DK*;k+>E!|a0#Q=j z?BxxzO;YqN*kp8PVeU)zYT_dhkfft~)Q8A5111I+lziB?|K9(*ewC5WZvIi?*hF!Q z?{%rkY-ykb-l30&d?O>>D>~Ma;CHC^0S||}i(#_=_nkagR2b;Z7m^Q&s<>z*P=q`N zZsTrT2C&M1xB^OgzkJJnT+|r)JG*#2)mScf@-~_O-&Okp433-scq&4HV`+4jP#}JX zfx#u<~hF9rLY zwvY()48O@voV~|)y^_CuekT* z{PX8rUK>ZAM<0_bXqPEr8fhE}0!1C@oFjN>YLzjM)FJ5|iiAY<;D!yoEk913l4o#% zk^>r5Lq}Z>MDkaK8hh(7*TGUfGWv%Px-QUgy3(6R>&Ry_k2@^dh^=CV6aoN6&VZwqmR z!HEvScMS*eF9RHM1HF3Vk=HSPMZAAY0bW9c3xf~k+OEMvhOz#jPC75GyO1)|7)GGT z#QUkmm+z&TiubcWFO=&EN-i=2i!3qC>yzPfPTNFdyZ7jyYp3ejl81iW?pm4Fk0qQ^ zuEjfPKevliqJvOWcQ9)1tES6_S$cG3)|fi}2rF-!VrA%R5T%3i?6zsGiWY9q)AC?# zPE~!7bq!d=6vKxul~dcXJXOCFh_AkYuH-VIUpX#GPOIm)?l8Dd;l0m*c#rogk6x%O z8SE9ZgZDz_Pt(Md1--t@-iyJpp0pT6D7Hb0oDmKN*+Jo&J5k> z0J3+b<0qK<8dmK>3x-;6VHfSW7wxTYAo_p<=BuEM=vWzL3!Vh{MVn^tEN(5|M>?u5 z{}URjex&!IDCs)sMf5zN&Q@&iAKbAUF^40mE1?Hg|($oAVo5O@HD zaera#;E<6zb_t2X2$Q`p`ESj9-v6$1ORd$mYj4hRT3X*1VhkeVPD36z&Z5Rcx{M8P z{`@=tW9W#&Yi+=Yyb08ZwC@9xDo~`Zt=BI@wJ`XqIakKLA6p0=uo{>l*U5ofaKZVT zE%yN#UuJZW+{0M@NQs4G>#1vN;_$Ec%->o3B1i}t=F^$?ey^|k;UxVC%T)v&JmAvI znR3uSiVlU~_g9ar)p#zwDJDQe$bQ0<&EEW50KKZf;lXmiu<4fvV1Kxv_SvD15x83w zEBg2YU@asoJ=FDUtXnEdcj}tmJORlod+%LBgA?tQ@eUHy1z{@Oe--ZU2}Q98L_lg1 z3|_6>L+EzgSJ5AS@(x2qPl)TjyNZB`Koq)UDSg9>v_VM)6vGNLm?z4o20nH2bOD78 zB%LS3I_LYPvf;8(tnu$B0niAB2exrJ+c>ZNTO-6XV7UlyD%y0AG=_W-!3`#G&Nin< zyg(zrC@wMmybBM6UQ3mmF~g7F}sYt#y|D`Mn5>;G8KtQTGq$fJ@dUOjaD zGM4)+o$8U5-)C|N^qjR?u&)Om2bc_YU-R``zlmqk)Y+!8R`g9gtVT(7A?OREI46D6 z_+T1%*q9|VXZO1`3umvsIFB-6*QDgxwpZ(7$jnL3oJ6FlV;o1)sQTjNvAgfyDcUR=ibqTh@z zbg2O%tdwjzjLWf|cM%$$J8v};d%|)PQl^B>nCkI$u$C&8DGqG;VA@pMG0rPM-AzTQ z)knZD93gU!irxpLU!zY6Ni7Vxm*Xay&uzI+Jed*_O;V!;V!_E5`F-AC#@sT>sd@2` z?+$F^)Jr%$V%L*wGn5VE)szZ$k8gtB*30JB7jAHfLu`m1HV-*;GtIoD;Og;&omHMg zvlQkZudSzXEQ*1~?PKYG#_Pgo2+{H&I^^)7U?M^J?RKCi36%rLV$4<672t%aT*-g# z`}%F=s{oM0yOa>6m}*2&ucR=tC`U+ZqA!|3=_>T_Ir4&Y35+&6v&mIVVy4C*u$H`?Dq*^uS$ zKxn1Pn(vDZf~ima1SUaH71u3&?Y{eWhv~#c@yvp0J0fQ9&>^u|w!^ty)$VGuwK5!7 zJmYK=ep7K>>Kylim*jwC75CeNM8?bEpub}Jg{?dACuU}$Boa|aY0lP+zpm;CE#A-_ z^TRrokyZ>b@QI+AvhUL{Bp2^+`C3LG1<^1vJFbwu!x0D4sTVU1l4@M@Y*|G6If}_8 zIh>PqgrXwHuI9C$u*G!ihaZDlq!>frWcOq3QFxVF^H7*$MuD#|f9!m=aecOrHF|id zl+Iai+NWUFcVbtsrW7MJLUAsL|F}z5Zmw4Kx*KB^?x48M6$`9p52Hh!>F&fn4LFVv z|K=3{){`Dm!yE+y;+uO1i0uQ!=W)1dA(0xEKkZumWu?L#C(%cc9VrK@Bajvd6$Wmv z98&fkC2&N?=Exg|`Ago|@E2d&G|TpC9D?cIKpgGBifNp#dN}=7eO(7y65`L)mCdU> zV52L^3^fCZE#`;$uPk5oO#d&%ISHs5Hg{pr<&r$U`%b@M{hCfrn2%`+ear6;$=&Jz0=_? zrP>h+j(A?We5<+b4s1WA=2k}W$c#?Uj#5k7hfTs7sR@>j^mmQOaV>_X#gD@mdBR9- z9ZUu`2cru!x(WGjL2yon%H{lGyK&o3U~7qeNM4YLpcu@{dNH@s^yA_yvUKD(=wOM5 zIK_}-c7l?ta?C`7akFfU32Uv{PJZiVGj79hFecorX7uOP!K#>vr=XssECHF&wjAuV>`>w0rn!yBn z$Oz_)G|G&TRJj~I674yOm(3Fh#H*Un6G$QW(%bV(!n3qEum|ly4dF;kU}a0n-O0rdy> zX1*p;V{Vd{G_ED`nncm8&$3l~b-j6V5k&(}#>hG>TxN zdSbkAA^LLW-|0+R;ke_yXrNN~;+&`7ijL6j?N@A=RHXg%ZkN7$z_l=h@;HMK4~)&hXsR>3=ArAN^%Nr3m6?4| z_W#&4h>EG;>CGMkOYR_I1s^N3=R7#O8rbp_J!fSXKM1Ccu@$b!N8WN%j>n;F1Ialj z@l_k3Nen4+293$qvxbT!)8zXehl3LEx}?qIK0Q;jeOCxAQ<;vir6O#&x9XnEXm5;7 zu(D$hIVnhKiD}lP2wrDkv_qn#Rjoq9?5Kb@MJ{z%YB!dTf&b;;u4-(z{A!6)UpKf{ z;$)XB7met7Y!#J*?8r`mSKncea`u`P!}fb$Xpcf}KfF}rk2Na6Md-Q&RE#u@a&=561`8sfDLG10mBNG+ z^cZo1QA?0tX=?z>h#XCNYtZV&NJbI5vgeI?VApW2o@?C%IM{w6J)%;TVdC`?xw~MX!t1`~W?32^)Qtw^sJe2jy z5I7NL)-!hhngn|<1fo9J0|QFZJ}9n6j5Lf3WJP`wnA|X7)+wdHxzK@fSR`KVV~8(rFPT?~f}> z@_J^;dbm|%C{Uu7qZkEj@y9aa z_RZZ{()n!vOnqtEC(U=jHC6SQO=1erlj~vmji(o*IE9T$!D#m7^{AzJB^9j^I#=$6 zJ3rx5Lh>*pB}SsaIH4Y+vzEEmf`-{n;6!ndnwD3gnSZi=ur>}1o63>uw&(`BJbPjk zzY%(tJeMt;i+z_cd!o@qywaZDd!FEIS(b|P7#7ZmlO=htw$IYq%t~;SZrSZ0;As3C zFsS0G(_Jf-1QObn^Mjrjm&;CM;CWaf?2jYLlg7M~pnT-54RBlal zL6Z1-uw@aP(cu|h&tq|o$T^fv=j&$BA4U3##ffFSj6uxg+H_I&E;5dVwL}XkfBPC7>2sgtxiiY z#T&4rqUHJAxm4_^mc#y`xDJ^7h*13h7RCWyvkg!BbgH6cG5s{TYEf8;U0G>gT3Rj){bkd9h>)N)Ls2e=-H~ZO63UgP_??+ua5AW zV7(tU}+t;k*Gm?;q^Q%oC^Xs|xOOFI~g=w6~&GkE_38GYJ6FOAO+vL^%c zv%q@XY8B4aq(Ttk8u8}4ulwoxKdd{J@6nX0xiJ_{k0@O26#EZ^m8Z?#VEoF|e}~CR zZ4y>HZigg;MIg?W-MC*Y+G=Km4)1K7Qc$RJFp5_I&7B(NIB&Ha=$Z8~AmZI_ex!?S zI=O><%w=eQ zmNh$NOF2CZbdUzNAxb`*BSN1xE1wCJBm(|dn2V#2R#m2(6$8*=D$Q9Ch`a|z$Kuu0 z9If=Y^OzXn3A=@@X(Khx2=X>HejyEYqO?NV6UzCGFyS0Mo9AK^tf0=sP=Saa9(6D# z^z&AE*LZU8N6BBZTpJXO)gM%$`N)6qY_Z;kpqg??2XY4|)*|eI?vctJd zkKq))?&c_L!~#wo3-Wwt0Dx)(m_<1S@#H<^_>!eZ225fFQha>rEv9rHI>)mjlB^h9 z7zvj|gojJYYSW^|+CQd6pJfpP%m8% z$Ntx$H!R)3%-#+*3iKNW>(-{{&ZfcVO40m}{ntN$ZBnyRi!+90Trte7N(Hgc#X18k zfnXqoP1~tRedrON@8SK4XZq@hi%{<>hxe=Qnb(#Pi$xWVgs4-`efOE}u(HUR3un9w z$3x=oQR+ZxNgGh*xp`FT`&jzDSnT5$C5~iWgBE&B`VrMTa)(%L!Gm8hjR#k=yJsc9O6 z4@y9FQSN!HHrnBmJgk$(8{O@=|MGX0f8y=fS(QKwldt(7xb)|oXL%QSvv)I#3nlXT zKj{%xGMSJCW9Z0FW(4!I@WtKR!#wlwu zEBG1mm@$=h}$2DvFQi7Km!%F&A^0hgoO*!TnND%zw>#? z;3e5VWYk#gC*WLs@-SU%RwGC^i?s3p_FdLCU&1&}%T}xJzvQT05si-DbjWayrO(oC z)2ywV+z$$Lpd!&zzURCH&3E}r?tk^M=MvF_y;>hDYi0!6FyhK{8W6D z=SIKlxK`c5Vb|GaP4l*m`UMx7ZaAS-8$)Xq&i0eb!x?vZ<_{}r3T>-+6T)H9GMyu!oR_>t@z1=NZ3}A;e3_OLk5&i>~9O(s%Dp3oag4v_9lqw&l^1%SR${7NT>^wOqGBE0`={x@OX?wYhykJAy$0qaTBG&IlIz&-1Q zF0GQI4XFO#E#0@@|G_O1``m1YmVnO%!Mizif+;-z?svbrb0~%}4mMIAqdVNN*y@MD zE6RHV?i4msb5w###UJZyCT!Mjd#!Krk74`yY68~vmziRp$K;k>PSHi*=jRaGb57o`g({0qaqX$CScB^@2bU82F>B#ZN@#n{~f-dvE*d z&F6%X^SD`3Af72c&!WA=Z3)CArl=6XpAmbKu;?_AKKGGEF-YR+(W?j&T!s_UJRpWY zIoJ)gMK3fA+ue-~fLl1Ky8%AYb&uRo*-6ZtH*`9la}-((qqC_dR>{dU_)jGNz;>;+ z;Cy%jF8NRAH7Dm=fKlqm=3V!Rs=sQH#R7%tfIHxa^e;O5U*!H!923kR}c)zB@Cs zv&0&*1yyAu)8qKotECse)=1q^eN5VhO4ba3m%!#RBj~&FmpcyKqfur9ODBJ?kPX9} zgIN3oimQM8)IostUYafcN{+D{lzzyXLJ~h89<_A%7Ve<+b(WGFx~erGM=ngMGfw?H zA{rToJo;~ojI^04^0UtHLEm&AE|cr`IU(X!uq*!%Qk4S7uJ`gxm6ujyZRg9GeaphT z#MOJZ{nI$L=#sI)2Lqj7%8px_;B+@aB|gpmjq3-Y&~GivsQfd*w9e3LvV>9%hq}i9 z62Eiz{2UjtXKZ|D)F04@I>OOJcFJzY91-=37p^ps5E>&m-INa=c%a!UlEJG${u^}g z1JU*)Cm-y5@M`_gEknu+vrGat)^NLuKF>i)WdKGlPwp9c2hb$r2$=wy%`1l1IVlHdgJ^yU} z;xymSM93^Uu`&$T7fhI)w@>la$o3t^pn^Y<#f+zk>wKiXlDP!&@`FFp|(rs@U1N6|!{qjO`J zFy^zWE5dH2QEg>8`0CMW1vbeYP$1#m|6ZqsoszO0;jeX9&S%y?C+KXi<8pwEYmUXM zHoUC8BC?%^!>T2frBtV|TQR&*-dR0P!#jrbNt!-u@*L@>CWyk}3S!i4==efAk!wOQPJASgrOKCN2^)yzB6*p<{8|h+k{umjgWx<|gex_`vZyexI zMS3wbRaL^?TR-Ma%m!hD=k-nn27*g^b zPtG!Fd1RFl1UpU5c0QqwXKLL< zN=c#UxC%3q!86=Z6CNV*abXn2&3XKl#z>%Ir%2z8!YWdpiV5CnKOyR(O3x%Xakk7X zU^Sx!)FAVgXmB{kJUsuqVWA&MNcJm!*Kt38e{3(_Byu0k_3CuP@2+B%&VLVcO#HBQ z2fOAlc7!s!V=2|Dctm(DD$BeJ!zv5>iTgqF*wMVD1TEZKd5^nBuzuiM`ikzD^51Pa z40S>31~l-Ic@nzp0bH0TBAqIM@hJ!)Wln4yW$iAlLrmZV<2AA2F>I!hD)KWK0wiu! zjQ+6!nI+pC$K%TD{Epar`qHfQ1P?^rMGsA<(N8uUN-u?);Y~f@ zCBZ8JjC2Oy(Ef%ZkBQVnH2Ht2Q>RkaM5#sX_#ti;z_r5`#IGjhw5Nm}aWI48NNxDy z&o2&Rwk30*BSQMDCyB3Q*0n6O_3%@}f8I7n5u76|mqWYD4W5a8;%VB6a8fVeiLu&) zsL>}N1gVvT(|#z4O74ZScJ-GV=w7|+!6N)SrkevBHWX9{jiqy1QdM7>EutP)UKz#F zY%bLS_$^R*-fX!KSd0E$n0N~GLLE*`m3mWutk3l{B<%wHuq3HICGz*i)?_5TaTvT5 ziNGI}Y4XRcl+U!jNLx@aeT;wETL{tG{{E$Ls+U3hLrMMkQoLY7d@n*u(FEu$MzuJG zgzQ~NzMN`%iY^1s4g?dfNuZkdN(LS>2^Pv}iu^D$B*K~muY-oq9m4%tI#LY^l_mH? zv@lC)PLR&GspS}cF)i&p9+UwZ16NaW_e+h%l}ydz-*JtFBaDma2smA*nCd)|dm2rJrmeBgFPr8ePtOc^7NPHKKCE#BJ{)pjj(-^v76(^mhtS{3?SzH(R$XyK zW7SKc!4ZL+_J%V+`W;0pn)2E$c|7ONEcYXe5(zI=vnC=e}wTkxRo0Z;pX#$Ipn)&VVpHmN2FJtm z3UabHgQTv;<&tCl@qRzk$c~blw|QKizw#21VbTQ(tIui5u8%pZgr=L@|1#hwub&xKBvM)y^6tP}kYbf2FY6CYyu&bbXF zXSR0FmRYQVKx4p$baI&KL`p1#B)ZUYBRA+&7IO(sQ@C{xooV^jlfWaXY?@05 zNH>U+%awjOR#vToijFaEjKvV-2FZ~zut=(5oNIMV;wPhJJ>b*XJSXihdZ2PtiBIhEr#IY@V`lSAuSd@eNaJ8%#jq=aCAy*_G=kYC z>o9z|-<2~)UIdL$-+$~l+2E~@+&%lw~1HR&qhr2^+Un8|Cvi8~AKuTxQ*HL8Zb zkgMVW$QPtDX)-xStof|muTJeu8jn|hFXg%lls?F5q7KqAgiP47dVjiJ&M`N@{cGD| zx@4n^@)=s0J$H}tcrTJ5uwN~%QDKOvJTrsEM$=5Tc6hw+e8+DL14*=8E_J+-ksK74 zIy6kzmEzyt#>p9}2=1#0rf62NWglqYDPg32nQT&<4(#7+;=nsPFbf%vb zn%GoWd-rilNv0mNi4k{lx$nmQ^hsA(=JCXvGt|juO6U!UAzP`=b59m0fb&x!b*23Z z$%ibE)scXkEfovbHdrZ)(K7CafkGv9DkJNuj0Ih!W0zsf+h5=T9abJ|q6x1qo}9>l zPJc%y)vS`7F{dY|3dO@wDASE&IJM=Y?S?08r~wOE&I^7t0UxTMNe*zIl%dO9`5;s> z5bqE0Z2RQb8PW9>g094;1CJ;Et_vP7O$qWAs@9POwE6GZ8NAIg!a1oP{T@UZw~qd= z9)AEdOdaXhDJkq&IkeJ9-97T%G4`bSUp!p$I&0WaXit<5^`XdOny5xIwD{ux%mK3i zik$<@Yb^Pg+-&uYtol`G#ubP_u0S^kK}If%HR?F|3#DGO$zF0&wjZe7&Fp3tQIxtj zsJYi|jm3C7>bw6xE-YDMy1FMOg{9)aq*02`ai%V&{(Jk@VIOM;w3fZ?np=IQln_U_ zANCuDeB(c2;m?_kg+s%-*1;=oxA97Q!t;p>i+`-T%{z2$Od+3)0!0+tj0I!UY3B_j zdOJBL921$RUI5m>;z&;S2Gk=haL<>s{y}OiV&sHydgC$^5L5VvI*`_FNAs+v$jY@=KaC8?F`M;bFlHIC+L8tNJPc{GkSvZs?r z1^h4`1sZx_QL(|6`$-ntY^TZSX<+9|P+f$Pt5OzNJSGtiSs_6igIjV?E`B~|KWMmm)xUrmam;}?>sd_wSftKK=#p^NCRI2MU;QT zzwi42KLi;3F}PJ>QA@Z87P0PrG%nv@q1k*wF9E*>ujjt&D&z?bE=!dV-dj{ ziTbYq-d{8p88flQp0-du%s4s2v?vtc+db2ScWb><08ICd7_@aQu$5?Qj_4gHWC%tC zW8OSr8n2~|0}moF1}@Vdlj$N%Ht)|VZo((OyZ!j(YL}_~w0_iAt6>fZTzq$6c~p_k zMf#W`!FNk8Ebzy#C@q~>D?1BK6x#DTJ?E#h(4V(g^sP7QOfQ&Et!^A0V;mw1dinr?_v(p)RIM*D|}M@k3%x(wHhKL zdFuPdxz6$Ga54mrx-{oA^n)HZHp?p7Y2=DoCL*fF>(3pCb(@dUFkkS^Qp(NdXe|-p z=Edc_+FMgqsuLWIKF~{J6R;~6TY?avx9&b|y!w~wKYtex7kIRy;+ZUVq?R14opK#9 z=ARo4;$2t_*zE^NfJljgCsP=Qj^t0TQquevD=L{Y^orxYV|wb^oV!AHrZQXQ$WQ7j zHHyk3zr}Hx=3n6JL)$r0raLeoMTnPTB~T9(NJ$rdAxxhAVLXFGNSXCYJYWcpy@611&~({r#d}#wmI`y3p-RKEov^RoU_Bm{ zbaW9jA2or)U7v~kd>nk#c>Qy8)4P5h*g?BsN`1T`4jtyn;6KNK;Dr>W4#eoWjo#0& zqJD%VQw^SzHni0zH2S@?6LG7{}SL z$Y->X4N7V90d~yw^l5sB1Re&WNW3w*>hV_16qNmclZZxL&senCjVwvuO7vT}OZz9k zx;1um5l$4i)kN2#=*zgk<{&TZ{jfC4BJd5KL_!WK)d(F zQ)SO%}a3oOlfbf`m6PTmG+_48Vr;c()}1l5v|9atn?obU~9Ka^0)T3IMI(D z*6o(|IeRhjNA#d<$)T@&!l0a*8HN_iM0~2psLx@~q#5HbcgUZ?o&4l&YmtVK7wWOB zcvWdwj%gg>qNsLx{!1}h_XPgJ=M@9>j=I{Yu*N-+xPU|=d<$fhI_HM|V$%q9R1F4D zIvw1iMWo?NLC40TaFxgyGO+OHGN3A_x_*jVB!A<+90;AF3csibNCo^XE~cm~FY_%8-Bb?x=e( z%6Ql2{xcAoc>KrG+qqWV?H_qYVNKfw6BUM=YV?|=ZOX7|n%stNV&66#N6Z9TRUV{L z2(#vbx@vRSLVR28vpB@+%D1 zcJ>^T)mM(R2}Qt;7R`IBr(gj8qFPc0(veVTY4@j&`|*Y6Mx5HpLV!W;8xk*%>j^k%T4`;XG;;%)dxzC3Ua;HxoEo^z_5yZJsGeFO17O>-Wr$_`w@h9xVuIA zR7aW!0>e!3%kwx=l)n9~6@UP#z$)*aC}r#A_QzQFAXh%}jF5EC`aRNne^PfK>zIbt z#dZ2*6E?PCXs(fk`qUCN}#sU%i4_39l#BH(T(cV zPF|NW%Do7p)`WQu4`tm2Or+b*!dgLYtP6CWrq?)8&$mFwg^olUn`Zi(QfL3|L?<+S z+YWUUHVIL!JkYzPW0N6q$ApEJJJjH=J~f>H+A(>q{D0VctDrdIHp~}iaCe8nb#S-A zA-EF=HZVwVhXi+bcXxLQ?(PtRL(l{X36K!ZWKZpGeY@xO+jmYNX^;k4NB<_+E2`6Eg8p}GWs0V2< zJv!F&zmBj8ju!i%+%K5Cby(sc|9e0s@eB$J#-=fE9=^5XJFS+XpSL8j$|2F&N#N~!RyPW&D@32JXmFhgfB(1I=koIQ^Th2 zT5A(1G0WC9qcfWfqR!Y%IE)N+zN(|qK54M^V08A+^~W`)L7+MKP^!0K>Jd>B^!x%v8 z|8nb0j*t1o6)&e5|8z}A%(FOrMu`*8%p+$+f~8U>O8tdpSTxftkVL&ND=qIr9gb!Z z?|1XN38lH3ACxO8B9`zs6{3mV0X`$g!@ADuWMg?|o(af4+qQ`jwlBw|!{o4*WasZU zD%f2`k}wNnj!X>;|B+xP(&K)OD}?Tx|J0+qTjHOt9Y#b8LW#7Ee|C$^-Gyn^Rr>{J z&3u+Bz6qTD15N!RhDCHK_qyngMlsX$9k^wie7xP^z!gk4tGHnnN_}i~glStp{(y}T zpLZyidpqUZ>$?B#v{2Jwaye8Y=X4r51O~=*kWbX<(wG^wR6>CXIfU&fxIU#!&OqGV z57$)OO%JgEiHj`GTDB(%6V10WxRKY!8i%Kj0LUQ`nHZ3G9h0F-6dp*jU>`ziak*!E&33j6fv(mPklHvNnQ`})Q| zDU8fphc79rhxgp(X*70XpiP9NKCXY9G+4r>%FjYwOWY$#Ck zgwJ6@8Xmc!_98wyVf%%2QqZ%2%^UlJ04I+l^poQ9s@QuBQLz76wx{O@D=r#dU9ASp zx9jzu^UXu{q$+Vky(&_jZ|K{HWTUm(PQ3DC=c_l{TUFr{3}O^Avywc~wpe?E+!e-{ zwtn~ODPS$h7xP{%7gl$D)=>PG#2Ex2!Uh}H`Z2#vG{7oRN)+wai+Vg*(up<@rxP2DkeyI7MCba*39jX89 zVZgbI|Bn-@!z^b`U;9>T|Bs7+eSL#}_RTw$zqa=OKH-0V*nixb|K73xctrmDDE`OG zwE5q6?0*eKus{AgD*oq)|KGvmzk|nrC+Po}Pyf5}_>WiOziY++oKOG%S_D&yn7w>@ zfiU0xfr+jzaew~`d=)dnHvgY20Ply_m%UdiP7^ac2qm>-rF`jhAPm=)xPnvykt=if zQG!T~JG;Y(H8*(ZFX1IK&wf47P$nkAW_7HPZT z8Lrb`g?WvC>PHuVypV*9iP_%j$vS$7Q0Wv2spv)w%7#mi5^v|LhsJC+wcD(=6g+U& z;Y2E{UtXeqEAvKb;LUr-=VEmFFJp6K(pithaW44d51k(wIJ=9cU>Q!T69q>SLsZGI8q$S?Pu71yq zs_!zg(O0sKG(~75tLLfKO9KHoD|X(LrYqC9Ix?Qy?KSG@yeRWA=U?=)^;!m8d8`4O zbmCVD#OZq21k{#klZX?%q}TOng9?CvIgj@Da#rbSa0u&b8v%10S6Jb|)7QUyuf3hk zk@ji5Pp?l>?|d_X0l&UA{_bf+q|4ka0;y->l{0q!ZhYBpe8a&{Ar?VX1D8nMYh^ZD zWF-C3!c8dFGAmcl0TfrQ0m;MmC11__969JXR)v{&UH)Fi{H${x{vd1Gu)I}eoxg!s zqK+euQR;Jh%?FD|;(`dc6w(@4GuOOQw;e;g`D+HRWpso`{CN1ZE|WebR(6gdX-nz~ z%U1O9JS0M$?yFf3?)Q)yI-68l`oy?m3z<6}(Yd*I!mKj|v08Y&DQ3Zy)XI)y zAgjd;fy$BJe5)TX8~M0#40X=KofIO8x`7?V^)Vyd-VN3sPv3G0Am;wI@Xo>mQ{0@< zFG{Q3B^|A3DKq;hi} zfDNCiRC9bok9a<(0N#Mido)&Ag-VE#TKLTeHogmoo`_lXtkxD7k){~JqQ=-igZO|l zNh-#j4ntp?8r3#bNmI+PBL20IoiZ!OzQi=|Gd)x{{B8_D`8q3My zI8&gEaU?Q=Ujb%tJY7vTh@PpUq1{`VO|M7C)%iaXcZ6mfT4#MWi?W(S2~%x zLo+uZkf)Wc2zDcnth85@KoAjltNdlD3ba((41>#kMH+C6E^s$p(?q~Pa{DwR;VN!L zY&~x0RCFW>U9G~&uCE_8nFwc2T5Zs0BBB=Qiw73pu#M7zZwk<%JB}gst zz%iUe(yp`8X0O^R#F643K7Y>lk-3Jqc*Ty&;XG($qOGRP$&szq%Ja?V4GMC0U(sXp z)nDCzArptoE*#fKKgxL)Cf;ze^`L@)oFp6Hu54Jx#kk zG2hQ%U(r!+pEWxfZH?RrfFH%br6RoB)Stu7!;T(Qu*_@sY@+a;mpgSzum7O)8**i- z0b-hlEnLjmHvBF;efGMVqHV0iD}7Z!z41kWIGdDQ*R4n4 z<5$yelqWtWkaO z_okXh`hT-wU4X*Zr@~Lyp9ijC;In-n%D-RsUh;!UdKSam2r42`Mk*6C!hH;2smNSh z<JExUM(2)wW zVW?BHq&V8CjaO@ZF{CD^R-Ue`E?YU@ zL^iP1AmAqW7vx&ymM99-9tgqU;xN&(XuWa)#W#2L(*WD$^DYan+BFAI?;Km$7ryqo z&7CQio9B~+mw611i6LKJxKeLeBhI8$w3g>60_%n!k*qWW{hXdhp`-wS6`|iy# zuAlXL&-EB77?roS+mk9g)FwJL{OB<*7p^1c9{X(tzCvEs zoUsrVs!eu%Ifsm3$;Rr@K+`f}wyRw(SFa^vDiF}KyR-g}9!maCU&U7*o3a(WUAhiJ zP(YhlB^6j-yV`_F`9=Pf2t|-K57^PV&Sx@q=G-S1@}1T`XgVp*#(xLXJSX*j?F^Xd z%=0;n-nD0E$8vOtwB3Xma9mA{G*R;_;730L0#cO$->Ft$5AQ(yn{`dqjF z==pwB0lF(lTdADk|E4ECWewy5WT$&t7?p{8V4v*;?8BVHU^S@}G}|sPm_)eVVAG$1 ze3@OEFUHe{te7td!ODZofkU~he8qG*F&Y!)x!?5p5IJp=X(JtWad-`8 zu!^1*6=m1x*|BG9S?_XrT~1#-Cwzaz77FG-L--x=+#;pR_}Z#IFMNdAspCK2_a?XJ z{VE`mqQ}QFD$FEb{KR9KSb4a6gpy~xI;arnAym39G5zuc<3axG6YrEqcIq4mJm}r! z9T%|A6Wzl6cPWt#!o4keHKS1UsG)z@6n{OxKerRT$qgxc9bSCp0FnEW5*N3+!ZrXO z(X~vb)^Ek{6yF~5dfF`-xC0I_av45c)ua`A@3JVDLoEFJu5c-khu%ORPFQGRfy^Wh? zmZ)1)EcOUm5ePJj_rV&0spRzZFKQNm8Ibe%{ zRq4_@%^!;;F(PPT8G3iV$37l-NfRt9s}>fkfFFB^0 zz89Ep)aHG!s3rq$(CFS8LLxsCT15DHo$?QYTRKnovW-iM%f(*m#ExpZX53Z}s(Y5UhN$><6=Bbaovu18zY@CU%Y5>Dw!(U5 zbW-0F^5EM|t3=P>E0w9bQ#*$DJ2NM7aKl$@V4E(iQr5wny4HEO58dO@c01`(`#c9v z3A;S5>TKeQ5jgr$BHC(sa_Yb9jIl2dyXn@`{d1<=h;r-j>`-h3hm^;bSStC2c}E)* z@O*`}htD8f?%77AFblc+KbVlaY@^i%-VIdQ>Fc+QINx~w>>X5zm_GRQMd59ts|zw% zeh5qq`f1oCQE!0hT*u2H()LSl6BB(dq1TKEvPTRrPNZ3 zwB^Ghvs;Us)8EDj`j#EaSUOxMOR+CdqEZ?GIbsh3cx#M&hvqvltUL1gTCepd9e+Df zY@OiiW6Jb6fB(z5m-6zbL3fE@8;(37`q0%Wp@N{e;Y!6l?4UfOw&Lj3_PdkcwP;L7 z;wAMCO$grCO#}oWI?`wwJpN5`au$6s5|ExQOm4{LXF!aJ-$(z~YZd?NO!uot;fy}N z#nqaIo{RF{TL1bFl(X%>V;24xI6c{fH5FJHBPCR`?)pD#YT~|4>GSN(1l&yb7+XxB z65DmN8dC>k5GU3iv%Y`I=XL#N=T|xd|Aw#v49>a?=xQ4j1W;eUnacF-#SHNc3ldHr zwfM$Sm_es=j*&@GFs7hZ(g2iWt1?svn&{#~RiDSd&RMj}&&F@DG;-hzeS`s%23}vj zW4IHW4LoF>x747rDEvC-{`4n>*nN75WkB{$q{xG12GgjSjYtjei`2F4im+O$wt5|V zruQFS`~1y&8*AYu#F>X}WK&t>-;#Qb?tbY}rQ76CwYa3%M1vrnlFZRA*GD z-i|{H?cpEudPATX6r*==AB^ZQ`#$P6GSUNL68hcc>nA(eODq|_)c`zdc{Pi_FzEs) z*Sf2ohJESSwa&L!JCsR*^1#sQ?E{M{^8BlQJ{@IZW(XFYhP1q<%%7V`sUMAKe}E2& zlXOC5AruM|ozCDvJY+DSFos{hZD@2uT&3vzo|))a45JV3UXEaUvJijHQ${&o_8?cR zO-DIzbROH(QiVJ6sYrde6SSmPqmzCU8`%PKa2P%Gj=9ZE`;E>+!nXkFZ`~HB>3xB%9YI;PQ4uRz^`WWLrE95z&Mq{PBRu;1-n#jRtxs;(^_B!_2G1yVzVMzw~YY z02;I{W5;wl9pHzIs6_C5@u2e52GYkY#$-}ZuBUpZS)!vKpv%Yc9d8Xw9Unf zgcIGG`Fh$y^Z^0Tt&=tLVOFidDeE@#Dd2)AWi9;y&sYeU{owgl=Wn&xzSQq4zPtWq zbNo~+B%0C5@~kk3Nrp9b%lfg3I=5+qN>rURG&;fVB1f^wuPN>>L&}TTcxLT6EKw!W zmtRHxQJA_Lt59lD96@H4AS-FgWJHh~+nCn+QK754O|u&-G6+bL?EGi_b)OcBH|ZW0 z&Wz^4+M@wJ%{A(Pj2>1-?}oA3c%(GZh0%^U^}n^072IFkjxr6j-Y{J;P5So7gWJ?s ztOqlC-ddo$gvGwW?Q`2Uc}NOzqQ}$J%)mY6gyZ3*I^FqECXT^U$I@X@*CD*xf5mt=VaE8sWpPbPu`zn~&DEQ@ zSA;zQ40Zad!Z$Z)_iz5bN6jTkh(eDPH|0-wk1#FzBVhtQ(Kq%Kt#onmD#rn;mWHK9 zSBeHV+n@9h0RYI!Ip)8ojFRHj1;*3uA~cT`NPa-yN^X`hC$Ys9EX`FMUcm%qMxJf0 zR`1`-5D>aJ*TY07Qkcac$W)5~$cN>C=RU79$CJx{+EjBT6)%k}EGf&3yp~j3mz)nv zPOsRKxTgej9{0GK1i@CxxWVslT2hF>hxLeh^f%aK(-!moqu#TuiS&>*JI4F7B?b)b z{doYAK|Q;~A6GbYz5K}Ogf#)MShMYZyY;^1nh)oh*ko@Wyu8kV&hwgZr$wbXH}E2- zq`HMYnF`KN&%||f$-%zdzDhy0$gc+da%q%ppr%!k$xqH;3sINLKTweToyKW7Inf0N zZ#|cov5&wk79m;6AwI`zzUbi?-6^S3mZ)s;uTY#mPfoqx_}sMth^Hbo?`qNdP7AL>m8&bU>SGqBNz_=x730Kd#{hrpR)c z&{&oOP6ki&j(P3X&_QTYJT?=1R-eEIkf1^dZq64wcG${5$m2T&y0G2E>diiublMA0neLv^wJGSnU>>+{9s`j-SS0RZT(eJ=& z*>d$Wp)|rnb)AH=N$!Xl_82tEQj$~_n$_^|hlGh{|3327GJ8(80=O8$U}ZmfJjWuB zpl@{=z)O&?tIB&bIz#%4sh?c0dV4=orT*q9^^Xo zd$d1Pr2=V19{ui}(BZZ{{!*gYzF{D*J55wK&~R~te~?*rR8Xxtfwz4~GBlnBNOglg zrih=OZ`p*wDf$*S1l+Dq+kI{D`d;>8Fpz?$J9@mxVdXXjg3m^aEsK@ZI3A7Dk-48M z-MK5<$&mvu;R1lwi3*Nq!Z=?vJn-_WUtpTPllpZ`%R4#fax81-V@}~WD(k3!no zg@f9HNi7Bzs&R9iK~}^L(td0>g>H`iDSBnL{E8)XhlNFr3l=Kk=hRd}mamh558vkg*#R8_DaTf+M)VUznqjNK7b(qC5xBTI$`IeHemEu`+! zZfLD{*O~`>CKvvl_RR6mT+(XqK*7QXNLf}5xiJHXo7Hfm&`NB6R<}KQOaaPd1FfRCE{HDkqd-B#tuE6p>_uljEYLra>3igER4Fu@y0I*^(!Dp ziqp?m&^QmJ9ETbu!50+4V-q6Wzx>)`wqucm4}z+1daP(Zl99VjJESf^~#<+0Z{7L1lJ}SxwX&a5{DkuDB5Z2JvZXq9?h}?<7%~EWFOFkYM-kBd96RnH}6uVz+eG*KVa?2Q+5;c)FCnw@|=cTpRl%Ovv zmPyB@qvaolK5qF{n{fwc{F|JQf6+W_gCy9sz z6)lz@PNgz0ADRT^Y6Bhw#mJ&V;xYp7a>S+CJ-g+C$_PP~#kbLRDTs4P9inl{v8s=1 zb>K{>W)h-C1kq{P3?$U5?~s{kFDe-{okZ9Fl#YJBW?~v}1lQxgRL9-}hmoWof$l=h z?b~(#3$xD=#q8UepV1G~zxi|0$OZzsTWi9MyYgQz8$bQs55-v>sOdiTU}MR=V6Z`6 z4VKGqcvh;-i7mype^e=U%j>zzcM%ZE2|83&&@fAUZE=^;Its0zABii)LI3y{j)d@Q z@42W>dxCZIZ@k&-_42oIvhH?Ai8}C{Ftj`_3!skYE2#@Jk1ku-Y1mE0>XtMUDzv`R zAi;G!eIu6Zm(nuwj0D?pzpkoXn<5*eq5%CeE{4QbJBH=w?yHD`5RIUO=aL12Yi3H4 z_VLf7Z^(kPGJ?}nJ0ogY^^|}e=zNk#h~Fw^@`jt^(RhBM8QRXsF%6-M7Z_|3*>bUz zWXGF$^^B+R_9WHi(RvbtYzWjwI`>5`In_`7qV=%iB>zoNMK!;QnCbq|1t*2&M`GG3 zteL6!Lgs)Qme;&h|5)OvHHZaV=5gzmVfLu-TR#luy>Wh7VF~)U3x_{%$MnaF@meN6 zJBieauNWsAQ$+No-^yreFVRHsglS+V*bo^t3}5-N;WR}54ACWJWWqU34$7ua6%nv~ zb9#Z8W0{&@L%6N5FPZpeCAG5ZD(d_T^S6U@)P#W7-vdR-hh<|`Pt_*yo8q_ z6~fWZy7*jc!3=+%67I18qdP{&H^H5*9O~2=@b1s&m4(LYNc_ZDProAVSxd%JDFza% z9qDf3MJHAzER#+qHlzazY;)k`ulr(gMoto=kN^XucmdLv-M}?J)$QZckdYAag%^gk z(VMAND4Gav8*#l*F)y6rV;E#GQe=)-6ND-9W{{kdnG#HK=i}~fWtyG8v#8VQMlIoI z{j(dsu-0nHic&4>kz+h=Ft|uGxIJmMO9T?$l+Mau>ai3e&vWfS;)g|$np#w6D+!Hv z+f2cACjP#4ALekm9kP^>SzIEeektu!%&eA`f%)&08 z?PJMu8)oOV#LMyVmBr{{A2oRQ0ivYkIf_I8Dy}w%&pQZrW_TQYO@~muGZb*iYgZ9YuFcl^O+gzNSotatmZPKKSc1TTM zcFtilTwd_Y^{;fDPWIhoc%tku`pY8Ers{{9EXB%%%WBbvS=L*shFvgaOpMiosU9P7 z2(LB9Sz+ftqfX08R~Z7y$qDf1MWvB`C7}6u9fw|MSvKzRh0VGT=Y5SsLq09uuLfeMVf>`w*vgV^sLDUm1(*S{tkMWYQ7LZA3#>QgrC z?p5JrcS{VC%)a9TIe6^{sdbvge71tXJ^Cd##g9_qef-|(yZk1-O94b0t zMdgZKy3*4$=LoehqxCbJ85X7-=Oeg~!rP^lC|BCXMe84Pjfkrb1=UkY22B3%KtD*Z zj>DBfI8@deg_4ZQLfEpY1nfi9ohe3T$bksTrx1UNFy`GUe+Ey!V8@|4Ej8iH!D^*& z9+_%c1SDwX{^I*8ps+{JMO4WBPb$*`r36QSKD8mthi0;-P)rB#=_KTa&$6~JuW~#C z3nEEPTp?G_z#|RBgM;e12ObX{HF>5|+caO+BrG(62c;P3imB}L7}I2nZA3gpNFMVmTFS`?!HP*41<-e!km4G$iq@0eSh}fWGNb}uPji?zWw1$ z4TuI0A6Rn}zt{Q~^MVht%#CFu?5P!X}Q|I;2K#y;^!OrjW`#5f}8=svQ{=;u#)$xJ5PCg_J@ zASaw1==*voyG<*O;H=hYCke}R0lU>NU0!msyFaUlban^X-e6+1zjmv_5swPU5~J1D z4Hph*X98aBSg+mt2Yak-$wxe(VFpAbPA=}Y2{FL>a!beoA7++{L{OHzG!vu?81j>Z*%-64 z1}4s2EZ0@>f_ZJe6~QMfPoS&0Ur;IU!Wsp;&FVf%$tu}$6aw5oX_5HVpj+}l+~FTt zwW*K76@#{#pQ3n&+Rko7O1p?k6_Mnbt~@Y8aYn`6q;33YD!#>)0tg83Rt-PlXn82d z%rNk?t9&@E+%l5EtWRW@h;8+G;^615<*l2J4sm*H!|g}EFmlzyjXfH?wE(2)wrQd< zm_LuMOuW(j%h~*r!uN)SGZFY8K=()lUojEF>k3FY* zIiMVIUX~zdT)?$S9qXy8*CR%aLCnaoqn8@PRbpz96%I))RVvblj_m#0j4>ka)h%rx zbcSA$u^>)o&vj$+jGw4=Qz^l6(Blzj>1mf<#ZlM!=|9h=~RX{1VDUdz_R~oqB<&p**bQ;roM0 zsKc(5H%?Uv?xpiV-QTRC9J~e|uBUA-1>Iln-kYIx)#K$3%7B7L8i1kp&_Rk59)gf% z5;w`Xn@OBI+%|yNhnMplyD{%0P`$Ven4*f3;`P=x>Wcz~zB=yRkyDw_D0sLW?+UaC z7IMN|brB3>AghC%i$WoS#&69hOkSg6<<*Tdj_fU~UE-R2kkz$D7W(7m!C{5dgg+I-43v@O0^mO`=*aOU@EQhpz%ua+)RCa7m|9nyZ7 z<%zZCpKrBf2@VsINZ0#7#p$nKO@q&gsbh!iC*oR2V9=z^$ydrqshSgG{yxB$rDuJ! z%R(2W4n(nt!9&EfLx^h22q4Lf+I)!O!3*zD*;U3#o8?bG!Z&Rq<7M+|lY|Gi4&Dgd zMI45%iH|SD%Z#THcg~(}=W=f1yRBPj>PJ!56n;=-84+o!Y5Iouu|VesOF2`lh_PVo zT9k5-6mHH6(Xo0D=}A~wS+=ly4yRUCQic$di5jkz@j}D8dYRJr1_46fGpwa)Ne3@; zM9|T;YW9Q|i`%0KevVMzWf(49 zFA2d_3h1xOP67!-Q?7aM1oMq8Y&bD1^(gX)2;wt9vFHfWe9P2B8`IYPT9lRoiM}5R z1N1LI%(Gz<&k>exLphoVlsEF@9Vc0O7nMpW_ILi@`y**_I&=;Dv2}V)#=83Z+m+&S zTvs?v7pQ!@U(<wv8skM}kmzx;ASkAjk2a%`%*ePQNhk5CZPUJ&!I6)3U-4;XxdVumx3M2rAF`x* zsvURlnbj<{=TI%>6%8m%9=9j&7yAHnY=?*|d34T-RxQvHli-##Ckb#pOIU_gRB1a{ zh{?-xVeoQp@aOY`-;#Hmw;&=_&#+<@B&OZoP%uI>nMf{?tpX>*nsQcM5fw!_oLhEC zPdZ&IQGFU5FHS9jezN-J&cXDXmM^lba4+|oLGQ#Gn9S^>G?cLDERSoRSuxy&@PlF| zEbtLyqIsIVGK3x3BGU~#?13wq7xTq}wgDR-0G(UB5OjBO$(muDhok(38?cXN`ZasE zlz|y%Mw^WnEe0-D4OLbGRX74H39klbWr}X_7604tYQ)M zldI4i4^8#r?u zSHc#0(z=Np9>I+UD`{?y`cM^GZ6Ocq4czV-HKMzkX_hy*y(fStNsxC^)@h^3Pka~n z(ysA`xQ`tv!Sr;bw%>fj&$+L^-!m51Ej4QYli7?%_LRnB7JgMv45PBsz7Dh1`%Rf$ zM~T1B8=tT(R#Fe(h-8|1t0zkod(N6!Tsp2)3hukDc{y=X%-v4H)bv@%5wKAaSpMud zoiJxCaX9UKrul^#@Eh~W`O}_c@;Rm^Ea++P zwI`#x<730^#2{ZzkD0P3+u&{Y>OnBTnzLebooy=nmKhb$rz=*A>>e{Pf>{KXe{6u7 zQoUt`!G%BJmhNsz@A@ie6MpjoOoSfATWUtl!+D1UR8%DcX_rA4$v}`qsp%UYI-e~s z{Ux_qpGF>)e8Xs2hOSZDJ2K~C98cGbWxAC8TibzZ2qXRwv;Ezef z5~s%f&RA3-#ML?uh|%Sb^n+B zG=}$s9qoahY?9k~aLL?5&E#HgCusc3i~ZT+bo`!Z?Tsu#(fLYc zn;|Uv=;H{bsf4dx^9*4IKIz_K^}RSK6tj6UCaotTBP_lN2$haQIn&_!kY-KwPURK$ z+*DL01~t`8R!$&9pVmYRet>slkWU1wYQ%i0$=b&hc`J|Qj7thf) zQc*^a7~#j}hc72)OhO_eqJJIh)dAhNE)q23@!3YFkF2$GMIKY>x|hh`q>l>}+oHm3 z4y~*EjF9${dCIWHw7Zq}9=W%QfQA#b}5j7}1gVTL9X<%o04cct7 zuUi+Ev4>7%s<|4Syy1XVUhpb@pgnJ>`&YsEr(Z2e9pWlSsunztYfuv@0yaI;q)VB4 z;&>dGdh()-1-n7qn$Av+5>-})fAL!4$(*^Lzsvzx2egUmpy#YN50S9$q1g_Es z?))4*q43%FYa22oIf{wYB7QO!qF&sChpb zzZ9Zj`tBnVAb5lT!Gd+|2&&L6BbhYRCq-AgDo=Oa0BEUMkn?hcrLLi6mH3+!bw6go&P zxjH{A@l;k08weA|%QnTnq%CQ%b1A>1fnv6%tGaGjWAgpo7@v@Y_ZFda5uvnvts{qQ zg8YV-uqx%v68pT;u5ys2Jh;i9YLs%Tn4`n*jm(p?g0Zu8U0adhy1Uy*-&Y2gCO?I1 zT5N+x!Heu9_mqg7Mr+|MQM%R+GDcdoDJ5dh|s|n4@l=D z!BI!~p)Cl!-Z2Z^MGuijgil|P0hv0*8w&>CcXwm|e1G6+4g#~1B_sb%d#T7HT{o6+90Gn5o`bw>Kkpo zG=thPF}oPqj=ZxR}Ue{pp;pDB9I@u zS0^3!YezvMRlnMEwBbpm+2jM`eVr=ZX9OnlIerjR+tOjQV(DUvR5*QJAzG#CzSU&o< zd-+W2cC|}tIL&SkeCd%}a^)+hNNU3p^88di>*DU|+EWB?6yXsHT z4x?$#k+q>iB9SO1X^3*YICO7S0wi(8eAQ^(kVl2gDetmIdGLg1waOS(Q;0d3$oljA z6AwEk=J729P*gj0_q z;)X-i`r98l^#HnW%Qo zyPO8+qo4(H(7G?D3m9q}=&eN-;N4Cy-W zcuNG9vOCQD`Ia09Jby^UZX3DT?!(qyA+55|kUGd{bg@0gl)xwa_x*da6^A3DVJIgP zs!v0;^M_t|)xo$pXG+>28Unfic`)AxDx~7uEoaMPPU<~nF}biJNK}^m@RSktVvlbq z{Q@o-Rv@Jh1pu};XI<-i?fFLumR(oCG5n+Pi7pu2W?`%@(|x8hjN# zN*z{l% zX`zzIU+~E&WZ!FJ%?*v}UC7JHBbI;F3qn;)q?|>-e7lF+Jp2Q|>=#9?;SaA(SWQlm z@JdP2vw&37fd00rX#5-Nd;~Zv!wjnD9DUf0ndLxy5^P5|BBvu{|GfDg13>vMInCEI z1y25=A{1RdkNs~FKf2BloUPw6elNKsnOVA*dmCD@BCkGE4t=^Z?_jd?=bm#iMuzMJ zy59}g(Qw1Q=R`&6K^3wJZeimfgp%=2Qgp00w9c#{5Un8?DK0R2ar{~I%g7F_knn3^ zP$>5NuW0XaLGV-%IbnT9wb2=u$wiJ>hJ9_-97XJ*UOZn7xgVXbXtd{YRYy#|tevh7 zxc?*6y#LYHSx&Q1R$Q{r!g9JHDRzUsVG$|Unah5Ro=eIobtoAraIEZ%{YTFI1577r zi;5ymsVHlUNNJlwV$2M0sxN1Xk{281DvF{atSp4de|Z&ho#dcESzN;v9zSHU<ph%22-QmhPsADAlEQjLY-0@;VFDMncDG3`o@ zK7GP?;4-FV5JYl*<4sVR_&w%l=ckwV*5ANN$?v}*7j*vl^#&;?pr_s4gaS8U0Aw75 z4;gKBbc@5n6WTX(ahH))q!lDP{Nj&i9U6qRqISXDbaJ4;h;wh3q> z2ngz*c5NNl{TR2Bx@_Fg;6~FV0*HST>4vw>U}pca9{tuH_W5HuFi{YQq-v0#tY<>Y z#DdUvyme(_w@$$Vka+`eMk}t!c2)oQV@YwCtnrZ>7X$A0>1&}X_A4xXaL?eEO3uK^ zd)4i?rw;l0cO?asIRHur9K~XDCCTCgm!L9aWjIBQiPCht0T1p%vvp!LT~1eo_`083GLUQzcLlTrNL(P~`h&x% zIh<)t??ggBeBABZTvZ2V=FJu>7CgPfrM0#|SGgkgCEh}aI_;y=fnk6|uo%&GEPk-_ zU*gNp&gNJKkie<>SzYoAufM7FwRnLs7DE@%g`^?dj?;biYO{pDBaJI*h3jSz>7};H zqfP^mr*Yv=t?M|NW#Q_#9!okldFr`lf)M{yVpQtD>6BV3g~pvyL<-Py_L-wK!Q4`4 za=)TlJ7E>Mq<_sLqLhx>fK+&S6_9oWE*qh&UVnG@?Q*$ebr4(ck&#F#SejMtH}=c% z6afhbB`>xQw`^pNci&c+lM+}AEf1BgT`Ljx4~~%N{wK4k`ACGXLA>>Ku_V`TZ&ky+ z{mb`y8|=BzPs{+{yokI3>g$qEswD7${+DMnjJLdEf{O`A@j(MU&K|t`dF}zCj+5T+ zT=(&QZnlGQ3j3}^8g2vZK^zRB7LpKGG&9`KikAGVfM+NO`iO9bU&rZA_OH zoR5-8z<1+9rfMg=fuO zq30uUQ3Y*4^C*?tDpN2=2xqL^m-}Y5A?v6ZIso87@qz;%0Z2u*UjHf{<}C4V1GFmo z9Ysh;h)5#qt{RQCP*TzEzUzkuo$ z2R;aXEd)?V#OEMYLZvrviE#vzzlxaPj}Zifp`bJQ_P2?xyxhu0Tye)}$4gC662xhz z*;|sA=-CbgJ?bhd$jjdsYu#~Fgz&@^K!VIn{=y9b^CGj_lQCVWHPFmlcO&H;K36SP zxoDU4MTu$%6o&`a3>++;3{GE-(0dZBIyZ5^!lI2Iodc`s@P|=Fgud@On|h|nvAH26 zWg$0y^h@)lfT#5YXCgSJG3#0{Pt*NB0L4H$zdUswjqlM=KF=aqoxeiU8WkU8z2vi~ znN4mBgbC{Z{`bFyCL()NXhy9f3rt_I_Oj?W* ziu3!>o`E$37jLDj&`$3exRe1kNbRc{c2<@3aX;5^d^muaUO51nMHzu3t;1B^O~Iv$KV zHvBAUhzc2iXm=6ma>!A13TDUCqwWaFsphgw%P*);pgVa*(*omXM$* zdIaU&W2T0%#5$E(fkuY^Msv7#Ft#8BH9b%GrZr^BC~7=$hj4W;1nnqtxZE6le->#s zTHrz|}~G(nn6gGV}3XRu0JrX` z-lF1U3z$DRk}NGg5G3F>EER5wMr-g9XJPNdq+SABuAqi9KI8C2(j>(DBuO)pOLVDR z3N7T^7#Z@WHSn1F2+UC-W^LHouK0GDWvpfrz;EFJQ8OinKVE?C0<}(Z9zA*_fo{pv zEwHTDTU40jRT#TWT2FX^l}gN(EDku3Vs@X$_Hq;1MiB^3h~K@yCMqwqu&`HCC)-*;d1yZk8*K#;fL6 zZ_8^CG?T&;zckX#)n*nrK7aI1ABKPs#@D#P8S3~(MM_<|m=4EKQ73k=biYOmpd!y< zE2@)79-PZoV zu;YBU890>H7w^{mQ5=;vr4X%nA(S>8$C%>ij8X5<6!ELk)q1C(`)I!xDS0kEgdH9V zqO;a<1;od~Ilte0>~jznQ4XNR;$U+{6T6d-+x#zHSFyxIM_nUaWncZ5CFi$jT25b?q%N7{ol+)+Q4VNsU2XtfO7NLB zlsbg7jTUn*vTQkbT_*A${wShH7pyR^y#0k zaDh1Z#Uh_BO zK@Rc&06+jqL_t(7Iky~}#X7WdT*bW7xc_k}G(t{g(PRoY1x$LRbaZH>lVWJUj!>J7UnU%AccQOlB z0-JYh2}kpo&mqWz+SU6EKb0?k81Hk>)2B~4(=9g1jWaVfg9e<83C+rAm6;JkNV^3f zA(n?XC|7^kARngukddn=Po5xbd4+|%O<}c5ls79r1F>7ZMrO#USY(oLm2fO3@b>g3 zUGtLQM{G`t+U#gv#^g#4O@tm_u-+Tnus8_0l^;WjfD~bRj8nX`o(UV(1fCXf1aP*f zctjg5?dSxOiEqdm?;xY<^fd2OOwiubaPXh=?>fre8no8kYW~vaEoH*^%0=W~D z*vOz;$v8kXp(S9c3&e`WaYL}tEPI)00T(GPvDAu=lNvsOIAd!=5TnQzd{ZVxI5LTY zZ(qnigl14brDlVJuReDO0-4d=@EDUnHOQFx+{C%_%(Rxa0h_4!=_w-`%B>r_xa* z08yoLZF2ZIJX`|EctLVD7^3bJ@fS68hxvjNZ{&fW8Xs0`6=I${FgWRi^V!nbhA&?7 zgwb<}B&oQIsK66|Nm!3LXN)0=I1nUCTGrzVp!=WR*-1uvOQ{=`q`{}6&sDi4;NoZ$ zW{n|6oyR00yWbGQbcw{#s0_jy@`fe@#En5*4zSJD6K3YpB3uqPY20A4E}YyhbsXzf zO;RdywT$UK;zNed4IIuRHn4X-3dE@|16H7)n)`(BWxy@mAN%=#x8j)j<;*1f6jQt_ zlULm1*t5JpA5E-TPCoxRZx=Af6IyJ^stGd7hGQSdbR!R06PkVM>S>lOvlVvauMqM-T69YkqnGk2_I8t%Ox{#_jOK>Loc|Ee zWA`+K8M-GCf&}JRv9m(GAgX2*>TxXDahM2~OfTN*vo+7PwI!QwwG11;CVZz7c ziSt73f{&RwlP(Wv#TE^OFrp}U<&j;DjD`@_Jj3#et6l8WiO0~tq3c3Z% zH&!ko3uGM~I>LH3gQisooCJ<}df%T?sF^kaSsMAK|!}IWy9I^Zis8Erdb9+=g8r9>h;p6yP4}`Eool7%{iJnJ_UJ*el#x=y%c z;lYHU8O`m5)u*hvRD=-xfNU7~G!W&!AlB}Oh ziozw>BYrZ2fCEHN#=qcIEqQ3Nd?763`h)!MTXycipWOhqZPfLr{& zQgORNVMzJ4a}grS+{W#<;Zc$+?{Y3Kero2TN$0%X@%O*~-M$uuzhad}t<;e$)$*h8 z>Mf_3vbu@R%Cy-{A~D%B$}}bwnMBBy4<9yQ<&HUS81>)!o^{KgmQ6kwMYK2b1f5;# z+W26!M=E66l{kvIz3&k+&lS2#B1LZaE=HJ2LCFa}{Vgsnxl)-qV@S);SRcD9k>q3# zlm!pXx>8+1m2(7{o0}VJ-f8ZUA?8vXz%g3j`M-ut5}=$b&SuBqWh%wOuH|YLD8jT{ z2|z0+2q2mvlErTJL76812`hhgj>(*g%E|F(B_K;9c~k7ulEu*a!9*}5D5=&nL0HXM z)ran=l^>}f)}m5}5j^B7CgCn@IJ~m7U?S)W4Z)%t5l@0;-3FJ6_uhLChZwYiw&~pF!{b85Z>kudsL0 zGR9oFr5*yNJEBVi!llQR+6|jnPCB-&NsrwWndk=-^Fl!QKmdlB18YR%(xmQmG>&+M zX*pR)UI;CC*PX5%1Ir5E`q02z%S^7MXl_NUj70?Ba5xEEZ#l^!E7r^^Hgrx(??lG| zb1?!Ng6HKhzDI3nXaSB^qwZw#Y$fxbW7e>gCL=%#iI(U%DZ_)7t|6SUQvUV$Jn3FF zDiPmdFC1Q3#^N<=LgS3}EHn;_&i4yRh<=TJmL-5Cs%2nWt3ny#h8jXr^EVEhY%0s4 zQH!wmqEpAjm7>1U3|05co%I7vm>NQBHz62e!bV10kXE}j?{{V^`SIh&__BnfM{{6I zST^i3=c`cXg-Xaz!YqQdRXxr+nMHEy|98ciiXlH5pE-%i^IVcH@gDMZ{2P`8&_*W@ zRPg0zW?_vK-J|K2Qt7qjn zaoZgE_JjjaQ&$_sRNTxm&mkM4A_R?}?f|@Un1;sN0MGg$x&(kEAaHe0sqrY=aAq$g z!z+cb$BMG5BYmK4B;zZ*>0j`)s@5G*ght9zM_WRHo-tk|^KxGt55mZ>! zO-*T5slEET#5i?aeS8k@x?|Kq_4Mge=a5^Yw_=AM#TP)|f(J*Mny%uDg{vw%nqR8F zQvq=d^zgccFzU<@tRxu1EE|{zZ+C`VR)~2t>W}~;a6KlDgU8{dcbRq-V3A5^Pqy?V zDVF4p8M#FmI&IuGN`-t;=((E1BOi{(6^E0>LcLo{G0rrN*5F(xM&98qG>NqsL`utX z>q_Ac;Uw-#!Dt!2j=32vEiv=f!R!r0AP#b8vm01H;FIAGdHLzj_06fj#w5%E-!*u` zfuf(~D9w+cHS>>64X)>g}M&G0eGECr` zwGkiF+pk%Wb1my!I{5KF|KrW`$3On@C4VOWppv12u(~W@(?I~IKI^F@JQ-l!hY!or zXGELJUD=zhL&4fsQn)vUC}slTEyem#4ZJu>tRid;5au)FaVcAGRWBU%#46>aZPi3n za!r$UV%ZDRrbZhTS~MX5B&w}Rwfe85On5An?;z&C~rno99G*KM3+5i5p z4Ze8z=~2Q*cP03vQ6Whl z4j)Pme2p*x8_hl?!}BCzYSyH^On}2c$Q$)&MF|8`x}X8!MrZ^mu0HJao`DC=z&EAm zLCx&odj`rtC9?Wk)l;LLF>XB040j-Z`Q~567*Yz=T-DhDU`WyQrs4Gi!dsY3FRL9h zJntY>==;e6r*}ebcZlk}cnHJ^f$0E{5Ql^?oo%9a?sjx_sE~KY05OWDNZ6xS9VQ&r zgr%#@L~y3l#gPNQj16&`K*%wv*K+KvN--wSA2&|-36Bm)r@NU#<0o*vKb=jAT0Y3C zcU^%mq0!A2=RI!h(!~sufKl`@)@Q#Mpi{BtU1FH99YpOavqG(61c1m10wygP<6+>d z7COazqs3doh>ci}#H<`L0xZMakohL}tV$+Wnb?7px`{N!j0Vn1kF@C4h6+q=%6It4e@=3#3nfoPnm-7mdIH+;UyHZUN-PuRQD!{K6+`0# z+T9gB z=l1s20igtWcj|RQWoXMV0a@Dw^)@_OsHTI5GW_)E)BG$neEt)Vb4kem{>xa=VlRX; zAZin5D83%76it|DIh2-?ro4Flz%;Y|F_BHJf_V>A^ej%9q)3u8#Lzs^Gr?S#8H#?{ zKA+QSz42DVaT>B^uEi23F)A9m&=3Vn7^g^2mu|XaqST2_sduJGg0krZ5pu{(fGSu% zIB@wRx6Wc;7!Bpa2#1f!hjBwPog+>|xrOY*G#pIxNS-q|zB0fUpbtHI;HG=x1qHJJ zO#v0RBOKqpL0J-wX8p8>{J{qwP-Ilu!i4&DAfEpsc{J?Zk3A^$o{6ckOOW#(_5C7ggf$UuN3aq!JGrhn1o zn(hX}J(ns&8hK8R%6Jl!t7d_in}H$B5m+DX3(croLeWuSooP%hEmoM*5SWm`N)CTW zdZa>9->8{#=MW~2wbx&9&&_7_1kEDTz~OL)3M=NqAEs_3=aFz?!#T}E&Wr#8r|1|e zz?dHSX1Q9rg5ra?R>kN6B=3SH1Nl%APUXKVj#TiFZ=4inA&YKSPBU(Zhh{3`b$Kf- z@b%Kd=sGX0ABJXKWt5CkU=85&_MZL-F!g40v>@PT7hsJ{!g?IN8Na=~<*Jt67SCGjVDW*s_N|p4>B^ z2gS5^JRx(bx50$R5OvmbAR-RO?)1dkOBm6g2?)QBry}JNE32Vs3rX2gFr~}q z70YZE)h3E&4b!W#>xw{R%<)7;jWsp4yfh-nv;oEIZ<@e%;5yxaWAc-e94J!A7PGxT z-PV8!A%Tq#^GMebOrF~mMZqFn4O$yeo_iN{%$fJmJxee*eC}~_9t*ssBBd4&51lNo z0sI%5o?V1SYxQ4>Gy>FT201_AJF#crp)>H!QhsPlJN}*l8Bk1>LM2b{{lQUDc~ywL zbtu-ALS@>)<9JZPl}HDPQI!;e%6^|S?+s_$H`GqeDdpn;|22;lcLE6E_$PLlhJ@oJ zaF+D*qr=0g=~yC(sh1*0i~9$8!ZJY1m8Fw|3<1X{COW%4q7C^L+qDMMyF>z;-Vj61 zHjZ_CQpCBGYPO2kQQBX6)AdTvs|~V8IH%0vz$wgCOhUS;z~_={A`b=bHsEWBdx=qZ z7gwdT76<$ik@I|J$RhOp1r95|zc|0~xuh`)K`P_{bI=^*)~aOEG9(*0pB6g&dRhcP znXjhL;wPC3nP7VqqahqUIs`*}z+6?B_6S*uX0k?o#f;Xw^rRwEYpgJ5bR}4a4;FrI z7-eXcIi(3>@;O5;UOC37Hj$=SK6BoX7KbDOefKu^DtT57x;TcD%*Fw?Pl=UbJIH21 zY<(2RVnOQC7KRM$EN;>J>jRJxut-sG~ED zW{6^@Py2%_$0iCXMM8k)P}7r#z^Ekjrt0w&WH@w7xSo3m_9=^K z@4b!+r;&=O{Ar4;>zaz`oK|z4BC4v3frt@I~D z0MViAlPxVLM3M=h9jIbAgr?pRQcn?6q7Y;ddYoLM;&63}JM*A1igVhyP2IdW@*Pwp z;W(>uptRG<2Lyc9nQ)A`(-TI!AK;&YU+08LJu$F6@u!bYQ%wH!>MQD8@+czbx4RnM zek-_c;&c?UMxDbhg1$Al_CfZYg%$FKBBDK#k48A)sYnjO`s5EBjX-IMPD_v9od93P z%C@;UIPpAiGFmE`2J5AWVrZ?aXb9~W>SQP)@~0_FIDQTg0LZ9ESD&mWd@%fC%MGKn zk`z~_Ule@diLhsj%^!3KL5eI1AmOZ{&@fqXy%vi{^4*PC$ziL}uJ4V}&C;lq*(d=G zoki-h>bB>t+jUBGlGYHh4Y9iM#%Mn9M5NYjX*3rc&QCsf9pq1=ugfuZ?IuZvZytN$ zkuJ_ODTr*9rCIaV4{y6GfByOBm%N%OD8SJ4OsMUU5gOR=^$}z=>PBa+W<^f;#_e{B zIpsDj^0`VLp!M&OjGiK9=Ja$61bCR3Z&fQ6PEN?}GO_0;!jOd_E<=KqfrKXyhg3eV z9wuuIu#~O>IF_8EYEEq>I|~|Esg|WaS4pPr=JSD_+kuNu-cHou1EB?Wd{~X%NzMs4EW1QN1f8_Xi@tVU+=wMfYnXa zQxQ}&RZ@S6s%kqn)M&3aa`*=niv@`Iff9LQ7XXEB@L_-E))0^HB2gU%QuY)%(J=j3g>>1ZvY+A?mZ@m zm9z4gZ0P}-P;Pi+D<*!5^LPmapV1j5&$MJ(0Pbq85$+$pTNA){MUV}PfkmgOdjk)x z5TRZ1bHG=9+8qSmfpl(~vw{&wjvc(i47N9zs((^d?<$z)JGEHZnB%;f(dB zSz6@zDY51G)Fu3e^oRL4He%Koix3b1#RbsdARP4$-O<#cptT*yyCriNU%GrPMVL^H zdTS{ty+APC<{7fE1*I}r2_6e@wK%5V1|sJz05QssT;U)%kMJlo<;bJB#O_h93V(va z!52=;1{B(;98UNiiQ$uMkEo)`ooq8F&jbkDXqF`b|NT5)r9qW&9ugxUr9dF-0gewI z@7btm6GdL>EikiOU9vlPf++)x>d_h#Ol$lhU%KMoX9Amh`uD!~Js;C8?12pd zxJc3RZ^)zkdz)2xe2BLuIo1p0SS%cm5uEQc4MpmSLDPDa^X5nhn9f`>D!k@vFY87% z4j(f0{GWty25tHXL{ld=ZfppWCX469T%yb-4*&i1s_Y<`b7*$tLEfnM$PjZpiQp|l z6gXNkV9k+U!ME=eYD-9r@GYd0Nz3ids1v~y^moxz%wT8`LO$!p&fGKba2eR&ejhHk z?=-T~reaEgI;g@b#|rV2Cr{MDPG2uJno8=FRRKv>?)@2RVii`$m0wE~Yv)R*NJUd^ z?Xg-+R}Ce6=~TE>d81ApC$!il+1Y~=RQYG!VgIXN{fb>G8ZjUoAUswO>Bgi5mLW#P zi`eW;04Sf~1zNp8WbK1KfbA#pk5i_h(dwTqYcx4pgIbfn#bM z^Ph+1Au+!(8g*&M?jCJ(*Gd@I!h(#*$CXFsbeA~!)3^P7?XyJ34JjL90b(2+d ztxu(%JTY|gj~_qwWk-zeVL}L&KuO43m~KVB&`YUH+b5rV;=f`-7&3v(%{s@C1ZZG| zU_;5$cpI$86v?=M$es@hZgQDR__S};@t?(<7loC8D?ZG@tOP{;_Q+%Q3Ijdr(d zmL8u)!4^W95u;Z=*|TlMXtB{HLzn~M;n2o__oWLE%&&R=^D)t!sC^c5+ft+CRaklJKIY5M!De<2_eApZ=x+! z$h;-cyc?p+C7j$V1iA?6#VpAbr6 zhWLq-UCS~LF)JW>6%JUVW1S&uhNSK9r!X>dWj68M*FE}oB2EDW@KwvLx~~S-7)&`7 z-ObGnUFLFNwVU;|FI~Y}kW^%875FyvW-eqYv_Jge4`*Qsm?=JJicvFV*_%)*4n%LEIgeJxIi@`J_DH}K^PB@@nr?(D4FoYI>cv*+cq9GQgQBMW}PtdZBMo7XO z0m76h!eUoNeb;~~8+^FEy_J6|fdp%`M|_;po=oNl?KIQp{~el>h9nsd zyR8AiT5=rl+dBcrnTp3H0_)d=EQSwutOMNJo`HwRz_VKX5YRj3yPN^#uM@)YqX?>f z3aAq4(OZqmt3;~8YTT4Lw{!Ox$AAK@Vk-JT6iS?Qo!QjUQ<_ytNrpqx5GXg`vuW>m zI6L%G%;KWL)*OMDI35MaVaJ%}#BzFdbwfLc42imfM9+DOg(D*8mYWCIm;ebRlZW^o z$!q8y8p0CG+H6vVSKB$4Dyo-c`b+C5*VFp4FoqQH(&x|ccq6Jku#7rV$|bo z#O5)hoz0lEW(XmKBhojsvyMh*LO7tT2gkp!A%^fn%MBclVzC$`$D{Rc$jUKsTJk)x zUkpi?H7^xz67Et)bJC>KE!${D{kmuc;s7^vB{@dhFR24exG+KGSfhkBb|cs_ah!Lf zwd?hadf(%J>d-WICQTBAAYLkVx8wi;JH5E90T+70wh~3wg_(PLe0hN51(*+BkjQrl z%b!g@iujgX4J<<&$sTj&>`=}?Kr|_5#bTB*R+>l5$ykN142w=Pq^5CIsT2`9_sl*; zd`XtH@ogqzFvk-^Qe=e~k`@Tp=RL;=Q`FeCT5fJ``hGF6LqL-O@|qa+eA9c?vEGRB zHjA5Cq=@ZXUh*T08>3y}h+6KK}S)Wx*mL;N46Qq!gKP+u!;JWl6a7<9l;c zL)dcR#vF62tqo0BGC8jS)C-yQjo7O_0?&s-f}ZaYezw$4F-0dvDuR+&6Q@>qEdL>p zp$J*eb4)!-ZYBOs^IUK$3ft6=#Ta2xCTSKExo)%k~gA<(((lR?D0X-&_?D8#n*{}!& zqNzf*l9jBO5H3+0=RG!G)VbHFdSc;u_2s8(Cdt;7-YCq7!Qm=8uZFo2*(Fwttf<=r zOo3CcYXc?(!J0&3+L~>;T7FU#vd2syN|8reK+Ls`U>|+?<(C!%CWp-?Adp2QX$m`|GH=5w3lbV?&Z!c4)pe`@1*IpQ2L#4 zj$UyUGCl!f9TlC364Du9$l1_AZU`(V4Kb}D(n2I-v}{X}U)w|)j$?_gAq}T_hxzyaqOT~HhtD@lUNAHAqG@(xt((v>cJibreDEQVG@Cfd}ydb;KkmMaUu z-AJl4M>xW0UP&q4R)|UGtx2)EEE_TcAz#b-rYDdC2XCGq>f|?K^C)T2L-a^DqG<|h zv^E5ML0LKGSESvrXZAxDorHT`lvT~`)2eodbAzF45{YRRkIO_0kGxSCvq6~1hhl+% zV1hi0dW#TvqZ{`&k8m{1grk|pK}%qCS}UNbkqymdk>^6-YA>b4!f_F|aV$}f?xXIP zE=em?Vy;I*W5~}|K-2*`iAXh!k#3SlFdbdyLdTS7)_wDe>&1eKWroM=5Q$uR{s z3$(0nd^>&Q4@`EsWWiXmRgY6ioWW#Hql;O2q+VjVgc)q^CB{lxZ1Ge1GMg;=*~)e% z^WJ;!nR;)wAp6&{TCRq;atrpCg5u@!C+r-~}02kL( z3F%@*-pI@}*+l-1#zw=DbDLCEuh?ctF{!hJT>P8}Mr|}( zEXS@}yyVHw-!t&-GqAt?e)}Bm2w61RNoc)6g@h$ z9I*;I4l5iJtRvGQ%_~LU>4L^|#*iUE%khs8K8_O(hd`%-bh&&8oFNEKjtEea<0NyU zFq9|y&T_rv_(V+}vxFeDTm_PcdWr2(uMsvKB+HUKYxo`wrC#4*77c_&-g+#?b#&Dl zn!qZHo`m30m~-urB0rzX8(qh+AugF3%S^(hckG@Lpa>Li zDbg46sW48SByZLj#q>zUqa-jbSBafPh-vY=>bbaE^Rj17@qvsYH{kG121Mj;;R5PT zZ(7R_0i#qH)mv#mU<#a9R6q?HD2|b|gXEc_okax*Nsm`tEg}MmSs`e>{1@q{LklPC z(qgltXG!zi1n~;fD1^`~8@rm^!gy2c5f_f$+`X z(1jp&tkpp&bO|%xEG=jW4@iphcuJ6`POJ@u&RcJWR=ZVEXpL4b2ymDpfAq=M z!(SaB)|%mPrcJ%F57{Q7E{7JFBwOchAU2&zVN_P}jiR-zkFB6iA5EX*NqQ+-R~gMQ zZfG+v5Jc+kCuwl}SM>Cjw0SQCk6MQIe$T+0$-oO&+c#r!xAHDzpt7#~c_|gITZ;>g zbF0%GUkUb#<-Nr{MlDt=jXDe6W)PiEdN*3{dzwn^yd!yLdjTJm$r5Q~I~rQ-2`bv#p9l>8&c;hvgruA!!ep(|G| zisS*hT(YZYx9GC6=@Vvj_>2OP1T$+H2$QSx4l_wx710`VK*%Gy2~l*tlQGvdi^MX3 z(>w5(+YrYD96BLuMuXjm&10civV@w<#4FR0^AzQZGlr-e;*3SB$y=UG-a^PFLpGFa zQ^8s=py05}F;j08h^+qq_U8m0Q?fNBOXMfu&OT@IpABs$sP;RWf?fd|<_zj4FkdV7Fuz`!j>d8>gsa2*Dh1r= zTC|<2R!4~fDE<;P(y1~FyJ8idev7trXFFB8E{e8y*{<12p7+!l>{O{t#{`g1S5~tY zQCCFflmFs#>Ym=~w8lFo;P(~Bvk23eTT`*sZ&$=zp=Y%;kZU`$sLH8gi=_3LMy<8x z+kH*5gx|8Uy#~_bhbUiJXQ*9BCl|n#ot!B8_TSy*2eqDC#PB+m9J#B zbyu@3q-W(nH@^m}a`b}gD(G`I81Gw~r8w)PsE%Qm@U8I7QJ;0u5axzi%jGl?eU4#g zgKf?9wM()a&P>)RUxCIx_n$Hgpf%Ys_tNsy=xe=KY+@L(TRi7$XKQnY+{`u$-2u8^ zHLYzyTZFaQ*)~q)cMt4T?U!HJr8^3vXdA+moJpK!%MY+M&?q{}tlzpT$ewqjQBKRR zdu;cL3U}YA>+XBa>{OY-o`A23T9`J=Y;o3gebs5A_`fUj3!l`=4X|R&F?rNvAwb^CZdhRL+q1kUC zb(@{viRM&i=a`3AW6Y-`|GH~N>ZsOA*Gos~qAI_VooW$qZq9Yl(brDpTo_d; zyz16$IkjB%fBoxUji{W)d%YuNdmGCcrUZfSAz=6WJyf;%;eD=~UAMxXJ>498N1wOh zXu4^3!|gq$hhH15?v*`~yA@Wf&-1`5zjQam&UUZtM%k`o{<_=ar+21#l-6Kn<`LCa zxu?-X{dp=pJoQ!p~l1y0FH)OCph>$9<(HPAay=^|zJ}IF3#`}o0xt4d^A&5P9TQ1`)p$n@)^(k> z<{Q_9+S!VAs!kht#in}sr@KNLx2}|4q3W);pMDo;%36mpPDiKwEfHX^DdCjg7;CD} zN!0JmLeUmaopyG*tIcvIzjVvB!u8maZU|MMaaFziCc5CK{8TA_5@nwJ0w`U{TFa_6 z*@&j2I&I<2+N;4CZ-pnk+lbEBYE`}ET0D^!TgTKrWlmyj&)1V^9;;H<1zdx5`mBYi zGG$87xcVAJ3#saL(TvPx*eLp1W@qQJ?NqPYBVVl;o0JwqXU{Mj&IxR~=(A$I-FAJI z=vwZmO1<9-Un`dG?2M>HF?39WW!_58*Fx%43#`SQKu33>a)RxhyQz<8dYpmRU}dJG z2Gt{=x*gq!x+vPB%r0e`jiU0clRy3GPkr{tNU16<&elVhTZLN^wRXyv?q*nzb=TOt zA=FDZwElKC0pC6=)=}mJQ=*Q#Qf|e*`0F%2)stxx+m+Zr#^imom2X-bPN#xu1{3IO ze|0S9%)*DD-@=|Nv;q~bVxy?#;{WSk|2mCVtSRWp*aGN@(;7&xo$ZuwYg(IScC@rw zekE$H2)q32x*41Y+_W}x*Jr0D=k{4^^*F^^0R7c>euHSHilJGneDU<`?m}+To1m_~ z=CpLvJxh8*>a(w9HpT|kF%QsMsP{5EDl@14_P4*4=u}B)e|UFz+hNY5Xr3SQ z%qagnv&%mfi?jIiNv1e^HJQhA1$rz^Yx6{WF^UPGPTLstEUCw;PYhG5oTAOGuXE~J zV4iNz&w3>sx>3ES()3!J*-v0<=0WO-&w{m*v4A13Vo{)Rb`7hA_w?_sC)RerrU zijHQmboi~$Vi@W&YwCUVa%RT+^Hi&%*M)x2Ddw|i-}gSl40)DE)hkouvo_wBuRy(2 zU=l^xi%3&kU|oj21v&lNw^0e0HiKN6eJJ__%H~n_1Fs9bHO}fwP*9vFVvaEckPUI(1;f|?0Ex-S4$PK6XyY2MeHgy|palS8H zPK6)lSFuO}Nv6-5ecg32H}+~(tnT`(N(-RbX;8%+bxQ5r`mAJUProhZc`L1>DmHP& z-^J3Zn|6+B?Q~2pRcYS43W{(F&)J%-R_W9JEpdt!bD?$&_^sG{t*LrXsnt@kx-K1N zjTGJtXG(UeMEQrJmK)lPt_XWI@4A?YsF4}k)T)c_NOjR>X*Owp{_~&b1^PraCHq?8 zUdpRpy2j_{uxGp9Q$7F0s+)h2pZu>(T%DWK#y&^gJ#4u$#n7n|b2QymX){yi45~!U z&UpK*e3ISkDp|9gDslQP(OtS97He$^G|YZ(NO--5#kfI@>^IIDM^TM_=n}uCglDNDgz9 zUrtAheB!ComQX`0_>L~pZp)pme2tf1v0A9gl)rse>Zq@IZcU<`>(}n%9k*hQrK4hnT4Hv3Rj*jjhDn^y61Ss~c2vEVucykKhkE(bay}I( zvtFjPQkGz}I*=`)X=x*fGM!PkC| z-ukWg+Uz~G7nGib^>wgX26`P!qm zlC7ds=GM_EIf;p<0oKv{f*x8`$|;P3?_ob&BX;7ywsX4ZrJU-vQ+PV6{VCkn zn(fpr!ueX=7G4ewq@2Etq3&iZC*WJf-VUe0L_5=07*%WqmRW;c zbX^E@s=!j}=t?Qn7ItM?t|82#?S;7V^-`H}I6W3atrb{*QmdWHx7ur`%;~)h49G#XJth)-;t4q|fUVRA!D&wSM|cyF@LQA5n$Nsa74$yp&%g zHQ3pZP3GscsLrex#UtPE$Oy zP<{oXu913d<+Qwe1k`nzRVo%}eFqNd9Z<9$mkq^%sTStg2O+1cDvgYxM%C!Hx;6rW#1y z&CRgb0_Cs0;l_9H{ZGM&x5)9 zau0@X!QDWo`V{zRQ!je?F{g@ueqQZOXr4XiQB(luargs*8P($HNUt!`>nN)8`EA3f zmpYyH>vaCjRqsaQ&|~;?_mvK{+yblFUK-|seOmtMi3a}0{Qr7i`Nu#0(cf5lm+C)( z*xUTKifv|g)>YlL_3gDI-j>|a_~`Z0Dt#DFuQd>k?!Ylm~4xC&JL_0)8u?q#~`sEWO+ z&inW4DiiA5Pg-mb3m4?%mk1<*K098l{~1?}_u%L#@^1cZ;U2bsGHWyQR(>nFMBA`) z>Y5me?igK+U6Mil|F{3V(?@--jr=z0s5_ea=Dmk(IDdNbUw*1gT-mlTx`(%_-y?j= zdU#)Ja#}^tPCd-OK1LrF_*t2)!?})T*EjpjhruMD*jzrRz}Y&UoRg^hq^r_=^>b=J z^mJWRrAg^?Oy+iRYW?^`;Y}p1mhMP_bp3XrJ&a`TlT)Q3%w>35cvwXyfiDxV`~5OE zwZ2@STYUGpZpyQF?Izwk(EJeiu+e!cbRV4eyzZWr?`tpg=b2ZVot**?8~KTu#nupt zt(Z#`e@DSj&fMP%uTvf8>vsycE^7G-_15iw{bGGR`whd$q1Mg-o5YsEtbyk7+G!zO zzkW&#HQsUBspUGl?x;VdksnZ|p-r6i)%!~S60EQN2K4Oh{0l;w3B5Yc_0k*aTrU;q zeInt?JaewV!>&)Wx3kaMxfNUc^~36XJDP%-g($iY3`lx&^tD8v&C#4XiTQd8oUfCZ zT9dBX-h=xWzwJ-@t6Ce&`TN_cQXtRf@w#|gHf?gcLOV{~TcYyUu4}oRJ|Av}Zgu-w z;Z|CoFUnt|EuOl25l@9H&@p{YFnvC|70iou`;))_{qMb<|K%@#>0e*JkMp!xa*gH9b)$j@{9LK<>>0SioQcqgehs zLG~;pzadN=fdhzfV<#fPfQ(nVk^VyKz`*}D-2hsTzgdb{? zSE@S!V17E8CvrRpUNO(2=%$-Im-?OYUmrDQ9zzn{7smnuOd@@w55FjCC0g5{pXFN2 zXK?^8zWlj3${18c354)h+&cA#f1&etg%3*pje^aXzqF^tMzgOil+_V|o5IL^u-;N4 z8d%isbc1GnuB@e!m99~PJjRckKgD_GGi(ab1=#wyUba@%{SUU1%;kD8DAjSa<3Zmv zxl0cEW@ySCk!Hj=nA9Sbn&e`a(X;_J$NZ?fEA%1dJ36zlD6fmHv@%n)QDs{Q5X1GO z!BMLLS6H;3HCbB;28h5PEu$Fjm<)AGMM(4k*Xt$p2ud0LGc0#&s?aECvt>ZFdISbJ zpHL32vkLII>QGQCuv_6?7G2AVpWc?Cd(vRz)Rt}zQ70IX=pd!46j1L}32_OY_}a3S zPH9qwpqr?T@OM}w1`Nd9oxuvHm7c-w#;lMy7tSJV81=n8x)e4w*%(G@AuxenPdG$+ zJU3z^x=3URHA$Nyl)J!h%ww4s0`w@$u&J4rYL1lmr0%dJwbv`xz1IUaDyp`HLmfHz zYHaN7_;MdSUdj%al{$^ZyzT4nKD;BA1bIPKc_N&>nA=fGszpF!fj3Fegw#ou!ArrGGd@A#yQqd!R8ZxeSV z(@5V9CwO?BX0g!q)XGX+vP&&7?h>Ib{J_>sCrV^hwZn!Bq?$W^g{b2d_*A%V{IiMT zj`sAV2U|E|m2)-|j5``TUG7C*Fmpdt7S64d9JmOKCz>YyP3Fxe=*E^@MOg~LU$m5u zPOcsa%6yH(A++#i1ST$0-2arBKabc(SDS zRa4MnXkSJCi_PfW?2`B5Hz*f6+_hdPkQ_c6v(R(GTpE6nd}g|^Tt-eBte;Cy_ml>W zz;^`98MBH{I;7ArT!ZD><~GL4tM$^J=1u-PWxG;g{dc~WPhjnP>6ON($DJ&PWpaff zdit3wTY-gv|5T@l>F_X(PS_;mnV_5rPuzcGJe>?4RSY zRcfEIW0%TwNWo3Y?umnb4GAcIsk|&Ir4%wx{Xl+A4RV*qcx`1cT4sj^MoSMr;rg|uQAB7Q}Ab%Gx=7 zM$Vop$8qQP$ena6YZ463=IDAY9~|Zno~Ik>&U#lsbl<`UWr1C#9d7AX8O>q7Y7;(y zOm65qiVxNl77NTq3IhR-$R5-lsZRNWs#7mFRJ&<2M|+Xd!*D)?ANjS@9dl|HAF{wy z9PjUL0NZm0K9hu7kX4&yLU_xy(d5S&rY4uLF>-nH%wa`Q5mNL}T*#=StH9TT*3mH7 zn=yqCMuY=zbSu^QH1(JKj;&Uh)hl_o^_wBXkNVx~Uni5#M6OA*Z%QgLX4EStS$)(1iW?yo?Wf*jpkl&2?k#Qk;I{Ygs>WW4%IP75d)7$BWmr&2}Y` z@lJaE2~-`Lmd>B9rNg32bT4l9qw;z(@3s<2;-#!2#DO?%XB|EUHgiO3xpm=O*eCj` zE%cW6ofb69*7Tu)?i^X)hp0C3ZU$*PD}iUz?;Oy z;%`!v_TsHkN^Xczwgu3q>c@UAuv;kP$%=D>BOIfK3B~hvP!!HBc`#lucWe)Bi9l+&l!#$TbGiuKODo}uw2lwf z{uLv(9iwA@IEVRP>@r`i`#_Nxn<^_fwjWRRZ|B&5zqu!LfBV9zJh5LuLvn~vYgYNKxao&oAdM^q&@*rmov3-ykNXnlcsazd^ZcDu9aBgb612)W-Tcs(!eF5M3yWo$Od)Smu^8TmxQL0?i z&g|<0Of2^OxA#Z!28-zW=ZL?_^mLShto|=7_WxCHTyFv(z_hjgI)Cd|PlX!A-_58o z&qp*F?9+hak0YW1g$5Lw`%!a0YDiPV3V*t18dlJ-f`%0|te|0q|C<$*&S~l|O5aeU R +- +- +- +- ### 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_deepseek_v4_dspark_n7_latest_main.md b/docs/design/sm70_deepseek_v4_dspark_n7_latest_main.md index 70563273f3..304038abdc 100644 --- a/docs/design/sm70_deepseek_v4_dspark_n7_latest_main.md +++ b/docs/design/sm70_deepseek_v4_dspark_n7_latest_main.md @@ -307,6 +307,33 @@ scoped greedy output remains byte-identical to control, and the quality prompt stops naturally after 76 coherent tokens. All task-owned services were stopped after validation. +## M=8 Real-Expert Grouping Candidate + +The fixed 48 one-row grouped verifier still reloads the same expert weights +when adjacent verifier rows select an overlapping expert. A default-off +candidate compacts those rows into the real expert segments already produced +on device, while retaining a fixed 48-entry graph contract. Empty tail groups +remain device-visible zero-row entries, so the route adds no host readback. +TurboMind keeps the single-group dispatch tactic but uses its ordinary +device-offset scheduler for the variable row bounds. + +An exclusively owned V100 CUDA Graph gate covered mixed overlap, 48 distinct +experts, and six hot experts. W13, activation, sorted W2 output, final output, +and changed-route replay were bitwise identical to the admitted slot-grouped +route in every case: + +| M=8 route distribution | Slot-grouped | Real-expert grouped | 43-layer projection | +|---|---:|---:|---:| +| Mixed, 34 unique experts | 0.42345 ms | 0.30437 ms | -5.120 ms | +| 48 unique experts | 0.42407 ms | 0.43113 ms | +0.304 ms | +| Six experts with eight rows each | 0.42183 ms | 0.15401 ms | -11.516 ms | + +The source gate released its GPU after completion. The route remains behind +`VLLM_SM70_MXFP4_MOE_GROUPED_M8_EXPERT_ROWS=1` until a synchronized TP8 +full-model profile proves that production routing overlap transfers the mixed +microbenchmark saving and the paired acceptance/quality gates pass. Evidence +is retained under `/data/models/dsv4-verifier-20ms-expert-grouped-m8-r2`. + ## Experiment Log | Date | Source | Test | Result | Decision | @@ -340,6 +367,7 @@ after validation. | 2026-08-03 | candidate | Exact-selector seeds 4210-4219 | TPOT 31.282 -> 25.533 ms; emitted/chunk +4.0% | Accept speed and acceptance gates | | 2026-08-03 | candidate | Greedy equality and natural stop | 256-token SHA256 equal; coherent 71-token stop | Accept quality gate and enable by default | | 2026-08-03 | candidate | Scoped-policy full-model smoke | Target 32.426 ms; matched emitted/chunk -0.23%; greedy equal; natural stop | Accept final scoped selector | +| 2026-08-24 | candidate | Real-expert-grouped M=8 CUDA Graph gate | Three route distributions bitwise; mixed projects -5.120 ms/43 layers; distinct-48 projects +0.304 ms | Admit only to matched TP8 profile and quality gate | ## Artifacts diff --git a/docs/design/sm70_deepseek_v4_dspark_verifier_20ms_long_context.md b/docs/design/sm70_deepseek_v4_dspark_verifier_20ms_long_context.md new file mode 100644 index 0000000000..66ebe2a582 --- /dev/null +++ b/docs/design/sm70_deepseek_v4_dspark_verifier_20ms_long_context.md @@ -0,0 +1,217 @@ +# SM70 DeepSeek V4 DSpark Verifier and Long-Context Ledger + +This ledger tracks the 2026-08-24 TP8/V100 campaign for +`/data/models/DeepSeek-V4-Flash-0731`. It complements the historical N7 ledger +and prevents short route-hit smokes from being reused as verifier, long-context, +or quality evidence. + +## Acceptance Contract + +- Eight V100 GPUs, TP8, one request, DeepSeek V4 Flash DSpark with seven draft + tokens, temperature 1.0, top-p 1.0, and the production strict rejection + sampler. +- Primary speed gate: at least 100 pure decode token/s on the pinned 120-token + Chinese HTML prompt with 3,500 generated tokens. TTFT and prefill are reported + separately. +- Primary latency gate: complete verifier at most 20 ms per speculative round. + Complete verifier means target forward, target logits, and strict + rejection/sampling. Draft work, accepted-state commit, CPU scheduling, and + total round wall are reported separately rather than hidden in that number. +- Long-context gate: report raw contexts 1K, 4K, 16K, 64K, 128K, and 252K with + identical output length and sampling. Report verifier/round cost, acceptance + by position, emitted tokens/round, and pure decode separately. +- Quality gate: compare no-speculation and DSpark on deterministic GSM8K, + HumanEval, a multilingual/multitask LongBench subset, and 128K/252K needle + retrieval. A speed result cannot promote a route that loses the paired quality + or output-health gates. +- Shared-machine rule: launch only after an event-driven free-GPU gate, never + stop unrelated owners, avoid frequent polling, and terminate every task-owned + worker after the artifact is complete. + +## Reproduced Baselines + +Current latest-main source plus the admitted SM70 route repairs reproduces the +historical endpoint target. Three 3,500-token runs measured 118.481, 115.975, +and 116.288 token/s; median throughput is 116.288 token/s. The third run +accepted 2,816 of 4,781 proposed draft tokens (0.589), emitted 5.123 tokens per +round, and had per-position unconditional acceptance +`0.8873/0.7818/0.6911/0.6047/0.5066/0.3880/0.2635`. + +The matched no-speculation route measured 66.014 token/s and 15.1483 ms/token +at 1,024 input plus 256 output tokens. A synchronized DSpark profile measured: + +| Component | Steady median per round | +|---|---:| +| Target forward, M=8 | 34.295 ms | +| Target logits | 0.326 ms | +| Strict rejection/sample | 0.499 ms | +| Complete verifier | 35.120 ms | +| Draft GPU work | 4.419 ms | +| Total round wall | 43.505 ms | + +The 20 ms goal therefore requires target-forward work, not sampler-only tuning. + +## Admitted Source Gates + +The following changes have passed focused current-source gates but still require +their stated end-to-end promotion gate: + +1. FP32 mHC staging now covers verifier rows M=1..8. M=8 falls from 0.05383 to + 0.01257 ms per call and projects 3.507 ms saved across 85 calls. Four V100 + numerical cases were bitwise exact. +2. The M=8 graph can use a bounded 2,048-token short-context bucket instead of + reserving the whole maximum context. Speculative buckets require an explicit + `VLLM_SM70_DSV4_DECODE_CONTEXT_BUCKETS`; CPU dispatcher tests pass, while the + synchronized TP8 transfer remains required. +3. A default-off `VLLM_SM70_DSV4_PRIVATE_COMPRESSOR_STATE=1` route stores + single-request SM70 compressor C4/C128 intermediates in bounded private + device rings when prefix caching, pipeline parallelism, KV transfer, and + parallel drafting are absent. Seven V100 ring/workspace tests and eleven CPU + route tests pass. It cannot become the default before the long-context + quality and capacity gates finish. +4. A default-off `VLLM_SM70_INDEXER_DECODE_CUBLAS=1` route gathers long-context + paged FP8 keys once and uses FP32 cuBLAS scores for bounded row shapes. It + retains per-head ReLU and FP8 scale semantics. The captured-graph test, + shorter dynamic replay, graph-tail mask, and workspace-width gate pass on + V100; the incomplete long-context quality gate keeps it opt-in. + +5. A default-off `VLLM_SM70_MXFP4_MOE_GROUPED_VERIFIER=1` route extends the + single-launch, one-row-per-slot MXFP4 contract from M=8 to M=2..M=8. This is + required before a shorter or confidence-scheduled DSpark block can reduce + verifier latency; the previous M=2..M=7 path launched one operation per + compact expert and erased the benefit of verifying fewer rows. + +The focused M=5 and M=6 CUDA Graph gates covered mixed overlap, all-distinct +slots, and six hot experts. Slot grouping and real-expert grouping were bitwise +at every pipeline stage and after a changed-route replay in all six cases. The +mixed-route results are: + +| Verifier rows | Active-expert loop | Slot grouped | Real-expert grouped | +|---:|---:|---:|---:| +| M=5, 22 unique of 30 slots | 1.98837 ms | 0.18552 ms | 0.18000 ms | +| M=6, 26 unique of 36 slots | 1.32993 ms | 0.34895 ms | 0.27260 ms | + +All-distinct expert grouping is slightly slower than slot grouping +(`+0.121 ms` projected over 43 layers at M=5 and `+0.210 ms` at M=6), so the +full-model route must report the observed production overlap. Evidence is +retained under `/data/models/dsv4-verifier-20ms-variable-grouped-r1`; the +single V100 was released after the gate. + +The 0731 checkpoint declares an official DSpark block size of five and stores +`mtp.2.confidence_head.proj.weight`. The runtime instantiates and loads that +optional head only for explicit confidence-threshold scheduling or alignment +diagnostics and fails cleanly if a requested projection is absent. Ordinary +serving performs no confidence projection. DeepSpec computes conditional step +probabilities with a sigmoid, +uses their cumulative product for prefix reliability reporting, and truncates +at the first conditional probability below a threshold. Before production +confidence scheduling, this runtime must first collect calibration data on the +pinned datasets and profile M=2..M=8 service; a raw uncalibrated threshold is +not an acceptance or latency proof. Static N=4/N=5 trials use the same strict +rejection sampler and therefore preserve the target distribution while the +hardware width crossover is measured. + +The indexer crossover gate uses FP32 scores, signed head weights, top-k 512, +M=8, H=64, D=128, and CUDA Graph replay: + +| Compressed keys | Fused paged Triton | Gather + FP32 cuBLAS | 21-layer projection | +|---:|---:|---:|---:| +| 1,024 | 0.06854 ms | 0.02096 ms | -0.999 ms | +| 2,048 | 0.13455 ms | 0.02519 ms | -2.297 ms | +| 4,096 | 0.23433 ms | 0.04738 ms | -3.926 ms | +| 8,192 | 0.45660 ms | 0.07680 ms | -7.976 ms | +| 16,384 | 0.85016 ms | 0.14493 ms | -14.810 ms | + +All 40 rows retained the exact top-k set; maximum logit absolute difference was +`1.221e-4`. The opt-in crossover is 1,024 compressed keys. Generic batches, +noncontiguous inputs, more than eight rows, and non-ReLU semantics retain the +existing fused paged kernel. + +The production block-size candidate verifies M=5 rows. Its matched crossover +gate retained the exact top-k set in all 30 rows and measured: + +| Compressed keys | Fused paged Triton | Gather + FP32 cuBLAS | 21-layer projection | +|---:|---:|---:|---:| +| 1,024 | 0.03548 ms | 0.01905 ms | -0.345 ms | +| 2,048 | 0.06881 ms | 0.02196 ms | -0.984 ms | +| 4,096 | 0.13532 ms | 0.03410 ms | -2.126 ms | +| 8,192 | 0.26614 ms | 0.05806 ms | -4.370 ms | +| 16,384 | 0.58173 ms | 0.10066 ms | -10.103 ms | +| 65,536 | 1.92835 ms | 0.32707 ms | -33.627 ms | + +The maximum logit absolute difference was `1.221e-4`. The microbenchmark owner +released its V100 immediately after writing the artifact. + +The speculative SM70 graph context candidate uses explicit raw widths 2K, 4K, +16K, 64K, and 128K when the maximum is 256K. Longer requests retain the generic +full-context graph. `VLLM_SM70_DSV4_DECODE_CONTEXT_BUCKETS` is required to +enable the candidate; an explicit MTP bucket override, including an empty +value, remains authoritative. Automatic non-speculative selection uses only +the compressed-index operator configuration and never architecture or +checkpoint identity. + +## Long-Context Diagnosis + +DeepSeek V4 does not use the repository's classic Flash-V100 dense-attention +decode route. Its production log identifies packed-FP8 SM70 sparse MLA with C4, +C128, SWA-128, and top-k 512. Sparse attention after selection is nearly fixed +work; the C4 indexer scan grows with compressed history and was the expected +linear decay source. The opt-in private compressor rings address capacity and +memory growth, while opt-in gather-once FP32 cuBLAS addresses the repeated +index scan. + +`benchmarks/benchmark_dsv4_dspark_long_context.py` records each completed +context incrementally. It uses server-reported prompt usage, separates TTFT +from pure decode, snapshots speculative counters around each request, and gates +suffix-marker retrieval, replacement characters, token diversity, and repeated +token runs. When given the TP0 server log it also records interior synchronized +profiler windows for target forward, target logits, strict rejection/sampling, +and their summed complete-verifier cost. The verifier row count and acceptance +positions follow the configured speculative width instead of assuming N=7. A +zero speculative width records the matched no-speculation long-context speed +and output-health curve without fabricating verifier counters. + +## Quality Gate + +`benchmarks/benchmark_dsv4_quality_api.py` provides deterministic HumanEval and +LongBench-subset API gates. HumanEval code executes with Landlock filesystem +isolation, seccomp network/process restrictions, a private temporary directory, +and CPU, address-space, output-file, and process limits. The sandbox has proven +that normal Python works while `/home`, `/data`, global `/tmp`, networking, +forking, and signalling other processes are denied. + +The first paired LongBench subset is `hotpotqa`, `multifieldqa_zh`, +`gov_report`, and `lcc`. Official LongBench metrics are loaded from the pinned +local checkout; `jieba==0.42.1`, `rouge==1.0.1`, and `fuzzywuzzy==0.18.0` are +isolated under `/data/models/dsv4-quality-deps` rather than installed into the +shared runtime environment. Every quality artifact includes SHA-256 hashes for +the HumanEval tasks, LongBench prompt/max-length configs, and each selected +dataset. LongBench records retain their source-row indices so the matched +no-speculation and DSpark scores cannot silently use different samples. +`benchmarks/compare_dsv4_quality_results.py` then rejects any reference-pass to +candidate-fail transition in HumanEval, GSM8K, an individual LongBench row, or +either final/anywhere needle retrieval. Aggregate scores alone cannot hide a +per-sample regression. + +## Rejected or Deferred Paths + +- A TP8 M=8 64 KiB hierarchical all-reduce was slower than NCCL in isolation + and saved only 0.028 ms across 87 joined calls, within noise. It was fully + reverted and must not be repeated without a different communication design. +- The expert-row grouped MXFP4 candidate is tracked separately. Its bitwise + microbenchmarks justify a full TP8 trial but not promotion by themselves. +- Disabling per-head ReLU, changing index-cache layout, reducing index top-k, + or weakening strict rejection are outside the quality contract. + +## Artifact Roots + +- 100+ endpoint: `/data/models/v100-dsv4-0731-tp8-dspark-prob-n7-latest-c4-20260824-r6` +- no-speculation: `/data/models/v100-dsv4-0731-tp8-nospec-latest-c4-20260824-r1` +- synchronized breakdown: `/data/models/v100-dsv4-0731-tp8-dspark-prob-n7-profile-latest-c4-20260824-r1` +- mHC source gate: `/data/models/dsv4-verifier-20ms-mhc-source-gate-r2-cuda128` +- private rings: `/data/models/dsv4-verifier-20ms-private-ring-source-gate-r1` +- indexer crossover: `/data/models/dsv4-verifier-20ms-indexer-crossover-r3` +- indexer source gate: `/data/models/dsv4-verifier-20ms-indexer-source-gate-r1` +- M=5 indexer crossover: `/data/models/dsv4-verifier-20ms-indexer-m5-r1` +- M=5/M=6 grouped verifier: `/data/models/dsv4-verifier-20ms-variable-grouped-r1` +- rejected TP8 all-reduce: `/data/models/dsv4-verifier-20ms-tp8-ar-screen-r1` diff --git a/docs/design/sm70_deepseek_v4_flash_bringup.md b/docs/design/sm70_deepseek_v4_flash_bringup.md index 750f8a896b..60be98350c 100644 --- a/docs/design/sm70_deepseek_v4_flash_bringup.md +++ b/docs/design/sm70_deepseek_v4_flash_bringup.md @@ -55,6 +55,40 @@ kernel module and userspace NVML versions differ. Local validation is limited to static and CPU tests until reboot; GPU correctness and performance run on the remote V100 host after the PR is merged and freshly cloned. +## PP Topology Follow-up (2026-08-22) + +The original bring-up named PP2 x TP4 as a required deployment candidate, but +the retained full-model evidence covers TP8 only. This follow-up removes the +configuration blocker on current `onecat/main` at +`7aede2cf010d92815c9d7bff25867b4fa009b6cb`; the full runtime evidence gap +remains explicit below. + +- The roughly 160 GB checkpoint cannot fit pure TP4 or pure PP4 on four + V100-32GB GPUs. The memory-feasible eight-GPU candidates are PP2 x TP4 and, + if four pipeline stages are required, PP4 x TP2. +- The default 43-layer partitions are `22+21` for PP2 and `11+11+11+10` for + PP4. DSpark target layers 40-42 remain entirely on the final pipeline stage + in both layouts, matching the runner's final-stage-only drafter ownership. +- PP intermediate tensors must retain the four mHC streams with shape + `[tokens, 4, 4096]`; collapsing them to a two-dimensional hidden state at a + stage boundary would change model semantics. +- The first runtime gate is PP2 x TP4 with the same model, FP8 MLA KV, DSpark + N7, 1024/256 request, sampling, graph mode, and fast-stack flags as the + accepted TP8 reproduction. The TP8-only hierarchical all-reduce is disabled + for each TP4 group. + +The first current-main PP2 x TP4 attempt failed before worker launch. DSpark's +draft config inherited target `pipeline_parallel_size=2`, so generic model +validation required the final-stage-local `DSparkDraftModel` to implement +`SupportsPP`. The runner intentionally constructs the complete drafter only on +the final target PP rank; partitioning that three-layer draft again would not +match its ownership model. The candidate therefore gives DSpark a draft-local +PP size of one while retaining the target TP/PP topology and draft TP size. It +also rejects a custom layer partition if any configured target layer falls +outside the final target pipeline stage. Exact configuration regressions pass +for PP2 x TP4 and PP4 x TP2. Post-fix runtime route, quality, speed, and cleanup +evidence remain pending. + ## Upstream Reuse Reuse these official implementations and semantics: diff --git a/docs/design/sm70_dflash2_acceptance_rd.md b/docs/design/sm70_dflash2_acceptance_rd.md new file mode 100644 index 0000000000..c190844ec0 --- /dev/null +++ b/docs/design/sm70_dflash2_acceptance_rd.md @@ -0,0 +1,77 @@ +# SM70 DFlash2 acceptance research without retraining + +> **Status:** the default-off alignment capture and offline analyzer are +> development tooling. They do not change the production proposal policy and +> do not claim an acceptance or performance improvement by themselves. + +## Scope and frozen contracts + +- Date: 2026-08-25. +- Integration base: `onecat/main@34403018d917054dd7765d5e820ad29c8d342348`. +- Branch: `codex/v100-v100-dflash2-acceptance-rd-20260825-113155`. +- No draft or target retraining and no checkpoint edits. +- Target validation remains the complete MRV2 block-eight workload: seven draft + tokens plus one target bonus row. DDTree index reuse and wider verifier trees + are outside this campaign. +- Target sampling remains exact standard rejection at temperature 1.0, top-p + 0.95, and top-k 20. A candidate proposal distribution is admissible only if + the selector draw and the rejection sampler consume the same exact `q`. + +The score gate remains the retained 128-row mixed32 dataset with SHA256 +`85ce10d84735ec981c001663d5748b1939a81892ef3ccd4549aec66177795607`. +Its private filesystem location is intentionally not committed. +It uses sequential B1, a fixed seed, xhigh reasoning, at most 2,048 generated +tokens, and natural EOS. Corrected MBPP-32 and WikiText-2 PPL retain their +existing independent gates. The historical native-grouped result is 83/96 on +GSM8K/MATH-500/HumanEval, 25/32 on corrected MBPP, request-mean completion +length 4.238430, and pooled acceptance length 3.749104. + +The production performance gate remains the PR #288 practical contract: TP4 +V100, 256K maximum context, `max_num_batched_tokens=4096`, prefix caching, +Mamba align mode, FP8 E5M2 target KV, FP16 draft KV, Flash-V100, probabilistic +DFlash2, tool/reasoning parsers, and FULL CUDA Graph. Its short-context complete +round is 18.465--18.603 ms. Acceptance work may not increase verifier rows or +regress that round by more than measurement noise. + +## First target and promotion gate + +The first research target is at least +0.5 request-mean completion tokens per +verification round on the unchanged mixed32 contract. Promotion additionally +requires: + +1. no aggregate task-score loss against the paired target-only and current + DFlash2 controls; +2. no corrected-MBPP or WikiText PPL regression; +3. exact rejection-distribution tests and normalized proposal rows; +4. unchanged eight-row target validation and at most 0.05 ms incremental + selector latency on SM70; +5. no material acceptance or latency regression at 1K, 32K, 128K, or 256K. + +## Selector headroom audit + +The first implementation is diagnostic-only. With both compact sparse +rejection and `VLLM_SPEC_DUMP_ALIGNMENT=1`, the speculator keeps the checkpoint +top16 IDs, unary logits, full `7 x 16 x 16` selector lattice, and exact realized +proposal rows. The target sampler records its top20 rows and the actual strict +rejection outcome. Disabled diagnostics allocate no shadow lattice and add no +copy or sampler operation. + +`benchmarks/analyze_dflash2_selector_alignment.py` reports, by draft depth: + +- current overlap `sum(min(p, q))`; +- target probability mass covered by the fixed top16 candidate support; +- the gap between support mass and current overlap; +- one-step counterfactual sweeps over proposal temperature, proposal nucleus, + unary/edge calibration, and backward log-sum-exp future messages. + +Even sampler steps tune the candidates and odd steps are a held-out report. +Counterfactual values after a changed path are explicitly only recorded-prefix +overlap proxies, not end-to-end acceptance claims. Survivors must be rerun on +the frozen dataset. + +## Initial verification + +- Focused Ruff lint and format: passed. +- Diagnostic analyzer synthetic checks: passed. +- CPU DFlash2 suite with CUDA hidden: 77 passed, 12 expected CUDA skips. +- No GPU performance or end-to-end acceptance result has been claimed yet. diff --git a/docs/design/sm70_dflash2_long_verify_decay.md b/docs/design/sm70_dflash2_long_verify_decay.md new file mode 100644 index 0000000000..7a8f161214 --- /dev/null +++ b/docs/design/sm70_dflash2_long_verify_decay.md @@ -0,0 +1,465 @@ +# SM70 DFlash2 long-context verification decay + +## Scope + +This task measures, attributes, and reduces the context-length growth of one +MRV2 DFlash2 verification round on V100. The historical study started on +private DFlash2 verifier head `e8daa778f55cf16df311f8ce8dd032e2281f1680`; +the public change set contains only the five retained verifier commits rebased +onto current `main`. It does not replace or retune DFlash1, `dflash_ddtree`, +Eagle, or MTP. + +DFlash2 names the measured workload, not a runtime activation identity. The +kernel entry admits only the explicit SM70, E5M2 paged-KV, Q1-Q8, Hq6/Hkv1, +D256, page-layout, causal-mask, and workspace contracts. It does not inspect a +model name, checkpoint, `model_type`, or architecture identity. + +Quality promotion remains score based. Greedy token identity is diagnostic; +the production gates are probabilistic task scores, PPL, acceptance length, +and output health against the same target-only checkpoint. + +## Frozen workload + +- Target: local Qwen3.8-27B mixed NVFP4/FP8 checkpoint, TP4. +- Draft: `incoai/Qwen3.8-27B-DFlash2`, FP16 draft KV. +- Target KV: `fp8_e5m2`. +- Attention: `FLASH_ATTN_V100` target and draft routes. +- Execution: FULL CUDA Graph, sequential batch one. +- Sampling: temperature 1.0, top-p 0.95, target top-k 20, probabilistic + proposals, fixed seed. +- Selector top-K: checkpoint contract 16; speculative tokens: 7. +- Context sweep: 1K, 32K, 128K, and 256K boundary. + +All comparisons must keep model, TP/GPU set, graph mode, KV dtypes, sampling, +and route gates identical. Prefill/TTFT are reported separately from steady +decode and verification-round wall time. + +## Existing evidence to reproduce + +The parent task recorded the following unprofiled paired measurements: + +| Context | Prior complete round | Grouped complete round | Prior target phase | Grouped target phase | +| --- | ---: | ---: | ---: | ---: | +| 32K | 36.311 ms | 26.361 ms | 30.248 ms | 20.452 ms | +| 128K | 69.855 ms | 40.778 ms | 63.708 ms | 34.647 ms | + +From 32K to 128K, the grouped complete round grows by 14.417 ms and its target +phase grows by 14.195 ms. This strongly localizes the observed growth to target +verification, but it does not yet distinguish paged attention service, +dependency gaps, launch geometry, communication placement, or trace overhead. + +The retained grouped-attention microbenchmark reported 0.0686/0.3564/1.2145/ +2.3962 ms at 1K/32K/128K/256K. These values must be re-run and reconciled with +the full-model critical-path wall before using them as a closed decomposition. + +## Closed trace attribution + +The retained 32K and 128K Nsight Systems traces use the same grouped one-pass +kernel and target CUDA Graph route. Summing the partial and combine kernels, +then dividing by the number of target graph replays across all four TP ranks, +gives the following service time per target verification graph: + +| Context | Partial average | Partial calls | Combine average | Attention per target graph | +| --- | ---: | ---: | ---: | ---: | +| 32K | 0.313606 ms | 1,344 | 0.008282 ms | 5.150214 ms | +| 128K | 1.198762 ms | 1,088 | 0.008295 ms | 19.312912 ms | + +There are sixteen full-attention layers per target graph. The attention service +therefore grows by 14.162698 ms between 32K and 128K, while the measured target +phase grows by 14.270586 ms. Grouped attention explains 99.24% of target-phase +growth and 95.83% of complete-round growth in this pair. Draft, selector, +sampling, TP exchange, and host residuals are not material causes of this +context-length slope. + +Both traces launch the partial kernel with grid `(2, 40, 1)`, block size 256, +128 registers per thread, and 39,424 bytes of static shared memory. The fixed +80-CTA grid maps one CTA to each V100 SM even as the KV length quadruples. The +kernel's `__launch_bounds__(256, 2)` and resource arithmetic permit two resident +CTAs per SM (65,536 registers and 78,848 shared-memory bytes), so the first +bounded experiment is 40 versus 80 context splits. This is a hypothesis about +exposed parallelism, not yet a performance claim: an exact-shape paired +microbenchmark must distinguish occupancy benefit from memory-bandwidth +saturation and extra split-reduction cost. + +The experiment follows the same split-KV principle used by the official +[FlashInfer scheduler](https://github.com/flashinfer-ai/flashinfer/blob/main/include/flashinfer/attention/scheduler.cuh), +which derives its maximum grid from CUDA occupancy, and by +[Flash-Decoding](https://princeton-nlp.github.io/flash-decoding/). Recent +[PersistentKV](https://arxiv.org/abs/2606.26666) results likewise motivate +adaptive sequence splitting and page-aware scheduling for long-context decode. +These sources guide the search space; only same-machine measurements can +promote a V100-specific default. + +## Independent upstream audit + +`sglang-V100` is derived from this project's V100 work. It remains useful for +route and result cross-checks, but it is not independent evidence for an +optimization decision. The bounded search therefore uses primary upstream +implementations and papers, pinned at the audited revisions: + +- Official FlashAttention revision + [`0251105a`](https://github.com/Dao-AILab/flash-attention/tree/0251105a2fb19d2957484b7f023cd8c115286ced) + packs `seqlen_q * qheads_per_kv_head` into the M dimension in + [`pack_gqa.h`](https://github.com/Dao-AILab/flash-attention/blob/0251105a2fb19d2957484b7f023cd8c115286ced/hopper/pack_gqa.h), + then computes dynamic Split-KV occupancy from the packed M work in + [`flash_prepare_scheduler.cu`](https://github.com/Dao-AILab/flash-attention/blob/0251105a2fb19d2957484b7f023cd8c115286ced/hopper/flash_prepare_scheduler.cu). + This is the direct algorithmic source for combining the six Q heads that + share the Qwen3.8 verifier KV head. Its SM80/SM90 implementation is not + copied; only the exact PackGQA scheduling principle is ported to Volta WMMA. +- FlashInfer revision + [`c18a974a`](https://github.com/flashinfer-ai/flashinfer/commit/c18a974a88b8c7da88be0338a7478636f1dffa17) + explicitly keeps head-dimension-256 prefill within the 64-KiB SM70/SM75 + shared-memory limit. Its release JIT successfully compiled on this V100, + but the grouped Q8/H6/D256 paged launch still ended in an unspecified launch + failure during CUDA Graph setup. The direct backend route is therefore + rejected; its packing and scheduler layout remain useful source material. +- TensorRT-LLM revision + [`958d651b`](https://github.com/NVIDIA/TensorRT-LLM/tree/958d651b18800da1da0677160e0277f701c0ff84) + supports multi-block and multi-query-token XQA, but + `supportConfigHMMA` rejects `SM < 80`. ONNX Runtime XQA likewise explicitly + requires compute capability 8.0 or newer. Neither is a drop-in V100 kernel; + both independently confirm head-group packing plus context splitting as the + intended speculative-verification topology. TensorRT-LLM's speculative + multi-block scheduler further caps the packed grid at one SM wave and avoids + scaling short histories, independently matching the bounded 80-CTA policy + tested below. +- [PersistentKV](https://arxiv.org/abs/2606.26666) maps native paged GQA work + by `(request, KV head, sequence split)`, explicitly reusing each K/V tile + across the grouped query heads. Its compact workqueue is valuable for ragged + concurrent batches; the frozen single-request Q8 verifier has no empty work + items to compact. The head-group mapping is directly applicable, while the + workqueue is retained as a later concurrency experiment rather than being + conflated with the B1 result. +- xFormers Split-K combines split states with LSE and supports small causal + query counts plus experimental paged attention. Its current Triton route + requires SM80, so it is algorithmic cross-check material, not a new V100 + dependency. +- [FlashAttention-3](https://papers.neurips.cc/paper_files/paper/2024/file/7ede97c3e082c6df10a8d6103a2eebd2-Paper-Conference.pdf), + [LeanAttention](https://arxiv.org/abs/2405.10480), and + [FlashDecoding++](https://arxiv.org/abs/2311.01282) provide the exact + split-softmax reduction, Stream-K work distribution, and GEMM-oriented + decode alternatives. They motivate bounded follow-ups, not unmeasured + performance claims on Volta. + +TinyFA revision `bc755512` was also inspected and excluded: its fast path +requires SM75-era asynchronous-copy machinery, lacks E5M2 paged KV support, +and its paged decode mapping does not pack all six query heads into one KV +load. This audit leaves native SM70 PackGQA as the smallest dependency-closed +candidate. + +## Split-count microbenchmark + +The first candidate doubles the maximum context splits from 40 to 80. With +two GQA head groups this changes the fixed partial launch from 80 to 160 CTAs, +matching the kernel's two-block-per-SM launch bound. No attention arithmetic, +mask, FP8 decode, or MRV2 scheduling code changes. + +An A/B/A CUDA-Graph run on physical GPU 1 used the production Qwen3.8 target +shape (Q8/H6/KV1/D256, E5M2 KV, page size 3,296), 40 warmups and 200 measured +rounds per point: + +| Context | Split 40 A | Split 80 | Split 40 A2 | Split 80 versus A/A2 | +| --- | ---: | ---: | ---: | ---: | +| 1K | 0.058368 ms | 0.067584 ms | 0.058368 ms | +15.79% / +15.79% | +| 32K | 0.319488 ms | 0.261120 ms | 0.322560 ms | -18.27% / -19.05% | +| 128K | 1.204736 ms | 0.955392 ms | 1.206272 ms | -20.70% / -20.80% | +| 256K | 2.377728 ms | 1.864704 ms | 2.378752 ms | -21.58% / -21.61% | + +At sixteen full-attention layers, the operator delta projects to -0.934 ms at +32K, -3.989 ms at 128K, and -8.208 ms at 256K per verification graph. The 1K +operator regression projects to +0.147 ms, or about +0.74% of the retained +approximately 20 ms complete short-context round. This remains a projection +until the TP4 full-model A/B is complete. + +The split-80 extension passes all six targeted grouped-verifier tests: five +random physical-page cases against the FP32 reference, and CUDA Graph replay +while mutating runtime sequence length. These are operator correctness gates; +acceptance length, PPL, and task scores remain required before promotion. + +## Full-model 32K and 128K confirmation + +A TP4 FULL-Graph Nsight run used the same 128K prompt, 128 generated tokens, +probabilistic sampling seed, model/KV dtypes, and route gates as retained +split-40 run `v80`. The only implementation change was the split-80 Flash-V100 +extension: + +| Phase | Split 40 | Split 80 | Delta | +| --- | ---: | ---: | ---: | +| Complete verification round | 40.778072 ms | 36.683712 ms | -4.094360 ms (-10.04%) | +| Target | 34.646866 ms | 30.557332 ms | -4.089534 ms (-11.80%) | +| Draft | 4.604373 ms | 4.627743 ms | +0.023370 ms | +| Draft to target | 0.230538 ms | 0.225575 ms | -0.004963 ms | +| Target to draft | 1.296296 ms | 1.273062 ms | -0.023233 ms | + +The partial kernel changed from grid `(2, 40, 1)` to `(2, 80, 1)`. Its mean +service time fell from 1.198762 ms to 0.944041 ms per full-attention layer +(-21.25%); combine rose from 0.008295 ms to 0.010430 ms. Across sixteen layers, +the net attention prediction is -4.041 ms, explaining 98.8% of the measured +target-phase improvement. The earlier operator projection was -3.989 ms, only +about 0.10 ms from the full target result. + +The fixed diagnostic request retained identical accepted-token counts +`[17, 17, 17, 16, 16, 15, 15]`, mean acceptance length 7.647059, and all 128 +output token IDs. This is trajectory evidence, not a substitute for the +required score/PPL gates. + +The matching 32K TP4 FULL-Graph trace independently confirms the shorter +context point: + +| Phase | Split 40 | Split 80 | Delta | +| --- | ---: | ---: | ---: | +| Complete verification round | 26.361447 ms | 25.238442 ms | -1.123005 ms (-4.26%) | +| Target | 20.451644 ms | 19.468326 ms | -0.983318 ms (-4.81%) | +| Draft | 4.365014 ms | 4.257568 ms | -0.107446 ms | +| Draft to target | 0.260419 ms | 0.235765 ms | -0.024654 ms | +| Target to draft | 1.284369 ms | 1.276783 ms | -0.007586 ms | + +The split-80 partial kernel averages 0.248999 ms at 32K, versus 0.313606 ms +for split 40. The measured target reduction is within 0.05 ms of the operator +projection. Across the paired full-model endpoints, the 32K-to-128K complete +round growth falls from 14.416625 ms to 11.445270 ms (-20.61%), while target +growth falls from 14.195222 ms to 11.089005 ms (-21.88%). The remaining slope +is still serious: target accounts for 96.89% of the optimized complete-round +growth, so subsequent work remains scoped to grouped paged attention. + +## Physical-page and tile-balanced follow-up + +Two address-equivalent scheduling changes pass the same six operator tests: + +- Specialize the actual FP8 kernel-page sizes (1,648 and 3,296) at compile + time, replacing runtime integer division with constant arithmetic. The + Qwen3.8 hybrid cache reports a 1,648-token logical attention block, while the + grouped verifier receives the 3,296-element physical FP8 view; the final + trace proves the `` specialization is selected. +- Divide the sequence by complete 32-token WMMA tiles, distributing the + remainder across splits. This removes per-split padded tail tiles without + changing visibility, loaded KV bytes, softmax, or WMMA arithmetic. + +Repeated CUDA-Graph A/B/A runs show no 32K regression and a further +approximately 1.7% operator reduction at 128K versus plain split 80. The final +TP4 FULL-Graph endpoints are: + +| Phase | Split 80 at 32K | Final at 32K | Split 80 at 128K | Final at 128K | +| --- | ---: | ---: | ---: | ---: | +| Complete verification round | 25.238442 ms | 25.237987 ms | 36.683712 ms | 36.482220 ms | +| Target | 19.468326 ms | 19.455442 ms | 30.557332 ms | 30.354906 ms | +| Draft | 4.257568 ms | 4.269400 ms | 4.627743 ms | 4.606652 ms | + +At 128K, the partial kernel falls from 0.944041 ms to 0.926568 ms per +full-attention layer; sixteen layers predict -0.280 ms and the measured target +delta is -0.202 ms. At 32K, the partial kernel is 0.245912 ms and complete wall +is unchanged within 0.001 ms. The final 32K-to-128K complete-round slope is +11.244233 ms, 22.01% below split 40 and 1.76% below plain split 80. The final +target slope is 10.899464 ms, 23.22% below split 40. Target attention still +accounts for 96.93% of the residual complete-round slope. + +A logical-work roofline proxy narrows the remaining operator bottleneck. The +two padded 32-row GQA groups perform about 2.15/8.59 GFLOP per layer at +32K/128K; the measured partial times correspond to only 8.73/9.27 logical +TFLOP/s. Their requested E5M2 K/V payload corresponds to about 136/145 GB/s +before accounting for duplicate-group L2 reuse. Both rates stay nearly flat as +context grows and are far below V100 peak figures. This is not an NCU hardware +counter claim, but it rules out treating the residual as unavoidable saturated +HBM traffic and prioritizes flat-QK/PV utilization, softmax synchronization, +and GQA packing over further page-address arithmetic. + +The final 128K diagnostic request preserves the split-80 accepted-token counts +and 7.647059 mean acceptance length. At 32K, the single-request diagnostic +changes from 6.142857 to 6.095238 (-0.047619), inside the 0.05 diagnostic gate; +dataset-level acceptance and model-quality gates remain authoritative. + +An attempted compact mapping for the two-head GQA tail passed the six operator +correctness tests but was 6.3-7.0% slower than split 80 across the micro sweep. +The run also observed a late external-worker ownership race, so its absolute +timings are not promotion evidence; the uniformly negative direction was +sufficient to revert the experiment completely. Nsight Compute hardware +counters are unavailable to the current user (`ERR_NVGPUCTRPERM`), so no NCU +counter claim is made. + +A second exact-address experiment replaced sixteen redundant page divisions, +modulos, and block-table loads per D=256 row with one half-warp calculation and +broadcast. It preserved register count and passed all six operator correctness +tests, but regressed the one-pass kernel by 3.51%/2.36%/2.47%/2.63% at +1K/32K/128K/256K in a production-page-size A/B/A run. On this Volta path, the +independent address arithmetic is hidden better than the added shuffle +dependency. The candidate was reverted completely. + +## Short-context floor and rejected split-K QK + +The missing 1K TP4 FULL-Graph endpoint now uses the same final binary and +sampling/configuration contract as the 32K and 128K traces: + +| Phase | 1K mean per verification round | +| --- | ---: | +| Complete | 21.634854 ms | +| Target | 16.348011 ms | +| Draft | 3.760873 ms | +| Draft to target | 0.242986 ms | +| Target to draft | 1.282984 ms | + +The grouped partial and combine kernels average 0.057043 and 0.005948 ms per +full-attention layer. Across sixteen layers they contribute about 1.007854 ms +to the target graph. Subtracting this measured service leaves approximately +15.340 ms of context-independent target work and 20.627 ms of complete-round +work. Thus long-attention work explains the decay, but eliminating it entirely +would not by itself satisfy a sub-20-ms short-context goal; the fixed target, +draft, and return path remain a separate optimization problem. + +An exact split-K experiment divided the D=256 QK accumulation between two +warps and joined the FP32 partials before the unchanged softmax. Limiting +unrolling removed all local stack spills (128/126 registers for the one-/two- +pass kernels), and all six grouped-verifier correctness tests passed. A/B/A +medians nevertheless changed by -0.40% at 32K, +0.11% at 128K, and +0.28% at +256K versus the second baseline. There is no stable long-context gain, so the +experiment was reverted and is not part of the performance result. + +## Pack-GQA48 follow-up + +The accepted grouped baseline launches two 256-thread CTAs per context split: +one four-head group and one two-head tail. Both CTAs load and decode the same +single-KV-head page. Following the official FlashAttention PackGQA mapping, +the follow-up packs all six heads and eight verifier tokens into one 48-row, +512-thread CTA. QK and PV remain Volta WMMA, online softmax remains FP32, and +the split outputs retain the same FP16-partial/FP32-LSE reduction contract. +The production partial kernel uses 128 registers, no local stack, and 50,944 +bytes of opt-in dynamic shared memory. + +The first packed prototype retained 128 tokens per split. It improved the +32K/128K/256K operator by approximately 20-21%, but regressed 1K by 19-21% +because only eight large CTAs were launched. Reducing the packed split quantum +to 64 tokens recovered the 1K parallelism, but the single-query utility shape +then crossed its strict FP32-reference mean-error gate. The retained dispatch +therefore keeps 128-token splits for Q1, forces one split for Q2-Q8 at total KV +length at most 128, and uses 64-token splits otherwise. This is a numerical +dispatch boundary, not a sampling or acceptance heuristic. + +The final production-page A/B/A CUDA-Graph medians are: + +| Context | Accepted grouped A/A2 | Pack-GQA48 | Delta versus A/A2 | +| --- | ---: | ---: | ---: | +| 1K | 0.057344 / 0.057344 ms | 0.046080 ms | -19.64% / -19.64% | +| 32K | 0.257024 / 0.257024 ms | 0.205824 ms | -19.92% / -19.92% | +| 128K | 0.939008 / 0.937984 ms | 0.745472 ms | -20.61% / -20.52% | +| 256K | 1.838080 / 1.838080 ms | 1.448960 ms | -21.17% / -21.17% | + +At the ultra-short 128/256-token fixed shapes the packed route costs about +0.002-0.003 ms more than the accepted grouped kernel; at 512 tokens it is +already 10.9-12.5% faster. These sub-millisecond boundary points do not affect +the frozen 1K-or-longer promotion gate, but remain covered by correctness +tests. + +The expanded operator suite covers Q1/Q3/Q4/Q5/Q6/Q8, 128/256/512/1K and +longer boundaries, random physical pages, one-/two-pass equivalence, and CUDA +Graph replay. All thirteen collected cases pass. The maximum absolute error is +unchanged from the accepted grouped baseline at every audited shape; the largest +positive mean-error delta is `1.60e-7`, below the per-shape `5e-7` guard. + +The TP4 FULL-Graph endpoints confirm that the operator gain lands on the full +target critical path without a short-context regression: + +| Context/phase | Accepted grouped | Pack-GQA48 | Delta | +| --- | ---: | ---: | ---: | +| 1K complete | 21.634854 ms | 21.303621 ms | -0.331233 ms (-1.53%) | +| 1K target | 16.348011 ms | 16.043125 ms | -0.304886 ms (-1.86%) | +| 32K complete | 25.237987 ms | 24.390177 ms | -0.847810 ms (-3.36%) | +| 32K target | 19.455442 ms | 18.633571 ms | -0.821871 ms (-4.22%) | +| 128K complete | 36.482220 ms | 33.334558 ms | -3.147662 ms (-8.63%) | +| 128K target | 30.354906 ms | 27.208276 ms | -3.146630 ms (-10.37%) | + +The traced packed partial averages 0.733521 ms per full-attention layer and its +combine averages 0.010906 ms. The diagnostic request keeps the exact 128 output +token IDs, accepted counts `[17, 17, 17, 16, 16, 15, 15]`, and mean acceptance +length 7.647059. Seventeen verification steps produce 128 output tokens, so +the conservative steady-round projection is about 225.9 token/s at 128K, +versus 206.4 token/s for the accepted grouped endpoint. + +The 32K-to-128K complete-round slope falls from 11.244233 to 8.944381 ms +(-20.45%), and the target slope falls from 10.899464 to 8.574705 ms (-21.33%). +The 32K/128K geometric-mean complete round falls from 30.343662 to 28.513782 +ms (-6.03%). Diagnostic mean acceptance changes by +0.030303 at 1K, +0.047619 +at 32K, and exactly zero at 128K, all inside the frozen 0.05 boundary. + +WikiText PPL is exactly equal to the accepted grouped candidate across all +16,376 prompt-logprob positions. Against target only, the maximum PPL delta is +0.003952 and the mean absolute prompt-logprob delta is 0.007519, exactly the +same quality envelope recorded before Pack-GQA. Mixed task-score gates remain +required before promotion. + +## Score and PPL promotion evidence + +The final split-80/page-specialized/tile-balanced binary was evaluated with +probabilistic proposals, target sampling at temperature 1.0/top-p 0.95/top-k +20, fixed seed, a 2,048-token cap, and natural EOS. The mixed dataset supplies +32 GSM8K, 32 HumanEval, and 32 MATH500 cases. Its embedded MBPP prompts have an +incompatible function-name contract and are excluded; the independent +`mbpp-tests-32` dataset supplies the valid MBPP result. +The tested extension SHA256 is +`4a925265cfed934ebd7772be2a1eccaa64163d5324f87d31ba8ee99d10c3a044`. + +| Suite | Target only | Prior grouped | Final candidate | +| --- | ---: | ---: | ---: | +| GSM8K | 29/32 | 30/32 | 30/32 | +| HumanEval | 28/32 | 28/32 | 27/32 | +| MATH500 | 26/32 | 25/32 | 26/32 | +| MBPP, independent tests | 25/32 | 25/32 | 26/32 | +| Valid total | 108/128 | 108/128 | 109/128 | + +The suite-level movement is one HumanEval case down and a net two MATH/MBPP +cases up; the valid aggregate is therefore one case above both controls, not a +quality-regression claim. WikiText prompt-logprob/PPL output is exactly equal +to the prior grouped candidate across all eight samples (zero maximum and mean +difference). Against target only, the maximum PPL difference is 0.003952 and +the mean absolute prompt-logprob difference is 0.007519. + +Acceptance also remains inside the frozen 0.05 gate. On the 128-request mixed +run, request-mean completion tokens per verification step change from 4.238430 +to 4.218840 (-0.019589), and pooled acceptance length changes from 3.749104 to +3.714972 (-0.034132). On independent MBPP they change from 3.527945 to +3.509266 (-0.018679), while pooled acceptance changes from 3.421634 to +3.431016 (+0.009382). The final runs produce 101,312 and 28,070 output tokens, +respectively; 103/128 and 24/32 requests terminate naturally before the cap. + +## Measurement order + +1. Reproduce low-overhead full-model B1 baselines at 1K/32K/128K/256K. +2. Record synchronized target, draft, selector/sampling, communication, and + host residual timing from the same requests. +3. Capture short Nsight Systems graph-node traces at 32K and 128K, aligned by + replay ordinal across all four TP workers. +4. Classify target graph nodes into attention, quantized GEMM, GDN/norm, + LM-head/selector, collectives, and dependency gaps. +5. Build the smallest exact-shape microbenchmark for the confirmed growing + kernel family. Run Nsight Compute only after the runtime route is proved. + +## Promotion gates + +- Reduce geometric-mean 32K/128K complete verification-round wall without a + regression greater than 1% at either length. +- Do not increase 1K complete-round wall by more than 1%. +- Preserve request-mean acceptance length within 0.05 and pooled acceptance + within sampling noise. +- Preserve the existing task-score and WikiText PPL gates before enabling a + numerical or scheduling change by default. +- Validate the 256K boundary for capacity, output health, and latency after a + candidate passes at 128K. + +## Current status + +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. + +The final source audit rebased all five patches without semantic change onto +`main@acc0f6fb92`, after the generic prefix-anchored and E4M3 batch-XQA merges. +The complete changed-file static suite passes and the combined SM70 extension +builds and links. Removing a redundant translation-unit-wide dynamic-shared +alignment attribute preserves the packed verifier layout while restoring all +121 pre-existing XQA kernel SASS instruction streams exactly; none are changed +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. diff --git a/docs/design/sm70_dflash2_ngram_hybrid.md b/docs/design/sm70_dflash2_ngram_hybrid.md new file mode 100644 index 0000000000..bb45e80d93 --- /dev/null +++ b/docs/design/sm70_dflash2_ngram_hybrid.md @@ -0,0 +1,130 @@ +# SM70 DFlash2 + ngram Hybrid + +## Purpose + +Add an optional prompt-ngram assistant to the MRV2 DFlash2 path. An ngram hit +must skip the DFlash2 query/selector while preserving DFlash2 context-KV state, +and a miss must fall back to the existing DFlash2 implementation unchanged. +The target verifier and the probabilistic rejection-sampling contract remain +authoritative, so the optimization cannot change the target distribution. + +This work is based on `onecat/main` at +`d62ef5cb20b48de93a91562e777ac48985f44b76` and is isolated on +`agent/v100-dflash2-ngram-hybrid-20260825-050018`. + +## Upstream audit + +- vLLM prompt lookup: PRs #12193, #22437, #24986, and #29184 provide the CPU + KMP and GPU-vectorized implementations. MRV2 still has no combined ngram + + model-drafter route. +- SGLang: PRs #17260, #21243, and #22737 show that overlap scheduling requires + complete request-token state and explicit accepted-token indexing. A stale + host output list is not a valid lookup source. +- llama.cpp: its comma-separated speculative configuration gives draftless + ngram proposers priority, falls back to DFlash, and still calls `process()` on + every proposer so model-drafter state remains synchronized. This is the + closest reference architecture for the first implementation here. +- Arctic Suffix Decoding and SAM-Decoding are useful follow-ups, but their + confidence policy, external dependency, and longer trees are intentionally + outside this first block-8 implementation. + +## Initial contract + +- The feature is opt-in and only valid for `method=dflash` with a DFlash2 + checkpoint. Standalone `ngram`, Eagle, MTP, DFlash1, and `dflash_ddtree` + routing is unchanged. +- Lookup reads the authoritative MRV2 request-token state and supports normal + synchronous and overlap scheduling without rebuilding history in Python. +- Structured-output requests bypass the ngram assistant until grammar-aware + proposal masking is proved correct. Tool/reasoning parsers without a grammar + are unaffected. +- A hit returns at most the configured DFlash2 draft width (seven tokens in the + official block-8 setup). The normal DFlash2 proposer handles misses. +- Context K/V materialization always runs. Only the DFlash2 query, candidate + projection, selector, and selector walk may be skipped. +- Greedy and probabilistic modes are supported. In probabilistic mode the + ngram proposal is represented as a one-hot draft distribution inside the + existing sparse target-rejection interface, preserving target sampling. +- Prefix-cache hits may rebuild missing DFlash draft K/V as before; this change + must not corrupt either target or draft cache state. + +## Test plan + +1. Unit tests: configuration/routing isolation; KMP hit, miss, truncation, and + overlap cases; per-request mixed-source state; one-hot sparse draft support; + grammar bypass; prefix/state transitions. +2. CPU microbenchmarks: lookup at 1K, 32K, 128K, and 256K contexts. +3. V100 correctness: eager before CUDA graph; batch 1/2/4; hit/miss/mixed; + compare proposed tokens, accepted trajectory, target output, and draft K/V. +4. V100 performance: report lookup, context-KV, query/selector, target verify, + full round, emitted tokens per round, and pure-decode tokens/s for baseline + DFlash2 and DFlash2+ngram at short and long contexts. +5. Quality: compare target-only, DFlash2, and hybrid on coding/general/tool + datasets under the same sampling configuration. Report scores, completion + counts, wall time, hit rate, conditional acceptance length, and failures. + +## Promotion gates + +- No statistically meaningful quality regression against target-only or the + existing DFlash2 route. +- DFlash2-miss acceptance trajectory remains unchanged in deterministic tests. +- Ngram hits actually skip query/selector work in a trace. +- Hybrid pure-decode throughput is non-regressing at every measured context; + otherwise the feature remains opt-in while the losing shape is investigated. + +## Status + +- 2026-08-25: Draft PR #287 implements opt-in MRV2 host lookup over the + authoritative UVA request-token state, full-hit query/selector bypass, + all-hit batch application, and one-hot dense/sparse rejection caches. + Structured-output batches bypass the assistant, and intermediate chunked + prefill materializes draft context K/V before returning without lookup. +- CPU split-history KMP microbenchmark (median/P95): 1K `2.66/2.74 us`, 32K + `50.05/53.22 us`, 128K `190.72/201.03 us`, and 256K `369.57/393.17 us`. + Focused validation reports `93 passed, 9 skipped` for the CPU DFlash2/ngram + set, `10 passed` for MRV2 routing, and `24 passed` for the V100 ngram/AOT + fullgraph set. +- PR #288 first restored the missing production optimization closure to main. + On that restored source, no-assist practical coding measured `18.660 ms` per + complete round, `135.70 tok/s`, and acceptance length `2.532`. Repeated runs + of the historical MBPP item 28 stabilized at `18.484--18.616 ms`, + `224.05--225.66 tok/s`, and acceptance length `4.184`. Older 27--31 ms + branch measurements predate this closure and are invalid baselines. +- With ngram `[5,5]` enabled under the identical TP4 practical contract, the + coding request measured `18.980 ms`, `131.14 tok/s`, and acceptance length + `2.494`; only about 2.2% of eligible rounds were full hits. The matched MBPP + item measured `18.700 ms`, `217.18 tok/s`, and acceptance length `4.071`. + Thus low-hit short requests currently pay roughly `0.1--0.3 ms` per round + and do not pass the default-enable performance gate. +- The 32-case MBPP natural-stop run (16K output cap, fixed per-item seeds) gave + the pure-DFlash control `65,934` output tokens in `389.542 s`, `19,767` + verification rounds, `19.707 ms` wall time per round, and acceptance length + `3.336`. Hybrid produced `53,710` tokens in `327.613 s`, `16,700` rounds, + `19.618 ms` per round, and acceptance length `3.216`. Cumulative full-hit + rate reached about 10.1% and lookup averaged `0.016--0.018 ms`; skipped + queries slightly reduced round cost, but the lower sampled acceptance path + left raw aggregate output throughput at `163.94` versus `169.26 tok/s`. +- EvalPlus on the mapped 31 MBPP cases reports pure DFlash Base `30/31` and + Plus `28/31`; hybrid reports Base `31/31` and Plus `28/31`. This is no score + regression, but it does not override the failed throughput gate. The + probabilistic one-hot proposal preserves the target distribution; it is not + expected to reproduce the same sample-by-sample random trajectory. +- A fresh hybrid 32K chunked-prefill request measured `10.688 s` cold and + `1.416 s` on an identical-prefix hit, versus the restored no-assist evidence + of `10.58--10.65 s` and `1.405 s`. The context-only skip log was present and + no cache corruption occurred. +- Promotion decision: keep `ngram_assist` opt-in. The merge audit removes the + mixed-hit hazard: an n-gram draft is now applied only when every active + request has a full-width hit and the complete DFlash2 query/selector can be + skipped. Mixed-hit batches retain the unchanged DFlash2 proposals because + overriding a row after paying the full query cost cannot improve latency and + can alter sampled acceptance. A future default-on policy still needs matched + evidence that full-query skips repay lookup/probe overhead; blindly reducing + the ngram length is not justified by the current data. +- After the paired run was stopped, fresh graph captures on GPUs 4--7 began + failing inside the draft paged-attention capture with + `cudaErrorStreamCaptureInvalidated`. A clean pre-ngram main worktree failed + at the same point, so this is not an ngram source regression. The NVLink + topology requires resetting all eight GPUs, which was deliberately not done + while the user-facing API remained live on GPUs 0--3. Do not count failed + startups as performance samples. diff --git a/docs/design/sm70_dflash2_nvfp4_17ms.md b/docs/design/sm70_dflash2_nvfp4_17ms.md new file mode 100644 index 0000000000..f29e7abc9d --- /dev/null +++ b/docs/design/sm70_dflash2_nvfp4_17ms.md @@ -0,0 +1,220 @@ +# SM70 Qwen3.8 NVFP4 DFlash2 17 ms Campaign + +> **Status:** planning record only. No source optimization or 17 ms result is +> claimed by this document. The historical baseline below is retained to make +> the next implementation reproducible. + +## Scope and frozen baseline + +This campaign targets the complete batch-one DFlash2 speculative round measured +with a Qwen3.8-27B NVFP4 workload on four V100 GPUs. Its historical baseline is +`onecat/main@34403018d917054dd7765d5e820ad29c8d342348`; the planning record was +integrated on `main@cadcf1d899b6d7511f815e7ee939b1e4676aff19`. It does not +credit route-hit logs as performance evidence. + +The retained benchmark artifacts identify a mixed NVFP4/channel-FP8 target +with a BF16 LM head and a DFlash2 draft at revision +`dedf8df68adfb1afeaf7b7480c0a0243108177b4`. Private filesystem locations are +intentionally not recorded. The frozen workload uses TP4 on an owned four-V100 +group, probabilistic DFlash2 with seven draft tokens, selector K=16, target +top-k=20, FP8 E5M2 target KV, FP16 draft KV, Flash-V100 for target and draft +attention, and FULL target and draft CUDA Graphs. Model/checkpoint labels are +benchmark evidence only and must not activate an optimized runtime route. + +Two runtime contracts are deliberately kept separate: + +1. The localization contract uses a 32K maximum model length, 512 maximum + batched tokens, and disables prefix caching, Mamba alignment, and parsers. + It is used only for short unprofiled repeats, graph-node traces, and + exact-shape microbenchmarks. +2. The promotion contract uses a 256K maximum model length, 4096 maximum + batched tokens, prefix caching, Mamba alignment, and the Qwen tool and + reasoning parsers. A speedup is not accepted until it reproduces under this + practical configuration. + +The current accepted practical baseline on PR #288 is `18.465--18.537 ms` per +complete round for 512-token coding requests and `18.587--18.603 ms` for +1,024-token coding requests. The high-acceptance MBPP item 28 baseline is +`18.567 ms`, acceptance length `4.686`, and `251.60 token/s`, with a natural +EOS and EvalPlus base/plus `1/1`. A fresh current-source 16-prompt GSM8K run +measures request-mean acceptance `4.45732`, pooled acceptance `4.07740`, and +`19.3688 ms` per round with diagnostic counters enabled. Diagnostic timing is +not a production baseline. + +## Acceptance and runtime-binary audit + +The apparent NVFP4 acceptance regression was reproduced on the frozen 64-row +GSM8K order, but it is not caused by NVFP4. With the same current source and +the same official stochastic contract, FP8 measured request-mean/pooled +acceptance `4.49887/3.89772` and 60/64 accuracy; NVFP4 measured +`4.53642/3.79374` and 59/64. The request mean is slightly higher for NVFP4. +Its lower pooled value comes from a longer generated trajectory and is not a +valid quantization comparison by itself. The corresponding request-mean +completion tokens per verification round are `4.49761` and `4.53214`; this +closely related metric was previously mislabeled as request-mean acceptance. +All future records report both fields explicitly. + +Both current paths are nevertheless below the repaired August 22 FP8 baseline +of `5.38880/4.56444` (`5.38884` request-mean completion tokens per round). +This is a real regression rather than only output-length mix: current FP8 is +lower on 61 of the 64 matched prompts, with paired +request-mean delta `-0.89123` and bootstrap 95% interval +`[-1.02152,-0.75597]`. Restricting each run to outputs of at most 512 tokens +still moves the mean from `5.60920` to `4.63777`. + +Single-variable runs exclude the target quantized operators and compact +sampler as the main cause. Disabling NVFP4 QPN2 on the fixed 16 rows measured +`4.37420` request mean instead of recovering the current `4.57520` control. +Disabling FP8 target QPN8 on all 64 rows measured `4.49095/3.80785`, and +disabling sparse target rejection measured `4.51820/4.02970`; none approaches +the repaired baseline. All completed 64-row variants retained 60/64 FP8 +accuracy. + +The causal fault is binary provenance. The low-acceptance runs loaded +`_C_stable_libtorch` SHA256 `fe986a...9ba8c`, even though the checked-out source +contains the batched-weight RMSNorm repair. That binary predates the per-layer +weight implementation: its focused GPU suite fails 10 of 13 cases, silently +uses weight row zero for the other draft layers, and mismatches 81--96% of +elements on multi-layer fixtures. The repaired binary SHA256 +`7689e5...d449` passes all 13 cases bitwise. The low result also numerically +matches the historical pre-repair `4.5395/3.9725` baseline. + +The source now verifies the loaded stable extension once with deterministic +nonzero per-layer inputs. A conforming binary continues to use the single +grouped kernel after warmup. A stale binary emits an explicit warning and uses +the bitwise per-layer fallback, so a wheel/source or worktree/symlink mismatch +cannot silently lower acceptance again. The new fallback tests pass for both +runtime behaviors; the full CPU DFlash2 suite reports 70 passed and 12 skipped. + +The repaired-binary FP8 closure uses the same frozen 64-row GSM8K contract and +the dense selector baseline. It records request-mean acceptance `5.37699`, +pooled acceptance `4.62217`, and request-mean completion tokens per round +`5.37814`, with 61/64 correct and 62/64 natural stops. Against the repaired +August 22 baseline, paired request-mean delta is `-0.01180` with bootstrap 95% +interval `[-0.08421,0.06282]`; every position-wise pooled acceptance value is +higher. The artifact records stable-extension SHA256 `7689e5...d449`. This +closes the large acceptance regression: the stale-binary run was lower by +`0.87812` request-mean tokens on the same rows, with bootstrap 95% interval +`[0.75099,1.00344]` after repair. Its concurrently collected throughput is not +a standalone performance baseline because NVFP4 compilation occupied the +other TP4 group. + +The matching repaired-binary NVFP4 run records request-mean/pooled acceptance +`5.39851/4.57051`, request-mean completion tokens per round `5.40178`, and +59/64 correct with 62/64 natural stops. Its request mean is `+0.02152` versus +the repaired current-source FP8 run; the paired 95% interval +`[-0.10021,0.14795]` does not distinguish the two quantizations. The lower +NVFP4 pooled value is explained by its longer output trajectory (35,366 versus +32,658 tokens), not lower per-request acceptance. Relative to the stale NVFP4 +run, the repair adds `0.86209` request-mean tokens with 95% interval +`[0.74547,0.98007]` and improves 62 of 64 matched rows. + +The experimental draft-MLP QPN8 route is rejected independently. On the same +fixed 16 requests it changed request mean `4.57520 -> 4.51123` and pooled +acceptance `4.39123 -> 4.09891`; the request delta `-0.06397` exceeds the +`-0.05` gate. Its code and benchmark hook were removed rather than retained as +a nominal 17 ms result. + +The corrected 16K coding-quality pair uses the same NVFP4 target, MBPP-32 +prompts, request seed, graph mode, E5M2 KV, and natural-EOS sampling on both +sides. The dataset tests score no-DFlash `31/32` and DFlash2 `32/32`. +EvalPlus maps 31 of those rows: base is `30/31` versus `31/31`, and plus is +`27/31` versus `27/31`. At the plus level, each route alone passes one +different row; exact paired McNemar `p=1.0`. Both routes naturally stop on 31 +rows and reach the 16K cap on one row. This provides no evidence of a DFlash2 +quality regression under the repaired runtime. + +The first score incorrectly reported one DFlash result as invalid syntax even +though the complete executable answer followed ``. The extractor had +searched illustrative reasoning-era code fences before checking an unfenced +final answer. The scorer now gives the final answer precedence, its regression +fixture parses the affected `comb_sort` result, and both routes were rescored +with the same corrected extractor. The corrected DFlash results are dataset +tests `32/32`, EvalPlus base `31/31`, and EvalPlus plus `27/31`. + +The opt-in selector-QPN8 candidate also passes the fixed 64-row acceptance +gate. Relative to the repaired dense selector, request-mean acceptance changes +`5.39851 -> 5.43951`, pooled acceptance `4.57051 -> 4.58679`, and GSM8K +`59/64 -> 60/64`. The paired request delta is `+0.04101` with bootstrap 95% +interval `[-0.02309,0.10669]`; 23 rows improve, 15 fall, and 26 tie. It +therefore has no acceptance-loss signal. + +Its first clean directional speed signal is positive but not promotable: the +dense run on physical GPUs 4--7 measures request-mean/median complete round +`19.083/19.046 ms`, while QPN8 on GPUs 0--3 measures `18.077/18.029 ms`. +Aggregate output is `218.58 -> 229.41 token/s` and mean steady decode is +`282.68 -> 300.12 token/s`. Because these are different physical groups, they +are not a valid paired speed result. QPN8 also lowers available KV capacity +from 529,616 to 459,817 tokens through its prepared rerank layout. An unrelated +eight-GPU job occupied both groups before the same-GPU repeat; the candidate +remains explicit opt-in and is not a default or a credited 17 ms result. + +## Acceptance gates + +A short-context candidate is accepted only when all of the following hold: + +- the no-diagnostic promotion contract measures a mean and median complete + speculative round at or below `17.0 ms` across at least three independent + steady-state requests; the p90 must be at or below `17.5 ms`; +- the same fixed prompt, seed, sampling parameters, output cap, CUDA Graph + shapes, and GPU set are used for the baseline/candidate pair; +- paired request-mean acceptance does not fall by more than `0.05`, the + per-position counters remain healthy, and no stale-buffer, graph-replay, or + draft-KV mismatch appears; +- the existing GSM8K, MATH-500, HumanEval, corrected MBPP, and WikiText PPL + gates do not regress versus target-only and the accepted DFlash2 baseline; +- every retained performance optimization has an exact engine-contract + admission predicate, a rollback switch, focused numerical coverage, and an + unprofiled endpoint win. Once these gates pass, the optimized path is the + default; rollback remains for diagnosis rather than permanent opt-in. + +Long-context work is evaluated at 32K, 128K, and 256K. The candidate must not +make complete-round latency or pure-decode throughput worse by more than 1% +at any length. Numerical probes require finite outputs and a dtype-appropriate +absolute/relative error envelope; official-sampling quality and acceptance +must remain non-regressing. Bitwise output or greedy-token identity is not a +promotion requirement when a changed but valid reduction order explains the +difference. The 17 ms target applies to the short round; long-context results +are reported as an explicit decay curve because the verifier attention cost +necessarily grows with resident context. + +## Trace-first implementation sequence + +1. Reproduce the current no-diagnostic NVFP4 baseline in this worktree and + record source SHA, extension hashes, route logs, GPU clocks, and per-request + round statistics. +2. Capture the smallest generation-only Nsight Systems trace that contains + steady CUDA Graph replays. Split every round into draft, draft-to-target, + target verification, and target-to-draft phases, then expand the target + phase into exact kernel categories and launch counts. +3. Microbenchmark only the largest current-source bucket at its production + M/N/K, TP, dtype, and graph shapes. A microbenchmark result is directional + evidence, never an endpoint claim. +4. Implement the smallest exact change that removes the measured bottleneck. + Rejected DDTree index reuse, skipped verification, generic auxiliary-stream + overlap, variable K, and approximate reranking are not retried: they either + changed output/acceptance or increased the complete round. +5. Promote a default only after focused CPU/GPU numerical tests, graph replay, + three clean short repeats, the practical 256K configuration, the quality + suite, and the 32K/128K/256K decay sweep pass. + +The first trace should determine whether the remaining roughly `1.6 ms` is in +NVFP4 QPN2/QPN8 projections, draft non-causal KV work, target verification, +or scheduler/metadata boundaries. No saving is assigned to a hypothesis +before that trace and its exact-shape microbenchmark exist. + +## Draft PR record + +Purpose: reduce the complete Qwen3.8-27B NVFP4 DFlash2 round to 17 ms on SM70 +without changing the sampling contract, acceptance, or task quality. + +Test plan: unprofiled paired endpoint repeats; Nsight graph-node phase +breakdown; exact-shape microbenchmarks; focused numerical and CUDA Graph replay +tests; GSM8K/MATH-500/HumanEval/MBPP/PPL quality gates; 32K/128K/256K decay +sweep. + +Test result: the runtime-provenance acceptance repair and paired MBPP quality +gate pass. Selector-QPN8 passes the acceptance gate and has a cross-group +directional speed signal, but remains opt-in pending a same-GPU endpoint pair. +The 17 ms endpoint and long-context performance gates remain pending; no +rejected or unpaired speed candidate is credited. diff --git a/docs/design/sm70_dflash2_target_graph_20ms.md b/docs/design/sm70_dflash2_target_graph_20ms.md new file mode 100644 index 0000000000..3e87b8aac3 --- /dev/null +++ b/docs/design/sm70_dflash2_target_graph_20ms.md @@ -0,0 +1,1229 @@ +# SM70 DFlash2 target-verifier graph: 20 ms control log + +## Scope and frozen baseline + +- Date: 2026-08-23. +- Integration base: `c62c1dcd833458054daa12e134695aa19a4ac609` + (`codex/v100-dflash2-gdn-metadata-20260822-122723`, Draft PR #257). +- Worktree: + `/home/ymzx/桌面/1cat-vllm/worktrees/v100-dflash2-target-graph-20ms-20260823-084427`. +- Branch: `codex/v100-dflash2-target-graph-20ms-20260823-084427`. +- Task cache: `/data/minimax-h3/task-cache/v100-dflash2-target-graph-20ms`. +- Frozen workload: Qwen3.8-27B-FP8 target, official DFlash2 draft, TP4, + batch one, block eight (seven draft tokens), Flash-V100, target E5M2 KV, + draft FP16/auto KV, and CUDA Graph. +- Baseline node trace: + `/data/minimax-h3/task-cache/v100-dflash2-pr257-gdn-metadata/profiles/latest-e2e/dflash2-fused-viewcache-b1-nodes-o128-v5.sqlite`. + +The 30-round, four-rank steady-state baseline (one edge round removed from +each side) is: + +| Phase | Critical-path wall | Nodes or launches per rank | +|---|---:|---:| +| Draft graph | 4.046 ms | - | +| Draft to target | 4.934 ms | 153 kernels | +| Target graph | **24.740 ms** | **2612 nodes** | +| Target to draft | 2.838 ms | - | +| Complete round | 36.300 ms | - | + +The target-graph objective is at most 20.000 ms under this same diagnostic. +Performance acceptance requires an improvement in the complete round. The +final quality contract uses probabilistic task scores, paired acceptance, and +PPL against target-only; greedy token identity is diagnostic evidence rather +than a production gate. The superseding measurements are recorded at the end +of this document. + +## Resource-use decomposition + +The target graph contains 23.365 ms of rank-average GPU kernel service inside +a 24.740 ms critical span, or 94.7% activity-envelope coverage. This rules out +a large idle bubble as the primary cause, but it does not imply efficient SM +or memory-pipeline utilization within each small kernel. Nsight Compute metrics +must remain separate from graph-span and Nsight Systems service time. + +| Kernel category | Rank-average service | Launches per rank | +|---|---:|---:| +| TurboMind FP8 dense GEMM | 10.831 ms | 256 | +| Copy/cast elementwise | 3.387 ms | 941 | +| Other kernels | 2.869 ms | 240 | +| TP all-reduce/communication | 2.316 ms | 128 | +| Other Torch/Triton elementwise | 2.494 ms | 739 | +| RMSNorm/residual | 0.616 ms | 128 | +| LM-head/sample/TP gather | 0.416 ms | 48 | +| Dense GEMV/GEMM/compressor | 0.257 ms | 65 | +| Fill/mask | 0.179 ms | 67 | + +The graph is therefore busy but fragmented: 1680 copy/cast/elementwise nodes +consume 5.881 ms, while the average node is only about 8.9 microseconds. The +20 ms target requires recovering at least 4.740 ms (19.2%); launch reduction +and work fusion are first-order requirements, not optional cleanup. + +## Ordered experiments and stop gates + +1. A/B the repaired, default-off DFlash2 packed GDN verifier with all other + flags frozen. It preserves the current gating values, recurrent arithmetic + order, and FP32-state contract while removing packed-QKV rearrangement and + the final output copy. Record route-hit logs, graph wall, graph nodes, + complete-round wall, exact output tokens, and acceptance trajectory. +2. If the packed route wins but remains above 20 ms, use its new node trace to + isolate residual GDN state/copy traffic. Compare with SGLang-V100's fused + recurrent verifier and all-layer state commit without importing unrelated + scheduler or Eagle/MTP changes. +3. Profile the exact TP4 all-reduce shape. The current 2.316 ms/rank service is + a second independent target; any replacement must preserve collective + ordering and exact target hidden states. +4. Only after the verifier contract is stable, assess a DFlash2-safe variant of + v100-skinny's small-M QPN8 projection work. It is not accepted on the basis + of its non-speculative result. + +Every failed experiment must be reverted or kept behind a default-off gate. +Do not report profiler-instrumented latency as unprofiled throughput. Do not +use target-only runs, and do not occupy a partial TP4 group or terminate an +unrelated process. + +## External references + +- SGLang-V100 fixed source and trace are compared locally at + `/data/minimax-h3/task-cache/v100-dflash2-pr257-gdn-metadata/sources/sglang-v100` + and + `/data/models/v100-dflash2-20260820/sglang-audit/perf-rootcause/sglang-dflash2-single1-step20-v2.nsys-rep`. +- v100-skinny is pinned at `5b589c0dc81223e0ba65bcb3e755874723f8b515`; + its 219.1 tok/s result is a Qwen3.8 mixed-NVFP4/FP8 MTP result, not a DFlash2 + target-verifier baseline. + +## Results + +### Accepted target-graph reduction + +The matched TP4 node trace at source `500130882ba68fb7bac3a0b9e4eb5872647fb045` +uses the same model, draft, request, sampling, KV dtypes, and graph policy as the +frozen baseline: + +- SQLite: + `/data/minimax-h3/task-cache/v100-dflash2-target-graph-20ms/profiles/dflash2-fused-gdn-norm-split-gemma-dynamic-b1-nodes-o128-gpu0123-v3.sqlite`. +- Result JSON: + `/data/minimax-h3/task-cache/v100-dflash2-target-graph-20ms/results/dflash2-fused-gdn-norm-split-gemma-dynamic-b1-nodes-o128-gpu0123-v3.json`. +- Four-rank synchronized target graph, with the first and last rounds removed: + **19.317 ms mean, 19.322 ms p50, 19.444 ms p90, 19.477 ms p99**, and + **1,257 nodes/rank**. +- Rank-average kernel service is 18.552 ms, or 96.04% of the critical span. +- Relative to the frozen 24.740 ms / 2,612-node baseline this removes + **5.423 ms (21.92%)** and **1,355 nodes (51.88%)**. The p99 is also below + the 20 ms objective. +- The benchmark endpoint inside this profiled run reports 125.060 steady + decode token/s; this is trajectory evidence, not an unprofiled throughput + claim. All 128 output + token IDs and hash `fe0300...` match the baseline exactly. Probabilistic + acceptance changed slightly from 4.129 (31 draft rounds) to 4.000 (32 draft + rounds), so the Gemma suffix fusion remains a Type-B candidate pending a + distribution-quality gate rather than being accepted from one trajectory. + +The reduction is cumulative: the one-pass GDN output norm first reached +22.646 ms / 2,084 nodes; fused GDN split materialization reached 21.625 ms / +1,892 nodes; the dynamic-shape-correct Gemma residual RMS fusion reached the +19.317 ms / 1,257-node result above. A prior `v2` launch omitted the installed +`flash_attn_v100_cuda` extension directory from `PYTHONPATH` and failed during +graph capture; it is not a performance result. + +### Remaining graph-boundary cost + +Across 31 steady boundaries per rank, the draft-graph end to target-graph +start interval is 5.776 ms mean (5.582 ms p50). Only 1.304 ms is covered by +non-graph GPU kernels, leaving about 4.47 ms of launch, synchronization, and +CPU queueing bubbles. Every boundary submits 153 kernels: + +- one TP `cross_device_reduce_1stage`: 0.785 ms; +- 20 `DeviceScanKernel`, 20 `DeviceScanInitKernel`, 20 `compute_cuda_kernel`, + and 21 `indexSelectSmallIndex` calls: 0.276 ms service but much larger + serialized submission cost; +- remaining input, slot, block-table, GDN metadata, and elementwise kernels: + about 0.243 ms service. + +The scan/select work arrives in five serial groups with scalar D2H ordering +points. Therefore the next boundary objective is not a faster individual +3-microsecond kernel. It is to construct persistent target metadata once and +capture or batch-submit the five repeated groups while preserving every +buffer/update dependency. + +The same parser applied to the retained SGLang-V100 trace gives a useful +matched structural target. Its draft-to-target interval is 3.398 ms mean with +30 non-graph kernels and 0.115 ms GPU service, versus this branch's 5.776 ms, +153 kernels, and 1.304 ms service. SGLang still pays a 2.715 ms host +`cudaStreamSynchronize`, so it is not a zero-overhead endpoint. The immediate +vLLM gap is nevertheless concrete: remove 123 repeated launches and move the +roughly 0.785 ms TP reduction adjacent to target input/embedding preparation +into the target replay dependency chain. This comparison comes from +`/data/models/v100-dflash2-20260820/sglang-audit/perf-rootcause/sglang-dflash2-single1-step20-v2.sqlite`. + +### Preliminary probabilistic quality pair + +A fixed 16-question sequential GSM8K pair used graph TP4, target E5M2 KV, +draft FP16 KV, block eight, and official `temperature=1.0/top_p=.95/top_k=20` +sampling. The control artifact is +`quality/gsm8k-16-control-gemma-off.json`; the candidate artifact is +`quality/gsm8k-16-candidate-gemma-on.json` under the task cache. + +- Control: 68.75% accuracy, aggregate acceptance length 4.470. +- Candidate: 75.00% accuracy, aggregate acceptance length 4.693. +- Eight of 16 full output-token trajectories match. Eight diverge under + probabilistic sampling, as expected from the accepted one-FP16-ULP numeric + bound; this small sample cannot establish a quality improvement. + +The Gemma fusion stays default-off until paired prompt perplexity and broader +dataset gates show no regression. The GDN norm and split stages retain their +Type-A classification and can be defaulted independently of that decision. + +The first deterministic distribution probe scores 1,850 prompt tokens from 32 +fixed GSM8K questions in eager target prefill (Graph is irrelevant to prompt +logprobs). Weighted perplexity is 4.20119 for the decomposed control and +4.20172 for the Gemma fusion, a +0.00053 / +0.013% change. All 32 next-token +argmaxes match. Mean absolute prompt-logprob difference is 0.00173 and the +worst token is 0.02524. This is small enough to continue dataset testing but +confirms that the kernel is Type B, so it remains default-off. + +### Fused Flash-V100 small-query metadata candidate + +The five repeated scan groups are generated by +`FlashAttnV100MetadataBuilder._update_smallq_decode_metadata`: each of five +target KV groups performs four `repeat_interleave` scans, then materializes and +copies block-table and sequence-length temporaries. The +`VLLM_SM70_DFLASH2_FUSED_SMALLQ_METADATA` candidate writes the persistent +block-table, sequence-length, and query-boundary buffers directly with one +Triton launch per group. It is limited to a DFlash2 target on SM70; draft, +DDTree, Eagle, MTP, and CPU metadata retain the existing path. + +- Exact V100 tests pass for B1 and mixed B3/B4 layouts, negative block IDs, + zero-length padded requests, and graph-token padding. +- The unchanged CPU persistent-buffer, padding, and overflow tests pass. +- A realistic `q=8`, three block columns, five-group microbenchmark reports + 1.278 ms legacy versus 0.114 ms fused wall time and 1.235 versus 0.100 ms GPU + service: **1.164 ms saved / 11.2x** for this isolated constructor. + +The matched TP4 end-to-end trace accepts the candidate as a Type-A default-on +optimization (explicit environment value `0` remains an opt-out): + +- Synchronized draft-to-target falls from 5.720 ms to 1.911 ms (-66.6%). +- Per-rank non-graph work falls from 153 nodes / 1.304 ms GPU service to 23 + nodes / 0.195 ms. All 20 `DeviceScan`, 20 `DeviceScanInit`, and 20 generic + scan-compute launches disappear; `indexSelectSmallIndex` falls from 21 to 1. +- The five replacement launches total 0.016 ms of GPU service. +- The full speculative round falls from 31.791 ms to 27.951 ms (-12.1%), while + the target graph remains effectively unchanged at 19.261 ms / 1,257 nodes. +- Profiled single-request steady decode rises from 125.06 to 142.37 token/s + (+13.8%). The 128-token output hash, aggregate acceptance length 4.0, and + per-position accepted counts `[28, 23, 18, 14, 8, 3, 2]` match exactly. + +The accepted trace is +`profiles/dflash2-fused-gdn-norm-split-gemma-smallqmeta-b1-nodes-o128-gpu0123-v1.sqlite` +under the task cache. Its 1.911 ms synchronized boundary is also below the +3.398 ms boundary in the retained SGLang-V100 audit trace, though that external +trace is a structural reference rather than a fully matched throughput run. + +### Next boundary: target-to-draft + +The accepted trace leaves a stable 2.738 ms target-to-next-draft interval on +rank 0 (33 steady intervals). A representative `M=8` interval and the aggregate +kernel service separate it as follows: + +| Stage | Representative wall / service | +| --- | ---: | +| Dense target LM head, local FP16 shard | 0.947 ms | +| TP full-vocabulary all-gather | 0.139 ms | +| Full-vocabulary top-k/top-p | 0.606 ms | +| Rejection statistics, rejection, and resample | 0.225 ms | +| Target-hidden-to-draft KV precompute | about 0.60 ms | +| Post-update and next-draft metadata | about 0.22 ms | + +The next candidate must therefore attack the first four rows as one semantic +unit. For the fixed no-penalty `top_k=20, top_p=0.95` contract, each TP rank can +compute its exact local top 20, exchange only 20 `(score, token-id)` pairs per +row, merge the global top 20, and perform rejection/residual sampling on that +sparse support. The repository already contains the SM70 TurboMind FP16 +LM-head top-20 epilogue, but the `_C` binary linked into this worktree predates +that op. The first TP4 microbenchmark was therefore rejected at route checking +without recording a timing; an incremental SM70 build is required before this +candidate can pass the microbenchmark gate. + +### Compact target sampling microbenchmarks and gated integration + +The production-shaped TP4 communication probe uses eight verifier rows, +vocabulary 248,320, target top-K 20, and four V100s. It compares the existing +full-vocabulary all-gather plus global top-K with local top-K followed by two +compact all-gathers and a global merge: + +- Exact top-K token IDs and values match. +- The full path moves 993,280 bytes per rank; the two-gather compact path moves + 1,920 bytes, a 517x reduction. +- Despite that byte reduction, p50 is 0.3318 ms for full gather plus top-K and + 0.3292 ms for the compact path. Local top-K (0.1208 ms), two collectives + (0.0911 ms), and compact merge (0.0896 ms) consume the saved transport time. +- Packing values and IDs into one collective is slower at 0.3538 ms. + +Compact communication alone is therefore rejected as a performance feature. +The retained artifact is +`results/tp4-compact-verifier-logits-m8-v248320-k20.json`. + +The second probe fuses compact target top-p, DFlash2 acceptance, log-domain +`relu(p-q)` recovery, and token-keyed Gumbel resampling into one Triton program +per request. At the real B1/block-eight shape it is token/count exact against +the dense rejection path for block 4/8 and `top_p` 1.0/0.95. Its p50 is: + +| Isolated stage | p50 | +| --- | ---: | +| Dense top-K/top-p | 0.8366 ms | +| Dense rejection only | 0.2202 ms | +| Dense top-K/top-p plus rejection | 0.8663 ms | +| Compact top-p plus sparse rejection | **0.0901 ms** | + +Combining this result with the measured TP4 compact-candidate transport gives +an expected real saving of roughly 0.5-0.6 ms per complete round. The artifact +is `results/dflash2-sparse-rejection-b8-v248320-k20-q16-v1.json`. + +The candidate is integrated behind default-off +`VLLM_SM70_DFLASH2_SPARSE_TARGET_REJECTION=1`. Its route is limited to SM70, +MRV2 DFlash2, one active decode request, probabilistic proposals, target +`top_k=20`, positive temperature, `0 < top_p <= 1`, and no grammar, min-p, +penalty, logit bias, bad words, NaN reporting, or logprob request. Every other +configuration computes the unchanged dense logits and uses the shared +rejection sampler. DFlash2 now keeps the 16 proposal IDs and FP32 realized +scores in persistent request-slot order in addition to the dense fallback +cache; a V100 poisoned/reordered-slot test verifies the mapping. + +Validation completed before the full-model run: + +- DFlash2 CPU suite: 46 passed, 11 CUDA-only skipped. +- Strict route and fallback tests: 13 passed. +- Sparse dense-equivalence plus selector/cache CUDA tests: 6 passed. +- Ruff, formatting, compileall, and `git diff --check`: passed. + +The matched single-request Graph/Nsight run remains pending because both TP4 +GPU groups became occupied by unrelated work after these tests. No external +process was interrupted. + +### Residual target-graph service after the 27.951 ms round + +The accepted 1,257-node target graph has 18.357 ms of rank-0 GPU service inside +its 19.261 ms critical span. Its largest exact buckets per replay are: + +| Target graph bucket | Nodes | GPU service | +| --- | ---: | ---: | +| TurboMind FP8 GEMM | 256 | 10.730 ms | +| vLLM one-stage TP4 all-reduce | 128 | 2.313 ms | +| recurrent GDN verifier | 48 | 1.213 ms | +| Flash-V100 partition/reduce | 32 | 0.802 ms | +| fused Gemma residual/RMS suffix | 127 | 0.643 ms | +| generic elementwise | 144 | 0.541 ms | + +The FP8 GEMMs split into three stable launch shapes: 64 calls at 4.394 ms, +128 calls at 3.924 ms, and 64 calls at 2.411 ms. They already match the retained +SGLang target GEMM total closely, so replacing them is a separate QPN8 quality +and throughput project rather than the next low-risk leaf. + +The TP4 collective remains the clearest next `>1 ms` gap. vLLM launches ten +512-thread CTAs and averages 18.07 microseconds per call. The matched +SGLang-V100 trace uses its JIT one-shot push kernel with eighty 128-thread CTAs +and averages 12.36 microseconds. Across 128 calls this structural difference is +about 0.73 ms of service; the broader matched audit found a 0.7-1.2 ms +per-graph TP reduction gap depending on critical-rank accounting. The next +microbenchmark should compare the exact `[8,5120]` FP16 target shape under one +captured 128-call chain, first sweeping the existing vLLM CTA/thread policy, +then porting SGLang's two-epoch push buffer only if the policy sweep cannot +close the gap. Bitwise FP32 accumulation order, CUDA-Graph registration, and +all 128 collective dependencies are hard gates. + +### Sparse rejection end-to-end acceptance + +The matched Graph/Nsight run completed after the GPU group became available: + +- Result: + `results/dflash2-sparse-target-rejection-b1-nodes-o128-gpu0123-v1.json`. +- SQLite: + `profiles/dflash2-sparse-target-rejection-b1-nodes-o128-gpu0123-v1.sqlite`. +- Full synchronized round: **27.166 ms**, down from 27.951 ms by 0.785 ms. +- Target-to-draft: **1.994 ms**, down from 2.738 ms by 0.744 ms. +- Draft, draft-to-target, and target graph remain 4.004, 1.977, and + 19.191 ms respectively. +- All 128 output IDs and hash `fe0300...` match the accepted baseline. The + run records acceptance length 4.129 and per-position counts + `[27, 22, 18, 14, 9, 4, 3]`. + +This is a Type-A result for the narrow sparse route contract. The environment +gate remains explicit until the default policy and paired dataset result are +landed together. + +### TP4 push all-reduce audit + +The existing vLLM CTA/thread-policy sweep did not improve the 128-collective +chain. A pinned SGLang-V100 checkout at `845b9fdf7a7e` measured its two-epoch +one-shot push collective at 0.846-0.886 ms for the same 128-call +`[8,5120]` FP16 chain, versus 1.854 ms for vLLM's pull path. A first vLLM port +measured 0.873 ms and matched the FP32 rank-order reference numerically across +eager, graph, and 128-call graph replay. + +Two issues prevented accepting that first port: + +1. Its temporary `_C` reused objects from the independent QPN8 worktree. Both + push-on and push-off runs first diverged from the accepted output at token + 24, proving that those runs were not a clean source comparison. +2. SGLang's positive-zero empty-slot sentinel converts an all-positive-zero + reduction to negative zero. A dedicated signed-zero probe found 40,960 + bit mismatches despite zero numerical error. This can perturb later finite + arithmetic and is below the required bitwise gate. + +A release/acquire CTA-ready protocol fixed every payload bit for exact, +random, model-like, positive-zero, and signed-zero inputs over a 128-call graph +chain, but regressed the chain to 1.884-1.901 ms. It is rejected. The accepted +microbenchmark candidate instead preinitializes the two data epochs with a +reserved FP16 NaN payload and retains SGLang's fine-grained polling. Finite +inputs, including both signed zeros, are never rewritten. + +The NaN-sentinel candidate passed the strict TP4 microbenchmark gate: + +- Artifact: + `results/vllm-push-ar-nan-sentinel-correctness-timing-m8-h5120-v1.json`. +- Exact-integer, rank-marker, random-small, model-like, positive-zero, and + signed-zero inputs are bitwise equal to the fixed rank-order FP32 reference + in eager, one-call Graph, and 128-call Graph-chain modes on every rank. +- The 128-call chain is **0.850-0.856 ms**, or about 1.00 ms faster than the + retained 1.854 ms vLLM pull chain. One call averages 6.64-6.68 microseconds + inside the chain. +- The task-owned `_C` contains no QPN8 operator and has SHA256 + `275594c0b38a358c1683efa1d6351e1c572d6da992789d3b5210a74ecbfce1e8`. + +The first unprofiled full-model run with that clean extension completed at +176.56 steady decode token/s. Its 128 token IDs, hash `0cc12...`, acceptance +length 4.467, and per-position accepted counts +`[28, 23, 14, 13, 11, 8, 7]` exactly match the retained push-off control built +with the current CUDA toolchain. They differ from the older `fe030...` release +extension trajectory at token 24; both push-on and push-off current-toolchain +builds make that same transition, so it is not attributed to the collective. +The result is +`results/dflash2-sparse-push-ar-nan-pure-b1-o128-gpu0123-v1.json`. + +A same-binary push-off rerun was invalidated when another task claimed the TP4 +group during model initialization, before KV-cache allocation; it produced no +generation or performance result. The push route therefore remains +default-off until that paired rerun and a clean Nsight complete-round trace are +recorded. The microbenchmark establishes an expected roughly 1 ms graph saving, +but the 176.56 token/s endpoint is not yet used as the paired throughput claim. + +### DFlash2 target QPN8 backport + +The accepted FP8 QPN8 operator and production route were cherry-picked from +`db978179de` and `e0547f1b4a`. The audited integration admits only explicit +QPN8 opt-in, SM70/TP4, block-FP8 layout, exact projection shapes, and bounded +concurrency. It does not inspect a model, architecture, or checkpoint identity; +unsupported tensor or runtime contracts retain TurboMind. + +The target LM head is not a QPN8 shape and remains on the dense FP16 path. The +only admitted target projections are fused gate/up, down, and GDN/full-attention +output. GDN input and full-attention QKV also remain TurboMind. This preserves +the DFlash2 selector and target-sampling contracts while attacking the three +measured FP8 verifier buckets. + +- The imported target-only evidence remains documented in + `sm70_qwen38_qpn8_decode.md`: WikiText perplexity and all fixed GSM8K, MMLU, + and C-Eval metrics pass, and the operator race passed 960 rows. +- The DFlash2-specific runtime-isolation tests pass for target-only, DFlash2, + DFlash1, and a generic non-DFlash speculative configuration. Together with + environment and warmup coverage, 70 focused CPU tests pass. +- A task-owned combined extension contains both the bitwise-gated push + collective and all six QPN8 operators. Its SHA256 is + `512b44cd49e99f26e0f091c68c11bf8c17e5575ffe2c42ba9af1de69938a29e9`. + +No DFlash2 throughput claim is made from the target-only result. The required +next gate is a production-shape M=8 operator smoke followed by paired +single-request DFlash2 generation and a clean Nsight phase trace. The working +hypothesis is that QPN8 plus push can remove enough of the 19.2 ms target graph +to approach a 23 ms complete round, but only the new trace may accept or reject +that hypothesis. + +### Quality-safe 23.678 ms checkpoint + +The next accepted stack keeps the sharded target-hidden projection disabled +and collapses the five full-attention metadata refreshes into one pointer-table +kernel. The paired single-request node traces measure 24.424136 ms for the +control and **23.678170 ms** for the candidate. Draft-to-target falls from +1.857847 to 1.213916 ms. All token IDs, token hash, acceptance length, and +per-position acceptance counts are exact. This grouped metadata route proved +the measured workload, but remains default-off in the audited source so the +kernel contract can be rolled out independently of a model identity. + +The retained target graph has 1,254 nodes and about 16.055 ms of rank-critical +GPU service. Its largest buckets are 4.235 ms gated-pair QPN8, 3.186 ms QPN8 +`(16,1)`, 2.059 ms QPN8 `(16,2)`, 1.277 ms TP4 push all-reduce, 1.205 ms +recurrent GDN, and 0.734 ms Flash-V100 partition attention. This decomposition +is the active short-context optimization budget; the 20 ms complete-round goal +requires at least another 3.679 ms from the measured checkpoint, not merely a +lower isolated kernel number. + +Caching the prepared GDN metadata tensor tuple removes repeated empty-tensor +allocation and tuple assembly. A ten-group/four-layer microbenchmark falls +from 50.519 to 2.663 microseconds and the targeted metadata test passes. The +post-convolution Q/K/V packer is bitwise exact in seven V100 tests and improves +the production-shape graph from 7.475 to 4.116 microseconds, projecting about +0.161 ms over 48 GDN layers. It stays default-off because the packed-GDN route +already subsumes the same rearrangement. + +The graph-internal text-embedding experiment is rejected and fully removed. +Full target graph capture failed in the AOT Qwen signature with +`AttributeError: 'NoneType' object has no attribute 'size'` after omitting the +previously compiled `inputs_embeds` tensor. It generated no tokens and has no +performance result; the failed run must never be counted toward the 20 ms +goal. + +### QPN8 plus TP4 push fusion rejection + +An exact fused epilogue was implemented for the 128 target out/down projection +calls. It rounded the local FP32 QPN8 accumulator to FP16, pushed all four +ranks' packets, accumulated ranks 0 through 3 in FP32, and rounded the result +to FP16. A TP4 128-call CUDA-Graph chain passed every per-rank bitwise check +against both the separate QPN8-plus-push path and the fixed rank-order +reference. + +The performance gate failed: the separate chain was 2.696 ms median, while the +fused chain was 3.249 ms, or 21.06 versus 25.38 microseconds per call. Keeping +the 512-thread QPN CTA resident while one warp polls peer packets increases SM +and NVLink contention instead of hiding communication. The entire production +route, environment switch, bindings, and tests were removed. Only the +standalone benchmark artifact is retained at +`results/qpn8-fused-push-tp4-m8-n5120-v1.json`; this negative 0.552 ms must not +appear in an endpoint speed claim. + +### 2026-08-24 external implementation audit + +The one permitted upstream-head refresh fetched SGLang-V100 at +`0083b9fd83a601b1fcd9a691f7240be4e6be111e`. Its new native long-context work +adds grouped q=1 exact split-KV decode, grouped multi-token verification, and +exact E5M2 byte-to-FP16 expansion. The implementation assigns a Qwen3.8 TP4 +verification layer to two GQA subgroups and at most 40 context splits, instead +of scanning the same cache independently for all eight query rows. + +The same redundant-read diagnosis was independently merged in +[vLLM-Metal PR #534](https://github.com/vllm-project/vllm-metal/pull/534). +That implementation shares each KV block across two verification-window rows +while retaining each row's original accumulation order and reports bitwise +parity. This is stronger evidence than copying a benchmark number: it +identifies redundant KV traffic as a kernel-shape problem specific to +multi-token verification and supplies a quality-preserving scheduling pattern. + +`dnv2003/v100-skinny` remains a useful QPN8 geometry inventory. Its graph +sweep reports `(8,2)` around 12.27 microseconds for M=8/K=1536/N=5120, while +the active target trace uses `(16,1)` around 24.9 microseconds. The kernels and +decoder contracts are not assumed identical. A task-owned sweep of the exact +current operator is therefore required before changing the table; this is now +the highest-value short-context microbenchmark after rejecting QPN/AR fusion. + +### Long-context verification microbenchmark + +The pinned SGLang TileLang 0.1.8 compiler required its matching +`apache-tvm-ffi==0.1.8.post2` and the CUDA 12.8 toolkit still mounted at +`/mnt/nvme3n1p2/usr/local/cuda-12.8`. These are off-tree measurement tools, +not new runtime dependencies. The production implementation will be native +Flash-V100 code. + +On the same V100 and shape `B=1, Q=8, Hq=6, Hkv=1, D=256, page=16, E5M2`, the +measured per-layer medians are: + +| Prefix | Flash-V100 scalar rows | Flash-V100 XQA rows | SGLang grouped verify | +| ---: | ---: | ---: | ---: | +| 1K | 0.238 ms | 0.315 ms | **0.122 ms** | +| 32K | 1.390 ms | 1.151 ms | **0.527 ms** | +| 128K | 5.119 ms | 3.561 ms | **1.957 ms** | +| 256K | 9.064 ms | 6.481 ms | **3.803 ms** | + +The actual current policy selects XQA for E5M2 at and above 16K. Relative to +that production long-context family, grouped verification is 2.18x, 1.82x, +and 1.70x faster at 32K, 128K, and 256K. Across the target's 16 full-attention +layers, the isolated projection is roughly 9.99, 25.67, and 42.84 ms removed +per complete verification round at those lengths. Short-context dispatch must +still use the existing scalar path when the captured workspace is small: the +accepted 8K-capacity trace measures only about 45.5 microseconds per layer, +below the grouped kernel's fixed cost. + +At 1K, both grouped and XQA outputs are within 1.526e-5 maximum and about +1.6e-6 mean absolute error of the FP32 reference. This is currently +`B-pending`, not a default-enablement result: the port must align block tables, +causal boundaries, E5M2 scale semantics, partition order, and shuffled physical +pages; then pass non-growing numeric bounds, greedy trajectory, acceptance, +PPL, and the fixed GSM8K/MATH-500/HumanEval/MBPP/MT-Bench quality suites. + +The first packed-GDN node trace on GPUs 4-7 produced no result. A separate +task claimed the group during initialization, the reservation guard reported +`TP4_HANDOFF_COLLISION_ABORT`, and the run exited before model generation. +No external process was stopped and no timing from that attempt is retained. + +### 2026-08-24 verifier follow-up: rejected seams and two bounded candidates + +The complete-round acceptance gate remains **strictly below 20 ms**. The +latest retained quality-safe result is still 23.678170 ms, so the unresolved +measured gap is at least 3.679 ms. Isolated microbenchmarks below are not +endpoint evidence and must not be presented as completion of that gate. + +Two short-path seam experiments were rejected before integration: + +- The existing fused pull-all-reduce plus Gemma RMS operator measured + 3.005-3.013 ms for a TP4 128-pair graph chain, versus 1.463-1.469 ms for the + current push-all-reduce plus RMS sequence. It was also not bitwise equal on + two ranks. No source route was enabled. +- Reusing the existing paged-prefill operator for the Q=8 verifier measured + 0.436/14.872/57.935/115.838 ms per layer at 1K/32K/128K/256K. Six head CTAs + still serially scan the whole prefix, so merely sharing the eight query rows + does not solve context parallelism. This route is rejected. + +The Flash-V100 split-KV framework was then extended narrowly to E5M2. The +implementation keeps the existing FP16 routes unchanged, dispatches the +generic paged loader with `KV_DTYPE=fp8_e5m2`, preserves separate K/V scales, +and uses the existing split merge order. The source compiles successfully; +the task build has SHA256 `e457f06f...9b0f23`. The production-shape graph +microbenchmark, non-unit scales, shuffled physical pages, and endpoint route +remain pending because all eight GPUs are occupied by external TP4 jobs. + +For the short-context target graph, the local +`Qwen3.8-27B-NVFP4` checkpoint was audited rather than assuming a foreign +checkpoint contract. Its config hash is `1b3c7186...3a5c`; layers 0-55 use +native NVFP4 only for the MLP trunk, while attention, GDN, layers 56-63, and +the full-vocabulary LM head remain FP8. This makes it a useful quality-gated +candidate: the DFlash2-sensitive head is not reduced to four bits, while most +of the verifier's MLP weight traffic is halved. The already accepted no-MTP +route for this exact checkpoint measures 14.017 ms/token, but that number is +not a DFlash2 verification-round result. + +The existing current-tree TurboMind NVFP4 implementation is the control. A +prior real-weight M=1 QPN4 experiment saved a projected 1.090 ms across the 56 +gate/up and 56 down calls, but it is M=1-only and therefore does not establish +an M=8 verifier gain. A new bounded benchmark, +`benchmarks/kernels/benchmark_sm70_nvfp4_qpn2.py`, now loads the real TP-local +layer-55 shards and compares M=8 CUDA-graph replay against the production +TurboMind route. The pinned v100-skinny QPN2 source compiled into a task-only +extension with SHA256 `09d96053...633a4`; no production dispatch has been +changed. The next free-GPU action is this M=8 race, followed by a full DFlash2 +boot only if the measured aggregate saving and numerical gate are positive. + +The mixed checkpoint may become a default performance profile only after +paired FP8-target comparisons pass WikiText PPL, fixed GSM8K, MATH-500, +HumanEval, MBPP, and MT-Bench gates, plus acceptance-length and per-position +acceptance checks. A low-bit microbenchmark alone cannot authorize that +change. + +### 2026-08-24 QPN2 production admission and split-KV rejection + +The generic Flash-V100 E5M2 split-KV follow-up is now rejected rather than +pending. With production shape `B=1,Q=8,Hq=6,Hkv=1,D=256`, non-unit scales, +and shuffled physical pages, its per-layer medians were 0.124/1.936/7.791/ +15.769 ms at 1K/32K/128K/256K. The 1K result is numerically healthy (maximum +absolute error `1.526e-5` versus FP32), but the long-context results are worse +than the retained XQA path. The generic split kernel still rereads one KV +stream per query head; it does not implement SGLang's grouped-GQA sharing and +therefore cannot be enabled. The raw artifact is +`results/flash-v100-splitkv-verify-e5m2-b1-q8-g6-d256-v1.json`; the task-only +binary SHA256 is `e457f06f...9b0f23`. + +The real-weight QPN2 operator race is positive at exact verifier width M=8, +but its control is the mixed NVFP4 checkpoint's existing TurboMind route, not +the FP8+QPN8 target used by the accepted 23.678170 ms complete-round baseline. +Across layer-0/rank-3 gate/up and down projections, replacing TurboMind with +QPN2 projects **3.814523 ms removed per complete mixed-target round**: +gate/up falls from 96.020 to 40.681 microseconds and down from 31.687 to +18.909 microseconds. Relative L2 error versus the FP32 checkpoint oracle is +`0.000569` for QPN2 gate/up versus `0.000598` for TurboMind, and `0.000300` +for QPN2 down versus `0.000291` for TurboMind. The retained artifacts are +`results/nvfp4-qpn2-vs-turbomind-real-layer0-rank3-m8-oracle-v1.json` and +`results/nvfp4-qpn2-vs-turbomind-real-layer55-rank0-m8-oracle-v1.json`. + +That 3.814523 ms value must not be subtracted directly from 23.678170 ms. +Against the actual FP8+QPN8 baseline buckets, the current evidence projects +only about 1.8 ms, with a plausible 1.8-2.5 ms range pending the production +fused-gate measurement. QPN2 alone therefore does not prove the sub-20 ms +gate; a new complete single-request DFlash2 trace must measure the movement, +and any remaining 1-2 ms gap must be removed from that trace's actual +critical-rank buckets. + +A production QPN2 implementation is linked behind +`VLLM_SM70_NVFP4_QPN2=1`. Admission uses TP4 and the checkpoint-native gate/up +or down tensor layout; it does not inspect model or speculator identity. M<=8 +uses QPN2, while larger dynamic M dispatches inside the opaque C++ operator to +the existing TurboMind layout. The switch remains default-off. + +The exact production source compiles for SM70, its four routing tests pass, +and a full `_C` candidate contains QPN8, graph-only push, and all four QPN2 +operators. Candidate path: +`build/push-qpn8-qpn2-graphonly-v3/_C.abi3.so`, SHA256 +`c133500e...3f068b5`. The earlier task-only standalone hash is obsolete after +the final production-source edits and is intentionally not used as endpoint +evidence. + +The production candidate passed the real-weight M=8 CUDA-Graph oracle on a +V100. Gate/up is 96.759 versus 36.801 microseconds and down is 31.830 versus +19.794 microseconds against TurboMind, for a 4.031684 ms mixed-target +projection. QPN2 relative L2 versus the FP32 checkpoint oracle is `0.000569` +for the fused gate/up output and `0.000300` for down. The result is +`results/nvfp4-qpn2-production-layer0-rank3-m8-v1.json`, SHA256 +`cfc0634d...1a5a6c`. This passes the operator gate but does not alter the +correct-baseline warning above: only a complete DFlash2 trace can prove the +net movement from 23.678170 ms. + +The first full-service launch exposed a checkpoint contract mismatch rather +than a QPN2 failure. The mixed checkpoint declares calibrated E4M3 KV scales, +which upstream vLLM correctly refuses to reinterpret as E5M2 scales. This is +the behavior retained by upstream PR #45040. A task-cache-only model overlay +therefore keeps every original weight/tokenizer file by symlink but removes +the incompatible `kv_cache_scheme` from its copied config, making the explicit +E5M2 run use the same unit-scale storage contract as the accepted FP8 target. +The original checkpoint is unchanged. Overlay config SHA256 is +`5b72ecd7...40f6e52`; the benchmark now records both that hash and the resolved +weight path. + +The overlay launch passed the KV contract, but an unrelated TP4 job acquired +GPUs 4-7 between the availability check and worker initialization. Free +memory fell to 24.41 GiB, below the matched 0.8 utilization contract, so vLLM +refused startup before model loading. No endpoint or timing result is claimed +from either failed launch. The next valid action remains the matched complete +single-request trace on a genuinely free TP4 group; the hard acceptance gate +is strictly below 20.000 ms. + +The next owned launch reached DFlash2 warmup and exposed a second checkpoint +contract rather than a timing result: the mixed checkpoint quantizes +`lm_head`, while DFlash2 candidate selection intentionally requires an +unquantized target LM head. The final task-cache overlay therefore maps only +`lm_head.weight` to the accepted FP8 target's BF16 `[248320,5120]` tensor and +removes `lm_head` from both the mixed quantization target list and its FP8 scale +entry. All other target weights still resolve to the original mixed +checkpoint. The combined overlay config and index SHA256 values are +`1b418636...01e8a` and `dacc00bf...836e`; the BF16 sidecar SHA256 is +`ddff1d66...b698ff`. The original FP8 and mixed model directories remain +unchanged. + +The launch audit also found that the production checkpoint selects +`CompressedTensorsW4A4Fp4` because its config contains dynamic NVFP4 input +activation metadata. The first QPN2 integration had been attached to the +W4A16 scheme, so its real-weight microbenchmark remained valid but the full +model did not actually dispatch QPN2. The exact admission and fused +gate/up/down dispatch now live in the W4A4 scheme, and the W4A16 file is back to +its original content. The routing, FP8-prefill, and DFlash2 targeted suites +pass **78 tests**. No endpoint speed claim is made until the newly queued +single-request node trace contains the QPN2 kernels and measures a complete +round strictly below the 20.000 ms gate. + +Because the default safetensors iterator visits every tensor in each referenced +shard, deleting only the index entry was insufficient: the original monolithic +body still exposed `lm_head.weight_scale` during loading. A streaming +task-cache rewrite now removes exactly `lm_head.weight` and +`lm_head.weight_scale` from the body without decoding or re-encoding any other +tensor. The slim body is 21,296,297,056 bytes with SHA256 +`9f49b437...389267`; all 1,951 retained tensor dtype/shape records match and +the first, middle, and last bytes of every retained tensor match the source. +The final index SHA256 is `94f500bf...c7ef2c`, and the benchmark records that +index plus every actually referenced weight realpath. This replaces neither +the original checkpoint nor the DFlash2 unquantized-LM-head guard. + +The first valid mixed-target node trace is a rejection, not a speed result to +retain. Its complete round is **40.775579 ms**: draft 3.941496 ms, +draft-to-target 1.228315 ms, target 33.567780 ms, and target-to-draft +2.037988 ms. The target graph has 1,462 nodes. QPN2 itself matches its +microbenchmark, consuming 2.521294 ms for 56 gate/up calls and 1.518048 ms for +56 down calls. The regression is instead 144 residual channel-FP8 projections +falling through to SM70 Marlin for **21.110017 ms**. The single-request output +has acceptance length 3.823529; neither its 91.451 tok/s steady decode nor that +one-prompt acceptance value is a quality or performance acceptance claim. + +The missing dependency was already implemented and DCO-signed in local commit +`5123a6a7b2` (`[Kernel] Recover Qwen3.8 NVFP4 SM70 decode`): its narrow +compressed-tensors W8A16 channel-FP8 adapter reuses the accepted TurboMind and +QPN8 operators. Only that adapter, its fused-SiLU delegation seam, and focused +tests were ported here; unrelated attention and tuning changes from the commit +were not duplicated. This also fixes the QPN2 production seam: without +`CompressedTensorsLinearMethod.apply_fused_silu_and_mul`, the full model ran a +separate Triton activation even though the scheme exposed the fused operator. +The combined compressed-FP8, QPN2, FP8-runtime, and DFlash2 suites pass **83 +tests**. A new complete trace must prove removal of the 21.110 ms Marlin bucket +and still meet the strict sub-20.000 ms gate. + +### 2026-08-24 mixed-target recovery and remaining 1.807 ms + +The channel-FP8 adapter and QPN2 route are now present in a clean combined +extension. The retained `v8` complete single-request trace uses the mixed +NVFP4 target, QPN2 MLP, QPN8 residual projections, graph-only TP4 push, +sharded target-hidden projection, E5M2 target KV, FP16 draft KV, and an +unquantized target LM head. Across 140 steady rank-cycles it measures: + +| Phase | Mean wall | +| --- | ---: | +| Draft graph | 3.914734 ms | +| Draft to target | 1.097309 ms | +| Target graph | 15.367610 ms | +| Target to draft | 1.729242 ms | +| Complete round | **22.108896 ms** | + +The target graph has 1,254 nodes. Its largest service buckets are QPN8 +projections (3.327 ms), QPN2 gate/up (2.511 ms), QPN2 down (1.568 ms), push +all-reduce (1.294 ms), recurrent GDN (1.200 ms), Flash-V100 attention +(0.738 ms), and the 127 Gemma suffix kernels (0.676 ms). Both the draft +selector and target rejection still execute an approximately 0.98 ms dense +FP16 LM head. + +The draft-to-target interval contained six four-byte device-to-host copies and +stream synchronizations. Five were accidental: the five Flash-V100 groups +recomputed `seq_lens_cpu` even though `InputBatch.seq_lens_cpu_upper_bound` +already held the required contract. Passing that tensor through +`MambaHybridModelState.prepare_attn()` removes exactly those five syncs. The +new focused unit test passes. The corrected `v10` trace, with the sharded +target-hidden projection explicitly restored after reboot, measures: + +| Phase | Mean wall | +| --- | ---: | +| Draft graph | 3.914056 ms | +| Draft to target | **0.790517 ms** | +| Target graph | 15.357529 ms | +| Target to draft | 1.745331 ms | +| Complete round | **21.807433 ms** | + +The complete median/min/max are 21.823373/21.557190/22.653044 ms. This is the +current accepted local best, but it **does not pass** the strict sub-20.000 ms +gate: the measured deficit remains 1.807433 ms. The one remaining D2H fence +is the batch-level GDN metadata ordering assertion. A retained historical +GPU-assert experiment measured 24.093540 ms versus 23.679850 ms for its +same-generation synchronized control, so it is rejected; the final ordering +edge will not be removed on the basis of a theoretical sync saving. + +At the exact production LM-head shape `M=8,K=5120,N_local=62080`, a new V100 +probe measures `torch.mm` at 1.018880 ms p50 and the existing SM70 dense +full-logits operator at 0.812032 ms p50. The 0.206848 ms per-call gain projects +to about 0.414 ms for the two heads in one complete round. The probe's top-16 +support and values are exact, but full-logit non-bitwise behavior still makes +this a quality-gated candidate rather than accepted endpoint evidence. Raw +artifact: `results/lmhead-dense-sm70-vs-torch-m8-n62080-k5120-v11.json`. + +The next bounded candidate preserves the accepted push transport and fuses +its following FP32 residual/Gemma RMSNorm suffix. It targets the 128 push +collectives plus 127 RMS kernels as a 128-node chain. It must first beat the +current production-shaped TP4 chain microbenchmark, keep the reduced residual +bitwise exact, and satisfy the FP16 normalized-output bound. Only then may it +be combined with the LM-head candidate in a full Graph/Nsight run; neither +microbenchmark may be subtracted from 21.807433 ms to claim completion. + +That push-suffix candidate is now rejected and fully removed from the +production source. The first centralized-row implementation measured +2.544-2.550 ms for the 128-pair TP4 graph chain versus 1.475-1.480 ms for the +current push-plus-Triton sequence. A second implementation fused each CTA's +partial square sum and distributed normalization across all five row CTAs; it +reduced the candidate to 1.638 ms but the matched control remained +1.470-1.472 ms. The best fused form is therefore still about 11.3% slower. +Both versions kept residual output bitwise exact; the second version was +normalized-output bitwise exact on three ranks and differed in two FP16 values +on the fourth, with maximum absolute error 0.000244. Artifacts are +`results/tp4-push-vs-push-fused-gemma-rms-m8-h5120-v2.json` and `v3.json`. +No endpoint run or speed claim is permitted from this failed branch. + +### 2026-08-24 synchronized-fence removal and 20.746 ms endpoint + +The final batch-level GDN metadata assertion was audited against both the +metadata builders and `ModelCudaGraphManager.run_fullgraph()`. All metadata +kernels, copies, and graph replay are enqueued on the same current CUDA stream; +the device `.item()` checked a debug invariant but supplied no ordering edge. +The assertion is separately gated by `VLLM_SM70_DFLASH2_GDN_SYNC_ASSERT`. +After the exact token/acceptance trajectory passed, the production default was +changed to off; it remains an explicit opt-in metadata audit. A structural CPU +invariant remains unconditional. + +The `v17` single-request Graph/Nsight trace disables only that debug fence and +enables the already quality-gated dense FP16 LM-head route. Across 140 steady +rank-cycles it measures: + +| Phase | Mean wall | +| --- | ---: | +| Draft graph | 3.756561 ms | +| Draft to target | 0.156985 ms | +| Target graph | 15.288068 ms | +| Target to draft | 1.543990 ms | +| Complete round | **20.745603 ms** | + +The complete median/min/max are 20.747498/20.523095/20.939874 ms. This is a +1.061830 ms improvement over `v10`, but it still misses the strict sub-20 ms +gate by 0.745603 ms and is therefore not an accepted endpoint. The output +token SHA256 remains +`22db7b2e81df17d3afdc6982fc1425769335c6cc7931da10e64bd6bac3459f28`; +the answer is correct and the complete acceptance contract is unchanged at +3.685714 with per-position counts `[28,22,17,12,8,5,2]`. Steady decode is +174.016 tokens/s. Result/report/sqlite SHA256 values are respectively +`1ee462b4...78a2e`, `7f686d51...f3ba`, and `1f82d151...3302`. + +The remaining two local LM-head passes are now the bounded next target. With +the real rank-0 FP16 shard `[62080,5120]`, channel-QPN8 local top-64 contains +every dense top-16 token in 512/512 random rows. Re-evaluating those 64 rows +with the original FP16 weight via `index_select` plus batched dot products +reduces the complete local selector from 0.904192 to 0.518144 ms p50. The +reranked top-1 matches 511/512 rows; maximum/mean absolute logit differences +are 0.00390625 and 6.68e-6. Artifact SHA256 is +`39fde704...718a7`. These random-hidden results authorize only a production +candidate, not a quality or endpoint claim. + +The implementation is isolated behind +`VLLM_SM70_DFLASH2_QPN8_RERANK=0`. It admits only SM70, TP4, the exact local +LM-head shape, FP16 weights, zero vocabulary padding, top-k 16/20, and at most +eight rows. QPN8 supplies only a coarse top-64 support; the returned logits +are always recomputed from the original FP16 rows. The same shared target LM +head serves both the draft selector and sparse target rejection. An eager-only +shadow gate executes the candidate, logs real-hidden support/set/top-1 +agreement, and returns the dense result so the audit trajectory cannot change. +The path remains default-off until four-rank real-weight coverage, real-hidden +shadow coverage, dataset quality/acceptance, and a new complete Graph trace all +pass. Eighteen focused CPU routing tests and static checks currently pass. + +The stricter four-shard `v18` probe uses local top-20, candidate width 64, +1,024 random rows per TP rank, and the same real checkpoint weights. All +4,096 dense top-20 sets are fully contained in the QPN8 support (minimum +recall 1.0 on every rank). Fixed split-8/non-fast QPN8 is 0.425-0.427 ms; +complete QPN8 plus original-FP16 rerank is 0.499-0.531 ms versus +0.893-0.923 ms for the dense selector. The reranked top-1 matches 4,081/4,096 +rows, while exact local top-20 set agreement is 3,971/4,096. The maximum +returned-logit difference is one FP16 step (0.00390625). Thus the coarse +support gate passes, but reduction-order differences still require real-hidden +and model-quality gates; `v18` alone does not authorize endpoint use. The four +result SHA256 prefixes are `b07428ec`, `85e94de8`, `e56ed9b0`, and `fb4c1016`. +The complete CPU DFlash2 targeted file now passes 52 tests with 12 GPU tests +skipped under an explicitly hidden CUDA environment. + +The `v19` eager shadow then exercised the same candidate on real hidden states +while always returning the dense local result. Across four TP ranks it +observed 8,564 local rows (draft top-16 and target top-20), with **zero** dense +tokens missing from the QPN8 top-64 support. Original-FP16 rerank set agreement +is 8,406/8,564 (98.1551%) and local top-1 agreement is 8,547/8,564 (99.8015%); +the largest ordered logit difference is 0.015625. This passes the coarse +support gate and quantifies the remaining reduction/tie-order effect. Because +shadow is eager-only and deliberately returns dense outputs, neither its +throughput nor its four truncated GSM8K answers are candidate quality evidence. +Log/result SHA256 values are `6dad2b23...482e` and `29c8ff88...753f`. The next +required artifact is the active single-request Graph/Nsight trace; the hard +gate remains a measured complete mean strictly below 20.000 ms. + +### 2026-08-24 exact-rerank endpoint and sub-20 ms replication + +The first active Graph launch, `v20`, is rejected before timing. Startup, +compilation, and QPN8 preparation succeeded, but the draft selector's TP +all-gather rejected a non-contiguous top-16 tensor. The implementation had +sliced the first 16 columns from one persistent `[8,20]` output, retaining a +row stride of 20. It now owns separate graph-stable `[8,16]` and `[8,20]` +value, position, and ID buffers. Row-only views preserve contiguity without a +runtime copy. Six layout cases (top-16/top-20 by 1/7/8 rows), the target +candidate route, Ruff, and `git diff --check` pass. `v20` produced no timing +or quality result. + +The corrected `v21` single-request probabilistic CUDA Graph/Nsight trace +crosses the strict complete-mean gate for the first time. Across 140 steady +rank-cycles on physical GPUs 4--7 it measures: + +| Phase | Mean wall | +| --- | ---: | +| Draft graph | 3.332724 ms | +| Draft to target | 0.158264 ms | +| Target graph | 15.281917 ms | +| Target to draft | 1.163312 ms | +| Complete round | **19.936216 ms** | + +The complete median/min/max are 19.956518/19.751641/20.129946 ms. Steady +single-request decode is 180.686 tokens/s. The 128-token SHA256 is exactly +`22db7b2e81df17d3afdc6982fc1425769335c6cc7931da10e64bd6bac3459f28`, +the answer is correct, and acceptance is exactly unchanged at 3.685714 with +per-position counts `[28,22,17,12,8,5,2]`. Result/report/sqlite SHA256 values +are `3d22d740...e1248`, `27175b2f...09c7`, and `5fe2773c...60e0`. + +An unchanged `v22` replication then ran on physical GPUs 0--3. Its 140 +steady rank-cycles measure 3.339601 ms draft, 0.155429 ms draft-to-target, +15.330372 ms target, 1.165597 ms target-to-draft, and **19.990999 ms** complete +mean. The median/min/max are 20.011181/19.785692/20.212695 ms, and steady +decode is 181.177 tokens/s. Token SHA256, correct answer, acceptance length, +and every per-position acceptance count again match the dense control. The +result/report/sqlite SHA256 values are `06fa010e...19b`, `af23b773...031`, and +`c34083f8...f62e`. + +Thus the complete mean is below 20.000 ms in two independent runs and on both +physical TP4 groups; this is direct Graph/Nsight evidence rather than a +microbenchmark projection. The margin is still only 0.009--0.064 ms and the +upper tail crosses 20 ms, so the candidate remains default-off pending a +larger stability margin and the multi-dataset/PPL quality gate. The next +bounded microbenchmark checks whether the QPN8 support top-64 can skip ordering +before exact FP16 reranking; candidate order is not consumed by the rerank. + +The `v23` sorted/unsorted support experiment rejects that optimization. On the +real rank-0 shard, removing QPN8 top-64 ordering changes complete rerank p50 +from 0.513024 to 0.501760 ms. In the production fixed split-8 core the saving +is only about 4.1 microseconds per LM-head call, or about 8 microseconds per +complete DFlash2 round. This is below run-to-run noise and changes tie order, +so production retains sorted support. The sorted/unsorted artifact SHA256 +values are `5eb1ad6f...31de` and `da3499cc...fb6`. + +The accepted next micro-optimization removes the 5.2 MiB FP16 +`index_select` temporary and its batched matrix multiplication. A fixed-grid +Triton kernel directly computes each selected original-FP16 vocabulary-row dot +product. For the production `[8,62080,5120]`, top-64 to top-20 shape, the +complete selector p50 falls from 0.517120 to 0.464896 ms. All 512/512 top-20 +sets and top-1 IDs match the gathered-bmm reference; maximum ordered-logit +difference is 0.001953125. The selected configuration is `BLOCK_K=1024`, four +warps, with a graph-stable 512-program grid masking inactive rows. The +artifact SHA256 is `b898c410...0631`. Dedicated V100 tests for one, seven, and +eight rows pass, as do the focused CPU routing/layout tests and static checks. + +Two unchanged full Graph/Nsight runs then validate the direct indexed-dot +endpoint on both TP4 groups: + +| Run / physical GPUs | Draft | D-to-T | Target | T-to-D | Complete mean | Median | Max | Steady decode | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| `v25` / 4--7 | 3.324374 | 0.157883 | 15.272063 | 1.139965 | **19.894284 ms** | 19.909407 | 20.087468 | 182.516 tok/s | +| `v26` / 0--3 | 3.332363 | 0.157819 | 15.313486 | 1.145328 | **19.948997 ms** | 19.960269 | 20.148380 | 181.216 tok/s | + +Each result contains 140 steady rank-cycles. Both target and draft are FULL +CUDA Graphs and the indexed-dot path is captured without a graph break. Both +runs reproduce the baseline 128-token SHA256 +`22db7b2e81df17d3afdc6982fc1425769335c6cc7931da10e64bd6bac3459f28`, +the correct answer, acceptance length 3.685714, and per-position counts +`[28,22,17,12,8,5,2]` exactly. The `v25` result/report/sqlite SHA256 values are +`c82ef281...0952`, `f9c26647...6c81`, and `7c7645ef...006a`; the corresponding +`v26` values are `04498a79...0cbc`, `0ebd0b86...4244`, and +`3de27587...9d1`. Parsed round artifacts have SHA256 +`a49fe2cd...be3b` and `bb9d09dc...5e2d`. + +The strict complete-mean and median target is therefore independently +replicated below 20 ms after the final kernel change. It is not yet an +all-sample tail claim: the observed maxima remain 20.087--20.148 ms. The +rerank and GDN debug-fence defaults remain conservative until paired +multi-dataset generation and prompt-logprob/PPL gates pass; a further +0.2--0.4 ms safety-margin reduction remains the immediate performance target. + +### 2026-08-24 paired quality rejection and dense-order correction + +The first paired 16-question GSM8K quality run (`v27`) rejects the active +QPN8 reranker despite its exact single-prompt trace. Both variants are +sequential single-request FULL Graph runs with probabilistic draft sampling, +temperature 1.0, top-p 0.95, top-k 20, identical shuffled rows and request +seed. Dense control scores 11/16 with aggregate acceptance 4.490716; the +candidate scores 9/16 with acceptance 4.412224. The acceptance delta is +-0.078493, below the fixed -0.05 gate, and only 5/16 token streams match. +Artifacts `gsm8k16-qpn8-indexeddot-v27-{control,candidate}.json` have SHA256 +`ced400d1...aee` and `e4c89772...4eeb`; the failed comparison has SHA256 +`9650f531...045c`. Consequently the candidate remains default-off and its +sub-20 traces are not accepted as a default quality-safe endpoint. + +The follow-up `v28` microbenchmark corrects the reference used during rerank +development. It compares against the production TurboMind +`sm70_f16_gemm_out`, not `index_select+bmm`. Across 1,024 random rows on the +four real TP shards, QPN top-64 contains every production dense top-20 token. +Depending on reduction configuration, 99.969--99.994% of candidate logits are +bitwise equal, yet candidate-order top-20 set equality is only 250--253/256 +rows per shard. This isolates a second source of drift: `torch.topk` tie +selection over a shuffled 64-element support does not reproduce its selection +over the original 62,080-element vocabulary order. + +`v29` restores that dense order by scattering the 64 re-evaluated FP16 logits +into a persistent full-shard buffer initialized to negative infinity before +local top-k. All four shards then match production dense output exactly for +all 1,024 rows: ordered token IDs, top-20 sets, top-1 IDs, and returned FP16 +values. Complete local selector p50 is 0.5007--0.5038 ms, versus +0.4608--0.4628 ms without dense-order restoration and 0.894--0.930 ms for the +production dense selector. The four result SHA256 values are +`93ebf045...15ec`, `6ec9f526...c064`, `9e4e9868...2a`, and +`904790b6...bb3`. Focused CPU tests (12) and V100 tests (4) pass. A new +same-source paired GSM8K run is required before this correction can re-enter a +complete Graph/Nsight trace. + +The same-source `v30` paired run rejects the corrected active path at the +acceptance gate. Both variants are sequential single-request FULL Graph runs +over the same 16 shuffled GSM8K questions, probabilistic draft sampling, +temperature 1.0, top-p 0.95, top-k 20, and identical request seeds; the only +environment difference is `VLLM_SM70_DFLASH2_QPN8_RERANK=0/1`. Control and +candidate both score 11/16 with zero invalid answers, but pooled acceptance +changes from 4.581867 to 4.498670 (delta -0.083197), below the fixed -0.05 +gate. Seven of 16 token streams are identical. Candidate mean steady decode +is 249.052 token/s versus 244.506 token/s for control, but that speed is not +accepted because the paired acceptance gate fails. Control, candidate, and +comparison SHA256 values are respectively `a3289717...db1`, +`d446ec34...9776`, and `6b32738b...71f`. The reranker remains default-off and +must not be included in a quality-safe sub-20 ms claim. + +This result also narrows the remaining numerical mechanism. The earlier +real-hidden shadow observed zero dense top-k tokens missing from QPN8 top-64, +so the failure is not coarse-support recall. The candidate recomputes only +the selected rows with a different reduction schedule from the production +TurboMind dense GEMM; real-hidden ordered logits can differ by up to 0.015625. +Restoring dense vocabulary tie order fixes exact random-row tests but cannot +restore those production reduction values. The next selector candidate must +either execute the selected rows with the same Tensor-Core reduction contract +as the dense kernel or pass a larger paired distribution/acceptance gate; it +cannot be promoted from random-hidden equality alone. + +### 2026-08-24 exact context-projection overlap rejection + +The default-off `v31` experiment moves the position-independent DFlash2 +target-hidden combine, TP4 sharded context projection, and fused five-layer K/V +projection onto an auxiliary CUDA stream before target sampling. The main +stream waits on one event only when RoPE and cache slots are available. It is +numerically exact in the full single-request probabilistic run: the 128-token +SHA256 remains `22db7b2e...9f28`, the answer remains correct, and acceptance is +exactly 3.685714 with per-position counts `[28,22,17,12,8,5,2]`. Both target +and draft remain FULL CUDA Graphs and there is no deadlock or graph break. + +Performance rejects the overlap. Across 140 steady rank-cycles, the trace +measures 3.802137 ms draft, 0.254715 ms draft-to-target, 16.041994 ms target, +1.748872 ms target-to-draft, and **21.847718 ms complete**. Complete +median/min/max are 21.858843/21.612554/22.046909 ms, and observed steady decode +is 165.686 token/s. This is about 1.10 ms slower than the 20.745603 ms exact +`v17` baseline, so the gate remains default-off. + +The dual-stream timeline explains why subtracting the isolated projection cost +was invalid. The 0.824 ms dense LM head occupies the SMs first; the 0.118 ms +context projection cannot start until its tail. It then overlaps the many +small target top-k/rejection kernels and stretches the main sampling tail by +about 0.25 ms. The auxiliary all-gather and K/V projection finish before the +event wait, but the contention and cross-stream launch schedule cost more than +the hidden work. Result, qdstrm, report, sqlite, and parsed-round SHA256 values +are `7a75a159...181`, `0cfa07d8...17d`, `83a05726...99e`, +`1891c169...6c1`, and `6b87003b...869`. No speed from this experiment is +accepted; a future overlap attempt needs a fused single-stream schedule or an +explicitly prioritized kernel, not another generic auxiliary stream. + +### 2026-08-24 native grouped verifier and long-context result + +The retained verifier optimization does not reuse DDTree indices or skip target +attention. It implements the MRV2 block-eight semantics directly in +Flash-V100: eight query rows share each E5M2 K/V page, accumulate all six GQA +heads, and write the eight verifier outputs in one pass. Admission requires +SM70, FP16 Q, E5M2 KV, B1, Q=8, Hq=6, Hkv=1, D=256, causal decode, and the +original single-request metadata. All other shapes keep the previous path. +The dedicated CUDA-Graph test poisons and replays the sequence length and cache +pages; six V100 cases pass against the FP32 oracle. + +The paired same-process CUDA-Graph microbenchmark at page size 3,296 measures: + +| Context | Native one-pass | Native two-pass | Previous XQA | +| ---: | ---: | ---: | ---: | +| 1K | 0.068608 ms | 0.097280 ms | 0.230400 ms | +| 32K | 0.356352 ms | 0.544768 ms | 1.024000 ms | +| 128K | 1.214464 ms | 1.864704 ms | 3.021312 ms | +| 256K | 2.396160 ms | 3.690496 ms | 6.011392 ms | + +The microbenchmark is +`results/flash-v100-grouped-paired-page3296-v72.json`, SHA256 +`aeccb306...657d9`. It establishes the isolated attention saving; endpoint +claims use the full-model traces below. + +Matched 32K and warmed 128K single-request FULL-Graph runs preserve the output +token stream and complete acceptance counters within each on/off pair: + +| Input | Grouped | Complete round | Target phase | Steady decode | +| ---: | :---: | ---: | ---: | ---: | +| 32K | off | 36.310865 ms | 30.248113 ms | 77.537 tok/s | +| 32K | on | **26.361447 ms** | **20.451644 ms** | **88.392 tok/s** | +| 128K | off | 69.854780 ms | 63.708214 ms | 106.428 tok/s | +| 128K | on | **40.778072 ms** | **34.646866 ms** | **182.294 tok/s** | + +At 32K the complete round falls 27.40% and the target phase 32.39%; at 128K +they fall 41.62% and 45.62%. The 128K pure-decode result rises 71.28%. The +repetitive fixed long-context prompt has acceptance 6.095 at 32K and 7.647 at +128K in both variants; these values are paired route checks, not corpus +acceptance estimates. The four phase artifacts and SHA256 values are +`v77`/`2d705aa9...0c0e`, `v76`/`de6d922d...4235`, +`v79`/`e61e858f...cc1c`, and `v80`/`76164c8c...8deb`. + +The short fixed 128-token trace is neutral rather than a claimed win: grouped +is 21.117097 ms versus 21.060181 ms off. The long-context admission threshold +therefore remains `max_model_len >= 32768`; the feature can be disabled with +`VLLM_FLASH_V100_DFLASH2_GROUPED_VERIFY=0`. Dataset generation below is the +deciding endpoint gate, where grouped improves aggregate throughput despite +mixed prompt and output lengths. + +### 2026-08-24 replicated sub-20 ms short-round bundle + +Before the grouped-attention addition, the complete candidate bundle was +replicated on both physical TP4 groups under FULL target and draft CUDA Graphs: + +| Run | Draft | D-to-T | Target | T-to-D | Complete mean / median / max | +| --- | ---: | ---: | ---: | ---: | ---: | +| `v52-r2` | 3.374848 | 0.156565 | 15.292271 | 1.073660 | **19.897343 / 19.895756 / 20.098654 ms** | +| `v53` | 3.385334 | 0.157777 | 15.320752 | 1.079131 | **19.942993 / 19.948144 / 20.145704 ms** | + +Thus mean and median are below 20 ms twice; this is not an all-sample-tail +claim because both maxima remain slightly above 20 ms. Phase artifacts have +SHA256 `eb8feded...5482` and `412ad1a0...7599`. These traces authorize the +performance candidate, while the larger score/PPL run below authorizes its +sampling and model-quality behavior. + +### 2026-08-24 probabilistic 2K score and PPL gate + +Per the final acceptance decision, greedy token-for-token equivalence is no +longer a production gate. Quality is evaluated with probabilistic DFlash2, +temperature 1.0, top-p 0.95, target top-k 20, reasoning effort `xhigh`, fixed +seed, `max_tokens=2048`, and natural EOS. Every run is sequential B1 and uses +the same target checkpoint, TP4, E5M2 target KV, FP16 draft KV, Flash-V100, and +FULL Graph. The target is the local mixed NVFP4/FP8 checkpoint with a BF16 +LM-head; it must not be described as the official unquantized checkpoint. + +The first 128-row dataset contains 32 GSM8K, 32 MATH-500, 32 HumanEval, and 32 +raw-MBPP prompts. The raw MBPP prompts copied the speed benchmark and omitted +the function name and tests, so their 0--2/32 results are invalid as task +scores and excluded. The remaining 96 scored cases are: + +| Variant | GSM8K | MATH-500 | HumanEval | Included total | +| --- | ---: | ---: | ---: | ---: | +| Target-only | 29/32 | 26/32 | 28/32 | **83/96 (86.458%)** | +| Dense verifier DFlash2 | 30/32 | 24/32 | 28/32 | 82/96 (85.417%) | +| Native grouped DFlash2 | 30/32 | 25/32 | 28/32 | **83/96 (86.458%)** | + +Target-only versus grouped has three wins on each side and McNemar exact +two-sided p=1.0. Dense versus grouped has two dense-only and three grouped-only +wins. Grouped therefore matches target-only aggregate score and improves one +case over the dense-verifier DFlash2 sample; no task-score regression is +observed. + +The corrected MBPP prompt includes the requested function name and all three +held-out tests, following the Google MBPP evaluation contract. Target-only and +grouped both score **25/32 (78.125%)**. They have two exclusive wins each, 23 +shared passes, five shared failures, and McNemar exact p=1.0. Target has eight +length finishes and grouped six. The grouped MBPP run overlapped an unrelated +process that appeared after reservation, so its throughput is discarded; +pass/fail scoring completed without OOM or engine error. + +Across the valid 96-case mixed run, performance is: + +| Variant | Aggregate output | Mean request steady decode | Mean completion tokens / verify step | +| --- | ---: | ---: | ---: | +| Target-only | 67.128 tok/s | 68.008 tok/s | n/a | +| Dense verifier DFlash2 | 176.341 tok/s | 218.436 tok/s | 4.2243 | +| Native grouped DFlash2 | **193.484 tok/s** | **231.126 tok/s** | **4.2384** | + +Grouped is +9.72% in aggregate output and +5.81% in mean steady decode versus +the dense verifier. Its pooled acceptance is 3.749104 versus 3.749153 dense; +the difference is negligible, while the request-mean completion length is ++0.0141. The summary artifact is +`quality/dflash2-mixed32-quality-2k-v83-summary.json`, SHA256 +`9ca37f6a...dab7`. Corrected MBPP score artifacts have SHA256 +`92908b64...08f3` and `326e559b...339e`. + +The supplementary WikiText-2 gate scores 16,376 prompt tokens in eight equal +2,048-token windows. Token-weighted/geometric PPL is 5.497872 target-only and +5.497376 grouped, a relative **-0.0090%** change. Mean absolute prompt-logprob +difference is 0.007519 and the maximum is 1.101103. The generic comparison tool +labels this as failure only because its historical bounds require every +logprob and every per-window PPL to agree within `1e-6`; that exactness contract +is intentionally superseded here by the user-approved score/PPL contract. +Aggregate PPL does not regress. Artifact SHA256 is `d2c54784...e0c7`. + +The official DFlash2 model card publishes lossless-distribution semantics plus +acceptance/throughput, but no GSM8K, MATH-500, HumanEval, or MBPP task accuracy. +Its published acceptance lengths are 5.46/5.28/4.39/4.79 respectively under +H200, SGLang, FlashAttention-3, block eight, maximum 4,096 tokens, and the full +official datasets. They are useful reference points, not a direct comparison +to this 32-row-per-suite V100/mixed-checkpoint/max-2K gate. See +. + +### Audit disposition and isolation + +The score, PPL, paired acceptance, microbenchmark, short-trace, and long-context +results remain evidence for the measured workload. They do not install a model +profile. Every new leaf is default-off and independently selected by explicit +environment opt-in plus hardware, algorithm configuration, dtype, tensor +layout, cache/graph mode, sampling, and bounded-concurrency contracts. + +The earlier clean 8-token TP4 auto-profile smoke remains historical route +evidence only; automatic profile injection was removed during merge audit. +Unsupported requests execute the previous dense/general path rather than +changing routing semantics. 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_prefix_anchored_swa.md b/docs/design/sm70_prefix_anchored_swa.md new file mode 100644 index 0000000000..3c0207f63b --- /dev/null +++ b/docs/design/sm70_prefix_anchored_swa.md @@ -0,0 +1,57 @@ +# SM70 Prefix-Anchored Sliding-Window Attention + +Prefix-anchored sliding-window attention is an explicit, model-agnostic +inference-engine option for bounded decode KV usage on V100/SM70. The request +prefix remains globally visible, while generated tokens attend only to the +most recent configured window: + +```text +visible(q, kv) = kv <= q and (kv < prefix_len or q - kv < window) +``` + +Enable it through `AttentionConfig`; checkpoint metadata and model architecture +names do not activate or gate the route: + +```bash +vllm serve \ + --attention-config '{"prefix_anchored_decode_window": 4096}' +``` + +The option is off by default. It changes attention semantics after generation +exceeds the configured window, so it is not an exact-output optimization for a +full-attention workload. + +## Current Engine Contract + +The engine admits the route only when all of these operator/runtime contracts +hold: + +- NVIDIA SM70 and the `FLASH_ATTN_V100` backend; +- causal, standard full decoder attention with fp16 KV; +- equal key and value head sizes, with no attention sinks or per-layer window; +- decode and prefill context-parallel sizes both equal to one; +- no speculative decoding, KV connector, or KV-cache offloading. + +Prefix caching is disabled, and full CUDA graphs are reduced to piecewise +graphs. Unsupported combinations fail during configuration or layer +construction. Once enabled, missing or mismatched per-request prefix metadata +is a hard runtime error; the engine never evicts gap blocks and then silently +runs an unmasked kernel. + +## Validation Status + +Validation on the current `main` source includes: + +- 32 CPU tests covering the mask, generic engine admission, fail-closed + metadata, cache-spec registration, and middle-gap block reclamation; +- a clean SM70 source build of the Flash-V100 extension; +- 19 V100 operator tests for paged decode and paged prefill, including GQA, + multiple head sizes, multi-partition decode, full-prompt and continuation + shapes, a negative control, and the default-off path. +- same-toolchain comparison against the current `origin/main` build: all 55 + default-off scalar decode kernel instances and all 64 default-off paged + prefill kernel instances are SASS-identical. + +No model end-to-end or throughput claim is attached to this route yet. Treat +performance and workload-level quality as separate acceptance gates before +enabling it in a production configuration. diff --git a/docs/design/sm70_qwen38_130_acceptance.md b/docs/design/sm70_qwen38_130_acceptance.md index 30fe226511..7da048b4c2 100644 --- a/docs/design/sm70_qwen38_130_acceptance.md +++ b/docs/design/sm70_qwen38_130_acceptance.md @@ -131,22 +131,28 @@ issue, not an idle-GPU or missing graph-shape issue. ## Wheel And Compatibility - Wheel: `1cat_vllm-1.3.0-cp312-cp312-linux_x86_64.whl`. -- SHA256: `e4ffba3bf61e16dcba0eeeafac4036b7cead5abe047e4474ecacb7e5abb35da3` - (75,765,498 bytes). +- Final release source: `6ada86ed64af6d1a7b3cb0f34df237fd86f06d48`. +- SHA256: `2bdb14a9c44f83ee6a766d88ed0d85b11390d6f5d65747e8dbe80a8e2d5d63e0` + (68,764,308 bytes). - A cloned environment imports vLLM 1.3.0 from site-packages with Torch 2.10.0+cu128. `_moe_C.moe_permute_sort_workspace_size` is present. - The wheel contains `_C`, `_C_stable_libtorch`, `_moe_C`, Flash-V100, vLLM FA2, and `flash_qla_sm70_gdn_strided.so`. - A clean-cache MTP4 plus FP8-KV boundary run does not create `TORCH_EXTENSIONS_DIR`, proving FlashQLA does not invoke NVCC at runtime. -- The 2026-08-17 release candidate was packaged from the accepted SM70 build - at source `383bbdd166`, installed into a fresh Python 3.12 environment, and - served Qwen3.8-27B-FP8 TP4 with the `fp8` shorthand resolving to E5M2. In - both no-MTP and MTP4 - lanes, repeated fixed-seed greedy responses were byte-identical and official - sampling stopped naturally without replacement characters. MTP's first - request may still JIT Triton helper kernels; this is a cold-latency concern, - not a numerical-output difference. +- The final 2026-08-17 wheel uses source-controlled Python from the release + source above and the accepted SM70 native payload at source `383bbdd166`. + Every `.so` member and `vllm-rs` is byte-identical to that accepted payload; + the precompiled packaging path preserves both Flash-V100 and the SM70 + FlashQLA extension. In an isolated Python 3.12 install, all native modules + and FlashQLA load successfully; the MTP mixed-batch guard and Flash-V100 + policy suites pass 87 tests. +- The accepted runtime evidence remains Qwen3.8-27B-FP8 TP4 with the `fp8` + shorthand resolving to E5M2. In both no-MTP and MTP4 lanes, repeated + fixed-seed greedy responses were byte-identical and official sampling + stopped naturally without replacement characters. MTP's first request may + still JIT Triton helper kernels; this is a cold-latency concern, not a + numerical-output difference. - Qwen3.6-35B-A3B-AWQ TP2 final-wheel checks pass. Official 1K/256 and 4K/1024 pure decode are 95.63 and 95.29 tok/s (10.457 and 10.494 ms TPOT), with complete, non-corrupted outputs and the TurboMind AWQ MoE route. 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 e36b263378..2ef310aaf6 100644 --- a/docs/design/sm70_qwen38_fp8_prefill_decay.md +++ b/docs/design/sm70_qwen38_fp8_prefill_decay.md @@ -123,8 +123,315 @@ The extra address-generation/scalar-publication issue cost exceeds the saved P-load replay. Do not retry this layout spelling or the earlier shuffle/repack form. +## 2026-08-24 Exact-D256 K-Stage Ping-Pong + +The accepted exact-N32 arithmetic order is unchanged. The candidate alternates +K D64 panels between two non-overlapping shared-memory stages. The second K +stage borrows bytes from the later V/P allocation: K is dead before V is +loaded and P is materialized, so those lifetimes do not overlap. A next-K +fragment is published to the stage that the current QK phase is not reading, +so the pre-overwrite full-CTA barrier is no longer required. The visibility +barrier before the next QK phase remains. + +The K layout's logical size is 2048 half elements, but its padded physical +`cosize` is 2188. The final stage stride is therefore 2240 half elements: it +is larger than the physical span, 16-byte aligned for `STS.128`, and keeps +both stages on the same 128-byte shared-bank phase. The final lifetime-aliased +layout leaves dynamic shared memory at 45,568 bytes. Dense/split-KV3 kernels +remain at 254/253 registers per thread with zero stack or spill. + +SASS for the exact dense nonpaged specialization changes as follows: + +- static `BAR.SYNC.DEFER_BLOCKING` count: 15 to 12; +- HMMA step counts: 128 each, unchanged; +- `LDG`, `LDS`, `STS`, `SHFL`, and `MUFU` counts: unchanged; +- `FADD` and `IADD3` static counts: 15 to 12 and 71 to 68; +- full SASS from the clean repository-patch build is byte-identical to the + measured lifetime-alias artifact. + +The initial fully-disjoint v5 stage layout was screened across the length +curve. Same-build, separate-process A/B/A CUDA-event measurements use FP16, +`Hq=6`, `Hkv=1`, D256, causal attention, five warmups, five queued calls per +sample, and seven samples. The control column is the mean of the two bracketing +control medians except for the first 8K chunk, where the first process was a +visible clock-ramp outlier and the hot second control is reported. + +| Q | KV | Control | Ping-pong | Latency change | Exact | +|---:|---:|---:|---:|---:|---:| +| 8192 | 8192 | 4.8259 ms | 4.8302 ms | +0.09% | bitwise | +| 8192 | 32768 | 32.5750 ms | 32.2132 ms | -1.11% | bitwise | +| 8192 | 65536 | 68.2283 ms | 67.7216 ms | -0.74% | bitwise | +| 8192 | 131072 | 142.3325 ms | 140.9176 ms | -0.99% | bitwise | +| 15680 | 125440 | 251.6402 ms | 250.2695 ms | -0.54% | bitwise | + +A second 128K run from the clean v5 build measured +`142.1020 -> 141.5512 ms` (`-0.39%`) and was again bitwise exact. The final +lifetime-alias layout was then compared with v5 in both execution orders, +each bracketed by fresh controls. Across the two orders, pooled control, v5, +and lifetime-alias medians are `142.4068`, `141.8076`, and `141.7669 ms`: +`-0.42%` for v5 and `-0.45%` for the final layout. The `0.03%` difference +between candidates is noise-sized; the final layout is selected because it +does not enlarge shared memory or alter the original V/P addresses. + +Direct paged and split-KV3 gates at `Q4096/KV32768` are bitwise equal over +6,291,456 elements each; the dense gate covers 12,582,912 elements. The gain +is small but length-directed: the hot first chunk is neutral while long-prefix +calls consistently improve. + +Closed intermediate variants must not be repeated: + +- 2048-half stage spacing overlaps the 2188-half K physical span and changes + almost every output element; +- 2188-half spacing removes overlap but misaligns the second stage for + `STS.128` and raises a CUDA misaligned-address fault; +- 2192-half spacing is exact but loses address/bank-phase efficiency; +- 2304-half spacing is exact but is slower than the minimal same-phase 2240 + spacing; +- allocating the two 2240-half stages as fully disjoint storage is exact and + performance-equivalent, but unnecessarily grows shared memory to 46,336 + bytes; the lifetime-alias form retains the original 45,568-byte envelope. + +Nsight Compute recognized the exact mangled kernel but the host driver denied +performance-counter access with `ERR_NVGPUCTRPERM`. No privilege or machine +security setting was changed. CUDA-event timing and static SASS are therefore +the timing and structural authorities for this experiment. + +The matched end-to-end TP4 128K A/B/A gate also passes. It fixes Qwen3.8-27B +FP8 weights, 131,072 input tokens, 256 sampled output tokens, E5M2 KV, chunk +8192, no MTP, prefix caching, Mamba align, Flash-V100, CUDA graphs, and +`temperature=1.0/top_p=0.95/top_k=20/seed=20260824`. The control is the mean +of the two bracketing baseline runs. + +| Metric | Bracketed control | Lifetime alias | Change | +|---|---:|---:|---:| +| prefill | 47.9233 s | 47.6687 s | -0.531% | +| prefill throughput | 2735.0 tok/s | 2749.6 tok/s | +0.534% | +| TTFT | 47.9498 s | 47.6937 s | -0.534% | +| pure decode | 5.1557 s | 5.1547 s | neutral (+0.020% speed) | + +All three runs produce the same 256-token output hash +`2db7ef09...503325a`. Every rank reports the same route counts, including 32 +exact dense D256 calls, 544 prefix-paged calls, and 544 E5M2 bridge calls. +This promotes the small dependency-chain win, but it also establishes that +barrier removal alone is not a double-digit long-context solution. + +## 2026-08-24 Q8000 Split-KV3 Tail-Wave Experiment + +The matched runtime does not present `Q=8192` to long prefix-attention calls. +Although `max_num_batched_tokens` is 8192, Mamba/attention page alignment makes +the repeated full chunks `Q=8000`; a 40K diagnostic observed +`Q=(1,8000,6,256)` and `KV=(1,40000,1,256)`. The first Q8192 experiment was +therefore a useful operator screen but had zero end-to-end route hits. Its +first control measurement was also invalidated by a second TP4 job starting on +GPUs 4-7 between warmup and measurement (`47.90 -> 64.45 s`). The uncontended +candidate/control-B pair used the same unsplit route and differed by only +0.15%, confirming that run did not measure split-KV3. + +The corrected candidate admits exactly `Q=(1,8000,6,256)` behind the +default-off +`VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3_Q8000_EXPERIMENTAL` gate. It leaves +the existing production Q4096 policy unchanged and does not admit Q8192 or +other nearby shapes. Three independent KV partitions turn 750 Q tiles into +2250 CTAs at Q8000, avoiding the short last wave of the 72-SM V100. Bracketed +CUDA-event operator results from the final K-stage binary are: + +| Q | KV | Unsplit | Split-KV3 | Throughput change | Causal TFLOP/s | +|---:|---:|---:|---:|---:|---:| +| 8000 | 40000 | 40.1736 ms | 38.6519 ms | +3.94% | 44.05 -> 45.78 | +| 8000 | 64000 | 66.5214 ms | 64.6753 ms | +2.85% | 44.33 -> 45.60 | +| 8000 | 96000 | 102.8820 ms | 99.4545 ms | +3.45% | 43.95 -> 45.47 | +| 8000 | 128000 | 138.9150 ms | 134.3961 ms | +3.36% | 43.87 -> 45.35 | + +The three FP32 partial-output buffers plus max/sum state require 141.72 MiB +per rank. Outputs are repeat-stable but not bitwise equal to unsplit because +the partition merge changes reduction order. Across the four shapes, maximum +absolute difference is `1.526e-5` to `3.052e-5`; mean absolute difference is +`7.51e-7` to `1.38e-6`. + +The corrected full-model A/B/A uses the same TP4 128K official-sampling +contract as the accepted K-stage endpoint. Control A/candidate/control B +prefill times are `47.5940/47.1074/47.6612 s`. The bracketed control is +`47.6276 s` or `2752.0 tok/s`; split-KV3 is `47.1074 s` or `2782.4 tok/s`, +giving `-1.092%` latency and `+1.104%` throughput. TTFT improves 1.078%. +Decode remains neutral: `5.1533 -> 5.1617 s` and +`49.483 -> 49.402 tok/s`. Every rank reports 384 split-KV3 kernel hits, which +matches 12 eligible long chunks times 16 full-attention layers times two +requests. KV capacity remains 2,424,439 tokens per rank and no OOM fallback +occurs. + +Quality does not yet justify a default-on promotion. Every official-sampling +lane is internally stable between warmup and measurement, and the candidate +is token-for-token identical to control B for both 256-token outputs. However, +identical unsplit control A selects a different stable sampled stream despite +the same official seed. The candidate therefore shows no unique output +divergence, but this cross-process sampling instability cannot prove +equivalence for a non-bitwise route. + +A follow-up deterministic greedy gate used three exact-128K natural-text +requests: two identical Chinese prompts and one Python task. Both control and +candidate reproduce the duplicate prompt exactly within their own process, +and all outputs are coherent; both Python outputs stop normally and return the +same correct `add(a, b)` implementation. The two routes are not token-for-token +equal, though. The Chinese outputs first differ at token 7 and align for +126/128 tokens; the Python outputs first differ at token 20, with all 67 +control tokens aligned inside the 70-token candidate stream. Candidate prefill +times are `47.1360/47.2292/47.3380 s` versus control +`47.6704/47.8330/47.9274 s`. This unbracketed quality run is directionally +consistent with the formal A/B/A but is not a replacement performance claim. +All four candidate ranks report exactly 576 split-KV3 hits, matching three +requests times 12 eligible chunks times 16 full-attention layers; controls +report zero. + +The semantic smoke passes stability and basic task correctness, but strict +greedy identity fails. Keep Q8000 split-KV3 explicit and default-off until a +fixed-text logprob/perplexity or dataset-level quality gate establishes that +the classified Type-B reduction-order drift does not reduce model quality. + +## 2026-08-25 Q8000/KV40K..128K GQA-packed Attention Architecture + +This experiment freezes the final long-prefill call at causal FP16 +`Q=8000`, `KV=128000`, `Hq=6`, `Hkv=1`, and `D=256` on one +V100-SXM2-32GB. Its metric is 6.094872576 useful causal TFLOPs divided by the +complete Attention elapsed time. It is neither whole-model TOPS nor prompt +tokens/s; the acceptance gate is at most 101.581 ms, or at least 60.0 useful +causal TFLOP/s. + +Tile and barrier tuning of the fused BM64/BN32 kernel was already exhausted: +the long-shape NCU record showed one CTA/SM, 12.5% occupancy, 62.14% of cycles +without an eligible warp, and only 9.06% DRAM throughput. The replacement is a +different scheduling and dataflow architecture: + +- split a fully visible 32K..120K prefix from the exact causal 8K tail; +- pack all six GQA query heads into GEMM M=48,000; +- run wide SM70 Tensor-Core QK GEMMs that emit FP16 scores plus row max/sum; +- transform scores into normalized probabilities during the PV + global-to-shared load, avoiding a materialized probability matrix; +- fold one reusable FP16 PV partial into a FP32 online prefix + numerator/max/sum state; +- export max/sum from the accepted exact 8K tail and merge the two online + softmax states. + +The retained score block is BN8192. Final frozen-extension Torch A/B/A for the +score-cache layout measures `140.00213 -> 100.76979 ms`, or +`43.53414 -> 60.48313` useful causal TFLOP/s: 28.0227% lower latency and a +1.38933x speedup. All 12,288,000 outputs are finite; max/mean absolute error +is `2.2888e-4 / 2.5430e-5` and relative L2 is `6.8629e-3` versus the exact +FP32-accumulator route. BN7168 reaches only 58.42057 TFLOP/s and BN8000 only +59.92825 TFLOP/s, so neither is promoted or rounded into a pass. + +The same bounded workspace was then generalized across the twelve observed +8K-chunk shapes. Each row is a strict A/B bracket against the active split-KV3 +control: + +| KV tokens | split-KV3 ms | architecture ms | latency change | useful TFLOP/s | +|---:|---:|---:|---:|---:| +| 40,000 | 38.41638 | 29.81530 | -22.3891% | 59.34862 | +| 48,000 | 47.25811 | 36.38477 | -23.0084% | 59.44005 | +| 56,000 | 56.53999 | 42.68715 | -24.5010% | 59.87584 | +| 64,000 | 65.17897 | 49.14552 | -24.5991% | 60.00842 | +| 72,000 | 74.00670 | 55.44960 | -25.0749% | 60.27745 | +| 80,000 | 82.74432 | 61.77792 | -25.3388% | 60.46783 | +| 88,000 | 91.98507 | 68.37111 | -25.6715% | 60.38797 | +| 96,000 | 100.88397 | 74.72862 | -25.9262% | 60.51241 | +| 104,000 | 109.80642 | 81.09448 | -26.1478% | 60.61108 | +| 112,000 | 118.42901 | 87.71959 | -25.9307% | 60.51602 | +| 120,000 | 127.35232 | 93.99757 | -26.1909% | 60.65749 | +| 128,000 | 135.93617 | 100.34193 | -26.1845% | 60.74103 | + +All 147,456,000 candidate output elements are finite. Across the family, the +worst max absolute difference is `5.0354e-4` and worst relative L2 is +`6.8983e-3` versus the exact FP32-accumulator route. The three shortest points +remain slightly below 60 because the 43-TFLOP/s exact causal tail is a larger +fixed fraction, but they are still 1.288x..1.325x faster than split-KV3. + +The architecture trace is compute-dense rather than launch-starved: 49 timed +kernels occupy 95.630493 ms of a 95.809850-ms GPU span, or 99.81%. Prefix QK +accounts for 54.22% and prefix PV 39.74%; the exact tail is 4.89%. The next +compute ceiling is therefore QK/PV instruction efficiency, not additional +host launch fusion. + +Real-model integration exposed a separate memory contract. Retaining all +partials or the entire workspace passes the operator gate but leaves too +little memory for the following 80-MiB TP all-reduce. The final layout retains +only the 750-MiB FP16 score matrix per device and packs about 100 MiB of other +state into one transient byte slab. Before first use it requires enough +driver-visible free memory for both allocations plus 128 MiB of downstream +headroom; otherwise it raises a typed OOM before touching the caching allocator +and falls back to the exact route. + +The route is default-on for its exact engine contract. The legacy +`VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL=0` setting is the +rollback, and the route rejects CUDA graph capture. It admits only `Q=8000`, +`KV=40000..128000` in 8000-token steps, `Hq=6/Hkv=1/D=256`, FP16, causal +attention, and scale 1/16. The earlier +endpoint-only `.875` TP4 +Qwen3.8-27B-FP8 control/candidate/control prefill times are +`46.45658 / 45.47797 / 46.11195 s`. Relative to the `46.28426-s` bracketed +control, the candidate lowers latency by 1.7421% and raises prompt throughput +from 2831.89 to 2882.10 token/s (1.7729%). Every candidate rank logs 16 +architecture hits, with no architecture OOM/fallback, and all three runs emit +the same 32 output token IDs and SHA256 +`df4fee7f5f0126fe6b391fe77b4fc19667831de5ef55fd69c28c2f52a3d7086e`. + +The endpoint-only `.88` run also succeeds at 45.53427-s prefill with the same +hash and 16 hits/rank. Its memory profiler left enough headroom, so the route +passed rather than exercising the fallback branch. + +The widened TP4 Qwen3.8-27B-FP8 gate brackets the candidate with `46.11195-s` +and `46.09222-s` controls. Relative to their `46.10208-s` mean, the +`41.51191-s` candidate lowers prefill latency by 9.9565% and raises prompt +throughput from 2843.08 to 3157.46 token/s (11.0575%). Every rank logs exactly +192 family-route hits with no architecture OOM/fallback. Both controls and the +candidate emit the same 32 token IDs and SHA256 +`df4fee7f5f0126fe6b391fe77b4fc19667831de5ef55fd69c28c2f52a3d7086e`. + +The 2026-08-26 merge audit rebuilt the four-patch vendored FA2 stack from its +locked clean baseline with CUDA 12.8 and `sm_70`. The production QK/PV kernels +again compiled at 254/119 registers with zero spill. On the rebuilt extension, +the 40K endpoint measured `39.29395 -> 30.27579 ms` (`1.29787x`) with all +outputs finite, max/mean absolute difference `5.4932e-4 / 4.4392e-5`, and +relative L2 `6.6098e-3`. The 128K endpoint measured +`136.07628 -> 100.20284 ms` (`1.35801x`) with all outputs finite, +max/mean absolute difference `2.2888e-4 / 2.5484e-5`, and relative L2 +`6.9166e-3`. Both points pass the FP16 merge envelope of max absolute error +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. + ## Artifacts +- D256 GQA architecture artifacts are retained outside Git; their private + filesystem location is intentionally not committed. +- Final operator A/B/A: + `results/torch-architecture-scorecache-final-v1-aba.json` under that task + root. +- Shape-family operator A/B: + `results/torch-architecture-shapefamily-v1-aba.json` under that task root. +- Shape-family TP4 model gate: + `results/tp4-128k-architecture-shapefamily-{candidate,control-after}-0875.json`, + bracketed with the endpoint gate's + `results/tp4-128k-architecture-scorecache-final-control-b-0875.json`. +- Final TP4 `.875` model gate: + `results/tp4-128k-architecture-scorecache-final-{control-a,candidate,control-b}-0875.json`. +- Final high-memory `.88` route-success run: + `results/tp4-128k-architecture-scorecache-final-preflight-fallback-088.json`; + despite the retained filename, the preflight passed and no fallback ran. +- K-stage task root: + `/data/minimax-h3/task-cache/qwen38-fp8-128k-flashattention-20260824`. +- Final clean FA2 binary: + `build/formal-v6-final-cmake/_vllm_fa2_C.abi3.so` under that task root. +- Endpoint A/B/A JSON: + `results/tp4-128k-{baseline-a-v6,candidate-v6,baseline-b-v6}.json`. +- Q8000 operator sweep: + `results/splitkv3-q8000-kv{40000,64000,96000,128000}-v6.json`. +- Q8000 endpoint A/B/A: + `results/tp4-128k-splitkv3-q8000-{control-a,candidate,control-b}.json`. +- Q8000 deterministic quality gate: + `results/tp4-128k-splitkv3-q8000-quality-{control,candidate}.json` and the + matching files under `logs/`. +- Shape diagnostic and the zero-hit Q8192 negative run: + `logs/tp4-splitkv3-shape-diag.log` and + `results/tp4-128k-splitkv3-{control-a,candidate,control-b}.json`. - Root: task-local `qwen38-fp8-prefill-decay-20260815` artifact directory. - Nsight Systems: `profiles/qwen38-fp8-tp4-i128k-chunk15680-r2.nsys-rep`. diff --git a/docs/design/sm70_qwen38_nvfp4_decode.md b/docs/design/sm70_qwen38_nvfp4_decode.md new file mode 100644 index 0000000000..e937564413 --- /dev/null +++ b/docs/design/sm70_qwen38_nvfp4_decode.md @@ -0,0 +1,711 @@ +# SM70 Qwen3.8-27B-NVFP4 Decode Recovery + +Date: 2026-08-23; acceptance updated 2026-08-24 + +## Decision + +The accepted Qwen3.8-27B-NVFP4 TP4, no-MTP decode route keeps every eligible +FP8 and NVFP4 projection on an SM70 native path. Channel-FP8 projections use +the memory-neutral QPN8 decode kernels for the accepted gate/up, down, and +output shapes. Remaining channel-FP8 projections, including the LM head, use +TurboMind W8A16. NVFP4 projections use TurboMind W4A16. No accepted result in +this document credits a Marlin fallback. + +The recovery route also adds E4M3 KV support to Flash-V100 XQA and selects a +p64 partition for the exact batch-one, G6/D256 layout. This is the change that +makes the frozen native-sampler result exceed 70 tok/s. The p64 route is +restricted to E4M3, batch one, G6/D256; the existing E5M2 and FP16 policies +are unchanged. The later no-MTP 80 tok/s acceptance adds exact QPN4, TP4 +all-reduce, QPN8 specialization, and an exact-Philox chunked sampler sidecar; +that acceptance is recorded at the end of this document. + +## Frozen Contract + +- Model: Qwen3.8-27B-NVFP4, config SHA256 + `1b3c71868d1299e52df6fc907deb202d5132b1ef0f72aae0ef6d15185dd53a5c`. +- Hardware/runtime: four V100-SXM2-32GB GPUs, TP4, Python 3.12, Torch + 2.10.0+cu128, CUDA 12.8. +- Engine: compressed-tensors, FP16 activation/output, E4M3 KV, + `FLASH_ATTN_V100`, `max_model_len=262144`, `max_num_batched_tokens=8192`, + `max_num_seqs=1`, prefix caching, aligned Mamba cache, chunked prefill, and + `FULL_AND_PIECEWISE` CUDA graphs. +- Request: input 1024, natural output cap 256, no `ignore_eos`, no MTP. +- Sampling: temperature 1.0, top-p 0.95, top-k 20, request seed 20260815, + engine seed 0. The 2026-08-23 recovery baseline uses the unmodified native + sampler. The 2026-08-24 acceptance uses the separately registered, + exact-Philox chunked sampler described below. +- Pure decode uses the 255 steady token intervals and excludes TTFT/prefill. + +## Route Map + +| Layer family | TP-local shape | Accepted decode route | +|---|---|---| +| FP8 gate/up | K5120 x N8704 | fused QPN8 gate/SiLU/up | +| FP8 down | K4352 x N5120 | QPN8 split-16 | +| FP8 output | K1536 x N5120 | QPN8 split-12 | +| FP8 GDN input/full-attention QKV and LM head | model shapes | TurboMind W8A16 | +| NVFP4 gate/up | K5120 x N8704 | TurboMind N32, lookahead 1 | +| NVFP4 down | K4352 x N5120 | TurboMind N32, lookahead 2 | +| E4M3 G6/D256 attention | B1, page 1568 | Flash-V100 XQA p64 | + +Channelwise `[N,1]` FP8 scales are admitted and packed without retaining a +second permanent weight copy. The existing QPN8 model/shape/concurrency gates +remain in force. The NVFP4 selector is exact-shape gated by +`VLLM_SM70_NVFP4_QWEN38_TP4_M1_FAST_SELECTOR`; setting it to zero restores +dynamic tuning. Setting `VLLM_FLASH_V100_DECODE_PARTITION_SIZE=128` restores +the previous E4M3 attention partition. + +## Pure Decode Result + +| Milestone | Steady decode | TPOT | +|---|---:|---:| +| Original mixed checkpoint route | 29.35 tok/s | about 34.07 ms | +| Channel-FP8 TurboMind | 57.285 tok/s | 17.457 ms | +| E4M3 XQA p256 | 62.615 tok/s | 15.971 ms | +| Exact channel-FP8 QPN8 | 65.522 tok/s | 15.262 ms | +| E4M3 XQA p128 | 67.880 tok/s | 14.732 ms | +| NVFP4 N32 selectors | 69.991 tok/s | 14.288 ms | +| Clean final binaries, native sampler, E4M3 p128 | 69.904 tok/s | 14.305 ms | +| Clean pre-merge binaries, native sampler, E4M3 p64 | 71.732 tok/s | 13.941 ms | +| Merged-source confirmation, native sampler, E4M3 p64 | **71.342 tok/s** | **14.017 ms** | + +Both p64 requests generated all 256 tokens, contained no EOS, and finished by +length. Relative to the immediately preceding clean-p128 control, the +pre-merge p64 run saves 0.365 ms/token and improves pure decode by 2.61%. +After merging `onecat/main` at `675a12dedc`, rebuilding Flash-V100, and +repeating the frozen request under an exclusive four-GPU reservation, p64 +measures 71.342 tok/s at 14.017 ms/token. This is the conservative 2026-08-23 +baseline and is superseded by the 2026-08-24 acceptance below. + +Sampled token identity is not an attention quality criterion: one-output-ULP +changes can flip a low-margin random sample. The final p64 stream is coherent +and, independently, reproduces an earlier accepted full-length stream. The +operator gate below is the numerical criterion. + +## Numerical and Operator Gates + +The E4M3 XQA p64/p128 race uses page size 1568, G6/D256, checkpoint-style +E4M3 bytes, K/V scales 0.04/0.25, and the scalar paged decoder as reference. + +| Sequence length | p64 | p128 | p64 gain | Maximum error | +|---:|---:|---:|---:|---:| +| 1025 | 46.633 us | 63.614 us | 16.981 us | 7.63e-6 | +| 1152 | 45.896 us | 59.346 us | 13.450 us | 7.63e-6 | +| 1280 | 46.072 us | 55.625 us | 9.553 us | 7.63e-6 | +| 2049 | 45.807 us | 56.751 us | 10.945 us | 7.63e-6 | + +Every p64 and p128 output is within one representable FP16 output ULP of the +scalar reference. The focused GPU regression covers p64/p128/p256 with unit +and non-unit K/V scales and passes 6/6. + +The retained NVFP4 selectors also passed independent FP32-oracle checks: +relative L2 is `2.701e-4` for gate/up and `2.924e-4` for down, with cosine +approximately one. The QPN8 source already passed its model-quality and +operator gates recorded in `sm70_qwen38_qpn8_decode.md`. + +## Per-Token Profile + +The latest short Nsight Systems node trace uses the merged p64 route. It is +composition evidence; its 14.635 ms traced TPOT is not used as the absolute +speed result because the unprofiled frozen request measures 14.017 ms. The +trace captures 63 decode replays on each TP rank and uses the 61 middle steps +for steady statistics. Graph-node kernel coverage is 93.48%. + +| Component | Time/token | +|---|---:| +| TurboMind NVFP4 gate/up | 2.762 ms | +| QPN8 split-16 projections | 2.234 ms | +| TurboMind NVFP4 down | 1.706 ms | +| TP all-reduce | 1.697 ms | +| QPN8 split-12 projections | 0.895 ms | +| E4M3 XQA p64 | 0.587 ms | +| QPN8 fused gate/up | 0.508 ms | +| TurboMind FP8 dense/LM head | 0.407 ms | + +Across rank-token samples, the replay interval is 14.621 ms and GPU activity +union is 14.055 ms, or 96.136% of the interval. Idle gaps total 0.565 ms/token. +The stream still launches 1140.8 kernels per rank per token; half are shorter +than 5 us. The grid-limited static occupancy ceiling places 69.40% of service +below 25% occupancy and 28.50% at 25-50%. Nearly continuous GPU work therefore +does not imply high useful compute utilization: batch-one launch geometry and +the serial projection/communication chain remain the main limit. + +Fifty-millisecond NVML samples report 99.47% mean GPU busy-window duty and +48.26% memory-active-window duty. Per-GPU power averages 175.9-180.2 W and +peaks at 221.4 W; runtime allocation peaks at 29.06-29.36 GiB/GPU. Model loading +accounts for 5.77 GiB/GPU, while the configured cache budget reports 21.29 +GiB/GPU; the remaining allocation includes CUDA graphs, workspaces, state, and +runtime context. + +NVML memory duty is not achieved HBM bandwidth. Current trace durations imply +effective minimum packed-weight rates of 451.2 GB/s for NVFP4 gate/up, 365.1 +GB/s for NVFP4 down, and 717.4 GB/s for QPN8 down, with useful arithmetic rates +of 1.805, 1.461, and 1.435 TFLOP/s/GPU respectively. These omit scales, +activation/output traffic, caches, and implementation-internal work. A current +NCU capture was attempted, but the host rejected non-root performance-counter +access with `ERR_NVGPUCTRPERM`; exact current SM, Tensor Core, occupancy, and +DRAM counters therefore require administrative profiling permission. The +previously accepted QPN8 counter evidence remains in +`sm70_qwen38_qpn8_decode.md`, but it is not credited as current NVFP4 or p64 +XQA counter evidence. + +## Rejected Experiments + +- NVFP4 N64 and K32 candidates lost to N32 and were removed. The retained + exact shapes use split 3/swizzle 4 with shape-specific lookahead. +- QPN8 split 17/20/24 lost. Split 12 won only for K1536 x N5120, saving about + 12.7 us/token across the 64 output projections. +- The first compact top-k20 Python sampler and first fused CUDA top20 + candidate did not preserve native Philox sampling and were removed. Both + 2026-08-23 p64 results use the native sampler. They are distinct from the + later canonicalized, generator-state-preserving implementation accepted on + 2026-08-24. +- Replacing only the full-vocabulary sort while retaining both native + full-vocabulary softmax operations saved 53.9 us in isolation, insufficient + for a stable acceptance margin. A one-full-softmax hybrid stayed + experimental because the p64 attention route solved the target without + changing sampling math. + +## Build, Tests, and Evidence + +- Final `_C` SHA256: + `e0ea14d0e40330b08a9951e67634e50b722540d5d1bbb48500f690323ab07624`. +- Final p64 Flash-V100 SHA256: + `b418fed86b9c1ab9297c8795c24732818239b9a3aaca5ec9efb60933853d8ce7`. +- Focused CPU policy/dispatch tests: 12/12 passed (nine FP8/QPN8 and three + E4M3 p64 policy cases). +- Focused E4M3 XQA GPU numerical tests: 6/6 passed. +- Ruff lint/format, Python byte compilation, shell syntax, and + `git diff --check` pass for the changed files. +- Retained task evidence is under + `.artifacts/qwen38_nvfp4_speed_20260823/`, notably the final p128/p64 JSON + results, the merged confirmation + `final_merged_qwen38_nvfp4_tm_e4m3_xqa_p64_native_sampler_full_graph_i1k_o256.json`, + `e4m3_xqa_p64_vs_p128_clean.json`, and the parsed per-token Nsight tables + under `profiles/`. + +## 2026-08-24 No-MTP 80 tok/s Acceptance + +### Accepted deployment and route + +The accepted deployment is deliberately a three-binary composition. It uses +the compatible primary `_C` module with SHA256 +`a0a0cd9ddeccc73fa3d920c7a869450c4b33d001f97637c35b75b966d89ad36d`, +the production CMake C++17 sampler sidecar with SHA256 +`cdbfdd87dfa9119e52acc88d5787063202561a879a63334145a5616344a549ae`, +and Flash-V100 p64 with SHA256 +`b418fed86b9c1ab9297c8795c24732818239b9a3aaca5ec9efb60933853d8ce7`. +MTP remains disabled. + +The exact FP8 gate/up, down, and output shapes stay on the proven-faster QPN8 +route. GDN input, full-attention QKV, and the LM head stay on TurboMind W8A16. +The accepted batch-one NVFP4 shapes use the native QPN4 decode route; unsupported +or non-admitted cases retain the TurboMind fallback. Thus every eligible FP8 +sublayer still uses TurboMind unless the frozen shape has a measured faster +QPN8 specialization. + +The sampler sidecar replaces a full-vocabulary sampling fragment with an +80-chunk top-20 reduction while preserving canonical sort order, top-p math, +Philox draws, and generator state. It is built separately so sampler operator +registration does not relink the quality- and speed-frozen primary `_C`. +`setup.py` extracts either CPython-SOABI or abi3 sidecar names, and +`vllm/_sm70_ops.py` loads the bundled module or the explicitly configured +`VLLM_SM70_SAMPLER_LIBRARY`. + +### Speed result + +All absolute results use the frozen TP4, input-1024/output-256 contract and +255 steady decode intervals. TTFT and prefill are excluded. + +| Binary/measurement | Steady decode | TPOT | Disposition | +|---|---:|---:|---| +| Merged native-sampler baseline | 71.342 tok/s | 14.017 ms | baseline | +| Hand sidecar repeat A2 | 80.177 tok/s | 12.472 ms | accepted | +| Hand sidecar repeat A3 | 80.164 tok/s | 12.474 ms | accepted | +| Hand sidecar repeat A4, physical GPUs 4-7 | 80.026 tok/s | 12.496 ms | accepted | +| Production CMake C++17 sidecar A2 | **80.624 tok/s** | **12.403 ms** | accepted | + +The hand-sidecar accepted repeats preserve the frozen 256-token SHA256 +`0b2d335ddce9b282e45eea1b6c86525bc61eeb5ba1655e8228e6ef3bd1ce823b`. +The production C++17 sidecar changes a low-margin full-model random trajectory, +so random token identity is not treated as the sole quality gate. A direct +cross-build sampler diagnostic covers 100 independent seeds and 256 sequential +draws: hand and production sidecars produce identical token selections and +identical final generator-state SHA256 +`4257d205138503840fea92fe6b15ddfa276df82c315513da4bbd785393c98c96`. + +### Output-quality gate + +The quality contract is the frozen first 250 GSM8K test questions, five-shot +prompting, greedy generation, maximum 256 output tokens, and per-question +record retention. It is intentionally independent of the random performance +prompt. + +| Route | Correct | Accuracy | Invalid | +|---|---:|---:|---:| +| Frozen pre-optimization baseline | 226/250 | 90.4% | 0 | +| Hand-sidecar candidate | 227/250 | 90.8% | 0 | +| Production CMake C++17 sidecar | **226/250** | **90.4%** | **0** | + +The production result therefore equals the frozen accuracy baseline with no +invalid outputs. Its JSON contains all questions, generated texts, extracted +answers, labels, correctness flags, token IDs, and per-item output hashes. The +result file SHA256 is +`143b2b345b830a858bff2568f9cf51148c8185acb42bb129e2dfc2d421a196d2`. + +### Accepted-route trace and resource use + +The latest Nsight Systems trace uses the accepted compatible main library and +the behavior-equivalent hand sidecar. It captures 63 graph replays on each of +four ranks and reports the middle 61 steps. Its 13.465 ms node-traced TPOT is +composition evidence; the unprofiled 12.474 ms result remains the absolute +speed evidence. Graph-node kernel coverage is 94.88%. + +| Critical component | GPU service/token | +|---|---:| +| QPN8 split-16 projections | 2.213 ms | +| QPN4 fused gate/up | 2.165 ms | +| TP4 pack32 all-reduce | 1.587 ms | +| QPN4 down | 1.441 ms | +| QPN8 output projection | 0.915 ms | +| Flash-V100 E4M3 XQA p64 | 0.618 ms | +| QPN8 fused gate/up | 0.496 ms | +| TurboMind FP8 dense/LM head | 0.410 ms | + +The mean replay interval is 13.430 ms and GPU activity union is 12.887 ms, +or 95.958%; idle gaps total 0.543 ms/token and the largest mean gap is only +9.538 us. The route still launches 1061.9 kernels per rank per token. The +grid-limited occupancy ceiling assigns 38.46% of service below 25% occupancy, +47.86% at 25-50%, and 13.63% at 50-75%, showing that batch-one kernel geometry +and the serial projection/communication chain remain the main headroom. + +NVML's coarse windows report 100% GPU-busy duty, 53.25% memory-active duty, +and 56.36% of the 300 W power limit. Runtime memory is about 28.97 GiB/GPU; +SM and memory clocks are 1530 and 877 MHz. These are duty indicators, not +achieved SM, Tensor Core, or HBM throughput. Nsight Compute counters remain +unavailable because the driver returns `ERR_NVGPUCTRPERM`. A payload-only +model estimate is about 491 GB/s/GPU and useful arithmetic is about 4.33 +TFLOP/s across TP4, but neither value is a hardware-counter measurement. + +### Reproducibility boundary and evidence + +The production sidecar is reproducible through the CMake target +`_sm70_sampler_C` and its focused microbenchmark selects exactly the reference +tokens across five distributions, 100 explicit-seed trials, and 100 default +generator trials. It reduces the measured fragment from 102.427 to 62.972 us. +The final Python regression run passes all 94 tests in the sampler, NVFP4 +admission, and SM70 TurboMind adapter files. Ruff lint/format, Python byte +compilation, the sidecar wheel-name gate, and `git diff --check` also pass. + +Freshly relinking the primary `_C` was tested separately. The first cold run +measured 77.314 tok/s, while the matching second run measured 80.181 tok/s at +12.472 ms/token. Both used a current-source C++17 sampler sidecar and produced +the exact same 256 token IDs as the accepted production C++17 run. The earlier +roughly 77-80 tok/s spread was therefore startup-state variation, not a decode +route or output regression. + +The clean source snapshot is `7bd1d783e4861af5f5396a721156c44f623f88b0`. +Its primary `_C` SHA256 is +`c5d7c6a9a0fa714c1e97427c3d16369404ae5af6ee80f48b165595081cc43e56`, +and its C++17 sampler SHA256 is +`3058051f59949d88f5982d0c5bd80121f3c75a0f5307ca901fa6aebdec4cbc74`. +The primary and sampler `.nv_fatbin` SHA256 values are respectively +`dbba740219f00db9a528310fa84c0defb97c03f5cf99b519af5975609eb926e6` +and +`9a89e8cc22d1f8e36378005c08e2b32b315d79728989d6b4b2fe44fedfa5cb2c`; +both are byte-identical to the accepted binaries. The full primary ELF differs +because merged main adds two unrelated ModelOpt NVFP4 MoE host entry points, +and full sidecar hashes include host build-ID differences. Those differences +do not enter this dense batch-one graph. + +The clean second-run result SHA256 is +`4e6ed3306772efca1a9c84b3681d50a268c644b8d6327df0fab2341581236bb1`. +Its serialized token-array SHA256 is +`f50e0ebab44ae350c7921bf91ba709fa896dcfb253c06842323b80de5b8bde32`, +exactly matching the accepted C++17 result. It records TP4, physical GPUs 0-3, +custom all-reduce enabled, both FP8 and NVFP4 TurboMind enabled, +`speculative_config=None`, and `spec_decoding_metrics=null`. Because both +device binaries and the full sampled stream are identical, the existing +production C++17 GSM8K result, 226/250 with zero invalid outputs, remains the +quality gate; repeating the same 250-question run would add no new path +coverage. Keeping the sampler in its own sidecar remains the packaging +boundary, but a current-source primary rebuild is now accepted. + +CMake component installation was also checked without rebuilding unrelated +wheel targets. It installs `vllm/_C.abi3.so` and +`vllm/_sm70_sampler_C.cpython-312-x86_64-linux-gnu.so`, removes their build-tree +RPATHs, and preserves both accepted fatbin hashes. The installed-file SHA256 +values are `d7ed490ef41e55f3858a45b3a92d62b2d527343fa14b04053cbe65ba16646749` +and `8ee8466359dc4432234dc3daa92918b2226879caad563d23a85b980c2a7c755e`. +A loader probe against the staged files registered the QPN4, QPN8, and +`_C::sm70_sample_packed_top20_out` schemas successfully. A broad wheel rebuild +of unrelated extensions is not required to establish this route's packaging +boundary. + +Primary evidence paths are: + +- `.artifacts/qwen38_nvfp4_speed_20260823/results/candidate_qpn4_fold_qpn8_m1_arpack32_official_cxx17_sidecar_chunk80_i1k_o256_a2.json` +- `.artifacts/qwen38_nvfp4_speed_20260823/results/candidate_qwen38_nvfp4_clean_head_cxx17_chunk80_i1k_o256_a2.json` +- `.artifacts/qwen38_nvfp4_speed_20260823/results/qwen38_nvfp4_gsm8k_250_official_cxx17_sidecar.json` +- `.artifacts/qwen38_nvfp4_speed_20260823/results/sampler_stream_hand.json` +- `.artifacts/qwen38_nvfp4_speed_20260823/results/sampler_stream_cxx17.json` +- `.artifacts/qwen38_nvfp4_speed_20260823/wheel_stage_clean_head_7bd1d783/` +- `.artifacts/qwen38_nvfp4_speed_20260823/profiles/candidate_qwen38_nvfp4_sidecar_chunk80_nsys_nvml_i1k_o64_per_token.md` +- `.artifacts/qwen38_nvfp4_speed_20260823/profiles/candidate_qwen38_nvfp4_sidecar_chunk80_resource_summary.md` + +## Quality-safe GDN decode fusion + +The next no-MTP batch-one optimization keeps the mixed checkpoint route map +unchanged and removes work at the GDN boundary. The exact Qwen3.8 TP4 decode +shape can issue its channel-FP8 QKV/Z projection and FP16 b/a projection from +one QPN8 launch, writing the four destination tensors directly. A separate +one-pass Triton kernel performs the 12-by-128 gated RMSNorm without changing +the accepted arithmetic. The accepted pair is opt-in through both +`VLLM_SM70_GDN_QPN8_BA_SPLIT=1` and +`VLLM_SM70_GDN_RMSNORM_ONEPASS=1`; the split projection cannot activate +without its paired RMSNorm gate. Other operator shapes and prefill retain the +existing operators. These are hardware/layout/operator gates; the measured +checkpoint identifies the evidence workload and does not select the route. +The archived JSON predates this audit naming cleanup and records the former +Qwen-prefixed flag names; kernel arithmetic is unchanged. The final source +also rechecks the loaded QPN8 code/scales, workspace pointer, bias, dtype, +contiguity, and same-device contracts before dispatch. + +The matched endpoint used the same four V100-SXM2-32GB GPUs, TP4, input 1024, +output cap 256, official random sampling, E4M3 KV, Flash-V100, full CUDA graph, +and no MTP. Each result is the stable median of three sequential requests and +pure decode excludes TTFT/prefill. + +| Route | Pure TPOT | Steady decode | Output gate | +|---|---:|---:|---| +| Current control | 12.249652 ms | 81.634970 tok/s | frozen 256-token stream | +| QPN8 QKV/Z plus b/a only | 11.941616 ms | 83.740763 tok/s | reject: first divergence at token 202 | +| QPN8 QKV/Z plus b/a and one-pass RMSNorm | **11.921648 ms** | **83.881019 tok/s** | pass: 3/3 exact | + +The accepted pair improves steady decode by 2.751% and TPOT by 2.678% versus +the matched control. All three accepted requests reproduce the control's full +256-token SHA256 +`8b37337f4c393711cb8550db6bae909b1e85de8df1cf5ba8c60d8c000749c0a2`. +The one-pass RMSNorm 48-layer microbenchmark is bitwise exact and saves +55.946 us/token. The QKV projection is exact on all layers; b/a differs only +at normal FP32-reduction roundoff (`6.368e-8` relative L2, maximum absolute +error `2.384e-7`) and the paired endpoint restores exact sampled output. + +The post-rebase admission recheck uses the final source-built extension and +the default Qwen3.5 model wrapper rather than a diagnostic GDN boundary. The +C++ split-route oracle fired on all four TP ranks. Its matching latest-main +control/candidate stable medians are `12.261727/11.956949 ms` and +`81.554583/83.633383 tok/s`, a 2.549% decode-throughput improvement. All +three candidate requests exactly match all three latest-main control +256-token streams with SHA256 +`ca77db3b032a1600a8567adea706108c1bd8c5472b3ace2318c41ab17c66c1f9`. + +A fused Q/K RMSNorm, MRoPE, and KV-write experiment reached 85.451 and +85.015 tok/s on its two stable repeats but is rejected: it diverged from the +control at token 1 and stopped after 208 output tokens. That source and its +route flag are intentionally absent from the submitted change. + +Primary local evidence paths are: + +- `.artifacts/qwen38_nvfp4_speed_20260823/results/gdn_endpoint_control_r3_i1k_o256.json` +- `.artifacts/qwen38_nvfp4_speed_20260823/results/gdn_endpoint_ba_split_rms_r3_i1k_o256.json` +- `.artifacts/qwen38_nvfp4_speed_20260823/results/current_main_gdn_control_r3_i1k_o256.json` +- `.artifacts/qwen38_nvfp4_speed_20260823/results/current_main_gdn_route_recheck_r3_i1k_o256.json` +- `.artifacts/qwen38_nvfp4_speed_20260823/results/gdn_rmsnorm_production_op_bound_48layers.json` +- `.artifacts/qwen38_nvfp4_speed_20260823/results/qpn8_fused_ba_split_bound_48layers.json` +- `.artifacts/qwen38_nvfp4_speed_20260823/results/qknorm_mrope_cache_gdn_ba_rms_r3_i1k_o256.json` + +## 2026-08-25 Long-context E4M3 page800 decode + +### Scope and route + +This optimization targets the exact SM70, E4M3, grouped-query decode route +used when aligned Mamba state makes the attention page 800 tokens. The +production KV allocation is interleaved as `[blocks, 2, 800, 1, 256]`; its K +and V views each have strides `[409600, 256, 256, 1]`. Admission additionally +requires G6/D256 and the existing dual-CTA batch route. Other page sizes, +layouts, head dimensions, KV dtypes, batch routes, and prefill keep the prior +implementation. + +The specialized kernel uses the compile-time 800-token page and physical +block stride, replaces repeated general page division with the bounded +one/two-page comparison for a 256-token partition, and loads adjacent shared +V values as aligned `half2`. Each output dimension retains the same token and +FMA order as the fallback. The strict layout predicate is checked before +dispatch; setting `VLLM_FLASH_V100_E4M3_PAGE800_FASTPATH=0` restores the +general-stride kernel. The fast path is enabled by default on an admitted +route. `VLLM_FLASH_V100_E4M3_PAGE800_FASTPATH_TRACE=1` emits a once-per-process +route diagnostic. + +The measured contract uses Qwen3.8-27B-NVFP4, four V100-SXM2-32GB GPUs with +TP4, compressed-tensors NVFP4 weights, FP16 activations, E4M3 KV, Flash-V100, +FlashQLA, prefix caching, aligned Mamba cache, FULL decode graphs, and no MTP. +The workload is NVIDIA SPEED-Bench revision +`487aa718444e816458d1a0a52bfce7a454285cf4`: the official 32K `low_entropy` +JSONL has SHA256 +`c3ee8a3f63b6cce18d063a2ff9992b6b96cb72c2b6408970fa78d2335b542f8e`. +The explicitly derived 64K workload pairs two official rows and has SHA256 +`0dcc877ecc7ccf66450c1f168ff7493dcf50f52da89d321dba394144a70e9800`. +Performance sampling is temperature 1.0, top-k 20, top-p 0.95, request seed +20260822, natural EOS, and a 512-token cap. Pure decode excludes queueing, +TTFT, and prefill; each cell has one warmup and three measured repeats. + +### Long-context result + +The matched rollback/candidate A/B was measured on source `05d5aa4e57`. B1 +references are 52.052 token/s at 32K and 40.152 token/s at 64K. The accepted +16K scaling references are 70.736% at B8 and 52.281% at B16. + +| Context | Batch | Rollback token/s | Candidate token/s | Candidate TPOT | Gain | Efficiency vs B1 | +|---:|---:|---:|---:|---:|---:|---:| +| 32K | 8 | 292.03 | **326.47** | 24.505 ms | **11.79%** | 78.40% | +| 32K | 16 | 392.84 | **451.92** | 35.405 ms | **15.04%** | 54.26% | +| 64K | 8 | 209.94 | **244.51** | 32.719 ms | **16.47%** | 76.12% | +| 64K | 16 | 260.36 | **311.74** | 51.325 ms | **19.74%** | 48.52% | + +B8's two-context geometric-mean gain is 14.11%, B16's is 17.36%, and the +four-cell geometric mean is 15.72%; no cell regresses. Relative to the 16K +references, 64K loses no efficiency at B8 and 3.76 percentage points at B16, +within the ten-point budget. Every one of the 144 candidate performance +requests completed with its full 512 tokens and without empty or replacement- +character output. Route logs show FULL graphs and the expected long-context +`dual_cta_split` B8 and `dual_cta` B16 dispatches on all four ranks. + +The endpoint reports that the exact D256 prefill extension is unavailable and +uses its established fallback. That warning is common to rollback and +candidate and is excluded from the pure-decode metric. This result must not be +cited as a prefill improvement. + +### Counter and correctness evidence + +An Nsight Compute rollback/candidate comparison on B16, 64K, p256 attributes +the endpoint gain to less address/control and shared-memory pressure rather +than added occupancy. Kernel duration falls 22.47%, executed instructions +11.92%, shared-bank conflicts 48.16%, MIO-throttle stalls 95.92%, long- +scoreboard stalls 8.61%, and wait stalls 11.06%. Eligible warps per active +cycle rise 0.565 to 0.699. Peak-relative DRAM and SM utilization rise from +24.49% to 28.11% and from 39.11% to 44.79%, respectively; registers rise only +from 166 to 168 and dynamic shared memory is unchanged. Thus the retained +change follows the measured copy/addressing bottleneck and does not depend on +an occupancy increase. + +The production-layout operator A/B is bitwise equal to the rollback path for +B8/B16 and p64/p128/p256. At p256, its direct kernel gains are 22.10% at +32K/B8, 21.37% at 32K/B16, 22.59% at 64K/B8, and 22.84% at 64K/B16. The two +focused GPU suites pass 15/15 and 14/14; the production interleaved cases use +`torch.equal`, not a tolerance. All 55 environment tests, Ruff lint/format, +the SM70 extension build, and `git diff --check` also pass. + +Because the production path uses official stochastic sampling, independent +server processes are not expected to reproduce every generated text. The +quality gate therefore pairs rollback and candidate over the same 32 GSM8K +main/test prompts at B8 and B16, temperature 1.0/top-k 20/top-p 0.95, natural +EOS, and sampling seeds 20260822, 20260823, and 20260824. This protocol was +fixed before the two additional seeds ran; both sides were rerun rather than +repeating only the candidate. + +| Batch | Rollback numeric EM | Candidate numeric EM | Delta | +|---:|---:|---:|---:| +| 8 | 71/96 | **74/96** | +3 | +| 16 | 74/96 | **74/96** | 0 | + +All 384 quality requests succeed, and the candidate has no empty or Unicode +replacement output. The first seed alone was one item lower in each batch; +the predeclared three-seed aggregate is recorded to make that sampling +variance visible rather than discarding it. The final source patch rebases +without conflict onto `onecat/main@d62ef5cb20`. Its rebuilt Flash-V100 ELF has +SHA256 `70816f9655e36fcd04b07d75be1b3abafabc81a30c261d4fd99b06389cbf82a2`. +The full SM70 SASS dump is byte-identical to the measured binary, with SHA256 +`b30031fa4f29181b8e1c757bc408c608bae97df380893e613b96a2cf7d3e9a51`; +the differing full-ELF hash is build metadata, not device code. + +## 2026-08-25 No-MTP Long-Context Decode Acceptance + +### Contract and route + +The long-context gate keeps the accepted Qwen3.8-27B-NVFP4 computation and +sampling paths unchanged. It uses four V100-SXM2-32GB GPUs with TP4, +compressed-tensors weights, FP16 activations/output, E4M3 KV, TurboMind for +the admitted FP8/NVFP4 projections, Flash-V100 attention, full CUDA Graph, +prefix caching, aligned Mamba state, and no MTP. Each request ends at an exact +128K or 256K context after 64 generated tokens. Pure decode is the 63 steady +token intervals and excludes prefill and TTFT. The acceptance thresholds are +60 tok/s at 128K and 45 tok/s at 256K. + +The new Flash-V100 route is restricted to batch one, G6/D256, 1,568-token +interleaved paged KV, and E4M3. It selects p64 below 12,288 tokens, p256 below +49,152, p512 below 98,304, p896 below 196,608, and p1664 thereafter. A single +long-wave partition kernel and matching reducer select p512/p896/p1664 from +device-resident sequence lengths, so CUDA Graph replay does not need a host +readback or three long inactive launch pairs. The existing short p64/p256 +launch remains separate. Other batch sizes, KV formats, head layouts, page +sizes, explicit partition overrides, and attention routes are unchanged. + +Full CUDA Graph capture has two semantic batch-one variants at the p512 +boundary. Contexts below 49,152 tokens replay the default graph, which contains +only the p64/p256 launch and reducer pair. Contexts at or above 49,152 tokens +replay the long-context graph, where `batch_context_routing` admits the wave +launch and reducer pair. The dispatcher uses its existing host scheduler +context length to select an already-captured graph; the attention kernel still +reads the live sequence length from device memory. This keeps the 1K decode +graph free of inactive long-wave launches without adding a device-to-host +sequence-length readback. + +E4M3 conversion uses a 256-entry exact half-bit lookup table initialized in +CTA shared memory. It adds 512 bytes of dynamic shared memory and avoids the +long integer conversion chain on every KV element. The final SM70 cubin uses +168 registers/thread and 32 bytes of stack for the 192-thread long-wave +partition kernel; total dynamic shared memory is about 44.8 KiB, retaining two +CTAs/SM. Its reducer uses 32 registers/thread and 80 bytes of static shared +memory. The route is enabled with +`VLLM_FLASH_V100_XQA_E4M3_G6_P64_P256_AUTO=1`, +`VLLM_FLASH_V100_XQA_E4M3_G6_WAVE_PARTITIONS=1`, and +`VLLM_FLASH_V100_XQA_E4M3_G6_MERGED_WAVE_LAUNCH=1`. Setting the wave flag to +zero restores the prior p64/p256 policy. Setting the merged-launch flag to zero +retains the wave thresholds but restores the separate long-wave launches. +Setting the auto flag to zero restores the original explicit-partition route. + +### Exact endpoint result + +The control and candidate use the same model, projections, sampler, TP/GPU +set, max length, input/output lengths, attention backend, KV type, graph mode, +and no-MTP state. Only the Flash-V100 E4M3 attention partition policy changes. + +| Final context | Route | Prefill | TTFT | Pure TPOT | Steady decode | Target | +|---:|---|---:|---:|---:|---:|---:| +| 128K | p64/p256 control | 283.296 s | 283.327 s | 24.654 ms | 40.561 tok/s | 60 tok/s | +| 128K | LUT merged-wave | 283.127 s | 283.154 s | **16.172 ms** | **61.834 tok/s** | pass, +3.06% | +| 256K | p64/p256 control | 1085.628 s | 1085.673 s | 36.422 ms | 27.456 tok/s | 45 tok/s | +| 256K | LUT merged-wave | 1085.519 s | 1085.568 s | **19.851 ms** | **50.376 tok/s** | pass, +11.95% | + +At 128K, throughput improves 52.45% and TPOT falls 34.40%. At 256K, +throughput improves 83.48% and TPOT falls 45.50%. Both requests emitted all +64 tokens, finished by length, and reported no corruption. The complete token +streams are exactly unchanged from control: 128K SHA256 is +`66a130ea9d68faaf40f7c7cfb942c579f759b9535394cce958699bd620df2d6f`; +256K SHA256 is +`e1bc9ab8a979e43cda22348c0ff0056b5a7c3212ba57e830b731f3cf51906524`. +This is an attention-only optimization and does not credit a sampler change. + +The post-integration 1K regression guard keeps the accepted TurboMind +projection routes, FlashQLA GDN BA-split route, official sampler, disabled +compile cache, and 256 output tokens. After excluding the first cold request, +the accepted GDN BA-split control is 83.761/83.721 tok/s and the +graph-specialized candidate is 83.184/83.187 tok/s, a 0.66% median delta. All +three candidate repetitions are token-for-token identical to one another and +to that accepted GDN control; their compact token-list SHA256 is +`c243818da5cced8b8653df4c441f713e978f8dfe7d8343d1b78c0eb3aeffe8ef`. +The short default graph therefore removes the inactive long-wave pair without +changing the admitted output stream or causing a material short-context +regression. A cache-reload diagnostic is not used as quality evidence. + +### Operator and resource evidence + +A same-GPU operator A/B uses the production interleaved K/V layout. At 128K, +the merged route takes 370.087 us versus 377.429 us for separate long-wave +dispatch; at 256K it takes 599.839 us versus 611.925 us. The selected explicit +p896/p1664 controls take 362.175/593.954 us respectively. Merged-route eager +and CUDA Graph outputs are bit-exact to those explicit controls. The small +remaining dispatch delta is preferable to three long kernel/reducer pairs in +the captured full-model graph. + +The final 256K Nsight Systems CUDA-profiler-range trace uses the same +production interleaved layout and final extension. The active merged p1664 +partition kernel takes 0.576603 ms and its reducer takes 0.033280 ms. The +device-routed short guard kernel/reducer take 0.004544/0.003552 ms, for +0.617979 ms across Flash attention kernels. The old p256 trace takes +1.391781 ms in the partition kernel and 0.316769 ms in the reducer, or +1.708550 ms total. The retained route therefore cuts traced Flash attention +time by 63.83% (2.76x), while remaining bit-exact to the explicit p1664 +control. The trace is an isolated per-rank attention operator capture, not a +full-model TPOT decomposition. + +During the exact decode windows NVML reports 100% GPU-busy duty on every rank. +At 128K, memory-active duty averages about 48.2%, allocation is +29.45-29.75 GiB/GPU, and mean power is 200.7-215.8 W. At 256K, memory-active +duty averages about 45.5%, allocation is unchanged, and mean power is +211.6-226.8 W. SM and memory clocks remain 1530/877 MHz. NVML memory-active +duty is not achieved HBM bandwidth; exact SM, Tensor Core, occupancy, and DRAM +counters remain unavailable because Nsight Compute returns +`ERR_NVGPUCTRPERM` on this host. + +The production-layout correction is material: treating the unbound K/V views +as contiguous produced full-model output corruption and is rejected. A QK +software pipeline slowed the p896/p1664 kernels by roughly 30-38%; a split +reducer did not beat the retained route; paired E4M3 conversion measured +409.958/683.428 us at 128K/256K and lost to the shared LUT. All three +experiments are absent from the admitted path. + +### Post-acceptance kernel headroom + +Three additional changes retain the same admitted route and the real unbound +interleaved K/V views. First, the long-wave CTA stages its partition page IDs +once instead of reloading them for every tile. Second, a strict host-side +stride gate specializes the production `(blocks, 2, 1568, 1, 256).unbind(1)` +layout: the physical block stride is `2 * 1568 * 256`, while token and head +strides remain 256. Any other view uses the generic stride path. Third, the +fixed-layout E4M3 PV loop reads adjacent shared-memory values as `half2` while +retaining one FP32 FMA per dimension and the original token accumulation +order. It only changes lane ownership and restores dimensions explicitly at +the output store. + +Each stage was measured as a four-arm control/candidate/candidate/control A/B +on one otherwise idle V100. Every arm contains five repetitions for each of +three page-table layouts; the table reports the median of the ten samples per +layout followed by the median across layouts. The exact partition-boundary +companions at 131,071 and 262,080 tokens pass as well. + +| Exact context | Stage | Control | Candidate | Kernel delta | +|---:|---|---:|---:|---:| +| 128K | partition page IDs | 373.076 us | 368.604 us | -1.199% | +| 256K | partition page IDs | 602.017 us | 594.632 us | -1.227% | +| 128K | fixed interleaved stride | 370.102 us | 364.585 us | -1.491% | +| 256K | fixed interleaved stride | 598.257 us | 587.909 us | -1.730% | +| 128K | shared-V `half2` PV | 362.739 us | 348.265 us | -3.990% | +| 256K | shared-V `half2` PV | 584.745 us | 558.515 us | -4.486% | + +All 12 sequence/layout cases have one output SHA256 across all four arms and +are bit-exact to their explicit p896/p1664 controls (`max_abs=0`). The final +partition kernel remains at 168 registers/thread and two CTAs/SM; the PV +change removes the prior 32-byte stack frame. Relative to the fixed-stride +binary, its static SASS instruction count falls from 2,424 to 2,248 and shared +loads fall from 118 to 90, with the same 57 FFMA instructions. These are +static cubin counts, not NCU throughput counters. The timed binary SHA256 is +`43f742cc228100dd0a58aed37bf1cb6007447129bcf3f1d28148bee71b6ff2a5`. +The clang-formatted final build SHA256 is +`5e66a5331aaf83e5e15080c5122c78047b94d3cb47a721572613bbbba9172d47`; +its target-kernel SASS is byte-identical to the timed build (normalized SASS +SHA256 `0684d453072ed630aa68477b618ab4c090eb18c87422a6de90af3fd43740a1ee`) +and its resource record is unchanged. The final build is bit-exact in eager +and CUDA Graph replay at 128K/p896 and 256K/p1664 (`max_abs=0`); all eight +parameterized page-1568 wave-boundary regressions pass. + +Applying the three independently matched operator ratios to the measured +5.921/9.597 ms attention slices, while leaving the measured 10.251/10.253 ms +non-attention residuals unchanged, projects 15.784 ms (63.354 tok/s) at 128K +and 19.151 ms (52.216 tok/s) at 256K. These are decomposition-based +projections, not replacement end-to-end measurements; the accepted measured +endpoint remains 61.834/50.376 tok/s. + +A packed pair-LUT initialization experiment is rejected because it is about +1.0-1.2% slower at 128K and 0.78% slower at 256K. Replacing constant division +with compare/subtract address normalization passed 2,801,664 exhaustive host +address checks and all GPU exactness cases but was statistically neutral. A +scalar PV-accumulator rewrite removed the stack frame but was neutral or +slower at 256K. None of these rejected variants is present in the admitted +path. + +Primary retained evidence is: + +- `.artifacts/qwen38_nvfp4_speed_20260823/results/e4m3_xqa_auto_single_p12288_long_i16k_64k_128k_256k_o64.json` +- `.artifacts/qwen38_nvfp4_speed_20260823/results/e4m3_xqa_lut_merged_wave_p12288_long_i128k_256k_o64.json` +- `.artifacts/qwen38_nvfp4_speed_20260823/telemetry/e4m3_xqa_lut_merged_wave_p12288_long_i128k_256k_o64.nvml.csv` +- `.artifacts/qwen38_e4m3_xqa_long_route/results/e4m3_xqa_lut_separate_same_gpu4.json` +- `.artifacts/qwen38_e4m3_xqa_long_route/results/e4m3_xqa_lut_merged_same_gpu4.json` +- `.artifacts/qwen38_e4m3_xqa_long_route/results/e4m3_xqa_wave_boundaries_graph_workspace.json` +- `.artifacts/qwen38_e4m3_xqa_long_route/results/page_ids/` +- `.artifacts/qwen38_e4m3_xqa_long_route/results/page_ids_fixed_stride/` +- `.artifacts/qwen38_e4m3_xqa_long_route/results/page_ids_fixed_stride_pv_half2/` +- `.artifacts/qwen38_e4m3_xqa_long_route/results/page_ids_packed_lut/` +- `.artifacts/qwen38_e4m3_xqa_long_route/results/page_ids_fast_address/` +- `.artifacts/qwen38_e4m3_xqa_long_route/profiles/e4m3_xqa_final_merged_wave_p1664_l262144_nsys.nsys-rep` diff --git a/docs/design/sm70_qwen38_qpn8_decode.md b/docs/design/sm70_qwen38_qpn8_decode.md new file mode 100644 index 0000000000..21ae39aafb --- /dev/null +++ b/docs/design/sm70_qwen38_qpn8_decode.md @@ -0,0 +1,245 @@ +# SM70 Qwen3.8-27B-FP8 QPN8 Decode Acceptance + +Date: 2026-08-23 + +## Decision + +The memory-neutral QPN8 dense route is available as an opt-in for the +Qwen3.8-27B-FP8 TP4 no-MTP model and projection shapes when configured +`max_num_seqs` is at most eight. Set `VLLM_SM70_FP8_QPN8=1` to request it. +The model-quality gate passes, but pure-FP8 default promotion still requires a +same-source speed rerun with the same number of steady decode intervals. An +explicit request with missing QPN8 operators fails closed. + +The gate checks the Qwen3.8 text architecture, 5120 hidden size, 17408 +intermediate size, 64 layers, full-attention interval 4, head dimension 256, +and TP4. Qwen3.6 and both 35B acceptance routes remain unchanged. +Configurations admitting more than eight live sequences retain TurboMind at +load time until a measured M>8 QPN8 route exists. MTP configurations also +retain TurboMind until their verification widths pass speed and quality gates. +These guards prevent the functional large-M dequantization fallback from +becoming a concurrency or speculative-decode regression. + +This acceptance is a no-MTP, M=1 decode speed result with numerical tests up +to M=8 and model-quality batches up to eight. M>8 remains functional through +one bounded 85 MiB FP16 weight workspace. Large-M fused gate/up also has a +transient `M x 8704` FP16 GEMM output, so that fallback is not an accepted +high-concurrency memory or speed result. + +## Contract + +- Model: `Qwen3.8-27B-FP8`, TP4 on physical V100-SXM2-32GB GPUs 4-7. +- Runtime: Python 3.12, Torch 2.10.0+cu128, CUDA 12.8, source checkout and + source-built SM70 extension; no wheel changes are used. +- Production speed case: no MTP, E5M2 KV, `FLASH_ATTN_V100`, CUDA graphs, + `max_model_len=262144`, input 1024, output 256, and one live request. +- Official sampling: `temperature=1.0`, `top_p=0.95`, `top_k=20`, request seed + 20260815. Pure decode excludes TTFT/prefill and uses the 255 steady token + intervals. +- Quality case: identical baseline/candidate datasets and seeds, 4096-token + rolling-PPL windows, and generation/log-likelihood batches 1, 2, and 8. + +## Source Route + +Checkpoint-native block-FP8 weights are repacked once at model load into one +QPN8 `[K, N]` code tensor plus grouped FP16 scales. The old TurboMind tensor +is replaced rather than retained, avoiding the approximately 6 GiB/rank +duplicate-weight experiment. The matched candidate reported 19.89 GiB +available KV memory versus 19.80 GiB for the control. + +The opt-in model and shape gate admits only: + +| Projection | TP-local K | TP-local N | Decode schedule | +|---|---:|---:|---| +| MLP gate/up with fused SiLU×up | 5120 | 8704 | split-K 8, two accumulator chains, prefetch | +| MLP down | 4352 | 5120 | split-K 16, one accumulator chain | +| GDN/full-attention output | 1536 | 5120 | split-K 16, one accumulator chain | + +GDN input and full-attention QKV remain on TurboMind. The M decision is inside +the opaque C++ operator so AOTInductor cannot incorrectly specialize a Python +M=1 branch for a later prefill. M=1-8 calls QPN8 directly; larger M +materializes one layer at a time in the shared weight workspace and calls FP16 +GEMM. The large-M gate/up temporary described above is bounded by the engine's +`max_num_batched_tokens`, not by the 85 MiB weight workspace. + +## Pure Decode Result + +| Route | Pure TPOT | Steady decode | Change | +|---|---:|---:|---:| +| Matched TurboMind control | 16.985327 ms | 58.874 tok/s | baseline | +| QPN8 accepted subset | 15.808000 ms | 63.259 tok/s | +7.448%, -1.177327 ms/token | +| Complete-source TurboMind control | 16.961295 ms | 58.958 tok/s | source A/B baseline | +| Complete-source QPN8 candidate | 16.347973 ms | 61.170 tok/s | +3.752%, -0.613322 ms/token | + +The matched operator-integration pair generated all 256 tokens and finished by +length. Their sampled token streams first differ at output token 91; sampled +token identity is not the quality gate because both trajectories are coherent +and the acceptance criterion is PPL plus fixed reasoning, knowledge, and +Chinese datasets. + +The complete-source control generated 256 tokens and finished by length. The +complete-source QPN8 request used the same prompt, sampling contract, +and source extension, but encountered the model's normal stop token 248044 at +token 233. Its TPOT therefore covers 232 steady intervals rather than the +control's 255. The earlier 63.259 tok/s matched operator-integration result is +retained as evidence of headroom. Because 61.170 tok/s uses 232 intervals while +the control uses 255, another same-source equal-interval run is required before +pure-FP8 default promotion. + +## Nsight Systems Per-Token Decomposition + +Node tracing adds profiler overhead, so these rows explain composition rather +than replace the unprofiled TPOT above. + +| Steady replay metric | Control | QPN8 | Delta | +|---|---:|---:|---:| +| Replay interval | 17.418 ms | 16.437 ms | -0.981 ms | +| GPU activity union | 16.870 ms | 15.894 ms | -0.976 ms | +| Total idle gaps | 0.548 ms | 0.543 ms | -0.005 ms | +| Kernel launches/rank/token | 1085.5 | 1085.9 | unchanged | + +The trace gain is therefore GPU kernel service, not a hidden CPU wait or +launch-gap reduction. Dense service changes from 10.335 ms/rank/token to: + +| Candidate dense component | Time | Launches/rank/token | +|---|---:|---:| +| Fused QPN8 gate/up | 4.018 ms | 64 | +| QPN8 down and output projections | 3.040 ms | 128 | +| Remaining TurboMind GDN/full-QKV input | 2.332 ms | 64 | +| Total dense | 9.391 ms | 256 | + +This accounts for about 0.944 ms/token of the trace improvement. TP +communication is 1.680 to 1.637 ms and LM-head/sample/gather remains about +1.33 ms, so neither explains the main win. + +## Resource Release and Peak-Band Residency + +Whole-request NVML averages are not used as achieved FLOPS or bandwidth. +Nsight Systems establishes how long grids can occupy the machine, and Nsight +Compute supplies achieved per-kernel counters. + +The grid-limited occupancy ceiling moves from 93.56% of service below 25% +occupancy in the control to 48.56% below 25%, 32.11% at 25-50%, and 19.26% at +50-75% with QPN8. This is a duration distribution, not an average-utilization +claim. + +For the accepted QPN8 kernels, NCU reports: + +| Shape | Registers/thread | Achieved occupancy | SM throughput | DRAM throughput | Tensor-pipe elapsed | +|---|---:|---:|---:|---:|---:| +| Down | 50 | 48.90% | 27.20% | 82.74% | 9.50% | +| Output | 50 | 48.17% | 27.69% | 71.32% | 8.23% | +| Fused gate/up | 64 | 42.33% | 29.98% | 87.74% | 10.08% | + +Counting the useful M=1 dense multiply-add work, throughput is 1.177 to +1.295 TFLOP/s per GPU, or 4.71 to 5.18 TFLOP/s across TP4. +Dense-service-weighted Tensor Core pipe activity is 7.97% to 9.15%, +equivalent to about 9.99 to 11.47 TFLOP/s per GPU, or 39.96 to 45.88 TFLOP/s +across TP4, relative to the V100's nominal FP16 Tensor Core peak. These answer +different questions: the first is model-useful arithmetic, while the second +includes tensor instructions spent inside the split-K implementation. Neither +is a whole-request average. + +Across the 9.391 ms candidate dense service, 7.059 ms resides in the 25-50% +SM-throughput band and 2.332 ms remains below 25%; no dense kernel reaches +50% SM throughput. For DRAM, 6.172 ms resides at 75-90% and 3.219 ms at +50-75%; no dense kernel reaches 90%. Dense-service-weighted achieved +occupancy rises from 9.96% to 36.25%, but the duration bands above are the +primary evidence. + +The result is not “the V100 cannot reach peak compute.” QPN8 makes a much +larger fraction of dense time sustain high memory duty and occupancy, while +the unchanged approximately 0.54 ms idle budget shows that host launch work +is not the present first-order limit. + +## SGLang-V100 Comparison Boundary + +The current [sglang-V100 README](https://github.com/haohervchb/sglang-V100/blob/main/README.md) +now reports a Qwen3.8-27B-FP8 target-only, E5M2 KV, TP4 row at 58.2 tok/s for +1K decode and 50.8 tok/s for 25K decode. Its table uses one cold request and +256 greedy output tokens. The linked +[target-only audit](https://github.com/haohervchb/sglang-V100/blob/main/benchmark/qwen38_27b_fp8_target_e5m2_v100_20260822/README.md) +records 17.180 ms TPOT, or 58.21 tok/s, at 1,024 input tokens. The 63.259 +tok/s QPN8 result here is 8.67% above that published number, so the public +evidence points to this route being faster. It is not a strict head-to-head +because this acceptance uses fixed random sampling and a different +prompt/client harness. A formal claim still requires running SGLang's +target-only command with the same prompt, sampling, graph, and pure-decode +accounting used here. + +The published 136.6 tok/s DFlash2-8 row is speculative and is not a no-MTP +comparison. + +## Quality Gate + +The gate compares the prior TurboMind route and QPN8 with identical prompts, +few-shot examples, seeds, and batching. Lower is better for the three +WikiText metrics; higher is better for the remaining accuracy metrics. + +| Dataset and metric | Samples | Control | QPN8 | Result | +|---|---:|---:|---:|---| +| WikiText word perplexity | 62 documents | 8.790013 | 8.789904 | pass, lower | +| WikiText byte perplexity | 62 documents | 1.501519 | 1.501515 | pass, lower | +| WikiText bits/byte | 62 documents | 0.586422 | 0.586419 | pass, lower | +| GSM8K 5-shot strict exact match | 128 | 69.531% | 70.312% | pass, +1 item | +| GSM8K 5-shot flexible extraction | 128 | 73.438% | 73.438% | pass, equal | +| MMLU abstract algebra 5-shot | 100 | 67.0% | 67.0% | pass, equal | +| MMLU high-school computer science 5-shot | 100 | 89.0% | 89.0% | pass, equal | +| MMLU professional law 5-shot | 100 | 63.0% | 66.0% | pass, +3 items | +| MMLU three-subject macro | 300 | 73.0% | 74.0% | pass | +| C-Eval computer network 5-shot | 19 | 73.684% | 73.684% | pass, equal | +| C-Eval high-school Chinese 5-shot | 19 | 73.684% | 73.684% | pass, equal | +| C-Eval advanced mathematics 5-shot | 19 | 57.895% | 57.895% | pass, equal | + +No evaluated metric regresses. This closes the quality side of the opt-in +route; it does not replace the missing equal-interval performance rerun. + +## Build and Tests + +- The complete source `_C` target built 47/47 objects for CUDA 12.8 and SM70; + the resulting 45 MiB extension registers all six prepare, dequantize, + prefill, dispatch, GEMM, and fused-gate QPN8 operators. +- The operator race passed 960 rows. Maximum relative L2 was `2.851e-4`, + minimum cosine was `0.99999988`, and maximum absolute error was `9.77e-4`; + CUDA Graph replay deltas were zero. +- Seventeen focused Python tests pass, covering opt-in/explicit-off policy, + model and shape gates, workspace reuse, M=1 and M>8 opaque dispatch, fused + gate/up dispatch, and warmup registration. +- Ruff lint/format, changed-line clang-format, Python byte compilation, shell + syntax checks, and `git diff --check` pass. + +## Next Optimization Order + +1. Extend QPN8 with a single-launch two-dimensional row-tile grid and race + M=9, 11, and 16 against TurboMind. Do not compose an 8-row launch with a + separate tail launch, and do not assume the batch-one winner stays best. +2. Race M=17, 32, and 64 and keep a measured route map. The external + [v100-skinny QPN study](https://github.com/dnv2003/v100-skinny/blob/main/docs/qpn_race_notes.md) + also finds a crossover to WMMA/TurboMind above the small-batch band. +3. Before high-concurrency acceptance, replace or chunk the large-M fused + gate/up temporary and record peak memory under the supported batched-token + limit. +4. Add a bounded large-M route for GDN input/full-QKV, then quality-gate it. + Those 64 remaining TurboMind launches consume 2.332 ms/token. +5. Continue shape-specific work on fused gate/up (4.018 ms) and down/output + (3.040 ms), measuring time above 90% DRAM and 50% SM rather than optimizing + a whole-run average. +6. Revisit the 1.637 ms communication and 1.33 ms LM-head only after the dense + residue falls. The current trace does not justify prioritizing Python launch + overhead. +7. Repeat the decomposition at 25K and 70K. SGLang's target-only 70K trace + reports 10.60 ms/token in FP8 projections and 10.35 ms/token in attention, + so the 1K optimization order must not be hard-coded for long context. + +The ordering is also consistent with external evidence, without treating a +newer-architecture implementation as a V100 drop-in. The +[v100-skinny M sweep](https://github.com/dnv2003/v100-skinny/blob/main/docs/REPRODUCE.md) +uses a measured SIMT/QPN/WMMA dispatch map rather than one kernel for every M. +The [CUTLASS grouped scheduler](https://github.com/NVIDIA/cutlass/blob/main/media/docs/cpp/grouped_scheduler.md) +shows how a bounded persistent CTA work queue can retain residency across many +independent small tiles; it becomes relevant only if independent problems can +be grouped without crossing transformer dependencies. The +[FlashInfer paper](https://arxiv.org/abs/2501.01005) similarly makes dynamic, +load-balanced scheduling compatible with a static CUDA Graph. For this route, +those are design constraints to test under concurrency, not reasons to replace +the measured SM70 kernels speculatively. 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 74b8425639..7bb276c394 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -24,6 +24,576 @@ Goal: sampling, and normal MoE architecture wherever possible. - Do not carry unsafe quality-regressing routes into the main path. +## Active MRV2 DFlash2 campaign, 2026-08-20 + +- Integration base: `onecat/main@7aede2cf010d92815c9d7bff25867b4fa009b6cb`. +- Review chain starts at Draft PR #252 on + `codex/v100-dflash-mrv2-base-20260820-041115`. +- Frozen sources: vLLM MRV2 DFlash base `9842d7014502`, vLLM DFlash2 PR + `66e5414c6d75`, SGLang `c14312a66420`, and sglang-V100 `5526ef1c6a82`. +- Route contract: `method=dflash` is MRV2-only; the experimental V1 flat path + is removed. `method=dflash_ddtree` retains the existing V1 DDTree scheduler, + payload, verifier, and diagnostics through a narrow model compatibility + adapter. Existing Eagle and MTP implementations are not replaced. +- Selector top-K is checkpoint-owned and fixed at 16 for Qwen3.8 DFlash2; + official bring-up uses seven draft tokens (block size eight). Target sampling + top-k 20 is a separate sampling parameter. +- PR1 evidence so far: focused routing/config/DDTree tests pass, and the V1 + DFlash input-preparation test passes on V100 after supplying the KV-group + initialization contract used in production. Two unrelated DDTree metadata + assertions fail identically on the frozen base and are recorded as pre-existing. +- Draft PR #253 is the PR2 review scope on + `codex/v100-dflash2-sm70-20260820-044011`, based only on PR1. Its worktree is + `/home/ymzx/桌面/1cat-vllm/worktrees/v100-dflash2-sm70-20260820-044011`. +- PR3 evaluated a gated SM70 dense-FP16 selector fusion on top of PR2. Its + required M=7 gate was negative, so the final PR #254 merge scope retains only + the correctness/performance evidence below; no selector implementation, + runtime flag, duplicate weight layout, or benchmark harness enters `main`. + +### PR2 DFlash2 correctness evidence, 2026-08-20 + +- Checkpoint contract is frozen at + `incoai/Qwen3.8-27B-DFlash2@dedf8df68adfb1afeaf7b7480c0a0243108177b4`. + The retained local `model.safetensors` SHA256 is + `67fc76d68dc5a9415511a4f394ef744d67510cd20e93b37cc2cc7d28e4bab65c`. + Its architecture is `DFlash2DraftModel`, selector K is 16, block size is 8, + attention is non-causal with sliding window 2048, and checkpoint layer IDs + `[5,19,33,47,61]` map to target hidden boundaries `[6,20,34,48,62]`. +- The dependency closure now contains the official grouped convolution, + selector codebooks, TP-aware candidate top-K, DFlash2 MRV2 walk/cache, FP32 + proposal scores, rejection `log1p`, and a separate selector compile-cache + tag. FlashInfer top-K is capability-gated and is never called on SM70. +- The SM70-only numeric adaptation keeps FP32 residual state, explicit BF16 + round-to-nearest-even points, scaled FP16 attention and gate-up projections, + and per-row power-of-two SwiGLU transport. The seventh selector slot uses a + separate Volta tail kernel; other architectures retain the upstream walk. +- Flash-V100 graph capture now binds DFlash's K+1 non-causal query directly to + persistent paged-prefix metadata. The causal small-query decode branch is + unavailable to a non-causal draft. Target E5M2 and draft FP16 KV pages are + unified without truncating hybrid Mamba physical padding. +- Single-request eager and CUDA Graph runs produced the exact same 16 token IDs, + hash `6de4c3f326c76ecc1b31b8026501ddf22f73a34b53b2fc46c622dee6b4a2d995`, + candidate lattice, acceptance trace, and draft KV. The graph artifact is + `/data/minimax-h3/task-cache/v100-dflash2-20260820/pr2/dflash2-graph-greedy16-fixed.json`; + the matched eager artifact is `dflash2-eager-noncausal-fixed16.json`. +- Matched batch-2 eager/graph runs generated 96 and 61 tokens with identical + hashes and all acceptance counters. CUDA Graph batch-4 generated mixed + lengths 64/61/64/64 without OOB, stale-buffer, NaN, or replay corruption; + each result matched the corresponding eager prefix. Retained artifacts are + `dflash2-{eager,graph}-gsm8k-mixed-b2-2x96-greedy.json` and + `dflash2-graph-gsm8k-mixed-b4-4x64-greedy.json` under the PR2 cache directory. +- On four fixed GSM8K prompts with the Qwen3.8 xhigh chat template and official + target sampling, greedy drafting emitted 421 tokens with mean accepted length + 4.1068 and per-position acceptance + `[.8544,.6505,.4854,.4175,.3301,.2233,.1456]`. Probabilistic drafting emitted + 452 tokens with mean 4.2407. These short V100/FP8 results are correctness and + sampling diagnostics, not a like-for-like replacement for the published + H200/BF16 acceptance figures. +- The fixed 32-to-16 graph diagnostic reported 80.11 steady decode tok/s versus + 18.72 eager with exact token/acceptance equality. The matched batch-2 run was + 116.67 tok/s graph versus 43.53 eager. These are short diagnostic shapes and + are not a target-only or dense-selector paired performance gate. +- DFlash1 remains on the official MRV2 speculator for the local 9B, 27B, and + 35B-A3B checkpoints. A real 9B-AWQ/FP16-draft eager smoke produced eight + tokens, hash `2a205303e4c2b54b62d7562643abcaa425952d7dbcf3f3c1a4b6a4e3e3634be9`, + accepted length 3.5, and confirmed Flash-V100 routing. The artifact is + `/data/minimax-h3/task-cache/v100-dflash2-20260820/pr2/dflash1-9b-awq-mrv2-eager-route.json`. +- Existing MTP was not replaced. A real Qwen3.8 TP4 MTP4 probabilistic smoke + confirmed native MTP, local argmax, Flash-V100, and E5M2 routes; it emitted 32 + tokens with accepted length 4.25 and per-position acceptance + `[1.0,1.0,.75,.5]`. The artifact is + `/data/minimax-h3/task-cache/v100-dflash2-20260820/pr2/dflash2-mtp4-qwen38-regression-eager.json`. +- Final changed-path gate: SM70 arguments plus KV layout, DFlash2/MRV2 routing, + and warmup reservation total `113 passed, 1 skipped` on CPU; the dedicated + V100 Mamba align kernels add `5 passed`. The new warmup matrix covers MRV2 + DFlash K+1, DDTree tree-budget lookahead, unchanged Eagle/MTP/DSpark values, + and align/non-align Mamba speculative tails. Earlier V100 rejection sampling + remains `14 passed` and probabilistic selector score-cache `1 passed`; Ruff + passes every changed Python file. The DDTree targeted matrix remains + `71 passed, 1 failed`, with + the one shared-storage assertion reproduced identically on PR1. Three stale + MTP step-index helper failures also reproduce on PR1; their 19 unaffected + cases pass. Four Eagle tests that require the gated Llama-3.1 config stop at + Hugging Face 401 on both PR1 and PR2; the local DFlash/Eagle compatibility + test itself passes. +- MRV2 now implements the upstream hybrid Mamba `align` state migration on the + GPU without changing the V1 DDTree state-selection path. The new-request + state column is seeded from the resolved `MambaSpec.block_size`, not the + smaller global cache block size; this is required when E5M2 target pages and + FP16 DFlash pages share a hybrid pool. Upstream hash-layout fix #51180 is also + included so aligned Mamba groups retain the finest divisible hash unit. +- Prefix correctness is accepted for the real Qwen3.8-27B-FP8 TP4 route. With + the resolved target Full/Mamba block size 3296 and draft SWA block size 1648, + a 3505-token prompt produced a 3296-token hit on the second request. Eager + kept the exact 16-token hash `a2c583fc17c7...`, accepted length 6.6667, and + per-position counts `[3,3,3,2,2,2,2]`; TTFT changed from 1.117s to 0.303s. + CUDA Graph kept the same token and acceptance trace with TTFT + 1.119s to 0.174s. Artifacts are `dflash2-prefix-tp4-{eager,graph}-4k.json` + under the PR2 cache directory. +- A paired target-only TP4 probe hit 3136 tokens with exact output and TTFT + 1.010s to 0.167s (`qwen38-prefix-tp4-target-only-eager-4k.json`). This + separates target Mamba-state restoration from DFlash proposal correctness. + No separate cross-request draft-prefix index or persistence layer was added; + the accepted same-process path uses the existing unified hybrid block cache, + while an ordinary miss rebuilds draft context KV during DFlash proposal. +- Numeric classification for the DFlash non-causal paged attention route is + `B-accept`, with its layout component separately passing Type A. Aligned + DFlash shapes at sequence lengths 8/69/512/2048 show semantic-boundary + paged-vs-dense maxima + `0/2.44e-4/1.22e-4/6.10e-5`, all below the FP16 `2e-3` bound and not growing + with length. Scaling V by 256 reproduces raw differences up to 0.0625, which + become at most 2.44e-4 after the model's required `/256` range restoration. + Reversing every physical KV page and applying the matching block-table + permutation is bitwise exact for every shape, proving page addressing, + layout, mask, and KV-source alignment. The bounded cross-kernel residual is + mechanically explained by the source arithmetic: for D != 256, dense + Flash-V100 scales QK by `log2(e)` and evaluates `exp2f`, whereas paged + Flash-V100 retains natural-log scores and evaluates `expf`. These are + mathematically equivalent reductions with different FP16 rounding. The + release extension was built from paged-kernel source SHA256 + `9be6b237fcd1c89653c26460054c1c25edd6e7d8f9685f7b9f4b83a2f1ba5e58`, + identical to the integration source; an older independently built extension + produced the same case-by-case results. + The retained evidence is + `/data/minimax-h3/task-cache/v100-dflash2-20260820/pr2/dflash2-noncausal-paged-attention-exactness.json`. + Together with exact eager/graph model tokens, candidates, acceptance trace, + and draft KV, the attention numeric gate no longer blocks PR3. This does not + relax the separate fused-selector performance and exact-token gates. + +### PR3 SM70 selector-fusion disposition, 2026-08-21 + +- Draft PR #254 evaluated reuse of the TurboMind dense-FP16 LM-head kernel + behind `VLLM_SM70_DFLASH2_FUSED_SELECTOR`. The experiment preserved the + checkpoint contract—local top-20 candidates reduced to the checkpoint-owned + global top-16, with the official seven-token draft block—but the production + implementation and flag were removed after the negative gate. +- The original top-20 epilogue assigned one row to each of four warps although + its CTA has `M=8`. The fixed epilogue assigns rows round-robin, covering the + seventh DFlash2 row and M=8 tail. A graph-safe operator variant accepts + persistent FP32 value and INT64 token partial workspaces; the allocation-owning + legacy entry remains available for non-graph callers. +- Equal FP16 logits exposed a second exactness issue: the fused kernel's + deterministic token-ID tie break did not reproduce `torch.topk`'s dense-shard + ordering. The accepted eager experiment scatters the fused top-20 into a + persistent shard-shaped FP16 buffer and runs local dense-order top-16 before + exchanging only 16 values and token IDs per TP rank. This restores exact + ordered IDs and values, not merely the same candidate set. +- The measured experimental SM70 extension SHA256 was + `077ca806e8fad6b452dbe59cd0da336ff88c9341386c8480bbbe009b8029c7b2`. + Persistent-workspace tests cover M=1/4/5/7/8, TP4 vocab offsets and padding, + equal-score ordering, and CUDA Graph replay. The M-shape set passed three + consecutive runs after the test stopped exercising the allocation-owning + legacy entry. +- The real Qwen3.8 TP4 selector microbenchmark used M=7, local N=62080, K=5120, + FP16 hidden/weight, local top-16 and TP global top-16. Ordered IDs and FP16 + values were exact with maximum difference zero. The dense + `torch.mm + local16 + TP global16` mean was 1.1585 ms; the production-shaped + fused path was 1.2254 ms, a 5.78% regression. Evidence is + `/data/minimax-h3/task-cache/v100-dflash2-20260820/pr3/sm70-top20-workspace-dense-order-m7-tp4.json`. +- The dedicated experimental-branch DFlash2 suite passed `30 passed, 11 skipped` without a + visible GPU and `41 passed` on V100. Ruff passes all eight changed Python + files, the rebuilt SM70 extension links successfully, and `git diff --check` + is clean. +- Earlier direct-top20 graph results that appeared much faster are rejected. + The dense 256-token artifact had hash `6477ffd9...` and accepted length + 2.2768, while the direct fused artifact had hash `1f7f1cdf...` and accepted + length 3.3462. The speed change therefore measured a different sampling and + acceptance trajectory, not an optimization of the same computation. +- Preparing the duplicate LM-head layout also raised model memory by about + 1.19 GiB per rank (10.23 versus 9.04 GiB in the final compiled-graph runs), + directly reducing long-context KV capacity. The fused selector is therefore + eager-only. Under the accepted Flash-V100 compiled CUDA Graph policy, an + explicit request is normalized to zero before workers fork; no duplicate + layout is prepared and startup prints the fallback decision. +- A paired post-rebuild TP4 graph check compared selector-off with a requested + selector that was normalized before worker startup. Both emitted the exact + same 32 token IDs, hash `471733c0...`, accepted-token count 13, mean accepted + length 1.6842, and per-position acceptance vector. Artifacts are + `dflash2-pr3-dense-rebuilt-graph-1k-32-r1.json` and + `dflash2-pr3-selector-normalized-graph-1k-32-r1.json` in the PR3 cache. +- Decision: do not merge the selector implementation; keep the official + draft-token default at seven. Because the required M=7 micro gate is already negative, + the 32K/128K confidence-interval sweep and 3/5/7 default scan are intentionally + not run. A future graph-capable fusion must first beat dense M=7 without + duplicate-weight capacity loss and then re-enter the full long-context gate. + +### DFlash acceptance root-cause repair, 2026-08-22 + +- A same-machine 64-prompt GSM8K audit held the FP8 target, E5M2 target KV, + FP16/auto draft KV, TP4, official probabilistic sampling, seven draft tokens, + Flash-V100, CUDA Graph, prompt order, and single-request sequencing fixed. + Before the repair, vLLM averaged 4.5395 completion tokens per verification + step versus 5.2489 for `haohervchb/sglang-V100`; the paired prompt delta was + 0.7094 with a 95% interval of [0.5882, 0.8306]. This excludes target FP8 as + the explanation for the acceptance gap. +- The MRV2 backport copied the grouped DFlash K-normalization call from upstream: + input `[L, num_ctx, num_kv_heads, head_dim]` with weights `[L, head_dim]`. + It omitted upstream dependency commit `02a1f23711`, whose stable RMSNorm + kernel uses the outer layer index to select each weight row. The local kernel + still treated every weight as one-dimensional and silently applied row zero + to all five draft layers. Checkpoint rows are not interchangeable: layers + 1-4 differ from layer zero by 0.468-0.621 mean absolute weight and + 0.383-0.443 relative L2. +- SGLang-V100 is unaffected because both paths preserve layer ownership: its + sequential fallback calls each layer's `k_norm`, and its fused Triton + materializer indexes `layer_id * k_norm_weight_stride_layer`. This is the + implementation difference that explains why SGLang acceptance was normal. +- Commit `c53b7974fb` restores the exact missing dependency semantics while + retaining the local SM70 RMSNorm additions. One-dimensional weights keep + stride zero; two-dimensional weights require `[input.size(0), hidden]` and + select the row from the outermost input dimension. Invalid rank and shape are + rejected instead of silently producing the wrong draft KV. +- The V100 kernel regression passes 13/13 cases for FP16, BF16, and FP32, + 3D/4D inputs, the real DFlash `[L, context, kv_heads, 128]` layout, a + non-power-of-two hidden size, and shape validation. Batched output is bitwise + identical to looping the original one-row RMSNorm. The repaired extension is + `7689e5e298be64a83c9f9743e08a073e857a5a4c89a79837396d05b37303d449`. +- The final 64-prompt repaired run averages 5.3888 per request and 4.5644 pooled, + versus 4.5395/3.9725 before repair and 5.2489/4.3571 for SGLang-V100. The + repaired-minus-control paired delta is 0.8493 with a 95% interval of + [0.7226, 0.9760]; it wins 60 of 64 matched prompts. Repaired per-position + acceptance is `[85.833,70.569,57.920,46.887,38.032,31.342,25.861]%`, above + SGLang-V100 at every position and within 0.0712 of the public official 5.46 + GSM8K mean. +- Acceptance recovery raises same-route aggregate output throughput from + 97.080 to 108.826 tok/s (+12.10%) and mean steady decode from 115.394 to + 134.940 tok/s (+16.94%). The repaired vLLM path remains slower than the + separately measured SGLang-V100 aggregate 156.063 tok/s, so acceptance is + closed but the remaining execution-speed gap is a distinct optimization item. +- Evidence is retained in + `/data/models/v100-dflash2-20260820/pr3/acceptance-rootcause/` as + `dflash2-probabilistic-graph-{zlabshuffle42,fixed-batchedrmsnorm-zlabshuffle42}-randomseed-first64-officialprompt-e5m2-o4096.json`; + the SGLang comparison is + `/data/models/v100-dflash2-20260820/sglang-audit/sglang-dflash2-graph-zlabshuffle42-randomseed-first64-officialprompt-e5m2-o4096.json`. + Earlier four-prompt and selector-pair acceptance values predate this repair + and remain route/equality diagnostics only; do not cite them as DFlash2 + acceptance baselines. A BF16 target rerun is intentionally unnecessary + because the causal A/B kept the same FP8 target and changed only K-norm row + selection. + +### DFlash2 single-request execution-cost root cause, 2026-08-22 + +- The repaired 64-prompt single-request runs close acceptance but not execution + cost. For requests with at least 512 completion tokens, the round-weighted + vLLM decode cost is 40.7148 ms/verification versus 26.8166 ms for the same + SGLang-V100 target, draft, TP4, E5M2-KV, block-8, sampling, prompt order, and + CUDA-Graph route. The 13.8982 ms gap is not caused by acceptance: vLLM's + per-request mean is 5.3888 versus SGLang's 5.2489. +- Matched CUDA-Graph node traces show that the one-pass DFlash2 proposal is not + the bottleneck. The vLLM and SGLang draft-model-plus-selector graphs take + 3.976 and 3.915 ms respectively under tracing. Their dense draft LM head, + five-layer draft forward, non-causal attention, selector top-K, and selector + walk therefore have effectively the same cost on this route. +- The raw traced round decomposition is: + - vLLM: draft graph 3.976 ms; draft-to-target transition 18.148 ms with 431 + eager kernels; target graph 24.845 ms; target-to-next-draft transition + 2.786 ms with 63 eager kernels. + - SGLang: draft graph 3.915 ms; draft-to-target transition 3.386 ms with 30 + eager kernels; target graph 19.622 ms; target-to-next-draft transition + 4.164 ms with 125 eager kernels. + - These traced wall times include launch-tracing overhead and must not replace + the unprofiled 40.7148/26.8166 ms round results. They are valid for + composition. SGLang captures its target LM head and TP all-gather inside the + target graph (about 1.11 ms combined), while vLLM executes those after its + target graph, so raw graph boundaries are not a perfectly semantic split. +- The dominant pre-verification gap is a hybrid-KV metadata amplification. The + target has 48 GDN and 16 full-attention layers, while the draft has five + layers. vLLM's unified grouping therefore creates ten padded GDN groups and + four padded full-attention groups. GDN metadata does not support cached + block-table updates, so `build_gdn_spec_decode_state_contract` repeats five + boolean-index selections for every GDN group. The trace contains exactly 50 + each of CUB reduce, compact-init, and select kernels per round, plus 281 other + metadata/input kernels: 431 launches and 18.148 ms of traced wall time for + only 2.452 ms of GPU service. SGLang keeps one persistent linear-attention + metadata object shared by all GDN layers; its corresponding transition has + 30 kernels and 0.117 ms of GPU service. +- SGLang also has a genuinely cheaper verifier graph: 1,130 nodes and 19.095 ms + GPU service versus vLLM's 2,612 nodes and 23.547 ms. Dense target FP8 GEMM is + already equal (10.798 versus 10.877 ms), and full-attention service is close, + so neither TurboMind GEMM nor Flash-V100 attention is the first DFlash2 + target. The main graph differences are: + - vLLM executes 1,422 generic vectorized/unrolled/elementwise nodes totaling + about 4.625 ms; SGLang executes 211 such nodes totaling about 0.677 ms. + - vLLM custom all-reduce costs about 2.350 ms per target graph versus 1.113 ms + for SGLang's one-stage push kernel. + - SGLang fuses sigmoid-gating computation, recurrent verification, and + intermediate-state writes. After rejection sampling, two all-layer fused + gather/scatter kernels commit the accepted SSM and convolution state. Its + target-to-draft interval also uses fused target-hidden-to-draft KV + materialization. +- Three narrow attempts to remove the repeated vLLM work were rejected and + fully reverted: reusing the DDTree one-kernel updater, reusing only its + pure-spec state path, and bypassing only the state-contract boolean + selections. All visible metadata regressions passed (34 focused tests), but + every shortcut changed the fixed-seed token trace at token 12, from the + 235-token/40-round control to a 309-token/57-round trajectory. This proves + that deleting the allocations/indexing before establishing explicit graph + buffer ownership is unsafe. The current evidence is consistent with a hidden + padding, alias, scratch-lifetime, or stream-order dependency; that precise + mechanism remains an inference rather than a proven field mismatch. +- Implementation decision and gates: + 1. Add a DFlash-only persistent GDN verify metadata/state object. Every live + and padded field is overwritten each epoch, and a debug shadow path compares + it with the existing generic builder before graph replay. + 2. Replace the ten repeated state contracts with one DFlash-specific fused + metadata kernel that writes shared structural metadata plus group-specific + state block IDs. Do not route Eagle, MTP, or DDTree through it. + 3. Port SGLang's fused target-verify GDN and two all-layer state-commit + kernels, including the per-step FP16 store/reload boundary needed for + recurrent-state numerical equivalence. + 4. Then port/capture fused target-hidden-to-draft KV materialization and tune + the one-stage TP4 all-reduce. Sparse rejection-sampling work is secondary. + 5. Do not reintroduce selector fusion until a new design clears the M=7 gate: + the measured proposal graphs are already equal, and the duplicate-layout + fused selector is slower and larger. + The first performance gate is an unprofiled round cost at or below 32 ms with + the fixed-seed token hash unchanged; the follow-up target is 28-29 ms before + claiming parity with SGLang's 26.8166 ms reference. Greedy tokens and accept + trajectories must be exact; probabilistic acceptance may differ by at most + 0.05 on the fixed corpus, with batch 1/2/4 mixed-length graph replay and + poisoned padding/state buffers included. +- Evidence is retained in + `/data/models/v100-dflash2-20260820/pr3/perf-rootcause/` as the vLLM baseline, + graph-node report, SQLite export, and rejected A/B JSON files. The matching + SGLang report is + `/data/models/v100-dflash2-20260820/sglang-audit/perf-rootcause/sglang-dflash2-single1-step20-v2.{qdstrm,nsys-rep,sqlite}`. + +### PR257 shared-metadata and packed-GDN checkpoint, 2026-08-22 + +- Draft PR #257 is isolated on + `codex/v100-dflash2-gdn-metadata-20260822-122723`; it is stacked on PR #254 + and changes only the MRV2 DFlash2 target-verification path. Eagle, MTP, + DFlash1, and `dflash_ddtree` retain their existing dispatch. Both candidate + leaves remain default-off while the TP4 profile and statistical-quality + gates are pending. +- The first leaf adapts vLLM PR #52297 at fixed head + `5f4c9678d1bb791ffb3ad23a9ae4129db9eb299d`: request classification, token + indices, and query offsets are computed once per batch, while cache-group + state block IDs remain group-owned. A same-hash/AAL 64-token Graph B1 probe + improved steady decode from 84.79 to 91.10 tok/s and decode time from 0.7430 + to 0.6915 s. The eager candidate regressed, so this evidence supports only + the production Graph route. An eight-request probabilistic probe changed the + aggregate acceptance trajectory and is not accepted as a compute-only speed + comparison. +- This port does not yet remove the dominant state-contract fan-out. Each of + ten padded GDN cache groups still performs group-specific block-table work + plus repeated boolean compaction for accepted-state metadata. The next + metadata step must be judged by the measured draft-to-target interval and + kernel count, not by the upstream H200 claim alone. Upstream reports its + builder moving from about 900 to 300 microseconds and B1 throughput +61%; the + V100 result remains an independent gate. +- A second DFlash2-only leaf removes packed-QKV rearrangement and the final GDN + output copy while preserving the existing split gating materialization. The + first implementation used `q * rsqrt(sum(q^2)+eps)` where the accepted + recurrent verifier uses `q / sqrt(sum(q^2)+eps)`. Although mathematically + equivalent, that changed at most one FP16 ULP in the operator test and moved + a fixed-seed probabilistic decision at output token 111. This explains why + the earlier fully fused and packed candidates followed the same alternate + trajectory. +- The repaired packed kernel now explicitly selects the recurrent verifier's + launch geometry and arithmetic order, including `/sqrt` normalization and + the same decay exponential helper. Gating remains precomputed in its original + dtype. On a physical V100, all eight production-shape parity cases + (`Hq=4`, `Hv=12`, `K=V=128` per TP4 rank; B1/B2; draft 3/7; FP16/FP32 + recurrent state) are bitwise equal for both output and persistent state. The + per-step low-precision state reload case also passes, and the complete shared + fused-sigmoid/recurrent file reports 27/27 passing tests. +- Before the arithmetic repair, the fully fused greedy one-request diagnostic + kept the exact token hash and acceptance length while improving steady decode + 99.15 to 103.45 tok/s (+4.34%) and decode time 1.2809 to 1.2277 s (-4.16%). + That number is retained only as evidence that this kernel family can reduce + work; it is not the accepted result for the repaired packed leaf. The repaired + leaf still needs a matched unprofiled B1 run and a control/candidate Nsight + pair before default enablement. +- The statistical quality gate is broader than greedy identity. Greedy hashes + remain a sensitive diagnostic, while default enablement also requires bounded + fixed-text logprob/perplexity change and no regression on fixed subsets of + GSM8K, MATH-500, HumanEval, and MBPP. The retained datasets and SHA256 values + are `GSM8K 3730d312...`, `MATH-500 35dc4108...`, + `HumanEval 2f2871a1...`, `MBPP e9e9efa2...`, and WikiText-2 validation + `204929b7...`; all live under + `/data/models/v100-dflash2-20260820/datasets/` and are not repository assets. +- The next Nsight pair must preserve TP4, Flash-V100, target E5M2 KV, draft + FP16 KV, block eight, single-request Graph execution, and the fixed request. + It will report draft Graph, draft-to-target metadata/input preparation, + target verifier Graph, rejection/target-to-draft, full round time, Graph node + count, rank-critical GPU service, and TP GPU-sum service. The existing control + trace lacks supported replay NVTX ranges and is therefore not fed through the + generic per-token parser; control and candidate will be recaptured with the + same replay markers rather than manufacturing a partial comparison. +- A stricter re-audit of the retained endpoint JSONs intersects the long-output + requests by `dataset_index` and requires both implementations to emit at + least 512 tokens. Across those 11 matched requests, vLLM uses 40.7194 ms per + verification round (`decode_time / num_drafts`) and SGLang uses at most + 26.7101 ms (`request_elapsed_seconds / spec_verify_ct`), leaving a gap of at + least 14.0093 ms. The SGLang field includes prefill and is therefore an upper + bound on its decode-round cost. The earlier 40.7148/26.8166 ms values use the + independently selected 13/14-request long-output subsets and remain useful + aggregate references, but they are not a strictly paired endpoint statistic. +- Re-parsing the existing node traces by complete round and taking the slowest + TP rank for every interval gives the following diagnostic wall table. These + values include graph-node tracing overhead and must not replace the endpoint + numbers above. + + | traced phase | vLLM mean / p50 / p90 / p99 (ms) | SGLang mean / p50 / p90 / p99 (ms) | vLLM / SGLang launches per rank | + | --- | --- | --- | ---: | + | draft Graph | 4.065 / 4.028 / 4.073 / 4.088 | 3.961 / 3.928 / 4.123 / 4.193 | 185 / 158 | + | draft to target | 18.382 / 17.660 / 21.258 / 22.690 | 3.517 / 3.423 / 3.660 / 4.066 | 431 / 30 | + | target Graph | 24.835 / 24.880 / 24.943 / 24.995 | 19.723 / 19.600 / 19.939 / 20.526 | 2612 / 1130 | + | target to next draft | 2.791 / 2.793 / 2.806 / 2.809 | 4.216 / 4.188 / 4.332 / 4.402 | 63 / 125 | + | complete traced round | 49.860 / 49.145 / 52.410 / 54.266 | 31.212 / 31.038 / 31.905 / 31.906 | n/a | + + Per-phase critical-rank means are diagnostic and are not strictly additive; + the complete-round row is measured directly from draft-Graph start to the + next draft-Graph start. +- The 431-kernel vLLM draft-to-target interval now has an exact launch + accounting. Repeated GDN state-contract compaction and its generic + indexing/copy/scan helpers consume 423 launches and about 1.276 ms of + rank-average GPU service; block-table, slot, RoPE, and input mapping consume + seven launches and about 0.044 ms; the final TP4 reduce/synchronization is one + launch and averages about 0.598 ms. That reduce is primarily a rank-skew wait + at this boundary (0.128 ms p50, 1.339 ms p90, 6.241 ms p99), not a stable + 0.598 ms compute kernel. SGLang reaches the target with 30 launches and about + 0.115 ms GPU service. This makes removal of the ten repeated group contracts + the first optimization, before tuning the boundary collective itself. +- Inside the target Graph, both implementations spend essentially the same + time in 256 target FP8 GEMMs (10.850 versus 10.570 ms of rank-average GPU + service). vLLM's excess is concentrated in 1,422 generic vectorized, + unrolled, and elementwise copy nodes (4.621 ms) versus SGLang's 211 nodes + (0.667 ms), plus the TP reduction path (2.350 ms on vLLM's slowest rank versus + 1.603 ms on SGLang's slowest rank). Attention and KV-cache service are close. + SGLang includes roughly 0.979 ms of target LM-head GEMM and 0.129 ms of TP + all-gather inside its target Graph, whereas vLLM launches the corresponding + work after the graph; semantic comparisons must retain this boundary + difference. +- With the repaired corpus-wide pooled acceptance length of 4.5644, the + current 40.7148 ms round cost corresponds to about 112 emitted tokens/s + before request-level residuals. A 32 ms round is only an intermediate gate + (about 143 tokens/s), and SGLang parity near 26.8 ms is about 170 tokens/s at + the same pooled acceptance. Reaching the requested 200 tokens/s without + relying on a higher acceptance length requires approximately 22.82 ms per + round. The optimization sequence is therefore persistent/fused GDN metadata, + removal of target-Graph copy/state fan-out, then the TP4 reduction path; + proposal/selector work is not on the critical gap. +- SGLang PR #26520's final-state recomputation path is not selected for this + B1 lane. Its own acceptance-aligned result is 2.5% lower throughput than the + cached-state control. vLLM already writes speculative intermediate states to + state-pool slots, so the nearer-term path is shared/persistent metadata and a + numerically equivalent verifier kernel. ReplaySSM remains a later high-batch + research lane, not a substitute for the current single-request verification + cost gate. +- The first persistent pointer-table implementation has now passed an isolated + SM70 microbenchmark at the production block-eight width and ten target GDN + cache groups. It overwrites group-specific state IDs plus every shared live + and padded field in one Triton launch and remains bitwise equal to the legacy + contract after poisoning all persistent outputs. On an otherwise-idle V100, + the B1 legacy fan-out uses 400 CUDA launches and 5.360 ms synchronized-wall + p50 versus one launch and 0.354 ms for the fused path, a 15.13x speedup and + 5.006 ms local saving. B2 is 5.360 versus 0.358 ms and B4 is 5.346 versus + 0.354 ms, both about 15x. A separate lower-bound probe measures the raw fused + kernel at 0.056 ms wall / 0.040 ms GPU p50; roughly 0.305 ms remains in Python + eligibility, pointer validation, and metadata-view construction. These are + microbenchmark results, not yet an endpoint round-cost claim. The benchmark + is `benchmarks/benchmark_sm70_dflash2_gdn_metadata.py`; an unprofiled TP4 + fixed-trajectory run remains required before default enablement. +- The persistent descriptor now also caches only the shape-specific + `GDNAttentionMetadata` views. Their backing state IDs, accepted counts, + selectors, masks, and query offsets are still overwritten by the fused + kernel on every replay. A poison-and-replay test proves that a same-shape + call reuses the view objects while observing new tensor contents, and a B2 + to B1 transition proves that a shape change rebuilds the views. On the same + B1/ten-group/block-eight microbenchmark, the cached fused path is 0.131 ms + synchronized-wall p50 versus 0.383 ms before view caching; the remaining + Python delta over the 0.053 ms raw-kernel wall is about 0.078 ms. Reinstating + the production scalar host fence costs only 0.028 ms p50, so it is not the + source of the former multi-millisecond Draft-to-Target gap and remains in + place pending the end-to-end trajectory gate. +- The first matched TP4 endpoint trace now validates that the microbenchmark + saving survives the real MRV2 DFlash2 path. Both runs use source + `11a190397abd7209ed1d3aeeeb8ed691e06087f7`, the same fixed shuffled GSM8K + request, Graph target and draft, block eight, probabilistic sampling, target + E5M2 KV, and draft FP16 KV. The control omitted the parent fast-path gate and + therefore exercised the generic metadata builder; the candidate enabled both + `VLLM_SM70_DFLASH2_VERIFY_FASTPATH=1` and + `VLLM_SM70_DFLASH2_FUSED_GDN_METADATA=1`. + - Control: 4.820868 s decode / 103 verification rounds = 46.8045 ms per + round, 88.99 steady-decode tok/s, and 86.11 aggregate tok/s. + - Fused candidate: 3.573145 s / 102 rounds = 35.0308 ms per round, 120.06 + steady-decode tok/s, and 114.20 aggregate tok/s. + - This is an 11.7737 ms (`25.15%`) reduction in traced endpoint round cost, + `+34.92%` steady-decode throughput, and `+32.61%` aggregate throughput. + The 430 output token IDs and SHA256 + `c3aa4bdf534d580213871cb17d31ee997db9cf2ae19ae4766646fe92d463a225` + are exactly equal and the answer remains correct. Acceptance length does + not regress: it moves from 4.1748 to 4.2157 (`+0.0409`). + - The candidate trace contains 412 fused metadata launches across TP4, + corresponding to 103 captured steps per rank (the request reports 102 + draft/verification rounds). Generic CUB `DeviceReduce` and `DeviceSelect` + disappear (`20,800 -> 0` each); `DeviceScan` falls from 20,008 to 19,816, + leaving scans owned by other paths. CUDA runtime fan-out falls from + 27,060 to 2,492 stream synchronizations, 90,524 to 16,328 async copies, + and 205,752 to 88,928 ordinary kernel launches across the four workers. + This proves that the wall saving is primarily removal of the repeated + host-synchronized metadata fan-out rather than the fused kernel's own + 1.528 ms aggregate GPU service. + - These endpoint values were collected under low-overhead Graph-level + Nsight tracing and are the accepted matched A/B for this leaf. The short + node-level trace below updates the semantic Draft-to-Target and + target-verifier phase table; it is used for composition only because node + tracing inflates absolute latency. + Artifacts are + `/data/minimax-h3/task-cache/v100-dflash2-pr257-gdn-metadata/results/dflash2-fused-viewcache-b1-graph-o512-{v2,v4}.json` and + `/data/minimax-h3/task-cache/v100-dflash2-pr257-gdn-metadata/profiles/latest-e2e/dflash2-fused-viewcache-b1-graph-v4.{qdstrm,nsys-rep,sqlite}`. +- A follow-up 128-token node trace at source + `9b8af0dbbe1d48214a555a96bbe762d26524e87c` supplies the missing semantic + phase split. The trace contains 32 complete draft-to-draft intervals; after + dropping the first and last edge intervals, 30 steady rounds remain. For + every phase and round the table takes the slowest of the four TP ranks, which + matches the earlier vLLM/SGLang audit methodology. + + | traced phase | old vLLM mean (ms) | fused vLLM mean / p50 / p90 / p99 (ms) | SGLang mean (ms) | fused launches per rank | + | --- | ---: | ---: | ---: | ---: | + | draft Graph | 4.065 | 4.046 / 4.048 / 4.078 / 4.090 | 3.961 | 185 | + | draft to target | 18.382 | 4.934 / 4.855 / 5.049 / 6.426 | 3.517 | 153 | + | target Graph | 24.835 | 24.740 / 24.747 / 24.892 / 24.907 | 19.723 | 2,612 | + | target to next draft | 2.791 | 2.838 / 2.840 / 2.849 / 2.934 | 4.216 | 63 | + | complete traced round | 49.860 | 36.300 / 36.257 / 36.424 / 37.827 | 31.212 | n/a | + + The persistent/fused metadata path therefore removes 13.448 ms (`73.16%`) + and 278 launches per rank (`64.50%`) from Draft-to-Target. The remaining gap + to SGLang in that phase is 1.417 ms and 123 launches, while the complete + traced-round gap is 5.088 ms. The dominant residual is now the target Graph: + 24.740 versus 19.723 ms (`+5.017 ms`), not the 4.046 ms draft Graph. +- The remaining 153 Draft-to-Target eager kernels per rank are stable across all + 30 steady rounds. They include one fused GDN metadata launch (0.0037 ms GPU + service), one TP4 one-stage reduce (0.2357 ms), 40 CUB scan/init launches + (0.1161 ms), 22 direct-copy launches (0.0686 ms), 20 index-select launches + (about 0.0926 ms), and 20 `compute_cuda_kernel` launches (0.0614 ms), + plus small slot/block/RoPE/input-mapping kernels. Critical-rank GPU service is + only 0.989 ms on average inside the 4.934 ms wall interval, so the next + Draft-to-Target leaf should share/capture the residual common-attention and + input-mapping fan-out rather than tune the 3.7-us fused kernel. In parallel, + the larger end-to-end opportunity is the verifier Graph's generic + elementwise/copy/state fan-out and TP4 reduction path. + Artifacts are + `/data/minimax-h3/task-cache/v100-dflash2-pr257-gdn-metadata/results/dflash2-fused-viewcache-b1-nodes-o128-v5.json` and + `/data/minimax-h3/task-cache/v100-dflash2-pr257-gdn-metadata/profiles/latest-e2e/dflash2-fused-viewcache-b1-nodes-o128-v5.{qdstrm,nsys-rep,sqlite}`. +- `dnv2003/v100-skinny` is useful as an optimization inventory, but its published + 219.1 tok/s result is a Qwen3.8 mixed-NVFP4/FP8 MTP-k7 route on its 1Cat-vLLM + 1.2.2 fork, not a DFlash2 result and therefore not a direct baseline. Its + chain-MTP GDN builder removes 21 device synchronizations and about 70 copies + for a reported 1.4 ms/step saving; PR257's persistent pointer table is the + stronger version of the same principle and now has direct DFlash2 evidence. + The fork's lagged CUDA-event/NVTX phase profiler is worth adapting to MRV2 so + verify, rejection, and proposal can be measured without adding a per-round + synchronization. Its partition pin is not active leverage at this run's + `max_model_len=8192`, and its FP16-KV win is specific to a route where FP8 KV + fell back to a scalar path; Flash-V100 reports the optimized E5M2 cache path + here, so any KV change still requires a matched A/B. The QPN8 verifier idea is + relevant, but the independently adapted block-FP8 implementation in Draft PR + #258 currently rejects speculative configurations. It must pass a dedicated + DFlash2 M=8 contract and quality gate before stacking onto this PR. + Main implementation priority: 1. Add a decoupled SM70 TurboMind backend for AWQ and FP8 base GEMM. @@ -1326,6 +1896,100 @@ Source behavior: - Decode path reads vLLM paged KV cache directly. - Prefill has direct and fallback paths, including paged KV gather fallback. +D256 8K-by-40K..128K GQA architecture checkpoint, 2026-08-25: + +- Frozen single-rank shape is causal FP16 Q/K/V, `Q=8000`, `KV=128000`, + `Hq=6`, `Hkv=1`, `D=256` on one V100-SXM2-32GB. The metric is useful + causal Attention FLOPs divided by complete Attention elapsed time; it is not + model TOPS or prompt throughput. +- The architecture separates a fully visible 32K..120K prefix from the causal + 8K tail. Six GQA heads are packed into GEMM M=48,000. QK writes FP16 logits and + row statistics; PV normalizes logits during its global-to-shared operand + transform. One reusable FP16 block partial is folded into a FP32 online + prefix numerator/state before the next block. The exact causal-tail kernel + exports max/sum state and a final kernel merges prefix and tail. +- The final BN8192 score-cache layout retains only the 750-MiB FP16 score + matrix per device. All other approximately 100 MiB of state is packed into + one transient aligned byte slab. A per-device event and mutex serialize + score reuse across streams. Before its first allocation, the route requires + driver-visible free memory for both workspaces plus 128 MiB of downstream + headroom; failure raises a typed OOM without touching the caching allocator + or launching a kernel, so the Python path can fall back safely. +- The final frozen-extension Torch A/B/A measured the existing exact operator + at `140.00213 ms / 43.53414 TFLOP/s` and the complete architecture at + `100.76979 ms / 60.48313 TFLOP/s`: `28.0227%` lower latency and `1.38933x` + speedup. All 12,288,000 outputs are finite; versus the exact operator, + max absolute difference is `2.2888e-4`, mean absolute difference + `2.5430e-5`, and relative L2 `6.8629e-3`. Peak allocated memory is + `1.11960 GiB`. +- The generalized operator admits all twelve observed `Q=8000` chunk shapes, + `KV=40000..128000` in 8000-token steps. Strict per-length A/B measured + `59.34862..60.74103` useful causal TFLOP/s. Against the active split-KV3 + control, every point is faster: latency is `22.3891%..26.1909%` lower and + speedup is `1.28848x..1.35485x`. All 147,456,000 candidate elements are + finite; the worst max absolute difference is `5.0354e-4` and worst relative + L2 is `6.8983e-3` versus the exact FP32-accumulator route. +- Smaller score blocks do not pass the target: BN7168 is + `104.32751 ms / 58.42057 TFLOP/s`; BN8000 is + `101.70283 ms / 59.92825 TFLOP/s` and must not be rounded to 60. BN8192 is + therefore retained while model KV reservation supplies downstream headroom. +- QK and PV compile at 254 and 119 registers/thread with no spill, stack, or + local storage. The complete operator is loadable through `torch.ops`; the + initial namespace registration mismatch was caught by a static load test and + fixed before timing. Synchronous device-symbol updates were changed to + current-stream asynchronous copies so the integrated measurement includes no + artificial whole-device fence. +- Integration is shape-family bounded and default-on. The legacy + `VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL=0` setting is the + rollback. It rejects CUDA-graph capture and falls back to the exact dense + kernel on workspace OOM. + The final external FA2 patch applies cleanly after the existing D256 + pipeline, split-KV3, and K-ping-pong patches, and a fresh CUDA 12.8 / Torch + cu128 SM70 build passes its operator-schema load test. The full Flash-V100 + policy file passes 112/112 tests on a real V100, including direct typed + architecture-OOM fallback to the exact dense route; Ruff and + `git diff --check` pass. +- Early real-model memory tests at `gpu_memory_utilization=0.88` rejected the + fully persistent and pre-preflight score-cache layouts after their first + successful route call: the following 80-MiB PYNCCL all-reduce had only + 62-74 MiB free. This is a downstream-headroom failure, not an Attention + compute failure. +- The final matched `.875` TP4 Qwen3.8-27B-FP8 + control/candidate/control prefill times are `46.45658 / 45.47797 / + 46.11195 s`. Against the `46.28426-s` bracketed control, the candidate is + `1.7421%` lower latency and `1.01773x` faster (`2831.89 -> 2882.10` + prompt token/s). Every candidate rank logs 16 architecture hits and no + architecture OOM/fallback. All three runs produce the same 32 output token + IDs and SHA256 + `df4fee7f5f0126fe6b391fe77b4fc19667831de5ef55fd69c28c2f52a3d7086e`. +- A final `.88` run also completes in `45.53427 s`, with the same token hash + and 16 route hits/rank. Its profiler left enough headroom for the route, so + this validates high-memory route success rather than directly exercising + the preflight fallback branch. +- The widened TP4 Qwen3.8-27B-FP8 gate uses a `46.11195-s` control before the + candidate and a `46.09222-s` control after it. Against their `46.10208-s` + mean, the `41.51191-s` candidate lowers prefill latency by `9.9565%` and + raises prompt throughput from `2843.08` to `3157.46` token/s (`11.0575%`). + Every rank logs exactly 192 family-route hits with no architecture OOM or + fallback. Both controls and the candidate emit the same 32 token IDs and + SHA256 + `df4fee7f5f0126fe6b391fe77b4fc19667831de5ef55fd69c28c2f52a3d7086e`. + Token identity is route-health evidence rather than a promotion + requirement. The stronger numerical gate covers 147,456,000 finite output + elements with worst max absolute error `5.0354e-4` and relative L2 + `6.8983e-3`, inside the accepted FP16 envelope of `1e-3` and `1e-2`. + Artifacts and the experiment ledger are retained outside Git; their private + filesystem location is intentionally not committed. +- The 2026-08-26 merge audit replayed all four patches from the locked clean + vendored baseline and rebuilt `_vllm_fa2_C` with CUDA 12.8 for `sm_70`. + QK/PV again compiled at 254/119 registers with zero spill. Current-source + direct A/B measured `39.29395 -> 30.27579 ms` (`1.29787x`) at KV40K and + `136.07628 -> 100.20284 ms` (`1.35801x`) at KV128K. Both outputs were fully + finite. KV40K had max absolute error `5.4932e-4` and relative L2 + `6.6098e-3`; KV128K had max absolute error `2.2888e-4` and relative L2 + `6.9166e-3`. These pass the FP16 merge envelope without imposing greedy or + bitwise identity and justify keeping the validated shape route default-on. + Latest target state: - No `FLASH_ATTN_V100` enum. @@ -41671,3 +42335,1418 @@ Interpretation: formatter changes from current main and keep large C/CUDA formatter work separate from accepted kernel repairs. Small policy substitutions still need their own current-source static/argument-equivalence checks before merge. + +## 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 + `MIXED_PRECISION`: dense projections use FP8, routed/shared experts use + `W4A16_NVFP4`, and 19 `mtp.*` tensors are present and unquantized. Exact SM70 + admission now requires both FP8 and NVFP4 TurboMind gates; unsupported model + shapes, TP8, EP, and all-to-all fail closed rather than taking a generic + fallback. +- The native Qwen3.6-35B-A3B MoE contract is H2048, expert I512, 256 experts, + top-k 8, replicated experts, and TP1/2/4. B1-B10 decode uses compact routed + slot groups. B11-B18 CUDA Graph/MTP decode and larger prefill use persistent + buffers with the full 256-expert grouped TurboMind operation. The grouped + prefill policy is default-on and may be disabled only for diagnosis with + `VLLM_SM70_NVFP4_MOE_GROUPED_PREFILL=0`. +- A real layer-0 TP4 checkpoint oracle matched the previous per-expert route: + M4096 W13/W2 were bitwise identical and improved 20.39x/9.48x. Full-grouped + B9 and B18 had cosine 1.0, maximum absolute error at most 1.221e-4, and + 21.59x-35.21x isolated-stage speedups. A rejected 144-group experiment + coalesced repeated expert IDs even though the forced-one-row scheduler + requires routed slots to remain independent; its roughly 0.495 W13 + relative-L2 error and catastrophic W2 output do not establish a general + 64-group correctness limit. The corrected slot-independent B9/B10 evidence + is recorded in the concurrency section below. +- Matching TP4, FP8-E5M2 KV-cache, Flash-V100/FlashQLA, full CUDA Graph, + input-4096/output-1024, max-length-8192 runs measured AWQ at 0.3813 s prefill + and 113.71 token/s steady decode. Default-grouped NVFP4 measured 0.4216 s and + 116.99 token/s: prefill is 10.56% slower than AWQ while decode is 2.89% + faster, placing both in the same performance class. Grouping reduced NVFP4 + prefill from 1.5479 s to 0.4216 s (3.67x). +- MTP4 resolved the bundled draft as `Qwen3_5MoeMTP`, captured FULL CUDA Graphs + through B18, and measured 0.8379 s prefill plus 174.76 token/s steady decode + under the same model contract. It accepted 792/924 draft tokens, with mean + acceptance length 4.43, and improved no-MTP decode by 1.49x. MTP prefill is + slower because the drafter participates. That original run still exposed + first-request sampler/Mamba JIT, which the later cold-start work below moves + ahead of request processing. +- The accepted cold-start patch warms the exact Mamba speculative postprocess, + all-greedy rejection sampler, and unquantized MTP-drafter MoE signatures. For + the checkpoint's 256-expert/top-8 draft, token counts 9, 33, and 257 cover the + sorted-assignment, unaligned-buffer, and SM70 small-to-large tile boundaries. + A fresh-cache ShareGPT16 run then produced no inference-time Triton JIT + warning: 97.69 end-to-end output token/s, 242.72 serial-prefill token/s, + 120.09 pure-decode token/s, and 8.827 ms mean TPOT. Moving these compilations + into startup adds about 4.33 s to model load versus the preceding + Mamba/sampler-only candidate. +- MTP quality is gated by dataset accuracy, validity, and repetition health, + not token equality with no-MTP. On pinned GSM8K test records 0-127 with + five-shot chat prompting, thinking disabled, greedy sampling, and a 512-token + cap, no-MTP scored 122/128 and three full-forward MTP4 runs scored 121, 122, + and 121. The post-rebase final at `d3b1ea3678` scored 121/128 with zero invalid + or repetitive records and the same wrong set as the preceding full MTP run: + `[2, 12, 37, 87, 89, 102, 119]`. Record 37 is a cap-induced truncation and is + correct at 1024 output tokens, so effective MTP4 accuracy is 122/128. A + standard-GDN diagnostic scored only 108/128 and is rejected. +- The bounded B10 source at `5422173e99` passes the same pinned GSM8K-128 + MTP4 gate at 122/128 (95.3125%), with zero invalid answers, zero repetitive + records, and wrong indices `[2, 12, 87, 89, 102, 119]`. It generated 19,434 + output tokens in 135.967 s. This is dataset-level evidence that the B10 + extension does not regress the accepted MTP4 quality band; no-MTP token + equality is not used as its oracle. +- A final default-grouped, no-MTP natural-text gate used two identical Chinese + prompts and one Python prompt. All three requests reached EOS with `stop`; + the duplicate outputs were token-for-token identical, and the code response + produced a correct `add(a, b)` implementation. This supersedes the earlier + 128-token truncated quality smoke for route acceptance. +- Reproducible JSON, logs, numeric oracles, and the rejected compact experiment + are under + `bench_results/modelopt_mixed_nvfp4_moe_20260822/{results,logs,telemetry}`. + The final dataset artifact is + `results/tp4_mtp4_fp8kv_gsm8k_chat_nothink_5shot_128_coldstart_rebased_final_graph.json`; + the all-JIT-warm ShareGPT artifact is + `results/tp4_gpu0123_mtp4_fp8kv_sharegpt16_seed20260822_all_jit_warmup_reserved_graph.json`. + The bounded-B10 dataset artifact is + `bench_results/qwen36_nvfp4_concurrency_scaling_20260824/results/candidate_final_mtp4_gsm8k_chat_nothink_5shot_128.json`. + These results are the accepted 35B evidence; earlier short route-hit and + quality smokes must not be substituted for the matched speed contract. + +## 2026-08-24 Qwen3.8 exact-D256 long-prefill K ping-pong + +- The Qwen3.8-27B-FP8 TP4 128K trace remains attention-led: exact-dense and + direct-paged D256 attention account for 35.55% and 6.22% of profiled wall. + The exact long shape sustains only 38.67% tensor-pipe active with 62.14% of + cycles lacking an eligible warp; DRAM is only 9.06%. This is an SM70 + shared-memory/dependency problem, not HBM saturation. +- Added a third vendored-FA2 patch that alternates the four D64 K panels across + two non-overlapping shared stages. The second stage borrows the later V/P + allocation because K and V/P lifetimes do not overlap. Its 2240-half stride + is above the padded 2188-half K `cosize`, aligned for 128-bit stores, and on + a common 128-byte bank phase. This removes three static full-CTA barriers + while preserving every HMMA/load/store count and the exact N32 arithmetic + order. +- Shared memory remains 45,568 bytes. Dense and split-KV3 remain at 254/253 + registers per thread with zero stack or spill. A clean CUDA 12.8 SM70 build + has byte-identical SASS to the measured candidate. The first clean-clone + build failed only because the CUTLASS submodule was absent (`cute/tensor.hpp` + not found); initializing pinned CUTLASS `62750a2b...` repaired the build. +- Exact dense A/B/A improves 0.74%-1.11% at Q8192/KV32K-128K and 0.54% at the + historical Q15680/KV125440 shape; a separate clean-build 128K run improves + 0.39%. The first Q8192/KV8192 chunk is neutral. Dense, randomized-paged, and + split-KV3 output gates are all bitwise equal to their matching controls. + Two order-reversed 128K comparisons pool to `142.4068 ms` control, + `141.8076 ms` fully-disjoint v5, and `141.7669 ms` lifetime alias. The final + form is `-0.45%` versus control and preserves the smaller shared envelope. +- Rejected 2048-half spacing because K stages overlap, 2188 because `STS.128` + becomes misaligned, and 2192/2304 because they are slower than the final + same-bank-phase layout. The exact fully-disjoint 2240 form is superseded + because it grows shared memory without a measurable speed advantage. Nsight + Compute counter collection is unavailable to this user + (`ERR_NVGPUCTRPERM`); no driver security setting was changed. +- The first full-model attempt failed before any request during FULL-graph XQA + metadata warmup. The source passed the new 19th + `batch_context_max_seq_len` argument, while the stale system Flash-V100 + binary exposed the old 18-argument ABI. Rebuilding the unchanged current + Flash-V100 source produced SHA `ac07dc9d...5580a8`; its 19-argument docstring + and a CUDA-graph E5M2 XQA scalar/padded bitwise gate pass. This harness repair + was used identically for both control arms and the candidate. +- Matched TP4 Qwen3.8 128K E5M2-KV, chunk-8192, official-sampling A/B/A passes. + Baseline A/candidate/baseline B prefill times are + `47.8992/47.6687/47.9475 s`. The bracketed control is `47.9233 s` or + `2735.0 tok/s`; the candidate is `47.6687 s` or `2749.6 tok/s`, giving + `-0.531%` latency and `+0.534%` throughput. TTFT changes + `47.9498 -> 47.6937 s` (`-0.534%`). Pure decode is neutral at + `5.1557 -> 5.1547 s` and `49.460 -> 49.470 tok/s`. +- All arms return 256 tokens with identical hash `2db7ef09...503325a`. + All four ranks report the same route counts: exact-dense D256 32, + prefix-paged 544, E5M2 bridge 544, scalar decode 176, and XQA decode 256. + The endpoint promotion gate is satisfied; the result is deliberately cited + as a small dependency-chain gain, not as completion of the broader 128K + utilization target. + +## 2026-08-24 Qwen3.8 Q8000 split-KV3 long-prefill experiment + +- A runtime-shape diagnostic closed the initial Q8192 assumption. With + chunk-8192 plus Mamba align, long prefix calls are actually Q8000; the first + Q8192-gated endpoint run had zero split-KV3 hits and is not performance + evidence. One control was additionally invalidated when another TP4 job + started between its 47.90-second warmup and 64.45-second measurement. +- Added an exact-shape, default-off + `VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3_Q8000_EXPERIMENTAL` policy. Existing + production Q4096 admission is unchanged; Q8192 and nearby shapes remain + rejected. The policy suite passes 89/89 and Ruff lint/format checks pass. +- Bracketed Q8000 operator sweeps at KV 40K/64K/96K/128K improve throughput by + 3.94%/2.85%/3.45%/3.36%. At KV128K the kernel moves from 138.9150 to + 134.3961 ms and 43.87 to 45.35 causal TFLOP/s. Workspace is 141.72 MiB per + rank. The FP16 outputs are repeat-stable but non-bitwise versus unsplit; + maximum absolute error is at most `3.052e-5` and mean error at most + `1.38e-6`. +- Matched TP4 128K A/B/A measures control A/candidate/control B at + `47.5940/47.1074/47.6612 s`. The bracketed control is `47.6276 s` or + `2752.0 tok/s`; candidate is `47.1074 s` or `2782.4 tok/s`, for `-1.092%` + latency and `+1.104%` throughput. TTFT improves 1.078%; decode is neutral. + All ranks report exactly 384 Q8000 split-KV3 hits and unchanged surrounding + prefill/decode route counts. KV capacity remains 2,424,439 tokens per rank. +- Each lane's warmup and measured 256-token stream is internally exact. The + candidate and control B are also token-for-token identical, but identical + unsplit control A selects a different stable sampled stream under the same + official seed. This cross-process sampling instability is not equivalence + evidence for a non-bitwise route. +- A deterministic greedy follow-up uses two identical Chinese exact-128K + prompts and one exact-128K Python task. Both routes are internally stable and + coherent, and both code responses stop normally with the same correct + `add(a, b)` implementation. Cross-route strict identity fails: Chinese first + differs at token 7 with 126/128 tokens aligned; Python first differs at token + 20 with all 67 control tokens aligned in the 70-token candidate response. + Candidate ranks each report the expected 576 split-KV3 hits and controls + report zero. Keep the Q8000 gate default-off until fixed-text + logprob/perplexity or dataset-level evidence proves no quality loss from the + classified Type-B reduction-order drift. Evidence is under + `/data/minimax-h3/task-cache/qwen38-fp8-128k-flashattention-20260824/`. + +## 2026-08-23 Qwen3.8-27B-FP8 TP4 no-MTP QPN8 decode acceptance + +- Integration base is `7aede2cf01`; the owned source worktree and Draft PR are + `codex/v100-qwen38-fp8-qpn8-operator-race-20260822-114943` and PR 258. +- Added a memory-neutral SM70 QPN8 weight layout and direct decode operators for + the accepted TP-local gate/up `(K,N)=(5120,8704)`, down `(4352,5120)`, and + output `(1536,5120)` projections. GDN input and full-attention QKV remain on + TurboMind. The old layout is replaced at load time rather than retained. +- Opt-in admission uses the engine contract only: SM70 block-FP8 layout, TP4, + exact projection shapes, no MTP, and configured `max_num_seqs <= 8`. + Native QPN8 handles M=1-8. + Wider-concurrency and MTP configurations retain TurboMind until their own + performance and quality gates pass; this is not a batch-one-only dispatch. +- The opaque C++ M dispatch prevents an AOTInductor M=1 trace from being reused + for prefill. M>8 remains functionally correct through one bounded 85 MiB FP16 + weight workspace, but fused gate/up additionally creates an `M x 8704` FP16 + temporary and is not accepted as a high-concurrency performance route. +- Complete-source no-MTP E5M2 TP4 A/B at input 1024 and output cap 256 measured + TurboMind at `16.961295 ms/token`, `58.958 tok/s`, and QPN8 at + `16.347973 ms/token`, `61.170 tok/s` (`+3.752%`). The candidate stopped + normally at token 233; the control finished all 256 tokens. An earlier + matched operator-integration run measured `15.808000 ms/token`, + `63.259 tok/s` (`+7.448%`) and remains headroom evidence, not the conservative + complete-source observation. Because these arms contain 255 versus 232 + steady intervals, respectively, this is not a pure-FP8 default-promotion + result; an equal-interval same-source rerun remains required. +- Nsight Systems attributes about `0.944 ms/token` of the trace improvement to + dense service (`10.335 -> 9.391 ms`), while launch count and GPU idle time are + unchanged. The QPN8 dense service spends 6.172 ms in the 75-90% DRAM band and + 7.059 ms in the 25-50% SM band; no dense kernel sustains at least 90% DRAM or + 50% SM. Useful dense arithmetic rises from 1.177 to 1.295 TFLOP/s per GPU. +- The quality gate passed without a measured regression: WikiText PPL/byte PPL + and bits/byte improved slightly, GSM8K flexible accuracy was equal and strict + accuracy improved by one item, all measured MMLU subjects were equal or + better, and all measured C-Eval subjects were equal. Sampled token identity + is therefore not used as the sole acceptance criterion. +- Full evidence and artifact paths are recorded in + `docs/design/sm70_qwen38_qpn8_decode.md`. Next work is a measured single-launch + M=9/11/16 row-tile route, followed by larger concurrency crossovers; do not + compose an eight-row launch with a separate tail or promote the current + large-M temporary as an accepted throughput implementation. + +## 2026-08-23 Qwen3.8-27B-NVFP4 TP4 decode recovery + +- Recovered the mixed NVFP4/channel-FP8 checkpoint without crediting Marlin: + accepted FP8 projection shapes use QPN8, remaining FP8 projections and the + LM head use TurboMind W8A16, and NVFP4 uses TurboMind W4A16. +- Added channelwise FP8 TurboMind packing, exact NVFP4 N32 selectors, E4M3 + Flash-V100 XQA, QPN8 split-12 for the K1536 x N5120 output projection, and + an exact B1/G6/D256 E4M3 p64 attention partition. +- Frozen TP4 I1024/O256, E4M3 KV, official random sampling, no-MTP, and full + CUDA graph decode first measured 71.732 tok/s at 13.941 ms/token with the + native sampler. After merging current `onecat/main` at `675a12dedc` and + rebuilding Flash-V100, the frozen request reconfirmed **71.342 tok/s** at + **14.017 ms/token**. The clean p128 control is 69.904 tok/s at + 14.305 ms/token. +- The p64 operator sweep saves 9.55-16.98 us per attention launch from sequence + lengths 1025-2049. All p64/p128 cases stay within one FP16 output ULP of the + scalar reference; the focused GPU suite passes 6/6. +- Compact and fused top-k20 sampler experiments were removed because their + sampled trajectories were not deterministic against the native route. They + are not part of either accepted p64 result. +- Post-merge checks pass: 12/12 focused CPU policy/dispatch cases, 6/6 E4M3 + XQA GPU numerical cases, Ruff lint/format, and `git diff --check`. The final + `_C` and Flash-V100 SHA256 values are `e0ea14d0...7624` and + `b418fed8...dce7`. +- The latest merged p64 Nsight Systems trace captures 63 decode replays on all + four ranks. Steady replay/GPU-union time is 14.621/14.055 ms, GPU activity + union is 96.136%, and the route launches 1140.8 kernels/rank/token. NVML + samples show 99.47% GPU busy-window duty but only 48.26% memory-active duty; + 69.40% of service has a grid-limited occupancy ceiling below 25%. Current NCU + counters remain unavailable because the driver rejects non-root counter + access with `ERR_NVGPUCTRPERM`; do not present NVML duty or effective payload + rates as achieved SM or HBM throughput. +- Full contract, route map, staged results, profile table, numerical evidence, + rejected paths, rollback controls, and artifacts are in + `docs/design/sm70_qwen38_nvfp4_decode.md`. + +## 2026-08-24 Qwen3.8-27B-NVFP4 TP4 no-MTP 80 tok/s acceptance + +- The frozen input-1024/output-256 TP4 contract now passes the requested + no-MTP speed gate. The production CMake C++17 sampler-sidecar run measures + **80.624 tok/s** at **12.403 ms/token**. Three behavior-equivalent hand + sidecar repeats measure 80.177, 80.164, and 80.026 tok/s, including a repeat + on physical GPUs 4-7. Pure decode excludes TTFT and prefill. +- The deployment retains the proven-faster exact-shape QPN8 projections. GDN + input, full-attention QKV, and the LM head continue through TurboMind W8A16; + non-admitted NVFP4 cases retain TurboMind fallback. The accepted batch-one + NVFP4 shapes use native QPN4, E4M3 attention uses Flash-V100 XQA p64, and TP4 + communication uses the pack32 one-stage all-reduce. MTP is disabled. +- The quality gate passes at the frozen baseline: the production sidecar scores + **226/250 (90.4%)** on the first 250 GSM8K questions with five-shot greedy + prompting and has zero invalid outputs. The baseline is also 226/250; the + hand-sidecar result is 227/250. Per-question prompts, texts, token IDs, + extracted answers, labels, correctness, and hashes are retained in the JSON. +- The production sidecar and hand-built validation library select identical + tokens for 100 independent seeds and 256 sequential Philox draws and finish + with the same generator state. This exact implementation supersedes the + earlier rejected compact/fused sampler experiments, which did not preserve + native Philox behavior. +- The accepted compatible main `_C`, production sampler sidecar, and + Flash-V100 SHA256 values are `a0a0cd9...d36d`, `cdbfdd87...9ae`, and + `b418fed8...dce7`. The sampler is intentionally a separate CMake/wheel module + so registering it does not relink the frozen main library. +- The latest 63-replay/rank Nsight Systems trace shows a 13.430 ms mean replay + interval, 12.887 ms GPU activity union (95.958%), 0.543 ms idle, and 1061.9 + launches/rank/token. QPN8 split-16, QPN4 gate/up, TP all-reduce, QPN4 down, + QPN8 output, XQA, QPN8 gate/up, and TurboMind FP8 service contribute 2.213, + 2.165, 1.587, 1.441, 0.915, 0.618, 0.496, and 0.410 ms/token respectively. +- NVML reports 100% busy-window duty, 53.25% memory-active duty, about 28.97 + GiB allocated per GPU, and 56.36% of the 300 W limit. Grid-limited occupancy + ceilings put 86.32% of service below 50% occupancy. These are not achieved + SM/Tensor/HBM counters; Nsight Compute remains blocked by + `ERR_NVGPUCTRPERM`. +- Current-source rebuild closure is complete at source snapshot `7bd1d783e4`. + A clean primary `_C` plus clean C++17 sampler measured 77.314 tok/s in the + cold A1 observation and **80.181 tok/s** at **12.472 ms/token** in A2. A2 + exactly matches the accepted C++17 256-token stream and records + `speculative_config=None` plus `spec_decoding_metrics=null`. +- The clean primary and sampler SHA256 values are `c5d7c6a9...3e56` and + `3058051f...bc74`. Their `.nv_fatbin` SHA256 values, `dbba7402...26e6` and + `9a89e8cc...b2c`, are byte-identical to the accepted binaries. Full ELF + hashes differ only across host/link content: merged main adds two unrelated + ModelOpt NVFP4 MoE host entry points, while the sampler includes host + build-ID differences. These do not enter the accepted dense batch-one graph. +- The clean A2 result SHA256 is `4e6ed330...bb1`; its serialized token-array + SHA256 is `f50e0eba...e32`, exactly matching the accepted production C++17 + run. Device-binary and sampled-stream identity carry forward the already + passed **226/250 (90.4%)**, zero-invalid GSM8K gate. The separate sampler + module remains the packaging boundary, but current-source primary rebuilds + are no longer rejected merely because the full ELF hash changes. +- Focused CMake installation puts the modules at `vllm/_C.abi3.so` and the + CPython-3.12 SOABI sampler filename, strips build-tree RPATHs, and preserves + both accepted fatbins. The installed SHA256 values are `d7ed490e...6749` and + `8ee84663...55e`. Loading both staged files registers QPN4, QPN8, and the + packed top-20 sampler schemas. This closes the relevant install boundary + without rebuilding 50 unrelated wheel-target objects. +- Final focused regressions pass 94/94 target sampler, NVFP4-admission, and + SM70 TurboMind-adapter tests; the post-merge four-file run including the + adjacent ModelOpt route passes 106/106. Ruff lint/format, Python byte + compilation, the sidecar wheel filename gate, and `git diff --check` also + pass. +- Full results, route details, quality artifacts, trace/resource tables, and + exact evidence paths are in `docs/design/sm70_qwen38_nvfp4_decode.md`. + +## 2026-08-24 Qwen3.8-27B-FP8 exact-8K CUTLASS projections + +- The accepted starting point is TP4 V100, input 8000, output 1, FP8-E5M2 KV, + Flash-V100/FlashQLA, FULL_AND_PIECEWISE CUDA Graph, bounded exact-dense FP8 + projections, and the fused TP4 reduce-scatter/Gemma norm/all-gather boundary. + Its request-wall median was `1.616075 s`, or `4950.27 prompt tok/s`. +- The fused graph trace put `506.081/254.947/93.510 ms` into 64 gate/up, down, + and attention-output FP16 GEMMs. Their current useful rates were only about + `90.2/89.5/86.1 TFLOP/s`, despite the V100 SXM2's 125-TFLOP/s Tensor Core + peak. QKVZ plus QKV retained another `241.625 ms` at exact 8K. +- NVIDIA CUTLASS profiler selected its generated SM70 + `f16_s884gemm_f16_128x256_32x2_nn_align8` operation for all five exact TP4 + shapes. It uses CTA `128x256x32`, warp `64x64x32`, native `m8n8k4`, two + stages, and split-K one. The compiled kernel uses 220 registers/thread, + 256 threads, and no stack or local spill. +- Real-checkpoint CUDA Graph operator tests passed bitwise equality. The key + complete-boundary medians improved gate/up `8.67876 -> 8.26086 ms`, down + `4.26622 -> 4.02442 ms`, and attention output `1.54365 -> 1.48992 ms`. + Gate/up reuses the same graph-captured full projection temporary and existing + interleaved SiLU kernel; it adds no resident weight or second workspace. +- The final source-matching accepted extension SHA256 is + `6a7e8493ca04bd7738c9b016a6a0b4914b2c582598c920f99bb67626ae3e11c1`. + It is default-on only at `M=8000` for TP4 shapes + `(K,N)={5120x4096,5120x3584,5120x8704,4352x5120,1536x5120}`. Decode, + tails, other M values, and other shapes retain existing dispatch. Rollback is + `VLLM_SM70_FP8_PREFILL_CUTLASS=0`. +- The source-matching final full-model run produced six steady request-wall + samples `1.559085/1.561460/1.560481/1.562664/1.564697/1.564643 s`. + Their `1.562062 s` median is `5121.44 prompt tok/s` with `0.132%` CV. The + pure prefill median is `1.547100 s`, or `5170.96 tok/s`. This saves + `54.013 ms`, raises wall throughput `3.46%`, and clears the 5000-tok/s + requirement by `121.44 tok/s`. +- Route proof contains exactly 20 one-time messages, five admitted shapes on + each of four TP ranks. The accepted 766 graph-address registration count and + fused collective backend remain intact. Seven of seven official-sampling + outputs remain token `[1061]`, text `" This"`. +- The side-stream two-workspace FP8 dequant prototype is rejected: placing its + two dequantizations across the RMSNorm window changed + `7.98956 -> 8.00819 ms` (`0.99767x`) despite bitwise equality. The + cuBLAS-algorithm-106 candidate + was also rejected. Both implementations were removed; do not repeat them as + substitutes for the accepted CUTLASS shape route. +- Final evidence is under task-cache + `qwen38-fp8-prefill-utilization-20260823/results/`, especially + `full_model_cutlass_all_projections_final_source_tp4_fused_graph_i8000.raw.json`. + +## 2026-08-24 Qwen3.6-35B NVFP4 MTP4 concurrency scaling + +- Audit disposition: the B10 compact-routing extension has its own matched + speed and GSM8K evidence and remains enabled. The broader high-width tune + limit, exact-shape drafter tile, all-reduce block policy, and eight-warp + sampler remain opt-in because their matched formal curve and combined + dataset-quality gate are incomplete. + +- The formal workload uses ModelOpt mixed FP8/NVFP4, TP4 on V100 GPU0-3, + Flash-V100/FlashQLA, FP8-E5M2 KV cache, MTP4 probabilistic sampling with + local argmax, max length/batched tokens 8192, and 48 fixed ShareGPT requests + totaling 9,115 prompt and 10,623 output tokens. Forty-eight requests give + C16 three complete waves and remove the tail-underfill artifact in the + earlier 16-request localization sweep. +- Current-main output throughput at C1/C2/C4/C8/C12/C16 is respectively + 98.265/146.124/276.600/414.315/491.270/542.982 tok/s. Scaling efficiency + relative to C1 is 100.00%/74.35%/70.37%/52.70%/41.66%/34.54%. Mean MTP + acceptance length remains 2.887-2.955 and draft acceptance remains + 47.18%-48.87%, so the scaling loss is not an acceptance collapse. +- The old AWQ scheduling pair is rejected for this route. Shared-expert stream + overlap plus `all_reduce_sum2` regressed corrected C16 output 4.10%; isolated + overlap and sum2 regressed 4.66% and 4.06%. The pair also regressed C1 from + 98.265 to 94.880 tok/s (3.45%). Leave both switches default-off for NVFP4. +- MTP4 verifier width is five times request concurrency. Before this work, + production warmup stopped at NVFP4 MoE M15 while captured C4/C8/C12/C16 + shapes require M20/M40/M60/M80, or 160/320/480/640 routed rows. +- Clean independent-process graph microbenchmarks on one exclusively claimed + V100 show measured routed-MoE W13/W2 speedups at M20/M40/M60/M80 of + 1.107x/1.311x, 1.200x/1.179x, 1.195x/1.181x, and 1.177x/1.117x. Shared-dense + W13/W2 speedups are 1.742x/1.261x, 2.749x/1.181x, 2.004x/1.179x, and + 1.519x/1.348x. All outputs are finite; maximum accumulation-order difference + is 6.104e-5, which still requires the dataset-level quality gate. +- The high-width candidate raises only the NVFP4 MoE tune limit to 640 routed + rows, derives warmup widths from the already-captured CUDA Graph sizes, and + builds balanced 256-expert offsets above 32 verifier tokens. The dense-NVFP4 + M80 extension was isolated and withdrawn; its default remains M16. In a + 48-request C16 candidate/control/candidate sandwich, routed-MoE tuning + measured `535.635/524.376/539.113 tok/s`; the candidate endpoint mean is + `+2.48%`, summed pure decode is `+2.85%`, and P50 TPOT improves about `4.0%` + with stable MTP acceptance. + Production keeps the prior 128-row default; reproducing this candidate + requires `VLLM_SM70_NVFP4_MOE_TUNE_MAX_TOKENS=640`. +- Two earlier low-width target-MoE candidates are rejected. Directly reusing + sorted expert IDs instead of the compact metadata copy regressed the C1 + sandwich; splitting B9/B10 into 64-slot compact chunks was numerically + healthy but `41.5%/10.3%` slower. The old conclusion that every direct + compact execution above 64 slots was invalid was too broad: that experiment + coalesced repeated expert IDs even though TurboMind's forced-one-row + scheduler requires each routed slot to remain an independent group. +- A C1 Nsight graph-node trace identified the bundled unquantized MTP drafter + MoE as a separate `2.825 ms` per-cycle hotspot. An initial oracle incorrectly + used the checkpoint-global I512 shape. The first candidate startup log proved + that TP4 shards it to local I128 (`E=256,N=128` config lookup); the candidate + therefore did not route-hit and its full-model run was stopped before + measurement. The global-I512 speedups are not TP4 selection evidence. The + corrected local E256/H2048/I128/top-k8 M1-M16 graph oracle is bitwise equal + at every width. The isolated M1 operation improved `1.290x` by retaining + N32/K64 and reducing to two warps; M2-M16 select M8/N128/K32/W4/S3 and + improve `1.543x/2.656x/2.879x/2.697x/2.625x` at the formal + C2/C4/C8/C12/C16 widths. +- The route-hit candidate initially exposed request-time old-tile MoE + signatures. The exact E256/H2048/I128/top-k8 TP4 contract can warm tuned + M1-M16 plus M33/M257 + and explicitly covers legacy M1/M2/M9/M10 naive/aligned signatures; a fresh + 48-request run then produced no inference-time `fused_moe_kernel` JIT. + Unmatched shapes retain the old `(9,33,257)` warmup set, and the tuned route + requires `VLLM_SM70_MTP_MOE_TUNED_CONFIG=1`. +- Full-model evidence rejects the isolated M1 result. A physical-GPU4-7 C1 + candidate/control/candidate sandwich measured endpoint-mean summed pure + decode `112.233` versus control `115.451 tok/s` (`-2.79%`) and endpoint-mean + P50 TPOT `8.874` versus `8.758 ms` (`+1.33%`). M1 therefore retains the exact + 0.0.3 tile; balanced single-GPU kernel means are not a substitute for TP4 + rank-imbalance and synchronization behavior. +- The narrowed M2-M16 candidate passes C2. Candidate/control/candidate output + was `154.865/147.180/163.480 tok/s`; endpoint mean is `+8.15%`. Summed pure + decode improves `+4.47%`, P50 TPOT is effectively flat (`+0.14%`), and mean + acceptance is lower than control rather than artificially favorable. M17+ + retains the old prefill tile. +- Every C1 graph all-reduce in the Nsight trace is the exact 10,240-element + FP16 verifier payload. A TP4 operator sweep found that one CTA is faster for + this 20-KiB payload, while eight CTAs are faster only for the exact + C8/C12/C16 verifier payloads; C2/C4 retain the default. A physical-GPU4-7 C1 + candidate/control/candidate sandwich with the one-CTA launch measured + endpoint-mean pure decode `116.405` versus `112.842 tok/s` (`+3.16%`), output + `103.070` versus `100.112 tok/s` (`+2.96%`), and P50 TPOT `8.678` versus + `8.750 ms` (`-0.82%`). The source heuristic is limited to fully connected + SM70 TP4, the exact MTP4 payloads, and ordinary all-reduce; the global block + override still takes precedence and + the policy is enabled only by + `VLLM_SM70_TP4_MTP_AR_BLOCK_TUNING=1`; unset/zero retains the legacy policy. +- The C1 graph also spends about `0.509 ms` per cycle in the combined + top-k/top-p kernel over five 248,320-vocabulary verifier rows. Keeping the + existing 8192/4096 tiles and masking algorithm but launching eight rather + than four warps improves isolated B5/B10/B20/B40/B60/B80 means by + `1.737x/1.943x/1.881x/1.732x/1.689x/1.646x`. Every candidate had zero finite + mask or retained-value mismatches against the four-warp result. The opt-in + source change is restricted to combined top-k plus top-p on SM70, vocabulary + 248,320, and measured row counts 5/10/20/40/60/80. Enable it with + `VLLM_SM70_TOPK_TOPP_8_WARPS=1`; unset/zero keeps Triton's default launch. +- Corrected slot-independent compact metadata is finite and route-correct on a + real TP4 layer through B20. B9 is bitwise equal to the dense route. B10 has + cosine at least `0.99999988`, relative L2 at most `1.0115e-4`, and maximum + absolute error `2.38e-7`; its isolated graph stage is `1.091x` faster. The + production candidate stops at B10/80 routed slots and prewarms the exact + B1-B10 metadata signatures before CUDA Graph capture. +- A physical-GPU4-7 C2 candidate/control/candidate sandwich accepts B10. + Output was `154.608/147.135/160.890 tok/s`; candidate-mean output improves + `7.21%`. Pure decode was `91.962/89.762/94.566 tok/s`, a `3.90%` candidate + mean gain. Candidate mean acceptance length was lower than control + (`2.909` versus `2.963`), so acceptance does not explain the win. Extending + compact routing through B20 is rejected by a matching C4 sandwich: candidate + mean output/pure decode regressed `1.90%/2.48%`, with P50 TPOT about `1.18%` + worse. B11+ therefore keeps the 256-expert dense-grouped route. +- The final B10 48-request curve on one exclusive physical-GPU4-7 claim is: + + | concurrency | output tok/s | efficiency vs C1 | pure decode tok/s | serial prefill tok/s | P50 TPOT ms | mean acceptance | + | ---: | ---: | ---: | ---: | ---: | ---: | ---: | + | 1 | 97.693 | 100.00% | 114.247 | 564.795 | 8.522 | 2.890 | + | 2 | 150.736 | 77.15% | 90.851 | 423.341 | 10.207 | 2.946 | + | 4 | 246.399 | 63.05% | 71.322 | 596.885 | 13.559 | 2.884 | + | 8 | 350.834 | 44.89% | 52.398 | 629.601 | 20.006 | 2.969 | + | 12 | 444.117 | 37.88% | 46.313 | 631.019 | 22.274 | 2.909 | + | 16 | 530.810 | 33.96% | 44.734 | 720.400 | 23.478 | 2.900 | + +- Relative to the preceding final-source single-run curve, output/pure-decode + movement was C1 `-2.73/-3.08%`, C2 `+3.13/+2.41%`, C4 + `-0.73/-0.18%`, C8 `-2.93/-2.99%`, C12 `+4.21/+3.93%`, and C16 + `-4.47/-6.38%`. This alternating high-width movement is not attributable to + the bounded B10 route; B11+ still uses the same dense-grouped main shape and + only drain-tail widths can reach B9/B10. Treat B10 as a low-concurrency win, + not as high-concurrency recovery. +- The next target-side candidate must compact GPU-resident active expert IDs + for B20/B40/B60/B80 and use a device-count guard with an exact full-group + fallback. It may reduce padded grouped scheduling without a host readback or + a probabilistic correctness assumption. A persistent tile scheduler remains + the larger follow-up if guarded active-group bucketing cannot beat its extra + CUDA Graph launch. B10 has passed the pinned GSM8K-128 MTP4 gate at 122/128, + with zero invalid or repetitive records; the next quality gate belongs to + any high-concurrency scheduler candidate that survives full-model A/B/A. + +### Matched no-MTP scaling control + +- A matched performance-only control disabled MTP while preserving the B10 + source, model, TP4 physical GPU4-7 claim, 48 fixed ShareGPT requests, + sampling, FP8-E5M2 KV cache, TurboMind ModelOpt mixed FP8/NVFP4 route, + Flash-V100/FlashQLA, max length/batched tokens 8192, and + `FULL_AND_PIECEWISE` CUDA Graph policy. Every point reported + `speculative_config=None` and null speculative metrics, generated the exact + expected 9,115 prompt and 10,623 output tokens, and had no timed-request + Triton JIT warning. This control is not an MTP quality oracle; dataset-level + MTP quality remains gated independently. +- The no-MTP curve is: + + | concurrency | output tok/s | efficiency vs C1 | pure decode tok/s | serial prefill tok/s | P50 TPOT ms | + | ---: | ---: | ---: | ---: | ---: | ---: | + | 1 | 92.260 | 100.00% | 111.256 | 454.959 | 9.007 | + | 2 | 141.946 | 76.93% | 84.615 | 432.510 | 10.197 | + | 4 | 210.221 | 56.96% | 61.679 | 385.377 | 12.333 | + | 8 | 305.539 | 41.40% | 46.288 | 330.254 | 20.327 | + | 12 | 388.944 | 35.13% | 41.609 | 288.067 | 25.511 | + | 16 | 448.918 | 30.41% | 35.490 | 688.532 | 29.957 | + +- Relative to this no-MTP control, MTP4 improves output throughput by + `5.89/6.19/17.21/14.82/14.19/18.24%` at C1/C2/C4/C8/C12/C16 and improves + pure-decode throughput by `2.69/7.37/15.63/13.20/11.30/26.05%`. + MTP4 scaling efficiency is `0.00/0.22/6.09/3.49/2.75/3.55` percentage + points higher. Disabling MTP therefore neither recovers the absolute + high-concurrency throughput nor improves scaling; verifier/drafter work is + not the primary scaling bottleneck. Continue on the shared target decode + path, especially high-width routed-MoE scheduling and TP4 synchronization, + while retaining the independent MTP dataset-quality gate. +- The first completely cold C1 process also exposed a startup-only anomaly: + FULL decode graph capture took 230.21 s, total graph capture took 294 s, and + recorded load time was 380.46 s. Subsequent independent processes captured + FULL decode graphs in less than two seconds and total graphs in 35-68 s. + This one-time cold-cache cost was outside every timed throughput interval; + investigate it separately rather than attributing it to steady-state decode. +- Reproducible JSON and route logs are under + `bench_results/qwen36_nvfp4_no_mtp_scaling_20260824/{results,logs,telemetry}`. + +## 2026-08-24 Qwen3.8-27B-NVFP4 GDN decode fusion + +- Frozen contract: TP4 on four V100-SXM2-32GB GPUs, input 1024, output cap + 256, official random sampling, E4M3 KV, Flash-V100, full CUDA graph, and no + MTP. Pure decode excludes TTFT/prefill and each endpoint is the stable median + of three sequential requests. +- The quality-safe QPN8 QKV/Z plus FP16 b/a split projection and one-pass GDN + RMSNorm improve the matched control from `12.249652 ms / 81.634970 tok/s` + to `11.921648 ms / 83.881019 tok/s`. All three requests exactly reproduce + the control's 256 token IDs and SHA256 + `8b37337f4c393711cb8550db6bae909b1e85de8df1cf5ba8c60d8c000749c0a2`. +- A final post-rebase default-wrapper admission recheck reports + `12.261727 ms / 81.554583 tok/s` for the matching control and + `11.956949 ms / 83.633383 tok/s` for the candidate. The C++ route oracle + fires on all four ranks and all three 256-token candidate streams exactly + match the current-main control hash + `ca77db3b032a1600a8567adea706108c1bd8c5472b3ace2318c41ab17c66c1f9`. +- The GDN RMSNorm 48-layer operator gate is bitwise exact and saves + `55.946 us/token`. QKV is exact on all 48 layers; b/a has relative L2 + `6.368e-8` and maximum absolute error `2.384e-7`. The split projection is + admitted only when its one-pass RMSNorm pair is also enabled. Runtime + selection uses SM70, exact tensor/layout, QPN8 workspace, bias, TP, and + no-speculation contracts; the measured model/checkpoint is evidence only. + Large-M and unsupported shapes retain the original paths. Archived evidence + uses the pre-audit Qwen-prefixed flag names; final source uses generic GDN + flags and leaves the measured kernels unchanged. +- QPN8 QKV/Z plus b/a alone reaches `83.740763 tok/s` but diverges at token + 202, so it is not an independent accepted configuration. It is accepted only + together with the exact one-pass RMSNorm, whose paired endpoint restores the + complete frozen output. +- The Q/K RMSNorm plus MRoPE plus KV-write candidate reaches stable repeats of + `85.451/85.015 tok/s` but diverges at token 1 and stops at 208 tokens. Its + source, schema, flag, and rejected QPN4 communication experiments were + removed before submission. +- Next target: quantify no-MTP pure-decode decay at exact final context lengths + 16K, 64K, 128K, and 256K, reporting TTFT/prefill separately and preserving + the same model, quantization, TP, attention, KV, sampling, and graph routes. + +## 2026-08-24 MRV2 DFlash2 score-gated verifier promotion + +- The native Flash-V100 grouped verifier keeps all eight MRV2 target rows and + does not reuse DDTree indices. Its one-pass paged-GQA kernel is 0.356/1.214/ + 2.396 ms at 32K/128K/256K versus 1.024/3.021/6.011 ms for the previous XQA + verifier. Matched FULL-Graph endpoints reduce complete rounds from 36.311 to + 26.361 ms at 32K and 69.855 to 40.778 ms at 128K; paired output and + acceptance counters match. +- Two independent short-context traces of the complete pre-grouped bundle + measure 19.897/19.943 ms mean and 19.896/19.948 ms median. Their maxima are + 20.099/20.146 ms, so this is a replicated mean/median sub-20 result, not an + all-tail claim. +- The final quality contract is probabilistic sampling at T=1/top-p=.95/top-k + 20/xhigh, maximum 2,048 tokens, and natural EOS. Greedy token identity is not + a hard gate. Grouped matches target-only at 83/96 across GSM8K, MATH-500, and + HumanEval, and at 25/32 on corrected MBPP. Paired McNemar p is 1.0 for both + aggregates. WikiText PPL changes 5.497872 -> 5.497376 (-0.0090%), so no + measured score or PPL regression remains. +- The invalid raw-MBPP speed prompts omitted function names/tests and are + permanently excluded from intelligence claims. Official DFlash2 publishes + acceptance and throughput but not task pass@1; local target-only is the + quality baseline, while official acceptance is reference-only. +- Audit disposition: the acceleration leaves remain explicit opt-ins and use + hardware, tensor layout, algorithm configuration, cache/graph, and bounded + concurrency contracts rather than model or checkpoint identity. The + recorded profile remains evidence for the measured workload, not a runtime + model whitelist. Full evidence is in + `docs/design/sm70_dflash2_target_graph_20ms.md`. + +## 2026-08-24 DeepSeek V4 DSpark verifier and long-context campaign + +- The active TP8/V100 campaign, exact baseline contracts, admitted focused + gates, rejected paths, long-context sweep, and paired quality requirements + are recorded in + `docs/design/sm70_deepseek_v4_dspark_verifier_20ms_long_context.md`. +- Latest-main DSpark has reproduced a 116.288 token/s median on the historical + 3,500-token endpoint contract. The current complete verifier is 35.120 ms per + round; the active target is at most 20 ms without acceptance or dataset loss. +- Do not cite indexer, mHC, compressor-ring, or MoE microbenchmarks as the final + verifier result. Promotion requires a synchronized TP8 full-model profile, + long-context 1K-252K evidence, paired no-speculation quality, and task-owned + GPU cleanup. +- Source-audit disposition: incomplete-risk leaves are retained only as + explicit opt-ins. `VLLM_SM70_INDEXER_DECODE_CUBLAS=1` enables gather-plus-FP32 + cuBLAS, `VLLM_SM70_DSV4_PRIVATE_COMPRESSOR_STATE=1` enables private state + rings, and speculative context buckets require an explicit list. Default + serving does not compute or retain confidence logits unless a threshold or + alignment diagnostic consumes them, and a requested missing confidence + projection fails before use. +- Runtime gates express hardware, operator configuration, tensor/layout, + graph, cache, and concurrency contracts. The architecture-name check was + removed; measured checkpoint details remain evidence only and never select + an optimization route. + +## 2026-08-24 PR #276 cache-block reuse audit + +- The reported corruption root cause is valid: a reused full-attention cache + block can retain non-finite values in slots that its new owner has not yet + written, and arithmetic masking does not make `0 * NaN` finite. The MRV2 + runner had omitted the cache-zero lifecycle hook already used by the legacy + runner. +- MRV2 now initializes the shared `KVBlockZeroer` metadata and consumes the + scheduler's `new_block_ids_to_zero` before either attention or a zero-token + early return. Work is proportional to newly allocated blocks, uses the + cache-group/operator layout metadata, and introduces no per-step sequence + length readback or request-by-layer launch loop. +- The PR's per-step tail walker and sampled-token vocabulary clamp are + rejected. The former performs device-to-host synchronization plus many + small zero launches on every step; the latter turns corrupted output into + an arbitrary boundary token and hides the producer. The proposed GDN PAD + replication and last-column gather clamp are also rejected: the align-mode + allocator already guarantees a live current block plus its speculative + scratch window, so those fallbacks would conceal a broken block-table + contract and repeat an unrelated state block. +- Route selection and correctness depend only on allocation lifecycle, + cache-group layout, block IDs, tensor shape/dtype, and operator contracts. + No model, checkpoint, or architecture identity participates in the fix. +- Focused validation covers MRV2 zeroer initialization and zero-before-return + ordering, dense and multi-pool interleaved zeroing, Mamba align allocation, + and the existing GDN state-window contracts. No full-model end-to-end run + was required for this source-level repair. + +## 2026-08-24 Qwen3.8-27B NVFP4 no-MTP concurrency scaling + +- The heading names the matched performance workload, not an activation + identity. Runtime admission does not inspect a model name, checkpoint, or + architecture: it uses only SM70 capability, FP8 format, batch and tensor + shapes, paged-KV layout, graph context, and sampler operator settings. +- Contract: four V100s with TP4, compressed-tensors NVFP4 weights, FP16 + activations, E4M3 FP8 KV cache, Flash-V100 attention, prefix caching, Mamba + align mode, no speculation, and CUDA Graph decode sizes 1/2/4/8/16. The + throughput workload is official SPEED-Bench `low_entropy` at 1K and 16K, + 512 output tokens, one warmup plus three measured repeats, temperature 1.0, + top-k 20, top-p 0.95, fixed seed 20260822, and natural EOS. Reported decode + throughput uses only the interval in which every request is actively + decoding, so TTFT and prefill are excluded. +- The admitted E4M3 G6/D256 XQA route covers exact B2-B16 only. It uses paired + 128-bit KV loads, packed exact E4M3-to-half2 conversion, the existing + partition workspace, and the established long-context graph variants. At + 16K, B8 captures `dual_cta_split` and B16 captures `dual_cta`; the baseline + CUDA Graph remains selected below 16K. `VLLM_FLASH_V100_E4M3_BATCH_XQA=0` + restores scalar paged attention. The separate + `VLLM_FLASH_V100_E4M3_BATCH_XQA_OPTIMIZED=0` switch retains XQA while + restoring scalar KV loads, the original E4M3 conversion, and baseline CTA + routing. B1 keeps the original converter regardless of this batch-only + optimization. +- The exact 248,320-column, no-spec B8/B16 sampler contract now selects eight + Triton warps for combined top-k/top-p masking. The operation and reduction + math are unchanged; `VLLM_SM70_TOPK_TOPP_B8_B16_8_WARPS=0` restores Triton's + launch heuristic. Existing MTP verifier-row tuning remains independently + opt-in. +- Development A/B on base `a8ce98ba34` measured the following median pure + decode rates. The candidate includes batched XQA, eight-warp sampling, and + the long-context graph variants. + + | Context | Batch | Control tok/s | Candidate tok/s | Speedup | Candidate scaling efficiency | FP8 reference efficiency | + |---:|---:|---:|---:|---:|---:|---:| + | 1K | 1 | 72.82 | 72.94 | 1.002x | 100.0% | 100.0% | + | 1K | 8 | 402.51 | 465.92 | 1.158x | 79.8% | 87.4% | + | 1K | 16 | 642.37 | 774.18 | 1.205x | 66.3% | 72.3% | + | 16K | 1 | 63.03 | 63.12 | 1.001x | 100.0% | 100.0% | + | 16K | 8 | 186.43 | 356.50 | 1.912x | 70.6% | 67.7% | + | 16K | 16 | 222.03 | 526.59 | 2.372x | 52.1% | 49.1% | + +- Thus the 16K NVFP4 scaling efficiency exceeds the same-criterion FP8 + reference by 4.2% at B8 and 6.2% at B16. The remaining short-context gap is + 8.6% at B8 and 8.2% at B16; do not repeat the attention experiment to close + it. At 1K the next target is non-attention concurrent work, primarily GDN + state traffic and dense projection scheduling. +- A clean production revalidation on base `febf1fc962` reproduced the result + with one warmup and three measured repeats per cell. Only the steady window + in which every request was decoding is reported; endpoint throughput, TTFT, + and prefill are excluded from the table. + + | Context | Batch | Control decode tok/s | Candidate decode tok/s | Speedup | Candidate scaling efficiency | FP8 reference efficiency | + |---:|---:|---:|---:|---:|---:|---:| + | 1K | 1 | 72.65 | 72.62 | 1.000x | 100.0% | 100.0% | + | 1K | 8 | 402.98 | 465.24 | 1.154x | 80.1% | 87.4% | + | 1K | 16 | 641.95 | 772.44 | 1.203x | 66.5% | 72.3% | + | 16K | 1 | 63.10 | 63.00 | 0.998x | 100.0% | 100.0% | + | 16K | 8 | 186.97 | 356.52 | 1.907x | 70.7% | 67.7% | + | 16K | 16 | 222.61 | 527.00 | 2.367x | 52.3% | 49.1% | + + The B8 geometric-mean efficiency gap to the FP8 reference is 2.2%, and the + B16 gap is 1.0%. Both therefore reach the revised near-FP8 scaling target; + at 16K they exceed FP8 by 4.4% and 6.5%, respectively. Every timed request + succeeded and produced its full 512-token cap without empty or replacement- + character output. The original implementation was rebased without conflict + onto `onecat/main@6fc9fe9492`; the Flash-V100 and sampler source blobs were + byte-identical to the measured tree. The final audit rebased the same two + commits onto `onecat/main@f6a5b57b64` after generic prefix-anchored SWA was + merged. `git range-diff` reports both patches as exactly equivalent. The + combined Flash-V100 extension builds for SM70, 24 focused E4M3/graph/sampler + policy tests pass, and all 32 prefix-anchored SWA CPU regressions remain + green. The audit also prevents E4M3-only Flash-V100 graph variants from + being captured for an explicitly selected FlashInfer backend. A + same-toolchain `cuobjdump` comparison against the prefix-anchored main build + finds all 115 common XQA kernel instances SASS-identical, with zero missing + or changed instances; the candidate adds only six batch-optimized instances. +- GPU operator validation covers B2/B16 at partition sizes 64/128/256 plus + B8/B16 long-context routing. All eight comparisons stay within one FP16 ULP + of scalar paged attention with mean absolute error at most `1e-5`. Separate + B8/B16 full-vocabulary tests show bitwise-identical top-k/top-p masks between + the eight-warp schedule and its rollback path. +- Matched GSM8K main/test quality uses 32 prompts per B1/B8/B16 with the same + official stochastic sampling and 512-token natural-EOS cap. Changed batches + move numeric exact match from 22/32 to 23/32 at B8 and remain 23/32 at B16; + all requests complete with no empty or replacement-character output. The + unchanged B1 route moves 25/32 to 24/32 while only 21/32 texts reproduce + across independent processes, so that one-sample movement is recorded as + stochastic run variance rather than attributed to the candidate. +- The first metadata-only long-context probe was invalid: the CUDA Graph + dispatcher still admitted specialized variants for E5M2 only. It produced + no E4M3 route trace and must not be cited. The accepted implementation gates + both metadata and dispatcher creation; capture logs prove eight FULL graphs + and explicit B8/B16 long-context route hits. + +## 2026-08-25 DFlash2 practical-performance closure + +- PRs #252, #253, #257, #266, and #284 are already merged. The missing + performance closure was twelve post-#284 files that remained only in the + measured production worktree, not an open historical PR. +- A matched single-request coding probe with FP8 E5M2 KV, 256K maximum + context, 4096 maximum batched tokens, prefix caching, Mamba align, CUDA + Graph, and tool/reasoning parsers measured `18.66--18.83 ms` per complete + round (`131.2--139.9 tok/s`) on that production tree. The clean + `onecat/main@05d5aa4e57` route hit the same visible QPN8, sharded projection, + grouped verifier, and CUDA Graph logs but measured `27.84--27.96 ms` + (`90.5--92.6 tok/s`). Route-hit logs alone are therefore not an acceptable + performance proof. +- Draft PR #288 restores the exact missing source closure on latest main. The + decode hot path enables fused GDN group metadata for `mamba_cache_mode=align` + and makes disabled GDN dump hooks return before CUDA runtime queries. The + prefill half skips unused DFlash queries on intermediate 4096-token chunks + and keeps the small-M TurboMind projection out of large-M prefill shapes. +- Focused pre-commit passes. On the latest-main worktree, DFlash2/GDN CPU + suites report `110 passed, 32 skipped`; the V100 align-metadata replay and + AOT fullgraph predicate tests both pass. The first end-to-end launch produced + no timing result because an external TP4 job acquired GPUs 4--7 between API + startup and worker memory admission. Wait for the cards and rerun the exact + practical contract before merging; do not lower memory utilization or cite + that failed launch as implementation evidence. +- The queued clean rerun on PR #288 used the exact practical contract: TP4 on + V100, FP8 E5M2 target KV, FP16 draft KV, 256K maximum context, 4096 maximum + batched tokens, prefix caching, Mamba align, probabilistic DFlash2, CUDA + Graph, and the Qwen tool/reasoning parsers. Startup logs prove QPN8 exact + rerank, the TP4 sharded context projection, the one-pass grouped verifier, + target and draft graph capture, and, critically, fused GDN metadata for all + 10 cache groups. +- After one explicitly excluded first-request JIT warmup, four 512-token + single-request runs of `给我写个自我介绍网页吧` measured `18.465--18.537 ms` + per complete round and `133.05--146.63 tok/s`; acceptance length was + `2.457--2.718`. Three 1,024-token practical coding runs measured + `18.587--18.603 ms`, `132.62--139.93 tok/s`, and acceptance length + `2.470--2.606`, reproducing the measured production worktree. +- The historical high-acceptance MBPP item 28 reproduced `251.60 tok/s`, + `18.567 ms` per complete round, acceptance length `4.686`, and a natural-EOS + 328-token completion. Its complete generated text has the same SHA256 as the + previous production result, and EvalPlus reports base `1/1` and plus `1/1`. + This closes both the roughly 18 ms and the 220+ token/s restoration gates. +- A fresh 32K chunked-prefill request logged the context-only DFlash path and + measured `10.58--10.65 s` cold; the identical prefix hit measured `1.405 s`. + The older public server's nominal cold result may have retained an identical + prefix, so only the fresh-server numbers are treated as strict cold evidence. + +## 2026-08-25 Qwen3.8-27B NVFP4 DFlash2 17 ms campaign + +- This section is an active planning record, not evidence that a 17 ms result + or source optimization already exists. The document itself is safe to merge + while implementation and measurements remain pending. +- The next isolated campaign starts from + `onecat/main@34403018d917054dd7765d5e820ad29c8d342348` and targets a stable + batch-one complete speculative round of at most `17.0 ms`. The frozen target + is the local mixed NVFP4/channel-FP8 Qwen3.8-27B checkpoint with BF16 LM + head, TP4 on V100 GPUs 0--3, E5M2 target KV, FP16 draft KV, seven draft + tokens, probabilistic sampling, and FULL target/draft CUDA Graphs. +- The accepted PR #288 practical baseline is `18.465--18.603 ms`, so the + remaining measured gap is approximately `1.5--1.6 ms`. A current-source + no-diagnostic graph-node trace is required before choosing the next source + change. Historical DDTree index reuse, skipped verification, generic + auxiliary-stream overlap, variable K, and approximate reranking remain + rejected. +- Development may use a 32K/512 minimal runtime only to localize kernels. The + promotion gate is the 256K/4096 configuration with prefix caching, Mamba + alignment, and tool/reasoning parsers enabled. Acceptance may not fall by + more than `0.05`; the existing scored datasets and PPL gates must not + regress. A separate 32K/128K/256K sweep protects long-context decay. +- A retained speedup becomes default-on after the engine-contract admission, + rollback, numerical, official-sampling quality, and matched endpoint gates + pass. A valid reduction-order change need not be bitwise or greedy-token + identical, but it must remain finite, within a dtype-appropriate error + envelope, and non-regressing on acceptance and scored quality. +- The full frozen contract, sequence, and pending Draft PR test record are in + `docs/design/sm70_dflash2_nvfp4_17ms.md`. +- The acceptance audit found a runtime/source mismatch rather than an NVFP4 + regression. Current FP8 and NVFP4 request-mean acceptance values are + `4.49887` and `4.53642` (request-mean completion tokens per round are + `4.49761` and `4.53214`), but both runs loaded stale stable-ABI RMSNorm SHA + `fe986a...9ba8c`. Its + batched-weight suite fails 10/13 cases and reproduces the old row-zero + DFlash K-norm defect; repaired SHA `7689e5...d449` passes 13/13. Target + QPN2/QPN8 and sparse-rejection ablations do not recover acceptance. +- DFlash now performs a one-time nonzero bitwise capability check before using + grouped per-layer K normalization. A stale binary is made correctness-safe + through an explicit per-layer fallback and warning. The two new behavior + tests and the full CPU DFlash2 suite pass (`70 passed, 12 skipped`). The + draft-MLP QPN8 experiment is removed after exceeding the `-0.05` paired + acceptance gate. +- The repaired-binary FP8 formal rerun closes at request/pooled acceptance + `5.37699/4.62217`, 61/64 GSM8K, and stable SHA `7689e5...d449`. Relative to + the repaired August 22 formal baseline, paired request delta is `-0.01180` + with bootstrap 95% interval `[-0.08421,0.06282]`; every pooled position is + higher. Relative to the stale current-source run, repair adds `0.87812` + request-mean tokens with 95% interval `[0.75099,1.00344]`. The concurrent + throughput is excluded from standalone speed evidence. +- The matching repaired NVFP4 formal run closes at request/pooled acceptance + `5.39851/4.57051` and 59/64 GSM8K. Its request mean is `+0.02152` versus + current repaired FP8 (paired 95% interval `[-0.10021,0.14795]`), while its + longer 35,366-token trajectory explains the lower pooled value. Repair adds + `0.86209` over stale NVFP4 with interval `[0.74547,0.98007]` and improves 62 + of 64 rows. Its GSM8K pass/fail set is exactly the same as the repaired + August 22 FP8 reference (59 pass, 5 fail). +- The repaired 16K executable MBPP-32 pair passes the quality gate. Dataset + tests are no-DFlash `31/32` versus DFlash2 `32/32`; EvalPlus base is + `30/31` versus `31/31`, and plus is `27/31` versus `27/31`. The plus + discordance is one win per route with exact McNemar `p=1.0`; both runs stop + naturally on 31 rows and hit the 16K limit on one. A scorer extraction bug + that preferred reasoning-era example fences over valid unfenced final code + caused the initial DFlash syntax false negative. Final-answer-first + extraction is regression-checked, and both sides use the corrected scorer. +- The selector-QPN8 candidate passes acceptance but not yet performance + promotion. Dense-to-QPN8 request/pooled acceptance is + `5.39851/4.57051 -> 5.43951/4.58679`; the paired request delta is + `+0.04101` with 95% interval `[-0.02309,0.10669]`, and GSM8K is + `59/64 -> 60/64`. A cross-physical-group run shows directional complete + round `19.083 -> 18.077 ms`, aggregate `218.58 -> 229.41 token/s`, and + steady decode `282.68 -> 300.12 token/s`, but it is not accepted as a speed + pair. The prepared rerank layout reduces available KV capacity from 529,616 + to 459,817 tokens. An unrelated eight-GPU job blocked the same-group repeat, + so the route remains opt-in and default-off. + +## 2026-08-25 Qwen3.8-27B NVFP4 long-context decode + +- The accepted candidate specializes only the SM70 E4M3 G6/D256 dual-CTA + decode route whose production K/V views come from an interleaved + `[blocks, 2, 800, 1, 256]` allocation. It replaces general page arithmetic + with a compile-time page800/physical-stride path and uses aligned shared-V + `half2` loads without changing per-output FMA order. Other layouts, dtypes, + dimensions, routes, and prefill retain the fallback. +- `VLLM_FLASH_V100_E4M3_PAGE800_FASTPATH=0` is the rollback. The default-on + path also has a once-per-process trace flag and is protected by exact shape + and stride checks. Production-layout B8/B16 p64/p128/p256 tests are bitwise + equal to rollback; the two focused GPU suites pass 15/15 and 14/14. +- The matched TP4, no-MTP, E4M3-KV, FULL-graph SPEED-Bench `low_entropy` + pure-decode A/B improves 32K B8/B16 by 11.79%/15.04% and 64K B8/B16 by + 16.47%/19.74%. The four-cell geometric mean is 15.72%, no cell regresses, + and 64K B16 scaling efficiency is 48.52%, only 3.76 percentage points below + the accepted 16K reference. All 144 candidate requests finish with healthy + 512-token output. +- NCU on 64K/B16 p256 attributes the 22.47% kernel reduction to 11.92% fewer + instructions, 48.16% fewer shared-bank conflicts, and 95.92% less MIO- + throttle stall. Peak-relative DRAM/SM utilization rises 24.49% to 28.11% + and 39.11% to 44.79%; occupancy is unchanged. This is copy/addressing and + shared-memory relief, not an occupancy-only claim. +- The official-sampling quality gate uses the same 32 GSM8K main/test prompts + at B8/B16 with temperature 1.0, top-k 20, top-p 0.95, natural EOS, and three + predeclared sampling seeds. Candidate/control aggregate numeric EM is 74/96 + versus 71/96 at B8 and 74/96 versus 74/96 at B16; all 384 requests succeed + with no candidate empty or replacement-character output. The initial single + seed's one-item deficits remain recorded rather than being silently rerun. +- The measured patch rebases without conflict onto `onecat/main@d62ef5cb20`. + The final-source rebuild and measured binary have byte-identical complete + SM70 SASS dumps. + The server's common D256 prefill fallback warning is outside the pure-decode + metric, so this change closes long decode only and must not be cited as a + prefill improvement. + +## 2026-08-25 generic PP profiling and compact MXFP4 audit + +- PR #283 mixed accepted engine work with rejected workload experiments. Its + QPN8 candidates had repeated baseline-correct-to-candidate-wrong quality + regressions, while the metadata-free PP transfer recovered only about + `0.034 ms/token` in the full endpoint. Neither path is promoted to main. +- The retained PP repair is model-independent. A non-last pipeline stage can + return `IntermediateTensors` during a dummy profile/capture; the runner now + selects the semantic `hidden_states` entry instead of treating the wrapper + as a tensor. A CPU unit test also proves that dictionary insertion order + cannot select a residual tensor accidentally. +- The retained MXFP4 route is admitted only by the existing SM70 operator, + group size 32, six one-row groups, and exact `(K,N)` tensor shapes. Runtime + code does not inspect checkpoint paths, model names, `model_type`, or + architecture identity, and the compact grouped route remains default-off. +- On a V100, both new shapes are bitwise equal to the per-expert fallback and + pass CUDA Graph replay. Median operator time changes from `0.139827` to + `0.037478 ms` for `(4096,1024)` and from `0.064696` to `0.021371 ms` for + `(512,4096)`, or `3.731x` and `3.027x` respectively. Final-source reruns on + latest main measured `0.151521` to `0.040428 ms` (`3.748x`) and `0.065004` + to `0.022538 ms` (`2.884x`). The full CUDA 12.8 + extension builds with target architecture 7.0 only; 28 CPU checks and the + changed-file static gate pass. + +## 2026-08-25 Qwen3.8-27B-NVFP4 no-MTP long-context decode + +- Frozen acceptance is TP4 on four V100-SXM2-32GB GPUs, E4M3 KV, + Flash-V100, TurboMind FP8/NVFP4 projections, full CUDA Graph, official + temperature-1/top-p-.95/top-k-20 sampling, no MTP, and exact final contexts. + Pure decode excludes prefill/TTFT and uses 63 intervals from 64 output + tokens. Required rates are 60 tok/s at 128K and 45 tok/s at 256K. +- The admitted batch-one G6/D256, page-1568 E4M3 route uses an exact shared + conversion LUT and device-selected p64/p256/p512/p896/p1664 waves with + thresholds 12,288/49,152/98,304/196,608. A merged long partition/reducer pair + replaces three inactive long launch pairs during graph replay. Admission is + based only on hardware, dtype, batch, head, page-layout, graph, and explicit + partition contracts; it does not inspect a model or checkpoint identity. +- Batch-one full CUDA Graph replay is specialized at 49,152 tokens. The + default graph contains only the p64/p256 launch pair; the semantic + long-context graph admits the merged wave pair. Selection uses scheduler + context metadata and does not synchronize a device sequence length to the + host. Eager execution continues to route from the live context. +- The matched control is 40.561 tok/s (24.654 ms) at 128K and 27.456 tok/s + (36.422 ms) at 256K. The candidate reaches **61.834 tok/s (16.172 ms)** and + **50.376 tok/s (19.851 ms)** respectively, exceeding both targets. Prefill + remains separately reported at 283.127/1085.519 s and TTFT at + 283.154/1085.568 s. +- Both candidate requests return all 64 tokens without corruption and exactly + preserve their control streams: hashes are `66a130ea9d68...` at 128K and + `e1bc9ab8a979...` at 256K. No sampling path or MTP state changes. +- The disabled-compile-cache post-integration 1K/256-output regression guard + reaches 83.184/83.187 tok/s on its two hot repetitions versus + 83.761/83.721 tok/s for the accepted GDN BA-split control, a 0.66% median + delta. All candidate repetitions preserve that control's exact token stream + (`c243818da5cc...`). The first cold request remains separately reported, and + cache-reload diagnostics are excluded from quality evidence. +- A final isolated 256K Nsight Systems trace measures 0.576603 ms for the + merged long partition kernel and 0.033280 ms for its reducer. Including the + inactive device-routed short guard pair gives 0.617979 ms of Flash kernels, + down 63.83% from the old p256 trace's 1.708550 ms. This is operator evidence, + not a full-model TPOT decomposition. +- Exact decode windows show 100% NVML GPU-busy duty, about 48.2%/45.5% + memory-active duty at 128K/256K, 29.45-29.75 GiB allocation per rank, and + 200.7-226.8 W mean per-rank power at 1530/877 MHz SM/memory clocks. Memory + duty is not HBM throughput. NCU counters are blocked by + `ERR_NVGPUCTRPERM` and are not inferred. +- A post-acceptance operator pass retains the same route and arithmetic. It + stages partition page IDs once, specializes only the actual page-1568 + interleaved unbound-view strides, and consumes shared V in adjacent `half2` + pairs while preserving per-dimension FP32 FMA order. Four-arm same-GPU A/B + gives staged exact-context kernel reductions of 1.20%, 1.49%, and 3.99% at + 128K, and 1.23%, 1.73%, and 4.49% at 256K. All 12 sequence/layout cases are + bit-exact with one hash per case across all arms. The final kernel retains + 168 registers and two CTAs/SM, removes its 32-byte stack frame, and reduces + static SASS from 2,424 to 2,248 instructions versus the fixed-stride stage. + The formatted final build is eager- and CUDA-Graph-exact at 128K/256K, and + all eight page-1568 wave-boundary GPU regressions pass. +- Scaling the measured attention slice by those independently matched ratios, + while freezing the measured non-attention residual, projects 63.354 tok/s + at 128K and 52.216 tok/s at 256K. This is an operator-decomposition + projection, not a new endpoint measurement; 61.834/50.376 tok/s remain the + accepted measured rates. +- Reject the contiguous-KV assumption, QK software pipeline, split reducer, + and paired E4M3 converter. The real unbound K/V views are interleaved; + violating that layout corrupts model output. The QK pipeline slows the long + kernels by about 30-38%, and the paired converter loses to the exact shared + LUT. Detailed evidence and rollback flags are in + `docs/design/sm70_qwen38_nvfp4_decode.md`. +- The merge audit rebases this route onto `main@1ce39810b` and retains the + existing page-800 E4M3 specialization beside page 1568. The combined CUDA + 12.8 build contains only SM70 cubins; three host CUDA Graph routing checks, + six page-800 V100 fallback comparisons, and eight page-1568 wave-boundary + comparisons pass. Both GPU suites require bitwise equality. Runtime source + admission remains limited to device capability, dtype, batch/head geometry, + cache strides, page size, graph semantics, and partition thresholds; no + model, checkpoint, `model_type`, or architecture identity is consulted. + +## 2026-08-26 generic in-flight KV eviction audit + +- PR #293 mixed a valid engine-level cache lifetime repair with + architecture-keyed serving recipes, model-specific generation defaults, and + OCR prompt/resolution behavior. The architecture/model identity defaults are + not admitted to main; they are unrelated to inference-engine route safety. +- The retained repair tracks tokens scheduled in batches whose outputs have + not settled yet. Sliding-window, chunked-local, prefix-anchored SWA, and + Mamba eviction now derive their safe freeing boundary from processed tokens, + rather than the scheduler's optimistic computed-token boundary. This avoids + recycling a KV block while an asynchronous or pipeline-parallel batch can + still read it, including when speculative tokens later roll back. +- Recycling-aware admission reserves for the same executor concurrency bound: + pipeline depth for PP, asynchronous queue depth otherwise. The contract uses + scheduler concurrency, attention/cache geometry, token state, and block + lifetime only; it does not inspect a model, checkpoint, `model_type`, or + architecture identity. +- Final-source validation passes 17 focused cache/accounting cases, all 11 + async-scheduler unit cases, and a simple CPU-offload store/load round trip. + All changed-file pre-commit hooks pass. This is a cache lifetime correctness + repair; no throughput claim is attached. + +## 2026-08-26 DeepSeek V4 PP2 TP4 no-DSpark decode recovery + +- The frozen endpoint contract is DeepSeek V4 Flash on eight V100-SXM2-32GB + GPUs, PP2 with TP4 ranks confined to GPUs 0--3 and 4--7, input 1024, output + cap 256, model-default sampling, no speculation/DSpark, E5M2 KV, sparse MLA, mHC, + TurboMind FP8/MXFP4 projections, and full CUDA Graph replay. Pure decode is + reported separately from prefill and TTFT; the target remains 100 tok/s. +- The latest uninstrumented control reaches 59.160 tok/s, or 16.903 ms/token. + A 26-step Nsight Systems trace is perturbed to 18.939 ms/token but localizes + the per-stage service leaders to 6.434 ms of FP8 dense work, 2.647 ms of + MXFP4 MoE, 2.582 ms of FP16 GEMV, 1.714 ms of mHC, 1.695 ms of sparse MLA, + 1.694 ms of routing, 1.688 ms of Q/KV work, and 1.476 ms of TP all-reduce. + Roughly 9 ms pipeline rows are dependency waits and are not added as PP + transfer cost. +- PR #299 extends the fully connected SM70 TP4 + push all-reduce to the exact 8-KiB FP16 CUDA Graph payload. Same-binary TP4 + A/B/B/A operator timing moves 9.286--9.425 us to 3.041--3.174 us. Both paths + are bitwise equal to the explicit rank-0-through-rank-3 FP32 accumulation + oracle. A stronger changing-input graph test covers 43 collective nodes, + eight input patterns, 64 replays, and all four ranks with zero mismatches; + every GPU claim is released after the bounded test. The route is default-on + only for SM70, fully connected TP4, active CUDA Graph capture, FP16, and the + exact 8-KiB or 80-KiB payload; `VLLM_SM70_TP4_PUSH_ALLREDUCE=0` rolls back. +- Matched full-model control and candidate medians are 59.160 and 61.272 + tok/s, or 16.903 and 16.321 ms/token. The candidate therefore recovers + 3.57% and 0.583 ms/token. This is accepted performance and operator evidence, + while two independent unchanged controls select different stable token + streams because of pre-existing cross-process TurboMind autotuning. Greedy + identity is diagnostic rather than an acceptance gate; the exact collective + oracle isolates the changed operator. +- The next arithmetic target is the FP8 dense/WO-A launch family, which is the + largest trace category. Each of the 43 layers currently dispatches the two + one-row `(K=4096,N=1024)` groups separately. The replacement prepares + two persistent TurboMind pointer rows and dispatches both groups through one + blocked GEMM only for batch-one decode; all other shapes and token counts + retain the existing loop. +- The accepted arithmetic contract fixes the grouped descriptor to the dense + `8x128x64`, split-K 7 schedule. An environment-injected selector screen and + a source-built rerun both pass eager and changing-input CUDA Graph checks for + exact-small, random-small, and model-like inputs with zero mismatches. The + measured source-built extension SHA256 is + `133724eb67e70e3b4ff098c5f32c00e3155cd7c5e8f7aeb6dff84934ac294e9e`. + The merge audit changes the absent-environment default to the measured + explicit `=1` behavior; it does not change the selected tactic or kernel. + Median isolated graph time changes from `74.943` to `22.450 us` per layer, + projecting `2.257 ms/token` of service reduction before model overlap. + Independently tuned flat/grouped schedules differ by up to one FP16 ULP and + remain rejected. +- The matched full-model pair keeps the source, extension, push all-reduce, + topology, lengths, and cache contract fixed. The control median is + `57.436 tok/s` (`17.411 ms/token`) and the grouped candidate reaches + `59.209 tok/s` (`16.889 ms/token`), recovering `0.521 ms/token` or 3.09%. + This is the measured endpoint reduction; the larger isolated service + projection is not used as a TPOT claim. Both sides are internally stable + across three requests. Their different hashes reflect unrelated + cross-process FP8 tactic selection and are not treated as a numerical gate. +- The grouped route is therefore default-on only for the exact engine + contract already proven above: SM70 block-FP8 `[128,128]`, two groups, + `K=4096`, `N=1024`, batch-one decode, the source-built grouped and pointer + operators, and the fixed `8x128x64` split-K-7 tactic. Other shapes, multi-row + calls, missing operators, and non-SM70 devices retain the two-launch path. + `VLLM_SM70_FP8_GROUPED_BMM_DECODE=0` is the rollback. Admission never reads + model name, checkpoint, `model_type`, or architecture identity. +- A multi-row FP16 GEMV screen is rejected. The only bitwise-exact candidate + (`BLOCK_N=2`, `BLOCK_K=1024`) is 0.5--12% slower across the five traced + `N=64/256/512/1024/2048` shapes; smaller-K/four-row variants are slower and + also change FP32 accumulation. The candidate is not retained. +- The next MXFP4 B1 screen removes redundant route sorting without changing + model arithmetic. The six unique top-k IDs become the compact TurboMind + group IDs directly, fixed offsets remain `[0, 1, ..., 6]`, W13 reads the + single physical input row through the accepted broadcast contract, and W2 + stays in original top-k order. The final weighted reduction therefore keeps + the same route-order FP32 FMA sequence without an inverse permutation. + Initial and changed-route CUDA Graph replays are bitwise at W13, clamped + SwiGLU, W2, and final output. Adding the exact top-k-6 `half2` reducer moves + 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/`. + 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 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 + +- Development is isolated on + `agent/v100-dflash2-ngram-hybrid-20260825-050018` and rebased onto the + restored `onecat/main@d62ef5cb20`; Draft PR #287 contains only the MRV2 + DFlash2 dependency closure. +- The opt-in assistant reads the authoritative UVA token state, uses the same + reverse-KMP policy as standalone vLLM n-gram lookup, skips DFlash2 query and + selector only on a full seven-token hit, and preserves context-KV + materialization. Misses remain on the unchanged DFlash2 route. +- V100 unit coverage is 24/24, including probabilistic one-hot rejection, + mixed rows, overlap state, and the CUDA override kernel. Practical TP4 graph + startup with prefix cache, FP8 E5M2 KV, 256K max context, 4096 batched-token + limit, and tool/reasoning parsers is proven. The earlier 27--31 ms clean-main + measurements predate the merged production closure and are invalid as a + speed baseline; collect `ngram_assist=false/true` numbers from the same + restored build. +- The restored paired result is now complete. No-assist practical coding is + `18.660 ms` per round, `135.70 tok/s`, acceptance `2.532`; ngram `[5,5]` is + `18.980 ms`, `131.14 tok/s`, acceptance `2.494`, with only about 2.2% full + hits. The feature therefore remains opt-in for short/low-hit traffic. +- On the fixed-seed, 16K-cap, natural-stop MBPP32 workload, pure DFlash records + `389.542 s`, `19,767` rounds, `19.707 ms/round`, acceptance `3.336`, and + EvalPlus Base/Plus `30/31` and `28/31`. Hybrid records `327.613 s`, `16,700` + rounds, `19.618 ms/round`, acceptance `3.216`, and Base/Plus `31/31` and + `28/31`. Its cumulative full-hit rate is about 10.1%, but lower sampled + acceptance leaves aggregate output throughput `163.94` versus + `169.26 tok/s`; the quality gate passes and the throughput gate does not. +- Hybrid 32K prefill is `10.688 s` cold and `1.416 s` on an identical-prefix + hit, matching restored no-assist evidence (`10.58--10.65 s`, `1.405 s`) + closely enough to rule out a material prefill regression in this probe. +- Post-run capture failures on GPUs 4--7 reproduce unchanged on the clean + pre-ngram main worktree at the same draft paged-attention capture. They are + runtime/GPU-state evidence, not an ngram regression or speed sample. A full + NVLink-group reset was intentionally deferred because GPUs 0--3 host the + live user API. +- The merge audit removes the mixed-hit quality/performance hazard. N-gram + drafts are now applied only when every active request has a full-width hit + and the complete DFlash2 query/selector is skipped. If any request misses, + the batch keeps the unchanged DFlash2 proposals because overriding a row + after paying the full query cost cannot improve latency and can alter + 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 + +- 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. +- 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. +- 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. +- 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. + Three matched endpoint runs improve QPN8 by about `0.22%` (`0.034 ms/token`) + without changing its output behavior. Non-admitted configurations use the + 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 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. +- 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. That single sample remains visible, but + it is not a greedy-identity gate and does not show that choosing the lowest + operator error improves aggregate task quality. The accepted repair + therefore preserves every numerically bounded speed winner and falls back + only the demonstrated gate/up outlier. 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 + +- PR #297 originally wrapped one named vision module by invoking the layer + offloader a second time. That form is not admitted: it is model-specific, + loses the tower prefix needed by `--cpu-offload-params visual`, and violates + the prefetch backend's single-call contract. +- The accepted path is an engine-level component hook on every module recorded + by the existing multimodal tower lifecycle. UVA qualifies relative parameter + names with the recorded module path before applying exact segment matching; + layer-prefetch and no-op backends inherit a single-component no-op. Parameters + already visited through `make_layers()` are not copied or counted twice. +- Admission uses only the configured offload backend and budget, the explicit + parameter-path selector, and the module's generic tower role. It never reads + a model name, checkpoint path, `model_type`, or architecture identity. With no + CPU-offload budget the global offloader remains a no-op. +- The submitted V100 evidence reports about 0.78 GiB additional device memory + per rank, a 152,056-token KV pool, and successful 151,296-token context on + two 16-GiB V100 ranks. Text-only decode does not execute the tower; image + prefill accepts the expected PCIe/UVA cost. This is capacity evidence, not a + general throughput claim, so the route remains explicitly configured rather + than becoming an unconditional default. + +## 2026-08-26 PR #299 8-KiB TP4 push all-reduce acceptance + +- The existing two-epoch SM70 TP4 push collective now also admits the exact + 8-KiB FP16 decode payload. Four CTAs cover its 512 packed 16-byte elements; + the existing 80-KiB verifier payload retains its 80-CTA launch. Buffer + sizing already uses the larger payload, so this adds no allocation growth. +- Same-binary TP4 A/B/B/A timing improves from 9.286--9.425 us to + 3.041--3.174 us. A dynamic CUDA Graph gate covers 43 collective nodes, + eight changing input patterns, 64 replays, and all four ranks with zero + element mismatches. Both paths are bitwise equal to an explicit fixed-order + rank-0-through-rank-3 FP32 accumulation oracle. +- The matched PP2 x TP4 endpoint improves from 59.160 to 61.272 token/s + (`+3.57%`) and from 16.903 to 16.321 ms/token (`-0.583 ms/token`). Independent + unchanged controls themselves select different stable greedy streams due to + pre-existing cross-process autotuning, so token hashes are diagnostic only; + the exact operator gate isolates this collective without imposing greedy + identity. +- `VLLM_SM70_TP4_PUSH_ALLREDUCE` is therefore default-on. Runtime admission is + limited to SM70, fully connected TP4, active CUDA Graph capture, FP16, and + the exact 8-KiB or 80-KiB payload. Every other device, topology, dtype, size, + 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. +- Production defaults on through `VLLM_SM70_DSV4_FP13_GEMV=1`; explicit `=0` + is the rollback. Admission checks SM70, CUDA device equality, dtype, shape, + layout, output contract, and the actual weight bits. It never reads a model + name or checkpoint identity, and an incompatible tensor retains FP16. Packing + occurs once after loading; capture and per-token execution perform no packing. + The retained buffers add 549.66 MiB without pipeline parallelism, or at most + 282.34 MiB per rank for the audited PP2 partition. diff --git a/docs/features/speculative_decoding/README.md b/docs/features/speculative_decoding/README.md index 768e9f78d4..bda540b88c 100644 --- a/docs/features/speculative_decoding/README.md +++ b/docs/features/speculative_decoding/README.md @@ -117,6 +117,33 @@ 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. + +```bash +vllm serve \ + --speculative-config '{ + "method": "dflash", + "model": "your-org/your-dflash2-selector-model", + "num_speculative_tokens": 7, + "draft_sample_method": "probabilistic", + "ngram_assist": true, + "prompt_lookup_min": 5, + "prompt_lookup_max": 5 + }' +``` + +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. + #### Suffix decoding | Key | Type | Default | Meaning | diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 0654a59caf..9ffaa6954b 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -539,6 +539,14 @@ See [this page](generative_models.md) for more information on how to use generat These models primarily accept the [`LLM.generate`](./generative_models.md#llmgenerate) API. Chat/Instruct models additionally support the [`LLM.chat`](./generative_models.md#llmchat) API. +!!! note "DeepSeek-OCR image resolution" + `DeepseekOCRForCausalLM` accepts explicit `mm_processor_kwargs` with + `image_mode` set to `tiny`, `small`, `base`, `large`, or `gundam`. + `min_crops` and `max_crops` may also be set with + `1 <= min_crops <= max_crops <= 9`. The default remains the existing + `gundam` processing contract; no mode is selected from a model or + checkpoint identity at runtime. + | Architecture | Models | Inputs | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) | | ------------ | ------ | ------ | ----------------- | -------------------- | ------------------------- | | `AriaForConditionalGeneration` | Aria | T + I+ | `rhymes-ai/Aria` | | | diff --git a/flash-attention-v100/flash_attn_v100/__init__.py b/flash-attention-v100/flash_attn_v100/__init__.py index 5d38754d54..b3560a3ca8 100644 --- a/flash-attention-v100/flash_attn_v100/__init__.py +++ b/flash-attention-v100/flash_attn_v100/__init__.py @@ -10,6 +10,7 @@ flash_attn_decode_paged_xqa_available, flash_attn_decode_qk_scores, flash_attn_func, + flash_attn_grouped_verify_paged, flash_attn_lse, flash_attn_prefill_paged, flash_attn_prefill_paged_bfla, @@ -30,6 +31,7 @@ "flash_attn_decode_paged_xqa", "flash_attn_decode_paged_xqa_available", "flash_attn_decode_paged_wmma", + "flash_attn_grouped_verify_paged", "flash_attn_turboquant_decode_paged", "flash_attn_turboquant_decode_paged_available", "flash_attn_bhmd_func", 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 a28f7571e4..c7042bbdb3 100644 --- a/flash-attention-v100/flash_attn_v100/flash_attn_interface.py +++ b/flash-attention-v100/flash_attn_v100/flash_attn_interface.py @@ -18,11 +18,13 @@ DEFAULT_DECODE_PARTITION_SIZE = 256 VALID_DECODE_PARTITION_SIZES = (256, 512, 1024) +E4M3_XQA_VALID_DECODE_PARTITION_SIZES = (64, 128, 256, 512, 896, 1024, 1664) _decode_plan_cache = {} _decode_workspace_cache = {} _xqa_staged_rescale_workspace_cache = {} _turboquant_decode_workspace_cache = {} _prefill_splitkv3_workspace_cache = {} +_grouped_verify_workspace_cache = {} logger = logging.getLogger(__name__) @@ -52,6 +54,12 @@ class _PrefillSplitkv3Workspace: softmax_lse: torch.Tensor +@dataclass +class _GroupedVerifyWorkspace: + partial_out: torch.Tensor + partial_lse: torch.Tensor + + def maybe_contiguous(x): return x.contiguous() if x is not None and not x.is_contiguous() else x @@ -152,6 +160,7 @@ def _get_decode_plan( workspace_seq_capacity_hint: int | None = None, active_num_partitions: torch.Tensor | None = None, partition_size_hint: int | None = None, + valid_partition_sizes: tuple[int, ...] = VALID_DECODE_PARTITION_SIZES, ) -> _DecodePlan: batch_capacity = batch_size_hint or block_table.shape[0] num_heads = q.shape[1] @@ -184,6 +193,7 @@ def _get_decode_plan( _validate_decode_partition_size( int(partition_size_hint), "partition_size_hint", + valid_partition_sizes, ) if partition_size_hint is not None else _get_decode_partition_size( @@ -493,6 +503,35 @@ def _get_prefill_splitkv3_workspace( return workspace +def _get_grouped_verify_workspace(q: torch.Tensor) -> _GroupedVerifyWorkspace: + device_index = q.device.index if q.device.index is not None else -1 + key = ( + q.device.type, + device_index, + _workspace_stream_id(q.device), + q.dtype, + ) + workspace = ( + _grouped_verify_workspace_cache.get(key) if _can_cache_workspace(q) else None + ) + if workspace is None: + workspace = _GroupedVerifyWorkspace( + partial_out=torch.empty( + (80, 8, 6, 256), + dtype=torch.float16, + device=q.device, + ), + partial_lse=torch.empty( + (80, 8, 6), + dtype=torch.float32, + device=q.device, + ), + ) + if _can_cache_workspace(q): + _grouped_verify_workspace_cache[key] = workspace + return workspace + + def _get_decode_partition_size( max_seq_capacity: int, head_dim: int, @@ -529,11 +568,13 @@ def _select_default_decode_partition_size( return DEFAULT_DECODE_PARTITION_SIZE -def _validate_decode_partition_size(value: int, name: str) -> int: - if value not in VALID_DECODE_PARTITION_SIZES: - raise ValueError( - f"{name} must be one of {VALID_DECODE_PARTITION_SIZES}, got {value}" - ) +def _validate_decode_partition_size( + value: int, + name: str, + valid_partition_sizes: tuple[int, ...] = VALID_DECODE_PARTITION_SIZES, +) -> int: + if value not in valid_partition_sizes: + raise ValueError(f"{name} must be one of {valid_partition_sizes}, got {value}") return value @@ -897,6 +938,8 @@ def flash_attn_decode_paged( workspace_seq_capacity_hint: int | None = None, active_num_partitions: torch.Tensor | None = None, partition_size_hint: int | None = None, + anchor_lens: torch.Tensor | None = None, + anchored_window: int = 0, ): if softmax_scale is None: softmax_scale = q.shape[-1] ** -0.5 @@ -905,9 +948,16 @@ def flash_attn_decode_paged( block_table = maybe_contiguous(block_table) seq_lens = maybe_contiguous(seq_lens) out = maybe_contiguous(out) + anchor_lens = maybe_contiguous(anchor_lens) window_size_left, window_size_right = window_size if window_size_left < -1 or window_size_right < -1: raise ValueError(f"Invalid window_size={window_size}; values must be >= -1") + if anchored_window < 0: + raise ValueError("anchored_window must be non-negative") + if (anchored_window > 0) != (anchor_lens is not None): + raise ValueError( + "anchor_lens and a positive anchored_window must be provided together" + ) batch_capacity = q.shape[0] num_heads = q.shape[1] head_dim = q.shape[2] @@ -960,6 +1010,8 @@ def flash_attn_decode_paged( float(v_scale), int(window_size_left), int(window_size_right), + anchor_lens, + int(anchored_window), ) @@ -971,6 +1023,49 @@ def flash_attn_decode_paged_xqa_staged_available() -> bool: return hasattr(flash_attn_v100_cuda, "decode_paged_xqa_staged_fwd") +def flash_attn_grouped_verify_paged( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + softmax_scale: float | None = None, + out: torch.Tensor | None = None, + kv_cache_dtype: str = "fp8_e5m2", + k_scale: float = 1.0, + v_scale: float = 1.0, + one_pass: bool = False, +) -> torch.Tensor: + """Exact grouped q8/H6/D256 DFlash2 verifier prototype 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. + """ + if softmax_scale is None: + softmax_scale = q.shape[-1] ** -0.5 + q = maybe_contiguous(q) + block_table = maybe_contiguous(block_table) + seq_lens = maybe_contiguous(seq_lens) + out = maybe_contiguous(out) + workspace = _get_grouped_verify_workspace(q) + return flash_attn_v100_cuda.grouped_verify_paged_fwd( + q, + k_cache, + v_cache, + out, + block_table, + seq_lens, + workspace.partial_out, + workspace.partial_lse, + float(softmax_scale), + kv_cache_dtype, + float(k_scale), + float(v_scale), + bool(one_pass), + ) + + def flash_attn_decode_paged_xqa( q: torch.Tensor, k_cache: torch.Tensor, @@ -987,6 +1082,7 @@ def flash_attn_decode_paged_xqa( workspace_seq_capacity_hint: int | None = None, active_num_partitions: torch.Tensor | None = None, partition_size_hint: int | None = None, + batch_context_routing: bool = False, ): if not flash_attn_decode_paged_xqa_available(): raise RuntimeError("flash_attn_v100 CUDA extension lacks XQA decode") @@ -1016,6 +1112,22 @@ def flash_attn_decode_paged_xqa( workspace_seq_capacity_hint=workspace_seq_capacity_hint, active_num_partitions=active_num_partitions, partition_size_hint=partition_size_hint, + valid_partition_sizes=( + E4M3_XQA_VALID_DECODE_PARTITION_SIZES + if kv_cache_dtype in ("fp8", "fp8_e4m3") + and q.ndim == 3 + and q.shape[1:] == (6, 256) + and ( + q.shape[0] == 1 + or ( + os.getenv("VLLM_FLASH_V100_E4M3_BATCH_XQA", "1") == "1" + and 1 < q.shape[0] <= 16 + ) + ) + and k_cache.dtype == torch.uint8 + and v_cache.dtype == torch.uint8 + else VALID_DECODE_PARTITION_SIZES + ), ) _assert_decode_launch_covers_seq_lens( plan, @@ -1074,6 +1186,22 @@ def flash_attn_decode_paged_xqa( int(window_size_right), ) + batch_context_max_seq_len = 0 + if batch_context_routing: + live_max_seq_len = int(max_seq_len_hint or 0) + workspace_max_seq_len = int(workspace_seq_capacity_hint or 0) + if _cuda_graph_capture_active(): + batch_context_max_seq_len = max( + live_max_seq_len, + workspace_max_seq_len, + ) + else: + batch_context_max_seq_len = live_max_seq_len + if batch_context_max_seq_len <= 0: + batch_context_max_seq_len = workspace_max_seq_len + if batch_context_max_seq_len <= 0: + batch_context_max_seq_len = plan.launch_num_partitions * plan.partition_size + return flash_attn_v100_cuda.decode_paged_xqa_fwd( q, k_cache, @@ -1093,6 +1221,7 @@ def flash_attn_decode_paged_xqa( float(v_scale), int(window_size_left), int(window_size_right), + int(batch_context_max_seq_len), ) @@ -1232,6 +1361,8 @@ def flash_attn_prefill_paged( v_scale: float = 1.0, causal: bool = True, window_size: tuple = (-1, -1), + anchor_lens: torch.Tensor | None = None, + anchored_window: int = 0, ): if softmax_scale is None: softmax_scale = q.shape[-1] ** -0.5 @@ -1241,9 +1372,16 @@ def flash_attn_prefill_paged( block_table = maybe_contiguous(block_table) seq_lens = maybe_contiguous(seq_lens) out = maybe_contiguous(out) + anchor_lens = maybe_contiguous(anchor_lens) window_size_left, window_size_right = window_size if window_size_left < -1 or window_size_right < -1: raise ValueError(f"Invalid window_size={window_size}; values must be >= -1") + if anchored_window < 0: + raise ValueError("anchored_window must be non-negative") + if (anchored_window > 0) != (anchor_lens is not None): + raise ValueError( + "anchor_lens and a positive anchored_window must be provided together" + ) q_ = q.permute(0, 2, 1, 3).contiguous() out_ = out.permute(0, 2, 1, 3).contiguous() if out is not None else None @@ -1262,6 +1400,8 @@ def flash_attn_prefill_paged( causal, int(window_size_left), int(window_size_right), + anchor_lens, + int(anchored_window), ) return _copy_bhmd_to_bmhd_out(out_, out_original) @@ -1433,6 +1573,8 @@ def flash_attn_prefill_paged_bhmd( causal, int(window_size_left), int(window_size_right), + None, + 0, ) @@ -1521,6 +1663,7 @@ def flash_attn_prefill_paged_d256_bm32_allp_pair_scratch_splitkv3( "flash_attn_decode_paged", "flash_attn_decode_paged_xqa", "flash_attn_decode_paged_xqa_available", + "flash_attn_grouped_verify_paged", "flash_attn_decode_paged_wmma", "flash_attn_decode_qk_scores", "flash_attn_turboquant_decode_paged", diff --git a/flash-attention-v100/include/fused_mha.h b/flash-attention-v100/include/fused_mha.h index e778ed0d8f..0987368cba 100644 --- a/flash-attention-v100/include/fused_mha.h +++ b/flash-attention-v100/include/fused_mha.h @@ -8,243 +8,132 @@ #include 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_ -); - -at::Tensor flash_attention_qk_scores( - const at::Tensor& q, - const at::Tensor& k, - const float softmax_scale, - const bool is_causal -); + 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_); + +at::Tensor flash_attention_qk_scores(const at::Tensor& q, const at::Tensor& k, + const float softmax_scale, + const bool is_causal); at::Tensor flash_attention_decode_paged( - const at::Tensor& q, - const at::Tensor& k_cache, - const at::Tensor& v_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& active_num_partitions, - const float softmax_scale, - const int partition_size, - const int launch_num_partitions, - const std::string& kv_cache_dtype, - const float k_scale, - const float v_scale, - const int window_size_left, - const int window_size_right -); + const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_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& active_num_partitions, + const float softmax_scale, const int partition_size, + const int launch_num_partitions, const std::string& kv_cache_dtype, + const float k_scale, const float v_scale, const int window_size_left, + const int window_size_right, const std::optional& anchor_lens, + const int64_t anchored_window); at::Tensor flash_attention_decode_paged_xqa( - const at::Tensor& q, - const at::Tensor& k_cache, - const at::Tensor& v_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& active_num_partitions, - const float softmax_scale, - const int partition_size, - const int launch_num_partitions, - const std::string& kv_cache_dtype, - const float k_scale, - const float v_scale, - const int window_size_left, - const int window_size_right -); + const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_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& active_num_partitions, + const float softmax_scale, const int partition_size, + const int launch_num_partitions, const std::string& kv_cache_dtype, + const float k_scale, const float v_scale, const int window_size_left, + const int window_size_right, const int batch_context_max_seq_len); at::Tensor flash_attention_decode_paged_xqa_staged( - const at::Tensor& q, - const at::Tensor& k_cache, - const at::Tensor& v_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, - at::Tensor& online_rescales, - const at::Tensor& active_num_partitions, - const float softmax_scale, - const int partition_size, - const int launch_num_partitions, - const std::string& kv_cache_dtype, - const float k_scale, - const float v_scale, - const int window_size_left, - const int window_size_right -); + const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_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, at::Tensor& online_rescales, + const at::Tensor& active_num_partitions, const float softmax_scale, + const int partition_size, const int launch_num_partitions, + const std::string& kv_cache_dtype, const float k_scale, const float v_scale, + const int window_size_left, const int window_size_right); + +at::Tensor flash_attention_grouped_verify_paged( + const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, + std::optional& out_, const at::Tensor& block_table, + const at::Tensor& seq_lens, at::Tensor& partial_out, + at::Tensor& partial_lse, const float softmax_scale, + const std::string& kv_cache_dtype, const float k_scale, const float v_scale, + const bool one_pass); at::Tensor flash_attention_decode_paged_wmma( - const at::Tensor& q, - const at::Tensor& k_cache, - const at::Tensor& v_cache, - std::optional& out_, - const at::Tensor& block_table, - const at::Tensor& seq_lens, - const float softmax_scale, - const std::string& kv_cache_dtype, - const float k_scale, - const float v_scale -); + const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, + std::optional& out_, const at::Tensor& block_table, + const at::Tensor& seq_lens, const float softmax_scale, + const std::string& kv_cache_dtype, const float k_scale, + const float v_scale); at::Tensor flash_attention_decode_qk_scores( - const at::Tensor& q, - const at::Tensor& k_cache, - const at::Tensor& block_table, - const at::Tensor& seq_lens, - const float softmax_scale, - const int partition_size, - const std::string& kv_cache_dtype, - const float k_scale -); + const at::Tensor& q, const at::Tensor& k_cache, + const at::Tensor& block_table, const at::Tensor& seq_lens, + const float softmax_scale, const int partition_size, + const std::string& kv_cache_dtype, const float k_scale); 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); at::Tensor flash_attention_prefill_paged( - const at::Tensor& q, - const at::Tensor& k_cache, - const at::Tensor& v_cache, - std::optional& out_, - const at::Tensor& block_table, - const at::Tensor& seq_lens, - const float softmax_scale, - const std::string& kv_cache_dtype, - const float k_scale, - const float v_scale, - const bool is_causal, - const int window_size_left, - const int window_size_right -); + const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, + std::optional& out_, const at::Tensor& block_table, + const at::Tensor& seq_lens, const float softmax_scale, + const std::string& kv_cache_dtype, const float k_scale, const float v_scale, + const bool is_causal, const int window_size_left, + const int window_size_right, const std::optional& anchor_lens, + const int64_t anchored_window); std::vector flash_attention_prefill_paged_d256_bm32_allp_pair_scratch( - const at::Tensor& q, - const at::Tensor& k_cache, - const at::Tensor& v_cache, - std::optional& out_, - std::optional& softmax_lse_, - const at::Tensor& block_table, - const at::Tensor& seq_lens, - const float softmax_scale -); + const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, + std::optional& out_, std::optional& softmax_lse_, + const at::Tensor& block_table, const at::Tensor& seq_lens, + const float softmax_scale); std::vector flash_attention_prefill_paged_d256_bm32_allp_pair_scratch_splitkv3( - const at::Tensor& q, - const at::Tensor& k_cache, - const at::Tensor& v_cache, - std::optional& out_, - std::optional& softmax_lse_, - at::Tensor& split_tmp_out, - at::Tensor& split_tmp_row_max, - at::Tensor& split_tmp_row_sum, - const at::Tensor& block_table, - const int64_t actual_n, - const float softmax_scale -); + const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, + std::optional& out_, std::optional& softmax_lse_, + at::Tensor& split_tmp_out, at::Tensor& split_tmp_row_max, + at::Tensor& split_tmp_row_sum, const at::Tensor& block_table, + const int64_t actual_n, const float softmax_scale); at::Tensor flash_attention_prefill_paged_bfla( - const at::Tensor& q, - const at::Tensor& k_cache, - const at::Tensor& v_cache, - std::optional& out_, - const at::Tensor& block_table, - const at::Tensor& seq_lens, - const at::Tensor& bfla_block_mask, - const int bfla_mask_block_n, - const float softmax_scale, - const std::string& kv_cache_dtype, - const float k_scale, - const float v_scale, - const bool is_causal, - const int window_size_left, - const int window_size_right -); + const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, + std::optional& out_, const at::Tensor& block_table, + const at::Tensor& seq_lens, const at::Tensor& bfla_block_mask, + const int bfla_mask_block_n, const float softmax_scale, + const std::string& kv_cache_dtype, const float k_scale, const float v_scale, + const bool is_causal, const int window_size_left, + const int window_size_right); at::Tensor flash_attention_prefill_paged_splitkv( - const at::Tensor& q, - const at::Tensor& k_cache, - const at::Tensor& v_cache, - std::optional& out_, - const at::Tensor& block_table, - const at::Tensor& seq_lens, - const float softmax_scale, - const std::string& kv_cache_dtype, - const float k_scale, - const float v_scale, - const bool is_causal, - const int window_size_left, - const int window_size_right, - const int split_kv_tokens, - const int max_seq_len_hint -); + const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, + std::optional& out_, const at::Tensor& block_table, + const at::Tensor& seq_lens, const float softmax_scale, + const std::string& kv_cache_dtype, const float k_scale, const float v_scale, + const bool is_causal, const int window_size_left, + const int window_size_right, const int split_kv_tokens, + const int max_seq_len_hint); void flash_attention_fp8_e5m2_paged_kv_to_fp16( - const at::Tensor& key_cache, - const at::Tensor& value_cache, - const at::Tensor& block_table, - const at::Tensor& seq_lens, - at::Tensor& key_out, - at::Tensor& value_out, - const float key_scale, - const float value_scale -); + const at::Tensor& key_cache, const at::Tensor& value_cache, + const at::Tensor& block_table, const at::Tensor& seq_lens, + at::Tensor& key_out, at::Tensor& value_out, const float key_scale, + const float value_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 -); + 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); #endif 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 f336475a35..b6417092d4 100644 --- a/flash-attention-v100/kernel/flash_decode_paged.cu +++ b/flash-attention-v100/kernel/flash_decode_paged.cu @@ -4,6 +4,8 @@ #include #include #include +#include +#include #include #include #include @@ -18,7 +20,8 @@ namespace { int kv_cache_dtype_code_from_string(const std::string& kv_cache_dtype) { - if (kv_cache_dtype == "auto" || kv_cache_dtype == "bfloat16") { + if (kv_cache_dtype == "auto" || kv_cache_dtype == "float16" || + kv_cache_dtype == "bfloat16") { return flash_v100::KV_CACHE_DTYPE_FP16; } if (kv_cache_dtype == "fp8" || kv_cache_dtype == "fp8_e4m3") { @@ -49,6 +52,8 @@ constexpr int kXQARouteP1024Sawtooth = 4; constexpr int kXQARouteP256Sawtooth = 5; constexpr int kXQARouteP1024SawtoothMid = 6; constexpr int kXQARouteP1024SawtoothFinal = 7; +constexpr int kXQARouteRangeSeqLens = 8; +constexpr int kXQARouteWaveLongSeqLens = 9; // The dense 256/128-half strides alias Volta shared-memory banks in the // WMMA A/B fragment loads. Both padded strides remain 16-byte aligned. constexpr int kXQATC256WidePaddedQStride = 264; @@ -145,6 +150,41 @@ bool xqa_g6_dual_cta_enabled() { return value != nullptr && value[0] == '1'; } +bool xqa_e4m3_batch_enabled() { + const char* value = std::getenv("VLLM_FLASH_V100_E4M3_BATCH_XQA"); + return value == nullptr || value[0] != '0'; +} + +bool xqa_e4m3_batch_optimized_enabled() { + const char* value = std::getenv("VLLM_FLASH_V100_E4M3_BATCH_XQA_OPTIMIZED"); + return value == nullptr || value[0] != '0'; +} + +bool xqa_e4m3_page800_fastpath_enabled() { + const char* value = std::getenv("VLLM_FLASH_V100_E4M3_PAGE800_FASTPATH"); + return value == nullptr || value[0] != '0'; +} + +bool xqa_e4m3_page800_fastpath_trace_enabled() { + const char* value = + std::getenv("VLLM_FLASH_V100_E4M3_PAGE800_FASTPATH_TRACE"); + return value != nullptr && value[0] == '1'; +} + +void trace_xqa_e4m3_page800_fastpath(const int batch_size, + const int partition_size, + const bool active) { + if (!active || !xqa_e4m3_page800_fastpath_trace_enabled()) { + return; + } + static std::atomic traced{false}; + if (!traced.exchange(true, std::memory_order_relaxed)) { + TORCH_WARN("Flash-V100 E4M3 page800 fast path active: batch=", batch_size, + ", partition_size=", partition_size, + ", standard interleaved Hkv=1, PV=half2"); + } +} + bool xqa_e5m2_g6_dual_cta_enabled() { const char* value = std::getenv("VLLM_FLASH_V100_XQA_E5M2_G6_DUAL_CTA"); return value == nullptr || value[0] != '0'; @@ -166,6 +206,11 @@ bool xqa_e5m2_pair_load_enabled() { return value == nullptr || value[0] != '0'; } +bool xqa_e5m2_batch_wide_load_enabled() { + const char* value = std::getenv("VLLM_FLASH_V100_XQA_E5M2_BATCH_WIDE_LOAD"); + 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)); @@ -213,6 +258,57 @@ bool xqa_g6_p1024_sawtooth_enabled() { return value == nullptr || value[0] != '0'; } +bool xqa_e4m3_g6_p64_p256_auto_enabled() { + const char* value = std::getenv("VLLM_FLASH_V100_XQA_E4M3_G6_P64_P256_AUTO"); + return value != nullptr && value[0] == '1'; +} + +bool xqa_e4m3_g6_p64_p256_auto_trace_enabled() { + static const bool enabled = [] { + const char* value = + std::getenv("VLLM_FLASH_V100_XQA_E4M3_G6_P64_P256_AUTO_TRACE"); + return value != nullptr && value[0] == '1'; + }(); + return enabled; +} + +int xqa_e4m3_g6_p256_begin() { + const char* value = std::getenv("VLLM_FLASH_V100_XQA_E4M3_G6_P256_BEGIN"); + return value == nullptr ? 12288 : std::max(1, std::atoi(value)); +} + +int xqa_e4m3_g6_dual_cta_begin() { + const char* value = std::getenv("VLLM_FLASH_V100_XQA_E4M3_G6_DUAL_CTA_BEGIN"); + return value == nullptr ? 32768 : std::max(1, std::atoi(value)); +} + +bool xqa_e4m3_g6_wave_partitions_enabled() { + const char* value = + std::getenv("VLLM_FLASH_V100_XQA_E4M3_G6_WAVE_PARTITIONS"); + return value != nullptr && value[0] == '1'; +} + +bool xqa_e4m3_g6_merged_wave_launch_enabled() { + const char* value = + std::getenv("VLLM_FLASH_V100_XQA_E4M3_G6_MERGED_WAVE_LAUNCH"); + return value != nullptr && value[0] == '1'; +} + +int xqa_e4m3_g6_p512_begin() { + const char* value = std::getenv("VLLM_FLASH_V100_XQA_E4M3_G6_P512_BEGIN"); + return value == nullptr ? 49152 : std::max(1, std::atoi(value)); +} + +int xqa_e4m3_g6_p896_begin() { + const char* value = std::getenv("VLLM_FLASH_V100_XQA_E4M3_G6_P896_BEGIN"); + return value == nullptr ? 98304 : std::max(1, std::atoi(value)); +} + +int xqa_e4m3_g6_p1664_begin() { + const char* value = std::getenv("VLLM_FLASH_V100_XQA_E4M3_G6_P1664_BEGIN"); + return value == nullptr ? 196608 : std::max(1, std::atoi(value)); +} + bool decode_partition_size_overridden() { return std::getenv("VLLM_FLASH_V100_DECODE_PARTITION_SIZE") != nullptr; } @@ -267,6 +363,102 @@ bool xqa_split_reduce_enabled() { return value != nullptr && value[0] == '1'; } +enum class XQABatchContextRoute : int { + kDisabled = -1, + kBaseline = 0, + kDualCta = 1, + kDualCtaSplit = 2, +}; + +bool xqa_batch_context_routing_enabled() { + const char* value = std::getenv("VLLM_FLASH_V100_XQA_BATCH_CONTEXT_ROUTING"); + return value == nullptr || value[0] != '0'; +} + +bool xqa_batch_context_routing_trace_enabled() { + const char* value = + std::getenv("VLLM_FLASH_V100_XQA_BATCH_CONTEXT_ROUTING_TRACE"); + return value != nullptr && value[0] == '1'; +} + +XQABatchContextRoute select_xqa_batch_context_route(const int batch_size, + const int max_seq_len, + const int partition_size) { + if (!xqa_batch_context_routing_enabled() || batch_size < 4 || + max_seq_len <= 0) { + return XQABatchContextRoute::kDisabled; + } + if (batch_size <= 4) { + if (max_seq_len <= 8191) { + return XQABatchContextRoute::kBaseline; + } + if (max_seq_len <= 12287) { + return XQABatchContextRoute::kDualCta; + } + return partition_size == 1024 ? XQABatchContextRoute::kBaseline + : XQABatchContextRoute::kDualCtaSplit; + } + if (batch_size <= 8) { + if (max_seq_len <= 4095) { + return XQABatchContextRoute::kBaseline; + } + return max_seq_len <= 12287 ? XQABatchContextRoute::kDualCta + : XQABatchContextRoute::kDualCtaSplit; + } + if (batch_size <= 12) { + if (max_seq_len <= 2047) { + return XQABatchContextRoute::kBaseline; + } + return max_seq_len <= 16000 ? XQABatchContextRoute::kDualCta + : XQABatchContextRoute::kDualCtaSplit; + } + return max_seq_len <= 12287 ? XQABatchContextRoute::kBaseline + : XQABatchContextRoute::kDualCta; +} + +const char* xqa_batch_context_route_name(const XQABatchContextRoute route) { + switch (route) { + case XQABatchContextRoute::kBaseline: + return "baseline"; + case XQABatchContextRoute::kDualCta: + return "dual_cta"; + case XQABatchContextRoute::kDualCtaSplit: + return "dual_cta_split"; + default: + return "disabled"; + } +} + +void trace_xqa_batch_context_route(const int batch_size, const int max_seq_len, + const int partition_size, + const int block_size, + const XQABatchContextRoute route) { + if (!xqa_batch_context_routing_trace_enabled() || + route == XQABatchContextRoute::kDisabled) { + return; + } + const int batch_class = batch_size <= 4 ? 0 + : batch_size <= 8 ? 1 + : batch_size <= 12 ? 2 + : 3; + const int context_class = max_seq_len <= 4095 ? 0 + : max_seq_len <= 8191 ? 1 + : max_seq_len <= 12287 ? 2 + : max_seq_len <= 16000 ? 3 + : 4; + const unsigned long long bit = + 1ULL << (batch_class * 15 + context_class * 3 + static_cast(route)); + static std::atomic traced_routes{0}; + const unsigned long long previous = + traced_routes.fetch_or(bit, std::memory_order_relaxed); + if ((previous & bit) == 0) { + TORCH_WARN("Flash-V100 XQA batch/context route active: batch=", batch_size, + ", max_seq_len=", max_seq_len, + ", partition_size=", partition_size, ", block_size=", block_size, + ", route=", xqa_batch_context_route_name(route)); + } +} + int xqa_block16_layout_mode() { const char* value = std::getenv("VLLM_FLASH_V100_XQA_BLOCK16_LAYOUT"); if (value == nullptr) { @@ -343,6 +535,10 @@ __device__ __forceinline__ bool xqa_seq_len_route_active( return seq_len >= route_seq_len_begin && seq_len < route_seq_len_end; } else if constexpr (SEQ_LEN_ROUTE == kXQARouteP1024SawtoothFinal) { return seq_len >= route_seq_len_final; + } else if constexpr (SEQ_LEN_ROUTE == kXQARouteRangeSeqLens) { + return seq_len >= route_seq_len_begin && seq_len < route_seq_len_end; + } else if constexpr (SEQ_LEN_ROUTE == kXQARouteWaveLongSeqLens) { + return seq_len >= route_seq_len_begin; } return true; } @@ -434,22 +630,115 @@ __device__ __forceinline__ uint4 fp8_e5m2_vector_to_half8(const uint64_t raw) { fp8_e5m2_pair_to_half2_bits(static_cast(raw >> 48))); } +__device__ __forceinline__ uint16_t fp8_e4m3fn_to_half_bits(const uint8_t raw) { + const uint16_t sign = static_cast(raw & 0x80u) << 8; + const uint8_t magnitude = raw & 0x7fu; + const uint8_t exponent = magnitude >> 3; + const uint8_t mantissa = magnitude & 0x07u; + if (magnitude == 0) { + return sign; + } + if (exponent == 0) { + // E4M3 subnormals are exact fp16 normals: mantissa * 2^-9. + const uint16_t magnitude_bits = + mantissa < 2 + ? 0x1800u + : (mantissa < 4 + ? static_cast(0x1c00u | ((mantissa - 2) << 9)) + : static_cast(0x2000u | ((mantissa - 4) << 8))); + return sign | magnitude_bits; + } + if (magnitude == 0x7fu) { + return sign | 0x7e00u; + } + return sign | static_cast((exponent + 8) << 10) | + static_cast(mantissa << 7); +} + +__device__ __forceinline__ uint32_t +fp8_e4m3fn_pair_to_half2_bits(const uint16_t raw_pair) { + return static_cast( + fp8_e4m3fn_to_half_bits(static_cast(raw_pair))) | + (static_cast( + fp8_e4m3fn_to_half_bits(static_cast(raw_pair >> 8))) + << 16); +} + +__device__ __forceinline__ uint32_t +fp8_e4m3fn_pair_to_half2_bits_fast(const uint16_t raw_pair) { + const uint8_t raw0 = static_cast(raw_pair); + const uint8_t raw1 = static_cast(raw_pair >> 8); + if ((raw0 & 0x7fu) == 0x7fu || (raw1 & 0x7fu) == 0x7fu) { + return fp8_e4m3fn_pair_to_half2_bits(raw_pair); + } + + // Moving a finite E4M3 encoding into the corresponding fp16 sign, + // exponent, and mantissa fields represents exactly value / 256. A packed + // half2 multiply restores both values without per-byte exponent branches. + const uint32_t expanded = (static_cast(raw_pair & 0x0080u) << 8) | + (static_cast(raw_pair & 0x007fu) << 7) | + (static_cast(raw_pair & 0x8000u) << 16) | + (static_cast(raw_pair & 0x7f00u) << 15); + union { + uint32_t u; + __half2 h2; + } converter; + converter.u = expanded; + converter.h2 = __hmul2(converter.h2, __float2half2_rn(256.0f)); + return converter.u; +} + +__device__ __forceinline__ uint4 +fp8_e4m3fn_vector_to_half8(const uint64_t raw) { + return make_uint4( + fp8_e4m3fn_pair_to_half2_bits(static_cast(raw)), + fp8_e4m3fn_pair_to_half2_bits(static_cast(raw >> 16)), + fp8_e4m3fn_pair_to_half2_bits(static_cast(raw >> 32)), + fp8_e4m3fn_pair_to_half2_bits(static_cast(raw >> 48))); +} + +__device__ __forceinline__ uint4 +fp8_e4m3fn_vector_to_half8_fast(const uint64_t raw) { + return make_uint4( + fp8_e4m3fn_pair_to_half2_bits_fast(static_cast(raw)), + fp8_e4m3fn_pair_to_half2_bits_fast(static_cast(raw >> 16)), + fp8_e4m3fn_pair_to_half2_bits_fast(static_cast(raw >> 32)), + fp8_e4m3fn_pair_to_half2_bits_fast(static_cast(raw >> 48))); +} + +__device__ __forceinline__ uint4 fp8_e4m3fn_vector_to_half8_lut( + const uint64_t raw, const uint16_t* __restrict__ lut) { + return make_uint4( + static_cast(lut[static_cast(raw)]) | + (static_cast(lut[static_cast(raw >> 8)]) << 16), + static_cast(lut[static_cast(raw >> 16)]) | + (static_cast(lut[static_cast(raw >> 24)]) << 16), + static_cast(lut[static_cast(raw >> 32)]) | + (static_cast(lut[static_cast(raw >> 40)]) << 16), + static_cast(lut[static_cast(raw >> 48)]) | + (static_cast(lut[static_cast(raw >> 56)]) << 16)); +} + template + int KV_DTYPE = flash_v100::KV_CACHE_DTYPE_FP16, + bool E4M3_SHARED_LUT = false> __device__ __forceinline__ uint4 load_xqa_tc_kv_vector( const void* __restrict__ kv_cache, const int* __restrict__ page_ids, const int copy_idx, const int panel_d_stride_uint4, const int tile_page_offset, const int kv_tile_start, const int block_size, const int kv_head_idx, const int64_t block_stride, const int64_t token_stride, const int64_t head_stride, - const int panel_offset) { + const int panel_offset, const uint16_t* __restrict__ e4m3_lut = nullptr) { const int row = copy_idx / panel_d_stride_uint4; const int vec_col = copy_idx % panel_d_stride_uint4; const int token_offset = tile_page_offset + kv_tile_start + row; - static_assert(BLOCK_SIZE == 0 || BLOCK_SIZE == 16 || BLOCK_SIZE == 784, + static_assert(BLOCK_SIZE == 0 || BLOCK_SIZE == 16 || BLOCK_SIZE == 784 || + BLOCK_SIZE == 800 || BLOCK_SIZE == 1568 || + BLOCK_SIZE == 1648 || BLOCK_SIZE == 3296, "Unsupported paged-KV block-size specialization"); - static_assert(!CONTIGUOUS_HKV1_LAYOUT || BLOCK_SIZE == 16, - "The contiguous Hkv=1 layout requires 16-token pages"); + static_assert(!CONTIGUOUS_HKV1_LAYOUT || BLOCK_SIZE == 16 || + BLOCK_SIZE == 800 || BLOCK_SIZE == 1568, + "The fixed-stride Hkv=1 layout requires a specialized page"); int logical_block; int block_offset; if constexpr (BLOCK_SIZE == 16) { @@ -458,6 +747,18 @@ __device__ __forceinline__ uint4 load_xqa_tc_kv_vector( } else if constexpr (BLOCK_SIZE == 784) { logical_block = token_offset / 784; block_offset = token_offset - logical_block * 784; + } else if constexpr (BLOCK_SIZE == 800) { + logical_block = token_offset >= 800; + block_offset = token_offset - logical_block * 800; + } else if constexpr (BLOCK_SIZE == 1568) { + logical_block = token_offset / 1568; + block_offset = token_offset - logical_block * 1568; + } else if constexpr (BLOCK_SIZE == 1648) { + logical_block = token_offset / 1648; + block_offset = token_offset - logical_block * 1648; + } else if constexpr (BLOCK_SIZE == 3296) { + logical_block = token_offset / 3296; + block_offset = token_offset - logical_block * 3296; } else { logical_block = token_offset / block_size; block_offset = token_offset % block_size; @@ -465,8 +766,9 @@ __device__ __forceinline__ uint4 load_xqa_tc_kv_vector( const int physical_block = page_ids[logical_block]; int64_t physical_offset; if constexpr (CONTIGUOUS_HKV1_LAYOUT) { - constexpr int kHeadDim = 256; - constexpr int kBlockStride = 16 * kHeadDim; + constexpr int64_t kHeadDim = 256; + constexpr int64_t kBlockStride = + BLOCK_SIZE == 16 ? 16 * kHeadDim : 2 * BLOCK_SIZE * kHeadDim; physical_offset = static_cast(physical_block) * kBlockStride + static_cast(block_offset) * kHeadDim + panel_offset; @@ -480,17 +782,28 @@ __device__ __forceinline__ uint4 load_xqa_tc_kv_vector( const uint4* cache_vec = reinterpret_cast(kv_cache); return __ldg(&cache_vec[physical_offset / 8 + vec_col]); } else { - static_assert(KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP8_E5M2, - "XQA only supports fp16 and FP8 E5M2 KV"); + static_assert(KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP8_E4M3 || + KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP8_E5M2, + "XQA only supports fp16, FP8 E4M3, and FP8 E5M2 KV"); const uint64_t* cache_vec = reinterpret_cast(kv_cache); - return fp8_e5m2_vector_to_half8( - __ldg(&cache_vec[physical_offset / 8 + vec_col])); + const uint64_t raw = __ldg(&cache_vec[physical_offset / 8 + vec_col]); + if constexpr (KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP8_E4M3) { + if constexpr (E4M3_SHARED_LUT) { + return fp8_e4m3fn_vector_to_half8_lut(raw, e4m3_lut); + } else { + return fp8_e4m3fn_vector_to_half8(raw); + } + } else { + static_assert(!E4M3_SHARED_LUT, + "The E4M3 conversion LUT requires E4M3 KV"); + return fp8_e5m2_vector_to_half8(raw); + } } } template + bool FP8_PAIR_LOAD = false, bool E4M3_SHARED_LUT = false> __device__ __forceinline__ void load_xqa_tc_kv_panel( __half* __restrict__ shared_kv, const void* __restrict__ kv_cache, const int* __restrict__ page_ids, const int valid_kv_tile_rows, @@ -498,11 +811,15 @@ __device__ __forceinline__ void load_xqa_tc_kv_panel( const int tile_page_offset, const int kv_tile_start, const int block_size, const int kv_head_idx, const int64_t block_stride, const int64_t token_stride, const int64_t head_stride, - const int panel_offset, const int copy_thread_idx = threadIdx.x) { + const int panel_offset, const int copy_thread_idx = threadIdx.x, + const uint16_t* __restrict__ e4m3_lut = nullptr) { uint4* shared_vec = reinterpret_cast(shared_kv); - if constexpr (E5M2_PAIR_LOAD) { - static_assert(KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP8_E5M2, - "Paired XQA loads are only valid for FP8 E5M2 KV"); + if constexpr (FP8_PAIR_LOAD) { + static_assert(KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP8_E4M3 || + KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP8_E5M2, + "Paired XQA loads require FP8 KV"); + static_assert(!E4M3_SHARED_LUT, + "Paired E4M3 conversion does not use the shared LUT"); const int pair_stride = panel_d_stride_uint4 / 2; const int pair_count = valid_kv_tile_rows * pair_stride; for (int pair_idx = copy_thread_idx; pair_idx < pair_count; @@ -518,6 +835,18 @@ __device__ __forceinline__ void load_xqa_tc_kv_panel( } else if constexpr (BLOCK_SIZE == 784) { logical_block = token_offset / 784; block_offset = token_offset - logical_block * 784; + } else if constexpr (BLOCK_SIZE == 800) { + logical_block = token_offset >= 800; + block_offset = token_offset - logical_block * 800; + } else if constexpr (BLOCK_SIZE == 1568) { + logical_block = token_offset / 1568; + block_offset = token_offset - logical_block * 1568; + } else if constexpr (BLOCK_SIZE == 1648) { + logical_block = token_offset / 1648; + block_offset = token_offset - logical_block * 1648; + } else if constexpr (BLOCK_SIZE == 3296) { + logical_block = token_offset / 3296; + block_offset = token_offset - logical_block * 3296; } else { logical_block = token_offset / block_size; block_offset = token_offset % block_size; @@ -525,11 +854,12 @@ __device__ __forceinline__ void load_xqa_tc_kv_panel( const int physical_block = page_ids[logical_block]; int64_t physical_offset; if constexpr (CONTIGUOUS_HKV1_LAYOUT) { - constexpr int kHeadDim = 256; - constexpr int kBlockStride = 16 * kHeadDim; - physical_offset = static_cast(physical_block) * kBlockStride + - static_cast(block_offset) * kHeadDim + - panel_offset; + constexpr int64_t kHeadDim = 256; + constexpr int64_t kPhysicalBlockStride = + BLOCK_SIZE == 16 ? 16 * kHeadDim : 2 * 800 * kHeadDim; + physical_offset = + static_cast(physical_block) * kPhysicalBlockStride + + static_cast(block_offset) * kHeadDim + panel_offset; } else { physical_offset = static_cast(physical_block) * block_stride + static_cast(block_offset) * token_stride + @@ -541,10 +871,16 @@ __device__ __forceinline__ void load_xqa_tc_kv_panel( const int shared_offset = row * kv_smem_stride_uint4 + vec_pair * 2; const uint64_t raw_lo = static_cast(raw.x) | (static_cast(raw.y) << 32); - shared_vec[shared_offset] = fp8_e5m2_vector_to_half8(raw_lo); + shared_vec[shared_offset] = + KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP8_E4M3 + ? fp8_e4m3fn_vector_to_half8_fast(raw_lo) + : fp8_e5m2_vector_to_half8(raw_lo); const uint64_t raw_hi = static_cast(raw.z) | (static_cast(raw.w) << 32); - shared_vec[shared_offset + 1] = fp8_e5m2_vector_to_half8(raw_hi); + shared_vec[shared_offset + 1] = + KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP8_E4M3 + ? fp8_e4m3fn_vector_to_half8_fast(raw_hi) + : fp8_e5m2_vector_to_half8(raw_hi); } } else { const int copy_count = valid_kv_tile_rows * panel_d_stride_uint4; @@ -553,16 +889,18 @@ __device__ __forceinline__ void load_xqa_tc_kv_panel( const int row = copy_idx / panel_d_stride_uint4; const int vec_col = copy_idx % panel_d_stride_uint4; shared_vec[row * kv_smem_stride_uint4 + vec_col] = - load_xqa_tc_kv_vector( + load_xqa_tc_kv_vector( kv_cache, page_ids, copy_idx, panel_d_stride_uint4, tile_page_offset, kv_tile_start, block_size, kv_head_idx, - block_stride, token_stride, head_stride, panel_offset); + block_stride, token_stride, head_stride, panel_offset, e4m3_lut); } } } template + int KV_DTYPE = flash_v100::KV_CACHE_DTYPE_FP16, + bool E4M3_SHARED_LUT = false> __device__ __forceinline__ void load_xqa_tc_kv_panel_and_zero( __half* __restrict__ shared_kv, const void* __restrict__ kv_cache, const int* __restrict__ page_ids, const int valid_kv_tile_rows, @@ -570,7 +908,8 @@ __device__ __forceinline__ void load_xqa_tc_kv_panel_and_zero( const int tile_page_offset, const int kv_tile_start, const int block_size, const int kv_head_idx, const int64_t block_stride, const int64_t token_stride, const int64_t head_stride, - const int panel_offset, const int copy_thread_idx) { + const int panel_offset, const int copy_thread_idx, + const uint16_t* __restrict__ e4m3_lut = nullptr) { const int copy_count = valid_kv_tile_rows * panel_d_stride_uint4; uint4* shared_vec = reinterpret_cast(shared_kv); constexpr int kLoadStages = 4; @@ -582,10 +921,12 @@ __device__ __forceinline__ void load_xqa_tc_kv_panel_and_zero( const int copy_idx = copy_base + stage * NUM_THREADS; if (copy_idx < copy_count) { staged[stage] = - load_xqa_tc_kv_vector( + load_xqa_tc_kv_vector( kv_cache, page_ids, copy_idx, panel_d_stride_uint4, tile_page_offset, kv_tile_start, block_size, kv_head_idx, - block_stride, token_stride, head_stride, panel_offset); + block_stride, token_stride, head_stride, panel_offset, + e4m3_lut); } } #pragma unroll @@ -663,7 +1004,7 @@ __device__ __forceinline__ float dot_qk_cache(const __half* __restrict__ q_ptr, } template + int SEQ_LEN_ROUTE = kXQARouteAllSeqLens, bool ANCHORED_SWA = false> __global__ void flash_attention_decode_partition_kernel( const __half* __restrict__ q, const void* __restrict__ k_cache, const void* __restrict__ v_cache, __half* __restrict__ tmp_out, @@ -681,7 +1022,12 @@ __global__ void flash_attention_decode_partition_kernel( const int64_t v_head_stride, const float softmax_scale, const float k_scale, const float v_scale, const int window_size_left, const int window_size_right, const int route_seq_len_begin, - const int route_seq_len_end, const int route_seq_len_final) { + const int route_seq_len_end, const int route_seq_len_final, + const int* __restrict__ anchor_lens, const int anchored_window) { + // The anchored decode-window mask variant is only generated for the fp16-KV + // configuration; every other instantiation keeps its original code path. + static_assert(!ANCHORED_SWA || KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP16, + "ANCHORED_SWA requires an fp16 KV cache"); const int batch_idx = blockIdx.x; const int head_idx = blockIdx.y; const int partition_idx = blockIdx.z; @@ -718,6 +1064,17 @@ __global__ void flash_attention_decode_partition_kernel( : seq_len - 1; const int part_start = max(start_token_idx, min_token_idx); const int part_end = min(start_token_idx + PARTITION_SIZE, max_token_idx + 1); + + // Anchored decode window: keys below the per-request prompt length stay + // globally visible; later (generated) keys must fall inside the sliding + // window over generated positions. Keys inside [gap_lo, gap_hi) are masked. + int gap_lo = seq_len; + int gap_hi = seq_len; + if constexpr (ANCHORED_SWA) { + const int anchor_len = anchor_lens[batch_idx]; + gap_lo = anchor_len; + gap_hi = max(anchor_len, query_pos - anchored_window + 1); + } const int q_per_kv = num_heads_q / num_heads_kv; const int kv_head_idx = head_idx / q_per_kv; const int lane = threadIdx.x % kWarpSize; @@ -777,9 +1134,21 @@ __global__ void flash_attention_decode_partition_kernel( float score = dot_qk_cache(q_shared, k_cache, k_index, lane); if (lane == 0) { - score *= score_scale; - scores_shared[token_local] = score; - local_max = fmaxf(local_max, score); + if constexpr (ANCHORED_SWA) { + const int token_idx = part_start + token_local; + if (token_idx >= gap_lo && token_idx < gap_hi) { + // Masked gap key: never contributes (nulled gap blocks included). + scores_shared[token_local] = -1.0e30f; + } else { + score *= score_scale; + scores_shared[token_local] = score; + local_max = fmaxf(local_max, score); + } + } else { + score *= score_scale; + scores_shared[token_local] = score; + local_max = fmaxf(local_max, score); + } } } @@ -787,6 +1156,14 @@ __global__ void flash_attention_decode_partition_kernel( float local_sum = 0.f; for (int i = threadIdx.x; i < part_tokens; i += blockDim.x) { + if constexpr (ANCHORED_SWA) { + const int token_idx = part_start + i; + if (token_idx >= gap_lo && token_idx < gap_hi) { + // Exact zero weight for masked keys (no reliance on expf underflow). + scores_shared[i] = 0.f; + continue; + } + } const float p = __expf(scores_shared[i] - part_max); scores_shared[i] = p; local_sum += p; @@ -827,7 +1204,8 @@ template + bool PARTITION_PAGE_IDS = false, bool FP8_PAIR_LOAD = false, + bool E4M3_SHARED_LUT = false> __global__ void __launch_bounds__(NUM_THREADS, MIN_BLOCKS_PER_SM) flash_attention_decode_xqa_tc_partition_kernel_256_wide( const __half* __restrict__ q, const void* __restrict__ k_cache, @@ -871,6 +1249,9 @@ __global__ void __launch_bounds__(NUM_THREADS, MIN_BLOCKS_PER_SM) QK_SW_PIPELINE ? SmemLayout::kQKStride / 8 : kv_smem_stride_uint4; constexpr int qk_panel_d_stride_uint4 = kQKPanelDim / 8; constexpr int kAccumsPerThread = D / kWarpSize; + constexpr bool kPVHalf2 = CONTIGUOUS_HKV1_LAYOUT && + (BLOCK_SIZE == 800 || BLOCK_SIZE == 1568) && + KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP8_E4M3; static_assert(GROUP_SIZE == 4 || GROUP_SIZE == 6 || GROUP_SIZE == 8, "Wide D=256 TC XQA kernel supports q_per_kv in {4, 6, 8}"); static_assert(NUM_THREADS >= GROUP_SIZE * kWarpSize, @@ -880,6 +1261,9 @@ __global__ void __launch_bounds__(NUM_THREADS, MIN_BLOCKS_PER_SM) NUM_THREADS == 8 * kWarpSize)), "The QK software pipeline requires a six- or eight-warp G6 " "kernel"); + static_assert( + !E4M3_SHARED_LUT || KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP8_E4M3, + "The shared conversion LUT is only valid for E4M3 KV"); const int batch_idx = blockIdx.x; const int kv_head_idx = blockIdx.y; const int partition_idx = blockIdx.z; @@ -895,13 +1279,27 @@ __global__ void __launch_bounds__(NUM_THREADS, MIN_BLOCKS_PER_SM) route_seq_len_final)) { return; } - const int start_token_idx = partition_idx * PARTITION_SIZE; + int partition_size; + if constexpr (PARTITION_SIZE == -1) { + static_assert(SEQ_LEN_ROUTE == kXQARouteWaveLongSeqLens, + "Runtime wave partitions require the wave-long route"); + partition_size = seq_len < route_seq_len_end + ? 512 + : (seq_len < route_seq_len_final ? 896 : 1664); + } else if constexpr (PARTITION_SIZE == 0) { + const int partition_size_begin = + route_seq_len_end > 0 ? route_seq_len_end : route_seq_len_begin; + partition_size = seq_len < partition_size_begin ? 64 : 256; + } else { + partition_size = PARTITION_SIZE; + } + const int start_token_idx = partition_idx * partition_size; if (seq_len <= 0 || start_token_idx >= seq_len) { return; } const int runtime_num_partitions = active_num_partitions[0]; const int seq_num_partitions = - (seq_len + PARTITION_SIZE - 1) / PARTITION_SIZE; + (seq_len + partition_size - 1) / partition_size; const int effective_num_partitions = min(max_num_partitions, max(runtime_num_partitions, seq_num_partitions)); if (partition_idx >= effective_num_partitions) { @@ -916,12 +1314,19 @@ __global__ void __launch_bounds__(NUM_THREADS, MIN_BLOCKS_PER_SM) const int tid = threadIdx.x; const int warp_id = tid / kWarpSize; const int lane_id = tid % kWarpSize; - const int part_tokens = min(PARTITION_SIZE, seq_len - start_token_idx); + const int part_tokens = min(partition_size, seq_len - start_token_idx); const int num_k_tiles = (part_tokens + kXQATCBlockN - 1) / kXQATCBlockN; const int* block_table_seq = block_table + batch_idx * max_num_blocks; extern __shared__ char smem_raw[]; auto& smem = *reinterpret_cast(smem_raw); + uint16_t* e4m3_lut = nullptr; + if constexpr (E4M3_SHARED_LUT) { + e4m3_lut = reinterpret_cast(smem_raw + sizeof(SmemLayout)); + for (int raw = tid; raw < 256; raw += NUM_THREADS) { + e4m3_lut[raw] = fp8_e4m3fn_to_half_bits(static_cast(raw)); + } + } __half* sQ = smem.q; __half* sK = smem.k_buffer(0); __half* sV = smem.v_buffer(); @@ -959,10 +1364,18 @@ __global__ void __launch_bounds__(NUM_THREADS, MIN_BLOCKS_PER_SM) int partition_start_page = 0; int partition_page_offset = 0; if constexpr (PARTITION_PAGE_IDS) { - partition_start_page = start_token_idx / block_size; - partition_page_offset = start_token_idx - partition_start_page * block_size; - const int partition_page_count = - (partition_page_offset + part_tokens + block_size - 1) / block_size; + int partition_page_count; + if constexpr (BLOCK_SIZE == 800) { + partition_start_page = start_token_idx / 800; + partition_page_offset = start_token_idx - partition_start_page * 800; + partition_page_count = 1 + (partition_page_offset + part_tokens > 800); + } else { + partition_start_page = start_token_idx / block_size; + partition_page_offset = + start_token_idx - partition_start_page * block_size; + partition_page_count = + (partition_page_offset + part_tokens + block_size - 1) / block_size; + } for (int idx = tid; idx < partition_page_count; idx += NUM_THREADS) { smem.page_ids[idx] = __ldg(&block_table_seq[partition_start_page + idx]); } @@ -986,6 +1399,14 @@ __global__ void __launch_bounds__(NUM_THREADS, MIN_BLOCKS_PER_SM) start_page = tile_token_start / 784; tile_page_offset = tile_token_start - start_page * 784; page_count = (tile_page_offset + valid_k_rows + 783) / 784; + } else if constexpr (BLOCK_SIZE == 800) { + start_page = tile_token_start / 800; + tile_page_offset = tile_token_start - start_page * 800; + page_count = 1 + (tile_page_offset + valid_k_rows > 800); + } else if constexpr (BLOCK_SIZE == 1568) { + start_page = tile_token_start / 1568; + tile_page_offset = tile_token_start - start_page * 1568; + page_count = (tile_page_offset + valid_k_rows + 1567) / 1568; } else { start_page = tile_token_start / block_size; tile_page_offset = tile_token_start - start_page * block_size; @@ -1025,11 +1446,12 @@ __global__ void __launch_bounds__(NUM_THREADS, MIN_BLOCKS_PER_SM) const int producer_tid = tid - kConsumerWarps * kWarpSize; if (producer_tid >= 0) { load_xqa_tc_kv_panel_and_zero( + kProducerThreads, KV_DTYPE, + E4M3_SHARED_LUT>( smem.k_buffer(0), k_cache, smem.page_ids, valid_kv_tile_rows, qk_panel_d_stride_uint4, qk_smem_stride_uint4, tile_page_offset, kv_tile_start, block_size, kv_head_idx, k_block_stride, - k_token_stride, k_head_stride, 0, producer_tid); + k_token_stride, k_head_stride, 0, producer_tid, e4m3_lut); } __syncthreads(); @@ -1038,12 +1460,14 @@ __global__ void __launch_bounds__(NUM_THREADS, MIN_BLOCKS_PER_SM) const int panel_offset = panel_idx * kQKPanelDim; if (producer_tid >= 0 && panel_idx + 1 < kNumQKPanels) { load_xqa_tc_kv_panel_and_zero( + kProducerThreads, KV_DTYPE, + E4M3_SHARED_LUT>( smem.k_buffer((panel_idx + 1) & 1), k_cache, smem.page_ids, valid_kv_tile_rows, qk_panel_d_stride_uint4, qk_smem_stride_uint4, tile_page_offset, kv_tile_start, block_size, kv_head_idx, k_block_stride, k_token_stride, - k_head_stride, panel_offset + kQKPanelDim, producer_tid); + k_head_stride, panel_offset + kQKPanelDim, producer_tid, + e4m3_lut); } if (warp_id < kConsumerWarps) { const int tile_n = warp_id * WMMA_N; @@ -1067,11 +1491,12 @@ __global__ void __launch_bounds__(NUM_THREADS, MIN_BLOCKS_PER_SM) for (int panel_idx = 0; panel_idx < kNumQKPanels; ++panel_idx) { const int panel_offset = panel_idx * kQKPanelDim; load_xqa_tc_kv_panel( + KV_DTYPE, FP8_PAIR_LOAD, E4M3_SHARED_LUT>( sK, k_cache, smem.page_ids, valid_kv_tile_rows, qk_panel_d_stride_uint4, qk_smem_stride_uint4, tile_page_offset, kv_tile_start, block_size, kv_head_idx, k_block_stride, - k_token_stride, k_head_stride, panel_offset); + k_token_stride, k_head_stride, panel_offset, threadIdx.x, + e4m3_lut); for (int idx = tid + valid_kv_tile_rows * qk_panel_d_stride_uint4; idx < kXQATCBlockN * qk_panel_d_stride_uint4; idx += NUM_THREADS) { @@ -1216,11 +1641,11 @@ __global__ void __launch_bounds__(NUM_THREADS, MIN_BLOCKS_PER_SM) const int valid_kv_tile_rows = min(kXQATCBlockN, valid_k_rows - kv_tile_start); load_xqa_tc_kv_panel( + KV_DTYPE, FP8_PAIR_LOAD, E4M3_SHARED_LUT>( sV, v_cache, smem.page_ids, valid_kv_tile_rows, pv_panel_d_stride_uint4, kv_smem_stride_uint4, tile_page_offset, kv_tile_start, block_size, kv_head_idx, v_block_stride, - v_token_stride, v_head_stride, panel_offset); + v_token_stride, v_head_stride, panel_offset, threadIdx.x, e4m3_lut); for (int idx = tid + valid_kv_tile_rows * pv_panel_d_stride_uint4; idx < kXQATCBlockN * pv_panel_d_stride_uint4; idx += NUM_THREADS) { const int row = idx / pv_panel_d_stride_uint4; @@ -1246,13 +1671,29 @@ __global__ void __launch_bounds__(NUM_THREADS, MIN_BLOCKS_PER_SM) prob = __half2float(sP_row[token]); } const __half* sV_row = sV + token * SmemLayout::kKVStride; + if constexpr (kPVHalf2) { #pragma unroll - for (int d_iter = 0; d_iter < (kPVPanelDim / kWarpSize); ++d_iter) { - const int local_d = lane_id + d_iter * kWarpSize; - const int acc_idx = - panel_idx * (kPVPanelDim / kWarpSize) + d_iter; - out_acc[acc_idx] = - fmaf(prob, __half2float(sV_row[local_d]), out_acc[acc_idx]); + for (int d_iter = 0; d_iter < (kPVPanelDim / (2 * kWarpSize)); + ++d_iter) { + const int local_d2 = lane_id + d_iter * kWarpSize; + const float2 value = __half22float2( + reinterpret_cast(sV_row)[local_d2]); + const int acc_idx = + panel_idx * (kPVPanelDim / kWarpSize) + d_iter * 2; + out_acc[acc_idx] = fmaf(prob, value.x, out_acc[acc_idx]); + out_acc[acc_idx + 1] = + fmaf(prob, value.y, out_acc[acc_idx + 1]); + } + } else { +#pragma unroll + for (int d_iter = 0; d_iter < (kPVPanelDim / kWarpSize); + ++d_iter) { + const int local_d = lane_id + d_iter * kWarpSize; + const int acc_idx = + panel_idx * (kPVPanelDim / kWarpSize) + d_iter; + out_acc[acc_idx] = + fmaf(prob, __half2float(sV_row[local_d]), out_acc[acc_idx]); + } } } } @@ -1281,11 +1722,25 @@ __global__ void __launch_bounds__(NUM_THREADS, MIN_BLOCKS_PER_SM) static_cast(batch_idx) * tmp_out_stride0 + static_cast(head_idx) * tmp_out_stride1 + static_cast(partition_idx) * tmp_out_stride2; - for (int d = thread_in_row; d < D; d += kXQATC256WideThreadsPerRow) { - const float output_scale = KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP16 - ? inv_row_sum - : inv_row_sum * v_scale; - tmp_out_ptr[d] = __float2half(out_acc[d / kWarpSize] * output_scale); + const float output_scale = KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP16 + ? inv_row_sum + : inv_row_sum * v_scale; + if constexpr (kPVHalf2) { +#pragma unroll + for (int acc_idx = 0; acc_idx < kAccumsPerThread; ++acc_idx) { + constexpr int kAccumsPerPanel = kPVPanelDim / kWarpSize; + const int panel_idx = acc_idx / kAccumsPerPanel; + const int panel_acc_idx = acc_idx % kAccumsPerPanel; + const int d_iter = panel_acc_idx / 2; + const int component = panel_acc_idx & 1; + const int d = panel_idx * kPVPanelDim + d_iter * (2 * kWarpSize) + + thread_in_row * 2 + component; + tmp_out_ptr[d] = __float2half(out_acc[acc_idx] * output_scale); + } + } else { + for (int d = thread_in_row; d < D; d += kXQATC256WideThreadsPerRow) { + tmp_out_ptr[d] = __float2half(out_acc[d / kWarpSize] * output_scale); + } } if (thread_in_row == 0) { const int64_t stats_index = @@ -1297,7 +1752,529 @@ __global__ void __launch_bounds__(NUM_THREADS, MIN_BLOCKS_PER_SM) } } -template +// Exact grouped selector-verifier kernel for the admitted TP4 tensor contract. +// +// Independent-row decode scans the same paged KV prefix once for every +// verifier token. This kernel instead assigns one packed GQA group to each of +// eighty context splits. Every CTA handles all six query heads for all eight +// verifier tokens, so K/V is shared across the 48 rows before the exact +// split-softmax states are merged. +// +// The diagnostic two-pass variant intentionally uses: +// 1. QK computes the exact FP32 max/sum for each query row. +// 2. QK is recomputed, then half P x half V is accumulated with Volta WMMA. +// 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; +constexpr int kGroupedVerifyHeads = 6; +constexpr int kGroupedVerifyHeadDim = 256; +constexpr int kGroupedVerifyHeadsPerCta = 6; +constexpr int kGroupedVerifyRows = + kGroupedVerifyMaxQ * kGroupedVerifyHeadsPerCta; +constexpr int kGroupedVerifyBlockN = 32; +// 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; +// 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 +// unchanged eighty splits. The single-query utility variant keeps the old +// 128-token split contract for its stricter FP32-reference error envelope. +constexpr int kGroupedVerifyMinTokensPerSplit = 64; +constexpr int kGroupedVerifySingleQueryMinTokensPerSplit = 128; +constexpr int kGroupedVerifyShortContextMaxTokens = 128; +constexpr int kGroupedVerifyThreads = 512; +constexpr int kGroupedVerifyWarps = kGroupedVerifyThreads / kWarpSize; +constexpr int kGroupedVerifyQKWarps = + (kGroupedVerifyRows / 16) * (kGroupedVerifyBlockN / 16); +constexpr int kGroupedVerifyOutputTiles = + (kGroupedVerifyRows / 16) * (kGroupedVerifyHeadDim / 16); +constexpr int kGroupedVerifyOutputTilesPerWarp = + kGroupedVerifyOutputTiles / kGroupedVerifyWarps; + +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]; + } 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]; +}; + +static_assert(sizeof(GroupedVerifySmem) <= 64 * 1024, + "packed grouped verifier must fit Volta's 64 KiB opt-in budget"); +static_assert(kGroupedVerifyOutputTiles % kGroupedVerifyWarps == 0, + "output tiles must divide evenly across warps"); + +template +__device__ __forceinline__ int grouped_verify_active_splits( + const int total_kv) { + constexpr int kMinTokensPerSplit = + SINGLE_QUERY ? kGroupedVerifySingleQueryMinTokensPerSplit + : kGroupedVerifyMinTokensPerSplit; + const int active_splits = + min(kGroupedVerifySplits, + max(1, (total_kv + kMinTokensPerSplit - 1) / kMinTokensPerSplit)); + if constexpr (SINGLE_QUERY) { + return active_splits; + } + return total_kv <= kGroupedVerifyShortContextMaxTokens ? 1 : active_splits; +} + +__device__ __forceinline__ void grouped_verify_qk( + const __half* __restrict__ shared_q, const __half* __restrict__ shared_k, + float* __restrict__ shared_scores, const float qk_scale) { + const int warp_id = threadIdx.x / kWarpSize; + if (warp_id >= kGroupedVerifyQKWarps) { + return; + } + + const int m_tile = warp_id / (kGroupedVerifyBlockN / 16); + const int n_tile = warp_id % (kGroupedVerifyBlockN / 16); + volta::fragment + q_fragment; + volta::fragment + k_fragment; + volta::fragment score_fragment; + volta::fill_fragment(score_fragment, 0.0f); + +#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); + // 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); + volta::mma_sync(score_fragment, q_fragment, k_fragment, score_fragment); + } +#pragma unroll + for (int i = 0; i < score_fragment.num_elements; ++i) { + 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); +} + +__device__ __forceinline__ void grouped_verify_scale_output_fragment( + volta::fragment& fragment, + const float* __restrict__ row_scale, const int tile_row_start) { + const int lane = threadIdx.x & 31; + const int row = (lane & 1) + ((lane >> 2) & 1) * 8 + ((lane >> 4) & 1) * 4; + const float first_scale = row_scale[tile_row_start + row]; + const float second_scale = row_scale[tile_row_start + row + 2]; + fragment.x[0] *= first_scale; + fragment.x[1] *= first_scale; + fragment.x[2] *= second_scale; + fragment.x[3] *= second_scale; + fragment.x[4] *= first_scale; + fragment.x[5] *= first_scale; + fragment.x[6] *= second_scale; + fragment.x[7] *= second_scale; +} + +template +__global__ +__launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_e5m2_partial_kernel( + const __half* __restrict__ q, const void* __restrict__ k_cache, + const void* __restrict__ v_cache, const int* __restrict__ block_table, + const int* __restrict__ seq_lens, __half* __restrict__ partial_out, + float* __restrict__ partial_lse, const int query_len, + const int max_num_blocks, const int page_block_size, + const int64_t k_block_stride, const int64_t k_token_stride, + 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) { + 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) { + return; + } + + const int total_kv = seq_lens[0]; + if (total_kv <= 0) { + return; + } + const int active_splits = + grouped_verify_active_splits(total_kv); + if (split_id >= active_splits) { + return; + } + const int total_tiles = + (total_kv + kGroupedVerifyBlockN - 1) / kGroupedVerifyBlockN; + const int base_tiles = total_tiles / active_splits; + const int extra_tiles = total_tiles - base_tiles * active_splits; + const int split_tile_start = + split_id * base_tiles + min(split_id, extra_tiles); + const int split_tiles = base_tiles + (split_id < extra_tiles ? 1 : 0); + const int split_start = split_tile_start * kGroupedVerifyBlockN; + 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 tid = threadIdx.x; + const int warp_id = tid / kWarpSize; + const int lane_id = tid % kWarpSize; + + extern __shared__ char grouped_verify_smem_raw[]; + GroupedVerifySmem& smem = + *reinterpret_cast(grouped_verify_smem_raw); + __half* shared_q = smem.storage.compute.q; + __half* shared_kv = smem.storage.compute.kv; + float* shared_scores = smem.storage.compute.scores; + __half* shared_probs = smem.storage.compute.probs; + + constexpr int kVecsPerRow = kGroupedVerifyHeadDim / 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 head_idx = head_start + local_head; + if (token_idx < query_len && head_idx < kGroupedVerifyHeads) { + shared_q_vec[idx] = __ldg( + q_vec + (token_idx * kGroupedVerifyHeads + head_idx) * kVecsPerRow + + vec_col); + } else { + shared_q_vec[idx] = make_uint4(0, 0, 0, 0); + } + } + if (tid < kGroupedVerifyRows) { + smem.row_max[tid] = kXQANegInf; + smem.row_sum[tid] = 0.0f; + smem.row_scale[tid] = 1.0f; + } + __syncthreads(); + + constexpr int kPanelStrideVec = kGroupedVerifyHeadDim / 8; + constexpr int kSharedStrideVec = kGroupedVerifyHeadDim / 8; + + volta::fragment + output_fragments[kGroupedVerifyOutputTilesPerWarp]; +#pragma unroll + for (int fragment_idx = 0; fragment_idx < kGroupedVerifyOutputTilesPerWarp; + ++fragment_idx) { + volta::fill_fragment(output_fragments[fragment_idx], 0.0f); + } + + // The conservative baseline computes exact per-split FP32 max/sum in a + // separate pass. The candidate below fuses this state update with P x V. + if constexpr (TWO_PASS) { + 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); + for (int idx = tid + valid_k_rows * kSharedStrideVec; + idx < kGroupedVerifyBlockN * kSharedStrideVec; + idx += kGroupedVerifyThreads) { + reinterpret_cast(shared_kv)[idx] = make_uint4(0, 0, 0, 0); + } + __syncthreads(); + + grouped_verify_qk(shared_q, shared_kv, shared_scores, qk_scale); + __syncthreads(); + +#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 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] + : kXQANegInf; + const float tile_max_lane = warp_reduce_max(score); + const float tile_max = __shfl_sync(0xffffffffu, tile_max_lane, 0); + const float probability = + visible ? __expf(fmaxf(score - tile_max, -80.0f)) : 0.0f; + const float tile_sum_lane = warp_reduce_sum(probability); + const float tile_sum = __shfl_sync(0xffffffffu, tile_sum_lane, 0); + if (lane_id == 0 && tile_sum > 0.0f) { + const float old_max = smem.row_max[row]; + const float old_sum = smem.row_sum[row]; + const float new_max = fmaxf(old_max, tile_max); + smem.row_sum[row] = + old_sum * __expf(fmaxf(old_max - new_max, -80.0f)) + + tile_sum * __expf(fmaxf(tile_max - new_max, -80.0f)); + smem.row_max[row] = new_max; + } + } + __syncthreads(); + } + } + + // Recompute QK for the conservative path, or consume it once while updating + // the online state for the fused path. Both form half P and use identical + // Volta WMMA P @ V arithmetic. + 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); + for (int idx = tid + valid_k_rows * kSharedStrideVec; + idx < kGroupedVerifyBlockN * kSharedStrideVec; + idx += kGroupedVerifyThreads) { + reinterpret_cast(shared_kv)[idx] = make_uint4(0, 0, 0, 0); + } + __syncthreads(); + + grouped_verify_qk(shared_q, shared_kv, shared_scores, qk_scale); + __syncthreads(); + + if constexpr (TWO_PASS) { + for (int idx = tid; idx < kGroupedVerifyRows * kGroupedVerifyBlockN; + 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 head_idx = head_start + local_head; + const int kv_idx = tile_start + col; + const bool visible = + token_idx < query_len && head_idx < kGroupedVerifyHeads && + 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); + } + __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 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] + : kXQANegInf; + const float tile_max_lane = warp_reduce_max(score); + const float tile_max = __shfl_sync(0xffffffffu, tile_max_lane, 0); + const float old_max = smem.row_max[row]; + const float new_max = fmaxf(old_max, tile_max); + const float probability = + visible ? __expf(fmaxf(score - new_max, -80.0f)) : 0.0f; + const float tile_sum_lane = warp_reduce_sum(probability); + 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] = + __float2half_rn(probability); + if (lane_id == 0) { + if (tile_sum > 0.0f) { + smem.row_sum[row] = smem.row_sum[row] * exp_diff + tile_sum; + smem.row_max[row] = new_max; + } + smem.row_scale[row] = exp_diff; + } + } + __syncthreads(); +#pragma unroll + for (int fragment_idx = 0; + fragment_idx < kGroupedVerifyOutputTilesPerWarp; ++fragment_idx) { + const int output_tile = warp_id + fragment_idx * kGroupedVerifyWarps; + const int m_tile = output_tile / (kGroupedVerifyHeadDim / 16); + grouped_verify_scale_output_fragment(output_fragments[fragment_idx], + smem.row_scale, m_tile * 16); + } + } + + 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); + for (int idx = tid + valid_k_rows * kSharedStrideVec; + idx < kGroupedVerifyBlockN * kSharedStrideVec; + idx += kGroupedVerifyThreads) { + reinterpret_cast(shared_kv)[idx] = make_uint4(0, 0, 0, 0); + } + __syncthreads(); + +#pragma unroll + for (int fragment_idx = 0; fragment_idx < kGroupedVerifyOutputTilesPerWarp; + ++fragment_idx) { + const int output_tile = warp_id + fragment_idx * kGroupedVerifyWarps; + const int m_tile = output_tile / (kGroupedVerifyHeadDim / 16); + const int d_tile = output_tile % (kGroupedVerifyHeadDim / 16); + volta::fragment + probability_fragment; + volta::fragment + value_fragment; +#pragma unroll + 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); + volta::load_matrix_sync( + value_fragment, + shared_kv + k_offset * kGroupedVerifyHeadDim + d_tile * 16, + kGroupedVerifyHeadDim); + volta::mma_sync(output_fragments[fragment_idx], probability_fragment, + value_fragment, output_fragments[fragment_idx]); + } + } + __syncthreads(); + } + + // The compute buffers are dead. Reuse their storage for the dense FP32 + // partial output, then normalize and write only real query/head rows. + __syncthreads(); + float* shared_output = smem.storage.output; +#pragma unroll + for (int fragment_idx = 0; fragment_idx < kGroupedVerifyOutputTilesPerWarp; + ++fragment_idx) { + const int output_tile = warp_id + fragment_idx * kGroupedVerifyWarps; + const int m_tile = output_tile / (kGroupedVerifyHeadDim / 16); + const int d_tile = output_tile % (kGroupedVerifyHeadDim / 16); + volta::store_matrix_sync( + shared_output + m_tile * 16 * kGroupedVerifyHeadDim + d_tile * 16, + output_fragments[fragment_idx], kGroupedVerifyHeadDim, + volta::mem_row_major); + } + __syncthreads(); + + for (int idx = tid; idx < kGroupedVerifyRows * kGroupedVerifyHeadDim; + 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 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) * + kGroupedVerifyHeads + + head_idx) * + kGroupedVerifyHeadDim + + d); + partial_out[output_idx] = __float2half_rn(shared_output[idx] * scale); + } + } + if (tid < kGroupedVerifyRows) { + const int token_idx = tid / kGroupedVerifyHeadsPerCta; + const int local_head = tid % kGroupedVerifyHeadsPerCta; + 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) * + kGroupedVerifyHeads + + head_idx; + partial_lse[lse_idx] = + sum > 0.0f ? smem.row_max[tid] + logf(sum) : kXQANegInf; + } + } +} + +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) { + 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]; + __shared__ float final_max; + __shared__ float final_inv_sum; + + if (threadIdx.x < kGroupedVerifySplits) { + const int64_t lse_idx = + (static_cast(threadIdx.x) * kGroupedVerifyMaxQ + token_idx) * + kGroupedVerifyHeads + + head_idx; + split_lse[threadIdx.x] = + threadIdx.x < active_splits ? partial_lse[lse_idx] : kXQANegInf; + } + __syncthreads(); + + if (threadIdx.x == 0) { + float max_lse = kXQANegInf; + for (int split = 0; split < active_splits; ++split) { + max_lse = fmaxf(max_lse, split_lse[split]); + } + float sum = 0.0f; + for (int split = 0; split < active_splits; ++split) { + if (split_lse[split] > -1.0e20f) { + sum += __expf(fmaxf(split_lse[split] - max_lse, -80.0f)); + } + } + final_max = max_lse; + final_inv_sum = sum > 0.0f ? 1.0f / sum : 0.0f; + } + __syncthreads(); + + for (int d = threadIdx.x; d < kGroupedVerifyHeadDim; + d += kGroupedVerifyThreads) { + float accumulator = 0.0f; + for (int split = 0; split < active_splits; ++split) { + if (split_lse[split] > -1.0e20f) { + 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) * + kGroupedVerifyHeads + + head_idx) * + kGroupedVerifyHeadDim + + d); + accumulator = + fmaf(weight, __half2float(partial_out[partial_idx]), accumulator); + } + } + out[(token_idx * kGroupedVerifyHeads + head_idx) * kGroupedVerifyHeadDim + + d] = __float2half_rn(accumulator); + } +} + +template __global__ void flash_attention_decode_reduce_kernel( const __half* __restrict__ tmp_out, const float* __restrict__ max_logits, const float* __restrict__ exp_sums, const int* __restrict__ seq_lens, @@ -1306,7 +2283,9 @@ __global__ void flash_attention_decode_reduce_kernel( 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 int64_t out_stride1, const int route_partition_size_begin, + const int route_seq_len_begin, const int route_seq_len_end, + const int route_seq_len_final) { const int batch_idx = blockIdx.x; const int head_idx = blockIdx.y; @@ -1315,8 +2294,25 @@ __global__ void flash_attention_decode_reduce_kernel( } const int seq_len = seq_lens[batch_idx]; + if (!xqa_seq_len_route_active(seq_len, route_seq_len_begin, + route_seq_len_end, + route_seq_len_final)) { + return; + } + int partition_size; + if constexpr (PARTITION_SIZE == -1) { + static_assert(SEQ_LEN_ROUTE == kXQARouteWaveLongSeqLens, + "Runtime wave partitions require the wave-long route"); + partition_size = seq_len < route_seq_len_end + ? 512 + : (seq_len < route_seq_len_final ? 896 : 1664); + } else if constexpr (PARTITION_SIZE == 0) { + partition_size = seq_len < route_partition_size_begin ? 64 : 256; + } else { + partition_size = PARTITION_SIZE; + } const int num_partitions = - min(max_num_partitions, (seq_len + PARTITION_SIZE - 1) / PARTITION_SIZE); + min(max_num_partitions, (seq_len + partition_size - 1) / partition_size); (void)active_num_partitions; if (seq_len <= 0 || num_partitions <= 0) { @@ -2045,7 +3041,8 @@ void launch_flash_attention_decode_paged( const int window_size_left, const int window_size_right, cudaStream_t stream, const int route_seq_len_begin = 0, const int route_seq_len_end = 0, const int route_seq_len_final = 0, - const bool launch_reduce = true) { + const bool launch_reduce = true, const int* anchor_lens = nullptr, + const int anchored_window = 0) { const int batch_size = q.size(0); const int num_heads_q = q.size(1); const int num_heads_kv = k_cache.size(2); @@ -2059,22 +3056,41 @@ void launch_flash_attention_decode_paged( const size_t reduce_shared_mem = static_cast(2 * max_num_partitions) * sizeof(float); - flash_attention_decode_partition_kernel - <<>>( - reinterpret_cast(q.data_ptr()), - k_cache.data_ptr(), v_cache.data_ptr(), - reinterpret_cast<__half*>(tmp_out.data_ptr()), - max_logits.data_ptr(), exp_sums.data_ptr(), - block_table.data_ptr(), seq_lens.data_ptr(), - active_num_partitions.data_ptr(), batch_size, max_num_blocks, - max_num_partitions, num_heads_q, num_heads_kv, block_size, - q.stride(0), q.stride(1), tmp_out.stride(0), tmp_out.stride(1), - tmp_out.stride(2), max_logits.stride(0), max_logits.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), - softmax_scale, k_scale, v_scale, window_size_left, window_size_right, - route_seq_len_begin, route_seq_len_end, route_seq_len_final); + const bool use_anchored = anchor_lens != nullptr && anchored_window > 0; + // Second kernel version: the anchored decode-window mask is a separate + // template instantiation, generated only for the fp16-KV configuration; + // the non-anchored instantiations stay untouched. + const auto launch_partition = [&](auto anchored_tag) { + constexpr bool kAnchored = decltype(anchored_tag)::value; + flash_attention_decode_partition_kernel + <<>>( + reinterpret_cast(q.data_ptr()), + k_cache.data_ptr(), v_cache.data_ptr(), + reinterpret_cast<__half*>(tmp_out.data_ptr()), + max_logits.data_ptr(), exp_sums.data_ptr(), + block_table.data_ptr(), seq_lens.data_ptr(), + active_num_partitions.data_ptr(), batch_size, max_num_blocks, + max_num_partitions, num_heads_q, num_heads_kv, block_size, + q.stride(0), q.stride(1), tmp_out.stride(0), tmp_out.stride(1), + tmp_out.stride(2), max_logits.stride(0), max_logits.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), + softmax_scale, k_scale, v_scale, window_size_left, + window_size_right, route_seq_len_begin, route_seq_len_end, + route_seq_len_final, anchor_lens, anchored_window); + }; + if constexpr (KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP16) { + if (use_anchored) { + launch_partition(std::true_type{}); + } else { + launch_partition(std::false_type{}); + } + } else { + TORCH_CHECK(!use_anchored, + "anchored decode window requires an fp16 KV cache"); + launch_partition(std::false_type{}); + } if (!launch_reduce) { return; @@ -2088,7 +3104,7 @@ void launch_flash_attention_decode_paged( 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)); + out.stride(0), out.stride(1), 0, 0, 0, 0); } template @@ -2150,7 +3166,8 @@ template + bool PARTITION_PAGE_IDS = false, bool FP8_PAIR_LOAD = false, + int KV_DTYPE_OVERRIDE = -1, bool E4M3_SHARED_LUT = false> void launch_flash_attention_decode_paged_xqa_tc_256_wide( const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, at::Tensor& out, const at::Tensor& block_table, const at::Tensor& seq_lens, @@ -2161,6 +3178,9 @@ void launch_flash_attention_decode_paged_xqa_tc_256_wide( cudaStream_t stream, const int route_seq_len_begin = 0, const int route_seq_len_end = 0, const int route_seq_len_final = 0, const bool launch_reduce = true) { + static_assert(!E4M3_SHARED_LUT || + KV_DTYPE_OVERRIDE == flash_v100::KV_CACHE_DTYPE_FP8_E4M3, + "The shared conversion LUT requires an E4M3 specialization"); const int batch_size = q.size(0); const int num_heads_q = q.size(1); const int num_heads_kv = k_cache.size(2); @@ -2170,7 +3190,8 @@ void launch_flash_attention_decode_paged_xqa_tc_256_wide( QK_SW_PIPELINE, XQATCQKPipelineSmem256WideLayout, XQATCSmem256WideLayout>; - const size_t shared_mem = sizeof(SmemLayout); + const size_t shared_mem = + sizeof(SmemLayout) + (E4M3_SHARED_LUT ? 256 * sizeof(uint16_t) : 0); #define LAUNCH_XQA_PARTITION(KV_DTYPE) \ do { \ auto partition_kernel = \ @@ -2178,7 +3199,7 @@ void launch_flash_attention_decode_paged_xqa_tc_256_wide( PARTITION_SIZE, GROUP_SIZE, PADDED_SMEM, NUM_THREADS, \ MIN_BLOCKS_PER_SM, BLOCK_SIZE, CONTIGUOUS_HKV1_LAYOUT, \ ALIGNED_PADDED_SMEM, KV_DTYPE, SEQ_LEN_ROUTE, QK_SW_PIPELINE, \ - PARTITION_PAGE_IDS, E5M2_PAIR_LOAD>; \ + PARTITION_PAGE_IDS, FP8_PAIR_LOAD, E4M3_SHARED_LUT>; \ cudaFuncSetAttribute(partition_kernel, \ cudaFuncAttributeMaxDynamicSharedMemorySize, \ shared_mem); \ @@ -2191,7 +3212,7 @@ void launch_flash_attention_decode_paged_xqa_tc_256_wide( PARTITION_SIZE, GROUP_SIZE, PADDED_SMEM, NUM_THREADS, \ MIN_BLOCKS_PER_SM, BLOCK_SIZE, CONTIGUOUS_HKV1_LAYOUT, \ ALIGNED_PADDED_SMEM, KV_DTYPE, SEQ_LEN_ROUTE, QK_SW_PIPELINE, \ - PARTITION_PAGE_IDS, E5M2_PAIR_LOAD> \ + PARTITION_PAGE_IDS, FP8_PAIR_LOAD, E4M3_SHARED_LUT> \ <<>>( \ reinterpret_cast(q.data_ptr()), \ k_cache.data_ptr(), v_cache.data_ptr(), \ @@ -2208,9 +3229,16 @@ void launch_flash_attention_decode_paged_xqa_tc_256_wide( route_seq_len_end, route_seq_len_final); \ } while (0) - if constexpr (E5M2_PAIR_LOAD) { + if constexpr (KV_DTYPE_OVERRIDE == flash_v100::KV_CACHE_DTYPE_FP8_E4M3) { + static_assert(!(FP8_PAIR_LOAD && E4M3_SHARED_LUT), + "Paired E4M3 conversion and the shared LUT are exclusive"); + TORCH_CHECK(k_cache.scalar_type() == at::kByte, + "E4M3 XQA requires uint8 KV cache"); + LAUNCH_XQA_PARTITION(flash_v100::KV_CACHE_DTYPE_FP8_E4M3); + } else if constexpr (FP8_PAIR_LOAD) { + static_assert(!E4M3_SHARED_LUT, "The E4M3 conversion LUT requires E4M3 KV"); TORCH_CHECK(k_cache.scalar_type() == at::kByte, - "Paired XQA loads require FP8 E5M2 KV cache"); + "Paired XQA loads require uint8 FP8 KV cache"); LAUNCH_XQA_PARTITION(flash_v100::KV_CACHE_DTYPE_FP8_E5M2); } else if (k_cache.scalar_type() == at::kByte) { LAUNCH_XQA_PARTITION(flash_v100::KV_CACHE_DTYPE_FP8_E5M2); @@ -2224,16 +3252,22 @@ void launch_flash_attention_decode_paged_xqa_tc_256_wide( return; } + if constexpr (PARTITION_SIZE == -1) { + TORCH_CHECK(!use_split_reduce, + "Runtime wave partitions do not support split reduction"); + } if (use_split_reduce) { - launch_flash_attention_decode_xqa_split_reduce( - out, seq_lens, tmp_out, max_logits, exp_sums, launch_num_partitions, - split_reduce_dim_tile, stream); + if constexpr (PARTITION_SIZE != -1) { + launch_flash_attention_decode_xqa_split_reduce( + out, seq_lens, tmp_out, max_logits, exp_sums, launch_num_partitions, + split_reduce_dim_tile, stream); + } } else { const dim3 reduce_grid(batch_size, num_heads_q, 1); const dim3 block(kThreadsPerBlock); const size_t reduce_shared_mem = static_cast(2 * launch_num_partitions) * sizeof(float); - flash_attention_decode_reduce_kernel<256, PARTITION_SIZE> + flash_attention_decode_reduce_kernel<256, PARTITION_SIZE, SEQ_LEN_ROUTE> <<>>( reinterpret_cast(tmp_out.data_ptr()), max_logits.data_ptr(), exp_sums.data_ptr(), @@ -2241,11 +3275,34 @@ void launch_flash_attention_decode_paged_xqa_tc_256_wide( reinterpret_cast<__half*>(out.data_ptr()), batch_size, launch_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)); + max_logits.stride(1), out.stride(0), out.stride(1), + route_seq_len_end > 0 ? route_seq_len_end : route_seq_len_begin, + route_seq_len_begin, route_seq_len_end, route_seq_len_final); } C10_CUDA_KERNEL_LAUNCH_CHECK(); } +template +void launch_flash_attention_decode_paged_xqa_e4m3_g6_page1568( + const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_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& active_num_partitions, const float softmax_scale, + const float k_scale, const float v_scale, const int launch_num_partitions, + cudaStream_t stream, const int route_seq_len_begin = 0, + const int route_seq_len_end = 0, const int route_seq_len_final = 0, + const bool launch_reduce = true) { + launch_flash_attention_decode_paged_xqa_tc_256_wide< + PARTITION_SIZE, 6, true, kXQATCG6DualCtaThreads, 2, 1568, + FIXED_INTERLEAVED_HKV1_LAYOUT, false, SEQ_LEN_ROUTE, false, true, false, + flash_v100::KV_CACHE_DTYPE_FP8_E4M3, true>( + q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, max_logits, + exp_sums, active_num_partitions, softmax_scale, k_scale, v_scale, + launch_num_partitions, false, 8, stream, route_seq_len_begin, + route_seq_len_end, route_seq_len_final, launch_reduce); +} + void launch_flash_attention_decode_paged_xqa_tc_256_staged( const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, at::Tensor& out, const at::Tensor& block_table, const at::Tensor& seq_lens, @@ -2346,7 +3403,7 @@ void launch_flash_attention_decode_paged_xqa_tc_256_staged( reinterpret_cast<__half*>(out.data_ptr()), batch_size, launch_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)); + max_logits.stride(1), out.stride(0), out.stride(1), 0, 0, 0, 0); } C10_CUDA_KERNEL_LAUNCH_CHECK(); } @@ -2380,6 +3437,158 @@ void launch_flash_attention_decode_qk_scores( } // namespace +at::Tensor flash_attention_grouped_verify_paged( + const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, + std::optional& out_, const at::Tensor& block_table, + const at::Tensor& seq_lens, at::Tensor& partial_out, + at::Tensor& partial_lse, const float softmax_scale, + const std::string& kv_cache_dtype, const float k_scale, const float v_scale, + const bool one_pass) { + TORCH_CHECK(q.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda(), + "grouped verify q/K/V must be CUDA tensors"); + TORCH_CHECK(block_table.is_cuda() && seq_lens.is_cuda(), + "grouped verify metadata must be CUDA tensors"); + TORCH_CHECK(partial_out.is_cuda() && partial_lse.is_cuda(), + "grouped verify workspaces must be CUDA tensors"); + TORCH_CHECK(q.dtype() == torch::kFloat16, "grouped verify q must be fp16"); + TORCH_CHECK(kv_cache_dtype == "fp8_e5m2", + "grouped verify prototype supports fp8_e5m2 KV only"); + TORCH_CHECK( + k_cache.dtype() == torch::kUInt8 && v_cache.dtype() == torch::kUInt8, + "grouped verify E5M2 cache must use uint8 storage"); + TORCH_CHECK( + block_table.dtype() == torch::kInt32 && seq_lens.dtype() == torch::kInt32, + "grouped verify block_table/seq_lens must be int32"); + TORCH_CHECK(partial_out.dtype() == torch::kFloat16 && + 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(1) == kGroupedVerifyHeads && + q.size(2) == kGroupedVerifyHeadDim, + "grouped verify q must have shape [1..8, 6, 256]"); + 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 && + k_cache.size(3) == kGroupedVerifyHeadDim, + "grouped verify KV must have shape [blocks, page, 1, 256]"); + TORCH_CHECK(block_table.dim() == 2 && block_table.size(0) == 1, + "grouped verify block_table must have shape [1, blocks]"); + TORCH_CHECK(seq_lens.dim() == 1 && seq_lens.size(0) >= 1, + "grouped verify seq_lens must cover one sequence"); + TORCH_CHECK(q.is_contiguous(), + "grouped verify q must be contiguous [M, H, D]"); + TORCH_CHECK(block_table.is_contiguous() && seq_lens.is_contiguous(), + "grouped verify metadata must be contiguous"); + TORCH_CHECK(k_cache.stride(-1) == 1 && v_cache.stride(-1) == 1, + "grouped verify KV head dimension must be contiguous"); + 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, + 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]"); + TORCH_CHECK(k_scale > 0.0f && v_scale > 0.0f, + "grouped verify E5M2 K/V scales must be positive"); + + at::Tensor out = out_.has_value() ? out_.value() : torch::empty_like(q); + TORCH_CHECK(out.is_cuda() && out.device() == q.device() && + out.dtype() == torch::kFloat16 && out.sizes() == q.sizes() && + out.is_contiguous(), + "grouped verify out must be contiguous fp16 and q-shaped"); + TORCH_CHECK( + q.device() == k_cache.device() && q.device() == v_cache.device() && + q.device() == block_table.device() && + q.device() == seq_lens.device() && + q.device() == partial_out.device() && + q.device() == partial_lse.device() && q.device() == out.device(), + "all grouped verify tensors must be on the same device"); + + c10::cuda::CUDAGuard device_guard(q.device()); + const auto* properties = at::cuda::getCurrentDeviceProperties(); + TORCH_CHECK(properties->major == 7 && properties->minor == 0, + "grouped verify prototype supports SM70 only"); + cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + + const dim3 partial_grid(1, kGroupedVerifySplits, 1); + const size_t partial_shared_mem = sizeof(GroupedVerifySmem); +#define LAUNCH_GROUPED_VERIFY_PARTIAL(TWO_PASS, PAGE_SIZE, SINGLE_QUERY) \ + do { \ + auto partial_kernel = \ + (void*)flash_attention_grouped_verify_e5m2_partial_kernel< \ + TWO_PASS, PAGE_SIZE, SINGLE_QUERY>; \ + const cudaError_t smem_status = cudaFuncSetAttribute( \ + partial_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, \ + partial_shared_mem); \ + TORCH_CHECK(smem_status == cudaSuccess, \ + "Failed to set packed grouped-verifier shared memory: ", \ + cudaGetErrorString(smem_status)); \ + const cudaError_t carveout_status = cudaFuncSetAttribute( \ + partial_kernel, cudaFuncAttributePreferredSharedMemoryCarveout, 100); \ + TORCH_CHECK(carveout_status == cudaSuccess, \ + "Failed to set packed grouped-verifier shared carveout: ", \ + cudaGetErrorString(carveout_status)); \ + flash_attention_grouped_verify_e5m2_partial_kernel \ + <<>>( \ + reinterpret_cast(q.data_ptr()), k_cache.data_ptr(), \ + v_cache.data_ptr(), block_table.data_ptr(), \ + seq_lens.data_ptr(), \ + reinterpret_cast<__half*>(partial_out.data_ptr()), \ + partial_lse.data_ptr(), static_cast(q.size(0)), \ + static_cast(block_table.size(1)), \ + static_cast(k_cache.size(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), softmax_scale * k_scale, \ + 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); \ + } \ + } while (0) + + const bool single_query = q.size(0) == 1; + if (one_pass && single_query) { + DISPATCH_GROUPED_VERIFY_PARTIAL(false, true); + } else if (one_pass) { + DISPATCH_GROUPED_VERIFY_PARTIAL(false, false); + } else if (single_query) { + DISPATCH_GROUPED_VERIFY_PARTIAL(true, true); + } else { + DISPATCH_GROUPED_VERIFY_PARTIAL(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()), \ + static_cast(q.size(0))) + if (single_query) { + LAUNCH_GROUPED_VERIFY_COMBINE(true); + } else { + LAUNCH_GROUPED_VERIFY_COMBINE(false); + } +#undef LAUNCH_GROUPED_VERIFY_COMBINE + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} + at::Tensor flash_attention_decode_paged( const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, std::optional& out_, const at::Tensor& block_table, @@ -2388,7 +3597,8 @@ at::Tensor flash_attention_decode_paged( const float softmax_scale, const int partition_size, const int launch_num_partitions, const std::string& kv_cache_dtype, const float k_scale, const float v_scale, const int window_size_left, - const int window_size_right) { + const int window_size_right, const std::optional& anchor_lens, + const int64_t anchored_window) { TORCH_CHECK(q.is_cuda(), "q must be on CUDA"); TORCH_CHECK(k_cache.is_cuda() && v_cache.is_cuda(), "k/v cache must be on CUDA"); @@ -2474,6 +3684,30 @@ at::Tensor flash_attention_decode_paged( TORCH_CHECK(window_size_left >= -1 && window_size_right >= -1, "window sizes must be >= -1"); + TORCH_CHECK(anchored_window >= 0 && anchored_window <= INT_MAX, + "anchored_window must fit in a non-negative int32"); + TORCH_CHECK((anchored_window > 0) == anchor_lens.has_value(), + "anchor_lens and a positive anchored_window must be provided " + "together"); + const bool use_anchored = anchor_lens.has_value(); + const int* anchor_lens_ptr = nullptr; + if (use_anchored) { + const at::Tensor& anchors = anchor_lens.value(); + TORCH_CHECK(anchors.is_cuda(), "anchor_lens must be on CUDA"); + TORCH_CHECK(anchors.device() == q.device(), + "anchor_lens must be on the q device"); + TORCH_CHECK(anchors.dtype() == torch::kInt32, "anchor_lens must be int32"); + TORCH_CHECK(anchors.dim() == 1 && anchors.size(0) == batch_size, + "anchor_lens must have shape [B]"); + TORCH_CHECK(anchors.is_contiguous(), "anchor_lens must be contiguous"); + TORCH_CHECK(kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP16, + "anchored decode window requires an fp16 KV cache"); + TORCH_CHECK(window_size_left == -1 && window_size_right == -1, + "anchored decode window cannot be combined with " + "sliding-window attention"); + anchor_lens_ptr = anchors.data_ptr(); + } + at::Tensor out = out_.has_value() ? out_.value() : torch::empty_like(q); TORCH_CHECK(out.is_cuda(), "out must be on CUDA"); TORCH_CHECK(out.dtype() == torch::kFloat16, "out must be fp16"); @@ -2483,11 +3717,12 @@ at::Tensor flash_attention_decode_paged( auto stream = at::cuda::getCurrentCUDAStream().stream(); c10::cuda::CUDAGuard device_guard(q.device()); -#define LAUNCH_TYPED(HDIM, PARTITION, KV_DTYPE_CODE) \ - launch_flash_attention_decode_paged( \ - q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, max_logits, \ - exp_sums, active_num_partitions, softmax_scale, launch_num_partitions, \ - k_scale, v_scale, window_size_left, window_size_right, stream) +#define LAUNCH_TYPED(HDIM, PARTITION, KV_DTYPE_CODE) \ + launch_flash_attention_decode_paged( \ + q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, max_logits, \ + exp_sums, active_num_partitions, softmax_scale, launch_num_partitions, \ + k_scale, v_scale, window_size_left, window_size_right, stream, 0, 0, 0, \ + true, anchor_lens_ptr, static_cast(anchored_window)) #define LAUNCH_BY_KV_DTYPE(HDIM, PARTITION) \ do { \ @@ -2563,7 +3798,7 @@ at::Tensor flash_attention_decode_paged_xqa( const float softmax_scale, const int partition_size, const int launch_num_partitions, const std::string& kv_cache_dtype, const float k_scale, const float v_scale, const int window_size_left, - const int window_size_right) { + const int window_size_right, const int batch_context_max_seq_len) { TORCH_CHECK(q.is_cuda(), "q must be on CUDA"); TORCH_CHECK(k_cache.is_cuda() && v_cache.is_cuda(), "k_cache and v_cache must be on CUDA"); @@ -2576,8 +3811,10 @@ at::Tensor flash_attention_decode_paged_xqa( TORCH_CHECK(q.dtype() == torch::kFloat16, "q must be fp16"); const int kv_dtype_code = kv_cache_dtype_code_from_string(kv_cache_dtype); TORCH_CHECK(kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP16 || + kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP8_E4M3 || kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP8_E5M2, - "XQA decode supports fp16 and fp8_e5m2 KV cache only"); + "XQA decode supports fp16, fp8_e4m3, and fp8_e5m2 KV cache " + "only"); if (kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP16) { TORCH_CHECK(k_cache.dtype() == torch::kFloat16 && v_cache.dtype() == torch::kFloat16, @@ -2585,7 +3822,7 @@ at::Tensor flash_attention_decode_paged_xqa( } else { TORCH_CHECK( k_cache.dtype() == torch::kUInt8 && v_cache.dtype() == torch::kUInt8, - "fp8_e5m2 XQA requires uint8 K/V tensors"); + "FP8 XQA requires uint8 K/V tensors"); } TORCH_CHECK(k_scale > 0.f && v_scale > 0.f, "XQA K/V scales must be positive"); @@ -2622,9 +3859,12 @@ at::Tensor flash_attention_decode_paged_xqa( const int q_per_kv = num_heads_q / num_heads_kv; TORCH_CHECK(q_per_kv == 4 || q_per_kv == 6 || q_per_kv == 8, "XQA decode supports q_per_kv in {4, 6, 8}, got ", q_per_kv); - TORCH_CHECK( - partition_size == 256 || partition_size == 512 || partition_size == 1024, - "Unsupported XQA decode partition_size: ", partition_size); + TORCH_CHECK(partition_size == 64 || partition_size == 128 || + partition_size == 256 || partition_size == 512 || + partition_size == 1024 || + ((partition_size == 896 || partition_size == 1664) && + kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP8_E4M3), + "Unsupported XQA decode partition_size: ", partition_size); TORCH_CHECK(launch_num_partitions > 0, "launch_num_partitions must be positive"); TORCH_CHECK(tmp_out.dtype() == torch::kFloat16, @@ -2648,6 +3888,345 @@ at::Tensor flash_attention_decode_paged_xqa( TORCH_CHECK(out.sizes() == q.sizes(), "out must have same shape as q"); TORCH_CHECK(out.stride(-1) == 1, "out last dim must be contiguous"); auto stream = at::cuda::getCurrentCUDAStream().stream(); + if (kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP8_E4M3) { + const bool e4m3_batch_allowed = + q.size(0) > 1 && q.size(0) <= 16 && xqa_e4m3_batch_enabled(); + const bool use_large_partition = + q.size(0) == 1 && + (partition_size == 512 || partition_size == 896 || + partition_size == 1024 || partition_size == 1664) && + k_cache.size(1) == 1568 && k_cache.size(2) == 1; + TORCH_CHECK((q.size(0) == 1 || e4m3_batch_allowed) && q_per_kv == 6 && + (partition_size == 64 || partition_size == 128 || + partition_size == 256 || use_large_partition), + "E4M3 XQA supports B=1, or B=2..16 when " + "VLLM_FLASH_V100_E4M3_BATCH_XQA=1; q_per_kv=6 and D=256 are " + "required. Page-1568/Hkv=1 B1 additionally supports partition " + "sizes 512, 896, 1024, and 1664"); + const bool e4m3_batch_optimized = + q.size(0) > 1 && xqa_e4m3_batch_optimized_enabled() && + k_cache.size(1) >= 256 && k_cache.size(1) % 16 == 0 && + k_cache.size(2) == 1 && k_cache.scalar_type() == at::kByte; + const XQABatchContextRoute e4m3_batch_route = + e4m3_batch_optimized && batch_context_max_seq_len > 0 && q.size(0) >= 4 + ? select_xqa_batch_context_route( + q.size(0), batch_context_max_seq_len, partition_size) + : XQABatchContextRoute::kDisabled; + const bool e4m3_dual_cta = + e4m3_batch_route == XQABatchContextRoute::kDualCta || + e4m3_batch_route == XQABatchContextRoute::kDualCtaSplit; + const bool e4m3_split_reduce = + e4m3_batch_route == XQABatchContextRoute::kDualCtaSplit; + constexpr int64_t kPage800HeadDim256Elements = 800 * 256; + const bool e4m3_page800_contiguous = + xqa_e4m3_page800_fastpath_enabled() && k_cache.size(1) == 800 && + k_cache.size(2) == 1 && k_cache.size(3) == 256 && + k_cache.stride(0) == 2 * kPage800HeadDim256Elements && + k_cache.stride(1) == 256 && k_cache.stride(2) == 256 && + k_cache.stride(3) == 1 && + v_cache.stride(0) == 2 * kPage800HeadDim256Elements && + v_cache.stride(1) == 256 && v_cache.stride(2) == 256 && + v_cache.stride(3) == 1; + trace_xqa_batch_context_route(q.size(0), batch_context_max_seq_len, + partition_size, k_cache.size(1), + e4m3_batch_route); + trace_xqa_e4m3_page800_fastpath(q.size(0), partition_size, + e4m3_dual_cta && e4m3_page800_contiguous); + const int split_reduce_dim_tile = xqa_split_reduce_dim_tile(); + +#define LAUNCH_E4M3_DUAL_CTA(PARTITION, PAGE_SIZE, CONTIGUOUS) \ + launch_flash_attention_decode_paged_xqa_tc_256_wide< \ + PARTITION, 6, true, kXQATCG6DualCtaThreads, 2, PAGE_SIZE, CONTIGUOUS, \ + false, kXQARouteAllSeqLens, false, true, true, \ + flash_v100::KV_CACHE_DTYPE_FP8_E4M3>( \ + q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, max_logits, \ + exp_sums, active_num_partitions, softmax_scale, k_scale, v_scale, \ + launch_num_partitions, e4m3_split_reduce, split_reduce_dim_tile, stream) + +#define LAUNCH_E4M3_BATCH_XQA(PARTITION) \ + do { \ + if (e4m3_dual_cta) { \ + if (e4m3_page800_contiguous) { \ + LAUNCH_E4M3_DUAL_CTA(PARTITION, 800, true); \ + } else { \ + LAUNCH_E4M3_DUAL_CTA(PARTITION, 0, false); \ + } \ + } else if (e4m3_batch_optimized) { \ + launch_flash_attention_decode_paged_xqa_tc_256_wide< \ + PARTITION, 6, true, kXQATC256WideThreads, 1, 0, false, false, \ + kXQARouteAllSeqLens, false, true, true, \ + flash_v100::KV_CACHE_DTYPE_FP8_E4M3>( \ + q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, \ + max_logits, exp_sums, active_num_partitions, softmax_scale, k_scale, \ + v_scale, launch_num_partitions, false, split_reduce_dim_tile, \ + stream); \ + } else { \ + launch_flash_attention_decode_paged_xqa_tc_256_wide< \ + PARTITION, 6, true, kXQATC256WideThreads, 1, 0, false, false, \ + kXQARouteAllSeqLens, false, false, false, \ + flash_v100::KV_CACHE_DTYPE_FP8_E4M3>( \ + q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, \ + max_logits, exp_sums, active_num_partitions, softmax_scale, k_scale, \ + v_scale, launch_num_partitions, false, 8, stream); \ + } \ + } while (0) + + const bool use_p64_p256_auto = + q.size(0) == 1 && partition_size == 64 && k_cache.size(1) == 1568 && + k_cache.size(2) == 1 && !decode_partition_size_overridden() && + xqa_e4m3_g6_p64_p256_auto_enabled(); + constexpr int64_t kPage1568HeadDim256Elements = 1568 * 256; + const bool fixed_interleaved_hkv1_layout = + k_cache.size(3) == 256 && + k_cache.stride(0) == 2 * kPage1568HeadDim256Elements && + k_cache.stride(1) == 256 && k_cache.stride(2) == 256 && + k_cache.stride(3) == 1 && + v_cache.stride(0) == 2 * kPage1568HeadDim256Elements && + v_cache.stride(1) == 256 && v_cache.stride(2) == 256 && + v_cache.stride(3) == 1; + if (use_p64_p256_auto) { + const int p256_begin = xqa_e4m3_g6_p256_begin(); + const int p512_begin = xqa_e4m3_g6_p512_begin(); + const bool use_wave_partitions = xqa_e4m3_g6_wave_partitions_enabled() && + batch_context_max_seq_len >= p512_begin; + if (use_wave_partitions) { + const int p896_begin = xqa_e4m3_g6_p896_begin(); + const int p1664_begin = xqa_e4m3_g6_p1664_begin(); + TORCH_CHECK(p256_begin < p512_begin && p512_begin < p896_begin && + p896_begin < p1664_begin, + "Invalid E4M3 wave-partition thresholds: ", p256_begin, + ", ", p512_begin, ", ", p896_begin, ", ", p1664_begin); + + const int p64_launch_num_partitions = std::max( + 1, std::min(launch_num_partitions, (p256_begin + 62) / 64)); + const int p256_launch_num_partitions = (launch_num_partitions + 3) / 4; + const int p512_launch_num_partitions = (launch_num_partitions + 7) / 8; + const int p896_launch_num_partitions = + (launch_num_partitions + 13) / 14; + const int p1664_launch_num_partitions = + (launch_num_partitions + 25) / 26; + const int p256_mid_launch_num_partitions = std::max( + 1, std::min(p256_launch_num_partitions, (p512_begin + 254) / 256)); + const int short_launch_num_partitions = + std::max(p64_launch_num_partitions, p256_mid_launch_num_partitions); + const int p512_route_launch_num_partitions = std::max( + 1, std::min(p512_launch_num_partitions, (p896_begin + 510) / 512)); + const int p896_route_launch_num_partitions = std::max( + 1, std::min(p896_launch_num_partitions, (p1664_begin + 894) / 896)); + const int wave_long_launch_num_partitions = std::max( + {p512_route_launch_num_partitions, p896_route_launch_num_partitions, + p1664_launch_num_partitions}); + const bool use_merged_wave_launch = + xqa_e4m3_g6_merged_wave_launch_enabled(); + + launch_flash_attention_decode_paged_xqa_tc_256_wide< + 0, 6, true, kXQATC256WideThreads, 1, 0, false, false, + kXQARouteShortSeqLens, false, false, false, + flash_v100::KV_CACHE_DTYPE_FP8_E4M3>( + q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, + max_logits, exp_sums, active_num_partitions, softmax_scale, k_scale, + v_scale, short_launch_num_partitions, false, 8, stream, p512_begin, + p256_begin, 0, false); + if (use_merged_wave_launch) { + if (fixed_interleaved_hkv1_layout) { + launch_flash_attention_decode_paged_xqa_e4m3_g6_page1568< + -1, kXQARouteWaveLongSeqLens, true>( + q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, + max_logits, exp_sums, active_num_partitions, softmax_scale, + k_scale, v_scale, wave_long_launch_num_partitions, stream, + p512_begin, p896_begin, p1664_begin, false); + } else { + launch_flash_attention_decode_paged_xqa_e4m3_g6_page1568< + -1, kXQARouteWaveLongSeqLens>( + q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, + max_logits, exp_sums, active_num_partitions, softmax_scale, + k_scale, v_scale, wave_long_launch_num_partitions, stream, + p512_begin, p896_begin, p1664_begin, false); + } + } else { + launch_flash_attention_decode_paged_xqa_e4m3_g6_page1568< + 512, kXQARouteRangeSeqLens>( + q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, + max_logits, exp_sums, active_num_partitions, softmax_scale, + k_scale, v_scale, p512_route_launch_num_partitions, stream, + p512_begin, p896_begin, 0, false); + launch_flash_attention_decode_paged_xqa_e4m3_g6_page1568< + 896, kXQARouteRangeSeqLens>( + q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, + max_logits, exp_sums, active_num_partitions, softmax_scale, + k_scale, v_scale, p896_route_launch_num_partitions, stream, + p896_begin, p1664_begin, 0, false); + launch_flash_attention_decode_paged_xqa_e4m3_g6_page1568< + 1664, kXQARouteLongSeqLens>( + q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, + max_logits, exp_sums, active_num_partitions, softmax_scale, + k_scale, v_scale, p1664_launch_num_partitions, stream, + p1664_begin, 0, 0, false); + } + + const dim3 reduce_grid(q.size(0), q.size(1), 1); + const dim3 reduce_block(kThreadsPerBlock); +#define LAUNCH_E4M3_WAVE_REDUCER(PARTITION_SIZE, SEQ_ROUTE, MAX_PARTITIONS, \ + PARTITION_BEGIN, ROUTE_BEGIN, ROUTE_END, \ + ROUTE_FINAL) \ + do { \ + const size_t reduce_shared_mem = \ + static_cast(2 * (MAX_PARTITIONS)) * sizeof(float); \ + flash_attention_decode_reduce_kernel<256, PARTITION_SIZE, SEQ_ROUTE> \ + <<>>( \ + reinterpret_cast(tmp_out.data_ptr()), \ + max_logits.data_ptr(), exp_sums.data_ptr(), \ + seq_lens.data_ptr(), active_num_partitions.data_ptr(), \ + reinterpret_cast<__half*>(out.data_ptr()), q.size(0), \ + MAX_PARTITIONS, q.size(1), 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), PARTITION_BEGIN, ROUTE_BEGIN, \ + ROUTE_END, ROUTE_FINAL); \ + } while (0) + LAUNCH_E4M3_WAVE_REDUCER(0, kXQARouteShortSeqLens, + short_launch_num_partitions, p256_begin, + p512_begin, 0, 0); + if (use_merged_wave_launch) { + LAUNCH_E4M3_WAVE_REDUCER(-1, kXQARouteWaveLongSeqLens, + wave_long_launch_num_partitions, 0, + p512_begin, p896_begin, p1664_begin); + } else { + LAUNCH_E4M3_WAVE_REDUCER(512, kXQARouteRangeSeqLens, + p512_route_launch_num_partitions, 0, + p512_begin, p896_begin, 0); + LAUNCH_E4M3_WAVE_REDUCER(896, kXQARouteRangeSeqLens, + p896_route_launch_num_partitions, 0, + p896_begin, p1664_begin, 0); + LAUNCH_E4M3_WAVE_REDUCER(1664, kXQARouteLongSeqLens, + p1664_launch_num_partitions, 0, p1664_begin, + 0, 0); + } +#undef LAUNCH_E4M3_WAVE_REDUCER + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + static bool traced_wave_partitions = false; + if (xqa_e4m3_g6_p64_p256_auto_trace_enabled() && + !traced_wave_partitions) { + TORCH_WARN( + "Flash-V100 XQA E4M3 G6 wave-partition route active; " + "thresholds p256/p512/p896/p1664=", + p256_begin, "/", p512_begin, "/", p896_begin, "/", p1664_begin, + ", launch partitions short/p512/p896/p1664=", + short_launch_num_partitions, "/", + p512_route_launch_num_partitions, "/", + p896_route_launch_num_partitions, "/", + p1664_launch_num_partitions, + ", merged_long=", use_merged_wave_launch, + ", converter=shared-lut"); + traced_wave_partitions = true; + } + return out; + } + const int dual_cta_begin = + std::max(p256_begin, xqa_e4m3_g6_dual_cta_begin()); + const int p64_launch_num_partitions = + std::max(1, std::min(launch_num_partitions, (p256_begin + 62) / 64)); + const int p256_launch_num_partitions = (launch_num_partitions + 3) / 4; + const int p256_mid_launch_num_partitions = std::max( + 1, + std::min(p256_launch_num_partitions, (dual_cta_begin + 254) / 256)); + const int short_launch_num_partitions = + std::max(p64_launch_num_partitions, p256_mid_launch_num_partitions); + const int reduce_num_partitions = + std::max(short_launch_num_partitions, p256_launch_num_partitions); + + launch_flash_attention_decode_paged_xqa_tc_256_wide< + 0, 6, true, kXQATC256WideThreads, 1, 0, false, false, + kXQARouteShortSeqLens, false, false, false, + flash_v100::KV_CACHE_DTYPE_FP8_E4M3>( + q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, max_logits, + exp_sums, active_num_partitions, softmax_scale, k_scale, v_scale, + short_launch_num_partitions, false, 8, stream, dual_cta_begin, + p256_begin, 0, false); + launch_flash_attention_decode_paged_xqa_tc_256_wide< + 256, 6, true, kXQATCG6DualCtaThreads, 2, 1568, false, false, + kXQARouteLongSeqLens, false, false, false, + flash_v100::KV_CACHE_DTYPE_FP8_E4M3, true>( + q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, max_logits, + exp_sums, active_num_partitions, softmax_scale, k_scale, v_scale, + p256_launch_num_partitions, false, 8, stream, dual_cta_begin, 0, 0, + false); + + const dim3 reduce_grid(q.size(0), q.size(1), 1); + const dim3 reduce_block(kThreadsPerBlock); + const size_t reduce_shared_mem = + static_cast(2 * reduce_num_partitions) * sizeof(float); + flash_attention_decode_reduce_kernel<256, 0> + <<>>( + reinterpret_cast(tmp_out.data_ptr()), + max_logits.data_ptr(), exp_sums.data_ptr(), + seq_lens.data_ptr(), active_num_partitions.data_ptr(), + reinterpret_cast<__half*>(out.data_ptr()), q.size(0), + reduce_num_partitions, q.size(1), 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), p256_begin, 0, + 0, 0); + + static bool traced_p64_p256_auto = false; + if (xqa_e4m3_g6_p64_p256_auto_trace_enabled() && !traced_p64_p256_auto) { + TORCH_WARN( + "Flash-V100 XQA E4M3 G6 device-side p64/p256 route active; " + "p256 threshold=", + p256_begin, ", p64 launch partitions=", p64_launch_num_partitions, + ", dual-CTA threshold=", dual_cta_begin, + ", short/mid launch partitions=", short_launch_num_partitions, + ", long p256 launch partitions=", p256_launch_num_partitions, + ", routed launch partitions=", reduce_num_partitions); + traced_p64_p256_auto = true; + } + } else if (q.size(0) == 1 && partition_size == 64) { + launch_flash_attention_decode_paged_xqa_tc_256_wide< + 64, 6, true, kXQATC256WideThreads, 1, 0, false, false, + kXQARouteAllSeqLens, false, false, false, + flash_v100::KV_CACHE_DTYPE_FP8_E4M3>( + q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, max_logits, + exp_sums, active_num_partitions, softmax_scale, k_scale, v_scale, + launch_num_partitions, false, 8, stream); + } else if (q.size(0) == 1 && partition_size == 128) { + launch_flash_attention_decode_paged_xqa_tc_256_wide< + 128, 6, true, kXQATC256WideThreads, 1, 0, false, false, + kXQARouteAllSeqLens, false, false, false, + flash_v100::KV_CACHE_DTYPE_FP8_E4M3>( + q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, max_logits, + exp_sums, active_num_partitions, softmax_scale, k_scale, v_scale, + launch_num_partitions, false, 8, stream); + } else if (q.size(0) == 1 && partition_size == 512) { + launch_flash_attention_decode_paged_xqa_e4m3_g6_page1568<512>( + q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, max_logits, + exp_sums, active_num_partitions, softmax_scale, k_scale, v_scale, + launch_num_partitions, stream); + } else if (q.size(0) == 1 && partition_size == 896) { + launch_flash_attention_decode_paged_xqa_e4m3_g6_page1568<896>( + q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, max_logits, + exp_sums, active_num_partitions, softmax_scale, k_scale, v_scale, + launch_num_partitions, stream); + } else if (q.size(0) == 1 && partition_size == 1024) { + launch_flash_attention_decode_paged_xqa_e4m3_g6_page1568<1024>( + q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, max_logits, + exp_sums, active_num_partitions, softmax_scale, k_scale, v_scale, + launch_num_partitions, stream); + } else if (q.size(0) == 1 && partition_size == 1664) { + launch_flash_attention_decode_paged_xqa_e4m3_g6_page1568<1664>( + q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, max_logits, + exp_sums, active_num_partitions, softmax_scale, k_scale, v_scale, + launch_num_partitions, stream); + } else if (partition_size == 64) { + LAUNCH_E4M3_BATCH_XQA(64); + } else if (partition_size == 128) { + LAUNCH_E4M3_BATCH_XQA(128); + } else { + LAUNCH_E4M3_BATCH_XQA(256); + } +#undef LAUNCH_E4M3_BATCH_XQA +#undef LAUNCH_E4M3_DUAL_CTA + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; + } // Default only the shape with exact direct and model-level evidence. const bool padded_smem_enabled = xqa_padded_smem_enabled(); const bool use_padded_smem = @@ -2702,6 +4281,20 @@ at::Tensor flash_attention_decode_paged_xqa( kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP8_E5M2 && !decode_partition_size_overridden() && xqa_e5m2_g6_dual_cta_enabled() && xqa_e5m2_g6_split_reduce_enabled(); + const bool batch_context_page_supported = + k_cache.size(1) == 16 || + (k_cache.size(1) >= 256 && k_cache.size(1) % 16 == 0); + const bool use_e5m2_g6_batch_context_route = + batch_context_max_seq_len > 0 && q.size(0) >= 4 && q_per_kv == 6 && + batch_context_page_supported && k_cache.size(2) == 1 && + k_cache.scalar_type() == at::kByte && + kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP8_E5M2 && + !decode_partition_size_overridden(); + const XQABatchContextRoute batch_context_route = + use_e5m2_g6_batch_context_route + ? select_xqa_batch_context_route(q.size(0), batch_context_max_seq_len, + partition_size) + : XQABatchContextRoute::kDisabled; const int e5m2_g6_dual_cta_seq_len = use_e5m2_g6_dual_cta ? at::cuda::getDeviceProperties(q.get_device())->multiProcessorCount * @@ -2729,13 +4322,27 @@ at::Tensor flash_attention_decode_paged_xqa( use_e5m2_g6_dual_cta && xqa_e5m2_partition_page_ids_enabled(); const bool use_e5m2_pair_load = use_e5m2_partition_page_ids && xqa_e5m2_pair_load_enabled(); + // The B1 E5M2 route already amortizes page-table lookups across a p256 + // partition and converts two half8 vectors from one aligned 128-bit load. + // Reuse that exact load path for the batch/context route without changing + // its partition, softmax, PV, or reduction order. Restrict the page size so + // one partition always fits the existing shared page-id capacity. + const bool use_e5m2_batch_wide_load = + use_e5m2_g6_batch_context_route && partition_size == 256 && + k_cache.size(1) >= 256 && k_cache.size(1) % 16 == 0 && + (batch_context_route == XQABatchContextRoute::kDualCta || + batch_context_route == XQABatchContextRoute::kDualCtaSplit) && + xqa_e5m2_batch_wide_load_enabled(); const bool use_g6_dual_cta = use_g6_p1024_auto || use_g6_p1024_sawtooth || use_mtp5_dual_cta || use_e5m2_g6_dual_cta || + batch_context_route == XQABatchContextRoute::kDualCta || + batch_context_route == XQABatchContextRoute::kDualCtaSplit || (xqa_g6_dual_cta_enabled() && (use_padded_smem || use_g6_dual_cta_dense)); - const bool use_split_reduce = use_g6_p1024_auto || use_g6_p1024_sawtooth || - use_e5m2_g6_dual_cta || - (use_g6_dual_cta && xqa_split_reduce_enabled()); + const bool use_split_reduce = + use_g6_p1024_auto || use_g6_p1024_sawtooth || use_e5m2_g6_dual_cta || + batch_context_route == XQABatchContextRoute::kDualCtaSplit || + (use_g6_dual_cta && xqa_split_reduce_enabled()); const bool supports_block16_index = use_g6_dual_cta && k_cache.size(1) == 16; const bool supports_block16_contiguous_layout = supports_block16_index && k_cache.size(2) == 1 && @@ -2762,6 +4369,17 @@ at::Tensor flash_attention_decode_paged_xqa( k_cache.stride(0), ",", k_cache.stride(1), ",", k_cache.stride(2), ",", k_cache.stride(3), "]"); } + trace_xqa_batch_context_route(q.size(0), batch_context_max_seq_len, + partition_size, k_cache.size(1), + batch_context_route); + static bool traced_e5m2_batch_wide_load = false; + if (xqa_batch_context_routing_trace_enabled() && use_e5m2_batch_wide_load && + !traced_e5m2_batch_wide_load) { + TORCH_WARN( + "Flash-V100 XQA batched E5M2 partition page reuse and paired " + "128-bit loads active"); + traced_e5m2_batch_wide_load = true; + } static bool traced_block16_layout = false; if (xqa_block16_layout_trace_enabled() && block16_layout_mode != 0 && !traced_block16_layout) { @@ -3082,12 +4700,22 @@ at::Tensor flash_attention_decode_paged_xqa( launch_num_partitions, use_split_reduce, split_reduce_dim_tile, stream); } else if (partition_size == 256) { - launch_flash_attention_decode_paged_xqa_tc_256_wide< - 256, 6, true, kXQATCG6DualCtaThreads, 2, 0, false>( - q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, max_logits, - exp_sums, active_num_partitions, softmax_scale, k_scale, v_scale, - launch_num_partitions, use_split_reduce, split_reduce_dim_tile, - stream); + if (use_e5m2_batch_wide_load) { + launch_flash_attention_decode_paged_xqa_tc_256_wide< + 256, 6, true, kXQATCG6DualCtaThreads, 2, 0, false, false, + kXQARouteAllSeqLens, false, true, true>( + q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, + max_logits, exp_sums, active_num_partitions, softmax_scale, k_scale, + v_scale, launch_num_partitions, use_split_reduce, + split_reduce_dim_tile, stream); + } else { + launch_flash_attention_decode_paged_xqa_tc_256_wide< + 256, 6, true, kXQATCG6DualCtaThreads, 2, 0, false>( + q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, + max_logits, exp_sums, active_num_partitions, softmax_scale, k_scale, + v_scale, launch_num_partitions, use_split_reduce, + split_reduce_dim_tile, stream); + } } else if (partition_size == 512) { launch_flash_attention_decode_paged_xqa_tc_256_wide< 512, 6, false, kXQATCG6DualCtaThreads, 2, 0, false>( 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_api.cpp b/flash-attention-v100/kernel/fused_mha_api.cpp index bbd31094df..ec1f942353 100644 --- a/flash-attention-v100/kernel/fused_mha_api.cpp +++ b/flash-attention-v100/kernel/fused_mha_api.cpp @@ -9,59 +9,41 @@ namespace py = pybind11; PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.doc() = "FlashAttention-2 implementation optimized for Volta"; - m.def("fwd", &flash_attention_forward, "FlashAttention-2 Forward Pass (Volta)"); - m.def( - "qk_scores_fwd", - &flash_attention_qk_scores, + m.doc() = "FlashAttention-2 implementation optimized for Volta"; + m.def("fwd", &flash_attention_forward, + "FlashAttention-2 Forward Pass (Volta)"); + m.def("qk_scores_fwd", &flash_attention_qk_scores, "Debug FlashAttention QK score dump before softmax (Volta)"); - m.def("bwd", &flash_attention_backward, "FlashAttention-2 Backward Pass (Volta)"); - m.def( - "decode_paged_fwd", - &flash_attention_decode_paged, + m.def("bwd", &flash_attention_backward, + "FlashAttention-2 Backward Pass (Volta)"); + m.def("decode_paged_fwd", &flash_attention_decode_paged, "FlashAttention decode over paged KV cache (Volta)"); - m.def( - "decode_paged_xqa_fwd", - &flash_attention_decode_paged_xqa, + m.def("decode_paged_xqa_fwd", &flash_attention_decode_paged_xqa, "FlashAttention XQA decode over paged KV cache (Volta)"); - m.def( - "decode_paged_xqa_staged_fwd", - &flash_attention_decode_paged_xqa_staged, + m.def("decode_paged_xqa_staged_fwd", &flash_attention_decode_paged_xqa_staged, "Staged FlashAttention XQA decode over paged KV cache (Volta)"); - m.def( - "decode_paged_wmma_fwd", - &flash_attention_decode_paged_wmma, - "FlashAttention single-query decode through paged-prefill WMMA order (Volta)"); - m.def( - "decode_qk_scores_fwd", - &flash_attention_decode_qk_scores, + m.def("grouped_verify_paged_fwd", &flash_attention_grouped_verify_paged, + "Exact grouped DFlash2 verification over paged KV cache (Volta)"); + m.def("decode_paged_wmma_fwd", &flash_attention_decode_paged_wmma, + "FlashAttention single-query decode through paged-prefill WMMA order " + "(Volta)"); + m.def("decode_qk_scores_fwd", &flash_attention_decode_qk_scores, "Debug scalar paged decode QK score dump before softmax (Volta)"); - m.def( - "decode_turboquant_paged_fwd", - &flash_attention_turboquant_decode_paged, + m.def("decode_turboquant_paged_fwd", &flash_attention_turboquant_decode_paged, "FlashAttention decode over TurboQuant paged KV cache (Volta)"); - m.def( - "prefill_paged_fwd", - &flash_attention_prefill_paged, + m.def("prefill_paged_fwd", &flash_attention_prefill_paged, "FlashAttention prefill over paged KV cache (Volta)"); - m.def( - "prefill_paged_d256_bm32_allp_pair_scratch_fwd", + m.def("prefill_paged_d256_bm32_allp_pair_scratch_fwd", &flash_attention_prefill_paged_d256_bm32_allp_pair_scratch, "Fixed causal D256 BM32 ALL_P pair-scratch paged prefill (SM70)"); - m.def( - "prefill_paged_d256_bm32_allp_pair_scratch_splitkv3_fwd", + m.def("prefill_paged_d256_bm32_allp_pair_scratch_splitkv3_fwd", &flash_attention_prefill_paged_d256_bm32_allp_pair_scratch_splitkv3, - "Fixed causal D256 BM32 ALL_P pair-scratch three-way split-KV paged prefill (SM70)"); - m.def( - "prefill_paged_bfla_fwd", - &flash_attention_prefill_paged_bfla, + "Fixed causal D256 BM32 ALL_P pair-scratch three-way split-KV paged " + "prefill (SM70)"); + m.def("prefill_paged_bfla_fwd", &flash_attention_prefill_paged_bfla, "BFLA sparse FlashAttention prefill over paged KV cache (Volta)"); - m.def( - "prefill_paged_splitkv_fwd", - &flash_attention_prefill_paged_splitkv, + m.def("prefill_paged_splitkv_fwd", &flash_attention_prefill_paged_splitkv, "FlashAttention split-KV prefill over paged KV cache (Volta)"); - m.def( - "fp8_e5m2_paged_kv_to_fp16", - &flash_attention_fp8_e5m2_paged_kv_to_fp16, + m.def("fp8_e5m2_paged_kv_to_fp16", &flash_attention_fp8_e5m2_paged_kv_to_fp16, "Expand paged FP8 E5M2 K/V into a preallocated FP16 paged workspace"); } 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 7a4d1056aa..538aaad265 100644 --- a/flash-attention-v100/kernel/fused_mha_forward_paged.cu +++ b/flash-attention-v100/kernel/fused_mha_forward_paged.cu @@ -3,6 +3,7 @@ #include #include #include +#include #include #include #include @@ -23,2296 +24,2064 @@ 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 namespace { int kv_cache_dtype_code_from_string(const std::string& kv_cache_dtype) { - if (kv_cache_dtype == "auto" || kv_cache_dtype == "bfloat16") { - return flash_v100::KV_CACHE_DTYPE_FP16; - } - if (kv_cache_dtype == "fp8" || kv_cache_dtype == "fp8_e4m3") { - return flash_v100::KV_CACHE_DTYPE_FP8_E4M3; - } - if (kv_cache_dtype == "fp8_e5m2") { - return flash_v100::KV_CACHE_DTYPE_FP8_E5M2; - } - return -1; + if (kv_cache_dtype == "auto" || kv_cache_dtype == "bfloat16") { + return flash_v100::KV_CACHE_DTYPE_FP16; + } + if (kv_cache_dtype == "fp8" || kv_cache_dtype == "fp8_e4m3") { + return flash_v100::KV_CACHE_DTYPE_FP8_E4M3; + } + if (kv_cache_dtype == "fp8_e5m2") { + return flash_v100::KV_CACHE_DTYPE_FP8_E5M2; + } + return -1; } } // namespace -#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 128 #define KV_STAGE_N_256_LOW_SMEM 1 -#define WARPS_256_LOW_SMEM 16 +#define WARPS_256_LOW_SMEM 16 #define BLOCK_N_256_LOW_SMEM_SCALAR_QK 32 #define KV_STAGE_N_256_LOW_SMEM_SCALAR_QK BLOCK_N_256_LOW_SMEM_SCALAR_QK -#define LOW_SMEM_PAGE_SIZE 16 +#define LOW_SMEM_PAGE_SIZE 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 : (LOW_SMEM ? (LOW_SMEM_BM32 ? BLOCK_M_256 : 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 ? (LOW_SMEM_SCALAR_QK ? BLOCK_N_256_LOW_SMEM_SCALAR_QK : 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 KV_STAGE_N = (D == 256 && LOW_SMEM) ? (LOW_SMEM_SCALAR_QK ? KV_STAGE_N_256_LOW_SMEM_SCALAR_QK : KV_STAGE_N_256_LOW_SMEM) : BLOCK_N; - 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 PIPELINE_S_ELEMENTS = - (D == 256 && LOW_SMEM && D256_SW_PIPELINE_QK - && D256_SW_PIPELINE_PV) - ? BLOCK_M * S_STRIDE - : 1; - static constexpr int PIPELINE_PAGE_ELEMENTS = - (D == 256 && LOW_SMEM && D256_SW_PIPELINE_QK - && D256_SW_PIPELINE_PV) - ? BLOCK_N / LOW_SMEM_PAGE_SIZE - : 1; - static constexpr bool PIPELINE_SWIZZLED_Q = - D == 256 && LOW_SMEM && D256_SW_PIPELINE_QK - && D256_SW_PIPELINE_PV; - static constexpr int Q_ELEMENTS = - PIPELINE_SWIZZLED_Q ? BLOCK_M * D : BLOCK_M * Q_STRIDE; - // The D256 low-SMEM candidate changes only the output accumulator pitch - // from 264 to 268 floats. This preserves all WMMA/softmax arithmetic and - // targets the largest 8-way shared-memory load/store groups. - static constexpr int OUTPUT_EXTRA_PAD = - (D == 256 && LOW_SMEM && D256_OUTPUT_STRIDE_268) ? 4 : 0; - static constexpr int O_STRIDE = D + PAD + OUTPUT_EXTRA_PAD; - static constexpr int PER_UINT4 = 8; - static constexpr int LOW_SMEM_PAGE_COUNT = - (D == 256 && LOW_SMEM) ? (BLOCK_N / LOW_SMEM_PAGE_SIZE) : 1; - struct alignas(128) SmemLayout { - alignas(16) __half q [Q_ELEMENTS]; + 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 ? (LOW_SMEM_BM32 ? BLOCK_M_256 : 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 ? (LOW_SMEM_SCALAR_QK ? BLOCK_N_256_LOW_SMEM_SCALAR_QK + : 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 KV_STAGE_N = + (D == 256 && LOW_SMEM) + ? (LOW_SMEM_SCALAR_QK ? KV_STAGE_N_256_LOW_SMEM_SCALAR_QK + : KV_STAGE_N_256_LOW_SMEM) + : BLOCK_N; + 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 PIPELINE_S_ELEMENTS = + (D == 256 && LOW_SMEM && D256_SW_PIPELINE_QK && D256_SW_PIPELINE_PV) + ? BLOCK_M * S_STRIDE + : 1; + static constexpr int PIPELINE_PAGE_ELEMENTS = + (D == 256 && LOW_SMEM && D256_SW_PIPELINE_QK && D256_SW_PIPELINE_PV) + ? BLOCK_N / LOW_SMEM_PAGE_SIZE + : 1; + static constexpr bool PIPELINE_SWIZZLED_Q = + D == 256 && LOW_SMEM && D256_SW_PIPELINE_QK && D256_SW_PIPELINE_PV; + static constexpr int Q_ELEMENTS = + PIPELINE_SWIZZLED_Q ? BLOCK_M * D : BLOCK_M * Q_STRIDE; + // The D256 low-SMEM candidate changes only the output accumulator pitch + // from 264 to 268 floats. This preserves all WMMA/softmax arithmetic and + // targets the largest 8-way shared-memory load/store groups. + static constexpr int OUTPUT_EXTRA_PAD = + (D == 256 && LOW_SMEM && D256_OUTPUT_STRIDE_268) ? 4 : 0; + static constexpr int O_STRIDE = D + PAD + OUTPUT_EXTRA_PAD; + static constexpr int PER_UINT4 = 8; + static constexpr int LOW_SMEM_PAGE_COUNT = + (D == 256 && LOW_SMEM) ? (BLOCK_N / LOW_SMEM_PAGE_SIZE) : 1; + struct alignas(128) SmemLayout { + alignas(16) __half q[Q_ELEMENTS]; union { - alignas(16) __half k [KV_STAGE_N * KV_STRIDE]; - alignas(16) __half v [KV_STAGE_N * KV_STRIDE]; + alignas(16) __half k[KV_STAGE_N * KV_STRIDE]; + alignas(16) __half v[KV_STAGE_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 pipeline_s[PIPELINE_S_ELEMENTS]; - alignas(16) int pipeline_page_idx[PIPELINE_PAGE_ELEMENTS]; - alignas(16) int pipeline_page_offset[PIPELINE_PAGE_ELEMENTS]; - // WMMA accumulator loads and stores require 256-bit alignment. - alignas(32) float o [BLOCK_M * O_STRIDE]; - alignas(16) float row_max[BLOCK_M]; - alignas(16) float row_sum[BLOCK_M]; - alignas(16) int page_idx[LOW_SMEM_PAGE_COUNT]; - alignas(16) int page_offset[LOW_SMEM_PAGE_COUNT]; - }; - - static constexpr size_t TOTAL_SMEM = ((sizeof(SmemLayout) + 127) & ~size_t(127)); - static_assert(offsetof(SmemLayout, o) % 32 == 0, - "WMMA accumulator storage must be 256-bit aligned"); + alignas(16) __half p_strict[P_STRICT_ELEMENTS]; + alignas(16) float pipeline_s[PIPELINE_S_ELEMENTS]; + alignas(16) int pipeline_page_idx[PIPELINE_PAGE_ELEMENTS]; + alignas(16) int pipeline_page_offset[PIPELINE_PAGE_ELEMENTS]; + // WMMA accumulator loads and stores require 256-bit alignment. + alignas(32) float o[BLOCK_M * O_STRIDE]; + alignas(16) float row_max[BLOCK_M]; + alignas(16) float row_sum[BLOCK_M]; + alignas(16) int page_idx[LOW_SMEM_PAGE_COUNT]; + alignas(16) int page_offset[LOW_SMEM_PAGE_COUNT]; + }; + + static constexpr size_t TOTAL_SMEM = + ((sizeof(SmemLayout) + 127) & ~size_t(127)); + static_assert(offsetof(SmemLayout, o) % 32 == 0, + "WMMA accumulator storage must be 256-bit aligned"); }; -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(); } __device__ __forceinline__ int pipeline_swizzled_q_row_slot(int row) { - return (row & 3) | ((row & 8) >> 1) | ((row & 4) << 1); + return (row & 3) | ((row & 8) >> 1) | ((row & 4) << 1); } -__device__ __forceinline__ int pipeline_swizzled_matrix_a_offset( - int row, - int col) { - const int tile = col / WMMA_K; - const int within_tile = col - tile * WMMA_K; - const int plane = within_tile >> 3; - const int inner = within_tile & 7; - return tile * WMMA_M * WMMA_K + plane * WMMA_M * 8 - + pipeline_swizzled_q_row_slot(row) * 8 + inner; +__device__ __forceinline__ int pipeline_swizzled_matrix_a_offset(int row, + int col) { + const int tile = col / WMMA_K; + const int within_tile = col - tile * WMMA_K; + const int plane = within_tile >> 3; + const int inner = within_tile & 7; + return tile * WMMA_M * WMMA_K + plane * WMMA_M * 8 + + pipeline_swizzled_q_row_slot(row) * 8 + inner; } __device__ __forceinline__ void load_pipeline_swizzled_matrix_a_fragment( fragment& frag, - const __half* __restrict__ matrix, - int k_offset) { - const int lane = threadIdx.x & 31; - const int row = - (lane & 3) + ((lane >> 4) & 1) * 4 + ((lane >> 2) & 1) * 8; - const int slot = pipeline_swizzled_q_row_slot(row); - const int tile_offset = (k_offset / WMMA_K) * WMMA_M * WMMA_K; - uint32_t address = static_cast( - __cvta_generic_to_shared(matrix + tile_offset + slot * 8)); - uint32_t* words = reinterpret_cast(&frag); - asm volatile( - "ld.shared.v4.u32 {%0, %1, %2, %3}, [%4];" - : "=r"(words[0]), "=r"(words[1]), "=r"(words[2]), "=r"(words[3]) - : "r"(address) - : "memory"); - address += WMMA_M * 8 * sizeof(__half); - asm volatile( - "ld.shared.v4.u32 {%0, %1, %2, %3}, [%4];" - : "=r"(words[4]), "=r"(words[5]), "=r"(words[6]), "=r"(words[7]) - : "r"(address) - : "memory"); + const __half* __restrict__ matrix, int k_offset) { + const int lane = threadIdx.x & 31; + const int row = (lane & 3) + ((lane >> 4) & 1) * 4 + ((lane >> 2) & 1) * 8; + const int slot = pipeline_swizzled_q_row_slot(row); + const int tile_offset = (k_offset / WMMA_K) * WMMA_M * WMMA_K; + uint32_t address = static_cast( + __cvta_generic_to_shared(matrix + tile_offset + slot * 8)); + uint32_t* words = reinterpret_cast(&frag); + asm volatile("ld.shared.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(words[0]), "=r"(words[1]), "=r"(words[2]), "=r"(words[3]) + : "r"(address) + : "memory"); + address += WMMA_M * 8 * sizeof(__half); + asm volatile("ld.shared.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(words[4]), "=r"(words[5]), "=r"(words[6]), "=r"(words[7]) + : "r"(address) + : "memory"); } -template +template __device__ __forceinline__ void pipeline_qk_slice( - const __half* __restrict__ sQ, - float* __restrict__ score, - const __half* __restrict__ k_cache, - const int* __restrict__ page_idx, - const int* __restrict__ page_offset, - int64_t k_block_stride, - int64_t k_token_stride, - int64_t k_head_stride, - int kv_head_id, - int start_col, - int valid_k_rows, - int start_row, - int valid_q_rows, - int causal_q_offset, - int tile_n, - int k_begin, - bool load_partial, - bool finalize, - float softmax_scale, - int window_size_left, - int window_size_right, - float neg_inf) { - if (tile_n >= valid_k_rows) { - return; - } + const __half* __restrict__ sQ, float* __restrict__ score, + const __half* __restrict__ k_cache, const int* __restrict__ page_idx, + const int* __restrict__ page_offset, int64_t k_block_stride, + int64_t k_token_stride, int64_t k_head_stride, int kv_head_id, + int start_col, int valid_k_rows, int start_row, int valid_q_rows, + int causal_q_offset, int tile_n, int k_begin, bool load_partial, + bool finalize, float softmax_scale, int window_size_left, + int window_size_right, float neg_inf, int anchor_len = 0, + int anchored_window = 0) { + if (tile_n >= valid_k_rows) { + return; + } + + const int page_slot = tile_n >> 4; + const int block_offset = page_offset[page_slot]; + const int physical_block_idx = page_idx[page_slot]; + const __half* k_tile_base = k_cache + + (int64_t)physical_block_idx * k_block_stride + + (int64_t)block_offset * k_token_stride + + (int64_t)kv_head_id * k_head_stride; + + fragment a_frag; + fragment b_frag; + fragment acc_frag; + if (load_partial) { + load_matrix_sync(acc_frag, score + tile_n, S_STRIDE, mem_row_major); + } else { + fill_fragment(acc_frag, 0.0f); + } + +#pragma unroll + for (int k_tile = 0; k_tile < K_TILES; ++k_tile) { + const int k_offset = k_begin + k_tile * WMMA_K; + if constexpr (SWIZZLED_Q) { + load_pipeline_swizzled_matrix_a_fragment(a_frag, sQ, k_offset); + } else { + load_matrix_sync(a_frag, sQ + k_offset, Q_STRIDE); + } + load_matrix_sync(b_frag, k_tile_base + k_offset, k_token_stride); + mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + if (finalize) { + const int lane_id = threadIdx.x & 31; + 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; + +#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 + 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 = 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; + } + bool is_anchor_valid = true; + if constexpr (ANCHORED_SWA) { + is_anchor_valid = (global_n < anchor_len) || + (global_q_pos - global_n < anchored_window); + } + acc_frag.x[i] = + (is_valid && is_causal_valid && is_window_valid && is_anchor_valid) + ? acc_frag.x[i] * softmax_scale + : neg_inf; + } + } + + store_matrix_sync(score + tile_n, acc_frag, S_STRIDE, mem_row_major); +} - const int page_slot = tile_n >> 4; - const int block_offset = page_offset[page_slot]; +template +__device__ __forceinline__ void pipeline_pv_tile( + const __half* __restrict__ sP, float* __restrict__ sO, + const __half* __restrict__ v_cache, const int* __restrict__ page_idx, + const int* __restrict__ page_offset, int64_t v_block_stride, + int64_t v_token_stride, int64_t v_head_stride, int kv_head_id, + int sub_start, int sub_valid_k_rows, int tile_d) { + fragment a_frag; + fragment b_frag; + fragment acc_frag; + load_matrix_sync(acc_frag, sO + tile_d, O_STRIDE, mem_row_major); + +#pragma unroll + for (int tile_k = 0; tile_k < 2; ++tile_k) { + const int k_offset = tile_k * WMMA_K; + if (k_offset >= sub_valid_k_rows) { + break; + } + const int token_offset = sub_start + k_offset; + const int page_slot = token_offset >> 4; + const int block_offset = page_offset[page_slot] + (token_offset & 15); const int physical_block_idx = page_idx[page_slot]; - const __half* k_tile_base = - k_cache + (int64_t)physical_block_idx * k_block_stride - + (int64_t)block_offset * k_token_stride - + (int64_t)kv_head_id * k_head_stride; - - fragment a_frag; - fragment b_frag; - fragment acc_frag; - if (load_partial) { - load_matrix_sync( - acc_frag, score + tile_n, S_STRIDE, mem_row_major); + const __half* v_tile_ptr = v_cache + + (int64_t)physical_block_idx * v_block_stride + + (int64_t)block_offset * v_token_stride + + (int64_t)kv_head_id * v_head_stride + tile_d; + if constexpr (SWIZZLED_P) { + load_pipeline_swizzled_matrix_a_fragment(a_frag, sP, k_offset); } else { - fill_fragment(acc_frag, 0.0f); + load_matrix_sync(a_frag, sP + k_offset, P_STRIDE); } + load_matrix_sync(b_frag, v_tile_ptr, v_token_stride); + mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + store_matrix_sync(sO + tile_d, acc_frag, O_STRIDE, mem_row_major); +} - #pragma unroll - for (int k_tile = 0; k_tile < K_TILES; ++k_tile) { - const int k_offset = k_begin + k_tile * WMMA_K; - if constexpr (SWIZZLED_Q) { - load_pipeline_swizzled_matrix_a_fragment( - a_frag, sQ, k_offset); - } else { - load_matrix_sync(a_frag, sQ + k_offset, Q_STRIDE); +template +__global__ void __launch_bounds__( + KernelConfig::THREADS_PER_BLOCK, + 2) + flash_attention_forward_kernel_paged( + const __half* __restrict__ Q, const void* __restrict__ K_cache, + const void* __restrict__ V_cache, __half* __restrict__ Out, + float* __restrict__ softmax_lse, const int* __restrict__ block_table, + const int* __restrict__ seqused_k, const int B, const int H, + const int M, const int N, const int* __restrict__ bfla_block_mask, + const int bfla_mask_block_n, const int64_t bfla_mask_stride_b, + const int64_t bfla_mask_stride_h, const int64_t bfla_mask_stride_q, + const int64_t bfla_mask_stride_k, const int page_block_size, + const int num_kv_heads, const int64_t k_block_stride, + const int64_t k_token_stride, 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 softmax_scale, + const float k_scale, const float v_scale, const int window_size_left, + const int window_size_right, float* __restrict__ split_tmp_out, + float* __restrict__ split_tmp_row_max, + float* __restrict__ split_tmp_row_sum, const int split_kv_tiles, + const int* __restrict__ anchor_lens, const int anchored_window) { + static_assert(!ANCHORED_SWA || + (IS_CAUSAL && KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP16), + "ANCHORED_SWA requires causal attention and an fp16 KV cache"); + using Config = KernelConfig; + using Traits = FlashV100Traits; + + 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; + constexpr bool USE_SW_PIPELINE = D256_SW_PIPELINE_QK || D256_SW_PIPELINE_PV; + if constexpr (USE_SW_PIPELINE) { + static_assert(D == 256 && LOW_SMEM && !LOW_SMEM_CONTIG_FAST && + !LOW_SMEM_SCALAR_QK && !LOW_SMEM_BM32 && !SPLIT_KV && + BLOCK_M == 16 && BLOCK_N == 128 && WARPS_PER_BLOCK == 16, + "software pipeline is specialized for D256 BM16 BN128"); + } + 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 / num_kv_heads; + 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; + + const int actual_N = seqused_k[batch_id]; + int num_n_tiles = (actual_N + BLOCK_N - 1) / BLOCK_N; + const int valid_q_rows = min(BLOCK_M, M - start_row); + const int causal_q_offset = max(actual_N - M, 0); + + int anchor_len = 0; + if constexpr (ANCHORED_SWA) { + anchor_len = anchor_lens[batch_id]; + } + + int max_key_pos = actual_N - 1; + if constexpr (IS_CAUSAL) { + max_key_pos = + min(max_key_pos, start_row + valid_q_rows - 1 + causal_q_offset); + } + if (window_size_right >= 0) { + max_key_pos = min(max_key_pos, start_row + valid_q_rows - 1 + + causal_q_offset + 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, start_row + causal_q_offset - 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 __half* q_ptr = Q + q_head_linear * M * D + start_row * 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; + + const int max_num_blocks_per_seq = + (N + page_block_size - 1) / page_block_size; + const int* block_table_seq = block_table + batch_id * max_num_blocks_per_seq; + + const int64_t k_row_stride = k_token_stride; + const int64_t v_row_stride = v_token_stride; + + extern __shared__ __align__(128) 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; + int* sPageIdx = smem.page_idx; + int* sPageOffset = smem.page_offset; + + 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; + sRowSum[tid] = 0.0f; + } + constexpr int O_UINT4_PER_ROW = D / 4; + constexpr int O_UINT4_STRIDE = O_STRIDE / 4; + for (int idx = tid; idx < BLOCK_M * O_UINT4_PER_ROW; + idx += THREADS_PER_BLOCK) { + const int row = idx / O_UINT4_PER_ROW; + const int col = idx % O_UINT4_PER_ROW; + reinterpret_cast(sO)[row * O_UINT4_STRIDE + col] = + make_uint4(0, 0, 0, 0); + } + __syncthreads(); + + 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]); + } + if constexpr (Config::PIPELINE_SWIZZLED_Q) { + const int k_tile = vec_col >> 1; + const int plane = vec_col & 1; + const int slot = pipeline_swizzled_q_row_slot(row); + sQ_vec[k_tile * (2 * BLOCK_M) + plane * BLOCK_M + slot] = q_val; + } else { + sQ_vec[row * q_stride_uint4 + vec_col] = q_val; + } + } + __syncthreads(); + + int cross_block_first_n_tile = 0; + if constexpr (D256_SW_PIPELINE_QK && D256_SW_PIPELINE_PV) { + const bool use_cross_block_pipeline = + valid_q_rows == BLOCK_M && actual_N > 0 && page_block_size == 784 && + bfla_block_mask == nullptr && window_size_left < 0 && + window_size_right < 0; + if (use_cross_block_pipeline) { + const __half* k_cache_h = reinterpret_cast(K_cache); + const __half* v_cache_h = reinterpret_cast(V_cache); + + if (tid < Config::LOW_SMEM_PAGE_COUNT) { + const int global_token_idx = tid * LOW_SMEM_PAGE_SIZE; + const int virtual_block_idx = global_token_idx / page_block_size; + sPageIdx[tid] = __ldg(&block_table_seq[virtual_block_idx]); + sPageOffset[tid] = + global_token_idx - virtual_block_idx * page_block_size; + } + __syncthreads(); + + if (warp_id < BLOCK_N / WMMA_N) { + const int valid_k_rows = min(BLOCK_N, actual_N); + pipeline_qk_slice( + sQ, sS, k_cache_h, sPageIdx, sPageOffset, k_block_stride, + k_token_stride, k_head_stride, kv_head_id, 0, valid_k_rows, + start_row, valid_q_rows, causal_q_offset, warp_id * WMMA_N, 0, + false, true, softmax_scale, window_size_left, window_size_right, + NEG_INF, anchor_len, anchored_window); + } + __syncthreads(); + + // Keep the hot loop on a compile-time steady-state path. The + // existing exact path below drains the final block, avoiding a + // loop-carried has_next predicate in every softmax/PV panel. + const int steady_n_tiles = num_n_tiles - 1; + for (int block_n = 0; block_n < steady_n_tiles; ++block_n) { + const int start_col = block_n * BLOCK_N; + const int valid_k_rows = BLOCK_N; + const int next_start_col = start_col + BLOCK_N; + const int next_valid_k_rows = min(BLOCK_N, actual_N - next_start_col); + const bool odd_block = (block_n & 1) != 0; + float* const score_current = odd_block ? smem.pipeline_s : sS; + float* const score_next = odd_block ? sS : smem.pipeline_s; + int* const page_idx_current = + odd_block ? smem.pipeline_page_idx : sPageIdx; + int* const page_offset_current = + odd_block ? smem.pipeline_page_offset : sPageOffset; + int* const page_idx_next = + odd_block ? sPageIdx : smem.pipeline_page_idx; + int* const page_offset_next = + odd_block ? sPageOffset : smem.pipeline_page_offset; + + if (tid < Config::LOW_SMEM_PAGE_COUNT) { + const int page_token_offset = tid * LOW_SMEM_PAGE_SIZE; + const int global_token_idx = next_start_col + page_token_offset; + const int virtual_block_idx = global_token_idx / page_block_size; + page_idx_next[tid] = __ldg(&block_table_seq[virtual_block_idx]); + page_offset_next[tid] = + global_token_idx - virtual_block_idx * page_block_size; } - load_matrix_sync( - b_frag, k_tile_base + k_offset, k_token_stride); - mma_sync(acc_frag, a_frag, b_frag, acc_frag); - } - - if (finalize) { - const int lane_id = threadIdx.x & 31; - 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; - - #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 + row; - const int global_n = start_col + tile_n + col; - const int global_q_pos = global_m + causal_q_offset; + __syncthreads(); - 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; + for (int sub_start = 0; sub_start < valid_k_rows; + sub_start += P_SUB_TILE) { + const int sub_valid_k_rows = + min(P_SUB_TILE, valid_k_rows - sub_start); + const int row = warp_id; + const int thread_in_row = lane_id; + const unsigned mask = 0xFFFFFFFFU; + float* sS_row_f = score_current + row * S_STRIDE + sub_start; + + 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) { + const int vc = thread_in_row + j * THREADS_PER_ROW; + if (vc < vec_cols) { + const float4 v4 = sS_vec4[vc]; + thread_max = fmaxf(thread_max, + fmaxf(fmaxf(v4.x, v4.y), fmaxf(v4.z, v4.w))); } - bool is_window_valid = true; - if (window_size_left >= 0) { - is_window_valid = - global_n >= global_q_pos - window_size_left; + } +#pragma unroll + 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)); + } + + const float row_max = __shfl_sync(mask, thread_max, 0); + 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; + int vc_base = thread_in_row; + +#pragma unroll 4 + for (int j = 0; j < vecs_per_thread; + ++j, vc_base += THREADS_PER_ROW) { + if (vc_base < vec_cols) { + const float4 v4 = sS_vec4[vc_base]; + const float e0 = __expf(fmaxf(v4.x - new_max, -80.0f)); + const float e1 = __expf(fmaxf(v4.y - new_max, -80.0f)); + const float e2 = __expf(fmaxf(v4.z - new_max, -80.0f)); + const float e3 = __expf(fmaxf(v4.w - new_max, -80.0f)); + thread_sum += (e0 + e1) + (e2 + e3); + const int p_col = vc_base * 4; + __half2* p_half2 = reinterpret_cast<__half2*>( + sP + pipeline_swizzled_matrix_a_offset(row, p_col)); + p_half2[0] = __float22half2_rn(make_float2(e0, e1)); + p_half2[1] = __float22half2_rn(make_float2(e2, e3)); } - if (window_size_right >= 0) { - is_window_valid = - is_window_valid - && global_n <= global_q_pos + window_size_right; + } +#pragma unroll 4 + for (int c = tail_start + thread_in_row; c < sub_valid_k_rows; + c += THREADS_PER_ROW) { + const float e = __expf(fmaxf(sS_row_f[c] - new_max, -80.0f)); + thread_sum += e; + sP[pipeline_swizzled_matrix_a_offset(row, c)] = __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); + } + const float row_sum = __shfl_sync(mask, thread_sum, 0); + if (thread_in_row == 0) { + sRowSum[row] = exp_diff * sRowSum[row] + row_sum; + sRowMax[row] = new_max; + } + +#pragma unroll 4 + for (int c = tail_start + thread_in_row; c < P_SUB_TILE; + c += THREADS_PER_ROW) { + if (c >= sub_valid_k_rows) { + sP[pipeline_swizzled_matrix_a_offset(row, c)] = + __float2half(0.0f); } - acc_frag.x[i] = - (is_valid && is_causal_valid && is_window_valid) - ? acc_frag.x[i] * softmax_scale - : neg_inf; - } - } - - store_matrix_sync(score + tile_n, acc_frag, S_STRIDE, mem_row_major); -} + } + + if (block_n > 0 || sub_start > 0) { + float4* sO_vec = reinterpret_cast(sO + row * O_STRIDE); +#pragma unroll 4 + for (int ov = thread_in_row; ov < D / 4; ov += THREADS_PER_ROW) { + float4 v = sO_vec[ov]; + v.x *= exp_diff; + v.y *= exp_diff; + v.z *= exp_diff; + v.w *= exp_diff; + sO_vec[ov] = v; + } + } + __syncthreads(); -template -__device__ __forceinline__ void pipeline_pv_tile( - const __half* __restrict__ sP, - float* __restrict__ sO, - const __half* __restrict__ v_cache, - const int* __restrict__ page_idx, - const int* __restrict__ page_offset, - int64_t v_block_stride, - int64_t v_token_stride, - int64_t v_head_stride, - int kv_head_id, - int sub_start, - int sub_valid_k_rows, - int tile_d) { - fragment a_frag; - fragment b_frag; - fragment acc_frag; - load_matrix_sync( - acc_frag, sO + tile_d, O_STRIDE, mem_row_major); - - #pragma unroll - for (int tile_k = 0; tile_k < 2; ++tile_k) { - const int k_offset = tile_k * WMMA_K; - if (k_offset >= sub_valid_k_rows) { - break; + if (sub_start == 0) { + if (warp_id < BLOCK_N / WMMA_N) { + pipeline_qk_slice( + sQ, score_next, k_cache_h, page_idx_next, page_offset_next, + k_block_stride, k_token_stride, k_head_stride, kv_head_id, + next_start_col, next_valid_k_rows, start_row, valid_q_rows, + causal_q_offset, warp_id * WMMA_N, 0, false, true, + softmax_scale, window_size_left, window_size_right, NEG_INF, + anchor_len, anchored_window); + } else { + const int pv_warp = warp_id - BLOCK_N / WMMA_N; + pipeline_pv_tile( + sP, sO, v_cache_h, page_idx_current, page_offset_current, + v_block_stride, v_token_stride, v_head_stride, kv_head_id, + sub_start, sub_valid_k_rows, pv_warp * WMMA_N); + pipeline_pv_tile( + sP, sO, v_cache_h, page_idx_current, page_offset_current, + v_block_stride, v_token_stride, v_head_stride, kv_head_id, + sub_start, sub_valid_k_rows, + (pv_warp + BLOCK_N / WMMA_N) * WMMA_N); + } + } else { + pipeline_pv_tile( + sP, sO, v_cache_h, page_idx_current, page_offset_current, + v_block_stride, v_token_stride, v_head_stride, kv_head_id, + sub_start, sub_valid_k_rows, warp_id * WMMA_N); + } + __syncthreads(); } - const int token_offset = sub_start + k_offset; - const int page_slot = token_offset >> 4; - const int block_offset = - page_offset[page_slot] + (token_offset & 15); - const int physical_block_idx = page_idx[page_slot]; - const __half* v_tile_ptr = - v_cache + (int64_t)physical_block_idx * v_block_stride - + (int64_t)block_offset * v_token_stride - + (int64_t)kv_head_id * v_head_stride + tile_d; - if constexpr (SWIZZLED_P) { - load_pipeline_swizzled_matrix_a_fragment( - a_frag, sP, k_offset); + } + cross_block_first_n_tile = steady_n_tiles; + } + } + + int first_n_tile = cross_block_first_n_tile; + int last_n_tile = num_n_tiles; + if constexpr (SPLIT_KV) { + const int partition_id = blockIdx.y; + const int tiles_per_partition = split_kv_tiles > 1 ? split_kv_tiles : 1; + first_n_tile = partition_id * tiles_per_partition; + last_n_tile = min(num_n_tiles, first_n_tile + tiles_per_partition); + } + + for (int block_n = first_n_tile; block_n < last_n_tile; ++block_n) { + const int start_col = block_n * BLOCK_N; + if (start_col >= actual_N) break; + const int valid_k_rows = min(BLOCK_N, actual_N - start_col); + + if (start_col + valid_k_rows <= min_key_pos) { + continue; + } + + if (bfla_block_mask != nullptr && bfla_mask_block_n > 0) { + const int mask_q_idx = start_row / bfla_mask_block_n; + const int mask_k_idx = start_col / bfla_mask_block_n; + const int keep_tile = + __ldg(bfla_block_mask + (int64_t)batch_id * bfla_mask_stride_b + + (int64_t)kv_head_id * bfla_mask_stride_h + + (int64_t)mask_q_idx * bfla_mask_stride_q + + (int64_t)mask_k_idx * bfla_mask_stride_k); + if (keep_tile == 0) { + continue; + } + } + + const int partial_block_size = + (block_n == num_n_tiles - 1 && actual_N % BLOCK_N != 0) + ? (actual_N % BLOCK_N) + : -1; + + uint4* sK_vec = reinterpret_cast(sK); + const int64_t row_stride_uint4 = k_row_stride / PER_UINT4; + const int start_page = start_col / page_block_size; + const int page_offset = start_col % page_block_size; + const bool single_page_tile = + (page_offset + valid_k_rows) <= page_block_size; + const bool two_page_tile = + !single_page_tile && + (page_offset + valid_k_rows) <= (page_block_size * 2); + const int first_page_rows = + single_page_tile ? valid_k_rows + : min(valid_k_rows, page_block_size - page_offset); + const int second_page_rows = valid_k_rows - first_page_rows; + const int physical_block_idx0 = __ldg(&block_table_seq[start_page]); + const int physical_block_idx1 = + second_page_rows > 0 ? __ldg(&block_table_seq[start_page + 1]) : -1; + const bool four_page_aligned_tile = + D == 256 && page_block_size == 16 && page_offset == 0 && + BLOCK_N == page_block_size * 4 && valid_k_rows == BLOCK_N && + k_block_stride == (int64_t)page_block_size * k_token_stride && + v_block_stride == (int64_t)page_block_size * v_token_stride; + const int physical_block_idx2 = + four_page_aligned_tile ? __ldg(&block_table_seq[start_page + 2]) : -1; + const int physical_block_idx3 = + four_page_aligned_tile ? __ldg(&block_table_seq[start_page + 3]) : -1; + const bool four_page_contiguous_tile = + four_page_aligned_tile && + physical_block_idx1 == physical_block_idx0 + 1 && + physical_block_idx2 == physical_block_idx0 + 2 && + physical_block_idx3 == physical_block_idx0 + 3; + bool low_smem_contiguous_page_tile = false; + if constexpr (LOW_SMEM) { + static_assert(D == 256, "low-smem paged prefill is D=256 only"); + static_assert(KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP16, + "low-smem paged prefill is fp16-KV only"); + static_assert(BLOCK_N % LOW_SMEM_PAGE_SIZE == 0, + "low-smem BLOCK_N must be page aligned"); + + constexpr int LOW_SMEM_PAGE_COUNT = Config::LOW_SMEM_PAGE_COUNT; + if (tid < LOW_SMEM_PAGE_COUNT) { + const int page_token_offset = tid * LOW_SMEM_PAGE_SIZE; + if (page_token_offset < valid_k_rows) { + const int global_token_idx = start_col + page_token_offset; + const int virtual_block_idx = global_token_idx / page_block_size; + sPageIdx[tid] = __ldg(&block_table_seq[virtual_block_idx]); + sPageOffset[tid] = + global_token_idx - virtual_block_idx * page_block_size; } else { - load_matrix_sync(a_frag, sP + k_offset, P_STRIDE); - } - load_matrix_sync(b_frag, v_tile_ptr, v_token_stride); - mma_sync(acc_frag, a_frag, b_frag, acc_frag); - } - store_matrix_sync(sO + tile_d, acc_frag, O_STRIDE, mem_row_major); -} - -template -__global__ void __launch_bounds__( - KernelConfig::THREADS_PER_BLOCK, 2) -flash_attention_forward_kernel_paged( - const __half* __restrict__ Q, - const void* __restrict__ K_cache, - const void* __restrict__ V_cache, - __half* __restrict__ Out, - float* __restrict__ softmax_lse, - const int* __restrict__ block_table, - const int* __restrict__ seqused_k, - const int B, - const int H, - const int M, - const int N, - const int* __restrict__ bfla_block_mask, - const int bfla_mask_block_n, - const int64_t bfla_mask_stride_b, - const int64_t bfla_mask_stride_h, - const int64_t bfla_mask_stride_q, - const int64_t bfla_mask_stride_k, - const int page_block_size, - const int num_kv_heads, - const int64_t k_block_stride, - const int64_t k_token_stride, - 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 softmax_scale, - const float k_scale, - const float v_scale, - const int window_size_left, - const int window_size_right, - float* __restrict__ split_tmp_out, - float* __restrict__ split_tmp_row_max, - float* __restrict__ split_tmp_row_sum, - const int split_kv_tiles -) { - using Config = KernelConfig; - using Traits = FlashV100Traits; - - 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; - constexpr bool USE_SW_PIPELINE = - D256_SW_PIPELINE_QK || D256_SW_PIPELINE_PV; - if constexpr (USE_SW_PIPELINE) { - static_assert(D == 256 && LOW_SMEM && !LOW_SMEM_CONTIG_FAST - && !LOW_SMEM_SCALAR_QK && !LOW_SMEM_BM32 - && !SPLIT_KV && BLOCK_M == 16 && BLOCK_N == 128 - && WARPS_PER_BLOCK == 16, - "software pipeline is specialized for D256 BM16 BN128"); - } - 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 / num_kv_heads; - 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; - - const int actual_N = seqused_k[batch_id]; - int num_n_tiles = (actual_N + BLOCK_N - 1) / BLOCK_N; - const int valid_q_rows = min(BLOCK_M, M - start_row); - const int causal_q_offset = max(actual_N - M, 0); - - int max_key_pos = actual_N - 1; - if constexpr (IS_CAUSAL) { - max_key_pos = min(max_key_pos, - start_row + valid_q_rows - 1 + causal_q_offset); - } - if (window_size_right >= 0) { - max_key_pos = min(max_key_pos, - start_row + valid_q_rows - 1 + causal_q_offset - + 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, start_row + causal_q_offset - - 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 __half* q_ptr = Q + q_head_linear * M * D + start_row * 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; - - const int max_num_blocks_per_seq = (N + page_block_size - 1) / page_block_size; - const int* block_table_seq = block_table + batch_id * max_num_blocks_per_seq; - - const int64_t k_row_stride = k_token_stride; - const int64_t v_row_stride = v_token_stride; - - extern __shared__ __align__(128) 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; - int* sPageIdx = smem.page_idx; - int* sPageOffset = smem.page_offset; - - 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; - sRowSum[tid] = 0.0f; - } - constexpr int O_UINT4_PER_ROW = D / 4; - constexpr int O_UINT4_STRIDE = O_STRIDE / 4; - for (int idx = tid; idx < BLOCK_M * O_UINT4_PER_ROW; - idx += THREADS_PER_BLOCK) { - const int row = idx / O_UINT4_PER_ROW; - const int col = idx % O_UINT4_PER_ROW; - reinterpret_cast(sO)[row * O_UINT4_STRIDE + col] = - make_uint4(0, 0, 0, 0); - } - __syncthreads(); - - 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]); + sPageIdx[tid] = -1; + sPageOffset[tid] = 0; } - if constexpr (Config::PIPELINE_SWIZZLED_Q) { - const int k_tile = vec_col >> 1; - const int plane = vec_col & 1; - const int slot = pipeline_swizzled_q_row_slot(row); - sQ_vec[k_tile * (2 * BLOCK_M) + plane * BLOCK_M + slot] = - q_val; - } else { - sQ_vec[row * q_stride_uint4 + vec_col] = q_val; + } + __syncthreads(); + + if constexpr (LOW_SMEM_CONTIG_FAST) { + low_smem_contiguous_page_tile = + valid_k_rows == BLOCK_N && + k_block_stride == (int64_t)page_block_size * k_token_stride && + v_block_stride == (int64_t)page_block_size * v_token_stride && + sPageIdx[0] >= 0; +#pragma unroll + for (int i = 1; i < LOW_SMEM_PAGE_COUNT; ++i) { + const int linear_offset = sPageOffset[0] + i * LOW_SMEM_PAGE_SIZE; + const int expected_page_delta = linear_offset / page_block_size; + const int expected_page_offset = + linear_offset - expected_page_delta * page_block_size; + low_smem_contiguous_page_tile = + low_smem_contiguous_page_tile && + sPageIdx[i] == sPageIdx[0] + expected_page_delta && + sPageOffset[i] == expected_page_offset; } - } - __syncthreads(); + } else { + low_smem_contiguous_page_tile = false; + } - int cross_block_first_n_tile = 0; - if constexpr (D256_SW_PIPELINE_QK && D256_SW_PIPELINE_PV) { - const bool use_cross_block_pipeline = - valid_q_rows == BLOCK_M && actual_N > 0 - && page_block_size == 784 && bfla_block_mask == nullptr - && window_size_left < 0 && window_size_right < 0; - if (use_cross_block_pipeline) { - const __half* k_cache_h = - reinterpret_cast(K_cache); - const __half* v_cache_h = - reinterpret_cast(V_cache); - - if (tid < Config::LOW_SMEM_PAGE_COUNT) { - const int global_token_idx = tid * LOW_SMEM_PAGE_SIZE; - const int virtual_block_idx = - global_token_idx / page_block_size; - sPageIdx[tid] = - __ldg(&block_table_seq[virtual_block_idx]); - sPageOffset[tid] = - global_token_idx - - virtual_block_idx * page_block_size; - } - __syncthreads(); + const __half* K_cache_h = reinterpret_cast(K_cache); - if (warp_id < BLOCK_N / WMMA_N) { - const int valid_k_rows = min(BLOCK_N, actual_N); - pipeline_qk_slice( - sQ, sS, k_cache_h, sPageIdx, sPageOffset, - k_block_stride, k_token_stride, - k_head_stride, kv_head_id, 0, valid_k_rows, start_row, - valid_q_rows, causal_q_offset, warp_id * WMMA_N, 0, false, - true, softmax_scale, window_size_left, window_size_right, - NEG_INF); + if constexpr (LOW_SMEM_SCALAR_QK) { + uint4* sK_vec = reinterpret_cast(sK); + 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) { + const int page_slot = row >> 4; + const int block_offset = + sPageOffset[page_slot] + (row - page_slot * LOW_SMEM_PAGE_SIZE); + const int physical_block_idx_direct = sPageIdx[page_slot]; + const __half* k_row_ptr; + if constexpr (LOW_SMEM_CONTIG_FAST) { + k_row_ptr = + low_smem_contiguous_page_tile + ? K_cache_h + (int64_t)sPageIdx[0] * k_block_stride + + (int64_t)(sPageOffset[0] + row) * k_token_stride + + (int64_t)kv_head_id * k_head_stride + : K_cache_h + + (int64_t)physical_block_idx_direct * + k_block_stride + + (int64_t)block_offset * k_token_stride + + (int64_t)kv_head_id * k_head_stride; + } else { + k_row_ptr = K_cache_h + + (int64_t)physical_block_idx_direct * k_block_stride + + (int64_t)block_offset * k_token_stride + + (int64_t)kv_head_id * k_head_stride; } - __syncthreads(); - - // Keep the hot loop on a compile-time steady-state path. The - // existing exact path below drains the final block, avoiding a - // loop-carried has_next predicate in every softmax/PV panel. - const int steady_n_tiles = num_n_tiles - 1; - for (int block_n = 0; block_n < steady_n_tiles; ++block_n) { - const int start_col = block_n * BLOCK_N; - const int valid_k_rows = BLOCK_N; - const int next_start_col = start_col + BLOCK_N; - const int next_valid_k_rows = - min(BLOCK_N, actual_N - next_start_col); - const bool odd_block = (block_n & 1) != 0; - float* const score_current = - odd_block ? smem.pipeline_s : sS; - float* const score_next = - odd_block ? sS : smem.pipeline_s; - int* const page_idx_current = - odd_block ? smem.pipeline_page_idx : sPageIdx; - int* const page_offset_current = - odd_block ? smem.pipeline_page_offset : sPageOffset; - int* const page_idx_next = - odd_block ? sPageIdx : smem.pipeline_page_idx; - int* const page_offset_next = - odd_block ? sPageOffset : smem.pipeline_page_offset; - - if (tid < Config::LOW_SMEM_PAGE_COUNT) { - const int page_token_offset = - tid * LOW_SMEM_PAGE_SIZE; - const int global_token_idx = - next_start_col + page_token_offset; - const int virtual_block_idx = - global_token_idx / page_block_size; - page_idx_next[tid] = - __ldg(&block_table_seq[virtual_block_idx]); - page_offset_next[tid] = - global_token_idx - - virtual_block_idx * page_block_size; - } - __syncthreads(); - - for (int sub_start = 0; sub_start < valid_k_rows; - sub_start += P_SUB_TILE) { - const int sub_valid_k_rows = - min(P_SUB_TILE, valid_k_rows - sub_start); - const int row = warp_id; - const int thread_in_row = lane_id; - const unsigned mask = 0xFFFFFFFFU; - float* sS_row_f = - score_current + row * S_STRIDE + sub_start; - - 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) { - const int vc = - thread_in_row + j * THREADS_PER_ROW; - if (vc < vec_cols) { - const float4 v4 = sS_vec4[vc]; - thread_max = fmaxf( - thread_max, - fmaxf( - fmaxf(v4.x, v4.y), - fmaxf(v4.z, v4.w))); - } - } - #pragma unroll - 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)); - } - - const float row_max = - __shfl_sync(mask, thread_max, 0); - 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; - int vc_base = thread_in_row; - - #pragma unroll 4 - for (int j = 0; j < vecs_per_thread; - ++j, vc_base += THREADS_PER_ROW) { - if (vc_base < vec_cols) { - const float4 v4 = sS_vec4[vc_base]; - const float e0 = - __expf(fmaxf(v4.x - new_max, -80.0f)); - const float e1 = - __expf(fmaxf(v4.y - new_max, -80.0f)); - const float e2 = - __expf(fmaxf(v4.z - new_max, -80.0f)); - const float e3 = - __expf(fmaxf(v4.w - new_max, -80.0f)); - thread_sum += (e0 + e1) + (e2 + e3); - const int p_col = vc_base * 4; - __half2* p_half2 = reinterpret_cast<__half2*>( - sP + pipeline_swizzled_matrix_a_offset( - row, p_col)); - p_half2[0] = - __float22half2_rn(make_float2(e0, e1)); - p_half2[1] = - __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) { - const float e = __expf( - fmaxf(sS_row_f[c] - new_max, -80.0f)); - thread_sum += e; - sP[pipeline_swizzled_matrix_a_offset(row, c)] = - __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); - } - const float row_sum = - __shfl_sync(mask, thread_sum, 0); - if (thread_in_row == 0) { - sRowSum[row] = exp_diff * sRowSum[row] + row_sum; - sRowMax[row] = new_max; - } - - #pragma unroll 4 - for (int c = tail_start + thread_in_row; - c < P_SUB_TILE; c += THREADS_PER_ROW) { - if (c >= sub_valid_k_rows) { - sP[pipeline_swizzled_matrix_a_offset(row, c)] = - __float2half(0.0f); - } - } - - if (block_n > 0 || sub_start > 0) { - float4* sO_vec = reinterpret_cast( - sO + row * O_STRIDE); - #pragma unroll 4 - for (int ov = thread_in_row; ov < D / 4; - ov += THREADS_PER_ROW) { - float4 v = sO_vec[ov]; - v.x *= exp_diff; - v.y *= exp_diff; - v.z *= exp_diff; - v.w *= exp_diff; - sO_vec[ov] = v; - } - } - __syncthreads(); - - if (sub_start == 0) { - if (warp_id < BLOCK_N / WMMA_N) { - pipeline_qk_slice< - Q_STRIDE, S_STRIDE, D / WMMA_K, IS_CAUSAL, - true>( - sQ, score_next, k_cache_h, page_idx_next, - page_offset_next, k_block_stride, - k_token_stride, k_head_stride, kv_head_id, - next_start_col, next_valid_k_rows, start_row, - valid_q_rows, causal_q_offset, - warp_id * WMMA_N, 0, false, true, - softmax_scale, window_size_left, - window_size_right, NEG_INF); - } else { - const int pv_warp = - warp_id - BLOCK_N / WMMA_N; - pipeline_pv_tile( - sP, sO, v_cache_h, page_idx_current, - page_offset_current, v_block_stride, - v_token_stride, v_head_stride, kv_head_id, - sub_start, sub_valid_k_rows, - pv_warp * WMMA_N); - pipeline_pv_tile( - sP, sO, v_cache_h, page_idx_current, - page_offset_current, v_block_stride, - v_token_stride, v_head_stride, kv_head_id, - sub_start, sub_valid_k_rows, - (pv_warp + BLOCK_N / WMMA_N) * WMMA_N); - } - } else { - pipeline_pv_tile( - sP, sO, v_cache_h, page_idx_current, - page_offset_current, v_block_stride, - v_token_stride, v_head_stride, kv_head_id, - sub_start, sub_valid_k_rows, warp_id * WMMA_N); - } - __syncthreads(); - } + const uint4* k_row_vec = reinterpret_cast(k_row_ptr); + k_val = __ldg(&k_row_vec[vec_col]); + } + sK_vec[row * kv_stride_uint4 + vec_col] = k_val; + } + __syncthreads(); + 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; + } + if constexpr (ANCHORED_SWA) { + is_valid = + is_valid && ((global_n < anchor_len) || + (global_q_pos - global_n < anchored_window)); + } + + 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); } - cross_block_first_n_tile = steady_n_tiles; + } + sS[row * S_STRIDE + col] = is_valid ? acc * softmax_scale : NEG_INF; } - } + } else if constexpr (D256_SW_PIPELINE_QK) { + constexpr int QK_TILES = BLOCK_N / WMMA_N; + if (warp_id < QK_TILES) { + pipeline_qk_slice( + sQ, sS, K_cache_h, sPageIdx, sPageOffset, k_block_stride, + k_token_stride, k_head_stride, kv_head_id, start_col, + valid_k_rows, start_row, valid_q_rows, causal_q_offset, + warp_id * WMMA_N, 0, false, true, softmax_scale, window_size_left, + window_size_right, NEG_INF, anchor_len, anchored_window); + } + } 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; - int first_n_tile = cross_block_first_n_tile; - int last_n_tile = num_n_tiles; - if constexpr (SPLIT_KV) { - const int partition_id = blockIdx.y; - const int tiles_per_partition = - split_kv_tiles > 1 ? split_kv_tiles : 1; - first_n_tile = partition_id * tiles_per_partition; - last_n_tile = min(num_n_tiles, first_n_tile + tiles_per_partition); - } + 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; - for (int block_n = first_n_tile; block_n < last_n_tile; ++block_n) { - const int start_col = block_n * BLOCK_N; - if (start_col >= actual_N) break; - const int valid_k_rows = min(BLOCK_N, actual_N - start_col); + 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 (start_col + valid_k_rows <= min_key_pos) { + if (tile_m >= valid_q_rows || tile_n >= valid_k_rows) { continue; - } + } - if (bfla_block_mask != nullptr && bfla_mask_block_n > 0) { - const int mask_q_idx = start_row / bfla_mask_block_n; - const int mask_k_idx = start_col / bfla_mask_block_n; - const int keep_tile = __ldg( - bfla_block_mask - + (int64_t)batch_id * bfla_mask_stride_b - + (int64_t)kv_head_id * bfla_mask_stride_h - + (int64_t)mask_q_idx * bfla_mask_stride_q - + (int64_t)mask_k_idx * bfla_mask_stride_k); - if (keep_tile == 0) { - continue; - } - } + const int page_slot = tile_n >> 4; + const int block_offset = sPageOffset[page_slot]; + const int physical_block_idx_direct = sPageIdx[page_slot]; - const int partial_block_size = (block_n == num_n_tiles - 1 && actual_N % BLOCK_N != 0) - ? (actual_N % BLOCK_N) : -1; + fragment a_frag; + fragment b_frag; + fragment acc_frag; + fill_fragment(acc_frag, 0.0f); - uint4* sK_vec = reinterpret_cast(sK); - const int64_t row_stride_uint4 = k_row_stride / PER_UINT4; - const int start_page = start_col / page_block_size; - const int page_offset = start_col % page_block_size; - const bool single_page_tile = - (page_offset + valid_k_rows) <= page_block_size; - const bool two_page_tile = - !single_page_tile && - (page_offset + valid_k_rows) <= (page_block_size * 2); - const int first_page_rows = - single_page_tile ? valid_k_rows - : min(valid_k_rows, page_block_size - page_offset); - const int second_page_rows = valid_k_rows - first_page_rows; - const int physical_block_idx0 = __ldg(&block_table_seq[start_page]); - const int physical_block_idx1 = - second_page_rows > 0 ? __ldg(&block_table_seq[start_page + 1]) : -1; - const bool four_page_aligned_tile = - D == 256 && page_block_size == 16 && page_offset == 0 - && BLOCK_N == page_block_size * 4 - && valid_k_rows == BLOCK_N - && k_block_stride == (int64_t)page_block_size * k_token_stride - && v_block_stride == (int64_t)page_block_size * v_token_stride; - const int physical_block_idx2 = - four_page_aligned_tile ? __ldg(&block_table_seq[start_page + 2]) : -1; - const int physical_block_idx3 = - four_page_aligned_tile ? __ldg(&block_table_seq[start_page + 3]) : -1; - const bool four_page_contiguous_tile = - four_page_aligned_tile - && physical_block_idx1 == physical_block_idx0 + 1 - && physical_block_idx2 == physical_block_idx0 + 2 - && physical_block_idx3 == physical_block_idx0 + 3; - bool low_smem_contiguous_page_tile = false; - if constexpr (LOW_SMEM) { - static_assert(D == 256, "low-smem paged prefill is D=256 only"); - static_assert(KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP16, - "low-smem paged prefill is fp16-KV only"); - static_assert(BLOCK_N % LOW_SMEM_PAGE_SIZE == 0, - "low-smem BLOCK_N must be page aligned"); - - constexpr int LOW_SMEM_PAGE_COUNT = - Config::LOW_SMEM_PAGE_COUNT; - if (tid < LOW_SMEM_PAGE_COUNT) { - const int page_token_offset = tid * LOW_SMEM_PAGE_SIZE; - if (page_token_offset < valid_k_rows) { - const int global_token_idx = - start_col + page_token_offset; - const int virtual_block_idx = - global_token_idx / page_block_size; - sPageIdx[tid] = __ldg( - &block_table_seq[virtual_block_idx]); - sPageOffset[tid] = - global_token_idx - - virtual_block_idx * page_block_size; - } else { - sPageIdx[tid] = -1; - sPageOffset[tid] = 0; - } - } - __syncthreads(); +#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; + const __half* k_tile_ptr; if constexpr (LOW_SMEM_CONTIG_FAST) { - low_smem_contiguous_page_tile = - valid_k_rows == BLOCK_N - && k_block_stride - == (int64_t)page_block_size * k_token_stride - && v_block_stride - == (int64_t)page_block_size * v_token_stride - && sPageIdx[0] >= 0; - #pragma unroll - for (int i = 1; i < LOW_SMEM_PAGE_COUNT; ++i) { - const int linear_offset = - sPageOffset[0] + i * LOW_SMEM_PAGE_SIZE; - const int expected_page_delta = - linear_offset / page_block_size; - const int expected_page_offset = - linear_offset - - expected_page_delta * page_block_size; - low_smem_contiguous_page_tile = - low_smem_contiguous_page_tile - && sPageIdx[i] == sPageIdx[0] + expected_page_delta - && sPageOffset[i] == expected_page_offset; - } + k_tile_ptr = + low_smem_contiguous_page_tile + ? K_cache_h + (int64_t)sPageIdx[0] * k_block_stride + + (int64_t)(sPageOffset[0] + tile_n) * + k_token_stride + + (int64_t)kv_head_id * k_head_stride + k_offset + : K_cache_h + + (int64_t)physical_block_idx_direct * + k_block_stride + + (int64_t)block_offset * k_token_stride + + (int64_t)kv_head_id * k_head_stride + k_offset; } else { - low_smem_contiguous_page_tile = false; + k_tile_ptr = K_cache_h + + (int64_t)physical_block_idx_direct * k_block_stride + + (int64_t)block_offset * k_token_stride + + (int64_t)kv_head_id * k_head_stride + k_offset; } - const __half* K_cache_h = reinterpret_cast(K_cache); - - if constexpr (LOW_SMEM_SCALAR_QK) { - uint4* sK_vec = reinterpret_cast(sK); - 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) { - const int page_slot = row >> 4; - const int block_offset = - sPageOffset[page_slot] - + (row - page_slot * LOW_SMEM_PAGE_SIZE); - const int physical_block_idx_direct = - sPageIdx[page_slot]; - const __half* k_row_ptr; - if constexpr (LOW_SMEM_CONTIG_FAST) { - k_row_ptr = - low_smem_contiguous_page_tile - ? K_cache_h - + (int64_t)sPageIdx[0] * k_block_stride - + (int64_t)(sPageOffset[0] + row) - * k_token_stride - + (int64_t)kv_head_id * k_head_stride - : K_cache_h - + (int64_t)physical_block_idx_direct - * k_block_stride - + (int64_t)block_offset - * k_token_stride - + (int64_t)kv_head_id * k_head_stride; - } else { - k_row_ptr = - K_cache_h - + (int64_t)physical_block_idx_direct - * k_block_stride - + (int64_t)block_offset * k_token_stride - + (int64_t)kv_head_id * k_head_stride; - } - const uint4* k_row_vec = - reinterpret_cast(k_row_ptr); - k_val = __ldg(&k_row_vec[vec_col]); - } - sK_vec[row * kv_stride_uint4 + vec_col] = k_val; - } - __syncthreads(); - - 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 * softmax_scale : NEG_INF; - } - } else if constexpr (D256_SW_PIPELINE_QK) { - constexpr int QK_TILES = BLOCK_N / WMMA_N; - if (warp_id < QK_TILES) { - pipeline_qk_slice( - sQ, sS, K_cache_h, sPageIdx, sPageOffset, - k_block_stride, k_token_stride, k_head_stride, - kv_head_id, start_col, valid_k_rows, start_row, - valid_q_rows, causal_q_offset, warp_id * WMMA_N, 0, - false, true, softmax_scale, window_size_left, - window_size_right, NEG_INF); - } - } 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; - } - - const int page_slot = tile_n >> 4; - const int block_offset = sPageOffset[page_slot]; - const int physical_block_idx_direct = sPageIdx[page_slot]; - - 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; - - const __half* k_tile_ptr; - if constexpr (LOW_SMEM_CONTIG_FAST) { - k_tile_ptr = - low_smem_contiguous_page_tile - ? K_cache_h - + (int64_t)sPageIdx[0] * k_block_stride - + (int64_t)(sPageOffset[0] + tile_n) - * k_token_stride - + (int64_t)kv_head_id * k_head_stride - + k_offset - : K_cache_h - + (int64_t)physical_block_idx_direct - * k_block_stride - + (int64_t)block_offset - * k_token_stride - + (int64_t)kv_head_id * k_head_stride - + k_offset; - } else { - k_tile_ptr = - K_cache_h - + (int64_t)physical_block_idx_direct - * k_block_stride - + (int64_t)block_offset * k_token_stride - + (int64_t)kv_head_id * k_head_stride - + k_offset; - } - - load_matrix_sync( - a_frag, - sQ + tile_m * Q_STRIDE + k_offset, - Q_STRIDE); - load_matrix_sync(b_frag, k_tile_ptr, k_token_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; - } - - acc_frag.x[i] = - (is_valid && is_causal_valid && is_window_valid) - ? acc_frag.x[i] * softmax_scale - : NEG_INF; - } + load_matrix_sync(a_frag, sQ + tile_m * Q_STRIDE + k_offset, + Q_STRIDE); + load_matrix_sync(b_frag, k_tile_ptr, k_token_stride); + mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } - store_matrix_sync( - sS + tile_m * S_STRIDE + tile_n, acc_frag, S_STRIDE, - mem_row_major); - } +#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; } - __syncthreads(); - } else { - if constexpr (KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP16) { - const __half* K_cache_h = reinterpret_cast(K_cache); - const uint4* k_page0_vec = reinterpret_cast( - K_cache_h + (int64_t)physical_block_idx0 * k_block_stride - + (int64_t)page_offset * k_token_stride - + (int64_t)kv_head_id * k_head_stride); - const uint4* k_page1_vec = - two_page_tile && second_page_rows > 0 - ? reinterpret_cast( - K_cache_h + (int64_t)physical_block_idx1 * k_block_stride - + (int64_t)kv_head_id * k_head_stride) - : nullptr; - - #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) { - if (four_page_contiguous_tile) { - k_val = __ldg( - &k_page0_vec[row * row_stride_uint4 + vec_col]); - } else if (single_page_tile) { - k_val = __ldg(&k_page0_vec[row * row_stride_uint4 + vec_col]); - } else if (two_page_tile) { - if (row < first_page_rows) { - k_val = __ldg( - &k_page0_vec[row * row_stride_uint4 + vec_col]); - } else { - const int row_page1 = row - first_page_rows; - k_val = __ldg( - &k_page1_vec[row_page1 * row_stride_uint4 + vec_col]); - } - } else { - const uint4* k_vec = reinterpret_cast(K_cache_h); - const int global_token_idx = start_col + row; - const int virtual_block_idx = - global_token_idx / page_block_size; - const int block_offset = global_token_idx % page_block_size; - const int physical_block_idx_slow = - __ldg(&block_table_seq[virtual_block_idx]); - const int64_t physical_offset_halfs = - (int64_t)physical_block_idx_slow * k_block_stride - + (int64_t)block_offset * k_token_stride - + (int64_t)kv_head_id * k_head_stride; - const int64_t physical_offset_uint4 = - (physical_offset_halfs / PER_UINT4) + vec_col; - k_val = __ldg(&k_vec[physical_offset_uint4]); - } - } - sK_vec[row * kv_stride_uint4 + vec_col] = k_val; + bool is_window_valid = true; + if (window_size_left >= 0) { + is_window_valid = is_window_valid && + global_n >= global_q_pos - window_size_left; } - } else { - for (int idx = tid; idx < valid_k_rows * D; idx += THREADS_PER_BLOCK) { - const int row = idx / D; - const int col = idx % D; - const int global_token_idx = start_col + row; - const int virtual_block_idx = global_token_idx / page_block_size; - const int block_offset = global_token_idx % page_block_size; - const int physical_block_idx_slow = - __ldg(&block_table_seq[virtual_block_idx]); - const int64_t physical_offset = - (int64_t)physical_block_idx_slow * k_block_stride - + (int64_t)block_offset * k_token_stride - + (int64_t)kv_head_id * k_head_stride + col; - sK[row * KV_STRIDE + col] = - flash_v100::load_kv_cache_half( - K_cache, physical_offset, k_scale); + if (window_size_right >= 0) { + is_window_valid = is_window_valid && + global_n <= global_q_pos + window_size_right; } - } - __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); + bool is_anchor_valid = true; + if constexpr (ANCHORED_SWA) { + is_anchor_valid = (global_n < anchor_len) || + (global_q_pos - global_n < anchored_window); + } - #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; + acc_frag.x[i] = (is_valid && is_causal_valid && is_window_valid && + is_anchor_valid) + ? acc_frag.x[i] * softmax_scale + : NEG_INF; + } - 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); + store_matrix_sync(sS + tile_m * S_STRIDE + tile_n, acc_frag, S_STRIDE, + mem_row_major); + } + } + __syncthreads(); + } else { + if constexpr (KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP16) { + const __half* K_cache_h = reinterpret_cast(K_cache); + const uint4* k_page0_vec = reinterpret_cast( + K_cache_h + (int64_t)physical_block_idx0 * k_block_stride + + (int64_t)page_offset * k_token_stride + + (int64_t)kv_head_id * k_head_stride); + const uint4* k_page1_vec = + two_page_tile && second_page_rows > 0 + ? reinterpret_cast( + K_cache_h + + (int64_t)physical_block_idx1 * k_block_stride + + (int64_t)kv_head_id * k_head_stride) + : nullptr; + +#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) { + if (four_page_contiguous_tile) { + k_val = __ldg(&k_page0_vec[row * row_stride_uint4 + vec_col]); + } else if (single_page_tile) { + k_val = __ldg(&k_page0_vec[row * row_stride_uint4 + vec_col]); + } else if (two_page_tile) { + if (row < first_page_rows) { + k_val = __ldg(&k_page0_vec[row * row_stride_uint4 + vec_col]); + } else { + const int row_page1 = row - first_page_rows; + k_val = + __ldg(&k_page1_vec[row_page1 * row_stride_uint4 + vec_col]); + } + } else { + const uint4* k_vec = reinterpret_cast(K_cache_h); + const int global_token_idx = start_col + row; + const int virtual_block_idx = global_token_idx / page_block_size; + const int block_offset = global_token_idx % page_block_size; + const int physical_block_idx_slow = + __ldg(&block_table_seq[virtual_block_idx]); + const int64_t physical_offset_halves = + (int64_t)physical_block_idx_slow * k_block_stride + + (int64_t)block_offset * k_token_stride + + (int64_t)kv_head_id * k_head_stride; + const int64_t physical_offset_uint4 = + (physical_offset_halves / PER_UINT4) + vec_col; + k_val = __ldg(&k_vec[physical_offset_uint4]); } + } + sK_vec[row * kv_stride_uint4 + vec_col] = k_val; + } + } else { + for (int idx = tid; idx < valid_k_rows * D; idx += THREADS_PER_BLOCK) { + const int row = idx / D; + const int col = idx % D; + const int global_token_idx = start_col + row; + const int virtual_block_idx = global_token_idx / page_block_size; + const int block_offset = global_token_idx % page_block_size; + const int physical_block_idx_slow = + __ldg(&block_table_seq[virtual_block_idx]); + const int64_t physical_offset = + (int64_t)physical_block_idx_slow * k_block_stride + + (int64_t)block_offset * k_token_stride + + (int64_t)kv_head_id * k_head_stride + col; + sK[row * KV_STRIDE + col] = flash_v100::load_kv_cache_half( + K_cache, physical_offset, k_scale); + } + } + __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 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; +#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; - 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; + 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); + } - 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; - } +#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; + } + + bool is_anchor_valid = true; + if constexpr (ANCHORED_SWA) { + is_anchor_valid = (global_n < anchor_len) || + (global_q_pos - global_n < anchored_window); + } + + acc_frag.x[i] = (is_valid && is_causal_valid && is_window_valid && + is_anchor_valid) + ? acc_frag.x[i] * softmax_scale + : NEG_INF; + } - acc_frag.x[i] = (is_valid && is_causal_valid && is_window_valid) - ? acc_frag.x[i] * softmax_scale - : NEG_INF; + store_matrix_sync(sS + tile_m * S_STRIDE + tile_n, acc_frag, S_STRIDE, + mem_row_major); + } + __syncthreads(); + } + + uint4* sV_vec = reinterpret_cast(sV); + const int64_t v_row_stride_uint4 = v_row_stride / PER_UINT4; + if constexpr (!LOW_SMEM) { + if constexpr (KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP16) { + const __half* V_cache_h = reinterpret_cast(V_cache); + const uint4* v_page0_vec = reinterpret_cast( + V_cache_h + (int64_t)physical_block_idx0 * v_block_stride + + (int64_t)page_offset * v_token_stride + + (int64_t)kv_head_id * v_head_stride); + const uint4* v_page1_vec = + two_page_tile && second_page_rows > 0 + ? reinterpret_cast( + V_cache_h + + (int64_t)physical_block_idx1 * v_block_stride + + (int64_t)kv_head_id * v_head_stride) + : nullptr; + +#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) { + if (four_page_contiguous_tile) { + v_val = __ldg(&v_page0_vec[row * v_row_stride_uint4 + vec_col]); + } else if (single_page_tile) { + v_val = __ldg(&v_page0_vec[row * v_row_stride_uint4 + vec_col]); + } else if (two_page_tile) { + if (row < first_page_rows) { + v_val = __ldg(&v_page0_vec[row * v_row_stride_uint4 + vec_col]); + } else { + const int row_page1 = row - first_page_rows; + v_val = __ldg( + &v_page1_vec[row_page1 * v_row_stride_uint4 + vec_col]); + } + } else { + const uint4* v_vec = reinterpret_cast(V_cache_h); + const int global_token_idx = start_col + row; + const int virtual_block_idx = global_token_idx / page_block_size; + const int block_offset = global_token_idx % page_block_size; + const int physical_block_idx_slow = + __ldg(&block_table_seq[virtual_block_idx]); + const int64_t physical_offset_halves = + (int64_t)physical_block_idx_slow * v_block_stride + + (int64_t)block_offset * v_token_stride + + (int64_t)kv_head_id * v_head_stride; + const int64_t physical_offset_uint4 = + (physical_offset_halves / PER_UINT4) + vec_col; + v_val = __ldg(&v_vec[physical_offset_uint4]); } - - store_matrix_sync(sS + tile_m * S_STRIDE + tile_n, acc_frag, S_STRIDE, mem_row_major); + } + sV_vec[row * kv_stride_uint4 + vec_col] = v_val; } - __syncthreads(); + } else { + for (int idx = tid; idx < valid_k_rows * D; idx += THREADS_PER_BLOCK) { + const int row = idx / D; + const int col = idx % D; + const int global_token_idx = start_col + row; + const int virtual_block_idx = global_token_idx / page_block_size; + const int block_offset = global_token_idx % page_block_size; + const int physical_block_idx_slow = + __ldg(&block_table_seq[virtual_block_idx]); + const int64_t physical_offset = + (int64_t)physical_block_idx_slow * v_block_stride + + (int64_t)block_offset * v_token_stride + + (int64_t)kv_head_id * v_head_stride + col; + sV[row * KV_STRIDE + col] = flash_v100::load_kv_cache_half( + V_cache, physical_offset, v_scale); } + } + __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); + + // ================================================================ + // Patch 0108 (surgical) — sS/sP union race fix, D-conditional. + // ISSUE-0007 / issues/0090. The race (W(__half2) P-commit racing + // R(float4) sS reads with no cross-warp barrier) exists ONLY for + // D != 256, where sP = reuse_sp.p aliases sS byte-for-byte. + // + // D == 256 is race-IMMUNE: sP = p_strict, a SEPARATE buffer + // (SmemLayout::p_strict, line ~134; selected at line ~491). So the + // D==256 arm below is the ORIGINAL monolithic block VERBATIM — no + // barrier, no variable hoisting. This is the ONLY difference from + // archived/patches/0011: 0011 split the guard and hoisted vars for + // EVERY D, which perturbed the D=256 register/spill allocation and + // cost the immune path +0.47..+2.79% for zero benefit. Guarding the + // whole transform under `if constexpr (D != 256)` makes the D=256 + // codegen byte-identical to baseline (proved via SASS diff) => the + // immune path pays EXACTLY zero. The D!=256 arm is 0011's + // two-phase-commit fix, verbatim (barrier ALONE, no riders). + // ================================================================ + if constexpr (D == 256) { + 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 + 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 = __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); + __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) { + 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); + + const __half2 p01 = __float22half2_rn(make_float2(e0, e1)); + const __half2 p23 = __float22half2_rn(make_float2(e2, e3)); + if constexpr (D256_SW_PIPELINE_PV) { + // D256 owns a separate probability buffer, so the + // values can leave registers before row reduction. + const int base_offset = vc_base * 2; + sP_half2[base_offset] = p01; + sP_half2[base_offset + 1] = p23; + } else { + half_buffer[h2_idx++] = p01; + half_buffer[h2_idx++] = p23; + } + } + } + +#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 = __expf(fmaxf(v - new_max, -80.0f)); + thread_sum += e; + if constexpr (D256_SW_PIPELINE_PV) { + sP_row_h[c] = __float2half_rn(e); + } else { + 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; + } + + if constexpr (!D256_SW_PIPELINE_PV) { + h2_idx = 0; + vc_base = thread_in_row; +#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++]; + } + } - uint4* sV_vec = reinterpret_cast(sV); - const int64_t v_row_stride_uint4 = v_row_stride / PER_UINT4; - if constexpr (!LOW_SMEM) { - if constexpr (KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP16) { - const __half* V_cache_h = reinterpret_cast(V_cache); - const uint4* v_page0_vec = reinterpret_cast( - V_cache_h + (int64_t)physical_block_idx0 * v_block_stride - + (int64_t)page_offset * v_token_stride - + (int64_t)kv_head_id * v_head_stride); - const uint4* v_page1_vec = - two_page_tile && second_page_rows > 0 - ? reinterpret_cast( - V_cache_h + (int64_t)physical_block_idx1 * v_block_stride - + (int64_t)kv_head_id * v_head_stride) - : nullptr; - - #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) { - if (four_page_contiguous_tile) { - v_val = __ldg( - &v_page0_vec[row * v_row_stride_uint4 + vec_col]); - } else if (single_page_tile) { - v_val = __ldg(&v_page0_vec[row * v_row_stride_uint4 + vec_col]); - } else if (two_page_tile) { - if (row < first_page_rows) { - v_val = __ldg( - &v_page0_vec[row * v_row_stride_uint4 + vec_col]); - } else { - const int row_page1 = row - first_page_rows; - v_val = __ldg( - &v_page1_vec[row_page1 * v_row_stride_uint4 + vec_col]); - } - } else { - const uint4* v_vec = reinterpret_cast(V_cache_h); - const int global_token_idx = start_col + row; - const int virtual_block_idx = - global_token_idx / page_block_size; - const int block_offset = global_token_idx % page_block_size; - const int physical_block_idx_slow = - __ldg(&block_table_seq[virtual_block_idx]); - const int64_t physical_offset_halfs = - (int64_t)physical_block_idx_slow * v_block_stride - + (int64_t)block_offset * v_token_stride - + (int64_t)kv_head_id * v_head_stride; - const int64_t physical_offset_uint4 = - (physical_offset_halfs / PER_UINT4) + vec_col; - v_val = __ldg(&v_vec[physical_offset_uint4]); - } - } - sV_vec[row * kv_stride_uint4 + vec_col] = v_val; + if (tail_col >= 0) { + sP_row_h[tail_col] = tail_value; } - } else { - for (int idx = tid; idx < valid_k_rows * D; idx += THREADS_PER_BLOCK) { - const int row = idx / D; - const int col = idx % D; - const int global_token_idx = start_col + row; - const int virtual_block_idx = global_token_idx / page_block_size; - const int block_offset = global_token_idx % page_block_size; - const int physical_block_idx_slow = - __ldg(&block_table_seq[virtual_block_idx]); - const int64_t physical_offset = - (int64_t)physical_block_idx_slow * v_block_stride - + (int64_t)block_offset * v_token_stride - + (int64_t)kv_head_id * v_head_stride + col; - sV[row * KV_STRIDE + col] = - flash_v100::load_kv_cache_half( - V_cache, physical_offset, v_scale); + } + +#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 = D / 4; + 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(); + } else { + // ---------------------------------------------------------------- + // ISSUE-0007 race-class fix (patch 0011, + // patches/0011-paged-race-fix-issue0007/ — PoC: poc/race_poc.cu). + // + // For D != 256, sP aliases sS byte-for-byte (union reuse_sp; + // sP = reuse_sp.p selected above with p_stride == S_STRIDE). The + // softmax phase reads sS as float4 (max pass, exp pass, scalar + // tail) and then commits P into the SAME union bytes as + // __half2/__half. Row r's half P row occupies union bytes + // [2*S_STRIDE*r, +2*S_STRIDE) — inside the FLOAT score row of row + // r/2 — so the P commit of warp w lands exactly on the score row + // still being read as float4 by warp w/2. The register staging + // (half_buffer) orders read-then-write only WITHIN one thread; + // cross-warp there was NO ordering until the __syncthreads() at + // the end of this phase, i.e. AFTER both: a W(__half2) x R(float4) + // data race. Evidence: compute-sanitizer racecheck 3 errors / + // 87,712 hazard instances on the PoC extraction of this exact + // access pattern, and the REAL kernel 10/10-runs bitwise + // non-deterministic at D=128 (q=512/kv=4096 and q=1024/kv=8192). + // + // Fix, proved racecheck-clean AND bit-equal in the PoC: two-phase + // P commit. Phase A = all sS reads + reductions with P staged in + // registers exactly as before; ONE uniform block-scope + // __syncthreads(); phase B = all sP stores. The former single + // guarded region is split in two so the barrier sits in UNIFORM + // control flow, and the per-thread state below is hoisted to + // survive across the barrier. No arithmetic is changed or + // reordered. D == 256 keeps its original barrier-free structure: + // sP = p_strict there, a separate buffer that never aliases sS + // (constexpr guard on the barrier below). + // ---------------------------------------------------------------- + const int row = tid / THREADS_PER_ROW; + const int thread_in_row = tid % THREADS_PER_ROW; + __half* sP_row_h = sP + row * p_stride; + __half2* sP_half2 = reinterpret_cast<__half2*>(sP_row_h); + 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 exp_diff = 1.0f; + __half2 half_buffer[20]; + int tail_col = -1; + __half tail_value = __float2half(0.f); + + if (tid < valid_q_rows * 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; + + 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 + 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); + // 0011: exp_diff hoisted above the phase split (used by the + // sO rescale in phase B). Same expression, same position in + // the arithmetic order. + exp_diff = __expf(old_max - new_max); + + float thread_sum = 0.0f; + int vc_base = thread_in_row; + int h2_idx = 0; + +#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); + + const __half2 p01 = __float22half2_rn(make_float2(e0, e1)); + const __half2 p23 = __float22half2_rn(make_float2(e2, e3)); + if constexpr (D256_SW_PIPELINE_PV) { + // D256 owns a separate probability buffer, so the + // values can leave registers before row reduction. + const int base_offset = vc_base * 2; + sP_half2[base_offset] = p01; + sP_half2[base_offset + 1] = p23; + } else { + half_buffer[h2_idx++] = p01; + half_buffer[h2_idx++] = p23; + } + } + } + +#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 = __expf(fmaxf(v - new_max, -80.0f)); + thread_sum += e; + if constexpr (D256_SW_PIPELINE_PV) { + sP_row_h[c] = __float2half_rn(e); + } else { + 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; + } + } // end phase A: every sS read of this sub-tile is done; P is + // staged in registers (half_buffer / tail_col / tail_value). + + // ISSUE-0007 fix (patch 0011): THE one uniform __syncthreads() + // between the sS read phase and the sP commit phase — the only + // behavioral change of this patch. UNIFORMITY: this point is at + // the body scope of the sub_start loop; its trip count + // (valid_k_rows, softmax_sub_tile) and every enclosing control + // transfer (kernel early returns on blockIdx.z/M; block_n loop + // break/continue on start_col/min_key_pos/bfla block mask) are + // block-uniform — the same uniformity the pre-existing + // __syncthreads() at the end of this loop body already requires. + // All threads of the block reach this barrier exactly once per + // sub_start iteration. D == 256 is immune (sP = p_strict, + // separate buffer) and keeps its barrier-free codegen. + if constexpr (D != 256) { + __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); - - // ================================================================ - // Patch 0108 (surgical) — sS/sP union race fix, D-conditional. - // ISSUE-0007 / issues/0090. The race (W(__half2) P-commit racing - // R(float4) sS reads with no cross-warp barrier) exists ONLY for - // D != 256, where sP = reuse_sp.p aliases sS byte-for-byte. - // - // D == 256 is race-IMMUNE: sP = p_strict, a SEPARATE buffer - // (SmemLayout::p_strict, line ~134; selected at line ~491). So the - // D==256 arm below is the ORIGINAL monolithic block VERBATIM — no - // barrier, no variable hoisting. This is the ONLY difference from - // archived/patches/0011: 0011 split the guard and hoisted vars for - // EVERY D, which perturbed the D=256 register/spill allocation and - // cost the immune path +0.47..+2.79% for zero benefit. Guarding the - // whole transform under `if constexpr (D != 256)` makes the D=256 - // codegen byte-identical to baseline (proved via SASS diff) => the - // immune path pays EXACTLY zero. The D!=256 arm is 0011's - // two-phase-commit fix, verbatim (barrier ALONE, no riders). - // ================================================================ - if constexpr (D == 256) { - 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 - 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 = __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); - __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) { - 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); - - const __half2 p01 = - __float22half2_rn(make_float2(e0, e1)); - const __half2 p23 = - __float22half2_rn(make_float2(e2, e3)); - if constexpr (D256_SW_PIPELINE_PV) { - // D256 owns a separate probability buffer, so the - // values can leave registers before row reduction. - const int base_offset = vc_base * 2; - sP_half2[base_offset] = p01; - sP_half2[base_offset + 1] = p23; - } else { - half_buffer[h2_idx++] = p01; - half_buffer[h2_idx++] = p23; - } - } - } - - #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 = __expf(fmaxf(v - new_max, -80.0f)); - thread_sum += e; - if constexpr (D256_SW_PIPELINE_PV) { - sP_row_h[c] = __float2half_rn(e); - } else { - 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; - } - - if constexpr (!D256_SW_PIPELINE_PV) { - h2_idx = 0; - vc_base = thread_in_row; - #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; - } - } + if (tid < valid_q_rows * THREADS_PER_ROW) { + // phase B: commit P to the union bytes. Bit-identical values + // (straight from the phase-A registers), same store order, + // same addresses as the pre-fix code. + if constexpr (!D256_SW_PIPELINE_PV) { + int h2_idx = 0; + int vc_base = thread_in_row; +#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++]; + } + } - #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 (tail_col >= 0) { + sP_row_h[tail_col] = tail_value; + } + } - 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 = D / 4; - 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 < p_tile_capacity; + c += THREADS_PER_ROW) { + if (c >= sub_valid_k_rows) { + sP_row_h[c] = __float2half(0.f); } - } else { - // ---------------------------------------------------------------- - // ISSUE-0007 race-class fix (patch 0011, - // patches/0011-paged-race-fix-issue0007/ — PoC: poc/race_poc.cu). - // - // For D != 256, sP aliases sS byte-for-byte (union reuse_sp; - // sP = reuse_sp.p selected above with p_stride == S_STRIDE). The - // softmax phase reads sS as float4 (max pass, exp pass, scalar - // tail) and then commits P into the SAME union bytes as - // __half2/__half. Row r's half P row occupies union bytes - // [2*S_STRIDE*r, +2*S_STRIDE) — inside the FLOAT score row of row - // r/2 — so the P commit of warp w lands exactly on the score row - // still being read as float4 by warp w/2. The register staging - // (half_buffer) orders read-then-write only WITHIN one thread; - // cross-warp there was NO ordering until the __syncthreads() at - // the end of this phase, i.e. AFTER both: a W(__half2) x R(float4) - // data race. Evidence: compute-sanitizer racecheck 3 errors / - // 87,712 hazard instances on the PoC extraction of this exact - // access pattern, and the REAL kernel 10/10-runs bitwise - // non-deterministic at D=128 (q=512/kv=4096 and q=1024/kv=8192). - // - // Fix, proved racecheck-clean AND bit-equal in the PoC: two-phase - // P commit. Phase A = all sS reads + reductions with P staged in - // registers exactly as before; ONE uniform block-scope - // __syncthreads(); phase B = all sP stores. The former single - // guarded region is split in two so the barrier sits in UNIFORM - // control flow, and the per-thread state below is hoisted to - // survive across the barrier. No arithmetic is changed or - // reordered. D == 256 keeps its original barrier-free structure: - // sP = p_strict there, a separate buffer that never aliases sS - // (constexpr guard on the barrier below). - // ---------------------------------------------------------------- - const int row = tid / THREADS_PER_ROW; - const int thread_in_row = tid % THREADS_PER_ROW; - __half* sP_row_h = sP + row * p_stride; - __half2* sP_half2 = reinterpret_cast<__half2*>(sP_row_h); - 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 exp_diff = 1.0f; - __half2 half_buffer[20]; - int tail_col = -1; - __half tail_value = __float2half(0.f); - - if (tid < valid_q_rows * 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; - - 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))); - } - } + } + + 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 = D / 4; + 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(); - #pragma unroll - 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]); - } + const int num_tiles_k_pv = (p_tile_capacity + WMMA_K - 1) / WMMA_K; - #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)); - } + 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 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); - // 0011: exp_diff hoisted above the phase split (used by the - // sO rescale in phase B). Same expression, same position in - // the arithmetic order. - exp_diff = __expf(old_max - new_max); - - float thread_sum = 0.0f; - int vc_base = thread_in_row; - int h2_idx = 0; - - #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); - - const __half2 p01 = - __float22half2_rn(make_float2(e0, e1)); - const __half2 p23 = - __float22half2_rn(make_float2(e2, e3)); - if constexpr (D256_SW_PIPELINE_PV) { - // D256 owns a separate probability buffer, so the - // values can leave registers before row reduction. - const int base_offset = vc_base * 2; - sP_half2[base_offset] = p01; - sP_half2[base_offset + 1] = p23; - } else { - half_buffer[h2_idx++] = p01; - half_buffer[h2_idx++] = p23; - } - } - } + const int tile_m_idx = global_tile_idx / num_tiles_n_pv; + const int tile_d_idx = global_tile_idx % num_tiles_n_pv; - #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 = __expf(fmaxf(v - new_max, -80.0f)); - thread_sum += e; - if constexpr (D256_SW_PIPELINE_PV) { - sP_row_h[c] = __float2half_rn(e); - } else { - tail_col = c; - tail_value = __float2half_rn(e); - } - } + const int tile_m = tile_m_idx * WMMA_M; + const int tile_d = tile_d_idx * WMMA_N; - #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); - } + if (tile_m >= valid_q_rows) continue; - float row_sum = - __shfl_sync(mask, thread_sum, row_leader, THREADS_PER_ROW); + fragment a_frag; + fragment b_frag; + fragment acc_frag; - if (thread_in_row == 0) { - sRowSum[row] = exp_diff * sRowSum[row] + row_sum; - sRowMax[row] = new_max; - } - } // end phase A: every sS read of this sub-tile is done; P is - // staged in registers (half_buffer / tail_col / tail_value). - - // ISSUE-0007 fix (patch 0011): THE one uniform __syncthreads() - // between the sS read phase and the sP commit phase — the only - // behavioral change of this patch. UNIFORMITY: this point is at - // the body scope of the sub_start loop; its trip count - // (valid_k_rows, softmax_sub_tile) and every enclosing control - // transfer (kernel early returns on blockIdx.z/M; block_n loop - // break/continue on start_col/min_key_pos/bfla block mask) are - // block-uniform — the same uniformity the pre-existing - // __syncthreads() at the end of this loop body already requires. - // All threads of the block reach this barrier exactly once per - // sub_start iteration. D == 256 is immune (sP = p_strict, - // separate buffer) and keeps its barrier-free codegen. - if constexpr (D != 256) { - __syncthreads(); - } + load_matrix_sync(acc_frag, sO + tile_m * O_STRIDE + tile_d, O_STRIDE, + mem_row_major); - if (tid < valid_q_rows * THREADS_PER_ROW) { - // phase B: commit P to the union bytes. Bit-identical values - // (straight from the phase-A registers), same store order, - // same addresses as the pre-fix code. - if constexpr (!D256_SW_PIPELINE_PV) { - int h2_idx = 0; - int vc_base = thread_in_row; - #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 + 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; - #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 = D / 4; - 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(); - - 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); - if constexpr (LOW_SMEM) { - const __half* V_cache_h = - reinterpret_cast(V_cache); - const int token_offset = sub_start + k_offset; - const int page_slot = token_offset >> 4; - const int block_offset = sPageOffset[page_slot]; - const int physical_block_idx_direct = - sPageIdx[page_slot]; - const __half* v_tile_ptr; - if constexpr (LOW_SMEM_CONTIG_FAST) { - v_tile_ptr = - low_smem_contiguous_page_tile - ? V_cache_h - + (int64_t)sPageIdx[0] * v_block_stride - + (int64_t)(sPageOffset[0] + token_offset) - * v_token_stride - + (int64_t)kv_head_id * v_head_stride - + tile_d - : V_cache_h - + (int64_t)physical_block_idx_direct - * v_block_stride - + (int64_t)block_offset * v_token_stride - + (int64_t)kv_head_id * v_head_stride - + tile_d; - } else { - v_tile_ptr = - V_cache_h - + (int64_t)physical_block_idx_direct - * v_block_stride - + (int64_t)block_offset * v_token_stride - + (int64_t)kv_head_id * v_head_stride + tile_d; - } - load_matrix_sync(b_frag, v_tile_ptr, v_token_stride); - } else { - 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); + load_matrix_sync(a_frag, sP + tile_m * p_stride + k_offset, p_stride); + if constexpr (LOW_SMEM) { + const __half* V_cache_h = reinterpret_cast(V_cache); + const int token_offset = sub_start + k_offset; + const int page_slot = token_offset >> 4; + const int block_offset = sPageOffset[page_slot]; + const int physical_block_idx_direct = sPageIdx[page_slot]; + const __half* v_tile_ptr; + if constexpr (LOW_SMEM_CONTIG_FAST) { + v_tile_ptr = + low_smem_contiguous_page_tile + ? V_cache_h + (int64_t)sPageIdx[0] * v_block_stride + + (int64_t)(sPageOffset[0] + token_offset) * + v_token_stride + + (int64_t)kv_head_id * v_head_stride + tile_d + : V_cache_h + + (int64_t)physical_block_idx_direct * + v_block_stride + + (int64_t)block_offset * v_token_stride + + (int64_t)kv_head_id * v_head_stride + tile_d; + } else { + v_tile_ptr = V_cache_h + + (int64_t)physical_block_idx_direct * v_block_stride + + (int64_t)block_offset * v_token_stride + + (int64_t)kv_head_id * v_head_stride + tile_d; } - __syncthreads(); + load_matrix_sync(b_frag, v_tile_ptr, v_token_stride); + } else { + 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(); } + } - if constexpr (SPLIT_KV) { - const int partition_id = blockIdx.y; - const int num_partitions = gridDim.y; - const int64_t row_base = - ((int64_t)batch_head_id * num_partitions + partition_id) * M - + start_row; - - for (int i = tid; i < valid_q_rows * D; i += THREADS_PER_BLOCK) { - const int row = i / D; - const int col = i - row * D; - split_tmp_out[(row_base + row) * D + col] = - sO[row * O_STRIDE + col]; - } - - if (tid < valid_q_rows) { - split_tmp_row_max[row_base + tid] = sRowMax[tid]; - split_tmp_row_sum[row_base + tid] = sRowSum[tid]; - } - return; - } + if constexpr (SPLIT_KV) { + const int partition_id = blockIdx.y; + const int num_partitions = gridDim.y; + const int64_t row_base = + ((int64_t)batch_head_id * num_partitions + partition_id) * M + + start_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" - ); + for (int i = tid; i < valid_q_rows * D; i += THREADS_PER_BLOCK) { + const int row = i / D; + const int col = i - row * D; + split_tmp_out[(row_base + row) * D + col] = sO[row * O_STRIDE + col]; } if (tid < valid_q_rows) { - const float sum = fmaxf(sRowSum[tid], 1e-24f); - softmax_lse_ptr[tid] = sRowMax[tid] + logf(sum); - } + split_tmp_row_max[row_base + tid] = sRowMax[tid]; + split_tmp_row_sum[row_base + tid] = sRowSum[tid]; + } + return; + } + + 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); + } } 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_default_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; } -template +template void launcher_flash_attention_forward_paged_impl( - const torch::Tensor& Q, - const torch::Tensor& K_cache, - const torch::Tensor& V_cache, - torch::Tensor& Out, - torch::Tensor& softmax_lse, - const torch::Tensor& block_table, - const torch::Tensor& seq_lens, - const int* bfla_mask_ptr, - int bfla_mask_block_n, - int64_t bfla_mask_stride_b, - int64_t bfla_mask_stride_h, - int64_t bfla_mask_stride_q, - int64_t bfla_mask_stride_k, - float softmax_scale, - bool is_causal, - float k_scale, - float v_scale, - 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 M = Q.size(2); - const int page_block_size = K_cache.size(1); - const int num_kv_heads = K_cache.size(2); - const int max_num_blocks = block_table.size(1); - const int N = max_num_blocks * page_block_size; - const int64_t k_block_stride = K_cache.stride(0); - const int64_t k_token_stride = K_cache.stride(1); - const int64_t k_head_stride = K_cache.stride(2); - const int64_t v_block_stride = V_cache.stride(0); - const int64_t v_token_stride = V_cache.stride(1); - const int64_t v_head_stride = V_cache.stride(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_paged< - D, LOW_SMEM, LOW_SMEM_CONTIG_FAST, - LOW_SMEM_SCALAR_QK, LOW_SMEM_BM32, false, true, - KV_DTYPE, D256_OUTPUT_STRIDE_268, - D256_SW_PIPELINE_QK, D256_SW_PIPELINE_PV> - : (void*)flash_attention_forward_kernel_paged< - D, LOW_SMEM, LOW_SMEM_CONTIG_FAST, - LOW_SMEM_SCALAR_QK, LOW_SMEM_BM32, false, false, - KV_DTYPE, D256_OUTPUT_STRIDE_268, - D256_SW_PIPELINE_QK, D256_SW_PIPELINE_PV>; - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, - smem); - - if (is_causal) { - flash_attention_forward_kernel_paged< - D, LOW_SMEM, LOW_SMEM_CONTIG_FAST, LOW_SMEM_SCALAR_QK, - LOW_SMEM_BM32, false, true, KV_DTYPE, D256_OUTPUT_STRIDE_268, - D256_SW_PIPELINE_QK, D256_SW_PIPELINE_PV> - <<>>( - reinterpret_cast(Q.data_ptr()), - K_cache.data_ptr(), - V_cache.data_ptr(), - reinterpret_cast<__half*>(Out.data_ptr()), - softmax_lse.data_ptr(), - block_table.data_ptr(), - seq_lens.data_ptr(), - B, - H, - M, - N, - bfla_mask_ptr, - bfla_mask_block_n, - bfla_mask_stride_b, - bfla_mask_stride_h, - bfla_mask_stride_q, - bfla_mask_stride_k, - page_block_size, - num_kv_heads, - k_block_stride, - k_token_stride, - k_head_stride, - v_block_stride, - v_token_stride, - v_head_stride, - softmax_scale, - k_scale, - v_scale, - window_size_left, - window_size_right, - nullptr, - nullptr, - nullptr, - 0 - ); - } else { - flash_attention_forward_kernel_paged< - D, LOW_SMEM, LOW_SMEM_CONTIG_FAST, LOW_SMEM_SCALAR_QK, - LOW_SMEM_BM32, false, false, KV_DTYPE, D256_OUTPUT_STRIDE_268, - D256_SW_PIPELINE_QK, D256_SW_PIPELINE_PV> - <<>>( - reinterpret_cast(Q.data_ptr()), - K_cache.data_ptr(), - V_cache.data_ptr(), - reinterpret_cast<__half*>(Out.data_ptr()), - softmax_lse.data_ptr(), - block_table.data_ptr(), - seq_lens.data_ptr(), - B, - H, - M, - N, - bfla_mask_ptr, - bfla_mask_block_n, - bfla_mask_stride_b, - bfla_mask_stride_h, - bfla_mask_stride_q, - bfla_mask_stride_k, - page_block_size, - num_kv_heads, - k_block_stride, - k_token_stride, - k_head_stride, - v_block_stride, - v_token_stride, - v_head_stride, - softmax_scale, - k_scale, - v_scale, - window_size_left, - window_size_right, - nullptr, - nullptr, - nullptr, - 0 - ); - } + const torch::Tensor& Q, const torch::Tensor& K_cache, + const torch::Tensor& V_cache, torch::Tensor& Out, + torch::Tensor& softmax_lse, const torch::Tensor& block_table, + const torch::Tensor& seq_lens, const int* bfla_mask_ptr, + int bfla_mask_block_n, int64_t bfla_mask_stride_b, + int64_t bfla_mask_stride_h, int64_t bfla_mask_stride_q, + int64_t bfla_mask_stride_k, float softmax_scale, bool is_causal, + float k_scale, float v_scale, int window_size_left, int window_size_right, + cudaStream_t stream, const int* anchor_lens = nullptr, + int anchored_window = 0) { + using Config = KernelConfig; + + const int B = Q.size(0); + const int H = Q.size(1); + const int M = Q.size(2); + const int page_block_size = K_cache.size(1); + const int num_kv_heads = K_cache.size(2); + const int max_num_blocks = block_table.size(1); + const int N = max_num_blocks * page_block_size; + const int64_t k_block_stride = K_cache.stride(0); + const int64_t k_token_stride = K_cache.stride(1); + const int64_t k_head_stride = K_cache.stride(2); + const int64_t v_block_stride = V_cache.stride(0); + const int64_t v_token_stride = V_cache.stride(1); + const int64_t v_head_stride = V_cache.stride(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"); + + const bool use_anchored = anchor_lens != nullptr && anchored_window > 0; + TORCH_CHECK(!use_anchored || is_causal, + "anchored decode window requires causal attention"); + if constexpr (KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP16) { + if (use_anchored) { + auto anchored_kernel = flash_attention_forward_kernel_paged< + D, LOW_SMEM, LOW_SMEM_CONTIG_FAST, LOW_SMEM_SCALAR_QK, LOW_SMEM_BM32, + false, true, KV_DTYPE, D256_OUTPUT_STRIDE_268, D256_SW_PIPELINE_QK, + D256_SW_PIPELINE_PV, true>; + cudaFuncSetAttribute((void*)anchored_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem); + anchored_kernel<<>>( + reinterpret_cast(Q.data_ptr()), K_cache.data_ptr(), + V_cache.data_ptr(), reinterpret_cast<__half*>(Out.data_ptr()), + softmax_lse.data_ptr(), block_table.data_ptr(), + seq_lens.data_ptr(), B, H, M, N, bfla_mask_ptr, + bfla_mask_block_n, bfla_mask_stride_b, bfla_mask_stride_h, + bfla_mask_stride_q, bfla_mask_stride_k, page_block_size, num_kv_heads, + k_block_stride, k_token_stride, k_head_stride, v_block_stride, + v_token_stride, v_head_stride, softmax_scale, k_scale, v_scale, + window_size_left, window_size_right, nullptr, nullptr, nullptr, 0, + anchor_lens, anchored_window); + return; + } + } else { + TORCH_CHECK(!use_anchored, + "anchored decode window requires an fp16 KV cache"); + } + + auto kernel = + is_causal + ? (void*)flash_attention_forward_kernel_paged< + D, LOW_SMEM, LOW_SMEM_CONTIG_FAST, LOW_SMEM_SCALAR_QK, + LOW_SMEM_BM32, false, true, KV_DTYPE, D256_OUTPUT_STRIDE_268, + D256_SW_PIPELINE_QK, D256_SW_PIPELINE_PV> + : (void*)flash_attention_forward_kernel_paged< + D, LOW_SMEM, LOW_SMEM_CONTIG_FAST, LOW_SMEM_SCALAR_QK, + LOW_SMEM_BM32, false, false, KV_DTYPE, D256_OUTPUT_STRIDE_268, + D256_SW_PIPELINE_QK, D256_SW_PIPELINE_PV>; + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + smem); + + if (is_causal) { + flash_attention_forward_kernel_paged< + D, LOW_SMEM, LOW_SMEM_CONTIG_FAST, LOW_SMEM_SCALAR_QK, LOW_SMEM_BM32, + false, true, KV_DTYPE, D256_OUTPUT_STRIDE_268, D256_SW_PIPELINE_QK, + D256_SW_PIPELINE_PV><<>>( + reinterpret_cast(Q.data_ptr()), K_cache.data_ptr(), + V_cache.data_ptr(), reinterpret_cast<__half*>(Out.data_ptr()), + softmax_lse.data_ptr(), block_table.data_ptr(), + seq_lens.data_ptr(), B, H, M, N, bfla_mask_ptr, bfla_mask_block_n, + bfla_mask_stride_b, bfla_mask_stride_h, bfla_mask_stride_q, + bfla_mask_stride_k, page_block_size, num_kv_heads, k_block_stride, + k_token_stride, k_head_stride, v_block_stride, v_token_stride, + v_head_stride, softmax_scale, k_scale, v_scale, window_size_left, + window_size_right, nullptr, nullptr, nullptr, 0, nullptr, 0); + } else { + flash_attention_forward_kernel_paged< + D, LOW_SMEM, LOW_SMEM_CONTIG_FAST, LOW_SMEM_SCALAR_QK, LOW_SMEM_BM32, + false, false, KV_DTYPE, D256_OUTPUT_STRIDE_268, D256_SW_PIPELINE_QK, + D256_SW_PIPELINE_PV><<>>( + reinterpret_cast(Q.data_ptr()), K_cache.data_ptr(), + V_cache.data_ptr(), reinterpret_cast<__half*>(Out.data_ptr()), + softmax_lse.data_ptr(), block_table.data_ptr(), + seq_lens.data_ptr(), B, H, M, N, bfla_mask_ptr, bfla_mask_block_n, + bfla_mask_stride_b, bfla_mask_stride_h, bfla_mask_stride_q, + bfla_mask_stride_k, page_block_size, num_kv_heads, k_block_stride, + k_token_stride, k_head_stride, v_block_stride, v_token_stride, + v_head_stride, softmax_scale, k_scale, v_scale, window_size_left, + window_size_right, nullptr, nullptr, nullptr, 0, nullptr, 0); + } } -template +template __global__ void flash_attention_forward_paged_splitkv_merge_kernel( const float* __restrict__ split_tmp_out, const float* __restrict__ split_tmp_row_max, - const float* __restrict__ split_tmp_row_sum, - __half* __restrict__ Out, - float* __restrict__ softmax_lse, - const int B, - const int H, - const int M, - const int num_partitions -) { - static_assert(D == 256, "split-KV paged prefill merge is D=256 only"); - constexpr int BLOCK_M = BLOCK_M_256_LOW_SMEM; - constexpr int THREADS = 512; - 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; - const int valid_q_rows = min(BLOCK_M, M - start_row); - const int tid = threadIdx.x; - - __shared__ float sFinalMax[BLOCK_M]; - __shared__ float sFinalSum[BLOCK_M]; - - if (tid < valid_q_rows) { - const int row = start_row + tid; - float final_max = NEG_INF; - #pragma unroll 1 - for (int p = 0; p < num_partitions; ++p) { - const int64_t state_idx = - ((int64_t)batch_head_id * num_partitions + p) * M + row; - final_max = fmaxf(final_max, split_tmp_row_max[state_idx]); - } - - float final_sum = 0.0f; - #pragma unroll 1 - for (int p = 0; p < num_partitions; ++p) { - const int64_t state_idx = - ((int64_t)batch_head_id * num_partitions + p) * M + row; - const float part_sum = split_tmp_row_sum[state_idx]; - if (part_sum > 0.0f) { - final_sum += - __expf(fmaxf(split_tmp_row_max[state_idx] - final_max, - -80.0f)) * part_sum; - } - } - sFinalMax[tid] = final_max; - sFinalSum[tid] = final_sum; - softmax_lse[(int64_t)batch_head_id * M + row] = - final_max + logf(fmaxf(final_sum, 1e-24f)); - } - __syncthreads(); - - for (int idx = tid; idx < valid_q_rows * D; idx += THREADS) { - const int row_local = idx / D; - const int col = idx - row_local * D; - const int row = start_row + row_local; - const float final_max = sFinalMax[row_local]; - const float inv_sum = 1.0f / fmaxf(sFinalSum[row_local], 1e-24f); - - float acc = 0.0f; - #pragma unroll 1 - for (int p = 0; p < num_partitions; ++p) { - const int64_t state_idx = - ((int64_t)batch_head_id * num_partitions + p) * M + row; - const float part_sum = split_tmp_row_sum[state_idx]; - if (part_sum > 0.0f) { - const float scale = - __expf(fmaxf(split_tmp_row_max[state_idx] - final_max, - -80.0f)); - const int64_t out_idx = state_idx * D + col; - acc = fmaf(scale, split_tmp_out[out_idx], acc); - } - } - Out[((int64_t)batch_head_id * M + row) * D + col] = - __float2half_rn(acc * inv_sum); - } + const float* __restrict__ split_tmp_row_sum, __half* __restrict__ Out, + float* __restrict__ softmax_lse, const int B, const int H, const int M, + const int num_partitions) { + static_assert(D == 256, "split-KV paged prefill merge is D=256 only"); + constexpr int BLOCK_M = BLOCK_M_256_LOW_SMEM; + constexpr int THREADS = 512; + 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; + const int valid_q_rows = min(BLOCK_M, M - start_row); + const int tid = threadIdx.x; + + __shared__ float sFinalMax[BLOCK_M]; + __shared__ float sFinalSum[BLOCK_M]; + + if (tid < valid_q_rows) { + const int row = start_row + tid; + float final_max = NEG_INF; +#pragma unroll 1 + for (int p = 0; p < num_partitions; ++p) { + const int64_t state_idx = + ((int64_t)batch_head_id * num_partitions + p) * M + row; + final_max = fmaxf(final_max, split_tmp_row_max[state_idx]); + } + + float final_sum = 0.0f; +#pragma unroll 1 + for (int p = 0; p < num_partitions; ++p) { + const int64_t state_idx = + ((int64_t)batch_head_id * num_partitions + p) * M + row; + const float part_sum = split_tmp_row_sum[state_idx]; + if (part_sum > 0.0f) { + final_sum += + __expf(fmaxf(split_tmp_row_max[state_idx] - final_max, -80.0f)) * + part_sum; + } + } + sFinalMax[tid] = final_max; + sFinalSum[tid] = final_sum; + softmax_lse[(int64_t)batch_head_id * M + row] = + final_max + logf(fmaxf(final_sum, 1e-24f)); + } + __syncthreads(); + + for (int idx = tid; idx < valid_q_rows * D; idx += THREADS) { + const int row_local = idx / D; + const int col = idx - row_local * D; + const int row = start_row + row_local; + const float final_max = sFinalMax[row_local]; + const float inv_sum = 1.0f / fmaxf(sFinalSum[row_local], 1e-24f); + + float acc = 0.0f; +#pragma unroll 1 + for (int p = 0; p < num_partitions; ++p) { + const int64_t state_idx = + ((int64_t)batch_head_id * num_partitions + p) * M + row; + const float part_sum = split_tmp_row_sum[state_idx]; + if (part_sum > 0.0f) { + const float scale = + __expf(fmaxf(split_tmp_row_max[state_idx] - final_max, -80.0f)); + const int64_t out_idx = state_idx * D + col; + acc = fmaf(scale, split_tmp_out[out_idx], acc); + } + } + Out[((int64_t)batch_head_id * M + row) * D + col] = + __float2half_rn(acc * inv_sum); + } } -template +template void launcher_flash_attention_forward_paged_splitkv_impl( - const torch::Tensor& Q, - const torch::Tensor& K_cache, - const torch::Tensor& V_cache, - torch::Tensor& Out, - torch::Tensor& softmax_lse, - torch::Tensor& split_tmp_out, - torch::Tensor& split_tmp_row_max, - torch::Tensor& split_tmp_row_sum, - const torch::Tensor& block_table, - const torch::Tensor& seq_lens, - float softmax_scale, - bool is_causal, - float k_scale, - float v_scale, - int window_size_left, - int window_size_right, - int split_kv_tokens, - int max_seq_len_hint, - cudaStream_t stream -) { - static_assert(D == 256, "split-KV paged prefill is D=256 only"); - using Config = KernelConfig; - constexpr int KV_DTYPE = flash_v100::KV_CACHE_DTYPE_FP16; - - const int B = Q.size(0); - const int H = Q.size(1); - const int M = Q.size(2); - const int page_block_size = K_cache.size(1); - const int num_kv_heads = K_cache.size(2); - const int max_num_blocks = block_table.size(1); - const int N = max_num_blocks * page_block_size; - const int64_t k_block_stride = K_cache.stride(0); - const int64_t k_token_stride = K_cache.stride(1); - const int64_t k_head_stride = K_cache.stride(2); - const int64_t v_block_stride = V_cache.stride(0); - const int64_t v_token_stride = V_cache.stride(1); - const int64_t v_head_stride = V_cache.stride(2); - - split_kv_tokens = std::max(split_kv_tokens, Config::BLOCK_N); - max_seq_len_hint = std::max(max_seq_len_hint, 1); - const int split_kv_tiles = - std::max(1, (split_kv_tokens + Config::BLOCK_N - 1) / Config::BLOCK_N); - const int max_kv_tiles = - std::max(1, (max_seq_len_hint + Config::BLOCK_N - 1) / Config::BLOCK_N); - const int num_partitions = - std::max(1, (max_kv_tiles + split_kv_tiles - 1) / split_kv_tiles); - - const int grid_x = (M + Config::BLOCK_M - 1) / Config::BLOCK_M; - const dim3 grid(grid_x, num_partitions, 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"); - TORCH_CHECK(split_tmp_out.size(2) == num_partitions, - "split_tmp_out partition mismatch"); - TORCH_CHECK(split_tmp_row_max.size(2) == num_partitions, - "split_tmp_row_max partition mismatch"); - TORCH_CHECK(split_tmp_row_sum.size(2) == num_partitions, - "split_tmp_row_sum partition mismatch"); - - auto kernel = is_causal - ? (void*)flash_attention_forward_kernel_paged< - D, true, LOW_SMEM_CONTIG_FAST, - LOW_SMEM_SCALAR_QK, false, true, true, KV_DTYPE> - : (void*)flash_attention_forward_kernel_paged< - D, true, LOW_SMEM_CONTIG_FAST, - LOW_SMEM_SCALAR_QK, false, true, false, KV_DTYPE>; - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, - smem); - - if (is_causal) { - flash_attention_forward_kernel_paged< - D, true, LOW_SMEM_CONTIG_FAST, LOW_SMEM_SCALAR_QK, false, true, true, - KV_DTYPE> - <<>>( - reinterpret_cast(Q.data_ptr()), - K_cache.data_ptr(), - V_cache.data_ptr(), - reinterpret_cast<__half*>(Out.data_ptr()), - softmax_lse.data_ptr(), - block_table.data_ptr(), - seq_lens.data_ptr(), - B, - H, - M, - N, - nullptr, - 0, - 0, - 0, - 0, - 0, - page_block_size, - num_kv_heads, - k_block_stride, - k_token_stride, - k_head_stride, - v_block_stride, - v_token_stride, - v_head_stride, - softmax_scale, - k_scale, - v_scale, - window_size_left, - window_size_right, - split_tmp_out.data_ptr(), - split_tmp_row_max.data_ptr(), - split_tmp_row_sum.data_ptr(), - split_kv_tiles - ); - } else { - flash_attention_forward_kernel_paged< - D, true, LOW_SMEM_CONTIG_FAST, LOW_SMEM_SCALAR_QK, false, true, false, - KV_DTYPE> - <<>>( - reinterpret_cast(Q.data_ptr()), - K_cache.data_ptr(), - V_cache.data_ptr(), - reinterpret_cast<__half*>(Out.data_ptr()), - softmax_lse.data_ptr(), - block_table.data_ptr(), - seq_lens.data_ptr(), - B, - H, - M, - N, - nullptr, - 0, - 0, - 0, - 0, - 0, - page_block_size, - num_kv_heads, - k_block_stride, - k_token_stride, - k_head_stride, - v_block_stride, - v_token_stride, - v_head_stride, - softmax_scale, - k_scale, - v_scale, - window_size_left, - window_size_right, - split_tmp_out.data_ptr(), + const torch::Tensor& Q, const torch::Tensor& K_cache, + const torch::Tensor& V_cache, torch::Tensor& Out, + torch::Tensor& softmax_lse, torch::Tensor& split_tmp_out, + torch::Tensor& split_tmp_row_max, torch::Tensor& split_tmp_row_sum, + const torch::Tensor& block_table, const torch::Tensor& seq_lens, + float softmax_scale, bool is_causal, float k_scale, float v_scale, + int window_size_left, int window_size_right, int split_kv_tokens, + int max_seq_len_hint, cudaStream_t stream) { + static_assert(D == 256, "split-KV paged prefill is D=256 only"); + static_assert(LOW_SMEM || !LOW_SMEM_CONTIG_FAST, + "contiguous low-smem loads require LOW_SMEM"); + static_assert(LOW_SMEM || !LOW_SMEM_SCALAR_QK, + "scalar low-smem QK requires LOW_SMEM"); + static_assert(!LOW_SMEM || KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP16, + "low-smem split-KV supports FP16 KV only"); + using Config = KernelConfig; + + const int B = Q.size(0); + const int H = Q.size(1); + const int M = Q.size(2); + const int page_block_size = K_cache.size(1); + const int num_kv_heads = K_cache.size(2); + const int max_num_blocks = block_table.size(1); + const int N = max_num_blocks * page_block_size; + const int64_t k_block_stride = K_cache.stride(0); + const int64_t k_token_stride = K_cache.stride(1); + const int64_t k_head_stride = K_cache.stride(2); + const int64_t v_block_stride = V_cache.stride(0); + const int64_t v_token_stride = V_cache.stride(1); + const int64_t v_head_stride = V_cache.stride(2); + + split_kv_tokens = std::max(split_kv_tokens, Config::BLOCK_N); + max_seq_len_hint = std::max(max_seq_len_hint, 1); + const int split_kv_tiles = + std::max(1, (split_kv_tokens + Config::BLOCK_N - 1) / Config::BLOCK_N); + const int max_kv_tiles = + std::max(1, (max_seq_len_hint + Config::BLOCK_N - 1) / Config::BLOCK_N); + const int num_partitions = + std::max(1, (max_kv_tiles + split_kv_tiles - 1) / split_kv_tiles); + + const int grid_x = (M + Config::BLOCK_M - 1) / Config::BLOCK_M; + const dim3 grid(grid_x, num_partitions, 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"); + TORCH_CHECK(split_tmp_out.size(2) == num_partitions, + "split_tmp_out partition mismatch"); + TORCH_CHECK(split_tmp_row_max.size(2) == num_partitions, + "split_tmp_row_max partition mismatch"); + TORCH_CHECK(split_tmp_row_sum.size(2) == num_partitions, + "split_tmp_row_sum partition mismatch"); + + auto kernel = is_causal + ? (void*)flash_attention_forward_kernel_paged< + D, LOW_SMEM, LOW_SMEM_CONTIG_FAST, LOW_SMEM_SCALAR_QK, + false, true, true, KV_DTYPE> + : (void*)flash_attention_forward_kernel_paged< + D, LOW_SMEM, LOW_SMEM_CONTIG_FAST, LOW_SMEM_SCALAR_QK, + false, true, false, KV_DTYPE>; + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + smem); + + if (is_causal) { + flash_attention_forward_kernel_paged + <<>>( + reinterpret_cast(Q.data_ptr()), K_cache.data_ptr(), + V_cache.data_ptr(), reinterpret_cast<__half*>(Out.data_ptr()), + softmax_lse.data_ptr(), block_table.data_ptr(), + seq_lens.data_ptr(), B, H, M, N, nullptr, 0, 0, 0, 0, 0, + page_block_size, num_kv_heads, k_block_stride, k_token_stride, + k_head_stride, v_block_stride, v_token_stride, v_head_stride, + softmax_scale, k_scale, v_scale, window_size_left, + window_size_right, split_tmp_out.data_ptr(), split_tmp_row_max.data_ptr(), - split_tmp_row_sum.data_ptr(), - split_kv_tiles - ); - } - - const dim3 merge_grid(grid_x, 1, B * H); - const dim3 merge_block(512); - flash_attention_forward_paged_splitkv_merge_kernel - <<>>( - split_tmp_out.data_ptr(), + split_tmp_row_sum.data_ptr(), split_kv_tiles, nullptr, 0); + } else { + flash_attention_forward_kernel_paged + <<>>( + reinterpret_cast(Q.data_ptr()), K_cache.data_ptr(), + V_cache.data_ptr(), reinterpret_cast<__half*>(Out.data_ptr()), + softmax_lse.data_ptr(), block_table.data_ptr(), + seq_lens.data_ptr(), B, H, M, N, nullptr, 0, 0, 0, 0, 0, + page_block_size, num_kv_heads, k_block_stride, k_token_stride, + k_head_stride, v_block_stride, v_token_stride, v_head_stride, + softmax_scale, k_scale, v_scale, window_size_left, + window_size_right, split_tmp_out.data_ptr(), split_tmp_row_max.data_ptr(), - split_tmp_row_sum.data_ptr(), - reinterpret_cast<__half*>(Out.data_ptr()), - softmax_lse.data_ptr(), - B, - H, - M, - num_partitions); + split_tmp_row_sum.data_ptr(), split_kv_tiles, nullptr, 0); + } + + const dim3 merge_grid(grid_x, 1, B * H); + const dim3 merge_block(512); + flash_attention_forward_paged_splitkv_merge_kernel + <<>>( + split_tmp_out.data_ptr(), split_tmp_row_max.data_ptr(), + split_tmp_row_sum.data_ptr(), + reinterpret_cast<__half*>(Out.data_ptr()), + softmax_lse.data_ptr(), B, H, M, num_partitions); } constexpr int D256_BM32_PHASE_BLOCK_M = 32; @@ -2331,41 +2100,41 @@ constexpr int D256_BM32_PHASE_PANELS = constexpr int D256_BM32_PHASE_PROBABILITY_ELEMENTS = D256_BM32_PHASE_PANEL_M * D256_BM32_PHASE_SOFTMAX_N; -template +template struct alignas(16) D256BM32PhaseSharedStorage { - __half query[D256_BM32_PHASE_BLOCK_M * D256_BM32_PHASE_D]; - float score[D256_BM32_PHASE_BLOCK_M * D256_BM32_PHASE_BLOCK_N]; - __half probability_top[PROBABILITY_PANELS - * D256_BM32_PHASE_PROBABILITY_ELEMENTS]; - __half probability_bottom[PROBABILITY_PANELS - * D256_BM32_PHASE_PROBABILITY_ELEMENTS]; - float row_max[D256_BM32_PHASE_BLOCK_M]; - float row_sum[D256_BM32_PHASE_BLOCK_M]; - float row_exp_diff[PROBABILITY_PANELS * D256_BM32_PHASE_BLOCK_M]; - int page_idx[D256_BM32_PHASE_PAGE_SLOTS]; - int page_offset[D256_BM32_PHASE_PAGE_SLOTS]; - uint64_t k_tile_ptr[D256_BM32_PHASE_PAGE_SLOTS]; - uint64_t v_tile_ptr[D256_BM32_PHASE_PAGE_SLOTS]; - int block_index; - int batch_id; - int kv_head_id; - int actual_n; + __half query[D256_BM32_PHASE_BLOCK_M * D256_BM32_PHASE_D]; + float score[D256_BM32_PHASE_BLOCK_M * D256_BM32_PHASE_BLOCK_N]; + __half probability_top[PROBABILITY_PANELS * + D256_BM32_PHASE_PROBABILITY_ELEMENTS]; + __half probability_bottom[PROBABILITY_PANELS * + D256_BM32_PHASE_PROBABILITY_ELEMENTS]; + float row_max[D256_BM32_PHASE_BLOCK_M]; + float row_sum[D256_BM32_PHASE_BLOCK_M]; + float row_exp_diff[PROBABILITY_PANELS * D256_BM32_PHASE_BLOCK_M]; + int page_idx[D256_BM32_PHASE_PAGE_SLOTS]; + int page_offset[D256_BM32_PHASE_PAGE_SLOTS]; + uint64_t k_tile_ptr[D256_BM32_PHASE_PAGE_SLOTS]; + uint64_t v_tile_ptr[D256_BM32_PHASE_PAGE_SLOTS]; + int block_index; + int batch_id; + int kv_head_id; + int actual_n; }; -template +template struct alignas(16) D256BM32PhaseSplitSharedStorage : D256BM32PhaseSharedStorage { - int split_end_block; + int split_end_block; }; -template +template struct D256BM32PhaseSharedStorageSelector { - using Type = D256BM32PhaseSharedStorage; + using Type = D256BM32PhaseSharedStorage; }; -template +template struct D256BM32PhaseSharedStorageSelector { - using Type = D256BM32PhaseSplitSharedStorage; + using Type = D256BM32PhaseSplitSharedStorage; }; using D256BM32PhaseSharedStorageSingle = D256BM32PhaseSharedStorage<1>; @@ -2378,8 +2147,8 @@ static_assert(sizeof(D256BM32PhaseSharedStorageAllP) == 41936, "D256 BM32 all-P shared layout changed unexpectedly"); static_assert(sizeof(D256BM32PhaseSharedStorageAllP) <= 48 * 1024, "D256 BM32 phase shared storage exceeds 48 KiB"); -static_assert(sizeof(D256BM32PhaseSplitSharedStorage< - D256_BM32_PHASE_PANELS>) == 41952, +static_assert(sizeof(D256BM32PhaseSplitSharedStorage) == + 41952, "D256 BM32 split shared layout changed unexpectedly"); using D256BM32PhaseMatrixAFragment = @@ -2392,2143 +2161,1928 @@ using D256BM32PhaseAccumulatorFragment = fragment; __device__ __forceinline__ int d256_bm32_phase_swizzled_row_slot(int row) { - return (row & 3) | ((row & 8) >> 1) | ((row & 4) << 1); + return (row & 3) | ((row & 8) >> 1) | ((row & 4) << 1); } -__device__ __forceinline__ int d256_bm32_phase_matrix_a_offset( - int row, - int column) { - const int tile = column / WMMA_K; - const int within_tile = column - tile * WMMA_K; - const int plane = within_tile >> 3; - const int inner = within_tile & 7; - return tile * WMMA_M * WMMA_K + plane * WMMA_M * 8 - + d256_bm32_phase_swizzled_row_slot(row) * 8 + inner; +__device__ __forceinline__ int d256_bm32_phase_matrix_a_offset(int row, + int column) { + const int tile = column / WMMA_K; + const int within_tile = column - tile * WMMA_K; + const int plane = within_tile >> 3; + const int inner = within_tile & 7; + return tile * WMMA_M * WMMA_K + plane * WMMA_M * 8 + + d256_bm32_phase_swizzled_row_slot(row) * 8 + inner; } __device__ __forceinline__ void d256_bm32_phase_stage_query( - const __half* __restrict__ source, - __half* __restrict__ destination) { - constexpr int HALF_PER_UINT4 = sizeof(uint4) / sizeof(__half); - constexpr int VECTORS_PER_ROW = D256_BM32_PHASE_D / HALF_PER_UINT4; - constexpr int VECTORS_PER_PANEL = - D256_BM32_PHASE_PANEL_M * VECTORS_PER_ROW; - constexpr int QUERY_VECTORS = - D256_BM32_PHASE_BLOCK_M * VECTORS_PER_ROW; - - const uint4* source_vectors = reinterpret_cast(source); - uint4* destination_vectors = reinterpret_cast(destination); - - #pragma unroll - for (int index = threadIdx.x; index < QUERY_VECTORS; - index += D256_BM32_PHASE_THREADS) { - const int row = index / VECTORS_PER_ROW; - const int vector_column = index % VECTORS_PER_ROW; - const int panel = row / D256_BM32_PHASE_PANEL_M; - const int panel_row = row - panel * D256_BM32_PHASE_PANEL_M; - const int k_tile = vector_column >> 1; - const int plane = vector_column & 1; - const int slot = d256_bm32_phase_swizzled_row_slot(panel_row); - destination_vectors[panel * VECTORS_PER_PANEL - + k_tile * (2 * D256_BM32_PHASE_PANEL_M) - + plane * D256_BM32_PHASE_PANEL_M + slot] = - __ldg(source_vectors + index); - } + const __half* __restrict__ source, __half* __restrict__ destination) { + constexpr int HALF_PER_UINT4 = sizeof(uint4) / sizeof(__half); + constexpr int VECTORS_PER_ROW = D256_BM32_PHASE_D / HALF_PER_UINT4; + constexpr int VECTORS_PER_PANEL = D256_BM32_PHASE_PANEL_M * VECTORS_PER_ROW; + constexpr int QUERY_VECTORS = D256_BM32_PHASE_BLOCK_M * VECTORS_PER_ROW; + + const uint4* source_vectors = reinterpret_cast(source); + uint4* destination_vectors = reinterpret_cast(destination); + +#pragma unroll + for (int index = threadIdx.x; index < QUERY_VECTORS; + index += D256_BM32_PHASE_THREADS) { + const int row = index / VECTORS_PER_ROW; + const int vector_column = index % VECTORS_PER_ROW; + const int panel = row / D256_BM32_PHASE_PANEL_M; + const int panel_row = row - panel * D256_BM32_PHASE_PANEL_M; + const int k_tile = vector_column >> 1; + const int plane = vector_column & 1; + const int slot = d256_bm32_phase_swizzled_row_slot(panel_row); + destination_vectors[panel * VECTORS_PER_PANEL + + k_tile * (2 * D256_BM32_PHASE_PANEL_M) + + plane * D256_BM32_PHASE_PANEL_M + slot] = + __ldg(source_vectors + index); + } } __device__ __forceinline__ void d256_bm32_phase_stage_query_partial( - const __half* __restrict__ source, - __half* __restrict__ destination, + const __half* __restrict__ source, __half* __restrict__ destination, int valid_rows) { - constexpr int HALF_PER_UINT4 = sizeof(uint4) / sizeof(__half); - constexpr int VECTORS_PER_ROW = D256_BM32_PHASE_D / HALF_PER_UINT4; - constexpr int VECTORS_PER_PANEL = - D256_BM32_PHASE_PANEL_M * VECTORS_PER_ROW; - constexpr int QUERY_VECTORS = - D256_BM32_PHASE_BLOCK_M * VECTORS_PER_ROW; - - const uint4* source_vectors = reinterpret_cast(source); - uint4* destination_vectors = reinterpret_cast(destination); - - #pragma unroll - for (int index = threadIdx.x; index < QUERY_VECTORS; - index += D256_BM32_PHASE_THREADS) { - const int row = index / VECTORS_PER_ROW; - const int vector_column = index % VECTORS_PER_ROW; - const int panel = row / D256_BM32_PHASE_PANEL_M; - const int panel_row = row - panel * D256_BM32_PHASE_PANEL_M; - const int k_tile = vector_column >> 1; - const int plane = vector_column & 1; - const int slot = d256_bm32_phase_swizzled_row_slot(panel_row); - uint4 value = {0, 0, 0, 0}; - if (row < valid_rows) { - value = __ldg(source_vectors + index); - } - destination_vectors[panel * VECTORS_PER_PANEL - + k_tile * (2 * D256_BM32_PHASE_PANEL_M) - + plane * D256_BM32_PHASE_PANEL_M + slot] = value; - } + constexpr int HALF_PER_UINT4 = sizeof(uint4) / sizeof(__half); + constexpr int VECTORS_PER_ROW = D256_BM32_PHASE_D / HALF_PER_UINT4; + constexpr int VECTORS_PER_PANEL = D256_BM32_PHASE_PANEL_M * VECTORS_PER_ROW; + constexpr int QUERY_VECTORS = D256_BM32_PHASE_BLOCK_M * VECTORS_PER_ROW; + + const uint4* source_vectors = reinterpret_cast(source); + uint4* destination_vectors = reinterpret_cast(destination); + +#pragma unroll + for (int index = threadIdx.x; index < QUERY_VECTORS; + index += D256_BM32_PHASE_THREADS) { + const int row = index / VECTORS_PER_ROW; + const int vector_column = index % VECTORS_PER_ROW; + const int panel = row / D256_BM32_PHASE_PANEL_M; + const int panel_row = row - panel * D256_BM32_PHASE_PANEL_M; + const int k_tile = vector_column >> 1; + const int plane = vector_column & 1; + const int slot = d256_bm32_phase_swizzled_row_slot(panel_row); + uint4 value = {0, 0, 0, 0}; + if (row < valid_rows) { + value = __ldg(source_vectors + index); + } + destination_vectors[panel * VECTORS_PER_PANEL + + k_tile * (2 * D256_BM32_PHASE_PANEL_M) + + plane * D256_BM32_PHASE_PANEL_M + slot] = value; + } } __device__ __forceinline__ void d256_bm32_phase_load_matrix_a( - D256BM32PhaseMatrixAFragment& fragment, - const __half* __restrict__ matrix, + D256BM32PhaseMatrixAFragment& fragment, const __half* __restrict__ matrix, int k_offset) { - const int lane = threadIdx.x & 31; - const int row = - (lane & 3) + ((lane >> 4) & 1) * 4 + ((lane >> 2) & 1) * 8; - const int slot = d256_bm32_phase_swizzled_row_slot(row); - const int tile_offset = (k_offset / WMMA_K) * WMMA_M * WMMA_K; - uint32_t address = static_cast( - __cvta_generic_to_shared(matrix + tile_offset + slot * 8)); - uint32_t* words = reinterpret_cast(&fragment); - asm volatile( - "ld.shared.v4.u32 {%0, %1, %2, %3}, [%4];" - : "=r"(words[0]), "=r"(words[1]), "=r"(words[2]), "=r"(words[3]) - : "r"(address) - : "memory"); - address += WMMA_M * 8 * sizeof(__half); - asm volatile( - "ld.shared.v4.u32 {%0, %1, %2, %3}, [%4];" - : "=r"(words[4]), "=r"(words[5]), "=r"(words[6]), "=r"(words[7]) - : "r"(address) - : "memory"); + const int lane = threadIdx.x & 31; + const int row = (lane & 3) + ((lane >> 4) & 1) * 4 + ((lane >> 2) & 1) * 8; + const int slot = d256_bm32_phase_swizzled_row_slot(row); + const int tile_offset = (k_offset / WMMA_K) * WMMA_M * WMMA_K; + uint32_t address = static_cast( + __cvta_generic_to_shared(matrix + tile_offset + slot * 8)); + uint32_t* words = reinterpret_cast(&fragment); + asm volatile("ld.shared.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(words[0]), "=r"(words[1]), "=r"(words[2]), "=r"(words[3]) + : "r"(address) + : "memory"); + address += WMMA_M * 8 * sizeof(__half); + asm volatile("ld.shared.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(words[4]), "=r"(words[5]), "=r"(words[6]), "=r"(words[7]) + : "r"(address) + : "memory"); } -__device__ __forceinline__ int d256_bm32_phase_accumulator_row( - int lane, - int element) { - const int row_base = - (lane & 1) + ((lane >> 2) & 1) * 8 + ((lane >> 4) & 1) * 4; - return row_base + ((element >> 1) & 1) * 2; +__device__ __forceinline__ int d256_bm32_phase_accumulator_row(int lane, + int element) { + const int row_base = + (lane & 1) + ((lane >> 2) & 1) * 8 + ((lane >> 4) & 1) * 4; + return row_base + ((element >> 1) & 1) * 2; } -__device__ __forceinline__ int d256_bm32_phase_accumulator_column( - int lane, - int element) { - const int column_base = ((lane >> 1) & 1) * 2 + ((lane >> 3) & 1) * 8; - return column_base + (element & 1) + ((element >> 2) & 1) * 4; +__device__ __forceinline__ int d256_bm32_phase_accumulator_column(int lane, + int element) { + const int column_base = ((lane >> 1) & 1) * 2 + ((lane >> 3) & 1) * 8; + return column_base + (element & 1) + ((element >> 2) & 1) * 4; } __device__ __forceinline__ void d256_bm32_phase_spill_pair_scratch( float* __restrict__ score, D256BM32PhaseAccumulatorFragment& accumulator_top, D256BM32PhaseAccumulatorFragment& accumulator_bottom) { - const int warp = threadIdx.x >> 5; - const int lane = threadIdx.x & 31; - const int warp_pair = warp >> 1; - const int warp_in_pair = warp & 1; - const int pair_column = warp_pair * 32 + (lane & 15) * 2; - const int lane_row = lane >> 4; - - #pragma unroll - for (int element_pair = 0; element_pair < 4; ++element_pair) { - const int element = element_pair * 2; - const int offset = - (warp_in_pair * 16 + element_pair * 2 + lane_row) - * D256_BM32_PHASE_BLOCK_N - + pair_column; - const uint32_t address = static_cast( - __cvta_generic_to_shared(score + offset)); - asm volatile( - "st.shared.v2.u32 [%0], {%1, %2};" - : - : "r"(address), - "r"(__float_as_uint(accumulator_top.x[element])), - "r"(__float_as_uint(accumulator_top.x[element + 1])) - : "memory"); - asm volatile( - "st.shared.v2.u32 [%0+4096], {%1, %2};" - : - : "r"(address), - "r"(__float_as_uint(accumulator_bottom.x[element])), - "r"(__float_as_uint(accumulator_bottom.x[element + 1])) - : "memory"); - } + const int warp = threadIdx.x >> 5; + const int lane = threadIdx.x & 31; + const int warp_pair = warp >> 1; + const int warp_in_pair = warp & 1; + const int pair_column = warp_pair * 32 + (lane & 15) * 2; + const int lane_row = lane >> 4; + +#pragma unroll + for (int element_pair = 0; element_pair < 4; ++element_pair) { + const int element = element_pair * 2; + const int offset = (warp_in_pair * 16 + element_pair * 2 + lane_row) * + D256_BM32_PHASE_BLOCK_N + + pair_column; + const uint32_t address = + static_cast(__cvta_generic_to_shared(score + offset)); + asm volatile("st.shared.v2.u32 [%0], {%1, %2};" + : + : "r"(address), + "r"(__float_as_uint(accumulator_top.x[element])), + "r"(__float_as_uint(accumulator_top.x[element + 1])) + : "memory"); + asm volatile("st.shared.v2.u32 [%0+4096], {%1, %2};" + : + : "r"(address), + "r"(__float_as_uint(accumulator_bottom.x[element])), + "r"(__float_as_uint(accumulator_bottom.x[element + 1])) + : "memory"); + } } __device__ __forceinline__ void d256_bm32_phase_reload_pair_scratch( const float* __restrict__ score, D256BM32PhaseAccumulatorFragment& accumulator_top, D256BM32PhaseAccumulatorFragment& accumulator_bottom) { - const int warp = threadIdx.x >> 5; - const int lane = threadIdx.x & 31; - const int warp_pair = warp >> 1; - const int warp_in_pair = warp & 1; - const int pair_column = warp_pair * 32 + (lane & 15) * 2; - const int lane_row = lane >> 4; - - #pragma unroll - for (int element_pair = 0; element_pair < 4; ++element_pair) { - const int element = element_pair * 2; - const int offset = - (warp_in_pair * 16 + element_pair * 2 + lane_row) - * D256_BM32_PHASE_BLOCK_N - + pair_column; - const uint32_t address = static_cast( - __cvta_generic_to_shared(score + offset)); - uint32_t first_word; - uint32_t second_word; - asm volatile( - "ld.shared.v2.u32 {%0, %1}, [%2];" - : "=r"(first_word), "=r"(second_word) - : "r"(address) - : "memory"); - accumulator_top.x[element] = __uint_as_float(first_word); - accumulator_top.x[element + 1] = __uint_as_float(second_word); - asm volatile( - "ld.shared.v2.u32 {%0, %1}, [%2+4096];" - : "=r"(first_word), "=r"(second_word) - : "r"(address) - : "memory"); - accumulator_bottom.x[element] = __uint_as_float(first_word); - accumulator_bottom.x[element + 1] = __uint_as_float(second_word); - } + const int warp = threadIdx.x >> 5; + const int lane = threadIdx.x & 31; + const int warp_pair = warp >> 1; + const int warp_in_pair = warp & 1; + const int pair_column = warp_pair * 32 + (lane & 15) * 2; + const int lane_row = lane >> 4; + +#pragma unroll + for (int element_pair = 0; element_pair < 4; ++element_pair) { + const int element = element_pair * 2; + const int offset = (warp_in_pair * 16 + element_pair * 2 + lane_row) * + D256_BM32_PHASE_BLOCK_N + + pair_column; + const uint32_t address = + static_cast(__cvta_generic_to_shared(score + offset)); + uint32_t first_word; + uint32_t second_word; + asm volatile("ld.shared.v2.u32 {%0, %1}, [%2];" + : "=r"(first_word), "=r"(second_word) + : "r"(address) + : "memory"); + accumulator_top.x[element] = __uint_as_float(first_word); + accumulator_top.x[element + 1] = __uint_as_float(second_word); + asm volatile("ld.shared.v2.u32 {%0, %1}, [%2+4096];" + : "=r"(first_word), "=r"(second_word) + : "r"(address) + : "memory"); + accumulator_bottom.x[element] = __uint_as_float(first_word); + accumulator_bottom.x[element + 1] = __uint_as_float(second_word); + } } __device__ __forceinline__ void d256_bm32_phase_sync_warp_pair(int warp_pair) { - const int barrier_id = warp_pair + 1; - asm volatile("bar.sync %0, 64;" : : "r"(barrier_id) : "memory"); + const int barrier_id = warp_pair + 1; + asm volatile("bar.sync %0, 64;" : : "r"(barrier_id) : "memory"); } __device__ __forceinline__ float d256_bm32_phase_make_probability_row( - const float* __restrict__ score_row, - __half* __restrict__ probability, - int probability_row, - float* __restrict__ row_max, - float* __restrict__ row_sum, - int state_row, - int panel, - float neg_inf) { - const int lane = threadIdx.x & 31; - const float* panel_score = - score_row + panel * D256_BM32_PHASE_SOFTMAX_N; - float thread_max = neg_inf; - if (lane < D256_BM32_PHASE_SOFTMAX_N / 4) { - const float4 values = - reinterpret_cast(panel_score)[lane]; - thread_max = fmaxf(fmaxf(values.x, values.y), - fmaxf(values.z, values.w)); - } - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) { - thread_max = fmaxf( - thread_max, __shfl_down_sync(0xffffffffU, thread_max, offset)); - } - const float panel_max = __shfl_sync(0xffffffffU, thread_max, 0); - const float old_max = row_max[state_row]; - const float new_max = fmaxf(old_max, panel_max); - const float exp_diff = __expf(old_max - new_max); - - float thread_sum = 0.0f; - if (lane < D256_BM32_PHASE_SOFTMAX_N / 4) { - float4 values = reinterpret_cast(panel_score)[lane]; - values.x = __expf(fmaxf(values.x - new_max, -80.0f)); - values.y = __expf(fmaxf(values.y - new_max, -80.0f)); - values.z = __expf(fmaxf(values.z - new_max, -80.0f)); - values.w = __expf(fmaxf(values.w - new_max, -80.0f)); - thread_sum = (values.x + values.y) + (values.z + values.w); - const int column = lane * 4; - __half2* first_pair = reinterpret_cast<__half2*>( - probability + d256_bm32_phase_matrix_a_offset( - probability_row, column)); - __half2* second_pair = reinterpret_cast<__half2*>( - probability + d256_bm32_phase_matrix_a_offset( - probability_row, column + 2)); - *first_pair = __float22half2_rn(make_float2(values.x, values.y)); - *second_pair = __float22half2_rn(make_float2(values.z, values.w)); - } - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) { - thread_sum += - __shfl_down_sync(0xffffffffU, thread_sum, offset); - } - const float panel_sum = __shfl_sync(0xffffffffU, thread_sum, 0); - if (lane == 0) { - row_sum[state_row] = exp_diff * row_sum[state_row] + panel_sum; - row_max[state_row] = new_max; - } - return exp_diff; + const float* __restrict__ score_row, __half* __restrict__ probability, + int probability_row, float* __restrict__ row_max, + float* __restrict__ row_sum, int state_row, int panel, float neg_inf) { + const int lane = threadIdx.x & 31; + const float* panel_score = score_row + panel * D256_BM32_PHASE_SOFTMAX_N; + float thread_max = neg_inf; + if (lane < D256_BM32_PHASE_SOFTMAX_N / 4) { + const float4 values = reinterpret_cast(panel_score)[lane]; + thread_max = fmaxf(fmaxf(values.x, values.y), fmaxf(values.z, values.w)); + } +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + thread_max = + fmaxf(thread_max, __shfl_down_sync(0xffffffffU, thread_max, offset)); + } + const float panel_max = __shfl_sync(0xffffffffU, thread_max, 0); + const float old_max = row_max[state_row]; + const float new_max = fmaxf(old_max, panel_max); + const float exp_diff = __expf(old_max - new_max); + + float thread_sum = 0.0f; + if (lane < D256_BM32_PHASE_SOFTMAX_N / 4) { + float4 values = reinterpret_cast(panel_score)[lane]; + values.x = __expf(fmaxf(values.x - new_max, -80.0f)); + values.y = __expf(fmaxf(values.y - new_max, -80.0f)); + values.z = __expf(fmaxf(values.z - new_max, -80.0f)); + values.w = __expf(fmaxf(values.w - new_max, -80.0f)); + thread_sum = (values.x + values.y) + (values.z + values.w); + const int column = lane * 4; + __half2* first_pair = reinterpret_cast<__half2*>( + probability + d256_bm32_phase_matrix_a_offset(probability_row, column)); + __half2* second_pair = reinterpret_cast<__half2*>( + probability + + d256_bm32_phase_matrix_a_offset(probability_row, column + 2)); + *first_pair = __float22half2_rn(make_float2(values.x, values.y)); + *second_pair = __float22half2_rn(make_float2(values.z, values.w)); + } +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + thread_sum += __shfl_down_sync(0xffffffffU, thread_sum, offset); + } + const float panel_sum = __shfl_sync(0xffffffffU, thread_sum, 0); + if (lane == 0) { + row_sum[state_row] = exp_diff * row_sum[state_row] + panel_sum; + row_max[state_row] = new_max; + } + return exp_diff; } -__device__ __forceinline__ float -d256_bm32_phase_make_split_probability_row( - const float* __restrict__ score_row, - __half* __restrict__ probability, - int probability_row, - float* __restrict__ row_max, - float* __restrict__ row_sum, - int state_row, - int panel, - float neg_inf, +__device__ __forceinline__ float d256_bm32_phase_make_split_probability_row( + const float* __restrict__ score_row, __half* __restrict__ probability, + int probability_row, float* __restrict__ row_max, + float* __restrict__ row_sum, int state_row, int panel, float neg_inf, bool has_visible_value) { - if (!has_visible_value) { - const int lane = threadIdx.x & 31; - if (lane < D256_BM32_PHASE_SOFTMAX_N / 4) { - const int column = lane * 4; - __half2* first_pair = reinterpret_cast<__half2*>( - probability + d256_bm32_phase_matrix_a_offset( - probability_row, column)); - __half2* second_pair = reinterpret_cast<__half2*>( - probability + d256_bm32_phase_matrix_a_offset( - probability_row, column + 2)); - const __half2 zero = __float2half2_rn(0.0f); - *first_pair = zero; - *second_pair = zero; - } - return 1.0f; - } - return d256_bm32_phase_make_probability_row( - score_row, probability, probability_row, row_max, row_sum, - state_row, panel, neg_inf); + if (!has_visible_value) { + const int lane = threadIdx.x & 31; + if (lane < D256_BM32_PHASE_SOFTMAX_N / 4) { + const int column = lane * 4; + __half2* first_pair = reinterpret_cast<__half2*>( + probability + + d256_bm32_phase_matrix_a_offset(probability_row, column)); + __half2* second_pair = reinterpret_cast<__half2*>( + probability + + d256_bm32_phase_matrix_a_offset(probability_row, column + 2)); + const __half2 zero = __float2half2_rn(0.0f); + *first_pair = zero; + *second_pair = zero; + } + return 1.0f; + } + return d256_bm32_phase_make_probability_row(score_row, probability, + probability_row, row_max, row_sum, + state_row, panel, neg_inf); } __device__ __forceinline__ void d256_bm32_phase_scale_accumulator_rows( - D256BM32PhaseAccumulatorFragment& accumulator, - float first_row_scale, + D256BM32PhaseAccumulatorFragment& accumulator, float first_row_scale, float second_row_scale) { - accumulator.x[0] *= first_row_scale; - accumulator.x[1] *= first_row_scale; - accumulator.x[2] *= second_row_scale; - accumulator.x[3] *= second_row_scale; - accumulator.x[4] *= first_row_scale; - accumulator.x[5] *= first_row_scale; - accumulator.x[6] *= second_row_scale; - accumulator.x[7] *= second_row_scale; + accumulator.x[0] *= first_row_scale; + accumulator.x[1] *= first_row_scale; + accumulator.x[2] *= second_row_scale; + accumulator.x[3] *= second_row_scale; + accumulator.x[4] *= first_row_scale; + accumulator.x[5] *= first_row_scale; + accumulator.x[6] *= second_row_scale; + accumulator.x[7] *= second_row_scale; } __device__ __forceinline__ void d256_bm32_phase_scale_accumulators( D256BM32PhaseAccumulatorFragment& accumulator_top, D256BM32PhaseAccumulatorFragment& accumulator_bottom, const float* __restrict__ row_exp_diff) { - const int row = d256_bm32_phase_accumulator_row(threadIdx.x & 31, 0); - d256_bm32_phase_scale_accumulator_rows( - accumulator_top, row_exp_diff[row], row_exp_diff[row + 2]); - d256_bm32_phase_scale_accumulator_rows( - accumulator_bottom, - row_exp_diff[D256_BM32_PHASE_PANEL_M + row], - row_exp_diff[D256_BM32_PHASE_PANEL_M + row + 2]); + const int row = d256_bm32_phase_accumulator_row(threadIdx.x & 31, 0); + d256_bm32_phase_scale_accumulator_rows(accumulator_top, row_exp_diff[row], + row_exp_diff[row + 2]); + d256_bm32_phase_scale_accumulator_rows( + accumulator_bottom, row_exp_diff[D256_BM32_PHASE_PANEL_M + row], + row_exp_diff[D256_BM32_PHASE_PANEL_M + row + 2]); } __device__ __forceinline__ void d256_bm32_phase_update_pv_panel( const __half* __restrict__ probability_top, const __half* __restrict__ probability_bottom, - const __half* __restrict__ value_k0, - const __half* __restrict__ value_k16, - int64_t value_token_stride, - int d_offset, - int valid_columns, + const __half* __restrict__ value_k0, const __half* __restrict__ value_k16, + int64_t value_token_stride, int d_offset, int valid_columns, D256BM32PhaseAccumulatorFragment& accumulator_top, D256BM32PhaseAccumulatorFragment& accumulator_bottom) { - D256BM32PhaseMatrixAFragment a_fragment; - D256BM32PhasePVMatrixBFragment b_fragment; - - load_matrix_sync(b_fragment, value_k0 + d_offset, value_token_stride); - d256_bm32_phase_load_matrix_a(a_fragment, probability_top, 0); + D256BM32PhaseMatrixAFragment a_fragment; + D256BM32PhasePVMatrixBFragment b_fragment; + + load_matrix_sync(b_fragment, value_k0 + d_offset, value_token_stride); + d256_bm32_phase_load_matrix_a(a_fragment, probability_top, 0); + mma_sync(accumulator_top, a_fragment, b_fragment, accumulator_top); + d256_bm32_phase_load_matrix_a(a_fragment, probability_bottom, 0); + mma_sync(accumulator_bottom, a_fragment, b_fragment, accumulator_bottom); + if (valid_columns > WMMA_K) { + load_matrix_sync(b_fragment, value_k16 + d_offset, value_token_stride); + d256_bm32_phase_load_matrix_a(a_fragment, probability_top, WMMA_K); mma_sync(accumulator_top, a_fragment, b_fragment, accumulator_top); - d256_bm32_phase_load_matrix_a(a_fragment, probability_bottom, 0); + d256_bm32_phase_load_matrix_a(a_fragment, probability_bottom, WMMA_K); mma_sync(accumulator_bottom, a_fragment, b_fragment, accumulator_bottom); - if (valid_columns > WMMA_K) { - load_matrix_sync( - b_fragment, value_k16 + d_offset, value_token_stride); - d256_bm32_phase_load_matrix_a(a_fragment, probability_top, WMMA_K); - mma_sync(accumulator_top, a_fragment, b_fragment, accumulator_top); - d256_bm32_phase_load_matrix_a(a_fragment, probability_bottom, WMMA_K); - mma_sync(accumulator_bottom, a_fragment, b_fragment, - accumulator_bottom); - } + } } __device__ __forceinline__ void d256_bm32_phase_store_output( const D256BM32PhaseAccumulatorFragment& accumulator, - __half* __restrict__ output, - const float* __restrict__ row_sum, - int row_offset, - int d_offset) { - const int lane = threadIdx.x & 31; - #pragma unroll - for (int element = 0; element < accumulator.num_elements; ++element) { - const int row = row_offset - + d256_bm32_phase_accumulator_row(lane, element); - const int column = d_offset - + d256_bm32_phase_accumulator_column(lane, element); - const float inverse_sum = 1.0f / fmaxf(row_sum[row], 1e-24f); - output[row * D256_BM32_PHASE_D + column] = - __float2half_rn(accumulator.x[element] * inverse_sum); - } + __half* __restrict__ output, const float* __restrict__ row_sum, + int row_offset, int d_offset) { + const int lane = threadIdx.x & 31; +#pragma unroll + for (int element = 0; element < accumulator.num_elements; ++element) { + const int row = row_offset + d256_bm32_phase_accumulator_row(lane, element); + const int column = + d_offset + d256_bm32_phase_accumulator_column(lane, element); + const float inverse_sum = 1.0f / fmaxf(row_sum[row], 1e-24f); + output[row * D256_BM32_PHASE_D + column] = + __float2half_rn(accumulator.x[element] * inverse_sum); + } } __device__ __forceinline__ void d256_bm32_phase_store_output_partial( const D256BM32PhaseAccumulatorFragment& accumulator, - __half* __restrict__ output, - const float* __restrict__ row_sum, - int row_offset, - int d_offset, - int valid_rows) { - const int lane = threadIdx.x & 31; - #pragma unroll - for (int element = 0; element < accumulator.num_elements; ++element) { - const int row = row_offset - + d256_bm32_phase_accumulator_row(lane, element); - if (row < valid_rows) { - const int column = - d_offset - + d256_bm32_phase_accumulator_column(lane, element); - const float inverse_sum = 1.0f / fmaxf(row_sum[row], 1e-24f); - output[row * D256_BM32_PHASE_D + column] = - __float2half_rn(accumulator.x[element] * inverse_sum); - } - } + __half* __restrict__ output, const float* __restrict__ row_sum, + int row_offset, int d_offset, int valid_rows) { + const int lane = threadIdx.x & 31; +#pragma unroll + for (int element = 0; element < accumulator.num_elements; ++element) { + const int row = row_offset + d256_bm32_phase_accumulator_row(lane, element); + if (row < valid_rows) { + const int column = + d_offset + d256_bm32_phase_accumulator_column(lane, element); + const float inverse_sum = 1.0f / fmaxf(row_sum[row], 1e-24f); + output[row * D256_BM32_PHASE_D + column] = + __float2half_rn(accumulator.x[element] * inverse_sum); + } + } } __device__ __forceinline__ void d256_bm32_phase_store_unnormalized_partial( const D256BM32PhaseAccumulatorFragment& accumulator, - float* __restrict__ output, - int row_offset, - int d_offset, - int valid_rows) { - const int lane = threadIdx.x & 31; - #pragma unroll - for (int element = 0; element < accumulator.num_elements; ++element) { - const int row = row_offset - + d256_bm32_phase_accumulator_row(lane, element); - if (row < valid_rows) { - const int column = - d_offset - + d256_bm32_phase_accumulator_column(lane, element); - output[row * D256_BM32_PHASE_D + column] = accumulator.x[element]; - } - } + float* __restrict__ output, int row_offset, int d_offset, int valid_rows) { + const int lane = threadIdx.x & 31; +#pragma unroll + for (int element = 0; element < accumulator.num_elements; ++element) { + const int row = row_offset + d256_bm32_phase_accumulator_row(lane, element); + if (row < valid_rows) { + const int column = + d_offset + d256_bm32_phase_accumulator_column(lane, element); + output[row * D256_BM32_PHASE_D + column] = accumulator.x[element]; + } + } } -template -__device__ __forceinline__ -void flash_attention_forward_paged_d256_bm32_phase_body( - const __half* __restrict__ Q, - const __half* __restrict__ K_cache, - const __half* __restrict__ V_cache, - __half* __restrict__ Out, - float* __restrict__ softmax_lse, - const int* __restrict__ block_table, - const int* __restrict__ seqused_k, - int H, - int M, - int max_num_blocks_per_seq, - int num_kv_heads, - int64_t k_block_stride, - int64_t k_token_stride, - int64_t k_head_stride, - int64_t v_block_stride, - int64_t v_token_stride, - int64_t v_head_stride, - float softmax_scale, - float* __restrict__ split_tmp_out, - float* __restrict__ split_tmp_row_max, - float* __restrict__ split_tmp_row_sum, - int split_actual_n) { - constexpr float NEG_INF = -1e30f; - static_assert(!SPLIT_KV || (ALL_P && PAIR_SCRATCH), - "BM32 split-KV requires the accepted all-P pair-scratch body"); - __shared__ __align__(16) - typename D256BM32PhaseSharedStorageSelector< - SPLIT_KV, - ALL_P ? D256_BM32_PHASE_PANELS : 1>::Type shared; - - const int start_row = blockIdx.x * D256_BM32_PHASE_BLOCK_M; - const int tid = threadIdx.x; - const int warp_id = tid >> 5; - const int lane_id = tid & 31; - const int valid_q_rows = - min(D256_BM32_PHASE_BLOCK_M, M - start_row); - - { - const __half* q_ptr = - Q + static_cast(blockIdx.z) * M * D256_BM32_PHASE_D - + start_row * D256_BM32_PHASE_D; - if (valid_q_rows == D256_BM32_PHASE_BLOCK_M) { - d256_bm32_phase_stage_query(q_ptr, shared.query); - } else { - d256_bm32_phase_stage_query_partial( - q_ptr, shared.query, valid_q_rows); - } - } - if (tid < D256_BM32_PHASE_BLOCK_M) { - shared.row_max[tid] = NEG_INF; - shared.row_sum[tid] = 0.0f; +template +__device__ __forceinline__ void +flash_attention_forward_paged_d256_bm32_phase_body( + const __half* __restrict__ Q, const __half* __restrict__ K_cache, + const __half* __restrict__ V_cache, __half* __restrict__ Out, + float* __restrict__ softmax_lse, const int* __restrict__ block_table, + const int* __restrict__ seqused_k, int H, int M, int max_num_blocks_per_seq, + int num_kv_heads, int64_t k_block_stride, int64_t k_token_stride, + int64_t k_head_stride, int64_t v_block_stride, int64_t v_token_stride, + int64_t v_head_stride, float softmax_scale, + float* __restrict__ split_tmp_out, float* __restrict__ split_tmp_row_max, + float* __restrict__ split_tmp_row_sum, int split_actual_n) { + constexpr float NEG_INF = -1e30f; + static_assert(!SPLIT_KV || (ALL_P && PAIR_SCRATCH), + "BM32 split-KV requires the accepted all-P pair-scratch body"); + __shared__ __align__(16) typename D256BM32PhaseSharedStorageSelector< + SPLIT_KV, ALL_P ? D256_BM32_PHASE_PANELS : 1>::Type shared; + + const int start_row = blockIdx.x * D256_BM32_PHASE_BLOCK_M; + const int tid = threadIdx.x; + const int warp_id = tid >> 5; + const int lane_id = tid & 31; + const int valid_q_rows = min(D256_BM32_PHASE_BLOCK_M, M - start_row); + + { + const __half* q_ptr = + Q + static_cast(blockIdx.z) * M * D256_BM32_PHASE_D + + start_row * D256_BM32_PHASE_D; + if (valid_q_rows == D256_BM32_PHASE_BLOCK_M) { + d256_bm32_phase_stage_query(q_ptr, shared.query); + } else { + d256_bm32_phase_stage_query_partial(q_ptr, shared.query, valid_q_rows); + } + } + if (tid < D256_BM32_PHASE_BLOCK_M) { + shared.row_max[tid] = NEG_INF; + shared.row_sum[tid] = 0.0f; + } + if (tid == 0) { + shared.block_index = 0; + shared.batch_id = blockIdx.z / H; + shared.kv_head_id = (blockIdx.z % H) / (H / num_kv_heads); + if constexpr (SPLIT_KV) { + shared.actual_n = split_actual_n; + } else { + shared.actual_n = seqused_k[shared.batch_id]; } + } + __syncthreads(); + + const int total_n_blocks = + (shared.actual_n + D256_BM32_PHASE_BLOCK_N - 1) / D256_BM32_PHASE_BLOCK_N; + if constexpr (SPLIT_KV) { if (tid == 0) { - shared.block_index = 0; - shared.batch_id = blockIdx.z / H; - shared.kv_head_id = - (blockIdx.z % H) / (H / num_kv_heads); - if constexpr (SPLIT_KV) { - shared.actual_n = split_actual_n; - } else { - shared.actual_n = seqused_k[shared.batch_id]; - } + shared.block_index = static_cast(blockIdx.y) * total_n_blocks / + D256_BM32_PHASE_SPLIT_PARTS; + shared.split_end_block = (static_cast(blockIdx.y) + 1) * + total_n_blocks / D256_BM32_PHASE_SPLIT_PARTS; } __syncthreads(); + } - const int total_n_blocks = - (shared.actual_n + D256_BM32_PHASE_BLOCK_N - 1) - / D256_BM32_PHASE_BLOCK_N; + D256BM32PhaseAccumulatorFragment accumulator_top; + D256BM32PhaseAccumulatorFragment accumulator_bottom; + fill_fragment(accumulator_top, 0.0f); + fill_fragment(accumulator_bottom, 0.0f); + + for (;;) { if constexpr (SPLIT_KV) { - if (tid == 0) { - shared.block_index = - static_cast(blockIdx.y) * total_n_blocks - / D256_BM32_PHASE_SPLIT_PARTS; - shared.split_end_block = - (static_cast(blockIdx.y) + 1) * total_n_blocks - / D256_BM32_PHASE_SPLIT_PARTS; - } - __syncthreads(); + if (shared.block_index >= shared.split_end_block) { + break; + } + } else { + if (shared.block_index >= total_n_blocks) { + break; + } + } + const int start_col = shared.block_index * D256_BM32_PHASE_BLOCK_N; + const int valid_k_rows = + min(D256_BM32_PHASE_BLOCK_N, shared.actual_n - start_col); + + if (tid < D256_BM32_PHASE_PAGE_SLOTS) { + const int token_offset = tid * D256_BM32_PHASE_PAGE_SIZE; + if (token_offset < valid_k_rows) { + const int global_token_idx = start_col + token_offset; + const int virtual_block_idx = + global_token_idx / D256_BM32_PHASE_PAGE_BLOCK_SIZE; + shared.page_idx[tid] = + __ldg(&block_table[shared.batch_id * max_num_blocks_per_seq + + virtual_block_idx]); + shared.page_offset[tid] = + global_token_idx - + virtual_block_idx * D256_BM32_PHASE_PAGE_BLOCK_SIZE; + shared.k_tile_ptr[tid] = reinterpret_cast( + K_cache + (int64_t)shared.page_idx[tid] * k_block_stride + + (int64_t)shared.page_offset[tid] * k_token_stride + + (int64_t)shared.kv_head_id * k_head_stride); + shared.v_tile_ptr[tid] = reinterpret_cast( + V_cache + (int64_t)shared.page_idx[tid] * v_block_stride + + (int64_t)shared.page_offset[tid] * v_token_stride + + (int64_t)shared.kv_head_id * v_head_stride); + } else { + shared.page_idx[tid] = -1; + shared.page_offset[tid] = 0; + shared.k_tile_ptr[tid] = 0; + shared.v_tile_ptr[tid] = 0; + } } + __syncthreads(); - D256BM32PhaseAccumulatorFragment accumulator_top; - D256BM32PhaseAccumulatorFragment accumulator_bottom; - fill_fragment(accumulator_top, 0.0f); - fill_fragment(accumulator_bottom, 0.0f); + if (warp_id < D256_BM32_PHASE_BLOCK_N / WMMA_N && + warp_id * WMMA_N < valid_k_rows) { + const int tile_n = warp_id * WMMA_N; + const __half* k_tile = + reinterpret_cast(shared.k_tile_ptr[warp_id]); + if constexpr (PAIR_SCRATCH) { + d256_bm32_phase_spill_pair_scratch(shared.score, accumulator_top, + accumulator_bottom); + } else { + store_matrix_sync(shared.score + tile_n, accumulator_top, + D256_BM32_PHASE_BLOCK_N, mem_row_major); + store_matrix_sync( + shared.score + D256_BM32_PHASE_PANEL_M * D256_BM32_PHASE_BLOCK_N + + tile_n, + accumulator_bottom, D256_BM32_PHASE_BLOCK_N, mem_row_major); + } + asm volatile("" ::: "memory"); + + D256BM32PhaseAccumulatorFragment qk_top; + D256BM32PhaseAccumulatorFragment qk_bottom; + { + D256BM32PhaseMatrixAFragment a_fragment; + D256BM32PhaseQKMatrixBFragment b_fragment; + fill_fragment(qk_top, 0.0f); + fill_fragment(qk_bottom, 0.0f); + +#pragma unroll + for (int k_offset = 0; k_offset < D256_BM32_PHASE_D; + k_offset += WMMA_K) { + load_matrix_sync(b_fragment, k_tile + k_offset, k_token_stride); + d256_bm32_phase_load_matrix_a(a_fragment, shared.query, k_offset); + mma_sync(qk_top, a_fragment, b_fragment, qk_top); + d256_bm32_phase_load_matrix_a( + a_fragment, + shared.query + D256_BM32_PHASE_PANEL_M * D256_BM32_PHASE_D, + k_offset); + mma_sync(qk_bottom, a_fragment, b_fragment, qk_bottom); + } + } + asm volatile("" ::: "memory"); + + if constexpr (PAIR_SCRATCH) { + d256_bm32_phase_reload_pair_scratch(shared.score, accumulator_top, + accumulator_bottom); + const int partner_warp = warp_id ^ 1; + if (partner_warp * WMMA_N < valid_k_rows) { + d256_bm32_phase_sync_warp_pair(warp_id >> 1); + } + } else { + load_matrix_sync(accumulator_top, shared.score + tile_n, + D256_BM32_PHASE_BLOCK_N, mem_row_major); + load_matrix_sync(accumulator_bottom, + shared.score + + D256_BM32_PHASE_PANEL_M * D256_BM32_PHASE_BLOCK_N + + tile_n, + D256_BM32_PHASE_BLOCK_N, mem_row_major); + } + + store_matrix_sync(shared.score + tile_n, qk_top, D256_BM32_PHASE_BLOCK_N, + mem_row_major); + store_matrix_sync(shared.score + + D256_BM32_PHASE_PANEL_M * D256_BM32_PHASE_BLOCK_N + + tile_n, + qk_bottom, D256_BM32_PHASE_BLOCK_N, mem_row_major); + } + __syncthreads(); - for (;;) { - if constexpr (SPLIT_KV) { - if (shared.block_index >= shared.split_end_block) { - break; - } + const int causal_q_offset = max(shared.actual_n - M, 0); +#pragma unroll 1 + for (int index = tid; + index < D256_BM32_PHASE_BLOCK_M * D256_BM32_PHASE_BLOCK_N; + index += D256_BM32_PHASE_THREADS) { + const int row = index / D256_BM32_PHASE_BLOCK_N; + const int column = index - row * D256_BM32_PHASE_BLOCK_N; + const int global_m = start_row + row; + const int global_n = start_col + column; + const bool is_valid = global_m < M && global_n < start_col + valid_k_rows; + if constexpr (IS_CAUSAL) { + if (is_valid && global_n <= global_m + causal_q_offset) { + shared.score[index] *= softmax_scale; } else { - if (shared.block_index >= total_n_blocks) { - break; - } + shared.score[index] = NEG_INF; } - const int start_col = - shared.block_index * D256_BM32_PHASE_BLOCK_N; - const int valid_k_rows = min(D256_BM32_PHASE_BLOCK_N, - shared.actual_n - start_col); - - if (tid < D256_BM32_PHASE_PAGE_SLOTS) { - const int token_offset = tid * D256_BM32_PHASE_PAGE_SIZE; - if (token_offset < valid_k_rows) { - const int global_token_idx = start_col + token_offset; - const int virtual_block_idx = - global_token_idx / D256_BM32_PHASE_PAGE_BLOCK_SIZE; - shared.page_idx[tid] = - __ldg(&block_table[shared.batch_id - * max_num_blocks_per_seq - + virtual_block_idx]); - shared.page_offset[tid] = - global_token_idx - - virtual_block_idx * D256_BM32_PHASE_PAGE_BLOCK_SIZE; - shared.k_tile_ptr[tid] = reinterpret_cast( - K_cache - + (int64_t)shared.page_idx[tid] * k_block_stride - + (int64_t)shared.page_offset[tid] * k_token_stride - + (int64_t)shared.kv_head_id * k_head_stride); - shared.v_tile_ptr[tid] = reinterpret_cast( - V_cache - + (int64_t)shared.page_idx[tid] * v_block_stride - + (int64_t)shared.page_offset[tid] * v_token_stride - + (int64_t)shared.kv_head_id * v_head_stride); - } else { - shared.page_idx[tid] = -1; - shared.page_offset[tid] = 0; - shared.k_tile_ptr[tid] = 0; - shared.v_tile_ptr[tid] = 0; - } + } else { + if (is_valid) { + shared.score[index] *= softmax_scale; + } else { + shared.score[index] = NEG_INF; } - __syncthreads(); - - if (warp_id < D256_BM32_PHASE_BLOCK_N / WMMA_N - && warp_id * WMMA_N < valid_k_rows) { - const int tile_n = warp_id * WMMA_N; - const __half* k_tile = - reinterpret_cast(shared.k_tile_ptr[warp_id]); - if constexpr (PAIR_SCRATCH) { - d256_bm32_phase_spill_pair_scratch( - shared.score, accumulator_top, accumulator_bottom); - } else { - store_matrix_sync(shared.score + tile_n, accumulator_top, - D256_BM32_PHASE_BLOCK_N, mem_row_major); - store_matrix_sync( - shared.score + D256_BM32_PHASE_PANEL_M - * D256_BM32_PHASE_BLOCK_N + tile_n, - accumulator_bottom, D256_BM32_PHASE_BLOCK_N, - mem_row_major); - } - asm volatile("" ::: "memory"); - - D256BM32PhaseAccumulatorFragment qk_top; - D256BM32PhaseAccumulatorFragment qk_bottom; - { - D256BM32PhaseMatrixAFragment a_fragment; - D256BM32PhaseQKMatrixBFragment b_fragment; - fill_fragment(qk_top, 0.0f); - fill_fragment(qk_bottom, 0.0f); - - #pragma unroll - for (int k_offset = 0; k_offset < D256_BM32_PHASE_D; - k_offset += WMMA_K) { - load_matrix_sync(b_fragment, k_tile + k_offset, - k_token_stride); - d256_bm32_phase_load_matrix_a( - a_fragment, shared.query, k_offset); - mma_sync(qk_top, a_fragment, b_fragment, qk_top); - d256_bm32_phase_load_matrix_a( - a_fragment, - shared.query - + D256_BM32_PHASE_PANEL_M - * D256_BM32_PHASE_D, - k_offset); - mma_sync(qk_bottom, a_fragment, b_fragment, qk_bottom); - } - } - asm volatile("" ::: "memory"); - - if constexpr (PAIR_SCRATCH) { - d256_bm32_phase_reload_pair_scratch( - shared.score, accumulator_top, accumulator_bottom); - const int partner_warp = warp_id ^ 1; - if (partner_warp * WMMA_N < valid_k_rows) { - d256_bm32_phase_sync_warp_pair(warp_id >> 1); - } - } else { - load_matrix_sync(accumulator_top, shared.score + tile_n, - D256_BM32_PHASE_BLOCK_N, mem_row_major); - load_matrix_sync( - accumulator_bottom, - shared.score + D256_BM32_PHASE_PANEL_M - * D256_BM32_PHASE_BLOCK_N + tile_n, - D256_BM32_PHASE_BLOCK_N, mem_row_major); - } + } + } + __syncthreads(); - store_matrix_sync(shared.score + tile_n, qk_top, - D256_BM32_PHASE_BLOCK_N, mem_row_major); - store_matrix_sync( - shared.score + D256_BM32_PHASE_PANEL_M - * D256_BM32_PHASE_BLOCK_N + tile_n, - qk_bottom, D256_BM32_PHASE_BLOCK_N, mem_row_major); + if constexpr (ALL_P) { +#pragma unroll + for (int panel = 0; panel < D256_BM32_PHASE_PANELS; ++panel) { + const int panel_start = panel * D256_BM32_PHASE_SOFTMAX_N; + const int valid_panel_columns = + min(D256_BM32_PHASE_SOFTMAX_N, valid_k_rows - panel_start); + if (valid_panel_columns > 0) { + __half* probability_top = + shared.probability_top + + panel * D256_BM32_PHASE_PROBABILITY_ELEMENTS; + __half* probability_bottom = + shared.probability_bottom + + panel * D256_BM32_PHASE_PROBABILITY_ELEMENTS; + float* row_exp_diff = + shared.row_exp_diff + panel * D256_BM32_PHASE_BLOCK_M; + float top_exp_diff; + float bottom_exp_diff; + if constexpr (SPLIT_KV && CHECK_SPLIT_EMPTY) { + const int first_global_n = start_col + panel_start; + const bool top_has_visible_value = + warp_id < valid_q_rows && + (!IS_CAUSAL || + first_global_n <= start_row + warp_id + causal_q_offset); + const bool bottom_has_visible_value = + D256_BM32_PHASE_PANEL_M + warp_id < valid_q_rows && + (!IS_CAUSAL || first_global_n <= start_row + + D256_BM32_PHASE_PANEL_M + + warp_id + causal_q_offset); + top_exp_diff = d256_bm32_phase_make_split_probability_row( + shared.score + warp_id * D256_BM32_PHASE_BLOCK_N, + probability_top, warp_id, shared.row_max, shared.row_sum, + warp_id, panel, NEG_INF, top_has_visible_value); + bottom_exp_diff = d256_bm32_phase_make_split_probability_row( + shared.score + (D256_BM32_PHASE_PANEL_M + warp_id) * + D256_BM32_PHASE_BLOCK_N, + probability_bottom, warp_id, shared.row_max, shared.row_sum, + D256_BM32_PHASE_PANEL_M + warp_id, panel, NEG_INF, + bottom_has_visible_value); + } else { + top_exp_diff = d256_bm32_phase_make_probability_row( + shared.score + warp_id * D256_BM32_PHASE_BLOCK_N, + probability_top, warp_id, shared.row_max, shared.row_sum, + warp_id, panel, NEG_INF); + bottom_exp_diff = d256_bm32_phase_make_probability_row( + shared.score + (D256_BM32_PHASE_PANEL_M + warp_id) * + D256_BM32_PHASE_BLOCK_N, + probability_bottom, warp_id, shared.row_max, shared.row_sum, + D256_BM32_PHASE_PANEL_M + warp_id, panel, NEG_INF); + } + if (lane_id == 0) { + row_exp_diff[warp_id] = top_exp_diff; + row_exp_diff[D256_BM32_PHASE_PANEL_M + warp_id] = bottom_exp_diff; + } } - __syncthreads(); - - const int causal_q_offset = max(shared.actual_n - M, 0); - #pragma unroll 1 - for (int index = tid; - index < D256_BM32_PHASE_BLOCK_M * D256_BM32_PHASE_BLOCK_N; - index += D256_BM32_PHASE_THREADS) { - const int row = index / D256_BM32_PHASE_BLOCK_N; - const int column = index - row * D256_BM32_PHASE_BLOCK_N; - const int global_m = start_row + row; - const int global_n = start_col + column; - const bool is_valid = global_m < M - && global_n < start_col + valid_k_rows; - if constexpr (IS_CAUSAL) { - if (is_valid && global_n <= global_m + causal_q_offset) { - shared.score[index] *= softmax_scale; - } else { - shared.score[index] = NEG_INF; - } - } else { - if (is_valid) { - shared.score[index] *= softmax_scale; - } else { - shared.score[index] = NEG_INF; - } - } + // Publish lane 0's online-softmax state to this warp before + // it advances to the next panel. + __syncwarp(); + } + __syncthreads(); + +#pragma unroll + for (int panel = 0; panel < D256_BM32_PHASE_PANELS; ++panel) { + const int panel_start = panel * D256_BM32_PHASE_SOFTMAX_N; + const int valid_panel_columns = + min(D256_BM32_PHASE_SOFTMAX_N, valid_k_rows - panel_start); + if (valid_panel_columns > 0) { + const __half* probability_top = + shared.probability_top + + panel * D256_BM32_PHASE_PROBABILITY_ELEMENTS; + const __half* probability_bottom = + shared.probability_bottom + + panel * D256_BM32_PHASE_PROBABILITY_ELEMENTS; + const float* row_exp_diff = + shared.row_exp_diff + panel * D256_BM32_PHASE_BLOCK_M; + if constexpr (SPLIT_KV) { + d256_bm32_phase_scale_accumulators( + accumulator_top, accumulator_bottom, row_exp_diff); + } else if (shared.block_index != 0 || panel != 0) { + d256_bm32_phase_scale_accumulators( + accumulator_top, accumulator_bottom, row_exp_diff); + } + const int page_slot_k0 = panel_start >> 4; + const __half* value_k0 = + reinterpret_cast(shared.v_tile_ptr[page_slot_k0]); + const __half* value_k16 = nullptr; + if (valid_panel_columns > WMMA_K) { + const int page_slot_k16 = (panel_start + WMMA_K) >> 4; + value_k16 = reinterpret_cast( + shared.v_tile_ptr[page_slot_k16]); + } + d256_bm32_phase_update_pv_panel(probability_top, probability_bottom, + value_k0, value_k16, v_token_stride, + warp_id * WMMA_N, valid_panel_columns, + accumulator_top, accumulator_bottom); } - __syncthreads(); - - if constexpr (ALL_P) { - #pragma unroll - for (int panel = 0; panel < D256_BM32_PHASE_PANELS; - ++panel) { - const int panel_start = - panel * D256_BM32_PHASE_SOFTMAX_N; - const int valid_panel_columns = - min(D256_BM32_PHASE_SOFTMAX_N, - valid_k_rows - panel_start); - if (valid_panel_columns > 0) { - __half* probability_top = - shared.probability_top - + panel * D256_BM32_PHASE_PROBABILITY_ELEMENTS; - __half* probability_bottom = - shared.probability_bottom - + panel * D256_BM32_PHASE_PROBABILITY_ELEMENTS; - float* row_exp_diff = - shared.row_exp_diff - + panel * D256_BM32_PHASE_BLOCK_M; - float top_exp_diff; - float bottom_exp_diff; - if constexpr (SPLIT_KV && CHECK_SPLIT_EMPTY) { - const int first_global_n = start_col + panel_start; - const bool top_has_visible_value = - warp_id < valid_q_rows - && (!IS_CAUSAL - || first_global_n - <= start_row + warp_id - + causal_q_offset); - const bool bottom_has_visible_value = - D256_BM32_PHASE_PANEL_M + warp_id < valid_q_rows - && (!IS_CAUSAL - || first_global_n - <= start_row - + D256_BM32_PHASE_PANEL_M - + warp_id + causal_q_offset); - top_exp_diff = - d256_bm32_phase_make_split_probability_row( - shared.score - + warp_id * D256_BM32_PHASE_BLOCK_N, - probability_top, warp_id, shared.row_max, - shared.row_sum, warp_id, panel, NEG_INF, - top_has_visible_value); - bottom_exp_diff = - d256_bm32_phase_make_split_probability_row( - shared.score - + (D256_BM32_PHASE_PANEL_M + warp_id) - * D256_BM32_PHASE_BLOCK_N, - probability_bottom, warp_id, shared.row_max, - shared.row_sum, - D256_BM32_PHASE_PANEL_M + warp_id, panel, - NEG_INF, bottom_has_visible_value); - } else { - top_exp_diff = d256_bm32_phase_make_probability_row( - shared.score - + warp_id * D256_BM32_PHASE_BLOCK_N, - probability_top, warp_id, shared.row_max, - shared.row_sum, warp_id, panel, NEG_INF); - bottom_exp_diff = - d256_bm32_phase_make_probability_row( - shared.score - + (D256_BM32_PHASE_PANEL_M + warp_id) - * D256_BM32_PHASE_BLOCK_N, - probability_bottom, warp_id, shared.row_max, - shared.row_sum, - D256_BM32_PHASE_PANEL_M + warp_id, panel, - NEG_INF); - } - if (lane_id == 0) { - row_exp_diff[warp_id] = top_exp_diff; - row_exp_diff[D256_BM32_PHASE_PANEL_M + warp_id] = - bottom_exp_diff; - } - } - // Publish lane 0's online-softmax state to this warp before - // it advances to the next panel. - __syncwarp(); - } - __syncthreads(); - - #pragma unroll - for (int panel = 0; panel < D256_BM32_PHASE_PANELS; - ++panel) { - const int panel_start = - panel * D256_BM32_PHASE_SOFTMAX_N; - const int valid_panel_columns = - min(D256_BM32_PHASE_SOFTMAX_N, - valid_k_rows - panel_start); - if (valid_panel_columns > 0) { - const __half* probability_top = - shared.probability_top - + panel * D256_BM32_PHASE_PROBABILITY_ELEMENTS; - const __half* probability_bottom = - shared.probability_bottom - + panel * D256_BM32_PHASE_PROBABILITY_ELEMENTS; - const float* row_exp_diff = - shared.row_exp_diff - + panel * D256_BM32_PHASE_BLOCK_M; - if constexpr (SPLIT_KV) { - d256_bm32_phase_scale_accumulators( - accumulator_top, accumulator_bottom, - row_exp_diff); - } else if (shared.block_index != 0 || panel != 0) { - d256_bm32_phase_scale_accumulators( - accumulator_top, accumulator_bottom, - row_exp_diff); - } - const int page_slot_k0 = panel_start >> 4; - const __half* value_k0 = - reinterpret_cast( - shared.v_tile_ptr[page_slot_k0]); - const __half* value_k16 = nullptr; - if (valid_panel_columns > WMMA_K) { - const int page_slot_k16 = - (panel_start + WMMA_K) >> 4; - value_k16 = - reinterpret_cast( - shared.v_tile_ptr[page_slot_k16]); - } - d256_bm32_phase_update_pv_panel( - probability_top, probability_bottom, - value_k0, value_k16, v_token_stride, - warp_id * WMMA_N, - valid_panel_columns, - accumulator_top, accumulator_bottom); - } - } - // Protect the current V pointers and block index from the next - // iteration while slower warps finish the last PV panel. - __syncthreads(); - } else { - #pragma unroll - for (int panel = 0; panel < D256_BM32_PHASE_PANELS; - ++panel) { - const int panel_start = - panel * D256_BM32_PHASE_SOFTMAX_N; - const int valid_panel_columns = - min(D256_BM32_PHASE_SOFTMAX_N, - valid_k_rows - panel_start); - if (valid_panel_columns > 0) { - const float top_exp_diff = - d256_bm32_phase_make_probability_row( - shared.score - + warp_id * D256_BM32_PHASE_BLOCK_N, - shared.probability_top, warp_id, shared.row_max, - shared.row_sum, warp_id, panel, NEG_INF); - const float bottom_exp_diff = - d256_bm32_phase_make_probability_row( - shared.score - + (D256_BM32_PHASE_PANEL_M + warp_id) - * D256_BM32_PHASE_BLOCK_N, - shared.probability_bottom, warp_id, - shared.row_max, shared.row_sum, - D256_BM32_PHASE_PANEL_M + warp_id, panel, - NEG_INF); - if (lane_id == 0) { - shared.row_exp_diff[warp_id] = top_exp_diff; - shared.row_exp_diff[ - D256_BM32_PHASE_PANEL_M + warp_id] = - bottom_exp_diff; - } - } - __syncthreads(); - - if (valid_panel_columns > 0) { - if constexpr (SPLIT_KV) { - d256_bm32_phase_scale_accumulators( - accumulator_top, accumulator_bottom, - shared.row_exp_diff); - } else if (shared.block_index != 0 || panel != 0) { - d256_bm32_phase_scale_accumulators( - accumulator_top, accumulator_bottom, - shared.row_exp_diff); - } - const int page_slot_k0 = panel_start >> 4; - const __half* value_k0 = - reinterpret_cast( - shared.v_tile_ptr[page_slot_k0]); - const __half* value_k16 = nullptr; - if (valid_panel_columns > WMMA_K) { - const int page_slot_k16 = - (panel_start + WMMA_K) >> 4; - value_k16 = - reinterpret_cast( - shared.v_tile_ptr[page_slot_k16]); - } - d256_bm32_phase_update_pv_panel( - shared.probability_top, - shared.probability_bottom, - value_k0, value_k16, v_token_stride, - warp_id * WMMA_N, - valid_panel_columns, - accumulator_top, accumulator_bottom); - } - __syncthreads(); - } + } + // Protect the current V pointers and block index from the next + // iteration while slower warps finish the last PV panel. + __syncthreads(); + } else { +#pragma unroll + for (int panel = 0; panel < D256_BM32_PHASE_PANELS; ++panel) { + const int panel_start = panel * D256_BM32_PHASE_SOFTMAX_N; + const int valid_panel_columns = + min(D256_BM32_PHASE_SOFTMAX_N, valid_k_rows - panel_start); + if (valid_panel_columns > 0) { + const float top_exp_diff = d256_bm32_phase_make_probability_row( + shared.score + warp_id * D256_BM32_PHASE_BLOCK_N, + shared.probability_top, warp_id, shared.row_max, shared.row_sum, + warp_id, panel, NEG_INF); + const float bottom_exp_diff = d256_bm32_phase_make_probability_row( + shared.score + + (D256_BM32_PHASE_PANEL_M + warp_id) * D256_BM32_PHASE_BLOCK_N, + shared.probability_bottom, warp_id, shared.row_max, + shared.row_sum, D256_BM32_PHASE_PANEL_M + warp_id, panel, + NEG_INF); + if (lane_id == 0) { + shared.row_exp_diff[warp_id] = top_exp_diff; + shared.row_exp_diff[D256_BM32_PHASE_PANEL_M + warp_id] = + bottom_exp_diff; + } } + __syncthreads(); - if (tid == 0) { - ++shared.block_index; + if (valid_panel_columns > 0) { + if constexpr (SPLIT_KV) { + d256_bm32_phase_scale_accumulators( + accumulator_top, accumulator_bottom, shared.row_exp_diff); + } else if (shared.block_index != 0 || panel != 0) { + d256_bm32_phase_scale_accumulators( + accumulator_top, accumulator_bottom, shared.row_exp_diff); + } + const int page_slot_k0 = panel_start >> 4; + const __half* value_k0 = + reinterpret_cast(shared.v_tile_ptr[page_slot_k0]); + const __half* value_k16 = nullptr; + if (valid_panel_columns > WMMA_K) { + const int page_slot_k16 = (panel_start + WMMA_K) >> 4; + value_k16 = reinterpret_cast( + shared.v_tile_ptr[page_slot_k16]); + } + d256_bm32_phase_update_pv_panel( + shared.probability_top, shared.probability_bottom, value_k0, + value_k16, v_token_stride, warp_id * WMMA_N, valid_panel_columns, + accumulator_top, accumulator_bottom); } __syncthreads(); + } } - { - const size_t q_head_linear = static_cast(blockIdx.z); - if constexpr (SPLIT_KV) { - const size_t partial_row_base = - (q_head_linear * D256_BM32_PHASE_SPLIT_PARTS - + static_cast(blockIdx.y)) - * M - + start_row; - float* partial_out = - split_tmp_out - + partial_row_base * D256_BM32_PHASE_D; - d256_bm32_phase_store_unnormalized_partial( - accumulator_top, partial_out, 0, - warp_id * WMMA_N, valid_q_rows); - d256_bm32_phase_store_unnormalized_partial( - accumulator_bottom, partial_out, - D256_BM32_PHASE_PANEL_M, warp_id * WMMA_N, valid_q_rows); - if (tid < valid_q_rows) { - split_tmp_row_max[partial_row_base + tid] = - shared.row_max[tid]; - split_tmp_row_sum[partial_row_base + tid] = - shared.row_sum[tid]; - } - } else { - __half* out_ptr = - Out + q_head_linear * M * D256_BM32_PHASE_D - + start_row * D256_BM32_PHASE_D; - float* softmax_lse_ptr = - softmax_lse + q_head_linear * M + start_row; - if (valid_q_rows == D256_BM32_PHASE_BLOCK_M) { - d256_bm32_phase_store_output( - accumulator_top, out_ptr, shared.row_sum, 0, - warp_id * WMMA_N); - d256_bm32_phase_store_output( - accumulator_bottom, out_ptr, shared.row_sum, - D256_BM32_PHASE_PANEL_M, warp_id * WMMA_N); - } else { - d256_bm32_phase_store_output_partial( - accumulator_top, out_ptr, shared.row_sum, 0, - warp_id * WMMA_N, valid_q_rows); - d256_bm32_phase_store_output_partial( - accumulator_bottom, out_ptr, shared.row_sum, - D256_BM32_PHASE_PANEL_M, warp_id * WMMA_N, valid_q_rows); - } - - if (tid < valid_q_rows) { - const float sum = fmaxf(shared.row_sum[tid], 1e-24f); - softmax_lse_ptr[tid] = shared.row_max[tid] + logf(sum); - } - } + if (tid == 0) { + ++shared.block_index; } + __syncthreads(); + } + + { + const size_t q_head_linear = static_cast(blockIdx.z); + if constexpr (SPLIT_KV) { + const size_t partial_row_base = + (q_head_linear * D256_BM32_PHASE_SPLIT_PARTS + + static_cast(blockIdx.y)) * + M + + start_row; + float* partial_out = split_tmp_out + partial_row_base * D256_BM32_PHASE_D; + d256_bm32_phase_store_unnormalized_partial( + accumulator_top, partial_out, 0, warp_id * WMMA_N, valid_q_rows); + d256_bm32_phase_store_unnormalized_partial( + accumulator_bottom, partial_out, D256_BM32_PHASE_PANEL_M, + warp_id * WMMA_N, valid_q_rows); + if (tid < valid_q_rows) { + split_tmp_row_max[partial_row_base + tid] = shared.row_max[tid]; + split_tmp_row_sum[partial_row_base + tid] = shared.row_sum[tid]; + } + } else { + __half* out_ptr = Out + q_head_linear * M * D256_BM32_PHASE_D + + start_row * D256_BM32_PHASE_D; + float* softmax_lse_ptr = softmax_lse + q_head_linear * M + start_row; + if (valid_q_rows == D256_BM32_PHASE_BLOCK_M) { + d256_bm32_phase_store_output(accumulator_top, out_ptr, shared.row_sum, + 0, warp_id * WMMA_N); + d256_bm32_phase_store_output(accumulator_bottom, out_ptr, + shared.row_sum, D256_BM32_PHASE_PANEL_M, + warp_id * WMMA_N); + } else { + d256_bm32_phase_store_output_partial(accumulator_top, out_ptr, + shared.row_sum, 0, + warp_id * WMMA_N, valid_q_rows); + d256_bm32_phase_store_output_partial( + accumulator_bottom, out_ptr, shared.row_sum, + D256_BM32_PHASE_PANEL_M, warp_id * WMMA_N, valid_q_rows); + } + + if (tid < valid_q_rows) { + const float sum = fmaxf(shared.row_sum[tid], 1e-24f); + softmax_lse_ptr[tid] = shared.row_max[tid] + logf(sum); + } + } + } } -template -__global__ __launch_bounds__(D256_BM32_PHASE_THREADS, 2) -void flash_attention_forward_paged_d256_bm32_phase_kernel( - const __half* __restrict__ Q, - const __half* __restrict__ K_cache, - const __half* __restrict__ V_cache, - __half* __restrict__ Out, - float* __restrict__ softmax_lse, - const int* __restrict__ block_table, - const int* __restrict__ seqused_k, - int H, - int M, - int max_num_blocks_per_seq, - int num_kv_heads, - int64_t k_block_stride, - int64_t k_token_stride, - int64_t k_head_stride, - int64_t v_block_stride, - int64_t v_token_stride, - int64_t v_head_stride, - float softmax_scale) { - flash_attention_forward_paged_d256_bm32_phase_body< - IS_CAUSAL, ALL_P, PAIR_SCRATCH, false>( - Q, K_cache, V_cache, Out, softmax_lse, block_table, seqused_k, - H, M, max_num_blocks_per_seq, num_kv_heads, k_block_stride, - k_token_stride, k_head_stride, v_block_stride, v_token_stride, - v_head_stride, softmax_scale, nullptr, nullptr, nullptr, 0); +template +__global__ +__launch_bounds__(D256_BM32_PHASE_THREADS, 2) void flash_attention_forward_paged_d256_bm32_phase_kernel( + const __half* __restrict__ Q, const __half* __restrict__ K_cache, + const __half* __restrict__ V_cache, __half* __restrict__ Out, + float* __restrict__ softmax_lse, const int* __restrict__ block_table, + const int* __restrict__ seqused_k, int H, int M, int max_num_blocks_per_seq, + int num_kv_heads, int64_t k_block_stride, int64_t k_token_stride, + int64_t k_head_stride, int64_t v_block_stride, int64_t v_token_stride, + int64_t v_head_stride, float softmax_scale) { + flash_attention_forward_paged_d256_bm32_phase_body( + Q, K_cache, V_cache, Out, softmax_lse, block_table, seqused_k, H, M, + max_num_blocks_per_seq, num_kv_heads, k_block_stride, k_token_stride, + k_head_stride, v_block_stride, v_token_stride, v_head_stride, + softmax_scale, nullptr, nullptr, nullptr, 0); } -template -__global__ __launch_bounds__(D256_BM32_PHASE_THREADS, 2) -void flash_attention_forward_paged_d256_bm32_splitkv3_partial_kernel( - const __half* __restrict__ Q, - const __half* __restrict__ K_cache, - const __half* __restrict__ V_cache, - const int* __restrict__ block_table, - int H, - int M, - int actual_n, - int max_num_blocks_per_seq, - int num_kv_heads, - int64_t k_block_stride, - int64_t k_token_stride, - int64_t k_head_stride, - int64_t v_block_stride, - int64_t v_token_stride, - int64_t v_head_stride, - float softmax_scale, - float* __restrict__ split_tmp_out, +template +__global__ +__launch_bounds__(D256_BM32_PHASE_THREADS, 2) void flash_attention_forward_paged_d256_bm32_splitkv3_partial_kernel( + const __half* __restrict__ Q, const __half* __restrict__ K_cache, + const __half* __restrict__ V_cache, const int* __restrict__ block_table, + int H, int M, int actual_n, int max_num_blocks_per_seq, int num_kv_heads, + int64_t k_block_stride, int64_t k_token_stride, int64_t k_head_stride, + int64_t v_block_stride, int64_t v_token_stride, int64_t v_head_stride, + float softmax_scale, float* __restrict__ split_tmp_out, float* __restrict__ split_tmp_row_max, float* __restrict__ split_tmp_row_sum) { - flash_attention_forward_paged_d256_bm32_phase_body< - true, true, true, true, CHECK_SPLIT_EMPTY>( - Q, K_cache, V_cache, nullptr, nullptr, block_table, nullptr, - H, M, max_num_blocks_per_seq, num_kv_heads, k_block_stride, - k_token_stride, k_head_stride, v_block_stride, v_token_stride, - v_head_stride, softmax_scale, split_tmp_out, - split_tmp_row_max, split_tmp_row_sum, actual_n); + flash_attention_forward_paged_d256_bm32_phase_body( + Q, K_cache, V_cache, nullptr, nullptr, block_table, nullptr, H, M, + max_num_blocks_per_seq, num_kv_heads, k_block_stride, k_token_stride, + k_head_stride, v_block_stride, v_token_stride, v_head_stride, + softmax_scale, split_tmp_out, split_tmp_row_max, split_tmp_row_sum, + actual_n); } -template +template void launch_flash_attention_forward_paged_d256_bm32_phase_kernel( - const torch::Tensor& Q, - const torch::Tensor& K_cache, - const torch::Tensor& V_cache, - torch::Tensor& Out, - torch::Tensor& softmax_lse, - const torch::Tensor& block_table, - const torch::Tensor& seq_lens, - float softmax_scale, - cudaStream_t stream) { - const int H = Q.size(1); - const int M = Q.size(2); - const int num_kv_heads = K_cache.size(2); - const int max_num_blocks_per_seq = block_table.size(1); - const int64_t k_block_stride = K_cache.stride(0); - const int64_t k_token_stride = K_cache.stride(1); - const int64_t k_head_stride = K_cache.stride(2); - const int64_t v_block_stride = V_cache.stride(0); - const int64_t v_token_stride = V_cache.stride(1); - const int64_t v_head_stride = V_cache.stride(2); - const dim3 grid((M + D256_BM32_PHASE_BLOCK_M - 1) - / D256_BM32_PHASE_BLOCK_M, - 1, Q.size(0) * H); - const dim3 block(D256_BM32_PHASE_THREADS); - flash_attention_forward_paged_d256_bm32_phase_kernel< - IS_CAUSAL, ALL_P, PAIR_SCRATCH><<>>( - reinterpret_cast(Q.data_ptr()), - reinterpret_cast(K_cache.data_ptr()), - reinterpret_cast(V_cache.data_ptr()), - reinterpret_cast<__half*>(Out.data_ptr()), - softmax_lse.data_ptr(), block_table.data_ptr(), - seq_lens.data_ptr(), H, M, max_num_blocks_per_seq, - num_kv_heads, k_block_stride, k_token_stride, k_head_stride, - v_block_stride, v_token_stride, v_head_stride, softmax_scale); + const torch::Tensor& Q, const torch::Tensor& K_cache, + const torch::Tensor& V_cache, torch::Tensor& Out, + torch::Tensor& softmax_lse, const torch::Tensor& block_table, + const torch::Tensor& seq_lens, float softmax_scale, cudaStream_t stream) { + const int H = Q.size(1); + const int M = Q.size(2); + const int num_kv_heads = K_cache.size(2); + const int max_num_blocks_per_seq = block_table.size(1); + const int64_t k_block_stride = K_cache.stride(0); + const int64_t k_token_stride = K_cache.stride(1); + const int64_t k_head_stride = K_cache.stride(2); + const int64_t v_block_stride = V_cache.stride(0); + const int64_t v_token_stride = V_cache.stride(1); + const int64_t v_head_stride = V_cache.stride(2); + const dim3 grid((M + D256_BM32_PHASE_BLOCK_M - 1) / D256_BM32_PHASE_BLOCK_M, + 1, Q.size(0) * H); + const dim3 block(D256_BM32_PHASE_THREADS); + flash_attention_forward_paged_d256_bm32_phase_kernel + <<>>( + reinterpret_cast(Q.data_ptr()), + reinterpret_cast(K_cache.data_ptr()), + reinterpret_cast(V_cache.data_ptr()), + reinterpret_cast<__half*>(Out.data_ptr()), + softmax_lse.data_ptr(), block_table.data_ptr(), + seq_lens.data_ptr(), H, M, max_num_blocks_per_seq, num_kv_heads, + k_block_stride, k_token_stride, k_head_stride, v_block_stride, + v_token_stride, v_head_stride, softmax_scale); } -__global__ __launch_bounds__(256) -void flash_attention_forward_paged_d256_bm32_splitkv3_merge_kernel( +__global__ +__launch_bounds__(256) void flash_attention_forward_paged_d256_bm32_splitkv3_merge_kernel( const float* __restrict__ split_tmp_out, const float* __restrict__ split_tmp_row_max, - const float* __restrict__ split_tmp_row_sum, - __half* __restrict__ Out, - float* __restrict__ softmax_lse, - int M) { - const int start_row = blockIdx.x * D256_BM32_PHASE_BLOCK_M; - const int row_local = threadIdx.x >> 3; - const int lane_in_row = threadIdx.x & 7; - const int row = start_row + row_local; - if (row >= M) { - return; - } - - const size_t q_head_linear = static_cast(blockIdx.z); - const size_t first_state = - (q_head_linear * D256_BM32_PHASE_SPLIT_PARTS) * M + row; - const float max_0 = split_tmp_row_max[first_state]; - const float max_1 = split_tmp_row_max[first_state + M]; - const float max_2 = split_tmp_row_max[first_state + 2 * M]; - const float sum_0 = split_tmp_row_sum[first_state]; - const float sum_1 = split_tmp_row_sum[first_state + M]; - const float sum_2 = split_tmp_row_sum[first_state + 2 * M]; - const float merged_max = fmaxf(max_0, fmaxf(max_1, max_2)); - const float weight_0 = - sum_0 > 0.0f ? __expf(fmaxf(max_0 - merged_max, -80.0f)) : 0.0f; - const float weight_1 = - sum_1 > 0.0f ? __expf(fmaxf(max_1 - merged_max, -80.0f)) : 0.0f; - const float weight_2 = - sum_2 > 0.0f ? __expf(fmaxf(max_2 - merged_max, -80.0f)) : 0.0f; - const float denominator = - weight_0 * sum_0 + weight_1 * sum_1 + weight_2 * sum_2; - const float inverse_denominator = - 1.0f / fmaxf(denominator, 1e-24f); - - if (lane_in_row == 0) { - softmax_lse[q_head_linear * M + row] = - merged_max + logf(fmaxf(denominator, 1e-24f)); - } - - const size_t output_base = - (q_head_linear * M + row) * D256_BM32_PHASE_D; - const size_t partial_base_0 = first_state * D256_BM32_PHASE_D; - const size_t partial_base_1 = - (first_state + M) * D256_BM32_PHASE_D; - const size_t partial_base_2 = - (first_state + 2 * M) * D256_BM32_PHASE_D; - #pragma unroll - for (int column = lane_in_row; column < D256_BM32_PHASE_D; column += 8) { - const float numerator = - weight_0 * split_tmp_out[partial_base_0 + column] - + weight_1 * split_tmp_out[partial_base_1 + column] - + weight_2 * split_tmp_out[partial_base_2 + column]; - Out[output_base + column] = - __float2half_rn(numerator * inverse_denominator); - } + const float* __restrict__ split_tmp_row_sum, __half* __restrict__ Out, + float* __restrict__ softmax_lse, int M) { + const int start_row = blockIdx.x * D256_BM32_PHASE_BLOCK_M; + const int row_local = threadIdx.x >> 3; + const int lane_in_row = threadIdx.x & 7; + const int row = start_row + row_local; + if (row >= M) { + return; + } + + const size_t q_head_linear = static_cast(blockIdx.z); + const size_t first_state = + (q_head_linear * D256_BM32_PHASE_SPLIT_PARTS) * M + row; + const float max_0 = split_tmp_row_max[first_state]; + const float max_1 = split_tmp_row_max[first_state + M]; + const float max_2 = split_tmp_row_max[first_state + 2 * M]; + const float sum_0 = split_tmp_row_sum[first_state]; + const float sum_1 = split_tmp_row_sum[first_state + M]; + const float sum_2 = split_tmp_row_sum[first_state + 2 * M]; + const float merged_max = fmaxf(max_0, fmaxf(max_1, max_2)); + const float weight_0 = + sum_0 > 0.0f ? __expf(fmaxf(max_0 - merged_max, -80.0f)) : 0.0f; + const float weight_1 = + sum_1 > 0.0f ? __expf(fmaxf(max_1 - merged_max, -80.0f)) : 0.0f; + const float weight_2 = + sum_2 > 0.0f ? __expf(fmaxf(max_2 - merged_max, -80.0f)) : 0.0f; + const float denominator = + weight_0 * sum_0 + weight_1 * sum_1 + weight_2 * sum_2; + const float inverse_denominator = 1.0f / fmaxf(denominator, 1e-24f); + + if (lane_in_row == 0) { + softmax_lse[q_head_linear * M + row] = + merged_max + logf(fmaxf(denominator, 1e-24f)); + } + + const size_t output_base = (q_head_linear * M + row) * D256_BM32_PHASE_D; + const size_t partial_base_0 = first_state * D256_BM32_PHASE_D; + const size_t partial_base_1 = (first_state + M) * D256_BM32_PHASE_D; + const size_t partial_base_2 = (first_state + 2 * M) * D256_BM32_PHASE_D; +#pragma unroll + for (int column = lane_in_row; column < D256_BM32_PHASE_D; column += 8) { + const float numerator = weight_0 * split_tmp_out[partial_base_0 + column] + + weight_1 * split_tmp_out[partial_base_1 + column] + + weight_2 * split_tmp_out[partial_base_2 + column]; + Out[output_base + column] = + __float2half_rn(numerator * inverse_denominator); + } } void launch_flash_attention_forward_paged_d256_bm32_splitkv3_kernel( - const torch::Tensor& Q, - const torch::Tensor& K_cache, - const torch::Tensor& V_cache, - torch::Tensor& Out, - torch::Tensor& softmax_lse, - torch::Tensor& split_tmp_out, - torch::Tensor& split_tmp_row_max, - torch::Tensor& split_tmp_row_sum, - const torch::Tensor& block_table, - int actual_n, - float softmax_scale, + const torch::Tensor& Q, const torch::Tensor& K_cache, + const torch::Tensor& V_cache, torch::Tensor& Out, + torch::Tensor& softmax_lse, torch::Tensor& split_tmp_out, + torch::Tensor& split_tmp_row_max, torch::Tensor& split_tmp_row_sum, + const torch::Tensor& block_table, int actual_n, float softmax_scale, cudaStream_t stream) { - const int H = Q.size(1); - const int M = Q.size(2); - const int num_kv_heads = K_cache.size(2); - const int max_num_blocks_per_seq = block_table.size(1); - const int64_t k_block_stride = K_cache.stride(0); - const int64_t k_token_stride = K_cache.stride(1); - const int64_t k_head_stride = K_cache.stride(2); - const int64_t v_block_stride = V_cache.stride(0); - const int64_t v_token_stride = V_cache.stride(1); - const int64_t v_head_stride = V_cache.stride(2); - const dim3 partial_grid( - (M + D256_BM32_PHASE_BLOCK_M - 1) / D256_BM32_PHASE_BLOCK_M, - D256_BM32_PHASE_SPLIT_PARTS, - Q.size(0) * H); - const dim3 partial_block(D256_BM32_PHASE_THREADS); - const int64_t total_n_blocks = - (static_cast(actual_n) + D256_BM32_PHASE_BLOCK_N - 1) - / D256_BM32_PHASE_BLOCK_N; - const int64_t last_split_begin = - (D256_BM32_PHASE_SPLIT_PARTS - 1) * total_n_blocks - / D256_BM32_PHASE_SPLIT_PARTS * D256_BM32_PHASE_BLOCK_N; - const bool check_split_empty = - total_n_blocks < D256_BM32_PHASE_SPLIT_PARTS - || last_split_begin > static_cast(actual_n) - M; - if (check_split_empty) { - flash_attention_forward_paged_d256_bm32_splitkv3_partial_kernel - <<>>( - reinterpret_cast(Q.data_ptr()), - reinterpret_cast(K_cache.data_ptr()), - reinterpret_cast(V_cache.data_ptr()), - block_table.data_ptr(), H, M, actual_n, - max_num_blocks_per_seq, num_kv_heads, k_block_stride, - k_token_stride, k_head_stride, v_block_stride, - v_token_stride, v_head_stride, softmax_scale, - split_tmp_out.data_ptr(), - split_tmp_row_max.data_ptr(), - split_tmp_row_sum.data_ptr()); - } else { - flash_attention_forward_paged_d256_bm32_splitkv3_partial_kernel - <<>>( + const int H = Q.size(1); + const int M = Q.size(2); + const int num_kv_heads = K_cache.size(2); + const int max_num_blocks_per_seq = block_table.size(1); + const int64_t k_block_stride = K_cache.stride(0); + const int64_t k_token_stride = K_cache.stride(1); + const int64_t k_head_stride = K_cache.stride(2); + const int64_t v_block_stride = V_cache.stride(0); + const int64_t v_token_stride = V_cache.stride(1); + const int64_t v_head_stride = V_cache.stride(2); + const dim3 partial_grid( + (M + D256_BM32_PHASE_BLOCK_M - 1) / D256_BM32_PHASE_BLOCK_M, + D256_BM32_PHASE_SPLIT_PARTS, Q.size(0) * H); + const dim3 partial_block(D256_BM32_PHASE_THREADS); + const int64_t total_n_blocks = + (static_cast(actual_n) + D256_BM32_PHASE_BLOCK_N - 1) / + D256_BM32_PHASE_BLOCK_N; + const int64_t last_split_begin = + (D256_BM32_PHASE_SPLIT_PARTS - 1) * total_n_blocks / + D256_BM32_PHASE_SPLIT_PARTS * D256_BM32_PHASE_BLOCK_N; + const bool check_split_empty = + total_n_blocks < D256_BM32_PHASE_SPLIT_PARTS || + last_split_begin > static_cast(actual_n) - M; + if (check_split_empty) { + flash_attention_forward_paged_d256_bm32_splitkv3_partial_kernel + <<>>( reinterpret_cast(Q.data_ptr()), reinterpret_cast(K_cache.data_ptr()), reinterpret_cast(V_cache.data_ptr()), - block_table.data_ptr(), H, M, actual_n, - max_num_blocks_per_seq, num_kv_heads, k_block_stride, - k_token_stride, k_head_stride, v_block_stride, - v_token_stride, v_head_stride, softmax_scale, - split_tmp_out.data_ptr(), - split_tmp_row_max.data_ptr(), - split_tmp_row_sum.data_ptr()); - } - - const dim3 merge_grid( - (M + D256_BM32_PHASE_BLOCK_M - 1) / D256_BM32_PHASE_BLOCK_M, - 1, - Q.size(0) * H); - flash_attention_forward_paged_d256_bm32_splitkv3_merge_kernel - <<>>( + block_table.data_ptr(), H, M, actual_n, max_num_blocks_per_seq, + num_kv_heads, k_block_stride, k_token_stride, k_head_stride, + v_block_stride, v_token_stride, v_head_stride, softmax_scale, + split_tmp_out.data_ptr(), + split_tmp_row_max.data_ptr(), + split_tmp_row_sum.data_ptr()); + } else { + flash_attention_forward_paged_d256_bm32_splitkv3_partial_kernel + <<>>( + reinterpret_cast(Q.data_ptr()), + reinterpret_cast(K_cache.data_ptr()), + reinterpret_cast(V_cache.data_ptr()), + block_table.data_ptr(), H, M, actual_n, max_num_blocks_per_seq, + num_kv_heads, k_block_stride, k_token_stride, k_head_stride, + v_block_stride, v_token_stride, v_head_stride, softmax_scale, split_tmp_out.data_ptr(), split_tmp_row_max.data_ptr(), - split_tmp_row_sum.data_ptr(), - reinterpret_cast<__half*>(Out.data_ptr()), - softmax_lse.data_ptr(), M); + split_tmp_row_sum.data_ptr()); + } + + const dim3 merge_grid( + (M + D256_BM32_PHASE_BLOCK_M - 1) / D256_BM32_PHASE_BLOCK_M, 1, + Q.size(0) * H); + flash_attention_forward_paged_d256_bm32_splitkv3_merge_kernel<<< + merge_grid, 256, 0, stream>>>(split_tmp_out.data_ptr(), + split_tmp_row_max.data_ptr(), + split_tmp_row_sum.data_ptr(), + reinterpret_cast<__half*>(Out.data_ptr()), + softmax_lse.data_ptr(), M); } void launcher_flash_attention_forward_paged_d256_bm32_phase( - const torch::Tensor& Q, - const torch::Tensor& K_cache, - const torch::Tensor& V_cache, - torch::Tensor& Out, - torch::Tensor& softmax_lse, - const torch::Tensor& block_table, - const torch::Tensor& seq_lens, - float softmax_scale, - bool is_causal, + const torch::Tensor& Q, const torch::Tensor& K_cache, + const torch::Tensor& V_cache, torch::Tensor& Out, + torch::Tensor& softmax_lse, const torch::Tensor& block_table, + const torch::Tensor& seq_lens, float softmax_scale, bool is_causal, cudaStream_t stream) { - const bool use_all_p = - env_flag_default_enabled("VLLM_FLASH_V100_PREFILL_D256_BM32_ALL_P"); - const bool use_pair_scratch = - use_all_p - && env_flag_default_enabled( - "VLLM_FLASH_V100_PREFILL_D256_BM32_PAIR_SCRATCH"); - - if (is_causal) { - if (use_all_p) { - if (use_pair_scratch) { - launch_flash_attention_forward_paged_d256_bm32_phase_kernel< - true, true, true>( - Q, K_cache, V_cache, Out, softmax_lse, block_table, - seq_lens, softmax_scale, stream); - } else { - launch_flash_attention_forward_paged_d256_bm32_phase_kernel< - true, true, false>( - Q, K_cache, V_cache, Out, softmax_lse, block_table, - seq_lens, softmax_scale, stream); - } - } else { - launch_flash_attention_forward_paged_d256_bm32_phase_kernel< - true, false, false>( - Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, - softmax_scale, stream); - } - return; - } + const bool use_all_p = + env_flag_default_enabled("VLLM_FLASH_V100_PREFILL_D256_BM32_ALL_P"); + const bool use_pair_scratch = + use_all_p && env_flag_default_enabled( + "VLLM_FLASH_V100_PREFILL_D256_BM32_PAIR_SCRATCH"); + if (is_causal) { if (use_all_p) { - if (use_pair_scratch) { - launch_flash_attention_forward_paged_d256_bm32_phase_kernel< - false, true, true>( - Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, - softmax_scale, stream); - } else { - launch_flash_attention_forward_paged_d256_bm32_phase_kernel< - false, true, false>( - Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, - softmax_scale, stream); - } - } else { - launch_flash_attention_forward_paged_d256_bm32_phase_kernel< - false, false, false>( + if (use_pair_scratch) { + launch_flash_attention_forward_paged_d256_bm32_phase_kernel( Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, softmax_scale, stream); - } + } else { + launch_flash_attention_forward_paged_d256_bm32_phase_kernel( + Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, + softmax_scale, stream); + } + } else { + launch_flash_attention_forward_paged_d256_bm32_phase_kernel( + Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, + softmax_scale, stream); + } + return; + } + + if (use_all_p) { + if (use_pair_scratch) { + launch_flash_attention_forward_paged_d256_bm32_phase_kernel( + Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, + softmax_scale, stream); + } else { + launch_flash_attention_forward_paged_d256_bm32_phase_kernel( + Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, + softmax_scale, stream); + } + } else { + launch_flash_attention_forward_paged_d256_bm32_phase_kernel( + Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, + softmax_scale, stream); + } } -template +template void launcher_flash_attention_forward_paged( - const torch::Tensor& Q, - const torch::Tensor& K_cache, - const torch::Tensor& V_cache, - torch::Tensor& Out, - torch::Tensor& softmax_lse, - const torch::Tensor& block_table, - const torch::Tensor& seq_lens, - float softmax_scale, - bool is_causal, - float k_scale, - float v_scale, - int window_size_left, - int window_size_right, - cudaStream_t stream, - const int* bfla_mask_ptr = nullptr, - int bfla_mask_block_n = 0, - int64_t bfla_mask_stride_b = 0, - int64_t bfla_mask_stride_h = 0, - int64_t bfla_mask_stride_q = 0, - int64_t bfla_mask_stride_k = 0 -) { - if constexpr (D == 256 && KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP16) { - const int M = Q.size(2); - const int page_block_size = K_cache.size(1); - const bool use_low_smem = - M > 1 && page_block_size >= 16 && (page_block_size % 16) == 0 - && env_flag_default_enabled( - "VLLM_FLASH_V100_PREFILL_D256_LOW_SMEM"); - if (use_low_smem) { - const bool use_d256_bm32_phase = - env_flag_default_enabled( - "VLLM_FLASH_V100_PREFILL_D256_BM32_PHASE") - && page_block_size == D256_BM32_PHASE_PAGE_BLOCK_SIZE - && M >= D256_BM32_PHASE_BLOCK_M - && bfla_mask_ptr == nullptr && window_size_left < 0 - && window_size_right < 0; - if (use_d256_bm32_phase) { - launcher_flash_attention_forward_paged_d256_bm32_phase( - Q, K_cache, V_cache, Out, softmax_lse, block_table, - seq_lens, softmax_scale, is_causal, stream); - return; - } - const bool use_low_smem_contig_fast = - page_block_size == 16 || - env_flag_enabled("VLLM_FLASH_V100_PREFILL_CONTIG_FAST"); - const bool use_low_smem_scalar_qk = - env_flag_enabled("VLLM_FLASH_V100_PREFILL_D256_SCALAR_QK"); - const bool use_low_smem_bm32 = - env_flag_enabled("VLLM_FLASH_V100_PREFILL_D256_BM32"); - // Default only for the measured hybrid-cache page size. Other - // page sizes retain the prior layout unless explicitly enabled. - const bool use_low_smem_output_stride_268 = - page_block_size == 784 - ? env_flag_default_enabled( - "VLLM_FLASH_V100_PREFILL_D256_OUTPUT_STRIDE_268") - : env_flag_enabled( - "VLLM_FLASH_V100_PREFILL_D256_OUTPUT_STRIDE_268"); - const bool use_sw_pipeline = - page_block_size == 784 - && env_flag_enabled( - "VLLM_FLASH_V100_PREFILL_D256_SOFTWARE_PIPELINE"); - const bool use_sw_pipeline_qk = - page_block_size == 784 - && (use_sw_pipeline - || env_flag_default_enabled( - "VLLM_FLASH_V100_PREFILL_D256_SW_PIPELINE_QK")); - const bool use_sw_pipeline_pv = - page_block_size == 784 - && (use_sw_pipeline - || env_flag_default_enabled( - "VLLM_FLASH_V100_PREFILL_D256_SW_PIPELINE_PV")); - if (use_low_smem_contig_fast) { - if (use_low_smem_scalar_qk) { - if (use_low_smem_bm32) { - launcher_flash_attention_forward_paged_impl< - D, KV_DTYPE, true, true, true, true>( - Q, K_cache, V_cache, Out, softmax_lse, block_table, - seq_lens, bfla_mask_ptr, bfla_mask_block_n, - bfla_mask_stride_b, bfla_mask_stride_h, - bfla_mask_stride_q, bfla_mask_stride_k, - softmax_scale, is_causal, k_scale, v_scale, - window_size_left, window_size_right, stream); - } else { - launcher_flash_attention_forward_paged_impl< - D, KV_DTYPE, true, true, true, false>( - Q, K_cache, V_cache, Out, softmax_lse, block_table, - seq_lens, bfla_mask_ptr, bfla_mask_block_n, - bfla_mask_stride_b, bfla_mask_stride_h, - bfla_mask_stride_q, bfla_mask_stride_k, - softmax_scale, is_causal, k_scale, v_scale, - window_size_left, window_size_right, stream); - } - } else { - if (use_low_smem_bm32) { - launcher_flash_attention_forward_paged_impl< - D, KV_DTYPE, true, true, false, true>( - Q, K_cache, V_cache, Out, softmax_lse, block_table, - seq_lens, bfla_mask_ptr, bfla_mask_block_n, - bfla_mask_stride_b, bfla_mask_stride_h, - bfla_mask_stride_q, bfla_mask_stride_k, - softmax_scale, is_causal, k_scale, v_scale, - window_size_left, window_size_right, stream); - } else { - launcher_flash_attention_forward_paged_impl< - D, KV_DTYPE, true, true, false, false>( - Q, K_cache, V_cache, Out, softmax_lse, block_table, - seq_lens, bfla_mask_ptr, bfla_mask_block_n, - bfla_mask_stride_b, bfla_mask_stride_h, - bfla_mask_stride_q, bfla_mask_stride_k, - softmax_scale, is_causal, k_scale, v_scale, - window_size_left, window_size_right, stream); - } - } - } else { - if (use_low_smem_scalar_qk) { - if (use_low_smem_bm32) { - launcher_flash_attention_forward_paged_impl< - D, KV_DTYPE, true, false, true, true>( - Q, K_cache, V_cache, Out, softmax_lse, block_table, - seq_lens, bfla_mask_ptr, bfla_mask_block_n, - bfla_mask_stride_b, bfla_mask_stride_h, - bfla_mask_stride_q, bfla_mask_stride_k, - softmax_scale, is_causal, k_scale, v_scale, - window_size_left, window_size_right, stream); - } else { - launcher_flash_attention_forward_paged_impl< - D, KV_DTYPE, true, false, true, false>( - Q, K_cache, V_cache, Out, softmax_lse, block_table, - seq_lens, bfla_mask_ptr, bfla_mask_block_n, - bfla_mask_stride_b, bfla_mask_stride_h, - bfla_mask_stride_q, bfla_mask_stride_k, - softmax_scale, is_causal, k_scale, v_scale, - window_size_left, window_size_right, stream); - } - } else { - if (use_low_smem_bm32) { - launcher_flash_attention_forward_paged_impl< - D, KV_DTYPE, true, false, false, true>( - Q, K_cache, V_cache, Out, softmax_lse, block_table, - seq_lens, bfla_mask_ptr, bfla_mask_block_n, - bfla_mask_stride_b, bfla_mask_stride_h, - bfla_mask_stride_q, bfla_mask_stride_k, - softmax_scale, is_causal, k_scale, v_scale, - window_size_left, window_size_right, stream); - } else { - if (use_low_smem_output_stride_268) { - if (use_sw_pipeline_qk) { - if (use_sw_pipeline_pv) { - launcher_flash_attention_forward_paged_impl< - D, KV_DTYPE, true, false, false, - false, true, true, true>( - Q, K_cache, V_cache, Out, - softmax_lse, block_table, seq_lens, - bfla_mask_ptr, bfla_mask_block_n, - bfla_mask_stride_b, - bfla_mask_stride_h, - bfla_mask_stride_q, - bfla_mask_stride_k, softmax_scale, - is_causal, k_scale, v_scale, - window_size_left, - window_size_right, stream); - } else { - launcher_flash_attention_forward_paged_impl< - D, KV_DTYPE, true, false, false, false, - true, true, false>( - Q, K_cache, V_cache, Out, softmax_lse, - block_table, seq_lens, bfla_mask_ptr, - bfla_mask_block_n, bfla_mask_stride_b, - bfla_mask_stride_h, - bfla_mask_stride_q, - bfla_mask_stride_k, softmax_scale, - is_causal, k_scale, v_scale, - window_size_left, window_size_right, - stream); - } - } else if (use_sw_pipeline_pv) { - launcher_flash_attention_forward_paged_impl< - D, KV_DTYPE, true, false, false, false, - true, false, true>( - Q, K_cache, V_cache, Out, softmax_lse, - block_table, seq_lens, bfla_mask_ptr, - bfla_mask_block_n, bfla_mask_stride_b, - bfla_mask_stride_h, bfla_mask_stride_q, - bfla_mask_stride_k, softmax_scale, - is_causal, k_scale, v_scale, - window_size_left, window_size_right, - stream); - } else { - launcher_flash_attention_forward_paged_impl< - D, KV_DTYPE, true, false, false, false, - true>( - Q, K_cache, V_cache, Out, softmax_lse, - block_table, seq_lens, bfla_mask_ptr, - bfla_mask_block_n, bfla_mask_stride_b, - bfla_mask_stride_h, bfla_mask_stride_q, - bfla_mask_stride_k, softmax_scale, - is_causal, k_scale, v_scale, - window_size_left, window_size_right, - stream); - } - } else { - launcher_flash_attention_forward_paged_impl< - D, KV_DTYPE, true, false, false, false>( - Q, K_cache, V_cache, Out, softmax_lse, - block_table, seq_lens, bfla_mask_ptr, - bfla_mask_block_n, bfla_mask_stride_b, - bfla_mask_stride_h, bfla_mask_stride_q, - bfla_mask_stride_k, softmax_scale, is_causal, - k_scale, v_scale, window_size_left, - window_size_right, stream); - } - } - } - } - } else { - launcher_flash_attention_forward_paged_impl< - D, KV_DTYPE, false, false, false, false>( + const torch::Tensor& Q, const torch::Tensor& K_cache, + const torch::Tensor& V_cache, torch::Tensor& Out, + torch::Tensor& softmax_lse, const torch::Tensor& block_table, + const torch::Tensor& seq_lens, float softmax_scale, bool is_causal, + float k_scale, float v_scale, int window_size_left, int window_size_right, + cudaStream_t stream, const int* bfla_mask_ptr = nullptr, + int bfla_mask_block_n = 0, int64_t bfla_mask_stride_b = 0, + int64_t bfla_mask_stride_h = 0, int64_t bfla_mask_stride_q = 0, + int64_t bfla_mask_stride_k = 0, const int* anchor_lens = nullptr, + int anchored_window = 0) { + if (anchor_lens != nullptr && anchored_window > 0) { + TORCH_CHECK(KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP16, + "anchored decode window requires an fp16 KV cache"); + TORCH_CHECK(is_causal, "anchored decode window requires causal attention"); + TORCH_CHECK(window_size_left < 0 && window_size_right < 0, + "anchored decode window cannot be combined with " + "sliding-window attention"); + TORCH_CHECK(bfla_mask_ptr == nullptr, + "anchored decode window cannot be combined with a " + "block sparsity mask"); + launcher_flash_attention_forward_paged_impl( + Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, + bfla_mask_ptr, bfla_mask_block_n, bfla_mask_stride_b, + bfla_mask_stride_h, bfla_mask_stride_q, bfla_mask_stride_k, + softmax_scale, is_causal, k_scale, v_scale, window_size_left, + window_size_right, stream, anchor_lens, anchored_window); + return; + } + if constexpr (D == 256 && KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP16) { + const int M = Q.size(2); + const int page_block_size = K_cache.size(1); + const bool use_low_smem = + M > 1 && page_block_size >= 16 && (page_block_size % 16) == 0 && + env_flag_default_enabled("VLLM_FLASH_V100_PREFILL_D256_LOW_SMEM"); + if (use_low_smem) { + const bool use_d256_bm32_phase = + env_flag_default_enabled("VLLM_FLASH_V100_PREFILL_D256_BM32_PHASE") && + page_block_size == D256_BM32_PHASE_PAGE_BLOCK_SIZE && + M >= D256_BM32_PHASE_BLOCK_M && bfla_mask_ptr == nullptr && + window_size_left < 0 && window_size_right < 0; + if (use_d256_bm32_phase) { + launcher_flash_attention_forward_paged_d256_bm32_phase( + Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, + softmax_scale, is_causal, stream); + return; + } + const bool use_low_smem_contig_fast = + page_block_size == 16 || + env_flag_enabled("VLLM_FLASH_V100_PREFILL_CONTIG_FAST"); + const bool use_low_smem_scalar_qk = + env_flag_enabled("VLLM_FLASH_V100_PREFILL_D256_SCALAR_QK"); + const bool use_low_smem_bm32 = + env_flag_enabled("VLLM_FLASH_V100_PREFILL_D256_BM32"); + // Default only for the measured hybrid-cache page size. Other + // page sizes retain the prior layout unless explicitly enabled. + const bool use_low_smem_output_stride_268 = + page_block_size == 784 + ? env_flag_default_enabled( + "VLLM_FLASH_V100_PREFILL_D256_OUTPUT_STRIDE_268") + : env_flag_enabled( + "VLLM_FLASH_V100_PREFILL_D256_OUTPUT_STRIDE_268"); + const bool use_sw_pipeline = + page_block_size == 784 && + env_flag_enabled("VLLM_FLASH_V100_PREFILL_D256_SOFTWARE_PIPELINE"); + const bool use_sw_pipeline_qk = + page_block_size == 784 && + (use_sw_pipeline || + env_flag_default_enabled( + "VLLM_FLASH_V100_PREFILL_D256_SW_PIPELINE_QK")); + const bool use_sw_pipeline_pv = + page_block_size == 784 && + (use_sw_pipeline || + env_flag_default_enabled( + "VLLM_FLASH_V100_PREFILL_D256_SW_PIPELINE_PV")); + if (use_low_smem_contig_fast) { + if (use_low_smem_scalar_qk) { + if (use_low_smem_bm32) { + launcher_flash_attention_forward_paged_impl( Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, bfla_mask_ptr, bfla_mask_block_n, bfla_mask_stride_b, bfla_mask_stride_h, bfla_mask_stride_q, bfla_mask_stride_k, softmax_scale, is_causal, k_scale, v_scale, window_size_left, window_size_right, stream); - } - } else { - launcher_flash_attention_forward_paged_impl< - D, KV_DTYPE, false, false, false, false>( - Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, - bfla_mask_ptr, bfla_mask_block_n, bfla_mask_stride_b, - bfla_mask_stride_h, bfla_mask_stride_q, bfla_mask_stride_k, - softmax_scale, is_causal, k_scale, v_scale, window_size_left, - window_size_right, stream); - } -} - -at::Tensor flash_attention_prefill_paged_splitkv( - const at::Tensor& q, - const at::Tensor& k_cache, - const at::Tensor& v_cache, - std::optional& out_, - const at::Tensor& block_table, - const at::Tensor& seq_lens, - const float softmax_scale, - const std::string& kv_cache_dtype, - const float k_scale, - const float v_scale, - const bool is_causal, - const int window_size_left, - const int window_size_right, - const int split_kv_tokens, - const int max_seq_len_hint -) { - TORCH_CHECK(q.dtype() == torch::kFloat16, "q must be fp16"); - const int kv_dtype_code = kv_cache_dtype_code_from_string(kv_cache_dtype); - TORCH_CHECK(kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP16, - "split-KV paged prefill supports fp16 KV cache only"); - TORCH_CHECK(k_cache.dtype() == torch::kFloat16, "k_cache must be fp16"); - TORCH_CHECK(v_cache.dtype() == torch::kFloat16, "v_cache must be fp16"); - TORCH_CHECK(q.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda(), - "Tensors must be on CUDA"); - TORCH_CHECK(block_table.is_cuda() && seq_lens.is_cuda(), - "block_table and seq_lens must be CUDA tensors"); - TORCH_CHECK(q.stride(-1) == 1 && k_cache.stride(-1) == 1 && - v_cache.stride(-1) == 1, - "Last dim must be contiguous"); - TORCH_CHECK(k_cache.stride(1) % 8 == 0 && k_cache.stride(2) % 8 == 0 && - v_cache.stride(1) % 8 == 0 && v_cache.stride(2) % 8 == 0, - "Paged KV strides must be divisible by 8 half elements"); - - 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 num_kv_heads = k_cache.size(2); - const int page_block_size = k_cache.size(1); - - TORCH_CHECK(D == 256, "split-KV paged prefill supports D=256 only"); - TORCH_CHECK(H % num_kv_heads == 0, - "num_heads must be divisible by num_kv_heads"); - TORCH_CHECK(M > 1, "split-KV paged prefill requires prefill M > 1"); - TORCH_CHECK(page_block_size >= 16 && (page_block_size % 16) == 0, - "split-KV paged prefill requires page block size multiple of 16"); - TORCH_CHECK(window_size_left >= -1 && window_size_right >= -1, - "window sizes must be >= -1"); - - at::Tensor out_fp16 = out_.has_value() ? out_.value() : torch::zeros_like(q); - auto softmax_lse = 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."); - - const bool use_low_smem_contig_fast = - page_block_size == 16 || - env_flag_enabled("VLLM_FLASH_V100_PREFILL_CONTIG_FAST"); - const bool use_low_smem_scalar_qk = - env_flag_enabled("VLLM_FLASH_V100_PREFILL_D256_SCALAR_QK"); - const int block_n = - use_low_smem_scalar_qk - ? BLOCK_N_256_LOW_SMEM_SCALAR_QK - : BLOCK_N_256_LOW_SMEM; - const int split_tokens_rounded = std::max(split_kv_tokens, block_n); - const int split_kv_tiles = - std::max(1, (split_tokens_rounded + block_n - 1) / block_n); - const int max_hint = std::max(max_seq_len_hint, 1); - const int max_kv_tiles = std::max(1, (max_hint + block_n - 1) / block_n); - const int num_partitions = - std::max(1, (max_kv_tiles + split_kv_tiles - 1) / split_kv_tiles); - - if (num_partitions <= 1) { - launcher_flash_attention_forward_paged<256, flash_v100::KV_CACHE_DTYPE_FP16>( - q, k_cache, v_cache, out_fp16, softmax_lse, block_table, seq_lens, - softmax_scale, is_causal, k_scale, v_scale, window_size_left, - window_size_right, stream); - return out_fp16; - } - - auto tmp_options = torch::dtype(torch::kFloat32).device(q.device()); - auto split_tmp_out = - torch::empty({B, H, num_partitions, M, D}, tmp_options); - auto split_tmp_row_max = - torch::empty({B, H, num_partitions, M}, tmp_options); - auto split_tmp_row_sum = - torch::empty({B, H, num_partitions, M}, tmp_options); - - if (use_low_smem_contig_fast) { - if (use_low_smem_scalar_qk) { - launcher_flash_attention_forward_paged_splitkv_impl< - 256, true, true>( - q, k_cache, v_cache, out_fp16, softmax_lse, split_tmp_out, - split_tmp_row_max, split_tmp_row_sum, block_table, seq_lens, + } else { + launcher_flash_attention_forward_paged_impl( + Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, + bfla_mask_ptr, bfla_mask_block_n, bfla_mask_stride_b, + bfla_mask_stride_h, bfla_mask_stride_q, bfla_mask_stride_k, softmax_scale, is_causal, k_scale, v_scale, window_size_left, - window_size_right, split_kv_tokens, max_seq_len_hint, stream); + window_size_right, stream); + } } else { - launcher_flash_attention_forward_paged_splitkv_impl< - 256, true, false>( - q, k_cache, v_cache, out_fp16, softmax_lse, split_tmp_out, - split_tmp_row_max, split_tmp_row_sum, block_table, seq_lens, + if (use_low_smem_bm32) { + launcher_flash_attention_forward_paged_impl( + Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, + bfla_mask_ptr, bfla_mask_block_n, bfla_mask_stride_b, + bfla_mask_stride_h, bfla_mask_stride_q, bfla_mask_stride_k, + softmax_scale, is_causal, k_scale, v_scale, window_size_left, + window_size_right, stream); + } else { + launcher_flash_attention_forward_paged_impl( + Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, + bfla_mask_ptr, bfla_mask_block_n, bfla_mask_stride_b, + bfla_mask_stride_h, bfla_mask_stride_q, bfla_mask_stride_k, softmax_scale, is_causal, k_scale, v_scale, window_size_left, - window_size_right, split_kv_tokens, max_seq_len_hint, stream); + window_size_right, stream); + } } - } else { + } else { if (use_low_smem_scalar_qk) { - launcher_flash_attention_forward_paged_splitkv_impl< - 256, false, true>( - q, k_cache, v_cache, out_fp16, softmax_lse, split_tmp_out, - split_tmp_row_max, split_tmp_row_sum, block_table, seq_lens, + if (use_low_smem_bm32) { + launcher_flash_attention_forward_paged_impl( + Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, + bfla_mask_ptr, bfla_mask_block_n, bfla_mask_stride_b, + bfla_mask_stride_h, bfla_mask_stride_q, bfla_mask_stride_k, softmax_scale, is_causal, k_scale, v_scale, window_size_left, - window_size_right, split_kv_tokens, max_seq_len_hint, stream); + window_size_right, stream); + } else { + launcher_flash_attention_forward_paged_impl( + Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, + bfla_mask_ptr, bfla_mask_block_n, bfla_mask_stride_b, + bfla_mask_stride_h, bfla_mask_stride_q, bfla_mask_stride_k, + softmax_scale, is_causal, k_scale, v_scale, window_size_left, + window_size_right, stream); + } } else { - launcher_flash_attention_forward_paged_splitkv_impl< - 256, false, false>( - q, k_cache, v_cache, out_fp16, softmax_lse, split_tmp_out, - split_tmp_row_max, split_tmp_row_sum, block_table, seq_lens, + if (use_low_smem_bm32) { + launcher_flash_attention_forward_paged_impl( + Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, + bfla_mask_ptr, bfla_mask_block_n, bfla_mask_stride_b, + bfla_mask_stride_h, bfla_mask_stride_q, bfla_mask_stride_k, softmax_scale, is_causal, k_scale, v_scale, window_size_left, - window_size_right, split_kv_tokens, max_seq_len_hint, stream); + window_size_right, stream); + } else { + if (use_low_smem_output_stride_268) { + if (use_sw_pipeline_qk) { + if (use_sw_pipeline_pv) { + launcher_flash_attention_forward_paged_impl< + D, KV_DTYPE, true, false, false, false, true, true, true>( + Q, K_cache, V_cache, Out, softmax_lse, block_table, + seq_lens, bfla_mask_ptr, bfla_mask_block_n, + bfla_mask_stride_b, bfla_mask_stride_h, + bfla_mask_stride_q, bfla_mask_stride_k, softmax_scale, + is_causal, k_scale, v_scale, window_size_left, + window_size_right, stream); + } else { + launcher_flash_attention_forward_paged_impl< + D, KV_DTYPE, true, false, false, false, true, true, + false>(Q, K_cache, V_cache, Out, softmax_lse, block_table, + seq_lens, bfla_mask_ptr, bfla_mask_block_n, + bfla_mask_stride_b, bfla_mask_stride_h, + bfla_mask_stride_q, bfla_mask_stride_k, + softmax_scale, is_causal, k_scale, v_scale, + window_size_left, window_size_right, stream); + } + } else if (use_sw_pipeline_pv) { + launcher_flash_attention_forward_paged_impl< + D, KV_DTYPE, true, false, false, false, true, false, true>( + Q, K_cache, V_cache, Out, softmax_lse, block_table, + seq_lens, bfla_mask_ptr, bfla_mask_block_n, + bfla_mask_stride_b, bfla_mask_stride_h, bfla_mask_stride_q, + bfla_mask_stride_k, softmax_scale, is_causal, k_scale, + v_scale, window_size_left, window_size_right, stream); + } else { + launcher_flash_attention_forward_paged_impl< + D, KV_DTYPE, true, false, false, false, true>( + Q, K_cache, V_cache, Out, softmax_lse, block_table, + seq_lens, bfla_mask_ptr, bfla_mask_block_n, + bfla_mask_stride_b, bfla_mask_stride_h, bfla_mask_stride_q, + bfla_mask_stride_k, softmax_scale, is_causal, k_scale, + v_scale, window_size_left, window_size_right, stream); + } + } else { + launcher_flash_attention_forward_paged_impl( + Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, + bfla_mask_ptr, bfla_mask_block_n, bfla_mask_stride_b, + bfla_mask_stride_h, bfla_mask_stride_q, bfla_mask_stride_k, + softmax_scale, is_causal, k_scale, v_scale, window_size_left, + window_size_right, stream); + } + } } - } - - return out_fp16; + } + } else { + launcher_flash_attention_forward_paged_impl( + Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, + bfla_mask_ptr, bfla_mask_block_n, bfla_mask_stride_b, + bfla_mask_stride_h, bfla_mask_stride_q, bfla_mask_stride_k, + softmax_scale, is_causal, k_scale, v_scale, window_size_left, + window_size_right, stream); + } + } else { + launcher_flash_attention_forward_paged_impl( + Q, K_cache, V_cache, Out, softmax_lse, block_table, seq_lens, + bfla_mask_ptr, bfla_mask_block_n, bfla_mask_stride_b, + bfla_mask_stride_h, bfla_mask_stride_q, bfla_mask_stride_k, + softmax_scale, is_causal, k_scale, v_scale, window_size_left, + window_size_right, stream); + } } -at::Tensor flash_attention_prefill_paged( - const at::Tensor& q, - const at::Tensor& k_cache, - const at::Tensor& v_cache, - std::optional& out_, - const at::Tensor& block_table, - const at::Tensor& seq_lens, - const float softmax_scale, - const std::string& kv_cache_dtype, - const float k_scale, - const float v_scale, - const bool is_causal, - const int window_size_left, - const int window_size_right -) { - TORCH_CHECK(q.dtype() == torch::kFloat16, "q must be fp16"); - const int kv_dtype_code = kv_cache_dtype_code_from_string(kv_cache_dtype); - TORCH_CHECK(kv_dtype_code >= 0, "Unsupported kv_cache_dtype: ", kv_cache_dtype); - if (kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP16) { - TORCH_CHECK(k_cache.dtype() == torch::kFloat16, "k_cache must be fp16"); - TORCH_CHECK(v_cache.dtype() == torch::kFloat16, "v_cache must be fp16"); +at::Tensor flash_attention_prefill_paged_splitkv( + const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, + std::optional& out_, const at::Tensor& block_table, + const at::Tensor& seq_lens, const float softmax_scale, + const std::string& kv_cache_dtype, const float k_scale, const float v_scale, + const bool is_causal, const int window_size_left, + const int window_size_right, const int split_kv_tokens, + const int max_seq_len_hint) { + TORCH_CHECK(q.dtype() == torch::kFloat16, "q must be fp16"); + const int kv_dtype_code = kv_cache_dtype_code_from_string(kv_cache_dtype); + TORCH_CHECK( + kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP16 || + kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP8_E5M2, + "split-KV paged prefill supports fp16 and fp8_e5m2 KV cache only"); + if (kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP16) { + TORCH_CHECK(k_cache.dtype() == torch::kFloat16, + "fp16 split-KV requires fp16 k_cache"); + TORCH_CHECK(v_cache.dtype() == torch::kFloat16, + "fp16 split-KV requires fp16 v_cache"); + } else { + TORCH_CHECK(k_cache.dtype() == torch::kUInt8, + "fp8_e5m2 split-KV requires uint8 k_cache"); + TORCH_CHECK(v_cache.dtype() == torch::kUInt8, + "fp8_e5m2 split-KV requires uint8 v_cache"); + TORCH_CHECK(k_scale > 0.f && v_scale > 0.f, + "fp8_e5m2 split-KV requires positive K/V scales"); + } + TORCH_CHECK(q.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda(), + "Tensors must be on CUDA"); + TORCH_CHECK(block_table.is_cuda() && seq_lens.is_cuda(), + "block_table and seq_lens must be CUDA tensors"); + TORCH_CHECK( + q.stride(-1) == 1 && k_cache.stride(-1) == 1 && v_cache.stride(-1) == 1, + "Last dim must be contiguous"); + TORCH_CHECK(k_cache.stride(1) % 8 == 0 && k_cache.stride(2) % 8 == 0 && + v_cache.stride(1) % 8 == 0 && v_cache.stride(2) % 8 == 0, + "Paged KV strides must be divisible by 8 half elements"); + + 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 num_kv_heads = k_cache.size(2); + const int page_block_size = k_cache.size(1); + + TORCH_CHECK(D == 256, "split-KV paged prefill supports D=256 only"); + TORCH_CHECK(H % num_kv_heads == 0, + "num_heads must be divisible by num_kv_heads"); + TORCH_CHECK(M > 1, "split-KV paged prefill requires prefill M > 1"); + TORCH_CHECK(page_block_size >= 16 && (page_block_size % 16) == 0, + "split-KV paged prefill requires page block size multiple of 16"); + TORCH_CHECK(window_size_left >= -1 && window_size_right >= -1, + "window sizes must be >= -1"); + + at::Tensor out_fp16 = out_.has_value() ? out_.value() : torch::zeros_like(q); + auto softmax_lse = + 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."); + + const bool fp16_kv = kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP16; + const bool use_low_smem_contig_fast = + fp16_kv && (page_block_size == 16 || + env_flag_enabled("VLLM_FLASH_V100_PREFILL_CONTIG_FAST")); + const bool use_low_smem_scalar_qk = + fp16_kv && env_flag_enabled("VLLM_FLASH_V100_PREFILL_D256_SCALAR_QK"); + const int block_n = + fp16_kv ? (use_low_smem_scalar_qk ? BLOCK_N_256_LOW_SMEM_SCALAR_QK + : BLOCK_N_256_LOW_SMEM) + : KernelConfig<256, false>::BLOCK_N; + const int split_tokens_rounded = std::max(split_kv_tokens, block_n); + const int split_kv_tiles = + std::max(1, (split_tokens_rounded + block_n - 1) / block_n); + const int max_hint = std::max(max_seq_len_hint, 1); + const int max_kv_tiles = std::max(1, (max_hint + block_n - 1) / block_n); + const int num_partitions = + std::max(1, (max_kv_tiles + split_kv_tiles - 1) / split_kv_tiles); + + if (num_partitions <= 1) { + if (fp16_kv) { + launcher_flash_attention_forward_paged<256, + flash_v100::KV_CACHE_DTYPE_FP16>( + q, k_cache, v_cache, out_fp16, softmax_lse, block_table, seq_lens, + softmax_scale, is_causal, k_scale, v_scale, window_size_left, + window_size_right, stream); } else { - TORCH_CHECK(k_cache.dtype() == torch::kUInt8, - "fp8 k_cache must be stored as uint8"); - TORCH_CHECK(v_cache.dtype() == torch::kUInt8, - "fp8 v_cache must be stored as uint8"); - TORCH_CHECK(k_scale > 0.f && v_scale > 0.f, - "fp8 k/v scales must be positive"); + launcher_flash_attention_forward_paged< + 256, flash_v100::KV_CACHE_DTYPE_FP8_E5M2>( + q, k_cache, v_cache, out_fp16, softmax_lse, block_table, seq_lens, + softmax_scale, is_causal, k_scale, v_scale, window_size_left, + window_size_right, stream); } - TORCH_CHECK(q.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda(), - "Tensors must be on CUDA"); - TORCH_CHECK(block_table.is_cuda() && seq_lens.is_cuda(), - "block_table and seq_lens must be CUDA tensors"); - TORCH_CHECK(q.stride(-1) == 1 && k_cache.stride(-1) == 1 && - v_cache.stride(-1) == 1, - "Last dim must be contiguous"); - TORCH_CHECK(k_cache.stride(1) % 8 == 0 && k_cache.stride(2) % 8 == 0 && - v_cache.stride(1) % 8 == 0 && v_cache.stride(2) % 8 == 0, - "Paged KV strides must be divisible by 8 half elements"); - - 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 num_kv_heads = k_cache.size(2); - - TORCH_CHECK(D <= 256 && D % 8 == 0 && D % 2 == 0, - "D must be even, <=256, multiple of 8"); - TORCH_CHECK(H % num_kv_heads == 0, - "num_heads must be divisible by num_kv_heads"); - TORCH_CHECK(window_size_left >= -1 && window_size_right >= -1, - "window sizes must be >= -1"); - - at::Tensor out_fp16 = out_.has_value() ? out_.value() : torch::zeros_like(q); - auto softmax_lse = 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."); - - #define LAUNCH_PAGED_TYPED(HDIM, KV_DTYPE_CODE) \ - launcher_flash_attention_forward_paged( \ - q, k_cache, v_cache, out_fp16, softmax_lse, block_table, seq_lens, \ - softmax_scale, is_causal, k_scale, v_scale, window_size_left, \ - window_size_right, stream) - - #define LAUNCH_PAGED_BY_KV(HDIM) \ - do { \ - switch (kv_dtype_code) { \ - case flash_v100::KV_CACHE_DTYPE_FP16: \ - LAUNCH_PAGED_TYPED(HDIM, flash_v100::KV_CACHE_DTYPE_FP16); \ - break; \ - case flash_v100::KV_CACHE_DTYPE_FP8_E4M3: \ - LAUNCH_PAGED_TYPED(HDIM, flash_v100::KV_CACHE_DTYPE_FP8_E4M3); \ - break; \ - case flash_v100::KV_CACHE_DTYPE_FP8_E5M2: \ - LAUNCH_PAGED_TYPED(HDIM, flash_v100::KV_CACHE_DTYPE_FP8_E5M2); \ - break; \ - default: \ - TORCH_CHECK(false, "Unsupported kv_cache_dtype: ", kv_cache_dtype); \ - } \ - } while (0) - - switch (D) { - case 16: - LAUNCH_PAGED_BY_KV(16); - break; - case 32: - LAUNCH_PAGED_BY_KV(32); - break; - case 64: - LAUNCH_PAGED_BY_KV(64); - break; - case 128: - LAUNCH_PAGED_BY_KV(128); - break; - case 256: - LAUNCH_PAGED_BY_KV(256); - break; - default: - TORCH_CHECK(false, "Unsupported D: ", D); + return out_fp16; + } + + auto tmp_options = torch::dtype(torch::kFloat32).device(q.device()); + auto split_tmp_out = torch::empty({B, H, num_partitions, M, D}, tmp_options); + auto split_tmp_row_max = torch::empty({B, H, num_partitions, M}, tmp_options); + auto split_tmp_row_sum = torch::empty({B, H, num_partitions, M}, tmp_options); + + if (!fp16_kv) { + launcher_flash_attention_forward_paged_splitkv_impl< + 256, flash_v100::KV_CACHE_DTYPE_FP8_E5M2, false, false, false>( + q, k_cache, v_cache, out_fp16, softmax_lse, split_tmp_out, + split_tmp_row_max, split_tmp_row_sum, block_table, seq_lens, + softmax_scale, is_causal, k_scale, v_scale, window_size_left, + window_size_right, split_kv_tokens, max_seq_len_hint, stream); + } else if (use_low_smem_contig_fast) { + if (use_low_smem_scalar_qk) { + launcher_flash_attention_forward_paged_splitkv_impl< + 256, flash_v100::KV_CACHE_DTYPE_FP16, true, true, true>( + q, k_cache, v_cache, out_fp16, softmax_lse, split_tmp_out, + split_tmp_row_max, split_tmp_row_sum, block_table, seq_lens, + softmax_scale, is_causal, k_scale, v_scale, window_size_left, + window_size_right, split_kv_tokens, max_seq_len_hint, stream); + } else { + launcher_flash_attention_forward_paged_splitkv_impl< + 256, flash_v100::KV_CACHE_DTYPE_FP16, true, true, false>( + q, k_cache, v_cache, out_fp16, softmax_lse, split_tmp_out, + split_tmp_row_max, split_tmp_row_sum, block_table, seq_lens, + softmax_scale, is_causal, k_scale, v_scale, window_size_left, + window_size_right, split_kv_tokens, max_seq_len_hint, stream); + } + } else { + if (use_low_smem_scalar_qk) { + launcher_flash_attention_forward_paged_splitkv_impl< + 256, flash_v100::KV_CACHE_DTYPE_FP16, true, false, true>( + q, k_cache, v_cache, out_fp16, softmax_lse, split_tmp_out, + split_tmp_row_max, split_tmp_row_sum, block_table, seq_lens, + softmax_scale, is_causal, k_scale, v_scale, window_size_left, + window_size_right, split_kv_tokens, max_seq_len_hint, stream); + } else { + launcher_flash_attention_forward_paged_splitkv_impl< + 256, flash_v100::KV_CACHE_DTYPE_FP16, true, false, false>( + q, k_cache, v_cache, out_fp16, softmax_lse, split_tmp_out, + split_tmp_row_max, split_tmp_row_sum, block_table, seq_lens, + softmax_scale, is_causal, k_scale, v_scale, window_size_left, + window_size_right, split_kv_tokens, max_seq_len_hint, stream); } + } - #undef LAUNCH_PAGED_BY_KV - #undef LAUNCH_PAGED_TYPED + return out_fp16; +} - return out_fp16; +at::Tensor flash_attention_prefill_paged( + const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, + std::optional& out_, const at::Tensor& block_table, + const at::Tensor& seq_lens, const float softmax_scale, + const std::string& kv_cache_dtype, const float k_scale, const float v_scale, + const bool is_causal, const int window_size_left, + const int window_size_right, const std::optional& anchor_lens, + const int64_t anchored_window) { + TORCH_CHECK(q.dtype() == torch::kFloat16, "q must be fp16"); + const int kv_dtype_code = kv_cache_dtype_code_from_string(kv_cache_dtype); + TORCH_CHECK(kv_dtype_code >= 0, + "Unsupported kv_cache_dtype: ", kv_cache_dtype); + if (kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP16) { + TORCH_CHECK(k_cache.dtype() == torch::kFloat16, "k_cache must be fp16"); + TORCH_CHECK(v_cache.dtype() == torch::kFloat16, "v_cache must be fp16"); + } else { + TORCH_CHECK(k_cache.dtype() == torch::kUInt8, + "fp8 k_cache must be stored as uint8"); + TORCH_CHECK(v_cache.dtype() == torch::kUInt8, + "fp8 v_cache must be stored as uint8"); + TORCH_CHECK(k_scale > 0.f && v_scale > 0.f, + "fp8 k/v scales must be positive"); + } + TORCH_CHECK(q.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda(), + "Tensors must be on CUDA"); + TORCH_CHECK(block_table.is_cuda() && seq_lens.is_cuda(), + "block_table and seq_lens must be CUDA tensors"); + TORCH_CHECK( + q.stride(-1) == 1 && k_cache.stride(-1) == 1 && v_cache.stride(-1) == 1, + "Last dim must be contiguous"); + TORCH_CHECK(k_cache.stride(1) % 8 == 0 && k_cache.stride(2) % 8 == 0 && + v_cache.stride(1) % 8 == 0 && v_cache.stride(2) % 8 == 0, + "Paged KV strides must be divisible by 8 half elements"); + + 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 num_kv_heads = k_cache.size(2); + + TORCH_CHECK(D <= 256 && D % 8 == 0 && D % 2 == 0, + "D must be even, <=256, multiple of 8"); + TORCH_CHECK(H % num_kv_heads == 0, + "num_heads must be divisible by num_kv_heads"); + TORCH_CHECK(window_size_left >= -1 && window_size_right >= -1, + "window sizes must be >= -1"); + + TORCH_CHECK(anchored_window >= 0 && anchored_window <= INT_MAX, + "anchored_window must fit in a non-negative int32"); + TORCH_CHECK((anchored_window > 0) == anchor_lens.has_value(), + "anchor_lens and a positive anchored_window must be provided " + "together"); + const bool use_anchored = anchor_lens.has_value(); + const int* anchor_lens_ptr = nullptr; + if (use_anchored) { + const at::Tensor& anchors = anchor_lens.value(); + TORCH_CHECK(anchors.is_cuda(), "anchor_lens must be on CUDA"); + TORCH_CHECK(anchors.device() == q.device(), + "anchor_lens must be on the q device"); + TORCH_CHECK(anchors.dtype() == torch::kInt32, "anchor_lens must be int32"); + TORCH_CHECK(anchors.dim() == 1 && anchors.size(0) == B, + "anchor_lens must have shape [B]"); + TORCH_CHECK(anchors.is_contiguous(), "anchor_lens must be contiguous"); + TORCH_CHECK(kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP16, + "anchored decode window requires an fp16 KV cache"); + TORCH_CHECK(is_causal, "anchored decode window requires causal attention"); + TORCH_CHECK(window_size_left == -1 && window_size_right == -1, + "anchored decode window cannot be combined with " + "sliding-window attention"); + anchor_lens_ptr = anchors.data_ptr(); + } + + at::Tensor out_fp16 = out_.has_value() ? out_.value() : torch::zeros_like(q); + auto softmax_lse = + 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."); + +#define LAUNCH_PAGED_TYPED(HDIM, KV_DTYPE_CODE) \ + launcher_flash_attention_forward_paged( \ + q, k_cache, v_cache, out_fp16, softmax_lse, block_table, seq_lens, \ + softmax_scale, is_causal, k_scale, v_scale, window_size_left, \ + window_size_right, stream, nullptr, 0, 0, 0, 0, 0, anchor_lens_ptr, \ + static_cast(anchored_window)) + +#define LAUNCH_PAGED_BY_KV(HDIM) \ + do { \ + switch (kv_dtype_code) { \ + case flash_v100::KV_CACHE_DTYPE_FP16: \ + LAUNCH_PAGED_TYPED(HDIM, flash_v100::KV_CACHE_DTYPE_FP16); \ + break; \ + case flash_v100::KV_CACHE_DTYPE_FP8_E4M3: \ + LAUNCH_PAGED_TYPED(HDIM, flash_v100::KV_CACHE_DTYPE_FP8_E4M3); \ + break; \ + case flash_v100::KV_CACHE_DTYPE_FP8_E5M2: \ + LAUNCH_PAGED_TYPED(HDIM, flash_v100::KV_CACHE_DTYPE_FP8_E5M2); \ + break; \ + default: \ + TORCH_CHECK(false, "Unsupported kv_cache_dtype: ", kv_cache_dtype); \ + } \ + } while (0) + + switch (D) { + case 16: + LAUNCH_PAGED_BY_KV(16); + break; + case 32: + LAUNCH_PAGED_BY_KV(32); + break; + case 64: + LAUNCH_PAGED_BY_KV(64); + break; + case 128: + LAUNCH_PAGED_BY_KV(128); + break; + case 256: + LAUNCH_PAGED_BY_KV(256); + break; + default: + TORCH_CHECK(false, "Unsupported D: ", D); + } + +#undef LAUNCH_PAGED_BY_KV +#undef LAUNCH_PAGED_TYPED + + return out_fp16; } std::vector flash_attention_prefill_paged_d256_bm32_allp_pair_scratch( - const at::Tensor& q, - const at::Tensor& k_cache, - const at::Tensor& v_cache, - std::optional& out_, - std::optional& softmax_lse_, - const at::Tensor& block_table, - const at::Tensor& seq_lens, - const float softmax_scale -) { - TORCH_CHECK(q.dim() == 4, "q must have shape [B, H, M, D]"); - TORCH_CHECK(k_cache.dim() == 4 && v_cache.dim() == 4, - "k_cache and v_cache must have shape [blocks, page, Hkv, D]"); - TORCH_CHECK(block_table.dim() == 2, - "block_table must have shape [B, max_num_blocks]"); - TORCH_CHECK(seq_lens.dim() == 1, "seq_lens must have shape [B]"); - TORCH_CHECK(q.dtype() == torch::kFloat16, "q must be fp16"); - TORCH_CHECK(k_cache.dtype() == torch::kFloat16, - "k_cache must be fp16"); - TORCH_CHECK(v_cache.dtype() == torch::kFloat16, - "v_cache must be fp16"); - 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.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda(), - "q, k_cache, and v_cache must be CUDA tensors"); - TORCH_CHECK(block_table.is_cuda() && seq_lens.is_cuda(), - "block_table and seq_lens must be CUDA tensors"); - TORCH_CHECK( - q.device() == k_cache.device() && q.device() == v_cache.device() - && q.device() == block_table.device() && q.device() == seq_lens.device(), - "all input tensors must be on the q device"); - TORCH_CHECK(q.is_contiguous(), "q must be contiguous [B, H, M, D]"); - TORCH_CHECK(block_table.is_contiguous(), "block_table must be contiguous"); - TORCH_CHECK(seq_lens.is_contiguous(), "seq_lens must be contiguous"); - TORCH_CHECK(k_cache.sizes() == v_cache.sizes(), - "k_cache and v_cache must have the same shape"); - TORCH_CHECK(k_cache.stride(-1) == 1 && v_cache.stride(-1) == 1, - "K/V head dimensions must be contiguous"); - TORCH_CHECK(k_cache.stride(1) % 8 == 0 && k_cache.stride(2) % 8 == 0 - && v_cache.stride(1) % 8 == 0 - && v_cache.stride(2) % 8 == 0, - "K/V strides must be divisible by 8 half elements"); - - 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 num_kv_heads = k_cache.size(2); - - TORCH_CHECK(B > 0 && H > 0, "q batch and head dimensions must be positive"); - TORCH_CHECK(M >= D256_BM32_PHASE_BLOCK_M, - "fixed BM32 prefill requires M >= ", - D256_BM32_PHASE_BLOCK_M); - TORCH_CHECK(D == D256_BM32_PHASE_D, - "fixed BM32 prefill requires D=", D256_BM32_PHASE_D); - TORCH_CHECK(k_cache.size(0) > 0 && num_kv_heads > 0, - "K/V cache block and head dimensions must be positive"); - TORCH_CHECK(k_cache.size(1) == D256_BM32_PHASE_PAGE_BLOCK_SIZE, - "fixed BM32 prefill requires page size ", - D256_BM32_PHASE_PAGE_BLOCK_SIZE); - TORCH_CHECK(k_cache.size(3) == D256_BM32_PHASE_D, - "K/V cache head dimension must be D=", - D256_BM32_PHASE_D); - TORCH_CHECK(H % num_kv_heads == 0, - "num_q_heads must be divisible by num_kv_heads"); - TORCH_CHECK(block_table.size(0) == B && seq_lens.size(0) == B, - "block_table and seq_lens batch dimensions must equal q batch"); - TORCH_CHECK(block_table.size(1) > 0, - "block_table must provide at least one page index"); - - if (out_.has_value()) { - const at::Tensor& out = out_.value(); - TORCH_CHECK(out.is_cuda() && out.device() == q.device(), - "out must be a CUDA tensor on the q device"); - TORCH_CHECK(out.dtype() == torch::kFloat16, "out must be fp16"); - TORCH_CHECK(out.sizes() == q.sizes(), "out must have shape [B, H, M, D]"); - TORCH_CHECK(out.is_contiguous(), "out must be contiguous [B, H, M, D]"); - } - if (softmax_lse_.has_value()) { - const at::Tensor& softmax_lse = softmax_lse_.value(); - TORCH_CHECK(softmax_lse.is_cuda() && softmax_lse.device() == q.device(), - "softmax_lse must be a CUDA tensor on the q device"); - TORCH_CHECK(softmax_lse.dtype() == torch::kFloat32, - "softmax_lse must be fp32"); - TORCH_CHECK(softmax_lse.dim() == 3 && softmax_lse.size(0) == B - && softmax_lse.size(1) == H && softmax_lse.size(2) == M, - "softmax_lse must have shape [B, H, M]"); - TORCH_CHECK(softmax_lse.is_contiguous(), - "softmax_lse must be contiguous [B, H, M]"); - } - - c10::cuda::CUDAGuard device_guard(q.device()); - auto stream = at::cuda::getCurrentCUDAStream().stream(); - cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; - const cudaError_t capture_error = cudaStreamIsCapturing(stream, &capture_status); - TORCH_CHECK(capture_error == cudaSuccess, "cudaStreamIsCapturing failed: ", - cudaGetErrorString(capture_error)); - TORCH_CHECK(capture_status == cudaStreamCaptureStatusNone, - "fixed BM32 prefill CUDA graph capture is not validated"); - - auto props = at::cuda::getCurrentDeviceProperties(); - TORCH_CHECK(props->major == 7 && props->minor == 0, - "fixed BM32 prefill supports only SM70 GPUs"); - - at::Tensor out_fp16 = out_.has_value() ? out_.value() : torch::empty_like(q); - at::Tensor softmax_lse = - softmax_lse_.has_value() - ? softmax_lse_.value() - : torch::empty({B, H, M}, - torch::dtype(torch::kFloat32).device(q.device())); - - launch_flash_attention_forward_paged_d256_bm32_phase_kernel( - q, k_cache, v_cache, out_fp16, softmax_lse, block_table, seq_lens, - softmax_scale, stream); - - return {out_fp16, softmax_lse}; + const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, + std::optional& out_, std::optional& softmax_lse_, + const at::Tensor& block_table, const at::Tensor& seq_lens, + const float softmax_scale) { + TORCH_CHECK(q.dim() == 4, "q must have shape [B, H, M, D]"); + TORCH_CHECK(k_cache.dim() == 4 && v_cache.dim() == 4, + "k_cache and v_cache must have shape [blocks, page, Hkv, D]"); + TORCH_CHECK(block_table.dim() == 2, + "block_table must have shape [B, max_num_blocks]"); + TORCH_CHECK(seq_lens.dim() == 1, "seq_lens must have shape [B]"); + TORCH_CHECK(q.dtype() == torch::kFloat16, "q must be fp16"); + TORCH_CHECK(k_cache.dtype() == torch::kFloat16, "k_cache must be fp16"); + TORCH_CHECK(v_cache.dtype() == torch::kFloat16, "v_cache must be fp16"); + 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.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda(), + "q, k_cache, and v_cache must be CUDA tensors"); + TORCH_CHECK(block_table.is_cuda() && seq_lens.is_cuda(), + "block_table and seq_lens must be CUDA tensors"); + TORCH_CHECK( + q.device() == k_cache.device() && q.device() == v_cache.device() && + q.device() == block_table.device() && q.device() == seq_lens.device(), + "all input tensors must be on the q device"); + TORCH_CHECK(q.is_contiguous(), "q must be contiguous [B, H, M, D]"); + TORCH_CHECK(block_table.is_contiguous(), "block_table must be contiguous"); + TORCH_CHECK(seq_lens.is_contiguous(), "seq_lens must be contiguous"); + TORCH_CHECK(k_cache.sizes() == v_cache.sizes(), + "k_cache and v_cache must have the same shape"); + TORCH_CHECK(k_cache.stride(-1) == 1 && v_cache.stride(-1) == 1, + "K/V head dimensions must be contiguous"); + TORCH_CHECK(k_cache.stride(1) % 8 == 0 && k_cache.stride(2) % 8 == 0 && + v_cache.stride(1) % 8 == 0 && v_cache.stride(2) % 8 == 0, + "K/V strides must be divisible by 8 half elements"); + + 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 num_kv_heads = k_cache.size(2); + + TORCH_CHECK(B > 0 && H > 0, "q batch and head dimensions must be positive"); + TORCH_CHECK(M >= D256_BM32_PHASE_BLOCK_M, + "fixed BM32 prefill requires M >= ", D256_BM32_PHASE_BLOCK_M); + TORCH_CHECK(D == D256_BM32_PHASE_D, + "fixed BM32 prefill requires D=", D256_BM32_PHASE_D); + TORCH_CHECK(k_cache.size(0) > 0 && num_kv_heads > 0, + "K/V cache block and head dimensions must be positive"); + TORCH_CHECK(k_cache.size(1) == D256_BM32_PHASE_PAGE_BLOCK_SIZE, + "fixed BM32 prefill requires page size ", + D256_BM32_PHASE_PAGE_BLOCK_SIZE); + TORCH_CHECK(k_cache.size(3) == D256_BM32_PHASE_D, + "K/V cache head dimension must be D=", D256_BM32_PHASE_D); + TORCH_CHECK(H % num_kv_heads == 0, + "num_q_heads must be divisible by num_kv_heads"); + TORCH_CHECK(block_table.size(0) == B && seq_lens.size(0) == B, + "block_table and seq_lens batch dimensions must equal q batch"); + TORCH_CHECK(block_table.size(1) > 0, + "block_table must provide at least one page index"); + + if (out_.has_value()) { + const at::Tensor& out = out_.value(); + TORCH_CHECK(out.is_cuda() && out.device() == q.device(), + "out must be a CUDA tensor on the q device"); + TORCH_CHECK(out.dtype() == torch::kFloat16, "out must be fp16"); + TORCH_CHECK(out.sizes() == q.sizes(), "out must have shape [B, H, M, D]"); + TORCH_CHECK(out.is_contiguous(), "out must be contiguous [B, H, M, D]"); + } + if (softmax_lse_.has_value()) { + const at::Tensor& softmax_lse = softmax_lse_.value(); + TORCH_CHECK(softmax_lse.is_cuda() && softmax_lse.device() == q.device(), + "softmax_lse must be a CUDA tensor on the q device"); + TORCH_CHECK(softmax_lse.dtype() == torch::kFloat32, + "softmax_lse must be fp32"); + TORCH_CHECK(softmax_lse.dim() == 3 && softmax_lse.size(0) == B && + softmax_lse.size(1) == H && softmax_lse.size(2) == M, + "softmax_lse must have shape [B, H, M]"); + TORCH_CHECK(softmax_lse.is_contiguous(), + "softmax_lse must be contiguous [B, H, M]"); + } + + c10::cuda::CUDAGuard device_guard(q.device()); + auto stream = at::cuda::getCurrentCUDAStream().stream(); + cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; + const cudaError_t capture_error = + cudaStreamIsCapturing(stream, &capture_status); + TORCH_CHECK(capture_error == cudaSuccess, "cudaStreamIsCapturing failed: ", + cudaGetErrorString(capture_error)); + TORCH_CHECK(capture_status == cudaStreamCaptureStatusNone, + "fixed BM32 prefill CUDA graph capture is not validated"); + + auto props = at::cuda::getCurrentDeviceProperties(); + TORCH_CHECK(props->major == 7 && props->minor == 0, + "fixed BM32 prefill supports only SM70 GPUs"); + + at::Tensor out_fp16 = out_.has_value() ? out_.value() : torch::empty_like(q); + at::Tensor softmax_lse = + softmax_lse_.has_value() + ? softmax_lse_.value() + : torch::empty({B, H, M}, + torch::dtype(torch::kFloat32).device(q.device())); + + launch_flash_attention_forward_paged_d256_bm32_phase_kernel( + q, k_cache, v_cache, out_fp16, softmax_lse, block_table, seq_lens, + softmax_scale, stream); + + return {out_fp16, softmax_lse}; } std::vector flash_attention_prefill_paged_d256_bm32_allp_pair_scratch_splitkv3( - const at::Tensor& q, - const at::Tensor& k_cache, - const at::Tensor& v_cache, - std::optional& out_, - std::optional& softmax_lse_, - at::Tensor& split_tmp_out, - at::Tensor& split_tmp_row_max, - at::Tensor& split_tmp_row_sum, - const at::Tensor& block_table, - const int64_t actual_n, - const float softmax_scale -) { - TORCH_CHECK(q.dim() == 4, "q must have shape [B, H, M, D]"); - TORCH_CHECK(k_cache.dim() == 4 && v_cache.dim() == 4, - "k_cache and v_cache must have shape [blocks, page, Hkv, D]"); - TORCH_CHECK(block_table.dim() == 2, - "block_table must have shape [B, max_num_blocks]"); - TORCH_CHECK(q.dtype() == torch::kFloat16, "q must be fp16"); - TORCH_CHECK(k_cache.dtype() == torch::kFloat16, - "k_cache must be fp16"); - TORCH_CHECK(v_cache.dtype() == torch::kFloat16, - "v_cache must be fp16"); - TORCH_CHECK(block_table.dtype() == torch::kInt32, - "block_table must be int32"); - TORCH_CHECK(q.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda(), - "q, k_cache, and v_cache must be CUDA tensors"); - TORCH_CHECK(block_table.is_cuda(), "block_table must be a CUDA tensor"); - TORCH_CHECK( - q.device() == k_cache.device() && q.device() == v_cache.device() - && q.device() == block_table.device(), - "all input tensors must be on the q device"); - TORCH_CHECK(q.is_contiguous(), "q must be contiguous [B, H, M, D]"); - TORCH_CHECK(block_table.is_contiguous(), "block_table must be contiguous"); - TORCH_CHECK(k_cache.sizes() == v_cache.sizes(), - "k_cache and v_cache must have the same shape"); - TORCH_CHECK(k_cache.stride(-1) == 1 && v_cache.stride(-1) == 1, - "K/V head dimensions must be contiguous"); - TORCH_CHECK(k_cache.stride(1) % 8 == 0 && k_cache.stride(2) % 8 == 0 - && v_cache.stride(1) % 8 == 0 - && v_cache.stride(2) % 8 == 0, - "K/V strides must be divisible by 8 half elements"); - - 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 num_kv_heads = k_cache.size(2); - - TORCH_CHECK(B == 1, - "fixed BM32 split-KV host-length path requires B=1"); - TORCH_CHECK(H > 0, "q head dimension must be positive"); - TORCH_CHECK(M >= D256_BM32_PHASE_BLOCK_M, - "fixed BM32 split-KV prefill requires M >= ", - D256_BM32_PHASE_BLOCK_M); - TORCH_CHECK(D == D256_BM32_PHASE_D, - "fixed BM32 split-KV prefill requires D=", - D256_BM32_PHASE_D); - TORCH_CHECK(k_cache.size(0) > 0 && num_kv_heads > 0, - "K/V cache block and head dimensions must be positive"); - TORCH_CHECK(k_cache.size(1) == D256_BM32_PHASE_PAGE_BLOCK_SIZE, - "fixed BM32 split-KV prefill requires page size ", - D256_BM32_PHASE_PAGE_BLOCK_SIZE); - TORCH_CHECK(k_cache.size(3) == D256_BM32_PHASE_D, - "K/V cache head dimension must be D=", - D256_BM32_PHASE_D); - TORCH_CHECK(H % num_kv_heads == 0, - "num_q_heads must be divisible by num_kv_heads"); - TORCH_CHECK(block_table.size(0) == B, - "block_table batch dimension must equal q batch"); - TORCH_CHECK(actual_n >= M && actual_n <= 2147483647LL, - "actual_n must satisfy M <= actual_n <= INT_MAX"); - const int64_t required_pages = - (actual_n + D256_BM32_PHASE_PAGE_BLOCK_SIZE - 1) - / D256_BM32_PHASE_PAGE_BLOCK_SIZE; - TORCH_CHECK(block_table.size(1) >= required_pages, - "block_table does not cover actual_n host length"); - - if (out_.has_value()) { - const at::Tensor& out = out_.value(); - TORCH_CHECK(out.is_cuda() && out.device() == q.device(), - "out must be a CUDA tensor on the q device"); - TORCH_CHECK(out.dtype() == torch::kFloat16, "out must be fp16"); - TORCH_CHECK(out.sizes() == q.sizes(), "out must have shape [B, H, M, D]"); - TORCH_CHECK(out.is_contiguous(), "out must be contiguous [B, H, M, D]"); - } - if (softmax_lse_.has_value()) { - const at::Tensor& softmax_lse = softmax_lse_.value(); - TORCH_CHECK(softmax_lse.is_cuda() && softmax_lse.device() == q.device(), - "softmax_lse must be a CUDA tensor on the q device"); - TORCH_CHECK(softmax_lse.dtype() == torch::kFloat32, - "softmax_lse must be fp32"); - TORCH_CHECK(softmax_lse.dim() == 3 && softmax_lse.size(0) == B - && softmax_lse.size(1) == H && softmax_lse.size(2) == M, - "softmax_lse must have shape [B, H, M]"); - TORCH_CHECK(softmax_lse.is_contiguous(), - "softmax_lse must be contiguous [B, H, M]"); - } - - const std::vector expected_partial_out = { - B, H, D256_BM32_PHASE_SPLIT_PARTS, M, D}; - const std::vector expected_partial_rows = { - B, H, D256_BM32_PHASE_SPLIT_PARTS, M}; - for (const at::Tensor* workspace : { - &split_tmp_out, &split_tmp_row_max, &split_tmp_row_sum}) { - TORCH_CHECK(workspace->is_cuda() && workspace->device() == q.device(), - "split-KV workspaces must be CUDA tensors on the q device"); - TORCH_CHECK(workspace->dtype() == torch::kFloat32, - "split-KV workspaces must be fp32"); - TORCH_CHECK(workspace->is_contiguous(), - "split-KV workspaces must be contiguous"); - } - TORCH_CHECK(split_tmp_out.sizes() == expected_partial_out, - "split_tmp_out must have shape [B, H, 3, M, D]"); - TORCH_CHECK(split_tmp_row_max.sizes() == expected_partial_rows, - "split_tmp_row_max must have shape [B, H, 3, M]"); - TORCH_CHECK(split_tmp_row_sum.sizes() == expected_partial_rows, - "split_tmp_row_sum must have shape [B, H, 3, M]"); - - c10::cuda::CUDAGuard device_guard(q.device()); - auto stream = at::cuda::getCurrentCUDAStream().stream(); - cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; - const cudaError_t capture_error = cudaStreamIsCapturing(stream, &capture_status); - TORCH_CHECK(capture_error == cudaSuccess, "cudaStreamIsCapturing failed: ", - cudaGetErrorString(capture_error)); - TORCH_CHECK(capture_status == cudaStreamCaptureStatusNone, - "fixed BM32 split-KV CUDA graph capture is not validated"); - - auto props = at::cuda::getCurrentDeviceProperties(); - TORCH_CHECK(props->major == 7 && props->minor == 0, - "fixed BM32 split-KV prefill supports only SM70 GPUs"); - - at::Tensor out_fp16 = out_.has_value() ? out_.value() : torch::empty_like(q); - at::Tensor softmax_lse = - softmax_lse_.has_value() - ? softmax_lse_.value() - : torch::empty({B, H, M}, - torch::dtype(torch::kFloat32).device(q.device())); - launch_flash_attention_forward_paged_d256_bm32_splitkv3_kernel( - q, k_cache, v_cache, out_fp16, softmax_lse, - split_tmp_out, split_tmp_row_max, split_tmp_row_sum, - block_table, static_cast(actual_n), - softmax_scale, stream); - - return {out_fp16, softmax_lse}; + const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, + std::optional& out_, std::optional& softmax_lse_, + at::Tensor& split_tmp_out, at::Tensor& split_tmp_row_max, + at::Tensor& split_tmp_row_sum, const at::Tensor& block_table, + const int64_t actual_n, const float softmax_scale) { + TORCH_CHECK(q.dim() == 4, "q must have shape [B, H, M, D]"); + TORCH_CHECK(k_cache.dim() == 4 && v_cache.dim() == 4, + "k_cache and v_cache must have shape [blocks, page, Hkv, D]"); + TORCH_CHECK(block_table.dim() == 2, + "block_table must have shape [B, max_num_blocks]"); + TORCH_CHECK(q.dtype() == torch::kFloat16, "q must be fp16"); + TORCH_CHECK(k_cache.dtype() == torch::kFloat16, "k_cache must be fp16"); + TORCH_CHECK(v_cache.dtype() == torch::kFloat16, "v_cache must be fp16"); + TORCH_CHECK(block_table.dtype() == torch::kInt32, + "block_table must be int32"); + TORCH_CHECK(q.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda(), + "q, k_cache, and v_cache must be CUDA tensors"); + TORCH_CHECK(block_table.is_cuda(), "block_table must be a CUDA tensor"); + TORCH_CHECK(q.device() == k_cache.device() && + q.device() == v_cache.device() && + q.device() == block_table.device(), + "all input tensors must be on the q device"); + TORCH_CHECK(q.is_contiguous(), "q must be contiguous [B, H, M, D]"); + TORCH_CHECK(block_table.is_contiguous(), "block_table must be contiguous"); + TORCH_CHECK(k_cache.sizes() == v_cache.sizes(), + "k_cache and v_cache must have the same shape"); + TORCH_CHECK(k_cache.stride(-1) == 1 && v_cache.stride(-1) == 1, + "K/V head dimensions must be contiguous"); + TORCH_CHECK(k_cache.stride(1) % 8 == 0 && k_cache.stride(2) % 8 == 0 && + v_cache.stride(1) % 8 == 0 && v_cache.stride(2) % 8 == 0, + "K/V strides must be divisible by 8 half elements"); + + 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 num_kv_heads = k_cache.size(2); + + TORCH_CHECK(B == 1, "fixed BM32 split-KV host-length path requires B=1"); + TORCH_CHECK(H > 0, "q head dimension must be positive"); + TORCH_CHECK( + M >= D256_BM32_PHASE_BLOCK_M, + "fixed BM32 split-KV prefill requires M >= ", D256_BM32_PHASE_BLOCK_M); + TORCH_CHECK(D == D256_BM32_PHASE_D, + "fixed BM32 split-KV prefill requires D=", D256_BM32_PHASE_D); + TORCH_CHECK(k_cache.size(0) > 0 && num_kv_heads > 0, + "K/V cache block and head dimensions must be positive"); + TORCH_CHECK(k_cache.size(1) == D256_BM32_PHASE_PAGE_BLOCK_SIZE, + "fixed BM32 split-KV prefill requires page size ", + D256_BM32_PHASE_PAGE_BLOCK_SIZE); + TORCH_CHECK(k_cache.size(3) == D256_BM32_PHASE_D, + "K/V cache head dimension must be D=", D256_BM32_PHASE_D); + TORCH_CHECK(H % num_kv_heads == 0, + "num_q_heads must be divisible by num_kv_heads"); + TORCH_CHECK(block_table.size(0) == B, + "block_table batch dimension must equal q batch"); + TORCH_CHECK(actual_n >= M && actual_n <= 2147483647LL, + "actual_n must satisfy M <= actual_n <= INT_MAX"); + const int64_t required_pages = + (actual_n + D256_BM32_PHASE_PAGE_BLOCK_SIZE - 1) / + D256_BM32_PHASE_PAGE_BLOCK_SIZE; + TORCH_CHECK(block_table.size(1) >= required_pages, + "block_table does not cover actual_n host length"); + + if (out_.has_value()) { + const at::Tensor& out = out_.value(); + TORCH_CHECK(out.is_cuda() && out.device() == q.device(), + "out must be a CUDA tensor on the q device"); + TORCH_CHECK(out.dtype() == torch::kFloat16, "out must be fp16"); + TORCH_CHECK(out.sizes() == q.sizes(), "out must have shape [B, H, M, D]"); + TORCH_CHECK(out.is_contiguous(), "out must be contiguous [B, H, M, D]"); + } + if (softmax_lse_.has_value()) { + const at::Tensor& softmax_lse = softmax_lse_.value(); + TORCH_CHECK(softmax_lse.is_cuda() && softmax_lse.device() == q.device(), + "softmax_lse must be a CUDA tensor on the q device"); + TORCH_CHECK(softmax_lse.dtype() == torch::kFloat32, + "softmax_lse must be fp32"); + TORCH_CHECK(softmax_lse.dim() == 3 && softmax_lse.size(0) == B && + softmax_lse.size(1) == H && softmax_lse.size(2) == M, + "softmax_lse must have shape [B, H, M]"); + TORCH_CHECK(softmax_lse.is_contiguous(), + "softmax_lse must be contiguous [B, H, M]"); + } + + const std::vector expected_partial_out = { + B, H, D256_BM32_PHASE_SPLIT_PARTS, M, D}; + const std::vector expected_partial_rows = { + B, H, D256_BM32_PHASE_SPLIT_PARTS, M}; + for (const at::Tensor* workspace : + {&split_tmp_out, &split_tmp_row_max, &split_tmp_row_sum}) { + TORCH_CHECK(workspace->is_cuda() && workspace->device() == q.device(), + "split-KV workspaces must be CUDA tensors on the q device"); + TORCH_CHECK(workspace->dtype() == torch::kFloat32, + "split-KV workspaces must be fp32"); + TORCH_CHECK(workspace->is_contiguous(), + "split-KV workspaces must be contiguous"); + } + TORCH_CHECK(split_tmp_out.sizes() == expected_partial_out, + "split_tmp_out must have shape [B, H, 3, M, D]"); + TORCH_CHECK(split_tmp_row_max.sizes() == expected_partial_rows, + "split_tmp_row_max must have shape [B, H, 3, M]"); + TORCH_CHECK(split_tmp_row_sum.sizes() == expected_partial_rows, + "split_tmp_row_sum must have shape [B, H, 3, M]"); + + c10::cuda::CUDAGuard device_guard(q.device()); + auto stream = at::cuda::getCurrentCUDAStream().stream(); + cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; + const cudaError_t capture_error = + cudaStreamIsCapturing(stream, &capture_status); + TORCH_CHECK(capture_error == cudaSuccess, "cudaStreamIsCapturing failed: ", + cudaGetErrorString(capture_error)); + TORCH_CHECK(capture_status == cudaStreamCaptureStatusNone, + "fixed BM32 split-KV CUDA graph capture is not validated"); + + auto props = at::cuda::getCurrentDeviceProperties(); + TORCH_CHECK(props->major == 7 && props->minor == 0, + "fixed BM32 split-KV prefill supports only SM70 GPUs"); + + at::Tensor out_fp16 = out_.has_value() ? out_.value() : torch::empty_like(q); + at::Tensor softmax_lse = + softmax_lse_.has_value() + ? softmax_lse_.value() + : torch::empty({B, H, M}, + torch::dtype(torch::kFloat32).device(q.device())); + launch_flash_attention_forward_paged_d256_bm32_splitkv3_kernel( + q, k_cache, v_cache, out_fp16, softmax_lse, split_tmp_out, + split_tmp_row_max, split_tmp_row_sum, block_table, + static_cast(actual_n), softmax_scale, stream); + + return {out_fp16, softmax_lse}; } at::Tensor flash_attention_prefill_paged_bfla( - const at::Tensor& q, - const at::Tensor& k_cache, - const at::Tensor& v_cache, - std::optional& out_, - const at::Tensor& block_table, - const at::Tensor& seq_lens, - const at::Tensor& bfla_block_mask, - const int bfla_mask_block_n, - const float softmax_scale, - const std::string& kv_cache_dtype, - const float k_scale, - const float v_scale, - const bool is_causal, - const int window_size_left, - const int window_size_right -) { - TORCH_CHECK(q.dtype() == torch::kFloat16, "q must be fp16"); - const int kv_dtype_code = kv_cache_dtype_code_from_string(kv_cache_dtype); - TORCH_CHECK(kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP16, - "BFLA paged prefill supports fp16 KV cache only"); - TORCH_CHECK(k_cache.dtype() == torch::kFloat16, "k_cache must be fp16"); - TORCH_CHECK(v_cache.dtype() == torch::kFloat16, "v_cache must be fp16"); - TORCH_CHECK(q.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda(), - "Tensors must be on CUDA"); - TORCH_CHECK(block_table.is_cuda() && seq_lens.is_cuda(), - "block_table and seq_lens must be CUDA tensors"); - TORCH_CHECK(bfla_block_mask.is_cuda(), - "bfla_block_mask must be a CUDA tensor"); - TORCH_CHECK(bfla_block_mask.dtype() == torch::kInt32, - "bfla_block_mask must be int32"); - TORCH_CHECK(bfla_block_mask.dim() == 4, - "bfla_block_mask must have shape [B, Hkv, q_tiles, kv_tiles]"); - TORCH_CHECK(q.stride(-1) == 1 && k_cache.stride(-1) == 1 && - v_cache.stride(-1) == 1, - "Last dim must be contiguous"); - TORCH_CHECK(k_cache.stride(1) % 8 == 0 && k_cache.stride(2) % 8 == 0 && - v_cache.stride(1) % 8 == 0 && v_cache.stride(2) % 8 == 0, - "Paged KV strides must be divisible by 8 half elements"); - - 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 num_kv_heads = k_cache.size(2); - const int page_block_size = k_cache.size(1); - - TORCH_CHECK(D == 256, "BFLA paged prefill supports D=256 only"); - TORCH_CHECK(H % num_kv_heads == 0, - "num_heads must be divisible by num_kv_heads"); - TORCH_CHECK(M > 1, "BFLA paged prefill requires prefill M > 1"); - TORCH_CHECK(page_block_size >= 16 && (page_block_size % 16) == 0, - "BFLA paged prefill requires page block size multiple of 16"); - TORCH_CHECK(is_causal, "BFLA paged prefill supports causal attention only"); - TORCH_CHECK(window_size_left == -1 && window_size_right == -1, - "BFLA paged prefill does not support sliding window"); - TORCH_CHECK(bfla_mask_block_n > 0, - "bfla_mask_block_n must be positive"); - TORCH_CHECK(B <= bfla_block_mask.size(0), - "bfla_block_mask batch dimension must cover q"); - TORCH_CHECK(num_kv_heads <= bfla_block_mask.size(1), - "bfla_block_mask head dimension must cover KV heads"); - - at::Tensor out_fp16 = out_.has_value() ? out_.value() : torch::zeros_like(q); - auto softmax_lse = 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."); - - launcher_flash_attention_forward_paged<256, flash_v100::KV_CACHE_DTYPE_FP16>( - q, k_cache, v_cache, out_fp16, softmax_lse, block_table, seq_lens, - softmax_scale, is_causal, k_scale, v_scale, window_size_left, - window_size_right, stream, bfla_block_mask.data_ptr(), - bfla_mask_block_n, bfla_block_mask.stride(0), - bfla_block_mask.stride(1), bfla_block_mask.stride(2), - bfla_block_mask.stride(3)); - - return out_fp16; + const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, + std::optional& out_, const at::Tensor& block_table, + const at::Tensor& seq_lens, const at::Tensor& bfla_block_mask, + const int bfla_mask_block_n, const float softmax_scale, + const std::string& kv_cache_dtype, const float k_scale, const float v_scale, + const bool is_causal, const int window_size_left, + const int window_size_right) { + TORCH_CHECK(q.dtype() == torch::kFloat16, "q must be fp16"); + const int kv_dtype_code = kv_cache_dtype_code_from_string(kv_cache_dtype); + TORCH_CHECK(kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP16, + "BFLA paged prefill supports fp16 KV cache only"); + TORCH_CHECK(k_cache.dtype() == torch::kFloat16, "k_cache must be fp16"); + TORCH_CHECK(v_cache.dtype() == torch::kFloat16, "v_cache must be fp16"); + TORCH_CHECK(q.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda(), + "Tensors must be on CUDA"); + TORCH_CHECK(block_table.is_cuda() && seq_lens.is_cuda(), + "block_table and seq_lens must be CUDA tensors"); + TORCH_CHECK(bfla_block_mask.is_cuda(), + "bfla_block_mask must be a CUDA tensor"); + TORCH_CHECK(bfla_block_mask.dtype() == torch::kInt32, + "bfla_block_mask must be int32"); + TORCH_CHECK(bfla_block_mask.dim() == 4, + "bfla_block_mask must have shape [B, Hkv, q_tiles, kv_tiles]"); + TORCH_CHECK( + q.stride(-1) == 1 && k_cache.stride(-1) == 1 && v_cache.stride(-1) == 1, + "Last dim must be contiguous"); + TORCH_CHECK(k_cache.stride(1) % 8 == 0 && k_cache.stride(2) % 8 == 0 && + v_cache.stride(1) % 8 == 0 && v_cache.stride(2) % 8 == 0, + "Paged KV strides must be divisible by 8 half elements"); + + 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 num_kv_heads = k_cache.size(2); + const int page_block_size = k_cache.size(1); + + TORCH_CHECK(D == 256, "BFLA paged prefill supports D=256 only"); + TORCH_CHECK(H % num_kv_heads == 0, + "num_heads must be divisible by num_kv_heads"); + TORCH_CHECK(M > 1, "BFLA paged prefill requires prefill M > 1"); + TORCH_CHECK(page_block_size >= 16 && (page_block_size % 16) == 0, + "BFLA paged prefill requires page block size multiple of 16"); + TORCH_CHECK(is_causal, "BFLA paged prefill supports causal attention only"); + TORCH_CHECK(window_size_left == -1 && window_size_right == -1, + "BFLA paged prefill does not support sliding window"); + TORCH_CHECK(bfla_mask_block_n > 0, "bfla_mask_block_n must be positive"); + TORCH_CHECK(B <= bfla_block_mask.size(0), + "bfla_block_mask batch dimension must cover q"); + TORCH_CHECK(num_kv_heads <= bfla_block_mask.size(1), + "bfla_block_mask head dimension must cover KV heads"); + + at::Tensor out_fp16 = out_.has_value() ? out_.value() : torch::zeros_like(q); + auto softmax_lse = + 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."); + + launcher_flash_attention_forward_paged<256, flash_v100::KV_CACHE_DTYPE_FP16>( + q, k_cache, v_cache, out_fp16, softmax_lse, block_table, seq_lens, + softmax_scale, is_causal, k_scale, v_scale, window_size_left, + window_size_right, stream, bfla_block_mask.data_ptr(), + bfla_mask_block_n, bfla_block_mask.stride(0), bfla_block_mask.stride(1), + bfla_block_mask.stride(2), bfla_block_mask.stride(3)); + + return out_fp16; } at::Tensor flash_attention_decode_paged_wmma( - const at::Tensor& q, - const at::Tensor& k_cache, - const at::Tensor& v_cache, - std::optional& out_, - const at::Tensor& block_table, - const at::Tensor& seq_lens, - const float softmax_scale, - const std::string& kv_cache_dtype, - const float k_scale, - const float v_scale -) { - TORCH_CHECK(q.dtype() == torch::kFloat16, "q must be fp16"); - const int kv_dtype_code = kv_cache_dtype_code_from_string(kv_cache_dtype); - TORCH_CHECK(kv_dtype_code >= 0, "Unsupported kv_cache_dtype: ", kv_cache_dtype); - if (kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP16) { - TORCH_CHECK(k_cache.dtype() == torch::kFloat16, "k_cache must be fp16"); - TORCH_CHECK(v_cache.dtype() == torch::kFloat16, "v_cache must be fp16"); - } else { - TORCH_CHECK(k_cache.dtype() == torch::kUInt8, - "fp8 k_cache must be stored as uint8"); - TORCH_CHECK(v_cache.dtype() == torch::kUInt8, - "fp8 v_cache must be stored as uint8"); - TORCH_CHECK(k_scale > 0.f && v_scale > 0.f, - "fp8 k/v scales must be positive"); - } - TORCH_CHECK(q.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda(), - "Tensors must be on CUDA"); - TORCH_CHECK(block_table.is_cuda() && seq_lens.is_cuda(), - "block_table and seq_lens must be CUDA tensors"); - TORCH_CHECK(q.dim() == 3, "q must have shape [B, H, D]"); - TORCH_CHECK(block_table.dim() == 2, - "block_table must have shape [B, max_num_blocks]"); - TORCH_CHECK(seq_lens.dim() == 1, "seq_lens must have shape [B]"); - TORCH_CHECK(q.stride(-1) == 1 && k_cache.stride(-1) == 1 && - v_cache.stride(-1) == 1, - "Last dim must be contiguous"); - TORCH_CHECK(k_cache.stride(1) % 8 == 0 && k_cache.stride(2) % 8 == 0 && - v_cache.stride(1) % 8 == 0 && v_cache.stride(2) % 8 == 0, - "Paged KV strides must be divisible by 8 half elements"); - - const int B = q.size(0); - const int H = q.size(1); - const int D = q.size(2); - const int num_kv_heads = k_cache.size(2); - - TORCH_CHECK(B <= block_table.size(0), "block_table batch size must cover q"); - TORCH_CHECK(B <= seq_lens.size(0), "seq_lens batch size must cover q"); - TORCH_CHECK(D <= 256 && D % 8 == 0 && D % 2 == 0, - "D must be even, <=256, multiple of 8"); - TORCH_CHECK(H % num_kv_heads == 0, - "num_heads must be divisible by num_kv_heads"); - - at::Tensor out_fp16 = out_.has_value() ? out_.value() : torch::zeros_like(q); - TORCH_CHECK(out_fp16.is_cuda(), "out must be on CUDA"); - TORCH_CHECK(out_fp16.dtype() == torch::kFloat16, "out must be fp16"); - TORCH_CHECK(out_fp16.sizes() == q.sizes(), "out must have same shape as q"); - TORCH_CHECK(out_fp16.stride(-1) == 1, "out last dim must be contiguous"); - - at::Tensor q_m1 = q.unsqueeze(2); - at::Tensor out_m1 = out_fp16.unsqueeze(2); - auto softmax_lse = torch::zeros({B, H, 1}, - 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."); - - #define LAUNCH_DECODE_WMMA_TYPED(HDIM, KV_DTYPE_CODE) \ - launcher_flash_attention_forward_paged( \ - q_m1, k_cache, v_cache, out_m1, softmax_lse, block_table, seq_lens, \ - softmax_scale, true, k_scale, v_scale, -1, -1, stream) - - #define LAUNCH_DECODE_WMMA_BY_KV(HDIM) \ - do { \ - switch (kv_dtype_code) { \ - case flash_v100::KV_CACHE_DTYPE_FP16: \ - LAUNCH_DECODE_WMMA_TYPED(HDIM, flash_v100::KV_CACHE_DTYPE_FP16); \ - break; \ - case flash_v100::KV_CACHE_DTYPE_FP8_E4M3: \ - LAUNCH_DECODE_WMMA_TYPED(HDIM, flash_v100::KV_CACHE_DTYPE_FP8_E4M3); \ - break; \ - case flash_v100::KV_CACHE_DTYPE_FP8_E5M2: \ - LAUNCH_DECODE_WMMA_TYPED(HDIM, flash_v100::KV_CACHE_DTYPE_FP8_E5M2); \ - break; \ - default: \ - TORCH_CHECK(false, "Unsupported kv_cache_dtype: ", kv_cache_dtype); \ - } \ - } while (0) - - switch (D) { - case 16: - LAUNCH_DECODE_WMMA_BY_KV(16); - break; - case 32: - LAUNCH_DECODE_WMMA_BY_KV(32); - break; - case 64: - LAUNCH_DECODE_WMMA_BY_KV(64); - break; - case 128: - LAUNCH_DECODE_WMMA_BY_KV(128); - break; - case 256: - LAUNCH_DECODE_WMMA_BY_KV(256); - break; - default: - TORCH_CHECK(false, "Unsupported D: ", D); - } - - #undef LAUNCH_DECODE_WMMA_BY_KV - #undef LAUNCH_DECODE_WMMA_TYPED - - C10_CUDA_KERNEL_LAUNCH_CHECK(); - return out_fp16; + const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, + std::optional& out_, const at::Tensor& block_table, + const at::Tensor& seq_lens, const float softmax_scale, + const std::string& kv_cache_dtype, const float k_scale, + const float v_scale) { + TORCH_CHECK(q.dtype() == torch::kFloat16, "q must be fp16"); + const int kv_dtype_code = kv_cache_dtype_code_from_string(kv_cache_dtype); + TORCH_CHECK(kv_dtype_code >= 0, + "Unsupported kv_cache_dtype: ", kv_cache_dtype); + if (kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP16) { + TORCH_CHECK(k_cache.dtype() == torch::kFloat16, "k_cache must be fp16"); + TORCH_CHECK(v_cache.dtype() == torch::kFloat16, "v_cache must be fp16"); + } else { + TORCH_CHECK(k_cache.dtype() == torch::kUInt8, + "fp8 k_cache must be stored as uint8"); + TORCH_CHECK(v_cache.dtype() == torch::kUInt8, + "fp8 v_cache must be stored as uint8"); + TORCH_CHECK(k_scale > 0.f && v_scale > 0.f, + "fp8 k/v scales must be positive"); + } + TORCH_CHECK(q.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda(), + "Tensors must be on CUDA"); + TORCH_CHECK(block_table.is_cuda() && seq_lens.is_cuda(), + "block_table and seq_lens must be CUDA tensors"); + TORCH_CHECK(q.dim() == 3, "q must have shape [B, H, D]"); + TORCH_CHECK(block_table.dim() == 2, + "block_table must have shape [B, max_num_blocks]"); + TORCH_CHECK(seq_lens.dim() == 1, "seq_lens must have shape [B]"); + TORCH_CHECK( + q.stride(-1) == 1 && k_cache.stride(-1) == 1 && v_cache.stride(-1) == 1, + "Last dim must be contiguous"); + TORCH_CHECK(k_cache.stride(1) % 8 == 0 && k_cache.stride(2) % 8 == 0 && + v_cache.stride(1) % 8 == 0 && v_cache.stride(2) % 8 == 0, + "Paged KV strides must be divisible by 8 half elements"); + + const int B = q.size(0); + const int H = q.size(1); + const int D = q.size(2); + const int num_kv_heads = k_cache.size(2); + + TORCH_CHECK(B <= block_table.size(0), "block_table batch size must cover q"); + TORCH_CHECK(B <= seq_lens.size(0), "seq_lens batch size must cover q"); + TORCH_CHECK(D <= 256 && D % 8 == 0 && D % 2 == 0, + "D must be even, <=256, multiple of 8"); + TORCH_CHECK(H % num_kv_heads == 0, + "num_heads must be divisible by num_kv_heads"); + + at::Tensor out_fp16 = out_.has_value() ? out_.value() : torch::zeros_like(q); + TORCH_CHECK(out_fp16.is_cuda(), "out must be on CUDA"); + TORCH_CHECK(out_fp16.dtype() == torch::kFloat16, "out must be fp16"); + TORCH_CHECK(out_fp16.sizes() == q.sizes(), "out must have same shape as q"); + TORCH_CHECK(out_fp16.stride(-1) == 1, "out last dim must be contiguous"); + + at::Tensor q_m1 = q.unsqueeze(2); + at::Tensor out_m1 = out_fp16.unsqueeze(2); + auto softmax_lse = + torch::zeros({B, H, 1}, 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."); + +#define LAUNCH_DECODE_WMMA_TYPED(HDIM, KV_DTYPE_CODE) \ + launcher_flash_attention_forward_paged( \ + q_m1, k_cache, v_cache, out_m1, softmax_lse, block_table, seq_lens, \ + softmax_scale, true, k_scale, v_scale, -1, -1, stream) + +#define LAUNCH_DECODE_WMMA_BY_KV(HDIM) \ + do { \ + switch (kv_dtype_code) { \ + case flash_v100::KV_CACHE_DTYPE_FP16: \ + LAUNCH_DECODE_WMMA_TYPED(HDIM, flash_v100::KV_CACHE_DTYPE_FP16); \ + break; \ + case flash_v100::KV_CACHE_DTYPE_FP8_E4M3: \ + LAUNCH_DECODE_WMMA_TYPED(HDIM, flash_v100::KV_CACHE_DTYPE_FP8_E4M3); \ + break; \ + case flash_v100::KV_CACHE_DTYPE_FP8_E5M2: \ + LAUNCH_DECODE_WMMA_TYPED(HDIM, flash_v100::KV_CACHE_DTYPE_FP8_E5M2); \ + break; \ + default: \ + TORCH_CHECK(false, "Unsupported kv_cache_dtype: ", kv_cache_dtype); \ + } \ + } while (0) + + switch (D) { + case 16: + LAUNCH_DECODE_WMMA_BY_KV(16); + break; + case 32: + LAUNCH_DECODE_WMMA_BY_KV(32); + break; + case 64: + LAUNCH_DECODE_WMMA_BY_KV(64); + break; + case 128: + LAUNCH_DECODE_WMMA_BY_KV(128); + break; + case 256: + LAUNCH_DECODE_WMMA_BY_KV(256); + break; + default: + TORCH_CHECK(false, "Unsupported D: ", D); + } + +#undef LAUNCH_DECODE_WMMA_BY_KV +#undef LAUNCH_DECODE_WMMA_TYPED + + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out_fp16; } 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/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 5331a3b79b..ffe58c6366 100644 --- a/requirements/docs.txt +++ b/requirements/docs.txt @@ -6,7 +6,7 @@ babel==2.18.0 # via # mkdocs-git-revision-date-localized-plugin # mkdocs-material -backrefs==7.0 +backrefs==8.0 # via mkdocs-material bracex==2.6 # via wcmatch @@ -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 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 831de114dd..e6a621d712 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,28 +552,28 @@ 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 orjson==3.11.5 # via genai-perf -packaging==24.2 +packaging==26.3 # via # accelerate # bitsandbytes @@ -604,19 +603,19 @@ pandas==2.2.3 # fastparquet # genai-perf # statsmodels -pathspec==0.12.1 +pathspec==1.1.1 # via black 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 @@ -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 6ca12d8387..a9c6a90ab7 100644 --- a/setup.py +++ b/setup.py @@ -891,6 +891,9 @@ def extract_precompiled_and_patch_package( flash_qla_sm70_ext_regex = re.compile( r"flash_qla/ops/gated_delta_rule/chunk/sm70/[^/]+\.so" ) + sm70_sampler_ext_regex = re.compile( + r"vllm/_sm70_sampler_C(?:\.[^/]+)?\.so$" + ) file_members = [] for member in wheel.filelist: if member.filename in exact_members: @@ -910,6 +913,7 @@ def extract_precompiled_and_patch_package( or deep_gemm_regex.match(member.filename) or flash_attn_v100_ext_regex.match(member.filename) or flash_qla_sm70_ext_regex.match(member.filename) + or sm70_sampler_ext_regex.match(member.filename) ): file_members.append(member) @@ -1235,6 +1239,8 @@ def _read_requirements(filename: str) -> list[str]: ext_modules.append(CMakeExtension(name="vllm._rocm_C")) if _is_cuda(): + if _cuda_arch_contains(7, 0): + ext_modules.append(CMakeExtension(name="vllm._sm70_sampler_C")) build_sm70_fa2 = _cuda_arch_contains(7, 0) and not _cuda_arch_at_least(8, 0) if _cuda_arch_at_least(8, 0) or build_sm70_fa2: ext_modules.append(CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa2_C")) @@ -1386,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", @@ -1402,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_dsv4_quality_tools.py b/tests/benchmarks/test_dsv4_quality_tools.py new file mode 100644 index 0000000000..09697faec9 --- /dev/null +++ b/tests/benchmarks/test_dsv4_quality_tools.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from copy import deepcopy + +from benchmarks.benchmark_dsv4_gsm8k_api import _last_integer +from benchmarks.compare_dsv4_quality_results import ( + _compare_gsm8k, + _compare_humaneval, + _compare_needle, +) + + +def test_last_integer_normalizes_signed_integral_decimal() -> None: + assert _last_integer("work... #### -1,024") == -1024 + assert _last_integer("answer: 12.0") == 12 + assert _last_integer("answer: 12.5") is None + assert _last_integer("no numeric answer") is None + + +def _gsm8k_artifact(predictions: list[int | None]) -> dict: + expected = [13, 40] + rows = [ + { + "index": index, + "question": f"question-{index}", + "expected": answer, + "predicted": prediction, + "correct": prediction is not None and prediction == answer, + "invalid": prediction is None, + } + for index, (answer, prediction) in enumerate(zip(expected, predictions)) + ] + return { + "contract": { + "questions": 2, + "few_shot": 5, + "temperature": 0.0, + "top_p": 1.0, + "seed": 42, + "max_tokens": 256, + "strictly_sequential": True, + "train_selection": "first_n", + "test_selection": "first_n", + "prompt_format": "gsm8k_question_answer_v1", + "answer_normalization": "last_signed_integral_decimal_v1", + "stop_sequences": ["Question"], + }, + "input_manifest": { + "train": {"sha256": "train-hash"}, + "test": {"sha256": "test-hash"}, + }, + "correct": sum(bool(row["correct"]) for row in rows), + "invalid": sum(bool(row["invalid"]) for row in rows), + "rows": rows, + } + + +def test_gsm8k_gate_allows_balanced_answer_flips_without_greedy_identity() -> None: + reference = _gsm8k_artifact([13, 42]) + candidate = _gsm8k_artifact([12, 40]) + + result = _compare_gsm8k(reference, candidate) + + assert result["passed"] + assert result["correct_delta"] == 0 + assert result["regressions"] == [0] + assert result["improvements"] == [1] + assert result["exact_prediction_matches"] == 0 + + +def test_gsm8k_gate_rejects_aggregate_drop_contract_drift_and_bad_summary() -> None: + reference = _gsm8k_artifact([13, 40]) + candidate = _gsm8k_artifact([12, 40]) + assert not _compare_gsm8k(reference, candidate)["passed"] + + candidate = deepcopy(reference) + candidate["contract"]["seed"] = 7 + assert not _compare_gsm8k(reference, candidate)["passed"] + + candidate = deepcopy(reference) + candidate["correct"] = 0 + assert not _compare_gsm8k(reference, candidate)["passed"] + + +def test_humaneval_gate_uses_aggregate_quality_and_reports_flips() -> None: + reference = { + "human_eval": { + "passed": 1, + "records": [ + {"task_id": "a", "passed": True, "response": "a0"}, + {"task_id": "b", "passed": False, "response": "b0"}, + ], + } + } + candidate = { + "human_eval": { + "passed": 1, + "records": [ + {"task_id": "a", "passed": False, "response": "a1"}, + {"task_id": "b", "passed": True, "response": "b1"}, + ], + } + } + + result = _compare_humaneval(reference, candidate) + + assert result["passed"] + assert result["regressions"] == ["a"] + assert result["improvements"] == ["b"] + + +def test_needle_gate_uses_aggregate_hits_and_reports_flips() -> None: + reference = [ + { + "sample_id": "a", + "target_tokens": 100, + "depth": 0.5, + "code": "x", + "hit": True, + "hit_anywhere": True, + }, + { + "sample_id": "b", + "target_tokens": 100, + "depth": 0.5, + "code": "y", + "hit": False, + "hit_anywhere": False, + }, + ] + candidate = [ + {**reference[0], "hit": False, "hit_anywhere": False}, + {**reference[1], "hit": True, "hit_anywhere": True}, + ] + + result = _compare_needle(reference, candidate) + + assert result["passed"] + assert result["final_hit_regressions"] == ["a"] + assert result["anywhere_hit_regressions"] == ["a"] diff --git a/tests/config/test_prefix_anchored_swa.py b/tests/config/test_prefix_anchored_swa.py new file mode 100644 index 0000000000..6a801406fc --- /dev/null +++ b/tests/config/test_prefix_anchored_swa.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +import vllm.platforms as platforms +from vllm.config.attention import AttentionConfig +from vllm.config.vllm import apply_prefix_anchored_swa_constraints +from vllm.model_executor.layers.attention.attention import Attention +from vllm.v1.attention.backend import AttentionType +from vllm.v1.attention.backends.registry import AttentionBackendEnum +from vllm.v1.kv_cache_interface import PrefixAnchoredSWASpec + + +def _config(window: int = 128) -> SimpleNamespace: + return SimpleNamespace( + attention_config=AttentionConfig( + prefix_anchored_decode_window=window, + ), + cache_config=SimpleNamespace( + cache_dtype="auto", + enable_prefix_caching=True, + kv_offloading_size=None, + ), + model_config=SimpleNamespace(dtype=torch.float16, use_mla=False), + parallel_config=SimpleNamespace( + decode_context_parallel_size=1, + prefill_context_parallel_size=1, + ), + kv_transfer_config=None, + speculative_config=None, + ) + + +@pytest.fixture +def sm70_platform(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + platforms, + "current_platform", + SimpleNamespace( + is_cuda=lambda: True, + is_device_capability=lambda capability: capability == (7, 0), + ), + ) + + +def test_engine_contract_is_model_identity_independent(sm70_platform: None): + cfg = _config() + + apply_prefix_anchored_swa_constraints(cfg) + + assert cfg.attention_config.backend == AttentionBackendEnum.FLASH_ATTN_V100 + assert not cfg.cache_config.enable_prefix_caching + + +@pytest.mark.parametrize("window", [0, -1]) +def test_engine_option_rejects_nonpositive_window(window: int): + with pytest.raises(ValueError): + AttentionConfig(prefix_anchored_decode_window=window) + + +def test_engine_contract_rejects_non_fp16_kv(sm70_platform: None): + cfg = _config() + cfg.cache_config.cache_dtype = "fp8_e5m2" + + with pytest.raises(ValueError, match="requires an fp16 KV cache"): + apply_prefix_anchored_swa_constraints(cfg) + + +@pytest.mark.parametrize( + ("path", "value", "message"), + [ + ("model_config.use_mla", True, "MLA attention"), + ("attention_config.use_non_causal", True, "causal decoder"), + ("parallel_config.decode_context_parallel_size", 2, "context parallelism"), + ("cache_config.kv_offloading_size", 1.0, "KV-cache offloading"), + ( + "kv_transfer_config", + SimpleNamespace(kv_connector="ExampleConnector"), + "KV connectors", + ), + ("speculative_config", object(), "speculative decoding"), + ( + "attention_config.backend", + AttentionBackendEnum.TRITON_ATTN, + "requires FLASH_ATTN_V100", + ), + ], +) +def test_engine_contract_rejects_unsupported_runtime_combinations( + sm70_platform: None, + path: str, + value: object, + message: str, +): + cfg = _config() + target = cfg + parts = path.split(".") + for part in parts[:-1]: + target = getattr(target, part) + setattr(target, parts[-1], value) + + with pytest.raises(ValueError, match=message): + apply_prefix_anchored_swa_constraints(cfg) + + +def test_engine_contract_rejects_non_sm70(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr( + platforms, + "current_platform", + SimpleNamespace( + is_cuda=lambda: True, + is_device_capability=lambda capability: False, + ), + ) + + with pytest.raises(ValueError, match="SM70 GPU"): + apply_prefix_anchored_swa_constraints(_config()) + + +def test_base_attention_emits_prefix_anchored_spec_without_model_subclass(): + layer = object.__new__(Attention) + object.__setattr__(layer, "attn_type", AttentionType.DECODER) + object.__setattr__( + layer, + "attn_backend", + SimpleNamespace(get_name=lambda: "FLASH_ATTN_V100"), + ) + object.__setattr__(layer, "kv_cache_dtype", "auto") + object.__setattr__(layer, "kv_cache_torch_dtype", torch.float16) + object.__setattr__(layer, "sliding_window", None) + object.__setattr__(layer, "num_kv_heads", 4) + object.__setattr__(layer, "head_size", 128) + object.__setattr__(layer, "head_size_v", 128) + object.__setattr__(layer, "has_sink", False) + vllm_config = SimpleNamespace( + cache_config=SimpleNamespace(block_size=16), + attention_config=SimpleNamespace(prefix_anchored_decode_window=256), + ) + + spec = Attention.get_kv_cache_spec(layer, vllm_config) + + assert isinstance(spec, PrefixAnchoredSWASpec) + assert spec.decode_sliding_window == 256 + + +def test_prefix_anchored_spec_rejects_nonpositive_window(): + with pytest.raises(ValueError, match="greater than zero"): + PrefixAnchoredSWASpec( + block_size=16, + num_kv_heads=1, + head_size=64, + dtype=torch.float16, + decode_sliding_window=0, + ) 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_allocator.py b/tests/distributed/test_custom_all_reduce_allocator.py new file mode 100644 index 0000000000..d455012fc9 --- /dev/null +++ b/tests/distributed/test_custom_all_reduce_allocator.py @@ -0,0 +1,224 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import regex as re + +from vllm.distributed.device_communicators import custom_all_reduce + +pytestmark = pytest.mark.skip_global_cleanup + + +def _patch_allocator_settings(monkeypatch: pytest.MonkeyPatch) -> list[str]: + settings: list[str] = [] + monkeypatch.setattr(custom_all_reduce.current_platform, "is_cuda", lambda: True) + monkeypatch.setattr( + custom_all_reduce.torch.cuda.memory, + "_set_allocator_settings", + settings.append, + ) + return settings + + +def test_cuda_ipc_capture_temporarily_disables_expandable_segments( + monkeypatch: pytest.MonkeyPatch, +) -> None: + allocator_conf = ( + "max_split_size_mb:512,garbage_collection_threshold:0.8," + "expandable_segments:True" + ) + monkeypatch.setenv( + "PYTORCH_CUDA_ALLOC_CONF", + allocator_conf, + ) + settings = _patch_allocator_settings(monkeypatch) + + with custom_all_reduce._disable_expandable_segments_for_cuda_ipc(True): + assert settings == [allocator_conf.replace("True", "False")] + + assert settings == [allocator_conf.replace("True", "False"), allocator_conf] + + +def test_cuda_ipc_capture_preserves_allocator_state( + monkeypatch: pytest.MonkeyPatch, +) -> None: + allocator_conf = ( + "max_split_size_mb:512,roundup_power2_divisions:4," + "garbage_collection_threshold:0.8,expandable_segments:True" + ) + monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", allocator_conf) + monkeypatch.setattr(custom_all_reduce.current_platform, "is_cuda", lambda: True) + allocator_state: dict[str, int | float | bool] = { + "max_split_size_mb": 512, + "roundup_power2_divisions": 4, + "garbage_collection_threshold": 0.8, + "expandable_segments": True, + } + + def set_allocator_settings(conf: str) -> None: + # Model PyTorch's parseArgs behavior: these options reset whenever + # they are omitted from a runtime settings update. + allocator_state.update( + max_split_size_mb=-1, + roundup_power2_divisions=0, + garbage_collection_threshold=0.0, + ) + for key in allocator_state: + match = re.search(rf"(?:^|,){key}:([^,]+)", conf) + if match is None: + continue + value = match.group(1) + if key == "garbage_collection_threshold": + allocator_state[key] = float(value) + elif key == "expandable_segments": + allocator_state[key] = value == "True" + else: + allocator_state[key] = int(value) + + monkeypatch.setattr( + custom_all_reduce.torch.cuda.memory, + "_set_allocator_settings", + set_allocator_settings, + ) + + with custom_all_reduce._disable_expandable_segments_for_cuda_ipc(True): + assert allocator_state == { + "max_split_size_mb": 512, + "roundup_power2_divisions": 4, + "garbage_collection_threshold": 0.8, + "expandable_segments": False, + } + + assert allocator_state == { + "max_split_size_mb": 512, + "roundup_power2_divisions": 4, + "garbage_collection_threshold": 0.8, + "expandable_segments": True, + } + + +def test_cuda_ipc_capture_supports_unified_allocator_conf( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("PYTORCH_CUDA_ALLOC_CONF", raising=False) + allocator_conf = "max_split_size_mb:512,expandable_segments:True" + monkeypatch.setenv("PYTORCH_ALLOC_CONF", allocator_conf) + settings = _patch_allocator_settings(monkeypatch) + + with custom_all_reduce._disable_expandable_segments_for_cuda_ipc(True): + pass + + assert settings == [allocator_conf.replace("True", "False"), allocator_conf] + + +@pytest.mark.parametrize("legacy_conf", ["", "expandable_segments:False"]) +def test_legacy_allocator_conf_takes_precedence_over_unified_conf( + monkeypatch: pytest.MonkeyPatch, + legacy_conf: str, +) -> None: + monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", legacy_conf) + monkeypatch.setenv("PYTORCH_ALLOC_CONF", "expandable_segments:True") + settings = _patch_allocator_settings(monkeypatch) + + with custom_all_reduce._disable_expandable_segments_for_cuda_ipc(True): + pass + + assert settings == [] + + +def test_enabled_legacy_allocator_conf_takes_precedence( + monkeypatch: pytest.MonkeyPatch, +) -> None: + legacy_conf = "max_split_size_mb:512,expandable_segments:True" + monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", legacy_conf) + monkeypatch.setenv("PYTORCH_ALLOC_CONF", "expandable_segments:False") + settings = _patch_allocator_settings(monkeypatch) + + with custom_all_reduce._disable_expandable_segments_for_cuda_ipc(True): + pass + + assert settings == [legacy_conf.replace("True", "False"), legacy_conf] + + +def test_cuda_ipc_capture_preserves_bracketed_allocator_syntax( + monkeypatch: pytest.MonkeyPatch, +) -> None: + allocator_conf = ( + "roundup_power2_divisions:[256:1,512:2,>:4], expandable_segments : True" + ) + monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", allocator_conf) + settings = _patch_allocator_settings(monkeypatch) + + with custom_all_reduce._disable_expandable_segments_for_cuda_ipc(True): + pass + + assert settings == [ + allocator_conf.replace("True", "False"), + allocator_conf, + ] + + +def test_cuda_ipc_capture_restores_allocator_after_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") + settings = _patch_allocator_settings(monkeypatch) + + with ( + pytest.raises(RuntimeError, match="capture failed"), + custom_all_reduce._disable_expandable_segments_for_cuda_ipc(True), + ): + raise RuntimeError("capture failed") + + assert settings == ["expandable_segments:False", "expandable_segments:True"] + + +def test_custom_allreduce_registers_graph_buffers_before_allocator_restore( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") + settings = _patch_allocator_settings(monkeypatch) + communicator = object.__new__(custom_all_reduce.CustomAllreduce) + communicator.disabled = False + communicator._IS_CAPTURING = False + communicator._ptr = None + registration_state: list[tuple[list[str], bool]] = [] + + def register_graph_buffers() -> None: + registration_state.append((list(settings), communicator._IS_CAPTURING)) + + monkeypatch.setattr( + communicator, + "register_graph_buffers", + register_graph_buffers, + ) + + with communicator.capture(): + assert communicator._IS_CAPTURING + assert settings == ["expandable_segments:False"] + + assert registration_state == [(["expandable_segments:False"], False)] + assert settings == ["expandable_segments:False", "expandable_segments:True"] + + +@pytest.mark.parametrize( + ("active", "allocator_conf"), + [ + (False, "expandable_segments:True"), + (True, "expandable_segments:False"), + (True, ""), + ], +) +def test_cuda_ipc_capture_allocator_guard_is_otherwise_a_noop( + monkeypatch: pytest.MonkeyPatch, + active: bool, + allocator_conf: str, +) -> None: + monkeypatch.delenv("PYTORCH_ALLOC_CONF", raising=False) + monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", allocator_conf) + settings = _patch_allocator_settings(monkeypatch) + + with custom_all_reduce._disable_expandable_segments_for_cuda_ipc(active): + pass + + assert settings == [] diff --git a/tests/distributed/test_static_tensor_transfer.py b/tests/distributed/test_static_tensor_transfer.py new file mode 100644 index 0000000000..4a1890f5dc --- /dev/null +++ b/tests/distributed/test_static_tensor_transfer.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Any + +import pytest +import torch + +from vllm.distributed.parallel_state import GroupCoordinator + + +class _DummyWork: + def wait(self) -> None: + pass + + +def _group(rank_in_group: int, world_size: int = 2) -> GroupCoordinator: + group = GroupCoordinator.__new__(GroupCoordinator) + group.world_size = world_size + group.rank_in_group = rank_in_group + group.ranks = list(range(world_size)) + group.use_cpu_custom_send_recv = False + group.device_group = None + group.cpu_group = None + return group + + +def test_static_tensor_dict_transfer_uses_caller_owned_buffers( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: list[tuple[str, int, torch.Tensor]] = [] + + def fake_isend(tensor: torch.Tensor, *, dst: int, group: Any) -> _DummyWork: + del group + calls.append(("send", dst, tensor)) + return _DummyWork() + + def fake_irecv(tensor: torch.Tensor, *, src: int, group: Any) -> _DummyWork: + del group + tensor.fill_(7) + calls.append(("recv", src, tensor)) + return _DummyWork() + + monkeypatch.setattr(torch.distributed, "isend", fake_isend) + monkeypatch.setattr(torch.distributed, "irecv", fake_irecv) + + sent = torch.arange(8, dtype=torch.float32) + send_handles = _group(0).isend_tensor_dict_static({"hidden_states": sent}) + + received = torch.empty_like(sent) + recv_handles = _group(1).irecv_tensor_dict_static({"hidden_states": received}) + + assert len(send_handles) == 1 + assert len(recv_handles) == 1 + assert calls == [("send", 1, sent), ("recv", 0, received)] + torch.testing.assert_close(received, torch.full_like(received, 7)) + + +def test_static_tensor_dict_transfer_rejects_non_tensors() -> None: + with pytest.raises(TypeError, match="only accepts tensors"): + _group(0).isend_tensor_dict_static( # type: ignore[arg-type] + {"hidden_states": object()} + ) diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 9f2d8819ef..541f3d5d42 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -237,19 +237,35 @@ def _fake_qwen_hybrid_model_config(): ) -def _apply_sm70_defaults(monkeypatch, *, env=None, speculative_config=None): +def _apply_sm70_defaults( + monkeypatch, + *, + env=None, + speculative_config=None, + enable_prefix_caching=None, + max_num_seqs=None, +): for key in ( "VLLM_1CAT_ENABLE_SM70_MTP_DEFAULTS", "VLLM_1CAT_ENABLE_QWEN35_MTP_DEFAULTS", "VLLM_1CAT_DISABLE_SM70_MTP_DEFAULTS", "VLLM_1CAT_DISABLE_QWEN35_MTP_DEFAULTS", + "VLLM_SM70_MTP_SPLIT_DRAFT_CUDAGRAPHS", ): monkeypatch.delenv(key, raising=False) for key, value in (env or {}).items(): monkeypatch.setenv(key, value) monkeypatch.setattr("vllm.engine.arg_utils.current_platform", _FakeSM70Platform()) + monkeypatch.setattr( + "vllm.engine.arg_utils.get_model_path", lambda model, revision: model + ) - args = EngineArgs(model="dummy", tensor_parallel_size=4) + args = EngineArgs( + model="dummy", + tensor_parallel_size=4, + enable_prefix_caching=enable_prefix_caching, + max_num_seqs=max_num_seqs, + ) args.speculative_config = speculative_config args._maybe_apply_sm70_mtp_defaults( UsageContext.OPENAI_API_SERVER, @@ -297,6 +313,39 @@ def test_sm70_mtp_defaults_require_env_opt_in(monkeypatch): ] +def test_sm70_mtp_split_cudagraphs_are_opt_in(monkeypatch): + args = _apply_sm70_defaults( + monkeypatch, + env={ + "VLLM_1CAT_ENABLE_SM70_MTP_DEFAULTS": "1", + "VLLM_SM70_MTP_SPLIT_DRAFT_CUDAGRAPHS": "1", + }, + ) + + assert args.compilation_config.cudagraph_capture_sizes == [5, 10, 20] + + +def test_sm70_mtp_split_cudagraphs_cover_production_batches(monkeypatch): + args = _apply_sm70_defaults( + monkeypatch, + env={ + "VLLM_1CAT_ENABLE_SM70_MTP_DEFAULTS": "1", + "VLLM_SM70_MTP_SPLIT_DRAFT_CUDAGRAPHS": "1", + }, + max_num_seqs=16, + ) + + assert args.compilation_config.cudagraph_capture_sizes == [ + 5, + 10, + 20, + 30, + 40, + 60, + 80, + ] + + def test_sm70_mtp_defaults_treat_zero_env_as_disabled(monkeypatch): args = _apply_sm70_defaults( monkeypatch, @@ -339,6 +388,37 @@ def test_sm70_explicit_mtp_still_gets_safe_defaults(monkeypatch): assert args.max_num_seqs == 4 +def test_sm70_explicit_dflash_preserves_probabilistic_default(monkeypatch): + args = _apply_sm70_defaults( + monkeypatch, + speculative_config={"method": "dflash", "num_speculative_tokens": 7}, + ) + + assert args.speculative_config == { + "method": "dflash", + "num_speculative_tokens": 7, + "draft_sample_method": "probabilistic", + } + assert args.enable_prefix_caching is True + assert args.mamba_cache_mode == "align" + + +def test_sm70_explicit_dflash_accepts_explicit_greedy(monkeypatch): + args = _apply_sm70_defaults( + monkeypatch, + speculative_config={ + "method": "dflash", + "num_speculative_tokens": 7, + "draft_sample_method": "greedy", + }, + enable_prefix_caching=True, + ) + + assert args.speculative_config["draft_sample_method"] == "greedy" + assert args.enable_prefix_caching is True + assert args.mamba_cache_mode == "align" + + @pytest.mark.parametrize( ("cli_args", "expected"), [ diff --git a/tests/kernels/attention/test_sm70_flash_v100_anchored_swa.py b/tests/kernels/attention/test_sm70_flash_v100_anchored_swa.py new file mode 100644 index 0000000000..d93cbc878e --- /dev/null +++ b/tests/kernels/attention/test_sm70_flash_v100_anchored_swa.py @@ -0,0 +1,288 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""SM70 Flash-V100 anchored decode-window mask kernel tests. + +Validates the masked (second) kernel instantiations of the paged decode and +paged prefill kernels against a plain torch reference: +``causal AND (kv < anchor_len OR q_abs - kv < window)``, element-wise. + +Requires a CUDA device and the flash_attn_v100 extension built with the +anchored kernel arguments; skipped otherwise. +""" + +from collections.abc import Callable +from typing import Any + +import pytest +import torch + +cuda_available = torch.cuda.is_available() +FlashOp = Callable[..., Any] +flash_ops: tuple[FlashOp, FlashOp] | None = None +if cuda_available: + try: + from flash_attn_v100 import ( + flash_attn_decode_paged, + flash_attn_prefill_paged, + ) + + flash_ops = (flash_attn_decode_paged, flash_attn_prefill_paged) + except ImportError: + flash_ops = None + +pytestmark = pytest.mark.skipif( + not cuda_available or flash_ops is None, + reason="requires CUDA and the flash_attn_v100 extension", +) + + +def _require_flash_ops() -> tuple[FlashOp, FlashOp]: + assert flash_ops is not None + return flash_ops + + +def _anchored_mask( + q_positions: torch.Tensor, + kv_len: int, + anchor_len: int, + window: int, + device: str, +) -> torch.Tensor: + """Boolean keep-mask [num_q, kv_len]; True = attend.""" + kv = torch.arange(kv_len, device=device)[None, :] + q = q_positions[:, None] + causal = kv <= q + in_prefix = kv < anchor_len + in_window = (q - kv) < window + return causal & (in_prefix | in_window) + + +def _ref_attention( + q: torch.Tensor, # [num_q, H, D] + k: torch.Tensor, # [kv_len, H_kv, D] + v: torch.Tensor, + mask: torch.Tensor, # [num_q, kv_len] + scale: float, +) -> torch.Tensor: + group = q.shape[1] // k.shape[1] + k_exp = k.repeat_interleave(group, dim=1) + v_exp = v.repeat_interleave(group, dim=1) + scores = torch.einsum("qhd,khd->hqk", q.float(), k_exp.float()) * scale + scores = scores.masked_fill(~mask[None], float("-inf")) + probs = torch.softmax(scores, dim=-1) + return torch.einsum("hqk,khd->qhd", probs, v_exp.float()) + + +def _build_paged_kv( + kv_len: int, + block_size: int, + num_kv_heads: int, + head_dim: int, + device: str, + dtype: torch.dtype, + seed: int, +): + torch.manual_seed(seed) + num_blocks = (kv_len + block_size - 1) // block_size + 2 + k_cache = torch.randn( + num_blocks, block_size, num_kv_heads, head_dim, device=device, dtype=dtype + ) + v_cache = torch.randn_like(k_cache) + block_table = torch.arange(num_blocks, dtype=torch.int32, device=device).unsqueeze( + 0 + ) + k_lin = k_cache.reshape(-1, num_kv_heads, head_dim)[:kv_len] + v_lin = v_cache.reshape(-1, num_kv_heads, head_dim)[:kv_len] + return k_cache, v_cache, block_table, k_lin, v_lin + + +@pytest.mark.parametrize( + "head_dim,num_heads,num_kv_heads", + [(256, 8, 4), (128, 8, 8), (128, 8, 4), (64, 4, 4)], +) +# kv_len=1100 spans three 512-token partitions with the middle partition +# entirely inside the masked gap (exercises the neutral-stats path). +@pytest.mark.parametrize( + "anchor,window,kv_len", [(32, 16, 96), (16, 8, 40), (32, 16, 1100)] +) +def test_decode_paged_anchored_matches_reference( + head_dim, num_heads, num_kv_heads, anchor, window, kv_len +): + device = "cuda" + dtype = torch.float16 + block_size = 16 + scale = head_dim**-0.5 + + k_cache, v_cache, block_table, k_lin, v_lin = _build_paged_kv( + kv_len, block_size, num_kv_heads, head_dim, device, dtype, seed=0 + ) + q = torch.randn(1, num_heads, head_dim, device=device, dtype=dtype) + seq_lens = torch.tensor([kv_len], dtype=torch.int32, device=device) + anchor_lens = torch.tensor([anchor], dtype=torch.int32, device=device) + + out = torch.empty_like(q) + decode_op, _ = _require_flash_ops() + decode_op( + q, + k_cache, + v_cache, + block_table, + seq_lens, + softmax_scale=scale, + out=out, + anchor_lens=anchor_lens, + anchored_window=window, + ) + + q_pos = torch.tensor([kv_len - 1], device=device) + mask = _anchored_mask(q_pos, kv_len, anchor, window, device) + assert not mask[0].all(), "test must exercise a non-trivial gap" + ref = _ref_attention(q[0].unsqueeze(0), k_lin, v_lin, mask, scale) + torch.testing.assert_close(out[0].float(), ref[0], atol=2e-2, rtol=2e-2) + + +def test_decode_paged_anchored_off_path_unchanged(): + """Without anchor arguments the decode op matches its own baseline.""" + device = "cuda" + dtype = torch.float16 + head_dim, num_heads, num_kv_heads = 128, 8, 8 + kv_len, block_size = 96, 16 + scale = head_dim**-0.5 + + k_cache, v_cache, block_table, k_lin, v_lin = _build_paged_kv( + kv_len, block_size, num_kv_heads, head_dim, device, dtype, seed=1 + ) + q = torch.randn(1, num_heads, head_dim, device=device, dtype=dtype) + seq_lens = torch.tensor([kv_len], dtype=torch.int32, device=device) + + out = torch.empty_like(q) + decode_op, _ = _require_flash_ops() + decode_op(q, k_cache, v_cache, block_table, seq_lens, softmax_scale=scale, out=out) + q_pos = torch.tensor([kv_len - 1], device=device) + causal = _anchored_mask(q_pos, kv_len, kv_len, kv_len, device) + assert causal[0].all() + ref = _ref_attention(q[0].unsqueeze(0), k_lin, v_lin, causal, scale) + torch.testing.assert_close(out[0].float(), ref[0], atol=2e-2, rtol=2e-2) + + +def test_decode_paged_uses_per_request_anchor_lengths(): + device = "cuda" + dtype = torch.float16 + head_dim, num_heads, num_kv_heads = 128, 8, 4 + kv_len, block_size, window = 96, 16, 16 + anchors = (16, 32) + scale = head_dim**-0.5 + + k_cache, v_cache, block_table, k_lin, v_lin = _build_paged_kv( + kv_len, block_size, num_kv_heads, head_dim, device, dtype, seed=3 + ) + block_table = block_table.repeat(2, 1) + q = torch.randn(2, num_heads, head_dim, device=device, dtype=dtype) + seq_lens = torch.full((2,), kv_len, dtype=torch.int32, device=device) + anchor_lens = torch.tensor(anchors, dtype=torch.int32, device=device) + + decode_op, _ = _require_flash_ops() + out = torch.empty_like(q) + decode_op( + q, + k_cache, + v_cache, + block_table, + seq_lens, + softmax_scale=scale, + out=out, + anchor_lens=anchor_lens, + anchored_window=window, + ) + + q_pos = torch.tensor([kv_len - 1], device=device) + for batch_idx, anchor in enumerate(anchors): + mask = _anchored_mask(q_pos, kv_len, anchor, window, device) + ref = _ref_attention(q[batch_idx].unsqueeze(0), k_lin, v_lin, mask, scale) + torch.testing.assert_close(out[batch_idx].float(), ref[0], atol=2e-2, rtol=2e-2) + + +@pytest.mark.parametrize( + "head_dim,num_heads,num_kv_heads", + [(256, 8, 4), (128, 8, 4)], +) +@pytest.mark.parametrize( + "anchor,window,kv_len,q_len", [(32, 16, 96, 8), (16, 8, 48, 48)] +) +def test_prefill_paged_anchored_matches_reference( + head_dim, num_heads, num_kv_heads, anchor, window, kv_len, q_len +): + """Masked paged prefill vs torch reference. + + Covers both the chunked-continuation shape (q_len < kv_len, queries in + the decode region) and the full-sequence shape (q_len == kv_len, where + prompt-region queries must stay purely causal). + """ + device = "cuda" + dtype = torch.float16 + block_size = 16 + scale = head_dim**-0.5 + + k_cache, v_cache, block_table, k_lin, v_lin = _build_paged_kv( + kv_len, block_size, num_kv_heads, head_dim, device, dtype, seed=2 + ) + q = torch.randn(1, q_len, num_heads, head_dim, device=device, dtype=dtype) + seq_lens = torch.tensor([kv_len], dtype=torch.int32, device=device) + anchor_lens = torch.tensor([anchor], dtype=torch.int32, device=device) + + _, prefill_op = _require_flash_ops() + out = prefill_op( + q, + k_cache, + v_cache, + block_table, + seq_lens, + softmax_scale=scale, + causal=True, + anchor_lens=anchor_lens, + anchored_window=window, + ) + + # Query absolute positions: the last q_len positions of the sequence. + q_pos = torch.arange(kv_len - q_len, kv_len, device=device) + mask = _anchored_mask(q_pos, kv_len, anchor, window, device) + ref = _ref_attention(q[0], k_lin, v_lin, mask, scale) + torch.testing.assert_close(out[0].float(), ref.float(), atol=2e-2, rtol=2e-2) + + +def test_prefill_paged_uses_per_request_anchor_lengths(): + device = "cuda" + dtype = torch.float16 + head_dim, num_heads, num_kv_heads = 128, 8, 4 + kv_len, q_len, block_size, window = 96, 8, 16, 16 + anchors = (16, 32) + scale = head_dim**-0.5 + + k_cache, v_cache, block_table, k_lin, v_lin = _build_paged_kv( + kv_len, block_size, num_kv_heads, head_dim, device, dtype, seed=4 + ) + block_table = block_table.repeat(2, 1) + q = torch.randn(2, q_len, num_heads, head_dim, device=device, dtype=dtype) + seq_lens = torch.full((2,), kv_len, dtype=torch.int32, device=device) + anchor_lens = torch.tensor(anchors, dtype=torch.int32, device=device) + + _, prefill_op = _require_flash_ops() + out = prefill_op( + q, + k_cache, + v_cache, + block_table, + seq_lens, + softmax_scale=scale, + causal=True, + anchor_lens=anchor_lens, + anchored_window=window, + ) + + q_pos = torch.arange(kv_len - q_len, kv_len, device=device) + for batch_idx, anchor in enumerate(anchors): + mask = _anchored_mask(q_pos, kv_len, anchor, window, device) + ref = _ref_attention(q[batch_idx], k_lin, v_lin, mask, scale) + torch.testing.assert_close( + out[batch_idx].float(), ref.float(), atol=2e-2, rtol=2e-2 + ) diff --git a/tests/kernels/attention/test_sm70_flash_v100_decode_planner.py b/tests/kernels/attention/test_sm70_flash_v100_decode_planner.py index d168663a14..9f7f609587 100644 --- a/tests/kernels/attention/test_sm70_flash_v100_decode_planner.py +++ b/tests/kernels/attention/test_sm70_flash_v100_decode_planner.py @@ -188,6 +188,70 @@ def test_static_decode_short_workspace_can_preserve_partition_boundaries( assert plan.workspace_num_partitions == 4 +@pytest.mark.parametrize( + ( + "seq_len", + "max_seq_len_hint", + "workspace_seq_capacity", + "routing_enabled", + "capture_active", + "expected_hint", + ), + ( + (4096, 4096, 8191, True, False, 4096), + (16001, 16001, 262144, True, False, 16001), + (4096, None, 8191, True, False, 8191), + (1, 1, 262144, True, True, 262144), + (4096, 4096, 8191, False, False, 0), + ), +) +def test_xqa_batch_context_route_prefers_live_context( + monkeypatch, + seq_len, + max_seq_len_hint, + workspace_seq_capacity, + routing_enabled, + capture_active, + expected_hint, +) -> None: + _clear_decode_caches() + captured_hints: list[int] = [] + monkeypatch.setattr( + flash_attn_v100, + "_cuda_graph_capture_active", + lambda: capture_active, + ) + + def fake_xqa(*args): + captured_hints.append(args[-1]) + return torch.empty_like(args[0]) + + monkeypatch.setattr( + flash_attn_v100.flash_attn_v100_cuda, + "decode_paged_xqa_fwd", + fake_xqa, + ) + q = torch.empty((4, 6, 256), dtype=torch.float16) + k_cache = torch.empty((2048, 16, 1, 256), dtype=torch.uint8) + v_cache = torch.empty_like(k_cache) + block_table = torch.zeros((4, 2048), dtype=torch.int32) + seq_lens = torch.full((4,), seq_len, dtype=torch.int32) + + flash_attn_v100.flash_attn_decode_paged_xqa( + q, + k_cache, + v_cache, + block_table, + seq_lens, + kv_cache_dtype="fp8_e5m2", + max_seq_len_hint=max_seq_len_hint, + workspace_seq_capacity_hint=workspace_seq_capacity, + batch_context_routing=routing_enabled, + ) + + assert captured_hints == [expected_hint] + + @torch.inference_mode() def test_stale_active_num_partitions_does_not_truncate_decode( monkeypatch, diff --git a/tests/kernels/attention/test_sm70_flash_v100_grouped_verify.py b/tests/kernels/attention/test_sm70_flash_v100_grouped_verify.py new file mode 100644 index 0000000000..84da54590b --- /dev/null +++ b/tests/kernels/attention/test_sm70_flash_v100_grouped_verify.py @@ -0,0 +1,245 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import math + +import pytest +import torch + +_EXPANDED_BASELINE_MEAN_ERROR = { + (16, 8, 120): 5.159488409844926e-6, + (16, 8, 121): 8.621820597909391e-6, + (16, 3, 253): 6.551763817697065e-6, + (16, 4, 252): 5.888879059057217e-6, + (16, 5, 507): 4.691817139246268e-6, + (16, 6, 506): 4.563758920994587e-6, + (16, 8, 1016): 3.1533952551399125e-6, +} + + +def _require_grouped_verify(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + if torch.cuda.get_device_capability() != (7, 0): + pytest.skip("grouped Flash-V100 verification is SM70-only") + flash_attn_v100 = pytest.importorskip("flash_attn_v100") + if not hasattr(flash_attn_v100, "flash_attn_grouped_verify_paged"): + pytest.skip("Flash-V100 extension lacks grouped verification") + return flash_attn_v100 + + +def _make_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, page_size, 1, 256) + key_source = torch.randn(source_shape, dtype=torch.float16, device="cuda").mul_( + 0.25 + ) + value_source = torch.randn_like(key_source).mul_(0.25) + key_cache = key_source.to(torch.float8_e5m2).view(torch.uint8) + value_cache = value_source.to(torch.float8_e5m2).view(torch.uint8) + 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, + value_cache: torch.Tensor, + block_table: torch.Tensor, + seq_len: int, + *, + k_scale: float, + v_scale: float, +) -> torch.Tensor: + query_len = query.shape[0] + prefix_len = seq_len - query_len + logical_pages = math.ceil(seq_len / key_cache.shape[1]) + physical_pages = block_table[0, :logical_pages].to(dtype=torch.long) + key = ( + key_cache.view(torch.float8_e5m2)[physical_pages] + .flatten(0, 1)[:seq_len, 0] + .float() + .mul_(k_scale) + ) + value = ( + value_cache.view(torch.float8_e5m2)[physical_pages] + .flatten(0, 1)[:seq_len, 0] + .float() + .mul_(v_scale) + ) + rows = [] + for token_idx in range(query_len): + visible = prefix_len + token_idx + 1 + scores = torch.einsum( + "hd,nd->hn", query[token_idx].float(), key[:visible] + ).mul_(256**-0.5) + probabilities = torch.softmax(scores, dim=-1) + rows.append(torch.einsum("hn,nd->hd", probabilities, value[:visible])) + return torch.stack(rows).half() + + +@pytest.mark.parametrize( + ("page_size", "query_len", "prefix_len", "strict_fp32_gate"), + [ + (16, 1, 127, True), + (16, 8, 120, False), + (16, 8, 121, False), + (16, 3, 253, False), + (16, 4, 252, False), + (16, 5, 507, False), + (16, 6, 506, False), + (16, 8, 1016, False), + (16, 3, 1025, True), + (784, 5, 2049, True), + (1648, 8, 4097, True), + (3296, 8, 8193, True), + ], +) +@torch.inference_mode() +def test_grouped_verify_matches_fp32_reference_with_random_pages( + page_size: int, + query_len: int, + prefix_len: int, + strict_fp32_gate: bool, +) -> None: + flash_attn_v100 = _require_grouped_verify() + torch.manual_seed(20260824 + page_size + query_len) + query, key_cache, value_cache, block_table, seq_lens = _make_case( + page_size=page_size, + query_len=query_len, + prefix_len=prefix_len, + ) + k_scale = 0.5 + v_scale = 2.0 + expected = _reference( + query, + key_cache, + value_cache, + block_table, + int(seq_lens.item()), + k_scale=k_scale, + v_scale=v_scale, + ) + two_pass = flash_attn_v100.flash_attn_grouped_verify_paged( + query, + key_cache, + value_cache, + block_table, + seq_lens, + k_scale=k_scale, + v_scale=v_scale, + one_pass=False, + ).clone() + one_pass = flash_attn_v100.flash_attn_grouped_verify_paged( + query, + key_cache, + value_cache, + block_table, + seq_lens, + k_scale=k_scale, + v_scale=v_scale, + one_pass=True, + ).clone() + 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, + ) + accepted_xqa = flash_attn_v100.flash_attn_decode_paged_xqa( + query, + key_cache, + value_cache, + decode_block_table, + decode_seq_lens, + kv_cache_dtype="fp8_e5m2", + k_scale=k_scale, + v_scale=v_scale, + max_seq_len_hint=prefix_len + query_len, + workspace_seq_capacity_hint=prefix_len + query_len, + ) + torch.accelerator.synchronize() + + for actual in (two_pass, one_pass): + difference = actual.float().sub(expected.float()).abs() + if strict_fp32_gate: + assert difference.max().item() <= 6.2e-5 + assert difference.mean().item() <= 6.0e-6 + one_pass_error = one_pass.float().sub(expected.float()).abs() + accepted_error = accepted_xqa.float().sub(expected.float()).abs() + assert one_pass_error.max().item() <= accepted_error.max().item() + if strict_fp32_gate: + mean_error_limit = accepted_error.mean().item() + else: + mean_error_limit = _EXPANDED_BASELINE_MEAN_ERROR[ + (page_size, query_len, prefix_len) + ] + assert one_pass_error.mean().item() <= mean_error_limit + 5.0e-7 + torch.testing.assert_close(one_pass, two_pass, atol=6.2e-5, rtol=2.0e-3) + + +@torch.inference_mode() +def test_grouped_verify_cuda_graph_replay_tracks_runtime_seq_len() -> 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, + prefix_len=4097, + ) + output = torch.empty_like(query) + flash_attn_v100.flash_attn_grouped_verify_paged( + query, + key_cache, + value_cache, + block_table, + seq_lens, + out=output, + one_pass=True, + ) + torch.accelerator.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + flash_attn_v100.flash_attn_grouped_verify_paged( + query, + key_cache, + value_cache, + block_table, + seq_lens, + out=output, + one_pass=True, + ) + + for prefix_len in (127, 2049, 4097): + seq_lens.fill_(prefix_len + query.shape[0]) + graph.replay() + torch.accelerator.synchronize() + expected = _reference( + query, + key_cache, + value_cache, + block_table, + prefix_len + query.shape[0], + k_scale=1.0, + v_scale=1.0, + ) + difference = output.float().sub(expected.float()).abs() + assert difference.max().item() <= 6.2e-5 + assert difference.mean().item() <= 6.0e-6 diff --git a/tests/kernels/attention/test_sm70_flash_v100_smallq_metadata.py b/tests/kernels/attention/test_sm70_flash_v100_smallq_metadata.py new file mode 100644 index 0000000000..3912f0f614 --- /dev/null +++ b/tests/kernels/attention/test_sm70_flash_v100_smallq_metadata.py @@ -0,0 +1,232 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm.v1.attention.backends.flash_attn_v100 import ( + _sm70_prepare_grouped_smallq_decode_metadata, + _sm70_prepare_smallq_decode_metadata, +) + + +def _reference_smallq_metadata( + block_table: torch.Tensor, + seq_lens: torch.Tensor, + query_start_loc: torch.Tensor, + *, + num_query_tokens: int, + real_num_query_tokens: int, +) -> tuple[torch.Tensor, torch.Tensor]: + query_lens = query_start_loc[1:] - query_start_loc[:-1] + repeat_query_lens = query_lens.clone() + repeat_query_lens[-1] += num_query_tokens - real_num_query_tokens + effective_seq_lens = torch.maximum(seq_lens, query_lens) + decode_block_table = torch.repeat_interleave( + block_table.clamp_min(0), + repeat_query_lens, + dim=0, + output_size=num_query_tokens, + ).contiguous() + seq_lens_rep = torch.repeat_interleave( + effective_seq_lens, + repeat_query_lens, + output_size=num_query_tokens, + ) + query_lens_rep = torch.repeat_interleave( + query_lens, + repeat_query_lens, + output_size=num_query_tokens, + ) + start_locs_rep = torch.repeat_interleave( + query_start_loc[:-1], + repeat_query_lens, + output_size=num_query_tokens, + ) + token_indices = torch.arange( + num_query_tokens, dtype=seq_lens.dtype, device=seq_lens.device + ) + decode_seq_lens = ( + seq_lens_rep - query_lens_rep + token_indices - start_locs_rep + 1 + ).contiguous() + if real_num_query_tokens < num_query_tokens: + padding_mask = token_indices >= real_num_query_tokens + decode_seq_lens = decode_seq_lens.masked_fill(padding_mask, 0) + decode_block_table = decode_block_table.masked_fill(padding_mask[:, None], 0) + return decode_block_table, decode_seq_lens + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize( + ("query_lens", "seq_lens", "num_query_tokens"), + [ + ([8], [128], 8), + ([3, 2, 0], [8, 7, 0], 6), + ([1, 4, 2, 0], [1, 10, 8, 0], 8), + ], +) +def test_sm70_fused_smallq_metadata_matches_repeat_interleave( + query_lens: list[int], + seq_lens: list[int], + num_query_tokens: int, +): + if torch.cuda.get_device_capability() != (7, 0): + pytest.skip("SM70 CUDA device required") + + device = torch.device("cuda") + num_reqs = len(query_lens) + query_start_loc = torch.tensor( + [0, *torch.tensor(query_lens).cumsum(0).tolist()], + dtype=torch.int32, + device=device, + ) + real_num_query_tokens = sum(query_lens) + seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int32, device=device) + block_table = torch.arange( + num_reqs * 7, + dtype=torch.int32, + device=device, + ).view(num_reqs, 7) + block_table[:, 0] = -1 + out_block_table = torch.empty( + (num_query_tokens, block_table.shape[1]), + dtype=torch.int32, + device=device, + ) + out_seq_lens = torch.empty(num_query_tokens, dtype=torch.int32, device=device) + out_query_start_loc = torch.empty(num_reqs + 1, dtype=torch.int32, device=device) + + _sm70_prepare_smallq_decode_metadata( + out_block_table, + out_seq_lens, + out_query_start_loc, + block_table, + seq_lens_tensor, + query_start_loc, + num_reqs=num_reqs, + num_query_tokens=num_query_tokens, + real_num_query_tokens=real_num_query_tokens, + ) + expected_block_table, expected_seq_lens = _reference_smallq_metadata( + block_table, + seq_lens_tensor, + query_start_loc, + num_query_tokens=num_query_tokens, + real_num_query_tokens=real_num_query_tokens, + ) + torch.accelerator.synchronize() + + assert torch.equal(out_block_table, expected_block_table) + assert torch.equal(out_seq_lens, expected_seq_lens) + assert torch.equal(out_query_start_loc, query_start_loc) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_sm70_grouped_smallq_metadata_matches_individual_launches(): + if torch.cuda.get_device_capability() != (7, 0): + pytest.skip("SM70 CUDA device required") + + device = torch.device("cuda") + num_groups = 5 + num_reqs = 3 + num_query_tokens = 8 + real_num_query_tokens = 6 + block_cols_by_group = [3, 5, 4, 5, 3] + query_start_loc = torch.tensor([0, 3, 5, 6], dtype=torch.int32, device=device) + seq_lens = torch.tensor([19, 13, 9], dtype=torch.int32, device=device) + block_tables = [ + ( + torch.arange( + num_reqs * block_cols, + dtype=torch.int32, + device=device, + ) + .view(num_reqs, block_cols) + .add_(group_id * 1000) + ) + for group_id, block_cols in enumerate(block_cols_by_group) + ] + for table in block_tables: + table[:, 0] = -1 + + grouped_tables = [ + torch.empty((num_query_tokens, block_cols), dtype=torch.int32, device=device) + for block_cols in block_cols_by_group + ] + grouped_seq_lens = [ + torch.empty(num_query_tokens, dtype=torch.int32, device=device) + for _ in range(num_groups) + ] + grouped_query_locs = [ + torch.empty(num_reqs + 1, dtype=torch.int32, device=device) + for _ in range(num_groups) + ] + descriptor = _sm70_prepare_grouped_smallq_decode_metadata( + grouped_tables, + grouped_seq_lens, + grouped_query_locs, + block_tables, + seq_lens, + query_start_loc, + num_reqs=num_reqs, + num_query_tokens=num_query_tokens, + real_num_query_tokens=real_num_query_tokens, + ) + + reference_tables = [] + reference_seq_lens = [] + for table in block_tables: + out_table = torch.empty( + (num_query_tokens, table.shape[1]), dtype=torch.int32, device=device + ) + out_seq = torch.empty_like(grouped_seq_lens[0]) + out_query = torch.empty_like(grouped_query_locs[0]) + _sm70_prepare_smallq_decode_metadata( + out_table, + out_seq, + out_query, + table, + seq_lens, + query_start_loc, + num_reqs=num_reqs, + num_query_tokens=num_query_tokens, + real_num_query_tokens=real_num_query_tokens, + ) + reference_tables.append(out_table) + reference_seq_lens.append(out_seq) + assert torch.equal(out_query, query_start_loc) + torch.accelerator.synchronize() + + for grouped, reference in zip(grouped_tables, reference_tables): + assert torch.equal(grouped, reference) + for grouped, reference in zip(grouped_seq_lens, reference_seq_lens): + assert torch.equal(grouped, reference) + for grouped in grouped_query_locs: + assert torch.equal(grouped, query_start_loc) + + # Reuse the persistent pointer tables and prove that no stale group output + # survives a replay with different sequence metadata. + seq_lens.add_(7) + replay_descriptor = _sm70_prepare_grouped_smallq_decode_metadata( + grouped_tables, + grouped_seq_lens, + grouped_query_locs, + block_tables, + seq_lens, + query_start_loc, + num_reqs=num_reqs, + num_query_tokens=num_query_tokens, + real_num_query_tokens=real_num_query_tokens, + descriptor=descriptor, + ) + torch.accelerator.synchronize() + assert replay_descriptor is descriptor + _, expected_seq_lens = _reference_smallq_metadata( + block_tables[0], + seq_lens, + query_start_loc, + num_query_tokens=num_query_tokens, + real_num_query_tokens=real_num_query_tokens, + ) + for grouped in grouped_seq_lens: + assert torch.equal(grouped, expected_seq_lens) 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/attention/test_sm70_flash_v100_varlen_layout.py b/tests/kernels/attention/test_sm70_flash_v100_varlen_layout.py index 83d0cd761b..df09e884c8 100644 --- a/tests/kernels/attention/test_sm70_flash_v100_varlen_layout.py +++ b/tests/kernels/attention/test_sm70_flash_v100_varlen_layout.py @@ -624,3 +624,333 @@ def test_sm70_flash_v100_fp8_e5m2_xqa_decode_matches_scalar( ).float() assert torch.all(delta <= torch.maximum(one_ulp, torch.full_like(delta, 1e-4))) assert float(delta.mean()) <= 1e-5 + + +@pytest.mark.parametrize(("k_scale", "v_scale"), ((1.0, 1.0), (0.04, 0.25))) +@pytest.mark.parametrize("partition_size", (64, 128, 256)) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@torch.inference_mode() +def test_sm70_flash_v100_fp8_e4m3_g6_xqa_decode_matches_scalar( + k_scale, + v_scale, + partition_size, +): + 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(20260823) + seq_len = 2049 + block_size = 1568 + head_dim = 256 + cache_shape = (3, block_size, 1, head_dim) + key_cache = ( + torch.randn(cache_shape, dtype=torch.float16, device="cuda") + .to(torch.float8_e4m3fn) + .view(torch.uint8) + ) + value_cache = ( + torch.randn(cache_shape, dtype=torch.float16, device="cuda") + .to(torch.float8_e4m3fn) + .view(torch.uint8) + ) + query = torch.randn(1, 6, head_dim, dtype=torch.float16, device="cuda") + block_table = torch.tensor([[2, 0]], dtype=torch.int32, device="cuda") + seq_lens = torch.tensor([seq_len], dtype=torch.int32, device="cuda") + kwargs = { + "kv_cache_dtype": "fp8_e4m3", + "k_scale": k_scale, + "v_scale": v_scale, + "max_seq_len_hint": seq_len, + } + + expected = flash_attn_v100.flash_attn_decode_paged( + query, + key_cache, + value_cache, + block_table, + seq_lens, + partition_size_hint=256, + **kwargs, + ) + actual = flash_attn_v100.flash_attn_decode_paged_xqa( + query, + key_cache, + value_cache, + block_table, + seq_lens, + partition_size_hint=partition_size, + **kwargs, + ) + + torch.accelerator.synchronize() + delta = (actual.float() - expected.float()).abs() + positive = torch.full_like(expected, torch.inf) + negative = torch.full_like(expected, -torch.inf) + one_ulp = torch.maximum( + (torch.nextafter(expected, positive) - expected).abs(), + (torch.nextafter(expected, negative) - expected).abs(), + ).float() + assert torch.all(delta <= torch.maximum(one_ulp, torch.full_like(delta, 1e-4))) + assert float(delta.mean()) <= 1e-5 + + +@pytest.mark.parametrize("batch_size", (2, 16)) +@pytest.mark.parametrize("partition_size", (64, 128, 256)) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@torch.inference_mode() +def test_sm70_flash_v100_fp8_e4m3_batched_g6_xqa_matches_scalar( + monkeypatch, + batch_size, + partition_size, +): + if torch.cuda.get_device_capability() != (7, 0): + pytest.skip("FlashAttention-V100 regression is SM70/V100 only") + + monkeypatch.setenv("VLLM_FLASH_V100_E4M3_BATCH_XQA", "1") + flash_attn_v100 = pytest.importorskip("flash_attn_v100") + torch.manual_seed(20260824) + block_size = 1568 + head_dim = 256 + blocks_per_seq = 2 + num_blocks = batch_size * blocks_per_seq + cache_shape = (num_blocks, block_size, 1, head_dim) + key_cache = ( + torch.randn(cache_shape, dtype=torch.float16, device="cuda") + .to(torch.float8_e4m3fn) + .view(torch.uint8) + ) + value_cache = ( + torch.randn(cache_shape, dtype=torch.float16, device="cuda") + .to(torch.float8_e4m3fn) + .view(torch.uint8) + ) + query = torch.randn( + batch_size, + 6, + head_dim, + dtype=torch.float16, + device="cuda", + ) + block_table = ( + torch.arange(num_blocks, dtype=torch.int32, device="cuda") + .view(batch_size, blocks_per_seq) + .flip(1) + ) + seq_lens = 1025 + (torch.arange(batch_size, device="cuda") * 61) % 900 + seq_lens = seq_lens.to(torch.int32) + kwargs = { + "kv_cache_dtype": "fp8_e4m3", + "k_scale": 0.04, + "v_scale": 0.25, + "max_seq_len_hint": int(seq_lens.max().item()), + } + + expected = flash_attn_v100.flash_attn_decode_paged( + query, + key_cache, + value_cache, + block_table, + seq_lens, + partition_size_hint=256, + **kwargs, + ) + actual = flash_attn_v100.flash_attn_decode_paged_xqa( + query, + key_cache, + value_cache, + block_table, + seq_lens, + partition_size_hint=partition_size, + **kwargs, + ) + + torch.accelerator.synchronize() + delta = (actual.float() - expected.float()).abs() + positive = torch.full_like(expected, torch.inf) + negative = torch.full_like(expected, -torch.inf) + one_ulp = torch.maximum( + (torch.nextafter(expected, positive) - expected).abs(), + (torch.nextafter(expected, negative) - expected).abs(), + ).float() + assert torch.all(delta <= torch.maximum(one_ulp, torch.full_like(delta, 1e-4))) + assert float(delta.mean()) <= 1e-5 + + +@pytest.mark.parametrize("batch_size", (8, 16)) +@pytest.mark.parametrize("partition_size", (64, 128, 256)) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@torch.inference_mode() +def test_sm70_flash_v100_fp8_e4m3_batched_long_context_route_matches_scalar( + monkeypatch, + batch_size, + partition_size, +): + if torch.cuda.get_device_capability() != (7, 0): + pytest.skip("FlashAttention-V100 regression is SM70/V100 only") + + monkeypatch.setenv("VLLM_FLASH_V100_E4M3_BATCH_XQA", "1") + monkeypatch.delenv("VLLM_FLASH_V100_E4M3_BATCH_XQA_OPTIMIZED", raising=False) + monkeypatch.setenv("VLLM_FLASH_V100_E4M3_PAGE800_FASTPATH_TRACE", "0") + flash_attn_v100 = pytest.importorskip("flash_attn_v100") + torch.manual_seed(20260824 + batch_size) + block_size = 800 + head_dim = 256 + seq_lens_cpu = 15360 + torch.arange(batch_size, dtype=torch.int32) * 113 + max_seq_len = int(seq_lens_cpu.max().item()) + blocks_per_seq = (max_seq_len + block_size - 1) // block_size + num_blocks = batch_size * blocks_per_seq + cache_shape = (num_blocks, 2, block_size, 1, head_dim) + kv_cache = ( + torch.randn(cache_shape, dtype=torch.float16, device="cuda") + .to(torch.float8_e4m3fn) + .view(torch.uint8) + ) + key_cache, value_cache = kv_cache.unbind(1) + query = torch.randn( + batch_size, + 6, + head_dim, + dtype=torch.float16, + device="cuda", + ) + block_table = ( + torch.arange(num_blocks, dtype=torch.int32, device="cuda") + .view(batch_size, blocks_per_seq) + .flip(1) + ) + seq_lens = seq_lens_cpu.to(device="cuda") + kwargs = { + "kv_cache_dtype": "fp8_e4m3", + "k_scale": 0.04, + "v_scale": 0.25, + "max_seq_len_hint": max_seq_len, + "workspace_seq_capacity_hint": 32768, + } + + expected = flash_attn_v100.flash_attn_decode_paged( + query, + key_cache, + value_cache, + block_table, + seq_lens, + partition_size_hint=256, + **kwargs, + ) + monkeypatch.setenv("VLLM_FLASH_V100_E4M3_PAGE800_FASTPATH", "0") + fallback = flash_attn_v100.flash_attn_decode_paged_xqa( + query, + key_cache, + value_cache, + block_table, + seq_lens, + partition_size_hint=partition_size, + batch_context_routing=True, + **kwargs, + ) + monkeypatch.setenv("VLLM_FLASH_V100_E4M3_PAGE800_FASTPATH", "1") + actual = flash_attn_v100.flash_attn_decode_paged_xqa( + query, + key_cache, + value_cache, + block_table, + seq_lens, + partition_size_hint=partition_size, + batch_context_routing=True, + **kwargs, + ) + + torch.accelerator.synchronize() + assert torch.equal(actual, fallback) + delta = (actual.float() - expected.float()).abs() + positive = torch.full_like(expected, torch.inf) + negative = torch.full_like(expected, -torch.inf) + one_ulp = torch.maximum( + (torch.nextafter(expected, positive) - expected).abs(), + (torch.nextafter(expected, negative) - expected).abs(), + ).float() + assert torch.all(delta <= torch.maximum(one_ulp, torch.full_like(delta, 1e-4))) + assert float(delta.mean()) <= 1e-5 + + +@pytest.mark.parametrize( + ("seq_len", "expected_partition"), + ( + (8191, 64), + (8192, 256), + (16383, 256), + (16384, 512), + (32767, 512), + (32768, 896), + (49151, 896), + (49152, 1664), + ), +) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@torch.inference_mode() +def test_sm70_flash_v100_fp8_e4m3_g6_routes_wave_partitions_exactly( + monkeypatch, + seq_len, + expected_partition, +): + 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") + monkeypatch.delenv("VLLM_FLASH_V100_DECODE_PARTITION_SIZE", raising=False) + monkeypatch.setenv("VLLM_FLASH_V100_XQA_E4M3_G6_P256_BEGIN", "8192") + monkeypatch.setenv("VLLM_FLASH_V100_XQA_E4M3_G6_P512_BEGIN", "16384") + monkeypatch.setenv("VLLM_FLASH_V100_XQA_E4M3_G6_P896_BEGIN", "32768") + monkeypatch.setenv("VLLM_FLASH_V100_XQA_E4M3_G6_P1664_BEGIN", "49152") + monkeypatch.setenv("VLLM_FLASH_V100_XQA_E4M3_G6_WAVE_PARTITIONS", "1") + monkeypatch.setenv("VLLM_FLASH_V100_XQA_E4M3_G6_MERGED_WAVE_LAUNCH", "1") + + torch.manual_seed(20260824) + block_size = 1568 + head_dim = 256 + num_blocks = (seq_len + block_size - 1) // block_size + cache_shape = (num_blocks, 2, block_size, 1, head_dim) + kv_cache = ( + torch.randn(cache_shape, dtype=torch.float16, device="cuda") + .to(torch.float8_e4m3fn) + .view(torch.uint8) + ) + key_cache, value_cache = kv_cache.unbind(1) + assert not key_cache.is_contiguous() + assert not value_cache.is_contiguous() + query = torch.randn(1, 6, head_dim, dtype=torch.float16, device="cuda") + block_table = torch.arange(num_blocks, dtype=torch.int32, device="cuda").view(1, -1) + seq_lens = torch.tensor([seq_len], dtype=torch.int32, device="cuda") + kwargs = { + "kv_cache_dtype": "fp8_e4m3", + "k_scale": 0.04, + "v_scale": 0.25, + "max_seq_len_hint": seq_len, + "workspace_seq_capacity_hint": 262144, + } + + monkeypatch.setenv("VLLM_FLASH_V100_XQA_E4M3_G6_P64_P256_AUTO", "0") + expected = flash_attn_v100.flash_attn_decode_paged_xqa( + query, + key_cache, + value_cache, + block_table, + seq_lens, + partition_size_hint=expected_partition, + **kwargs, + ).clone() + + monkeypatch.setenv("VLLM_FLASH_V100_XQA_E4M3_G6_P64_P256_AUTO", "1") + actual = flash_attn_v100.flash_attn_decode_paged_xqa( + query, + key_cache, + value_cache, + block_table, + seq_lens, + partition_size_hint=64, + batch_context_routing=True, + **kwargs, + ) + + torch.accelerator.synchronize() + assert torch.equal(actual, expected) diff --git a/tests/kernels/core/test_batched_weight_rms_norm.py b/tests/kernels/core/test_batched_weight_rms_norm.py new file mode 100644 index 0000000000..790367eaad --- /dev/null +++ b/tests/kernels/core/test_batched_weight_rms_norm.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the batched-weight RMS norm kernel. + +The outermost input dimension selects the corresponding weight row. DFlash +uses this for its fused per-layer K normalization. +""" + +import pytest +import torch + +from vllm import _custom_ops as ops +from vllm.platforms import current_platform +from vllm.utils.torch_utils import set_random_seed + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="rms_norm requires a CUDA/ROCm device", +) + + +@pytest.mark.parametrize( + "shape", + [ + (28, 17, 128), + (1, 5, 2, 128), + (28, 13, 8, 128), + (6, 3, 4, 769), + ], +) +@pytest.mark.parametrize("dtype", [torch.half, torch.bfloat16, torch.float]) +@pytest.mark.parametrize("seed", [42]) +@torch.inference_mode() +def test_rms_norm_matches_loop( + shape: tuple[int, ...], dtype: torch.dtype, seed: int +) -> None: + set_random_seed(seed) + torch.set_default_device("cuda") + + num_rows, hidden = shape[0], shape[-1] + eps = 1e-6 + x = torch.randn(*shape, dtype=dtype) * 0.1 + weight = torch.randn(num_rows, hidden, dtype=dtype) * 0.1 + 1.0 + + out_ref = torch.empty_like(x) + for i in range(x.shape[0]): + ops.rms_norm(out_ref[i], x[i], weight[i], eps) + + out = torch.empty_like(x) + ops.rms_norm(out, x, weight, eps) + torch.testing.assert_close(out, out_ref, atol=0, rtol=0) + + +@torch.inference_mode() +def test_rms_norm_validates_shapes() -> None: + torch.set_default_device("cuda") + + x = torch.randn(4, 8, 128, dtype=torch.float) + out = torch.empty_like(x) + with pytest.raises(RuntimeError): + ops.rms_norm(out, x, torch.randn(3, 128), 1e-6) + with pytest.raises(RuntimeError): + ops.rms_norm(out, x, torch.randn(4, 64), 1e-6) diff --git a/tests/kernels/core/test_sm70_dflash2_gemma_rms.py b/tests/kernels/core/test_sm70_dflash2_gemma_rms.py new file mode 100644 index 0000000000..87c0eda3da --- /dev/null +++ b/tests/kernels/core/test_sm70_dflash2_gemma_rms.py @@ -0,0 +1,121 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm.model_executor.layers.layernorm import ( + GemmaRMSNorm, + _sm70_dflash2_gemma_fused_add_rms_norm, + _use_sm70_dflash2_gemma_fused_add_rms, +) + + +@pytest.mark.parametrize("num_tokens", [1, 8, 32, 137, 512]) +@pytest.mark.parametrize("weight_dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_sm70_dflash2_gemma_fused_add_rms_is_within_one_fp16_ulp( + num_tokens: int, + weight_dtype: torch.dtype, +): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0): + pytest.skip("SM70 CUDA device required") + + torch.manual_seed(20260823) + x = torch.randn((num_tokens, 5120), dtype=torch.float16, device="cuda").mul_(0.125) + residual = torch.randn((num_tokens, 5120), dtype=torch.float32, device="cuda").mul_( + 0.125 + ) + weight = torch.randn(5120, dtype=torch.float32, device="cuda").mul_(0.05) + weight = weight.to(weight_dtype) + + expected_normalized, expected_residual = GemmaRMSNorm._forward_static_with_residual( + weight, 1e-6, x, residual + ) + actual_normalized, actual_residual = _sm70_dflash2_gemma_fused_add_rms_norm( + x, residual, weight, 1e-6 + ) + + lower = torch.nextafter( + expected_normalized, + torch.full_like(expected_normalized, -float("inf")), + ) + upper = torch.nextafter( + expected_normalized, + torch.full_like(expected_normalized, float("inf")), + ) + assert torch.all((actual_normalized >= lower) & (actual_normalized <= upper)) + assert torch.equal(actual_residual, expected_residual) + + +def test_sm70_dflash2_gemma_fused_add_rms_graph_replay_reads_current_inputs(): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0): + pytest.skip("SM70 CUDA device required") + + torch.manual_seed(20260823) + x = torch.randn((8, 5120), dtype=torch.float16, device="cuda").mul_(0.125) + residual = torch.randn((8, 5120), dtype=torch.float32, device="cuda").mul_(0.125) + weight = torch.randn(5120, dtype=torch.float32, device="cuda").mul_(0.05) + _sm70_dflash2_gemma_fused_add_rms_norm(x, residual, weight, 1e-6) + torch.accelerator.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual_normalized, actual_residual = _sm70_dflash2_gemma_fused_add_rms_norm( + x, residual, weight, 1e-6 + ) + + x.copy_(torch.randn_like(x).mul_(0.125)) + residual.copy_(torch.randn_like(residual).mul_(0.125)) + weight.copy_(torch.randn_like(weight).mul_(0.05)) + graph.replay() + torch.accelerator.synchronize() + + expected_normalized, expected_residual = GemmaRMSNorm._forward_static_with_residual( + weight, 1e-6, x, residual + ) + lower = torch.nextafter( + expected_normalized, + torch.full_like(expected_normalized, -float("inf")), + ) + upper = torch.nextafter( + expected_normalized, + torch.full_like(expected_normalized, float("inf")), + ) + assert torch.all((actual_normalized >= lower) & (actual_normalized <= upper)) + assert torch.equal(actual_residual, expected_residual) + + +def test_sm70_dflash2_gemma_fused_add_rms_allows_aot_warmup_shape(monkeypatch): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0): + pytest.skip("SM70 CUDA device required") + + monkeypatch.setenv("VLLM_SM70_DFLASH2_FUSED_GEMMA_RMS", "1") + monkeypatch.setenv("VLLM_SM70_FLASH_V100_0DOT3_COMPILE_GRAPH", "1") + x = torch.empty((512, 5120), dtype=torch.float16, device="cuda") + residual = torch.empty_like(x, dtype=torch.float32) + weight = torch.empty(5120, dtype=torch.float16, device="cuda") + + assert _use_sm70_dflash2_gemma_fused_add_rms(x, residual, weight) + + +def test_sm70_dflash2_gemma_fused_predicate_is_fullgraph_safe(monkeypatch): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0): + pytest.skip("SM70 CUDA device required") + + monkeypatch.setenv("VLLM_SM70_DFLASH2_FUSED_GEMMA_RMS", "1") + monkeypatch.setenv("VLLM_SM70_FLASH_V100_0DOT3_COMPILE_GRAPH", "1") + x = torch.zeros((8, 5120), dtype=torch.float16, device="cuda") + residual = torch.empty_like(x, dtype=torch.float32) + weight = torch.empty(5120, dtype=torch.float16, device="cuda") + + @torch.compile(backend="eager", fullgraph=True) + def apply_predicate( + values: torch.Tensor, + residual_values: torch.Tensor, + norm_weight: torch.Tensor, + ) -> torch.Tensor: + if _use_sm70_dflash2_gemma_fused_add_rms(values, residual_values, norm_weight): + return values + 1 + return values + + torch.testing.assert_close(apply_predicate(x, residual, weight), x + 1) diff --git a/tests/kernels/mamba/test_mrv2_mamba_align.py b/tests/kernels/mamba/test_mrv2_mamba_align.py new file mode 100644 index 0000000000..f6bb99f974 --- /dev/null +++ b/tests/kernels/mamba/test_mrv2_mamba_align.py @@ -0,0 +1,196 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + +import pytest +import torch + +from vllm.platforms import current_platform +from vllm.triton_utils import triton +from vllm.v1.worker.gpu.mamba_align import ( + preprocess_mamba_align_fused_kernel, + run_mamba_align_postprocess, + run_mamba_align_precopy, +) + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda(), reason="MRV2 Mamba align needs CUDA/Triton" +) + + +def _make_context( + conv: torch.Tensor, + temporal: torch.Tensor, + block_table: torch.Tensor, + *, + block_size: int = 4, + max_num_reqs: int | None = None, +) -> SimpleNamespace: + states = (conv, temporal) + device = block_table.device + return SimpleNamespace( + is_initialized=True, + num_layers=1, + num_state_types=2, + block_table_ptrs=torch.tensor( + [block_table.data_ptr()], dtype=torch.int64, device=device + ), + block_table_stride_req=block_table.stride(0), + state_base_addrs=torch.tensor( + [state.data_ptr() for state in states], + dtype=torch.int64, + device=device, + ), + state_block_strides=torch.tensor( + [state.stride(0) * state.element_size() for state in states], + dtype=torch.int64, + device=device, + ), + state_elem_sizes=torch.tensor( + [state.element_size() for state in states], + dtype=torch.int32, + device=device, + ), + state_inner_sizes=torch.tensor( + [conv.stride(1), temporal[0].numel()], + dtype=torch.int64, + device=device, + ), + state_conv_widths=torch.tensor( + [conv.size(1), 0], dtype=torch.int32, device=device + ), + state_group_indices=torch.zeros(2, dtype=torch.int32, device=device), + block_size=block_size, + num_accepted_tokens_out=torch.empty( + max_num_reqs or block_table.shape[0], + dtype=torch.int32, + device=device, + ), + ) + + +def test_mrv2_align_prefix_seed_uses_resolved_mamba_block_size() -> None: + device = torch.device("cuda") + # Batch order is intentionally different from persistent request-slot order. + idx_mapping = torch.tensor([1, 0], dtype=torch.int32, device=device) + state_idx = torch.full((2,), -2, dtype=torch.int32, device=device) + num_computed = torch.tensor([8, 0], dtype=torch.int32, device=device) + query_start_loc = torch.tensor([0, 1, 2], dtype=torch.int32, device=device) + num_accepted = torch.ones(2, dtype=torch.int32, device=device) + src_col = torch.empty(2, dtype=torch.int32, device=device) + token_bias = torch.empty(2, dtype=torch.int32, device=device) + + preprocess_mamba_align_fused_kernel[(1,)]( + idx_mapping, + state_idx, + num_computed, + query_start_loc, + num_accepted, + src_col, + token_bias, + 2, + BLOCK_SIZE=triton.next_power_of_2(2), + MAMBA_BLOCK_SIZE=8, + ) + torch.accelerator.synchronize() + + # req0 resumes after one complete Mamba block and crosses into column 1; + # req1 is fresh and has no source state before entering column 0. + assert src_col.cpu().tolist() == [0, -1] + assert state_idx.cpu().tolist() == [1, 0] + assert token_bias.cpu().tolist() == [0, 0] + + +@pytest.mark.parametrize("token_bias", [0, 1, 2]) +def test_mrv2_align_precopy_matches_v1_sd_semantics(token_bias: int) -> None: + torch.manual_seed(42) + device = torch.device("cuda") + num_reqs, num_cols = 4, 6 + conv_width, conv_dim = 6, 128 + num_blocks = num_reqs * num_cols + 1 + block_table = torch.arange(1, num_blocks, dtype=torch.int32, device=device).reshape( + num_reqs, num_cols + ) + conv = torch.randn( + num_blocks, conv_width, conv_dim, dtype=torch.float16, device=device + ) + temporal = torch.randn(num_blocks, 4, 16, 16, dtype=torch.float32, device=device) + conv_before = conv.clone() + temporal_before = temporal.clone() + + # State arrays use persistent request-slot order; block-table rows use + # current batch order. The non-identity mapping proves the two are not mixed. + src_col = torch.tensor([-1, 1, 1, 2], dtype=torch.int32, device=device) + dst_col = torch.tensor([0, 1, 2, 0], dtype=torch.int32, device=device) + bias = torch.tensor( + [0, 0, token_bias, token_bias], dtype=torch.int32, device=device + ) + idx_mapping = torch.tensor([2, 0, 3, 1], dtype=torch.int32, device=device) + + run_mamba_align_precopy( + _make_context(conv, temporal, block_table), + num_reqs, + dst_col, + src_col, + bias, + idx_mapping, + ) + torch.accelerator.synchronize() + + conv_ref = conv_before.clone() + temporal_ref = temporal_before.clone() + for batch_idx, req_idx in enumerate(idx_mapping.cpu().tolist()): + src = int(src_col[req_idx]) + dst = int(dst_col[req_idx]) + accepted_bias = int(bias[req_idx]) + if src < 0 or src == dst: + continue + src_block = int(block_table[batch_idx, src]) + dst_block = int(block_table[batch_idx, dst]) + temporal_block = int(block_table[batch_idx, src + accepted_bias]) + conv_ref[dst_block, : conv_width - accepted_bias] = conv_before[ + src_block, accepted_bias: + ] + temporal_ref[dst_block] = temporal_before[temporal_block] + + torch.testing.assert_close(conv, conv_ref, rtol=0, atol=0) + torch.testing.assert_close(temporal, temporal_ref, rtol=0, atol=0) + + +def test_mrv2_align_postprocess_same_block_shift_is_exact() -> None: + torch.manual_seed(43) + device = torch.device("cuda") + conv_width, conv_dim = 6, 128 + block_table = torch.tensor([[1, 2, 3, 4, 5]], dtype=torch.int32, device=device) + conv = torch.randn(6, conv_width, conv_dim, dtype=torch.float16, device=device) + temporal = torch.randn(6, 4, 16, 16, dtype=torch.float32, device=device) + conv_before = conv.clone() + temporal_before = temporal.clone() + + # accepted=3 and new_computed=8 means running=6, aligned=8, bias=2. + # src_col == dst_col == 1 exercises the overlapping conv memmove. + accepted = torch.tensor([3, 1, 1, 1], dtype=torch.int32, device=device) + state_idx = torch.tensor([1, 0, 0, 0], dtype=torch.int32, device=device) + new_computed = torch.tensor([8, 0, 0, 0], dtype=torch.int32, device=device) + idx_mapping = torch.tensor([0], dtype=torch.int32, device=device) + run_mamba_align_postprocess( + _make_context( + conv, temporal, block_table, block_size=4, max_num_reqs=accepted.numel() + ), + 1, + accepted, + state_idx, + new_computed, + idx_mapping, + ) + torch.accelerator.synchronize() + + dst_block = int(block_table[0, 1]) + temporal_src_block = int(block_table[0, 3]) + conv_ref = conv_before.clone() + temporal_ref = temporal_before.clone() + conv_ref[dst_block, : conv_width - 2] = conv_before[dst_block, 2:] + temporal_ref[dst_block] = temporal_before[temporal_src_block] + torch.testing.assert_close(conv, conv_ref, rtol=0, atol=0) + torch.testing.assert_close(temporal, temporal_ref, rtol=0, atol=0) + assert int(accepted[0]) == 1 diff --git a/tests/kernels/moe/test_moe_permute_unpermute.py b/tests/kernels/moe/test_moe_permute_unpermute.py index 662da0fb62..dc596f8795 100644 --- a/tests/kernels/moe/test_moe_permute_unpermute.py +++ b/tests/kernels/moe/test_moe_permute_unpermute.py @@ -463,6 +463,7 @@ def test_sm70_single_token_weighted_reduce_matches_moe_unpermute(): ] cases = [ ("topk2_h256_contiguous", 2, 256, 0), + ("topk6_h4096_contiguous", 6, 4096, 0), ("topk8_h2048_contiguous", 8, 2048, 0), ("topk8_h7168_padded_stride", 8, 7168, 64), ] @@ -473,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 ) @@ -483,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/moe/test_sm70_unquantized_moe_config.py b/tests/kernels/moe/test_sm70_unquantized_moe_config.py new file mode 100644 index 0000000000..ed024de854 --- /dev/null +++ b/tests/kernels/moe/test_sm70_unquantized_moe_config.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +import vllm.envs as envs +import vllm.model_executor.layers.fused_moe.fused_moe as fused_moe_module +from vllm.model_executor.layers.fused_moe.fused_moe import ( + _get_sm70_mtp_moe_decode_config, + force_sm70_mtp_moe_legacy_config, +) + + +@pytest.fixture(autouse=True) +def _enable_tuned_mtp_config(monkeypatch): + monkeypatch.setattr(envs, "VLLM_SM70_MTP_MOE_TUNED_CONFIG", True) + + +def test_mtp_sm70_decode_config_keeps_legacy_tile_at_m1(): + assert _get_sm70_mtp_moe_decode_config(1, 256, 128, 2048, 8) is None + + +@pytest.mark.parametrize("m", range(2, 17)) +def test_mtp_sm70_decode_config_uses_exact_local_tile(m): + config = _get_sm70_mtp_moe_decode_config(m, 256, 128, 2048, 8) + + assert config is not None + assert config["BLOCK_SIZE_M"] == 8 + assert config["BLOCK_SIZE_N"] == 128 + assert config["BLOCK_SIZE_K"] == 32 + + +@pytest.mark.parametrize( + "shape", + [ + (17, 256, 128, 2048, 8), + (2, 128, 128, 2048, 8), + (2, 256, 512, 2048, 8), + (2, 256, 128, 4096, 8), + (2, 256, 128, 2048, 4), + ], +) +def test_mtp_sm70_decode_config_is_shape_bounded(shape): + assert _get_sm70_mtp_moe_decode_config(*shape) is None + + +def test_mtp_sm70_decode_config_can_be_disabled(monkeypatch): + monkeypatch.setattr(envs, "VLLM_SM70_MTP_MOE_TUNED_CONFIG", False) + + assert _get_sm70_mtp_moe_decode_config(2, 256, 128, 2048, 8) is None + + +def test_mtp_sm70_decode_config_can_be_forced_to_legacy_for_warmup(): + with force_sm70_mtp_moe_legacy_config(): + config = _get_sm70_mtp_moe_decode_config(2, 256, 128, 2048, 8) + + assert config is None + assert _get_sm70_mtp_moe_decode_config(2, 256, 128, 2048, 8) is not None + + +class _FakeSM70Platform: + @staticmethod + def is_cuda(): + return True + + @staticmethod + def is_rocm(): + return False + + @staticmethod + def has_device_capability(capability): + return capability == 70 + + +def test_mtp_sm70_decode_config_is_selected_when_opted_in(monkeypatch): + monkeypatch.setattr(fused_moe_module, "current_platform", _FakeSM70Platform()) + + config = fused_moe_module.get_default_config(2, 256, 128, 2048, 8, None) + + assert config["BLOCK_SIZE_M"] == 8 + assert config["BLOCK_SIZE_N"] == 128 + assert config["BLOCK_SIZE_K"] == 32 + + +def test_mtp_sm70_decode_config_rolls_back_to_0dot3(monkeypatch): + monkeypatch.setattr(fused_moe_module, "current_platform", _FakeSM70Platform()) + monkeypatch.setattr(envs, "VLLM_SM70_MTP_MOE_TUNED_CONFIG", False) + + config = fused_moe_module.get_default_config(1, 256, 128, 2048, 8, None) + + assert config["BLOCK_SIZE_N"] == 32 + assert config["BLOCK_SIZE_K"] == 64 + assert "num_warps" not in config diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index c6daab2d86..dd7e475549 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -12,6 +12,7 @@ """ import math +from types import SimpleNamespace import pytest import torch @@ -24,6 +25,11 @@ from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import ( _fused_kv_compress_norm_rope_insert_indexer_attn, _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn, + compress_norm_rope_store_triton, +) +from vllm.models.deepseek_v4.common.ops.save_partial_states import ( + save_partial_states_to_ring, + stage_partial_states_from_ring, ) from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4 @@ -607,6 +613,12 @@ def test_fused_kv_insert_indexer(num_tokens: int, kv_block_size: int, use_fp4: b kv_cache, slot_mapping, kv_block_size, + state_cache, + state_cache.stride(0), + state_cache, + state_cache.stride(0), + rms_weight, + rms_weight.stride(0), HEAD_SIZE=HEAD_DIM, TRITON_BLOCK_SIZE=HEAD_DIM, STATE_WIDTH=coff * HEAD_DIM, @@ -618,6 +630,10 @@ def test_fused_kv_insert_indexer(num_tokens: int, kv_block_size: int, use_fp4: b TOKEN_STRIDE=TOKEN_STRIDE, SCALE_DIM=SCALE_DIM, KV_BLOCK_STRIDE=kv_cache.stride(0), + USE_PRIVATE_STATE=False, + USE_DENSE_PRIVATE_STATE=False, + RING_SIZE=1, + PRIVATE_HISTORY_SIZE=1, num_warps=1, ) @@ -667,3 +683,212 @@ def test_fused_kv_insert_indexer(num_tokens: int, kv_block_size: int, use_fp4: b assert torch.equal(actual_scale, scale[i : i + 1]), ( f"token {i}: scale {actual_scale.item()} != {scale[i].item()}" ) + + +@pytest.mark.parametrize( + ("head_dim", "compress_ratio", "use_fp4_cache"), + [(512, 4, False), (512, 128, False), (128, 4, False)], +) +def test_private_compressor_ring_matches_paged_state_bytes( + head_dim: int, compress_ratio: int, use_fp4_cache: bool +): + """The SM70 private-state gather must preserve final cache bytes exactly.""" + torch.manual_seed(19) + device = "cuda" + overlap = compress_ratio == 4 + coff = 1 + overlap + state_width = coff * head_dim + state_block_size = 4 if compress_ratio == 4 else 8 + chunk_start = 12 if compress_ratio == 4 else 120 + num_tokens = 12 if compress_ratio == 4 else 136 + total_tokens = chunk_start + num_tokens + ring_size = coff * compress_ratio + 8 + + all_kv = torch.randn(total_tokens, state_width, dtype=torch.float32, device=device) + all_score = torch.randn_like(all_kv) + ape = torch.randn(compress_ratio, state_width, dtype=torch.float32, device=device) + absolute_positions = torch.arange(total_tokens, device=device) + stored_score = all_score + ape[absolute_positions % compress_ratio] + dense_state = torch.cat((all_kv, stored_score), dim=-1) + + num_state_blocks = math.ceil(total_tokens / state_block_size) + state_cache = torch.zeros( + num_state_blocks, + state_block_size, + 2 * state_width, + dtype=torch.float32, + device=device, + ) + state_cache.view(-1, 2 * state_width)[:total_tokens].copy_(dense_state) + block_table = torch.arange( + num_state_blocks, dtype=torch.int32, device=device + ).unsqueeze(0) + + state_ring = torch.empty( + 1, ring_size, 2 * state_width, dtype=torch.float32, device=device + ) + history_start = max(0, chunk_start - ring_size) + history_positions = torch.arange(history_start, chunk_start, device=device) + state_ring[0, history_positions % ring_size] = dense_state[history_positions] + + positions = torch.arange( + chunk_start, total_tokens, dtype=torch.int64, device=device + ) + valid_slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device) + token_to_req = torch.zeros(num_tokens, dtype=torch.int32, device=device) + compressed_slot_mapping = torch.where( + (positions + 1) % compress_ratio == 0, + positions // compress_ratio, + -1, + ) + kv_block_size = 16 + num_compressed_tokens = math.ceil(total_tokens / compress_ratio) + num_kv_blocks = math.ceil(num_compressed_tokens / kv_block_size) + 1 + + if head_dim == 512: + token_stride, scale_dim, quant_block = 576, 8, 64 + elif use_fp4_cache: + token_stride, scale_dim, quant_block = 64, 4, 32 + else: + token_stride, scale_dim, quant_block = 128, 4, 128 + bytes_per_token = token_stride + scale_dim + paged_output = torch.zeros( + num_kv_blocks, + kv_block_size, + bytes_per_token, + dtype=torch.uint8, + device=device, + ) + private_output = torch.zeros_like(paged_output) + k_cache_metadata = SimpleNamespace(slot_mapping=compressed_slot_mapping) + rms_weight = torch.randn(head_dim, dtype=torch.bfloat16, device=device) + cos_sin_cache = torch.randn(total_tokens + 1, 64, device=device) + current_kv = all_kv[chunk_start:] + current_score = all_score[chunk_start:] + + common_kwargs = dict( + current_kv=current_kv, + current_score=current_score, + ape=ape, + num_actual=num_tokens, + token_to_req_indices=token_to_req, + positions=positions, + slot_mapping=valid_slot_mapping, + block_table=block_table, + block_size=state_block_size, + state_width=state_width, + cos_sin_cache=cos_sin_cache, + k_cache_metadata=k_cache_metadata, + pdl_kwargs={}, + head_dim=head_dim, + rope_head_dim=64, + compress_ratio=compress_ratio, + overlap=overlap, + use_fp4_cache=use_fp4_cache, + rms_norm_weight=rms_weight, + rms_norm_eps=1e-6, + quant_block=quant_block, + token_stride=token_stride, + scale_dim=scale_dim, + ) + compress_norm_rope_store_triton( + state_cache=state_cache, + kv_cache=paged_output, + use_private_state=False, + **common_kwargs, + ) + compress_norm_rope_store_triton( + state_cache=state_ring, + kv_cache=private_output, + use_private_state=True, + ring_size=ring_size, + **common_kwargs, + ) + + assert torch.equal(private_output, paged_output) + + if compress_ratio == 128: + dense_state = torch.empty( + compress_ratio + num_tokens, + 2 * state_width, + dtype=torch.float32, + device=device, + ) + stage_partial_states_from_ring( + kv=current_kv, + score=current_score, + ape=ape, + positions=positions, + state_ring=state_ring, + dense_state=dense_state, + slot_mapping=valid_slot_mapping, + state_width=state_width, + history_size=compress_ratio, + compress_ratio=compress_ratio, + ) + dense_output = torch.zeros_like(paged_output) + compress_norm_rope_store_triton( + state_cache=dense_state.unsqueeze(0), + kv_cache=dense_output, + use_dense_private_state=True, + private_history_size=compress_ratio, + **common_kwargs, + ) + assert torch.equal(dense_output, paged_output) + + +@pytest.mark.parametrize( + ("head_dim", "compress_ratio"), [(512, 4), (512, 128), (128, 4)] +) +def test_private_compressor_ring_keeps_only_chunk_tail( + head_dim: int, compress_ratio: int +): + overlap = compress_ratio == 4 + state_width = (1 + overlap) * head_dim + ring_size = (1 + overlap) * compress_ratio + 8 + chunk_start = 4096 + num_tokens = 8192 + positions = torch.arange( + chunk_start, + chunk_start + num_tokens, + dtype=torch.int64, + device="cuda", + ) + kv = torch.randn(num_tokens, state_width, dtype=torch.float32, device="cuda") + score = torch.randn_like(kv) + ape = torch.randn(compress_ratio, state_width, dtype=torch.float32, device="cuda") + state_ring = torch.full( + (1, ring_size, 2 * state_width), + float("nan"), + dtype=torch.float32, + device="cuda", + ) + slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device="cuda") + token_to_req = torch.zeros(num_tokens, dtype=torch.int32, device="cuda") + seq_lens = torch.tensor( + [chunk_start + num_tokens], dtype=torch.int32, device="cuda" + ) + + save_partial_states_to_ring( + kv=kv, + score=score, + ape=ape, + positions=positions, + state_ring=state_ring, + slot_mapping=slot_mapping, + token_to_req_indices=token_to_req, + seq_lens=seq_lens, + state_width=state_width, + compress_ratio=compress_ratio, + ) + + tail_positions = positions[-ring_size:] + actual = state_ring[0, tail_positions % ring_size] + expected = torch.cat( + ( + kv[-ring_size:], + score[-ring_size:] + ape[tail_positions % compress_ratio], + ), + dim=-1, + ) + assert torch.equal(actual, expected) diff --git a/tests/kernels/test_deepseek_v4_sm70_indexer.py b/tests/kernels/test_deepseek_v4_sm70_indexer.py index ab98194cff..2d30d4bf6a 100644 --- a/tests/kernels/test_deepseek_v4_sm70_indexer.py +++ b/tests/kernels/test_deepseek_v4_sm70_indexer.py @@ -4,33 +4,322 @@ import pytest import torch -from vllm.models.deepseek_v4.sm70.indexer import sm70_indexer_prefill_logits - +from vllm.models.deepseek_v4.sm70 import indexer as sm70_indexer +from vllm.models.deepseek_v4.sm70.indexer import ( + sm70_indexer_decode_logits, + sm70_indexer_prefill_logits, +) -@pytest.mark.skipif( +requires_sm70 = pytest.mark.skipif( not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0), reason="requires NVIDIA V100/SM70", ) -def test_sm70_indexer_prefill_dequantizes_fp8_storage_as_uint8() -> None: - torch.manual_seed(20260803) - num_queries = 3 + +_HEAD_DIM = 128 +_BLOCK_SIZE = 64 + +# E4M3FN bit patterns with exact values. Building keys from a palette keeps the +# reference bit-exact and needs no native FP8 conversion, which SM70 lacks. +_FP8_BITS = torch.tensor([0x00, 0x30, 0x38, 0x3C, 0x40, 0xB0, 0xB8, 0xC0]) +_FP8_VALUES = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, -0.5, -1.0, -2.0]) + + +def _random_fp8_keys(num_keys: int, generator: torch.Generator): + """Return (uint8 bit patterns, exact fp32 values) for `num_keys` keys.""" + picks = torch.randint(len(_FP8_BITS), (num_keys, _HEAD_DIM), generator=generator) + return _FP8_BITS[picks].to(torch.uint8).cuda(), _FP8_VALUES[picks].cuda() + + +def _reference_index_logits( + q: torch.Tensor, weights: torch.Tensor, k: torch.Tensor +) -> torch.Tensor: + """I[t, s] = sum_h weights[t, h] * relu(q[t, h] . k[s]). + + The model's actual indexer scoring function (reference `Indexer.forward`, + inference/model.py:427; the in-repo `fp8_mqa_logits` reference at + v1/attention/ops/rocm_aiter_mla_sparse.py:510 applies the same relu). It + does NOT factor to (sum_h weights[t, h] * q[t, h]) . k[s] - the relu sits + between the weighting and the head sum. + """ + per_head = torch.einsum("mhd,nd->mhn", q.float(), k.float()) + return torch.einsum("mhn,mh->mn", torch.relu(per_head), weights.float()) + + +def _reference_factored_logits( + q: torch.Tensor, weights: torch.Tensor, k: torch.Tensor +) -> torch.Tensor: + """The relu-free form, kept reachable via VLLM_SM70_INDEXER_RELU=0.""" + weighted_q = torch.einsum("thd,th->td", q.float(), weights.float()) + return weighted_q @ k.float().T + + +def _build_paged_index_cache( + value_bits: torch.Tensor, scales: torch.Tensor +) -> torch.Tensor: + """Pack keys into the block-major layout the C++ cache kernels write. + + `indexer_k_quant_and_cache` / `cp_gather_indexer_k_quant_cache` + (csrc/libtorch_stable/cache_kernels.cu) store, per block, all `block_size` + rows of 128 FP8 values first and only then that block's FP32 scales: + + value: block_base + pos_in_block * 128 + d + scale: block_base + block_size * 128 + pos_in_block * 4 + + It is NOT a per-token [128 values][4 scale] record, which is why only + token 0 of each block used to read back correctly. + """ + num_blocks, block_size, head_dim = value_bits.shape + assert scales.shape == (num_blocks, block_size) + value_bytes = value_bits.reshape(num_blocks, block_size * head_dim) + scale_bytes = scales.contiguous().view(torch.uint8) + return ( + torch.cat([value_bytes, scale_bytes], dim=1) + .reshape(num_blocks, block_size, head_dim + 4) + .contiguous() + ) + + +def _make_queries(rows: int, num_heads: int): + q = ( + torch.randn((rows, num_heads, _HEAD_DIM), device="cuda", dtype=torch.float16) + * 0.25 + ) + weights = torch.randn((rows, num_heads), device="cuda", dtype=torch.float32) + return q, weights + + +@requires_sm70 +@pytest.mark.parametrize("use_cublas", [True, False]) +def test_prefill_matches_the_reference_scoring_function(monkeypatch, use_cublas): + torch.manual_seed(20260819) + generator = torch.Generator().manual_seed(20260819) + num_queries, num_heads, num_keys = 5, 8, 37 + monkeypatch.setattr(sm70_indexer, "_PREFILL_CUBLAS", use_cublas) + + q, weights = _make_queries(num_queries, num_heads) + bits, values = _random_fp8_keys(num_keys, generator) + scales = torch.linspace(0.5, 2.0, num_keys, device="cuda", dtype=torch.float32) + k = values * scales[:, None] + + actual = sm70_indexer_prefill_logits( + q, bits.view(torch.float8_e4m3fn), scales, weights + ) + expected = _reference_index_logits(q, weights, k) + torch.testing.assert_close(actual, expected, rtol=1e-2, atol=1e-2) + + # The relu is not a rounding detail: without it the scoring is a different + # function, which is what broke long-context top-k selection. + factored = _reference_factored_logits(q, weights, k) + assert not torch.allclose(expected, factored, rtol=1e-2, atol=1e-2) + + +@requires_sm70 +def test_prefill_relu_off_restores_the_factored_form(monkeypatch): + """VLLM_SM70_INDEXER_RELU=0 keeps the old scoring reachable for A/B.""" + torch.manual_seed(20260819) + generator = torch.Generator().manual_seed(20260819) + num_queries, num_heads, num_keys = 4, 8, 23 + monkeypatch.setattr(sm70_indexer, "_RELU_LOGITS", False) + + q, weights = _make_queries(num_queries, num_heads) + bits, values = _random_fp8_keys(num_keys, generator) + scales = torch.linspace(0.5, 2.0, num_keys, device="cuda", dtype=torch.float32) + + actual = sm70_indexer_prefill_logits( + q, bits.view(torch.float8_e4m3fn), scales, weights + ) + expected = _reference_factored_logits(q, weights, values * scales[:, None]) + torch.testing.assert_close(actual, expected, rtol=1e-2, atol=1e-2) + + +@requires_sm70 +@pytest.mark.parametrize( + "relu,fused", + [ + pytest.param(True, True, id="relu"), + pytest.param(False, True, id="factored-fused"), + pytest.param(False, False, id="factored-gather"), + ], +) +def test_decode_reads_the_block_major_paged_cache(monkeypatch, relu, fused): + """Every decode path must address values and scales the way the cache is + written. Reading a token as a [128 values][4 scale] record lands on the + wrong bytes for every token but the first of each block, and reinterprets + four FP8 value bytes as the FP32 scale. + """ + torch.manual_seed(20260819) + generator = torch.Generator().manual_seed(20260819) + monkeypatch.setattr(sm70_indexer, "_RELU_LOGITS", relu) + monkeypatch.setattr(sm70_indexer, "_FUSED_DECODE_LOGITS", fused) + + num_heads = 8 + seq_lens_list = [1, _BLOCK_SIZE + 1, 3 * _BLOCK_SIZE - 5] + rows = len(seq_lens_list) + max_seq_len = max(seq_lens_list) + blocks_per_row = (max_seq_len + _BLOCK_SIZE - 1) // _BLOCK_SIZE + num_blocks = rows * blocks_per_row + 2 # +2 so no row can start at block 0 + + bits, values = _random_fp8_keys(num_blocks * _BLOCK_SIZE, generator) + value_bits = bits.reshape(num_blocks, _BLOCK_SIZE, _HEAD_DIM) + block_values = values.reshape(num_blocks, _BLOCK_SIZE, _HEAD_DIM) + # Distinct per-token scales inside every block: with one scale per block + # the layout bug is invisible. + scales = torch.rand((num_blocks, _BLOCK_SIZE), generator=generator).cuda() + 0.5 + cache = _build_paged_index_cache(value_bits, scales) + + # Shuffled, non-identity mapping so a wrong block stride cannot pass. + perm = torch.randperm(num_blocks, generator=generator)[: rows * blocks_per_row] + block_table = perm.to(torch.int32).cuda().reshape(rows, blocks_per_row) + seq_lens = torch.tensor(seq_lens_list, device="cuda", dtype=torch.int32) + q, weights = _make_queries(rows, num_heads) + + actual = sm70_indexer_decode_logits( + q, cache, weights, seq_lens, block_table, max_seq_len + ) + + reference = _reference_index_logits if relu else _reference_factored_logits + for row, seq_len in enumerate(seq_lens_list): + positions = torch.arange(seq_len, device="cuda") + physical = block_table[row][positions // _BLOCK_SIZE].long() + pos_in_block = positions % _BLOCK_SIZE + k = ( + block_values[physical, pos_in_block] + * scales[physical, pos_in_block][:, None] + ) + expected = reference(q[row : row + 1], weights[row : row + 1], k) + torch.testing.assert_close( + actual[row, :seq_len].unsqueeze(0), expected, rtol=1e-2, atol=1e-2 + ) + + +@requires_sm70 +def test_decode_cublas_keeps_the_fp32_topk_set_and_masks_graph_tail(monkeypatch): + torch.manual_seed(20260824) + generator = torch.Generator().manual_seed(20260824) + num_rows, num_heads = 8, 64 + live_seq_len = 1017 + graph_width = 2048 + blocks = (live_seq_len + _BLOCK_SIZE - 1) // _BLOCK_SIZE + + bits, _ = _random_fp8_keys(blocks * _BLOCK_SIZE, generator) + value_bits = bits.reshape(blocks, _BLOCK_SIZE, _HEAD_DIM) + scales = torch.rand((blocks, _BLOCK_SIZE), generator=generator).cuda() + 0.5 + cache = _build_paged_index_cache(value_bits, scales) + block_table = torch.full( + (1, graph_width // _BLOCK_SIZE), + -1, + dtype=torch.int32, + device="cuda", + ) + block_table[0, :blocks] = torch.randperm(blocks, generator=generator).to( + device="cuda", dtype=torch.int32 + ) + seq_lens = torch.arange( + live_seq_len - num_rows + 1, + live_seq_len + 1, + dtype=torch.int32, + device="cuda", + ).view(1, num_rows) + q, weights = _make_queries(num_rows, num_heads) + + monkeypatch.setattr(sm70_indexer, "_DECODE_CUBLAS", False) + baseline = sm70_indexer_decode_logits( + q, cache, weights, seq_lens, block_table, graph_width + ) + + static_workspace = ( + torch.empty((graph_width, _HEAD_DIM), dtype=torch.float16, device="cuda"), + torch.empty((graph_width,), dtype=torch.float32, device="cuda"), + torch.empty( + (num_rows * num_heads, graph_width), + dtype=torch.float32, + device="cuda", + ), + ) + + class StaticWorkspace: + @staticmethod + def get_simultaneous(*specs): + assert len(specs) == len(static_workspace) + return static_workspace + + monkeypatch.setattr(sm70_indexer, "_DECODE_CUBLAS", True) + monkeypatch.setattr(sm70_indexer, "_DECODE_CUBLAS_MIN_KEYS", 1) + monkeypatch.setattr( + sm70_indexer, "current_workspace_manager", lambda: StaticWorkspace() + ) + + def candidate_call(): + return sm70_indexer_decode_logits( + q, cache, weights, seq_lens, block_table, graph_width + ) + + capture_stream = torch.cuda.Stream() + capture_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(capture_stream): + for _ in range(3): + candidate_call() + capture_stream.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=capture_stream): + candidate = candidate_call() + graph.replay() + torch.accelerator.synchronize() + + for row, row_len in enumerate(seq_lens.reshape(-1).tolist()): + expected_topk = torch.topk(baseline[row, :row_len], 512).indices.sort().values + actual_topk = torch.topk(candidate[row, :row_len], 512).indices.sort().values + assert torch.equal(actual_topk, expected_topk) + + # The graph bucket has no allocated blocks after the live request tail. + # Reaching this assertion proves the gather masked those -1 entries. + assert torch.isfinite(candidate[:, :live_seq_len]).all() + + # Replay the same graph with a shorter live request. This exercises the + # device-derived bound instead of merely validating the capture-time row. + seq_lens.sub_(127) + monkeypatch.setattr(sm70_indexer, "_DECODE_CUBLAS", False) + replay_baseline = sm70_indexer_decode_logits( + q, cache, weights, seq_lens, block_table, graph_width + ) + graph.replay() + torch.accelerator.synchronize() + for row, row_len in enumerate(seq_lens.reshape(-1).tolist()): + expected_topk = ( + torch.topk(replay_baseline[row, :row_len], 512).indices.sort().values + ) + actual_topk = torch.topk(candidate[row, :row_len], 512).indices.sort().values + assert torch.equal(actual_topk, expected_topk) + + +@requires_sm70 +def test_sm70_indexer_decode_valid_logits_do_not_depend_on_workspace_width() -> None: + torch.manual_seed(20260804) + num_queries = 8 num_heads = 4 num_keys = 7 head_dim = 128 + cache_block_size = 64 q = torch.randn( (num_queries, num_heads, head_dim), device="cuda", dtype=torch.float16 ) - weights = torch.randn((num_queries, num_heads), device="cuda", dtype=torch.float32) - # E4M3FN bit pattern 0x38 is exactly 1.0. Constructing the tensor through - # uint8 also avoids requiring native FP8 conversion support on SM70. - k_bits = torch.full((num_keys, head_dim), 0x38, device="cuda", dtype=torch.uint8) - k_quant = k_bits.view(torch.float8_e4m3fn) - scales = torch.linspace(0.5, 2.0, num_keys, device="cuda", dtype=torch.float32) + weights = torch.randn((num_queries, num_heads), device="cuda", dtype=torch.float16) + cache = torch.zeros( + (1, cache_block_size, head_dim + 4), device="cuda", dtype=torch.uint8 + ) + cache[..., :head_dim].fill_(0x38) + cache[..., head_dim:].view(torch.float32).fill_(1.0) + seq_lens = torch.full((1, num_queries), num_keys, device="cuda", dtype=torch.int32) + block_table = torch.zeros((1, 1), device="cuda", dtype=torch.int32) - actual = sm70_indexer_prefill_logits(q, k_quant, scales, weights) + narrow = sm70_indexer_decode_logits( + q, cache, weights, seq_lens, block_table, max_seq_len=16 + ) + wide = sm70_indexer_decode_logits( + q, cache, weights, seq_lens, block_table, max_seq_len=64 + ) - weighted_q = torch.sum(q.float() * weights[:, :, None], dim=1).half() - k_reference = scales[:, None].expand(num_keys, head_dim).half() - expected = torch.mm(weighted_q, k_reference.T, out_dtype=torch.float32) - torch.testing.assert_close(actual, expected, rtol=2e-3, atol=2e-3) + assert narrow.shape == (num_queries, 16) + assert wide.shape == (num_queries, 64) + torch.testing.assert_close(narrow[:, :num_keys], wide[:, :num_keys], rtol=0, atol=0) 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 8dff20c955..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 @@ -25,10 +25,7 @@ def _make_cos_sin_cache(max_position: int, device: str) -> torch.Tensor: inv_freq = 1.0 / ( 10000.0 - ** ( - torch.arange(0, ROPE_DIM, 2, dtype=torch.float32, device=device) - / ROPE_DIM - ) + ** (torch.arange(0, ROPE_DIM, 2, dtype=torch.float32, device=device) / ROPE_DIM) ) positions = torch.arange(max_position, dtype=torch.float32, device=device) frequencies = torch.outer(positions, inv_freq) @@ -120,3 +117,45 @@ def test_sm70_kv_rope_insert_is_repeatable_across_physical_blocks(): torch.testing.assert_close(gathered, gathered_outputs[0], rtol=0, atol=0) for q_out in q_outputs[1:]: torch.testing.assert_close(q_out, q_outputs[0], rtol=0, atol=0) + + +def test_sm70_tp4_fused_qnorm_kv_insert_matches_legacy( + monkeypatch: pytest.MonkeyPatch, +): + torch.manual_seed(20260826) + device = "cuda" + num_heads = 16 + block_size = 256 + block_stride = round_up(block_size * HEAD_BYTES, TOKEN_DATA_BYTES) + storage_shape = (3 * block_stride,) + q_input = torch.randn(1, num_heads, HEAD_DIM, dtype=torch.float16, device=device) + kv = torch.randn(1, HEAD_DIM, dtype=torch.float16, device=device) + positions = torch.tensor([1024], dtype=torch.int64, device=device) + slot_mapping = torch.tensor([block_size + 3], dtype=torch.int64, device=device) + cos_sin_cache = _make_cos_sin_cache(2048, device) + + def run(enabled: bool) -> tuple[torch.Tensor, torch.Tensor]: + monkeypatch.setenv("VLLM_SM70_DSV4_QNORM_KV_FUSED_TP4", "1" if enabled else "0") + q = q_input.clone() + storage = torch.full(storage_shape, 0xA5, dtype=torch.uint8, device=device) + cache = storage.as_strided( + (3, block_size, HEAD_BYTES), + (block_stride, HEAD_BYTES, 1), + ) + sm70_qnorm_rope_kv_fp8_insert( + q, + kv, + cache, + slot_mapping, + positions, + cos_sin_cache, + eps=1e-6, + block_size=block_size, + ) + torch.accelerator.synchronize() + return q, storage + + legacy_q, legacy_storage = run(False) + fused_q, fused_storage = run(True) + torch.testing.assert_close(fused_q, legacy_q, rtol=0, atol=0) + torch.testing.assert_close(fused_storage, legacy_storage, rtol=0, atol=0) 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_fused_sigmoid_gating_delta_rule.py b/tests/kernels/test_fused_sigmoid_gating_delta_rule.py index 2b03e83c30..2ac073b410 100644 --- a/tests/kernels/test_fused_sigmoid_gating_delta_rule.py +++ b/tests/kernels/test_fused_sigmoid_gating_delta_rule.py @@ -7,7 +7,12 @@ from vllm.model_executor.layers.fla.ops import ( fused_recurrent_gated_delta_rule, + fused_recurrent_gated_delta_rule_packed_decode, fused_sigmoid_gating_delta_rule_update, + fused_sigmoid_gating_delta_rule_update_mixed_qkv_out, +) +from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( + fused_gdn_gating, ) from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed @@ -15,6 +20,149 @@ DEVICE = current_platform.device_type +def _bitwise_diff_stats( + actual: torch.Tensor, + expected: torch.Tensor, +) -> tuple[int, float, float]: + diff = actual.float() - expected.float() + return ( + torch.count_nonzero(actual != expected).item(), + diff.abs().max().item(), + diff.abs().mean().item(), + ) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_split_verify_matches_fused_decode_beta_semantics( + dtype: torch.dtype, +) -> None: + torch.set_default_device(DEVICE) + if dtype is torch.bfloat16 and not torch.cuda.is_bf16_supported(): + pytest.skip("BF16 requires a supported accelerator") + set_random_seed(0) + + # Qwen3.8-27B TP4 per-rank head geometry (16 QK heads, 48 V heads). + num_qk_heads = 4 + num_v_heads = 12 + head_dim = 128 + scale = head_dim**-0.5 + # Make the decay exactly zero in both kernels so this regression isolates + # beta materialization instead of also measuring unrelated fused/split + # decay rounding. softplus(-100) rounds to zero in the kernel. + A_log = torch.zeros(num_v_heads, dtype=torch.float32) + dt_bias = torch.zeros(num_v_heads, dtype=dtype) + a = torch.full((1, num_v_heads), -100, dtype=dtype) + # sigmoid(-0.5) is not exactly representable in FP16 or BF16, exposing + # whether the split verifier keeps the fused decode FP32 beta semantics. + b = torch.full((1, num_v_heads), -0.5, dtype=dtype) + query = torch.rand(1, 1, num_qk_heads, head_dim, dtype=dtype) + key = torch.rand_like(query) + value = torch.rand(1, 1, num_v_heads, head_dim, dtype=dtype) + # A zero state also removes decay and dot-product reduction-order effects + # from this one-step comparison. The state transition then differs only + # through the materialized beta value. + initial_state = torch.zeros( + 1, + num_v_heads, + head_dim, + head_dim, + dtype=torch.float32, + ) + state_indices = torch.zeros(1, dtype=torch.int32) + cu_seqlens = torch.tensor([0, 1], dtype=torch.int32) + + fused_state = initial_state.clone() + fused_output, _ = fused_sigmoid_gating_delta_rule_update( + A_log=A_log, + a=a, + b=b, + dt_bias=dt_bias, + q=query, + k=key, + v=value, + scale=scale, + initial_state=fused_state, + inplace_final_state=True, + cu_seqlens=cu_seqlens, + ssm_state_indices=state_indices, + use_qk_l2norm_in_kernel=False, + ) + + packed_state = initial_state.clone() + packed_output = torch.empty_like(fused_output) + mixed_qkv = torch.cat( + [ + query.squeeze(0).reshape(1, -1), + key.squeeze(0).reshape(1, -1), + value.squeeze(0).reshape(1, -1), + ], + dim=-1, + ).contiguous() + fused_recurrent_gated_delta_rule_packed_decode( + mixed_qkv=mixed_qkv, + a=a, + b=b, + A_log=A_log, + dt_bias=dt_bias, + scale=scale, + initial_state=packed_state, + out=packed_output, + ssm_state_indices=state_indices, + use_qk_l2norm_in_kernel=False, + ) + + def run_split( + beta_dtype: torch.dtype | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + g, beta = fused_gdn_gating( + A_log, + a, + b, + dt_bias, + beta_dtype=beta_dtype, + ) + state = initial_state.clone() + output, _ = fused_recurrent_gated_delta_rule( + q=query, + k=key, + v=value, + g=g, + beta=beta, + scale=scale, + initial_state=state, + inplace_final_state=True, + ssm_state_indices=state_indices, + cu_seqlens=cu_seqlens, + use_qk_l2norm_in_kernel=False, + ) + return output, state, beta + + fixed_output, fixed_state, fixed_beta = run_split(torch.float32) + legacy_output, legacy_state, legacy_beta = run_split(None) + + assert fixed_beta.dtype is torch.float32 + assert legacy_beta.dtype is dtype + fixed_output_diff = _bitwise_diff_stats(fixed_output, fused_output) + fixed_state_diff = _bitwise_diff_stats(fixed_state, fused_state) + legacy_output_diff = _bitwise_diff_stats(legacy_output, fused_output) + legacy_state_diff = _bitwise_diff_stats(legacy_state, fused_state) + assert torch.equal(packed_output, fused_output), _bitwise_diff_stats( + packed_output, fused_output + ) + assert torch.equal(packed_state, fused_state), _bitwise_diff_stats( + packed_state, fused_state + ) + assert torch.equal(fixed_output, fused_output), fixed_output_diff + assert torch.equal(fixed_state, fused_state), ( + "fixed", + fixed_state_diff, + "legacy", + legacy_state_diff, + ) + assert legacy_state_diff[0] > 0, legacy_state_diff + assert legacy_output_diff[0] > 0, legacy_output_diff + + @pytest.mark.parametrize("tp_size", [1]) @pytest.mark.parametrize("num_reqs", [1, 2, 4]) @pytest.mark.parametrize("num_k_heads", [16]) @@ -194,3 +342,176 @@ def test_fused_sigmoid_gating_delta_rule_update_spec( torch.testing.assert_close( last_recurrent_state, last_recurrent_state_ref, atol=1e-2, rtol=1e-2 ) + + +@pytest.mark.parametrize("num_reqs", [1, 2]) +@pytest.mark.parametrize("num_speculative_tokens", [3, 7]) +@pytest.mark.parametrize("state_dtype", [torch.float16, torch.float32]) +def test_dflash2_packed_verify_matches_split_fp16_contract( + num_reqs: int, + num_speculative_tokens: int, + state_dtype: torch.dtype, +) -> None: + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0): + pytest.skip("the DFlash2 verifier fast path is SM70-only") + + torch.set_default_device(DEVICE) + set_random_seed(11) + dtype = torch.float16 + # Production GDN verifier shape used by the recorded TP4 parity evidence. + num_k_heads, num_v_heads = 4, 12 + head_k_dim = head_v_dim = 128 + tokens_per_req = num_speculative_tokens + 1 + num_tokens = num_reqs * tokens_per_req + qkv_width = 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim + + mixed_qkv = torch.randn(num_tokens, qkv_width, dtype=dtype) + query, key, value = torch.split( + mixed_qkv, + [ + num_k_heads * head_k_dim, + num_k_heads * head_k_dim, + num_v_heads * head_v_dim, + ], + dim=-1, + ) + query = query.contiguous().view(1, num_tokens, num_k_heads, head_k_dim) + key = key.contiguous().view(1, num_tokens, num_k_heads, head_k_dim) + value = value.contiguous().view(1, num_tokens, num_v_heads, head_v_dim) + a = torch.randn(num_tokens, num_v_heads, dtype=dtype) + b = torch.randn(num_tokens, num_v_heads, dtype=dtype) + A_log = torch.randn(num_v_heads, dtype=torch.float32) + dt_bias = torch.randn(num_v_heads, dtype=dtype) + state_indices = torch.arange(num_tokens, dtype=torch.int32).view( + num_reqs, tokens_per_req + ) + num_accepted_tokens = torch.randint( + 1, tokens_per_req + 1, (num_reqs,), dtype=torch.int32 + ) + cu_seqlens = torch.arange(0, num_tokens + 1, tokens_per_req, dtype=torch.int32) + initial_state = torch.randn( + num_tokens + 3, + num_v_heads, + head_v_dim, + head_k_dim, + dtype=state_dtype, + ) + + g, beta = fused_gdn_gating(A_log, a, b, dt_bias) + reference_state = initial_state.clone() + reference_out, _ = fused_recurrent_gated_delta_rule( + q=query, + k=key, + v=value, + g=g, + beta=beta, + initial_state=reference_state, + inplace_final_state=True, + cu_seqlens=cu_seqlens, + ssm_state_indices=state_indices, + num_accepted_tokens=num_accepted_tokens, + use_qk_l2norm_in_kernel=True, + ) + + fused_state = initial_state.clone() + fused_out = torch.empty(num_tokens, 1, num_v_heads, head_v_dim, dtype=dtype) + fused_sigmoid_gating_delta_rule_update_mixed_qkv_out( + A_log=A_log, + a=a, + b=b, + dt_bias=dt_bias, + mixed_qkv=mixed_qkv, + num_q_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + out=fused_out, + initial_state=fused_state, + cu_seqlens=cu_seqlens, + ssm_state_indices=state_indices, + num_accepted_tokens=num_accepted_tokens, + use_qk_l2norm_in_kernel=True, + precomputed_g=g, + precomputed_beta=beta, + match_recurrent_schedule=True, + match_recurrent_numerics=True, + ) + torch.accelerator.synchronize() + + torch.testing.assert_close(fused_out.transpose(0, 1), reference_out, rtol=0, atol=0) + torch.testing.assert_close(fused_state, reference_state, rtol=0, atol=0) + + +def test_dflash2_quantized_state_reload_matches_repeated_decode() -> None: + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0): + pytest.skip("the DFlash2 verifier fast path is SM70-only") + + torch.set_default_device(DEVICE) + set_random_seed(12) + dtype = torch.float16 + num_tokens = 4 + num_k_heads, num_v_heads = 2, 4 + head_k_dim = head_v_dim = 16 + qkv_width = 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim + mixed_qkv = torch.randn(num_tokens, qkv_width, dtype=dtype) + a = torch.randn(num_tokens, num_v_heads, dtype=dtype) + b = torch.randn(num_tokens, num_v_heads, dtype=dtype) + A_log = torch.randn(num_v_heads, dtype=torch.float32) + dt_bias = torch.randn(num_v_heads, dtype=dtype) + initial_state = torch.randn(2, num_v_heads, head_v_dim, head_k_dim, dtype=dtype) + shared_state_indices = torch.zeros((1, num_tokens), dtype=torch.int32) + accepted = torch.ones(1, dtype=torch.int32) + + fused_state = initial_state.clone() + fused_out = torch.empty(num_tokens, 1, num_v_heads, head_v_dim, dtype=dtype) + fused_sigmoid_gating_delta_rule_update_mixed_qkv_out( + A_log=A_log, + a=a, + b=b, + dt_bias=dt_bias, + mixed_qkv=mixed_qkv, + num_q_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + out=fused_out, + initial_state=fused_state, + cu_seqlens=torch.tensor([0, num_tokens], dtype=torch.int32), + ssm_state_indices=shared_state_indices, + num_accepted_tokens=accepted, + quantize_beta=True, + quantize_state_each_step=True, + match_recurrent_schedule=True, + use_qk_l2norm_in_kernel=True, + ) + + sequential_state = initial_state.clone() + sequential_outputs = [] + one_token_cu_seqlens = torch.tensor([0, 1], dtype=torch.int32) + one_state_index = torch.zeros(1, dtype=torch.int32) + for token_idx in range(num_tokens): + token_out = torch.empty(1, 1, num_v_heads, head_v_dim, dtype=dtype) + fused_sigmoid_gating_delta_rule_update_mixed_qkv_out( + A_log=A_log, + a=a[token_idx : token_idx + 1], + b=b[token_idx : token_idx + 1], + dt_bias=dt_bias, + mixed_qkv=mixed_qkv[token_idx : token_idx + 1], + num_q_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + out=token_out, + initial_state=sequential_state, + cu_seqlens=one_token_cu_seqlens, + ssm_state_indices=one_state_index, + num_accepted_tokens=accepted, + quantize_beta=True, + match_recurrent_schedule=True, + use_qk_l2norm_in_kernel=True, + ) + sequential_outputs.append(token_out) + torch.accelerator.synchronize() + + assert torch.equal(fused_out, torch.cat(sequential_outputs, dim=0)) + assert torch.equal(fused_state, sequential_state) diff --git a/tests/kernels/test_mhc_sm70_fp16.py b/tests/kernels/test_mhc_sm70_fp16.py index 75b5ae7242..a0086f77e7 100644 --- a/tests/kernels/test_mhc_sm70_fp16.py +++ b/tests/kernels/test_mhc_sm70_fp16.py @@ -154,7 +154,10 @@ def test_mhc_sm70_fp16_block_m_prenorm_keeps_rows_independent() -> None: ), reason="NVIDIA V100/SM70 and TileLang required", ) -def test_mhc_sm70_fp32_stage_matches_fused_decode_bitwise() -> None: +@pytest.mark.parametrize("num_tokens", [1, 4, 7, 8]) +def test_mhc_sm70_fp32_stage_matches_fused_decode_bitwise( + num_tokens: int, +) -> None: from vllm.model_executor.kernels.mhc.tilelang_kernels import ( mhc_fused_tilelang, sm70_mhc_dot_from_fp32_stage_tilelang, @@ -165,21 +168,26 @@ def test_mhc_sm70_fp32_stage_matches_fused_decode_bitwise() -> None: hidden_size = 4096 hc_mult = 4 hc_out = 24 - n_splits = 8 + tile_n = 2 if num_tokens < 8 else 3 + n_splits = 8 if num_tokens < 8 else 4 device = mhc_tilelang.current_platform.device_type - x = torch.randn((1, hidden_size), device=device, dtype=torch.float16) + x = torch.randn((num_tokens, hidden_size), device=device, dtype=torch.float16) residual = torch.randn( - (1, hc_mult, hidden_size), device=device, dtype=torch.float16 + (num_tokens, hc_mult, hidden_size), device=device, dtype=torch.float16 + ) + post_mix = torch.randn((num_tokens, hc_mult), device=device, dtype=torch.float32) + comb_mix = torch.randn( + (num_tokens, hc_mult, hc_mult), device=device, dtype=torch.float32 ) - post_mix = torch.randn((1, hc_mult), device=device, dtype=torch.float32) - comb_mix = torch.randn((1, hc_mult, hc_mult), device=device, dtype=torch.float32) fn = torch.randn((hc_out, hc_mult, hidden_size), device=device, dtype=torch.float32) baseline_gemm = torch.empty( - (n_splits, 1, hc_out), device=device, dtype=torch.float32 + (n_splits, num_tokens, hc_out), device=device, dtype=torch.float32 + ) + baseline_sqrsum = torch.empty( + (n_splits, num_tokens), device=device, dtype=torch.float32 ) - baseline_sqrsum = torch.empty((n_splits, 1), device=device, dtype=torch.float32) baseline_residual = torch.empty_like(residual) candidate_gemm = torch.empty_like(baseline_gemm) candidate_sqrsum = torch.empty_like(baseline_sqrsum) @@ -198,7 +206,7 @@ def test_mhc_sm70_fp32_stage_matches_fused_decode_bitwise() -> None: hc_mult, hidden_size, hc_out, - tile_n=2, + tile_n=tile_n, n_splits=n_splits, use_fp16=True, ) @@ -221,10 +229,10 @@ def test_mhc_sm70_fp32_stage_matches_fused_decode_bitwise() -> None: hidden_size, hc_mult, hc_out, - tile_n=2, + tile_n=tile_n, n_splits=n_splits, ) - torch.cuda.synchronize() + torch.accelerator.synchronize() torch.testing.assert_close(candidate_residual, baseline_residual, rtol=0, atol=0) torch.testing.assert_close(candidate_gemm, baseline_gemm, rtol=0, atol=0) 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/kernels/test_sm70_qwen35_gdn_split.py b/tests/kernels/test_sm70_qwen35_gdn_split.py new file mode 100644 index 0000000000..e8a65c0002 --- /dev/null +++ b/tests/kernels/test_sm70_qwen35_gdn_split.py @@ -0,0 +1,127 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( + _sm70_pack_qwen_gdn_qkv, +) +from vllm.model_executor.models.qwen3_5 import ( + _sm70_materialize_qwen35_gdn_splits, +) + + +@pytest.mark.parametrize("num_rows", [1, 8, 32]) +def test_qwen35_gdn_split_materialization_is_bitwise_exact(num_rows: int): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for the Qwen3.5 GDN split kernel") + + torch.manual_seed(7) + qkv_size = 2560 + z_size = 1536 + ba_size = 12 + + # Slice wider allocations so the source row strides are deliberately + # larger than their logical widths, matching the projection-view contract. + qkvz_storage = torch.randn( + (num_rows, qkv_size + z_size + 7), + dtype=torch.float16, + device="cuda", + ) + ba_storage = torch.randn( + (num_rows, 2 * ba_size + 5), + dtype=torch.float16, + device="cuda", + ) + mixed_qkvz = qkvz_storage[:, : qkv_size + z_size] + mixed_ba = ba_storage[:, : 2 * ba_size] + + actual_z, actual_b, actual_a = _sm70_materialize_qwen35_gdn_splits( + mixed_qkvz, + mixed_ba, + qkv_size, + z_size, + ba_size, + ) + + expected_z = mixed_qkvz[:, qkv_size : qkv_size + z_size].contiguous() + expected_b = mixed_ba[:, :ba_size].contiguous() + expected_a = mixed_ba[:, ba_size : 2 * ba_size].contiguous() + assert torch.equal(actual_z, expected_z) + assert torch.equal(actual_b, expected_b) + assert torch.equal(actual_a, expected_a) + + +def test_qwen35_gdn_split_graph_replay_reads_current_projection_values(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for the Qwen3.5 GDN split kernel") + + num_rows, qkv_size, z_size, ba_size = 8, 2560, 1536, 12 + mixed_qkvz = torch.randn( + (num_rows, qkv_size + z_size), dtype=torch.float16, device="cuda" + ) + mixed_ba = torch.randn((num_rows, 2 * ba_size), dtype=torch.float16, device="cuda") + _sm70_materialize_qwen35_gdn_splits(mixed_qkvz, mixed_ba, qkv_size, z_size, ba_size) + torch.accelerator.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual_z, actual_b, actual_a = _sm70_materialize_qwen35_gdn_splits( + mixed_qkvz, mixed_ba, qkv_size, z_size, ba_size + ) + + mixed_qkvz.copy_(torch.randn_like(mixed_qkvz)) + mixed_ba.copy_(torch.randn_like(mixed_ba)) + graph.replay() + torch.accelerator.synchronize() + + assert torch.equal(actual_z, mixed_qkvz[:, qkv_size:].contiguous()) + assert torch.equal(actual_b, mixed_ba[:, :ba_size].contiguous()) + assert torch.equal(actual_a, mixed_ba[:, ba_size:].contiguous()) + + +@pytest.mark.parametrize("num_rows", [1, 8]) +def test_qwen35_gdn_qkv_pack_is_bitwise_exact(num_rows: int): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for the Qwen3.5 GDN Q/K/V pack kernel") + + torch.manual_seed(11) + q_dim, k_dim, v_dim = 512, 512, 1536 + storage = torch.randn( + (num_rows, q_dim + k_dim + v_dim + 13), + dtype=torch.float16, + device="cuda", + ) + mixed_qkv = storage[:, : q_dim + k_dim + v_dim] + actual = _sm70_pack_qwen_gdn_qkv(mixed_qkv, q_dim, k_dim, v_dim) + + q, k, v = torch.split(mixed_qkv, [q_dim, k_dim, v_dim], dim=-1) + expected = torch.cat([q.reshape(-1), k.reshape(-1), v.reshape(-1)]) + assert torch.equal(actual, expected) + + +def test_qwen35_gdn_qkv_pack_graph_replay_reads_current_values(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for the Qwen3.5 GDN Q/K/V pack kernel") + + num_rows, q_dim, k_dim, v_dim = 8, 512, 512, 1536 + mixed_qkv = torch.randn( + (num_rows, q_dim + k_dim + v_dim), + dtype=torch.float16, + device="cuda", + ) + _sm70_pack_qwen_gdn_qkv(mixed_qkv, q_dim, k_dim, v_dim) + torch.accelerator.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = _sm70_pack_qwen_gdn_qkv(mixed_qkv, q_dim, k_dim, v_dim) + + mixed_qkv.copy_(torch.randn_like(mixed_qkv)) + graph.replay() + torch.accelerator.synchronize() + + q, k, v = torch.split(mixed_qkv, [q_dim, k_dim, v_dim], dim=-1) + expected = torch.cat([q.reshape(-1), k.reshape(-1), v.reshape(-1)]) + assert torch.equal(actual, expected) diff --git a/tests/model_executor/test_deepseek_v4_swa_rope.py b/tests/model_executor/test_deepseek_v4_swa_rope.py new file mode 100644 index 0000000000..e494d16cfc --- /dev/null +++ b/tests/model_executor/test_deepseek_v4_swa_rope.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""YaRN must be off on DeepSeek-V4's sliding-window-only layers. + +For compress_ratio == 0 layers the reference passes original_seq_len=0 to +precompute_freqs_cis (inference/model.py:481-486), which skips the correction +range entirely and leaves plain rope at the base rope_theta. Swapping only +rope_theta and leaving rope_type="deepseek_yarn" with factor=16 keeps the +interpolation on for those layers. +""" + +import math + +import pytest +import torch + +from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.platforms import current_platform + +# DeepSeek-V4-Flash config.json. +_HEAD_SIZE = 512 +_ROPE_DIM = 64 +_MAX_POSITION = 262144 # served context; the model passes max_position_embeddings +_ORIGINAL_MAX_POSITION = 65536 +_YARN_FACTOR = 16 +_ROPE_THETA = 10000 # SWA-only layers; compressing layers use 160000 +_SWA_WINDOW = 128 + + +@pytest.fixture(autouse=True) +def _default_device(): + """DeepseekV4ScalingRotaryEmbedding builds its position grid on + current_platform.device_type but its inv_freq on the default device, so it + only assembles under the device context the model loader already sets.""" + with torch.device(current_platform.device_type): + yield + + +def _rope_parameters(factor: int, original_max_position: int) -> dict: + """What DeepseekV4Attention hands get_rope for a compress_ratio<=1 layer.""" + return { + "rope_type": "deepseek_yarn", + "beta_fast": 32, + "beta_slow": 1, + "factor": factor, + "original_max_position_embeddings": original_max_position, + "rope_theta": _ROPE_THETA, + "mscale": 0, + "mscale_all_dim": 0, + "is_deepseek_v4": True, + "rope_dim": _ROPE_DIM, + } + + +def _get_rope(factor: int, original_max_position: int): + return get_rope( + _HEAD_SIZE, + max_position=_MAX_POSITION, + rope_parameters=_rope_parameters(factor, original_max_position), + is_neox_style=False, + ) + + +def _plain_inv_freq() -> torch.Tensor: + exponents = torch.arange(0, _ROPE_DIM, 2, dtype=torch.float32) / _ROPE_DIM + return 1.0 / (_ROPE_THETA**exponents) + + +def test_factor_one_is_exactly_plain_rope(default_vllm_config) -> None: + """factor=1 collapses YaRN to an identity, so no class swap is needed. + + DeepseekV4ScalingRotaryEmbedding is what the fused kernels expect: it + rotates the *last* rotary_dim and keeps cos/sin in fp32. Neither is true of + the base RotaryEmbedding, so "turn YaRN off" has to mean factor=1, not + rope_type="default". + """ + swa = _get_rope(1, _MAX_POSITION) + + assert swa.mscale == 1.0 + # inv_freq_interpolation == inv_freq_extrapolation at factor=1, so the ramp + # blend `a * (1 - m) + a * m` is algebraically `a`; in fp32 it lands within + # a couple of ulp on the two pairs where the mask is neither 0 nor 1. + torch.testing.assert_close( + swa._compute_inv_freq(swa.scaling_factor), + _plain_inv_freq(), + rtol=1e-6, + atol=0, + ) + # Cache length is original_max_position * factor, so original_max_position + # has to carry the full range once factor is 1. + assert swa.cos_sin_cache.shape[0] == _MAX_POSITION + + +def test_yarn_config_still_interpolates_the_low_frequency_pairs( + default_vllm_config, +) -> None: + """The unfixed config is a real divergence, not a no-op.""" + yarn = _get_rope(_YARN_FACTOR, _ORIGINAL_MAX_POSITION) + + yarn_inv_freq = yarn._compute_inv_freq(yarn.scaling_factor) + plain = _plain_inv_freq() + assert not torch.allclose(yarn_inv_freq, plain) + + # Only the low-frequency pairs move: high-frequency pairs sit above the + # correction range and extrapolate unchanged. + moved = (yarn_inv_freq != plain).nonzero().flatten().tolist() + assert moved and min(moved) >= _ROPE_DIM // 4 + + # Accumulated rotation error across one 128-token SWA window. + worst_radians = ((plain - yarn_inv_freq).abs() * (_SWA_WINDOW - 1)).max().item() + assert 0.0 < math.degrees(worst_radians) < 5.0 + + +def test_shared_rope_parameters_dict_must_not_be_mutated(default_vllm_config) -> None: + """A shared rope_parameters dict must not be mutated in place. + + DeepseekV4Attention reads config.rope_parameters itself, and every later + layer reads it back, so setting factor=1 in place would disable YaRN + model-wide. + """ + shared = _rope_parameters(_YARN_FACTOR, _ORIGINAL_MAX_POSITION) + before = dict(shared) + + swa_parameters = dict(shared) + swa_parameters["factor"] = 1 + swa_parameters["original_max_position_embeddings"] = _MAX_POSITION + get_rope( + _HEAD_SIZE, + max_position=_MAX_POSITION, + rope_parameters=swa_parameters, + is_neox_style=False, + ) + + assert shared == before + compressing = _get_rope(_YARN_FACTOR, _ORIGINAL_MAX_POSITION) + assert compressing.scaling_factor == _YARN_FACTOR diff --git a/tests/model_executor/test_qwen3_5_quantization.py b/tests/model_executor/test_qwen3_5_quantization.py index ae7e7a79c3..74914bd765 100644 --- a/tests/model_executor/test_qwen3_5_quantization.py +++ b/tests/model_executor/test_qwen3_5_quantization.py @@ -1,13 +1,102 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace from unittest.mock import Mock, patch +import pytest +import torch + class _QuantConfig: def __init__(self) -> None: self.ignore: list[str] = [] self.config: dict[str, object] = {} + self.quantized_layers: dict[str, dict[str, str]] = {} + + +@pytest.mark.parametrize( + ("split_enabled", "norm_enabled"), + [ + (False, False), + (False, True), + ], +) +def test_sm70_gdn_qpn8_ba_split_is_disabled_without_split_flag( + monkeypatch, + split_enabled: bool, + norm_enabled: bool, +) -> None: + from vllm.model_executor.layers.mamba.gdn import qwen_gdn_linear_attn as gdn + + monkeypatch.setattr(gdn.envs, "VLLM_SM70_GDN_QPN8_BA_SPLIT", split_enabled) + monkeypatch.setattr(gdn.envs, "VLLM_SM70_GDN_RMSNORM_ONEPASS", norm_enabled) + + assert not gdn._sm70_gdn_qpn8_ba_split_enabled() + + +def test_sm70_gdn_qpn8_ba_split_rejects_unpaired_route(monkeypatch) -> None: + from vllm.model_executor.layers.mamba.gdn import qwen_gdn_linear_attn as gdn + + monkeypatch.setattr(gdn.envs, "VLLM_SM70_GDN_QPN8_BA_SPLIT", True) + monkeypatch.setattr(gdn.envs, "VLLM_SM70_GDN_RMSNORM_ONEPASS", False) + + with pytest.raises(RuntimeError, match="requires the accepted"): + gdn._sm70_gdn_qpn8_ba_split_enabled() + + +def test_sm70_gdn_qpn8_ba_split_requires_source_built_ops(monkeypatch) -> None: + from vllm.model_executor.layers.mamba.gdn import qwen_gdn_linear_attn as gdn + + monkeypatch.setattr(gdn.envs, "VLLM_SM70_GDN_QPN8_BA_SPLIT", True) + monkeypatch.setattr(gdn.envs, "VLLM_SM70_GDN_RMSNORM_ONEPASS", True) + monkeypatch.setattr( + gdn, + "_missing_sm70_gdn_qpn8_ba_ops", + lambda: ["fp8_qpn8_dispatch_ba_split_sm70_out"], + ) + + with pytest.raises(RuntimeError, match="requires the source-built"): + gdn._sm70_gdn_qpn8_ba_split_enabled() + + +def test_sm70_gdn_qpn8_ba_split_accepts_complete_contract(monkeypatch) -> None: + from vllm.model_executor.layers.mamba.gdn import qwen_gdn_linear_attn as gdn + + monkeypatch.setattr(gdn.envs, "VLLM_SM70_GDN_QPN8_BA_SPLIT", True) + monkeypatch.setattr(gdn.envs, "VLLM_SM70_GDN_RMSNORM_ONEPASS", True) + monkeypatch.setattr(gdn, "_missing_sm70_gdn_qpn8_ba_ops", lambda: []) + + assert gdn._sm70_gdn_qpn8_ba_split_enabled() + + +def test_sm70_gdn_qpn8_ba_weight_contract_is_layout_based() -> 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_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"), + bias=None, + ) + layer = SimpleNamespace( + enable_sm70_gdn_qpn8_ba_split=True, + in_proj_qkvz=qkvz, + in_proj_ba=ba, + ) + + assert gdn._sm70_gdn_qpn8_ba_weight_contract(layer) + + 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") + assert not gdn._sm70_gdn_qpn8_ba_weight_contract(layer) def test_qwen3_5_split_gdn_detects_compressed_tensors_ignore(): @@ -41,6 +130,53 @@ def test_qwen3_5_split_gdn_detects_compressed_tensors_config_ignore(): assert _uses_split_gdn_input_projections(quant_config) +def test_qwen3_5_split_gdn_detects_modelopt_mixed_unquantized_ba(): + from vllm.model_executor.models.qwen3_5 import ( + _uses_split_gdn_input_projections, + ) + + quant_config = _QuantConfig() + prefix = "model.language_model.layers.0.linear_attn" + quant_config.quantized_layers = { + f"{prefix}.in_proj_qkv": {"quant_algo": "FP8"}, + f"{prefix}.in_proj_z": {"quant_algo": "FP8"}, + } + + assert _uses_split_gdn_input_projections(quant_config) + + +def test_qwen3_5_split_gdn_keeps_same_precision_projections_fused(): + from vllm.model_executor.models.qwen3_5 import ( + _uses_split_gdn_input_projections, + ) + + quant_config = _QuantConfig() + prefix = "model.language_model.layers.0.linear_attn" + quant_config.quantized_layers = { + f"{prefix}.{name}": {"quant_algo": "FP8"} + for name in ("in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a") + } + + assert not _uses_split_gdn_input_projections(quant_config) + + +def test_qwen3_5_split_gdn_detects_different_modelopt_algorithms(): + from vllm.model_executor.models.qwen3_5 import ( + _uses_split_gdn_input_projections, + ) + + quant_config = _QuantConfig() + prefix = "model.language_model.layers.0.linear_attn" + quant_config.quantized_layers = { + f"{prefix}.in_proj_qkv": {"quant_algo": "FP8"}, + f"{prefix}.in_proj_z": {"quant_algo": "FP8"}, + f"{prefix}.in_proj_b": {"quant_algo": "W4A16_NVFP4"}, + f"{prefix}.in_proj_a": {"quant_algo": "W4A16_NVFP4"}, + } + + assert _uses_split_gdn_input_projections(quant_config) + + def test_qwen3_5_lm_head_receives_quant_config(): from vllm.model_executor.models.qwen3_5 import Qwen3_5ForCausalLMBase diff --git a/tests/model_executor/test_sm70_fp8_qpn8_pp2_tp4.py b/tests/model_executor/test_sm70_fp8_qpn8_pp2_tp4.py new file mode 100644 index 0000000000..17580656d8 --- /dev/null +++ b/tests/model_executor/test_sm70_fp8_qpn8_pp2_tp4.py @@ -0,0 +1,229 @@ +# 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 import envs +from vllm.model_executor.layers.quantization import fp8 + + +def _layer( + suffix: str, + tp_size: int, + k_dim: int, + n_dim: int, + *, + output_partition_sizes: list[int] | None = None, +) -> SimpleNamespace: + return SimpleNamespace( + prefix=f"arbitrary.engine.graph.{suffix}", + tp_size=tp_size, + input_size_per_partition=k_dim, + output_size_per_partition=n_dim, + output_partition_sizes=output_partition_sizes, + weight_block_size=[128, 128], + weight=torch.empty((n_dim, k_dim), device="meta"), + ) + + +def test_pp2_tp4_qpn8_is_default_on_with_explicit_rollback(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 fp8._sm70_fp8_qpn8_pp2_tp4_enabled() + + monkeypatch.setenv("VLLM_SM70_FP8_QPN8", "0") + envs.disable_envs_cache() + assert not fp8._sm70_fp8_qpn8_pp2_tp4_enabled() + + monkeypatch.delenv("VLLM_SM70_FP8_QPN8") + monkeypatch.setenv("VLLM_SM70_FP8_QPN8_PP2_TP4", "0") + envs.disable_envs_cache() + assert not fp8._sm70_fp8_qpn8_pp2_tp4_enabled() + finally: + envs.disable_envs_cache() + + +@pytest.mark.parametrize( + ("layer", "gated_silu", "expected"), + [ + (_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( + layer: SimpleNamespace, + gated_silu: bool, + expected: tuple[int, int, bool], +) -> None: + assert fp8._sm70_fp8_qpn8_pp2_tp4_config(layer, gated_silu=gated_silu) == expected + + +def test_pp2_tp4_qpn8_rejects_wrong_tensor_and_concurrency_roles() -> None: + wrong_tp = _layer("wq_b", 8, 1024, 8192) + assert fp8._sm70_fp8_qpn8_pp2_tp4_config(wrong_tp, gated_silu=False) is None + + concurrent_indexer = _layer("wq_b", 1, 1024, 8192) + assert ( + fp8._sm70_fp8_qpn8_pp2_tp4_config(concurrent_indexer, gated_silu=False) is None + ) + + excluded_gate = _layer( + "gate_up_proj", + 4, + 4096, + 1024, + output_partition_sizes=[512, 512], + ) + assert fp8._sm70_fp8_qpn8_pp2_tp4_config(excluded_gate, gated_silu=False) is None + assert fp8._sm70_fp8_qpn8_pp2_tp4_config(excluded_gate, gated_silu=True) is None + + wrong_layout = _layer("wo_b", 4, 2048, 4096) + wrong_layout.weight_block_size = [64, 128] + assert fp8._sm70_fp8_qpn8_pp2_tp4_config(wrong_layout, gated_silu=False) is None + + +def test_pp2_tp4_qpn8_runtime_and_workspace_contract() -> None: + 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, + ) + with patch.object(fp8, "get_current_vllm_config", return_value=config): + assert fp8._is_sm70_fp8_qpn8_pp2_tp4_runtime_contract() + config.speculative_config = object() + assert not fp8._is_sm70_fp8_qpn8_pp2_tp4_runtime_contract() + config.speculative_config = None + config.parallel_config.pipeline_parallel_size = 1 + assert not fp8._is_sm70_fp8_qpn8_pp2_tp4_runtime_contract() + config.parallel_config.pipeline_parallel_size = 2 + config.scheduler_config.max_num_seqs = 2 + assert not fp8._is_sm70_fp8_qpn8_pp2_tp4_runtime_contract() + config.scheduler_config.max_num_seqs = 1 + config.parallel_config.enable_dbo = True + assert not fp8._is_sm70_fp8_qpn8_pp2_tp4_runtime_contract() + config.parallel_config.enable_dbo = False + config.parallel_config.ubatch_size = 2 + assert not fp8._is_sm70_fp8_qpn8_pp2_tp4_runtime_contract() + + assert fp8._SM70_FP8_QPN8_PP2_TP4_WORKSPACE_ELEMENTS * 2 == 16 * 1024 * 1024 + + +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) == (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 = 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") + + calls: list[tuple[tuple[int, ...], int, int, bool, bool]] = [] + + def fake_dispatch( + out: torch.Tensor, + workspace_ptr: int, + input_: torch.Tensor, + codes: torch.Tensor, + scales: torch.Tensor, + split_k: int, + nacc: int, + prefetch: bool, + gated_silu: bool, + ) -> None: + del codes, scales + calls.append( + (tuple(input_.shape), workspace_ptr, split_k, prefetch, gated_silu) + ) + out.fill_(len(calls)) + 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, 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: + monkeypatch.delenv("VLLM_SM70_FP8_QPN8", raising=False) + monkeypatch.delenv("VLLM_SM70_FP8_QPN8_PP2_TP4", raising=False) + envs.disable_envs_cache() + layer = _layer("fused_wqa_wkv", 1, 4096, 1536) + layer.orig_dtype = torch.float16 + layer.is_bmm = False + layer.weight_scale_inv = torch.empty((12, 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, 1536), dtype=torch.uint8, device="meta"), + torch.empty((256, 48), 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 diff --git a/tests/model_executor/test_uva_tower_offload.py b/tests/model_executor/test_uva_tower_offload.py new file mode 100644 index 0000000000..91e3d5de93 --- /dev/null +++ b/tests/model_executor/test_uva_tower_offload.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Generator +from types import SimpleNamespace + +import torch.nn as nn + +from vllm.model_executor.models.interfaces import SupportsMultiModal +from vllm.model_executor.offloader import ( + BaseOffloader, + PrefetchOffloader, + get_offloader, + set_offloader, +) +from vllm.model_executor.offloader.uva import UVAOffloader + + +class _RecordingOffloader(BaseOffloader): + def __init__(self) -> None: + self.component_calls: list[tuple[nn.Module, str]] = [] + + def wrap_modules( + self, modules_generator: Generator[nn.Module, None, None] + ) -> list[nn.Module]: + raise AssertionError("tower registration must not reuse the layer wrapper") + + def wrap_module( + self, module: nn.Module, *, parameter_prefix: str = "" + ) -> nn.Module: + self.component_calls.append((module, parameter_prefix)) + return module + + +class _MultiModalModel(nn.Module): + _mark_tower_model = SupportsMultiModal._mark_tower_model + + +def _selector_only_offloader(selectors: set[str]) -> UVAOffloader: + offloader = UVAOffloader.__new__(UVAOffloader) + offloader.cpu_offload_params = selectors + return offloader + + +def test_uva_tower_selector_uses_qualified_parameter_path() -> None: + offloader = _selector_only_offloader({"visual"}) + + assert offloader._is_parameter_selected("patch_embed.proj.weight", "visual") + assert not offloader._is_parameter_selected("patch_embed.proj.weight", "audio") + assert not offloader._is_parameter_selected("patch_embed.proj.weight", "visualizer") + assert _selector_only_offloader(set())._is_parameter_selected( + "patch_embed.proj.weight", "visual" + ) + + +def test_uva_multi_segment_selector_remains_segment_exact() -> None: + offloader = _selector_only_offloader({"visual.patch_embed"}) + + assert offloader._is_parameter_selected("patch_embed.proj.weight", "visual") + assert not offloader._is_parameter_selected("patch_embedding.proj.weight", "visual") + + +def test_tower_registration_uses_single_component_wrapper() -> None: + model = _MultiModalModel() + offloader = _RecordingOffloader() + previous_offloader = get_offloader() + config = SimpleNamespace( + model_config=SimpleNamespace( + multimodal_config=SimpleNamespace(get_limit_per_prompt=lambda _modality: 1) + ) + ) + + set_offloader(offloader) + try: + with model._mark_tower_model(config, "image"): + model.visual = nn.Linear(4, 4) + finally: + set_offloader(previous_offloader) + + assert model._tower_model_names == ["visual"] + assert offloader.component_calls == [(model.visual, "visual")] + + +def test_prefetch_inherits_noop_component_wrapper() -> None: + tower = nn.Linear(4, 4) + offloader = PrefetchOffloader.__new__(PrefetchOffloader) + + assert PrefetchOffloader.wrap_module is BaseOffloader.wrap_module + assert offloader.wrap_module(tower, parameter_prefix="visual") is tower 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/test_deepseek_v4_pipeline.py b/tests/models/test_deepseek_v4_pipeline.py new file mode 100644 index 0000000000..2920cd5c2c --- /dev/null +++ b/tests/models/test_deepseek_v4_pipeline.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.distributed.utils import get_pp_indices +from vllm.models.deepseek_v4.nvidia.model import DeepseekV4Model + + +@pytest.mark.parametrize( + ("pp_size", "expected_partitions"), + [ + (2, [(0, 22), (22, 43)]), + (4, [(0, 11), (11, 22), (22, 33), (33, 43)]), + ], +) +def test_deepseek_v4_pipeline_partitions_keep_dspark_targets_on_last_stage( + monkeypatch: pytest.MonkeyPatch, + pp_size: int, + expected_partitions: list[tuple[int, int]], +) -> None: + monkeypatch.delenv("VLLM_PP_LAYER_PARTITION", raising=False) + + partitions = [get_pp_indices(43, rank, pp_size) for rank in range(pp_size)] + + assert partitions == expected_partitions + last_start, last_end = partitions[-1] + assert all(last_start <= layer_id < last_end for layer_id in (40, 41, 42)) + + +def test_deepseek_v4_pipeline_intermediate_preserves_mhc_streams() -> None: + model = SimpleNamespace( + hc_mult=4, + config=SimpleNamespace(hidden_size=4096), + ) + + intermediate = DeepseekV4Model.make_empty_intermediate_tensors( + model, + batch_size=3, + dtype=torch.float16, + device=torch.device("cpu"), + ) + + assert set(intermediate.tensors) == {"hidden_states"} + assert intermediate["hidden_states"].shape == (3, 4, 4096) + assert intermediate["hidden_states"].dtype == torch.float16 + assert torch.count_nonzero(intermediate["hidden_states"]) == 0 diff --git a/tests/models/test_deepseek_v4_sm70_routes.py b/tests/models/test_deepseek_v4_sm70_routes.py index 0c5ea28f90..a983b2b206 100644 --- a/tests/models/test_deepseek_v4_sm70_routes.py +++ b/tests/models/test_deepseek_v4_sm70_routes.py @@ -101,6 +101,14 @@ def test_sm70_sparse_qk_dsplit_uses_graph_workspace(): kwargs = qk_dsplit.call_args.kwargs assert kwargs["partial_qk"].shape == (1, 8, 8, 8, 16) assert kwargs["partial_probs"].shape == (1, 8, 8, 16) + assert kwargs["stage1_block_h"] == 8 + + +def test_sm70_sparse_qk_dsplit_uses_one_tp4_head_group(): + from vllm.models.deepseek_v4.sm70.sparse import _qk_dsplit_block_h + + assert _qk_dsplit_block_h(16) == 16 + assert _qk_dsplit_block_h(8) == 8 def test_sm75_does_not_select_sm70_impl(): @@ -197,6 +205,52 @@ def make_metadata(starts: list[int]) -> CommonAttentionMetadata: assert _get_c128_boundary(make_metadata([127, 10])) is True +def test_sm70_private_compressor_state_requires_a_contiguous_single_request(): + from vllm.models.deepseek_v4 import compressor + + platform = MagicMock() + platform.is_cuda.return_value = True + platform.is_device_capability.return_value = True + + config = MagicMock() + config.scheduler_config.max_num_seqs = 1 + config.cache_config.enable_prefix_caching = False + config.parallel_config.pipeline_parallel_size = 1 + config.kv_transfer_config = None + config.speculative_config.parallel_drafting = False + + with ( + patch.object(compressor, "current_platform", platform), + patch.object( + compressor.envs, + "VLLM_SM70_DSV4_PRIVATE_COMPRESSOR_STATE", + False, + ), + ): + assert not compressor._can_use_sm70_private_compressor_state(config) + + with ( + patch.object(compressor, "current_platform", platform), + patch.object( + compressor.envs, + "VLLM_SM70_DSV4_PRIVATE_COMPRESSOR_STATE", + True, + ), + ): + assert compressor._can_use_sm70_private_compressor_state(config) + + config.speculative_config.parallel_drafting = True + assert not compressor._can_use_sm70_private_compressor_state(config) + + config.speculative_config.parallel_drafting = False + config.scheduler_config.max_num_seqs = 2 + assert not compressor._can_use_sm70_private_compressor_state(config) + + config.scheduler_config.max_num_seqs = 1 + config.cache_config.enable_prefix_caching = True + assert not compressor._can_use_sm70_private_compressor_state(config) + + def test_v4_c128_metadata_keeps_graph_stable_row_stride(): from vllm.v1.attention.backends.mla import flashmla_sparse diff --git a/tests/quantization/test_nvfp4_mxfp4_gating.py b/tests/quantization/test_nvfp4_mxfp4_gating.py index 6f03dc5ade..d3f0ed641c 100644 --- a/tests/quantization/test_nvfp4_mxfp4_gating.py +++ b/tests/quantization/test_nvfp4_mxfp4_gating.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + import torch from torch.nn.parameter import Parameter @@ -18,6 +20,9 @@ from vllm.model_executor.layers.quantization.compressed_tensors.schemes import ( compressed_tensors_w4a4_mxfp4 as mxfp4_scheme, ) +from vllm.model_executor.layers.quantization.compressed_tensors.schemes import ( + compressed_tensors_w4a4_nvfp4 as nvfp4_scheme, +) from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_w4a4_mxfp4 import ( # noqa: E501 CompressedTensorsW4A4Mxfp4, ) @@ -93,6 +98,92 @@ def test_flashinfer_nvfp4_backends_reject_sm70(): assert "sm_100" in cudnn_reason +def _make_sm70_tp4_nvfp4_gate_up_layer() -> torch.nn.Module: + layer = torch.nn.Module() + layer.tp_size = 4 + layer.prefix = "model.language_model.layers.0.mlp.gate_up_proj" + layer.input_size_per_partition = 5120 + layer.output_size_per_partition = 8704 + layer.logical_widths = [4352, 4352] + return layer + + +def _make_sm70_tp4_nvfp4_down_layer() -> torch.nn.Module: + layer = torch.nn.Module() + layer.tp_size = 4 + layer.prefix = "model.language_model.layers.0.mlp.down_proj" + layer.input_size_per_partition = 4352 + layer.output_size_per_partition = 5120 + return layer + + +def test_sm70_tp4_nvfp4_gate_up_match_is_shape_exact(): + layer = _make_sm70_tp4_nvfp4_gate_up_layer() + + assert nvfp4_scheme._is_sm70_tp4_nvfp4_gate_up(layer) + + layer.tp_size = 2 + assert not nvfp4_scheme._is_sm70_tp4_nvfp4_gate_up(layer) + layer.tp_size = 4 + layer.output_size_per_partition = 8703 + assert not nvfp4_scheme._is_sm70_tp4_nvfp4_gate_up(layer) + layer.output_size_per_partition = 8704 + layer.prefix = "model.language_model.layers.0.mlp.down_proj" + assert not nvfp4_scheme._is_sm70_tp4_nvfp4_gate_up(layer) + + +def test_sm70_tp4_nvfp4_down_match_is_shape_exact(): + layer = _make_sm70_tp4_nvfp4_down_layer() + + assert nvfp4_scheme._is_sm70_tp4_nvfp4_down(layer) + + layer.input_size_per_partition = 4351 + assert not nvfp4_scheme._is_sm70_tp4_nvfp4_down(layer) + layer.input_size_per_partition = 4352 + layer.output_size_per_partition = 5119 + assert not nvfp4_scheme._is_sm70_tp4_nvfp4_down(layer) + + +def test_sm70_nvfp4_qpn4_runtime_contract_is_single_sequence_no_mtp( + monkeypatch, +): + config = SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_seqs=1), + speculative_config=None, + ) + monkeypatch.setattr( + nvfp4_scheme, + "get_current_vllm_config", + lambda: config, + ) + + assert nvfp4_scheme._is_sm70_nvfp4_qpn4_runtime_contract() + + config.scheduler_config.max_num_seqs = 2 + assert not nvfp4_scheme._is_sm70_nvfp4_qpn4_runtime_contract() + config.scheduler_config.max_num_seqs = 1 + config.speculative_config = object() + assert not nvfp4_scheme._is_sm70_nvfp4_qpn4_runtime_contract() + + +def test_nvfp4_scheme_delegates_fused_silu_only_for_prepared_layer(monkeypatch): + scheme = CompressedTensorsW4A4Fp4() + layer = torch.nn.Module() + x = torch.ones((1, 4), dtype=torch.float16) + expected = torch.full((1, 2), 7.0, dtype=torch.float16) + + monkeypatch.setattr(nvfp4_scheme.sm70_tm, "has_prepared_linear", lambda _: True) + monkeypatch.setattr( + nvfp4_scheme.sm70_tm, + "apply_prepared_fused_silu_and_mul", + lambda prepared_layer, prepared_x: expected + if prepared_layer is layer and prepared_x is x + else None, + ) + + assert scheme.apply_fused_silu_and_mul(layer, x) is expected + + def _make_mxfp4_layer() -> torch.nn.Module: layer = torch.nn.Module() layer.weight_packed = Parameter( diff --git a/tests/quantization/test_sm70_compressed_tensors_fastpaths.py b/tests/quantization/test_sm70_compressed_tensors_fastpaths.py new file mode 100644 index 0000000000..2512a5864c --- /dev/null +++ b/tests/quantization/test_sm70_compressed_tensors_fastpaths.py @@ -0,0 +1,293 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import torch +from compressed_tensors.quantization import ( + QuantizationArgs, + QuantizationStrategy, + QuantizationType, +) + +import vllm.envs as envs +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501 + CompressedTensorsLinearMethod, +) +from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_w8a16_fp8 import ( # noqa: E501 + CompressedTensorsW8A16Fp8, + _sm70_channel_fp8_qpn8_config, +) + + +def _make_channel_fp8_scheme(monkeypatch, *, static_input: bool = False): + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.compressed_tensors.schemes." + "compressed_tensors_w8a16_fp8.get_current_vllm_config", + lambda: SimpleNamespace(model_config=SimpleNamespace(dtype=torch.float16)), + ) + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.compressed_tensors.schemes." + "compressed_tensors_w8a16_fp8.sm70_tm.is_exact_sm70_cuda_platform", + lambda: True, + ) + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.compressed_tensors.schemes." + "compressed_tensors_w8a16_fp8.sm70_tm.use_turbomind", + lambda enabled: enabled, + ) + weight_quant = QuantizationArgs( + num_bits=8, + type=QuantizationType.FLOAT, + symmetric=True, + strategy=QuantizationStrategy.CHANNEL, + dynamic=False, + ) + return CompressedTensorsW8A16Fp8(weight_quant, static_input) + + +def test_compressed_tensors_channel_fp8_selects_turbomind_on_sm70(monkeypatch): + monkeypatch.setenv("VLLM_SM70_FP8_TURBOMIND", "1") + envs.disable_envs_cache() + try: + assert _make_channel_fp8_scheme(monkeypatch).use_sm70_fp8_turbomind + assert not _make_channel_fp8_scheme( + monkeypatch, static_input=True + ).use_sm70_fp8_turbomind + finally: + envs.disable_envs_cache() + + +def test_compressed_tensors_channel_fp8_prepares_and_dispatches_turbomind( + monkeypatch, +): + scheme = object.__new__(CompressedTensorsW8A16Fp8) + scheme.use_sm70_fp8_turbomind = True + layer = torch.nn.Module() + layer.register_parameter( + "weight", + torch.nn.Parameter( + torch.empty((6, 4), dtype=torch.float8_e4m3fn), requires_grad=False + ), + ) + layer.register_parameter( + "weight_scale", + torch.nn.Parameter(torch.ones((6, 1)), requires_grad=False), + ) + layer.orig_dtype = torch.float16 + layer.output_size_per_partition = 6 + prepare_calls = [] + + monkeypatch.setattr(torch.ops._C, "fp8_sm70_prepare", object(), raising=False) + + def fake_prepare(weight, scales, group_size, gated_silu): + prepare_calls.append( + (tuple(weight.shape), tuple(scales.shape), group_size, gated_silu) + ) + return ( + torch.empty((4, 6), dtype=torch.uint8), + torch.ones((1, 6), dtype=torch.float16), + torch.tensor([4, 6]), + ) + + module = ( + "vllm.model_executor.layers.quantization.compressed_tensors.schemes." + "compressed_tensors_w8a16_fp8" + ) + monkeypatch.setattr(f"{module}.sm70_ops.fp8_sm70_prepare", fake_prepare) + scheme.process_weights_after_loading(layer) + + assert prepare_calls == [((6, 4), (6, 1), 128, False)] + assert layer.sm70_fp8_turbomind + assert layer.sm70_fp8_channel_scale + assert tuple(layer.weight.shape) == (4, 6) + assert tuple(layer.weight_scale_inv.shape) == (1, 6) + assert not hasattr(layer, "weight_scale") + + gemm_calls = [] + + def fake_gemm(out, x, weight, scales, group_size, k_ld, q_ld, gated_silu): + gemm_calls.append( + ( + tuple(out.shape), + tuple(x.shape), + tuple(weight.shape), + tuple(scales.shape), + group_size, + k_ld, + q_ld, + gated_silu, + ) + ) + out.fill_(2) + + monkeypatch.setattr(f"{module}.sm70_ops.fp8_gemm_sm70_out", fake_gemm) + output = scheme.apply_weights( + layer, + torch.ones((2, 4), dtype=torch.float16), + torch.ones(6, dtype=torch.float16), + ) + + assert torch.equal(output, torch.full((2, 6), 3, dtype=torch.float16)) + assert gemm_calls == [((2, 6), (2, 4), (4, 6), (1, 6), 128, 4, 6, False)] + + +def test_compressed_tensors_channel_fp8_qpn8_shape_gate_is_exact(): + layer = SimpleNamespace( + tp_size=4, + prefix="model.language_model.layers.0.linear_attn.in_proj_qkvz", + weight=SimpleNamespace(shape=(4096, 5120)), + ) + assert _sm70_channel_fp8_qpn8_config(layer) == (16, 2, False) + + layer.prefix = "model.language_model.layers.3.self_attn.qkv_proj" + layer.weight = SimpleNamespace(shape=(3584, 5120)) + assert _sm70_channel_fp8_qpn8_config(layer) == (16, 2, False) + + layer.prefix = "model.language_model.layers.3.self_attn.o_proj" + layer.weight = SimpleNamespace(shape=(5120, 1536)) + assert _sm70_channel_fp8_qpn8_config(layer) == (12, 2, False) + + layer.prefix = "model.language_model.lm_head" + layer.weight = SimpleNamespace(shape=(62080, 5120)) + assert _sm70_channel_fp8_qpn8_config(layer) is None + + layer.tp_size = 2 + layer.prefix = "model.language_model.layers.3.self_attn.o_proj" + layer.weight = SimpleNamespace(shape=(5120, 1536)) + assert _sm70_channel_fp8_qpn8_config(layer) is None + + +def test_compressed_tensors_channel_fp8_qpn8_prepares_and_dispatches(monkeypatch): + scheme = object.__new__(CompressedTensorsW8A16Fp8) + scheme.use_sm70_fp8_turbomind = True + scheme.use_sm70_fp8_qpn8 = True + layer = torch.nn.Module() + layer.register_parameter( + "weight", + torch.nn.Parameter( + torch.empty((6, 4), dtype=torch.float8_e4m3fn), requires_grad=False + ), + ) + layer.register_parameter( + "weight_scale", + torch.nn.Parameter(torch.ones((6, 1)), requires_grad=False), + ) + layer.orig_dtype = torch.float16 + layer.output_size_per_partition = 6 + layer.prefix = "model.language_model.layers.63.mlp.gate_up_proj" + workspace = torch.empty(1, dtype=torch.float16) + prepare_calls = [] + + module = ( + "vllm.model_executor.layers.quantization.compressed_tensors.schemes." + "compressed_tensors_w8a16_fp8" + ) + monkeypatch.setattr(torch.ops._C, "fp8_sm70_prepare", object(), raising=False) + monkeypatch.setattr( + f"{module}._sm70_channel_fp8_qpn8_config", lambda layer: (16, 2, False) + ) + monkeypatch.setattr(f"{module}._is_sm70_fp8_qpn8_runtime_contract", lambda: True) + monkeypatch.setattr(f"{module}._missing_sm70_fp8_qpn8_ops", lambda: []) + monkeypatch.setattr( + f"{module}._get_sm70_fp8_prefill_exact_dense_workspace", + lambda weight: workspace, + ) + + def fake_prepare(weight, scales): + prepare_calls.append((tuple(weight.shape), tuple(scales.shape))) + return ( + torch.empty((4, 6), dtype=torch.uint8), + torch.ones((1, 6), dtype=torch.float16), + ) + + monkeypatch.setattr(f"{module}.sm70_ops.fp8_qpn8_prepare_sm70", fake_prepare) + scheme.process_weights_after_loading(layer) + + assert prepare_calls == [((6, 4), (6, 1))] + assert layer.sm70_fp8_turbomind + assert layer.sm70_fp8_qpn8 + assert layer.sm70_fp8_channel_scale + assert layer.sm70_fp8_qpn8_split_k == 16 + assert layer.sm70_fp8_qpn8_nacc == 2 + assert not layer.sm70_fp8_qpn8_prefetch + assert layer.sm70_fp8_gated_silu + assert layer.sm70_fp8_qpn8_gated_split_k == 8 + assert layer.sm70_fp8_qpn8_gated_nacc == 2 + assert not layer.sm70_fp8_qpn8_gated_prefetch + assert tuple(layer.weight.shape) == (4, 6) + assert tuple(layer.weight_scale_inv.shape) == (1, 6) + assert not hasattr(layer, "weight_scale") + + dispatch_calls = [] + + def fake_dispatch( + out, workspace_ptr, x, weight, scales, split_k, nacc, prefetch, gated_silu + ): + dispatch_calls.append( + ( + tuple(out.shape), + tuple(x.shape), + workspace_ptr, + tuple(weight.shape), + tuple(scales.shape), + split_k, + nacc, + prefetch, + gated_silu, + ) + ) + out.fill_(2) + + monkeypatch.setattr(f"{module}.sm70_ops.fp8_qpn8_dispatch_sm70_out", fake_dispatch) + output = scheme.apply_weights( + layer, + torch.ones((2, 4), dtype=torch.float16), + torch.ones(6, dtype=torch.float16), + ) + + assert torch.equal(output, torch.full((2, 6), 3, dtype=torch.float16)) + fused_output = scheme.apply_fused_silu_and_mul( + layer, torch.ones((2, 4), dtype=torch.float16) + ) + assert fused_output is not None + assert torch.equal(fused_output, torch.full((2, 3), 2, dtype=torch.float16)) + assert dispatch_calls == [ + ( + (2, 6), + (2, 4), + workspace.data_ptr(), + (4, 6), + (1, 6), + 16, + 2, + False, + False, + ), + ( + (2, 3), + (2, 4), + workspace.data_ptr(), + (4, 6), + (1, 6), + 8, + 2, + False, + True, + ), + ] + + +def test_compressed_tensors_linear_method_delegates_fused_apply(): + expected = torch.ones((1, 3), dtype=torch.float16) + layer = SimpleNamespace( + scheme=SimpleNamespace( + apply_fused_silu_and_mul=lambda layer, x: expected, + ) + ) + method = object.__new__(CompressedTensorsLinearMethod) + + assert ( + method.apply_fused_silu_and_mul(layer, torch.ones((1, 4), dtype=torch.float16)) + is expected + ) diff --git a/tests/quantization/test_sm70_fp8_prefill_exact_dense.py b/tests/quantization/test_sm70_fp8_prefill_exact_dense.py index 800e658dc7..166f871954 100644 --- a/tests/quantization/test_sm70_fp8_prefill_exact_dense.py +++ b/tests/quantization/test_sm70_fp8_prefill_exact_dense.py @@ -4,18 +4,344 @@ from types import SimpleNamespace import torch +from compressed_tensors.quantization import ( + QuantizationArgs, + QuantizationStrategy, + QuantizationType, +) import vllm.envs as envs +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501 + CompressedTensorsLinearMethod, +) +from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_w8a16_fp8 import ( # noqa: E501 + CompressedTensorsW8A16Fp8, + _sm70_channel_fp8_qpn8_config, + _sm70_fp8_qpn8_enabled, +) from vllm.model_executor.layers.quantization.fp8 import ( + _SM70_FP8_EXACT_8K_PREFILL_M, _SM70_FP8_PREFILL_DENSE_MIN_M, _SM70_FP8_PREFILL_DENSE_WORKSPACE_BYTES, Fp8LinearMethod, _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, ) +def _make_channel_fp8_scheme( + monkeypatch, + *, + static_input: bool = False, + enable_qpn8_by_default: bool = False, +): + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.compressed_tensors.schemes." + "compressed_tensors_w8a16_fp8.get_current_vllm_config", + lambda: SimpleNamespace(model_config=SimpleNamespace(dtype=torch.float16)), + ) + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.compressed_tensors.schemes." + "compressed_tensors_w8a16_fp8.sm70_tm.is_exact_sm70_cuda_platform", + lambda: True, + ) + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.compressed_tensors.schemes." + "compressed_tensors_w8a16_fp8.sm70_tm.use_turbomind", + lambda enabled: enabled, + ) + weight_quant = QuantizationArgs( + num_bits=8, + type=QuantizationType.FLOAT, + symmetric=True, + strategy=QuantizationStrategy.CHANNEL, + dynamic=False, + ) + return CompressedTensorsW8A16Fp8( + weight_quant, + static_input, + enable_sm70_qpn8_by_default=enable_qpn8_by_default, + ) + + +def test_compressed_tensors_channel_fp8_selects_turbomind_on_sm70(monkeypatch): + monkeypatch.setenv("VLLM_SM70_FP8_TURBOMIND", "1") + envs.disable_envs_cache() + try: + assert _make_channel_fp8_scheme(monkeypatch).use_sm70_fp8_turbomind + assert not _make_channel_fp8_scheme( + monkeypatch, static_input=True + ).use_sm70_fp8_turbomind + finally: + envs.disable_envs_cache() + + +def test_channel_fp8_qpn8_uses_mixed_nvfp4_default_only(monkeypatch): + monkeypatch.delenv("VLLM_SM70_FP8_QPN8", raising=False) + envs.disable_envs_cache() + try: + assert not _make_channel_fp8_scheme(monkeypatch).use_sm70_fp8_qpn8 + assert _make_channel_fp8_scheme( + monkeypatch, + enable_qpn8_by_default=True, + ).use_sm70_fp8_qpn8 + + monkeypatch.setenv("VLLM_SM70_FP8_QPN8", "0") + envs.disable_envs_cache() + assert not _make_channel_fp8_scheme( + monkeypatch, + enable_qpn8_by_default=True, + ).use_sm70_fp8_qpn8 + finally: + envs.disable_envs_cache() + + +def test_compressed_tensors_channel_fp8_prepares_and_dispatches_turbomind( + monkeypatch, +): + scheme = object.__new__(CompressedTensorsW8A16Fp8) + scheme.use_sm70_fp8_turbomind = True + layer = torch.nn.Module() + layer.register_parameter( + "weight", + torch.nn.Parameter( + torch.empty((6, 4), dtype=torch.float8_e4m3fn), requires_grad=False + ), + ) + layer.register_parameter( + "weight_scale", + torch.nn.Parameter(torch.ones((6, 1)), requires_grad=False), + ) + layer.orig_dtype = torch.float16 + layer.output_size_per_partition = 6 + prepare_calls = [] + + monkeypatch.setattr(torch.ops._C, "fp8_sm70_prepare", object(), raising=False) + + def fake_prepare(weight, scales, group_size, gated_silu): + prepare_calls.append( + (tuple(weight.shape), tuple(scales.shape), group_size, gated_silu) + ) + return ( + torch.empty((4, 6), dtype=torch.uint8), + torch.ones((1, 6), dtype=torch.float16), + torch.tensor([4, 6]), + ) + + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.compressed_tensors.schemes." + "compressed_tensors_w8a16_fp8.sm70_ops.fp8_sm70_prepare", + fake_prepare, + ) + scheme.process_weights_after_loading(layer) + + assert prepare_calls == [((6, 4), (6, 1), 128, False)] + assert layer.sm70_fp8_turbomind + assert layer.sm70_fp8_channel_scale + assert tuple(layer.weight.shape) == (4, 6) + assert tuple(layer.weight_scale_inv.shape) == (1, 6) + assert not hasattr(layer, "weight_scale") + + gemm_calls = [] + + def fake_gemm(out, x, weight, scales, group_size, k_ld, q_ld, gated_silu): + gemm_calls.append( + ( + tuple(out.shape), + tuple(x.shape), + tuple(weight.shape), + tuple(scales.shape), + group_size, + k_ld, + q_ld, + gated_silu, + ) + ) + out.fill_(2) + + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.compressed_tensors.schemes." + "compressed_tensors_w8a16_fp8.sm70_ops.fp8_gemm_sm70_out", + fake_gemm, + ) + output = scheme.apply_weights( + layer, + torch.ones((2, 4), dtype=torch.float16), + torch.ones(6, dtype=torch.float16), + ) + + assert torch.equal(output, torch.full((2, 6), 3, dtype=torch.float16)) + assert gemm_calls == [((2, 6), (2, 4), (4, 6), (1, 6), 128, 4, 6, False)] + + +def test_compressed_tensors_channel_fp8_qpn8_shape_gate_is_exact(): + layer = SimpleNamespace( + tp_size=4, + prefix="model.language_model.layers.0.linear_attn.in_proj_qkvz", + weight=SimpleNamespace(shape=(4096, 5120)), + ) + assert _sm70_channel_fp8_qpn8_config(layer) == (16, 2, False) + + layer.prefix = "model.language_model.layers.3.self_attn.qkv_proj" + layer.weight = SimpleNamespace(shape=(3584, 5120)) + assert _sm70_channel_fp8_qpn8_config(layer) == (16, 2, False) + + layer.prefix = "model.language_model.layers.3.self_attn.o_proj" + layer.weight = SimpleNamespace(shape=(5120, 1536)) + assert _sm70_channel_fp8_qpn8_config(layer) == (12, 2, False) + + layer.prefix = "model.language_model.lm_head" + layer.weight = SimpleNamespace(shape=(62080, 5120)) + assert _sm70_channel_fp8_qpn8_config(layer) is None + + layer.tp_size = 2 + layer.prefix = "model.language_model.layers.3.self_attn.o_proj" + layer.weight = SimpleNamespace(shape=(5120, 1536)) + assert _sm70_channel_fp8_qpn8_config(layer) is None + + +def test_compressed_tensors_channel_fp8_qpn8_prepares_and_dispatches(monkeypatch): + scheme = object.__new__(CompressedTensorsW8A16Fp8) + scheme.use_sm70_fp8_turbomind = True + scheme.use_sm70_fp8_qpn8 = True + layer = torch.nn.Module() + layer.register_parameter( + "weight", + torch.nn.Parameter( + torch.empty((6, 4), dtype=torch.float8_e4m3fn), requires_grad=False + ), + ) + layer.register_parameter( + "weight_scale", + torch.nn.Parameter(torch.ones((6, 1)), requires_grad=False), + ) + layer.orig_dtype = torch.float16 + layer.output_size_per_partition = 6 + layer.prefix = "model.language_model.layers.63.mlp.gate_up_proj" + workspace = torch.empty(1, dtype=torch.float16) + prepare_calls = [] + + module = ( + "vllm.model_executor.layers.quantization.compressed_tensors.schemes." + "compressed_tensors_w8a16_fp8" + ) + monkeypatch.setattr(torch.ops._C, "fp8_sm70_prepare", object(), raising=False) + monkeypatch.setattr( + f"{module}._sm70_channel_fp8_qpn8_config", lambda layer: (16, 2, False) + ) + monkeypatch.setattr(f"{module}._is_sm70_fp8_qpn8_runtime_contract", lambda: True) + monkeypatch.setattr(f"{module}._missing_sm70_fp8_qpn8_ops", lambda: []) + monkeypatch.setattr( + f"{module}._get_sm70_fp8_prefill_exact_dense_workspace", + lambda weight: workspace, + ) + + def fake_prepare(weight, scales): + prepare_calls.append((tuple(weight.shape), tuple(scales.shape))) + return ( + torch.empty((4, 6), dtype=torch.uint8), + torch.ones((1, 6), dtype=torch.float16), + ) + + monkeypatch.setattr(f"{module}.sm70_ops.fp8_qpn8_prepare_sm70", fake_prepare) + scheme.process_weights_after_loading(layer) + + assert prepare_calls == [((6, 4), (6, 1))] + assert layer.sm70_fp8_turbomind + assert layer.sm70_fp8_qpn8 + assert layer.sm70_fp8_channel_scale + assert layer.sm70_fp8_qpn8_split_k == 16 + assert layer.sm70_fp8_qpn8_nacc == 2 + assert not layer.sm70_fp8_qpn8_prefetch + assert layer.sm70_fp8_gated_silu + assert layer.sm70_fp8_qpn8_gated_split_k == 8 + assert layer.sm70_fp8_qpn8_gated_nacc == 2 + assert not layer.sm70_fp8_qpn8_gated_prefetch + assert tuple(layer.weight.shape) == (4, 6) + assert tuple(layer.weight_scale_inv.shape) == (1, 6) + assert not hasattr(layer, "weight_scale") + + dispatch_calls = [] + + def fake_dispatch( + out, workspace_ptr, x, weight, scales, split_k, nacc, prefetch, gated_silu + ): + dispatch_calls.append( + ( + tuple(out.shape), + tuple(x.shape), + workspace_ptr, + tuple(weight.shape), + tuple(scales.shape), + split_k, + nacc, + prefetch, + gated_silu, + ) + ) + out.fill_(2) + + monkeypatch.setattr(f"{module}.sm70_ops.fp8_qpn8_dispatch_sm70_out", fake_dispatch) + output = scheme.apply_weights( + layer, + torch.ones((2, 4), dtype=torch.float16), + torch.ones(6, dtype=torch.float16), + ) + + assert torch.equal(output, torch.full((2, 6), 3, dtype=torch.float16)) + fused_output = scheme.apply_fused_silu_and_mul( + layer, torch.ones((2, 4), dtype=torch.float16) + ) + assert fused_output is not None + assert torch.equal(fused_output, torch.full((2, 3), 2, dtype=torch.float16)) + assert dispatch_calls == [ + ( + (2, 6), + (2, 4), + workspace.data_ptr(), + (4, 6), + (1, 6), + 16, + 2, + False, + False, + ), + ( + (2, 3), + (2, 4), + workspace.data_ptr(), + (4, 6), + (1, 6), + 8, + 2, + False, + True, + ), + ] + + +def test_compressed_tensors_linear_method_delegates_fused_apply(): + expected = torch.ones((1, 3), dtype=torch.float16) + layer = SimpleNamespace( + scheme=SimpleNamespace( + apply_fused_silu_and_mul=lambda layer, x: expected, + ) + ) + method = object.__new__(CompressedTensorsLinearMethod) + + assert ( + method.apply_fused_silu_and_mul(layer, torch.ones((1, 4), dtype=torch.float16)) + is expected + ) + + def test_fp8_prefill_exact_dense_is_default_on(monkeypatch): monkeypatch.delenv("VLLM_SM70_FP8_PREFILL_EXACT_DENSE", raising=False) envs.disable_envs_cache() @@ -25,9 +351,49 @@ def test_fp8_prefill_exact_dense_is_default_on(monkeypatch): envs.disable_envs_cache() +def test_fp8_exact_8k_prefill_controls_are_generic_defaults(monkeypatch): + for name in ( + "VLLM_SM70_FP8_PREFILL_FAST_SELECTOR", + "VLLM_SM70_FP8_PREFILL_PRESCALED", + "VLLM_SM70_FP8_PREFILL_CUTLASS", + "VLLM_SM70_FP8_PREFILL_VISIBLE_DENSE_MM", + ): + monkeypatch.delenv(name, raising=False) + envs.disable_envs_cache() + try: + assert envs.VLLM_SM70_FP8_PREFILL_FAST_SELECTOR + assert envs.VLLM_SM70_FP8_PREFILL_PRESCALED + assert envs.VLLM_SM70_FP8_PREFILL_CUTLASS + assert not envs.VLLM_SM70_FP8_PREFILL_VISIBLE_DENSE_MM + finally: + envs.disable_envs_cache() + + +def test_fp8_qpn8_is_opt_in_except_for_mixed_nvfp4(monkeypatch): + monkeypatch.delenv("VLLM_SM70_FP8_QPN8", raising=False) + monkeypatch.delenv("VLLM_SM70_FP8_QPN8_LIBRARY", raising=False) + envs.disable_envs_cache() + try: + assert not envs.VLLM_SM70_FP8_QPN8 + assert envs.VLLM_SM70_FP8_QPN8_LIBRARY is None + assert not _sm70_fp8_qpn8_enabled(False) + assert _sm70_fp8_qpn8_enabled(True) + monkeypatch.setenv("VLLM_SM70_FP8_QPN8", "0") + envs.disable_envs_cache() + assert not envs.VLLM_SM70_FP8_QPN8 + assert not _sm70_fp8_qpn8_enabled(True) + monkeypatch.setenv("VLLM_SM70_FP8_QPN8", "1") + envs.disable_envs_cache() + assert envs.VLLM_SM70_FP8_QPN8 + assert _sm70_fp8_qpn8_enabled(False) + finally: + envs.disable_envs_cache() + + def test_fp8_prefill_exact_dense_shape_gate_is_narrow(): layer = SimpleNamespace( tp_size=4, + weight_block_size=[128, 128], prefix="model.language_model.layers.1.mlp.gate_up_proj", weight=SimpleNamespace(shape=(5120, 8704)), ) @@ -39,13 +405,89 @@ def test_fp8_prefill_exact_dense_shape_gate_is_narrow(): layer.tp_size = 4 layer.prefix = "model.language_model.layers.1.self_attn.qkv_proj" layer.weight = SimpleNamespace(shape=(5120, 3584)) - assert not _is_sm70_fp8_prefill_exact_dense_layer(layer) + assert _is_sm70_fp8_prefill_exact_dense_layer(layer) layer.prefix = "model.language_model.layers.1.linear_attn.in_proj_qkvz" layer.weight = SimpleNamespace(shape=(5120, 4096)) - assert not _is_sm70_fp8_prefill_exact_dense_layer(layer) + assert _is_sm70_fp8_prefill_exact_dense_layer(layer) layer.prefix = "model.language_model.layers.1.mlp.gate_up_proj" layer.weight = SimpleNamespace(shape=(5120, 8192)) assert not _is_sm70_fp8_prefill_exact_dense_layer(layer) + layer.weight_block_size = [64, 128] + layer.weight = SimpleNamespace(shape=(5120, 8704)) + assert not _is_sm70_fp8_prefill_exact_dense_layer(layer) + + +def test_fp8_exact_8k_prefill_gate_uses_operator_contract_only(): + layer = SimpleNamespace( + tp_size=4, + weight_block_size=[128, 128], + prefix="engine.layers.1.linear_attn.in_proj_qkvz", + weight=SimpleNamespace(shape=(5120, 4096)), + ) + + assert _is_sm70_fp8_exact_8k_prefill_layer(layer) + layer.prefix = "engine.layers.3.self_attn.qkv_proj" + layer.weight = SimpleNamespace(shape=(5120, 3584)) + assert _is_sm70_fp8_exact_8k_prefill_layer(layer) + + layer.tp_size = 2 + assert not _is_sm70_fp8_exact_8k_prefill_layer(layer) + layer.tp_size = 4 + layer.weight_block_size = [64, 128] + assert not _is_sm70_fp8_exact_8k_prefill_layer(layer) + layer.weight_block_size = [128, 128] + layer.weight = SimpleNamespace(shape=(5120, 4096)) + assert not _is_sm70_fp8_exact_8k_prefill_layer(layer) + + +def test_fp8_qpn8_shape_gate_uses_checkpoint_native_layout(): + layer = SimpleNamespace( + tp_size=4, + weight_block_size=[128, 128], + prefix="model.language_model.layers.1.mlp.gate_up_proj", + input_size_per_partition=5120, + output_size_per_partition=8704, + weight=SimpleNamespace(shape=(8704, 5120)), + ) + + assert _is_sm70_fp8_qpn8_layer(layer) + + layer.tp_size = 2 + assert not _is_sm70_fp8_qpn8_layer(layer) + layer.tp_size = 4 + layer.prefix = "model.language_model.layers.1.linear_attn.in_proj_qkvz" + layer.output_size_per_partition = 4096 + layer.weight = SimpleNamespace(shape=(4096, 5120)) + assert _is_sm70_fp8_qpn8_layer(layer) + layer.prefix = "model.language_model.layers.3.self_attn.qkv_proj" + layer.output_size_per_partition = 3584 + layer.weight = SimpleNamespace(shape=(3584, 5120)) + assert _is_sm70_fp8_qpn8_layer(layer) + layer.output_size_per_partition = 4096 + assert not _is_sm70_fp8_qpn8_layer(layer) + layer.output_size_per_partition = 3584 + layer.weight_block_size = [64, 128] + assert not _is_sm70_fp8_qpn8_layer(layer) + + +def test_fp8_qpn8_runtime_gate_uses_engine_contract_only(monkeypatch): + vllm_config = SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_seqs=8), + speculative_config=None, + ) + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.fp8.get_current_vllm_config", + lambda: vllm_config, + ) + + assert _is_sm70_fp8_qpn8_runtime_contract() + vllm_config.scheduler_config.max_num_seqs = 16 + assert not _is_sm70_fp8_qpn8_runtime_contract() + vllm_config.scheduler_config.max_num_seqs = 8 + vllm_config.speculative_config = object() + assert not _is_sm70_fp8_qpn8_runtime_contract() + monkeypatch.setenv("VLLM_SM70_FP8_QPN8", "1") + assert _is_sm70_fp8_qpn8_runtime_contract() def test_fp8_prefill_exact_dense_workspace_is_bounded(): @@ -75,6 +517,51 @@ def fake_empty(shape, *, dtype, device): _sm70_fp8_prefill_dense_workspaces.clear() +def test_fp8_prefill_visible_dense_mm_is_long_prefill_only(monkeypatch): + workspace = torch.empty(24, dtype=torch.float16) + weight = torch.empty((4, 6), dtype=torch.uint8) + scales = torch.empty((1, 6), dtype=torch.float16) + calls = [] + + def fake_dequantize(out, qweight, factors, group_size): + calls.append((qweight, factors, group_size)) + out.fill_(1) + + monkeypatch.setenv("VLLM_SM70_FP8_PREFILL_VISIBLE_DENSE_MM", "1") + monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 0) + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.fp8.sm70_ops.fp8_sm70_dequantize_out", + fake_dequantize, + ) + envs.disable_envs_cache() + _sm70_fp8_prefill_dense_workspaces[(0, torch.float16)] = workspace + try: + long_prefill = _sm70_fp8_prefill_visible_dense_mm( + torch.ones((_SM70_FP8_PREFILL_DENSE_MIN_M, 4), dtype=torch.float16), + weight, + scales, + workspace.data_ptr(), + gated_silu=False, + min_prefill_m=_SM70_FP8_PREFILL_DENSE_MIN_M, + ) + tail = _sm70_fp8_prefill_visible_dense_mm( + torch.ones((1, 4), dtype=torch.float16), + weight, + scales, + workspace.data_ptr(), + gated_silu=False, + min_prefill_m=_SM70_FP8_PREFILL_DENSE_MIN_M, + ) + finally: + _sm70_fp8_prefill_dense_workspaces.clear() + envs.disable_envs_cache() + + assert long_prefill is not None + assert long_prefill.shape == (_SM70_FP8_PREFILL_DENSE_MIN_M, 6) + assert tail is None + assert calls == [(weight, scales, 128)] + + def test_fp8_prefill_dispatch_reaches_runtime_op_for_small_and_large_m(monkeypatch): calls = [] @@ -125,3 +612,245 @@ def fake_dispatch( False, ), ] + + +def test_fp8_prefill_prescaled_scales_only_reach_exact_8k_route(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_prefill_prescaled_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_prefill_prescaled_scales=prescaled_scales, + sm70_fp8_k_ld=4, + sm70_fp8_q_ld=6, + ) + method = SimpleNamespace() + + monkeypatch.setenv("VLLM_SM70_FP8_PREFILL_FAST_SELECTOR", "1") + monkeypatch.setenv("VLLM_SM70_FP8_PREFILL_PRESCALED", "1") + envs.disable_envs_cache() + try: + Fp8LinearMethod.apply(method, layer, torch.empty((1, 4), dtype=torch.float16)) + Fp8LinearMethod.apply( + method, + layer, + torch.empty((_SM70_FP8_EXACT_8K_PREFILL_M, 4), dtype=torch.float16), + ) + monkeypatch.setenv("VLLM_SM70_FP8_PREFILL_PRESCALED", "0") + envs.disable_envs_cache() + Fp8LinearMethod.apply( + method, + layer, + torch.empty((_SM70_FP8_EXACT_8K_PREFILL_M, 4), dtype=torch.float16), + ) + finally: + envs.disable_envs_cache() + + assert [(route, m) for route, m, _ in calls] == [ + ("default", 1), + ("prescaled", _SM70_FP8_EXACT_8K_PREFILL_M), + ("default", _SM70_FP8_EXACT_8K_PREFILL_M), + ] + assert calls[0][2] is normal_scales + assert calls[1][2] is prescaled_scales + 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 = [] + + def fake_dispatch( + out, workspace, input, codes, scales, split_k, nacc, prefetch, gated_silu + ): + calls.append( + ("dispatch", input.shape[0], workspace, split_k, nacc, prefetch, gated_silu) + ) + out.zero_() + + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.fp8.sm70_ops." + "fp8_qpn8_dispatch_sm70_out", + fake_dispatch, + ) + layer = SimpleNamespace( + sm70_fp8_turbomind=True, + sm70_fp8_qpn8=True, + output_size_per_partition=6, + weight=torch.empty((4, 6), dtype=torch.uint8), + weight_scale_inv=torch.empty((1, 1), dtype=torch.float16), + sm70_fp8_qpn8_split_k=16, + sm70_fp8_qpn8_nacc=1, + sm70_fp8_qpn8_prefetch=False, + sm70_fp8_prefill_exact_dense_workspace_ptr=42, + ) + method = SimpleNamespace() + + for m in (1, 9): + output = Fp8LinearMethod.apply( + method, layer, torch.empty((m, 4), dtype=torch.float16) + ) + assert output.shape == (m, 6) + + assert calls == [ + ("dispatch", 1, 42, 16, 1, False, False), + ("dispatch", 9, 42, 16, 1, False, False), + ] + + +def test_fp8_qpn8_fused_gate_dispatches_without_intermediate(monkeypatch): + calls = [] + + def fake_dispatch( + out, workspace, input, codes, scales, split_k, nacc, prefetch, gated_silu + ): + calls.append( + ("dispatch", input.shape[0], workspace, split_k, nacc, prefetch, gated_silu) + ) + out.zero_() + + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.fp8.sm70_ops." + "fp8_qpn8_dispatch_sm70_out", + fake_dispatch, + ) + layer = SimpleNamespace( + sm70_fp8_qpn8=True, + sm70_fp8_gated_silu=True, + output_size_per_partition=12, + weight=torch.empty((4, 12), dtype=torch.uint8), + weight_scale_inv=torch.empty((1, 1), dtype=torch.float16), + sm70_fp8_qpn8_gated_split_k=8, + sm70_fp8_qpn8_gated_nacc=2, + sm70_fp8_qpn8_gated_prefetch=True, + sm70_fp8_prefill_exact_dense_workspace_ptr=42, + ) + method = SimpleNamespace() + + for m in (8, 16): + output = Fp8LinearMethod.apply_fused_silu_and_mul( + method, layer, torch.empty((m, 4), dtype=torch.float16) + ) + assert output is not None + assert output.shape == (m, 6) + + assert calls == [ + ("dispatch", 8, 42, 8, 2, True, True), + ("dispatch", 16, 42, 8, 2, True, True), + ] diff --git a/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py b/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py new file mode 100644 index 0000000000..065b447cc9 --- /dev/null +++ b/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py @@ -0,0 +1,192 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Narrow SM70 admission and shape gates for ModelOpt mixed NVFP4.""" + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from vllm import envs +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 ( + ModelOptFp8Config, + ModelOptMixedPrecisionConfig, + ModelOptNvFp4Config, +) +from vllm.model_executor.layers.quantization.nvfp4_sm70_moe import ( + ModelOptNvFp4SM70MoEMethod, + _prepare_compact_slot_groups, + _validate_weight_layout, + validate_nvfp4_sm70_moe_contract, +) + + +def _mixed_config() -> ModelOptMixedPrecisionConfig: + fp8 = ModelOptFp8Config("FP8", True, None, []) + nvfp4 = ModelOptNvFp4Config( + quant_method="NVFP4", + is_checkpoint_nvfp4_serialized=True, + ) + w4a16 = ModelOptNvFp4Config( + quant_method="W4A16_NVFP4", + is_checkpoint_nvfp4_serialized=True, + ) + return ModelOptMixedPrecisionConfig( + kv_cache_quant_method=None, + exclude_modules=[], + quantized_layers={"model.layers.0.mlp.experts": {"quant_algo": "W4A16_NVFP4"}}, + fp8_config=fp8, + nvfp4_config=nvfp4, + w4a16_nvfp4_config=w4a16, + ) + + +def _moe_contract(**overrides): + values = { + "num_experts": 256, + "experts_per_token": 8, + "hidden_dim": 2048, + "intermediate_size_per_partition": 128, + "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), + patch.object(sm70_tm, "use_turbomind", side_effect=[True, True]), + ): + assert ModelOptMixedPrecisionConfig.get_min_capability() == 70 + + with ( + patch.object(sm70_tm, "is_exact_sm70_cuda_platform", return_value=True), + patch.object(sm70_tm, "use_turbomind", side_effect=[True, False]), + ): + assert ModelOptMixedPrecisionConfig.get_min_capability() == 89 + + with patch.object(sm70_tm, "is_exact_sm70_cuda_platform", return_value=False): + assert ModelOptMixedPrecisionConfig.get_min_capability() == 89 + + +def test_nvfp4_grouped_prefill_defaults_on_and_can_be_disabled(monkeypatch): + monkeypatch.delenv("VLLM_SM70_NVFP4_MOE_GROUPED_PREFILL", raising=False) + assert envs.VLLM_SM70_NVFP4_MOE_GROUPED_PREFILL + + monkeypatch.setenv("VLLM_SM70_NVFP4_MOE_GROUPED_PREFILL", "0") + assert not envs.VLLM_SM70_NVFP4_MOE_GROUPED_PREFILL + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("num_experts", 128), + ("experts_per_token", 6), + ("hidden_dim", 4096), + ("intermediate_size_per_partition", 96), + ("tp_size", 2), + ], +) +def test_nvfp4_moe_contract_rejects_non_qwen36_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_rejects_shape_consistent_unvalidated_tp8(): + with pytest.raises(NotImplementedError, match="tensor parallel"): + validate_nvfp4_sm70_moe_contract( + _moe_contract(tp_size=8, intermediate_size_per_partition=64) + ) + + +def _meta_layer() -> SimpleNamespace: + experts, hidden, intermediate = 256, 2048, 128 + return SimpleNamespace( + local_num_experts=experts, + moe_config=_moe_contract(), + w13_weight=torch.empty( + experts, 2 * intermediate, hidden // 2, dtype=torch.uint8, device="meta" + ), + w13_weight_scale=torch.empty( + experts, + 2 * intermediate, + hidden // 16, + dtype=torch.float8_e4m3fn, + device="meta", + ), + w13_weight_scale_2=torch.empty(experts, 2, device="meta"), + w2_weight=torch.empty( + experts, hidden, intermediate // 2, dtype=torch.uint8, device="meta" + ), + w2_weight_scale=torch.empty( + experts, + hidden, + intermediate // 16, + dtype=torch.float8_e4m3fn, + device="meta", + ), + w2_weight_scale_2=torch.empty(experts, device="meta"), + ) + + +def test_nvfp4_moe_weight_layout_is_explicit(): + layer = _meta_layer() + _validate_weight_layout(layer) + + layer.w2_weight = torch.empty(256, 2048, 32, dtype=torch.uint8, device="meta") + with pytest.raises(ValueError, match="w2_weight"): + _validate_weight_layout(layer) + + +def test_nvfp4_sm70_moe_owns_routing_without_generic_modular_wrapper(): + method = ModelOptNvFp4SM70MoEMethod( + quant_config=ModelOptNvFp4Config( + quant_method="W4A16_NVFP4", + is_checkpoint_nvfp4_serialized=True, + ), + moe_config=_moe_contract(), + ) + + assert method.maybe_make_prepare_finalize() is None + + +@pytest.mark.skipif( + 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)) +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 + ) % 256 + compact_offsets = torch.empty(total_slots + 1, dtype=torch.int32, device="cuda") + active_expert_ids = torch.empty(total_slots, dtype=torch.int32, device="cuda") + + _prepare_compact_slot_groups(sorted_expert_ids, compact_offsets, active_expert_ids) + + assert torch.equal( + compact_offsets.cpu(), + torch.arange(total_slots + 1, dtype=torch.int32, device="cpu"), + ) + assert torch.equal(active_expert_ids.cpu(), sorted_expert_ids.cpu()) + + +def test_mixed_w4a16_moe_requires_turbomind_on_sm70(): + config = _mixed_config() + + class FakeRoutedExperts: + moe_config = _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=False), + pytest.raises(NotImplementedError, match="TurboMind"), + ): + config.get_quant_method(FakeRoutedExperts(), "model.layers.0.mlp.experts") diff --git a/tests/quantization/test_sm70_modelopt_nvfp4.py b/tests/quantization/test_sm70_modelopt_nvfp4.py new file mode 100644 index 0000000000..98a7db2be3 --- /dev/null +++ b/tests/quantization/test_sm70_modelopt_nvfp4.py @@ -0,0 +1,255 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Exact-SM70 ModelOpt NVFP4 admission: capability, provenance, routing. + +These tests exist so PR 114/166 cannot recur: +- the class-wide gate is not blindly lowered to 70 +- scale convention is ModelOpt provenance (amax/2688), never max(|scale|) < 1 +- TurboMind prepare is exact SM70 only +""" + +from unittest.mock import MagicMock, patch + +import torch +from torch.nn.parameter import Parameter + +from vllm.config import KernelConfig, VllmConfig, set_current_vllm_config +from vllm.model_executor.layers.quantization import sm70_turbomind as sm70_tm +from vllm.model_executor.layers.quantization.modelopt import ( + ModelOptNvFp4Config, + ModelOptNvFp4LinearMethod, + ModelOptNvFp4W4A16LinearMethod, +) + + +def _modelopt_dequant_oracle( + packed_nibble: float, block_scale: float, weight_global_scale: float +) -> float: + """Independent ModelOpt NVFP4 dequant: nibble * block * (amax/2688).""" + return packed_nibble * block_scale * weight_global_scale + + +def _ct_dequant_oracle( + packed_nibble: float, block_scale: float, disk_global: float +) -> float: + """Independent CT NVFP4 dequant: nibble * block / disk_global.""" + return packed_nibble * block_scale / disk_global + + +def _rejected_magnitude_heuristic_dequant( + packed_nibble: float, block_scale: float, disk_global: float +) -> float: + """PR 166 heuristic: max(|block|) < 1 => drop the global factor.""" + global_scale = 1.0 if abs(block_scale) < 1 else 1.0 / disk_global + return packed_nibble * block_scale * global_scale + + +def test_modelopt_nvfp4_min_capability_requires_proven_backend(monkeypatch): + monkeypatch.delenv("VLLM_USE_NVFP4_CT_EMULATIONS", raising=False) + monkeypatch.delenv("VLLM_NVFP4_GEMM_BACKEND", raising=False) + monkeypatch.delenv("VLLM_SM70_QUANT_BACKEND", raising=False) + + with set_current_vllm_config( + VllmConfig(kernel_config=KernelConfig(linear_backend="auto")) + ): + assert ModelOptNvFp4Config.get_min_capability() == 70 + + monkeypatch.setenv("VLLM_SM70_NVFP4_TURBOMIND", "0") + with set_current_vllm_config( + VllmConfig(kernel_config=KernelConfig(linear_backend="auto")) + ): + assert ModelOptNvFp4Config.get_min_capability() == 75 + + monkeypatch.delenv("VLLM_SM70_NVFP4_TURBOMIND", raising=False) + monkeypatch.setenv("VLLM_SM70_QUANT_BACKEND", "marlin") + with set_current_vllm_config( + VllmConfig(kernel_config=KernelConfig(linear_backend="auto")) + ): + assert ModelOptNvFp4Config.get_min_capability() == 70 + + monkeypatch.setenv("VLLM_SM70_QUANT_BACKEND", "turbomind") + with set_current_vllm_config( + VllmConfig(kernel_config=KernelConfig(linear_backend="auto")) + ): + assert ModelOptNvFp4Config.get_min_capability() == 70 + + monkeypatch.setenv("VLLM_SM70_NVFP4_TURBOMIND", "0") + monkeypatch.setenv("VLLM_SM70_QUANT_BACKEND", "auto") + with set_current_vllm_config( + VllmConfig(kernel_config=KernelConfig(linear_backend="emulation")) + ): + assert ModelOptNvFp4Config.get_min_capability() == 70 + + +def test_modelopt_scale_provenance_is_not_the_pr166_heuristic(): + # yangzhuxinyzx 100x counterexample: valid CT encoding, FP8 block 0.5, + # disk global 100, nibble 2. Correct CT dequant is 0.01. The rejected + # magnitude rule rewrites it to 1.0. + nibble, block, disk_global = 2.0, 0.5, 100.0 + assert _ct_dequant_oracle(nibble, block, disk_global) == 0.01 + assert _rejected_magnitude_heuristic_dequant(nibble, block, disk_global) == 1.0 + + # ModelOpt provenance is explicit: global is already amax/2688. + # Same numeric tensors must keep the on-disk global (no rewrite). + modelopt_global = 100.0 + assert _modelopt_dequant_oracle(nibble, block, modelopt_global) == 100.0 + assert abs(block) < 1 + assert _modelopt_dequant_oracle(nibble, block, modelopt_global) != 1.0 + + +def _w4a16_layer(global_scale: float, block_scale: float) -> torch.nn.Module: + layer = torch.nn.Module() + layer.weight = Parameter( + torch.zeros((16, 8), dtype=torch.uint8), + requires_grad=False, + ) + layer.weight_scale = Parameter( + torch.full((16, 1), block_scale, dtype=torch.float8_e4m3fn), + requires_grad=False, + ) + layer.weight_scale_2 = Parameter( + torch.tensor([global_scale], dtype=torch.float32), + requires_grad=False, + ) + layer.input_scale = Parameter( + torch.tensor([1.0], dtype=torch.float32), + requires_grad=False, + ) + return layer + + +def test_w4a16_keeps_modelopt_global_scale_without_reciprocation_or_heuristic(): + layer = _w4a16_layer(global_scale=100.0, block_scale=0.5) + fake_kernel = MagicMock() + + with patch( + "vllm.model_executor.layers.quantization.modelopt.MarlinNvFp4LinearKernel", + return_value=fake_kernel, + ): + method = ModelOptNvFp4W4A16LinearMethod( + ModelOptNvFp4Config( + quant_method="W4A16_NVFP4", + is_checkpoint_nvfp4_serialized=True, + ) + ) + + with ( + patch.object(sm70_tm, "should_prepare_turbomind", return_value=False), + patch.object( + sm70_tm, + "prepare_nvfp4_linear", + side_effect=AssertionError("prepare must not run"), + ), + ): + method.process_weights_after_loading(layer) + + # Provenance: keep on-disk ModelOpt global. Do not invert. Do not + # force 1.0 because max(|weight_scale|) < 1. + assert torch.equal(layer.weight_global_scale.cpu(), torch.tensor(100.0)) + assert not hasattr(layer, "weight_scale_2") + fake_kernel.process_weights_after_loading.assert_called_once_with(layer) + + +def test_w4a16_turbomind_prepare_only_on_exact_sm70(): + layer = _w4a16_layer(global_scale=0.25, block_scale=1.0) + fake_kernel = MagicMock() + + with patch( + "vllm.model_executor.layers.quantization.modelopt.MarlinNvFp4LinearKernel", + return_value=fake_kernel, + ): + method = ModelOptNvFp4W4A16LinearMethod( + ModelOptNvFp4Config( + quant_method="W4A16_NVFP4", + is_checkpoint_nvfp4_serialized=True, + ) + ) + + def fake_prepare(prepared: torch.nn.Module) -> None: + setattr(prepared, sm70_tm.STATE_ATTR, object()) + + with ( + patch.object(sm70_tm, "should_prepare_turbomind", return_value=True), + patch.object(sm70_tm, "prepare_nvfp4_linear", side_effect=fake_prepare), + ): + method.process_weights_after_loading(layer) + + assert sm70_tm.has_prepared_linear(layer) + assert layer.weight.numel() == 0 + assert layer.weight_scale.numel() == 0 + fake_kernel.process_weights_after_loading.assert_not_called() + + +def test_w4a4_turbomind_prepare_skipped_off_sm70(): + layer = torch.nn.Module() + layer.input_scale = Parameter( + torch.ones(1, dtype=torch.float32), + requires_grad=False, + ) + layer.weight_scale_2 = Parameter( + torch.tensor([0.25], dtype=torch.float32), + requires_grad=False, + ) + layer.weight = Parameter( + torch.zeros((16, 8), dtype=torch.uint8), + requires_grad=False, + ) + layer.weight_scale = Parameter( + torch.ones((16, 1), dtype=torch.float8_e4m3fn), + requires_grad=False, + ) + fake_kernel = MagicMock() + + with patch( + "vllm.model_executor.layers.quantization.modelopt.init_nvfp4_linear_kernel", + return_value=fake_kernel, + ): + method = ModelOptNvFp4LinearMethod( + ModelOptNvFp4Config( + quant_method="NVFP4", + is_checkpoint_nvfp4_serialized=True, + ) + ) + + with ( + patch.object(sm70_tm, "should_prepare_turbomind", return_value=False), + patch.object( + sm70_tm, + "prepare_nvfp4_linear", + side_effect=AssertionError("prepare must not run off SM70"), + ), + ): + method.process_weights_after_loading(layer) + + fake_kernel.process_weights_after_loading.assert_called_once_with(layer) + assert not sm70_tm.has_prepared_linear(layer) + + +def test_apply_uses_prepared_turbomind_path(): + layer = torch.nn.Module() + fake_kernel = MagicMock() + fake_out = torch.zeros((1, 4), dtype=torch.float16) + + with patch( + "vllm.model_executor.layers.quantization.modelopt.MarlinNvFp4LinearKernel", + return_value=fake_kernel, + ): + method = ModelOptNvFp4W4A16LinearMethod( + ModelOptNvFp4Config( + quant_method="W4A16_NVFP4", + is_checkpoint_nvfp4_serialized=True, + ) + ) + + x = torch.ones((1, 16), dtype=torch.float16) + with ( + patch.object(sm70_tm, "has_prepared_linear", return_value=True), + patch.object( + sm70_tm, "apply_prepared_linear", return_value=fake_out + ) as apply_tm, + ): + out = method.apply(layer, x) + + assert out is fake_out + apply_tm.assert_called_once() + fake_kernel.apply_weights.assert_not_called() diff --git a/tests/quantization/test_sm70_modelopt_nvfp4_kernel_oracle.py b/tests/quantization/test_sm70_modelopt_nvfp4_kernel_oracle.py new file mode 100644 index 0000000000..ddffa37532 --- /dev/null +++ b/tests/quantization/test_sm70_modelopt_nvfp4_kernel_oracle.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Exact-SM70 ModelOpt NVFP4: real kernel path vs an independent fp32 reference. + +Runs only on a Volta (7,0) CUDA device. Builds a synthetic linear in the exact +on-disk ModelOpt layout (``weight`` uint8 E2M1 pairs, ``weight_scale`` fp8 e4m3 +per 16-group, ``weight_scale_2`` fp32 = amax/(6*448)), loads it through the real +ModelOpt W4A4 / W4A16 quant methods (so the SM70 TurboMind prepare/apply route is +exercised end to end), and compares against a pure-torch dequant + fp32 GEMM that +shares no code with vLLM's scale handling. A convention error (reciprocated or +dropped global, PR 166 class) shows up as an output-norm ratio of ~2688x or +~1/2688x and fails loudly; kernel arithmetic error shows up as low cosine. +""" + +import os + +import pytest +import torch + +from vllm.config import VllmConfig, set_current_vllm_config +from vllm.distributed import ( + ensure_model_parallel_initialized, + init_distributed_environment, +) +from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.layers.quantization import sm70_turbomind as sm70_tm +from vllm.model_executor.layers.quantization.modelopt import ModelOptNvFp4Config + +E2M1 = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32) + + +def _is_exact_sm70() -> bool: + return torch.cuda.is_available() and torch.cuda.get_device_capability(0) == (7, 0) + + +pytestmark = pytest.mark.skipif( + not _is_exact_sm70(), reason="requires an exact SM70 (Volta) CUDA device" +) + + +def _decode_e2m1(packed: torch.Tensor) -> torch.Tensor: + lo = packed & 0x0F + hi = (packed >> 4) & 0x0F + nib = torch.stack([lo, hi], dim=-1).reshape(packed.shape[0], -1).long() + return E2M1[nib & 0x7] * torch.where((nib & 0x8) != 0, -1.0, 1.0) + + +def _make_modelopt_layer(n: int, k: int, seed: int): + """Synthetic ModelOpt NVFP4 tensors: nibbles, fp8 block scales, amax global.""" + g = torch.Generator().manual_seed(seed) + packed = torch.randint(0, 256, (n, k // 2), generator=g, dtype=torch.uint8) + amax = 0.35 + weight_scale_2 = torch.tensor(amax / (6.0 * 448.0), dtype=torch.float32) + # block scales in fp8: values so that block*global lands in a realistic weight range + block = torch.rand(n, k // 16, generator=g) * 400.0 + 20.0 + weight_scale = block.to(torch.float8_e4m3fn) + w_ref = ( + _decode_e2m1(packed) + * weight_scale.to(torch.float32).repeat_interleave(16, dim=1) + * weight_scale_2 + ) + return packed, weight_scale, weight_scale_2, w_ref + + +def _build_layer(quant_method: str, n: int, k: int, tensors, device): + packed, weight_scale, weight_scale_2, _ = tensors + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29581") + with set_current_vllm_config(VllmConfig()): + try: + init_distributed_environment( + world_size=1, + rank=0, + local_rank=0, + distributed_init_method="env://", + backend="gloo", + ) + ensure_model_parallel_initialized(1, 1) + except Exception: + pass # already initialised by an earlier test in this process + qc = ModelOptNvFp4Config( + is_checkpoint_nvfp4_serialized=True, + kv_cache_quant_algo=None, + exclude_modules=[], + group_size=16, + quant_method=quant_method, + ) + lin = ReplicatedLinear( + k, + n, + bias=False, + quant_config=qc, + prefix="oracle", + params_dtype=torch.float16, + ).to(device) + for pname, param in lin.named_parameters(): + src = { + "weight": packed, + "weight_scale": weight_scale, + "weight_scale_2": weight_scale_2.reshape(1), + "input_scale": torch.tensor([1.0]), + }.get(pname) + if src is None: + continue + loader = getattr(param, "weight_loader", None) + if loader is not None: + loader( + param, src if param.dtype == torch.uint8 else src.to(param.dtype) + ) + else: + param.data.copy_(src.reshape(param.shape)) + lin.quant_method.process_weights_after_loading(lin) + return lin + + +@pytest.mark.parametrize("quant_method", ["NVFP4", "W4A16_NVFP4"]) +@pytest.mark.parametrize("n,k", [(512, 1024), (1024, 512)]) +def test_sm70_modelopt_nvfp4_kernel_matches_independent_reference( + monkeypatch, quant_method, n, k +): + monkeypatch.setenv("VLLM_SM70_NVFP4_TURBOMIND", "1") + monkeypatch.delenv("VLLM_SM70_QUANT_BACKEND", raising=False) + assert ModelOptNvFp4Config.get_min_capability() == 70 + + device = torch.device("cuda:0") + tensors = _make_modelopt_layer(n, k, seed=n + k) + w_ref = tensors[3] + lin = _build_layer(quant_method, n, k, tensors, device) + assert sm70_tm.has_prepared_linear(lin), "expected the exact-SM70 TurboMind route" + + torch.manual_seed(0) + x = (torch.randn(64, k) * 0.5).to(torch.float16) + y_ref = x.float() @ w_ref.T # independent fp32 reference + y_floor = (x @ w_ref.to(torch.float16).T).float() # fp16 accumulation floor + with torch.no_grad(): + y_dut, _ = lin(x.to(device)) + y_dut = y_dut.float().cpu() + + ratio = float(y_dut.norm() / y_ref.norm()) + cos = float( + torch.nn.functional.cosine_similarity(y_dut.flatten(), y_ref.flatten(), dim=0) + ) + rel = float((y_dut - y_ref).norm() / y_ref.norm()) + rel_floor = float((y_floor - y_ref).norm() / y_ref.norm()) + # A reciprocated/dropped global would put ratio at ~2688 or ~1/2688. + assert abs(ratio - 1.0) < 0.02, f"output scale ratio {ratio} (convention error?)" + assert cos > 0.999, f"cosine {cos} (kernel arithmetic error?)" + assert rel < max(5.0 * rel_floor, 5e-3), f"rel err {rel} vs fp16 floor {rel_floor}" diff --git a/tests/quantization/test_sm70_mxfp4_moe.py b/tests/quantization/test_sm70_mxfp4_moe.py index ebb2e1785d..9e3e9fa18f 100644 --- a/tests/quantization/test_sm70_mxfp4_moe.py +++ b/tests/quantization/test_sm70_mxfp4_moe.py @@ -6,6 +6,7 @@ import pytest import torch +from vllm import _sm70_ops as sm70_ops from vllm import envs from vllm.model_executor.layers.fused_moe import ( FusedMoEConfig, @@ -21,6 +22,7 @@ from vllm.model_executor.layers.quantization.mxfp4_sm70_moe import ( Mxfp4SM70MoEMethod, _compact_mxfp4_active_experts, + _select_mxfp4_direct_order_offsets, _select_mxfp4_stage_dispatch, validate_mxfp4_sm70_moe_contract, validate_mxfp4_sm70_moe_weight_layout, @@ -209,6 +211,21 @@ def test_mxfp4_sm70_b1_dispatch_selects_six_runtime_experts(monkeypatch): assert count == 6 +def test_mxfp4_sm70_direct_order_ignores_compacted_offsets(): + buffers = { + # 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 + ), + "slot_expert_offsets": torch.arange(7, dtype=torch.int32), + } + + offsets = _select_mxfp4_direct_order_offsets(buffers) + + assert offsets is buffers["slot_expert_offsets"] + torch.testing.assert_close(offsets, torch.arange(7, dtype=torch.int32)) + + def test_mxfp4_sm70_b1_dispatch_rejects_incompatible_permute_fastpath( monkeypatch, ): @@ -308,6 +325,9 @@ def test_mxfp4_sm70_m8_grouped_dispatch_keeps_one_row_slots(monkeypatch): } monkeypatch.setattr(mxfp4_moe, "_mxfp4_active_expert_max_tokens", lambda: 8) monkeypatch.setattr(mxfp4_moe, "_mxfp4_grouped_m8_enabled", lambda: True) + monkeypatch.setattr( + mxfp4_moe, "_mxfp4_grouped_m8_expert_rows_enabled", lambda: False + ) offsets, expert_ids, count = _select_mxfp4_stage_dispatch( buffers, @@ -321,6 +341,72 @@ def test_mxfp4_sm70_m8_grouped_dispatch_keeps_one_row_slots(monkeypatch): assert count == 48 +def test_mxfp4_sm70_m8_expert_grouped_dispatch_keeps_real_segments(monkeypatch): + buffers = { + "compact_expert_offsets": torch.arange(49, dtype=torch.int32) * 2, + "slot_expert_offsets": torch.arange(49, dtype=torch.int32), + "permuted_experts_id": torch.arange(48, dtype=torch.int32).flip(0), + "active_expert_ids": torch.arange(48, dtype=torch.int32), + "expert_offsets": torch.arange(257, dtype=torch.int32), + "dense_expert_ids": torch.arange(256, dtype=torch.int32), + } + monkeypatch.setattr(mxfp4_moe, "_mxfp4_active_expert_max_tokens", lambda: 8) + monkeypatch.setattr(mxfp4_moe, "_mxfp4_grouped_m8_enabled", lambda: True) + monkeypatch.setattr( + mxfp4_moe, "_mxfp4_grouped_m8_expert_rows_enabled", lambda: True + ) + + offsets, expert_ids, count = _select_mxfp4_stage_dispatch( + buffers, + num_tokens=8, + num_experts=256, + fully_replicated_experts=True, + ) + + assert offsets is buffers["compact_expert_offsets"] + assert expert_ids is buffers["active_expert_ids"] + assert count == 48 + + +@pytest.mark.parametrize("expert_rows", [False, True]) +def test_mxfp4_sm70_m5_grouped_dispatch(expert_rows, monkeypatch): + buffers = { + "compact_expert_offsets": torch.arange(31, dtype=torch.int32) * 2, + "slot_expert_offsets": torch.arange(31, dtype=torch.int32), + "permuted_experts_id": torch.arange(30, dtype=torch.int32).flip(0), + "active_expert_ids": torch.arange(30, dtype=torch.int32), + "expert_offsets": torch.arange(257, dtype=torch.int32), + "dense_expert_ids": torch.arange(256, dtype=torch.int32), + } + monkeypatch.setattr(mxfp4_moe, "_mxfp4_active_expert_max_tokens", lambda: 8) + monkeypatch.setattr(mxfp4_moe, "_mxfp4_grouped_m8_enabled", lambda: False) + monkeypatch.setattr(mxfp4_moe, "_mxfp4_grouped_verifier_enabled", lambda: True) + monkeypatch.setattr( + mxfp4_moe, + "_mxfp4_grouped_m8_expert_rows_enabled", + lambda: expert_rows, + ) + + offsets, expert_ids, count = _select_mxfp4_stage_dispatch( + buffers, + num_tokens=5, + num_experts=256, + fully_replicated_experts=True, + ) + + expected_offsets = ( + buffers["compact_expert_offsets"] + if expert_rows + else buffers["slot_expert_offsets"] + ) + expected_ids = ( + buffers["active_expert_ids"] if expert_rows else buffers["permuted_experts_id"] + ) + assert offsets is expected_offsets + assert expert_ids is expected_ids + assert count == 30 + + @pytest.mark.skipif( not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0), reason="requires NVIDIA V100/SM70", @@ -363,6 +449,117 @@ def test_mxfp4_sm70_active_expert_compaction_replays_dynamic_routes(): ) +@pytest.mark.parametrize(("k", "n"), [(4096, 512), (4096, 1024), (512, 4096)]) +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0), + reason="requires NVIDIA V100/SM70", +) +def test_mxfp4_sm70_tp4_compact_shapes_match_per_expert( + monkeypatch: pytest.MonkeyPatch, + k: int, + n: int, +): + num_experts = 6 + group_size = 32 + torch.manual_seed(20260825) + packed = torch.randint(0, 16, (k, n), dtype=torch.uint8, device="cuda") + scales = torch.full((k // group_size, n), 127, dtype=torch.uint8, device="cuda") + weight, scale, meta = sm70_ops.mxfp4_sm70_prepare(packed, scales, group_size) + weights = weight.unsqueeze(0).repeat(num_experts, *([1] * weight.ndim)) + expert_scales = scale.unsqueeze(0).repeat(num_experts, *([1] * scale.ndim)) + ptrs_w, ptrs_s = torch.ops._C.awq_moe_build_strided_ptrs( + weights, + expert_scales, + int(meta[0].item()), + int(meta[1].item()), + num_experts, + ) + x = torch.randn((1, k), dtype=torch.float16, device="cuda") * 0.01 + grouped_input = x.expand(num_experts, -1).contiguous() + offsets = torch.arange(num_experts + 1, dtype=torch.int32, device="cuda") + expert_ids = torch.arange(num_experts, dtype=torch.int32, device="cuda") + reference = torch.empty((num_experts, n), dtype=torch.float16, device="cuda") + actual = torch.empty_like(reference) + + monkeypatch.setenv("VLLM_SM70_MXFP4_MOE_COMPACT_GROUPED_DECODE", "0") + torch.ops._C.mxfp4_moe_dense_stage_sm70_out( + reference, + grouped_input, + offsets, + expert_ids, + ptrs_w, + ptrs_s, + num_experts, + k, + n, + group_size, + ) + monkeypatch.setenv("VLLM_SM70_MXFP4_MOE_COMPACT_GROUPED_DECODE", "1") + torch.ops._C.mxfp4_moe_dense_stage_sm70_out( + actual, + grouped_input, + offsets, + expert_ids, + ptrs_w, + ptrs_s, + num_experts, + k, + n, + group_size, + ) + torch.accelerator.synchronize() + torch.testing.assert_close(actual, reference, rtol=0.0, atol=0.0) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + torch.ops._C.mxfp4_moe_dense_stage_sm70_out( + actual, + grouped_input, + offsets, + expert_ids, + ptrs_w, + ptrs_s, + num_experts, + k, + n, + group_size, + ) + graph.replay() + torch.accelerator.synchronize() + torch.testing.assert_close(actual, reference, rtol=0.0, atol=0.0) + + if k == 4096: + direct = torch.empty_like(actual) + legacy_direct = torch.empty_like(actual) + direct_input = torch.empty_like(grouped_input) + direct_offsets = torch.empty_like(offsets) + inverse_indices = torch.empty(num_experts, dtype=torch.int32, device="cuda") + direct_expert_ids = torch.empty_like(expert_ids) + for broadcast, output in ((False, legacy_direct), (True, direct)): + monkeypatch.setenv( + "VLLM_SM70_MXFP4_MOE_BROADCAST_INPUT_DECODE", + "1" if broadcast else "0", + ) + torch.ops._C.mxfp4_moe_single_token_prepare_w13_sm70_out( + output, + direct_input, + x, + expert_ids.view(1, num_experts), + ptrs_w, + ptrs_s, + direct_offsets, + inverse_indices, + direct_expert_ids, + k, + n, + group_size, + k, + ) + torch.accelerator.synchronize() + torch.testing.assert_close(legacy_direct, reference, rtol=0.0, atol=0.0) + torch.testing.assert_close(direct, reference, rtol=0.0, atol=0.0) + + def test_mxfp4_sm70_post_load_reads_bias_from_method_config(monkeypatch): for op_name in ( "mxfp4_sm70_prepare", diff --git a/tests/quantization/test_sm70_nvfp4_qpn2.py b/tests/quantization/test_sm70_nvfp4_qpn2.py new file mode 100644 index 0000000000..84f8174158 --- /dev/null +++ b/tests/quantization/test_sm70_nvfp4_qpn2.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import torch +from torch.nn.parameter import Parameter + +import vllm.envs as envs +from vllm.model_executor.layers.quantization import sm70_turbomind as sm70_tm +from vllm.model_executor.layers.quantization.compressed_tensors.schemes import ( + compressed_tensors_w4a4_nvfp4 as nvfp4_scheme, +) +from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_w4a4_nvfp4 import ( # noqa: E501 + CompressedTensorsW4A4Fp4, +) + + +def test_nvfp4_qpn2_is_default_off_with_explicit_on(monkeypatch): + monkeypatch.delenv("VLLM_SM70_NVFP4_QPN2", raising=False) + envs.disable_envs_cache() + try: + assert not envs.VLLM_SM70_NVFP4_QPN2 + monkeypatch.setenv("VLLM_SM70_NVFP4_QPN2", "1") + envs.disable_envs_cache() + assert envs.VLLM_SM70_NVFP4_QPN2 + finally: + envs.disable_envs_cache() + + +def test_nvfp4_qpn2_shape_gate_is_exact_tp4(): + layer = SimpleNamespace( + tp_size=4, + prefix="model.language_model.layers.0.mlp.gate_up_proj", + input_size_per_partition=5120, + output_size_per_partition=8704, + weight=SimpleNamespace(shape=(8704, 2560)), + ) + assert nvfp4_scheme._is_qpn2_layer(layer) + + layer.tp_size = 2 + assert not nvfp4_scheme._is_qpn2_layer(layer) + layer.tp_size = 4 + layer.prefix = "model.language_model.layers.0.self_attn.qkv_proj" + assert not nvfp4_scheme._is_qpn2_layer(layer) + layer.prefix = "model.language_model.layers.0.mlp.down_proj" + layer.input_size_per_partition = 4352 + layer.output_size_per_partition = 5120 + layer.weight = SimpleNamespace(shape=(5120, 2176)) + assert nvfp4_scheme._is_qpn2_layer(layer) + layer.input_size_per_partition = 4096 + assert not nvfp4_scheme._is_qpn2_layer(layer) + + +def _make_small_layer() -> torch.nn.Module: + layer = torch.nn.Module() + layer.prefix = "model.language_model.layers.0.mlp.gate_up_proj" + layer.tp_size = 4 + layer.input_size_per_partition = 64 + layer.output_size_per_partition = 64 + layer.weight_packed = Parameter( + torch.zeros((64, 32), dtype=torch.uint8), requires_grad=False + ) + layer.weight_scale = Parameter( + torch.ones((64, 4), dtype=torch.float8_e4m3fn), requires_grad=False + ) + layer.weight_global_scale = Parameter( + torch.tensor([2.0, 1.0], dtype=torch.float32), requires_grad=False + ) + layer.input_global_scale = Parameter( + torch.tensor([4.0, 4.0], dtype=torch.float32), requires_grad=False + ) + return layer + + +def test_nvfp4_qpn2_prepare_and_dispatch_contract(monkeypatch): + monkeypatch.setenv("VLLM_SM70_NVFP4_QPN2", "1") + envs.disable_envs_cache() + layer = _make_small_layer() + calls = [] + + monkeypatch.setattr(nvfp4_scheme.sm70_tm, "use_turbomind", lambda enabled: True) + scheme = CompressedTensorsW4A4Fp4() + monkeypatch.setattr( + nvfp4_scheme.sm70_tm, + "should_prepare_turbomind", + lambda tensor, enabled: enabled, + ) + monkeypatch.setattr(nvfp4_scheme, "_is_qpn2_layer", lambda layer: True) + monkeypatch.setattr(nvfp4_scheme, "_missing_qpn2_ops", lambda: []) + 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( + nvfp4_scheme.sm70_ops, + "nvfp4_qpn2_prepare_sm70", + lambda weight, scales: ( + torch.empty_like(weight), + torch.empty(scales.shape, dtype=torch.uint8), + ), + ) + + def fake_prepare(prepared_layer, *, interleave_gated_silu=False): + assert not interleave_gated_silu + state = sm70_tm.SM70TurboMindLinearState( + weight=torch.empty((1,), dtype=torch.int32), + scales=torch.empty((1,), dtype=torch.float16), + group_size=16, + k_ld=64, + q_ld=64, + output_size=64, + op_kind="nvfp4", + ) + setattr(prepared_layer, sm70_tm.STATE_ATTR, state) + + monkeypatch.setattr(nvfp4_scheme.sm70_tm, "prepare_nvfp4_linear", fake_prepare) + + def fake_dispatch(*args): + calls.append(args) + args[0].fill_(3) + + monkeypatch.setattr( + nvfp4_scheme.sm70_ops, "nvfp4_qpn2_dispatch_sm70_out", fake_dispatch + ) + + 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.weight.numel() == 0 + assert layer.weight_scale.numel() == 0 + + x = torch.ones((8, 64), dtype=torch.float16) + raw = scheme.apply_weights(layer, x) + fused = scheme.apply_fused_silu_and_mul(layer, x) + assert raw.shape == (8, 64) + assert fused is not None and fused.shape == (8, 32) + assert torch.equal(raw, torch.full_like(raw, 3)) + assert torch.equal(fused, torch.full_like(fused, 3)) + assert calls[0][-1] is False + assert calls[1][-1] is True + finally: + envs.disable_envs_cache() diff --git a/tests/quantization/test_sm70_turbomind_adapter.py b/tests/quantization/test_sm70_turbomind_adapter.py index af022393fe..c04d961caf 100644 --- a/tests/quantization/test_sm70_turbomind_adapter.py +++ b/tests/quantization/test_sm70_turbomind_adapter.py @@ -91,3 +91,224 @@ def test_symmetric_int4_zero_points_are_eight(): assert zeros.dtype == torch.float16 assert zeros.tolist() == [[8, 8, 8], [8, 8, 8]] + + +def _make_nvfp4_state(tm, *, gated_silu: bool): + return tm.SM70TurboMindLinearState( + weight=torch.empty((3, 1), dtype=torch.uint8), + scales=torch.empty((1, 4), dtype=torch.float16), + group_size=16, + k_ld=3, + q_ld=4, + output_size=4, + op_kind="nvfp4", + gated_silu=gated_silu, + ) + + +def _make_nvfp4_qpn4_state(tm, *, gated_silu: bool): + output_size = 4 + return tm.SM70TurboMindLinearState( + weight=torch.empty((3, output_size // 2), dtype=torch.uint8), + scales=torch.empty( + (1, output_size), + dtype=torch.uint8 if gated_silu else torch.float16, + ), + group_size=16, + k_ld=0, + q_ld=0, + output_size=output_size, + op_kind="nvfp4_qpn4", + gated_silu=gated_silu, + dense_weight_ptr=1234, + global_scale=0.25 if gated_silu else 0.0, + use_scale_code=gated_silu, + ) + + +def test_nvfp4_gated_layout_is_deinterleaved_for_unfused_apply(monkeypatch): + tm = _load_adapter() + layer = torch.nn.Module() + setattr(layer, tm.STATE_ATTR, _make_nvfp4_state(tm, gated_silu=True)) + + from vllm import _sm70_ops as sm70_ops + + calls = [] + + def fake_gemm( + out, + x, + weight, + scales, + group_size, + k_ld, + q_ld, + gated_silu=False, + ): + del x, weight, scales, group_size, k_ld, q_ld + calls.append(gated_silu) + out.copy_(torch.tensor([[1.0, 10.0, 2.0, 20.0]], dtype=out.dtype)) + + monkeypatch.setattr(sm70_ops, "nvfp4_gemm_sm70_out", fake_gemm) + + output = tm.apply_prepared_linear( + layer, + torch.ones((1, 3), dtype=torch.float16), + bias=None, + ) + + assert calls == [False] + assert output.tolist() == [[1.0, 2.0, 10.0, 20.0]] + + +def test_nvfp4_gated_layout_uses_fused_silu_epilogue(monkeypatch): + tm = _load_adapter() + layer = torch.nn.Module() + setattr(layer, tm.STATE_ATTR, _make_nvfp4_state(tm, gated_silu=True)) + + from vllm import _sm70_ops as sm70_ops + + calls = [] + + def fake_gemm( + out, + x, + weight, + scales, + group_size, + k_ld, + q_ld, + gated_silu=False, + ): + del weight, scales + calls.append( + { + "out_shape": tuple(out.shape), + "x_shape": tuple(x.shape), + "group_size": group_size, + "k_ld": k_ld, + "q_ld": q_ld, + "gated_silu": gated_silu, + } + ) + out.fill_(3.0) + + monkeypatch.setattr(sm70_ops, "nvfp4_gemm_sm70_out", fake_gemm) + + output = tm.apply_prepared_fused_silu_and_mul( + layer, + torch.ones((2, 1, 3), dtype=torch.float16), + ) + + assert output is not None + assert output.shape == (2, 1, 2) + assert output.tolist() == [[[3.0, 3.0]], [[3.0, 3.0]]] + assert calls == [ + { + "out_shape": (2, 2), + "x_shape": (2, 3), + "group_size": 16, + "k_ld": 3, + "q_ld": 4, + "gated_silu": True, + } + ] + + +def test_nvfp4_fused_silu_rejects_non_gated_state(): + tm = _load_adapter() + layer = torch.nn.Module() + setattr(layer, tm.STATE_ATTR, _make_nvfp4_state(tm, gated_silu=False)) + + output = tm.apply_prepared_fused_silu_and_mul( + layer, + torch.ones((1, 3), dtype=torch.float16), + ) + + assert output is None + + +def test_nvfp4_qpn4_regular_apply_uses_opaque_dynamic_m_dispatch(monkeypatch): + tm = _load_adapter() + layer = torch.nn.Module() + setattr(layer, tm.STATE_ATTR, _make_nvfp4_qpn4_state(tm, gated_silu=False)) + + from vllm import _sm70_ops as sm70_ops + + calls = [] + + def fake_dispatch( + out, + dense_weight_ptr, + x, + weight, + scales, + global_scale, + use_scale_code, + gated_silu, + ): + del weight, scales + calls.append( + ( + tuple(out.shape), + dense_weight_ptr, + tuple(x.shape), + global_scale, + use_scale_code, + gated_silu, + ) + ) + out.fill_(2.0) + + monkeypatch.setattr(sm70_ops, "nvfp4_qpn4_dispatch_sm70_out", fake_dispatch) + output = tm.apply_prepared_linear( + layer, + torch.ones((2, 3), dtype=torch.float16), + bias=None, + ) + + assert output.tolist() == [[2.0] * 4, [2.0] * 4] + assert calls == [((2, 4), 1234, (2, 3), 0.0, False, False)] + + +def test_nvfp4_qpn4_fused_gate_uses_scale_code_dispatch(monkeypatch): + tm = _load_adapter() + layer = torch.nn.Module() + setattr(layer, tm.STATE_ATTR, _make_nvfp4_qpn4_state(tm, gated_silu=True)) + + from vllm import _sm70_ops as sm70_ops + + calls = [] + + def fake_dispatch( + out, + dense_weight_ptr, + x, + weight, + scales, + global_scale, + use_scale_code, + gated_silu, + ): + del weight, scales + calls.append( + ( + tuple(out.shape), + dense_weight_ptr, + tuple(x.shape), + global_scale, + use_scale_code, + gated_silu, + ) + ) + out.fill_(4.0) + + monkeypatch.setattr(sm70_ops, "nvfp4_qpn4_dispatch_sm70_out", fake_dispatch) + output = tm.apply_prepared_fused_silu_and_mul( + layer, + torch.ones((1, 3), dtype=torch.float16), + ) + + assert output is not None + assert output.tolist() == [[4.0, 4.0]] + assert calls == [((1, 2), 1234, (1, 3), 0.25, True, True)] diff --git a/tests/quantization/test_sm70_warmup.py b/tests/quantization/test_sm70_warmup.py index 781c01d7e7..e6e12b922a 100644 --- a/tests/quantization/test_sm70_warmup.py +++ b/tests/quantization/test_sm70_warmup.py @@ -6,6 +6,8 @@ import torch from torch import nn +import vllm.envs as envs +from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod from vllm.model_executor.warmup import awq_sm70_warmup as warmup @@ -26,6 +28,14 @@ def _grouped_fp8_layer() -> nn.Module: return layer +def test_fp8_grouped_bmm_decode_defaults_on_with_rollback(monkeypatch): + monkeypatch.delenv("VLLM_SM70_FP8_GROUPED_BMM_DECODE", raising=False) + assert envs.VLLM_SM70_FP8_GROUPED_BMM_DECODE + + monkeypatch.setenv("VLLM_SM70_FP8_GROUPED_BMM_DECODE", "0") + assert not envs.VLLM_SM70_FP8_GROUPED_BMM_DECODE + + def test_fp8_warmup_discovers_grouped_bmm_by_per_group_shape(): layer = _grouped_fp8_layer() model = nn.Sequential(layer) @@ -35,6 +45,14 @@ def test_fp8_warmup_discovers_grouped_bmm_by_per_group_shape(): assert discovered == [(layer, False)] +def test_fp8_warmup_skips_static_qpn8_dispatch(): + layer = _grouped_fp8_layer() + layer.sm70_fp8_qpn8 = True + model = nn.Sequential(layer) + + assert list(warmup._iter_unique_fp8_dense_layers(model)) == [] + + def test_fp8_warmup_matches_grouped_bmm_runtime_slice(monkeypatch): layer = _grouped_fp8_layer() calls = [] @@ -68,6 +86,344 @@ def record_call(out, x, weight, scales, group_size, k_ld, q_ld, gated_silu): assert all(not call.gated_silu for call in calls) +def test_fp8_warmup_includes_one_launch_grouped_decode(monkeypatch): + layer = _grouped_fp8_layer() + layer.sm70_fp8_bmm_groups = 2 + layer.sm70_fp8_bmm_grouped_decode = True + layer.weight = nn.Parameter( + torch.empty((2, 128, 64), dtype=torch.uint8), requires_grad=False + ) + layer.weight_scale_inv = nn.Parameter( + torch.empty((2, 1, 64), dtype=torch.float16), requires_grad=False + ) + layer.sm70_fp8_bmm_grouped_offsets = torch.arange(3, dtype=torch.int32) + layer.sm70_fp8_bmm_grouped_ptrs_w = torch.empty(2, dtype=torch.int64) + layer.sm70_fp8_bmm_grouped_ptrs_s = torch.empty(2, dtype=torch.int64) + dense_calls = [] + grouped_calls = [] + monkeypatch.setattr(torch.ops._C, "fp8_gemm_sm70_out_meta", object(), raising=False) + monkeypatch.setattr( + warmup.sm70_ops, + "fp8_gemm_sm70_out", + lambda *args: dense_calls.append(args), + ) + monkeypatch.setattr( + warmup.sm70_ops, + "fp8_moe_gemm_sm70_per_expert_dispatch_out", + lambda *args: grouped_calls.append(args), + ) + + count = warmup._warmup_fp8_dense_layers([(layer, False)], [1, 4]) + + assert count == 3 + assert len(dense_calls) == 2 + assert len(grouped_calls) == 1 + call = grouped_calls[0] + assert tuple(call[0].shape) == (2, 64) + assert tuple(call[1].shape) == (2, 128) + assert call[2] is layer.sm70_fp8_bmm_grouped_offsets + assert call[3] is layer.sm70_fp8_bmm_grouped_ptrs_w + assert call[4] is layer.sm70_fp8_bmm_grouped_ptrs_s + assert call[5:] == (2, 128, 64, 128, False) + + +def test_fp8_grouped_bmm_decode_uses_one_dispatch(monkeypatch): + layer = _grouped_fp8_layer() + layer.sm70_fp8_bmm_groups = 2 + layer.sm70_fp8_bmm_grouped_decode = True + layer.sm70_fp8_bmm_grouped_offsets = torch.arange(3, dtype=torch.int32) + layer.sm70_fp8_bmm_grouped_ptrs_w = torch.empty(2, dtype=torch.int64) + layer.sm70_fp8_bmm_grouped_ptrs_s = torch.empty(2, dtype=torch.int64) + layer.output_size_per_partition = 128 + calls = [] + + def grouped(out, x, *args): + calls.append((out, x, args)) + out.copy_(torch.arange(out.numel(), dtype=out.dtype).reshape_as(out)) + + def fail_dense(*args): + raise AssertionError("dense fallback used") + + monkeypatch.setattr( + warmup.sm70_ops, + "fp8_moe_gemm_sm70_per_expert_dispatch_out", + grouped, + ) + monkeypatch.setattr( + warmup.sm70_ops, + "fp8_gemm_sm70_out", + fail_dense, + ) + x = torch.empty((1, 2, 128), dtype=torch.float16) + + out = object.__new__(Fp8LinearMethod).apply(layer, x) + + assert tuple(out.shape) == (1, 2, 64) + assert len(calls) == 1 + grouped_out, grouped_x, args = calls[0] + assert tuple(grouped_out.shape) == (2, 64) + assert tuple(grouped_x.shape) == (2, 128) + assert args[0] is layer.sm70_fp8_bmm_grouped_offsets + assert args[1] is layer.sm70_fp8_bmm_grouped_ptrs_w + assert args[2] is layer.sm70_fp8_bmm_grouped_ptrs_s + assert args[3:] == (2, 128, 64, 128, False) + torch.testing.assert_close( + out, + torch.arange(128, dtype=torch.float16).reshape(1, 2, 64), + ) + + +def test_fp8_grouped_bmm_decode_retains_multirow_fallback(monkeypatch): + layer = _grouped_fp8_layer() + layer.sm70_fp8_bmm_groups = 2 + layer.sm70_fp8_bmm_grouped_decode = True + layer.sm70_fp8_bmm_output_size = 64 + dense_calls = [] + + def dense(out, x, *args): + dense_calls.append((tuple(out.shape), tuple(x.shape), args)) + out.zero_() + + def fail_grouped(*args): + raise AssertionError("grouped decode must remain batch-one only") + + monkeypatch.setattr(warmup.sm70_ops, "fp8_gemm_sm70_out", dense) + monkeypatch.setattr( + warmup.sm70_ops, + "fp8_moe_gemm_sm70_per_expert_dispatch_out", + fail_grouped, + ) + x = torch.empty((2, 2, 128), dtype=torch.float16) + + out = object.__new__(Fp8LinearMethod).apply(layer, x) + + assert tuple(out.shape) == (2, 2, 64) + assert [call[:2] for call in dense_calls] == [ + ((2, 64), (2, 128)), + ((2, 64), (2, 128)), + ] + + +def test_fp8_warmup_supports_modelopt_turbomind_layout(monkeypatch): + layer = nn.Module() + layer.sm70_modelopt_fp8_turbomind = True + layer.sm70_modelopt_fp8_k_ld = 128 + layer.sm70_modelopt_fp8_q_ld = 64 + layer.output_size_per_partition = 64 + layer.weight = nn.Parameter( + torch.empty((128, 64), dtype=torch.uint8), requires_grad=False + ) + layer.weight_scale = nn.Parameter( + torch.empty((1, 64), dtype=torch.float16), requires_grad=False + ) + model = nn.Sequential(layer) + calls = [] + monkeypatch.setattr(torch.ops._C, "fp8_gemm_sm70_out_meta", object(), raising=False) + monkeypatch.setattr( + warmup.sm70_ops, + "fp8_gemm_sm70_out", + lambda out, x, weight, scales, group_size, k_ld, q_ld, gated_silu: calls.append( + (out.shape, x.shape, scales, group_size, k_ld, q_ld) + ), + ) + + discovered = list(warmup._iter_unique_fp8_dense_layers(model)) + count = warmup._warmup_fp8_dense_layers(discovered, [1, 4]) + + assert discovered == [(layer, False)] + assert count == 2 + assert [tuple(call[0]) for call in calls] == [(1, 64), (4, 64)] + assert all(call[2] is layer.weight_scale for call in calls) + assert all(call[3:] == (128, 128, 64) for call in calls) + + +def _nvfp4_moe_layer() -> nn.Module: + layer = nn.Module() + layer.sm70_nvfp4_moe = True + layer.moe_config = SimpleNamespace(experts_per_token=8) + layer.swiglu_limit = None + layer.sm70_nvfp4_num_experts = 256 + layer.sm70_nvfp4_w13_k_dim = 2048 + layer.sm70_nvfp4_w13_n_dim = 256 + layer.sm70_nvfp4_w2_k_dim = 128 + layer.sm70_nvfp4_w2_n_dim = 2048 + layer.sm70_nvfp4_group_size = 16 + layer.sm70_nvfp4_graph_safe_max_tokens = 18 + layer.sm70_nvfp4_compact_grouped_max_tokens = 10 + layer.w13_tm_weight = nn.Parameter( + torch.empty((1, 1), dtype=torch.uint8), requires_grad=False + ) + layer.w13_strided_ptrs_w = torch.empty(1, dtype=torch.uint8) + layer.w13_strided_ptrs_s = torch.empty(1, dtype=torch.uint8) + layer.w2_strided_ptrs_w = torch.empty(1, dtype=torch.uint8) + layer.w2_strided_ptrs_s = torch.empty(1, dtype=torch.uint8) + layer._nvfp4_sm70_dense_expert_ids = torch.arange(256, dtype=torch.int32) + return layer + + +def test_nvfp4_moe_warmup_discovers_and_uses_compact_decode_shapes(monkeypatch): + layer = _nvfp4_moe_layer() + model = nn.Sequential(layer) + calls = [] + monkeypatch.setattr( + torch.ops._C, "nvfp4_moe_dense_stage_sm70_out", object(), raising=False + ) + + def record_call( + out, + x, + offsets, + expert_ids, + ptrs_w, + ptrs_s, + num_experts, + k, + n, + group_size, + ): + calls.append( + ( + tuple(out.shape), + tuple(x.shape), + offsets.tolist(), + expert_ids.tolist(), + num_experts, + k, + n, + group_size, + ) + ) + + monkeypatch.setattr(warmup.sm70_ops, "nvfp4_moe_dense_stage_sm70_out", record_call) + monkeypatch.setattr( + warmup, + "_prepare_compact_slot_groups", + lambda sorted_ids, offsets, active_ids: ( + offsets.copy_(torch.arange(offsets.numel(), dtype=torch.int32)), + active_ids.copy_(sorted_ids), + ), + ) + monkeypatch.setattr( + torch.ops._C, + "silu_and_mul", + lambda out, gate_up: out.zero_(), + raising=False, + ) + + discovered = list(warmup._iter_unique_nvfp4_moe_layers(model)) + count = warmup._warmup_nvfp4_moe_decode_layers(discovered, [1, 2]) + + assert discovered == [layer] + assert count == 4 + assert [call[4] for call in calls] == [8, 8, 16, 16] + assert calls[0][2] == list(range(9)) + assert calls[0][3] == list(range(8)) + assert calls[2][2] == list(range(17)) + assert calls[2][3] == list(range(16)) + + +def test_nvfp4_moe_warmup_includes_opted_in_cuda_graph_shapes(monkeypatch): + monkeypatch.setattr(warmup.envs, "VLLM_SM70_NVFP4_MOE_TUNE_MAX_TOKENS", 640) + worker = SimpleNamespace( + vllm_config=SimpleNamespace( + compilation_config=SimpleNamespace( + cudagraph_capture_sizes=[1, 2, 4, 5, 8, 9, 18, 20, 40, 60, 80, 81] + ) + ) + ) + + assert warmup._get_nvfp4_moe_token_counts( + worker, [_nvfp4_moe_layer()], [1, 2, 4, 8] + ) == [*range(1, 11), 18, 20, 40, 60, 80] + + +def test_nvfp4_moe_warmup_uses_slot_compact_through_b10(monkeypatch): + layer = _nvfp4_moe_layer() + calls = [] + monkeypatch.setattr( + torch.ops._C, "nvfp4_moe_dense_stage_sm70_out", object(), raising=False + ) + monkeypatch.setattr( + warmup.sm70_ops, + "nvfp4_moe_dense_stage_sm70_out", + lambda *args: calls.append(args), + ) + monkeypatch.setattr( + warmup, + "_prepare_compact_slot_groups", + lambda sorted_ids, offsets, active_ids: ( + offsets.copy_(torch.arange(offsets.numel(), dtype=torch.int32)), + active_ids.copy_(sorted_ids), + ), + ) + monkeypatch.setattr( + torch.ops._C, + "silu_and_mul", + lambda out, gate_up: out.zero_(), + raising=False, + ) + + assert warmup._warmup_nvfp4_moe_decode_layers([layer], [9, 10]) == 4 + assert [call[6] for call in calls] == [72, 72, 80, 80] + assert all(call[2].numel() == call[6] + 1 for call in calls) + assert all(call[3].numel() == call[6] for call in calls) + assert calls[0][2].tolist() == list(range(73)) + assert calls[2][2].tolist() == list(range(81)) + + +def test_nvfp4_moe_warmup_uses_full_expert_groups_above_compact_b10(monkeypatch): + layer = _nvfp4_moe_layer() + calls = [] + monkeypatch.setattr( + torch.ops._C, "nvfp4_moe_dense_stage_sm70_out", object(), raising=False + ) + monkeypatch.setattr( + warmup.sm70_ops, + "nvfp4_moe_dense_stage_sm70_out", + lambda *args: calls.append(args), + ) + monkeypatch.setattr( + torch.ops._C, + "silu_and_mul", + lambda out, gate_up: out.zero_(), + raising=False, + ) + + assert warmup._warmup_nvfp4_moe_decode_layers([layer], [11]) == 2 + assert [call[6] for call in calls] == [256, 256] + assert all(call[2].numel() == 257 for call in calls) + assert all(call[3].numel() == 256 for call in calls) + assert calls[0][2][:89].tolist() == list(range(89)) + assert calls[0][2][89:].tolist() == [88] * (257 - 89) + + +def test_nvfp4_moe_warmup_supports_mtp4_c16_width(monkeypatch): + layer = _nvfp4_moe_layer() + calls = [] + monkeypatch.setattr( + torch.ops._C, "nvfp4_moe_dense_stage_sm70_out", object(), raising=False + ) + monkeypatch.setattr( + warmup.sm70_ops, + "nvfp4_moe_dense_stage_sm70_out", + lambda *args: calls.append(args), + ) + monkeypatch.setattr( + torch.ops._C, + "silu_and_mul", + lambda out, gate_up: out.zero_(), + raising=False, + ) + + assert warmup._warmup_nvfp4_moe_decode_layers([layer], [80]) == 2 + assert [tuple(call[0].shape) for call in calls] == [(640, 256), (640, 2048)] + assert [call[6] for call in calls] == [256, 256] + offsets = calls[0][2] + assert offsets.numel() == 257 + assert offsets[:4].tolist() == [0, 3, 6, 9] + assert offsets[-1].item() == 640 + + def test_fp8_coordinated_warmup_leader_broadcasts_rank0_lut(monkeypatch): import vllm.distributed.parallel_state as parallel_state diff --git a/tests/test_envs.py b/tests/test_envs.py index db50deb241..ffc79b49e4 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -121,6 +121,60 @@ def test_precompiled_install_flags_are_orthogonal() -> None: assert environment_variables["VLLM_USE_PRECOMPILED_RUST"]() is False +def test_sm70_concurrency_tuning_envs( + monkeypatch: pytest.MonkeyPatch, +) -> None: + names = ( + "VLLM_SM70_TP4_MTP_AR_BLOCK_TUNING", + "VLLM_SM70_TOPK_TOPP_8_WARPS", + "VLLM_SM70_MTP_MOE_TUNED_CONFIG", + ) + for name in names: + monkeypatch.delenv(name, raising=False) + assert environment_variables[name]() is False + monkeypatch.setenv(name, "1") + assert environment_variables[name]() is True + + name = "VLLM_SM70_TOPK_TOPP_B8_B16_8_WARPS" + monkeypatch.delenv(name, raising=False) + assert environment_variables[name]() is True + monkeypatch.setenv(name, "0") + assert environment_variables[name]() is False + + name = "VLLM_SM70_DSV4_QNORM_KV_FUSED_TP4" + monkeypatch.delenv(name, raising=False) + assert environment_variables[name]() is True + 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") + assert environment_variables["VLLM_SM70_NVFP4_MOE_TUNE_MAX_TOKENS"]() == 640 + + def test_flash_v100_g6_sawtooth_pipeline_envs( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -150,6 +204,20 @@ def test_flash_v100_g6_sawtooth_pipeline_envs( assert environment_variables[name]() == expected_int + 1 +def test_flash_v100_e4m3_page800_fastpath_envs( + monkeypatch: pytest.MonkeyPatch, +) -> None: + bool_defaults = { + "VLLM_FLASH_V100_E4M3_PAGE800_FASTPATH": True, + "VLLM_FLASH_V100_E4M3_PAGE800_FASTPATH_TRACE": False, + } + for name, expected_bool in bool_defaults.items(): + monkeypatch.delenv(name, raising=False) + assert environment_variables[name]() is expected_bool + monkeypatch.setenv(name, "0" if expected_bool else "1") + assert environment_variables[name]() is not expected_bool + + 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..9167493688 100644 --- a/tests/tool_parsers/test_qwen3coder_tool_parser.py +++ b/tests/tool_parsers/test_qwen3coder_tool_parser.py @@ -1384,7 +1384,7 @@ def test_streaming_multi_param_single_chunk(qwen3_tool_parser, qwen3_tokenizer): @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/test_deepseek_ocr_resolution_modes.py b/tests/transformers_utils/test_deepseek_ocr_resolution_modes.py new file mode 100644 index 0000000000..b7f8888ecd --- /dev/null +++ b/tests/transformers_utils/test_deepseek_ocr_resolution_modes.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek-OCR resolution modes: per-request ``image_mode`` selection, +crop parametrization, and — critically — the producer/counter consistency +contract (`tokenize_with_images` versus +`count_image_tokens`), whose violation is the "N multimodal tokens vs M +placeholders" crash class.""" + +import pytest +from transformers.processing_utils import ProcessorMixin + +from vllm.transformers_utils.processors.deepseek_ocr import ( + BASE_SIZE, + CROP_MODE, + IMAGE_SIZE, + RESOLUTION_MODES, + DeepseekOCRProcessor, + count_image_tokens_for, +) + +# --------------------------------------------------------------------------- +# Pure-function tests (no tokenizer). +# --------------------------------------------------------------------------- + + +def test_mode_table_matches_official_config(): + # Verbatim from the model repo's config.py (captured in the issue refs). + assert RESOLUTION_MODES == { + "tiny": {"base_size": 512, "image_size": 512, "crop_mode": False}, + "small": {"base_size": 640, "image_size": 640, "crop_mode": False}, + "base": {"base_size": 1024, "image_size": 1024, "crop_mode": False}, + "large": {"base_size": 1280, "image_size": 1280, "crop_mode": False}, + "gundam": {"base_size": 1024, "image_size": 640, "crop_mode": True}, + } + # The module defaults ARE the gundam mode (byte-compat anchor). + assert (BASE_SIZE, IMAGE_SIZE, CROP_MODE) == (1024, 640, True) + + +def _count(mode: str, width: int, height: int, **kw) -> int: + cfg = RESOLUTION_MODES[mode] + return count_image_tokens_for( + image_width=width, + image_height=height, + base_size=cfg["base_size"], + image_size=cfg["image_size"], + cropping=cfg["crop_mode"], + **kw, + ) + + +def test_no_crop_modes_are_size_independent(): + for mode in ("tiny", "small", "base", "large"): + assert _count(mode, 100, 100) == _count(mode, 4000, 3000) + + +def test_gundam_small_image_equals_no_crop_arithmetic(): + # <= image_size on both dims: never cropped, regardless of crop_mode. + assert _count("gundam", 640, 480) == count_image_tokens_for( + image_width=640, + image_height=480, + base_size=1024, + image_size=640, + cropping=False, + ) + + +def test_gundam_large_image_counts_tiles(): + assert _count("gundam", 1700, 2200) > _count("gundam", 100, 100) + + +def test_max_crops_bounds_the_tile_count(): + big = _count("gundam", 4000, 4000, max_crops=9) + small = _count("gundam", 4000, 4000, max_crops=2) + assert big > small + + +# --------------------------------------------------------------------------- +# Processor tests with a minimal tokenizer (CPU-only, no model files). +# --------------------------------------------------------------------------- + + +class _TokenizerStub: + bos_token_id = 1 + eos_token_id = 2 + pad_token_id = 0 + pad_token = None + padding_side = "right" + vocab = {"": 99} + + def add_special_tokens(self, tokens): + self.pad_token = tokens["pad_token"] + + def encode(self, text, add_special_tokens=False): + del add_special_tokens + return [10] if text else [] + + def decode(self, tokens, **kwargs): + del kwargs + return " ".join(map(str, tokens)) + + +@pytest.fixture +def tokenizer(monkeypatch): + monkeypatch.setattr(ProcessorMixin, "__init__", lambda *args, **kwargs: None) + + return _TokenizerStub() + + +def _images(sizes): + from PIL import Image + + return [Image.new("RGB", s, color="white") for s in sizes] + + +def test_default_construction_is_gundam_byte_compat(tokenizer): + default = DeepseekOCRProcessor(tokenizer=tokenizer) + explicit = DeepseekOCRProcessor(tokenizer=tokenizer, image_mode="gundam") + for attr in ("base_size", "image_size", "crop_mode", "min_crops", "max_crops"): + assert getattr(default, attr) == getattr(explicit, attr) + assert default.crop_mode is True and default.image_size == 640 + + +def test_invalid_mode_and_crops_raise(tokenizer): + with pytest.raises(ValueError, match="image_mode"): + DeepseekOCRProcessor(tokenizer=tokenizer, image_mode="giant") + with pytest.raises(ValueError, match="crop bounds"): + DeepseekOCRProcessor(tokenizer=tokenizer, min_crops=5, max_crops=2) + with pytest.raises(ValueError, match="crop bounds"): + DeepseekOCRProcessor(tokenizer=tokenizer, min_crops=2, max_crops=10) + + +def test_image_mode_authoritative_over_trio(tokenizer): + # The processing info always forwards the constant trio; a named mode + # must win over it. + proc = DeepseekOCRProcessor( + tokenizer=tokenizer, + image_size=640, + base_size=1024, + crop_mode=True, + image_mode="base", + ) + assert (proc.base_size, proc.image_size, proc.crop_mode) == (1024, 1024, False) + + +@pytest.mark.parametrize("mode", [None, "base", "large", "tiny", "gundam"]) +@pytest.mark.parametrize("size", [(100, 100), (1200, 800), (1000, 2200)]) +def test_producer_counter_consistency(tokenizer, mode, size): + kwargs = {} if mode is None else {"image_mode": mode} + proc = DeepseekOCRProcessor(tokenizer=tokenizer, **kwargs) + out = proc(prompt="\nFree OCR. ", images=_images([size])) + produced = int(out["num_image_tokens"][0]) + counted = proc.count_image_tokens(image_width=size[0], image_height=size[1]) + assert produced == counted + + +def test_multi_image_producer_counter_consistency(tokenizer): + proc = DeepseekOCRProcessor(tokenizer=tokenizer) + sizes = [(1200, 800), (1000, 2200)] + out = proc(prompt="ab", images=_images(sizes)) + for i, size in enumerate(sizes): + produced = int(out["num_image_tokens"][i]) + counted = proc.count_image_tokens(image_width=size[0], image_height=size[1]) + assert produced == counted diff --git a/tests/v1/attention/test_gdn_metadata_builder.py b/tests/v1/attention/test_gdn_metadata_builder.py index e43f1c0fa7..76626418e1 100644 --- a/tests/v1/attention/test_gdn_metadata_builder.py +++ b/tests/v1/attention/test_gdn_metadata_builder.py @@ -12,6 +12,7 @@ from pathlib import Path from types import SimpleNamespace +import numpy as np import pytest import torch @@ -36,14 +37,82 @@ build_gdn_spec_decode_state_contract, gdn_spec_metadata_tensors, get_registered_gdn_spec_metadata_tensors, + prepare_dflash2_gdn_group_metadata, ) from vllm.v1.attention.backends.utils import PAD_SLOT_ID from vllm.v1.kv_cache_interface import MambaSpec +from vllm.v1.worker.gpu.attn_utils import compute_common_gdn_attn_metadata +from vllm.v1.worker.gpu.model_states import mamba_hybrid +from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState BLOCK_SIZE = 16 DEVICE = torch.device("cpu") +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) + monkeypatch.delenv("VLLM_SM70_DUMP_GDN_GRAPH_DIR", raising=False) + + def fail_debug_work(*_args, **_kwargs): + raise AssertionError("disabled debug hook performed debug work") + + cuda_module = vars(qwen_gdn.torch)["cuda"] + monkeypatch.setattr(cuda_module, "is_current_stream_capturing", fail_debug_work) + monkeypatch.setattr(qwen_gdn, "_sm70_gdn_graph_buffer_copy", fail_debug_work) + qwen_gdn._sm70_dump_gdn_core_tensor( + "input_qkv", + "layer.0", + torch.zeros(1), + ) + + +def test_mamba_hybrid_passes_seq_lens_cpu_upper_bound(monkeypatch): + """Flash-V100 metadata must not lazily copy seq_lens back from the GPU.""" + + captured_kwargs = None + + def capture_build_attn_metadata(**kwargs): + nonlocal captured_kwargs + captured_kwargs = kwargs + return {} + + monkeypatch.setattr( + mamba_hybrid, + "build_attn_metadata", + capture_build_attn_metadata, + ) + state = object.__new__(MambaHybridModelState) + state.max_model_len = 4096 + seq_lens_cpu_upper_bound = torch.tensor([130], dtype=torch.int32) + input_batch = SimpleNamespace( + num_reqs=1, + num_tokens=8, + num_reqs_after_padding=1, + num_tokens_after_padding=8, + query_start_loc_np=np.array([0, 8], dtype=np.int32), + num_scheduled_tokens=np.array([8], dtype=np.int32), + is_prefilling_np=np.array([False]), + query_start_loc=torch.tensor([0, 8], dtype=torch.int32), + seq_lens=torch.tensor([123], dtype=torch.int32), + seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, + dcp_local_seq_lens=None, + prefix_anchor_lens=None, + ) + state.prepare_attn( + input_batch, + CUDAGraphMode.NONE, + block_tables=(), + slot_mappings=torch.empty((0, 8), dtype=torch.int32), + attn_groups=[], + kv_cache_config=SimpleNamespace(kv_cache_groups=[]), + for_capture=True, + ) + + assert captured_kwargs is not None + assert captured_kwargs["seq_lens_cpu_upper_bound"] is seq_lens_cpu_upper_bound + + @pytest.fixture def local_gdn_model(tmp_path: Path) -> str: model_dir = tmp_path / "gdn-test-model" @@ -172,6 +241,7 @@ def _create_gdn_builder( use_full_cuda_graph: bool = False, mamba_cache_mode: str = "none", max_cudagraph_capture_size: int = 4, + device: torch.device = DEVICE, ) -> GDNAttentionMetadataBuilder: """Create a GDNAttentionMetadataBuilder with minimal config.""" vllm_config = create_vllm_config(model_name=model_name, block_size=BLOCK_SIZE) @@ -181,6 +251,11 @@ def _create_gdn_builder( vllm_config.compilation_config.max_cudagraph_capture_size = ( max_cudagraph_capture_size ) + else: + # Keep the helper's default contract explicit. The repository-wide + # compilation default may select FULL graphs on CUDA platforms, which + # would pad otherwise eager metadata and make these tests route-dependent. + vllm_config.compilation_config.cudagraph_mode = CUDAGraphMode.NONE if num_speculative_tokens > 0: vllm_config.speculative_config = SpeculativeConfig( method="ngram", @@ -196,7 +271,7 @@ def _create_gdn_builder( kv_cache_spec=mamba_spec, layer_names=["layer.0"], vllm_config=vllm_config, - device=DEVICE, + device=device, ) @@ -212,20 +287,73 @@ def _build( builder: GDNAttentionMetadataBuilder, batch_spec: BatchSpec, num_decode_draft_tokens: list[int] | None = None, + use_common_metadata: bool = False, ) -> GDNAttentionMetadata: """Build GDN attention metadata, optionally with spec-decode kwargs.""" common = create_common_attn_metadata(batch_spec, BLOCK_SIZE, DEVICE) kwargs: dict = {} if num_decode_draft_tokens is not None: - kwargs["num_decode_draft_tokens_cpu"] = torch.tensor( + num_decode_draft_tokens_cpu = torch.tensor( num_decode_draft_tokens, dtype=torch.int32, device="cpu" ) - kwargs["num_accepted_tokens"] = torch.ones( + num_accepted_tokens = torch.ones( batch_spec.batch_size, dtype=torch.int32, device=DEVICE ) + kwargs["num_decode_draft_tokens_cpu"] = num_decode_draft_tokens_cpu + kwargs["num_accepted_tokens"] = num_accepted_tokens + if use_common_metadata: + kwargs["common_gdn_metadata"] = compute_common_gdn_attn_metadata( + num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, + query_start_loc=common.query_start_loc, + query_start_loc_cpu=common.query_start_loc_cpu, + num_spec_state_tokens=builder.num_spec_state_tokens, + legacy_mixed_decode_routing=( + qwen_gdn.envs.VLLM_SM70_MTP_LEGACY_GDN_MIXED_DECODE_ROUTING + ), + ) return builder.build(common_prefix_len=0, common_attn_metadata=common, **kwargs) +def _assert_gdn_metadata_equal( + actual: GDNAttentionMetadata, + expected: GDNAttentionMetadata, +) -> None: + scalar_fields = ( + "num_prefills", + "num_prefill_tokens", + "num_decodes", + "num_decode_tokens", + "num_spec_decodes", + "num_spec_decode_tokens", + "num_actual_tokens", + ) + for field in scalar_fields: + assert getattr(actual, field) == getattr(expected, field), field + + tensor_fields = ( + "has_initial_state", + "spec_query_start_loc", + "non_spec_query_start_loc", + "spec_state_indices_tensor", + "non_spec_state_indices_tensor", + "spec_sequence_masks", + "spec_token_indx", + "non_spec_token_indx", + "num_accepted_tokens", + "spec_state_slot_selectors", + "chunk_indices", + "chunk_offsets", + "batch_ptr", + "token_chunk_offset_ptr", + ) + for field in tensor_fields: + actual_tensor = getattr(actual, field) + expected_tensor = getattr(expected, field) + assert (actual_tensor is None) == (expected_tensor is None), field + if actual_tensor is not None: + assert torch.equal(actual_tensor, expected_tensor), field + + def _effective_spec_initial_state_slots( meta: GDNAttentionMetadata, ) -> list[int]: @@ -260,6 +388,51 @@ def test_gdn_build_classification( assert meta.num_spec_decodes == test_case.expected_num_spec_decodes +@pytest.mark.parametrize( + "test_case", + [ + case + for case in GDN_BUILD_TEST_CASES.values() + if case.num_decode_draft_tokens is not None + ], + ids=[ + name + for name, case in GDN_BUILD_TEST_CASES.items() + if case.num_decode_draft_tokens is not None + ], +) +def test_common_gdn_metadata_matches_per_group_builder( + test_case: GDNBuildTestCase, + local_gdn_model: str, +): + generic_builder = _create_gdn_builder( + local_gdn_model, test_case.num_speculative_tokens + ) + common_builder = _create_gdn_builder( + local_gdn_model, test_case.num_speculative_tokens + ) + batch = BatchSpec( + seq_lens=test_case.seq_lens, + query_lens=test_case.query_lens, + ) + + torch.manual_seed(0) + expected = _build( + generic_builder, + batch, + test_case.num_decode_draft_tokens, + ) + torch.manual_seed(0) + actual = _build( + common_builder, + batch, + test_case.num_decode_draft_tokens, + use_common_metadata=True, + ) + + _assert_gdn_metadata_equal(actual, expected) + + def test_mixed_decode_stays_decode_fastpath(local_gdn_model): """A non-spec query_len=1 row should stay on the decode fastpath even when another row is doing speculative verification.""" @@ -337,6 +510,419 @@ def test_full_cuda_graph_spec_replay_tail_uses_pad_slot(local_gdn_model): ] +def test_common_gdn_metadata_matches_full_graph_padding(local_gdn_model): + generic_builder = _create_gdn_builder( + local_gdn_model, + num_speculative_tokens=2, + use_full_cuda_graph=True, + ) + common_builder = _create_gdn_builder( + local_gdn_model, + num_speculative_tokens=2, + use_full_cuda_graph=True, + ) + batch = BatchSpec(seq_lens=[4097], query_lens=[3]) + common = create_common_attn_metadata(batch, BLOCK_SIZE, DEVICE).replace( + block_table_tensor=torch.tensor( + [[10, 11, 12, 13, 14]], + dtype=torch.int32, + device=DEVICE, + ), + num_actual_tokens=4, + ) + num_accepted_tokens = torch.tensor([3], dtype=torch.int32, device=DEVICE) + num_decode_draft_tokens_cpu = torch.tensor([2], dtype=torch.int32, device="cpu") + + expected = generic_builder.build( + common_prefix_len=0, + common_attn_metadata=common, + num_accepted_tokens=num_accepted_tokens, + num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, + ) + common_metadata = compute_common_gdn_attn_metadata( + num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, + query_start_loc=common.query_start_loc, + query_start_loc_cpu=common.query_start_loc_cpu, + num_spec_state_tokens=common_builder.num_spec_state_tokens, + legacy_mixed_decode_routing=( + qwen_gdn.envs.VLLM_SM70_MTP_LEGACY_GDN_MIXED_DECODE_ROUTING + ), + ) + assert common_metadata is not None + actual = common_builder.build( + common_prefix_len=0, + common_attn_metadata=common, + num_accepted_tokens=num_accepted_tokens, + num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, + common_gdn_metadata=common_metadata, + ) + + _assert_gdn_metadata_equal(actual, expected) + + +def test_prepared_dflash2_metadata_belongs_to_exact_builder(local_gdn_model): + builder = _create_gdn_builder( + local_gdn_model, + num_speculative_tokens=2, + use_full_cuda_graph=True, + ) + batch = BatchSpec(seq_lens=[4097], query_lens=[3]) + common = create_common_attn_metadata(batch, BLOCK_SIZE, DEVICE).replace( + block_table_tensor=torch.tensor( + [[10, 11, 12, 13]], dtype=torch.int32, device=DEVICE + ), + num_actual_tokens=4, + ) + num_accepted_tokens = torch.tensor([2], dtype=torch.int32, device=DEVICE) + num_decode_draft_tokens_cpu = torch.tensor([2], dtype=torch.int32, device="cpu") + common_metadata = compute_common_gdn_attn_metadata( + num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, + query_start_loc=common.query_start_loc, + query_start_loc_cpu=common.query_start_loc_cpu, + num_spec_state_tokens=builder.num_spec_state_tokens, + legacy_mixed_decode_routing=False, + ) + assert common_metadata is not None + prepared = builder.build( + common_prefix_len=0, + common_attn_metadata=common, + num_accepted_tokens=num_accepted_tokens, + num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, + common_gdn_metadata=common_metadata, + ) + + actual = builder.build( + common_prefix_len=0, + common_attn_metadata=common, + num_accepted_tokens=num_accepted_tokens, + num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, + common_gdn_metadata=common_metadata, + prepared_dflash2_metadata=prepared, + ) + assert actual is prepared + cached_contract = prepared._prepared_spec_metadata_tensors + assert cached_contract is not None + + replayed = builder.build( + common_prefix_len=0, + common_attn_metadata=common, + num_accepted_tokens=num_accepted_tokens, + num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, + common_gdn_metadata=common_metadata, + prepared_dflash2_metadata=prepared, + ) + assert replayed is prepared + assert prepared._prepared_spec_metadata_tensors is cached_contract + + other_builder = _create_gdn_builder( + local_gdn_model, + num_speculative_tokens=2, + use_full_cuda_graph=True, + ) + with pytest.raises(ValueError, match="does not belong"): + other_builder.build( + common_prefix_len=0, + common_attn_metadata=common, + num_accepted_tokens=num_accepted_tokens, + num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, + common_gdn_metadata=common_metadata, + prepared_dflash2_metadata=prepared, + ) + + +def test_dflash2_fused_gdn_group_metadata_replay( + monkeypatch, + local_gdn_model, +): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for the fused GDN metadata kernel") + device = torch.device("cuda") + if torch.cuda.get_device_capability(device) != (7, 0): + pytest.skip("the fused GDN metadata kernel is SM70-only") + + # The accepted DFlash2 verifier keeps the legacy spec-core operator off. + # Fused metadata must therefore own its graph-stable buffers independently. + monkeypatch.setenv("VLLM_SM70_QWEN_GDN_SPEC_CORE_OP", "0") + monkeypatch.setenv("VLLM_SM70_DFLASH2_FUSED_GDN_METADATA", "1") + monkeypatch.setenv("VLLM_SM70_DFLASH2_GDN_METADATA_SHADOW", "1") + qwen_gdn.envs.disable_envs_cache() + + builders = [ + _create_gdn_builder( + local_gdn_model, + num_speculative_tokens=7, + use_full_cuda_graph=True, + max_cudagraph_capture_size=16, + device=device, + ) + for _ in range(2) + ] + width = builders[0].num_spec_state_tokens + 1 + assert width == 8 + query_start_loc = torch.tensor([0, 8, 16], dtype=torch.int32, device=device) + query_start_loc_cpu = query_start_loc.cpu() + common_metadata = compute_common_gdn_attn_metadata( + num_decode_draft_tokens_cpu=torch.tensor([7, 7], dtype=torch.int32), + query_start_loc=query_start_loc, + query_start_loc_cpu=query_start_loc_cpu, + num_spec_state_tokens=7, + legacy_mixed_decode_routing=False, + ) + assert common_metadata is not None + + block_tables = ( + torch.arange(10, 10 + 2 * width, dtype=torch.int32, device=device).view( + 2, width + ), + torch.arange(100, 100 + 2 * width, dtype=torch.int32, device=device).view( + 2, width + ), + ) + accepted = torch.tensor([3, 5], dtype=torch.int32, device=device) + for builder in builders: + builder.spec_state_indices_tensor.fill_(777) + + first_result = prepare_dflash2_gdn_group_metadata( + builders_by_group=[(0, builders[0]), (1, builders[1])], + block_tables=block_tables, + common_gdn_metadata=common_metadata, + num_accepted_tokens=accepted, + num_actual_tokens=16, + descriptor=None, + ) + assert first_result is not None + prepared, descriptor = first_result + torch.accelerator.synchronize(device) + + for group_id, builder in enumerate(builders): + state = prepared[id(builder)].spec_state_indices_tensor + assert state is not None + torch.testing.assert_close(state[:2], block_tables[group_id], rtol=0, atol=0) + torch.testing.assert_close( + state[2:], torch.full_like(state[2:], PAD_SLOT_ID), rtol=0, atol=0 + ) + first = prepared[id(builders[0])] + assert first.spec_query_start_loc is not None + assert first.spec_query_start_loc.tolist() == [0, 8, 16] + [16] * 14 + assert first.num_accepted_tokens is not None + assert first.num_accepted_tokens.tolist() == [3, 5] + [1] * 14 + assert first.spec_state_slot_selectors is not None + assert first.spec_state_slot_selectors.tolist() == [3, 5] + [1] * 14 + assert first.spec_sequence_masks is not None + assert first.spec_sequence_masks.tolist() == [True, True] + [False] * 14 + + # Poison every persistent output and replay with new source values. This + # catches stale graph-tail data and accidental reuse of one group's table. + for builder in builders: + builder.spec_state_indices_tensor.fill_(888) + common_buffers = builders[0]._ddtree_fast_common_buffers + assert common_buffers is not None + common_buffers.spec_sequence_masks.fill_(True) + common_buffers.spec_query_start_loc.fill_(-1) + common_buffers.num_accepted_tokens.fill_(-1) + common_buffers.spec_state_slot_selectors.fill_(-1) + block_tables[0].add_(1000) + block_tables[1].add_(2000) + accepted.copy_(torch.tensor([7, 2], dtype=torch.int32, device=device)) + + second_result = prepare_dflash2_gdn_group_metadata( + builders_by_group=[(0, builders[0]), (1, builders[1])], + block_tables=block_tables, + common_gdn_metadata=common_metadata, + num_accepted_tokens=accepted, + num_actual_tokens=16, + descriptor=descriptor, + ) + assert second_result is not None + replayed, replay_descriptor = second_result + assert replay_descriptor is descriptor + assert replayed is prepared + torch.accelerator.synchronize(device) + for group_id, builder in enumerate(builders): + state = replayed[id(builder)].spec_state_indices_tensor + assert state is not None + torch.testing.assert_close(state[:2], block_tables[group_id], rtol=0, atol=0) + torch.testing.assert_close( + state[2:], torch.full_like(state[2:], PAD_SLOT_ID), rtol=0, atol=0 + ) + assert replayed[id(builders[0])].num_accepted_tokens is not None + assert replayed[id(builders[0])].num_accepted_tokens.tolist() == [7, 2] + [1] * 14 + + # The pointer tables remain valid when the graph batch shape changes, but + # the graph-buffer views do not. Ensure the cache key rebuilds their + # lengths instead of returning the steady-state B2 metadata. + single_common = compute_common_gdn_attn_metadata( + num_decode_draft_tokens_cpu=torch.tensor([7], dtype=torch.int32), + query_start_loc=torch.tensor([0, 8], dtype=torch.int32, device=device), + query_start_loc_cpu=torch.tensor([0, 8], dtype=torch.int32), + num_spec_state_tokens=7, + legacy_mixed_decode_routing=False, + ) + assert single_common is not None + single_result = prepare_dflash2_gdn_group_metadata( + builders_by_group=[(0, builders[0]), (1, builders[1])], + block_tables=(block_tables[0][:1], block_tables[1][:1]), + common_gdn_metadata=single_common, + num_accepted_tokens=accepted[:1], + num_actual_tokens=8, + descriptor=descriptor, + ) + assert single_result is not None + single_prepared, single_descriptor = single_result + assert single_descriptor is descriptor + assert single_prepared is not replayed + torch.accelerator.synchronize(device) + single = single_prepared[id(builders[0])] + assert single.spec_state_indices_tensor is not None + assert single.spec_state_indices_tensor.shape == (8, 8) + assert single.spec_query_start_loc is not None + assert single.spec_query_start_loc.tolist() == [0, 8] + [8] * 7 + assert single.num_accepted_tokens is not None + assert single.num_accepted_tokens.tolist() == [7] + [1] * 7 + + +def test_dflash2_fused_gdn_group_metadata_align_replay( + monkeypatch, + local_gdn_model, +): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for the fused GDN metadata kernel") + device = torch.device("cuda") + if torch.cuda.get_device_capability(device) != (7, 0): + pytest.skip("the fused GDN metadata kernel is SM70-only") + + monkeypatch.setenv("VLLM_SM70_QWEN_GDN_SPEC_CORE_OP", "0") + monkeypatch.setenv("VLLM_SM70_DFLASH2_FUSED_GDN_METADATA", "1") + monkeypatch.setenv("VLLM_SM70_DFLASH2_GDN_METADATA_SHADOW", "1") + qwen_gdn.envs.disable_envs_cache() + + builders = [ + _create_gdn_builder( + local_gdn_model, + num_speculative_tokens=7, + use_full_cuda_graph=True, + mamba_cache_mode="align", + max_cudagraph_capture_size=16, + device=device, + ) + for _ in range(2) + ] + width = builders[0].num_spec_state_tokens + 1 + query_start_loc = torch.tensor([0, 8, 16], dtype=torch.int32, device=device) + common_metadata = compute_common_gdn_attn_metadata( + num_decode_draft_tokens_cpu=torch.tensor([7, 7], dtype=torch.int32), + query_start_loc=query_start_loc, + query_start_loc_cpu=query_start_loc.cpu(), + num_spec_state_tokens=7, + legacy_mixed_decode_routing=False, + ) + assert common_metadata is not None + + table_columns = 32 + block_tables = ( + torch.arange( + 10, + 10 + 2 * table_columns, + dtype=torch.int32, + device=device, + ).view(2, table_columns), + torch.arange( + 100, + 100 + 2 * table_columns, + dtype=torch.int32, + device=device, + ).view(2, table_columns), + ) + # Batch rows deliberately map to non-adjacent persistent request slots. + req_index_mapping = torch.tensor([2, 0], dtype=torch.int32, device=device) + state_start_indices = torch.tensor([3, 12, 7, 18], dtype=torch.int32, device=device) + accepted = torch.tensor([3, 5], dtype=torch.int32, device=device) + + # Align metadata is incomplete without both persistent-state inputs. + assert ( + prepare_dflash2_gdn_group_metadata( + builders_by_group=[(0, builders[0]), (1, builders[1])], + block_tables=block_tables, + common_gdn_metadata=common_metadata, + num_accepted_tokens=accepted, + num_actual_tokens=16, + descriptor=None, + state_start_indices=state_start_indices, + ) + is None + ) + + first_result = prepare_dflash2_gdn_group_metadata( + builders_by_group=[(0, builders[0]), (1, builders[1])], + block_tables=block_tables, + common_gdn_metadata=common_metadata, + num_accepted_tokens=accepted, + num_actual_tokens=16, + descriptor=None, + state_start_indices=state_start_indices, + req_index_mapping=req_index_mapping, + ) + assert first_result is not None + prepared, descriptor = first_result + torch.accelerator.synchronize(device) + + starts = [7, 3] + for group_id, builder in enumerate(builders): + state = prepared[id(builder)].spec_state_indices_tensor + assert state is not None + expected = torch.stack( + [ + block_tables[group_id][row, start : start + width] + for row, start in enumerate(starts) + ] + ) + torch.testing.assert_close(state[:2], expected, rtol=0, atol=0) + torch.testing.assert_close( + state[2:], torch.full_like(state[2:], PAD_SLOT_ID), rtol=0, atol=0 + ) + + # Replay with the same persistent pointers but a different request mapping + # and state columns. CUDA graph replay must observe the live values. + for builder in builders: + builder.spec_state_indices_tensor.fill_(888) + # The model runner allocates this short mapping per batch. Its pointer is + # intentionally allowed to change without rebuilding the group descriptor. + req_index_mapping = torch.tensor([1, 3], dtype=torch.int32, device=device) + state_start_indices.copy_( + torch.tensor([2, 9, 5, 20], dtype=torch.int32, device=device) + ) + second_result = prepare_dflash2_gdn_group_metadata( + builders_by_group=[(0, builders[0]), (1, builders[1])], + block_tables=block_tables, + common_gdn_metadata=common_metadata, + num_accepted_tokens=accepted, + num_actual_tokens=16, + descriptor=descriptor, + state_start_indices=state_start_indices, + req_index_mapping=req_index_mapping, + ) + assert second_result is not None + replayed, replay_descriptor = second_result + assert replay_descriptor is descriptor + assert replayed is prepared + torch.accelerator.synchronize(device) + + starts = [9, 20] + for group_id, builder in enumerate(builders): + state = replayed[id(builder)].spec_state_indices_tensor + assert state is not None + expected = torch.stack( + [ + block_tables[group_id][row, start : start + width] + for row, start in enumerate(starts) + ] + ) + torch.testing.assert_close(state[:2], expected, rtol=0, atol=0) + torch.testing.assert_close( + state[2:], torch.full_like(state[2:], PAD_SLOT_ID), rtol=0, atol=0 + ) + + def test_full_cuda_graph_capture_single_token_decode_is_not_spec(local_gdn_model): builder = _create_gdn_builder( local_gdn_model, @@ -468,9 +1054,7 @@ def test_spec_state_slot_selector_can_differ_from_accepted_count(local_gdn_model common_prefix_len=0, common_attn_metadata=common, num_accepted_tokens=torch.tensor([2], dtype=torch.int32, device=DEVICE), - spec_state_slot_selectors=torch.tensor( - [4], dtype=torch.int32, device=DEVICE - ), + spec_state_slot_selectors=torch.tensor([4], dtype=torch.int32, device=DEVICE), num_decode_draft_tokens_cpu=torch.tensor([4], dtype=torch.int32, device="cpu"), ) @@ -749,9 +1333,7 @@ def test_gdn_state_contract_debug_assert_rejects_pad_accepted_slot( spec_sequence_masks_cpu=torch.tensor( [True], dtype=torch.bool, device="cpu" ), - num_accepted_tokens=torch.tensor( - [5], dtype=torch.int32, device=DEVICE - ), + num_accepted_tokens=torch.tensor([5], dtype=torch.int32, device=DEVICE), current_state_block_ids=None, is_mamba_cache_all=False, ) @@ -778,9 +1360,7 @@ def test_spec_commit_pure_decode_consumes_padded_graph_rows_without_metadata_pat num_spec_decodes=2, num_spec_decode_tokens=10, num_actual_tokens=10, - spec_query_start_loc=torch.tensor( - [0, 5, 10], dtype=torch.int32, device=DEVICE - ), + spec_query_start_loc=torch.tensor([0, 5, 10], dtype=torch.int32, device=DEVICE), spec_state_indices_tensor=torch.tensor( [[10, 11, 12, 13, 14], [20, 21, 22, 23, 24]], dtype=torch.int32, @@ -821,7 +1401,6 @@ def test_spec_commit_pure_decode_consumes_padded_graph_rows_without_metadata_pat seen: dict[str, object] = {} class FakeLayer: - def __init__(self) -> None: self.conv1d = torch.nn.Identity() self.conv1d.weight = torch.empty((8, 1, 1), device=DEVICE) @@ -862,10 +1441,11 @@ def fake_conv_update( seen["conv"] = True return x + 1 - def fake_gating(a_log, a_tensor, b_tensor, dt_bias): + def fake_gating(a_log, a_tensor, b_tensor, dt_bias, *, beta_dtype=None): del a_log, dt_bias assert a_tensor is a assert b_tensor is b + assert beta_dtype is torch.float32 seen["gating"] = True return a_tensor + 2, b_tensor + 3 @@ -1023,9 +1603,7 @@ def test_sm70_qwen_gdn_blocks_003_spec_core_only_for_deep_native_mtp( expected, ): monkeypatch.setenv("VLLM_SM70_QWEN_GDN_003_SPEC_CORE_OP", "1") - monkeypatch.delenv( - "VLLM_SM70_QWEN_GDN_003_SPEC_ALLOW_DEEP_MTP", raising=False - ) + monkeypatch.delenv("VLLM_SM70_QWEN_GDN_003_SPEC_ALLOW_DEEP_MTP", raising=False) _clear_envs_cache() vllm_config = SimpleNamespace( @@ -1036,9 +1614,7 @@ def test_sm70_qwen_gdn_blocks_003_spec_core_only_for_deep_native_mtp( ) assert ( - qwen_gdn._sm70_qwen_gdn_block_003_spec_for_deep_native_mtp( - vllm_config - ) + qwen_gdn._sm70_qwen_gdn_block_003_spec_for_deep_native_mtp(vllm_config) is expected ) @@ -1057,9 +1633,7 @@ def test_sm70_qwen_gdn_003_spec_core_deep_mtp_override_is_diagnostic( ) ) - assert not qwen_gdn._sm70_qwen_gdn_block_003_spec_for_deep_native_mtp( - vllm_config - ) + assert not qwen_gdn._sm70_qwen_gdn_block_003_spec_for_deep_native_mtp(vllm_config) def test_sm70_qwen_gdn_003_spec_core_route_has_priority( diff --git a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py index 159bb8af3f..7c7d535f08 100644 --- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py +++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py @@ -63,6 +63,9 @@ def test_indexer_builder_deepseek_v4_compressed_slot_mapping_uses_storage_block_ md = builder.build(common_prefix_len=0, common_attn_metadata=common) + assert md.max_seq_len == 280 + assert md.compressed_max_seq_len == 70 + # The compressed slot_mapping retains the original uncompressed size (40). # Only every compress_ratio-th position gets a valid slot; the rest are -1. assert md.slot_mapping.numel() == 40 diff --git a/tests/v1/attention/test_prefix_anchored_swa_mask.py b/tests/v1/attention/test_prefix_anchored_swa_mask.py new file mode 100644 index 0000000000..0e6a744e3a --- /dev/null +++ b/tests/v1/attention/test_prefix_anchored_swa_mask.py @@ -0,0 +1,181 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the prefix-anchored sliding-window attention mask. + +Semantics under test: keys inside the prompt/prefix (``kv < anchor_len``) +stay globally visible; generated keys are additionally restricted to a +sliding window (``q_abs - kv < window``); everything is AND-ed with the +causal mask. + +The CPU tests validate the mask formula (the same one the masked kernels +implement) against plain causal attention with a torch reference. +""" + +from types import SimpleNamespace +from typing import Any, cast + +import pytest +import torch + +from vllm.v1.attention.backend import AttentionType +from vllm.v1.attention.backends.flash_attn_v100 import ( + FlashAttnV100Impl, + FlashAttnV100MetadataBuilder, +) + + +def anchored_swa_mask( + seq_len: int, anchor_len: int, window: int, device: str = "cpu" +) -> torch.Tensor: + """Boolean keep-mask [seq_len, seq_len]; True = attend. + + Mirrors the kernel formula: + ``causal AND (kv < anchor_len OR q_abs - kv < window)``. + """ + q = torch.arange(seq_len, device=device)[:, None] + kv = torch.arange(seq_len, device=device)[None, :] + causal = kv <= q + in_prefix = kv < anchor_len + in_window = (q - kv) < window + return causal & (in_prefix | in_window) + + +def causal_mask(seq_len: int, device: str = "cpu") -> torch.Tensor: + q = torch.arange(seq_len, device=device)[:, None] + kv = torch.arange(seq_len, device=device)[None, :] + return kv <= q + + +def ref_attention( + q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor +) -> torch.Tensor: + """Plain fp32 reference attention. q/k/v: [heads, seq, dim].""" + scale = q.shape[-1] ** -0.5 + scores = torch.einsum("hqd,hkd->hqk", q.float(), k.float()) * scale + scores = scores.masked_fill(~mask[None], float("-inf")) + probs = torch.softmax(scores, dim=-1) + return torch.einsum("hqk,hkd->hqd", probs, v.float()) + + +def _flash_impl(window: int | None) -> FlashAttnV100Impl: + impl = object.__new__(FlashAttnV100Impl) + object.__setattr__(impl, "prefix_anchored_decode_window", window) + object.__setattr__(impl, "attn_type", AttentionType.DECODER) + object.__setattr__(impl, "kv_cache_dtype", "auto") + return impl + + +@pytest.mark.cpu_test +def test_backend_default_off_does_not_inspect_metadata(): + class NoMetadataAccess: + def __getattr__(self, name: str): + raise AssertionError(f"default-off path inspected {name}") + + metadata = cast(Any, NoMetadataAccess()) + assert _flash_impl(None)._anchored_swa_params(metadata) == (None, 0) + + +@pytest.mark.cpu_test +def test_backend_enabled_fails_closed_without_matching_metadata(): + with pytest.raises(RuntimeError, match="metadata does not match"): + _flash_impl(128)._anchored_swa_params(cast(Any, object())) + + +@pytest.mark.cpu_test +def test_metadata_builder_attaches_prefix_lengths_for_eager_path(): + builder = object.__new__(FlashAttnV100MetadataBuilder) + object.__setattr__(builder, "decode_sliding_window", 128) + object.__setattr__(builder, "device", torch.device("cpu")) + object.__setattr__( + builder, + "persistent_prefix_anchor_lens", + torch.empty(4, dtype=torch.int32), + ) + metadata = SimpleNamespace() + common = SimpleNamespace( + prefix_anchor_lens=torch.tensor([7, 11], dtype=torch.int64), + num_reqs=2, + ) + + builder._attach_prefix_anchored_metadata( + cast(Any, metadata), + cast(Any, common), + ) + + assert metadata.decode_sliding_window == 128 + assert metadata.prefix_anchor_lens.dtype == torch.int32 + assert metadata.prefix_anchor_lens.tolist() == [7, 11] + + +@pytest.mark.cpu_test +def test_metadata_builder_fails_closed_without_prefix_lengths(): + builder = object.__new__(FlashAttnV100MetadataBuilder) + object.__setattr__(builder, "decode_sliding_window", 128) + object.__setattr__(builder, "device", torch.device("cpu")) + object.__setattr__( + builder, + "persistent_prefix_anchor_lens", + torch.empty(4, dtype=torch.int32), + ) + + with pytest.raises(RuntimeError, match="requires per-request prefix lengths"): + builder._attach_prefix_anchored_metadata( + cast(Any, SimpleNamespace()), + cast(Any, SimpleNamespace(prefix_anchor_lens=None, num_reqs=2)), + ) + + +@pytest.mark.cpu_test +def test_mask_equals_causal_while_within_window(): + # While the number of generated tokens is <= window, no key falls in the + # gap, so the mask must be exactly the causal mask. + anchor, window = 6, 8 + for seq_len in range(1, anchor + window + 1): + assert torch.equal( + anchored_swa_mask(seq_len, anchor, window), + causal_mask(seq_len), + ), f"seq_len={seq_len}" + + +@pytest.mark.cpu_test +def test_mask_prunes_gap_beyond_window(): + anchor, window = 6, 8 + seq_len = anchor + window + 3 # 3 generated keys fall out of the window + mask = anchored_swa_mask(seq_len, anchor, window) + causal = causal_mask(seq_len) + assert not torch.equal(mask, causal) + # The last query keeps: the full prefix and the window of recent keys... + last = mask[-1] + assert last[:anchor].all(), "prefix must stay globally visible" + assert last[seq_len - window :].all(), "recent window must be visible" + # ...and drops exactly the gap keys in between. + gap = last[anchor : seq_len - window] + assert not gap.any(), "gap keys must be masked" + # Prefix queries (inside the prompt) are never affected. + assert torch.equal(mask[: anchor + window], causal[: anchor + window]) + + +@pytest.mark.cpu_test +def test_attention_output_equivalence_and_divergence(): + torch.manual_seed(0) + heads, dim = 2, 32 + anchor, window = 6, 8 + + # Within the window: identical outputs. + seq_len = anchor + window + q, k, v = (torch.randn(heads, seq_len, dim) for _ in range(3)) + out_anchored = ref_attention(q, k, v, anchored_swa_mask(seq_len, anchor, window)) + out_causal = ref_attention(q, k, v, causal_mask(seq_len)) + torch.testing.assert_close(out_anchored, out_causal) + + # Beyond the window: the late queries must diverge from full attention + # (negative control proving the mask actually removes information). + seq_len = anchor + 2 * window + q, k, v = (torch.randn(heads, seq_len, dim) for _ in range(3)) + out_anchored = ref_attention(q, k, v, anchored_swa_mask(seq_len, anchor, window)) + out_causal = ref_attention(q, k, v, causal_mask(seq_len)) + assert not torch.allclose(out_anchored[:, -1], out_causal[:, -1]) + # Queries that see no gap still match exactly. + torch.testing.assert_close( + out_anchored[:, : anchor + window], out_causal[:, : anchor + window] + ) diff --git a/tests/v1/attention/test_sm70_flash_v100_policy.py b/tests/v1/attention/test_sm70_flash_v100_policy.py index f16ce3800c..fb2e195796 100644 --- a/tests/v1/attention/test_sm70_flash_v100_policy.py +++ b/tests/v1/attention/test_sm70_flash_v100_policy.py @@ -237,6 +237,35 @@ def test_sm70_fa2_d256_prefill_env_is_default_on(monkeypatch): assert envs.VLLM_FLASH_V100_FA2_D256_PREFILL is False +def test_sm70_d256_gqa_architecture_env_is_default_on(monkeypatch): + import vllm.envs as envs + + name = "VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL" + monkeypatch.delenv(name, raising=False) + envs.disable_envs_cache() + assert envs.VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL is True + + monkeypatch.setenv(name, "0") + envs.disable_envs_cache() + assert envs.VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL is False + + +def test_sm70_e4m3_batch_xqa_env_contract(monkeypatch): + import vllm.envs as envs + + monkeypatch.delenv("VLLM_FLASH_V100_E4M3_BATCH_XQA", raising=False) + monkeypatch.delenv("VLLM_FLASH_V100_E4M3_BATCH_XQA_OPTIMIZED", raising=False) + envs.disable_envs_cache() + assert envs.VLLM_FLASH_V100_E4M3_BATCH_XQA is True + assert envs.VLLM_FLASH_V100_E4M3_BATCH_XQA_OPTIMIZED is True + + monkeypatch.setenv("VLLM_FLASH_V100_E4M3_BATCH_XQA", "0") + monkeypatch.setenv("VLLM_FLASH_V100_E4M3_BATCH_XQA_OPTIMIZED", "0") + envs.disable_envs_cache() + assert envs.VLLM_FLASH_V100_E4M3_BATCH_XQA is False + assert envs.VLLM_FLASH_V100_E4M3_BATCH_XQA_OPTIMIZED is False + + def test_sm70_e5m2_decode_fast_route_envs_are_default_on(monkeypatch): import vllm.envs as envs @@ -245,6 +274,7 @@ def test_sm70_e5m2_decode_fast_route_envs_are_default_on(monkeypatch): "VLLM_FLASH_V100_XQA_E5M2_G6_SPLIT_REDUCE", "VLLM_FLASH_V100_XQA_E5M2_PARTITION_PAGE_IDS", "VLLM_FLASH_V100_XQA_E5M2_PAIR_LOAD", + "VLLM_FLASH_V100_XQA_E5M2_BATCH_WIDE_LOAD", ) for name in names: monkeypatch.delenv(name, raising=False) @@ -254,6 +284,7 @@ def test_sm70_e5m2_decode_fast_route_envs_are_default_on(monkeypatch): assert envs.VLLM_FLASH_V100_XQA_E5M2_G6_SPLIT_REDUCE is True assert envs.VLLM_FLASH_V100_XQA_E5M2_PARTITION_PAGE_IDS is True assert envs.VLLM_FLASH_V100_XQA_E5M2_PAIR_LOAD is True + assert envs.VLLM_FLASH_V100_XQA_E5M2_BATCH_WIDE_LOAD is True assert envs.VLLM_FLASH_V100_XQA_E5M2_P1024_BEGIN == 61633 for name in names: @@ -263,6 +294,7 @@ def test_sm70_e5m2_decode_fast_route_envs_are_default_on(monkeypatch): assert envs.VLLM_FLASH_V100_XQA_E5M2_G6_SPLIT_REDUCE is False assert envs.VLLM_FLASH_V100_XQA_E5M2_PARTITION_PAGE_IDS is False assert envs.VLLM_FLASH_V100_XQA_E5M2_PAIR_LOAD is False + assert envs.VLLM_FLASH_V100_XQA_E5M2_BATCH_WIDE_LOAD is False monkeypatch.setenv("VLLM_FLASH_V100_XQA_E5M2_P1024_BEGIN", "49152") envs.disable_envs_cache() @@ -376,10 +408,15 @@ def test_sm70_prefill_dense_splitkv3_env_is_default_on(monkeypatch): monkeypatch.delenv("VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3", raising=False) monkeypatch.delenv("VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3_MIN_KV", raising=False) + monkeypatch.delenv( + "VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3_Q8000_EXPERIMENTAL", + raising=False, + ) envs.disable_envs_cache() assert envs.VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3 is True assert envs.VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3_MIN_KV == 32768 + assert envs.VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3_Q8000_EXPERIMENTAL is False monkeypatch.setenv("VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3", "0") envs.disable_envs_cache() @@ -580,6 +617,7 @@ def test_prefix_prefill_prioritizes_gathered_exact_dense_over_paged( impl.head_size = 256 impl.scale = 0.0625 impl.sliding_window = None + impl.prefix_anchored_decode_window = None impl.kv_cache_dtype = "auto" impl.use_flash_v100_decode = False impl.use_decode_paged_prefill = False @@ -700,6 +738,36 @@ def test_sm70_splitd_d256_loader_requires_exact_ops(monkeypatch): assert flash_v100._get_sm70_splitd_d256_ops() == (dense, paged, splitkv3) +def test_sm70_d256_gqa_architecture_loader_is_optional(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, + ) + + architecture = object() + fake_ops = SimpleNamespace( + _vllm_fa2_C=SimpleNamespace( + sm70_d256_gqa_architecture_fwd=architecture, + ) + ) + monkeypatch.setattr(flash_v100, "torch", SimpleNamespace(ops=fake_ops)) + monkeypatch.setattr( + flash_v100, + "_sm70_d256_gqa_architecture_op_checked", + False, + ) + monkeypatch.setattr(flash_v100, "_sm70_d256_gqa_architecture_op", None) + + assert flash_v100._get_sm70_d256_gqa_architecture_op() is architecture + + def test_prefill_dense_splitkv3_workspace_reuses_exact_shape(monkeypatch): import vllm.v1.attention.backends.flash_attn_v100 as flash_v100 @@ -743,6 +811,8 @@ def test_prefill_dense_splitkv3_policy_is_exact_shape_bounded(monkeypatch): key = torch.empty((1, 65536, 1, 256), dtype=torch.float16, device="meta") key_32k = torch.empty((1, 32768, 1, 256), dtype=torch.float16, device="meta") key_16k = torch.empty((1, 16384, 1, 256), dtype=torch.float16, device="meta") + query_8k = torch.empty((1, 8000, 6, 256), dtype=torch.float16, device="meta") + query_8192 = torch.empty((1, 8192, 6, 256), dtype=torch.float16, device="meta") assert flash_v100._should_use_prefill_dense_splitkv3( query, @@ -772,6 +842,168 @@ def test_prefill_dense_splitkv3_policy_is_exact_shape_bounded(monkeypatch): max_seqlen_k=65536, splitkv3_op=object(), ) + assert not flash_v100._should_use_prefill_dense_splitkv3( + query_8k, + key, + max_seqlen_q=8000, + max_seqlen_k=65536, + splitkv3_op=object(), + ) + + monkeypatch.setenv("VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3_Q8000_EXPERIMENTAL", "1") + envs.disable_envs_cache() + assert flash_v100._should_use_prefill_dense_splitkv3( + query_8k, + key, + max_seqlen_q=8000, + max_seqlen_k=65536, + splitkv3_op=object(), + ) + assert not flash_v100._should_use_prefill_dense_splitkv3( + query_8192, + key, + max_seqlen_q=8192, + max_seqlen_k=65536, + splitkv3_op=object(), + ) + + +def test_prefill_d256_gqa_architecture_policy_is_shape_family_bounded(monkeypatch): + import vllm.envs as envs + import vllm.v1.attention.backends.flash_attn_v100 as flash_v100 + + name = "VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL" + query = torch.empty((1, 8000, 6, 256), dtype=torch.float16, device="meta") + key = torch.empty((1, 128000, 1, 256), dtype=torch.float16, device="meta") + value = torch.empty_like(key) + + monkeypatch.setenv(name, "0") + envs.disable_envs_cache() + assert not flash_v100._should_use_prefill_d256_gqa_architecture( + query, + key, + value, + max_seqlen_q=8000, + max_seqlen_k=128000, + softmax_scale=0.0625, + architecture_op=object(), + ) + + monkeypatch.delenv(name, raising=False) + envs.disable_envs_cache() + for kv_len in range(40000, 128001, 8000): + family_key = torch.empty( + (1, kv_len, 1, 256), dtype=torch.float16, device="meta" + ) + assert flash_v100._should_use_prefill_d256_gqa_architecture( + query, + family_key, + torch.empty_like(family_key), + max_seqlen_q=8000, + max_seqlen_k=kv_len, + softmax_scale=0.0625, + architecture_op=object(), + ) + assert not flash_v100._should_use_prefill_d256_gqa_architecture( + query[:, :7999], + key, + value, + max_seqlen_q=7999, + max_seqlen_k=128000, + softmax_scale=0.0625, + architecture_op=object(), + ) + assert not flash_v100._should_use_prefill_d256_gqa_architecture( + query, + key[:, :127999], + value[:, :127999], + max_seqlen_q=8000, + max_seqlen_k=127999, + softmax_scale=0.0625, + architecture_op=object(), + ) + assert not flash_v100._should_use_prefill_d256_gqa_architecture( + query, + key[:, :44000], + value[:, :44000], + max_seqlen_q=8000, + max_seqlen_k=44000, + softmax_scale=0.0625, + architecture_op=object(), + ) + assert not flash_v100._should_use_prefill_d256_gqa_architecture( + query, + key, + value, + max_seqlen_q=8000, + max_seqlen_k=128000, + softmax_scale=1.0, + architecture_op=object(), + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_prefill_d256_gqa_architecture_oom_uses_dense_fallback(monkeypatch): + import vllm.envs as envs + import vllm.v1.attention.backends.flash_attn_v100 as flash_v100 + + if torch.cuda.get_device_capability() != (7, 0): + pytest.skip("SM70/V100 is required") + + name = "VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL" + monkeypatch.delenv(name, raising=False) + envs.disable_envs_cache() + monkeypatch.setattr(flash_v100, "_is_cuda_graph_capturing", lambda _: False) + + dense_calls = 0 + + def dense_op(query, key, value, out, softmax_scale, causal): + nonlocal dense_calls + dense_calls += 1 + out.fill_(3) + return out + + def architecture_oom(*args, **kwargs): + raise torch.OutOfMemoryError("expected architecture workspace OOM") + + monkeypatch.setattr( + flash_v100, + "_get_sm70_splitd_d256_ops", + lambda: (dense_op, object(), None), + ) + monkeypatch.setattr( + flash_v100, + "_get_sm70_d256_gqa_architecture_op", + lambda: architecture_oom, + ) + monkeypatch.setattr(flash_v100, "_warned_prefill_d256_gqa_architecture_oom", False) + + query = torch.zeros((1, 8000, 6, 256), dtype=torch.float16, device="cuda") + key = torch.zeros((1, 40000, 1, 256), dtype=torch.float16, device="cuda") + value = torch.zeros_like(key) + out = torch.zeros_like(query) + cu_seqlens_q = torch.tensor([0, 8000], dtype=torch.int32, device="cuda") + cu_seqlens_k = torch.tensor([0, 40000], dtype=torch.int32, device="cuda") + + result = flash_v100._try_sm70_fa2_d256_prefill( + query, + key, + value, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=8000, + max_seqlen_k=40000, + softmax_scale=0.0625, + causal=True, + window_size=(-1, -1), + out=out, + ) + + assert result is not None + assert result.data_ptr() == out.data_ptr() + assert dense_calls == 1 + assert torch.all(out == 3) + assert flash_v100._warned_prefill_d256_gqa_architecture_oom is True def test_dense_prefill_restores_uniform_batch_view_for_exact_splitd(monkeypatch): @@ -1296,6 +1528,27 @@ def test_sm70_nomtp_cudagraph_capture_sizes_cover_concurrency( assert _sm70_nomtp_cudagraph_capture_sizes(max_num_seqs) == expected +@pytest.mark.parametrize( + ("max_num_seqs", "expected"), + [ + (1, [5]), + (2, [5, 10]), + (4, [5, 10, 20]), + (6, [5, 10, 20, 30]), + (12, [5, 10, 20, 30, 40, 60]), + (16, [5, 10, 20, 30, 40, 60, 80]), + (32, [5, 10, 20, 30, 40, 60, 80]), + ], +) +def test_sm70_mtp_cudagraph_capture_sizes_cover_production_concurrency( + max_num_seqs: int, + expected: list[int], +): + from vllm.config.vllm import _sm70_mtp_cudagraph_capture_sizes + + assert _sm70_mtp_cudagraph_capture_sizes(max_num_seqs, 5) == expected + + def test_flash_v100_decode_query_does_not_attach_smallq_metadata( monkeypatch, local_flash_v100_model, @@ -1402,6 +1655,89 @@ def fake_decode( assert torch.all(output == 1) +def test_flash_v100_dflash2_grouped_verify_uses_original_request_metadata(): + from vllm.v1.attention.backends.flash_attn_v100 import FlashAttnV100Impl + + impl = FlashAttnV100Impl( + num_heads=6, + head_size=256, + scale=1.0, + num_kv_heads=1, + alibi_slopes=None, + sliding_window=None, + kv_cache_dtype="fp8_e5m2", + ) + impl.use_dflash2_grouped_verify = True + captured: dict[str, object] = {} + + def grouped_verify( + query, + key_cache, + value_cache, + block_table, + seq_lens, + **kwargs, + ): + captured["block_table"] = block_table + captured["seq_lens"] = seq_lens + captured["one_pass"] = kwargs["one_pass"] + kwargs["out"].fill_(1) + + impl.flash_attn_grouped_verify_paged = grouped_verify + impl.flash_attn_decode_paged = lambda *args, **kwargs: pytest.fail( + "exact DFlash2 grouped route fell through to independent rows" + ) + 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, + causal=True, + is_dflash_selector_target=True, + max_model_len=32768, + query_start_loc=torch.tensor([0, 8], 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), + ) + layer = SimpleNamespace(_k_scale_float=0.5, _v_scale_float=2.0) + query = torch.zeros((8, 6, 256), dtype=torch.float16) + output = torch.zeros_like(query) + key_cache = torch.zeros((2, 3296, 1, 256), dtype=torch.uint8) + value_cache = torch.zeros_like(key_cache) + + attn_metadata.max_model_len = 8192 + assert not impl._dflash2_grouped_verify_allowed( + query, + key_cache, + value_cache, + attn_metadata, + num_query_tokens=8, + ) + attn_metadata.max_model_len = 32768 + result = impl._flash_v100_small_query_prefill_as_decode( + layer, + query, + key_cache, + value_cache, + attn_metadata, + output, + attn_metadata.query_start_loc, + original_seq_lens, + ) + + assert result is output + captured_block_table = captured["block_table"] + captured_seq_lens = captured["seq_lens"] + assert isinstance(captured_block_table, torch.Tensor) + assert isinstance(captured_seq_lens, torch.Tensor) + assert captured_block_table.data_ptr() == original_block_table.data_ptr() + assert captured_seq_lens.data_ptr() == original_seq_lens.data_ptr() + assert captured["one_pass"] is True + assert torch.all(output == 1) + + def test_flash_v100_decode_forwards_shape_hints(): from vllm.v1.attention.backends.flash_attn_v100 import FlashAttnV100Impl @@ -1516,6 +1852,126 @@ def fail_scalar(*args, **kwargs): assert torch.all(output == 1) +def test_flash_v100_decode_uses_xqa_for_e4m3_g6_d256(monkeypatch): + from vllm.v1.attention.backends.flash_attn_v100 import FlashAttnV100Impl + + monkeypatch.delenv("VLLM_FLASH_V100_DECODE_USE_XQA", raising=False) + monkeypatch.delenv("VLLM_FLASH_V100_DECODE_PARTITION_SIZE", raising=False) + + impl = FlashAttnV100Impl( + num_heads=6, + head_size=256, + scale=1.0, + num_kv_heads=1, + alibi_slopes=None, + sliding_window=None, + kv_cache_dtype="fp8_e4m3", + ) + calls: list[tuple[str, int | None, str | None]] = [] + + def hit_xqa(*args, **kwargs): + calls.append( + ( + "xqa", + kwargs.get("partition_size_hint"), + kwargs.get("kv_cache_dtype"), + ) + ) + kwargs["out"].fill_(1) + + def fail_scalar(*args, **kwargs): + raise AssertionError("E4M3 G6/D256 decode should select XQA") + + impl.flash_attn_decode_paged_xqa = hit_xqa # type: ignore[method-assign] + impl.flash_attn_decode_paged = fail_scalar # type: ignore[method-assign] + attn_metadata = SimpleNamespace( + num_actual_tokens=1, + block_table=torch.tensor([[0]], dtype=torch.int32), + seq_lens=torch.tensor([1025], dtype=torch.int32), + flash_v100_decode_max_seq_len_hint=1025, + flash_v100_decode_workspace_seq_capacity_hint=262144, + flash_v100_decode_active_num_partitions=torch.tensor([5], dtype=torch.int32), + ) + layer = SimpleNamespace(_k_scale_float=0.04, _v_scale_float=0.25) + query = torch.zeros((1, 6, 256), dtype=torch.float16) + output = torch.zeros_like(query) + kv_cache = torch.zeros((2, 2, 1568, 1, 256), dtype=torch.uint8) + + result = impl._flash_v100_decode( + layer, + query, + query, + query, + kv_cache, + attn_metadata, + output, + ) + + assert result is output + assert calls == [("xqa", 64, "fp8_e4m3")] + assert torch.all(output == 1) + + +@pytest.mark.parametrize( + ("batch_size", "enabled", "expected_route"), + ((2, False, "scalar"), (2, True, "xqa"), (16, True, "xqa"), (17, True, "scalar")), +) +def test_flash_v100_e4m3_batched_xqa_is_exactly_gated( + monkeypatch, batch_size, enabled, expected_route +): + from vllm.v1.attention.backends.flash_attn_v100 import FlashAttnV100Impl + + monkeypatch.setenv("VLLM_FLASH_V100_E4M3_BATCH_XQA", "1" if enabled else "0") + monkeypatch.delenv("VLLM_FLASH_V100_DECODE_PARTITION_SIZE", raising=False) + impl = FlashAttnV100Impl( + num_heads=6, + head_size=256, + scale=1.0, + num_kv_heads=1, + alibi_slopes=None, + sliding_window=None, + kv_cache_dtype="fp8_e4m3", + ) + calls: list[str] = [] + + def hit_xqa(*args, **kwargs): + calls.append("xqa") + kwargs["out"].fill_(1) + + def hit_scalar(*args, **kwargs): + calls.append("scalar") + kwargs["out"].fill_(1) + + impl.flash_attn_decode_paged_xqa = hit_xqa # type: ignore[method-assign] + impl.flash_attn_decode_paged = hit_scalar # type: ignore[method-assign] + attn_metadata = SimpleNamespace( + num_actual_tokens=batch_size, + block_table=torch.zeros((batch_size, 2), dtype=torch.int32), + seq_lens=torch.full((batch_size,), 16384, dtype=torch.int32), + flash_v100_decode_max_seq_len_hint=16384, + flash_v100_decode_workspace_seq_capacity_hint=32768, + flash_v100_decode_active_num_partitions=torch.tensor([64], dtype=torch.int32), + ) + layer = SimpleNamespace(_k_scale_float=0.04, _v_scale_float=0.25) + query = torch.zeros((batch_size, 6, 256), dtype=torch.float16) + output = torch.zeros_like(query) + kv_cache = torch.zeros((2, 2, 1568, 1, 256), dtype=torch.uint8) + + result = impl._flash_v100_decode( + layer, + query, + query, + query, + kv_cache, + attn_metadata, + output, + ) + + assert result is output + assert calls == [expected_route] + assert torch.all(output == 1) + + @pytest.mark.parametrize( ("num_heads", "seq_len", "expected_route"), ( @@ -1723,6 +2179,113 @@ def test_flash_v100_fp8_xqa_graph_capture_uses_static_context_hint( assert mod._decode_fp8_xqa_allowed(metadata, torch.empty(1)) is expected +@pytest.mark.parametrize( + ("routing_enabled", "graph_variant", "expected"), + ( + (True, None, True), + (True, 0, False), + (True, -1, True), + (False, None, False), + (False, -1, False), + ), +) +def test_flash_v100_batch_context_routing_isolated_by_graph_variant( + routing_enabled, + graph_variant, + expected, +): + from vllm.v1.attention.backends import flash_attn_v100 as mod + + assert ( + mod._batch_context_routing_for_graph_variant( + routing_enabled, + graph_variant, + ) + is expected + ) + + +@pytest.mark.parametrize( + ("cache_dtype", "e4m3_enabled", "expected"), + ( + ("fp8_e5m2", False, True), + ("fp8_e4m3", True, True), + ("fp8", True, True), + ("fp8_e4m3", False, False), + ("auto", True, False), + ), +) +def test_flash_v100_batch_context_routing_accepts_exact_fp8_xqa_formats( + monkeypatch, + cache_dtype, + e4m3_enabled, + expected, +): + from vllm.v1.attention.backends import flash_attn_v100 as mod + + monkeypatch.setenv("VLLM_FLASH_V100_E4M3_BATCH_XQA", "1" if e4m3_enabled else "0") + assert mod._batch_context_routing_cache_dtype_supported(cache_dtype) is expected + + +@pytest.mark.parametrize("routing_enabled", (True, False)) +def test_flash_v100_fp8_xqa_full_capacity_graph_preserves_baseline_xqa( + monkeypatch, + routing_enabled, +): + from vllm.v1.attention.backends.flash_attn_v100 import FlashAttnV100Impl + + impl = FlashAttnV100Impl( + num_heads=6, + head_size=256, + scale=1.0, + num_kv_heads=1, + alibi_slopes=None, + sliding_window=None, + kv_cache_dtype="fp8_e5m2", + ) + calls: list[tuple[str, bool]] = [] + + def hit_xqa(*args, **kwargs): + calls.append(("xqa", kwargs.get("batch_context_routing", False))) + kwargs["out"].fill_(1) + + def hit_scalar(*args, **kwargs): + calls.append(("scalar", False)) + kwargs["out"].fill_(1) + + impl.flash_attn_decode_paged_xqa = hit_xqa # type: ignore[method-assign] + impl.flash_attn_decode_paged = hit_scalar # type: ignore[method-assign] + attn_metadata = SimpleNamespace( + num_actual_tokens=4, + block_table=torch.zeros((4, 256), dtype=torch.int32), + seq_lens=torch.full((4,), 1, dtype=torch.int32), + flash_v100_cudagraph_capture=True, + flash_v100_batch_context_routing=routing_enabled, + flash_v100_decode_max_seq_len_hint=1, + flash_v100_static_decode_seq_hint=262144, + flash_v100_decode_workspace_seq_capacity_hint=262144, + flash_v100_decode_active_num_partitions=torch.tensor([16], dtype=torch.int32), + ) + layer = SimpleNamespace(_k_scale_float=1.0, _v_scale_float=1.0) + query = torch.zeros((4, 6, 256), dtype=torch.float16) + output = torch.zeros_like(query) + kv_cache = torch.zeros((2, 2, 16, 1, 256), dtype=torch.uint8) + + result = impl._flash_v100_decode( + layer, + query, + query, + query, + kv_cache, + attn_metadata, + output, + ) + + assert result is output + assert calls == [("xqa", routing_enabled)] + assert torch.all(output == 1) + + def test_flash_v100_mtp5_dual_cta_partition_policy(monkeypatch): from vllm.v1.attention.backends import flash_attn_v100 as mod @@ -2079,7 +2642,7 @@ def test_flash_v100_decode_e5m2_keeps_small_pages_on_default_planner(monkeypatch ) -def test_flash_v100_decode_e4m3_does_not_use_e5m2_partition_hint(monkeypatch): +def test_flash_v100_decode_e4m3_uses_exact_p64_partition_hint(monkeypatch): from vllm.v1.attention.backends.flash_attn_v100 import ( _g6_aligned_page_partition_size_hint, ) @@ -2091,7 +2654,7 @@ def test_flash_v100_decode_e4m3_does_not_use_e5m2_partition_hint(monkeypatch): assert ( _g6_aligned_page_partition_size_hint(query, key_cache, key_cache, "fp8_e4m3") - is None + == 64 ) diff --git a/tests/v1/core/test_async_scheduler.py b/tests/v1/core/test_async_scheduler.py index eaf4ef82ad..e4ea646ed9 100644 --- a/tests/v1/core/test_async_scheduler.py +++ b/tests/v1/core/test_async_scheduler.py @@ -306,6 +306,9 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): request.status = RequestStatus.RUNNING request.num_computed_tokens = request.num_tokens request.num_output_placeholders = 1 + # This fixture calls update_from_output without schedule(); mirror the + # scheduled token that would normally be recorded by _update_after_schedule. + request.num_in_flight_tokens = 1 scheduler.perf_metrics = None scheduler.connector = None diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 94e9f6f4c1..1ec844fead 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -3,6 +3,7 @@ import hashlib import importlib from collections.abc import Callable +from types import SimpleNamespace from typing import Any import pytest @@ -182,6 +183,37 @@ def new_mamba_spec( ) +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( + num_blocks=1, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec(["attention"], new_kv_cache_spec(block_size=16)), + KVCacheGroupSpec( + ["mamba"], + new_mamba_spec(block_size=1648, mamba_cache_mode="align"), + ), + ], + ) + vllm_config = SimpleNamespace( + cache_config=SimpleNamespace( + block_size=16, + enable_prefix_caching=True, + hash_block_size=None, + ), + parallel_config=SimpleNamespace( + decode_context_parallel_size=1, + prefill_context_parallel_size=1, + ), + kv_transfer_config=None, + ) + + assert kv_cache_utils.resolve_kv_cache_block_sizes( + kv_cache_config, vllm_config + ) == (1648, 16) + + @pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor]) def test_none_hash(monkeypatch, hash_fn): import vllm.v1.core.kv_cache_utils @@ -1362,6 +1394,8 @@ def test_get_max_concurrency_for_kv_cache_config(): enable_chunked_prefill=True, max_model_len=model_config.max_model_len, is_encoder_decoder=model_config.is_encoder_decoder, + # Pin to sync: SWA per-request bounds grow with overlapping batches. + async_scheduling=False, ) vllm_config = VllmConfig( @@ -1898,6 +1932,60 @@ def test_get_kv_cache_spec_kind_prefers_specific_attention_subclasses(): ) +def test_deepseek_v4_tuple_width_minimizes_physical_pool_pages(): + full_spec = MLAAttentionSpec( + block_size=256, + num_kv_heads=1, + head_size=512, + dtype=torch.uint8, + cache_dtype_str="fp8_ds_mla", + compress_ratio=4, + alignment=576, + model_version="deepseek_v4", + ) + swa_spec = SlidingWindowMLASpec( + block_size=64, + num_kv_heads=1, + head_size=512, + dtype=torch.uint8, + sliding_window=128, + cache_dtype_str="fp8_ds_mla", + alignment=576, + model_version="deepseek_v4", + ) + grouped_specs = [ + UniformTypeKVCacheSpecs( + block_size=256, + kv_cache_specs={f"full_{i}": full_spec for i in range(21)}, + ), + UniformTypeKVCacheSpecs( + block_size=64, + kv_cache_specs={f"swa_{i}": swa_spec for i in range(46)}, + ), + ] + vllm_config = SimpleNamespace( + model_config=SimpleNamespace(max_model_len=1_048_576), + scheduler_config=SimpleNamespace(max_num_batched_tokens=8192), + parallel_config=SimpleNamespace( + decode_context_parallel_size=1, + prefill_context_parallel_size=1, + ), + ) + + assert ( + kv_cache_utils._select_deepseek_v4_tuple_width( + vllm_config, # type: ignore[arg-type] + grouped_specs, + ) + == 21 + ) + groups = kv_cache_utils._get_kv_cache_groups_uniform_groups( + vllm_config, # type: ignore[arg-type] + grouped_specs, + ) + assert [len(group.layer_names) for group in groups] == [21, 16, 15, 15] + + def test_get_kv_cache_spec_kind_unwraps_uniform_type_specs(): uniform_mla_spec = UniformTypeKVCacheSpecs( block_size=16, diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index 546412b1d2..5b1d53cd8a 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -1009,6 +1009,98 @@ def test_prefill_hybrid_model_mamba_align(): manager.free(req0) +def test_mamba_align_offload_handoff_tracks_relocated_mtp_boundary(): + """Expose the committed state block, not its stale logical position. + + A large second prefill chunk relocates the four MTP scratch blocks. The + first scratch block moves from logical position 1 to position 5 and becomes + the state after token 24. An append-only connector mirror would still see + that physical block at position 1 and could save the token-24 state under + the token-8 key. + """ + block_size = 4 + kv_cache_config = KVCacheConfig( + num_blocks=64, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=block_size, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + num_speculative_blocks=4, + ), + ), + ], + ) + manager = KVCacheManager( + kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=block_size, + ) + request = make_request("0", [0] * 24, block_size, sha256) + computed_blocks, num_computed_tokens = manager.get_computed_blocks(request) + + assert ( + manager.allocate_slots( + request, + num_new_tokens=4, + num_new_computed_tokens=num_computed_tokens, + new_computed_blocks=computed_blocks, + ) + is not None + ) + mamba_manager = manager.coordinator.single_type_managers[1] + mamba_blocks = mamba_manager.req_to_blocks[request.request_id] + first_boundary_block_id = mamba_blocks[0].block_id + first_scratch_block = mamba_blocks[1] + + assert manager.take_boundary_state_offloads() == { + request.request_id: [(1, first_boundary_block_id, 4)] + } + + request.num_computed_tokens = 4 + assert manager.allocate_slots(request, num_new_tokens=20) is not None + mamba_blocks = mamba_manager.req_to_blocks[request.request_id] + + assert mamba_blocks[1].is_null + assert mamba_blocks[5] is first_scratch_block + assert manager.take_boundary_state_offloads() == { + request.request_id: [(1, first_scratch_block.block_id, 24)] + } + assert manager.take_boundary_state_offloads() == {} + + manager.free(request) + + # A request freed before the scheduler drains its offer must not leave a + # block ID that can be recycled and later misidentified as its boundary. + freed_request = make_request("freed", [1] * 4, block_size, sha256) + computed_blocks, num_computed_tokens = manager.get_computed_blocks(freed_request) + assert ( + manager.allocate_slots( + freed_request, + num_new_tokens=4, + num_new_computed_tokens=num_computed_tokens, + new_computed_blocks=computed_blocks, + ) + is not None + ) + manager.free(freed_request) + assert manager.take_boundary_state_offloads() == {} + + def test_prefill_plp(): """Test prefill with APC and some prompt logprobs (plp) requests. @@ -2784,7 +2876,8 @@ def test_can_fit_full_sequence_swa_cap_admits_long_prompt(): manager = KVCacheManager( config, max_model_len=max_model_len, - max_num_batched_tokens=max_num_batched_tokens, + # Single (sync) batch in flight, so in-flight tokens == batched tokens. + max_in_flight_tokens=max_num_batched_tokens, enable_caching=True, hash_block_size=block_size, ) @@ -2840,7 +2933,8 @@ def test_can_fit_full_sequence_full_attention_still_gates_oversized(): manager = KVCacheManager( config, max_model_len=max_model_len, - max_num_batched_tokens=max_num_batched_tokens, + # Single (sync) batch in flight, so in-flight tokens == batched tokens. + max_in_flight_tokens=max_num_batched_tokens, enable_caching=True, hash_block_size=block_size, ) diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index db33b4c6df..cfd91b83da 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -23,6 +23,7 @@ from vllm.sampling_params import SamplingParams, StructuredOutputsParams from vllm.utils.hashing import sha256 from vllm.v1.core.encoder_cache_manager import EncoderCacheManager +from vllm.v1.core.kv_cache_manager import KVCacheManager from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput from vllm.v1.core.sched.scheduler import Scheduler @@ -41,6 +42,46 @@ pytestmark = pytest.mark.cpu_test +@pytest.mark.parametrize( + "num_computed_tokens,expected_attention_ids", + [ + (0, []), + (1, [10]), + (16, [10]), + (17, [10, 11]), + (64, [10, 11, 12, 13]), + ], +) +def test_block_ids_for_computed_tokens_exclude_attention_lookahead( + num_computed_tokens: int, + expected_attention_ids: list[int], +): + """Only attention blocks beyond the computed-token frontier are clipped. + + Non-attention state groups retain their allocator-defined block table. + """ + manager = object.__new__(KVCacheManager) + attention_spec = FullAttentionSpec( + block_size=16, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ) + manager.kv_cache_config = Mock( + kv_cache_groups=[ + Mock(kv_cache_spec=attention_spec), + Mock(kv_cache_spec=object()), + ] + ) + manager.get_block_ids = Mock(return_value=([10, 11, 12, 13], [20, 21, 22])) + + block_ids = manager.get_block_ids_for_computed_tokens( + "request", num_computed_tokens + ) + + assert block_ids == (expected_attention_ids, [20, 21, 22]) + + def test_add_requests(): scheduler = create_scheduler() requests = create_requests(num_requests=10) diff --git a/tests/v1/core/test_single_type_kv_cache_manager.py b/tests/v1/core/test_single_type_kv_cache_manager.py index f59830dcd7..3af8c830cb 100644 --- a/tests/v1/core/test_single_type_kv_cache_manager.py +++ b/tests/v1/core/test_single_type_kv_cache_manager.py @@ -14,9 +14,14 @@ ) from vllm.v1.core.single_type_kv_cache_manager import ( ChunkedLocalAttentionManager, + PrefixAnchoredSWAManager, SlidingWindowManager, ) -from vllm.v1.kv_cache_interface import ChunkedLocalAttentionSpec, SlidingWindowSpec +from vllm.v1.kv_cache_interface import ( + ChunkedLocalAttentionSpec, + PrefixAnchoredSWASpec, + SlidingWindowSpec, +) pytestmark = pytest.mark.cpu_test @@ -481,3 +486,103 @@ def test_predictor_matches_allocator_blocks_calculation_with_admission_cap(): f"but allocator pulled {len(new_blocks)}" ) total_computed = num_tokens + + +def test_prefix_anchored_swa_remove_skipped_blocks_gap_range(): + block_size = 4 + spec = PrefixAnchoredSWASpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + decode_sliding_window=8, + ) + block_pool = BlockPool( + num_gpu_blocks=2000, enable_caching=True, hash_block_size=block_size + ) + manager = PrefixAnchoredSWAManager( + spec, + block_pool=block_pool, + enable_caching=True, + kv_cache_group_id=0, + ) + + null_block_id = block_pool.null_block.block_id + original_block_ids = list(range(1000, 1010)) + block_table = [KVCacheBlock(id_) for id_ in original_block_ids] + manager.req_to_blocks["test"] = block_table + + prefix_len = 16 + + # Without num_prompt_tokens, gap blocks are not evicted (full-attention + # default path). + manager.remove_skipped_blocks("test", 28) + assert [b.block_id for b in block_table] == original_block_ids + + # Gap = block 4 only (tokens [16, 20) fall in the gap). + manager.remove_skipped_blocks("test", 28, num_prompt_tokens=prefix_len) + expected = original_block_ids.copy() + expected[4] = null_block_id + assert [b.block_id for b in block_table] == expected + + # Window moves: blocks 5 and 6 also enter the gap; block 4 is already + # null and stays null. + manager.remove_skipped_blocks("test", 36, num_prompt_tokens=prefix_len) + expected[5] = null_block_id + expected[6] = null_block_id + assert [b.block_id for b in block_table] == expected + + # Blocks holding the prefix and the live decode window are never touched. + assert [b.block_id for b in block_table[:4]] == original_block_ids[:4] + assert [b.block_id for b in block_table[7:]] == original_block_ids[7:] + + +def test_prefix_anchored_swa_spec_merge(): + def make_spec(window: int) -> PrefixAnchoredSWASpec: + return PrefixAnchoredSWASpec( + block_size=16, + num_kv_heads=2, + head_size=64, + dtype=torch.float16, + decode_sliding_window=window, + ) + + merged = PrefixAnchoredSWASpec.merge([make_spec(128), make_spec(128)]) + assert isinstance(merged, PrefixAnchoredSWASpec) + assert merged.decode_sliding_window == 128 + assert merged.block_size == 16 + assert merged.num_kv_heads == 2 + assert merged.head_size == 64 + assert merged.head_size_v == 64 + assert merged.dtype == torch.float16 + + with pytest.raises(AssertionError): + PrefixAnchoredSWASpec.merge([make_spec(128), make_spec(64)]) + + +def test_prefix_anchored_swa_manager_registered(): + from vllm.v1.core.single_type_kv_cache_manager import ( + get_manager_for_kv_cache_spec, + spec_manager_map, + ) + + assert spec_manager_map[PrefixAnchoredSWASpec] is PrefixAnchoredSWAManager + + spec = PrefixAnchoredSWASpec( + block_size=16, + num_kv_heads=1, + head_size=64, + dtype=torch.float16, + decode_sliding_window=128, + ) + block_pool = BlockPool(num_gpu_blocks=100, enable_caching=False, hash_block_size=16) + manager = get_manager_for_kv_cache_spec( + spec, + max_in_flight_tokens=1024, + max_model_len=4096, + block_pool=block_pool, + enable_caching=False, + kv_cache_group_id=0, + ) + assert isinstance(manager, PrefixAnchoredSWAManager) + assert manager.decode_sliding_window == 128 diff --git a/tests/v1/core/test_swa_inflight_window_free.py b/tests/v1/core/test_swa_inflight_window_free.py new file mode 100644 index 0000000000..3c66fd533b --- /dev/null +++ b/tests/v1/core/test_swa_inflight_window_free.py @@ -0,0 +1,351 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Out-of-window block frees vs in-flight GPU steps. + +With async scheduling / PP, `num_computed_tokens` optimistically includes +tokens of unprocessed steps whose attention windows still read the blocks +just below the optimistic boundary (and rejected spec tokens can roll it +back), so `allocate_slots` frees on the processed-token basis: +`num_computed_tokens - num_in_flight_tokens`. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm import envs +from vllm.config import VllmConfig +from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.kv_cache_interface import ( + ChunkedLocalAttentionSpec, + PrefixAnchoredSWASpec, + SlidingWindowSpec, +) +from vllm.v1.outputs import ModelRunnerOutput + +from .utils import create_requests, create_scheduler, mock_kv + +NUM_PROMPT_TOKENS = 100 +BLOCK_SIZE = 16 +SLIDING_WINDOW = 16 +# Tokens 0..84 are outside the window of the next token to compute +# (100 - 16 + 1 = 85), i.e. 5 full blocks. +NUM_OUT_OF_WINDOW_BLOCKS = 85 // BLOCK_SIZE +# Chunked-local skips whole chunks left of the current one: +# (100 // 32) * 32 = 96 settled tokens -> 6 full blocks. +CHUNK_SIZE = 32 +NUM_OUT_OF_CHUNK_BLOCKS = (NUM_PROMPT_TOKENS // CHUNK_SIZE) * CHUNK_SIZE // BLOCK_SIZE + + +@pytest.mark.parametrize( + ("pp_size", "async_scheduling", "queue_depth", "expected_batches"), + [ + (1, False, 0, 1), + (1, True, 0, 2), + (1, True, 4, 4), + (3, False, 0, 3), + (3, True, 4, 3), + ], +) +def test_max_in_flight_tokens_matches_executor_concurrency( + monkeypatch, + pp_size: int, + async_scheduling: bool, + queue_depth: int, + expected_batches: int, +): + monkeypatch.setattr(envs, "VLLM_SM70_ASYNC_SCHEDULING_QUEUE_DEPTH", queue_depth) + config = object.__new__(VllmConfig) + config.parallel_config = SimpleNamespace(pipeline_parallel_size=pp_size) + config.scheduler_config = SimpleNamespace( + async_scheduling=async_scheduling, + max_num_batched_tokens=1024, + ) + + assert config.max_concurrent_batches == expected_batches + assert config.max_in_flight_tokens == expected_batches * 1024 + + +def _make_model_runner_output( + scheduler_output: SchedulerOutput, + token_id: int = 0, +) -> ModelRunnerOutput: + req_ids = list(scheduler_output.num_scheduled_tokens.keys()) + return ModelRunnerOutput( + req_ids=req_ids, + req_id_to_index={req_id: i for i, req_id in enumerate(req_ids)}, + sampled_token_ids=[[token_id] for _ in req_ids], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ) + + +def _create_swa_scheduler(async_scheduling: bool): + return create_scheduler( + block_size=BLOCK_SIZE, + async_scheduling=async_scheduling, + kv_cache_spec=SlidingWindowSpec( + block_size=BLOCK_SIZE, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=SLIDING_WINDOW, + ), + ) + + +def _create_chunked_scheduler(async_scheduling: bool): + return create_scheduler( + block_size=BLOCK_SIZE, + async_scheduling=async_scheduling, + kv_cache_spec=ChunkedLocalAttentionSpec( + block_size=BLOCK_SIZE, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + attention_chunk_size=CHUNK_SIZE, + ), + ) + + +def _create_anchored_scheduler(async_scheduling: bool): + return create_scheduler( + block_size=BLOCK_SIZE, + async_scheduling=async_scheduling, + enable_prefix_caching=False, + kv_cache_spec=PrefixAnchoredSWASpec( + block_size=BLOCK_SIZE, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + decode_sliding_window=SLIDING_WINDOW, + ), + ) + + +def _num_null_blocks(scheduler, request_id: str) -> int: + manager = scheduler.kv_cache_manager.coordinator.single_type_managers[0] + null_block = manager._null_block + return sum(1 for b in manager.req_to_blocks[request_id] if b is null_block) + + +def test_num_in_flight_tokens_accounting(): + scheduler = create_scheduler(async_scheduling=True) + request = create_requests(num_requests=1, num_tokens=NUM_PROMPT_TOKENS)[0] + scheduler.add_request(request) + + out0 = scheduler.schedule() + assert request.num_in_flight_tokens == NUM_PROMPT_TOKENS + # Async: decode scheduled before the prefill output is processed. + out1 = scheduler.schedule() + assert request.num_in_flight_tokens == NUM_PROMPT_TOKENS + 1 + + scheduler.update_from_output(out0, _make_model_runner_output(out0)) + assert request.num_in_flight_tokens == 1 + scheduler.update_from_output(out1, _make_model_runner_output(out1)) + assert request.num_in_flight_tokens == 0 + + +def test_swa_free_waits_for_in_flight_step(): + """Async: out-of-window blocks stay allocated until the step that still + reads them has been processed.""" + scheduler = _create_swa_scheduler(async_scheduling=True) + request = create_requests( + num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE + )[0] + scheduler.add_request(request) + req_id = request.request_id + block_pool = scheduler.kv_cache_manager.block_pool + + out0 = scheduler.schedule() # prefill, in flight from here on + free_after_prefill = block_pool.get_num_free_blocks() + + # Decode scheduled while the prefill still reads the out-of-window blocks: + # they must not be freed yet. + out1 = scheduler.schedule() + assert _num_null_blocks(scheduler, req_id) == 0 + assert block_pool.get_num_free_blocks() == free_after_prefill + + # Prefill output processed; the next allocate frees the out-of-window + # blocks. + scheduler.update_from_output(out0, _make_model_runner_output(out0)) + scheduler.schedule() + assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_WINDOW_BLOCKS + assert ( + block_pool.get_num_free_blocks() + == free_after_prefill + NUM_OUT_OF_WINDOW_BLOCKS + ) + # Not double-freed on the following steps. + scheduler.update_from_output(out1, _make_model_runner_output(out1)) + scheduler.schedule() + assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_WINDOW_BLOCKS + + +def test_swa_free_immediate_when_sync(): + """Sync: no in-flight step at schedule time, frees happen at the first + decode allocation as before.""" + scheduler = _create_swa_scheduler(async_scheduling=False) + request = create_requests( + num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE + )[0] + scheduler.add_request(request) + req_id = request.request_id + + out0 = scheduler.schedule() + scheduler.update_from_output(out0, _make_model_runner_output(out0)) + assert request.num_in_flight_tokens == 0 + + scheduler.schedule() + assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_WINDOW_BLOCKS + + +def test_anchored_swa_gap_free_waits_for_in_flight_step(): + """Prefix-anchored SWA (windowed OCR decode): gap blocks between the + prefill tail and the decode window are freed on the processed-token + basis under async scheduling — an in-flight step whose attention still + reads the gap must hold the blocks.""" + scheduler = _create_anchored_scheduler(async_scheduling=True) + request = create_requests( + num_requests=1, + num_tokens=NUM_PROMPT_TOKENS, + block_size=BLOCK_SIZE, + max_tokens=200, + ignore_eos=True, + )[0] + scheduler.add_request(request) + req_id = request.request_id + + # Prefill, then decode until a full gap block exists on the PROCESSED + # basis: gap = [ceil(P/bs), (T - W) // bs) needs T >= 144 for block 7 + # (P=100 -> first_gap_block 7; (144-16)//16 = 8). + out_prev = scheduler.schedule() # prefill (in flight) + scheduler.update_from_output(out_prev, _make_model_runner_output(out_prev)) + for _ in range(44): # T: 100 -> 144, settled step by step + out = scheduler.schedule() + scheduler.update_from_output(out, _make_model_runner_output(out)) + assert request.num_in_flight_tokens == 0 + assert request.num_computed_tokens >= 144 + + # Async over-scheduling: the settled-basis allocate on this schedule may + # evict gap blocks for the SETTLED T. + out_a = scheduler.schedule() + null_settled = _num_null_blocks(scheduler, req_id) + assert null_settled >= 1 # gap eviction engaged at all + + # Second in-flight step: the optimistic boundary advances by 1 token but + # the processed basis does not — no additional frees may happen. + out_b = scheduler.schedule() + assert _num_null_blocks(scheduler, req_id) == null_settled + + scheduler.update_from_output(out_a, _make_model_runner_output(out_a)) + scheduler.update_from_output(out_b, _make_model_runner_output(out_b)) + + +def test_swa_admission_cap_accounts_for_overlapping_batches(): + spec = SlidingWindowSpec( + block_size=16, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=1024, + ) + base = spec.max_admission_blocks_per_request( + max_in_flight_tokens=1024, max_model_len=16384 + ) + # (1024 - 1 + 1024) tokens -> 128 blocks, +1 for window misalignment. + assert base == 129 + overlapped = spec.max_admission_blocks_per_request( + max_in_flight_tokens=2 * 1024, max_model_len=16384 + ) + # One extra in-flight chunk is held back: (1024 - 1 + 2 * 1024) tokens. + assert overlapped == 193 + + +def test_chunked_local_free_waits_for_in_flight_step(): + """Chunked-local attention frees whole chunks left of the current one, and + is exposed to the same load-WAR: with async scheduling those chunks must + stay allocated until the in-flight step that still reads them settles.""" + scheduler = _create_chunked_scheduler(async_scheduling=True) + request = create_requests( + num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE + )[0] + scheduler.add_request(request) + req_id = request.request_id + block_pool = scheduler.kv_cache_manager.block_pool + + out0 = scheduler.schedule() # prefill, in flight from here on + free_after_prefill = block_pool.get_num_free_blocks() + + # Decode scheduled while the prefill still reads the out-of-chunk blocks. + out1 = scheduler.schedule() + assert _num_null_blocks(scheduler, req_id) == 0 + assert block_pool.get_num_free_blocks() == free_after_prefill + + # Prefill output processed; the next allocate frees the out-of-chunk blocks. + scheduler.update_from_output(out0, _make_model_runner_output(out0)) + scheduler.schedule() + assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_CHUNK_BLOCKS + assert ( + block_pool.get_num_free_blocks() == free_after_prefill + NUM_OUT_OF_CHUNK_BLOCKS + ) + # Not double-freed on the following steps. + scheduler.update_from_output(out1, _make_model_runner_output(out1)) + scheduler.schedule() + assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_CHUNK_BLOCKS + + +def test_chunked_local_admission_cap_accounts_for_overlapping_batches(): + spec = ChunkedLocalAttentionSpec( + block_size=16, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + attention_chunk_size=1024, + ) + base = spec.max_admission_blocks_per_request( + max_in_flight_tokens=1024, max_model_len=16384 + ) + # (1024 + 1024) tokens -> 128 blocks. + assert base == 128 + overlapped = spec.max_admission_blocks_per_request( + max_in_flight_tokens=2 * 1024, max_model_len=16384 + ) + # One extra in-flight chunk is held back: (1024 + 2 * 1024) tokens. + assert overlapped == 192 + + +def test_connector_finish_frees_on_settled_basis(): + """The out-of-window prune done at request finish, before the block table + is handed to a KV connector, must use the same processed-token basis.""" + scheduler = create_scheduler( + block_size=BLOCK_SIZE, + async_scheduling=True, + use_kv_connector=mock_kv(matched_tokens=0, is_async=False), + kv_cache_spec=SlidingWindowSpec( + block_size=BLOCK_SIZE, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=SLIDING_WINDOW, + ), + ) + request = create_requests( + num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE + )[0] + scheduler.add_request(request) + req_id = request.request_id + + out0 = scheduler.schedule() # prefill, in flight from here on + scheduler.schedule() # decode over-scheduled: num_computed_tokens optimistic + + # Finishing now (connector store) must NOT prune out-of-window blocks the + # in-flight prefill still reads. + scheduler._connector_finished(request) + assert _num_null_blocks(scheduler, req_id) == 0 + + # Once the in-flight step settles, the same prune releases them. + scheduler.update_from_output(out0, _make_model_runner_output(out0)) + scheduler._connector_finished(request) + assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_WINDOW_BLOCKS diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py index 2d9834d2e3..472c1876f4 100644 --- a/tests/v1/core/utils.py +++ b/tests/v1/core/utils.py @@ -28,6 +28,7 @@ FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, + KVCacheSpec, ) from vllm.v1.request import Request from vllm.v1.structured_output import StructuredOutputManager @@ -57,6 +58,7 @@ def create_scheduler( pipeline_parallel_size: int = 1, use_ec_connector: bool = False, ec_role: str | None = None, + kv_cache_spec: KVCacheSpec | None = None, ) -> Scheduler | AsyncScheduler: """Create scheduler under test. @@ -144,20 +146,17 @@ def create_scheduler( speculative_config=speculative_config, ec_transfer_config=ec_transfer_config, ) + if kv_cache_spec is None: + kv_cache_spec = FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ) kv_cache_config = KVCacheConfig( num_blocks=num_blocks, # A large number of blocks to hold all requests kv_cache_tensors=[], - kv_cache_groups=[ - KVCacheGroupSpec( - ["layer"], - FullAttentionSpec( - block_size=block_size, - num_kv_heads=1, - head_size=1, - dtype=torch.float32, - ), - ) - ], + kv_cache_groups=[KVCacheGroupSpec(["layer"], kv_cache_spec)], ) cache_config.num_gpu_blocks = num_blocks scheduler_cls = AsyncScheduler if async_scheduling else Scheduler diff --git a/tests/v1/cudagraph/test_cudagraph_dispatch.py b/tests/v1/cudagraph/test_cudagraph_dispatch.py index 0c959f6b71..e78c865a8c 100644 --- a/tests/v1/cudagraph/test_cudagraph_dispatch.py +++ b/tests/v1/cudagraph/test_cudagraph_dispatch.py @@ -19,7 +19,12 @@ VllmConfig, ) from vllm.config.lora import LoRAConfig -from vllm.forward_context import BatchDescriptor, set_forward_context +from vllm.forward_context import ( + CUDAGRAPH_VARIANT_DEFAULT, + CUDAGRAPH_VARIANT_LONG_CONTEXT, + BatchDescriptor, + set_forward_context, +) from vllm.platforms import current_platform from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher @@ -373,7 +378,9 @@ def test_dsv4_decode_context_bucket_is_derived_on_sm70(self, monkeypatch): cudagraph_capture_sizes=[1, 2], ) config = _create_vllm_config(comp_config, max_num_seqs=2) - config.model_config.architectures = ["DeepseekV4ForCausalLM"] + # Route by the compressed-index operator contract, not architecture or + # checkpoint identity. + config.model_config.architectures = ["UnrelatedForCausalLM"] config.model_config.max_model_len = 4096 config.model_config.hf_config.index_topk = 512 config.model_config.hf_config.compress_ratios = [128, 4, 128] @@ -401,6 +408,119 @@ def test_dsv4_decode_context_bucket_is_derived_on_sm70(self, monkeypatch): disabled = CudagraphDispatcher(config) assert not disabled.sm70_dsv4_decode_context_buckets + def test_dsv4_context_bucket_requires_explicit_override_for_mtp_on_sm70( + self, monkeypatch + ): + monkeypatch.delenv("VLLM_SM70_DSV4_DECODE_CONTEXT_BUCKETS", raising=False) + monkeypatch.delenv("VLLM_SM70_MTP_CONTEXT_BUCKETS", raising=False) + comp_config = CompilationConfig( + cudagraph_mode="FULL_DECODE_ONLY", + mode=CompilationMode.NONE, + cudagraph_capture_sizes=[8, 16], + ) + config = _create_vllm_config(comp_config, max_num_seqs=2) + config.speculative_config = MagicMock(num_speculative_tokens=7) + config.model_config.architectures = ["DeepseekV4ForCausalLM"] + config.model_config.max_model_len = 4096 + config.model_config.hf_config.index_topk = 512 + config.model_config.hf_config.compress_ratios = [128, 4, 128] + + with ( + patch.object(current_platform, "is_cuda", return_value=True), + patch.object(current_platform, "is_device_capability", return_value=True), + ): + default_dispatcher = CudagraphDispatcher(config) + assert not default_dispatcher.has_attention_context_buckets + + monkeypatch.setenv("VLLM_SM70_DSV4_DECODE_CONTEXT_BUCKETS", "2048") + dispatcher = CudagraphDispatcher(config) + dispatcher.initialize_cudagraph_keys( + cudagraph_mode=comp_config.cudagraph_mode, + uniform_decode_query_len=8, + ) + + bounded = BatchDescriptor( + num_tokens=8, + num_reqs=1, + uniform=True, + attention_context_bucket=2048, + ) + assert dispatcher.has_attention_context_buckets + assert bounded in dispatcher.cudagraph_keys[CUDAGraphMode.FULL] + + mode, desc = dispatcher.dispatch( + num_tokens=8, + uniform_decode=True, + attention_context_len=1024, + ) + assert mode == CUDAGraphMode.FULL + assert desc == bounded + + monkeypatch.setenv("VLLM_SM70_MTP_CONTEXT_BUCKETS", "") + explicitly_disabled = CudagraphDispatcher(config) + explicitly_disabled.initialize_cudagraph_keys( + cudagraph_mode=comp_config.cudagraph_mode, + uniform_decode_query_len=8, + ) + assert not explicitly_disabled.has_attention_context_buckets + + def test_dsv4_long_context_buckets_bound_mtp_graph_width(self, monkeypatch): + monkeypatch.setenv( + "VLLM_SM70_DSV4_DECODE_CONTEXT_BUCKETS", + "2048,4096,16384,65536,131072", + ) + monkeypatch.delenv("VLLM_SM70_MTP_CONTEXT_BUCKETS", raising=False) + comp_config = CompilationConfig( + cudagraph_mode="FULL_DECODE_ONLY", + mode=CompilationMode.NONE, + cudagraph_capture_sizes=[8, 16], + ) + config = _create_vllm_config(comp_config, max_num_seqs=2) + config.speculative_config = MagicMock(num_speculative_tokens=7) + config.model_config.architectures = ["DeepseekV4ForCausalLM"] + config.model_config.max_model_len = 262144 + config.model_config.hf_config.index_topk = 512 + config.model_config.hf_config.compress_ratios = [128, 4, 128] + + with ( + patch.object(current_platform, "is_cuda", return_value=True), + patch.object(current_platform, "is_device_capability", return_value=True), + ): + dispatcher = CudagraphDispatcher(config) + dispatcher.initialize_cudagraph_keys( + cudagraph_mode=comp_config.cudagraph_mode, + uniform_decode_query_len=8, + ) + + assert dispatcher.sm70_dsv4_decode_context_buckets == ( + 2048, + 4096, + 16384, + 65536, + 131072, + ) + expected_buckets = { + 3000: 4096, + 5000: 16384, + 20000: 65536, + 100000: 131072, + } + for context_len, expected_bucket in expected_buckets.items(): + mode, descriptor = dispatcher.dispatch( + num_tokens=8, + uniform_decode=True, + attention_context_len=context_len, + ) + assert mode == CUDAGraphMode.FULL + assert descriptor.attention_context_bucket == expected_bucket + + _, unbounded = dispatcher.dispatch( + num_tokens=8, + uniform_decode=True, + attention_context_len=200000, + ) + assert unbounded.attention_context_bucket is None + def test_fp8_e5m2_decode_context_bucket_is_shape_derived_on_sm70(self, monkeypatch): monkeypatch.delenv("VLLM_SM70_FP8_KV_DECODE_CONTEXT_BUCKETS", raising=False) comp_config = CompilationConfig( @@ -422,6 +542,8 @@ def test_fp8_e5m2_decode_context_bucket_is_shape_derived_on_sm70(self, monkeypat patch("vllm.v1.cudagraph_dispatcher.envs") as mock_envs, ): mock_envs.VLLM_SM70_FLASH_ATTN_V100 = True + mock_envs.VLLM_FLASH_V100_XQA_BATCH_CONTEXT_ROUTING = False + mock_envs.VLLM_FLASH_V100_DECODE_PARTITION_SIZE = None dispatcher = CudagraphDispatcher(config) dispatcher.initialize_cudagraph_keys( cudagraph_mode=comp_config.cudagraph_mode, @@ -453,6 +575,312 @@ def test_fp8_e5m2_decode_context_bucket_is_shape_derived_on_sm70(self, monkeypat assert mode == CUDAGraphMode.FULL assert desc == BatchDescriptor(num_tokens=1, num_reqs=1, uniform=True) + @pytest.mark.parametrize( + ("num_tokens", "context_len", "expected_variant"), + ( + (4, 8191, CUDAGRAPH_VARIANT_DEFAULT), + (4, 12288, CUDAGRAPH_VARIANT_DEFAULT), + (4, 16383, CUDAGRAPH_VARIANT_DEFAULT), + (4, 16384, CUDAGRAPH_VARIANT_LONG_CONTEXT), + (8, 4095, CUDAGRAPH_VARIANT_DEFAULT), + (8, 12288, CUDAGRAPH_VARIANT_DEFAULT), + (8, 16383, CUDAGRAPH_VARIANT_DEFAULT), + (8, 16384, CUDAGRAPH_VARIANT_LONG_CONTEXT), + (16, 1024, CUDAGRAPH_VARIANT_DEFAULT), + (16, 16001, CUDAGRAPH_VARIANT_DEFAULT), + (16, 16383, CUDAGRAPH_VARIANT_DEFAULT), + (16, 16384, CUDAGRAPH_VARIANT_LONG_CONTEXT), + ), + ) + def test_fp8_e5m2_batch_context_graph_routing_on_sm70( + self, + monkeypatch, + num_tokens, + context_len, + expected_variant, + ): + monkeypatch.delenv("VLLM_SM70_FP8_KV_DECODE_CONTEXT_BUCKETS", raising=False) + comp_config = CompilationConfig( + cudagraph_mode="FULL_DECODE_ONLY", + mode=CompilationMode.NONE, + cudagraph_capture_sizes=[1, 2, 4, 8, 16], + ) + config = _create_vllm_config(comp_config, max_num_seqs=16) + config.cache_config.cache_dtype = "fp8_e5m2" + config.attention_config.backend = None + config.model_config.max_model_len = 262144 + config.model_config.hf_text_config.num_attention_heads = 24 + config.model_config.hf_text_config.num_key_value_heads = 4 + config.model_config.hf_text_config.head_dim = 256 + + with ( + patch.object(current_platform, "is_cuda", return_value=True), + patch.object(current_platform, "is_device_capability", return_value=True), + patch("vllm.v1.cudagraph_dispatcher.envs") as mock_envs, + ): + mock_envs.VLLM_SM70_FLASH_ATTN_V100 = True + mock_envs.VLLM_FLASH_V100_XQA_BATCH_CONTEXT_ROUTING = True + mock_envs.VLLM_FLASH_V100_DECODE_PARTITION_SIZE = None + dispatcher = CudagraphDispatcher(config) + dispatcher.initialize_cudagraph_keys( + cudagraph_mode=comp_config.cudagraph_mode, + uniform_decode_query_len=1, + ) + + mode, desc = dispatcher.dispatch( + num_tokens=num_tokens, + uniform_decode=True, + attention_context_len=context_len, + ) + + assert dispatcher.sm70_fp8_kv_batch_context_routing + assert dispatcher.has_attention_context_specialization + assert mode == CUDAGraphMode.FULL + assert desc == BatchDescriptor( + num_tokens=num_tokens, + num_reqs=num_tokens, + uniform=True, + graph_variant=expected_variant, + ) + assert ( + BatchDescriptor( + num_tokens=num_tokens, + num_reqs=num_tokens, + uniform=True, + graph_variant=CUDAGRAPH_VARIANT_LONG_CONTEXT, + ) + in dispatcher.cudagraph_keys[CUDAGraphMode.FULL] + ) + + def test_fp8_e5m2_long_context_variants_capture_after_baseline(self, monkeypatch): + comp_config = CompilationConfig( + cudagraph_mode="FULL_DECODE_ONLY", + mode=CompilationMode.NONE, + cudagraph_capture_sizes=[1, 2, 4, 8, 16], + ) + config = _create_vllm_config(comp_config, max_num_seqs=16) + config.cache_config.cache_dtype = "fp8_e5m2" + config.attention_config.backend = None + config.model_config.max_model_len = 262144 + config.model_config.hf_text_config.num_attention_heads = 24 + config.model_config.hf_text_config.num_key_value_heads = 4 + config.model_config.hf_text_config.head_dim = 256 + + with ( + patch.object(current_platform, "is_cuda", return_value=True), + patch.object(current_platform, "is_device_capability", return_value=True), + patch("vllm.v1.cudagraph_dispatcher.envs") as mock_envs, + ): + mock_envs.VLLM_SM70_FLASH_ATTN_V100 = True + mock_envs.VLLM_FLASH_V100_XQA_BATCH_CONTEXT_ROUTING = True + mock_envs.VLLM_FLASH_V100_DECODE_PARTITION_SIZE = None + dispatcher = CudagraphDispatcher(config) + dispatcher.initialize_cudagraph_keys( + cudagraph_mode=comp_config.cudagraph_mode, + uniform_decode_query_len=1, + ) + + full = dict(dispatcher.get_capture_descs())[CUDAGraphMode.FULL] + variants = [desc.graph_variant for desc in full] + first_long = variants.index(CUDAGRAPH_VARIANT_LONG_CONTEXT) + assert all( + variant == CUDAGRAPH_VARIANT_DEFAULT for variant in variants[:first_long] + ) + assert variants[first_long:] == [CUDAGRAPH_VARIANT_LONG_CONTEXT] * 3 + + @pytest.mark.parametrize( + ("e4m3_enabled", "expected_routing"), + ((True, True), (False, False)), + ) + def test_fp8_e4m3_batch_context_graph_routing_on_sm70( + self, + monkeypatch, + e4m3_enabled, + expected_routing, + ): + comp_config = CompilationConfig( + cudagraph_mode="FULL_DECODE_ONLY", + mode=CompilationMode.NONE, + cudagraph_capture_sizes=[1, 2, 4, 8, 16], + ) + config = _create_vllm_config(comp_config, max_num_seqs=16) + config.cache_config.cache_dtype = "fp8_e4m3" + config.attention_config.backend = None + config.model_config.max_model_len = 32768 + config.model_config.hf_text_config.num_attention_heads = 24 + config.model_config.hf_text_config.num_key_value_heads = 4 + config.model_config.hf_text_config.head_dim = 256 + + with ( + patch.object(current_platform, "is_cuda", return_value=True), + patch.object(current_platform, "is_device_capability", return_value=True), + patch("vllm.v1.cudagraph_dispatcher.envs") as mock_envs, + ): + mock_envs.VLLM_SM70_FLASH_ATTN_V100 = True + mock_envs.VLLM_FLASH_V100_XQA_BATCH_CONTEXT_ROUTING = True + mock_envs.VLLM_FLASH_V100_DECODE_PARTITION_SIZE = None + mock_envs.VLLM_FLASH_V100_E4M3_BATCH_XQA = e4m3_enabled + dispatcher = CudagraphDispatcher(config) + dispatcher.initialize_cudagraph_keys( + cudagraph_mode=comp_config.cudagraph_mode, + uniform_decode_query_len=1, + ) + + mode, desc = dispatcher.dispatch( + num_tokens=8, + uniform_decode=True, + attention_context_len=16384, + ) + + assert dispatcher.sm70_fp8_kv_batch_context_routing is expected_routing + assert mode == CUDAGraphMode.FULL + assert desc == BatchDescriptor( + num_tokens=8, + num_reqs=8, + uniform=True, + graph_variant=( + CUDAGRAPH_VARIANT_LONG_CONTEXT + if expected_routing + else CUDAGRAPH_VARIANT_DEFAULT + ), + ) + + @pytest.mark.parametrize( + ("auto_enabled", "wave_enabled", "expected_variant"), + ( + (True, True, CUDAGRAPH_VARIANT_LONG_CONTEXT), + (False, True, CUDAGRAPH_VARIANT_DEFAULT), + (True, False, CUDAGRAPH_VARIANT_DEFAULT), + ), + ) + def test_fp8_e4m3_b1_wave_graph_routing_on_sm70( + self, + monkeypatch, + auto_enabled, + wave_enabled, + expected_variant, + ): + monkeypatch.delenv("VLLM_SM70_FP8_KV_DECODE_CONTEXT_BUCKETS", raising=False) + monkeypatch.setenv( + "VLLM_FLASH_V100_XQA_E4M3_G6_P64_P256_AUTO", + "1" if auto_enabled else "0", + ) + monkeypatch.setenv( + "VLLM_FLASH_V100_XQA_E4M3_G6_WAVE_PARTITIONS", + "1" if wave_enabled else "0", + ) + monkeypatch.setenv("VLLM_FLASH_V100_XQA_E4M3_G6_P512_BEGIN", "49152") + comp_config = CompilationConfig( + cudagraph_mode="FULL_DECODE_ONLY", + mode=CompilationMode.NONE, + cudagraph_capture_sizes=[1, 2, 4, 8, 16], + ) + config = _create_vllm_config(comp_config, max_num_seqs=16) + config.cache_config.cache_dtype = "fp8_e4m3" + config.attention_config.backend = None + config.model_config.max_model_len = 262144 + config.model_config.hf_text_config.num_attention_heads = 24 + config.model_config.hf_text_config.num_key_value_heads = 4 + config.model_config.hf_text_config.head_dim = 256 + + with ( + patch.object(current_platform, "is_cuda", return_value=True), + patch.object(current_platform, "is_device_capability", return_value=True), + patch("vllm.v1.cudagraph_dispatcher.envs") as mock_envs, + ): + mock_envs.VLLM_SM70_FLASH_ATTN_V100 = True + mock_envs.VLLM_FLASH_V100_XQA_BATCH_CONTEXT_ROUTING = True + mock_envs.VLLM_FLASH_V100_DECODE_PARTITION_SIZE = None + mock_envs.VLLM_FLASH_V100_E4M3_BATCH_XQA = True + dispatcher = CudagraphDispatcher(config) + dispatcher.initialize_cudagraph_keys( + cudagraph_mode=comp_config.cudagraph_mode, + uniform_decode_query_len=1, + ) + + mode, below = dispatcher.dispatch( + num_tokens=1, + uniform_decode=True, + attention_context_len=49151, + ) + assert mode == CUDAGraphMode.FULL + assert below.graph_variant == CUDAGRAPH_VARIANT_DEFAULT + + mode, at_boundary = dispatcher.dispatch( + num_tokens=1, + uniform_decode=True, + attention_context_len=49152, + ) + assert mode == CUDAGraphMode.FULL + assert at_boundary.graph_variant == expected_variant + assert dispatcher.sm70_e4m3_b1_wave_context_routing is ( + expected_variant == CUDAGRAPH_VARIANT_LONG_CONTEXT + ) + + def test_fp8_e4m3_batch_context_graph_routing_rejects_flashinfer(self, monkeypatch): + comp_config = CompilationConfig( + cudagraph_mode="FULL_DECODE_ONLY", + mode=CompilationMode.NONE, + cudagraph_capture_sizes=[1, 2, 4, 8, 16], + ) + config = _create_vllm_config(comp_config, max_num_seqs=16) + config.cache_config.cache_dtype = "fp8_e4m3" + config.attention_config.backend = "FLASHINFER_SM70" + config.model_config.max_model_len = 32768 + config.model_config.hf_text_config.num_attention_heads = 24 + config.model_config.hf_text_config.num_key_value_heads = 4 + config.model_config.hf_text_config.head_dim = 256 + + with ( + patch.object(current_platform, "is_cuda", return_value=True), + patch.object(current_platform, "is_device_capability", return_value=True), + patch("vllm.v1.cudagraph_dispatcher.envs") as mock_envs, + ): + mock_envs.VLLM_SM70_FLASH_ATTN_V100 = True + mock_envs.VLLM_FLASH_V100_XQA_BATCH_CONTEXT_ROUTING = True + mock_envs.VLLM_FLASH_V100_DECODE_PARTITION_SIZE = None + mock_envs.VLLM_FLASH_V100_E4M3_BATCH_XQA = True + dispatcher = CudagraphDispatcher(config) + + assert not dispatcher.sm70_fp8_kv_batch_context_routing + + def test_fp8_e5m2_batch_context_graph_routing_rollback(self, monkeypatch): + comp_config = CompilationConfig( + cudagraph_mode="FULL_DECODE_ONLY", + mode=CompilationMode.NONE, + cudagraph_capture_sizes=[1, 2, 4, 8, 16], + ) + config = _create_vllm_config(comp_config, max_num_seqs=16) + config.cache_config.cache_dtype = "fp8_e5m2" + config.model_config.max_model_len = 262144 + config.model_config.hf_text_config.num_attention_heads = 24 + config.model_config.hf_text_config.num_key_value_heads = 4 + config.model_config.hf_text_config.head_dim = 256 + + with ( + patch.object(current_platform, "is_cuda", return_value=True), + patch.object(current_platform, "is_device_capability", return_value=True), + patch("vllm.v1.cudagraph_dispatcher.envs") as mock_envs, + ): + mock_envs.VLLM_SM70_FLASH_ATTN_V100 = True + mock_envs.VLLM_FLASH_V100_XQA_BATCH_CONTEXT_ROUTING = False + mock_envs.VLLM_FLASH_V100_DECODE_PARTITION_SIZE = None + dispatcher = CudagraphDispatcher(config) + dispatcher.initialize_cudagraph_keys( + cudagraph_mode=comp_config.cudagraph_mode, + uniform_decode_query_len=1, + ) + + mode, desc = dispatcher.dispatch( + num_tokens=8, + uniform_decode=True, + attention_context_len=4096, + ) + + assert not dispatcher.sm70_fp8_kv_batch_context_routing + assert mode == CUDAGraphMode.FULL + assert desc == BatchDescriptor(num_tokens=8, num_reqs=8, uniform=True) + @pytest.mark.parametrize("cache_dtype", ("auto", "fp8_e4m3")) def test_fp8_e5m2_decode_context_bucket_does_not_match_other_cache_dtypes( self, monkeypatch, cache_dtype @@ -490,6 +918,35 @@ def test_fp8_e5m2_decode_context_bucket_explicit_empty_disables(self, monkeypatc assert not dispatcher.sm70_fp8_kv_decode_context_buckets +def test_negative_graph_variant_uses_an_isolated_pool(): + wrapper = object.__new__(CUDAGraphWrapper) + wrapper.graph_pool = object() + wrapper.concrete_cudagraph_entries = {} + wrapper.semantic_variant_graph_pools = {} + default_desc = BatchDescriptor( + num_tokens=4, + graph_variant=CUDAGRAPH_VARIANT_DEFAULT, + ) + long_context_desc = replace( + default_desc, + graph_variant=CUDAGRAPH_VARIANT_LONG_CONTEXT, + ) + isolated_pool = object() + + with patch.object( + current_platform, + "graph_pool_handle", + return_value=isolated_pool, + ) as graph_pool_handle: + assert wrapper._graph_pool_for_descriptor(default_desc) is wrapper.graph_pool + assert wrapper._graph_pool_for_descriptor(long_context_desc) is isolated_pool + assert wrapper._graph_pool_for_descriptor(long_context_desc) is isolated_pool + + graph_pool_handle.assert_called_once_with() + wrapper.clear_graphs() + assert not wrapper.semantic_variant_graph_pools + + @pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda") class TestCUDAGraphWrapper: def setup_method(self): diff --git a/tests/v1/cudagraph/test_sm70_mtp_split_cudagraphs.py b/tests/v1/cudagraph/test_sm70_mtp_split_cudagraphs.py new file mode 100644 index 0000000000..23c58085e8 --- /dev/null +++ b/tests/v1/cudagraph/test_sm70_mtp_split_cudagraphs.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch + +from vllm.config import CompilationConfig, CUDAGraphMode +from vllm.v1.worker.gpu.cudagraph_utils import CudaGraphManager + + +def _make_config(max_num_seqs: int, verifier_sizes: list[int]): + speculative_config = SimpleNamespace( + method="mtp", + num_speculative_state_tokens=lambda: 4, + ) + return SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_seqs=max_num_seqs), + parallel_config=SimpleNamespace( + tensor_parallel_size=4, + data_parallel_size=1, + ), + compilation_config=CompilationConfig( + cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE, + cudagraph_capture_sizes=verifier_sizes.copy(), + max_cudagraph_capture_size=max(verifier_sizes), + ), + speculative_config=speculative_config, + ) + + +@pytest.mark.parametrize( + ("max_num_seqs", "request_sizes", "verifier_sizes"), + [ + (16, [1, 2, 4, 6, 8, 12, 16], [5, 10, 20, 30, 40, 60, 80]), + (32, [1, 2, 4, 6, 8, 12, 16, 32], [5, 10, 20, 30, 40, 60, 80, 160]), + ], +) +def test_split_managers_keep_exact_full_graph_shapes( + monkeypatch, + max_num_seqs: int, + request_sizes: list[int], + verifier_sizes: list[int], +): + monkeypatch.setenv("VLLM_SM70_MTP_SPLIT_DRAFT_CUDAGRAPHS", "1") + monkeypatch.setattr( + "vllm.v1.worker.gpu.cudagraph_utils.current_platform.is_cuda", + lambda: True, + ) + monkeypatch.setattr( + "vllm.v1.worker.gpu.cudagraph_utils.current_platform.is_device_capability", + lambda capability: capability == (7, 0), + ) + 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 = _make_config(max_num_seqs, verifier_sizes) + shared_sizes = config.compilation_config.cudagraph_capture_sizes.copy() + shared_max_size = config.compilation_config.max_cudagraph_capture_size + target = CudaGraphManager( + config, + torch.device("cpu"), + CUDAGraphMode.FULL_DECODE_ONLY, + decode_query_len=5, + ) + draft = CudaGraphManager( + config, + torch.device("cpu"), + CUDAGraphMode.FULL_DECODE_ONLY, + decode_query_len=1, + ) + + assert target._capture_sizes == verifier_sizes + assert draft._capture_sizes == request_sizes + assert config.compilation_config.cudagraph_capture_sizes == shared_sizes + assert config.compilation_config.max_cudagraph_capture_size == shared_max_size + + target._graphs_captured = True + draft._graphs_captured = True + for batch_size in request_sizes: + target_desc = target.dispatch(batch_size, batch_size * 5, 5) + draft_desc = draft.dispatch(batch_size, batch_size, 1) + assert target_desc.cg_mode == CUDAGraphMode.FULL + assert target_desc.num_tokens == batch_size * 5 + assert target_desc.num_reqs == batch_size + assert draft_desc.cg_mode == CUDAGraphMode.FULL + assert draft_desc.num_tokens == batch_size + assert draft_desc.num_reqs == batch_size diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 5cc7065c1a..f377f244b2 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Iterable +from types import SimpleNamespace from unittest.mock import MagicMock import pytest @@ -13,15 +14,21 @@ from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID from vllm.distributed.kv_events import BlockRemoved, BlockStored from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import ( + GroupOffloadConfig, OffloadingConnectorScheduler, + RequestOffloadState, + SchedulerOffloadConfig, ) from vllm.v1.core.kv_cache_utils import BlockHash from vllm.v1.kv_cache_interface import ( FullAttentionSpec, + KVCacheConfig, KVCacheGroupSpec, + MambaSpec, SlidingWindowSpec, ) from vllm.v1.kv_offload.base import ( + GPULoadStoreSpec, OffloadingEvent, OffloadingManager, OffloadPolicy, @@ -806,6 +813,192 @@ def test_request_level_policy_stores_all_blocks(request_runner, async_scheduling # --------------------------------------------------------------------------- +def _make_exact_boundary_scheduler() -> OffloadingConnectorScheduler: + scheduler = object.__new__(OffloadingConnectorScheduler) + scheduler.config = SchedulerOffloadConfig( + kv_group_configs=( + GroupOffloadConfig( + group_idx=0, + gpu_block_size=16, + offloaded_block_size=16, + hash_block_size_factor=4, + sliding_window_size_in_blocks=None, + ), + GroupOffloadConfig( + group_idx=1, + gpu_block_size=16, + offloaded_block_size=16, + hash_block_size_factor=4, + sliding_window_size_in_blocks=1, + requires_exact_boundary_source=True, + ), + ), + block_size_factor=1, + num_workers=1, + offload_prompt_only=False, + ) + scheduler.manager = MagicMock(spec=OffloadingManager) + scheduler.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + request = MagicMock() + request.request_id = "req" + request.kv_transfer_params = None + request.block_hashes = [BlockHash(f"h{i}".encode()) for i in range(4)] + request.num_computed_tokens = 0 + request.num_prompt_tokens = 16 + request.num_tokens = 16 + request.is_finished.return_value = False + req_context = ReqContext(req_id=request.request_id) + req_status = RequestOffloadState( + config=scheduler.config, + req=request, + req_context=req_context, + offloading_context=RequestOffloadingContext(), + ) + req_status.update_offload_keys() + req_status.group_states[0].block_ids[:] = [11] + # This is the connector's stale append-only view of the Mamba table. + req_status.group_states[1].block_ids[:] = [21] + + scheduler._req_status = {request.request_id: req_status} + scheduler._current_batch_load_jobs = {} + scheduler._current_batch_jobs_to_flush = set() + scheduler._blocks_being_loaded = set() + scheduler._job_counter = 0 + scheduler._stale_job_threshold = 0 + scheduler._jobs = {} + scheduler._block_id_to_pending_jobs = {} + return scheduler + + +def test_mamba_align_group_requires_exact_boundary_source(): + block_size = 16 + kv_cache_config = KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=block_size, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + ) + spec = SimpleNamespace( + vllm_config=SimpleNamespace(parallel_config=SimpleNamespace(world_size=4)), + kv_cache_config=kv_cache_config, + gpu_block_size=(block_size, block_size), + block_size_factor=1, + hash_block_size=block_size, + offload_prompt_only=False, + ) + + config = SchedulerOffloadConfig.from_spec(spec) + + assert not config.kv_group_configs[0].requires_exact_boundary_source + assert config.kv_group_configs[1].requires_exact_boundary_source + + +def _empty_store_output(**overrides): + output = SimpleNamespace( + scheduled_new_reqs=[], + scheduled_cached_reqs=SimpleNamespace( + req_ids=[], new_block_ids=[], resumed_req_ids=set() + ), + num_scheduled_tokens={}, + finished_req_ids=set(), + preempted_req_ids=set(), + boundary_state_offloads=None, + ) + for name, value in overrides.items(): + setattr(output, name, value) + return output + + +def test_mamba_boundary_store_uses_exact_source_not_positional_mirror(): + scheduler = _make_exact_boundary_scheduler() + output = _empty_store_output( + num_scheduled_tokens={"req": 16}, + boundary_state_offloads={"req": [(1, 99, 16)]}, + ) + + jobs = scheduler._build_boundary_state_store_jobs(output) + + assert len(jobs) == 1 + [job_id] = jobs + src_spec, _ = jobs[job_id].transfer_spec + assert isinstance(src_spec, GPULoadStoreSpec) + assert src_spec.block_ids.tolist() == [99] + assert src_spec.block_ids.tolist() != [21] + assert src_spec.group_sizes == [0, 1] + assert src_spec.block_indices == [0, 0] + assert scheduler._block_id_to_pending_jobs == {99: {job_id}} + + +def test_normal_store_excludes_positional_mamba_align_source(): + scheduler = _make_exact_boundary_scheduler() + output = _empty_store_output( + scheduled_new_reqs=[SimpleNamespace(req_id="req", block_ids=None)], + num_scheduled_tokens={"req": 16}, + ) + + jobs = scheduler._build_store_jobs(output) + + assert len(jobs) == 1 + [job] = jobs.values() + src_spec, _ = job.transfer_spec + assert isinstance(src_spec, GPULoadStoreSpec) + assert src_spec.block_ids.tolist() == [11] + assert src_spec.group_sizes == [1, 0] + assert 21 not in src_spec.block_ids + + +def test_mamba_boundary_store_flushes_before_source_block_reuse(): + scheduler = _make_exact_boundary_scheduler() + first = _empty_store_output( + num_scheduled_tokens={"req": 16}, + boundary_state_offloads={"req": [(1, 99, 16)]}, + ) + first_meta = scheduler.build_connector_meta(first) + [job_id] = first_meta.store_jobs + + second = _empty_store_output( + scheduled_cached_reqs=SimpleNamespace( + req_ids=["req"], + new_block_ids=[([12], [99])], + resumed_req_ids=set(), + ), + num_scheduled_tokens={"req": 0}, + ) + second_meta = scheduler.build_connector_meta(second) + + assert second_meta.jobs_to_flush == {job_id} + + +def test_pending_job_cleanup_deduplicates_block_ids(): + """MTP lookahead may alias a physical block within one completed job.""" + scheduler = object.__new__(OffloadingConnectorScheduler) + scheduler._block_id_to_pending_jobs = {7: {42}} + + scheduler._remove_pending_job(42, [7, 7]) + + assert scheduler._block_id_to_pending_jobs == {} + + def test_loads_do_not_populate_fence_index(request_runner): """Loads don't populate _block_id_to_pending_jobs (protected by delay_free_blocks while in flight).""" @@ -960,8 +1153,8 @@ def setup(r, max_offload_tokens): token_ids=[0] * offloaded_block_size * 3, kv_transfer_params={"max_offload_tokens": max_offload_tokens}, ) - r.manager.prepare_store.side_effect = ( - lambda keys, req_context: generate_store_output(keys) + r.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) ) # With sync scheduling, the connector flushes completed stores when the @@ -1059,8 +1252,8 @@ def test_offload_prompt_only(request_runner, async_scheduling: bool): extra_config_overrides={"offload_prompt_only": True}, ) - runner.manager.prepare_store.side_effect = ( - lambda keys, req_context: generate_store_output(keys) + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) ) runner.new_request(token_ids=[0] * offloaded_block_size * num_prompt_blocks) diff --git a/tests/v1/kv_connector/unit/test_config.py b/tests/v1/kv_connector/unit/test_config.py index 33c9abd09e..05189639c8 100644 --- a/tests/v1/kv_connector/unit/test_config.py +++ b/tests/v1/kv_connector/unit/test_config.py @@ -104,6 +104,16 @@ def test_kv_connector_rejects_expandable_segments(monkeypatch, kv_connector): _build_config(kv_connector=kv_connector) +def test_native_offloading_allows_expandable_segments(monkeypatch): + """Native CPU offload resolves GPU virtual addresses for each CUDA copy + and explicitly opts into the VMM-safe transfer contract.""" + monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") + + cfg = _build_config(kv_connector="OffloadingConnector") + + assert cfg._connector_supports_vmm_safe_transfers() + + def test_kv_connector_allows_expandable_segments_with_sleep_mode(monkeypatch): """Sleep mode routes KV allocations through CuMemAllocator's pool, which auto-disables expandable_segments (see #40812).""" diff --git a/tests/v1/sample/test_sampler.py b/tests/v1/sample/test_sampler.py index c67199fa40..f8c0a5bc65 100644 --- a/tests/v1/sample/test_sampler.py +++ b/tests/v1/sample/test_sampler.py @@ -5,6 +5,7 @@ import pytest import torch +import vllm.envs as envs from tests.v1.sample.utils import create_allowed_token_ids from vllm.platforms import current_platform from vllm.utils.platform_utils import is_pin_memory_available @@ -25,6 +26,30 @@ MAX_NUM_PROMPT_TOKENS = 64 +def test_sm70_compact_topk20_metadata_gate(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(envs, "VLLM_SM70_COMPACT_TOPK20_SAMPLER", True) + metadata = _create_default_sampling_metadata( + num_output_tokens=0, + batch_size=1, + vocab_size=VOCAB_SIZE, + device=torch.device("cpu"), + ) + metadata.all_greedy = False + metadata.all_random = True + metadata.max_num_logprobs = None + metadata.top_k_cpu = (20,) + metadata.top_p_cpu = (0.95,) + metadata.temperature_cpu = (1.0,) + + assert Sampler._sm70_compact_topk20_metadata_supported(metadata, "raw_logprobs") + metadata.top_k_cpu = (19,) + assert not Sampler._sm70_compact_topk20_metadata_supported(metadata, "raw_logprobs") + metadata.top_k_cpu = (20,) + assert not Sampler._sm70_compact_topk20_metadata_supported( + metadata, "processed_logits" + ) + + def _create_fake_logits(batch_size: int, vocab_size: int) -> torch.Tensor: fake_logits = torch.full((batch_size, vocab_size), 1e-2, dtype=torch.float) return fake_logits diff --git a/tests/v1/sample/test_topk_topp_sampler.py b/tests/v1/sample/test_topk_topp_sampler.py index 2fc26bfe93..3619d75dae 100644 --- a/tests/v1/sample/test_topk_topp_sampler.py +++ b/tests/v1/sample/test_topk_topp_sampler.py @@ -54,6 +54,69 @@ def reset_default_device(): torch.set_default_device(original_device) +@pytest.mark.skipif(not HAS_TRITON, reason="Triton is not available") +def test_sm70_topk_topp_8_warp_policy(monkeypatch: pytest.MonkeyPatch): + from vllm.v1.sample.ops import topk_topp_triton + + monkeypatch.delenv("VLLM_SM70_TOPK_TOPP_B8_B16_8_WARPS", raising=False) + monkeypatch.delenv("VLLM_SM70_TOPK_TOPP_8_WARPS", raising=False) + monkeypatch.setattr(topk_topp_triton.current_platform, "is_cuda", lambda: True) + monkeypatch.setattr( + topk_topp_triton.current_platform, + "is_device_capability", + lambda capability: capability == (7, 0), + ) + + use_8_warps = topk_topp_triton._use_sm70_topk_topp_8_warps + cuda = torch.device("cuda:0") + assert use_8_warps(cuda, 8, 248_320, True, True) + assert use_8_warps(cuda, 16, 248_320, True, True) + assert not use_8_warps(cuda, 5, 248_320, True, True) + + monkeypatch.setenv("VLLM_SM70_TOPK_TOPP_8_WARPS", "1") + assert use_8_warps(cuda, 5, 248_320, True, True) + monkeypatch.setenv("VLLM_SM70_TOPK_TOPP_B8_B16_8_WARPS", "0") + assert not use_8_warps(cuda, 8, 248_320, True, True) + + assert not use_8_warps(cuda, 16, 128_000, True, True) + assert not use_8_warps(cuda, 16, 248_320, False, True) + assert not use_8_warps(torch.device("cpu"), 16, 248_320, True, True) + + +@pytest.mark.parametrize("batch_size", (8, 16)) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_sm70_topk_topp_b8_b16_8_warps_matches_rollback( + monkeypatch: pytest.MonkeyPatch, + batch_size: int, +): + if torch.cuda.get_device_capability() != (7, 0): + pytest.skip("The eight-warp production schedule is SM70-only") + + import vllm.envs as envs + from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton + + generator = Generator(device="cuda").manual_seed(20260824 + batch_size) + logits = torch.randn( + (batch_size, 248_320), + generator=generator, + device="cuda", + dtype=torch.float32, + ) + k = torch.full((batch_size,), 20, device="cuda", dtype=torch.int32) + p = torch.full((batch_size,), 0.95, device="cuda", dtype=torch.float32) + + monkeypatch.setenv("VLLM_SM70_TOPK_TOPP_B8_B16_8_WARPS", "0") + envs.disable_envs_cache() + rollback = apply_top_k_top_p_triton(logits.clone(), k, p) + + monkeypatch.setenv("VLLM_SM70_TOPK_TOPP_B8_B16_8_WARPS", "1") + envs.disable_envs_cache() + candidate = apply_top_k_top_p_triton(logits.clone(), k, p) + + torch.accelerator.synchronize() + assert torch.equal(candidate, rollback) + + def test_flashinfer_sampler_default_falls_back_when_import_fails(monkeypatch): from vllm.v1.sample.ops import topk_topp_sampler as sampler_mod 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 new file mode 100644 index 0000000000..98cda4fe8c --- /dev/null +++ b/tests/v1/spec_decode/test_dflash2.py @@ -0,0 +1,1310 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import pytest +import torch + +import vllm.model_executor.layers.logits_processor as logits_processor_module +import vllm.model_executor.models.qwen3_dflash as dflash_model +import vllm.model_executor.models.qwen3_dflash2 as dflash2_model +import vllm.v1.attention.backends.flash_attn_v100 as flash_v100 +import vllm.v1.worker.gpu.attn_utils as attn_utils +import vllm.v1.worker.gpu.spec_decode.dflash.speculator as dflash_speculator +from vllm import envs +from vllm.config.speculative import SpeculativeConfig +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( + _is_dflash2_spec_config, +) +from vllm.model_executor.layers.vocab_parallel_embedding import ( + UnquantizedEmbeddingMethod, + VocabParallelEmbedding, + _sm70_dflash2_dense_order_topk, + _sm70_dflash2_rerank_output_buffers, +) +from vllm.model_executor.models.dflash_sm70 import ( + DFLASH_SM70_GATE_UP_INPUT_SCALE, + DFLASH_SM70_WIDE_OUTPUT_SCALE, + DFlashSM70RMSNorm, + dflash_scale_output_sm70, + dflash_silu_and_mul_sm70, +) +from vllm.model_executor.models.qwen3_dflash import ( + DFlashQwen3ForCausalLM, + DFlashQwen3Model, + _dflash_layer_causal, +) +from vllm.model_executor.models.qwen3_dflash2 import ( + DFlash2Qwen3ForCausalLM, + DFlash2Qwen3Model, + _grouped_conv, + _score_edges, +) +from vllm.v1.attention.backends.flash_attn_v100 import FlashAttnV100Impl +from vllm.v1.core.kv_cache_utils import unify_kv_cache_spec_page_size +from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec +from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample +from vllm.v1.worker.gpu.spec_decode import init_speculator +from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator +from vllm.v1.worker.gpu.spec_decode.dflash2.sparse_rejection import ( + _parse_alignment_steps, + _supports_sparse_sampling_contract, +) +from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import ( + DFlash2Speculator, + _requires_sm70_tail, + _selector_walk_kernel, +) + + +def test_dflash2_gdn_fastpaths_are_default_off(monkeypatch): + names = ( + "VLLM_FLASH_V100_DFLASH2_GROUPED_VERIFY", + "VLLM_SM70_DFLASH2_QPN8_RERANK", + "VLLM_SM70_DFLASH2_VERIFY_FASTPATH", + "VLLM_SM70_DFLASH2_FUSED_GDN_METADATA", + "VLLM_SM70_DFLASH2_GDN_METADATA_SHADOW", + "VLLM_SM70_DFLASH2_FUSED_GDN_VERIFY", + "VLLM_SM70_DFLASH2_FUSED_GDN_NORM", + "VLLM_SM70_DFLASH2_FUSED_GDN_SPLIT", + "VLLM_SM70_DFLASH2_FUSED_SMALLQ_METADATA", + "VLLM_SM70_DFLASH2_GROUPED_SMALLQ_METADATA", + "VLLM_SM70_DFLASH2_FUSED_QKV_PACK", + "VLLM_SM70_DFLASH2_FUSED_GEMMA_RMS", + "VLLM_SM70_DFLASH2_SPARSE_TARGET_REJECTION", + "VLLM_SM70_DFLASH2_SHARDED_CONTEXT_FC", + ) + for name in names: + monkeypatch.delenv(name, raising=False) + envs.disable_envs_cache() + try: + assert not any(getattr(envs, name) for name in names) + 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() + try: + assert envs.VLLM_SM70_TP4_PUSH_ALLREDUCE + monkeypatch.setenv("VLLM_SM70_TP4_PUSH_ALLREDUCE", "0") + envs.disable_envs_cache() + assert not envs.VLLM_SM70_TP4_PUSH_ALLREDUCE + finally: + envs.disable_envs_cache() + + +def _bare_dflash2_model() -> DFlash2Qwen3Model: + model = DFlash2Qwen3Model.__new__(DFlash2Qwen3Model) + torch.nn.Module.__init__(model) + model.quant_config = None + return model + + +def _fake_rms_norm( + output: torch.Tensor, + input_: torch.Tensor, + weight: torch.Tensor, + epsilon: float, +) -> None: + normalized = input_.float() * torch.rsqrt( + input_.float().pow(2).mean(dim=-1, keepdim=True) + epsilon + ) + if weight.ndim == 2: + weight = weight.view(weight.shape[0], *([1] * (input_.ndim - 2)), -1) + output.copy_((normalized * weight.float()).to(output.dtype)) + + +def _bare_context_k_model() -> DFlashQwen3Model: + model = DFlashQwen3Model.__new__(DFlashQwen3Model) + torch.nn.Module.__init__(model) + model._rms_norm_eps = 1e-6 + model._k_norm_weights = torch.tensor( + [ + [1.0, 1.0, 1.0, 1.0], + [0.5, 1.5, 0.75, 1.25], + [1.5, 0.5, 1.25, 0.75], + ], + dtype=torch.float32, + ) + model._batched_k_norm_runtime_verified = None + return model + + +def test_context_k_norm_trusts_verified_batched_runtime(monkeypatch): + model = _bare_context_k_model() + all_k = torch.randn(3, 2, 1, 4) + calls = [] + + def grouped_runtime(output, input_, weight, epsilon): + calls.append(weight.ndim) + _fake_rms_norm(output, input_, weight, epsilon) + + monkeypatch.setattr(dflash_model.ops, "rms_norm", grouped_runtime) + first = model._normalize_context_k(all_k) + assert model._batched_k_norm_runtime_verified is True + assert calls == [2, 1, 1, 1, 2] + + calls.clear() + second = model._normalize_context_k(all_k) + assert calls == [2] + assert torch.equal(second, first) + + +def test_context_k_norm_falls_back_for_stale_stable_binary(monkeypatch): + model = _bare_context_k_model() + # An all-zero profiling input would make the stale grouped output appear + # correct unless the capability probe uses its own nonzero fixture. + all_k = torch.zeros(3, 2, 1, 4) + calls = [] + + def stale_runtime(output, input_, weight, epsilon): + calls.append(weight.ndim) + # Historical stable-ABI binaries silently used row zero for a 2-D + # weight tensor. One-dimensional calls remained correct. + effective_weight = weight[0] if weight.ndim == 2 else weight + _fake_rms_norm(output, input_, effective_weight, epsilon) + + monkeypatch.setattr(dflash_model.ops, "rms_norm", stale_runtime) + actual = model._normalize_context_k(all_k) + expected = torch.empty_like(all_k) + for layer_idx in range(all_k.shape[0]): + _fake_rms_norm( + expected[layer_idx], + all_k[layer_idx], + model._k_norm_weights[layer_idx], + model._rms_norm_eps, + ) + + assert model._batched_k_norm_runtime_verified is False + assert calls == [2, 1, 1, 1, 1, 1, 1] + assert torch.equal(actual, expected) + + calls.clear() + repeated = model._normalize_context_k(all_k) + assert calls == [1, 1, 1] + assert torch.equal(repeated, expected) + + +def test_sm70_tp4_shards_only_compatible_dflash2_context_projection(monkeypatch): + model = _bare_dflash2_model() + config = SimpleNamespace( + parallel_config=SimpleNamespace(tensor_parallel_size=4), + model_config=SimpleNamespace(dtype=torch.float16), + ) + created = {} + + def fake_column_parallel(**kwargs): + created.update(kwargs) + return SimpleNamespace() + + monkeypatch.setattr( + dflash2_model.envs, + "VLLM_SM70_DFLASH2_SHARDED_CONTEXT_FC", + True, + ) + monkeypatch.setattr(dflash2_model.current_platform, "is_cuda", lambda: True) + monkeypatch.setattr( + dflash2_model.current_platform, + "is_device_capability", + lambda capability: capability == 70, + ) + monkeypatch.setattr(dflash2_model, "ColumnParallelLinear", fake_column_parallel) + + projection = model._make_context_projection( + vllm_config=config, + input_size=25600, + output_size=5120, + prefix="model.fc", + ) + + assert created["gather_output"] is True + assert created["input_size"] == 25600 + assert created["output_size"] == 5120 + assert projection._sm70_f16_force_enable is True + assert projection._sm70_f16_max_m == 64 + + +@pytest.mark.parametrize( + ("enabled", "tp_size", "input_size", "output_size"), + [ + (False, 4, 25600, 5120), + (True, 2, 25600, 5120), + (True, 4, 5120, 5120), + (True, 4, 25600, 4096), + ], +) +def test_sharded_context_projection_falls_back_outside_exact_contract( + monkeypatch, enabled, tp_size, input_size, output_size +): + model = _bare_dflash2_model() + config = SimpleNamespace( + parallel_config=SimpleNamespace(tensor_parallel_size=tp_size), + model_config=SimpleNamespace(dtype=torch.float16), + ) + sentinel = object() + monkeypatch.setattr( + dflash2_model.envs, + "VLLM_SM70_DFLASH2_SHARDED_CONTEXT_FC", + enabled, + ) + monkeypatch.setattr(dflash2_model.current_platform, "is_cuda", lambda: True) + monkeypatch.setattr( + dflash2_model.current_platform, + "is_device_capability", + lambda capability: capability == 70, + ) + monkeypatch.setattr( + DFlashQwen3Model, + "_make_context_projection", + lambda *_args, **_kwargs: sentinel, + ) + + projection = model._make_context_projection( + vllm_config=config, + input_size=input_size, + output_size=output_size, + prefix="model.fc", + ) + + assert projection is sentinel + + +@pytest.mark.parametrize("block_size", [4, 6, 8]) +def test_grouped_conv_matches_reference(block_size: int): + torch.manual_seed(0) + batch, taps, num_groups, group_size = 3, 3, 4, 2 + hidden = torch.randn(batch * block_size, num_groups * group_size) + delta = torch.randn(batch * block_size, taps, num_groups) + base = torch.randn(taps, num_groups * group_size) + + actual = _grouped_conv( + hidden, delta, base, block_size, num_groups, group_size, taps + ) + hidden_blocks = hidden.view(batch, block_size, num_groups, group_size) + expected = torch.zeros_like(hidden_blocks) + base = base.view(taps, num_groups, group_size) + delta = delta.view(batch, block_size, taps, num_groups) + for position in range(block_size): + for tap in range(min(taps, position + 1)): + expected[:, position] += ( + base[tap] + delta[:, position, tap, :, None] + ) * hidden_blocks[:, position - tap] + + torch.testing.assert_close(actual, expected.flatten(0, 1).flatten(-2)) + + +def test_selector_edges_match_sequential_reference(): + torch.manual_seed(1) + batch, steps, top_k, rank = 2, 4, 3, 5 + vocab = 17 + predecessors = torch.randn(vocab, rank) + successors = torch.randn(vocab, rank) + candidate_ids = torch.randint(vocab, (batch, steps, top_k)) + unary = torch.randn(batch, steps, top_k) + hidden = torch.randn(batch, steps, rank) + anchors = torch.randint(vocab, (batch,)) + + actual = _score_edges( + predecessors, + successors, + candidate_ids, + unary, + hidden, + anchors, + top_k, + ) + expected = torch.empty_like(actual) + for step in range(steps): + pred = ( + anchors[:, None].expand(-1, top_k) + if step == 0 + else candidate_ids[:, step - 1] + ) + expected[:, step] = unary[:, step, None] + torch.einsum( + "bpr,bcr->bpc", + predecessors[pred] * hidden[:, step, None], + successors[candidate_ids[:, step]], + ) + + torch.testing.assert_close(actual, expected) + + +def _stub_base(monkeypatch: pytest.MonkeyPatch, draft_logits): + def init_base(self, _vllm_config, device): + self.draft_model_config = SimpleNamespace( + hf_config=SimpleNamespace(dflash_config={"selector_top_k": 16}) + ) + self.max_num_reqs = 2 + self.num_query_per_req = 8 + self.num_speculative_steps = 7 + self.vocab_size = 31 + self.draft_tokens = torch.empty((2, 7), dtype=torch.int64, device=device) + self.draft_logits = draft_logits + + monkeypatch.setattr(DFlashSpeculator, "__init__", init_base) + + +def test_selector_leaves_greedy_without_proposal_logits(monkeypatch): + _stub_base(monkeypatch, None) + speculator = DFlash2Speculator(None, torch.device("cpu")) + assert speculator.draft_logits is None + + +def test_selector_default_path_does_not_allocate_sparse_score_cache(monkeypatch): + allocated = torch.zeros((2, 7, 31), dtype=torch.float32) + _stub_base(monkeypatch, allocated) + monkeypatch.setattr(envs, "VLLM_SM70_DFLASH2_SPARSE_TARGET_REJECTION", False) + monkeypatch.setattr(envs, "VLLM_SPEC_DUMP_ALIGNMENT", False) + speculator = DFlash2Speculator(None, torch.device("cpu")) + assert speculator.draft_logits is allocated + assert torch.isneginf(speculator.draft_logits).all() + assert speculator.get_sparse_draft_logits() is None + assert speculator.get_selector_alignment_shadow() is None + + +def test_selector_opt_in_allocates_sparse_score_cache(monkeypatch): + allocated = torch.zeros((2, 7, 31), dtype=torch.float32) + _stub_base(monkeypatch, allocated) + monkeypatch.setattr(envs, "VLLM_SM70_DFLASH2_SPARSE_TARGET_REJECTION", True) + speculator = DFlash2Speculator(None, torch.device("cpu")) + sparse_logits = speculator.get_sparse_draft_logits() + assert sparse_logits is not None + candidate_ids, candidate_scores = sparse_logits + assert candidate_ids.shape == (2, 7, 16) + assert candidate_scores.shape == (2, 7, 16) + assert candidate_scores.dtype is torch.float32 + + +def test_selector_alignment_shadow_is_explicit_and_keeps_full_lattice(monkeypatch): + allocated = torch.zeros((2, 7, 31), dtype=torch.float32) + _stub_base(monkeypatch, allocated) + monkeypatch.setattr(envs, "VLLM_SM70_DFLASH2_SPARSE_TARGET_REJECTION", True) + monkeypatch.setattr(envs, "VLLM_SPEC_DUMP_ALIGNMENT", True) + + speculator = DFlash2Speculator(None, torch.device("cpu")) + shadow = speculator.get_selector_alignment_shadow() + + assert shadow is not None + candidate_ids, unary_logits, lattice_scores = shadow + assert candidate_ids.shape == (2, 7, 16) + assert unary_logits.shape == (2, 7, 16) + assert unary_logits.dtype is torch.float32 + assert lattice_scores.shape == (2, 7, 16, 16) + assert lattice_scores.dtype is torch.float32 + + +@pytest.mark.parametrize( + ("raw", "expected"), + [ + (None, None), + ("", None), + ("1,3-5", {1, 3, 4, 5}), + ("5-3", set()), + ("bad", set()), + ("-1", set()), + ], +) +def test_selector_alignment_step_filter(raw, expected): + assert _parse_alignment_steps(raw) == expected + + +def test_selector_uses_checkpoint_top16_and_fp32_proposal_cache(monkeypatch): + _stub_base(monkeypatch, None) + speculator = DFlash2Speculator(None, torch.device("cpu")) + dtype, fill = DFlash2Speculator.draft_logits_spec(None, None) + assert speculator.selector_top_k == 16 + assert dtype is torch.float32 + assert fill == float("-inf") + + +def _sparse_sampling_contract_fixture(): + idx = np.array([0], dtype=np.int32) + sampling_states = SimpleNamespace( + temperature=SimpleNamespace(np=np.array([1.0], dtype=np.float32)), + top_k=SimpleNamespace(np=np.array([20], dtype=np.int32)), + top_p=SimpleNamespace(np=np.array([0.95], dtype=np.float32)), + min_p=SimpleNamespace(np=np.array([0.0], dtype=np.float32)), + max_num_logprobs=Mock(return_value=-1), + ) + sampler = SimpleNamespace( + sampling_states=sampling_states, + penalties_state=SimpleNamespace(use_penalty=np.array([False])), + logit_bias_state=SimpleNamespace(use_logit_bias=np.array([False])), + bad_words_state=SimpleNamespace( + num_bad_words=SimpleNamespace(np=np.array([0], dtype=np.int32)) + ), + logprob_token_ids_state=SimpleNamespace(max_num_token_ids=Mock(return_value=0)), + compute_nans=False, + ) + rejection_sampler = SimpleNamespace( + rejection_sample_method="standard", + sampler=sampler, + ) + input_batch = SimpleNamespace( + num_reqs=1, + is_prefilling_np=np.array([False]), + idx_mapping_np=idx, + ) + return rejection_sampler, input_batch + + +def test_sparse_target_rejection_accepts_official_sampling_contract(): + rejection_sampler, input_batch = _sparse_sampling_contract_fixture() + assert _supports_sparse_sampling_contract(rejection_sampler, input_batch) + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("temperature", 0.0), + ("top_k", 16), + ("top_p", 0.0), + ("min_p", 0.05), + ("penalty", True), + ("logit_bias", True), + ("bad_words", 1), + ("logprobs", 1), + ("custom_logprobs", 1), + ("compute_nans", True), + ], +) +def test_sparse_target_rejection_falls_back_for_unsupported_sampling( + field: str, + value: float | int | bool, +): + rejection_sampler, input_batch = _sparse_sampling_contract_fixture() + sampler = rejection_sampler.sampler + if field in {"temperature", "top_k", "top_p", "min_p"}: + getattr(sampler.sampling_states, field).np[0] = value + elif field == "penalty": + sampler.penalties_state.use_penalty[0] = value + elif field == "logit_bias": + sampler.logit_bias_state.use_logit_bias[0] = value + elif field == "bad_words": + sampler.bad_words_state.num_bad_words.np[0] = value + elif field == "logprobs": + sampler.sampling_states.max_num_logprobs.return_value = value + elif field == "custom_logprobs": + sampler.logprob_token_ids_state.max_num_token_ids.return_value = value + else: + sampler.compute_nans = value + + assert not _supports_sparse_sampling_contract(rejection_sampler, input_batch) + + +def test_probabilistic_selector_caches_temperature_applied_scores(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for the DFlash2 selector kernel") + + device = torch.device("cuda") + num_steps, top_k = 2, 4 + scores = torch.tensor( + [ + [ + [ + [0.0, 0.5, 1.0, 1.5], + [2.0, 2.5, 3.0, 3.5], + [4.0, 4.5, 5.0, 5.5], + [6.0, 6.5, 7.0, 7.5], + ], + [ + [8.0, 8.5, 9.0, 9.5], + [10.0, 10.5, 11.0, 11.5], + [12.0, 12.5, 13.0, 13.5], + [14.0, 14.5, 15.0, 15.5], + ], + ] + ], + dtype=torch.float32, + device=device, + ) + candidates = torch.arange(top_k, dtype=torch.int64, device=device).repeat( + 1, num_steps, 1 + ) + sample_pos = torch.tensor([10, 11], dtype=torch.int64, device=device) + req_state = torch.zeros(num_steps, dtype=torch.int32, device=device) + temperature = torch.tensor([0.5], dtype=torch.float32, device=device) + seeds = torch.tensor([123], dtype=torch.int64, device=device) + tokens = torch.full((num_steps,), -1, dtype=torch.int64, device=device) + realized = torch.full( + (1, num_steps, top_k), + float("nan"), + dtype=torch.float32, + device=device, + ) + path_state = torch.empty(1, dtype=torch.int32, device=device) + + _selector_walk_kernel[(1,)]( + scores, + candidates, + sample_pos, + req_state, + temperature, + seeds, + tokens, + realized, + path_state, + num_steps=num_steps, + walk_steps=num_steps, + top_k=top_k, + BLOCK_K=top_k, + SAMPLE_PROBABILISTIC=True, + USE_FP64=False, + num_warps=1, + ) + + first_index = int(tokens[0].item()) + torch.testing.assert_close(realized[0, 0], scores[0, 0, 0] / temperature[0]) + torch.testing.assert_close( + realized[0, 1], scores[0, 1, first_index] / temperature[0] + ) + + +def test_probabilistic_cache_respects_column_stride(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for the Gumbel kernel") + + device = torch.device("cuda") + vocab_size = 17 + padded_vocab_size = 23 + logits = torch.arange(vocab_size, dtype=torch.float32, device=device)[None] + storage = torch.full( + (1, 2, padded_vocab_size), + float("nan"), + dtype=torch.float32, + device=device, + ) + cache = storage[:, :, :vocab_size] + + gumbel_sample( + logits, + expanded_idx_mapping=torch.tensor([0], dtype=torch.int32, device=device), + temperature=torch.tensor([1.0], dtype=torch.float32, device=device), + seed=torch.tensor([123], dtype=torch.int64, device=device), + pos=torch.tensor([7], dtype=torch.int64, device=device), + apply_temperature=True, + output_processed_logits=cache, + output_processed_logits_col=torch.tensor(1, device=device), + ) + + assert torch.isnan(cache[0, 0]).all() + torch.testing.assert_close(cache[0, 1], logits[0]) + assert torch.isnan(storage[:, :, vocab_size:]).all() + + +def test_probabilistic_cache_keeps_ids_and_scores_in_request_slot_order(monkeypatch): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for the DFlash2 cache kernel") + + device = torch.device("cuda") + dense_cache = torch.zeros((2, 7, 31), dtype=torch.float32, device=device) + _stub_base(monkeypatch, dense_cache) + monkeypatch.setattr(envs, "VLLM_SM70_DFLASH2_SPARSE_TARGET_REJECTION", True) + speculator = DFlash2Speculator(None, device) + speculator.sample_idx_mapping = torch.tensor( + [1] * 7 + [0] * 7, + dtype=torch.int32, + device=device, + ) + candidate_ids = torch.stack( + ( + torch.arange(16, dtype=torch.int64, device=device).repeat(7, 1), + torch.arange(15, 31, dtype=torch.int64, device=device).repeat(7, 1), + ) + ) + selector_scores = torch.arange( + 2 * 7 * 16, + dtype=torch.float32, + device=device, + ).view(2, 7, 16) + speculator._selector_scores.copy_(selector_scores) + + speculator._cache_draft_logits(candidate_ids, num_sample=14) + sparse_logits = speculator.get_sparse_draft_logits() + assert sparse_logits is not None + cached_ids, cached_scores = sparse_logits + + assert torch.equal(cached_ids[1], candidate_ids[0]) + assert torch.equal(cached_ids[0], candidate_ids[1]) + assert torch.equal(cached_scores[1], selector_scores[0]) + assert torch.equal(cached_scores[0], selector_scores[1]) + assert torch.equal(dense_cache[1].gather(1, candidate_ids[0]), selector_scores[0]) + assert torch.equal(dense_cache[0].gather(1, candidate_ids[1]), selector_scores[1]) + + +def test_dflash2_selector_contract_dispatches_to_mrv2(monkeypatch): + monkeypatch.setattr(DFlash2Speculator, "__init__", lambda self, *_args: None) + config = SimpleNamespace( + speculative_config=SimpleNamespace( + method="dflash", + draft_model_config=SimpleNamespace( + hf_config=SimpleNamespace(dflash_config={"selector_top_k": 16}) + ), + ) + ) + assert isinstance(init_speculator(config, torch.device("cpu")), DFlash2Speculator) + + +def test_dflash_without_selector_stays_on_official_mrv2_speculator(monkeypatch): + monkeypatch.setattr(DFlashSpeculator, "__init__", lambda self, *_args: None) + config = SimpleNamespace( + speculative_config=SimpleNamespace( + method="dflash", + draft_model_config=SimpleNamespace( + hf_config=SimpleNamespace(dflash_config={}) + ), + ) + ) + assert isinstance(init_speculator(config, torch.device("cpu")), DFlashSpeculator) + + +@pytest.mark.parametrize( + ("method", "selector_top_k", "expected"), + [ + ("dflash", 16, True), + ("dflash", 0, False), + ("dflash_ddtree", 16, False), + ("mtp", 16, False), + ("eagle3", 16, False), + ], +) +def test_fused_gdn_verify_config_uses_selector_engine_contract( + method: str, + selector_top_k: int, + expected: bool, +): + config = SimpleNamespace( + speculative_config=SimpleNamespace( + method=method, + draft_model_config=SimpleNamespace( + hf_config=SimpleNamespace( + dflash_config={"selector_top_k": selector_top_k} + ) + ), + ) + ) + + assert _is_dflash2_spec_config(config) is expected + + +@pytest.mark.parametrize( + ("method", "expected"), + [ + ("dflash", False), + ("dflash_ddtree", True), + ("eagle3", True), + ("mtp", True), + ], +) +def test_only_mrv2_dflash_skips_eagle_prefix_block_drop(method, expected): + config = SimpleNamespace( + method=method, + use_eagle=lambda: True, + use_dflash=lambda: method == "dflash", + ) + assert SpeculativeConfig.use_eagle_kv_cache(config) is expected + + +@pytest.mark.parametrize("method", ["eagle3", "mtp"]) +def test_non_dflash_speculators_keep_eagle_dispatch(monkeypatch, method): + from vllm.v1.worker.gpu.spec_decode.eagle.speculator import EagleSpeculator + + monkeypatch.setattr(EagleSpeculator, "__init__", lambda self, *_args: None) + config = SimpleNamespace( + speculative_config=SimpleNamespace( + method=method, + use_eagle=lambda: True, + ) + ) + assert isinstance(init_speculator(config, torch.device("cpu")), EagleSpeculator) + + +def test_top_level_noncausal_override_wins_over_sliding_layer_default(): + config = SimpleNamespace( + is_causal=False, + dflash_config={}, + layer_types=["sliding_attention"], + ) + assert _dflash_layer_causal(config, 0) is False + + +def test_aux_hidden_states_follow_loaded_draft_projection_dtype(): + fc = torch.nn.Linear(10, 2, bias=False, dtype=torch.float16) + fc.input_size = 10 + model = SimpleNamespace(use_aux_hidden_state=True, fc=fc) + outer = SimpleNamespace(model=model) + hidden_states = torch.randn(3, 10, dtype=torch.float32) + + output = DFlashQwen3ForCausalLM.combine_hidden_states(outer, hidden_states) + + assert output.dtype is torch.float16 + + +@pytest.mark.parametrize("mamba_page_size_padded", [None, 16 * 512]) +def test_fp16_draft_cache_grows_padded_fp8_hybrid_pages( + mamba_page_size_padded: int | None, +): + block_size = 16 + target_page_size = 16 * 512 + specs = { + "target.attn": FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=256, + dtype=torch.float8_e5m2, + ), + "target.mamba": MambaSpec( + block_size=block_size, + shapes=((target_page_size,),), + dtypes=(torch.uint8,), + page_size_padded=mamba_page_size_padded, + ), + "draft.attn": FullAttentionSpec( + block_size=block_size, + num_kv_heads=2, + head_size=128, + dtype=torch.float16, + ), + } + + unified = unify_kv_cache_spec_page_size(specs) + + expected_page_size = specs["draft.attn"].page_size_bytes + assert {spec.page_size_bytes for spec in unified.values()} == {expected_page_size} + assert unified["target.attn"].block_size == 2 * block_size + assert unified["target.mamba"].block_size == 2 * block_size + assert unified["target.mamba"].page_size_padded == expected_page_size + + +def test_flashinfer_topk_is_capability_gated_on_sm70(monkeypatch): + dflash2_model._flashinfer_topk.cache_clear() + monkeypatch.setattr(dflash2_model.current_platform, "is_cuda", lambda: True) + monkeypatch.setattr( + dflash2_model.current_platform, + "has_device_capability", + lambda capability: capability <= 70, + ) + monkeypatch.setattr(dflash2_model, "has_flashinfer", lambda: True) + assert dflash2_model._flashinfer_topk() is None + dflash2_model._flashinfer_topk.cache_clear() + + +def test_target_topk_uses_reranked_local_candidates_without_dense_logits(monkeypatch): + dense_apply = Mock(side_effect=AssertionError("dense logits must not run")) + values = torch.linspace(5.0, -5.0, 20, dtype=torch.float16).reshape(1, 20) + ids = torch.arange(100, 120, dtype=torch.int64).reshape(1, 20) + lm_head = SimpleNamespace( + quant_method=SimpleNamespace(apply=dense_apply), + weight=torch.empty((62080, 1), dtype=torch.float16), + maybe_get_sm70_dflash2_top20=lambda hidden, top_k, bias: (values, ids), + ) + processor = LogitsProcessor(vocab_size=248320, scale=0.5, soft_cap=3.0) + monkeypatch.setattr( + logits_processor_module, + "get_tensor_model_parallel_world_size", + lambda: 1, + ) + + actual_ids, actual_values = processor.get_topk_tokens_and_logits( + lm_head, + torch.empty((1, 1), dtype=torch.float16), + 20, + ) + + expected_values = (torch.tanh(values / 3.0) * 3.0 * 0.5).float() + assert torch.equal(actual_ids, ids) + torch.testing.assert_close(actual_values, expected_values) + dense_apply.assert_not_called() + + +def test_draft_candidates_use_reranked_lm_head_without_dense_logits(monkeypatch): + dense_apply = Mock(side_effect=AssertionError("dense logits must not run")) + values = torch.linspace(4.0, -4.0, 16, dtype=torch.float16).reshape(1, 16) + ids = torch.arange(200, 216, dtype=torch.int64).reshape(1, 16) + candidate_path = Mock(return_value=(values, ids)) + lm_head = SimpleNamespace( + quant_method=UnquantizedEmbeddingMethod(), + maybe_get_sm70_dflash2_top20=candidate_path, + ) + lm_head.quant_method.apply = dense_apply + model = SimpleNamespace(candidate_selector=SimpleNamespace(top_k=16)) + dflash2 = SimpleNamespace( + lm_head=lm_head, + model=model, + output_multiplier=1.0, + final_logit_softcapping=None, + ) + hidden_states = torch.empty((1, 8), dtype=torch.float16) + monkeypatch.setattr( + dflash2_model, + "get_tensor_model_parallel_world_size", + lambda: 1, + ) + + actual_ids, actual_values = DFlash2Qwen3ForCausalLM.compute_candidates( + dflash2, hidden_states + ) + + candidate_path.assert_called_once_with(hidden_states, 16) + dense_apply.assert_not_called() + assert torch.equal(actual_ids, ids) + assert torch.equal(actual_values, values.float()) + + +def test_lm_head_candidate_interface_falls_back_when_rerank_is_disabled(monkeypatch): + monkeypatch.setattr(envs, "VLLM_SM70_DFLASH2_QPN8_RERANK", False) + layer = SimpleNamespace() + + assert ( + VocabParallelEmbedding.maybe_get_sm70_dflash2_top20( + layer, + torch.empty((1, 8)), + 20, + ) + is None + ) + + +@pytest.mark.parametrize("selector_k", [16, 20]) +@pytest.mark.parametrize("num_rows", [1, 7, 8]) +def test_qpn8_rerank_output_buffers_are_contiguous(selector_k, num_rows): + layer = SimpleNamespace() + for top_k in (16, 20): + setattr( + layer, + f"_sm70_dflash2_rerank_values_{top_k}", + torch.empty((8, top_k), dtype=torch.float16), + ) + setattr( + layer, + f"_sm70_dflash2_rerank_positions_{top_k}", + torch.empty((8, top_k), dtype=torch.int64), + ) + setattr( + layer, + f"_sm70_dflash2_rerank_ids_{top_k}", + torch.empty((8, top_k), dtype=torch.int64), + ) + + buffers = _sm70_dflash2_rerank_output_buffers(layer, num_rows, selector_k) + + assert all(buffer.shape == (num_rows, selector_k) for buffer in buffers) + assert all(buffer.is_contiguous() for buffer in buffers) + + +def test_qpn8_rerank_restores_dense_vocab_tie_order(): + vocab_start = 96 + candidate_ids = torch.tensor( + [[9, 1, 7, 3, 12, 4], [18, 2, 15, 5, 11, 8]], dtype=torch.int64 + ) + candidate_logits = torch.tensor( + [[5, 5, 5, 5, 4, 3], [7, 7, 7, 7, 6, 5]], dtype=torch.float16 + ) + sparse_logits = torch.empty((2, 24), dtype=torch.float16) + actual_values = torch.empty((2, 4), dtype=torch.float16) + actual_ids = torch.empty((2, 4), dtype=torch.int64) + + _sm70_dflash2_dense_order_topk( + sparse_logits, + candidate_ids, + candidate_logits, + actual_values, + actual_ids, + 4, + vocab_start, + ) + + reference = torch.full_like(sparse_logits, -float("inf")) + reference.scatter_(1, candidate_ids, candidate_logits) + expected_values, expected_ids = torch.topk(reference, 4, dim=-1, sorted=True) + assert torch.equal(actual_values, expected_values) + assert torch.equal(actual_ids, expected_ids + vocab_start) + + +@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]) +def test_sm70_f16_rerank_topk_matches_composite_key_contract( + selector_k: int, num_rows: int +): + if torch.cuda.get_device_capability() != (7, 0): + pytest.skip("the exact rerank top-k kernel is SM70-only") + required_ops = ( + "sm70_f16_rerank_keys_out", + "sm70_f16_rerank_topk_out", + ) + if any(not hasattr(torch.ops._C, name) for name in required_ops): + pytest.skip("the exact rerank top-k operators are not built") + + from vllm import _sm70_ops + + generator = torch.Generator().manual_seed(20260824) + local_vocab = 62080 + candidate_ids = torch.stack( + [torch.randperm(local_vocab, generator=generator)[:64] for _ in range(num_rows)] + ).to(device="cuda", dtype=torch.int64) + torch.manual_seed(20260824) + candidate_logits = torch.randn( + (num_rows, 64), dtype=torch.float16, device="cuda" + ).clamp_(-3, 3) + # Force more equal maxima than either requested top-k so ID-ascending tie + # precedence is part of every comparison, not just an incidental edge. + candidate_logits[:, :24] = 4 + if num_rows > 1: + candidate_logits[1, 24:28] = torch.tensor( + [-float("inf"), -0.0, 0.0, float("inf")], + dtype=torch.float16, + device="cuda", + ) + + actual_values = torch.empty( + (num_rows, selector_k), dtype=torch.float16, device="cuda" + ) + actual_ids = torch.empty((num_rows, selector_k), dtype=torch.int64, device="cuda") + vocab_start = 186240 + _sm70_ops.sm70_f16_rerank_topk_out( + actual_values, + actual_ids, + candidate_logits, + candidate_ids, + vocab_start, + ) + + keys = torch.empty_like(candidate_ids) + _sm70_ops.sm70_f16_rerank_keys_out(keys, candidate_logits, candidate_ids) + _, positions = torch.topk(keys, selector_k, dim=-1, sorted=True) + expected_values = candidate_logits.gather(1, positions) + expected_ids = candidate_ids.gather(1, positions).add_(vocab_start) + torch.accelerator.synchronize() + + assert torch.equal( + actual_values.view(torch.int16), expected_values.view(torch.int16) + ) + assert torch.equal(actual_ids, expected_ids) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("num_rows", [1, 7, 8]) +def test_sm70_dflash2_exact_rerank_matches_gathered_bmm(num_rows): + if torch.cuda.get_device_capability() != (7, 0): + pytest.skip("the exact rerank kernel is SM70-only") + if not hasattr(torch.ops._C, "sm70_f16_indexed_rerank_out"): + pytest.skip("the exact TurboMind rerank op is not built") + + from vllm import _sm70_ops + + torch.manual_seed(20260824) + device = torch.device("cuda") + candidates, vocab, hidden_size = 64, 256, 5120 + hidden = torch.randn((num_rows, hidden_size), dtype=torch.float16, device=device) + weight = torch.randn((vocab, hidden_size), dtype=torch.float16, device=device) + candidate_ids = torch.randint( + 0, vocab, (num_rows, candidates), dtype=torch.int64, device=device + ) + actual = torch.empty((num_rows, candidates), dtype=torch.float16, device=device) + selected_raw = torch.empty( + (num_rows * candidates, hidden_size), dtype=torch.float16, device=device + ) + selected_packed = torch.empty_like(selected_raw) + expanded = torch.empty( + (num_rows, num_rows * candidates), dtype=torch.float16, device=device + ) + partials = torch.empty( + (num_rows, num_rows * candidates), dtype=torch.float32, device=device + ) + barriers = torch.zeros(64, dtype=torch.int32, device=device) + + _sm70_ops.sm70_f16_indexed_rerank_out( + actual, + hidden, + weight, + candidate_ids, + selected_raw, + selected_packed, + expanded, + partials, + barriers, + 128, + 10, + ) + gathered = weight.index_select(0, candidate_ids.reshape(-1)) + expected = torch.bmm( + gathered.view(num_rows, candidates, hidden_size), + hidden.unsqueeze(-1), + ).squeeze(-1) + + torch.testing.assert_close(actual, expected, atol=0.125, rtol=0.002) + + +@pytest.mark.parametrize( + ("capability", "num_steps", "expected"), + [((7, 0), 7, True), ((7, 0), 1, False), ((8, 0), 7, False)], +) +def test_selector_tail_split_is_sm70_only(monkeypatch, capability, num_steps, expected): + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: capability) + assert _requires_sm70_tail(torch.device("cuda:0"), num_steps) is expected + + +def test_noncausal_draft_cannot_enter_flash_v100_small_query_fast_path(): + impl = SimpleNamespace( + use_flash_v100_decode=True, + smallq_decode_max_query_len=8, + smallq_decode_max_model_len=4096, + ) + metadata = SimpleNamespace( + causal=False, + query_start_loc=torch.tensor([0, 8], dtype=torch.int32), + max_model_len=128, + ) + assert not FlashAttnV100Impl._small_query_decode_enabled(impl, metadata) + + +def test_dflash_attention_builders_receive_the_draft_model_config(monkeypatch): + def fake_replace(config, **updates): + return SimpleNamespace(source=config, **updates) + + monkeypatch.setattr(dflash_speculator, "replace", fake_replace) + target_model_config = object() + draft_model_config = object() + attention_config = object() + speculator = SimpleNamespace( + vllm_config=SimpleNamespace( + model_config=target_model_config, + attention_config=attention_config, + ), + draft_model_config=draft_model_config, + requires_non_causal=True, + ) + + config = DFlashSpeculator.attn_vllm_config.fget(speculator) + + assert config.model_config is draft_model_config + assert config.attention_config.source is attention_config + assert config.attention_config.use_non_causal is True + + +@pytest.mark.parametrize( + ("mask", "expected"), + [ + (np.array([True], dtype=np.bool_), True), + (np.array([True, True], dtype=np.bool_), True), + (np.array([True, False], dtype=np.bool_), False), + (np.array([False], dtype=np.bool_), False), + (None, False), + ], +) +def test_dflash_context_only_prefill_requires_every_live_request(mask, expected): + batch = SimpleNamespace( + num_reqs=1 if mask is None else mask.size, + is_incomplete_prefilling_np=mask, + ) + assert dflash_speculator._is_context_only_prefill(batch) is expected + + +def test_dflash_intermediate_prefill_materializes_context_without_query(monkeypatch): + prepare_inputs = Mock() + monkeypatch.setattr(dflash_speculator, "prepare_dflash_inputs", prepare_inputs) + + speculator = DFlashSpeculator.__new__(DFlashSpeculator) + speculator.max_model_len = 256 + speculator.num_query_per_req = 8 + speculator.num_speculative_steps = 7 + speculator.max_num_reqs = 1 + speculator.max_num_tokens = 8 + speculator.hidden_states = torch.zeros(8, 3) + speculator.draft_tokens = torch.zeros(1, 7, dtype=torch.int64) + speculator.input_buffers = object() + speculator.context_positions = torch.zeros(8, dtype=torch.int64) + speculator.sample_indices = torch.zeros(7, dtype=torch.int64) + speculator.sample_pos = torch.zeros(7, dtype=torch.int64) + speculator.sample_idx_mapping = torch.zeros(7, dtype=torch.int32) + speculator.temperature = torch.ones(1) + speculator.seeds = torch.zeros(1, dtype=torch.int64) + speculator.parallel_drafting_token_id = 0 + speculator.sample_from_anchor = False + speculator.draft_kv_cache_group_id = 0 + speculator.draft_kv_cache_group_ids = [0] + speculator.block_tables = SimpleNamespace( + slot_mappings=[torch.zeros(8, dtype=torch.int64)], + input_block_tables=[torch.zeros(1, 8, dtype=torch.int32)], + kernel_block_sizes=[1], + cp_rank=0, + cp_size=1, + cp_interleave=1, + ) + speculator._context_slot_mappings = torch.zeros(1, 8, dtype=torch.int64) + speculator._layer_group_idx = None + speculator._context_only_prefill_logged = False + speculator._prepare_ngram_assist = Mock( + side_effect=AssertionError("ngram lookup must not run mid-prefill") + ) + speculator.model = SimpleNamespace( + precompute_and_store_context_kv=Mock(), + ) + + input_batch = SimpleNamespace( + num_reqs=1, + num_tokens=4, + seq_lens_cpu_upper_bound=torch.tensor([4], dtype=torch.int32), + is_incomplete_prefilling_np=np.array([True], dtype=np.bool_), + ) + output = speculator.propose( + input_batch=input_batch, + attn_metadata={}, + slot_mappings={}, + last_hidden_states=torch.arange(12, dtype=torch.float32).view(4, 3), + aux_hidden_states=None, + num_sampled=torch.zeros(1, dtype=torch.int32), + num_rejected=torch.zeros(1, dtype=torch.int32), + last_sampled=torch.zeros(1, dtype=torch.int64), + next_prefill_tokens=torch.zeros(1, dtype=torch.int64), + temperature=torch.ones(1), + seeds=torch.zeros(1, dtype=torch.int64), + ) + + prepare_inputs.assert_called_once() + speculator.model.precompute_and_store_context_kv.assert_called_once() + torch.testing.assert_close( + speculator.hidden_states[:4], + torch.arange(12, dtype=torch.float32).view(4, 3), + ) + assert output.tolist() == [[-1] * 7] + + +def test_noncausal_dflash_capture_binds_paged_prefix_attention(monkeypatch): + monkeypatch.setattr(flash_v100, "_is_cuda_graph_capturing", lambda _query: True) + output = torch.empty(8, 1, 1) + paged_prefix = Mock(return_value=output) + impl = SimpleNamespace( + _supports_flash_v100_path=lambda: True, + _layer_debug_info=lambda _layer: { + "layer_name": "draft", + "is_dflash_draft_attn": True, + }, + use_triton_prefill=False, + use_decode_scalar_paged=True, + use_decode_paged_prefill=False, + use_flash_v100_prefill_paged=True, + _small_query_decode_enabled=lambda _metadata: False, + _flash_v100_prefill_with_prefix=paged_prefix, + ) + metadata = SimpleNamespace( + max_query_len=8, + max_seq_len=1024, + num_actual_tokens=8, + causal=False, + query_start_loc=torch.tensor([0, 8], dtype=torch.int32), + query_start_loc_cpu=torch.tensor([0, 8], dtype=torch.int32), + # Capture-time CPU metadata looks like no-prefix prefill, while the + # persistent device metadata is updated before replay. + seq_lens=torch.tensor([17], dtype=torch.int32), + seq_lens_cpu=torch.tensor([8], dtype=torch.int32), + block_table=torch.tensor([[0]], dtype=torch.int32), + ) + query = torch.empty(8, 1, 1) + layer = SimpleNamespace(is_dflash_draft_attn=True) + + result = FlashAttnV100Impl.forward( + impl, + layer, + query, + query, + query, + torch.empty(1), + metadata, + output, + ) + + assert result is output + paged_prefix.assert_called_once() + + +def test_draft_attention_causality_is_resolved_per_kv_group(monkeypatch): + observed_causality = [] + + class CapturedCommonAttentionMetadata: + def __init__(self, **kwargs): + observed_causality.append(kwargs["causal"]) + + monkeypatch.setattr( + attn_utils, + "CommonAttentionMetadata", + CapturedCommonAttentionMetadata, + ) + kv_cache_config = SimpleNamespace( + kv_cache_groups=[SimpleNamespace(), SimpleNamespace()] + ) + attn_utils.build_attn_metadata( + attn_groups=[[], []], + num_reqs=1, + num_tokens=8, + query_start_loc_gpu=torch.tensor([0, 8], dtype=torch.int32), + query_start_loc_cpu=torch.tensor([0, 8], dtype=torch.int32), + max_query_len=8, + seq_lens=torch.tensor([8], dtype=torch.int32), + max_seq_len=8, + block_tables=[ + torch.zeros((1, 1), dtype=torch.int32), + torch.zeros((1, 1), dtype=torch.int32), + ], + slot_mappings=torch.zeros((2, 8), dtype=torch.int64), + kv_cache_config=kv_cache_config, + causal={0: False, 1: True}, + ) + + assert observed_causality == [False, True] + + +def test_sm70_rmsnorm_keeps_bf16_residual_in_fp32(): + norm = DFlashSM70RMSNorm(8, 1e-6, torch.float16) + norm.weight.data.copy_(torch.linspace(0.5, 1.5, 8, dtype=torch.float16)) + x = torch.full((2, 8), 300.0, dtype=torch.float16) + residual = torch.full((2, 8), 70000.0, dtype=torch.float32) + + output, residual_output = norm(x, residual) + expected_residual = ( + (x.float() * DFLASH_SM70_WIDE_OUTPUT_SCALE + residual) + .to(torch.bfloat16) + .float() + ) + + assert residual_output.dtype is torch.float32 + assert torch.isfinite(output).all() + assert residual_output.max() > torch.finfo(torch.float16).max + torch.testing.assert_close(residual_output, expected_residual, rtol=0, atol=0) + + +def test_sm70_swiglu_uses_power_of_two_row_scale(): + gate_up = ( + torch.tensor( + [ + [2000.0, -1500.0, 1000.0, 800.0, 1800.0, 900.0, -700.0, 600.0], + [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], + ], + dtype=torch.float16, + ) + / DFLASH_SM70_GATE_UP_INPUT_SCALE + ) + transported, row_scales = dflash_silu_and_mul_sm70(gate_up) + + assert torch.all(row_scales >= 1) + torch.testing.assert_close( + torch.log2(row_scales), + torch.log2(row_scales).round(), + rtol=0, + atol=0, + ) + assert torch.isfinite(transported).all() + + down = torch.ones((2, 8), dtype=torch.float16) + restored = dflash_scale_output_sm70(down, row_scales) + expected = ( + down.float() * (row_scales[:, None] / DFLASH_SM70_WIDE_OUTPUT_SCALE) + ).half() + torch.testing.assert_close(restored, expected, rtol=0, atol=0) diff --git a/tests/v1/spec_decode/test_dflash2_ngram_assist.py b/tests/v1/spec_decode/test_dflash2_ngram_assist.py new file mode 100644 index 0000000000..cbf44d08f4 --- /dev/null +++ b/tests/v1/spec_decode/test_dflash2_ngram_assist.py @@ -0,0 +1,386 @@ +# 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 + +from vllm.config.speculative import SpeculativeConfig +from vllm.v1.spec_decode.ngram_proposer import ( + _find_longest_matched_ngram_and_propose_tokens, +) +from vllm.v1.worker.gpu.spec_decode.dflash2.ngram_assist import ( + DFlash2NgramAssist, + find_split_ngram_proposal, +) +from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import ( + DFlash2Speculator, + _apply_ngram_draft_kernel, +) +from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import ( + dflash2_sparse_topk_rejection_sample, + rejection_sample, +) + + +@pytest.mark.parametrize("split", [0, 1, 4, 7, 10]) +@pytest.mark.parametrize(("min_ngram", "max_ngram"), [(1, 1), (2, 3), (3, 5)]) +def test_split_lookup_matches_standalone_ngram( + split: int, min_ngram: int, max_ngram: int +) -> None: + tokens = np.array([1, 2, 3, 4, 1, 2, 3, 5, 1, 2], dtype=np.int32) + expected = _find_longest_matched_ngram_and_propose_tokens( + tokens, + min_ngram, + max_ngram, + max_model_len=64, + k=3, + ) + prefix = tokens[:split] + suffix = tokens[split:].astype(np.int64) + actual = find_split_ngram_proposal( + prefix, + len(prefix), + suffix, + len(suffix), + min_ngram, + max_ngram, + max_model_len=64, + num_draft_tokens=3, + ) + np.testing.assert_array_equal(actual, expected) + + +def test_batch_assist_reports_only_full_width_hits() -> None: + assist = DFlash2NgramAssist(2, 3, num_draft_tokens=2, max_model_len=64) + token_ids = np.zeros((3, 16), dtype=np.int32) + token_ids[2, :6] = [1, 2, 3, 4, 1, 2] + token_ids[0, :4] = [8, 9, 10, 11] + sampled = np.array([[3, -1], [12, -1]], dtype=np.int64) + sampled_lens = np.array([1, 1], dtype=np.int32) + output = np.zeros((2, 2), dtype=np.int64) + output_lens = np.zeros(2, dtype=np.int32) + + hits = assist.propose( + token_ids, + req_state_indices=np.array([2, 0], dtype=np.int32), + prior_lengths=np.array([6, 4], dtype=np.int32), + sampled_token_ids=sampled, + num_sampled_tokens=sampled_lens, + eligible=np.array([True, True]), + output_tokens=output, + output_lengths=output_lens, + ) + + assert hits == 1 + np.testing.assert_array_equal(output[0], [4, 1]) + np.testing.assert_array_equal(output_lens, [2, 0]) + assert assist.num_eligible == 2 + assert assist.num_full_hits == 1 + + +class _CopyEvent: + def __init__(self) -> None: + self.synchronized = False + + def synchronize(self) -> None: + self.synchronized = True + + +def _host_only_speculator() -> DFlash2Speculator: + speculator = object.__new__(DFlash2Speculator) + speculator._ngram_assist = DFlash2NgramAssist( + 2, 3, num_draft_tokens=2, max_model_len=64 + ) + speculator._ngram_num_hits = 0 + speculator._ngram_rounds = 0 + speculator._ngram_skipped_rounds = 0 + speculator.num_speculative_steps = 2 + speculator._ngram_tokens_cpu_tensor = torch.zeros((2, 2), dtype=torch.int64) + speculator._ngram_lengths_cpu_tensor = torch.zeros(2, dtype=torch.int32) + speculator._ngram_tokens_cpu = speculator._ngram_tokens_cpu_tensor.numpy() + speculator._ngram_lengths_cpu = speculator._ngram_lengths_cpu_tensor.numpy() + speculator._ngram_tokens = torch.zeros((2, 2), dtype=torch.int64) + speculator._ngram_lengths = torch.zeros(2, dtype=torch.int32) + return speculator + + +def test_host_token_state_is_requested_only_for_enabled_assist() -> None: + speculator = object.__new__(DFlash2Speculator) + speculator._ngram_assist = None + assert not speculator.requires_host_token_state + + speculator._ngram_assist = object() # type: ignore[assignment] + assert speculator.requires_host_token_state + + +def test_prepare_assist_uses_request_slots_and_skips_only_all_hit() -> None: + speculator = _host_only_speculator() + token_ids = np.zeros((3, 16), dtype=np.int32) + token_ids[2, :6] = [1, 2, 3, 4, 1, 2] + token_ids[0, :6] = [7, 8, 9, 10, 7, 8] + batch = SimpleNamespace( + has_structured_output_reqs=False, + num_reqs=2, + num_draft_tokens_per_req=np.array([2, 2], dtype=np.int32), + seq_lens_cpu_upper_bound=torch.tensor([8, 8], dtype=torch.int32), + is_prefilling_np=np.array([False, False]), + idx_mapping_np=np.array([2, 0], dtype=np.int32), + ) + event = _CopyEvent() + + skip = speculator._prepare_ngram_assist( + batch, + event, + sampled_token_ids_cpu=np.array([[3], [9]], dtype=np.int64), + num_sampled_tokens_cpu=np.array([1, 1], dtype=np.int32), + all_token_ids_cpu=token_ids, + ) + + assert event.synchronized + assert skip + assert speculator._ngram_num_hits == 2 + assert torch.equal( + speculator._ngram_tokens, + torch.tensor([[4, 1], [10, 7]], dtype=torch.int64), + ) + + +def test_prepare_assist_keeps_dflash_draft_for_mixed_hits() -> None: + speculator = _host_only_speculator() + token_ids = np.zeros((3, 16), dtype=np.int32) + token_ids[2, :6] = [1, 2, 3, 4, 1, 2] + token_ids[0, :6] = [7, 8, 9, 10, 11, 12] + batch = SimpleNamespace( + has_structured_output_reqs=False, + num_reqs=2, + num_draft_tokens_per_req=np.array([2, 2], dtype=np.int32), + seq_lens_cpu_upper_bound=torch.tensor([8, 8], dtype=torch.int32), + is_prefilling_np=np.array([False, False]), + idx_mapping_np=np.array([2, 0], dtype=np.int32), + ) + event = _CopyEvent() + + skip = speculator._prepare_ngram_assist( + batch, + event, + sampled_token_ids_cpu=np.array([[3], [13]], dtype=np.int64), + num_sampled_tokens_cpu=np.array([1, 1], dtype=np.int32), + all_token_ids_cpu=token_ids, + ) + + assert event.synchronized + assert not skip + assert speculator._ngram_assist is not None + assert speculator._ngram_assist.num_full_hits == 1 + assert speculator._ngram_num_hits == 0 + assert torch.count_nonzero(speculator._ngram_tokens) == 0 + + +def test_prepare_assist_bypasses_structured_output_without_sync() -> None: + speculator = _host_only_speculator() + batch = SimpleNamespace(has_structured_output_reqs=True) + event = _CopyEvent() + + assert not speculator._prepare_ngram_assist( + batch, + event, + sampled_token_ids_cpu=np.zeros((1, 1), dtype=np.int64), + num_sampled_tokens_cpu=np.ones(1, dtype=np.int32), + all_token_ids_cpu=np.zeros((1, 8), dtype=np.int32), + ) + assert not event.synchronized + + +def test_ngram_assist_rejects_non_dflash_method() -> None: + with pytest.raises(ValueError, match="only supported with method='dflash'"): + SpeculativeConfig( + method="ngram", + num_speculative_tokens=2, + ngram_assist=True, + ) + + +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") + + device = torch.device("cuda") + num_reqs, num_steps, top_k, vocab_size = 2, 3, 4, 32 + ngram_tokens = torch.tensor( + [[11, 12, 13], [21, 22, 23]], dtype=torch.int64, device=device + ) + ngram_lengths = torch.tensor([3, 0], dtype=torch.int32, device=device) + # Batch row 0 belongs to request-state slot 1. + sample_req_state = torch.tensor( + [1, 1, 1, 0, 0, 0], dtype=torch.int32, device=device + ) + draft_tokens = torch.full( + (num_reqs, num_steps), 7, dtype=torch.int64, device=device + ) + cached_ids = ( + torch.arange( + num_reqs * num_steps * top_k, dtype=torch.int64, device=device + ).view(num_reqs, num_steps, top_k) + % vocab_size + ) + cached_scores = torch.randn( + num_reqs, num_steps, top_k, dtype=torch.float32, device=device + ) + draft_logits = torch.full( + (num_reqs, num_steps, vocab_size), + -float("inf"), + dtype=torch.float32, + device=device, + ) + draft_logits.scatter_(2, cached_ids, cached_scores) + miss_ids = cached_ids[0].clone() + miss_scores = cached_scores[0].clone() + miss_logits = draft_logits[0].clone() + + _apply_ngram_draft_kernel[(num_reqs * num_steps,)]( + ngram_tokens, + ngram_lengths, + sample_req_state, + draft_tokens, + draft_tokens.stride(0), + cached_ids, + cached_scores, + cached_ids.stride(0), + cached_ids.stride(1), + draft_logits, + draft_logits.stride(0), + draft_logits.stride(1), + num_steps=num_steps, + top_k=top_k, + BLOCK_K=top_k, + CACHE_DRAFT_LOGITS=True, + CACHE_SCORES=True, + num_warps=1, + ) + + assert torch.equal(draft_tokens[0], ngram_tokens[0]) + assert torch.equal(draft_tokens[1], torch.full((num_steps,), 7, device=device)) + assert torch.equal(cached_ids[0], miss_ids) + assert torch.equal(cached_scores[0], miss_scores) + assert torch.equal(draft_logits[0], miss_logits) + assert torch.equal(cached_ids[1, :, 0], ngram_tokens[0]) + assert torch.equal(cached_scores[1, :, 0], torch.zeros(num_steps, device=device)) + assert torch.isneginf(cached_scores[1, :, 1:]).all() + assert torch.equal( + draft_logits[1].gather(1, ngram_tokens[0, :, None]).squeeze(1), + torch.zeros(num_steps, device=device), + ) + + +@pytest.mark.parametrize("top_p", [1.0, 0.95]) +def test_ngram_one_hot_sparse_rejection_matches_dense(top_p: float) -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for rejection sampling") + + from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p + + torch.manual_seed(20260825) + device = torch.device("cuda") + num_reqs, num_steps = 16, 3 + target_top_k, draft_top_k, vocab_size = 20, 16, 97 + rows_per_req = num_steps + 1 + num_logits = num_reqs * rows_per_req + + target_ids = torch.stack( + [ + torch.randperm(vocab_size, device=device)[:target_top_k] + for _ in range(num_logits) + ] + ) + target_logits = ( + torch.randn(num_logits, target_top_k, dtype=torch.float32, device=device) + .sort(dim=-1, descending=True) + .values + ) + target_dense = torch.full( + (num_logits, vocab_size), + -float("inf"), + dtype=torch.float32, + device=device, + ) + target_dense.scatter_(1, target_ids, target_logits) + target_dense = apply_top_k_top_p( + target_dense, + torch.full((num_logits,), target_top_k, dtype=torch.int32, device=device), + torch.full((num_logits,), top_p, dtype=torch.float32, device=device), + ) + + proposals = torch.randint( + 0, vocab_size, (num_reqs, num_steps), dtype=torch.int64, device=device + ) + draft_ids = torch.zeros( + num_reqs, num_steps, draft_top_k, dtype=torch.int64, device=device + ) + draft_ids[:, :, 0] = proposals + draft_logits = torch.full( + (num_reqs, num_steps, draft_top_k), + -float("inf"), + dtype=torch.float32, + device=device, + ) + draft_logits[:, :, 0] = 0.0 + draft_dense = torch.full( + (num_reqs, num_steps, vocab_size), + -float("inf"), + dtype=torch.float32, + device=device, + ) + draft_dense.scatter_(2, proposals[:, :, None], 0.0) + + draft_sampled_2d = torch.zeros( + num_reqs, rows_per_req, dtype=torch.int64, device=device + ) + draft_sampled_2d[:, 1:] = proposals + draft_sampled = draft_sampled_2d.flatten() + cu_num_logits = ( + torch.arange(num_reqs + 1, dtype=torch.int32, device=device) * rows_per_req + ) + idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device) + expanded_idx_mapping = idx_mapping.repeat_interleave(rows_per_req) + expanded_local_pos = torch.arange( + rows_per_req, dtype=torch.int32, device=device + ).repeat(num_reqs) + positions = torch.arange(num_logits, dtype=torch.int64, device=device) + 4096 + temperature = torch.ones(num_reqs, dtype=torch.float32, device=device) + seeds = torch.arange(100, 100 + num_reqs, dtype=torch.int64, device=device) + + dense_sampled, dense_lengths = rejection_sample( + target_dense, + draft_dense, + draft_sampled, + cu_num_logits, + positions, + idx_mapping, + expanded_idx_mapping, + expanded_local_pos, + temperature, + seeds, + num_steps, + ) + sparse_sampled, sparse_lengths = dflash2_sparse_topk_rejection_sample( + target_ids, + target_logits, + draft_ids, + draft_logits, + draft_sampled, + cu_num_logits, + positions, + idx_mapping, + temperature, + torch.full((num_reqs,), top_p, dtype=torch.float32, device=device), + seeds, + num_steps, + ) + + assert torch.equal(sparse_lengths, dense_lengths) + steps = torch.arange(rows_per_req, device=device).unsqueeze(0) + valid = steps < dense_lengths[:, None] + assert torch.equal(sparse_sampled[valid], dense_sampled[valid]) diff --git a/tests/v1/spec_decode/test_dflash2_selector_analysis.py b/tests/v1/spec_decode/test_dflash2_selector_analysis.py new file mode 100644 index 0000000000..3316e78b6c --- /dev/null +++ b/tests/v1/spec_decode/test_dflash2_selector_analysis.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import importlib.util +import sys +from pathlib import Path + +import torch + +_ANALYZER_PATH = ( + Path(__file__).parents[3] / "benchmarks/analyze_dflash2_selector_alignment.py" +) +_SPEC = importlib.util.spec_from_file_location( + "analyze_dflash2_selector_alignment", _ANALYZER_PATH +) +assert _SPEC is not None and _SPEC.loader is not None +_ANALYZER = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = _ANALYZER +_SPEC.loader.exec_module(_ANALYZER) + + +def _record(): + candidate_ids = torch.tensor([[10, 11], [20, 21]]) + lattice = torch.tensor( + [ + [[0.0, 0.0], [0.0, 0.0]], + [[0.0, 4.0], [0.0, -4.0]], + ], + dtype=torch.float64, + ) + return _ANALYZER.AlignmentRecord( + path=Path("synthetic.pt"), + step=1, + temperature=1.0, + top_p=1.0, + target_topk_ids=candidate_ids, + target_topk_logits=torch.zeros((2, 2), dtype=torch.float64), + candidate_ids=candidate_ids, + realized_logits=torch.stack((lattice[0, 0], lattice[1, 0])), + unary_logits=torch.zeros((2, 2), dtype=torch.float64), + lattice_scores=lattice, + draft_sampled=torch.tensor([1, 10, 20]), + num_sampled=2, + ) + + +def test_compact_top_p_keeps_the_crossing_token(): + probs = _ANALYZER._compact_probs(torch.log(torch.tensor([0.6, 0.3, 0.1])), 0.7) + torch.testing.assert_close( + probs, torch.tensor([2 / 3, 1 / 3, 0.0], dtype=torch.float64) + ) + + +def test_beta_zero_reconstructs_realized_selector_rows(): + record = _record() + current = _ANALYZER._proposal_probs( + record, + _ANALYZER.ProposalConfig(name="current", use_cached_logits=True), + ) + reconstructed = _ANALYZER._proposal_probs( + record, + _ANALYZER.ProposalConfig(name="beta-zero"), + ) + + for actual, expected in zip(reconstructed, current): + torch.testing.assert_close(actual, expected) + + +def test_future_message_can_prefer_a_better_supported_branch(): + record = _record() + local = _ANALYZER._proposal_probs( + record, + _ANALYZER.ProposalConfig(name="local"), + )[0] + global_chain = _ANALYZER._proposal_probs( + record, + _ANALYZER.ProposalConfig(name="global", future_beta=1.0), + )[0] + + assert local[0] == local[1] + assert global_chain[0] > global_chain[1] + + +def test_greedy_mixture_is_normalized_and_moves_exact_mass(): + record = _record() + baseline = _ANALYZER._proposal_probs( + record, + _ANALYZER.ProposalConfig(name="baseline", use_cached_logits=True), + )[1] + mixed = _ANALYZER._proposal_probs( + record, + _ANALYZER.ProposalConfig( + name="mixture", + greedy_mix=0.25, + use_cached_logits=True, + ), + )[1] + + expected = baseline * 0.75 + expected[torch.argmax(baseline)] += 0.25 + torch.testing.assert_close(mixed, expected) + torch.testing.assert_close(mixed.sum(), torch.tensor(1.0, dtype=torch.float64)) diff --git a/tests/v1/spec_decode/test_dflash_mrv2_config.py b/tests/v1/spec_decode/test_dflash_mrv2_config.py new file mode 100644 index 0000000000..7a0292d9d2 --- /dev/null +++ b/tests/v1/spec_decode/test_dflash_mrv2_config.py @@ -0,0 +1,199 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from dataclasses import dataclass +from types import SimpleNamespace + +import pytest +import torch + +from vllm.config.speculative import SpeculativeConfig +from vllm.config.vllm import VllmConfig +from vllm.model_executor.models.qwen3_dflash import DFlashQwen3ForCausalLM +from vllm.v1.spec_decode.dflash import DFlashProposer +from vllm.v1.spec_decode.llm_base_proposer import SpecDecodeBaseProposer +from vllm.v1.worker.gpu.attn_utils import build_attn_metadata +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 + + +def _config(method: str, num_speculative_tokens: int = 7) -> SpeculativeConfig: + config = object.__new__(SpeculativeConfig) + config.method = method + config.num_speculative_tokens = num_speculative_tokens + config.parallel_drafting = True + return config + + +def test_dflash_and_ddtree_routes_are_disjoint() -> None: + mrv2 = _config("dflash") + ddtree = _config("dflash_ddtree") + + assert mrv2.use_dflash() + assert not mrv2.use_dflash_ddtree() + assert mrv2.use_dflash_family() + assert not ddtree.use_dflash() + assert ddtree.use_dflash_ddtree() + assert ddtree.use_dflash_family() + + +def test_mrv2_dflash_reserves_all_mask_slots() -> None: + assert _config("dflash").max_num_new_slots_for_drafting == 7 + # The retained DDTree path keeps its existing flat-parallel slot contract. + assert _config("dflash_ddtree").max_num_new_slots_for_drafting == 6 + + +def test_dflash_forces_v2_and_rejects_explicit_v1(monkeypatch) -> None: + fake = SimpleNamespace( + speculative_config=SimpleNamespace(use_dflash=lambda: True), + ) + monkeypatch.setattr("vllm.config.vllm.envs.VLLM_USE_V2_MODEL_RUNNER", None) + assert VllmConfig.use_v2_model_runner.fget(fake) + + monkeypatch.setattr("vllm.config.vllm.envs.VLLM_USE_V2_MODEL_RUNNER", False) + with pytest.raises(ValueError, match="implemented only by Model Runner V2"): + VllmConfig.use_v2_model_runner.fget(fake) + + +def test_dflash_target_layers_use_boundary_indices() -> None: + spec_config = SimpleNamespace( + draft_model_config=SimpleNamespace( + hf_config=SimpleNamespace(target_layer_ids=[5, 19, 33, 47, 61]) + ) + ) + assert get_eagle3_aux_layers_from_config(spec_config) == (6, 20, 34, 48, 62) + + +def test_parallel_drafting_token_id_prefers_dflash_config() -> None: + config = SimpleNamespace( + dflash_config={"mask_token_id": 248070}, + pard_token=1, + ) + assert get_parallel_drafting_token_id(config) == 248070 + + +def test_ddtree_topk_adapter_uses_official_logits_processor() -> None: + hidden_states = torch.randn(2, 4) + expected_ids = torch.tensor([[3, 2], [1, 0]]) + expected_logprobs = torch.tensor([[-0.1, -0.2], [-0.3, -0.4]]) + + class Processor: + def get_topk_tokens_and_logprobs(self, head, states, top_k): + assert head == "head" + assert states is hidden_states + assert top_k == 2 + return expected_ids, expected_logprobs + + model = SimpleNamespace( + draft_id_to_target_id=None, + logits_processor=Processor(), + lm_head="head", + ) + actual = DFlashQwen3ForCausalLM.get_topk_tokens_and_logprobs( + model, hidden_states, 2 + ) + assert actual is not None + assert actual[0] is expected_ids + assert actual[1] is expected_logprobs + + +def test_dflash_aux_projection_matches_fc_weight_dtype() -> None: + class FC: + input_size = 4 + weight = torch.empty((4, 4), dtype=torch.float16) + + def __call__(self, hidden_states: torch.Tensor) -> torch.Tensor: + assert hidden_states.dtype == self.weight.dtype + return hidden_states + + model = SimpleNamespace( + model=SimpleNamespace(use_aux_hidden_state=True, fc=FC()), + ) + hidden_states = torch.ones((2, 4), dtype=torch.float32) + + output = DFlashQwen3ForCausalLM.combine_hidden_states(model, hidden_states) + + assert output.dtype == torch.float16 + + +def test_draft_kv_dtype_is_public_and_defaults_to_inherit() -> None: + assert SpeculativeConfig.kv_cache_dtype is None + + +def test_ddtree_draft_config_combines_sm70_kv_and_non_causal(monkeypatch) -> None: + @dataclass(frozen=True) + class CacheConfig: + cache_dtype: str + + @dataclass(frozen=True) + class AttentionConfig: + use_non_causal: bool + + @dataclass(frozen=True) + class DraftConfig: + cache_config: CacheConfig + attention_config: AttentionConfig + + base = DraftConfig( + cache_config=CacheConfig(cache_dtype="fp8_e5m2"), + attention_config=AttentionConfig(use_non_causal=False), + ) + proposer = object.__new__(DFlashProposer) + proposer.speculative_config = SimpleNamespace(kv_cache_dtype=None) + proposer.dflash_causal = False + + monkeypatch.setattr( + SpecDecodeBaseProposer, + "_create_draft_vllm_config", + lambda _self: base, + ) + monkeypatch.setattr( + "vllm.v1.spec_decode.dflash.current_platform", + SimpleNamespace( + is_cuda=lambda: True, + is_device_capability=lambda capability: capability == 70, + ), + ) + + draft = DFlashProposer._create_draft_vllm_config(proposer) + + assert draft.cache_config.cache_dtype == "auto" + assert draft.attention_config.use_non_causal is True + + +def test_draft_attention_metadata_resolves_causal_per_kv_group() -> None: + class Builder: + def build(self, common_prefix_len, common_attn_metadata, **_kwargs): + assert common_prefix_len == 0 + return common_attn_metadata + + class Group: + def __init__(self, layer_name: str): + self.layer_names = [layer_name] + self.builder = Builder() + + def get_metadata_builder(self, _index: int): + return self.builder + + attn_metadata = build_attn_metadata( + attn_groups=[[Group("layer.0")], [Group("layer.1")]], + num_reqs=1, + num_tokens=1, + query_start_loc_gpu=torch.tensor([0, 1], dtype=torch.int32), + query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.int32), + max_query_len=1, + seq_lens=torch.tensor([1], dtype=torch.int32), + max_seq_len=1, + block_tables=[ + torch.zeros((1, 1), dtype=torch.int32), + torch.zeros((1, 1), dtype=torch.int32), + ], + slot_mappings=torch.zeros((2, 1), dtype=torch.int64), + kv_cache_config=SimpleNamespace(kv_cache_groups=[object(), object()]), + causal={0: False, 1: True}, + ) + + assert attn_metadata["layer.0"].causal is False + assert attn_metadata["layer.1"].causal is True diff --git a/tests/v1/spec_decode/test_dspark.py b/tests/v1/spec_decode/test_dspark.py index aaddd48e98..3f3251c7b9 100644 --- a/tests/v1/spec_decode/test_dspark.py +++ b/tests/v1/spec_decode/test_dspark.py @@ -1,20 +1,100 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from pathlib import Path from types import MethodType, SimpleNamespace import pytest import torch import torch.nn as nn +from vllm.config import ModelConfig, ParallelConfig, SpeculativeConfig from vllm.models.deepseek_v4.nvidia.dspark import ( DSparkDeepseekV4ForCausalLM, DSparkDeepseekV4Model, DSparkMarkovHead, ) +from vllm.platforms import current_platform +from vllm.transformers_utils.configs.deepseek_v4 import DeepseekV4Config from vllm.v1.spec_decode.dspark import DSparkProposer +def _make_dspark_speculative_config( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + *, + pp_size: int, + tp_size: int, +) -> tuple[ParallelConfig, SpeculativeConfig]: + monkeypatch.setattr(current_platform, "device_count", lambda: 8) + DeepseekV4Config( + architectures=["DeepseekV4ForCausalLM"], + hidden_size=4096, + num_hidden_layers=43, + num_attention_heads=64, + num_key_value_heads=1, + vocab_size=129280, + num_nextn_predict_layers=1, + dspark_target_layer_ids=[40, 41, 42], + dspark_block_size=5, + dspark_noise_token_id=128799, + dspark_markov_rank=256, + ).save_pretrained(tmp_path) + target_model_config = ModelConfig( + model=str(tmp_path), + tokenizer=str(tmp_path), + runner="generate", + dtype=torch.float16, + max_model_len=2048, + ) + target_parallel_config = ParallelConfig( + pipeline_parallel_size=pp_size, + tensor_parallel_size=tp_size, + ) + speculative_config = SpeculativeConfig( + target_model_config=target_model_config, + target_parallel_config=target_parallel_config, + method="dspark", + num_speculative_tokens=7, + ) + return target_parallel_config, speculative_config + + +@pytest.mark.parametrize(("pp_size", "tp_size"), [(2, 4), (4, 2)]) +def test_dspark_draft_is_local_to_last_pipeline_stage( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + pp_size: int, + tp_size: int, +) -> None: + target_parallel_config, speculative_config = _make_dspark_speculative_config( + tmp_path, + monkeypatch, + pp_size=pp_size, + tp_size=tp_size, + ) + + assert target_parallel_config.pipeline_parallel_size == pp_size + assert speculative_config.draft_parallel_config.pipeline_parallel_size == 1 + assert speculative_config.draft_parallel_config.tensor_parallel_size == tp_size + assert speculative_config.draft_model_config.architectures == ["DSparkDraftModel"] + + +def test_dspark_rejects_target_layers_outside_final_pipeline_stage( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("VLLM_PP_LAYER_PARTITION", "42,1") + + with pytest.raises(ValueError, match="final target pipeline stage"): + _make_dspark_speculative_config( + tmp_path, + monkeypatch, + pp_size=2, + tp_size=4, + ) + + def test_deepseek_v4_dspark_checkpoint_name_mapping() -> None: remap = DSparkDeepseekV4ForCausalLM._remap_dspark_name @@ -27,7 +107,9 @@ def test_deepseek_v4_dspark_checkpoint_name_mapping() -> None: assert remap("mtp.2.markov_head.markov_w1.weight") == ( "model.markov_head.markov_w1.weight" ) - assert remap("mtp.2.confidence_head.proj.weight") is None + assert remap("mtp.2.confidence_head.proj.weight") == ( + "model.confidence_head.proj.weight" + ) class _FakeDSparkModel: @@ -47,6 +129,13 @@ def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor: bias.scatter_(1, ((previous + 1) % self.vocab_size).unsqueeze(1), 10.0) return bias + @staticmethod + def confidence_logits( + hidden_states: torch.Tensor, + markov_embeds: torch.Tensor, + ) -> torch.Tensor: + return hidden_states.sum(dim=-1) + markov_embeds.sum(dim=-1) + @staticmethod def map_draft_to_target(token_ids: torch.Tensor) -> torch.Tensor: return token_ids @@ -61,6 +150,8 @@ def test_dspark_markov_sampling_is_sequential( proposer.model = _FakeDSparkModel() proposer._enable_probabilistic_draft_probs = False proposer._static_draft_vocab = None + proposer.collect_confidence_logits = False + proposer.confidence_scheduling_enabled = False proposer.input_ids = torch.zeros(2 * num_speculative_tokens, dtype=torch.int32) proposer.input_ids[0] = 1 proposer.input_ids[num_speculative_tokens] = 4 @@ -83,6 +174,8 @@ def test_dspark_markov_sampling_is_sequential( for anchor in (1, 4) ] assert output.view(2, num_speculative_tokens).tolist() == expected + assert proposer._last_confidence_logits is None + assert proposer._last_verification_lengths is None @pytest.mark.parametrize("num_speculative_tokens", [5, 7]) @@ -94,6 +187,8 @@ def test_dspark_probabilistic_sampling_returns_sequential_probs( proposer.model = _FakeDSparkModel() proposer._enable_probabilistic_draft_probs = True proposer._static_draft_vocab = None + proposer.collect_confidence_logits = True + proposer.confidence_scheduling_enabled = True proposer.input_ids = torch.zeros(2 * num_speculative_tokens, dtype=torch.int32) proposer.input_ids[0] = 1 proposer.input_ids[num_speculative_tokens] = 4 @@ -132,6 +227,37 @@ def sample_from_logits( for anchor in (1, 4) ] assert output.view(2, num_speculative_tokens).tolist() == expected + assert proposer._last_confidence_logits.shape == (2, num_speculative_tokens) + assert proposer._last_verification_lengths.shape == (2,) + confidence_logits = proposer.take_last_confidence_logits() + verification_lengths = proposer.take_last_verification_lengths() + assert confidence_logits is not None + assert verification_lengths is not None + assert proposer.take_last_confidence_logits() is None + assert proposer.take_last_verification_lengths() is None + + +def test_dspark_static_verification_cap_skips_confidence_projection() -> None: + proposer = object.__new__(DSparkProposer) + proposer.num_speculative_tokens = 5 + proposer.model = _FakeDSparkModel() + proposer._enable_probabilistic_draft_probs = False + proposer._static_draft_vocab = None + proposer.collect_confidence_logits = False + proposer.confidence_scheduling_enabled = True + proposer._max_verification_tokens = 2 + proposer.input_ids = torch.zeros(5, dtype=torch.int32) + proposer._anchor_indices = torch.tensor([0], dtype=torch.int64) + + output, draft_probs = proposer._sample_draft_tokens( + torch.zeros(5, 2), + sampling_metadata=None, # type: ignore[arg-type] + ) + + assert output.shape == (5,) + assert draft_probs is None + assert proposer._last_confidence_logits is None + assert proposer._last_verification_lengths.tolist() == [2] def test_dspark_replicated_linears_return_tensors() -> None: @@ -154,3 +280,42 @@ def test_dspark_replicated_linears_return_tensors() -> None: markov_head.markov_w2 = nn.Identity() markov_embed = torch.randn(3, 4) assert torch.equal(markov_head.bias(markov_embed), markov_embed) + + +def test_dspark_confidence_projection_rejects_missing_weight() -> None: + draft_model = object.__new__(DSparkDeepseekV4ForCausalLM) + nn.Module.__init__(draft_model) + draft_model._confidence_head_loaded = False + + with pytest.raises(RuntimeError, match="did not provide confidence_head"): + draft_model.confidence_logits(torch.empty(1, 1), torch.empty(1, 1)) + + +def test_dspark_confidence_scheduler_keeps_only_confident_prefix() -> None: + proposer = object.__new__(DSparkProposer) + proposer._confidence_temperatures = torch.tensor([[1.0, 2.0, 1.0, 1.0, 1.0]]) + proposer._confidence_threshold = 0.7 + proposer._max_verification_tokens = 4 + logits = torch.tensor( + [ + [2.0, 2.0, 0.0, 4.0, 4.0], + [2.0, 4.0, 4.0, 4.0, 4.0], + [0.0, 4.0, 4.0, 4.0, 4.0], + ] + ) + + lengths = proposer._select_verification_lengths(logits) + + assert lengths.dtype == torch.int32 + assert lengths.tolist() == [2, 4, 0] + + +def test_dspark_confidence_scheduler_zero_threshold_uses_cap() -> None: + proposer = object.__new__(DSparkProposer) + proposer._confidence_temperatures = torch.ones((1, 5)) + proposer._confidence_threshold = 0.0 + proposer._max_verification_tokens = 2 + + lengths = proposer._select_verification_lengths(torch.randn(3, 5)) + + assert lengths.tolist() == [2, 2, 2] diff --git a/tests/v1/spec_decode/test_eagle.py b/tests/v1/spec_decode/test_eagle.py index c13de6d4f7..4de8315c06 100644 --- a/tests/v1/spec_decode/test_eagle.py +++ b/tests/v1/spec_decode/test_eagle.py @@ -63,7 +63,7 @@ def _create_proposer( elif method == "draft_model": target_model_dir = model_dir draft_model_dir = ar_draft_model_dir - elif method == "dflash": + elif method == "dflash_ddtree": target_model_dir = dflash_target_dir draft_model_dir = dflash_dir else: @@ -73,7 +73,7 @@ def _create_proposer( model=target_model_dir, runner="generate", max_model_len=100, - trust_remote_code=(method == "dflash"), + trust_remote_code=(method == "dflash_ddtree"), ) speculative_config = SpeculativeConfig( @@ -105,7 +105,7 @@ def _create_proposer( attention_config=AttentionConfig(backend=attention_backend), ) - if method == "dflash": + if method == "dflash_ddtree": proposer = DFlashProposer(vllm_config=vllm_config, device=device) elif "eagle" in method: proposer = EagleProposer(vllm_config=vllm_config, device=device) @@ -1132,7 +1132,10 @@ def test_set_inputs_first_pass_dflash(): device = torch.device(current_platform.device_type) num_speculative_tokens = 3 - proposer = _create_proposer("dflash", num_speculative_tokens) + proposer = _create_proposer("dflash_ddtree", num_speculative_tokens) + # Production initialization assigns the draft KV cache group before this + # method runs. This isolated input-preparation test supplies that contract. + proposer.kv_cache_gid = 0 mask_token_id = proposer.parallel_drafting_token_id # Setup batch with 3 requests diff --git a/tests/v1/spec_decode/test_mtp.py b/tests/v1/spec_decode/test_mtp.py index 7c478f81d8..ec3fca1d0c 100644 --- a/tests/v1/spec_decode/test_mtp.py +++ b/tests/v1/spec_decode/test_mtp.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace from unittest import mock import pytest @@ -12,6 +13,7 @@ create_standard_kv_cache_spec, try_get_attention_backend, ) +from vllm import envs from vllm.config import ( CacheConfig, DeviceConfig, @@ -26,11 +28,121 @@ from vllm.platforms import current_platform from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.spec_decode.eagle import EagleProposer +from vllm.v1.spec_decode.llm_base_proposer import ( + SpecDecodeBaseProposer, + _sm70_mtp_hotpath_warmup_batch_sizes, + _sm70_mtp_moe_warmup_sizes, +) mimo_7b_dir = "XiaomiMiMo/MiMo-7B-Base" DEVICE_TYPE = current_platform.device_type +def test_sm70_mtp_moe_warmup_sizes(): + assert _sm70_mtp_moe_warmup_sizes(256, 8, 8192) == (9, 33, 257) + assert _sm70_mtp_moe_warmup_sizes(256, 8, 32) == (9,) + assert _sm70_mtp_moe_warmup_sizes(256, 8, 8192, 16) == ( + *range(1, 17), + 33, + 257, + ) + assert _sm70_mtp_moe_warmup_sizes(0, 8, 8192) == () + + +def test_sm70_mtp_hotpath_warmup_batch_sizes(): + assert _sm70_mtp_hotpath_warmup_batch_sizes(1) == (1,) + assert _sm70_mtp_hotpath_warmup_batch_sizes(4) == (4,) + assert _sm70_mtp_hotpath_warmup_batch_sizes(4, include_alternate=True) == (1, 4) + assert _sm70_mtp_hotpath_warmup_batch_sizes(64, include_alternate=True) == (1, 4) + + +def test_sm70_mtp_concurrency_warmup_is_default_off(monkeypatch): + monkeypatch.delenv("VLLM_SM70_MTP_CONCURRENCY_WARMUP", raising=False) + envs.disable_envs_cache() + try: + assert not envs.VLLM_SM70_MTP_CONCURRENCY_WARMUP + finally: + envs.disable_envs_cache() + + +def test_sm70_mtp_moe_warmup_runs_each_shape_once(): + dummy_run = mock.Mock() + draft_model_config = SimpleNamespace( + is_moe=True, + hf_text_config=SimpleNamespace( + num_experts_per_tok=8, + moe_intermediate_size=512, + ), + get_num_experts=lambda: 256, + get_hidden_size=lambda: 4096, + ) + proposer = SimpleNamespace( + method="mtp", + device=torch.device("cuda"), + draft_model_config=draft_model_config, + vllm_config=SimpleNamespace( + parallel_config=SimpleNamespace(tensor_parallel_size=4) + ), + max_num_tokens=8192, + dummy_run=dummy_run, + _sm70_mtp_moe_warmed=False, + ) + + with ( + mock.patch.object(current_platform, "is_device_capability", return_value=True), + mock.patch.object(torch.accelerator, "synchronize"), + ): + assert SpecDecodeBaseProposer.warmup_sm70_mtp_moe_kernels(proposer) == ( + "mtp_draft_moe", + ) + + assert [call.args[0] for call in dummy_run.call_args_list] == [9, 33, 257] + assert all(call.kwargs["spec_step_idx"] == 0 for call in dummy_run.call_args_list) + + +def test_sm70_mtp_moe_warmup_covers_opted_in_exact_tiles(): + dummy_run = mock.Mock() + draft_model_config = SimpleNamespace( + is_moe=True, + hf_text_config=SimpleNamespace( + num_experts_per_tok=8, + moe_intermediate_size=512, + ), + get_num_experts=lambda: 256, + get_hidden_size=lambda: 2048, + ) + proposer = SimpleNamespace( + method="mtp", + device=torch.device("cuda"), + draft_model_config=draft_model_config, + vllm_config=SimpleNamespace( + parallel_config=SimpleNamespace(tensor_parallel_size=4) + ), + max_num_tokens=8192, + dummy_run=dummy_run, + _sm70_mtp_moe_warmed=False, + ) + + with ( + mock.patch.object(current_platform, "is_device_capability", return_value=True), + mock.patch.object(torch.accelerator, "synchronize"), + mock.patch.object(envs, "VLLM_SM70_MTP_MOE_TUNED_CONFIG", True), + ): + assert SpecDecodeBaseProposer.warmup_sm70_mtp_moe_kernels(proposer) == ( + "mtp_draft_moe", + ) + + assert [call.args[0] for call in dummy_run.call_args_list] == [ + *range(1, 17), + 33, + 257, + 1, + 2, + 9, + 10, + ] + + def _create_mtp_proposer(num_speculative_tokens: int) -> EagleProposer: """Create an MTP proposer with unified model configuration.""" model_config = ModelConfig( diff --git a/tests/v1/spec_decode/test_mtp_draft_prob_alignment.py b/tests/v1/spec_decode/test_mtp_draft_prob_alignment.py index 66ddb3429c..26316d35d8 100644 --- a/tests/v1/spec_decode/test_mtp_draft_prob_alignment.py +++ b/tests/v1/spec_decode/test_mtp_draft_prob_alignment.py @@ -4,7 +4,10 @@ import pytest import torch -from vllm.v1.spec_decode.draft_prob_alignment import get_aligned_draft_probs +from vllm.v1.spec_decode.draft_prob_alignment import ( + get_aligned_draft_probs, + get_aligned_draft_scalar_values, +) from vllm.v1.spec_decode.metadata import SpecDecodeMetadata @@ -115,3 +118,27 @@ def test_get_spec_decode_draft_probs_rejects_cached_token_mismatch(): draft_prob_token_ids=torch.tensor([[1, 9]], dtype=torch.int32), spec_decode_metadata=_metadata([[1, 2]]), ) + + +def test_get_aligned_draft_scalar_values_binds_prefixes_by_req_id(): + values = torch.tensor([[0.1, 0.2, 0.3], [1.1, 1.2, 1.3]]) + actual = get_aligned_draft_scalar_values( + req_ids=["b", "a"], + values=values, + value_req_ids=["a", "b"], + value_token_ids=torch.tensor([[10, 11, 12], [20, 21, 22]], dtype=torch.int32), + spec_decode_metadata=_metadata([[20], [10, 11]]), + ) + + assert torch.equal(actual, torch.tensor([1.1, 0.1, 0.2])) + + +def test_get_aligned_draft_scalar_values_rejects_token_mismatch(): + with pytest.raises(RuntimeError, match="do not match verifier draft tokens"): + get_aligned_draft_scalar_values( + req_ids=["a"], + values=torch.tensor([[0.1, 0.2]]), + value_req_ids=["a"], + value_token_ids=[[1, 9]], + spec_decode_metadata=_metadata([[1, 2]]), + ) diff --git a/tests/v1/spec_decode/test_rejection_sampler_utils.py b/tests/v1/spec_decode/test_rejection_sampler_utils.py index 28f6044de8..83bb10599e 100644 --- a/tests/v1/spec_decode/test_rejection_sampler_utils.py +++ b/tests/v1/spec_decode/test_rejection_sampler_utils.py @@ -7,6 +7,7 @@ import torch from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import ( + dflash2_sparse_topk_rejection_sample, rejection_sample, ) @@ -129,6 +130,130 @@ def _assert_distribution_match( ) +@pytest.mark.parametrize("num_speculative_steps", [3, 7]) +@pytest.mark.parametrize("top_p", [1.0, 0.95]) +def test_dflash2_sparse_topk_matches_dense_rejection( + num_speculative_steps: int, + top_p: float, +): + """Compact p/q support must preserve the dense DFlash2 decision path.""" + torch.manual_seed(20260823) + device = torch.device("cuda") + num_reqs = 4 + target_top_k = 20 + draft_top_k = 16 + rows_per_req = num_speculative_steps + 1 + num_logits = num_reqs * rows_per_req + + # Unique logits avoid making the reference depend on unspecified top-k + # ordering for equal FP16 values. Adversarial tie behavior is a separate + # candidate-set policy gate rather than a rejection-kernel property. + target_topk_ids = torch.stack( + [ + torch.randperm(VOCAB_SIZE, device=device)[:target_top_k] + for _ in range(num_logits) + ] + ) + target_topk_logits = ( + torch.randn(num_logits, target_top_k, device=device) + .sort(dim=-1, descending=True) + .values + + torch.arange(target_top_k, device=device).flip(0) * 1e-3 + ) + target_dense = torch.full( + (num_logits, VOCAB_SIZE), + -float("inf"), + dtype=torch.float32, + device=device, + ) + target_dense.scatter_(1, target_topk_ids, target_topk_logits) + + top_k_tensor = torch.full( + (num_logits,), target_top_k, dtype=torch.int32, device=device + ) + 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) + + draft_topk_ids = torch.stack( + [ + torch.stack( + [ + torch.randperm(VOCAB_SIZE, device=device)[:draft_top_k] + for _ in range(num_speculative_steps) + ] + ) + for _ in range(num_reqs) + ] + ) + draft_topk_logits = torch.randn( + num_reqs, + num_speculative_steps, + draft_top_k, + dtype=torch.float32, + device=device, + ) + draft_dense = torch.full( + (num_reqs, num_speculative_steps, VOCAB_SIZE), + -float("inf"), + dtype=torch.float32, + device=device, + ) + draft_dense.scatter_(2, draft_topk_ids, draft_topk_logits) + + draft_sampled_2d = torch.zeros( + num_reqs, rows_per_req, dtype=torch.int64, device=device + ) + draft_sampled_2d[:, 1:] = draft_topk_ids[:, :, 0] + draft_sampled = draft_sampled_2d.reshape(-1) + cu_num_logits = ( + torch.arange(num_reqs + 1, dtype=torch.int32, device=device) * rows_per_req + ) + idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device) + expanded_idx_mapping = idx_mapping.repeat_interleave(rows_per_req) + expanded_local_pos = torch.arange( + 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) + 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) + + dense_sampled, dense_num_sampled = rejection_sample( + processed_target, + draft_dense, + draft_sampled, + cu_num_logits, + pos, + idx_mapping, + expanded_idx_mapping, + expanded_local_pos, + temperature, + seeds, + num_speculative_steps, + ) + sparse_sampled, sparse_num_sampled = dflash2_sparse_topk_rejection_sample( + target_topk_ids, + target_topk_logits, + draft_topk_ids, + draft_topk_logits, + draft_sampled, + cu_num_logits, + pos, + idx_mapping, + temperature, + top_p_per_req, + seeds, + num_speculative_steps, + ) + + assert torch.equal(sparse_num_sampled, dense_num_sampled) + steps = torch.arange(rows_per_req, device=device).unsqueeze(0) + valid = steps < dense_num_sampled.unsqueeze(1) + assert torch.equal(sparse_sampled[valid], dense_sampled[valid]) + + @pytest.mark.parametrize( "num_speculative_steps,temperature", [ diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 8d7826116e..b2c7d06de1 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -26,6 +26,7 @@ from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2 from vllm.platforms import current_platform from vllm.sampling_params import SamplingParams +from vllm.sequence import IntermediateTensors from vllm.utils.mem_constants import GiB_bytes from vllm.utils.system_utils import update_environment_variables from vllm.utils.torch_utils import set_random_seed @@ -51,6 +52,27 @@ DEVICE_TYPE = current_platform.device_type +def test_unwrap_pipeline_intermediate_hidden_states_uses_semantic_key(): + hidden_states = torch.zeros((2, 3)) + residual = torch.ones_like(hidden_states) + output = IntermediateTensors({"residual": residual, "hidden_states": hidden_states}) + + assert ( + gpu_model_runner_module._unwrap_pipeline_intermediate_hidden_states(output) + is hidden_states + ) + assert ( + gpu_model_runner_module._unwrap_pipeline_intermediate_hidden_states( + hidden_states + ) + is hidden_states + ) + with pytest.raises(RuntimeError, match="must contain a 'hidden_states' tensor"): + gpu_model_runner_module._unwrap_pipeline_intermediate_hidden_states( + IntermediateTensors({"residual": residual}) + ) + + def initialize_kv_cache(runner: GPUModelRunner): """ Only perform necessary steps in GPUModelRunner.initialize_kv_cache() @@ -795,6 +817,256 @@ def test_sample_passes_reordered_draft_probs_to_rejection_sampler(): assert torch.equal(passed_draft_probs, expected_draft_probs) +def test_dspark_verification_lengths_trim_scheduler_drafts(monkeypatch): + runner = object.__new__(GPUModelRunner) + runner._draft_token_ids = torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]) + runner._draft_token_req_ids = ["a", "b"] + runner.draft_token_ids_event = object() + runner.draft_token_ids_cpu = runner._draft_token_ids.clone() + runner.dspark_confidence_scheduling = True + runner._dspark_verification_lengths_cpu = torch.tensor([2, 0], dtype=torch.int32) + monkeypatch.setattr( + "vllm.v1.worker.gpu_model_runner.sm70_trace_event_sync", + lambda *_args, **_kwargs: None, + ) + + token_ids, req_ids = runner._get_draft_token_ids_cpu() + + assert req_ids == ["a", "b"] + assert token_ids == [[1, 2], []] + + +@pytest.mark.skip_global_cleanup +def test_sync_mamba_accepted_token_state_tracks_request_ownership(): + runner = object.__new__(GPUModelRunner) + move_state = object() + resumed_state = object() + replaced_old_state = object() + replaced_new_state = object() + runner._mamba_accepted_token_state_rows = { + "old": (0, object()), + "resumed": (1, resumed_state), + "move": (2, move_state), + "same-id": (3, replaced_old_state), + } + runner.num_accepted_tokens = SimpleNamespace( + np=np.array([5, 3, 4, 5, 0], dtype=np.int32) + ) + runner.spec_state_slot_selectors = SimpleNamespace( + np=np.array([5, 4, 2, 5, 0], dtype=np.int32) + ) + runner.input_batch = SimpleNamespace( + req_ids=["new", "move", "resumed", "same-id"], + num_accepted_tokens_cpu=np.full(5, 9, dtype=np.int32), + spec_num_accepted_tokens_cpu=np.full(5, 9, dtype=np.int32), + ) + runner.requests = { + "new": object(), + "move": move_state, + "resumed": resumed_state, + "same-id": replaced_new_state, + } + scheduler_output = SimpleNamespace( + scheduled_new_reqs=[SimpleNamespace(req_id="new")], + scheduled_cached_reqs=SimpleNamespace(resumed_req_ids={"resumed"}), + ) + + GPUModelRunner._sync_mamba_accepted_token_state( + runner, scheduler_output, num_reqs=4 + ) + + # Only the continuing request keeps its prior accepted-token state. A new + # request, a resumed request, and a replacement using the same request ID + # all start with zero speculative bias (the encoded value 1). + expected_counts = np.array([1, 4, 1, 1], dtype=np.int32) + expected_selectors = np.array([1, 2, 1, 1], dtype=np.int32) + np.testing.assert_array_equal(runner.num_accepted_tokens.np[:4], expected_counts) + np.testing.assert_array_equal( + runner.spec_state_slot_selectors.np[:4], expected_selectors + ) + np.testing.assert_array_equal( + runner.input_batch.num_accepted_tokens_cpu[:4], expected_counts + ) + np.testing.assert_array_equal( + runner.input_batch.spec_num_accepted_tokens_cpu[:4], expected_selectors + ) + + +@pytest.mark.skip_global_cleanup +def test_mamba_align_postprocess_uses_runner_owned_d2h_buffers(monkeypatch): + runner = object.__new__(GPUModelRunner) + req_state = object() + runner.speculative_config = object() + runner.model_config = SimpleNamespace(is_hybrid=True) + runner.cache_config = SimpleNamespace(mamba_cache_mode="align") + runner.num_accepted_tokens = SimpleNamespace( + cpu=torch.zeros(1, dtype=torch.int32), + gpu=torch.zeros(1, dtype=torch.int32), + ) + runner.spec_state_slot_selectors = SimpleNamespace( + cpu=torch.zeros(1, dtype=torch.int32), + gpu=torch.zeros(1, dtype=torch.int32), + ) + runner.input_batch = SimpleNamespace( + req_ids=["req"], + num_accepted_tokens_cpu_tensor=torch.full((1,), 9, dtype=torch.int32), + spec_num_accepted_tokens_cpu_tensor=torch.full((1,), 9, dtype=torch.int32), + ) + runner.requests = {"req": req_state} + runner._get_mamba_bufs = Mock( + return_value=SimpleNamespace(postprocess_align=object()) + ) + runner.kv_cache_config = object() + runner.compilation_config = SimpleNamespace(static_forward_context=object()) + runner.model = SimpleNamespace(get_mamba_state_copy_func=Mock(return_value=())) + runner.mamba_state_idx = {} + runner.num_accepted_tokens_event = Mock() + postprocess = Mock() + monkeypatch.setattr( + gpu_model_runner_module.envs, + "VLLM_MAMBA_ALIGN_CPU_POSTPROCESS", + False, + ) + monkeypatch.setattr( + gpu_model_runner_module.mamba_utils, + "stage_postprocess_inputs_to_gpu", + Mock(), + ) + monkeypatch.setattr( + gpu_model_runner_module.mamba_utils, + "postprocess_mamba_align_gpu", + postprocess, + ) + + GPUModelRunner._update_states_after_model_execute( + runner, + output_token_ids=torch.tensor([[10, 11, -1]], dtype=torch.int32), + scheduler_output=SimpleNamespace(scheduled_ddtree_payloads={}), + ) + + kwargs = postprocess.call_args.kwargs + assert kwargs["num_accepted_tokens_cpu_tensor"] is runner.num_accepted_tokens.cpu + assert ( + kwargs["spec_num_accepted_tokens_cpu_tensor"] + is runner.spec_state_slot_selectors.cpu + ) + assert ( + kwargs["num_accepted_tokens_cpu_tensor"] + is not runner.input_batch.num_accepted_tokens_cpu_tensor + ) + assert ( + kwargs["spec_num_accepted_tokens_cpu_tensor"] + is not runner.input_batch.spec_num_accepted_tokens_cpu_tensor + ) + # The GPU path must not leave any asynchronous DMA targeting mutable + # InputBatch rows; both accepted-state arrays are remapped only after the + # event is synchronized in _prepare_inputs(). + assert torch.equal( + runner.input_batch.num_accepted_tokens_cpu_tensor, + torch.full((1,), 9, dtype=torch.int32), + ) + assert torch.equal( + runner.input_batch.spec_num_accepted_tokens_cpu_tensor, + torch.full((1,), 9, dtype=torch.int32), + ) + + +@pytest.mark.skip_global_cleanup +def test_mamba_align_cpu_postprocess_mirrors_runner_owned_snapshot(monkeypatch): + runner = object.__new__(GPUModelRunner) + req_state = object() + runner.speculative_config = object() + runner.model_config = SimpleNamespace(is_hybrid=True) + runner.cache_config = SimpleNamespace(mamba_cache_mode="align") + runner.num_accepted_tokens = SimpleNamespace( + cpu=torch.zeros(1, dtype=torch.int32), + gpu=torch.zeros(1, dtype=torch.int32), + ) + runner.spec_state_slot_selectors = SimpleNamespace( + cpu=torch.zeros(1, dtype=torch.int32), + gpu=torch.zeros(1, dtype=torch.int32), + ) + runner.input_batch = SimpleNamespace( + req_ids=["req"], + num_accepted_tokens_cpu_tensor=torch.zeros(1, dtype=torch.int32), + spec_num_accepted_tokens_cpu_tensor=torch.zeros(1, dtype=torch.int32), + ) + runner.requests = {"req": req_state} + runner._get_mamba_bufs = Mock(return_value=SimpleNamespace(preprocess=object())) + runner.kv_cache_config = object() + runner.compilation_config = SimpleNamespace(static_forward_context=object()) + runner.model = SimpleNamespace(get_mamba_state_copy_func=Mock(return_value=())) + runner.mamba_state_idx = {} + runner.num_accepted_tokens_event = Mock() + + def cpu_postprocess(*args, **kwargs): + runner.input_batch.num_accepted_tokens_cpu_tensor.fill_(3) + runner.input_batch.spec_num_accepted_tokens_cpu_tensor.fill_(2) + + monkeypatch.setattr( + gpu_model_runner_module.envs, + "VLLM_MAMBA_ALIGN_CPU_POSTPROCESS", + True, + ) + monkeypatch.setattr( + gpu_model_runner_module.mamba_utils, + "postprocess_mamba", + cpu_postprocess, + ) + + GPUModelRunner._update_states_after_model_execute( + runner, + output_token_ids=torch.tensor([[10, 11, -1]], dtype=torch.int32), + scheduler_output=SimpleNamespace(scheduled_ddtree_payloads={}), + ) + + assert runner.num_accepted_tokens.cpu.tolist() == [3] + assert runner.spec_state_slot_selectors.cpu.tolist() == [2] + + +@pytest.mark.skip_global_cleanup +def test_mamba_all_mode_uses_runner_owned_d2h_buffers(monkeypatch): + runner = object.__new__(GPUModelRunner) + req_state = object() + runner.speculative_config = object() + runner.model_config = SimpleNamespace(is_hybrid=True) + runner.cache_config = SimpleNamespace(mamba_cache_mode="all") + runner.num_spec_tokens = 4 + runner.num_accepted_tokens = SimpleNamespace( + cpu=torch.zeros(1, dtype=torch.int32), + gpu=torch.zeros(1, dtype=torch.int32), + ) + runner.spec_state_slot_selectors = SimpleNamespace( + cpu=torch.zeros(1, dtype=torch.int32), + gpu=torch.zeros(1, dtype=torch.int32), + ) + runner.input_batch = SimpleNamespace( + req_ids=["req"], + num_accepted_tokens_cpu_tensor=torch.full((1,), 9, dtype=torch.int32), + spec_num_accepted_tokens_cpu_tensor=torch.full((1,), 9, dtype=torch.int32), + ) + runner.requests = {"req": req_state} + runner.kv_cache_config = object() + runner.mamba_state_idx = {} + runner.num_accepted_tokens_event = Mock() + monkeypatch.setattr( + gpu_model_runner_module.mamba_utils, + "postprocess_mamba_all", + Mock(), + ) + + GPUModelRunner._update_states_after_model_execute( + runner, + output_token_ids=torch.tensor([[10, 11, -1]], dtype=torch.int32), + scheduler_output=SimpleNamespace(scheduled_ddtree_payloads={}), + ) + + assert runner.num_accepted_tokens.cpu.tolist() == [2] + assert runner.spec_state_slot_selectors.cpu.tolist() == [2] + assert runner.input_batch.num_accepted_tokens_cpu_tensor.tolist() == [9] + assert runner.input_batch.spec_num_accepted_tokens_cpu_tensor.tolist() == [9] + + def test_init_kv_cache_with_kv_sharing_invalid_target_layer_order(default_vllm_config): torch.set_default_dtype(torch.float16) layer_0 = "model.layers.0.self_attn.attn" diff --git a/tests/v1/worker/test_gpu_model_runner_v2_kv_zero.py b/tests/v1/worker/test_gpu_model_runner_v2_kv_zero.py new file mode 100644 index 0000000000..19e7c4fe03 --- /dev/null +++ b/tests/v1/worker/test_gpu_model_runner_v2_kv_zero.py @@ -0,0 +1,87 @@ +# 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.v1.worker.gpu import model_runner as mrv2 + + +class FakeKVBlockZeroer: + def __init__(self, device: torch.device, pin_memory: bool): + self.device = device + self.pin_memory = pin_memory + self.init_kwargs: dict[str, Any] | None = None + + def init_meta(self, **kwargs: Any) -> None: + kwargs["attn_groups_iter"] = list(kwargs["attn_groups_iter"]) + self.init_kwargs = kwargs + + +def test_v2_initializes_kv_zeroer_from_all_attention_groups(monkeypatch) -> None: + monkeypatch.setattr(mrv2, "KVBlockZeroer", FakeKVBlockZeroer) + monkeypatch.setattr(mrv2, "is_pin_memory_available", lambda: False) + + groups = [[object()], [object(), object()]] + static_forward_context = object() + runner = mrv2.GPUModelRunner.__new__(mrv2.GPUModelRunner) + runner.device = torch.device("cpu") + runner.attn_groups = groups + runner._kernel_block_sizes = [16, 32] + runner.cache_config = SimpleNamespace(cache_dtype="auto") + runner.compilation_config = SimpleNamespace( + static_forward_context=static_forward_context + ) + + runner._init_kv_zero_meta() + + zeroer = runner._kv_block_zeroer + assert isinstance(zeroer, FakeKVBlockZeroer) + assert zeroer.device == runner.device + assert zeroer.pin_memory is False + assert zeroer.init_kwargs == { + "attn_groups_iter": [*groups[0], *groups[1]], + "kernel_block_sizes": [16, 32], + "cache_dtype": "auto", + "runner_only_attn_layers": set(), + "static_forward_context": static_forward_context, + } + + +def test_v2_clears_new_cache_blocks_before_zero_token_return() -> None: + events: list[object] = [] + empty_output = object() + + def no_forward(_: Any) -> object: + events.append("no_forward") + return empty_output + + runner = mrv2.GPUModelRunner.__new__(mrv2.GPUModelRunner) + runner.finish_requests = lambda _: events.append("finish") + runner.free_states = lambda _: events.append("free") + runner.add_requests = lambda _: events.append("add") + runner.update_requests = lambda _: events.append("update") + runner.block_tables = SimpleNamespace( + apply_staged_writes=lambda: events.append("apply") + ) + runner._zero_block_ids = lambda block_ids: events.append(("zero", block_ids)) + runner.kv_connector = SimpleNamespace(no_forward=no_forward) + scheduler_output = SimpleNamespace( + total_num_scheduled_tokens=0, + new_block_ids_to_zero=[3, 7], + ) + + output = runner.execute_model(scheduler_output) + + assert output is empty_output + assert events == [ + "finish", + "free", + "add", + "update", + "apply", + ("zero", [3, 7]), + "no_forward", + ] diff --git a/tests/v1/worker/test_gpu_warmup_blocks.py b/tests/v1/worker/test_gpu_warmup_blocks.py new file mode 100644 index 0000000000..b4fcb8bd0a --- /dev/null +++ b/tests/v1/worker/test_gpu_warmup_blocks.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MRV2 warmup must reserve the same speculative KV tail as the scheduler.""" + +import pytest +import torch + +from vllm.config.speculative import SpeculativeConfig +from vllm.config.vllm import VllmConfig +from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec +from vllm.v1.worker.gpu.warmup import _reserved_block_count + +BLOCK_SIZE = 16 +MAX_MODEL_LEN = 1024 +NUM_SPEC_TOKENS = 7 + + +def _speculative_config(method: str) -> SpeculativeConfig: + config = object.__new__(SpeculativeConfig) + object.__setattr__(config, "method", method) + object.__setattr__(config, "num_speculative_tokens", NUM_SPEC_TOKENS) + object.__setattr__(config, "ddtree_disable_tree_verify", False) + object.__setattr__(config, "ddtree_budget", 24) + return config + + +class _Config: + speculative_config: SpeculativeConfig | None = None + diffusion_config = None + num_speculative_tokens = VllmConfig.num_speculative_tokens + num_lookahead_tokens = VllmConfig.num_lookahead_tokens + + +@pytest.mark.parametrize( + ("method", "expected"), + [ + ("dflash", NUM_SPEC_TOKENS + 1), + ("dflash_ddtree", 25), + ("eagle3", NUM_SPEC_TOKENS), + ("mtp", NUM_SPEC_TOKENS), + ("dspark", NUM_SPEC_TOKENS), + ("draft_model", NUM_SPEC_TOKENS), + ("ngram", 0), + ], +) +def test_num_lookahead_tokens_per_method(method: str, expected: int) -> None: + config = _Config() + config.speculative_config = _speculative_config(method) + + assert config.num_lookahead_tokens == expected + + +def test_num_lookahead_tokens_without_speculation() -> None: + assert _Config().num_lookahead_tokens == 0 + + +def _full_attention_spec() -> FullAttentionSpec: + return FullAttentionSpec( + block_size=BLOCK_SIZE, + num_kv_heads=1, + head_size=1, + dtype=torch.float16, + ) + + +def _mamba_spec(mode: str) -> MambaSpec: + return MambaSpec( + block_size=BLOCK_SIZE, + shapes=((1,),), + dtypes=(torch.float16,), + mamba_cache_mode=mode, + num_speculative_blocks=NUM_SPEC_TOKENS, + ) + + +@pytest.mark.parametrize( + ("spec", "expected"), + [ + (_full_attention_spec(), 2), + (_mamba_spec("align"), 8), + (_mamba_spec("none"), 9), + (_mamba_spec("all"), 9), + ], +) +def test_reserved_block_count_matches_speculative_tail(spec, expected) -> None: + # 14 scheduled tokens plus DFlash's eight lookahead positions crosses the + # second attention block. Align-mode Mamba excludes lookahead from its + # token range but always retains seven speculative state blocks. + assert ( + _reserved_block_count( + 14, + spec, + num_lookahead_tokens=NUM_SPEC_TOKENS + 1, + max_model_len=MAX_MODEL_LEN, + max_encoder_len=0, + ) + == expected + ) diff --git a/tests/v1/worker/test_gpu_worker_static_pp.py b/tests/v1/worker/test_gpu_worker_static_pp.py new file mode 100644 index 0000000000..b7e9e8da3a --- /dev/null +++ b/tests/v1/worker/test_gpu_worker_static_pp.py @@ -0,0 +1,153 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +import vllm.envs as envs +from vllm.config import CUDAGraphMode +from vllm.sequence import IntermediateTensors +from vllm.v1.worker import gpu_worker +from vllm.v1.worker.gpu_model_runner import GPUModelRunner +from vllm.v1.worker.gpu_worker import Worker + + +def _worker() -> SimpleNamespace: + model_config = SimpleNamespace( + dtype=torch.float16, + ) + return SimpleNamespace( + model_runner=SimpleNamespace( + model=SimpleNamespace( + make_empty_intermediate_tensors=lambda batch_size, dtype, device: ( + IntermediateTensors( + { + "hidden_states": torch.empty( + (batch_size, 4, 4096), dtype=dtype, device=device + ) + } + ) + ) + ) + ), + vllm_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, + compilation_config=SimpleNamespace( + cudagraph_mode=CUDAGraphMode.FULL, + pass_config=SimpleNamespace(enable_sp=False), + ), + model_config=model_config, + ), + ) + + +def test_static_pp_admission_uses_engine_and_tensor_geometry(monkeypatch) -> None: + monkeypatch.setenv("VLLM_SM70_PP_STATIC_HIDDEN_TRANSFER", "1") + envs.disable_envs_cache() + worker = _worker() + try: + with ( + patch.object(gpu_worker.current_platform, "is_cuda", return_value=True), + patch.object( + gpu_worker.current_platform, + "is_device_capability", + return_value=True, + ), + ): + assert Worker._use_sm70_static_pp_hidden_transfer(worker, 1) + assert not Worker._use_sm70_static_pp_hidden_transfer(worker, 2) + + wrong_schema = _worker() + wrong_schema.model_runner.model.make_empty_intermediate_tensors = ( + lambda batch_size, dtype, device: IntermediateTensors( + { + "hidden_states": torch.empty( + (batch_size, 4096), dtype=dtype, device=device + ) + } + ) + ) + assert not Worker._use_sm70_static_pp_hidden_transfer(wrong_schema, 1) + finally: + envs.disable_envs_cache() + + +def test_static_pp_admission_rejects_unsafe_concurrency(monkeypatch) -> None: + monkeypatch.setenv("VLLM_SM70_PP_STATIC_HIDDEN_TRANSFER", "1") + envs.disable_envs_cache() + try: + for field, value in ( + ("max_num_seqs", 2), + ("pipeline_parallel_size", 1), + ("tensor_parallel_size", 8), + ("enable_dbo", True), + ("ubatch_size", 2), + ): + worker = _worker() + target = ( + worker.vllm_config.scheduler_config + if field == "max_num_seqs" + else worker.vllm_config.parallel_config + ) + setattr(target, field, value) + with ( + patch.object(gpu_worker.current_platform, "is_cuda", return_value=True), + patch.object( + gpu_worker.current_platform, + "is_device_capability", + return_value=True, + ), + ): + assert not Worker._use_sm70_static_pp_hidden_transfer(worker, 1) + finally: + envs.disable_envs_cache() + + +def test_static_pp_schema_is_exact() -> None: + hidden_states = SimpleNamespace( + shape=(1, 4, 4096), + dtype=torch.float16, + is_cuda=True, + is_contiguous=lambda: True, + ) + assert Worker._is_static_pp_hidden_tensor_dict( + {"hidden_states": hidden_states}, + 1, # type: ignore[dict-item] + ) + hidden_states.shape = (1, 4096) + assert not Worker._is_static_pp_hidden_tensor_dict( + {"hidden_states": hidden_states}, + 1, # type: ignore[dict-item] + ) + + +def test_static_pp_receive_buffer_skips_self_copy() -> None: + hidden_states = torch.ones((4, 8), dtype=torch.float16) + runner = SimpleNamespace( + intermediate_tensors=IntermediateTensors({"hidden_states": hidden_states}), + vllm_config=SimpleNamespace( + parallel_config=SimpleNamespace(tensor_parallel_size=4) + ), + ) + source = IntermediateTensors({"hidden_states": hidden_states}) + version = hidden_states._version + + with patch( + "vllm.v1.worker.gpu_model_runner.is_residual_scattered_for_sp", + return_value=False, + ): + result = GPUModelRunner.sync_and_gather_intermediate_tensors( + runner, 1, source, True + ) + + assert hidden_states._version == version + assert result["hidden_states"].data_ptr() == hidden_states.data_ptr() diff --git a/tests/v1/worker/test_kv_block_zeroer.py b/tests/v1/worker/test_kv_block_zeroer.py new file mode 100644 index 0000000000..332111b7f4 --- /dev/null +++ b/tests/v1/worker/test_kv_block_zeroer.py @@ -0,0 +1,161 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""KVBlockZeroer must zero the right blocks under a block-major interleaved KV +layout (n_segs segments one cell apart, so the block stride is +n_segs * page_size_el). Build such a layout directly and check that zeroing +block 1 clears only block 1's cells; the old single-page_size_el stride zeros +the wrong cells, so this fails before the fix and passes after. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.v1.kv_cache_interface import MambaSpec +from vllm.v1.worker.utils import KVBlockZeroer, _infer_segment_block_strides + +CELL_ELS = 8 # int32 elements zeroed per (block, segment) +CELL_BYTES = CELL_ELS * 4 +N_LAYERS = 3 # -> n_segs = 3 interleaved segments +N_BLOCKS = 4 + + +@pytest.mark.parametrize( + ("seg_addrs", "page_size_el", "expected"), + [ + ([0x1000], 8, [8]), + ([0x1000, 0x2000], 8, [8, 8]), + ([0x1000, 0x1020, 0x1040], 8, [24, 24, 24]), + # Two independent two-segment interleaved pools. Input order differs + # from address order to prove strides stay bound to their segments. + ([0x3020, 0x1000, 0x3000, 0x1020], 8, [16, 16, 16, 16]), + ], +) +def test_infer_segment_block_strides( + seg_addrs: list[int], + page_size_el: int, + expected: list[int], +) -> None: + assert _infer_segment_block_strides(seg_addrs, page_size_el) == expected + + +def test_infer_segment_block_strides_rejects_duplicate_addresses() -> None: + with pytest.raises(ValueError, match="must be unique"): + _infer_segment_block_strides([0x1000, 0x1000], 8) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_zero_block_ids_block_major_interleaved() -> None: + device = torch.device("cuda") + + # Block-major interleaved buffer: cell k holds (block B, layer i) where + # k == B * N_LAYERS + i. Fill with a sentinel so any write is visible. + flat = torch.ones(N_BLOCKS * N_LAYERS * CELL_ELS, dtype=torch.int32, device=device) + + # One mamba spec whose page is exactly one cell; page_size_padded lets us + # pin page_size_bytes without depending on real state shapes. + spec = MambaSpec( + block_size=16, + shapes=((1,),), + dtypes=(torch.int32,), + page_size_padded=CELL_BYTES, + ) + layer_names = [f"mamba.{i}" for i in range(N_LAYERS)] + group = SimpleNamespace( + kv_cache_spec=spec, + kv_cache_group_id=0, + layer_names=layer_names, + backend=None, # unused on the mamba path + ) + # Layer i's state view starts one cell into the buffer, so the segment base + # addresses are exactly one cell apart -> the interleaved layout. + static_forward_context = { + name: SimpleNamespace(kv_cache=[flat[i * CELL_ELS :]]) + for i, name in enumerate(layer_names) + } + + zeroer = KVBlockZeroer(device, pin_memory=False) + zeroer.init_meta( + [group], + kernel_block_sizes=[spec.block_size], + cache_dtype="auto", + runner_only_attn_layers=set(), + static_forward_context=static_forward_context, + ) + + # Zero block 1 (block 0 is correct under both old and new strides, so it + # does not discriminate; block 1 does). + target = 1 + zeroer.zero_block_ids([target]) + torch.accelerator.synchronize() + + cells = flat.view(N_BLOCKS * N_LAYERS, CELL_ELS) + zeroed = {int(k) for k in range(cells.shape[0]) if bool((cells[k] == 0).all())} + expected = {target * N_LAYERS + i for i in range(N_LAYERS)} + assert zeroed == expected, ( + f"expected only block {target}'s cells {sorted(expected)} to be " + f"zeroed, got {sorted(zeroed)}" + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_zero_block_ids_multiple_interleaved_pools() -> None: + device = torch.device("cuda") + layers_per_pool = 2 + buffers = [ + torch.ones( + N_BLOCKS * layers_per_pool * CELL_ELS, + dtype=torch.int32, + device=device, + ) + for _ in range(2) + ] + spec = MambaSpec( + block_size=16, + shapes=((1,),), + dtypes=(torch.int32,), + page_size_padded=CELL_BYTES, + ) + groups = [] + static_forward_context = {} + for pool_index, buffer in enumerate(buffers): + layer_names = [ + f"mamba.{pool_index}.{layer_index}" + for layer_index in range(layers_per_pool) + ] + groups.append( + SimpleNamespace( + kv_cache_spec=spec, + kv_cache_group_id=pool_index, + layer_names=layer_names, + backend=None, + ) + ) + static_forward_context.update( + { + name: SimpleNamespace(kv_cache=[buffer[index * CELL_ELS :]]) + for index, name in enumerate(layer_names) + } + ) + + zeroer = KVBlockZeroer(device, pin_memory=False) + zeroer.init_meta( + groups, + kernel_block_sizes=[spec.block_size] * len(groups), + cache_dtype="auto", + runner_only_attn_layers=set(), + static_forward_context=static_forward_context, + ) + + target = 1 + zeroer.zero_block_ids([target]) + torch.accelerator.synchronize() + + expected = {target * layers_per_pool + index for index in range(layers_per_pool)} + for buffer in buffers: + cells = buffer.view(N_BLOCKS * layers_per_pool, CELL_ELS) + zeroed = { + index for index in range(cells.shape[0]) if bool((cells[index] == 0).all()) + } + assert zeroed == expected diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index 0a64067c1a..fc2ecea746 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -543,6 +543,32 @@ def test_warmup_batch_memcpy_kernel(): assert worker_mamba_utils.warmup_batch_memcpy_kernel(torch.device("cuda:0")) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_warmup_fused_postprocess_does_not_modify_state(): + device = torch.device("cuda:0") + cfg = _TestConfig() + layer_names = ["layer_0"] + kv_cache_config = _make_kv_cache_config(cfg, layer_names) + gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device) + _, _, conv_state, temporal_state, _, forward_context = _make_dual_layer_state( + cfg, device + ) + block_table = torch.arange( + cfg.max_num_reqs * cfg.num_blocks, + dtype=torch.int32, + device=device, + ).reshape(cfg.max_num_reqs, cfg.num_blocks) + gpu_ctx.initialize_from_forward_context( + kv_cache_config, forward_context, _COPY_FUNCS, [block_table] + ) + conv_before = conv_state.clone() + temporal_before = temporal_state.clone() + + assert gpu_ctx.warmup_fused_postprocess() + assert torch.equal(conv_state, conv_before) + assert torch.equal(temporal_state, temporal_before) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") class TestPostprocessMambaFusedKernel: """Tests for postprocess_mamba_fused_kernel comparing GPU vs CPU paths.""" @@ -689,9 +715,7 @@ def test_mtp4_boundary_rollover_counts_1_to_5(self, device, test_config): torch.manual_seed(20260617) accepted = [1, 2, 3, 4, 5] - req_ids = [f"same_{i}" for i in accepted] + [ - f"cross_{i}" for i in accepted - ] + req_ids = [f"same_{i}" for i in accepted] + [f"cross_{i}" for i in accepted] num_accepted_tokens = accepted + accepted # scheduled=5 and draft=4 gives running_state = computed + 1. num_scheduled_tokens = {req_id: 5 for req_id in req_ids} @@ -742,9 +766,7 @@ def test_mtp4_boundary_rollover_counts_1_to_5(self, device, test_config): torch.accelerator.synchronize() gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device) - block_table_gpu = torch.zeros( - len(req_ids), 8, dtype=torch.int32, device=device - ) + block_table_gpu = torch.zeros(len(req_ids), 8, dtype=torch.int32, device=device) for i, block_ids in enumerate(block_ids_per_req): block_table_gpu[i, : len(block_ids)] = torch.tensor( block_ids, dtype=torch.int32, device=device 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 ec62b2101e..8d8d0b188b 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -3032,6 +3032,32 @@ def sm70_tp4_all_reduce_gemma_rms_norm( ) +def sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather( + fa: int, + inp: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + normalized_out: torch.Tensor, + residual_out: torch.Tensor, + reg_input_buffer: int, + reg_output_buffer: int, + reg_buffer_sz_bytes: int, + epsilon: float, +) -> None: + torch.ops._C_custom_ar.sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather( + fa, + inp, + residual, + weight, + normalized_out, + residual_out, + reg_input_buffer, + reg_output_buffer, + reg_buffer_sz_bytes, + epsilon, + ) + + def all_reduce_sum2( fa: int, inp_a: torch.Tensor, @@ -3123,10 +3149,18 @@ def meta_size() -> int: return torch.ops._C_custom_ar.meta_size() +def sm70_tp4_push_allreduce_buffer_size() -> int: + return torch.ops._C_custom_ar.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) +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) + + 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) diff --git a/vllm/_sm70_ops.py b/vllm/_sm70_ops.py index 0bcd044544..5870e325e4 100644 --- a/vllm/_sm70_ops.py +++ b/vllm/_sm70_ops.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os +from pathlib import Path from typing import TYPE_CHECKING import torch @@ -9,6 +11,44 @@ current_platform.import_kernels() + +def _maybe_load_fp8_qpn8_library() -> None: + """Load an explicitly selected source-built QPN8 extension. + + Production builds register these operators in ``vllm._C``. This opt-in + path lets source experiments add only the QPN8 operators to an otherwise + compatible installed build, including in spawned TP workers. + """ + library_path = os.getenv("VLLM_SM70_FP8_QPN8_LIBRARY") + if library_path is None: + return + generic_override = os.getenv("VLLM_SM70_FP8_QPN8") + specific_override = os.getenv("VLLM_SM70_FP8_QPN8_PP2_TP4") + generic_enabled = generic_override == "1" + default_route_enabled = generic_override != "0" and specific_override != "0" + if generic_enabled or default_route_enabled: + torch.ops.load_library(library_path) + + +_maybe_load_fp8_qpn8_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"): + return + + library_path = os.getenv("VLLM_SM70_SAMPLER_LIBRARY") + if library_path is None: + bundled = sorted(Path(__file__).resolve().parent.glob("_sm70_sampler_C*.so")) + if bundled: + library_path = str(bundled[-1]) + if library_path is not None: + torch.ops.load_library(library_path) + + +_maybe_load_sm70_sampler_library() + if TYPE_CHECKING: def register_fake(fn): @@ -421,6 +461,647 @@ def _fp8_gemm_sm70_out_fake( return None +def fp8_gemm_sm70_prefill_prescaled_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_prefill_prescaled_out")( + out, input, qweight, prescaled_factors, group_size, k_ld, q_ld + ) + + +if hasattr(torch.ops._C, "fp8_gemm_sm70_prefill_prescaled_out"): + + @register_fake("_C::fp8_gemm_sm70_prefill_prescaled_out") + def _fp8_gemm_sm70_prefill_prescaled_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_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) + + +if hasattr(torch.ops._C, "fp8_qpn8_prepare_sm70"): + + @register_fake("_C::fp8_qpn8_prepare_sm70") + def _fp8_qpn8_prepare_sm70_fake( + qweight: torch.Tensor, + scales: torch.Tensor, + ) -> list[torch.Tensor]: + n = qweight.size(0) + k = qweight.size(1) + codes = torch.empty((k, n), dtype=torch.uint8, device=qweight.device) + scale_shape = (1, n) if scales.shape == (n, 1) else (k // 128, n // 32) + group_scales = torch.empty( + scale_shape, dtype=torch.float16, device=scales.device + ) + return [codes, group_scales] + + +def fp8_qpn8_dequantize_sm70_out( + out: torch.Tensor, + codes: torch.Tensor, + 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) + + +if hasattr(torch.ops._C, "fp8_qpn8_dequantize_sm70_out"): + + @register_fake("_C::fp8_qpn8_dequantize_sm70_out") + def _fp8_qpn8_dequantize_sm70_out_fake( + out: torch.Tensor, + codes: torch.Tensor, + group_scales: torch.Tensor, + ) -> None: + return None + + +def fp8_qpn8_prefill_sm70_out( + out: torch.Tensor, + dense_weight_ptr: int, + input: torch.Tensor, + codes: torch.Tensor, + group_scales: torch.Tensor, + gated_silu: bool, +) -> None: + """Dequantize QPN8 into bounded workspace and run a large-M FP16 GEMM.""" + _op("fp8_qpn8_prefill_sm70_out")( + out, + dense_weight_ptr, + input, + codes, + group_scales, + gated_silu, + ) + + +if hasattr(torch.ops._C, "fp8_qpn8_prefill_sm70_out"): + + @register_fake("_C::fp8_qpn8_prefill_sm70_out") + def _fp8_qpn8_prefill_sm70_out_fake( + out: torch.Tensor, + dense_weight_ptr: int, + input: torch.Tensor, + codes: torch.Tensor, + group_scales: torch.Tensor, + gated_silu: bool, + ) -> None: + return None + + +def fp8_qpn8_dispatch_sm70_out( + out: torch.Tensor, + dense_weight_ptr: int, + input: torch.Tensor, + codes: torch.Tensor, + group_scales: torch.Tensor, + split_k: int, + accumulator_chains: int, + prefetch_codes: bool, + gated_silu: bool, +) -> None: + """Runtime-dispatch dynamic M without specializing a Python branch.""" + _op("fp8_qpn8_dispatch_sm70_out")( + out, + dense_weight_ptr, + input, + codes, + group_scales, + split_k, + accumulator_chains, + prefetch_codes, + gated_silu, + ) + + +if hasattr(torch.ops._C, "fp8_qpn8_dispatch_sm70_out"): + + @register_fake("_C::fp8_qpn8_dispatch_sm70_out") + def _fp8_qpn8_dispatch_sm70_out_fake( + out: torch.Tensor, + dense_weight_ptr: int, + input: torch.Tensor, + codes: torch.Tensor, + group_scales: torch.Tensor, + split_k: int, + accumulator_chains: int, + prefetch_codes: bool, + gated_silu: bool, + ) -> None: + return None + + +def fp8_qpn8_gemm_sm70_out( + out: torch.Tensor, + input: torch.Tensor, + codes: torch.Tensor, + group_scales: torch.Tensor, + split_k: int, + accumulator_chains: int, + fast_decoder: bool, + prefetch_codes: bool = False, +) -> None: + """Run the SM70 QPN8 FP8 GEMM into ``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")( + out, + input, + codes, + group_scales, + split_k, + accumulator_chains, + fast_decoder, + prefetch_codes, + ) + + +if hasattr(torch.ops._C, "fp8_qpn8_gemm_sm70_out"): + + @register_fake("_C::fp8_qpn8_gemm_sm70_out") + def _fp8_qpn8_gemm_sm70_out_fake( + out: torch.Tensor, + input: torch.Tensor, + codes: torch.Tensor, + group_scales: torch.Tensor, + split_k: int, + accumulator_chains: int, + fast_decoder: bool, + prefetch_codes: bool, + ) -> None: + return None + + +def fp8_qpn8_gemm_ba_split_sm70_out( + qkv_out: torch.Tensor, + z_out: torch.Tensor, + b_out: torch.Tensor, + a_out: torch.Tensor, + input: torch.Tensor, + codes: torch.Tensor, + group_scales: torch.Tensor, + 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")( + qkv_out, + z_out, + b_out, + a_out, + input, + codes, + group_scales, + ba_weight, + ) + + +if hasattr(torch.ops._C, "fp8_qpn8_gemm_ba_split_sm70_out"): + + @register_fake("_C::fp8_qpn8_gemm_ba_split_sm70_out") + def _fp8_qpn8_gemm_ba_split_sm70_out_fake( + qkv_out: torch.Tensor, + z_out: torch.Tensor, + b_out: torch.Tensor, + a_out: torch.Tensor, + input: torch.Tensor, + codes: torch.Tensor, + group_scales: torch.Tensor, + ba_weight: torch.Tensor, + ) -> None: + return None + + +def fp8_qpn8_dispatch_ba_split_sm70_out( + qkv_out: torch.Tensor, + z_out: torch.Tensor, + b_out: torch.Tensor, + a_out: torch.Tensor, + qkvz_staging: torch.Tensor, + ba_staging: torch.Tensor, + dense_weight_ptr: int, + input: torch.Tensor, + codes: torch.Tensor, + group_scales: torch.Tensor, + 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")( + qkv_out, + z_out, + b_out, + a_out, + qkvz_staging, + ba_staging, + dense_weight_ptr, + input, + codes, + group_scales, + ba_weight, + ) + + +if hasattr(torch.ops._C, "fp8_qpn8_dispatch_ba_split_sm70_out"): + + @register_fake("_C::fp8_qpn8_dispatch_ba_split_sm70_out") + def _fp8_qpn8_dispatch_ba_split_sm70_out_fake( + qkv_out: torch.Tensor, + z_out: torch.Tensor, + b_out: torch.Tensor, + a_out: torch.Tensor, + qkvz_staging: torch.Tensor, + ba_staging: torch.Tensor, + dense_weight_ptr: int, + input: torch.Tensor, + codes: torch.Tensor, + group_scales: torch.Tensor, + ba_weight: torch.Tensor, + ) -> None: + return None + + +def fp8_qpn8_gated_pair_sm70_out( + out: torch.Tensor, + input: torch.Tensor, + codes: torch.Tensor, + group_scales: torch.Tensor, + split_k: int, + accumulator_chains: int, + fast_decoder: bool, + prefetch_codes: bool = False, +) -> None: + """Run the single-kernel paired-tile QPN8 gated SiLU experiment.""" + _op("fp8_qpn8_gated_pair_sm70_out")( + out, + input, + codes, + group_scales, + split_k, + accumulator_chains, + fast_decoder, + prefetch_codes, + ) + + +if hasattr(torch.ops._C, "fp8_qpn8_gated_pair_sm70_out"): + + @register_fake("_C::fp8_qpn8_gated_pair_sm70_out") + def _fp8_qpn8_gated_pair_sm70_out_fake( + out: torch.Tensor, + input: torch.Tensor, + codes: torch.Tensor, + group_scales: torch.Tensor, + split_k: int, + accumulator_chains: int, + fast_decoder: bool, + prefetch_codes: bool, + ) -> None: + return None + + +def nvfp4_qpn4_prepare_sm70( + qweight: torch.Tensor, + scales: torch.Tensor, +) -> list[torch.Tensor]: + """Pack unpacked FP4 weights and FP16 group scales into QPN4 layout.""" + return _op("nvfp4_qpn4_prepare_sm70")(qweight, scales) + + +if hasattr(torch.ops._C, "nvfp4_qpn4_prepare_sm70"): + + @register_fake("_C::nvfp4_qpn4_prepare_sm70") + def _nvfp4_qpn4_prepare_sm70_fake( + qweight: torch.Tensor, + scales: torch.Tensor, + ) -> list[torch.Tensor]: + k, n = qweight.shape + codes = torch.empty((k, n // 2), dtype=torch.uint8, device=qweight.device) + packed_scales = torch.empty_like(scales) + return [codes, packed_scales] + + +def nvfp4_qpn4_prepare_scale_code_sm70( + qweight: torch.Tensor, + scale_codes: torch.Tensor, +) -> list[torch.Tensor]: + """Pack unpacked FP4 weights and checkpoint E4M3 group-scale codes.""" + return _op("nvfp4_qpn4_prepare_scale_code_sm70")(qweight, scale_codes) + + +if hasattr(torch.ops._C, "nvfp4_qpn4_prepare_scale_code_sm70"): + + @register_fake("_C::nvfp4_qpn4_prepare_scale_code_sm70") + def _nvfp4_qpn4_prepare_scale_code_sm70_fake( + qweight: torch.Tensor, + scale_codes: torch.Tensor, + ) -> list[torch.Tensor]: + k, n = qweight.shape + codes = torch.empty((k, n // 2), dtype=torch.uint8, device=qweight.device) + packed_scale_codes = torch.empty( + (k // 16, n), + dtype=torch.uint8, + device=scale_codes.device, + ) + return [codes, packed_scale_codes] + + +def nvfp4_qpn4_dequantize_sm70_out( + out: torch.Tensor, + codes: torch.Tensor, + scales: torch.Tensor, + global_scale: float, + use_scale_code: bool, +) -> None: + _op("nvfp4_qpn4_dequantize_sm70_out")( + out, codes, scales, global_scale, use_scale_code + ) + + +if hasattr(torch.ops._C, "nvfp4_qpn4_dequantize_sm70_out"): + + @register_fake("_C::nvfp4_qpn4_dequantize_sm70_out") + def _nvfp4_qpn4_dequantize_sm70_out_fake( + out: torch.Tensor, + codes: torch.Tensor, + scales: torch.Tensor, + global_scale: float, + use_scale_code: bool, + ) -> None: + return None + + +def nvfp4_qpn4_prefill_sm70_out( + out: torch.Tensor, + dense_weight_ptr: int, + input: torch.Tensor, + codes: torch.Tensor, + scales: torch.Tensor, + global_scale: float, + use_scale_code: bool, + gated_silu: bool, +) -> None: + _op("nvfp4_qpn4_prefill_sm70_out")( + out, + dense_weight_ptr, + input, + codes, + scales, + global_scale, + use_scale_code, + gated_silu, + ) + + +if hasattr(torch.ops._C, "nvfp4_qpn4_prefill_sm70_out"): + + @register_fake("_C::nvfp4_qpn4_prefill_sm70_out") + def _nvfp4_qpn4_prefill_sm70_out_fake( + out: torch.Tensor, + dense_weight_ptr: int, + input: torch.Tensor, + codes: torch.Tensor, + scales: torch.Tensor, + global_scale: float, + use_scale_code: bool, + gated_silu: bool, + ) -> None: + return None + + +def nvfp4_qpn4_dispatch_sm70_out( + out: torch.Tensor, + dense_weight_ptr: int, + input: torch.Tensor, + codes: torch.Tensor, + scales: torch.Tensor, + global_scale: float, + use_scale_code: bool, + gated_silu: bool, +) -> None: + """Dispatch exact M=1 decode or bounded-workspace large-M prefill.""" + _op("nvfp4_qpn4_dispatch_sm70_out")( + out, + dense_weight_ptr, + input, + codes, + scales, + global_scale, + use_scale_code, + gated_silu, + ) + + +if hasattr(torch.ops._C, "nvfp4_qpn4_dispatch_sm70_out"): + + @register_fake("_C::nvfp4_qpn4_dispatch_sm70_out") + def _nvfp4_qpn4_dispatch_sm70_out_fake( + out: torch.Tensor, + dense_weight_ptr: int, + input: torch.Tensor, + codes: torch.Tensor, + scales: torch.Tensor, + global_scale: float, + use_scale_code: bool, + gated_silu: bool, + ) -> None: + return None + + +def nvfp4_qpn2_prepare_sm70( + weight_packed: torch.Tensor, + weight_scale: torch.Tensor, +) -> list[torch.Tensor]: + """Pack checkpoint-native NVFP4 tensors into the QPN2 layout.""" + return _op("nvfp4_qpn2_prepare_sm70")(weight_packed, weight_scale) + + +if hasattr(torch.ops._C, "nvfp4_qpn2_prepare_sm70"): + + @register_fake("_C::nvfp4_qpn2_prepare_sm70") + def _nvfp4_qpn2_prepare_sm70_fake( + weight_packed: torch.Tensor, + weight_scale: torch.Tensor, + ) -> list[torch.Tensor]: + codes = torch.empty_like(weight_packed) + scales = torch.empty( + weight_scale.shape, dtype=torch.uint8, device=weight_scale.device + ) + return [codes, scales] + + +def nvfp4_qpn2_gemm_sm70_out( + out: torch.Tensor, + input: torch.Tensor, + codes: torch.Tensor, + scales: torch.Tensor, + global_scale: float, + split_k: int, + accumulator_chains: int, +) -> None: + _op("nvfp4_qpn2_gemm_sm70_out")( + out, + input, + codes, + scales, + global_scale, + split_k, + accumulator_chains, + ) + + +if hasattr(torch.ops._C, "nvfp4_qpn2_gemm_sm70_out"): + + @register_fake("_C::nvfp4_qpn2_gemm_sm70_out") + def _nvfp4_qpn2_gemm_sm70_out_fake( + out: torch.Tensor, + input: torch.Tensor, + codes: torch.Tensor, + scales: torch.Tensor, + global_scale: float, + split_k: int, + accumulator_chains: int, + ) -> None: + return None + + +def nvfp4_qpn2_gated_sm70_out( + out: torch.Tensor, + input: torch.Tensor, + codes: torch.Tensor, + scales: torch.Tensor, + global_scale: float, + split_k: int, + accumulator_chains: int, +) -> None: + _op("nvfp4_qpn2_gated_sm70_out")( + out, + input, + codes, + scales, + global_scale, + split_k, + accumulator_chains, + ) + + +if hasattr(torch.ops._C, "nvfp4_qpn2_gated_sm70_out"): + + @register_fake("_C::nvfp4_qpn2_gated_sm70_out") + def _nvfp4_qpn2_gated_sm70_out_fake( + out: torch.Tensor, + input: torch.Tensor, + codes: torch.Tensor, + scales: torch.Tensor, + global_scale: float, + split_k: int, + accumulator_chains: int, + ) -> None: + return None + + +def nvfp4_qpn2_dispatch_sm70_out( + out: torch.Tensor, + input: torch.Tensor, + codes: torch.Tensor, + scales: torch.Tensor, + global_scale: float, + split_k: int, + accumulator_chains: int, + tm_weight: torch.Tensor, + tm_scales: torch.Tensor, + tm_group_size: int, + tm_k_ld: int, + tm_q_ld: int, + gated_silu: bool, +) -> None: + """Select QPN2 for M<=8 and TurboMind for larger dynamic M.""" + _op("nvfp4_qpn2_dispatch_sm70_out")( + out, + input, + codes, + scales, + global_scale, + split_k, + accumulator_chains, + tm_weight, + tm_scales, + tm_group_size, + tm_k_ld, + tm_q_ld, + gated_silu, + ) + + +if hasattr(torch.ops._C, "nvfp4_qpn2_dispatch_sm70_out"): + + @register_fake("_C::nvfp4_qpn2_dispatch_sm70_out") + def _nvfp4_qpn2_dispatch_sm70_out_fake( + out: torch.Tensor, + input: torch.Tensor, + codes: torch.Tensor, + scales: torch.Tensor, + global_scale: float, + split_k: int, + accumulator_chains: int, + tm_weight: torch.Tensor, + tm_scales: torch.Tensor, + tm_group_size: int, + tm_k_ld: int, + tm_q_ld: int, + gated_silu: bool, + ) -> None: + return None + + def fp8_gemm_sm70_prefill_dispatch_out( out: torch.Tensor, dense_weight_ptr: int, @@ -541,6 +1222,50 @@ def _mxfp4_moe_dense_stage_sm70_out_fake( return None +def nvfp4_moe_dense_stage_sm70_out( + out: torch.Tensor, + input: torch.Tensor, + expert_offsets: torch.Tensor, + dense_expert_ids: torch.Tensor, + ptrs_w: torch.Tensor, + ptrs_s: torch.Tensor, + num_experts: int, + k: int, + n: int, + group_size: int, +) -> None: + _op("nvfp4_moe_dense_stage_sm70_out")( + out, + input, + expert_offsets, + dense_expert_ids, + ptrs_w, + ptrs_s, + num_experts, + k, + n, + group_size, + ) + + +if hasattr(torch.ops._C, "nvfp4_moe_dense_stage_sm70_out"): + + @register_fake("_C::nvfp4_moe_dense_stage_sm70_out") + def _nvfp4_moe_dense_stage_sm70_out_fake( + out: torch.Tensor, + input: torch.Tensor, + expert_offsets: torch.Tensor, + dense_expert_ids: torch.Tensor, + ptrs_w: torch.Tensor, + ptrs_s: torch.Tensor, + num_experts: int, + k: int, + n: int, + group_size: int, + ) -> None: + return None + + def mxfp4_moe_single_token_prepare_w13_sm70_out( gate_up: torch.Tensor, compact_input: torch.Tensor, @@ -767,6 +1492,145 @@ def _sm70_f16_gemm_out_fake( return None +def sm70_f16_indexed_rerank_out( + out: torch.Tensor, + input: torch.Tensor, + weight: torch.Tensor, + candidate_ids: torch.Tensor, + selected_raw: torch.Tensor, + selected_packed: torch.Tensor, + expanded: torch.Tensor, + partials: torch.Tensor, + barriers: torch.Tensor, + cta_n: int, + split_k: int, +) -> None: + _op("sm70_f16_indexed_rerank_out")( + out, + input, + weight, + candidate_ids, + selected_raw, + selected_packed, + expanded, + partials, + barriers, + cta_n, + split_k, + ) + + +if hasattr(torch.ops._C, "sm70_f16_indexed_rerank_out"): + + @register_fake("_C::sm70_f16_indexed_rerank_out") + def _sm70_f16_indexed_rerank_out_fake( + out: torch.Tensor, + input: torch.Tensor, + weight: torch.Tensor, + candidate_ids: torch.Tensor, + selected_raw: torch.Tensor, + selected_packed: torch.Tensor, + expanded: torch.Tensor, + partials: torch.Tensor, + barriers: torch.Tensor, + cta_n: int, + split_k: int, + ) -> None: + return None + + +def sm70_f16_indexed_rerank_packed_out( + out: torch.Tensor, + input: torch.Tensor, + packed_weight: torch.Tensor, + candidate_ids: torch.Tensor, + selected_packed: torch.Tensor, + expanded: torch.Tensor, + partials: torch.Tensor, + barriers: torch.Tensor, + cta_n: int, + split_k: int, +) -> None: + _op("sm70_f16_indexed_rerank_packed_out")( + out, + input, + packed_weight, + candidate_ids, + selected_packed, + expanded, + partials, + barriers, + cta_n, + split_k, + ) + + +if hasattr(torch.ops._C, "sm70_f16_indexed_rerank_packed_out"): + + @register_fake("_C::sm70_f16_indexed_rerank_packed_out") + def _sm70_f16_indexed_rerank_packed_out_fake( + out: torch.Tensor, + input: torch.Tensor, + packed_weight: torch.Tensor, + candidate_ids: torch.Tensor, + selected_packed: torch.Tensor, + expanded: torch.Tensor, + partials: torch.Tensor, + barriers: torch.Tensor, + cta_n: int, + split_k: int, + ) -> None: + return None + + +def sm70_f16_rerank_keys_out( + keys: torch.Tensor, + logits: torch.Tensor, + candidate_ids: torch.Tensor, +) -> None: + _op("sm70_f16_rerank_keys_out")(keys, logits, candidate_ids) + + +if hasattr(torch.ops._C, "sm70_f16_rerank_keys_out"): + + @register_fake("_C::sm70_f16_rerank_keys_out") + def _sm70_f16_rerank_keys_out_fake( + keys: torch.Tensor, + logits: torch.Tensor, + candidate_ids: torch.Tensor, + ) -> None: + return None + + +def sm70_f16_rerank_topk_out( + values_out: torch.Tensor, + ids_out: torch.Tensor, + logits: torch.Tensor, + candidate_ids: torch.Tensor, + vocab_start_index: int, +) -> None: + _op("sm70_f16_rerank_topk_out")( + values_out, + ids_out, + logits, + candidate_ids, + vocab_start_index, + ) + + +if hasattr(torch.ops._C, "sm70_f16_rerank_topk_out"): + + @register_fake("_C::sm70_f16_rerank_topk_out") + def _sm70_f16_rerank_topk_out_fake( + values_out: torch.Tensor, + ids_out: torch.Tensor, + logits: torch.Tensor, + candidate_ids: torch.Tensor, + vocab_start_index: int, + ) -> None: + return None + + def sm70_f16_lm_head_top1_out( values_out: torch.Tensor, indices_out: torch.Tensor, @@ -939,6 +1803,114 @@ def _sm70_sample_packed_top20_out_fake( return None +def sm70_sample_sorted_top20_philox_out( + sampled_token_out: torch.Tensor, + sparse_ids_out: torch.Tensor, + sparse_probs_out: torch.Tensor, + top_values: torch.Tensor, + top_indices: torch.Tensor, + generator: torch.Generator | None, + vocab_size: int, + top_p: float, +) -> None: + _op("sm70_sample_sorted_top20_philox_out")( + sampled_token_out, + sparse_ids_out, + sparse_probs_out, + top_values, + top_indices, + generator, + vocab_size, + top_p, + ) + + +if hasattr(torch.ops._C, "sm70_sample_sorted_top20_philox_out"): + + @register_fake("_C::sm70_sample_sorted_top20_philox_out") + def _sm70_sample_sorted_top20_philox_out_fake( + sampled_token_out: torch.Tensor, + sparse_ids_out: torch.Tensor, + sparse_probs_out: torch.Tensor, + top_values: torch.Tensor, + top_indices: torch.Tensor, + generator: torch.Generator | None, + vocab_size: int, + top_p: float, + ) -> None: + return None + + +def sm70_sample_sorted_top20_philox_token_out( + sampled_token_out: torch.Tensor, + top_values: torch.Tensor, + top_indices: torch.Tensor, + generator: torch.Generator | None, + vocab_size: int, + top_p: float, +) -> None: + _op("sm70_sample_sorted_top20_philox_token_out")( + sampled_token_out, + top_values, + top_indices, + generator, + vocab_size, + top_p, + ) + + +if hasattr(torch.ops._C, "sm70_sample_sorted_top20_philox_token_out"): + + @register_fake("_C::sm70_sample_sorted_top20_philox_token_out") + def _sm70_sample_sorted_top20_philox_token_out_fake( + sampled_token_out: torch.Tensor, + top_values: torch.Tensor, + top_indices: torch.Tensor, + generator: torch.Generator | None, + vocab_size: int, + top_p: float, + ) -> None: + return None + + +def sm70_sample_chunked_top20_philox_token_out( + sampled_token_out: torch.Tensor, + global_values: torch.Tensor, + local_indices: torch.Tensor, + global_positions: torch.Tensor, + generator: torch.Generator | None, + vocab_size: int, + top_p: float, + chunk_size: int, +) -> None: + _op("sm70_sample_chunked_top20_philox_token_out")( + sampled_token_out, + global_values, + local_indices, + global_positions, + generator, + vocab_size, + top_p, + chunk_size, + ) + + +if hasattr(torch.ops._C, "sm70_sample_chunked_top20_philox_token_out"): + + @register_fake("_C::sm70_sample_chunked_top20_philox_token_out") + def _sm70_sample_chunked_top20_philox_token_out_fake( + sampled_token_out: torch.Tensor, + global_values: torch.Tensor, + local_indices: torch.Tensor, + global_positions: torch.Tensor, + generator: torch.Generator | None, + vocab_size: int, + top_p: float, + chunk_size: int, + ) -> None: + return None + + def sm70_dynamic_draft_vocab_update_tail_out( lru_token_ids: torch.Tensor, local_tail_token_ids: torch.Tensor, diff --git a/vllm/compilation/cuda_graph.py b/vllm/compilation/cuda_graph.py index 5664cb7361..5e80444996 100644 --- a/vllm/compilation/cuda_graph.py +++ b/vllm/compilation/cuda_graph.py @@ -192,8 +192,7 @@ def __init__( self.first_run_finished = False self.is_debugging_mode = envs.VLLM_LOGGING_LEVEL == "DEBUG" self.check_input_addresses = ( - self.is_debugging_mode - or envs.VLLM_CUDAGRAPH_INPUT_ADDR_DEBUG + self.is_debugging_mode or envs.VLLM_CUDAGRAPH_INPUT_ADDR_DEBUG ) self._runnable_str = str(runnable) if self.is_debugging_mode else None @@ -211,6 +210,7 @@ def __init__( # the entries for different batch descriptors that we need to capture # cudagraphs for. self.concrete_cudagraph_entries: dict[BatchDescriptor, CUDAGraphEntry] = {} + self.semantic_variant_graph_pools: dict[int, Any] = {} CUDAGraphWrapper._all_instances.add(self) @@ -264,6 +264,20 @@ def _collect_tensor_debug( def clear_graphs(self) -> None: self.concrete_cudagraph_entries.clear() + self.semantic_variant_graph_pools.clear() + + def _graph_pool_for_descriptor(self, batch_descriptor: BatchDescriptor) -> Any: + if batch_descriptor.graph_variant >= 0: + return self.graph_pool + # Negative variants are mutually exclusive semantic paths for an + # otherwise identical graph shape. They can replay in arbitrary order, + # so do not let a later capture reuse the default graph's intermediates. + graph_variant = batch_descriptor.graph_variant + if graph_variant not in self.semantic_variant_graph_pools: + self.semantic_variant_graph_pools[graph_variant] = ( + current_platform.graph_pool_handle() + ) + return self.semantic_variant_graph_pools[graph_variant] def __call__(self, *args: Any, **kwargs: Any) -> Any | None: if not is_forward_context_available(): @@ -335,8 +349,9 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any | None: ) ) - if self.graph_pool is not None: - set_graph_pool_id(self.graph_pool) + graph_pool = self._graph_pool_for_descriptor(entry.batch_descriptor) + if graph_pool is not None: + set_graph_pool_id(graph_pool) else: set_graph_pool_id(current_platform.graph_pool_handle()) @@ -347,7 +362,7 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any | None: # mind-exploding: carefully manage the reference and memory. with torch.cuda.graph( cudagraph, - pool=self.graph_pool, + pool=graph_pool, stream=current_stream(), ): # `output` is managed by pytorch's cudagraph pool diff --git a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py index aa0f68fa85..7726c72238 100644 --- a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py +++ b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py @@ -442,6 +442,89 @@ def _sm70_ar_gemma_rms_match(match: pm.Match) -> bool: return False +def _sm70_tp4_fusion_diagnostic(graph: fx.Graph) -> str: + """Summarize real post-grad users when the fail-closed pattern misses.""" + all_reduces: list[fx.Node] = [] + direct_users: dict[str, int] = {} + for node in graph.nodes: + if node.target != _ALL_REDUCE_OP: + continue + value = node.meta.get("val") + if value is None or value.dim() != 2 or value.shape[-1] != 5120: + continue + all_reduces.append(node) + for user in node.users: + target = str(user.target) + direct_users[target] = direct_users.get(target, 0) + 1 + + traces: list[str] = [] + for root in all_reduces: + queue: list[tuple[fx.Node, int]] = [(root, 0)] + seen = {root} + trace_nodes: list[tuple[fx.Node, int]] = [] + while queue and len(trace_nodes) < 32: + node, depth = queue.pop(0) + trace_nodes.append((node, depth)) + if depth == 10: + continue + for user in node.users: + if user not in seen: + seen.add(user) + queue.append((user, depth + 1)) + targets = " ".join(str(node.target) for node, _ in trace_nodes) + if "mean" not in targets and "rsqrt" not in targets and "rms" not in targets: + continue + lines = [] + for node, depth in trace_nodes: + value = node.meta.get("val") + shape = getattr(value, "shape", None) + dtype = getattr(value, "dtype", None) + lines.append( + f"{depth}:{node.name}:{node.target}:shape={shape}:dtype={dtype}:" + f"args={node.args}:kwargs={node.kwargs}" + ) + traces.append("\n".join(lines)) + if len(traces) == 4: + break + + return ( + f"all_reduce_f16_5120={len(all_reduces)}; " + f"direct_users={direct_users}; candidate_traces=\n" + "\n---\n".join(traces) + ) + + +def _sm70_tp4_expected_graph_matches(graph: fx.Graph) -> int: + """Count supported FP32-residual boundaries in one post-grad graph.""" + count = 0 + for node in graph.nodes: + if ( + node.target + != torch.ops.vllm.sm70_gemma_long_prefill_fused_add_rms_norm.default + ): + continue + reduced = node.args[0] + residual = node.args[1] + if not isinstance(reduced, fx.Node) or reduced.target != _ALL_REDUCE_OP: + continue + if not isinstance(residual, fx.Node): + continue + residual_value = residual.meta.get("val") + if residual_value is None or residual_value.dtype != torch.float32: + continue + input_node = reduced.args[0] + if not isinstance(input_node, fx.Node): + continue + value = input_node.meta.get("val") + if ( + value is not None + and value.dtype == torch.float16 + and value.dim() == 2 + and value.shape[-1] == 5120 + ): + count += 1 + return count + + class Sm70AllReduceGemmaRMSNormPattern(BasePattern): def __init__( self, @@ -513,6 +596,73 @@ def replacement( ) +class Sm70Tp4LongPrefillFusedNormPattern(BasePattern): + """Fuse TP4 AR plus the accepted opaque mixed-dtype Gemma norm.""" + + def __init__( + self, + epsilon: float, + dtype: torch.dtype, + device: str | None, + ) -> None: + super().__init__(dtype, device) + self.epsilon = epsilon + self.group_name = get_tp_group().unique_name + + def get_inputs(self) -> list[torch.Tensor]: + return [ + self.empty_f32(5, 5120), + self.empty(5, 5120), + self.empty(5120), + ] + + def register(self, pm_pass: PatternMatcherPass) -> None: + def pattern( + residual: torch.Tensor, + input: torch.Tensor, + weight: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + reduced = tensor_model_parallel_all_reduce(input) + return torch.ops.vllm.sm70_gemma_long_prefill_fused_add_rms_norm( + reduced, + residual, + weight, + self.epsilon, + ) + + def replacement( + residual: torch.Tensor, + input: torch.Tensor, + weight: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + return torch.ops.vllm.sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather( + input, + residual, + weight, + self.epsilon, + group_name=self.group_name, + ) + + pm.register_replacement( + pattern, + replacement, + self.get_inputs(), + pm.fwd_only, + pm_pass, + extra_check=_sm70_ar_gemma_rms_match, + ) + + first_return_only = lambda fn: lambda a, b, c: fn(a, b, c)[0] + pm.register_replacement( + first_return_only(pattern), # type: ignore[no-untyped-call] + first_return_only(replacement), # type: ignore[no-untyped-call] + self.get_inputs(), + pm.fwd_only, + pm_pass, + extra_check=_sm70_ar_gemma_rms_match, + ) + + class AllReduceFusedRMSNormStaticQuantFP8Pattern(BasePattern): """ This pattern replaces the allreduce + rms norm (without residual) @@ -874,6 +1024,7 @@ def __init__(self, config: VllmConfig) -> None: self.hidden_dim = config.model_config.get_hidden_size() self.group = get_tp_group().device_group rank = get_tensor_model_parallel_rank() + self.rank = rank sm70_common = ( self.hidden_dim == 5120 and self.model_dtype == torch.float16 @@ -882,7 +1033,20 @@ def __init__(self, config: VllmConfig) -> None: self.sm70_tp2_mode = ( envs.VLLM_SM70_TP2_AR_GEMMA_RMS_FUSION and self.tp_size == 2 and sm70_common ) - self.sm70_mode = self.sm70_tp2_mode + self.sm70_tp4_long_mode = ( + envs.VLLM_SM70_TP4_LONG_PREFILL_FUSED_NORM + and self.tp_size == 4 + and sm70_common + and config.parallel_config.pipeline_parallel_size == 1 + and config.speculative_config is None + ) + self.sm70_tp4_model_expected_patterns = ( + 2 * config.model_config.get_num_layers(config.parallel_config) - 1 + if self.sm70_tp4_long_mode + else 0 + ) + self.sm70_tp4_total_matches = 0 + self.sm70_mode = self.sm70_tp2_mode or self.sm70_tp4_long_mode if self.sm70_mode: device_comm = get_tp_group().device_communicator ca_comm = ( @@ -894,6 +1058,25 @@ def __init__(self, config: VllmConfig) -> None: self.tp_size, ) return + if self.sm70_tp4_long_mode and ca_comm.long_prefill_output_ptrs is None: + logger.warning_once( + "SM70 TP4 long-prefill collective-norm buffers are unavailable." + ) + return + if self.sm70_tp4_long_mode and not hasattr( + torch.ops.vllm, + "sm70_gemma_long_prefill_fused_add_rms_norm", + ): + logger.warning_once( + "SM70 long-prefill Gemma norm op is unavailable for TP4 fusion." + ) + return + if self.sm70_tp4_long_mode and not hasattr( + torch.ops.vllm, + "sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather", + ): + logger.warning_once("SM70 TP4 fused collective-norm op is unavailable.") + return self.max_token_num = config.scheduler_config.max_num_batched_tokens self.register_sm70_patterns() self.dump_patterns(config, self.patterns) @@ -968,6 +1151,14 @@ def __init__(self, config: VllmConfig) -> None: @enable_fake_mode def register_sm70_patterns(self) -> None: + if self.sm70_tp4_long_mode: + Sm70Tp4LongPrefillFusedNormPattern( + 1e-6, + self.model_dtype, + self.device, + ).register(self.patterns) + self.disabled = False + return for residual_dtype in (torch.float16, torch.float32): Sm70AllReduceGemmaRMSNormPattern( 1e-6, @@ -1029,6 +1220,12 @@ def is_applicable_for_range(self, compile_range: Range) -> bool: if self.disabled: logger.warning_once("AllReduce fusion pass is disabled.") return False + if getattr(self, "sm70_tp4_long_mode", False): + return bool( + compile_range.is_single_size() + and 8000 <= compile_range.start <= min(8192, self.max_token_num) + and compile_range.start % self.tp_size == 0 + ) return bool(compile_range.end <= self.max_token_num) @VllmInductorPass.time_and_log @@ -1037,7 +1234,36 @@ def __call__(self, graph: fx.Graph) -> None: logger.debug("AllReduceFusionPass disabled") return + expected_graph_matches = ( + _sm70_tp4_expected_graph_matches(graph) if self.sm70_tp4_long_mode else 0 + ) self.matched_count = self.patterns.apply(graph) + if self.sm70_tp4_long_mode and self.matched_count != expected_graph_matches: + if self.rank == 0: + logger.error( + "SM70 TP4 fusion post-grad diagnostic:\n%s", + _sm70_tp4_fusion_diagnostic(graph), + ) + raise RuntimeError( + "SM70 TP4 long-prefill fused collective-norm requires a " + "complete residual chain in every piecewise graph: expected " + f"{expected_graph_matches} matches, got " + f"{self.matched_count}." + ) + if self.sm70_tp4_long_mode: + self.sm70_tp4_total_matches += self.matched_count + if self.sm70_tp4_total_matches > self.sm70_tp4_model_expected_patterns: + raise RuntimeError( + "SM70 TP4 long-prefill fused collective-norm exceeded the " + "model residual-chain count: expected at most " + f"{self.sm70_tp4_model_expected_patterns}, got " + f"{self.sm70_tp4_total_matches}." + ) + if self.sm70_tp4_total_matches == self.sm70_tp4_model_expected_patterns: + logger.info_once( + "SM70 TP4 fused the complete %d-boundary model residual chain.", + self.sm70_tp4_total_matches, + ) logger.debug("Replaced %s patterns", self.matched_count) if getattr(self, "sm70_mode", False): logger.info( 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/fusion/sequence_parallelism.py b/vllm/compilation/passes/fusion/sequence_parallelism.py index 8d0f40e2c7..be8bff058b 100644 --- a/vllm/compilation/passes/fusion/sequence_parallelism.py +++ b/vllm/compilation/passes/fusion/sequence_parallelism.py @@ -256,6 +256,71 @@ def replacement( ) +class MiddleAllReduceSm70GemmaLongPrefillPattern(_SequenceParallelPatternHelper): + """Keep the accepted mixed-dtype SM70 norm while applying SP.""" + + def __init__( + self, + epsilon: float, + dtype: torch.dtype, + device: str | None, + weight_dtype: torch.dtype, + ) -> None: + super().__init__(epsilon, dtype, device) + self.weight_dtype = weight_dtype + + def get_inputs(self) -> list[torch.Tensor]: + residual = torch.empty([4, 4], device=self.device, dtype=torch.float32) + mm_1 = torch.empty([4, 4], device=self.device, dtype=self.dtype) + weight = torch.empty([4], device=self.device, dtype=self.weight_dtype) + return [residual, mm_1, weight] + + def register(self, pm_pass: PatternMatcherPass) -> None: + def pattern( + residual: torch.Tensor, + mm_1: torch.Tensor, + weight: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + all_reduce = self._all_reduce(mm_1) + return torch.ops.vllm.sm70_gemma_long_prefill_fused_add_rms_norm( + all_reduce, + residual, + weight, + self.epsilon, + ) + + def replacement( + residual: torch.Tensor, + mm_1: torch.Tensor, + weight: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + reduce_scatter = self._reduce_scatter(mm_1) + local_len = reduce_scatter.size(0) + residual = residual[ + self.tp_rank * local_len : self.tp_rank * local_len + local_len, ... + ] + normalized, residual_out = ( + torch.ops.vllm.sm70_gemma_long_prefill_fused_add_rms_norm( + reduce_scatter, + residual, + weight, + self.epsilon, + ) + ) + return self._all_gather(normalized), residual_out + + pm.register_replacement( + pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass + ) + pm.register_replacement( + get_first_out_wrapper(pattern), + get_first_out_wrapper(replacement), + self.get_inputs(), + pm.fwd_only, + pm_pass, + ) + + class FirstAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper): def __init__( self, @@ -583,6 +648,16 @@ def __init__(self, config: VllmConfig) -> None: MiddleAllReduceRMSNormPattern( epsilon, self.model_dtype, self.device ).register(self.patterns) + if epsilon == 1e-6 and hasattr( + torch.ops.vllm, + "sm70_gemma_long_prefill_fused_add_rms_norm", + ): + MiddleAllReduceSm70GemmaLongPrefillPattern( + epsilon, + self.model_dtype, + self.device, + self.model_dtype, + ).register(self.patterns) self.dump_patterns(config, self.patterns) 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/attention.py b/vllm/config/attention.py index 52ce9f102a..487cafe7fb 100644 --- a/vllm/config/attention.py +++ b/vllm/config/attention.py @@ -3,7 +3,7 @@ from typing import Any, Literal -from pydantic import field_validator +from pydantic import Field, field_validator from vllm.config.utils import config from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum @@ -53,6 +53,12 @@ class AttentionConfig: use_non_causal: bool = False """Whether to use non-causal (bidirectional) attention.""" + prefix_anchored_decode_window: int | None = Field(default=None, gt=0, le=2**31 - 1) + """Bound generated-token KV to this sliding window while keeping the + request prefix globally visible. This is an explicit, model-agnostic + engine option. It is currently supported only by the SM70 + ``FLASH_ATTN_V100`` fp16 causal path. ``None`` disables it.""" + flex_attn_block_m: int | None = None """Triton kernel BLOCK_M tile size for flex attention. Must be a power of 2 >= 16. If None and VLLM_BATCH_INVARIANT=1, diff --git a/vllm/config/offload.py b/vllm/config/offload.py index ad65e8acf3..75aa758536 100644 --- a/vllm/config/offload.py +++ b/vllm/config/offload.py @@ -40,6 +40,9 @@ class UVAOffloadConfig: - For parameter name "mlp.experts.w2_weight": - "experts" or "experts.w2_weight" will match. - "expert" or "w2" will NOT match (must be exact segments). + - A tower component name such as "visual" selects that complete + multimodal tower, including parameters whose names are relative to + the tower module. This allows distinguishing parameters like "w2_weight" and "w2_weight_scale". """ diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 27b1de7f73..c74aeb882e 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -2,12 +2,14 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import copy +import math from typing import TYPE_CHECKING, Any, Literal, get_args from pydantic import Field, SkipValidation, field_validator, model_validator from typing_extensions import Self from vllm.config import LoadConfig +from vllm.config.cache import CacheDType from vllm.config.kernel import MoEBackend from vllm.config.model import ModelConfig from vllm.config.parallel import ParallelConfig @@ -52,7 +54,8 @@ "gemma4_mtp", ] NgramGPUTypes = Literal["ngram_gpu"] -DFlashModelTypes = Literal["dflash", "dflash_ddtree"] +DFlashModelTypes = Literal["dflash"] +DFlashDDTreeModelTypes = Literal["dflash_ddtree"] DSparkModelTypes = Literal["dspark"] DDTreeBuildMode = Literal["best_first", "root_leaf", "spine_leaf"] EagleModelTypes = Literal[ @@ -61,6 +64,7 @@ "extract_hidden_states", MTPModelTypes, DFlashModelTypes, + DFlashDDTreeModelTypes, DSparkModelTypes, ] SpeculativeMethod = Literal[ @@ -119,6 +123,10 @@ class SpeculativeConfig: """Attention backend to use for the draft model. When `None`, the backend is automatically selected. Useful when the drafter requires a different attention backend (e.g. DFlash needs a non-causal-capable backend like FLASH_ATTN).""" + kv_cache_dtype: CacheDType | None = None + """KV cache dtype for the draft model. When ``None``, the draft inherits + the target cache dtype, except for SM70 DFlash with a quantized target, + which is resolved to ``auto`` during draft model loading.""" max_model_len: int | None = Field(default=None, ge=1) """The maximum model length of the draft model. Used when testing the ability to skip speculation for some sequences.""" @@ -150,6 +158,11 @@ class SpeculativeConfig: prompt_lookup_min: int | None = Field(default=None, ge=1) """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.""" # Alternative drafting strategies parallel_drafting: bool = False @@ -286,6 +299,20 @@ def _resolve_synthetic_acceptance_rates( during rejection sampling. This comes at the cost of additional GPU memory usage.""" + dspark_confidence_threshold: float = Field(default=0.0, ge=0.0, le=1.0) + """Conditional-acceptance threshold for DSpark prefix scheduling. A value + of zero keeps every proposed position. Positive values retain only the + leading positions whose calibrated confidence is at least the threshold.""" + + dspark_confidence_temperatures: list[float] | None = None + """Optional positive per-position temperatures applied to DSpark confidence + logits before prefix scheduling. The list must match num_speculative_tokens.""" + + dspark_max_verification_tokens: int | None = Field(default=None, ge=0) + """Optional cap on DSpark draft tokens submitted to target verification. + DSpark still generates the checkpoint's complete block; only a prefix is + scheduled, so values below the checkpoint block size remain lossless.""" + def compute_hash(self) -> str: """ WARNING: Whenever a new field is added to this config, @@ -307,7 +334,7 @@ def compute_hash(self) -> str: "eagle3", "extract_hidden_states", ) - or self.use_dflash() + or self.use_dflash_family() or self.use_dspark() ) factors.append(uses_aux_hidden_states) @@ -643,6 +670,22 @@ def __post_init__(self): if self.method in ("ngram", "[ngram]"): self.method = "ngram" + if self.ngram_assist: + if self.prompt_lookup_min is None and self.prompt_lookup_max is None: + self.prompt_lookup_min = 5 + self.prompt_lookup_max = 5 + elif self.prompt_lookup_min is None: + self.prompt_lookup_min = self.prompt_lookup_max + elif self.prompt_lookup_max is None: + self.prompt_lookup_max = self.prompt_lookup_min + assert self.prompt_lookup_min is not None + assert self.prompt_lookup_max is not None + if self.prompt_lookup_min > self.prompt_lookup_max: + raise ValueError( + f"prompt_lookup_min={self.prompt_lookup_min} must " + f"be <= prompt_lookup_max={self.prompt_lookup_max}" + ) + if self.method in ("ngram", "ngram_gpu"): # Set default values if not provided if self.prompt_lookup_min is None and self.prompt_lookup_max is None: @@ -719,8 +762,9 @@ def __post_init__(self): self.draft_parallel_config = self.target_parallel_config else: - self.prompt_lookup_max = 0 - self.prompt_lookup_min = 0 + if not self.ngram_assist: + self.prompt_lookup_max = 0 + self.prompt_lookup_min = 0 if self.model is not None: self.draft_model_config = ModelConfig( @@ -748,7 +792,7 @@ def __post_init__(self): # Automatically detect the method if ( self.method in ("eagle", "eagle3") - or self.use_dflash() + or self.use_dflash_family() or self.use_dspark() ): pass @@ -788,7 +832,7 @@ def __post_init__(self): ) # Replace hf_config for EAGLE draft_model - if self.method in ("eagle", "eagle3") or self.use_dflash(): + if self.method in ("eagle", "eagle3") or self.use_dflash_family(): from vllm.transformers_utils.configs.eagle import EAGLEConfig from vllm.transformers_utils.configs.speculators import ( SpeculatorsConfig, @@ -800,7 +844,9 @@ def __post_init__(self): ): pass else: - eagle_method = "dflash" if self.use_dflash() else self.method + eagle_method = ( + "dflash" if self.use_dflash_family() else self.method + ) eagle_config = EAGLEConfig( self.draft_model_config.hf_config, method=eagle_method, @@ -815,7 +861,7 @@ def __post_init__(self): draft_hf_config.architectures = ["DSparkDraftModel"] self.update_arch_() - if self.use_dflash() or self.use_dspark(): + if self.use_dflash_family() or self.use_dspark(): self.parallel_drafting = True if self.num_speculative_tokens is not None and hasattr( @@ -864,6 +910,32 @@ def __post_init__(self): f"{self.num_speculative_tokens}. Smaller values " "produce incorrect output." ) + if ( + self.dspark_max_verification_tokens is not None + and self.dspark_max_verification_tokens + > self.num_speculative_tokens + ): + raise ValueError( + "dspark_max_verification_tokens must be <= " + f"num_speculative_tokens ({self.num_speculative_tokens}); " + f"got {self.dspark_max_verification_tokens}." + ) + temperatures = self.dspark_confidence_temperatures + if temperatures is not None: + if len(temperatures) != self.num_speculative_tokens: + raise ValueError( + "dspark_confidence_temperatures must have length " + f"{self.num_speculative_tokens}; got {temperatures}." + ) + if not all( + math.isfinite(value) and value > 0.0 + for value in temperatures + ): + raise ValueError( + "dspark_confidence_temperatures entries must be " + f"finite and positive; got {temperatures}." + ) + self._verify_dspark_final_stage_ownership() if self.use_dflash_ddtree() and self.ddtree_budget is None: self.ddtree_budget = self.num_speculative_tokens @@ -886,16 +958,48 @@ def __post_init__(self): self.draft_parallel_config = ( SpeculativeConfig.create_draft_parallel_config( - self.target_parallel_config, self.draft_tensor_parallel_size + self.target_parallel_config, + self.draft_tensor_parallel_size, + # The runner constructs the complete DSpark drafter on + # the target model's final PP rank. It is local to that + # stage rather than partitioned across target PP ranks. + speculative_draft_pipeline_parallel_size=( + 1 if self.use_dspark() else None + ), ) ) return self + def _verify_dspark_final_stage_ownership(self) -> None: + """Validate the layer contract for a final-stage-local drafter.""" + from vllm.distributed.utils import get_pp_indices + + pp_size = self.target_parallel_config.pipeline_parallel_size + num_layers = self.target_model_config.get_total_num_hidden_layers() + last_start, last_end = get_pp_indices(num_layers, pp_size - 1, pp_size) + layer_ids = tuple( + getattr( + self.draft_model_config.hf_config, + "dspark_target_layer_ids", + (), + ) + ) + outside_last_stage = [ + layer_id for layer_id in layer_ids if not last_start <= layer_id < last_end + ] + if not layer_ids or outside_last_stage: + raise ValueError( + "DSpark's final-stage-local drafter requires every " + "dspark_target_layer_id to belong to the final target pipeline " + f"stage [{last_start}, {last_end}); got {layer_ids}. Adjust " + "VLLM_PP_LAYER_PARTITION or the DSpark layer contract." + ) + 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 @@ -1029,13 +1133,19 @@ def update_arch_(self): def create_draft_parallel_config( target_parallel_config: ParallelConfig, speculative_draft_tensor_parallel_size: int, + speculative_draft_pipeline_parallel_size: int | None = None, ) -> ParallelConfig: """Create a parallel config for use by the draft worker. - This is mostly a copy of the target parallel config, except the tp_size. + This is mostly a copy of the target parallel config, except the draft + tensor size and, for final-stage-local drafters, pipeline size. """ draft_parallel_config = ParallelConfig( - pipeline_parallel_size=target_parallel_config.pipeline_parallel_size, + pipeline_parallel_size=( + target_parallel_config.pipeline_parallel_size + if speculative_draft_pipeline_parallel_size is None + else speculative_draft_pipeline_parallel_size + ), tensor_parallel_size=speculative_draft_tensor_parallel_size, distributed_executor_backend=target_parallel_config.distributed_executor_backend, max_parallel_loading_workers=target_parallel_config.max_parallel_loading_workers, @@ -1075,6 +1185,13 @@ def _verify_args(self) -> Self: "Expected num_speculative_tokens to be greater " f"than zero ({self.num_speculative_tokens})." ) + if self.ngram_assist: + if not self.use_dflash(): + raise ValueError("ngram_assist is only supported with method='dflash'.") + draft_hf_config = getattr(self.draft_model_config, "hf_config", None) + dflash_config = getattr(draft_hf_config, "dflash_config", None) or {} + if int(dflash_config.get("selector_top_k", 0) or 0) <= 0: + raise ValueError("ngram_assist requires DFlash2 selector capability.") if self.use_dflash_ddtree(): if self.ddtree_budget is None: self.ddtree_budget = self.num_speculative_tokens @@ -1135,6 +1252,12 @@ def max_num_new_slots_for_drafting(self) -> int: Calculate the maximum number of new slots that might be added to the batch when drafting. """ + if self.use_dflash(): + # MRV2 DFlash uses one bonus query followed by K mask queries. The + # scheduler already owns the target query slot, so K more slots + # must be available to the drafter. + return self.num_speculative_tokens + slots_per_req = 0 # for serial non-draft-model methods, no change needed if self.parallel_drafting: # For parallel drafting, we need one new slot per 'masked' token @@ -1164,16 +1287,28 @@ def use_step3p5_mtp(self) -> bool: def use_eagle(self) -> bool: return ( self.method in ("eagle", "eagle3", "mtp") - or self.use_dflash() + or self.use_dflash_family() or self.use_dspark() ) + def use_eagle_kv_cache(self) -> bool: + """Whether prefix hits need Eagle's final target-block recompute.""" + # MRV2 DFlash caches its projected draft context KV alongside the + # target KV. It only needs target hidden states for the uncached suffix, + # so dropping an additional target block is both unnecessary and can + # eliminate all prefix hits when hybrid page unification makes blocks + # large. Keep DDTree and DSpark on their established V1 semantics. + return self.use_eagle() and not self.use_dflash() + def use_dflash(self) -> bool: - return self.method in get_args(DFlashModelTypes) + return self.method == "dflash" def use_dflash_ddtree(self) -> bool: return self.method == "dflash_ddtree" + def use_dflash_family(self) -> bool: + return self.use_dflash() or self.use_dflash_ddtree() + def use_dspark(self) -> bool: return self.method == "dspark" diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index ea67803abb..961b979c12 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -68,6 +68,7 @@ DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset({"Qwen3ForCausalLM"}) _SM70_NOMTP_CUDAGRAPH_CAPTURE_SIZES = (1, 2, 4, 8, 16) +_SM70_MTP_CUDAGRAPH_REQUEST_SIZES = (1, 2, 4, 6, 8, 12, 16) def _sm70_nomtp_cudagraph_capture_sizes(max_num_seqs: int) -> list[int]: @@ -79,6 +80,19 @@ def _sm70_nomtp_cudagraph_capture_sizes(max_num_seqs: int) -> list[int]: return sorted(capture_sizes) +def _sm70_mtp_cudagraph_capture_sizes( + max_num_seqs: int, + decode_query_len: int, +) -> list[int]: + """Return exact SM70 MTP verifier token shapes for production requests.""" + max_graph_reqs = min(max(int(max_num_seqs), 1), 16) + request_sizes = { + size for size in _SM70_MTP_CUDAGRAPH_REQUEST_SIZES if size <= max_graph_reqs + } + request_sizes.add(max_graph_reqs) + return [decode_query_len * size for size in sorted(request_sizes)] + + class OptimizationLevel(IntEnum): """Optimization level enum.""" @@ -216,6 +230,81 @@ def enable_mla_dual_rms_norm_fusion(cfg: "VllmConfig") -> bool: return rocm_aiter_ops.is_enabled() and check_aiter_fused_qk_rmsnorm() +def apply_prefix_anchored_swa_constraints(cfg: "VllmConfig") -> None: + """Validate the model-agnostic prefix-anchored SWA engine contract.""" + window = cfg.attention_config.prefix_anchored_decode_window + if window is None: + return + + from vllm.platforms import current_platform + from vllm.v1.attention.backends.registry import AttentionBackendEnum + + if not ( + current_platform.is_cuda() and current_platform.is_device_capability((7, 0)) + ): + raise ValueError("prefix_anchored_decode_window requires an NVIDIA SM70 GPU") + + model_config = cfg.model_config + if model_config is None: + raise ValueError("prefix_anchored_decode_window requires a model config") + if model_config.use_mla: + raise ValueError("prefix_anchored_decode_window does not support MLA attention") + if cfg.attention_config.use_non_causal: + raise ValueError( + "prefix_anchored_decode_window requires causal decoder attention" + ) + + cache_dtype = cfg.cache_config.cache_dtype + effective_cache_dtype = model_config.dtype if cache_dtype == "auto" else cache_dtype + if effective_cache_dtype not in (torch.float16, "float16"): + raise ValueError( + "prefix_anchored_decode_window requires an fp16 KV cache; " + f"got cache_dtype={cache_dtype!r} and model dtype={model_config.dtype}" + ) + + parallel_config = cfg.parallel_config + if ( + parallel_config.decode_context_parallel_size != 1 + or parallel_config.prefill_context_parallel_size != 1 + ): + raise ValueError( + "prefix_anchored_decode_window does not support decode or prefill " + "context parallelism" + ) + if cfg.cache_config.kv_offloading_size is not None: + raise ValueError( + "prefix_anchored_decode_window does not support KV-cache offloading" + ) + if ( + cfg.kv_transfer_config is not None + and cfg.kv_transfer_config.kv_connector is not None + ): + raise ValueError("prefix_anchored_decode_window does not support KV connectors") + if cfg.speculative_config is not None: + raise ValueError( + "prefix_anchored_decode_window does not yet support speculative decoding" + ) + + backend = cfg.attention_config.backend + if backend is None: + cfg.attention_config.backend = AttentionBackendEnum.FLASH_ATTN_V100 + logger.info( + "Prefix-anchored sliding-window attention: auto-selecting the " + "FLASH_ATTN_V100 backend." + ) + elif backend != AttentionBackendEnum.FLASH_ATTN_V100: + raise ValueError( + "prefix_anchored_decode_window requires FLASH_ATTN_V100; " + f"got {backend.name}" + ) + if cfg.cache_config.enable_prefix_caching: + cfg.cache_config.enable_prefix_caching = False + logger.info( + "Prefix-anchored sliding-window attention: disabling prefix " + "caching (windowed decode KV is not reusable across requests)." + ) + + OPTIMIZATION_LEVEL_00 = { "compilation_config": { "pass_config": { @@ -511,6 +600,28 @@ def compute_hash(self) -> str: ] return hash_str + @property + def max_concurrent_batches(self) -> int: + # Scheduler-side upper bound mirroring the executors' + # `max_concurrent_batches` (multiproc: queue depth under async + # scheduling, else the PP pipeline depth). + pp_size = self.parallel_config.pipeline_parallel_size + if pp_size <= 1 and self.scheduler_config.async_scheduling: + return max(2, envs.VLLM_SM70_ASYNC_SCHEDULING_QUEUE_DEPTH) + return pp_size + + @property + def max_in_flight_tokens(self) -> int: + # Upper bound on tokens that are scheduled but not yet settled (freed): + # every concurrent batch may hold up to a full `max_num_batched_tokens`. + # Recycling-aware KV cache specs (sliding-window, chunked-local, + # prefix-anchored SWA) reserve for this because out-of-window blocks + # are freed on the processed-token basis, so in-flight steps + # transiently keep their blocks. + return ( + self.max_concurrent_batches * self.scheduler_config.max_num_batched_tokens + ) + @property def num_speculative_tokens(self) -> int: if ( @@ -520,12 +631,39 @@ def num_speculative_tokens(self) -> int: return self.speculative_config.num_speculative_tokens return 0 + @property + def num_lookahead_tokens(self) -> int: + """KV positions reserved beyond the target model's scheduled query.""" + speculative_config = self.speculative_config + if speculative_config is None: + return 0 + if speculative_config.use_dflash_family(): + # DFlash has one anchor query plus the checkpoint-owned draft state. + return speculative_config.num_speculative_state_tokens() + 1 + if speculative_config.use_eagle() or speculative_config.uses_draft_model(): + return self.num_speculative_tokens + return 0 + @property def use_v2_model_runner(self) -> bool: use_v2_model_runner = envs.VLLM_USE_V2_MODEL_RUNNER + is_mrv2_dflash = ( + self.speculative_config is not None and self.speculative_config.use_dflash() + ) if use_v2_model_runner is not None: + if is_mrv2_dflash and not use_v2_model_runner: + raise ValueError( + "method='dflash' is implemented only by Model Runner V2. " + "Use method='dflash_ddtree' for the retained V1 DDTree route." + ) return use_v2_model_runner + # The old V1 flat-DFlash route is intentionally removed. Force every + # method='dflash' checkpoint through the MRV2 speculator and fail in + # _validate_v2_model_runner if another configured feature is unsupported. + if is_mrv2_dflash: + return True + if not self._is_default_v2_model_runner_model(): return False @@ -826,8 +964,8 @@ def _verify_kv_transfer_compat(self) -> None: # producing RDMA failures like IBV_WC_REM_ACCESS_ERR / # NIXL_ERR_REMOTE_DISCONNECT at the first inter-node KV transfer. # We can't enumerate every in-tree and out-of-tree connector that - # pins memory, so we conservatively reject the combination whenever - # any KV connector is configured. + # pins memory, so we conservatively reject the combination unless the + # connector explicitly opts into the VMM-safe transfer contract. # # CuMem allocator is exempt: CuMemAllocator.use_memory_pool toggles # expandable_segments off around its pool (see #40812), so the KV @@ -839,6 +977,14 @@ def _verify_kv_transfer_compat(self) -> None: return if self.model_config is not None and (self.model_config.enable_cumem_allocator): return + if self._connector_supports_vmm_safe_transfers(): + logger.warning_once( + "Allowing KV connector %s with " + "PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True because " + "it implements SupportsVmmSafeTransfers.", + self.kv_transfer_config.kv_connector, + ) + return raise ValueError( f"KV connector {self.kv_transfer_config.kv_connector} is " @@ -853,6 +999,25 @@ def _verify_kv_transfer_compat(self) -> None: "expandable_segments is automatically disabled)." ) + def _connector_supports_vmm_safe_transfers(self) -> bool: + """Resolve the connector and check its explicit VMM-safety marker.""" + from vllm.distributed.kv_transfer.kv_connector.factory import ( + KVConnectorFactory, + ) + from vllm.distributed.kv_transfer.kv_connector.v1 import ( + supports_vmm_safe_transfers, + ) + + kv_transfer_config = self.kv_transfer_config + if kv_transfer_config is None: + return False + + try: + connector_cls = KVConnectorFactory.get_connector_class(kv_transfer_config) + return supports_vmm_safe_transfers(connector_cls) + except (ValueError, AttributeError, ImportError, TypeError): + return False + def __post_init__(self): """Verify configs are valid & consistent with each other.""" @@ -897,6 +1062,8 @@ def __post_init__(self): if self.lora_config is not None: self.lora_config.verify_with_model_config(self.model_config) + apply_prefix_anchored_swa_constraints(self) + if ( self.mamba_config.enable_stochastic_rounding and self.cache_config.mamba_ssm_cache_dtype != "float16" @@ -1307,7 +1474,6 @@ def __post_init__(self): env_name, env_value, ) - sm70_flash_0dot3_compile_graph = envs.VLLM_SM70_FLASH_V100_0DOT3_COMPILE_GRAPH sm70_flash_no_compile_graph = ( envs.VLLM_SM70_FLASH_V100_DECODE_GRAPH_NO_COMPILE @@ -1341,11 +1507,6 @@ def __post_init__(self): self.speculative_config is not None and self.speculative_config.num_speculative_tokens ): - cudagraph_capture_sizes = ( - [1, 2, 4, 8, 9, 18] - if self.parallel_config.tensor_parallel_size >= 4 - else [1, 2, 4, 8, 9] - ) decode_query_len = ( self.speculative_config.num_speculative_state_tokens() + 1 ) @@ -1363,26 +1524,47 @@ def __post_init__(self): smallq_env, decode_query_len, ) - max_graph_reqs = ( - 4 if self.parallel_config.tensor_parallel_size >= 4 else 1 - ) - max_graph_reqs = min( - max(int(self.scheduler_config.max_num_seqs), 1), - max_graph_reqs, - ) - cudagraph_capture_sizes = sorted( - set(cudagraph_capture_sizes) - | { - decode_query_len * num_reqs - for num_reqs in range(1, max_graph_reqs + 1) - } - ) - logger.info_once( - "Using SM70 speculative verifier cudagraph shapes " - "%sx1..%s for Flash-V100 compile graph.", - decode_query_len, - max_graph_reqs, - ) + if ( + envs.VLLM_SM70_MTP_SPLIT_DRAFT_CUDAGRAPHS + and self.speculative_config.method == "mtp" + ): + cudagraph_capture_sizes = _sm70_mtp_cudagraph_capture_sizes( + self.scheduler_config.max_num_seqs, + decode_query_len, + ) + logger.info_once( + "Using split SM70 MTP verifier cudagraph token " + "shapes %s for Flash-V100 compile graph.", + tuple(cudagraph_capture_sizes), + ) + else: + cudagraph_capture_sizes = ( + [1, 2, 4, 8, 9, 18] + if self.parallel_config.tensor_parallel_size >= 4 + else [1, 2, 4, 8, 9] + ) + max_graph_reqs = ( + 4 + if self.parallel_config.tensor_parallel_size >= 4 + else 1 + ) + max_graph_reqs = min( + max(int(self.scheduler_config.max_num_seqs), 1), + max_graph_reqs, + ) + cudagraph_capture_sizes = sorted( + set(cudagraph_capture_sizes) + | { + decode_query_len * num_reqs + for num_reqs in range(1, max_graph_reqs + 1) + } + ) + logger.info_once( + "Using SM70 speculative verifier cudagraph shapes " + "%sx1..%s for Flash-V100 compile graph.", + decode_query_len, + max_graph_reqs, + ) elif cudagraph_capture_sizes != [1, 2]: logger.info_once( "Using SM70 no-MTP decode cudagraph request shapes %s.", @@ -1689,6 +1871,22 @@ def enable_quant_fp8_custom_op_for_blocked_weights() -> bool: CUDAGraphMode.FULL_DECODE_ONLY ) + # Prefix-anchored sliding-window attention: full-graph capture + # bakes attention metadata that carries no per-request anchor + # lengths, so the anchored decode-window mask cannot run inside a + # full cudagraph. Piecewise graphs keep attention outside capture. + if ( + self.attention_config.prefix_anchored_decode_window is not None + and self.compilation_config.cudagraph_mode.has_full_cudagraphs() + ): + logger.info( + "Prefix-anchored sliding-window attention does not " + "support full cudagraphs. Overriding cudagraph_mode " + "from %s to PIECEWISE.", + self.compilation_config.cudagraph_mode.name, + ) + self.compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE + # Check if KV connector requires PIECEWISE mode for CUDA graphs if ( self.kv_transfer_config is not None @@ -2455,13 +2653,6 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]: model_config = self.model_config speculative_config = self.speculative_config - if ( - model_config is not None - and model_config.has_inner_state - and self.cache_config.mamba_cache_mode == "align" - ): - unsupported.append("hybrid/mamba models with align cache mode") - if self.parallel_config.prefill_context_parallel_size > 1: unsupported.append("prefill context parallelism") @@ -2478,11 +2669,19 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]: # TODO: ngram / ngram_gpu are not supported by the v2 model runner yet if speculative_config.method in ("ngram", "ngram_gpu"): unsupported.append("ngram/ngram_gpu speculative decoding") - elif speculative_config.method not in ("eagle", "eagle3", "mtp"): + elif speculative_config.method not in ( + "eagle", + "eagle3", + "mtp", + "dflash", + ): unsupported.append(f"speculative method '{speculative_config.method}'") # V2 EagleSpeculator does not support parallel_drafting (required by PEagle) - if speculative_config.parallel_drafting: + if ( + speculative_config.parallel_drafting + and speculative_config.method != "dflash" + ): unsupported.append("parallel drafting for speculative decoding") if ( @@ -2490,6 +2689,11 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]: and self.parallel_config.pipeline_parallel_size > 1 ): unsupported.append("EAGLE3 with pipeline parallelism") + if ( + speculative_config.method == "dflash" + and self.parallel_config.pipeline_parallel_size > 1 + ): + unsupported.append("DFlash with pipeline parallelism") if self.parallel_config.enable_dbo: unsupported.append("dual batch overlap") @@ -2592,10 +2796,6 @@ def validate_block_size(self) -> None: "to schedule a multiple of block_size tokens even if they are " "in the middle of a mm input" ) - # TODO: support align mamba cache mode for model runner v2 - assert not envs.VLLM_USE_V2_MODEL_RUNNER, ( - "Model Runner V2 has not yet supported mamba_cache_mode='align'. " - ) @model_validator(mode="after") def validate_nvfp4_kv_cache_with_mla(self) -> "VllmConfig": diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index 34b999e5d5..fc46b2419e 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -21,9 +21,10 @@ from .base_device_communicator import DeviceCommunicatorBase logger = init_logger(__name__) -_SEEN_TP_ALLREDUCE_PATHS: set[ - tuple[str, str, tuple[int, ...], torch.dtype, int] -] = set() +_SM70_TP4_LONG_PREFILL_BUFFER_BYTES = 8192 * 5120 * 2 +_SEEN_TP_ALLREDUCE_PATHS: set[tuple[str, str, tuple[int, ...], torch.dtype, int]] = ( + set() +) def _trace_all_reduce_path( @@ -82,6 +83,7 @@ def __init__( use_top1_custom_ar = False use_sm70_awq_mlp_down_tile_ar = False use_sm70_awq_mlp_down_tile_overlap = False + use_sm70_tp4_long_prefill_fused_norm = False use_torch_symm_mem = False use_flashinfer_allreduce = False else: @@ -93,15 +95,19 @@ def __init__( use_sm70_awq_mlp_down_tile_overlap = ( envs.VLLM_SM70_AWQ_MLP_DOWN_TILE_OVERLAP ) + use_sm70_tp4_long_prefill_fused_norm = ( + envs.VLLM_SM70_TP4_LONG_PREFILL_FUSED_NORM + and self.world_size == 4 + and current_platform.is_device_capability(70) + ) use_torch_symm_mem = envs.VLLM_ALLREDUCE_USE_SYMM_MEM use_flashinfer_allreduce = envs.VLLM_ALLREDUCE_USE_FLASHINFER self.use_custom_allreduce = use_custom_allreduce self.use_top1_custom_ar = use_top1_custom_ar self.use_sm70_awq_mlp_down_tile_ar = use_sm70_awq_mlp_down_tile_ar - self.use_sm70_awq_mlp_down_tile_overlap = ( - use_sm70_awq_mlp_down_tile_overlap - ) + self.use_sm70_awq_mlp_down_tile_overlap = use_sm70_awq_mlp_down_tile_overlap + self.use_sm70_tp4_long_prefill_fused_norm = use_sm70_tp4_long_prefill_fused_norm self.use_torch_symm_mem = use_torch_symm_mem self.use_flashinfer_allreduce = use_flashinfer_allreduce @@ -149,6 +155,7 @@ def __init__( or use_top1_custom_ar or use_sm70_awq_mlp_down_tile_ar or use_sm70_awq_mlp_down_tile_overlap + or use_sm70_tp4_long_prefill_fused_norm ) and self.world_size > 1: # Initialize a custom fast all-reduce implementation. # `use_top1_custom_ar` deliberately only provisions the IPC/signaling @@ -161,6 +168,12 @@ def __init__( symm_mem_enabled=( self.symm_mem_comm is not None and not self.symm_mem_comm.disabled ), + max_size=( + _SM70_TP4_LONG_PREFILL_BUFFER_BYTES + if use_sm70_tp4_long_prefill_fused_norm + else 8192 * 1024 + ), + long_prefill_fusion_enabled=use_sm70_tp4_long_prefill_fused_norm, ) if current_platform.is_rocm(): @@ -435,6 +448,31 @@ def sm70_tp2_all_reduce_gemma_rms_norm( input_, residual, weight, epsilon ) + def sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather( + self, + input_: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + epsilon: float, + ) -> tuple[torch.Tensor, torch.Tensor]: + ca_comm = self.ca_comm + if ( + not self.use_sm70_tp4_long_prefill_fused_norm + or ca_comm is None + or ca_comm.disabled + or not ca_comm.can_sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather( + input_, residual, weight + ) + ): + raise RuntimeError( + "SM70 TP4 long-prefill fused collective-norm route was " + "compiled for an unsupported runtime shape or communicator." + ) + _trace_all_reduce_path(self, "sm70_tp4_fused_rs_norm_ag", input_) + return ca_comm.sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather( + input_, residual, weight, epsilon + ) + def sm70_awq_mlp_down_tile_all_reduce(self, input_): if not self.use_sm70_awq_mlp_down_tile_ar: return None @@ -468,18 +506,14 @@ def sm70_awq_mlp_down_tile_gemm_reduce( k_ld, q_ld, tile_numel=envs.VLLM_SM70_AWQ_MLP_DOWN_TILE_OVERLAP_TILE_NUMEL, - reducer_blocks=( - envs.VLLM_SM70_AWQ_MLP_DOWN_TILE_OVERLAP_REDUCER_BLOCKS - ), + reducer_blocks=(envs.VLLM_SM70_AWQ_MLP_DOWN_TILE_OVERLAP_REDUCER_BLOCKS), kernel_reducer_blocks=( envs.VLLM_SM70_AWQ_MLP_DOWN_TILE_OVERLAP_KERNEL_REDUCER_BLOCKS ), overlap=envs.VLLM_SM70_AWQ_MLP_DOWN_TILE_OVERLAP_SIDE_STREAM, ) if out is not None: - _trace_all_reduce_path( - self, "sm70_awq_mlp_down_tile_overlap", input_ - ) + _trace_all_reduce_path(self, "sm70_awq_mlp_down_tile_overlap", input_) return out def reduce_scatter(self, input_: torch.Tensor, dim: int = -1): diff --git a/vllm/distributed/device_communicators/custom_all_reduce.py b/vllm/distributed/device_communicators/custom_all_reduce.py index a06dfebc26..8d55ceeaff 100644 --- a/vllm/distributed/device_communicators/custom_all_reduce.py +++ b/vllm/distributed/device_communicators/custom_all_reduce.py @@ -1,9 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os from contextlib import contextmanager from typing import cast +import regex as re import torch import torch.distributed as dist from torch.distributed import ProcessGroup @@ -28,6 +30,49 @@ logger = init_logger(__name__) _SM70_TP8_HIERARCHICAL_ELEMENTS = 4096 +_EXPANDABLE_SEGMENTS_TRUE_PATTERN = re.compile( + r"((?:^|,)\s*expandable_segments\s*:\s*)True(?=\s*(?:,|$))" +) + + +def _get_effective_allocator_conf() -> str: + """Return the allocator config selected by PyTorch on CUDA. + + PyTorch 2.10 prefers the legacy CUDA-specific variable whenever it is + present, even when its value is empty, and otherwise falls back to the + unified accelerator variable. + """ + if "PYTORCH_CUDA_ALLOC_CONF" in os.environ: + return os.environ["PYTORCH_CUDA_ALLOC_CONF"] + return os.environ.get("PYTORCH_ALLOC_CONF", "") + + +@contextmanager +def _disable_expandable_segments_for_cuda_ipc(active: bool): + """Keep CUDA-IPC graph buffers on cudaMalloc-backed allocations. + + ``cudaIpcGetMemHandle`` cannot export CUDA VMM allocations created when + ``expandable_segments:True`` is active through either supported allocator + environment variable. Custom all-reduce exports graph buffers through that + legacy IPC API, so disable expandable segments for the complete capture + and registration window while preserving every other allocator option. + """ + allocator_conf = _get_effective_allocator_conf() + disabled_conf, replacements = _EXPANDABLE_SEGMENTS_TRUE_PATTERN.subn( + r"\1False", allocator_conf + ) + should_disable = active and current_platform.is_cuda() and replacements > 0 + if should_disable: + logger.info_once( + "Temporarily disabling expandable_segments during custom allreduce " + "CUDA graph capture for CUDA IPC compatibility." + ) + torch.cuda.memory._set_allocator_settings(disabled_conf) + try: + yield + finally: + if should_disable: + torch.cuda.memory._set_allocator_settings(allocator_conf) def _sm70_tp8_hierarchical_peer_ranks(rank: int) -> tuple[int, ...]: @@ -68,6 +113,7 @@ def __init__( device: int | str | torch.device, max_size=8192 * 1024, symm_mem_enabled=False, + long_prefill_fusion_enabled=False, ) -> None: """ Args: @@ -81,6 +127,8 @@ def __init__( """ self._IS_CAPTURING = False self.disabled = True + self.long_prefill_output_ptrs: list[int] | None = None + self.sm70_tp4_push_buffer_ptrs: list[int] | None = None if not custom_ar: # disable because of missing custom allreduce library @@ -133,6 +181,7 @@ def __init__( if ( current_platform.is_cuda() and symm_mem_enabled + and not long_prefill_fusion_enabled and device_capability is not None ): device_capability_str = device_capability.as_version_str() @@ -251,6 +300,12 @@ def __init__( 8 * 1024 * 1024, dtype=torch.uint8, device=self.device ) self.max_size = max_size + # Provisioning the long-prefill fusion buffers must not widen ordinary + # custom-AR dispatch beyond its established 8-MiB policy. The fused + # predicate below uses the full allocation capacity explicitly. + self.dispatch_max_size = ( + min(max_size, 8192 * 1024) if long_prefill_fusion_enabled else max_size + ) self.rank = rank self.world_size = world_size self.fully_connected = fully_connected @@ -259,6 +314,33 @@ def __init__( self.meta_ptrs, self.rank_data, rank, self.fully_connected ) ops.register_buffer(self._ptr, self.buffer_ptrs) + if long_prefill_fusion_enabled: + self.long_prefill_output_ptrs = self.create_shared_buffer( + max_size, + group=group, + ) + ops.register_buffer(self._ptr, self.long_prefill_output_ptrs) + if ( + envs.VLLM_SM70_TP4_PUSH_ALLREDUCE + and world_size == 4 + and fully_connected + and current_platform.is_cuda() + and device_capability is not None + and device_capability.major == 7 + and device_capability.minor == 0 + ): + push_buffer_size = ops.sm70_tp4_push_allreduce_buffer_size() + self.sm70_tp4_push_buffer_ptrs = self.create_shared_buffer( + push_buffer_size, + group=group, + ) + ops.register_sm70_tp4_push_allreduce_buffer( + self._ptr, self.sm70_tp4_push_buffer_ptrs + ) + logger.info( + "SM70 TP4 SGLang-style push all-reduce enabled for the " + "FP16 80-KiB verifier and 8-KiB decode payloads." + ) @contextmanager def capture(self): @@ -267,13 +349,14 @@ def capture(self): `register_graph_buffers` call at the end of the context. It records all the buffer addresses used in the CUDA graph. """ - try: - self._IS_CAPTURING = True - yield - finally: - self._IS_CAPTURING = False - if not self.disabled: - self.register_graph_buffers() + with _disable_expandable_segments_for_cuda_ipc(not self.disabled): + try: + self._IS_CAPTURING = True + yield + finally: + self._IS_CAPTURING = False + if not self.disabled: + self.register_graph_buffers() def register_graph_buffers(self): handle, offset = ops.get_graph_buffer_ipc_meta(self._ptr) @@ -311,7 +394,7 @@ def should_custom_ar(self, inp: torch.Tensor): # for 4 or more non NVLink-capable GPUs, custom allreduce provides # little performance improvement over NCCL. if self.world_size == 2 or self.fully_connected: - return inp_size < self.max_size + return inp_size < self.dispatch_max_size return False def all_reduce( @@ -453,6 +536,72 @@ def can_sm70_tp4_all_reduce_gemma_rms_norm( and inp.device == residual.device == weight.device ) + def sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather( + self, + inp: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + epsilon: float, + *, + normalized_out: torch.Tensor | None = None, + residual_out: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Benchmark-only long-prefill fused RS + Gemma RMSNorm + AG.""" + if not self.can_sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather( + inp, residual, weight + ): + raise RuntimeError("SM70 TP4 long-prefill fused norm is unavailable") + assert self.long_prefill_output_ptrs is not None + if normalized_out is None: + normalized_out = torch.empty_like(inp) + if residual_out is None: + residual_out = torch.empty_like(residual, dtype=torch.float32) + graph_registered = ( + self._IS_CAPTURING and torch.cuda.is_current_stream_capturing() + ) + ops.sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather( + self._ptr, + inp, + residual, + weight, + normalized_out, + residual_out, + 0 if graph_registered else self.buffer_ptrs[self.rank], + 0 if graph_registered else self.long_prefill_output_ptrs[self.rank], + self.max_size, + epsilon, + ) + return normalized_out, residual_out + + def can_sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather( + self, + inp: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + ) -> bool: + return ( + not self.disabled + and self.long_prefill_output_ptrs is not None + and self.world_size == 4 + and self.fully_connected + and inp.is_cuda + and inp.dtype == torch.float16 + and inp.ndim == 2 + and inp.shape[0] % self.world_size == 0 + and 1 <= inp.shape[0] // self.world_size <= 2048 + and inp.shape[1] == 5120 + and inp.numel() * inp.element_size() <= self.max_size + and residual.shape == inp.shape + and residual.dtype == torch.float32 + and weight.ndim == 1 + and weight.numel() == 5120 + and weight.dtype in (torch.float16, torch.float32) + and inp.is_contiguous() + and residual.is_contiguous() + and weight.is_contiguous() + and inp.device == residual.device == weight.device + ) + def top1_argmax( self, input_pair: torch.Tensor, @@ -718,6 +867,15 @@ def close(self): self._ptr = 0 self.free_shared_buffer(self.meta_ptrs, rank=self.rank) self.free_shared_buffer(self.buffer_ptrs, rank=self.rank) + if self.long_prefill_output_ptrs is not None: + self.free_shared_buffer( + self.long_prefill_output_ptrs, + rank=self.rank, + ) + self.long_prefill_output_ptrs = None + if self.sm70_tp4_push_buffer_ptrs is not None: + self.free_shared_buffer(self.sm70_tp4_push_buffer_ptrs, rank=self.rank) + self.sm70_tp4_push_buffer_ptrs = None def __del__(self): self.close() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/__init__.py b/vllm/distributed/kv_transfer/kv_connector/v1/__init__.py index 47329207f4..96ab863d02 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/__init__.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/__init__.py @@ -4,7 +4,9 @@ KVConnectorBase_V1, KVConnectorRole, SupportsHMA, + SupportsVmmSafeTransfers, supports_hma, + supports_vmm_safe_transfers, ) from vllm.distributed.kv_transfer.kv_connector.v1.decode_bench_connector import ( # noqa: E501 DecodeBenchConnector, @@ -15,5 +17,7 @@ "KVConnectorBase_V1", "supports_hma", "SupportsHMA", + "SupportsVmmSafeTransfers", + "supports_vmm_safe_transfers", "DecodeBenchConnector", ] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/base.py index fb5658da88..bf98b69872 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/base.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/base.py @@ -121,6 +121,21 @@ def supports_hma(connector: Any) -> bool: return isinstance(connector, SupportsHMA) +class SupportsVmmSafeTransfers(ABC): # noqa: B024 + """Marker for connectors safe under CUDA VMM remapping. + + Implement this marker only when every KV transfer dereferences GPU + virtual addresses at transfer time (for example, cuMemcpy/cudaMemcpy), + and the connector never registers KV device memory with RDMA, GDS, + CUDA IPC, or another mechanism that depends on stable physical pages. + """ + + +def supports_vmm_safe_transfers(connector_cls: type) -> bool: + """Return whether a connector class explicitly declares VMM safety.""" + return issubclass(connector_cls, SupportsVmmSafeTransfers) + + class KVConnectorRole(enum.Enum): # Connector running in the scheduler process SCHEDULER = 0 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 6560bcf9de..d0f2124b22 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -15,7 +15,7 @@ TransferJob, ) from vllm.logger import init_logger -from vllm.utils.math_utils import cdiv +from vllm.utils.math_utils import cdiv, round_down from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_cache_interface import ( @@ -73,6 +73,10 @@ class GroupOffloadConfig(NamedTuple): # than the MLA full-attention group). # None for full-attention groups or when the optimization doesn't apply. alignment_block_count: int | None = None + # Mamba ``align`` tables mutate existing logical positions. Such groups + # must be stored from exact core-provided boundary block IDs instead of the + # connector's append-only positional mirror. + requires_exact_boundary_source: bool = False def get_sliding_window_size_in_blocks( @@ -90,6 +94,24 @@ def get_sliding_window_size_in_blocks( return None +def resolve_mamba_align_size(spec: "OffloadingSpec") -> int | None: + """Scan all KV cache groups in *spec* and return the single mamba alignment + size, or None if no group requires mamba alignment. + + For MambaSpec groups in "align" cache mode the hit window must be rounded + down to a multiple of the offloaded block size. Asserts that all such + groups agree on the same value. + """ + mamba_align_size: int | None = None + for idx, gpu_block_size in enumerate(spec.gpu_block_size): + kv_spec = spec.kv_cache_config.kv_cache_groups[idx].kv_cache_spec + if isinstance(kv_spec, MambaSpec) and kv_spec.mamba_cache_mode == "align": + offload_block_size = gpu_block_size * spec.block_size_factor + assert mamba_align_size is None or mamba_align_size == offload_block_size + mamba_align_size = offload_block_size + return mamba_align_size + + class SchedulerOffloadConfig(NamedTuple): kv_group_configs: tuple[GroupOffloadConfig, ...] block_size_factor: int @@ -133,6 +155,12 @@ def _alignment_block_count( return None return per_segment + def _requires_exact_boundary_source(group_idx: int) -> bool: + kv_spec = spec.kv_cache_config.kv_cache_groups[group_idx].kv_cache_spec + return ( + isinstance(kv_spec, MambaSpec) and kv_spec.mamba_cache_mode == "align" + ) + return cls( num_workers=spec.vllm_config.parallel_config.world_size, kv_group_configs=tuple( @@ -153,6 +181,9 @@ def _alignment_block_count( alignment_block_count=_alignment_block_count( gpu_block_size * spec.block_size_factor, sw ), + requires_exact_boundary_source=( + _requires_exact_boundary_source(idx) + ), ) for idx, gpu_block_size in enumerate(spec.gpu_block_size) ), @@ -282,6 +313,7 @@ def _sliding_window_sort_key(i: int) -> int: # used by _lookup self._sliding_window_groups: tuple[int, ...] = tuple(sliding_window_groups) self._lookup_groups = tuple(full_attention_groups) + self._sliding_window_groups + self._mamba_align_size: int | None = resolve_mamba_align_size(spec) self._req_status: dict[ReqId, RequestOffloadState] = {} self._current_batch_load_jobs: dict[int, TransferJob] = {} @@ -312,7 +344,10 @@ def _generate_job_id(self) -> int: return job_id def _remove_pending_job(self, job_id: int, block_ids: list[int] | None) -> None: - for bid in block_ids or (): + # Registration is idempotent because each block maps to a set of job IDs. + # MTP lookahead may report the same physical block ID more than once, so + # completion must use the same set semantics instead of removing twice. + for bid in set(block_ids or ()): pending = self._block_id_to_pending_jobs[bid] pending.remove(job_id) if not pending: @@ -398,6 +433,11 @@ def _lookup(self, req_status: RequestOffloadState) -> int | None: # for sliding window attention, we must reduce by 1 to make sure # we still have a hit after reduction max_hit_size_tokens -= 1 + if self._mamba_align_size is not None: + # Constrain hit-window to the mamba block size. + max_hit_size_tokens = round_down( + max_hit_size_tokens, self._mamba_align_size + ) num_hit_tokens: int = 0 defer_lookup = False lookup_groups = self._lookup_groups @@ -727,6 +767,8 @@ def _build_store_jobs( for group_config, group_state in zip( self.config.kv_group_configs, req_status.group_states ): + if group_config.requires_exact_boundary_source: + continue num_blocks = num_offloadable_tokens // group_config.offloaded_block_size start_block_idx = group_state.next_stored_block_idx if num_blocks <= start_block_idx: @@ -795,6 +837,10 @@ def _build_store_jobs( for group_config, group_state in zip( self.config.kv_group_configs, req_status.group_states ): + if group_config.requires_exact_boundary_source: + group_sizes.append(0) + block_indices.append(0) + continue is_sliding_window = ( group_config.sliding_window_size_in_blocks is not None ) @@ -872,6 +918,99 @@ def _build_store_jobs( return store_jobs + def _build_boundary_state_store_jobs( + self, + scheduler_output: SchedulerOutput, + ) -> dict[int, TransferJob]: + """Build stores from exact committed Mamba boundary-state blocks.""" + handoffs = scheduler_output.boundary_state_offloads + if not handoffs: + return {} + + store_jobs: dict[int, TransferJob] = {} + num_groups = len(self.config.kv_group_configs) + for req_id, entries in handoffs.items(): + req_status = self._req_status.get(req_id) + if req_status is None: + continue + req_status.update_offload_keys() + req = req_status.req + num_scheduled_tokens = scheduler_output.num_scheduled_tokens.get(req_id, 0) + num_offloadable_tokens = min( + req.num_computed_tokens + num_scheduled_tokens, + req.num_tokens, + ) + if req_status.max_offload_tokens is not None: + num_offloadable_tokens = min( + num_offloadable_tokens, req_status.max_offload_tokens + ) + if self.config.offload_prompt_only: + num_offloadable_tokens = min( + num_offloadable_tokens, req.num_prompt_tokens + ) + + for group_idx, block_id, boundary_tokens in entries: + group_config = self.config.kv_group_configs[group_idx] + if not group_config.requires_exact_boundary_source: + continue + if ( + block_id == 0 + or boundary_tokens > num_offloadable_tokens + or boundary_tokens % group_config.offloaded_block_size != 0 + ): + continue + + offload_block_idx = ( + boundary_tokens // group_config.offloaded_block_size - 1 + ) + group_state = req_status.group_states[group_idx] + if offload_block_idx >= len(group_state.offload_keys): + continue + offload_key = group_state.offload_keys[offload_block_idx] + store_output = self.manager.prepare_store( + [offload_key], req_status.req_context + ) + if store_output is None: + logger.warning( + "Request %s: cannot store Mamba boundary at %d tokens", + req_id, + boundary_tokens, + ) + continue + keys_to_store = set(store_output.keys_to_store) + if not keys_to_store: + continue + assert keys_to_store == {offload_key} + + job_id = self._generate_job_id() + req_status.transfer_jobs.add(job_id) + self._block_id_to_pending_jobs.setdefault(block_id, set()).add(job_id) + self._jobs[job_id] = TransferJobStatus( + req_id=req_id, + pending_count=self.config.num_workers, + keys=keys_to_store, + is_store=True, + sliding_window_block_ids=[block_id], + ) + + group_sizes = [0] * num_groups + group_sizes[group_idx] = 1 + block_indices = [0] * num_groups + block_indices[group_idx] = ( + boundary_tokens // group_config.gpu_block_size - 1 + ) + src_spec = GPULoadStoreSpec( + [block_id], + group_sizes=group_sizes, + block_indices=block_indices, + ) + store_jobs[job_id] = TransferJob( + req_id=req_id, + transfer_spec=(src_spec, store_output.store_spec), + ) + + return store_jobs + def build_connector_meta( self, scheduler_output: SchedulerOutput ) -> KVConnectorMetadata: @@ -891,9 +1030,11 @@ def build_connector_meta( ): self._current_batch_jobs_to_flush.update(self._jobs.keys()) + boundary_store_jobs = self._build_boundary_state_store_jobs(scheduler_output) + normal_store_jobs = self._build_store_jobs(scheduler_output) meta = OffloadingConnectorMetadata( load_jobs=self._current_batch_load_jobs, - store_jobs=self._build_store_jobs(scheduler_output), + store_jobs=boundary_store_jobs | normal_store_jobs, jobs_to_flush=self._current_batch_jobs_to_flush, ) self._current_batch_load_jobs = {} diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py index 20888c71f8..e1be60bcd4 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py @@ -11,6 +11,7 @@ KVConnectorBase_V1, KVConnectorRole, SupportsHMA, + SupportsVmmSafeTransfers, ) from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( @@ -43,7 +44,15 @@ from vllm.v1.request import Request -class OffloadingConnector(KVConnectorBase_V1, SupportsHMA): +class OffloadingConnector(KVConnectorBase_V1, SupportsHMA, SupportsVmmSafeTransfers): + """Native CPU KV offload connector. + + Transfers resolve the current tensor virtual address for every operation + and use CUDA memcpy. The connector does not register or export GPU KV + memory through RDMA, GDS, or CUDA IPC, so it is safe under CUDA VMM + remapping. + """ + @property def prefer_cross_layer_blocks(self) -> bool: return True diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index a567922199..8449f958cb 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -182,6 +182,33 @@ def sm70_tp2_all_reduce_gemma_rms_norm_fake( return torch.empty_like(tensor), torch.empty_like(residual, dtype=torch.float32) +def sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather( + tensor: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + epsilon: float, + group_name: str, +) -> tuple[torch.Tensor, torch.Tensor]: + assert group_name in _groups, f"Group {group_name} is not found." + group = _groups[group_name]() + if group is None: + raise ValueError(f"Group {group_name} is destroyed.") + return group._sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather_out_place( + tensor, residual, weight, epsilon + ) + + +def sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather_fake( + tensor: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + epsilon: float, + group_name: str, +) -> tuple[torch.Tensor, torch.Tensor]: + del weight, epsilon, group_name + return torch.empty_like(tensor), torch.empty_like(residual, dtype=torch.float32) + + def sm70_awq_mlp_down_tile_all_reduce( tensor: torch.Tensor, group_name: str ) -> torch.Tensor: @@ -371,6 +398,12 @@ def patched_fused_scaled_matmul_reduce_scatter( fake_impl=sm70_tp2_all_reduce_gemma_rms_norm_fake, ) +direct_register_custom_op( + op_name="sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather", + op_func=sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather, + fake_impl=sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather_fake, +) + direct_register_custom_op( op_name="sm70_awq_mlp_down_tile_all_reduce", op_func=sm70_awq_mlp_down_tile_all_reduce, @@ -695,9 +728,27 @@ def _sm70_tp2_all_reduce_gemma_rms_norm_out_place( raise RuntimeError("Device communicator lacks SM70 fused AR RMSNorm") return fused_op(input_, residual, weight, epsilon) - def sm70_awq_mlp_down_tile_all_reduce( - self, input_: torch.Tensor - ) -> torch.Tensor: + def _sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather_out_place( + self, + input_: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + epsilon: float, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.device_communicator is None: + raise ValueError("No device communicator found") + fused_op = getattr( + self.device_communicator, + "sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather", + None, + ) + if fused_op is None: + raise RuntimeError( + "Device communicator lacks SM70 TP4 fused collective-norm" + ) + return fused_op(input_, residual, weight, epsilon) + + def sm70_awq_mlp_down_tile_all_reduce(self, input_: torch.Tensor) -> torch.Tensor: if self.world_size == 1: return input_ @@ -1172,6 +1223,68 @@ def isend_tensor_dict( return handles + def isend_tensor_dict_static( + self, + tensor_dict: dict[str, torch.Tensor], + dst: int | None = None, + ) -> list[Handle]: + """Send tensors without metadata when both peers know the schema. + + The receiver must provide the same ordered keys, shapes, dtypes, and + devices to :meth:`irecv_tensor_dict_static`. Callers must fall back to + :meth:`isend_tensor_dict` whenever that contract can change. + """ + if self.world_size <= 1: + return [] + if self.use_cpu_custom_send_recv: + raise RuntimeError("static tensor transfer requires a process group") + + if dst is None: + dst = (self.rank_in_group + 1) % self.world_size + assert 0 <= dst < self.world_size, f"Invalid dst rank ({dst})" + + handles: list[Handle] = [] + for tensor in tensor_dict.values(): + if not isinstance(tensor, torch.Tensor): + raise TypeError("static tensor transfer only accepts tensors") + if tensor.numel() == 0: + continue + comm_group = self.cpu_group if tensor.is_cpu else self.device_group + handle = torch.distributed.isend( + tensor, dst=self.ranks[dst], group=comm_group + ) + if tensor.is_cuda: + tensor.record_stream(torch.cuda.current_stream(tensor.device)) + handles.append(handle) + return handles + + def irecv_tensor_dict_static( + self, + tensor_dict: dict[str, torch.Tensor], + src: int | None = None, + ) -> list[Handle]: + """Receive metadata-free tensors into caller-owned buffers.""" + if self.world_size <= 1: + return [] + if self.use_cpu_custom_send_recv: + raise RuntimeError("static tensor transfer requires a process group") + + if src is None: + src = (self.rank_in_group - 1) % self.world_size + assert 0 <= src < self.world_size, f"Invalid src rank ({src})" + + handles: list[Handle] = [] + for tensor in tensor_dict.values(): + if not isinstance(tensor, torch.Tensor): + raise TypeError("static tensor transfer only accepts tensors") + if tensor.numel() == 0: + continue + comm_group = self.cpu_group if tensor.is_cpu else self.device_group + handles.append( + torch.distributed.irecv(tensor, src=self.ranks[src], group=comm_group) + ) + return handles + def recv_tensor_dict( self, src: int | None = None, diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 121d2a30f8..31850497d3 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -93,7 +93,11 @@ ) from vllm.config.scheduler import SchedulerPolicy from vllm.config.utils import get_field -from vllm.config.vllm import OptimizationLevel, PerformanceMode +from vllm.config.vllm import ( + OptimizationLevel, + PerformanceMode, + _sm70_mtp_cudagraph_capture_sizes, +) from vllm.logger import init_logger, suppress_logging from vllm.platforms import CpuArchEnum, current_platform from vllm.plugins import load_general_plugins @@ -1765,10 +1769,11 @@ def _maybe_apply_sm70_mtp_defaults( has_linear_attention = any( layer_type == "linear_attention" for layer_type in layer_types ) - if is_server and self.enable_prefix_caching is None and has_linear_attention: self.enable_prefix_caching = True - profile_updates.append("enable_prefix_caching=True") + profile_updates.append( + f"enable_prefix_caching={self.enable_prefix_caching}" + ) if ( self.enable_prefix_caching and has_linear_attention @@ -1814,12 +1819,9 @@ def _maybe_apply_sm70_mtp_defaults( return if "draft_sample_method" not in self.speculative_config: - # For SM70 native MTP, official Qwen sampling is non-greedy - # (temperature=1.0/top_p=0.95/top_k=20). Probabilistic draft - # sampling provides draft_probs for standard rejection sampling and - # is the validated high-acceptance path. Greedy requests still take - # the argmax fast path inside the proposer when sampling metadata is - # all-greedy, so this default does not penalize deterministic runs. + # 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" self.speculative_config["draft_sample_method"] = draft_sample_method profile_updates.append( @@ -1860,11 +1862,6 @@ def _maybe_apply_sm70_mtp_defaults( and self.compilation_config.cudagraph_capture_sizes is None and self.compilation_config.max_cudagraph_capture_size is None ): - cudagraph_capture_sizes = ( - [1, 2, 4, 8, 9, 18] - if self.tensor_parallel_size >= 4 - else [1, 2, 4, 8, 9] - ) num_speculative_tokens = int( self.speculative_config.get("num_speculative_tokens") or 0 ) @@ -1881,9 +1878,22 @@ def _maybe_apply_sm70_mtp_defaults( num_speculative_state_tokens = max( num_speculative_tokens, ddtree_budget ) - if num_speculative_state_tokens > 0: - decode_query_len = num_speculative_state_tokens + 1 - max_num_seqs = max(int(self.max_num_seqs or 1), 1) + decode_query_len = num_speculative_state_tokens + 1 + max_num_seqs = max(int(self.max_num_seqs or 1), 1) + if envs.VLLM_SM70_MTP_SPLIT_DRAFT_CUDAGRAPHS and spec_method == "mtp": + cudagraph_capture_sizes = _sm70_mtp_cudagraph_capture_sizes( + max_num_seqs, + decode_query_len, + ) + profile_updates.append( + f"mtp_split_verifier_cudagraph_shapes={cudagraph_capture_sizes}" + ) + else: + cudagraph_capture_sizes = ( + [1, 2, 4, 8, 9, 18] + if self.tensor_parallel_size >= 4 + else [1, 2, 4, 8, 9] + ) cudagraph_capture_sizes = sorted( set(cudagraph_capture_sizes) | { diff --git a/vllm/envs.py b/vllm/envs.py index 499eb2f577..2d3e04aa45 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -136,6 +136,9 @@ VLLM_SM70_AWQ_TP2_FAST_SELECTOR: bool = True VLLM_SM70_AWQ_TP2_FAST_TARGETS: str | None = None VLLM_SM70_TP2_AR_GEMMA_RMS_FUSION: bool = False + VLLM_SM70_TP4_LONG_PREFILL_FUSED_NORM: bool = False + VLLM_SM70_TP4_LONG_FUSED_NORM_THREADS: int = 512 + VLLM_SM70_TP4_LONG_FUSED_NORM_BLOCKS: int = 80 VLLM_SM70_AWQ_MLP_ENGINE: bool = False VLLM_SM70_AWQ_PREFILL_EXACT_DENSE: bool = True VLLM_SM70_AWQ_MLP_DOWN_TILE_AR: bool = False @@ -153,11 +156,23 @@ VLLM_SM70_FP8_COORDINATED_TUNING: bool = True VLLM_SM70_FP8_REUSE_IMPORTED_CACHE: bool = False VLLM_SM70_FP8_SAFE_FAST_SELECTOR: bool = False + 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_FP8_QPN8_LIBRARY: str | None = None + VLLM_SM70_SAMPLER_LIBRARY: str | None = None + VLLM_SM70_FP8_PREFILL_VISIBLE_DENSE_MM: bool = False + VLLM_SM70_NVFP4_QPN2: bool = False 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_AWQ_REUSE_IMPORTED_CACHE: bool = False VLLM_SM70_AWQ_WARMUP: bool = True VLLM_SM70_AWQ_WARMUP_MAX_M: int = 16 @@ -169,22 +184,50 @@ 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 VLLM_SM70_AWQ_MOE_TUNE_MAX_TOKENS: int = 128 + VLLM_SM70_NVFP4_MOE_TUNE_MAX_TOKENS: int = 128 VLLM_SM70_ENABLE_DENSE_F16_FASTPATH: bool = False VLLM_SM70_ENABLE_LM_HEAD_FASTPATH: bool = False VLLM_SM70_LM_HEAD_TOP1: bool = True VLLM_SM70_LM_HEAD_TOP1_TC: bool = False + 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_VERIFY_FASTPATH: bool = False + VLLM_SM70_DFLASH2_FUSED_GDN_METADATA: bool = False + VLLM_SM70_DFLASH2_GDN_METADATA_SHADOW: bool = False + VLLM_SM70_DFLASH2_GDN_SYNC_ASSERT: bool = False + VLLM_SM70_DFLASH2_FUSED_GDN_VERIFY: bool = False + VLLM_SM70_DFLASH2_FUSED_GDN_NORM: bool = False + VLLM_SM70_DFLASH2_FUSED_GDN_SPLIT: bool = False + VLLM_SM70_DFLASH2_FUSED_SMALLQ_METADATA: bool = False + VLLM_SM70_DFLASH2_GROUPED_SMALLQ_METADATA: bool = False + VLLM_SM70_DFLASH2_FUSED_QKV_PACK: bool = False + VLLM_SM70_DFLASH2_FUSED_GEMMA_RMS: bool = False + 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_TOP1_CUSTOM_AR: bool = False VLLM_SM70_GREEDY_TOKEN_FASTPATH: bool = True VLLM_SM70_GREEDY_TOKEN_FASTPATH_TRACE: bool = False + VLLM_SM70_COMPACT_TOPK20_SAMPLER: bool = False + VLLM_SM70_CHUNKED_TOPK20_CHUNKS: int = 0 + VLLM_SM70_TP_LOCAL_TOPK20_SAMPLER: bool = False + VLLM_SM70_TOPK_TOPP_8_WARPS: bool = False + VLLM_SM70_TOPK_TOPP_B8_B16_8_WARPS: bool = True VLLM_SM70_ASYNC_SCHEDULING_QUEUE_DEPTH: int = 0 VLLM_SM70_ASYNC_STAGED_INPUT_PREP: bool = False VLLM_SM70_ASYNC_CPU_TRACE: bool = False VLLM_SM70_ASYNC_CPU_TRACE_EVERY: int = 16 VLLM_TP_ALLREDUCE_TRACE: bool = False VLLM_CUSTOM_ALLREDUCE_BLOCK_LIMIT: int | None = None + VLLM_SM70_TP4_MTP_AR_BLOCK_TUNING: bool = False VLLM_SM70_TP4_M5_AR_THREADS: int | None = None + VLLM_SM70_TP4_SMALL_AR_PACK32: bool = False VLLM_SM70_F16_DENSE_ALLOWLIST: str | None = None VLLM_SM70_MOE_DENSE_ALLOWLIST: str | None = None VLLM_SM70_F16_DENSE_MAX_M: int = 64 @@ -198,13 +241,21 @@ VLLM_SM70_FP8_TURBOMIND: bool = True VLLM_SM70_FP8_DENSE_GATED_SILU: bool = True VLLM_SM70_NVFP4_TURBOMIND: bool = True + VLLM_SM70_NVFP4_MOE_GROUPED_PREFILL: bool = True + VLLM_SM70_NVFP4_DENSE_GATED_SILU: bool = True + VLLM_SM70_NVFP4_QPN4: bool = True + VLLM_SM70_NVFP4_QPN4_DOWN_SCALE_CODE: bool = False 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_TOP6_DECODE: bool = True + VLLM_SM70_MXFP4_MOE_DIRECT_ORDER_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 VLLM_SM70_DSV4_SPARSE_MLA_SPLITK_C128: bool = False @@ -237,8 +288,12 @@ VLLM_SM70_DECODE_EVENT_TRACE_EVERY: int = 16 VLLM_SM70_MTP_PROFILE: bool = False VLLM_SM70_MTP_PROFILE_INTERVAL: int = 16 + VLLM_SM70_MTP_SPLIT_DRAFT_CUDAGRAPHS: bool = False + VLLM_SM70_MTP_CONCURRENCY_WARMUP: bool = False VLLM_SM70_MTP_CONTEXT_BUCKETS: str | None = None VLLM_SM70_DSV4_DECODE_CONTEXT_BUCKETS: str | None = None + VLLM_SM70_DSV4_PRIVATE_COMPRESSOR_STATE: bool = False + VLLM_SM70_FP8_KV_DECODE_CONTEXT_BUCKETS: str | None = None VLLM_SM70_MTP_CONTEXT_BUCKET_PARTITION_SIZE: str | None = None VLLM_SM70_REJECTION_PROFILE: bool = False VLLM_SM70_REJECTION_PROFILE_INTERVAL: int = 20 @@ -298,6 +353,8 @@ VLLM_FLASH_V100_PREFILL_GATHER_DENSE_MIN_KV: int = 8192 VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3: bool = True VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3_MIN_KV: int = 32768 + VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3_Q8000_EXPERIMENTAL: bool = False + VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL: bool = True VLLM_FLASH_V100_PREFILL_SPLIT_KV: bool = False VLLM_FLASH_V100_PREFILL_SPLIT_KV_TOKENS: int = 32768 VLLM_FLASH_V100_PREFILL_SPLIT_KV_MIN_Q: int = 1 @@ -335,6 +392,8 @@ 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_MIN_MODEL_LEN: int = 32768 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 @@ -344,11 +403,20 @@ VLLM_FLASH_V100_XQA_G6_QK_PIPELINE: bool = True VLLM_FLASH_V100_XQA_G6_QK_PIPELINE_WARPS: int = 8 VLLM_FLASH_V100_XQA_G6_QK_PIPELINE_TRACE: bool = False + VLLM_FLASH_V100_XQA_G6_DUAL_CTA: bool = False + VLLM_FLASH_V100_XQA_SPLIT_REDUCE: bool = False + VLLM_FLASH_V100_E4M3_BATCH_XQA: bool = True + VLLM_FLASH_V100_E4M3_BATCH_XQA_OPTIMIZED: bool = True + VLLM_FLASH_V100_E4M3_PAGE800_FASTPATH: bool = True + VLLM_FLASH_V100_E4M3_PAGE800_FASTPATH_TRACE: bool = False + VLLM_FLASH_V100_XQA_BATCH_CONTEXT_ROUTING: bool = True + VLLM_FLASH_V100_XQA_BATCH_CONTEXT_ROUTING_TRACE: bool = False VLLM_FLASH_V100_XQA_E5M2_G6_DUAL_CTA: bool = True VLLM_FLASH_V100_XQA_E5M2_G6_SPLIT_REDUCE: bool = True VLLM_FLASH_V100_XQA_E5M2_P1024_BEGIN: int = 61633 VLLM_FLASH_V100_XQA_E5M2_PARTITION_PAGE_IDS: bool = True VLLM_FLASH_V100_XQA_E5M2_PAIR_LOAD: bool = True + VLLM_FLASH_V100_XQA_E5M2_BATCH_WIDE_LOAD: bool = True VLLM_FLASH_V100_XQA_E5M2_G6_DUAL_CTA_TRACE: bool = False VLLM_FLASH_V100_TRACE_DECODE_ACTIVE: bool = False VLLM_FLASH_V100_DECODE_USE_SCALAR_PAGED: bool = True @@ -443,6 +511,8 @@ VLLM_SM70_QWEN_GDN_DISABLE_INPUT_CORE_OP: bool = False VLLM_SM70_QWEN_GDN_INPUT_PROJECTION_OP: bool = False VLLM_SM70_QWEN_GDN_OUTPUT_PROJECTION_OP: bool = False + VLLM_SM70_GDN_QPN8_BA_SPLIT: bool = False + VLLM_SM70_GDN_RMSNORM_ONEPASS: bool = False VLLM_SM70_GEMMA_RMS_NORM_EAGER: bool = False VLLM_SM70_GEMMA_RMS_NORM_COMPILE_NATIVE: bool = False VLLM_SM70_GEMMA_LONG_PREFILL_FUSED: bool = True @@ -460,6 +530,7 @@ VLLM_QWEN3NEXT_ENABLE_SHARED_MOE_OVERLAP: bool = False VLLM_SM70_DISABLE_QWEN3NEXT_SHARED_MOE_OVERLAP: bool = False VLLM_SM70_UNQUANTIZED_MOE_0DOT3_CONFIG: bool = True + VLLM_SM70_MTP_MOE_TUNED_CONFIG: bool = False VLLM_SM70_DENSE_CUDAGRAPH_CAPTURE: bool = False VLLM_SM70_USE_BREAKABLE_CUDAGRAPH: bool = False VLLM_SM70_FLASH_V100_0DOT3_COMPILE_GRAPH: bool = False @@ -1584,6 +1655,17 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_TP2_AR_GEMMA_RMS_FUSION": lambda: bool( int(os.getenv("VLLM_SM70_TP2_AR_GEMMA_RMS_FUSION", "0")) ), + # Default-off exact-8K TP4 experiment: fuse the row-parallel reduction, + # mixed-FP32-residual Gemma RMSNorm, and normalized-output all-gather. + "VLLM_SM70_TP4_LONG_PREFILL_FUSED_NORM": lambda: bool( + int(os.getenv("VLLM_SM70_TP4_LONG_PREFILL_FUSED_NORM", "0")) + ), + "VLLM_SM70_TP4_LONG_FUSED_NORM_THREADS": lambda: int( + os.getenv("VLLM_SM70_TP4_LONG_FUSED_NORM_THREADS", "512") + ), + "VLLM_SM70_TP4_LONG_FUSED_NORM_BLOCKS": lambda: int( + os.getenv("VLLM_SM70_TP4_LONG_FUSED_NORM_BLOCKS", "80") + ), # Experimental TileRT-inspired dense MLP lane for SM70 AWQ decode. The # first stage fuses gate_up_proj + SiluAndMul through the TurboMind GEMM # epilogue for M=1/TP2 while keeping the existing down_proj/reduce path. @@ -1601,6 +1683,31 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_FP8_PREFILL_EXACT_DENSE": lambda: bool( int(os.getenv("VLLM_SM70_FP8_PREFILL_EXACT_DENSE", "1")) ), + # Memory-neutral QPN8 layout for shape- and runtime-gated TP4 block-FP8 + # dense projections. Pure-FP8 checkpoints must opt in; + # 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_FP8_QPN8_PP2_TP4": lambda: bool( + int(os.getenv("VLLM_SM70_FP8_QPN8_PP2_TP4", "1")) + ), + # 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_FP8_PREFILL_CUTLASS": lambda: bool( + int(os.getenv("VLLM_SM70_FP8_PREFILL_CUTLASS", "1")) + ), + "VLLM_SM70_FP8_PREFILL_VISIBLE_DENSE_MM": lambda: bool( + int(os.getenv("VLLM_SM70_FP8_PREFILL_VISIBLE_DENSE_MM", "0")) + ), + # 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"))), # 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. @@ -1652,6 +1759,18 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_FP8_SAFE_FAST_SELECTOR": lambda: bool( int(os.getenv("VLLM_SM70_FP8_SAFE_FAST_SELECTOR", "0")) ), + "VLLM_SM70_FP8_GROUPED_BMM_DECODE": lambda: bool( + int(os.getenv("VLLM_SM70_FP8_GROUPED_BMM_DECODE", "1")) + ), + "VLLM_SM70_FP8_PREFILL_FAST_SELECTOR": lambda: bool( + int(os.getenv("VLLM_SM70_FP8_PREFILL_FAST_SELECTOR", "1")) + ), + "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")) ), @@ -1664,6 +1783,9 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_NVFP4_TUNE_SMALL_SHAPES": lambda: bool( int(os.getenv("VLLM_SM70_NVFP4_TUNE_SMALL_SHAPES", "1")) ), + "VLLM_SM70_NVFP4_QWEN38_TP4_M1_FAST_SELECTOR": lambda: bool( + int(os.getenv("VLLM_SM70_NVFP4_QWEN38_TP4_M1_FAST_SELECTOR", "1")) + ), "VLLM_SM70_AWQ_REUSE_IMPORTED_CACHE": lambda: bool( int(os.getenv("VLLM_SM70_AWQ_REUSE_IMPORTED_CACHE", "0")) ), @@ -1696,12 +1818,26 @@ 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")) ), + "VLLM_SM70_DSV4_QNORM_KV_FUSED_TP4": lambda: bool( + int(os.getenv("VLLM_SM70_DSV4_QNORM_KV_FUSED_TP4", "1")) + ), + # Skip PP metadata and TP reconstruction only for the exact, replicated + # SM70 B1 hidden-state schema validated by the worker on both stages. + "VLLM_SM70_PP_STATIC_HIDDEN_TRANSFER": lambda: bool( + int(os.getenv("VLLM_SM70_PP_STATIC_HIDDEN_TRANSFER", "1")) + ), "VLLM_SM70_AWQ_MOE_TUNE_MAX_TOKENS": lambda: int( os.getenv("VLLM_SM70_AWQ_MOE_TUNE_MAX_TOKENS", "128") ), + "VLLM_SM70_NVFP4_MOE_TUNE_MAX_TOKENS": lambda: int( + os.getenv("VLLM_SM70_NVFP4_MOE_TUNE_MAX_TOKENS", "128") + ), # Experimental unquantized FP16 SM70 TurboMind fast paths. Keep default-off: # these must pass the numeric policy and model-level token gate before # becoming part of the default V100 route. @@ -1737,6 +1873,110 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_LM_HEAD_TOP1_TC": lambda: bool( int(os.getenv("VLLM_SM70_LM_HEAD_TOP1_TC", "0")) ), + # Candidate-only QPN8 LM head. QPN8 selects a conservative local top-64 + # support, then directly re-evaluates the original FP16 rows and restores + # dense-vocabulary top-k tie order. Keep default-off until real-hidden, + # quality, and complete Graph/Nsight gates all pass. + "VLLM_SM70_DFLASH2_QPN8_RERANK": lambda: bool( + int(os.getenv("VLLM_SM70_DFLASH2_QPN8_RERANK", "0")) + ), + # Audit-only eager mode: execute QPN8+rerank, compare it with the dense + # local top-k, and return the dense result so the baseline trajectory is + # unchanged. This intentionally synchronizes for diagnostics. + "VLLM_SM70_DFLASH2_QPN8_RERANK_SHADOW": lambda: bool( + int(os.getenv("VLLM_SM70_DFLASH2_QPN8_RERANK_SHADOW", "0")) + ), + # Preserve the production full-vocabulary torch.topk tie contract by + # default. The candidate-order experiment avoids a host-blocking SM70 + # multi-block top-k, but may be enabled only for paired quality tests. + "VLLM_SM70_DFLASH2_QPN8_DENSE_ORDER": lambda: bool( + int(os.getenv("VLLM_SM70_DFLASH2_QPN8_DENSE_ORDER", "1")) + ), + # 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( + int(os.getenv("VLLM_SM70_DFLASH2_VERIFY_FASTPATH", "0")) + ), + # Build all selector-based DFlash target GDN state-index metadata with one + # pointer-table Triton launch. Keep separate from the shared-classification + # gate until the fixed-trajectory and mixed-batch Graph checks pass. + "VLLM_SM70_DFLASH2_FUSED_GDN_METADATA": lambda: bool( + int(os.getenv("VLLM_SM70_DFLASH2_FUSED_GDN_METADATA", "0")) + ), + # Debug-only oracle: materialize the legacy advanced-indexing contract and + # compare it with the fused persistent buffers before graph replay. + "VLLM_SM70_DFLASH2_GDN_METADATA_SHADOW": lambda: bool( + int(os.getenv("VLLM_SM70_DFLASH2_GDN_METADATA_SHADOW", "0")) + ), + # Debug-only device-value assertion. The metadata producers and FULL graph + # replay share the current CUDA stream, so production ordering does not + # require a host-blocking Tensor.item(). The no-sync route preserves the + # accepted token/acceptance trajectory; opt in only when auditing metadata. + "VLLM_SM70_DFLASH2_GDN_SYNC_ASSERT": lambda: bool( + int(os.getenv("VLLM_SM70_DFLASH2_GDN_SYNC_ASSERT", "0")) + ), + # Selector-based DFlash packed GDN verification kernel. This is deliberately + # independent from the shared-metadata umbrella gate so each optimization + # can be paired against the unchanged verifier in isolation. + "VLLM_SM70_DFLASH2_FUSED_GDN_VERIFY": lambda: bool( + int(os.getenv("VLLM_SM70_DFLASH2_FUSED_GDN_VERIFY", "0")) + ), + # Route compatible target GDN output gates through the existing one-pass + # CUDA RMSNormGated implementation. This remains an explicit opt-in. + "VLLM_SM70_DFLASH2_FUSED_GDN_NORM": lambda: bool( + int(os.getenv("VLLM_SM70_DFLASH2_FUSED_GDN_NORM", "0")) + ), + # Fuse compatible nonzero-offset GDN z/b/a materialization into one + # copy kernel. This must stay separate from the plain-view path because + # nonzero-offset views are unsafe under the SM70 compile/full-graph route. + "VLLM_SM70_DFLASH2_FUSED_GDN_SPLIT": lambda: bool( + int(os.getenv("VLLM_SM70_DFLASH2_FUSED_GDN_SPLIT", "0")) + ), + # Build Flash-V100 small-query verifier rows directly in their persistent + # graph buffers. This replaces four repeat_interleave scans per KV group. + # The matched TP4 trace is token/acceptance exact and cuts the synchronized + # DFlash2 draft-to-target interval from 5.720 ms to 1.911 ms on V100. + "VLLM_SM70_DFLASH2_FUSED_SMALLQ_METADATA": lambda: bool( + int(os.getenv("VLLM_SM70_DFLASH2_FUSED_SMALLQ_METADATA", "0")) + ), + # Collapse five compatible target small-query metadata launches into one + # heterogeneous-width pointer-table kernel. The paired + # TP4 node trace is token/acceptance exact and reduces synchronized D2T. + "VLLM_SM70_DFLASH2_GROUPED_SMALLQ_METADATA": lambda: bool( + int(os.getenv("VLLM_SM70_DFLASH2_GROUPED_SMALLQ_METADATA", "0")) + ), + # Copy the post-convolution Q/K/V row slices into the recurrent kernel's + # packed contiguous layout with one bitwise Triton launch. This replaces + # three reshape copies plus one cat in the exact block-eight verifier. + "VLLM_SM70_DFLASH2_FUSED_QKV_PACK": lambda: bool( + int(os.getenv("VLLM_SM70_DFLASH2_FUSED_QKV_PACK", "0")) + ), + # Fuse the FP16 projection + FP32 residual + Gemma RMSNorm suffix used by + # small DFlash2 verifier graphs. Default-off pending numeric/quality gates. + "VLLM_SM70_DFLASH2_FUSED_GEMMA_RMS": lambda: bool( + int(os.getenv("VLLM_SM70_DFLASH2_FUSED_GEMMA_RMS", "0")) + ), + # Avoid materializing/gathering full-vocabulary target logits when the + # DFlash2 proposal and target sampling distributions both have compact + # top-k support. Default-off until paired output/acceptance and end-to-end + # V100 gates pass; unsupported sampling features fall back to dense logits. + "VLLM_SM70_DFLASH2_SPARSE_TARGET_REJECTION": lambda: bool( + int(os.getenv("VLLM_SM70_DFLASH2_SPARSE_TARGET_REJECTION", "0")) + ), + # Compute the compatible 25600->5120 target-hidden projection as four output + # shards, then all-gather only the 80-KiB block-eight result. This + # remains opt-in until the acceptance/quality gate follows the positive + # production-weight TP4 microbenchmark. + "VLLM_SM70_DFLASH2_SHARDED_CONTEXT_FC": lambda: bool( + int(os.getenv("VLLM_SM70_DFLASH2_SHARDED_CONTEXT_FC", "0")) + ), + # Default-on SGLang-style push collective for the validated FP16 80-KiB + # verifier and 8-KiB decode payloads on fully-connected SM70 TP4 CUDA + # Graphs. Other devices, topologies, sizes, and eager calls retain the + # ordinary pull path; explicit 0 is the rollback. + "VLLM_SM70_TP4_PUSH_ALLREDUCE": lambda: bool( + int(os.getenv("VLLM_SM70_TP4_PUSH_ALLREDUCE", "1")) + ), # 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 @@ -1748,6 +1988,16 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_GREEDY_TOKEN_FASTPATH_TRACE": lambda: bool( int(os.getenv("VLLM_SM70_GREEDY_TOKEN_FASTPATH_TRACE", "0")) ), + # Opt-in V100 launch for the exact validated combined top-k/top-p shapes. + "VLLM_SM70_TOPK_TOPP_8_WARPS": lambda: bool( + int(os.getenv("VLLM_SM70_TOPK_TOPP_8_WARPS", "0")) + ), + # The exact B8/B16, 248320-column sampler contract uses one logits row per + # request. Eight warps preserves Qrita's masking math while improving SM70 + # reduction parallelism. Set to 0 to restore Triton's launch heuristic. + "VLLM_SM70_TOPK_TOPP_B8_B16_8_WARPS": lambda: bool( + int(os.getenv("VLLM_SM70_TOPK_TOPP_B8_B16_8_WARPS", "1")) + ), # Diagnostic SM70 async scheduling depth override. Default 0 preserves # upstream behavior. Values >2 let no-PP async scheduling enqueue more real # decode steps before collecting CPU-visible outputs; this changes queueing @@ -1775,11 +2025,18 @@ def _resolve_rust_frontend_path() -> str | None: if "VLLM_CUSTOM_ALLREDUCE_BLOCK_LIMIT" in os.environ else None ), + # Opt in exact MTP4 verifier payloads on fully-connected SM70 TP4. + "VLLM_SM70_TP4_MTP_AR_BLOCK_TUNING": lambda: bool( + int(os.getenv("VLLM_SM70_TP4_MTP_AR_BLOCK_TUNING", "0")) + ), "VLLM_SM70_TP4_M5_AR_THREADS": lambda: ( int(os.environ["VLLM_SM70_TP4_M5_AR_THREADS"]) if "VLLM_SM70_TP4_M5_AR_THREADS" in os.environ else None ), + "VLLM_SM70_TP4_SMALL_AR_PACK32": lambda: bool( + int(os.getenv("VLLM_SM70_TP4_SMALL_AR_PACK32", "0")) + ), # Optional custom allreduce for the tiny per-rank top1 pair. Keep default # off until the communicator path has same-criterion model evidence. "VLLM_SM70_TOP1_CUSTOM_AR": lambda: bool( @@ -1832,6 +2089,34 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_NVFP4_TURBOMIND": lambda: bool( int(os.getenv("VLLM_SM70_NVFP4_TURBOMIND", "1")) ), + # Dispatch all 256 routed experts in one grouped TurboMind call for the + # exact Qwen3.6-35B-A3B TP1/2/4 NVFP4 prefill shapes. B1-B8 decode keeps + # the compact active-expert route; larger graph shapes use full groups. + "VLLM_SM70_NVFP4_MOE_GROUPED_PREFILL": lambda: bool( + int(os.getenv("VLLM_SM70_NVFP4_MOE_GROUPED_PREFILL", "1")) + ), + "VLLM_SM70_NVFP4_DENSE_GATED_SILU": lambda: bool( + int(os.getenv("VLLM_SM70_NVFP4_DENSE_GATED_SILU", "1")) + ), + # Memory-neutral native QPN4 M=1 decode for the accepted Qwen3.8-27B + # NVFP4 TP4 no-MTP contract. Larger-M prefill dequantizes into one bounded + # shared FP16 workspace. Set to 0 to retain the TurboMind NVFP4 route. + "VLLM_SM70_NVFP4_QPN4": lambda: bool(int(os.getenv("VLLM_SM70_NVFP4_QPN4", "1"))), + "VLLM_SM70_NVFP4_QPN4_DOWN_SCALE_CODE": lambda: bool( + int(os.getenv("VLLM_SM70_NVFP4_QPN4_DOWN_SCALE_CODE", "0")) + ), + # Exact no-MTP random sampler for the Qwen3.8 TP4 batch-one contract. + # The base flag enables exact compact sampling after full-logits gather. + # The separate TP-local flag opts into rank-local top-20 reduction. + "VLLM_SM70_COMPACT_TOPK20_SAMPLER": lambda: bool( + int(os.getenv("VLLM_SM70_COMPACT_TOPK20_SAMPLER", "0")) + ), + "VLLM_SM70_CHUNKED_TOPK20_CHUNKS": lambda: int( + os.getenv("VLLM_SM70_CHUNKED_TOPK20_CHUNKS", "0") + ), + "VLLM_SM70_TP_LOCAL_TOPK20_SAMPLER": lambda: bool( + int(os.getenv("VLLM_SM70_TP_LOCAL_TOPK20_SAMPLER", "0")) + ), "VLLM_SM70_MXFP4_TURBOMIND": lambda: bool( int(os.getenv("VLLM_SM70_MXFP4_TURBOMIND", "1")) ), @@ -1900,13 +2185,24 @@ 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. "VLLM_SM70_MXFP4_MOE_GROUPED_M8": lambda: bool( int(os.getenv("VLLM_SM70_MXFP4_MOE_GROUPED_M8", "0")) ), + # Extend the one-launch grouped verifier route from M=8 to M=2..M=8. + # Kept independent until each width passes its CUDA Graph and model gates. + "VLLM_SM70_MXFP4_MOE_GROUPED_VERIFIER": lambda: bool( + int(os.getenv("VLLM_SM70_MXFP4_MOE_GROUPED_VERIFIER", "0")) + ), + # Group verifier slots routed to the same expert into one multi-row group. + # This changes the MXFP4 reduction tactic, so it stays behind an + # independent acceptance/quality gate. + "VLLM_SM70_MXFP4_MOE_GROUPED_M8_EXPERT_ROWS": lambda: bool( + int(os.getenv("VLLM_SM70_MXFP4_MOE_GROUPED_M8_EXPERT_ROWS", "0")) + ), # Use deterministic W13/W2 tactics for the fixed 48 one-row # verifier groups instead of relying on capture-time autotune stability. "VLLM_SM70_MXFP4_MOE_GROUPED_M8_FAST_SELECTOR": lambda: bool( @@ -1915,7 +2211,17 @@ 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. 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", "1")) + ), + "VLLM_SM70_MXFP4_MOE_BROADCAST_INPUT_DECODE": lambda: bool( + int(os.getenv("VLLM_SM70_MXFP4_MOE_BROADCAST_INPUT_DECODE", "1")) ), # FP8 caller for the generic SM70 TurboMind active-source-group compact # decode path. The backend scheduler keeps source expert group semantics @@ -2005,10 +2311,26 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_MTP_PROFILE_INTERVAL": lambda: max( 1, int(os.getenv("VLLM_SM70_MTP_PROFILE_INTERVAL", "16")) ), + "VLLM_SM70_MTP_SPLIT_DRAFT_CUDAGRAPHS": lambda: bool( + int(os.getenv("VLLM_SM70_MTP_SPLIT_DRAFT_CUDAGRAPHS", "0")) + ), + # Compile alternate single/concurrent MTP helper signatures at startup. + # Default-off until matched cold-start and steady-state evidence is complete. + "VLLM_SM70_MTP_CONCURRENCY_WARMUP": lambda: bool( + int(os.getenv("VLLM_SM70_MTP_CONCURRENCY_WARMUP", "0")) + ), "VLLM_SM70_MTP_CONTEXT_BUCKETS": lambda: os.getenv("VLLM_SM70_MTP_CONTEXT_BUCKETS"), "VLLM_SM70_DSV4_DECODE_CONTEXT_BUCKETS": lambda: os.getenv( "VLLM_SM70_DSV4_DECODE_CONTEXT_BUCKETS" ), + # The private ring changes compressor-state ownership and lifetime. Keep it + # opt-in until its long-context quality and capacity gates are complete. + "VLLM_SM70_DSV4_PRIVATE_COMPRESSOR_STATE": lambda: bool( + int(os.getenv("VLLM_SM70_DSV4_PRIVATE_COMPRESSOR_STATE", "0")) + ), + "VLLM_SM70_FP8_KV_DECODE_CONTEXT_BUCKETS": lambda: os.getenv( + "VLLM_SM70_FP8_KV_DECODE_CONTEXT_BUCKETS" + ), "VLLM_SM70_MTP_CONTEXT_BUCKET_PARTITION_SIZE": lambda: os.getenv( "VLLM_SM70_MTP_CONTEXT_BUCKET_PARTITION_SIZE" ), @@ -2169,6 +2491,22 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3_MIN_KV": lambda: int( os.getenv("VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3_MIN_KV", "32768") ), + "VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3_Q8000_EXPERIMENTAL": lambda: bool( + int( + os.getenv( + "VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3_Q8000_EXPERIMENTAL", + "0", + ) + ) + ), + "VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL": lambda: bool( + int( + os.getenv( + "VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL", + "1", + ) + ) + ), "VLLM_FLASH_V100_PREFILL_SPLIT_KV": lambda: bool( int(os.getenv("VLLM_FLASH_V100_PREFILL_SPLIT_KV", "0")) ), @@ -2291,6 +2629,12 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_FLASH_V100_DECODE_USE_XQA": lambda: bool( 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")) + ), + "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_DECODE_XQA_Q4_MIN_SEQ_LEN": lambda: int( os.getenv("VLLM_FLASH_V100_DECODE_XQA_Q4_MIN_SEQ_LEN", "32768") ), @@ -2327,6 +2671,32 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_FLASH_V100_XQA_G6_QK_PIPELINE_TRACE": lambda: bool( int(os.getenv("VLLM_FLASH_V100_XQA_G6_QK_PIPELINE_TRACE", "0")) ), + "VLLM_FLASH_V100_XQA_G6_DUAL_CTA": lambda: bool( + int(os.getenv("VLLM_FLASH_V100_XQA_G6_DUAL_CTA", "0")) + ), + "VLLM_FLASH_V100_XQA_SPLIT_REDUCE": lambda: bool( + int(os.getenv("VLLM_FLASH_V100_XQA_SPLIT_REDUCE", "0")) + ), + # Exact SM70 G6/D256 E4M3 XQA for B2-B16. Set to 0 to restore the scalar + # paged decoder for matched control runs or emergency rollback. + "VLLM_FLASH_V100_E4M3_BATCH_XQA": lambda: bool( + int(os.getenv("VLLM_FLASH_V100_E4M3_BATCH_XQA", "1")) + ), + "VLLM_FLASH_V100_E4M3_BATCH_XQA_OPTIMIZED": lambda: bool( + int(os.getenv("VLLM_FLASH_V100_E4M3_BATCH_XQA_OPTIMIZED", "1")) + ), + "VLLM_FLASH_V100_E4M3_PAGE800_FASTPATH": lambda: bool( + int(os.getenv("VLLM_FLASH_V100_E4M3_PAGE800_FASTPATH", "1")) + ), + "VLLM_FLASH_V100_E4M3_PAGE800_FASTPATH_TRACE": lambda: bool( + int(os.getenv("VLLM_FLASH_V100_E4M3_PAGE800_FASTPATH_TRACE", "0")) + ), + "VLLM_FLASH_V100_XQA_BATCH_CONTEXT_ROUTING": lambda: bool( + int(os.getenv("VLLM_FLASH_V100_XQA_BATCH_CONTEXT_ROUTING", "1")) + ), + "VLLM_FLASH_V100_XQA_BATCH_CONTEXT_ROUTING_TRACE": lambda: bool( + int(os.getenv("VLLM_FLASH_V100_XQA_BATCH_CONTEXT_ROUTING_TRACE", "0")) + ), "VLLM_FLASH_V100_XQA_E5M2_G6_DUAL_CTA": lambda: bool( int(os.getenv("VLLM_FLASH_V100_XQA_E5M2_G6_DUAL_CTA", "1")) ), @@ -2342,6 +2712,9 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_FLASH_V100_XQA_E5M2_PAIR_LOAD": lambda: bool( int(os.getenv("VLLM_FLASH_V100_XQA_E5M2_PAIR_LOAD", "1")) ), + "VLLM_FLASH_V100_XQA_E5M2_BATCH_WIDE_LOAD": lambda: bool( + int(os.getenv("VLLM_FLASH_V100_XQA_E5M2_BATCH_WIDE_LOAD", "1")) + ), "VLLM_FLASH_V100_XQA_E5M2_G6_DUAL_CTA_TRACE": lambda: bool( int(os.getenv("VLLM_FLASH_V100_XQA_E5M2_G6_DUAL_CTA_TRACE", "0")) ), @@ -2675,6 +3048,14 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_QWEN_GDN_OUTPUT_PROJECTION_OP": lambda: bool( int(os.getenv("VLLM_SM70_QWEN_GDN_OUTPUT_PROJECTION_OP", "0")) ), + # Exact-shape SM70 GDN decode fusions. The measured workload is evidence, + # not a model-identity selector; runtime admission uses operator contracts. + "VLLM_SM70_GDN_QPN8_BA_SPLIT": lambda: bool( + int(os.getenv("VLLM_SM70_GDN_QPN8_BA_SPLIT", "0")) + ), + "VLLM_SM70_GDN_RMSNORM_ONEPASS": lambda: bool( + int(os.getenv("VLLM_SM70_GDN_RMSNORM_ONEPASS", "0")) + ), # Diagnostic-only: keep Qwen3.5/Gemma RMSNorm arithmetic behind an opaque # custom-op boundary under the SM70 compile/FULL graph lane. "VLLM_SM70_GEMMA_RMS_NORM_EAGER": lambda: bool( @@ -2748,6 +3129,11 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_UNQUANTIZED_MOE_0DOT3_CONFIG": lambda: bool( int(os.getenv("VLLM_SM70_UNQUANTIZED_MOE_0DOT3_CONFIG", "1")) ), + # Opt-in decode tiles for the exact E256/N128/K2048/top-k8 SM70 contract. + # Larger or unmatched token shapes retain the 0.0.3 config. + "VLLM_SM70_MTP_MOE_TUNED_CONFIG": lambda: bool( + int(os.getenv("VLLM_SM70_MTP_MOE_TUNED_CONFIG", "0")) + ), # Legacy SM70 CUDA-graph capture-size tuning from 0.0.3. Default-off # because dense capture can increase startup/compile cost; when enabled on # SM70 it restores the old fine-grained small-batch capture list. diff --git a/vllm/forward_context.py b/vllm/forward_context.py index bc42ef9104..154a85dd03 100644 --- a/vllm/forward_context.py +++ b/vllm/forward_context.py @@ -19,6 +19,9 @@ logger = init_logger(__name__) +CUDAGRAPH_VARIANT_DEFAULT = 0 +CUDAGRAPH_VARIANT_LONG_CONTEXT = -1 + track_batchsize: bool = envs.VLLM_LOG_BATCHSIZE_INTERVAL >= 0 last_logging_time: float = 0 forward_start_time: float = 0 @@ -56,11 +59,13 @@ class BatchDescriptor: (like fused_moe_lora) whose grid size depends on num_active_loras to be properly captured. """ - graph_variant: int = 0 + graph_variant: int = CUDAGRAPH_VARIANT_DEFAULT """ Optional semantic specialization for CUDA graphs that share the same batch shape but execute different model branches. The default keeps existing graph - keys unchanged. MTP drafter uses this to separate per-step layer selection. + keys unchanged. Negative values are reserved for target-model attention + specializations; MTP drafter uses non-negative values to separate per-step + layer selection. """ attention_context_bucket: int | None = None """ diff --git a/vllm/model_executor/kernels/mhc/tilelang.py b/vllm/model_executor/kernels/mhc/tilelang.py index de5f02d8ce..e591018d4f 100644 --- a/vllm/model_executor/kernels/mhc/tilelang.py +++ b/vllm/model_executor/kernels/mhc/tilelang.py @@ -626,12 +626,14 @@ def mhc_fused_post_pre_tilelang( and use_fp16 and capability is not None and capability.to_int() == 70 - and num_tokens == 1 + and 1 <= num_tokens <= 8 and hidden_size == 4096 and hc_mult == 4 and hc_mult3 == 24 - and tile_n == 2 - and n_splits == 8 + and ( + (num_tokens < 8 and tile_n == 2 and n_splits == 8) + or (num_tokens == 8 and tile_n == 3 and n_splits == 4) + ) ) gemm_out_mul = torch.empty( diff --git a/vllm/model_executor/kernels/mhc/tilelang_kernels.py b/vllm/model_executor/kernels/mhc/tilelang_kernels.py index 777804072f..bf4a94e556 100644 --- a/vllm/model_executor/kernels/mhc/tilelang_kernels.py +++ b/vllm/model_executor/kernels/mhc/tilelang_kernels.py @@ -539,7 +539,7 @@ def sm70_mhc_post_fp32_stage_tilelang( n_thr: int = 256, n_splits: int = 8, ) -> tilelang.JITKernel: - """Materialize the M=1 post mapping without rounding the dot input.""" + """Materialize a small-token post mapping without rounding the dot input.""" num_tokens = T.dynamic("num_tokens") h_per_split = hidden // n_splits h_iters = h_per_split // n_thr @@ -608,7 +608,7 @@ def sm70_mhc_dot_from_fp32_stage_tilelang( tile_n: int = 2, n_splits: int = 8, ) -> tilelang.JITKernel: - """Run the M=1 dot accumulation from the exact FP32 post stage.""" + """Run small-token dot accumulation from the exact FP32 post stage.""" num_tokens = T.dynamic("num_tokens") h_per_split = hidden // n_splits h_iters = h_per_split // n_thr diff --git a/vllm/model_executor/layers/attention/attention.py b/vllm/model_executor/layers/attention/attention.py index 9a85a75ddf..bfedffcffd 100644 --- a/vllm/model_executor/layers/attention/attention.py +++ b/vllm/model_executor/layers/attention/attention.py @@ -41,6 +41,8 @@ from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheSpec, + KVQuantMode, + PrefixAnchoredSWASpec, SlidingWindowSpec, get_kv_quant_mode, ) @@ -372,6 +374,42 @@ def __init__( f"but got {self.attn_backend.get_name()}." ) + prefix_anchored_decode_window = getattr( + vllm_config.attention_config, + "prefix_anchored_decode_window", + None, + ) + if ( + prefix_anchored_decode_window is not None + and attn_type == AttentionType.DECODER + ): + if self.attn_backend.get_name() != "FLASH_ATTN_V100": + raise ValueError( + "prefix-anchored SWA requires standard full decoder " + "attention on FLASH_ATTN_V100" + ) + if self.sliding_window is not None: + raise ValueError( + "prefix-anchored SWA cannot be combined with per-layer " + "sliding-window attention" + ) + if ( + get_kv_quant_mode(self.kv_cache_dtype) != KVQuantMode.NONE + or self.kv_cache_torch_dtype != torch.float16 + ): + raise ValueError("prefix-anchored SWA requires an fp16 KV cache") + if self.head_size_v != self.head_size: + raise ValueError( + "prefix-anchored SWA requires identical key/value head sizes" + ) + if self.has_sink: + raise ValueError( + "prefix-anchored SWA cannot be combined with attention sinks" + ) + extra_impl_args["prefix_anchored_decode_window"] = ( + prefix_anchored_decode_window + ) + if self.attn_backend.get_name() == "FLEX_ATTENTION": block_m = vllm_config.attention_config.flex_attn_block_m block_n = vllm_config.attention_config.flex_attn_block_n @@ -601,6 +639,45 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: # Should not be called for enc-dec or encoder-only attention. assert self.attn_type == AttentionType.DECODER quant_mode = get_kv_quant_mode(self.kv_cache_dtype) + anchored_window = getattr( + vllm_config.attention_config, + "prefix_anchored_decode_window", + None, + ) + if anchored_window is not None: + if self.attn_backend.get_name() != "FLASH_ATTN_V100": + raise ValueError( + "prefix-anchored SWA requires the FLASH_ATTN_V100 backend" + ) + if self.sliding_window is not None: + raise ValueError( + "prefix-anchored SWA cannot be combined with per-layer " + "sliding-window attention" + ) + if ( + quant_mode != KVQuantMode.NONE + or self.kv_cache_torch_dtype != torch.float16 + ): + raise ValueError("prefix-anchored SWA requires an fp16 KV cache") + if self.kv_cache_dtype.startswith("turboquant_"): + raise ValueError("prefix-anchored SWA does not support TurboQuant KV") + if self.head_size_v != self.head_size: + raise ValueError( + "prefix-anchored SWA requires identical key/value head sizes" + ) + if self.has_sink: + raise ValueError( + "prefix-anchored SWA cannot be combined with attention sinks" + ) + return PrefixAnchoredSWASpec( + block_size=block_size, + num_kv_heads=self.num_kv_heads, + head_size=self.head_size, + head_size_v=self.head_size_v, + dtype=self.kv_cache_torch_dtype, + kv_quant_mode=quant_mode, + decode_sliding_window=anchored_window, + ) if self.sliding_window is not None: assert not vllm_config.model_config.use_mla, ( "MLA is not supported for slidingwindow" 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..35491a71e0 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") diff --git a/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py b/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py index 83f599b62b..49f1e3aad1 100644 --- a/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py +++ b/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py @@ -13,6 +13,8 @@ from vllm.triton_utils import tl, triton +from .op import exp + def _parse_positive_int_env(name: str) -> int | None: value = os.getenv(name) @@ -83,6 +85,37 @@ def _select_fused_sigmoid_schedule( return BV, num_warps, num_stages +def _select_fused_sigmoid_launch( + V: int, + N: int, + HV: int, + T: int, + device: torch.device, + *, + match_recurrent_schedule: bool, +) -> tuple[int, int, int]: + if not _use_sm70_fused_sigmoid_schedule(device): + return min(triton.next_power_of_2(V), 32), 4, 3 + if not match_recurrent_schedule: + return _select_fused_sigmoid_schedule(V, N, HV, device) + + # The unfused verifier uses this launch geometry. Keeping it for the first + # fused rollout removes schedule-induced reduction-order changes from the + # correctness comparison; schedule tuning is a separate performance gate. + from .fused_recurrent import ( + _select_sm70_bv, + _select_sm70_num_stages, + _select_sm70_num_warps, + ) + + BV = _select_sm70_bv(V, N, HV, device) + return ( + BV, + _select_sm70_num_warps(BV, N, HV), + _select_sm70_num_stages(T), + ) + + @triton.heuristics( { "USE_INITIAL_STATE": lambda args: args["h0"] is not None, @@ -140,6 +173,10 @@ def fused_sigmoid_gating_delta_rule_update_kernel( IS_DDTREE: tl.constexpr, IS_KDA: tl.constexpr, MIXED_QKV: tl.constexpr, + PRECOMPUTED_GATING: tl.constexpr, + QUANTIZE_BETA: tl.constexpr, + QUANTIZE_STATE_EACH_STEP: tl.constexpr, + MATCH_RECURRENT_NUMERICS: 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 @@ -237,25 +274,44 @@ def fused_sigmoid_gating_delta_rule_update_kernel( b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) - b_b = tl.load(p_b).to(tl.float32) + if PRECOMPUTED_GATING: + b_g = tl.load(p_a).to(tl.float32) + b_beta = tl.load(p_b).to(tl.float32) + else: + b_b = tl.load(p_b).to(tl.float32) - # If the model is loaded in fp16, without the .float() here, A might be -inf - x = tl.load(p_a).to(tl.float32) + tl.load(p_dt_bias).to(tl.float32) - softplus_x = tl.where( - beta * x <= threshold, (1 / beta) * tl.log(1 + tl.exp(beta * x)), x - ) - b_g = -tl.exp(tl.load(p_A_log).to(tl.float32)) * softplus_x + # If the model is loaded in fp16, without the .float() here, A + # might be -inf. + x = tl.load(p_a).to(tl.float32) + tl.load(p_dt_bias).to(tl.float32) + softplus_x = tl.where( + beta * x <= threshold, (1 / beta) * tl.log(1 + tl.exp(beta * x)), x + ) + b_g = -tl.exp(tl.load(p_A_log).to(tl.float32)) * softplus_x - # compute beta_output = sigmoid(b) - b_beta = tl.sigmoid(b_b.to(tl.float32)) + # compute beta_output = sigmoid(b) + b_beta = tl.sigmoid(b_b.to(tl.float32)) + if QUANTIZE_BETA: + # The split GDN verifier materializes beta in b.dtype before + # the recurrent kernel consumes it. Preserve that boundary. + b_beta = b_beta.to(p_b.dtype.element_ty).to(tl.float32) if USE_QK_L2NORM_IN_KERNEL: - b_q = b_q * (tl.rsqrt(tl.sum(b_q * b_q) + 1e-6)) - b_k = b_k * (tl.rsqrt(tl.sum(b_k * b_k) + 1e-6)) + if MATCH_RECURRENT_NUMERICS: + # Preserve the split recurrent verifier's operation order. + # Multiplication by rsqrt is mathematically equivalent but can + # move a probabilistic rejection decision by one FP16 ULP. + b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6) + b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6) + else: + b_q = b_q * (tl.rsqrt(tl.sum(b_q * b_q) + 1e-6)) + b_k = b_k * (tl.rsqrt(tl.sum(b_k * b_k) + 1e-6)) b_q = b_q * scale # [BV, BK] if not IS_KDA: - b_h *= tl.exp(b_g) + if MATCH_RECURRENT_NUMERICS: + b_h *= exp(b_g) + else: + b_h *= tl.exp(b_g) else: b_h *= tl.exp(b_g[None, :]) # [BV] @@ -277,10 +333,17 @@ def fused_sigmoid_gating_delta_rule_update_kernel( p_ht = ht + final_state_idx * stride_final_state_token p_ht = p_ht + i_hv * V * K + o_v[:, None] * K + o_k[None, :] tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) + if QUANTIZE_STATE_EACH_STEP: + # Target verification handles several linear tokens in one + # invocation. Reloading the just-persisted state reproduces + # the low-precision store/load boundary of repeated decode. + b_h = tl.load(p_ht, mask=mask_h, other=0).to(tl.float32) else: p_ht = ht + (bos + i_t) * stride_final_state_token p_ht = p_ht + i_hv * V * K + o_v[:, None] * K + o_k[None, :] tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) + if QUANTIZE_STATE_EACH_STEP: + b_h = tl.load(p_ht, mask=mask_h, other=0).to(tl.float32) if MIXED_QKV: p_q += QKV_STRIDE @@ -410,6 +473,10 @@ def fused_sigmoid_gating_delta_rule_update( USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, IS_KDA=is_kda, MIXED_QKV=False, + PRECOMPUTED_GATING=False, + QUANTIZE_BETA=False, + QUANTIZE_STATE_EACH_STEP=False, + MATCH_RECURRENT_NUMERICS=False, num_warps=num_warps, num_stages=num_stages, ) @@ -541,6 +608,10 @@ def fused_sigmoid_gating_delta_rule_update_mixed_qkv( USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, IS_KDA=False, MIXED_QKV=True, + PRECOMPUTED_GATING=False, + QUANTIZE_BETA=False, + QUANTIZE_STATE_EACH_STEP=False, + MATCH_RECURRENT_NUMERICS=False, num_warps=num_warps, num_stages=num_stages, ) @@ -564,9 +635,21 @@ def fused_sigmoid_gating_delta_rule_update_mixed_qkv_out( initial_state: torch.Tensor | None = None, cu_seqlens: torch.Tensor | None = None, ssm_state_indices: torch.Tensor | None = None, + num_accepted_tokens: torch.Tensor | None = None, use_qk_l2norm_in_kernel: bool = False, + quantize_beta: bool = False, + quantize_state_each_step: bool = False, + match_recurrent_schedule: bool = False, + match_recurrent_numerics: bool = False, + precomputed_g: torch.Tensor | None = None, + precomputed_beta: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - """Mixed-QKV fused update that writes into a caller-provided decode output.""" + """Mixed-QKV update that writes into a caller-provided output buffer. + + ``precomputed_g`` and ``precomputed_beta`` retain the split verifier's + exact gating materialization while still removing the packed-QKV rearrange. + Omitting both computes gating inside the recurrent kernel. + """ if mixed_qkv.ndim != 2: raise ValueError("mixed_qkv must have shape [T, qkv_hidden].") if not mixed_qkv.is_contiguous(): @@ -575,6 +658,10 @@ def fused_sigmoid_gating_delta_rule_update_mixed_qkv_out( raise ValueError("cu_seqlens is required for mixed_qkv_out.") if ssm_state_indices is None: raise ValueError("ssm_state_indices is required for mixed_qkv_out.") + if (precomputed_g is None) != (precomputed_beta is None): + raise ValueError( + "precomputed_g and precomputed_beta must be provided together." + ) T = mixed_qkv.shape[0] N = cu_seqlens.numel() - 1 @@ -597,6 +684,28 @@ def fused_sigmoid_gating_delta_rule_update_mixed_qkv_out( else: assert scale > 0, "scale must be positive" + use_precomputed_gating = precomputed_g is not None + if use_precomputed_gating: + assert precomputed_g is not None + assert precomputed_beta is not None + if precomputed_g.numel() != T * HV: + raise ValueError( + f"precomputed_g must contain {T * HV} values, " + f"got {precomputed_g.numel()}." + ) + if precomputed_beta.numel() != T * HV: + raise ValueError( + f"precomputed_beta must contain {T * HV} values, " + f"got {precomputed_beta.numel()}." + ) + if quantize_beta: + raise ValueError("precomputed_beta is already materialized in its dtype.") + kernel_a = precomputed_g.contiguous().view(T, HV) + kernel_b = precomputed_beta.contiguous().view(T, HV) + else: + kernel_a = a.contiguous() + kernel_b = b.contiguous() + assert initial_state is not None final_state = initial_state stride_init_state_token = initial_state.stride(0) @@ -608,18 +717,21 @@ def fused_sigmoid_gating_delta_rule_update_mixed_qkv_out( stride_indices_seq, stride_indices_tok = ssm_state_indices.stride() BK = triton.next_power_of_2(K) - BV, num_warps, num_stages = ( - _select_fused_sigmoid_schedule(V, N, HV, mixed_qkv.device) - if _use_sm70_fused_sigmoid_schedule(mixed_qkv.device) - else (min(triton.next_power_of_2(V), 32), 4, 3) + BV, num_warps, num_stages = _select_fused_sigmoid_launch( + V, + N, + HV, + T, + mixed_qkv.device, + match_recurrent_schedule=match_recurrent_schedule, ) NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) assert NK == 1, "NK > 1 is not supported yet" fused_sigmoid_gating_delta_rule_update_kernel[(NK, NV, N * HV)]( A_log=A_log, - a=a.contiguous(), - b=b.contiguous(), + a=kernel_a, + b=kernel_b, dt_bias=dt_bias, beta=beta, threshold=threshold, @@ -632,7 +744,7 @@ def fused_sigmoid_gating_delta_rule_update_mixed_qkv_out( ht=final_state, cu_seqlens=cu_seqlens, ssm_state_indices=ssm_state_indices, - num_accepted_tokens=None, + num_accepted_tokens=num_accepted_tokens, ddtree_parent_ids=None, scale=scale, N=N, @@ -658,6 +770,10 @@ def fused_sigmoid_gating_delta_rule_update_mixed_qkv_out( USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, IS_KDA=False, MIXED_QKV=True, + PRECOMPUTED_GATING=use_precomputed_gating, + QUANTIZE_BETA=quantize_beta, + QUANTIZE_STATE_EACH_STEP=quantize_state_each_step, + MATCH_RECURRENT_NUMERICS=match_recurrent_numerics, num_warps=num_warps, num_stages=num_stages, ) diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py index fb94ed46a1..922309a5dc 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe.py @@ -5,6 +5,8 @@ import functools import json import os +from collections.abc import Iterator +from contextlib import contextmanager from typing import Any import torch @@ -34,6 +36,20 @@ logger = init_logger(__name__) +_force_sm70_mtp_moe_legacy_config = False + + +@contextmanager +def force_sm70_mtp_moe_legacy_config() -> Iterator[None]: + """Temporarily select the legacy SM70 MoE tile during startup warmup.""" + global _force_sm70_mtp_moe_legacy_config + previous = _force_sm70_mtp_moe_legacy_config + _force_sm70_mtp_moe_legacy_config = True + try: + yield + finally: + _force_sm70_mtp_moe_legacy_config = previous + @triton.jit def write_zeros_to_output( @@ -1200,6 +1216,32 @@ def should_moe_wna16_use_cuda( ) +def _get_sm70_mtp_moe_decode_config( + M: int, + E: int, + N: int, + K: int, + topk: int, +) -> dict[str, int] | None: + """Return the graph-tuned exact-shape SM70 MTP tile for M2-M16.""" + if _force_sm70_mtp_moe_legacy_config or not envs.VLLM_SM70_MTP_MOE_TUNED_CONFIG: + return None + if (E, N, K, topk) != (256, 128, 2048, 8) or not 2 <= M <= 16: + return None + + # TP4 shards the checkpoint-global I512 expert width to I128 per rank. + # M2-M16 all select this tile in the exact local-shape graph oracle. + return { + "BLOCK_SIZE_M": 8, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "SPLIT_K": 1, + "num_warps": 4, + "num_stages": 3, + } + + def get_default_config( M: int, E: int, @@ -1262,7 +1304,11 @@ def get_default_config( # upstream's larger small-batch tiles are poor for single-token SM70 # decode, which is the FP8-MoE-dequant fallback route used by the # old 35B-FP8 baseline. - if M <= E: + sm70_mtp_config = _get_sm70_mtp_moe_decode_config(M, E, N, K, topk) + if sm70_mtp_config is not None: + config = sm70_mtp_config + logger.info_once(f"Using tuned exact-shape SM70 MTP MoE config: {config}") + elif M <= E: config = { "BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 32, @@ -1278,9 +1324,7 @@ def get_default_config( "GROUP_SIZE_M": 8, "SPLIT_K": 1, } - logger.info_once( - f"Using SM70 0.0.3 unquantized MoE default config: {config}" - ) + logger.info_once(f"Using SM70 0.0.3 unquantized MoE default config: {config}") else: # General defaults for bf16/fp16 and fp8 per-tensor. # Tile sizes scale with batch: small batches are memory-bound 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/layernorm.py b/vllm/model_executor/layers/layernorm.py index 0a72c8a363..0106f989bc 100644 --- a/vllm/model_executor/layers/layernorm.py +++ b/vllm/model_executor/layers/layernorm.py @@ -14,11 +14,100 @@ from vllm.model_executor.custom_op import CustomOp from vllm.model_executor.layers.batch_invariant import rms_norm_batch_invariant from vllm.platforms import current_platform +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_dflash2_gemma_fused_add_rms_kernel( + x, + residual, + weight, + normalized_out, + residual_out, + hidden_size: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + epsilon, +): + row = tl.program_id(0) + cols = tl.arange(0, BLOCK_SIZE) + mask = cols < hidden_size + values = tl.load(x + row * hidden_size + cols, mask=mask, other=0.0).to(tl.float32) + values += tl.load(residual + row * hidden_size + cols, mask=mask, other=0.0).to( + tl.float32 + ) + tl.store(residual_out + row * hidden_size + cols, values, mask=mask) + + variance = tl.sum(tl.where(mask, values * values, 0.0), axis=0) + inverse_rms = tl.rsqrt(variance / hidden_size + epsilon) + gemma_weight = tl.load(weight + cols, mask=mask, other=0.0).to(tl.float32) + 1.0 + tl.store( + normalized_out + row * hidden_size + cols, + values * inverse_rms * gemma_weight, + mask=mask, + ) + + +def _sm70_dflash2_gemma_fused_add_rms_norm( + x: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + variance_epsilon: float, + *, + num_warps: int = 8, +) -> tuple[torch.Tensor, torch.Tensor]: + normalized_out = torch.empty_like(x) + residual_out = torch.empty_like(residual) + _sm70_dflash2_gemma_fused_add_rms_kernel[(x.shape[0],)]( + x, + residual, + weight, + normalized_out, + residual_out, + hidden_size=x.shape[1], + BLOCK_SIZE=triton.next_power_of_2(x.shape[1]), + epsilon=variance_epsilon, + num_warps=num_warps, + num_stages=1, + ) + return normalized_out, residual_out + + +def _use_sm70_dflash2_gemma_fused_add_rms( + x: torch.Tensor, + residual: torch.Tensor | None, + weight: torch.Tensor, +) -> bool: + # Keep the dynamic token dimension out of this Python predicate. AOT traces + # the target once at a large warmup shape; a decode-only row bound would be + # constant-folded there and would leave the M=8 replay on the decomposed path. + return bool( + envs.VLLM_SM70_DFLASH2_FUSED_GEMMA_RMS + and envs.VLLM_SM70_FLASH_V100_0DOT3_COMPILE_GRAPH + # Keep the cached platform query out of the AOT fullgraph. The helper + # below is explicitly constant-foldable by Dynamo. + and _sm70_gemma_long_prefill_available() + and residual is not None + and x.is_cuda + and x.dtype == torch.float16 + and residual.dtype == torch.float32 + and weight.dtype in (torch.float16, torch.bfloat16, torch.float32) + and x.ndim == 2 + and x.shape[0] > 0 + and x.shape[1] == 5120 + and residual.shape == x.shape + and residual.device == x.device + and weight.device == x.device + and weight.ndim == 1 + and weight.numel() == 5120 + and x.is_contiguous() + and residual.is_contiguous() + and weight.is_contiguous() + ) + + @torch.compiler.assume_constant_result def _sm70_gemma_long_prefill_available() -> bool: return current_platform.is_device_capability(70) @@ -338,6 +427,14 @@ def forward_native( residual: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """PyTorch-native implementation equivalent to forward().""" + if _use_sm70_dflash2_gemma_fused_add_rms(x, residual, self.weight): + assert residual is not None + return _sm70_dflash2_gemma_fused_add_rms_norm( + x, + residual, + self.weight, + self.variance_epsilon, + ) if self._use_sm70_long_prefill_fused(x, residual): assert residual is not None if not torch.compiler.is_compiling(): @@ -384,6 +481,14 @@ def forward_cuda( x: torch.Tensor, residual: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if _use_sm70_dflash2_gemma_fused_add_rms(x, residual, self.weight): + assert residual is not None + return _sm70_dflash2_gemma_fused_add_rms_norm( + x, + residual, + self.weight, + self.variance_epsilon, + ) if self._use_sm70_long_prefill_fused(x, residual): assert residual is not None if not torch.compiler.is_compiling(): diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index 3632f22c47..68f99ee3eb 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -64,6 +64,9 @@ def _maybe_sm70_dense_forward( return None if not hasattr(torch.ops._C, "sm70_f16_gemm"): return None + max_m = getattr(layer, "_sm70_f16_max_m", None) + if max_m is not None and x.numel() // x.shape[-1] > max_m: + return None if envs.VLLM_SM70_F16_DENSE_DEBUG or envs.VLLM_QWEN3_NEXT_SM70_TRACE: rows = x.numel() // x.shape[-1] @@ -366,6 +369,15 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: if not current_platform.is_cuda_alike(): 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 @@ -1793,14 +1805,10 @@ def forward( if output is None: output_parallel = _maybe_sm70_dense_forward(self, input_parallel, bias_) if output_parallel is None: - output_parallel = self.quant_method.apply( - self, input_parallel, bias_ - ) + output_parallel = self.quant_method.apply(self, input_parallel, bias_) if self.reduce_results and self.tp_size > 1: - output = _maybe_sm70_awq_mlp_down_tile_all_reduce( - self, output_parallel - ) + output = _maybe_sm70_awq_mlp_down_tile_all_reduce(self, output_parallel) if output is None: output = tensor_model_parallel_all_reduce(output_parallel) else: diff --git a/vllm/model_executor/layers/logits_processor.py b/vllm/model_executor/layers/logits_processor.py index 1b49aa4f6e..8bf4ef2a28 100644 --- a/vllm/model_executor/layers/logits_processor.py +++ b/vllm/model_executor/layers/logits_processor.py @@ -151,8 +151,9 @@ def _cuda_graph_capture_active() -> bool: return False -def _maybe_sync_top1_all_gather(module: torch.nn.Module, - local_pair: torch.Tensor) -> None: +def _maybe_sync_top1_all_gather( + module: torch.nn.Module, local_pair: torch.Tensor +) -> None: raw_steps = envs.VLLM_SM70_SYNC_TOP1_ALLGATHER_STEPS if not raw_steps or not local_pair.is_cuda: return @@ -256,8 +257,8 @@ def _get_logits( pre_stream_sync_ms = 0.0 pre_device_extra_sync_ms = 0.0 if profile_enabled: - pre_stream_sync_ms, pre_device_extra_sync_ms = ( - _cuda_sync_breakdown_ms(hidden_states) + pre_stream_sync_ms, pre_device_extra_sync_ms = _cuda_sync_breakdown_ms( + hidden_states ) pre_sync_ms = pre_stream_sync_ms + pre_device_extra_sync_ms @@ -315,9 +316,7 @@ def get_top_tokens( ) tp_size = get_tensor_model_parallel_world_size() - local_top1 = lm_head.maybe_get_sm70_lm_head_top1( - hidden_states, embedding_bias - ) + local_top1 = lm_head.maybe_get_sm70_lm_head_top1(hidden_states, embedding_bias) if local_top1 is None: logits = lm_head.quant_method.apply( lm_head, hidden_states, bias=embedding_bias @@ -396,14 +395,12 @@ def get_topk_tokens_and_logprobs( pre_stream_sync_ms = 0.0 pre_device_extra_sync_ms = 0.0 if profile_enabled: - pre_stream_sync_ms, pre_device_extra_sync_ms = ( - _cuda_sync_breakdown_ms(hidden_states) + pre_stream_sync_ms, pre_device_extra_sync_ms = _cuda_sync_breakdown_ms( + hidden_states ) pre_sync_ms = pre_stream_sync_ms + pre_device_extra_sync_ms stage_start = _cuda_stage_start(profile_enabled) - logits = lm_head.quant_method.apply( - lm_head, hidden_states, bias=embedding_bias - ) + logits = lm_head.quant_method.apply(lm_head, hidden_states, bias=embedding_bias) local_lm_head_ms = _cuda_stage_ms(profile_enabled, stage_start) stage_start = _cuda_stage_start(profile_enabled) @@ -530,43 +527,66 @@ def get_topk_tokens_and_logits( pre_stream_sync_ms = 0.0 pre_device_extra_sync_ms = 0.0 if profile_enabled: - pre_stream_sync_ms, pre_device_extra_sync_ms = ( - _cuda_sync_breakdown_ms(hidden_states) + pre_stream_sync_ms, pre_device_extra_sync_ms = _cuda_sync_breakdown_ms( + hidden_states ) pre_sync_ms = pre_stream_sync_ms + pre_device_extra_sync_ms stage_start = _cuda_stage_start(profile_enabled) - logits = lm_head.quant_method.apply( - lm_head, hidden_states, bias=embedding_bias + local_candidates = lm_head.maybe_get_sm70_dflash2_top20( + hidden_states, + top_k, + embedding_bias, ) - local_lm_head_ms = _cuda_stage_ms(profile_enabled, stage_start) + if local_candidates is None: + logits = lm_head.quant_method.apply( + lm_head, hidden_states, bias=embedding_bias + ) + local_lm_head_ms = _cuda_stage_ms(profile_enabled, stage_start) - stage_start = _cuda_stage_start(profile_enabled) - if self.soft_cap is not None: - logits = torch.tanh(logits / self.soft_cap) * self.soft_cap - if self.scale != 1.0: - logits = logits * self.scale + stage_start = _cuda_stage_start(profile_enabled) + if self.soft_cap is not None: + logits = torch.tanh(logits / self.soft_cap) * self.soft_cap + if self.scale != 1.0: + logits = logits * self.scale - num_pad = lm_head.shard_indices.num_org_vocab_padding - if num_pad > 0: - logits[..., -num_pad:] = -float("inf") + num_pad = lm_head.shard_indices.num_org_vocab_padding + if num_pad > 0: + logits[..., -num_pad:] = -float("inf") - logits_float = logits.float() - postprocess_ms = _cuda_stage_ms(profile_enabled, stage_start) + logits_float = logits.float() + postprocess_ms = _cuda_stage_ms(profile_enabled, stage_start) - local_k = min(top_k, logits_float.shape[-1]) - stage_start = _cuda_stage_start(profile_enabled) - local_vals, local_indices = torch.topk( - logits_float, - k=local_k, - dim=-1, - ) - local_topk_ms = _cuda_stage_ms(profile_enabled, stage_start) + local_k = min(top_k, logits_float.shape[-1]) + stage_start = _cuda_stage_start(profile_enabled) + local_vals, local_indices = torch.topk( + logits_float, + k=local_k, + dim=-1, + ) + local_topk_ms = _cuda_stage_ms(profile_enabled, stage_start) - stage_start = _cuda_stage_start(profile_enabled) - vocab_start = lm_head.shard_indices.org_vocab_start_index - local_global_indices = local_indices + vocab_start - local_indices_ms = _cuda_stage_ms(profile_enabled, stage_start) + stage_start = _cuda_stage_start(profile_enabled) + vocab_start = lm_head.shard_indices.org_vocab_start_index + local_global_indices = local_indices + vocab_start + local_indices_ms = _cuda_stage_ms(profile_enabled, stage_start) + local_vocab_size = int(logits_float.shape[-1]) + else: + local_vals, local_global_indices = local_candidates + local_lm_head_ms = _cuda_stage_ms(profile_enabled, stage_start) + stage_start = _cuda_stage_start(profile_enabled) + # Positive scale and tanh soft-capping are monotonic, so applying + # them after the exact local top-k preserves the dense ordering. + if self.soft_cap is not None: + local_vals = torch.tanh(local_vals / self.soft_cap) * self.soft_cap + if self.scale != 1.0: + local_vals = local_vals * self.scale + local_vals = local_vals.float() + postprocess_ms = _cuda_stage_ms(profile_enabled, stage_start) + local_k = int(local_vals.shape[-1]) + local_topk_ms = 0.0 + local_indices_ms = 0.0 + local_vocab_size = int(lm_head.weight.shape[0]) tp_size = get_tensor_model_parallel_world_size() gather_vals_ms = 0.0 @@ -615,7 +635,7 @@ def get_topk_tokens_and_logits( gather_indices_ms, merge_topk_ms, int(hidden_states.shape[0]), - int(logits_float.shape[-1]), + local_vocab_size, top_k, local_k, tp_size, 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 c28dd507f6..5f7c5ce0bf 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 @@ -12,6 +12,7 @@ from einops import rearrange from torch import nn +from vllm import _sm70_ops as sm70_ops from vllm import envs from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig, get_current_vllm_config @@ -74,6 +75,7 @@ get_registered_gdn_spec_metadata_tensors, ) from vllm.v1.attention.backends.utils import compute_causal_conv1d_metadata +from vllm.v1.worker.gpu.spec_decode import uses_dflash_selector_engine # Optional ROCm AITER Triton kernels for the GDN decode fast-path. # Availability is checked centrally via rocm_aiter_ops; the actual function @@ -103,6 +105,80 @@ _DFLASH_DDTREE_PATH_PROBE_REPORTS = 0 +@triton.jit +def _sm70_pack_qwen_gdn_qkv_kernel( + mixed_qkv, + packed, + input_row_stride: tl.int64, + num_rows: tl.constexpr, + q_dim: tl.constexpr, + k_dim: tl.constexpr, + v_dim: tl.constexpr, + BLOCK_Q: tl.constexpr, + BLOCK_K: tl.constexpr, + BLOCK_V: tl.constexpr, +): + """Bitwise row-slice copy into [all Q][all K][all V] storage.""" + row = tl.program_id(0) + input_row = mixed_qkv + row * input_row_stride + + q_cols = tl.arange(0, BLOCK_Q) + q_mask = q_cols < q_dim + q = tl.load(input_row + q_cols, mask=q_mask) + tl.store(packed + row * q_dim + q_cols, q, mask=q_mask) + + k_cols = tl.arange(0, BLOCK_K) + k_mask = k_cols < k_dim + k = tl.load(input_row + q_dim + k_cols, mask=k_mask) + q_numel = num_rows * q_dim + tl.store(packed + q_numel + row * k_dim + k_cols, k, mask=k_mask) + + v_cols = tl.arange(0, BLOCK_V) + v_mask = v_cols < v_dim + v = tl.load(input_row + q_dim + k_dim + v_cols, mask=v_mask) + qk_numel = q_numel + num_rows * k_dim + tl.store(packed + qk_numel + row * v_dim + v_cols, v, mask=v_mask) + + +def _sm70_pack_qwen_gdn_qkv( + mixed_qkv: torch.Tensor, + q_dim: int, + k_dim: int, + v_dim: int, +) -> torch.Tensor: + """Materialize the recurrent Q/K/V layout with one copy launch.""" + if mixed_qkv.ndim != 2: + raise ValueError("mixed_qkv must be rank two") + if mixed_qkv.stride(1) != 1: + raise ValueError("mixed_qkv must be contiguous by row") + if min(q_dim, k_dim, v_dim) <= 0: + raise ValueError("Q/K/V widths must be positive") + if mixed_qkv.shape[1] < q_dim + k_dim + v_dim: + raise ValueError("mixed_qkv is narrower than the Q/K/V widths") + + num_rows = mixed_qkv.shape[0] + packed = torch.empty( + num_rows * (q_dim + k_dim + v_dim), + dtype=mixed_qkv.dtype, + device=mixed_qkv.device, + ) + _sm70_pack_qwen_gdn_qkv_kernel[(num_rows,)]( + mixed_qkv, + packed, + mixed_qkv.stride(0), + num_rows=num_rows, + q_dim=q_dim, + k_dim=k_dim, + v_dim=v_dim, + BLOCK_Q=triton.next_power_of_2(q_dim), + BLOCK_K=triton.next_power_of_2(k_dim), + BLOCK_V=triton.next_power_of_2(v_dim), + num_warps=8, + num_stages=1, + ) + return packed + + def _ddtree_parent_ids_require_branch( parent_ids: torch.Tensor | None, num_tree_tokens_cpu: torch.Tensor | None, @@ -427,11 +503,19 @@ def _sm70_dump_gdn_core_tensor( tensor: torch.Tensor, source: str = "core", ) -> None: - _sm70_gdn_graph_buffer_copy(label, layer_name, tensor, source) - if torch.cuda.is_current_stream_capturing(): - return dump_dir = os.getenv("VLLM_SM70_DUMP_GDN_CORE_DIR") - if not dump_dir: + graph_dump = os.getenv("VLLM_SM70_DUMP_GDN_GRAPH_BUFFERS") == "1" and bool( + os.getenv("VLLM_SM70_DUMP_GDN_GRAPH_DIR") + ) + # Keep the production path free of CUDA runtime queries. This hook runs + # several times per GDN layer, while core dumps are normally disabled. + if not dump_dir and not graph_dump: + return + if graph_dump: + _sm70_gdn_graph_buffer_copy(label, layer_name, tensor, source) + if not dump_dir or not tensor.is_cuda: + return + if torch.cuda.is_current_stream_capturing(): return raw_layer_ids = os.getenv("VLLM_SM70_DUMP_GDN_CORE_LAYER_IDS") if raw_layer_ids: @@ -590,6 +674,10 @@ def _sm70_qwen_gdn_spec_method(vllm_config: object) -> str | None: return getattr(spec_config, "method", None) +def _is_dflash2_spec_config(vllm_config: object) -> bool: + return uses_dflash_selector_engine(vllm_config) + + def _sm70_qwen_gdn_full_forward_enabled( layer_name: LayerNameType, *, @@ -1250,6 +1338,60 @@ def _sm70_compile_graph_slice_dim( return tensor.index_select(dim, indices) +def _sm70_gdn_rmsnorm_onepass_enabled() -> bool: + return envs.VLLM_SM70_GDN_RMSNORM_ONEPASS + + +_SM70_GDN_QPN8_BA_REQUIRED_OPS = ( + "fp8_qpn8_gemm_ba_split_sm70_out", + "fp8_qpn8_dispatch_ba_split_sm70_out", +) + + +def _missing_sm70_gdn_qpn8_ba_ops() -> list[str]: + return [ + name + for name in _SM70_GDN_QPN8_BA_REQUIRED_OPS + if not hasattr(torch.ops._C, name) + ] + + +def _sm70_gdn_qpn8_ba_split_enabled() -> bool: + if envs.VLLM_SM70_GDN_QPN8_BA_SPLIT and not envs.VLLM_SM70_GDN_RMSNORM_ONEPASS: + raise RuntimeError( + "VLLM_SM70_GDN_QPN8_BA_SPLIT=1 requires the accepted " + "VLLM_SM70_GDN_RMSNORM_ONEPASS=1 pair." + ) + if envs.VLLM_SM70_GDN_QPN8_BA_SPLIT: + missing_ops = _missing_sm70_gdn_qpn8_ba_ops() + if missing_ops: + raise RuntimeError( + "VLLM_SM70_GDN_QPN8_BA_SPLIT=1 requires the source-built " + f"SM70 GDN extension; missing ops: {missing_ops}." + ) + return envs.VLLM_SM70_GDN_QPN8_BA_SPLIT + + +@triton.jit +def _sm70_gdn_rmsnorm_gated_onepass_kernel( + x_ptr, + z_ptr, + weight_ptr, + out_ptr, + D: tl.constexpr, + EPS: tl.constexpr, +): + row = tl.program_id(0) + offsets = tl.arange(0, D) + x = tl.load(x_ptr + row * D + offsets).to(tl.float32) + z = tl.load(z_ptr + row * D + offsets).to(tl.float32) + weight = tl.load(weight_ptr + offsets).to(tl.float32) + sumsq = tl.sum(x * x, axis=0) + rstd = tl.rsqrt(sumsq / D + EPS) + gated = z * tl.sigmoid(z) + tl.store(out_ptr + row * D + offsets, x * rstd * weight * gated) + + def _sm70_qwen_gdn_rmsnorm_gated_impl( x: torch.Tensor, z: torch.Tensor, @@ -1259,6 +1401,39 @@ def _sm70_qwen_gdn_rmsnorm_gated_impl( norm_before_gate: bool, activation: str, ) -> torch.Tensor: + if ( + _sm70_gdn_rmsnorm_onepass_enabled() + and current_platform.is_device_capability(70) + and x.is_cuda + and x.dtype == torch.float16 + and x.shape == (12, 128) + and x.is_contiguous() + and z.is_cuda + and z.device == x.device + and z.dtype == x.dtype + and z.shape == x.shape + and z.is_contiguous() + and weight.is_cuda + and weight.device == x.device + and weight.dtype == x.dtype + and weight.shape == (128,) + and weight.is_contiguous() + and group_size <= 0 + and norm_before_gate + and activation in ("silu", "swish") + ): + out = torch.empty_like(x) + _sm70_gdn_rmsnorm_gated_onepass_kernel[(12,)]( + x, + z, + weight, + out, + D=128, + EPS=eps, + num_warps=2, + ) + return out + from vllm.model_executor.layers.fla.ops.layernorm_guard import rmsnorm_fn return rmsnorm_fn( @@ -2195,12 +2370,54 @@ def __init__( 1, 1 + int(vllm_config.speculative_config.num_speculative_state_tokens()), ) + self.enable_sm70_gdn_rmsnorm_onepass = ( + _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.tp_size == 4 + and self.num_v_heads == 48 + and self.head_v_dim == 128 + and not self.gqa_interleaved_layout + ) + 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.tp_size == 4 + and self.num_v_heads == 48 + and self.head_v_dim == 128 + and not self.gqa_interleaved_layout + and not self.disable_tp_for_ba_proj + and self.in_proj_ba is not None + and _sm70_gdn_qpn8_ba_split_enabled() + ) self.enable_packed_recurrent_decode = ( envs.VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE ) self.enable_sm70_fused_sigmoid_mixed_qkv = ( envs.VLLM_SM70_FUSED_SIGMOID_MIXED_QKV ) + self.enable_sm70_dflash2_fused_gdn_verify = bool( + envs.VLLM_SM70_DFLASH2_FUSED_GDN_VERIFY + and current_platform.is_device_capability(70) + and _is_dflash2_spec_config(vllm_config) + ) + self.enable_sm70_dflash2_fused_gdn_norm = bool( + envs.VLLM_SM70_DFLASH2_FUSED_GDN_NORM + and current_platform.is_device_capability(70) + and _is_dflash2_spec_config(vllm_config) + ) + self.enable_sm70_dflash2_fused_gdn_split = bool( + envs.VLLM_SM70_DFLASH2_FUSED_GDN_SPLIT + and current_platform.is_device_capability(70) + and _is_dflash2_spec_config(vllm_config) + ) + self.enable_sm70_dflash2_fused_qkv_pack = bool( + envs.VLLM_SM70_DFLASH2_FUSED_QKV_PACK + and current_platform.is_device_capability(70) + and _is_dflash2_spec_config(vllm_config) + ) self.compare_sm70_fused_sigmoid_mixed_qkv = ( envs.VLLM_SM70_FUSED_SIGMOID_MIXED_QKV_COMPARE ) @@ -2289,6 +2506,26 @@ def __init__( "SM70 FlashQLA GDN decode route enabled.", scope="local", ) + if self.enable_sm70_dflash2_fused_gdn_verify: + logger.info_once( + "SM70 DFlash2 packed GDN target-verification route enabled.", + scope="local", + ) + if self.enable_sm70_dflash2_fused_gdn_norm: + logger.info_once( + "SM70 DFlash2 one-pass GDN output RMSNorm route enabled.", + scope="local", + ) + if self.enable_sm70_dflash2_fused_gdn_split: + logger.info_once( + "SM70 DFlash2 fused GDN z/b/a materialization route enabled.", + scope="local", + ) + if self.enable_sm70_dflash2_fused_qkv_pack: + logger.info_once( + "SM70 DFlash2 fused post-convolution Q/K/V packing enabled.", + scope="local", + ) if envs.is_set("VLLM_QWEN3_NEXT_FUSED_SIGMOID_GATING"): logger.info_once( "VLLM_QWEN3_NEXT_FUSED_SIGMOID_GATING is upstream-split " @@ -2648,11 +2885,19 @@ def rearrange_mixed_qkv(self, mixed_qkv): k_dim = self.key_dim // self.tp_size v_dim = self.value_dim // self.tp_size - query, key, value = torch.split(mixed_qkv, [q_dim, k_dim, v_dim], dim=-1) - - fused = torch.cat( - [query.reshape(-1), key.reshape(-1), value.reshape(-1)], dim=0 - ) + if ( + self.enable_sm70_dflash2_fused_qkv_pack + and seq_len == 8 + and mixed_qkv.is_cuda + and mixed_qkv.dtype == torch.float16 + and mixed_qkv.stride(1) == 1 + ): + fused = _sm70_pack_qwen_gdn_qkv(mixed_qkv, q_dim, k_dim, v_dim) + else: + query, key, value = torch.split(mixed_qkv, [q_dim, k_dim, v_dim], dim=-1) + fused = torch.cat( + [query.reshape(-1), key.reshape(-1), value.reshape(-1)], dim=0 + ) q_size = seq_len * q_dim k_size = seq_len * k_dim @@ -2745,6 +2990,7 @@ def _probe_ddtree_path_replay_states( a_spec.index_select(0, path_token_idx), b_spec.index_select(0, path_token_idx), self.dt_bias, + beta_dtype=torch.float32, ) local_ssm_state = initial_ssm_state[req_row : req_row + 1].clone() _, root_final_state = fused_recurrent_gated_delta_rule( @@ -2841,6 +3087,7 @@ def _probe_ddtree_path_replay_states( a_spec.index_select(0, path_token_idx), b_spec.index_select(0, path_token_idx), self.dt_bias, + beta_dtype=torch.float32, ) local_ssm_state = initial_ssm_state[req_row : req_row + 1].clone() _, path_final_state = fused_recurrent_gated_delta_rule( @@ -3134,6 +3381,7 @@ def _forward_ddtree_gdn_pure_spec( a_spec, b_spec, self.dt_bias, + beta_dtype=torch.float32, ) core_attn_out_linear, _ = fused_recurrent_gated_delta_rule( q=query_linear, @@ -3249,6 +3497,7 @@ def _forward_ddtree_gdn_pure_spec( a_spec.index_select(0, token_idx), b_spec.index_select(0, token_idx), self.dt_bias, + beta_dtype=torch.float32, ) cu_seqlens = torch.arange( input_state_idx.shape[0] + 1, @@ -3336,6 +3585,7 @@ def _forward_ddtree_gdn_pure_spec( a_spec.index_select(0, token_idx), b_spec.index_select(0, token_idx), self.dt_bias, + beta_dtype=torch.float32, ) cu_seqlens = torch.arange( conv_state_indices.shape[0] + 1, @@ -3560,7 +3810,7 @@ def _forward_core_decode_flashqla( def forward( self, hidden_states: torch.Tensor, - output: torch.Tensor, + output: torch.Tensor | None, ): if self.maybe_sm70_qwen_gdn_full_forward: layer_name = _encode_layer_name(self.prefix) @@ -3582,14 +3832,14 @@ def forward( layer_name, ) return - self._forward_method(hidden_states, output) + return self._forward_method(hidden_states, output) def _full_forward( self, hidden_states: torch.Tensor, output: torch.Tensor, ): - self._forward_method(hidden_states, output) + return self._forward_method(hidden_states, output) def _compute_output_projection( self, @@ -3612,7 +3862,44 @@ def _compute_output_projection( ) z = _sm70_dump_gdn_projection_tensor("proj_z", layer_name, z) profile_start = _sm70_gdn_prefill_profile_start() - core_attn_out = self.norm(core_attn_out, z) + use_sm70_onepass_norm = ( + self.enable_sm70_gdn_rmsnorm_onepass + and num_tokens == 1 + and core_attn_out.dtype == torch.float16 + and core_attn_out.shape == (12, 128) + and core_attn_out.is_contiguous() + and z.dtype == core_attn_out.dtype + and z.shape == core_attn_out.shape + and z.is_contiguous() + and self.norm.weight.dtype == core_attn_out.dtype + and self.norm.weight.shape == (128,) + and self.norm.weight.is_contiguous() + and self.norm.group_size is None + and self.norm.norm_before_gate + and self.norm.activation in ("silu", "swish") + ) + if self.enable_sm70_dflash2_fused_gdn_norm: + # RMSNormGated is decomposed into a nine-kernel native chain when + # it is reached from the opaque Qwen GDN op under the default + # inductor custom-op policy. Its canonical CUDA implementation is + # a one-pass Triton kernel with the same FP32 accumulation and + # FP16 output contract, so call it directly on this isolated path. + core_attn_out = self.norm.forward_cuda(core_attn_out, z) + elif use_sm70_onepass_norm: + _log_runtime_route_once( + "SM70 GDN 12x128 RMSNormGated one-pass route enabled." + ) + core_attn_out = torch.ops.vllm.sm70_qwen_gdn_rmsnorm_gated( + core_attn_out, + z, + self.norm.weight, + self.norm.eps, + -1, + True, + self.norm.activation, + ) + else: + core_attn_out = self.norm(core_attn_out, z) _sm70_gdn_prefill_profile_end( layer_name, "projection_norm", @@ -3645,18 +3932,19 @@ def _output_projection( self, core_attn_out: torch.Tensor, z: torch.Tensor, - output: torch.Tensor, + output: torch.Tensor | None, num_tokens: int, ) -> torch.Tensor: layer_name = _encode_layer_name(self.prefix) proj_out = self._compute_output_projection(core_attn_out, z, num_tokens) - output[:num_tokens] = proj_out - _sm70_gdn_graph_buffer_copy( - "proj_output_after_write", - layer_name, - output[:num_tokens], - "proj", - ) + if output is not None: + output[:num_tokens] = proj_out + _sm70_gdn_graph_buffer_copy( + "proj_output_after_write", + layer_name, + output[:num_tokens], + "proj", + ) return proj_out def forward_hip( @@ -3715,7 +4003,7 @@ def forward_hip( def forward_cuda( self, hidden_states: torch.Tensor, - output: torch.Tensor, + output: torch.Tensor | None, ): """ Forward pass with three parts: @@ -3753,8 +4041,8 @@ def forward_cuda( layer_name, ) else: - self._output_projection(core_attn_out, z, output, num_tokens) - return + return self._output_projection(core_attn_out, z, output, num_tokens) + return None if envs.VLLM_SM70_QWEN_GDN_INPUT_PROJECTION_OP: mixed_qkv = torch.empty( @@ -3805,8 +4093,8 @@ def forward_cuda( layer_name, ) else: - self._output_projection(core_attn_out, z, output, num_tokens) - return + return self._output_projection(core_attn_out, z, output, num_tokens) + return None # ============================================================ # Part 1: Input Projection @@ -3904,7 +4192,8 @@ def forward_cuda( layer_name, ) else: - self._output_projection(core_attn_out, z, output, num_tokens) + return self._output_projection(core_attn_out, z, output, num_tokens) + return None def forward_xpu( self, @@ -4813,6 +5102,85 @@ def _forward_core_rocm( kv_cache=kv_cache, ) + def _can_use_dflash2_packed_gdn_verify( + self, + *, + mixed_qkv: torch.Tensor | None, + a: torch.Tensor, + b: torch.Tensor, + core_attn_out: torch.Tensor, + ssm_state: torch.Tensor, + attn_metadata: GDNAttentionMetadata, + ) -> bool: + return bool( + self.enable_sm70_dflash2_fused_gdn_verify + and mixed_qkv is not None + and attn_metadata.spec_sequence_masks is not None + and attn_metadata.num_spec_decodes > 0 + and attn_metadata.num_prefills == 0 + and attn_metadata.num_decodes == 0 + and attn_metadata.ddtree_parent_ids is None + and attn_metadata.spec_query_start_loc is not None + and attn_metadata.spec_state_indices_tensor is not None + and attn_metadata.spec_state_slot_selectors is not None + and mixed_qkv.is_cuda + and mixed_qkv.dtype == torch.float16 + and a.dtype == torch.float16 + and b.dtype == torch.float16 + # The supported verifier contract keeps recurrent state in FP32; + # an explicit FP16 cache override is also supported. + and ssm_state.dtype in (torch.float16, torch.float32) + and mixed_qkv.is_contiguous() + and a.is_contiguous() + and b.is_contiguous() + and core_attn_out.is_contiguous() + ) + + def _forward_dflash2_packed_gdn_verify( + self, + *, + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + core_attn_out: torch.Tensor, + ssm_state: torch.Tensor, + spec_query_start_loc: torch.Tensor, + spec_state_indices_tensor: torch.Tensor, + spec_state_slot_selectors: torch.Tensor, + num_spec_decodes: int, + ) -> torch.Tensor: + num_tokens = mixed_qkv.shape[0] + out = core_attn_out[:num_tokens].unsqueeze(1) + g, beta = fused_gdn_gating(self.A_log, a, b, self.dt_bias) + fused_sigmoid_gating_delta_rule_update_mixed_qkv_out( + A_log=self.A_log, + a=a, + b=b, + dt_bias=self.dt_bias, + mixed_qkv=mixed_qkv, + num_q_heads=self.num_k_heads // self.tp_size, + num_v_heads=self.num_v_heads // self.tp_size, + head_k_dim=self.head_k_dim, + head_v_dim=self.head_v_dim, + scale=self.head_k_dim**-0.5, + initial_state=ssm_state, + out=out, + cu_seqlens=spec_query_start_loc[: num_spec_decodes + 1], + ssm_state_indices=spec_state_indices_tensor[:num_spec_decodes], + num_accepted_tokens=spec_state_slot_selectors[:num_spec_decodes], + use_qk_l2norm_in_kernel=True, + # Keep gating materialization and launch geometry bitwise aligned + # with the current verifier. This stage removes packed-QKV + # rearrangement and the final output copy without changing the + # target recurrence. Fully fused gating stays separately testable. + precomputed_g=g, + precomputed_beta=beta, + quantize_state_each_step=False, + match_recurrent_schedule=True, + match_recurrent_numerics=True, + ) + return out.transpose(0, 1) + def _forward_core( self, mixed_qkv: torch.Tensor, @@ -5165,7 +5533,18 @@ def _forward_core( core_attn_out[:num_non_spec_tokens] = core_attn_out_non_spec.squeeze(0) return - query_spec, key_spec, value_spec = self.rearrange_mixed_qkv(mixed_qkv_spec) + use_dflash2_packed_gdn_verify = self._can_use_dflash2_packed_gdn_verify( + mixed_qkv=mixed_qkv_spec, + a=a, + b=b, + core_attn_out=core_attn_out, + ssm_state=ssm_state, + attn_metadata=attn_metadata, + ) + if use_dflash2_packed_gdn_verify: + query_spec, key_spec, value_spec = None, None, None + else: + query_spec, key_spec, value_spec = self.rearrange_mixed_qkv(mixed_qkv_spec) prefill_l2norm_in_kernel = False use_original_flashqla_prefill = _sm70_flashqla_original_prefill_enabled() prefill_gate_is_exp = False @@ -5280,7 +5659,27 @@ def _forward_core( assert spec_token_indx is not None a_spec = a.index_select(0, spec_token_indx) b_spec = b.index_select(0, spec_token_indx) - if ddtree_tree_gdn_pure_spec: + if use_dflash2_packed_gdn_verify: + assert mixed_qkv_spec is not None + assert spec_query_start_loc is not None + assert spec_state_indices_tensor is not None + assert spec_state_slot_selectors is not None + core_attn_out_spec = self._forward_dflash2_packed_gdn_verify( + mixed_qkv=mixed_qkv_spec, + a=a_spec, + b=b_spec, + core_attn_out=core_attn_out, + ssm_state=ssm_state, + spec_query_start_loc=spec_query_start_loc, + spec_state_indices_tensor=spec_state_indices_tensor, + spec_state_slot_selectors=spec_state_slot_selectors, + num_spec_decodes=attn_metadata.num_spec_decodes, + ) + last_recurrent_state = ssm_state + _log_runtime_route_once( + "SM70 DFlash2 packed GDN target-verification route hit." + ) + elif ddtree_tree_gdn_pure_spec: assert mixed_qkv_spec is not None assert spec_query_start_loc is not None assert spec_state_indices_tensor is not None @@ -5313,8 +5712,15 @@ def _forward_core( ) last_recurrent_state = ssm_state else: + assert query_spec is not None + assert key_spec is not None + assert value_spec is not None g_spec, beta_spec = fused_gdn_gating( - self.A_log, a_spec, b_spec, self.dt_bias + self.A_log, + a_spec, + b_spec, + self.dt_bias, + beta_dtype=torch.float32, ) core_attn_out_spec, last_recurrent_state = ( fused_recurrent_gated_delta_rule( @@ -5515,7 +5921,8 @@ def _forward_core( merged_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec) core_attn_out[:num_actual_tokens] = merged_out.squeeze(0) elif spec_sequence_masks is not None: - core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0) + if not use_dflash2_packed_gdn_verify: + core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0) elif core_attn_out_non_spec is not None and not ( attn_metadata.num_prefills > 0 and spec_sequence_masks is None @@ -6298,12 +6705,41 @@ def _patch_fallback_metadata(name: str, value: object) -> None: max_query_len=spec_state_indices_tensor.size(-1), validate_data=False, ) + if isinstance(self, QwenGatedDeltaNetAttention) and ( + self._can_use_dflash2_packed_gdn_verify( + mixed_qkv=mixed_qkv_spec, + a=a, + b=b, + core_attn_out=core_attn_out, + ssm_state=ssm_state, + attn_metadata=attn_metadata, + ) + ): + self._forward_dflash2_packed_gdn_verify( + mixed_qkv=mixed_qkv_spec, + a=a, + b=b, + core_attn_out=core_attn_out, + ssm_state=ssm_state, + spec_query_start_loc=spec_query_start_loc, + spec_state_indices_tensor=spec_state_indices_tensor, + spec_state_slot_selectors=spec_state_slot_selectors, + num_spec_decodes=attn_metadata.num_spec_decodes, + ) + _log_runtime_route_once( + "SM70 DFlash2 packed GDN target-verification route hit." + ) + _sm70_dump_gdn_core_tensor( + "core_out", layer_name, core_attn_out, metadata_source + ) + return core_attn_out query_spec, key_spec, value_spec = self.rearrange_mixed_qkv(mixed_qkv_spec) g_spec, beta_spec = fused_gdn_gating( self.A_log, a, b, self.dt_bias, + beta_dtype=torch.float32, ) core_attn_out_spec, _ = fused_recurrent_gated_delta_rule( q=query_spec, @@ -6682,6 +7118,77 @@ def qwen_gdn_output_projection_fake( """Fake implementation for torch.compile.""" +def _sm70_gdn_qpn8_ba_weight_contract( + self: "QwenGatedDeltaNetAttention", +) -> bool: + """Check the loaded operator layout without using model identity.""" + qkvz = self.in_proj_qkvz + ba = self.in_proj_ba + 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 + return bool( + self.enable_sm70_gdn_qpn8_ba_split + and ba is not None + and getattr(qkvz, "sm70_fp8_qpn8", False) + 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 qkvz_weight.is_contiguous() + and isinstance(qkvz_scales, torch.Tensor) + and qkvz_scales.dtype == torch.float16 + and qkvz_scales.shape == (1, 4096) + 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.is_contiguous() + and qkvz_weight.device == qkvz_scales.device == ba_weight.device + and getattr(qkvz, "bias", None) is None + and getattr(ba, "bias", None) is None + ) + + +def _sm70_gdn_qpn8_ba_dispatch_eligible( + self: "QwenGatedDeltaNetAttention", + hidden_states: torch.Tensor, + layer_name: LayerNameType, +) -> bool: + qkvz = self.in_proj_qkvz + ba = self.in_proj_ba + return bool( + _sm70_gdn_qpn8_ba_weight_contract(self) + and ba is not None + and not _sm70_gdn_projection_dump_requested(layer_name) + and hasattr(torch.ops._C, "fp8_qpn8_dispatch_ba_split_sm70_out") + and hasattr(torch.ops._C, "fp8_qpn8_gemm_ba_split_sm70_out") + 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.is_contiguous() + and hidden_states.device == qkvz.weight.device == ba.weight.device + ) + + +def _sm70_gdn_qpn8_ba_split_eligible( + self: "QwenGatedDeltaNetAttention", + hidden_states: torch.Tensor, + z_out: torch.Tensor, + layer_name: LayerNameType, +) -> bool: + return ( + _sm70_gdn_qpn8_ba_dispatch_eligible(self, hidden_states, layer_name) + and hidden_states.shape == (1, 5120) + and z_out.is_cuda + and z_out.device == hidden_states.device + and z_out.dtype == hidden_states.dtype + and z_out.shape == (1, 12, 128) + and z_out.is_contiguous() + ) + + def qwen_gdn_input_projection_core( hidden_states: torch.Tensor, z_out: torch.Tensor, @@ -6700,46 +7207,75 @@ def qwen_gdn_input_projection_core( layer_name, hidden_states, ) - mixed_qkvz, _ = self.in_proj_qkvz(hidden_states) - ba, _ = self.in_proj_ba(hidden_states) - mixed_qkvz = _sm70_dump_gdn_projection_tensor( - "input_core_in_proj_qkvz", + if _sm70_gdn_qpn8_ba_split_eligible( + self, + hidden_states, + z_out, layer_name, - mixed_qkvz, - ) - ba = _sm70_dump_gdn_projection_tensor("input_core_in_proj_ba", layer_name, ba) - - if self.gqa_interleaved_layout: - query, key, value, z, b, a = self.fix_query_key_value_ordering( - mixed_qkvz, - ba, + ): + mixed_qkv = torch.empty( + (1, 2560), dtype=hidden_states.dtype, device=hidden_states.device ) - query, key, value = map( - lambda x: rearrange(x, "l p d -> l (p d)"), (query, key, value) + b = torch.empty((1, 12), dtype=hidden_states.dtype, device=hidden_states.device) + a = torch.empty_like(b) + sm70_ops.fp8_qpn8_gemm_ba_split_sm70_out( + mixed_qkv, + z_out, + b, + a, + hidden_states, + self.in_proj_qkvz.weight, + self.in_proj_qkvz.weight_scale_inv, + self.in_proj_ba.weight, ) - mixed_qkv = torch.cat((query, key, value), dim=-1) + _log_runtime_route_once( + "SM70 GDN QPN8 K5120/N4096 plus FP16 b/a N24 split route enabled." + ) + z = z_out else: - qkv_size = (self.key_dim * 2 + self.value_dim) // self.tp_size - z_size = self.value_dim // self.tp_size - mixed_qkv = mixed_qkvz[..., :qkv_size] - if envs.VLLM_SM70_GDN_MIXED_QKV_CONTIGUOUS: - mixed_qkv = mixed_qkv.contiguous() - z = _sm70_compile_graph_slice_dim(mixed_qkvz, -1, qkv_size, z_size) - z = z.reshape(z.size(0), -1, self.head_v_dim) - ba_size = ba.shape[-1] // 2 - b = ba[..., :ba_size] - a = _sm70_compile_graph_slice_dim(ba, -1, ba_size, ba_size) - if self.disable_tp_for_ba_proj and self.tp_size > 1: - ba_chunk = self.num_v_heads // self.tp_size - ba_start = self.tp_rank * ba_chunk - b = b[:, ba_start : ba_start + ba_chunk] - a = a[:, ba_start : ba_start + ba_chunk] - b = b.contiguous() - a = a.contiguous() + mixed_qkvz, _ = self.in_proj_qkvz(hidden_states) + ba, _ = self.in_proj_ba(hidden_states) + mixed_qkvz = _sm70_dump_gdn_projection_tensor( + "input_core_in_proj_qkvz", + layer_name, + mixed_qkvz, + ) + ba = _sm70_dump_gdn_projection_tensor("input_core_in_proj_ba", layer_name, ba) - if envs.VLLM_SM70_GDN_Z_CONTIGUOUS and current_platform.is_device_capability(70): - z = z.contiguous() - z_out.copy_(z) + if self.gqa_interleaved_layout: + query, key, value, z, b, a = self.fix_query_key_value_ordering( + mixed_qkvz, + ba, + ) + query, key, value = map( + lambda x: rearrange(x, "l p d -> l (p d)"), + (query, key, value), + ) + mixed_qkv = torch.cat((query, key, value), dim=-1) + else: + qkv_size = (self.key_dim * 2 + self.value_dim) // self.tp_size + z_size = self.value_dim // self.tp_size + mixed_qkv = mixed_qkvz[..., :qkv_size] + if envs.VLLM_SM70_GDN_MIXED_QKV_CONTIGUOUS: + mixed_qkv = mixed_qkv.contiguous() + z = _sm70_compile_graph_slice_dim(mixed_qkvz, -1, qkv_size, z_size) + z = z.reshape(z.size(0), -1, self.head_v_dim) + ba_size = ba.shape[-1] // 2 + b = ba[..., :ba_size] + a = _sm70_compile_graph_slice_dim(ba, -1, ba_size, ba_size) + if self.disable_tp_for_ba_proj and self.tp_size > 1: + ba_chunk = self.num_v_heads // self.tp_size + ba_start = self.tp_rank * ba_chunk + b = b[:, ba_start : ba_start + ba_chunk] + a = a[:, ba_start : ba_start + ba_chunk] + b = b.contiguous() + a = a.contiguous() + + if envs.VLLM_SM70_GDN_Z_CONTIGUOUS and current_platform.is_device_capability( + 70 + ): + z = z.contiguous() + z_out.copy_(z) _qwen_gdn_run_recurrent_core( self, @@ -6784,7 +7320,29 @@ def qwen_gdn_input_projection( layer_name, hidden_states, ) + if _sm70_gdn_qpn8_ba_split_eligible( + self, + hidden_states, + z_out, + layer_name, + ): + sm70_ops.fp8_qpn8_gemm_ba_split_sm70_out( + mixed_qkv_out, + z_out, + b_out, + a_out, + hidden_states, + self.in_proj_qkvz.weight, + self.in_proj_qkvz.weight_scale_inv, + self.in_proj_ba.weight, + ) + _log_runtime_route_once( + "SM70 GDN QPN8 K5120/N4096 plus FP16 b/a N24 split route enabled." + ) + return + mixed_qkvz, _ = self.in_proj_qkvz(hidden_states) + assert self.in_proj_ba is not None ba, _ = self.in_proj_ba(hidden_states) if self.gqa_interleaved_layout: @@ -6984,18 +7542,25 @@ def fused_gdn_gating( dt_bias: torch.Tensor, beta: float = 1.0, threshold: float = 20.0, + *, + beta_dtype: torch.dtype | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """ Fused computation of g and beta for Gated Delta Net. g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias) - beta_output = b.sigmoid() + beta_output = b.sigmoid(), materialized in beta_dtype when provided. + + Speculative recurrent updates request FP32 beta to match the fused packed + decode transition. Other callers retain the historical b.dtype output. TODO maybe use torch.compile to replace this triton kernel """ batch, num_heads = a.shape seq_len = 1 grid = (batch, seq_len, triton.cdiv(num_heads, 8)) g = torch.empty(1, batch, num_heads, dtype=torch.float32, device=a.device) - beta_output = torch.empty(1, batch, num_heads, dtype=b.dtype, device=b.device) + if beta_dtype is None: + beta_dtype = b.dtype + beta_output = torch.empty(1, batch, num_heads, dtype=beta_dtype, device=b.device) fused_gdn_gating_kernel[grid]( g, beta_output, 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/compressed_tensors.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py index 406546d90c..7791dc2f1a 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py @@ -380,6 +380,12 @@ def _is_nvfp4_format(quant_args: QuantizationArgs): and is_symmetric ) + def _contains_nvfp4_weights(self) -> bool: + return any( + scheme is not None and self._is_nvfp4_format(scheme.get("weights")) + for scheme in self.target_scheme_map.values() + ) + @staticmethod def _is_mxfp4(quant_args: QuantizationArgs) -> bool: if quant_args is None: @@ -676,6 +682,7 @@ def _get_scheme_from_parts( return CompressedTensorsW8A16Fp8( weight_quant=weight_quant, is_static_input_scheme=not input_quant.dynamic, + enable_sm70_qpn8_by_default=self._contains_nvfp4_weights(), ) # note: input_quant can be None @@ -684,6 +691,7 @@ def _get_scheme_from_parts( return CompressedTensorsW8A16Fp8( weight_quant=weight_quant, is_static_input_scheme=is_static_input_scheme, + enable_sm70_qpn8_by_default=self._contains_nvfp4_weights(), ) if self._is_static_tensor_w8a8(weight_quant, input_quant): @@ -854,6 +862,19 @@ def apply( raise ValueError("A scheme must be defined for each layer") return scheme.apply_weights(layer, x, bias=bias) + def apply_fused_silu_and_mul( + self, + layer: torch.nn.Module, + x: torch.Tensor, + ) -> torch.Tensor | None: + scheme = layer.scheme + if scheme is None: + raise ValueError("A scheme must be defined for each layer") + fused_apply = getattr(scheme, "apply_fused_silu_and_mul", None) + if fused_apply is None: + return None + return fused_apply(layer, x) + class CompressedTensorsKVCacheMethod(BaseKVCacheMethod): """ 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 e740664b4e..c47e313ba9 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 @@ -5,7 +5,9 @@ import torch from torch.nn.parameter import Parameter +from vllm import _sm70_ops as sm70_ops from vllm import envs +from vllm.config import get_current_vllm_config from vllm.logger import init_logger from vllm.model_executor.kernels.linear import init_nvfp4_linear_kernel from vllm.model_executor.layers.quantization import sm70_turbomind as sm70_tm @@ -21,14 +23,96 @@ logger = init_logger(__name__) +def _is_sm70_tp4_nvfp4_gate_up(layer: torch.nn.Module) -> bool: + return bool( + getattr(layer, "tp_size", 1) == 4 + and getattr(layer, "prefix", "").rsplit(".", 1)[-1] == "gate_up_proj" + and getattr(layer, "input_size_per_partition", 0) == 5120 + and getattr(layer, "output_size_per_partition", 0) == 8704 + and getattr(layer, "logical_widths", None) == [4352, 4352] + ) + + +def _is_sm70_tp4_nvfp4_down(layer: torch.nn.Module) -> bool: + return bool( + getattr(layer, "tp_size", 1) == 4 + and getattr(layer, "prefix", "").rsplit(".", 1)[-1] == "down_proj" + and getattr(layer, "input_size_per_partition", 0) == 4352 + and getattr(layer, "output_size_per_partition", 0) == 5120 + ) + + +def _is_sm70_nvfp4_qpn4_runtime_contract() -> bool: + """Admit only the measured single-sequence, no-MTP decode contract.""" + vllm_config = get_current_vllm_config() + scheduler_config = getattr(vllm_config, "scheduler_config", None) + max_num_seqs = int(getattr(scheduler_config, "max_num_seqs", 1)) + speculative_config = getattr(vllm_config, "speculative_config", None) + return max_num_seqs == 1 and speculative_config is None + + +_SM70_NVFP4_QPN4_REQUIRED_OPS = ( + "nvfp4_qpn4_prepare_sm70", + "nvfp4_qpn4_prepare_scale_code_sm70", + "nvfp4_qpn4_dequantize_sm70_out", + "nvfp4_qpn4_prefill_sm70_out", + "nvfp4_qpn4_dispatch_sm70_out", +) + + +def _missing_sm70_nvfp4_qpn4_ops() -> list[str]: + return [ + name + for name in _SM70_NVFP4_QPN4_REQUIRED_OPS + if not hasattr(torch.ops._C, name) + ] + + __all__ = ["CompressedTensorsW4A4Fp4"] +_SM70_NVFP4_QPN2_CONFIGS = { + # (K, N, fused gated-SiLU): (split-K, independent accumulator chains) + (5120, 8704, False): (8, 2), + (5120, 8704, True): (8, 2), + (4352, 5120, False): (16, 2), +} +_SM70_NVFP4_QPN2_SHAPES = { + # Checkpoint-native packed tensors are [N, K/2]. + "gate_up_proj": (8704, 2560), + "down_proj": (5120, 2176), +} +_SM70_NVFP4_QPN2_REQUIRED_OPS = ( + "nvfp4_qpn2_prepare_sm70", + "nvfp4_qpn2_gemm_sm70_out", + "nvfp4_qpn2_gated_sm70_out", + "nvfp4_qpn2_dispatch_sm70_out", +) + + +def _is_qpn2_layer(layer: torch.nn.Module) -> bool: + if getattr(layer, "tp_size", 1) != 4: + return False + suffix = getattr(layer, "prefix", "").rsplit(".", 1)[-1] + expected = _SM70_NVFP4_QPN2_SHAPES.get(suffix) + if expected is None or tuple(layer.weight.shape) != expected: + return False + expected_n, expected_packed_k = expected + return bool( + getattr(layer, "input_size_per_partition", 0) == expected_packed_k * 2 + and getattr(layer, "output_size_per_partition", 0) == expected_n + ) + + +def _missing_qpn2_ops() -> list[str]: + return [ + name + for name in _SM70_NVFP4_QPN2_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" - ): + if envs.VLLM_USE_NVFP4_CT_EMULATIONS or envs.VLLM_NVFP4_GEMM_BACKEND == "emulation": return True from vllm.config import get_current_vllm_config_or_none @@ -153,7 +237,106 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: logger.info_once( "SM70 compressed-tensors NVFP4 TurboMind W4A16 dense path enabled." ) - sm70_tm.prepare_nvfp4_linear(layer) + is_qpn4_gate = _is_sm70_tp4_nvfp4_gate_up(layer) + is_qpn4_down = _is_sm70_tp4_nvfp4_down(layer) + qpn4_model_layer = envs.VLLM_SM70_NVFP4_QPN4 and ( + is_qpn4_down or (is_qpn4_gate and envs.VLLM_SM70_NVFP4_DENSE_GATED_SILU) + ) + qpn4_runtime = ( + _is_sm70_nvfp4_qpn4_runtime_contract() if qpn4_model_layer else False + ) + if qpn4_model_layer and not qpn4_runtime: + logger.info_once( + "The SM70 NVFP4 QPN4 route retains TurboMind unless the " + "runtime contract is max_num_seqs=1 with no MTP." + ) + if qpn4_model_layer and qpn4_runtime: + missing_ops = _missing_sm70_nvfp4_qpn4_ops() + if missing_ops: + logger.warning_once( + "The automatic SM70 NVFP4 QPN4 route is unavailable " + "in the loaded vllm._C; retaining TurboMind. Missing " + f"ops: {missing_ops}." + ) + workspace = ( + None + if missing_ops + else sm70_tm.get_nvfp4_qpn4_dense_workspace(layer.weight) + ) + if not missing_ops and workspace is not None: + sm70_tm.prepare_nvfp4_qpn4_linear( + layer, + workspace, + gated_silu=is_qpn4_gate, + ) + layer.weight = Parameter( + torch.empty(0, dtype=torch.uint8, device=layer.weight.device), + requires_grad=False, + ) + layer.weight_scale = Parameter( + torch.empty( + 0, + dtype=torch.float8_e4m3fn, + device=layer.weight_scale.device, + ), + requires_grad=False, + ) + logger.info_once( + "Memory-neutral SM70 NVFP4 QPN4 M=1 decode " + "path enabled with bounded FP16 prefill workspace." + ) + return + if not missing_ops: + logger.warning_once( + "Insufficient memory for the bounded SM70 NVFP4 QPN4 " + "prefill workspace; retaining TurboMind." + ) + use_qpn2 = bool(envs.VLLM_SM70_NVFP4_QPN2 and _is_qpn2_layer(layer)) + if use_qpn2: + missing_ops = _missing_qpn2_ops() + if missing_ops: + logger.warning_once( + "The requested SM70 NVFP4 QPN2 route is unavailable; " + f"retaining TurboMind. Missing ops: {missing_ops}." + ) + use_qpn2 = False + if use_qpn2: + qpn2_codes, qpn2_scales = sm70_ops.nvfp4_qpn2_prepare_sm70( + layer.weight.data, layer.weight_scale.data + ) + qpn2_global_scale = float(layer.weight_global_scale.item()) + + use_gated_silu = bool( + envs.VLLM_SM70_NVFP4_DENSE_GATED_SILU and is_qpn4_gate and not use_qpn2 + ) + sm70_tm.prepare_nvfp4_linear( + layer, + interleave_gated_silu=use_gated_silu, + ) + if use_qpn2: + suffix = layer.prefix.rsplit(".", 1)[-1] + k = layer.input_size_per_partition + n = layer.output_size_per_partition + split_k, nacc = _SM70_NVFP4_QPN2_CONFIGS[(k, n, False)] + layer.register_buffer( + "sm70_nvfp4_qpn2_codes", qpn2_codes, persistent=False + ) + layer.register_buffer( + "sm70_nvfp4_qpn2_scales", qpn2_scales, persistent=False + ) + layer.sm70_nvfp4_qpn2 = True + layer.sm70_nvfp4_qpn2_global_scale = qpn2_global_scale + layer.sm70_nvfp4_qpn2_split_k = split_k + layer.sm70_nvfp4_qpn2_nacc = nacc + layer.sm70_nvfp4_qpn2_gated_silu = suffix == "gate_up_proj" + logger.info_once( + "SM70 NVFP4 QPN2 M<=8 route enabled for a compatible " + "TP4 projection contract." + ) + elif use_gated_silu: + logger.info_once( + "SM70 NVFP4 TurboMind gated-SiLU single-layout path enabled." + ) layer.weight = Parameter( torch.empty(0, dtype=torch.uint8, device=layer.weight.device), requires_grad=False, @@ -180,6 +363,68 @@ def apply_weights( x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: + if getattr(layer, "sm70_nvfp4_qpn2", False): + return self._apply_qpn2(layer, x, bias, gated_silu=False) if sm70_tm.has_prepared_linear(layer): return sm70_tm.apply_prepared_linear(layer, x, bias) return self._fallback_kernel().apply_weights(layer=layer, x=x, bias=bias) + + def apply_fused_silu_and_mul( + self, + layer: torch.nn.Module, + x: torch.Tensor, + ) -> torch.Tensor | None: + if getattr(layer, "sm70_nvfp4_qpn2_gated_silu", False): + return self._apply_qpn2(layer, x, None, gated_silu=True) + if not sm70_tm.has_prepared_linear(layer): + return None + return sm70_tm.apply_prepared_fused_silu_and_mul(layer, x) + + @staticmethod + def _apply_qpn2( + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None, + *, + gated_silu: bool, + ) -> torch.Tensor: + if x.dtype != torch.float16: + raise RuntimeError( + f"SM70 NVFP4 QPN2 requires float16 activations, got {x.dtype}." + ) + x_2d = x.reshape(-1, x.shape[-1]) + if x_2d.stride(-1) != 1: + x_2d = x_2d.contiguous() + output_size = layer.output_size_per_partition + if gated_silu: + output_size //= 2 + out_2d = torch.empty( + (x_2d.shape[0], output_size), dtype=x.dtype, device=x.device + ) + if x_2d.shape[0] == 0: + 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) + if gated_silu: + split_k, nacc = _SM70_NVFP4_QPN2_CONFIGS[ + (x_2d.shape[1], output_size * 2, True) + ] + sm70_ops.nvfp4_qpn2_dispatch_sm70_out( + out_2d, + x_2d, + layer.sm70_nvfp4_qpn2_codes, + layer.sm70_nvfp4_qpn2_scales, + float(layer.sm70_nvfp4_qpn2_global_scale), + split_k, + nacc, + state.weight, + state.scales, + state.group_size, + state.k_ld, + state.q_ld, + gated_silu, + ) + if bias is not None: + out_2d.add_(bias) + return out_2d.reshape(*x.shape[:-1], output_size) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py index 61f0f5b923..2b7ca24ff1 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py @@ -1,15 +1,20 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os from collections.abc import Callable import torch from compressed_tensors.quantization import QuantizationArgs, QuantizationStrategy +from vllm import _sm70_ops as sm70_ops +from vllm import envs from vllm.config import get_current_vllm_config +from vllm.logger import init_logger from vllm.model_executor.kernels.linear import ( init_wfp8_a16_linear_kernel, ) +from vllm.model_executor.layers.quantization import sm70_turbomind as sm70_tm from vllm.model_executor.layers.quantization.compressed_tensors.schemes import ( CompressedTensorsScheme, ) @@ -17,6 +22,11 @@ STRATEGY_TO_PARAMETER_TYPE, STRATEGY_TO_WEIGHT_QUANT_KEY, ) +from vllm.model_executor.layers.quantization.fp8 import ( + _get_sm70_fp8_prefill_exact_dense_workspace, + _is_sm70_fp8_qpn8_runtime_contract, + _missing_sm70_fp8_qpn8_ops, +) from vllm.model_executor.layers.quantization.utils.fp8_utils import ( create_fp8_scale_parameter, create_fp8_weight_parameter, @@ -34,9 +44,53 @@ __all__ = ["CompressedTensorsW8A16Fp8"] +logger = init_logger(__name__) + +_SM70_CHANNEL_FP8_QPN8_SHAPES = { + "in_proj_qkvz": (4096, 5120), + "qkv_proj": (3584, 5120), + "out_proj": (5120, 1536), + "o_proj": (5120, 1536), + "gate_up_proj": (8704, 5120), + "down_proj": (5120, 4352), +} +_SM70_CHANNEL_FP8_QPN8_CONFIGS = { + # (K, N): (split-K, accumulator chains, prefetch codes) + (5120, 4096): (16, 2, False), + (5120, 3584): (16, 2, False), + (1536, 5120): (12, 2, False), + (5120, 8704): (16, 2, False), + (4352, 5120): (16, 2, False), +} +_SM70_CHANNEL_FP8_QPN8_GATED_CONFIG = (8, 2, False) + + +def _sm70_fp8_qpn8_enabled(enable_by_default: bool) -> bool: + """Resolve QPN8 while preserving the validated mixed-NVFP4 default.""" + if os.getenv("VLLM_SM70_FP8_QPN8") is None: + return enable_by_default + return envs.VLLM_SM70_FP8_QPN8 + + +def _sm70_channel_fp8_qpn8_config( + layer: torch.nn.Module, +) -> tuple[int, int, bool] | None: + if getattr(layer, "tp_size", 1) != 4: + return None + suffix = getattr(layer, "prefix", "").rsplit(".", 1)[-1] + if tuple(layer.weight.shape) != _SM70_CHANNEL_FP8_QPN8_SHAPES.get(suffix): + return None + n_dim, k_dim = (int(dim) for dim in layer.weight.shape) + return _SM70_CHANNEL_FP8_QPN8_CONFIGS.get((k_dim, n_dim)) + class CompressedTensorsW8A16Fp8(CompressedTensorsScheme): - def __init__(self, weight_quant: QuantizationArgs, is_static_input_scheme: bool): + def __init__( + self, + weight_quant: QuantizationArgs, + is_static_input_scheme: bool, + enable_sm70_qpn8_by_default: bool = False, + ): self.weight_quant = weight_quant self.strategy = weight_quant.strategy self.out_dtype = torch.get_default_dtype() @@ -48,6 +102,15 @@ def __init__(self, weight_quant: QuantizationArgs, is_static_input_scheme: bool) self.activation_quant_key = ( kFp8StaticTensorSym if is_static_input_scheme else kFp8DynamicTensorSym ) + self.use_sm70_fp8_turbomind = ( + self.strategy == QuantizationStrategy.CHANNEL + and not self.is_static_input_scheme + and sm70_tm.is_exact_sm70_cuda_platform() + and sm70_tm.use_turbomind(envs.VLLM_SM70_FP8_TURBOMIND) + ) + self.use_sm70_fp8_qpn8 = self.use_sm70_fp8_turbomind and _sm70_fp8_qpn8_enabled( + enable_sm70_qpn8_by_default + ) @classmethod def get_min_capability(cls) -> int: @@ -108,6 +171,9 @@ def create_weights( ) layer.register_parameter("input_scale", input_scale) + if self.use_sm70_fp8_turbomind: + return + self.linear_kernel = init_wfp8_a16_linear_kernel( weight_quant_key=self.weight_quant_key, activation_quant_key=self.activation_quant_key, @@ -117,6 +183,112 @@ def create_weights( ) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + if self.use_sm70_fp8_turbomind: + if layer.orig_dtype != torch.float16: + raise RuntimeError( + "SM70 TurboMind channel-FP8 requires fp16 model weights, " + f"got {layer.orig_dtype}." + ) + if not hasattr(torch.ops._C, "fp8_sm70_prepare"): + raise RuntimeError( + "VLLM_SM70_FP8_TURBOMIND=1 requires a build with CUDA " + "arch 7.0 and the SM70 TurboMind extension." + ) + weight_scale = layer.weight_scale.to(torch.float32).contiguous() + expected_scale_shape = (layer.weight.shape[0], 1) + if tuple(weight_scale.shape) != expected_scale_shape: + raise RuntimeError( + "SM70 TurboMind channel-FP8 expected weight scales with " + f"shape {expected_scale_shape}, got {tuple(weight_scale.shape)}." + ) + qpn8_config = ( + _sm70_channel_fp8_qpn8_config(layer) + if getattr(self, "use_sm70_fp8_qpn8", False) + else None + ) + qpn8_concurrency = ( + _is_sm70_fp8_qpn8_runtime_contract() + if qpn8_config is not None + else False + ) + if qpn8_config is not None and not qpn8_concurrency: + logger.info_once( + "The SM70 channel-FP8 QPN8 route retains TurboMind when " + "MTP is enabled or max_num_seqs exceeds 8." + ) + if qpn8_config is not None and qpn8_concurrency: + missing_ops = _missing_sm70_fp8_qpn8_ops() + if missing_ops and os.getenv("VLLM_SM70_FP8_QPN8") is not None: + raise RuntimeError( + "VLLM_SM70_FP8_QPN8=1 requires the source-built SM70 " + f"QPN8 extension; missing ops: {missing_ops}." + ) + workspace = ( + None + if missing_ops + else _get_sm70_fp8_prefill_exact_dense_workspace(layer.weight) + ) + if not missing_ops and workspace is not None: + qpn8_codes, qpn8_scales = sm70_ops.fp8_qpn8_prepare_sm70( + layer.weight, weight_scale + ) + replace_parameter(layer, "weight", qpn8_codes) + del layer._parameters["weight_scale"] + replace_parameter(layer, "weight_scale_inv", qpn8_scales) + split_k, nacc, prefetch = qpn8_config + layer.input_scale = None + layer.sm70_fp8_turbomind = True + layer.sm70_fp8_qpn8 = True + layer.sm70_fp8_channel_scale = True + layer.sm70_fp8_qpn8_split_k = split_k + layer.sm70_fp8_qpn8_nacc = nacc + layer.sm70_fp8_qpn8_prefetch = prefetch + layer.sm70_fp8_prefill_exact_dense_workspace_ptr = ( + workspace.data_ptr() + ) + if getattr(layer, "prefix", "").rsplit(".", 1)[-1] == ( + "gate_up_proj" + ): + gated_split_k, gated_nacc, gated_prefetch = ( + _SM70_CHANNEL_FP8_QPN8_GATED_CONFIG + ) + layer.sm70_fp8_gated_silu = True + layer.sm70_fp8_qpn8_gated_split_k = gated_split_k + layer.sm70_fp8_qpn8_gated_nacc = gated_nacc + layer.sm70_fp8_qpn8_gated_prefetch = gated_prefetch + logger.info_once( + "Memory-neutral SM70 channel-FP8 QPN8 path enabled " + "for accepted Qwen3.8-27B TP4 dense shapes." + ) + return + if missing_ops: + logger.warning_once( + "The SM70 channel-FP8 QPN8 operators are unavailable; " + "retaining the TurboMind layout." + ) + else: + logger.warning_once( + "Insufficient memory for the SM70 channel-FP8 QPN8 " + "prefill workspace; retaining the TurboMind layout." + ) + tm_weight, tm_scales, meta = sm70_ops.fp8_sm70_prepare( + layer.weight, weight_scale, 128, False + ) + replace_parameter(layer, "weight", tm_weight) + del layer._parameters["weight_scale"] + replace_parameter(layer, "weight_scale_inv", tm_scales) + layer.input_scale = None + layer.sm70_fp8_turbomind = True + layer.sm70_fp8_channel_scale = True + 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()) + logger.info_once( + "SM70 compressed-tensors channel-FP8 TurboMind W8A16 dense " + "path enabled." + ) + return + if self.strategy == QuantizationStrategy.BLOCK: assert self.is_static_input_scheme is False # MarlinFP8ScaledMMLinearKernel uses "weight_scale_inv" for block @@ -145,4 +317,83 @@ def apply_weights( x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: + if getattr(layer, "sm70_fp8_turbomind", False): + if x.dtype != torch.float16: + raise RuntimeError( + "SM70 TurboMind channel-FP8 requires float16 activations, " + f"got {x.dtype}." + ) + out_shape = (*x.shape[:-1], layer.output_size_per_partition) + x_2d = x.reshape(-1, x.shape[-1]) + if x_2d.stride(-1) != 1: + x_2d = x_2d.contiguous() + out_2d = torch.empty( + (x_2d.shape[0], layer.output_size_per_partition), + device=x.device, + dtype=x.dtype, + ) + if x_2d.shape[0] == 0: + return out_2d.reshape(out_shape) + if getattr(layer, "sm70_fp8_qpn8", False): + sm70_ops.fp8_qpn8_dispatch_sm70_out( + out_2d, + int(layer.sm70_fp8_prefill_exact_dense_workspace_ptr), + x_2d, + layer.weight, + layer.weight_scale_inv, + int(layer.sm70_fp8_qpn8_split_k), + int(layer.sm70_fp8_qpn8_nacc), + bool(layer.sm70_fp8_qpn8_prefetch), + False, + ) + else: + sm70_ops.fp8_gemm_sm70_out( + out_2d, + x_2d, + layer.weight, + layer.weight_scale_inv, + 128, + layer.sm70_fp8_k_ld, + layer.sm70_fp8_q_ld, + False, + ) + if bias is not None: + out_2d.add_(bias) + return out_2d.reshape(out_shape) return self.linear_kernel.apply_weights(layer, x, bias) + + def apply_fused_silu_and_mul( + self, + layer: torch.nn.Module, + x: torch.Tensor, + ) -> torch.Tensor | None: + if not getattr(layer, "sm70_fp8_qpn8", False): + return None + if not getattr(layer, "sm70_fp8_gated_silu", False): + return None + if x.dtype != torch.float16: + raise RuntimeError( + "SM70 channel-FP8 QPN8 gated-SiLU requires float16 " + f"activations, got {x.dtype}." + ) + x_2d = x.reshape(-1, x.shape[-1]) + if x_2d.stride(-1) != 1: + x_2d = x_2d.contiguous() + out_features = layer.output_size_per_partition // 2 + out_2d = torch.empty( + (x_2d.shape[0], out_features), device=x.device, dtype=x.dtype + ) + if x_2d.shape[0] == 0: + return out_2d.reshape(*x.shape[:-1], out_features) + sm70_ops.fp8_qpn8_dispatch_sm70_out( + out_2d, + int(layer.sm70_fp8_prefill_exact_dense_workspace_ptr), + x_2d, + layer.weight, + layer.weight_scale_inv, + int(layer.sm70_fp8_qpn8_gated_split_k), + int(layer.sm70_fp8_qpn8_gated_nacc), + bool(layer.sm70_fp8_qpn8_gated_prefetch), + True, + ) + return out_2d.reshape(*x.shape[:-1], out_features) diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 25eb4e363f..8f8c975d85 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os from typing import TYPE_CHECKING, Any import torch @@ -100,6 +101,11 @@ logger = init_logger(__name__) _SM70_FP8_PREFILL_DENSE_MIN_M = 3920 +_SM70_FP8_EXACT_8K_PREFILL_M = 8000 +_SM70_FP8_EXACT_8K_PREFILL_SHAPES = { + "in_proj_qkvz": (5120, 4096), + "qkv_proj": (5120, 3584), +} _SM70_FP8_PREFILL_DENSE_SHAPES = { "gate_up_proj": (5120, 8704), "down_proj": (4352, 5120), @@ -112,20 +118,237 @@ _SM70_FP8_PREFILL_DENSE_WORKSPACE_BYTES = ( _SM70_FP8_PREFILL_DENSE_WORKSPACE_ELEMENTS * torch.float16.itemsize ) +_SM70_FP8_QPN8_CONFIGS = { + # (K, N, fused gated-SiLU): (split-K, accumulator chains, prefetch codes) + (4352, 5120, False): (16, 1, False), + (1536, 5120, False): (16, 1, False), + (5120, 4096, False): (16, 2, False), + (5120, 3584, False): (16, 2, False), + (5120, 8704, False): (16, 1, True), + (5120, 8704, True): (8, 2, True), +} +_SM70_FP8_QPN8_EXTRA_SHAPES = { + "in_proj_qkvz": (5120, 4096), + "qkv_proj": (5120, 3584), +} +_SM70_FP8_QPN8_PP2_TP4_CONFIGS = { + # Real-weight speed winners remain numerically bounded for every admitted + # projection. The shared-expert gate/up role is deliberately absent below + # after its fused activation matched only 53.5% of TurboMind FP16 elements. + (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. + "fused_wqa_wkv": {(1, 4096, 1536)}, + "wq_b": {(4, 1024, 8192)}, + "wo_b": {(4, 2048, 4096)}, + "down_proj": {(4, 512, 4096)}, +} +_SM70_FP8_QPN8_PP2_TP4_WORKSPACE_ELEMENTS = max( + k * n for k, n, _ in _SM70_FP8_QPN8_PP2_TP4_CONFIGS +) +_SM70_FP8_QPN8_REQUIRED_OPS = ( + "fp8_qpn8_prepare_sm70", + "fp8_qpn8_dequantize_sm70_out", + "fp8_qpn8_prefill_sm70_out", + "fp8_qpn8_dispatch_sm70_out", + "fp8_qpn8_gemm_sm70_out", + "fp8_qpn8_gated_pair_sm70_out", +) +_SM70_FP8_QPN8_MAX_NUM_SEQS = 8 # Layers retain only data_ptr(), so this cache owns each allocation's lifetime. _sm70_fp8_prefill_dense_workspaces: dict[tuple[int, torch.dtype], torch.Tensor] = {} +_sm70_fp8_qpn8_pp2_tp4_workspaces: dict[tuple[int, torch.dtype], torch.Tensor] = {} + + +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 + if getattr(layer, "weight_block_size", None) != [128, 128]: + return False + suffix = getattr(layer, "prefix", "").rsplit(".", 1)[-1] + expected = _SM70_FP8_EXACT_8K_PREFILL_SHAPES.get(suffix) + return expected is not None and tuple(layer.weight.shape) == expected def _is_sm70_fp8_prefill_exact_dense_layer(layer: torch.nn.Module) -> bool: if getattr(layer, "tp_size", 1) != 4: return False + if getattr(layer, "weight_block_size", None) != [128, 128]: + return False suffix = getattr(layer, "prefix", "").rsplit(".", 1)[-1] expected = _SM70_FP8_PREFILL_DENSE_SHAPES.get(suffix) + if expected is None: + expected = _SM70_FP8_EXACT_8K_PREFILL_SHAPES.get(suffix) if expected is None: return False return tuple(layer.weight.shape) == expected +def _is_sm70_fp8_qpn8_layer(layer: torch.nn.Module) -> bool: + """Admit only shapes with an accepted bounded-workspace prefill route.""" + if getattr(layer, "tp_size", 1) != 4: + return False + if getattr(layer, "weight_block_size", None) != [128, 128]: + return False + suffix = getattr(layer, "prefix", "").rsplit(".", 1)[-1] + expected_kn = _SM70_FP8_PREFILL_DENSE_SHAPES.get(suffix) + if expected_kn is None: + expected_kn = _SM70_FP8_QPN8_EXTRA_SHAPES.get(suffix) + if expected_kn is None: + return False + # Checkpoint-native block-FP8 weights are [N, K]; the shared prefill + # workspace and QPN8 replacement parameter are [K, N]. + if tuple(reversed(layer.weight.shape)) != expected_kn: + return False + expected_k, expected_n = expected_kn + return bool( + getattr(layer, "input_size_per_partition", 0) == expected_k + and getattr(layer, "output_size_per_partition", 0) == expected_n + ) + + +def _is_sm70_fp8_qpn8_runtime_contract() -> bool: + """Admit only the no-MTP M range with a measured native decode kernel.""" + vllm_config = get_current_vllm_config() + scheduler_config = getattr(vllm_config, "scheduler_config", None) + max_num_seqs = int(getattr(scheduler_config, "max_num_seqs", 1)) + speculative_config = getattr(vllm_config, "speculative_config", None) + if max_num_seqs > _SM70_FP8_QPN8_MAX_NUM_SEQS: + return False + if speculative_config is None: + return True + # Speculative workloads require an explicit operator opt-in; admission is + # still based only on bounded concurrency and the layer tensor contract. + return os.getenv("VLLM_SM70_FP8_QPN8") is not None + + +def _sm70_fp8_qpn8_pp2_tp4_enabled() -> bool: + """Resolve the validated default while retaining explicit rollback.""" + generic_override = os.getenv("VLLM_SM70_FP8_QPN8") + if generic_override is not None and not envs.VLLM_SM70_FP8_QPN8: + return False + specific_override = os.getenv("VLLM_SM70_FP8_QPN8_PP2_TP4") + if specific_override is not None: + return envs.VLLM_SM70_FP8_QPN8_PP2_TP4 + if generic_override is not None: + return envs.VLLM_SM70_FP8_QPN8 + return envs.VLLM_SM70_FP8_QPN8_PP2_TP4 + + +def _is_sm70_fp8_qpn8_pp2_tp4_runtime_contract() -> bool: + """Require the measured serialized PP2 x TP4 single-request route.""" + 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 _sm70_fp8_qpn8_pp2_tp4_config( + layer: torch.nn.Module, *, gated_silu: bool +) -> tuple[int, int, bool] | None: + """Select by operator/tensor contract, never model or checkpoint identity.""" + if getattr(layer, "weight_block_size", None) != [128, 128]: + return None + suffix = getattr(layer, "prefix", "").rsplit(".", 1)[-1] + accepted = _SM70_FP8_QPN8_PP2_TP4_SHAPES.get(suffix) + if accepted is None: + return None + k_dim = int(getattr(layer, "input_size_per_partition", 0)) + n_dim = int(getattr(layer, "output_size_per_partition", 0)) + layer_contract = (int(getattr(layer, "tp_size", 1)), k_dim, n_dim) + if layer_contract not in accepted: + return None + if tuple(reversed(layer.weight.shape)) != (k_dim, n_dim): + return None + if gated_silu: + output_partitions = getattr(layer, "output_partition_sizes", None) + if ( + suffix != "gate_up_proj" + or not isinstance(output_partitions, list) + or output_partitions != [n_dim // 2, n_dim // 2] + ): + return None + return _SM70_FP8_QPN8_PP2_TP4_CONFIGS.get((k_dim, n_dim, gated_silu)) + + +def _sm70_fp8_qpn8_pp2_tp4_bmm_config( + layer: torch.nn.Module, +) -> tuple[int, int, bool] | None: + if ( + getattr(layer, "prefix", "").rsplit(".", 1)[-1] != "wo_a" + or getattr(layer, "weight_block_size", None) != [128, 128] + or int(getattr(layer, "tp_size", 1)) != 4 + or int(getattr(layer, "bmm_batch_size", 0)) != 2 + or int(getattr(layer, "input_size_per_partition", 0)) != 4096 + or int(getattr(layer, "output_size_per_partition", 0)) != 2048 + or tuple(layer.weight.shape) != (2048, 4096) + ): + return None + return _SM70_FP8_QPN8_PP2_TP4_CONFIGS[(4096, 1024, False)] + + +def _missing_sm70_fp8_qpn8_ops() -> list[str]: + return [ + name for name in _SM70_FP8_QPN8_REQUIRED_OPS if not hasattr(torch.ops._C, name) + ] + + def _get_sm70_fp8_prefill_exact_dense_workspace( weight: torch.Tensor, ) -> torch.Tensor | None: @@ -152,6 +375,78 @@ def _get_sm70_fp8_prefill_exact_dense_workspace( return workspace +def _get_sm70_fp8_qpn8_pp2_tp4_workspace( + weight: torch.Tensor, +) -> torch.Tensor | None: + """Allocate one bounded FP16 prefill fallback per device.""" + device_index = weight.device.index + if device_index is None: + device_index = torch.accelerator.current_device_index() + cache_key = (device_index, torch.float16) + workspace = _sm70_fp8_qpn8_pp2_tp4_workspaces.get(cache_key) + if workspace is not None: + return workspace + try: + workspace = torch.empty( + (_SM70_FP8_QPN8_PP2_TP4_WORKSPACE_ELEMENTS,), + dtype=torch.float16, + device=weight.device, + ) + except torch.OutOfMemoryError: + logger.warning_once( + "Insufficient memory for the bounded SM70 PP2 x TP4 QPN8 " + "prefill workspace; retaining TurboMind FP8." + ) + return None + _sm70_fp8_qpn8_pp2_tp4_workspaces[cache_key] = workspace + return workspace + + +def _sm70_fp8_prefill_visible_dense_mm( + input: torch.Tensor, + weight: torch.Tensor, + scales: torch.Tensor, + dense_weight_ptr: int | None, + *, + gated_silu: bool, + min_prefill_m: int, +) -> torch.Tensor | None: + """Expose the long-prefill dense MM to AsyncTP pattern matching. + + This diagnostic route intentionally keeps the accepted dequantization and + FP16 MM arithmetic while moving ``aten.mm`` out of the opaque C++ wrapper. + """ + if not envs.VLLM_SM70_FP8_PREFILL_VISIBLE_DENSE_MM: + return None + if dense_weight_ptr is None: + return None + if input.dtype != torch.float16 or input.shape[0] < min_prefill_m: + return None + + device_index = input.device.index + if device_index is None: + device_index = torch.accelerator.current_device_index() + workspace = _sm70_fp8_prefill_dense_workspaces.get((device_index, torch.float16)) + if workspace is None: + return None + if not torch.compiler.is_compiling() and workspace.data_ptr() != dense_weight_ptr: + return None + + dense_weight = workspace.narrow(0, 0, weight.numel()).view(weight.shape) + sm70_ops.fp8_sm70_dequantize_out(dense_weight, weight, scales, 128) + dense_out = torch.mm(input, dense_weight) + if not gated_silu: + return dense_out + + out = torch.empty( + (input.shape[0], dense_out.shape[1] // 2), + dtype=input.dtype, + device=input.device, + ) + sm70_ops.silu_and_mul_interleaved(out, dense_out) + return out + + class Fp8Config(QuantizationConfig): """Config class for FP8.""" @@ -534,6 +829,15 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: if weight_scale_inv.dtype != torch.float32: weight_scale_inv = weight_scale_inv.to(torch.float32) if getattr(layer, "is_bmm", False): + qpn8_bmm_config = ( + _sm70_fp8_qpn8_pp2_tp4_bmm_config(layer) + if _sm70_fp8_qpn8_pp2_tp4_enabled() + else None + ) + qpn8_bmm_runtime = bool( + qpn8_bmm_config is not None + and _is_sm70_fp8_qpn8_pp2_tp4_runtime_contract() + ) group_count = int(getattr(layer, "bmm_batch_size", 0)) if group_count <= 0 or weight.shape[0] % group_count != 0: raise RuntimeError( @@ -557,6 +861,73 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: f"{weight_scale_inv.shape[0]}." ) + if qpn8_bmm_config is not None and not qpn8_bmm_runtime: + logger.info_once( + "Grouped SM70 QPN8 retains TurboMind outside the " + "serialized PP2 x TP4 single-request contract." + ) + if qpn8_bmm_runtime: + missing_ops = _missing_sm70_fp8_qpn8_ops() + explicitly_enabled = ( + os.getenv("VLLM_SM70_FP8_QPN8") == "1" + or os.getenv("VLLM_SM70_FP8_QPN8_PP2_TP4") == "1" + ) + if missing_ops and explicitly_enabled: + raise RuntimeError( + "The explicitly enabled SM70 PP2 x TP4 QPN8 route " + f"requires source-built operators; missing: {missing_ops}." + ) + if missing_ops: + logger.warning_once( + "The default SM70 PP2 x TP4 QPN8 route is unavailable " + "in the loaded vllm._C; retaining TurboMind FP8." + ) + workspace = ( + None + if missing_ops + else _get_sm70_fp8_qpn8_pp2_tp4_workspace(weight) + ) + if workspace is not None: + qpn8_weights = [] + qpn8_scales = [] + for group_idx in range(group_count): + row_start = group_idx * rows_per_group + scale_start = group_idx * scale_rows_per_group + qpn8_weight, qpn8_scale = sm70_ops.fp8_qpn8_prepare_sm70( + weight[ + row_start : row_start + rows_per_group + ].contiguous(), + weight_scale_inv[ + scale_start : scale_start + scale_rows_per_group + ].contiguous(), + ) + qpn8_weights.append(qpn8_weight) + qpn8_scales.append(qpn8_scale) + replace_parameter(layer, "weight", torch.stack(qpn8_weights)) + replace_parameter( + layer, "weight_scale_inv", torch.stack(qpn8_scales) + ) + assert qpn8_bmm_config is not None + split_k, nacc, prefetch = qpn8_bmm_config + layer.input_scale = None + layer.sm70_fp8_turbomind = True + layer.sm70_fp8_qpn8 = True + layer.sm70_fp8_qpn8_bmm = True + layer.sm70_fp8_bmm = True + layer.sm70_fp8_bmm_groups = group_count + layer.sm70_fp8_bmm_output_size = rows_per_group + layer.sm70_fp8_qpn8_split_k = split_k + layer.sm70_fp8_qpn8_nacc = nacc + layer.sm70_fp8_qpn8_prefetch = prefetch + layer.sm70_fp8_prefill_exact_dense_workspace_ptr = ( + workspace.data_ptr() + ) + logger.info_once( + "Default SM70 grouped QPN8 enabled for the validated " + "serialized PP2 x TP4 tensor contract." + ) + return + prepared_weights = [] prepared_scales = [] metas = [] @@ -596,12 +967,182 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: layer.register_buffer("sm70_fp8_meta", first_meta, persistent=False) layer.sm70_fp8_k_ld = int(first_meta[0].item()) layer.sm70_fp8_q_ld = int(first_meta[1].item()) + if ( + envs.VLLM_SM70_FP8_GROUPED_BMM_DECODE + and group_count == 2 + and rows_per_group == 1024 + and int(layer.weight.shape[1]) == 4096 + and hasattr( + torch.ops._C, + "fp8_moe_gemm_sm70_per_expert_dispatch_out", + ) + and hasattr(torch.ops._C, "awq_moe_build_strided_ptrs") + ): + ptrs_w, ptrs_s = sm70_ops.awq_moe_build_strided_ptrs( + layer.weight, + layer.weight_scale_inv, + layer.sm70_fp8_k_ld, + layer.sm70_fp8_q_ld, + group_count, + ) + layer.register_buffer( + "sm70_fp8_bmm_grouped_ptrs_w", ptrs_w, persistent=False + ) + layer.register_buffer( + "sm70_fp8_bmm_grouped_ptrs_s", ptrs_s, persistent=False + ) + layer.register_buffer( + "sm70_fp8_bmm_grouped_offsets", + torch.arange( + group_count + 1, + dtype=torch.int32, + device=layer.weight.device, + ), + persistent=False, + ) + layer.sm70_fp8_bmm_grouped_decode = True + logger.info_once( + "SM70 FP8 one-launch grouped-BMM decode path enabled." + ) logger.info_once( "SM70 FP8 TurboMind grouped-BMM path enabled for DeepSeek V4." ) return is_gated_silu_layer = self._is_sm70_gated_silu_layer(layer) use_gated_silu = is_gated_silu_layer and envs.VLLM_SM70_FP8_DENSE_GATED_SILU + generic_qpn8_candidate = ( + envs.VLLM_SM70_FP8_QPN8 and _is_sm70_fp8_qpn8_layer(layer) + ) + pp2_tp4_qpn8_config = ( + _sm70_fp8_qpn8_pp2_tp4_config(layer, gated_silu=False) + if _sm70_fp8_qpn8_pp2_tp4_enabled() + else None + ) + pp2_tp4_qpn8_candidate = pp2_tp4_qpn8_config is not None + qpn8_candidate_layer = generic_qpn8_candidate or pp2_tp4_qpn8_candidate + qpn8_runtime = bool( + ( + pp2_tp4_qpn8_candidate + and _is_sm70_fp8_qpn8_pp2_tp4_runtime_contract() + ) + or (generic_qpn8_candidate and _is_sm70_fp8_qpn8_runtime_contract()) + ) + if qpn8_candidate_layer and not qpn8_runtime: + logger.info_once( + "The SM70 FP8 QPN8 route retains TurboMind unless its " + "bounded-concurrency runtime contract is explicit." + ) + if qpn8_candidate_layer and qpn8_runtime: + pp2_tp4_gated_config = None + if pp2_tp4_qpn8_candidate and use_gated_silu: + pp2_tp4_gated_config = _sm70_fp8_qpn8_pp2_tp4_config( + layer, gated_silu=True + ) + if pp2_tp4_gated_config is None: + raise RuntimeError( + "The SM70 PP2 x TP4 QPN8 gate/up layer violated " + "its fused-SiLU tensor contract." + ) + missing_ops = _missing_sm70_fp8_qpn8_ops() + if missing_ops: + explicitly_enabled = ( + os.getenv("VLLM_SM70_FP8_QPN8") == "1" + or os.getenv("VLLM_SM70_FP8_QPN8_PP2_TP4") == "1" + ) + if explicitly_enabled: + raise RuntimeError( + "The explicitly enabled SM70 QPN8 route requires " + f"source-built operators; missing: {missing_ops}." + ) + logger.warning_once( + "The default SM70 FP8 QPN8 route is unavailable in " + "the loaded vllm._C; retaining the TurboMind layout." + ) + + if missing_ops: + workspace = None + elif pp2_tp4_qpn8_candidate: + workspace = _get_sm70_fp8_qpn8_pp2_tp4_workspace(weight) + else: + workspace = _get_sm70_fp8_prefill_exact_dense_workspace(weight) + if not missing_ops and workspace is not None: + qpn8_codes, qpn8_scales = sm70_ops.fp8_qpn8_prepare_sm70( + weight, weight_scale_inv + ) + k_dim, n_dim = (int(dim) for dim in qpn8_codes.shape) + if pp2_tp4_qpn8_candidate: + assert pp2_tp4_qpn8_config is not None + split_k, nacc, prefetch = pp2_tp4_qpn8_config + else: + split_k, nacc, prefetch = _SM70_FP8_QPN8_CONFIGS[ + (k_dim, n_dim, False) + ] + replace_parameter(layer, "weight", qpn8_codes) + replace_parameter(layer, "weight_scale_inv", qpn8_scales) + layer.input_scale = None + layer.sm70_fp8_turbomind = True + layer.sm70_fp8_qpn8 = True + layer.sm70_fp8_qpn8_split_k = split_k + layer.sm70_fp8_qpn8_nacc = nacc + layer.sm70_fp8_qpn8_prefetch = prefetch + layer.sm70_fp8_prefill_exact_dense_workspace_ptr = ( + workspace.data_ptr() + ) + if use_gated_silu: + if pp2_tp4_qpn8_candidate: + assert pp2_tp4_gated_config is not None + gated_split_k, gated_nacc, gated_prefetch = ( + pp2_tp4_gated_config + ) + else: + gated_split_k, gated_nacc, gated_prefetch = ( + _SM70_FP8_QPN8_CONFIGS[(k_dim, n_dim, True)] + ) + layer.sm70_fp8_gated_silu = True + layer.sm70_fp8_gated_silu_primary = True + layer.sm70_fp8_qpn8_gated_split_k = gated_split_k + layer.sm70_fp8_qpn8_gated_nacc = gated_nacc + layer.sm70_fp8_qpn8_gated_prefetch = gated_prefetch + if pp2_tp4_qpn8_candidate: + logger.info_once( + "Default SM70 QPN8 enabled for the validated " + "serialized PP2 x TP4 operator contract." + ) + else: + logger.info_once( + "Memory-neutral SM70 FP8 QPN8 path enabled for " + "accepted TP4 block-FP8 operator shapes." + ) + return + if not missing_ops: + logger.warning_once( + "Insufficient memory for the SM70 FP8 QPN8 prefill " + "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, @@ -629,16 +1170,66 @@ 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 + and hasattr(torch.ops._C, "fp8_gemm_sm70_prefill_prescaled_out") + and _is_sm70_fp8_exact_8k_prefill_layer(layer) + ): + layer.register_buffer( + "sm70_fp8_prefill_prescaled_scales", + tm_scales.mul(256), + persistent=False, + ) + logger.info_once( + "SM70 block-FP8 exact-8K pre-scaled projection path enabled." + ) if ( envs.VLLM_SM70_FP8_PREFILL_EXACT_DENSE and hasattr(torch.ops._C, "fp8_gemm_sm70_prefill_dispatch_out") and _is_sm70_fp8_prefill_exact_dense_layer(layer) ): + is_exact_8k_projection = _is_sm70_fp8_exact_8k_prefill_layer(layer) workspace = _get_sm70_fp8_prefill_exact_dense_workspace(tm_weight) if workspace is not None: layer.sm70_fp8_prefill_exact_dense_workspace_ptr = ( workspace.data_ptr() ) + layer.sm70_fp8_prefill_exact_dense_min_m = ( + _SM70_FP8_EXACT_8K_PREFILL_M + if is_exact_8k_projection + else _SM70_FP8_PREFILL_DENSE_MIN_M + ) logger.info_once( "SM70 FP8 exact-dense prefill path enabled with a bounded " "85 MiB workspace." @@ -758,6 +1349,77 @@ def apply( bias: torch.Tensor | None = None, ) -> torch.Tensor: if getattr(layer, "sm70_fp8_turbomind", False): + if getattr(layer, "sm70_fp8_qpn8", False): + if x.dtype != torch.float16: + raise RuntimeError( + "SM70 FP8 QPN8 currently requires float16 activations, " + f"got {x.dtype}." + ) + if getattr(layer, "sm70_fp8_qpn8_bmm", False): + group_count = int(layer.sm70_fp8_bmm_groups) + output_size = int(layer.sm70_fp8_bmm_output_size) + if x.ndim < 2 or x.shape[-2] != group_count: + raise RuntimeError( + "SM70 grouped QPN8 input must end in [groups, K], got " + f"{tuple(x.shape)} for groups={group_count}." + ) + x_grouped = x.reshape(-1, group_count, x.shape[-1]) + x_by_group = x_grouped.transpose(0, 1).contiguous() + out_by_group = torch.empty( + (group_count, x_grouped.shape[0], output_size), + device=x.device, + dtype=x.dtype, + ) + if x_grouped.shape[0] == 0: + return out_by_group.transpose(0, 1).reshape( + *x.shape[:-2], group_count, output_size + ) + for group_idx in range(group_count): + sm70_ops.fp8_qpn8_dispatch_sm70_out( + out_by_group[group_idx], + int(layer.sm70_fp8_prefill_exact_dense_workspace_ptr), + x_by_group[group_idx], + layer.weight[group_idx], + layer.weight_scale_inv[group_idx], + int(layer.sm70_fp8_qpn8_split_k), + int(layer.sm70_fp8_qpn8_nacc), + bool(layer.sm70_fp8_qpn8_prefetch), + False, + ) + out = out_by_group.transpose(0, 1).reshape( + *x.shape[:-2], group_count, output_size + ) + if bias is not None: + out.add_(bias.view(group_count, output_size)) + return out + + out_shape = (*x.shape[:-1], layer.output_size_per_partition) + x_2d = x.reshape(-1, x.shape[-1]) + if x_2d.stride(-1) != 1: + x_2d = x_2d.contiguous() + out_2d = torch.empty( + (x_2d.shape[0], layer.output_size_per_partition), + device=x.device, + dtype=x.dtype, + ) + if x_2d.shape[0] == 0: + return out_2d.reshape(out_shape) + sm70_ops.fp8_qpn8_dispatch_sm70_out( + out_2d, + int(layer.sm70_fp8_prefill_exact_dense_workspace_ptr), + x_2d, + layer.weight, + layer.weight_scale_inv, + int(layer.sm70_fp8_qpn8_split_k), + int(layer.sm70_fp8_qpn8_nacc), + bool(layer.sm70_fp8_qpn8_prefetch), + False, + ) + out = out_2d.reshape(out_shape) + if bias is not None: + out.add_(bias) + return out + if getattr(layer, "sm70_fp8_bmm", False): group_count = int(layer.sm70_fp8_bmm_groups) output_size = int(layer.sm70_fp8_bmm_output_size) @@ -773,17 +1435,34 @@ def apply( device=x.device, dtype=x.dtype, ) - for group_idx in range(group_count): - sm70_ops.fp8_gemm_sm70_out( - out_by_group[group_idx], - x_by_group[group_idx], - layer.weight[group_idx], - layer.weight_scale_inv[group_idx], + if ( + getattr(layer, "sm70_fp8_bmm_grouped_decode", False) + and x_grouped.shape[0] == 1 + ): + sm70_ops.fp8_moe_gemm_sm70_per_expert_dispatch_out( + out_by_group.reshape(group_count, output_size), + x_by_group.reshape(group_count, x.shape[-1]), + layer.sm70_fp8_bmm_grouped_offsets, + layer.sm70_fp8_bmm_grouped_ptrs_w, + layer.sm70_fp8_bmm_grouped_ptrs_s, + group_count, + x.shape[-1], + output_size, 128, - layer.sm70_fp8_k_ld, - layer.sm70_fp8_q_ld, False, ) + else: + for group_idx in range(group_count): + sm70_ops.fp8_gemm_sm70_out( + out_by_group[group_idx], + x_by_group[group_idx], + layer.weight[group_idx], + layer.weight_scale_inv[group_idx], + 128, + layer.sm70_fp8_k_ld, + layer.sm70_fp8_q_ld, + False, + ) out = out_by_group.transpose(0, 1).reshape( *x.shape[:-2], group_count, output_size ) @@ -803,7 +1482,42 @@ def apply( prefill_workspace_ptr = getattr( layer, "sm70_fp8_prefill_exact_dense_workspace_ptr", None ) - if prefill_workspace_ptr is not None and x_2d.dtype == torch.float16: + 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", + _SM70_FP8_PREFILL_DENSE_MIN_M, + ) + visible_dense_out = _sm70_fp8_prefill_visible_dense_mm( + x_2d, + layer.weight, + layer.weight_scale_inv, + prefill_workspace_ptr, + gated_silu=False, + min_prefill_m=prefill_min_m, + ) + 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( out_2d, prefill_workspace_ptr, @@ -814,7 +1528,22 @@ def apply( layer.sm70_fp8_k_ld, layer.sm70_fp8_q_ld, False, - _SM70_FP8_PREFILL_DENSE_MIN_M, + prefill_min_m, + ) + elif ( + prefill_prescaled_scales is not None + and envs.VLLM_SM70_FP8_PREFILL_FAST_SELECTOR + and envs.VLLM_SM70_FP8_PREFILL_PRESCALED + and x_2d.shape[0] == _SM70_FP8_EXACT_8K_PREFILL_M + ): + sm70_ops.fp8_gemm_sm70_prefill_prescaled_out( + out_2d, + x_2d, + layer.weight, + prefill_prescaled_scales, + 128, + layer.sm70_fp8_k_ld, + layer.sm70_fp8_q_ld, ) else: sm70_ops.fp8_gemm_sm70_out( @@ -889,6 +1618,36 @@ def apply_fused_silu_and_mul( layer: torch.nn.Module, x: torch.Tensor, ) -> torch.Tensor | None: + if getattr(layer, "sm70_fp8_qpn8", False): + if not getattr(layer, "sm70_fp8_gated_silu", False): + return None + if x.dtype != torch.float16: + raise RuntimeError( + "SM70 FP8 QPN8 gated-SiLU requires float16 activations, " + f"got {x.dtype}." + ) + x_2d = x.reshape(-1, x.shape[-1]) + if x_2d.stride(-1) != 1: + x_2d = x_2d.contiguous() + out_features = layer.output_size_per_partition // 2 + out_2d = torch.empty( + (x_2d.shape[0], out_features), device=x.device, dtype=x.dtype + ) + if x_2d.shape[0] == 0: + return out_2d.reshape(*x.shape[:-1], out_features) + sm70_ops.fp8_qpn8_dispatch_sm70_out( + out_2d, + int(layer.sm70_fp8_prefill_exact_dense_workspace_ptr), + x_2d, + layer.weight, + layer.weight_scale_inv, + int(layer.sm70_fp8_qpn8_gated_split_k), + int(layer.sm70_fp8_qpn8_gated_nacc), + bool(layer.sm70_fp8_qpn8_gated_prefetch), + True, + ) + return out_2d.reshape(*x.shape[:-1], out_features) + if not getattr(layer, "sm70_fp8_gated_silu", False): return None if not getattr(layer, "sm70_fp8_turbomind", False): @@ -917,6 +1676,25 @@ def apply_fused_silu_and_mul( layer, "sm70_fp8_prefill_exact_dense_workspace_ptr", None ) if prefill_workspace_ptr is not None and x_2d.dtype == torch.float16: + visible_dense_out = _sm70_fp8_prefill_visible_dense_mm( + x_2d, + weight, + scales, + prefill_workspace_ptr, + gated_silu=True, + min_prefill_m=getattr( + layer, + "sm70_fp8_prefill_exact_dense_min_m", + _SM70_FP8_PREFILL_DENSE_MIN_M, + ), + ) + if visible_dense_out is not None: + return visible_dense_out.reshape(*x.shape[:-1], out_features) + min_prefill_m = getattr( + layer, + "sm70_fp8_prefill_exact_dense_min_m", + _SM70_FP8_PREFILL_DENSE_MIN_M, + ) sm70_ops.fp8_gemm_sm70_prefill_dispatch_out( out_2d, prefill_workspace_ptr, @@ -927,7 +1705,7 @@ def apply_fused_silu_and_mul( k_ld, q_ld, True, - _SM70_FP8_PREFILL_DENSE_MIN_M, + min_prefill_m, ) return out_2d.reshape(*x.shape[:-1], out_features) sm70_ops.fp8_gemm_sm70_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 eabaf62be7..09b3049800 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -8,6 +8,7 @@ from torch.nn.parameter import Parameter import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm import envs from vllm.config import get_current_vllm_config from vllm.logger import init_logger from vllm.model_executor.kernels.linear import ( @@ -48,6 +49,7 @@ UnquantizedLinearMethod, ) from vllm.model_executor.layers.quantization import QuantizationMethods +from vllm.model_executor.layers.quantization import sm70_turbomind as sm70_tm from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, QuantizeMethodBase, @@ -451,6 +453,10 @@ def __init__(self, quant_config: ModelOptFp8Config) -> None: self.quant_config = quant_config self.out_dtype = torch.get_default_dtype() self.input_dtype = get_current_vllm_config().model_config.dtype + self.use_sm70_fp8_turbomind = ( + sm70_tm.is_exact_sm70_cuda_platform() + and sm70_tm.use_turbomind(envs.VLLM_SM70_FP8_TURBOMIND) + ) def create_weights( self, @@ -500,6 +506,9 @@ def create_weights( scale[:] = torch.finfo(torch.float32).min layer.register_parameter("input_scale", scale) + if self.use_sm70_fp8_turbomind: + return + self.fp8_linear = init_fp8_linear_kernel( activation_quant_key=kFp8StaticTensorSym, weight_quant_key=kFp8StaticTensorSym, @@ -516,6 +525,40 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: max_w_scale, weight = requantize_with_max_scale( layer.weight, layer.weight_scale, layer.logical_widths ) + + if self.use_sm70_fp8_turbomind: + if not hasattr(torch.ops._C, "fp8_sm70_prepare"): + raise RuntimeError( + "ModelOpt FP8 on SM70 requires the TurboMind extension " + "with fp8_sm70_prepare." + ) + if self.input_dtype != torch.float16: + raise RuntimeError( + "ModelOpt FP8 TurboMind on SM70 requires FP16 activations, " + f"got {self.input_dtype}." + ) + + from vllm import _sm70_ops as sm70_ops + + block_size = 128 + out_blocks = (weight.shape[0] + block_size - 1) // block_size + in_blocks = (weight.shape[1] + block_size - 1) // block_size + block_scales = ( + max_w_scale.reshape(1, 1).expand(out_blocks, in_blocks).contiguous() + ) + tm_weight, tm_scales, meta = sm70_ops.fp8_sm70_prepare( + weight, block_scales, block_size + ) + replace_parameter(layer, "weight", tm_weight) + replace_parameter(layer, "weight_scale", tm_scales) + layer.input_scale = None + layer.sm70_modelopt_fp8_turbomind = True + layer.sm70_modelopt_fp8_group_size = block_size + layer.sm70_modelopt_fp8_k_ld = int(meta[0].item()) + layer.sm70_modelopt_fp8_q_ld = int(meta[1].item()) + logger.info_once("SM70 ModelOpt FP8 TurboMind W8A16 dense path enabled.") + return + layer.weight = Parameter(weight.t(), requires_grad=False) layer.weight_scale = Parameter(max_w_scale, requires_grad=False) layer.input_scale = Parameter(layer.input_scale.max(), requires_grad=False) @@ -527,6 +570,31 @@ def apply( x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: + if getattr(layer, "sm70_modelopt_fp8_turbomind", False): + if x.dtype != torch.float16: + raise TypeError( + "ModelOpt FP8 TurboMind on SM70 requires FP16 activations." + ) + from vllm import _sm70_ops as sm70_ops + + x_2d = x.reshape(-1, x.shape[-1]) + out = torch.empty( + (x_2d.shape[0], layer.output_size_per_partition), + dtype=x.dtype, + device=x.device, + ) + sm70_ops.fp8_gemm_sm70_out( + out, + x_2d, + layer.weight, + layer.weight_scale, + layer.sm70_modelopt_fp8_group_size, + layer.sm70_modelopt_fp8_k_ld, + layer.sm70_modelopt_fp8_q_ld, + ) + if bias is not None: + out.add_(bias) + return out.reshape(*x.shape[:-1], layer.output_size_per_partition) return self.fp8_linear.apply_weights(layer, x, bias) @@ -1002,6 +1070,51 @@ def apply( ModelOptFp8Config.KVCacheMethodCls = ModelOptKVCacheMethod +def _explicit_nvfp4_emulation_requested() -> bool: + """True when an explicit NVFP4 software-emulation backend is selected. + + Matches compressed-tensors W4A4 NVFP4: emulation is a proven software + path, so SM70 may load. Cutlass/FlashInfer are not. + """ + if envs.VLLM_USE_NVFP4_CT_EMULATIONS or envs.VLLM_NVFP4_GEMM_BACKEND == "emulation": + return True + + from vllm.config import get_current_vllm_config_or_none + + vllm_config = get_current_vllm_config_or_none() + return ( + vllm_config is not None + and vllm_config.kernel_config.linear_backend == "emulation" + ) + + +def _try_prepare_sm70_modelopt_nvfp4(layer: torch.nn.Module) -> bool: + """Prepare exact-SM70 TurboMind NVFP4 weights for a ModelOpt linear. + + ModelOpt stores ``weight_scale_2`` / ``weight_global_scale`` as + ``amax / (6 * 448)`` (the Marlin multiplier). TurboMind combine is + ``block * global``. Do not infer convention from scale magnitude. + """ + if not sm70_tm.should_prepare_turbomind( + layer.weight, envs.VLLM_SM70_NVFP4_TURBOMIND + ): + return False + logger.info_once( + "SM70 ModelOpt NVFP4 TurboMind dense path enabled " + "(weight-only; activations remain half)." + ) + sm70_tm.prepare_nvfp4_linear(layer) + layer.weight = Parameter( + torch.empty(0, dtype=torch.uint8, device=layer.weight.device), + requires_grad=False, + ) + layer.weight_scale = Parameter( + torch.empty(0, dtype=torch.float8_e4m3fn, device=layer.weight_scale.device), + requires_grad=False, + ) + return True + + class ModelOptNvFp4Config(ModelOptQuantConfigBase): """Config class for ModelOpt FP4.""" @@ -1050,6 +1163,18 @@ def get_supported_act_dtypes(self) -> list[torch.dtype]: @classmethod def get_min_capability(cls) -> int: + # Do not unconditionally lower the class-wide gate to 70. W4A4 + # (quant_algo=NVFP4) and W4A16_NVFP4 share this config; SM70 is + # admitted only when a proven software backend is selected. + # Exact-device routing still happens later via + # sm70_tm.should_prepare_turbomind (capability == (7, 0)). + if ( + sm70_tm.use_turbomind(envs.VLLM_SM70_NVFP4_TURBOMIND) + or sm70_tm.forces_marlin() + ): + return 70 + if _explicit_nvfp4_emulation_requested(): + return 70 return 75 @classmethod @@ -1220,6 +1345,9 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: (1.0 / layer.input_global_scale).to(torch.float32), requires_grad=False ) + if _try_prepare_sm70_modelopt_nvfp4(layer): + return + # Convert layer to NVFP4 linear kernel format self.kernel.process_weights_after_loading(layer) @@ -1229,6 +1357,8 @@ def apply( x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: + if sm70_tm.has_prepared_linear(layer): + return sm70_tm.apply_prepared_linear(layer, x, bias) return self.kernel.apply_weights(layer=layer, x=x, bias=bias) @@ -1360,13 +1490,16 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # Rename weight_scale_2 -> weight_global_scale. NO reciprocation: # ModelOpt already stores amax/2688, which is exactly what Marlin - # consumes via nvfp4_marlin_process_global_scale (called inside the - # Marlin adapter's process_weights_after_loading). + # and SM70 TurboMind consume. This is checkpoint-format provenance, + # not a magnitude heuristic (see PR 166 rejection). layer.weight_global_scale = Parameter( layer.weight_scale_2.max().to(torch.float32), requires_grad=False ) del layer.weight_scale_2 + if _try_prepare_sm70_modelopt_nvfp4(layer): + return + self.kernel.process_weights_after_loading(layer) def apply( @@ -1375,6 +1508,8 @@ def apply( x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: + if sm70_tm.has_prepared_linear(layer): + return sm70_tm.apply_prepared_linear(layer, x, bias) return self.kernel.apply_weights(layer=layer, x=x, bias=bias) @@ -2217,6 +2352,12 @@ def get_supported_act_dtypes(self) -> list[torch.dtype]: @classmethod def get_min_capability(cls) -> int: + if ( + sm70_tm.is_exact_sm70_cuda_platform() + and sm70_tm.use_turbomind(envs.VLLM_SM70_FP8_TURBOMIND) + and sm70_tm.use_turbomind(envs.VLLM_SM70_NVFP4_TURBOMIND) + ): + return 70 return 89 @classmethod @@ -2412,6 +2553,20 @@ def get_quant_method( moe_config=layer.moe_config, ) if quant_algo == "W4A16_NVFP4": + if sm70_tm.is_exact_sm70_cuda_platform(): + if not sm70_tm.should_use_nvfp4_moe_turbomind(): + raise NotImplementedError( + "ModelOpt W4A16 NVFP4 MoE on SM70 requires the " + "TurboMind backend." + ) + from vllm.model_executor.layers.quantization.nvfp4_sm70_moe import ( + ModelOptNvFp4SM70MoEMethod, + ) + + return ModelOptNvFp4SM70MoEMethod( + quant_config=self.w4a16_nvfp4_config, + moe_config=layer.moe_config, + ) return ModelOptNvFp4FusedMoE( quant_config=self.w4a16_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 6e255a5150..dc39f8426a 100644 --- a/vllm/model_executor/layers/quantization/mxfp4_sm70_moe.py +++ b/vllm/model_executor/layers/quantization/mxfp4_sm70_moe.py @@ -64,6 +64,27 @@ def _mxfp4_grouped_m8_enabled() -> bool: return bool(envs.VLLM_SM70_MXFP4_MOE_GROUPED_M8) +def _mxfp4_grouped_verifier_enabled() -> bool: + return bool(envs.VLLM_SM70_MXFP4_MOE_GROUPED_VERIFIER) + + +def _mxfp4_grouped_verifier_for_tokens(num_tokens: int) -> bool: + return bool( + (num_tokens == _GRAPH_SAFE_MAX_TOKENS and _mxfp4_grouped_m8_enabled()) + or ( + 1 < num_tokens <= _GRAPH_SAFE_MAX_TOKENS + and _mxfp4_grouped_verifier_enabled() + ) + ) + + +def _mxfp4_grouped_m8_expert_rows_enabled() -> bool: + return bool( + (_mxfp4_grouped_m8_enabled() or _mxfp4_grouped_verifier_enabled()) + and envs.VLLM_SM70_MXFP4_MOE_GROUPED_M8_EXPERT_ROWS + ) + + @triton.jit def _compact_sorted_experts_kernel( sorted_expert_ids_ptr, @@ -140,7 +161,13 @@ def _select_mxfp4_stage_dispatch( # tail entries as zero-row experts, avoiding a host readback of the # dynamic unique-expert count. graph_expert_slots = num_tokens * _DEEPSEEK_V4_FLASH_TOP_K - if num_tokens == _GRAPH_SAFE_MAX_TOKENS and _mxfp4_grouped_m8_enabled(): + if _mxfp4_grouped_verifier_for_tokens(num_tokens): + if _mxfp4_grouped_m8_expert_rows_enabled(): + return ( + buffers["compact_expert_offsets"], + buffers["active_expert_ids"], + graph_expert_slots, + ) return ( buffers["slot_expert_offsets"], buffers["permuted_experts_id"], @@ -158,6 +185,18 @@ def _select_mxfp4_stage_dispatch( return buffers["expert_offsets"], buffers["dense_expert_ids"], num_experts +def _select_mxfp4_direct_order_offsets( + buffers: dict[str, torch.Tensor], +) -> torch.Tensor: + """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 one row per route, + so it must use the separately maintained slot offsets instead. + """ + return buffers["slot_expert_offsets"] + + def validate_mxfp4_sm70_moe_contract( *, global_num_experts: int, @@ -316,6 +355,11 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: ) if envs.VLLM_SM70_MXFP4_MOE_DIRECT_TOP6_DECODE: required_ops += ("mxfp4_moe_single_token_prepare_w13_sm70_out",) + if ( + envs.VLLM_SM70_MXFP4_MOE_DIRECT_TOP6_DECODE + and envs.VLLM_SM70_MXFP4_MOE_DIRECT_ORDER_DECODE + ): + required_ops += ("awq_moe_single_token_weighted_reduce_out",) missing_ops = [name for name in required_ops if not hasattr(torch.ops._C, name)] if missing_ops: raise RuntimeError( @@ -699,6 +743,49 @@ def apply( and layer.expert_map is None and layer.local_num_experts == layer.global_num_experts ) + direct_order = ( + direct_top6 + and envs.VLLM_SM70_MXFP4_MOE_DIRECT_ORDER_DECODE + and topk_ids.dtype == torch.int32 + and topk_ids.is_contiguous() + ) + if direct_order: + route_ids = topk_ids.view(-1) + route_offsets = _select_mxfp4_direct_order_offsets(buffers) + sm70_ops.mxfp4_moe_dense_stage_sm70_out( + buffers["gate_up"], + x, + route_offsets, + route_ids, + layer.w13_strided_ptrs_w, + layer.w13_strided_ptrs_s, + _DEEPSEEK_V4_FLASH_TOP_K, + layer.sm70_mxfp4_w13_k_dim, + layer.sm70_mxfp4_w13_n_dim, + layer.sm70_mxfp4_group_size, + ) + self._apply_swiglu(layer, buffers["intermediate"], buffers["gate_up"]) + sm70_ops.mxfp4_moe_dense_stage_sm70_out( + buffers["sorted_output"], + buffers["intermediate"], + route_offsets, + route_ids, + layer.w2_strided_ptrs_w, + layer.w2_strided_ptrs_s, + _DEEPSEEK_V4_FLASH_TOP_K, + layer.sm70_mxfp4_w2_k_dim, + layer.sm70_mxfp4_w2_n_dim, + layer.sm70_mxfp4_group_size, + ) + 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, + ) + return output if direct_top6: sm70_ops.mxfp4_moe_single_token_prepare_w13_sm70_out( buffers["gate_up"], @@ -767,7 +854,8 @@ def apply( num_tokens > 1 and num_tokens <= _mxfp4_active_expert_max_tokens() and not ( - num_tokens == _GRAPH_SAFE_MAX_TOKENS and _mxfp4_grouped_m8_enabled() + _mxfp4_grouped_verifier_for_tokens(num_tokens) + and not _mxfp4_grouped_m8_expert_rows_enabled() ) and layer.expert_map is None and layer.local_num_experts == layer.global_num_experts diff --git a/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py b/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py new file mode 100644 index 0000000000..7eb31d9ce3 --- /dev/null +++ b/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py @@ -0,0 +1,635 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Native SM70 TurboMind NVFP4 MoE for Qwen3.6-35B-A3B. + +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 +load time, repacks both tensors for TurboMind, and never materializes an FP16 +expert-weight copy. +""" + +from __future__ import annotations + +from typing import Final + +import torch +from torch.nn import Parameter + +from vllm import _sm70_ops as sm70_ops +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe import ( + FusedMoEConfig, + FusedMoEMethodBase, + FusedMoEQuantConfig, + MoEActivation, + RoutedExperts, + SharedExperts, +) +from vllm.model_executor.layers.quantization.modelopt import ( + ModelOptNvFp4Config, + ModelOptNvFp4FusedMoE, +) +from vllm.model_executor.layers.quantization.sm70_turbomind import ( + NVFP4_GROUP_SIZE, + is_exact_sm70_cuda, + unpack_mxfp4_weight, +) +from vllm.triton_utils import tl, triton + +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) +_GRAPH_SAFE_MAX_TOKENS: Final = 18 +_COMPACT_GROUPED_MAX_TOKENS: Final = 10 + + +@triton.jit +def _prepare_compact_slot_groups_kernel( + sorted_expert_ids_ptr, + compact_offsets_ptr, + active_expert_ids_ptr, + TOTAL_SLOTS: tl.constexpr, + BLOCK: tl.constexpr, +): + offsets = tl.arange(0, BLOCK) + valid = offsets < TOTAL_SLOTS + expert_ids = tl.load( + sorted_expert_ids_ptr + offsets, + mask=valid, + other=-1, + ) + tl.store( + compact_offsets_ptr + offsets, + offsets, + mask=offsets <= TOTAL_SLOTS, + ) + tl.store( + active_expert_ids_ptr + offsets, + expert_ids, + mask=valid, + ) + + +def _prepare_compact_slot_groups( + sorted_expert_ids: torch.Tensor, + compact_offsets: torch.Tensor, + active_expert_ids: torch.Tensor, +) -> None: + total_slots = sorted_expert_ids.numel() + max_slots = _COMPACT_GROUPED_MAX_TOKENS * _QWEN36_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) + # TurboMind's compact grouped dispatch forces one row per group. Keep each + # routed slot independent even when adjacent slots select the same expert; + # coalescing duplicate expert IDs would make the forced one-row scheduler + # silently skip or miscompute the additional rows. + _prepare_compact_slot_groups_kernel[(1,)]( + sorted_expert_ids, + compact_offsets, + active_expert_ids, + TOTAL_SLOTS=total_slots, + BLOCK=block, + num_warps=1, + ) + + +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: + raise NotImplementedError( + "SM70 TurboMind NVFP4 MoE currently supports tensor parallel " + f"sizes {_QWEN36_SUPPORTED_TP_SIZES}, got {moe.tp_size}." + ) + local_intermediate = moe.intermediate_size_per_partition + if local_intermediate <= 0 or local_intermediate % NVFP4_GROUP_SIZE: + raise NotImplementedError( + "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: + 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}." + ) + if moe.moe_parallel_config.use_all2all_kernels: + raise NotImplementedError( + "SM70 TurboMind NVFP4 MoE does not support DP+EP all-to-all." + ) + + +def _validate_weight_layout(layer: RoutedExperts) -> None: + local_experts = int(layer.local_num_experts) + hidden = int(layer.moe_config.hidden_dim) + intermediate = int(layer.moe_config.intermediate_size_per_partition) + expected = { + "w13_weight": (local_experts, 2 * intermediate, hidden // 2), + "w13_weight_scale": ( + local_experts, + 2 * intermediate, + hidden // NVFP4_GROUP_SIZE, + ), + "w13_weight_scale_2": (local_experts, 2), + "w2_weight": (local_experts, hidden, intermediate // 2), + "w2_weight_scale": ( + local_experts, + hidden, + intermediate // NVFP4_GROUP_SIZE, + ), + "w2_weight_scale_2": (local_experts,), + } + tensors = {name: getattr(layer, name) for name in expected} + for name, shape in expected.items(): + if tuple(tensors[name].shape) != shape: + raise ValueError( + f"SM70 NVFP4 MoE layout mismatch for {name}: " + f"expected {shape}, got {tuple(tensors[name].shape)}." + ) + if layer.w13_weight.dtype != torch.uint8 or layer.w2_weight.dtype != torch.uint8: + raise TypeError("SM70 NVFP4 MoE requires packed uint8 expert weights.") + if ( + layer.w13_weight_scale.dtype != torch.float8_e4m3fn + or layer.w2_weight_scale.dtype != torch.float8_e4m3fn + ): + raise TypeError("SM70 NVFP4 MoE requires FP8 E4M3 block scales.") + + +class ModelOptNvFp4SM70MoEMethod(ModelOptNvFp4FusedMoE): + """Qwen3.6 ModelOpt W4A16_NVFP4 experts on native TurboMind SM70.""" + + def __init__( + self, + quant_config: ModelOptNvFp4Config, + moe_config: FusedMoEConfig, + ) -> None: + FusedMoEMethodBase.__init__(self, moe_config) + if quant_config.quant_method != "W4A16_NVFP4": + raise NotImplementedError( + "SM70 TurboMind ModelOpt NVFP4 MoE currently requires " + "W4A16_NVFP4 checkpoint weights." + ) + self.quant_config = quant_config + self.use_a16 = True + self.use_global_sf = False + validate_nvfp4_sm70_moe_contract(moe_config) + + @property + def supports_eplb(self) -> bool: + return False + + def maybe_make_prepare_finalize( + self, + routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, + ) -> None: + # This method owns routing, TurboMind expert GEMMs, and unpermutation. + # Do not wrap it in the generic ModelOpt modular-kernel path. + del routing_tables + return None + + def process_weights_after_loading(self, layer: RoutedExperts) -> None: + required_ops = ( + "nvfp4_sm70_prepare", + "nvfp4_moe_dense_stage_sm70_out", + "awq_moe_build_strided_ptrs", + ) + missing = [name for name in required_ops if not hasattr(torch.ops._C, name)] + if missing: + raise RuntimeError( + "Qwen3.6 NVFP4 MoE on SM70 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." + ) + if self.moe.has_bias: + raise NotImplementedError("SM70 NVFP4 MoE does not support expert bias.") + if layer.activation != MoEActivation.SILU: + raise NotImplementedError( + "SM70 NVFP4 MoE currently supports SwiGLU/SILU only." + ) + if layer.apply_router_weight_on_input: + raise NotImplementedError( + "SM70 NVFP4 MoE does not support router weights on input." + ) + if layer.expert_map is not None: + raise NotImplementedError( + "SM70 NVFP4 MoE currently requires fully replicated experts." + ) + if layer.local_num_experts != layer.global_num_experts: + raise NotImplementedError( + "SM70 NVFP4 MoE currently requires local and global experts to match." + ) + + validate_nvfp4_sm70_moe_contract(layer.moe_config) + _validate_weight_layout(layer) + num_experts = int(layer.local_num_experts) + hidden = int(layer.moe_config.hidden_dim) + intermediate = int(layer.moe_config.intermediate_size_per_partition) + + w13_tm_weights: list[torch.Tensor] = [] + w13_tm_scales: list[torch.Tensor] = [] + w13_meta: list[torch.Tensor] = [] + w2_tm_weights: list[torch.Tensor] = [] + w2_tm_scales: list[torch.Tensor] = [] + w2_meta: list[torch.Tensor] = [] + for expert_id in range(num_experts): + w13_packed = unpack_mxfp4_weight(layer.w13_weight[expert_id].data) + w13_scales = layer.w13_weight_scale[expert_id].float().clone() + w13_global = layer.w13_weight_scale_2[expert_id].float() + w13_scales[:intermediate].mul_(w13_global[0]) + w13_scales[intermediate:].mul_(w13_global[1]) + prepared_w13 = sm70_ops.nvfp4_sm70_prepare( + w13_packed, + w13_scales.half().t().contiguous(), + NVFP4_GROUP_SIZE, + ) + w13_tm_weights.append(prepared_w13[0]) + w13_tm_scales.append(prepared_w13[1]) + w13_meta.append(prepared_w13[2]) + + w2_packed = unpack_mxfp4_weight(layer.w2_weight[expert_id].data) + w2_scales = ( + layer.w2_weight_scale[expert_id].float() + * layer.w2_weight_scale_2[expert_id].float() + ) + prepared_w2 = sm70_ops.nvfp4_sm70_prepare( + w2_packed, + w2_scales.half().t().contiguous(), + NVFP4_GROUP_SIZE, + ) + w2_tm_weights.append(prepared_w2[0]) + w2_tm_scales.append(prepared_w2[1]) + w2_meta.append(prepared_w2[2]) + + layer.w13_tm_weight = Parameter( + torch.stack(w13_tm_weights), requires_grad=False + ) + layer.w13_tm_scales = Parameter(torch.stack(w13_tm_scales), requires_grad=False) + layer.w2_tm_weight = Parameter(torch.stack(w2_tm_weights), requires_grad=False) + layer.w2_tm_scales = Parameter(torch.stack(w2_tm_scales), requires_grad=False) + + w13_k_ld = int(w13_meta[0][0].item()) + 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()) + w13_ptrs = sm70_ops.awq_moe_build_strided_ptrs( + layer.w13_tm_weight, + layer.w13_tm_scales, + w13_k_ld, + w13_q_ld, + num_experts, + ) + w2_ptrs = sm70_ops.awq_moe_build_strided_ptrs( + layer.w2_tm_weight, + layer.w2_tm_scales, + w2_k_ld, + w2_q_ld, + num_experts, + ) + layer.w13_strided_ptrs_w = Parameter(w13_ptrs[0], requires_grad=False) + layer.w13_strided_ptrs_s = Parameter(w13_ptrs[1], requires_grad=False) + layer.w2_strided_ptrs_w = Parameter(w2_ptrs[0], requires_grad=False) + layer.w2_strided_ptrs_s = Parameter(w2_ptrs[1], requires_grad=False) + + layer.sm70_nvfp4_moe = True + layer.sm70_nvfp4_num_experts = num_experts + layer.sm70_nvfp4_hidden_size = hidden + layer.sm70_nvfp4_intermediate_size = intermediate + layer.sm70_nvfp4_w13_k_dim = hidden + layer.sm70_nvfp4_w13_n_dim = 2 * intermediate + layer.sm70_nvfp4_w2_k_dim = intermediate + layer.sm70_nvfp4_w2_n_dim = hidden + layer.sm70_nvfp4_group_size = NVFP4_GROUP_SIZE + layer.sm70_nvfp4_graph_safe_max_tokens = _GRAPH_SAFE_MAX_TOKENS + layer.sm70_nvfp4_compact_grouped_max_tokens = _COMPACT_GROUPED_MAX_TOKENS + self._allocate_graph_safe_decode_buffers(layer) + + del layer.w13_weight + del layer.w13_weight_scale + del layer.w13_weight_scale_2 + del layer.w13_input_scale + del layer.w2_weight + del layer.w2_weight_scale + 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).", + num_experts, + _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 + experts = int(layer.sm70_nvfp4_num_experts) + hidden = int(layer.sm70_nvfp4_hidden_size) + intermediate = int(layer.sm70_nvfp4_intermediate_size) + + layer._nvfp4_sm70_output = torch.empty( + _GRAPH_SAFE_MAX_TOKENS, hidden, dtype=torch.float16, device=device + ) + layer._nvfp4_sm70_permuted_input = torch.empty( + max_slots, hidden, dtype=torch.float16, device=device + ) + layer._nvfp4_sm70_gate_up = torch.empty( + max_slots, 2 * intermediate, dtype=torch.float16, device=device + ) + layer._nvfp4_sm70_intermediate = torch.empty( + max_slots, intermediate, dtype=torch.float16, device=device + ) + layer._nvfp4_sm70_sorted_output = torch.empty( + max_slots, hidden, dtype=torch.float16, device=device + ) + layer._nvfp4_sm70_expert_offsets = torch.empty( + experts + 1, dtype=torch.int32, device=device + ) + layer._nvfp4_sm70_expert_offsets64 = torch.empty( + experts + 1, dtype=torch.int64, device=device + ) + layer._nvfp4_sm70_inv_permuted_idx = torch.empty( + _GRAPH_SAFE_MAX_TOKENS, + _QWEN36_TOP_K, + dtype=torch.int32, + device=device, + ) + layer._nvfp4_sm70_topk_ids = torch.empty( + _GRAPH_SAFE_MAX_TOKENS, + _QWEN36_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) + layer._nvfp4_sm70_permuted_idx = torch.empty( + max_slots, dtype=torch.int32, device=device + ) + layer._nvfp4_sm70_permuted_experts_id = torch.empty( + max_slots, dtype=torch.int32, device=device + ) + layer._nvfp4_sm70_sorted_row_idx = torch.empty( + max_slots, dtype=torch.int32, device=device + ) + layer._nvfp4_sm70_topk_ids_for_sort = torch.empty( + max_slots, dtype=torch.int32, device=device + ) + workspace_size = torch.ops._moe_C.moe_permute_sort_workspace_size( + max_slots, layer.global_num_experts + ) + layer._nvfp4_sm70_sort_workspace = torch.empty( + workspace_size, dtype=torch.int8, device=device + ) + layer._nvfp4_sm70_dense_expert_ids = torch.arange( + experts, dtype=torch.int32, device=device + ) + layer._nvfp4_sm70_compact_offsets = torch.empty( + max_slots + 1, dtype=torch.int32, device=device + ) + layer._nvfp4_sm70_active_expert_ids = torch.empty( + max_slots, dtype=torch.int32, device=device + ) + + @staticmethod + def _persistent_buffers( + layer: RoutedExperts, num_tokens: int + ) -> dict[str, torch.Tensor]: + slots = num_tokens * _QWEN36_TOP_K + return { + "output": layer._nvfp4_sm70_output[:num_tokens], + "permuted_input": layer._nvfp4_sm70_permuted_input[:slots], + "gate_up": layer._nvfp4_sm70_gate_up[:slots], + "intermediate": layer._nvfp4_sm70_intermediate[:slots], + "sorted_output": layer._nvfp4_sm70_sorted_output[:slots], + "expert_offsets": layer._nvfp4_sm70_expert_offsets, + "expert_offsets64": layer._nvfp4_sm70_expert_offsets64, + "inv_permuted_idx": layer._nvfp4_sm70_inv_permuted_idx[:num_tokens], + "topk_ids": layer._nvfp4_sm70_topk_ids[:num_tokens], + "token_expert_indices": ( + layer._nvfp4_sm70_token_expert_indices[:num_tokens] + ), + "permuted_idx": layer._nvfp4_sm70_permuted_idx[:slots], + "sort_workspace": layer._nvfp4_sm70_sort_workspace, + "permuted_experts_id": layer._nvfp4_sm70_permuted_experts_id[:slots], + "sorted_row_idx": layer._nvfp4_sm70_sorted_row_idx[:slots], + "topk_ids_for_sort": layer._nvfp4_sm70_topk_ids_for_sort[:slots], + "dense_expert_ids": layer._nvfp4_sm70_dense_expert_ids, + "compact_offsets": layer._nvfp4_sm70_compact_offsets[: slots + 1], + "active_expert_ids": layer._nvfp4_sm70_active_expert_ids[:slots], + } + + @staticmethod + def _eager_buffers( + layer: RoutedExperts, num_tokens: int + ) -> dict[str, torch.Tensor]: + device = layer.w13_tm_weight.device + slots = num_tokens * _QWEN36_TOP_K + experts = int(layer.sm70_nvfp4_num_experts) + hidden = int(layer.sm70_nvfp4_hidden_size) + intermediate = int(layer.sm70_nvfp4_intermediate_size) + workspace_size = torch.ops._moe_C.moe_permute_sort_workspace_size( + slots, layer.global_num_experts + ) + return { + "output": torch.empty( + num_tokens, hidden, dtype=torch.float16, device=device + ), + "permuted_input": torch.empty( + slots, hidden, dtype=torch.float16, device=device + ), + "gate_up": torch.empty( + slots, 2 * intermediate, dtype=torch.float16, device=device + ), + "intermediate": torch.empty( + slots, intermediate, dtype=torch.float16, device=device + ), + "sorted_output": torch.empty( + slots, hidden, dtype=torch.float16, device=device + ), + "expert_offsets": torch.empty( + experts + 1, dtype=torch.int32, device=device + ), + "expert_offsets64": torch.empty( + experts + 1, dtype=torch.int64, device=device + ), + "inv_permuted_idx": torch.empty( + num_tokens, _QWEN36_TOP_K, dtype=torch.int32, device=device + ), + "topk_ids": torch.empty( + num_tokens, _QWEN36_TOP_K, dtype=torch.int32, device=device + ), + "token_expert_indices": torch.arange( + slots, dtype=torch.int32, device=device + ).view(num_tokens, _QWEN36_TOP_K), + "permuted_idx": torch.empty(slots, dtype=torch.int32, device=device), + "sort_workspace": torch.empty( + workspace_size, dtype=torch.int8, device=device + ), + "permuted_experts_id": torch.empty(slots, dtype=torch.int32, device=device), + "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), + "active_expert_ids": torch.empty(slots, dtype=torch.int32, device=device), + } + + def _get_buffers( + self, layer: RoutedExperts, num_tokens: int + ) -> dict[str, torch.Tensor]: + if 0 < num_tokens <= _GRAPH_SAFE_MAX_TOKENS: + return self._persistent_buffers(layer, num_tokens) + return self._eager_buffers(layer, num_tokens) + + @staticmethod + def _apply_swiglu( + layer: RoutedExperts, out: torch.Tensor, gate_up: torch.Tensor + ) -> None: + if layer.swiglu_limit is None: + torch.ops._C.silu_and_mul(out, gate_up) + else: + torch.ops._C.silu_and_mul_with_clamp( + out, gate_up, float(layer.swiglu_limit) + ) + + def apply( + self, + layer: RoutedExperts, + x: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + shared_experts: SharedExperts | None, + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor: + del shared_experts, shared_experts_input + if not x.is_cuda or x.dtype != torch.float16 or x.ndim != 2: + 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: + raise ValueError( + "SM70 NVFP4 MoE activation hidden size mismatch: expected " + f"{_QWEN36_HIDDEN_SIZE}, got {x.shape[1]}." + ) + 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: + raise TypeError("SM70 NVFP4 MoE requires float32 top-k weights.") + + num_tokens = x.shape[0] + if num_tokens == 0: + return x.new_empty((0, _QWEN36_HIDDEN_SIZE)) + 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) + + if num_tokens <= _COMPACT_GROUPED_MAX_TOKENS: + _prepare_compact_slot_groups( + buffers["permuted_experts_id"], + buffers["compact_offsets"], + buffers["active_expert_ids"], + ) + stage_offsets = buffers["compact_offsets"] + stage_expert_ids = buffers["active_expert_ids"] + stage_experts = slots + else: + stage_offsets = buffers["expert_offsets"] + stage_expert_ids = buffers["dense_expert_ids"] + stage_experts = int(layer.sm70_nvfp4_num_experts) + + sm70_ops.nvfp4_moe_dense_stage_sm70_out( + buffers["gate_up"], + buffers["permuted_input"], + stage_offsets, + stage_expert_ids, + layer.w13_strided_ptrs_w, + layer.w13_strided_ptrs_s, + stage_experts, + layer.sm70_nvfp4_w13_k_dim, + layer.sm70_nvfp4_w13_n_dim, + layer.sm70_nvfp4_group_size, + ) + self._apply_swiglu(layer, buffers["intermediate"], buffers["gate_up"]) + sm70_ops.nvfp4_moe_dense_stage_sm70_out( + buffers["sorted_output"], + buffers["intermediate"], + stage_offsets, + stage_expert_ids, + layer.w2_strided_ptrs_w, + layer.w2_strided_ptrs_s, + stage_experts, + layer.sm70_nvfp4_w2_k_dim, + 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, + ) + return output + + def apply_monolithic( + self, + layer: RoutedExperts, + x: torch.Tensor, + router_logits: torch.Tensor, + input_ids: torch.Tensor | None = None, + ) -> torch.Tensor: + del layer, x, router_logits, input_ids + raise NotImplementedError("SM70 NVFP4 MoE is not a monolithic route.") + + def get_fused_moe_quant_config( # type: ignore[override] + self, layer: RoutedExperts + ) -> FusedMoEQuantConfig | None: + del layer + return None diff --git a/vllm/model_executor/layers/quantization/sm70_turbomind.py b/vllm/model_executor/layers/quantization/sm70_turbomind.py index d948a94a7f..a2482692d7 100644 --- a/vllm/model_executor/layers/quantization/sm70_turbomind.py +++ b/vllm/model_executor/layers/quantization/sm70_turbomind.py @@ -14,6 +14,7 @@ COMPRESSED_UINT4_GROUP_SIZES = (32, 128) MXFP4_GROUP_SIZE = 32 NVFP4_GROUP_SIZE = 16 +NVFP4_QPN4_DENSE_WORKSPACE_ELEMENTS = 5120 * 8704 STATE_ATTR = "_sm70_turbomind_linear" SM70QuantBackend = Literal["auto", "marlin", "turbomind"] @@ -26,7 +27,15 @@ class SM70TurboMindLinearState: k_ld: int q_ld: int output_size: int - op_kind: Literal["uint4", "mxfp4", "nvfp4"] + op_kind: Literal["uint4", "mxfp4", "nvfp4", "nvfp4_qpn4"] + gated_silu: bool = False + dense_weight_ptr: int = 0 + global_scale: float = 0.0 + use_scale_code: bool = False + + +# States retain only data_ptr(), so this cache owns the bounded allocation. +_nvfp4_qpn4_dense_workspaces: dict[tuple[int, torch.dtype], torch.Tensor] = {} def quant_backend() -> SM70QuantBackend: @@ -64,6 +73,13 @@ def should_use_mxfp4_moe_turbomind() -> bool: ) +def should_use_nvfp4_moe_turbomind() -> bool: + """Select the native NVFP4 MoE path only on exact SM70.""" + return is_exact_sm70_cuda_platform() and use_turbomind( + envs.VLLM_SM70_NVFP4_TURBOMIND + ) + + def should_prepare_turbomind( tensor: torch.Tensor, default_enabled: bool, @@ -134,19 +150,27 @@ def _store_state( layer: torch.nn.Module, weight: torch.Tensor, scales: torch.Tensor, - meta: torch.Tensor, + meta: torch.Tensor | None, group_size: int, output_size: int, - op_kind: Literal["uint4", "mxfp4", "nvfp4"], + op_kind: Literal["uint4", "mxfp4", "nvfp4", "nvfp4_qpn4"], + gated_silu: bool = False, + dense_weight_ptr: int = 0, + global_scale: float = 0.0, + use_scale_code: bool = False, ) -> None: state = SM70TurboMindLinearState( weight=weight, scales=scales, group_size=group_size, - k_ld=int(meta[0]), - q_ld=int(meta[1]), + k_ld=0 if meta is None else int(meta[0]), + q_ld=0 if meta is None else int(meta[1]), output_size=output_size, op_kind=op_kind, + gated_silu=gated_silu, + dense_weight_ptr=dense_weight_ptr, + global_scale=global_scale, + use_scale_code=use_scale_code, ) setattr(layer, STATE_ATTR, state) @@ -284,6 +308,71 @@ def prepare_nvfp4_linear( NVFP4_GROUP_SIZE, qweight.size(1), "nvfp4", + interleave_gated_silu, + ) + + +def get_nvfp4_qpn4_dense_workspace(weight: torch.Tensor) -> torch.Tensor | None: + device_index = weight.device.index + if device_index is None: + device_index = torch.accelerator.current_device_index() + cache_key = (device_index, torch.float16) + workspace = _nvfp4_qpn4_dense_workspaces.get(cache_key) + if workspace is not None: + return workspace + try: + workspace = torch.empty( + (NVFP4_QPN4_DENSE_WORKSPACE_ELEMENTS,), + dtype=torch.float16, + device=weight.device, + ) + except torch.OutOfMemoryError: + return None + _nvfp4_qpn4_dense_workspaces[cache_key] = workspace + return workspace + + +def prepare_nvfp4_qpn4_linear( + layer: torch.nn.Module, + workspace: torch.Tensor, + gated_silu: bool, +) -> None: + """Replace one accepted NVFP4 dense weight with memory-neutral QPN4.""" + from vllm import _sm70_ops as sm70_ops + + qweight = unpack_mxfp4_weight(layer.weight.data) + use_scale_code = gated_silu or envs.VLLM_SM70_NVFP4_QPN4_DOWN_SCALE_CODE + global_scale = 0.0 + if use_scale_code: + global_scale = float(layer.weight_global_scale.detach().float().item()) + raw_scale_codes = layer.weight_scale.data.t().contiguous() + packed_weight, packed_scales = sm70_ops.nvfp4_qpn4_prepare_scale_code_sm70( + qweight, raw_scale_codes + ) + else: + fp16_scales = ( + ( + layer.weight_scale.data.t().to(torch.float32) + * layer.weight_global_scale.to(torch.float32) + ) + .to(torch.float16) + .contiguous() + ) + packed_weight, packed_scales = sm70_ops.nvfp4_qpn4_prepare_sm70( + qweight, fp16_scales + ) + _store_state( + layer, + packed_weight, + packed_scales, + None, + NVFP4_GROUP_SIZE, + qweight.size(1), + "nvfp4_qpn4", + gated_silu=gated_silu, + dense_weight_ptr=workspace.data_ptr(), + global_scale=global_scale, + use_scale_code=use_scale_code, ) @@ -332,8 +421,88 @@ def apply_prepared_linear( state.k_ld, state.q_ld, ) + elif state.op_kind == "nvfp4_qpn4": + if reshaped_x.dtype != torch.float16: + raise RuntimeError( + f"SM70 NVFP4 QPN4 requires float16 activations, got {reshaped_x.dtype}." + ) + if reshaped_x.stride(-1) != 1: + reshaped_x = reshaped_x.contiguous() + sm70_ops.nvfp4_qpn4_dispatch_sm70_out( + out, + state.dense_weight_ptr, + reshaped_x, + state.weight, + state.scales, + state.global_scale, + state.use_scale_code, + False, + ) else: raise AssertionError(f"unknown SM70 TurboMind op kind: {state.op_kind}") + if state.gated_silu and state.op_kind == "nvfp4": + out_features = state.output_size // 2 + out = ( + out.reshape(reshaped_x.shape[0], out_features, 2) + .transpose(1, 2) + .reshape(reshaped_x.shape[0], state.output_size) + ) if bias is not None: out.add_(bias) return out.reshape(out_shape) + + +def apply_prepared_fused_silu_and_mul( + layer: torch.nn.Module, + x: torch.Tensor, +) -> torch.Tensor | None: + state = getattr(layer, STATE_ATTR, None) + if ( + state is None + or state.op_kind not in ("nvfp4", "nvfp4_qpn4") + or not state.gated_silu + ): + return None + if x.dtype != torch.float16: + raise RuntimeError( + "SM70 TurboMind NVFP4 gated-SiLU requires float16 activations, " + f"got {x.dtype}." + ) + + reshaped_x = x.reshape(-1, x.shape[-1]) + if reshaped_x.stride(-1) != 1: + reshaped_x = reshaped_x.contiguous() + out_features = state.output_size // 2 + out = torch.empty( + (reshaped_x.shape[0], out_features), + dtype=x.dtype, + device=x.device, + ) + if reshaped_x.shape[0] == 0: + return out.reshape(*x.shape[:-1], out_features) + + from vllm import _sm70_ops as sm70_ops + + if state.op_kind == "nvfp4_qpn4": + sm70_ops.nvfp4_qpn4_dispatch_sm70_out( + out, + state.dense_weight_ptr, + reshaped_x, + state.weight, + state.scales, + state.global_scale, + state.use_scale_code, + True, + ) + else: + sm70_ops.nvfp4_gemm_sm70_out( + out, + reshaped_x, + state.weight, + state.scales, + state.group_size, + state.k_ld, + state.q_ld, + True, + ) + return out.reshape(*x.shape[:-1], out_features) 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/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.py b/vllm/model_executor/layers/sparse_attn_indexer.py index d3b10092a6..89f7f2d457 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -355,7 +355,7 @@ def sparse_attn_indexer( padded_weights.reshape(num_padded_tokens, -1), seq_lens, decode_metadata.block_table, - attn_metadata_narrowed.max_seq_len, + attn_metadata_narrowed.compressed_max_seq_len, ) elif current_platform.is_xpu(): if padded_q_scale is not None: @@ -397,7 +397,7 @@ def sparse_attn_indexer( topk_indices, topk_workspace, topk_tokens, - attn_metadata_narrowed.max_seq_len, + logits.shape[1], ) else: ops.top_k_per_row_decode( diff --git a/vllm/model_executor/layers/vocab_parallel_embedding.py b/vllm/model_executor/layers/vocab_parallel_embedding.py index 73b9e24167..9f69084c62 100644 --- a/vllm/model_executor/layers/vocab_parallel_embedding.py +++ b/vllm/model_executor/layers/vocab_parallel_embedding.py @@ -35,6 +35,14 @@ DEFAULT_VOCAB_PADDING_SIZE = 64 logger = init_logger(__name__) +_SM70_DFLASH2_QPN8_LM_HEAD_SHAPE = (62080, 5120) +_SM70_DFLASH2_QPN8_MAX_ROWS = 8 +_SM70_DFLASH2_QPN8_CANDIDATES = 64 +_SM70_DFLASH2_QPN8_SPLIT_K = 8 +_SM70_DFLASH2_QPN8_ACCUMULATOR_CHAINS = 1 +_SM70_DFLASH2_RERANK_CTA_N = 128 +_SM70_DFLASH2_RERANK_SPLIT_K = 10 + def _sm70_env_bool(name: str, default: bool) -> bool: value = os.getenv(name) @@ -44,8 +52,16 @@ def _sm70_env_bool(name: str, default: bool) -> bool: def _sm70_lm_head_top1_default() -> bool: - return not _sm70_env_bool( - "VLLM_SM70_FLASH_V100_0DOT3_COMPILE_GRAPH", False + return not _sm70_env_bool("VLLM_SM70_FLASH_V100_0DOT3_COMPILE_GRAPH", False) + + +def _sm70_dflash2_qpn8_rerank_enabled() -> bool: + return envs.VLLM_SM70_DFLASH2_QPN8_RERANK + + +def _sm70_dflash2_qpn8_rerank_requested() -> bool: + return ( + _sm70_dflash2_qpn8_rerank_enabled() or envs.VLLM_SM70_DFLASH2_QPN8_RERANK_SHADOW ) @@ -65,6 +81,7 @@ def _is_sm70_lm_head_fastpath_eligible(layer: torch.nn.Module) -> bool: _sm70_env_bool("VLLM_SM70_ENABLE_LM_HEAD_FASTPATH", False) or _sm70_env_bool("VLLM_SM70_LM_HEAD_TOP1", _sm70_lm_head_top1_default()) or _sm70_env_bool("VLLM_SM70_LM_HEAD_TOP1_TC", False) + or _sm70_dflash2_qpn8_rerank_requested() ): _trace_sm70_lm_head_skip("disabled") return False @@ -94,6 +111,229 @@ def _is_sm70_lm_head_fastpath_eligible(layer: torch.nn.Module) -> bool: return True +def _is_sm70_dflash2_qpn8_rerank_eligible(layer: torch.nn.Module) -> bool: + if not _sm70_dflash2_qpn8_rerank_requested(): + return False + if tuple(layer.weight.shape) != _SM70_DFLASH2_QPN8_LM_HEAD_SHAPE: + logger.warning_once( + "SM70 DFlash2 QPN8 rerank requires LM-head shape %s; got %s. " + "Using the dense LM head.", + _SM70_DFLASH2_QPN8_LM_HEAD_SHAPE, + tuple(layer.weight.shape), + ) + return False + if getattr(layer, "tp_size", 1) != 4: + logger.warning_once( + "SM70 DFlash2 QPN8 rerank requires TP4; using the dense LM head." + ) + return False + if layer.shard_indices.num_org_vocab_padding != 0: + logger.warning_once( + "SM70 DFlash2 QPN8 rerank requires an unpadded local vocabulary; " + "using the dense LM head." + ) + return False + required_ops = ( + "fp8_qpn8_prepare_sm70", + "fp8_qpn8_gemm_sm70_out", + "sm70_f16_indexed_rerank_packed_out", + "sm70_f16_rerank_keys_out", + "sm70_f16_rerank_topk_out", + ) + missing = [name for name in required_ops if not hasattr(torch.ops._C, name)] + if missing: + logger.warning_once( + "SM70 DFlash2 QPN8 rerank operators are unavailable (%s); using " + "the dense LM head.", + ", ".join(missing), + ) + return False + return True + + +@torch.inference_mode() +def _prepare_sm70_dflash2_qpn8_rerank(layer: torch.nn.Module) -> bool: + if not _is_sm70_dflash2_qpn8_rerank_eligible(layer): + return False + + weight = layer.weight + rows, hidden = weight.shape + qweight = torch.empty_like(weight, dtype=torch.float8_e4m3fn) + channel_scales = torch.empty((rows, 1), dtype=torch.float32, device=weight.device) + # Quantize in bounded row chunks so startup never materializes another + # full-size FP32 LM head. The original FP16 parameter remains the oracle + # used by the exact candidate rerank. + chunk_rows = 4096 + for begin in range(0, rows, chunk_rows): + end = min(begin + chunk_rows, rows) + weight_f32 = weight[begin:end].float() + scales = weight_f32.abs().amax(dim=1, keepdim=True).div_(448.0) + scales.clamp_(min=torch.finfo(torch.float32).tiny) + channel_scales[begin:end].copy_(scales) + qweight[begin:end].copy_( + weight_f32.div_(scales).clamp_(-448.0, 448.0).to(torch.float8_e4m3fn) + ) + + codes, packed_scales = sm70_ops.fp8_qpn8_prepare_sm70(qweight, channel_scales) + del qweight, channel_scales, weight_f32, scales + torch.accelerator.empty_cache() + + device = weight.device + max_rows = _SM70_DFLASH2_QPN8_MAX_ROWS + candidates = _SM70_DFLASH2_QPN8_CANDIDATES + layer.register_buffer("_sm70_dflash2_qpn8_codes", codes, persistent=False) + layer.register_buffer("_sm70_dflash2_qpn8_scales", packed_scales, persistent=False) + layer.register_buffer( + "_sm70_dflash2_qpn8_logits", + torch.empty((max_rows, rows), dtype=torch.float16, device=device), + persistent=False, + ) + layer.register_buffer( + "_sm70_dflash2_qpn8_values", + torch.empty((max_rows, candidates), dtype=torch.float16, device=device), + persistent=False, + ) + layer.register_buffer( + "_sm70_dflash2_qpn8_ids", + torch.empty((max_rows, candidates), dtype=torch.int64, device=device), + persistent=False, + ) + layer.register_buffer( + "_sm70_dflash2_rerank_logits", + torch.empty((max_rows, candidates), dtype=torch.float16, device=device), + persistent=False, + ) + selected_rows = max_rows * candidates + layer.register_buffer( + "_sm70_dflash2_rerank_selected_raw", + torch.empty((selected_rows, hidden), dtype=torch.float16, device=device), + persistent=False, + ) + layer.register_buffer( + "_sm70_dflash2_rerank_selected_packed", + torch.empty((selected_rows, hidden), dtype=torch.float16, device=device), + persistent=False, + ) + layer.register_buffer( + "_sm70_dflash2_rerank_expanded", + torch.empty((max_rows, selected_rows), dtype=torch.float16, device=device), + persistent=False, + ) + layer.register_buffer( + "_sm70_dflash2_rerank_partials", + torch.empty((max_rows, selected_rows), dtype=torch.float32, device=device), + persistent=False, + ) + layer.register_buffer( + "_sm70_dflash2_rerank_barriers", + torch.zeros(64, dtype=torch.int32, device=device), + persistent=False, + ) + layer.register_buffer( + "_sm70_dflash2_rerank_dense_logits", + torch.empty((max_rows, rows), dtype=torch.float16, device=device), + persistent=False, + ) + layer.register_buffer( + "_sm70_dflash2_rerank_keys", + torch.empty((max_rows, candidates), dtype=torch.int64, device=device), + persistent=False, + ) + # Keep distinct top-16 and top-20 outputs. Slicing the columns of one + # [max_rows, 20] allocation for top-16 leaves a row stride of 20 and makes + # the result non-contiguous. The TP all-gather requires contiguous inputs, + # and inserting a runtime contiguous() copy would add work to both graphs. + for selector_k in (16, 20): + layer.register_buffer( + f"_sm70_dflash2_rerank_values_{selector_k}", + torch.empty((max_rows, selector_k), dtype=torch.float16, device=device), + persistent=False, + ) + layer.register_buffer( + f"_sm70_dflash2_rerank_positions_{selector_k}", + torch.empty((max_rows, selector_k), dtype=torch.int64, device=device), + persistent=False, + ) + layer.register_buffer( + f"_sm70_dflash2_rerank_ids_{selector_k}", + torch.empty((max_rows, selector_k), dtype=torch.int64, device=device), + persistent=False, + ) + layer.register_buffer( + f"_sm70_dflash2_rerank_key_values_{selector_k}", + torch.empty((max_rows, selector_k), dtype=torch.int64, device=device), + persistent=False, + ) + layer._sm70_dflash2_qpn8_rerank_prepared = True + logger.info_once( + "SM70 DFlash2 QPN8 top-64 plus exact TurboMind FP16 rerank layout prepared." + ) + return True + + +def _sm70_dflash2_rerank_output_buffers( + layer: torch.nn.Module, + num_rows: int, + selector_k: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Select graph-stable, contiguous rerank outputs for top-16 or top-20.""" + if selector_k == 16: + values = layer._sm70_dflash2_rerank_values_16[:num_rows] + positions = layer._sm70_dflash2_rerank_positions_16[:num_rows] + ids = layer._sm70_dflash2_rerank_ids_16[:num_rows] + elif selector_k == 20: + values = layer._sm70_dflash2_rerank_values_20[:num_rows] + positions = layer._sm70_dflash2_rerank_positions_20[:num_rows] + ids = layer._sm70_dflash2_rerank_ids_20[:num_rows] + else: + raise ValueError(f"Unsupported DFlash2 rerank top-k: {selector_k}") + return values, positions, ids + + +def _sm70_dflash2_dense_order_topk( + sparse_logits: torch.Tensor, + candidate_ids: torch.Tensor, + candidate_logits: torch.Tensor, + values: torch.Tensor, + ids: torch.Tensor, + selector_k: int, + vocab_start_index: int, +) -> None: + """Restore dense-vocabulary tie order after sparse candidate reranking.""" + sparse_logits.fill_(-float("inf")) + sparse_logits.scatter_(1, candidate_ids, candidate_logits) + torch.topk( + sparse_logits, + selector_k, + dim=-1, + sorted=True, + out=(values, ids), + ) + ids.add_(vocab_start_index) + + +def _sm70_dflash2_candidate_order_topk( + candidate_ids: torch.Tensor, + candidate_logits: torch.Tensor, + values: torch.Tensor, + ids: torch.Tensor, + selector_k: int, + vocab_start_index: int, +) -> None: + """Select exact values with original-vocabulary tie precedence.""" + if values.size(1) != selector_k: + raise ValueError( + f"DFlash2 rerank output width {values.size(1)} != top-k {selector_k}" + ) + sm70_ops.sm70_f16_rerank_topk_out( + values, + ids, + candidate_logits, + candidate_ids, + vocab_start_index, + ) + + def maybe_prepare_sm70_lm_head_top1(layer: torch.nn.Module) -> bool: if getattr(layer, "_sm70_f16_prepared", False): return True @@ -103,7 +343,9 @@ def maybe_prepare_sm70_lm_head_top1(layer: torch.nn.Module) -> bool: layer._sm70_f16_tm_weight = prepared[0] layer._sm70_f16_k_ld = int(prepared[1][0].item()) layer._sm70_f16_prepared = True - if envs.VLLM_SM70_ENABLE_LM_HEAD_FASTPATH: + if _prepare_sm70_dflash2_qpn8_rerank(layer): + pass + elif envs.VLLM_SM70_ENABLE_LM_HEAD_FASTPATH: logger.info_once("SM70 dense fp16 fast path enabled for LM head.") else: logger.info_once("SM70 LM head top1 layout prepared.") @@ -164,8 +406,7 @@ def _maybe_sm70_lm_head_top1( or hasattr(torch.ops._C, "sm70_f16_lm_head_top1_tc_out") ): logger.warning_once( - "SM70 LM head top1 requested, but no top1 op is available; " - "falling back." + "SM70 LM head top1 requested, but no top1 op is available; falling back." ) return None if x.dtype != torch.float16 or not x.is_cuda: @@ -193,9 +434,7 @@ def _maybe_sm70_lm_head_top1( values = torch.empty((x_2d.size(0),), dtype=torch.float32, device=x_2d.device) indices = torch.empty((x_2d.size(0),), dtype=torch.int64, device=x_2d.device) - if lm_head_top1_tc and hasattr( - torch.ops._C, "sm70_f16_lm_head_top1_tc_out" - ): + if lm_head_top1_tc and hasattr(torch.ops._C, "sm70_f16_lm_head_top1_tc_out"): tm_weight = getattr(layer, "_sm70_f16_tm_weight", None) k_ld = getattr(layer, "_sm70_f16_k_ld", None) if tm_weight is not None and k_ld is not None: @@ -216,9 +455,7 @@ def _maybe_sm70_lm_head_top1( if num_rows != 1: return None - if not lm_head_top1 or not hasattr( - torch.ops._C, "sm70_f16_lm_head_top1_out" - ): + if not lm_head_top1 or not hasattr(torch.ops._C, "sm70_f16_lm_head_top1_out"): return None sm70_ops.sm70_f16_lm_head_top1_out( @@ -234,6 +471,146 @@ def _maybe_sm70_lm_head_top1( return values.reshape(*x.shape[:-1]), indices.reshape(*x.shape[:-1]) +def _maybe_sm70_dflash2_qpn8_rerank( + layer: torch.nn.Module, + x: torch.Tensor, + selector_k: int, + bias: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor] | None: + """Return shard-local top-k after QPN8 support search and FP16 rerank.""" + if not _sm70_dflash2_qpn8_rerank_requested(): + return None + if not getattr(layer, "_sm70_dflash2_qpn8_rerank_prepared", False): + return None + if selector_k not in (16, 20) or bias is not None: + return None + if x.dtype != torch.float16 or not x.is_cuda: + return None + if torch.cuda.get_device_capability(x.device) != (7, 0): + return None + + x_2d = x.reshape(-1, x.shape[-1]) + num_rows = x_2d.size(0) + if not 1 <= num_rows <= _SM70_DFLASH2_QPN8_MAX_ROWS: + return None + if not x_2d.is_contiguous(): + x_2d = x_2d.contiguous() + + qpn8_logits = layer._sm70_dflash2_qpn8_logits[:num_rows] + qpn8_values = layer._sm70_dflash2_qpn8_values[:num_rows] + qpn8_ids = layer._sm70_dflash2_qpn8_ids[:num_rows] + sm70_ops.fp8_qpn8_gemm_sm70_out( + qpn8_logits, + x_2d, + layer._sm70_dflash2_qpn8_codes, + layer._sm70_dflash2_qpn8_scales, + _SM70_DFLASH2_QPN8_SPLIT_K, + _SM70_DFLASH2_QPN8_ACCUMULATOR_CHAINS, + False, + False, + ) + # The exact FP16 rerank is permutation-invariant over this approximate + # support, so skip the unnecessary 64-element result sort. This keeps the + # official PyTorch multiblock selector while avoiding its final bitonic + # kernel and leaves candidate quality unchanged. + torch.topk( + qpn8_logits, + _SM70_DFLASH2_QPN8_CANDIDATES, + dim=-1, + sorted=False, + out=(qpn8_values, qpn8_ids), + ) + + sm70_ops.sm70_f16_indexed_rerank_packed_out( + layer._sm70_dflash2_rerank_logits[:num_rows], + x_2d, + layer._sm70_f16_tm_weight, + qpn8_ids, + layer._sm70_dflash2_rerank_selected_packed, + layer._sm70_dflash2_rerank_expanded, + layer._sm70_dflash2_rerank_partials, + layer._sm70_dflash2_rerank_barriers, + _SM70_DFLASH2_RERANK_CTA_N, + _SM70_DFLASH2_RERANK_SPLIT_K, + ) + rerank_logits = layer._sm70_dflash2_rerank_logits[:num_rows] + values, _positions, ids = _sm70_dflash2_rerank_output_buffers( + layer, num_rows, selector_k + ) + if envs.VLLM_SM70_DFLASH2_QPN8_DENSE_ORDER: + _sm70_dflash2_dense_order_topk( + layer._sm70_dflash2_rerank_dense_logits[:num_rows], + qpn8_ids, + rerank_logits, + values, + ids, + selector_k, + layer.shard_indices.org_vocab_start_index, + ) + else: + _sm70_dflash2_candidate_order_topk( + qpn8_ids, + rerank_logits, + values, + ids, + selector_k, + layer.shard_indices.org_vocab_start_index, + ) + + if envs.VLLM_SM70_DFLASH2_QPN8_RERANK_SHADOW: + if torch.compiler.is_compiling() or torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "VLLM_SM70_DFLASH2_QPN8_RERANK_SHADOW is eager-only; disable " + "CUDA Graph/torch.compile for the real-hidden coverage audit." + ) + dense_logits = layer.quant_method.apply(layer, x_2d, bias=None) + dense_values, dense_ids = torch.topk( + dense_logits, + selector_k, + dim=-1, + sorted=True, + ) + support_matches = (dense_ids[:, :, None] == qpn8_ids[:, None, :]).any(dim=-1) + dense_global_ids = dense_ids + layer.shard_indices.org_vocab_start_index + rerank_matches = (dense_global_ids[:, :, None] == ids[:, None, :]).any(dim=-1) + call = int(getattr(layer, "_sm70_dflash2_qpn8_shadow_calls", 0)) + 1 + layer._sm70_dflash2_qpn8_shadow_calls = call + logger.info( + "SM70_DFLASH2_QPN8_SHADOW call=%d rows=%d top_k=%d " + "support_missing=%d exact_set_rows=%d top1_match_rows=%d " + "ordered_value_max_abs=%.6g", + call, + num_rows, + selector_k, + int((~support_matches).sum().item()), + int(rerank_matches.all(dim=-1).sum().item()), + int((dense_global_ids[:, 0] == ids[:, 0]).sum().item()), + float((dense_values.float() - values.float()).abs().max().item()), + ) + output_shape = (*x.shape[:-1], selector_k) + return dense_values.reshape(output_shape), dense_global_ids.reshape( + output_shape + ) + + 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, + ) + output_shape = (*x.shape[:-1], selector_k) + return values.reshape(output_shape), ids.reshape(output_shape) + + +def _maybe_sm70_dflash2_lm_head_top20( + layer: torch.nn.Module, + x: torch.Tensor, + selector_k: int, + bias: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor] | None: + """Return selector candidates from the opt-in QPN8 plus FP16 rerank.""" + return _maybe_sm70_dflash2_qpn8_rerank(layer, x, selector_k, bias) + + class UnquantizedEmbeddingMethod(QuantizeMethodBase): """Unquantized method for embeddings.""" @@ -712,6 +1089,14 @@ def maybe_get_sm70_lm_head_top1( ) -> tuple[torch.Tensor, torch.Tensor] | None: return _maybe_sm70_lm_head_top1(self, hidden_states, bias) + def maybe_get_sm70_dflash2_top20( + self, + hidden_states: torch.Tensor, + selector_k: int, + bias: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor] | None: + return _maybe_sm70_dflash2_lm_head_top20(self, hidden_states, selector_k, bias) + def extra_repr(self) -> str: s = f"num_embeddings={self.num_embeddings_per_partition}" s += f", embedding_dim={self.embedding_dim}" diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index d48ec981af..de9786c777 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 " diff --git a/vllm/model_executor/models/deepseek_ocr.py b/vllm/model_executor/models/deepseek_ocr.py index 2575d3dcd4..78ad964a3c 100644 --- a/vllm/model_executor/models/deepseek_ocr.py +++ b/vllm/model_executor/models/deepseek_ocr.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Inference-only Deepseek-OCR model compatible with HuggingFace weights.""" -import math from collections.abc import Iterable, Mapping, Sequence from typing import Annotated, Literal @@ -53,7 +52,7 @@ BASE_SIZE, CROP_MODE, DeepseekOCRProcessor, - count_tiles, + count_image_tokens_for, ) from vllm.utils.tensor_schema import TensorSchema, TensorShape from vllm.v1.sample.logits_processor import ( @@ -190,12 +189,15 @@ def get_hf_config(self): return self.ctx.get_hf_config(DeepseekVLV2Config) def get_hf_processor(self, **kwargs: object): - v1_processor_config = dict( - image_size=IMAGE_SIZE, - base_size=BASE_SIZE, - crop_mode=CROP_MODE, - strategy="v1", - ) + v1_processor_config: dict[str, object] = dict(strategy="v1") + if "image_mode" not in kwargs: + # Only inject the constant trio when no named mode is requested; + # the processor treats ``image_mode`` as authoritative. + v1_processor_config.update( + image_size=IMAGE_SIZE, + base_size=BASE_SIZE, + crop_mode=CROP_MODE, + ) return self.ctx.get_hf_processor( DeepseekOCRProcessor, @@ -208,35 +210,13 @@ def get_supported_mm_limits(self) -> Mapping[str, int | None]: def get_num_image_tokens( self, *, image_width: int, image_height: int, cropping: bool = True ) -> int: - image_size = IMAGE_SIZE - base_size = BASE_SIZE - patch_size = 16 - downsample_ratio = 4 - - if CROP_MODE: - if image_width <= 640 and image_height <= 640: - crop_ratio = [1, 1] - else: - # find the closest aspect ratio to the target - crop_ratio = count_tiles( - image_width, image_height, image_size=IMAGE_SIZE - ) - - num_width_tiles, num_height_tiles = crop_ratio - else: - num_width_tiles = num_height_tiles = 1 - - h = w = math.ceil((base_size // patch_size) / downsample_ratio) - - h2 = w2 = math.ceil((image_size // patch_size) / downsample_ratio) - - global_views_tokens = h * (w + 1) - if num_width_tiles > 1 or num_height_tiles > 1: - local_views_tokens = (num_height_tiles * h2) * (num_width_tiles * w2 + 1) - else: - local_views_tokens = 0 - - return global_views_tokens + local_views_tokens + 1 + return count_image_tokens_for( + image_width=image_width, + image_height=image_height, + base_size=BASE_SIZE, + image_size=IMAGE_SIZE, + cropping=cropping and CROP_MODE, + ) def get_image_size_with_most_features(self) -> ImageSize: if IMAGE_SIZE == 1024 and BASE_SIZE == 1280: @@ -334,10 +314,12 @@ def get_replacement_deepseek_vl2(item_idx: int): else: size = images.get_image_size(item_idx) - num_image_tokens = self.info.get_num_image_tokens( + # Count with the SAME processor instance the request's + # mm kwargs construct (mode/crops) — the producer/counter + # pair must never diverge. + num_image_tokens = hf_processor.count_image_tokens( image_width=size.width, image_height=size.height, - cropping=CROP_MODE, ) return [image_token_id] * num_image_tokens diff --git a/vllm/model_executor/models/dflash_sm70.py b/vllm/model_executor/models/dflash_sm70.py new file mode 100644 index 0000000000..039ea71a27 --- /dev/null +++ b/vllm/model_executor/models/dflash_sm70.py @@ -0,0 +1,300 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Range-preserving BF16 arithmetic for DFlash2 drafts on SM70. + +The checkpoint was trained with BF16 activations, but Volta executes its dense +GEMMs in FP16. These kernels retain FP16 Tensor Core transport while preserving +BF16 residual range and rounding at the model's trained rounding points. + +Adapted from haohervchb/sglang-V100@5526ef1c6a82. +""" + +from __future__ import annotations + +import torch +from torch import nn + +from vllm.triton_utils import tl, triton + +DFLASH_SM70_GATE_UP_INPUT_SCALE = 32.0 +DFLASH_SM70_WIDE_OUTPUT_SCALE = 256.0 + + +@triton.jit +def _round_fp32_to_bf16_rne(value): + """Round FP32 to BF16 RNE and keep the result represented as FP32.""" + bits = value.to(tl.uint32, bitcast=True) + rounded_bits = (bits + 0x7FFF + ((bits >> 16) & 1)) & 0xFFFF0000 + return rounded_bits.to(tl.float32, bitcast=True) + + +@triton.jit +def _dflash_silu_and_mul_sm70_kernel( + output_ptr, + output_scales_ptr, + gate_up_ptr, + n_cols: tl.constexpr, + gate_up_scale: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + offsets = tl.arange(0, BLOCK_SIZE) + mask = offsets < n_cols + input_offsets = row * 2 * n_cols + offsets + output_offsets = row * n_cols + offsets + + gate = tl.load(gate_up_ptr + input_offsets, mask=mask, other=0.0).to(tl.float32) + up = tl.load(gate_up_ptr + input_offsets + n_cols, mask=mask, other=0.0).to( + tl.float32 + ) + gate = _round_fp32_to_bf16_rne(gate * gate_up_scale) + up = _round_fp32_to_bf16_rne(up * gate_up_scale) + activated_gate = _round_fp32_to_bf16_rne(gate * tl.sigmoid(gate)) + output = _round_fp32_to_bf16_rne(activated_gate * up) + + # A per-row power-of-two divisor retains exact transport scaling. + max_abs = tl.max(tl.abs(output), axis=0) + output_scale = tl.maximum(max_abs / 32752.0, 1.0) + output_scale = tl.exp2(tl.ceil(tl.log2(output_scale))) + tl.store(output_scales_ptr + row, output_scale) + tl.store(output_ptr + output_offsets, output / output_scale, mask=mask) + + +@triton.jit +def _dflash_scale_output_sm70_kernel( + output_ptr, + row_scales_ptr, + n_cols: tl.constexpr, + residual_scale: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + offsets = tl.arange(0, BLOCK_SIZE) + mask = offsets < n_cols + row_offsets = row * n_cols + offsets + value = tl.load(output_ptr + row_offsets, mask=mask, other=0.0).to(tl.float32) + row_scale = tl.load(row_scales_ptr + row).to(tl.float32) + tl.store( + output_ptr + row_offsets, + value * (row_scale / residual_scale), + mask=mask, + ) + + +@triton.jit +def _dflash_rmsnorm_sm70_kernel( + output_ptr, + residual_output_ptr, + x_ptr, + residual_ptr, + weight_ptr, + n_cols: tl.constexpr, + eps: tl.constexpr, + residual_scale: tl.constexpr, + HAS_RESIDUAL: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + offsets = tl.arange(0, BLOCK_SIZE) + mask = offsets < n_cols + row_offsets = row * n_cols + offsets + + value = tl.load(x_ptr + row_offsets, mask=mask, other=0.0).to(tl.float32) + if HAS_RESIDUAL: + residual = tl.load(residual_ptr + row_offsets, mask=mask, other=0.0).to( + tl.float32 + ) + value = _round_fp32_to_bf16_rne(value * residual_scale + residual) + tl.store(residual_output_ptr + row_offsets, value, mask=mask) + + variance = tl.sum(value * value, axis=0) / n_cols + normalized = value * tl.rsqrt(variance + eps) + weight = tl.load(weight_ptr + offsets, mask=mask, other=0.0).to(tl.float32) + weight = _round_fp32_to_bf16_rne(weight) + output = _round_fp32_to_bf16_rne(normalized * weight) + tl.store(output_ptr + row_offsets, output, mask=mask) + + +def _num_warps(n_cols: int) -> int: + return max(min(triton.next_power_of_2(triton.cdiv(n_cols, 256)), 16), 4) + + +def _silu_and_mul_reference( + gate_up: torch.Tensor, gate_up_scale: float +) -> tuple[torch.Tensor, torch.Tensor]: + gate, up = gate_up.float().chunk(2, dim=-1) + gate = (gate * gate_up_scale).to(torch.bfloat16).float() + up = (up * gate_up_scale).to(torch.bfloat16).float() + activated = (gate * torch.sigmoid(gate)).to(torch.bfloat16).float() + output = (activated * up).to(torch.bfloat16).float() + max_abs = output.abs().amax(dim=-1) + row_scales = torch.maximum(max_abs / 32752.0, torch.ones_like(max_abs)) + row_scales = torch.pow(2.0, torch.ceil(torch.log2(row_scales))) + return (output / row_scales.unsqueeze(-1)).to(gate_up.dtype), row_scales + + +def dflash_silu_and_mul_sm70( + gate_up: torch.Tensor, + gate_up_scale: float = DFLASH_SM70_GATE_UP_INPUT_SCALE, +) -> tuple[torch.Tensor, torch.Tensor]: + if gate_up.ndim < 2 or gate_up.shape[-1] % 2: + raise ValueError( + "DFlash SM70 SwiGLU expects [..., 2 * intermediate_size], " + f"got {tuple(gate_up.shape)}." + ) + if gate_up_scale <= 0: + raise ValueError(f"gate_up_scale must be positive, got {gate_up_scale}.") + if not gate_up.is_cuda: + return _silu_and_mul_reference(gate_up, gate_up_scale) + if gate_up.dtype != torch.float16: + raise ValueError( + f"DFlash SM70 SwiGLU requires FP16 CUDA transport, got {gate_up.dtype}." + ) + if not gate_up.is_contiguous(): + gate_up = gate_up.contiguous() + + n_cols = gate_up.shape[-1] // 2 + n_rows = gate_up.numel() // (2 * n_cols) + output = torch.empty( + (*gate_up.shape[:-1], n_cols), + dtype=gate_up.dtype, + device=gate_up.device, + ) + output_scales = torch.empty(n_rows, dtype=torch.float32, device=gate_up.device) + block_size = triton.next_power_of_2(n_cols) + _dflash_silu_and_mul_sm70_kernel[(n_rows,)]( + output, + output_scales, + gate_up, + n_cols=n_cols, + gate_up_scale=gate_up_scale, + BLOCK_SIZE=block_size, + num_warps=_num_warps(n_cols), + ) + return output, output_scales + + +def dflash_scale_output_sm70( + output: torch.Tensor, + row_scales: torch.Tensor, + residual_scale: float = DFLASH_SM70_WIDE_OUTPUT_SCALE, +) -> torch.Tensor: + if output.ndim < 2: + raise ValueError(f"DFlash SM70 output must have ndim >= 2, got {output.ndim}.") + if residual_scale <= 0: + raise ValueError(f"residual_scale must be positive, got {residual_scale}.") + n_cols = output.shape[-1] + n_rows = output.numel() // n_cols + if row_scales.numel() != n_rows: + raise ValueError( + f"DFlash SM70 row-scale mismatch: rows={n_rows}, " + f"scales={row_scales.numel()}." + ) + if not output.is_cuda: + return (output.float() * (row_scales.reshape(-1, 1) / residual_scale)).to( + output.dtype + ) + if output.dtype != torch.float16 or not output.is_contiguous(): + raise ValueError("DFlash SM70 W2 scaling requires contiguous CUDA FP16 output.") + block_size = triton.next_power_of_2(n_cols) + _dflash_scale_output_sm70_kernel[(n_rows,)]( + output, + row_scales, + n_cols=n_cols, + residual_scale=residual_scale, + BLOCK_SIZE=block_size, + num_warps=_num_warps(n_cols), + ) + return output + + +class DFlashSM70RMSNorm(nn.Module): + """RMSNorm with FP32 residual transport and explicit BF16 RNE.""" + + def __init__( + self, + hidden_size: int, + eps: float, + dtype: torch.dtype, + residual_scale: float = DFLASH_SM70_WIDE_OUTPUT_SCALE, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.variance_epsilon = eps + self.residual_scale = residual_scale + self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype)) + + def _reference( + self, x: torch.Tensor, residual: torch.Tensor | None + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + orig_dtype = x.dtype + value = x.float() + residual_output = None + if residual is not None: + value = ( + (value * self.residual_scale + residual.float()) + .to(torch.bfloat16) + .float() + ) + residual_output = value + variance = value.square().mean(dim=-1, keepdim=True) + normalized = value * torch.rsqrt(variance + self.variance_epsilon) + weight = self.weight.float().to(torch.bfloat16).float() + output = (normalized * weight).to(torch.bfloat16).to(orig_dtype) + if residual_output is None: + return output + return output, residual_output + + def forward( + self, x: torch.Tensor, residual: torch.Tensor | None = None + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if not x.is_cuda: + return self._reference(x, residual) + if x.dtype != torch.float16: + raise ValueError( + f"DFlash SM70 RMSNorm requires FP16 CUDA transport, got {x.dtype}." + ) + if not x.is_contiguous(): + x = x.contiguous() + if residual is not None and not residual.is_contiguous(): + residual = residual.contiguous() + + n_rows = x.numel() // self.hidden_size + output = torch.empty_like(x) + has_residual = residual is not None + residual_output = ( + torch.empty_like(x, dtype=torch.float32) if has_residual else output + ) + residual_ptr = residual if residual is not None else x + block_size = triton.next_power_of_2(self.hidden_size) + _dflash_rmsnorm_sm70_kernel[(n_rows,)]( + output, + residual_output, + x, + residual_ptr, + self.weight, + n_cols=self.hidden_size, + eps=self.variance_epsilon, + residual_scale=self.residual_scale, + HAS_RESIDUAL=has_residual, + BLOCK_SIZE=block_size, + num_warps=_num_warps(self.hidden_size), + ) + if has_residual: + return output, residual_output + return output + + +class DFlashSM70MLP(nn.Module): + """Wrap the official Qwen MLP with overflow-safe BF16-equivalent math.""" + + def __init__(self, mlp: nn.Module) -> None: + super().__init__() + self.gate_up_proj = mlp.gate_up_proj + self.down_proj = mlp.down_proj + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate_up, _ = self.gate_up_proj(x / DFLASH_SM70_GATE_UP_INPUT_SCALE) + x, row_scales = dflash_silu_and_mul_sm70(gate_up) + x, _ = self.down_proj(x) + return dflash_scale_output_sm70(x, row_scales) 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 46b3269f8b..d6011a3e58 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -265,6 +265,8 @@ def _mark_tower_model( If `targets` is set, instead include descendants that are an instance of `targets`, even if they aren't direct children. """ + from vllm.model_executor.offloader import get_offloader + from .utils import StageMissingLayer, collect_children, no_init_weights if isinstance(modalities, str): @@ -290,6 +292,14 @@ def _mark_tower_model( yield self._tower_model_names = children_names + offloader = get_offloader() + for name in children_names: + if not name: + continue + tower = self.get_submodule(name) + wrapped = offloader.wrap_module(tower, parameter_prefix=name) + if wrapped is not tower: + self.set_submodule(name, wrapped) @contextmanager def _mark_composite_model( diff --git a/vllm/model_executor/models/qwen2_moe.py b/vllm/model_executor/models/qwen2_moe.py index f27c4f3037..ecfbf4f064 100644 --- a/vllm/model_executor/models/qwen2_moe.py +++ b/vllm/model_executor/models/qwen2_moe.py @@ -149,9 +149,7 @@ 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) diff --git a/vllm/model_executor/models/qwen3_5.py b/vllm/model_executor/models/qwen3_5.py index b91c77dbc0..dbc891fee5 100644 --- a/vllm/model_executor/models/qwen3_5.py +++ b/vllm/model_executor/models/qwen3_5.py @@ -48,6 +48,7 @@ _resolve_qwen_gdn_kv_cache_args, _sm70_compile_graph_slice_dim, _sm70_dump_gdn_projection_tensor, + _sm70_gdn_qpn8_ba_dispatch_eligible, ) from vllm.model_executor.layers.mamba.mamba_utils import ( MambaStateCopyFunc, @@ -74,6 +75,7 @@ Qwen3_5MoeConfig, Qwen3_5MoeTextConfig, ) +from vllm.triton_utils import tl, triton from vllm.utils.torch_utils import ( _encode_layer_name, ) @@ -116,6 +118,98 @@ logger = init_logger(__name__) + +@triton.jit +def _sm70_qwen35_gdn_split_kernel( + mixed_qkvz, + mixed_ba, + output, + stride_qkvz_row: tl.int64, + stride_ba_row: tl.int64, + num_rows: tl.constexpr, + qkv_size: tl.constexpr, + z_size: tl.constexpr, + ba_size: tl.constexpr, + BLOCK_Z: tl.constexpr, + BLOCK_BA: tl.constexpr, +): + row = tl.program_id(0) + + z_cols = tl.arange(0, BLOCK_Z) + z_mask = z_cols < z_size + z = tl.load( + mixed_qkvz + row * stride_qkvz_row + qkv_size + z_cols, + mask=z_mask, + ) + tl.store(output + row * z_size + z_cols, z, mask=z_mask) + + ba_cols = tl.arange(0, BLOCK_BA) + ba_mask = ba_cols < ba_size + b = tl.load(mixed_ba + row * stride_ba_row + ba_cols, mask=ba_mask) + a = tl.load( + mixed_ba + row * stride_ba_row + ba_size + ba_cols, + mask=ba_mask, + ) + z_numel = num_rows * z_size + ba_numel = num_rows * ba_size + tl.store(output + z_numel + row * ba_size + ba_cols, b, mask=ba_mask) + tl.store( + output + z_numel + ba_numel + row * ba_size + ba_cols, + a, + mask=ba_mask, + ) + + +def _sm70_materialize_qwen35_gdn_splits( + mixed_qkvz: torch.Tensor, + mixed_ba: torch.Tensor, + qkv_size: int, + z_size: int, + ba_size: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Materialize Qwen3.5 z/b/a tails with one bitwise copy kernel.""" + if mixed_qkvz.ndim != 2 or mixed_ba.ndim != 2: + raise ValueError("mixed_qkvz and mixed_ba must both be rank two") + if mixed_qkvz.shape[0] != mixed_ba.shape[0]: + raise ValueError("mixed_qkvz and mixed_ba row counts must match") + if mixed_qkvz.dtype != mixed_ba.dtype: + raise ValueError("mixed_qkvz and mixed_ba dtypes must match") + if mixed_qkvz.stride(1) != 1 or mixed_ba.stride(1) != 1: + raise ValueError("mixed_qkvz and mixed_ba must be contiguous by row") + if mixed_qkvz.shape[1] < qkv_size + z_size: + raise ValueError("mixed_qkvz is narrower than qkv_size + z_size") + if mixed_ba.shape[1] < 2 * ba_size: + raise ValueError("mixed_ba is narrower than 2 * ba_size") + + num_rows = mixed_qkvz.shape[0] + z_numel = num_rows * z_size + ba_numel = num_rows * ba_size + packed = torch.empty( + z_numel + 2 * ba_numel, + dtype=mixed_qkvz.dtype, + device=mixed_qkvz.device, + ) + _sm70_qwen35_gdn_split_kernel[(num_rows,)]( + mixed_qkvz, + mixed_ba, + packed, + mixed_qkvz.stride(0), + mixed_ba.stride(0), + num_rows=num_rows, + qkv_size=qkv_size, + z_size=z_size, + ba_size=ba_size, + BLOCK_Z=triton.next_power_of_2(z_size), + BLOCK_BA=triton.next_power_of_2(ba_size), + num_warps=8, + num_stages=1, + ) + z = packed[:z_numel].view(num_rows, z_size) + b = packed[z_numel : z_numel + ba_numel].view(num_rows, ba_size) + a = packed[z_numel + ba_numel :].view(num_rows, ba_size) + return z, b, a + + def _parse_sm70_moe_dense_allowlist() -> set[str] | None: raw = envs.VLLM_SM70_MOE_DENSE_ALLOWLIST if raw is None: @@ -146,15 +240,63 @@ def add_ignored_modules(value: object) -> None: raw_config = getattr(quant_config, "config", None) if isinstance(raw_config, dict): add_ignored_modules(raw_config.get("ignore")) - if not ignored_modules: - return False - return any( + if any( module_name == "linear_attn" or module_name.endswith(".linear_attn") or ("linear_attn.in_proj_a" in module_name) or ("linear_attn.in_proj_b" in module_name) for module_name in ignored_modules - ) + ): + return True + + # ModelOpt mixed checkpoints describe quantized modules positively in + # ``quantized_layers`` instead of listing the unquantized b/a projections + # in an ignore list. A fused qkv/z/b/a Linear would then allocate the b/a + # rows in the qkv/z quant format and leave their scale slots uninitialized. + # Compare algorithms per GDN layer and split whenever b/a is absent or + # differs from the quantized qkv/z group. + quantized_layers = getattr(quant_config, "quantized_layers", None) + if not isinstance(quantized_layers, dict): + return False + + per_layer: dict[str, dict[str, str]] = {} + projection_names = { + "in_proj_qkv", + "in_proj_z", + "in_proj_b", + "in_proj_a", + } + marker = ".linear_attn." + for module_name, layer_info in quantized_layers.items(): + if marker not in module_name or not isinstance(layer_info, dict): + continue + layer_prefix, projection_name = module_name.rsplit(marker, 1) + if projection_name not in projection_names: + continue + quant_algo = str(layer_info.get("quant_algo", "")).upper() + if quant_algo: + per_layer.setdefault(layer_prefix, {})[projection_name] = quant_algo + + for projection_algos in per_layer.values(): + qkvz_algos = { + projection_algos[name] + for name in ("in_proj_qkv", "in_proj_z") + if name in projection_algos + } + if not qkvz_algos: + continue + ba_algos = { + projection_algos[name] + for name in ("in_proj_b", "in_proj_a") + if name in projection_algos + } + has_ba_pair = all( + name in projection_algos for name in ("in_proj_b", "in_proj_a") + ) + if not has_ba_pair or ba_algos != qkvz_algos: + return True + + return False def _get_default_sm70_dense_force_suffixes(tp_size: int) -> set[str]: @@ -248,7 +390,7 @@ def create_ba_proj( def forward_cuda( self, hidden_states: torch.Tensor, - output: torch.Tensor, + output: torch.Tensor | None, ): num_tokens = hidden_states.size(0) layer_name = _encode_layer_name(self.prefix) @@ -260,7 +402,52 @@ def forward_cuda( "gdn_hidden_states", layer_name, hidden_states ) - if self.use_split_input_projections: + if _sm70_gdn_qpn8_ba_dispatch_eligible( + self, + hidden_states, + layer_name, + ): + mixed_qkv = torch.empty( + (num_tokens, qkv_size), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + z = torch.empty( + (num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + b = torch.empty( + (num_tokens, ba_size), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + a = torch.empty_like(b) + qkvz_staging = torch.empty( + (num_tokens, qkv_size + z_size), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + ba_staging = torch.empty( + (num_tokens, ba_size * 2), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + assert self.in_proj_ba is not None + torch.ops._C.fp8_qpn8_dispatch_ba_split_sm70_out( + mixed_qkv, + z, + b, + a, + qkvz_staging, + ba_staging, + int(self.in_proj_qkvz.sm70_fp8_prefill_exact_dense_workspace_ptr), + hidden_states, + self.in_proj_qkvz.weight, + self.in_proj_qkvz.weight_scale_inv, + self.in_proj_ba.weight, + ) + elif self.use_split_input_projections: mixed_qkvz, _ = self.in_proj_qkvz(hidden_states) assert self.in_proj_ba is not None ba, _ = self.in_proj_ba(hidden_states) @@ -269,9 +456,18 @@ def forward_cuda( ) ba = _sm70_dump_gdn_projection_tensor("in_proj_ba", layer_name, ba) mixed_qkv = mixed_qkvz[..., :qkv_size] - z = _sm70_compile_graph_slice_dim(mixed_qkvz, -1, qkv_size, z_size) - b = ba[..., :ba_size] - a = _sm70_compile_graph_slice_dim(ba, -1, ba_size, ba_size) + if self.enable_sm70_dflash2_fused_gdn_split: + z, b, a = _sm70_materialize_qwen35_gdn_splits( + mixed_qkvz, + ba, + qkv_size, + z_size, + ba_size, + ) + else: + z = _sm70_compile_graph_slice_dim(mixed_qkvz, -1, qkv_size, z_size) + b = ba[..., :ba_size] + a = _sm70_compile_graph_slice_dim(ba, -1, ba_size, ba_size) else: mixed_qkvzba, _ = self.in_proj_qkvz(hidden_states) mixed_qkvzba = _sm70_dump_gdn_projection_tensor( @@ -281,15 +477,9 @@ def forward_cuda( ba_start = z_start + z_size a_start = ba_start + ba_size mixed_qkv = mixed_qkvzba[..., :qkv_size] - z = _sm70_compile_graph_slice_dim( - mixed_qkvzba, -1, z_start, z_size - ) - b = _sm70_compile_graph_slice_dim( - mixed_qkvzba, -1, ba_start, ba_size - ) - a = _sm70_compile_graph_slice_dim( - mixed_qkvzba, -1, a_start, ba_size - ) + z = _sm70_compile_graph_slice_dim(mixed_qkvzba, -1, z_start, z_size) + b = _sm70_compile_graph_slice_dim(mixed_qkvzba, -1, ba_start, ba_size) + a = _sm70_compile_graph_slice_dim(mixed_qkvzba, -1, a_start, ba_size) mixed_qkv = _sm70_dump_gdn_projection_tensor( "split_mixed_qkv", layer_name, mixed_qkv @@ -320,7 +510,7 @@ def forward_cuda( conv_state_cache=conv_state_cache, ssm_state_cache=ssm_state_cache, ) - self._output_projection(core_attn_out, z, output, num_tokens) + return self._output_projection(core_attn_out, z, output, num_tokens) class Qwen3_5DecoderLayer(Qwen3NextDecoderLayer): 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_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index 63ddb51248..554d0459a8 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import os +import io from collections.abc import Iterable import torch @@ -10,10 +10,12 @@ from transformers import Qwen3Config from vllm import _custom_ops as ops -from vllm import envs from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config -from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) from vllm.logger import init_logger from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.layernorm import RMSNorm @@ -35,12 +37,17 @@ ) from vllm.multimodal.inputs import NestedTensors from vllm.transformers_utils.config import set_default_rope_theta +from vllm.transformers_utils.repo_utils import get_hf_file_bytes from vllm.v1.attention.backend import AttentionType +from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( + get_eagle3_aux_layers_from_config, +) from .qwen2 import Qwen2MLP as Qwen3MLP from .qwen3 import Qwen3ForCausalLM from .utils import ( AutoWeightsLoader, + WeightsMapper, get_draft_quant_config, maybe_prefix, process_eagle_weight, @@ -49,43 +56,121 @@ logger = init_logger(__name__) -def _get_dflash_per_layer_sliding_window( - config: Qwen3Config, - layer_idx: int, -) -> int | None: - layer_types = getattr(config, "layer_types", None) - if layer_types is None: - return None - layer_type = layer_types[layer_idx] - if layer_type == "sliding_attention": - return getattr(config, "sliding_window", None) - if layer_type == "full_attention": - return None - raise ValueError(f"Invalid DFlash layer_type {layer_type}") - - -def _debug_sync_context_kv_enabled() -> bool: - return envs.VLLM_DFLASH_SYNC_CONTEXT_KV - +_SLIDING_ATTENTION = "sliding_attention" -def _skip_context_kv_precompute_enabled() -> bool: - return envs.VLLM_DFLASH_SKIP_CONTEXT_KV_PRECOMPUTE +def _dflash_layer_causal(config: Qwen3Config, layer_idx: int) -> bool: + """Resolve explicit causality before falling back to legacy layer defaults.""" + is_causal = getattr(config, "is_causal", None) + if is_causal is not None: + return bool(is_causal) + override = (getattr(config, "dflash_config", None) or {}).get("causal") + if override is not None: + return bool(override) + layer_types = getattr(config, "layer_types", None) + return bool(layer_types) and layer_types[layer_idx] == _SLIDING_ATTENTION + + +def dflash_has_any_non_causal(config: Qwen3Config) -> bool: + """Whether the draft needs a non-causal-capable backend, resolved from config + (config mirror of the model's ``get_draft_attn_causal``, usable pre-build).""" + return not all( + _dflash_layer_causal(config, i) for i in range(config.num_hidden_layers) + ) + + +def dflash_target_rope_is_neox_style(target_model: nn.Module) -> bool | None: + """The target's RoPE layout, from its first attention layer. + + A DFlash head must rotate Q/K the way the target it was distilled against + does, and a mismatch is silent — acceptance collapses but nothing errors and + the output stays correct. Draft checkpoints do not carry this, so take it + from the target. None if the target uses no RoPE. + """ + language_model = ( + target_model.get_language_model() + if hasattr(target_model, "get_language_model") + else target_model + ) + for module in language_model.modules(): + style = getattr(module, "is_neox_style", None) + if isinstance(style, bool): + return style + return None + + +def _get_dflash_fc_input_size(vllm_config: VllmConfig) -> int: + spec_config = vllm_config.speculative_config + config = spec_config.draft_model_config.hf_config + aux_layers = get_eagle3_aux_layers_from_config(spec_config) + num_features_to_use = len(aux_layers) if aux_layers else config.num_hidden_layers + target_hidden_size = ( + getattr(config, "target_hidden_size", None) or config.hidden_size + ) + return target_hidden_size * num_features_to_use + + +def _resolve_layer_attention( + config: Qwen3Config, layer_idx: int +) -> tuple[int | None, bool]: + """Resolve ``(sliding_window, causal)`` for one DFlash draft layer. + + +----------------------+-------------------------+--------------------------------+ + | Config | ``layer_type`` | *``causal`` | + +======================+=========================+================================+ + | ``layer_types`` | SWA if ``use_swa`` | True if ``layer_types[i]=SWA`` | + | | else ``layer_types[i]`` | else False | + +----------------------+-------------------------+--------------------------------+ + | ``layer_types=None`` | SWA | False | + | + ``use_swa=True`` | | | + +----------------------+-------------------------+--------------------------------+ + | ``layer_types=None`` | Full | False | + | + ``use_swa=False`` | | | + +----------------------+-------------------------+--------------------------------+ + * If ``dflash_config.causal`` is set, its value overrides ``causal`` for all layers. + + This is to support a varied ecosystem of checkpoints, including: + - XiaomiMiMo/MiMo-V2.5-Pro-FP4-DFlash (sets "use_swa", assumes non-causal) + - z-lab/gemma-4-31B-it-DFlash (has mixed layer types, assumes causal only for SWA) + - z-lab/Qwen3.5-9B-DFlash ("standard" DFlash, all full attn, assumes non-causal) + """ + dflash_config = getattr(config, "dflash_config", None) or {} + layer_types = getattr(config, "layer_types", None) + use_swa = dflash_config.get("use_swa", False) -def _debug_context_kv_enabled() -> bool: - return envs.VLLM_DFLASH_DEBUG_CONTEXT_KV - - -def _dump_draft_layer_hiddens_enabled() -> bool: - return envs.VLLM_DFLASH_DUMP_LAYER_HIDDENS - + any_sliding = False + if layer_types is not None: + num_sliding = sum(lt == _SLIDING_ATTENTION for lt in layer_types) + any_sliding = num_sliding > 0 + # Mixed sliding/full attention needs multiple KV groups (V2 runner only). + if ( + 0 < num_sliding < len(layer_types) + and not get_current_vllm_config().use_v2_model_runner + ): + raise NotImplementedError( + "DFlash drafters with mixed sliding/full attention require " + "the V2 model runner; relaunch with " + "VLLM_USE_V2_MODEL_RUNNER=1." + ) -def _dump_draft_layer0_components_enabled() -> bool: - return envs.VLLM_DFLASH_DUMP_LAYER0_COMPONENTS + # ``use_swa`` forces SWA on every layer, even an all-full ``layer_types``. + if layer_types is None or (use_swa and not any_sliding): + is_sliding = use_swa + else: + is_sliding = layer_types[layer_idx] == _SLIDING_ATTENTION + sliding_window = None + if is_sliding: + sliding_window = dflash_config.get( + "swa_window_size", getattr(config, "sliding_window", None) + ) + if sliding_window is None: + raise ValueError( + "DFlash sliding attention requires a window size configured in " + "dflash_config.swa_window_size or the top-level sliding_window." + ) -def _dump_draft_attn_components_enabled() -> bool: - return envs.VLLM_DFLASH_DUMP_ATTN_COMPONENTS + return sliding_window, _dflash_layer_causal(config, layer_idx) class DFlashQwen3Attention(nn.Module): @@ -105,16 +190,18 @@ def __init__( head_dim: int | None = None, rms_norm_eps: float = 1e-06, attention_bias: bool = False, + add_swa_attention_sink_bias: bool = False, + sliding_window: int | None = None, + causal: bool = False, + is_neox_style: bool = True, cache_config: CacheConfig | None = None, quant_config: QuantizationConfig | None = None, prefix: str = "", attn_type: str = AttentionType.DECODER, - per_layer_sliding_window: int | None = None, ) -> None: super().__init__() self.layer_name = prefix self.hidden_size = hidden_size - self.per_layer_sliding_window = per_layer_sliding_window tp_size = get_tensor_model_parallel_world_size() self.total_num_heads = num_heads assert self.total_num_heads % tp_size == 0 @@ -150,8 +237,17 @@ def __init__( self.rotary_emb = get_rope( self.head_dim, max_position=max_position, + is_neox_style=is_neox_style, rope_parameters=rope_parameters, ) + + self.attention_sink_bias = ( + torch.nn.Parameter(torch.empty(self.num_heads), requires_grad=False) + if add_swa_attention_sink_bias + else None + ) + + self.sliding_window = sliding_window self.attn = Attention( self.num_heads, self.head_dim, @@ -159,14 +255,17 @@ def __init__( num_kv_heads=self.num_kv_heads, cache_config=cache_config, quant_config=quant_config, + per_layer_sliding_window=sliding_window, prefix=f"{prefix}.attn", attn_type=attn_type, - per_layer_sliding_window=per_layer_sliding_window, + sinks=self.attention_sink_bias, ) + # The in-tree FlashAttention-V100 backend uses this marker to select + # its DFlash-compatible non-causal paged-prefill fallback. self.attn.is_dflash_draft_attn = True + self.causal = causal self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) - self._attn_components_dumped = False def forward( self, @@ -177,7 +276,7 @@ def forward( with the context K/V from the target model's hidden states. This forward op computes attention for the query tokens only. See also: precompute_and_store_context_kv""" - qkv = F.linear(hidden_states, self.qkv_proj.weight, self.qkv_proj.bias) + qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) # Per-head RMSNorm @@ -192,64 +291,67 @@ def forward( q, k = self.rotary_emb(positions, q, k) attn_output = self.attn(q, k, v) + output_input_scale = getattr(self, "output_input_scale", 1.0) + if output_input_scale != 1.0: + attn_output = attn_output / output_input_scale output, _ = self.o_proj(attn_output) - if ( - _dump_draft_attn_components_enabled() - and not self._attn_components_dumped - and positions.shape[-1] == 17 - and int(positions.max().item()) > 0 - and ".layers.32.self_attn" in self.layer_name - ): - dump_path = f"/tmp/dflash_draft_attn_components_pid{os.getpid()}.pt" - torch.save( - { - "positions": positions.detach().cpu(), - "q": q.detach().to(torch.float16).cpu(), - "k": k.detach().to(torch.float16).cpu(), - "v": v.detach().to(torch.float16).cpu(), - "attn_output": attn_output.detach().to(torch.float16).cpu(), - "o_proj_output": output.detach().to(torch.float16).cpu(), - }, - dump_path, - ) - self._attn_components_dumped = True - logger.warning("Saved DFlash draft attn components to %s", dump_path) return output class DFlashQwen3DecoderLayer(nn.Module): + attention_cls = DFlashQwen3Attention + def __init__( self, vllm_config: VllmConfig, *, config: Qwen3Config, + layer_idx: int, cache_config: CacheConfig | None = None, quant_config: QuantizationConfig | None = None, - layer_idx: int, prefix: str = "", ) -> None: super().__init__() self.hidden_size = config.hidden_size set_default_rope_theta(config, default_theta=1000000) attn_type = AttentionType.DECODER - per_layer_sliding_window = _get_dflash_per_layer_sliding_window( - config, layer_idx + + # DFlash drafts store the sink-bias flag inside dflash_config; fall back + # to the top-level attribute used by other (e.g. MiMo) configs. + dflash_config = getattr(config, "dflash_config", None) or {} + add_swa_attention_sink_bias = dflash_config.get( + "attention_sink_bias", + getattr(config, "add_swa_attention_sink_bias", False), ) - self.self_attn = DFlashQwen3Attention( + # Resolve this layer's attention mode (full vs sliding window, causal vs + # non-causal) from the draft config. + sliding_window, causal = _resolve_layer_attention(config, layer_idx) + + # RoPE layout, copied off the target at load time by the draft loader + # (see `dflash_target_rope_is_neox_style`). Checkpoints do not carry it: + # a head distilled from an interleaved-RoPE target must rotate the way + # that target does, or every drafted Q/K is wrong and acceptance + # collapses with no error raised. + is_neox_style = getattr(config, "is_neox_style", True) + + self.self_attn = self.attention_cls( hidden_size=self.hidden_size, num_heads=config.num_attention_heads, max_position=config.max_position_embeddings, num_kv_heads=config.num_key_value_heads, rms_norm_eps=config.rms_norm_eps, attention_bias=getattr(config, "attention_bias", False), + add_swa_attention_sink_bias=add_swa_attention_sink_bias, + sliding_window=sliding_window, + causal=causal, + is_neox_style=is_neox_style, head_dim=getattr(config, "head_dim", None), cache_config=cache_config, quant_config=quant_config, rope_parameters=config.rope_parameters, prefix=f"{prefix}.self_attn", attn_type=attn_type, - per_layer_sliding_window=per_layer_sliding_window, ) self.mlp = Qwen3MLP( hidden_size=self.hidden_size, @@ -262,7 +364,6 @@ def __init__( self.post_attention_layernorm = RMSNorm( config.hidden_size, eps=config.rms_norm_eps ) - self._layer0_components_dumped = False def forward( self, @@ -275,49 +376,50 @@ def forward( else: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) - ln_out = hidden_states hidden_states = self.self_attn( positions=positions, hidden_states=hidden_states, ) - attn_out = hidden_states hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) - post_attn_ln_out = hidden_states hidden_states = self.mlp(hidden_states) - if ( - _dump_draft_layer0_components_enabled() - and not self._layer0_components_dumped - and positions.shape[-1] == 17 - and int(positions.max().item()) > 0 - and ".layers.32.self_attn" in self.self_attn.layer_name - ): - mlp_gate_up, _ = self.mlp.gate_up_proj(post_attn_ln_out) - mlp_act_out = self.mlp.act_fn(mlp_gate_up) - dump_path = f"/tmp/dflash_draft_layer0_components_pid{os.getpid()}.pt" - torch.save( - { - "positions": positions.detach().cpu(), - "input_layernorm_out": ln_out.detach().to(torch.float16).cpu(), - "self_attn_out": attn_out.detach().to(torch.float16).cpu(), - "post_attention_layernorm_out": post_attn_ln_out.detach() - .to(torch.float16) - .cpu(), - "mlp_gate_up": mlp_gate_up.detach().to(torch.float16).cpu(), - "mlp_act_out": mlp_act_out.detach().to(torch.float16).cpu(), - "mlp_out": hidden_states.detach().to(torch.float16).cpu(), - "residual": residual.detach().to(torch.float16).cpu(), - }, - dump_path, - ) - self._layer0_components_dumped = True - logger.warning("Saved DFlash draft layer0 components to %s", dump_path) return hidden_states, residual @support_torch_compile class DFlashQwen3Model(nn.Module): + decoder_layer_cls = DFlashQwen3DecoderLayer + + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_substr={ + "midlayer.": "layers.0.", + # Muse-Glimmer-30B-assistant names the aux-hidden-state encoder + # `encoder.fc` / `encoder.output_norm_enc`; this head calls them + # `fc` / `hidden_norm`. Same tensors and shapes, different names. + "encoder.output_norm_enc.": "hidden_norm.", + "encoder.fc.": "fc.", + }, + ) + + def _make_context_projection( + self, + *, + vllm_config: VllmConfig, + input_size: int, + output_size: int, + prefix: str, + ) -> nn.Module: + return ReplicatedLinear( + input_size=input_size, + output_size=output_size, + bias=False, + params_dtype=vllm_config.model_config.dtype, + quant_config=self.quant_config, + prefix=prefix, + return_bias=False, + ) + def __init__( self, *, @@ -346,37 +448,37 @@ def __init__( prefix=maybe_prefix(prefix, "embed_tokens"), ) + # Masked query slots are fed to the draft as `mask_token_id`. Most DFlash + # checkpoints will have the mask embedding in the vocabulary embedding table + # at that slot id. Some checkpoints (XiaomiMiMo/MiMo-V2.5-Pro-FP4-DFlash) ship + # with a separate mask embedding tensor to use instead. When present, we load it + # and substitute it for embed_tokens[mask_token_id] when computing embeddings. + self.mask_token_id = drafter_config.get("mask_token_id") + self.mask_embedding = nn.Parameter( + torch.zeros(self.config.hidden_size, dtype=vllm_config.model_config.dtype), + requires_grad=False, + ) + self.has_separate_mask_embedding = False + self.layers = nn.ModuleList( [ - DFlashQwen3DecoderLayer( + self.decoder_layer_cls( current_vllm_config, config=self.config, + layer_idx=layer_idx, cache_config=current_vllm_config.cache_config, quant_config=self.quant_config, - layer_idx=layer_idx, prefix=maybe_prefix(prefix, f"layers.{layer_idx + start_layer_id}"), ) for layer_idx in range(self.config.num_hidden_layers) ] ) if self.use_aux_hidden_state: - num_features_to_use = self.config.num_hidden_layers - if "target_layer_ids" in drafter_config: - num_features_to_use = len(drafter_config["target_layer_ids"]) - elif "layer_ids" in drafter_config: - num_features_to_use = len(drafter_config["layer_ids"]) - if hasattr(self.config, "target_hidden_size"): - fc_input_size = self.config.target_hidden_size * num_features_to_use - else: - fc_input_size = self.config.hidden_size * num_features_to_use - self.fc = ReplicatedLinear( - input_size=fc_input_size, + self.fc = self._make_context_projection( + vllm_config=vllm_config, + input_size=_get_dflash_fc_input_size(vllm_config), output_size=self.config.hidden_size, - bias=False, - params_dtype=vllm_config.model_config.dtype, - quant_config=self.quant_config, prefix=maybe_prefix(prefix, "fc"), - return_bias=False, ) self.hidden_norm = RMSNorm( self.config.hidden_size, @@ -386,11 +488,43 @@ def __init__( self.config.hidden_size, eps=self.config.rms_norm_eps, ) - self._debug_context_kv_dumped = False - self._draft_layer_hiddens_dumped = False def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.embed_tokens(input_ids) + embeds = self.embed_tokens(input_ids) + if self.has_separate_mask_embedding and self.mask_token_id is not None: + # Replace masked slots with the dedicated mask embedding. + is_mask = (input_ids == self.mask_token_id).unsqueeze(-1) + embeds = torch.where(is_mask, self.mask_embedding.to(embeds.dtype), embeds) + return embeds + + def _build_context_kv_buffers( + self, + layers_attn: list[nn.Module], + has_bias: bool, + ) -> None: + self._hidden_norm_weight = self.hidden_norm.weight.data + + # KV projection weights: [num_layers * 2 * kv_size, hidden_size] + kv_weights = [a.qkv_proj.weight[a.q_size :] for a in layers_attn] + self._fused_kv_weight = torch.cat(kv_weights, dim=0) + if has_bias: + kv_biases = [a.qkv_proj.bias[a.q_size :] for a in layers_attn] + self._fused_kv_bias: torch.Tensor | None = torch.cat(kv_biases, dim=0) + else: + self._fused_kv_bias = None + + # K-norm weights stacked into one contiguous [num_layers, head_dim] + # tensor so the per-layer K-norm runs as a single grouped kernel. + self._k_norm_weights = torch.stack( + [a.k_norm.weight.data for a in layers_attn], dim=0 + ).contiguous() + # The grouped call below relies on the stable-ABI RMSNorm extension + # selecting one weight row per outer (layer) index. Older binaries + # silently accepted the 2-D tensor but applied row zero to every + # layer, which severely lowers DFlash acceptance without affecting + # target-only output quality. Verify the loaded binary once before + # trusting its grouped result. + self._batched_k_norm_runtime_verified: bool | None = None def _build_fused_kv_buffers(self) -> None: """Build fused weight buffers for precompute_and_store_context_kv. @@ -404,19 +538,7 @@ def _build_fused_kv_buffers(self) -> None: attn0 = layers_attn[0] has_bias = attn0.qkv_proj.bias is not None - self._hidden_norm_weight = self.hidden_norm.weight.data - - # KV projection weights: [num_layers * 2 * kv_size, hidden_size] - kv_weights = [a.qkv_proj.weight[a.q_size :] for a in layers_attn] - self._fused_kv_weight = torch.cat(kv_weights, dim=0) - if has_bias: - kv_biases = [a.qkv_proj.bias[a.q_size :] for a in layers_attn] - self._fused_kv_bias: torch.Tensor | None = torch.cat(kv_biases, dim=0) - else: - self._fused_kv_bias = None - - # K-norm weights: list of [head_dim] tensors, one per layer. - self._k_norm_weights = [a.k_norm.weight.data for a in layers_attn] + self._build_context_kv_buffers(layers_attn, has_bias) # RoPE parameters self._rope_head_size = attn0.rotary_emb.head_size @@ -447,11 +569,112 @@ def _build_fused_kv_buffers(self) -> None: # References to inner Attention layers for direct cache writes self._attn_layers = [layer.self_attn.attn for layer in self.layers] + def _project_context_kv( + self, + context_states: torch.Tensor, + num_ctx: int, + num_layers: int, + num_kv_heads: int, + head_dim: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + # --- Fused KV projection (one GEMM for all layers) --- + normed_context_states = self._normalize_context_states(context_states) + all_kv_flat = F.linear( + normed_context_states, self._fused_kv_weight, self._fused_kv_bias + ) + # Single contiguous copy that separates K/V and transposes to + # layer-major layout. Result: [2, L, num_ctx, nkv, hd] contiguous. + # Indexing dim-0 gives contiguous [L, num_ctx, nkv, hd] for K and V. + all_kv = ( + all_kv_flat.view(num_ctx, num_layers, 2, num_kv_heads, head_dim) + .permute(2, 1, 0, 3, 4) + .contiguous() + ) + all_k = all_kv[0] # [L, num_ctx, nkv, hd], contiguous + all_v = all_kv[1] # [L, num_ctx, nkv, hd], contiguous + return all_k, all_v + + def _normalize_context_states(self, context_states: torch.Tensor) -> torch.Tensor: + normed_context_states = torch.empty_like(context_states) + ops.rms_norm( + normed_context_states, + context_states, + self._hidden_norm_weight, + self._rms_norm_eps, + ) + return normed_context_states + + def _normalize_context_k_per_layer(self, all_k: torch.Tensor) -> torch.Tensor: + """Reference path that works with pre-batched-weight RMSNorm binaries.""" + all_k_normed = torch.empty_like(all_k) + for layer_idx in range(all_k.shape[0]): + ops.rms_norm( + all_k_normed[layer_idx], + all_k[layer_idx], + self._k_norm_weights[layer_idx], + self._rms_norm_eps, + ) + return all_k_normed + + def _verify_batched_context_k_norm_runtime( + self, dtype: torch.dtype, device: torch.device + ) -> bool: + """Check the loaded stable extension with nonzero per-layer inputs.""" + num_layers, hidden_size = self._k_norm_weights.shape + probe = torch.ones( + (num_layers, 1, 1, hidden_size), + dtype=dtype, + device=device, + ) + grouped = torch.empty_like(probe) + ops.rms_norm( + grouped, + probe, + self._k_norm_weights, + self._rms_norm_eps, + ) + per_layer = self._normalize_context_k_per_layer(probe) + return torch.equal(grouped, per_layer) + + def _normalize_context_k(self, all_k: torch.Tensor) -> torch.Tensor: + # --- Grouped RMSNorm K across all layers ([L, num_ctx, nkv, hd]) --- + # The weight is selected per layer by the outermost (layer) index. + runtime_verified = getattr(self, "_batched_k_norm_runtime_verified", None) + if runtime_verified is None: + # This one-time comparison is normally consumed by model profiling + # or warmup. Use synthetic nonzero inputs so an all-zero dummy run + # cannot falsely approve a stale binary. The equality check + # deliberately synchronizes once: silently using the wrong K + # weights is a much larger failure than a cold startup fence. + runtime_verified = self._verify_batched_context_k_norm_runtime( + all_k.dtype, all_k.device + ) + self._batched_k_norm_runtime_verified = runtime_verified + if not runtime_verified: + logger.warning_once( + "The loaded stable-ABI RMSNorm extension does not preserve " + "DFlash per-layer batched weights. Falling back to the " + "bitwise per-layer path; rebuild the vLLM stable extension " + "to restore the grouped K-normalization fast path." + ) + + if not runtime_verified: + return self._normalize_context_k_per_layer(all_k) + + grouped = torch.empty_like(all_k) + ops.rms_norm( + grouped, + all_k, + self._k_norm_weights, + self._rms_norm_eps, + ) + return grouped + def precompute_and_store_context_kv( self, context_states: torch.Tensor, context_positions: torch.Tensor, - context_slot_mapping: torch.Tensor | dict[str, torch.Tensor] | None = None, + context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None, ) -> None: """Precompute K/V for context states write them into each layer's KV cache. @@ -465,10 +688,6 @@ def precompute_and_store_context_kv( When context_slot_mapping is None (e.g. during dummy_run) only the computation runs, and no K/V is written to cache. """ - if _skip_context_kv_precompute_enabled(): - logger.warning("Skipping DFlash context KV precompute for debugging.") - return - if not hasattr(self, "_num_attn_layers"): logger.warning_once( "DFlash buffer initialization was skipped. If dummy weights are not " @@ -482,35 +701,8 @@ def precompute_and_store_context_kv( hd = self._head_dim nkv = self._num_kv_heads - # --- Fused KV projection (one GEMM for all layers) --- - normed_context_states = torch.empty_like(context_states) - ops.rms_norm( - normed_context_states, - context_states, - self._hidden_norm_weight, - self._rms_norm_eps, - ) - all_kv_flat = F.linear( - normed_context_states, self._fused_kv_weight, self._fused_kv_bias - ) - # Single contiguous copy that separates K/V and transposes to - # layer-major layout. Result: [2, L, num_ctx, nkv, hd] contiguous. - # Indexing dim-0 gives contiguous [L, num_ctx, nkv, hd] for K and V. - all_kv = ( - all_kv_flat.view(num_ctx, L, 2, nkv, hd).permute(2, 1, 0, 3, 4).contiguous() - ) - all_k = all_kv[0] # [L, num_ctx, nkv, hd], contiguous - all_v = all_kv[1] # [L, num_ctx, nkv, hd], contiguous - - # --- Per-layer RMSNorm K (3D: [num_ctx, nkv, hd] per layer) --- - all_k_normed = torch.empty_like(all_k) - for i in range(L): - ops.rms_norm( - all_k_normed[i], - all_k[i], - self._k_norm_weights[i], - self._rms_norm_eps, - ) + all_k, all_v = self._project_context_kv(context_states, num_ctx, L, nkv, hd) + all_k_normed = self._normalize_context_k(all_k) # --- Fused RoPE across all layers --- # View as [L * num_ctx, kv] so RoPE sees one big batch (no copy). @@ -534,145 +726,22 @@ def precompute_and_store_context_kv( # --- Per-layer cache insert --- all_k_final = all_k_flat.view(L, num_ctx, nkv, hd) - debug_direct_refs: dict[ - int, list[tuple[str, int, torch.Tensor, torch.Tensor]] - ] = {} - if _debug_context_kv_enabled() and not self._debug_context_kv_dumped: - sample_tokens = min(8, num_ctx) - debug_ranges: list[tuple[str, int, int]] = [("head", 0, sample_tokens)] - if num_ctx > sample_tokens: - tail_start = max(0, num_ctx - sample_tokens) - if tail_start > 0: - debug_ranges.append(("tail", tail_start, num_ctx)) - - debug_cos_sin_cache = self._rope_cos_sin_cache - for layer_idx in sorted({0, L // 2, L - 1}): - attn_layer = self.layers[layer_idx].self_attn - layer_refs: list[tuple[str, int, torch.Tensor, torch.Tensor]] = [] - for sample_name, sample_start, sample_end in debug_ranges: - sample_len = sample_end - sample_start - direct_kv = F.linear( - normed_context_states[sample_start:sample_end], - attn_layer.qkv_proj.weight[attn_layer.q_size :], - ( - None - if attn_layer.qkv_proj.bias is None - else attn_layer.qkv_proj.bias[attn_layer.q_size :] - ), - ) - direct_k, direct_v = direct_kv.split([kv, kv], dim=-1) - direct_k = direct_k.view(sample_len, nkv, hd).contiguous() - direct_v = direct_v.view(sample_len, nkv, hd).contiguous() - direct_k_norm = torch.empty_like(direct_k) - ops.rms_norm( - direct_k_norm, - direct_k, - attn_layer.k_norm.weight.data, - self._rms_norm_eps, - ) - direct_k_flat = direct_k_norm.view(sample_len, kv) - if debug_cos_sin_cache.dtype != direct_k_flat.dtype: - debug_cos_sin_cache = debug_cos_sin_cache.to( - dtype=direct_k_flat.dtype - ) - ops.rotary_embedding( - context_positions[sample_start:sample_end], - direct_k_flat, - None, - self._rope_head_size, - debug_cos_sin_cache, - self._rope_is_neox, - ) - direct_k = direct_k_flat.view(sample_len, nkv, hd) - layer_refs.append((sample_name, sample_start, direct_k, direct_v)) - - fused_k_diff = ( - all_k_final[layer_idx, sample_start:sample_end] - direct_k - ).abs() - fused_v_diff = ( - all_v[layer_idx, sample_start:sample_end] - direct_v - ).abs() - logger.warning( - "DFlash context KV fused-vs-direct layer=%d sample=%s " - "start=%d tokens=%d k_max=%.6f k_mean=%.6f " - "v_max=%.6f v_mean=%.6f", - layer_idx, - sample_name, - sample_start, - sample_len, - float(fused_k_diff.max().item()), - float(fused_k_diff.mean().item()), - float(fused_v_diff.max().item()), - float(fused_v_diff.mean().item()), - ) - debug_direct_refs[layer_idx] = layer_refs - + per_layer = isinstance(context_slot_mapping, (list, tuple)) for i in range(L): + slot_mapping = ( + context_slot_mapping[i] if per_layer else context_slot_mapping + ) + if slot_mapping is None: + continue # dummy run: skip cache ops attn = self._attn_layers[i] - layer_context_slot_mapping = context_slot_mapping - if isinstance(context_slot_mapping, dict): - layer_context_slot_mapping = context_slot_mapping[attn.layer_name] - assert layer_context_slot_mapping is not None kv_cache = attn.kv_cache - if isinstance(kv_cache, list): - kv_cache = kv_cache[0] - if _debug_sync_context_kv_enabled(): - logger.warning( - "DFlash context KV update layer=%s kv_cache_shape=%s " - "slot_min=%s slot_max=%s num_ctx=%d", - attn.layer_name, - tuple(kv_cache.shape), - int(layer_context_slot_mapping.min().item()), - int(layer_context_slot_mapping.max().item()), - num_ctx, - ) attn.impl.do_kv_cache_update( attn, all_k_final[i], all_v[i], kv_cache, - layer_context_slot_mapping, + slot_mapping, ) - if ( - i in debug_direct_refs - and _debug_context_kv_enabled() - and not self._debug_context_kv_dumped - ): - block_size = kv_cache.shape[2] - for ( - sample_name, - sample_start, - direct_k, - direct_v, - ) in debug_direct_refs[i]: - sample_len = direct_k.shape[0] - sample_end = sample_start + sample_len - slots = layer_context_slot_mapping[sample_start:sample_end].to( - torch.long - ) - block_idx = torch.div(slots, block_size, rounding_mode="floor") - block_off = torch.remainder(slots, block_size) - cache_k = kv_cache[block_idx, 0, block_off] - cache_v = kv_cache[block_idx, 1, block_off] - cache_k_diff = (cache_k - direct_k).abs() - cache_v_diff = (cache_v - direct_v).abs() - logger.warning( - "DFlash context KV cache-readback layer=%d sample=%s " - "start=%d tokens=%d k_max=%.6f k_mean=%.6f " - "v_max=%.6f v_mean=%.6f", - i, - sample_name, - sample_start, - sample_len, - float(cache_k_diff.max().item()), - float(cache_k_diff.mean().item()), - float(cache_v_diff.max().item()), - float(cache_v_diff.mean().item()), - ) - if i == max(debug_direct_refs): - self._debug_context_kv_dumped = True - if _debug_sync_context_kv_enabled(): - torch.accelerator.synchronize(all_k_final.device) def forward( self, @@ -686,44 +755,34 @@ def forward( hidden_states = input_embeds residual = None - layer_hidden_states: list[torch.Tensor] | None = None - should_dump_layer_hiddens = ( - _dump_draft_layer_hiddens_enabled() - and not self._draft_layer_hiddens_dumped - and positions.shape[-1] == 17 - and int(positions.max().item()) > 0 - ) - if should_dump_layer_hiddens: - layer_hidden_states = [] for layer in self.layers: hidden_states, residual = layer( positions=positions, hidden_states=hidden_states, residual=residual, ) - if layer_hidden_states is not None: - layer_hidden_states.append( - hidden_states.detach().to(torch.float16).cpu() - ) hidden_states, _ = self.norm(hidden_states, residual) - if layer_hidden_states is not None: - dump_path = f"/tmp/dflash_draft_layers_pid{os.getpid()}.pt" - torch.save( - { - "positions": positions.detach().cpu(), - "input_embeds": input_embeds.detach().to(torch.float16).cpu(), - "layer_hidden_states": layer_hidden_states, - "final_hidden_states": hidden_states.detach() - .to(torch.float16) - .cpu(), - }, - dump_path, - ) - self._draft_layer_hiddens_dumped = True - logger.warning("Saved DFlash draft layer hiddens to %s", dump_path) return hidden_states + def _preprocess( + self, weights: Iterable[tuple[str, torch.Tensor]] + ) -> Iterable[tuple[str, torch.Tensor]]: + tp_size = get_tensor_model_parallel_world_size() + tp_rank = get_tensor_model_parallel_rank() + for name, loaded_weight in weights: + if "attention_sink_bias" in name: + # Sink bias is per-head; shard it across TP ranks like the + # attention heads themselves. + heads_per_rank = loaded_weight.shape[0] // tp_size + loaded_weight = loaded_weight.narrow( + 0, tp_rank * heads_per_rank, heads_per_rank + ) + yield name, loaded_weight + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # onecat's current AutoWeightsLoader predates stacked mapper metadata, + # so preserve the upstream naming contract while passing shard ids + # directly to the existing linear weight loaders. stacked_params_mapping = [ (".qkv_proj", ".q_proj", "q"), (".qkv_proj", ".k_proj", "k"), @@ -733,9 +792,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ] params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() - for name, loaded_weight in weights: - if "midlayer." in name: - name = name.replace("midlayer.", "layers.0.") + for name, loaded_weight in self._preprocess(weights): + mapped_name = self.hf_to_vllm_mapper._map_name(name) + if mapped_name is None: + continue + name = mapped_name if self.quant_config is not None and ( scale_name := self.quant_config.get_cache_scale(name) ): @@ -756,8 +817,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: continue name = name.replace(weight_name, param_name) param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) + param.weight_loader(param, loaded_weight, shard_id) break else: param = params_dict[name] @@ -768,15 +828,18 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: class DFlashQwen3ForCausalLM(Qwen3ForCausalLM): + model_cls = DFlashQwen3Model + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): nn.Module.__init__(self) - self.config = vllm_config.speculative_config.draft_model_config.hf_config + self.draft_model_config = vllm_config.speculative_config.draft_model_config + self.config = self.draft_model_config.hf_config if getattr(self.config, "draft_vocab_size", None) is None: self.config.draft_vocab_size = getattr(self.config, "vocab_size", None) target_layer_num = vllm_config.model_config.get_num_layers( vllm_config.parallel_config ) - self.model = DFlashQwen3Model( + self.model = self.model_cls( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"), start_layer_id=target_layer_num, @@ -816,6 +879,14 @@ def forward( ) -> torch.Tensor: return self.model(input_ids, positions, inputs_embeds) + def get_draft_kv_cache_layer_names(self) -> list[str]: + return [layer.self_attn.attn.layer_name for layer in self.model.layers] + + def get_draft_attn_causal(self) -> list[bool]: + """Per-layer attention causality, aligned with + get_draft_kv_cache_layer_names.""" + return [layer.self_attn.causal for layer in self.model.layers] + def compute_logits( self, hidden_states: torch.Tensor, @@ -838,6 +909,12 @@ def get_topk_tokens_and_logprobs( hidden_states: torch.Tensor, top_k: int, ) -> tuple[torch.Tensor, torch.Tensor] | None: + """Compatibility adapter for the retained V1 DDTree proposer. + + MRV2 DFlash does not use this method. Keeping the adapter on the + official model avoids retaining the removed experimental flat-DFlash + implementation solely for DDTree payload construction. + """ if self.draft_id_to_target_id is not None: return None return self.logits_processor.get_topk_tokens_and_logprobs( @@ -850,7 +927,7 @@ def precompute_and_store_context_kv( self, context_states: torch.Tensor, context_positions: torch.Tensor, - context_slot_mapping: torch.Tensor | dict[str, torch.Tensor] | None = None, + context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None, ) -> None: """Precompute projected + RoPE'd K/V and write to cache.""" self.model.precompute_and_store_context_kv( @@ -866,9 +943,17 @@ def combine_hidden_states( needs_squeeze = hidden_states.dim() == 1 if needs_squeeze: hidden_states = hidden_states.unsqueeze(0) - # Keep the auxiliary hidden-state projection on the same dtype as the - # DFlash FC weights. AWQ target paths may surface fp32 hidden states - # even when the draft projection is initialized in fp16. + expected = self.model.fc.input_size + if hidden_states.shape[-1] != expected: + raise ValueError( + f"DFlash drafter expects {expected} concatenated aux hidden " + f"features but received {hidden_states.shape[-1]}. This usually " + "means the draft model's target_layer_ids reference layers that " + "do not exist in the target model (incompatible draft/target pair)." + ) + # AWQ target paths can surface fp32 auxiliary states even when the + # draft projection is initialized in fp16. Match the FC weight dtype + # before dispatching the linear kernel. fc_weight = getattr(self.model.fc, "weight", None) if fc_weight is not None and hidden_states.dtype != fc_weight.dtype: hidden_states = hidden_states.to(dtype=fc_weight.dtype) @@ -883,7 +968,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): includes_embed_tokens = False for name, loaded_weight in weights: assert "mask_hidden" not in name, ( - "DFlash should use mask_token_id to embed the padding hidden state" + "DFlash embeds masked slots via mask_token_id (optionally " + "overridden by a mask_embedding.pt file); it should not ship a " + "mask_hidden weight." ) if "t2d" in name: continue @@ -897,6 +984,13 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): model_weights[name] = loaded_weight process_eagle_weight(self, name) + # Route the separately-trained mask embedding (if shipped) through the + # standard weight loader alongside the rest of the draft weights. + mask_embedding = self._read_mask_embedding() + if mask_embedding is not None: + model_weights["model.mask_embedding"] = mask_embedding + self.model.has_separate_mask_embedding = True + skip_substrs = [] if not includes_draft_id_mapping: skip_substrs.append("draft_id_to_target_id") @@ -904,6 +998,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): skip_substrs.append("embed_tokens") if not self.model.use_aux_hidden_state: skip_substrs.append("fc.") + if not self.model.has_separate_mask_embedding: + skip_substrs.append("mask_embedding") loader = AutoWeightsLoader( self, skip_prefixes=None, @@ -911,3 +1007,42 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): ) loader.load_weights(model_weights.items()) self.model._build_fused_kv_buffers() + + def _read_mask_embedding(self) -> torch.Tensor | None: + """Checks for an override mask embedding in `mask_embedding.pt` and returns it. + + Some checkpoints ship a separately-trained mask embedding for the mask token, + which we use to overwrite the embedding for `mask_token_id`. This helper + checks for the file, loads the pytorch tensor, and returns the embedding to use. + + Returns None if the override file is not present. + """ + mask_token_id = self.model.mask_token_id + if mask_token_id is None: + return None + + MASK_EMBEDDING_FILENAME = "mask_embedding.pt" + data = get_hf_file_bytes( + MASK_EMBEDDING_FILENAME, + self.draft_model_config.model, + self.draft_model_config.revision, + ) + if data is None: + return None + + state = torch.load(io.BytesIO(data), weights_only=True) + if isinstance(state, dict): + if state.get("mask_token_id", mask_token_id) != mask_token_id: + raise ValueError( + f"{MASK_EMBEDDING_FILENAME} mask_token_id does not match " + f"dflash_config.mask_token_id ({mask_token_id}). " + f"Got {state.get('mask_token_id')}." + ) + state = state["embedding"] + + logger.info( + "Loaded DFlash mask embedding for mask_token_id %s from %s", + mask_token_id, + MASK_EMBEDDING_FILENAME, + ) + return state.reshape(-1) diff --git a/vllm/model_executor/models/qwen3_dflash2.py b/vllm/model_executor/models/qwen3_dflash2.py new file mode 100644 index 0000000000..af6560385c --- /dev/null +++ b/vllm/model_executor/models/qwen3_dflash2.py @@ -0,0 +1,484 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Callable +from functools import cache + +import torch +import torch.nn.functional as F +from torch import nn + +import vllm.envs as envs +from vllm.compilation.backends import set_model_tag +from vllm.compilation.decorators import support_torch_compile +from vllm.config import CacheConfig, VllmConfig +from vllm.distributed import ( + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_gather, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + ReplicatedLinear, + UnquantizedLinearMethod, +) +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +from vllm.model_executor.layers.vocab_parallel_embedding import ( + UnquantizedEmbeddingMethod, +) +from vllm.platforms import current_platform +from vllm.utils.flashinfer import has_flashinfer + +from .dflash_sm70 import ( + DFLASH_SM70_WIDE_OUTPUT_SCALE, + DFlashSM70MLP, + DFlashSM70RMSNorm, +) +from .qwen3_dflash import ( + DFlashQwen3DecoderLayer, + DFlashQwen3ForCausalLM, + DFlashQwen3Model, +) +from .utils import maybe_prefix + +logger = init_logger(__name__) + + +def _use_sm70_bf16_emulation(config) -> bool: + declared_dtype = getattr(config, "dtype", None) + if declared_dtype is None: + declared_dtype = getattr(config, "torch_dtype", None) + is_bf16 = declared_dtype is torch.bfloat16 or str(declared_dtype).lower() in { + "bfloat16", + "bf16", + "torch.bfloat16", + } + return ( + is_bf16 + and current_platform.is_cuda() + and current_platform.is_device_capability(70) + ) + + +@cache +def _flashinfer_topk() -> Callable[..., tuple[torch.Tensor, torch.Tensor]] | None: + """Return FlashInfer radix top-k only where the installed kernel can run. + + Presence is not a sufficient capability check: FlashInfer can be installed + in the V100 environment while its radix top-k has no SM70 implementation. + """ + if not current_platform.is_cuda(): + return None + if not current_platform.has_device_capability(80): + logger.info_once( + "DFlash2 disables FlashInfer top-k below SM80; using torch.topk." + ) + return None + if not has_flashinfer(): + logger.info_once( + "FlashInfer is unavailable; the DFlash2 selector uses torch.topk." + ) + return None + from flashinfer import top_k + + return top_k + + +def _topk(scores: torch.Tensor, k: int) -> tuple[torch.Tensor, torch.Tensor]: + impl = _flashinfer_topk() + if impl is None or not scores.is_cuda: + return torch.topk(scores, k, dim=-1) + return impl(scores, k, sorted=True, deterministic=True) + + +def _grouped_conv( + hidden_states: torch.Tensor, + delta: torch.Tensor, + base: torch.Tensor, + block_size: int, + num_groups: int, + group_size: int, + taps: int, +) -> torch.Tensor: + blocks = hidden_states.unflatten(-1, (num_groups, group_size)) + coefficients = base.view(1, taps, num_groups, group_size) + delta.unsqueeze(-1) + output = coefficients[:, 0] * blocks + position = torch.arange(hidden_states.shape[0], device=hidden_states.device) + if block_size & (block_size - 1) == 0: + position = position & (block_size - 1) + else: + position = position % block_size + for tap in range(1, taps): + shifted = F.pad(blocks[:-tap], (0, 0, 0, 0, tap, 0)) + output += coefficients[:, tap] * shifted * (position >= tap).view(-1, 1, 1) + return output.flatten(-2) + + +class DFlashGroupedConv(nn.Module): + def __init__( + self, + hidden_size: int, + taps: int, + group_size: int, + block_size: int, + params_dtype: torch.dtype, + prefix: str, + ) -> None: + super().__init__() + if hidden_size % group_size: + raise ValueError( + f"conv_group_size={group_size} must divide hidden_size={hidden_size}." + ) + self.block_size = block_size + self.taps = taps + self.group_size = group_size + self.num_groups = hidden_size // group_size + self.base_kernel = nn.Parameter( + torch.empty(2, taps, hidden_size, dtype=params_dtype), + requires_grad=False, + ) + self.kernel_projection = ReplicatedLinear( + hidden_size, + 2 * taps * self.num_groups, + bias=False, + params_dtype=params_dtype, + quant_config=None, + prefix=maybe_prefix(prefix, "kernel_projection"), + return_bias=False, + ) + + def _convolve( + self, hidden_states: torch.Tensor, delta: torch.Tensor, side: int + ) -> torch.Tensor: + return _grouped_conv( + hidden_states, + delta, + self.base_kernel[side], + self.block_size, + self.num_groups, + self.group_size, + self.taps, + ) + + def prepare(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + coefficients = self.kernel_projection(hidden_states).reshape( + hidden_states.shape[0], 2, self.taps, self.num_groups + ) + return self._convolve(hidden_states, coefficients[:, 0], 0), coefficients[:, 1] + + def finish( + self, hidden_states: torch.Tensor, coefficients: torch.Tensor + ) -> torch.Tensor: + return self._convolve(hidden_states, coefficients, 1) + + +class DFlash2Qwen3DecoderLayer(DFlashQwen3DecoderLayer): + def __init__( + self, + vllm_config: VllmConfig, + *, + config, + layer_idx: int, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__( + vllm_config, + config=config, + layer_idx=layer_idx, + cache_config=cache_config, + quant_config=quant_config, + prefix=prefix, + ) + self.use_sm70_bf16_emulation = _use_sm70_bf16_emulation(config) + if self.use_sm70_bf16_emulation: + runtime_dtype = vllm_config.model_config.dtype + if runtime_dtype != torch.float16: + raise ValueError( + "BF16 DFlash2 on SM70 requires FP16 runtime transport; " + f"got runtime dtype {runtime_dtype}." + ) + self.self_attn.output_input_scale = DFLASH_SM70_WIDE_OUTPUT_SCALE + self.mlp = DFlashSM70MLP(self.mlp) + self.input_layernorm = DFlashSM70RMSNorm( + config.hidden_size, + config.rms_norm_eps, + runtime_dtype, + ) + self.post_attention_layernorm = DFlashSM70RMSNorm( + config.hidden_size, + config.rms_norm_eps, + runtime_dtype, + ) + draft_config = config.dflash_config + speculative_config = vllm_config.speculative_config + assert speculative_config is not None + conv_args = dict( + 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, + params_dtype=vllm_config.model_config.dtype, + ) + self.attention_conv = DFlashGroupedConv( + **conv_args, prefix=maybe_prefix(prefix, "attention_conv") + ) + self.mlp_conv = DFlashGroupedConv( + **conv_args, prefix=maybe_prefix(prefix, "mlp_conv") + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + + hidden_states, coefficients = self.attention_conv.prepare(hidden_states) + hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states) + hidden_states = self.attention_conv.finish(hidden_states, coefficients) + + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states, coefficients = self.mlp_conv.prepare(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = self.mlp_conv.finish(hidden_states, coefficients) + return hidden_states, residual + + +def _score_edges( + predecessor_table: torch.Tensor, + successor_table: torch.Tensor, + candidate_ids: torch.Tensor, + unary_logits: torch.Tensor, + hidden: torch.Tensor, + anchor_token_ids: torch.Tensor, + top_k: int, +) -> torch.Tensor: + successors = successor_table[candidate_ids] + predecessor_ids = torch.cat( + ( + anchor_token_ids[:, None, None].expand(-1, 1, top_k), + candidate_ids[:, :-1], + ), + dim=1, + ) + predecessors = predecessor_table[predecessor_ids] + return unary_logits[:, :, None] + torch.einsum( + "blpr,blcr->blpc", predecessors * hidden[:, :, None], successors + ) + + +@support_torch_compile +class CandidateSelector(nn.Module): + def __init__( + self, + hidden_size: int, + vocab_size: int, + rank: int, + top_k: int, + params_dtype: torch.dtype, + prefix: str, + ) -> None: + super().__init__() + self.top_k = top_k + self.predecessor_codebook = nn.Parameter( + torch.empty(vocab_size, rank, dtype=params_dtype), requires_grad=False + ) + self.successor_codebook = nn.Parameter( + torch.empty(vocab_size, rank, dtype=params_dtype), requires_grad=False + ) + self.hidden_projection = ReplicatedLinear( + hidden_size, + rank, + bias=False, + params_dtype=params_dtype, + quant_config=None, + prefix=maybe_prefix(prefix, "hidden_projection"), + return_bias=False, + ) + + def forward( + self, + candidate_ids: torch.Tensor, + unary_logits: torch.Tensor, + hidden_states: torch.Tensor, + anchor_token_ids: torch.Tensor, + ) -> torch.Tensor: + hidden = self.hidden_projection(hidden_states) + return _score_edges( + self.predecessor_codebook, + self.successor_codebook, + candidate_ids, + unary_logits, + hidden, + anchor_token_ids, + self.top_k, + ) + + +class DFlash2Qwen3Model(DFlashQwen3Model): + decoder_layer_cls = DFlash2Qwen3DecoderLayer + + def _make_context_projection( + self, + *, + vllm_config: VllmConfig, + input_size: int, + output_size: int, + prefix: str, + ) -> nn.Module: + use_sharded_fc = ( + envs.VLLM_SM70_DFLASH2_SHARDED_CONTEXT_FC + and self.quant_config is None + and current_platform.is_cuda() + and current_platform.is_device_capability(70) + and vllm_config.parallel_config.tensor_parallel_size == 4 + and input_size == 25600 + and output_size == 5120 + ) + if not use_sharded_fc: + return super()._make_context_projection( + vllm_config=vllm_config, + input_size=input_size, + output_size=output_size, + prefix=prefix, + ) + + projection = ColumnParallelLinear( + input_size=input_size, + output_size=output_size, + bias=False, + gather_output=True, + params_dtype=vllm_config.model_config.dtype, + quant_config=None, + prefix=prefix, + return_bias=False, + ) + projection._sm70_f16_force_enable = True + # TurboMind wins at the 8-row decode shape, while cuBLAS is faster for + # 128+ context rows. Keep the exact small-M path and let ordinary + # F.linear handle chunked-prefill projection shapes. + projection._sm70_f16_max_m = 64 + logger.info_once( + "Using TP4 output-sharded DFlash2 target-hidden projection on SM70 " + "(25600->5120 plus compact all-gather)." + ) + return projection + + def __init__( + self, + *, + vllm_config: VllmConfig, + start_layer_id: int = 0, + prefix: str = "", + ) -> None: + super().__init__( + vllm_config=vllm_config, + start_layer_id=start_layer_id, + prefix=prefix, + ) + draft_config = self.config.dflash_config + self.use_sm70_bf16_emulation = _use_sm70_bf16_emulation(self.config) + if self.use_sm70_bf16_emulation: + runtime_dtype = vllm_config.model_config.dtype + self.hidden_norm = DFlashSM70RMSNorm( + self.config.hidden_size, + self.config.rms_norm_eps, + runtime_dtype, + residual_scale=1.0, + ) + self.norm = DFlashSM70RMSNorm( + self.config.hidden_size, + self.config.rms_norm_eps, + runtime_dtype, + ) + logger.info_once( + "Using range-preserving BF16 DFlash2 arithmetic on SM70 " + "(FP32 residual, BF16 RNE, scaled FP16 projections)." + ) + self.input_embedding_scale = float( + draft_config.get("input_embedding_scale", 1.0) + ) + # The selector is compiled separately from the draft backbone. A unique + # model tag prevents the two incompatible input signatures from sharing + # a persistent compile-cache namespace. + with set_model_tag("dflash2_candidate_selector"): + self.candidate_selector = CandidateSelector( + hidden_size=self.config.hidden_size, + vocab_size=self.config.vocab_size, + rank=int(draft_config["selector_rank"]), + top_k=int(draft_config["selector_top_k"]), + params_dtype=vllm_config.model_config.dtype, + prefix=maybe_prefix(prefix, "candidate_selector"), + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return super().embed_input_ids(input_ids) * self.input_embedding_scale + + def _normalize_context_states(self, context_states: torch.Tensor) -> torch.Tensor: + if self.use_sm70_bf16_emulation: + return self.hidden_norm(context_states) + return super()._normalize_context_states(context_states) + + +class DFlash2Qwen3ForCausalLM(DFlashQwen3ForCausalLM): + model_cls = DFlash2Qwen3Model + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__(vllm_config=vllm_config, prefix=prefix) + draft_config = self.config.dflash_config + self.output_multiplier = float(draft_config.get("output_multiplier", 1.0)) + softcap = float(draft_config.get("final_logit_softcapping") or 0.0) + self.final_logit_softcapping = softcap if softcap > 0 else None + + def compute_candidates( + self, hidden_states: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + if not isinstance( + self.lm_head.quant_method, + (UnquantizedEmbeddingMethod, UnquantizedLinearMethod), + ): + raise ValueError( + "DFlash2 requires an unquantized target LM head for candidate TopK; " + f"got {type(self.lm_head.quant_method).__name__}." + ) + + selector = self.model.candidate_selector + local_candidates = self.lm_head.maybe_get_sm70_dflash2_top20( + hidden_states, selector.top_k + ) + if local_candidates is None: + logits = self.lm_head.quant_method.apply( + self.lm_head, hidden_states, bias=None + ) + num_pad = self.lm_head.shard_indices.num_org_vocab_padding + if num_pad > 0: + logits[..., -num_pad:] = -float("inf") + values, ids = _topk(logits, selector.top_k) + ids = ids.to(torch.int64) + self.lm_head.shard_indices.org_vocab_start_index + else: + values, ids = local_candidates + + if get_tensor_model_parallel_world_size() > 1: + values = tensor_model_parallel_all_gather(values, dim=-1) + ids = tensor_model_parallel_all_gather(ids, dim=-1) + + if values.shape[-1] > selector.top_k: + values, selected = _topk(values, selector.top_k) + ids = ids.gather(-1, selected) + + values = values.float() * self.output_multiplier + if self.final_logit_softcapping is not None: + cap = self.final_logit_softcapping + values = torch.tanh(values / cap) * cap + return ids, values + + +EntryClass = DFlash2Qwen3ForCausalLM diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index cf4034724a..b54454adad 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/vllm/model_executor/models/qwen3_next.py @@ -549,7 +549,7 @@ def __init__( def forward( self, positions: torch.Tensor, - output: torch.Tensor, + output: torch.Tensor | None, hidden_states: torch.Tensor, ): qkv, _ = self.qkv_proj(hidden_states) @@ -623,13 +623,16 @@ def forward( attn_output, ) - output[:], _ = self.o_proj(attn_output) + projected_output, _ = self.o_proj(attn_output) + if output is not None: + output[:] = projected_output _sm70_dump_qwen_layer_tensor( "full_attn_o_proj_out", self.layer_idx, "full_attention", - output, + projected_output, ) + return projected_output class Qwen3NextDecoderLayer(nn.Module): @@ -744,21 +747,33 @@ def forward( hidden_states, ) - self_attention_output = torch.empty_like(hidden_states) + use_direct_attention_output = ( + envs.VLLM_SM70_TP4_LONG_PREFILL_FUSED_NORM and torch.compiler.is_compiling() + ) + self_attention_output = ( + None if use_direct_attention_output else torch.empty_like(hidden_states) + ) if self.layer_type == "linear_attention": - self.linear_attn( + projected_attention_output = self.linear_attn( hidden_states=hidden_states, output=self_attention_output, ) elif self.layer_type == "full_attention": - self.self_attn( + projected_attention_output = self.self_attn( hidden_states=hidden_states, output=self_attention_output, positions=positions, ) else: raise ValueError("Invalid layer_type") - hidden_states = self_attention_output + if use_direct_attention_output: + if projected_attention_output is None: + raise RuntimeError( + "SM70 TP4 fused prefill requires a direct attention output" + ) + hidden_states = projected_attention_output + else: + hidden_states = self_attention_output hidden_states = _sm70_dump_qwen_layer_tensor( "attn_out", self.layer_idx, diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 3a7a7b4ba0..5723dbe285 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -598,6 +598,7 @@ "EagleLlama4ForCausalLM": ("llama4_eagle", "EagleLlama4ForCausalLM"), "EagleMiniCPMForCausalLM": ("minicpm_eagle", "EagleMiniCPMForCausalLM"), "DFlashDraftModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"), + "DFlash2DraftModel": ("qwen3_dflash2", "DFlash2Qwen3ForCausalLM"), "DSparkDraftModel": ( "vllm.models.deepseek_v4", "DSparkDeepseekV4ForCausalLM", diff --git a/vllm/model_executor/offloader/base.py b/vllm/model_executor/offloader/base.py index ceff60cd4c..b29b57d34f 100644 --- a/vllm/model_executor/offloader/base.py +++ b/vllm/model_executor/offloader/base.py @@ -76,6 +76,18 @@ def post_init(self): """ return + def wrap_module( + self, module: nn.Module, *, parameter_prefix: str = "" + ) -> nn.Module: + """Optionally offload one non-layer component. + + Layer-prefetch backends intentionally inherit this no-op. UVA + overrides it so model components created outside ``make_layers`` can + participate without invoking a backend's multi-layer wrapper twice. + """ + del parameter_prefix + return module + def sync_prev_onload(self) -> None: # noqa: B027 """Sync previous onload operations. Override in subclasses.""" pass diff --git a/vllm/model_executor/offloader/uva.py b/vllm/model_executor/offloader/uva.py index 51eb1a14fc..67218d5dd6 100644 --- a/vllm/model_executor/offloader/uva.py +++ b/vllm/model_executor/offloader/uva.py @@ -61,7 +61,30 @@ def wrap_modules( ) return modules - def _maybe_offload_to_cpu(self, module: nn.Module) -> nn.Module: + def wrap_module( + self, module: nn.Module, *, parameter_prefix: str = "" + ) -> nn.Module: + """Apply UVA offloading to one component with its qualified prefix.""" + previous_offload_bytes = self.cpu_offload_bytes + module = self._maybe_offload_to_cpu(module, parameter_prefix=parameter_prefix) + if self.cpu_offload_bytes > previous_offload_bytes: + logger.info( + "Total CPU offloaded parameters: %s", + format_gib(self.cpu_offload_bytes), + ) + return module + + def _is_parameter_selected(self, name: str, parameter_prefix: str = "") -> bool: + if not self.cpu_offload_params: + return True + qualified_name = f"{parameter_prefix}.{name}" if parameter_prefix else name + return any( + f".{param}." in f".{qualified_name}." for param in self.cpu_offload_params + ) + + def _maybe_offload_to_cpu( + self, module: nn.Module, *, parameter_prefix: str = "" + ) -> nn.Module: """Offload module parameters to CPU using UVA if budget allows.""" if (params := next(module.parameters(), None)) is None: return module @@ -83,16 +106,17 @@ def _maybe_offload_to_cpu(self, module: nn.Module) -> nn.Module: # one module might have some parameters offloaded and some not break - if self.cpu_offload_params: - # Check if parameter belongs to the offloading set - # Add dots here to ensure we match full segments only - # e.g., "experts.w2_weight" matches "mlp.experts.w2_weight" - # but not "mlp.experts.w2_weight_scale" - should_offload = any( - f".{param}." in f".{name}." for param in self.cpu_offload_params - ) - if not should_offload: - continue + # Add dots to require exact path-segment matches. A qualified + # prefix lets selectors such as ``visual`` target an entire tower + # even though ``module.named_parameters()`` returns relative names. + if not self._is_parameter_selected(name, parameter_prefix): + continue + + # A component may contain layers already visited by make_layers(). + # Do not copy or account for the same UVA parameter twice when the + # enclosing component is registered afterwards. + if getattr(p, "_vllm_is_uva_offloaded", False): + continue cpu_data = p.data.to(device="cpu") if self.pin_memory: diff --git a/vllm/model_executor/warmup/awq_sm70_warmup.py b/vllm/model_executor/warmup/awq_sm70_warmup.py index e12589a2c7..9bc11324d9 100644 --- a/vllm/model_executor/warmup/awq_sm70_warmup.py +++ b/vllm/model_executor/warmup/awq_sm70_warmup.py @@ -15,6 +15,9 @@ from vllm import _sm70_ops as sm70_ops from vllm.logger import init_logger from vllm.model_executor.layers.quantization import sm70_turbomind as sm70_tm +from vllm.model_executor.layers.quantization.nvfp4_sm70_moe import ( + _prepare_compact_slot_groups, +) if TYPE_CHECKING: from vllm.v1.worker.gpu_worker import Worker @@ -44,6 +47,7 @@ def _lut_cache_disabled_for_dynamic_quant_dispatch( has_fp8_dense: bool, fp4_kinds: set[str], has_mxfp4_moe: bool = False, + has_nvfp4_moe: bool = False, ) -> bool: awq_no_preserve = ( has_awq_dense @@ -59,7 +63,9 @@ def _lut_cache_disabled_for_dynamic_quant_dispatch( mxfp4_dynamic = ( "mxfp4" in fp4_kinds or has_mxfp4_moe ) and envs.VLLM_SM70_MXFP4_TUNE_SMALL_SHAPES - nvfp4_dynamic = "nvfp4" in fp4_kinds and envs.VLLM_SM70_NVFP4_TUNE_SMALL_SHAPES + nvfp4_dynamic = ( + "nvfp4" in fp4_kinds or has_nvfp4_moe + ) and envs.VLLM_SM70_NVFP4_TUNE_SMALL_SHAPES return awq_no_preserve or fp8_dynamic or mxfp4_dynamic or nvfp4_dynamic @@ -259,7 +265,15 @@ def _iter_unique_fp8_dense_layers( ) -> Iterable[tuple[torch.nn.Module, bool]]: seen: set[tuple[int, int, bool]] = set() for layer in model.modules(): - if not getattr(layer, "sm70_fp8_turbomind", False): + if not ( + getattr(layer, "sm70_fp8_turbomind", False) + or getattr(layer, "sm70_modelopt_fp8_turbomind", False) + ): + continue + # QPN8 has a static source-selected dispatch and does not populate the + # TurboMind LUT warmed below. CUDA graph capture exercises its admitted + # M<=8 kernels separately. + if getattr(layer, "sm70_fp8_qpn8", False): continue if getattr(layer, "sm70_fp8_bmm", False): @@ -293,14 +307,20 @@ def _iter_unique_fp8_dense_layers( def _iter_unique_fp4_dense_layers(model: torch.nn.Module) -> Iterable[Any]: - seen: set[tuple[str, int, int, int]] = set() + seen: set[tuple[str, int, int, int, bool]] = set() for layer in model.modules(): state = getattr(layer, sm70_tm.STATE_ATTR, None) if state is None or state.op_kind not in ("mxfp4", "nvfp4"): continue k_dim = int(state.weight.shape[0]) n_dim = int(state.output_size) - key = (state.op_kind, k_dim, n_dim, int(state.group_size)) + key = ( + state.op_kind, + k_dim, + n_dim, + int(state.group_size), + bool(state.gated_silu), + ) if key in seen: continue seen.add(key) @@ -387,6 +407,45 @@ def _iter_unique_mxfp4_moe_layers( yield layer +def _is_nvfp4_moe_sm70_layer(layer: torch.nn.Module) -> bool: + return getattr(layer, "sm70_nvfp4_moe", False) and all( + hasattr(layer, name) + for name in ( + "w13_strided_ptrs_w", + "w13_strided_ptrs_s", + "w2_strided_ptrs_w", + "w2_strided_ptrs_s", + "sm70_nvfp4_num_experts", + "sm70_nvfp4_w13_k_dim", + "sm70_nvfp4_w13_n_dim", + "sm70_nvfp4_w2_k_dim", + "sm70_nvfp4_w2_n_dim", + "sm70_nvfp4_group_size", + "_nvfp4_sm70_dense_expert_ids", + ) + ) + + +def _iter_unique_nvfp4_moe_layers( + model: torch.nn.Module, +) -> Iterable[torch.nn.Module]: + seen: set[tuple[int, int, int, int, int]] = set() + for layer in model.modules(): + if not _is_nvfp4_moe_sm70_layer(layer): + continue + key = ( + int(layer.sm70_nvfp4_w13_k_dim), + int(layer.sm70_nvfp4_w13_n_dim), + int(layer.sm70_nvfp4_w2_k_dim), + int(layer.sm70_nvfp4_w2_n_dim), + int(layer.sm70_nvfp4_num_experts), + ) + if key in seen: + continue + seen.add(key) + yield layer + + def _warmup_dense_layers( dense_layers: list[torch.nn.Module], m_values: list[int], @@ -423,6 +482,7 @@ def _warmup_fp8_dense_layers( calls = 0 for layer, gated_silu in dense_layers: + is_modelopt = getattr(layer, "sm70_modelopt_fp8_turbomind", False) if getattr(layer, "sm70_fp8_bmm", False): assert not gated_silu weight = layer.weight[0] @@ -444,9 +504,14 @@ def _warmup_fp8_dense_layers( n_dim = int(weight.shape[1]) // 2 else: weight = layer.weight - scales = layer.weight_scale_inv - k_ld = int(layer.sm70_fp8_k_ld) - q_ld = int(layer.sm70_fp8_q_ld) + if is_modelopt: + scales = layer.weight_scale + k_ld = int(layer.sm70_modelopt_fp8_k_ld) + q_ld = int(layer.sm70_modelopt_fp8_q_ld) + else: + scales = layer.weight_scale_inv + k_ld = int(layer.sm70_fp8_k_ld) + q_ld = int(layer.sm70_fp8_q_ld) n_dim = int(layer.output_size_per_partition) device = weight.device @@ -458,6 +523,27 @@ def _warmup_fp8_dense_layers( out, x, weight, scales, 128, k_ld, q_ld, gated_silu ) calls += 1 + if getattr(layer, "sm70_fp8_bmm_grouped_decode", False): + group_count = int(layer.sm70_fp8_bmm_groups) + grouped_x = torch.empty( + (group_count, k_dim), dtype=torch.float16, device=device + ) + grouped_out = torch.empty( + (group_count, n_dim), dtype=torch.float16, device=device + ) + sm70_ops.fp8_moe_gemm_sm70_per_expert_dispatch_out( + grouped_out, + grouped_x, + layer.sm70_fp8_bmm_grouped_offsets, + layer.sm70_fp8_bmm_grouped_ptrs_w, + layer.sm70_fp8_bmm_grouped_ptrs_s, + group_count, + k_dim, + n_dim, + 128, + False, + ) + calls += 1 return calls @@ -469,7 +555,8 @@ def _warmup_fp4_dense_layers( for state in dense_layers: device = state.weight.device k_dim = int(state.weight.shape[0]) - n_dim = int(state.output_size) + gated_silu = bool(state.gated_silu) + n_dim = int(state.output_size) // (2 if gated_silu else 1) for m_dim in m_values: x = torch.empty((m_dim, k_dim), dtype=torch.float16, device=device) out = torch.empty((m_dim, n_dim), dtype=torch.float16, device=device) @@ -484,7 +571,7 @@ def _warmup_fp4_dense_layers( int(state.group_size), int(state.k_ld), int(state.q_ld), - False, + gated_silu, ) elif state.op_kind == "nvfp4": if not hasattr(torch.ops._C, "nvfp4_gemm_sm70_out"): @@ -497,7 +584,7 @@ def _warmup_fp4_dense_layers( int(state.group_size), int(state.k_ld), int(state.q_ld), - False, + gated_silu, ) else: continue @@ -644,6 +731,133 @@ def _warmup_mxfp4_moe_b1_layers(moe_layers: list[torch.nn.Module]) -> int: return calls +def _warmup_nvfp4_moe_decode_layers( + moe_layers: list[torch.nn.Module], token_counts: list[int] +) -> int: + """Tune compact and full-group NVFP4 graph shapes.""" + if not hasattr(torch.ops._C, "nvfp4_moe_dense_stage_sm70_out"): + return 0 + + calls = 0 + for layer in moe_layers: + device = layer.w13_tm_weight.device + top_k = int(layer.moe_config.experts_per_token) + max_experts = int(layer.sm70_nvfp4_num_experts) + compact_max_tokens = int( + getattr(layer, "sm70_nvfp4_compact_grouped_max_tokens", 8) + ) + for num_tokens in token_counts: + total_slots = num_tokens * top_k + if num_tokens <= compact_max_tokens: + stage_experts = total_slots + sorted_expert_ids = layer._nvfp4_sm70_dense_expert_ids[:total_slots] + expert_offsets = torch.empty( + total_slots + 1, dtype=torch.int32, device=device + ) + active_expert_ids = torch.empty( + total_slots, dtype=torch.int32, device=device + ) + _prepare_compact_slot_groups( + sorted_expert_ids, + expert_offsets, + active_expert_ids, + ) + else: + stage_experts = max_experts + expert_offsets = _build_balanced_offsets( + total_slots, max_experts, device + ) + active_expert_ids = layer._nvfp4_sm70_dense_expert_ids + permuted_input = torch.empty( + total_slots, + int(layer.sm70_nvfp4_w13_k_dim), + dtype=torch.float16, + device=device, + ) + gate_up = torch.empty( + total_slots, + int(layer.sm70_nvfp4_w13_n_dim), + dtype=torch.float16, + device=device, + ) + intermediate = torch.empty( + total_slots, + int(layer.sm70_nvfp4_w2_k_dim), + dtype=torch.float16, + device=device, + ) + sorted_output = torch.empty( + total_slots, + int(layer.sm70_nvfp4_w2_n_dim), + dtype=torch.float16, + device=device, + ) + sm70_ops.nvfp4_moe_dense_stage_sm70_out( + gate_up, + permuted_input, + expert_offsets, + active_expert_ids, + layer.w13_strided_ptrs_w, + layer.w13_strided_ptrs_s, + stage_experts, + int(layer.sm70_nvfp4_w13_k_dim), + int(layer.sm70_nvfp4_w13_n_dim), + int(layer.sm70_nvfp4_group_size), + ) + if layer.swiglu_limit is None: + torch.ops._C.silu_and_mul(intermediate, gate_up) + else: + torch.ops._C.silu_and_mul_with_clamp( + intermediate, gate_up, float(layer.swiglu_limit) + ) + sm70_ops.nvfp4_moe_dense_stage_sm70_out( + sorted_output, + intermediate, + expert_offsets, + active_expert_ids, + layer.w2_strided_ptrs_w, + layer.w2_strided_ptrs_s, + stage_experts, + int(layer.sm70_nvfp4_w2_k_dim), + int(layer.sm70_nvfp4_w2_n_dim), + int(layer.sm70_nvfp4_group_size), + ) + calls += 2 + return calls + + +def _get_nvfp4_moe_token_counts( + worker: Worker, + moe_layers: list[torch.nn.Module], + base_token_counts: list[int], +) -> list[int]: + """Include graph-safe and tuned NVFP4 CUDA-graph token widths.""" + if not moe_layers: + return base_token_counts + graph_safe_max_tokens = max( + int(getattr(layer, "sm70_nvfp4_graph_safe_max_tokens", 8)) + for layer in moe_layers + ) + compact_max_tokens = max( + int(getattr(layer, "sm70_nvfp4_compact_grouped_max_tokens", 8)) + for layer in moe_layers + ) + max_top_k = max(int(layer.moe_config.experts_per_token) for layer in moe_layers) + tuned_max_tokens = max(0, int(envs.VLLM_SM70_NVFP4_MOE_TUNE_MAX_TOKENS)) // max( + max_top_k, 1 + ) + max_tokens = max(graph_safe_max_tokens, tuned_max_tokens) + compilation_config = getattr( + getattr(worker, "vllm_config", None), "compilation_config", None + ) + capture_sizes = getattr(compilation_config, "cudagraph_capture_sizes", None) or [] + return sorted( + set(base_token_counts) + | set(range(1, compact_max_tokens + 1)) + | {int(size) for size in capture_sizes if 0 < int(size) <= max_tokens} + ) + + def _warmup_moe_single_token_layers(moe_layers: list[torch.nn.Module]) -> int: if not ( hasattr(torch.ops._C, "awq_moe_single_token_dense_w13_sm70_out") @@ -797,12 +1011,14 @@ def sm70_awq_warmup(worker: Worker) -> None: fp4_dense_layers = list(_iter_unique_fp4_dense_layers(model)) moe_layers = list(_iter_unique_moe_layers(model)) mxfp4_moe_layers = list(_iter_unique_mxfp4_moe_layers(model)) + nvfp4_moe_layers = list(_iter_unique_nvfp4_moe_layers(model)) if not ( dense_layers or fp8_dense_layers or fp4_dense_layers or moe_layers or mxfp4_moe_layers + or nvfp4_moe_layers ): return @@ -812,6 +1028,7 @@ def sm70_awq_warmup(worker: Worker) -> None: bool(fp8_dense_layers), fp4_kinds, bool(mxfp4_moe_layers), + bool(nvfp4_moe_layers), ) imported_records = _load_lut_cache(device, skip_import=skip_lut_cache) if imported_records > 0: @@ -823,20 +1040,26 @@ def sm70_awq_warmup(worker: Worker) -> None: m_values = _get_decode_m_values(worker) moe_token_counts = _get_moe_token_counts(worker) + nvfp4_moe_token_counts = _get_nvfp4_moe_token_counts( + worker, nvfp4_moe_layers, moe_token_counts + ) lut_path = _resolve_lut_path(device) logger.info( "Warming up SM70 TurboMind accepted routes (%d AWQ dense layer " "shapes, %d FP8 dense layer shapes, %d FP4 dense layer shapes, " - "%d MoE layer shapes, %d MXFP4 MoE B1 layer shapes, dense_m=%s, " - "moe_tokens=%s, lut_path=%s).", + "%d MoE layer shapes, %d MXFP4 MoE B1 layer shapes, " + "%d NVFP4 MoE layer shapes, dense_m=%s, " + "moe_tokens=%s, nvfp4_moe_tokens=%s, lut_path=%s).", len(dense_layers), len(fp8_dense_layers), len(fp4_dense_layers), len(moe_layers), len(mxfp4_moe_layers), + len(nvfp4_moe_layers), m_values, moe_token_counts, + nvfp4_moe_token_counts, lut_path, ) with torch.inference_mode(): @@ -850,17 +1073,22 @@ def sm70_awq_warmup(worker: Worker) -> None: moe_stage_calls = _warmup_moe_dense_stage_layers(moe_layers, moe_token_counts) single_token_calls = _warmup_moe_single_token_layers(moe_layers) mxfp4_moe_stage_calls = _warmup_mxfp4_moe_b1_layers(mxfp4_moe_layers) + nvfp4_moe_stage_calls = _warmup_nvfp4_moe_decode_layers( + nvfp4_moe_layers, nvfp4_moe_token_counts + ) torch.accelerator.synchronize(device) logger.info( "SM70 TurboMind warmup finished (%d AWQ dense calls, %d FP8 dense " "calls, %d FP4 dense calls, %d MoE stage calls, " - "%d single-token active-expert calls, %d MXFP4 MoE B1 stage calls).", + "%d single-token active-expert calls, %d MXFP4 MoE B1 stage calls, " + "%d NVFP4 MoE stage calls).", dense_calls, fp8_dense_calls, fp4_dense_calls, moe_stage_calls, single_token_calls, mxfp4_moe_stage_calls, + nvfp4_moe_stage_calls, ) exported_records = _save_lut_cache(device, skip_export=skip_lut_cache) if exported_records > 0: diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index 0920e6db4d..bbdd5b1053 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -308,6 +308,7 @@ def __init__( rotate=True, prefix=f"{prefix}.compressor", k_cache_prefix=self.mla_attn.prefix, + validity_cache_prefix=self.swa_cache_layer.prefix, ) def forward( @@ -415,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 @@ -434,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 @@ -446,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 @@ -850,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" @@ -923,7 +940,7 @@ def forward( attn_metadata = get_forward_context().attn_metadata if isinstance(attn_metadata, dict): indexer_metadata = cast(Any, attn_metadata[self.k_cache.prefix]) - if indexer_metadata.max_seq_len // self.compress_ratio <= self.topk_tokens: + if indexer_metadata.compressed_max_seq_len <= self.topk_tokens: compressor(compressed_kv_score, positions, rotary_emb) assert self.topk_indices_buffer is not None num_tokens = ( diff --git a/vllm/models/deepseek_v4/common/ops/__init__.py b/vllm/models/deepseek_v4/common/ops/__init__.py index dc6f3c608d..ae26189162 100644 --- a/vllm/models/deepseek_v4/common/ops/__init__.py +++ b/vllm/models/deepseek_v4/common/ops/__init__.py @@ -11,7 +11,11 @@ from .fused_inv_rope_fp8_quant import fused_inv_rope_fp8_quant from .fused_mtp_input_rmsnorm import fused_mtp_input_rmsnorm, mtp_shared_head_rmsnorm from .fused_qk_rmsnorm import fused_q_kv_rmsnorm -from .save_partial_states import save_partial_states +from .save_partial_states import ( + save_partial_states, + save_partial_states_to_ring, + stage_partial_states_from_ring, +) __all__ = [ "MXFP4_BLOCK_SIZE", @@ -25,4 +29,6 @@ "mtp_shared_head_rmsnorm", "quantize_and_insert_k_cache", "save_partial_states", + "save_partial_states_to_ring", + "stage_partial_states_from_ring", ] diff --git a/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py b/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py index f5e0886972..9cec99251f 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py +++ b/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py @@ -30,8 +30,148 @@ from .fused_indexer_q import _fp32x2_to_fp4x2 +@triton.jit +def _gather_compressor_window( + state_cache_ptr, + state_cache_stride0, + state_cache_stride1, + current_kv_ptr, + current_kv_stride, + current_score_ptr, + current_score_stride, + ape_ptr, + ape_stride, + block_table_ptr, + block_table_stride, + positions_ptr, + position, + token_idx, + req_idx, + block_size, + block, + mask, + HEAD_SIZE: tl.constexpr, + STATE_WIDTH: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + OVERLAP: tl.constexpr, + USE_PRIVATE_STATE: tl.constexpr, + USE_DENSE_PRIVATE_STATE: tl.constexpr, + RING_SIZE: tl.constexpr, + PRIVATE_HISTORY_SIZE: tl.constexpr, +): + """Gather a compression window from paged state or current input + ring.""" + tokens = tl.arange(0, (1 + OVERLAP) * COMPRESS_RATIO) + pos = position - (1 + OVERLAP) * COMPRESS_RATIO + 1 + tokens + mask_pos = pos >= 0 + head_offset = (tokens >= COMPRESS_RATIO).to(tl.int32) * HEAD_SIZE + + if USE_DENSE_PRIVATE_STATE: + current_start = position - token_idx + dense_rows = pos - current_start + PRIVATE_HISTORY_SIZE + row_base = ( + state_cache_ptr + + req_idx.to(tl.int64) * state_cache_stride0 + + dense_rows.to(tl.int64) * state_cache_stride1 + + head_offset + ) + combined_mask = mask_pos[:, None] & mask[None, :] + score = tl.load( + row_base[:, None] + STATE_WIDTH + block[None, :], + mask=combined_mask, + other=float("-inf"), + ) + kv = tl.load( + row_base[:, None] + block[None, :], + mask=combined_mask, + other=0.0, + ) + return kv, score + + if USE_PRIVATE_STATE: + # This route is gated to one request with contiguous positions. + current_start = position - token_idx + current_idx = pos - current_start + from_current = mask_pos & (current_idx >= 0) + from_history = mask_pos & (current_idx < 0) + current_mask = from_current[:, None] & mask[None, :] + history_mask = from_history[:, None] & mask[None, :] + offsets = head_offset[:, None] + block[None, :] + + current_kv = tl.load( + current_kv_ptr + + current_idx[:, None].to(tl.int64) * current_kv_stride + + offsets, + mask=current_mask, + other=0.0, + ).to(tl.float32) + current_score = tl.load( + current_score_ptr + + current_idx[:, None].to(tl.int64) * current_score_stride + + offsets, + mask=current_mask, + other=0.0, + ).to(tl.float32) + ape = tl.load( + ape_ptr + + (pos % COMPRESS_RATIO)[:, None].to(tl.int64) * ape_stride + + offsets, + mask=current_mask, + other=0.0, + ).to(tl.float32) + + safe_pos = tl.maximum(pos, 0) + ring_base = ( + state_cache_ptr + + req_idx.to(tl.int64) * state_cache_stride0 + + (safe_pos % RING_SIZE).to(tl.int64) * state_cache_stride1 + + head_offset + ) + history_kv = tl.load( + ring_base[:, None] + block[None, :], + mask=history_mask, + other=0.0, + ) + history_score = tl.load( + ring_base[:, None] + STATE_WIDTH + block[None, :], + mask=history_mask, + other=float("-inf"), + ) + kv = tl.where(current_mask, current_kv, history_kv) + score = tl.where(current_mask, current_score + ape, history_score) + return kv, score + + block_indices = pos // block_size + block_numbers = tl.load( + block_table_ptr + req_idx * block_table_stride + block_indices, + mask=mask_pos, + other=0, + ) + block_offsets = pos % block_size + row_base = ( + state_cache_ptr + + block_numbers.to(tl.int64) * state_cache_stride0 + + block_offsets * state_cache_stride1 + + head_offset + ) + combined_mask = mask_pos[:, None] & mask[None, :] + score = tl.load( + row_base[:, None] + STATE_WIDTH + block[None, :], + mask=combined_mask, + other=float("-inf"), + ) + kv = tl.load( + row_base[:, None] + block[None, :], + mask=combined_mask, + other=0.0, + ) + return kv, score + + def compress_norm_rope_store_triton( state_cache: torch.Tensor, + current_kv: torch.Tensor, + current_score: torch.Tensor, + ape: torch.Tensor, num_actual: int, token_to_req_indices: torch.Tensor, positions: torch.Tensor, @@ -53,6 +193,10 @@ def compress_norm_rope_store_triton( quant_block: int, token_stride: int, scale_dim: int, + use_private_state: bool = False, + use_dense_private_state: bool = False, + ring_size: int = 1, + private_history_size: int = 1, ) -> None: """Shared triton launcher for the fused compress+norm+RoPE+insert path. @@ -91,6 +235,12 @@ def compress_norm_rope_store_triton( kv_cache, k_cache_metadata.slot_mapping, kv_cache.shape[1], # paged KV cache block size (tokens per block) + current_kv, + current_kv.stride(0), + current_score, + current_score.stride(0), + ape, + ape.stride(0), # constexprs HEAD_SIZE=head_dim, TRITON_BLOCK_SIZE=triton.next_power_of_2(head_dim), @@ -106,6 +256,10 @@ def compress_norm_rope_store_triton( USE_SOFTWARE_FP8=( current_platform.is_cuda() and current_platform.is_device_capability((7, 0)) ), + USE_PRIVATE_STATE=use_private_state, + USE_DENSE_PRIVATE_STATE=use_dense_private_state, + RING_SIZE=ring_size, + PRIVATE_HISTORY_SIZE=private_history_size, num_warps=num_warps, **pdl_kwargs, ) @@ -137,6 +291,13 @@ def _fused_kv_compress_norm_rope_insert_sparse_attn( k_cache_ptr, kv_slot_mapping_ptr, kv_cache_block_size, + # current chunk (private-state path) + current_kv_ptr, + current_kv_stride, + current_score_ptr, + current_score_stride, + ape_ptr, + ape_stride, # ── constexprs ── HEAD_SIZE: tl.constexpr, TRITON_BLOCK_SIZE: tl.constexpr, @@ -150,6 +311,10 @@ def _fused_kv_compress_norm_rope_insert_sparse_attn( SCALE_DIM: tl.constexpr, # 8 for DeepseekV4 (7 real + 1 pad) KV_BLOCK_STRIDE: tl.constexpr, USE_SOFTWARE_FP8: tl.constexpr, + USE_PRIVATE_STATE: tl.constexpr, + USE_DENSE_PRIVATE_STATE: tl.constexpr, + RING_SIZE: tl.constexpr, + PRIVATE_HISTORY_SIZE: tl.constexpr, ): """Fused compress → RMSNorm → FP8 quant (nope) → RoPE → bf16 store (rope). @@ -172,48 +337,40 @@ def _fused_kv_compress_norm_rope_insert_sparse_attn( req_idx = tl.load(token_to_req_indices_ptr + token_idx) # ── Gather state cache entries ──────────────────────────────────── - start = position - (1 + OVERLAP) * COMPRESS_RATIO + 1 - tokens = tl.arange(0, (1 + OVERLAP) * COMPRESS_RATIO) - pos = start + tokens - mask_pos = pos >= 0 - - block_indices = pos // block_size - block_numbers = tl.load( - block_table_ptr + req_idx * block_table_stride + block_indices, - mask=mask_pos, - other=0, - ) - block_offsets = pos % block_size - head_offset = (tokens >= COMPRESS_RATIO).to(tl.int32) * HEAD_SIZE - block = tl.arange(0, TRITON_BLOCK_SIZE) mask = block < HEAD_SIZE - block_numbers_i64 = block_numbers.to(tl.int64) - - # Precomputed row base shared by score and kv loads - row_base = ( - state_cache_ptr - + block_numbers_i64 * state_cache_stride0 - + block_offsets * state_cache_stride1 - + head_offset + kv, score = _gather_compressor_window( + state_cache_ptr, + state_cache_stride0, + state_cache_stride1, + current_kv_ptr, + current_kv_stride, + current_score_ptr, + current_score_stride, + ape_ptr, + ape_stride, + block_table_ptr, + block_table_stride, + positions_ptr, + position, + token_idx, + req_idx, + block_size, + block, + mask, + HEAD_SIZE, + STATE_WIDTH, + COMPRESS_RATIO, + OVERLAP, + USE_PRIVATE_STATE, + USE_DENSE_PRIVATE_STATE, + RING_SIZE, + PRIVATE_HISTORY_SIZE, ) - combined_mask = mask_pos[:, None] & mask[None, :] - # ── Softmax + weighted sum ─────────────────────────────────────── - score = tl.load( - row_base[:, None] + STATE_WIDTH + block[None, :], - mask=combined_mask, - other=float("-inf"), - ) score = tl.softmax(score, dim=0) - kv = tl.load( - row_base[:, None] + block[None, :], - mask=combined_mask, - other=0.0, - ) - compressed_kv = tl.sum(kv * score, axis=0) # [TRITON_BLOCK_SIZE] fp32 # ── RMSNorm (fp32 throughout) ────────────────────────────────────── @@ -331,6 +488,13 @@ def _fused_kv_compress_norm_rope_insert_indexer_attn( k_cache_ptr, kv_slot_mapping_ptr, kv_cache_block_size, + # current chunk (private-state path) + current_kv_ptr, + current_kv_stride, + current_score_ptr, + current_score_stride, + ape_ptr, + ape_stride, # ── constexprs ── HEAD_SIZE: tl.constexpr, TRITON_BLOCK_SIZE: tl.constexpr, @@ -344,6 +508,10 @@ def _fused_kv_compress_norm_rope_insert_indexer_attn( SCALE_DIM: tl.constexpr, # 4 for indexer (1 float32) KV_BLOCK_STRIDE: tl.constexpr, USE_SOFTWARE_FP8: tl.constexpr, + USE_PRIVATE_STATE: tl.constexpr, + USE_DENSE_PRIVATE_STATE: tl.constexpr, + RING_SIZE: tl.constexpr, + PRIVATE_HISTORY_SIZE: tl.constexpr, ): """Fused compress → RMSNorm → RoPE → FP8 quant → store. @@ -370,46 +538,38 @@ def _fused_kv_compress_norm_rope_insert_indexer_attn( req_idx = tl.load(token_to_req_indices_ptr + token_idx) # ── Gather state cache entries ──────────────────────────────────── - start = position - (1 + OVERLAP) * COMPRESS_RATIO + 1 - tokens = tl.arange(0, (1 + OVERLAP) * COMPRESS_RATIO) - pos = start + tokens - mask_pos = pos >= 0 - - block_indices = pos // block_size - block_numbers = tl.load( - block_table_ptr + req_idx * block_table_stride + block_indices, - mask=mask_pos, - other=0, - ) - block_offsets = pos % block_size - head_offset = (tokens >= COMPRESS_RATIO).to(tl.int32) * HEAD_SIZE - block = tl.arange(0, TRITON_BLOCK_SIZE) mask = block < HEAD_SIZE - block_numbers_i64 = block_numbers.to(tl.int64) - - row_base = ( - state_cache_ptr - + block_numbers_i64 * state_cache_stride0 - + block_offsets * state_cache_stride1 - + head_offset - ) - - combined_mask = mask_pos[:, None] & mask[None, :] - - score = tl.load( - row_base[:, None] + STATE_WIDTH + block[None, :], - mask=combined_mask, - other=float("-inf"), + kv, score = _gather_compressor_window( + state_cache_ptr, + state_cache_stride0, + state_cache_stride1, + current_kv_ptr, + current_kv_stride, + current_score_ptr, + current_score_stride, + ape_ptr, + ape_stride, + block_table_ptr, + block_table_stride, + positions_ptr, + position, + token_idx, + req_idx, + block_size, + block, + mask, + HEAD_SIZE, + STATE_WIDTH, + COMPRESS_RATIO, + OVERLAP, + USE_PRIVATE_STATE, + USE_DENSE_PRIVATE_STATE, + RING_SIZE, + PRIVATE_HISTORY_SIZE, ) score = tl.softmax(score, dim=0) - kv = tl.load( - row_base[:, None] + block[None, :], - mask=combined_mask, - other=0.0, - ) - compressed_kv = tl.sum(kv * score, axis=0) # [TRITON_BLOCK_SIZE] fp32 # ── RMSNorm (fp32 throughout) ────────────────────────────────────── @@ -512,6 +672,13 @@ def _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn( k_cache_ptr, kv_slot_mapping_ptr, kv_cache_block_size, + # current chunk (private-state path) + current_kv_ptr, + current_kv_stride, + current_score_ptr, + current_score_stride, + ape_ptr, + ape_stride, # ── constexprs ── HEAD_SIZE: tl.constexpr, TRITON_BLOCK_SIZE: tl.constexpr, @@ -525,6 +692,10 @@ def _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn( SCALE_DIM: tl.constexpr, # HEAD_SIZE // QUANT_BLOCK = 4 ue8m0 bytes/token KV_BLOCK_STRIDE: tl.constexpr, USE_SOFTWARE_FP8: tl.constexpr, + USE_PRIVATE_STATE: tl.constexpr, + USE_DENSE_PRIVATE_STATE: tl.constexpr, + RING_SIZE: tl.constexpr, + PRIVATE_HISTORY_SIZE: tl.constexpr, ): """Fused compress → RMSNorm → RoPE → MXFP4 quant → store. @@ -553,46 +724,38 @@ def _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn( req_idx = tl.load(token_to_req_indices_ptr + token_idx) # ── Gather state cache entries ──────────────────────────────────── - start = position - (1 + OVERLAP) * COMPRESS_RATIO + 1 - tokens = tl.arange(0, (1 + OVERLAP) * COMPRESS_RATIO) - pos = start + tokens - mask_pos = pos >= 0 - - block_indices = pos // block_size - block_numbers = tl.load( - block_table_ptr + req_idx * block_table_stride + block_indices, - mask=mask_pos, - other=0, - ) - block_offsets = pos % block_size - head_offset = (tokens >= COMPRESS_RATIO).to(tl.int32) * HEAD_SIZE - block = tl.arange(0, TRITON_BLOCK_SIZE) mask = block < HEAD_SIZE - block_numbers_i64 = block_numbers.to(tl.int64) - - row_base = ( - state_cache_ptr - + block_numbers_i64 * state_cache_stride0 - + block_offsets * state_cache_stride1 - + head_offset - ) - - combined_mask = mask_pos[:, None] & mask[None, :] - - score = tl.load( - row_base[:, None] + STATE_WIDTH + block[None, :], - mask=combined_mask, - other=float("-inf"), + kv, score = _gather_compressor_window( + state_cache_ptr, + state_cache_stride0, + state_cache_stride1, + current_kv_ptr, + current_kv_stride, + current_score_ptr, + current_score_stride, + ape_ptr, + ape_stride, + block_table_ptr, + block_table_stride, + positions_ptr, + position, + token_idx, + req_idx, + block_size, + block, + mask, + HEAD_SIZE, + STATE_WIDTH, + COMPRESS_RATIO, + OVERLAP, + USE_PRIVATE_STATE, + USE_DENSE_PRIVATE_STATE, + RING_SIZE, + PRIVATE_HISTORY_SIZE, ) score = tl.softmax(score, dim=0) - kv = tl.load( - row_base[:, None] + block[None, :], - mask=combined_mask, - other=0.0, - ) - compressed_kv = tl.sum(kv * score, axis=0) # [TRITON_BLOCK_SIZE] fp32 # ── RMSNorm (fp32 throughout) ────────────────────────────────────── diff --git a/vllm/models/deepseek_v4/common/ops/save_partial_states.py b/vllm/models/deepseek_v4/common/ops/save_partial_states.py index e3d7d38f45..07b40cfcea 100644 --- a/vllm/models/deepseek_v4/common/ops/save_partial_states.py +++ b/vllm/models/deepseek_v4/common/ops/save_partial_states.py @@ -45,6 +45,93 @@ def save_partial_states( ) +def save_partial_states_to_ring( + kv: torch.Tensor, + score: torch.Tensor, + ape: torch.Tensor, + positions: torch.Tensor, + state_ring: torch.Tensor, + slot_mapping: torch.Tensor, + token_to_req_indices: torch.Tensor, + seq_lens: torch.Tensor, + state_width: int, + compress_ratio: int, + pdl_kwargs: dict | None = None, +) -> None: + """Keep only the compressor history that can cross a forward boundary.""" + num_slots = slot_mapping.shape[0] + head_size = kv.shape[-1] + _save_partial_states_to_ring_kernel[(num_slots,)]( + kv, + kv.stride(0), + score, + score.stride(0), + ape, + ape.stride(0), + positions, + state_ring, + state_ring.stride(0), + state_ring.stride(1), + slot_mapping, + token_to_req_indices, + seq_lens, + HEAD_SIZE=head_size, + TRITON_BLOCK_SIZE=triton.next_power_of_2(head_size), + STATE_WIDTH=state_width, + RING_SIZE=state_ring.shape[1], + COMPRESS_RATIO=compress_ratio, + **(pdl_kwargs or {}), + ) + + +def stage_partial_states_from_ring( + kv: torch.Tensor, + score: torch.Tensor, + ape: torch.Tensor, + positions: torch.Tensor, + state_ring: torch.Tensor, + dense_state: torch.Tensor, + slot_mapping: torch.Tensor, + state_width: int, + history_size: int, + compress_ratio: int, + pdl_kwargs: dict | None = None, +) -> None: + """Stage private history and the current chunk into one dense state view.""" + num_slots = slot_mapping.shape[0] + copy_block_size = triton.next_power_of_2(2 * state_width) + _copy_ring_history_to_dense_kernel[(history_size,)]( + positions, + state_ring, + state_ring.stride(1), + dense_state, + dense_state.stride(0), + STATE_DIM=2 * state_width, + HISTORY_SIZE=history_size, + RING_SIZE=state_ring.shape[1], + BLOCK_SIZE=copy_block_size, + **(pdl_kwargs or {}), + ) + _save_current_states_to_dense_kernel[(num_slots,)]( + kv, + kv.stride(0), + score, + score.stride(0), + ape, + ape.stride(0), + positions, + dense_state, + dense_state.stride(0), + slot_mapping, + HEAD_SIZE=kv.shape[-1], + BLOCK_SIZE=triton.next_power_of_2(kv.shape[-1]), + STATE_WIDTH=state_width, + HISTORY_SIZE=history_size, + COMPRESS_RATIO=compress_ratio, + **(pdl_kwargs or {}), + ) + + @triton.jit def _save_partial_states_kernel( kv_ptr, @@ -99,3 +186,123 @@ def _save_partial_states_kernel( score + ape, mask=mask, ) + + +@triton.jit +def _save_partial_states_to_ring_kernel( + kv_ptr, + kv_stride, + score_ptr, + score_stride, + ape_ptr, + ape_stride, + positions_ptr, + state_ring_ptr, + state_ring_stride0, + state_ring_stride1, + slot_mapping_ptr, + token_to_req_indices_ptr, + seq_lens_ptr, + HEAD_SIZE: tl.constexpr, + TRITON_BLOCK_SIZE: tl.constexpr, + STATE_WIDTH: tl.constexpr, + RING_SIZE: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, +): + token_idx = tl.program_id(0) + if tl.load(slot_mapping_ptr + token_idx) < 0: + return + + req_idx = tl.load(token_to_req_indices_ptr + token_idx) + position = tl.load(positions_ptr + token_idx) + seq_len = tl.load(seq_lens_ptr + req_idx) + if position < seq_len - RING_SIZE: + return + + ring_row = position % RING_SIZE + base_ptr = ( + state_ring_ptr + + req_idx.to(tl.int64) * state_ring_stride0 + + ring_row.to(tl.int64) * state_ring_stride1 + ) + block = tl.arange(0, TRITON_BLOCK_SIZE) + mask = block < HEAD_SIZE + + kv = tl.load(kv_ptr + token_idx * kv_stride + block, mask=mask).to(tl.float32) + tl.store(base_ptr + block, kv, mask=mask) + + ape_row = position % COMPRESS_RATIO + ape = tl.load(ape_ptr + ape_row * ape_stride + block, mask=mask).to(tl.float32) + score = tl.load(score_ptr + token_idx * score_stride + block, mask=mask).to( + tl.float32 + ) + tl.store(base_ptr + STATE_WIDTH + block, score + ape, mask=mask) + + +@triton.jit +def _copy_ring_history_to_dense_kernel( + positions_ptr, + state_ring_ptr, + state_ring_stride1, + dense_state_ptr, + dense_state_stride0, + STATE_DIM: tl.constexpr, + HISTORY_SIZE: tl.constexpr, + RING_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + history_idx = tl.program_id(0) + current_start = tl.load(positions_ptr) + position = current_start - HISTORY_SIZE + history_idx + if position < 0: + return + + block = tl.arange(0, BLOCK_SIZE) + mask = block < STATE_DIM + ring_row = position % RING_SIZE + values = tl.load( + state_ring_ptr + ring_row.to(tl.int64) * state_ring_stride1 + block, + mask=mask, + ) + tl.store( + dense_state_ptr + history_idx * dense_state_stride0 + block, + values, + mask=mask, + ) + + +@triton.jit +def _save_current_states_to_dense_kernel( + kv_ptr, + kv_stride, + score_ptr, + score_stride, + ape_ptr, + ape_stride, + positions_ptr, + dense_state_ptr, + dense_state_stride0, + slot_mapping_ptr, + HEAD_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + STATE_WIDTH: tl.constexpr, + HISTORY_SIZE: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, +): + token_idx = tl.program_id(0) + if tl.load(slot_mapping_ptr + token_idx) < 0: + return + + block = tl.arange(0, BLOCK_SIZE) + mask = block < HEAD_SIZE + row_ptr = dense_state_ptr + (HISTORY_SIZE + token_idx) * dense_state_stride0 + kv = tl.load(kv_ptr + token_idx * kv_stride + block, mask=mask).to(tl.float32) + tl.store(row_ptr + block, kv, mask=mask) + + position = tl.load(positions_ptr + token_idx) + ape_row = position % COMPRESS_RATIO + ape = tl.load(ape_ptr + ape_row * ape_stride + block, mask=mask).to(tl.float32) + score = tl.load(score_ptr + token_idx * score_stride + block, mask=mask).to( + tl.float32 + ) + tl.store(row_ptr + STATE_WIDTH + block, score + ape, mask=mask) diff --git a/vllm/models/deepseek_v4/compressor.py b/vllm/models/deepseek_v4/compressor.py index 2738efc293..f9d8d43e4c 100644 --- a/vllm/models/deepseek_v4/compressor.py +++ b/vllm/models/deepseek_v4/compressor.py @@ -1,14 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Callable from dataclasses import dataclass from typing import Any, ClassVar, cast import torch from torch import nn +from vllm import envs from vllm.config import CUDAGraphMode, VllmConfig, get_current_vllm_config from vllm.forward_context import get_forward_context +from vllm.logger import init_logger from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import MergedColumnParallelLinear @@ -18,6 +21,8 @@ from vllm.models.deepseek_v4.common.ops.fused_indexer_q import MXFP4_BLOCK_SIZE from vllm.models.deepseek_v4.common.ops.save_partial_states import ( save_partial_states, + save_partial_states_to_ring, + stage_partial_states_from_ring, ) from vllm.platforms import current_platform from vllm.v1.attention.backend import ( @@ -32,6 +37,28 @@ MLAAttentionSpec, SlidingWindowMLASpec, ) +from vllm.v1.worker.workspace import current_workspace_manager + +logger = init_logger(__name__) + + +def _can_use_sm70_private_compressor_state(vllm_config: VllmConfig) -> bool: + kv_transfer_config = vllm_config.kv_transfer_config + spec_config = vllm_config.speculative_config + return ( + envs.VLLM_SM70_DSV4_PRIVATE_COMPRESSOR_STATE + and current_platform.is_cuda() + and current_platform.is_device_capability((7, 0)) + and vllm_config.scheduler_config.max_num_seqs == 1 + and not vllm_config.cache_config.enable_prefix_caching + and vllm_config.parallel_config.pipeline_parallel_size == 1 + # The private gather assumes one contiguous token chain. Parallel + # drafting uses multiple branches and must retain the paged fallback. + and (spec_config is None or not spec_config.parallel_drafting) + and ( + kv_transfer_config is None or not kv_transfer_config.is_kv_transfer_instance + ) + ) def _get_c128_boundary(metadata: CommonAttentionMetadata) -> bool | None: @@ -141,12 +168,27 @@ def __init__( dtype: torch.dtype, compress_ratio: int, prefix: str, + private_ring_size: int = 0, + max_num_reqs: int = 1, ): super().__init__() self.state_dim = state_dim self.dtype = dtype self.prefix = prefix self.kv_cache = torch.tensor([]) + self.private_ring_size = private_ring_size + if private_ring_size > 0: + self.register_buffer( + "private_state", + torch.empty( + (max_num_reqs, private_ring_size, state_dim), + dtype=dtype, + device=current_platform.device_type, + ), + persistent=False, + ) + else: + self.private_state = None compilation_config = get_current_vllm_config().compilation_config if prefix in compilation_config.static_forward_context: raise ValueError(f"Duplicate layer name: {prefix}") @@ -170,7 +212,9 @@ def __init__( else: raise ValueError(f"Invalid compress ratio: {compress_ratio}") - def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: + if self.private_state is not None: + return None return SlidingWindowMLASpec( # only has one vector instead of K + V block_size=self.block_size, num_kv_heads=1, @@ -206,6 +250,7 @@ def __init__( rotate: bool = False, prefix: str = "", k_cache_prefix="", + validity_cache_prefix="", use_fp4_cache: bool = False, ): super().__init__() @@ -215,6 +260,7 @@ def __init__( self.rotate = rotate self.prefix = prefix self.k_cache_prefix = k_cache_prefix + self.validity_cache_prefix = validity_cache_prefix self.use_fp4_cache = use_fp4_cache config = vllm_config.model_config.hf_config @@ -227,6 +273,17 @@ def __init__( self.overlap = compress_ratio == 4 self.coff = 1 + self.overlap + self.use_private_state = _can_use_sm70_private_compressor_state(vllm_config) + num_speculative_tokens = ( + vllm_config.speculative_config.num_speculative_tokens + if vllm_config.speculative_config is not None + else 0 + ) + private_ring_size = ( + self.coff * compress_ratio + max(1, num_speculative_tokens + 1) + if self.use_private_state + else 0 + ) state_dtype = torch.float32 self.ape = nn.Parameter( @@ -247,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( @@ -254,7 +312,15 @@ def __init__( dtype=state_dtype, compress_ratio=compress_ratio, prefix=f"{prefix}.state_cache", + private_ring_size=private_ring_size, + max_num_reqs=self.max_num_reqs, ) + if self.use_private_state: + logger.info_once( + "Using SM70 private DeepSeek V4 compressor state rings " + "(max_num_seqs=1, prefix cache disabled); paged compressor " + "state allocation is disabled." + ) # Save reference to static_forward_context for forward-time KV cache lookup. # get_current_vllm_config() is only available during __init__, not forward. @@ -303,19 +369,33 @@ def forward( if not isinstance(attn_metadata, dict): return - state_metadata = cast( - CompressorMetadata, attn_metadata[self.state_cache.prefix] - ) - token_to_req_indices = state_metadata.token_to_req_indices - slot_mapping = state_metadata.slot_mapping + seq_lens = None + if self.use_private_state: + assert self.validity_cache_prefix, ( + "Private compressor state requires an SWA metadata source." + ) + validity_metadata = cast(Any, attn_metadata[self.validity_cache_prefix]) + token_to_req_indices = validity_metadata.token_to_req_indices + slot_mapping = validity_metadata.slot_mapping + block_table = validity_metadata.block_table + block_size = validity_metadata.block_size + seq_lens = validity_metadata.seq_lens + state_cache = self.state_cache.private_state + assert state_cache is not None + assert token_to_req_indices is not None and seq_lens is not None + state_width = self.state_cache.state_dim // 2 + state_metadata = None + else: + state_metadata = cast( + CompressorMetadata, attn_metadata[self.state_cache.prefix] + ) + token_to_req_indices = state_metadata.token_to_req_indices + slot_mapping = state_metadata.slot_mapping + block_table = state_metadata.block_table + block_size = state_metadata.block_size + state_cache = self.state_cache.kv_cache + state_width = state_cache.shape[-1] // 2 num_actual = slot_mapping.shape[0] - block_table = state_metadata.block_table - block_size = state_metadata.block_size - - # [num_blocks, block_size, kv_dim+score_dim], where kv_dim == score_dim - state_cache = self.state_cache.kv_cache - # kv_state stored in first half, score_state stored in second half - state_width = state_cache.shape[-1] // 2 pdl_kwargs = ( {} if current_platform.is_rocm() or current_platform.is_xpu() @@ -328,21 +408,48 @@ def forward( # GEMM; state_cache from this kernel) but neither emits/waits on PDL # grid dependency primitives, so launch_pdl=True caused a # read-after-write race and non-deterministic output. - save_partial_states( - kv=kv, - score=score, - ape=self.ape, - positions=positions, - state_cache=state_cache, - slot_mapping=slot_mapping, - block_size=block_size, - state_width=state_width, - compress_ratio=self.compress_ratio, - pdl_kwargs=pdl_kwargs, - ) + if not self.use_private_state: + save_partial_states( + kv=kv, + score=score, + ape=self.ape, + positions=positions, + state_cache=state_cache, + slot_mapping=slot_mapping, + block_size=block_size, + state_width=state_width, + compress_ratio=self.compress_ratio, + pdl_kwargs=pdl_kwargs, + ) + + use_dense_private_state = self.use_private_state and self.compress_ratio == 128 + compression_state_cache = state_cache + private_history_size = self.coff * self.compress_ratio + if use_dense_private_state: + (dense_state,) = current_workspace_manager().get_simultaneous( + ( + (private_history_size + num_actual, 2 * state_width), + torch.float32, + ) + ) + stage_partial_states_from_ring( + kv=kv, + score=score, + ape=self.ape, + positions=positions, + state_ring=state_cache, + dense_state=dense_state, + slot_mapping=slot_mapping, + state_width=state_width, + history_size=private_history_size, + compress_ratio=self.compress_ratio, + pdl_kwargs=pdl_kwargs, + ) + compression_state_cache = dense_state.unsqueeze(0) if ( current_platform.is_cuda() + and not self.use_private_state and self.head_dim == 512 and self.compress_ratio == 128 and forward_context.cudagraph_runtime_mode != CUDAGraphMode.FULL @@ -361,6 +468,7 @@ def forward( k_cache_metadata = cast(Any, attn_metadata[self.k_cache_prefix]) kv_cache = self._static_forward_context[self.k_cache_prefix].kv_cache + compress_norm_rope_store_fn: Callable[..., None] if current_platform.is_cuda() and not current_platform.is_device_capability( (7, 0) ): @@ -382,8 +490,8 @@ def forward( # Always use a triton kernel. compress_norm_rope_store_fn = compress_norm_rope_store_triton - compress_norm_rope_store_fn( - state_cache=state_cache, + compress_kwargs = dict( + state_cache=compression_state_cache, num_actual=num_actual, token_to_req_indices=token_to_req_indices, positions=positions, @@ -406,3 +514,31 @@ def forward( token_stride=self._token_stride, scale_dim=self._scale_dim, ) + if compress_norm_rope_store_fn is compress_norm_rope_store_triton: + compress_kwargs.update( + current_kv=kv, + current_score=score, + ape=self.ape, + use_private_state=self.use_private_state + and not use_dense_private_state, + use_dense_private_state=use_dense_private_state, + ring_size=self.state_cache.private_ring_size or 1, + private_history_size=private_history_size, + ) + compress_norm_rope_store_fn(**compress_kwargs) + + if self.use_private_state: + assert seq_lens is not None + save_partial_states_to_ring( + kv=kv, + score=score, + ape=self.ape, + positions=positions, + state_ring=state_cache, + slot_mapping=slot_mapping, + token_to_req_indices=token_to_req_indices, + seq_lens=seq_lens, + state_width=state_width, + compress_ratio=self.compress_ratio, + pdl_kwargs=pdl_kwargs, + ) diff --git a/vllm/models/deepseek_v4/nvidia/dspark.py b/vllm/models/deepseek_v4/nvidia/dspark.py index da3319124a..84f711947c 100644 --- a/vllm/models/deepseek_v4/nvidia/dspark.py +++ b/vllm/models/deepseek_v4/nvidia/dspark.py @@ -15,6 +15,7 @@ import torch import torch.nn as nn +from vllm import envs from vllm.config import VllmConfig, get_current_vllm_config from vllm.distributed import ( get_tensor_model_parallel_rank, @@ -71,8 +72,40 @@ def bias(self, markov_embed: torch.Tensor) -> torch.Tensor: return self.markov_w2(markov_embed) +class DSparkConfidenceHead(nn.Module): + """Replicated FP32 conditional-acceptance predictor.""" + + def __init__(self, input_size: int, prefix: str) -> None: + super().__init__() + # The checkpoint tensor is BF16, but the official implementation keeps + # this projection and its input in FP32 for calibration accuracy. + self.proj = ReplicatedLinear( + input_size, + 1, + bias=False, + params_dtype=torch.float32, + quant_config=None, + return_bias=False, + prefix=maybe_prefix(prefix, "proj"), + ) + + def forward( + self, + hidden_states: torch.Tensor, + markov_embeds: torch.Tensor, + ) -> torch.Tensor: + features = torch.cat((hidden_states.float(), markov_embeds.float()), dim=-1) + return self.proj(features).squeeze(-1) + + class DSparkDeepseekV4Model(nn.Module): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + def __init__( + self, + *, + vllm_config: VllmConfig, + prefix: str = "", + use_confidence_head: bool = False, + ) -> None: super().__init__() assert vllm_config.speculative_config is not None config = vllm_config.speculative_config.draft_model_config.hf_config @@ -147,6 +180,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: config.dspark_markov_rank, prefix=maybe_prefix(prefix, "markov_head"), ) + self.confidence_head = ( + DSparkConfidenceHead( + config.hidden_size + config.dspark_markov_rank, + prefix=maybe_prefix(prefix, "confidence_head"), + ) + if use_confidence_head + else None + ) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) @@ -263,8 +304,16 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: assert vllm_config.speculative_config is not None self.draft_model_config = vllm_config.speculative_config.draft_model_config self.config = self.draft_model_config.hf_config + speculative_config = vllm_config.speculative_config + self._confidence_head_required = bool( + speculative_config.dspark_confidence_threshold > 0.0 + or envs.VLLM_SPEC_DUMP_ALIGNMENT + ) + self._confidence_head_loaded = False self.model = DSparkDeepseekV4Model( - vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "model"), + use_confidence_head=self._confidence_head_required, ) self.lm_head = ParallelLMHead( self.config.vocab_size, @@ -311,6 +360,21 @@ def markov_embed(self, token_ids: torch.Tensor) -> torch.Tensor: def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor: return self.model.markov_head.bias(markov_embed) + def confidence_logits( + self, + hidden_states: torch.Tensor, + markov_embeds: torch.Tensor, + ) -> torch.Tensor: + if not self._confidence_head_loaded: + raise RuntimeError( + "DSpark confidence logits were requested, but the draft " + "checkpoint did not provide confidence_head.proj.weight." + ) + confidence_head = self.model.confidence_head + if confidence_head is None: + raise RuntimeError("DSpark confidence head was not initialized.") + return confidence_head(hidden_states, markov_embeds) + def map_draft_to_target(self, draft_ids: torch.Tensor) -> torch.Tensor: return draft_ids @@ -352,6 +416,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: name = self._remap_dspark_name(checkpoint_name) if name is None: continue + # Ordinary serving does not instantiate the optional diagnostic / + # scheduler head, so ignore its checkpoint tensor without adding + # an unused parameter to the runtime model. + if name.startswith("model.confidence_head.") and name not in params_dict: + continue if name.endswith(".scale"): suffix = ( expert_scale_suffix @@ -416,6 +485,13 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: for layer in self.model.layers: layer.ffn.finalize_mega_moe_weights() + confidence_weight_name = "model.confidence_head.proj.weight" + self._confidence_head_loaded = confidence_weight_name in loaded_params + if self._confidence_head_required and not self._confidence_head_loaded: + raise ValueError( + "DSpark confidence collection requires the draft checkpoint " + f"parameter {confidence_weight_name!r}." + ) logger.info_once("DSpark draft model loaded: %d params", len(loaded_params)) return loaded_params @@ -426,14 +502,13 @@ def _remap_dspark_name(name: str) -> str | None: return None stage = int(match.group(1)) rest = match.group(2) - if rest.startswith("confidence_head."): - return None head_prefixes = ( "norm.", "hc_head_fn", "hc_head_base", "hc_head_scale", "markov_head.", + "confidence_head.", ) if rest.startswith(("main_proj.", "main_norm.")) or rest.startswith( head_prefixes diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index fb165aa730..750b3179ea 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -718,6 +718,27 @@ def __init__( rope_parameters["mscale_all_dim"] = 0 # Disable mscale rope_parameters["is_deepseek_v4"] = True rope_parameters["rope_dim"] = self.rope_head_dim + if self.compress_ratio <= 1: + # The sliding-window-only layers get base rope_theta above, but the + # reference also turns YaRN itself off for them: it passes + # original_seq_len=0, and precompute_freqs_cis then skips the + # correction range entirely (inference/model.py:481-486, :227). + # Swapping only theta leaves the interpolation on, which shifts the + # low-frequency pairs by up to ~2 degrees across the 128-token + # window. factor=1 makes the interpolation an exact identity + # (inv_freq_interpolation == inv_freq_extrapolation) and mscale + # collapses to 1.0, so this is plain rope without swapping the + # class -- DeepseekV4ScalingRotaryEmbedding still rotates the last + # rotary_dim and keeps the cache in fp32, which the fused kernels + # require. Copy first: the dict is config.rope_parameters itself + # and every later layer reads it back. + rope_parameters = dict(rope_parameters) + rope_parameters["factor"] = 1 + # Cache length is original_max_position * factor, so it has to + # carry the full position range now that factor is 1. + rope_parameters["original_max_position_embeddings"] = ( + self.max_position_embeddings + ) self.rotary_emb = get_rope( self.head_dim, max_position=self.max_position_embeddings, @@ -780,6 +801,10 @@ def __init__( quant_config=quant_config, prefix=prefix, ) + if self.indexer is not None: + self.indexer.compressor.validity_cache_prefix = ( + self.mla_attn.swa_cache_layer.prefix + ) def forward( self, 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/indexer.py b/vllm/models/deepseek_v4/sm70/indexer.py index 369dc2e147..291f150cc3 100644 --- a/vllm/models/deepseek_v4/sm70/indexer.py +++ b/vllm/models/deepseek_v4/sm70/indexer.py @@ -3,16 +3,100 @@ """DeepSeek V4 C4 indexer fallback using FP16 HMMA on SM70.""" +import os + import torch from vllm.models.deepseek_v4.common.ops.fp8_software import ( fp8_e4m3fn_bits_to_fp32, ) +from vllm.platforms import current_platform from vllm.triton_utils import tl, triton +from vllm.v1.worker.workspace import current_workspace_manager _INDEX_HEAD_DIM = 128 _INDEX_CACHE_BYTES = _INDEX_HEAD_DIM + 4 +# Fuse dequant+scoring into one paged kernel instead of materialising +# [rows, max_seq_len, 128] fp16 and running a dense bmm. Set to 0 to get the +# old gather+bmm path back for an A/B. Only affects decode; prefill scores a +# contiguous chunk whose width is the real chunk length already. +_FUSED_DECODE_LOGITS = os.getenv("VLLM_SM70_INDEXER_FUSED_LOGITS", "1") == "1" + +# Score with the model's actual indexer function, +# I[t, s] = sum_h weights[t, h] * relu(q[t, h] . k[s]), +# instead of the cheaper (sum_h weights[t, h] * q[t, h]) . k[s]. Those are +# equal only without the relu, and the relu is not optional: it is what the +# checkpoint was trained with (reference `Indexer.forward`, +# inference/model.py:427, and the in-repo `fp8_mqa_logits` reference at +# v1/attention/ops/rocm_aiter_mla_sparse.py:510 -- every non-SM70 backend +# applies it inside its logits kernel). Dropping it costs 64x fewer flops and +# ranks the wrong keys, which only shows up once top-k is actually selective. +# Set 0 to restore the factored form for an A/B. +_RELU_LOGITS = os.getenv("VLLM_SM70_INDEXER_RELU", "1") == "1" + +# Key-axis splits are chosen so the launch is roughly this many blocks, which +# keeps all 80 SMs busy at long context without paying for thousands of +# no-op programs at short context. +_LOGITS_TARGET_BLOCKS = 512 +_LOGITS_BLOCK_N = 32 +# The relu path multiplies q against k with an MMA instead of a broadcast +# multiply, so the key tile has to be at least one 16x16x16 HMMA fragment +# wide; 64 keeps the [64 heads, 128] x [128, N] shape efficient. +_RELU_BLOCK_N = 64 +_RELU_BLOCK_M = 32 + +# Route prefill scoring through cuBLAS instead of the Triton MMA. Set 0 for an +# A/B against the Triton kernel above (which stays as the reference shape and +# is what the decode path still uses, since decode gathers from a paged cache). +_PREFILL_CUBLAS = os.getenv("VLLM_SM70_INDEXER_PREFILL_CUBLAS", "1") == "1" +# Cap (MiB) on the [tokens*heads, key_tile] fp16 score tile. At the default +# 2048-token prefill chunk and 64 heads this is 128 KiB per key, so the tile +# is what keeps a long-context chunk from asking for gigabytes in one go. +_PREFILL_TILE_MB = int(os.getenv("VLLM_SM70_INDEXER_PREFILL_TILE_MB", "192")) +_EPILOGUE_BLOCK_K = 256 +_EPILOGUE_BLOCK_H = 8 + +# At long context the paged Triton decode kernel rescans the same FP8 MQA keys +# once for every verifier row and head. For a bounded one-request verifier +# shape, gather the keys once and let cuBLAS score all rows/heads together. +# FP32 output is deliberate: an FP16 score intermediate produced rare top-k +# set changes at 16K/64K compressed tokens. Short and generic shapes retain the +# fused paged kernel, whose launch overhead is lower there. +_DECODE_CUBLAS = os.getenv("VLLM_SM70_INDEXER_DECODE_CUBLAS", "0") == "1" +_DECODE_CUBLAS_MIN_KEYS = int( + os.getenv("VLLM_SM70_INDEXER_DECODE_CUBLAS_MIN_KEYS", "1024") +) +_DECODE_CUBLAS_MAX_ROWS = 8 + + +@triton.jit +def _index_cache_ptrs( + cache_ptr, + physical_block, + pos_in_block, + block_stride, + cache_block_size, + head_dim: tl.constexpr, +): + """Value and scale pointers for one token of the paged indexer cache. + + The layout is block-major, NOT token-interleaved: within a block come all + `cache_block_size` rows of FP8 values, and only then the FP32 scales, one + per token (csrc/libtorch_stable/cache_kernels.cu, the + `indexer_k_quant_and_cache` store and the `cp_gather_indexer_k_quant_cache` + load both address it this way). Reading a token as + `[128 values][4 scale]` at `pos * cache.stride(1)` therefore lands on the + wrong bytes for every token but the first of each block, and the "scale" + it picks up is really four FP8 value bytes. + """ + base = cache_ptr + physical_block.to(tl.int64) * block_stride + value_ptr = base + pos_in_block.to(tl.int64) * head_dim + scale_ptr = (base + cache_block_size * head_dim + pos_in_block * 4).to( + tl.pointer_type(tl.float32) + ) + return value_ptr, scale_ptr + @triton.jit def _weighted_query_kernel( @@ -79,7 +163,6 @@ def _dequant_paged_index_k_kernel( seq_lens_ptr, out_ptr, cache_stride0, - cache_stride1, block_table_stride0, out_stride0, out_stride1, @@ -100,10 +183,13 @@ def _dequant_paged_index_k_kernel( mask=valid, other=0, ) - token_ptr = ( - cache_ptr - + physical_block.to(tl.int64) * cache_stride0 - + pos_in_block * cache_stride1 + token_ptr, scale_ptr = _index_cache_ptrs( + cache_ptr, + physical_block, + pos_in_block, + cache_stride0, + cache_block_size, + head_dim, ) packed = tl.load( token_ptr[:, None] + dim_offsets[None, :], @@ -111,7 +197,6 @@ def _dequant_paged_index_k_kernel( other=0, ) fp8 = fp8_e4m3fn_bits_to_fp32(packed) - scale_ptr = (token_ptr + head_dim).to(tl.pointer_type(tl.float32)) scale = tl.load(scale_ptr, mask=valid, other=0.0).to(tl.float32) dequant = fp8 * scale[:, None] tl.store( @@ -124,6 +209,285 @@ def _dequant_paged_index_k_kernel( ) +@triton.jit +def _gather_paged_index_k_unscaled_kernel( + cache_ptr, + block_table_ptr, + seq_lens_ptr, + out_k_ptr, + out_scale_ptr, + cache_stride0, + out_k_stride0, + cache_block_size, + max_seq_len, + num_rows, + head_dim: tl.constexpr, + MAX_ROWS: tl.constexpr, + BLOCK_N: tl.constexpr, +): + """Gather one request's live paged keys once for all verifier rows. + + ``max_seq_len`` is a CUDA-graph bucket and can be much wider than the + currently allocated request. Derive the live bound on device from the + row lengths so graph replay never dereferences a ``-1`` block-table tail. + Rows outside the live bound are zeroed to keep the workspace deterministic; + the downstream top-k still applies each row's own sequence length. + """ + key_offsets = tl.program_id(0) * BLOCK_N + tl.arange(0, BLOCK_N) + dim_offsets = tl.arange(0, head_dim) + row_offsets = tl.arange(0, MAX_ROWS) + row_lens = tl.load( + seq_lens_ptr + row_offsets, + mask=row_offsets < num_rows, + other=0, + ) + live_seq_len = tl.max(row_lens, axis=0) + in_bucket = key_offsets < max_seq_len + valid = in_bucket & (key_offsets < live_seq_len) + + block_in_seq = key_offsets // cache_block_size + pos_in_block = key_offsets % cache_block_size + physical_block = tl.load( + block_table_ptr + block_in_seq, + mask=valid, + other=0, + ) + token_ptr, scale_ptr = _index_cache_ptrs( + cache_ptr, + physical_block, + pos_in_block, + cache_stride0, + cache_block_size, + head_dim, + ) + packed = tl.load( + token_ptr[:, None] + dim_offsets[None, :], + mask=valid[:, None], + other=0, + ) + scales = tl.load(scale_ptr, mask=valid, other=0.0).to(tl.float32) + tl.store( + out_k_ptr + key_offsets[:, None] * out_k_stride0 + dim_offsets[None, :], + fp8_e4m3fn_bits_to_fp32(packed).to(tl.float16), + mask=in_bucket[:, None], + ) + tl.store(out_scale_ptr + key_offsets, scales, mask=in_bucket) + + +@triton.jit +def _paged_index_logits_kernel( + q_ptr, + cache_ptr, + block_table_ptr, + seq_lens_ptr, + out_ptr, + q_stride0, + cache_stride0, + block_table_stride0, + out_stride0, + cache_block_size, + head_dim: tl.constexpr, + BLOCK_N: tl.constexpr, + NUM_SPLITS: tl.constexpr, +): + """Index logits straight out of the paged FP8 cache, no gather buffer. + + `_dequant_paged_index_k_kernel` + bmm has a grid and a GEMM shape derived + from `max_seq_len`, and under a captured decode CUDA Graph that is frozen + at `max_model_len` forever. Here the grid is static -- (rows, NUM_SPLITS), + sized from `max_seq_len` on the host -- and only each program's trip count + follows the live `seq_len`, so one captured graph serves any context while + a short context costs almost nothing. + """ + row_idx = tl.program_id(0) + split_id = tl.program_id(1) + dim_offsets = tl.arange(0, head_dim) + q = tl.load(q_ptr + row_idx * q_stride0 + dim_offsets).to(tl.float32) + + seq_len = tl.load(seq_lens_ptr + row_idx) + chunks = (seq_len + BLOCK_N - 1) // BLOCK_N + chunks_per_split = (chunks + NUM_SPLITS - 1) // NUM_SPLITS + lo = split_id * chunks_per_split * BLOCK_N + hi = tl.minimum(lo + chunks_per_split * BLOCK_N, seq_len) + + for start in range(lo, hi, BLOCK_N): + key_offsets = start + tl.arange(0, BLOCK_N) + valid = key_offsets < seq_len + block_in_seq = key_offsets // cache_block_size + pos_in_block = key_offsets % cache_block_size + physical_block = tl.load( + block_table_ptr + row_idx * block_table_stride0 + block_in_seq, + mask=valid, + other=0, + ) + token_ptr, scale_ptr = _index_cache_ptrs( + cache_ptr, + physical_block, + pos_in_block, + cache_stride0, + cache_block_size, + head_dim, + ) + packed = tl.load( + token_ptr[:, None] + dim_offsets[None, :], + mask=valid[:, None], + other=0, + ) + fp8 = fp8_e4m3fn_bits_to_fp32(packed) + scale = tl.load(scale_ptr, mask=valid, other=0.0).to(tl.float32) + # Per-token scale factors out of the dot product. + logits = tl.sum(fp8 * q[None, :], axis=1) * scale + tl.store(out_ptr + row_idx * out_stride0 + key_offsets, logits, mask=valid) + + +@triton.jit +def _contiguous_index_logits_relu_kernel( + q_ptr, + w_ptr, + k_ptr, + k_scale_ptr, + out_ptr, + q_stride0, + q_stride1, + w_stride0, + k_stride0, + out_stride0, + num_q, + num_k, + num_heads, + head_dim: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + """Prefill index scores, per head, with the relu inside the head sum. + + Reference shape for the scoring, and the A/B fallback for the cuBLAS path + (`VLLM_SM70_INDEXER_PREFILL_CUBLAS=0`). It is not the fast path: Triton's + `tl.dot` does not reach Volta's tensor cores, so this tops out near 3.3 + TFLOP/s no matter how the tile is shaped -- folding heads into the M + dimension to get a [256, 128] x [128, 64] MMA instead of a per-head + [32, 128] x [128, 64] one measured 3.34 vs 3.39 TFLOP/s, i.e. nothing. + cuBLAS does the same shape at 54. + + The key tile is loaded once and reused by every head (the indexer is MQA). + FP8 keys are widened to FP16 *unscaled* -- e4m3 maxes out at 448, which + fp16 holds exactly -- and the per-key scale is applied to the finished + accumulator. That keeps `k * scale` out of fp16 entirely, where a large + ue8m0 scale could otherwise overflow. + """ + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + dim_offsets = tl.arange(0, head_dim) + m_valid = offs_m < num_q + n_valid = offs_n < num_k + + packed = tl.load( + k_ptr + offs_n[:, None] * k_stride0 + dim_offsets[None, :], + mask=n_valid[:, None], + other=0, + ) + k_t = tl.trans(fp8_e4m3fn_bits_to_fp32(packed).to(tl.float16)) + scale = tl.load(k_scale_ptr + offs_n, mask=n_valid, other=0.0).to(tl.float32) + + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + for head in range(num_heads): + q = tl.load( + q_ptr + + offs_m[:, None] * q_stride0 + + head * q_stride1 + + dim_offsets[None, :], + mask=m_valid[:, None], + other=0.0, + ) + w = tl.load(w_ptr + offs_m * w_stride0 + head, mask=m_valid, other=0.0) + acc += tl.maximum(tl.dot(q, k_t), 0.0) * w[:, None].to(tl.float32) + + tl.store( + out_ptr + offs_m[:, None] * out_stride0 + offs_n[None, :], + acc * scale[None, :], + mask=m_valid[:, None] & n_valid[None, :], + ) + + +@triton.jit +def _paged_index_logits_relu_kernel( + q_ptr, + w_ptr, + cache_ptr, + block_table_ptr, + seq_lens_ptr, + out_ptr, + q_stride0, + q_stride1, + w_stride0, + cache_stride0, + block_table_stride0, + out_stride0, + cache_block_size, + head_dim: tl.constexpr, + NUM_HEADS: tl.constexpr, + BLOCK_N: tl.constexpr, + NUM_SPLITS: tl.constexpr, +): + """Decode counterpart of `_contiguous_index_logits_relu_kernel`. + + Same static-grid property as `_paged_index_logits_kernel`: the grid is + (rows, NUM_SPLITS) and only the trip count follows the live `seq_len`, so + one captured CUDA graph serves any context length. + """ + row_idx = tl.program_id(0) + split_id = tl.program_id(1) + head_offsets = tl.arange(0, NUM_HEADS) + dim_offsets = tl.arange(0, head_dim) + q = tl.load( + q_ptr + + row_idx * q_stride0 + + head_offsets[:, None] * q_stride1 + + dim_offsets[None, :] + ) + w = tl.load(w_ptr + row_idx * w_stride0 + head_offsets).to(tl.float32) + + seq_len = tl.load(seq_lens_ptr + row_idx) + chunks = (seq_len + BLOCK_N - 1) // BLOCK_N + chunks_per_split = (chunks + NUM_SPLITS - 1) // NUM_SPLITS + lo = split_id * chunks_per_split * BLOCK_N + hi = tl.minimum(lo + chunks_per_split * BLOCK_N, seq_len) + + for start in range(lo, hi, BLOCK_N): + key_offsets = start + tl.arange(0, BLOCK_N) + valid = key_offsets < seq_len + block_in_seq = key_offsets // cache_block_size + pos_in_block = key_offsets % cache_block_size + physical_block = tl.load( + block_table_ptr + row_idx * block_table_stride0 + block_in_seq, + mask=valid, + other=0, + ) + token_ptr, scale_ptr = _index_cache_ptrs( + cache_ptr, + physical_block, + pos_in_block, + cache_stride0, + cache_block_size, + head_dim, + ) + packed = tl.load( + token_ptr[:, None] + dim_offsets[None, :], + mask=valid[:, None], + other=0, + ) + k_t = tl.trans(fp8_e4m3fn_bits_to_fp32(packed).to(tl.float16)) + scale = tl.load(scale_ptr, mask=valid, other=0.0).to(tl.float32) + # relu(q_h . k) is per head; only then does the weighted head sum + # collapse it to one logit. The per-key scale is positive, so pulling + # it out past the relu and the sum is exact. + scores = tl.maximum(tl.dot(q, k_t), 0.0) * w[:, None] + logits = tl.sum(scores, axis=0) * scale + tl.store(out_ptr + row_idx * out_stride0 + key_offsets, logits, mask=valid) + + def _combine_index_queries(q: torch.Tensor, weights: torch.Tensor) -> torch.Tensor: assert q.dtype == torch.float16 and q.ndim == 3 assert q.shape[-1] == _INDEX_HEAD_DIM @@ -146,6 +510,134 @@ def _combine_index_queries(q: torch.Tensor, weights: torch.Tensor) -> torch.Tens return out +@triton.jit +def _relu_weight_headsum_kernel( + s_ptr, + w_ptr, + scale_ptr, + out_ptr, + s_stride0, + s_stride1, + w_stride0, + out_stride0, + num_k, + num_heads: tl.constexpr, + BLOCK_K: tl.constexpr, + BLOCK_H: tl.constexpr, +): + """relu -> weight -> head sum -> key scale, in a single pass over scores. + + Epilogue for the cuBLAS prefill path. Doing it as torch ops instead costs + three extra round trips through the [tokens, heads, keys] score tile, which + is the largest tensor in the whole indexer; fusing them measured 13 -> 26 + TFLOP/s end to end. + """ + token = tl.program_id(0) + k_base = tl.program_id(1) * BLOCK_K + k_offsets = k_base + tl.arange(0, BLOCK_K) + k_valid = k_offsets < num_k + + acc = tl.zeros((BLOCK_K,), dtype=tl.float32) + for head in range(0, num_heads, BLOCK_H): + h_offsets = head + tl.arange(0, BLOCK_H) + s = tl.load( + s_ptr + + token * s_stride0 + + h_offsets[:, None] * s_stride1 + + k_offsets[None, :], + mask=k_valid[None, :], + other=0.0, + ).to(tl.float32) + w = tl.load(w_ptr + token * w_stride0 + h_offsets).to(tl.float32) + acc += tl.sum(tl.maximum(s, 0.0) * w[:, None], axis=0) + + scale = tl.load(scale_ptr + k_offsets, mask=k_valid, other=0.0).to(tl.float32) + tl.store(out_ptr + token * out_stride0 + k_offsets, acc * scale, mask=k_valid) + + +@triton.jit +def _dequant_index_k_unscaled_kernel( + k_ptr, + out_ptr, + k_stride0, + out_stride0, + head_dim: tl.constexpr, +): + """Widen FP8 index keys to FP16 without applying the per-key scale. + + e4m3 tops out at 448 and the scale is folded into the epilogue instead, so + the GEMM never sees `k * scale` -- a large ue8m0 scale would otherwise + overflow fp16 on the way in. + """ + row_idx = tl.program_id(0) + offsets = tl.arange(0, head_dim) + values = tl.load( + k_ptr.to(tl.pointer_type(tl.uint8)) + row_idx * k_stride0 + offsets + ) + tl.store( + out_ptr + row_idx * out_stride0 + offsets, + fp8_e4m3fn_bits_to_fp32(values).to(tl.float16), + ) + + +def _prefill_logits_cublas(q, k_quant, k_scales, weights): + """Prefill index scores through cuBLAS, with the relu inside the head sum. + + Triton's `tl.dot` does not reach Volta's tensor cores: on the exact shape + this needs it measures 3.3 TFLOP/s, against 54 TFLOP/s for cuBLAS fp16 and + 13 TFLOP/s for plain fp32 CUDA cores. The relu blocks the one-GEMM + factorisation, but it does not block cuBLAS -- stacking heads into the M + dimension turns the scoring into a single [tokens*heads, 128] x [128, keys] + GEMM whose output is reduced over heads afterwards. Keys are tiled so that + intermediate stays bounded; it is the largest allocation in the indexer. + """ + num_q, num_heads, head_dim = q.shape + num_k = k_quant.shape[0] + out = torch.empty((num_q, num_k), dtype=torch.float32, device=q.device) + + k_fp16 = torch.empty((num_k, head_dim), dtype=torch.float16, device=q.device) + _dequant_index_k_unscaled_kernel[(num_k,)]( + k_quant.view(torch.uint8), + k_fp16, + k_quant.stride(0), + k_fp16.stride(0), + head_dim=head_dim, + num_warps=4, + ) + + q2 = q.reshape(num_q * num_heads, head_dim) + if not q2.is_contiguous(): + q2 = q2.contiguous() + + rows = num_q * num_heads + budget = max(1, _PREFILL_TILE_MB * 2**20 // max(1, rows * 2)) + k_tile = max(_RELU_BLOCK_N, min(num_k, budget // _RELU_BLOCK_N * _RELU_BLOCK_N)) + + for start in range(0, num_k, k_tile): + stop = min(start + k_tile, num_k) + width = stop - start + # fp16 out with fp32 accumulate: the epilogue reduces in fp32 and only + # the top-k *set* is consumed downstream, so the 5e-4 rounding here is + # far below the gap between neighbouring index scores. + scores = torch.mm(q2, k_fp16[start:stop].t()) + _relu_weight_headsum_kernel[(num_q, triton.cdiv(width, _EPILOGUE_BLOCK_K))]( + scores, + weights, + k_scales[start:stop], + out[:, start:stop], + scores.stride(0) * num_heads, + scores.stride(0), + weights.stride(0), + out.stride(0), + width, + num_heads=num_heads, + BLOCK_K=_EPILOGUE_BLOCK_K, + BLOCK_H=_EPILOGUE_BLOCK_H, + num_warps=4, + ) + return out + + def sm70_indexer_prefill_logits( q: torch.Tensor, k_quant: torch.Tensor, @@ -158,6 +650,37 @@ def sm70_indexer_prefill_logits( k_scales = k_scale_storage.view(torch.float32).reshape(-1) assert k_scales.shape[0] == k_quant.shape[0] + if _RELU_LOGITS and _PREFILL_CUBLAS: + return _prefill_logits_cublas(q, k_quant, k_scales, weights) + + if _RELU_LOGITS: + num_q, num_heads, _ = q.shape + num_k = k_quant.shape[0] + out = torch.empty((num_q, num_k), dtype=torch.float32, device=q.device) + block_m, block_n = _RELU_BLOCK_M, _RELU_BLOCK_N + _contiguous_index_logits_relu_kernel[ + (triton.cdiv(num_q, block_m), triton.cdiv(num_k, block_n)) + ]( + q, + weights, + k_quant.view(torch.uint8), + k_scales, + out, + q.stride(0), + q.stride(1), + weights.stride(0), + k_quant.stride(0), + out.stride(0), + num_q, + num_k, + num_heads, + head_dim=_INDEX_HEAD_DIM, + BLOCK_M=block_m, + BLOCK_N=block_n, + num_warps=8, + ) + return out + weighted_q = _combine_index_queries(q, weights) k_fp16 = torch.empty(k_quant.shape, dtype=torch.float16, device=k_quant.device) _dequant_contiguous_index_k_kernel[(k_quant.shape[0],)]( @@ -172,6 +695,68 @@ def sm70_indexer_prefill_logits( return torch.mm(weighted_q, k_fp16.t(), out_dtype=torch.float32) +def _decode_logits_cublas( + q: torch.Tensor, + cache: torch.Tensor, + weights: torch.Tensor, + flat_lens: torch.Tensor, + block_table: torch.Tensor, + max_seq_len: int, +) -> torch.Tensor: + """Score a one-request decode by sharing one paged gather across rows.""" + total_rows, num_heads, head_dim = q.shape + workspace = current_workspace_manager() + gathered_k, k_scales, scores = workspace.get_simultaneous( + ((max_seq_len, head_dim), torch.float16), + ((max_seq_len,), torch.float32), + ((total_rows * num_heads, max_seq_len), torch.float32), + ) + block_n = 16 + _gather_paged_index_k_unscaled_kernel[(triton.cdiv(max_seq_len, block_n),)]( + cache, + block_table, + flat_lens, + gathered_k, + k_scales, + cache.stride(0), + gathered_k.stride(0), + cache.shape[1], + max_seq_len, + total_rows, + head_dim=head_dim, + MAX_ROWS=_DECODE_CUBLAS_MAX_ROWS, + BLOCK_N=block_n, + num_warps=4, + ) + + q2 = q.reshape(total_rows * num_heads, head_dim) + torch.mm( + q2, + gathered_k.t(), + out=scores, + out_dtype=torch.float32, + ) + out = torch.empty((total_rows, max_seq_len), dtype=torch.float32, device=q.device) + _relu_weight_headsum_kernel[ + (total_rows, triton.cdiv(max_seq_len, _EPILOGUE_BLOCK_K)) + ]( + scores, + weights, + k_scales, + out, + scores.stride(0) * num_heads, + scores.stride(0), + weights.stride(0), + out.stride(0), + max_seq_len, + num_heads=num_heads, + BLOCK_K=_EPILOGUE_BLOCK_K, + BLOCK_H=_EPILOGUE_BLOCK_H, + num_warps=4, + ) + return out + + def sm70_indexer_decode_logits( q: torch.Tensor, cache: torch.Tensor, @@ -183,33 +768,142 @@ def sm70_indexer_decode_logits( """Gather paged FP8 index keys and compute batched decode scores.""" assert cache.dtype == torch.uint8 and cache.ndim == 3 assert cache.shape[-1] >= _INDEX_CACHE_BYTES - weighted_q = _combine_index_queries(q, weights) + # The relu path scores per head, so it needs q as [rows, heads, dim]; the + # factored path collapses the head axis up front. + weighted_q = q if _RELU_LOGITS else _combine_index_queries(q, weights) - if seq_lens.ndim == 2: + single_request = block_table.shape[0] == 1 + native_rows = seq_lens.ndim == 2 + if native_rows: next_n = seq_lens.shape[1] flat_lens = seq_lens.reshape(-1).to(torch.int32) - block_table = block_table.repeat_interleave(next_n, dim=0) else: flat_lens = seq_lens.reshape(-1).to(torch.int32) assert flat_lens.shape[0] == weighted_q.shape[0] - assert block_table.shape[0] == weighted_q.shape[0] max_seq_len = max(1, int(max_seq_len)) + total_rows = weighted_q.shape[0] + + if ( + _RELU_LOGITS + and _DECODE_CUBLAS + and current_platform.is_cuda() + and current_platform.is_device_capability((7, 0)) + and max_seq_len >= _DECODE_CUBLAS_MIN_KEYS + and single_request + and (native_rows or total_rows == 1) + and 0 < total_rows <= _DECODE_CUBLAS_MAX_ROWS + and q.ndim == 3 + and q.dtype == torch.float16 + and q.shape[2] == _INDEX_HEAD_DIM + and q.shape[1] > 0 + and q.shape[1] % _EPILOGUE_BLOCK_H == 0 + and weights.ndim == 2 + and weights.shape == q.shape[:2] + and weights.dtype == torch.float32 + and cache.ndim == 3 + and cache.dtype == torch.uint8 + and cache.shape[1] > 0 + and cache.shape[2] == _INDEX_CACHE_BYTES + and cache.stride(1) == _INDEX_CACHE_BYTES + and cache.stride(2) == 1 + and block_table.ndim == 2 + and block_table.dtype == torch.int32 + and flat_lens.dtype == torch.int32 + and triton.cdiv(max_seq_len, cache.shape[1]) <= block_table.shape[1] + and q.device + == cache.device + == weights.device + == flat_lens.device + == block_table.device + and q.is_contiguous() + and weights.is_contiguous() + and cache.is_contiguous() + and flat_lens.is_contiguous() + and block_table.is_contiguous() + ): + return _decode_logits_cublas( + q, + cache, + weights, + flat_lens, + block_table, + max_seq_len, + ) + + if native_rows: + block_table = block_table.repeat_interleave(next_n, dim=0) + assert block_table.shape[0] == weighted_q.shape[0] + + if _RELU_LOGITS: + out = torch.empty( + (total_rows, max_seq_len), dtype=torch.float32, device=q.device + ) + max_chunks = triton.cdiv(max_seq_len, _RELU_BLOCK_N) + splits = max(1, min(max_chunks, _LOGITS_TARGET_BLOCKS // total_rows)) + _paged_index_logits_relu_kernel[(total_rows, splits)]( + q, + weights, + cache, + block_table, + flat_lens, + out, + q.stride(0), + q.stride(1), + weights.stride(0), + cache.stride(0), + block_table.stride(0), + out.stride(0), + cache.shape[1], + head_dim=_INDEX_HEAD_DIM, + NUM_HEADS=q.shape[1], + BLOCK_N=_RELU_BLOCK_N, + NUM_SPLITS=splits, + num_warps=8, + ) + # Positions >= seq_len are left uninitialised; every consumer masks by + # seq_lens (see the note on the factored path below). + return out + + if _FUSED_DECODE_LOGITS: + out = torch.empty( + (total_rows, max_seq_len), dtype=torch.float32, device=q.device + ) + max_chunks = triton.cdiv(max_seq_len, _LOGITS_BLOCK_N) + splits = max(1, min(max_chunks, _LOGITS_TARGET_BLOCKS // total_rows)) + _paged_index_logits_kernel[(total_rows, splits)]( + weighted_q, + cache, + block_table, + flat_lens, + out, + weighted_q.stride(0), + cache.stride(0), + block_table.stride(0), + out.stride(0), + cache.shape[1], + head_dim=_INDEX_HEAD_DIM, + BLOCK_N=_LOGITS_BLOCK_N, + NUM_SPLITS=splits, + num_warps=4, + ) + # Positions >= seq_len are left uninitialised, exactly as the + # single-chunk gather path leaves them; every consumer masks by + # seq_lens. + return out + + block_n = 16 gathered_k = torch.empty( - (weighted_q.shape[0], max_seq_len, _INDEX_HEAD_DIM), + (total_rows, max_seq_len, _INDEX_HEAD_DIM), dtype=torch.float16, device=q.device, ) - block_n = 16 - _dequant_paged_index_k_kernel[ - (weighted_q.shape[0], triton.cdiv(max_seq_len, block_n)) - ]( + _dequant_paged_index_k_kernel[(total_rows, triton.cdiv(max_seq_len, block_n))]( cache, block_table, flat_lens, gathered_k, cache.stride(0), - cache.stride(1), block_table.stride(0), gathered_k.stride(0), gathered_k.stride(1), 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 ac88ab5d3d..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 @@ -5,9 +5,13 @@ import torch +import vllm.envs as envs from vllm.models.deepseek_v4.common.ops.cache_utils import ( quantize_and_insert_k_cache, ) +from vllm.models.deepseek_v4.common.ops.fp8_software import ( + fp32_to_fp8_e4m3fn_bits, +) from vllm.triton_utils import tl, triton _HEAD_DIM = 512 @@ -93,6 +97,111 @@ def _sm70_qnorm_rope_kernel( ) +@triton.jit +def _sm70_qnorm_rope_parallel_kv_insert_kernel( + q_ptr, + kv_ptr, + cache_ptr, + slot_mapping_ptr, + positions_ptr, + cos_sin_cache_ptr, + cache_stride0, + eps: tl.constexpr, + num_tokens, + num_heads: tl.constexpr, + cache_block_size: tl.constexpr, + HEAD_DIM: tl.constexpr, + ROPE_DIM: tl.constexpr, + NOPE_DIM: tl.constexpr, + HALF_ROPE: tl.constexpr, + TOKEN_DATA_SIZE: tl.constexpr, +): + """Fuse TP4 Q transform with eight exact packed-KV writer CTAs.""" + token_idx = tl.program_id(0) + worker_idx = tl.program_id(1) + if token_idx >= num_tokens: + return + + if worker_idx < num_heads: + q_position = tl.load(positions_ptr + token_idx).to(tl.int64) + q_pair_idx = tl.arange(0, HALF_ROPE) + q_cos = tl.load(cos_sin_cache_ptr + q_position * ROPE_DIM + q_pair_idx).to( + tl.float32 + ) + q_sin = tl.load( + cos_sin_cache_ptr + q_position * ROPE_DIM + HALF_ROPE + q_pair_idx + ).to(tl.float32) + q_base = q_ptr + token_idx * num_heads * HEAD_DIM + worker_idx * HEAD_DIM + q_offsets = tl.arange(0, HEAD_DIM) + q_values = tl.load(q_base + q_offsets).to(tl.float32) + q_rrms = tl.rsqrt(tl.sum(q_values * q_values, axis=0) / HEAD_DIM + eps) + q_values *= q_rrms + tl.store( + q_base + q_offsets, + q_values.to(q_ptr.type.element_ty), + mask=q_offsets < NOPE_DIM, + ) + + q_even_offsets = NOPE_DIM + q_pair_idx * 2 + q_odd_offsets = q_even_offsets + 1 + q_even = tl.load(q_base + q_even_offsets).to(tl.float32) * q_rrms + q_odd = tl.load(q_base + q_odd_offsets).to(tl.float32) * q_rrms + tl.store( + q_base + q_even_offsets, + (q_even * q_cos - q_odd * q_sin).to(q_ptr.type.element_ty), + ) + tl.store( + q_base + q_odd_offsets, + (q_even * q_sin + q_odd * q_cos).to(q_ptr.type.element_ty), + ) + else: + part_idx = worker_idx - num_heads + slot_idx = tl.load(slot_mapping_ptr + token_idx) + if slot_idx < 0: + return + + block_idx = slot_idx // cache_block_size + pos_in_block = slot_idx % cache_block_size + cache_block = cache_ptr + block_idx.to(tl.int64) * cache_stride0 + token_data = cache_block + pos_in_block * TOKEN_DATA_SIZE + token_scales = ( + cache_block + cache_block_size * TOKEN_DATA_SIZE + pos_in_block * 8 + ) + kv_row = kv_ptr + token_idx * HEAD_DIM + kv_offsets = tl.arange(0, 64) + if part_idx < 7: + input_offsets = part_idx * 64 + kv_offsets + kv_values = tl.load(kv_row + input_offsets) + block_max = tl.maximum(tl.max(tl.abs(kv_values), axis=0), 1.0e-4) + exponent = tl.ceil(tl.log2(block_max / 448.0)) + scale = tl.exp2(exponent) + scaled = tl.clamp(kv_values / scale, -448.0, 448.0) + packed = fp32_to_fp8_e4m3fn_bits(scaled.to(tl.float32)) + tl.store(token_data + input_offsets, packed) + encoded = tl.maximum(tl.minimum(exponent + 127.0, 255.0), 0.0) + tl.store(token_scales + part_idx, encoded.to(tl.uint8)) + else: + tl.store(token_scales + 7, tl.zeros((), dtype=tl.uint8)) + kv_position = tl.load(positions_ptr + token_idx).to(tl.int64) + kv_pair_idx = tl.arange(0, HALF_ROPE) + kv_cos = tl.load( + cos_sin_cache_ptr + kv_position * ROPE_DIM + kv_pair_idx + ).to(tl.float32) + kv_sin = tl.load( + cos_sin_cache_ptr + kv_position * ROPE_DIM + HALF_ROPE + kv_pair_idx + ).to(tl.float32) + kv_even = tl.load(kv_row + NOPE_DIM + kv_pair_idx * 2).to(tl.float32) + kv_odd = tl.load(kv_row + NOPE_DIM + kv_pair_idx * 2 + 1).to(tl.float32) + rotated_even = (kv_even * kv_cos - kv_odd * kv_sin).to(tl.float16) + rotated_odd = (kv_even * kv_sin + kv_odd * kv_cos).to(tl.float16) + bf16_out = (token_data + NOPE_DIM).to(tl.pointer_type(tl.bfloat16)) + tl.store(bf16_out + kv_pair_idx * 2, rotated_even.to(tl.bfloat16)) + tl.store( + bf16_out + kv_pair_idx * 2 + 1, + rotated_odd.to(tl.bfloat16), + ) + + def sm70_qnorm_rope_kv_fp8_insert( q: torch.Tensor, kv: torch.Tensor, @@ -112,6 +221,29 @@ def sm70_qnorm_rope_kv_fp8_insert( assert q.is_contiguous() and kv.is_contiguous() num_tokens, num_heads, _ = q.shape + if envs.VLLM_SM70_DSV4_QNORM_KV_FUSED_TP4 and num_tokens == 1 and num_heads == 16: + cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1) + _sm70_qnorm_rope_parallel_kv_insert_kernel[(num_tokens, num_heads + 8)]( + q, + kv, + cache_2d, + slot_mapping, + positions, + cos_sin_cache, + cache_2d.stride(0), + eps, + num_tokens, + num_heads=num_heads, + cache_block_size=block_size, + HEAD_DIM=_HEAD_DIM, + ROPE_DIM=_ROPE_DIM, + NOPE_DIM=_NOPE_DIM, + HALF_ROPE=_HALF_ROPE, + TOKEN_DATA_SIZE=576, + num_warps=4, + ) + return q + kv_roped = torch.empty_like(kv) _sm70_qnorm_rope_kernel[(num_tokens, num_heads + 1)]( q, diff --git a/vllm/models/deepseek_v4/sm70/sparse.py b/vllm/models/deepseek_v4/sm70/sparse.py index f507372031..1eec99a606 100644 --- a/vllm/models/deepseek_v4/sm70/sparse.py +++ b/vllm/models/deepseek_v4/sm70/sparse.py @@ -40,6 +40,11 @@ logger = init_logger(__name__) +def _qk_dsplit_block_h(num_heads: int) -> int: + """Use one CTA head group for the exact TP4 decode shape.""" + return 16 if num_heads == 16 else 8 + + class DeepseekV4SM70SparseBackend(DeepseekV4FlashMLASparseBackend): supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16] @@ -258,6 +263,7 @@ def _forward_decode( sm70_sparse_attention_paged_fp8_splitk_qk_dsplit( partial_qk=partial_qk, partial_probs=partial_probs, + stage1_block_h=_qk_dsplit_block_h(q.shape[1]), **common_kwargs, ) else: diff --git a/vllm/tool_parsers/qwen3coder_tool_parser.py b/vllm/tool_parsers/qwen3coder_tool_parser.py index 2b21f8f6ae..7a3a63ad53 100644 --- a/vllm/tool_parsers/qwen3coder_tool_parser.py +++ b/vllm/tool_parsers/qwen3coder_tool_parser.py @@ -358,10 +358,7 @@ def extract_tool_calls_streaming( if pos != -1 ] tool_start = min(tool_start_candidates) if tool_start_candidates else -1 - if ( - self.tool_call_start_token_id in delta_token_ids - or tool_start != -1 - ): + if self.tool_call_start_token_id in delta_token_ids or tool_start != -1: self.is_tool_call_started = True # Return any content before the tool call if tool_start > len(previous_text): diff --git a/vllm/transformers_utils/processors/deepseek_ocr.py b/vllm/transformers_utils/processors/deepseek_ocr.py index 68a2b1aaaa..62972ad242 100644 --- a/vllm/transformers_utils/processors/deepseek_ocr.py +++ b/vllm/transformers_utils/processors/deepseek_ocr.py @@ -3,7 +3,7 @@ # adapted from https://github.com/deepseek-ai/DeepSeek-OCR/blob/main/DeepSeek-OCR-master/DeepSeek-OCR-vllm/process/image_process.py # and https://github.com/deepseek-ai/DeepSeek-OCR-2/blob/main/DeepSeek-OCR2-master/DeepSeek-OCR2-vllm/process/image_process.py import math -from typing import Literal +from typing import Literal, TypedDict import torch import torchvision.transforms as T @@ -11,18 +11,30 @@ from transformers import BatchFeature, LlamaTokenizerFast from transformers.processing_utils import ProcessorMixin -# TODO(Isotr0py): change modes for variants -# see: https://github.com/deepseek-ai/DeepSeek-OCR/blob/8cf003d38821fa1b19c73da3bd1b0dc262ea8136/DeepSeek-OCR-master/DeepSeek-OCR-vllm/config.py#L1-L6 -# Tiny: base_size = 512, image_size = 512, crop_mode = False -# Small: base_size = 640, image_size = 640, crop_mode = False -# Base: base_size = 1024, image_size = 1024, crop_mode = False -# Large: base_size = 1280, image_size = 1280, crop_mode = False -# Gundam: base_size = 1024, image_size = 640, crop_mode = True + +# Official resolution modes, verbatim from the model's config.py: +# https://github.com/deepseek-ai/DeepSeek-OCR/blob/8cf003d38821fa1b19c73da3bd1b0dc262ea8136/DeepSeek-OCR-master/DeepSeek-OCR-vllm/config.py#L1-L6 +# Selectable per request via the ``image_mode`` mm_processor_kwarg; the +# default stays the Gundam configuration below (byte-compatible with the +# previous hardcoded constants). +class ResolutionMode(TypedDict): + base_size: int + image_size: int + crop_mode: bool + + +RESOLUTION_MODES: dict[str, ResolutionMode] = { + "tiny": {"base_size": 512, "image_size": 512, "crop_mode": False}, + "small": {"base_size": 640, "image_size": 640, "crop_mode": False}, + "base": {"base_size": 1024, "image_size": 1024, "crop_mode": False}, + "large": {"base_size": 1280, "image_size": 1280, "crop_mode": False}, + "gundam": {"base_size": 1024, "image_size": 640, "crop_mode": True}, +} + BASE_SIZE = 1024 IMAGE_SIZE = 640 CROP_MODE = True -# TODO(Isotr0py): Expose as mm_kwargs MIN_CROPS = 2 MAX_CROPS = 6 # max:9; If your GPU memory is small, it is recommended to set it to 6. @@ -117,6 +129,53 @@ def dynamic_preprocess( return processed_images, target_aspect_ratio +def count_image_tokens_for( + *, + image_width: int, + image_height: int, + base_size: int, + image_size: int, + cropping: bool, + min_crops: int = MIN_CROPS, + max_crops: int = MAX_CROPS, + strategy: Literal["v1", "v2"] = "v1", + patch_size: int = 16, + downsample_ratio: int = 4, +) -> int: + """Placeholder-token count mirroring ``tokenize_with_images``.""" + if image_width <= image_size and image_height <= image_size: + crop_ratio = (1, 1) + elif cropping: + crop_ratio = count_tiles( + image_width, + image_height, + min_num=min_crops, + max_num=max_crops, + image_size=image_size, + ) + else: + crop_ratio = (1, 1) + num_width_tiles, num_height_tiles = crop_ratio + + num_queries = math.ceil((image_size // patch_size) / downsample_ratio) + num_queries_base = math.ceil((base_size // patch_size) / downsample_ratio) + + num_tokens_base = ( + (num_queries_base * (num_queries_base + 1)) + if strategy == "v1" + else num_queries_base * num_queries_base + ) + total = num_tokens_base + 1 + if num_width_tiles > 1 or num_height_tiles > 1: + num_tokens_per_row = ( + num_queries * num_width_tiles + 1 + if strategy == "v1" + else num_queries * num_width_tiles + ) + total += num_tokens_per_row * (num_queries * num_height_tiles) + return total + + class ImageTransform: def __init__( self, @@ -160,22 +219,63 @@ def __init__( ignore_id: int = -100, image_size: int = IMAGE_SIZE, base_size: int = BASE_SIZE, + crop_mode: bool = CROP_MODE, + image_mode: str | None = None, + min_crops: int = MIN_CROPS, + max_crops: int = MAX_CROPS, strategy: Literal["v1", "v2"] = "v1", **kwargs, ): + if image_mode is not None: + # The named mode is authoritative over any individually passed + # base_size/image_size/crop_mode (the vLLM processing info always + # forwards the defaults, so a raise-on-conflict would never let a + # mode through). + if not isinstance(image_mode, str): + raise ValueError( + f"image_mode must be a string, got {type(image_mode).__name__}." + ) + try: + mode = RESOLUTION_MODES[image_mode] + except KeyError: + raise ValueError( + f"Unknown image_mode {image_mode!r}; valid modes: " + f"{sorted(RESOLUTION_MODES)}" + ) from None + base_size = mode["base_size"] + image_size = mode["image_size"] + crop_mode = mode["crop_mode"] + if not ( + type(min_crops) is int + and type(max_crops) is int + and 1 <= min_crops <= max_crops <= 9 + ): + raise ValueError( + f"Invalid crop bounds: min_crops={min_crops!r}, " + f"max_crops={max_crops!r} (need 1 <= min <= max <= 9)." + ) + if type(crop_mode) is not bool: + raise ValueError(f"crop_mode must be a bool, got {crop_mode!r}.") + if patch_size <= 0 or downsample_ratio <= 0: + raise ValueError("patch_size and downsample_ratio must be positive.") self.image_size = image_size self.base_size = base_size + self.crop_mode = crop_mode + self.image_mode = image_mode + self.min_crops = min_crops + self.max_crops = max_crops # image token calculation strategy for # Deepseek-OCR and Deepseek-OCR-2 self.strategy = strategy - assert strategy in ["v1", "v2"], "Only 'v1' and 'v2' strategies are supported." + if strategy not in ("v1", "v2"): + raise ValueError("Only 'v1' and 'v2' strategies are supported.") - self.patch_size = 16 + self.patch_size = patch_size self.image_mean = image_mean self.image_std = image_std self.normalize = normalize - self.downsample_ratio = 4 + self.downsample_ratio = downsample_ratio self.image_transform = ImageTransform( mean=image_mean, std=image_std, normalize=normalize @@ -203,6 +303,28 @@ def __init__( **kwargs, ) + def count_image_tokens(self, *, image_width: int, image_height: int) -> int: + """Number of placeholder tokens for one image. + + This mirrors ``tokenize_with_images`` exactly (same instance + parameters and same crop arithmetic) so the prompt-update counting + can never diverge from what processing + produces — a mismatch here is the "N multimodal tokens vs M + placeholders" crash class. + """ + return count_image_tokens_for( + image_width=image_width, + image_height=image_height, + base_size=self.base_size, + image_size=self.image_size, + cropping=self.crop_mode, + min_crops=self.min_crops, + max_crops=self.max_crops, + strategy=self.strategy, + patch_size=self.patch_size, + downsample_ratio=self.downsample_ratio, + ) + @property def bos_id(self): return self.tokenizer.bos_token_id @@ -230,7 +352,7 @@ def process_one( self, prompt: str, images: list[Image.Image], - crop_mode: bool = CROP_MODE, + crop_mode: bool | None = None, ): """ @@ -253,6 +375,8 @@ def process_one( ) sft_format = prompt + if crop_mode is None: + crop_mode = self.crop_mode ( input_ids, @@ -288,7 +412,7 @@ def __call__( *, prompt: str, images: list[Image.Image], - crop_mode: bool = CROP_MODE, + crop_mode: bool | None = None, **kwargs, ): prepare = self.process_one( @@ -332,7 +456,10 @@ def tokenize_with_images( crop_ratio = [1, 1] elif cropping: images_crop_raw, crop_ratio = dynamic_preprocess( - image, image_size=self.image_size + image, + min_num=self.min_crops, + max_num=self.max_crops, + image_size=self.image_size, ) else: crop_ratio = [1, 1] diff --git a/vllm/utils/cache.py b/vllm/utils/cache.py index 4338983f90..5f7647dce6 100644 --- a/vllm/utils/cache.py +++ b/vllm/utils/cache.py @@ -118,10 +118,8 @@ def stat(self, *, delta: bool = False) -> CacheInfo: return info def touch(self, key: _K) -> None: - try: + if key in self: self._LRUCache__order.move_to_end(key) # type: ignore - except KeyError: - self._LRUCache__order[key] = None # type: ignore @overload def get(self, key: _K, /) -> _V | None: ... diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index b1cb946d33..4425fae48b 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -388,6 +388,9 @@ class CommonAttentionMetadata: causal: bool = True + cudagraph_graph_variant: int | None = None + """Optional semantic variant for CUDA graph capture dummy runs.""" + # Needed by FastPrefillAttentionBuilder logits_indices_padded: torch.Tensor | None = None num_logits_indices: int | None = None @@ -416,6 +419,14 @@ class CommonAttentionMetadata: decode rows (assumes every draft was accepted). Not safe for kernels that need exact per-row context lengths on decode rows.""" + prefix_anchor_lens: torch.Tensor | None = None + """(batch_size,) per-request prefix length (prompt token count) for + prefix-anchored sliding-window attention. Tokens with logical index below + this stay globally visible; later (generated) tokens additionally see a + fixed sliding window. None disables the mechanism. The attention backend + copies this into its own persistent buffer and reads the window size from + the cache group's ``PrefixAnchoredSWASpec``.""" + # WARNING: Deprecated fields. Will be removed in a future release (v0.15.0) _seq_lens_cpu: torch.Tensor | None = None _num_computed_tokens_cpu: torch.Tensor | None = None @@ -516,6 +527,7 @@ def unpadded( block_table_tensor=self.block_table_tensor[:num_actual_reqs], slot_mapping=self.slot_mapping[:num_actual_tokens], causal=self.causal, + cudagraph_graph_variant=self.cudagraph_graph_variant, logits_indices_padded=self.logits_indices_padded, num_logits_indices=self.num_logits_indices, encoder_seq_lens=maybe_slice_reqs(self.encoder_seq_lens), @@ -523,6 +535,7 @@ def unpadded( dcp_local_seq_lens=maybe_slice_reqs(self.dcp_local_seq_lens), dcp_local_seq_lens_cpu=maybe_slice_reqs(self.dcp_local_seq_lens_cpu), is_prefilling=maybe_slice_reqs(self.is_prefilling), + prefix_anchor_lens=maybe_slice_reqs(self.prefix_anchor_lens), ) 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 1ecbbdb308..e6a6f15e65 100644 --- a/vllm/v1/attention/backends/flash_attn_v100.py +++ b/vllm/v1/attention/backends/flash_attn_v100.py @@ -16,14 +16,17 @@ import os import time from collections.abc import Callable +from dataclasses import dataclass from functools import partial from typing import cast import torch import vllm.envs as envs +from vllm.forward_context import CUDAGRAPH_VARIANT_LONG_CONTEXT from vllm.logger import init_logger from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton from vllm.v1.attention.backend import AttentionCGSupport, AttentionType from vllm.v1.attention.backends.triton_attn import ( TritonAttentionBackend, @@ -31,10 +34,396 @@ TritonAttentionMetadata, TritonAttentionMetadataBuilder, ) +from vllm.v1.kv_cache_interface import PrefixAnchoredSWASpec +from vllm.v1.worker.gpu.spec_decode import uses_dflash_selector_engine logger = init_logger(__name__) +@triton.jit +def _sm70_prepare_smallq_decode_metadata_kernel( + out_block_table_ptr, + out_seq_lens_ptr, + out_query_start_loc_ptr, + block_table_ptr, + seq_lens_ptr, + query_start_loc_ptr, + block_table_stride, + out_block_table_stride, + num_reqs, + num_query_tokens, + real_num_query_tokens, + block_cols, + REQ_BLOCK: tl.constexpr, + BLOCK_COLS: tl.constexpr, +): + token_idx = tl.program_id(0) + + # Find the request that owns this flattened query token. Equal trailing + # boundaries are CUDA-graph padding; clamping maps them to the final padded + # request, matching repeat_query_lens[-1] += padding_tokens. + req_offsets = tl.arange(0, REQ_BLOCK) + req_mask = req_offsets < num_reqs + query_ends = tl.load( + query_start_loc_ptr + req_offsets + 1, + mask=req_mask, + other=0x7FFFFFFF, + ) + req_idx = tl.sum((token_idx >= query_ends).to(tl.int32), axis=0) + req_idx = tl.minimum(req_idx, num_reqs - 1) + + query_start = tl.load(query_start_loc_ptr + req_idx) + query_end = tl.load(query_start_loc_ptr + req_idx + 1) + query_len = query_end - query_start + seq_len = tl.load(seq_lens_ptr + req_idx) + effective_seq_len = tl.maximum(seq_len, query_len) + decode_seq_len = effective_seq_len - query_len + token_idx - query_start + 1 + is_padding = token_idx >= real_num_query_tokens + tl.store( + out_seq_lens_ptr + token_idx, + tl.where(is_padding, 0, decode_seq_len), + ) + + cols = tl.arange(0, BLOCK_COLS) + col_mask = cols < block_cols + block_ids = tl.load( + block_table_ptr + req_idx * block_table_stride + cols, + mask=col_mask, + other=0, + ) + block_ids = tl.maximum(block_ids, 0) + block_ids = tl.where(is_padding, 0, block_ids) + tl.store( + out_block_table_ptr + token_idx * out_block_table_stride + cols, + block_ids, + mask=col_mask, + ) + + # The same launch also refreshes the graph-stable query boundaries. + if token_idx == 0: + boundary_offsets = tl.arange(0, REQ_BLOCK) + boundary_mask = boundary_offsets < num_reqs + 1 + boundaries = tl.load( + query_start_loc_ptr + boundary_offsets, + mask=boundary_mask, + other=0, + ) + tl.store( + out_query_start_loc_ptr + boundary_offsets, + boundaries, + mask=boundary_mask, + ) + + +def _sm70_prepare_smallq_decode_metadata( + out_block_table: torch.Tensor, + out_seq_lens: torch.Tensor, + out_query_start_loc: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + query_start_loc: torch.Tensor, + *, + num_reqs: int, + num_query_tokens: int, + real_num_query_tokens: int, +) -> None: + """Materialize persistent Flash-V100 verifier metadata in one launch.""" + if num_reqs <= 0 or num_query_tokens <= 0: + raise ValueError("small-query metadata requires positive request/token counts") + block_cols = int(block_table.shape[1]) + if block_cols <= 0: + raise ValueError("small-query block table must have at least one column") + if out_block_table.shape[0] < num_query_tokens: + raise ValueError("small-query output block table is too small") + if out_seq_lens.numel() < num_query_tokens: + raise ValueError("small-query output sequence buffer is too small") + if out_query_start_loc.numel() < num_reqs + 1: + raise ValueError("small-query output boundary buffer is too small") + + _sm70_prepare_smallq_decode_metadata_kernel[(num_query_tokens,)]( + out_block_table, + out_seq_lens, + out_query_start_loc, + block_table, + seq_lens, + query_start_loc, + block_table.stride(0), + out_block_table.stride(0), + num_reqs, + num_query_tokens, + real_num_query_tokens, + block_cols, + REQ_BLOCK=triton.next_power_of_2(num_reqs + 1), + BLOCK_COLS=triton.next_power_of_2(block_cols), + num_warps=1, + ) + + +@triton.jit +def _load_sm70_smallq_i32_ptr(ptrs, group_id): + ptr = tl.load(ptrs + group_id) + return tl.cast(ptr, tl.pointer_type(tl.int32)) + + +@triton.jit +def _sm70_prepare_grouped_smallq_decode_metadata_kernel( + out_block_table_ptrs, + out_seq_lens_ptrs, + out_query_start_loc_ptrs, + block_table_ptrs, + block_table_strides, + out_block_table_strides, + block_col_counts, + seq_lens_ptr, + query_start_loc_ptr, + num_reqs, + real_num_query_tokens, + REQ_BLOCK: tl.constexpr, + BLOCK_COLS: tl.constexpr, +): + """Materialize every target full-attention group's verifier metadata.""" + group_id = tl.program_id(0) + token_idx = tl.program_id(1) + out_block_table = _load_sm70_smallq_i32_ptr(out_block_table_ptrs, group_id) + out_seq_lens = _load_sm70_smallq_i32_ptr(out_seq_lens_ptrs, group_id) + out_query_start_loc = _load_sm70_smallq_i32_ptr(out_query_start_loc_ptrs, group_id) + block_table = _load_sm70_smallq_i32_ptr(block_table_ptrs, group_id) + block_table_stride = tl.load(block_table_strides + group_id) + out_block_table_stride = tl.load(out_block_table_strides + group_id) + block_cols = tl.load(block_col_counts + group_id) + + req_offsets = tl.arange(0, REQ_BLOCK) + req_mask = req_offsets < num_reqs + query_ends = tl.load( + query_start_loc_ptr + req_offsets + 1, + mask=req_mask, + other=0x7FFFFFFF, + ) + req_idx = tl.sum((token_idx >= query_ends).to(tl.int32), axis=0) + req_idx = tl.minimum(req_idx, num_reqs - 1) + + query_start = tl.load(query_start_loc_ptr + req_idx) + query_end = tl.load(query_start_loc_ptr + req_idx + 1) + query_len = query_end - query_start + seq_len = tl.load(seq_lens_ptr + req_idx) + effective_seq_len = tl.maximum(seq_len, query_len) + decode_seq_len = effective_seq_len - query_len + token_idx - query_start + 1 + is_padding = token_idx >= real_num_query_tokens + tl.store( + out_seq_lens + token_idx, + tl.where(is_padding, 0, decode_seq_len), + ) + + cols = tl.arange(0, BLOCK_COLS) + col_mask = cols < block_cols + block_ids = tl.load( + block_table + req_idx * block_table_stride + cols, + mask=col_mask, + other=0, + ) + block_ids = tl.maximum(block_ids, 0) + block_ids = tl.where(is_padding, 0, block_ids) + tl.store( + out_block_table + token_idx * out_block_table_stride + cols, + block_ids, + mask=col_mask, + ) + + if token_idx == 0: + boundary_offsets = tl.arange(0, REQ_BLOCK) + boundary_mask = boundary_offsets < num_reqs + 1 + boundaries = tl.load( + query_start_loc_ptr + boundary_offsets, + mask=boundary_mask, + other=0, + ) + tl.store( + out_query_start_loc + boundary_offsets, + boundaries, + mask=boundary_mask, + ) + + +@dataclass +class DFlash2SmallQGroupDescriptor: + """Persistent pointer tables for grouped Flash-V100 verifier metadata.""" + + key: tuple[object, ...] + block_table_ptrs: torch.Tensor + out_block_table_ptrs: torch.Tensor + out_seq_lens_ptrs: torch.Tensor + out_query_start_loc_ptrs: torch.Tensor + block_table_strides: torch.Tensor + out_block_table_strides: torch.Tensor + block_col_counts: torch.Tensor + + +@dataclass(frozen=True) +class DFlash2SmallQPreparedMetadata: + """Graph-stable buffers already refreshed by the grouped launch.""" + + builder_id: int + num_reqs: int + num_query_tokens: int + max_seq_len_hint: int + workspace_seq_capacity_hint: int + partition_size_hint: int | None = None + + +def _sm70_prepare_grouped_smallq_decode_metadata( + out_block_tables: list[torch.Tensor], + out_seq_lens: list[torch.Tensor], + out_query_start_locs: list[torch.Tensor], + block_tables: list[torch.Tensor], + seq_lens: torch.Tensor, + query_start_loc: torch.Tensor, + *, + num_reqs: int, + num_query_tokens: int, + real_num_query_tokens: int, + descriptor: DFlash2SmallQGroupDescriptor | None = None, +) -> DFlash2SmallQGroupDescriptor: + """Refresh N full-attention cache groups in one Triton launch.""" + num_groups = len(block_tables) + if num_groups <= 0: + raise ValueError("grouped small-query metadata requires at least one group") + if not ( + len(out_block_tables) + == len(out_seq_lens) + == len(out_query_start_locs) + == num_groups + ): + raise ValueError("grouped small-query metadata lists must have equal lengths") + if num_reqs <= 0 or num_query_tokens <= 0: + raise ValueError("grouped small-query metadata requires positive sizes") + if not 0 <= real_num_query_tokens <= num_query_tokens: + raise ValueError("real query token count exceeds grouped launch size") + + block_col_counts = [int(table.shape[1]) for table in block_tables] + if any(cols <= 0 for cols in block_col_counts): + raise ValueError("grouped small-query block tables cannot be empty") + max_block_cols = max(block_col_counts) + device = block_tables[0].device + if ( + seq_lens.device != device + or query_start_loc.device != device + or seq_lens.dtype != torch.int32 + or query_start_loc.dtype != torch.int32 + or not seq_lens.is_contiguous() + or not query_start_loc.is_contiguous() + or seq_lens.numel() < num_reqs + or query_start_loc.numel() < num_reqs + 1 + ): + raise ValueError("grouped small-query input metadata contract mismatch") + key: tuple[object, ...] = ( + device.type, + device.index, + tuple(block_col_counts), + tuple(table.data_ptr() for table in block_tables), + tuple(table.data_ptr() for table in out_block_tables), + tuple(tensor.data_ptr() for tensor in out_seq_lens), + tuple(tensor.data_ptr() for tensor in out_query_start_locs), + tuple(table.stride(0) for table in block_tables), + tuple(table.stride(0) for table in out_block_tables), + ) + if descriptor is None or descriptor.key != key: + for group, ( + block_table, + out_block_table, + out_seq_len, + out_query_start, + ) in enumerate( + zip( + block_tables, + out_block_tables, + out_seq_lens, + out_query_start_locs, + strict=True, + ) + ): + block_cols = block_col_counts[group] + if ( + block_table.device != device + or out_block_table.device != device + or out_seq_len.device != device + or out_query_start.device != device + or block_table.dtype != torch.int32 + or out_block_table.dtype != torch.int32 + or out_seq_len.dtype != torch.int32 + or out_query_start.dtype != torch.int32 + or block_table.ndim != 2 + or out_block_table.ndim != 2 + or block_table.shape[0] < num_reqs + or out_block_table.shape[0] < num_query_tokens + or out_block_table.shape[1] < block_cols + or block_table.stride(1) != 1 + or out_block_table.stride(1) != 1 + or not out_seq_len.is_contiguous() + or not out_query_start.is_contiguous() + or out_seq_len.numel() < num_query_tokens + or out_query_start.numel() < num_reqs + 1 + ): + raise ValueError( + f"grouped small-query metadata contract mismatch for group {group}" + ) + descriptor = DFlash2SmallQGroupDescriptor( + key=key, + block_table_ptrs=torch.tensor( + [table.data_ptr() for table in block_tables], + dtype=torch.uint64, + device=device, + ), + out_block_table_ptrs=torch.tensor( + [table.data_ptr() for table in out_block_tables], + dtype=torch.uint64, + device=device, + ), + out_seq_lens_ptrs=torch.tensor( + [tensor.data_ptr() for tensor in out_seq_lens], + dtype=torch.uint64, + device=device, + ), + out_query_start_loc_ptrs=torch.tensor( + [tensor.data_ptr() for tensor in out_query_start_locs], + dtype=torch.uint64, + device=device, + ), + block_table_strides=torch.tensor( + [table.stride(0) for table in block_tables], + dtype=torch.int64, + device=device, + ), + out_block_table_strides=torch.tensor( + [table.stride(0) for table in out_block_tables], + dtype=torch.int64, + device=device, + ), + block_col_counts=torch.tensor( + block_col_counts, + dtype=torch.int32, + device=device, + ), + ) + + _sm70_prepare_grouped_smallq_decode_metadata_kernel[(num_groups, num_query_tokens)]( + descriptor.out_block_table_ptrs, + descriptor.out_seq_lens_ptrs, + descriptor.out_query_start_loc_ptrs, + descriptor.block_table_ptrs, + descriptor.block_table_strides, + descriptor.out_block_table_strides, + descriptor.block_col_counts, + seq_lens, + query_start_loc, + num_reqs, + real_num_query_tokens, + REQ_BLOCK=triton.next_power_of_2(num_reqs + 1), + BLOCK_COLS=triton.next_power_of_2(max_block_cols), + num_warps=1, + ) + return descriptor + + class FlashAttnV100Metadata(TritonAttentionMetadata): """Static view of Flash-V100 fields attached to Triton metadata.""" @@ -43,7 +432,10 @@ class FlashAttnV100Metadata(TritonAttentionMetadata): causal: bool max_model_len: int flash_v100_cudagraph_capture: bool + flash_v100_batch_context_routing: bool flash_v100_contig_dense_cache: dict[tuple[int, int, int, int, int], int] + prefix_anchor_lens: torch.Tensor | None = None + decode_sliding_window: int | None = None flash_v100_decode_max_seq_len_hint: int | None flash_v100_decode_workspace_seq_capacity_hint: int | None flash_v100_static_decode_seq_hint: int | None @@ -59,6 +451,7 @@ class FlashAttnV100Metadata(TritonAttentionMetadata): smallq_decode_max_seq_len_hint: int | None smallq_decode_workspace_seq_capacity_hint: int | None smallq_decode_partition_size_hint: int | None + is_dflash_selector_target: bool def _as_flash_v100_metadata( @@ -69,6 +462,25 @@ def _as_flash_v100_metadata( return cast(FlashAttnV100Metadata, attn_metadata) +def _batch_context_routing_for_graph_variant( + routing_enabled: bool, + graph_variant: int | None, +) -> bool: + if not routing_enabled: + return False + if graph_variant is None: + # Eager execution can route directly from the live batch and context. + return True + return graph_variant == CUDAGRAPH_VARIANT_LONG_CONTEXT + + +def _batch_context_routing_cache_dtype_supported(cache_dtype: str | None) -> bool: + """Admit the exact FP8 XQA formats implemented by Flash-V100.""" + return cache_dtype == "fp8_e5m2" or ( + cache_dtype in ("fp8", "fp8_e4m3") and envs.VLLM_FLASH_V100_E4M3_BATCH_XQA + ) + + def _sm70_profile_trace(message: str, *args: object) -> None: if envs.VLLM_SM70_PROFILE_TRACE: if args: @@ -82,12 +494,16 @@ def _sm70_profile_trace(message: str, *args: object) -> None: _flash_attn_decode_paged = None _flash_attn_decode_paged_xqa = None _flash_attn_decode_paged_wmma = None +_flash_attn_grouped_verify_paged = None +_flash_attn_grouped_verify_checked = False _flash_attn_prefill_paged = None _flash_attn_prefill_paged_bhmd = None _flash_attn_prefill_paged_bfla = None _flash_attn_prefill_paged_splitkv = None _sm70_splitd_d256_ops = None _sm70_splitd_d256_ops_checked = False +_sm70_d256_gqa_architecture_op = None +_sm70_d256_gqa_architecture_op_checked = False _sm70_fa2_cu_seqlens_cache: dict[ tuple[int, int, int, int], tuple[torch.Tensor, torch.Tensor] ] = {} @@ -101,6 +517,7 @@ def _sm70_profile_trace(message: str, *args: object) -> None: _warned_decode_strict_fallback = False _warned_prefill_gather_oom = False _warned_prefill_dense_splitkv3_oom = False +_warned_prefill_d256_gqa_architecture_oom = False _logged_prefill_flash = False _logged_prefill_prefix_flash = False _logged_prefill_prefix_contig_dense = False @@ -109,8 +526,11 @@ def _sm70_profile_trace(message: str, *args: object) -> None: _logged_prefill_paged_cache = False _logged_prefill_smallq_decode = False _logged_prefill_smallq_decode_xqa = False +_logged_prefill_smallq_grouped_verify = False +_logged_prefill_smallq_grouped_verify_gate = False _logged_prefill_fa2_d256 = False _logged_prefill_dense_splitkv3 = False +_logged_prefill_d256_gqa_architecture = False _logged_prefill_triton_safe = False _logged_decode_flash = False _logged_decode_dense_reference = False @@ -316,7 +736,7 @@ def _g6_aligned_page_partition_size_hint( ): return None if ( - kv_cache_dtype in ("auto", "bfloat16") + kv_cache_dtype in ("auto", "float16", "bfloat16") and key_cache.dtype == torch.float16 and key_cache.shape[1] == 784 ): @@ -324,6 +744,12 @@ def _g6_aligned_page_partition_size_hint( # selects between them from device seq_lens. Plan the p256 workspace # envelope once. return 256 + if kv_cache_dtype in ("fp8", "fp8_e4m3") and key_cache.dtype == torch.uint8: + # Qwen3.8 TP4 has G6/D256 attention and checkpoint-provided E4M3 + # scales. The p64 route wins the accepted 1K-2K operator sweep while + # preserving the p256 scalar route's E4M3 conversion within one fp16 + # output ULP. + return 64 if kv_cache_dtype == "fp8_e5m2" and key_cache.dtype == torch.uint8: # Plan the largest p256 workspace once. The extension selects p256 or # p1024 from device seq_lens, so CUDA graph replay keeps one @@ -491,6 +917,15 @@ def _decode_xqa_allowed_for_q_per_kv( return int(seq_hint) >= _decode_xqa_q4_min_seq_len() +def _e4m3_batch_xqa_allowed(query: torch.Tensor) -> bool: + """Gate the exact SM70 E4M3 G6 batched XQA route.""" + return ( + envs.VLLM_FLASH_V100_E4M3_BATCH_XQA + and 1 < query.shape[0] <= 16 + and query.shape[1:] == (6, 256) + ) + + def _same_storage(left: torch.Tensor, right: torch.Tensor) -> bool: return left.untyped_storage().data_ptr() == right.untyped_storage().data_ptr() @@ -785,7 +1220,12 @@ def _get_flash_ops(): _flash_attn_prefill_paged_splitkv = flash_attn_prefill_paged_splitkv except ImportError: _flash_attn_prefill_paged_splitkv = None - except ImportError: + except ImportError as exc: + logger.warning_once( + "Flash-V100 Python operators could not be imported (%s: %s).", + type(exc).__name__, + exc, + ) _flash_attn_func = None _flash_attn_bhmd_func = None _flash_attn_decode_paged = None @@ -808,6 +1248,23 @@ def _get_flash_ops(): ) +def _get_flash_grouped_verify_op(): + """Load the optional exact SM70 DFlash2 grouped verifier.""" + global _flash_attn_grouped_verify_paged + global _flash_attn_grouped_verify_checked + if _flash_attn_grouped_verify_checked: + return _flash_attn_grouped_verify_paged + + _flash_attn_grouped_verify_checked = True + try: + from flash_attn_v100 import flash_attn_grouped_verify_paged + + _flash_attn_grouped_verify_paged = flash_attn_grouped_verify_paged + except ImportError: + _flash_attn_grouped_verify_paged = None + return _flash_attn_grouped_verify_paged + + def _get_sm70_splitd_d256_ops(): """Load the exact SM70 Split-D dense and paged prefill operators.""" global _sm70_splitd_d256_ops @@ -842,6 +1299,35 @@ def _get_sm70_splitd_d256_ops(): return _sm70_splitd_d256_ops +def _get_sm70_d256_gqa_architecture_op(): + """Load the optional SM70 GQA long-prefill architecture operator.""" + global _sm70_d256_gqa_architecture_op + global _sm70_d256_gqa_architecture_op_checked + if _sm70_d256_gqa_architecture_op_checked: + return _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 + + _sm70_d256_gqa_architecture_op = getattr( + torch.ops._vllm_fa2_C, + "sm70_d256_gqa_architecture_fwd", + None, + ) + except (AttributeError, ImportError, RuntimeError) as exc: + _sm70_d256_gqa_architecture_op = None + if envs.VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL: + logger.warning_once( + "SM70 D256 GQA architecture operator is unavailable " + "(%s: %s); using the exact dense prefill kernel.", + type(exc).__name__, + exc, + ) + return _sm70_d256_gqa_architecture_op + + def _uniform_cu_seqlens( tensor: torch.Tensor, *, @@ -937,18 +1423,60 @@ def _should_use_prefill_dense_splitkv3( return ( envs.VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3 and splitkv3_op is not None - and query.shape == (1, 4096, 6, 256) + and ( + query.shape == (1, 4096, 6, 256) + or ( + envs.VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3_Q8000_EXPERIMENTAL + and query.shape == (1, 8000, 6, 256) + ) + ) and key.ndim == 4 and key.shape[0] == 1 and key.shape[1] == max_seqlen_k and key.shape[2:] == (1, 256) - and max_seqlen_q == 4096 + and max_seqlen_q == query.shape[1] and max_seqlen_k >= envs.VLLM_FLASH_V100_PREFILL_DENSE_SPLITKV3_MIN_KV and max_seqlen_k > max_seqlen_q and not _is_cuda_graph_capturing(query) ) +def _should_use_prefill_d256_gqa_architecture( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + *, + max_seqlen_q: int, + max_seqlen_k: int, + softmax_scale: float, + architecture_op: Callable[..., torch.Tensor] | None, +) -> bool: + """Gate the measured Q8000/KV40K..128K/Hq6/Hkv1/D256 family.""" + return ( + envs.VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL + and architecture_op is not None + and query.shape == (1, 8000, 6, 256) + and key.ndim == 4 + and key.shape[0] == 1 + and key.shape[2:] == (1, 256) + and value.shape == key.shape + and max_seqlen_q == 8000 + and max_seqlen_k == key.shape[1] + and 40000 <= max_seqlen_k <= 128000 + and max_seqlen_k % 8000 == 0 + and query.dtype == torch.float16 + and key.dtype == query.dtype + and value.dtype == query.dtype + and query.device == key.device + and query.device == value.device + and query.is_contiguous() + and key.is_contiguous() + and value.is_contiguous() + and abs(softmax_scale - 0.0625) <= 1.0e-8 + and not _is_cuda_graph_capturing(query) + ) + + def _try_sm70_fa2_d256_prefill( query: torch.Tensor, key: torch.Tensor, @@ -965,6 +1493,9 @@ def _try_sm70_fa2_d256_prefill( seqused_k: torch.Tensor | None = None, block_table: torch.Tensor | None = None, ) -> torch.Tensor | None: + global _logged_prefill_d256_gqa_architecture + global _warned_prefill_d256_gqa_architecture_oom + int32_max = torch.iinfo(torch.int32).max if not envs.VLLM_FLASH_V100_FA2_D256_PREFILL: return None @@ -1065,7 +1596,47 @@ def _try_sm70_fa2_d256_prefill( ) if splitd_eligible: splitd_out = out if out is not None else torch.empty_like(query) - if _should_use_prefill_dense_splitkv3( + architecture_op = ( + _get_sm70_d256_gqa_architecture_op() + if envs.VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL + else None + ) + if _should_use_prefill_d256_gqa_architecture( + query, + key, + value, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + softmax_scale=softmax_scale, + architecture_op=architecture_op, + ): + assert architecture_op is not None + try: + splitd_result = architecture_op( + query, + key, + value, + splitd_out, + softmax_scale, + True, + ) + except torch.OutOfMemoryError: + if not _warned_prefill_d256_gqa_architecture_oom: + logger.warning( + "Insufficient memory for the default-on " + "SM70 D256 GQA long-prefill architecture; " + "falling back to the exact dense kernel." + ) + _warned_prefill_d256_gqa_architecture_oom = True + if splitd_result is not None: + if not _logged_prefill_d256_gqa_architecture: + logger.info( + "FLASH_ATTN_V100 SM70 D256 GQA " + "8K-by-40K..128K architecture route active." + ) + _logged_prefill_d256_gqa_architecture = True + _record_route("prefill_dense_d256_gqa_arch_long") + if splitd_result is None and _should_use_prefill_dense_splitkv3( query, key, max_seqlen_q=max_seqlen_q, @@ -2461,13 +3032,76 @@ class FlashAttnV100MetadataBuilder(TritonAttentionMetadataBuilder): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) spec_config = getattr(self.vllm_config, "speculative_config", None) + cache_config = getattr(self.vllm_config, "cache_config", None) + model_config = self.vllm_config.model_config + hf_text_config = getattr(model_config, "hf_text_config", None) + num_attention_heads = getattr(hf_text_config, "num_attention_heads", None) + num_key_value_heads = getattr(hf_text_config, "num_key_value_heads", None) + head_dim = getattr(hf_text_config, "head_dim", None) + batch_context_shape_supported = ( + isinstance(num_attention_heads, int) + and isinstance(num_key_value_heads, int) + and num_key_value_heads > 0 + and num_attention_heads == 6 * num_key_value_heads + and head_dim == 256 + ) self._is_speculative_draft_model = ( spec_config is not None and getattr(spec_config, "draft_model_config", None) is self.vllm_config.model_config ) + self._batch_context_routing_enabled = ( + envs.VLLM_FLASH_V100_XQA_BATCH_CONTEXT_ROUTING + and envs.VLLM_FLASH_V100_DECODE_PARTITION_SIZE is None + and spec_config is None + and _batch_context_routing_cache_dtype_supported( + getattr(cache_config, "cache_dtype", None) + ) + and batch_context_shape_supported + ) + self._is_dflash_draft_model = self._is_speculative_draft_model and ( + getattr(spec_config, "method", None) == "dflash" + ) + use_dflash = bool( + spec_config is not None + and callable(getattr(spec_config, "use_dflash", None)) + and spec_config.use_dflash() + ) + selector_engine = uses_dflash_selector_engine(self.vllm_config) + 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 selector_engine + ) + self._use_sm70_dflash2_fused_smallq_metadata = bool( + envs.VLLM_SM70_DFLASH2_FUSED_SMALLQ_METADATA + and self.device.type == "cuda" + and current_platform.is_device_capability(70) + and use_dflash + and selector_engine + ) + if self._use_sm70_dflash2_fused_smallq_metadata: + logger.info_once( + "SM70 DFlash2 fused Flash-V100 small-query metadata active." + ) self._draft_block_table: torch.Tensor | None = None self._draft_seq_lens: torch.Tensor | None = None + # Prefix-anchored SWA: persistent per-request prompt-length buffer so + # the device address stays stable across steps. + kv_cache_spec = self.kv_cache_spec + self.decode_sliding_window = ( + kv_cache_spec.decode_sliding_window + if isinstance(kv_cache_spec, PrefixAnchoredSWASpec) + else None + ) + self.persistent_prefix_anchor_lens: torch.Tensor | None = None + if self.decode_sliding_window is not None: + self.persistent_prefix_anchor_lens = torch.empty( + self.vllm_config.scheduler_config.max_num_seqs, + dtype=torch.int32, + device=self.device, + ) self._draft_query_start_loc: torch.Tensor | None = None self._flash_draft_buffer_shape: tuple[int, int] | None = None self._smallq_decode_block_table: torch.Tensor | None = None @@ -2477,6 +3111,40 @@ def __init__(self, *args, **kwargs): self._smallq_buffer_shape: tuple[int, int, int] | None = None self._decode_active_num_partitions: torch.Tensor | None = None + def _attach_prefix_anchored_metadata( + self, + attn_metadata: TritonAttentionMetadata, + common_attn_metadata, + ) -> None: + window = self.decode_sliding_window + if window is None: + return + + prefix_anchor_lens = common_attn_metadata.prefix_anchor_lens + if prefix_anchor_lens is None: + raise RuntimeError( + "prefix-anchored SWA requires per-request prefix lengths" + ) + assert self.persistent_prefix_anchor_lens is not None + anchor_reqs = common_attn_metadata.num_reqs + if prefix_anchor_lens.ndim != 1 or prefix_anchor_lens.numel() < anchor_reqs: + raise RuntimeError( + "prefix-anchored SWA prefix lengths must have shape [num_reqs]" + ) + if anchor_reqs > self.persistent_prefix_anchor_lens.numel(): + raise RuntimeError( + "prefix-anchored SWA request count exceeds metadata capacity" + ) + + prefix_anchor_lens = prefix_anchor_lens.to( + device=self.device, dtype=torch.int32, non_blocking=True + ) + persistent_anchor_lens = self.persistent_prefix_anchor_lens[:anchor_reqs] + persistent_anchor_lens.copy_(prefix_anchor_lens[:anchor_reqs]) + flash_metadata = _as_flash_v100_metadata(attn_metadata) + flash_metadata.prefix_anchor_lens = persistent_anchor_lens + flash_metadata.decode_sliding_window = window + def _attach_common_flash_metadata( self, attn_metadata: TritonAttentionMetadata, @@ -2502,8 +3170,15 @@ def _attach_common_flash_metadata( else common_attn_metadata.seq_lens_cpu ) flash_metadata.causal = common_attn_metadata.causal + flash_metadata.is_dflash_selector_target = self._is_dflash_selector_target flash_metadata.max_model_len = self.vllm_config.model_config.max_model_len flash_metadata.flash_v100_cudagraph_capture = False + flash_metadata.flash_v100_batch_context_routing = ( + _batch_context_routing_for_graph_variant( + self._batch_context_routing_enabled, + getattr(common_attn_metadata, "cudagraph_graph_variant", None), + ) + ) def _attach_ddtree_metadata( self, @@ -2893,6 +3568,45 @@ def _clear_smallq_decode_metadata( flash_metadata.smallq_decode_workspace_seq_capacity_hint = None flash_metadata.smallq_decode_partition_size_hint = None + def _attach_prepared_dflash2_smallq_metadata( + self, + attn_metadata: TritonAttentionMetadata, + prepared: DFlash2SmallQPreparedMetadata, + ) -> None: + """Attach buffers refreshed by the cross-cache-group launch.""" + if prepared.builder_id != id(self): + raise ValueError("grouped small-query metadata belongs to another builder") + if ( + self._smallq_decode_block_table is None + or self._smallq_decode_seq_lens is None + or self._smallq_query_start_loc is None + or self._smallq_buffer_shape is None + ): + raise RuntimeError("grouped small-query metadata has no persistent buffers") + token_capacity, req_capacity, _ = self._smallq_buffer_shape + if ( + prepared.num_query_tokens > token_capacity + or prepared.num_reqs > req_capacity + ): + raise RuntimeError("grouped small-query metadata exceeds captured capacity") + + self._clear_smallq_decode_metadata(attn_metadata) + flash_metadata = _as_flash_v100_metadata(attn_metadata) + flash_metadata.smallq_decode_block_table = self._smallq_decode_block_table[ + : prepared.num_query_tokens + ] + flash_metadata.smallq_decode_seq_lens = self._smallq_decode_seq_lens[ + : prepared.num_query_tokens + ] + flash_metadata.smallq_query_start_loc = self._smallq_query_start_loc[ + : prepared.num_reqs + 1 + ] + flash_metadata.smallq_decode_max_seq_len_hint = prepared.max_seq_len_hint + flash_metadata.smallq_decode_workspace_seq_capacity_hint = ( + prepared.workspace_seq_capacity_hint + ) + flash_metadata.smallq_decode_partition_size_hint = prepared.partition_size_hint + def _update_smallq_decode_metadata( self, attn_metadata: TritonAttentionMetadata, @@ -2990,92 +3704,110 @@ def _update_smallq_decode_metadata( profile_stage_t0 = time.perf_counter() if profile_enabled else 0.0 query_start_loc = attn_metadata.query_start_loc[: num_reqs + 1] - query_lens = query_start_loc[1:] - query_start_loc[:-1] - real_query_lens = query_lens real_num_query_tokens = int(query_start_loc_cpu[-1].item()) if real_num_query_tokens > num_query_tokens: return - - repeat_query_lens = query_lens padding_tokens = num_query_tokens - real_num_query_tokens - if padding_tokens > 0: - repeat_query_lens = query_lens.clone() - repeat_query_lens[-1] += padding_tokens - seq_lens = attn_metadata.seq_lens[:num_reqs] - effective_seq_lens = torch.maximum( - seq_lens, - real_query_lens.to(dtype=seq_lens.dtype), - ) - block_table = block_table.clamp_min(0) prep_ms = ( (time.perf_counter() - profile_stage_t0) * 1000.0 if profile_enabled else 0.0 ) profile_stage_t0 = time.perf_counter() if profile_enabled else 0.0 - decode_block_table = torch.repeat_interleave( - block_table, - repeat_query_lens, - dim=0, - output_size=num_query_tokens, - ).contiguous() - seq_lens_rep = torch.repeat_interleave( - effective_seq_lens, - repeat_query_lens, - output_size=num_query_tokens, - ) - query_lens_rep = torch.repeat_interleave( - real_query_lens.to(dtype=seq_lens.dtype), - repeat_query_lens, - output_size=num_query_tokens, - ) - start_locs_rep = torch.repeat_interleave( - query_start_loc[:-1].to(dtype=seq_lens.dtype), - repeat_query_lens, - output_size=num_query_tokens, - ) - token_indices = self._smallq_token_indices[:num_query_tokens].to( - dtype=seq_lens.dtype - ) - offsets = token_indices - start_locs_rep + 1 - decode_seq_lens = (seq_lens_rep - query_lens_rep + offsets).contiguous() - if padding_tokens > 0: - padding_mask = token_indices >= real_num_query_tokens - decode_seq_lens = torch.where( - padding_mask, - torch.zeros_like(decode_seq_lens), - decode_seq_lens, - ).contiguous() - decode_block_table = torch.where( - padding_mask[:, None], - torch.zeros_like(decode_block_table), - decode_block_table, + if self._use_sm70_dflash2_fused_smallq_metadata: + _sm70_prepare_smallq_decode_metadata( + self._smallq_decode_block_table, + self._smallq_decode_seq_lens, + self._smallq_query_start_loc, + block_table, + seq_lens, + query_start_loc, + num_reqs=num_reqs, + num_query_tokens=num_query_tokens, + real_num_query_tokens=real_num_query_tokens, + ) + expand_ms = ( + (time.perf_counter() - profile_stage_t0) * 1000.0 + if profile_enabled + else 0.0 + ) + copy_ms = 0.0 + else: + query_lens = query_start_loc[1:] - query_start_loc[:-1] + real_query_lens = query_lens + repeat_query_lens = query_lens + if padding_tokens > 0: + repeat_query_lens = query_lens.clone() + repeat_query_lens[-1] += padding_tokens + + effective_seq_lens = torch.maximum( + seq_lens, + real_query_lens.to(dtype=seq_lens.dtype), + ) + clamped_block_table = block_table.clamp_min(0) + decode_block_table = torch.repeat_interleave( + clamped_block_table, + repeat_query_lens, + dim=0, + output_size=num_query_tokens, ).contiguous() - expand_ms = ( - (time.perf_counter() - profile_stage_t0) * 1000.0 - if profile_enabled - else 0.0 - ) + seq_lens_rep = torch.repeat_interleave( + effective_seq_lens, + repeat_query_lens, + output_size=num_query_tokens, + ) + query_lens_rep = torch.repeat_interleave( + real_query_lens.to(dtype=seq_lens.dtype), + repeat_query_lens, + output_size=num_query_tokens, + ) + start_locs_rep = torch.repeat_interleave( + query_start_loc[:-1].to(dtype=seq_lens.dtype), + repeat_query_lens, + output_size=num_query_tokens, + ) + token_indices = self._smallq_token_indices[:num_query_tokens].to( + dtype=seq_lens.dtype + ) + offsets = token_indices - start_locs_rep + 1 + decode_seq_lens = (seq_lens_rep - query_lens_rep + offsets).contiguous() + if padding_tokens > 0: + padding_mask = token_indices >= real_num_query_tokens + decode_seq_lens = torch.where( + padding_mask, + torch.zeros_like(decode_seq_lens), + decode_seq_lens, + ).contiguous() + decode_block_table = torch.where( + padding_mask[:, None], + torch.zeros_like(decode_block_table), + decode_block_table, + ).contiguous() + expand_ms = ( + (time.perf_counter() - profile_stage_t0) * 1000.0 + if profile_enabled + else 0.0 + ) - profile_stage_t0 = time.perf_counter() if profile_enabled else 0.0 - self._smallq_decode_block_table[:num_query_tokens].copy_( - decode_block_table, - non_blocking=True, - ) - self._smallq_decode_seq_lens[:num_query_tokens].copy_( - decode_seq_lens, - non_blocking=True, - ) - self._smallq_query_start_loc[: num_reqs + 1].copy_( - query_start_loc, - non_blocking=True, - ) - copy_ms = ( - (time.perf_counter() - profile_stage_t0) * 1000.0 - if profile_enabled - else 0.0 - ) + profile_stage_t0 = time.perf_counter() if profile_enabled else 0.0 + self._smallq_decode_block_table[:num_query_tokens].copy_( + decode_block_table, + non_blocking=True, + ) + self._smallq_decode_seq_lens[:num_query_tokens].copy_( + decode_seq_lens, + non_blocking=True, + ) + self._smallq_query_start_loc[: num_reqs + 1].copy_( + query_start_loc, + non_blocking=True, + ) + copy_ms = ( + (time.perf_counter() - profile_stage_t0) * 1000.0 + if profile_enabled + else 0.0 + ) profile_stage_t0 = time.perf_counter() if profile_enabled else 0.0 flash_metadata.smallq_decode_block_table = self._smallq_decode_block_table[ @@ -3118,7 +3850,7 @@ def _update_smallq_decode_metadata( "total_ms=%.3f clear_ms=%.3f guard_ms=%.3f ensure_ms=%.3f " "prep_ms=%.3f expand_ms=%.3f copy_ms=%.3f hint_ms=%.3f " "num_reqs=%d num_query_tokens=%d real_query_tokens=%d " - "padding_tokens=%d block_cols=%d", + "padding_tokens=%d block_cols=%d fused=%s", (time.perf_counter() - profile_t0) * 1000.0, clear_ms, guard_ms, @@ -3132,6 +3864,7 @@ def _update_smallq_decode_metadata( real_num_query_tokens, padding_tokens, int(block_table.shape[1]), + self._use_sm70_dflash2_fused_smallq_metadata, ) if _draft_graph_debug_enabled(): _graph_metadata_debug_log( @@ -3175,6 +3908,7 @@ def build_for_cudagraph_capture(self, common_attn_metadata): attn_metadata = super().build_for_cudagraph_capture(common_attn_metadata) self._attach_common_flash_metadata(attn_metadata, common_attn_metadata) flash_metadata = _as_flash_v100_metadata(attn_metadata) + self._attach_prefix_anchored_metadata(attn_metadata, common_attn_metadata) flash_metadata.seq_lens_cpu = capture_seq_lens_cpu # The Triton builder shortens capture seq_lens to 1 so full graph @@ -3203,12 +3937,13 @@ def build_for_cudagraph_capture(self, common_attn_metadata): workspace_seq_capacity_cap=workspace_seq_capacity_cap, partition_size_hint=partition_size_hint, ) - else: + if max_query_len == 1 or self._is_dflash_draft_model: # PIECEWISE graph replay captures the q=1 decode kernel arguments # during metadata warmup. Runtime drafting updates the persistent - # draft metadata buffers in build_for_drafting(), so capture must - # bind the graph to the same buffers instead of transient dummy - # capture tensors. + # draft metadata buffers, so capture must bind the graph to the + # same buffers instead of transient dummy capture tensors. DFlash + # parallel drafting has q=K+1; its non-causal paged-prefill graph + # consumes the same dynamic block table and sequence lengths. self._stabilize_draft_graph_metadata( attn_metadata, common_attn_metadata, @@ -3235,11 +3970,13 @@ def build( fast_build: bool = False, ddtree_parent_ids: torch.Tensor | None = None, ddtree_num_tree_tokens_cpu: torch.Tensor | None = None, + prepared_dflash2_smallq_metadata: (DFlash2SmallQPreparedMetadata | None) = None, ): attn_metadata = super().build( common_prefix_len, common_attn_metadata, fast_build ) self._attach_common_flash_metadata(attn_metadata, common_attn_metadata) + self._attach_prefix_anchored_metadata(attn_metadata, common_attn_metadata) self._attach_ddtree_metadata( attn_metadata, ddtree_parent_ids=ddtree_parent_ids, @@ -3252,13 +3989,14 @@ def build( and getattr(attn_metadata, "max_query_len", 1) > 1 and bool(torch.any(ddtree_num_tree_tokens_cpu[:num_reqs] > 0).item()) ) - if ( + if self._flash_draft_buffer_shape is not None and ( getattr(attn_metadata, "max_query_len", 1) == 1 - and self._flash_draft_buffer_shape is not None + or self._is_dflash_draft_model ): # FULL graph capture binds q=1 decode to these persistent buffers. - # Refresh them on every runtime decode step so replay sees the - # current request's block table and sequence metadata. + # DFlash binds its q=K+1 paged-prefill graph to the same buffers. + # Refresh them on every runtime step so replay sees the current + # request's block table and sequence metadata. self._stabilize_draft_graph_metadata( attn_metadata, common_attn_metadata, @@ -3274,13 +4012,20 @@ def build( # max_model_len=262144 / S=4224 / ps=256). The cap keeps launch equal # to the runtime coverage; the interface floors it at effective # max_seq_len (_get_decode_plan:165-169), so it can never under-cover. - self._update_smallq_decode_metadata( - attn_metadata, - common_attn_metadata, - workspace_seq_capacity_cap=( - int(getattr(common_attn_metadata, "max_seq_len", 0) or 0) or None - ), - ) + if prepared_dflash2_smallq_metadata is not None: + self._attach_prepared_dflash2_smallq_metadata( + attn_metadata, + prepared_dflash2_smallq_metadata, + ) + else: + self._update_smallq_decode_metadata( + attn_metadata, + common_attn_metadata, + workspace_seq_capacity_cap=( + int(getattr(common_attn_metadata, "max_seq_len", 0) or 0) + or None + ), + ) self._attach_decode_shape_hints(attn_metadata, common_attn_metadata) self._update_decode_active_num_partitions(attn_metadata, stage="build") self._debug_draft_metadata("build", attn_metadata, common_attn_metadata) @@ -3381,10 +4126,177 @@ def build_for_drafting(self, common_attn_metadata, draft_index: int): return attn_metadata +def prepare_dflash2_smallq_group_metadata( + *, + builders_by_group: list[tuple[int, FlashAttnV100MetadataBuilder]], + block_tables: tuple[torch.Tensor, ...], + seq_lens: torch.Tensor, + query_start_loc: torch.Tensor, + query_start_loc_cpu: torch.Tensor, + num_reqs: int, + num_query_tokens: int, + max_seq_len_hint: int, + workspace_seq_capacity_cap: int | None, + descriptor: DFlash2SmallQGroupDescriptor | None, +) -> ( + tuple[ + dict[int, DFlash2SmallQPreparedMetadata], + DFlash2SmallQGroupDescriptor, + ] + | None +): + """Refresh all pure-DFlash2 target full-attention metadata in one launch.""" + + def fallback(reason: str) -> None: + logger.info_once( + "DFlash2 grouped small-query metadata fell back to per-group launches: %s", + reason, + ) + + if ( + not envs.VLLM_SM70_DFLASH2_FUSED_SMALLQ_METADATA + or not envs.VLLM_SM70_DFLASH2_GROUPED_SMALLQ_METADATA + or not builders_by_group + or num_reqs <= 0 + or num_query_tokens <= 1 + or max_seq_len_hint <= 0 + ): + fallback("route-or-shape guard") + return None + if ( + seq_lens.device.type != "cuda" + or seq_lens.dtype != torch.int32 + or query_start_loc.device != seq_lens.device + or query_start_loc.dtype != torch.int32 + or query_start_loc_cpu.device.type != "cpu" + or query_start_loc_cpu.dtype != torch.int32 + or seq_lens.numel() < num_reqs + or query_start_loc.numel() < num_reqs + 1 + or query_start_loc_cpu.numel() < num_reqs + 1 + ): + fallback("common tensor contract") + return None + + real_num_query_tokens = int(query_start_loc_cpu[num_reqs].item()) + if not 0 < real_num_query_tokens <= num_query_tokens: + fallback("real query-token count") + return None + + input_tables: list[torch.Tensor] = [] + output_tables: list[torch.Tensor] = [] + output_seq_lens: list[torch.Tensor] = [] + output_query_start_locs: list[torch.Tensor] = [] + builder_ids: list[int] = [] + workspace_hints: list[int] = [] + seen_builders: set[int] = set() + for group_id, builder in builders_by_group: + builder_id = id(builder) + if builder_id in seen_builders: + continue + seen_builders.add(builder_id) + if ( + not builder._use_sm70_dflash2_fused_smallq_metadata + or group_id < 0 + or group_id >= len(block_tables) + or builder._smallq_decode_block_table is None + or builder._smallq_decode_seq_lens is None + or builder._smallq_query_start_loc is None + or builder._smallq_buffer_shape is None + ): + fallback("builder route or persistent buffers") + return None + + input_table = block_tables[group_id] + output_table = builder._smallq_decode_block_table + output_seq = builder._smallq_decode_seq_lens + output_query = builder._smallq_query_start_loc + token_capacity, req_capacity, builder_block_cols = builder._smallq_buffer_shape + input_block_cols = int(input_table.shape[1]) + if ( + input_table.device != seq_lens.device + or input_table.dtype != torch.int32 + or input_table.ndim != 2 + or input_table.shape[0] < num_reqs + or not input_table.is_contiguous() + or input_block_cols <= 0 + or builder_block_cols != input_block_cols + or token_capacity < num_query_tokens + or req_capacity < num_reqs + or output_table.dtype != torch.int32 + or output_seq.dtype != torch.int32 + or output_query.dtype != torch.int32 + or not output_table.is_contiguous() + or not output_seq.is_contiguous() + or not output_query.is_contiguous() + ): + fallback( + "cache-group tensor contract " + f"(group={group_id}, input_shape={tuple(input_table.shape)}, " + f"input_dtype={input_table.dtype}, input_device={input_table.device}, " + f"input_stride={input_table.stride()}, " + f"input_contiguous={input_table.is_contiguous()}, " + f"buffer_shape={builder._smallq_buffer_shape}, " + f"output_dtype={output_table.dtype}, " + f"output_stride={output_table.stride()}, " + f"output_contiguous={output_table.is_contiguous()}, " + f"seq_dtype={output_seq.dtype}, " + f"seq_contiguous={output_seq.is_contiguous()}, " + f"query_dtype={output_query.dtype}, " + f"query_contiguous={output_query.is_contiguous()}, " + f"num_reqs={num_reqs}, num_query_tokens={num_query_tokens}, " + f"input_block_cols={input_block_cols}, " + f"builder_block_cols={builder_block_cols})" + ) + return None + + raw_seq_capacity = input_block_cols * int(builder.block_size) + if workspace_seq_capacity_cap is not None: + raw_seq_capacity = min( + raw_seq_capacity, + max(max_seq_len_hint, int(workspace_seq_capacity_cap)), + ) + input_tables.append(input_table) + output_tables.append(output_table) + output_seq_lens.append(output_seq) + output_query_start_locs.append(output_query) + builder_ids.append(builder_id) + workspace_hints.append(raw_seq_capacity) + + if not input_tables: + fallback("no distinct full-attention builders") + return None + descriptor = _sm70_prepare_grouped_smallq_decode_metadata( + output_tables, + output_seq_lens, + output_query_start_locs, + input_tables, + seq_lens[:num_reqs], + query_start_loc[: num_reqs + 1], + num_reqs=num_reqs, + num_query_tokens=num_query_tokens, + real_num_query_tokens=real_num_query_tokens, + descriptor=descriptor, + ) + prepared = { + builder_id: DFlash2SmallQPreparedMetadata( + builder_id=builder_id, + num_reqs=num_reqs, + num_query_tokens=num_query_tokens, + max_seq_len_hint=max_seq_len_hint, + workspace_seq_capacity_hint=workspace_hint, + ) + for builder_id, workspace_hint in zip(builder_ids, workspace_hints) + } + return prepared, descriptor + + class FlashAttnV100Impl(TritonAttentionImpl): """Flash Attention V100 implementation with explicit fallback policy.""" def __init__(self, *args, **kwargs): + self.prefix_anchored_decode_window = kwargs.pop( + "prefix_anchored_decode_window", None + ) super().__init__(*args, **kwargs) self.kv_cache_dtype = _normalize_flash_v100_kv_cache_dtype(self.kv_cache_dtype) _log_kv_dtype_contract(self.kv_cache_dtype) @@ -3399,6 +4311,7 @@ def __init__(self, *args, **kwargs): self.flash_attn_prefill_paged_bfla, self.flash_attn_prefill_paged_splitkv, ) = _get_flash_ops() + self.flash_attn_grouped_verify_paged = _get_flash_grouped_verify_op() self.fp8_e5m2_paged_kv_to_fp16 = _get_fp8_e5m2_paged_kv_bridge_op() # V100 FA2 kernels consume fp16 Q. FP8 KV cache support is implemented # as storage compression only, with K/V dequantized inside FA2 kernels. @@ -3413,10 +4326,16 @@ def __init__(self, *args, **kwargs): "workspace_seq_capacity_hint", "active_num_partitions", "partition_size_hint", + "anchor_lens", + "anchored_window", ) if self.flash_attn_decode_paged is not None and _callable_accepts_keyword(self.flash_attn_decode_paged, name) } + self._flash_prefill_paged_supports_anchor = ( + self.flash_attn_prefill_paged is not None + and _callable_accepts_keyword(self.flash_attn_prefill_paged, "anchor_lens") + ) paged_prefill_enable = os.getenv("VLLM_FLASH_V100_ENABLE_PAGED_PREFILL") paged_prefill_disable = ( os.getenv("VLLM_FLASH_V100_DISABLE_PAGED_PREFILL", "0") == "1" @@ -3511,6 +4430,18 @@ def __init__(self, *args, **kwargs): self.use_decode_xqa and os.getenv("VLLM_FLASH_V100_SMALLQ_DECODE_USE_XQA", "1") == "1" ) + self.use_dflash2_grouped_verify = ( + self.flash_attn_grouped_verify_paged is not None + and envs.VLLM_FLASH_V100_DFLASH2_GROUPED_VERIFY + and current_platform.is_device_capability(70) + ) + self.dflash2_grouped_verify_min_model_len = ( + envs.VLLM_FLASH_V100_DFLASH2_GROUPED_VERIFY_MIN_MODEL_LEN + ) + if self.dflash2_grouped_verify_min_model_len < 1: + raise ValueError( + "VLLM_FLASH_V100_DFLASH2_GROUPED_VERIFY_MIN_MODEL_LEN must be positive" + ) decode_scalar_paged_env = os.getenv("VLLM_FLASH_V100_DECODE_USE_SCALAR_PAGED") self.use_decode_scalar_paged = decode_scalar_paged_env != "0" self.compare_bhmd_out_dir = os.getenv("VLLM_FLASH_V100_COMPARE_BHMD_OUT_DIR") @@ -3536,6 +4467,52 @@ def __init__(self, *args, **kwargs): self._decode_cache_len = 0 self._decode_cache_capacity = 0 + if self.prefix_anchored_decode_window is not None: + if ( + self.prefix_anchored_decode_window <= 0 + or self.attn_type != AttentionType.DECODER + or self.kv_cache_dtype != "auto" + ): + raise ValueError( + "prefix-anchored SWA requires a positive window, causal " + "decoder attention, and an fp16 KV cache" + ) + if self.use_triton_prefill: + raise ValueError( + "prefix-anchored SWA cannot use the Triton prefill fallback" + ) + if ( + not self.use_flash_v100_decode + or not self.use_decode_scalar_paged + or not {"anchor_lens", "anchored_window"} + <= self._flash_decode_paged_kwargs + ): + raise RuntimeError( + "prefix-anchored SWA requires the masked scalar paged " + "decode extension" + ) + if ( + not self.use_flash_v100_prefill_paged + or not self._flash_prefill_paged_supports_anchor + ): + raise RuntimeError( + "prefix-anchored SWA requires the masked paged prefill extension" + ) + + # Select the only two routes that carry the anchored mask once at + # construction time. The default-off hot path therefore retains + # its existing route predicates without extra metadata parsing. + self.smallq_decode_max_query_len = 0 + self.use_decode_paged_prefill = False + self.use_decode_dense_cache = False + self.use_decode_dense_reference = False + self.use_decode_xqa = False + self.use_smallq_decode_xqa = False + self.use_flash_v100_prefill_splitkv = False + self.use_flash_v100_prefill_bfla = False + self.use_flash_v100_prefill_contig_dense = False + self.use_flash_v100_prefill_gather_dense = False + def _reset_decode_cache(self) -> None: self._decode_cache_k = None self._decode_cache_v = None @@ -4108,6 +5085,8 @@ def _call_flash_attn_decode_paged( workspace_seq_capacity_hint: int | None = None, active_num_partitions: int | None = None, partition_size_hint: int | None = None, + anchor_lens: torch.Tensor | None = None, + anchored_window: int = 0, ) -> None: kwargs: dict[str, object] = { "softmax_scale": softmax_scale, @@ -4123,6 +5102,15 @@ def _call_flash_attn_decode_paged( "FLASH_ATTN_V100 decode op does not support sliding-window " "attention with this extension build." ) + if anchor_lens is not None and anchored_window > 0: + if "anchor_lens" not in self._flash_decode_paged_kwargs: + raise RuntimeError( + "FLASH_ATTN_V100 decode op does not support the anchored " + "decode-window mask with this extension build; rebuild " + "flash_attn_v100." + ) + kwargs["anchor_lens"] = anchor_lens + kwargs["anchored_window"] = anchored_window optional_kwargs = { "max_seq_len_hint": max_seq_len_hint, "workspace_seq_capacity_hint": workspace_seq_capacity_hint, @@ -4141,6 +5129,122 @@ def _call_flash_attn_decode_paged( **kwargs, ) + def _dflash2_grouped_verify_allowed( + self, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + attn_metadata: TritonAttentionMetadata, + *, + num_query_tokens: int, + ) -> bool: + """Gate the verifier on its hardware and tensor-layout contract.""" + global _logged_prefill_smallq_grouped_verify_gate + block_table = getattr(attn_metadata, "block_table", None) + seq_lens = getattr(attn_metadata, "seq_lens", None) + allowed = bool( + self.use_dflash2_grouped_verify + and self.flash_attn_grouped_verify_paged is not None + and getattr(attn_metadata, "is_dflash_selector_target", False) + and getattr(attn_metadata, "max_model_len", 0) + >= 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 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) + and tuple(key_cache.shape[2:]) == (1, 256) + and tuple(value_cache.shape) == tuple(key_cache.shape) + and key_cache.dtype == torch.uint8 + and value_cache.dtype == torch.uint8 + and key_cache.stride(-1) == 1 + and value_cache.stride(-1) == 1 + and self.kv_cache_dtype == "fp8_e5m2" + and block_table is not None + and block_table.ndim == 2 + and block_table.shape[0] == 1 + and block_table.device == query.device + and block_table.dtype == torch.int32 + and block_table.is_contiguous() + and seq_lens is not None + and seq_lens.ndim == 1 + and seq_lens.shape[0] == 1 + and seq_lens.device == query.device + and seq_lens.dtype == torch.int32 + and seq_lens.is_contiguous() + ) + if ( + self.use_dflash2_grouped_verify + and not allowed + and not _logged_prefill_smallq_grouped_verify_gate + ): + logger.info( + "FLASH_ATTN_V100 DFlash2 grouped verifier gate rejected: " + "op=%s marker=%s max_model_len=%s min_model_len=%s " + "causal=%s window=%s actual=%d q=%s/%s " + "k=%s/%s v=%s/%s kv_dtype=%s block_table=%s/%s " + "seq_lens=%s/%s.", + self.flash_attn_grouped_verify_paged is not None, + getattr(attn_metadata, "is_dflash_selector_target", False), + getattr(attn_metadata, "max_model_len", None), + self.dflash2_grouped_verify_min_model_len, + getattr(attn_metadata, "causal", True), + self._flash_v100_window_size(causal=True), + num_query_tokens, + tuple(query.shape), + query.dtype, + tuple(key_cache.shape), + key_cache.dtype, + tuple(value_cache.shape), + value_cache.dtype, + self.kv_cache_dtype, + None if block_table is None else tuple(block_table.shape), + None if block_table is None else block_table.dtype, + None if seq_lens is None else tuple(seq_lens.shape), + None if seq_lens is None else seq_lens.dtype, + ) + _logged_prefill_smallq_grouped_verify_gate = True + return allowed + + def _call_dflash2_grouped_verify( + self, + layer: torch.nn.Module, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + attn_metadata: TritonAttentionMetadata, + *, + out: torch.Tensor, + ) -> None: + global _logged_prefill_smallq_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)." + ) + _logged_prefill_smallq_grouped_verify = True + self.flash_attn_grouped_verify_paged( + query, + key_cache, + value_cache, + attn_metadata.block_table[:1], + attn_metadata.seq_lens[:1], + softmax_scale=self.scale, + out=out, + kv_cache_dtype=self.kv_cache_dtype, + k_scale=float(layer._k_scale_float), + v_scale=float(layer._v_scale_float), + one_pass=True, + ) + _log_fp8_kv_cache_route("decode", self.kv_cache_dtype, "dflash2_grouped_verify") + _record_route("prefill_smallq_dflash2_grouped_verify") + def _smallq_decode_xqa_allowed( self, query: torch.Tensor, @@ -4171,7 +5275,7 @@ def _smallq_decode_xqa_allowed( return False fp16_kv = ( - self.kv_cache_dtype in ("auto", "bfloat16") + self.kv_cache_dtype in ("auto", "float16", "bfloat16") and key_cache.dtype == torch.float16 and value_cache.dtype == torch.float16 ) @@ -4288,6 +5392,34 @@ def _call_flash_attn_smallq_decode_paged( ) _record_route("prefill_smallq_decode_scalar") + def _anchored_swa_params( + self, + attn_metadata: TritonAttentionMetadata, + ) -> tuple[torch.Tensor | None, int]: + """Anchored decode-window mask parameters, when active. + + Returns ``(prefix_anchor_lens, decode_sliding_window)`` when this + decoder cache group carries the engine's prefix-anchored spec and + per-request prompt lengths; otherwise ``(None, 0)``. + """ + window = self.prefix_anchored_decode_window + if window is None: + return None, 0 + + metadata_window = getattr(attn_metadata, "decode_sliding_window", None) + anchor_lens = getattr(attn_metadata, "prefix_anchor_lens", None) + if ( + self.attn_type != AttentionType.DECODER + or self.kv_cache_dtype != "auto" + or metadata_window != window + or anchor_lens is None + ): + raise RuntimeError( + "FLASH_ATTN_V100 prefix-anchored SWA metadata does not match " + "the enabled decoder-layer contract" + ) + return anchor_lens, int(window) + def _small_query_decode_enabled( self, attn_metadata: TritonAttentionMetadata, @@ -4298,7 +5430,6 @@ def _small_query_decode_enabled( or self.smallq_decode_max_query_len <= 0 ): return False - query_start_loc_cpu = getattr(attn_metadata, "query_start_loc_cpu", None) query_start_loc = ( query_start_loc_cpu @@ -4357,7 +5488,14 @@ def forward( "Flash op is unavailable or the attention features/KV cache dtype " "are unsupported. Select TRITON_ATTN for a full Triton route, or " "set VLLM_FLASH_V100_ALLOW_TRITON_FALLBACK=1 for explicit " - "diagnostic fallback." + "diagnostic fallback. " + f"Details: layer={layer_info.get('layer_name')!r}, " + f"flash_ops_available={self.use_flash_v100}, " + f"attn_type={self.attn_type!r}, " + f"has_alibi={self.alibi_slopes is not None}, " + f"logits_soft_cap={self.logits_soft_cap!r}, " + f"has_sinks={self.sinks is not None}, " + f"kv_cache_dtype={self.kv_cache_dtype!r}." ) if not (self.allow_triton_fallback or is_dflash_draft_attn): raise RuntimeError(message) @@ -4488,6 +5626,26 @@ def forward( # look like no-prefix prefill, while replayed MTP verification # is a uniform small-query decode over an existing KV prefix. # Capture the same small-query kernel branch that replay needs. + is_dflash_non_causal = bool( + getattr(layer, "is_dflash_draft_attn", False) + ) and not bool(getattr(attn_metadata, "causal", True)) + if is_dflash_non_causal: + # DFlash pre-inserts target context K/V before replay. Its + # dummy capture has seq_len == query_len and would + # otherwise freeze the no-prefix dense branch into the + # graph. Bind directly to the non-causal paged-prefix + # kernel; runtime updates its persistent sequence and + # block-table buffers before every replay. + _record_route("prefill_capture_dflash_noncausal_paged") + return self._flash_v100_prefill_with_prefix( + layer, + query, + key, + value, + kv_cache, + attn_metadata, + output, + ) smallq_decode = self._small_query_decode_enabled(attn_metadata) if smallq_decode: if _draft_graph_debug_enabled(): @@ -5321,6 +6479,10 @@ def _flash_v100_decode( ) -> torch.Tensor: """Decode path using Flash V100 directly over paged KV cache.""" window_size = self._flash_v100_window_size(causal=True) + if self.prefix_anchored_decode_window is None: + anchor_lens, anchored_window = None, 0 + else: + anchor_lens, anchored_window = self._anchored_swa_params(attn_metadata) num_actual_tokens = attn_metadata.num_actual_tokens query = query[:num_actual_tokens] out_view = output[:num_actual_tokens] @@ -5335,13 +6497,21 @@ def _flash_v100_decode( else 0 ) xqa_kv_supported = ( - self.kv_cache_dtype in ("auto", "bfloat16") - and key_cache.dtype == torch.float16 - and value_cache.dtype == torch.float16 - ) or ( - self.kv_cache_dtype == "fp8_e5m2" - and key_cache.dtype == torch.uint8 - and value_cache.dtype == torch.uint8 + ( + self.kv_cache_dtype in ("auto", "float16", "bfloat16") + and key_cache.dtype == torch.float16 + and value_cache.dtype == torch.float16 + ) + or ( + self.kv_cache_dtype == "fp8_e5m2" + and key_cache.dtype == torch.uint8 + and value_cache.dtype == torch.uint8 + ) + or ( + self.kv_cache_dtype in ("fp8", "fp8_e4m3") + and key_cache.dtype == torch.uint8 + and value_cache.dtype == torch.uint8 + ) ) # FP8 G4 XQA had no end-to-end gain on 35B-A3B TP4 and has no accepted @@ -5355,6 +6525,13 @@ def _flash_v100_decode( and key_cache.shape[2] > 0 and query.shape[1] % key_cache.shape[2] == 0 and _decode_xqa_allowed_for_q_per_kv(q_per_kv, attn_metadata) + and ( + self.kv_cache_dtype not in ("fp8", "fp8_e4m3") + or ( + q_per_kv == 6 + and (query.shape[0] == 1 or _e4m3_batch_xqa_allowed(query)) + ) + ) and ( self.kv_cache_dtype != "fp8_e5m2" or (q_per_kv != 4 and _decode_fp8_xqa_allowed(attn_metadata, query)) @@ -5408,6 +6585,13 @@ def _flash_v100_decode( None, ), partition_size_hint=partition_size_hint, + batch_context_routing=bool( + getattr( + attn_metadata, + "flash_v100_batch_context_routing", + False, + ) + ), ) _record_route("decode_xqa_paged") return output @@ -5448,6 +6632,8 @@ def _flash_v100_decode( "flash_v100_decode_active_num_partitions", None, ), + anchor_lens=anchor_lens, + anchored_window=anchored_window, ) _record_route("decode_scalar_paged") return output @@ -5795,6 +6981,25 @@ def _flash_v100_small_query_prefill_as_decode( int(query.shape[0]), int(output.shape[0]), ) + if self.use_dflash2_grouped_verify and self._dflash2_grouped_verify_allowed( + query, + key_cache, + value_cache, + attn_metadata, + num_query_tokens=num_query_tokens, + ): + query = query[:num_query_tokens] + out_view = output[:num_query_tokens] + self._call_dflash2_grouped_verify( + layer, + query, + key_cache, + value_cache, + attn_metadata, + out=out_view, + ) + return output + persistent_decode_block_table = getattr( attn_metadata, "smallq_decode_block_table", @@ -6316,6 +7521,25 @@ def _flash_v100_prefill_with_prefix( global _logged_prefill_compare, _logged_prefill_smallq_decode causal = getattr(attn_metadata, "causal", True) window_size = self._flash_v100_window_size(causal) + if self.prefix_anchored_decode_window is None: + anchor_lens, anchored_window = None, 0 + else: + anchor_lens, anchored_window = self._anchored_swa_params(attn_metadata) + if anchor_lens is not None: + # Fail closed: with the anchored decode-window mask active the + # KV cache manager evicts gap blocks, so running any unmasked + # prefill route would silently produce wrong output. + if not self.use_flash_v100_prefill_paged: + raise RuntimeError( + "FLASH_ATTN_V100 anchored decode-window mask requires " + "the paged prefill kernel; it is disabled or unavailable." + ) + if not self._flash_prefill_paged_supports_anchor: + raise RuntimeError( + "FLASH_ATTN_V100 prefill op does not support the " + "anchored decode-window mask with this extension build; " + "rebuild flash_attn_v100." + ) num_actual_tokens = attn_metadata.num_actual_tokens query = query[:num_actual_tokens] out_view = output[:num_actual_tokens] @@ -6351,6 +7575,11 @@ def _flash_v100_prefill_with_prefix( attn_metadata, query_start_loc, ): + if anchor_lens is not None: + raise RuntimeError( + "FLASH_ATTN_V100 anchored decode-window mask does not " + "support ddtree drafting metadata." + ) return self._flash_v100_ddtree_small_query_prefill_dense( layer, query, @@ -6366,6 +7595,7 @@ def _flash_v100_prefill_with_prefix( if ( causal + and anchor_lens is None and self.use_flash_v100_decode and self.smallq_decode_max_query_len > 0 and max_query_len <= self.smallq_decode_max_query_len @@ -6405,6 +7635,36 @@ def _flash_v100_prefill_with_prefix( q_len = end - start seq_len = int(seq_lens[i].item()) q_seq = query[start:end].unsqueeze(0) + if anchor_lens is not None: + # Anchored decode-window mask: single masked paged + # prefill route; every unmasked fast path is bypassed. + _record_route("prefill_prefix_paged_anchored") + out_seq = self._run_prefill_paged_call( + route="prefill_prefix_paged_anchored", + q_len=q_len, + seq_len=seq_len, + heads_q=query.shape[1], + heads_kv=num_kv_heads, + head_dim=head_dim, + block_size=block_size, + fn=lambda q_seq=q_seq, i=i: self.flash_attn_prefill_paged( # type: ignore[misc] + q_seq, + key_cache, + value_cache, + attn_metadata.block_table[i : i + 1], + attn_metadata.seq_lens[i : i + 1], + softmax_scale=self.scale, + kv_cache_dtype=self.kv_cache_dtype, + k_scale=float(layer._k_scale_float), + v_scale=float(layer._v_scale_float), + causal=causal, + window_size=window_size, + anchor_lens=anchor_lens[i : i + 1], + anchored_window=anchored_window, + ), + ) + out_view[start:end].copy_(out_seq.squeeze(0)) + continue bfla_block_mask = None use_bfla = self._should_use_prefill_bfla( q_len=q_len, diff --git a/vllm/v1/attention/backends/flashinfer_sm70.py b/vllm/v1/attention/backends/flashinfer_sm70.py index 182234612a..42762bb5e8 100644 --- a/vllm/v1/attention/backends/flashinfer_sm70.py +++ b/vllm/v1/attention/backends/flashinfer_sm70.py @@ -15,6 +15,7 @@ from vllm.platforms.interface import DeviceCapability from vllm.v1.attention.backend import AttentionType from vllm.v1.attention.backends.flash_attn_v100 import ( + DFlash2SmallQPreparedMetadata, FlashAttnV100Backend, FlashAttnV100Impl, FlashAttnV100MetadataBuilder, @@ -262,6 +263,7 @@ def build( fast_build: bool = False, ddtree_parent_ids=None, ddtree_num_tree_tokens_cpu=None, + prepared_dflash2_smallq_metadata: DFlash2SmallQPreparedMetadata | None = None, ): attn_metadata = super().build( common_prefix_len, @@ -269,6 +271,7 @@ def build( fast_build, ddtree_parent_ids, ddtree_num_tree_tokens_cpu, + prepared_dflash2_smallq_metadata, ) return self._attach_planner_decision( attn_metadata, diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index b52a06422c..33f496826c 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -27,6 +27,7 @@ split_decodes_and_prefills, ) from vllm.v1.kv_cache_interface import AttentionSpec, MambaSpec +from vllm.v1.worker.gpu.attn_utils import CommonGDNSpecMetadata logger = init_logger(__name__) @@ -59,6 +60,18 @@ class _GDNDdTreeFastCommonBuffers: token_index_initialized_size: int = 0 +@dataclass +class DFlash2GDNGroupDescriptor: + """Persistent pointer tables for one target runner's GDN cache groups.""" + + key: tuple[object, ...] + block_table_ptrs: torch.Tensor + state_output_ptrs: torch.Tensor + block_table_strides: torch.Tensor + prepared_key: tuple[int, int, int, int] | None = None + prepared_metadata: dict[int, "GDNAttentionMetadata"] | None = None + + _GDN_DDTREE_FAST_COMMON_BUFFERS: dict[ tuple[str, int | None, int, int], _GDNDdTreeFastCommonBuffers ] = {} @@ -84,8 +97,9 @@ def _get_ddtree_gdn_fast_common_buffers( dtype=torch.bool, device=device, ) - spec_token_indx = torch.empty( - (decode_cudagraph_max_bs * width,), + spec_token_capacity = decode_cudagraph_max_bs * width + spec_token_indx = torch.arange( + spec_token_capacity, dtype=torch.int32, device=device, ) @@ -116,6 +130,7 @@ def _get_ddtree_gdn_fast_common_buffers( spec_query_start_loc=spec_query_start_loc, num_accepted_tokens=num_accepted_tokens, spec_state_slot_selectors=spec_state_slot_selectors, + token_index_initialized_size=spec_token_capacity, ) _GDN_DDTREE_FAST_COMMON_BUFFERS[key] = buffers return buffers @@ -149,6 +164,90 @@ def _dflash_ddtree_gdn_fast_build_triton_enabled() -> bool: ).strip().lower() not in ("0", "false", "no", "off", "") +@triton.jit +def _load_gdn_i32_ptr(ptr_to_ptr): + ptr = tl.load(ptr_to_ptr) + ptr = tl.cast(ptr, tl.pointer_type(tl.int32)) + return tl.multiple_of(ptr, 16) + + +@triton.jit +def _dflash2_gdn_group_metadata_kernel( + block_table_ptrs, + state_output_ptrs, + block_table_strides, + state_start_indices, + req_index_mapping, + spec_query_start_loc_src, + num_accepted_src, + state_selector_src, + spec_sequence_masks_out, + spec_query_start_loc_out, + num_accepted_out, + state_selector_out, + num_spec_decodes, + batch_size, + WIDTH: tl.constexpr, + PAD_ID: tl.constexpr, + BLOCK: tl.constexpr, + USE_STATE_START: tl.constexpr, +): + """Write every GDN group's state IDs and the shared graph metadata.""" + group_id = tl.program_id(0) + block_table = _load_gdn_i32_ptr(block_table_ptrs + group_id) + state_output = _load_gdn_i32_ptr(state_output_ptrs + group_id) + block_table_stride = tl.load(block_table_strides + group_id) + + offsets = tl.arange(0, BLOCK) + rows = offsets // WIDTH + columns = offsets % WIDTH + output_mask = offsets < batch_size * WIDTH + live_state_mask = output_mask & (rows < num_spec_decodes) + state_columns = columns + if USE_STATE_START: + req_indices = tl.load( + req_index_mapping + rows, + mask=live_state_mask, + other=0, + ) + live_state_mask &= req_indices >= 0 + state_starts = tl.load( + state_start_indices + req_indices, + mask=live_state_mask, + other=-1, + ) + state_columns = columns + state_starts + live_state_mask &= (state_starts >= 0) & (state_columns < block_table_stride) + state_ids = tl.load( + block_table + rows * block_table_stride + state_columns, + mask=live_state_mask, + other=PAD_ID, + ) + tl.store(state_output + offsets, state_ids, mask=output_mask) + + # All groups share these buffers. Restrict the stores to group zero so the + # launch has one writer without introducing a second metadata kernel. + is_common_writer = group_id == 0 + row_mask = (offsets < batch_size) & is_common_writer + live_row_mask = row_mask & (offsets < num_spec_decodes) + accepted = tl.load(num_accepted_src + offsets, mask=live_row_mask, other=1) + selectors = tl.load(state_selector_src + offsets, mask=live_row_mask, other=1) + tl.store( + spec_sequence_masks_out + offsets, offsets < num_spec_decodes, mask=row_mask + ) + tl.store(num_accepted_out + offsets, accepted, mask=row_mask) + tl.store(state_selector_out + offsets, selectors, mask=row_mask) + + query_mask = (offsets < batch_size + 1) & is_common_writer + query_src_offsets = tl.minimum(offsets, num_spec_decodes) + query_offsets = tl.load( + spec_query_start_loc_src + query_src_offsets, + mask=query_mask, + other=0, + ) + tl.store(spec_query_start_loc_out + offsets, query_offsets, mask=query_mask) + + @triton.jit def _ddtree_gdn_fast_metadata_kernel( state_src, @@ -353,6 +452,12 @@ class GDNAttentionMetadata: batch_ptr: torch.Tensor | None = None token_chunk_offset_ptr: torch.Tensor | None = None + # Pure DFlash2 full-graph replay updates these persistent tensors in place. + # Cache the tuple consumed by the layer registry so the prepared fast path + # does not allocate two empty CUDA tensors and rebuild the same tuple for + # every GDN cache group on every verification round. + _prepared_spec_metadata_tensors: GDN_SPEC_METADATA_TENSORS | None = None + @dataclass class GDNSpecDecodeStateContract: @@ -716,7 +821,10 @@ def __init__( self._ddtree_fast_tail_key: tuple[int, int, int] | None = None if ( self.use_spec_decode - and envs.VLLM_SM70_QWEN_GDN_SPEC_CORE_OP + and ( + envs.VLLM_SM70_QWEN_GDN_SPEC_CORE_OP + or envs.VLLM_SM70_DFLASH2_FUSED_GDN_METADATA + ) and _dflash_ddtree_gdn_shared_common_enabled() ): self._ddtree_fast_common_buffers = _get_ddtree_gdn_fast_common_buffers( @@ -1183,6 +1291,8 @@ def build( # type: ignore[override] spec_state_slot_selectors: torch.Tensor | None = None, num_decode_draft_tokens_cpu: torch.Tensor | None = None, spec_sequence_masks_cpu: torch.Tensor | None = None, + common_gdn_metadata: CommonGDNSpecMetadata | None = None, + prepared_dflash2_metadata: GDNAttentionMetadata | None = None, current_state_block_ids: torch.Tensor | None = None, ddtree_parent_ids: torch.Tensor | None = None, ddtree_num_tree_tokens_cpu: torch.Tensor | None = None, @@ -1199,6 +1309,37 @@ def build( # type: ignore[override] query_start_loc = m.query_start_loc query_start_loc_cpu = m.query_start_loc_cpu + if prepared_dflash2_metadata is not None: + if ( + for_cudagraph_capture + or common_gdn_metadata is None + or ddtree_parent_ids is not None + or current_state_block_ids is not None + ): + raise ValueError( + "Prepared DFlash2 GDN metadata is valid only for MRV2 " + "runtime graph replay" + ) + prepared_state = prepared_dflash2_metadata.spec_state_indices_tensor + if ( + prepared_state is None + or prepared_state.data_ptr() + != self.spec_state_indices_tensor.data_ptr() + ): + raise ValueError( + "Prepared DFlash2 GDN metadata does not belong to this builder" + ) + prepared_tensors = prepared_dflash2_metadata._prepared_spec_metadata_tensors + if prepared_tensors is None: + prepared_tensors = gdn_spec_metadata_tensors( + prepared_dflash2_metadata, + query_start_loc.device, + ) + prepared_dflash2_metadata._prepared_spec_metadata_tensors = ( + prepared_tensors + ) + register_gdn_spec_metadata_tensors(self.layer_names, prepared_tensors) + return prepared_dflash2_metadata if fast_build and _dflash_ddtree_gdn_fast_build_enabled(): fast_metadata = self._build_fast_pure_ddtree_full_graph( common_attn_metadata=common_attn_metadata, @@ -1227,6 +1368,12 @@ def build( # type: ignore[override] is_mamba_cache_all = self.vllm_config.cache_config.mamba_cache_mode == "all" num_reqs = query_start_loc_cpu.numel() - 1 + if common_gdn_metadata is not None: + if ddtree_parent_ids is not None or current_state_block_ids is not None: + raise ValueError( + "Common MRV2 GDN metadata cannot be used with DDTree state" + ) + spec_sequence_masks_cpu = common_gdn_metadata.spec_sequence_masks_cpu if spec_sequence_masks_cpu is not None: assert spec_sequence_masks_cpu.dtype == torch.bool assert spec_sequence_masks_cpu.ndim == 1 @@ -1263,6 +1410,9 @@ def build( # type: ignore[override] if not self.use_spec_decode: spec_sequence_masks = None num_spec_decodes = 0 + elif common_gdn_metadata is not None: + spec_sequence_masks = common_gdn_metadata.spec_sequence_masks + num_spec_decodes = common_gdn_metadata.num_spec_decodes else: if spec_sequence_masks_cpu is None: if num_decode_draft_tokens_cpu is None: @@ -1473,9 +1623,11 @@ def build( # type: ignore[override] else: # query_start_loc may be padded for CUDA graph replay. The CPU # metadata is authoritative for the live request count here. - query_lens = query_lens_cpu.to( - query_start_loc.device, non_blocking=True - ) + query_lens = None + if common_gdn_metadata is None: + query_lens = query_lens_cpu.to( + query_start_loc.device, non_blocking=True + ) profile_state_contract_t0 = ( time.perf_counter() if metadata_profile else 0.0 ) @@ -1495,113 +1647,142 @@ def build( # type: ignore[override] time.perf_counter() - profile_state_contract_t0 ) * 1000.0 - if envs.VLLM_SM70_MTP_LEGACY_GDN_MIXED_DECODE_ROUTING: - # 0.0.3 kept ordinary query_len==1 rows on the decode path even - # when another row was running speculative verification. This - # is an A/B guard for MTP-only recurrent-state corruption. - num_decodes = (non_spec_query_lens_cpu == 1).sum().item() - num_prefills = ( - non_spec_query_lens_cpu.size(0) - num_decodes - num_zero_len - ) - num_decode_tokens = num_decodes - num_prefill_tokens = ( - non_spec_query_lens_cpu.sum().item() - num_decode_tokens - ) - else: - # When active spec decodes are present, route non-spec requests - # through the prefill path so mixed batches keep separate GDN - # state metadata for spec and non-spec tokens. - num_decodes = 0 - num_prefills = non_spec_query_lens_cpu.size(0) - num_zero_len - num_decode_tokens = 0 - num_prefill_tokens = non_spec_query_lens_cpu.sum().item() - num_spec_decode_tokens = ( - query_lens_cpu.sum().item() - num_prefill_tokens - num_decode_tokens - ) - - if num_prefills == 0 and num_decodes == 0: - spec_token_size = min( - num_spec_decodes * (self.num_spec_state_tokens + 1), - query_start_loc_cpu[-1].item(), + if common_gdn_metadata is not None: + num_prefills = common_gdn_metadata.num_prefills + num_prefill_tokens = common_gdn_metadata.num_prefill_tokens + num_decodes = common_gdn_metadata.num_decodes + num_decode_tokens = common_gdn_metadata.num_decode_tokens + num_spec_decode_tokens = common_gdn_metadata.num_spec_decode_tokens + spec_token_indx = common_gdn_metadata.spec_token_indx + non_spec_token_indx = common_gdn_metadata.non_spec_token_indx + spec_query_start_loc = common_gdn_metadata.spec_query_start_loc + non_spec_query_start_loc = ( + common_gdn_metadata.non_spec_query_start_loc ) - spec_token_indx = torch.arange( - spec_token_size, - dtype=torch.int32, - device=query_start_loc.device, - ) - non_spec_token_indx = torch.empty( - 0, dtype=torch.int32, device=query_start_loc.device + non_spec_query_start_loc_cpu = ( + common_gdn_metadata.non_spec_query_start_loc_cpu ) spec_state_indices_tensor = state_contract.spec_state_indices_tensor - if for_cudagraph_capture: - spec_state_indices_tensor = torch.full_like( - spec_state_indices_tensor, PAD_SLOT_ID - ) - non_spec_state_indices_tensor = None - # Padded sequences are always at the back, so the first - # num_spec_decodes + 1 entries of query_start_loc already - # contain the correct cumulative token counts. - spec_query_start_loc = query_start_loc[: num_spec_decodes + 1] - non_spec_query_start_loc = None - non_spec_query_start_loc_cpu = None + non_spec_state_indices_tensor = ( + None + if num_prefills == 0 and num_decodes == 0 + else state_contract.non_spec_state_indices_tensor + ) else: - spec_token_masks = torch.repeat_interleave( - spec_sequence_masks, - query_lens, - output_size=query_start_loc_cpu[-1].item(), + if envs.VLLM_SM70_MTP_LEGACY_GDN_MIXED_DECODE_ROUTING: + # 0.0.3 kept ordinary query_len==1 rows on the decode + # path even when another row was running speculative + # verification. + num_decodes = (non_spec_query_lens_cpu == 1).sum().item() + num_prefills = ( + non_spec_query_lens_cpu.size(0) - num_decodes - num_zero_len + ) + num_decode_tokens = num_decodes + num_prefill_tokens = ( + non_spec_query_lens_cpu.sum().item() - num_decode_tokens + ) + else: + # Mixed non-spec rows use the prefill path so their GDN + # state metadata stays separate from verification rows. + num_decodes = 0 + num_prefills = non_spec_query_lens_cpu.size(0) - num_zero_len + num_decode_tokens = 0 + num_prefill_tokens = non_spec_query_lens_cpu.sum().item() + num_spec_decode_tokens = ( + query_lens_cpu.sum().item() + - num_prefill_tokens + - num_decode_tokens ) - index = torch.argsort(spec_token_masks, stable=True) - num_non_spec_tokens = num_prefill_tokens + num_decode_tokens - non_spec_token_indx = index[:num_non_spec_tokens] - spec_token_indx = index[num_non_spec_tokens:] - spec_state_indices_tensor = state_contract.spec_state_indices_tensor - non_spec_state_indices_tensor = ( - state_contract.non_spec_state_indices_tensor - ) - if for_cudagraph_capture: - spec_state_indices_tensor = torch.full_like( - spec_state_indices_tensor, PAD_SLOT_ID + if num_prefills == 0 and num_decodes == 0: + spec_token_size = min( + num_spec_decodes * (self.num_spec_state_tokens + 1), + query_start_loc_cpu[-1].item(), + ) + spec_token_indx = torch.arange( + spec_token_size, + dtype=torch.int32, + device=query_start_loc.device, ) - non_spec_state_indices_tensor = torch.full_like( - non_spec_state_indices_tensor, PAD_SLOT_ID + non_spec_token_indx = torch.empty( + 0, + dtype=torch.int32, + device=query_start_loc.device, ) + spec_state_indices_tensor = ( + state_contract.spec_state_indices_tensor + ) + non_spec_state_indices_tensor = None + # Padded sequences are always at the back. + 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: + assert query_lens is not None + spec_token_masks = torch.repeat_interleave( + spec_sequence_masks, + query_lens, + output_size=query_start_loc_cpu[-1].item(), + ) + index = torch.argsort(spec_token_masks, stable=True) + num_non_spec_tokens = num_prefill_tokens + num_decode_tokens + non_spec_token_indx = index[:num_non_spec_tokens] + spec_token_indx = index[num_non_spec_tokens:] - spec_query_start_loc = torch.zeros( - num_spec_decodes + 1, - dtype=torch.int32, - device=query_start_loc.device, - ) - torch.cumsum( - query_lens[spec_sequence_masks], - dim=0, - out=spec_query_start_loc[1:], - ) - non_spec_query_start_loc = torch.zeros( - query_lens.size(0) - num_spec_decodes + 1, - dtype=torch.int32, - device=query_start_loc.device, - ) - torch.cumsum( - query_lens[~spec_sequence_masks], - dim=0, - out=non_spec_query_start_loc[1:], - ) - non_spec_query_start_loc_cpu = torch.zeros( - query_lens_cpu.size(0) - num_spec_decodes + 1, - dtype=torch.int32, - device="cpu", - ) - torch.cumsum( - query_lens_cpu[~spec_sequence_masks_cpu], - dim=0, - out=non_spec_query_start_loc_cpu[1:], - ) + spec_state_indices_tensor = ( + state_contract.spec_state_indices_tensor + ) + non_spec_state_indices_tensor = ( + state_contract.non_spec_state_indices_tensor + ) + if for_cudagraph_capture: + spec_state_indices_tensor = torch.full_like( + spec_state_indices_tensor, PAD_SLOT_ID + ) + non_spec_state_indices_tensor = torch.full_like( + non_spec_state_indices_tensor, PAD_SLOT_ID + ) + + spec_query_start_loc = torch.zeros( + num_spec_decodes + 1, + dtype=torch.int32, + device=query_start_loc.device, + ) + torch.cumsum( + query_lens[spec_sequence_masks], + dim=0, + out=spec_query_start_loc[1:], + ) + non_spec_query_start_loc = torch.zeros( + query_lens.size(0) - num_spec_decodes + 1, + dtype=torch.int32, + device=query_start_loc.device, + ) + torch.cumsum( + query_lens[~spec_sequence_masks], + dim=0, + out=non_spec_query_start_loc[1:], + ) + non_spec_query_start_loc_cpu = torch.zeros( + query_lens_cpu.size(0) - num_spec_decodes + 1, + dtype=torch.int32, + device="cpu", + ) + torch.cumsum( + query_lens_cpu[~spec_sequence_masks_cpu], + dim=0, + out=non_spec_query_start_loc_cpu[1:], + ) num_accepted_tokens = state_contract.num_accepted_tokens spec_state_slot_selectors = state_contract.spec_state_slot_selectors assert spec_query_start_loc is not None - assert spec_query_start_loc[-1].item() == num_spec_decode_tokens + if common_gdn_metadata is None: + assert spec_query_start_loc[-1].item() == num_spec_decode_tokens + else: + assert ( + common_gdn_metadata.num_spec_decode_tokens == num_spec_decode_tokens + ) assert spec_state_indices_tensor is not None assert spec_state_indices_tensor.shape[0] == num_spec_decodes @@ -1995,3 +2176,319 @@ def build_for_cudagraph_capture( spec_sequence_masks_cpu=spec_sequence_masks_cpu, for_cudagraph_capture=True, ) + + +def prepare_dflash2_gdn_group_metadata( + *, + builders_by_group: list[tuple[int, GDNAttentionMetadataBuilder]], + block_tables: tuple[torch.Tensor, ...], + common_gdn_metadata: CommonGDNSpecMetadata, + num_accepted_tokens: torch.Tensor, + num_actual_tokens: int, + descriptor: DFlash2GDNGroupDescriptor | None, + state_start_indices: torch.Tensor | None = None, + req_index_mapping: torch.Tensor | None = None, +) -> ( + tuple[ + dict[int, GDNAttentionMetadata], + DFlash2GDNGroupDescriptor, + ] + | None +): + """Prepare all pure-MRV2 DFlash2 GDN graph metadata in one launch. + + ``mamba_cache_mode=none`` reads the first speculative state columns. + ``mamba_cache_mode=align`` supplies the authoritative, post-precopy state + column for each live request. DFlash2 batches keep live speculative rows at + the front and CUDA-graph padding at the back, so one pointer-table kernel can + perform the same state selection and tail fill without ten independent + gather/copy pipelines. + """ + if not envs.VLLM_SM70_DFLASH2_FUSED_GDN_METADATA: + return None + if not builders_by_group or num_actual_tokens <= 0: + return None + if num_accepted_tokens.device.type != "cuda": + return None + if num_accepted_tokens.dtype != torch.int32 or num_accepted_tokens.ndim != 1: + return None + + use_state_start = state_start_indices is not None + if use_state_start != (req_index_mapping is not None): + return None + if use_state_start: + assert state_start_indices is not None + assert req_index_mapping is not None + if ( + state_start_indices.device != num_accepted_tokens.device + or state_start_indices.dtype != torch.int32 + or state_start_indices.ndim != 1 + or req_index_mapping.device != num_accepted_tokens.device + or req_index_mapping.dtype != torch.int32 + or req_index_mapping.ndim != 1 + ): + return None + + spec_mask_cpu = common_gdn_metadata.spec_sequence_masks_cpu + num_spec_decodes = common_gdn_metadata.num_spec_decodes + if ( + common_gdn_metadata.num_prefills != 0 + or common_gdn_metadata.num_decodes != 0 + or num_spec_decodes <= 0 + or num_spec_decodes > num_actual_tokens + or num_spec_decodes > spec_mask_cpu.numel() + or common_gdn_metadata.num_spec_decode_tokens > num_actual_tokens + or ( + req_index_mapping is not None + and req_index_mapping.numel() < num_spec_decodes + ) + ): + return None + if not bool(torch.all(spec_mask_cpu[:num_spec_decodes]).item()): + return None + if bool(torch.any(spec_mask_cpu[num_spec_decodes:]).item()): + return None + + query_start_loc = common_gdn_metadata.spec_query_start_loc + if ( + query_start_loc.device != num_accepted_tokens.device + or query_start_loc.dtype != torch.int32 + or query_start_loc.ndim != 1 + or query_start_loc.numel() != num_spec_decodes + 1 + ): + return None + if num_accepted_tokens.numel() < num_spec_decodes: + return None + + first_builder = builders_by_group[0][1] + mamba_cache_mode = first_builder.vllm_config.cache_config.mamba_cache_mode + if mamba_cache_mode not in ("none", "align"): + return None + if use_state_start != (mamba_cache_mode == "align"): + return None + width = first_builder.num_spec_state_tokens + 1 + common_buffers = first_builder._ddtree_fast_common_buffers + if common_buffers is None: + return None + if ( + num_actual_tokens > first_builder.decode_cudagraph_max_bs + or common_gdn_metadata.num_spec_decode_tokens + > first_builder.decode_cudagraph_max_bs + ): + return None + + input_tables: list[torch.Tensor] = [] + output_states: list[torch.Tensor] = [] + builder_ids: list[int] = [] + group_ids: list[int] = [] + seen_builders: set[int] = set() + for group_id, builder in builders_by_group: + builder_id = id(builder) + if builder_id in seen_builders: + continue + seen_builders.add(builder_id) + if ( + builder.num_spec_state_tokens + 1 != width + or not builder.use_full_cuda_graph + or builder.vllm_config.cache_config.mamba_cache_mode != mamba_cache_mode + or builder.decode_cudagraph_max_bs < num_actual_tokens + or builder._ddtree_fast_common_buffers is not common_buffers + or group_id < 0 + or group_id >= len(block_tables) + ): + return None + block_table = block_tables[group_id] + state_output = builder.spec_state_indices_tensor + if ( + block_table.device != num_accepted_tokens.device + or block_table.dtype != torch.int32 + or block_table.ndim != 2 + or block_table.shape[0] < num_spec_decodes + or block_table.shape[1] < width + or not block_table.is_contiguous() + or state_output.device != num_accepted_tokens.device + or state_output.dtype != torch.int32 + or state_output.ndim != 2 + or state_output.shape[0] < num_actual_tokens + or state_output.shape[1] != width + or not state_output.is_contiguous() + ): + return None + input_tables.append(block_table) + output_states.append(state_output) + builder_ids.append(builder_id) + group_ids.append(group_id) + + if not input_tables: + return None + spec_token_size = common_gdn_metadata.spec_token_indx.numel() + if ( + spec_token_size > common_buffers.spec_token_indx.numel() + or num_actual_tokens > common_buffers.spec_sequence_masks.numel() + or num_actual_tokens + 1 > common_buffers.spec_query_start_loc.numel() + ): + return None + if spec_token_size > common_buffers.token_index_initialized_size: + return None + + descriptor_key: tuple[object, ...] = ( + num_accepted_tokens.device.type, + num_accepted_tokens.device.index, + width, + tuple(group_ids), + tuple(builder_ids), + tuple(table.data_ptr() for table in input_tables), + tuple(state.data_ptr() for state in output_states), + tuple(table.stride(0) for table in input_tables), + use_state_start, + common_buffers.spec_sequence_masks.data_ptr(), + common_buffers.spec_token_indx.data_ptr(), + common_buffers.non_spec_token_indx.data_ptr(), + common_buffers.spec_query_start_loc.data_ptr(), + common_buffers.num_accepted_tokens.data_ptr(), + common_buffers.spec_state_slot_selectors.data_ptr(), + ) + if descriptor is None or descriptor.key != descriptor_key: + descriptor = DFlash2GDNGroupDescriptor( + key=descriptor_key, + block_table_ptrs=torch.tensor( + [table.data_ptr() for table in input_tables], + dtype=torch.uint64, + device=num_accepted_tokens.device, + ), + state_output_ptrs=torch.tensor( + [state.data_ptr() for state in output_states], + dtype=torch.uint64, + device=num_accepted_tokens.device, + ), + block_table_strides=torch.tensor( + [table.stride(0) for table in input_tables], + dtype=torch.int64, + device=num_accepted_tokens.device, + ), + ) + + block = triton.next_power_of_2( + max(num_actual_tokens * width, num_actual_tokens + 1) + ) + _dflash2_gdn_group_metadata_kernel[(len(input_tables),)]( + descriptor.block_table_ptrs, + descriptor.state_output_ptrs, + descriptor.block_table_strides, + num_accepted_tokens if state_start_indices is None else state_start_indices, + num_accepted_tokens if req_index_mapping is None else req_index_mapping, + query_start_loc, + num_accepted_tokens, + num_accepted_tokens, + common_buffers.spec_sequence_masks, + common_buffers.spec_query_start_loc, + common_buffers.num_accepted_tokens, + common_buffers.spec_state_slot_selectors, + num_spec_decodes, + num_actual_tokens, + WIDTH=width, + PAD_ID=PAD_SLOT_ID, + BLOCK=block, + USE_STATE_START=use_state_start, + num_warps=1, + ) + common_buffers.initialized_key = ( + num_actual_tokens, + common_gdn_metadata.num_spec_decode_tokens, + width, + ) + + prepared_key = ( + num_spec_decodes, + num_actual_tokens, + common_gdn_metadata.num_spec_decode_tokens, + spec_token_size, + ) + prepared = descriptor.prepared_metadata + if descriptor.prepared_key != prepared_key or prepared is None: + prepared = {} + for builder_id, state_output in zip(builder_ids, output_states): + prepared[builder_id] = GDNAttentionMetadata( + num_prefills=0, + num_prefill_tokens=0, + num_decodes=0, + num_decode_tokens=0, + num_spec_decodes=num_spec_decodes, + num_spec_decode_tokens=common_gdn_metadata.num_spec_decode_tokens, + num_actual_tokens=num_actual_tokens, + has_initial_state=None, + chunk_indices=None, + chunk_offsets=None, + spec_query_start_loc=common_buffers.spec_query_start_loc[ + : num_actual_tokens + 1 + ], + non_spec_query_start_loc=None, + spec_state_indices_tensor=state_output[:num_actual_tokens], + non_spec_state_indices_tensor=None, + spec_sequence_masks=common_buffers.spec_sequence_masks[ + :num_actual_tokens + ], + spec_token_indx=common_buffers.spec_token_indx[:spec_token_size], + non_spec_token_indx=common_buffers.non_spec_token_indx[:0], + num_accepted_tokens=common_buffers.num_accepted_tokens[ + :num_actual_tokens + ], + spec_state_slot_selectors=common_buffers.spec_state_slot_selectors[ + :num_actual_tokens + ], + ddtree_parent_ids=None, + ddtree_num_tree_tokens_cpu=None, + nums_dict=None, + batch_ptr=None, + token_chunk_offset_ptr=None, + ) + descriptor.prepared_key = prepared_key + descriptor.prepared_metadata = prepared + + assert prepared is not None + if envs.VLLM_SM70_DFLASH2_GDN_METADATA_SHADOW: + spec_mask = common_gdn_metadata.spec_sequence_masks + expected_accepted = num_accepted_tokens[spec_mask] + for group_id, builder in builders_by_group: + actual = prepared[id(builder)] + actual_state = actual.spec_state_indices_tensor + assert actual_state is not None + source_table = block_tables[group_id][spec_mask] + if use_state_start: + assert state_start_indices is not None + assert req_index_mapping is not None + req_indices = req_index_mapping[:num_spec_decodes].to(torch.long) + starts = state_start_indices[req_indices].to(torch.long) + columns = starts[:, None] + torch.arange( + width, + dtype=torch.long, + device=source_table.device, + ) + expected_state = torch.gather(source_table, 1, columns) + else: + expected_state = source_table[:, :width] + torch.testing.assert_close( + actual_state[:num_spec_decodes], + expected_state, + rtol=0, + atol=0, + ) + torch.testing.assert_close( + actual_state[num_spec_decodes:], + torch.full_like( + actual_state[num_spec_decodes:], + PAD_SLOT_ID, + ), + rtol=0, + atol=0, + ) + torch.testing.assert_close( + common_buffers.num_accepted_tokens[:num_spec_decodes], + expected_accepted, + rtol=0, + atol=0, + ) + if torch.any(common_buffers.spec_sequence_masks[num_spec_decodes:]).item(): + raise AssertionError("DFlash2 fused GDN metadata left a live padded row") + + return prepared, descriptor diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 2870ec9a15..2171f74647 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -203,6 +203,11 @@ class DeepseekV32IndexerMetadata: # hacky way to access the data now, need to be in chunked meta seq_lens: torch.Tensor max_seq_len: int + # Decode seq_lens and the indexer KV cache use compressed positions for + # DeepSeek V4. Keep the matching host bound explicit: CUDA Graph capture + # needs a static logits width, but using the uncompressed attention bound + # launches 4x/128x too much work for C4/C128 layers. + compressed_max_seq_len: int slot_mapping: torch.Tensor # New for MLA (compared to FlashAttention) @@ -627,6 +632,9 @@ def build( attn_metadata = DeepseekV32IndexerMetadata( seq_lens=common_attn_metadata.seq_lens, max_seq_len=common_attn_metadata.max_seq_len, + compressed_max_seq_len=max( + 1, common_attn_metadata.max_seq_len // self.compress_ratio + ), slot_mapping=compressed_slot_mapping, num_decodes=num_decodes, num_decode_tokens=num_decode_tokens, diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index a96c0dd115..9688a42b64 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" ) 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 c5e8953745..224b4d98d7 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -34,7 +34,7 @@ def __init__( self, kv_cache_config: KVCacheConfig, max_model_len: int, - max_num_batched_tokens: int, + max_in_flight_tokens: int, use_eagle: bool, enable_caching: bool, enable_kv_cache_events: bool, @@ -66,7 +66,7 @@ def __init__( self.single_type_managers = tuple( get_manager_for_kv_cache_spec( kv_cache_spec=kv_cache_group.kv_cache_spec, - max_num_batched_tokens=max_num_batched_tokens, + max_in_flight_tokens=max_in_flight_tokens, max_model_len=max_model_len, block_pool=self.block_pool, enable_caching=enable_caching, @@ -236,7 +236,10 @@ def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]: ] def remove_skipped_blocks( - self, request_id: str, total_computed_tokens: int + self, + request_id: str, + processed_computed_tokens: int, + num_prompt_tokens: int | None = None, ) -> None: """ Remove the blocks that are no longer needed from `blocks` and replace @@ -244,11 +247,16 @@ def remove_skipped_blocks( Args: request_id: The request ID. - total_computed_tokens: The total number of computed tokens, including - local computed tokens and external computed tokens. + processed_computed_tokens: Computed-token prefix length covering + fully processed and committed tokens only (safe to free). + num_prompt_tokens: Optional prompt length. Prefix-anchored SWA + managers use this to free gap blocks between the prefill tail + and the decode window; other manager types ignore it. """ for manager in self.single_type_managers: - manager.remove_skipped_blocks(request_id, total_computed_tokens) + manager.remove_skipped_blocks( + request_id, processed_computed_tokens, num_prompt_tokens + ) def get_blocks(self, request_id: str) -> tuple[list[KVCacheBlock], ...]: """ @@ -285,7 +293,7 @@ def __init__( self, kv_cache_config: KVCacheConfig, max_model_len: int, - max_num_batched_tokens: int, + max_in_flight_tokens: int, use_eagle: bool, enable_kv_cache_events: bool, dcp_world_size: int, @@ -296,7 +304,7 @@ def __init__( super().__init__( kv_cache_config, max_model_len, - max_num_batched_tokens, + max_in_flight_tokens, use_eagle, False, enable_kv_cache_events, @@ -332,7 +340,7 @@ def __init__( self, kv_cache_config: KVCacheConfig, max_model_len: int, - max_num_batched_tokens: int, + max_in_flight_tokens: int, use_eagle: bool, enable_caching: bool, enable_kv_cache_events: bool, @@ -344,7 +352,7 @@ def __init__( super().__init__( kv_cache_config, max_model_len, - max_num_batched_tokens, + max_in_flight_tokens, use_eagle, enable_caching, enable_kv_cache_events, @@ -399,7 +407,7 @@ def __init__( self, kv_cache_config: KVCacheConfig, max_model_len: int, - max_num_batched_tokens: int, + max_in_flight_tokens: int, use_eagle: bool, enable_caching: bool, enable_kv_cache_events: bool, @@ -411,7 +419,7 @@ def __init__( super().__init__( kv_cache_config, max_model_len, - max_num_batched_tokens, + max_in_flight_tokens, use_eagle, enable_caching, enable_kv_cache_events, @@ -610,7 +618,7 @@ def _get_block_hashes(kv_cache_spec: KVCacheSpec) -> BlockHashList: def get_kv_cache_coordinator( kv_cache_config: KVCacheConfig, max_model_len: int, - max_num_batched_tokens: int, + max_in_flight_tokens: int, use_eagle: bool, enable_caching: bool, enable_kv_cache_events: bool, @@ -623,7 +631,7 @@ def get_kv_cache_coordinator( return KVCacheCoordinatorNoPrefixCache( kv_cache_config, max_model_len, - max_num_batched_tokens, + max_in_flight_tokens, use_eagle, enable_kv_cache_events, dcp_world_size=dcp_world_size, @@ -635,7 +643,7 @@ def get_kv_cache_coordinator( return UnitaryKVCacheCoordinator( kv_cache_config, max_model_len, - max_num_batched_tokens, + max_in_flight_tokens, use_eagle, enable_caching, enable_kv_cache_events, @@ -647,7 +655,7 @@ def get_kv_cache_coordinator( return HybridKVCacheCoordinator( kv_cache_config, max_model_len, - max_num_batched_tokens, + max_in_flight_tokens, use_eagle, enable_caching, enable_kv_cache_events, diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 9359d8843a..947d8cb0ad 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -8,10 +8,14 @@ from vllm.distributed.kv_events import BlockStored, KVCacheEvent from vllm.logger import init_logger +from vllm.utils.math_utils import cdiv from vllm.v1.core.kv_cache_coordinator import get_kv_cache_coordinator from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector from vllm.v1.core.kv_cache_utils import KVCacheBlock from vllm.v1.kv_cache_interface import ( + AttentionSpec, + CrossAttentionSpec, + EncoderOnlyAttentionSpec, KVCacheConfig, get_kv_cache_spec_kind, get_kv_cache_spec_sliding_window, @@ -113,7 +117,7 @@ def __init__( kv_cache_config: KVCacheConfig, max_model_len: int, hash_block_size: int, - max_num_batched_tokens: int | None = None, + max_in_flight_tokens: int | None = None, enable_caching: bool = True, use_eagle: bool = False, log_stats: bool = False, @@ -126,8 +130,8 @@ def __init__( # When unset, fall back to `max_model_len` so the recycling-aware cap # collapses to the prior (uncapped) admission behavior. The scheduler # always supplies the real value at runtime. - if max_num_batched_tokens is None: - max_num_batched_tokens = max_model_len + if max_in_flight_tokens is None: + max_in_flight_tokens = max_model_len self.enable_caching = enable_caching self.use_eagle = use_eagle @@ -141,7 +145,7 @@ def __init__( self.coordinator = get_kv_cache_coordinator( kv_cache_config=kv_cache_config, max_model_len=self.max_model_len, - max_num_batched_tokens=max_num_batched_tokens, + max_in_flight_tokens=max_in_flight_tokens, use_eagle=self.use_eagle, enable_caching=self.enable_caching, enable_kv_cache_events=enable_kv_cache_events, @@ -370,8 +374,13 @@ def allocate_slots( # insufficient free blocks. # Should call this function before allocating new blocks to reduce # the number of evicted blocks. + # Free on the processed-token basis: in-flight steps' attention windows + # still read blocks below the optimistic boundary, and rejected spec + # tokens can roll it back. self.coordinator.remove_skipped_blocks( - request.request_id, total_computed_tokens + request.request_id, + max(0, total_computed_tokens - request.num_in_flight_tokens), + num_prompt_tokens=request.num_prompt_tokens, ) num_blocks_to_allocate = self.coordinator.get_num_blocks_to_allocate( @@ -437,17 +446,24 @@ def free(self, request: Request) -> None: self.coordinator.free(request.request_id) def remove_skipped_blocks( - self, request_id: str, total_computed_tokens: int + self, + request_id: str, + processed_computed_tokens: int, + num_prompt_tokens: int | None = None, ) -> None: """Remove the blocks that are no longer needed from `blocks` and replace the removed blocks with null_block. Args: request_id: The request ID. - total_computed_tokens: The total number of computed tokens, including - local computed tokens and external computed tokens. + processed_computed_tokens: Computed-token prefix length covering + fully processed and committed tokens only (safe to free). + num_prompt_tokens: Optional prompt length for prefix-anchored SWA + gap eviction. """ - self.coordinator.remove_skipped_blocks(request_id, total_computed_tokens) + self.coordinator.remove_skipped_blocks( + request_id, processed_computed_tokens, num_prompt_tokens + ) def evict_blocks(self, block_ids: set[int]) -> None: """evict blocks from the prefix cache by their block IDs. @@ -541,6 +557,26 @@ def get_block_ids(self, request_id: str) -> tuple[list[int], ...]: """Get the block ids of a request.""" return self.get_blocks(request_id).get_block_ids() + def get_block_ids_for_computed_tokens( + self, + request_id: str, + num_computed_tokens: int, + ) -> tuple[list[int], ...]: + """Get block ids covering the request's computed tokens.""" + block_ids = self.get_block_ids(request_id) + clipped_block_ids: list[list[int]] = [] + for group, ids in zip(self.kv_cache_config.kv_cache_groups, block_ids): + spec = group.kv_cache_spec + if not isinstance(spec, AttentionSpec) or isinstance( + spec, (CrossAttentionSpec, EncoderOnlyAttentionSpec) + ): + clipped_block_ids.append(ids) + continue + + num_valid_blocks = cdiv(num_computed_tokens, spec.block_size) + clipped_block_ids.append(ids[:num_valid_blocks]) + return tuple(clipped_block_ids) + def cache_blocks(self, request: Request, num_computed_tokens: int) -> None: """Cache the blocks for the request, if enabled. @@ -565,6 +601,28 @@ def take_new_block_ids(self) -> list[int]: ids.extend(mgr.take_new_block_ids()) return ids + def take_boundary_state_offloads( + self, + ) -> dict[str, list[tuple[int, int, int]]]: + """Drain exact Mamba ``align`` state blocks for KV connectors. + + Returns ``{request_id: [(group_id, block_id, boundary_tokens), ...]}``. + These scheduler-local handoffs avoid reconstructing mutable Mamba block + tables from a connector's append-only block-ID mirror. + """ + offloads: dict[str, list[tuple[int, int, int]]] = {} + for mgr in self.coordinator.single_type_managers: + for ( + request_id, + group_id, + block, + boundary_tokens, + ) in mgr.take_pending_boundary_state_offloads(): + offloads.setdefault(request_id, []).append( + (group_id, block.block_id, boundary_tokens) + ) + return offloads + def new_step_starts(self) -> None: """Called when a new step is started.""" self.coordinator.new_step_starts() diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 7f3a5e4fdf..9f1f447340 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -611,12 +611,12 @@ def resolve_kv_cache_block_sizes( if not (cache_config.enable_prefix_caching or connector_enabled): return scheduler_block_size, scheduler_block_size - # Mamba groups with block_size != cache_config.block_size - # (mamba_cache_mode != "align") break divisibility; back off to the - # scheduler block size. + # Mamba groups outside align mode break divisibility; back off to the + # scheduler block size. Read the mode from the resolved group spec because + # its block size may have been updated independently of cache_config. if any( isinstance(g.kv_cache_spec, MambaSpec) - and g.kv_cache_spec.block_size != cache_config.block_size + and g.kv_cache_spec.mamba_cache_mode != "align" for g in groups ): return scheduler_block_size, scheduler_block_size @@ -1043,8 +1043,23 @@ def unify_kv_cache_spec_page_size( ) ratio = max_page_size // layer_page_size new_block_size = layer_spec.block_size * ratio - new_spec = replace(layer_spec, block_size=new_block_size) - assert new_spec.page_size_bytes == max_page_size + replace_args = {"block_size": new_block_size} + # A padded page does not grow when only block_size changes. This + # happens for hybrid Mamba targets when a higher-precision draft + # cache has a larger page than the target cache. Keep the logical + # block-size adjustment and grow the physical padding with it. + if ( + isinstance(layer_spec, MambaSpec) + or getattr(layer_spec, "page_size_padded", None) is not None + ): + replace_args["page_size_padded"] = max_page_size + new_spec = replace(layer_spec, **replace_args) + assert new_spec.page_size_bytes == max_page_size, ( + f"Failed to unify KV page for {layer_name}: " + f"spec={layer_spec!r}, old_page={layer_page_size}, " + f"target_page={max_page_size}, replacement={replace_args!r}, " + f"new_page={new_spec.page_size_bytes}" + ) new_kv_cache_spec[layer_name] = new_spec return new_kv_cache_spec @@ -1455,42 +1470,36 @@ def group_and_unify_kv_cache_specs( return [mla_uniform_spec, *swa_uniform_specs] -def _approximate_gcd(values: Sequence[int], *, lower_bound: int | None = None) -> int: - """Pick a chunk size that minimizes total upward padding. - - Each x is rounded up to a multiple of d: - - x -> ceil(x / d) * d - - Total padding is: - - pad(d) = sum_i (ceil(x_i / d) * d - x_i) - - We brute-force d in [lower_bound, max(values)] (fine for small lists / small - maxima) and return the d with minimum padding. Ties prefer larger d. - """ - if not values: - raise ValueError("values must be non-empty") - if any(x <= 0 for x in values): - raise ValueError(f"values must be positive, got: {list(values)!r}") - - min_d = max(1, lower_bound if lower_bound is not None else 1) - max_d = max(values) - if min_d > max_d: - return min_d - - best_d = min_d - best_pad: int | None = None - for d in range(min_d, max_d + 1): - pad = sum((d - (x % d)) % d for x in values) - if best_pad is None or pad < best_pad or (pad == best_pad and d > best_d): - best_pad = pad - best_d = d +def _select_deepseek_v4_tuple_width( + vllm_config: VllmConfig, + grouped_specs: list[UniformTypeKVCacheSpecs], +) -> int: + """Pick the DeepSeek V4 tuple width with the smallest physical pool. + + Padding layer counts alone is insufficient: full/compressed MLA may need + thousands more pages than a short sliding-window group at long context. + Widening every physical block to save one SWA subgroup can therefore cost + substantially more memory than it saves. + """ + tuple_counts = [spec.get_num_layer_tuples() for spec in grouped_specs] + page_counts = [spec.max_memory_usage_pages(vllm_config) for spec in grouped_specs] + min_width = tuple_counts[0] + max_width = max(tuple_counts) + + def pool_pages(width: int) -> int: + return width * sum( + cdiv(tuple_count, width) * page_count + for tuple_count, page_count in zip(tuple_counts, page_counts) + ) - return best_d + return min( + range(min_width, max_width + 1), + key=lambda width: (pool_pages(width), width), + ) def _get_kv_cache_groups_uniform_groups( + vllm_config: VllmConfig, grouped_specs: list[UniformTypeKVCacheSpecs], ) -> list[KVCacheGroupSpec]: """ @@ -1521,10 +1530,8 @@ def _get_kv_cache_groups_uniform_groups( num_layer_tuples_per_group: list[int] = [ g_spec.get_num_layer_tuples() for g_spec in grouped_specs ] - # Choose `num_layer_tuples` to minimize total padding across groups. - num_layer_tuples = _approximate_gcd( - num_layer_tuples_per_group, lower_bound=num_layer_tuples_per_group[0] - ) + # Choose `num_layer_tuples` by physical bytes, not layer-count padding. + num_layer_tuples = _select_deepseek_v4_tuple_width(vllm_config, grouped_specs) # Round up to the nearest multiple of `num_layer_tuples` (i.e., padding) num_layer_tuples_per_group = [ round_up(x, num_layer_tuples) for x in num_layer_tuples_per_group @@ -1651,7 +1658,9 @@ def get_kv_cache_groups( # yet some layers are full attention while others are sliding window # attention in different sizes. Need to group layers into multiple # UniformTypeKVCacheSpecs. - kv_cache_groups = _get_kv_cache_groups_uniform_groups(grouped_specs) + kv_cache_groups = _get_kv_cache_groups_uniform_groups( + vllm_config, grouped_specs + ) _annotate_eagle_groups_deepseek_v4(vllm_config, kv_cache_spec, kv_cache_groups) return kv_cache_groups diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py index c153d1103c..88da600d24 100644 --- a/vllm/v1/core/sched/output.py +++ b/vllm/v1/core/sched/output.py @@ -246,6 +246,11 @@ class SchedulerOutput: # preventing stale NaN/data from corrupting attention or SSM computation. new_block_ids_to_zero: list[int] | None = None + # Exact Mamba ``align`` boundary-state sources offered to a KV connector: + # {request_id: [(group_id, block_id, boundary_tokens), ...]}. + # Scheduler-local; cleared after connector metadata is built. + boundary_state_offloads: dict[str, list[tuple[int, int, int]]] | None = None + @classmethod def make_empty(cls) -> "SchedulerOutput": return cls( diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 7fa7757ea7..86b5a1aeff 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -247,25 +247,10 @@ def __init__( ) self.use_eagle = False - self.num_spec_tokens = self.num_lookahead_tokens = 0 - if speculative_config: - self.num_spec_tokens = speculative_config.num_speculative_tokens - if speculative_config.use_eagle(): - self.use_eagle = True - self.num_lookahead_tokens = self.num_spec_tokens - if speculative_config.uses_draft_model(): - self.num_lookahead_tokens = self.num_spec_tokens - if speculative_config.use_dflash(): - # DFlash requires an extra lookahead slot since it uses in-fill-style - # decoding instead of standard next-token sampling, so it has a query - # for the last sampled token plus queries for each draft token. - self.num_lookahead_tokens = ( - speculative_config.num_speculative_state_tokens() + 1 - ) - if speculative_config.use_dspark(): - # The anchor itself is the first prediction position, with no - # separate bonus query, so DSpark needs exactly N slots. - self.num_lookahead_tokens = self.num_spec_tokens + self.num_spec_tokens = vllm_config.num_speculative_tokens + self.num_lookahead_tokens = vllm_config.num_lookahead_tokens + if speculative_config and speculative_config.use_eagle_kv_cache(): + self.use_eagle = True # Create the KV cache manager. if hash_block_size is None: @@ -273,7 +258,7 @@ def __init__( self.kv_cache_manager = KVCacheManager( kv_cache_config=kv_cache_config, max_model_len=self.max_model_len, - max_num_batched_tokens=self.scheduler_config.max_num_batched_tokens, + max_in_flight_tokens=vllm_config.max_in_flight_tokens, enable_caching=self.cache_config.enable_prefix_caching, use_eagle=self.use_eagle, log_stats=self.log_stats, @@ -415,9 +400,6 @@ def _mamba_block_aligned_split( num_new_local_computed_tokens: int = 0, num_external_computed_tokens: int = 0, ) -> int: - assert num_external_computed_tokens == 0, ( - "External KV connector is not verified yet" - ) num_computed_tokens = ( request.num_computed_tokens + num_new_local_computed_tokens @@ -923,7 +905,8 @@ def schedule(self) -> SchedulerOutput: # The request cannot be scheduled. break - if self.need_mamba_block_aligned_split: + # Skip block alignment when setting up async receive (no local work). + if self.need_mamba_block_aligned_split and not load_kv_async: num_new_tokens = self._mamba_block_aligned_split( request, num_new_tokens, @@ -1120,6 +1103,9 @@ def schedule(self) -> SchedulerOutput: if self.needs_kv_cache_zeroing else None ) + # A Mamba ``align`` table is not append-only. Drain exact committed + # boundary-state block IDs before the connector builds store jobs. + boundary_state_offloads = self.kv_cache_manager.take_boundary_state_offloads() scheduler_output = SchedulerOutput( scheduled_new_reqs=new_reqs_data, @@ -1138,6 +1124,9 @@ def schedule(self) -> SchedulerOutput: finished_req_ids=self.finished_req_ids, free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(), new_block_ids_to_zero=new_block_ids_to_zero, + boundary_state_offloads=( + boundary_state_offloads if self.connector is not None else None + ), ) # NOTE(Kuntai): this function is designed for multiple purposes: @@ -1148,6 +1137,9 @@ def schedule(self) -> SchedulerOutput: meta = self._build_kv_connector_meta(self.connector, scheduler_output) scheduler_output.kv_connector_metadata = meta + # Exact source IDs are scheduler-side connector inputs only. + scheduler_output.boundary_state_offloads = None + # Build the connector meta for ECConnector if self.ec_connector is not None: ec_meta: ECConnectorMetadata = self.ec_connector.build_connector_meta( @@ -1200,6 +1192,7 @@ def _update_after_schedule(self, scheduler_output: SchedulerOutput) -> None: for req_id, num_scheduled_token in num_scheduled_tokens.items(): request = self.requests[req_id] request.num_computed_tokens += num_scheduled_token + request.num_in_flight_tokens += num_scheduled_token request.is_prefill_chunk = request.num_computed_tokens < ( request.num_tokens + request.num_output_placeholders ) @@ -1584,10 +1577,12 @@ def update_from_output( stopped_preempted_reqs: set[Request] = set() for req_id, num_tokens_scheduled in num_scheduled_tokens.items(): assert num_tokens_scheduled > 0 + request = self.requests.get(req_id) + if request is not None: + request.num_in_flight_tokens -= num_tokens_scheduled if failed_kv_load_req_ids and req_id in failed_kv_load_req_ids: # skip failed or rescheduled requests from KV load failure continue - request = self.requests.get(req_id) if request is None or request.is_finished(): # The request is already finished. This can happen if the # request is aborted while the model is executing it (e.g., @@ -2391,13 +2386,19 @@ def _connector_finished( return False, None # Free any out-of-window prefix blocks before we hand the block table to - # the connector. + # the connector, on the processed-token basis (see `allocate_slots`). self.kv_cache_manager.remove_skipped_blocks( request_id=request.request_id, - total_computed_tokens=request.num_computed_tokens, + processed_computed_tokens=max( + 0, request.num_computed_tokens - request.num_in_flight_tokens + ), + num_prompt_tokens=request.num_prompt_tokens, ) - block_ids = self.kv_cache_manager.get_block_ids(request.request_id) + block_ids = self.kv_cache_manager.get_block_ids_for_computed_tokens( + request_id=request.request_id, + num_computed_tokens=request.num_computed_tokens, + ) if not isinstance(self.connector, SupportsHMA): # NOTE(Kuntai): We should deprecate this code path after we enforce diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index cd000dc849..b3980c592f 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -20,6 +20,7 @@ KVCacheSpec, MambaSpec, MLAAttentionSpec, + PrefixAnchoredSWASpec, SinkFullAttentionSpec, SlidingWindowMLASpec, SlidingWindowSpec, @@ -68,6 +69,14 @@ def __init__( self._max_admission_blocks_per_request = max_admission_blocks_per_request self.new_block_ids: list[int] = [] + # Mamba ``align`` block tables are sparse and mutable: committed + # boundary states can be followed by speculative scratch blocks that + # are relocated in place. A connector therefore cannot recover an + # exact boundary-state source from an append-only block-ID mirror. + self._pending_boundary_state_offloads: list[ + tuple[str, int, KVCacheBlock, int] + ] = [] + # Mapping from request ID to blocks to track the blocks allocated # for each request, so that we can free the blocks when the request # is finished. @@ -82,6 +91,18 @@ def __init__( self.kv_cache_group_id = kv_cache_group_id self._null_block = block_pool.null_block + def take_pending_boundary_state_offloads( + self, + ) -> list[tuple[str, int, KVCacheBlock, int]]: + """Drain exact Mamba ``align`` boundary-state handoffs. + + Entries are ``(request_id, group_id, block, boundary_tokens)``. + Non-Mamba managers leave the list empty. + """ + pending = self._pending_boundary_state_offloads + self._pending_boundary_state_offloads = [] + return pending + @classmethod def _get_num_evictable_blocks(cls, blocks: Sequence[KVCacheBlock]): return sum(blk.ref_cnt == 0 and not blk.is_null for blk in blocks) @@ -425,8 +446,41 @@ def find_longest_cache_hit( raise NotImplementedError + def _remove_blocks_in_range( + self, + request_id: str, + first_block: int, + last_block: int, + ) -> None: + """Free blocks in ``[first_block, last_block)`` and replace with null_block. + + Iterates backward so newly-evictable tail blocks are reached even after + earlier blocks in the range were nulled in a prior call. + """ + if request_id not in self.req_to_blocks: + return + if first_block >= last_block: + return + blocks = self.req_to_blocks[request_id] + last_block = min(last_block, len(blocks)) + + freed: list[KVCacheBlock] = [] + for i in range(last_block - 1, first_block - 1, -1): + if blocks[i] == self._null_block: + # If the block is already a null block, the blocks before it + # should also have been set to null blocks by the previous calls + # to this function. + break + freed.append(blocks[i]) + blocks[i] = self._null_block + if freed: + self.block_pool.free_blocks(freed) + def remove_skipped_blocks( - self, request_id: str, total_computed_tokens: int + self, + request_id: str, + processed_computed_tokens: int, + num_prompt_tokens: int | None = None, ) -> None: """ Remove and free the blocks that are no longer needed for attention computation. @@ -437,36 +491,27 @@ def remove_skipped_blocks( Args: request_id: The request ID. - total_computed_tokens: The total number of computed tokens, including - local computed tokens and external computed tokens. + processed_computed_tokens: Computed-token prefix length covering + fully processed and committed tokens only (safe to free). + num_prompt_tokens: Optional prompt length for attention types (e.g. + prefix-anchored SWA) that evict a middle gap rather than a head + prefix. Ignored by the default implementation. """ + del num_prompt_tokens # Remove the blocks that will be skipped during attention computation. - num_skipped_tokens = self.get_num_skipped_tokens(total_computed_tokens) + num_skipped_tokens = self.get_num_skipped_tokens(processed_computed_tokens) if num_skipped_tokens <= 0: # This indicates that ALL tokens are inside attention window. # Thus we do not need to free any blocks outside attention window. # A typical case is full attention that we never free any token # before the request is finished. return - blocks = self.req_to_blocks[request_id] num_skipped_blocks = num_skipped_tokens // self.block_size # `num_skipped_tokens` may include tokens that haven't been allocated yet # (e.g., when the attention window moves into the external computed tokens - # range), so we must cap to the number of blocks that currently exist for - # this request. - num_skipped_blocks = min(num_skipped_blocks, len(blocks)) - removed_blocks: list[KVCacheBlock] = [] - # Because the block starts from index 0, the num_skipped_block-th block - # corresponds to index num_skipped_blocks - 1. - for i in range(num_skipped_blocks - 1, -1, -1): - if blocks[i] == self._null_block: - # If the block is already a null block, the blocks before it - # should also have been set to null blocks by the previous calls - # to this function. - break - removed_blocks.append(blocks[i]) - blocks[i] = self._null_block - self.block_pool.free_blocks(removed_blocks) + # range); `_remove_blocks_in_range` caps to the number of blocks that + # currently exist for this request. + self._remove_blocks_in_range(request_id, 0, num_skipped_blocks) def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: """ @@ -547,6 +592,53 @@ def get_num_common_prefix_blocks(self, running_request_id: str) -> int: return num_common_blocks +class PrefixAnchoredSWAManager(FullAttentionManager): + """KV cache manager for prefix-anchored sliding-window attention. + + When ``num_prompt_tokens`` is supplied to ``remove_skipped_blocks``, frees + gap blocks between the prefill tail and the current decode window. This + bounds per-request KV memory at O(prefix_len + decode_sliding_window) + instead of growing linearly with decode length. + """ + + def __init__(self, kv_cache_spec: PrefixAnchoredSWASpec, **kwargs) -> None: + super().__init__(kv_cache_spec, **kwargs) + self.decode_sliding_window: int = kv_cache_spec.decode_sliding_window + + def remove_skipped_blocks( + self, + request_id: str, + processed_computed_tokens: int, + num_prompt_tokens: int | None = None, + ) -> None: + """Free gap blocks that are no longer needed for attention. + + Gap = blocks entirely within + [ceil(prefix_len / block_size) * block_size, + max(prefix_len, processed_computed_tokens - decode_sliding_window)) + + Freed blocks are replaced with null_block in req_to_blocks so the + block_table handed to the attention backend stays valid (null_block + KV is all-zero; the backend mask marks gap positions as non-visible + so they never contribute to the output). + """ + if num_prompt_tokens is None: + super().remove_skipped_blocks( + request_id, processed_computed_tokens, num_prompt_tokens + ) + return + + bs = self.block_size + # First block fully after the prefill boundary. + first_gap_block = cdiv(num_prompt_tokens, bs) + # Decode window start position; blocks before this are evictable. + window_start = max( + num_prompt_tokens, processed_computed_tokens - self.decode_sliding_window + ) + last_gap_block = window_start // bs # exclusive upper bound + self._remove_blocks_in_range(request_id, first_gap_block, last_gap_block) + + class SlidingWindowManager(SingleTypeKVCacheManager): def __init__(self, kv_cache_spec: SlidingWindowSpec, **kwargs) -> None: super().__init__(kv_cache_spec, **kwargs) @@ -910,17 +1002,17 @@ def find_longest_cache_hit( return computed_blocks - def remove_skipped_blocks(self, request_id: str, num_computed_tokens: int) -> None: + def remove_skipped_blocks( + self, + request_id: str, + processed_computed_tokens: int, + num_prompt_tokens: int | None = None, + ) -> None: assert isinstance(self.kv_cache_spec, MambaSpec) - # NOTE (tdoublep) with async scheduling, the num_computed_tokens can contain - # draft tokens from the previous step that may or may not be rejected later. - # This can make us think we are further ahead in the sequence than we actually - # are, so let's assume that all tokens are rejected so we don't free blocks - # that we might actually need. - num_computed_tokens = max(0, num_computed_tokens - self.num_speculative_blocks) - - super().remove_skipped_blocks(request_id, num_computed_tokens) + super().remove_skipped_blocks( + request_id, processed_computed_tokens, num_prompt_tokens + ) if self.mamba_cache_mode == "align": # `last_state_block_idx` refers to the block index allocated two steps ago. # The block allocated in the previous step is used to copy Mamba states @@ -933,7 +1025,7 @@ def remove_skipped_blocks(self, request_id: str, num_computed_tokens: int) -> No if ( last_state_block_idx is not None and last_state_block_idx - < cdiv(num_computed_tokens, self.block_size) - 1 + < cdiv(processed_computed_tokens, self.block_size) - 1 ): blocks = self.req_to_blocks[request_id] if blocks[last_state_block_idx] != self._null_block: @@ -1095,6 +1187,14 @@ def free(self, request_id: str) -> None: if self.mamba_cache_mode == "align": self._allocated_block_reqs.discard(request_id) self.last_state_block_idx.pop(request_id, None) + # A hand-off is valid only while its source still belongs to this + # request. Drop offers that have not reached the connector before + # returning the request's blocks to the pool. + self._pending_boundary_state_offloads = [ + entry + for entry in self._pending_boundary_state_offloads + if entry[0] != request_id + ] super().free(request_id) def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: @@ -1115,13 +1215,22 @@ def cache_blocks( super().cache_blocks(request, num_tokens, alignment_tokens=alignment_tokens) num_cached_blocks_after = self.num_cached_block.get(request.request_id, 0) if num_cached_blocks_after > num_cached_blocks_before: - for block in self.req_to_blocks[request.request_id][ - num_cached_blocks_before:num_cached_blocks_after - ]: + blocks = self.req_to_blocks[request.request_id] + for block_idx in range(num_cached_blocks_before, num_cached_blocks_after): + block = blocks[block_idx] if block.is_null: continue assert block.block_hash is not None self.cached_blocks_this_step.add(block.block_hash) + if self.mamba_cache_mode == "align": + self._pending_boundary_state_offloads.append( + ( + request.request_id, + self.kv_cache_group_id, + block, + (block_idx + 1) * self.block_size, + ) + ) def new_step_starts(self) -> None: self.cached_blocks_this_step.clear() @@ -1210,6 +1319,7 @@ def __init__( TQFullAttentionSpec: FullAttentionManager, MLAAttentionSpec: FullAttentionManager, HiddenStateCacheSpec: FullAttentionManager, + PrefixAnchoredSWASpec: PrefixAnchoredSWAManager, SlidingWindowSpec: SlidingWindowManager, SlidingWindowMLASpec: SlidingWindowManager, ChunkedLocalAttentionSpec: ChunkedLocalAttentionManager, @@ -1221,7 +1331,7 @@ def __init__( def get_manager_for_kv_cache_spec( kv_cache_spec: KVCacheSpec, - max_num_batched_tokens: int, + max_in_flight_tokens: int, max_model_len: int, **kwargs, ) -> SingleTypeKVCacheManager: @@ -1232,7 +1342,7 @@ def get_manager_for_kv_cache_spec( if isinstance(kv_cache_spec, (SlidingWindowSpec, ChunkedLocalAttentionSpec)): kwargs["max_admission_blocks_per_request"] = ( kv_cache_spec.max_admission_blocks_per_request( - max_num_batched_tokens=max_num_batched_tokens, + max_in_flight_tokens=max_in_flight_tokens, max_model_len=max_model_len, ) ) diff --git a/vllm/v1/cudagraph_dispatcher.py b/vllm/v1/cudagraph_dispatcher.py index 862c3f2855..b50f362594 100644 --- a/vllm/v1/cudagraph_dispatcher.py +++ b/vllm/v1/cudagraph_dispatcher.py @@ -7,13 +7,21 @@ from vllm import envs from vllm.config import CUDAGraphMode, VllmConfig -from vllm.forward_context import BatchDescriptor +from vllm.forward_context import ( + CUDAGRAPH_VARIANT_DEFAULT, + CUDAGRAPH_VARIANT_LONG_CONTEXT, + BatchDescriptor, +) from vllm.logger import init_logger from vllm.lora.utils import get_captured_lora_counts from vllm.platforms import current_platform logger = init_logger(__name__) +_SM70_FP8_KV_BATCH_CONTEXT_SIZES = frozenset((4, 8, 16)) +_SM70_FP8_KV_LONG_CONTEXT_MIN_SEQ_LEN = 16384 +_SM70_E4M3_B1_WAVE_LONG_CONTEXT_MIN_SEQ_LEN = 49152 + def _get_sm70_context_buckets(env_name: str) -> tuple[int, ...]: """Parse optional one-request attention graph context buckets.""" @@ -39,15 +47,13 @@ def _get_sm70_dsv4_decode_context_buckets( env_name = "VLLM_SM70_DSV4_DECODE_CONTEXT_BUCKETS" if env_name in os.environ: return _get_sm70_context_buckets(env_name) + # Speculative graphs cover multiple query rows and require a separate + # end-to-end long-context gate. They remain available through the explicit + # environment override above. if vllm_config.speculative_config is not None: return () model_config = vllm_config.model_config - architectures = getattr(model_config, "architectures", ()) - if not isinstance(architectures, (list, tuple)) or ( - "DeepseekV4ForCausalLM" not in architectures - ): - return () if not ( current_platform.is_cuda() and current_platform.is_device_capability((7, 0)) ): @@ -66,64 +72,88 @@ def _get_sm70_dsv4_decode_context_buckets( if not positive_ratios: return () - bucket = topk_tokens * min(positive_ratios) - if bucket <= 0 or model_config.max_model_len <= bucket: + short_bucket = topk_tokens * min(positive_ratios) + if short_bucket <= 0 or model_config.max_model_len <= short_bucket: return () + # Keep the graph width close enough to the live context that the C4 + # indexer does not score the model's entire maximum length immediately + # after crossing the short-context bypass. The sparse progression limits + # startup graph count while covering the first long-indexer bucket and the + # 16K/64K/128K service points used by the long-context quality gates. + bucket_multipliers = (1, 2, 8, 32, 64) + buckets = tuple( + short_bucket * multiplier + for multiplier in bucket_multipliers + if short_bucket * multiplier < model_config.max_model_len + ) logger.info_once( - "Auto-enabling the SM70 DeepSeek V4 short-context decode CUDA graph " - "at %d tokens. Set %s explicitly to override or to an empty value " - "to disable.", - bucket, + "Auto-enabling SM70 compressed-index decode CUDA graph context buckets " + "%s. Set %s explicitly to override or to an empty value to disable.", + buckets, env_name, ) - return (bucket,) + return buckets -def _get_sm70_fp8_kv_decode_context_buckets( +def _is_sm70_fp8_kv_decode_shape( vllm_config: VllmConfig, -) -> tuple[int, ...]: - """Select a scalar-only short graph for the SM70 E5M2 D256 route.""" - env_name = "VLLM_SM70_FP8_KV_DECODE_CONTEXT_BUCKETS" - if env_name in os.environ: - return _get_sm70_context_buckets(env_name) + *, + q_per_kv_values: tuple[int, ...], + cache_dtypes: tuple[str, ...] = ("fp8_e5m2",), + backend_names: tuple[str, ...] = ( + "FLASH_ATTN_V100", + "FLASHINFER_SM70", + ), +) -> bool: if vllm_config.speculative_config is not None: - return () + return False if not ( current_platform.is_cuda() and current_platform.is_device_capability((7, 0)) ): - return () + return False if not envs.VLLM_SM70_FLASH_ATTN_V100: - return () + return False cache_config = getattr(vllm_config, "cache_config", None) - if getattr(cache_config, "cache_dtype", None) != "fp8_e5m2": - return () + if getattr(cache_config, "cache_dtype", None) not in cache_dtypes: + return False attention_config = getattr(vllm_config, "attention_config", None) attention_backend = getattr(attention_config, "backend", None) attention_backend_name = getattr(attention_backend, "name", attention_backend) - if attention_backend is not None and attention_backend_name not in ( - "FLASH_ATTN_V100", - "FLASHINFER_SM70", - ): - return () + if attention_backend is not None and attention_backend_name not in backend_names: + return False model_config = vllm_config.model_config hf_text_config = getattr(model_config, "hf_text_config", None) num_attention_heads = getattr(hf_text_config, "num_attention_heads", None) num_key_value_heads = getattr(hf_text_config, "num_key_value_heads", None) head_dim = getattr(hf_text_config, "head_dim", None) - if not ( + return bool( isinstance(num_attention_heads, int) and isinstance(num_key_value_heads, int) and num_key_value_heads > 0 and num_attention_heads % num_key_value_heads == 0 - and num_attention_heads // num_key_value_heads in (6, 8) + and num_attention_heads // num_key_value_heads in q_per_kv_values and head_dim == 256 + ) + + +def _get_sm70_fp8_kv_decode_context_buckets( + vllm_config: VllmConfig, +) -> tuple[int, ...]: + """Select a scalar-only B1 graph for the SM70 E5M2 D256 route.""" + env_name = "VLLM_SM70_FP8_KV_DECODE_CONTEXT_BUCKETS" + if env_name in os.environ: + return _get_sm70_context_buckets(env_name) + if not _is_sm70_fp8_kv_decode_shape( + vllm_config, + q_per_kv_values=(6, 8), ): return () bucket = 8192 + model_config = vllm_config.model_config max_model_len = getattr(model_config, "max_model_len", None) if not isinstance(max_model_len, int) or max_model_len <= bucket: return () @@ -138,6 +168,35 @@ def _get_sm70_fp8_kv_decode_context_buckets( return (bucket,) +def _sm70_fp8_kv_batch_context_routing_enabled( + vllm_config: VllmConfig, +) -> bool: + if ( + not envs.VLLM_FLASH_V100_XQA_BATCH_CONTEXT_ROUTING + or envs.VLLM_FLASH_V100_DECODE_PARTITION_SIZE is not None + ): + return False + cache_dtype = getattr(vllm_config.cache_config, "cache_dtype", None) + cache_dtypes: tuple[str, ...] + backend_names: tuple[str, ...] + if cache_dtype in ("fp8", "fp8_e4m3"): + if not envs.VLLM_FLASH_V100_E4M3_BATCH_XQA: + return False + cache_dtypes = ("fp8", "fp8_e4m3") + # E4M3 batch-context variants exist to select Flash-V100 XQA launch + # routes. An explicit FlashInfer backend cannot consume that metadata. + backend_names = ("FLASH_ATTN_V100",) + else: + cache_dtypes = ("fp8_e5m2",) + backend_names = ("FLASH_ATTN_V100", "FLASHINFER_SM70") + return _is_sm70_fp8_kv_decode_shape( + vllm_config, + q_per_kv_values=(6,), + cache_dtypes=cache_dtypes, + backend_names=backend_names, + ) + + class CudagraphDispatcher: """ Runtime cudagraph dispatcher to dispatch keys for multiple set of @@ -174,7 +233,27 @@ def __init__(self, vllm_config: VllmConfig): self.sm70_fp8_kv_decode_context_buckets = ( _get_sm70_fp8_kv_decode_context_buckets(vllm_config) ) + self.sm70_fp8_kv_batch_context_routing = ( + _sm70_fp8_kv_batch_context_routing_enabled(vllm_config) + ) + cache_dtype = getattr(vllm_config.cache_config, "cache_dtype", None) + self.sm70_e4m3_b1_wave_context_routing = bool( + self.sm70_fp8_kv_batch_context_routing + and cache_dtype in ("fp8", "fp8_e4m3") + and os.getenv("VLLM_FLASH_V100_XQA_E4M3_G6_P64_P256_AUTO", "0") == "1" + and os.getenv("VLLM_FLASH_V100_XQA_E4M3_G6_WAVE_PARTITIONS", "0") == "1" + ) + self.sm70_e4m3_b1_wave_context_min_seq_len = max( + 1, + int( + os.getenv( + "VLLM_FLASH_V100_XQA_E4M3_G6_P512_BEGIN", + str(_SM70_E4M3_B1_WAVE_LONG_CONTEXT_MIN_SEQ_LEN), + ) + ), + ) self._logged_sm70_context_bucket = False + self._logged_sm70_batch_context_variant = False # Dict to store valid cudagraph dispatching keys. self.cudagraph_keys: dict[CUDAGraphMode, set[BatchDescriptor]] = { @@ -309,7 +388,9 @@ def add_cudagraph_key( def _active_context_buckets(self) -> tuple[int, ...]: if self.uniform_decode_query_len > 1: - return self.sm70_mtp_context_buckets + if "VLLM_SM70_MTP_CONTEXT_BUCKETS" in os.environ: + return self.sm70_mtp_context_buckets + return self.sm70_dsv4_decode_context_buckets return tuple( sorted( set(self.sm70_dsv4_decode_context_buckets) @@ -317,39 +398,102 @@ def _active_context_buckets(self) -> tuple[int, ...]: ) ) + def _is_batch_context_variant_descriptor( + self, + batch_descriptor: BatchDescriptor, + ) -> bool: + return bool( + self.sm70_fp8_kv_batch_context_routing + and batch_descriptor.uniform + and batch_descriptor.num_reqs is not None + and batch_descriptor.num_tokens == batch_descriptor.num_reqs + and ( + batch_descriptor.num_reqs in _SM70_FP8_KV_BATCH_CONTEXT_SIZES + or ( + self.sm70_e4m3_b1_wave_context_routing + and batch_descriptor.num_reqs == 1 + ) + ) + ) + + def _batch_context_variant_min_seq_len( + self, + batch_descriptor: BatchDescriptor, + ) -> int: + if self.sm70_e4m3_b1_wave_context_routing and batch_descriptor.num_reqs == 1: + return self.sm70_e4m3_b1_wave_context_min_seq_len + return _SM70_FP8_KV_LONG_CONTEXT_MIN_SEQ_LEN + + def _context_buckets_for_descriptor( + self, + batch_descriptor: BatchDescriptor, + ) -> tuple[int, ...]: + if not batch_descriptor.uniform or batch_descriptor.num_reqs is None: + return () + + if self.uniform_decode_query_len > 1: + if batch_descriptor.num_tokens == self.uniform_decode_query_len: + if "VLLM_SM70_MTP_CONTEXT_BUCKETS" in os.environ: + return self.sm70_mtp_context_buckets + return self.sm70_dsv4_decode_context_buckets + return () + + buckets: set[int] = set() + if batch_descriptor.num_tokens == 1: + buckets.update(self.sm70_dsv4_decode_context_buckets) + buckets.update(self.sm70_fp8_kv_decode_context_buckets) + + return tuple(sorted(buckets)) + @property def has_attention_context_buckets(self) -> bool: return bool(self._active_context_buckets()) - def _add_context_bucket_keys(self, batch_descriptor: BatchDescriptor) -> None: - """Add bounded graphs only for a single uniform-decode request.""" - context_buckets = self._active_context_buckets() - if ( - not context_buckets - or not batch_descriptor.uniform - or batch_descriptor.num_tokens != self.uniform_decode_query_len - ): - return + @property + def has_attention_context_specialization(self) -> bool: + return bool( + self.has_attention_context_buckets or self.sm70_fp8_kv_batch_context_routing + ) + def _add_context_bucket_keys(self, batch_descriptor: BatchDescriptor) -> None: + """Add attention-specialized graphs for this exact descriptor.""" + context_buckets = self._context_buckets_for_descriptor(batch_descriptor) for bucket in context_buckets: self.add_cudagraph_key( CUDAGraphMode.FULL, replace(batch_descriptor, attention_context_bucket=bucket), ) + if self._is_batch_context_variant_descriptor(batch_descriptor): + self.add_cudagraph_key( + CUDAGraphMode.FULL, + replace( + batch_descriptor, + graph_variant=CUDAGRAPH_VARIANT_LONG_CONTEXT, + ), + ) def _get_context_bucket_descriptor( self, batch_descriptor: BatchDescriptor, attention_context_len: int | None, ) -> BatchDescriptor: - """Return a bounded graph key only when its capacity is sufficient.""" - context_buckets = self._active_context_buckets() + """Return the graph specialization selected by the active context.""" + if ( + attention_context_len is not None + and attention_context_len + >= self._batch_context_variant_min_seq_len(batch_descriptor) + and self._is_batch_context_variant_descriptor(batch_descriptor) + ): + return replace( + batch_descriptor, + graph_variant=CUDAGRAPH_VARIANT_LONG_CONTEXT, + ) + + context_buckets = self._context_buckets_for_descriptor(batch_descriptor) if ( attention_context_len is None or attention_context_len <= 0 or not context_buckets - or not batch_descriptor.uniform - or batch_descriptor.num_tokens != self.uniform_decode_query_len ): return batch_descriptor @@ -537,6 +681,17 @@ def dispatch( batch_desc_to_check, ) self._logged_sm70_context_bucket = True + if ( + batch_desc_to_check.graph_variant == CUDAGRAPH_VARIANT_LONG_CONTEXT + and not self._logged_sm70_batch_context_variant + ): + logger.info( + "Using SM70 long-context XQA CUDA graph: " + "context=%d descriptor=%s", + attention_context_len, + batch_desc_to_check, + ) + self._logged_sm70_batch_context_variant = True return CUDAGraphMode.FULL, batch_desc_to_check # A bounded graph might be absent because capture-size policy # excluded it. Preserve the existing full-context graph then. @@ -575,16 +730,32 @@ def get_capture_descs(self) -> list[tuple[CUDAGraphMode, list[BatchDescriptor]]] if descs: # Capture the generic full-context graph before smaller # context-specialized graphs so its allocations seed the pool. - descs.sort( - key=lambda d: ( - d.num_tokens, - d.attention_context_bucket - if d.attention_context_bucket is not None - else float("inf"), - d.num_active_loras, - ), - reverse=True, - ) + if self.sm70_fp8_kv_batch_context_routing: + # Preserve the rollback graph capture order before adding + # long-context variants. This keeps baseline graph memory + # allocation and warmup behavior reproducible. + descs.sort( + key=lambda d: ( + d.graph_variant == CUDAGRAPH_VARIANT_DEFAULT, + d.num_tokens, + d.attention_context_bucket + if d.attention_context_bucket is not None + else float("inf"), + d.num_active_loras, + ), + reverse=True, + ) + else: + descs.sort( + key=lambda d: ( + d.num_tokens, + d.attention_context_bucket + if d.attention_context_bucket is not None + else float("inf"), + d.num_active_loras, + ), + reverse=True, + ) result.append((mode, descs)) return result diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 4e094eb6a6..c1f11a69ba 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -461,9 +461,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 +476,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 +493,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 +605,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 31ee89bc72..83f32dfccf 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -403,12 +403,58 @@ class HiddenStateCacheSpec(MLAAttentionSpec): pass +@dataclass(frozen=True, kw_only=True) +class PrefixAnchoredSWASpec(FullAttentionSpec): + """KV cache spec for prefix-anchored sliding-window attention. + + Prefill (prompt/prefix) tokens are always globally visible. + Only the last ``decode_sliding_window`` generated tokens are kept in the + KV cache; gap blocks (between the prefill tail and the current decode + window) are evicted during each decode step to bound per-request memory + at O(prefix_blocks + window_blocks). + """ + + decode_sliding_window: int + + def __post_init__(self) -> None: + super().__post_init__() + if self.decode_sliding_window <= 0: + raise ValueError("decode_sliding_window must be greater than zero") + + @classmethod + def merge(cls, specs: list[Self]) -> Self: + assert all(isinstance(spec, PrefixAnchoredSWASpec) for spec in specs), ( + "All attention layers in the same KV cache group must be " + "PrefixAnchoredSWASpec." + ) + windows = {spec.decode_sliding_window for spec in specs} + assert len(windows) == 1, ( + "All prefix-anchored SWA layers must share the same " + f"decode_sliding_window, got {windows}" + ) + # Delegate common field merging to the parent, then reattach the + # decode window. + base = FullAttentionSpec.merge(specs) # type: ignore[arg-type] + return cls( + block_size=base.block_size, + num_kv_heads=base.num_kv_heads, + head_size=base.head_size, + head_size_v=base.head_size_v, + dtype=base.dtype, + kv_quant_mode=base.kv_quant_mode, + page_size_padded=base.page_size_padded, + sliding_window=base.sliding_window, + attention_chunk_size=base.attention_chunk_size, + decode_sliding_window=windows.pop(), + ) + + @dataclass(frozen=True, kw_only=True) class ChunkedLocalAttentionSpec(AttentionSpec): attention_chunk_size: int def max_admission_blocks_per_request( - self, max_num_batched_tokens: int, max_model_len: int + self, max_in_flight_tokens: int, max_model_len: int ) -> int: """Per-request admission cap, in blocks. @@ -418,15 +464,15 @@ def max_admission_blocks_per_request( """ # During chunked prefill, we hold KV for at most one chunk window. num_tokens = min( - self.attention_chunk_size + max_num_batched_tokens, max_model_len + self.attention_chunk_size + max_in_flight_tokens, max_model_len ) return cdiv(num_tokens, self.block_size) def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: max_model_len = vllm_config.model_config.max_model_len - max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens + max_in_flight_tokens = vllm_config.max_in_flight_tokens max_blocks = self.max_admission_blocks_per_request( - max_num_batched_tokens=max_num_batched_tokens, max_model_len=max_model_len + max_in_flight_tokens=max_in_flight_tokens, max_model_len=max_model_len ) return max_blocks * self.page_size_bytes @@ -461,7 +507,7 @@ def real_page_size_bytes(self) -> int: ) def max_admission_blocks_per_request( - self, max_num_batched_tokens: int, max_model_len: int + self, max_in_flight_tokens: int, max_model_len: int ) -> int: """Per-request admission cap, in blocks. @@ -474,9 +520,7 @@ def max_admission_blocks_per_request( # During chunked prefill, we hold KV for the last `sliding_window-1` # computed tokens plus the newly scheduled tokens, and never more # than `max_model_len`. - num_tokens = min( - self.sliding_window - 1 + max_num_batched_tokens, max_model_len - ) + num_tokens = min(self.sliding_window - 1 + max_in_flight_tokens, max_model_len) # +1 because the sliding window may not start from the beginning of # the block. E.g. block size 4 and num_token 4 needs two blocks # [XXCD][EF] to store the 6-token window [CDEF]. @@ -487,9 +531,9 @@ def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: "DCP not support sliding window." ) max_model_len = vllm_config.model_config.max_model_len - max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens + max_in_flight_tokens = vllm_config.max_in_flight_tokens max_blocks = self.max_admission_blocks_per_request( - max_num_batched_tokens=max_num_batched_tokens, max_model_len=max_model_len + max_in_flight_tokens=max_in_flight_tokens, max_model_len=max_model_len ) return max_blocks * self.page_size_bytes diff --git a/vllm/v1/request.py b/vllm/v1/request.py index 26cc82fc4a..095ee6e975 100644 --- a/vllm/v1/request.py +++ b/vllm/v1/request.py @@ -141,6 +141,11 @@ def __init__( self.num_output_placeholders = 0 self.async_tokens_to_discard = 0 + # Tokens of steps whose output is not yet processed (async scheduling + # and PP run ahead of the GPU); `num_computed_tokens` counts them + # optimistically. + self.num_in_flight_tokens = 0 + self.spec_token_ids: list[int] = [] self.num_computed_tokens = 0 self.cache_salt: str | None = cache_salt diff --git a/vllm/v1/sample/ops/topk_topp_triton.py b/vllm/v1/sample/ops/topk_topp_triton.py index 4b17831dfc..55965bed40 100644 --- a/vllm/v1/sample/ops/topk_topp_triton.py +++ b/vllm/v1/sample/ops/topk_topp_triton.py @@ -11,6 +11,8 @@ import torch +from vllm import envs +from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.utils.math_utils import next_power_of_2 from vllm.utils.platform_utils import num_compute_units @@ -18,6 +20,29 @@ _TRITON_TABLE_CACHE: dict[tuple[torch.device], tuple[torch.Tensor, torch.Tensor]] = {} _TRITON_BUFFER_CACHE: dict[tuple[torch.device, torch.dtype, int], torch.Tensor] = {} + +def _use_sm70_topk_topp_8_warps( + device: torch.device, + batch_size: int, + vocab_size: int, + topk_enabled: bool, + topp_enabled: bool, +) -> bool: + if ( + device.type != "cuda" + or not current_platform.is_cuda() + or not current_platform.is_device_capability((7, 0)) + or vocab_size != 248_320 + or not topk_enabled + or not topp_enabled + ): + return False + + if batch_size in (8, 16) and envs.VLLM_SM70_TOPK_TOPP_B8_B16_8_WARPS: + return True + return envs.VLLM_SM70_TOPK_TOPP_8_WARPS and batch_size in (5, 10, 20, 40, 60, 80) + + # fmt: off _NORMAL_CDF_TO_SIGMA_TABLE = [ 3.656, 3.650, 3.650, 3.650, 3.626, 3.626, 3.626, 3.514, 3.514, 3.503, @@ -1044,6 +1069,19 @@ def apply_top_k_top_p_triton( else: block_size, block_size_trunc = 8192, 4096 + launch_kwargs: dict[str, int] = {} + if _use_sm70_topk_topp_8_warps( + logits.device, + batch_size, + vocab_size, + topk_enabled, + topp_enabled, + ): + # Eight warps preserve the tile and masking algorithm while using the + # otherwise idle lanes in its reduction-heavy passes. B8/B16 no-MTP + # rows are default-on; the measured MTP verifier rows remain opt-in. + launch_kwargs["num_warps"] = 8 + _topk_topp_kernel[(NUM_PROGRAMS,)]( logits, logits.stride(0), @@ -1059,6 +1097,7 @@ def apply_top_k_top_p_triton( BLOCK_SIZE_TRUNC=block_size_trunc, TOPK_ENABLED=topk_enabled, TOPP_ENABLED=topp_enabled, + **launch_kwargs, ) return logits diff --git a/vllm/v1/sample/rejection_sampler.py b/vllm/v1/sample/rejection_sampler.py index ba05ef5db5..7933896f7c 100644 --- a/vllm/v1/sample/rejection_sampler.py +++ b/vllm/v1/sample/rejection_sampler.py @@ -270,6 +270,7 @@ def forward( # [num_tokens + batch_size, vocab_size] logits: torch.Tensor, sampling_metadata: SamplingMetadata, + draft_confidence_logits: torch.Tensor | None = None, ) -> SamplerOutput: """ Args: @@ -331,6 +332,7 @@ def forward( bonus_logits_indices=bonus_logits_indices, profile_events=profile_events, profile_total_start=profile_total_start, + draft_confidence_logits=draft_confidence_logits, ) # When indexing with a tensor (bonus_logits_indices), PyTorch # creates a new tensor with separate storage from the original @@ -426,6 +428,7 @@ def forward( bonus_token_ids=bonus_token_ids, output_token_ids=output_token_ids, sampling_metadata=sampling_metadata, + draft_confidence_logits=draft_confidence_logits, ) logprobs_tensors = None @@ -455,6 +458,7 @@ def _forward_combined_bonus( bonus_logits_indices: torch.Tensor, profile_events: list[tuple[str, torch.cuda.Event, torch.cuda.Event]] | None, profile_total_start: torch.cuda.Event | None, + draft_confidence_logits: torch.Tensor | None, ) -> SamplerOutput: profile_constraints_start = _rejection_profile_start(profile_events) sampled_logits = logits.to(torch.float32) @@ -510,6 +514,7 @@ def _forward_combined_bonus( bonus_token_ids=bonus_token_ids, output_token_ids=output_token_ids, sampling_metadata=sampling_metadata, + draft_confidence_logits=draft_confidence_logits, ) return SamplerOutput( sampled_token_ids=output_token_ids, @@ -574,6 +579,7 @@ def _maybe_dump_alignment( bonus_token_ids: torch.Tensor, output_token_ids: torch.Tensor, sampling_metadata: SamplingMetadata, + draft_confidence_logits: torch.Tensor | None, ) -> None: global _SPEC_ALIGNMENT_DUMP_COUNT, _SPEC_ALIGNMENT_STEP_COUNTER _SPEC_ALIGNMENT_STEP_COUNTER += 1 @@ -651,6 +657,9 @@ def _maybe_dump_alignment( "sampling_output_token_ids_tail": [ list(ids[-32:]) for ids in sampling_metadata.output_token_ids ], + "sampling_output_lengths": [ + len(ids) for ids in sampling_metadata.output_token_ids + ], "sampling_spec_token_ids": [ list(ids) for ids in (sampling_metadata.spec_token_ids or []) ], @@ -672,9 +681,23 @@ def _maybe_dump_alignment( .cpu() ) payload["recovered_residual_mass"] = residual_mass.detach().cpu() + payload["conditional_acceptance_targets"] = ( + torch.minimum(target_probs, draft_probs).sum(dim=-1).detach().cpu() + ) draft_topk = torch.topk(draft_probs, k=k, dim=-1) payload["draft_topk_ids"] = draft_topk.indices.detach().cpu() payload["draft_topk_values"] = draft_topk.values.detach().cpu() + if draft_confidence_logits is not None: + if draft_confidence_logits.shape != (target_logits.shape[0],): + raise RuntimeError( + "Aligned DSpark confidence logits must have one value per " + "draft token: " + f"confidence={tuple(draft_confidence_logits.shape)}, " + f"target={target_logits.shape[0]}." + ) + payload["draft_confidence_logits"] = ( + draft_confidence_logits.detach().float().cpu() + ) _SPEC_ALIGNMENT_DUMP_COUNT += 1 dump_dir = os.getenv("VLLM_SPEC_DUMP_ALIGNMENT_DIR", "/tmp") os.makedirs(dump_dir, exist_ok=True) diff --git a/vllm/v1/sample/sampler.py b/vllm/v1/sample/sampler.py index 9f565ce458..8d960ed297 100644 --- a/vllm/v1/sample/sampler.py +++ b/vllm/v1/sample/sampler.py @@ -8,7 +8,9 @@ import torch import torch.nn as nn +from vllm import envs from vllm.config.model import LogprobsMode +from vllm.logger import init_logger from vllm.utils.platform_utils import is_pin_memory_available from vllm.v1.outputs import LogprobsTensors, SamplerOutput from vllm.v1.sample.metadata import SamplingMetadata @@ -19,6 +21,7 @@ _SAMPLING_EPS = 1e-5 _SM70_LOGITS_DUMP_COUNTER = 0 +logger = init_logger(__name__) def _sm70_cuda_graph_capture_active() -> bool: @@ -77,10 +80,7 @@ def tensor_cpu(value: torch.Tensor | None) -> torch.Tensor | None: ], }, path - / ( - f"sampler_logits_pid{os.getpid()}" - f"_step{_SM70_LOGITS_DUMP_COUNTER:04d}.pt" - ), + / (f"sampler_logits_pid{os.getpid()}_step{_SM70_LOGITS_DUMP_COUNTER:04d}.pt"), ) @@ -303,6 +303,217 @@ def apply_temperature( def greedy_sample(logits: torch.Tensor) -> torch.Tensor: return logits.argmax(dim=-1).view(-1) + @staticmethod + def _try_sm70_compact_topk20_sample( + logits: torch.Tensor, + sampling_metadata: SamplingMetadata, + logprobs_mode: LogprobsMode, + ) -> torch.Tensor | None: + if not Sampler._sm70_compact_topk20_metadata_supported( + sampling_metadata, logprobs_mode + ): + return None + if ( + logits.dtype != torch.float32 + or not logits.is_cuda + or logits.shape != (1, 248320) + or torch.cuda.get_device_capability(logits.device) != (7, 0) + ): + return None + + chunks = envs.VLLM_SM70_CHUNKED_TOPK20_CHUNKS + if ( + chunks > 0 + and 248320 % chunks == 0 + and hasattr(torch.ops._C, "sm70_sample_chunked_top20_philox_token_out") + and len(sampling_metadata.generators) <= 1 + ): + from vllm import _sm70_ops as sm70_ops + + chunk_size = 248320 // chunks + local_values, local_indices = torch.topk( + logits.view(chunks, chunk_size), + k=20, + dim=-1, + largest=True, + sorted=False, + ) + global_values, global_positions = torch.topk( + local_values.view(1, -1), + k=20, + dim=-1, + largest=True, + sorted=False, + ) + sampled = torch.empty((1,), dtype=torch.int64, device=logits.device) + generator = next(iter(sampling_metadata.generators.values()), None) + sm70_ops.sm70_sample_chunked_top20_philox_token_out( + sampled, + global_values, + local_indices, + global_positions, + generator, + 248320, + 0.95, + chunk_size, + ) + logger.info_once( + "SM70 exact %d-way chunked top-k20 Philox sampler path enabled.", + chunks, + ) + return sampled + + # The sparse-Philox kernel canonicalizes the selected pairs itself, so + # skip torch.topk's final ordering pass only when that exact path will + # consume the result. Keep sorted output for the legacy fallback. + use_sparse_philox = ( + hasattr(torch.ops._C, "sm70_sample_sorted_top20_philox_token_out") + or hasattr(torch.ops._C, "sm70_sample_sorted_top20_philox_out") + ) and len(sampling_metadata.generators) <= 1 + top_values, top_indices = torch.topk( + logits, + k=20, + dim=-1, + largest=True, + sorted=not use_sparse_philox, + ) + return Sampler._sm70_sample_top20_pairs( + top_values, top_indices, sampling_metadata + ) + + @staticmethod + def _sm70_compact_topk20_metadata_supported( + sampling_metadata: SamplingMetadata, + logprobs_mode: LogprobsMode, + ) -> bool: + if not envs.VLLM_SM70_COMPACT_TOPK20_SAMPLER: + return False + if logprobs_mode in ("processed_logits", "processed_logprobs"): + return False + if sampling_metadata.max_num_logprobs is not None: + return False + if sampling_metadata.logprob_token_ids: + return False + if not sampling_metadata.all_random: + return False + if sampling_metadata.top_k_cpu != (20,): + return False + top_p_cpu = sampling_metadata.top_p_cpu + temperature_cpu = sampling_metadata.temperature_cpu + return not ( + top_p_cpu is None + or len(top_p_cpu) != 1 + or abs(top_p_cpu[0] - 0.95) > 1e-6 + or temperature_cpu is None + or len(temperature_cpu) != 1 + or abs(temperature_cpu[0] - 1.0) > 1e-6 + ) + + @staticmethod + def _sm70_sample_top20_pairs( + top_values: torch.Tensor, + top_indices: torch.Tensor, + sampling_metadata: SamplingMetadata, + ) -> torch.Tensor | None: + if ( + top_values.dtype != torch.float32 + or top_indices.dtype != torch.int64 + or not top_values.is_cuda + or not top_indices.is_cuda + or top_values.shape != (1, 20) + or top_indices.shape != (1, 20) + or top_values.device != top_indices.device + or torch.cuda.get_device_capability(top_values.device) != (7, 0) + ): + return None + + from vllm import _sm70_ops as sm70_ops + + has_token_philox = hasattr( + torch.ops._C, "sm70_sample_sorted_top20_philox_token_out" + ) + has_metadata_philox = hasattr( + torch.ops._C, "sm70_sample_sorted_top20_philox_out" + ) + if not ( + has_token_philox + or has_metadata_philox + or hasattr(torch.ops._C, "sm70_sample_packed_top20_out") + ): + return None + + sampled = torch.empty((1,), dtype=torch.int64, device=top_values.device) + generators = sampling_metadata.generators + if has_token_philox and len(generators) <= 1: + generator = next(iter(generators.values()), None) + sm70_ops.sm70_sample_sorted_top20_philox_token_out( + sampled, + top_values, + top_indices, + generator, + 248320, + 0.95, + ) + logger.info_once( + "SM70 exact token-only Philox compact top-k20 sampler path enabled." + ) + return sampled + + sparse_ids = torch.empty((20,), dtype=torch.int64, device=top_values.device) + sparse_probs = torch.empty((20,), dtype=torch.float32, device=top_values.device) + if has_metadata_philox and len(generators) <= 1: + generator = next(iter(generators.values()), None) + sm70_ops.sm70_sample_sorted_top20_philox_out( + sampled, + sparse_ids, + sparse_probs, + top_values, + top_indices, + generator, + 248320, + 0.95, + ) + logger.info_once( + "SM70 exact sparse-Philox compact top-k20 sampler path enabled." + ) + return sampled + + top_indices_float = top_indices[0].to(torch.float32) + pairs = torch.stack( + (top_values[0], top_indices_float, top_indices_float), dim=-1 + ) + exponentials = torch.empty( + (1, 248320), dtype=torch.float32, device=top_values.device + ) + if len(generators) != 1: + exponentials.exponential_() + for row, generator in generators.items(): + exponentials[row].exponential_(generator=generator) + + sm70_ops.sm70_sample_packed_top20_out( + sampled, + sparse_ids, + sparse_probs, + pairs, + exponentials, + 0.95, + ) + logger.info_once("SM70 exact compact top-k20 random sampler path enabled.") + return sampled + + def sm70_compact_topk20_pairs_sample( + self, + top_values: torch.Tensor, + top_indices: torch.Tensor, + sampling_metadata: SamplingMetadata, + ) -> torch.Tensor | None: + """Sample exact global top-20 pairs while preserving native RNG use.""" + if not self._sm70_compact_topk20_metadata_supported( + sampling_metadata, self.logprobs_mode + ): + return None + return self._sm70_sample_top20_pairs(top_values, top_indices, sampling_metadata) + def sample( self, logits: torch.Tensor, @@ -345,6 +556,12 @@ def sample( for processor in sampling_metadata.logitsprocs.argmax_invariant: logits = processor.apply(logits) + compact_sampled = self._try_sm70_compact_topk20_sample( + logits, sampling_metadata, logprobs_mode + ) + if compact_sampled is not None: + return compact_sampled, None + # Apply top_k and/or top_p. random_sampled, processed_logprobs = self.topk_topp_sampler( logits, 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/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index 24b6a178ce..6cdd5e5b41 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -119,9 +119,7 @@ def __init__( self.cpu_coordinator: KVCacheCoordinator = get_kv_cache_coordinator( kv_cache_config=self.cpu_kv_cache_config, max_model_len=vllm_config.model_config.max_model_len, - max_num_batched_tokens=( - vllm_config.scheduler_config.max_num_batched_tokens - ), + max_in_flight_tokens=vllm_config.max_in_flight_tokens, use_eagle=False, enable_caching=True, enable_kv_cache_events=self.enable_kv_cache_events, 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/dflash.py b/vllm/v1/spec_decode/dflash.py index f90928b8b7..ffabcdc576 100644 --- a/vllm/v1/spec_decode/dflash.py +++ b/vllm/v1/spec_decode/dflash.py @@ -55,6 +55,32 @@ def _ddtree_worker_profile_enabled() -> bool: class DFlashProposer(SpecDecodeBaseProposer): + @override + def _create_draft_vllm_config(self) -> VllmConfig: + base = super()._create_draft_vllm_config() + if ( + self.speculative_config.kv_cache_dtype is None + and str(base.cache_config.cache_dtype).startswith("fp8") + and current_platform.is_cuda() + and current_platform.is_device_capability(70) + ): + logger.info_once( + "Using FP16 draft KV cache for SM70 DFlash DDTree while the " + "target uses %s.", + base.cache_config.cache_dtype, + ) + base = replace( + base, + cache_config=replace(base.cache_config, cache_dtype="auto"), + ) + return replace( + base, + attention_config=replace( + base.attention_config, + use_non_causal=not self.dflash_causal, + ), + ) + def __init__( self, vllm_config: VllmConfig, @@ -63,7 +89,7 @@ def __init__( ): assert vllm_config.speculative_config is not None assert ( - vllm_config.speculative_config.use_dflash() + vllm_config.speculative_config.use_dflash_ddtree() or vllm_config.speculative_config.use_dspark() ) super().__init__( @@ -323,17 +349,6 @@ def warmup_sm70_dflash_hotpath_kernels(self) -> tuple[str, ...]: "dflash_expand_inputs", ) - @override - def _create_draft_vllm_config(self) -> VllmConfig: - base = super()._create_draft_vllm_config() - return replace( - base, - attention_config=replace( - base.attention_config, - use_non_causal=not self.dflash_causal, - ), - ) - @override def _warn_if_multimodal(self): # Override to allow multimodal inputs since DFlash supports Qwen3.5 models diff --git a/vllm/v1/spec_decode/draft_prob_alignment.py b/vllm/v1/spec_decode/draft_prob_alignment.py index 8a0ad46c7e..21e6dd8b64 100644 --- a/vllm/v1/spec_decode/draft_prob_alignment.py +++ b/vllm/v1/spec_decode/draft_prob_alignment.py @@ -160,3 +160,108 @@ def get_aligned_draft_probs( if not draft_probs_rows: return None return torch.cat(draft_probs_rows, dim=0).contiguous() + + +def get_aligned_draft_scalar_values( + *, + req_ids: Sequence[str], + values: torch.Tensor | None, + value_req_ids: Sequence[str] | None, + value_token_ids: DraftProbTokenIds | None, + spec_decode_metadata: SpecDecodeMetadata, +) -> torch.Tensor | None: + """Align a cached per-draft scalar tensor to verifier request order.""" + if values is None or value_req_ids is None: + return None + if values.ndim != 2: + raise RuntimeError( + "Cached draft scalar values must have shape " + f"[num_requests, max_draft_tokens], got {tuple(values.shape)}." + ) + if values.shape[0] != len(value_req_ids): + raise RuntimeError( + "Cached draft scalar row count does not match request ids: " + f"rows={values.shape[0]}, req_ids={len(value_req_ids)}." + ) + if value_token_ids is None: + raise RuntimeError( + "Cached draft scalar values are missing their draft token ids; " + "cannot verify alignment." + ) + if isinstance(value_token_ids, torch.Tensor): + if value_token_ids.ndim != 2 or value_token_ids.shape[0] != len(value_req_ids): + raise RuntimeError( + "Cached draft scalar token ids must have one 2-D row per " + f"request, got shape={tuple(value_token_ids.shape)}, " + f"req_ids={len(value_req_ids)}." + ) + elif len(value_token_ids) != len(value_req_ids): + raise RuntimeError( + "Cached draft scalar token row count does not match request ids: " + f"rows={len(value_token_ids)}, req_ids={len(value_req_ids)}." + ) + if len(req_ids) != len(spec_decode_metadata.num_draft_tokens): + raise RuntimeError( + "Spec decode draft metadata batch size does not match current " + f"requests: metadata={len(spec_decode_metadata.num_draft_tokens)}, " + f"req_ids={len(req_ids)}." + ) + expected_num_draft_tokens = sum(spec_decode_metadata.num_draft_tokens) + if spec_decode_metadata.draft_token_ids.numel() != expected_num_draft_tokens: + raise RuntimeError( + "Spec decode draft token count does not match per-request counts: " + f"draft_token_ids={spec_decode_metadata.draft_token_ids.numel()}, " + f"num_draft_tokens={expected_num_draft_tokens}." + ) + + row_by_req_id = {req_id: idx for idx, req_id in enumerate(value_req_ids)} + aligned_rows: list[torch.Tensor] = [] + draft_token_offset = 0 + for req_id, num_draft in zip(req_ids, spec_decode_metadata.num_draft_tokens): + if num_draft == 0: + continue + row_idx = row_by_req_id.get(req_id) + if row_idx is None: + raise RuntimeError( + f"Missing cached draft scalar values for request {req_id}; " + "cannot verify alignment." + ) + if values.shape[1] < num_draft: + raise RuntimeError( + "Cached draft scalar values do not have enough positions for " + f"request {req_id}: available={values.shape[1]}, needed={num_draft}." + ) + expected_tokens = spec_decode_metadata.draft_token_ids[ + draft_token_offset : draft_token_offset + num_draft + ] + if isinstance(value_token_ids, torch.Tensor): + if value_token_ids.shape[1] < num_draft: + raise RuntimeError( + "Cached draft scalar token ids do not have enough positions " + f"for request {req_id}: available={value_token_ids.shape[1]}, " + f"needed={num_draft}." + ) + cached_tokens = value_token_ids[row_idx, :num_draft].to( + device=expected_tokens.device, + dtype=expected_tokens.dtype, + ) + else: + cached_token_row = value_token_ids[row_idx] + if len(cached_token_row) < num_draft: + raise RuntimeError( + "Cached draft scalar token ids do not have enough positions " + f"for request {req_id}: available={len(cached_token_row)}, " + f"needed={num_draft}." + ) + cached_tokens = torch.tensor( + cached_token_row[:num_draft], + device=expected_tokens.device, + dtype=expected_tokens.dtype, + ) + _assert_cached_tokens_match(cached_tokens, expected_tokens, req_id) + aligned_rows.append(values[row_idx, :num_draft]) + draft_token_offset += num_draft + + if not aligned_rows: + return None + return torch.cat(aligned_rows, dim=0).contiguous() diff --git a/vllm/v1/spec_decode/dspark.py b/vllm/v1/spec_decode/dspark.py index 93d8fb2abb..cb013f86e4 100644 --- a/vllm/v1/spec_decode/dspark.py +++ b/vllm/v1/spec_decode/dspark.py @@ -5,6 +5,7 @@ import torch from typing_extensions import override +from vllm import envs from vllm.config import VllmConfig from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.spec_decode.dflash import DFlashProposer @@ -27,6 +28,76 @@ def __init__( torch.arange(self.max_batch_size, dtype=torch.int64, device=device) * self.num_query_per_req ) + self._confidence_threshold = self.speculative_config.dspark_confidence_threshold + self._confidence_temperatures: torch.Tensor | None = None + if self._confidence_threshold > 0.0: + temperatures = self.speculative_config.dspark_confidence_temperatures + if temperatures is None: + temperatures = [1.0] * self.num_speculative_tokens + self._confidence_temperatures = torch.tensor( + temperatures, + dtype=torch.float32, + device=device, + ).view(1, -1) + configured_cap = self.speculative_config.dspark_max_verification_tokens + self._max_verification_tokens = ( + self.num_speculative_tokens if configured_cap is None else configured_cap + ) + self.confidence_scheduling_enabled = ( + self._confidence_threshold > 0.0 + or self._max_verification_tokens < self.num_speculative_tokens + ) + # A fixed cap needs verification lengths but not a confidence + # projection. Collect logits only when the calibrated threshold or the + # explicit alignment diagnostic consumes them. + self.collect_confidence_logits = ( + self._confidence_threshold > 0.0 or envs.VLLM_SPEC_DUMP_ALIGNMENT + ) + self._last_confidence_logits: torch.Tensor | None = None + self._last_verification_lengths: torch.Tensor | None = None + + def take_last_confidence_logits(self) -> torch.Tensor | None: + confidence_logits = self._last_confidence_logits + self._last_confidence_logits = None + return confidence_logits + + def take_last_verification_lengths(self) -> torch.Tensor | None: + verification_lengths = self._last_verification_lengths + self._last_verification_lengths = None + return verification_lengths + + def _select_verification_lengths( + self, + confidence_logits: torch.Tensor, + ) -> torch.Tensor: + temperatures = getattr(self, "_confidence_temperatures", None) + if temperatures is None: + temperatures = torch.ones( + (1, confidence_logits.shape[1]), + dtype=torch.float32, + device=confidence_logits.device, + ) + threshold = float(getattr(self, "_confidence_threshold", 0.0)) + cap = int( + getattr( + self, + "_max_verification_tokens", + confidence_logits.shape[1], + ) + ) + if threshold <= 0.0: + return torch.full( + (confidence_logits.shape[0],), + cap, + dtype=torch.int32, + device=confidence_logits.device, + ) + + conditional_probs = torch.sigmoid(confidence_logits.float() / temperatures) + confident_prefix = ( + (conditional_probs >= threshold).to(torch.int32).cumprod(dim=1) + ) + return confident_prefix.sum(dim=1).clamp_max_(cap).to(torch.int32) @override def _sample_draft_tokens( @@ -54,9 +125,19 @@ def _sample_draft_tokens( # contents across asynchronous scheduling and CUDA Graph replay. prev = self.input_ids[self._anchor_indices[:batch_size]].to(torch.long) draft_tokens: list[torch.Tensor] = [] + self._last_confidence_logits = None + self._last_verification_lengths = None + collect_confidence_logits = bool( + getattr(self, "collect_confidence_logits", False) + ) + markov_embeds: list[torch.Tensor] | None = ( + [] if collect_confidence_logits else None + ) draft_probs: list[torch.Tensor] | None = None for step in range(self.num_speculative_tokens): markov_embed = self.model.markov_embed(prev) + if markov_embeds is not None: + markov_embeds.append(markov_embed) step_logits = base_logits[:, step] + self.model.markov_bias(markov_embed) sampled, step_probs = self._sample_from_logits( step_logits, @@ -69,6 +150,26 @@ def _sample_draft_tokens( draft_probs = [] draft_probs.append(step_probs) + if markov_embeds is not None: + block_hidden_states = hidden_states.view( + batch_size, self.num_speculative_tokens, -1 + ) + self._last_confidence_logits = self.model.confidence_logits( + block_hidden_states, + torch.stack(markov_embeds, dim=1), + ) + if bool(getattr(self, "confidence_scheduling_enabled", False)): + if self._last_confidence_logits is None: + self._last_verification_lengths = torch.full( + (batch_size,), + int(self._max_verification_tokens), + dtype=torch.int32, + device=hidden_states.device, + ) + else: + self._last_verification_lengths = self._select_verification_lengths( + self._last_confidence_logits + ) flat_tokens = torch.stack(draft_tokens, dim=1).reshape(-1) if draft_probs is None: return flat_tokens, None diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index 85811c4c75..98bed3ac9d 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -87,6 +87,47 @@ def _sm70_mtp_profile_interval() -> int: return envs.VLLM_SM70_MTP_PROFILE_INTERVAL +def _sm70_mtp_moe_warmup_sizes( + num_experts: int, + top_k: int, + max_num_tokens: int, + decode_tile_max_tokens: int = 0, +) -> tuple[int, ...]: + """Pick one token count for each SM70 Triton MoE dispatch region.""" + if num_experts <= 0 or top_k <= 0 or max_num_tokens <= 0: + return () + + def prefer_unaligned(size: int) -> int: + # Triton otherwise specializes another kernel when token_count * top_k + # first loses 16-byte divisibility. This cannot happen when top_k is + # itself a multiple of 16. + if top_k % 16: + while size * top_k % 16 == 0: + size += 1 + return size + + sizes = { + # First sorted-assignment shape after the naive small-token path. + prefer_unaligned(num_experts // (top_k * 4) + 1), + # First shape whose sorted-token allocation is not block aligned. + prefer_unaligned((num_experts + top_k - 1) // top_k), + # First shape using the large-M SM70 0.0.3 MoE tile. + prefer_unaligned(num_experts + 1), + } + sizes.update(range(1, min(decode_tile_max_tokens, max_num_tokens) + 1)) + return tuple(sorted(size for size in sizes if size <= max_num_tokens)) + + +def _sm70_mtp_hotpath_warmup_batch_sizes( + max_batch_size: int, include_alternate: bool = False +) -> tuple[int, ...]: + """Preserve the primary warmup and optionally cover its alternate shape.""" + concurrent_batch_size = max(1, min(int(max_batch_size), 4)) + if include_alternate: + return tuple(sorted({1, concurrent_batch_size})) + return (concurrent_batch_size,) + + def _is_dflash_method(method: str | None) -> bool: return method in ("dflash", "dflash_ddtree", "dspark") @@ -835,6 +876,125 @@ def _sm70_mtp_profile_report( summary, ) + def _warmup_sm70_mtp_hotpath_batch( + self, batch_size: int, vocab_size: int, warm_draft_topk: bool + ) -> None: + valid_sampled_tokens_count = None + for num_sampled_tokens in sorted({1, self.num_speculative_tokens + 1}): + sampled_token_ids = torch.zeros( + (batch_size, num_sampled_tokens), + dtype=torch.int32, + device=self.device, + ) + discard_request_mask = torch.zeros( + batch_size, dtype=torch.bool, device=self.device + ) + backup_next_token_ids = torch.zeros( + batch_size, dtype=torch.int32, device=self.device + ) + next_token_ids = torch.empty( + batch_size, dtype=torch.int32, device=self.device + ) + next_valid_count = torch.empty_like(next_token_ids) + eagle_prepare_next_token_padded_kernel[(batch_size,)]( + sampled_token_ids, + discard_request_mask, + backup_next_token_ids, + next_token_ids, + next_valid_count, + vocab_size, + num_sampled_tokens, + batch_size, + sampled_token_ids.stride(0), + BLOCK_SIZE_TOKENS=next_power_of_2(num_sampled_tokens), + ) + if num_sampled_tokens == self.num_speculative_tokens + 1: + valid_sampled_tokens_count = next_valid_count + assert valid_sampled_tokens_count is not None + + next_token_ids = torch.empty(batch_size, dtype=torch.int32, device=self.device) + cu_num_draft_tokens = ( + torch.arange( + 1, + batch_size + 1, + dtype=torch.int32, + device=self.device, + ) + * self.num_speculative_tokens + ) + query_start_loc = torch.arange( + batch_size + 1, dtype=torch.int32, device=self.device + ) + token_indices_to_sample = torch.empty_like(next_token_ids) + num_rejected_tokens_gpu = torch.empty_like(next_token_ids) + eagle_prepare_inputs_padded_kernel[(batch_size,)]( + cu_num_draft_tokens, + valid_sampled_tokens_count, + query_start_loc, + token_indices_to_sample, + num_rejected_tokens_gpu, + batch_size, + ) + + for dtype, replace_from, replace_to in ( + (torch.float32, 0, 1), # temperature + (torch.float32, 0, 0), # top_p + (torch.int32, 0, 0), # top_k + ): + rejection_expand_input = torch.ones( + batch_size, dtype=dtype, device=self.device + ) + rejection_expand_output = torch.empty( + batch_size * self.num_speculative_tokens, + dtype=dtype, + device=self.device, + ) + rejection_expand_kernel[(batch_size,)]( + rejection_expand_output, + rejection_expand_input, + cu_num_draft_tokens, + replace_from, + replace_to, + MAX_NUM_TOKENS=MAX_SPEC_LEN, + ) + + n_blocks_per_req = max( + 1, (self.max_model_len + self.block_size - 1) // self.block_size + ) + positions = torch.zeros(batch_size, dtype=torch.int64, device=self.device) + block_table = torch.zeros( + (batch_size, n_blocks_per_req), + dtype=torch.int32, + device=self.device, + ) + seq_lens = torch.ones(batch_size, dtype=torch.int32, device=self.device) + out_positions = torch.empty_like(positions) + out_slot_mapping = torch.empty( + batch_size, dtype=torch.int64, device=self.device + ) + eagle_step_update_slot_mapping_and_metadata( + positions_1d=positions, + block_table_tensor=block_table, + seq_lens=seq_lens, + block_size=self.block_size, + max_model_len=self.max_model_len, + out_clamped_positions=out_positions, + out_slot_mapping=out_slot_mapping, + input_batch_size=batch_size, + ) + + if warm_draft_topk: + draft_logits = torch.zeros( + (batch_size, vocab_size), dtype=torch.float32, device=self.device + ) + draft_top_k = torch.full( + (batch_size,), + min(20, vocab_size), + dtype=torch.int32, + device=self.device, + ) + apply_top_k_top_p(draft_logits, draft_top_k, None) + def warmup_sm70_mtp_hotpath_kernels(self) -> tuple[str, ...]: """Warm MTP helper kernels that otherwise JIT on the first request.""" if ( @@ -850,126 +1010,101 @@ def warmup_sm70_mtp_hotpath_kernels(self) -> tuple[str, ...]: return () try: - batch_size = max(1, min(self.max_batch_size, 4)) vocab_size = max(2, self.draft_model_config.get_vocab_size()) - - valid_sampled_tokens_count = None - for num_sampled_tokens in sorted({1, self.num_speculative_tokens + 1}): - sampled_token_ids = torch.zeros( - (batch_size, num_sampled_tokens), - dtype=torch.int32, - device=self.device, - ) - discard_request_mask = torch.zeros( - batch_size, dtype=torch.bool, device=self.device - ) - backup_next_token_ids = torch.zeros( - batch_size, dtype=torch.int32, device=self.device - ) - next_token_ids = torch.empty( - batch_size, dtype=torch.int32, device=self.device - ) - next_valid_count = torch.empty_like(next_token_ids) - eagle_prepare_next_token_padded_kernel[(batch_size,)]( - sampled_token_ids, - discard_request_mask, - backup_next_token_ids, - next_token_ids, - next_valid_count, - vocab_size, - num_sampled_tokens, - batch_size, - sampled_token_ids.stride(0), - BLOCK_SIZE_TOKENS=next_power_of_2(num_sampled_tokens), - ) - if num_sampled_tokens == self.num_speculative_tokens + 1: - valid_sampled_tokens_count = next_valid_count - assert valid_sampled_tokens_count is not None - - next_token_ids = torch.empty( - batch_size, dtype=torch.int32, device=self.device - ) - - cu_num_draft_tokens = ( - torch.arange( - 1, - batch_size + 1, - dtype=torch.int32, - device=self.device, - ) - * self.num_speculative_tokens - ) - query_start_loc = torch.arange( - batch_size + 1, dtype=torch.int32, device=self.device - ) - token_indices_to_sample = torch.empty_like(next_token_ids) - num_rejected_tokens_gpu = torch.empty_like(next_token_ids) - eagle_prepare_inputs_padded_kernel[(batch_size,)]( - cu_num_draft_tokens, - valid_sampled_tokens_count, - query_start_loc, - token_indices_to_sample, - num_rejected_tokens_gpu, - batch_size, - ) - - for dtype, replace_from, replace_to in ( - (torch.float32, 0, 1), # temperature - (torch.float32, 0, 0), # top_p - (torch.int32, 0, 0), # top_k + expanded_warmup = envs.VLLM_SM70_MTP_CONCURRENCY_WARMUP + for batch_size in _sm70_mtp_hotpath_warmup_batch_sizes( + self.max_batch_size, include_alternate=expanded_warmup ): - rejection_expand_input = torch.ones( - batch_size, dtype=dtype, device=self.device - ) - rejection_expand_output = torch.empty( - batch_size * self.num_speculative_tokens, - dtype=dtype, - device=self.device, + self._warmup_sm70_mtp_hotpath_batch( + batch_size, vocab_size, warm_draft_topk=expanded_warmup ) - rejection_expand_kernel[(batch_size,)]( - rejection_expand_output, - rejection_expand_input, - cu_num_draft_tokens, - replace_from, - replace_to, - MAX_NUM_TOKENS=MAX_SPEC_LEN, - ) - - n_blocks_per_req = max( - 1, (self.max_model_len + self.block_size - 1) // self.block_size - ) - positions = torch.zeros(batch_size, dtype=torch.int64, device=self.device) - block_table = torch.zeros( - (batch_size, n_blocks_per_req), - dtype=torch.int32, - device=self.device, - ) - seq_lens = torch.ones(batch_size, dtype=torch.int32, device=self.device) - out_positions = torch.empty_like(positions) - out_slot_mapping = torch.empty( - batch_size, dtype=torch.int64, device=self.device - ) - eagle_step_update_slot_mapping_and_metadata( - positions_1d=positions, - block_table_tensor=block_table, - seq_lens=seq_lens, - block_size=self.block_size, - max_model_len=self.max_model_len, - out_clamped_positions=out_positions, - out_slot_mapping=out_slot_mapping, - input_batch_size=batch_size, - ) torch.accelerator.synchronize() except Exception as err: # pragma: no cover - best-effort warmup logger.warning_once("SM70 MTP hotpath warmup skipped: %s", err) return () - return ( + warmed_kernels = [ "mtp_prepare_next_token", "mtp_prepare_inputs", "mtp_rejection_expand", "mtp_step_slot_mapping", + ] + if envs.VLLM_SM70_MTP_CONCURRENCY_WARMUP: + warmed_kernels.append("mtp_draft_topk") + return tuple(warmed_kernels) + + def warmup_sm70_mtp_moe_kernels(self) -> tuple[str, ...]: + """Warm the finite Triton-MoE shape specializations used by MTP.""" + if ( + self.method != "mtp" + or self.device.type != "cuda" + or not current_platform.is_device_capability(70) + or not envs.VLLM_SM70_UNQUANTIZED_MOE_0DOT3_CONFIG + or not self.draft_model_config.is_moe + ): + return () + if getattr(self, "_sm70_mtp_moe_warmed", False): + return () + self._sm70_mtp_moe_warmed = True + + text_config = self.draft_model_config.hf_text_config + top_k = int(getattr(text_config, "num_experts_per_tok", 0)) + num_experts = self.draft_model_config.get_num_experts() + tp_size = self.vllm_config.parallel_config.tensor_parallel_size + use_mtp_decode_tiles = ( + envs.VLLM_SM70_MTP_MOE_TUNED_CONFIG + and num_experts == 256 + and top_k == 8 + and self.draft_model_config.get_hidden_size() == 2048 + and int(getattr(text_config, "moe_intermediate_size", 0)) == 512 + and tp_size == 4 ) + sizes = _sm70_mtp_moe_warmup_sizes( + num_experts, + top_k, + self.max_num_tokens, + decode_tile_max_tokens=16 if use_mtp_decode_tiles else 0, + ) + legacy_sizes = (1, 2, 9, 10) if use_mtp_decode_tiles else () + if not sizes: + return () + + try: + for num_tokens in sizes: + self.dummy_run( + num_tokens, + use_cudagraphs=False, + spec_step_idx=0, + ) + if legacy_sizes: + # The tuned drafter tile replaces legacy graph + # specializations that used to be compiled incidentally. + # Other unquantized MoE calls can reuse those legacy Triton + # signatures during serving. Cover both token divisibility + # cases for the naive (M1/M2) and aligned (M9/M10) assignment + # paths instead of allowing a first-request JIT spike. + from vllm.model_executor.layers.fused_moe.fused_moe import ( + force_sm70_mtp_moe_legacy_config, + ) + + with force_sm70_mtp_moe_legacy_config(): + for num_tokens in legacy_sizes: + self.dummy_run( + num_tokens, + use_cudagraphs=False, + spec_step_idx=0, + ) + torch.accelerator.synchronize() + except Exception as err: # pragma: no cover - best-effort warmup + logger.warning_once("SM70 MTP MoE warmup skipped: %s", err) + return () + + logger.info_once( + "SM70 MTP MoE warmup finished: token_counts=%s, legacy_token_counts=%s", + sizes, + legacy_sizes, + ) + return ("mtp_draft_moe",) def take_last_draft_probs(self) -> torch.Tensor | None: return self._last_draft_probs @@ -1969,6 +2104,15 @@ def _create_draft_vllm_config(self) -> VllmConfig: ), ) + if spec_cfg.kv_cache_dtype is not None: + base = replace( + base, + cache_config=replace( + base.cache_config, + cache_dtype=spec_cfg.kv_cache_dtype, + ), + ) + return base def _get_model(self) -> nn.Module: diff --git a/vllm/v1/worker/block_table.py b/vllm/v1/worker/block_table.py index 616fdd15bd..8329a2c196 100644 --- a/vllm/v1/worker/block_table.py +++ b/vllm/v1/worker/block_table.py @@ -104,9 +104,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 c2ba12437f..771f06bff5 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Iterable, Sequence +from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass from typing import Any, cast @@ -372,6 +372,159 @@ def build_slot_mappings_by_layer( return slot_mappings_by_layer +@dataclass(frozen=True) +class CommonGDNSpecMetadata: + """Batch-level GDN speculative metadata shared by every cache group.""" + + num_prefills: int + num_prefill_tokens: int + num_decodes: int + num_decode_tokens: int + num_spec_decodes: int + num_spec_decode_tokens: int + spec_query_start_loc: torch.Tensor + non_spec_query_start_loc: torch.Tensor | None + non_spec_query_start_loc_cpu: torch.Tensor | None + spec_sequence_masks_cpu: torch.Tensor + spec_sequence_masks: torch.Tensor + spec_token_indx: torch.Tensor + non_spec_token_indx: torch.Tensor + + +def compute_common_gdn_attn_metadata( + *, + num_decode_draft_tokens_cpu: torch.Tensor | None, + query_start_loc: torch.Tensor, + query_start_loc_cpu: torch.Tensor, + num_spec_state_tokens: int, + legacy_mixed_decode_routing: bool, +) -> CommonGDNSpecMetadata | None: + """Compute GDN speculative batch metadata once instead of per group. + + State block IDs remain group-specific and are deliberately excluded. This + is the dependency boundary needed to share request classification, token + indices, and query offsets without changing recurrent-state ownership. + """ + if num_spec_state_tokens <= 0 or num_decode_draft_tokens_cpu is None: + return None + + num_reqs = query_start_loc_cpu.numel() - 1 + if num_decode_draft_tokens_cpu.ndim != 1: + raise ValueError("num_decode_draft_tokens_cpu must be one-dimensional") + if num_decode_draft_tokens_cpu.numel() != num_reqs: + raise ValueError("num_decode_draft_tokens_cpu must align with query_start_loc") + + spec_sequence_masks_cpu = num_decode_draft_tokens_cpu >= 0 + num_spec_decodes = int(spec_sequence_masks_cpu.sum().item()) + if num_spec_decodes == 0: + return None + num_spec_draft_tokens = int( + num_decode_draft_tokens_cpu[spec_sequence_masks_cpu].sum().item() + ) + if num_spec_draft_tokens == 0: + return None + + spec_sequence_masks = spec_sequence_masks_cpu.to( + query_start_loc.device, non_blocking=True + ) + query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] + non_spec_query_lens_cpu = query_lens_cpu[~spec_sequence_masks_cpu] + num_zero_len = int((non_spec_query_lens_cpu == 0).sum().item()) + + if legacy_mixed_decode_routing: + num_decodes = int((non_spec_query_lens_cpu == 1).sum().item()) + num_prefills = non_spec_query_lens_cpu.numel() - num_decodes - num_zero_len + num_decode_tokens = num_decodes + num_prefill_tokens = ( + int(non_spec_query_lens_cpu.sum().item()) - num_decode_tokens + ) + else: + num_decodes = 0 + num_prefills = non_spec_query_lens_cpu.numel() - num_zero_len + num_decode_tokens = 0 + num_prefill_tokens = int(non_spec_query_lens_cpu.sum().item()) + + num_spec_decode_tokens = ( + int(query_lens_cpu.sum().item()) - num_prefill_tokens - num_decode_tokens + ) + if num_prefills == 0 and num_decodes == 0: + spec_token_size = min( + num_spec_decodes * (num_spec_state_tokens + 1), + int(query_start_loc_cpu[-1].item()), + ) + spec_token_indx = torch.arange( + spec_token_size, + dtype=torch.int32, + device=query_start_loc.device, + ) + non_spec_token_indx = torch.empty( + 0, + dtype=torch.int32, + device=query_start_loc.device, + ) + 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: + query_lens = query_lens_cpu.to(query_start_loc.device, non_blocking=True) + spec_token_masks = torch.repeat_interleave( + spec_sequence_masks, + query_lens, + output_size=int(query_start_loc_cpu[-1].item()), + ) + index = torch.argsort(spec_token_masks, stable=True) + num_non_spec_tokens = num_prefill_tokens + num_decode_tokens + non_spec_token_indx = index[:num_non_spec_tokens] + spec_token_indx = index[num_non_spec_tokens:] + + spec_query_start_loc = torch.zeros( + num_spec_decodes + 1, + dtype=torch.int32, + device=query_start_loc.device, + ) + torch.cumsum( + query_lens[spec_sequence_masks], + dim=0, + out=spec_query_start_loc[1:], + ) + non_spec_query_start_loc = torch.zeros( + query_lens.numel() - num_spec_decodes + 1, + dtype=torch.int32, + device=query_start_loc.device, + ) + torch.cumsum( + query_lens[~spec_sequence_masks], + dim=0, + out=non_spec_query_start_loc[1:], + ) + non_spec_query_start_loc_cpu = torch.zeros( + query_lens_cpu.numel() - num_spec_decodes + 1, + dtype=torch.int32, + device="cpu", + ) + torch.cumsum( + query_lens_cpu[~spec_sequence_masks_cpu], + dim=0, + out=non_spec_query_start_loc_cpu[1:], + ) + + return CommonGDNSpecMetadata( + num_prefills=num_prefills, + num_prefill_tokens=num_prefill_tokens, + num_decodes=num_decodes, + num_decode_tokens=num_decode_tokens, + num_spec_decodes=num_spec_decodes, + num_spec_decode_tokens=num_spec_decode_tokens, + spec_query_start_loc=spec_query_start_loc, + non_spec_query_start_loc=non_spec_query_start_loc, + non_spec_query_start_loc_cpu=non_spec_query_start_loc_cpu, + spec_sequence_masks_cpu=spec_sequence_masks_cpu, + spec_sequence_masks=spec_sequence_masks, + spec_token_indx=spec_token_indx, + non_spec_token_indx=non_spec_token_indx, + ) + + def build_attn_metadata( attn_groups: list[list[AttentionGroup]], num_reqs: int, @@ -389,6 +542,8 @@ def build_attn_metadata( positions: torch.Tensor | None = None, model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None, for_cudagraph_capture: bool = False, + causal: bool | torch.Tensor | Mapping[int, bool] = True, + prefix_anchor_lens: torch.Tensor | None = None, ) -> dict[str, Any]: seq_lens = seq_lens[:num_reqs] if dcp_local_seq_lens is not None: @@ -401,6 +556,12 @@ def build_attn_metadata( for i in range(num_kv_cache_groups): block_table = block_tables[i] slot_mapping = slot_mappings[i] + # Hybrid drafters can assign different causality to each KV group. + # A Mapping must be resolved here; passing the dict itself into the + # backend metadata makes it truthy and silently selects causal kernels. + group_causal = ( + causal if isinstance(causal, (bool, torch.Tensor)) else causal.get(i, True) + ) common_attn_metadata_extra_kwargs = ( model_specific_attn_metadata.get_extra_common_attn_kwargs(i, num_reqs) @@ -418,9 +579,10 @@ def build_attn_metadata( max_query_len=max_query_len, block_table_tensor=block_table, slot_mapping=slot_mapping, - causal=True, + causal=group_causal, dcp_local_seq_lens=dcp_local_seq_lens, positions=positions, + prefix_anchor_lens=prefix_anchor_lens, **common_attn_metadata_extra_kwargs, ) diff --git a/vllm/v1/worker/gpu/buffer_utils.py b/vllm/v1/worker/gpu/buffer_utils.py index 5963790a77..2dd6bf92b8 100644 --- a/vllm/v1/worker/gpu/buffer_utils.py +++ b/vllm/v1/worker/gpu/buffer_utils.py @@ -137,6 +137,12 @@ def __init__( self.write_starts = new_buffer(self.num_rows, dtype=torch.int32) self.write_cu_lens = new_buffer(self.num_rows, dtype=torch.int32) + def get_cpu_view(self) -> torch.Tensor: + """Return the host view when this tensor is backed by mapped memory.""" + if not hasattr(self, "_uva_buf"): + raise RuntimeError("This StagedWriteTensor is not UVA-backed") + return self._uva_buf.cpu + def stage_write( self, index: int, start: int, x: Iterable[int] | Iterable[float] ) -> None: diff --git a/vllm/v1/worker/gpu/cp_utils.py b/vllm/v1/worker/gpu/cp_utils.py index 6dd8fd3474..08a1a258ed 100644 --- a/vllm/v1/worker/gpu/cp_utils.py +++ b/vllm/v1/worker/gpu/cp_utils.py @@ -59,3 +59,24 @@ def _dcp_local_seq_lens_kernel( # For [num_reqs, max_num_reqs), pad with 0 local_seq_lens = tl.where(block < num_reqs, local_seq_lens, 0) tl.store(out_ptr + block, local_seq_lens, mask=block < max_num_reqs) + + +@triton.jit +def cp_local_slot( + positions, + block_numbers, + block_size, + cp_rank, + CP_SIZE: tl.constexpr, + CP_INTERLEAVE: tl.constexpr, + PAD_ID: tl.constexpr, +): + """Return rank-local KV slots, or PAD_ID for non-local positions.""" + block_offsets = positions % (block_size * CP_SIZE) + if CP_SIZE == 1: + return block_numbers * block_size + block_offsets + is_local = block_offsets // CP_INTERLEAVE % CP_SIZE == cp_rank + rounds = block_offsets // (CP_INTERLEAVE * CP_SIZE) + remainder = block_offsets % CP_INTERLEAVE + local_offsets = rounds * CP_INTERLEAVE + remainder + return tl.where(is_local, block_numbers * block_size + local_offsets, PAD_ID) diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 4ecbb9f97d..a0953c44c5 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -34,6 +34,36 @@ logger = init_logger(__name__) +def _use_split_sm70_mtp_cudagraphs(vllm_config: VllmConfig) -> bool: + speculative_config = vllm_config.speculative_config + return bool( + envs.VLLM_SM70_MTP_SPLIT_DRAFT_CUDAGRAPHS + and speculative_config is not None + and speculative_config.method == "mtp" + and current_platform.is_cuda() + and current_platform.is_device_capability((7, 0)) + ) + + +def _split_sm70_mtp_manager_capture_sizes( + vllm_config: VllmConfig, + decode_query_len: int, +) -> list[int]: + capture_sizes = sorted(vllm_config.compilation_config.cudagraph_capture_sizes or []) + if decode_query_len != 1: + return capture_sizes + + speculative_config = vllm_config.speculative_config + assert speculative_config is not None + verifier_query_len = speculative_config.num_speculative_state_tokens() + 1 + return [ + size // verifier_query_len + for size in capture_sizes + if size % verifier_query_len == 0 + and size // verifier_query_len <= vllm_config.scheduler_config.max_num_seqs + ] + + class CapturedAttentionState(NamedTuple): attn_metadata: dict[str, Any] | None slot_mappings: dict[str, torch.Tensor] @@ -111,15 +141,32 @@ def __init__( self._graphs_captured = False self._candidates: list[list[BatchExecutionDescriptor]] = [] self._capture_descs: dict[CUDAGraphMode, list[BatchExecutionDescriptor]] = {} - # adjust the cudagraph sizes to be a multiple of the uniform decode query length - self.compilation_config.adjust_cudagraph_sizes_for_spec_decode( - self.decode_query_len, self.tp_size - ) + if _use_split_sm70_mtp_cudagraphs(vllm_config): + self._capture_sizes = _split_sm70_mtp_manager_capture_sizes( + vllm_config, + decode_query_len, + ) + logger.info_once( + "Using split SM70 MTP CUDA graph manager shapes: " + "decode_query_len=%d capture_sizes=%s shared_verifier_shapes=%s.", + decode_query_len, + tuple(self._capture_sizes), + tuple(self.compilation_config.cudagraph_capture_sizes or ()), + ) + else: + # Legacy MRV1 behavior. The rollback path deliberately preserves + # the shared in-place normalization used before split managers. + self.compilation_config.adjust_cudagraph_sizes_for_spec_decode( + self.decode_query_len, self.tp_size + ) + self._capture_sizes = list( + self.compilation_config.cudagraph_capture_sizes or [] + ) self._init_candidates() def _init_candidates(self) -> None: """Build priority-ordered candidate lists for each token count.""" - capture_sizes = self.compilation_config.cudagraph_capture_sizes + capture_sizes = self._capture_sizes if not (self.cudagraph_mode and capture_sizes): return diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index b253d7d8c0..1e117e499d 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -84,6 +84,14 @@ class InputBatch: # Whether any requests in batch use structured output. has_structured_output_reqs: bool + # [num_reqs_after_padding] per-request prompt length for prefix-anchored + # sliding-window attention (optional). + prefix_anchor_lens: torch.Tensor | None = None + # CPU mask for requests whose scheduled chunk does not finish prefill. + # Model-specific runners may use this to omit work whose outputs cannot be + # consumed until a later chunk. None is equivalent to all False. + is_incomplete_prefilling_np: np.ndarray | None = None + @classmethod def make_dummy( cls, diff --git a/vllm/v1/worker/gpu/mamba_align.py b/vllm/v1/worker/gpu/mamba_align.py new file mode 100644 index 0000000000..52a03e83c3 --- /dev/null +++ b/vllm/v1/worker/gpu/mamba_align.py @@ -0,0 +1,319 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MRV2 GPU state migration for hybrid Mamba ``align`` prefix caching. + +This is the narrow Model Runner V2 portion of upstream vLLM PR #42406. It is +kept separate from :mod:`vllm.v1.worker.mamba_utils` because that module also +owns this tree's V1 DDTree state-selection compatibility path. +""" + +import torch + +from vllm.triton_utils import tl, triton +from vllm.v1.worker.mamba_utils import MambaSpecDecodeGPUContext + +# Temporal states are much larger than conv states. Splitting their byte range +# keeps block-boundary copies from serializing on one CTA at small batch sizes. +_TEMPORAL_TILES = 16 + + +@triton.jit +def _copy_mamba_state_block( + state_idx, + block_table_row, + src_col, + dst_col, + token_bias, + block_table_ptrs_ptr, + block_table_stride_req, + state_base_addrs_ptr, + state_block_strides_ptr, + state_elem_sizes_ptr, + state_inner_sizes_ptr, + state_conv_widths_ptr, + state_group_indices_ptr, + tile_idx, + COPY_BLOCK_SIZE: tl.constexpr, + TEMPORAL_TILES: tl.constexpr, +): + """Copy one SD-layout conv or temporal state using V1 align semantics.""" + state_base_addr = tl.load(state_base_addrs_ptr + state_idx) + state_block_stride = tl.load(state_block_strides_ptr + state_idx) + state_elem_size = tl.load(state_elem_sizes_ptr + state_idx).to(tl.int64) + state_inner_size = tl.load(state_inner_sizes_ptr + state_idx).to(tl.int64) + conv_width = tl.load(state_conv_widths_ptr + state_idx) + + group_idx = tl.load(state_group_indices_ptr + state_idx).to(tl.int64) + group_base_addr = tl.load(block_table_ptrs_ptr + group_idx) + block_table = group_base_addr.to(tl.pointer_type(tl.int32)) + block_table += block_table_row * block_table_stride_req + + dst_block_id = tl.load(block_table + dst_col).to(tl.int64) + dst_addr = state_base_addr + dst_block_id * state_block_stride + is_conv_state = conv_width > 0 + + if is_conv_state: + # Conv states are small and may be shifted left within the same block. + # One CTA owns the copy so a block-local barrier can make each chunk + # memmove-safe before lower addresses overwrite overlapping input. + if tile_idx > 0: + return + src_block_id = tl.load(block_table + src_col).to(tl.int64) + token_bytes = state_inner_size * state_elem_size + src_addr = ( + state_base_addr + + src_block_id * state_block_stride + + token_bias.to(tl.int64) * token_bytes + ) + copy_size = (conv_width - token_bias).to(tl.int64) * token_bytes + tile_start = copy_size * 0 + tile_end = copy_size + is_left_overlap = (dst_addr < src_addr) & (dst_addr + copy_size > src_addr) + else: + # Temporal state chooses the accepted speculative column. + src_block_id = tl.load(block_table + src_col + token_bias).to(tl.int64) + src_addr = state_base_addr + src_block_id * state_block_stride + copy_size = state_inner_size * state_elem_size + bytes_per_tile = (copy_size + TEMPORAL_TILES - 1) // TEMPORAL_TILES + tile_start = tile_idx.to(tl.int64) * bytes_per_tile + tile_end = tl.minimum(tile_start + bytes_per_tile, copy_size) + is_left_overlap = False + + offsets = tl.arange(0, COPY_BLOCK_SIZE) + for chunk_start in range(tile_start, tile_end, COPY_BLOCK_SIZE): + byte_offsets = chunk_start + offsets + mask = byte_offsets < tile_end + src = (src_addr + byte_offsets).to(tl.pointer_type(tl.uint8)) + dst = (dst_addr + byte_offsets).to(tl.pointer_type(tl.uint8)) + data = tl.load(src, mask=mask) + if is_left_overlap: + tl.debug_barrier() + tl.store(dst, data, mask=mask) + + +@triton.jit +def preprocess_mamba_align_fused_kernel( + idx_mapping_ptr, + state_idx_ptr, + num_computed_tokens_ptr, + query_start_loc_ptr, + num_accepted_tokens_ptr, + src_col_ptr, + token_bias_ptr, + num_reqs, + BLOCK_SIZE: tl.constexpr, + MAMBA_BLOCK_SIZE: tl.constexpr, +): + """Advance running state columns and emit any required pre-copy inputs.""" + rows = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = rows < num_reqs + req_indices = tl.load(idx_mapping_ptr + rows, mask=mask, other=0) + + state_idx = tl.load(state_idx_ptr + req_indices, mask=mask, other=-1) + num_computed = tl.load(num_computed_tokens_ptr + req_indices, mask=mask, other=0) + # A new request is seeded here, after the resolved MambaSpec is available. + # cache_config.block_size can be smaller than the actual Mamba block when + # an FP16 DFlash cache shares the hybrid page pool, so add_request cannot + # derive the state column from the global cache block size. + seeded_state_idx = (num_computed + MAMBA_BLOCK_SIZE - 1) // MAMBA_BLOCK_SIZE - 1 + state_idx = tl.where(state_idx == -2, seeded_state_idx, state_idx) + num_accepted = tl.load(num_accepted_tokens_ptr + req_indices, mask=mask, other=1) + token_bias = tl.maximum(num_accepted - 1, 0) + tl.store(src_col_ptr + req_indices, state_idx, mask=mask) + tl.store(token_bias_ptr + req_indices, token_bias, mask=mask) + + query_start = tl.load(query_start_loc_ptr + rows, mask=mask, other=0) + query_end = tl.load(query_start_loc_ptr + rows + 1, mask=mask, other=0) + computed_after = num_computed + query_end - query_start + new_state_idx = (computed_after + MAMBA_BLOCK_SIZE - 1) // MAMBA_BLOCK_SIZE - 1 + tl.store(state_idx_ptr + req_indices, new_state_idx, mask=mask) + crossed_boundary = (state_idx >= 0) & (state_idx != new_state_idx) + tl.store( + num_accepted_tokens_ptr + req_indices, + 1, + mask=mask & crossed_boundary, + ) + + +@triton.jit +def _precopy_mamba_align_kernel( + state_idx_ptr, + src_col_ptr, + token_bias_ptr, + block_table_ptrs_ptr, + block_table_stride_req, + state_base_addrs_ptr, + state_block_strides_ptr, + state_elem_sizes_ptr, + state_inner_sizes_ptr, + state_conv_widths_ptr, + state_group_indices_ptr, + idx_mapping_ptr, + num_reqs, + COPY_BLOCK_SIZE: tl.constexpr, + TEMPORAL_TILES: tl.constexpr, +): + batch_idx = tl.program_id(0) + state_idx = tl.program_id(1) + tile_idx = tl.program_id(2) + if batch_idx >= num_reqs: + return + req_idx = tl.load(idx_mapping_ptr + batch_idx) + if req_idx < 0: + return + + src_col = tl.load(src_col_ptr + req_idx) + dst_col = tl.load(state_idx_ptr + req_idx) + if src_col < 0 or src_col == dst_col: + return + token_bias = tl.load(token_bias_ptr + req_idx) + _copy_mamba_state_block( + state_idx, + batch_idx, + src_col, + dst_col, + token_bias, + block_table_ptrs_ptr, + block_table_stride_req, + state_base_addrs_ptr, + state_block_strides_ptr, + state_elem_sizes_ptr, + state_inner_sizes_ptr, + state_conv_widths_ptr, + state_group_indices_ptr, + tile_idx, + COPY_BLOCK_SIZE, + TEMPORAL_TILES, + ) + + +@triton.jit +def _postprocess_mamba_align_kernel( + num_accepted_snapshot_ptr, + num_accepted_out_ptr, + state_idx_ptr, + new_num_computed_tokens_ptr, + block_table_ptrs_ptr, + block_table_stride_req, + state_base_addrs_ptr, + state_block_strides_ptr, + state_elem_sizes_ptr, + state_inner_sizes_ptr, + state_conv_widths_ptr, + state_group_indices_ptr, + idx_mapping_ptr, + num_reqs, + MAMBA_BLOCK_SIZE: tl.constexpr, + COPY_BLOCK_SIZE: tl.constexpr, + TEMPORAL_TILES: tl.constexpr, +): + batch_idx = tl.program_id(0) + state_meta_idx = tl.program_id(1) + tile_idx = tl.program_id(2) + if batch_idx >= num_reqs: + return + req_idx = tl.load(idx_mapping_ptr + batch_idx) + if req_idx < 0: + return + + num_accepted = tl.load(num_accepted_snapshot_ptr + req_idx) + src_col = tl.load(state_idx_ptr + req_idx) + new_num_computed = tl.load(new_num_computed_tokens_ptr + req_idx) + num_tokens_running_state = new_num_computed - num_accepted + 1 + aligned_new_computed = (new_num_computed // MAMBA_BLOCK_SIZE) * MAMBA_BLOCK_SIZE + if aligned_new_computed < num_tokens_running_state: + return + + token_bias = aligned_new_computed - num_tokens_running_state + dst_col = aligned_new_computed // MAMBA_BLOCK_SIZE - 1 + if src_col == dst_col and state_meta_idx == 0 and tile_idx == 0: + tl.store(num_accepted_out_ptr + req_idx, 1) + if src_col == dst_col and token_bias == 0: + return + + _copy_mamba_state_block( + state_meta_idx, + batch_idx, + src_col, + dst_col, + token_bias, + block_table_ptrs_ptr, + block_table_stride_req, + state_base_addrs_ptr, + state_block_strides_ptr, + state_elem_sizes_ptr, + state_inner_sizes_ptr, + state_conv_widths_ptr, + state_group_indices_ptr, + tile_idx, + COPY_BLOCK_SIZE, + TEMPORAL_TILES, + ) + + +def run_mamba_align_precopy( + ctx: MambaSpecDecodeGPUContext, + num_reqs: int, + state_idx: torch.Tensor, + src_col: torch.Tensor, + token_bias: torch.Tensor, + idx_mapping: torch.Tensor, +) -> 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) + _precopy_mamba_align_kernel[grid]( + state_idx, + src_col, + token_bias, + ctx.block_table_ptrs, + ctx.block_table_stride_req, + ctx.state_base_addrs, + ctx.state_block_strides, + ctx.state_elem_sizes, + ctx.state_inner_sizes, + ctx.state_conv_widths, + ctx.state_group_indices, + idx_mapping, + num_reqs, + COPY_BLOCK_SIZE=1024, + TEMPORAL_TILES=_TEMPORAL_TILES, + ) + + +def run_mamba_align_postprocess( + ctx: MambaSpecDecodeGPUContext, + num_reqs: int, + num_accepted_tokens: torch.Tensor, + state_idx: torch.Tensor, + new_num_computed_tokens: torch.Tensor, + idx_mapping: torch.Tensor, +) -> None: + if num_reqs == 0 or not ctx.is_initialized: + return + # Snapshot the decision array because one program may reset a request's + # 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) + _postprocess_mamba_align_kernel[grid]( + snapshot, + num_accepted_tokens, + state_idx, + new_num_computed_tokens, + ctx.block_table_ptrs, + ctx.block_table_stride_req, + ctx.state_base_addrs, + ctx.state_block_strides, + ctx.state_elem_sizes, + ctx.state_inner_sizes, + ctx.state_conv_widths, + ctx.state_group_indices, + idx_mapping, + num_reqs, + MAMBA_BLOCK_SIZE=ctx.block_size, + COPY_BLOCK_SIZE=1024, + TEMPORAL_TILES=_TEMPORAL_TILES, + ) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 5cba66e5c9..47d1da107e 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -47,6 +47,7 @@ from vllm.tasks import SupportedTask from vllm.utils.math_utils import cdiv from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib +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 @@ -95,14 +96,19 @@ from vllm.v1.worker.gpu.sample.sampler import Sampler from vllm.v1.worker.gpu.shutdown import free_before_shutdown from vllm.v1.worker.gpu.spec_decode import init_speculator +from vllm.v1.worker.gpu.spec_decode.dflash2.sparse_rejection import ( + try_dflash2_sparse_target_rejection, +) from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( set_eagle3_aux_hidden_state_layers, ) from vllm.v1.worker.gpu.spec_decode.rejection_sampler import RejectionSampler +from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator from vllm.v1.worker.gpu.spec_decode.utils import DraftTokensHandler from vllm.v1.worker.gpu.states import RequestState from vllm.v1.worker.gpu.structured_outputs import StructuredOutputsWorker from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin +from vllm.v1.worker.utils import KVBlockZeroer logger = init_logger(__name__) @@ -175,11 +181,14 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): if self.is_last_pp_rank: self.speculator = init_speculator(self.vllm_config, self.device) - if self.speculative_config.method == "eagle3": - # EAGLE3 may require auxiliary hidden states from target model outputs. + if self.speculative_config.method in ("eagle3", "dflash"): + # EAGLE3 and DFlash may require auxiliary target hidden states. self.use_aux_hidden_state_outputs = True if self.use_pp: - raise ValueError("EAGLE3 with pipeline parallel is not supported.") + raise ValueError( + f"{self.speculative_config.method} with pipeline parallel " + "is not supported." + ) # Draft tokens propagation - for spec-dec + struct outputs. self.draft_tokens_handler = DraftTokensHandler(self.device) @@ -203,6 +212,20 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): max_num_tokens=self.max_num_tokens, device=self.device, ) + # Prefix-anchored SWA: persistent GPU buffer for per-request prompt + # lengths (stable device address across steps). + self.prefix_anchor_lens_buffer: torch.Tensor | None = None + if ( + getattr( + self.vllm_config.attention_config, + "prefix_anchored_decode_window", + None, + ) + is not None + ): + self.prefix_anchor_lens_buffer = torch.zeros( + self.max_num_reqs, dtype=torch.int32, device=self.device + ) self.sampler: Sampler | None = None self.rejection_sampler: RejectionSampler | None = None @@ -396,6 +419,7 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: self.attn_groups, attn_cg_support, kernel_block_sizes = init_attn_backend( self.kv_cache_config, self.vllm_config, self.device ) + self._kernel_block_sizes = kernel_block_sizes self.block_tables = BlockTables( block_sizes=block_sizes, max_num_reqs=self.max_num_reqs, @@ -424,11 +448,21 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: cudagraph_mode, decode_query_len=self.decode_query_len, ) - if self.speculator is not None: - self.speculator.init_cudagraph_manager(cudagraph_mode) - check_attention_cp_compatibility(self.vllm_config) - if self.speculator is not None: + if isinstance(self.speculator, DraftModelSpeculator): + self.speculator.set_attn( + self.model_state, + self.kv_cache_config, + self.block_tables, + self.input_buffers, + self.attn_groups, + ) + # DFlash sizes its graph mode from the selected draft attention + # backend, which is available only after set_attn(). + self.speculator.init_cudagraph_manager(cudagraph_mode) + elif self.speculator is not None: + # Preserve the existing Eagle initialization order. + self.speculator.init_cudagraph_manager(cudagraph_mode) # HACK(woosuk) self.speculator.set_attn( self.model_state, self.kv_cache_config, self.block_tables @@ -447,6 +481,24 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: ) self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict) + def _init_kv_zero_meta(self) -> None: + """Precompute metadata used to clear newly allocated cache blocks.""" + self._kv_block_zeroer = KVBlockZeroer( + self.device, pin_memory=is_pin_memory_available() + ) + self._kv_block_zeroer.init_meta( + attn_groups_iter=(group for groups in self.attn_groups for group in groups), + kernel_block_sizes=self._kernel_block_sizes, + cache_dtype=self.cache_config.cache_dtype, + runner_only_attn_layers=set(), + static_forward_context=self.compilation_config.static_forward_context, + ) + + def _zero_block_ids(self, block_ids: list[int]) -> None: + """Clear cache blocks before their first use after allocation.""" + if hasattr(self, "_kv_block_zeroer"): + self._kv_block_zeroer.zero_block_ids(block_ids) + @torch.inference_mode() @step_eplb_after(is_dummy=True) def _dummy_run( @@ -650,7 +702,10 @@ def capture_model(self) -> int: use_aux_hidden_state_outputs=self.use_aux_hidden_state_outputs, ) if self.speculator is not None: - self.speculator.capture(captured_attn_states) + if isinstance(self.speculator, DraftModelSpeculator): + self.speculator.capture() + else: + self.speculator.capture(captured_attn_states) end_time = time.perf_counter() end_free_gpu_memory = torch.cuda.mem_get_info()[0] @@ -820,6 +875,13 @@ def prepare_inputs( query_start_loc_np = query_start_loc_np[: num_reqs_padded + 1] query_start_loc = self.input_buffers.query_start_loc[: num_reqs_padded + 1] is_prefilling_np = self.req_states.is_prefilling(idx_mapping_np) + computed_prefill_lens = self.req_states.num_computed_prefill_tokens[ + idx_mapping_np + ] + prefill_lens = self.req_states.prefill_len.np[idx_mapping_np] + is_incomplete_prefilling_np = is_prefilling_np & ( + computed_prefill_lens + num_scheduled_tokens < prefill_lens + ) # Get prefill tokens if any. if np.any(is_prefilling_np): @@ -878,6 +940,16 @@ def prepare_inputs( out=seq_lens_cpu_upper_bound_np[:num_reqs], ) seq_lens_cpu_upper_bound = torch.from_numpy(seq_lens_cpu_upper_bound_np) + + prefix_anchor_lens = None + if self.prefix_anchor_lens_buffer is not None: + prefix_anchor_lens = self.prefix_anchor_lens_buffer[:num_reqs_padded] + prefix_anchor_lens[:num_reqs] = self.req_states.prompt_len.gpu[ + idx_mapping[:num_reqs] + ] + if num_reqs_padded > num_reqs: + prefix_anchor_lens[num_reqs:].zero_() + return InputBatch( req_ids=req_ids, num_reqs=num_reqs, @@ -903,6 +975,8 @@ def prepare_inputs( cu_num_logits=cu_num_logits, cu_num_logits_np=cu_num_logits_np, has_structured_output_reqs=scheduler_output.has_structured_output_requests, + prefix_anchor_lens=prefix_anchor_lens, + is_incomplete_prefilling_np=is_incomplete_prefilling_np, ) def prepare_attn( @@ -939,31 +1013,44 @@ def sample( grammar_output: GrammarOutput | None, ) -> tuple[SamplerOutput, torch.Tensor, torch.Tensor]: sample_hidden_states = hidden_states[input_batch.logits_indices] - logits = self.model.compute_logits(sample_hidden_states) - if grammar_output is not None: - # Apply grammar bitmask to the logits in-place. - assert self.structured_outputs_worker is not None - self.structured_outputs_worker.apply_grammar_bitmask( - logits, - input_batch, - grammar_output.structured_output_request_ids, - grammar_output.grammar_bitmask, - ) - - if input_batch.num_draft_tokens == 0: - # No draft tokens (common case). - assert self.sampler is not None - sampler_output = self.sampler(logits, input_batch) - else: - # Rejection sampling for spec decoding. + sampler_output = None + if input_batch.num_draft_tokens > 0: assert self.rejection_sampler is not None assert self.speculator is not None - sampler_output = self.rejection_sampler( - logits, + sampler_output = try_dflash2_sparse_target_rejection( + self.model, + self.speculator, + self.rejection_sampler, + sample_hidden_states, input_batch, - # Draft logits are needed for probabilistic rejection sampling. - self.speculator.draft_logits, + grammar_output, ) + if sampler_output is None: + logits = self.model.compute_logits(sample_hidden_states) + if grammar_output is not None: + # Apply grammar bitmask to the logits in-place. + assert self.structured_outputs_worker is not None + self.structured_outputs_worker.apply_grammar_bitmask( + logits, + input_batch, + grammar_output.structured_output_request_ids, + grammar_output.grammar_bitmask, + ) + + if input_batch.num_draft_tokens == 0: + # No draft tokens (common case). + assert self.sampler is not None + sampler_output = self.sampler(logits, input_batch) + else: + # Rejection sampling for spec decoding. + assert self.rejection_sampler is not None + assert self.speculator is not None + sampler_output = self.rejection_sampler( + logits, + input_batch, + # Draft logits are needed for probabilistic rejection sampling. + self.speculator.draft_logits, + ) # Get the number of sampled and rejected tokens. # For chunked prefills, num_sampled and num_rejected are both 0. @@ -1002,7 +1089,11 @@ def postprocess( self.req_states.total_len.gpu, ) - self.model_state.postprocess_state(input_batch, num_sampled) + self.model_state.postprocess_state( + input_batch, + num_sampled, + self.req_states.num_computed_tokens.gpu, + ) @torch.inference_mode() def execute_model( @@ -1020,6 +1111,8 @@ def execute_model( self.add_requests(scheduler_output) self.update_requests(scheduler_output) self.block_tables.apply_staged_writes() + if scheduler_output.new_block_ids_to_zero: + self._zero_block_ids(scheduler_output.new_block_ids_to_zero) if scheduler_output.total_num_scheduled_tokens == 0: # No need to run the model. empty_output = self.kv_connector.no_forward(scheduler_output) @@ -1058,6 +1151,14 @@ def execute_model( # Prepare all the inputs and copy to the input buffers. input_batch = self.prepare_inputs(scheduler_output, batch_desc) block_tables, slot_mappings = self.prepare_attn(input_batch) + # Hybrid Mamba align-mode prefix caching migrates recurrent state + # across block boundaries before attention metadata consumes it. + self.model_state.preprocess_state( + input_batch, + block_tables, + self.kv_cache_config, + self.req_states.num_computed_tokens.gpu, + ) if self.lora_config: # Activate LoRA adapters. @@ -1319,6 +1420,9 @@ def sample_tokens( if self.speculator is not None: assert self.sampler is not None + requires_host_token_state = bool( + getattr(self.speculator, "requires_host_token_state", False) + ) # Let the target override the hidden state fed to the drafter # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The # target returns a persistent buffer sized at max_num_batched_tokens; @@ -1340,6 +1444,24 @@ def sample_tokens( self.sampler.sampling_states.temperature.gpu, self.sampler.sampling_states.seeds.gpu, mm_inputs=mm_inputs, + output_copy_event=( + async_output.copy_event if requires_host_token_state else None + ), + sampled_token_ids_cpu=( + async_output.sampled_token_ids + if requires_host_token_state + else None + ), + num_sampled_tokens_cpu=( + async_output.num_sampled_tokens_np + if requires_host_token_state + else None + ), + all_token_ids_cpu=( + self.req_states.all_token_ids.get_cpu_view().numpy() + if requires_host_token_state + else None + ), ) self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens self.draft_tokens_handler.set_draft_tokens(input_batch, draft_tokens) diff --git a/vllm/v1/worker/gpu/model_states/default.py b/vllm/v1/worker/gpu/model_states/default.py index 7f7955a58a..f6b62dfcd5 100644 --- a/vllm/v1/worker/gpu/model_states/default.py +++ b/vllm/v1/worker/gpu/model_states/default.py @@ -195,5 +195,6 @@ def prepare_attn( dcp_local_seq_lens=input_batch.dcp_local_seq_lens, positions=input_batch.positions, for_cudagraph_capture=for_capture, + prefix_anchor_lens=input_batch.prefix_anchor_lens, ) return attn_metadata diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py index 721e5c2013..0e268af17b 100644 --- a/vllm/v1/worker/gpu/model_states/interface.py +++ b/vllm/v1/worker/gpu/model_states/interface.py @@ -56,10 +56,21 @@ def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: def apply_staged_writes(self) -> None: return None + def preprocess_state( + self, + input_batch: InputBatch, + block_tables: tuple[torch.Tensor, ...], + kv_cache_config: KVCacheConfig, + num_computed_tokens: torch.Tensor, + ) -> None: + """Hook for model-specific state movement before a real forward.""" + return None + def postprocess_state( self, input_batch: InputBatch, num_sampled: torch.Tensor, + num_computed_tokens: torch.Tensor | None = None, ) -> None: return None diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 93115fdf64..140cf97413 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -7,24 +7,64 @@ import torch import torch.nn as nn +from vllm import envs from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode -from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder +from vllm.logger import init_logger +from vllm.model_executor.layers.mamba.mamba_utils import ( + get_conv_copy_spec, + is_conv_state_dim_first, +) +from vllm.platforms import current_platform +from vllm.triton_utils import triton +from vllm.utils.platform_utils import is_pin_memory_available +from vllm.v1.attention.backends.flash_attn_v100 import ( + DFlash2SmallQGroupDescriptor, + DFlash2SmallQPreparedMetadata, + FlashAttnV100MetadataBuilder, + prepare_dflash2_smallq_group_metadata, +) +from vllm.v1.attention.backends.gdn_attn import ( + DFlash2GDNGroupDescriptor, + GDNAttentionMetadata, + GDNAttentionMetadataBuilder, + prepare_dflash2_gdn_group_metadata, +) from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadataBuilder -from vllm.v1.kv_cache_interface import KVCacheConfig -from vllm.v1.worker.gpu.attn_utils import build_attn_metadata +from vllm.v1.core.sched.output import NewRequestData +from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec +from vllm.v1.utils import CpuGpuBuffer +from vllm.v1.worker.gpu.attn_utils import ( + CommonGDNSpecMetadata, + build_attn_metadata, + compute_common_gdn_attn_metadata, +) from vllm.v1.worker.gpu.input_batch import InputBatch +from vllm.v1.worker.gpu.mamba_align import ( + preprocess_mamba_align_fused_kernel, + run_mamba_align_postprocess, + run_mamba_align_precopy, +) from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache 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.utils import AttentionGroup +logger = init_logger(__name__) + @dataclass class MambaHybridAttnMetadata(ModelSpecificAttnMetadata): is_prefilling: torch.Tensor num_accepted_tokens: torch.Tensor | None = None num_decode_draft_tokens_cpu: torch.Tensor | None = None + common_gdn_metadata: CommonGDNSpecMetadata | None = None + prepared_dflash2_gdn_metadata: dict[int, GDNAttentionMetadata] | None = None + prepared_dflash2_smallq_metadata: ( + dict[int, DFlash2SmallQPreparedMetadata] | None + ) = None def get_extra_common_attn_kwargs( self, @@ -38,12 +78,22 @@ def get_extra_attn_kwargs( attn_metadata_builder: Any, num_reqs: int, ) -> dict[str, Any]: + if isinstance(attn_metadata_builder, FlashAttnV100MetadataBuilder): + return { + "prepared_dflash2_smallq_metadata": ( + None + if self.prepared_dflash2_smallq_metadata is None + else self.prepared_dflash2_smallq_metadata.get( + id(attn_metadata_builder) + ) + ) + } if not isinstance( attn_metadata_builder, (Mamba2AttentionMetadataBuilder, GDNAttentionMetadataBuilder), ): return {} - return { + kwargs = { "num_accepted_tokens": None if self.num_accepted_tokens is None else self.num_accepted_tokens[:num_reqs], @@ -51,6 +101,13 @@ def get_extra_attn_kwargs( if self.num_decode_draft_tokens_cpu is None else self.num_decode_draft_tokens_cpu[:num_reqs], } + if isinstance(attn_metadata_builder, GDNAttentionMetadataBuilder): + kwargs["common_gdn_metadata"] = self.common_gdn_metadata + if self.prepared_dflash2_gdn_metadata is not None: + kwargs["prepared_dflash2_metadata"] = ( + self.prepared_dflash2_gdn_metadata.get(id(attn_metadata_builder)) + ) + return kwargs class MambaHybridModelState(DefaultModelState): @@ -64,9 +121,187 @@ def __init__( device: torch.device, ) -> None: super().__init__(vllm_config, model, encoder_cache, device) + self.cache_config = vllm_config.cache_config self.num_accepted_tokens_gpu = torch.ones( self.max_num_reqs, dtype=torch.int32, device=self.device ) + self._use_dflash2_common_gdn_metadata = bool( + envs.VLLM_SM70_DFLASH2_VERIFY_FASTPATH + and uses_dflash_selector_engine(vllm_config) + ) + if self._use_dflash2_common_gdn_metadata: + logger.info_once("DFlash2 shared GDN batch metadata fast path enabled.") + self._use_dflash2_fused_gdn_metadata = bool( + self._use_dflash2_common_gdn_metadata + and envs.VLLM_SM70_DFLASH2_FUSED_GDN_METADATA + and self.cache_config.mamba_cache_mode in ("none", "align") + and device.type == "cuda" + and current_platform.is_device_capability(70) + ) + self._dflash2_gdn_builders: ( + list[tuple[int, GDNAttentionMetadataBuilder]] | None + ) = None + self._dflash2_gdn_group_descriptor: DFlash2GDNGroupDescriptor | None = None + self._dflash2_fused_gdn_metadata_logged = False + self._use_dflash2_grouped_smallq_metadata = bool( + self._use_dflash2_common_gdn_metadata + and envs.VLLM_SM70_DFLASH2_FUSED_SMALLQ_METADATA + and envs.VLLM_SM70_DFLASH2_GROUPED_SMALLQ_METADATA + and device.type == "cuda" + and current_platform.is_device_capability(70) + ) + self._dflash2_smallq_builders: ( + list[tuple[int, FlashAttnV100MetadataBuilder]] | None + ) = None + self._dflash2_smallq_group_descriptor: DFlash2SmallQGroupDescriptor | None = ( + None + ) + self._dflash2_grouped_smallq_metadata_logged = False + self._align_mode = self.cache_config.mamba_cache_mode == "align" + if self._align_mode: + self._mamba_state_idx_gpu = torch.full( + (self.max_num_reqs,), -2, dtype=torch.int32, device=self.device + ) + self._mamba_src_col_gpu = torch.full( + (self.max_num_reqs,), -1, dtype=torch.int32, device=self.device + ) + self._mamba_token_bias_gpu = torch.zeros( + self.max_num_reqs, dtype=torch.int32, device=self.device + ) + self._mamba_ctx: MambaSpecDecodeGPUContext | None = None + self._mamba_group_ids: list[int] = [] + self._mamba_spec: MambaSpec | None = None + + def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: + super().add_request(req_index, new_req_data) + # A recycled request slot may contain the previous request's accepted + # count. The neutral value is required in both align and non-align mode. + self.num_accepted_tokens_gpu[req_index].fill_(1) + if self._align_mode: + # Delay state-column seeding until preprocess_state has the + # resolved MambaSpec block size. -2 is distinct from -1 (fresh + # request with no computed state). + self._mamba_state_idx_gpu[req_index].fill_(-2) + + 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] + return self._mamba_group_ids, self._mamba_spec + + def _ensure_align_ctx( + self, + kv_cache_config: KVCacheConfig, + mamba_group_ids: list[int], + block_tables: tuple[torch.Tensor, ...], + ) -> MambaSpecDecodeGPUContext: + if self._mamba_ctx is None: + copy_funcs = self.model.get_mamba_state_copy_func() + # 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(): + raise ValueError( + "MRV2 align prefix caching with speculative decoding requires " + "the default SD Mamba conv-state layout on this V100 path." + ) + self._mamba_ctx = MambaSpecDecodeGPUContext.create( + max_num_reqs=self.max_num_reqs, + kv_cache_config=kv_cache_config, + num_state_types=len(copy_funcs), + device=self.device, + make_buffer=lambda n, dtype: CpuGpuBuffer( + n, + dtype=dtype, + device=self.device, + pin_memory=is_pin_memory_available(), + ), + ) + ctx = self._mamba_ctx + if not ctx.is_initialized: + 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(), + [block_tables[group_id] for group_id in mamba_group_ids], + ) + return ctx + + def preprocess_state( + self, + input_batch: InputBatch, + block_tables: tuple[torch.Tensor, ...], + kv_cache_config: KVCacheConfig, + num_computed_tokens: torch.Tensor, + ) -> None: + """Move align-mode recurrent state before a real MRV2 forward.""" + if not self._align_mode or input_batch.num_reqs == 0: + return + mamba_group_ids, mamba_spec = self._get_mamba_group_info(kv_cache_config) + ctx = self._ensure_align_ctx(kv_cache_config, mamba_group_ids, block_tables) + + block = 256 + preprocess_mamba_align_fused_kernel[ + (triton.cdiv(input_batch.num_reqs, block),) + ]( + input_batch.idx_mapping, + self._mamba_state_idx_gpu, + num_computed_tokens, + input_batch.query_start_loc, + self.num_accepted_tokens_gpu, + self._mamba_src_col_gpu, + self._mamba_token_bias_gpu, + input_batch.num_reqs, + BLOCK_SIZE=block, + MAMBA_BLOCK_SIZE=mamba_spec.block_size, + ) + run_mamba_align_precopy( + ctx, + input_batch.num_reqs, + self._mamba_state_idx_gpu, + self._mamba_src_col_gpu, + self._mamba_token_bias_gpu, + input_batch.idx_mapping, + ) + + def _get_dflash2_gdn_builders( + self, + attn_groups: list[list[AttentionGroup]], + ) -> list[tuple[int, GDNAttentionMetadataBuilder]]: + if self._dflash2_gdn_builders is None: + builders: list[tuple[int, GDNAttentionMetadataBuilder]] = [] + for kv_cache_group_id, groups in enumerate(attn_groups): + for group in groups: + builder = group.get_metadata_builder(0) + if isinstance(builder, GDNAttentionMetadataBuilder): + builders.append((kv_cache_group_id, builder)) + self._dflash2_gdn_builders = builders + return self._dflash2_gdn_builders + + def _get_dflash2_smallq_builders( + self, + attn_groups: list[list[AttentionGroup]], + ) -> list[tuple[int, FlashAttnV100MetadataBuilder]]: + if self._dflash2_smallq_builders is None: + builders: list[tuple[int, FlashAttnV100MetadataBuilder]] = [] + for kv_cache_group_id, groups in enumerate(attn_groups): + for group in groups: + builder = group.get_metadata_builder(0) + if isinstance(builder, FlashAttnV100MetadataBuilder): + builders.append((kv_cache_group_id, builder)) + self._dflash2_smallq_builders = builders + return self._dflash2_smallq_builders def prepare_attn( self, @@ -96,6 +331,9 @@ def prepare_attn( # compute them during actual (non-capture) forward execution. num_accepted_tokens = None num_decode_draft_tokens_cpu = None + common_gdn_metadata = None + prepared_dflash2_gdn_metadata = None + prepared_dflash2_smallq_metadata = None if not for_capture: num_accepted_tokens = self.num_accepted_tokens_gpu.new_ones(num_reqs) num_accepted_tokens[: input_batch.num_reqs] = self.num_accepted_tokens_gpu[ @@ -115,13 +353,93 @@ def prepare_attn( -1, ) num_decode_draft_tokens_cpu = torch.from_numpy(num_decode_draft_tokens_np) + if self._use_dflash2_common_gdn_metadata: + speculative_config = self.vllm_config.speculative_config + assert speculative_config is not None + common_gdn_metadata = compute_common_gdn_attn_metadata( + num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, + query_start_loc=input_batch.query_start_loc, + query_start_loc_cpu=query_start_loc_cpu, + num_spec_state_tokens=( + speculative_config.num_speculative_state_tokens() + ), + legacy_mixed_decode_routing=( + envs.VLLM_SM70_MTP_LEGACY_GDN_MIXED_DECODE_ROUTING + ), + ) + + if ( + self._use_dflash2_fused_gdn_metadata + and cudagraph_mode == CUDAGraphMode.FULL + and common_gdn_metadata is not None + ): + prepared_result = prepare_dflash2_gdn_group_metadata( + builders_by_group=self._get_dflash2_gdn_builders(attn_groups), + block_tables=block_tables, + common_gdn_metadata=common_gdn_metadata, + num_accepted_tokens=num_accepted_tokens, + num_actual_tokens=num_tokens, + descriptor=self._dflash2_gdn_group_descriptor, + state_start_indices=( + self._mamba_state_idx_gpu if self._align_mode else None + ), + req_index_mapping=( + input_batch.idx_mapping if self._align_mode else None + ), + ) + if prepared_result is not None: + ( + prepared_dflash2_gdn_metadata, + self._dflash2_gdn_group_descriptor, + ) = prepared_result + if not self._dflash2_fused_gdn_metadata_logged: + logger.info( + "DFlash2 fused GDN metadata active for %d cache groups.", + len(prepared_dflash2_gdn_metadata), + ) + self._dflash2_fused_gdn_metadata_logged = True + + if ( + self._use_dflash2_grouped_smallq_metadata + and cudagraph_mode == CUDAGraphMode.FULL + and common_gdn_metadata is not None + ): + seq_lens_cpu = input_batch.seq_lens_cpu_upper_bound[:num_reqs] + max_seq_len_hint = int(seq_lens_cpu.max().item()) + prepared_smallq_result = prepare_dflash2_smallq_group_metadata( + builders_by_group=self._get_dflash2_smallq_builders(attn_groups), + block_tables=block_tables, + seq_lens=input_batch.seq_lens, + query_start_loc=input_batch.query_start_loc, + query_start_loc_cpu=query_start_loc_cpu, + num_reqs=num_reqs, + num_query_tokens=num_tokens, + max_seq_len_hint=max_seq_len_hint, + workspace_seq_capacity_cap=self.max_model_len, + descriptor=self._dflash2_smallq_group_descriptor, + ) + if prepared_smallq_result is not None: + ( + prepared_dflash2_smallq_metadata, + self._dflash2_smallq_group_descriptor, + ) = prepared_smallq_result + if not self._dflash2_grouped_smallq_metadata_logged: + logger.info( + "DFlash2 grouped small-query metadata active for %d " + "full-attention cache groups.", + len(prepared_dflash2_smallq_metadata), + ) + self._dflash2_grouped_smallq_metadata_logged = True mamba_attn_metadata = MambaHybridAttnMetadata( is_prefilling=is_prefilling, num_accepted_tokens=num_accepted_tokens, num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, + common_gdn_metadata=common_gdn_metadata, + prepared_dflash2_gdn_metadata=prepared_dflash2_gdn_metadata, + prepared_dflash2_smallq_metadata=prepared_dflash2_smallq_metadata, ) - return build_attn_metadata( + attn_metadata = build_attn_metadata( attn_groups=attn_groups, num_reqs=num_reqs, num_tokens=num_tokens, @@ -133,18 +451,49 @@ def prepare_attn( block_tables=block_tables, slot_mappings=slot_mappings, kv_cache_config=kv_cache_config, + seq_lens_cpu_upper_bound=input_batch.seq_lens_cpu_upper_bound, dcp_local_seq_lens=input_batch.dcp_local_seq_lens, model_specific_attn_metadata=mamba_attn_metadata, for_cudagraph_capture=for_capture, + prefix_anchor_lens=input_batch.prefix_anchor_lens, ) + if common_gdn_metadata is not None: + assert common_gdn_metadata.spec_query_start_loc.numel() == ( + common_gdn_metadata.num_spec_decodes + 1 + ) + # Metadata kernels/copies and FULL CUDA Graph replay are enqueued on + # the same current stream, which already supplies the required + # device ordering. Retain the old value check as an opt-in debug + # fence while the no-sync route is quality- and trace-gated. + if envs.VLLM_SM70_DFLASH2_GDN_SYNC_ASSERT: + assert ( + common_gdn_metadata.spec_query_start_loc[-1].item() + == common_gdn_metadata.num_spec_decode_tokens + ) + return attn_metadata def postprocess_state( self, input_batch: InputBatch, num_sampled: torch.Tensor, + num_computed_tokens: torch.Tensor | None = None, ) -> None: # Chunked prefill does not sample a token, so num_sampled can be 0. # Mamba treats num_accepted_tokens=1 as the neutral non-spec value. self.num_accepted_tokens_gpu[input_batch.idx_mapping] = torch.clamp( num_sampled, min=1 ) + if ( + self._align_mode + and num_computed_tokens is not None + and self._mamba_ctx is not None + and input_batch.num_reqs > 0 + ): + run_mamba_align_postprocess( + self._mamba_ctx, + input_batch.num_reqs, + self.num_accepted_tokens_gpu, + self._mamba_state_idx_gpu, + num_computed_tokens, + input_batch.idx_mapping, + ) diff --git a/vllm/v1/worker/gpu/model_states/whisper.py b/vllm/v1/worker/gpu/model_states/whisper.py index 0ef3cadc87..191b930a2e 100644 --- a/vllm/v1/worker/gpu/model_states/whisper.py +++ b/vllm/v1/worker/gpu/model_states/whisper.py @@ -161,6 +161,7 @@ def prepare_attn( dcp_local_seq_lens=input_batch.dcp_local_seq_lens, model_specific_attn_metadata=whisper_attn_metadata, for_cudagraph_capture=for_capture, + prefix_anchor_lens=input_batch.prefix_anchor_lens, ) return attn_metadata diff --git a/vllm/v1/worker/gpu/sample/gumbel.py b/vllm/v1/worker/gpu/sample/gumbel.py index aaa49283d3..77fb4a16f2 100644 --- a/vllm/v1/worker/gpu/sample/gumbel.py +++ b/vllm/v1/worker/gpu/sample/gumbel.py @@ -2,18 +2,16 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch -from vllm.triton_utils import HAS_TRITON, tl, triton +from vllm.triton_utils import HAS_TRITON, tl, tldevice, triton -# Smallest positive normal fp32 value. Used to clamp the uniform draw so that -# `log(u)` cannot produce -inf (and thus `-log(-log(u))` stays finite). +# Smallest positive value produced by Triton's fp32 `tl.rand`. Used to clamp +# zero draws before the flipped Gumbel transform below. # # Triton requires globals accessed from `@triton.jit` functions to be wrapped # in `tl.constexpr(...)`. We can only do that when Triton is actually # available — on the CPU worker path `tl` is a placeholder whose `constexpr` # attribute is `None`, and `tl.constexpr(...)` would crash at import time. -_FP32_TINY = ( - tl.constexpr(float.fromhex("0x1p-126")) if HAS_TRITON else float.fromhex("0x1p-126") -) +_TL_RAND_MIN = tl.constexpr(4.6566127342e-10) if HAS_TRITON else 4.6566127342e-10 @triton.jit @@ -75,6 +73,46 @@ def tl_rand64(seed, offset, includes_zero: tl.constexpr): return u +@triton.jit +def tl_rand32(seed, offset, includes_zero: tl.constexpr): + u = tl.rand(seed, offset) + if not includes_zero: + u = tl.maximum(u, _TL_RAND_MIN) + return u + + +@triton.jit +def gumbel_noised_argmax( + logits, + keys, + mask, + seed, + pos, + temp, + USE_FP64: tl.constexpr, + APPLY_TEMPERATURE: tl.constexpr = True, +): + """Argmax under token-keyed Gumbel noise, or plain argmax at temp zero.""" + if temp != 0.0 and APPLY_TEMPERATURE: + logits = logits / temp + + if USE_FP64: + logits = logits.to(tl.float64) + if temp != 0.0: + gumbel_seed = tl.randint(seed, pos) + if USE_FP64: + u = tl_rand64(gumbel_seed, keys, includes_zero=False) + gumbel_noise = -tl.log(-tl.log(u)) + else: + u = tl_rand32(gumbel_seed, keys, includes_zero=False) + # Flip the transform so the deciding tail lies near zero, where + # fp32 uniforms retain resolution. `log1p` avoids cancellation. + gumbel_noise = -tl.log(-tldevice.log1p(-u)) + logits = tl.where(mask, logits + gumbel_noise, float("-inf")) + + return tl.max(logits, axis=0, return_indices=True) + + @triton.jit def gumbel_block_argmax( logits, @@ -87,21 +125,24 @@ def gumbel_block_argmax( pos_ptr, processed_logits_ptr, processed_logits_stride, + processed_logits_col_stride, processed_logits_col_ptr, vocab_size, APPLY_TEMPERATURE: tl.constexpr, USE_FP64: tl.constexpr, ): - req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx) - temp = tl.load(temp_ptr + req_state_idx).to(tl.float32) + req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx).to(tl.int64) + is_valid_req = req_state_idx >= 0 + temp = tl.load(temp_ptr + req_state_idx, mask=is_valid_req, other=0.0).to( + tl.float32 + ) if temp != 0.0 and APPLY_TEMPERATURE: - # Apply temperature. - # NOTE(woosuk): Match the behavior of _temperature_kernel. - # E.g., if the kernel uses tl.div_rn, we should use tl.div_rn here too. + # Preserve the established MRV2 contract: proposal logits handed to + # rejection sampling have already had request temperature applied. logits = logits / temp if processed_logits_ptr is not None: - # Store the temperature-applied logits. + # Store the temperature-applied proposal distribution. if processed_logits_col_ptr is not None: col = tl.load(processed_logits_col_ptr) else: @@ -109,34 +150,24 @@ def gumbel_block_argmax( tl.store( processed_logits_ptr + req_state_idx * processed_logits_stride - + col * vocab_size + + col * processed_logits_col_stride + block, logits, - mask=mask, + mask=mask & is_valid_req, ) - # fp32 is the default reduction dtype; fp64 is ~1/32–1/64x the throughput - # on H100/Ada/Blackwell and empirically indistinguishable for Gumbel-max. - if USE_FP64: - logits = logits.to(tl.float64) - if temp != 0.0: - # Calculate the seed for gumbel noise. - seed = tl.load(seeds_ptr + req_state_idx) - pos = tl.load(pos_ptr + token_idx) - gumbel_seed = tl.randint(seed, pos) - - if USE_FP64: - u = tl_rand64(gumbel_seed, block, includes_zero=False) - else: - u = tl.rand(gumbel_seed, block) - u = tl.maximum(u, _FP32_TINY) - gumbel_noise = -tl.log(-tl.log(u)) - - # Apply gumbel noise. - logits = tl.where(mask, logits + gumbel_noise, float("-inf")) - - value, idx = tl.max(logits, axis=0, return_indices=True) - return value, idx + seed = tl.load(seeds_ptr + req_state_idx, mask=is_valid_req, other=0) + pos = tl.load(pos_ptr + token_idx) + return gumbel_noised_argmax( + logits, + block, + mask & is_valid_req, + seed, + pos, + temp, + USE_FP64=USE_FP64, + APPLY_TEMPERATURE=False, + ) @triton.jit @@ -147,6 +178,7 @@ def _gumbel_sample_kernel( local_max_stride, processed_logits_ptr, processed_logits_stride, + processed_logits_col_stride, processed_logits_col_ptr, logits_ptr, logits_stride, @@ -181,6 +213,7 @@ def _gumbel_sample_kernel( pos_ptr, processed_logits_ptr, processed_logits_stride, + processed_logits_col_stride, processed_logits_col_ptr, vocab_size, APPLY_TEMPERATURE=APPLY_TEMPERATURE, @@ -215,6 +248,7 @@ def gumbel_sample( local_max.stride(0), output_processed_logits, output_processed_logits.stride(0) if output_processed_logits is not None else 0, + output_processed_logits.stride(1) if output_processed_logits is not None else 0, output_processed_logits_col, logits, logits.stride(0), diff --git a/vllm/v1/worker/gpu/spec_decode/__init__.py b/vllm/v1/worker/gpu/spec_decode/__init__.py index 536b7526bd..ee25697099 100644 --- a/vllm/v1/worker/gpu/spec_decode/__init__.py +++ b/vllm/v1/worker/gpu/spec_decode/__init__.py @@ -1,13 +1,51 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Mapping + import torch from vllm.config import VllmConfig +def uses_dflash_selector_engine(vllm_config: object) -> bool: + """Return whether the configured DFlash algorithm exposes selector inputs. + + Dispatch is based on the inference algorithm's configuration contract, not + a model, architecture, or checkpoint identity. + """ + speculative_config = getattr(vllm_config, "speculative_config", None) + if ( + speculative_config is None + or getattr(speculative_config, "method", None) != "dflash" + ): + return False + 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 {} + if not isinstance(dflash_config, Mapping): + return False + try: + return int(dflash_config.get("selector_top_k", 0) or 0) > 0 + except (TypeError, ValueError): + return False + + def init_speculator(vllm_config: VllmConfig, device: torch.device): speculative_config = vllm_config.speculative_config assert speculative_config is not None + if speculative_config.method == "dflash": + draft_config = speculative_config.draft_model_config + if draft_config is None: + raise ValueError("method='dflash' requires a draft model config") + if uses_dflash_selector_engine(vllm_config): + from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import ( + DFlash2Speculator, + ) + + return DFlash2Speculator(vllm_config, device) + from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator + + return DFlashSpeculator(vllm_config, device) if speculative_config.use_eagle(): from vllm.v1.worker.gpu.spec_decode.eagle.speculator import EagleSpeculator diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/__init__.py b/vllm/v1/worker/gpu/spec_decode/dflash/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/vllm/v1/worker/gpu/spec_decode/dflash/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py b/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py new file mode 100644 index 0000000000..cfbfc5aeb9 --- /dev/null +++ b/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Callable, Mapping + +import torch + +from vllm.config.compilation import CUDAGraphMode +from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.worker.gpu.attn_utils import ( + build_attn_metadata, + build_slot_mappings_by_layer, +) +from vllm.v1.worker.gpu.block_table import BlockTables +from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens +from vllm.v1.worker.gpu.cudagraph_utils import ( + BatchExecutionDescriptor, + CapturedAttentionState, + CudaGraphManager, +) +from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers +from vllm.v1.worker.utils import AttentionGroup + + +def _prepare_dflash_inputs_to_capture( + num_reqs: int, + num_tokens: int, + input_buffers: InputBuffers, + block_tables: BlockTables, + attn_groups: list[list[AttentionGroup]], + kv_cache_config: KVCacheConfig, + max_model_len: int, + skip_attn: bool, + causal: bool | Mapping[int, bool], +) -> CapturedAttentionState: + input_batch = InputBatch.make_dummy(num_reqs, num_tokens, input_buffers) + input_block_tables = block_tables.get_dummy_block_tables(num_reqs) + slot_mappings = block_tables.get_dummy_slot_mappings(num_tokens) + slot_mappings_by_layer = build_slot_mappings_by_layer( + slot_mappings, kv_cache_config + ) + + attn_metadata = None + if not skip_attn: + query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np) + dcp_local_seq_lens = None + if block_tables.cp_size > 1: + prepare_dcp_local_seq_lens( + input_buffers.dcp_local_seq_lens, + input_buffers.seq_lens, + num_reqs, + block_tables.cp_size, + block_tables.cp_rank, + block_tables.cp_interleave, + ) + dcp_local_seq_lens = input_buffers.dcp_local_seq_lens + attn_metadata = build_attn_metadata( + attn_groups=attn_groups, + num_reqs=num_reqs, + num_tokens=num_tokens, + query_start_loc_gpu=input_batch.query_start_loc, + query_start_loc_cpu=query_start_loc_cpu, + max_query_len=num_tokens // num_reqs, + seq_lens=input_batch.seq_lens, + dcp_local_seq_lens=dcp_local_seq_lens, + max_seq_len=max_model_len, + block_tables=input_block_tables, + slot_mappings=slot_mappings, + kv_cache_config=kv_cache_config, + for_cudagraph_capture=True, + causal=causal, + ) + return CapturedAttentionState(attn_metadata, slot_mappings_by_layer) + + +class DFlashCudaGraphManager(CudaGraphManager): + """DFlash CudaGraphManager for the parallel-drafting query forward, + building its own attention metadata from scratch.""" + + def capture( + self, + forward_fn: Callable, + input_buffers: InputBuffers, + block_tables: BlockTables, + attn_groups: list[list[AttentionGroup]], + kv_cache_config: KVCacheConfig, + max_model_len: int, + causal: bool | Mapping[int, bool], + progress_bar_desc: str = "Capturing CUDA graphs", + ) -> None: + def create_forward_fn( + desc: BatchExecutionDescriptor, + ) -> tuple[Callable[[CUDAGraphMode], None], CapturedAttentionState]: + num_tokens = desc.num_tokens + num_reqs = desc.num_reqs or min(num_tokens, self.max_num_reqs) + num_tokens_across_dp = ( + torch.full((self.dp_size,), num_tokens, dtype=torch.int32, device="cpu") + if self.dp_size > 1 + else None + ) + attn_state = _prepare_dflash_inputs_to_capture( + num_reqs, + num_tokens, + input_buffers, + block_tables, + attn_groups, + kv_cache_config, + max_model_len, + skip_attn=(desc.cg_mode == CUDAGraphMode.PIECEWISE), + causal=causal, + ) + attn_metadata, slot_mappings = attn_state + + return ( + lambda cg_mode: forward_fn( + num_reqs, + num_tokens, + attn_metadata, + slot_mappings, + num_tokens_across_dp, + cg_mode, + ), + attn_state, + ) + + super().capture(create_forward_fn, progress_bar_desc) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py new file mode 100644 index 0000000000..9d08b77ee5 --- /dev/null +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -0,0 +1,829 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Mapping +from typing import Any + +import numpy as np +import torch +import torch.nn as nn + +from vllm.config import VllmConfig, replace +from vllm.config.compilation import CUDAGraphMode +from vllm.forward_context import BatchDescriptor, set_forward_context +from vllm.logger import init_logger +from vllm.triton_utils import tl, triton +from vllm.v1.attention.backend import AttentionCGSupport +from vllm.v1.attention.backends.utils import PAD_SLOT_ID +from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.utils import record_function_or_nullcontext +from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer +from vllm.v1.worker.gpu.block_table import BlockTables +from vllm.v1.worker.gpu.cp_utils import cp_local_slot, prepare_dcp_local_seq_lens +from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp +from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers +from vllm.v1.worker.gpu.model_states.interface import ModelState +from vllm.v1.worker.gpu.spec_decode.dflash.cudagraph import DFlashCudaGraphManager +from vllm.v1.worker.gpu.spec_decode.dflash.utils import load_dflash_model +from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator +from vllm.v1.worker.gpu.spec_decode.utils import get_parallel_drafting_token_id +from vllm.v1.worker.utils import AttentionGroup + +logger = init_logger(__name__) + + +def _is_context_only_prefill(input_batch: InputBatch) -> bool: + incomplete_prefill = input_batch.is_incomplete_prefilling_np + return bool( + incomplete_prefill is not None + and incomplete_prefill.size == input_batch.num_reqs + and np.all(incomplete_prefill) + ) + + +class DFlashSpeculator(DraftModelSpeculator): + _speculator_name = "DFlash" # For logging, so we can share methods with subclasses + + def __init__(self, vllm_config: VllmConfig, device: torch.device): + super().__init__(vllm_config, device) + + self.hidden_states = torch.zeros( + self.max_num_tokens, self.hidden_size, dtype=self.dtype, device=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 + + self.parallel_drafting_token_id = get_parallel_drafting_token_id( + self.draft_model_config.hf_config + ) + + from vllm.model_executor.models.qwen3_dflash import dflash_has_any_non_causal + + self.requires_non_causal = dflash_has_any_non_causal( + self.draft_model_config.hf_config + ) + + # Whether the anchor query position is itself a prediction. DFlash default uses + # the anchor as the bonus token (only mask tokens predict); DSpark samples from + # the anchor and the N-1 mask token positions. See _prepare_dflash_inputs_kernel + dflash_config = ( + getattr(self.draft_model_config.hf_config, "dflash_config", None) or {} + ) + if dflash_config.get("sample_from_anchor", False): + raise ValueError( + "sample_from_anchor=True is not supported for DFlash. " + "DFlash uses a fixed 1+N query layout where the anchor " + "is the bonus token." + ) + self.sample_from_anchor = False + + # Context positions for the K/V precompute. Populated by + # prepare_dflash_inputs, and processed by the model's + # precompute_and_store_context_kv method. NOT captured by CUDA graphs. + self.context_positions = torch.zeros( + self.max_num_tokens, dtype=torch.int64, device=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 + self.sample_indices = torch.zeros( + max_num_sampled_tokens, dtype=torch.int64, device=device + ) + self.sample_pos = torch.zeros( + max_num_sampled_tokens, dtype=torch.int64, device=device + ) + # -1 marks an inert sampling row. CUDA graph capture can execute the + # full buffer before a real batch has populated it, so zero would make + # every padding row scatter into request slot 0. + self.sample_idx_mapping = torch.full( + (max_num_sampled_tokens,), -1, dtype=torch.int32, device=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 + ).repeat(self.max_num_reqs) + + self.query_cudagraph_manager: DFlashCudaGraphManager | None = None + self.draft_kv_cache_group_id: int = -1 + self._context_only_prefill_logged = False + + @property + def attn_vllm_config(self) -> VllmConfig: + # The draft's attention differs from the target's in causality. + return replace( + self.vllm_config, + model_config=self.draft_model_config, + attention_config=replace( + self.vllm_config.attention_config, + use_non_causal=self.requires_non_causal, + ), + ) + + def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: + wants_full = cudagraph_mode.decode_mode() == CUDAGraphMode.FULL + supports_full = ( + self.attn_cg_support.min_cg_support.value + >= AttentionCGSupport.UNIFORM_BATCH.value + ) + if wants_full and not supports_full: + logger.warning( + "%s draft attention (%s) does not support full CUDA graphs; " + "running the draft eagerly.", + self._speculator_name, + self.attn_cg_support.min_cg_attn_backend, + ) + # PIECEWISE cudagraphs are not supported for dflash. + if wants_full and supports_full: + cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY + else: + cudagraph_mode = CUDAGraphMode.NONE + + self.query_cudagraph_manager = DFlashCudaGraphManager( + self.vllm_config, + self.device, + cudagraph_mode, + decode_query_len=self.num_query_per_req, + ) + + def capture(self) -> None: + logger.info("Capturing model for %s speculator...", self._speculator_name) + # Padded sample rows must not scatter into a live request during capture. + self.sample_indices.zero_() + self.sample_pos.zero_() + self.sample_idx_mapping.fill_(-1) + assert self.query_cudagraph_manager is not None + self.query_cudagraph_manager.capture( + self._generate_draft, + self.input_buffers, + self.block_tables, + self.attn_groups, + self.kv_cache_config, + self.max_model_len, + causal=self._group_causal, + progress_bar_desc=f"Capturing {self._speculator_name.lower()} CUDA graphs", + ) + + def load_draft_model( + self, + target_model: nn.Module, + target_attn_layer_names: set[str], + ) -> nn.Module: + return load_dflash_model(target_model, self.vllm_config) + + def set_attn( + self, + model_state: ModelState, + kv_cache_config: KVCacheConfig, + block_tables: BlockTables, + target_input_buffers: InputBuffers, + target_attn_groups: list[list[AttentionGroup]], + ) -> None: + super().set_attn( + model_state, + kv_cache_config, + block_tables, + target_input_buffers, + target_attn_groups, + ) + + self.draft_kv_cache_group_ids = [ + gid for gid, g in enumerate(self.attn_groups) if g + ] + assert self.draft_kv_cache_group_ids, "No draft attention groups found." + self.draft_kv_cache_group_id = self.draft_kv_cache_group_ids[0] + + # Per-group context slot buffers for the precompute (one row per group). + self._context_slot_mappings = torch.zeros( + len(self.draft_kv_cache_group_ids), + self.max_num_tokens, + dtype=torch.int64, + device=self.device, + ) + + # Map each draft decoder layer to the index (within draft_kv_cache_group_ids) + # of the kv-cache group its cache belongs to. Models that share a single group + # leave this as None and share one context slot mapping. + self._layer_group_idx: list[int] | None = None + # Per-KV-group causal, falling back to whether the drafter is all-causal. + self._group_causal: dict[int, bool] | bool = not self.requires_non_causal + if hasattr(self.model, "get_draft_kv_cache_layer_names"): + layer_names = self.model.get_draft_kv_cache_layer_names() + name_to_gid = { + ln: gid + for gid, group in enumerate(kv_cache_config.kv_cache_groups) + for ln in group.layer_names + } + gid_to_idx = {gid: i for i, gid in enumerate(self.draft_kv_cache_group_ids)} + self._layer_group_idx = [ + gid_to_idx[name_to_gid[name]] for name in layer_names + ] + if hasattr(self.model, "get_draft_attn_causal"): + self._group_causal = { + name_to_gid[name]: layer_causal + for name, layer_causal in zip( + layer_names, self.model.get_draft_attn_causal() + ) + } + + @torch.inference_mode() + def _run_model( + self, + num_tokens: int, + attn_metadata: dict[str, Any] | None, + slot_mappings: dict[str, torch.Tensor] | None, + num_tokens_across_dp: torch.Tensor | None, + cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE, + ) -> torch.Tensor: + batch_descriptor = BatchDescriptor(num_tokens=num_tokens) + with set_forward_context( + attn_metadata, + self.vllm_config, + num_tokens=num_tokens, + cudagraph_runtime_mode=cudagraph_runtime_mode, + num_tokens_across_dp=num_tokens_across_dp, + slot_mapping=slot_mappings, + batch_descriptor=batch_descriptor, + ): + last_hidden_states = self.model( + input_ids=self.input_buffers.input_ids[:num_tokens], + positions=self.input_buffers.positions[:num_tokens], + inputs_embeds=None, + ) + return last_hidden_states + + def _generate_draft( + self, + num_reqs: int, + num_tokens_padded: int, + attn_metadata: dict[str, Any] | None, + slot_mappings: dict[str, torch.Tensor] | None, + num_tokens_across_dp: torch.Tensor | None, + cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE, + ) -> None: + last_hidden_states = self._run_model( + num_tokens_padded, + attn_metadata, + slot_mappings, + num_tokens_across_dp, + cudagraph_runtime_mode, + ) + num_sample = num_reqs * self.num_speculative_steps + 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. + draft_tokens = self.sample_draft( + sample_hidden_states, + self.sample_pos[:num_sample] - 2, + self.sample_idx_mapping[:num_sample], + self.temperature, + self.seeds, + self.sample_col[:num_sample], + self.draft_logits, + ) + self.draft_tokens[:num_reqs] = draft_tokens.view( + num_reqs, self.num_speculative_steps + ) + + def _build_draft_attn_metadata( + self, + num_reqs: int, + num_reqs_padded: int, + num_tokens_padded: int, + seq_lens_cpu_upper_bound: torch.Tensor, + step: int, + num_query_per_req: int | None = None, + causal: bool | Mapping[int, bool] = False, + query_start_loc_np: np.ndarray | None = None, + dcp_local_seq_lens: torch.Tensor | None = None, + ) -> dict[str, Any] | None: + if not self.draft_attn_layer_names: + return None + assert num_query_per_req is None # Omitted for DFlash, read from self instead + if dcp_local_seq_lens is None and self.block_tables.cp_size > 1: + prepare_dcp_local_seq_lens( + self.input_buffers.dcp_local_seq_lens, + self.input_buffers.seq_lens, + num_reqs, + self.block_tables.cp_size, + self.block_tables.cp_rank, + self.block_tables.cp_interleave, + ) + dcp_local_seq_lens = self.input_buffers.dcp_local_seq_lens + return super()._build_draft_attn_metadata( + num_reqs, + num_reqs_padded, + num_tokens_padded, + seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, + step=step, + num_query_per_req=self.num_query_per_req, + causal=causal, + query_start_loc_np=query_start_loc_np, + dcp_local_seq_lens=dcp_local_seq_lens, + ) + + def _prepare_ngram_assist( + self, + input_batch: InputBatch, + output_copy_event: torch.cuda.Event | None, + sampled_token_ids_cpu: np.ndarray | None, + num_sampled_tokens_cpu: np.ndarray | None, + all_token_ids_cpu: np.ndarray | None, + ) -> bool: + """Prepare an optional draftless proposal; return whether it is complete.""" + return False + + def _apply_ngram_assist(self, num_reqs: int) -> None: + """Override model proposals for ngram-hit rows, if configured.""" + + @torch.inference_mode() + def propose( + self, + input_batch: InputBatch, + attn_metadata: dict[str, Any], + slot_mappings: dict[str, torch.Tensor], + # [num_tokens, hidden_size] + last_hidden_states: torch.Tensor, + # num_layers x [num_tokens, hidden_size] + aux_hidden_states: list[torch.Tensor] | None, + # [num_reqs] + num_sampled: torch.Tensor, + # [num_reqs] + num_rejected: torch.Tensor, + # [max_num_reqs] + last_sampled: torch.Tensor, + # [max_num_reqs] + next_prefill_tokens: torch.Tensor, + # [max_num_reqs] + temperature: torch.Tensor, + # [max_num_reqs] + seeds: torch.Tensor, + num_tokens_across_dp: torch.Tensor | None = None, + dummy_run: bool = False, + skip_attn_for_dummy_run: bool = False, + mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, + is_profile: bool = False, + output_copy_event: torch.cuda.Event | None = None, + sampled_token_ids_cpu: np.ndarray | None = None, + num_sampled_tokens_cpu: np.ndarray | None = None, + all_token_ids_cpu: np.ndarray | None = None, + ) -> torch.Tensor: + num_reqs = input_batch.num_reqs + num_target_tokens = input_batch.num_tokens + num_query_tokens = num_reqs * self.num_query_per_req + max_seq_len = input_batch.seq_lens_cpu_upper_bound[:num_reqs].max().item() + self.draft_max_seq_len = min( + max_seq_len + self.num_query_per_req, self.max_model_len + ) + + # NOTE: To avoid CPU-GPU synchronization without CPU knowing the + # number of rejected tokens, we maintain the size of input_ids and + # hidden_states the same as the target model's. This means, we pad each + # request's query length to include any rejected positions. + if aux_hidden_states: + with record_function_or_nullcontext("dflash: concatenate target hidden"): + combined_target_hidden = torch.cat(aux_hidden_states, dim=-1) + with record_function_or_nullcontext("dflash: project target hidden"): + hidden_states = self.model.combine_hidden_states(combined_target_hidden) + else: + hidden_states = last_hidden_states + with record_function_or_nullcontext("dflash: stage target hidden"): + self.hidden_states[:num_target_tokens].copy_( + hidden_states[:num_target_tokens] + ) + + if dummy_run and skip_attn_for_dummy_run: + # Memory profiling path: block_tables / kv_cache_config are not initialized. + # Since DFlash needs to build its own attention metadata, we must skip the + # preparation in this path and run a minimal forward pass. + self.model.precompute_and_store_context_kv( + self.hidden_states[:num_target_tokens], + self.context_positions[:num_target_tokens], + ) + self._generate_draft( + num_reqs, + num_query_tokens, + attn_metadata=None, + slot_mappings=None, + num_tokens_across_dp=num_tokens_across_dp, + cudagraph_runtime_mode=CUDAGraphMode.NONE, + ) + return self.draft_tokens[:num_reqs] + + # The query slot mapping is written into the shared BlockTables slot_mappings. + # That buffer's address is what the captured CUDA graph reads from at replay. + assert self.draft_kv_cache_group_id >= 0 + # Support multiple draft KV cache groups by preparing inputs once for each + with record_function_or_nullcontext("dflash: prepare inputs"): + for i, gid in enumerate(self.draft_kv_cache_group_ids): + prepare_dflash_inputs( + self.input_buffers, + self.block_tables.slot_mappings[gid], + self.context_positions, + self._context_slot_mappings[i], + self.sample_indices, + self.sample_pos, + self.sample_idx_mapping, + self.temperature, + self.seeds, + input_batch, + num_sampled, + num_rejected, + last_sampled, + next_prefill_tokens, + temperature, + seeds, + self.block_tables.input_block_tables[gid], + self.block_tables.kernel_block_sizes[gid], + self.block_tables.cp_rank, + self.block_tables.cp_size, + self.block_tables.cp_interleave, + self.parallel_drafting_token_id, + self.num_query_per_req, + self.num_speculative_steps, + self.max_num_reqs, + self.max_num_tokens, + self.max_model_len, + self.sample_from_anchor, + ) + + # Pre-insert context K/V into the cache. Runs eagerly outside the captured graph + # because the context shape varies per step. During dummy runs the block tables + # are placeholders, so we skip the cache write to avoid clobbering real entries. + # Each layer uses the context slots of its own kv-cache group. + if dummy_run: + context_slots: torch.Tensor | list[torch.Tensor | None] | None = None + elif self._layer_group_idx is not None: + context_slots = [ + self._context_slot_mappings[gidx][:num_target_tokens] + for gidx in self._layer_group_idx + ] + else: + context_slots = self._context_slot_mappings[0][:num_target_tokens] + with record_function_or_nullcontext("dflash: materialize context kv"): + self.model.precompute_and_store_context_kv( + self.hidden_states[:num_target_tokens], + self.context_positions[:num_target_tokens], + context_slots, + ) + + if not dummy_run and _is_context_only_prefill(input_batch): + # Intermediate chunked-prefill steps only need to materialize the + # target hidden states into draft KV. Their proposed tokens cannot + # be consumed before the final prompt chunk, which will run the + # normal query graph and replace this sentinel output. + self.draft_tokens[:num_reqs].fill_(-1) + if not self._context_only_prefill_logged: + logger.info( + "%s context-only path is skipping unused draft queries " + "for intermediate prefill chunks.", + self._speculator_name, + ) + self._context_only_prefill_logged = True + return self.draft_tokens[:num_reqs] + + if self._prepare_ngram_assist( + input_batch, + output_copy_event, + sampled_token_ids_cpu, + num_sampled_tokens_cpu, + all_token_ids_cpu, + ): + self._apply_ngram_assist(num_reqs) + return self.draft_tokens[:num_reqs] + + with record_function_or_nullcontext("dflash: query and selector"): + # Every DFlash step has exactly num_query_per_req tokens, so we can + # use FULL CUDA graphs. + batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp( + self.query_cudagraph_manager, + num_reqs, + num_query_tokens, + uniform_token_count=self.num_query_per_req, + dp_size=self.dp_size, + dp_rank=self.dp_rank, + need_eager=is_profile, + ) + + num_reqs_padded = batch_desc.num_reqs or num_reqs + num_tokens_padded = batch_desc.num_tokens + + # Rebuild the draft attention metadata even when replaying the FULL + # graph so that any attention metadata builder state is updated. + draft_attn_metadata = self._build_draft_attn_metadata( + num_reqs=num_reqs, + num_reqs_padded=num_reqs_padded, + num_tokens_padded=num_tokens_padded, + seq_lens_cpu_upper_bound=input_batch.seq_lens_cpu_upper_bound, + step=self.num_query_per_req, + causal=self._group_causal, + ) + draft_slot_mappings_by_layer = build_slot_mappings_by_layer( + self.block_tables.slot_mappings[:, :num_tokens_padded], + self.kv_cache_config, + ) + + if batch_desc.cg_mode == CUDAGraphMode.FULL: + assert self.query_cudagraph_manager is not None + self.query_cudagraph_manager.run_fullgraph(batch_desc) + else: + self._generate_draft( + num_reqs, + num_tokens_padded, + draft_attn_metadata, + draft_slot_mappings_by_layer, + num_tokens_across_dp=num_tokens_across_dp, + cudagraph_runtime_mode=batch_desc.cg_mode, + ) + + self._apply_ngram_assist(num_reqs) + + return self.draft_tokens[:num_reqs] + + +@triton.jit +def _prepare_dflash_inputs_kernel( + # Outputs + out_input_ids_ptr, + out_query_positions_ptr, + out_query_start_loc_ptr, + out_seq_lens_ptr, + out_query_slot_mapping_ptr, + out_context_positions_ptr, + out_context_slot_mapping_ptr, + out_sample_indices_ptr, + out_sample_pos_ptr, + out_sample_idx_mapping_ptr, + out_temperature_ptr, + out_seeds_ptr, + # Inputs from target batch + target_positions_ptr, + target_query_start_loc_ptr, + idx_mapping_ptr, + last_sampled_ptr, + next_prefill_tokens_ptr, + num_sampled_ptr, + num_rejected_ptr, + # Sampling params + temperature_ptr, + seeds_ptr, + # Block table for slot mapping lookup. + block_table_ptr, + block_table_stride, + # Scalars + parallel_drafting_token_id, + block_size, + num_query_per_req, + num_speculative_steps, + max_num_reqs, + max_num_tokens, + max_model_len, + cp_rank, + SAMPLE_FROM_ANCHOR: tl.constexpr, + PAD_SLOT_ID: tl.constexpr, + CP_SIZE: tl.constexpr, + CP_INTERLEAVE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + req_idx = tl.program_id(0) + block_idx = tl.program_id(1) + num_reqs = tl.num_programs(0) + req_state_idx = tl.load(idx_mapping_ptr + req_idx) + + ctx_start = tl.load(target_query_start_loc_ptr + req_idx) + ctx_end = tl.load(target_query_start_loc_ptr + req_idx + 1) + num_ctx = ctx_end - ctx_start + + num_rejected = tl.load(num_rejected_ptr + req_idx) + valid_ctx_end = ctx_end - num_rejected + num_valid_ctx = valid_ctx_end - ctx_start + + num_sampled = tl.load(num_sampled_ptr + req_idx) + if num_sampled > 0: + bonus_token = tl.load(last_sampled_ptr + req_state_idx).to(tl.int32) + else: + # Chunked prefilling: splice in the next prefill token. + bonus_token = tl.load(next_prefill_tokens_ptr + req_state_idx).to(tl.int32) + + last_valid_pos = tl.load(target_positions_ptr + valid_ctx_end - 1) + query_base = req_idx * num_query_per_req + + j = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + is_ctx = j < num_ctx + is_valid_ctx = j < num_valid_ctx + is_query = (j >= num_valid_ctx) & (j < num_valid_ctx + num_query_per_req) + query_off = j - num_valid_ctx + + # --- Context positions / slots --- + ctx_pos_idx = ctx_start + tl.where(is_ctx, j, 0) + ctx_pos = tl.load(target_positions_ptr + ctx_pos_idx, mask=is_valid_ctx, other=0) + ctx_block_num = ctx_pos // (block_size * CP_SIZE) + ctx_block_num = tl.minimum(ctx_block_num, block_table_stride - 1) + ctx_block_id = tl.load( + block_table_ptr + req_idx * block_table_stride + ctx_block_num, + mask=is_valid_ctx, + other=0, + ).to(tl.int64) + # Block 0 is the null block. Old sliding-window context positions can map + # to it after eviction; rejected suffix rows are invalid context as well. + # Neither kind of row may write draft KV into physical block 0. + ctx_resident = is_valid_ctx & (ctx_block_id != 0) + local_ctx_slot = cp_local_slot( + ctx_pos, ctx_block_id, block_size, cp_rank, CP_SIZE, CP_INTERLEAVE, PAD_SLOT_ID + ) + ctx_slot = tl.where( + ctx_resident, + local_ctx_slot, + PAD_SLOT_ID, + ) + # Stored over the full [0, num_ctx) span while the loads above are masked to + # [0, num_valid_ctx): the rejected suffix rows in between get position 0 and + # PAD_SLOT_ID. That is intentional — those rows write no KV and their + # positions are never consumed, but the span must stay fully initialized so + # a replayed graph cannot observe a stale value from an earlier batch. + tl.store(out_context_positions_ptr + ctx_start + j, ctx_pos, mask=is_ctx) + tl.store(out_context_slot_mapping_ptr + ctx_start + j, ctx_slot, mask=is_ctx) + + # --- Query positions / input_ids / slots --- + query_pos = last_valid_pos + 1 + query_off + query_idx = query_base + query_off + is_bonus = is_query & (query_off == 0) + input_id = tl.where(is_bonus, bonus_token, parallel_drafting_token_id) + + q_block_num = query_pos // (block_size * CP_SIZE) + q_block_num = tl.minimum(q_block_num, block_table_stride - 1) + q_block_id = tl.load( + block_table_ptr + req_idx * block_table_stride + q_block_num, + mask=is_query, + other=0, + ).to(tl.int64) + # A null block is never a writable cache slot. This can occur when a + # sliding-window block table contains evicted/global padding entries. + q_resident = is_query & (q_block_id != 0) + local_q_slot = cp_local_slot( + query_pos, + q_block_id, + block_size, + cp_rank, + CP_SIZE, + CP_INTERLEAVE, + PAD_SLOT_ID, + ) + q_slot = tl.where( + q_resident, + local_q_slot, + PAD_SLOT_ID, + ) + + tl.store(out_input_ids_ptr + query_idx, input_id, mask=is_query) + clamped_query_pos = tl.minimum(query_pos, max_model_len - 1) + tl.store(out_query_positions_ptr + query_idx, clamped_query_pos, mask=is_query) + tl.store(out_query_slot_mapping_ptr + query_idx, q_slot, mask=is_query) + + # --- Sample indices / positions / idx_mapping --- + # When SAMPLE_FROM_ANCHOR (DSpark), so we sample at EVERY query position + # and each position k predicts the NEXT token (sampled position = query_pos + 1). + # Otherwise (DFlash default) the anchor is the bonus token and only the mask tokens + # at offsets > 0 are sampled from, each AT its own position. + sample_off = 0 if SAMPLE_FROM_ANCHOR else 1 + is_sample = is_query & (query_off >= sample_off) + sample_idx = req_idx * num_speculative_steps + (query_off - sample_off) + sample_pos = query_pos + 1 if SAMPLE_FROM_ANCHOR else query_pos + tl.store(out_sample_indices_ptr + sample_idx, query_idx, mask=is_sample) + tl.store(out_sample_pos_ptr + sample_idx, sample_pos, mask=is_sample) + tl.store(out_sample_idx_mapping_ptr + sample_idx, req_state_idx, mask=is_sample) + + if block_idx == 0: + tl.store(out_query_start_loc_ptr + req_idx, query_base) + # seq_lens is the absolute sequence length the draft attention + # reads up to (context + query), not just the count of accepted + # tokens this step. + tl.store( + out_seq_lens_ptr + req_idx, + tl.minimum(last_valid_pos + 1 + num_query_per_req, max_model_len), + ) + # Copy sampling state. + tl.store( + out_temperature_ptr + req_state_idx, + tl.load(temperature_ptr + req_state_idx), + ) + tl.store(out_seeds_ptr + req_state_idx, tl.load(seeds_ptr + req_state_idx)) + if req_idx == num_reqs - 1: + # Pad per-request buffers to max_num_reqs for CUDA graph safety. + last_query_end = num_reqs * num_query_per_req + for i in range(num_reqs, max_num_reqs + 1, BLOCK_SIZE): + block = i + tl.arange(0, BLOCK_SIZE) + mask = block < max_num_reqs + 1 + tl.store(out_query_start_loc_ptr + block, last_query_end, mask=mask) + for i in range(num_reqs, max_num_reqs, BLOCK_SIZE): + block = i + tl.arange(0, BLOCK_SIZE) + mask = block < max_num_reqs + tl.store(out_seq_lens_ptr + block, 0, mask=mask) + # Padded sample slots point at query index 0 (a valid row in + # last_hidden_states) so CG replay never reads OOB. Padded + # sample idx mappings point to -1, which is ignored during + # sampling to prevent writing stale values to draft logits. + pad_start = num_reqs * num_speculative_steps + pad_end = max_num_reqs * num_speculative_steps + for i in range(pad_start, pad_end, BLOCK_SIZE): + block = i + tl.arange(0, BLOCK_SIZE) + mask = block < pad_end + tl.store(out_sample_indices_ptr + block, 0, mask=mask) + tl.store(out_sample_pos_ptr + block, 0, mask=mask) + tl.store(out_sample_idx_mapping_ptr + block, -1, mask=mask) + # Pad query slot mappings past num_query_tokens with PAD so the + # captured CG sees PAD slots (no K/V write) for replay sizes + # larger than the current request count. + q_pad_start = num_reqs * num_query_per_req + for i in range(q_pad_start, max_num_tokens, BLOCK_SIZE): + block = i + tl.arange(0, BLOCK_SIZE) + mask = block < max_num_tokens + tl.store(out_query_slot_mapping_ptr + block, PAD_SLOT_ID, mask=mask) + + +def prepare_dflash_inputs( + input_buffers: InputBuffers, + query_slot_mapping: torch.Tensor, + context_positions: torch.Tensor, + context_slot_mapping: torch.Tensor, + sample_indices: torch.Tensor, + sample_pos: torch.Tensor, + sample_idx_mapping: torch.Tensor, + temperature: torch.Tensor, + seeds: torch.Tensor, + input_batch: InputBatch, + # [num_reqs] + num_sampled: torch.Tensor, + # [num_reqs] + num_rejected: torch.Tensor, + # [max_num_reqs] + last_sampled: torch.Tensor, + # [max_num_reqs] + next_prefill_tokens: torch.Tensor, + # [max_num_reqs] + input_temperature: torch.Tensor, + # [max_num_reqs] + input_seeds: torch.Tensor, + # [max_num_reqs, max_num_blocks] + block_table: torch.Tensor, + block_size: int, + cp_rank: int, + cp_size: int, + cp_interleave: int, + parallel_drafting_token_id: int, + num_query_per_req: int, + num_speculative_steps: int, + max_num_reqs: int, + max_num_tokens: int, + max_model_len: int, + sample_from_anchor: bool = False, +) -> None: + num_reqs = input_batch.num_reqs + assert num_reqs > 0 + # Cover the longest possible per-request span (ctx + query). Use the max + # per-request query length, not the total token count across the batch. + max_target_query_len = int(input_batch.num_scheduled_tokens.max()) + max_tokens_per_req = max_target_query_len + num_query_per_req + BLOCK_SIZE = min(256, triton.next_power_of_2(max(1, max_tokens_per_req))) + num_blocks = triton.cdiv(max_tokens_per_req, BLOCK_SIZE) + _prepare_dflash_inputs_kernel[(num_reqs, num_blocks)]( + input_buffers.input_ids, + input_buffers.positions, + input_buffers.query_start_loc, + input_buffers.seq_lens, + query_slot_mapping, + context_positions, + context_slot_mapping, + sample_indices, + sample_pos, + sample_idx_mapping, + temperature, + seeds, + input_batch.positions, + input_batch.query_start_loc, + input_batch.idx_mapping, + last_sampled, + next_prefill_tokens, + num_sampled, + num_rejected, + input_temperature, + input_seeds, + block_table, + block_table.stride(0), + parallel_drafting_token_id, + block_size, + num_query_per_req, + num_speculative_steps, + max_num_reqs, + max_num_tokens, + max_model_len, + cp_rank, + SAMPLE_FROM_ANCHOR=sample_from_anchor, + PAD_SLOT_ID=PAD_SLOT_ID, + CP_SIZE=cp_size, + CP_INTERLEAVE=cp_interleave, + BLOCK_SIZE=BLOCK_SIZE, + ) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/utils.py b/vllm/v1/worker/gpu/spec_decode/dflash/utils.py new file mode 100644 index 0000000000..f714b3db30 --- /dev/null +++ b/vllm/v1/worker/gpu/spec_decode/dflash/utils.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch.nn as nn + +from vllm.config import VllmConfig, replace +from vllm.distributed.parallel_state import get_pp_group +from vllm.logger import init_logger +from vllm.model_executor.model_loader import get_model +from vllm.platforms import current_platform +from vllm.v1.worker.gpu.spec_decode.eagle.utils import ( + _should_share, + get_target_lm_head, +) + +logger = init_logger(__name__) + + +def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module: + from vllm.compilation.backends import set_model_tag + from vllm.model_executor.models.qwen3_dflash import ( + dflash_has_any_non_causal, + dflash_target_rope_is_neox_style, + ) + + speculative_config = vllm_config.speculative_config + assert speculative_config is not None + draft_model_config = speculative_config.draft_model_config + draft_cache_dtype = speculative_config.kv_cache_dtype + if ( + draft_cache_dtype is None + and str(vllm_config.cache_config.cache_dtype).startswith("fp8") + and current_platform.is_cuda() + and current_platform.is_device_capability(70) + ): + draft_cache_dtype = "auto" + logger.info_once( + "Using FP16 draft KV cache for SM70 DFlash while the target uses %s.", + vllm_config.cache_config.cache_dtype, + ) + # The drafter must rotate Q/K the way its target does. Take that from the + # built target before super() constructs the draft. + is_neox_style = dflash_target_rope_is_neox_style(target_model) + if is_neox_style is not None: + draft_model_config.hf_config.is_neox_style = is_neox_style + # Select an attention backend that supports the drafter's attention: mixing + # a non-causal layer onto a causal-only backend would fail. + draft_vllm_config = replace( + vllm_config, + attention_config=replace( + vllm_config.attention_config, + use_non_causal=dflash_has_any_non_causal(draft_model_config.hf_config), + backend=speculative_config.attention_backend, + ), + cache_config=( + replace( + vllm_config.cache_config, + cache_dtype=draft_cache_dtype, + ) + if draft_cache_dtype is not None + else vllm_config.cache_config + ), + ) + with set_model_tag("dflash_head"): + dflash_model = get_model( + vllm_config=draft_vllm_config, model_config=draft_model_config + ) + + target_language_model = ( + target_model.get_language_model() + if hasattr(target_model, "get_language_model") + else target_model + ) + # MuseGlimmerForCausalLM marks its inner MuseGlimmerModel as the language + # model, so get_language_model() already returns the inner module and has + # no .model of its own. + target_inner = getattr(target_language_model, "model", target_language_model) + draft_inner = dflash_model.model + + # Skip embedding sharing under PP — each rank owns its own embedding. + if get_pp_group().world_size == 1: + target_embed = getattr(target_inner, "embed_tokens", None) or getattr( + target_inner, "embedding", None + ) + draft_embed = getattr(draft_inner, "embed_tokens", None) + if target_embed is not None and _should_share( + dflash_model, "has_own_embed_tokens", draft_embed, target_embed + ): + if draft_embed is not None: + del draft_inner.embed_tokens + draft_inner.embed_tokens = target_embed + + target_lm_head = get_target_lm_head(target_model, target_language_model) + draft_lm_head = getattr(dflash_model, "lm_head", None) + if target_lm_head is not None and _should_share( + dflash_model, "has_own_lm_head", draft_lm_head, target_lm_head + ): + if draft_lm_head is not None: + del dflash_model.lm_head + dflash_model.lm_head = target_lm_head + + return dflash_model diff --git a/vllm/v1/worker/gpu/spec_decode/dflash2/__init__.py b/vllm/v1/worker/gpu/spec_decode/dflash2/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/vllm/v1/worker/gpu/spec_decode/dflash2/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/v1/worker/gpu/spec_decode/dflash2/ngram_assist.py b/vllm/v1/worker/gpu/spec_decode/dflash2/ngram_assist.py new file mode 100644 index 0000000000..fe89c55512 --- /dev/null +++ b/vllm/v1/worker/gpu/spec_decode/dflash2/ngram_assist.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Host prompt-ngram lookup for the MRV2 DFlash2 assistant.""" + +from __future__ import annotations + +import time + +import numpy as np +from numba import njit + + +@njit(inline="always") +def _split_token( + prefix: np.ndarray, + prefix_len: int, + suffix: np.ndarray, + index: int, +) -> int: + if index < prefix_len: + return int(prefix[index]) + return int(suffix[index - prefix_len]) + + +@njit(cache=True) +def find_split_ngram_proposal( + prefix: np.ndarray, + prefix_len: int, + suffix: np.ndarray, + suffix_len: int, + min_ngram: int, + max_ngram: int, + max_model_len: int, + num_draft_tokens: int, +) -> np.ndarray: + """Run the prompt-lookup KMP over ``prefix + suffix`` without copying it.""" + total_tokens = prefix_len + suffix_len + if total_tokens < min_ngram: + return np.empty(0, dtype=np.int32) + + num_draft_tokens = min(num_draft_tokens, max_model_len - total_tokens) + if num_draft_tokens <= 0: + return np.empty(0, dtype=np.int32) + + # This is the same reverse-KMP policy as the standalone vLLM ngram + # proposer: prefer the longest suffix match and the earliest occurrence. + lps = np.zeros(max_ngram, dtype=np.int32) + longest_ngram = 0 + position = 0 + prev_lps = 0 + i = 1 + while i < total_tokens: + prefix_token = _split_token( + prefix, + prefix_len, + suffix, + total_tokens - 1 - prev_lps, + ) + current_token = _split_token( + prefix, + prefix_len, + suffix, + total_tokens - 1 - i, + ) + if prefix_token == current_token: + prev_lps += 1 + if prev_lps >= longest_ngram: + longest_ngram = prev_lps + position = i + if i < max_ngram: + lps[i] = prev_lps + if prev_lps == max_ngram: + prev_lps = lps[max_ngram - 1] + i += 1 + elif prev_lps != 0: + prev_lps = lps[prev_lps - 1] + else: + i += 1 + + if longest_ngram < min_ngram: + return np.empty(0, dtype=np.int32) + + start = total_tokens - 1 - position + longest_ngram + proposal_len = min(num_draft_tokens, total_tokens - start) + proposal = np.empty(proposal_len, dtype=np.int32) + for j in range(proposal_len): + proposal[j] = _split_token(prefix, prefix_len, suffix, start + j) + return proposal + + +class DFlash2NgramAssist: + """Batch adapter and counters around the split prompt lookup.""" + + def __init__( + self, + min_ngram: int, + max_ngram: int, + num_draft_tokens: int, + max_model_len: int, + ) -> None: + self.min_ngram = min_ngram + self.max_ngram = max_ngram + self.num_draft_tokens = num_draft_tokens + self.max_model_len = max_model_len + self.num_eligible = 0 + self.num_full_hits = 0 + self.lookup_seconds = 0.0 + + # Compile the small split-input kernel at startup, not on the first + # user-visible decode token. + find_split_ngram_proposal( + np.zeros(1, dtype=np.int32), + 1, + np.zeros(1, dtype=np.int64), + 0, + min_ngram, + max_ngram, + max_model_len, + num_draft_tokens, + ) + + def propose( + self, + token_ids_cpu: np.ndarray, + req_state_indices: np.ndarray, + prior_lengths: np.ndarray, + sampled_token_ids: np.ndarray, + num_sampled_tokens: np.ndarray, + eligible: np.ndarray, + output_tokens: np.ndarray, + output_lengths: np.ndarray, + ) -> int: + """Fill fixed-width outputs and return the number of full-width hits.""" + num_reqs = len(req_state_indices) + output_tokens[:num_reqs].fill(0) + output_lengths[:num_reqs].fill(0) + full_hits = 0 + started = time.perf_counter() + for batch_idx in range(num_reqs): + if not eligible[batch_idx]: + continue + self.num_eligible += 1 + req_state_idx = int(req_state_indices[batch_idx]) + proposal = find_split_ngram_proposal( + token_ids_cpu[req_state_idx], + int(prior_lengths[batch_idx]), + sampled_token_ids[batch_idx], + int(num_sampled_tokens[batch_idx]), + self.min_ngram, + self.max_ngram, + self.max_model_len, + self.num_draft_tokens, + ) + proposal_len = len(proposal) + output_lengths[batch_idx] = proposal_len + if proposal_len: + output_tokens[batch_idx, :proposal_len] = proposal + if proposal_len == self.num_draft_tokens: + full_hits += 1 + self.num_full_hits += full_hits + self.lookup_seconds += time.perf_counter() - started + return full_hits diff --git a/vllm/v1/worker/gpu/spec_decode/dflash2/sparse_rejection.py b/vllm/v1/worker/gpu/spec_decode/dflash2/sparse_rejection.py new file mode 100644 index 0000000000..01b42e3968 --- /dev/null +++ b/vllm/v1/worker/gpu/spec_decode/dflash2/sparse_rejection.py @@ -0,0 +1,261 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Strictly gated compact target rejection for DFlash2 on SM70.""" + +from __future__ import annotations + +import os +from typing import TYPE_CHECKING, Any + +import numpy as np +import torch + +import vllm.envs as envs +from vllm.logger import init_logger +from vllm.v1.worker.gpu.sample.output import SamplerOutput +from vllm.v1.worker.gpu.sample.states import NO_LOGPROBS +from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import DFlash2Speculator +from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import ( + dflash2_sparse_topk_rejection_sample, +) + +if TYPE_CHECKING: + from vllm.v1.core.sched.output import GrammarOutput + from vllm.v1.worker.gpu.input_batch import InputBatch + from vllm.v1.worker.gpu.spec_decode.rejection_sampler import RejectionSampler + +logger = init_logger(__name__) + +_TARGET_TOP_K = 20 +_SELECTOR_ALIGNMENT_DUMP_COUNT = 0 +_SELECTOR_ALIGNMENT_STEP = 0 + + +def _parse_alignment_steps(raw_steps: str | None) -> set[int] | None: + if not raw_steps: + return None + steps: set[int] = set() + try: + for item in raw_steps.split(","): + item = item.strip() + if not item: + continue + if "-" in item: + start_text, end_text = item.split("-", 1) + start = int(start_text) + end = int(end_text) + if start < 0 or end < start: + return set() + steps.update(range(start, end + 1)) + else: + step = int(item) + if step < 0: + return set() + steps.add(step) + except ValueError: + return set() + return steps + + +def _safe_dump_tag(raw_tag: str) -> str: + return "".join(ch if ch.isalnum() or ch in "._-" else "_" for ch in raw_tag) + + +def _diagnostic_rank() -> int: + return int(os.getenv("RANK", os.getenv("LOCAL_RANK", "0"))) + + +def _maybe_dump_selector_alignment( + *, + speculator: DFlash2Speculator, + rejection_sampler: RejectionSampler, + input_batch: InputBatch, + target_topk_ids: torch.Tensor, + target_topk_logits: torch.Tensor, + draft_topk_ids: torch.Tensor, + draft_topk_logits: torch.Tensor, + draft_sampled: torch.Tensor, + pos: torch.Tensor, + sampled: torch.Tensor, + num_sampled: torch.Tensor, +) -> None: + """Dump one exact B1 selector/target alignment record when requested.""" + if not envs.VLLM_SPEC_DUMP_ALIGNMENT: + return + if _diagnostic_rank() != 0: + return + + global _SELECTOR_ALIGNMENT_DUMP_COUNT, _SELECTOR_ALIGNMENT_STEP + _SELECTOR_ALIGNMENT_STEP += 1 + if _SELECTOR_ALIGNMENT_DUMP_COUNT >= envs.VLLM_SPEC_DUMP_ALIGNMENT_LIMIT: + return + selected_steps = _parse_alignment_steps(envs.VLLM_SPEC_DUMP_ALIGNMENT_STEPS) + if selected_steps is not None and _SELECTOR_ALIGNMENT_STEP not in selected_steps: + return + + shadow = speculator.get_selector_alignment_shadow() + if shadow is None: + return + shadow_ids, unary_logits, lattice_scores = shadow + req_state = int(input_batch.idx_mapping_np[0]) + packed_row = 0 + sampling_states = rejection_sampler.sampler.sampling_states + + with torch.no_grad(): + draft_sampled_cpu = draft_sampled.detach().cpu() + if bool(torch.all(draft_sampled_cpu == 0).item()): + return + payload = { + "format": "dflash2_selector_alignment_v1", + "rank": _diagnostic_rank(), + "step": _SELECTOR_ALIGNMENT_STEP, + "request_state": req_state, + "selector_top_k": speculator.selector_top_k, + "num_speculative_steps": speculator.num_speculative_steps, + "target_topk_ids": target_topk_ids.detach().cpu(), + "target_topk_logits": target_topk_logits.detach().float().cpu(), + "draft_candidate_ids": draft_topk_ids[req_state].detach().cpu(), + "draft_realized_logits": ( + draft_topk_logits[req_state].detach().float().cpu() + ), + "selector_candidate_ids": shadow_ids[packed_row].detach().cpu(), + "selector_unary_logits": (unary_logits[packed_row].detach().float().cpu()), + "selector_lattice_scores": ( + lattice_scores[packed_row].detach().float().cpu() + ), + "draft_sampled": draft_sampled_cpu, + "positions": pos.detach().cpu(), + "cu_num_logits": input_batch.cu_num_logits.detach().cpu(), + "idx_mapping": input_batch.idx_mapping.detach().cpu(), + "temperature": float(sampling_states.temperature.np[req_state]), + "top_p": float(sampling_states.top_p.np[req_state]), + "top_k": int(sampling_states.top_k.np[req_state]), + "sampled_token_ids": sampled.detach().cpu(), + "num_sampled": num_sampled.detach().cpu(), + } + _SELECTOR_ALIGNMENT_DUMP_COUNT += 1 + dump_dir = os.getenv("VLLM_SPEC_DUMP_ALIGNMENT_DIR", "/tmp") + os.makedirs(dump_dir, exist_ok=True) + tag = _safe_dump_tag(os.getenv("VLLM_SPEC_DUMP_ALIGNMENT_TAG", "")) + tag_part = f"{tag}_" if tag else "" + dump_path = os.path.join( + dump_dir, + f"spec_alignment_dflash2_selector_{tag_part}pid{os.getpid()}_" + f"step{_SELECTOR_ALIGNMENT_STEP:06d}_" + f"{_SELECTOR_ALIGNMENT_DUMP_COUNT}.pt", + ) + torch.save(payload, dump_path) + logger.warning("Dumped DFlash2 selector alignment diagnostics to %s", dump_path) + + +def _supports_sparse_sampling_contract( + rejection_sampler: RejectionSampler, + input_batch: InputBatch, +) -> bool: + """Whether compact logits preserve every requested sampling transform.""" + if rejection_sampler.rejection_sample_method != "standard": + return False + # Start with the single-request path used by the latency target. The + # kernel supports batches, but mixed-request graph validation is a + # separate promotion gate. + if input_batch.num_reqs != 1 or np.any(input_batch.is_prefilling_np): + return False + + sampler = rejection_sampler.sampler + idx = input_batch.idx_mapping_np + states = sampler.sampling_states + temperatures = states.temperature.np[idx] + top_k = states.top_k.np[idx] + top_p = states.top_p.np[idx] + if np.any(temperatures <= 0.0): + return False + if np.any(top_k != _TARGET_TOP_K): + return False + if np.any((top_p <= 0.0) | (top_p > 1.0)): + return False + if np.any(states.min_p.np[idx] != 0.0): + return False + + if np.any(sampler.penalties_state.use_penalty[idx]): + return False + if np.any(sampler.logit_bias_state.use_logit_bias[idx]): + return False + if np.any(sampler.bad_words_state.num_bad_words.np[idx] != 0): + return False + if states.max_num_logprobs(idx) != NO_LOGPROBS: + return False + if sampler.logprob_token_ids_state.max_num_token_ids(idx) != 0: + return False + return not sampler.compute_nans + + +def try_dflash2_sparse_target_rejection( + model: Any, + speculator: Any, + rejection_sampler: RejectionSampler, + sample_hidden_states: torch.Tensor, + input_batch: InputBatch, + grammar_output: GrammarOutput | None, +) -> SamplerOutput | None: + """Sample from compact target/draft supports, or return ``None`` safely.""" + if not envs.VLLM_SM70_DFLASH2_SPARSE_TARGET_REJECTION: + return None + if not isinstance(speculator, DFlash2Speculator): + return None + if grammar_output is not None or input_batch.has_structured_output_reqs: + return None + if sample_hidden_states.device.type != "cuda": + return None + if torch.cuda.get_device_capability(sample_hidden_states.device) != (7, 0): + return None + if not hasattr(model, "get_topk_tokens_and_logits"): + return None + if not _supports_sparse_sampling_contract(rejection_sampler, input_batch): + return None + + sparse_draft_logits = speculator.get_sparse_draft_logits() + if sparse_draft_logits is None: + return None + draft_topk_ids, draft_topk_logits = sparse_draft_logits + target_topk_ids, target_topk_logits = model.get_topk_tokens_and_logits( + sample_hidden_states, + _TARGET_TOP_K, + ) + draft_sampled = input_batch.input_ids[input_batch.logits_indices] + pos = input_batch.positions[input_batch.logits_indices] + sampled, num_sampled = dflash2_sparse_topk_rejection_sample( + target_topk_ids, + target_topk_logits, + draft_topk_ids, + draft_topk_logits, + draft_sampled, + input_batch.cu_num_logits, + pos, + input_batch.idx_mapping, + rejection_sampler.sampler.sampling_states.temperature.gpu, + rejection_sampler.sampler.sampling_states.top_p.gpu, + rejection_sampler.sampler.sampling_states.seeds.gpu, + rejection_sampler.num_speculative_steps, + use_fp64=rejection_sampler.sampler.use_fp64_gumbel, + ) + if envs.VLLM_SPEC_DUMP_ALIGNMENT: + _maybe_dump_selector_alignment( + speculator=speculator, + rejection_sampler=rejection_sampler, + input_batch=input_batch, + target_topk_ids=target_topk_ids, + target_topk_logits=target_topk_logits, + draft_topk_ids=draft_topk_ids, + draft_topk_logits=draft_topk_logits, + draft_sampled=draft_sampled, + pos=pos, + sampled=sampled, + num_sampled=num_sampled, + ) + logger.info_once("Using SM70 DFlash2 compact target top-k rejection sampling.") + return SamplerOutput( + sampled_token_ids=sampled, + logprobs_tensors=None, + num_nans=None, + num_sampled=num_sampled, + ) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py new file mode 100644 index 0000000000..64f987238a --- /dev/null +++ b/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py @@ -0,0 +1,621 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Any + +import numpy as np +import torch + +from vllm import envs +from vllm.config import VllmConfig +from vllm.config.compilation import CUDAGraphMode +from vllm.logger import init_logger +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.ngram_assist import DFlash2NgramAssist + +logger = init_logger(__name__) + + +def _requires_sm70_tail(device: torch.device, num_steps: int) -> bool: + """Whether the final dependent selector slot needs its own kernel.""" + return ( + device.type == "cuda" + and num_steps > 1 + and torch.cuda.get_device_capability(device) == (7, 0) + ) + + +@triton.jit +def _selector_walk_kernel( + scores_ptr, + candidate_ptr, + sample_pos_ptr, + req_state_ptr, + temperature_ptr, + seeds_ptr, + tokens_ptr, + realized_scores_ptr, + path_state_ptr, + num_steps: tl.constexpr, + walk_steps: tl.constexpr, + top_k: tl.constexpr, + BLOCK_K: tl.constexpr, + SAMPLE_PROBABILISTIC: tl.constexpr, + USE_FP64: tl.constexpr, +): + row = tl.program_id(0) + offsets = tl.arange(0, BLOCK_K) + mask = offsets < top_k + req_state = tl.load(req_state_ptr + row * num_steps) + valid = req_state >= 0 + temperature = tl.load(temperature_ptr + req_state, mask=valid, other=0.0) + seed = tl.load(seeds_ptr + req_state, mask=valid, other=0) + previous = 0 + for step in range(walk_steps): + flat = row * num_steps + step + score_base = (flat * top_k + previous) * top_k + scores = tl.load( + scores_ptr + score_base + offsets, + mask=mask & valid, + other=float("-inf"), + ).to(tl.float64 if USE_FP64 else tl.float32) + if SAMPLE_PROBABILISTIC and temperature != 0.0: + # Cache the exact temperature-applied proposal scores expected by + # the shared rejection sampler. This keeps Eagle/MTP's established + # contract unchanged while matching the DFlash2 selector draw. + scores = scores / temperature + candidate_base = flat * top_k + candidates = tl.load( + candidate_ptr + candidate_base + offsets, + mask=mask & valid, + other=0, + ) + + # Candidate token IDs key the noise, matching the target sampler. + position = tl.load(sample_pos_ptr + flat) - 1 + _, index = gumbel_noised_argmax( + scores, + candidates, + mask & valid, + seed, + position, + temperature if SAMPLE_PROBABILISTIC else 0.0, + USE_FP64=USE_FP64, + APPLY_TEMPERATURE=False, + ) + + tl.store( + realized_scores_ptr + candidate_base + offsets, + scores, + mask=mask & valid, + ) + token = tl.load(candidate_ptr + candidate_base + index, mask=valid, other=0) + tl.store(tokens_ptr + flat, token, mask=valid) + previous = index + + if walk_steps < num_steps: + tl.store(path_state_ptr + row, previous, mask=valid) + + +@triton.jit +def _selector_walk_tail_kernel( + scores_ptr, + candidate_ptr, + sample_pos_ptr, + req_state_ptr, + temperature_ptr, + seeds_ptr, + tokens_ptr, + realized_scores_ptr, + path_state_ptr, + num_steps: tl.constexpr, + top_k: tl.constexpr, + BLOCK_K: tl.constexpr, + SAMPLE_PROBABILISTIC: tl.constexpr, + USE_FP64: tl.constexpr, +): + """Write the final dependent slot separately on SM70. + + Triton can drop the seventh store of a fully unrolled selector walk during + CUDA Graph replay on Volta. The first six slots remain fused; this tiny tail + consumes their persistent path state and guarantees the seventh write. + """ + row = tl.program_id(0) + offsets = tl.arange(0, BLOCK_K) + mask = offsets < top_k + req_state = tl.load(req_state_ptr + row * num_steps) + valid = req_state >= 0 + temperature = tl.load(temperature_ptr + req_state, mask=valid, other=0.0) + seed = tl.load(seeds_ptr + req_state, mask=valid, other=0) + previous = tl.load(path_state_ptr + row, mask=valid, other=0) + step: tl.constexpr = num_steps - 1 + flat = row * num_steps + step + score_base = (flat * top_k + previous) * top_k + scores = tl.load( + scores_ptr + score_base + offsets, + mask=mask & valid, + other=float("-inf"), + ).to(tl.float64 if USE_FP64 else tl.float32) + if SAMPLE_PROBABILISTIC and temperature != 0.0: + scores = scores / temperature + candidate_base = flat * top_k + candidates = tl.load( + candidate_ptr + candidate_base + offsets, + mask=mask & valid, + other=0, + ) + position = tl.load(sample_pos_ptr + flat) - 1 + _, index = gumbel_noised_argmax( + scores, + candidates, + mask & valid, + seed, + position, + temperature if SAMPLE_PROBABILISTIC else 0.0, + USE_FP64=USE_FP64, + APPLY_TEMPERATURE=False, + ) + tl.store( + realized_scores_ptr + candidate_base + offsets, + scores, + mask=mask & valid, + ) + token = tl.load(candidate_ptr + candidate_base + index, mask=valid, other=0) + tl.store(tokens_ptr + flat, token, mask=valid) + + +@triton.jit +def _cache_draft_logits_kernel( + draft_logits_ptr, + cached_candidate_ptr, + cached_score_ptr, + candidate_ptr, + scores_ptr, + req_state_ptr, + 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_state = tl.load(req_state_ptr + flat) + step = flat % num_steps + 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 + old_token_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_token_ids, -float("inf"), mask=mask) + token_ids = tl.load(candidate_ptr + candidate_base + offsets, mask=mask) + scores = tl.load(scores_ptr + candidate_base + offsets, mask=mask) + tl.store(logits_base + token_ids, scores, mask=mask) + tl.store(cached_candidate_ptr + cache_base + offsets, token_ids, mask=mask) + if CACHE_SCORES: + tl.store(cached_score_ptr + cache_base + offsets, scores, mask=mask) + + +@triton.jit +def _apply_ngram_draft_kernel( + ngram_tokens_ptr, + ngram_lengths_ptr, + sample_req_state_ptr, + draft_tokens_ptr, + draft_tokens_stride, + cached_candidate_ptr, + cached_score_ptr, + cache_stride_0, + cache_stride_1, + draft_logits_ptr, + draft_logits_stride_0, + draft_logits_stride_1, + num_steps: tl.constexpr, + top_k: tl.constexpr, + BLOCK_K: tl.constexpr, + CACHE_DRAFT_LOGITS: tl.constexpr, + CACHE_SCORES: tl.constexpr, +): + flat = tl.program_id(0) + batch_idx = flat // num_steps + step = flat % num_steps + req_state = tl.load(sample_req_state_ptr + flat) + valid = (req_state >= 0) & (tl.load(ngram_lengths_ptr + batch_idx) == num_steps) + token = tl.load(ngram_tokens_ptr + flat, mask=valid, other=0).to(tl.int64) + tl.store( + draft_tokens_ptr + batch_idx * draft_tokens_stride + step, + token, + mask=valid, + ) + + if CACHE_DRAFT_LOGITS: + offsets = tl.arange(0, BLOCK_K) + topk_mask = valid & (offsets < top_k) + cache_base = ( + cached_candidate_ptr + req_state * cache_stride_0 + step * cache_stride_1 + ) + old_ids = tl.load(cache_base + offsets, mask=topk_mask, other=0) + 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=topk_mask) + + is_proposal = offsets == 0 + new_ids = tl.where(is_proposal, token, 0) + new_scores = tl.where(is_proposal, 0.0, -float("inf")) + tl.store(cache_base + offsets, new_ids, mask=topk_mask) + if CACHE_SCORES: + score_base = ( + cached_score_ptr + req_state * cache_stride_0 + step * cache_stride_1 + ) + tl.store(score_base + offsets, new_scores, mask=topk_mask) + tl.store(logits_base + token, 0.0, mask=valid) + + +class DFlash2Speculator(DFlashSpeculator): + _speculator_name = "DFlash2" + + def __init__(self, vllm_config: VllmConfig, device: torch.device): + super().__init__(vllm_config, device) + draft_config = self.draft_model_config.hf_config.dflash_config + self.selector_top_k = int(draft_config["selector_top_k"]) + self._anchor_indices = ( + torch.arange(self.max_num_reqs, dtype=torch.int64, device=device) + * self.num_query_per_req + ) + self._selector_scores = torch.empty( + self.max_num_reqs, + self.num_speculative_steps, + 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._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._selector_path_state = torch.empty( + self.max_num_reqs, dtype=torch.int32, device=device + ) + self._alignment_candidate_ids: torch.Tensor | None = None + self._alignment_unary_logits: torch.Tensor | None = None + self._alignment_lattice_scores: torch.Tensor | None = None + if ( + envs.VLLM_SPEC_DUMP_ALIGNMENT + and envs.VLLM_SM70_DFLASH2_SPARSE_TARGET_REJECTION + ): + packed_shape = ( + self.max_num_reqs, + self.num_speculative_steps, + self.selector_top_k, + ) + self._alignment_candidate_ids = torch.empty( + packed_shape, dtype=torch.int64, device=device + ) + self._alignment_unary_logits = torch.empty( + packed_shape, dtype=torch.float32, device=device + ) + self._alignment_lattice_scores = torch.empty( + (*packed_shape, self.selector_top_k), + dtype=torch.float32, + device=device, + ) + self._use_sm70_tail = _requires_sm70_tail(device, self.num_speculative_steps) + if self.draft_logits is not None: + # The cache kernel writes only K columns; all other vocabulary + # columns must remain impossible. + self.draft_logits.fill_(-float("inf")) + + self._ngram_assist: DFlash2NgramAssist | None = None + self._ngram_num_hits = 0 + 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 + ): + min_ngram = speculative_config.prompt_lookup_min + max_ngram = speculative_config.prompt_lookup_max + assert min_ngram is not None and max_ngram is not None + self._ngram_assist = DFlash2NgramAssist( + min_ngram=min_ngram, + max_ngram=max_ngram, + num_draft_tokens=self.num_speculative_steps, + max_model_len=self.max_model_len, + ) + self._ngram_tokens_cpu_tensor = torch.zeros( + self.max_num_reqs, + self.num_speculative_steps, + dtype=torch.int64, + device="cpu", + pin_memory=True, + ) + self._ngram_lengths_cpu_tensor = torch.zeros( + self.max_num_reqs, + dtype=torch.int32, + device="cpu", + pin_memory=True, + ) + self._ngram_tokens_cpu = self._ngram_tokens_cpu_tensor.numpy() + self._ngram_lengths_cpu = self._ngram_lengths_cpu_tensor.numpy() + self._ngram_tokens = torch.zeros( + self.max_num_reqs, + self.num_speculative_steps, + dtype=torch.int64, + device=device, + ) + self._ngram_lengths = torch.zeros( + self.max_num_reqs, dtype=torch.int32, device=device + ) + logger.info( + "Enabled DFlash2 ngram assist with prompt lookup [%d, %d] " + "and draft width %d.", + min_ngram, + max_ngram, + self.num_speculative_steps, + ) + + @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 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. + return torch.float32, -float("inf") + + def _sample_path( + self, + candidate_ids: torch.Tensor, + scores: torch.Tensor, + num_reqs: int, + ) -> None: + # The SM70 tail must consume the exact same packed lattice as the + # prefix walk. Keep one persistent view instead of creating temporary + # contiguous inputs only for the first launch. + 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 + ) + _selector_walk_kernel[(num_reqs,)]( + scores, + candidate_ids, + self.sample_pos, + self.sample_idx_mapping, + self.temperature, + self.seeds, + self.draft_tokens, + self._selector_scores, + self._selector_path_state, + num_steps=self.num_speculative_steps, + walk_steps=walk_steps, + top_k=self.selector_top_k, + BLOCK_K=block_k, + SAMPLE_PROBABILISTIC=self.draft_logits is not None, + USE_FP64=self.use_fp64_gumbel, + num_warps=1, + ) + if self._use_sm70_tail: + _selector_walk_tail_kernel[(num_reqs,)]( + scores, + candidate_ids, + self.sample_pos, + self.sample_idx_mapping, + self.temperature, + self.seeds, + self.draft_tokens, + self._selector_scores, + self._selector_path_state, + num_steps=self.num_speculative_steps, + top_k=self.selector_top_k, + BLOCK_K=block_k, + SAMPLE_PROBABILISTIC=self.draft_logits is not None, + USE_FP64=self.use_fp64_gumbel, + num_warps=1, + ) + + def _cache_draft_logits(self, candidate_ids: torch.Tensor, num_sample: int) -> None: + draft_logits = self.draft_logits + assert draft_logits is not None + cached_scores = self._cached_candidate_scores + block_k = triton.next_power_of_2(self.selector_top_k) + _cache_draft_logits_kernel[(num_sample,)]( + draft_logits, + self._cached_candidate_ids, + self._selector_scores if cached_scores is None else cached_scores, + candidate_ids, + self._selector_scores, + self.sample_idx_mapping, + 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 get_sparse_draft_logits( + self, + ) -> tuple[torch.Tensor, torch.Tensor] | None: + """Return request-slot proposal candidates for sparse rejection.""" + if self.draft_logits is None or self._cached_candidate_scores is None: + return None + return self._cached_candidate_ids, self._cached_candidate_scores + + def get_selector_alignment_shadow( + self, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None: + """Return packed selector tensors for an explicitly enabled diagnostic.""" + if self._alignment_candidate_ids is None: + return None + assert self._alignment_unary_logits is not None + assert self._alignment_lattice_scores is not None + return ( + self._alignment_candidate_ids, + self._alignment_unary_logits, + self._alignment_lattice_scores, + ) + + def _prepare_ngram_assist( + self, + input_batch, + output_copy_event: torch.cuda.Event | None, + sampled_token_ids_cpu: np.ndarray | None, + num_sampled_tokens_cpu: np.ndarray | None, + all_token_ids_cpu: np.ndarray | None, + ) -> bool: + assist = self._ngram_assist + self._ngram_num_hits = 0 + if ( + assist is None + or input_batch.has_structured_output_reqs + or output_copy_event is None + or sampled_token_ids_cpu is None + or num_sampled_tokens_cpu is None + or all_token_ids_cpu is None + ): + return False + + # The copy stream only depends on target sampling. The main stream can + # materialize DFlash context K/V while the host waits here, so lookup + # does not serialize the context projection. + output_copy_event.synchronize() + num_reqs = input_batch.num_reqs + num_draft_tokens = input_batch.num_draft_tokens_per_req + if num_draft_tokens is None: + num_draft_tokens = np.zeros(num_reqs, dtype=np.int32) + prior_lengths = ( + input_batch.seq_lens_cpu_upper_bound[:num_reqs].numpy() - num_draft_tokens + ) + eligible = (~input_batch.is_prefilling_np[:num_reqs]) & ( + num_sampled_tokens_cpu[:num_reqs] > 0 + ) + full_hits = assist.propose( + all_token_ids_cpu, + input_batch.idx_mapping_np, + prior_lengths, + sampled_token_ids_cpu, + num_sampled_tokens_cpu, + eligible, + self._ngram_tokens_cpu, + self._ngram_lengths_cpu, + ) + self._ngram_rounds += 1 + skip_query = num_reqs > 0 and full_hits == num_reqs + self._ngram_num_hits = full_hits if skip_query else 0 + if skip_query: + self._ngram_tokens[:num_reqs].copy_( + self._ngram_tokens_cpu_tensor[:num_reqs], non_blocking=True + ) + self._ngram_lengths[:num_reqs].copy_( + self._ngram_lengths_cpu_tensor[:num_reqs], non_blocking=True + ) + + self._ngram_skipped_rounds += int(skip_query) + if ( + envs.VLLM_DFLASH_PROFILE + and self._ngram_rounds % envs.VLLM_DFLASH_PROFILE_LOG_INTERVAL == 0 + ): + eligible_count = max(assist.num_eligible, 1) + logger.info( + "DFLASH2_NGRAM_PROFILE rounds=%d eligible=%d full_hits=%d " + "hit_rate=%.4f skipped_query_rounds=%d lookup_avg_ms=%.4f", + self._ngram_rounds, + assist.num_eligible, + assist.num_full_hits, + assist.num_full_hits / eligible_count, + self._ngram_skipped_rounds, + assist.lookup_seconds * 1000.0 / self._ngram_rounds, + ) + return skip_query + + def _apply_ngram_assist(self, num_reqs: int) -> None: + if self._ngram_assist is None or self._ngram_num_hits == 0: + return + draft_logits = self.draft_logits + cached_scores = self._cached_candidate_scores + block_k = triton.next_power_of_2(self.selector_top_k) + _apply_ngram_draft_kernel[(num_reqs * self.num_speculative_steps,)]( + self._ngram_tokens, + self._ngram_lengths, + self.sample_idx_mapping, + self.draft_tokens, + self.draft_tokens.stride(0), + self._cached_candidate_ids, + self._selector_scores if cached_scores is None else cached_scores, + self._cached_candidate_ids.stride(0), + self._cached_candidate_ids.stride(1), + self._selector_scores if draft_logits is None else draft_logits, + 0 if draft_logits is None else draft_logits.stride(0), + 0 if draft_logits is None else draft_logits.stride(1), + num_steps=self.num_speculative_steps, + top_k=self.selector_top_k, + BLOCK_K=block_k, + CACHE_DRAFT_LOGITS=draft_logits is not None, + CACHE_SCORES=cached_scores is not None, + num_warps=1, + ) + + def _generate_draft( + self, + num_reqs: int, + num_tokens_padded: int, + attn_metadata: dict[str, Any] | None, + slot_mappings: dict[str, torch.Tensor] | None, + num_tokens_across_dp: torch.Tensor | None, + cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE, + ) -> None: + last_hidden_states = self._run_model( + num_tokens_padded, + attn_metadata, + slot_mappings, + num_tokens_across_dp, + cudagraph_runtime_mode, + ) + num_sample = num_reqs * self.num_speculative_steps + hidden_states = last_hidden_states[self.sample_indices[:num_sample]].view( + num_reqs, self.num_speculative_steps, -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 + ) + unary_logits = unary_logits.view_as(candidate_ids) + anchor_token_ids = self.input_buffers.input_ids[self._anchor_indices[:num_reqs]] + scores = self.model.model.candidate_selector( + candidate_ids, + unary_logits, + hidden_states, + anchor_token_ids, + ) + if self._alignment_candidate_ids is not None: + assert self._alignment_unary_logits is not None + assert self._alignment_lattice_scores is not None + self._alignment_candidate_ids[:num_reqs].copy_(candidate_ids) + 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) + if self.draft_logits is not None: + self._cache_draft_logits(candidate_ids, num_sample) diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py b/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py index d805c88582..8c8d56ba31 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py @@ -38,9 +38,17 @@ def get_eagle3_aux_layers_from_config( if not (spec_config and spec_config.draft_model_config): return None hf_config = spec_config.draft_model_config.hf_config - if not hasattr(hf_config, "eagle_aux_hidden_state_layer_ids"): - return None - layer_ids = hf_config.eagle_aux_hidden_state_layer_ids + layer_ids = getattr(hf_config, "eagle_aux_hidden_state_layer_ids", None) + if not layer_ids: + dflash_config = getattr(hf_config, "dflash_config", None) + if dflash_config and isinstance(dflash_config, dict): + # DFlash checkpoint indices identify decoder layers, while target + # hidden-state capture uses layer boundaries. + layer_ids = [i + 1 for i in (dflash_config.get("target_layer_ids") or [])] + if not layer_ids: + target_layer_ids = getattr(hf_config, "target_layer_ids", None) + if target_layer_ids: + layer_ids = [i + 1 for i in target_layer_ids] if layer_ids and isinstance(layer_ids, (list, tuple)): return tuple(layer_ids) return None diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py index 6ae3fe793b..8376a51558 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import Any +import numpy as np import torch import torch.nn as nn @@ -481,6 +482,10 @@ def propose( skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, is_profile: bool = False, + output_copy_event: torch.cuda.Event | None = None, + sampled_token_ids_cpu: np.ndarray | None = None, + num_sampled_tokens_cpu: np.ndarray | None = None, + all_token_ids_cpu: np.ndarray | None = None, ) -> torch.Tensor: num_tokens = input_batch.num_tokens_after_padding num_reqs = input_batch.num_reqs diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py index fcbfc5569e..6826a234f0 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py @@ -24,6 +24,13 @@ def _should_share(eagle: nn.Module, flag: str, draft, target) -> bool: return torch.equal(w, target.weight) +def get_target_lm_head(target_model: nn.Module, target_language_model: nn.Module): + """Return the language-model head for plain and conditional targets.""" + return getattr(target_language_model, "lm_head", None) or getattr( + target_model, "lm_head", None + ) + + def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module: from vllm.compilation.backends import set_model_tag diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py index 0cfbdf4182..da9cca92d3 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py @@ -2,8 +2,12 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch -from vllm.triton_utils import tl, triton -from vllm.v1.worker.gpu.sample.gumbel import gumbel_block_argmax, tl_rand64 +from vllm.triton_utils import tl, tldevice, triton +from vllm.v1.worker.gpu.sample.gumbel import ( + gumbel_block_argmax, + gumbel_noised_argmax, + tl_rand32, +) @triton.jit @@ -243,7 +247,7 @@ def _rejection_kernel( if SYNTHETIC_MODE: pos = tl.load(pos_ptr + logit_idx) - u = tl_rand64(seed, pos, includes_zero=False) + u = tl_rand32(seed, pos, includes_zero=False) rate = tl.load(synthetic_conditional_rates_ptr + i) accepted &= u < rate else: @@ -267,7 +271,7 @@ def _rejection_kernel( ) target_log_prob = target_logit - target_lse pos = tl.load(pos_ptr + logit_idx) - u = tl_rand64(seed, pos, includes_zero=False) + u = tl_rand32(seed, pos, includes_zero=False) if HAS_DRAFT_LOGITS: draft_logit = tl.load( draft_logits_ptr @@ -388,7 +392,7 @@ def _resample_kernel( ratio = tl.exp(draft_log_probs - target_log_probs) residual_logits = tl.where( ratio < 1.0, - target_log_probs + tl.log(1 - ratio), + target_log_probs + tldevice.log1p(-ratio), float("-inf"), ).to(tl.float32) else: @@ -413,6 +417,7 @@ def _resample_kernel( pos_ptr, None, # processed_logits_ptr 0, # processed_logits_stride + 0, # processed_logits_col_stride None, # processed_logits_col_ptr vocab_size, APPLY_TEMPERATURE=False, @@ -490,6 +495,327 @@ def _insert_resampled_kernel( ) +@triton.jit +def _dflash2_compact_target_row( + target_ids_ptr, + target_logits_ptr, + target_stride, + row, + temperature, + top_p, + TOP_K: tl.constexpr, + BLOCK_K: tl.constexpr, +): + offsets = tl.arange(0, BLOCK_K) + mask = offsets < TOP_K + token_ids = tl.load( + target_ids_ptr + row * target_stride + offsets, + mask=mask, + other=0, + ).to(tl.int64) + logits = tl.load( + target_logits_ptr + row * target_stride + offsets, + mask=mask, + other=float("-inf"), + ).to(tl.float32) + if temperature != 1.0: + logits = logits / temperature + + # get_topk_tokens_and_logits returns scores in descending order. Apply + # nucleus filtering only within that exact top-k support, matching the + # target's top-k-first sampling contract without materializing V columns. + row_max = tl.max(logits, axis=0) + unnormalized = tl.where(mask, tl.exp(logits - row_max), 0.0) + normalizer = tl.sum(unnormalized, axis=0) + probs = unnormalized / normalizer + cumulative_before = tl.cumsum(probs, axis=0) - probs + keep = mask & ((top_p >= 1.0) | (cumulative_before < top_p)) + logits = tl.where(keep, logits, float("-inf")) + + kept_max = tl.max(logits, axis=0) + kept_sumexp = tl.sum(tl.where(keep, tl.exp(logits - kept_max), 0.0), axis=0) + logsumexp = kept_max + tl.log(kept_sumexp) + return token_ids, logits, keep, logsumexp + + +@triton.jit +def _dflash2_sparse_topk_rejection_kernel( + # [num_reqs, num_speculative_steps + 1] + sampled_ptr, + sampled_stride, + # [num_reqs] + num_sampled_ptr, + # [num_logits, target_top_k] + target_ids_ptr, + target_logits_ptr, + target_stride, + # [max_num_reqs, num_speculative_steps, draft_top_k] + draft_ids_ptr, + draft_logits_ptr, + draft_stride_0, + draft_stride_1, + # [num_logits] + draft_sampled_ptr, + # [num_reqs + 1] + cu_num_logits_ptr, + # [num_reqs] + idx_mapping_ptr, + # [max_num_reqs] + temperature_ptr, + top_p_ptr, + seed_ptr, + # [num_logits] + pos_ptr, + num_speculative_steps: tl.constexpr, + TARGET_TOP_K: tl.constexpr, + TARGET_BLOCK_K: tl.constexpr, + DRAFT_TOP_K: tl.constexpr, + DRAFT_BLOCK_K: tl.constexpr, + USE_FP64: tl.constexpr, +): + """DFlash2 chain rejection on compact target/draft supports. + + The target distribution is zero outside its top-k/top-p support and the + DFlash2 proposal is zero outside selector top-k. Therefore both the + acceptance ratio and ``relu(p - q)`` recovery distribution are exact on + these compact sets; scanning or storing a full-vocabulary row is unnecessary. + """ + req_idx = tl.program_id(0) + req_state_idx = tl.load(idx_mapping_ptr + req_idx) + start_idx = tl.load(cu_num_logits_ptr + req_idx) + end_idx = tl.load(cu_num_logits_ptr + req_idx + 1) + num_tokens = end_idx - start_idx + temperature = tl.load(temperature_ptr + req_state_idx).to(tl.float32) + top_p = tl.load(top_p_ptr + req_state_idx).to(tl.float32) + seed = tl.load(seed_ptr + req_state_idx) + + target_offsets = tl.arange(0, TARGET_BLOCK_K) + target_mask = target_offsets < TARGET_TOP_K + draft_offsets = tl.arange(0, DRAFT_BLOCK_K) + draft_mask = draft_offsets < DRAFT_TOP_K + + accepted = True + accepted_steps = tl.zeros((), dtype=tl.int32) + for step in range(num_speculative_steps): + if accepted and step < num_tokens - 1: + row = start_idx + step + target_ids, target_logits, target_keep, target_lse = ( + _dflash2_compact_target_row( + target_ids_ptr, + target_logits_ptr, + target_stride, + row, + temperature, + top_p, + TOP_K=TARGET_TOP_K, + BLOCK_K=TARGET_BLOCK_K, + ) + ) + proposed = tl.load(draft_sampled_ptr + row + 1).to(tl.int64) + target_proposed_logit = tl.max( + tl.where( + target_keep & (target_ids == proposed), + target_logits, + float("-inf"), + ), + axis=0, + ) + + draft_base = ( + draft_logits_ptr + + req_state_idx * draft_stride_0 + + step * draft_stride_1 + ) + draft_id_base = ( + draft_ids_ptr + req_state_idx * draft_stride_0 + step * draft_stride_1 + ) + draft_ids = tl.load( + draft_id_base + draft_offsets, + mask=draft_mask, + other=0, + ).to(tl.int64) + draft_logits = tl.load( + draft_base + draft_offsets, + mask=draft_mask, + other=float("-inf"), + ).to(tl.float32) + draft_max, draft_sumexp = _compute_block_max_and_sumexp(draft_logits) + draft_lse = draft_max + tl.log(draft_sumexp) + draft_proposed_logit = tl.max( + tl.where( + draft_mask & (draft_ids == proposed), + draft_logits, + float("-inf"), + ), + axis=0, + ) + + pos = tl.load(pos_ptr + row) + uniform = tl_rand32(seed, pos, includes_zero=False) + accepted &= (target_proposed_logit - target_lse) > ( + tl.log(uniform) + draft_proposed_logit - draft_lse + ) + tl.store( + sampled_ptr + req_idx * sampled_stride + step, + proposed, + ) + accepted_steps += accepted.to(tl.int32) + + resample_step = accepted_steps + resample_row = start_idx + resample_step + target_ids, target_logits, target_keep, target_lse = _dflash2_compact_target_row( + target_ids_ptr, + target_logits_ptr, + target_stride, + resample_row, + temperature, + top_p, + TOP_K=TARGET_TOP_K, + BLOCK_K=TARGET_BLOCK_K, + ) + is_bonus = resample_row == end_idx - 1 + residual_logits = target_logits + if not is_bonus: + draft_base = ( + draft_logits_ptr + + req_state_idx * draft_stride_0 + + resample_step * draft_stride_1 + ) + draft_id_base = ( + draft_ids_ptr + + req_state_idx * draft_stride_0 + + resample_step * draft_stride_1 + ) + draft_ids = tl.load( + draft_id_base + draft_offsets, + mask=draft_mask, + other=0, + ).to(tl.int64) + draft_logits = tl.load( + draft_base + draft_offsets, + mask=draft_mask, + other=float("-inf"), + ).to(tl.float32) + draft_max, draft_sumexp = _compute_block_max_and_sumexp(draft_logits) + draft_lse = draft_max + tl.log(draft_sumexp) + + id_match = target_ids[:, None] == draft_ids[None, :] + draft_logits_at_target = tl.max( + tl.where( + id_match & target_mask[:, None] & draft_mask[None, :], + draft_logits[None, :], + float("-inf"), + ), + axis=1, + ) + target_log_probs = target_logits - target_lse + draft_log_probs = draft_logits_at_target - draft_lse + ratio = tl.exp(draft_log_probs - target_log_probs) + residual_logits = tl.where( + target_keep & (ratio < 1.0), + target_log_probs + tldevice.log1p(-ratio), + float("-inf"), + ) + + pos = tl.load(pos_ptr + resample_row) + _, selected_offset = gumbel_noised_argmax( + residual_logits, + target_ids, + target_keep, + seed, + pos, + temperature, + USE_FP64=USE_FP64, + APPLY_TEMPERATURE=False, + ) + selected_token = tl.load( + target_ids_ptr + resample_row * target_stride + selected_offset + ) + tl.store( + sampled_ptr + req_idx * sampled_stride + resample_step, + selected_token, + ) + tl.store(num_sampled_ptr + req_idx, resample_step + 1) + + +def dflash2_sparse_topk_rejection_sample( + target_topk_ids: torch.Tensor, + target_topk_logits: torch.Tensor, + draft_topk_ids: torch.Tensor, + draft_topk_logits: torch.Tensor, + draft_sampled: torch.Tensor, + cu_num_logits: torch.Tensor, + pos: torch.Tensor, + idx_mapping: torch.Tensor, + temperature: torch.Tensor, + top_p: torch.Tensor, + seed: torch.Tensor, + num_speculative_steps: int, + use_fp64: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run exact DFlash2 rejection on compact top-k distributions. + + Callers must enforce the no-penalty/no-logit-bias contract and provide + target candidates after a global top-k merge. Draft scores must be the + temperature-applied scores used to draw the corresponding proposal. + """ + if target_topk_ids.ndim != 2 or target_topk_logits.ndim != 2: + raise ValueError("target top-k tensors must be two-dimensional") + if target_topk_ids.shape != target_topk_logits.shape: + raise ValueError("target top-k ids/logits shapes must match") + if draft_topk_ids.ndim != 3 or draft_topk_logits.ndim != 3: + raise ValueError("draft top-k tensors must be three-dimensional") + if draft_topk_ids.shape != draft_topk_logits.shape: + raise ValueError("draft top-k ids/logits shapes must match") + if draft_topk_ids.shape[1] != num_speculative_steps: + raise ValueError("draft top-k step count does not match configuration") + if target_topk_ids.shape[0] != draft_sampled.shape[0]: + raise ValueError("target rows and sampled-token rows must match") + if not 0 < target_topk_ids.shape[1] <= 64: + raise ValueError("target top-k width must be in [1, 64]") + if not 0 < draft_topk_ids.shape[2] <= 64: + raise ValueError("draft top-k width must be in [1, 64]") + + target_topk_ids = target_topk_ids.contiguous() + target_topk_logits = target_topk_logits.contiguous() + draft_topk_ids = draft_topk_ids.contiguous() + draft_topk_logits = draft_topk_logits.contiguous() + num_reqs = cu_num_logits.shape[0] - 1 + sampled = draft_sampled.new_empty( + num_reqs, + num_speculative_steps + 1, + dtype=torch.int64, + ) + num_sampled = sampled.new_empty(num_reqs, dtype=torch.int32) + _dflash2_sparse_topk_rejection_kernel[(num_reqs,)]( + sampled, + sampled.stride(0), + num_sampled, + target_topk_ids, + target_topk_logits, + target_topk_ids.stride(0), + draft_topk_ids, + draft_topk_logits, + draft_topk_ids.stride(0), + draft_topk_ids.stride(1), + draft_sampled, + cu_num_logits, + idx_mapping, + temperature, + top_p, + seed, + pos, + num_speculative_steps=num_speculative_steps, + TARGET_TOP_K=target_topk_ids.shape[1], + TARGET_BLOCK_K=triton.next_power_of_2(target_topk_ids.shape[1]), + DRAFT_TOP_K=draft_topk_ids.shape[2], + DRAFT_BLOCK_K=triton.next_power_of_2(draft_topk_ids.shape[2]), + USE_FP64=use_fp64, + num_warps=1, + ) + return sampled, num_sampled + + def rejection_sample( # [num_logits, V] target_logits: torch.Tensor, diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py new file mode 100644 index 0000000000..0775cfdca6 --- /dev/null +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -0,0 +1,371 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from abc import ABC, abstractmethod +from collections.abc import Mapping +from typing import Any + +import numpy as np +import torch +import torch.nn as nn + +from vllm.config import VllmConfig, get_layers_from_vllm_config +from vllm.config.compilation import CUDAGraphMode +from vllm.logger import init_logger +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.model_executor.models import supports_multimodal +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.worker.gpu.attn_utils import ( + build_attn_metadata, + init_attn_backend, +) +from vllm.v1.worker.gpu.block_table import BlockTables +from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers +from vllm.v1.worker.gpu.model_states.interface import ModelState +from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample +from vllm.v1.worker.utils import AttentionGroup + +logger = init_logger(__name__) + + +class BaseSpeculator(ABC): + @abstractmethod + def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: + pass + + @abstractmethod + def capture(self) -> None: + pass + + @abstractmethod + def propose( + self, + input_batch: InputBatch, + attn_metadata: dict[str, Any], + slot_mappings: dict[str, torch.Tensor], + # [num_tokens, hidden_size] + last_hidden_states: torch.Tensor, + # num_layers x [num_tokens, hidden_size] + aux_hidden_states: list[torch.Tensor] | None, + # [num_reqs] + num_sampled: torch.Tensor, + # [num_reqs] + num_rejected: torch.Tensor, + # [max_num_reqs] + last_sampled: torch.Tensor, + # [num_prefill_lookahead, max_num_reqs] + next_prefill_tokens: torch.Tensor, + # [max_num_reqs] + temperature: torch.Tensor, + # [max_num_reqs] + seeds: torch.Tensor, + num_tokens_across_dp: torch.Tensor | None = None, + dummy_run: bool = False, + skip_attn_for_dummy_run: bool = False, + mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, + is_profile: bool = False, + output_copy_event: torch.cuda.Event | None = None, + sampled_token_ids_cpu: np.ndarray | None = None, + num_sampled_tokens_cpu: np.ndarray | None = None, + all_token_ids_cpu: np.ndarray | None = None, + ) -> torch.Tensor: + pass + + +class DraftModelSpeculator(BaseSpeculator): + def __init__(self, vllm_config: VllmConfig, device: torch.device): + self.vllm_config = vllm_config + self.device = device + + assert vllm_config.speculative_config is not None + self.speculative_config = vllm_config.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 + + self.scheduler_config = vllm_config.scheduler_config + self.max_num_reqs = self.scheduler_config.max_num_seqs + self.max_num_tokens = self.scheduler_config.max_num_batched_tokens + self.max_model_len = vllm_config.model_config.max_model_len + self.draft_max_seq_len = self.max_model_len + # 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() + # Widen for HC-multiplexed residuals (e.g. DeepSeek V4 feeds the MTP + # draft the target's pre-hc_head (T, hc_mult * hidden_size) residual). + # Non-HC models default to hc_mult=1 and are unaffected. + hc_mult = getattr(self.draft_model_config.hf_config, "hc_mult", 1) + self.hidden_size = self.hidden_size * hc_mult + 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 + ) + + # DP configuration + self.dp_size = vllm_config.parallel_config.data_parallel_size + self.dp_rank = vllm_config.parallel_config.data_parallel_rank + + self.input_buffers = InputBuffers( + max_num_reqs=self.max_num_reqs, + max_num_tokens=self.max_num_tokens, + device=device, + ) + self.idx_mapping = torch.zeros( + self.max_num_reqs, dtype=torch.int32, device=device + ) + self.temperature = torch.zeros( + self.max_num_reqs, dtype=torch.float32, device=device + ) + self.seeds = torch.zeros(self.max_num_reqs, dtype=torch.int64, device=device) + self.draft_tokens = torch.zeros( + self.max_num_reqs, + self.num_speculative_steps, + dtype=torch.int64, + device=device, + ) + self.arange = torch.arange( + self.max_num_reqs + 1, dtype=torch.int32, device="cpu" + ) + + self.draft_logits: torch.Tensor | None = None + if self.speculative_config.draft_sample_method == "probabilistic": + # Temperature-applied logits, cached from the previous decode step. + dtype, fill = self.draft_logits_spec(vllm_config) + self.draft_logits = torch.full( + ( + self.max_num_reqs, + self.num_speculative_steps, + self.vocab_size, + ), + fill, + dtype=dtype, + device=device, + ) + + self.supports_mm_inputs = False + + @abstractmethod + def load_draft_model( + self, + target_model: nn.Module, + target_attn_layer_names: set[str], + ) -> nn.Module: + pass + + def load_model(self, target_model: nn.Module) -> None: + target_attn_layer_names = set( + get_layers_from_vllm_config( + self.vllm_config, + AttentionLayerBase, # type: ignore[type-abstract] + ).keys() + ) + + self.model = self.load_draft_model(target_model, target_attn_layer_names) + self._validate_local_argmax_reduction() + + all_attn_layers = set[str]( + get_layers_from_vllm_config( + self.vllm_config, + AttentionLayerBase, # type: ignore[type-abstract] + ).keys() + ) + self.draft_attn_layer_names = all_attn_layers - target_attn_layer_names + + target_supports_mm = MULTIMODAL_REGISTRY.supports_multimodal_inputs( + self.vllm_config.model_config + ) + draft_supports_mm = supports_multimodal(self.model) + self.supports_mm_inputs = target_supports_mm and draft_supports_mm + if target_supports_mm and not draft_supports_mm: + logger.warning_once( + "Draft model %s does not support external multimodal embeddings. " + "Embeddings from the target model will not be passed to the " + "drafter; using text-only draft inputs instead.", + type(self.model).__name__, + ) + + @property + def attn_vllm_config(self) -> VllmConfig: + """Config for the draft's attention metadata builders. Overridden by + speculators whose attention mode differs from the target's.""" + return self.vllm_config + + def set_attn( + self, + model_state: ModelState, + kv_cache_config: KVCacheConfig, + block_tables: BlockTables, + target_input_buffers: InputBuffers, + target_attn_groups: list[list[AttentionGroup]], + ) -> None: + self.model_state = model_state + self.kv_cache_config = kv_cache_config + self.attn_groups, self.attn_cg_support, _ = init_attn_backend( + kv_cache_config, + self.attn_vllm_config, + self.device, + active_layer_names=self.draft_attn_layer_names, + ) + self.block_tables = block_tables + # The target model runner's buffers and attention groups. Draft + # prefill reuses the target model's attention metadata, so its + # cudagraph capture must build dummy metadata through the same + # builders and buffers. + self.target_input_buffers = target_input_buffers + self.target_attn_groups = target_attn_groups + + def _build_draft_attn_metadata( + self, + num_reqs: int, + num_reqs_padded: int, + num_tokens_padded: int, + seq_lens_cpu_upper_bound: torch.Tensor, + step: int, + num_query_per_req: int = 1, + causal: bool | Mapping[int, bool] = True, + query_start_loc_np: np.ndarray | None = None, + dcp_local_seq_lens: torch.Tensor | None = None, + ) -> dict[str, Any] | None: + if query_start_loc_np is not None: + # Non-uniform query layout (e.g. multi-module MTP's mixed + # prefill/decode queries); num_query_per_req is ignored. + query_start_loc_cpu = torch.empty(num_reqs_padded + 1, dtype=torch.int32) + query_start_loc_cpu[: num_reqs + 1] = torch.from_numpy( + query_start_loc_np[: num_reqs + 1] + ) + query_start_loc_cpu[num_reqs:] = query_start_loc_cpu[num_reqs] + max_query_len = int( + (query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]).max() + ) + else: + # Uniform query: query_start_loc[i] = min(i, num_reqs) * num_query_per_req. + # Clamp keeps the series non-decreasing past num_reqs, which some + # attention backends require. + query_start_loc_cpu = ( + torch.clamp(self.arange[: num_reqs_padded + 1], max=num_reqs) + * num_query_per_req + ) + max_query_len = num_query_per_req + block_tables = [ + x[:num_reqs_padded] for x in self.block_tables.input_block_tables + ] + slot_mappings = self.block_tables.slot_mappings[:, :num_tokens_padded] + draft_seq_lens_cpu_upper_bound = torch.zeros( + num_reqs_padded, dtype=torch.int32, device="cpu" + ) + torch.add( + seq_lens_cpu_upper_bound[:num_reqs], + step, + out=draft_seq_lens_cpu_upper_bound[:num_reqs], + ) + draft_seq_lens_cpu_upper_bound[:num_reqs].clamp_(max=self.max_model_len) + attn_metadata = build_attn_metadata( + attn_groups=self.attn_groups, + num_reqs=num_reqs_padded, + num_tokens=num_tokens_padded, + query_start_loc_gpu=self.input_buffers.query_start_loc[ + : num_reqs_padded + 1 + ], + query_start_loc_cpu=query_start_loc_cpu, + max_query_len=max_query_len, + seq_lens=self.input_buffers.seq_lens[:num_reqs_padded], + dcp_local_seq_lens=( + None + if dcp_local_seq_lens is None + else dcp_local_seq_lens[:num_reqs_padded] + ), + max_seq_len=self.draft_max_seq_len, + block_tables=block_tables, + slot_mappings=slot_mappings, + kv_cache_config=self.kv_cache_config, + causal=causal, + seq_lens_cpu_upper_bound=draft_seq_lens_cpu_upper_bound, + ) + return attn_metadata + + def draft_logits_spec(self, vllm_config: VllmConfig) -> tuple[torch.dtype, float]: + """Dtype and initial value for the cached proposal distribution. + + Most speculators overwrite every vocabulary column and use the target + head dtype. DFlash2 overrides this because it writes only K candidates + and its selector/rejection path must share FP32 scores exactly. + """ + return vllm_config.model_config.head_dtype, 0.0 + + 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 _greedy_sample_draft(self, hidden_states: torch.Tensor) -> torch.Tensor: + 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 sample_draft( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + idx_mapping: torch.Tensor, + temperature: torch.Tensor, + seeds: 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( + logits, + idx_mapping, + temperature, + seeds, + positions + 1, + apply_temperature=True, + output_processed_logits=draft_logits, + output_processed_logits_col=draft_step, + use_fp64=self.use_fp64_gumbel, + ) + return self._greedy_sample_draft(hidden_states) + + def _copy_request_inputs( + self, + num_reqs: int, + # [num_reqs] + idx_mapping: torch.Tensor, + # [max_num_reqs] + temperature: torch.Tensor, + # [max_num_reqs] + seeds: torch.Tensor, + ) -> None: + # Copy temperature, seeds, and idx mapping to the pre-allocated buffers. + # NOTE(woosuk): For draft sampling, we only consider the temperature + # and ignore the other sampling parameters such as top_k and top_p, + # for simplicity and performance. + # While this may slightly degrade the acceptance rate, it does not + # affect the output distribution after rejection sampling. + self.temperature.copy_(temperature) + self.seeds.copy_(seeds) + self.idx_mapping[:num_reqs].copy_(idx_mapping) + # idx_mapping for CG padded requests points to -1, which is ignored + # during sampling to prevent writing stale values to draft logits. + self.idx_mapping[num_reqs:].fill_(-1) diff --git a/vllm/v1/worker/gpu/spec_decode/utils.py b/vllm/v1/worker/gpu/spec_decode/utils.py index e1fa21aeb8..e721e85b35 100644 --- a/vllm/v1/worker/gpu/spec_decode/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/utils.py @@ -45,3 +45,18 @@ def get_draft_tokens(self) -> DraftTokenIds | None: # This case only happens when async scheduling is disabled. draft_token_ids = [[-1] * self.num_draft_tokens for _ in self.req_ids] return DraftTokenIds(self.req_ids, draft_token_ids) + + +def get_parallel_drafting_token_id(hf_config) -> int: + """Resolve the mask token used by a parallel-drafting checkpoint.""" + dflash_config = getattr(hf_config, "dflash_config", None) or {} + if "mask_token_id" in dflash_config: + return int(dflash_config["mask_token_id"]) + if hasattr(hf_config, "pard_token"): + return int(hf_config.pard_token) + if hasattr(hf_config, "ptd_token_id"): + return int(hf_config.ptd_token_id) + raise ValueError( + "Model config must specify dflash_config.mask_token_id, pard_token, " + "or ptd_token_id for parallel drafting." + ) diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index 83d87c74a4..1cdcf567d6 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -15,10 +15,46 @@ NewRequestData, SchedulerOutput, ) +from vllm.v1.kv_cache_interface import CrossAttentionSpec, KVCacheSpec, MambaSpec from vllm.v1.request import Request from vllm.v1.worker.gpu.model_runner import GPUModelRunner +def _reserved_block_count( + num_tokens: int, + kv_cache_spec: KVCacheSpec, + *, + num_lookahead_tokens: int, + max_model_len: int, + max_encoder_len: int, +) -> int: + """Match the scheduler's block reservation in hand-built warmup batches.""" + if isinstance(kv_cache_spec, CrossAttentionSpec): + return cdiv(max_encoder_len, kv_cache_spec.block_size) + num_speculative_blocks = 0 + if isinstance(kv_cache_spec, MambaSpec): + num_speculative_blocks = kv_cache_spec.num_speculative_blocks + if kv_cache_spec.mamba_cache_mode == "align": + return cdiv(num_tokens, kv_cache_spec.block_size) + num_speculative_blocks + num_tokens = min(num_tokens + num_lookahead_tokens, max_model_len) + return cdiv(num_tokens, kv_cache_spec.block_size) + num_speculative_blocks + + +def _warmup_block_counter( + model_runner: GPUModelRunner, +) -> Callable[[int, KVCacheSpec], int]: + def block_count(num_tokens: int, kv_cache_spec: KVCacheSpec) -> int: + return _reserved_block_count( + num_tokens, + kv_cache_spec, + num_lookahead_tokens=model_runner.vllm_config.num_lookahead_tokens, + max_model_len=model_runner.max_model_len, + max_encoder_len=getattr(model_runner.model_state, "max_encoder_len", 0), + ) + + return block_count + + @torch.inference_mode() def warmup_kernels( model_runner: GPUModelRunner, @@ -46,9 +82,10 @@ def warmup_kernels( num_kv_cache_groups = len(kv_cache_groups) # Compute per-request block counts for each KV cache group. - group_block_sizes = [g.kv_cache_spec.block_size for g in kv_cache_groups] - prefill_block_counts = [cdiv(prompt_len, bs) for bs in group_block_sizes] - decode_block_counts = [cdiv(decode_len, bs) for bs in group_block_sizes] + block_count = _warmup_block_counter(model_runner) + kv_cache_specs = [group.kv_cache_spec for group in kv_cache_groups] + prefill_block_counts = [block_count(prompt_len, spec) for spec in kv_cache_specs] + decode_block_counts = [block_count(decode_len, spec) for spec in kv_cache_specs] decode_block_deltas = [ d - p for d, p in zip(decode_block_counts, prefill_block_counts) ] diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index e6dbb57163..c150cf791e 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -204,6 +204,7 @@ from vllm.v1.spec_decode.draft_prob_alignment import ( clone_draft_prob_token_ids, get_aligned_draft_probs, + get_aligned_draft_scalar_values, ) from vllm.v1.spec_decode.dspark import DSparkProposer from vllm.v1.spec_decode.eagle import EagleProposer @@ -270,6 +271,19 @@ _SM70_DECODE_EVENT_TRACE_CONFIG_LOGGED = False +def _unwrap_pipeline_intermediate_hidden_states( + output: torch.Tensor | IntermediateTensors, +) -> torch.Tensor: + if not isinstance(output, IntermediateTensors): + return output + try: + return output.tensors["hidden_states"] + except KeyError as exc: + raise RuntimeError( + "Pipeline intermediate output must contain a 'hidden_states' tensor." + ) from exc + + def _dflash_ddtree_debug_enabled() -> bool: return os.getenv("VLLM_DFLASH_DDTREE_DEBUG", "0") == "1" @@ -391,7 +405,10 @@ def _dflash_ddtree_target_forward_nvtx_enabled() -> bool: def _dflash_ddtree_target_forward_profiler_step() -> int: - raw = os.getenv("VLLM_DFLASH_DDTREE_TARGET_FORWARD_PROFILER_STEP", "0") + raw = os.getenv( + "VLLM_SM70_SPEC_TARGET_FORWARD_PROFILER_STEP", + os.getenv("VLLM_DFLASH_DDTREE_TARGET_FORWARD_PROFILER_STEP", "0"), + ) try: return max(0, int(raw)) except ValueError: @@ -1141,13 +1158,19 @@ def _dflash_ddtree_target_forward_profile_scope( if ( not use_spec_decode or self.speculative_config is None - or not self.speculative_config.use_dflash_ddtree() + or not ( + self.speculative_config.use_dflash_ddtree() + or self.speculative_config.use_dspark() + ) or self.device.type != "cuda" ): yield return - nvtx_enabled = _dflash_ddtree_target_forward_nvtx_enabled() + nvtx_enabled = bool( + _dflash_ddtree_target_forward_nvtx_enabled() + or os.getenv("VLLM_SM70_SPEC_TARGET_FORWARD_NVTX", "0") == "1" + ) profiler_step = _dflash_ddtree_target_forward_profiler_step() if not nvtx_enabled and profiler_step <= 0: yield @@ -1155,9 +1178,13 @@ def _dflash_ddtree_target_forward_profile_scope( step = getattr(self, "_dflash_ddtree_target_forward_profile_step", 0) + 1 self._dflash_ddtree_target_forward_profile_step = step - label = ( + label_prefix = ( "ddtree_target_forward" - f":step={step}:tokens={num_tokens}:reqs={num_reqs}" + if self.speculative_config.use_dflash_ddtree() + else "dspark_target_forward" + ) + label = ( + f"{label_prefix}:step={step}:tokens={num_tokens}:reqs={num_reqs}" f":mode={cudagraph_mode.name}:pid={os.getpid()}" ) profile_this_step = profiler_step > 0 and step == profiler_step @@ -1373,6 +1400,23 @@ def __init__( # Async scheduling self.use_async_scheduling = bool(self.scheduler_config.async_scheduling) + self.dspark_confidence_scheduling = bool( + self.speculative_config is not None + and self.speculative_config.use_dspark() + and ( + self.speculative_config.dspark_confidence_threshold > 0.0 + or ( + self.speculative_config.dspark_max_verification_tokens is not None + and self.speculative_config.dspark_max_verification_tokens + < self.num_spec_tokens + ) + ) + ) + if self.dspark_confidence_scheduling and self.use_async_scheduling: + raise ValueError( + "DSpark confidence prefix scheduling currently requires " + "synchronous scheduling." + ) self._sm70_async_worker_execute_trace_step = 0 self._sm70_async_worker_sample_trace_step = 0 self._sm70_async_worker_input_prep_trace_step = 0 @@ -1499,7 +1543,7 @@ def __init__( elif self.speculative_config.use_dspark(): self.drafter = DSparkProposer(self.vllm_config, self.device, self) self.use_aux_hidden_state_outputs = True - elif self.speculative_config.use_dflash(): + elif self.speculative_config.use_dflash_ddtree(): self.drafter = DFlashProposer(self.vllm_config, self.device, self) self.use_aux_hidden_state_outputs = ( self.drafter.eagle3_use_aux_hidden_state @@ -1767,6 +1811,18 @@ def __init__( self._draft_probs: torch.Tensor | None = None self._draft_prob_req_ids: list[str] | None = None self._draft_prob_token_ids: list[list[int]] | torch.Tensor | None = None + self._draft_confidence_logits: torch.Tensor | None = None + self._draft_confidence_req_ids: list[str] | None = None + self._draft_confidence_token_ids: list[list[int]] | torch.Tensor | None = None + self._dspark_verification_lengths: torch.Tensor | None = None + self._dspark_verification_lengths_cpu: torch.Tensor | None = None + if self.dspark_confidence_scheduling: + self._dspark_verification_lengths_cpu = torch.empty( + self.max_num_reqs, + dtype=torch.int32, + device="cpu", + pin_memory=self.pin_memory, + ) self._dflash_ddtree_payloads: tuple[DDTreeDraftPayload, ...] | None = None self._ddtree_parent_metadata: DDTreeParentMetadata | None = None self._ddtree_accepted_rows_cpu_sidecar: list[list[int]] | None = None @@ -1808,6 +1864,12 @@ def __init__( self.valid_sampled_token_count_cpu: torch.Tensor | None = None self.draft_token_ids_cpu: torch.Tensor | None = None self.num_accepted_tokens_event: torch.Event | None = None + # Row ownership for the runner-owned accepted-token D2H snapshots. + # Keep the request object as well as its ID so abort-and-resubmit with + # the same ID is still treated as a new request. + self._mamba_accepted_token_state_rows: dict[ + str, tuple[int, CachedRequestState] + ] = {} if self.num_spec_tokens: self.draft_token_ids_event = torch.Event() self.num_accepted_tokens_event = torch.Event() @@ -2126,10 +2188,33 @@ def _warmup_sm70_aux_kernels(self) -> None: and mamba_utils.warmup_batch_memcpy_kernel(self.device) ): warmed.append("mamba_batch_memcpy") + if ( + self.cache_config.mamba_cache_mode == "align" + and self.speculative_config is not None + and self.model_config.is_hybrid + ): + mamba_bufs = self._get_mamba_bufs() + postprocess_ctx = mamba_bufs.postprocess_align + assert postprocess_ctx is not None + if not postprocess_ctx.is_initialized: + postprocess_ctx.initialize_from_forward_context( + self.kv_cache_config, + self.compilation_config.static_forward_context, + self.model.get_mamba_state_copy_func(), + [ + self.input_batch.block_table[group_id].get_device_tensor(1) + for group_id in postprocess_ctx.mamba_group_ids + ], + ) + if postprocess_ctx.warmup_fused_postprocess(): + warmed.append("mamba_spec_postprocess") drafter = getattr(self, "drafter", None) mtp_warmup = getattr(drafter, "warmup_sm70_mtp_hotpath_kernels", None) if mtp_warmup is not None: warmed.extend(mtp_warmup()) + mtp_moe_warmup = getattr(drafter, "warmup_sm70_mtp_moe_kernels", None) + if mtp_moe_warmup is not None: + warmed.extend(mtp_moe_warmup()) dflash_warmup = getattr(drafter, "warmup_sm70_dflash_hotpath_kernels", None) if dflash_warmup is not None: warmed.extend(dflash_warmup()) @@ -2589,6 +2674,10 @@ def _update_states_after_model_execute( # Count only the contiguous accepted prefix. Values after the first -1 # are rejected/padding slots and may contain stale token ids. num_reqs = output_token_ids.size(0) + self._mamba_accepted_token_state_rows = { + req_id: (i, self.requests[req_id]) + for i, req_id in enumerate(self.input_batch.req_ids[:num_reqs]) + } profile_count_t0 = time.perf_counter() if profile_enabled else 0.0 sidecar_values = None if ( @@ -2715,19 +2804,16 @@ def _update_states_after_model_execute( else None ), ) + # CPU postprocess updates InputBatch in place. Mirror its final + # values into the runner-owned snapshot used by the next step. + self.num_accepted_tokens.cpu[:num_reqs].copy_( + self.input_batch.num_accepted_tokens_cpu_tensor[:num_reqs] + ) + self.spec_state_slot_selectors.cpu[:num_reqs].copy_( + self.input_batch.spec_num_accepted_tokens_cpu_tensor[:num_reqs] + ) else: profile_mamba_stage_t0 = time.perf_counter() if profile_enabled else 0.0 - if spec_state_slot_selectors_cpu is not None: - self.input_batch.spec_num_accepted_tokens_cpu_tensor[ - :num_reqs - ].copy_(spec_state_slot_selectors_cpu) - else: - self.input_batch.spec_num_accepted_tokens_cpu_tensor[ - :num_reqs - ].copy_( - self.spec_state_slot_selectors.gpu[:num_reqs], - non_blocking=True, - ) assert mamba_bufs.postprocess_align is not None mamba_utils.stage_postprocess_inputs_to_gpu( mamba_bufs.postprocess_align, @@ -2749,11 +2835,9 @@ def _update_states_after_model_execute( num_reqs=num_reqs, num_accepted_tokens_gpu=self.num_accepted_tokens.gpu, spec_state_slot_selectors_gpu=(self.spec_state_slot_selectors.gpu), - num_accepted_tokens_cpu_tensor=( - self.input_batch.num_accepted_tokens_cpu_tensor - ), + num_accepted_tokens_cpu_tensor=self.num_accepted_tokens.cpu, spec_num_accepted_tokens_cpu_tensor=( - self.input_batch.spec_num_accepted_tokens_cpu_tensor + self.spec_state_slot_selectors.cpu ), input_batch=self.input_batch, kv_cache_config=self.kv_cache_config, @@ -2786,10 +2870,10 @@ def _update_states_after_model_execute( spec_state_slot_selectors_cpu ) else: - self.input_batch.num_accepted_tokens_cpu_tensor[:num_reqs].copy_( + self.num_accepted_tokens.cpu[:num_reqs].copy_( self.num_accepted_tokens.gpu[:num_reqs], non_blocking=True ) - self.input_batch.spec_num_accepted_tokens_cpu_tensor[:num_reqs].copy_( + self.spec_state_slot_selectors.cpu[:num_reqs].copy_( self.spec_state_slot_selectors.gpu[:num_reqs], non_blocking=True, ) @@ -4658,6 +4742,45 @@ def _get_encoder_seq_lens( return encoder_seq_lens, encoder_seq_lens_cpu + def _sync_mamba_accepted_token_state( + self, + scheduler_output: "SchedulerOutput", + num_reqs: int, + ) -> None: + """Remap the previous step's accepted-token state by request. + + The GPU postprocess copies into runner-owned CPU buffers. InputBatch + rows can be removed, reused, or condensed before this synchronization, + so row-for-row writeback is not safe even with synchronous scheduling. + """ + previous_counts = self.num_accepted_tokens.np.copy() + previous_selectors = self.spec_state_slot_selectors.np.copy() + previous_rows = self._mamba_accepted_token_state_rows + reset_req_ids = set(scheduler_output.scheduled_cached_reqs.resumed_req_ids) + reset_req_ids.update( + req_data.req_id for req_data in scheduler_output.scheduled_new_reqs + ) + + current_counts = self.num_accepted_tokens.np[:num_reqs] + current_selectors = self.spec_state_slot_selectors.np[:num_reqs] + for current_idx, req_id in enumerate(self.input_batch.req_ids[:num_reqs]): + previous = previous_rows.get(req_id) + if ( + previous is not None + and req_id not in reset_req_ids + and self.requests.get(req_id) is previous[1] + ): + previous_idx = previous[0] + current_counts[current_idx] = previous_counts[previous_idx] + current_selectors[current_idx] = previous_selectors[previous_idx] + else: + # A new/restored request has no accepted speculative prefix. + current_counts[current_idx] = 1 + current_selectors[current_idx] = 1 + + self.input_batch.num_accepted_tokens_cpu[:num_reqs] = current_counts + self.input_batch.spec_num_accepted_tokens_cpu[:num_reqs] = current_selectors + def _prepare_inputs( self, scheduler_output: "SchedulerOutput", @@ -4869,32 +4992,7 @@ def _prepare_inputs( self.num_accepted_tokens_event, "GPUModelRunner.num_accepted_tokens_event.synchronize", ) - # Async mode: condense() reordered indices, use prev_positions mapping - if self.use_async_scheduling and prev_req_id_to_index: - prev_idx = self.prev_positions.np[:num_reqs] - new_mask = prev_idx < 0 - prev_idx_or_zero = np.where(new_mask, 0, prev_idx) - align_counts = self.input_batch.num_accepted_tokens_cpu[ - prev_idx_or_zero - ].copy() - align_counts[new_mask] = 1 - self.input_batch.num_accepted_tokens_cpu[:num_reqs] = align_counts - spec_counts = self.input_batch.spec_num_accepted_tokens_cpu[ - prev_idx_or_zero - ] - spec_counts = spec_counts.copy() - spec_counts[new_mask] = 1 - self.input_batch.spec_num_accepted_tokens_cpu[:num_reqs] = spec_counts - self.num_accepted_tokens.np[:num_reqs] = align_counts - self.spec_state_slot_selectors.np[:num_reqs] = spec_counts - else: - # Non-async mode: use values directly - self.num_accepted_tokens.np[:num_reqs] = ( - self.input_batch.num_accepted_tokens_cpu[:num_reqs] - ) - self.spec_state_slot_selectors.np[:num_reqs] = ( - self.input_batch.spec_num_accepted_tokens_cpu[:num_reqs] - ) + self._sync_mamba_accepted_token_state(scheduler_output, num_reqs) self.num_accepted_tokens.np[num_reqs:].fill(1) self.spec_state_slot_selectors.np[num_reqs:].fill(1) self._copy_buffer_to_gpu(self.num_accepted_tokens) @@ -5153,6 +5251,7 @@ def _build_attention_metadata( slot_mappings: dict[int, torch.Tensor] | None = None, ddtree_parent_metadata: DDTreeParentMetadata | None = None, cudagraph_capture_max_seq_len: int | None = None, + cudagraph_graph_variant: int | None = None, ) -> tuple[PerLayerAttnMetadata, CommonAttentionMetadata | None]: """ :return: tuple[attn_metadata, spec_decode_common_attn_metadata] @@ -5257,6 +5356,20 @@ def _get_block_table(kv_cache_gid: int): seq_lens_cpu = None num_computed_tokens_cpu = None + # Prefix-anchored SWA: pass per-request prompt lengths so the + # attention backend can keep the prefix globally visible. The backend + # owns the persistent device buffer. + prefix_anchor_lens = None + if ( + getattr( + self.vllm_config.attention_config, + "prefix_anchored_decode_window", + None, + ) + is not None + ): + prefix_anchor_lens = num_prompt_tokens_cpu + cm_base = CommonAttentionMetadata( query_start_loc=self.query_start_loc.gpu[: num_reqs_padded + 1], query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs_padded + 1], @@ -5271,8 +5384,10 @@ def _get_block_table(kv_cache_gid: int): block_table_tensor=block_table_gid_0, slot_mapping=slot_mapping_gid_0, causal=True, + cudagraph_graph_variant=cudagraph_graph_variant, is_prefilling=is_prefilling, positions=self.positions[:num_tokens_padded], + prefix_anchor_lens=prefix_anchor_lens, ) current_mamba_state_block_ids_by_gid: dict[int, torch.Tensor] = {} @@ -6350,9 +6465,14 @@ def sync_and_gather_intermediate_tensors( local_len = num_tokens // tp v = get_tp_group().all_gather(v[:local_len], dim=0) - self.intermediate_tensors[k][:num_tokens].copy_( - v[:num_tokens], non_blocking=True - ) + destination = self.intermediate_tensors[k][:num_tokens] + source = v[:num_tokens] + if ( + destination.data_ptr() != source.data_ptr() + or destination.shape != source.shape + or destination.stride() != source.stride() + ): + destination.copy_(source, non_blocking=True) return IntermediateTensors( {k: v[:num_tokens] for k, v in self.intermediate_tensors.items()} @@ -6888,6 +7008,26 @@ def _sample( _maybe_dump_sm70_qwen_layer_graph_buffers("pre_sample") _maybe_sync_sm70_sample_tensors(sample_hidden_states) if spec_decode_metadata is None: + if logits is None and self._can_use_sm70_compact_topk20_tokens( + scheduler_output, spec_decode_metadata + ): + topk_token_ids, topk_logits = self.model.get_topk_tokens_and_logits( + sample_hidden_states, + 20, + ) + compact_sampled = self.sampler.sm70_compact_topk20_pairs_sample( + topk_logits, + topk_token_ids, + sampling_metadata, + ) + if compact_sampled is not None: + logger.info_once( + "SM70 TP-local exact top-k20 random sampling path enabled." + ) + return SamplerOutput( + sampled_token_ids=compact_sampled.to(torch.int32).unsqueeze(-1), + logprobs_tensors=None, + ) sampler_output = self._sample_greedy_token_fastpath( logits, sampling_metadata, @@ -7320,6 +7460,9 @@ def _sample( ) draft_probs = self._get_spec_decode_draft_probs(spec_decode_metadata) + draft_confidence_logits = self._get_spec_decode_confidence_logits( + spec_decode_metadata + ) if ( self.speculative_config is not None and self.speculative_config.method == "mtp" @@ -7339,6 +7482,7 @@ def _sample( draft_probs, logits, sampling_metadata, + draft_confidence_logits=draft_confidence_logits, ) target_candidate_ids = self.rejection_sampler.take_last_target_candidate_ids() if target_candidate_ids is not None and hasattr( @@ -7438,6 +7582,88 @@ def _can_use_greedy_token_fastpath( self._trace_greedy_token_fastpath("enabled") return True + def _can_use_sm70_compact_topk20_tokens( + self, + scheduler_output: "SchedulerOutput", + spec_decode_metadata: SpecDecodeMetadata | None, + ) -> bool: + if not envs.VLLM_SM70_COMPACT_TOPK20_SAMPLER: + return False + if not envs.VLLM_SM70_TP_LOCAL_TOPK20_SAMPLER: + return False + if spec_decode_metadata is not None: + return self._reject_sm70_compact_topk20_tokens("spec_decode") + if self.is_pooling_model or self.broadcast_pp_output: + return self._reject_sm70_compact_topk20_tokens("pooling_or_pp_broadcast") + if scheduler_output.has_structured_output_requests: + return self._reject_sm70_compact_topk20_tokens("structured_output") + if self.device.type != "cuda": + return self._reject_sm70_compact_topk20_tokens("non_cuda") + if torch.cuda.get_device_capability(self.device) != (7, 0): + return self._reject_sm70_compact_topk20_tokens("non_sm70") + if self.model_config.get_vocab_size() != 248320: + return self._reject_sm70_compact_topk20_tokens("vocab_size") + if not hasattr(self.model, "get_topk_tokens_and_logits"): + return self._reject_sm70_compact_topk20_tokens("missing_model_topk") + if self.input_batch.num_reqs != 1: + return self._reject_sm70_compact_topk20_tokens("batch_size") + + sampling_metadata = self.input_batch.sampling_metadata + if not sampling_metadata.all_random or sampling_metadata.all_greedy: + return self._reject_sm70_compact_topk20_tokens("sampling_mode") + if sampling_metadata.max_num_logprobs is not None: + return self._reject_sm70_compact_topk20_tokens("logprobs") + if sampling_metadata.logprob_token_ids: + return self._reject_sm70_compact_topk20_tokens("logprob_token_ids") + if not sampling_metadata.no_penalties: + return self._reject_sm70_compact_topk20_tokens("penalties") + if sampling_metadata.allowed_token_ids_mask is not None: + return self._reject_sm70_compact_topk20_tokens("allowed_token_ids") + if sampling_metadata.bad_words_token_ids: + return self._reject_sm70_compact_topk20_tokens("bad_words") + holder = sampling_metadata.thinking_budget_state_holder + if holder is not None and holder.has_tracked_requests(): + return self._reject_sm70_compact_topk20_tokens("thinking_budget") + if not self._non_argmax_logits_processors_inactive(sampling_metadata): + return self._reject_sm70_compact_topk20_tokens("logits_processor") + if not self._sm70_argmax_logits_processors_inactive(sampling_metadata): + return self._reject_sm70_compact_topk20_tokens( + "argmax_invariant_logits_processor" + ) + if sampling_metadata.top_k_cpu != (20,): + return self._reject_sm70_compact_topk20_tokens("top_k") + + top_p_cpu = sampling_metadata.top_p_cpu + if top_p_cpu is None or len(top_p_cpu) != 1 or abs(top_p_cpu[0] - 0.95) > 1e-6: + return self._reject_sm70_compact_topk20_tokens("top_p") + temperature_cpu = sampling_metadata.temperature_cpu + if ( + temperature_cpu is None + or len(temperature_cpu) != 1 + or abs(temperature_cpu[0] - 1.0) > 1e-6 + ): + return self._reject_sm70_compact_topk20_tokens("temperature") + return True + + @staticmethod + def _reject_sm70_compact_topk20_tokens(reason: str) -> bool: + logger.info_once("SM70 TP-local top-k20 route rejected: %s", reason) + return False + + @staticmethod + def _sm70_argmax_logits_processors_inactive( + sampling_metadata: SamplingMetadata, + ) -> bool: + from vllm.v1.sample.logits_processor.builtin import MinPLogitsProcessor + + for processor in sampling_metadata.logitsprocs.argmax_invariant: + if isinstance(processor, MinPLogitsProcessor): + if processor.min_p_count: + return False + continue + return False + return True + def _trace_greedy_token_fastpath(self, reason: str) -> None: if not self.sm70_greedy_token_fastpath_trace: return @@ -7786,7 +8012,7 @@ def _determine_batch_execution_and_padding( attention_context_len is None and uniform_decode and num_reqs > 0 - and self.cudagraph_dispatcher.has_attention_context_buckets + and self.cudagraph_dispatcher.has_attention_context_specialization ): attention_context_len = int( self.optimistic_seq_lens_cpu[:num_reqs].max().item() @@ -8438,6 +8664,9 @@ def execute_model( self._can_use_greedy_token_fastpath( scheduler_output, spec_decode_metadata ) + or self._can_use_sm70_compact_topk20_tokens( + scheduler_output, spec_decode_metadata + ) or self._can_use_ddtree_greedy_top_tokens( scheduler_output, spec_decode_metadata ) @@ -8792,6 +9021,10 @@ def sample_tokens( self._draft_probs = None self._draft_prob_req_ids = None self._draft_prob_token_ids = None + self._draft_confidence_logits = None + self._draft_confidence_req_ids = None + self._draft_confidence_token_ids = None + self._dspark_verification_lengths = None self._dflash_ddtree_payloads = None self._ddtree_parent_metadata = None self._ddtree_accepted_rows_cpu_sidecar = None @@ -8935,6 +9168,10 @@ def propose_draft_token_ids(sampled_token_ids): self._draft_probs = None self._draft_prob_req_ids = None self._draft_prob_token_ids = None + self._draft_confidence_logits = None + self._draft_confidence_req_ids = None + self._draft_confidence_token_ids = None + self._dspark_verification_lengths = None self._copy_draft_token_ids_to_cpu(scheduler_output) with record_function_or_nullcontext("gpu_model_runner: bookkeep"): @@ -9239,6 +9476,15 @@ def _copy_draft_token_ids_to_cpu( self.draft_token_ids_cpu[:num_reqs].copy_( draft_token_ids, non_blocking=True ) + if self.dspark_confidence_scheduling: + lengths = self._dspark_verification_lengths + lengths_cpu = self._dspark_verification_lengths_cpu + if lengths is None or lengths_cpu is None: + raise RuntimeError( + "DSpark confidence scheduling is enabled but the " + "proposer did not return verification lengths." + ) + lengths_cpu[:num_reqs].copy_(lengths[:num_reqs], non_blocking=True) else: # No copy needed, just zero-out cpu tensor. self.draft_token_ids_cpu[:num_reqs] = 0 @@ -9259,7 +9505,16 @@ def _get_draft_token_ids_cpu(self) -> tuple[list[list[int]], list[str]]: self.draft_token_ids_event, "GPUModelRunner.draft_token_ids_event.synchronize", ) - return self.draft_token_ids_cpu[: len(req_ids)].tolist(), req_ids + draft_token_ids = self.draft_token_ids_cpu[: len(req_ids)].tolist() + if self.dspark_confidence_scheduling: + lengths_cpu = self._dspark_verification_lengths_cpu + assert lengths_cpu is not None + lengths = lengths_cpu[: len(req_ids)].tolist() + draft_token_ids = [ + token_ids[: max(0, min(int(length), len(token_ids)))] + for token_ids, length in zip(draft_token_ids, lengths, strict=True) + ] + return draft_token_ids, req_ids def _copy_valid_sampled_token_count( self, next_token_ids: torch.Tensor, valid_sampled_tokens_count: torch.Tensor @@ -9309,6 +9564,22 @@ def _get_spec_decode_draft_probs( spec_decode_metadata=spec_decode_metadata, ) + def _get_spec_decode_confidence_logits( + self, spec_decode_metadata: SpecDecodeMetadata + ) -> torch.Tensor | None: + # Alignment has a small but non-zero launch/copy cost. Keep it out of + # ordinary serving until confidence scheduling is explicitly enabled; + # alignment dumps are the calibration path used before that gate. + if not envs.VLLM_SPEC_DUMP_ALIGNMENT: + return None + return get_aligned_draft_scalar_values( + req_ids=self.input_batch.req_ids, + values=self._draft_confidence_logits, + value_req_ids=self._draft_confidence_req_ids, + value_token_ids=self._draft_confidence_token_ids, + spec_decode_metadata=spec_decode_metadata, + ) + def propose_draft_token_ids( self, scheduler_output: "SchedulerOutput", @@ -9327,6 +9598,10 @@ def propose_draft_token_ids( self._draft_probs = None self._draft_prob_req_ids = None self._draft_prob_token_ids = None + self._draft_confidence_logits = None + self._draft_confidence_req_ids = None + self._draft_confidence_token_ids = None + self._dspark_verification_lengths = None self._dflash_ddtree_payloads = None self._ddtree_parent_metadata = None self._ddtree_accepted_rows_cpu_sidecar = None @@ -9451,7 +9726,7 @@ def propose_draft_token_ids( elif ( spec_config.use_eagle() - or spec_config.use_dflash() + or spec_config.use_dflash_ddtree() or spec_config.uses_draft_model() ): assert isinstance( @@ -9643,6 +9918,18 @@ def propose_draft_token_ids( self._draft_prob_token_ids = clone_draft_prob_token_ids( draft_token_ids ) + if hasattr(self.drafter, "take_last_confidence_logits"): + confidence_logits = self.drafter.take_last_confidence_logits() + if confidence_logits is not None: + self._draft_confidence_logits = confidence_logits + self._draft_confidence_req_ids = self.input_batch.req_ids.copy() + self._draft_confidence_token_ids = clone_draft_prob_token_ids( + draft_token_ids + ) + if hasattr(self.drafter, "take_last_verification_lengths"): + verification_lengths = self.drafter.take_last_verification_lengths() + if verification_lengths is not None: + self._dspark_verification_lengths = verification_lengths if hasattr(self.drafter, "take_last_ddtree_payloads"): self._dflash_ddtree_payloads = self.drafter.take_last_ddtree_payloads() first_payload = ( @@ -10501,6 +10788,11 @@ def _dummy_run( if is_graph_capturing else None ), + cudagraph_graph_variant=( + batch_desc.graph_variant + if batch_descriptor_override is not None + else None + ), ) _sm70_profile_trace( "_dummy_run attention metadata built num_tokens=%s " @@ -10615,6 +10907,15 @@ def _dummy_run( else: hidden_states = outputs + if isinstance(hidden_states, IntermediateTensors): + assert not get_pp_group().is_last_rank + # Non-last PP ranks return intermediate tensors instead of + # logits-bearing hidden states. Dummy-run callers only need a + # tensor to preserve the profiling/capture return contract. + hidden_states = _unwrap_pipeline_intermediate_hidden_states( + hidden_states + ) + if self.speculative_config and ( self.speculative_config.use_eagle() or self.speculative_config.uses_draft_model() @@ -10801,6 +11102,21 @@ def _dummy_sampler_run( logits, dummy_metadata, ) + # The mixed-sampling warmup above passes a non-null is_greedy + # tensor. Pure greedy serving passes None and otherwise triggers a + # separate Triton compile during the first decode verifier step. + greedy_metadata = replace( + dummy_metadata, + temperature=torch.zeros_like(dummy_metadata.temperature), + all_greedy=True, + all_random=False, + ) + self.rejection_sampler( + dummy_spec_decode_metadata, + None, + logits, + greedy_metadata, + ) return sampler_output def _dummy_pooler_run_task( diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 6ca3c687f7..c131a019b1 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -220,9 +220,9 @@ def _scoped_allocator_max_split(self, max_split_size_mb: int): yield return - set_allocator_settings = getattr(torch._C, - "_accelerator_setAllocatorSettings", - None) + set_allocator_settings = getattr( + torch._C, "_accelerator_setAllocatorSettings", None + ) if set_allocator_settings is None: yield return @@ -788,6 +788,92 @@ def sample_tokens( ) -> ModelRunnerOutput | AsyncModelRunnerOutput: return self.model_runner.sample_tokens(grammar_output) + def _use_sm70_static_pp_hidden_transfer(self, num_tokens: int) -> bool: + """Whether this step has the exact metadata-free PP tensor contract.""" + if not envs.VLLM_SM70_PP_STATIC_HIDDEN_TRANSFER or num_tokens != 1: + return False + cached = getattr(self, "_sm70_static_pp_hidden_contract", None) + if cached is not None: + return bool(cached) + + config = self.vllm_config + parallel_config = config.parallel_config + compilation_config = config.compilation_config + model_config = config.model_config + enabled = not ( + parallel_config.pipeline_parallel_size != 2 + or parallel_config.tensor_parallel_size != 4 + or parallel_config.enable_dbo + or parallel_config.ubatch_size > 1 + or config.scheduler_config.max_num_seqs != 1 + or getattr(config, "speculative_config", None) is not None + or compilation_config.cudagraph_mode == CUDAGraphMode.NONE + or compilation_config.pass_config.enable_sp + or model_config.dtype != torch.float16 + ) + + schema_matches = False + if enabled: + try: + schema = self.model_runner.model.make_empty_intermediate_tensors( + batch_size=1, + dtype=model_config.dtype, + device=torch.device("meta"), + ).tensors + hidden_states = schema.get("hidden_states") + schema_matches = bool( + tuple(schema) == ("hidden_states",) + and isinstance(hidden_states, torch.Tensor) + and hidden_states.shape == (1, 4, 4096) + and hidden_states.dtype == torch.float16 + and hidden_states.is_contiguous() + ) + except (AttributeError, RuntimeError, TypeError): + # This is an optional fast path; unsupported model-provided + # schema construction retains the metadata protocol. + schema_matches = False + enabled = bool( + enabled + and schema_matches + and current_platform.is_cuda() + and current_platform.is_device_capability((7, 0)) + ) + self._sm70_static_pp_hidden_contract = enabled + return enabled + + @staticmethod + def _is_static_pp_hidden_tensor_dict( + tensor_dict: dict[str, torch.Tensor], num_tokens: int + ) -> bool: + if tuple(tensor_dict) != ("hidden_states",): + return False + hidden_states = tensor_dict["hidden_states"] + return bool( + hidden_states.shape == (num_tokens, 4, 4096) + and hidden_states.dtype == torch.float16 + and hidden_states.is_cuda + and hidden_states.is_contiguous() + ) + + def _static_pp_hidden_recv_buffers( + self, num_tokens: int + ) -> dict[str, torch.Tensor] | None: + if not self._use_sm70_static_pp_hidden_transfer(num_tokens): + return None + buffers = self.model_runner.intermediate_tensors + if buffers is None or tuple(buffers.tensors) != ("hidden_states",): + raise RuntimeError( + "static PP hidden transfer requires one persistent " + "hidden_states input buffer" + ) + hidden_states = buffers.tensors["hidden_states"][:num_tokens] + tensor_dict = {"hidden_states": hidden_states} + if not self._is_static_pp_hidden_tensor_dict(tensor_dict, num_tokens): + raise RuntimeError( + "static PP hidden transfer input buffer violates its fixed schema" + ) + return tensor_dict + @torch.inference_mode() def execute_model( self, scheduler_output: "SchedulerOutput" @@ -835,12 +921,20 @@ def execute_model( } if forward_pass and not get_pp_group().is_first_rank: - tensor_dict, comm_handles, comm_postprocess = ( - get_pp_group().irecv_tensor_dict( - all_gather_group=get_tp_group(), - all_gather_tensors=all_gather_tensors, + tensor_dict = self._static_pp_hidden_recv_buffers(num_scheduled_tokens) + if tensor_dict is not None: + comm_handles = get_pp_group().irecv_tensor_dict_static(tensor_dict) + comm_postprocess: list[Callable[[], None]] = [] + logger.info_once( + "Default SM70 metadata-free PP hidden transfer enabled." + ) + else: + tensor_dict, comm_handles, comm_postprocess = ( + get_pp_group().irecv_tensor_dict( + all_gather_group=get_tp_group(), + all_gather_tensors=all_gather_tensors, + ) ) - ) assert tensor_dict is not None intermediate_tensors = AsyncIntermediateTensors( tensor_dict, @@ -870,12 +964,23 @@ def execute_model( and not get_pp_group().is_last_rank ) - # launch non-blocking send of intermediate tensors - self._pp_send_work = get_pp_group().isend_tensor_dict( - output.tensors, - all_gather_group=get_tp_group(), - all_gather_tensors=all_gather_tensors, - ) + # The exact SM70 B1 path sends the replicated tensor directly into + # the next stage's persistent graph-input buffer without CPU metadata. + if self._use_sm70_static_pp_hidden_transfer(num_scheduled_tokens): + if not self._is_static_pp_hidden_tensor_dict( + output.tensors, num_scheduled_tokens + ): + raise RuntimeError( + "static PP hidden transfer output violates its fixed schema" + ) + self._pp_send_work = get_pp_group().isend_tensor_dict_static(output.tensors) + logger.info_once("Default SM70 metadata-free PP hidden transfer enabled.") + else: + self._pp_send_work = get_pp_group().isend_tensor_dict( + output.tensors, + all_gather_group=get_tp_group(), + all_gather_tensors=all_gather_tensors, + ) return None diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 24f7536887..7ffb0c28f9 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -659,6 +659,63 @@ def run_fused_postprocess( HAS_DDTREE_ACCEPTED_NODES=has_ddtree, ) + def warmup_fused_postprocess(self) -> bool: + """Compile MTP postprocess signatures without copying model state.""" + if not self.is_initialized: + 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]) + warmup_sizes.add(max(1, min(max_num_reqs, 4))) + for num_reqs in sorted(warmup_sizes): + num_accepted_tokens = torch.ones(num_reqs, dtype=torch.int32, device=device) + spec_state_slot_selectors = torch.ones( + num_reqs, dtype=torch.int32, device=device + ) + mamba_state_idx = torch.zeros(num_reqs, dtype=torch.int32, device=device) + num_scheduled_tokens = torch.ones( + num_reqs, dtype=torch.int32, device=device + ) + # Put every request one token past a block boundary. The kernel + # exits before dereferencing state or block-table pointers. + num_computed_tokens = torch.full( + (num_reqs,), self.block_size, dtype=torch.int32, device=device + ) + num_draft_tokens = torch.zeros(num_reqs, dtype=torch.int32, device=device) + ddtree_accepted_node_indices = torch.empty( + (1, 1), dtype=torch.int32, device=device + ) + + postprocess_mamba_fused_kernel[(num_reqs, total_states)]( + num_accepted_tokens, + spec_state_slot_selectors, + ddtree_accepted_node_indices, + ddtree_accepted_node_indices.stride(0), + mamba_state_idx, + num_scheduled_tokens, + num_computed_tokens, + num_draft_tokens, + self.block_table_ptrs, + self.block_table_stride_req, + self.state_base_addrs, + self.state_block_strides, + self.state_elem_sizes, + self.state_inner_sizes, + self.state_conv_widths, + self.state_group_indices, + self.num_accepted_tokens_out, + self.spec_state_slot_selectors_out, + num_reqs, + block_size=self.block_size, + COPY_BLOCK_SIZE=1024, + HAS_DDTREE_ACCEPTED_NODES=False, + ) + torch.accelerator.synchronize() + return True + @dataclasses.dataclass class MambaBuffers: diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index fb02fdc935..6b634557a5 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -37,9 +37,48 @@ logger = init_logger(__name__) +def _infer_segment_block_strides( + seg_addrs: list[int], + page_size_el: int, +) -> list[int]: + """Infer logical block strides from block-major segment address runs. + + Dense segments advance by one page. An interleaved pool exposes one segment + base per cell, exactly ``page_size_el * 4`` bytes apart, and advances to the + next logical block after every segment in that address run. Treat separate + runs independently so multiple pools do not disable one another. + """ + if page_size_el <= 0: + raise ValueError(f"page_size_el must be positive, got {page_size_el}.") + if len(set(seg_addrs)) != len(seg_addrs): + raise ValueError("Segment addresses must be unique.") + + block_strides = [page_size_el] * len(seg_addrs) + if len(seg_addrs) < 2: + return block_strides + + cell_bytes = page_size_el * 4 + ordered = sorted((addr, index) for index, addr in enumerate(seg_addrs)) + run_start = 0 + for run_end in range(1, len(ordered) + 1): + if ( + run_end < len(ordered) + and ordered[run_end][0] - ordered[run_end - 1][0] == cell_bytes + ): + continue + run = ordered[run_start:run_end] + if len(run) > 1: + run_stride = page_size_el * len(run) + for _, original_index in run: + block_strides[original_index] = run_stride + run_start = run_end + return block_strides + + @triton.jit(do_not_specialize=["n_blocks"]) def _zero_kv_blocks_kernel( seg_addrs_ptr, + seg_block_strides_ptr, block_ids_ptr, n_blocks, N_SEGS: tl.constexpr, @@ -55,6 +94,9 @@ def _zero_kv_blocks_kernel( seg_addrs_ptr holds absolute byte addresses (int64) for each segment, allowing segments to live in different CUDA allocations. + seg_block_strides_ptr holds the per-segment block-to-block stride in + 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). """ @@ -69,9 +111,11 @@ def _zero_kv_blocks_kernel( 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) * PAGE_SIZE_EL + chunk_index.to(tl.int64) * BLOCK_SIZE + block_id.to(tl.int64) * block_stride_el.to(tl.int64) + + chunk_index.to(tl.int64) * BLOCK_SIZE ) cols = tl.arange(0, BLOCK_SIZE).to(tl.int64) tl.store(ptr + offset + cols, tl.zeros([BLOCK_SIZE], dtype=tl.int32)) @@ -88,7 +132,7 @@ class KVBlockZeroer: def __init__(self, device: torch.device, pin_memory: bool): self.device = device self.pin_memory = pin_memory - self._meta: tuple[torch.Tensor, int, int, int] | None = None + self._meta: tuple[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 @@ -155,8 +199,7 @@ def init_meta( page_size_el = cur_page_el else: assert page_size_el == cur_page_el, ( - f"Non-uniform page sizes: {page_size_el} vs " - f"{cur_page_el}" + f"Non-uniform page sizes: {page_size_el} vs {cur_page_el}" ) block_stride_bytes = cur_bytes @@ -187,8 +230,7 @@ def init_meta( page_size_el = cur_page_el else: assert page_size_el == cur_page_el, ( - f"Non-uniform page sizes: {page_size_el} vs " - f"{cur_page_el}" + f"Non-uniform page sizes: {page_size_el} vs {cur_page_el}" ) seg_addrs.append(dp) @@ -197,6 +239,18 @@ def init_meta( return blk_size = min(largest_power_of_2_divisor(page_size_el), 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) + seg_block_strides = torch.tensor( + _infer_segment_block_strides(seg_addrs, page_size_el), + dtype=torch.int64, + device=self.device, + ) + self._id_cap = 8192 self._ids_pinned = torch.empty( self._id_cap, @@ -206,16 +260,17 @@ def init_meta( self._ids_gpu = torch.empty(self._id_cap, dtype=torch.int64, device=self.device) self._meta = ( torch.tensor(seg_addrs, dtype=torch.uint64, device=self.device), + seg_block_strides, page_size_el, blk_size, - len(seg_addrs), + n_segs, ) 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, page_size_el, blk_size, n_segs = self._meta + seg_addrs, seg_block_strides, page_size_el, blk_size, n_segs = self._meta n_blocks = len(block_ids) if n_blocks > self._id_cap: self._id_cap = n_blocks * 2 @@ -234,6 +289,7 @@ def zero_block_ids(self, block_ids: list[int]) -> None: grid = (n_blocks * n_segs * (page_size_el // blk_size),) _zero_kv_blocks_kernel[grid]( seg_addrs, + seg_block_strides, idx, n_blocks, N_SEGS=n_segs, From 854a4e07b7cf7e04b1e9847b60ceeea4af478ad4 Mon Sep 17 00:00:00 2001 From: Leonccaa <166551845+Leonccaa@users.noreply.github.com> Date: Thu, 27 Aug 2026 07:28:43 -0700 Subject: [PATCH 079/105] perf(kv-cache): avoid singleton hybrid group sizing Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com> --- tests/v1/core/test_kv_cache_utils.py | 20 ++++++++++ vllm/v1/core/kv_cache_utils.py | 56 ++++++++++++++++++++++------ 2 files changed, 64 insertions(+), 12 deletions(-) diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 1ec844fead..2a0b8906a0 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -214,6 +214,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/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 9f1f447340..25e476b042 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -1069,6 +1069,42 @@ 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 _get_kv_cache_groups_uniform_page_size( kv_cache_spec: dict[str, KVCacheSpec], ) -> list[KVCacheGroupSpec]: @@ -1153,18 +1189,14 @@ 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) + 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 From cf5c44a1aa65bd8556ff9050bf9cb19d1c5f439d Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 22:48:32 +0800 Subject: [PATCH 080/105] [Bugfix] Load fused Qwen4Exp MTP experts Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/models/qwen4_exp/test_weight_loading.py | 71 +++++++++++++++++++ vllm/model_executor/layers/fused_moe/layer.py | 47 +++++++++++- vllm/models/qwen4_exp/nvidia/model.py | 1 + 3 files changed, 118 insertions(+), 1 deletion(-) diff --git a/tests/models/qwen4_exp/test_weight_loading.py b/tests/models/qwen4_exp/test_weight_loading.py index fb5d9903c7..09a975f196 100644 --- a/tests/models/qwen4_exp/test_weight_loading.py +++ b/tests/models/qwen4_exp/test_weight_loading.py @@ -2,8 +2,10 @@ # 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, @@ -144,6 +146,9 @@ def fake_qwen3_next_init(self, vllm_config, prefix="") -> None: 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"), @@ -153,6 +158,72 @@ def fake_qwen3_next_init(self, vllm_config, prefix="") -> None: ] +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]) + + @pytest.mark.parametrize( ("checkpoint_name", "model_name"), [ 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/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py index 7eddc1bb35..06ce82e8cf 100644 --- a/vllm/models/qwen4_exp/nvidia/model.py +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -208,6 +208,7 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: ckpt_up_proj_name="up_proj", num_experts=self.n_routed_experts, num_redundant_experts=self.n_redundant_experts, + include_fused=True, ) From ceb543c05573c9bcbf1cc9563bab0c2e8c746ae3 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 22:56:37 +0800 Subject: [PATCH 081/105] [Bugfix] Reject incomplete Qwen4Exp MTP loads Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/models/qwen4_exp/test_weight_loading.py | 19 ++++++++++++++ vllm/models/qwen4_exp/nvidia/mtp.py | 26 ++++++++++++++++++- 2 files changed, 44 insertions(+), 1 deletion(-) diff --git a/tests/models/qwen4_exp/test_weight_loading.py b/tests/models/qwen4_exp/test_weight_loading.py index 09a975f196..ce95ac3441 100644 --- a/tests/models/qwen4_exp/test_weight_loading.py +++ b/tests/models/qwen4_exp/test_weight_loading.py @@ -13,6 +13,7 @@ Qwen4ExpSparseMoeBlock, _remap_qsa_cache_scale_name, ) +from vllm.models.qwen4_exp.nvidia.mtp import _validate_mtp_expert_weights_loaded @pytest.mark.parametrize( @@ -224,6 +225,24 @@ def weight_loader( 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"), [ diff --git a/vllm/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py index 3945ebc663..1b1df75d37 100644 --- a/vllm/models/qwen4_exp/nvidia/mtp.py +++ b/vllm/models/qwen4_exp/nvidia/mtp.py @@ -116,6 +116,28 @@ def _remap_mtp_weight_name(name: str) -> str | None: 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, @@ -459,7 +481,9 @@ def remap_weight_names(): skip_substrs=["hyper_connection_mixer.block_inject_weight"], ignore_unexpected_suffixes=_QWEN4_EXP_IGNORED_MISSING_SUFFIXES.copy(), ) - return loader.load_weights(remap_weight_names()) + loaded_weights = loader.load_weights(remap_weight_names()) + _validate_mtp_expert_weights_loaded(self, loaded_weights) + return loaded_weights __all__ = ["Qwen4ExpMTP", "Qwen4ExpMultiTokenPredictor"] From 06724d0fd4c021a1fc33c50269d894d7c9f2a2bd Mon Sep 17 00:00:00 2001 From: 1CatTCat <117812519+TCat2003@users.noreply.github.com> Date: Thu, 27 Aug 2026 23:12:52 +0800 Subject: [PATCH 082/105] Update README.md --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 From e9aef4433477d9fb549f1c73a231712bb10a35b0 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 23:21:36 +0800 Subject: [PATCH 083/105] [Core][SM70] Start DFlash2 lookup drafting campaign Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> From 5bf433c96c5c5ec2ec669855ebee1e076370a712 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 23:29:11 +0800 Subject: [PATCH 084/105] [Core] Account for DFlash2 hybrid cache capacity Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/v1/core/test_kv_cache_utils.py | 30 +++++++++++++ .../test_gpu_cudagraph_memory_reserve.py | 29 ++++++++++++ vllm/envs.py | 7 +++ vllm/v1/core/kv_cache_utils.py | 44 +++++++++++++++++++ vllm/v1/worker/gpu/cudagraph_utils.py | 17 +++++++ vllm/v1/worker/gpu/model_runner.py | 5 ++- 6 files changed, 131 insertions(+), 1 deletion(-) create mode 100644 tests/v1/worker/test_gpu_cudagraph_memory_reserve.py diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 1ec844fead..7b2ec25972 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_heuristic_rejects_excessive_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(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) == 61 + assert max(len(group.layer_names) for group in groups) == 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( 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/vllm/envs.py b/vllm/envs.py index 887121a739..d77a1bb7c7 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -4188,6 +4188,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/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 9f1f447340..bcc573cd37 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -1069,6 +1069,48 @@ def is_kv_cache_type_attention_free(kv_cache_spec: dict[str, KVCacheSpec]) -> bo return not kv_cache_spec +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 ( + 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]: @@ -1165,6 +1207,8 @@ def _get_kv_cache_groups_uniform_page_size( # layers while accommodating speculative decoding drafters that add # extra layers to one attention type. group_size = max_num_layers + else: + group_size = _prefer_padding_sliding_window_layers(same_type_layers, group_size) grouped_layers = [] for layers in same_type_layers.values(): num_padding_layers = group_size - len(layers) % group_size diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index a0953c44c5..24aad335c5 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -34,6 +34,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( diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 47d1da107e..faaeb8dae2 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -66,6 +66,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 @@ -670,7 +671,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: From 8800477f5196017e4223371b727cec9c404bf219 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 23:39:14 +0800 Subject: [PATCH 085/105] [Bugfix][SM70] Keep bounded FP13 path enabled Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_v100_migration_control.md | 24 +++++++++++-------- .../test_sm70_deepseek_v4_fp16_gemv.py | 8 +++---- vllm/envs.py | 4 ++-- 3 files changed, 20 insertions(+), 16 deletions(-) diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index b960d02e7a..d28877d5fc 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -43740,14 +43740,16 @@ Interpretation: 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 now defaults off because these operator and overlap screens do not - constitute a matched model-level dataset gate. Explicit - `VLLM_SM70_DSV4_FP13_GEMV=1` remains available for an isolated future A/B. +- 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 would add - 549.66 MiB without pipeline parallelism, or at most 282.34 MiB per rank for - the audited PP2 partition. + 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 @@ -43760,10 +43762,12 @@ Interpretation: 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 current quality-safe source-`59adf8d89a` control explicitly disables all - FP8 QPN8 routes and packed FP13 while retaining 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 +- 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 diff --git a/tests/kernels/test_sm70_deepseek_v4_fp16_gemv.py b/tests/kernels/test_sm70_deepseek_v4_fp16_gemv.py index ae42aad9f3..f3ed2c5669 100644 --- a/tests/kernels/test_sm70_deepseek_v4_fp16_gemv.py +++ b/tests/kernels/test_sm70_deepseek_v4_fp16_gemv.py @@ -23,15 +23,15 @@ def reset_env_cache() -> Iterator[None]: envs.disable_envs_cache() -def test_sm70_dsv4_fp13_gemv_is_default_off_with_opt_in(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 not envs.VLLM_SM70_DSV4_FP13_GEMV + assert envs.VLLM_SM70_DSV4_FP13_GEMV - monkeypatch.setenv("VLLM_SM70_DSV4_FP13_GEMV", "1") + monkeypatch.setenv("VLLM_SM70_DSV4_FP13_GEMV", "0") envs.disable_envs_cache() - assert envs.VLLM_SM70_DSV4_FP13_GEMV + assert not envs.VLLM_SM70_DSV4_FP13_GEMV def test_sm70_dsv4_fp13_enables_fp16_fallback(monkeypatch) -> None: diff --git a/vllm/envs.py b/vllm/envs.py index 43a506020e..e779d5e3ca 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -185,7 +185,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 = 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 @@ -1829,7 +1829,7 @@ def _resolve_rust_frontend_path() -> str | None: int(os.getenv("VLLM_SM70_DSV4_FP16_GEMV", "0")) ), "VLLM_SM70_DSV4_FP13_GEMV": lambda: bool( - int(os.getenv("VLLM_SM70_DSV4_FP13_GEMV", "0")) + 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")) From 1889515028bc7cdecccdd0a86f0409e2ef632b88 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 23:44:15 +0800 Subject: [PATCH 086/105] [Kernel][SM70] Add exact q16 DFlash2 verification Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../benchmark_sm70_dflash2_grouped_verify.py | 163 +++++++++++++++ .../flash_attn_v100/flash_attn_interface.py | 17 +- .../kernel/flash_decode_paged.cu | 194 +++++++++++------- .../test_sm70_flash_v100_grouped_verify.py | 15 +- .../attention/test_sm70_flash_v100_policy.py | 19 +- vllm/v1/attention/backends/flash_attn_v100.py | 7 +- 6 files changed, 321 insertions(+), 94 deletions(-) create mode 100644 benchmarks/kernels/benchmark_sm70_dflash2_grouped_verify.py 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/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/kernel/flash_decode_paged.cu b/flash-attention-v100/kernel/flash_decode_paged.cu index 7d704a9aaa..95ef338821 100644 --- a/flash-attention-v100/kernel/flash_decode_paged.cu +++ b/flash-attention-v100/kernel/flash_decode_paged.cu @@ -1777,16 +1777,16 @@ __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; @@ -1796,7 +1796,9 @@ 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 @@ -1814,6 +1816,25 @@ 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 { @@ -1835,14 +1856,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; @@ -1905,7 +1927,8 @@ __device__ __forceinline__ void grouped_verify_scale_output_fragment( fragment.x[7] *= second_scale; } -template __global__ @@ -1919,11 +1942,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; } @@ -1932,7 +1956,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; } @@ -1947,7 +1971,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; @@ -1990,8 +2014,8 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ 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[row * kSharedQVecsPerRow + vec_col] = __ldg( @@ -2049,8 +2073,8 @@ __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 = @@ -2109,8 +2133,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 = @@ -2131,8 +2155,8 @@ __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 = @@ -2233,14 +2257,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 + @@ -2249,13 +2273,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] = @@ -2264,26 +2288,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] = @@ -2315,7 +2340,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 + @@ -3518,10 +3543,17 @@ 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 && @@ -3540,13 +3572,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]"); + "partial_out must have shape [80, 8, 6, 256] or " + "[40, 16, 6, 256]"); TORCH_CHECK(partial_lse.sizes() == - at::IntArrayRef({kGroupedVerifySplits, kGroupedVerifyMaxQ, + at::IntArrayRef({grouped_splits, max_query_tokens, kGroupedVerifyHeads}), - "partial_lse must have shape [80, 8, 6]"); + "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"); @@ -3569,15 +3602,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, \ - CONTIGUOUS_LAYOUT, STAGE_PAGE_IDS) \ +#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, CONTIGUOUS_LAYOUT, \ - STAGE_PAGE_IDS>; \ + 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); \ @@ -3590,7 +3624,8 @@ at::Tensor flash_attention_grouped_verify_paged( "Failed to set packed grouped-verifier shared carveout: ", \ cudaGetErrorString(carveout_status)); \ flash_attention_grouped_verify_e5m2_partial_kernel< \ - TWO_PASS, PAGE_SIZE, SINGLE_QUERY, CONTIGUOUS_LAYOUT, STAGE_PAGE_IDS> \ + 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(), \ @@ -3604,7 +3639,8 @@ at::Tensor flash_attention_grouped_verify_paged( v_scale); \ } while (0) -#define DISPATCH_GROUPED_VERIFY_PARTIAL(TWO_PASS, 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; \ @@ -3620,51 +3656,61 @@ at::Tensor flash_attention_grouped_verify_paged( 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(TWO_PASS, 1648, SINGLE_QUERY, true, true); \ + LAUNCH_GROUPED_VERIFY_PARTIAL(MAX_QUERY_TOKENS, TWO_PASS, 1648, \ + SINGLE_QUERY, true, true); \ } else if (stage_page_ids) { \ - LAUNCH_GROUPED_VERIFY_PARTIAL(TWO_PASS, 3296, SINGLE_QUERY, true, true); \ + 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(TWO_PASS, 1648, SINGLE_QUERY, true, \ - false); \ + LAUNCH_GROUPED_VERIFY_PARTIAL(MAX_QUERY_TOKENS, TWO_PASS, 1648, \ + SINGLE_QUERY, true, false); \ } else if (fixed_interleaved_layout) { \ - LAUNCH_GROUPED_VERIFY_PARTIAL(TWO_PASS, 3296, SINGLE_QUERY, true, \ - false); \ + LAUNCH_GROUPED_VERIFY_PARTIAL(MAX_QUERY_TOKENS, TWO_PASS, 3296, \ + SINGLE_QUERY, true, false); \ } else if (page_size == 1648) { \ - LAUNCH_GROUPED_VERIFY_PARTIAL(TWO_PASS, 1648, SINGLE_QUERY, false, \ - false); \ + LAUNCH_GROUPED_VERIFY_PARTIAL(MAX_QUERY_TOKENS, TWO_PASS, 1648, \ + SINGLE_QUERY, false, false); \ } else if (page_size == 3296) { \ - LAUNCH_GROUPED_VERIFY_PARTIAL(TWO_PASS, 3296, SINGLE_QUERY, false, \ - false); \ + LAUNCH_GROUPED_VERIFY_PARTIAL(MAX_QUERY_TOKENS, TWO_PASS, 3296, \ + SINGLE_QUERY, false, false); \ } else { \ - LAUNCH_GROUPED_VERIFY_PARTIAL(TWO_PASS, 0, SINGLE_QUERY, false, false); \ + 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/tests/kernels/attention/test_sm70_flash_v100_grouped_verify.py b/tests/kernels/attention/test_sm70_flash_v100_grouped_verify.py index c7fab4fdb3..8000984b07 100644 --- a/tests/kernels/attention/test_sm70_flash_v100_grouped_verify.py +++ b/tests/kernels/attention/test_sm70_flash_v100_grouped_verify.py @@ -127,10 +127,13 @@ 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), (3296, 8, 8193, True), + (3296, 16, 8185, True), ], ) @torch.inference_mode() @@ -217,13 +220,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) @@ -269,15 +275,16 @@ def test_grouped_verify_cuda_graph_replay_tracks_runtime_seq_len() -> None: @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, monkeypatch: pytest.MonkeyPatch + 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=8, + query_len=query_len, prefix_len=8193, ) diff --git a/tests/v1/attention/test_sm70_flash_v100_policy.py b/tests/v1/attention/test_sm70_flash_v100_policy.py index 419c4d847c..d8deb3913e 100644 --- a/tests/v1/attention/test_sm70_flash_v100_policy.py +++ b/tests/v1/attention/test_sm70_flash_v100_policy.py @@ -1665,7 +1665,10 @@ 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]) +def test_flash_v100_dflash2_grouped_verify_uses_original_request_metadata( + query_len: int, +): from vllm.v1.attention.backends.flash_attn_v100 import FlashAttnV100Impl impl = FlashAttnV100Impl( @@ -1700,19 +1703,19 @@ 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) value_cache = torch.zeros_like(key_cache) @@ -1723,7 +1726,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/vllm/v1/attention/backends/flash_attn_v100.py b/vllm/v1/attention/backends/flash_attn_v100.py index e42b3df8ae..c9c9f25979 100644 --- a/vllm/v1/attention/backends/flash_attn_v100.py +++ b/vllm/v1/attention/backends/flash_attn_v100.py @@ -5150,8 +5150,8 @@ 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 @@ -5226,7 +5226,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( From 9d62e5db9865d7bd94ef03d7c2d4bbd8e39cdea3 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 28 Aug 2026 00:00:14 +0800 Subject: [PATCH 087/105] [Perf][SM70] Accelerate Qwen3.8 QSA indexing Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_v100_migration_control.md | 34 ++ tests/models/qwen4_exp/test_qsa_ops.py | 49 +++ vllm/models/qwen4_exp/nvidia/ops/qsa.py | 345 ++++++++++++++++++++- 3 files changed, 416 insertions(+), 12 deletions(-) diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index 1425b153b6..3d4b0e96b2 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -68,6 +68,40 @@ Goal: 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 diff --git a/tests/models/qwen4_exp/test_qsa_ops.py b/tests/models/qwen4_exp/test_qsa_ops.py index 11f4de01c7..ebc3f3273d 100644 --- a/tests/models/qwen4_exp/test_qsa_ops.py +++ b/tests/models/qwen4_exp/test_qsa_ops.py @@ -1,7 +1,11 @@ # 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, ) @@ -15,3 +19,48 @@ def test_sm70_qsa_prefill_uses_narrow_tiles_and_four_warps(): 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/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py index 5cd2f60d3a..c3678eba85 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py @@ -5,14 +5,28 @@ 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 @@ -118,6 +132,124 @@ def _qsa_mqa_paged_kernel( ) +@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, @@ -682,6 +814,144 @@ def qsa_mqa_paged( 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, @@ -762,9 +1032,44 @@ def qsa_select_paged_tokens( if not rows: return out - columns = page_table.shape[1] * k_cache.shape[1] + capacity_columns = page_table.shape[1] * k_cache.shape[1] + score_columns = capacity_columns block_topk = token_topk // compress_ratio - rows_per_chunk = max(1, _LOGITS_WORKSPACE_BYTES // max(columns * 4, 1)) + 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 @@ -775,15 +1080,24 @@ def qsa_select_paged_tokens( 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) - 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, - ) + 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 @@ -796,7 +1110,14 @@ def qsa_select_paged_tokens( if use_cooperative_topk else torch.ops._C.persistent_topk ) - topk_op(logits, visible_blocks, blocks, topk_workspace, block_topk, columns) + topk_op( + logits, + visible_blocks, + blocks, + topk_workspace, + block_topk, + score_columns, + ) expand_qsa_block_indices_cuda( blocks, query_positions[row_slice], From be44103c65928d574701fdee452ed9fac084e07c Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 28 Aug 2026 00:10:40 +0800 Subject: [PATCH 088/105] [CI] Use approved regex import in SM70 benchmarks Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- benchmarks/benchmark_sm70_tool_datasets.py | 2 +- benchmarks/benchmark_sm70_tool_protocol.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/benchmarks/benchmark_sm70_tool_datasets.py b/benchmarks/benchmark_sm70_tool_datasets.py index 095cfb2702..b85077ba97 100644 --- a/benchmarks/benchmark_sm70_tool_datasets.py +++ b/benchmarks/benchmark_sm70_tool_datasets.py @@ -14,12 +14,12 @@ import ast import hashlib import json -import re 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:" diff --git a/benchmarks/benchmark_sm70_tool_protocol.py b/benchmarks/benchmark_sm70_tool_protocol.py index 83bfb648a2..b7cccb1425 100644 --- a/benchmarks/benchmark_sm70_tool_protocol.py +++ b/benchmarks/benchmark_sm70_tool_protocol.py @@ -13,12 +13,12 @@ import argparse import copy import json -import re import time from pathlib import Path from typing import Any import jsonschema +import regex as re import requests TOOLS: list[dict[str, Any]] = [ From 95fdf660b325ef1627a01b414518d1c45f7e67d0 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 28 Aug 2026 00:38:05 +0800 Subject: [PATCH 089/105] [Bugfix][SM70] Repair Qwen4Exp integration gates Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/models/qwen4_exp/test_ple.py | 2 +- tests/models/qwen4_exp/test_qsa_cache.py | 5 +++-- .../test_sm70_modelopt_mixed_nvfp4.py | 4 ++-- tests/v1/worker/test_model_state_init.py | 2 +- tests/v1/worker/test_qwen4_exp_ngram.py | 3 ++- tests/v1/worker/test_qwen4_exp_v2.py | 2 +- vllm/model_executor/layers/linear.py | 12 ++++++++---- vllm/models/qwen4_exp/nvidia/model.py | 13 +++++++------ vllm/models/qwen4_exp/nvidia/mtp.py | 4 ++-- vllm/models/qwen4_exp/nvidia/ple_layer.py | 19 +++++++++++++------ vllm/models/qwen4_exp/nvidia/qsa.py | 7 +++++-- vllm/v1/kv_cache_interface.py | 2 +- .../gpu/spec_decode/eagle/speculator.py | 7 ++++--- vllm/v1/worker/gpu_model_runner.py | 4 +--- 14 files changed, 51 insertions(+), 35 deletions(-) diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index 102d1045da..6777828c76 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -101,7 +101,7 @@ def test_pinned_host_ple_fp8_rows_are_gatherable_on_sm70( layer.prepare_accelerator_weight() output = layer(torch.tensor([0, 7], dtype=torch.int64, device="cuda")) - torch.cuda.synchronize() + torch.accelerator.synchronize() assert output.dtype == torch.float16 expected = raw.view(torch.float8_e4m3fn).float() * 0.25 diff --git a/tests/models/qwen4_exp/test_qsa_cache.py b/tests/models/qwen4_exp/test_qsa_cache.py index 9a5ed1daca..d3cfea8d5f 100644 --- a/tests/models/qwen4_exp/test_qsa_cache.py +++ b/tests/models/qwen4_exp/test_qsa_cache.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from types import SimpleNamespace +from typing import Any import torch @@ -12,7 +13,7 @@ def test_bind_qsa_key_cache_builds_key_and_mrope_views() -> None: prefix = "model.layers.0.self_attn.raw_key_cache" - static_forward_context = {} + static_forward_context: dict[str, Any] = {} layer = QSAKeyStateCache( head_size=128, dtype=torch.float16, @@ -27,7 +28,7 @@ def test_bind_qsa_key_cache_builds_key_and_mrope_views() -> None: ), ) cache = torch.empty(2, 8, 1, layer.head_size, dtype=torch.float16) - runner_kv_caches = [] + runner_kv_caches: list[torch.Tensor] = [] bind_kv_cache({prefix: cache}, static_forward_context, runner_kv_caches) diff --git a/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py b/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py index 469e8aadb9..da2cd7dc1f 100644 --- a/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py +++ b/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py @@ -222,7 +222,7 @@ def test_nvfp4_single_token_direct_routing_is_exact_and_graph_dynamic(): torch.ops._moe_C.moe_unpermute( expert_output, weights, identity, None, top_k, reference ) - torch.cuda.synchronize() + 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) @@ -240,7 +240,7 @@ def test_nvfp4_single_token_direct_routing_is_exact_and_graph_dynamic(): torch.ops._moe_C.moe_unpermute( expert_output, weights, identity, None, top_k, reference ) - torch.cuda.synchronize() + 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) diff --git a/tests/v1/worker/test_model_state_init.py b/tests/v1/worker/test_model_state_init.py index 7535a70b6c..44200b5806 100644 --- a/tests/v1/worker/test_model_state_init.py +++ b/tests/v1/worker/test_model_state_init.py @@ -10,7 +10,7 @@ def test_model_can_select_custom_model_state() -> None: - captured = {} + captured: dict[str, object] = {} class CustomModelState: def __init__(self, vllm_config, model, encoder_cache, device) -> None: diff --git a/tests/v1/worker/test_qwen4_exp_ngram.py b/tests/v1/worker/test_qwen4_exp_ngram.py index c3904e1b65..6c2459c897 100644 --- a/tests/v1/worker/test_qwen4_exp_ngram.py +++ b/tests/v1/worker/test_qwen4_exp_ngram.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from types import SimpleNamespace +from typing import Any import numpy as np import torch @@ -58,7 +59,7 @@ def test_v1_runner_prepares_committed_ngram_context() -> None: def test_v1_runner_uses_stable_dummy_ngram_buffers() -> None: runner = _runner() - model_kwargs = {} + model_kwargs: dict[str, Any] = {} runner._maybe_add_ngram_kwargs( model_kwargs, diff --git a/tests/v1/worker/test_qwen4_exp_v2.py b/tests/v1/worker/test_qwen4_exp_v2.py index 1e6d9d5af2..954779157f 100644 --- a/tests/v1/worker/test_qwen4_exp_v2.py +++ b/tests/v1/worker/test_qwen4_exp_v2.py @@ -219,7 +219,7 @@ def test_qsa_circular_group_emits_no_generic_slots_on_sm70() -> None: positions=torch.tensor([153797, 165757], dtype=torch.int64, device=device), num_tokens_padded=2, ) - torch.cuda.synchronize() + torch.accelerator.synchronize() assert slot_mappings[0].tolist() == [-1, -1] assert slot_mappings[1].tolist() == [ diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index 1d1f726b45..7f549fed61 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -1224,11 +1224,13 @@ def load_weights( self.validate_shard_id(shard_id) if "." in name: submodule, _, attr = name.rpartition(".") - param = getattr(self.get_submodule(submodule), attr, self) + param = getattr(self.get_submodule(submodule), attr, None) else: - param = getattr(self, name, self) + 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", @@ -1346,11 +1348,13 @@ def load_weights( self.validate_shard_id(shard_id) if "." in name: submodule, _, attr = name.rpartition(".") - param = getattr(self.get_submodule(submodule), attr, self) + param = getattr(self.get_submodule(submodule), attr, None) else: - param = getattr(self, name, self) + 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", diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py index 06ce82e8cf..526acbe0a4 100644 --- a/vllm/models/qwen4_exp/nvidia/model.py +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -375,6 +375,7 @@ def forward( attn_out = self.self_attn( hidden_states=block_input, positions=positions, + output=None, ) else: raise ValueError("Invalid layer_type") @@ -691,7 +692,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: prefix=maybe_prefix(prefix, "lm_head"), ) self.logits_processor = LogitsProcessor(config.vocab_size) - self.make_empty_intermediate_tensors = ( + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] self.model.make_empty_intermediate_tensors ) self.set_moe_parameters(self.model.layers) @@ -769,7 +770,7 @@ def get_gdn_mamba_state_dtype_from_config( @classmethod def get_gdn_mamba_state_shape_from_config( cls, vllm_config: VllmConfig - ) -> tuple[tuple[int, int], tuple[int, int]]: + ) -> 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 @@ -798,7 +799,7 @@ def get_mamba_state_dtype_from_config( @classmethod def get_mamba_state_shape_from_config( cls, vllm_config: VllmConfig - ) -> tuple[tuple[int, int], tuple[int, int]]: + ) -> tuple[tuple[int, int], tuple[int, int, int]]: return cls.get_gdn_mamba_state_shape_from_config(vllm_config) @classmethod @@ -967,7 +968,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model") -> None: prefix=maybe_prefix(prefix, "language_model"), ) - self.make_empty_intermediate_tensors = ( + 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: @@ -1069,10 +1070,10 @@ def get_mamba_state_dtype_from_config( return Qwen4ExpForCausalLM.get_mamba_state_dtype_from_config(vllm_config) @classmethod - def get_mamba_state_shape_from_config( + def get_mamba_state_shape_from_config( # type: ignore[override] cls, vllm_config: VllmConfig, - ) -> tuple[tuple[int, int], tuple[int, int]]: + ) -> tuple[tuple[int, int], tuple[int, int, int]]: return Qwen4ExpForCausalLM.get_mamba_state_shape_from_config(vllm_config) @classmethod diff --git a/vllm/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py index 1b1df75d37..18f4575f54 100644 --- a/vllm/models/qwen4_exp/nvidia/mtp.py +++ b/vllm/models/qwen4_exp/nvidia/mtp.py @@ -432,7 +432,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: self.lm_head = PPMissingLayer() self.logits_processor = LogitsProcessor(config.vocab_size) - self.make_empty_intermediate_tensors = ( + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] self.model.make_empty_intermediate_tensors ) self.set_moe_parameters(self.model.layers) @@ -441,7 +441,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model.embed_input_ids(input_ids) - def forward( + def forward( # type: ignore[override] self, input_ids: torch.Tensor | None, positions: torch.Tensor, diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py index 1ccad20a4b..b4902d6c46 100644 --- a/vllm/models/qwen4_exp/nvidia/ple_layer.py +++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py @@ -334,10 +334,13 @@ def __init__( 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(self.weight.shape), - dtype=self.weight.dtype, + tuple(meta_weight.shape), + dtype=meta_weight.dtype, device="cpu", pin_memory=True, ), @@ -369,7 +372,9 @@ def get_accelerator_weight(self, device: torch.device) -> torch.Tensor: f"Qwen4Exp pinned-host PLE requires a CUDA input, got {device}" ) device_index = ( - torch.cuda.current_device() if device.index is None else 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: @@ -377,20 +382,22 @@ def get_accelerator_weight(self, device: torch.device) -> torch.Tensor: raise RuntimeError( "Qwen4Exp PLE UVA view must be prepared before CUDA graph capture" ) - with torch.cuda.device(device_index): + 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.cuda.current_device())) + 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.cuda.current_device() + torch.accelerator.current_device_index() if input_.device.index is None else input_.device.index ) diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py index b6b68b45cb..25e9ab967d 100644 --- a/vllm/models/qwen4_exp/nvidia/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/qsa.py @@ -429,6 +429,7 @@ def _run_qsa( def forward( self, positions: torch.Tensor, + output: torch.Tensor | None, hidden_states: torch.Tensor, ) -> torch.Tensor: qkv, _ = self.qkv_proj(hidden_states) @@ -462,8 +463,10 @@ def forward( flat_output = attn_output.view(num_tokens, -1) if gate is not None: flat_output = flat_output * torch.sigmoid(gate) - output, _ = self.o_proj(flat_output) - return output + 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( diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index d6ae8e0fc2..10edbf4eb5 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -634,7 +634,7 @@ def merge(cls, specs: list[Self]) -> Self: @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" diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py index 07afbd9c04..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 diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 8ada303e4d..237861e3d7 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -1557,9 +1557,7 @@ def __init__( 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 - ) + 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 From 5aad3a09f13d4315f16139626a8f551f8bf5d428 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 28 Aug 2026 00:59:54 +0800 Subject: [PATCH 090/105] [Kernel] Accelerate Qwen3.8 Flash Next decode on SM70 Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- benchmarks/benchmark_sm70_decode.py | 32 +- csrc/custom_all_reduce.cuh | 6 +- csrc/ops.h | 12 + csrc/sm70_turbomind/ops/fp8_qpn8_sm70.cu | 576 +++++++++++++++++- csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu | 197 ++++++ csrc/torch_bindings.cpp | 16 + docs/design/sm70_qwen38_flash_next_nvfp4.md | 280 ++++++++- tests/kernels/moe/test_fused_topk.py | 54 +- .../test_qwen3_5_quantization.py | 9 +- .../test_sm70_modelopt_mixed_nvfp4.py | 27 + tests/quantization/test_sm70_online_qpn8.py | 171 ++++++ .../worker/test_gpu_model_runner_v2_greedy.py | 134 ++++ vllm/_custom_ops.py | 72 ++- vllm/_sm70_ops.py | 191 +++++- .../device_communicators/custom_all_reduce.py | 3 +- vllm/envs.py | 33 + .../fused_moe/router/fused_topk_router.py | 101 +++ vllm/model_executor/layers/linear.py | 14 + .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 18 +- .../layers/quantization/nvfp4_sm70_moe.py | 59 ++ .../layers/quantization/sm70_online_qpn8.py | 287 +++++++++ vllm/model_executor/models/qwen2_moe.py | 32 +- .../qwen4_exp/nvidia/hyperconnection.py | 50 +- vllm/models/qwen4_exp/nvidia/model.py | 6 + vllm/models/qwen4_exp/nvidia/ops/hc.py | 6 +- vllm/models/qwen4_exp/nvidia/ops/qsa.py | 16 +- vllm/v1/worker/gpu/model_runner.py | 27 +- vllm/v1/worker/gpu/sample/sampler.py | 18 + 28 files changed, 2333 insertions(+), 114 deletions(-) create mode 100644 tests/quantization/test_sm70_online_qpn8.py create mode 100644 tests/v1/worker/test_gpu_model_runner_v2_greedy.py create mode 100644 vllm/model_executor/layers/quantization/sm70_online_qpn8.py diff --git a/benchmarks/benchmark_sm70_decode.py b/benchmarks/benchmark_sm70_decode.py index f10eac8d3d..93f4596100 100644 --- a/benchmarks/benchmark_sm70_decode.py +++ b/benchmarks/benchmark_sm70_decode.py @@ -1195,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 @@ -1226,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" ), @@ -1272,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. " @@ -1290,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, @@ -1302,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" ), @@ -1342,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 " @@ -1349,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/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/ops.h b/csrc/ops.h index 67170c205d..d1e6250b33 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); @@ -489,6 +496,11 @@ 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/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 index 5d65a8267f..cf4458e165 100644 --- a/csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu +++ b/csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu @@ -138,6 +138,108 @@ __global__ void mxfp4_qpn_m1_sm70_kernel(const half* __restrict__ input, } } +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, @@ -153,6 +255,50 @@ void launch_mxfp4_qpn_m1(torch::Tensor out, torch::Tensor input, 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, @@ -201,3 +347,54 @@ void mxfp4_moe_qpn_m1_sm70_out(torch::Tensor out, torch::Tensor input, } 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/torch_bindings.cpp b/csrc/torch_bindings.cpp index 4621db264c..3cd72c4217 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); @@ -651,6 +660,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { 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/sm70_qwen38_flash_next_nvfp4.md b/docs/design/sm70_qwen38_flash_next_nvfp4.md index 169425e151..3027d31caf 100644 --- a/docs/design/sm70_qwen38_flash_next_nvfp4.md +++ b/docs/design/sm70_qwen38_flash_next_nvfp4.md @@ -2,13 +2,19 @@ ## Status and ownership -- Status: source bring-up, native Qwen4Exp MTP, and prefix-cache configuration - are implemented; focused CPU/configuration gates pass and the local - ModelScope snapshot is fully verified. Native MTP4 now completes TP4 model - load, graph capture, warmup, and two 1024x256 requests. Acceptance, repeated - token equality, memory, and steady decode are measured below; matched no-MTP - token equality and verifier cost remain pending. -- Integration line: `private/main`. +- 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`. @@ -120,6 +126,155 @@ 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 @@ -153,6 +308,110 @@ prompts with 256 forced output tokens each. The artifact is 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, @@ -166,9 +425,10 @@ prompts with 256 forced output tokens each. The artifact is 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 target is at - least 100 emitted tokens/s (at most 10 ms/token) with CUDA graphs enabled. - Record an otherwise identical eager control. + 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. 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/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/quantization/test_sm70_modelopt_mixed_nvfp4.py b/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py index da2cd7dc1f..5b6406f0f1 100644 --- a/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py +++ b/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py @@ -21,6 +21,7 @@ _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, ) @@ -117,6 +118,32 @@ 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( diff --git a/tests/quantization/test_sm70_online_qpn8.py b/tests/quantization/test_sm70_online_qpn8.py new file mode 100644 index 0000000000..2122a7edab --- /dev/null +++ b/tests/quantization/test_sm70_online_qpn8.py @@ -0,0 +1,171 @@ +# 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_is_opt_in(monkeypatch): + monkeypatch.delenv("VLLM_SM70_QWEN4_EXP_ONLINE_QPN8", raising=False) + envs.disable_envs_cache() + try: + assert not envs.VLLM_SM70_QWEN4_EXP_ONLINE_QPN8 + monkeypatch.setenv("VLLM_SM70_QWEN4_EXP_ONLINE_QPN8", "1") + envs.disable_envs_cache() + assert envs.VLLM_SM70_QWEN4_EXP_ONLINE_QPN8 + finally: + envs.disable_envs_cache() + + +@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="qwen4_exp_text", + hidden_size=2560, + num_hidden_layers=48, + num_experts=512, + ) + 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() + + +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/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..bbfec35b0e --- /dev/null +++ b/tests/v1/worker/test_gpu_model_runner_v2_greedy.py @@ -0,0 +1,134 @@ +# 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.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] + + +@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/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 a00b80e614..4395d864e7 100644 --- a/vllm/_sm70_ops.py +++ b/vllm/_sm70_ops.py @@ -33,6 +33,17 @@ def _maybe_load_fp8_qpn8_library() -> None: _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" + 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 +80,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) @@ -524,7 +555,7 @@ def fp8_qpn8_prepare_sm70( 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"): @@ -550,7 +581,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"): @@ -573,7 +604,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, @@ -609,7 +640,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, @@ -654,7 +685,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, @@ -693,7 +724,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, @@ -735,7 +766,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, @@ -780,7 +811,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, @@ -792,6 +823,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") @@ -808,6 +910,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, @@ -1249,6 +1374,56 @@ def _mxfp4_moe_qpn_m1_sm70_out_fake( 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, diff --git a/vllm/distributed/device_communicators/custom_all_reduce.py b/vllm/distributed/device_communicators/custom_all_reduce.py index 790922fd6f..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 diff --git a/vllm/envs.py b/vllm/envs.py index 8fdffd50cc..3233a41804 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -165,6 +165,8 @@ 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_QWEN4_EXP_ONLINE_QPN8: bool = False + VLLM_SM70_QWEN3NEXT_SHARED_GATE_FUSION: bool = False 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 @@ -177,6 +179,9 @@ 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 @@ -216,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 @@ -1696,6 +1702,15 @@ 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"))), + "VLLM_SM70_QWEN4_EXP_ONLINE_QPN8": lambda: bool( + int(os.getenv("VLLM_SM70_QWEN4_EXP_ONLINE_QPN8", "0")) + ), + # 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", "0")) + ), # 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 @@ -1810,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")) ), @@ -2007,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 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/linear.py b/vllm/model_executor/layers/linear.py index 7f549fed61..19f74ab5f7 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -370,6 +370,13 @@ 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 ): @@ -435,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 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..647b4eda4d 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 @@ -2374,7 +2374,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 +2383,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 +7127,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 +7137,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 +7146,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 +7170,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 +7233,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 +7341,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/quantization/nvfp4_sm70_moe.py b/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py index 938439ddb9..460134a96c 100644 --- a/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py +++ b/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py @@ -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, @@ -47,6 +48,30 @@ _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 @@ -294,6 +319,11 @@ 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( "SM70 NVFP4 MoE requires the TurboMind extension " @@ -639,6 +669,35 @@ def apply( output = buffers["output"] 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, 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..58031fc110 --- /dev/null +++ b/vllm/model_executor/layers/quantization/sm70_online_qpn8.py @@ -0,0 +1,287 @@ +# 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 getattr(text_config, "model_type", None) == "qwen4_exp_text" + 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 + ) + + +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/models/qwen2_moe.py b/vllm/model_executor/models/qwen2_moe.py index ecfbf4f064..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) @@ -155,7 +166,26 @@ def forward(self, x): 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/models/qwen4_exp/nvidia/hyperconnection.py b/vllm/models/qwen4_exp/nvidia/hyperconnection.py index 8ca503d214..dbe637d80d 100644 --- a/vllm/models/qwen4_exp/nvidia/hyperconnection.py +++ b/vllm/models/qwen4_exp/nvidia/hyperconnection.py @@ -29,6 +29,9 @@ 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 ( @@ -124,6 +127,27 @@ def __init__( 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]: @@ -134,18 +158,7 @@ def mix( self.hc_count, ) - if self.use_combine: - # produce injection logits for combine - 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) # [M, D] - block_input = hc_gate_mix(xn, gate, self.hc_count) + block_input, injection = self._project(xn) return hidden_states, block_input, injection @@ -170,18 +183,7 @@ def combine_and_mix( self.hc_count, ) - if self.use_combine: - # produce injection logits for combine - 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) # [M, D] - block_input = hc_gate_mix(xn, gate, self.hc_count) + block_input, injection = self._project(xn) return hidden_states, block_input, injection diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py index 526acbe0a4..de665095dc 100644 --- a/vllm/models/qwen4_exp/nvidia/model.py +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -851,6 +851,9 @@ def get_mamba_specs_from_config( 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 @@ -1023,6 +1026,9 @@ def embed_input_ids( 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, diff --git a/vllm/models/qwen4_exp/nvidia/ops/hc.py b/vllm/models/qwen4_exp/nvidia/ops/hc.py index a3487b916e..cde145f6fb 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/hc.py +++ b/vllm/models/qwen4_exp/nvidia/ops/hc.py @@ -352,7 +352,10 @@ def _hc_combine_norm( out = residual.new_empty(residual.shape) y = residual.new_empty(residual.shape) - BLOCK_SIZE = 512 + # 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, @@ -371,6 +374,7 @@ def _hc_combine_norm( EPS=eps, BLOCK_SIZE=BLOCK_SIZE, launch_pdl=current_platform.is_arch_support_pdl(), + num_warps=4, ) return out, y diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py index b1fda40d18..8b5ebf099c 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py @@ -637,7 +637,11 @@ def qsa_mqa_paged( 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 - BLOCK_N = 64 + 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. @@ -873,6 +877,16 @@ def qsa_sparse_paged_attention( else: block_n, target_splits, partial_warps = 64, 1, 2 + 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) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index b85b6d8e9b..c12dbc8081 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 @@ -503,7 +504,9 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: def _init_kv_zero_meta(self) -> None: """Precompute metadata used to clear newly allocated cache blocks.""" self._kv_block_zeroer = KVBlockZeroer( - self.device, pin_memory=is_pin_memory_available() + self.device, + pin_memory=is_pin_memory_available(), + max_concurrency=self.vllm_config.max_concurrent_batches, ) self._kv_block_zeroer.init_meta( attn_groups_iter=(group for groups in self.attn_groups for group in groups), @@ -1044,6 +1047,28 @@ 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 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: 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, From a42459fe0c46e662279c237b13252f16747d3419 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 28 Aug 2026 01:12:57 +0800 Subject: [PATCH 091/105] [CI] Fix Qwen3.8 public integration hooks Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu | 18 ++++++++---------- vllm/v1/worker/gpu/model_runner.py | 4 +--- 2 files changed, 9 insertions(+), 13 deletions(-) diff --git a/csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu b/csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu index cf4458e165..18574d8add 100644 --- a/csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu +++ b/csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu @@ -194,8 +194,8 @@ __global__ void nvfp4_qpn_m1_sm70_kernel(const half* __restrict__ input, // 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)); + const half2 scale = + __hmul2(__halves2half2(scalar, scalar), __float2half2_rn(16384.0f)); half2 decoded[8]; dequant_e2m1x8(packed0, scale, decoded); @@ -262,8 +262,7 @@ void launch_nvfp4_qpn_m1(torch::Tensor out, torch::Tensor 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()), @@ -276,10 +275,10 @@ 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); \ +#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); @@ -292,8 +291,7 @@ void dispatch_nvfp4_qpn_m1(torch::Tensor out, torch::Tensor input, 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 ", + TORCH_CHECK(false, "nvfp4_moe_qpn_m1_sm70_out: unsupported split_k ", split_k); } #undef VLLM_NVFP4_QPN_M1_CASE diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index c12dbc8081..32708ff825 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -504,9 +504,7 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: def _init_kv_zero_meta(self) -> None: """Precompute metadata used to clear newly allocated cache blocks.""" self._kv_block_zeroer = KVBlockZeroer( - self.device, - pin_memory=is_pin_memory_available(), - max_concurrency=self.vllm_config.max_concurrent_batches, + self.device, pin_memory=is_pin_memory_available() ) self._kv_block_zeroer.init_meta( attn_groups_iter=(group for groups in self.attn_groups for group in groups), From 32b6178e57d4527ce1f0f778e1502233038ad79c Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 28 Aug 2026 01:30:52 +0800 Subject: [PATCH 092/105] [CI][SM70] Start #361 audit repair Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> From 992f883f9a6677b51537da3ea01a3267b2befaab Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 28 Aug 2026 01:33:43 +0800 Subject: [PATCH 093/105] [Bugfix][SM70] Harden Qwen3.8 default fast paths Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/quantization/test_sm70_online_qpn8.py | 46 +++++++++++++++++-- .../worker/test_gpu_model_runner_v2_greedy.py | 5 ++ vllm/_sm70_ops.py | 6 ++- vllm/envs.py | 8 ++-- .../layers/quantization/sm70_online_qpn8.py | 10 +++- vllm/v1/worker/gpu/model_runner.py | 1 + 6 files changed, 65 insertions(+), 11 deletions(-) diff --git a/tests/quantization/test_sm70_online_qpn8.py b/tests/quantization/test_sm70_online_qpn8.py index 2122a7edab..cd2aa75285 100644 --- a/tests/quantization/test_sm70_online_qpn8.py +++ b/tests/quantization/test_sm70_online_qpn8.py @@ -10,18 +10,44 @@ from vllm.model_executor.layers.quantization import sm70_online_qpn8 as online_qpn8 -def test_qwen4_exp_online_qpn8_is_opt_in(monkeypatch): +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 - monkeypatch.setenv("VLLM_SM70_QWEN4_EXP_ONLINE_QPN8", "1") + finally: envs.disable_envs_cache() - assert envs.VLLM_SM70_QWEN4_EXP_ONLINE_QPN8 + + +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"), [ @@ -50,10 +76,19 @@ def test_online_qpn8_shape_gate(prefix, k, n, expected): def test_online_qpn8_runtime_contract_is_tp4_no_mtp(monkeypatch): text_config = SimpleNamespace( - model_type="qwen4_exp_text", + 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), @@ -68,6 +103,9 @@ def test_online_qpn8_runtime_contract_is_tp4_no_mtp(monkeypatch): 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): diff --git a/tests/v1/worker/test_gpu_model_runner_v2_greedy.py b/tests/v1/worker/test_gpu_model_runner_v2_greedy.py index bbfec35b0e..92f8b2d9bc 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_greedy.py +++ b/tests/v1/worker/test_gpu_model_runner_v2_greedy.py @@ -54,6 +54,7 @@ def compute_logits(self, _hidden_states: torch.Tensor) -> torch.Tensor: runner = object.__new__(GPUModelRunner) runner.model = Model() + runner.lora_config = None runner.sampler = SimpleNamespace( can_use_sm70_greedy_token_fastpath=lambda _input_batch: True ) @@ -91,6 +92,10 @@ def compute_logits(self, _hidden_states: torch.Tensor) -> torch.Tensor: 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", diff --git a/vllm/_sm70_ops.py b/vllm/_sm70_ops.py index 4395d864e7..97f34f4241 100644 --- a/vllm/_sm70_ops.py +++ b/vllm/_sm70_ops.py @@ -24,9 +24,11 @@ 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" specific_enabled = generic_override != "0" and specific_override == "1" - if generic_enabled or specific_enabled: + online_enabled = online_override != "0" + if generic_enabled or specific_enabled or online_enabled: torch.ops.load_library(library_path) @@ -36,7 +38,7 @@ def _maybe_load_fp8_qpn8_library() -> None: 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" + 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) diff --git a/vllm/envs.py b/vllm/envs.py index 3233a41804..2891323f90 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -165,8 +165,8 @@ 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_QWEN4_EXP_ONLINE_QPN8: bool = False - VLLM_SM70_QWEN3NEXT_SHARED_GATE_FUSION: bool = False + 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 @@ -1703,13 +1703,13 @@ def _resolve_rust_frontend_path() -> str | None: # the compressed-tensors scheme. Explicit 0 disables both routes. "VLLM_SM70_FP8_QPN8": lambda: bool(int(os.getenv("VLLM_SM70_FP8_QPN8", "0"))), "VLLM_SM70_QWEN4_EXP_ONLINE_QPN8": lambda: bool( - int(os.getenv("VLLM_SM70_QWEN4_EXP_ONLINE_QPN8", "0")) + 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", "0")) + 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 diff --git a/vllm/model_executor/layers/quantization/sm70_online_qpn8.py b/vllm/model_executor/layers/quantization/sm70_online_qpn8.py index 58031fc110..6ab67cc784 100644 --- a/vllm/model_executor/layers/quantization/sm70_online_qpn8.py +++ b/vllm/model_executor/layers/quantization/sm70_online_qpn8.py @@ -70,10 +70,18 @@ def _exact_runtime_contract() -> bool: return bool( tp_size == 4 and config.speculative_config is None - and getattr(text_config, "model_type", None) == "qwen4_exp_text" 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 ) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 32708ff825..2c32775e96 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1054,6 +1054,7 @@ def sample( 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) From a64245df3d4f6804eabc05b5e2075390027873e9 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 28 Aug 2026 01:54:08 +0800 Subject: [PATCH 094/105] [WIP] Repair #355 against latest main Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> From 62b0979337e6777f02fafc66c1692c5d907dde16 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 28 Aug 2026 02:01:32 +0800 Subject: [PATCH 095/105] [CI] Format q16 grouped verifier Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../kernel/flash_decode_paged.cu | 113 +++++++++--------- 1 file changed, 54 insertions(+), 59 deletions(-) diff --git a/flash-attention-v100/kernel/flash_decode_paged.cu b/flash-attention-v100/kernel/flash_decode_paged.cu index 95ef338821..6d92630656 100644 --- a/flash-attention-v100/kernel/flash_decode_paged.cu +++ b/flash-attention-v100/kernel/flash_decode_paged.cu @@ -1821,12 +1821,9 @@ 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 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, @@ -1928,8 +1925,7 @@ __device__ __forceinline__ void grouped_verify_scale_output_fragment( } template __global__ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_e5m2_partial_kernel( @@ -3550,8 +3546,7 @@ at::Tensor flash_attention_grouped_verify_paged( 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 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 && @@ -3576,10 +3571,10 @@ at::Tensor flash_attention_grouped_verify_paged( kGroupedVerifyHeads, kGroupedVerifyHeadDim}), "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( + 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"); @@ -3604,7 +3599,7 @@ at::Tensor flash_attention_grouped_verify_paged( const dim3 partial_grid(head_groups, grouped_splits, 1); const size_t partial_shared_mem = sizeof(GroupedVerifySmem); -#define LAUNCH_GROUPED_VERIFY_PARTIAL(MAX_QUERY_TOKENS, TWO_PASS, PAGE_SIZE, \ +#define LAUNCH_GROUPED_VERIFY_PARTIAL(MAX_QUERY_TOKENS, TWO_PASS, PAGE_SIZE, \ SINGLE_QUERY, CONTIGUOUS_LAYOUT, \ STAGE_PAGE_IDS) \ do { \ @@ -3640,43 +3635,43 @@ at::Tensor flash_attention_grouped_verify_paged( } while (0) #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); \ - } \ + 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; @@ -3697,13 +3692,13 @@ at::Tensor flash_attention_grouped_verify_paged( #undef LAUNCH_GROUPED_VERIFY_PARTIAL const dim3 combine_grid(static_cast(q.size(0)), kGroupedVerifyHeads, 1); -#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()), \ +#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 (wide_query) { LAUNCH_GROUPED_VERIFY_COMBINE(kGroupedVerifyQ16MaxQ, false); From 4ba7cb53c447bac9457d968e22cec9eefd058346 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 28 Aug 2026 02:13:54 +0800 Subject: [PATCH 096/105] [WIP] Sync #328 with latest main Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> From c689e06028a2d35f7e98c45f1abc1b5c0edc3dd2 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 28 Aug 2026 00:51:20 +0800 Subject: [PATCH 097/105] [Core][SM70] Add adaptive DFlash2 lookup drafting Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../kernels/benchmark_dflash2_lookup.py | 162 +++++++ docs/design/sm70_dflash2_lookup_augmented.md | 105 +++++ docs/features/speculative_decoding/README.md | 37 +- .../test_sm70_flash_v100_grouped_verify.py | 4 + .../attention/test_sm70_flash_v100_policy.py | 5 +- .../test_sm70_mtp_split_cudagraphs.py | 50 +++ tests/v1/spec_decode/test_dflash2.py | 2 + tests/v1/spec_decode/test_dflash2_lookup.py | 220 ++++++++++ .../spec_decode/test_dflash2_ngram_assist.py | 89 ++++ .../v1/spec_decode/test_dflash_mrv2_config.py | 114 ++++- vllm/config/speculative.py | 47 +- vllm/config/vllm.py | 26 +- vllm/envs.py | 30 ++ vllm/model_executor/models/qwen3_dflash2.py | 6 +- vllm/v1/attention/backends/flash_attn_v100.py | 9 +- vllm/v1/worker/gpu/cudagraph_utils.py | 54 ++- vllm/v1/worker/gpu/model_runner.py | 11 +- .../gpu/spec_decode/dflash/speculator.py | 37 +- .../worker/gpu/spec_decode/dflash2/lookup.py | 338 +++++++++++++++ .../gpu/spec_decode/dflash2/speculator.py | 409 +++++++++++++++++- vllm/v1/worker/gpu/spec_decode/utils.py | 14 +- 21 files changed, 1701 insertions(+), 68 deletions(-) create mode 100644 benchmarks/kernels/benchmark_dflash2_lookup.py create mode 100644 docs/design/sm70_dflash2_lookup_augmented.md create mode 100644 tests/v1/spec_decode/test_dflash2_lookup.py create mode 100644 vllm/v1/worker/gpu/spec_decode/dflash2/lookup.py diff --git a/benchmarks/kernels/benchmark_dflash2_lookup.py b/benchmarks/kernels/benchmark_dflash2_lookup.py new file mode 100644 index 0000000000..52a07e7507 --- /dev/null +++ b/benchmarks/kernels/benchmark_dflash2_lookup.py @@ -0,0 +1,162 @@ +# 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.cuda.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.cuda.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/docs/design/sm70_dflash2_lookup_augmented.md b/docs/design/sm70_dflash2_lookup_augmented.md new file mode 100644 index 0000000000..1443c4316a --- /dev/null +++ b/docs/design/sm70_dflash2_lookup_augmented.md @@ -0,0 +1,105 @@ +# 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 earlier GPU LABD suite +also passes suffix hit/miss/overlap, request eligibility, 7-to-15 fusion, 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. + +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. + +## Remaining promotion gates + +1. Benchmark generic 3,456-page grouped verification against independent XQA + at q8/q16 and 1K/32K/128K. Add a fixed CUDA specialization only if the + generic path loses materially. +2. Run a normal miss workload and a context-copy workload with synchronous + scheduling. Scheduler metrics and profile transitions must prove q8 remains + the default and q16 is entered only for sustained hits. +3. 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. +4. Run probabilistic quality gates against target-only and the accepted DFlash2 + baseline, including coding, tool calling, structured JSON, long context, + and 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/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/tests/kernels/attention/test_sm70_flash_v100_grouped_verify.py b/tests/kernels/attention/test_sm70_flash_v100_grouped_verify.py index 8000984b07..c895d4f040 100644 --- a/tests/kernels/attention/test_sm70_flash_v100_grouped_verify.py +++ b/tests/kernels/attention/test_sm70_flash_v100_grouped_verify.py @@ -132,8 +132,12 @@ def _reference( (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() diff --git a/tests/v1/attention/test_sm70_flash_v100_policy.py b/tests/v1/attention/test_sm70_flash_v100_policy.py index 57a7d0eb08..8627c64663 100644 --- a/tests/v1/attention/test_sm70_flash_v100_policy.py +++ b/tests/v1/attention/test_sm70_flash_v100_policy.py @@ -1704,8 +1704,9 @@ def fake_decode( @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, + query_len: int, page_size: int ): from vllm.v1.attention.backends.flash_attn_v100 import FlashAttnV100Impl @@ -1755,7 +1756,7 @@ def grouped_verify( layer = SimpleNamespace(_k_scale_float=0.5, _v_scale_float=2.0) 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 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_dflash2.py b/tests/v1/spec_decode/test_dflash2.py index 6fa8aa81ff..6cf705973a 100644 --- a/tests/v1/spec_decode/test_dflash2.py +++ b/tests/v1/spec_decode/test_dflash2.py @@ -357,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 @@ -1146,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..d589689362 --- /dev/null +++ b/tests/v1/spec_decode/test_dflash2_lookup.py @@ -0,0 +1,220 @@ +# 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_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/vllm/config/speculative.py b/vllm/config/speculative.py index 9d18c98185..41dbed8515 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -82,6 +82,43 @@ 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.""" @@ -163,10 +200,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 diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 48c994f894..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 @@ -1125,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 " @@ -1168,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/envs.py b/vllm/envs.py index 264644551b..9940e5b3ad 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -2895,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")) ), 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/v1/attention/backends/flash_attn_v100.py b/vllm/v1/attention/backends/flash_attn_v100.py index 7a324bb319..cea126b91b 100644 --- a/vllm/v1/attention/backends/flash_attn_v100.py +++ b/vllm/v1/attention/backends/flash_attn_v100.py @@ -24,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 @@ -3099,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( @@ -5186,7 +5188,10 @@ def _dflash2_grouped_verify_allowed( 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 diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 24aad335c5..d7d1dc6ca5 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, @@ -146,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 = (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 @@ -179,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: @@ -188,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() @@ -199,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/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 1b61c4a508..6d688a9c3a 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -214,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, @@ -1510,7 +1512,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/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..384c9457f3 --- /dev/null +++ b/vllm/v1/worker/gpu/spec_decode/dflash2/lookup.py @@ -0,0 +1,338 @@ +# 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..d80ff31675 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,9 +398,11 @@ 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 ngram_assist and self.draft_block == self.num_speculative_steps: min_ngram = speculative_config.prompt_lookup_min max_ngram = speculative_config.prompt_lookup_max assert min_ngram is not None and max_ngram is not None @@ -369,11 +444,233 @@ 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. @@ -392,9 +689,9 @@ def _sample_path( candidate_ids = candidate_ids.contiguous() block_k = triton.next_power_of_2(self.selector_top_k) walk_steps = ( - self.num_speculative_steps - 1 + self.draft_block - 1 if self._use_sm70_tail - else self.num_speculative_steps + else self.draft_block ) _selector_walk_kernel[(num_reqs,)]( scores, @@ -403,10 +700,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 +719,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 +744,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 +874,73 @@ 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 +957,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 +982,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/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. From 4cd1f0e72995c7f30e4f075fba80fb1c3166baae Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 28 Aug 2026 02:10:38 +0800 Subject: [PATCH 098/105] [Bugfix] Preserve burst tool-call arguments Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../test_qwen3coder_tool_parser.py | 54 +++++++++++++++++++ vllm/tool_parsers/qwen3coder_tool_parser.py | 29 ++++++++-- 2 files changed, 80 insertions(+), 3 deletions(-) diff --git a/tests/tool_parsers/test_qwen3coder_tool_parser.py b/tests/tool_parsers/test_qwen3coder_tool_parser.py index 2df97f2536..dbeeb61e1d 100644 --- a/tests/tool_parsers/test_qwen3coder_tool_parser.py +++ b/tests/tool_parsers/test_qwen3coder_tool_parser.py @@ -1571,6 +1571,60 @@ def test_streaming_final_param_and_function_close_single_chunk( } +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", [ diff --git a/vllm/tool_parsers/qwen3coder_tool_parser.py b/vllm/tool_parsers/qwen3coder_tool_parser.py index 6ac6b53b5f..a14d8b2e22 100644 --- a/vllm/tool_parsers/qwen3coder_tool_parser.py +++ b/vllm/tool_parsers/qwen3coder_tool_parser.py @@ -711,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 "{}", } ) @@ -726,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( @@ -735,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", ) From 7849d8b9bb2987a5f705abe84ca063e0c398102d Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 28 Aug 2026 02:10:44 +0800 Subject: [PATCH 099/105] [Core][SM70] Warm MRV2 GDN decode fallback Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../v1/attention/test_gdn_metadata_builder.py | 23 +++++++++++++++++++ .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 23 +++++++++++++++++++ vllm/v1/worker/gpu/model_runner.py | 11 +++++++++ 3 files changed, 57 insertions(+) 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/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 647b4eda4d..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, diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 6d688a9c3a..ca73473d5c 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -522,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( From e0067ad933de1236064208a46e477ebe2f27f872 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 28 Aug 2026 02:10:48 +0800 Subject: [PATCH 100/105] [Doc] Record adaptive DFlash2 validation Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_dflash2_lookup_augmented.md | 94 +++++++++++++++++--- 1 file changed, 80 insertions(+), 14 deletions(-) diff --git a/docs/design/sm70_dflash2_lookup_augmented.md b/docs/design/sm70_dflash2_lookup_augmented.md index 1443c4316a..c8409e3ea3 100644 --- a/docs/design/sm70_dflash2_lookup_augmented.md +++ b/docs/design/sm70_dflash2_lookup_augmented.md @@ -62,9 +62,9 @@ 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 earlier GPU LABD suite -also passes suffix hit/miss/overlap, request eligibility, 7-to-15 fusion, and -sparse point-mass cache invariants. +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: @@ -79,27 +79,93 @@ 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. Benchmark generic 3,456-page grouped verification against independent XQA - at q8/q16 and 1K/32K/128K. Add a fixed CUDA specialization only if the - generic path loses materially. -2. Run a normal miss workload and a context-copy workload with synchronous - scheduling. Scheduler metrics and profile transitions must prove q8 remains - the default and q16 is entered only for sustained hits. -3. Compare q7 DFlash2, adaptive LABD, and fixed q16 at B1/B2/B4. Report target +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. -4. Run probabilistic quality gates against target-only and the accepted DFlash2 - baseline, including coding, tool calling, structured JSON, long context, - and all 128 relevant CUDA-Graph/prefix-cache residues. No task-score or - parser-validity regression is acceptable. +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. From 1d96b35660f392fe49da08da65b00055b7dca5e0 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 28 Aug 2026 02:15:29 +0800 Subject: [PATCH 101/105] [CI] Format adaptive DFlash2 lookup Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../kernels/benchmark_dflash2_lookup.py | 7 ++---- tests/v1/spec_decode/test_dflash2_lookup.py | 7 ++---- vllm/config/speculative.py | 13 ++++------ .../worker/gpu/spec_decode/dflash2/lookup.py | 24 +++++-------------- .../gpu/spec_decode/dflash2/speculator.py | 18 ++++---------- 5 files changed, 18 insertions(+), 51 deletions(-) diff --git a/benchmarks/kernels/benchmark_dflash2_lookup.py b/benchmarks/kernels/benchmark_dflash2_lookup.py index 52a07e7507..e4df4bedf0 100644 --- a/benchmarks/kernels/benchmark_dflash2_lookup.py +++ b/benchmarks/kernels/benchmark_dflash2_lookup.py @@ -36,8 +36,7 @@ def _make_history(batch_size: int, context_len: int) -> UvaBuffer: history.cpu[row, match_start : match_start + suffix.numel()].copy_(suffix) history.cpu[ row, - match_start - + suffix.numel() : match_start + match_start + suffix.numel() : match_start + suffix.numel() + continuation.numel(), ].copy_(continuation) @@ -54,9 +53,7 @@ def _benchmark_case( 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 - ) + 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( diff --git a/tests/v1/spec_decode/test_dflash2_lookup.py b/tests/v1/spec_decode/test_dflash2_lookup.py index d589689362..13001125d8 100644 --- a/tests/v1/spec_decode/test_dflash2_lookup.py +++ b/tests/v1/spec_decode/test_dflash2_lookup.py @@ -30,8 +30,7 @@ def _reference_lookup( 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] + and tokens[candidate_end - match_len] == tokens[suffix_end - match_len] ): match_len += 1 if match_len >= nmin and (match_len, candidate_end) > best: @@ -89,9 +88,7 @@ def test_suffix_lookup_matches_reference_and_supports_overlap(k: int) -> None: 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 - ) + 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) diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 41dbed8515..6cfd519856 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -88,12 +88,9 @@ def get_dflash_model_draft_tokens(speculative_config: Any) -> int: 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) + 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 @@ -109,9 +106,7 @@ def get_dflash_model_draft_tokens(speculative_config: Any) -> int: 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 - ) + 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)) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash2/lookup.py b/vllm/v1/worker/gpu/spec_decode/dflash2/lookup.py index 384c9457f3..0bc065907f 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash2/lookup.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash2/lookup.py @@ -86,8 +86,7 @@ def _suffix_lookup_kernel( 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), + match_len.to(tl.int64) * _SCORE_STRIDE + candidate_end.to(tl.int64), -1, ) best = tl.maximum(best, tl.max(score, axis=0)) @@ -155,24 +154,15 @@ def _fuse_draft_kernel( ).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 + (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 - ) + 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, @@ -183,9 +173,7 @@ def _fuse_draft_kernel( # 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 - ) + 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)) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py index d80ff31675..26d2498735 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py @@ -486,9 +486,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): 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_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() @@ -664,9 +662,7 @@ def next_num_draft_tokens(self) -> int: long_active=self._lookup_long_active, ) width = ( - self.num_speculative_steps - if self._lookup_long_active - else self.draft_block + 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) @@ -688,11 +684,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.draft_block - 1 - if self._use_sm70_tail - else self.draft_block - ) + walk_steps = self.draft_block - 1 if self._use_sm70_tail else self.draft_block _selector_walk_kernel[(num_reqs,)]( scores, candidate_ids, @@ -919,9 +911,7 @@ def _apply_lookup(self, num_reqs: int) -> 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,) - ]( + _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, From ee67907c3379f569942287eb95b44e87b4654d46 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 28 Aug 2026 02:16:04 +0800 Subject: [PATCH 102/105] [WIP] Sync #302 with latest main Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> From 808504a0d85ad0b3b5a32567807690ae88033474 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 28 Aug 2026 02:31:20 +0800 Subject: [PATCH 103/105] [CI][SM70] Start adaptive DFlash2 repair Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> From 0c8ed2ac04895f83ad91cdaf4e09bb379c97c545 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 28 Aug 2026 02:53:31 +0800 Subject: [PATCH 104/105] [CI][SM70] Repair adaptive DFlash2 gates Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- benchmarks/kernels/benchmark_dflash2_lookup.py | 4 ++-- vllm/v1/worker/gpu/cudagraph_utils.py | 2 +- vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py | 6 +++++- 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/benchmarks/kernels/benchmark_dflash2_lookup.py b/benchmarks/kernels/benchmark_dflash2_lookup.py index e4df4bedf0..65c8300ad0 100644 --- a/benchmarks/kernels/benchmark_dflash2_lookup.py +++ b/benchmarks/kernels/benchmark_dflash2_lookup.py @@ -15,7 +15,7 @@ def _time_ms(fn: Callable[[], None], repeats: int, warmup: int) -> float: for _ in range(warmup): fn() - torch.cuda.synchronize() + torch.accelerator.synchronize() start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) start.record() @@ -114,7 +114,7 @@ def lookup_and_fuse() -> None: for _ in range(3): lookup_and_fuse() torch.cuda.current_stream().wait_stream(capture_stream) - torch.cuda.synchronize() + torch.accelerator.synchronize() with torch.cuda.graph(graph): lookup_and_fuse() graph_ms = _time_ms(graph.replay, repeats, warmup) diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index d7d1dc6ca5..0a09cab818 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -150,7 +150,7 @@ 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 = (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 diff --git a/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py index 26d2498735..9de53db15e 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py @@ -402,7 +402,11 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): speculative_config is not None and getattr(speculative_config, "ngram_assist", False) ) - if ngram_assist and self.draft_block == self.num_speculative_steps: + 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 assert min_ngram is not None and max_ngram is not None From 7e0e68ab4beb75d2614dd8c5685a2ae8301b13b4 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 28 Aug 2026 03:14:59 +0800 Subject: [PATCH 105/105] [CI][SM70] Start PR #308 latest-main sync Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>