diff --git a/benchmarks/benchmark_sm70_decode.py b/benchmarks/benchmark_sm70_decode.py index bf6d364c19..f10eac8d3d 100644 --- a/benchmarks/benchmark_sm70_decode.py +++ b/benchmarks/benchmark_sm70_decode.py @@ -52,20 +52,42 @@ def _sm70_fa2_d256_prefill_status(torch: Any) -> dict[str, Any]: "sm70_d256_splitd_n32_paged_fwd", ) optional_ops = ("sm70_d256_splitd_n32_dense_splitkv3_fwd",) - try: - importlib.import_module("vllm.vllm_flash_attn.flash_attn_interface") - except (AttributeError, ImportError, RuntimeError) as exc: - return { - "available": False, - "error": f"{type(exc).__name__}: {exc}", - "extension_file": None, - "extension_realpath": None, - "required_ops": {name: False for name in required_ops}, - "optional_ops": {name: False for name in optional_ops}, - } - extension_file = _module_file("vllm.vllm_flash_attn._vllm_fa2_C") namespace = getattr(torch.ops, "_vllm_fa2_C", None) + import_error: Exception | None = None + if namespace is None or not all(hasattr(namespace, name) for name in required_ops): + try: + importlib.import_module("vllm.vllm_flash_attn.flash_attn_interface") + except (AttributeError, ImportError, RuntimeError) as exc: + import_error = exc + namespace = getattr(torch.ops, "_vllm_fa2_C", None) + if namespace is None or not all(hasattr(namespace, name) for name in required_ops): + library_path = os.getenv("VLLM_SM70_FA2_D256_LIBRARY") + if library_path is not None: + try: + torch.ops.load_library(library_path) + except (OSError, RuntimeError) as load_exc: + return { + "available": False, + "error": f"{type(load_exc).__name__}: {load_exc}", + "extension_file": library_path, + "extension_realpath": str(Path(library_path).resolve()), + "required_ops": {name: False for name in required_ops}, + "optional_ops": {name: False for name in optional_ops}, + } + extension_file = library_path + elif import_error is not None: + return { + "available": False, + "error": f"{type(import_error).__name__}: {import_error}", + "extension_file": None, + "extension_realpath": None, + "required_ops": {name: False for name in required_ops}, + "optional_ops": {name: False for name in optional_ops}, + } + namespace = getattr(torch.ops, "_vllm_fa2_C", None) + if extension_file is None: + extension_file = os.getenv("VLLM_SM70_FA2_D256_LIBRARY") required_status = { name: namespace is not None and hasattr(namespace, name) for name in required_ops 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 new file mode 100644 index 0000000000..1621b4c077 --- /dev/null +++ b/docs/design/sm70_dflash2_prefill_closure.md @@ -0,0 +1,171 @@ +# SM70 DFlash2 Prefill Closure + +## Scope and integration base + +This private campaign is stacked on the quality-audited DFlash2 branch at +`ee4ac48a479c3dbd458d5f7c09a59f39fd271d82`. It keeps the accepted NVFP4 +target, official BF16 DFlash2 draft, FP8 E5M2 target KV, FP16 draft KV, +prefix caching, Mamba alignment, and CUDA Graph decode contract. Draft-MLP +QPN8 remains disabled. + +The first objective is to restore already accepted SM70 prefill operators in +source-overlay deployments. Kernel arithmetic is not changed by that repair. +Any later shape expansion is a separate gate. + +## Project PR audit + +The historical short- and long-prefill numbers use different contracts: + +| Evidence | Contract | Accepted result | Boundary | +|---|---|---:|---| +| Public PR #271 | Qwen3.8-27B-FP8, TP4, exact input 8000, target-only | 5121.44 request-wall and 5170.96 pure-prefill tok/s | Exact-8K only | +| Public PR #324 | Same exact-8K FP8 contract | 5500 tok/s is a campaign target, not a measured implementation result | Documentation-only PR | +| Public PR #224 | Qwen3.8-27B-FP8, TP4, exact input 65536, target-only | 2798.6 to 3496.4 prompt tok/s after the exact D256 operator | Closest retained 64K target-only reference | +| Private PR #8/#13 | Qwen3.8-27B-FP8, TP4, input 261888, chunk 8192 with FP16 Mamba/SSM cache and Q8000 aligned chunks, target-only | 2438.89 prompt tok/s | Stable max-aware D256 architecture | +| Rejected public PR #315 lane | Same 256K FP8 contract | 2971.51 prompt tok/s | Rejected: 32 output token IDs were zero | + +The 2438.89 tok/s route uses max-shifted exponentiation and max-aware online +softmax merging. Its output hash exactly matched the exact control. It also +explicitly overrides the checkpoint's FP32 SSM-cache contract to FP16; that +override is a separate quality variable and is not inherited automatically by +DFlash2. The removed raw-logit half2 polynomial must not be restored. + +## Current DFlash2 baseline + +The same-card cold benchmark resets prefix cache before every warmup and +measurement. It uses the practical 256K API contract, including chunk 4096. + +| Input | Mean pure prefill | Prompt throughput | +|---:|---:|---:| +| 32768 | 10.485347 s | 3125.12 tok/s | +| 65536 | 25.317796 s | 2588.54 tok/s | + +All three repeats at each length emitted the same first-token hash. Artifact: +`/data/minimax-h3/task-cache/v100-dflash2-prefill-32k64k-20260827/current-dflash2-cold-prefill-v1/`. + +## Confirmed root cause + +Every retained practical DFlash2 long-prefill log reports that +`_vllm_fa2_C` cannot be imported. The source checkout contains the D256 +dispatch and quality-safe architecture, but the source overlay shadows the +installed package containing the native extension. Therefore the merged D256 +path is not merely underperforming; it has never executed in these runs. + +The accepted stable binary is retained at +`/data/minimax-h3/task-cache/qwen38-d256-attn-80tflops-20260825/build/exact-stat-256k-py312-v2/_vllm_fa2_C.abi3.so` +with SHA256 `f9f9acbc610c87fce9984e8fbd93fe0c8fa59887542123a74b3eaef6d3b8abf9`. +It loads against the active Torch 2.10/CUDA 12.8 environment and registers the +required dense, paged, split-KV3, and stable GQA architecture operators. + +This branch adds an explicit `VLLM_SM70_FA2_D256_LIBRARY` source-overlay +sidecar. It is opt-in and follows the existing SM70 native-sidecar convention. +Bundled-wheel behavior remains unchanged, and missing or incompatible +operators still fail closed to the existing fallback with a warning. +Both the benchmark preflight and runtime loader validate registered operators, +not merely a successful Python-interface import. This covers partially cached +interfaces that otherwise appear importable while exposing no native kernels. + +## Shape boundary and next measurements + +The practical chunk-4096 contract can use the exact D256 Split-D operators, +but it cannot enter the stable long GQA architecture, whose validated kernel +contract is Q8000 with KV16K..256K in 8K steps. With seven speculative slots, +the scheduler first reduces the configured 4096-token budget to 4089. The +checkpoint's FP32 SSM state makes one aligned attention/Mamba block 1648 +tokens, so the observed steady prefill query is Q3296. A configured 8192-token +budget retains that FP32 state and yields Q6592, not Q8000. The historical +target-only FP16 Mamba/SSM contract used an 800-token block, but DFlash2 grows +the convolution state by seven speculative slots. Its FP16 block is therefore +880 tokens and the same budget yields Q7920, so it cannot enter the existing +Q8000 architecture either. The next paired measurements are therefore +deliberately separated: + +1. chunk 4096 plus the stable sidecar, to measure the dependency-closure gain; +2. chunk 8192 plus the same sidecar and FP32 SSM state, to measure Q6592; +3. chunk 8192 plus FP16 Mamba/SSM cache, to measure the actual Q7920 DFlash2 + geometry rather than crediting the target-only Q8000 route; +4. profile the remaining NVFP4 projection cost before considering a new + Q6592/Q7920 attention architecture. + +Each candidate must prove operator-route hits, preserve output validity, fit +the 256K DFlash2 memory contract, and retain the quality-audit PPL and scored +coding gates. Prefix-hit time is reported separately and never counted as cold +prefill throughput. + +## Dependency-closure A/B + +The first candidate changes only native-extension resolution and keeps chunk +4096. All six measured requests are cold. The stable sidecar loads on all four +ranks and the benchmark reports both required D256 operators as available. + +| Input | Missing-extension control | Stable-sidecar candidate | Throughput gain | +|---:|---:|---:|---:| +| 32768 | 3125.12 tok/s | 3476.53 tok/s | +11.24% | +| 65536 | 2588.54 tok/s | 3103.02 tok/s | +19.87% | + +Candidate pure-prefill means are 9.425490 s and 21.120102 s. The three repeats +at each length retain the control first-token hash +`54363ddee68f4a5db81c9d37e5fb738d28f5b67dc7f725ad7333172b1ea157da`. +Artifact: +`/data/minimax-h3/task-cache/v100-dflash2-prefill-32k64k-20260827/candidate-stable-fa2-q4096-v1/`. + +## Chunk and Mamba dtype closure + +Increasing the configured budget to 8192 while retaining the checkpoint's +FP32 SSM state produces Q6592. It is neutral at 32K and slightly slower at 64K: + +| Input | Q3296 sidecar | Q6592 sidecar | Change | +|---:|---:|---:|---:| +| 32768 | 3476.53 tok/s | 3481.83 tok/s | +0.15% | +| 65536 | 3103.02 tok/s | 3077.03 tok/s | -0.84% | + +The FP16 Mamba/SSM arm produces Q7920, not Q8000, and loses at both lengths: + +| Input | Q3296 sidecar | Q7920 FP16 state | Change | +|---:|---:|---:|---:| +| 32768 | 3476.53 tok/s | 3271.49 tok/s | -5.90% | +| 65536 | 3103.02 tok/s | 2827.03 tok/s | -8.90% | + +All twelve measured requests retain the same first-token hash. Q6592 is not a +promotion, and Q7920 is rejected on speed before spending a dataset-quality +run. Artifacts are `candidate-stable-fa2-q8192-fp32-v4` and +`candidate-stable-fa2-q8000-mamba-fp16-v1` under the task root above. + +## NVFP4 projection candidate + +The accepted DFlash2 target keeps M<=8 verification on the existing QPN2 +route, while large-M prefill currently falls through to TurboMind W4A16. A +single-V100 M3296 microbenchmark shows that the already present bounded FP16 +QPN4 prefill operator is materially faster even though its timing includes +weight dequantization on every call: + +| Projection | TurboMind | bounded QPN4 | Latency reduction | +|---|---:|---:|---:| +| fused gate/up, 5120x8704 | 5.351 ms | 3.896 ms | 27.19% | +| down, 4352x5120 | 2.646 ms | 1.866 ms | 29.47% | + +The candidate reuses the QPN2 code and E4M3 scale-code buffers that are already +resident for verification, plus the existing shared 85-MiB FP16 workspace. It +does not retain a third packed weight layout: QPN2 and QPN4 expose different +2-D shapes but use the same flattened physical tile order, so the bridge uses +zero-copy views. Admission is default-off and requires +`VLLM_SM70_NVFP4_QPN2_PREFILL=1`; M below the separately recorded crossover +threshold keeps TurboMind, and M<=8 verification remains QPN2. The current +campaign stops at 64K; no 128K/256K throughput run is required for promotion. +The single-V100 bridge probe confirms byte-identical flattened code and scale +storage despite shapes `[N,K/2]` and `[K,N/2]`; both ordinary and fused +gate-SiLU QPN4-prefill calls are bitwise equal with maximum absolute error +zero. + +## Rejected SWA tail-only shortcut + +The DFlash2 draft uses a 2048-token sliding window, but projecting and writing +only the currently readable tail is not valid with automatic prefix caching. +Draft-group blocks are registered by content for later requests; a block whose +target tokens were computed without the matching draft-context KV write can be +reused later at a sequence length where those tokens are inside the window. +That produces a cache hit backed by uninitialized draft KV and collapses +acceptance. Therefore this campaign keeps full coverage of every newly +computed target token. A future tail-only design would first need to decouple +draft-group cache registration from target-prefix hits and prove rebuild +semantics; that is outside this prefill dependency closure. diff --git a/tests/benchmarks/test_benchmark_sm70_decode.py b/tests/benchmarks/test_benchmark_sm70_decode.py index b9fe6ebb78..929630db88 100644 --- a/tests/benchmarks/test_benchmark_sm70_decode.py +++ b/tests/benchmarks/test_benchmark_sm70_decode.py @@ -46,3 +46,31 @@ def test_sm70_fa2_d256_prefill_status_requires_dense_and_paged_ops(monkeypatch): assert status["error"] is None assert all(status["required_ops"].values()) assert all(status["optional_ops"].values()) + + +def test_sm70_fa2_d256_prefill_status_accepts_explicit_sidecar(monkeypatch): + namespace = types.SimpleNamespace() + loaded: list[str] = [] + + def load_library(path: str) -> None: + loaded.append(path) + namespace.sm70_d256_splitd_n32_dense_fwd = object() + namespace.sm70_d256_splitd_n32_paged_fwd = object() + + fake_ops = types.SimpleNamespace( + _vllm_fa2_C=namespace, + load_library=load_library, + ) + fake_torch = types.SimpleNamespace(ops=fake_ops) + monkeypatch.setattr( + benchmark_sm70_decode.importlib, + "import_module", + lambda _name: types.SimpleNamespace(), + ) + monkeypatch.setenv("VLLM_SM70_FA2_D256_LIBRARY", "/tmp/stable-fa2.so") + + status = benchmark_sm70_decode._sm70_fa2_d256_prefill_status(fake_torch) + + assert status["available"] is True + assert status["extension_file"] == "/tmp/stable-fa2.so" + assert loaded == ["/tmp/stable-fa2.so"] diff --git a/tests/quantization/test_sm70_nvfp4_qpn2.py b/tests/quantization/test_sm70_nvfp4_qpn2.py index 84f8174158..c9084e7cd6 100644 --- a/tests/quantization/test_sm70_nvfp4_qpn2.py +++ b/tests/quantization/test_sm70_nvfp4_qpn2.py @@ -18,12 +18,20 @@ def test_nvfp4_qpn2_is_default_off_with_explicit_on(monkeypatch): monkeypatch.delenv("VLLM_SM70_NVFP4_QPN2", raising=False) + monkeypatch.delenv("VLLM_SM70_NVFP4_QPN2_PREFILL", raising=False) + monkeypatch.delenv("VLLM_SM70_NVFP4_QPN2_PREFILL_MIN_M", raising=False) envs.disable_envs_cache() try: assert not envs.VLLM_SM70_NVFP4_QPN2 + assert not envs.VLLM_SM70_NVFP4_QPN2_PREFILL + assert envs.VLLM_SM70_NVFP4_QPN2_PREFILL_MIN_M == 1024 monkeypatch.setenv("VLLM_SM70_NVFP4_QPN2", "1") + monkeypatch.setenv("VLLM_SM70_NVFP4_QPN2_PREFILL", "1") + monkeypatch.setenv("VLLM_SM70_NVFP4_QPN2_PREFILL_MIN_M", "9") envs.disable_envs_cache() assert envs.VLLM_SM70_NVFP4_QPN2 + assert envs.VLLM_SM70_NVFP4_QPN2_PREFILL + assert envs.VLLM_SM70_NVFP4_QPN2_PREFILL_MIN_M == 9 finally: envs.disable_envs_cache() @@ -75,6 +83,8 @@ def _make_small_layer() -> torch.nn.Module: def test_nvfp4_qpn2_prepare_and_dispatch_contract(monkeypatch): monkeypatch.setenv("VLLM_SM70_NVFP4_QPN2", "1") + monkeypatch.setenv("VLLM_SM70_NVFP4_QPN2_PREFILL", "1") + monkeypatch.setenv("VLLM_SM70_NVFP4_QPN2_PREFILL_MIN_M", "9") envs.disable_envs_cache() layer = _make_small_layer() calls = [] @@ -88,6 +98,13 @@ def test_nvfp4_qpn2_prepare_and_dispatch_contract(monkeypatch): ) monkeypatch.setattr(nvfp4_scheme, "_is_qpn2_layer", lambda layer: True) monkeypatch.setattr(nvfp4_scheme, "_missing_qpn2_ops", lambda: []) + monkeypatch.setattr(nvfp4_scheme, "_missing_qpn2_prefill_ops", lambda: []) + workspace = torch.empty((1,), dtype=torch.float16) + monkeypatch.setattr( + nvfp4_scheme.sm70_tm, + "get_nvfp4_qpn4_dense_workspace", + lambda _weight: workspace, + ) monkeypatch.setitem(nvfp4_scheme._SM70_NVFP4_QPN2_CONFIGS, (64, 64, False), (8, 2)) monkeypatch.setitem(nvfp4_scheme._SM70_NVFP4_QPN2_CONFIGS, (64, 64, True), (8, 2)) monkeypatch.setattr( @@ -121,12 +138,34 @@ def fake_dispatch(*args): monkeypatch.setattr( nvfp4_scheme.sm70_ops, "nvfp4_qpn2_dispatch_sm70_out", fake_dispatch ) + prefill_calls = [] + + def fake_prefill(*args): + prefill_calls.append(args) + args[0].fill_(5) + + monkeypatch.setattr( + nvfp4_scheme.sm70_ops, + "nvfp4_qpn4_prefill_sm70_out", + fake_prefill, + ) try: scheme.process_weights_after_loading(layer) assert layer.sm70_nvfp4_qpn2 assert layer.sm70_nvfp4_qpn2_gated_silu assert layer.sm70_nvfp4_qpn2_global_scale == 0.5 + assert layer.sm70_nvfp4_qpn2_prefill_dense_weight_ptr == workspace.data_ptr() + assert layer.sm70_nvfp4_qpn2_prefill_codes.shape == (64, 32) + assert layer.sm70_nvfp4_qpn2_prefill_scales.shape == (4, 64) + assert ( + layer.sm70_nvfp4_qpn2_prefill_codes.data_ptr() + == layer.sm70_nvfp4_qpn2_codes.data_ptr() + ) + assert ( + layer.sm70_nvfp4_qpn2_prefill_scales.data_ptr() + == layer.sm70_nvfp4_qpn2_scales.data_ptr() + ) assert layer.weight.numel() == 0 assert layer.weight_scale.numel() == 0 @@ -139,5 +178,16 @@ def fake_dispatch(*args): assert torch.equal(fused, torch.full_like(fused, 3)) assert calls[0][-1] is False assert calls[1][-1] is True + + large_x = torch.ones((9, 64), dtype=torch.float16) + large_raw = scheme.apply_weights(layer, large_x) + large_fused = scheme.apply_fused_silu_and_mul(layer, large_x) + assert torch.equal(large_raw, torch.full_like(large_raw, 5)) + assert large_fused is not None + assert torch.equal(large_fused, torch.full_like(large_fused, 5)) + assert prefill_calls[0][-2:] == (True, False) + assert prefill_calls[1][-2:] == (True, True) + assert prefill_calls[0][3].shape == (64, 32) + assert prefill_calls[0][4].shape == (4, 64) finally: envs.disable_envs_cache() diff --git a/tests/v1/attention/test_sm70_flash_v100_policy.py b/tests/v1/attention/test_sm70_flash_v100_policy.py index 419c4d847c..74bca1faef 100644 --- a/tests/v1/attention/test_sm70_flash_v100_policy.py +++ b/tests/v1/attention/test_sm70_flash_v100_policy.py @@ -738,6 +738,44 @@ def test_sm70_splitd_d256_loader_requires_exact_ops(monkeypatch): assert flash_v100._get_sm70_splitd_d256_ops() == (dense, paged, splitkv3) +def test_sm70_splitd_d256_loader_accepts_explicit_sidecar(monkeypatch): + import vllm.v1.attention.backends.flash_attn_v100 as flash_v100 + + fake_interface = types.ModuleType("vllm.vllm_flash_attn.flash_attn_interface") + fake_package = types.ModuleType("vllm.vllm_flash_attn") + fake_package.__dict__["flash_attn_interface"] = fake_interface + monkeypatch.setitem(sys.modules, "vllm.vllm_flash_attn", fake_package) + monkeypatch.setitem( + sys.modules, + "vllm.vllm_flash_attn.flash_attn_interface", + fake_interface, + ) + + namespace = SimpleNamespace() + loaded: list[str] = [] + + def load_library(path: str) -> None: + loaded.append(path) + namespace.sm70_d256_splitd_n32_dense_fwd = "dense" + namespace.sm70_d256_splitd_n32_paged_fwd = "paged" + + fake_ops = SimpleNamespace( + _vllm_fa2_C=namespace, + load_library=load_library, + ) + monkeypatch.setattr(flash_v100, "torch", SimpleNamespace(ops=fake_ops)) + monkeypatch.setenv("VLLM_SM70_FA2_D256_LIBRARY", "/tmp/stable-fa2.so") + monkeypatch.setattr(flash_v100, "_sm70_splitd_d256_ops_checked", False) + monkeypatch.setattr(flash_v100, "_sm70_splitd_d256_ops", None) + + assert flash_v100._get_sm70_splitd_d256_ops() == ( + "dense", + "paged", + None, + ) + assert loaded == ["/tmp/stable-fa2.so"] + + def test_sm70_d256_gqa_architecture_loader_is_optional(monkeypatch): import vllm.v1.attention.backends.flash_attn_v100 as flash_v100 diff --git a/vllm/envs.py b/vllm/envs.py index 092f02a80e..8fdffd50cc 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -169,8 +169,11 @@ VLLM_SM70_FP8_QPN8_PP2_TP4_SHARED_GATE: bool = False VLLM_SM70_FP8_QPN8_LIBRARY: str | None = None VLLM_SM70_SAMPLER_LIBRARY: str | None = None + VLLM_SM70_FA2_D256_LIBRARY: str | None = None VLLM_SM70_FP8_PREFILL_VISIBLE_DENSE_MM: bool = False VLLM_SM70_NVFP4_QPN2: bool = False + VLLM_SM70_NVFP4_QPN2_PREFILL: bool = False + VLLM_SM70_NVFP4_QPN2_PREFILL_MIN_M: int = 1024 VLLM_SM70_MXFP4_TUNE_SMALL_SHAPES: bool = True VLLM_SM70_NVFP4_TUNE_SMALL_SHAPES: bool = True VLLM_SM70_NVFP4_QWEN38_TP4_M1_FAST_SELECTOR: bool = True @@ -1710,6 +1713,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")) ), @@ -1719,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..0ff88ba608 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py @@ -87,6 +87,7 @@ def _missing_sm70_nvfp4_qpn4_ops() -> list[str]: "nvfp4_qpn2_gated_sm70_out", "nvfp4_qpn2_dispatch_sm70_out", ) +_SM70_NVFP4_QPN2_PREFILL_REQUIRED_OPS = ("nvfp4_qpn4_prefill_sm70_out",) def _is_qpn2_layer(layer: torch.nn.Module) -> bool: @@ -111,6 +112,14 @@ def _missing_qpn2_ops() -> list[str]: ] +def _missing_qpn2_prefill_ops() -> list[str]: + return [ + name + for name in _SM70_NVFP4_QPN2_PREFILL_REQUIRED_OPS + if not hasattr(torch.ops._C, name) + ] + + def _explicit_nvfp4_emulation_requested() -> bool: if envs.VLLM_USE_NVFP4_CT_EMULATIONS or envs.VLLM_NVFP4_GEMM_BACKEND == "emulation": return True @@ -305,6 +314,25 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer.weight.data, layer.weight_scale.data ) qpn2_global_scale = float(layer.weight_global_scale.item()) + qpn2_prefill_workspace = None + if envs.VLLM_SM70_NVFP4_QPN2_PREFILL: + missing_prefill_ops = _missing_qpn2_prefill_ops() + if missing_prefill_ops: + logger.warning_once( + "The requested SM70 NVFP4 QPN2-packed prefill " + "route is unavailable; retaining TurboMind for " + f"large M. Missing ops: {missing_prefill_ops}." + ) + else: + qpn2_prefill_workspace = sm70_tm.get_nvfp4_qpn4_dense_workspace( + layer.weight + ) + if qpn2_prefill_workspace is None: + logger.warning_once( + "Insufficient memory for the bounded SM70 " + "NVFP4 QPN2-packed prefill workspace; " + "retaining TurboMind for large M." + ) use_gated_silu = bool( envs.VLLM_SM70_NVFP4_DENSE_GATED_SILU and is_qpn4_gate and not use_qpn2 @@ -329,10 +357,28 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer.sm70_nvfp4_qpn2_split_k = split_k layer.sm70_nvfp4_qpn2_nacc = nacc layer.sm70_nvfp4_qpn2_gated_silu = suffix == "gate_up_proj" + layer.sm70_nvfp4_qpn2_prefill_dense_weight_ptr = ( + 0 + if qpn2_prefill_workspace is None + else qpn2_prefill_workspace.data_ptr() + ) + if qpn2_prefill_workspace is not None: + # QPN2 and QPN4 use the same physical tile order, but + # expose checkpoint-native [N, K/2] and GEMM-native + # [K, N/2] shapes respectively. Keep zero-copy views so + # prefill does not retain a third multi-GB weight layout. + layer.sm70_nvfp4_qpn2_prefill_codes = qpn2_codes.view(k, n // 2) + layer.sm70_nvfp4_qpn2_prefill_scales = qpn2_scales.view(k // 16, n) logger.info_once( "SM70 NVFP4 QPN2 M<=8 route enabled for a compatible " "TP4 projection contract." ) + if qpn2_prefill_workspace is not None: + logger.info_once( + "SM70 NVFP4 QPN2-packed bounded FP16 prefill route " + "enabled for M>=%d.", + envs.VLLM_SM70_NVFP4_QPN2_PREFILL_MIN_M, + ) elif use_gated_silu: logger.info_once( "SM70 NVFP4 TurboMind gated-SiLU single-layout path enabled." @@ -403,6 +449,26 @@ def _apply_qpn2( ) if x_2d.shape[0] == 0: return out_2d.reshape(*x.shape[:-1], output_size) + prefill_dense_weight_ptr = int( + getattr(layer, "sm70_nvfp4_qpn2_prefill_dense_weight_ptr", 0) + ) + if ( + x_2d.shape[0] >= envs.VLLM_SM70_NVFP4_QPN2_PREFILL_MIN_M + and prefill_dense_weight_ptr + ): + sm70_ops.nvfp4_qpn4_prefill_sm70_out( + out_2d, + prefill_dense_weight_ptr, + x_2d, + layer.sm70_nvfp4_qpn2_prefill_codes, + layer.sm70_nvfp4_qpn2_prefill_scales, + float(layer.sm70_nvfp4_qpn2_global_scale), + True, + gated_silu, + ) + if bias is not None: + out_2d.add_(bias) + return out_2d.reshape(*x.shape[:-1], output_size) state = getattr(layer, sm70_tm.STATE_ATTR) split_k = int(layer.sm70_nvfp4_qpn2_split_k) nacc = int(layer.sm70_nvfp4_qpn2_nacc) diff --git a/vllm/v1/attention/backends/flash_attn_v100.py b/vllm/v1/attention/backends/flash_attn_v100.py index e42b3df8ae..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,8 +1275,30 @@ def _get_sm70_splitd_d256_ops(): _sm70_splitd_d256_ops_checked = True try: - # Importing the interface loads the vendored FA2 torch library. - from vllm.vllm_flash_attn import flash_attn_interface # noqa: F401 + required_ops = ( + "sm70_d256_splitd_n32_dense_fwd", + "sm70_d256_splitd_n32_paged_fwd", + ) + with suppress(ImportError): + # Importing the interface loads the bundled FA2 torch library. + from vllm.vllm_flash_attn import flash_attn_interface # noqa: F401 + + namespace = getattr(torch.ops, "_vllm_fa2_C", None) + if namespace is None or not all( + hasattr(namespace, op_name) for op_name in required_ops + ): + # A partially cached Python interface can import successfully + # without registering its native operators. Source-overlay + # deployments can also intentionally keep the extension outside + # the checkout. In both cases, load only an explicitly selected + # sidecar and then validate the actual operator capability below. + library_path = os.getenv("VLLM_SM70_FA2_D256_LIBRARY") + if library_path is not None: + torch.ops.load_library(library_path) + logger.info( + "Loaded external SM70 D256 prefill library from %s.", + library_path, + ) dense = torch.ops._vllm_fa2_C.sm70_d256_splitd_n32_dense_fwd paged = torch.ops._vllm_fa2_C.sm70_d256_splitd_n32_paged_fwd @@ -1285,7 +1308,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 +1331,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,