From 65f988cad3c53959fc8e6c999673802d8208eae6 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 3 Sep 2026 22:12:50 +0800 Subject: [PATCH 1/9] [Kernel][SM70] Extend NVFP4 QPN2 verifier batches Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../kernels/benchmark_sm70_nvfp4_qpn2.py | 144 ++++++++++++++++-- csrc/sm70_turbomind/ops/nvfp4_qpn2_sm70.cu | 47 ++++-- vllm/_sm70_ops.py | 2 +- .../schemes/compressed_tensors_w4a4_nvfp4.py | 4 +- 4 files changed, 167 insertions(+), 30 deletions(-) diff --git a/benchmarks/kernels/benchmark_sm70_nvfp4_qpn2.py b/benchmarks/kernels/benchmark_sm70_nvfp4_qpn2.py index 4393434cc6..69ba8b3850 100644 --- a/benchmarks/kernels/benchmark_sm70_nvfp4_qpn2.py +++ b/benchmarks/kernels/benchmark_sm70_nvfp4_qpn2.py @@ -4,9 +4,9 @@ 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. +the current TurboMind layout and the QPN2 fragment layout, and measures a +CUDA-graph replay at verifier batch shapes through M=64. 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 @@ -44,6 +44,9 @@ class Projection: QPN2_CONFIGS = { # (K, N): (split K, independent accumulator chains) + (1536, 5120): (8, 2), + (5120, 3584): (16, 2), + (5120, 4128): (16, 2), (5120, 8704): (8, 2), (4352, 5120): (16, 2), } @@ -64,6 +67,9 @@ def _load_projection_shards( tp_size: int, ) -> tuple[Projection, Projection]: path = model / "model.safetensors" + config = json.loads((model / "config.json").read_text()) + text_config = config.get("text_config", config) + num_hidden_layers = int(text_config["num_hidden_layers"]) prefix = f"model.language_model.layers.{layer_index}.mlp" intermediate_size = 17408 hidden_size = 5120 @@ -111,7 +117,7 @@ def _load_projection_shards( 1.0 / torch.cat((gate_global, up_global)).max().float() ), gated_silu=True, - calls_per_round=56, + calls_per_round=num_hidden_layers, ) down = Projection( name="down_proj", @@ -119,11 +125,40 @@ def _load_projection_shards( scales=down_scales.contiguous(), inverse_global_scale=float(1.0 / down_global.max().float()), gated_silu=False, - calls_per_round=56, + calls_per_round=num_hidden_layers, ) return gate_up, down +def _synthetic_projections(model: Path) -> tuple[Projection, ...]: + """Build deterministic tensors for target shapes absent from local weights.""" + config = json.loads((model / "config.json").read_text()) + text_config = config.get("text_config", config) + layer_types = text_config["layer_types"] + num_attention_layers = layer_types.count("full_attention") + num_linear_attention_layers = layer_types.count("linear_attention") + shapes = ( + ("attention_out_proj", 1536, 5120, len(layer_types)), + ("attention_qkv_proj", 5120, 3584, num_attention_layers), + ("gdn_in_proj_qkvz", 5120, 4128, num_linear_attention_layers), + ) + projections = [] + for name, k, n, calls_per_round in shapes: + packed = torch.randint(0, 256, (n, k // 2), dtype=torch.uint8) + scales = (torch.rand((n, k // 16)) * 0.5 + 0.25).to(torch.float8_e4m3fn) + projections.append( + Projection( + name=name, + packed=packed, + scales=scales, + inverse_global_scale=0.01, + gated_silu=False, + calls_per_round=calls_per_round, + ) + ) + return tuple(projections) + + def _unpack_nibbles(weight_packed: torch.Tensor) -> torch.Tensor: """Convert checkpoint [N,K/2] bytes to TurboMind's [K,N] codes.""" return ( @@ -262,6 +297,7 @@ def _run_projection( projection: Projection, extension: Any | None, production: bool, + production_source_candidate: bool, m: int, device: torch.device, warmup: int, @@ -271,7 +307,6 @@ def _run_projection( 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() @@ -334,9 +369,38 @@ def run_qpn2() -> torch.Tensor: ) return qpn_final + elif production_source_candidate: + qpn_codes, qpn_scales = torch.ops._qpn2_candidate.prepare( + packed, projection.scales.to(device) + ) + + def run_qpn2() -> torch.Tensor: + if projection.gated_silu: + torch.ops._qpn2_candidate.gated( + qpn_final, + x, + qpn_codes, + qpn_scales, + projection.inverse_global_scale, + split_k, + accumulator_chains, + ) + else: + torch.ops._qpn2_candidate.gemm( + 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") + scale_codes = projection.scales.view(torch.uint8).to(device) qpn_codes, qpn_scales = _qpn2_prepack(packed, scale_codes) def run_qpn2() -> torch.Tensor: @@ -397,7 +461,17 @@ def _parse_args() -> argparse.Namespace: 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( + "--synthetic-shapes", + action="store_true", + help="Race deterministic tensors for non-MLP Qwen3.8 projection shapes.", + ) parser.add_argument("--qpn2-source", type=Path) + parser.add_argument( + "--production-source-candidate", + type=Path, + help="Compile this repository's QPN2 source under a private namespace.", + ) parser.add_argument( "--production-library", type=Path, @@ -409,6 +483,9 @@ def _parse_args() -> argparse.Namespace: def _load_production_library(path: Path) -> None: + stable_path = path.with_name("_C_stable_libtorch.abi3.so") + if stable_path.exists(): + torch.ops.load_library(stable_path) 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}") @@ -425,15 +502,30 @@ def main() -> int: 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: + if not 1 <= args.m <= 64: + raise ValueError("QPN2 benchmark supports M in [1, 64]") + if args.production_library is not None: _load_production_library(args.production_library) - extension = None - else: - if args.qpn2_source is None: - raise ValueError("--qpn2-source is required without --production-library") + if args.qpn2_source is not None and args.production_source_candidate is not None: + raise ValueError("QPN2 source modes are mutually exclusive") + production_source_candidate = args.production_source_candidate is not None + extension = None + if production_source_candidate: + load( + name=args.extension_name, + sources=[str(args.production_source_candidate)], + extra_cuda_cflags=[ + "-O3", + "--use_fast_math", + "-lineinfo", + "-gencode=arch=compute_70,code=sm_70", + "-DVLLM_NVFP4_QPN2_STANDALONE", + "-DVLLM_NVFP4_QPN2_BENCHMARK_CANDIDATE", + ], + verbose=False, + is_python_module=False, + ) + elif args.qpn2_source is not None: extension = load( name=args.extension_name, sources=[str(args.qpn2_source)], @@ -445,15 +537,24 @@ def main() -> int: ], verbose=False, ) + elif args.production_library is None: + raise ValueError( + "one of --qpn2-source, --production-source-candidate, or " + "--production-library is required" + ) + production = not production_source_candidate and args.qpn2_source is None torch.manual_seed(20260824) - projections = _load_projection_shards( - args.model, args.layer, args.tp_rank, args.tp_size + projections = ( + _synthetic_projections(args.model) + if args.synthetic_shapes + else _load_projection_shards(args.model, args.layer, args.tp_rank, args.tp_size) ) rows = [ _run_projection( projection, extension, production, + production_source_candidate, args.m, device, args.warmup, @@ -471,10 +572,21 @@ def main() -> int: "layer": args.layer, "tp_rank": args.tp_rank, "tp_size": args.tp_size, + "synthetic_shapes": args.synthetic_shapes, "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_source_candidate": ( + str(args.production_source_candidate) + if args.production_source_candidate + else None + ), + "production_source_candidate_sha256": ( + _sha256_file(args.production_source_candidate) + if args.production_source_candidate + else None + ), "production_library": ( str(args.production_library) if args.production_library else None ), diff --git a/csrc/sm70_turbomind/ops/nvfp4_qpn2_sm70.cu b/csrc/sm70_turbomind/ops/nvfp4_qpn2_sm70.cu index 641949e4b4..e9ea30ff01 100644 --- a/csrc/sm70_turbomind/ops/nvfp4_qpn2_sm70.cu +++ b/csrc/sm70_turbomind/ops/nvfp4_qpn2_sm70.cu @@ -25,6 +25,9 @@ void nvfp4_gemm_sm70_out(torch::Tensor out, torch::Tensor input, namespace { constexpr int kPrepareThreads = 256; +constexpr int kQpn2RowsPerCta = 8; +constexpr int kQpn2MaxRows = 64; +constexpr int kQpn2DispatchMaxRows = 32; __device__ __forceinline__ int qpn2_col_from_lane(int lane) { return ((lane >> 2) & 3) * 8 + (lane & 3) + ((lane & 16) ? 4 : 0); @@ -131,7 +134,8 @@ __global__ void nvfp4_qpn2_sm70_kernel(const uint8_t* __restrict__ codes, 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 local_row = (lane & 3) + ((lane & 16) ? 4 : 0); + const int row = blockIdx.y * kQpn2RowsPerCta + local_row; const int groups_k16 = k >> 4; const int groups_per_warp = groups_k16 / SplitK; const int group_begin = warp * groups_per_warp; @@ -201,7 +205,7 @@ __global__ void nvfp4_qpn2_sm70_kernel(const uint8_t* __restrict__ codes, for (int k_warp = 0; k_warp < SplitK; ++k_warp) { value += partials[k_warp][element]; } - const int output_row = element >> 5; + const int output_row = blockIdx.y * kQpn2RowsPerCta + (element >> 5); const int output_col = element & 31; if (output_row < m) { output[static_cast(output_row) * n + tile * 32 + output_col] = @@ -224,7 +228,8 @@ __global__ void nvfp4_qpn2_gated_sm70_kernel( 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 local_row = (lane & 3) + ((lane & 16) ? 4 : 0); + const int row = blockIdx.y * kQpn2RowsPerCta + local_row; const int groups_k16 = k >> 4; const int groups_per_warp = groups_k16 / SplitK; const int group_begin = warp * groups_per_warp; @@ -296,7 +301,7 @@ __global__ void nvfp4_qpn2_gated_sm70_kernel( gate += partials[0][k_warp][element]; up += partials[1][k_warp][element]; } - const int output_row = element >> 5; + const int output_row = blockIdx.y * kQpn2RowsPerCta + (element >> 5); const int output_col = element & 31; if (output_row < m) { // Match the existing SM70 silu_and_mul contract: round both GEMM @@ -317,7 +322,8 @@ 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>>>( + const dim3 grid(n / 32, (m + kQpn2RowsPerCta - 1) / kQpn2RowsPerCta); + nvfp4_qpn2_sm70_kernel<<>>( codes, scales, input, output, n, k, m, global_scale); } @@ -325,9 +331,10 @@ 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) { + const dim3 grid(hidden / 32, (m + kQpn2RowsPerCta - 1) / kQpn2RowsPerCta); nvfp4_qpn2_gated_sm70_kernel - <<<(hidden / 32), (64 * SplitK), 0, stream>>>( - codes, scales, input, output, hidden, k, m, global_scale); + <<>>(codes, scales, input, output, hidden, + k, m, global_scale); } void check_qpn2_tensors(const torch::Tensor& out, const torch::Tensor& input, @@ -353,8 +360,8 @@ void check_qpn2_tensors(const torch::Tensor& out, const torch::Tensor& input, 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(m >= 1 && m <= kQpn2MaxRows && out.size(0) == m, + "NVFP4 QPN2 requires M in [1, ", kQpn2MaxRows, "]"); 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, @@ -498,7 +505,7 @@ void nvfp4_qpn2_dispatch_sm70_out(torch::Tensor out, torch::Tensor input, 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 (input.size(0) <= kQpn2DispatchMaxRows) { if (gated_silu) { nvfp4_qpn2_gated_sm70_out(out, input, codes, scales, global_scale, split_k, accumulator_chains); @@ -522,7 +529,8 @@ void nvfp4_qpn2_dispatch_sm70_out(torch::Tensor out, torch::Tensor input, } #endif -#ifdef VLLM_NVFP4_QPN2_STANDALONE +#if defined(VLLM_NVFP4_QPN2_STANDALONE) && \ + !defined(VLLM_NVFP4_QPN2_BENCHMARK_CANDIDATE) // 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. @@ -544,3 +552,20 @@ TORCH_LIBRARY_FRAGMENT(_C, ops) { &nvfp4_qpn2_gated_sm70_out); } #endif + +#ifdef VLLM_NVFP4_QPN2_BENCHMARK_CANDIDATE +// Register a private namespace so the extended-M candidate can race the +// installed production operators without replacing vllm._C. +TORCH_LIBRARY_FRAGMENT(_qpn2_candidate, ops) { + ops.def("prepare(Tensor weight_packed, Tensor weight_scale) -> Tensor[]"); + ops.impl("prepare", torch::kCUDA, &nvfp4_qpn2_prepare_sm70); + ops.def( + "gemm(Tensor(a!) out, Tensor input, Tensor codes, Tensor scales, " + "float global_scale, int split_k, int accumulator_chains) -> ()"); + ops.impl("gemm", torch::kCUDA, &nvfp4_qpn2_gemm_sm70_out); + ops.def( + "gated(Tensor(a!) out, Tensor input, Tensor codes, Tensor scales, " + "float global_scale, int split_k, int accumulator_chains) -> ()"); + ops.impl("gated", torch::kCUDA, &nvfp4_qpn2_gated_sm70_out); +} +#endif diff --git a/vllm/_sm70_ops.py b/vllm/_sm70_ops.py index 6fc4440ca1..42941be1e0 100644 --- a/vllm/_sm70_ops.py +++ b/vllm/_sm70_ops.py @@ -1249,7 +1249,7 @@ def nvfp4_qpn2_dispatch_sm70_out( tm_q_ld: int, gated_silu: bool, ) -> None: - """Select QPN2 for M<=8 and TurboMind for larger dynamic M.""" + """Select QPN2 for M<=32 and TurboMind for larger dynamic M.""" _op("nvfp4_qpn2_dispatch_sm70_out")( out, input, 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 478b045096..d57c134d89 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 @@ -56,7 +56,7 @@ def _is_sm70_dflash2_nvfp4_qpn2_runtime_contract() -> bool: """Admit the quality-audited DFlash2 TP4 operator contract. Scheduler capacity is intentionally not part of this model-load decision. - The opaque dispatcher selects QPN2 only from the live ``M <= 8`` shape and + The opaque dispatcher selects QPN2 only from live ``M <= 32`` shapes and retains the existing TurboMind path for larger dynamic M. A server that can hold many requests must therefore load the same small-M layout as a server configured with ``max_num_seqs=1``. @@ -425,7 +425,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer.sm70_nvfp4_qpn2_gated_silu = suffix == "gate_up_proj" layer.sm70_nvfp4_qpn2_prefill_enabled = qpn2_prefill_enabled logger.info_once( - "SM70 NVFP4 QPN2 M<=8 route enabled for a compatible " + "SM70 NVFP4 QPN2 M<=32 route enabled for a compatible " "TP4 projection contract." ) if qpn2_prefill_enabled: From 028bca7e14a06671ec8fb2db276ddebcfa95d38f Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 3 Sep 2026 22:13:04 +0800 Subject: [PATCH 2/9] [Kernel][SM70] Batch exact DFlash2 grouped verifier Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- ...ark_sm70_dflash2_batched_grouped_verify.py | 174 ++++++++++++++++++ .../flash_attn_v100/__init__.py | 2 + .../flash_attn_v100/flash_attn_interface.py | 55 ++++-- flash-attention-v100/include/fused_mha.h | 2 + .../kernel/flash_decode_paged.cu | 92 ++++++--- flash-attention-v100/kernel/fused_mha_api.cpp | 3 + .../test_sm70_flash_v100_grouped_verify.py | 152 +++++++++++++++ .../attention/test_sm70_flash_v100_policy.py | 122 ++++++++++++ vllm/envs.py | 6 + vllm/v1/attention/backends/flash_attn_v100.py | 67 ++++++- 10 files changed, 624 insertions(+), 51 deletions(-) create mode 100644 benchmarks/benchmark_sm70_dflash2_batched_grouped_verify.py diff --git a/benchmarks/benchmark_sm70_dflash2_batched_grouped_verify.py b/benchmarks/benchmark_sm70_dflash2_batched_grouped_verify.py new file mode 100644 index 0000000000..38b1a7053e --- /dev/null +++ b/benchmarks/benchmark_sm70_dflash2_batched_grouped_verify.py @@ -0,0 +1,174 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Race request-major grouped DFlash2 verification against independent XQA.""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path + +import torch + + +def _make_case( + *, batch_size: int, seq_len: int, page_size: int +) -> tuple[torch.Tensor, ...]: + query_len = 8 + seq_lens = torch.tensor( + [seq_len - (req_idx % 3) * 17 for req_idx in range(batch_size)], + dtype=torch.int32, + device="cuda", + ) + max_pages = math.ceil(seq_len / page_size) + physical_pages = batch_size * max_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) + del source + key_cache, value_cache = cache.unbind(1) + block_table = torch.randperm(physical_pages, dtype=torch.int32, device="cuda")[ + : batch_size * max_pages + ].view(batch_size, max_pages) + query = torch.randn( + (batch_size * query_len, 6, 256), + dtype=torch.float16, + device="cuda", + ).mul_(0.25) + return query, key_cache, value_cache, block_table, seq_lens + + +def _measure_ms(fn, *, warmups: int, repeats: int) -> float: + for _ in range(warmups): + 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 float(start.elapsed_time(end)) / repeats + + +def _capture(fn) -> 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("--batch-size", type=int, choices=(1, 2, 4, 8), required=True) + parser.add_argument("--seq-len", type=int, required=True) + parser.add_argument("--page-size", type=int, default=3296) + parser.add_argument("--warmups", type=int, default=20) + parser.add_argument("--repeats", type=int, default=100) + parser.add_argument("--json-out", type=Path) + args = parser.parse_args() + if args.seq_len < 42: + parser.error("--seq-len must be at least 42 for the varied batch case") + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0): + raise RuntimeError("This benchmark requires one SM70 GPU") + + import flash_attn_v100 + + torch.manual_seed(20260903 + args.batch_size + args.seq_len) + query, key_cache, value_cache, block_table, seq_lens = _make_case( + batch_size=args.batch_size, + seq_len=args.seq_len, + page_size=args.page_size, + ) + grouped_out = torch.empty_like(query) + xqa_out = torch.empty_like(query) + query_len = 8 + decode_block_table = block_table.repeat_interleave(query_len, dim=0).contiguous() + decode_seq_lens = ( + seq_lens[:, None] + - query_len + + torch.arange(1, query_len + 1, dtype=torch.int32, device="cuda") + ).flatten() + + def grouped() -> None: + flash_attn_v100.flash_attn_grouped_verify_paged( + query, + key_cache, + value_cache, + block_table, + seq_lens, + out=grouped_out, + one_pass=True, + ) + + def xqa() -> None: + flash_attn_v100.flash_attn_decode_paged_xqa( + query, + key_cache, + value_cache, + decode_block_table, + decode_seq_lens, + out=xqa_out, + kv_cache_dtype="fp8_e5m2", + max_seq_len_hint=args.seq_len, + workspace_seq_capacity_hint=args.seq_len, + ) + + grouped() + xqa() + per_request = torch.cat( + [ + flash_attn_v100.flash_attn_grouped_verify_paged( + query[req_idx * query_len : (req_idx + 1) * query_len], + key_cache, + value_cache, + block_table[req_idx : req_idx + 1], + seq_lens[req_idx : req_idx + 1], + one_pass=True, + ).clone() + for req_idx in range(args.batch_size) + ] + ) + torch.accelerator.synchronize() + grouped_vs_xqa = grouped_out.float().sub(xqa_out.float()).abs() + + grouped_eager_ms = _measure_ms(grouped, warmups=args.warmups, repeats=args.repeats) + xqa_eager_ms = _measure_ms(xqa, warmups=args.warmups, repeats=args.repeats) + grouped_graph = _capture(grouped) + xqa_graph = _capture(xqa) + grouped_graph_ms = _measure_ms( + grouped_graph.replay, warmups=args.warmups, repeats=args.repeats + ) + xqa_graph_ms = _measure_ms( + xqa_graph.replay, warmups=args.warmups, repeats=args.repeats + ) + + result = { + "batch_size": args.batch_size, + "seq_len": args.seq_len, + "page_size": args.page_size, + "batched_is_bitwise_per_request": bool(torch.equal(grouped_out, per_request)), + "grouped_vs_xqa_max_abs": float(grouped_vs_xqa.max().item()), + "grouped_vs_xqa_mean_abs": float(grouped_vs_xqa.mean().item()), + "grouped_eager_ms": grouped_eager_ms, + "xqa_eager_ms": xqa_eager_ms, + "grouped_graph_ms": grouped_graph_ms, + "xqa_graph_ms": xqa_graph_ms, + "graph_speedup": xqa_graph_ms / grouped_graph_ms, + } + 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") + + +if __name__ == "__main__": + main() diff --git a/flash-attention-v100/flash_attn_v100/__init__.py b/flash-attention-v100/flash_attn_v100/__init__.py index 9acd9962c0..e0557019fe 100644 --- a/flash-attention-v100/flash_attn_v100/__init__.py +++ b/flash-attention-v100/flash_attn_v100/__init__.py @@ -12,6 +12,7 @@ flash_attn_func, flash_attn_grouped_verify_max_query_tokens, flash_attn_grouped_verify_paged, + flash_attn_grouped_verify_request_major_abi_version, flash_attn_lse, flash_attn_prefill_paged, flash_attn_prefill_paged_bfla, @@ -34,6 +35,7 @@ "flash_attn_decode_paged_wmma", "flash_attn_grouped_verify_max_query_tokens", "flash_attn_grouped_verify_paged", + "flash_attn_grouped_verify_request_major_abi_version", "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 433cd1fee7..a48d14f288 100644 --- a/flash-attention-v100/flash_attn_v100/flash_attn_interface.py +++ b/flash-attention-v100/flash_attn_v100/flash_attn_interface.py @@ -503,26 +503,44 @@ def _get_prefill_splitkv3_workspace( return workspace -def _get_grouped_verify_workspace(q: torch.Tensor) -> _GroupedVerifyWorkspace: +def _get_grouped_verify_workspace( + q: torch.Tensor, + batch_size: int, +) -> _GroupedVerifyWorkspace: + query_len = q.shape[0] // batch_size + max_query_tokens = 16 if query_len > 8 else 8 + grouped_splits = 640 // max_query_tokens 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), + batch_size, + max_query_tokens, q.dtype, ) workspace = ( _grouped_verify_workspace_cache.get(key) if _can_cache_workspace(q) else None ) if workspace is None: + partial_out_shape = ( + (grouped_splits, max_query_tokens, 6, 256) + if batch_size == 1 + else (batch_size, grouped_splits, max_query_tokens, 6, 256) + ) + partial_lse_shape = ( + (grouped_splits, max_query_tokens, 6) + if batch_size == 1 + else (batch_size, grouped_splits, max_query_tokens, 6) + ) workspace = _GroupedVerifyWorkspace( partial_out=torch.empty( - (80, 8, 6, 256), + partial_out_shape, dtype=torch.float16, device=q.device, ), partial_lse=torch.empty( - (80, 8, 6), + partial_lse_shape, dtype=torch.float32, device=q.device, ), @@ -1040,6 +1058,16 @@ def flash_attn_grouped_verify_max_query_tokens() -> int: return int(get_max_query_tokens()) +def flash_attn_grouped_verify_request_major_abi_version() -> int: + """Return zero for binaries that only support single-request grouping.""" + get_abi_version = getattr( + flash_attn_v100_cuda, + "grouped_verify_request_major_abi_version", + None, + ) + return 0 if get_abi_version is None else int(get_abi_version()) + + def flash_attn_grouped_verify_paged( q: torch.Tensor, k_cache: torch.Tensor, @@ -1053,12 +1081,13 @@ def flash_attn_grouped_verify_paged( v_scale: float = 1.0, one_pass: bool = False, ) -> torch.Tensor: - """Exact grouped q8/q16 H6/D256 DFlash2 verifier for SM70. + """Exact request-major grouped q8/q16 H6/D256 DFlash2 verifier for SM70. + Each request has one block-table row and a uniform contiguous query span. The native entry keeps all causal verifier rows together and reuses each - 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. + paged-KV scan across a packed GQA group. Single-request q16 uses two + three-head groups; batched requests use request-major q8 groups. Workspaces + are stream- and batch-local and CUDA-graph safe. """ if softmax_scale is None: softmax_scale = q.shape[-1] ** -0.5 @@ -1066,13 +1095,7 @@ def flash_attn_grouped_verify_paged( block_table = maybe_contiguous(block_table) 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 + workspace = _get_grouped_verify_workspace(q, int(block_table.shape[0])) return flash_attn_v100_cuda.grouped_verify_paged_fwd( q, k_cache, @@ -1080,8 +1103,8 @@ def flash_attn_grouped_verify_paged( out, block_table, seq_lens, - partial_out, - partial_lse, + workspace.partial_out, + workspace.partial_lse, float(softmax_scale), kv_cache_dtype, float(k_scale), diff --git a/flash-attention-v100/include/fused_mha.h b/flash-attention-v100/include/fused_mha.h index 4004701800..de71125c76 100644 --- a/flash-attention-v100/include/fused_mha.h +++ b/flash-attention-v100/include/fused_mha.h @@ -59,6 +59,8 @@ at::Tensor flash_attention_grouped_verify_paged( int64_t flash_attention_grouped_verify_max_query_tokens(); +int64_t flash_attention_grouped_verify_request_major_abi_version(); + int64_t flash_attention_grouped_sparse_page4_abi_version(); at::Tensor flash_attention_grouped_sparse_page4( diff --git a/flash-attention-v100/kernel/flash_decode_paged.cu b/flash-attention-v100/kernel/flash_decode_paged.cu index 7384538f75..757a261a26 100644 --- a/flash-attention-v100/kernel/flash_decode_paged.cu +++ b/flash-attention-v100/kernel/flash_decode_paged.cu @@ -1980,14 +1980,22 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ using Traits = GroupedVerifyTraits; const int head_group = blockIdx.x; const int split_id = blockIdx.y; - const int group_idx = SPARSE_PAGE4 ? blockIdx.z : 0; + const int group_idx = blockIdx.z; if (head_group >= Traits::kHeadGroups || split_id >= Traits::kSplits || group_idx >= num_groups || query_len <= 0 || query_len > MAX_QUERY_TOKENS) { return; } - const int total_kv = seq_lens[SPARSE_PAGE4 ? group_idx : 0]; + if constexpr (!SPARSE_PAGE4) { + partial_out += static_cast(group_idx) * Traits::kSplits * + MAX_QUERY_TOKENS * kGroupedVerifyHeads * + kGroupedVerifyHeadDim; + partial_lse += static_cast(group_idx) * Traits::kSplits * + MAX_QUERY_TOKENS * kGroupedVerifyHeads; + } + + const int total_kv = seq_lens[group_idx]; if (total_kv <= 0) { if constexpr (SPARSE_PAGE4) { constexpr int kGroupOutputElements = @@ -2054,7 +2062,7 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ 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]); + smem.page_ids[idx] = __ldg(&page_ids[split_start_page + idx]); } } if (use_staged_page_ids) { @@ -2437,9 +2445,17 @@ __launch_bounds__(kGroupedVerifyThreads) void flash_attention_grouped_verify_e5m using Traits = GroupedVerifyTraits; const int token_idx = blockIdx.x; const int head_idx = blockIdx.y; + const int request_idx = blockIdx.z; if (token_idx >= query_len || head_idx >= kGroupedVerifyHeads) { return; } + partial_out += static_cast(request_idx) * Traits::kSplits * + MAX_QUERY_TOKENS * kGroupedVerifyHeads * kGroupedVerifyHeadDim; + partial_lse += static_cast(request_idx) * Traits::kSplits * + MAX_QUERY_TOKENS * kGroupedVerifyHeads; + seq_lens += request_idx; + out += static_cast(request_idx) * query_len * kGroupedVerifyHeads * + kGroupedVerifyHeadDim; const int active_splits = grouped_verify_active_splits(seq_lens[0]); __shared__ float split_lse[Traits::kSplits]; @@ -4034,6 +4050,8 @@ int64_t flash_attention_grouped_verify_max_query_tokens() { return kGroupedVerifyMaxSupportedQ; } +int64_t flash_attention_grouped_verify_request_major_abi_version() { return 1; } + int64_t flash_attention_grouped_sparse_page4_abi_version() { // Version 1 accepted FP16 K/V through the nine-argument forward binding. // Version 2 adds kv_cache_dtype and calibrated K/V scales. @@ -4065,12 +4083,19 @@ at::Tensor flash_attention_grouped_verify_paged( 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) <= kGroupedVerifyMaxSupportedQ && + TORCH_CHECK(block_table.dim() == 2 && block_table.size(0) > 0, + "grouped verify block_table must have shape [batch, blocks]"); + const int64_t batch_size = block_table.size(0); + TORCH_CHECK(q.dim() == 3 && q.size(0) > 0 && q.size(0) % batch_size == 0 && q.size(1) == kGroupedVerifyHeads && q.size(2) == kGroupedVerifyHeadDim, - "grouped verify q must have shape [1..16, 6, 256]"); - const bool wide_query = q.size(0) > kGroupedVerifyQ8MaxQ; + "grouped verify q must have shape [batch * query, 6, 256]"); + const int64_t query_len = q.size(0) / batch_size; + TORCH_CHECK( + batch_size == 1 ? query_len <= kGroupedVerifyMaxSupportedQ + : query_len == kGroupedVerifyQ8MaxQ, + "grouped verify requires q1..q16 for B1 or request-major q8 for batches"); + const bool wide_query = query_len > kGroupedVerifyQ8MaxQ; const int max_query_tokens = wide_query ? kGroupedVerifyQ16MaxQ : kGroupedVerifyQ8MaxQ; const int grouped_splits = kGroupedVerifyWorkspaceRows / max_query_tokens; @@ -4081,10 +4106,8 @@ at::Tensor flash_attention_grouped_verify_paged( 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(seq_lens.dim() == 1 && seq_lens.size(0) == batch_size, + "grouped verify seq_lens must cover every request"); TORCH_CHECK(q.is_contiguous(), "grouped verify q must be contiguous [M, H, D]"); TORCH_CHECK(block_table.is_contiguous() && seq_lens.is_contiguous(), @@ -4093,15 +4116,31 @@ at::Tensor flash_attention_grouped_verify_paged( "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({grouped_splits, max_query_tokens, - kGroupedVerifyHeads, kGroupedVerifyHeadDim}), - "partial_out must have shape [80, 8, 6, 256] or " - "[40, 16, 6, 256]"); - TORCH_CHECK( + const bool single_out = + batch_size == 1 && + partial_out.sizes() == + at::IntArrayRef({grouped_splits, max_query_tokens, + kGroupedVerifyHeads, kGroupedVerifyHeadDim}); + const bool batched_out = + batch_size > 1 && + partial_out.sizes() == + at::IntArrayRef({batch_size, grouped_splits, max_query_tokens, + kGroupedVerifyHeads, kGroupedVerifyHeadDim}); + TORCH_CHECK(single_out || batched_out, + "partial_out must be [splits, query, 6, 256] for B1 or " + "[batch, splits, query, 6, 256]"); + const bool single_lse = + batch_size == 1 && partial_lse.sizes() == at::IntArrayRef({grouped_splits, max_query_tokens, - kGroupedVerifyHeads}), - "partial_lse must have shape [80, 8, 6] or [40, 16, 6]"); + kGroupedVerifyHeads}); + const bool batched_lse = + batch_size > 1 && + partial_lse.sizes() == + at::IntArrayRef({batch_size, grouped_splits, max_query_tokens, + kGroupedVerifyHeads}); + TORCH_CHECK(single_lse || batched_lse, + "partial_lse must be [splits, query, 6] for B1 or " + "[batch, splits, query, 6]"); TORCH_CHECK(k_scale > 0.0f && v_scale > 0.0f, "grouped verify E5M2 K/V scales must be positive"); @@ -4124,7 +4163,8 @@ at::Tensor flash_attention_grouped_verify_paged( "grouped verify prototype supports SM70 only"); cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); - const dim3 partial_grid(head_groups, grouped_splits, 1); + const dim3 partial_grid(head_groups, grouped_splits, + static_cast(batch_size)); const size_t partial_shared_mem = sizeof(GroupedVerifySmem); #define LAUNCH_GROUPED_VERIFY_PARTIAL(MAX_QUERY_TOKENS, TWO_PASS, PAGE_SIZE, \ SINGLE_QUERY, CONTIGUOUS_LAYOUT, \ @@ -4153,12 +4193,12 @@ at::Tensor flash_attention_grouped_verify_paged( 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)), \ + partial_lse.data_ptr(), static_cast(query_len), \ 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, nullptr, 1); \ + v_scale, nullptr, static_cast(batch_size)); \ } while (0) #define DISPATCH_GROUPED_VERIFY_PARTIAL(MAX_QUERY_TOKENS, TWO_PASS, \ @@ -4201,7 +4241,7 @@ at::Tensor flash_attention_grouped_verify_paged( } \ } while (0) - const bool single_query = q.size(0) == 1; + const bool single_query = query_len == 1; if (wide_query && one_pass) { DISPATCH_GROUPED_VERIFY_PARTIAL(kGroupedVerifyQ16MaxQ, false, false); } else if (wide_query) { @@ -4217,8 +4257,8 @@ at::Tensor flash_attention_grouped_verify_paged( } #undef DISPATCH_GROUPED_VERIFY_PARTIAL #undef LAUNCH_GROUPED_VERIFY_PARTIAL - const dim3 combine_grid(static_cast(q.size(0)), kGroupedVerifyHeads, - 1); + const dim3 combine_grid(static_cast(query_len), kGroupedVerifyHeads, + static_cast(batch_size)); #define LAUNCH_GROUPED_VERIFY_COMBINE(MAX_QUERY_TOKENS, SINGLE_QUERY) \ flash_attention_grouped_verify_e5m2_combine_kernel \ @@ -4226,7 +4266,7 @@ at::Tensor flash_attention_grouped_verify_paged( 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))) + static_cast(query_len)) if (wide_query) { LAUNCH_GROUPED_VERIFY_COMBINE(kGroupedVerifyQ16MaxQ, false); } else if (single_query) { diff --git a/flash-attention-v100/kernel/fused_mha_api.cpp b/flash-attention-v100/kernel/fused_mha_api.cpp index 76859871ac..4ef81fc58a 100644 --- a/flash-attention-v100/kernel/fused_mha_api.cpp +++ b/flash-attention-v100/kernel/fused_mha_api.cpp @@ -27,6 +27,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("grouped_verify_max_query_tokens", &flash_attention_grouped_verify_max_query_tokens, "Maximum query length supported by grouped DFlash2 verification"); + m.def("grouped_verify_request_major_abi_version", + &flash_attention_grouped_verify_request_major_abi_version, + "Request-major grouped DFlash2 forward ABI version"); m.def("grouped_sparse_page4_fwd", &flash_attention_grouped_sparse_page4, "Grouped exact QSA page4 attention over paged KV cache (Volta)"); m.def("grouped_sparse_page4_abi_version", 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 c895d4f040..9ec98526e8 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,52 @@ def _make_case( return query, key_cache, value_cache, block_table, seq_lens +def _make_batched_case( + *, + page_size: int, + batch_size: int, + query_len: int, + prefix_len: int, + interleaved: bool, +) -> tuple[torch.Tensor, ...]: + seq_lens = torch.tensor( + [prefix_len + (req_idx % 3) * 17 + query_len for req_idx in range(batch_size)], + dtype=torch.int32, + device="cuda", + ) + max_logical_pages = math.ceil(int(seq_lens.max().item()) / page_size) + physical_pages = batch_size * max_logical_pages + 3 + if interleaved: + 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) + else: + 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", + )[: batch_size * max_logical_pages].view(batch_size, max_logical_pages) + query = torch.randn( + (batch_size * query_len, 6, 256), + dtype=torch.float16, + device="cuda", + ).mul_(0.25) + return query, key_cache, value_cache, block_table, seq_lens + + def _make_interleaved_case( *, page_size: int, @@ -116,6 +162,49 @@ def _reference( return torch.stack(rows).half() +@pytest.mark.parametrize("interleaved", [False, True]) +@pytest.mark.parametrize("batch_size", [2, 4, 8]) +@torch.inference_mode() +def test_batched_grouped_verify_is_bitwise_per_request( + batch_size: int, + interleaved: bool, +) -> None: + flash_attn_v100 = _require_grouped_verify() + torch.manual_seed(20260903 + batch_size) + query, key_cache, value_cache, block_table, seq_lens = _make_batched_case( + page_size=1648, + batch_size=batch_size, + query_len=8, + prefix_len=4097, + interleaved=interleaved, + ) + + batched = flash_attn_v100.flash_attn_grouped_verify_paged( + query, + key_cache, + value_cache, + block_table, + seq_lens, + one_pass=True, + ).clone() + per_request = torch.cat( + [ + flash_attn_v100.flash_attn_grouped_verify_paged( + query[req_idx * 8 : (req_idx + 1) * 8], + key_cache, + value_cache, + block_table[req_idx : req_idx + 1], + seq_lens[req_idx : req_idx + 1], + one_pass=True, + ).clone() + for req_idx in range(batch_size) + ] + ) + torch.accelerator.synchronize() + + assert torch.equal(batched, per_request) + + @pytest.mark.parametrize( ("page_size", "query_len", "prefix_len", "strict_fp32_gate"), [ @@ -278,6 +367,69 @@ def test_grouped_verify_cuda_graph_replay_tracks_runtime_seq_len( assert difference.mean().item() <= 6.0e-6 +@torch.inference_mode() +def test_batched_grouped_verify_cuda_graph_tracks_each_request_seq_len() -> None: + flash_attn_v100 = _require_grouped_verify() + torch.manual_seed(20260903) + batch_size = 4 + query, key_cache, value_cache, block_table, seq_lens = _make_batched_case( + page_size=1648, + batch_size=batch_size, + query_len=8, + prefix_len=4097, + interleaved=True, + ) + 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.copy_( + torch.tensor( + [prefix_len + req_idx * 17 + 8 for req_idx in range(batch_size)], + dtype=torch.int32, + device="cuda", + ) + ) + graph.replay() + torch.accelerator.synchronize() + expected = torch.cat( + [ + flash_attn_v100.flash_attn_grouped_verify_paged( + query[req_idx * 8 : (req_idx + 1) * 8], + key_cache, + value_cache, + block_table[req_idx : req_idx + 1], + seq_lens[req_idx : req_idx + 1], + one_pass=True, + ).clone() + for req_idx in range(batch_size) + ] + ) + torch.accelerator.synchronize() + assert torch.equal(output, expected) + + @pytest.mark.parametrize("page_size", [1648, 3296]) @pytest.mark.parametrize("query_len", [8, 16]) @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 46e1f49fd2..387575ab97 100644 --- a/tests/v1/attention/test_sm70_flash_v100_policy.py +++ b/tests/v1/attention/test_sm70_flash_v100_policy.py @@ -1809,6 +1809,128 @@ def grouped_verify( assert torch.all(output == 1) +def test_flash_v100_batched_grouped_workspace_preserves_single_request_layout(): + from flash_attn_v100.flash_attn_interface import _get_grouped_verify_workspace + + q8 = torch.empty((8, 6, 256), dtype=torch.float16, device="meta") + q16 = torch.empty((16, 6, 256), dtype=torch.float16, device="meta") + batched = torch.empty((32, 6, 256), dtype=torch.float16, device="meta") + + single_q8 = _get_grouped_verify_workspace(q8, 1) + single_q16 = _get_grouped_verify_workspace(q16, 1) + batch_q8 = _get_grouped_verify_workspace(batched, 4) + + assert tuple(single_q8.partial_out.shape) == (80, 8, 6, 256) + assert tuple(single_q8.partial_lse.shape) == (80, 8, 6) + assert tuple(single_q16.partial_out.shape) == (40, 16, 6, 256) + assert tuple(single_q16.partial_lse.shape) == (40, 16, 6) + assert tuple(batch_q8.partial_out.shape) == (4, 80, 8, 6, 256) + assert tuple(batch_q8.partial_lse.shape) == (4, 80, 8, 6) + + +@pytest.mark.parametrize("batch_size", [2, 4, 8]) +def test_flash_v100_dflash2_batched_grouped_verify_uses_exact_requests( + batch_size: int, +): + 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 + impl.dflash2_grouped_verify_max_query_tokens = 16 + captured: dict[str, object] = {} + + def grouped_verify( + query, + key_cache, + value_cache, + block_table, + seq_lens, + **kwargs, + ): + captured["query"] = query + captured["block_table"] = block_table + captured["seq_lens"] = seq_lens + kwargs["out"].fill_(1) + + impl.flash_attn_grouped_verify_paged = grouped_verify + query_start_loc = torch.arange(0, (batch_size + 1) * 8, 8, dtype=torch.int32) + original_block_table = torch.arange(batch_size * 2, dtype=torch.int32).view( + batch_size, 2 + ) + original_seq_lens = torch.arange(2056, 2056 + batch_size, dtype=torch.int32) + attn_metadata = SimpleNamespace( + num_actual_tokens=batch_size * 8, + num_reqs=batch_size, + max_query_len=8, + causal=True, + is_dflash_selector_target=True, + max_model_len=32768, + query_start_loc=query_start_loc, + seq_lens=original_seq_lens, + block_table=original_block_table, + ) + layer = SimpleNamespace(_k_scale_float=0.5, _v_scale_float=2.0) + query = torch.zeros((batch_size * 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) + + impl.use_dflash2_batched_grouped_verify = False + impl.dflash2_grouped_verify_request_major_abi_version = 1 + assert not impl._dflash2_grouped_verify_allowed( + query, + key_cache, + value_cache, + attn_metadata, + num_query_tokens=batch_size * 8, + ) + + impl.use_dflash2_batched_grouped_verify = True + impl.dflash2_grouped_verify_request_major_abi_version = 0 + assert not impl._dflash2_grouped_verify_allowed( + query, + key_cache, + value_cache, + attn_metadata, + num_query_tokens=batch_size * 8, + ) + + impl.dflash2_grouped_verify_request_major_abi_version = 1 + result = impl._flash_v100_small_query_prefill_as_decode( + layer, + query, + key_cache, + value_cache, + attn_metadata, + output, + query_start_loc, + original_seq_lens, + ) + + assert result is output + captured_query = captured["query"] + captured_block_table = captured["block_table"] + captured_seq_lens = captured["seq_lens"] + assert isinstance(captured_query, torch.Tensor) + assert isinstance(captured_block_table, torch.Tensor) + assert isinstance(captured_seq_lens, torch.Tensor) + assert captured_query.data_ptr() == query.data_ptr() + assert captured_block_table.data_ptr() == original_block_table.data_ptr() + assert captured_seq_lens.data_ptr() == original_seq_lens.data_ptr() + assert tuple(captured_query.shape) == (batch_size * 8, 6, 256) + assert tuple(captured_block_table.shape) == (batch_size, 2) + assert tuple(captured_seq_lens.shape) == (batch_size,) + assert torch.all(output == 1) + + def test_flash_v100_dflash2_q16_falls_back_for_q8_native_binary(): from vllm.v1.attention.backends.flash_attn_v100 import FlashAttnV100Impl diff --git a/vllm/envs.py b/vllm/envs.py index e63f2353e5..034b37df2d 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -415,6 +415,7 @@ VLLM_FLASH_V100_DECODE_USE_WMMA_WRAPPER: bool = False VLLM_FLASH_V100_DECODE_USE_XQA: bool = True VLLM_FLASH_V100_DFLASH2_GROUPED_VERIFY: bool = True + VLLM_FLASH_V100_DFLASH2_BATCHED_GROUPED_VERIFY: bool = False 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 @@ -2775,6 +2776,11 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_FLASH_V100_DFLASH2_GROUPED_VERIFY": lambda: bool( int(os.getenv("VLLM_FLASH_V100_DFLASH2_GROUPED_VERIFY", "1")) ), + # Keep batched admission independent until the request-major kernel has + # passed B2/B4/B8 operator, graph, endpoint, and quality gates. + "VLLM_FLASH_V100_DFLASH2_BATCHED_GROUPED_VERIFY": lambda: bool( + int(os.getenv("VLLM_FLASH_V100_DFLASH2_BATCHED_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") ), diff --git a/vllm/v1/attention/backends/flash_attn_v100.py b/vllm/v1/attention/backends/flash_attn_v100.py index c6404fda5b..62a8953bcd 100644 --- a/vllm/v1/attention/backends/flash_attn_v100.py +++ b/vllm/v1/attention/backends/flash_attn_v100.py @@ -498,6 +498,7 @@ def _sm70_profile_trace(message: str, *args: object) -> None: _flash_attn_decode_paged_wmma = None _flash_attn_grouped_verify_paged = None _flash_attn_grouped_verify_max_query_tokens = 8 +_flash_attn_grouped_verify_request_major_abi_version = 0 _flash_attn_grouped_verify_checked = False _flash_attn_prefill_paged = None _flash_attn_prefill_paged_bhmd = None @@ -1264,6 +1265,7 @@ 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_max_query_tokens + global _flash_attn_grouped_verify_request_major_abi_version global _flash_attn_grouped_verify_checked if _flash_attn_grouped_verify_checked: return _flash_attn_grouped_verify_paged @@ -1283,6 +1285,16 @@ def _get_flash_grouped_verify_op(): ) except (ImportError, RuntimeError, TypeError, ValueError): _flash_attn_grouped_verify_max_query_tokens = 8 + try: + from flash_attn_v100 import ( + flash_attn_grouped_verify_request_major_abi_version, + ) + + _flash_attn_grouped_verify_request_major_abi_version = int( + flash_attn_grouped_verify_request_major_abi_version() + ) + except (ImportError, RuntimeError, TypeError, ValueError): + _flash_attn_grouped_verify_request_major_abi_version = 0 except ImportError: _flash_attn_grouped_verify_paged = None return _flash_attn_grouped_verify_paged @@ -4366,6 +4378,9 @@ def __init__(self, *args, **kwargs): self.dflash2_grouped_verify_max_query_tokens = ( _flash_attn_grouped_verify_max_query_tokens ) + self.dflash2_grouped_verify_request_major_abi_version = ( + _flash_attn_grouped_verify_request_major_abi_version + ) 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. @@ -4489,6 +4504,10 @@ def __init__(self, *args, **kwargs): and envs.VLLM_FLASH_V100_DFLASH2_GROUPED_VERIFY and current_platform.is_device_capability(70) ) + self.use_dflash2_batched_grouped_verify = ( + self.use_dflash2_grouped_verify + and envs.VLLM_FLASH_V100_DFLASH2_BATCHED_GROUPED_VERIFY + ) self.dflash2_grouped_verify_min_model_len = ( envs.VLLM_FLASH_V100_DFLASH2_GROUPED_VERIFY_MIN_MODEL_LEN ) @@ -5196,16 +5215,41 @@ def _dflash2_grouped_verify_allowed( global _logged_prefill_smallq_grouped_verify_gate block_table = getattr(attn_metadata, "block_table", None) seq_lens = getattr(attn_metadata, "seq_lens", None) + num_reqs = int( + getattr( + attn_metadata, + "num_reqs", + 0 if block_table is None else block_table.shape[0], + ) + ) + max_query_len = int( + getattr( + attn_metadata, + "max_query_len", + num_query_tokens if num_reqs == 1 else 0, + ) + ) + single_request_shape = bool( + num_reqs == 1 + and num_query_tokens in (8, 16) + and num_query_tokens <= self.dflash2_grouped_verify_max_query_tokens + ) + batched_request_shape = bool( + self.use_dflash2_batched_grouped_verify + and self.dflash2_grouped_verify_request_major_abi_version >= 1 + and num_reqs in (2, 4, 8) + and max_query_len == 8 + and num_query_tokens == num_reqs * 8 + ) allowed = bool( self.use_dflash2_grouped_verify + and (single_request_shape or batched_request_shape) 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 in (8, 16) - and num_query_tokens <= self.dflash2_grouped_verify_max_query_tokens and tuple(query.shape) == (num_query_tokens, 6, 256) and query.dtype == torch.float16 and query.is_contiguous() @@ -5226,13 +5270,13 @@ def _dflash2_grouped_verify_allowed( 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.shape[0] == num_reqs 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.shape[0] == num_reqs and seq_lens.device == query.device and seq_lens.dtype == torch.int32 and seq_lens.is_contiguous() @@ -5245,7 +5289,8 @@ def _dflash2_grouped_verify_allowed( 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 native_max_q=%d q=%s/%s " + "causal=%s window=%s reqs=%d max_q=%d actual=%d " + "native_max_q=%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, @@ -5254,6 +5299,8 @@ def _dflash2_grouped_verify_allowed( self.dflash2_grouped_verify_min_model_len, getattr(attn_metadata, "causal", True), self._flash_v100_window_size(causal=True), + num_reqs, + max_query_len, num_query_tokens, self.dflash2_grouped_verify_max_query_tokens, tuple(query.shape), @@ -5282,19 +5329,21 @@ def _call_dflash2_grouped_verify( out: torch.Tensor, ) -> None: global _logged_prefill_smallq_grouped_verify + num_reqs = int(attn_metadata.block_table.shape[0]) if not _logged_prefill_smallq_grouped_verify: logger.info( "FLASH_ATTN_V100 DFlash2 exact grouped verifier active " - "(q%d/H6/Hkv1/D256, FP8 E5M2 KV, one-pass).", - query.shape[0], + "(request-major B%d/q%d/H6/Hkv1/D256, FP8 E5M2 KV, one-pass).", + num_reqs, + query.shape[0] // num_reqs, ) _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], + attn_metadata.block_table[:num_reqs], + attn_metadata.seq_lens[:num_reqs], softmax_scale=self.scale, out=out, kv_cache_dtype=self.kv_cache_dtype, From bb487afcad3eabd944ec28a7371d9381079981ba Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 3 Sep 2026 22:13:13 +0800 Subject: [PATCH 3/9] [Core][SM70] Admit batched compact DFlash2 rejection Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- ...benchmark_sm70_dflash2_sparse_rejection.py | 66 +++++++++++++------ tests/v1/spec_decode/test_dflash2.py | 49 ++++++++++---- .../test_rejection_sampler_utils.py | 3 +- .../spec_decode/dflash2/sparse_rejection.py | 7 +- 4 files changed, 87 insertions(+), 38 deletions(-) diff --git a/benchmarks/benchmark_sm70_dflash2_sparse_rejection.py b/benchmarks/benchmark_sm70_dflash2_sparse_rejection.py index 51cc7ecf35..9f2df17d5a 100644 --- a/benchmarks/benchmark_sm70_dflash2_sparse_rejection.py +++ b/benchmarks/benchmark_sm70_dflash2_sparse_rejection.py @@ -1,6 +1,6 @@ # 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. +"""Benchmark batched 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. @@ -27,6 +27,7 @@ def _parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--vocab-size", type=int, default=248320) + parser.add_argument("--num-reqs", type=int, default=1) 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) @@ -78,6 +79,8 @@ def main() -> int: 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]") + if args.num_reqs < 1: + raise ValueError("--num-reqs must be positive") device = torch.device(args.device) torch.accelerator.set_device_index(device) @@ -86,8 +89,10 @@ def main() -> int: raise RuntimeError(f"Expected SM70, got sm_{capability[0]}{capability[1]}.") torch.manual_seed(args.seed) + num_reqs = args.num_reqs num_steps = args.num_speculative_steps - num_logits = num_steps + 1 + rows_per_req = num_steps + 1 + num_logits = num_reqs * rows_per_req raw_target = torch.randn( num_logits, args.vocab_size, @@ -111,37 +116,55 @@ def main() -> int: ) 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 + target_topk_ids_by_req = target_topk_ids.view( + num_reqs, rows_per_req, args.target_top_k ) + target_topk_logits_by_req = target_topk_logits.view( + num_reqs, rows_per_req, args.target_top_k + ) + draft_topk_ids = target_topk_ids_by_req[ + :, :num_steps, : args.draft_top_k + ].contiguous() draft_topk_logits = ( - target_topk_logits[:num_steps, : args.draft_top_k] + target_topk_logits_by_req[:, :num_steps, : args.draft_top_k] + torch.randn( + num_reqs, num_steps, args.draft_top_k, dtype=torch.float32, device=device, ) * 0.2 - ).view(1, num_steps, args.draft_top_k) + ).contiguous() dense_draft = torch.full( - (1, num_steps, args.vocab_size), + (num_reqs, 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) + 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.flatten() + cu_num_logits = ( + torch.arange(num_reqs + 1, dtype=torch.int32, device=device) * rows_per_req + ) 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) + 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) + temperature = torch.ones(num_reqs, dtype=torch.float32, device=device) + top_p_per_req = torch.full( + (num_reqs,), args.top_p, dtype=torch.float32, device=device + ) + seeds = torch.arange( + args.seed, args.seed + num_reqs, dtype=torch.int64, device=device + ) def dense_rejection() -> tuple[torch.Tensor, torch.Tensor]: return rejection_sample( @@ -197,13 +220,16 @@ def dense_finalize() -> tuple[torch.Tensor, torch.Tensor]: 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])) + valid = ( + torch.arange(rows_per_req, device=device).unsqueeze(0) < dense_count[:, None] + ) + tokens_equal = bool(torch.equal(dense_out[valid], sparse_out[valid])) result = { "device": torch.cuda.get_device_name(device), "device_capability": list(capability), "shape": { + "num_reqs": num_reqs, "num_logits": num_logits, "vocab_size": args.vocab_size, "target_top_k": args.target_top_k, @@ -213,8 +239,8 @@ def dense_finalize() -> tuple[torch.Tensor, torch.Tensor]: "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()), + "dense_num_sampled": dense_count.tolist(), + "sparse_num_sampled": sparse_count.tolist(), }, "timings": { "dense_topk_topp_only": _time_cuda( diff --git a/tests/v1/spec_decode/test_dflash2.py b/tests/v1/spec_decode/test_dflash2.py index 85ff6c2d69..c0710bbd3a 100644 --- a/tests/v1/spec_decode/test_dflash2.py +++ b/tests/v1/spec_decode/test_dflash2.py @@ -207,6 +207,19 @@ def test_dflash2_grouped_verify_is_default_on_with_rollback(monkeypatch): envs.disable_envs_cache() +def test_dflash2_batched_grouped_verify_is_default_off(monkeypatch): + name = "VLLM_FLASH_V100_DFLASH2_BATCHED_GROUPED_VERIFY" + monkeypatch.delenv(name, raising=False) + envs.disable_envs_cache() + try: + assert not getattr(envs, name) + monkeypatch.setenv(name, "1") + envs.disable_envs_cache() + assert 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() @@ -574,21 +587,23 @@ def test_probabilistic_dense_fallback_is_allocated_after_initialization(monkeypa assert torch.isneginf(speculator.draft_logits).all() -def _sparse_sampling_contract_fixture(): - idx = np.array([0], dtype=np.int32) +def _sparse_sampling_contract_fixture(num_reqs: int = 1): + idx = np.arange(num_reqs, 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)), + temperature=SimpleNamespace(np=np.ones(num_reqs, dtype=np.float32)), + top_k=SimpleNamespace(np=np.full(num_reqs, 20, dtype=np.int32)), + top_p=SimpleNamespace(np=np.full(num_reqs, 0.95, dtype=np.float32)), + min_p=SimpleNamespace(np=np.zeros(num_reqs, 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])), + penalties_state=SimpleNamespace(use_penalty=np.zeros(num_reqs, dtype=np.bool_)), + logit_bias_state=SimpleNamespace( + use_logit_bias=np.zeros(num_reqs, dtype=np.bool_) + ), bad_words_state=SimpleNamespace( - num_bad_words=SimpleNamespace(np=np.array([0], dtype=np.int32)) + num_bad_words=SimpleNamespace(np=np.zeros(num_reqs, dtype=np.int32)) ), logprob_token_ids_state=SimpleNamespace(max_num_token_ids=Mock(return_value=0)), compute_nans=False, @@ -598,18 +613,26 @@ def _sparse_sampling_contract_fixture(): sampler=sampler, ) input_batch = SimpleNamespace( - num_reqs=1, - is_prefilling_np=np.array([False]), + num_reqs=num_reqs, + is_prefilling_np=np.zeros(num_reqs, dtype=np.bool_), 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() +@pytest.mark.parametrize("num_reqs", [1, 2, 4, 8]) +def test_sparse_target_rejection_accepts_official_sampling_contract(num_reqs): + rejection_sampler, input_batch = _sparse_sampling_contract_fixture(num_reqs) assert _supports_sparse_sampling_contract(rejection_sampler, input_batch) +def test_sparse_target_rejection_rejects_mixed_prefill_batch(): + rejection_sampler, input_batch = _sparse_sampling_contract_fixture(4) + input_batch.is_prefilling_np[2] = True + + assert not _supports_sparse_sampling_contract(rejection_sampler, input_batch) + + @pytest.mark.parametrize( ("field", "value"), [ diff --git a/tests/v1/spec_decode/test_rejection_sampler_utils.py b/tests/v1/spec_decode/test_rejection_sampler_utils.py index 6ca495528b..8d93745638 100644 --- a/tests/v1/spec_decode/test_rejection_sampler_utils.py +++ b/tests/v1/spec_decode/test_rejection_sampler_utils.py @@ -133,7 +133,9 @@ 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]) +@pytest.mark.parametrize("num_reqs", [2, 4, 8]) def test_dflash2_sparse_topk_matches_dense_rejection( + num_reqs: int, num_speculative_steps: int, top_p: float, temperature: float, @@ -141,7 +143,6 @@ def test_dflash2_sparse_topk_matches_dense_rejection( """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 diff --git a/vllm/v1/worker/gpu/spec_decode/dflash2/sparse_rejection.py b/vllm/v1/worker/gpu/spec_decode/dflash2/sparse_rejection.py index 01b42e3968..caeb3bfda2 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash2/sparse_rejection.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash2/sparse_rejection.py @@ -155,10 +155,9 @@ def _supports_sparse_sampling_contract( """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): + # Every request must be in the uniform decode verifier phase. The compact + # kernel is request-indexed and preserves each request's sampling state. + if input_batch.num_reqs < 1 or np.any(input_batch.is_prefilling_np): return False sampler = rejection_sampler.sampler From 51541062b883ca4c6057e95a7e9f99dfa227c97f Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 3 Sep 2026 22:13:21 +0800 Subject: [PATCH 4/9] [Doc] Record Qwen3.8 DFlash2 concurrency audit Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../sm70_quasar_nvfp4_dflash2_acceptance.md | 48 +++++++++++++++++++ 1 file changed, 48 insertions(+) diff --git a/docs/design/sm70_quasar_nvfp4_dflash2_acceptance.md b/docs/design/sm70_quasar_nvfp4_dflash2_acceptance.md index 8103daf21c..f16f239942 100644 --- a/docs/design/sm70_quasar_nvfp4_dflash2_acceptance.md +++ b/docs/design/sm70_quasar_nvfp4_dflash2_acceptance.md @@ -287,3 +287,51 @@ explicit-history Responses replay. Its aggregate flag remains false only for because the current store has no eviction policy. The first stored response is valid and the expected follow-up receives HTTP 404; ordinary Chat Completions tool calling and explicit-history Responses are unaffected. + +## 2026-09-03 concurrency operator campaign + +This campaign extends the q7 DFlash2 verifier contract to B2/B4/B8 without +presenting operator projections as endpoint throughput. The fully QUASAR +checkpoint is not available on this host, so the NVFP4 MLP race uses real +layer-55 TP4 weights from the local mixed checkpoint; the three remaining +QUASAR projection shapes use deterministic native-E2M1 tensors. All timings +are CUDA-graph medians on one V100-SXM2-32GB with Torch 2.10/CUDA 12.8. + +The QPN2 kernel now tiles verifier rows in independent eight-row CTAs. The +per-row reduction order is unchanged. Across all 64 MLP projections, 16 full +attention projections, and 48 GDN projections, the weighted operator saving +versus TurboMind is `3.115 ms` at M=16 and `2.430 ms` at M=32. M=64 is a +`4.601 ms` regression, so opaque production dispatch admits only M<=32 and +retains TurboMind for B8/M=64. Every measured output is finite; QPN2-versus- +FP32 relative L2 is `3.3e-4--5.5e-4` with cosine approximately one. + +The Flash-V100 grouped verifier now accepts request-major q8 batches while +preserving the existing single-request q8/q16 and sparse-page4 paths. B2/B4/B8 +interleaved and non-interleaved KV outputs are bitwise equal to concatenated +single-request calls, and a captured B4 graph remains bitwise equal while each +request's runtime sequence length changes. The graph timings are: + +| Context | Batch | Grouped | Independent XQA | Speedup | +|---:|---:|---:|---:|---:| +| 1,024 | 2 | 0.0739 ms | 0.0571 ms | 0.77x | +| 1,024 | 4 | 0.0353 ms | 0.1058 ms | 3.00x | +| 1,024 | 8 | 0.0680 ms | 0.2188 ms | 3.22x | +| 16,384 | 2 | 0.1362 ms | 0.7056 ms | 5.18x | +| 16,384 | 4 | 0.2530 ms | 1.3867 ms | 5.48x | +| 16,384 | 8 | 0.4975 ms | 2.7351 ms | 5.50x | + +The B2/1K loss prevents unconditional promotion. Batched grouped verification +therefore remains behind the default-off +`VLLM_FLASH_V100_DFLASH2_BATCHED_GROUPED_VERIFY` switch until an endpoint run +can establish an actual-length admission policy. The compact target-rejection +path is also still opt-in, but its Python gate now accepts uniform decode-only +B2/B4/B8 batches. Against dense top-k/top-p plus rejection, compact p50 time is +`0.1290/0.0604/0.1300 ms` versus `0.5437/0.9810/1.2360 ms`, respectively. All +24 B2/B4/B8 combinations of q3/q7, top-p 1.0/0.95, and temperature 0.6/1.0 +produce the exact same valid tokens and accepted lengths as dense rejection. + +Before either opt-in becomes a default, run a matched TP4 endpoint matrix with +the fully QUASAR checkpoint and report pure decode separately from prefill and +TTFT. The required rows are B1/B2/B4/B8 at the same prompt/output lengths, +sampling seed policy, q7 draft, FP8 target KV, FP16 draft KV, attention backend, +and CUDA-graph state. From 45079079283e86a8a55e3a02ba7c9c3e674f9b9f Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 3 Sep 2026 23:44:10 +0800 Subject: [PATCH 5/9] [Doc] Record DFlash2 endpoint scaling Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../sm70_quasar_nvfp4_dflash2_acceptance.md | 55 +++++++++++++++++++ docs/design/sm70_v100_migration_control.md | 32 +++++++++++ 2 files changed, 87 insertions(+) diff --git a/docs/design/sm70_quasar_nvfp4_dflash2_acceptance.md b/docs/design/sm70_quasar_nvfp4_dflash2_acceptance.md index f16f239942..0649b23afd 100644 --- a/docs/design/sm70_quasar_nvfp4_dflash2_acceptance.md +++ b/docs/design/sm70_quasar_nvfp4_dflash2_acceptance.md @@ -335,3 +335,58 @@ the fully QUASAR checkpoint and report pure decode separately from prefill and TTFT. The required rows are B1/B2/B4/B8 at the same prompt/output lengths, sampling seed policy, q7 draft, FP8 target KV, FP16 draft KV, attention backend, and CUDA-graph state. + +### Mixed-checkpoint endpoint scaling probe + +A follow-up endpoint probe used the locally complete mixed-NVFP4 target because +the fully QUASAR checkpoint is still absent. The contract was TP4 on four +V100-SXM2-32GB GPUs, the official BF16 LM head, q7 probabilistic DFlash2, FP8 +E5M2 target KV, FP16 draft KV, Flash-V100 target and draft attention, +FULL_AND_PIECEWISE CUDA Graphs, official temperature 1.0/top-p 0.95/top-k 20 +sampling, sixteen fixed low-entropy SPEED-Bench 1K prompts, and 512 output +tokens per request. Batch-specific sampling shapes were warmed before the +reported B2/B4/B8 rows. The B1 row came from the same source, model, sampling, +and graph contract in the immediately preceding service boot. + +| Concurrency | Output token/s | Versus B1 | Ideal scaling efficiency | p50 TTFT | p50 TPOT | Mean accepted length | +|---:|---:|---:|---:|---:|---:|---:| +| 1 | 187.77 | 1.000x | 100.0% | 318.83 ms | 4.67 ms | 4.16 | +| 2 | 175.26 | 0.933x | 46.7% | 393.04 ms | 10.33 ms | 4.33 | +| 4 | 247.13 | 1.316x | 32.9% | 457.74 ms | 14.28 ms | 4.32 | +| 8 | 362.53 | 1.931x | 24.1% | 1147.63 ms | 18.62 ms | 4.49 | + +B2 is a real negative scaling point: steady aggregate output throughput is +6.7% below B1. B4 is 31.6% above B1 and 41.0% above B2. B8 is 93.1% above B1 +and 46.7% above B4, but still only 24.1% efficient relative to ideal linear +scaling. The earlier 1K operator result already showed that grouped B2 is +slower than independent XQA, but this endpoint matrix does not isolate that +operator from all other batch-dependent work. A matched grouped-verifier-off +arm is therefore required before changing the admission policy. + +The first formal B2 attempt measured only 139.26 output token/s because its +first batch triggered a target-sampling Triton JIT and p99 TTFT reached +11.46 seconds. It is retained as a cold-shape observation, not a steady +baseline. The repeated row above measured 175.26 token/s with p99 TTFT +0.73 seconds. Future concurrency harnesses must warm at least two output steps +for every measured batch shape; a one-token prefix warmup does not compile the +steady sampling path. + +The nominal prefix-warm pass queried the exact same input-length sequence, but +Prometheus recorded zero prefix-cache hits in every speed row. These numbers +are consequently 1K-input plus 512-output endpoint measurements, not pure +decode measurements. The source overlay also lacked the optional +`_vllm_fa2_C` exact D256 prefill operators and logged the slower prefill +fallback. Neither caveat invalidates the decode concurrency route hit, but +both prevent using this table as a final prefill or TTFT baseline. + +Runtime audit records show FULL target and DFlash CUDA Graph dispatch at +B2/q8=16 tokens, B4/q8=32 tokens, and B8/q8=64 tokens on every TP rank. Worker +logs also confirm QPN2 M<=32, FlashQLA GDN decode, FP8 E5M2 KV decode, compact +target rejection, and the request-major B8 grouped verifier. All four official +sampling rows completed 16/16 requests with zero errors, zero empty outputs, +and the requested 512 tokens. Sampled text is not byte-identical across +concurrency. A separate greedy B1/B8 smoke was byte-identical for 2/8 prompts; +the other six followed different but coherent trajectories, with 8/8 requests +complete and no runtime corruption signal. This is a text-health pass, not a +semantic-quality equivalence claim; normal benchmark quality gates remain +required before either concurrency switch is promoted. diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index ba36266f7c..14184a57c8 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -44895,3 +44895,35 @@ Interpretation: arithmetic path: it makes default capacity select the same previously quality-audited, checkpoint-code-preserving B1 operator path. Test services were stopped after collection and all four V100s returned to idle memory. + +## 2026-09-03 mixed-NVFP4 DFlash2 concurrency endpoint probe + +- Draft PR #476 at `51541062b8` was exercised with the local mixed-NVFP4 27B + target, BF16 LM head, q7 probabilistic DFlash2, TP4, Flash-V100, FlashQLA, + FP8 E5M2 target KV, FP16 draft KV, FULL_AND_PIECEWISE graphs, and sixteen + fixed SPEED-Bench 1K-by-512 requests per concurrency row. +- After batch-specific sampling warmup, aggregate output throughput was + B1 `187.77`, B2 `175.26`, B4 `247.13`, and B8 `362.53` token/s. Relative to + B1 this is `0.933x/1.316x/1.931x`, or `46.7%/32.9%/24.1%` ideal scaling + efficiency at B2/B4/B8. B2 is a 6.7% regression and must not be promoted + without a matched grouped-verifier-off endpoint arm. +- The cold B2 row was only `139.26` token/s because the first formal batch + JIT-compiled its sampling kernel and reached 11.46-second p99 TTFT. Retain it + as cold-shape evidence only. A one-output-token warmup is insufficient; + concurrency harnesses must warm the steady sampling path for each batch. +- Prefix-cache queries matched the prewarmed prompt lengths but recorded zero + hits, and the source overlay lacked optional exact D256 prefill operators. + The table is therefore an endpoint scaling measurement, not a pure-decode, + prefill, or final TTFT baseline. +- Target and DFlash graph audits hit FULL B2/q8, B4/q8, and B8/q8 descriptors + on all four ranks. QPN2 M<=32, FlashQLA decode, FP8 E5M2 KV, compact + rejection, and request-major grouped verification were present in worker + logs. Every official row completed 16/16 full-length outputs without an + error or empty response. Greedy B1/B8 was byte-identical on 2/8 prompts and + otherwise diverged into coherent text, so this closes text health but not + semantic-quality equivalence. +- The fully QUASAR checkpoint remains unavailable. Do not use this mixed-target + probe to close the final QUASAR acceptance gate. Raw artifacts are retained + under the PR worktree's `.artifacts/runtime/endpoint-b1-b2-b4-b8-v2` and + `.artifacts/runtime/endpoint-b2-b4-b8-v3` directories. GPU 0-3 returned to + 4 MiB per rank after graceful shutdown. From 086cdeebf5670a802d7cd4211f0d2701483cd075 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 4 Sep 2026 03:27:19 +0800 Subject: [PATCH 6/9] [Kernel][SM70] Optimize DFlash2 concurrency projections Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- benchmarks/benchmark_sm70_fp8_qpn8.py | 178 ++++++- csrc/custom_all_reduce.cuh | 32 +- csrc/sm70_turbomind/ops/fp8_qpn8_sm70.cu | 445 +++++++++++++++--- .../sm70_quasar_nvfp4_dflash2_acceptance.md | 128 +++++ docs/design/sm70_v100_migration_control.md | 58 +++ tests/test_envs.py | 4 + tests/v1/worker/test_gpu_model_runner_v2.py | 4 +- .../device_communicators/custom_all_reduce.py | 9 +- vllm/envs.py | 25 + vllm/v1/worker/gpu/model_runner.py | 2 +- 10 files changed, 779 insertions(+), 106 deletions(-) diff --git a/benchmarks/benchmark_sm70_fp8_qpn8.py b/benchmarks/benchmark_sm70_fp8_qpn8.py index 29584b4082..b164da7686 100644 --- a/benchmarks/benchmark_sm70_fp8_qpn8.py +++ b/benchmarks/benchmark_sm70_fp8_qpn8.py @@ -3,8 +3,8 @@ """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. +dispatch. It covers the exact Qwen3.8-27B-FP8 TP4 decode shapes through M=32 +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 @@ -92,8 +92,16 @@ def _load_keys(model: Path, keys: list[str]) -> dict[str, torch.Tensor]: return result -def _weight_keys(prefix: str) -> tuple[str, str]: - return f"{prefix}.weight", f"{prefix}.weight_scale_inv" +def _weight_keys(model: Path, prefix: str) -> tuple[str, str]: + weight_key = f"{prefix}.weight" + weight_map = json.loads((model / "model.safetensors.index.json").read_text())[ + "weight_map" + ] + for suffix in ("weight_scale_inv", "weight_scale"): + scale_key = f"{prefix}.{suffix}" + if scale_key in weight_map: + return weight_key, scale_key + raise KeyError(f"no FP8 scale tensor found for {prefix}") def _column_shard_raw( @@ -107,9 +115,11 @@ def _column_shard_raw( raise ValueError(f"N={n} is not divisible by TP={tp_size}") begin = tp_rank * (n // tp_size) end = begin + n // tp_size + raw = weight.view(torch.uint8)[begin:end].contiguous() + if scales.shape == (n, 1): + return raw, scales[begin:end].contiguous() 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() @@ -124,9 +134,11 @@ def _row_shard_raw( raise ValueError(f"K={k} is not divisible by TP={tp_size}") begin = tp_rank * (k // tp_size) end = begin + k // tp_size + raw = weight.view(torch.uint8)[:, begin:end].contiguous() + if scales.shape == (weight.shape[0], 1): + return raw, scales.contiguous() 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() @@ -138,10 +150,14 @@ def _load_case( device: torch.device, ) -> tuple[torch.Tensor, torch.Tensor, str]: root = "model.language_model.layers" + weight_map = json.loads((model / "model.safetensors.index.json").read_text())[ + "weight_map" + ] + mlp_layer = 1 if f"{root}.1.mlp.down_proj.weight" in weight_map else 56 if case == "down": - prefixes = [f"{root}.1.mlp.down_proj"] + prefixes = [f"{root}.{mlp_layer}.mlp.down_proj"] shard = "row" - note = "production down_proj TP row shard" + note = f"production layer-{mlp_layer} down_proj TP row shard" elif case == "gdn_in": prefixes = [ f"{root}.1.linear_attn.in_proj_qkv", @@ -163,8 +179,8 @@ def _load_case( 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", + f"{root}.{mlp_layer}.mlp.gate_proj", + f"{root}.{mlp_layer}.mlp.up_proj", ] shard = "column" note = ( @@ -175,13 +191,13 @@ def _load_case( else: raise ValueError(f"unknown case: {case}") - keys = [key for prefix in prefixes for key in _weight_keys(prefix)] + keys = [key for prefix in prefixes for key in _weight_keys(model, 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_key, scale_key = _weight_keys(model, prefix) weight = loaded[weight_key] scales = loaded[scale_key].float() if fp8_dtype is None: @@ -232,6 +248,8 @@ def _qpn8_prepack(raw_weight: torch.Tensor) -> torch.Tensor: def _qpn8_group_scales(scales: torch.Tensor, n: int, k: int) -> torch.Tensor: + if tuple(scales.shape) == (n, 1): + return scales.t().mul(256.0).half().contiguous() expected = (math.ceil(n / 128), math.ceil(k / 128)) if tuple(scales.shape) != expected or n % 128 or k % 128: raise ValueError( @@ -250,7 +268,11 @@ def _qpn8_group_scales(scales: torch.Tensor, n: int, k: int) -> torch.Tensor: 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) + expanded = ( + scales.expand(n, k) + if tuple(scales.shape) == (n, 1) + else scales.repeat_interleave(128, 0).repeat_interleave(128, 1) + ) return qweight.float().mul(expanded[:n, :k]) @@ -446,6 +468,7 @@ def _run_case(args: argparse.Namespace, case: str) -> dict[str, Any]: if (k // 16) % split_k == 0 and not (prefetch == "on" and decoder != "fast") and not (gated_silu and split_k == 32) + and not (split_k == 12 and (nacc != 2 or decoder != "fast" or prefetch == "on")) ] for m in args.m: torch.manual_seed(args.seed + m) @@ -520,7 +543,87 @@ def launch_tm( } ) + if m > 8: + dense_workspace = torch.empty((k, n), device=device, dtype=torch.float16) + prefill_out = torch.empty((m, output_n), device=device, dtype=torch.float16) + + def launch_prefill( + prefill_out: torch.Tensor = prefill_out, + dense_workspace: torch.Tensor = dense_workspace, + inputs: torch.Tensor = inputs, + codes: torch.Tensor = codes, + group_scales: torch.Tensor = group_scales, + gated_silu: bool = gated_silu, + ) -> None: + sm70_ops.fp8_qpn8_prefill_sm70_out( + prefill_out, + dense_workspace.data_ptr(), + inputs, + codes, + group_scales, + gated_silu, + ) + + launch_prefill() + torch.accelerator.synchronize(device) + prefill_quality = _error_stats(prefill_out, reference) + prefill_bytes = 5 * n * k + 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_prefill, + args.warmup, + args.iters, + args.trials, + scrub, + ) + if mode == "eager" + else benchmark( + launch_prefill, + prefill_out, + args.warmup, + args.iters, + args.trials, + scrub, + ) + ) + case_result["rows"].append( + { + "backend": "qpn8_prefill_current", + "mode": mode, + "cache_state": cache_state, + "m": m, + "config": None, + "quality": prefill_quality, + "quality_pass": _quality_pass( + prefill_quality, + args.relative_l2_limit, + args.cosine_limit, + ), + **measured, + **_derived_metrics( + measured["timing"], m, n, k, prefill_bytes + ), + } + ) + del dense_workspace, prefill_out + for split_k, nacc, fast_decoder, prefetch_codes in valid_configs: + if m > 8 and (split_k > 16 or (gated_silu and split_k > 8)): + continue + if m > 16 and ( + gated_silu + or split_k not in (12, 16) + or nacc != 2 + or not fast_decoder + or prefetch_codes + ): + continue qpn_out = torch.empty((m, output_n), device=device, dtype=torch.float16) def launch_qpn( @@ -560,6 +663,42 @@ def launch_qpn( launch_qpn() torch.accelerator.synchronize(device) quality = _error_stats(qpn_out, reference) + batch_invariance_equal: bool | None = None + batch_invariance_max_abs: float | None = None + if m > 8: + chunked_out = torch.empty_like(qpn_out) + for row_begin in range(0, m, 8): + row_end = min(row_begin + 8, m) + chunk_input = inputs[row_begin:row_end] + chunk_output = chunked_out[row_begin:row_end] + if gated_silu: + sm70_ops.fp8_qpn8_gated_pair_sm70_out( + chunk_output, + chunk_input, + codes, + group_scales, + split_k, + nacc, + fast_decoder, + prefetch_codes, + ) + else: + sm70_ops.fp8_qpn8_gemm_sm70_out( + chunk_output, + chunk_input, + codes, + group_scales, + split_k, + nacc, + fast_decoder, + prefetch_codes, + ) + torch.accelerator.synchronize(device) + batch_invariance_equal = bool(torch.equal(qpn_out, chunked_out)) + batch_invariance_max_abs = float( + (qpn_out - chunked_out).abs().max().item() + ) + del chunked_out quality_pass = _quality_pass( quality, args.relative_l2_limit, args.cosine_limit ) @@ -605,6 +744,8 @@ def launch_qpn( "config": config, "quality": quality, "quality_pass": quality_pass, + "batch_invariance_equal": batch_invariance_equal, + "batch_invariance_max_abs": batch_invariance_max_abs, **measured, **_derived_metrics(measured["timing"], m, n, k, qpn_bytes), } @@ -638,13 +779,16 @@ def _summarize(cases: list[dict[str, Any]]) -> list[dict[str, Any]]: 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_backend = ( + "qpn8_prefill_current" if m > 8 else "turbomind_current" + ) 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" + and row["backend"] == current_backend ) candidates = [ row @@ -654,6 +798,7 @@ def _summarize(cases: list[dict[str, Any]]) -> list[dict[str, Any]]: and row["cache_state"] == cache_state and row["backend"] == "qpn8_experimental" and row["quality_pass"] + and row.get("batch_invariance_equal") is not False and (row["replay_max_abs"] in (None, 0.0)) ] best = min(candidates, key=lambda row: row["timing"]["median_us"]) @@ -666,6 +811,7 @@ def _summarize(cases: list[dict[str, Any]]) -> list[dict[str, Any]]: "m": m, "mode": mode, "cache_state": cache_state, + "current_backend": current_backend, "current_us": current_us, "best_qpn8_us": best_us, "speedup": current_us / best_us, @@ -691,8 +837,8 @@ def main() -> int: 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 any(m < 1 or m > 32 for m in args.m): + raise ValueError("QPN8 operator supports M=1..32") 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") diff --git a/csrc/custom_all_reduce.cuh b/csrc/custom_all_reduce.cuh index 66e6765ff3..66197e05f5 100644 --- a/csrc/custom_all_reduce.cuh +++ b/csrc/custom_all_reduce.cuh @@ -67,8 +67,13 @@ constexpr int kSm70Tp4PushAllreduceThreads = 128; constexpr int kSm70Tp4PushAllreduceEpochs = 2; constexpr uint16_t kSm70Tp4PushAllreduceSentinel = 0x7f7f; constexpr int kSm70Tp4PushAllreduceSentinelByte = 0x7f; -constexpr size_t kSm70Tp4PushAllreduceBytes = +constexpr size_t kSm70Tp4PushAllreduceM8Bytes = 8 * kSm70GemmaRmsNormHiddenSize * sizeof(half); +constexpr size_t kSm70Tp4PushAllreduceM16Bytes = + 16 * kSm70GemmaRmsNormHiddenSize * sizeof(half); +constexpr size_t kSm70Tp4PushAllreduceM32Bytes = + 32 * kSm70GemmaRmsNormHiddenSize * sizeof(half); +constexpr size_t kSm70Tp4PushAllreduceMaxBytes = kSm70Tp4PushAllreduceM32Bytes; constexpr size_t kSm70Tp4PushAllreduce8KiBBytes = 4096 * sizeof(half); constexpr size_t kSm70Tp4PushAllreduceQwen4ExpBytes = 2560 * sizeof(half); constexpr size_t kSm70Tp4PushAllreduceQwen4ExpMtp5Bytes = @@ -78,10 +83,17 @@ constexpr size_t kSm70Tp4PushAllreduceSignalBytes = constexpr size_t kSm70Tp4PushAllreduceBufferBytes = kSm70Tp4PushAllreduceSignalBytes + kSm70Tp4PushAllreduceEpochs * kSm70Tp4PushAllreduceWorldSize * - kSm70Tp4PushAllreduceBytes; + kSm70Tp4PushAllreduceMaxBytes; inline int sm70_tp4_push_allreduce_blocks(size_t bytes) { - if (bytes == kSm70Tp4PushAllreduceBytes) { + if (bytes == kSm70Tp4PushAllreduceM8Bytes) { + return kSm70Tp4PushAllreduceBlocks; + } + const char* concurrency = + std::getenv("VLLM_SM70_TP4_PUSH_ALLREDUCE_CONCURRENCY"); + if (concurrency != nullptr && std::strcmp(concurrency, "1") == 0 && + (bytes == kSm70Tp4PushAllreduceM16Bytes || + bytes == kSm70Tp4PushAllreduceM32Bytes)) { return kSm70Tp4PushAllreduceBlocks; } if (bytes == kSm70Tp4PushAllreduce8KiBBytes) { @@ -655,11 +667,12 @@ __global__ void __launch_bounds__(1024, 1) 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); + constexpr int packed_stride = kSm70Tp4PushAllreduceMaxBytes / sizeof(P); const int epoch_offset = epoch * ngpus * packed_stride; - const int offset = blockIdx.x * blockDim.x + threadIdx.x; + const int first_offset = blockIdx.x * blockDim.x + threadIdx.x; - if (offset < packed_size) { + for (int offset = first_offset; offset < packed_size; + offset += gridDim.x * blockDim.x) { P value = reinterpret_cast(input)[offset]; #pragma unroll for (int element = 0; element < P::size; ++element) { @@ -732,11 +745,12 @@ __global__ void __launch_bounds__(1024, 1) 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); + constexpr int packed_stride = kSm70Tp4PushAllreduceMaxBytes / sizeof(P); const int epoch_offset = epoch * ngpus * packed_stride; - const int offset = blockIdx.x * blockDim.x + threadIdx.x; + const int first_offset = blockIdx.x * blockDim.x + threadIdx.x; - if (offset < packed_size) { + for (int offset = first_offset; offset < packed_size; + offset += gridDim.x * blockDim.x) { P value_a = reinterpret_cast(input_a)[offset]; const P value_b = reinterpret_cast(input_b)[offset]; packed_assign_add(value_a, value_b); diff --git a/csrc/sm70_turbomind/ops/fp8_qpn8_sm70.cu b/csrc/sm70_turbomind/ops/fp8_qpn8_sm70.cu index c91adc4449..0914a4d77c 100644 --- a/csrc/sm70_turbomind/ops/fp8_qpn8_sm70.cu +++ b/csrc/sm70_turbomind/ops/fp8_qpn8_sm70.cu @@ -16,6 +16,8 @@ #include #include +#include +#include #include #include @@ -203,7 +205,8 @@ __global__ void fp8_qpn8_ba_split_copy_sm70_kernel( } template + bool M1Only = false, bool FusedBA = false, bool SplitOutputs = false, + int RowTiles = 1> __global__ void fp8_qpn8_sm70_kernel( const uint8_t* __restrict__ codes, const half* __restrict__ group_scales, const half* __restrict__ input, half* __restrict__ output, @@ -211,7 +214,11 @@ __global__ void fp8_qpn8_sm70_kernel( 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]; + static_assert(RowTiles == 1 || RowTiles == 2, + "QPN8 supports one or two 8-row tiles"); + static_assert(!M1Only || RowTiles == 1, + "QPN8 M=1 specialization uses one row tile"); + __shared__ float partials[SplitK][M1Only ? 32 : RowTiles * 256]; const int lane = threadIdx.x & 31; const int warp = threadIdx.x >> 5; @@ -275,12 +282,15 @@ __global__ void fp8_qpn8_sm70_kernel( static_cast(tile) * groups_k16 * 32 + lane; const half* scale_ptr = group_scales + tile; - float accum[NAcc][8]; + float accum[RowTiles][NAcc][8]; #pragma unroll - for (int chain = 0; chain < NAcc; ++chain) { + for (int row_tile = 0; row_tile < RowTiles; ++row_tile) { #pragma unroll - for (int i = 0; i < 8; ++i) { - accum[chain][i] = 0.0f; + for (int chain = 0; chain < NAcc; ++chain) { +#pragma unroll + for (int i = 0; i < 8; ++i) { + accum[row_tile][chain][i] = 0.0f; + } } } int loaded_scale_group = -1; @@ -327,31 +337,38 @@ __global__ void fp8_qpn8_sm70_kernel( 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]); +#pragma unroll + for (int row_tile = 0; row_tile < RowTiles; ++row_tile) { + uint4 input01 = make_uint4(0, 0, 0, 0); + uint4 input23 = make_uint4(0, 0, 0, 0); + const int input_row_idx = row_tile * 8 + row; + if (input_row_idx < m) { + const half* input_row = input + static_cast(input_row_idx) * 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); + VLLM_SM70_MMA_8N8K4(accum[row_tile][0], a0[0], a0[1], b[0], b[1]); + VLLM_SM70_MMA_8N8K4(accum[row_tile][1 % NAcc], a0[2], a0[3], b[2], b[3]); + VLLM_SM70_MMA_8N8K4(accum[row_tile][2 % NAcc], a1[0], a1[1], b[4], b[5]); + VLLM_SM70_MMA_8N8K4(accum[row_tile][3 % NAcc], a1[2], a1[3], b[6], b[7]); + } if constexpr (PrefetchCodes) { prefetched = next; } } #pragma unroll - for (int chain = 1; chain < NAcc; ++chain) { + for (int row_tile = 0; row_tile < RowTiles; ++row_tile) { #pragma unroll - for (int i = 0; i < 8; ++i) { - accum[0][i] += accum[chain][i]; + for (int chain = 1; chain < NAcc; ++chain) { +#pragma unroll + for (int i = 0; i < 8; ++i) { + accum[row_tile][0][i] += accum[row_tile][chain][i]; + } } } @@ -364,22 +381,27 @@ __global__ void fp8_qpn8_sm70_kernel( 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]; + partials[warp][quadpair * 8 + output_col] = accum[0][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]; + for (int row_tile = 0; row_tile < RowTiles; ++row_tile) { +#pragma unroll + for (int i = 0; i < 8; ++i) { + const int output_row = + row_tile * 8 + (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[row_tile][0][i]; + } } } __syncthreads(); - constexpr int kOutputElements = M1Only ? 32 : 256; + constexpr int kOutputElements = M1Only ? 32 : RowTiles * 256; for (int element = threadIdx.x; element < kOutputElements; element += blockDim.x) { float value = 0.0f; @@ -410,14 +432,144 @@ __global__ void fp8_qpn8_sm70_kernel( } template + bool M1Only = false, int RowTiles = 1> 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); + 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); +} + +// M32 needs four independent 8-row accumulator tiles. Keeping all logical +// split-K warps resident would require 64 KiB of static reduction storage for +// split-16. Instead, half as many physical warps execute the original logical +// warp ranges in two ordered phases. The compact first-half sum lets the final +// reduction retain the exact p0 + ... + p(SplitK-1) order while using only +// 28 KiB (split-12) or 36 KiB (split-16) of shared memory. Each output CTA +// streams its packed weight tile once and reuses it across all 32 rows. +template +__global__ void fp8_qpn8_m32_twophase_sm70_kernel( + const uint8_t* __restrict__ codes, const half* __restrict__ channel_scales, + const half* __restrict__ input, half* __restrict__ output, int n, int k, + int m) { + static_assert(SplitK == 12 || SplitK == 16, + "M32 two-phase QPN8 supports split-12 or split-16"); + constexpr int kPhysicalWarps = SplitK / 2; + constexpr int kRowTiles = 4; + constexpr int kOutputElements = kRowTiles * 256; + __shared__ float reduction_storage[kPhysicalWarps + 1][kOutputElements]; + + 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 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); + +#pragma unroll + for (int phase = 0; phase < 2; ++phase) { + float accum[kRowTiles][2][8]; +#pragma unroll + for (int row_tile = 0; row_tile < kRowTiles; ++row_tile) { +#pragma unroll + for (int chain = 0; chain < 2; ++chain) { +#pragma unroll + for (int index = 0; index < 8; ++index) { + accum[row_tile][chain][index] = 0.0f; + } + } + } + + const int logical_warp = warp + phase * kPhysicalWarps; + const int group_begin = logical_warp * groups_per_warp; +#pragma unroll 4 + for (int group = group_begin; group < group_begin + groups_per_warp; + ++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); + } + + const unsigned* b = reinterpret_cast(weights); +#pragma unroll + for (int row_tile = 0; row_tile < kRowTiles; ++row_tile) { + uint4 input01 = make_uint4(0, 0, 0, 0); + uint4 input23 = make_uint4(0, 0, 0, 0); + const int input_row_idx = row_tile * 8 + row; + if (input_row_idx < m) { + const half* input_row = + input + static_cast(input_row_idx) * 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); + VLLM_SM70_MMA_8N8K4(accum[row_tile][0], a0[0], a0[1], b[0], b[1]); + VLLM_SM70_MMA_8N8K4(accum[row_tile][1], a0[2], a0[3], b[2], b[3]); + VLLM_SM70_MMA_8N8K4(accum[row_tile][0], a1[0], a1[1], b[4], b[5]); + VLLM_SM70_MMA_8N8K4(accum[row_tile][1], a1[2], a1[3], b[6], b[7]); + } + } + +#pragma unroll + for (int row_tile = 0; row_tile < kRowTiles; ++row_tile) { +#pragma unroll + for (int index = 0; index < 8; ++index) { + accum[row_tile][0][index] += accum[row_tile][1][index]; + const int output_row = + row_tile * 8 + (index & 2) + ((lane & 16) ? 4 : 0) + (lane & 1); + const int output_col = + (index & 1) | (((lane >> 1) & 1) << 1) | ((index >> 2) << 2); + reduction_storage[warp][output_row * 32 + quadpair * 8 + output_col] = + accum[row_tile][0][index]; + } + } + __syncthreads(); + + for (int element = threadIdx.x; element < kOutputElements; + element += blockDim.x) { + float value = + phase == 0 ? 0.0f : reduction_storage[kPhysicalWarps][element]; +#pragma unroll + for (int k_warp = 0; k_warp < kPhysicalWarps; ++k_warp) { + value += reduction_storage[k_warp][element]; + } + if (phase == 0) { + reduction_storage[kPhysicalWarps][element] = 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); + } + } + } + __syncthreads(); + } +} + +template +void launch_fp8_qpn8_m32_twophase_sm70(const uint8_t* codes, + const half* channel_scales, + const half* input, half* output, int n, + int k, int m, cudaStream_t stream) { + constexpr int kPhysicalWarps = SplitK / 2; + fp8_qpn8_m32_twophase_sm70_kernel + <<<(n / 32), (32 * kPhysicalWarps), 0, stream>>>(codes, channel_scales, + input, output, n, k, m); } void launch_fp8_qpn8_ba_split_sm70(const uint8_t* codes, @@ -434,12 +586,16 @@ void launch_fp8_qpn8_ba_split_sm70(const uint8_t* codes, } template + bool M1Only = false, int RowTiles = 1> __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]; + static_assert(RowTiles == 1 || RowTiles == 2, + "QPN8 gated pair supports one or two 8-row tiles"); + static_assert(!M1Only || RowTiles == 1, + "QPN8 gated M=1 specialization uses one row tile"); + __shared__ float partials[2][SplitK][M1Only ? 32 : RowTiles * 256]; const int lane = threadIdx.x & 31; const int warp_in_block = threadIdx.x >> 5; @@ -457,12 +613,15 @@ __global__ void fp8_qpn8_gated_pair_sm70_kernel( static_cast(tile) * groups_k16 * 32 + lane; const half* scale_ptr = group_scales + tile; - float accum[NAcc][8]; + float accum[RowTiles][NAcc][8]; #pragma unroll - for (int chain = 0; chain < NAcc; ++chain) { + for (int row_tile = 0; row_tile < RowTiles; ++row_tile) { #pragma unroll - for (int i = 0; i < 8; ++i) { - accum[chain][i] = 0.0f; + for (int chain = 0; chain < NAcc; ++chain) { +#pragma unroll + for (int i = 0; i < 8; ++i) { + accum[row_tile][chain][i] = 0.0f; + } } } int loaded_scale_group = -1; @@ -508,30 +667,37 @@ __global__ void fp8_qpn8_gated_pair_sm70_kernel( 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]); +#pragma unroll + for (int row_tile = 0; row_tile < RowTiles; ++row_tile) { + uint4 input01 = make_uint4(0, 0, 0, 0); + uint4 input23 = make_uint4(0, 0, 0, 0); + const int input_row_idx = row_tile * 8 + row; + if (input_row_idx < m) { + const half* input_row = input + static_cast(input_row_idx) * 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); + VLLM_SM70_MMA_8N8K4(accum[row_tile][0], a0[0], a0[1], b[0], b[1]); + VLLM_SM70_MMA_8N8K4(accum[row_tile][1 % NAcc], a0[2], a0[3], b[2], b[3]); + VLLM_SM70_MMA_8N8K4(accum[row_tile][2 % NAcc], a1[0], a1[1], b[4], b[5]); + VLLM_SM70_MMA_8N8K4(accum[row_tile][3 % NAcc], a1[2], a1[3], b[6], b[7]); + } if constexpr (PrefetchCodes) { prefetched = next; } } #pragma unroll - for (int chain = 1; chain < NAcc; ++chain) { + for (int row_tile = 0; row_tile < RowTiles; ++row_tile) { #pragma unroll - for (int i = 0; i < 8; ++i) { - accum[0][i] += accum[chain][i]; + for (int chain = 1; chain < NAcc; ++chain) { +#pragma unroll + for (int i = 0; i < 8; ++i) { + accum[row_tile][0][i] += accum[row_tile][chain][i]; + } } } if constexpr (M1Only) { @@ -543,23 +709,28 @@ __global__ void fp8_qpn8_gated_pair_sm70_kernel( 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]; + partials[projection][warp][quadpair * 8 + output_col] = + accum[0][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]; + for (int row_tile = 0; row_tile < RowTiles; ++row_tile) { +#pragma unroll + for (int i = 0; i < 8; ++i) { + const int output_row = + row_tile * 8 + (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[row_tile][0][i]; + } } } __syncthreads(); - constexpr int kOutputElements = M1Only ? 32 : 256; + constexpr int kOutputElements = M1Only ? 32 : RowTiles * 256; for (int element = threadIdx.x; element < kOutputElements; element += blockDim.x) { float gate = 0.0f; @@ -585,14 +756,14 @@ __global__ void fp8_qpn8_gated_pair_sm70_kernel( } template + bool M1Only = false, int RowTiles = 1> 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 + M1Only, RowTiles> <<<(hidden / 32), (64 * SplitK), 0, stream>>>( codes, group_scales, input, output, hidden, k, m, channel_scales); } @@ -1190,7 +1361,8 @@ void fp8_qpn8_gemm_sm70_out(torch::Tensor out, torch::Tensor input, 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(m >= 1 && m <= 32, + "fp8_qpn8_gemm_sm70_out: M must be in [1, 32]"); 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"); @@ -1205,11 +1377,15 @@ void fp8_qpn8_gemm_sm70_out(torch::Tensor out, torch::Tensor input, "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(m <= 8 || channel_scales, + "fp8_qpn8_gemm_sm70_out: M=9..32 requires channel scales"); 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(m <= 8 || split_k <= 16, + "fp8_qpn8_gemm_sm70_out: M=9..32 does not support split_k 32"); 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, @@ -1220,6 +1396,11 @@ void fp8_qpn8_gemm_sm70_out(torch::Tensor out, torch::Tensor input, (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"); + TORCH_CHECK( + m <= 16 || ((split_k == 12 || split_k == 16) && accumulator_chains == 2 && + fast_decoder && !prefetch_codes), + "fp8_qpn8_gemm_sm70_out: M=17..32 requires split_k 12/16, " + "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(); @@ -1230,6 +1411,20 @@ void fp8_qpn8_gemm_sm70_out(torch::Tensor out, torch::Tensor input, reinterpret_cast(input.data_ptr()); auto* output_ptr = reinterpret_cast(out.data_ptr()); + if (m > 16) { + if (split_k == 12) { + launch_fp8_qpn8_m32_twophase_sm70<12>( + code_ptr, scale_ptr, input_ptr, output_ptr, static_cast(n), + static_cast(k), static_cast(m), stream); + } else { + launch_fp8_qpn8_m32_twophase_sm70<16>( + code_ptr, scale_ptr, input_ptr, output_ptr, static_cast(n), + static_cast(k), static_cast(m), stream); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return; + } + // 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. @@ -1248,10 +1443,20 @@ void fp8_qpn8_gemm_sm70_out(torch::Tensor out, torch::Tensor input, 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) +#define VLLM_LAUNCH_QPN8(SPLIT, NACC, FAST, PREFETCH) \ + do { \ + if (m <= 8) { \ + launch_fp8_qpn8_sm70( \ + code_ptr, scale_ptr, input_ptr, output_ptr, static_cast(n), \ + static_cast(k), static_cast(m), channel_scales, stream); \ + } else if constexpr (SPLIT <= 16) { \ + launch_fp8_qpn8_sm70( \ + code_ptr, scale_ptr, input_ptr, output_ptr, static_cast(n), \ + static_cast(k), static_cast(m), true, stream); \ + } else { \ + TORCH_CHECK(false, "QPN8 M=9..16 does not support split_k ", SPLIT); \ + } \ + } while (0) if (prefetch_codes) { if (split_k == 4 && accumulator_chains == 1) { @@ -1509,8 +1714,8 @@ void fp8_qpn8_gated_pair_sm70_out(torch::Tensor out, torch::Tensor input, 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(m >= 1 && m <= 16 && out.size(0) == m, + "fp8_qpn8_gated_pair_sm70_out: M must be in [1, 16]"); const bool channel_scales = group_scales.size(0) == 1 && group_scales.size(1) == n; const bool block_scales = @@ -1524,10 +1729,14 @@ void fp8_qpn8_gated_pair_sm70_out(torch::Tensor out, torch::Tensor input, "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(m <= 8 || channel_scales, + "fp8_qpn8_gated_pair_sm70_out: M=9..16 requires channel scales"); 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(m <= 8 || split_k <= 8, + "fp8_qpn8_gated_pair_sm70_out: M=9..16 supports split_k 4 or 8"); TORCH_CHECK( accumulator_chains == 1 || accumulator_chains == 2, "fp8_qpn8_gated_pair_sm70_out: accumulator_chains must be 1 or 2"); @@ -1554,10 +1763,23 @@ void fp8_qpn8_gated_pair_sm70_out(torch::Tensor out, torch::Tensor input, 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) +#define VLLM_LAUNCH_QPN8_GATED_PAIR(SPLIT, NACC, FAST, PREFETCH) \ + do { \ + if (m <= 8) { \ + 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); \ + } else if constexpr (SPLIT <= 8) { \ + launch_fp8_qpn8_gated_pair_sm70( \ + code_ptr, scale_ptr, input_ptr, output_ptr, \ + static_cast(hidden), static_cast(k), static_cast(m), \ + true, stream); \ + } else { \ + TORCH_CHECK(false, "QPN8 gated M=9..16 does not support split_k ", \ + SPLIT); \ + } \ + } while (0) if (prefetch_codes) { if (split_k == 4 && accumulator_chains == 1) { @@ -1622,6 +1844,73 @@ void fp8_qpn8_dispatch_sm70_out(torch::Tensor out, int64_t dense_weight_ptr, } return; } + + // DFlash2 verifies eight speculative tokens per request, so B2 reaches + // M=16. The old path reconstructed the complete channel-FP8 weight before + // every GEMM. Keep the two-row-tile QPN8 candidate opt-in and restricted to + // the exact Qwen3.8 attention/GDN/last-eight-MLP shapes measured below. + const auto env_enabled = [](const char* name) { + const char* value = std::getenv(name); + return value != nullptr && value[0] == '1' && value[1] == '\0'; + }; + const bool qpn8_m16_enabled = env_enabled("VLLM_SM70_FP8_QPN8_M16"); + const bool qpn8_m32_chunked_enabled = + env_enabled("VLLM_SM70_FP8_QPN8_M32_CHUNKED"); + const bool qpn8_m32_native_enabled = + env_enabled("VLLM_SM70_FP8_QPN8_M32_NATIVE"); + const int64_t m = input.size(0); + const int64_t k = input.size(1); + const int64_t packed_n = codes.size(1); + const bool channel_scales = group_scales.dim() == 2 && + group_scales.size(0) == 1 && + group_scales.size(1) == packed_n; + const bool qwen38_dense_shape = + !gated_silu && + ((k == 5120 && (packed_n == 4096 || packed_n == 3584)) || + (k == 1536 && packed_n == 5120) || (k == 4352 && packed_n == 5120)); + const bool qwen38_gated_shape = gated_silu && k == 5120 && packed_n == 8704; + const bool native_m32 = + qpn8_m32_native_enabled && m > 16 && m <= 32 && qwen38_dense_shape; + const bool admitted_m = (qpn8_m16_enabled && m <= 16) || + (qpn8_m32_chunked_enabled && m <= 32) || native_m32; + if (admitted_m && channel_scales && + (qwen38_dense_shape || qwen38_gated_shape)) { + if (native_m32) { + static std::once_flag qpn8_m32_native_log_once; + std::call_once(qpn8_m32_native_log_once, []() { + std::fprintf(stderr, + "INFO SM70 channel-FP8 QPN8 native M=17..32 " + "two-phase dense candidate enabled.\n"); + }); + fp8_qpn8_gemm_sm70_out(out, input, codes, group_scales, split_k, + accumulator_chains, true, prefetch_codes); + return; + } + static std::once_flag qpn8_m16_log_once; + std::call_once(qpn8_m16_log_once, []() { + std::fprintf(stderr, + "INFO SM70 channel-FP8 QPN8 M=9..32 batch-invariant " + "two-row-tile/chunked candidate enabled.\n"); + }); + // M32 reuses the accepted M16 body in contiguous row chunks. Each chunk + // still reads the packed weight directly, avoiding the old full matrix + // reconstruction while preserving the same reduction order. + for (int64_t row = 0; row < m; row += 16) { + const int64_t rows = std::min(16, m - row); + auto input_chunk = input.narrow(0, row, rows); + auto out_chunk = out.narrow(0, row, rows); + if (qwen38_gated_shape) { + fp8_qpn8_gated_pair_sm70_out(out_chunk, input_chunk, codes, + group_scales, split_k, accumulator_chains, + true, prefetch_codes); + } else { + fp8_qpn8_gemm_sm70_out(out_chunk, input_chunk, 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); } diff --git a/docs/design/sm70_quasar_nvfp4_dflash2_acceptance.md b/docs/design/sm70_quasar_nvfp4_dflash2_acceptance.md index 0649b23afd..44dd40a6aa 100644 --- a/docs/design/sm70_quasar_nvfp4_dflash2_acceptance.md +++ b/docs/design/sm70_quasar_nvfp4_dflash2_acceptance.md @@ -390,3 +390,131 @@ the other six followed different but coherent trajectories, with 8/8 requests complete and no runtime corruption signal. This is a text-health pass, not a semantic-quality equivalence claim; normal benchmark quality gates remain required before either concurrency switch is promoted. + +### M16/M32 channel-FP8 concurrency optimization + +The explicit scaling targets use the steady B1 result above as the denominator: +B2 must reach `300.43 token/s` (80%), B4 `525.76 token/s` (70%), and B8 +`901.30 token/s` (60%). The work below remains a mixed-checkpoint optimization +screen; it does not close the unavailable fully-QUASAR quality gate. + +The old channel-FP8 path reconstructs a complete FP16 weight before every +M>8 GEMM. A default-off QPN8 candidate now handles M=9--16 with two 8-row +tiles, and a separate M=17--32 dense candidate executes the original logical +split-K ranges in two ordered phases. The M32 design reduces static reduction +storage from the naive 64 KiB to 28/36 KiB for split-12/16 and streams each +packed weight tile once across all 32 rows. Admission is restricted to the +five measured Qwen3.8 TP4 channel-FP8 shapes. The controls are +`VLLM_SM70_FP8_QPN8_M16`, `VLLM_SM70_FP8_QPN8_M32_CHUNKED`, and +`VLLM_SM70_FP8_QPN8_M32_NATIVE`; all default to off. + +Actual-checkpoint operator tests show that M16 reduces the weighted QPN8 +projection bucket from `20.276` to `6.959 ms` per target round. At M32, native +dense graph timings are `82.68 us` for GDN input, `27.38 us` for output, +`81.86 us` for full-attention QKV, and `73.71 us` for down projection. Across +M=17/18/24/31/32, every production split is bitwise equal to concatenated M8 +calls with maximum difference zero; CUDA Graph replay is also stable. The +retained operator artifacts are `.artifacts/runtime/qpn8-m32-native-dense-r2.json` +and `.artifacts/runtime/qpn8-m32-native-dense-tails-r1.json`. + +Same-contract endpoint results are: + +| Candidate | B2 token/s | B2 efficiency | B4 token/s | B4 efficiency | +|---|---:|---:|---:|---:| +| steady baseline | 175.26 | 46.7% | 247.13 | 32.9% | +| exact M16 + chunked M32 | 253.95 | 67.6% | 309.06 | 41.2% | +| exact M16 + native dense M32 | - | - | 322.68 | 43.0% | + +The exact B2 candidate improves the old B2 row by 44.9%; native M32 improves +the old B4 row by 30.6%. All reported rows completed 16/16 requests, generated +the requested 512 tokens, and contained no empty output or replacement +character. B2 acceptance was `47.62%` with mean accepted length `4.33`; B4 +native acceptance was `51.68%` with mean length `4.62`. These are output-health +and distribution-correction checks, not a semantic benchmark. Raw endpoint +artifacts are under `.artifacts/runtime/endpoint-qpn8-m16-m32-exact-b2-b4-v1` +and `.artifacts/runtime/endpoint-qpn8-m32-native-b4-v1`. + +Three experiments were rejected: + +- Moving channel scale to the epilogue and changing split configurations made + M16 faster, but B4/B8 acceptance length fell by about 8%/10%; it is not the + retained quality-first implementation. +- Draft proposal temperature scale `0.85` measured B2 `258.38 token/s` and B4 + `313.46 token/s`; it did not improve native-M32 B4 and remains off. +- A bitwise-exact native M64 kernel looked positive with warm operator caches, + but the endpoint fell to `283.94 token/s`, 21.7% below the steady B8 baseline. + It was removed. The retained negative artifact is + `.artifacts/runtime/endpoint-qpn8-m64-native-b8-v1`. + +The gates remain open. Native M32 leaves B4 at about 43% efficiency, and B8 +still needs a structural change rather than more row tiling. At observed +accepted lengths, the remaining gap cannot be closed by selector or epilogue +micro-tuning alone. The next measurement should split the unprofiled native +candidate into target forward, target logits/rejection, draft, and host +bookkeeping, then evaluate either multi-stream/request partitioning or a +deployment topology that adds independent replicas. Two Nsight Systems runs +were attempted, but CUPTI crashed during multiprocess shutdown before writing +a report; do not repeat that capture path unchanged. + +The built-in CUDA-event profiler initially produced no output because the MRV2 +gate admitted `mtp` only, while this service resolves the same diagnostic path +with `method=dflash`. The default-off profiler now admits `mtp`, `dflash`, and +`dspark`, matching the legacy runner. A short B2/B4 diagnostic then completed. +The profiler synchronizes every round, so its endpoint throughput is not a +performance result; use only the per-phase CUDA-event split. Median stable +full-batch intervals were: + +| Batch | Target forward | Target sample + state | Draft | Total GPU | +|---:|---:|---:|---:|---:| +| B2 / M16 | 37.64 ms | 1.29 ms | 7.50 ms | 46.65 ms | +| B4 / M32 | 47.41 ms | 1.54 ms | 9.05 ms | 58.10 ms | + +Target forward is about 81% of the measured GPU interval in both rows and +accounts for nearly all B2-to-B4 growth. Rejection/sampling is only about +1.3--1.5 ms, which rules out more selector micro-tuning as the primary scaling +project. The diagnostic artifact is +`.artifacts/runtime/endpoint-qpn8-mrv2-profile-b2-b4-v1`. + +A q3 B8 screen tested whether shrinking the verifier from M64 to M32 could +avoid the remaining large-batch fallback. It reached only `249.14 token/s` +with `72.85%` draft acceptance and `3.19` emitted tokens per round, versus the +q7 steady B8 result of `362.53 token/s` and mean accepted length `4.49`. The +31.3% throughput loss rejects q3 despite its higher per-position acceptance; +the shorter proposal cannot amortize the target round. The artifact is +`.artifacts/runtime/endpoint-qpn8-q3-b8-screen-v2`. + +The historical B4 trace also attributed `5.31 ms` and about 133 launches per +round to TP all-reduce. The accepted M8 push collective handled only the +80-KiB `[8,5120]` payload. A default-off +`VLLM_SM70_TP4_PUSH_ALLREDUCE_CONCURRENCY` candidate expands its IPC slot and +uses a grid-stride loop for M16/M32 without changing rank-ordered FP32 +accumulation. Across 128 consecutive collectives per CUDA Graph replay and +four input patterns, every rank is bitwise equal to the current custom-order +reference. Per-collective medians are: + +| Payload | Current pull | Push candidate | Saving | +|---:|---:|---:|---:| +| M16 / 160 KiB | 18.45 us | 11.03 us | 40.2% | +| M32 / 320 KiB | 26.78 us | 18.36 us | 31.5% | + +M64 measured `34.08 us` versus `30.79 us` current and was removed from the +admission set. The retained operator artifacts are +`.artifacts/runtime/tp4-push-concurrency-control-r1.json` and +`.artifacts/runtime/tp4-push-concurrency-final-r1.json`. + +With exact QPN8 and the push candidate together, the endpoint measured B2 +`258.04 token/s` (68.7% efficiency) and B4 `317.11 token/s` (42.2%). B2 is +1.6% above the prior exact-M16 row. B4 raw throughput is below the prior +`322.68 token/s`, while mean emitted tokens per round also moved from `4.62` +to `4.39`; throughput divided by that acceptance length improves by about +3.4%, consistent with the operator saving but not enough to claim an absolute +B4 endpoint win. Both rows completed 16/16 requests at 512 output tokens with +no errors or invalid text. The switch remains default-off and the scaling +gates remain open. Raw results are in +`.artifacts/runtime/endpoint-qpn8-push-ar-b2-b4-v1`. + +A third Nsight Systems attempt changed capture termination from +`stop-shutdown` to `stop`, completed the B4 workload, and still crashed in +`cuptiActivityFlushAll` while exiting without generating a report. CUPTI is +therefore unsuitable for this process topology until the external tool/runtime +issue changes; do not spend another run on capture-end variations. diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index 14184a57c8..c11f69fab1 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -44927,3 +44927,61 @@ Interpretation: under the PR worktree's `.artifacts/runtime/endpoint-b1-b2-b4-b8-v2` and `.artifacts/runtime/endpoint-b2-b4-b8-v3` directories. GPU 0-3 returned to 4 MiB per rank after graceful shutdown. + +## 2026-09-04 mixed-NVFP4 DFlash2 concurrency optimization follow-up + +- The acceptance targets are B2/B4/B8 scaling efficiencies of 80%/70%/60% + relative to the `187.77 token/s` B1 row, or absolute throughput gates of + `300.43/525.76/901.30 token/s`. +- Default-off channel-FP8 QPN8 routes now cover exact M16 plus native dense + M32. The M32 kernel time-multiplexes the original split-12/16 warp ranges, + retains their FP32 reduction order, uses 28/36 KiB shared memory, and reads + each packed weight tile once for all 32 rows. M=17/18/24/31/32 is bitwise + equal to concatenated M8 calls for every production dense projection split. +- Formal mixed-checkpoint endpoint results are B2 `253.95 token/s` (67.6%) + with exact M16, B4 `309.06 token/s` (41.2%) with chunked M32, and B4 + `322.68 token/s` (43.0%) with native dense M32. Every row completed 16/16 + full outputs. The B2/B4 acceptance statistics remain healthy at + `47.62%/4.33` and `51.68%/4.62` for rate/mean length. +- Proposal temperature scale 0.85 was rejected: B2 was `258.38 token/s`, but + B4 fell to `313.46 token/s` versus native-M32 default proposal. A native M64 + route was also rejected and removed: despite bitwise operator outputs and a + warm-cache microbenchmark win, B8 regressed to `283.94 token/s`, 21.7% below + the steady `362.53 token/s` baseline. +- Retained source controls are `VLLM_SM70_FP8_QPN8_M16`, + `VLLM_SM70_FP8_QPN8_M32_CHUNKED`, and + `VLLM_SM70_FP8_QPN8_M32_NATIVE`, all default off. The benchmark now records + batch-invariance equality and supports real channel-scale checkpoint data. +- The throughput gates are not closed. Further row tiling is stopped. The next + high-yield branch must split draft versus target cost without CUPTI, then + test request/stream partitioning or replica topology; two Nsight attempts + crashed in `cuptiActivityFlushAll` during multiprocess shutdown and produced + no report. +- MRV2's default-off phase profiler incorrectly admitted only `method=mtp`, so + the DFlash2 service emitted no phase records. Its gate now also admits + `dflash` and `dspark`, matching the legacy runner. Targeted tests pass. +- The resulting synchronized B2/B4 diagnostic is not a throughput result, but + its stable full-batch CUDA-event medians localize the work: B2 is target + forward `37.64 ms`, target sample/state `1.29 ms`, draft `7.50 ms`, and total + GPU `46.65 ms`; B4 is `47.41/1.54/9.05/58.10 ms`, respectively. Target + forward consumes about 81% of both intervals and causes nearly all + B2-to-B4 growth. Selector/rejection micro-tuning is therefore not the next + primary optimization. +- A q3 B8 screen reduced the target shape from M64 to M32 but reached only + `249.14 token/s`, 31.3% below the q7 steady B8 result. Draft acceptance was + `72.85%`, yet the shorter proposal emitted only `3.19` tokens per round + versus q7's `4.49`; q3 is rejected as a scaling workaround. +- A default-off TP4 push all-reduce extension is bitwise equal to current + custom-order output over 128 consecutive graph collectives and four input + patterns. M16 improves `18.45 -> 11.03 us` per collective and M32 + `26.78 -> 18.36 us`; M64 regresses and is not admitted. The control is + `VLLM_SM70_TP4_PUSH_ALLREDUCE_CONCURRENCY`. +- Combined QPN8 plus push-AR endpoint results are B2 `258.04 token/s` (68.7%) + and B4 `317.11 token/s` (42.2%). B2 improves 1.6% over exact M16. B4's raw + value is below the prior `322.68 token/s`, but its mean acceptance length is + also lower (`4.39` versus `4.62`); acceptance-normalized round rate improves + about 3.4%. Both rows completed 16/16 full outputs without errors. Keep the + switch default-off because the absolute B4 endpoint gate did not improve. +- A third Nsight run used stop-only capture termination, completed the B4 + workload, and still crashed in `cuptiActivityFlushAll` without a report. + Do not retry this multiprocess CUPTI path until the external runtime changes. diff --git a/tests/test_envs.py b/tests/test_envs.py index 473155db0f..a28a84e8ed 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -127,6 +127,10 @@ def test_sm70_concurrency_tuning_envs( names = ( "VLLM_SM70_TP4_MTP_AR_BLOCK_TUNING", "VLLM_SM70_TOPK_TOPP_8_WARPS", + "VLLM_SM70_FP8_QPN8_M16", + "VLLM_SM70_FP8_QPN8_M32_CHUNKED", + "VLLM_SM70_FP8_QPN8_M32_NATIVE", + "VLLM_SM70_TP4_PUSH_ALLREDUCE_CONCURRENCY", ) for name in names: monkeypatch.delenv(name, raising=False) diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index ff140fa681..d96bde7f29 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -27,7 +27,9 @@ def test_sm70_v2_mtp_profile_gate(monkeypatch): assert runner._sm70_v2_mtp_profile_enabled() runner.speculative_config = SimpleNamespace(method="dflash") - assert not runner._sm70_v2_mtp_profile_enabled() + assert runner._sm70_v2_mtp_profile_enabled() + runner.speculative_config = SimpleNamespace(method="dspark") + assert runner._sm70_v2_mtp_profile_enabled() runner.speculative_config = SimpleNamespace(method="mtp") runner.is_last_pp_rank = False assert not runner._sm70_v2_mtp_profile_enabled() diff --git a/vllm/distributed/device_communicators/custom_all_reduce.py b/vllm/distributed/device_communicators/custom_all_reduce.py index 3bf6f6cf31..bcda9eaaf3 100644 --- a/vllm/distributed/device_communicators/custom_all_reduce.py +++ b/vllm/distributed/device_communicators/custom_all_reduce.py @@ -340,10 +340,17 @@ def __init__( mtp5_status = ( "enabled" if envs.VLLM_SM70_TP4_PUSH_ALLREDUCE_MTP5 else "disabled" ) + concurrency_status = ( + "enabled" + if envs.VLLM_SM70_TP4_PUSH_ALLREDUCE_CONCURRENCY + else "disabled" + ) logger.info( "SM70 TP4 SGLang-style push all-reduce enabled for the " "FP16 80-KiB verifier, 8-KiB decode, and 5-KiB Qwen4Exp " - "payloads; opt-in 25-KiB Qwen4Exp MTP4 payload is %s.", + "payloads; Qwen3.8 M16/M32 concurrency is %s and the " + "25-KiB Qwen4Exp MTP4 payload is %s.", + concurrency_status, mtp5_status, ) diff --git a/vllm/envs.py b/vllm/envs.py index 034b37df2d..ede27a29c8 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -166,6 +166,9 @@ 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_M16: bool = False + VLLM_SM70_FP8_QPN8_M32_CHUNKED: bool = False + VLLM_SM70_FP8_QPN8_M32_NATIVE: bool = False VLLM_SM70_QWEN4_EXP_ONLINE_QPN8: bool = False VLLM_SM70_QWEN38_FP16_GEMV: bool = False VLLM_SM70_QWEN38_FUSED_GDN_INPUT_FP16: bool = False @@ -230,6 +233,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_TP4_PUSH_ALLREDUCE_CONCURRENCY: bool = False VLLM_SM70_TP4_PUSH_ALLREDUCE_MTP5: bool = False VLLM_SM70_CUSTOM_AR_LIBRARY: str | None = None VLLM_SM70_TOP1_CUSTOM_AR: bool = False @@ -1722,6 +1726,22 @@ def _resolve_rust_frontend_path() -> str | None: # online route also stays opt-in because it requantizes checkpoint BF16 # attention, GDN, QSA, and mHC weights without calibration. "VLLM_SM70_FP8_QPN8": lambda: bool(int(os.getenv("VLLM_SM70_FP8_QPN8", "0"))), + # Opt-in Qwen3.8 DFlash2 B2 candidate. It keeps channel-FP8 weights in + # QPN8 form for exact M=9..16 projection shapes instead of reconstructing + # a full FP16 matrix before every GEMM. + "VLLM_SM70_FP8_QPN8_M16": lambda: bool( + int(os.getenv("VLLM_SM70_FP8_QPN8_M16", "0")) + ), + # Follow-on B4 experiment: replay the accepted M16 body in contiguous row + # chunks through M=32 while retaining the same strict shape allowlist. + "VLLM_SM70_FP8_QPN8_M32_CHUNKED": lambda: bool( + int(os.getenv("VLLM_SM70_FP8_QPN8_M32_CHUNKED", "0")) + ), + # Native dense-only B4 verifier experiment. It keeps the logical split-K + # reduction order while one CTA reuses each packed weight tile for M<=32. + "VLLM_SM70_FP8_QPN8_M32_NATIVE": lambda: bool( + int(os.getenv("VLLM_SM70_FP8_QPN8_M32_NATIVE", "0")) + ), "VLLM_SM70_QWEN4_EXP_ONLINE_QPN8": lambda: bool( int(os.getenv("VLLM_SM70_QWEN4_EXP_ONLINE_QPN8", "0")) ), @@ -2107,6 +2127,11 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_TP4_PUSH_ALLREDUCE": lambda: bool( int(os.getenv("VLLM_SM70_TP4_PUSH_ALLREDUCE", "1")) ), + # Opt-in Qwen3.8 DFlash2 extension of the TP4 push collective from the + # accepted M8 payload to M16/M32 verifier payloads. + "VLLM_SM70_TP4_PUSH_ALLREDUCE_CONCURRENCY": lambda: bool( + int(os.getenv("VLLM_SM70_TP4_PUSH_ALLREDUCE_CONCURRENCY", "0")) + ), # Exact Qwen3.8 MTP4 verifier payload: FP16 [5, 2560] (25 KiB). The # existing push allocation is sized for 80 KiB, so this changes dispatch # only. Keep opt-in until the TP4 dynamic-graph gate is recorded. diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index af27d8c3bc..25bad9cab2 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -298,7 +298,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): def _sm70_v2_mtp_profile_enabled(self) -> bool: return ( self.speculative_config is not None - and self.speculative_config.method == "mtp" + and self.speculative_config.method in ("mtp", "dflash", "dspark") and self.is_last_pp_rank and self.device.type == "cuda" and envs.VLLM_SM70_MTP_PROFILE From d444517f38a0796d85999bba678bc4f98f8ca16e Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 4 Sep 2026 04:44:40 +0800 Subject: [PATCH 7/9] [Kernel][SM70] Reuse NVFP4 weights across B2 verifier rows Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../kernels/benchmark_sm70_nvfp4_qpn2.py | 69 ++++- csrc/sm70_turbomind/ops/nvfp4_qpn2_sm70.cu | 247 +++++++++++------- .../sm70_quasar_nvfp4_dflash2_acceptance.md | 59 ++++- docs/design/sm70_v100_migration_control.md | 34 ++- tests/test_envs.py | 1 + vllm/envs.py | 6 + 6 files changed, 315 insertions(+), 101 deletions(-) diff --git a/benchmarks/kernels/benchmark_sm70_nvfp4_qpn2.py b/benchmarks/kernels/benchmark_sm70_nvfp4_qpn2.py index 69ba8b3850..c3b56c5e6f 100644 --- a/benchmarks/kernels/benchmark_sm70_nvfp4_qpn2.py +++ b/benchmarks/kernels/benchmark_sm70_nvfp4_qpn2.py @@ -303,6 +303,8 @@ def _run_projection( warmup: int, iterations: int, trials: int, + split_k_override: int | None, + accumulator_chains_override: int | None, ) -> dict[str, object]: from vllm import _sm70_ops as sm70_ops @@ -318,6 +320,10 @@ def _run_projection( ) k, n = qweight.shape split_k, accumulator_chains = QPN2_CONFIGS[(k, n)] + if split_k_override is not None: + split_k = split_k_override + if accumulator_chains_override is not None: + accumulator_chains = accumulator_chains_override x = torch.randn((m, k), dtype=torch.float16, device=device) * 0.1 tm_weight, tm_scales, meta = sm70_ops.nvfp4_sm70_prepare( @@ -341,6 +347,7 @@ def run_tm() -> torch.Tensor: return tm_final qpn_final = torch.empty((m, final_n), dtype=torch.float16, device=device) + run_qpn2_chunked = None if production: qpn_codes, qpn_scales = sm70_ops.nvfp4_qpn2_prepare_sm70( packed, projection.scales.to(device) @@ -374,11 +381,13 @@ def run_qpn2() -> torch.Tensor: packed, projection.scales.to(device) ) - def run_qpn2() -> torch.Tensor: + def run_candidate( + candidate_out: torch.Tensor, candidate_input: torch.Tensor + ) -> torch.Tensor: if projection.gated_silu: torch.ops._qpn2_candidate.gated( - qpn_final, - x, + candidate_out, + candidate_input, qpn_codes, qpn_scales, projection.inverse_global_scale, @@ -387,15 +396,30 @@ def run_qpn2() -> torch.Tensor: ) else: torch.ops._qpn2_candidate.gemm( - qpn_final, - x, + candidate_out, + candidate_input, qpn_codes, qpn_scales, projection.inverse_global_scale, split_k, accumulator_chains, ) - return qpn_final + return candidate_out + + def run_qpn2() -> torch.Tensor: + return run_candidate(qpn_final, x) + + if m > 8: + qpn_chunked = torch.empty_like(qpn_final) + + def run_qpn2_chunked() -> torch.Tensor: + for row in range(0, m, 8): + rows = min(8, m - row) + run_candidate( + qpn_chunked.narrow(0, row, rows), + x.narrow(0, row, rows), + ) + return qpn_chunked else: if extension is None: @@ -422,11 +446,26 @@ def run_qpn2() -> torch.Tensor: 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) + chunked_timing = None + qpn_chunked_output = None + if run_qpn2_chunked is not None: + chunked_graph, qpn_chunked_output = _capture(run_qpn2_chunked) + chunked_timing = _time_graph(chunked_graph, warmup, iterations, trials) tm_graph.replay() qpn_graph.replay() torch.accelerator.synchronize() reference = _fp32_reference(projection, x, device) + batch_invariance = None + if qpn_chunked_output is not None: + run_qpn2_chunked() + torch.accelerator.synchronize() + batch_invariance = _quality(qpn_output, qpn_chunked_output) saved_us = float(tm_timing["median_us"]) - float(qpn_timing["median_us"]) + native_saved_us = ( + float(chunked_timing["median_us"]) - float(qpn_timing["median_us"]) + if chunked_timing is not None + else None + ) return { "name": projection.name, "m": m, @@ -440,11 +479,23 @@ def run_qpn2() -> torch.Tensor: }, "turbomind": tm_timing, "qpn2": qpn_timing, + "qpn2_chunked_m8": chunked_timing, "speedup": float(tm_timing["median_us"]) / float(qpn_timing["median_us"]), "saved_ms_per_round": saved_us * projection.calls_per_round / 1000.0, + "native_vs_chunked_speedup": ( + float(chunked_timing["median_us"]) / float(qpn_timing["median_us"]) + if chunked_timing is not None + else None + ), + "native_saved_ms_per_round": ( + native_saved_us * projection.calls_per_round / 1000.0 + if native_saved_us is not None + else None + ), "quality_vs_turbomind": _quality(qpn_output, tm_output), "turbomind_quality_vs_fp32": _quality(tm_output, reference), "qpn2_quality_vs_fp32": _quality(qpn_output, reference), + "qpn2_batch_invariance": batch_invariance, } @@ -461,6 +512,8 @@ def _parse_args() -> argparse.Namespace: 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("--split-k-override", type=int, choices=(8, 16, 32)) + parser.add_argument("--accumulator-chains-override", type=int, choices=(1, 2)) parser.add_argument( "--synthetic-shapes", action="store_true", @@ -560,6 +613,8 @@ def main() -> int: args.warmup, args.iterations, args.trials, + args.split_k_override, + args.accumulator_chains_override, ) for projection in projections ] @@ -596,6 +651,8 @@ def main() -> int: "warmup": args.warmup, "iterations": args.iterations, "trials": args.trials, + "split_k_override": args.split_k_override, + "accumulator_chains_override": args.accumulator_chains_override, "rows": rows, "projected_total_saved_ms_per_round": sum( float(row["saved_ms_per_round"]) for row in rows diff --git a/csrc/sm70_turbomind/ops/nvfp4_qpn2_sm70.cu b/csrc/sm70_turbomind/ops/nvfp4_qpn2_sm70.cu index e9ea30ff01..668bf61a10 100644 --- a/csrc/sm70_turbomind/ops/nvfp4_qpn2_sm70.cu +++ b/csrc/sm70_turbomind/ops/nvfp4_qpn2_sm70.cu @@ -13,6 +13,10 @@ #include #include +#include +#include +#include + #ifndef VLLM_NVFP4_QPN2_STANDALONE void silu_and_mul(torch::Tensor& out, torch::Tensor& input); @@ -122,20 +126,22 @@ __device__ __forceinline__ void dequant_e2m1x8(unsigned packed, half2 scale, "+f"(C[5]), "+f"(C[6]), "+f"(C[7]) \ : "r"(A0), "r"(A1), "r"(B0), "r"(B1)) -template +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]; + static_assert(RowTiles == 1 || RowTiles == 2, + "NVFP4 QPN2 supports one or two 8-row tiles"); + __shared__ float partials[SplitK][RowTiles * 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 local_row = (lane & 3) + ((lane & 16) ? 4 : 0); - const int row = blockIdx.y * kQpn2RowsPerCta + local_row; + const int row_base = blockIdx.y * kQpn2RowsPerCta * RowTiles; const int groups_k16 = k >> 4; const int groups_per_warp = groups_k16 / SplitK; const int group_begin = warp * groups_per_warp; @@ -145,12 +151,15 @@ __global__ void nvfp4_qpn2_sm70_kernel(const uint8_t* __restrict__ codes, group_scales + static_cast(tile) * groups_k16 * 32 + lane; const half2 global_scale2 = __float2half2_rn(global_scale * 16384.0f); - float accum[NAcc][8]; + float accum[RowTiles][NAcc][8]; #pragma unroll - for (int chain = 0; chain < NAcc; ++chain) { + for (int row_tile = 0; row_tile < RowTiles; ++row_tile) { #pragma unroll - for (int index = 0; index < 8; ++index) { - accum[chain][index] = 0.0f; + for (int chain = 0; chain < NAcc; ++chain) { +#pragma unroll + for (int index = 0; index < 8; ++index) { + accum[row_tile][chain][index] = 0.0f; + } } } @@ -165,47 +174,59 @@ __global__ void nvfp4_qpn2_sm70_kernel(const uint8_t* __restrict__ codes, 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 row_tile = 0; row_tile < RowTiles; ++row_tile) { + uint4 input01 = make_uint4(0, 0, 0, 0); + uint4 input23 = make_uint4(0, 0, 0, 0); + const int row = row_base + row_tile * kQpn2RowsPerCta + local_row; + 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); + VLLM_SM70_QPN2_MMA(accum[row_tile][0], a0[0], a0[1], b[0], b[1]); + VLLM_SM70_QPN2_MMA(accum[row_tile][1 % NAcc], a0[2], a0[3], b[2], b[3]); + VLLM_SM70_QPN2_MMA(accum[row_tile][2 % NAcc], a1[0], a1[1], b[4], b[5]); + VLLM_SM70_QPN2_MMA(accum[row_tile][3 % NAcc], a1[2], a1[3], b[6], b[7]); + } } #pragma unroll - for (int chain = 1; chain < NAcc; ++chain) { + for (int row_tile = 0; row_tile < RowTiles; ++row_tile) { #pragma unroll - for (int index = 0; index < 8; ++index) { - accum[0][index] += accum[chain][index]; + for (int chain = 1; chain < NAcc; ++chain) { +#pragma unroll + for (int index = 0; index < 8; ++index) { + accum[row_tile][0][index] += accum[row_tile][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]; + for (int row_tile = 0; row_tile < RowTiles; ++row_tile) { +#pragma unroll + for (int index = 0; index < 8; ++index) { + const int output_row = row_tile * kQpn2RowsPerCta + (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[row_tile][0][index]; + } } __syncthreads(); - for (int element = threadIdx.x; element < 256; element += blockDim.x) { + for (int element = threadIdx.x; element < RowTiles * 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 = blockIdx.y * kQpn2RowsPerCta + (element >> 5); + const int output_row = row_base + (element >> 5); const int output_col = element & 31; if (output_row < m) { output[static_cast(output_row) * n + tile * 32 + output_col] = @@ -214,12 +235,14 @@ __global__ void nvfp4_qpn2_sm70_kernel(const uint8_t* __restrict__ codes, } } -template +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]; + static_assert(RowTiles == 1 || RowTiles == 2, + "NVFP4 gated QPN2 supports one or two 8-row tiles"); + __shared__ float partials[2][SplitK][RowTiles * 256]; const int lane = threadIdx.x & 31; const int warp_in_block = threadIdx.x >> 5; @@ -229,7 +252,7 @@ __global__ void nvfp4_qpn2_gated_sm70_kernel( const int tile = blockIdx.x + projection * hidden_tiles; const int quadpair = (lane >> 2) & 3; const int local_row = (lane & 3) + ((lane & 16) ? 4 : 0); - const int row = blockIdx.y * kQpn2RowsPerCta + local_row; + const int row_base = blockIdx.y * kQpn2RowsPerCta * RowTiles; const int groups_k16 = k >> 4; const int groups_per_warp = groups_k16 / SplitK; const int group_begin = warp * groups_per_warp; @@ -239,12 +262,15 @@ __global__ void nvfp4_qpn2_gated_sm70_kernel( group_scales + static_cast(tile) * groups_k16 * 32 + lane; const half2 global_scale2 = __float2half2_rn(global_scale * 16384.0f); - float accum[NAcc][8]; + float accum[RowTiles][NAcc][8]; #pragma unroll - for (int chain = 0; chain < NAcc; ++chain) { + for (int row_tile = 0; row_tile < RowTiles; ++row_tile) { #pragma unroll - for (int index = 0; index < 8; ++index) { - accum[chain][index] = 0.0f; + for (int chain = 0; chain < NAcc; ++chain) { +#pragma unroll + for (int index = 0; index < 8; ++index) { + accum[row_tile][chain][index] = 0.0f; + } } } @@ -259,41 +285,53 @@ __global__ void nvfp4_qpn2_gated_sm70_kernel( 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 row_tile = 0; row_tile < RowTiles; ++row_tile) { + uint4 input01 = make_uint4(0, 0, 0, 0); + uint4 input23 = make_uint4(0, 0, 0, 0); + const int row = row_base + row_tile * kQpn2RowsPerCta + local_row; + 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); + VLLM_SM70_QPN2_MMA(accum[row_tile][0], a0[0], a0[1], b[0], b[1]); + VLLM_SM70_QPN2_MMA(accum[row_tile][1 % NAcc], a0[2], a0[3], b[2], b[3]); + VLLM_SM70_QPN2_MMA(accum[row_tile][2 % NAcc], a1[0], a1[1], b[4], b[5]); + VLLM_SM70_QPN2_MMA(accum[row_tile][3 % NAcc], a1[2], a1[3], b[6], b[7]); + } } #pragma unroll - for (int chain = 1; chain < NAcc; ++chain) { + for (int row_tile = 0; row_tile < RowTiles; ++row_tile) { #pragma unroll - for (int index = 0; index < 8; ++index) { - accum[0][index] += accum[chain][index]; + for (int chain = 1; chain < NAcc; ++chain) { +#pragma unroll + for (int index = 0; index < 8; ++index) { + accum[row_tile][0][index] += accum[row_tile][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]; + for (int row_tile = 0; row_tile < RowTiles; ++row_tile) { +#pragma unroll + for (int index = 0; index < 8; ++index) { + const int output_row = row_tile * kQpn2RowsPerCta + (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[row_tile][0][index]; + } } __syncthreads(); - for (int element = threadIdx.x; element < 256; element += blockDim.x) { + for (int element = threadIdx.x; element < RowTiles * 256; + element += blockDim.x) { float gate = 0.0f; float up = 0.0f; #pragma unroll @@ -301,7 +339,7 @@ __global__ void nvfp4_qpn2_gated_sm70_kernel( gate += partials[0][k_warp][element]; up += partials[1][k_warp][element]; } - const int output_row = blockIdx.y * kQpn2RowsPerCta + (element >> 5); + const int output_row = row_base + (element >> 5); const int output_col = element & 31; if (output_row < m) { // Match the existing SM70 silu_and_mul contract: round both GEMM @@ -318,25 +356,43 @@ __global__ void nvfp4_qpn2_gated_sm70_kernel( } } -template +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) { - const dim3 grid(n / 32, (m + kQpn2RowsPerCta - 1) / kQpn2RowsPerCta); - nvfp4_qpn2_sm70_kernel<<>>( - codes, scales, input, output, n, k, m, global_scale); + constexpr int kRowsPerCta = kQpn2RowsPerCta * RowTiles; + const dim3 grid(n / 32, (m + kRowsPerCta - 1) / kRowsPerCta); + nvfp4_qpn2_sm70_kernel + <<>>(codes, scales, input, output, n, k, + m, global_scale); } -template +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) { - const dim3 grid(hidden / 32, (m + kQpn2RowsPerCta - 1) / kQpn2RowsPerCta); - nvfp4_qpn2_gated_sm70_kernel + constexpr int kRowsPerCta = kQpn2RowsPerCta * RowTiles; + const dim3 grid(hidden / 32, (m + kRowsPerCta - 1) / kRowsPerCta); + nvfp4_qpn2_gated_sm70_kernel <<>>(codes, scales, input, output, hidden, k, m, global_scale); } +bool qpn2_m16_native_enabled(int m) { + const char* value = std::getenv("VLLM_SM70_NVFP4_QPN2_M16_NATIVE"); + const bool enabled = m > kQpn2RowsPerCta && m <= 16 && value != nullptr && + value[0] == '1' && value[1] == '\0'; + if (enabled) { + static std::once_flag m16_log_once; + std::call_once(m16_log_once, []() { + std::fprintf(stderr, + "INFO SM70 NVFP4 QPN2 native M=9..16 two-row-tile " + "candidate enabled.\n"); + }); + } + return enabled; +} + void check_qpn2_tensors(const torch::Tensor& out, const torch::Tensor& input, const torch::Tensor& codes, const torch::Tensor& scales, bool gated_silu) { @@ -436,21 +492,31 @@ void nvfp4_qpn2_gemm_sm70_out(torch::Tensor out, torch::Tensor input, 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); +#define VLLM_LAUNCH_QPN2(ROWS, SPLIT, NACC) \ + launch_qpn2(code_ptr, scale_ptr, input_ptr, output_ptr, \ + n, k, m, static_cast(global_scale), \ + stream) + const bool native_two_tile = qpn2_m16_native_enabled(m); + if (native_two_tile && split_k == 8 && accumulator_chains == 1) { + VLLM_LAUNCH_QPN2(2, 8, 1); + } else if (native_two_tile && split_k == 8) { + VLLM_LAUNCH_QPN2(2, 8, 2); + } else if (native_two_tile && split_k == 16 && accumulator_chains == 1) { + VLLM_LAUNCH_QPN2(2, 16, 1); + } else if (native_two_tile && split_k == 16) { + VLLM_LAUNCH_QPN2(2, 16, 2); + } else if (split_k == 8 && accumulator_chains == 1) { + VLLM_LAUNCH_QPN2(1, 8, 1); } else if (split_k == 8) { - VLLM_LAUNCH_QPN2(8, 2); + VLLM_LAUNCH_QPN2(1, 8, 2); } else if (split_k == 16 && accumulator_chains == 1) { - VLLM_LAUNCH_QPN2(16, 1); + VLLM_LAUNCH_QPN2(1, 16, 1); } else if (split_k == 16) { - VLLM_LAUNCH_QPN2(16, 2); + VLLM_LAUNCH_QPN2(1, 16, 2); } else if (accumulator_chains == 1) { - VLLM_LAUNCH_QPN2(32, 1); + VLLM_LAUNCH_QPN2(1, 32, 1); } else { - VLLM_LAUNCH_QPN2(32, 2); + VLLM_LAUNCH_QPN2(1, 32, 2); } #undef VLLM_LAUNCH_QPN2 C10_CUDA_KERNEL_LAUNCH_CHECK(); @@ -479,18 +545,23 @@ void nvfp4_qpn2_gated_sm70_out(torch::Tensor out, torch::Tensor input, 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); +#define VLLM_LAUNCH_QPN2_GATED(ROWS, SPLIT, NACC) \ + launch_qpn2_gated( \ + code_ptr, scale_ptr, input_ptr, output_ptr, hidden, k, m, \ + static_cast(global_scale), stream) + const bool native_two_tile = qpn2_m16_native_enabled(m); + if (native_two_tile && split_k == 8 && accumulator_chains == 1) { + VLLM_LAUNCH_QPN2_GATED(2, 8, 1); + } else if (native_two_tile && split_k == 8) { + VLLM_LAUNCH_QPN2_GATED(2, 8, 2); + } else if (split_k == 8 && accumulator_chains == 1) { + VLLM_LAUNCH_QPN2_GATED(1, 8, 1); } else if (split_k == 8) { - VLLM_LAUNCH_QPN2_GATED(8, 2); + VLLM_LAUNCH_QPN2_GATED(1, 8, 2); } else if (accumulator_chains == 1) { - VLLM_LAUNCH_QPN2_GATED(16, 1); + VLLM_LAUNCH_QPN2_GATED(1, 16, 1); } else { - VLLM_LAUNCH_QPN2_GATED(16, 2); + VLLM_LAUNCH_QPN2_GATED(1, 16, 2); } #undef VLLM_LAUNCH_QPN2_GATED C10_CUDA_KERNEL_LAUNCH_CHECK(); diff --git a/docs/design/sm70_quasar_nvfp4_dflash2_acceptance.md b/docs/design/sm70_quasar_nvfp4_dflash2_acceptance.md index 44dd40a6aa..c3d2b9cfc4 100644 --- a/docs/design/sm70_quasar_nvfp4_dflash2_acceptance.md +++ b/docs/design/sm70_quasar_nvfp4_dflash2_acceptance.md @@ -451,10 +451,10 @@ still needs a structural change rather than more row tiling. At observed accepted lengths, the remaining gap cannot be closed by selector or epilogue micro-tuning alone. The next measurement should split the unprofiled native candidate into target forward, target logits/rejection, draft, and host -bookkeeping, then evaluate either multi-stream/request partitioning or a -deployment topology that adds independent replicas. Two Nsight Systems runs -were attempted, but CUPTI crashed during multiprocess shutdown before writing -a report; do not repeat that capture path unchanged. +bookkeeping, then evaluate multi-stream or request partitioning inside the +same TP4 instance. Two Nsight Systems runs were attempted, but CUPTI crashed +during multiprocess shutdown before writing a report; do not repeat that +capture path unchanged. The built-in CUDA-event profiler initially produced no output because the MRV2 gate admitted `mtp` only, while this service resolves the same diagnostic path @@ -518,3 +518,54 @@ A third Nsight Systems attempt changed capture termination from `cuptiActivityFlushAll` while exiting without generating a report. CUPTI is therefore unsuitable for this process topology until the external tool/runtime issue changes; do not spend another run on capture-end variations. + +The synchronized MRV2 profiler was then extended to the official q7 B8/M64 +shape. Across 43 stable full-batch rounds, median target forward was +`54.103 ms`, target sampling `1.980 ms`, state update `0.044 ms`, draft +`12.041 ms`, and total GPU interval `68.290 ms`. Target forward therefore +accounts for 79.2% of the interval, draft 17.6%, and sampling plus state only +3.0%. The profiler's synchronized `262.45 token/s` endpoint result is not a +speed measurement. The retained phase records are in +`.artifacts/runtime/endpoint-qpn8-mrv2-profile-b8-v1`; they confirm that target +forward, rather than selector or host bookkeeping, remains the B8 bottleneck. + +For B2/M16, the NVFP4 QPN2 MLP path previously issued two independent M8 +kernels and loaded the same packed weights twice. A default-off +`VLLM_SM70_NVFP4_QPN2_M16_NATIVE` candidate assigns both eight-row groups to +one CTA, reuses each packed weight tile, and preserves the existing split-K +and FP32 reduction order. On real layer-55 TP4-rank-0 weights, gate/up improves +from `72.30` to `66.76 us` and down projection from `38.06` to `31.31 us`, a +weighted `0.786 ms` saving per target round. M9, M15, and M16 outputs are +bitwise equal to concatenated M8 calls for both projections with maximum +difference zero. The operator artifacts are +`.artifacts/runtime/qpn2-m16-native-real-final-r2.json` and +`.artifacts/runtime/qpn2-m16-native-tails-m9-r1.json` / +`.artifacts/runtime/qpn2-m16-native-tails-m15-r1.json`. + +With exact QPN8, TP4 push all-reduce, and native NVFP4 QPN2 M16 enabled, two +same-contract TP4/B2 endpoint runs measured `271.46` and `270.20 token/s`. +The final source-matched row is `270.20 token/s`, or 72.0% ideal scaling from +the fixed `187.77 token/s` B1 denominator. It completed 16/16 requests and all +8,192 requested output tokens, with no request errors, empty strings, or +replacement characters. Draft acceptance was `45.30%` and mean accepted +length `4.17`. The final artifact is +`.artifacts/runtime/endpoint-qpn2-m16-native-b2-final-v1` and records extension +SHA256 `10440281536d1364a2faa3b6c71129189aa1ccd6ee91aec598477a7d7d2b35f7`. + +Three additional branches were rejected. A bitwise-exact M32 NVFP4 two-row +CTA was slower than the retained path, including down projection +`140.94 us` versus approximately `71.59 us`, and was removed. A q5 TP4/B8 +endpoint screen reached only `335.71 token/s`, 7.4% below q7, because mean +emission fell to `3.93` tokens per round despite `58.58%` acceptance. A +single-accumulator channel-FP8 QPN8 M16 variant reached `274.28 token/s` in +one B2 sample, but acceptance-normalized round rate was about 0.6% worse and +the altered reduction order weakened its quality contract; that source was +also removed. + +The B2/B4/B8 gates remain open at `300.43/525.76/901.30 token/s`. B2 is now +about 10.1% below its gate, while the retained B4 and B8 rows remain +`322.68/362.53 token/s`. Generic vLLM DBO is not directly applicable: it +targets DP+EP/DeepEP, is unsupported by ModelRunnerV2, and the present service +is dense TP4/DP1. The next structural branch must therefore stay within the +single TP4 instance and prototype MRV2-local two-way verifier microbatching or +request partitioning, with target arithmetic and output quality held fixed. diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index c11f69fab1..782b6c7a69 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -44954,9 +44954,9 @@ Interpretation: batch-invariance equality and supports real channel-scale checkpoint data. - The throughput gates are not closed. Further row tiling is stopped. The next high-yield branch must split draft versus target cost without CUPTI, then - test request/stream partitioning or replica topology; two Nsight attempts - crashed in `cuptiActivityFlushAll` during multiprocess shutdown and produced - no report. + test request/stream partitioning inside the same TP4 instance; two Nsight + attempts crashed in `cuptiActivityFlushAll` during multiprocess shutdown and + produced no report. - MRV2's default-off phase profiler incorrectly admitted only `method=mtp`, so the DFlash2 service emitted no phase records. Its gate now also admits `dflash` and `dspark`, matching the legacy runner. Targeted tests pass. @@ -44985,3 +44985,31 @@ Interpretation: - A third Nsight run used stop-only capture termination, completed the B4 workload, and still crashed in `cuptiActivityFlushAll` without a report. Do not retry this multiprocess CUPTI path until the external runtime changes. +- The TP4/B8 synchronized MRV2 phase probe now has 43 stable full-batch q7/M64 + rounds. Median target forward is `54.103 ms`, target sample plus state + `2.024 ms`, draft `12.041 ms`, and total GPU `68.290 ms`. Target forward is + 79.2% of the interval, so selector work is not the primary B8 bottleneck. + The profiler-synchronized endpoint throughput is deliberately excluded. +- A default-off `VLLM_SM70_NVFP4_QPN2_M16_NATIVE` route reuses each packed + NVFP4 weight tile across the two verifier row groups. On real layer-55 TP4 + shards, gate/up improves `72.30 -> 66.76 us` and down projection + `38.06 -> 31.31 us`, saving a projected `0.786 ms` per target round versus + concatenated M8 calls. M9/M15/M16 gate/up and down outputs are bitwise equal + to that existing order with maximum difference zero. +- Two same-contract single-instance TP4/B2 endpoint runs with retained QPN8, + push all-reduce, and native QPN2 M16 measured `271.46` and + `270.20 token/s`. The final source-matched result is `270.20 token/s`, or + 72.0% scaling efficiency versus fixed B1 `187.77 token/s`. It completed + 16/16 requests and 8,192/8,192 output tokens without errors, empty text, or + replacement characters; acceptance was `45.30%` with mean length `4.17`. +- NVFP4 QPN2 M32 was removed because the exact candidate regressed down + projection to `140.94 us` versus approximately `71.59 us`. q5 B8 was also + rejected at `335.71 token/s`, 7.4% below q7, and a single-accumulator QPN8 + candidate was removed because its acceptance-normalized B2 rate regressed + about 0.6% while changing reduction order. +- The absolute gates remain open: B2 is about 10.1% below `300.43 token/s`, + while retained B4/B8 are `322.68/362.53 token/s` versus + `525.76/901.30`. Do not change tensor parallelism or use replica topology for + this acceptance task. Generic DBO targets DP+EP/DeepEP and is unsupported by + ModelRunnerV2, so the next branch is a TP4/DP1 MRV2-local verifier + microbatch/request-partitioning prototype. diff --git a/tests/test_envs.py b/tests/test_envs.py index a28a84e8ed..0bff8cf31b 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -130,6 +130,7 @@ def test_sm70_concurrency_tuning_envs( "VLLM_SM70_FP8_QPN8_M16", "VLLM_SM70_FP8_QPN8_M32_CHUNKED", "VLLM_SM70_FP8_QPN8_M32_NATIVE", + "VLLM_SM70_NVFP4_QPN2_M16_NATIVE", "VLLM_SM70_TP4_PUSH_ALLREDUCE_CONCURRENCY", ) for name in names: diff --git a/vllm/envs.py b/vllm/envs.py index ede27a29c8..f0c66176a7 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -181,6 +181,7 @@ 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_M16_NATIVE: bool = False VLLM_SM70_NVFP4_QPN2_PREFILL: bool = False VLLM_SM70_NVFP4_QPN2_PREFILL_LIBRARY: str | None = None VLLM_SM70_NVFP4_QPN2_PREFILL_MIN_M: int = 1024 @@ -1795,6 +1796,11 @@ 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 each packed NVFP4 tile across two eight-row verifier groups in one + # CTA. This is a default-off Qwen3.8 DFlash2 B2 operator candidate. + "VLLM_SM70_NVFP4_QPN2_M16_NATIVE": lambda: bool( + int(os.getenv("VLLM_SM70_NVFP4_QPN2_M16_NATIVE", "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. From 8df5cc7b174105ba5fa818a46cb342d0b911b568 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 4 Sep 2026 18:25:40 +0800 Subject: [PATCH 8/9] [Core][SM70] Trace DFlash2 concurrency boundaries Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../kernels/benchmark_sm70_nvfp4_qpn2.py | 7 + docs/design/sm70_v100_migration_control.md | 59 ++++++++ tests/test_envs.py | 9 ++ .../engine/test_sm70_dflash2_batch_trace.py | 63 ++++++++ vllm/envs.py | 11 ++ vllm/v1/engine/core.py | 143 ++++++++++++++++++ .../gpu/spec_decode/dflash/speculator.py | 80 ++++++++++ 7 files changed, 372 insertions(+) create mode 100644 tests/v1/engine/test_sm70_dflash2_batch_trace.py diff --git a/benchmarks/kernels/benchmark_sm70_nvfp4_qpn2.py b/benchmarks/kernels/benchmark_sm70_nvfp4_qpn2.py index c3b56c5e6f..3fd56a8566 100644 --- a/benchmarks/kernels/benchmark_sm70_nvfp4_qpn2.py +++ b/benchmarks/kernels/benchmark_sm70_nvfp4_qpn2.py @@ -60,6 +60,11 @@ def _sha256_file(path: Path) -> str: return digest.hexdigest() +def _sha256_tensor(tensor: torch.Tensor) -> str: + raw = tensor.detach().contiguous().view(torch.uint8).cpu().numpy().tobytes() + return hashlib.sha256(raw).hexdigest() + + def _load_projection_shards( model: Path, layer_index: int, @@ -493,6 +498,8 @@ def run_qpn2() -> torch.Tensor: else None ), "quality_vs_turbomind": _quality(qpn_output, tm_output), + "turbomind_output_sha256": _sha256_tensor(tm_output), + "qpn2_output_sha256": _sha256_tensor(qpn_output), "turbomind_quality_vs_fp32": _quality(tm_output, reference), "qpn2_quality_vs_fp32": _quality(qpn_output, reference), "qpn2_batch_invariance": batch_invariance, diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index 98dd44dca0..684dcaa02f 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -45013,6 +45013,65 @@ Interpretation: this acceptance task. Generic DBO targets DP+EP/DeepEP and is unsupported by ModelRunnerV2, so the next branch is a TP4/DP1 MRV2-local verifier microbatch/request-partitioning prototype. + +## 2026-09-04 DFlash2 concurrency scheduling and graph-boundary audit + +- A default-off completed-iteration trace now records decode request occupancy, + target M, queue depth, accepted draft length, scheduled/emitted tokens, and + completions from values already visible on the CPU. It creates no CUDA event + or synchronization. This avoids interpreting profiler-synchronized endpoint + throughput as a production result. +- Sixteen-request short traces show average request occupancy of about 91.6%, + 84.4%, and 85.9% for B2/B4/B8. The fraction of completely full rounds is + lower (83.2%, 59.6%, and 50.7%), but the B4/B8 occupancy loss relative to B2 + is only 6--8%. Acceptance length is also close across the three rows. Tail + batches and acceptance variation are real endpoint costs, but they are not + large enough to explain the missing scaling by themselves. +- Raising the async batch queue depth from two to four changed the short B8 + result by only +0.8% and made B2/B4 slower. An adaptive depth-four policy was + then measured on 16 requests by 512 outputs: B8 moved only + `366.47 -> 367.83 token/s` (+0.37%) while doing 5.1% more target work because + acceptance length fell `4.563 -> 4.384`. Both queue candidates were removed. +- DFlash2 q7 is one parallel query layout and one draft-model call, not seven + serial draft forwards. Query attention, candidate generation, selector walk, + and sampling are captured in one DFlash2 FULL CUDA graph. The eager work + before replay is target-hidden staging, input/slot preparation, and context + K/V materialization. +- A batched CUDA-event diagnostic sampled only pure-decode proposals and + synchronized once per 16-round window. Stable B2/B4/near-full-B8 windows + measured total DFlash proposal stages of `5.075/6.864/10.191 ms`. Their FULL + query graphs were `4.795/6.551/9.732 ms`; all graph-external work was only + `0.280/0.313/0.459 ms`. Thus 96.5% of the measured B2-to-B8 draft-stage + growth is inside the FULL graph, not metadata or host launch gaps. +- The engine trace shows schedule, execute submission, sample submission, and + scheduler update together remain below 1 ms in sampled steady steps, while + waiting for the GPU future is about 43--58 ms. Combined with the synchronized + phase medians, B2-to-B8 round growth is approximately 76% target forward, + 21% draft, and 3% target sampling/state. Scheduling/queue work is therefore + not the primary wall. +- Two current B8 graph-node Nsight attempts completed inference but crashed in + `cuptiActivityFlushAll` during multiprocess shutdown and emitted no report. + They are retained only as profiler-failure evidence; do not repeat this + capture until the external CUPTI runtime changes. +- A narrowly tuned M64 NVFP4 down projection improved the real layer-55 + operator from `137.53` to `78.84 us`, projecting about 3.76 ms saved per + target round. It changed the reduction path, however: B8 acceptance length + fell `4.563 -> 4.00` and endpoint throughput regressed + `366.47 -> 341.90 token/s`. A bitwise-preserving kernel/split variant then + measured `124.365 us` versus `124.346 us` control after clock warmup. Both + variants were removed; the operator benchmark now records raw output SHA256 + so future tactic races expose numerical changes directly. +- Extending the accepted B1-only DFlash2 QPN8 LM-head rerank to B2 by executing + its 14 rows as 8+6 preserved healthy acceptance (`4.36`) but reached only + `174.46 token/s` versus the retained `270.20 token/s` B2 result. Re-reading + the 318-MiB QPN8 layout and reranking twice loses decisively to the batched + dense LM head, so the prototype was removed without B4/B8 testing. +- The next viable work must reduce target FULL-graph time with a numerically + stable M64 path, or redesign draft candidate generation to process multiple + rows in one compressed-weight pass and rerank only row-local candidates. + Queue depth, tail-only policy, repeated 8-row rerank chunks, and unconstrained + TurboMind tactic tuning are closed directions for this contract. + ## 2026-09-03 Qwen3.8 unified prefill/decode compilation - The matched TP4/no-MTP evidence separated the regression from PLE residency. diff --git a/tests/test_envs.py b/tests/test_envs.py index 0bff8cf31b..5b7eb43ebc 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -132,6 +132,7 @@ def test_sm70_concurrency_tuning_envs( "VLLM_SM70_FP8_QPN8_M32_NATIVE", "VLLM_SM70_NVFP4_QPN2_M16_NATIVE", "VLLM_SM70_TP4_PUSH_ALLREDUCE_CONCURRENCY", + "VLLM_SM70_DFLASH2_BATCH_TRACE", ) for name in names: monkeypatch.delenv(name, raising=False) @@ -139,6 +140,14 @@ def test_sm70_concurrency_tuning_envs( monkeypatch.setenv(name, "1") assert environment_variables[name]() is True + name = "VLLM_SM70_DFLASH2_BATCH_TRACE_EVERY" + monkeypatch.delenv(name, raising=False) + assert environment_variables[name]() == 16 + monkeypatch.setenv(name, "0") + assert environment_variables[name]() == 1 + monkeypatch.setenv(name, "7") + assert environment_variables[name]() == 7 + name = "VLLM_SM70_MTP_MOE_TUNED_CONFIG" monkeypatch.delenv(name, raising=False) assert environment_variables[name]() is True diff --git a/tests/v1/engine/test_sm70_dflash2_batch_trace.py b/tests/v1/engine/test_sm70_dflash2_batch_trace.py new file mode 100644 index 0000000000..3a9862686d --- /dev/null +++ b/tests/v1/engine/test_sm70_dflash2_batch_trace.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest.mock import patch + +from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs +from vllm.v1.engine.core import EngineCore +from vllm.v1.outputs import ModelRunnerOutput + + +def test_sm70_dflash2_batch_trace_records_completed_decode_iteration() -> None: + core = EngineCore.__new__(EngineCore) + core._sm70_dflash2_batch_trace_enabled = True + core._sm70_dflash2_batch_trace_every = 16 + core.scheduler = SimpleNamespace(get_num_unfinished_requests=lambda: 0) + core._reset_sm70_dflash2_batch_trace() + + scheduled_spec_tokens = { + "r0": [10, 11, 12, 13, 14, 15, 16], + "r1": [20, 21, 22, 23, 24, 25, 26], + } + scheduler_output = SimpleNamespace( + scheduled_new_reqs=[], + scheduled_cached_reqs=SimpleNamespace(is_context_phase=lambda _req_id: False), + num_scheduled_tokens={"r0": 8, "r1": 8}, + scheduled_spec_decode_tokens=scheduled_spec_tokens, + num_invalid_spec_tokens={"r1": 1}, + ) + model_output = ModelRunnerOutput( + req_ids=["r0", "r1"], + req_id_to_index={"r0": 0, "r1": 1}, + sampled_token_ids=[[30, 31, 32, 33], [40, 41]], + ) + engine_core_outputs = { + 0: EngineCoreOutputs( + outputs=[ + EngineCoreOutput("r0", [30, 31, 32, 33]), + EngineCoreOutput("r1", [40, 41]), + ] + ) + } + + with patch("vllm.v1.engine.core.logger.info") as log_info: + core._record_sm70_dflash2_batch_trace( + scheduler_output, + model_output, + engine_core_outputs, + queue_len=0, + ) + + log_info.assert_called_once() + log_args = log_info.call_args.args + message = log_args[0] % log_args[1:] + assert "steps=1" in message + assert "phase_hist=decode:1" in message + assert "gen_req_hist=2:1" in message + assert "target_m_hist=16:1" in message + assert "accept_len_hist=2:1,4:1" in message + assert "queue_len_hist=0:1" in message + assert "draft_tokens=13" in message + assert "accepted_draft_tokens=4" in message + assert "sampled_tokens=6 emitted_tokens=6 finished_reqs=0" in message diff --git a/vllm/envs.py b/vllm/envs.py index 4026c4f0f7..7c04a3082a 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -250,6 +250,8 @@ VLLM_SM70_ASYNC_STAGED_INPUT_PREP: bool = False VLLM_SM70_ASYNC_CPU_TRACE: bool = False VLLM_SM70_ASYNC_CPU_TRACE_EVERY: int = 16 + VLLM_SM70_DFLASH2_BATCH_TRACE: bool = False + VLLM_SM70_DFLASH2_BATCH_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 @@ -2186,6 +2188,15 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_ASYNC_CPU_TRACE_EVERY": lambda: max( 1, int(os.getenv("VLLM_SM70_ASYNC_CPU_TRACE_EVERY", "16")) ), + # Completed-iteration trace for DFlash2 concurrency diagnostics. This only + # reads scheduler/model outputs that are already CPU-visible and never adds + # a CUDA event or synchronization to the measured path. + "VLLM_SM70_DFLASH2_BATCH_TRACE": lambda: bool( + int(os.getenv("VLLM_SM70_DFLASH2_BATCH_TRACE", "0")) + ), + "VLLM_SM70_DFLASH2_BATCH_TRACE_EVERY": lambda: max( + 1, int(os.getenv("VLLM_SM70_DFLASH2_BATCH_TRACE_EVERY", "16")) + ), # Legacy 0.0.3 diagnostic gate. Logs one TP all-reduce backend decision per # group/backend/shape/dtype so route-hit data can distinguish custom AR, # pynccl, flashinfer, symmetric-memory, and torch fallback paths. diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 9210b37eec..15010d94f3 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -201,6 +201,9 @@ def __init__( self.batch_queue_size, ) self._sm70_async_cpu_trace_step = 0 + self._sm70_dflash2_batch_trace_enabled = envs.VLLM_SM70_DFLASH2_BATCH_TRACE + self._sm70_dflash2_batch_trace_every = envs.VLLM_SM70_DFLASH2_BATCH_TRACE_EVERY + self._reset_sm70_dflash2_batch_trace() self.is_ec_consumer = ( vllm_config.ec_transfer_config is None @@ -448,6 +451,137 @@ def log_iteration_details(self, scheduler_output: SchedulerOutput | None): ) self._iteration_index += 1 + def _reset_sm70_dflash2_batch_trace(self) -> None: + self._sm70_dflash2_trace_steps = 0 + self._sm70_dflash2_trace_started_at = 0.0 + self._sm70_dflash2_trace_phase_hist: defaultdict[str, int] = defaultdict(int) + self._sm70_dflash2_trace_gen_req_hist: defaultdict[int, int] = defaultdict(int) + self._sm70_dflash2_trace_target_m_hist: defaultdict[int, int] = defaultdict(int) + self._sm70_dflash2_trace_accept_len_hist: defaultdict[int, int] = defaultdict( + int + ) + self._sm70_dflash2_trace_queue_len_hist: defaultdict[int, int] = defaultdict( + int + ) + self._sm70_dflash2_trace_unfinished_hist: defaultdict[int, int] = defaultdict( + int + ) + self._sm70_dflash2_trace_scheduled_gen_tokens = 0 + self._sm70_dflash2_trace_draft_tokens = 0 + self._sm70_dflash2_trace_accepted_draft_tokens = 0 + self._sm70_dflash2_trace_sampled_tokens = 0 + self._sm70_dflash2_trace_emitted_tokens = 0 + self._sm70_dflash2_trace_finished_reqs = 0 + + @staticmethod + def _format_sm70_dflash2_trace_hist(hist: dict[Any, int]) -> str: + return ",".join(f"{key}:{hist[key]}" for key in sorted(hist)) or "none" + + def _record_sm70_dflash2_batch_trace( + self, + scheduler_output: SchedulerOutput, + model_output: ModelRunnerOutput, + engine_core_outputs: dict[int, EngineCoreOutputs], + queue_len: int, + ) -> None: + """Aggregate completed scheduler iterations without synchronizing CUDA.""" + if not self._sm70_dflash2_batch_trace_enabled: + return + + now = time.perf_counter() + if self._sm70_dflash2_trace_steps == 0: + self._sm70_dflash2_trace_started_at = now + + details = compute_iteration_details(scheduler_output) + if details.num_ctx_requests and details.num_generation_requests: + phase = "mixed" + elif details.num_ctx_requests: + phase = "prefill" + elif details.num_generation_requests: + phase = "decode" + else: + phase = "empty" + + invalid_spec_tokens = scheduler_output.num_invalid_spec_tokens or {} + scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens + valid_draft_tokens = sum( + max(0, len(token_ids) - invalid_spec_tokens.get(req_id, 0)) + for req_id, token_ids in scheduled_spec_tokens.items() + ) + accepted_draft_tokens = 0 + sampled_tokens = 0 + for req_id, req_index in model_output.req_id_to_index.items(): + generated_tokens = ( + model_output.sampled_token_ids[req_index] + if model_output.sampled_token_ids + else [] + ) + sampled_tokens += len(generated_tokens) + if req_id in scheduled_spec_tokens: + accept_len = len(generated_tokens) + self._sm70_dflash2_trace_accept_len_hist[accept_len] += 1 + accepted_draft_tokens += max(0, accept_len - 1) + + emitted_tokens = 0 + finished_reqs = 0 + for client_output in engine_core_outputs.values(): + for request_output in client_output.outputs: + emitted_tokens += len(request_output.new_token_ids) + finished_reqs += int(request_output.finished) + + unfinished_reqs = self.scheduler.get_num_unfinished_requests() + self._sm70_dflash2_trace_steps += 1 + self._sm70_dflash2_trace_phase_hist[phase] += 1 + self._sm70_dflash2_trace_gen_req_hist[details.num_generation_requests] += 1 + self._sm70_dflash2_trace_target_m_hist[details.num_generation_tokens] += 1 + self._sm70_dflash2_trace_queue_len_hist[queue_len] += 1 + self._sm70_dflash2_trace_unfinished_hist[unfinished_reqs] += 1 + self._sm70_dflash2_trace_scheduled_gen_tokens += details.num_generation_tokens + self._sm70_dflash2_trace_draft_tokens += valid_draft_tokens + self._sm70_dflash2_trace_accepted_draft_tokens += accepted_draft_tokens + self._sm70_dflash2_trace_sampled_tokens += sampled_tokens + self._sm70_dflash2_trace_emitted_tokens += emitted_tokens + self._sm70_dflash2_trace_finished_reqs += finished_reqs + + flush = ( + self._sm70_dflash2_trace_steps >= self._sm70_dflash2_batch_trace_every + or (unfinished_reqs == 0 and queue_len == 0) + ) + if not flush: + return + + completion_span_ms = (now - self._sm70_dflash2_trace_started_at) * 1000.0 + logger.info( + "SM70 DFlash2 batch trace steps=%d completion_span_ms=%.3f " + "phase_hist=%s gen_req_hist=%s target_m_hist=%s " + "accept_len_hist=%s queue_len_hist=%s unfinished_hist=%s " + "scheduled_gen_tokens=%d draft_tokens=%d accepted_draft_tokens=%d " + "sampled_tokens=%d emitted_tokens=%d finished_reqs=%d", + self._sm70_dflash2_trace_steps, + completion_span_ms, + self._format_sm70_dflash2_trace_hist(self._sm70_dflash2_trace_phase_hist), + self._format_sm70_dflash2_trace_hist(self._sm70_dflash2_trace_gen_req_hist), + self._format_sm70_dflash2_trace_hist( + self._sm70_dflash2_trace_target_m_hist + ), + self._format_sm70_dflash2_trace_hist( + self._sm70_dflash2_trace_accept_len_hist + ), + self._format_sm70_dflash2_trace_hist( + self._sm70_dflash2_trace_queue_len_hist + ), + self._format_sm70_dflash2_trace_hist( + self._sm70_dflash2_trace_unfinished_hist + ), + self._sm70_dflash2_trace_scheduled_gen_tokens, + self._sm70_dflash2_trace_draft_tokens, + self._sm70_dflash2_trace_accepted_draft_tokens, + self._sm70_dflash2_trace_sampled_tokens, + self._sm70_dflash2_trace_emitted_tokens, + self._sm70_dflash2_trace_finished_reqs, + ) + self._reset_sm70_dflash2_batch_trace() + def step(self) -> tuple[dict[int, EngineCoreOutputs], bool]: """Schedule, execute, and make output. @@ -508,6 +642,9 @@ def step(self) -> tuple[dict[int, EngineCoreOutputs], bool]: engine_core_outputs = self.scheduler.update_from_output( scheduler_output, model_output ) + self._record_sm70_dflash2_batch_trace( + scheduler_output, model_output, engine_core_outputs, queue_len=0 + ) if profile_ddtree_engine: profile_update_ms = (time.perf_counter() - profile_part_t0) * 1000.0 logger.info( @@ -706,6 +843,12 @@ def step_with_batch_queue( engine_core_outputs = self.scheduler.update_from_output( scheduler_output, model_output ) + self._record_sm70_dflash2_batch_trace( + scheduler_output, + model_output, + engine_core_outputs, + queue_len=len(batch_queue), + ) if trace_log: trace_update_ms = (time.perf_counter() - trace_update_t0) * 1000.0 diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index c006e980f8..187d9408aa 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -7,6 +7,7 @@ import torch import torch.nn as nn +from vllm import envs from vllm.config import VllmConfig, replace from vllm.config.compilation import CUDAGraphMode from vllm.config.speculative import get_dflash_model_draft_tokens @@ -133,6 +134,63 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.draft_kv_cache_group_id: int = -1 self._context_only_prefill_logged = False self._query_slot_mappings: torch.Tensor | None = None + self._stage_profile_records: list[ + tuple[int, int, tuple[torch.cuda.Event, ...]] + ] = [] + + def _flush_stage_profile( + self, + num_reqs: int, + num_target_tokens: int, + events: tuple[torch.cuda.Event, ...] | None, + ) -> None: + """Report batched stream timing without synchronizing every decode step.""" + if events is None: + return + + self._stage_profile_records.append((num_reqs, num_target_tokens, events)) + interval = envs.VLLM_DFLASH_PROFILE_LOG_INTERVAL + if len(self._stage_profile_records) < interval: + return + + # One synchronization closes the whole window. Event-to-event elapsed + # time includes any exposed host launch gap on the current CUDA stream, + # which is exactly what this diagnostic is intended to find. + self._stage_profile_records[-1][2][-1].synchronize() + stage_names = ("hidden", "prepare", "context_kv", "metadata", "query_graph") + stage_values = {name: [] for name in stage_names} + total_values: list[float] = [] + reqs = 0 + target_tokens = 0 + for record_reqs, record_tokens, record_events in self._stage_profile_records: + reqs += record_reqs + target_tokens += record_tokens + for name, start, end in zip( + stage_names, record_events[:-1], record_events[1:] + ): + stage_values[name].append(start.elapsed_time(end)) + total_values.append(record_events[0].elapsed_time(record_events[-1])) + + rounds = len(self._stage_profile_records) + means = { + name: sum(values) / len(values) for name, values in stage_values.items() + } + total_mean = sum(total_values) / len(total_values) + logger.info( + "DFLASH_STAGE_PROFILE rounds=%d avg_reqs=%.2f avg_target_tokens=%.2f " + "hidden_ms=%.3f prepare_ms=%.3f context_kv_ms=%.3f " + "metadata_ms=%.3f query_graph_ms=%.3f total_ms=%.3f", + rounds, + reqs / rounds, + target_tokens / rounds, + means["hidden"], + means["prepare"], + means["context_kv"], + means["metadata"], + means["query_graph"], + total_mean, + ) + self._stage_profile_records.clear() @property def attn_vllm_config(self) -> VllmConfig: @@ -422,6 +480,17 @@ def propose( max_seq_len + self.num_query_per_req, self.max_model_len ) self._prepare_proposal_runtime(input_batch, num_sampled, num_rejected) + profile_events = None + if ( + envs.VLLM_DFLASH_PROFILE + and self.device.type == "cuda" + and not dummy_run + and not np.any(input_batch.is_prefilling_np[:num_reqs]) + ): + profile_events = tuple( + torch.cuda.Event(enable_timing=True) for _ in range(6) + ) + profile_events[0].record() # NOTE: To avoid CPU-GPU synchronization without CPU knowing the # number of rejected tokens, we maintain the size of input_ids and @@ -438,6 +507,8 @@ def propose( self.hidden_states[:num_target_tokens].copy_( hidden_states[:num_target_tokens] ) + if profile_events is not None: + profile_events[1].record() if dummy_run and skip_attn_for_dummy_run: # Memory profiling path: block_tables / kv_cache_config are not initialized. @@ -495,6 +566,8 @@ def propose( self.max_model_len, self.sample_from_anchor, ) + if profile_events is not None: + profile_events[2].record() # 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 @@ -515,6 +588,8 @@ def propose( self.context_positions[:num_target_tokens], context_slots, ) + if profile_events is not None: + profile_events[3].record() if not dummy_run and _is_context_only_prefill(input_batch): # Intermediate chunked-prefill steps only need to materialize the @@ -571,6 +646,8 @@ def propose( query_slot_mappings[:, :num_tokens_padded], self.kv_cache_config, ) + if profile_events is not None: + profile_events[4].record() if batch_desc.cg_mode == CUDAGraphMode.FULL: assert self.query_cudagraph_manager is not None @@ -584,8 +661,11 @@ def propose( num_tokens_across_dp=num_tokens_across_dp, cudagraph_runtime_mode=batch_desc.cg_mode, ) + if profile_events is not None: + profile_events[5].record() self._apply_ngram_assist(num_reqs) + self._flush_stage_profile(num_reqs, num_target_tokens, profile_events) return self.draft_tokens[:num_reqs] From 35092639ee7ad29ac8519a36f6fa9f874c6ee763 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 4 Sep 2026 18:36:23 +0800 Subject: [PATCH 9/9] Revert "[Core][SM70] Trace DFlash2 concurrency boundaries" This reverts commit 8df5cc7b174105ba5fa818a46cb342d0b911b568. --- .../kernels/benchmark_sm70_nvfp4_qpn2.py | 7 - docs/design/sm70_v100_migration_control.md | 59 -------- tests/test_envs.py | 9 -- .../engine/test_sm70_dflash2_batch_trace.py | 63 -------- vllm/envs.py | 11 -- vllm/v1/engine/core.py | 143 ------------------ .../gpu/spec_decode/dflash/speculator.py | 80 ---------- 7 files changed, 372 deletions(-) delete mode 100644 tests/v1/engine/test_sm70_dflash2_batch_trace.py diff --git a/benchmarks/kernels/benchmark_sm70_nvfp4_qpn2.py b/benchmarks/kernels/benchmark_sm70_nvfp4_qpn2.py index 3fd56a8566..c3b56c5e6f 100644 --- a/benchmarks/kernels/benchmark_sm70_nvfp4_qpn2.py +++ b/benchmarks/kernels/benchmark_sm70_nvfp4_qpn2.py @@ -60,11 +60,6 @@ def _sha256_file(path: Path) -> str: return digest.hexdigest() -def _sha256_tensor(tensor: torch.Tensor) -> str: - raw = tensor.detach().contiguous().view(torch.uint8).cpu().numpy().tobytes() - return hashlib.sha256(raw).hexdigest() - - def _load_projection_shards( model: Path, layer_index: int, @@ -498,8 +493,6 @@ def run_qpn2() -> torch.Tensor: else None ), "quality_vs_turbomind": _quality(qpn_output, tm_output), - "turbomind_output_sha256": _sha256_tensor(tm_output), - "qpn2_output_sha256": _sha256_tensor(qpn_output), "turbomind_quality_vs_fp32": _quality(tm_output, reference), "qpn2_quality_vs_fp32": _quality(qpn_output, reference), "qpn2_batch_invariance": batch_invariance, diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index 684dcaa02f..98dd44dca0 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -45013,65 +45013,6 @@ Interpretation: this acceptance task. Generic DBO targets DP+EP/DeepEP and is unsupported by ModelRunnerV2, so the next branch is a TP4/DP1 MRV2-local verifier microbatch/request-partitioning prototype. - -## 2026-09-04 DFlash2 concurrency scheduling and graph-boundary audit - -- A default-off completed-iteration trace now records decode request occupancy, - target M, queue depth, accepted draft length, scheduled/emitted tokens, and - completions from values already visible on the CPU. It creates no CUDA event - or synchronization. This avoids interpreting profiler-synchronized endpoint - throughput as a production result. -- Sixteen-request short traces show average request occupancy of about 91.6%, - 84.4%, and 85.9% for B2/B4/B8. The fraction of completely full rounds is - lower (83.2%, 59.6%, and 50.7%), but the B4/B8 occupancy loss relative to B2 - is only 6--8%. Acceptance length is also close across the three rows. Tail - batches and acceptance variation are real endpoint costs, but they are not - large enough to explain the missing scaling by themselves. -- Raising the async batch queue depth from two to four changed the short B8 - result by only +0.8% and made B2/B4 slower. An adaptive depth-four policy was - then measured on 16 requests by 512 outputs: B8 moved only - `366.47 -> 367.83 token/s` (+0.37%) while doing 5.1% more target work because - acceptance length fell `4.563 -> 4.384`. Both queue candidates were removed. -- DFlash2 q7 is one parallel query layout and one draft-model call, not seven - serial draft forwards. Query attention, candidate generation, selector walk, - and sampling are captured in one DFlash2 FULL CUDA graph. The eager work - before replay is target-hidden staging, input/slot preparation, and context - K/V materialization. -- A batched CUDA-event diagnostic sampled only pure-decode proposals and - synchronized once per 16-round window. Stable B2/B4/near-full-B8 windows - measured total DFlash proposal stages of `5.075/6.864/10.191 ms`. Their FULL - query graphs were `4.795/6.551/9.732 ms`; all graph-external work was only - `0.280/0.313/0.459 ms`. Thus 96.5% of the measured B2-to-B8 draft-stage - growth is inside the FULL graph, not metadata or host launch gaps. -- The engine trace shows schedule, execute submission, sample submission, and - scheduler update together remain below 1 ms in sampled steady steps, while - waiting for the GPU future is about 43--58 ms. Combined with the synchronized - phase medians, B2-to-B8 round growth is approximately 76% target forward, - 21% draft, and 3% target sampling/state. Scheduling/queue work is therefore - not the primary wall. -- Two current B8 graph-node Nsight attempts completed inference but crashed in - `cuptiActivityFlushAll` during multiprocess shutdown and emitted no report. - They are retained only as profiler-failure evidence; do not repeat this - capture until the external CUPTI runtime changes. -- A narrowly tuned M64 NVFP4 down projection improved the real layer-55 - operator from `137.53` to `78.84 us`, projecting about 3.76 ms saved per - target round. It changed the reduction path, however: B8 acceptance length - fell `4.563 -> 4.00` and endpoint throughput regressed - `366.47 -> 341.90 token/s`. A bitwise-preserving kernel/split variant then - measured `124.365 us` versus `124.346 us` control after clock warmup. Both - variants were removed; the operator benchmark now records raw output SHA256 - so future tactic races expose numerical changes directly. -- Extending the accepted B1-only DFlash2 QPN8 LM-head rerank to B2 by executing - its 14 rows as 8+6 preserved healthy acceptance (`4.36`) but reached only - `174.46 token/s` versus the retained `270.20 token/s` B2 result. Re-reading - the 318-MiB QPN8 layout and reranking twice loses decisively to the batched - dense LM head, so the prototype was removed without B4/B8 testing. -- The next viable work must reduce target FULL-graph time with a numerically - stable M64 path, or redesign draft candidate generation to process multiple - rows in one compressed-weight pass and rerank only row-local candidates. - Queue depth, tail-only policy, repeated 8-row rerank chunks, and unconstrained - TurboMind tactic tuning are closed directions for this contract. - ## 2026-09-03 Qwen3.8 unified prefill/decode compilation - The matched TP4/no-MTP evidence separated the regression from PLE residency. diff --git a/tests/test_envs.py b/tests/test_envs.py index 5b7eb43ebc..0bff8cf31b 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -132,7 +132,6 @@ def test_sm70_concurrency_tuning_envs( "VLLM_SM70_FP8_QPN8_M32_NATIVE", "VLLM_SM70_NVFP4_QPN2_M16_NATIVE", "VLLM_SM70_TP4_PUSH_ALLREDUCE_CONCURRENCY", - "VLLM_SM70_DFLASH2_BATCH_TRACE", ) for name in names: monkeypatch.delenv(name, raising=False) @@ -140,14 +139,6 @@ def test_sm70_concurrency_tuning_envs( monkeypatch.setenv(name, "1") assert environment_variables[name]() is True - name = "VLLM_SM70_DFLASH2_BATCH_TRACE_EVERY" - monkeypatch.delenv(name, raising=False) - assert environment_variables[name]() == 16 - monkeypatch.setenv(name, "0") - assert environment_variables[name]() == 1 - monkeypatch.setenv(name, "7") - assert environment_variables[name]() == 7 - name = "VLLM_SM70_MTP_MOE_TUNED_CONFIG" monkeypatch.delenv(name, raising=False) assert environment_variables[name]() is True diff --git a/tests/v1/engine/test_sm70_dflash2_batch_trace.py b/tests/v1/engine/test_sm70_dflash2_batch_trace.py deleted file mode 100644 index 3a9862686d..0000000000 --- a/tests/v1/engine/test_sm70_dflash2_batch_trace.py +++ /dev/null @@ -1,63 +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 - -from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs -from vllm.v1.engine.core import EngineCore -from vllm.v1.outputs import ModelRunnerOutput - - -def test_sm70_dflash2_batch_trace_records_completed_decode_iteration() -> None: - core = EngineCore.__new__(EngineCore) - core._sm70_dflash2_batch_trace_enabled = True - core._sm70_dflash2_batch_trace_every = 16 - core.scheduler = SimpleNamespace(get_num_unfinished_requests=lambda: 0) - core._reset_sm70_dflash2_batch_trace() - - scheduled_spec_tokens = { - "r0": [10, 11, 12, 13, 14, 15, 16], - "r1": [20, 21, 22, 23, 24, 25, 26], - } - scheduler_output = SimpleNamespace( - scheduled_new_reqs=[], - scheduled_cached_reqs=SimpleNamespace(is_context_phase=lambda _req_id: False), - num_scheduled_tokens={"r0": 8, "r1": 8}, - scheduled_spec_decode_tokens=scheduled_spec_tokens, - num_invalid_spec_tokens={"r1": 1}, - ) - model_output = ModelRunnerOutput( - req_ids=["r0", "r1"], - req_id_to_index={"r0": 0, "r1": 1}, - sampled_token_ids=[[30, 31, 32, 33], [40, 41]], - ) - engine_core_outputs = { - 0: EngineCoreOutputs( - outputs=[ - EngineCoreOutput("r0", [30, 31, 32, 33]), - EngineCoreOutput("r1", [40, 41]), - ] - ) - } - - with patch("vllm.v1.engine.core.logger.info") as log_info: - core._record_sm70_dflash2_batch_trace( - scheduler_output, - model_output, - engine_core_outputs, - queue_len=0, - ) - - log_info.assert_called_once() - log_args = log_info.call_args.args - message = log_args[0] % log_args[1:] - assert "steps=1" in message - assert "phase_hist=decode:1" in message - assert "gen_req_hist=2:1" in message - assert "target_m_hist=16:1" in message - assert "accept_len_hist=2:1,4:1" in message - assert "queue_len_hist=0:1" in message - assert "draft_tokens=13" in message - assert "accepted_draft_tokens=4" in message - assert "sampled_tokens=6 emitted_tokens=6 finished_reqs=0" in message diff --git a/vllm/envs.py b/vllm/envs.py index 7c04a3082a..4026c4f0f7 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -250,8 +250,6 @@ VLLM_SM70_ASYNC_STAGED_INPUT_PREP: bool = False VLLM_SM70_ASYNC_CPU_TRACE: bool = False VLLM_SM70_ASYNC_CPU_TRACE_EVERY: int = 16 - VLLM_SM70_DFLASH2_BATCH_TRACE: bool = False - VLLM_SM70_DFLASH2_BATCH_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 @@ -2188,15 +2186,6 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_ASYNC_CPU_TRACE_EVERY": lambda: max( 1, int(os.getenv("VLLM_SM70_ASYNC_CPU_TRACE_EVERY", "16")) ), - # Completed-iteration trace for DFlash2 concurrency diagnostics. This only - # reads scheduler/model outputs that are already CPU-visible and never adds - # a CUDA event or synchronization to the measured path. - "VLLM_SM70_DFLASH2_BATCH_TRACE": lambda: bool( - int(os.getenv("VLLM_SM70_DFLASH2_BATCH_TRACE", "0")) - ), - "VLLM_SM70_DFLASH2_BATCH_TRACE_EVERY": lambda: max( - 1, int(os.getenv("VLLM_SM70_DFLASH2_BATCH_TRACE_EVERY", "16")) - ), # Legacy 0.0.3 diagnostic gate. Logs one TP all-reduce backend decision per # group/backend/shape/dtype so route-hit data can distinguish custom AR, # pynccl, flashinfer, symmetric-memory, and torch fallback paths. diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 15010d94f3..9210b37eec 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -201,9 +201,6 @@ def __init__( self.batch_queue_size, ) self._sm70_async_cpu_trace_step = 0 - self._sm70_dflash2_batch_trace_enabled = envs.VLLM_SM70_DFLASH2_BATCH_TRACE - self._sm70_dflash2_batch_trace_every = envs.VLLM_SM70_DFLASH2_BATCH_TRACE_EVERY - self._reset_sm70_dflash2_batch_trace() self.is_ec_consumer = ( vllm_config.ec_transfer_config is None @@ -451,137 +448,6 @@ def log_iteration_details(self, scheduler_output: SchedulerOutput | None): ) self._iteration_index += 1 - def _reset_sm70_dflash2_batch_trace(self) -> None: - self._sm70_dflash2_trace_steps = 0 - self._sm70_dflash2_trace_started_at = 0.0 - self._sm70_dflash2_trace_phase_hist: defaultdict[str, int] = defaultdict(int) - self._sm70_dflash2_trace_gen_req_hist: defaultdict[int, int] = defaultdict(int) - self._sm70_dflash2_trace_target_m_hist: defaultdict[int, int] = defaultdict(int) - self._sm70_dflash2_trace_accept_len_hist: defaultdict[int, int] = defaultdict( - int - ) - self._sm70_dflash2_trace_queue_len_hist: defaultdict[int, int] = defaultdict( - int - ) - self._sm70_dflash2_trace_unfinished_hist: defaultdict[int, int] = defaultdict( - int - ) - self._sm70_dflash2_trace_scheduled_gen_tokens = 0 - self._sm70_dflash2_trace_draft_tokens = 0 - self._sm70_dflash2_trace_accepted_draft_tokens = 0 - self._sm70_dflash2_trace_sampled_tokens = 0 - self._sm70_dflash2_trace_emitted_tokens = 0 - self._sm70_dflash2_trace_finished_reqs = 0 - - @staticmethod - def _format_sm70_dflash2_trace_hist(hist: dict[Any, int]) -> str: - return ",".join(f"{key}:{hist[key]}" for key in sorted(hist)) or "none" - - def _record_sm70_dflash2_batch_trace( - self, - scheduler_output: SchedulerOutput, - model_output: ModelRunnerOutput, - engine_core_outputs: dict[int, EngineCoreOutputs], - queue_len: int, - ) -> None: - """Aggregate completed scheduler iterations without synchronizing CUDA.""" - if not self._sm70_dflash2_batch_trace_enabled: - return - - now = time.perf_counter() - if self._sm70_dflash2_trace_steps == 0: - self._sm70_dflash2_trace_started_at = now - - details = compute_iteration_details(scheduler_output) - if details.num_ctx_requests and details.num_generation_requests: - phase = "mixed" - elif details.num_ctx_requests: - phase = "prefill" - elif details.num_generation_requests: - phase = "decode" - else: - phase = "empty" - - invalid_spec_tokens = scheduler_output.num_invalid_spec_tokens or {} - scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens - valid_draft_tokens = sum( - max(0, len(token_ids) - invalid_spec_tokens.get(req_id, 0)) - for req_id, token_ids in scheduled_spec_tokens.items() - ) - accepted_draft_tokens = 0 - sampled_tokens = 0 - for req_id, req_index in model_output.req_id_to_index.items(): - generated_tokens = ( - model_output.sampled_token_ids[req_index] - if model_output.sampled_token_ids - else [] - ) - sampled_tokens += len(generated_tokens) - if req_id in scheduled_spec_tokens: - accept_len = len(generated_tokens) - self._sm70_dflash2_trace_accept_len_hist[accept_len] += 1 - accepted_draft_tokens += max(0, accept_len - 1) - - emitted_tokens = 0 - finished_reqs = 0 - for client_output in engine_core_outputs.values(): - for request_output in client_output.outputs: - emitted_tokens += len(request_output.new_token_ids) - finished_reqs += int(request_output.finished) - - unfinished_reqs = self.scheduler.get_num_unfinished_requests() - self._sm70_dflash2_trace_steps += 1 - self._sm70_dflash2_trace_phase_hist[phase] += 1 - self._sm70_dflash2_trace_gen_req_hist[details.num_generation_requests] += 1 - self._sm70_dflash2_trace_target_m_hist[details.num_generation_tokens] += 1 - self._sm70_dflash2_trace_queue_len_hist[queue_len] += 1 - self._sm70_dflash2_trace_unfinished_hist[unfinished_reqs] += 1 - self._sm70_dflash2_trace_scheduled_gen_tokens += details.num_generation_tokens - self._sm70_dflash2_trace_draft_tokens += valid_draft_tokens - self._sm70_dflash2_trace_accepted_draft_tokens += accepted_draft_tokens - self._sm70_dflash2_trace_sampled_tokens += sampled_tokens - self._sm70_dflash2_trace_emitted_tokens += emitted_tokens - self._sm70_dflash2_trace_finished_reqs += finished_reqs - - flush = ( - self._sm70_dflash2_trace_steps >= self._sm70_dflash2_batch_trace_every - or (unfinished_reqs == 0 and queue_len == 0) - ) - if not flush: - return - - completion_span_ms = (now - self._sm70_dflash2_trace_started_at) * 1000.0 - logger.info( - "SM70 DFlash2 batch trace steps=%d completion_span_ms=%.3f " - "phase_hist=%s gen_req_hist=%s target_m_hist=%s " - "accept_len_hist=%s queue_len_hist=%s unfinished_hist=%s " - "scheduled_gen_tokens=%d draft_tokens=%d accepted_draft_tokens=%d " - "sampled_tokens=%d emitted_tokens=%d finished_reqs=%d", - self._sm70_dflash2_trace_steps, - completion_span_ms, - self._format_sm70_dflash2_trace_hist(self._sm70_dflash2_trace_phase_hist), - self._format_sm70_dflash2_trace_hist(self._sm70_dflash2_trace_gen_req_hist), - self._format_sm70_dflash2_trace_hist( - self._sm70_dflash2_trace_target_m_hist - ), - self._format_sm70_dflash2_trace_hist( - self._sm70_dflash2_trace_accept_len_hist - ), - self._format_sm70_dflash2_trace_hist( - self._sm70_dflash2_trace_queue_len_hist - ), - self._format_sm70_dflash2_trace_hist( - self._sm70_dflash2_trace_unfinished_hist - ), - self._sm70_dflash2_trace_scheduled_gen_tokens, - self._sm70_dflash2_trace_draft_tokens, - self._sm70_dflash2_trace_accepted_draft_tokens, - self._sm70_dflash2_trace_sampled_tokens, - self._sm70_dflash2_trace_emitted_tokens, - self._sm70_dflash2_trace_finished_reqs, - ) - self._reset_sm70_dflash2_batch_trace() - def step(self) -> tuple[dict[int, EngineCoreOutputs], bool]: """Schedule, execute, and make output. @@ -642,9 +508,6 @@ def step(self) -> tuple[dict[int, EngineCoreOutputs], bool]: engine_core_outputs = self.scheduler.update_from_output( scheduler_output, model_output ) - self._record_sm70_dflash2_batch_trace( - scheduler_output, model_output, engine_core_outputs, queue_len=0 - ) if profile_ddtree_engine: profile_update_ms = (time.perf_counter() - profile_part_t0) * 1000.0 logger.info( @@ -843,12 +706,6 @@ def step_with_batch_queue( engine_core_outputs = self.scheduler.update_from_output( scheduler_output, model_output ) - self._record_sm70_dflash2_batch_trace( - scheduler_output, - model_output, - engine_core_outputs, - queue_len=len(batch_queue), - ) if trace_log: trace_update_ms = (time.perf_counter() - trace_update_t0) * 1000.0 diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 187d9408aa..c006e980f8 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -7,7 +7,6 @@ import torch import torch.nn as nn -from vllm import envs from vllm.config import VllmConfig, replace from vllm.config.compilation import CUDAGraphMode from vllm.config.speculative import get_dflash_model_draft_tokens @@ -134,63 +133,6 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.draft_kv_cache_group_id: int = -1 self._context_only_prefill_logged = False self._query_slot_mappings: torch.Tensor | None = None - self._stage_profile_records: list[ - tuple[int, int, tuple[torch.cuda.Event, ...]] - ] = [] - - def _flush_stage_profile( - self, - num_reqs: int, - num_target_tokens: int, - events: tuple[torch.cuda.Event, ...] | None, - ) -> None: - """Report batched stream timing without synchronizing every decode step.""" - if events is None: - return - - self._stage_profile_records.append((num_reqs, num_target_tokens, events)) - interval = envs.VLLM_DFLASH_PROFILE_LOG_INTERVAL - if len(self._stage_profile_records) < interval: - return - - # One synchronization closes the whole window. Event-to-event elapsed - # time includes any exposed host launch gap on the current CUDA stream, - # which is exactly what this diagnostic is intended to find. - self._stage_profile_records[-1][2][-1].synchronize() - stage_names = ("hidden", "prepare", "context_kv", "metadata", "query_graph") - stage_values = {name: [] for name in stage_names} - total_values: list[float] = [] - reqs = 0 - target_tokens = 0 - for record_reqs, record_tokens, record_events in self._stage_profile_records: - reqs += record_reqs - target_tokens += record_tokens - for name, start, end in zip( - stage_names, record_events[:-1], record_events[1:] - ): - stage_values[name].append(start.elapsed_time(end)) - total_values.append(record_events[0].elapsed_time(record_events[-1])) - - rounds = len(self._stage_profile_records) - means = { - name: sum(values) / len(values) for name, values in stage_values.items() - } - total_mean = sum(total_values) / len(total_values) - logger.info( - "DFLASH_STAGE_PROFILE rounds=%d avg_reqs=%.2f avg_target_tokens=%.2f " - "hidden_ms=%.3f prepare_ms=%.3f context_kv_ms=%.3f " - "metadata_ms=%.3f query_graph_ms=%.3f total_ms=%.3f", - rounds, - reqs / rounds, - target_tokens / rounds, - means["hidden"], - means["prepare"], - means["context_kv"], - means["metadata"], - means["query_graph"], - total_mean, - ) - self._stage_profile_records.clear() @property def attn_vllm_config(self) -> VllmConfig: @@ -480,17 +422,6 @@ def propose( max_seq_len + self.num_query_per_req, self.max_model_len ) self._prepare_proposal_runtime(input_batch, num_sampled, num_rejected) - profile_events = None - if ( - envs.VLLM_DFLASH_PROFILE - and self.device.type == "cuda" - and not dummy_run - and not np.any(input_batch.is_prefilling_np[:num_reqs]) - ): - profile_events = tuple( - torch.cuda.Event(enable_timing=True) for _ in range(6) - ) - profile_events[0].record() # NOTE: To avoid CPU-GPU synchronization without CPU knowing the # number of rejected tokens, we maintain the size of input_ids and @@ -507,8 +438,6 @@ def propose( self.hidden_states[:num_target_tokens].copy_( hidden_states[:num_target_tokens] ) - if profile_events is not None: - profile_events[1].record() if dummy_run and skip_attn_for_dummy_run: # Memory profiling path: block_tables / kv_cache_config are not initialized. @@ -566,8 +495,6 @@ def propose( self.max_model_len, self.sample_from_anchor, ) - if profile_events is not None: - profile_events[2].record() # 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 @@ -588,8 +515,6 @@ def propose( self.context_positions[:num_target_tokens], context_slots, ) - if profile_events is not None: - profile_events[3].record() if not dummy_run and _is_context_only_prefill(input_batch): # Intermediate chunked-prefill steps only need to materialize the @@ -646,8 +571,6 @@ def propose( query_slot_mappings[:, :num_tokens_padded], self.kv_cache_config, ) - if profile_events is not None: - profile_events[4].record() if batch_desc.cg_mode == CUDAGraphMode.FULL: assert self.query_cudagraph_manager is not None @@ -661,11 +584,8 @@ def propose( num_tokens_across_dp=num_tokens_across_dp, cudagraph_runtime_mode=batch_desc.cg_mode, ) - if profile_events is not None: - profile_events[5].record() self._apply_ngram_assist(num_reqs) - self._flush_stage_profile(num_reqs, num_target_tokens, profile_events) return self.draft_tokens[:num_reqs]